From eac99628bb8bb7243739b916beda6cdeef03d27f Mon Sep 17 00:00:00 2001 From: Matthew Meszaros Date: Tue, 2 Jun 2026 16:38:49 +0200 Subject: [PATCH] feat: enforce worker pool assignment safety Reserve dedicated worker allocation for the control plane, auto-promote spare capacity when needed, and keep risky or quarantined mailboxes off clean shared workers. --- docs/WORKER_ASSIGNMENT.md | 94 +++++++++ internal/api/handler/admin_provisioning.go | 28 +++ internal/api/handler/admin_workers_ssh.go | 7 + .../api/handler/internal_worker_config.go | 14 +- internal/app/worker/assignment.go | 166 ++++++++++++---- internal/app/worker/assignment_test.go | 182 +++++++++++++++++- internal/repository/pg_worker.go | 44 +++++ internal/repository/pg_worker_heartbeat.go | 19 +- internal/repository/pg_worker_risk.go | 65 +++++++ 9 files changed, 580 insertions(+), 39 deletions(-) create mode 100644 docs/WORKER_ASSIGNMENT.md diff --git a/docs/WORKER_ASSIGNMENT.md b/docs/WORKER_ASSIGNMENT.md new file mode 100644 index 00000000..c7036332 --- /dev/null +++ b/docs/WORKER_ASSIGNMENT.md @@ -0,0 +1,94 @@ +# Worker Assignment, Tiers, and Risk Pools + +How the control plane decides which worker a mailbox lands on. All of this is +backend/consumer logic — workers never make placement decisions and never +touch Postgres. The hot path is `internal/app/worker/assignment.go` +(`AssignWorkerToEmail`); the steady-state corrector is the hourly risk +rebalancer (`internal/app/consumer/risk_rebalancer.go`). + +## Tier model + +Three capacities, encoded on the `workers` row as `(worker_type, free_tier)`: + +| Tier | `worker_type` | `free_tier` | Who picks it | +| ---------------- | ------------- | ----------- | ----------------------------- | +| shared free | `shared` | `true` | free-trial orgs | +| shared premium | `shared` | `false` | paying orgs (no dedicated) | +| dedicated | `dedicated` | `false` | **control plane only** | + +Tier separation is strict: a free-tier mailbox never lands on a premium +worker and vice versa. + +### Dedicated is auto-allocated, never hand-picked + +Admins and customers only ever choose **free or premium**. Dedicated capacity +is created by the control plane: + +- `dedicated` is rejected server-side at every request boundary — the worker + heartbeat (`InternalWorkerHeartbeat`) and admin provisioning + template/job creation/update (`admin_provisioning.go`). The reject carries a + stable `tier_not_allowed` code. The single authority for this rule is + `repository.IsClientRequestableTier` (next to `tierToColumns`). A blank tier + is still allowed and maps to the shared-premium default, so existing workers + that don't set `WORKER_TIER` keep auto-registering. +- The admin provisioning UI only offers the two shared tiers + (`admin/.../ProvisioningTemplateForm.tsx`). + +When an org needs a dedicated worker — on a plan upgrade +(`MigrateOrgToDedicated`, fired from the Stripe webhook) or when a mailbox is +added for a dedicated-plan org that has none bound yet — `AssignDedicatedWorker`: + +1. tries `GetAvailableDedicatedWorker` (a pre-provisioned, unbound dedicated box); +2. if none is free, **promotes a spare** idle premium shared worker to + dedicated (`PromoteIdlePremiumWorkerToDedicated`: `worker_type = 'shared'`, + `free_tier = false`, `account_count = 0`, selected `FOR UPDATE SKIP LOCKED` + so concurrent promotions can't collide). Only idle workers are eligible, so + a promotion never strands existing mailboxes on a box that suddenly belongs + to one org; +3. binds it via `CreateDedicatedAssignmentIfNotExists`. If the bind race is + lost (the org was bound concurrently) and we had just promoted a worker, the + promotion is reverted back to `shared` so it isn't stranded as an unbound + dedicated box; +4. only if there's nothing to promote either does it surface + `ErrNoDedicatedWorkers` — in the hot path that degrades gracefully to shared + premium placement (the rebalancer / next onboarding retries). + +## Risk-band placement (health segregation) + +Shared workers are bucketed into risk pools (`workers.risk_pool`: +`clean` / `risky` / `quarantine`); mailboxes carry a matching +`email_accounts.risk_band` derived from warmup health by the rebalancer +(`RiskBandFromHealth`). The invariant is +`email.risk_band.MatchingRiskPool() == worker.risk_pool`. + +**Initial placement is strict.** `AssignWorkerToEmail` reads the mailbox's band +(`GetEmailAccountRiskBand`) and places via `selectSharedWorkerForBandWeight`: + +- **clean band** → the capacity-aware path (`selectSharedWorkerForWeight`), + unchanged: honours per-mailbox weight and worker headroom. A fresh mailbox is + `clean` (column default until the warmup sweep classifies it), so onboarding + takes this path. +- **risky / quarantine band** → placed **only** on a worker whose `risk_pool` + matches. If that pool has no worker, an idle clean worker is **promoted** into + the pool (`PromoteWorkerToPool`, idle-only + `FOR UPDATE SKIP LOCKED`) rather + than diluting the clean pool. If there's nothing to promote, placement + **refuses** (`ErrNoAvailableWorkers`) — a risky/quarantine inbox is never + co-located with trusted ones. Onboarding treats the refusal as non-fatal and + the rebalancer retries next tick. + +`SelectSharedWorkerForBand` (used by the rebalancer, which has no per-mailbox +weight) now delegates to the same strict logic with the default weight, so +**initial placement and rebalancing share identical rules** and cannot fight. +Existing risky mailboxes already sitting on a clean worker are left for the +rebalancer to migrate; `AssignWorkerToEmail` only governs new placement. + +Promotions (shared→dedicated, clean→risk-pool) are logged at info level for ops +visibility; the rebalancer additionally writes an admin audit log per mailbox +migration. + +## Backwards compatibility + +Installs that never enable risk pools leave every worker in `risk_pool = +'clean'` and every mailbox in `risk_band = 'clean'`, so placement always takes +the clean capacity-aware path — behaviour is unchanged. Dedicated workers carry +no risk-pool semantics (one customer per worker, no cross-tenant contamination). diff --git a/internal/api/handler/admin_provisioning.go b/internal/api/handler/admin_provisioning.go index 485f303f..0aa2353c 100644 --- a/internal/api/handler/admin_provisioning.go +++ b/internal/api/handler/admin_provisioning.go @@ -442,6 +442,13 @@ func (h *Handler) AdminCreateProvisioningTemplate(c *gin.Context) { c.JSON(http.StatusBadRequest, gin.H{"error": "name, provider, location, server_type and worker_tier are required"}) return } + if !repository.IsClientRequestableTier(t.Tier) { + c.JSON(http.StatusBadRequest, gin.H{ + "error": "dedicated tier templates cannot be created; dedicated workers are allocated automatically by the control plane", + "code": "tier_not_allowed", + }) + return + } if err := h.ProvisioningTemplateRepo.Create(c.Request.Context(), t); err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) return @@ -466,6 +473,13 @@ func (h *Handler) AdminUpdateProvisioningTemplate(c *gin.Context) { c.JSON(http.StatusBadRequest, gin.H{"error": "name, provider, location, server_type and worker_tier are required"}) return } + if !repository.IsClientRequestableTier(t.Tier) { + c.JSON(http.StatusBadRequest, gin.H{ + "error": "templates cannot be set to the dedicated tier; dedicated workers are allocated automatically by the control plane", + "code": "tier_not_allowed", + }) + return + } if err := h.ProvisioningTemplateRepo.Update(c.Request.Context(), t); err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) return @@ -626,6 +640,13 @@ func (h *Handler) AdminCreateProvisioningJob(c *gin.Context) { c.JSON(http.StatusNotFound, gin.H{"error": "template not found"}) return } + if !repository.IsClientRequestableTier(t.Tier) { + c.JSON(http.StatusBadRequest, gin.H{ + "error": "this template targets the dedicated tier, which is allocated automatically by the control plane; pick a shared-tier template", + "code": "tier_not_allowed", + }) + return + } // Snapshot the flat template into the job's config column. Its field // names line up with provisioning.JobConfig, so the state machine reads // it directly. @@ -651,6 +672,13 @@ func (h *Handler) AdminCreateProvisioningJob(c *gin.Context) { return } snap := fromTemplateDTO(&provTemplateDTO{Name: "custom", Config: cfgDTO}) + if !repository.IsClientRequestableTier(snap.Tier) { + c.JSON(http.StatusBadRequest, gin.H{ + "error": "dedicated tier cannot be provisioned directly; dedicated workers are allocated automatically by the control plane", + "code": "tier_not_allowed", + }) + return + } b, _ := json.Marshal(snap) config = b provider = cfgDTO.Provider diff --git a/internal/api/handler/admin_workers_ssh.go b/internal/api/handler/admin_workers_ssh.go index a4c45373..a46df85d 100644 --- a/internal/api/handler/admin_workers_ssh.go +++ b/internal/api/handler/admin_workers_ssh.go @@ -76,6 +76,13 @@ func (h *Handler) AdminCreateWorker(c *gin.Context) { errx.JSON(c, errx.New(errx.BadRequest, "invalid request body")) return } + // A worker is never born dedicated: dedicated capacity is allocated by the + // control plane (by promoting a spare shared worker). Admins create shared + // workers and may convert one later via AdminConvertWorkerToDedicated. + if !repository.IsClientRequestableTier(req.WorkerType) { + errx.JSON(c, errx.New(errx.BadRequest, "dedicated workers cannot be created directly; dedicated capacity is allocated automatically by the control plane — create a shared worker and convert it if needed")) + return + } if req.SSHPort == 0 { req.SSHPort = 22 } diff --git a/internal/api/handler/internal_worker_config.go b/internal/api/handler/internal_worker_config.go index 739acea4..8fd365e0 100644 --- a/internal/api/handler/internal_worker_config.go +++ b/internal/api/handler/internal_worker_config.go @@ -9,6 +9,7 @@ import ( "github.com/gin-gonic/gin" "github.com/google/uuid" "github.com/warmbly/warmbly/internal/infrastructure/kafka" + "github.com/warmbly/warmbly/internal/repository" ) // Internal worker bootstrap + config endpoints. A worker process starts with @@ -103,7 +104,7 @@ func (h *Handler) InternalWorkerConfig(c *gin.Context) { type HeartbeatPayload struct { WorkerID string `json:"worker_id"` BindIP string `json:"bind_ip"` - Tier string `json:"tier,omitempty"` // shared_free | shared_premium | dedicated + Tier string `json:"tier,omitempty"` // shared_free | shared_premium (dedicated is rejected; allocated by the control plane) EgressKind string `json:"egress_kind,omitempty"` // cold_smtp | oauth_api | warmup_only } @@ -127,6 +128,17 @@ func (h *Handler) InternalWorkerHeartbeat(c *gin.Context) { c.JSON(http.StatusBadRequest, gin.H{"error": "bind_ip required"}) return } + // A worker may only register as a shared tier. Dedicated capacity is + // allocated by the control plane (by promoting a spare shared worker), so + // a worker must never self-designate as dedicated. A blank tier is fine — + // it maps to the shared-premium default. + if !repository.IsClientRequestableTier(p.Tier) { + c.JSON(http.StatusBadRequest, gin.H{ + "error": "dedicated tier cannot be requested by a worker; dedicated capacity is allocated automatically by the control plane", + "code": "tier_not_allowed", + }) + return + } if h.WorkerRepo == nil { // Worker repository not wired (e.g. tests). Treat as 204 noop. c.Status(http.StatusNoContent) diff --git a/internal/app/worker/assignment.go b/internal/app/worker/assignment.go index 7ed73de0..744c2ef8 100644 --- a/internal/app/worker/assignment.go +++ b/internal/app/worker/assignment.go @@ -7,6 +7,7 @@ import ( "time" "github.com/google/uuid" + "github.com/rs/zerolog/log" "github.com/warmbly/warmbly/internal/models" "github.com/warmbly/warmbly/internal/repository" ) @@ -32,10 +33,11 @@ type WorkerAssignmentService interface { // SelectSharedWorker selects the least loaded shared worker for the given tier SelectSharedWorker(ctx context.Context, freeTier bool) (*models.Worker, error) - // SelectSharedWorkerForBand selects the least-loaded shared worker whose - // risk_pool matches the mailbox's risk band. Falls back to the clean - // pool if no worker exists in the target pool, and to any worker of the - // tier as a last resort. + // SelectSharedWorkerForBand selects the shared worker whose risk_pool + // matches the mailbox's risk band. Strict: a risky/quarantine mailbox is + // never placed in the clean pool — if the matching pool is empty an idle + // clean worker is promoted into it, and if there's nothing to promote the + // call refuses rather than co-locating risky traffic with trusted inboxes. SelectSharedWorkerForBand(ctx context.Context, freeTier bool, band models.EmailRiskBand) (*models.Worker, error) // Dedicated worker management @@ -98,6 +100,26 @@ func (s *workerAssignmentService) AssignWorkerToEmail(ctx context.Context, email if err != nil { return nil, err } + // No dedicated worker bound yet — e.g. this mailbox was added + // before the subscription-upgrade migration ran, or that + // migration found no free dedicated box. Allocate one on demand: + // AssignDedicatedWorker promotes a spare premium shared worker to + // dedicated when the dedicated pool is empty. If there's nothing + // to promote (ErrNoDedicatedWorkers) we fall through to shared + // placement rather than failing the add — the rebalancer/next + // onboarding will retry. + if dedicatedWorker == nil { + aerr := s.AssignDedicatedWorker(ctx, orgID, sub.ID) + if aerr != nil && !errors.Is(aerr, ErrNoDedicatedWorkers) && !errors.Is(aerr, ErrOrgAlreadyAssigned) { + return nil, aerr + } + if aerr == nil || errors.Is(aerr, ErrOrgAlreadyAssigned) { + dedicatedWorker, err = s.workerRepo.GetDedicatedWorkerByUserID(ctx, orgID) + if err != nil { + return nil, err + } + } + } if dedicatedWorker != nil { // Assign to dedicated worker if err := s.workerRepo.UpdateEmailAccountWorker(ctx, emailAccountID, dedicatedWorker.ID); err != nil { @@ -111,14 +133,33 @@ func (s *workerAssignmentService) AssignWorkerToEmail(ctx context.Context, email // worker), but keeping the column accurate makes the // capacity view useful for ops dashboards. _ = s.workerRepo.AddLoadScore(ctx, dedicatedWorker.ID, weight) + // Dedicated workers are paid-org only, so the mailbox belongs + // to the premium warmup pool. Set it explicitly to match the + // shared-premium path and MigrateOrgToDedicated — otherwise the + // account keeps the schema default 'free' and reporting/ + // filtering misclassifies it. Non-fatal. + if err := s.workerRepo.UpdateEmailAccountWarmupPoolType(ctx, emailAccountID, "premium"); err != nil { + // Log but don't fail + } return &dedicatedWorker.ID, nil } } } - // 5. Assign to shared worker (strict tier separation) - freeTier := !isPaidOrg // Free trial = free workers, Paid = premium workers - worker, err := s.selectSharedWorkerForWeight(ctx, freeTier, weight) + // 5. Assign to a shared worker, strict on BOTH axes: + // - tier separation: free trial → free workers, paid → premium workers + // - risk segregation: the mailbox's risk_band must match the worker's + // risk_pool, so a risky/quarantine inbox never lands on a clean + // worker next to trusted ones. A fresh mailbox is 'clean' (the + // column default until the warmup health sweep classifies it), so + // onboarding takes the capacity-aware clean path; only already + // degraded mailboxes hit the strict risky/quarantine branch. + freeTier := !isPaidOrg + band, err := s.workerRepo.GetEmailAccountRiskBand(ctx, emailAccountID) + if err != nil { + return nil, err + } + worker, err := s.selectSharedWorkerForBandWeight(ctx, freeTier, band, weight) if err != nil { return nil, err } @@ -289,21 +330,34 @@ func (s *workerAssignmentService) selectSharedWorkerLegacy(ctx context.Context, return &workers[0], nil } -// SelectSharedWorkerForBand picks the least-loaded shared worker whose -// risk_pool matches band.MatchingRiskPool(). Fallback chain: -// -// 1. Worker in the matching pool of the right tier -// 2. Worker in the clean pool of the right tier (better to land risky -// mailboxes on clean workers than to refuse, but log + audit this -// since it dilutes the clean pool — operator should provision a -// risky/quarantine worker) -// 3. Any worker of the right tier (existing SelectSharedWorker behavior) -// -// Step 3 maintains backwards compatibility with installations that don't -// run risk pools yet — they leave everything in risk_pool='clean' and -// behavior is unchanged. +// SelectSharedWorkerForBand picks the shared worker that should host a mailbox +// of the given risk band. It is strict: a risky/quarantine mailbox is NEVER +// placed in the clean pool. Used by the risk rebalancer (which has no per- +// mailbox weight) — it delegates to selectSharedWorkerForBandWeight with the +// default weight so initial placement and rebalancing share identical logic +// and can't fight each other. func (s *workerAssignmentService) SelectSharedWorkerForBand(ctx context.Context, freeTier bool, band models.EmailRiskBand) (*models.Worker, error) { + return s.selectSharedWorkerForBandWeight(ctx, freeTier, band, defaultMailboxWeight) +} + +// selectSharedWorkerForBandWeight is the strict, capacity-aware band selector +// shared by initial placement and the rebalancer. +// +// - clean band: route through the capacity-aware path +// (selectSharedWorkerForWeight), which honours per-mailbox weight and +// worker headroom. This is the unchanged behaviour for the common case +// and for installs that never enable risk pools (everything stays clean). +// - risky / quarantine band: place ONLY on a worker whose risk_pool matches. +// If the matching pool is empty, promote an idle clean worker into it +// rather than diluting the clean pool. If there's nothing to promote, +// refuse (ErrNoAvailableWorkers) — never co-locate a risky/quarantine +// inbox with trusted ones. The caller (onboarding) treats this as +// non-fatal and the rebalancer retries on the next tick. +func (s *workerAssignmentService) selectSharedWorkerForBandWeight(ctx context.Context, freeTier bool, band models.EmailRiskBand, weight float64) (*models.Worker, error) { target := band.MatchingRiskPool() + if target == models.WorkerRiskPoolClean { + return s.selectSharedWorkerForWeight(ctx, freeTier, weight) + } workers, err := s.workerRepo.GetSharedWorkersByTierAndPool(ctx, freeTier, target) if err != nil { @@ -313,23 +367,34 @@ func (s *workerAssignmentService) SelectSharedWorkerForBand(ctx context.Context, return &workers[0], nil } - // Step 2: fall back to the clean pool. Only kicks in for risky/quarantine - // bands when no matching-pool worker exists. - if target != models.WorkerRiskPoolClean { - workers, err = s.workerRepo.GetSharedWorkersByTierAndPool(ctx, freeTier, models.WorkerRiskPoolClean) - if err != nil { - return nil, err - } - if len(workers) > 0 { - return &workers[0], nil - } + // No worker in the matching pool. Promote an idle clean worker into it so + // we keep risky/quarantine traffic strictly segregated from trusted + // inboxes instead of falling back onto the clean pool. + promoted, err := s.workerRepo.PromoteWorkerToPool(ctx, freeTier, target) + if err != nil { + return nil, err + } + if promoted != nil { + log.Info(). + Str("worker_id", promoted.ID.String()). + Bool("free_tier", freeTier). + Str("risk_pool", string(target)). + Msg("assignment: promoted idle clean worker into risk pool") + return promoted, nil } - // Step 3: last-resort, any tier worker. Same as legacy SelectSharedWorker. - return s.SelectSharedWorker(ctx, freeTier) + // Nothing to promote. Refuse rather than dilute the clean pool. + return nil, ErrNoAvailableWorkers } -// AssignDedicatedWorker assigns a dedicated worker to an organization +// AssignDedicatedWorker assigns a dedicated worker to an organization. +// +// Dedicated capacity is allocated automatically: admins/customers only ever +// pick free or premium, and the control plane creates dedicated workers as +// needed. If the dedicated pool has no free worker, we promote a spare idle +// premium shared worker to dedicated (the same SetWorkerType + bind sequence +// the admin "convert to dedicated" action uses). Only when there's nothing to +// promote do we surface ErrNoDedicatedWorkers. func (s *workerAssignmentService) AssignDedicatedWorker(ctx context.Context, orgID, subscriptionID uuid.UUID) error { // Use atomic insert with conflict check to prevent race conditions. // Two concurrent requests could both pass the "check if exists" step and @@ -338,8 +403,24 @@ func (s *workerAssignmentService) AssignDedicatedWorker(ctx context.Context, org if err != nil { return err } + + // promoted tracks whether we flipped a shared worker to dedicated in this + // call, so we can undo it if we then lose the bind race below. + promoted := false if worker == nil { - return ErrNoDedicatedWorkers + spare, perr := s.workerRepo.PromoteIdlePremiumWorkerToDedicated(ctx) + if perr != nil { + return perr + } + if spare == nil { + return ErrNoDedicatedWorkers + } + log.Info(). + Str("worker_id", spare.ID.String()). + Str("org_id", orgID.String()). + Msg("assignment: promoted idle premium shared worker to dedicated") + worker = spare + promoted = true } assignment := &models.DedicatedWorkerAssignment{ @@ -355,6 +436,23 @@ func (s *workerAssignmentService) AssignDedicatedWorker(ctx context.Context, org return err } if !created { + // Lost the bind race: the org already has a dedicated worker. If we + // promoted a worker just now, revert it to the shared pool so it + // isn't stranded as an unbound dedicated box. A pre-existing + // dedicated worker (promoted == false) is left untouched. + if promoted { + if rerr := s.workerRepo.SetWorkerType(ctx, worker.ID, models.WorkerTypeShared); rerr != nil { + // The promotion couldn't be undone: the worker stays marked + // dedicated with no binding, so GetAvailableDedicatedWorker + // will keep re-selecting it. Log loudly so ops can reconcile. + // Still return ErrOrgAlreadyAssigned — the org IS bound (by the + // race winner), so failing the assignment here would be wrong. + log.Error().Err(rerr). + Str("worker_id", worker.ID.String()). + Str("org_id", orgID.String()). + Msg("assignment: failed to revert promoted worker to shared after losing dedicated bind race; worker stranded as dedicated") + } + } return ErrOrgAlreadyAssigned } return nil diff --git a/internal/app/worker/assignment_test.go b/internal/app/worker/assignment_test.go index 7388a74a..e83dc534 100644 --- a/internal/app/worker/assignment_test.go +++ b/internal/app/worker/assignment_test.go @@ -14,6 +14,7 @@ package worker import ( "context" + "errors" "testing" "github.com/google/uuid" @@ -37,6 +38,17 @@ type stubWorkerRepo struct { lastEmailPoolTypeSet string incrementedWorkerCounts map[uuid.UUID]int loadScoreDeltas map[uuid.UUID]float64 + + // Dedicated auto-promotion knobs. + availableDedicated *models.Worker // GetAvailableDedicatedWorker result + promotableDedicated *models.Worker // PromoteIdlePremiumWorkerToDedicated result + dedicatedAssignCreated bool // CreateDedicatedAssignmentIfNotExists result + setWorkerTypeCalls map[uuid.UUID]models.WorkerType // records SetWorkerType calls + + // Risk-band placement knobs. + riskBand models.EmailRiskBand // GetEmailAccountRiskBand result ("" → clean) + sharedByPool map[models.WorkerRiskPool][]models.Worker // GetSharedWorkersByTierAndPool result + promotedToPool *models.Worker // PromoteWorkerToPool result } func (r *stubWorkerRepo) GetDedicatedWorkerByUserID(_ context.Context, _ uuid.UUID) (*models.Worker, error) { @@ -95,6 +107,50 @@ func (r *stubWorkerRepo) AddLoadScore(_ context.Context, workerID uuid.UUID, del return nil } +func (r *stubWorkerRepo) GetEmailAccountRiskBand(_ context.Context, _ uuid.UUID) (models.EmailRiskBand, error) { + if r.riskBand == "" { + return models.EmailRiskBandClean, nil + } + return r.riskBand, nil +} + +func (r *stubWorkerRepo) GetAvailableDedicatedWorker(_ context.Context) (*models.Worker, error) { + return r.availableDedicated, nil +} + +func (r *stubWorkerRepo) PromoteIdlePremiumWorkerToDedicated(_ context.Context) (*models.Worker, error) { + if r.promotableDedicated == nil { + return nil, nil + } + w := *r.promotableDedicated + w.WorkerType = models.WorkerTypeDedicated + return &w, nil +} + +func (r *stubWorkerRepo) CreateDedicatedAssignmentIfNotExists(_ context.Context, a *models.DedicatedWorkerAssignment) (bool, error) { + if r.dedicatedAssignCreated { + // Bind it so the post-assign re-fetch in AssignWorkerToEmail finds it. + r.dedicatedForOrg = &models.Worker{ID: a.WorkerID, WorkerType: models.WorkerTypeDedicated, Active: true} + } + return r.dedicatedAssignCreated, nil +} + +func (r *stubWorkerRepo) SetWorkerType(_ context.Context, id uuid.UUID, t models.WorkerType) error { + if r.setWorkerTypeCalls == nil { + r.setWorkerTypeCalls = map[uuid.UUID]models.WorkerType{} + } + r.setWorkerTypeCalls[id] = t + return nil +} + +func (r *stubWorkerRepo) GetSharedWorkersByTierAndPool(_ context.Context, _ bool, pool models.WorkerRiskPool) ([]models.Worker, error) { + return r.sharedByPool[pool], nil +} + +func (r *stubWorkerRepo) PromoteWorkerToPool(_ context.Context, _ bool, _ models.WorkerRiskPool) (*models.Worker, error) { + return r.promotedToPool, nil +} + type stubSubRepo struct { repository.SubscriptionRepository sub *models.Subscription @@ -189,10 +245,15 @@ func TestAssign_PaidOrgWithDedicatedPlan_LandsOnDedicatedWorker(t *testing.T) { } func TestAssign_PaidOrgWithDedicatedPlanButNoAssignment_FallsBackToPremium(t *testing.T) { + // Dedicated plan, no bound worker, no free dedicated worker, AND nothing + // to promote (no idle premium spare). The add must still succeed by + // falling back to a premium shared worker rather than failing. premium := newWorker(uuid.New(), false, models.WorkerTypeShared) wr := &stubWorkerRepo{ - dedicatedForOrg: nil, // org has the plan but no worker assigned yet - sharedPremium: []models.Worker{premium}, + dedicatedForOrg: nil, // org has the plan but no worker assigned yet + availableDedicated: nil, // dedicated pool is empty + promotableDedicated: nil, // and there's no idle premium spare to promote + sharedPremium: []models.Worker{premium}, } sub := paidSub() plan := &models.Plan{ID: sub.PlanID, DedicatedWorkers: 1} @@ -203,7 +264,122 @@ func TestAssign_PaidOrgWithDedicatedPlanButNoAssignment_FallsBackToPremium(t *te t.Fatalf("AssignWorkerToEmail: %v", err) } if *got != premium.ID { - t.Errorf("org with dedicated plan but no assignment should fall back to premium, got %s", got) + t.Errorf("org with dedicated plan but nothing to promote should fall back to premium, got %s", got) + } +} + +func TestAssign_PaidOrgWithDedicatedPlanButNoAssignment_PromotesSpareWorker(t *testing.T) { + // Dedicated plan, no bound worker, dedicated pool empty — but an idle + // premium shared worker is available to promote. The mailbox must land on + // the promoted (now dedicated) worker, not on the shared pool. + spare := newWorker(uuid.New(), false, models.WorkerTypeShared) + wr := &stubWorkerRepo{ + dedicatedForOrg: nil, + availableDedicated: nil, + promotableDedicated: &spare, + dedicatedAssignCreated: true, // bind succeeds + sharedPremium: []models.Worker{newWorker(uuid.New(), false, models.WorkerTypeShared)}, + } + sub := paidSub() + plan := &models.Plan{ID: sub.PlanID, DedicatedWorkers: 1} + svc := NewAssignmentService(wr, &stubSubRepo{sub: sub}, &stubPlanRepo{plan: plan}) + + got, err := svc.AssignWorkerToEmail(context.Background(), uuid.New(), uuid.New()) + if err != nil { + t.Fatalf("AssignWorkerToEmail: %v", err) + } + if *got != spare.ID { + t.Errorf("org with dedicated plan should land on the promoted spare worker %s, got %s", spare.ID, got) + } +} + +func TestAssignDedicatedWorker_PromoteThenLoseBind_RevertsToShared(t *testing.T) { + // Promotion succeeds but the bind race is lost (a concurrent add bound the + // org first). The just-promoted worker must be reverted to shared so it + // isn't stranded as an unbound dedicated box. + spare := newWorker(uuid.New(), false, models.WorkerTypeShared) + wr := &stubWorkerRepo{ + availableDedicated: nil, // dedicated pool empty → promote + promotableDedicated: &spare, + dedicatedAssignCreated: false, // lose the bind race + } + svc := NewAssignmentService(wr, &stubSubRepo{}, &stubPlanRepo{}) + + err := svc.AssignDedicatedWorker(context.Background(), uuid.New(), uuid.New()) + if !errors.Is(err, ErrOrgAlreadyAssigned) { + t.Fatalf("expected ErrOrgAlreadyAssigned on lost bind race, got %v", err) + } + if got, ok := wr.setWorkerTypeCalls[spare.ID]; !ok || got != models.WorkerTypeShared { + t.Errorf("promoted worker must be reverted to shared after losing the bind race, got %v (called=%v)", got, ok) + } +} + +func TestAssign_RiskyMailbox_LandsOnRiskyPoolWorker(t *testing.T) { + // A risky mailbox must land on a worker in the risky pool, never on a + // clean-pool worker, so it can't damage the reputation of trusted inboxes. + riskyWorker := newWorker(uuid.New(), false, models.WorkerTypeShared) + riskyWorker.RiskPool = models.WorkerRiskPoolRisky + cleanWorker := newWorker(uuid.New(), false, models.WorkerTypeShared) + + wr := &stubWorkerRepo{ + riskBand: models.EmailRiskBandRisky, + sharedByPool: map[models.WorkerRiskPool][]models.Worker{ + models.WorkerRiskPoolRisky: {riskyWorker}, + models.WorkerRiskPoolClean: {cleanWorker}, + }, + } + sub := paidSub() + plan := &models.Plan{ID: sub.PlanID, DedicatedWorkers: 0} + svc := NewAssignmentService(wr, &stubSubRepo{sub: sub}, &stubPlanRepo{plan: plan}) + + got, err := svc.AssignWorkerToEmail(context.Background(), uuid.New(), uuid.New()) + if err != nil { + t.Fatalf("AssignWorkerToEmail: %v", err) + } + if *got != riskyWorker.ID { + t.Errorf("risky mailbox must land on a risky-pool worker %s, got %s", riskyWorker.ID, got) + } +} + +func TestAssign_RiskyMailbox_NoRiskyWorker_PromotesCleanWorker(t *testing.T) { + // When no risky-pool worker exists, we promote an idle clean worker into + // the risky pool rather than co-locating with trusted inboxes. + promoted := newWorker(uuid.New(), false, models.WorkerTypeShared) + promoted.RiskPool = models.WorkerRiskPoolRisky + + wr := &stubWorkerRepo{ + riskBand: models.EmailRiskBandRisky, + sharedByPool: map[models.WorkerRiskPool][]models.Worker{}, // risky pool empty + promotedToPool: &promoted, + } + sub := paidSub() + plan := &models.Plan{ID: sub.PlanID, DedicatedWorkers: 0} + svc := NewAssignmentService(wr, &stubSubRepo{sub: sub}, &stubPlanRepo{plan: plan}) + + got, err := svc.AssignWorkerToEmail(context.Background(), uuid.New(), uuid.New()) + if err != nil { + t.Fatalf("AssignWorkerToEmail: %v", err) + } + if *got != promoted.ID { + t.Errorf("risky mailbox with empty pool should land on the promoted worker %s, got %s", promoted.ID, got) + } +} + +func TestAssign_RiskyMailbox_NoWorkerNoPromotion_Refuses(t *testing.T) { + // Strict invariant: if there's no risky-pool worker and nothing to + // promote, refuse rather than place a risky inbox next to trusted ones. + wr := &stubWorkerRepo{ + riskBand: models.EmailRiskBandQuarantine, + sharedByPool: map[models.WorkerRiskPool][]models.Worker{}, + promotedToPool: nil, + } + sub := paidSub() + plan := &models.Plan{ID: sub.PlanID, DedicatedWorkers: 0} + svc := NewAssignmentService(wr, &stubSubRepo{sub: sub}, &stubPlanRepo{plan: plan}) + + _, err := svc.AssignWorkerToEmail(context.Background(), uuid.New(), uuid.New()) + if !errors.Is(err, ErrNoAvailableWorkers) { + t.Errorf("strict placement should refuse with ErrNoAvailableWorkers, got %v", err) } } diff --git a/internal/repository/pg_worker.go b/internal/repository/pg_worker.go index 478194d4..8fabe5b3 100644 --- a/internal/repository/pg_worker.go +++ b/internal/repository/pg_worker.go @@ -34,6 +34,7 @@ type WorkerRepository interface { GetSharedWorkersByTier(ctx context.Context, freeTier bool) ([]models.Worker, error) GetAllActiveWorkers(ctx context.Context) ([]models.Worker, error) GetAvailableDedicatedWorker(ctx context.Context) (*models.Worker, error) + PromoteIdlePremiumWorkerToDedicated(ctx context.Context) (*models.Worker, error) IncrementAccountCount(ctx context.Context, workerID uuid.UUID) error DecrementAccountCount(ctx context.Context, workerID uuid.UUID) error SetWorkerType(ctx context.Context, workerID uuid.UUID, workerType models.WorkerType) error @@ -74,7 +75,9 @@ type WorkerRepository interface { // Threat-level segregation SetWorkerRiskPool(ctx context.Context, workerID uuid.UUID, pool models.WorkerRiskPool) error SetEmailAccountRiskBand(ctx context.Context, emailAccountID uuid.UUID, band models.EmailRiskBand) error + GetEmailAccountRiskBand(ctx context.Context, emailAccountID uuid.UUID) (models.EmailRiskBand, error) GetSharedWorkersByTierAndPool(ctx context.Context, freeTier bool, pool models.WorkerRiskPool) ([]models.Worker, error) + PromoteWorkerToPool(ctx context.Context, freeTier bool, target models.WorkerRiskPool) (*models.Worker, error) ListRiskCandidates(ctx context.Context, limit int) ([]RiskCandidate, error) // Tags @@ -217,6 +220,47 @@ func (r *workerRepository) GetAvailableDedicatedWorker(ctx context.Context) (*mo return &w, nil } +// PromoteIdlePremiumWorkerToDedicated flips the oldest idle premium shared +// worker to dedicated and returns it, so the control plane can seed dedicated +// capacity on demand when GetAvailableDedicatedWorker finds none free. Only +// premium (free_tier = false), shared, active workers with account_count = 0 +// are eligible — promoting a loaded worker would strand its mailboxes on a box +// that suddenly belongs to one org. Returns nil if no idle premium worker +// exists. +// +// The candidate is selected and flipped in a single statement +// (UPDATE ... WHERE id = (SELECT ... FOR UPDATE SKIP LOCKED)), so two +// concurrent promotions never pick the same row. free/premium capacity +// separation is preserved: free-tier workers are never promoted. +func (r *workerRepository) PromoteIdlePremiumWorkerToDedicated(ctx context.Context) (*models.Worker, error) { + query := ` + UPDATE workers SET worker_type = 'dedicated', updated_at = NOW() + WHERE id = ( + SELECT id FROM workers + WHERE worker_type = 'shared' + AND active = true + AND free_tier = false + AND account_count = 0 + ORDER BY created_at ASC + LIMIT 1 + FOR UPDATE SKIP LOCKED + ) + RETURNING id, ip_addr, active, free_tier, worker_type, account_count, risk_pool, created_at, updated_at + ` + var w models.Worker + err := r.db.QueryRow(ctx, query).Scan( + &w.ID, &w.IPAddr, &w.Active, &w.FreeTier, &w.WorkerType, &w.AccountCount, &w.RiskPool, + &w.CreatedAt, &w.UpdatedAt, + ) + if err == pgx.ErrNoRows { + return nil, nil + } + if err != nil { + return nil, err + } + return &w, nil +} + func (r *workerRepository) IncrementAccountCount(ctx context.Context, workerID uuid.UUID) error { query := `UPDATE workers SET account_count = account_count + 1, updated_at = NOW() WHERE id = $1` _, err := r.db.Exec(ctx, query, workerID) diff --git a/internal/repository/pg_worker_heartbeat.go b/internal/repository/pg_worker_heartbeat.go index 889be2d0..73fc5959 100644 --- a/internal/repository/pg_worker_heartbeat.go +++ b/internal/repository/pg_worker_heartbeat.go @@ -38,9 +38,26 @@ func (r *workerRepository) UpsertOnHeartbeat(ctx context.Context, id uuid.UUID, return err } +// TierDedicated is the reserved tier the control plane allocates on its own +// (see worker.AssignDedicatedWorker, which promotes a spare shared worker to +// dedicated on demand). It must never be requested by a worker heartbeat or +// an admin provisioning template/job — clients only pick a shared tier. +const TierDedicated = "dedicated" + +// IsClientRequestableTier reports whether a tier string may be supplied by a +// client: a worker heartbeat or an admin provisioning request. Everything +// except the reserved "dedicated" tier is allowed; a blank tier is treated as +// the shared-premium default by tierToColumns, so existing workers that don't +// set WORKER_TIER keep auto-registering normally. +func IsClientRequestableTier(tier string) bool { + return tier != TierDedicated +} + // tierToColumns converts the higher-level tier name used by templates and // the heartbeat API into the (worker_type, free_tier) tuple stored on the -// workers row. +// workers row. The "dedicated" case stays for defensive completeness, but the +// request paths now reject that tier before it reaches here (dedicated workers +// are created by promoting a shared worker, not by self-designation). func tierToColumns(tier string) (workerType string, freeTier bool) { switch tier { case "dedicated": diff --git a/internal/repository/pg_worker_risk.go b/internal/repository/pg_worker_risk.go index a83b20e1..732109c6 100644 --- a/internal/repository/pg_worker_risk.go +++ b/internal/repository/pg_worker_risk.go @@ -2,9 +2,11 @@ package repository import ( "context" + "errors" "time" "github.com/google/uuid" + "github.com/jackc/pgx/v5" "github.com/warmbly/warmbly/internal/models" ) @@ -30,6 +32,69 @@ func (r *workerRepository) SetEmailAccountRiskBand(ctx context.Context, emailAcc return err } +// GetEmailAccountRiskBand returns the stored risk_band for a mailbox so the +// assignment service can place it on a worker whose risk_pool matches. A +// missing row, or a NULL band (shouldn't happen — the column is NOT NULL with +// a 'clean' default), is treated as clean: assume innocent until the warmup +// health sweep proves otherwise. +func (r *workerRepository) GetEmailAccountRiskBand(ctx context.Context, emailAccountID uuid.UUID) (models.EmailRiskBand, error) { + var band models.EmailRiskBand + err := r.db.QueryRow(ctx, ` + SELECT COALESCE(risk_band, 'clean'::email_risk_band) + FROM email_accounts + WHERE id = $1 + `, emailAccountID).Scan(&band) + if errors.Is(err, pgx.ErrNoRows) { + return models.EmailRiskBandClean, nil + } + if err != nil { + return models.EmailRiskBandClean, err + } + return band, nil +} + +// PromoteWorkerToPool converts the oldest idle clean-pool shared worker of the +// given tier into the target risk pool and returns it, so risky/quarantine +// mailboxes never have to share a worker with clean ones. Only idle workers +// (account_count = 0) are eligible — relabelling a worker that already holds +// clean mailboxes would itself co-locate. Returns nil if no eligible worker +// exists, in which case the caller should refuse placement rather than dilute +// the clean pool. +// +// Select-and-convert happen in one statement +// (UPDATE ... WHERE id = (SELECT ... FOR UPDATE SKIP LOCKED)), so concurrent +// callers never grab the same row. tier separation is preserved: free_tier is +// matched, so a free-tier mailbox only ever promotes a free-tier worker. +func (r *workerRepository) PromoteWorkerToPool(ctx context.Context, freeTier bool, target models.WorkerRiskPool) (*models.Worker, error) { + var w models.Worker + err := r.db.QueryRow(ctx, ` + UPDATE workers SET risk_pool = $2, updated_at = NOW() + WHERE id = ( + SELECT id FROM workers + WHERE worker_type = 'shared' + AND active = true + AND free_tier = $1 + AND risk_pool = 'clean' + AND account_count = 0 + ORDER BY created_at ASC + LIMIT 1 + FOR UPDATE SKIP LOCKED + ) + RETURNING id, ip_addr, active, free_tier, worker_type, account_count, created_at, updated_at + `, freeTier, target).Scan( + &w.ID, &w.IPAddr, &w.Active, &w.FreeTier, &w.WorkerType, &w.AccountCount, + &w.CreatedAt, &w.UpdatedAt, + ) + if errors.Is(err, pgx.ErrNoRows) { + return nil, nil + } + if err != nil { + return nil, err + } + w.RiskPool = target // just set it; the row is now in the target pool + return &w, nil +} + // GetSharedWorkersByTierAndPool is the assignment service's primary lookup: // give me the least-loaded shared worker for this tier AND risk pool. // Falls back to "any pool" via separate caller logic when no match exists.