diff --git a/internal/server/health_internal_test.go b/internal/server/health_internal_test.go index 11a5203f..cb6e57ff 100644 --- a/internal/server/health_internal_test.go +++ b/internal/server/health_internal_test.go @@ -4,6 +4,8 @@ package server import ( + "context" + "sync/atomic" "testing" "time" @@ -11,6 +13,26 @@ import ( "github.com/stretchr/testify/assert" ) +// fakeHealthCheckStream is a minimal pb.Healthz_CheckServer whose lifetime is +// driven entirely by ctx, so tests can control when Check() unblocks. sendCount +// tracks how many health messages Check() has sent, so tests can verify the +// notifyHealthChan branch does (or doesn't) fire for a given stream. +type fakeHealthCheckStream struct { + pb.Healthz_CheckServer + ctx context.Context + sendCount atomic.Int32 +} + +func (f *fakeHealthCheckStream) Context() context.Context { return f.ctx } +func (f *fakeHealthCheckStream) Send(*pb.Health) error { + f.sendCount.Add(1) + return nil +} +func (f *fakeHealthCheckStream) Recv() (*pb.Health, error) { + <-f.ctx.Done() + return nil, f.ctx.Err() +} + // clearHealthNotifiers empties the package-level notifier registry so each test // starts from a known state and MonitorHealthChan does not fan out to notifiers // left over from other tests. @@ -37,11 +59,15 @@ func TestNotifyHealthRegistersAndReplacesReceiver(t *testing.T) { second := make(chan pb.HealthStatus_Code, 1) notifyHealth(addr, second) - // Registering a second notifier for the same client must close the prior one. - _, open := <-first - assert.False(t, open, "prior receiver channel should be closed on replacement") + // Registering a second notifier for the same client must leave only the + // new receiver registered, without closing the prior channel: the prior + // owner (a still-running Check() call) is responsible for closing it. + select { + case _, open := <-first: + assert.True(t, open, "prior receiver channel should not be closed by replacement") + default: + } - // ...and leave only the new receiver registered. got, ok = healthNotifiers.Get(addr) assert.True(t, ok) assert.Equal(t, second, got) @@ -90,3 +116,103 @@ func TestMonitorHealthChanFansOutToNotifiers(t *testing.T) { t.Fatal("MonitorHealthChan did not exit after its receiver was closed") } } + +// waitForNotifier polls healthNotifiers until addr's registered value differs +// from prior, returning the new value. +func waitForNotifier(t *testing.T, addr string, prior interface{}) interface{} { + t.Helper() + deadline := time.After(time.Second) + for { + if got, ok := healthNotifiers.Get(addr); ok && got != prior { + return got + } + select { + case <-deadline: + t.Fatalf("timed out waiting for a new notifier registration for %s", addr) + return nil + case <-time.After(time.Millisecond): + } + } +} + +// TestCheckReconnectDoesNotPanicOrLoseNewRegistration reproduces a client +// reconnecting while the server hasn't yet noticed the old stream is gone: +// two Check() calls register for the same clientAddr, the second replacing +// the first, and the first stream is then torn down. +func TestCheckReconnectDoesNotPanicOrLoseNewRegistration(t *testing.T) { + clearHealthNotifiers() + defer clearHealthNotifiers() + + addr := "client-reconnect" + oldGetClientAddr := GetClientAddr + GetClientAddr = func(context.Context) (string, error) { return addr, nil } + defer func() { GetClientAddr = oldGetClientAddr }() + + hlthSrv := &HealthzServer{} + + firstCtx, cancelFirst := context.WithCancel(context.Background()) + defer cancelFirst() + firstStream := &fakeHealthCheckStream{ctx: firstCtx} + firstDone := make(chan struct{}) + var firstPanic interface{} + go func() { + defer close(firstDone) + defer func() { firstPanic = recover() }() + _ = hlthSrv.Check(firstStream) + }() + + firstNotifier := waitForNotifier(t, addr, nil) + // initial connect message + initialFirstSends := firstStream.sendCount.Load() + + secondCtx, cancelSecond := context.WithCancel(context.Background()) + defer cancelSecond() + secondStream := &fakeHealthCheckStream{ctx: secondCtx} + secondDone := make(chan struct{}) + var secondPanic interface{} + go func() { + defer close(secondDone) + defer func() { secondPanic = recover() }() + _ = hlthSrv.Check(secondStream) + }() + + // The second Check() call replaces the registration for the same clientAddr. + secondNotifier := waitForNotifier(t, addr, firstNotifier) + initialSecondSends := secondStream.sendCount.Load() + + // Broadcast a code the way MonitorHealthChan would; only the current + // registrant (second) should receive it and reach the notifyHealthChan + // send branch, proving the stale (first) channel is neither written to + // nor spun on. + secondNotifier.(chan pb.HealthStatus_Code) <- pb.HealthStatus_UNHEALTHY + + assert.Eventually(t, func() bool { + return secondStream.sendCount.Load() > initialSecondSends + }, time.Second, time.Millisecond, "second stream should have received the broadcast health message") + assert.Equal(t, initialFirstSends, firstStream.sendCount.Load(), "stale first stream must not receive or spin on the broadcast") + + // Tearing down the superseded first stream must not panic (double close) + // or remove the second stream's registration. + cancelFirst() + select { + case <-firstDone: + case <-time.After(time.Second): + t.Fatal("first Check() call did not return after its context was canceled") + } + assert.Nil(t, firstPanic, "Check() must not panic when a superseded stream ends") + + got, ok := healthNotifiers.Get(addr) + assert.True(t, ok, "the second stream's registration should remain") + assert.Equal(t, secondNotifier, got, "the second stream's registration must be untouched") + + cancelSecond() + select { + case <-secondDone: + case <-time.After(time.Second): + t.Fatal("second Check() call did not return after its context was canceled") + } + assert.Nil(t, secondPanic, "Check() must not panic on normal shutdown") + + _, ok = healthNotifiers.Get(addr) + assert.False(t, ok, "registration should be removed once the second stream ends") +} diff --git a/internal/server/server.go b/internal/server/server.go index 4d545be7..e054b0ee 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -768,12 +768,9 @@ func clientExists(ctx context.Context) bool { } func notifyHealth(clientAddr string, receiver chan pb.HealthStatus_Code) { - // only allow one notifier per client - if recInt, ok := healthNotifiers.Get(clientAddr); ok { - rec := recInt.(chan pb.HealthStatus_Code) - close(rec) - healthNotifiers.Delete(clientAddr) - } + // only allow one notifier per client; replacing the map entry (rather than + // closing the previous channel) avoids racing with the prior Check() call, + // which still owns that channel and closes it itself once its stream ends. healthNotifiers.Add(clientAddr, receiver) } @@ -800,7 +797,9 @@ func (s *HealthzServer) Check(stream pb.Healthz_CheckServer) error { notifyHealthChan := make(chan pb.HealthStatus_Code) notifyHealth(clientAddr, notifyHealthChan) defer func() { - healthNotifiers.Delete(clientAddr) + // Only remove our own registration; a newer Check() call for the same + // clientAddr may have already replaced it. + _ = healthNotifiers.DeleteIfEqual(clientAddr, notifyHealthChan) close(notifyHealthChan) }() diff --git a/internal/util/concurrentMap.go b/internal/util/concurrentMap.go index cb78ba2c..7d95b698 100644 --- a/internal/util/concurrentMap.go +++ b/internal/util/concurrentMap.go @@ -34,6 +34,20 @@ func (cm *ConcurrentMap) Delete(key string) { delete(cm.items, key) } +// DeleteIfEqual deletes key only if its current value is still expected, +// avoiding removal of an entry that was replaced concurrently. Returns +// true if the key was deleted, false if it was not present or its value was +// not equal to expected. +func (cm *ConcurrentMap) DeleteIfEqual(key string, expected interface{}) bool { + cm.Lock() + defer cm.Unlock() + if item, ok := cm.items[key]; ok && item == expected { + delete(cm.items, key) + return true + } + return false +} + // Get Return the value for a given key in the map func (cm *ConcurrentMap) Get(key string) (interface{}, bool) { cm.RLock() diff --git a/internal/util/concurrentMap_test.go b/internal/util/concurrentMap_test.go index 97badc4b..2864d1ef 100644 --- a/internal/util/concurrentMap_test.go +++ b/internal/util/concurrentMap_test.go @@ -52,6 +52,39 @@ func TestConcurrentMapDelete(t *testing.T) { assert.False(t, ok) } +func TestConcurrentMapDeleteIfEqual(t *testing.T) { + cMap := NewConcurrentMap() + testItem := TestItem{"test item"} + staleItem := TestItem{"stale item"} + replacementItem := TestItem{"replacement item"} + + cMap.Add("testItem", testItem) + + // a stale expected value must not delete the entry + assert.False(t, cMap.DeleteIfEqual("testItem", staleItem)) + cItem, ok := cMap.Get("testItem") + assert.True(t, ok) + assert.Equal(t, testItem, cItem) + + // a newer value must not be removed by an old expected value + cMap.Add("testItem", replacementItem) + assert.False(t, cMap.DeleteIfEqual("testItem", testItem)) + cItem, ok = cMap.Get("testItem") + assert.True(t, ok) + assert.Equal(t, replacementItem, cItem) + + // a matching expected value deletes the entry + assert.True(t, cMap.DeleteIfEqual("testItem", replacementItem)) + cItem, ok = cMap.Get("testItem") + assert.Nil(t, cItem) + assert.False(t, ok) + + // deleting a missing key is a no-op + assert.False(t, cMap.DeleteIfEqual("missingItem", testItem)) + _, ok = cMap.Get("missingItem") + assert.False(t, ok) +} + func TestConcurrentMapGetList(t *testing.T) { cMap := NewConcurrentMap() assert.NotNil(t, cMap)