fix: reconnecting clients cause a panic in arke

internal/server/health_internal_test.go

@@ -4,13 +4,35 @@
 package server
 
 import (
+	"context"
+	"sync/atomic"
 	"testing"
 	"time"
 
 	pb "github.com/sassoftware/arke/api"
 	"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")
+}

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)
 	}()
 

internal/util/concurrentMap.go

@@ -34,6 +34,16 @@ 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.
+func (cm *ConcurrentMap) DeleteIfEqual(key string, expected interface{}) {
+	cm.Lock()
+	defer cm.Unlock()
+	if cm.items[key] == expected {
+		delete(cm.items, key)
+	}
+}
+
 // Get Return the value for a given key in the map
 func (cm *ConcurrentMap) Get(key string) (interface{}, bool) {
 	cm.RLock()

internal/util/concurrentMap_test.go

@@ -52,6 +52,30 @@ func TestConcurrentMapDelete(t *testing.T) {
 	assert.False(t, ok)
 }
 
+func TestConcurrentMapDeleteIfEqual(t *testing.T) {
+	cMap := NewConcurrentMap()
+	testItem := TestItem{"test item"}
+	staleItem := TestItem{"stale item"}
+	cMap.Add("testItem", testItem)
+
+	// a stale expected value must not delete the entry
+	cMap.DeleteIfEqual("testItem", staleItem)
+	cItem, ok := cMap.Get("testItem")
+	assert.True(t, ok)
+	assert.Equal(t, testItem, cItem)
+
+	// a matching expected value deletes the entry
+	cMap.DeleteIfEqual("testItem", testItem)
+	cItem, ok = cMap.Get("testItem")
+	assert.Nil(t, cItem)
+	assert.False(t, ok)
+
+	// deleting a missing key is a no-op
+	cMap.DeleteIfEqual("missingItem", testItem)
+	_, ok = cMap.Get("missingItem")
+	assert.False(t, ok)
+}
+
 func TestConcurrentMapGetList(t *testing.T) {
 	cMap := NewConcurrentMap()
 	assert.NotNil(t, cMap)