Skip to content

Commit

Permalink
Add clientcreatorpool abstraction for utilizing multiple github apps …
Browse files Browse the repository at this point in the history
…for extra ratelimit (#756)
  • Loading branch information
smonero authored May 21, 2024
1 parent b007de4 commit 54a9851
Show file tree
Hide file tree
Showing 2 changed files with 154 additions and 0 deletions.
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]
}
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))
}

0 comments on commit 54a9851

Please sign in to comment.