feat(load-balancer): suspend and resume a pool around a sleeping container

internal/server/load_balancer.go

@@ -151,7 +151,14 @@ func (lb *LoadBalancer) HealthyTargets() TargetList {
 }
 
 func (lb *LoadBalancer) WaitUntilHealthy(timeout time.Duration) error {
-	ctx, cancel := context.WithTimeout(lb.waitForHealthyContext, timeout)
+	// Locked because ResumeFromSleep replaces this context at runtime. Until
+	// scale-to-zero existed it was written once in NewLoadBalancer before the
+	// value was published, and reading it unlocked was safe.
+	lb.lock.Lock()
+	waitForHealthy := lb.waitForHealthyContext
+	lb.lock.Unlock()
+
+	ctx, cancel := context.WithTimeout(waitForHealthy, timeout)
 	defer cancel()
 
 	<-ctx.Done()
@@ -183,6 +190,59 @@ func (lb *LoadBalancer) RecheckHealth() {
 	lb.all.BeginHealthChecks(lb)
 }
 
+// SuspendForSleep empties the pool and stops probing, so a container that is
+// deliberately stopped is neither routed to nor dialled once a second for the
+// whole nap.
+//
+// Called before the containers go down. The idle controller holds arriving
+// requests while this runs, so the empty pool is never observable to a client.
+func (lb *LoadBalancer) SuspendForSleep() {
+	lb.all.StopHealthChecks()
+
+	for _, target := range lb.all {
+		target.updateState(TargetStateUnhealthy)
+	}
+
+	lb.lock.Lock()
+	defer lb.lock.Unlock()
+
+	lb.writers = TargetList{}
+	lb.readers = TargetList{}
+}
+
+// ResumeFromSleep puts the pool back in the state a fresh deployment starts in --
+// unverified, health-checked -- and re-arms WaitUntilHealthy so a caller can wait
+// for the woken containers to actually answer.
+//
+// Re-arming is the whole point. A single-target pool stops probing once its
+// target first goes healthy, so markHealthy has already fired and the context is
+// already cancelled; without this, WaitUntilHealthy returns nil instantly against
+// a container that has not started, and the wake forwards its held request into a
+// connection refused.
+//
+// The context is replaced rather than cancelled: WaitUntilHealthy reports any
+// non-deadline cancellation as success, so cancelling would tell a waiter that
+// parked before the resume "healthy" at the exact moment every target was marked
+// unverified.
+func (lb *LoadBalancer) ResumeFromSleep() {
+	lb.lock.Lock()
+	lb.waitForHealthyContext, lb.markHealthy = context.WithCancel(context.Background())
+	lb.lock.Unlock()
+
+	// Adding, not Healthy: a successful probe promotes Adding to Healthy, while a
+	// failed one only ever demotes Healthy to Unhealthy. A container that never
+	// comes up therefore stays out of the pool instead of flapping into it.
+	for _, target := range lb.all {
+		target.updateState(TargetStateAdding)
+	}
+
+	// Safe to reuse rather than needing a separate restart path: BeginHealthChecks
+	// now assigns stateConsumer under the inflight lock. NewHealthCheck runs one
+	// immediate probe before it starts ticking, so readiness costs a round trip
+	// rather than a whole check interval.
+	lb.all.BeginHealthChecks(lb)
+}
+
 func (lb *LoadBalancer) Dispose() {
 	lb.all.StopHealthChecks()
 }

internal/server/load_balancer_sleep_test.go

@@ -0,0 +1,187 @@
+package server
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// testSleepableLoadBalancer builds a SINGLE-target pool on purpose. That is the
+// shape scale-to-zero actually targets, and the only shape where the health
+// checks stop themselves once the target first goes healthy -- which is what
+// makes the resume path subtle.
+func testSleepableLoadBalancer(t *testing.T, handler http.HandlerFunc) (*LoadBalancer, *httptest.Server) {
+	t.Helper()
+
+	backend := httptest.NewServer(handler)
+	t.Cleanup(backend.Close)
+
+	targets, err := NewTargetList([]string{backend.Listener.Addr().String()}, []string{}, defaultTargetOptions)
+	require.NoError(t, err)
+
+	lb := NewLoadBalancer(targets, DefaultWriterAffinityTimeout, false)
+	t.Cleanup(lb.Dispose)
+
+	require.NoError(t, lb.WaitUntilHealthy(5*time.Second))
+	require.Len(t, lb.HealthyTargets(), 1)
+
+	return lb, backend
+}
+
+func TestLoadBalancer_SuspendForSleepEmptiesThePoolAndStopsProbing(t *testing.T) {
+	var probes atomic.Int64
+
+	lb, _ := testSleepableLoadBalancer(t, func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path == defaultTargetOptions.HealthCheckConfig.Path {
+			probes.Add(1)
+		}
+		w.WriteHeader(http.StatusOK)
+	})
+
+	lb.SuspendForSleep()
+
+	assert.Empty(t, lb.HealthyTargets(), "a deliberately stopped container must not be routed to")
+
+	before := probes.Load()
+	time.Sleep(3 * defaultTargetOptions.HealthCheckConfig.Interval)
+
+	assert.Equal(t, before, probes.Load(),
+		"a sleeping service must not be dialled once a second for the whole nap")
+}
+
+// The defect this whole pair exists for. A single-target pool stops probing at
+// first-healthy, so markHealthy has already fired and waitForHealthyContext is
+// already cancelled. Without re-arming, WaitUntilHealthy returns nil INSTANTLY
+// against a container that has not started -- the wake would report ready and
+// forward the held request into a connection refused.
+func TestLoadBalancer_ResumeFromSleepRearmsWaitUntilHealthy(t *testing.T) {
+	var answering atomic.Bool
+
+	// Answering during setup so the pool reaches healthy the ordinary way, then
+	// silenced to stand in for a container that is starting but not yet listening.
+	answering.Store(true)
+
+	lb, _ := testSleepableLoadBalancer(t, func(w http.ResponseWriter, r *http.Request) {
+		if !answering.Load() {
+			w.WriteHeader(http.StatusServiceUnavailable)
+			return
+		}
+		w.WriteHeader(http.StatusOK)
+	})
+
+	answering.Store(false)
+
+	lb.SuspendForSleep()
+	lb.ResumeFromSleep()
+
+	// Nothing is answering yet, so readiness must NOT be reported.
+	err := lb.WaitUntilHealthy(300 * time.Millisecond)
+	require.Error(t, err, "readiness was reported for a target that never answered")
+	assert.ErrorIs(t, err, ErrorTargetFailedToBecomeHealthy)
+	assert.Empty(t, lb.HealthyTargets())
+
+	// Once the backend answers, readiness arrives through the machinery that
+	// already exists: probe -> HealthCheckCompleted -> updateHealthyTargets ->
+	// markHealthy.
+	answering.Store(true)
+
+	require.NoError(t, lb.WaitUntilHealthy(5*time.Second))
+	assert.Len(t, lb.HealthyTargets(), 1)
+}
+
+// Cancelling the old context instead of replacing it would tell a waiter that
+// parked before the resume "healthy" at the exact moment every target was marked
+// unverified: WaitUntilHealthy reports any non-deadline cancellation as success.
+func TestLoadBalancer_ResumeFromSleepDoesNotReleaseAPreviousWaiter(t *testing.T) {
+	var answering atomic.Bool
+
+	answering.Store(true)
+
+	lb, _ := testSleepableLoadBalancer(t, func(w http.ResponseWriter, r *http.Request) {
+		if !answering.Load() {
+			w.WriteHeader(http.StatusServiceUnavailable)
+			return
+		}
+		w.WriteHeader(http.StatusOK)
+	})
+
+	answering.Store(false)
+
+	lb.SuspendForSleep()
+	lb.ResumeFromSleep()
+
+	parked := make(chan error, 1)
+	go func() { parked <- lb.WaitUntilHealthy(500 * time.Millisecond) }()
+
+	// Give the waiter time to park on the context the resume installed.
+	time.Sleep(50 * time.Millisecond)
+
+	// A second resume -- a wake retried after a failure -- must not hand that
+	// waiter a success.
+	lb.ResumeFromSleep()
+
+	select {
+	case err := <-parked:
+		t.Fatalf("a resume released a parked waiter instead of an actual health check: %v", err)
+	case <-time.After(150 * time.Millisecond):
+	}
+
+	// It times out rather than ever being released, because the resume replaced
+	// the context it parked on along with the markHealthy that would close it.
+	// That is the deliberate trade: a waiter spanning a resume gets a timeout it
+	// can retry, never a false "healthy" for targets that were just marked
+	// unverified. It is also unreachable in the real flow -- the controller
+	// serializes wakes behind its generation counter, so each wake calls resume
+	// and then waits on the context that same resume installed.
+	err := <-parked
+	require.Error(t, err)
+	assert.ErrorIs(t, err, ErrorTargetFailedToBecomeHealthy)
+}
+
+// A woken target is held to exactly the standard a freshly deployed one is: it
+// re-enters unverified and only joins the pool after a probe succeeds, rather
+// than being assumed healthy the way a restored target is.
+func TestLoadBalancer_ResumeFromSleepReentersUnverified(t *testing.T) {
+	lb, _ := testSleepableLoadBalancer(t, func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	})
+
+	lb.SuspendForSleep()
+	for _, target := range lb.Targets() {
+		assert.Equal(t, TargetStateUnhealthy, target.State(), "suspended targets leave the pool")
+	}
+
+	lb.ResumeFromSleep()
+
+	require.NoError(t, lb.WaitUntilHealthy(5*time.Second))
+	assert.Len(t, lb.HealthyTargets(), 1)
+}
+
+// Sleep and wake are driven from the controller's goroutines while requests are
+// still being routed, so the pool mutation has to be safe under -race.
+func TestLoadBalancer_SuspendAndResumeAreSafeUnderConcurrentRouting(t *testing.T) {
+	lb, _ := testSleepableLoadBalancer(t, func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	})
+
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		for range 200 {
+			_ = lb.HealthyTargets()
+			lb.WaitUntilHealthy(time.Millisecond)
+		}
+	}()
+
+	for range 20 {
+		lb.SuspendForSleep()
+		lb.ResumeFromSleep()
+	}
+
+	<-done
+}