From b3618997745686f6687715d9464249ec6203c494 Mon Sep 17 00:00:00 2001 From: mhenrixon Date: Wed, 29 Jul 2026 15:01:13 +0200 Subject: [PATCH] feat(load-balancer): suspend and resume a pool around a sleeping container The pool half of scale-to-zero (#19). Nothing calls it yet -- the idle controller is wired up when the request gate lands -- so proxy behavior is unchanged. SuspendForSleep empties the pool and stops probing, so a container that was deliberately stopped is neither routed to nor dialled once a second for the whole nap. It runs before the containers go down, and the controller holds arriving requests while it does, so the empty pool is never observable. ResumeFromSleep is the subtle half, and it exists because of how a single-target pool behaves -- which is the only shape scale-to-zero actually targets. updateHealthyTargets stops health checking once a lone target first goes healthy, so markHealthy has already fired and waitForHealthyContext is already cancelled. Without re-arming, WaitUntilHealthy returns nil INSTANTLY for a container that has not started, and the wake would report ready and forward its held request straight into a connection refused. Verified by mutation: drop the re-arm and TestLoadBalancer_ResumeFromSleepRearmsWaitUntilHealthy fails in 0.00s with "readiness was reported for a target that never answered". The context is replaced rather than cancelled. WaitUntilHealthy treats any non-deadline cancellation as success, so cancelling would tell a waiter that parked before the resume "healthy" at the very moment every target was marked unverified. The cost is that such a waiter times out instead of ever being released; that is a timeout it can retry rather than a false ready, and it is unreachable in the real flow because the controller serializes wakes behind its generation counter. Targets re-enter as Adding, not Healthy: a successful probe promotes Adding to Healthy while a failed one only ever demotes Healthy to Unhealthy, so a container that never comes up stays out of the pool instead of flapping into it. A woken target is therefore held to exactly the standard a freshly deployed one is. WaitUntilHealthy now reads waitForHealthyContext under the lock. This is the third race from the review that landed in 9e6f28b, deliberately left then: the field was written once in NewLoadBalancer before publication, so it was not a race until something wrote it at runtime. ResumeFromSleep is that writer, so the lock lands in the commit where the second writer appears rather than ahead of it. Reuses BeginHealthChecks rather than adding a restart variant. The plan called for a separate method because BeginHealthChecks assigned stateConsumer outside the inflight lock, but 9e6f28b moved that assignment inside, so there is nothing left to work around. Refs #19 --- internal/server/load_balancer.go | 62 ++++++- internal/server/load_balancer_sleep_test.go | 187 ++++++++++++++++++++ 2 files changed, 248 insertions(+), 1 deletion(-) create mode 100644 internal/server/load_balancer_sleep_test.go diff --git a/internal/server/load_balancer.go b/internal/server/load_balancer.go index b4d4782..6d14fbc 100644 --- a/internal/server/load_balancer.go +++ b/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() } diff --git a/internal/server/load_balancer_sleep_test.go b/internal/server/load_balancer_sleep_test.go new file mode 100644 index 0000000..1b271c5 --- /dev/null +++ b/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 +}