Merge branch 'main' into preetam/orch-stages

Author: behinddwalls

entity/BUILD.bazel

@@ -8,6 +8,7 @@ go_library(
         "build.go",
         "change_provider.go",
         "request.go",
+        "speculation_tree.go",
     ],
     importpath = "github.com/uber/submitqueue/entity",
     visibility = ["//visibility:public"],

entity/speculation_tree.go

@@ -0,0 +1,37 @@
+package entity
+
+// SpeculationPathAction defines the possible actions for a speculation path.
+type SpeculationPathAction string
+
+const (
+	// SpeculationPathActionUnknown is the default zero value for SpeculationPathAction.
+	SpeculationPathActionUnknown SpeculationPathAction = ""
+	// TODO: Add comprehensive list of actions
+)
+
+// SpeculationInfo represents metadata about a single speculation path, including the path through the dependency graph, its current state, and the predicted build score.
+type SpeculationInfo struct {
+	// Path represents the speculation path; which is an ordered list of batches.
+	Path []string
+	// Action is a state that this path is in.
+	Action SpeculationPathAction
+	// Score is score for this speculation path.
+	Score float32
+}
+
+// SpeculationTree represents the set of speculation paths constructed for a batch based on its dependency graph.
+type SpeculationTree struct {
+	// BatchID is the batch for which this speculation tree is constructed.
+	BatchID string
+	// Speculations is a list of speculation paths for this batch based on a graph of its
+	// dependents.
+	//
+	// For e.g - Consider batches - queueA/batch/1, queueA/batch/2, queueA/batch/3
+	// such that - queueA/batch/2 and queueA/batch/3 depend on queueA/batch/1
+	//
+	// Speculations for queueA/batch/1 - [{Path: []string{"queueA/batch/1"}, State: "scheduled", Score: 0.1}]
+	// Speculations for queueA/batch/2 - [{Path: []string{"queueA/batch/2"}, State: "scheduled", Score: 0.9}, {Path: []string{"queueA/batch/1", "queueA/batch/2"}, State: "scheduled", Score: 0.3}]
+	// Speculations for queueA/batch/3 - [{Path: []string{"queueA/batch/3"}, State: "scheduled", Score: 0.9}, {Path: []string{"queueA/batch/1", "queueA/batch/3"}, State: "scheduled", Score: 0.3}]
+	//
+	Speculations []SpeculationInfo
+}

example/server/docker-compose.yml

@@ -18,7 +18,10 @@ services:
     ports:
       - "3306"  # Random ephemeral port to avoid conflicts
     healthcheck:
-      test: ["CMD", "mysqladmin", "ping", "-h", "localhost", "-proot"]
+      # Use 127.0.0.1 (TCP) instead of localhost (Unix socket). MySQL treats
+      # "localhost" as a socket connection, which can be ready before the TCP
+      # listener — causing dependent services that connect over TCP to fail.
+      test: ["CMD", "mysqladmin", "ping", "-h", "127.0.0.1", "-proot"]
       interval: 5s
       timeout: 5s
       retries: 10
@@ -33,7 +36,7 @@ services:
     ports:
       - "3306"  # Random ephemeral port to avoid conflicts
     healthcheck:
-      test: ["CMD", "mysqladmin", "ping", "-h", "localhost", "-proot"]
+      test: ["CMD", "mysqladmin", "ping", "-h", "127.0.0.1", "-proot"]
       interval: 5s
       timeout: 5s
       retries: 10

example/server/gateway/docker-compose.yml

@@ -18,7 +18,10 @@ services:
     ports:
       - "3306"  # Random ephemeral port to avoid conflicts
     healthcheck:
-      test: ["CMD", "mysqladmin", "ping", "-h", "localhost", "-proot"]
+      # Use 127.0.0.1 (TCP) instead of localhost (Unix socket). MySQL treats
+      # "localhost" as a socket connection, which can be ready before the TCP
+      # listener — causing dependent services that connect over TCP to fail.
+      test: ["CMD", "mysqladmin", "ping", "-h", "127.0.0.1", "-proot"]
       interval: 5s
       timeout: 5s
       retries: 10
@@ -33,7 +36,7 @@ services:
     ports:
       - "3306"  # Random ephemeral port to avoid conflicts
     healthcheck:
-      test: ["CMD", "mysqladmin", "ping", "-h", "localhost", "-proot"]
+      test: ["CMD", "mysqladmin", "ping", "-h", "127.0.0.1", "-proot"]
       interval: 5s
       timeout: 5s
       retries: 10

example/server/orchestrator/docker-compose.yml

@@ -18,7 +18,10 @@ services:
     ports:
       - "3306"  # Random ephemeral port to avoid conflicts
     healthcheck:
-      test: ["CMD", "mysqladmin", "ping", "-h", "localhost", "-proot"]
+      # Use 127.0.0.1 (TCP) instead of localhost (Unix socket). MySQL treats
+      # "localhost" as a socket connection, which can be ready before the TCP
+      # listener — causing dependent services that connect over TCP to fail.
+      test: ["CMD", "mysqladmin", "ping", "-h", "127.0.0.1", "-proot"]
       interval: 5s
       timeout: 5s
       retries: 10
@@ -33,7 +36,7 @@ services:
     ports:
       - "3306"  # Random ephemeral port to avoid conflicts
     healthcheck:
-      test: ["CMD", "mysqladmin", "ping", "-h", "localhost", "-proot"]
+      test: ["CMD", "mysqladmin", "ping", "-h", "127.0.0.1", "-proot"]
       interval: 5s
       timeout: 5s
       retries: 10

extension/storage/BUILD.bazel

@@ -3,6 +3,9 @@ load("@rules_go//go:def.bzl", "go_library")
 go_library(
     name = "storage",
     srcs = [
+        "batch_dependent_store.go",
+        "batch_store.go",
+        "build_store.go",
         "change_provider_store.go",
         "request_store.go",
         "storage.go",

extension/storage/batch_dependent_store.go

@@ -0,0 +1,21 @@
+package storage
+
+import (
+	"context"
+
+	"github.com/uber/submitqueue/entity"
+)
+
+// BatchDependentStore is an interface that defines methods for managing batch dependent information in the database.
+type BatchDependentStore interface {
+	// Get retrieves the batch dependent by batch ID.
+	// Returns ErrNotFound if the batch dependent is not found.
+	Get(ctx context.Context, batchID string) (entity.BatchDependent, error)
+
+	// Create creates a new batch dependent.
+	// Returns ErrAlreadyExists if the entry already exists.
+	Create(ctx context.Context, batchDependent entity.BatchDependent) error
+
+	// There is no update function since once created, data is only ever read from this
+	// store.
+}

extension/storage/batch_store.go

@@ -0,0 +1,21 @@
+package storage
+
+import (
+	"context"
+
+	"github.com/uber/submitqueue/entity"
+)
+
+// BatchStore is an interface that defines methods for managing batches in the database.
+type BatchStore interface {
+	// Get retrieves a batch by ID. Returns ErrNotFound if the batch is not found.
+	Get(ctx context.Context, id string) (entity.Batch, error)
+
+	// Create creates a new batch. The batch must have a unique ID already assigned.
+	// Returns ErrAlreadyExists if a batch with the same ID already exists.
+	Create(ctx context.Context, batch entity.Batch) error
+
+	// UpdateState updates the state of a batch if the current version matches the expected version. If versions do not match, returns ErrVersionMismatch.
+	// The implementation should increment the version by 1 atomically with the state update.
+	UpdateState(ctx context.Context, id string, version int32, newState entity.BatchState) error
+}

extension/storage/build_store.go

@@ -0,0 +1,20 @@
+package storage
+
+import (
+	"context"
+
+	"github.com/uber/submitqueue/entity"
+)
+
+// BuildStore is an interface that defines methods for managing builds in the database.
+type BuildStore interface {
+	// Get retrieves a build by ID. Returns ErrNotFound if the build is not found.
+	Get(ctx context.Context, id string) (entity.Build, error)
+
+	// Create creates a new build. The build must have a unique ID already assigned.
+	// Returns ErrAlreadyExists if a build with the same ID already exists.
+	Create(ctx context.Context, build entity.Build) error
+
+	// UpdateStatus updates the status of a build.
+	UpdateStatus(ctx context.Context, id string, newStatus entity.BuildStatus) error
+}

extension/storage/mysql/BUILD.bazel

@@ -3,6 +3,9 @@ load("@rules_go//go:def.bzl", "go_library")
 go_library(
     name = "mysql",
     srcs = [
+        "batch_dependent_store.go",
+        "batch_store.go",
+        "build_store.go",
         "change_provider_store.go",
         "request_store.go",
         "storage.go",