internal/postgres: limit imported-by-count updates

Add a query parameter that limits the total number of counts updated.
The intent is to limit the work that the endpoint does, to avoid
overloading the system.

Change-Id: I28b0031f171192870e84111714c72d996a6a6964
Reviewed-on: https://go-review.googlesource.com/c/pkgsite/+/854068
Reviewed-by: Nicholas Husin <nsh@golang.org>
Reviewed-by: Ethan Lee <ethanalee@google.com>
LUCI-TryBot-Result: golang-scoped@luci-project-accounts.iam.gserviceaccount.com <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Nicholas Husin <husin@google.com>

internal/postgres/import_counts.go

diff --git a/internal/postgres/import_counts.go b/internal/postgres/import_counts.go
index 89f1652..64a8da2 100644
--- a/internal/postgres/import_counts.go
+++ b/internal/postgres/import_counts.go
@@ -21,6 +21,9 @@
 	// (in a single transaction, with the rows locked).
 	// If zero, a reasonable default is used.
 	BatchSize int
+	// A hard limit on the number of rows to update (per kind: applies separately to modules
+	// and packages). If zero, then no limit.
+	Limit int
 	// Maximum time to spend updating. If zero, then no max.
 	MaxTime time.Duration
 	// How old a count has to be before updating it.
@@ -221,9 +224,17 @@
 	cutoff := time.Now().Add(-opts.Staleness)
 	deadline := time.Now().Add(opts.MaxTime)
 	for time.Now().Before(deadline) {
+		batchSize := opts.BatchSize
+		if opts.Limit > 0 {
+			rem := opts.Limit - int(nUpdated)
+			if rem <= 0 {
+				break
+			}
+			batchSize = min(batchSize, rem)
+		}
 		var nu int64
 		err = db.db.Transact(ctx, sql.LevelDefault, func(tx *database.DB) error {
-			if _, err := tx.Exec(ctx, query, cutoff, opts.BatchSize); err != nil {
+			if _, err := tx.Exec(ctx, query, cutoff, batchSize); err != nil {
 				return fmt.Errorf("creating temp table: %w", err)
 			}
 			var err error

internal/postgres/import_counts_test.go

diff --git a/internal/postgres/import_counts_test.go b/internal/postgres/import_counts_test.go
index ec23d42..6a72e5f 100644
--- a/internal/postgres/import_counts_test.go
+++ b/internal/postgres/import_counts_test.go
@@ -235,6 +235,24 @@
 		// counted in imported_by_module_count.
 		_ = validateImportedByCountAndGetSearchDocument(t, testDB, "fmt", 1, 2)
 	})
+	t.Run("limit", func(t *testing.T) {
+		testDB, release := acquire(t)
+		defer release()
+
+		m := sample.Module("mod.com/A", "v1.0.0", "p1", "p2", "p3", "p4")
+		testDB.MustInsertModule(t, m)
+
+		np, nm, err := testDB.UpdateSearchDocumentsImportedByCount(ctx, &ImporterCountOptions{
+			BatchSize: 2,
+			Limit:     3,
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+		if np != 3 || nm != 3 {
+			t.Errorf("got (%d, %d) updated, want (3, 3)", np, nm)
+		}
+	})
 }
 
 func TestUpdateSearchDocumentsImportedByCountWithCounts(t *testing.T) {

internal/worker/server.go

diff --git a/internal/worker/server.go b/internal/worker/server.go
index eec758b..1d1d915 100644
--- a/internal/worker/server.go
+++ b/internal/worker/server.go
@@ -278,15 +278,19 @@
 	// The Cloud Scheduler max timeout is 30 minutes.
 	// Default to less for wiggle room.
 	defaultMaxTime = 26 * time.Minute
+	// Update as many counts as possible in the given time.
+	defaultLimit = 0
 )
 
 // handleUpdateImportedByCount updates imported_by_count for all packages.
 func (s *Server) handleUpdateImportedByCount(w http.ResponseWriter, r *http.Request) error {
 	batchSize := parseIntParam(r, "batch", defaultBatchSize)
+	limit := parseIntParam(r, "limit", defaultLimit)
 	np, nm, err := s.db.UpdateSearchDocumentsImportedByCount(r.Context(), &postgres.ImporterCountOptions{
 		BatchSize: batchSize,
 		Staleness: defaultStaleness,
 		MaxTime:   defaultMaxTime,
+		Limit:     limit,
 	})
 	if err != nil {
 		return err