chore: Replace QueuedCount -> CountByState with bitset parameter (#64302)

Previously, the QueuedCount method was confusing because:

1. By default, it actually returned the count for both the 'queued' and
   'errored' states (despite the name just have 'Queued').
2. There was an additional boolean flag for also returning entries in
    the 'processing' state, but reduced clarity at call-sites.

So I've changed the method to take a bitset instead, mirroring the
just-added Exists API, and renamed the method to a more
generic 'CountByState'.

While this does make call-sites a bit more verbose, I think the
clarity win makes the change an overall positive one.

Author: varungandhi-src

cmd/worker/internal/executormultiqueue/multiqueue_metrics_reporter.go

@@ -47,8 +47,8 @@ func (j *multiqueueMetricsReporterJob) Routines(_ context.Context, observationCt
 	multiqueueMetricsReporter, err := executorqueue.NewMultiqueueMetricReporter(
 		executortypes.ValidQueueNames,
 		configInst.MetricsConfig,
-		codeIntelStore.QueuedCount,
-		batchesStore.QueuedCount,
+		codeIntelStore.CountByState,
+		batchesStore.CountByState,
 	)
 	if err != nil {
 		return nil, err

cmd/worker/internal/executorqueue/external_emitter.go

@@ -7,12 +7,13 @@ import (
 	uploadsshared "github.com/sourcegraph/sourcegraph/internal/codeintel/uploads/shared"
 	"github.com/sourcegraph/sourcegraph/internal/goroutine"
 	"github.com/sourcegraph/sourcegraph/internal/workerutil"
+	"github.com/sourcegraph/sourcegraph/internal/workerutil/dbworker/store"
 	"github.com/sourcegraph/sourcegraph/lib/errors"
 )
 
 type externalEmitter[T workerutil.Record] struct {
 	queueName  string
-	countFuncs []func(ctx context.Context, includeProcessing bool) (int, error)
+	countFuncs []func(ctx context.Context, bitset store.RecordState) (int, error)
 	reporters  []reporter
 	allocation QueueAllocation
 }
@@ -27,9 +28,9 @@ type reporter interface {
 func (r *externalEmitter[T]) Handle(ctx context.Context) error {
 	var count int
 	for _, countFunc := range r.countFuncs {
-		subCount, err := countFunc(context.Background(), true)
+		subCount, err := countFunc(context.Background(), store.StateQueued|store.StateErrored|store.StateProcessing)
 		if err != nil {
-			return errors.Wrap(err, "dbworkerstore.QueuedCount")
+			return errors.Wrap(err, "dbworkerstore.CountByState")
 		}
 		count += subCount
 	}

cmd/worker/internal/executorqueue/reporter.go

@@ -19,7 +19,7 @@ func NewMetricReporter[T workerutil.Record](observationCtx *observation.Context,
 	return initExternalMetricReporters(queueName, store, metricsConfig)
 }
 
-func initExternalMetricReporters[T workerutil.Record](queueName string, store store.Store[T], metricsConfig *Config) (goroutine.BackgroundRoutine, error) {
+func initExternalMetricReporters[T workerutil.Record](queueName string, store_ store.Store[T], metricsConfig *Config) (goroutine.BackgroundRoutine, error) {
 	reporters, err := configureReporters(metricsConfig)
 	if err != nil {
 		return nil, err
@@ -30,7 +30,7 @@ func initExternalMetricReporters[T workerutil.Record](queueName string, store st
 		ctx,
 		&externalEmitter[T]{
 			queueName:  queueName,
-			countFuncs: []func(ctx context.Context, includeProcessing bool) (int, error){store.QueuedCount},
+			countFuncs: []func(context.Context, store.RecordState) (int, error){store_.CountByState},
 			reporters:  reporters,
 			allocation: metricsConfig.Allocations[queueName],
 		},
@@ -43,7 +43,7 @@ func initExternalMetricReporters[T workerutil.Record](queueName string, store st
 // NewMultiqueueMetricReporter returns a periodic background routine that reports the sum of the lengths all configured queues.
 // This does not reinitialise Prometheus metrics as is done in NewMetricReporter, as this only needs to be done once and is
 // already done for the single queue metrics.
-func NewMultiqueueMetricReporter(queueNames []string, metricsConfig *Config, countFuncs ...func(ctx context.Context, includeProcessing bool) (int, error)) (goroutine.BackgroundRoutine, error) {
+func NewMultiqueueMetricReporter(queueNames []string, metricsConfig *Config, countFuncs ...func(_ context.Context, bitset store.RecordState) (int, error)) (goroutine.BackgroundRoutine, error) {
 	reporters, err := configureReporters(metricsConfig)
 	if err != nil {
 		return nil, err

cmd/worker/internal/permissions/BUILD.bazel

@@ -79,6 +79,7 @@ go_test(
         "//internal/observation",
         "//internal/timeutil",
         "//internal/types",
+        "//internal/workerutil/dbworker/store",
         "//lib/pointers",
         "//schema",
         "@com_github_google_go_cmp//cmp",

cmd/worker/internal/permissions/bitbucket_projects_test.go

@@ -22,6 +22,7 @@ import (
 	"github.com/sourcegraph/sourcegraph/internal/extsvc/bitbucketserver"
 	"github.com/sourcegraph/sourcegraph/internal/observation"
 	"github.com/sourcegraph/sourcegraph/internal/types"
+	dbworker "github.com/sourcegraph/sourcegraph/internal/workerutil/dbworker/store"
 	"github.com/sourcegraph/sourcegraph/schema"
 )
 
@@ -39,7 +40,7 @@ func TestStore(t *testing.T) {
 	require.NotZero(t, jobID)
 
 	store := createBitbucketProjectPermissionsStore(observation.TestContextTB(t), db, &config{})
-	count, err := store.QueuedCount(ctx, true)
+	count, err := store.CountByState(ctx, dbworker.StateQueued|dbworker.StateErrored|dbworker.StateProcessing)
 	require.NoError(t, err)
 	require.Equal(t, 1, count)
 }

cmd/worker/internal/search/job.go

@@ -20,6 +20,7 @@ import (
 	"github.com/sourcegraph/sourcegraph/internal/search/exhaustive"
 	"github.com/sourcegraph/sourcegraph/internal/search/exhaustive/service"
 	"github.com/sourcegraph/sourcegraph/internal/search/exhaustive/store"
+	dbworkerstore "github.com/sourcegraph/sourcegraph/internal/workerutil/dbworker/store"
 )
 
 // config stores shared config we can override in each worker. We don't expose
@@ -39,7 +40,7 @@ type searchJob struct {
 	once         sync.Once
 	err          error
 	workerStores []interface {
-		QueuedCount(context.Context, bool) (int, error)
+		CountByState(_ context.Context, bitset dbworkerstore.RecordState) (int, error)
 	}
 	workers []goroutine.BackgroundRoutine
 }
@@ -136,8 +137,9 @@ func (j *searchJob) newSearchJobRoutines(
 // hasWork returns true if any of the workers have work in its queue or is
 // processing something. This is only exposed for tests.
 func (j *searchJob) hasWork(ctx context.Context) bool {
+	statesBitset := dbworkerstore.StateQueued | dbworkerstore.StateErrored | dbworkerstore.StateProcessing
 	for _, w := range j.workerStores {
-		if count, _ := w.QueuedCount(ctx, true); count > 0 {
+		if count, _ := w.CountByState(ctx, statesBitset); count > 0 {
 			return true
 		}
 	}

internal/codeintel/syntactic_indexing/BUILD.bazel

@@ -65,6 +65,7 @@ go_test(
         "//internal/gitserver",
         "//internal/gitserver/gitdomain",
         "//internal/observation",
+        "//internal/workerutil/dbworker/store",
         "@com_github_keegancsmith_sqlf//:sqlf",
         "@com_github_stretchr_testify//require",
     ],

internal/codeintel/syntactic_indexing/scheduler_test.go

@@ -23,6 +23,7 @@ import (
 	"github.com/sourcegraph/sourcegraph/internal/gitserver"
 	"github.com/sourcegraph/sourcegraph/internal/gitserver/gitdomain"
 	"github.com/sourcegraph/sourcegraph/internal/observation"
+	dbworker "github.com/sourcegraph/sourcegraph/internal/workerutil/dbworker/store"
 )
 
 func TestSyntacticIndexingScheduler(t *testing.T) {
@@ -81,7 +82,7 @@ func TestSyntacticIndexingScheduler(t *testing.T) {
 
 	err := scheduler.Schedule(observationCtx, ctx, time.Now())
 	require.NoError(t, err)
-	require.Equal(t, 2, unwrap(jobStore.DBWorkerStore().QueuedCount(ctx, false))(t))
+	require.Equal(t, 2, unwrap(jobStore.DBWorkerStore().CountByState(ctx, dbworker.StateQueued|dbworker.StateErrored))(t))
 
 	job1, recordReturned, err := jobStore.DBWorkerStore().Dequeue(ctx, "worker-1", []*sqlf.Query{})
 	require.NoError(t, err)

internal/codeintel/syntactic_indexing/store_test.go

@@ -13,6 +13,7 @@ import (
 	"github.com/sourcegraph/sourcegraph/internal/database"
 	"github.com/sourcegraph/sourcegraph/internal/database/dbtest"
 	"github.com/sourcegraph/sourcegraph/internal/observation"
+	dbworker "github.com/sourcegraph/sourcegraph/internal/workerutil/dbworker/store"
 )
 
 func TestSyntacticIndexingStoreDequeue(t *testing.T) {
@@ -38,7 +39,7 @@ func TestSyntacticIndexingStoreDequeue(t *testing.T) {
 
 	ctx := context.Background()
 
-	initCount, _ := store.QueuedCount(ctx, true)
+	initCount, _ := store.CountByState(ctx, dbworker.StateQueued|dbworker.StateErrored|dbworker.StateProcessing)
 
 	require.Equal(t, 0, initCount)
 
@@ -83,7 +84,7 @@ func TestSyntacticIndexingStoreDequeue(t *testing.T) {
 		},
 	)
 
-	afterCount, _ := store.QueuedCount(ctx, true)
+	afterCount, _ := store.CountByState(ctx, dbworker.StateQueued|dbworker.StateErrored|dbworker.StateProcessing)
 
 	require.Equal(t, 3, afterCount)
 
@@ -166,7 +167,7 @@ func TestSyntacticIndexingStoreEnqueue(t *testing.T) {
 
 	// Assertions below verify that records inserted by InsertJobs are
 	// still visible by DB Worker interface
-	afterCount, _ := store.QueuedCount(ctx, true)
+	afterCount, _ := store.CountByState(ctx, dbworker.StateQueued|dbworker.StateErrored|dbworker.StateProcessing)
 
 	require.Equal(t, 2, afterCount)
 

internal/insights/background/queryrunner/worker.go

@@ -61,7 +61,7 @@ func NewWorker(ctx context.Context, logger log.Logger, workerStore *workerStoreE
 		Name: "src_query_runner_worker_total",
 		Help: "Total number of jobs in the queued state.",
 	}, func() float64 {
-		count, err := workerStore.QueuedCount(context.Background(), false)
+		count, err := workerStore.CountByState(context.Background(), dbworkerstore.StateQueued|dbworkerstore.StateErrored)
 		if err != nil {
 			logger.Error("Failed to get queued job count", log.Error(err))
 		}