feat: add worker distributor and model registry

internal/worker/queries.sql

@@ -2,14 +2,23 @@
 SELECT 1;
 
 -- name: RegisterWorker :exec
-INSERT INTO worker_register (worker_id)
+INSERT INTO worker_register (worker_id, capacity)
 VALUES (
-    $1
+      $1
+    , $2
 );
 
+-- name: GetWorkers :many
+SELECT 
+      worker_id
+    , capacity
+FROM
+    worker_register;
+
 -- name: UpdateWorkerHeartbeat :exec
 UPDATE worker_register
 SET
     heart_beat = now()
 WHERE
     worker_id = $1;
+

internal/worker/repositories/models.go

@@ -9,7 +9,22 @@ import (
 	"github.com/jackc/pgx/v5/pgtype"
 )
 
+type ModelSchedule struct {
+	ModelName string             `json:"model_name"`
+	LastRun   pgtype.Timestamptz `json:"last_run"`
+}
+
+type WorkSchedule struct {
+	ScheduleID uuid.UUID          `json:"schedule_id"`
+	WorkerID   uuid.UUID          `json:"worker_id"`
+	StartRun   pgtype.Timestamptz `json:"start_run"`
+	EndRun     pgtype.Timestamptz `json:"end_run"`
+	UpdatedAt  pgtype.Timestamptz `json:"updated_at"`
+	State      string             `json:"state"`
+}
+
 type WorkerRegister struct {
 	WorkerID  uuid.UUID          `json:"worker_id"`
+	Capacity  int32              `json:"capacity"`
 	HeartBeat pgtype.Timestamptz `json:"heart_beat"`
 }

internal/worker/repositories/querier.go

@@ -11,8 +11,9 @@ import (
 )
 
 type Querier interface {
+	GetWorkers(ctx context.Context) ([]*GetWorkersRow, error)
 	Ping(ctx context.Context) (int32, error)
-	RegisterWorker(ctx context.Context, workerID uuid.UUID) error
+	RegisterWorker(ctx context.Context, arg *RegisterWorkerParams) error
 	UpdateWorkerHeartbeat(ctx context.Context, workerID uuid.UUID) error
 }
 

internal/worker/repositories/queries.sql.go

@@ -11,6 +11,39 @@ import (
 	"github.com/google/uuid"
 )
 
+const getWorkers = `-- name: GetWorkers :many
+SELECT 
+      worker_id
+    , capacity
+FROM
+    worker_register
+`
+
+type GetWorkersRow struct {
+	WorkerID uuid.UUID `json:"worker_id"`
+	Capacity int32     `json:"capacity"`
+}
+
+func (q *Queries) GetWorkers(ctx context.Context) ([]*GetWorkersRow, error) {
+	rows, err := q.db.Query(ctx, getWorkers)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+	items := []*GetWorkersRow{}
+	for rows.Next() {
+		var i GetWorkersRow
+		if err := rows.Scan(&i.WorkerID, &i.Capacity); err != nil {
+			return nil, err
+		}
+		items = append(items, &i)
+	}
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+	return items, nil
+}
+
 const ping = `-- name: Ping :one
 SELECT 1
 `
@@ -23,14 +56,20 @@ func (q *Queries) Ping(ctx context.Context) (int32, error) {
 }
 
 const registerWorker = `-- name: RegisterWorker :exec
-INSERT INTO worker_register (worker_id)
+INSERT INTO worker_register (worker_id, capacity)
 VALUES (
-    $1
+      $1
+    , $2
 )
 `
 
-func (q *Queries) RegisterWorker(ctx context.Context, workerID uuid.UUID) error {
-	_, err := q.db.Exec(ctx, registerWorker, workerID)
+type RegisterWorkerParams struct {
+	WorkerID uuid.UUID `json:"worker_id"`
+	Capacity int32     `json:"capacity"`
+}
+
+func (q *Queries) RegisterWorker(ctx context.Context, arg *RegisterWorkerParams) error {
+	_, err := q.db.Exec(ctx, registerWorker, arg.WorkerID, arg.Capacity)
 	return err
 }
 

internal/worker/worker.go

@@ -11,36 +11,94 @@ import (
 	"github.com/jackc/pgx/v5/pgxpool"
 )
 
+type workProcessor interface {
+	ProcessNext(ctx context.Context, worker_id uuid.UUID) error
+}
+
 //go:generate sqlc generate
 
 type Worker struct {
 	workerID uuid.UUID
 
-	db     *pgxpool.Pool
-	logger *slog.Logger
+	db            *pgxpool.Pool
+	workProcessor workProcessor
+	logger        *slog.Logger
+
+	capacity uint
 }
 
 func NewWorker(
 	db *pgxpool.Pool,
 	logger *slog.Logger,
+	workProcessor workProcessor,
 ) *Worker {
 	return &Worker{
-		workerID: uuid.New(),
-		db:       db,
-		logger:   logger,
+		workerID:      uuid.New(),
+		db:            db,
+		workProcessor: workProcessor,
+		logger:        logger,
+
+		capacity: 50,
 	}
 }
 
 func (w *Worker) Setup(ctx context.Context) error {
 	repo := repositories.New(w.db)
 
-	if err := repo.RegisterWorker(ctx, w.workerID); err != nil {
+	w.logger.InfoContext(ctx, "setting up worker", "worker_id", w.workerID)
+	if err := repo.RegisterWorker(
+		ctx,
+		&repositories.RegisterWorkerParams{
+			WorkerID: w.workerID,
+			Capacity: int32(w.capacity),
+		},
+	); err != nil {
 		return nil
 	}
 
 	return nil
 }
 
+type Workers struct {
+	Instances []WorkerInstance
+}
+
+func (w *Workers) Capacity() uint {
+	capacity := uint(0)
+
+	for _, worker := range w.Instances {
+		capacity += worker.Capacity
+	}
+
+	return capacity
+}
+
+type WorkerInstance struct {
+	WorkerID uuid.UUID
+	Capacity uint
+}
+
+func (w *Worker) GetWorkers(ctx context.Context) (*Workers, error) {
+	repo := repositories.New(w.db)
+
+	dbInstances, err := repo.GetWorkers(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("failed to find workers: %w", err)
+	}
+
+	instances := make([]WorkerInstance, 0, len(dbInstances))
+	for _, dbInstance := range dbInstances {
+		instances = append(instances, WorkerInstance{
+			WorkerID: dbInstance.WorkerID,
+			Capacity: uint(dbInstance.Capacity),
+		})
+	}
+
+	return &Workers{
+		Instances: instances,
+	}, nil
+}
+
 func (w *Worker) Start(ctx context.Context) error {
 	heartBeatCtx, heartBeatCancel := context.WithCancel(context.Background())
 	go func() {
@@ -69,10 +127,15 @@ func (w *Worker) Start(ctx context.Context) error {
 	}()
 
 	for {
-		if err := w.processWorkQueue(ctx); err != nil {
-			// FIXME: dead letter item, right now we just log and continue
+		select {
+		case <-ctx.Done():
+			return nil
+		default:
+			if err := w.processWorkQueue(ctx); err != nil {
+				// FIXME: dead letter item, right now we just log and continue
 
-			w.logger.WarnContext(ctx, "failed to handle work item", "error", err)
+				w.logger.WarnContext(ctx, "failed to handle work item", "error", err)
+			}
 		}
 	}
 }
@@ -84,8 +147,6 @@ func (w *Worker) updateHeartBeat(ctx context.Context) error {
 	return repo.UpdateWorkerHeartbeat(ctx, w.workerID)
 }
 
-func (w *Worker) processWorkQueue(_ context.Context) error {
-	time.Sleep(time.Second)
-
-	return nil
+func (w *Worker) processWorkQueue(ctx context.Context) error {
+	return w.workProcessor.ProcessNext(ctx, w.workerID)
 }