fix(sse): cancel active streams on user logout

clubanderson · clubanderson · commit 93d79190ce82 · 2026-04-10T09:44:59.000-04:00
Fixes #6029. SSE streaming endpoints (streamClusters in sse.go) ran inside SetBodyStreamWriter callbacks that blocked for up to sseOverallDeadline (~30s), so when a user logged out any in-flight SSE streams continued to emit cluster_data events until the deadline fired or the client disconnected. Mirrors the exec session cancellation pattern added in #6031 for /ws/exec sessions (#6024): - Add a per-user SSE session registry in sse.go with registerSSESession / unregisterSSESession / CancelUserSSEStreams, guarded by a plain sync.Mutex. - In streamClusters, capture the authenticated user ID before the deferred SetBodyStreamWriter callback runs (the fiber.Ctx may be reused by then) and register the stream context's cancel func in the registry, with a deferred unregister for normal stream end. - In auth.Logout, after CancelUserExecSessions, call CancelUserSSEStreams(claims.UserID) so logout tears down any live SSE streams the user had open. - Add sse_test.go with the four lifecycle tests mirroring exec_test.go: CancelsRegisteredContexts, OtherUserUnaffected, UnregisterSSESession_RemovesEntry, and NoSessions (no panic on empty map). streamDemoSSE is left unchanged — demo streams are instant, they do not block, so they need no cancellation path. Signed-off-by: Andrew Anderson <andy@clubanderson.com>
diff --git a/pkg/api/handlers/auth.go b/pkg/api/handlers/auth.go
@@ -587,6 +587,13 @@ func (h *AuthHandler) Logout(c *fiber.Ctx) error {
 	// unblocks the stream and tears the connection down promptly.
 	if claims.UserID != uuid.Nil {
 		CancelUserExecSessions(claims.UserID)
+		// Cancel any active SSE streams for this user (#6029). SSE streams
+		// run inside SetBodyStreamWriter callbacks that block for up to
+		// sseOverallDeadline (~30s); without this, a logged-out user would
+		// continue to receive cluster_data events until the deadline fires.
+		// streamClusters registers each stream's cancel func in a per-user
+		// registry on start; cancelling those funcs here ends the stream.
+		CancelUserSSEStreams(claims.UserID)
 	}
 
 	slog.Info("[Auth] token revoked, WS sessions closed", "user", claims.GitHubLogin, "jti", claims.ID)
diff --git a/pkg/api/handlers/sse.go b/pkg/api/handlers/sse.go
@@ -11,6 +11,8 @@ import (
 	"time"
 
 	"github.com/gofiber/fiber/v2"
+	"github.com/google/uuid"
+	"github.com/kubestellar/console/pkg/api/middleware"
 	"github.com/kubestellar/console/pkg/k8s"
 )
 
@@ -91,6 +93,91 @@ const sseCacheTTL = 15 * time.Second
 // to remove expired entries and prevent unbounded memory growth.
 const sseCacheEvictInterval = 30 * time.Second
 
+// sseSessionRegistry tracks active SSE streams per user so that
+// CancelUserSSEStreams can tear them down on logout (#6029).
+//
+// SSE streams run inside c.Context().SetBodyStreamWriter callbacks that block
+// until either the client disconnects or sseOverallDeadline fires. Without a
+// per-user registry, a logged-out user's in-flight streams continue emitting
+// "cluster_data" events for up to ~30s because nothing actively cancels the
+// stream context. This registry mirrors the exec session registry in exec.go:
+// when a stream's context is created, its cancel func is recorded keyed by
+// userID; on logout, CancelUserSSEStreams runs every recorded cancel for that
+// user, which causes the SetBodyStreamWriter callback to exit promptly.
+//
+// A regular sync.Mutex is used (not RWMutex) because writes (add/remove on
+// stream start/end) and reads (CancelUserSSEStreams on logout) are both
+// infrequent and always short; an RWMutex would add complexity for no gain.
+var (
+	sseSessionsMu  sync.Mutex
+	sseSessions    = make(map[uuid.UUID]map[int64]context.CancelFunc)
+	sseSessionSeq  int64 // monotonic id generator, guarded by sseSessionsMu
+)
+
+// registerSSESession records cancel under userID and returns the assigned
+// session id. The session id is used by unregisterSSESession to remove the
+// specific entry when the stream ends normally, so the map does not grow
+// unbounded across many streams by the same user.
+func registerSSESession(userID uuid.UUID, cancel context.CancelFunc) int64 {
+	sseSessionsMu.Lock()
+	defer sseSessionsMu.Unlock()
+	sseSessionSeq++
+	id := sseSessionSeq
+	sessions, ok := sseSessions[userID]
+	if !ok {
+		sessions = make(map[int64]context.CancelFunc)
+		sseSessions[userID] = sessions
+	}
+	sessions[id] = cancel
+	return id
+}
+
+// unregisterSSESession removes a single stream entry. Called from the SSE
+// handler's deferred cleanup on normal stream end so the registry stays
+// bounded by the number of concurrently live streams, not the total lifetime
+// count.
+func unregisterSSESession(userID uuid.UUID, id int64) {
+	sseSessionsMu.Lock()
+	defer sseSessionsMu.Unlock()
+	sessions, ok := sseSessions[userID]
+	if !ok {
+		return
+	}
+	delete(sessions, id)
+	if len(sessions) == 0 {
+		delete(sseSessions, userID)
+	}
+}
+
+// CancelUserSSEStreams cancels every active SSE stream belonging to the given
+// user and clears the entries from the registry. Called from the auth Logout
+// handler after revoking the JWT so that any streaming endpoint the user had
+// open stops emitting events promptly (#6029). Safe to call with a userID
+// that has no live streams.
+func CancelUserSSEStreams(userID uuid.UUID) {
+	sseSessionsMu.Lock()
+	sessions, ok := sseSessions[userID]
+	if !ok {
+		sseSessionsMu.Unlock()
+		return
+	}
+	// Take ownership of the cancel funcs under the lock, then release the
+	// lock before invoking them. Calling cancel() itself is cheap but the
+	// goroutines it unblocks may contend for other locks; holding
+	// sseSessionsMu across those is unnecessary and risks deadlock.
+	cancels := make([]context.CancelFunc, 0, len(sessions))
+	for _, c := range sessions {
+		cancels = append(cancels, c)
+	}
+	delete(sseSessions, userID)
+	sseSessionsMu.Unlock()
+
+	for _, cancel := range cancels {
+		cancel()
+	}
+	slog.Info("[SSE] cancelled SSE streams for user", "user", userID, "count", len(cancels))
+}
+
 // SSE response cache — avoids re-fetching when the user navigates away and back.
 var (
 	sseCache     = map[string]*sseCacheEntry{}
@@ -207,6 +294,11 @@ func streamClusters(
 		offline = filteredOffline
 	}
 
+	// Capture the authenticated user ID before entering the deferred
+	// SetBodyStreamWriter callback. The fiber.Ctx may be reused by the time
+	// the callback runs, so c.Locals is not safe to read inside it (#6029).
+	userID := middleware.GetUserID(c)
+
 	c.Set("Content-Type", "text/event-stream")
 	c.Set("Cache-Control", "no-cache")
 	c.Set("Connection", "keep-alive")
@@ -220,6 +312,16 @@ func streamClusters(
 		streamCtx, streamCancel := context.WithTimeout(context.Background(), sseOverallDeadline)
 		defer streamCancel()
 
+		// Register this stream's cancel with the per-user SSE session
+		// registry so a later Logout call can tear the stream down promptly
+		// instead of waiting for sseOverallDeadline (#6029). Only register
+		// when we have a real userID — in dev/demo without a valid UserID
+		// claim there is nothing to key on.
+		if userID != uuid.Nil {
+			sessionID := registerSSESession(userID, streamCancel)
+			defer unregisterSSESession(userID, sessionID)
+		}
+
 		var mu sync.Mutex
 		totalClusters := len(healthy) + len(offline)
 		completedClusters := 0
diff --git a/pkg/api/handlers/sse_test.go b/pkg/api/handlers/sse_test.go
@@ -1,12 +1,16 @@
 package handlers
 
 import (
+	"context"
 	"errors"
 	"io"
 	"net/http"
 	"strings"
+	"sync/atomic"
 	"testing"
+	"time"
 
+	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	corev1 "k8s.io/api/core/v1"
@@ -174,3 +178,128 @@ func TestStreamClusters_EmitsClusterErrorOnFailure(t *testing.T) {
 	assert.True(t, strings.Contains(body, "cluster_data"))
 	assert.True(t, strings.Contains(body, "done"))
 }
+
+// testSSECancelWait is how long tests wait for CancelUserSSEStreams to
+// propagate through the registered cancel funcs before asserting the context
+// is Done. The registry just calls cancel() synchronously, so in practice a
+// few milliseconds is enough; we use a generous budget to keep CI flake-free.
+const testSSECancelWait = 250 * time.Millisecond
+
+// TestCancelUserSSEStreams_CancelsRegisteredContexts verifies the core
+// lifecycle invariant for #6029: a stream context registered for a user is
+// Done() after CancelUserSSEStreams runs for that user, and the registry no
+// longer holds a reference to it.
+func TestCancelUserSSEStreams_CancelsRegisteredContexts(t *testing.T) {
+	userID := uuid.New()
+	ctxA, cancelA := context.WithCancel(context.Background())
+	t.Cleanup(cancelA)
+	ctxB, cancelB := context.WithCancel(context.Background())
+	t.Cleanup(cancelB)
+
+	idA := registerSSESession(userID, cancelA)
+	idB := registerSSESession(userID, cancelB)
+	require.NotEqual(t, idA, idB, "session ids must be unique within a user")
+
+	CancelUserSSEStreams(userID)
+
+	// Both contexts should see cancellation essentially immediately.
+	select {
+	case <-ctxA.Done():
+	case <-time.After(testSSECancelWait):
+		t.Fatalf("stream A context was not cancelled within %s", testSSECancelWait)
+	}
+	select {
+	case <-ctxB.Done():
+	case <-time.After(testSSECancelWait):
+		t.Fatalf("stream B context was not cancelled within %s", testSSECancelWait)
+	}
+
+	// The registry entry should be gone so it can't leak across users/logouts.
+	sseSessionsMu.Lock()
+	_, stillThere := sseSessions[userID]
+	sseSessionsMu.Unlock()
+	assert.False(t, stillThere, "registry entry for user should be cleared after cancellation")
+}
+
+// TestCancelUserSSEStreams_OtherUserUnaffected confirms that cancelling one
+// user's streams does not touch another user's streams. This is important
+// because logout of user A must not drop user B's live SSE subscriptions.
+func TestCancelUserSSEStreams_OtherUserUnaffected(t *testing.T) {
+	userA := uuid.New()
+	userB := uuid.New()
+
+	ctxA, cancelA := context.WithCancel(context.Background())
+	t.Cleanup(cancelA)
+	ctxB, cancelB := context.WithCancel(context.Background())
+	t.Cleanup(cancelB)
+
+	registerSSESession(userA, cancelA)
+	idB := registerSSESession(userB, cancelB)
+	t.Cleanup(func() { unregisterSSESession(userB, idB) })
+
+	CancelUserSSEStreams(userA)
+
+	// userA's context should be cancelled.
+	select {
+	case <-ctxA.Done():
+	case <-time.After(testSSECancelWait):
+		t.Fatalf("userA context was not cancelled")
+	}
+
+	// userB's context must still be alive — a different user logging out
+	// must not tear down unrelated SSE streams.
+	select {
+	case <-ctxB.Done():
+		t.Fatal("userB context was cancelled despite only userA logging out")
+	case <-time.After(testSSECancelWait / 2):
+		// expected: context still alive
+	}
+}
+
+// TestUnregisterSSESession_RemovesEntry verifies the deferred cleanup path
+// for normal stream end: unregisterSSESession should drop the specific
+// session id without touching sibling sessions, and drop the per-user map
+// entry when the user's last stream ends.
+func TestUnregisterSSESession_RemovesEntry(t *testing.T) {
+	userID := uuid.New()
+
+	var cancelCalledA int32
+	cancelA := func() { atomic.StoreInt32(&cancelCalledA, 1) }
+	var cancelCalledB int32
+	cancelB := func() { atomic.StoreInt32(&cancelCalledB, 1) }
+
+	idA := registerSSESession(userID, cancelA)
+	idB := registerSSESession(userID, cancelB)
+
+	unregisterSSESession(userID, idA)
+
+	// Removing one entry must not invoke any cancel funcs — cancellation is
+	// a separate concern from registry cleanup.
+	assert.Equal(t, int32(0), atomic.LoadInt32(&cancelCalledA))
+	assert.Equal(t, int32(0), atomic.LoadInt32(&cancelCalledB))
+
+	// Entry for B must still be present.
+	sseSessionsMu.Lock()
+	sessions, ok := sseSessions[userID]
+	remaining := len(sessions)
+	sseSessionsMu.Unlock()
+	require.True(t, ok)
+	assert.Equal(t, 1, remaining, "session B should still be registered")
+
+	// Removing the last entry should drop the whole per-user map slot.
+	unregisterSSESession(userID, idB)
+	sseSessionsMu.Lock()
+	_, stillThere := sseSessions[userID]
+	sseSessionsMu.Unlock()
+	assert.False(t, stillThere, "per-user map entry should be removed when empty")
+}
+
+// TestCancelUserSSEStreams_NoSessions verifies that calling the cancel
+// function for a user with no registered streams is a no-op and does not
+// panic — logout must always be safe to call whether or not the user had an
+// open SSE stream.
+func TestCancelUserSSEStreams_NoSessions(t *testing.T) {
+	userID := uuid.New()
+	// Should not panic, should not block.
+	CancelUserSSEStreams(userID)
+}
