From 8b83062effbe2333a3ffa9b88eadd25a43679149 Mon Sep 17 00:00:00 2001 From: Matthew Meszaros Date: Wed, 9 Sep 2026 05:21:17 -0700 Subject: [PATCH] feat: fix ten defects the review found, the worst being that every seeder still wrote the worker columns migration 000141 moved to fleet_nodes so make dev could not seed, that the generated systemd unit used a command substitution systemd never expands so a joined machine restart-looped while the script reported success, that the state directory was not bind-mounted so the node wrote its update target inside the container where the host timer never saw it, and that placement stopped assigning warmup pool membership so paid mailboxes silently warmed in the free pool --- admin/src/app/dashboard/WorkerDetailPage.tsx | 8 ++-- .../fleet/ConvertDedicatedDialog.tsx | 6 +-- admin/src/lib/api/client/admin/fleet.ts | 16 ++++--- admin/src/lib/api/client/admin/workers.ts | 16 ++++--- cmd/seed/main.go | 18 +++++--- internal/api/handler/nodescript/join.sh | 12 ++++- internal/app/fleet/rotation_live_test.go | 14 ++++++ internal/app/worker/assignment.go | 45 +++++++++++++++++-- internal/app/worker/rotation.go | 7 ++- internal/app/worker/rotation_test.go | 21 +++++++-- internal/repository/pg_worker.go | 3 +- internal/sandbox/seed.go | 13 ++++-- internal/seed/workers.go | 22 ++++++--- 13 files changed, 156 insertions(+), 45 deletions(-) diff --git a/admin/src/app/dashboard/WorkerDetailPage.tsx b/admin/src/app/dashboard/WorkerDetailPage.tsx index cfe1260c5..b84f784a8 100644 --- a/admin/src/app/dashboard/WorkerDetailPage.tsx +++ b/admin/src/app/dashboard/WorkerDetailPage.tsx @@ -261,9 +261,11 @@ export default function WorkerDetailPage() { {statsQ.data && (
{statsQ.data.emails_sent_today} - {statsQ.data.emails_sent_total} - {statsQ.data.active_campaigns} - {statsQ.data.warmup_emails} + {statsQ.data.emails_sent_this_week} + {statsQ.data.total_emails_sent} + + {`${Math.round(statsQ.data.success_rate * 100)}%`} +
)} diff --git a/admin/src/app/dashboard/fleet/ConvertDedicatedDialog.tsx b/admin/src/app/dashboard/fleet/ConvertDedicatedDialog.tsx index a2841b556..c4903de8f 100644 --- a/admin/src/app/dashboard/fleet/ConvertDedicatedDialog.tsx +++ b/admin/src/app/dashboard/fleet/ConvertDedicatedDialog.tsx @@ -73,9 +73,9 @@ export function ConvertDedicatedDialog({ }), onSuccess: (res) => { toast.success( - res.new_assignment - ? `Worker is now dedicated to ${org?.name}${res.accounts_drained ? ` (${res.accounts_drained} mailboxes drained)` : ""}` - : "Binding already existed; worker type set to dedicated", + res.new_reservation + ? `Worker reserved for ${org?.name}. Other tenants drift off it on the rotation loop.` + : "That workspace already had this worker reserved.", ); qc.invalidateQueries({ queryKey: ["admin", "workers"] }); qc.invalidateQueries({ queryKey: ["admin", "fleet"] }); diff --git a/admin/src/lib/api/client/admin/fleet.ts b/admin/src/lib/api/client/admin/fleet.ts index c48a76ce4..c182c59ac 100644 --- a/admin/src/lib/api/client/admin/fleet.ts +++ b/admin/src/lib/api/client/admin/fleet.ts @@ -67,10 +67,12 @@ export interface AdminConvertDedicatedRequest { drain_to_worker_id?: string | null; } -export interface AdminConvertDedicatedResponse { +// Mirrors the /reserve handler. Reserving no longer drains anything: the +// rotation loop moves other tenants off on its own schedule. +export interface AdminReserveWorkerResponse { ok: boolean; - accounts_drained: number; - new_assignment: boolean; + worker_id: string; + new_reservation: boolean; } export interface AdminReleaseDedicatedResponse { @@ -127,15 +129,15 @@ export function releaseDedicatedWorker(orgId: string): Promise { +): Promise { return Request({ method: "POST", - url: `/admin/workers/${workerId}/convert-dedicated`, + url: `/admin/workers/${workerId}/reserve`, authorization: true, data: body, }); diff --git a/admin/src/lib/api/client/admin/workers.ts b/admin/src/lib/api/client/admin/workers.ts index e4c1205e2..8d71c1c8a 100644 --- a/admin/src/lib/api/client/admin/workers.ts +++ b/admin/src/lib/api/client/admin/workers.ts @@ -22,12 +22,16 @@ export function getWorkerEmails( }); } +// Mirrors models.WorkerStats. Renaming a field here does not rename it on the +// wire; it just renders blank. export interface WorkerStats { + worker_id: string; + total_emails_sent: number; emails_sent_today: number; - emails_sent_total: number; - active_campaigns: number; - connected_emails: number; - warmup_emails: number; + emails_sent_this_week: number; + average_delivery_time_ms: number; + success_rate: number; + queue_depth: number; } export function getWorkerStats(id: string): Promise { @@ -48,7 +52,9 @@ export function reassignWorkerEmails( return Request({ method: "POST", url: `/admin/workers/${targetWorkerId}/reassign`, - data: { email_ids: emailIds }, + // The target is in the URL and in the body: models.ReassignEmailsRequest + // marks new_worker_id required, so omitting it is a 400. + data: { email_ids: emailIds, new_worker_id: targetWorkerId }, authorization: true, }); } diff --git a/cmd/seed/main.go b/cmd/seed/main.go index c6b29be5e..791cee269 100644 --- a/cmd/seed/main.go +++ b/cmd/seed/main.go @@ -406,16 +406,22 @@ func upsertOrg(ctx context.Context, pool *pgxpool.Pool, id uuid.UUID, name, slug } func upsertWorker(ctx context.Context, pool *pgxpool.Pool, id uuid.UUID, name, ip, region string, active bool) error { - _, err := pool.Exec(ctx, ` - INSERT INTO workers (id, name, ip_addr, region, active) - VALUES ($1, $2, $3, $4, $5) + if _, err := pool.Exec(ctx, ` + INSERT INTO fleet_nodes (id, role, name, address, region, active, last_seen_at) + VALUES ($1, 'worker', $2, $3, $4, $5, now()) ON CONFLICT (id) DO UPDATE SET name = EXCLUDED.name, - ip_addr = EXCLUDED.ip_addr, + address = EXCLUDED.address, region = EXCLUDED.region, active = EXCLUDED.active, - updated_at = NOW()`, - id, name, ip, region, active) + last_seen_at = now(), + updated_at = now()`, + id, name, ip, region, active); err != nil { + return err + } + _, err := pool.Exec(ctx, ` + INSERT INTO workers (id, account_count) VALUES ($1, 0) + ON CONFLICT (id) DO NOTHING`, id) return err } diff --git a/internal/api/handler/nodescript/join.sh b/internal/api/handler/nodescript/join.sh index ff11ddba9..222d6d08d 100755 --- a/internal/api/handler/nodescript/join.sh +++ b/internal/api/handler/nodescript/join.sh @@ -185,6 +185,10 @@ write_config() { printf '%s\n' "$DESIRED_VERSION" > "$STATE_DIR/target-version" printf '%s\n' "$WARMBLY_IMAGE_REPO/$WARMBLY_ROLE" > "$STATE_DIR/image" + # systemd performs no command substitution, so the image reference has to + # reach the unit as an environment variable it can expand itself. + printf 'WARMBLY_IMAGE_REF=%s/%s:%s\n' \ + "$WARMBLY_IMAGE_REPO" "$WARMBLY_ROLE" "$DESIRED_VERSION" > "$STATE_DIR/image-ref" log "Wrote $CONFIG_DIR/node.env" } @@ -201,6 +205,10 @@ Requires=docker.service [Service] Restart=always RestartSec=5 +# systemd does not run a shell, so the image reference comes from a file it +# reads as environment rather than from a command substitution. \${VAR} expands +# to exactly one argument, which is what an image:tag needs. +EnvironmentFile=$STATE_DIR/image-ref # The container is replaced rather than reconfigured, so start always removes # any previous one first: a name collision after an unclean stop would # otherwise wedge the service in a restart loop. @@ -208,7 +216,8 @@ ExecStartPre=-/usr/bin/docker rm -f $service ExecStart=/usr/bin/docker run --rm --name $service \\ --env-file $CONFIG_DIR/node.env \\ --network host \\ - \$(cat $STATE_DIR/image):\$(cat $STATE_DIR/target-version) + -v $STATE_DIR:$STATE_DIR \\ + \${WARMBLY_IMAGE_REF} ExecStop=/usr/bin/docker stop $service [Install] @@ -246,6 +255,7 @@ if ! docker pull "$image:$target" >/dev/null 2>&1; then fi sed -i "s|^WARMBLY_VERSION=.*|WARMBLY_VERSION=$target|" "$CONFIG_DIR/node.env" +printf 'WARMBLY_IMAGE_REF=%s:%s\n' "$image" "$target" > "$STATE_DIR/image-ref" echo "warmbly-node-update: $current -> $target" systemctl restart "warmbly-$role" UPDATER diff --git a/internal/app/fleet/rotation_live_test.go b/internal/app/fleet/rotation_live_test.go index 99783146c..e5ec2b0b8 100644 --- a/internal/app/fleet/rotation_live_test.go +++ b/internal/app/fleet/rotation_live_test.go @@ -120,6 +120,20 @@ func TestLivePlacementAndRotation(t *testing.T) { } } + // 1b. Placement also settles warmup pool membership. It used to fall out of + // tier placement; with tiers gone it has to be set explicitly, and + // leaving it unset silently warms paying customers in the free pool. + for i, mb := range mailboxes { + var poolType *string + if err := pool.QueryRow(ctx, + `SELECT warmup_pool_type FROM email_accounts WHERE id = $1`, mb).Scan(&poolType); err != nil { + t.Fatalf("mailbox %d: read warmup pool: %v", i, err) + } + if poolType == nil || *poolType == "" { + t.Fatalf("mailbox %d: placement left warmup_pool_type unset", i) + } + } + // 2. Gather them onto this test's own worker, so the rotation assertions // below are about this test's fleet and not whatever else is running. placed := map[uuid.UUID]uuid.UUID{} diff --git a/internal/app/worker/assignment.go b/internal/app/worker/assignment.go index 628a7fc66..e9adcd4e6 100644 --- a/internal/app/worker/assignment.go +++ b/internal/app/worker/assignment.go @@ -7,6 +7,7 @@ import ( "github.com/google/uuid" "github.com/rs/zerolog/log" + "github.com/warmbly/warmbly/internal/config" "github.com/warmbly/warmbly/internal/models" "github.com/warmbly/warmbly/internal/repository" ) @@ -73,6 +74,10 @@ type PlacementLookup struct { // CurrentWorkerID is the incumbent, if any. Present means the caller is // considering a move and the incumbent should get its stickiness bonus. CurrentWorkerID *uuid.UUID + // ExcludeWorkerID drops one worker from consideration entirely. Set when + // draining: without it the drained worker is still the incumbent, wins on + // stickiness, and the drain is a silent no-op. + ExcludeWorkerID *uuid.UUID // Region is where the mailbox signs in from today. Supplied on a // re-placement so a move keeps the sign-in geography its provider has // already seen; empty on first placement, which scores neutral. @@ -137,6 +142,14 @@ func (s *workerAssignmentService) AssignWorkerToEmail(ctx context.Context, email if err := s.workerRepo.AddLoadScore(ctx, res.Worker.ID, s.resolveMailboxWeight(ctx, emailAccountID)); err != nil { log.Warn().Err(err).Str("worker_id", res.Worker.ID.String()).Msg("placement: load score bump failed") } + + // Warmup pool membership still follows the subscription. It used to be a + // by-product of tier placement; with tiers gone it has to be set + // explicitly, or every new mailbox keeps the column default and paid + // customers silently warm in the free pool. + if err := s.workerRepo.UpdateEmailAccountWarmupPoolType(ctx, emailAccountID, s.warmupPoolFor(ctx, orgID)); err != nil { + log.Warn().Err(err).Str("account_id", emailAccountID.String()).Msg("placement: warmup pool assignment failed") + } return &res.Worker.ID, nil } @@ -174,6 +187,9 @@ func (s *workerAssignmentService) SelectWorkerFor(ctx context.Context, lookup Pl candidates := make([]PlacementCandidate, 0, len(rows)) for _, row := range rows { + if lookup.ExcludeWorkerID != nil && row.WorkerID == *lookup.ExcludeWorkerID { + continue + } candidates = append(candidates, PlacementCandidate{ WorkerID: row.WorkerID, Region: row.Region, @@ -263,6 +279,28 @@ func (s *workerAssignmentService) selectFallback(ctx context.Context, req Placem return &PlacementResult{Worker: &w}, nil } +// warmupPoolFor resolves which warmup pool a mailbox joins. Unrelated to +// placement: it is a property of the organization's subscription, and workers +// no longer carry a tier to infer it from. Anything unknown answers "free", +// which is the conservative direction - a paid mailbox in the free pool warms +// more slowly, where the reverse would put unproven mail in front of paying +// customers. +func (s *workerAssignmentService) warmupPoolFor(ctx context.Context, orgID uuid.UUID) string { + if s.subRepo == nil { + return "free" + } + // With billing disabled there is no free/paid split to enforce, so every + // org gets the premium pool. Mirrors feature.gate's self-host unlock. + if config.BillingProvider() == "none" { + return "premium" + } + sub, err := s.subRepo.GetByOrganizationID(ctx, orgID) + if err != nil || sub == nil || !sub.HasPaidSubscription() { + return "free" + } + return "premium" +} + // hasIsolatedEgress reports whether the org's plan entitles it to a worker // nobody else sends from. Any lookup failure answers false: the entitlement // only ever adds preference, so losing it degrades to ordinary placement. @@ -413,9 +451,10 @@ func (s *workerAssignmentService) MigrateEmailsFromWorker(ctx context.Context, w continue } res, err := s.SelectWorkerFor(ctx, PlacementLookup{ - EmailAccountID: accountID, - OrgID: *state.OrganizationID, - Region: state.WorkerRegion, + EmailAccountID: accountID, + OrgID: *state.OrganizationID, + Region: state.WorkerRegion, + ExcludeWorkerID: &workerID, }) if err != nil || res == nil || res.Worker == nil || res.Worker.ID == workerID { continue diff --git a/internal/app/worker/rotation.go b/internal/app/worker/rotation.go index 8d391a4b7..5f669cb9d 100644 --- a/internal/app/worker/rotation.go +++ b/internal/app/worker/rotation.go @@ -102,7 +102,12 @@ func EvaluateRotation(in RotationInput) (RotationUrgency, string) { } if in.OnSomeoneElsesReservedWorker { - return RotationOpportunistic, "worker is reserved for another organization" + // Elevated, not opportunistic: this mailbox has to leave for the + // reservation to mean anything, and an opportunistic move would be + // weighed against the incumbent's stickiness bonus and refused every + // tick, so the stranger would never go and the row would re-enter the + // scan budget forever. + return RotationElevated, "worker is reserved for another organization" } if in.AwayFromOwnReservedWorker { return RotationOpportunistic, "organization has a reserved worker elsewhere" diff --git a/internal/app/worker/rotation_test.go b/internal/app/worker/rotation_test.go index 8d886caa3..85538af40 100644 --- a/internal/app/worker/rotation_test.go +++ b/internal/app/worker/rotation_test.go @@ -107,16 +107,29 @@ func TestUrgentMovesTakeAnythingEligible(t *testing.T) { } func TestReservedWorkerDriftRotatesBothWays(t *testing.T) { + // A stranger has to leave for the reservation to mean anything, and the + // destination does not have to be better than where it is. Weighing it + // against the incumbent's stickiness bonus would refuse the move forever. stranger := healthyInput() stranger.OnSomeoneElsesReservedWorker = true - if urgency, _ := EvaluateRotation(stranger); urgency != RotationOpportunistic { - t.Fatal("a mailbox on someone else's reserved worker should be moved off") + urgency, _ := EvaluateRotation(stranger) + if urgency != RotationElevated { + t.Fatalf("a mailbox on someone else's reserved worker must be evicted, got %v", urgency) + } + if !WorthMoving(urgency, 5.0, 0.1, false) { + t.Fatal("evicting a stranger must accept any eligible destination") } + // Pulling the owner back is opportunistic, but the placer marks the + // reserved worker as mandated so the score comparison does not block it. owner := healthyInput() owner.AwayFromOwnReservedWorker = true - if urgency, _ := EvaluateRotation(owner); urgency != RotationOpportunistic { - t.Fatal("a mailbox away from its own reserved worker should be pulled back") + urgency, _ = EvaluateRotation(owner) + if urgency != RotationOpportunistic { + t.Fatalf("a mailbox away from its own reserved worker should be pulled back, got %v", urgency) + } + if !WorthMoving(urgency, 5.0, 0.1, true) { + t.Fatal("the reserved worker is mandated, so the move must not be score-gated") } } diff --git a/internal/repository/pg_worker.go b/internal/repository/pg_worker.go index f56198e62..4f2f0c1d1 100644 --- a/internal/repository/pg_worker.go +++ b/internal/repository/pg_worker.go @@ -211,13 +211,14 @@ func (r *workerRepository) GetAllActiveWorkers(ctx context.Context) ([]models.Wo func (r *workerRepository) GetIdleUnboundWorker(ctx context.Context) (*models.Worker, error) { w, err := scanWorker(r.db.QueryRow(ctx, workerSelect+` WHERE n.active + AND n.last_seen_at > now() - $1::interval AND w.account_count = 0 AND NOT EXISTS ( SELECT 1 FROM dedicated_worker_assignments dwa WHERE dwa.worker_id = w.id AND dwa.released_at IS NULL ) ORDER BY w.created_at ASC - LIMIT 1`)) + LIMIT 1`, WorkerLivenessWindow)) if errors.Is(err, pgx.ErrNoRows) { return nil, nil } diff --git a/internal/sandbox/seed.go b/internal/sandbox/seed.go index ade00421d..b8d3b0f62 100644 --- a/internal/sandbox/seed.go +++ b/internal/sandbox/seed.go @@ -406,11 +406,16 @@ func seedIdentity(ctx context.Context, pool *pgxpool.Pool) error { // placement and the reconciler treat it as live; the real worker process adopts // the row on its first heartbeat. func seedWorker(ctx context.Context, pool *pgxpool.Pool) error { + if _, err := pool.Exec(ctx, ` + INSERT INTO fleet_nodes (id, role, name, notes, address, active, last_seen_at) + VALUES ($1, 'worker', 'worker-sandbox-1', 'Sandbox worker (make sandbox / make worker)', '127.0.0.1', TRUE, now()) + ON CONFLICT (id) DO UPDATE SET active = TRUE, last_seen_at = now(), updated_at = now()`, + sandboxWorker); err != nil { + return err + } _, err := pool.Exec(ctx, ` - INSERT INTO workers (id, name, notes, ip_addr, active, account_count) - VALUES ($1, 'worker-sandbox-1', 'Sandbox worker (make sandbox / make worker)', '127.0.0.1', TRUE, 0) - ON CONFLICT (id) DO UPDATE SET active = TRUE, updated_at = NOW()`, - sandboxWorker) + INSERT INTO workers (id, account_count) VALUES ($1, 0) + ON CONFLICT (id) DO NOTHING`, sandboxWorker) return err } diff --git a/internal/seed/workers.go b/internal/seed/workers.go index 7f30a6355..710578d61 100644 --- a/internal/seed/workers.go +++ b/internal/seed/workers.go @@ -23,18 +23,26 @@ func seedWorkers(ctx context.Context, pool *pgxpool.Pool, r *Result) error { } for _, w := range workers { - _, err := pool.Exec(ctx, ` - INSERT INTO workers (id, name, notes, ip_addr, active, region, account_count, created_at, updated_at) - VALUES ($1,$2,$3,$4,TRUE,$5,0,$6,$6) + // A worker is a node (the machine) plus a placement row (the mail it + // carries). Seeding both is what an enrolling node does. + if _, err := pool.Exec(ctx, ` + INSERT INTO fleet_nodes (id, role, name, notes, address, active, region, last_seen_at, created_at, updated_at) + VALUES ($1,'worker',$2,$3,$4,TRUE,$5,now(),$6,$6) ON CONFLICT (id) DO UPDATE SET name = EXCLUDED.name, notes = EXCLUDED.notes, - ip_addr = EXCLUDED.ip_addr, + address = EXCLUDED.address, active = TRUE, region = EXCLUDED.region, - updated_at = NOW() - `, w.id, w.name, w.notes, w.ipAddr, w.region, time.Now()) - if err != nil { + last_seen_at = now(), + updated_at = now() + `, w.id, w.name, w.notes, w.ipAddr, w.region, time.Now()); err != nil { + return err + } + if _, err := pool.Exec(ctx, ` + INSERT INTO workers (id, account_count) VALUES ($1, 0) + ON CONFLICT (id) DO NOTHING + `, w.id); err != nil { return err } r.Workers = append(r.Workers, SeededWorker{Name: w.name, Region: w.region, ID: w.id.String()})