artin/sourcegraph
1
0
Fork 0
mirror of https://github.com/sourcegraph/sourcegraph.git synced 2026-02-06 17:31:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

Worker: use string ID instead of int for job tracking (#52454)

Browse Source 
- Prerequisite for
https://github.com/sourcegraph/sourcegraph/issues/51658, see [thread for
context](https://sourcegraph.slack.com/archives/C04K1NSCK6K/p1684841039315679)

This PR lets the worker keep track of it's running job IDs using a
string ID rather than int IDs. To achieve this, the `Record` interface
is extended with a second method `RecordUID()` which returns a string.

All the types that implement `Record` add a method `RecordUID()` which
simply converts its int value to a string representation. The only
exception to this rule is `executortypes.Job`, which includes the queue
name if it is set. The end goal here is that an executor can process
jobs from multiple queues by retrieving the required context from the
heartbeat IDs without leaking implementation details to the worker
level.

## Test plan
Unit tests.
Looking for input on other ways to validate nothing is broken.
<!-- All pull requests REQUIRE a test plan:
https://docs.sourcegraph.com/dev/background-information/testing_principles
-->
This commit is contained in:
Sander Ginn 2023-05-26 16:26:42 +02:00 committed by GitHub
parent 33cac9fae1
commit a1a2684610
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
45 changed files with 333 additions and 417 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
deps.bzl
View File

@ -6356,8 +6356,8 @@ def go_dependencies():
"//third_party/com_github_sourcegraph_zoekt:zoekt_webserver.patch",
"//third_party/com_github_sourcegraph_zoekt:zoekt_indexserver.patch",
],
sum = "h1:/5s1HW1DdlGpgr9PkOIgLcdFMHR1IHAJibXqT2Op5fk=",
version = "v0.0.0-20230523175034-5250e0e52a1b",
sum = "h1:BNbBGAoGT2+nSzad7cXmQ2OC6tpWZXBcNFlk7YIQljU=",
version = "v0.0.0-20230524133412-579a9a1eb19b",
)
go_repository(

42
enterprise/cmd/executor/internal/apiclient/queue/client.go
View File

@ -149,10 +149,10 @@ func (c *Client) MarkFailed(ctx context.Context, job types.Job, failureMessage s
return true, nil
}
func (c *Client) Heartbeat(ctx context.Context, jobIDs []int) (knownIDs, cancelIDs []int, err error) {
func (c *Client) Heartbeat(ctx context.Context, jobIDs []string) (knownIDs, cancelIDs []string, err error) {
ctx, _, endObservation := c.operations.heartbeat.With(ctx, &err, observation.Args{Attrs: []attribute.KeyValue{
attribute.String("queueName", c.options.QueueName),
attribute.IntSlice("jobIDs", jobIDs),
attribute.StringSlice("jobIDs", jobIDs),
}})
defer endObservation(1, observation.Args{})
@ -163,9 +163,6 @@ func (c *Client) Heartbeat(ctx context.Context, jobIDs []int) (knownIDs, cancelI
}
req, err := c.client.NewJSONRequest(http.MethodPost, fmt.Sprintf("%s/heartbeat", c.options.QueueName), types.HeartbeatRequest{
// Request the new-fashioned payload.
Version: types.ExecutorAPIVersion2,
ExecutorName: c.options.ExecutorName,
JobIDs: jobIDs,
@ -203,23 +200,7 @@ func (c *Client) Heartbeat(ctx context.Context, jobIDs []int) (knownIDs, cancelI
// If that works, we can return the data.
return respV2.KnownIDs, respV2.CancelIDs, nil
}
// If unmarshalling fails, try to parse it as a V1 payload.
var respV1 []int
if err := json.Unmarshal(bodyBytes, &respV1); err != nil {
return nil, nil, err
}
// If that works, we also have to fetch canceled jobs separately, as we
// are talking to a pre-4.3 Sourcegraph API and that doesn't return canceled
// jobs as part of heartbeats.
cancelIDs, err = c.CanceledJobs(ctx, jobIDs)
if err != nil {
return nil, nil, err
}
return respV1, cancelIDs, nil
return nil, nil, err
}
func gatherMetrics(logger log.Logger, gatherer prometheus.Gatherer) (string, error) {
@ -246,23 +227,6 @@ func gatherMetrics(logger log.Logger, gatherer prometheus.Gatherer) (string, err
return buf.String(), nil
}
// TODO: Remove this in Sourcegraph 4.4.
func (c *Client) CanceledJobs(ctx context.Context, knownIDs []int) (canceledIDs []int, err error) {
req, err := c.client.NewJSONRequest(http.MethodPost, fmt.Sprintf("%s/canceledJobs", c.options.QueueName), types.CanceledJobsRequest{
KnownJobIDs: knownIDs,
ExecutorName: c.options.ExecutorName,
})
if err != nil {
return nil, err
}
if _, err := c.client.DoAndDecode(ctx, req, &canceledIDs); err != nil {
return nil, err
}
return canceledIDs, nil
}
func (c *Client) Ping(ctx context.Context) (err error) {
req, err := c.client.NewJSONRequest(http.MethodPost, fmt.Sprintf("%s/heartbeat", c.options.QueueName), types.HeartbeatRequest{
ExecutorName: c.options.ExecutorName,

36
enterprise/cmd/executor/internal/apiclient/queue/client_test.go
View File

@ -301,26 +301,6 @@ func TestClient_MarkFailed(t *testing.T) {
}
}
func TestCanceledJobs(t *testing.T) {
spec := routeSpec{
expectedMethod: "POST",
expectedPath: "/.executors/queue/test_queue/canceledJobs",
expectedUsername: "test",
expectedToken: "hunter2",
expectedPayload: `{"executorName": "deadbeef","knownJobIds":[1]}`,
responseStatus: http.StatusOK,
responsePayload: `[1]`,
}
testRoute(t, spec, func(client *queue.Client) {
if ids, err := client.CanceledJobs(context.Background(), []int{1}); err != nil {
t.Fatalf("unexpected error completing job: %s", err)
} else if diff := cmp.Diff(ids, []int{1}); diff != "" {
t.Fatalf("unexpected set of IDs returned: %s", diff)
}
})
}
func TestHeartbeat(t *testing.T) {
spec := routeSpec{
expectedMethod: "POST",
@ -329,8 +309,7 @@ func TestHeartbeat(t *testing.T) {
expectedToken: "hunter2",
expectedPayload: `{
"executorName": "deadbeef",
"jobIds": [1,2,3],
"version": "V2",
"jobIds": ["1","2","3"],
"os": "test-os",
"architecture": "test-architecture",
@ -343,20 +322,20 @@ func TestHeartbeat(t *testing.T) {
"prometheusMetrics": ""
}`,
responseStatus: http.StatusOK,
responsePayload: `{"knownIDs": [1], "cancelIDs": [1]}`,
responsePayload: `{"knownIDs": ["1"], "cancelIDs": ["1"]}`,
}
testRoute(t, spec, func(client *queue.Client) {
unknownIDs, cancelIDs, err := client.Heartbeat(context.Background(), []int{1, 2, 3})
unknownIDs, cancelIDs, err := client.Heartbeat(context.Background(), []string{"1", "2", "3"})
if err != nil {
t.Fatalf("unexpected error performing heartbeat: %s", err)
}
if diff := cmp.Diff([]int{1}, unknownIDs); diff != "" {
if diff := cmp.Diff([]string{"1"}, unknownIDs); diff != "" {
t.Errorf("unexpected unknown ids (-want +got):\n%s", diff)
}
if diff := cmp.Diff([]int{1}, cancelIDs); diff != "" {
if diff := cmp.Diff([]string{"1"}, cancelIDs); diff != "" {
t.Errorf("unexpected unknown cancel ids (-want +got):\n%s", diff)
}
})
@ -370,8 +349,7 @@ func TestHeartbeatBadResponse(t *testing.T) {
expectedToken: "hunter2",
expectedPayload: `{
"executorName": "deadbeef",
"jobIds": [1,2,3],
"version": "V2",
"jobIds": ["1","2","3"],
"os": "test-os",
"architecture": "test-architecture",
@ -388,7 +366,7 @@ func TestHeartbeatBadResponse(t *testing.T) {
}
testRoute(t, spec, func(client *queue.Client) {
if _, _, err := client.Heartbeat(context.Background(), []int{1, 2, 3}); err == nil {
if _, _, err := client.Heartbeat(context.Background(), []string{"1", "2", "3"}); err == nil {
t.Fatalf("expected an error")
}
})

1
enterprise/cmd/frontend/internal/executorqueue/handler/BUILD.bazel generated
View File

@ -48,7 +48,6 @@ go_test(
"//internal/database",
"//internal/executor",
"//internal/metrics/store",
"//internal/types",
"//internal/workerutil",
"//internal/workerutil/dbworker/store",
"//internal/workerutil/dbworker/store/mocks",

38
enterprise/cmd/frontend/internal/executorqueue/handler/handler.go
View File

@ -44,8 +44,6 @@ type ExecutorHandler interface {
HandleMarkFailed(w http.ResponseWriter, r *http.Request)
// HandleHeartbeat handles the heartbeat of an executor.
HandleHeartbeat(w http.ResponseWriter, r *http.Request)
// HandleCanceledJobs cancels the specified executor.Jobs.
HandleCanceledJobs(w http.ResponseWriter, r *http.Request)
}
var _ ExecutorHandler = &handler[workerutil.Record]{}
@ -377,16 +375,11 @@ func (h *handler[T]) HandleHeartbeat(w http.ResponseWriter, r *http.Request) {
knownIDs, cancelIDs, err := h.heartbeat(r.Context(), e, payload.JobIDs)
if payload.Version == executortypes.ExecutorAPIVersion2 {
return http.StatusOK, executortypes.HeartbeatResponse{KnownIDs: knownIDs, CancelIDs: cancelIDs}, err
}
// TODO: Remove in Sourcegraph 4.4.
return http.StatusOK, knownIDs, err
return http.StatusOK, executortypes.HeartbeatResponse{KnownIDs: knownIDs, CancelIDs: cancelIDs}, err
})
}
func (h *handler[T]) heartbeat(ctx context.Context, executor types.Executor, ids []int) ([]int, []int, error) {
func (h *handler[T]) heartbeat(ctx context.Context, executor types.Executor, ids []string) ([]string, []string, error) {
if err := validateWorkerHostname(executor.Hostname); err != nil {
return nil, nil, err
}
@ -407,16 +400,6 @@ func (h *handler[T]) heartbeat(ctx context.Context, executor types.Executor, ids
return knownIDs, cancelIDs, errors.Wrap(err, "dbworkerstore.UpsertHeartbeat")
}
// TODO: This handler can be removed in Sourcegraph 4.4.
func (h *handler[T]) HandleCanceledJobs(w http.ResponseWriter, r *http.Request) {
var payload executortypes.CanceledJobsRequest
wrapHandler(w, r, &payload, h.logger, func() (int, any, error) {
canceledIDs, err := h.cancelJobs(r.Context(), payload.ExecutorName, payload.KnownJobIDs)
return http.StatusOK, canceledIDs, err
})
}
// wrapHandler decodes the request body into the given payload pointer, then calls the given
// handler function. If the body cannot be decoded, a 400 BadRequest is returned and the handler
// function is not called. If the handler function returns an error, a 500 Internal Server Error
@ -505,23 +488,6 @@ type errorResponse struct {
Error string `json:"error"`
}
// cancelJobs reaches to the queueHandlers.FetchCanceled to determine jobs that need to be canceled.
// This endpoint is deprecated and should be removed in Sourcegraph 4.4.
func (h *handler[T]) cancelJobs(ctx context.Context, executorName string, knownIDs []int) ([]int, error) {
if err := validateWorkerHostname(executorName); err != nil {
return nil, err
}
// The Heartbeat method now handles both heartbeats and cancellation. For backcompat,
// we fall back to this method.
_, canceledIDs, err := h.queueHandler.Store.Heartbeat(ctx, knownIDs, store.HeartbeatOptions{
// 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,
})
return canceledIDs, errors.Wrap(err, "dbworkerstore.CanceledJobs")
}
// validateWorkerHostname validates the WorkerHostname field sent for all the endpoints.
// We don't allow empty hostnames, as it would bypass the hostname verification, which
// could lead to stray workers updating records they no longer own.

163
enterprise/cmd/frontend/internal/executorqueue/handler/handler_test.go
View File

@ -7,6 +7,7 @@ import (
"io"
"net/http"
"net/http/httptest"
"strconv"
"strings"
"testing"
"time"
@ -23,7 +24,6 @@ import (
"github.com/sourcegraph/sourcegraph/internal/database"
internalexecutor "github.com/sourcegraph/sourcegraph/internal/executor"
metricsstore "github.com/sourcegraph/sourcegraph/internal/metrics/store"
"github.com/sourcegraph/sourcegraph/internal/types"
dbworkerstore "github.com/sourcegraph/sourcegraph/internal/workerutil/dbworker/store"
dbworkerstoremocks "github.com/sourcegraph/sourcegraph/internal/workerutil/dbworker/store/mocks"
"github.com/sourcegraph/sourcegraph/lib/errors"
@ -830,46 +830,15 @@ func TestHandler_HandleHeartbeat(t *testing.T) {
expectedResponseBody string
assertionFunc func(t *testing.T, metricsStore *metricsstore.MockDistributedStore, executorStore *database.MockExecutorStore, mockStore *dbworkerstoremocks.MockStore[testRecord])
}{
{
name: "Heartbeat",
body: `{"executorName": "test-executor", "jobIds": [42, 7], "os": "test-os", "architecture": "test-arch", "dockerVersion": "1.0", "executorVersion": "2.0", "gitVersion": "3.0", "igniteVersion": "4.0", "srcCliVersion": "5.0", "prometheusMetrics": ""}`,
mockFunc: func(metricsStore *metricsstore.MockDistributedStore, executorStore *database.MockExecutorStore, mockStore *dbworkerstoremocks.MockStore[testRecord]) {
executorStore.UpsertHeartbeatFunc.PushReturn(nil)
mockStore.HeartbeatFunc.PushReturn([]int{42, 7}, nil, nil)
},
expectedStatusCode: http.StatusOK,
expectedResponseBody: "[42,7]",
assertionFunc: func(t *testing.T, metricsStore *metricsstore.MockDistributedStore, executorStore *database.MockExecutorStore, mockStore *dbworkerstoremocks.MockStore[testRecord]) {
require.Len(t, executorStore.UpsertHeartbeatFunc.History(), 1)
assert.Equal(
t,
types.Executor{
Hostname: "test-executor",
QueueName: "test",
OS: "test-os",
Architecture: "test-arch",
DockerVersion: "1.0",
ExecutorVersion: "2.0",
GitVersion: "3.0",
IgniteVersion: "4.0",
SrcCliVersion: "5.0",
},
executorStore.UpsertHeartbeatFunc.History()[0].Arg1,
)
require.Len(t, mockStore.HeartbeatFunc.History(), 1)
assert.Equal(t, []int{42, 7}, mockStore.HeartbeatFunc.History()[0].Arg1)
assert.Equal(t, dbworkerstore.HeartbeatOptions{WorkerHostname: "test-executor"}, mockStore.HeartbeatFunc.History()[0].Arg2)
},
},
{
name: "V2 Heartbeat",
body: `{"version":"V2", "executorName": "test-executor", "jobIds": [42, 7], "os": "test-os", "architecture": "test-arch", "dockerVersion": "1.0", "executorVersion": "2.0", "gitVersion": "3.0", "igniteVersion": "4.0", "srcCliVersion": "5.0", "prometheusMetrics": ""}`,
body: `{"version":"V2", "executorName": "test-executor", "jobIds": ["42", "7"], "os": "test-os", "architecture": "test-arch", "dockerVersion": "1.0", "executorVersion": "2.0", "gitVersion": "3.0", "igniteVersion": "4.0", "srcCliVersion": "5.0", "prometheusMetrics": ""}`,
mockFunc: func(metricsStore *metricsstore.MockDistributedStore, executorStore *database.MockExecutorStore, mockStore *dbworkerstoremocks.MockStore[testRecord]) {
executorStore.UpsertHeartbeatFunc.PushReturn(nil)
mockStore.HeartbeatFunc.PushReturn([]int{42, 7}, nil, nil)
mockStore.HeartbeatFunc.PushReturn([]string{"42", "7"}, nil, nil)
},
expectedStatusCode: http.StatusOK,
expectedResponseBody: `{"knownIds":[42,7],"cancelIds":null}`,
expectedResponseBody: `{"knownIds":["42","7"],"cancelIds":null}`,
assertionFunc: func(t *testing.T, metricsStore *metricsstore.MockDistributedStore, executorStore *database.MockExecutorStore, mockStore *dbworkerstoremocks.MockStore[testRecord]) {
require.Len(t, executorStore.UpsertHeartbeatFunc.History(), 1)
require.Len(t, mockStore.HeartbeatFunc.History(), 1)
@ -877,7 +846,7 @@ func TestHandler_HandleHeartbeat(t *testing.T) {
},
{
name: "Invalid worker hostname",
body: `{"executorName": "", "jobIds": [42, 7], "os": "test-os", "architecture": "test-arch", "dockerVersion": "1.0", "executorVersion": "2.0", "gitVersion": "3.0", "igniteVersion": "4.0", "srcCliVersion": "5.0", "prometheusMetrics": ""}`,
body: `{"executorName": "", "jobIds": ["42", "7"], "os": "test-os", "architecture": "test-arch", "dockerVersion": "1.0", "executorVersion": "2.0", "gitVersion": "3.0", "igniteVersion": "4.0", "srcCliVersion": "5.0", "prometheusMetrics": ""}`,
expectedStatusCode: http.StatusInternalServerError,
expectedResponseBody: `{"error":"worker hostname cannot be empty"}`,
assertionFunc: func(t *testing.T, metricsStore *metricsstore.MockDistributedStore, executorStore *database.MockExecutorStore, mockStore *dbworkerstoremocks.MockStore[testRecord]) {
@ -887,13 +856,13 @@ func TestHandler_HandleHeartbeat(t *testing.T) {
},
{
name: "Failed to upsert heartbeat",
body: `{"executorName": "test-executor", "jobIds": [42, 7], "os": "test-os", "architecture": "test-arch", "dockerVersion": "1.0", "executorVersion": "2.0", "gitVersion": "3.0", "igniteVersion": "4.0", "srcCliVersion": "5.0", "prometheusMetrics": ""}`,
body: `{"executorName": "test-executor", "jobIds": ["42", "7"], "os": "test-os", "architecture": "test-arch", "dockerVersion": "1.0", "executorVersion": "2.0", "gitVersion": "3.0", "igniteVersion": "4.0", "srcCliVersion": "5.0", "prometheusMetrics": ""}`,
mockFunc: func(metricsStore *metricsstore.MockDistributedStore, executorStore *database.MockExecutorStore, mockStore *dbworkerstoremocks.MockStore[testRecord]) {
executorStore.UpsertHeartbeatFunc.PushReturn(errors.New("failed"))
mockStore.HeartbeatFunc.PushReturn([]int{42, 7}, nil, nil)
mockStore.HeartbeatFunc.PushReturn([]string{"42", "7"}, nil, nil)
},
expectedStatusCode: http.StatusOK,
expectedResponseBody: `[42,7]`,
expectedResponseBody: `{"knownIds":["42","7"],"cancelIds":null}`,
assertionFunc: func(t *testing.T, metricsStore *metricsstore.MockDistributedStore, executorStore *database.MockExecutorStore, mockStore *dbworkerstoremocks.MockStore[testRecord]) {
require.Len(t, executorStore.UpsertHeartbeatFunc.History(), 1)
require.Len(t, mockStore.HeartbeatFunc.History(), 1)
@ -901,7 +870,7 @@ func TestHandler_HandleHeartbeat(t *testing.T) {
},
{
name: "Failed to heartbeat",
body: `{"executorName": "test-executor", "jobIds": [42, 7], "os": "test-os", "architecture": "test-arch", "dockerVersion": "1.0", "executorVersion": "2.0", "gitVersion": "3.0", "igniteVersion": "4.0", "srcCliVersion": "5.0", "prometheusMetrics": ""}`,
body: `{"executorName": "test-executor", "jobIds": ["42", "7"], "os": "test-os", "architecture": "test-arch", "dockerVersion": "1.0", "executorVersion": "2.0", "gitVersion": "3.0", "igniteVersion": "4.0", "srcCliVersion": "5.0", "prometheusMetrics": ""}`,
mockFunc: func(metricsStore *metricsstore.MockDistributedStore, executorStore *database.MockExecutorStore, mockStore *dbworkerstoremocks.MockStore[testRecord]) {
executorStore.UpsertHeartbeatFunc.PushReturn(nil)
mockStore.HeartbeatFunc.PushReturn(nil, nil, errors.New("failed"))
@ -913,29 +882,15 @@ func TestHandler_HandleHeartbeat(t *testing.T) {
require.Len(t, mockStore.HeartbeatFunc.History(), 1)
},
},
{
name: "Has cancelled ids",
body: `{"executorName": "test-executor", "jobIds": [42, 7], "os": "test-os", "architecture": "test-arch", "dockerVersion": "1.0", "executorVersion": "2.0", "gitVersion": "3.0", "igniteVersion": "4.0", "srcCliVersion": "5.0", "prometheusMetrics": ""}`,
mockFunc: func(metricsStore *metricsstore.MockDistributedStore, executorStore *database.MockExecutorStore, mockStore *dbworkerstoremocks.MockStore[testRecord]) {
executorStore.UpsertHeartbeatFunc.PushReturn(nil)
mockStore.HeartbeatFunc.PushReturn(nil, []int{42, 7}, nil)
},
expectedStatusCode: http.StatusOK,
expectedResponseBody: `null`,
assertionFunc: func(t *testing.T, metricsStore *metricsstore.MockDistributedStore, executorStore *database.MockExecutorStore, mockStore *dbworkerstoremocks.MockStore[testRecord]) {
require.Len(t, executorStore.UpsertHeartbeatFunc.History(), 1)
require.Len(t, mockStore.HeartbeatFunc.History(), 1)
},
},
{
name: "V2 has cancelled ids",
body: `{"version": "V2", "executorName": "test-executor", "jobIds": [42, 7], "os": "test-os", "architecture": "test-arch", "dockerVersion": "1.0", "executorVersion": "2.0", "gitVersion": "3.0", "igniteVersion": "4.0", "srcCliVersion": "5.0", "prometheusMetrics": ""}`,
body: `{"version": "V2", "executorName": "test-executor", "jobIds": ["42", "7"], "os": "test-os", "architecture": "test-arch", "dockerVersion": "1.0", "executorVersion": "2.0", "gitVersion": "3.0", "igniteVersion": "4.0", "srcCliVersion": "5.0", "prometheusMetrics": ""}`,
mockFunc: func(metricsStore *metricsstore.MockDistributedStore, executorStore *database.MockExecutorStore, mockStore *dbworkerstoremocks.MockStore[testRecord]) {
executorStore.UpsertHeartbeatFunc.PushReturn(nil)
mockStore.HeartbeatFunc.PushReturn(nil, []int{42, 7}, nil)
mockStore.HeartbeatFunc.PushReturn(nil, []string{"42", "7"}, nil)
},
expectedStatusCode: http.StatusOK,
expectedResponseBody: `{"knownIds":null,"cancelIds":[42,7]}`,
expectedResponseBody: `{"knownIds":null,"cancelIds":["42","7"]}`,
assertionFunc: func(t *testing.T, metricsStore *metricsstore.MockDistributedStore, executorStore *database.MockExecutorStore, mockStore *dbworkerstoremocks.MockStore[testRecord]) {
require.Len(t, executorStore.UpsertHeartbeatFunc.History(), 1)
require.Len(t, mockStore.HeartbeatFunc.History(), 1)
@ -999,102 +954,16 @@ func encodeMetrics(t *testing.T, data ...*dto.MetricFamily) string {
return buf.String()
}
func TestHandler_HandleCanceledJobs(t *testing.T) {
tests := []struct {
name string
body string
mockFunc func(mockStore *dbworkerstoremocks.MockStore[testRecord])
expectedStatusCode int
expectedResponseBody string
assertionFunc func(t *testing.T, mockStore *dbworkerstoremocks.MockStore[testRecord])
}{
{
name: "Cancel Jobs",
body: `{"knownJobIds": [42,7], "executorName": "test-executor"}`,
mockFunc: func(mockStore *dbworkerstoremocks.MockStore[testRecord]) {
mockStore.HeartbeatFunc.PushReturn(nil, []int{42, 7}, nil)
},
expectedStatusCode: http.StatusOK,
expectedResponseBody: "[42,7]",
assertionFunc: func(t *testing.T, mockStore *dbworkerstoremocks.MockStore[testRecord]) {
require.Len(t, mockStore.HeartbeatFunc.History(), 1)
assert.Equal(t, []int{42, 7}, mockStore.HeartbeatFunc.History()[0].Arg1)
assert.Equal(t, dbworkerstore.HeartbeatOptions{WorkerHostname: "test-executor"}, mockStore.HeartbeatFunc.History()[0].Arg2)
},
},
{
name: "Invalid worker hostname",
body: `{"knownJobIds": [42,7], "executorName": ""}`,
mockFunc: func(mockStore *dbworkerstoremocks.MockStore[testRecord]) {
},
expectedStatusCode: http.StatusInternalServerError,
expectedResponseBody: `{"error":"worker hostname cannot be empty"}`,
assertionFunc: func(t *testing.T, mockStore *dbworkerstoremocks.MockStore[testRecord]) {
require.Len(t, mockStore.HeartbeatFunc.History(), 0)
},
},
{
name: "Failed to cancel Jobs",
body: `{"knownJobIds": [42,7], "executorName": "test-executor"}`,
mockFunc: func(mockStore *dbworkerstoremocks.MockStore[testRecord]) {
mockStore.HeartbeatFunc.PushReturn(nil, nil, errors.New("failed"))
},
expectedStatusCode: http.StatusInternalServerError,
expectedResponseBody: `{"error":"dbworkerstore.CanceledJobs: failed"}`,
assertionFunc: func(t *testing.T, mockStore *dbworkerstoremocks.MockStore[testRecord]) {
require.Len(t, mockStore.HeartbeatFunc.History(), 1)
},
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
mockStore := dbworkerstoremocks.NewMockStore[testRecord]()
h := handler.NewHandler(
database.NewMockExecutorStore(),
executorstore.NewMockJobTokenStore(),
metricsstore.NewMockDistributedStore(),
handler.QueueHandler[testRecord]{Store: mockStore},
)
router := mux.NewRouter()
router.HandleFunc("/{queueName}", h.HandleCanceledJobs)
req, err := http.NewRequest(http.MethodPost, "/test", strings.NewReader(test.body))
require.NoError(t, err)
rw := httptest.NewRecorder()
if test.mockFunc != nil {
test.mockFunc(mockStore)
}
router.ServeHTTP(rw, req)
assert.Equal(t, test.expectedStatusCode, rw.Code)
b, err := io.ReadAll(rw.Body)
require.NoError(t, err)
if len(test.expectedResponseBody) > 0 {
assert.JSONEq(t, test.expectedResponseBody, string(b))
} else {
assert.Empty(t, string(b))
}
if test.assertionFunc != nil {
test.assertionFunc(t, mockStore)
}
})
}
}
type testRecord struct {
id int
}
func (r testRecord) RecordID() int { return r.id }
func (r testRecord) RecordUID() string {
return strconv.Itoa(r.id)
}
func newIntPtr(i int) *int {
return &i
}

1
enterprise/cmd/frontend/internal/executorqueue/handler/routes.go
View File

@ -14,7 +14,6 @@ func SetupRoutes(handler ExecutorHandler, router *mux.Router) {
subRouter := router.PathPrefix(fmt.Sprintf("/{queueName:(?:%s)}", regexp.QuoteMeta(handler.Name()))).Subrouter()
subRouter.Path("/dequeue").Methods(http.MethodPost).HandlerFunc(handler.HandleDequeue)
subRouter.Path("/heartbeat").Methods(http.MethodPost).HandlerFunc(handler.HandleHeartbeat)
subRouter.Path("/canceledJobs").Methods(http.MethodPost).HandlerFunc(handler.HandleCanceledJobs)
}
// SetupJobRoutes registers all route handlers required for all configured executor

9
enterprise/cmd/frontend/internal/executorqueue/handler/routes_test.go
View File

@ -39,15 +39,6 @@ func TestSetupRoutes(t *testing.T) {
h.On("HandleHeartbeat").Once()
},
},
{
name: "CanceledJobs",
method: http.MethodPost,
path: "/test/canceledJobs",
expectedStatusCode: http.StatusOK,
expectationsFunc: func(h *testExecutorHandler) {
h.On("HandleCanceledJobs").Once()
},
},
{
name: "Invalid root",
method: http.MethodPost,

5
enterprise/internal/batches/types/batch_spec_resolution_job.go
View File

@ -1,6 +1,7 @@
package types
import (
"strconv"
"strings"
"time"
@ -65,3 +66,7 @@ type BatchSpecResolutionJob struct {
func (j *BatchSpecResolutionJob) RecordID() int {
return int(j.ID)
}
func (j *BatchSpecResolutionJob) RecordUID() string {
return strconv.FormatInt(j.ID, 10)
}

5
enterprise/internal/batches/types/batch_spec_workspace_execution_job.go
View File

@ -1,6 +1,7 @@
package types
import (
"strconv"
"strings"
"time"
@ -74,3 +75,7 @@ type BatchSpecWorkspaceExecutionJob struct {
}
func (j *BatchSpecWorkspaceExecutionJob) RecordID() int { return int(j.ID) }
func (j *BatchSpecWorkspaceExecutionJob) RecordUID() string {
return strconv.FormatInt(j.ID, 10)
}

4
enterprise/internal/batches/types/changeset.go
View File

@ -311,6 +311,10 @@ type Changeset struct {
// RecordID is needed to implement the workerutil.Record interface.
func (c *Changeset) RecordID() int { return int(c.ID) }
func (c *Changeset) RecordUID() string {
return strconv.FormatInt(c.ID, 10)
}
// Clone returns a clone of a Changeset.
func (c *Changeset) Clone() *Changeset {
tt := *c

5
enterprise/internal/batches/types/changeset_job.go
View File

@ -1,6 +1,7 @@
package types
import (
"strconv"
"strings"
"time"
)
@ -96,3 +97,7 @@ type ChangesetJob struct {
func (j *ChangesetJob) RecordID() int {
return int(j.ID)
}
func (j *ChangesetJob) RecordUID() string {
return strconv.FormatInt(j.ID, 10)
}

30
enterprise/internal/codeintel/autoindexing/internal/background/dependencies/mocks_test.go generated
View File

@ -4559,7 +4559,7 @@ func NewMockWorkerStore[T workerutil.Record]() *MockWorkerStore[T] {
},
},
HeartbeatFunc: &WorkerStoreHeartbeatFunc[T]{
defaultHook: func(context.Context, []int, store1.HeartbeatOptions) (r0 []int, r1 []int, r2 error) {
defaultHook: func(context.Context, []string, store1.HeartbeatOptions) (r0 []string, r1 []string, r2 error) {
return
},
},
@ -4631,7 +4631,7 @@ func NewStrictMockWorkerStore[T workerutil.Record]() *MockWorkerStore[T] {
},
},
HeartbeatFunc: &WorkerStoreHeartbeatFunc[T]{
defaultHook: func(context.Context, []int, store1.HeartbeatOptions) ([]int, []int, error) {
defaultHook: func(context.Context, []string, store1.HeartbeatOptions) ([]string, []string, error) {
panic("unexpected invocation of MockWorkerStore.Heartbeat")
},
},
@ -5062,15 +5062,15 @@ func (c WorkerStoreHandleFuncCall[T]) Results() []interface{} {
// WorkerStoreHeartbeatFunc describes the behavior when the Heartbeat method
// of the parent MockWorkerStore instance is invoked.
type WorkerStoreHeartbeatFunc[T workerutil.Record] struct {
defaultHook func(context.Context, []int, store1.HeartbeatOptions) ([]int, []int, error)
hooks []func(context.Context, []int, store1.HeartbeatOptions) ([]int, []int, error)
defaultHook func(context.Context, []string, store1.HeartbeatOptions) ([]string, []string, error)
hooks []func(context.Context, []string, store1.HeartbeatOptions) ([]string, []string, error)
history []WorkerStoreHeartbeatFuncCall[T]
mutex sync.Mutex
}
// Heartbeat delegates to the next hook function in the queue and stores the
// parameter and result values of this invocation.
func (m *MockWorkerStore[T]) Heartbeat(v0 context.Context, v1 []int, v2 store1.HeartbeatOptions) ([]int, []int, error) {
func (m *MockWorkerStore[T]) Heartbeat(v0 context.Context, v1 []string, v2 store1.HeartbeatOptions) ([]string, []string, error) {
r0, r1, r2 := m.HeartbeatFunc.nextHook()(v0, v1, v2)
m.HeartbeatFunc.appendCall(WorkerStoreHeartbeatFuncCall[T]{v0, v1, v2, r0, r1, r2})
return r0, r1, r2
@ -5079,7 +5079,7 @@ func (m *MockWorkerStore[T]) Heartbeat(v0 context.Context, v1 []int, v2 store1.H
// SetDefaultHook sets function that is called when the Heartbeat method of
// the parent MockWorkerStore instance is invoked and the hook queue is
// empty.
func (f *WorkerStoreHeartbeatFunc[T]) SetDefaultHook(hook func(context.Context, []int, store1.HeartbeatOptions) ([]int, []int, error)) {
func (f *WorkerStoreHeartbeatFunc[T]) SetDefaultHook(hook func(context.Context, []string, store1.HeartbeatOptions) ([]string, []string, error)) {
f.defaultHook = hook
}
@ -5087,7 +5087,7 @@ func (f *WorkerStoreHeartbeatFunc[T]) SetDefaultHook(hook func(context.Context,
// Heartbeat method of the parent MockWorkerStore instance invokes the hook
// at the front of the queue and discards it. After the queue is empty, the
// default hook function is invoked for any future action.
func (f *WorkerStoreHeartbeatFunc[T]) PushHook(hook func(context.Context, []int, store1.HeartbeatOptions) ([]int, []int, error)) {
func (f *WorkerStoreHeartbeatFunc[T]) PushHook(hook func(context.Context, []string, store1.HeartbeatOptions) ([]string, []string, error)) {
f.mutex.Lock()
f.hooks = append(f.hooks, hook)
f.mutex.Unlock()
@ -5095,20 +5095,20 @@ func (f *WorkerStoreHeartbeatFunc[T]) PushHook(hook func(context.Context, []int,
// SetDefaultReturn calls SetDefaultHook with a function that returns the
// given values.
func (f *WorkerStoreHeartbeatFunc[T]) SetDefaultReturn(r0 []int, r1 []int, r2 error) {
f.SetDefaultHook(func(context.Context, []int, store1.HeartbeatOptions) ([]int, []int, error) {
func (f *WorkerStoreHeartbeatFunc[T]) SetDefaultReturn(r0 []string, r1 []string, r2 error) {
f.SetDefaultHook(func(context.Context, []string, store1.HeartbeatOptions) ([]string, []string, error) {
return r0, r1, r2
})
}
// PushReturn calls PushHook with a function that returns the given values.
func (f *WorkerStoreHeartbeatFunc[T]) PushReturn(r0 []int, r1 []int, r2 error) {
f.PushHook(func(context.Context, []int, store1.HeartbeatOptions) ([]int, []int, error) {
func (f *WorkerStoreHeartbeatFunc[T]) PushReturn(r0 []string, r1 []string, r2 error) {
f.PushHook(func(context.Context, []string, store1.HeartbeatOptions) ([]string, []string, error) {
return r0, r1, r2
})
}
func (f *WorkerStoreHeartbeatFunc[T]) nextHook() func(context.Context, []int, store1.HeartbeatOptions) ([]int, []int, error) {
func (f *WorkerStoreHeartbeatFunc[T]) nextHook() func(context.Context, []string, store1.HeartbeatOptions) ([]string, []string, error) {
f.mutex.Lock()
defer f.mutex.Unlock()
@ -5146,16 +5146,16 @@ type WorkerStoreHeartbeatFuncCall[T workerutil.Record] struct {
Arg0 context.Context
// Arg1 is the value of the 2nd argument passed to this method
// invocation.
Arg1 []int
Arg1 []string
// Arg2 is the value of the 3rd argument passed to this method
// invocation.
Arg2 store1.HeartbeatOptions
// Result0 is the value of the 1st result returned from this method
// invocation.
Result0 []int
Result0 []string
// Result1 is the value of the 2nd result returned from this method
// invocation.
Result1 []int
Result1 []string
// Result2 is the value of the 3rd result returned from this method
// invocation.
Result2 error

13
enterprise/internal/codeintel/autoindexing/internal/background/dependencies/workerutil.go
View File

@ -1,6 +1,9 @@
package dependencies
import "time"
import (
"strconv"
"time"
)
// dependencyIndexingJob is a subset of the lsif_dependency_indexing_jobs table and acts as the
// queue and execution record for indexing the dependencies of a particular completed upload.
@ -22,6 +25,10 @@ func (u dependencyIndexingJob) RecordID() int {
return u.ID
}
func (u dependencyIndexingJob) RecordUID() string {
return strconv.Itoa(u.ID)
}
// dependencySyncingJob is a subset of the lsif_dependency_syncing_jobs table and acts as the
// queue and execution record for indexing the dependencies of a particular completed upload.
type dependencySyncingJob struct {
@ -39,3 +46,7 @@ type dependencySyncingJob struct {
func (u dependencySyncingJob) RecordID() int {
return u.ID
}
func (u dependencySyncingJob) RecordUID() string {
return strconv.Itoa(u.ID)
}

9
enterprise/internal/codeintel/sentinel/shared/types.go
View File

@ -1,6 +1,9 @@
package shared
import "time"
import (
"strconv"
"time"
)
type Vulnerability struct {
ID int // internal ID
@ -26,6 +29,10 @@ func (v Vulnerability) RecordID() int {
return v.ID
}
func (v Vulnerability) RecordUID() string {
return strconv.Itoa(v.ID)
}
// Data that varies across instances of a vulnerability
// Need to decide if this will be flat inside Vulnerability (and have multiple duplicate vulns)
// or a separate struct/table

30
enterprise/internal/codeintel/uploads/internal/background/expirer/mocks_test.go generated
View File

@ -10093,7 +10093,7 @@ func NewMockWorkerStore[T workerutil.Record]() *MockWorkerStore[T] {
},
},
HeartbeatFunc: &WorkerStoreHeartbeatFunc[T]{
defaultHook: func(context.Context, []int, store1.HeartbeatOptions) (r0 []int, r1 []int, r2 error) {
defaultHook: func(context.Context, []string, store1.HeartbeatOptions) (r0 []string, r1 []string, r2 error) {
return
},
},
@ -10165,7 +10165,7 @@ func NewStrictMockWorkerStore[T workerutil.Record]() *MockWorkerStore[T] {
},
},
HeartbeatFunc: &WorkerStoreHeartbeatFunc[T]{
defaultHook: func(context.Context, []int, store1.HeartbeatOptions) ([]int, []int, error) {
defaultHook: func(context.Context, []string, store1.HeartbeatOptions) ([]string, []string, error) {
panic("unexpected invocation of MockWorkerStore.Heartbeat")
},
},
@ -10596,15 +10596,15 @@ func (c WorkerStoreHandleFuncCall[T]) Results() []interface{} {
// WorkerStoreHeartbeatFunc describes the behavior when the Heartbeat method
// of the parent MockWorkerStore instance is invoked.
type WorkerStoreHeartbeatFunc[T workerutil.Record] struct {
defaultHook func(context.Context, []int, store1.HeartbeatOptions) ([]int, []int, error)
hooks []func(context.Context, []int, store1.HeartbeatOptions) ([]int, []int, error)
defaultHook func(context.Context, []string, store1.HeartbeatOptions) ([]string, []string, error)
hooks []func(context.Context, []string, store1.HeartbeatOptions) ([]string, []string, error)
history []WorkerStoreHeartbeatFuncCall[T]
mutex sync.Mutex
}
// Heartbeat delegates to the next hook function in the queue and stores the
// parameter and result values of this invocation.
func (m *MockWorkerStore[T]) Heartbeat(v0 context.Context, v1 []int, v2 store1.HeartbeatOptions) ([]int, []int, error) {
func (m *MockWorkerStore[T]) Heartbeat(v0 context.Context, v1 []string, v2 store1.HeartbeatOptions) ([]string, []string, error) {
r0, r1, r2 := m.HeartbeatFunc.nextHook()(v0, v1, v2)
m.HeartbeatFunc.appendCall(WorkerStoreHeartbeatFuncCall[T]{v0, v1, v2, r0, r1, r2})
return r0, r1, r2
@ -10613,7 +10613,7 @@ func (m *MockWorkerStore[T]) Heartbeat(v0 context.Context, v1 []int, v2 store1.H
// SetDefaultHook sets function that is called when the Heartbeat method of
// the parent MockWorkerStore instance is invoked and the hook queue is
// empty.
func (f *WorkerStoreHeartbeatFunc[T]) SetDefaultHook(hook func(context.Context, []int, store1.HeartbeatOptions) ([]int, []int, error)) {
func (f *WorkerStoreHeartbeatFunc[T]) SetDefaultHook(hook func(context.Context, []string, store1.HeartbeatOptions) ([]string, []string, error)) {
f.defaultHook = hook
}
@ -10621,7 +10621,7 @@ func (f *WorkerStoreHeartbeatFunc[T]) SetDefaultHook(hook func(context.Context,
// Heartbeat method of the parent MockWorkerStore instance invokes the hook
// at the front of the queue and discards it. After the queue is empty, the
// default hook function is invoked for any future action.
func (f *WorkerStoreHeartbeatFunc[T]) PushHook(hook func(context.Context, []int, store1.HeartbeatOptions) ([]int, []int, error)) {
func (f *WorkerStoreHeartbeatFunc[T]) PushHook(hook func(context.Context, []string, store1.HeartbeatOptions) ([]string, []string, error)) {
f.mutex.Lock()
f.hooks = append(f.hooks, hook)
f.mutex.Unlock()
@ -10629,20 +10629,20 @@ func (f *WorkerStoreHeartbeatFunc[T]) PushHook(hook func(context.Context, []int,
// SetDefaultReturn calls SetDefaultHook with a function that returns the
// given values.
func (f *WorkerStoreHeartbeatFunc[T]) SetDefaultReturn(r0 []int, r1 []int, r2 error) {
f.SetDefaultHook(func(context.Context, []int, store1.HeartbeatOptions) ([]int, []int, error) {
func (f *WorkerStoreHeartbeatFunc[T]) SetDefaultReturn(r0 []string, r1 []string, r2 error) {
f.SetDefaultHook(func(context.Context, []string, store1.HeartbeatOptions) ([]string, []string, error) {
return r0, r1, r2
})
}
// PushReturn calls PushHook with a function that returns the given values.
func (f *WorkerStoreHeartbeatFunc[T]) PushReturn(r0 []int, r1 []int, r2 error) {
f.PushHook(func(context.Context, []int, store1.HeartbeatOptions) ([]int, []int, error) {
func (f *WorkerStoreHeartbeatFunc[T]) PushReturn(r0 []string, r1 []string, r2 error) {
f.PushHook(func(context.Context, []string, store1.HeartbeatOptions) ([]string, []string, error) {
return r0, r1, r2
})
}
func (f *WorkerStoreHeartbeatFunc[T]) nextHook() func(context.Context, []int, store1.HeartbeatOptions) ([]int, []int, error) {
func (f *WorkerStoreHeartbeatFunc[T]) nextHook() func(context.Context, []string, store1.HeartbeatOptions) ([]string, []string, error) {
f.mutex.Lock()
defer f.mutex.Unlock()
@ -10680,16 +10680,16 @@ type WorkerStoreHeartbeatFuncCall[T workerutil.Record] struct {
Arg0 context.Context
// Arg1 is the value of the 2nd argument passed to this method
// invocation.
Arg1 []int
Arg1 []string
// Arg2 is the value of the 3rd argument passed to this method
// invocation.
Arg2 store1.HeartbeatOptions
// Result0 is the value of the 1st result returned from this method
// invocation.
Result0 []int
Result0 []string
// Result1 is the value of the 2nd result returned from this method
// invocation.
Result1 []int
Result1 []string
// Result2 is the value of the 3rd result returned from this method
// invocation.
Result2 error

30
enterprise/internal/codeintel/uploads/internal/background/processor/mocks_test.go generated
View File

@ -9907,7 +9907,7 @@ func NewMockWorkerStore[T workerutil.Record]() *MockWorkerStore[T] {
},
},
HeartbeatFunc: &WorkerStoreHeartbeatFunc[T]{
defaultHook: func(context.Context, []int, store1.HeartbeatOptions) (r0 []int, r1 []int, r2 error) {
defaultHook: func(context.Context, []string, store1.HeartbeatOptions) (r0 []string, r1 []string, r2 error) {
return
},
},
@ -9979,7 +9979,7 @@ func NewStrictMockWorkerStore[T workerutil.Record]() *MockWorkerStore[T] {
},
},
HeartbeatFunc: &WorkerStoreHeartbeatFunc[T]{
defaultHook: func(context.Context, []int, store1.HeartbeatOptions) ([]int, []int, error) {
defaultHook: func(context.Context, []string, store1.HeartbeatOptions) ([]string, []string, error) {
panic("unexpected invocation of MockWorkerStore.Heartbeat")
},
},
@ -10410,15 +10410,15 @@ func (c WorkerStoreHandleFuncCall[T]) Results() []interface{} {
// WorkerStoreHeartbeatFunc describes the behavior when the Heartbeat method
// of the parent MockWorkerStore instance is invoked.
type WorkerStoreHeartbeatFunc[T workerutil.Record] struct {
defaultHook func(context.Context, []int, store1.HeartbeatOptions) ([]int, []int, error)
hooks []func(context.Context, []int, store1.HeartbeatOptions) ([]int, []int, error)
defaultHook func(context.Context, []string, store1.HeartbeatOptions) ([]string, []string, error)
hooks []func(context.Context, []string, store1.HeartbeatOptions) ([]string, []string, error)
history []WorkerStoreHeartbeatFuncCall[T]
mutex sync.Mutex
}
// Heartbeat delegates to the next hook function in the queue and stores the
// parameter and result values of this invocation.
func (m *MockWorkerStore[T]) Heartbeat(v0 context.Context, v1 []int, v2 store1.HeartbeatOptions) ([]int, []int, error) {
func (m *MockWorkerStore[T]) Heartbeat(v0 context.Context, v1 []string, v2 store1.HeartbeatOptions) ([]string, []string, error) {
r0, r1, r2 := m.HeartbeatFunc.nextHook()(v0, v1, v2)
m.HeartbeatFunc.appendCall(WorkerStoreHeartbeatFuncCall[T]{v0, v1, v2, r0, r1, r2})
return r0, r1, r2
@ -10427,7 +10427,7 @@ func (m *MockWorkerStore[T]) Heartbeat(v0 context.Context, v1 []int, v2 store1.H
// SetDefaultHook sets function that is called when the Heartbeat method of
// the parent MockWorkerStore instance is invoked and the hook queue is
// empty.
func (f *WorkerStoreHeartbeatFunc[T]) SetDefaultHook(hook func(context.Context, []int, store1.HeartbeatOptions) ([]int, []int, error)) {
func (f *WorkerStoreHeartbeatFunc[T]) SetDefaultHook(hook func(context.Context, []string, store1.HeartbeatOptions) ([]string, []string, error)) {
f.defaultHook = hook
}
@ -10435,7 +10435,7 @@ func (f *WorkerStoreHeartbeatFunc[T]) SetDefaultHook(hook func(context.Context,
// Heartbeat method of the parent MockWorkerStore instance invokes the hook
// at the front of the queue and discards it. After the queue is empty, the
// default hook function is invoked for any future action.
func (f *WorkerStoreHeartbeatFunc[T]) PushHook(hook func(context.Context, []int, store1.HeartbeatOptions) ([]int, []int, error)) {
func (f *WorkerStoreHeartbeatFunc[T]) PushHook(hook func(context.Context, []string, store1.HeartbeatOptions) ([]string, []string, error)) {
f.mutex.Lock()
f.hooks = append(f.hooks, hook)
f.mutex.Unlock()
@ -10443,20 +10443,20 @@ func (f *WorkerStoreHeartbeatFunc[T]) PushHook(hook func(context.Context, []int,
// SetDefaultReturn calls SetDefaultHook with a function that returns the
// given values.
func (f *WorkerStoreHeartbeatFunc[T]) SetDefaultReturn(r0 []int, r1 []int, r2 error) {
f.SetDefaultHook(func(context.Context, []int, store1.HeartbeatOptions) ([]int, []int, error) {
func (f *WorkerStoreHeartbeatFunc[T]) SetDefaultReturn(r0 []string, r1 []string, r2 error) {
f.SetDefaultHook(func(context.Context, []string, store1.HeartbeatOptions) ([]string, []string, error) {
return r0, r1, r2
})
}
// PushReturn calls PushHook with a function that returns the given values.
func (f *WorkerStoreHeartbeatFunc[T]) PushReturn(r0 []int, r1 []int, r2 error) {
f.PushHook(func(context.Context, []int, store1.HeartbeatOptions) ([]int, []int, error) {
func (f *WorkerStoreHeartbeatFunc[T]) PushReturn(r0 []string, r1 []string, r2 error) {
f.PushHook(func(context.Context, []string, store1.HeartbeatOptions) ([]string, []string, error) {
return r0, r1, r2
})
}
func (f *WorkerStoreHeartbeatFunc[T]) nextHook() func(context.Context, []int, store1.HeartbeatOptions) ([]int, []int, error) {
func (f *WorkerStoreHeartbeatFunc[T]) nextHook() func(context.Context, []string, store1.HeartbeatOptions) ([]string, []string, error) {
f.mutex.Lock()
defer f.mutex.Unlock()
@ -10494,16 +10494,16 @@ type WorkerStoreHeartbeatFuncCall[T workerutil.Record] struct {
Arg0 context.Context
// Arg1 is the value of the 2nd argument passed to this method
// invocation.
Arg1 []int
Arg1 []string
// Arg2 is the value of the 3rd argument passed to this method
// invocation.
Arg2 store1.HeartbeatOptions
// Result0 is the value of the 1st result returned from this method
// invocation.
Result0 []int
Result0 []string
// Result1 is the value of the 2nd result returned from this method
// invocation.
Result1 []int
Result1 []string
// Result2 is the value of the 3rd result returned from this method
// invocation.
Result2 error

30
enterprise/internal/codeintel/uploads/mocks_test.go generated
View File

@ -9793,7 +9793,7 @@ func NewMockWorkerStore[T workerutil.Record]() *MockWorkerStore[T] {
},
},
HeartbeatFunc: &WorkerStoreHeartbeatFunc[T]{
defaultHook: func(context.Context, []int, store1.HeartbeatOptions) (r0 []int, r1 []int, r2 error) {
defaultHook: func(context.Context, []string, store1.HeartbeatOptions) (r0 []string, r1 []string, r2 error) {
return
},
},
@ -9865,7 +9865,7 @@ func NewStrictMockWorkerStore[T workerutil.Record]() *MockWorkerStore[T] {
},
},
HeartbeatFunc: &WorkerStoreHeartbeatFunc[T]{
defaultHook: func(context.Context, []int, store1.HeartbeatOptions) ([]int, []int, error) {
defaultHook: func(context.Context, []string, store1.HeartbeatOptions) ([]string, []string, error) {
panic("unexpected invocation of MockWorkerStore.Heartbeat")
},
},
@ -10296,15 +10296,15 @@ func (c WorkerStoreHandleFuncCall[T]) Results() []interface{} {
// WorkerStoreHeartbeatFunc describes the behavior when the Heartbeat method
// of the parent MockWorkerStore instance is invoked.
type WorkerStoreHeartbeatFunc[T workerutil.Record] struct {
defaultHook func(context.Context, []int, store1.HeartbeatOptions) ([]int, []int, error)
hooks []func(context.Context, []int, store1.HeartbeatOptions) ([]int, []int, error)
defaultHook func(context.Context, []string, store1.HeartbeatOptions) ([]string, []string, error)
hooks []func(context.Context, []string, store1.HeartbeatOptions) ([]string, []string, error)
history []WorkerStoreHeartbeatFuncCall[T]
mutex sync.Mutex
}
// Heartbeat delegates to the next hook function in the queue and stores the
// parameter and result values of this invocation.
func (m *MockWorkerStore[T]) Heartbeat(v0 context.Context, v1 []int, v2 store1.HeartbeatOptions) ([]int, []int, error) {
func (m *MockWorkerStore[T]) Heartbeat(v0 context.Context, v1 []string, v2 store1.HeartbeatOptions) ([]string, []string, error) {
r0, r1, r2 := m.HeartbeatFunc.nextHook()(v0, v1, v2)
m.HeartbeatFunc.appendCall(WorkerStoreHeartbeatFuncCall[T]{v0, v1, v2, r0, r1, r2})
return r0, r1, r2
@ -10313,7 +10313,7 @@ func (m *MockWorkerStore[T]) Heartbeat(v0 context.Context, v1 []int, v2 store1.H
// SetDefaultHook sets function that is called when the Heartbeat method of
// the parent MockWorkerStore instance is invoked and the hook queue is
// empty.
func (f *WorkerStoreHeartbeatFunc[T]) SetDefaultHook(hook func(context.Context, []int, store1.HeartbeatOptions) ([]int, []int, error)) {
func (f *WorkerStoreHeartbeatFunc[T]) SetDefaultHook(hook func(context.Context, []string, store1.HeartbeatOptions) ([]string, []string, error)) {
f.defaultHook = hook
}
@ -10321,7 +10321,7 @@ func (f *WorkerStoreHeartbeatFunc[T]) SetDefaultHook(hook func(context.Context,
// Heartbeat method of the parent MockWorkerStore instance invokes the hook
// at the front of the queue and discards it. After the queue is empty, the
// default hook function is invoked for any future action.
func (f *WorkerStoreHeartbeatFunc[T]) PushHook(hook func(context.Context, []int, store1.HeartbeatOptions) ([]int, []int, error)) {
func (f *WorkerStoreHeartbeatFunc[T]) PushHook(hook func(context.Context, []string, store1.HeartbeatOptions) ([]string, []string, error)) {
f.mutex.Lock()
f.hooks = append(f.hooks, hook)
f.mutex.Unlock()
@ -10329,20 +10329,20 @@ func (f *WorkerStoreHeartbeatFunc[T]) PushHook(hook func(context.Context, []int,
// SetDefaultReturn calls SetDefaultHook with a function that returns the
// given values.
func (f *WorkerStoreHeartbeatFunc[T]) SetDefaultReturn(r0 []int, r1 []int, r2 error) {
f.SetDefaultHook(func(context.Context, []int, store1.HeartbeatOptions) ([]int, []int, error) {
func (f *WorkerStoreHeartbeatFunc[T]) SetDefaultReturn(r0 []string, r1 []string, r2 error) {
f.SetDefaultHook(func(context.Context, []string, store1.HeartbeatOptions) ([]string, []string, error) {
return r0, r1, r2
})
}
// PushReturn calls PushHook with a function that returns the given values.
func (f *WorkerStoreHeartbeatFunc[T]) PushReturn(r0 []int, r1 []int, r2 error) {
f.PushHook(func(context.Context, []int, store1.HeartbeatOptions) ([]int, []int, error) {
func (f *WorkerStoreHeartbeatFunc[T]) PushReturn(r0 []string, r1 []string, r2 error) {
f.PushHook(func(context.Context, []string, store1.HeartbeatOptions) ([]string, []string, error) {
return r0, r1, r2
})
}
func (f *WorkerStoreHeartbeatFunc[T]) nextHook() func(context.Context, []int, store1.HeartbeatOptions) ([]int, []int, error) {
func (f *WorkerStoreHeartbeatFunc[T]) nextHook() func(context.Context, []string, store1.HeartbeatOptions) ([]string, []string, error) {
f.mutex.Lock()
defer f.mutex.Unlock()
@ -10380,16 +10380,16 @@ type WorkerStoreHeartbeatFuncCall[T workerutil.Record] struct {
Arg0 context.Context
// Arg1 is the value of the 2nd argument passed to this method
// invocation.
Arg1 []int
Arg1 []string
// Arg2 is the value of the 3rd argument passed to this method
// invocation.
Arg2 store1.HeartbeatOptions
// Result0 is the value of the 1st result returned from this method
// invocation.
Result0 []int
Result0 []string
// Result1 is the value of the 2nd result returned from this method
// invocation.
Result1 []int
Result1 []string
// Result2 is the value of the 3rd result returned from this method
// invocation.
Result2 error

9
enterprise/internal/codeintel/uploads/shared/types.go
View File

@ -3,6 +3,7 @@ package shared
import (
"database/sql/driver"
"encoding/json"
"strconv"
"time"
"github.com/sourcegraph/sourcegraph/internal/executor"
@ -40,6 +41,10 @@ func (u Upload) RecordID() int {
return u.ID
}
func (u Upload) RecordUID() string {
return strconv.Itoa(u.ID)
}
// TODO - unify with Upload
// Dump is a subset of the lsif_uploads table (queried via the lsif_dumps_with_repository_name view)
// and stores only processed records.
@ -110,6 +115,10 @@ func (i Index) RecordID() int {
return i.ID
}
func (i Index) RecordUID() string {
return strconv.Itoa(i.ID)
}
type DockerStep struct {
Root string `json:"root"`
Image string `json:"image"`

5
enterprise/internal/database/code_monitor_action_jobs.go
View File

@ -4,6 +4,7 @@ import (
"context"
"database/sql"
"encoding/json"
"strconv"
"time"
"github.com/keegancsmith/sqlf"
@ -34,6 +35,10 @@ func (a *ActionJob) RecordID() int {
return int(a.ID)
}
func (a *ActionJob) RecordUID() string {
return strconv.FormatInt(int64(a.ID), 10)
}
type ActionJobMetadata struct {
Description string
MonitorID int64

5
enterprise/internal/database/code_monitor_trigger_jobs.go
View File

@ -4,6 +4,7 @@ import (
"context"
"database/sql"
"encoding/json"
"strconv"
"time"
"github.com/keegancsmith/sqlf"
@ -36,6 +37,10 @@ func (r *TriggerJob) RecordID() int {
return int(r.ID)
}
func (r *TriggerJob) RecordUID() string {
return strconv.FormatInt(int64(r.ID), 10)
}
const enqueueTriggerQueryFmtStr = `
WITH due AS (
SELECT cm_queries.id as id

5
enterprise/internal/embeddings/background/contextdetection/types.go
View File

@ -1,6 +1,7 @@
package contextdetection
import (
"strconv"
"time"
"github.com/sourcegraph/sourcegraph/internal/executor"
@ -25,3 +26,7 @@ type ContextDetectionEmbeddingJob struct {
func (j *ContextDetectionEmbeddingJob) RecordID() int {
return j.ID
}
func (j *ContextDetectionEmbeddingJob) RecordUID() string {
return strconv.Itoa(j.ID)
}

5
enterprise/internal/embeddings/background/repo/types.go
View File

@ -1,6 +1,7 @@
package repo
import (
"strconv"
"time"
"github.com/sourcegraph/sourcegraph/internal/api"
@ -30,6 +31,10 @@ func (j *RepoEmbeddingJob) RecordID() int {
return j.ID
}
func (j *RepoEmbeddingJob) RecordUID() string {
return strconv.Itoa(j.ID)
}
func (j *RepoEmbeddingJob) IsRepoEmbeddingJobScheduledOrCompleted() bool {
return j != nil && (j.State == "completed" || j.State == "processing" || j.State == "queued")
}

16
enterprise/internal/executor/types/http.go
View File

@ -39,14 +39,10 @@ type MarkErroredRequest struct {
}
type HeartbeatRequest struct {
// TODO: This field is set to become unneccesary in Sourcegraph 4.4.
Version ExecutorAPIVersion `json:"version"`
ExecutorName string `json:"executorName"`
JobIDs []int `json:"jobIds"`
ExecutorName string `json:"executorName"`
JobIDs []string `json:"jobIds"`
// Telemetry data.
OS string `json:"os"`
Architecture string `json:"architecture"`
DockerVersion string `json:"dockerVersion"`
@ -65,12 +61,12 @@ const (
)
type HeartbeatResponse struct {
KnownIDs []int `json:"knownIds"`
CancelIDs []int `json:"cancelIds"`
KnownIDs []string `json:"knownIds"`
CancelIDs []string `json:"cancelIds"`
}
// TODO: Deprecated. Can be removed in Sourcegraph 4.4.
type CanceledJobsRequest struct {
KnownJobIDs []int `json:"knownJobIds"`
ExecutorName string `json:"executorName"`
KnownJobIDs []string `json:"knownJobIds"`
ExecutorName string `json:"executorName"`
}

10
enterprise/internal/executor/types/job.go
View File

@ -2,6 +2,7 @@ package types
import (
"encoding/json"
"strconv"
"time"
)
@ -254,6 +255,15 @@ func (j Job) RecordID() int {
return j.ID
}
func (j Job) RecordUID() string {
uid := strconv.Itoa(j.ID)
// outside of multi-queue executors, jobs aren't guaranteed to have a queue specified
if j.Queue != "" {
uid += "-" + j.Queue
}
return uid
}
type DockerStep struct {
// Key is a unique identifier of the step. It can be used to retrieve the
// associated log entry.

5
enterprise/internal/insights/background/queryrunner/worker.go
View File

@ -3,6 +3,7 @@ package queryrunner
import (
"context"
"database/sql"
"strconv"
"time"
"github.com/keegancsmith/sqlf"
@ -463,6 +464,10 @@ func (j *Job) RecordID() int {
return j.ID
}
func (j *Job) RecordUID() string {
return strconv.Itoa(j.ID)
}
func scanJobs(rows *sql.Rows, err error) ([]*Job, error) {
if err != nil {
return nil, err

5
enterprise/internal/insights/background/retention/job.go
View File

@ -1,6 +1,7 @@
package retention
import (
"strconv"
"time"
"github.com/keegancsmith/sqlf"
@ -48,6 +49,10 @@ func (j *DataRetentionJob) RecordID() int {
return j.ID
}
func (j *DataRetentionJob) RecordUID() string {
return strconv.Itoa(j.ID)
}
func scanDataRetentionJob(s dbutil.Scanner) (*DataRetentionJob, error) {
var job DataRetentionJob
var executionLogs []executor.ExecutionLogEntry

5
enterprise/internal/insights/scheduler/scheduler.go
View File

@ -2,6 +2,7 @@ package scheduler
import (
"context"
"strconv"
"time"
"github.com/lib/pq"
@ -49,6 +50,10 @@ func (b *BaseJob) RecordID() int {
return b.ID
}
func (b *BaseJob) RecordUID() string {
return strconv.Itoa(b.ID)
}
var baseJobColumns = []*sqlf.Query{
sqlf.Sprintf("id"),
sqlf.Sprintf("state"),

6
enterprise/internal/own/background/background.go
View File

@ -3,6 +3,7 @@ package background
import (
"context"
"fmt"
"strconv"
"time"
"github.com/keegancsmith/sqlf"
@ -10,6 +11,7 @@ import (
"golang.org/x/time/rate"
"github.com/sourcegraph/log"
"github.com/sourcegraph/sourcegraph/enterprise/internal/own/types"
"github.com/sourcegraph/sourcegraph/enterprise/internal/codeintel/shared/background"
@ -59,6 +61,10 @@ func (b *Job) RecordID() int {
return b.ID
}
func (b *Job) RecordUID() string {
return strconv.Itoa(b.ID)
}
var jobColumns = []*sqlf.Query{
sqlf.Sprintf("id"),
sqlf.Sprintf("state"),

4
internal/database/permission_sync_jobs.go
View File

@ -777,6 +777,10 @@ type PermissionSyncJob struct {
func (j *PermissionSyncJob) RecordID() int { return j.ID }
func (j *PermissionSyncJob) RecordUID() string {
return strconv.Itoa(j.ID)
}
var PermissionSyncJobColumns = []*sqlf.Query{
sqlf.Sprintf("permission_sync_jobs.id"),
sqlf.Sprintf("permission_sync_jobs.state"),

5
internal/repos/sync_worker.go
View File

@ -3,6 +3,7 @@ package repos
import (
"context"
"database/sql"
"strconv"
"time"
"github.com/keegancsmith/sqlf"
@ -154,3 +155,7 @@ type SyncJob struct {
func (s *SyncJob) RecordID() int {
return s.ID
}
func (s *SyncJob) RecordUID() string {
return strconv.Itoa(s.ID)
}

5
internal/types/bitbucket_permissions.go
View File

@ -2,6 +2,7 @@ package types
import (
"database/sql/driver"
"strconv"
"time"
)
@ -41,6 +42,10 @@ func (g *BitbucketProjectPermissionJob) RecordID() int {
return g.ID
}
func (g *BitbucketProjectPermissionJob) RecordUID() string {
return strconv.Itoa(g.ID)
}
type UserPermission struct {
BindID string `json:"bindID"`
Permission string `json:"permission"`

5
internal/types/outbound_webhook_jobs.go
View File

@ -1,6 +1,7 @@
package types
import (
"strconv"
"time"
"github.com/sourcegraph/sourcegraph/internal/encryption"
@ -31,3 +32,7 @@ type OutboundWebhookJob struct {
func (j *OutboundWebhookJob) RecordID() int {
return int(j.ID)
}
func (j *OutboundWebhookJob) RecordUID() string {
return strconv.FormatInt(j.ID, 10)
}

5
internal/workerutil/dbworker/resetter_test.go
View File

@ -1,6 +1,7 @@
package dbworker
import (
"strconv"
"testing"
"time"
@ -21,6 +22,10 @@ func (v TestRecord) RecordID() int {
return v.ID
}
func (v TestRecord) RecordUID() string {
return strconv.Itoa(v.ID)
}
func TestResetter(t *testing.T) {
logger := logtest.Scoped(t)
s := storemocks.NewMockStore[*TestRecord]()

13
internal/workerutil/dbworker/store/helpers_test.go
View File

@ -2,6 +2,7 @@ package store
import (
"database/sql"
"strconv"
"testing"
"time"
@ -34,6 +35,10 @@ func (v TestRecord) RecordID() int {
return v.ID
}
func (v TestRecord) RecordUID() string {
return strconv.Itoa(v.ID)
}
func testScanRecord(sc dbutil.Scanner) (*TestRecord, error) {
var record TestRecord
return &record, sc.Scan(&record.ID, &record.State, pq.Array(&record.ExecutionLogs))
@ -49,6 +54,10 @@ func (v TestRecordView) RecordID() int {
return v.ID
}
func (v TestRecordView) RecordUID() string {
return strconv.Itoa(v.ID)
}
func testScanRecordView(sc dbutil.Scanner) (*TestRecordView, error) {
var record TestRecordView
return &record, sc.Scan(&record.ID, &record.State, &record.NewField)
@ -64,6 +73,10 @@ func (v TestRecordRetry) RecordID() int {
return v.ID
}
func (v TestRecordRetry) RecordUID() string {
return strconv.Itoa(v.ID)
}
func testScanRecordRetry(sc dbutil.Scanner) (*TestRecordRetry, error) {
var record TestRecordRetry
return &record, sc.Scan(&record.ID, &record.State, &record.NumResets)

30
internal/workerutil/dbworker/store/mocks/mocks_temp.go
View File

@ -84,7 +84,7 @@ func NewMockStore[T workerutil.Record]() *MockStore[T] {
},
},
HeartbeatFunc: &StoreHeartbeatFunc[T]{
defaultHook: func(context.Context, []int, store.HeartbeatOptions) (r0 []int, r1 []int, r2 error) {
defaultHook: func(context.Context, []string, store.HeartbeatOptions) (r0 []string, r1 []string, r2 error) {
return
},
},
@ -156,7 +156,7 @@ func NewStrictMockStore[T workerutil.Record]() *MockStore[T] {
},
},
HeartbeatFunc: &StoreHeartbeatFunc[T]{
defaultHook: func(context.Context, []int, store.HeartbeatOptions) ([]int, []int, error) {
defaultHook: func(context.Context, []string, store.HeartbeatOptions) ([]string, []string, error) {
panic("unexpected invocation of MockStore.Heartbeat")
},
},
@ -582,15 +582,15 @@ func (c StoreHandleFuncCall[T]) Results() []interface{} {
// StoreHeartbeatFunc describes the behavior when the Heartbeat method of
// the parent MockStore instance is invoked.
type StoreHeartbeatFunc[T workerutil.Record] struct {
defaultHook func(context.Context, []int, store.HeartbeatOptions) ([]int, []int, error)
hooks []func(context.Context, []int, store.HeartbeatOptions) ([]int, []int, error)
defaultHook func(context.Context, []string, store.HeartbeatOptions) ([]string, []string, error)
hooks []func(context.Context, []string, store.HeartbeatOptions) ([]string, []string, error)
history []StoreHeartbeatFuncCall[T]
mutex sync.Mutex
}
// Heartbeat delegates to the next hook function in the queue and stores the
// parameter and result values of this invocation.
func (m *MockStore[T]) Heartbeat(v0 context.Context, v1 []int, v2 store.HeartbeatOptions) ([]int, []int, error) {
func (m *MockStore[T]) Heartbeat(v0 context.Context, v1 []string, v2 store.HeartbeatOptions) ([]string, []string, error) {
r0, r1, r2 := m.HeartbeatFunc.nextHook()(v0, v1, v2)
m.HeartbeatFunc.appendCall(StoreHeartbeatFuncCall[T]{v0, v1, v2, r0, r1, r2})
return r0, r1, r2
@ -598,7 +598,7 @@ func (m *MockStore[T]) Heartbeat(v0 context.Context, v1 []int, v2 store.Heartbea
// SetDefaultHook sets function that is called when the Heartbeat method of
// the parent MockStore instance is invoked and the hook queue is empty.
func (f *StoreHeartbeatFunc[T]) SetDefaultHook(hook func(context.Context, []int, store.HeartbeatOptions) ([]int, []int, error)) {
func (f *StoreHeartbeatFunc[T]) SetDefaultHook(hook func(context.Context, []string, store.HeartbeatOptions) ([]string, []string, error)) {
f.defaultHook = hook
}
@ -606,7 +606,7 @@ func (f *StoreHeartbeatFunc[T]) SetDefaultHook(hook func(context.Context, []int,
// Heartbeat method of the parent MockStore instance invokes the hook at the
// front of the queue and discards it. After the queue is empty, the default
// hook function is invoked for any future action.
func (f *StoreHeartbeatFunc[T]) PushHook(hook func(context.Context, []int, store.HeartbeatOptions) ([]int, []int, error)) {
func (f *StoreHeartbeatFunc[T]) PushHook(hook func(context.Context, []string, store.HeartbeatOptions) ([]string, []string, error)) {
f.mutex.Lock()
f.hooks = append(f.hooks, hook)
f.mutex.Unlock()
@ -614,20 +614,20 @@ func (f *StoreHeartbeatFunc[T]) PushHook(hook func(context.Context, []int, store
// SetDefaultReturn calls SetDefaultHook with a function that returns the
// given values.
func (f *StoreHeartbeatFunc[T]) SetDefaultReturn(r0 []int, r1 []int, r2 error) {
f.SetDefaultHook(func(context.Context, []int, store.HeartbeatOptions) ([]int, []int, error) {
func (f *StoreHeartbeatFunc[T]) SetDefaultReturn(r0 []string, r1 []string, r2 error) {
f.SetDefaultHook(func(context.Context, []string, store.HeartbeatOptions) ([]string, []string, error) {
return r0, r1, r2
})
}
// PushReturn calls PushHook with a function that returns the given values.
func (f *StoreHeartbeatFunc[T]) PushReturn(r0 []int, r1 []int, r2 error) {
f.PushHook(func(context.Context, []int, store.HeartbeatOptions) ([]int, []int, error) {
func (f *StoreHeartbeatFunc[T]) PushReturn(r0 []string, r1 []string, r2 error) {
f.PushHook(func(context.Context, []string, store.HeartbeatOptions) ([]string, []string, error) {
return r0, r1, r2
})
}
func (f *StoreHeartbeatFunc[T]) nextHook() func(context.Context, []int, store.HeartbeatOptions) ([]int, []int, error) {
func (f *StoreHeartbeatFunc[T]) nextHook() func(context.Context, []string, store.HeartbeatOptions) ([]string, []string, error) {
f.mutex.Lock()
defer f.mutex.Unlock()
@ -665,16 +665,16 @@ type StoreHeartbeatFuncCall[T workerutil.Record] struct {
Arg0 context.Context
// Arg1 is the value of the 2nd argument passed to this method
// invocation.
Arg1 []int
Arg1 []string
// Arg2 is the value of the 3rd argument passed to this method
// invocation.
Arg2 store.HeartbeatOptions
// Result0 is the value of the 1st result returned from this method
// invocation.
Result0 []int
Result0 []string
// Result1 is the value of the 2nd result returned from this method
// invocation.
Result1 []int
Result1 []string
// Result2 is the value of the 3rd result returned from this method
// invocation.
Result2 error

30
internal/workerutil/dbworker/store/store.go
View File

@ -4,6 +4,7 @@ import (
"context"
"database/sql"
"fmt"
"strconv"
"strings"
"time"
@ -96,7 +97,7 @@ type Store[T workerutil.Record] interface {
// Heartbeat marks the given records as currently being processed and returns the list of records that are
// still known to the database (to detect lost jobs) and jobs that are marked as to be canceled.
Heartbeat(ctx context.Context, ids []int, options HeartbeatOptions) (knownIDs, cancelIDs []int, err error)
Heartbeat(ctx context.Context, ids []string, options HeartbeatOptions) (knownIDs, cancelIDs []string, 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.
@ -605,12 +606,12 @@ func (s *store[T]) makeDequeueUpdateStatements(updatedColumns map[string]*sqlf.Q
return updateStatements
}
func (s *store[T]) Heartbeat(ctx context.Context, ids []int, options HeartbeatOptions) (knownIDs, cancelIDs []int, err error) {
func (s *store[T]) Heartbeat(ctx context.Context, ids []string, options HeartbeatOptions) (knownIDs, cancelIDs []string, err error) {
ctx, _, endObservation := s.operations.heartbeat.With(ctx, &err, observation.Args{})
defer endObservation(1, observation.Args{})
if len(ids) == 0 {
return []int{}, []int{}, nil
return []string{}, []string{}, nil
}
quotedTableName := quote(s.options.TableName)
@ -621,7 +622,7 @@ func (s *store[T]) Heartbeat(ctx context.Context, ids []int, options HeartbeatOp
}
conds = append(conds, options.ToSQLConds(s.formatQuery)...)
scanner := basestore.NewMapScanner(func(scanner dbutil.Scanner) (id int, cancel bool, err error) {
scanner := basestore.NewMapScanner(func(scanner dbutil.Scanner) (id string, cancel bool, err error) {
err = scanner.Scan(&id, &cancel)
return
})
@ -646,15 +647,22 @@ func (s *store[T]) Heartbeat(ctx context.Context, ids []int, options HeartbeatOp
}
}
debug, debugErr := s.fetchDebugInformationForJob(ctx, recordID)
if debugErr != nil {
s.logger.Error("failed to fetch debug information for job",
log.Int("recordID", recordID),
log.Error(debugErr),
)
var debug string
intId, convErr := strconv.Atoi(recordID)
if convErr != nil {
debug = fmt.Sprintf("can't fetch debug information for job, failed to convert recordID to int: %s", convErr.Error())
} else {
var debugErr error
debug, debugErr = s.fetchDebugInformationForJob(ctx, intId)
if debugErr != nil {
s.logger.Error("failed to fetch debug information for job",
log.String("recordID", recordID),
log.Error(debugErr),
)
}
}
s.logger.Error("heartbeat lost a job",
log.Int("recordID", recordID),
log.String("recordID", recordID),
log.String("debug", debug),
log.String("options.workerHostname", options.WorkerHostname),
)

10
internal/workerutil/dbworker/store/store_test.go
View File

@ -1045,7 +1045,7 @@ func TestStoreHeartbeat(t *testing.T) {
clock.Advance(5 * time.Second)
if _, _, err := store.Heartbeat(context.Background(), []int{1, 2, 3}, HeartbeatOptions{}); err != nil {
if _, _, err := store.Heartbeat(context.Background(), []string{"1", "2", "3"}, HeartbeatOptions{}); err != nil {
t.Fatalf("unexpected error updating heartbeat: %s", err)
}
readAndCompareTimes(map[int]time.Duration{
@ -1062,7 +1062,7 @@ func TestStoreHeartbeat(t *testing.T) {
clock.Advance(5 * time.Second)
// Only one worker
if _, _, err := store.Heartbeat(context.Background(), []int{1, 2, 3}, HeartbeatOptions{WorkerHostname: "worker1"}); err != nil {
if _, _, err := store.Heartbeat(context.Background(), []string{"1", "2", "3"}, HeartbeatOptions{WorkerHostname: "worker1"}); err != nil {
t.Fatalf("unexpected error updating heartbeat: %s", err)
}
readAndCompareTimes(map[int]time.Duration{
@ -1074,7 +1074,7 @@ func TestStoreHeartbeat(t *testing.T) {
clock.Advance(5 * time.Second)
// Multiple workers
if _, _, err := store.Heartbeat(context.Background(), []int{1, 3}, HeartbeatOptions{}); err != nil {
if _, _, err := store.Heartbeat(context.Background(), []string{"1", "3"}, HeartbeatOptions{}); err != nil {
t.Fatalf("unexpected error updating heartbeat: %s", err)
}
readAndCompareTimes(map[int]time.Duration{
@ -1102,10 +1102,10 @@ func TestStoreCanceledJobs(t *testing.T) {
t.Fatalf("unexpected error inserting records: %s", err)
}
_, toCancel, err := testStore(db, defaultTestStoreOptions(nil, testScanRecord)).Heartbeat(context.Background(), []int{1, 2, 3}, HeartbeatOptions{WorkerHostname: "worker1"})
_, toCancel, err := testStore(db, defaultTestStoreOptions(nil, testScanRecord)).Heartbeat(context.Background(), []string{"1", "2", "3"}, HeartbeatOptions{WorkerHostname: "worker1"})
if err != nil {
t.Fatalf("unexpected error fetching canceled jobs: %s", err)
}
require.ElementsMatch(t, toCancel, []int{3}, "invalid set of jobs returned")
require.ElementsMatch(t, toCancel, []string{"3"}, "invalid set of jobs returned")
}

2
internal/workerutil/dbworker/store_shim.go
View File

@ -41,7 +41,7 @@ func (s *storeShim[T]) Dequeue(ctx context.Context, workerHostname string, extra
return s.Store.Dequeue(ctx, workerHostname, conditions)
}
func (s *storeShim[T]) Heartbeat(ctx context.Context, ids []int) (knownIDs, cancelIDs []int, err error) {
func (s *storeShim[T]) Heartbeat(ctx context.Context, ids []string) (knownIDs, cancelIDs []string, err error) {
return s.Store.Heartbeat(ctx, ids, store.HeartbeatOptions{})
}

18
internal/workerutil/idset.go
View File

@ -8,17 +8,17 @@ import (
type IDSet struct {
sync.RWMutex
ids map[int]context.CancelFunc
ids map[string]context.CancelFunc
}
func newIDSet() *IDSet {
return &IDSet{ids: map[int]context.CancelFunc{}}
return &IDSet{ids: map[string]context.CancelFunc{}}
}
// Add associates the given identifier with the given cancel function
// in the set. If the identifier was already present then the set is
// unchanged.
func (i *IDSet) Add(id int, cancel context.CancelFunc) bool {
func (i *IDSet) Add(id string, cancel context.CancelFunc) bool {
i.Lock()
defer i.Unlock()
@ -33,7 +33,7 @@ func (i *IDSet) Add(id int, cancel context.CancelFunc) bool {
// Remove invokes the cancel function associated with the given identifier
// in the set and removes the identifier from the set. If the identifier is
// not a member of the set, then no action is performed.
func (i *IDSet) Remove(id int) bool {
func (i *IDSet) Remove(id string) bool {
i.Lock()
cancel, ok := i.ids[id]
delete(i.ids, id)
@ -49,7 +49,7 @@ func (i *IDSet) Remove(id int) bool {
// Remove invokes the cancel function associated with the given identifier
// in the set. If the identifier is not a member of the set, then no action
// is performed.
func (i *IDSet) Cancel(id int) {
func (i *IDSet) Cancel(id string) {
i.RLock()
cancel, ok := i.ids[id]
i.RUnlock()
@ -60,21 +60,21 @@ func (i *IDSet) Cancel(id int) {
}
// Slice returns an ordered copy of the identifiers composing the set.
func (i *IDSet) Slice() []int {
func (i *IDSet) Slice() []string {
i.RLock()
defer i.RUnlock()
ids := make([]int, 0, len(i.ids))
ids := make([]string, 0, len(i.ids))
for id := range i.ids {
ids = append(ids, id)
}
sort.Ints(ids)
sort.Strings(ids)
return ids
}
// Has returns whether the IDSet contains the given id.
func (i *IDSet) Has(id int) bool {
func (i *IDSet) Has(id string) bool {
for _, have := range i.Slice() {
if id == have {
return true

22
internal/workerutil/idset_test.go
View File

@ -10,17 +10,17 @@ func TestIDAddRemove(t *testing.T) {
var called1, called2, called3 bool
idSet := newIDSet()
if !idSet.Add(1, func() { called1 = true }) {
if !idSet.Add("1", func() { called1 = true }) {
t.Fatalf("expected add to succeed")
}
if !idSet.Add(2, func() { called2 = true }) {
if !idSet.Add("2", func() { called2 = true }) {
t.Fatalf("expected add to succeed")
}
if idSet.Add(1, func() { called3 = true }) {
if idSet.Add("1", func() { called3 = true }) {
t.Fatalf("expected duplicate add to fail")
}
idSet.Remove(1)
idSet.Remove("1")
if !called1 {
t.Fatalf("expected first function to be called")
@ -32,20 +32,20 @@ func TestIDAddRemove(t *testing.T) {
t.Fatalf("did not expect third function to be called")
}
if diff := cmp.Diff([]int{2}, idSet.Slice()); diff != "" {
if diff := cmp.Diff([]string{"2"}, idSet.Slice()); diff != "" {
t.Errorf("unexpected slice (-want +got):\n%s", diff)
}
}
func TestIDSetSlice(t *testing.T) {
idSet := newIDSet()
idSet.Add(2, nil)
idSet.Add(4, nil)
idSet.Add(5, nil)
idSet.Add(1, nil)
idSet.Add(3, nil)
idSet.Add("2", nil)
idSet.Add("4", nil)
idSet.Add("5", nil)
idSet.Add("1", nil)
idSet.Add("3", nil)
if diff := cmp.Diff([]int{1, 2, 3, 4, 5}, idSet.Slice()); diff != "" {
if diff := cmp.Diff([]string{"1", "2", "3", "4", "5"}, idSet.Slice()); diff != "" {
t.Errorf("unexpected slice (-want +got):\n%s", diff)
}
}

30
internal/workerutil/mocks_test.go generated
View File

@ -197,7 +197,7 @@ func NewMockStore[T Record]() *MockStore[T] {
},
},
HeartbeatFunc: &StoreHeartbeatFunc[T]{
defaultHook: func(context.Context, []int) (r0 []int, r1 []int, r2 error) {
defaultHook: func(context.Context, []string) (r0 []string, r1 []string, r2 error) {
return
},
},
@ -234,7 +234,7 @@ func NewStrictMockStore[T Record]() *MockStore[T] {
},
},
HeartbeatFunc: &StoreHeartbeatFunc[T]{
defaultHook: func(context.Context, []int) ([]int, []int, error) {
defaultHook: func(context.Context, []string) ([]string, []string, error) {
panic("unexpected invocation of MockStore.Heartbeat")
},
},
@ -402,15 +402,15 @@ func (c StoreDequeueFuncCall[T]) Results() []interface{} {
// StoreHeartbeatFunc describes the behavior when the Heartbeat method of
// the parent MockStore instance is invoked.
type StoreHeartbeatFunc[T Record] struct {
defaultHook func(context.Context, []int) ([]int, []int, error)
hooks []func(context.Context, []int) ([]int, []int, error)
defaultHook func(context.Context, []string) ([]string, []string, error)
hooks []func(context.Context, []string) ([]string, []string, error)
history []StoreHeartbeatFuncCall[T]
mutex sync.Mutex
}
// Heartbeat delegates to the next hook function in the queue and stores the
// parameter and result values of this invocation.
func (m *MockStore[T]) Heartbeat(v0 context.Context, v1 []int) ([]int, []int, error) {
func (m *MockStore[T]) Heartbeat(v0 context.Context, v1 []string) ([]string, []string, error) {
r0, r1, r2 := m.HeartbeatFunc.nextHook()(v0, v1)
m.HeartbeatFunc.appendCall(StoreHeartbeatFuncCall[T]{v0, v1, r0, r1, r2})
return r0, r1, r2
@ -418,7 +418,7 @@ func (m *MockStore[T]) Heartbeat(v0 context.Context, v1 []int) ([]int, []int, er
// SetDefaultHook sets function that is called when the Heartbeat method of
// the parent MockStore instance is invoked and the hook queue is empty.
func (f *StoreHeartbeatFunc[T]) SetDefaultHook(hook func(context.Context, []int) ([]int, []int, error)) {
func (f *StoreHeartbeatFunc[T]) SetDefaultHook(hook func(context.Context, []string) ([]string, []string, error)) {
f.defaultHook = hook
}
@ -426,7 +426,7 @@ func (f *StoreHeartbeatFunc[T]) SetDefaultHook(hook func(context.Context, []int)
// Heartbeat method of the parent MockStore instance invokes the hook at the
// front of the queue and discards it. After the queue is empty, the default
// hook function is invoked for any future action.
func (f *StoreHeartbeatFunc[T]) PushHook(hook func(context.Context, []int) ([]int, []int, error)) {
func (f *StoreHeartbeatFunc[T]) PushHook(hook func(context.Context, []string) ([]string, []string, error)) {
f.mutex.Lock()
f.hooks = append(f.hooks, hook)
f.mutex.Unlock()
@ -434,20 +434,20 @@ func (f *StoreHeartbeatFunc[T]) PushHook(hook func(context.Context, []int) ([]in
// SetDefaultReturn calls SetDefaultHook with a function that returns the
// given values.
func (f *StoreHeartbeatFunc[T]) SetDefaultReturn(r0 []int, r1 []int, r2 error) {
f.SetDefaultHook(func(context.Context, []int) ([]int, []int, error) {
func (f *StoreHeartbeatFunc[T]) SetDefaultReturn(r0 []string, r1 []string, r2 error) {
f.SetDefaultHook(func(context.Context, []string) ([]string, []string, error) {
return r0, r1, r2
})
}
// PushReturn calls PushHook with a function that returns the given values.
func (f *StoreHeartbeatFunc[T]) PushReturn(r0 []int, r1 []int, r2 error) {
f.PushHook(func(context.Context, []int) ([]int, []int, error) {
func (f *StoreHeartbeatFunc[T]) PushReturn(r0 []string, r1 []string, r2 error) {
f.PushHook(func(context.Context, []string) ([]string, []string, error) {
return r0, r1, r2
})
}
func (f *StoreHeartbeatFunc[T]) nextHook() func(context.Context, []int) ([]int, []int, error) {
func (f *StoreHeartbeatFunc[T]) nextHook() func(context.Context, []string) ([]string, []string, error) {
f.mutex.Lock()
defer f.mutex.Unlock()
@ -485,13 +485,13 @@ type StoreHeartbeatFuncCall[T Record] struct {
Arg0 context.Context
// Arg1 is the value of the 2nd argument passed to this method
// invocation.
Arg1 []int
Arg1 []string
// Result0 is the value of the 1st result returned from this method
// invocation.
Result0 []int
Result0 []string
// Result1 is the value of the 2nd result returned from this method
// invocation.
Result1 []int
Result1 []string
// Result2 is the value of the 3rd result returned from this method
// invocation.
Result2 error

4
internal/workerutil/store.go
View File

@ -8,6 +8,8 @@ import (
type Record interface {
// RecordID returns the integer primary key of the record.
RecordID() int
// RecordUID returns a UID of the record, of which the format is defined by the concrete type.
RecordUID() string
}
// Store is the persistence layer for the workerutil package that handles worker-side operations.
@ -22,7 +24,7 @@ type Store[T Record] interface {
// Heartbeat updates last_heartbeat_at of all the given jobs, when they're processing. All IDs of records that were
// touched are returned. Additionally, jobs in the working set that are flagged as to be canceled are returned.
Heartbeat(ctx context.Context, jobIDs []int) (knownIDs, cancelIDs []int, err error)
Heartbeat(ctx context.Context, jobIDs []string) (knownIDs, cancelIDs []string, err error)
// MarkComplete attempts to update the state of the record to complete. This method returns a boolean flag indicating
// if the record was updated.

25
internal/workerutil/worker.go
View File

@ -3,6 +3,7 @@ package workerutil
import (
"context"
"fmt"
"strconv"
"sync"
"time"
@ -49,6 +50,10 @@ type dummyType struct{}
func (d dummyType) RecordID() int { return 0 }
func (d dummyType) RecordUID() string {
return strconv.Itoa(0)
}
var _ recorder.Recordable = &Worker[dummyType]{}
type WorkerOptions struct {
@ -162,12 +167,12 @@ func (w *Worker[T]) Start() {
knownIDs, canceledIDs, err := w.store.Heartbeat(w.rootCtx, ids)
if err != nil {
w.options.Metrics.logger.Error("Failed to refresh heartbeats",
log.Ints("ids", ids),
log.Strings("ids", ids),
log.Error(err))
// Bail out and restart the for loop.
continue
}
knownIDsMap := map[int]struct{}{}
knownIDsMap := map[string]struct{}{}
for _, id := range knownIDs {
knownIDsMap[id] = struct{}{}
}
@ -176,13 +181,13 @@ func (w *Worker[T]) Start() {
if _, ok := knownIDsMap[id]; !ok {
if w.runningIDSet.Remove(id) {
w.options.Metrics.logger.Error("Removed unknown job from running set",
log.Int("id", id))
log.String("id", id))
}
}
}
if len(canceledIDs) > 0 {
w.options.Metrics.logger.Info("Found jobs to cancel", log.Ints("IDs", canceledIDs))
w.options.Metrics.logger.Info("Found jobs to cancel", log.Strings("IDs", canceledIDs))
}
for _, id := range canceledIDs {
@ -316,7 +321,7 @@ func (w *Worker[T]) dequeueAndHandle() (dequeued bool, err error) {
processLog := trace.Logger(workerCtxWithSpan, w.options.Metrics.logger)
// Register the record as running so it is included in heartbeat updates.
if !w.runningIDSet.Add(record.RecordID(), cancel) {
if !w.runningIDSet.Add(record.RecordUID(), cancel) {
workerSpan.LogFields(otlog.Error(ErrJobAlreadyExists))
workerSpan.Finish()
return false, ErrJobAlreadyExists
@ -324,9 +329,9 @@ func (w *Worker[T]) dequeueAndHandle() (dequeued bool, err error) {
// Set up observability
w.options.Metrics.numJobs.Inc()
processLog.Info("Dequeued record for processing", log.Int("id", record.RecordID()))
processLog.Info("Dequeued record for processing", log.String("id", record.RecordUID()))
processArgs := observation.Args{
Attrs: []attribute.KeyValue{attribute.Int("record.id", record.RecordID())},
Attrs: []attribute.KeyValue{attribute.String("record.id", record.RecordUID())},
}
if hook, ok := w.handler.(WithHooks[T]); ok {
@ -353,7 +358,7 @@ func (w *Worker[T]) dequeueAndHandle() (dequeued bool, err error) {
// Remove the record from the set of running jobs, so it is not included
// in heartbeat updates anymore.
defer w.runningIDSet.Remove(record.RecordID())
defer w.runningIDSet.Remove(record.RecordUID())
w.options.Metrics.numJobs.Dec()
w.handlerSemaphore <- struct{}{}
w.wg.Done()
@ -400,7 +405,7 @@ func (w *Worker[T]) handle(ctx, workerContext context.Context, record T) (err er
go w.recorder.LogRun(w, duration, handleErr)
}
if errcode.IsNonRetryable(handleErr) || handleErr != nil && w.isJobCanceled(record.RecordID(), handleErr, ctx.Err()) {
if errcode.IsNonRetryable(handleErr) || handleErr != nil && w.isJobCanceled(record.RecordUID(), handleErr, ctx.Err()) {
if marked, markErr := w.store.MarkFailed(workerContext, record, handleErr.Error()); markErr != nil {
return errors.Wrap(markErr, "store.MarkFailed")
} else if marked {
@ -427,7 +432,7 @@ func (w *Worker[T]) handle(ctx, workerContext context.Context, record T) (err er
// isJobCanceled returns true if the job has been canceled through the Cancel interface.
// If the context is canceled, and the job is still part of the running ID set,
// we know that it has been canceled for that reason.
func (w *Worker[T]) isJobCanceled(id int, handleErr, ctxErr error) bool {
func (w *Worker[T]) isJobCanceled(id string, handleErr, ctxErr error) bool {
return errors.Is(handleErr, ctxErr) && w.runningIDSet.Has(id) && !errors.Is(handleErr, context.DeadlineExceeded)
}

11
internal/workerutil/worker_test.go
View File

@ -3,6 +3,7 @@ package workerutil
import (
"context"
"fmt"
"strconv"
"strings"
"sync"
"testing"
@ -25,6 +26,10 @@ func (v TestRecord) RecordID() int {
return v.ID
}
func (v TestRecord) RecordUID() string {
return strconv.Itoa(v.ID)
}
func TestWorkerHandlerSuccess(t *testing.T) {
store := NewMockStore[*TestRecord]()
handler := NewMockHandler[*TestRecord]()
@ -372,7 +377,7 @@ func TestWorkerDequeueHeartbeat(t *testing.T) {
}
heartbeats := make(chan struct{})
store.HeartbeatFunc.SetDefaultHook(func(c context.Context, i []int) ([]int, []int, error) {
store.HeartbeatFunc.SetDefaultHook(func(c context.Context, i []string) ([]string, []string, error) {
heartbeats <- struct{}{}
return i, nil, nil
})
@ -524,7 +529,7 @@ func TestWorkerCancelJobs(t *testing.T) {
}
canceledJobsCalled := make(chan struct{})
store.HeartbeatFunc.SetDefaultHook(func(c context.Context, i []int) ([]int, []int, error) {
store.HeartbeatFunc.SetDefaultHook(func(c context.Context, i []string) ([]string, []string, error) {
close(canceledJobsCalled)
// Cancel all jobs.
return i, i, nil
@ -603,7 +608,7 @@ func TestWorkerDeadline(t *testing.T) {
}
heartbeats := make(chan struct{})
store.HeartbeatFunc.SetDefaultHook(func(c context.Context, i []int) ([]int, []int, error) {
store.HeartbeatFunc.SetDefaultHook(func(c context.Context, i []string) ([]string, []string, error) {
heartbeats <- struct{}{}
return i, nil, nil
})