Merge pull request #127 from whywaita/feat/metrics-rate-limit

Add rate limit metrics

pkg/gh/github.go

@@ -8,6 +8,7 @@ import (
 	"net/url"
 	"path"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/bradleyfalzon/ghinstallation/v2"
@@ -25,11 +26,19 @@ var (
 
 	// ResponseCache is cache variable
 	responseCache *cache.Cache
+
+	// rateLimitRemain is remaining of Rate limit, for metrics
+	rateLimitRemain = sync.Map{}
+	// rateLimitLimit is limit of Rate limit, for metrics
+	rateLimitLimit = sync.Map{}
 )
 
 func init() {
 	c := cache.New(5*time.Minute, 10*time.Minute)
 	responseCache = c
+
+	rateLimitRemain = sync.Map{}
+	rateLimitLimit = sync.Map{}
 }
 
 // NewClient create a client of GitHub
@@ -170,6 +179,7 @@ func ListRunners(ctx context.Context, client *github.Client, owner, repo string)
 		if err != nil {
 			return nil, fmt.Errorf("failed to list runners: %w", err)
 		}
+		storeRateLimit(getRateLimitKey(owner, repo), resp.Rate)
 
 		rs = append(rs, runners.Runners...)
 		if resp.NextPage == 0 {

pkg/gh/jwt.go

@@ -22,11 +22,12 @@ var (
 
 // GenerateGitHubAppsToken generate token of GitHub Apps using private key
 // clientApps needs to response of `NewClientGitHubApps()`
-func GenerateGitHubAppsToken(ctx context.Context, clientApps *github.Client, installationID int64) (string, *time.Time, error) {
-	token, _, err := clientApps.Apps.CreateInstallationToken(ctx, installationID, nil)
+func GenerateGitHubAppsToken(ctx context.Context, clientApps *github.Client, installationID int64, scope string) (string, *time.Time, error) {
+	token, resp, err := clientApps.Apps.CreateInstallationToken(ctx, installationID, nil)
 	if err != nil {
 		return "", nil, fmt.Errorf("failed to generate token from API: %w", err)
 	}
+	storeRateLimit(scope, resp.Rate)
 	return *token.Token, token.ExpiresAt, nil
 }
 
@@ -97,7 +98,7 @@ func IsInstalledGitHubApp(ctx context.Context, gheDomain, inputScope string) (in
 }
 
 func isInstalledGitHubAppSelected(ctx context.Context, gheDomain, inputScope string, installationID int64) error {
-	lr, err := GHlistAppsInstalledRepo(ctx, gheDomain, installationID)
+	lr, err := GHlistAppsInstalledRepo(ctx, gheDomain, installationID, inputScope)
 	if err != nil {
 		return fmt.Errorf("failed to get list of installed repositories: %w", err)
 	}
@@ -123,12 +124,12 @@ func isInstalledGitHubAppSelected(ctx context.Context, gheDomain, inputScope str
 	return fmt.Errorf("not found")
 }
 
-func listAppsInstalledRepo(ctx context.Context, gheDomain string, installationID int64) (*github.ListRepositories, error) {
+func listAppsInstalledRepo(ctx context.Context, gheDomain string, installationID int64, inputScope string) (*github.ListRepositories, error) {
 	clientApps, err := NewClientGitHubApps(gheDomain, config.Config.GitHub.AppID, config.Config.GitHub.PEMByte)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create github.Client from installationID: %w", err)
 	}
-	token, _, err := GenerateGitHubAppsToken(ctx, clientApps, installationID)
+	token, _, err := GenerateGitHubAppsToken(ctx, clientApps, installationID, inputScope)
 	if err != nil {
 		return nil, fmt.Errorf("failed to generate GitHub Apps Token: %w", err)
 	}

pkg/gh/jwt_test.go

@@ -49,7 +49,7 @@ func setStubFunctions() {
 		}, nil
 	}
 
-	GHlistAppsInstalledRepo = func(ctx context.Context, gheDomain string, installationID int64) (*github.ListRepositories, error) {
+	GHlistAppsInstalledRepo = func(ctx context.Context, gheDomain string, installationID int64, scope string) (*github.ListRepositories, error) {
 		total := 1
 		fullName1 := "example-selected/sample-registered"
 		return &github.ListRepositories{

pkg/gh/ratelimit.go

@@ -0,0 +1,75 @@
+package gh
+
+import (
+	"fmt"
+	"sync"
+
+	"github.com/google/go-github/v35/github"
+)
+
+func storeRateLimit(scope string, rateLimit github.Rate) {
+	if rateLimit.Reset.IsZero() {
+		// Not configure rate limit, don't need to store (e.g. GHES)
+		return
+	}
+
+	rateLimitLimit.Store(scope, rateLimit.Limit)
+	rateLimitRemain.Store(scope, rateLimit.Remaining)
+}
+
+func getRateLimitKey(org, repo string) string {
+	if repo == "" {
+		return org
+	}
+	return fmt.Sprintf("%s/%s", org, repo)
+}
+
+// GetRateLimitRemain get a list of rate limit remaining
+// key: scope, value: remain
+func GetRateLimitRemain() map[string]int {
+	m := map[string]int{}
+	mu := sync.Mutex{}
+
+	rateLimitRemain.Range(func(key, value interface{}) bool {
+		k, ok := key.(string)
+		if !ok {
+			return false
+		}
+		v, ok := value.(int)
+		if !ok {
+			return false
+		}
+
+		mu.Lock()
+		m[k] = v
+		mu.Unlock()
+		return true
+	})
+
+	return m
+}
+
+// GetRateLimitLimit get a list of rate limit
+// key: scope, value: remain
+func GetRateLimitLimit() map[string]int {
+	m := map[string]int{}
+	mu := sync.Mutex{}
+
+	rateLimitLimit.Range(func(key, value interface{}) bool {
+		k, ok := key.(string)
+		if !ok {
+			return false
+		}
+		v, ok := value.(int)
+		if !ok {
+			return false
+		}
+
+		mu.Lock()
+		m[k] = v
+		mu.Unlock()
+		return true
+	})
+
+	return m
+}

pkg/metric/scrape_memory.go

@@ -8,6 +8,7 @@ import (
 
 	"github.com/whywaita/myshoes/internal/config"
 	"github.com/whywaita/myshoes/pkg/datastore"
+	"github.com/whywaita/myshoes/pkg/gh"
 	"github.com/whywaita/myshoes/pkg/starter"
 )
 
@@ -29,6 +30,16 @@ var (
 		"waiting queue in starter",
 		[]string{"starter"}, nil,
 	)
+	memoryGitHubRateLimitRemaining = prometheus.NewDesc(
+		prometheus.BuildFQName(namespace, memoryName, "github_rate_limit_remaining"),
+		"The number of rate limit remaining",
+		[]string{"scope"}, nil,
+	)
+	memoryGitHubRateLimitLimiting = prometheus.NewDesc(
+		prometheus.BuildFQName(namespace, memoryName, "github_rate_limit_limiting"),
+		"The number of rate limit max",
+		[]string{"scope"}, nil,
+	)
 )
 
 // ScraperMemory is scraper implement for memory
@@ -49,6 +60,9 @@ func (ScraperMemory) Scrape(ctx context.Context, ds datastore.Datastore, ch chan
 	if err := scrapeStarterValues(ch); err != nil {
 		return fmt.Errorf("failed to scrape starter values: %w", err)
 	}
+	if err := scrapeGitHubValues(ch); err != nil {
+		return fmt.Errorf("failed to scrape GitHub values: %w", err)
+	}
 
 	return nil
 }
@@ -72,4 +86,22 @@ func scrapeStarterValues(ch chan<- prometheus.Metric) error {
 	return nil
 }
 
+func scrapeGitHubValues(ch chan<- prometheus.Metric) error {
+	rateLimitRemain := gh.GetRateLimitRemain()
+	for scope, remain := range rateLimitRemain {
+		ch <- prometheus.MustNewConstMetric(
+			memoryGitHubRateLimitRemaining, prometheus.GaugeValue, float64(remain), scope,
+		)
+	}
+
+	rateLimitLimit := gh.GetRateLimitLimit()
+	for scope, limit := range rateLimitLimit {
+		ch <- prometheus.MustNewConstMetric(
+			memoryGitHubRateLimitLimiting, prometheus.GaugeValue, float64(limit), scope,
+		)
+	}
+
+	return nil
+}
+
 var _ Scraper = ScraperMemory{}

pkg/runner/token_update.go

@@ -38,7 +38,7 @@ func (m *Manager) doTargetToken(ctx context.Context) error {
 			logger.Logf(false, "failed to create client of GitHub installation: %+v", err)
 			continue
 		}
-		token, expiredAt, err := gh.GenerateGitHubAppsToken(ctx, clientApps, installationID)
+		token, expiredAt, err := gh.GenerateGitHubAppsToken(ctx, clientApps, installationID, target.Scope)
 		if err != nil {
 			logger.Logf(false, "failed to get Apps Token: %+v", err)
 			continue

pkg/web/target_create.go

@@ -44,7 +44,7 @@ func handleTargetCreate(w http.ResponseWriter, r *http.Request, ds datastore.Dat
 		outputErrorMsg(w, http.StatusInternalServerError, "failed to client GitHub Apps")
 		return
 	}
-	token, expiredAt, err := GHGenerateGitHubAppsToken(ctx, clientApps, installationID)
+	token, expiredAt, err := GHGenerateGitHubAppsToken(ctx, clientApps, installationID, inputTarget.Scope)
 	if err != nil {
 		logger.Logf(false, "failed to generate GitHub Apps Token: %+v", err)
 		outputErrorMsg(w, http.StatusInternalServerError, "failed to generate GitHub Apps token")

pkg/web/target_test.go

@@ -47,7 +47,7 @@ func setStubFunctions() {
 		return testInstallationID, nil
 	}
 
-	web.GHGenerateGitHubAppsToken = func(ctx context.Context, clientInstallation *github.Client, installationID int64) (string, *time.Time, error) {
+	web.GHGenerateGitHubAppsToken = func(ctx context.Context, clientInstallation *github.Client, installationID int64, scope string) (string, *time.Time, error) {
 		return testGitHubAppToken, &testTime, nil
 	}