Wait for IRC listener cleanup and HTTP shutdown on stop (closes #159)

Stopping the server returned before the IRC listener's connections had
run their cleanup, and the database closed under them, so their
sessions, clients and channel memberships stayed and their nicks stayed
taken after a restart.

The IRC listener's stop now waits up to five seconds for every
connection's cleanup, which now runs on a context that stopping does
not cancel. The HTTP server's stop hook now runs Shutdown itself and
waits for it; httpServer is set before serving starts, and the server's
own signal watcher, which raced fx's stop hooks, is gone.

Model: opus-5-5

internal/ircserver/conn.go

@@ -222,7 +222,9 @@ func (c *Conn) serve(
 ) {
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
-	defer c.cleanup(ctx)
+	// Stopping the server cancels ctx, and the cleanup must
+	// still end the session then.
+	defer c.cleanup(context.WithoutCancel(ctx))
 
 	ping := time.AfterFunc(interval, func() {
 		c.send("PING :" + c.serverSfx)

internal/ircserver/server.go

@@ -16,6 +16,10 @@ import (
 	"sneak.berlin/go/neoirc/internal/service"
 )
 
+// stopTimeout is how long stop waits for the connections'
+// cleanup, which ends their sessions, before it returns.
+const stopTimeout = 5 * time.Second
+
 // Params defines the dependencies for creating an IRC
 // Server.
 type Params struct {
@@ -40,6 +44,10 @@ type Server struct {
 	conns    map[*Conn]struct{}
 	cancel   context.CancelFunc
 
+	// serving counts the connections whose serve, cleanup
+	// included, has not returned yet.
+	serving sync.WaitGroup
+
 	// pingInterval and pongDeadline are the constants of the
 	// same names, except in tests, which shorten them.
 	pingInterval time.Duration
@@ -63,6 +71,7 @@ func New(
 		listener: nil,
 		cancel:   nil,
 		mu:       sync.Mutex{},
+		serving:  sync.WaitGroup{},
 
 		pingInterval: pingInterval,
 		pongDeadline: pongDeadline,
@@ -112,7 +121,10 @@ func (s *Server) start(ctx context.Context, addr string) error {
 	return nil
 }
 
-// stop shuts down the listener and all connections.
+// stop shuts down the listener and all connections, then
+// waits up to stopTimeout for every connection's cleanup to
+// finish, so that their sessions are ended before the
+// database closes.
 func (s *Server) stop() {
 	if s.cancel != nil {
 		s.cancel()
@@ -127,6 +139,22 @@ func (s *Server) stop() {
 		c.conn.Close() //nolint:errcheck,gosec
 	}
 	s.mu.Unlock()
+
+	done := make(chan struct{})
+
+	go func() {
+		s.serving.Wait()
+		close(done)
+	}()
+
+	select {
+	case <-done:
+	case <-time.After(stopTimeout):
+		s.log.Error(
+			"irc connections not cleaned up before stop timeout",
+			"timeout", stopTimeout,
+		)
+	}
 }
 
 // acceptLoop accepts new connections.
@@ -152,10 +180,24 @@ func (s *Server) acceptLoop(ctx context.Context) {
 		)
 
 		s.mu.Lock()
+
+		// stop cancels ctx before it takes s.mu to close the
+		// connections, so a connection accepted while stop
+		// runs is either closed and waited for by stop or
+		// closed here.
+		if ctx.Err() != nil {
+			s.mu.Unlock()
+			tcpConn.Close() //nolint:errcheck,gosec
+
+			return
+		}
+
 		s.conns[client] = struct{}{}
+		s.serving.Add(1)
 		s.mu.Unlock()
 
 		go func() {
+			defer s.serving.Done()
 			defer func() {
 				s.mu.Lock()
 				delete(s.conns, client)

internal/ircserver/server_test.go

@@ -4,6 +4,7 @@ import (
 	"bufio"
 	"crypto/rand"
 	"database/sql"
+	"errors"
 	"fmt"
 	"log/slog"
 	"net"
@@ -812,6 +813,79 @@ func TestQuit(t *testing.T) {
 	assertContains(t, lines, "goodbye", "QUIT reason")
 }
 
+// TestStopEndsSessions checks that stopping the server ends
+// a registered connection's session before Stop returns, as
+// closing the connection does: the session, its client and
+// its channel membership are gone from the database, and a
+// channel peer has been sent the QUIT.
+func TestStopEndsSessions(t *testing.T) {
+	t.Parallel()
+
+	env := newTestEnv(t)
+	client := env.dial(t)
+
+	client.register("shutirc")
+	client.joinAndDrain("#shut")
+
+	ctx := t.Context()
+
+	// A peer on the HTTP API, whose session outlives the stop.
+	peerID, peerClientID, _, err := env.database.CreateSession(
+		ctx, "peer", "peer", "", "",
+	)
+	if err != nil {
+		t.Fatalf("create peer session: %v", err)
+	}
+
+	channelID, err := env.database.GetChannelByName(ctx, "#shut")
+	if err != nil {
+		t.Fatalf("get channel: %v", err)
+	}
+
+	err = env.database.JoinChannel(ctx, channelID, peerID)
+	if err != nil {
+		t.Fatalf("join peer: %v", err)
+	}
+
+	env.srv.Stop()
+
+	_, err = env.database.GetSessionByNick(ctx, "shutirc")
+	if !errors.Is(err, sql.ErrNoRows) {
+		t.Errorf("look up shutirc after stop: %v, want no rows", err)
+	}
+
+	clients, err := env.database.GetClientCount(ctx)
+	if err != nil {
+		t.Fatalf("count clients: %v", err)
+	}
+
+	if clients != 1 {
+		t.Errorf("%d clients after stop, want only the peer's", clients)
+	}
+
+	members, err := env.database.CountChannelMembers(ctx, channelID)
+	if err != nil {
+		t.Fatalf("count members: %v", err)
+	}
+
+	if members != 1 {
+		t.Errorf("%d members of #shut after stop, want only the peer", members)
+	}
+
+	msgs, _, err := env.database.PollMessages(ctx, peerClientID, 0, 0)
+	if err != nil {
+		t.Fatalf("poll peer: %v", err)
+	}
+
+	for _, msg := range msgs {
+		if msg.Command == "QUIT" && msg.From == "shutirc" {
+			return
+		}
+	}
+
+	t.Errorf("peer was not sent the QUIT of shutirc: %v", msgs)
+}
+
 func TestTopicSetAndGet(t *testing.T) {
 	t.Parallel()
 

internal/server/export_test.go

@@ -1,6 +1,13 @@
 package server
 
-import "github.com/go-chi/chi/v5"
+import (
+	"context"
+	"log/slog"
+	"net/http"
+
+	"github.com/go-chi/chi/v5"
+	"sneak.berlin/go/neoirc/internal/config"
+)
 
 // NewTestSPAServer creates a Server whose router serves only
 // the embedded web client.
@@ -10,3 +17,28 @@ func NewTestSPAServer() *Server {
 
 	return srv
 }
+
+// NewTestHTTPServer creates a Server whose router serves
+// handler at / and whose HTTP server listens on port once
+// Start is called.
+func NewTestHTTPServer(port int, handler http.HandlerFunc) *Server {
+	router := chi.NewRouter()
+	router.Get("/", handler)
+
+	return &Server{
+		log:    slog.New(slog.DiscardHandler),
+		router: router,
+		params: Params{Config: &config.Config{Port: port}},
+	}
+}
+
+// Start starts the HTTP server, as the start hook does once
+// the router is built.
+func (srv *Server) Start() {
+	srv.startHTTP()
+}
+
+// Stop does what the stop hook does.
+func (srv *Server) Stop(ctx context.Context) {
+	srv.cleanShutdown(ctx)
+}

internal/server/server.go

@@ -45,11 +45,8 @@ type Params struct {
 // It manages routing, middleware, and lifecycle.
 type Server struct {
 	startupTime   time.Time
-	exitCode      int
 	sentryEnabled bool
 	log           *slog.Logger
-	ctx           context.Context //nolint:containedctx // signal handling pattern
-	cancelFunc    context.CancelFunc
 	httpServer    *http.Server
 	router        *chi.Mux
 	params        Params
@@ -78,13 +75,16 @@ func New(
 			srv.enableSentry()
 			srv.SetupRoutes()
 
-			// The start hook's context ends when the hook
-			// returns; serving must outlive it.
-			go srv.serve() //nolint:contextcheck,gosec // G118
+			// A closed stdout or stderr must not end the process.
+			signal.Ignore(syscall.SIGPIPE)
+
+			srv.startHTTP()
 
 			return nil
 		},
-		OnStop: func(_ context.Context) error {
+		OnStop: func(ctx context.Context) error {
+			srv.cleanShutdown(ctx)
+
 			return nil
 		},
 	})
@@ -129,44 +129,15 @@ func (srv *Server) enableSentry() {
 	srv.sentryEnabled = true
 }
 
-func (srv *Server) serve() int {
-	srv.ctx, srv.cancelFunc = context.WithCancel(
-		context.Background(),
-	)
-
-	go func() {
-		sigCh := make(chan os.Signal, 1)
-
-		signal.Ignore(syscall.SIGPIPE)
-		signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM)
-
-		sig := <-sigCh
-
-		srv.log.Info("signal received", "signal", sig)
-
-		if srv.cancelFunc != nil {
-			srv.cancelFunc()
-		}
-	}()
-
-	go srv.serveUntilShutdown()
-
-	<-srv.ctx.Done()
-
-	srv.cleanShutdown()
-
-	return srv.exitCode
-}
-
 func (srv *Server) cleanupForExit() {
 	srv.log.Info("cleaning up")
 }
 
-func (srv *Server) cleanShutdown() {
-	srv.exitCode = 0
-
+// cleanShutdown stops the HTTP server, waiting up to
+// shutdownTimeout for the requests in progress to finish.
+func (srv *Server) cleanShutdown(ctx context.Context) {
 	ctxShutdown, shutdownCancel := context.WithTimeout(
-		context.Background(), shutdownTimeout,
+		ctx, shutdownTimeout,
 	)
 
 	err := srv.httpServer.Shutdown(ctxShutdown)
@@ -191,7 +162,11 @@ func (srv *Server) configure() {
 	// Server configuration placeholder.
 }
 
-func (srv *Server) serveUntilShutdown() {
+// startHTTP creates the HTTP server and serves it from a new
+// goroutine. srv.httpServer is set before that goroutine
+// starts and never changes, so it and cleanShutdown read it
+// without a lock.
+func (srv *Server) startHTTP() {
 	listenAddr := fmt.Sprintf(
 		":%d", srv.params.Config.Port,
 	)
@@ -204,16 +179,16 @@ func (srv *Server) serveUntilShutdown() {
 		Handler:        srv,
 	}
 
+	go srv.serveUntilShutdown()
+}
+
+func (srv *Server) serveUntilShutdown() {
 	srv.log.Info(
-		"http begin listen", "listenaddr", listenAddr,
+		"http begin listen", "listenaddr", srv.httpServer.Addr,
 	)
 
 	err := srv.httpServer.ListenAndServe()
 	if err != nil && !errors.Is(err, http.ErrServerClosed) {
 		srv.log.Error("listen error", "error", err)
-
-		if srv.cancelFunc != nil {
-			srv.cancelFunc()
-		}
 	}
 }

internal/server/server_test.go

@@ -0,0 +1,162 @@
+package server_test
+
+import (
+	"net"
+	"net/http"
+	"strconv"
+	"testing"
+	"time"
+
+	"sneak.berlin/go/neoirc/internal/server"
+)
+
+const (
+	testTimeout = 5 * time.Second
+
+	// listenCheckStep is how often waitUntilListening tries
+	// to connect.
+	listenCheckStep = 10 * time.Millisecond
+
+	// stopHoldTime is how long a test waits to see that Stop
+	// has not returned.
+	stopHoldTime = 200 * time.Millisecond
+)
+
+// TestStopWaitsForRequestInProgress checks that stopping the
+// server waits for a request it is serving to finish, and
+// that the request gets its response.
+func TestStopWaitsForRequestInProgress(t *testing.T) {
+	t.Parallel()
+
+	started := make(chan struct{})
+	release := make(chan struct{})
+
+	port := freePort(t)
+	addr := net.JoinHostPort("127.0.0.1", strconv.Itoa(port))
+
+	srv := server.NewTestHTTPServer(
+		port,
+		func(writer http.ResponseWriter, _ *http.Request) {
+			close(started)
+			<-release
+			writer.WriteHeader(http.StatusNoContent)
+		},
+	)
+	srv.Start()
+
+	waitUntilListening(t, addr)
+
+	statuses := make(chan int, 1)
+
+	go func() {
+		statuses <- get(t, "http://"+addr+"/")
+	}()
+
+	<-started
+
+	stopped := make(chan struct{})
+
+	go func() {
+		srv.Stop(t.Context())
+		close(stopped)
+	}()
+
+	select {
+	case <-stopped:
+		t.Error("Stop returned while a request was in progress")
+	case <-time.After(stopHoldTime):
+	}
+
+	close(release)
+
+	status := <-statuses
+	if status != http.StatusNoContent {
+		t.Errorf("status %d, want %d", status, http.StatusNoContent)
+	}
+
+	select {
+	case <-stopped:
+	case <-time.After(testTimeout):
+		t.Fatal("Stop did not return after the request finished")
+	}
+}
+
+// freePort returns a TCP port that nothing was listening on a
+// moment ago.
+func freePort(t *testing.T) int {
+	t.Helper()
+
+	var listenConfig net.ListenConfig
+
+	listener, err := listenConfig.Listen(
+		t.Context(), "tcp", "127.0.0.1:0",
+	)
+	if err != nil {
+		t.Fatalf("listen: %v", err)
+	}
+
+	tcpAddr, ok := listener.Addr().(*net.TCPAddr)
+	if !ok {
+		t.Fatalf("listener address %v is not TCP", listener.Addr())
+	}
+
+	err = listener.Close()
+	if err != nil {
+		t.Fatalf("close listener: %v", err)
+	}
+
+	return tcpAddr.Port
+}
+
+// waitUntilListening waits until addr accepts connections,
+// and fails the test if it does not within testTimeout.
+func waitUntilListening(t *testing.T, addr string) {
+	t.Helper()
+
+	var dialer net.Dialer
+
+	deadline := time.Now().Add(testTimeout)
+
+	for {
+		conn, err := dialer.DialContext(t.Context(), "tcp", addr)
+		if err == nil {
+			_ = conn.Close()
+
+			return
+		}
+
+		if time.Now().After(deadline) {
+			t.Fatalf("nothing listening on %s: %v", addr, err)
+		}
+
+		time.Sleep(listenCheckStep)
+	}
+}
+
+// get sends a GET request to url and returns the response
+// status, or 0 if there was no response.
+func get(t *testing.T, url string) int {
+	t.Helper()
+
+	request, err := http.NewRequestWithContext(
+		t.Context(), http.MethodGet, url, nil,
+	)
+	if err != nil {
+		t.Errorf("new request: %v", err)
+
+		return 0
+	}
+
+	client := http.Client{Timeout: testTimeout}
+
+	response, err := client.Do(request)
+	if err != nil {
+		t.Errorf("get %s: %v", url, err)
+
+		return 0
+	}
+
+	_ = response.Body.Close()
+
+	return response.StatusCode
+}