Merge pull request #1678 from spring-financial-group/feat/status-breakup

feat: batched graphQL for keeper status queries

Author: jenkins-x-bot

pkg/keeper/keeper.go

@@ -1889,40 +1889,8 @@ func restAPISearch(spc scmProviderClient, log *logrus.Entry, queries keeper.Quer
 }
 
 func bucketedGraphQLSearch(querier querier, query keeper.Query, log *logrus.Entry) ([]PullRequest, error) {
-	bucketedQueries := query.BucketedQueries(100)
-	var wg sync.WaitGroup
-	resultsChan := make(chan []PullRequest, len(bucketedQueries))
-	errsChan := make(chan error, len(bucketedQueries))
-
-	for idx, q := range bucketedQueries {
-		wg.Add(1)
-		go func(idx int, q string) {
-			defer wg.Done()
-			bucketResults, err := graphQLSearch(querier, log, q, time.Time{}, time.Now())
-			if err != nil {
-				errsChan <- fmt.Errorf("graphQLSearch failed for bucket %d with query %s: %w", idx, q, err)
-				return
-			}
-			resultsChan <- bucketResults
-		}(idx, q)
-	}
-
-	wg.Wait()
-	close(resultsChan)
-	close(errsChan)
-
-	var results []PullRequest
-	for r := range resultsChan {
-		results = append(results, r...)
-	}
-	var errs []error
-	for err := range errsChan {
-		errs = append(errs, err)
-	}
-	if len(errs) > 0 {
-		return results, fmt.Errorf("one or more bucketed GraphQL searches failed: %w", errorutil.NewAggregate(errs...))
-	}
-	return results, nil
+	bucketedQueries := query.BucketedQueries(repoBucketSize)
+	return executeBucketedQueries(querier, log, bucketedQueries, time.Time{}, time.Now())
 }
 
 func loadMissingLabels(spc scmProviderClient, pr *scm.PullRequest) error {

pkg/keeper/keeper_test.go

@@ -25,6 +25,7 @@ import (
 	"net/http/httptest"
 	"reflect"
 	"strings"
+	"sync"
 	"testing"
 	"text/template"
 	"time"
@@ -512,6 +513,9 @@ type fgc struct {
 	ignoreExpected bool
 	combinedStatus map[string]map[string]commitStatus
 	fakeClient     *scm.Client
+
+	mu       sync.Mutex
+	queryLog []string
 }
 
 func (f *fgc) ListPullRequestComments(owner, repo string, number int) ([]*scm.Comment, error) {
@@ -563,6 +567,11 @@ func (f *fgc) Query(ctx context.Context, q interface{}, vars map[string]interfac
 	if !ok {
 		return errors.New("unexpected query type")
 	}
+	if qs, ok := vars["query"]; ok {
+		f.mu.Lock()
+		f.queryLog = append(f.queryLog, string(qs.(githubql.String)))
+		f.mu.Unlock()
+	}
 	for _, pr := range f.prs {
 		sq.Search.Nodes = append(
 			sq.Search.Nodes,

pkg/keeper/search.go

@@ -19,13 +19,18 @@ package keeper
 import (
 	"context"
 	"fmt"
+	"sync"
 	"time"
 
+	"github.com/jenkins-x/lighthouse/pkg/errorutil"
 	"github.com/jenkins-x/lighthouse/pkg/scmprovider"
 	githubql "github.com/shurcooL/githubv4"
 	"github.com/sirupsen/logrus"
 )
 
+// repoBucketSize is the maximum number of repos to include in a single GraphQL search query.
+const repoBucketSize = 100
+
 type querier func(ctx context.Context, result interface{}, vars map[string]interface{}) error
 
 func datedQuery(q string, start, end time.Time) string {
@@ -82,6 +87,43 @@ func graphQLSearch(query querier, log *logrus.Entry, q string, start, end time.T
 	return ret, nil
 }
 
+// executeBucketedQueries runs bucketed queries through graphQLSearch in parallel
+func executeBucketedQueries(q querier, log *logrus.Entry, bucketQueries []string, start, end time.Time) ([]PullRequest, error) {
+	var wg sync.WaitGroup
+	resultsChan := make(chan []PullRequest, len(bucketQueries))
+	errsChan := make(chan error, len(bucketQueries))
+
+	for idx, bq := range bucketQueries {
+		wg.Add(1)
+		go func(idx int, bq string) {
+			defer wg.Done()
+			bucketResults, err := graphQLSearch(q, log, bq, start, end)
+			if err != nil {
+				errsChan <- fmt.Errorf("graphQLSearch failed for bucket %d with query %s: %w", idx, bq, err)
+				return
+			}
+			resultsChan <- bucketResults
+		}(idx, bq)
+	}
+
+	wg.Wait()
+	close(resultsChan)
+	close(errsChan)
+
+	var results []PullRequest
+	for r := range resultsChan {
+		results = append(results, r...)
+	}
+	var errs []error
+	for err := range errsChan {
+		errs = append(errs, err)
+	}
+	if len(errs) > 0 {
+		return results, fmt.Errorf("one or more bucketed GraphQL searches failed: %w", errorutil.NewAggregate(errs...))
+	}
+	return results, nil
+}
+
 // dateToken generates a GitHub graphQLSearch query token for the specified date range.
 // See: https://help.github.com/articles/understanding-the-search-syntax/#query-for-dates
 func dateToken(start, end time.Time) string {

pkg/keeper/search_test.go

@@ -20,6 +20,8 @@ import (
 	"context"
 	"errors"
 	"reflect"
+	"sort"
+	"sync"
 	"testing"
 	"time"
 
@@ -171,3 +173,139 @@ func TestSearch(t *testing.T) {
 		})
 	}
 }
+
+func TestExecuteBucketedQueries(t *testing.T) {
+	now := time.Now()
+	earlier := now.Add(-5 * time.Hour)
+	log := logrus.WithField("test", "TestExecuteBucketedQueries")
+
+	makePRs := func(numbers ...int) []PullRequest {
+		var prs []PullRequest
+		for _, n := range numbers {
+			prs = append(prs, PullRequest{Number: githubql.Int(n)})
+		}
+		return prs
+	}
+
+	sortPRs := func(prs []PullRequest) {
+		sort.Slice(prs, func(i, j int) bool {
+			return prs[i].Number < prs[j].Number
+		})
+	}
+
+	cases := []struct {
+		name        string
+		bucketQ     []string
+		resultsByQ  map[string][]PullRequest
+		errsByQ     map[string]error
+		expectedPRs []PullRequest
+		expectErr   bool
+	}{
+		{
+			name:        "empty bucket list",
+			bucketQ:     nil,
+			resultsByQ:  map[string][]PullRequest{},
+			errsByQ:     map[string]error{},
+			expectedPRs: nil,
+			expectErr:   false,
+		},
+		{
+			name:    "single bucket succeeds",
+			bucketQ: []string{"bucket-q-1"},
+			resultsByQ: map[string][]PullRequest{
+				"bucket-q-1": makePRs(1, 2),
+			},
+			errsByQ:     map[string]error{},
+			expectedPRs: makePRs(1, 2),
+			expectErr:   false,
+		},
+		{
+			name:    "multiple buckets succeed",
+			bucketQ: []string{"bucket-q-1", "bucket-q-2"},
+			resultsByQ: map[string][]PullRequest{
+				"bucket-q-1": makePRs(1, 2),
+				"bucket-q-2": makePRs(3, 4),
+			},
+			errsByQ:     map[string]error{},
+			expectedPRs: makePRs(1, 2, 3, 4),
+			expectErr:   false,
+		},
+		{
+			name:    "one bucket fails, partial results",
+			bucketQ: []string{"bucket-q-1", "bucket-q-2"},
+			resultsByQ: map[string][]PullRequest{
+				"bucket-q-1": makePRs(1, 2),
+			},
+			errsByQ: map[string]error{
+				"bucket-q-2": errors.New("bucket-q-2 failed"),
+			},
+			expectedPRs: makePRs(1, 2),
+			expectErr:   true,
+		},
+		{
+			name:       "all buckets fail",
+			bucketQ:    []string{"bucket-q-1", "bucket-q-2"},
+			resultsByQ: map[string][]PullRequest{},
+			errsByQ: map[string]error{
+				"bucket-q-1": errors.New("bucket-q-1 failed"),
+				"bucket-q-2": errors.New("bucket-q-2 failed"),
+			},
+			expectedPRs: nil,
+			expectErr:   true,
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			// graphQLSearch wraps queries with datedQuery; build the dated form so
+			// our mock can match on it.
+			datedResults := make(map[string][]PullRequest, len(tc.resultsByQ))
+			datedErrs := make(map[string]error, len(tc.errsByQ))
+			datedBuckets := make([]string, len(tc.bucketQ))
+			for i, q := range tc.bucketQ {
+				dq := datedQuery(q, floor(earlier), floor(now))
+				datedBuckets[i] = q // pass bare query to executeBucketedQueries
+				if v, ok := tc.resultsByQ[q]; ok {
+					datedResults[dq] = v
+				}
+				if v, ok := tc.errsByQ[q]; ok {
+					datedErrs[dq] = v
+				}
+			}
+
+			var mu sync.Mutex
+			mockQuery := func(_ context.Context, result interface{}, vars map[string]interface{}) error {
+				q := string(vars["query"].(githubql.String))
+				mu.Lock()
+				prs, hasPRs := datedResults[q]
+				err, hasErr := datedErrs[q]
+				mu.Unlock()
+				if hasErr {
+					return err
+				}
+				if hasPRs {
+					sq := result.(*searchQuery)
+					for _, pr := range prs {
+						sq.Search.Nodes = append(sq.Search.Nodes, PRNode{pr})
+					}
+				}
+				return nil
+			}
+
+			prs, err := executeBucketedQueries(mockQuery, log, datedBuckets, earlier, now)
+
+			if tc.expectErr && err == nil {
+				t.Error("expected error but got nil")
+			}
+			if !tc.expectErr && err != nil {
+				t.Errorf("unexpected error: %v", err)
+			}
+
+			sortPRs(prs)
+			sortPRs(tc.expectedPRs)
+			if !reflect.DeepEqual(tc.expectedPRs, prs) {
+				t.Errorf("prs do not match:\n%s", cmp.Diff(tc.expectedPRs, prs))
+			}
+		})
+	}
+}

pkg/keeper/status.go

@@ -316,7 +316,7 @@ func (sc *statusController) setStatuses(all []PullRequest, pool map[string]prWit
 	// queryMap caches which queries match a repo.
 	// Make a new one each sync loop as queries will change.
 	queryMap := sc.config().Keeper.Queries.QueryMap()
-	processed := sets.NewString()
+	processed := sets.New[string]()
 
 	process := func(pr *PullRequest) {
 		processed.Insert(pr.prKey())
@@ -461,8 +461,9 @@ func (sc *statusController) search() []PullRequest {
 	}
 
 	orgExceptions, repos := queries.OrgExceptionsAndRepos()
-	orgs := sets.StringKeySet(orgExceptions)
-	query := openPRsQuery(orgs.List(), repos.List(), orgExceptions)
+	orgs := sets.KeySet(orgExceptions)
+	reposSet := sets.Set[string](repos)
+	query := openPRsQuery(sets.List(orgs), sets.List(reposSet), orgExceptions)
 	now := time.Now()
 	log := sc.logger.WithField("query", query)
 	if query != sc.PreviousQuery {
@@ -476,7 +477,7 @@ func (sc *statusController) search() []PullRequest {
 	var err error
 
 	if sc.spc.SupportsGraphQL() {
-		prs, err = graphQLSearch(sc.spc.Query, sc.logger, query, sc.LatestPR.Time, now)
+		prs, err = sc.bucketedGraphQLStatusSearch(orgs, reposSet, orgExceptions, now)
 	} else {
 		kq := keeper.Query{}
 		kq.Repos = append(kq.Repos, repos.List()...)
@@ -497,7 +498,13 @@ func (sc *statusController) search() []PullRequest {
 		return nil
 	}
 
-	latest := prs[len(prs)-1].UpdatedAt
+	// Find the latest UpdatedAt across all results (bucketed queries may not be globally sorted)
+	var latest metav1.Time
+	for _, pr := range prs {
+		if pr.UpdatedAt.After(latest.Time) {
+			latest = metav1.Time{Time: pr.UpdatedAt.Time}
+		}
+	}
 	if latest.IsZero() {
 		log.WithField("latestPR", sc.LatestPR).Debug("latest PR has zero time")
 		return prs
@@ -507,6 +514,27 @@ func (sc *statusController) search() []PullRequest {
 	return prs
 }
 
+func (sc *statusController) bucketedGraphQLStatusSearch(orgs sets.Set[string], repos sets.Set[string], orgExceptions map[string]sets.String, now time.Time) ([]PullRequest, error) {
+	var bucketQueries []string
+
+	if orgs.Len() > 0 {
+		// One query covering all orgs with their exclusions
+		bucketQueries = append(bucketQueries, openPRsQuery(sets.List(orgs), nil, orgExceptions))
+	}
+
+	repoList := sets.List(repos)
+	for i := 0; i < len(repoList); i += repoBucketSize {
+		end := i + repoBucketSize
+		if end > len(repoList) {
+			end = len(repoList)
+		}
+		// Each repo bucket query has no orgs and no exceptions.
+		bucketQueries = append(bucketQueries, openPRsQuery(nil, repoList[i:end], nil))
+	}
+
+	return executeBucketedQueries(sc.spc.Query, sc.logger, bucketQueries, sc.LatestPR.Time, now)
+}
+
 func openPRsQuery(orgs, repos []string, orgExceptions map[string]sets.String) string {
 	return "is:pr state:open sort:updated-asc " + orgRepoQueryString(orgs, repos, orgExceptions)
 }