fix(engine): end recreate drain on exit

docs/onebox.run-v1.schema.json

@@ -1995,7 +1995,7 @@
                 "type": "string"
               },
               "wait": {
-                "description": "Time allowed for the proxy to stop routing before shutdown begins, at most 7d. Expects a duration such as 30s, 5m, 1h30m or 14d.",
+                "description": "Maximum drain window before shutdown continues, at most 7d. Recreate workloads continue sooner when every old container exits. Rolling workloads wait the full interval before stopping each container when their health check supports drain guarding. Expects a duration such as 30s, 5m, 1h30m or 14d.",
                 "examples": [
                   "10s"
                 ],

internal/app/runtime.go

@@ -145,8 +145,10 @@ func (w Workload) ReadyTiming() (within, interval time.Duration) {
 	return within, interval
 }
 
-// DrainWait is how long to leave a container marked unhealthy before stopping
-// it, so the proxy has time to notice and stop sending it traffic.
+// DrainWait is the authored drain interval. Rolling workloads with a
+// drain-guardable health check use it as a fixed pre-stop window after the drain
+// attempt; recreate workloads use it as the maximum time allowed for signalled
+// containers to exit.
 //
 // Only an authored wait counts. The derived value this used to fall back to
 // was unreachable — every caller checks `drain.wait` was written before asking

internal/app/types.go

@@ -216,7 +216,7 @@ type Health struct {
 
 type Drain struct {
 	Signal string `json:"signal" description:"Signal sent to begin graceful shutdown." default:"TERM"`
-	Wait   string `json:"wait,omitempty" description:"Time allowed for the proxy to stop routing before shutdown begins, at most 7d." example:"10s"`
+	Wait   string `json:"wait,omitempty" description:"Maximum drain window before shutdown continues, at most 7d. Recreate workloads continue sooner when every old container exits. Rolling workloads wait the full interval before stopping each container when their health check supports drain guarding." example:"10s"`
 	Grace  string `json:"grace,omitempty" description:"Maximum graceful-shutdown time before forced termination, at most 7d." example:"30s"`
 }
 

internal/engine/deploy_test.go

@@ -149,6 +149,10 @@ func happyFake() *transport.Fake {
 				return transport.Result{Stdout: "unhealthy\n"}, true
 			}
 			return transport.Result{Stdout: "healthy\n"}, true
+		case strings.Contains(cmd, "{{.State.Running}}") && strings.Contains(cmd, "W1"):
+			// Recreate drain observes the old worker after signalling it. The happy
+			// fixture models a worker that exits promptly and can be replaced.
+			return transport.Result{Stdout: "false\n"}, true
 		case strings.Contains(cmd, "service='worker'") && strings.Contains(cmd, "ob.release="):
 			if !workerGone || recreateCount > initialRecreateCount {
 				return transport.Result{Stdout: "W1\n"}, true

internal/engine/engine.go

@@ -21,8 +21,10 @@ import (
 type Options struct {
 	Verbose bool
 	Out     io.Writer
-	// Sleep and Now are injectable for tests.
+	// Sleep, Wait and Now are injectable for tests. Wait is the context-aware
+	// form used by lifecycle polling; Sleep remains for fixed protocol delays.
 	Sleep func(time.Duration)
+	Wait  func(context.Context, time.Duration) error
 	Now   func() time.Time
 	// SecretGeneration returns a fresh opaque identifier for a secret
 	// transaction. It is injectable so crash-boundary tests remain deterministic.
@@ -124,6 +126,21 @@ func New(a *app.Resolved, c *ctypes.Project, t transport.Transport, o Options) *
 	if o.Sleep == nil {
 		o.Sleep = time.Sleep
 	}
+	if o.Wait == nil {
+		o.Wait = func(ctx context.Context, d time.Duration) error {
+			if d <= 0 {
+				return ctx.Err()
+			}
+			timer := time.NewTimer(d)
+			defer timer.Stop()
+			select {
+			case <-ctx.Done():
+				return ctx.Err()
+			case <-timer.C:
+				return nil
+			}
+		}
+	}
 	if o.Now == nil {
 		o.Now = time.Now
 	}

internal/engine/plan.go

@@ -571,14 +571,14 @@ func (e *Engine) DescribeWorkloadPlans(remoteCompose string, plans map[string]Wo
 				scale = fmt.Sprintf(" --scale %s=%d", svc, n)
 			}
 			// Gated exactly as recreate gates it — an authored `drain.wait`,
-			// nothing else. recreate signals every container and then sleeps,
-			// whatever the signal is, so excluding the default TERM here hid a
-			// kill and a pause the deploy certainly takes. A plan that shows a
-			// step execution skips, or hides one it takes, is a plan nobody can
-			// check against.
+			// nothing else. recreate signals every container and waits up to the
+			// bound for all of them to exit, whatever the signal is. Excluding the
+			// default TERM here hid a kill and a wait the deploy certainly takes.
+			// A plan that shows a step execution skips, or hides one it takes, is
+			// a plan nobody can check against.
 			if wait := role.DrainWait(); role.Drain != nil && role.Drain.Wait != "" && wait > 0 {
 				out = append(out,
-					fmt.Sprintf("  docker kill --signal=%s <current %s>; wait %s", role.DrainSignal(), svc, wait),
+					fmt.Sprintf("  docker kill --signal=%s <current %s>; wait up to %s for exit", role.DrainSignal(), svc, wait),
 				)
 			}
 			out = append(out,

internal/engine/plan_drain_test.go

@@ -31,7 +31,7 @@ func TestPlanPromisesADrainWaitOnlyWhenTheDeployTakesOne(t *testing.T) {
 	withWait.Drain = &app.Drain{Signal: "USR1", Wait: "12s"}
 	config.Workloads["web"] = withWait
 	lines = strings.Join(newPlanEngine(t, config).Describe("/var/lib/ob/sample/releases/R1/compose.yaml"), "\n")
-	if !strings.Contains(lines, "--signal=USR1 <current web>; wait 12s") {
+	if !strings.Contains(lines, "--signal=USR1 <current web>; wait up to 12s for exit") {
 		t.Fatalf("the plan omits the drain step the deploy does take:\n%s", lines)
 	}
 }
@@ -42,8 +42,8 @@ func newPlanEngine(t *testing.T, config *app.Resolved) *Engine {
 }
 
 // recreate sends the drain signal for every authored wait, TERM included, and
-// then sleeps. A plan that shows the step only for a non-default signal hides
-// a kill and a pause the deploy will certainly take.
+// then waits up to the bound for exit. A plan that shows the step only for a
+// non-default signal hides a kill and wait the deploy will certainly take.
 func TestPlanShowsTheDrainStepForTheDefaultSignalToo(t *testing.T) {
 	config := testConfig()
 	workload := config.Workloads["web"]
@@ -53,7 +53,7 @@ func TestPlanShowsTheDrainStepForTheDefaultSignalToo(t *testing.T) {
 	config.Workloads["web"] = workload
 
 	lines := strings.Join(newPlanEngine(t, config).Describe("/var/lib/ob/sample/releases/R1/compose.yaml"), "\n")
-	if !strings.Contains(lines, "--signal=TERM <current web>; wait 15s") {
+	if !strings.Contains(lines, "--signal=TERM <current web>; wait up to 15s for exit") {
 		t.Fatalf("the plan hides the drain step recreate will take:\n%s", lines)
 	}
 }

internal/engine/recreate.go

@@ -8,8 +8,11 @@ import (
 	"os/exec"
 	"path/filepath"
 	"strings"
+	"time"
 )
 
+const recreateDrainPollInterval = 250 * time.Millisecond
+
 // RecreateRole replaces a role's containers in place: a stated brief gap, the
 // mode for workers and anything that can't roll. Honors replicas —
 // recreates the whole fleet at the desired count and gives each a clean slot
@@ -33,7 +36,7 @@ func (e *Engine) recreateRoleForRelease(ctx context.Context, roleName, remoteCom
 	if err := e.pullBeforeRelease(ctx, svc, cc); err != nil {
 		return err
 	}
-	// Signal before recreate whenever the contract declares a fixed drain wait.
+	// Signal before recreate whenever the contract declares a bounded drain wait.
 	// Compose sends TERM during replacement too, but doing it only then skipped
 	// drain.wait entirely and left recreate workers at Compose's default timeout.
 	if wait := role.DrainWait(); role.Drain != nil && role.Drain.Wait != "" && wait > 0 {
@@ -53,7 +56,9 @@ func (e *Engine) recreateRoleForRelease(ctx context.Context, roleName, remoteCom
 			}
 		}
 		if len(ids) > 0 {
-			e.Opts.Sleep(wait)
+			if err := e.waitForContainersExit(ctx, svc, ids, wait); err != nil {
+				return err
+			}
 		}
 	}
 	scaleArg := ""
@@ -92,6 +97,68 @@ func (e *Engine) recreateRoleForRelease(ctx context.Context, roleName, remoteCom
 	return e.reslot(ctx, svc, releaseID, desired)
 }
 
+// waitForContainersExit gives the exact containers signalled above up to wait
+// to finish. It observes only runtime lifecycle state: an exited or vanished
+// container is done, while an unknown inspection failure aborts rather than
+// pretending graceful shutdown succeeded.
+func (e *Engine) waitForContainersExit(ctx context.Context, svc string, ids []string, wait time.Duration) error {
+	// Consume successful waits directly. Options.Now may deliberately be fixed for
+	// deterministic operation metadata while Wait still uses a real timer.
+	remaining := wait
+	pending := append([]string(nil), ids...)
+	for len(pending) > 0 {
+		if err := ctx.Err(); err != nil {
+			return err
+		}
+		stillRunning := pending[:0]
+		for _, id := range pending {
+			running, err := e.drainingContainerRunning(ctx, id)
+			if err != nil {
+				return fmt.Errorf("wait for %s drain: %w", svc, err)
+			}
+			if running {
+				stillRunning = append(stillRunning, id)
+			}
+		}
+		pending = stillRunning
+		if len(pending) == 0 {
+			return nil
+		}
+		if remaining <= 0 {
+			return nil
+		}
+		delay := min(remaining, recreateDrainPollInterval)
+		if err := e.Opts.Wait(ctx, delay); err != nil {
+			return err
+		}
+		remaining -= delay
+	}
+	return nil
+}
+
+func (e *Engine) drainingContainerRunning(ctx context.Context, id string) (bool, error) {
+	res, err := e.T.Run(ctx, "docker inspect -f '{{.State.Running}}' "+id)
+	if err != nil {
+		return false, err
+	}
+	if res.ExitCode != 0 {
+		message := strings.TrimSpace(res.Stderr)
+		lower := strings.ToLower(message)
+		if strings.Contains(lower, "no such object") || strings.Contains(lower, "no such container") {
+			return false, nil
+		}
+		return false, fmt.Errorf("inspect draining container %s failed (exit %d): %s", id, res.ExitCode, message)
+	}
+	switch strings.TrimSpace(res.Stdout) {
+	case "true":
+		return true, nil
+	case "false":
+		return false, nil
+	default:
+		return false, fmt.Errorf("inspect draining container %s returned an invalid running state", id)
+	}
+}
+
 // RunHook executes a user hook verbatim. Hooks are unplannable
 // commands — the operator's own, same trust level as their shell). Release
 // hooks get compose env exported so `docker compose ...` targets that release;