Cleanup executor queue and move state to DB (sourcegraph#23165)

This PR updates the executor queue to do not hold state in-memory and removes now redundant cleanup logic. We don't need to manually keep track of records anymore, as we now have proper heartbeats coming from the executor itself. So we could get rid of most of the logic and just rely on the resetter now. Additionally, we now have one handler per queue now which makes it easier to track records as their ID is unique now (wasn't across queues before). That way, using the executorName as the worker_hostname, we can get all the relevant information from the DB and don't need to hold the state in-memory. This also extends the dbworker.Store by additional options to validate the worker is still the correct owner of the record (we did that before checking the in-memory map).

enterprise/cmd/executor-queue/config.go

@@ -1,35 +1,22 @@
 package main
 
 import (
-	"time"
-
 	apiserver "github.com/sourcegraph/sourcegraph/enterprise/cmd/executor-queue/internal/server"
 	"github.com/sourcegraph/sourcegraph/internal/env"
 )
 
 type Config struct {
 	env.BaseConfig
 
-	Port                       int
-	JobRequeueDelay            time.Duration
-	JobCleanupInterval         time.Duration
-	MaximumNumMissedHeartbeats int
+	Port int
 }
 
 func (c *Config) Load() {
 	c.Port = c.GetInt("EXECUTOR_QUEUE_API_PORT", "3191", "The port to listen on.")
-	c.JobRequeueDelay = c.GetInterval("EXECUTOR_QUEUE_JOB_REQUEUE_DELAY", "1m", "The requeue delay of jobs assigned to an unreachable executor.")
-	c.JobCleanupInterval = c.GetInterval("EXECUTOR_QUEUE_JOB_CLEANUP_INTERVAL", "10s", "Interval between cleanup runs.")
-	c.MaximumNumMissedHeartbeats = c.GetInt("EXECUTOR_QUEUE_MAXIMUM_NUM_MISSED_HEARTBEATS", "5", "The number of heartbeats an executor must miss to be considered unreachable.")
 }
 
-func (c *Config) ServerOptions(queueOptions map[string]apiserver.QueueOptions) apiserver.Options {
-	return apiserver.Options{
-		Port:             c.Port,
-		QueueOptions:     queueOptions,
-		RequeueDelay:     c.JobRequeueDelay,
-		UnreportedMaxAge: c.JobCleanupInterval * time.Duration(c.MaximumNumMissedHeartbeats),
-		DeathThreshold:   c.JobCleanupInterval * time.Duration(c.MaximumNumMissedHeartbeats),
-		CleanupInterval:  c.JobCleanupInterval,
+func (c *Config) ServerOptions() apiserver.ServerOptions {
+	return apiserver.ServerOptions{
+		Port: c.Port,
 	}
 }

enterprise/cmd/executor-queue/internal/server/handler.go

@@ -3,50 +3,17 @@ package server
 import (
 	"context"
 	"fmt"
-	"sync"
-	"time"
 
 	"github.com/cockroachdb/errors"
-	"github.com/derision-test/glock"
 	"github.com/inconshreveable/log15"
 
 	apiclient "github.com/sourcegraph/sourcegraph/enterprise/internal/executor"
-	"github.com/sourcegraph/sourcegraph/internal/observation"
 	"github.com/sourcegraph/sourcegraph/internal/workerutil"
 	"github.com/sourcegraph/sourcegraph/internal/workerutil/dbworker/store"
 )
 
 type handler struct {
-	options      Options
-	clock        glock.Clock
-	executors    map[string]*executorMeta
-	m            sync.Mutex // protects executors
-	queueMetrics *QueueMetrics
-}
-
-type Options struct {
-	// Port is the port on which to listen for HTTP connections.
-	Port int
-
-	// QueueOptions is a map from queue name to options specific to that queue.
-	QueueOptions map[string]QueueOptions
-
-	// RequeueDelay controls how far into the future to make a job record visible to the job
-	// queue once the currently processing executor has become unresponsive.
-	RequeueDelay time.Duration
-
-	// UnreportedMaxAge is the maximum time between a record being dequeued and it appearing
-	// in the executor's heartbeat requests before it being considered lost.
-	UnreportedMaxAge time.Duration
-
-	// DeathThreshold is the minimum time since the last heartbeat of an executor before that
-	// executor can be considered as unresponsive. This should be configured to be longer than
-	// the duration between heartbeat interval.
-	DeathThreshold time.Duration
-
-	// CleanupInterval is the duration between periodic invocations of Cleanup, which will
-	// requeue any records that are "lost" according to the thresholds described above.
-	CleanupInterval time.Duration
+	QueueOptions
 }
 
 type QueueOptions struct {
@@ -58,206 +25,89 @@ type QueueOptions struct {
 	RecordTransformer func(ctx context.Context, record workerutil.Record) (apiclient.Job, error)
 }
 
-type executorMeta struct {
-	lastUpdate time.Time
-	jobs       []jobMeta
-}
-
-type jobMeta struct {
-	queueName string
-	record    workerutil.Record
-	started   time.Time
-}
-
-func newHandler(options Options, clock glock.Clock) *handler {
-	return newHandlerWithMetrics(options, clock, &observation.TestContext)
-}
-
-func newHandlerWithMetrics(options Options, clock glock.Clock, observationContext *observation.Context) *handler {
+func newHandler(queueOptions QueueOptions) *handler {
 	return &handler{
-		options:      options,
-		clock:        clock,
-		executors:    map[string]*executorMeta{},
-		queueMetrics: newQueueMetrics(observationContext),
+		QueueOptions: queueOptions,
 	}
 }
 
-var (
-	ErrUnknownQueue = errors.New("unknown queue")
-	ErrUnknownJob   = errors.New("unknown job")
-)
+var ErrUnknownJob = errors.New("unknown job")
 
 // dequeue selects a job record from the database and stashes metadata including
 // the job record and the locking transaction. If no job is available for processing,
-// or the server has hit its maximum transactions, a false-valued flag is returned.
-func (m *handler) dequeue(ctx context.Context, queueName, executorName, executorHostname string) (_ apiclient.Job, dequeued bool, _ error) {
-	queueOptions, ok := m.options.QueueOptions[queueName]
-	if !ok {
-		return apiclient.Job{}, false, ErrUnknownQueue
-	}
-
-	record, dequeued, err := queueOptions.Store.Dequeue(context.Background(), executorHostname, nil)
+// a false-valued flag is returned.
+func (h *handler) dequeue(ctx context.Context, executorName, executorHostname string) (_ apiclient.Job, dequeued bool, _ error) {
+	// We explicitly DON'T want to use executorHostname here, it is NOT guaranteed to be unique.
+	record, dequeued, err := h.Store.Dequeue(ctx, executorName, nil)
 	if err != nil {
 		return apiclient.Job{}, false, err
 	}
 	if !dequeued {
 		return apiclient.Job{}, false, nil
 	}
 
-	job, err := queueOptions.RecordTransformer(ctx, record)
+	job, err := h.RecordTransformer(ctx, record)
 	if err != nil {
-		if _, err := queueOptions.Store.MarkFailed(ctx, record.RecordID(), fmt.Sprintf("failed to transform record: %s", err)); err != nil {
+		if _, err := h.Store.MarkFailed(ctx, record.RecordID(), fmt.Sprintf("failed to transform record: %s", err), store.MarkFinalOptions{}); err != nil {
 			log15.Error("Failed to mark record as failed", "recordID", record.RecordID(), "error", err)
 		}
 
 		return apiclient.Job{}, false, err
 	}
 
-	now := m.clock.Now()
-	m.addMeta(executorName, jobMeta{queueName: queueName, record: record, started: now})
 	return job, true, nil
 }
 
-// addExecutionLogEntry calls AddExecutionLogEntry for the given job. If the job identifier
-// is not known, a false-valued flag is returned.
-func (m *handler) addExecutionLogEntry(ctx context.Context, queueName, executorName string, jobID int, entry workerutil.ExecutionLogEntry) error {
-	queueOptions, ok := m.options.QueueOptions[queueName]
-	if !ok {
-		return ErrUnknownQueue
-	}
-
-	_, err := m.findMeta(queueName, executorName, jobID, false)
-	if err != nil {
-		return err
-	}
-
-	if err := queueOptions.Store.AddExecutionLogEntry(ctx, jobID, entry); err != nil {
-		return err
-	}
-
-	return nil
+// addExecutionLogEntry calls AddExecutionLogEntry for the given job.
+func (h *handler) addExecutionLogEntry(ctx context.Context, executorName string, jobID int, entry workerutil.ExecutionLogEntry) error {
+	return h.Store.AddExecutionLogEntry(ctx, jobID, entry, store.AddExecutionLogEntryOptions{
+		// We pass the WorkerHostname, so the store enforces the record to be owned by this executor. When
+		// the previous executor didn't report heartbeats anymore, but is still alive and reporting logs,
+		// both executors that ever got the job would be writing to the same record. This prevents it.
+		WorkerHostname: executorName,
+		// We pass state to enforce adding log entries is only possible while the record is still dequeued.
+		State: "processing",
+	})
 }
 
-// markComplete calls MarkComplete for the given job, then commits the job's transaction.
-// The job is removed from the executor's job list on success.
-func (m *handler) markComplete(ctx context.Context, queueName, executorName string, jobID int) error {
-	queueOptions, ok := m.options.QueueOptions[queueName]
+// markComplete calls MarkComplete for the given job.
+func (h *handler) markComplete(ctx context.Context, executorName string, jobID int) error {
+	ok, err := h.Store.MarkComplete(ctx, jobID, store.MarkFinalOptions{
+		// We pass the WorkerHostname, so the store enforces the record to be owned by this executor. When
+		// the previous executor didn't report heartbeats anymore, but is still alive and reporting state,
+		// both executors that ever got the job would be writing to the same record. This prevents it.
+		WorkerHostname: executorName,
+	})
 	if !ok {
-		return ErrUnknownQueue
-	}
-
-	job, err := m.findMeta(queueName, executorName, jobID, true)
-	if err != nil {
-		return err
+		return ErrUnknownJob
 	}
-
-	_, err = queueOptions.Store.MarkComplete(ctx, job.record.RecordID())
 	return err
 }
 
-// markErrored calls MarkErrored for the given job, then commits the job's transaction.
-// The job is removed from the executor's job list on success.
-func (m *handler) markErrored(ctx context.Context, queueName, executorName string, jobID int, errorMessage string) error {
-	queueOptions, ok := m.options.QueueOptions[queueName]
+// markErrored calls MarkErrored for the given job.
+func (h *handler) markErrored(ctx context.Context, executorName string, jobID int, errorMessage string) error {
+	ok, err := h.Store.MarkErrored(ctx, jobID, errorMessage, store.MarkFinalOptions{
+		// We pass the WorkerHostname, so the store enforces the record to be owned by this executor. When
+		// the previous executor didn't report heartbeats anymore, but is still alive and reporting state,
+		// both executors that ever got the job would be writing to the same record. This prevents it.
+		WorkerHostname: executorName,
+	})
 	if !ok {
-		return ErrUnknownQueue
-	}
-
-	job, err := m.findMeta(queueName, executorName, jobID, true)
-	if err != nil {
-		return err
+		return ErrUnknownJob
 	}
-
-	_, err = queueOptions.Store.MarkErrored(ctx, job.record.RecordID(), errorMessage)
 	return err
 }
 
-// markFailed calls MarkFailed for the given job, then commits the job's transaction.
-// The job is removed from the executor's job list on success.
-func (m *handler) markFailed(ctx context.Context, queueName, executorName string, jobID int, errorMessage string) error {
-	queueOptions, ok := m.options.QueueOptions[queueName]
+// markFailed calls MarkFailed for the given job.
+func (h *handler) markFailed(ctx context.Context, executorName string, jobID int, errorMessage string) error {
+	ok, err := h.Store.MarkFailed(ctx, jobID, errorMessage, store.MarkFinalOptions{
+		// We pass the WorkerHostname, so the store enforces the record to be owned by this executor. When
+		// the previous executor didn't report heartbeats anymore, but is still alive and reporting state,
+		// both executors that ever got the job would be writing to the same record. This prevents it.
+		WorkerHostname: executorName,
+	})
 	if !ok {
-		return ErrUnknownQueue
-	}
-
-	job, err := m.findMeta(queueName, executorName, jobID, true)
-	if err != nil {
-		return err
+		return ErrUnknownJob
 	}
-
-	_, err = queueOptions.Store.MarkFailed(ctx, job.record.RecordID(), errorMessage)
 	return err
 }
-
-// findMeta returns the job with the given id and executor name. If the job is
-// unknown, an error is returned. If the remove parameter is true, the job will
-// be removed from the executor's job list on success.
-func (m *handler) findMeta(queueName, executorName string, jobID int, remove bool) (jobMeta, error) {
-	m.m.Lock()
-	defer m.m.Unlock()
-
-	executor, ok := m.executors[executorName]
-	if !ok {
-		return jobMeta{}, ErrUnknownJob
-	}
-
-	for i, job := range executor.jobs {
-		if job.queueName == queueName && job.record.RecordID() == jobID {
-			if remove {
-				l := len(executor.jobs) - 1
-				executor.jobs[i] = executor.jobs[l]
-				executor.jobs = executor.jobs[:l]
-				m.updateMetrics()
-			}
-
-			return job, nil
-		}
-	}
-
-	return jobMeta{}, ErrUnknownJob
-}
-
-// addMeta adds a job to the given executor's job list.
-func (m *handler) addMeta(executorName string, job jobMeta) {
-	m.m.Lock()
-	defer m.m.Unlock()
-
-	executor, ok := m.executors[executorName]
-	if !ok {
-		executor = &executorMeta{}
-		m.executors[executorName] = executor
-	}
-
-	now := m.clock.Now()
-	executor.jobs = append(executor.jobs, job)
-	executor.lastUpdate = now
-	m.updateMetrics()
-}
-
-func (m *handler) updateMetrics() {
-	type queueStat struct {
-		JobIDs        []int
-		ExecutorNames map[string]struct{}
-	}
-	queueStats := map[string]queueStat{}
-
-	for executorName, meta := range m.executors {
-		for _, job := range meta.jobs {
-			stat, ok := queueStats[job.queueName]
-			if !ok {
-				stat = queueStat{
-					ExecutorNames: map[string]struct{}{},
-				}
-			}
-
-			stat.JobIDs = append(stat.JobIDs, job.record.RecordID())
-			stat.ExecutorNames[executorName] = struct{}{}
-			queueStats[job.queueName] = stat
-		}
-	}
-
-	for queueName, temp := range queueStats {
-		m.queueMetrics.NumJobs.WithLabelValues(queueName).Set(float64(len(temp.JobIDs)))
-		m.queueMetrics.NumExecutors.WithLabelValues(queueName).Set(float64(len(temp.ExecutorNames)))
-	}
-}

enterprise/cmd/executor-queue/internal/server/handler_test.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"testing"
 
-	"github.com/derision-test/glock"
 	"github.com/google/go-cmp/cmp"
 
 	apiclient "github.com/sourcegraph/sourcegraph/enterprise/internal/executor"
@@ -35,14 +34,9 @@ func TestDequeue(t *testing.T) {
 		return transformedJob, nil
 	}
 
-	options := Options{
-		QueueOptions: map[string]QueueOptions{
-			"test_queue": {Store: store, RecordTransformer: recordTransformer},
-		},
-	}
-	handler := newHandler(options, glock.NewMockClock())
+	handler := newHandler(QueueOptions{Store: store, RecordTransformer: recordTransformer})
 
-	job, dequeued, err := handler.dequeue(context.Background(), "test_queue", "deadbeef", "test")
+	job, dequeued, err := handler.dequeue(context.Background(), "deadbeef", "test")
 	if err != nil {
 		t.Fatalf("unexpected error dequeueing job: %s", err)
 	}
@@ -58,14 +52,9 @@ func TestDequeue(t *testing.T) {
 }
 
 func TestDequeueNoRecord(t *testing.T) {
-	options := Options{
-		QueueOptions: map[string]QueueOptions{
-			"test_queue": {Store: workerstoremocks.NewMockStore()},
-		},
-	}
-	handler := newHandler(options, glock.NewMockClock())
+	handler := newHandler(QueueOptions{Store: workerstoremocks.NewMockStore()})
 
-	_, dequeued, err := handler.dequeue(context.Background(), "test_queue", "deadbeef", "test")
+	_, dequeued, err := handler.dequeue(context.Background(), "deadbeef", "test")
 	if err != nil {
 		t.Fatalf("unexpected error dequeueing job: %s", err)
 	}
@@ -74,30 +63,16 @@ func TestDequeueNoRecord(t *testing.T) {
 	}
 }
 
-func TestDequeueUnknownQueue(t *testing.T) {
-	options := Options{}
-	handler := newHandler(options, glock.NewMockClock())
-
-	if _, _, err := handler.dequeue(context.Background(), "test_queue", "deadbeef", "test"); err != ErrUnknownQueue {
-		t.Fatalf("unexpected error. want=%q have=%q", ErrUnknownQueue, err)
-	}
-}
-
 func TestAddExecutionLogEntry(t *testing.T) {
 	store := workerstoremocks.NewMockStore()
 	store.DequeueFunc.SetDefaultReturn(testRecord{ID: 42}, true, nil)
 	recordTransformer := func(ctx context.Context, record workerutil.Record) (apiclient.Job, error) {
 		return apiclient.Job{ID: 42}, nil
 	}
 
-	options := Options{
-		QueueOptions: map[string]QueueOptions{
-			"test_queue": {Store: store, RecordTransformer: recordTransformer},
-		},
-	}
-	handler := newHandler(options, glock.NewMockClock())
+	handler := newHandler(QueueOptions{Store: store, RecordTransformer: recordTransformer})
 
-	job, dequeued, err := handler.dequeue(context.Background(), "test_queue", "deadbeef", "test")
+	job, dequeued, err := handler.dequeue(context.Background(), "deadbeef", "test")
 	if err != nil {
 		t.Fatalf("unexpected error dequeueing job: %s", err)
 	}
@@ -109,7 +84,7 @@ func TestAddExecutionLogEntry(t *testing.T) {
 		Command: []string{"ls", "-a"},
 		Out:     "<log payload>",
 	}
-	if err := handler.addExecutionLogEntry(context.Background(), "test_queue", "deadbeef", job.ID, entry); err != nil {
+	if err := handler.addExecutionLogEntry(context.Background(), "deadbeef", job.ID, entry); err != nil {
 		t.Fatalf("unexpected error updating log contents: %s", err)
 	}
 
@@ -125,59 +100,39 @@ func TestAddExecutionLogEntry(t *testing.T) {
 	}
 }
 
-func TestAddExecutionLogEntryUnknownQueue(t *testing.T) {
-	options := Options{}
-	handler := newHandler(options, glock.NewMockClock())
-
-	entry := workerutil.ExecutionLogEntry{
-		Command: []string{"ls", "-a"},
-		Out:     "<log payload>",
-	}
-	if err := handler.addExecutionLogEntry(context.Background(), "test_queue", "deadbjeef", 42, entry); err != ErrUnknownQueue {
-		t.Fatalf("unexpected error. want=%q have=%q", ErrUnknownQueue, err)
-	}
-}
-
 func TestAddExecutionLogEntryUnknownJob(t *testing.T) {
-	options := Options{
-		QueueOptions: map[string]QueueOptions{
-			"test_queue": {Store: workerstoremocks.NewMockStore()},
-		},
-	}
-	handler := newHandler(options, glock.NewMockClock())
+	store := workerstoremocks.NewMockStore()
+	store.AddExecutionLogEntryFunc.SetDefaultReturn(ErrUnknownJob)
+	handler := newHandler(QueueOptions{Store: store})
 
 	entry := workerutil.ExecutionLogEntry{
 		Command: []string{"ls", "-a"},
 		Out:     "<log payload>",
 	}
-	if err := handler.addExecutionLogEntry(context.Background(), "test_queue", "deadbeef", 42, entry); err != ErrUnknownJob {
+	if err := handler.addExecutionLogEntry(context.Background(), "deadbeef", 42, entry); err != ErrUnknownJob {
 		t.Fatalf("unexpected error. want=%q have=%q", ErrUnknownJob, err)
 	}
 }
 
 func TestMarkComplete(t *testing.T) {
 	store := workerstoremocks.NewMockStore()
 	store.DequeueFunc.SetDefaultReturn(testRecord{ID: 42}, true, nil)
+	store.MarkCompleteFunc.SetDefaultReturn(true, nil)
 	recordTransformer := func(ctx context.Context, record workerutil.Record) (apiclient.Job, error) {
 		return apiclient.Job{ID: 42}, nil
 	}
 
-	options := Options{
-		QueueOptions: map[string]QueueOptions{
-			"test_queue": {Store: store, RecordTransformer: recordTransformer},
-		},
-	}
-	handler := newHandler(options, glock.NewMockClock())
+	handler := newHandler(QueueOptions{Store: store, RecordTransformer: recordTransformer})
 
-	job, dequeued, err := handler.dequeue(context.Background(), "test_queue", "deadbeef", "test")
+	job, dequeued, err := handler.dequeue(context.Background(), "deadbeef", "test")
 	if err != nil {
 		t.Fatalf("unexpected error dequeueing job: %s", err)
 	}
 	if !dequeued {
 		t.Fatalf("expected a job to be dequeued")
 	}
 
-	if err := handler.markComplete(context.Background(), "test_queue", "deadbeef", job.ID); err != nil {
+	if err := handler.markComplete(context.Background(), "deadbeef", job.ID); err != nil {
 		t.Fatalf("unexpected error completing job: %s", err)
 	}
 
@@ -191,50 +146,34 @@ func TestMarkComplete(t *testing.T) {
 }
 
 func TestMarkCompleteUnknownJob(t *testing.T) {
-	options := Options{
-		QueueOptions: map[string]QueueOptions{
-			"test_queue": {Store: workerstoremocks.NewMockStore()},
-		},
-	}
-	handler := newHandler(options, glock.NewMockClock())
+	store := workerstoremocks.NewMockStore()
+	store.MarkCompleteFunc.SetDefaultReturn(false, nil)
+	handler := newHandler(QueueOptions{Store: store})
 
-	if err := handler.markComplete(context.Background(), "test_queue", "deadbeef", 42); err != ErrUnknownJob {
+	if err := handler.markComplete(context.Background(), "deadbeef", 42); err != ErrUnknownJob {
 		t.Fatalf("unexpected error. want=%q have=%q", ErrUnknownJob, err)
 	}
 }
 
-func TestMarkCompleteUnknownQueue(t *testing.T) {
-	options := Options{}
-	handler := newHandler(options, glock.NewMockClock())
-
-	if err := handler.markComplete(context.Background(), "test_queue", "deadbeef", 42); err != ErrUnknownQueue {
-		t.Fatalf("unexpected error. want=%q have=%q", ErrUnknownQueue, err)
-	}
-}
-
 func TestMarkErrored(t *testing.T) {
 	store := workerstoremocks.NewMockStore()
 	store.DequeueFunc.SetDefaultReturn(testRecord{ID: 42}, true, nil)
+	store.MarkErroredFunc.SetDefaultReturn(true, nil)
 	recordTransformer := func(ctx context.Context, record workerutil.Record) (apiclient.Job, error) {
 		return apiclient.Job{ID: 42}, nil
 	}
 
-	options := Options{
-		QueueOptions: map[string]QueueOptions{
-			"test_queue": {Store: store, RecordTransformer: recordTransformer},
-		},
-	}
-	handler := newHandler(options, glock.NewMockClock())
+	handler := newHandler(QueueOptions{Store: store, RecordTransformer: recordTransformer})
 
-	job, dequeued, err := handler.dequeue(context.Background(), "test_queue", "deadbeef", "test")
+	job, dequeued, err := handler.dequeue(context.Background(), "deadbeef", "test")
 	if err != nil {
 		t.Fatalf("unexpected error dequeueing job: %s", err)
 	}
 	if !dequeued {
 		t.Fatalf("expected a job to be dequeued")
 	}
 
-	if err := handler.markErrored(context.Background(), "test_queue", "deadbeef", job.ID, "OH NO"); err != nil {
+	if err := handler.markErrored(context.Background(), "deadbeef", job.ID, "OH NO"); err != nil {
 		t.Fatalf("unexpected error completing job: %s", err)
 	}
 
@@ -251,50 +190,34 @@ func TestMarkErrored(t *testing.T) {
 }
 
 func TestMarkErroredUnknownJob(t *testing.T) {
-	options := Options{
-		QueueOptions: map[string]QueueOptions{
-			"test_queue": {Store: workerstoremocks.NewMockStore()},
-		},
-	}
-	handler := newHandler(options, glock.NewMockClock())
+	store := workerstoremocks.NewMockStore()
+	store.MarkErroredFunc.SetDefaultReturn(false, nil)
+	handler := newHandler(QueueOptions{Store: store})
 
-	if err := handler.markErrored(context.Background(), "test_queue", "deadbeef", 42, "OH NO"); err != ErrUnknownJob {
+	if err := handler.markErrored(context.Background(), "deadbeef", 42, "OH NO"); err != ErrUnknownJob {
 		t.Fatalf("unexpected error. want=%q have=%q", ErrUnknownJob, err)
 	}
 }
 
-func TestMarkErroredUnknownQueue(t *testing.T) {
-	options := Options{}
-	handler := newHandler(options, glock.NewMockClock())
-
-	if err := handler.markErrored(context.Background(), "test_queue", "deadbeef", 42, "OH NO"); err != ErrUnknownQueue {
-		t.Fatalf("unexpected error. want=%q have=%q", ErrUnknownQueue, err)
-	}
-}
-
 func TestMarkFailed(t *testing.T) {
 	store := workerstoremocks.NewMockStore()
 	store.DequeueFunc.SetDefaultReturn(testRecord{ID: 42}, true, nil)
+	store.MarkFailedFunc.SetDefaultReturn(true, nil)
 	recordTransformer := func(ctx context.Context, record workerutil.Record) (apiclient.Job, error) {
 		return apiclient.Job{ID: 42}, nil
 	}
 
-	options := Options{
-		QueueOptions: map[string]QueueOptions{
-			"test_queue": {Store: store, RecordTransformer: recordTransformer},
-		},
-	}
-	handler := newHandler(options, glock.NewMockClock())
+	handler := newHandler(QueueOptions{Store: store, RecordTransformer: recordTransformer})
 
-	job, dequeued, err := handler.dequeue(context.Background(), "test_queue", "deadbeef", "test")
+	job, dequeued, err := handler.dequeue(context.Background(), "deadbeef", "test")
 	if err != nil {
 		t.Fatalf("unexpected error dequeueing job: %s", err)
 	}
 	if !dequeued {
 		t.Fatalf("expected a job to be dequeued")
 	}
 
-	if err := handler.markFailed(context.Background(), "test_queue", "deadbeef", job.ID, "OH NO"); err != nil {
+	if err := handler.markFailed(context.Background(), "deadbeef", job.ID, "OH NO"); err != nil {
 		t.Fatalf("unexpected error completing job: %s", err)
 	}
 
@@ -310,6 +233,16 @@ func TestMarkFailed(t *testing.T) {
 	}
 }
 
+func TestMarkFailedUnknownJob(t *testing.T) {
+	store := workerstoremocks.NewMockStore()
+	store.MarkFailedFunc.SetDefaultReturn(false, nil)
+	handler := newHandler(QueueOptions{Store: store})
+
+	if err := handler.markFailed(context.Background(), "deadbeef", 42, "OH NO"); err != ErrUnknownJob {
+		t.Fatalf("unexpected error. want=%q have=%q", ErrUnknownJob, err)
+	}
+}
+
 type testRecord struct {
 	ID      int
 	Payload string

enterprise/cmd/executor-queue/internal/server/lifecycle.go

@@ -3,123 +3,30 @@ package server
 import (
 	"context"
 
-	"github.com/hashicorp/go-multierror"
+	"github.com/sourcegraph/sourcegraph/internal/workerutil/dbworker/store"
 )
 
-// heartbeat will release the transaction for any job that is not confirmed to be in-progress
-// by the given executor. This method is called when the executor POSTs its in-progress job
-// identifiers to the /heartbeat route. This method returns the set of identifiers which the
-// executor erroneously claims to hold (and are sent back as a hint to stop processing).
-func (h *handler) heartbeat(ctx context.Context, executorName string, jobIDs []int) ([]int, error) {
-	deadJobs, unknownIDs, err := h.heartbeatJobs(ctx, executorName, jobIDs)
-	err2 := h.requeueJobs(ctx, deadJobs)
-	if err != nil && err2 != nil {
-		err = multierror.Append(err, err2)
+// heartbeat is called when the executor POSTs its in-progress job identifiers to the /heartbeat route.
+// This method returns the set of identifiers which the executor erroneously claims to hold (and are sent
+// back as a hint to stop processing).
+// The set of job identifiers assigned to the given executor are getting a heartbeat, indicating they're
+// still being worked on.
+func (h *handler) heartbeat(ctx context.Context, executorName string, ids []int) (unknownInQueueJobs []int, errs error) {
+	knownIDs, err := h.Store.Heartbeat(ctx, ids, store.HeartbeatOptions{WorkerHostname: executorName})
+	if err != nil {
+		return nil, err
 	}
-	if err2 != nil {
-		err = err2
+	knownIDsMap := map[int]struct{}{}
+	for _, id := range knownIDs {
+		knownIDsMap[id] = struct{}{}
 	}
-	return unknownIDs, err
-}
-
-// cleanup will release the transactions held by any executor that has not sent a heartbeat
-// in a while. This method is called periodically in the background.
-func (h *handler) cleanup(ctx context.Context) error {
-	return h.requeueJobs(ctx, h.pruneExecutors())
-}
 
-// heartbeatJobs updates the set of job identifiers assigned to the given executor and returns
-// any job that was known to us but not reported by the executor, plus the set of job identifiers
-// reported by the executor which do not have an associated record held by this instance of the
-// executor queue. This can occur when the executor-queue restarts and loses its in-memory state.
-// We send these identifiers back to the executor as a hint to stop processing.
-func (h *handler) heartbeatJobs(ctx context.Context, executorName string, ids []int) (dead []jobMeta, unknownIDs []int, errs error) {
-	now := h.clock.Now()
-
-	executorIDsMap := map[int]struct{}{}
+	unknownInQueueJobs = make([]int, 0)
 	for _, id := range ids {
-		executorIDsMap[id] = struct{}{}
-	}
-
-	h.m.Lock()
-	defer h.m.Unlock()
-
-	executor, ok := h.executors[executorName]
-	if !ok {
-		executor = &executorMeta{}
-		h.executors[executorName] = executor
-	}
-
-	executorQueueIDsMap := map[int]struct{}{}
-	var live []jobMeta
-	for _, job := range executor.jobs {
-		executorQueueIDsMap[job.record.RecordID()] = struct{}{}
-		if _, ok := executorIDsMap[job.record.RecordID()]; ok || now.Sub(job.started) < h.options.UnreportedMaxAge {
-			live = append(live, job)
-			if err := h.heartbeatJob(ctx, job); err != nil {
-				errs = multierror.Append(errs, err)
-			}
-		} else {
-			dead = append(dead, job)
+		if _, ok := knownIDsMap[id]; !ok {
+			unknownInQueueJobs = append(unknownInQueueJobs, id)
 		}
 	}
-	executor.jobs = live
-	executor.lastUpdate = now
-
-	unknownIDs = make([]int, 0, len(ids))
-	for _, id := range ids {
-		if _, ok := executorQueueIDsMap[id]; !ok {
-			unknownIDs = append(unknownIDs, id)
-		}
-	}
-
-	return dead, unknownIDs, errs
-}
-
-// pruneExecutors will release the transactions held by any executor that has not sent a
-// heartbeat in a while and return the attached jobs.
-func (h *handler) pruneExecutors() (jobs []jobMeta) {
-	h.m.Lock()
-	defer h.m.Unlock()
-
-	for name, executor := range h.executors {
-		if h.clock.Now().Sub(executor.lastUpdate) <= h.options.DeathThreshold {
-			continue
-		}
-
-		jobs = append(jobs, executor.jobs...)
-		delete(h.executors, name)
-	}
-
-	return jobs
-}
-
-func (h *handler) heartbeatJob(ctx context.Context, job jobMeta) error {
-	queueOptions, ok := h.options.QueueOptions[job.queueName]
-	if !ok {
-		return ErrUnknownQueue
-	}
-
-	return queueOptions.Store.Heartbeat(ctx, job.record.RecordID())
-}
-
-// requeueJobs releases and requeues each of the given jobs.
-func (h *handler) requeueJobs(ctx context.Context, jobs []jobMeta) (errs error) {
-	for _, job := range jobs {
-		if err := h.requeueJob(ctx, job); err != nil {
-			errs = multierror.Append(errs, err)
-		}
-	}
-
-	return errs
-}
-
-// requeueJob requeues the given job and releases the associated transaction.
-func (h *handler) requeueJob(ctx context.Context, job jobMeta) error {
-	queueOptions, ok := h.options.QueueOptions[job.queueName]
-	if !ok {
-		return ErrUnknownQueue
-	}
 
-	return queueOptions.Store.Requeue(ctx, job.record.RecordID(), h.clock.Now().Add(h.options.RequeueDelay))
+	return unknownInQueueJobs, errs
 }

enterprise/cmd/executor-queue/internal/server/lifecycle_test.go

@@ -3,135 +3,81 @@ package server
 import (
 	"context"
 	"testing"
-	"time"
 
-	"github.com/derision-test/glock"
 	"github.com/google/go-cmp/cmp"
 
 	apiclient "github.com/sourcegraph/sourcegraph/enterprise/internal/executor"
 	"github.com/sourcegraph/sourcegraph/internal/workerutil"
+	"github.com/sourcegraph/sourcegraph/internal/workerutil/dbworker/store"
 	workerstoremocks "github.com/sourcegraph/sourcegraph/internal/workerutil/dbworker/store/mocks"
 )
 
 func TestHeartbeat(t *testing.T) {
 	store1 := workerstoremocks.NewMockStore()
-	store2 := workerstoremocks.NewMockStore()
 	recordTransformer := func(ctx context.Context, record workerutil.Record) (apiclient.Job, error) {
 		return apiclient.Job{ID: record.RecordID()}, nil
 	}
 
 	store1.DequeueFunc.PushReturn(testRecord{ID: 41}, true, nil)
 	store1.DequeueFunc.PushReturn(testRecord{ID: 42}, true, nil)
-	store2.DequeueFunc.PushReturn(testRecord{ID: 43}, true, nil)
-	store2.DequeueFunc.PushReturn(testRecord{ID: 44}, true, nil)
+	store1.DequeueFunc.PushReturn(testRecord{ID: 43}, true, nil)
+	store1.DequeueFunc.PushReturn(testRecord{ID: 44}, true, nil)
+	store1.HeartbeatFunc.SetDefaultHook(func(ctx context.Context, ids []int, options store.HeartbeatOptions) ([]int, error) {
+		knownIDs := make([]int, 0)
+		for _, id := range ids {
+			if id >= 41 && id <= 44 {
+				knownIDs = append(knownIDs, id)
+			}
+		}
 
-	options := Options{
-		QueueOptions: map[string]QueueOptions{
-			"q1": {Store: store1, RecordTransformer: recordTransformer},
-			"q2": {Store: store2, RecordTransformer: recordTransformer},
-		},
-		UnreportedMaxAge: time.Second,
-	}
-	clock := glock.NewMockClock()
-	handler := newHandler(options, clock)
+		return knownIDs, nil
+	})
+
+	handler := newHandler(QueueOptions{Store: store1, RecordTransformer: recordTransformer})
 
-	_, dequeued1, _ := handler.dequeue(context.Background(), "q1", "deadbeef", "test")
-	_, dequeued2, _ := handler.dequeue(context.Background(), "q1", "deadveal", "test")
-	_, dequeued3, _ := handler.dequeue(context.Background(), "q2", "deadbeef", "test")
-	_, dequeued4, _ := handler.dequeue(context.Background(), "q2", "deadveal", "test")
+	_, dequeued1, _ := handler.dequeue(context.Background(), "deadbeef", "test")
+	_, dequeued2, _ := handler.dequeue(context.Background(), "deadveal", "test")
+	_, dequeued3, _ := handler.dequeue(context.Background(), "deadbeef", "test")
+	_, dequeued4, _ := handler.dequeue(context.Background(), "deadveal", "test")
 	if !dequeued1 || !dequeued2 || !dequeued3 || !dequeued4 {
 		t.Fatalf("failed to dequeue records")
 	}
 
-	// missing all jobs, but they're less than UnreportedMaxAge
-	clock.Advance(time.Second / 2)
-	if _, err := handler.heartbeat(context.Background(), "deadbeef", []int{}); err != nil {
-		t.Fatalf("unexpected error performing heartbeat: %s", err)
-	}
-	if _, err := handler.heartbeat(context.Background(), "deadveal", []int{}); err != nil {
-		t.Fatalf("unexpected error performing heartbeat: %s", err)
-	}
-
 	// missing no jobs
-	clock.Advance(time.Minute * 2)
-	if _, err := handler.heartbeat(context.Background(), "deadbeef", []int{41, 43}); err != nil {
+	if ids, err := handler.heartbeat(context.Background(), "deadbeef", []int{41, 43}); err != nil {
 		t.Fatalf("unexpected error performing heartbeat: %s", err)
+	} else if diff := cmp.Diff(ids, []int{}); diff != "" {
+		t.Fatalf("invalid unknownIDs returned diff=%s", diff)
 	}
-	if _, err := handler.heartbeat(context.Background(), "deadveal", []int{42, 44}); err != nil {
+	if ids, err := handler.heartbeat(context.Background(), "deadveal", []int{42, 44}); err != nil {
 		t.Fatalf("unexpected error performing heartbeat: %s", err)
+	} else if diff := cmp.Diff(ids, []int{}); diff != "" {
+		t.Fatalf("invalid unknownIDs returned diff=%s", diff)
 	}
 
 	// missing one deadbeef jobs
-	clock.Advance(time.Minute * 2)
-	if _, err := handler.heartbeat(context.Background(), "deadbeef", []int{41}); err != nil {
-		t.Fatalf("unexpected error performing heartbeat: %s", err)
-	}
-	if _, err := handler.heartbeat(context.Background(), "deadveal", []int{42, 44}); err != nil {
+	if ids, err := handler.heartbeat(context.Background(), "deadbeef", []int{41}); err != nil {
 		t.Fatalf("unexpected error performing heartbeat: %s", err)
+	} else if diff := cmp.Diff(ids, []int{}); diff != "" {
+		t.Fatalf("invalid unknownIDs returned diff=%s", diff)
 	}
 
 	// missing two deadveal jobs
-	clock.Advance(time.Minute * 2)
-	if _, err := handler.heartbeat(context.Background(), "deadbeef", []int{41}); err != nil {
-		t.Fatalf("unexpected error performing heartbeat: %s", err)
-	}
-	if _, err := handler.heartbeat(context.Background(), "deadveal", []int{}); err != nil {
+	if ids, err := handler.heartbeat(context.Background(), "deadbeef", []int{}); err != nil {
 		t.Fatalf("unexpected error performing heartbeat: %s", err)
+	} else if diff := cmp.Diff(ids, []int{}); diff != "" {
+		t.Fatalf("invalid unknownIDs returned diff=%s", diff)
 	}
 
 	// unknown jobs
-	clock.Advance(time.Minute * 2)
 	if unknownIDs, err := handler.heartbeat(context.Background(), "deadbeef", []int{41, 43, 45}); err != nil {
 		t.Fatalf("unexpected error performing heartbeat: %s", err)
-	} else if diff := cmp.Diff([]int{43, 45}, unknownIDs); diff != "" {
+	} else if diff := cmp.Diff([]int{45}, unknownIDs); diff != "" {
 		t.Errorf("unexpected unknown ids (-want +got):\n%s", diff)
 	}
 	if unknownIDs, err := handler.heartbeat(context.Background(), "deadveal", []int{42, 44, 45}); err != nil {
 		t.Fatalf("unexpected error performing heartbeat: %s", err)
-	} else if diff := cmp.Diff([]int{42, 44, 45}, unknownIDs); diff != "" {
+	} else if diff := cmp.Diff([]int{45}, unknownIDs); diff != "" {
 		t.Errorf("unexpected unknown ids (-want +got):\n%s", diff)
 	}
 }
-
-func TestCleanup(t *testing.T) {
-	store1 := workerstoremocks.NewMockStore()
-	store2 := workerstoremocks.NewMockStore()
-	recordTransformer := func(ctx context.Context, record workerutil.Record) (apiclient.Job, error) {
-		return apiclient.Job{ID: record.RecordID()}, nil
-	}
-
-	store1.DequeueFunc.PushReturn(testRecord{ID: 41}, true, nil)
-	store1.DequeueFunc.PushReturn(testRecord{ID: 42}, true, nil)
-	store2.DequeueFunc.PushReturn(testRecord{ID: 43}, true, nil)
-	store2.DequeueFunc.PushReturn(testRecord{ID: 44}, true, nil)
-
-	options := Options{
-		QueueOptions: map[string]QueueOptions{
-			"q1": {Store: store1, RecordTransformer: recordTransformer},
-			"q2": {Store: store2, RecordTransformer: recordTransformer},
-		},
-		DeathThreshold: time.Minute * 5,
-	}
-	clock := glock.NewMockClock()
-	handler := newHandler(options, clock)
-
-	_, dequeued1, _ := handler.dequeue(context.Background(), "q1", "deadbeef", "test")
-	_, dequeued2, _ := handler.dequeue(context.Background(), "q1", "deadveal", "test")
-	_, dequeued3, _ := handler.dequeue(context.Background(), "q2", "deadbeef", "test")
-	_, dequeued4, _ := handler.dequeue(context.Background(), "q2", "deadveal", "test")
-	if !dequeued1 || !dequeued2 || !dequeued3 || !dequeued4 {
-		t.Fatalf("failed to dequeue records")
-	}
-
-	for i := 0; i < 6; i++ {
-		clock.Advance(time.Minute)
-
-		if _, err := handler.heartbeat(context.Background(), "deadbeef", []int{41, 43}); err != nil {
-			t.Fatalf("unexpected error performing heartbeat: %s", err)
-		}
-	}
-
-	if err := handler.cleanup(context.Background()); err != nil {
-		t.Fatalf("unexpected error performing cleanup: %s", err)
-	}
-}

enterprise/cmd/executor-queue/internal/server/routes.go

@@ -7,40 +7,40 @@ import (
 	"io"
 	"net/http"
 	"regexp"
-	"strings"
 
 	"github.com/gorilla/mux"
 	"github.com/inconshreveable/log15"
 
 	apiclient "github.com/sourcegraph/sourcegraph/enterprise/internal/executor"
 )
 
-func (h *handler) setupRoutes(router *mux.Router) {
-	var names []string
-	for queueName := range h.options.QueueOptions {
-		names = append(names, regexp.QuoteMeta(queueName))
-	}
-
-	routes := map[string]func(w http.ResponseWriter, r *http.Request){
-		"dequeue":              h.handleDequeue,
-		"addExecutionLogEntry": h.handleAddExecutionLogEntry,
-		"markComplete":         h.handleMarkComplete,
-		"markErrored":          h.handleMarkErrored,
-		"markFailed":           h.handleMarkFailed,
-	}
-	for path, handler := range routes {
-		router.Path(fmt.Sprintf("/{queueName:(?:%s)}/%s", strings.Join(names, "|"), path)).Methods("POST").HandlerFunc(handler)
+func setupRoutes(queueOptionsMap map[string]QueueOptions) func(router *mux.Router) {
+	return func(router *mux.Router) {
+		for name, queueOptions := range queueOptionsMap {
+			h := newHandler(queueOptions)
+
+			subRouter := router.PathPrefix(fmt.Sprintf("/{queueName:(?:%s)}/", regexp.QuoteMeta(name))).Subrouter()
+			routes := map[string]func(w http.ResponseWriter, r *http.Request){
+				"dequeue":              h.handleDequeue,
+				"addExecutionLogEntry": h.handleAddExecutionLogEntry,
+				"markComplete":         h.handleMarkComplete,
+				"markErrored":          h.handleMarkErrored,
+				"markFailed":           h.handleMarkFailed,
+				"heartbeat":            h.handleHeartbeat,
+			}
+			for path, handler := range routes {
+				subRouter.Path(fmt.Sprintf("/%s", path)).Methods("POST").HandlerFunc(handler)
+			}
+		}
 	}
-
-	router.Path("/heartbeat").Methods("POST").HandlerFunc(h.handleHeartbeat)
 }
 
 // POST /{queueName}/dequeue
 func (h *handler) handleDequeue(w http.ResponseWriter, r *http.Request) {
 	var payload apiclient.DequeueRequest
 
 	h.wrapHandler(w, r, &payload, func() (int, interface{}, error) {
-		job, dequeued, err := h.dequeue(r.Context(), mux.Vars(r)["queueName"], payload.ExecutorName, payload.ExecutorHostname)
+		job, dequeued, err := h.dequeue(r.Context(), payload.ExecutorName, payload.ExecutorHostname)
 		if !dequeued {
 			return http.StatusNoContent, nil, err
 		}
@@ -54,7 +54,7 @@ func (h *handler) handleAddExecutionLogEntry(w http.ResponseWriter, r *http.Requ
 	var payload apiclient.AddExecutionLogEntryRequest
 
 	h.wrapHandler(w, r, &payload, func() (int, interface{}, error) {
-		err := h.addExecutionLogEntry(r.Context(), mux.Vars(r)["queueName"], payload.ExecutorName, payload.JobID, payload.ExecutionLogEntry)
+		err := h.addExecutionLogEntry(r.Context(), payload.ExecutorName, payload.JobID, payload.ExecutionLogEntry)
 		return http.StatusNoContent, nil, err
 	})
 }
@@ -64,7 +64,7 @@ func (h *handler) handleMarkComplete(w http.ResponseWriter, r *http.Request) {
 	var payload apiclient.MarkCompleteRequest
 
 	h.wrapHandler(w, r, &payload, func() (int, interface{}, error) {
-		err := h.markComplete(r.Context(), mux.Vars(r)["queueName"], payload.ExecutorName, payload.JobID)
+		err := h.markComplete(r.Context(), payload.ExecutorName, payload.JobID)
 		if err == ErrUnknownJob {
 			return http.StatusNotFound, nil, nil
 		}
@@ -78,7 +78,7 @@ func (h *handler) handleMarkErrored(w http.ResponseWriter, r *http.Request) {
 	var payload apiclient.MarkErroredRequest
 
 	h.wrapHandler(w, r, &payload, func() (int, interface{}, error) {
-		err := h.markErrored(r.Context(), mux.Vars(r)["queueName"], payload.ExecutorName, payload.JobID, payload.ErrorMessage)
+		err := h.markErrored(r.Context(), payload.ExecutorName, payload.JobID, payload.ErrorMessage)
 		if err == ErrUnknownJob {
 			return http.StatusNotFound, nil, nil
 		}
@@ -92,7 +92,7 @@ func (h *handler) handleMarkFailed(w http.ResponseWriter, r *http.Request) {
 	var payload apiclient.MarkErroredRequest
 
 	h.wrapHandler(w, r, &payload, func() (int, interface{}, error) {
-		err := h.markFailed(r.Context(), mux.Vars(r)["queueName"], payload.ExecutorName, payload.JobID, payload.ErrorMessage)
+		err := h.markFailed(r.Context(), payload.ExecutorName, payload.JobID, payload.ErrorMessage)
 		if err == ErrUnknownJob {
 			return http.StatusNotFound, nil, nil
 		}
@@ -101,7 +101,7 @@ func (h *handler) handleMarkFailed(w http.ResponseWriter, r *http.Request) {
 	})
 }
 
-// POST /heartbeat
+// POST /{queueName}/heartbeat
 func (h *handler) handleHeartbeat(w http.ResponseWriter, r *http.Request) {
 	var payload apiclient.HeartbeatRequest
 

enterprise/cmd/executor-queue/internal/server/server.go

@@ -1,32 +1,25 @@
 package server
 
 import (
-	"context"
 	"fmt"
 	"net/http"
 
-	"github.com/derision-test/glock"
-	"github.com/inconshreveable/log15"
-
 	"github.com/sourcegraph/sourcegraph/internal/goroutine"
 	"github.com/sourcegraph/sourcegraph/internal/httpserver"
-	"github.com/sourcegraph/sourcegraph/internal/observation"
 	"github.com/sourcegraph/sourcegraph/internal/trace/ot"
 )
 
+// ServerOptions captures the options required for setting up an executor queue
+// server.
+type ServerOptions struct {
+	Port int
+}
+
 // NewServer returns an HTTP job queue server.
-func NewServer(options Options, observationContext *observation.Context) goroutine.BackgroundRoutine {
+func NewServer(options ServerOptions, queueOptions map[string]QueueOptions) goroutine.BackgroundRoutine {
 	addr := fmt.Sprintf(":%d", options.Port)
-	handler := newHandlerWithMetrics(options, glock.NewRealClock(), observationContext)
-	httpHandler := ot.Middleware(httpserver.NewHandler(handler.setupRoutes))
+	router := setupRoutes(queueOptions)
+	httpHandler := ot.Middleware(httpserver.NewHandler(router))
 	server := httpserver.NewFromAddr(addr, &http.Server{Handler: httpHandler})
-	janitor := goroutine.NewPeriodicGoroutine(context.Background(), options.CleanupInterval, &handlerWrapper{handler})
-	return goroutine.CombinedRoutine{server, janitor}
+	return server
 }
-
-type handlerWrapper struct{ handler *handler }
-
-var _ goroutine.Handler = &handlerWrapper{}
-
-func (hw *handlerWrapper) Handle(ctx context.Context) error { return hw.handler.cleanup(ctx) }
-func (hw *handlerWrapper) HandleError(err error)            { log15.Error("Failed to requeue jobs", "err", err) }

enterprise/cmd/executor-queue/main.go

@@ -97,7 +97,7 @@ func main() {
 		}))
 	}
 
-	server := apiserver.NewServer(serviceConfig.ServerOptions(queueOptions), observationContext)
+	server := apiserver.NewServer(serviceConfig.ServerOptions(), queueOptions)
 	goroutine.MonitorBackgroundRoutines(context.Background(), server)
 }
 

enterprise/cmd/executor/config.go

@@ -10,6 +10,7 @@ import (
 	"github.com/sourcegraph/sourcegraph/enterprise/cmd/executor/internal/command"
 	apiworker "github.com/sourcegraph/sourcegraph/enterprise/cmd/executor/internal/worker"
 	"github.com/sourcegraph/sourcegraph/internal/env"
+	"github.com/sourcegraph/sourcegraph/internal/hostname"
 	"github.com/sourcegraph/sourcegraph/internal/workerutil"
 )
 
@@ -97,8 +98,12 @@ func (c *Config) ResourceOptions() command.ResourceOptions {
 }
 
 func (c *Config) ClientOptions(transport http.RoundTripper) apiclient.Options {
+	hn := hostname.Get()
+
 	return apiclient.Options{
-		ExecutorName:      uuid.New().String(),
+		// Be unique but also descriptive.
+		ExecutorName:      hn + "-" + uuid.New().String(),
+		ExecutorHostname:  hn,
 		PathPrefix:        "/.executors/queue",
 		EndpointOptions:   c.EndpointOptions(),
 		BaseClientOptions: c.BaseClientOptions(transport),

enterprise/cmd/executor/internal/apiclient/client.go

@@ -151,8 +151,8 @@ func (c *Client) MarkFailed(ctx context.Context, queueName string, jobID int, er
 	return c.client.DoAndDrop(ctx, req)
 }
 
-func (c *Client) Ping(ctx context.Context, jobIDs []int) (err error) {
-	req, err := c.makeRequest("POST", "heartbeat", executor.HeartbeatRequest{
+func (c *Client) Ping(ctx context.Context, queueName string, jobIDs []int) (err error) {
+	req, err := c.makeRequest("POST", fmt.Sprintf("%s/heartbeat", queueName), executor.HeartbeatRequest{
 		ExecutorName: c.options.ExecutorName,
 	})
 	if err != nil {
@@ -162,13 +162,14 @@ func (c *Client) Ping(ctx context.Context, jobIDs []int) (err error) {
 	return c.client.DoAndDrop(ctx, req)
 }
 
-func (c *Client) Heartbeat(ctx context.Context, jobIDs []int) (unknownIDs []int, err error) {
+func (c *Client) Heartbeat(ctx context.Context, queueName string, jobIDs []int) (unknownIDs []int, err error) {
 	ctx, endObservation := c.operations.heartbeat.With(ctx, &err, observation.Args{LogFields: []log.Field{
+		log.String("queueName", queueName),
 		log.String("jobIDs", intsToString(jobIDs)),
 	}})
 	defer endObservation(1, observation.Args{})
 
-	req, err := c.makeRequest("POST", "heartbeat", executor.HeartbeatRequest{
+	req, err := c.makeRequest("POST", fmt.Sprintf("%s/heartbeat", queueName), executor.HeartbeatRequest{
 		ExecutorName: c.options.ExecutorName,
 		JobIDs:       jobIDs,
 	})

enterprise/cmd/executor/internal/apiclient/client_test.go

@@ -247,7 +247,7 @@ func TestMarkFailed(t *testing.T) {
 func TestHeartbeat(t *testing.T) {
 	spec := routeSpec{
 		expectedMethod:   "POST",
-		expectedPath:     "/.executors/queue/heartbeat",
+		expectedPath:     "/.executors/queue/test_queue/heartbeat",
 		expectedUsername: "test",
 		expectedPassword: "hunter2",
 		expectedPayload:  `{"executorName": "deadbeef", "jobIds": [1, 2, 3]}`,
@@ -256,7 +256,7 @@ func TestHeartbeat(t *testing.T) {
 	}
 
 	testRoute(t, spec, func(client *Client) {
-		unknownIDs, err := client.Heartbeat(context.Background(), []int{1, 2, 3})
+		unknownIDs, err := client.Heartbeat(context.Background(), "test_queue", []int{1, 2, 3})
 		if err != nil {
 			t.Fatalf("unexpected error performing heartbeat: %s", err)
 		}
@@ -270,7 +270,7 @@ func TestHeartbeat(t *testing.T) {
 func TestHeartbeatBadResponse(t *testing.T) {
 	spec := routeSpec{
 		expectedMethod:   "POST",
-		expectedPath:     "/.executors/queue/heartbeat",
+		expectedPath:     "/.executors/queue/test_queue/heartbeat",
 		expectedUsername: "test",
 		expectedPassword: "hunter2",
 		expectedPayload:  `{"executorName": "deadbeef", "jobIds": [1, 2, 3]}`,
@@ -279,7 +279,7 @@ func TestHeartbeatBadResponse(t *testing.T) {
 	}
 
 	testRoute(t, spec, func(client *Client) {
-		if _, err := client.Heartbeat(context.Background(), []int{1, 2, 3}); err == nil {
+		if _, err := client.Heartbeat(context.Background(), "test_queue", []int{1, 2, 3}); err == nil {
 			t.Fatalf("expected an error")
 		}
 	})

enterprise/cmd/executor/internal/worker/store.go

@@ -38,9 +38,9 @@ func (s *storeShim) Dequeue(ctx context.Context, workerHostname string, extraArg
 	return job, dequeued, nil
 }
 
-func (s *storeShim) Heartbeat(ctx context.Context, id int) error {
+func (s *storeShim) Heartbeat(ctx context.Context, ids []int) ([]int, error) {
 	// Not needed, we do bulk updates from the executor.
-	return nil
+	return nil, nil
 }
 
 func (s *storeShim) AddExecutionLogEntry(ctx context.Context, id int, entry workerutil.ExecutionLogEntry) error {

enterprise/cmd/executor/internal/worker/worker.go

@@ -75,9 +75,9 @@ func NewWorker(options Options, observationContext *observation.Context) gorouti
 		runnerFactory: command.NewRunner,
 	}
 
-	indexer := workerutil.NewWorker(context.Background(), store, handler, options.WorkerOptions)
+	worker := workerutil.NewWorker(context.Background(), store, handler, options.WorkerOptions)
 	heartbeat := goroutine.NewHandlerWithErrorMessage("heartbeat", func(ctx context.Context) error {
-		unknownIDs, err := queueStore.Heartbeat(ctx, idSet.Slice())
+		unknownIDs, err := queueStore.Heartbeat(ctx, options.QueueName, idSet.Slice())
 		if err != nil {
 			return err
 		}
@@ -90,7 +90,7 @@ func NewWorker(options Options, observationContext *observation.Context) gorouti
 	})
 
 	return goroutine.CombinedRoutine{
-		indexer,
+		worker,
 		goroutine.NewPeriodicGoroutine(context.Background(), options.HeartbeatInterval, heartbeat),
 	}
 }
@@ -111,7 +111,7 @@ func connectToFrontend(queueStore *apiclient.Client, options Options) bool {
 	defer signal.Stop(signals)
 
 	for {
-		err := queueStore.Ping(context.Background(), nil)
+		err := queueStore.Ping(context.Background(), options.QueueName, nil)
 		if err == nil {
 			log15.Info("Connected to Sourcegraph instance")
 			return true

enterprise/cmd/frontend/internal/executor/internal_proxy_handler.go

@@ -63,7 +63,7 @@ func newInternalProxyHandler(uploadHandler http.Handler) (func() http.Handler, e
 		base.Path("/git/{rest:.*/(?:info/refs|git-upload-pack)}").Handler(reverseProxy(frontendOrigin))
 
 		// Proxy only the known routes in the executor queue API
-		base.Path("/queue/{rest:heartbeat|.*/(?:dequeue|addExecutionLogEntry|markComplete|markErrored|markFailed)}").Handler(reverseProxy(queueOrigin))
+		base.Path("/queue/{rest:.*/(?:dequeue|addExecutionLogEntry|markComplete|markErrored|markFailed|heartbeat)}").Handler(reverseProxy(queueOrigin))
 
 		// Upload LSIF indexes without a sudo access token or github tokens
 		base.Path("/lsif/upload").Methods("POST").Handler(uploadHandler)

enterprise/internal/batches/background/executor_store.go

@@ -64,20 +64,20 @@ type executorStore struct {
 const markCompleteQuery = `
 UPDATE batch_spec_executions
 SET state = 'completed', finished_at = clock_timestamp(), batch_spec_id = (SELECT id FROM batch_specs WHERE rand_id = %s)
-WHERE id = %s AND state = 'processing'
+WHERE id = %s AND state = 'processing' AND worker_hostname = %s
 RETURNING id
 `
 
-func (s *executorStore) MarkComplete(ctx context.Context, id int) (_ bool, err error) {
+func (s *executorStore) MarkComplete(ctx context.Context, id int, options dbworkerstore.MarkFinalOptions) (_ bool, err error) {
 	batchesStore := store.New(s.Store.Handle().DB(), nil)
 
 	batchSpecRandID, err := loadAndExtractBatchSpecRandID(ctx, batchesStore, int64(id))
 	if err != nil {
 		// If we couldn't extract the batch spec rand id, we mark the job as failed
-		return s.Store.MarkFailed(ctx, id, fmt.Sprintf("failed to extract batch spec ID: %s", err))
+		return s.Store.MarkFailed(ctx, id, fmt.Sprintf("failed to extract batch spec ID: %s", err), options)
 	}
 
-	_, ok, err := basestore.ScanFirstInt(batchesStore.Query(ctx, sqlf.Sprintf(markCompleteQuery, batchSpecRandID, id)))
+	_, ok, err := basestore.ScanFirstInt(batchesStore.Query(ctx, sqlf.Sprintf(markCompleteQuery, batchSpecRandID, id, options.WorkerHostname)))
 	return ok, err
 }
 

enterprise/internal/batches/background/executor_store_test.go

@@ -56,7 +56,7 @@ stdout: {"operation":"CREATING_BATCH_SPEC","timestamp":"2021-07-06T09:38:51.535Z
 		}
 
 		for _, e := range entries {
-			err := workStore.AddExecutionLogEntry(context.Background(), int(specExec.ID), e)
+			err := workStore.AddExecutionLogEntry(context.Background(), int(specExec.ID), e, dbworkerstore.AddExecutionLogEntryOptions{})
 			if err != nil {
 				t.Fatal(err)
 			}

internal/database/schema.md

@@ -318,6 +318,7 @@ Referenced by:
  trigger_event     | integer                  |           |          | 
  worker_hostname   | text                     |           | not null | ''::text
  last_heartbeat_at | timestamp with time zone |           |          | 
+ execution_logs    | json[]                   |           |          | 
 Indexes:
     "cm_action_jobs_pkey" PRIMARY KEY, btree (id)
 Foreign-key constraints:
@@ -437,6 +438,7 @@ Foreign-key constraints:
  num_results       | integer                  |           |          | 
  worker_hostname   | text                     |           | not null | ''::text
  last_heartbeat_at | timestamp with time zone |           |          | 
+ execution_logs    | json[]                   |           |          | 
 Indexes:
     "cm_trigger_jobs_pkey" PRIMARY KEY, btree (id)
 Foreign-key constraints:

internal/workerutil/dbworker/store/helpers_test.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/derision-test/glock"
 	"github.com/keegancsmith/sqlf"
+	"github.com/lib/pq"
 
 	"github.com/sourcegraph/sourcegraph/internal/database/basestore"
 	"github.com/sourcegraph/sourcegraph/internal/database/dbtesting"
@@ -20,8 +21,9 @@ func testStore(db dbutil.DB, options Options) *store {
 }
 
 type TestRecord struct {
-	ID    int
-	State string
+	ID            int
+	State         string
+	ExecutionLogs []ExecutionLogEntry
 }
 
 func (v TestRecord) RecordID() int {
@@ -36,7 +38,7 @@ func testScanFirstRecord(rows *sql.Rows, queryErr error) (v workerutil.Record, _
 
 	if rows.Next() {
 		var record TestRecord
-		if err := rows.Scan(&record.ID, &record.State); err != nil {
+		if err := rows.Scan(&record.ID, &record.State, pq.Array(&record.ExecutionLogs)); err != nil {
 			return nil, false, err
 		}
 
@@ -146,6 +148,7 @@ func defaultTestStoreOptions(clock glock.Clock) Options {
 		ColumnExpressions: []*sqlf.Query{
 			sqlf.Sprintf("w.id"),
 			sqlf.Sprintf("w.state"),
+			sqlf.Sprintf("w.execution_logs"),
 		},
 		StalledMaxAge: time.Second * 5,
 		MaxNumResets:  5,
@@ -170,6 +173,12 @@ func assertDequeueRecordResult(t *testing.T, expectedID int, record interface{},
 	}
 }
 
+func assertDequeueRecordResultLogCount(t *testing.T, expectedLogCount int, record interface{}) {
+	if val := len(record.(TestRecord).ExecutionLogs); val != expectedLogCount {
+		t.Errorf("unexpected count of logs. want=%d have=%d", expectedLogCount, val)
+	}
+}
+
 func assertDequeueRecordViewResult(t *testing.T, expectedID, expectedNewField int, record interface{}, ok bool, err error) {
 	if err != nil {
 		t.Fatalf("unexpected error: %s", err)

internal/workerutil/dbworker/store/observability.go

@@ -16,6 +16,7 @@ type operations struct {
 	markErrored          *observation.Operation
 	markFailed           *observation.Operation
 	resetStalled         *observation.Operation
+	heartbeat            *observation.Operation
 }
 
 func newOperations(storeName string, observationContext *observation.Context) *operations {
@@ -43,5 +44,6 @@ func newOperations(storeName string, observationContext *observation.Context) *o
 		markErrored:          op("MarkErrored"),
 		markFailed:           op("MarkFailed"),
 		resetStalled:         op("ResetStalled"),
+		heartbeat:            op("Heartbeat"),
 	}
 }

internal/workerutil/dbworker/store/store.go

@@ -19,6 +19,50 @@ import (
 	"github.com/sourcegraph/sourcegraph/internal/workerutil"
 )
 
+type HeartbeatOptions struct {
+	// WorkerHostname, if set, enforces worker_hostname to be set to a specific value.
+	WorkerHostname string
+}
+
+func (o *HeartbeatOptions) ToSQLConds(formatQuery func(query string, args ...interface{}) *sqlf.Query) []*sqlf.Query {
+	conds := []*sqlf.Query{}
+	if o.WorkerHostname != "" {
+		conds = append(conds, formatQuery("{worker_hostname} = %s", o.WorkerHostname))
+	}
+	return conds
+}
+
+type AddExecutionLogEntryOptions struct {
+	// WorkerHostname, if set, enforces worker_hostname to be set to a specific value.
+	WorkerHostname string
+	// State, if set, enforces state to be set to a specific value.
+	State string
+}
+
+func (o *AddExecutionLogEntryOptions) ToSQLConds(formatQuery func(query string, args ...interface{}) *sqlf.Query) []*sqlf.Query {
+	conds := []*sqlf.Query{}
+	if o.WorkerHostname != "" {
+		conds = append(conds, formatQuery("{worker_hostname} = %s", o.WorkerHostname))
+	}
+	if o.State != "" {
+		conds = append(conds, formatQuery("{state} = %s", o.State))
+	}
+	return conds
+}
+
+type MarkFinalOptions struct {
+	// WorkerHostname, if set, enforces worker_hostname to be set to a specific value.
+	WorkerHostname string
+}
+
+func (o *MarkFinalOptions) ToSQLConds(formatQuery func(query string, args ...interface{}) *sqlf.Query) []*sqlf.Query {
+	conds := []*sqlf.Query{}
+	if o.WorkerHostname != "" {
+		conds = append(conds, formatQuery("{worker_hostname} = %s", o.WorkerHostname))
+	}
+	return conds
+}
+
 // Store is the persistence layer for the dbworker package that handles worker-side operations backed by a Postgres
 // database. See Options for details on the required shape of the database tables (e.g. table column names/types).
 type Store interface {
@@ -35,29 +79,29 @@ type Store interface {
 	Dequeue(ctx context.Context, workerHostname string, conditions []*sqlf.Query) (workerutil.Record, bool, error)
 
 	// Heartbeat marks the given record as currently being processed.
-	Heartbeat(ctx context.Context, id int) error
+	Heartbeat(ctx context.Context, ids []int, options HeartbeatOptions) (knownIDs []int, err error)
 
 	// Requeue updates the state of the record with the given identifier to queued and adds a processing delay before
 	// the next dequeue of this record can be performed.
 	Requeue(ctx context.Context, id int, after time.Time) error
 
 	// AddExecutionLogEntry adds an executor log entry to the record.
-	AddExecutionLogEntry(ctx context.Context, id int, entry workerutil.ExecutionLogEntry) error
+	AddExecutionLogEntry(ctx context.Context, id int, entry workerutil.ExecutionLogEntry, options AddExecutionLogEntryOptions) error
 
 	// MarkComplete attempts to update the state of the record to complete. If this record has already been moved from
 	// the processing state to a terminal state, this method will have no effect. This method returns a boolean flag
 	// indicating if the record was updated.
-	MarkComplete(ctx context.Context, id int) (bool, error)
+	MarkComplete(ctx context.Context, id int, options MarkFinalOptions) (bool, error)
 
 	// MarkErrored attempts to update the state of the record to errored. This method will only have an effect
 	// if the current state of the record is processing or completed. A requeued record or a record already marked
 	// with an error will not be updated. This method returns a boolean flag indicating if the record was updated.
-	MarkErrored(ctx context.Context, id int, failureMessage string) (bool, error)
+	MarkErrored(ctx context.Context, id int, failureMessage string, options MarkFinalOptions) (bool, error)
 
 	// MarkFailed attempts to update the state of the record to failed. This method will only have an effect
 	// if the current state of the record is processing or completed. A requeued record or a record already marked
 	// with an error will not be updated. This method returns a boolean flag indicating if the record was updated.
-	MarkFailed(ctx context.Context, id int, failureMessage string) (bool, error)
+	MarkFailed(ctx context.Context, id int, failureMessage string, options MarkFinalOptions) (bool, error)
 
 	// ResetStalled moves all processing records that have not received a heartbeat within `StalledMaxAge` back to the
 	// queued state. In order to prevent input that continually crashes worker instances, records that have been reset
@@ -379,6 +423,7 @@ SET
 	{last_heartbeat_at} = %s,
 	{finished_at} = NULL,
 	{failure_message} = NULL,
+	{execution_logs} = NULL,
 	{worker_hostname} = %s
 WHERE {id} IN (SELECT {id} FROM candidate)
 RETURNING {id}
@@ -389,21 +434,51 @@ const selectRecordQuery = `
 SELECT %s FROM %s WHERE {id} = %s
 `
 
-func (s *store) Heartbeat(ctx context.Context, id int) error {
-	err := s.Exec(ctx, s.formatQuery(updateCandidateQuery, quote(s.options.TableName), s.now(), id))
-	if err != nil {
-		if err != ctx.Err() {
-			return err
-		}
+func (s *store) Heartbeat(ctx context.Context, ids []int, options HeartbeatOptions) (knownIDs []int, err error) {
+	ctx, endObservation := s.operations.heartbeat.With(ctx, &err, observation.Args{})
+	defer endObservation(1, observation.Args{})
+
+	if len(ids) == 0 {
+		return []int{}, nil
+	}
+
+	sqlIDs := make([]*sqlf.Query, 0, len(ids))
+	for _, id := range ids {
+		sqlIDs = append(sqlIDs, sqlf.Sprintf("%s", id))
+	}
+
+	quotedTableName := quote(s.options.TableName)
+
+	conds := []*sqlf.Query{
+		s.formatQuery("{id} in (%s)", sqlf.Join(sqlIDs, "")),
+		s.formatQuery("{state} = 'processing'"),
 	}
-	return nil
+	conds = append(conds, options.ToSQLConds(s.formatQuery)...)
+
+	knownIDs, err = basestore.ScanInts(s.Query(ctx, s.formatQuery(updateCandidateQuery, quotedTableName, sqlf.Join(conds, "AND"), quotedTableName, s.now())))
+	return knownIDs, err
 }
 
 const updateCandidateQuery = `
 -- source: internal/workerutil/store.go:Heartbeat
-UPDATE %s
-SET {last_heartbeat_at} = %s
-WHERE {id} = %s AND {state} = 'processing'
+WITH alive_candidates AS (
+	SELECT
+		id
+	FROM
+		%s
+	WHERE
+		%s
+	ORDER BY
+		id ASC
+	FOR UPDATE
+)
+UPDATE
+	%s
+SET
+	{last_heartbeat_at} = %s
+WHERE
+	id IN (SELECT id FROM alive_candidates)
+RETURNING id
 `
 
 // Requeue updates the state of the record with the given identifier to queued and adds a processing delay before
@@ -431,58 +506,75 @@ WHERE {id} = %s
 `
 
 // AddExecutionLogEntry adds an executor log entry to the record.
-func (s *store) AddExecutionLogEntry(ctx context.Context, id int, entry workerutil.ExecutionLogEntry) (err error) {
+func (s *store) AddExecutionLogEntry(ctx context.Context, id int, entry workerutil.ExecutionLogEntry, options AddExecutionLogEntryOptions) (err error) {
 	ctx, endObservation := s.operations.addExecutionLogEntry.With(ctx, &err, observation.Args{LogFields: []log.Field{
 		log.Int("id", id),
 	}})
 	defer endObservation(1, observation.Args{})
 
+	conds := []*sqlf.Query{
+		s.formatQuery("{id} = %s", id),
+	}
+	conds = append(conds, options.ToSQLConds(s.formatQuery)...)
+
 	return s.Exec(ctx, s.formatQuery(
 		addExecutionLogEntryQuery,
 		quote(s.options.TableName),
 		ExecutionLogEntry(entry),
-		id,
+		sqlf.Join(conds, "AND"),
 	))
 }
 
 const addExecutionLogEntryQuery = `
 -- source: internal/workerutil/store.go:AddExecutionLogEntry
 UPDATE %s
 SET {execution_logs} = {execution_logs} || %s::json
-WHERE {id} = %s
+WHERE %s
 `
 
 // MarkComplete attempts to update the state of the record to complete. If this record has already been moved from
 // the processing state to a terminal state, this method will have no effect. This method returns a boolean flag
 // indicating if the record was updated.
-func (s *store) MarkComplete(ctx context.Context, id int) (_ bool, err error) {
+func (s *store) MarkComplete(ctx context.Context, id int, options MarkFinalOptions) (_ bool, err error) {
 	ctx, endObservation := s.operations.markComplete.With(ctx, &err, observation.Args{LogFields: []log.Field{
 		log.Int("id", id),
 	}})
 	defer endObservation(1, observation.Args{})
 
-	_, ok, err := basestore.ScanFirstInt(s.Query(ctx, s.formatQuery(markCompleteQuery, quote(s.options.TableName), id)))
+	conds := []*sqlf.Query{
+		s.formatQuery("{id} = %s", id),
+		s.formatQuery("{state} = 'processing'"),
+	}
+	conds = append(conds, options.ToSQLConds(s.formatQuery)...)
+
+	_, ok, err := basestore.ScanFirstInt(s.Query(ctx, s.formatQuery(markCompleteQuery, quote(s.options.TableName), sqlf.Join(conds, "AND"))))
 	return ok, err
 }
 
 const markCompleteQuery = `
 -- source: internal/workerutil/store.go:MarkComplete
 UPDATE %s
 SET {state} = 'completed', {finished_at} = clock_timestamp()
-WHERE {id} = %s AND {state} = 'processing'
+WHERE %s
 RETURNING {id}
 `
 
 // MarkErrored attempts to update the state of the record to errored. This method will only have an effect
 // if the current state of the record is processing or completed. A requeued record or a record already marked
 // with an error will not be updated. This method returns a boolean flag indicating if the record was updated.
-func (s *store) MarkErrored(ctx context.Context, id int, failureMessage string) (_ bool, err error) {
+func (s *store) MarkErrored(ctx context.Context, id int, failureMessage string, options MarkFinalOptions) (_ bool, err error) {
 	ctx, endObservation := s.operations.markErrored.With(ctx, &err, observation.Args{LogFields: []log.Field{
 		log.Int("id", id),
 	}})
 	defer endObservation(1, observation.Args{})
 
-	q := s.formatQuery(markErroredQuery, quote(s.options.TableName), s.options.MaxNumRetries, failureMessage, id)
+	conds := []*sqlf.Query{
+		s.formatQuery("{id} = %s", id),
+		s.formatQuery("{state} = 'processing'"),
+	}
+	conds = append(conds, options.ToSQLConds(s.formatQuery)...)
+
+	q := s.formatQuery(markErroredQuery, quote(s.options.TableName), s.options.MaxNumRetries, failureMessage, sqlf.Join(conds, "AND"))
 	_, ok, err := basestore.ScanFirstInt(s.Query(ctx, q))
 	return ok, err
 }
@@ -494,20 +586,26 @@ SET {state} = CASE WHEN {num_failures} + 1 = %d THEN 'failed' ELSE 'errored' END
 	{finished_at} = clock_timestamp(),
 	{failure_message} = %s,
 	{num_failures} = {num_failures} + 1
-WHERE {id} = %s AND ({state} = 'processing' OR {state} = 'completed')
+WHERE %s
 RETURNING {id}
 `
 
 // MarkFailed attempts to update the state of the record to failed. This method will only have an effect
 // if the current state of the record is processing or completed. A requeued record or a record already marked
 // with an error will not be updated. This method returns a boolean flag indicating if the record was updated.
-func (s *store) MarkFailed(ctx context.Context, id int, failureMessage string) (_ bool, err error) {
+func (s *store) MarkFailed(ctx context.Context, id int, failureMessage string, options MarkFinalOptions) (_ bool, err error) {
 	ctx, endObservation := s.operations.markFailed.With(ctx, &err, observation.Args{LogFields: []log.Field{
 		log.Int("id", id),
 	}})
 	defer endObservation(1, observation.Args{})
 
-	q := s.formatQuery(markFailedQuery, quote(s.options.TableName), failureMessage, id)
+	conds := []*sqlf.Query{
+		s.formatQuery("{id} = %s", id),
+		s.formatQuery("{state} = 'processing'"),
+	}
+	conds = append(conds, options.ToSQLConds(s.formatQuery)...)
+
+	q := s.formatQuery(markFailedQuery, quote(s.options.TableName), failureMessage, sqlf.Join(conds, "AND"))
 	_, ok, err := basestore.ScanFirstInt(s.Query(ctx, q))
 	return ok, err
 }
@@ -519,7 +617,7 @@ SET {state} = 'failed',
 	{finished_at} = clock_timestamp(),
 	{failure_message} = %s,
 	{num_failures} = {num_failures} + 1
-WHERE {id} = %s AND ({state} = 'processing' OR {state} = 'completed')
+WHERE %s
 RETURNING {id}
 `
 

internal/workerutil/dbworker/store/store_test.go

@@ -149,6 +149,22 @@ func TestStoreDequeueConditions(t *testing.T) {
 	assertDequeueRecordResult(t, 3, record, ok, err)
 }
 
+func TestStoreDequeueResetExecutionLogs(t *testing.T) {
+	db := setupStoreTest(t)
+
+	if _, err := db.ExecContext(context.Background(), `
+		INSERT INTO workerutil_test (id, state, execution_logs, uploaded_at)
+		VALUES
+			(1, 'queued', E'{"{\\"key\\": \\"test\\"}"}', NOW() - '1 minute'::interval)
+	`); err != nil {
+		t.Fatalf("unexpected error inserting records: %s", err)
+	}
+
+	record, ok, err := testStore(db, defaultTestStoreOptions(nil)).Dequeue(context.Background(), "test", nil)
+	assertDequeueRecordResult(t, 1, record, ok, err)
+	assertDequeueRecordResultLogCount(t, 0, record)
+}
+
 func TestStoreDequeueDelay(t *testing.T) {
 	db := setupStoreTest(t)
 
@@ -384,7 +400,7 @@ func TestStoreAddExecutionLogEntry(t *testing.T) {
 			Command: command,
 			Out:     payload,
 		}
-		if err := testStore(db, defaultTestStoreOptions(nil)).AddExecutionLogEntry(context.Background(), 1, entry); err != nil {
+		if err := testStore(db, defaultTestStoreOptions(nil)).AddExecutionLogEntry(context.Background(), 1, entry, AddExecutionLogEntryOptions{}); err != nil {
 			t.Fatalf("unexpected error adding executor log entry: %s", err)
 		}
 	}
@@ -424,7 +440,7 @@ func TestStoreMarkComplete(t *testing.T) {
 		t.Fatalf("unexpected error inserting records: %s", err)
 	}
 
-	marked, err := testStore(db, defaultTestStoreOptions(nil)).MarkComplete(context.Background(), 1)
+	marked, err := testStore(db, defaultTestStoreOptions(nil)).MarkComplete(context.Background(), 1, MarkFinalOptions{})
 	if err != nil {
 		t.Fatalf("unexpected error marking record as completed: %s", err)
 	}
@@ -466,7 +482,7 @@ func TestStoreMarkCompleteNotProcessing(t *testing.T) {
 		t.Fatalf("unexpected error inserting records: %s", err)
 	}
 
-	marked, err := testStore(db, defaultTestStoreOptions(nil)).MarkComplete(context.Background(), 1)
+	marked, err := testStore(db, defaultTestStoreOptions(nil)).MarkComplete(context.Background(), 1, MarkFinalOptions{})
 	if err != nil {
 		t.Fatalf("unexpected error marking record as completed: %s", err)
 	}
@@ -508,7 +524,7 @@ func TestStoreMarkErrored(t *testing.T) {
 		t.Fatalf("unexpected error inserting records: %s", err)
 	}
 
-	marked, err := testStore(db, defaultTestStoreOptions(nil)).MarkErrored(context.Background(), 1, "new message")
+	marked, err := testStore(db, defaultTestStoreOptions(nil)).MarkErrored(context.Background(), 1, "new message", MarkFinalOptions{})
 	if err != nil {
 		t.Fatalf("unexpected error marking record as errored: %s", err)
 	}
@@ -550,7 +566,7 @@ func TestStoreMarkFailed(t *testing.T) {
 		t.Fatalf("unexpected error inserting records: %s", err)
 	}
 
-	marked, err := testStore(db, defaultTestStoreOptions(nil)).MarkFailed(context.Background(), 1, "new message")
+	marked, err := testStore(db, defaultTestStoreOptions(nil)).MarkFailed(context.Background(), 1, "new message", MarkFinalOptions{})
 	if err != nil {
 		t.Fatalf("unexpected error marking upload as completed: %s", err)
 	}
@@ -592,12 +608,12 @@ func TestStoreMarkErroredAlreadyCompleted(t *testing.T) {
 		t.Fatalf("unexpected error inserting records: %s", err)
 	}
 
-	marked, err := testStore(db, defaultTestStoreOptions(nil)).MarkErrored(context.Background(), 1, "new message")
+	marked, err := testStore(db, defaultTestStoreOptions(nil)).MarkErrored(context.Background(), 1, "new message", MarkFinalOptions{})
 	if err != nil {
 		t.Fatalf("unexpected error marking record as errored: %s", err)
 	}
-	if !marked {
-		t.Fatalf("expected record to be marked")
+	if marked {
+		t.Fatalf("expected record not to be marked errired")
 	}
 
 	rows, err := db.QueryContext(context.Background(), `SELECT state, failure_message FROM workerutil_test WHERE id = 1`)
@@ -615,11 +631,11 @@ func TestStoreMarkErroredAlreadyCompleted(t *testing.T) {
 	if err := rows.Scan(&state, &failureMessage); err != nil {
 		t.Fatalf("unexpected error scanning record: %s", err)
 	}
-	if state != "errored" {
-		t.Errorf("unexpected state. want=%q have=%q", "errored", state)
+	if state != "completed" {
+		t.Errorf("unexpected state. want=%q have=%q", "completed", state)
 	}
-	if failureMessage == nil || *failureMessage != "new message" {
-		t.Errorf("unexpected failure message. want=%v have=%v", "new message", failureMessage)
+	if failureMessage != nil {
+		t.Errorf("unexpected non-empty failure message")
 	}
 }
 
@@ -634,7 +650,7 @@ func TestStoreMarkErroredAlreadyErrored(t *testing.T) {
 		t.Fatalf("unexpected error inserting records: %s", err)
 	}
 
-	marked, err := testStore(db, defaultTestStoreOptions(nil)).MarkErrored(context.Background(), 1, "new message")
+	marked, err := testStore(db, defaultTestStoreOptions(nil)).MarkErrored(context.Background(), 1, "new message", MarkFinalOptions{})
 	if err != nil {
 		t.Fatalf("unexpected error marking record as errored: %s", err)
 	}
@@ -682,7 +698,7 @@ func TestStoreMarkErroredRetriesExhausted(t *testing.T) {
 	store := testStore(db, options)
 
 	for i := 1; i < 3; i++ {
-		marked, err := store.MarkErrored(context.Background(), i, "new message")
+		marked, err := store.MarkErrored(context.Background(), i, "new message", MarkFinalOptions{})
 		if err != nil {
 			t.Fatalf("unexpected error marking record as errored: %s", err)
 		}
@@ -814,7 +830,7 @@ func TestStoreHeartbeat(t *testing.T) {
 	clock := glock.NewMockClockAt(now)
 	store := testStore(db, defaultTestStoreOptions(clock))
 
-	if err := store.Heartbeat(context.Background(), 1); err != nil {
+	if _, err := store.Heartbeat(context.Background(), []int{1}, HeartbeatOptions{}); err != nil {
 		t.Fatalf("unexpected error updating heartbeat: %s", err)
 	}
 
@@ -849,7 +865,7 @@ func TestStoreHeartbeat(t *testing.T) {
 	`); err != nil {
 		t.Fatalf("unexpected error updating records: %s", err)
 	}
-	if err := store.Heartbeat(context.Background(), 1); err != nil {
+	if _, err := store.Heartbeat(context.Background(), []int{1}, HeartbeatOptions{}); err != nil {
 		t.Fatalf("unexpected error updating heartbeat: %s", err)
 	}
 

internal/workerutil/dbworker/store_shim.go

@@ -46,6 +46,26 @@ func (s *storeShim) Dequeue(ctx context.Context, workerHostname string, extraArg
 	return s.Store.Dequeue(ctx, workerHostname, conditions)
 }
 
+func (s *storeShim) Heartbeat(ctx context.Context, ids []int) (knownIDs []int, err error) {
+	return s.Store.Heartbeat(ctx, ids, store.HeartbeatOptions{})
+}
+
+func (s *storeShim) AddExecutionLogEntry(ctx context.Context, id int, entry workerutil.ExecutionLogEntry) error {
+	return s.Store.AddExecutionLogEntry(ctx, id, entry, store.AddExecutionLogEntryOptions{})
+}
+
+func (s *storeShim) MarkComplete(ctx context.Context, id int) (bool, error) {
+	return s.Store.MarkComplete(ctx, id, store.MarkFinalOptions{})
+}
+
+func (s *storeShim) MarkFailed(ctx context.Context, id int, failureMessage string) (bool, error) {
+	return s.Store.MarkFailed(ctx, id, failureMessage, store.MarkFinalOptions{})
+}
+
+func (s *storeShim) MarkErrored(ctx context.Context, id int, errorMessage string) (bool, error) {
+	return s.Store.MarkErrored(ctx, id, errorMessage, store.MarkFinalOptions{})
+}
+
 // ErrNotConditions occurs when a PreDequeue handler returns non-sql query extra arguments.
 var ErrNotConditions = errors.New("expected slice of *sqlf.Query values")
 

internal/workerutil/store.go

@@ -22,8 +22,9 @@ type Store interface {
 	// flag indicating the existence of a processable record.
 	Dequeue(ctx context.Context, workerHostname string, extraArguments interface{}) (Record, bool, error)
 
-	// Heartbeat marks the given record as currently being processed.:2
-	Heartbeat(ctx context.Context, id int) error
+	// Heartbeat updates last_heartbeat_at of all the given jobs, when they're processing. All IDs of records that were
+	// touched are returned.
+	Heartbeat(ctx context.Context, jobIDs []int) (knownIDs []int, err error)
 
 	// AddExecutionLogEntry adds an executor log entry to the record.
 	AddExecutionLogEntry(ctx context.Context, id int, entry ExecutionLogEntry) error

internal/workerutil/worker.go

@@ -183,7 +183,7 @@ func (w *Worker) dequeueAndHandle() (dequeued bool, err error) {
 			}
 
 			id := record.RecordID()
-			if err := w.store.Heartbeat(heartbeatCtx, id); err != nil {
+			if _, err := w.store.Heartbeat(heartbeatCtx, []int{id}); err != nil {
 				log15.Error("Failed to refresh last_heartbeat_at", "name", w.options.Name, "id", id, "error", err)
 			}
 		}

internal/workerutil/worker_test.go

@@ -355,9 +355,9 @@ func TestWorkerDequeueHeartbeat(t *testing.T) {
 	}
 
 	heartbeats := make(chan struct{})
-	store.HeartbeatFunc.SetDefaultHook(func(c context.Context, i int) error {
+	store.HeartbeatFunc.SetDefaultHook(func(c context.Context, i []int) ([]int, error) {
 		heartbeats <- struct{}{}
-		return nil
+		return i, nil
 	})
 
 	worker := newWorker(context.Background(), store, handler, options, clock)

migrations/frontend/1528395855_add-missing-execution-logs-cols.down.sql

@@ -0,0 +1,6 @@
+BEGIN;
+
+ALTER TABLE IF EXISTS cm_trigger_jobs DROP COLUMN IF EXISTS execution_logs;
+ALTER TABLE IF EXISTS cm_action_jobs  DROP COLUMN IF EXISTS execution_logs;
+
+COMMIT;

migrations/frontend/1528395855_add-missing-execution-logs-cols.up.sql

@@ -0,0 +1,6 @@
+BEGIN;
+
+ALTER TABLE IF EXISTS cm_trigger_jobs ADD COLUMN IF NOT EXISTS execution_logs JSON[];
+ALTER TABLE IF EXISTS cm_action_jobs  ADD COLUMN IF NOT EXISTS execution_logs JSON[];
+
+COMMIT;