-
Notifications
You must be signed in to change notification settings - Fork 6
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add clientcreatorpool abstraction for utilizing multiple github apps …
…for extra ratelimit (#756)
- Loading branch information
Showing
2 changed files
with
154 additions
and
0 deletions.
There are no files selected for viewing
80 changes: 80 additions & 0 deletions
80
server/neptune/temporalworker/clientcreatorpool/clientcreatorpool.go
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,80 @@ | ||
package clientcreatorpool | ||
|
||
import ( | ||
"fmt" | ||
|
||
"github.com/palantir/go-githubapp/githubapp" | ||
"github.com/pkg/errors" | ||
ghClient "github.com/runatlantis/atlantis/server/neptune/workflows/activities/github" | ||
"github.com/uber-go/tally/v4" | ||
) | ||
|
||
type ClientCreatorPool struct { | ||
// keys are app ids aka integration ids | ||
// Note: It would make a lot more sense to use something like a name, or the slug, right? | ||
// The reason why that is not the case, is that then we would have to pass those around. We can't | ||
// modify the clientCreator (it is part of the githubapp library), and you can see that inside, its | ||
// fields are private, and there is no way to associate a given clientCreator with a name. There is | ||
// no slug inside the githubapp.Config, only private keys and app ids, and we don't want to use | ||
// private keys as keys. | ||
idToClientCreator map[int]githubapp.ClientCreator | ||
idToRateLimitRemaning map[int]int | ||
// Note that integration id is NOT installation id. Those are 2 separate things. | ||
} | ||
|
||
type ClientCreatorPoolConfig struct { | ||
id int | ||
config githubapp.Config | ||
} | ||
|
||
// This function is only called once, when the server starts | ||
func (c *ClientCreatorPool) Initialize(configs []ClientCreatorPoolConfig, scope tally.Scope) error { | ||
c.idToClientCreator = make(map[int]githubapp.ClientCreator) | ||
c.idToRateLimitRemaning = make(map[int]int) | ||
|
||
for _, config := range configs { | ||
t := fmt.Sprintf("github.app.%d", config.id) | ||
clientCreator, err := githubapp.NewDefaultCachingClientCreator( | ||
config.config, | ||
githubapp.WithClientMiddleware( | ||
// combine the id with app | ||
ghClient.ClientMetrics(scope.SubScope(t)), | ||
)) | ||
if err != nil { | ||
return errors.Wrap(err, "client creator") | ||
} | ||
c.idToClientCreator[config.id] = clientCreator | ||
// For the rate limit remaining, initially needs to be non-zero, the actual rate limit number value will be | ||
// updated within 60 seconds by the cron that checks the rate limit | ||
c.idToRateLimitRemaning[config.id] = 100 | ||
} | ||
|
||
return nil | ||
} | ||
|
||
func (c *ClientCreatorPool) GetClientCreatorWithMaxRemainingRateLimit() (githubapp.ClientCreator, error) { | ||
maxSeenSoFar := 0 | ||
theId := 0 | ||
for id, num := range c.idToRateLimitRemaning { | ||
if num > maxSeenSoFar { | ||
maxSeenSoFar = num | ||
theId = id | ||
} | ||
} | ||
|
||
clientCreator, ok := c.idToClientCreator[theId] | ||
if !ok { | ||
return nil, errors.New("client creator not found") | ||
} | ||
|
||
return clientCreator, nil | ||
} | ||
|
||
// this func will be used in the crons to update the rate limit remaining | ||
func (c *ClientCreatorPool) SetRateLimitRemaining(id int, remaining int) { | ||
c.idToRateLimitRemaning[id] = remaining | ||
} | ||
|
||
func (c *ClientCreatorPool) GetRateLimitRemaining(id int) int { | ||
return c.idToRateLimitRemaning[id] | ||
} |
74 changes: 74 additions & 0 deletions
74
server/neptune/temporalworker/clientcreatorpool/clientcreatorpool_test.go
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,74 @@ | ||
package clientcreatorpool | ||
|
||
import ( | ||
"testing" | ||
|
||
"github.com/palantir/go-githubapp/githubapp" | ||
"github.com/runatlantis/atlantis/server/logging" | ||
"github.com/runatlantis/atlantis/server/metrics" | ||
"github.com/stretchr/testify/assert" | ||
) | ||
|
||
func initialize(t *testing.T) ClientCreatorPool { | ||
configs := []ClientCreatorPoolConfig{ | ||
{ | ||
id: 564754, | ||
config: githubapp.Config{}, | ||
}, | ||
{ | ||
id: 243643, | ||
config: githubapp.Config{}, | ||
}, | ||
} | ||
|
||
configs[0].config.App.IntegrationID = 564754 | ||
configs[0].config.App.PrivateKey = "key1" | ||
configs[0].config.App.WebhookSecret = "secret1" | ||
|
||
configs[1].config.App.IntegrationID = 243643 | ||
configs[1].config.App.PrivateKey = "key2" | ||
configs[1].config.App.WebhookSecret = "secret2" | ||
|
||
c := ClientCreatorPool{} | ||
ctxLogger := logging.NewNoopCtxLogger(t) | ||
scope, _, _ := metrics.NewLoggingScope(ctxLogger, "null") | ||
_ = c.Initialize(configs, scope) | ||
return c | ||
} | ||
|
||
func TestInitialize(t *testing.T) { | ||
configs := []ClientCreatorPoolConfig{ | ||
{ | ||
id: 1, | ||
config: githubapp.Config{}, | ||
}, | ||
{ | ||
id: 2, | ||
config: githubapp.Config{}, | ||
}, | ||
} | ||
|
||
configs[0].config.App.IntegrationID = 1 | ||
configs[0].config.App.PrivateKey = "key1" | ||
configs[0].config.App.WebhookSecret = "secret1" | ||
|
||
configs[1].config.App.IntegrationID = 2 | ||
configs[1].config.App.PrivateKey = "key2" | ||
configs[1].config.App.WebhookSecret = "secret2" | ||
|
||
c := ClientCreatorPool{} | ||
ctxLogger := logging.NewNoopCtxLogger(t) | ||
scope, _, _ := metrics.NewLoggingScope(ctxLogger, "null") | ||
err := c.Initialize(configs, scope) | ||
assert.NoError(t, err) | ||
} | ||
|
||
func TestGetClientCreatorWithMaxRemainingRateLimit(t *testing.T) { | ||
c := initialize(t) | ||
c.SetRateLimitRemaining(564754, 9000) | ||
clientCreator, err := c.GetClientCreatorWithMaxRemainingRateLimit() | ||
assert.NoError(t, err) | ||
assert.NotNil(t, clientCreator) | ||
assert.Equal(t, clientCreator, c.idToClientCreator[564754]) | ||
assert.Equal(t, 9000, c.GetRateLimitRemaining(564754)) | ||
} |