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 +}