Merge pull request #41 from mcarmonaa/improvement/support-multi-org

Improvement/support multi org

Author: mcarmonaa

cmd/gitcollector/subcmd/download.go

@@ -5,6 +5,8 @@ import (
 	"io/ioutil"
 	"os"
 	"runtime"
+	"strings"
+	"sync"
 	"time"
 
 	"github.com/src-d/gitcollector"
@@ -27,7 +29,7 @@ type DownloadCmd struct {
 	TmpPath         string `long:"tmp" description:"directory to place generated temporal files" default:"/tmp" env:"GITCOLLECTOR_TMP"`
 	Workers         int    `long:"workers" description:"number of workers, default to GOMAXPROCS" env:"GITCOLLECTOR_WORKERS"`
 	NotAllowUpdates bool   `long:"no-updates" description:"don't allow updates on already downloaded repositories" env:"GITCOLLECTOR_NO_UPDATES"`
-	Org             string `long:"org" env:"GITHUB_ORGANIZATION" description:"github organization" required:"true"`
+	Orgs            string `long:"orgs" env:"GITHUB_ORGANIZATIONS" description:"list of github organization names separated by comma" required:"true"`
 	Token           string `long:"token" env:"GITHUB_TOKEN" description:"github token"`
 	MetricsDBURI    string `long:"metrics-db" env:"GITCOLLECTOR_METRICS_DB_URI" description:"uri to a database where metrics will be sent"`
 	MetricsDBTable  string `long:"metrics-db-table" env:"GITCOLLECTOR_METRICS_DB_TABLE" default:"gitcollector_metrics" description:"table name where the metrics will be added"`
@@ -36,6 +38,10 @@ type DownloadCmd struct {
 
 // Execute runs the command.
 func (c *DownloadCmd) Execute(args []string) error {
+	start := time.Now()
+
+	orgs := strings.Split(c.Orgs, ",")
+
 	info, err := os.Stat(c.LibPath)
 	check(err, "wrong path to locate the library")
 
@@ -73,7 +79,9 @@ func (c *DownloadCmd) Execute(args []string) error {
 	authTokens := map[string]string{}
 	if c.Token != "" {
 		log.Debugf("acces token found")
-		authTokens[c.Org] = c.Token
+		for _, org := range orgs {
+			authTokens[org] = c.Token
+		}
 	}
 
 	workers := c.Workers
@@ -103,16 +111,12 @@ func (c *DownloadCmd) Execute(args []string) error {
 
 	var mc gitcollector.MetricsCollector
 	if c.MetricsDBURI != "" {
-		db, err := metrics.PrepareDB(
-			c.MetricsDBURI, c.MetricsDBTable, c.Org,
+		mc = setupMetrics(
+			c.MetricsDBURI,
+			c.MetricsDBTable,
+			orgs,
+			c.MetricsSync,
 		)
-		check(err, "metrics database")
-
-		mc = metrics.NewCollector(&metrics.CollectorOpts{
-			Log:      log.New(nil),
-			Send:     metrics.SendToDB(db, c.MetricsDBTable, c.Org),
-			SyncTime: time.Duration(c.MetricsSync) * time.Second,
-		})
 
 		log.Debugf("metrics collection activated: sync timeout %d",
 			c.MetricsSync)
@@ -125,23 +129,13 @@ func (c *DownloadCmd) Execute(args []string) error {
 	wp.Run()
 	log.Debugf("worker pool is running")
 
-	dp := discovery.NewGHProvider(
-		c.Org,
-		download,
-		&discovery.GHProviderOpts{
-			AuthToken: c.Token,
-		},
-	)
+	go runGHOrgProviders(log.New(nil), orgs, c.Token, download)
 
-	log.Debugf("github provider started")
-	if err := dp.Start(); err != nil &&
-		!gitcollector.ErrProviderStopped.Is(err) {
-		check(err, "github provider failed")
-	}
-
-	close(download)
 	wp.Wait()
 	log.Debugf("worker pool stopped successfully")
+
+	elapsed := time.Since(start).String()
+	log.Infof("collection finished in %s", elapsed)
 	return nil
 }
 
@@ -151,3 +145,64 @@ func check(err error, message string) {
 		os.Exit(1)
 	}
 }
+
+func setupMetrics(
+	uri, table string,
+	orgs []string,
+	metricSync int64,
+) gitcollector.MetricsCollector {
+	db, err := metrics.PrepareDB(uri, table, orgs)
+	check(err, "metrics database")
+
+	mcs := make(map[string]*metrics.Collector, len(orgs))
+	for _, org := range orgs {
+		mc := metrics.NewCollector(&metrics.CollectorOpts{
+			Log:      log.New(log.Fields{"org": org}),
+			Send:     metrics.SendToDB(db, table, org),
+			SyncTime: time.Duration(metricSync) * time.Second,
+		})
+
+		mcs[org] = mc
+	}
+
+	return metrics.NewCollectorByOrg(mcs)
+}
+
+func runGHOrgProviders(
+	logger log.Logger,
+	orgs []string,
+	token string,
+	download chan gitcollector.Job,
+) {
+	var wg sync.WaitGroup
+	wg.Add(len(orgs))
+	for _, o := range orgs {
+		org := o
+		p := discovery.NewGHProvider(
+			download,
+			discovery.NewGHOrgReposIter(
+				org,
+				&discovery.GHReposIterOpts{
+					AuthToken: token,
+				},
+			),
+			&discovery.GHProviderOpts{},
+		)
+
+		go func() {
+			err := p.Start()
+			if err != nil &&
+				!discovery.ErrNewRepositoriesNotFound.Is(err) {
+				logger.Warningf(err.Error())
+			}
+
+			logger.Debugf("%s organization provider stopped", org)
+			wg.Done()
+		}()
+
+		logger.Debugf("%s organization provider started", org)
+	}
+
+	wg.Wait()
+	close(download)
+}

discovery/iterator.go

@@ -0,0 +1,159 @@
+package discovery
+
+import (
+	"context"
+	"net/http"
+	"time"
+
+	"github.com/google/go-github/github"
+	"golang.org/x/oauth2"
+)
+
+// GHReposIterOpts represents configuration options for a GHReposIter.
+type GHReposIterOpts struct {
+	HTTPTimeout    time.Duration
+	ResultsPerPage int
+	TimeNewRepos   time.Duration
+	AuthToken      string
+}
+
+const (
+	httpTimeout    = 30 * time.Second
+	resultsPerPage = 100
+	waitNewRepos   = 24 * time.Hour
+)
+
+// GHOrgReposIter is a GHRepositoriesIter by organization name.
+type GHOrgReposIter struct {
+	org          string
+	client       *github.Client
+	repos        []*github.Repository
+	checkpoint   int
+	opts         *github.RepositoryListByOrgOptions
+	waitNewRepos time.Duration
+}
+
+var _ GHRepositoriesIter = (*GHOrgReposIter)(nil)
+
+// NewGHOrgReposIter builds a new GHOrgReposIter.
+func NewGHOrgReposIter(org string, opts *GHReposIterOpts) *GHOrgReposIter {
+	if opts == nil {
+		opts = &GHReposIterOpts{}
+	}
+
+	to := opts.HTTPTimeout
+	if to <= 0 {
+		to = httpTimeout
+	}
+
+	rpp := opts.ResultsPerPage
+	if rpp <= 0 || rpp > 100 {
+		rpp = resultsPerPage
+	}
+
+	wnr := opts.TimeNewRepos
+	if wnr <= 0 {
+		wnr = waitNewRepos
+	}
+
+	return &GHOrgReposIter{
+		org:    org,
+		client: newGithubClient(opts.AuthToken, to),
+		opts: &github.RepositoryListByOrgOptions{
+			ListOptions: github.ListOptions{PerPage: rpp},
+		},
+		waitNewRepos: wnr,
+	}
+}
+
+func newGithubClient(token string, timeout time.Duration) *github.Client {
+	var client *http.Client
+	if token == "" {
+		client = &http.Client{}
+	} else {
+		client = oauth2.NewClient(
+			context.Background(),
+			oauth2.StaticTokenSource(
+				&oauth2.Token{AccessToken: token},
+			),
+		)
+	}
+
+	client.Timeout = timeout
+	return github.NewClient(client)
+}
+
+// Next implements the GHRepositoriesIter interface.
+func (p *GHOrgReposIter) Next() (*github.Repository, time.Duration, error) {
+	if len(p.repos) == 0 {
+		retry, err := p.requestRepos()
+		if err != nil && len(p.repos) == 0 {
+			return nil, retry, err
+		}
+	}
+
+	var next *github.Repository
+	next, p.repos = p.repos[0], p.repos[1:]
+	return next, 0, nil
+}
+
+func (p *GHOrgReposIter) requestRepos() (time.Duration, error) {
+	repos, res, err := p.client.Repositories.ListByOrg(
+		context.Background(),
+		p.org,
+		p.opts,
+	)
+
+	if err != nil {
+		if _, ok := err.(*github.RateLimitError); !ok {
+			return -1, err
+		}
+
+		return timeToRetry(res), ErrRateLimitExceeded.Wrap(err)
+	}
+
+	bufRepos := repos
+	if p.checkpoint > 0 {
+		i := p.checkpoint
+		if len(repos) < p.checkpoint {
+			// return err?
+			i = 0
+		}
+
+		bufRepos = repos[i:]
+	}
+
+	if len(repos) < p.opts.PerPage {
+		p.checkpoint = len(repos)
+	}
+
+	err = nil
+	if res.NextPage == 0 {
+		if len(repos) == p.opts.PerPage {
+			p.opts.Page++
+		}
+
+		err = ErrNewRepositoriesNotFound.New()
+	} else {
+		p.opts.Page = res.NextPage
+	}
+
+	p.repos = bufRepos
+	return p.waitNewRepos, err
+}
+
+func timeToRetry(res *github.Response) time.Duration {
+	now := time.Now().UTC().Unix()
+	resetTime := res.Rate.Reset.UTC().Unix()
+	timeToReset := time.Duration(resetTime-now) * time.Second
+	remaining := res.Rate.Remaining
+	if timeToReset < 0 || timeToReset > 1*time.Hour {
+		// If this happens, the system clock is probably wrong, so we
+		// assume we are at the beginning of the window and consider
+		// only total requests per hour.
+		timeToReset = 1 * time.Hour
+		remaining = res.Rate.Limit
+	}
+
+	return timeToReset / time.Duration(remaining+1)
+}