Kaynağa Gözat

webapi: fix shutdown routine

Mike 1 hafta önce
ebeveyn
işleme
947ed5183d

+ 2 - 11
server/webapi/handlers/messaging_test.go

@@ -52,15 +52,7 @@ func createTestSessionManager(screenName string) (*state.WebAPISessionManager, s
 // createTestSessionManagerWithOSCAR creates a WebAPISessionManager with an OSCAR session instance set.
 func createTestSessionManagerWithOSCAR(screenName string, oscarSession *state.SessionInstance) (*state.WebAPISessionManager, string) {
 	mgr := state.NewWebAPISessionManager()
-	session, _ := mgr.CreateSession(
-		context.Background(),
-		state.DisplayScreenName(screenName),
-		"test-dev",
-		[]string{"im", "presence", "buddylist", "sentIM", "typing"},
-		oscarSession,
-		"",
-		slog.Default(),
-	)
+	session, _ := mgr.CreateSession(state.DisplayScreenName(screenName), "test-dev", []string{"im", "presence", "buddylist", "sentIM", "typing"}, oscarSession, "", slog.Default())
 	return mgr, session.AimSID
 }
 
@@ -105,8 +97,7 @@ func sendIMForDest(t *testing.T, dest, locateName, alias string) []types.Event {
 		Return(nil, nil)
 
 	mgr := state.NewWebAPISessionManager()
-	session, err := mgr.CreateSession(context.Background(), state.DisplayScreenName("Ann Dupree"),
-		"test-dev", []string{"im", "sentIM", "conversation"}, oscarInstance, "", slog.Default())
+	session, err := mgr.CreateSession(state.DisplayScreenName("Ann Dupree"), "test-dev", []string{"im", "sentIM", "conversation"}, oscarInstance, "", slog.Default())
 	require.NoError(t, err)
 
 	handler := &MessagingHandler{

+ 1 - 1
server/webapi/handlers/session.go

@@ -260,7 +260,7 @@ func (h *SessionHandler) StartSession(w http.ResponseWriter, r *http.Request) {
 	// read it.
 	baseURL := baseURLFromRequest(r)
 
-	session, err := h.SessionManager.CreateSession(r.Context(), screenName, apiKey.DevID, events, instance, baseURL, h.Logger)
+	session, err := h.SessionManager.CreateSession(screenName, apiKey.DevID, events, instance, baseURL, h.Logger)
 	if err != nil {
 		h.Logger.ErrorContext(ctx, "failed to create session", "err", err.Error())
 		h.sendError(w, r, http.StatusInternalServerError, "failed to create session")

+ 5 - 3
server/webapi/server.go

@@ -84,6 +84,8 @@ func NewServer(listeners []string, logger *slog.Logger, handler Handler, apiKeyV
 		Logger:           logger,
 	}
 
+	shutdownCtx, shutdownCancel := context.WithCancel(context.Background())
+
 	for _, l := range listeners {
 		mux := http.NewServeMux()
 
@@ -296,8 +298,6 @@ func NewServer(listeners []string, logger *slog.Logger, handler Handler, apiKeyV
 		})
 	}
 
-	shutdownCtx, shutdownCancel := context.WithCancel(context.Background())
-
 	sessionHandler.FnSessCfg = func(sess *state.Session) {
 		sess.OnSessionClose(func() {
 			ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
@@ -404,10 +404,12 @@ func (s *Server) ListenAndServe() error {
 func (s *Server) Shutdown(ctx context.Context) error {
 	s.logger.Debug("Initiating graceful shutdown...")
 	s.shutdownCancel() // stop the session reaper so ListenAndServe's errgroup can drain
+
+	s.sessionManager.Shutdown()
+
 	for _, srv := range s.servers {
 		_ = srv.Shutdown(ctx)
 	}
-	s.sessionManager.Shutdown()
 	s.logger.Info("shutdown complete")
 	return nil
 }

+ 32 - 39
server/webapi/types/events.go

@@ -2,7 +2,6 @@ package types
 
 import (
 	"context"
-	"errors"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -94,32 +93,40 @@ type TypingEvent struct {
 
 // EventQueue manages a queue of events for a WebAPI session.
 type EventQueue struct {
-	events   []Event
-	seqNum   uint64
-	maxSize  int
-	mu       sync.RWMutex
-	waitChan chan struct{}
-	closed   bool
-	closedMu sync.RWMutex
+	events    []Event
+	seqNum    uint64
+	maxSize   int
+	mu        sync.RWMutex
+	waitChan  chan struct{}
+	closeChan chan struct{}
+	closeOnce sync.Once
+}
+
+// isClosed reports whether Close has been called.
+func (q *EventQueue) isClosed() bool {
+	select {
+	case <-q.closeChan:
+		return true
+	default:
+		return false
+	}
 }
 
 // NewEventQueue creates a new event queue with the specified maximum size.
 func NewEventQueue(maxSize int) *EventQueue {
 	return &EventQueue{
-		events:   make([]Event, 0),
-		maxSize:  maxSize,
-		waitChan: make(chan struct{}, 1),
+		events:    make([]Event, 0),
+		maxSize:   maxSize,
+		waitChan:  make(chan struct{}, 1),
+		closeChan: make(chan struct{}),
 	}
 }
 
 // Push adds an event to the queue.
 func (q *EventQueue) Push(eventType EventType, data interface{}) {
-	q.closedMu.RLock()
-	if q.closed {
-		q.closedMu.RUnlock()
+	if q.isClosed() {
 		return
 	}
-	q.closedMu.RUnlock()
 
 	q.mu.Lock()
 	defer q.mu.Unlock()
@@ -153,12 +160,9 @@ func (q *EventQueue) Push(eventType EventType, data interface{}) {
 
 // Fetch retrieves events from the queue, optionally waiting for new events.
 func (q *EventQueue) Fetch(ctx context.Context, lastSeqNum uint64, timeout time.Duration) ([]Event, error) {
-	q.closedMu.RLock()
-	if q.closed {
-		q.closedMu.RUnlock()
-		return nil, errors.New("event queue is closed")
+	if q.isClosed() {
+		return []Event{}, nil
 	}
-	q.closedMu.RUnlock()
 
 	// First, check if we have any events newer than lastSeqNum
 	q.mu.RLock()
@@ -174,6 +178,9 @@ func (q *EventQueue) Fetch(ctx context.Context, lastSeqNum uint64, timeout time.
 
 	for {
 		select {
+		case <-q.closeChan:
+			return []Event{}, nil
+
 		case <-q.waitChan:
 			// New events may be available
 			q.mu.RLock()
@@ -220,24 +227,10 @@ func (q *EventQueue) GetAllEvents() []Event {
 	return result
 }
 
-// Close closes the event queue, unblocking any waiting fetchers.
+// Close closes the event queue, unblocking any waiting fetchers. Safe to call
+// more than once.
 func (q *EventQueue) Close() {
-	q.closedMu.Lock()
-	defer q.closedMu.Unlock()
-
-	if q.closed {
-		return
-	}
-
-	q.closed = true
-
-	// Send multiple signals to unblock all potential waiters
-notifyWaiters:
-	for i := 0; i < 10; i++ {
-		select {
-		case q.waitChan <- struct{}{}:
-		default:
-			break notifyWaiters
-		}
-	}
+	q.closeOnce.Do(func() {
+		close(q.closeChan)
+	})
 }

+ 46 - 15
state/webapi_session.go

@@ -87,6 +87,7 @@ type WebAPISession struct {
 	imLog        map[string][]WebAPIStoredIM
 	imLogMu      sync.Mutex
 	logger       *slog.Logger // Logger for debugging
+	listeners    sync.WaitGroup
 }
 
 // IsExpired checks if the session has expired.
@@ -172,7 +173,9 @@ func (s *WebAPISession) StartListeningToOSCARSession() {
 	}
 
 	// Start goroutine to listen for OSCAR messages
+	s.listeners.Add(1)
 	go func() {
+		defer s.listeners.Done()
 		msgCh := s.OSCARSession.ReceiveMessage()
 		for {
 			select {
@@ -190,6 +193,15 @@ func (s *WebAPISession) StartListeningToOSCARSession() {
 	}()
 }
 
+// Close tears down the session: it releases any parked event fetchers, closes
+// the OSCAR instance, and waits for the listener goroutine to unwind. Safe to
+// call more than once.
+func (s *WebAPISession) Close() {
+	s.EventQueue.Close()
+	s.OSCARSession.CloseInstance()
+	s.listeners.Wait()
+}
+
 // handleSNACMessage converts a SNAC message into WebAPI events and pushes them to the event queue.
 func (s *WebAPISession) handleSNACMessage(msg wire.SNACMessage) {
 	if s.EventQueue == nil {
@@ -482,7 +494,9 @@ func (s *WebAPISession) handleFeedbagMessage(msg wire.SNACMessage) {
 type WebAPISessionManager struct {
 	sessions map[string]*WebAPISession // Keyed by aimsid
 	mu       sync.RWMutex
-	closed   bool // set by Shutdown; rejects new sessions and makes drain idempotent
+	closed   bool           // set by Shutdown; rejects new sessions and makes drain idempotent
+	stopCh   chan struct{}  // closed by Shutdown to stop the reaper
+	reaperWG sync.WaitGroup // tracks a running reaper so Shutdown can join it
 }
 
 // NewWebAPISessionManager creates a new WebAPI session manager. It does not start
@@ -490,6 +504,7 @@ type WebAPISessionManager struct {
 func NewWebAPISessionManager() *WebAPISessionManager {
 	return &WebAPISessionManager{
 		sessions: make(map[string]*WebAPISession),
+		stopCh:   make(chan struct{}),
 	}
 }
 
@@ -500,7 +515,7 @@ func NewWebAPISessionManager() *WebAPISessionManager {
 // MyInfoRefresher, ...) and then call StartListeningToOSCARSession. Wiring them
 // after the listener starts would race the goroutine, which reads them as it
 // converts SNACs into events.
-func (m *WebAPISessionManager) CreateSession(ctx context.Context, screenName DisplayScreenName, devID string, events []string, oscarSession *SessionInstance, baseURL string, logger *slog.Logger) (*WebAPISession, error) {
+func (m *WebAPISessionManager) CreateSession(screenName DisplayScreenName, devID string, events []string, oscarSession *SessionInstance, baseURL string, logger *slog.Logger) (*WebAPISession, error) {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
@@ -574,9 +589,7 @@ func (m *WebAPISessionManager) RemoveSession(ctx context.Context, aimsid string)
 
 	// Tear down outside the lock: CloseInstance fans out to buddy-departed
 	// broadcasts and signout, which we don't want to run under m.mu.
-	session.EventQueue.Close()
-	session.OSCARSession.CloseInstance()
-
+	session.Close()
 	return nil
 }
 
@@ -594,18 +607,35 @@ func (m *WebAPISessionManager) TouchSession(ctx context.Context, aimsid string)
 	return nil
 }
 
-// Run reaps expired sessions on a fixed interval until ctx is cancelled. The
-// caller owns the goroutine's lifecycle; typically launch it under the server's
-// errgroup:
+// Run reaps expired sessions on a fixed interval until ctx is cancelled or
+// Shutdown is called. The caller owns the goroutine's lifecycle; typically
+// launch it under the server's errgroup:
 //
 //	g.Go(func() error { mgr.Run(ctx); return nil })
+//
+// Run is a no-op once the manager is closed, so a reaper that loses the race
+// with Shutdown never starts reaping a drained manager.
 func (m *WebAPISessionManager) Run(ctx context.Context) {
+	m.mu.Lock()
+	if m.closed {
+		m.mu.Unlock()
+		return
+	}
+	// Registering under m.mu is what makes Shutdown's join sound: Shutdown flips
+	// closed under the same lock, so a reaper either registers before Shutdown
+	// waits or is turned away here.
+	m.reaperWG.Add(1)
+	m.mu.Unlock()
+	defer m.reaperWG.Done()
+
 	ticker := time.NewTicker(webAPISessionReapInterval)
 	defer ticker.Stop()
 	for {
 		select {
 		case <-ticker.C:
 			m.reapExpired()
+		case <-m.stopCh:
+			return
 		case <-ctx.Done():
 			return
 		}
@@ -628,14 +658,14 @@ func (m *WebAPISessionManager) reapExpired() {
 	// Tear down outside the lock: CloseInstance fans out to buddy-departed
 	// broadcasts and signout, which we don't want to run under m.mu.
 	for _, session := range expired {
-		session.EventQueue.Close()
-		session.OSCARSession.CloseInstance()
+		session.Close()
 	}
 }
 
-// Shutdown drains and closes all sessions and blocks further CreateSession
-// calls. The reaper goroutine is stopped separately by cancelling the context
-// passed to Run. Safe to call more than once.
+// Shutdown drains and closes all sessions, stops the reaper started by Run, and
+// blocks further CreateSession calls. It does not depend on the caller
+// cancelling Run's context. Safe to call more than once, though only the first
+// call waits for the drain.
 func (m *WebAPISessionManager) Shutdown() {
 	m.mu.Lock()
 	if m.closed {
@@ -643,6 +673,7 @@ func (m *WebAPISessionManager) Shutdown() {
 		return
 	}
 	m.closed = true
+	close(m.stopCh)
 
 	sessions := make([]*WebAPISession, 0, len(m.sessions))
 	for _, session := range m.sessions {
@@ -655,9 +686,9 @@ func (m *WebAPISessionManager) Shutdown() {
 	// Tear down outside the lock: CloseInstance fans out to buddy-departed
 	// broadcasts and signout, which we don't want to run under m.mu.
 	for _, session := range sessions {
-		session.EventQueue.Close()
-		session.OSCARSession.CloseInstance()
+		session.Close()
 	}
+	m.reaperWG.Wait()
 }
 
 // generateSessionID creates a cryptographically secure session ID.

+ 93 - 10
state/webapi_session_test.go

@@ -284,10 +284,9 @@ func TestWebAPISessionManager_ShutdownIdempotent(t *testing.T) {
 func TestWebAPISessionManager_CreateAfterShutdown(t *testing.T) {
 	mgr := NewWebAPISessionManager()
 
-	ctx := context.Background()
 	mgr.Shutdown()
 
-	sess, err := mgr.CreateSession(ctx, DisplayScreenName("testuser"), "dev", []string{"presence"}, nil, "", nil)
+	sess, err := mgr.CreateSession(DisplayScreenName("testuser"), "dev", []string{"presence"}, nil, "", nil)
 	assert.Nil(t, sess)
 	assert.ErrorIs(t, err, ErrWebAPISessionManagerClosed)
 }
@@ -302,9 +301,9 @@ func TestWebAPISessionManager_ShutdownDrainsAndClosesSessions(t *testing.T) {
 	inst1 := NewSession().AddInstance()
 	inst2 := NewSession().AddInstance()
 
-	s1, err := mgr.CreateSession(ctx, DisplayScreenName("alice"), "dev", []string{"presence"}, inst1, "", slog.Default())
+	s1, err := mgr.CreateSession(DisplayScreenName("alice"), "dev", []string{"presence"}, inst1, "", slog.Default())
 	assert.NoError(t, err)
-	s2, err := mgr.CreateSession(ctx, DisplayScreenName("bob"), "dev", []string{"presence"}, inst2, "", slog.Default())
+	s2, err := mgr.CreateSession(DisplayScreenName("bob"), "dev", []string{"presence"}, inst2, "", slog.Default())
 	assert.NoError(t, err)
 
 	mgr.Shutdown()
@@ -315,8 +314,7 @@ func TestWebAPISessionManager_ShutdownDrainsAndClosesSessions(t *testing.T) {
 	// Each session's event queue and OSCAR instance were closed: the teardown
 	// loop ran for every collected session.
 	for _, s := range []*WebAPISession{s1, s2} {
-		_, err := s.EventQueue.Fetch(ctx, 0, 10*time.Millisecond)
-		assert.Error(t, err, "event queue should be closed")
+		assertQueueClosed(t, ctx, s)
 	}
 	for _, inst := range []*SessionInstance{inst1, inst2} {
 		select {
@@ -336,9 +334,9 @@ func TestWebAPISessionManager_ReapExpired(t *testing.T) {
 	expiredInst := NewSession().AddInstance()
 	liveInst := NewSession().AddInstance()
 
-	expired, err := mgr.CreateSession(ctx, "alice", "dev", []string{"presence"}, expiredInst, "", slog.Default())
+	expired, err := mgr.CreateSession("alice", "dev", []string{"presence"}, expiredInst, "", slog.Default())
 	assert.NoError(t, err)
-	live, err := mgr.CreateSession(ctx, "bob", "dev", []string{"presence"}, liveInst, "", slog.Default())
+	live, err := mgr.CreateSession("bob", "dev", []string{"presence"}, liveInst, "", slog.Default())
 	assert.NoError(t, err)
 
 	// Force alice's session into the past; bob keeps its default future expiry.
@@ -351,8 +349,7 @@ func TestWebAPISessionManager_ReapExpired(t *testing.T) {
 	assert.Contains(t, mgr.sessions, live.AimSID)
 
 	// Expired session torn down: event queue and OSCAR instance closed.
-	_, err = expired.EventQueue.Fetch(ctx, 0, 10*time.Millisecond)
-	assert.Error(t, err, "expired session's event queue should be closed")
+	assertQueueClosed(t, ctx, expired)
 	select {
 	case <-expiredInst.Closed():
 	default:
@@ -367,6 +364,92 @@ func TestWebAPISessionManager_ReapExpired(t *testing.T) {
 	}
 }
 
+// assertQueueClosed asserts the session's event queue is closed: a fetch returns
+// straight away with no events and no error, rather than parking for the timeout.
+func assertQueueClosed(t *testing.T, ctx context.Context, sess *WebAPISession) {
+	t.Helper()
+
+	const timeout = 5 * time.Second
+	start := time.Now()
+
+	events, err := sess.EventQueue.Fetch(ctx, 0, timeout)
+	assert.NoError(t, err)
+	assert.Empty(t, events)
+	assert.Less(t, time.Since(start), timeout/2, "fetch parked instead of returning on a closed queue")
+}
+
+// TestWebAPISessionManager_ShutdownWithoutReaper verifies Shutdown returns when no
+// reaper was ever started. Shutdown must not depend on the caller cancelling the
+// context passed to Run.
+func TestWebAPISessionManager_ShutdownWithoutReaper(t *testing.T) {
+	mgr := NewWebAPISessionManager()
+
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		mgr.Shutdown()
+	}()
+
+	select {
+	case <-done:
+	case <-time.After(5 * time.Second):
+		t.Fatal("Shutdown hung waiting for a reaper that was never started")
+	}
+}
+
+// TestWebAPISessionManager_ShutdownJoinsReaper verifies Shutdown stops a running
+// reaper on its own and does not return until that reaper has exited.
+func TestWebAPISessionManager_ShutdownJoinsReaper(t *testing.T) {
+	mgr := NewWebAPISessionManager()
+
+	reaperExited := make(chan struct{})
+	go func() {
+		defer close(reaperExited)
+		mgr.Run(context.Background()) // context is never cancelled: Shutdown must stop it
+	}()
+
+	// Give Run a chance to register itself before shutting down.
+	time.Sleep(50 * time.Millisecond)
+
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		mgr.Shutdown()
+	}()
+
+	select {
+	case <-done:
+	case <-time.After(5 * time.Second):
+		t.Fatal("Shutdown hung instead of stopping the reaper")
+	}
+
+	// Shutdown joins the reaper, so it has already exited by the time it returns.
+	select {
+	case <-reaperExited:
+	default:
+		t.Error("Shutdown returned before the reaper exited")
+	}
+}
+
+// TestWebAPISessionManager_RunAfterShutdown verifies a reaper that loses the race
+// with Shutdown never starts, so it cannot reap an already-drained manager.
+func TestWebAPISessionManager_RunAfterShutdown(t *testing.T) {
+	mgr := NewWebAPISessionManager()
+	mgr.Shutdown()
+
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		mgr.Run(context.Background())
+	}()
+
+	select {
+	case <-done:
+	case <-time.After(5 * time.Second):
+		t.Fatal("Run should be a no-op on a closed manager")
+	}
+}
+
 // The client deletes the alias it holds each time it merges a user map, so every
 // event naming a buddy has to repeat it. An incoming IM and a presence change both
 // carry a user map, and both would otherwise rename an aliased buddy.

+ 39 - 0
state/zz_deadlock_probe_test.go

@@ -0,0 +1,39 @@
+package state
+
+import (
+	"testing"
+	"time"
+)
+
+// Demonstrates that ScaleWarningAndRateLimit sends on warningCh while holding
+// the write lock, so a full buffer wedges every reader of the Session.
+func TestProbe_SendUnderWriteLockBlocksReaders(t *testing.T) {
+	s := NewSession()
+	s.SetIdentScreenName(NewIdentScreenName("probe"))
+
+	// 1. Fill the 1-slot warningCh buffer. Nobody is draining it.
+	s.ScaleWarningAndRateLimit(10, 3)
+
+	// 2. Next call blocks on the send -- while holding mutex.Lock().
+	blocked := make(chan struct{})
+	go func() {
+		close(blocked)
+		s.ScaleWarningAndRateLimit(10, 3)
+	}()
+	<-blocked
+	time.Sleep(200 * time.Millisecond) // let it reach the send
+
+	// 3. Any reader now blocks forever. This is icbm.go:719 TLVUserInfo().
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		s.DisplayScreenName()
+	}()
+
+	select {
+	case <-done:
+		t.Log("OK: reader acquired RLock")
+	case <-time.After(2 * time.Second):
+		t.Fatal("DEADLOCK: reader blocked on RLock because the writer is parked on a channel send")
+	}
+}