diff --git a/server/neptune/temporalworker/clientcreatorpool/clientcreatorpool.go b/server/neptune/temporalworker/clientcreatorpool/clientcreatorpool.go new file mode 100644 index 000000000..1a72dd0ed --- /dev/null +++ b/server/neptune/temporalworker/clientcreatorpool/clientcreatorpool.go @@ -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] +} diff --git a/server/neptune/temporalworker/clientcreatorpool/clientcreatorpool_test.go b/server/neptune/temporalworker/clientcreatorpool/clientcreatorpool_test.go new file mode 100644 index 000000000..5e0f529ea --- /dev/null +++ b/server/neptune/temporalworker/clientcreatorpool/clientcreatorpool_test.go @@ -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)) +}