Add rate limit publisher. (#84)

server/scheduled_executor_service.go

@@ -1,6 +1,7 @@
 package server
 
 import (
+	"context"
 	"io"
 	"os"
 	"os/signal"
@@ -9,6 +10,7 @@ import (
 	"text/template"
 	"time"
 
+	"github.com/google/go-github/v31/github"
 	stats "github.com/lyft/gostats"
 	"github.com/runatlantis/atlantis/server/events"
 	"github.com/runatlantis/atlantis/server/events/metrics"
@@ -20,36 +22,54 @@ type ScheduledExecutorService struct {
 	log logging.SimpleLogging
 
 	// jobs
-	garbageCollector CronDefinition
+	garbageCollector   JobDefinition
+	rateLimitPublisher JobDefinition
 }
 
 func NewScheduledExecutorService(
 	workingDirIterator events.WorkDirIterator,
 	statsScope stats.Scope,
 	log logging.SimpleLogging,
 	pullCleaner events.PullCleaner,
+	githubClient *github.Client,
 ) *ScheduledExecutorService {
+
+	scheduledScope := statsScope.Scope("scheduled")
 	garbageCollector := &GarbageCollector{
 		workingDirIterator: workingDirIterator,
-		stats:              statsScope.Scope("scheduled.garbagecollector"),
+		stats:              scheduledScope.Scope("garbagecollector"),
 		log:                log,
 		cleaner:            pullCleaner,
 	}
 
-	garbageCollectorCron := CronDefinition{
+	garbageCollectorJob := JobDefinition{
 		Job: garbageCollector,
 
 		// 5 minutes should probably be the lowest to prevent GH rate limits
 		Period: 5 * time.Minute,
 	}
 
+	rateLimitPublisher := &RateLimitStatsPublisher{
+		client: githubClient,
+		stats:  scheduledScope.Scope("ratelimitpublisher"),
+		log:    log,
+	}
+
+	rateLimitPublisherJob := JobDefinition{
+		Job: rateLimitPublisher,
+
+		// since rate limit api doesn't contribute to the rate limit we can call this every minute
+		Period: 1 * time.Minute,
+	}
+
 	return &ScheduledExecutorService{
-		log:              log,
-		garbageCollector: garbageCollectorCron,
+		log:                log,
+		garbageCollector:   garbageCollectorJob,
+		rateLimitPublisher: rateLimitPublisherJob,
 	}
 }
 
-type CronDefinition struct {
+type JobDefinition struct {
 	Job    Job
 	Period time.Duration
 }
@@ -61,6 +81,9 @@ func (s *ScheduledExecutorService) Run() {
 	garbageCollectorTicker := time.NewTicker(s.garbageCollector.Period)
 	defer garbageCollectorTicker.Stop()
 
+	rateLimitPublisherTicker := time.NewTicker(s.rateLimitPublisher.Period)
+	defer rateLimitPublisherTicker.Stop()
+
 	interrupt := make(chan os.Signal, 1)
 
 	// Stop on SIGINTs and SIGTERMs.
@@ -73,6 +96,8 @@ func (s *ScheduledExecutorService) Run() {
 			return
 		case <-garbageCollectorTicker.C:
 			go s.garbageCollector.Job.Run()
+		case <-rateLimitPublisherTicker.C:
+			go s.rateLimitPublisher.Job.Run()
 		}
 	}
 }
@@ -81,6 +106,38 @@ type Job interface {
 	Run()
 }
 
+type RateLimitStatsPublisher struct {
+	log    logging.SimpleLogging
+	stats  stats.Scope
+	client *github.Client
+}
+
+func (r *RateLimitStatsPublisher) Run() {
+
+	// since we are calling this at timed intervals it's probably ok to create our own context here.
+	// it's not critical to cancel this context
+	ctx := context.Background()
+
+	errCounter := r.stats.NewCounter(metrics.ExecutionErrorMetric)
+	rateLimitRemainingCounter := r.stats.NewCounter("ratelimitremaining")
+
+	rateLimits, resp, err := r.client.RateLimits(ctx)
+
+	if err != nil {
+		r.log.Err("error retrieving rate limits: %s", err)
+		errCounter.Inc()
+		return
+	}
+
+	if resp.StatusCode != 200 {
+		r.log.Err("error retrieving rate limits: %s", resp.Status)
+		errCounter.Inc()
+		return
+	}
+
+	rateLimitRemainingCounter.Add(uint64(rateLimits.GetCore().Remaining))
+}
+
 var gcStaleClosedPullTemplate = template.Must(template.New("").Parse(
 	"Pull Request has been closed for 30 days. Atlantis GC has deleted the locks and plans for the following projects and workspaces:\n" +
 		"{{ range . }}\n" +

server/server.go

@@ -32,6 +32,7 @@ import (
 	"syscall"
 	"time"
 
+	"github.com/google/go-github/v31/github"
 	"github.com/mitchellh/go-homedir"
 	"github.com/runatlantis/atlantis/server/core/db"
 	"github.com/runatlantis/atlantis/server/events/yaml/valid"
@@ -142,6 +143,11 @@ func NewServer(userConfig UserConfig, config Config) (*Server, error) {
 	statsScope.Store().AddStatGenerator(stats.NewRuntimeStats(statsScope.Scope("go")))
 
 	var supportedVCSHosts []models.VCSHostType
+
+	// not to be used directly, currently this is just used
+	// for reporting rate limits
+	var rawGithubClient *github.Client
+
 	var githubClient vcs.IGithubClient
 	var githubAppEnabled bool
 	var githubCredentials vcs.GithubCredentials
@@ -696,6 +702,7 @@ func NewServer(userConfig UserConfig, config Config) (*Server, error) {
 			// using a specific template to signal that this is from an async process
 			PullClosedTemplate: &GCStalePullTemplate{},
 		},
+		rawGithubClient,
 	)
 
 	return &Server{