refactor(storage): give the orchestrator its own aggregate, keeping store contracts at domain level

Author: roychying

Makefile

@@ -386,7 +386,7 @@ local-submitqueue-gateway-stop: ## Stop Gateway service
 
 local-init-submitqueue-schemas: ## Manually apply all database schemas
 	@echo "Applying storage schema to mysql-app..."
-	@for file in submitqueue/extension/storage/mysql/schema/*.sql; do \
+	@for file in submitqueue/orchestrator/extension/storage/mysql/schema/*.sql; do \
 		echo "  - Applying $$(basename $$file)..."; \
 		docker exec -i $(SUBMITQUEUE_LOCAL_PROJECT)-mysql-app-1 mysql -uroot -proot submitqueue < $$file 2>&1 | grep -v "Using a password" || true; \
 	done

service/submitqueue/orchestrator/server/BUILD.bazel

@@ -60,11 +60,11 @@ go_library(
         "//submitqueue/extension/speculation/scorer/heuristic:go_default_library",
         "//submitqueue/extension/speculation/speculator:go_default_library",
         "//submitqueue/extension/speculation/speculator/standard:go_default_library",
-        "//submitqueue/extension/storage:go_default_library",
-        "//submitqueue/extension/storage/mysql:go_default_library",
         "//submitqueue/extension/validator:go_default_library",
         "//submitqueue/extension/validator/fake:go_default_library",
         "//submitqueue/orchestrator:go_default_library",
+        "//submitqueue/orchestrator/extension/storage:go_default_library",
+        "//submitqueue/orchestrator/extension/storage/mysql:go_default_library",
         "@com_github_go_sql_driver_mysql//:go_default_library",
         "@com_github_uber_go_tally//:go_default_library",
         "@in_gopkg_yaml_v3//:go_default_library",
@@ -125,7 +125,7 @@ go_test(
         "//submitqueue/extension/conflict:go_default_library",
         "//submitqueue/extension/speculation/scorer:go_default_library",
         "//submitqueue/extension/speculation/speculator:go_default_library",
-        "//submitqueue/extension/storage:go_default_library",
+        "//submitqueue/orchestrator/extension/storage:go_default_library",
         "@com_github_stretchr_testify//assert:go_default_library",
         "@com_github_stretchr_testify//require:go_default_library",
         "@com_github_uber_go_tally//:go_default_library",

service/submitqueue/orchestrator/server/main.go

@@ -26,6 +26,8 @@ import (
 	"syscall"
 	"time"
 
+	orchstorage "github.com/uber/submitqueue/submitqueue/orchestrator/extension/storage"
+
 	_ "github.com/go-sql-driver/mysql"
 
 	"github.com/uber-go/tally"
@@ -44,11 +46,10 @@ import (
 	"github.com/uber/submitqueue/platform/pipeline"
 	servicemq "github.com/uber/submitqueue/service/messagequeue"
 	"github.com/uber/submitqueue/submitqueue/core/changeset"
-	"github.com/uber/submitqueue/submitqueue/extension/storage"
-	mysqlstorage "github.com/uber/submitqueue/submitqueue/extension/storage/mysql"
 	"github.com/uber/submitqueue/submitqueue/extension/validator"
 	validatorfake "github.com/uber/submitqueue/submitqueue/extension/validator/fake"
 	"github.com/uber/submitqueue/submitqueue/orchestrator"
+	mysqlstorage "github.com/uber/submitqueue/submitqueue/orchestrator/extension/storage/mysql"
 	"go.uber.org/zap"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/reflection"
@@ -438,15 +439,15 @@ func getEnv(key, defaultVal string) string {
 }
 
 // storageFactory adapts the MySQL storage backend's queue binding to the
-// storage.Factory seam. Routing every queue to the single shared backend is
+// orchstorage.Factory seam. Routing every queue to the single shared backend is
 // this host's policy; a deployment that splits queues across backends swaps
 // this adapter for a routing one.
 type storageFactory struct {
 	backend *mysqlstorage.Storage
 }
 
 // For returns the queue-scoped store aggregate bound to the queue named in config.
-func (f storageFactory) For(config storage.Config) (storage.Storage, error) {
+func (f storageFactory) For(config orchstorage.Config) (orchstorage.Storage, error) {
 	return f.backend.For(config.QueueName)
 }
 

service/submitqueue/orchestrator/server/profiles.go

@@ -18,6 +18,8 @@ import (
 	"fmt"
 	nethttp "net/http"
 
+	orchstorage "github.com/uber/submitqueue/submitqueue/orchestrator/extension/storage"
+
 	"github.com/uber-go/tally"
 	"go.uber.org/zap"
 	"golang.org/x/oauth2"
@@ -54,7 +56,6 @@ import (
 	"github.com/uber/submitqueue/submitqueue/extension/speculation/scorer/heuristic"
 	"github.com/uber/submitqueue/submitqueue/extension/speculation/speculator"
 	specstandard "github.com/uber/submitqueue/submitqueue/extension/speculation/speculator/standard"
-	"github.com/uber/submitqueue/submitqueue/extension/storage"
 )
 
 // Profile holds the per-queue extension implementations. Grouping them per
@@ -74,7 +75,7 @@ type Profile struct {
 	// Storage resolves the queue-scoped store aggregate for this queue. Every
 	// profile points at the shared backend by default; a deployment that
 	// splits queues across storage backends overrides this per queue.
-	Storage storage.Factory
+	Storage orchstorage.Factory
 
 	// Scorer holds this queue's ranking profile. There is no scoring stage: the
 	// scorer feeds the queue's speculator, which ranks candidate paths by how
@@ -143,10 +144,10 @@ func (p Profiles) ScorerFactory() scorer.Factory {
 	})
 }
 
-// StorageFactory returns a storage.Factory that routes each queue to its
+// StorageFactory returns a orchstorage.Factory that routes each queue to its
 // profile's storage backend before binding the queue-scoped store aggregate.
-func (p Profiles) StorageFactory() storage.Factory {
-	return storageFunc(func(c storage.Config) (storage.Storage, error) {
+func (p Profiles) StorageFactory() orchstorage.Factory {
+	return storageFunc(func(c orchstorage.Config) (orchstorage.Storage, error) {
 		return p.For(c.QueueName).Storage.For(c)
 	})
 }
@@ -169,9 +170,9 @@ type analyzerFunc func(conflict.Config) (conflict.Analyzer, error)
 
 func (f analyzerFunc) For(c conflict.Config) (conflict.Analyzer, error) { return f(c) }
 
-type storageFunc func(storage.Config) (storage.Storage, error)
+type storageFunc func(orchstorage.Config) (orchstorage.Storage, error)
 
-func (f storageFunc) For(c storage.Config) (storage.Storage, error) { return f(c) }
+func (f storageFunc) For(c orchstorage.Config) (orchstorage.Storage, error) { return f(c) }
 
 type scorerFunc func(scorer.Config) (scorer.Scorer, error)
 
@@ -192,7 +193,7 @@ func newProfiles(
 	logger *zap.Logger,
 	scope tally.Scope,
 	resolver changeset.Resolver,
-	stores storage.Factory,
+	stores orchstorage.Factory,
 	cfg profilesConfig,
 ) (Profiles, error) {
 	b := &profileBuilder{
@@ -245,7 +246,7 @@ type profileBuilder struct {
 	logger   *zap.Logger
 	scope    tally.Scope
 	resolver changeset.Resolver
-	stores   storage.Factory
+	stores   orchstorage.Factory
 	built    map[string]any
 }
 

service/submitqueue/orchestrator/server/profiles_test.go

@@ -19,6 +19,8 @@ import (
 	"errors"
 	"testing"
 
+	orchstorage "github.com/uber/submitqueue/submitqueue/orchestrator/extension/storage"
+
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
@@ -28,7 +30,6 @@ import (
 	"github.com/uber/submitqueue/submitqueue/extension/conflict"
 	"github.com/uber/submitqueue/submitqueue/extension/speculation/scorer"
 	"github.com/uber/submitqueue/submitqueue/extension/speculation/speculator"
-	"github.com/uber/submitqueue/submitqueue/extension/storage"
 )
 
 // recorder captures the queue name each seam's factory was handed, so a test
@@ -66,7 +67,7 @@ func profileRecording(rec *recorder) Profile {
 			rec.analyzer = c.QueueName
 			return nil, nil
 		}),
-		Storage: storageFunc(func(c storage.Config) (storage.Storage, error) {
+		Storage: storageFunc(func(c orchstorage.Config) (orchstorage.Storage, error) {
 			rec.storage = c.QueueName
 			return nil, nil
 		}),
@@ -109,7 +110,7 @@ func TestProfilesForwardQueueNameToFactories(t *testing.T) {
 			require.NoError(t, err)
 			_, err = profiles.AnalyzerFactory().For(conflict.Config{QueueName: tt.queue})
 			require.NoError(t, err)
-			_, err = profiles.StorageFactory().For(storage.Config{QueueName: tt.queue})
+			_, err = profiles.StorageFactory().For(orchstorage.Config{QueueName: tt.queue})
 			require.NoError(t, err)
 			_, err = profiles.ScorerFactory().For(scorer.Config{QueueName: tt.queue})
 			require.NoError(t, err)

submitqueue/core/batch/BUILD.bazel

@@ -12,6 +12,7 @@ go_library(
     deps = [
         "//submitqueue/entity:go_default_library",
         "//submitqueue/extension/storage:go_default_library",
+        "//submitqueue/orchestrator/extension/storage:go_default_library",
         "@org_golang_x_sync//errgroup:go_default_library",
     ],
 )
@@ -28,6 +29,7 @@ go_test(
         "//submitqueue/entity:go_default_library",
         "//submitqueue/extension/storage:go_default_library",
         "//submitqueue/extension/storage/mock:go_default_library",
+        "//submitqueue/orchestrator/extension/storage/mock:go_default_library",
         "@com_github_stretchr_testify//assert:go_default_library",
         "@com_github_stretchr_testify//require:go_default_library",
         "@org_uber_go_mock//gomock:go_default_library",

submitqueue/core/batch/find.go

@@ -21,7 +21,8 @@ import (
 	"sort"
 
 	"github.com/uber/submitqueue/submitqueue/entity"
-	"github.com/uber/submitqueue/submitqueue/extension/storage"
+	storage "github.com/uber/submitqueue/submitqueue/extension/storage"
+	orchstorage "github.com/uber/submitqueue/submitqueue/orchestrator/extension/storage"
 )
 
 // FindByRequestID resolves every batch attempt associated with a request,
@@ -35,7 +36,7 @@ import (
 //
 // Unlike ListByStates, which treats a dangling membership record as store
 // corruption, a dangling association is an expected retry artifact.
-func FindByRequestID(ctx context.Context, store storage.Storage, requestID string) ([]entity.Batch, int, error) {
+func FindByRequestID(ctx context.Context, store orchstorage.Storage, requestID string) ([]entity.Batch, int, error) {
 	associations, err := store.GetRequestBatchStore().GetByRequestID(ctx, requestID)
 	if err != nil {
 		return nil, 0, fmt.Errorf("failed to get batch associations for request %s: %w", requestID, err)

submitqueue/core/batch/find_test.go

@@ -26,6 +26,7 @@ import (
 	"github.com/uber/submitqueue/submitqueue/entity"
 	"github.com/uber/submitqueue/submitqueue/extension/storage"
 	storagemock "github.com/uber/submitqueue/submitqueue/extension/storage/mock"
+	orchstoragemock "github.com/uber/submitqueue/submitqueue/orchestrator/extension/storage/mock"
 )
 
 const testRequestID = "monorepo/4"
@@ -36,11 +37,11 @@ func association(batchID string) entity.RequestBatch {
 }
 
 // findStores wires a MockStorage over a batch store and a request-batch store.
-func findStores(t *testing.T) (*storagemock.MockStorage, *storagemock.MockBatchStore, *storagemock.MockRequestBatchStore) {
+func findStores(t *testing.T) (*orchstoragemock.MockStorage, *storagemock.MockBatchStore, *storagemock.MockRequestBatchStore) {
 	t.Helper()
 
 	ctrl := gomock.NewController(t)
-	mockStorage := storagemock.NewMockStorage(ctrl)
+	mockStorage := orchstoragemock.NewMockStorage(ctrl)
 	mockBatchStore := storagemock.NewMockBatchStore(ctrl)
 	mockAssociationStore := storagemock.NewMockRequestBatchStore(ctrl)
 	mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes()

submitqueue/core/batch/list.go

@@ -21,7 +21,7 @@ import (
 	"golang.org/x/sync/errgroup"
 
 	"github.com/uber/submitqueue/submitqueue/entity"
-	"github.com/uber/submitqueue/submitqueue/extension/storage"
+	orchstorage "github.com/uber/submitqueue/submitqueue/orchestrator/extension/storage"
 )
 
 // hydrateConcurrency bounds the parallel per-key batch reads a single
@@ -39,7 +39,7 @@ const hydrateConcurrency = 16
 // A candidate ID whose batch does not exist is returned as an error rather than
 // skipped: batch rows are never deleted, so a dangling record means the store is
 // inconsistent, not that the batch concluded.
-func ListByStates(ctx context.Context, store storage.Storage, states []entity.BatchState) ([]entity.Batch, error) {
+func ListByStates(ctx context.Context, store orchstorage.Storage, states []entity.BatchState) ([]entity.Batch, error) {
 	wanted := make(map[entity.BatchState]bool, len(states))
 	seen := make(map[string]bool)
 	var ids []string

submitqueue/core/batch/transition.go

@@ -36,7 +36,7 @@ import (
 	"fmt"
 
 	"github.com/uber/submitqueue/submitqueue/entity"
-	"github.com/uber/submitqueue/submitqueue/extension/storage"
+	orchstorage "github.com/uber/submitqueue/submitqueue/orchestrator/extension/storage"
 )
 
 // Transition moves a batch to newState: it performs the optimistic-locking CAS on
@@ -52,7 +52,7 @@ import (
 // applied — the CAS may have committed with the record move incomplete — and the
 // caller is expected to let redelivery retry; the retry's already-in-target-state
 // branch repairs the record via EnsureRecord.
-func Transition(ctx context.Context, store storage.Storage, batch entity.Batch, newState entity.BatchState) (entity.Batch, error) {
+func Transition(ctx context.Context, store orchstorage.Storage, batch entity.Batch, newState entity.BatchState) (entity.Batch, error) {
 	oldState := batch.State
 	newVersion := batch.Version + 1
 	updated := batch
@@ -78,7 +78,7 @@ func Transition(ctx context.Context, store storage.Storage, batch entity.Batch,
 // the repair half of the transition protocol: idempotent redelivery branches that
 // skip the CAS because the batch is already in the target state call this instead,
 // covering a prior attempt that crashed between the CAS and the record move.
-func EnsureRecord(ctx context.Context, store storage.Storage, batch entity.Batch) error {
+func EnsureRecord(ctx context.Context, store orchstorage.Storage, batch entity.Batch) error {
 	record := entity.QueueBatchState{Queue: batch.Queue, State: batch.State, BatchID: batch.ID}
 	if err := store.GetQueueBatchStateStore().Put(ctx, record); err != nil {
 		return fmt.Errorf("failed to put queue batch state record for batch %s under state %s: %w", batch.ID, batch.State, err)