fix(chat): treat server shutdown as clean stop

pkgs/chat/main.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"log"
 	"os"
@@ -59,22 +60,40 @@ func newApp() *app {
 	return a
 }
 
+func runServer(serve func() error, shutdown func(context.Context) error, done <-chan os.Signal, timeout time.Duration) error {
+	serveErr := make(chan error, 1)
+	go func() {
+		serveErr <- serve()
+	}()
+
+	select {
+	case err := <-serveErr:
+		if err != nil && !errors.Is(err, ssh.ErrServerClosed) {
+			return err
+		}
+		return nil
+	case <-done:
+		log.Println("Stopping SSH server")
+		ctx, cancel := context.WithTimeout(context.Background(), timeout)
+		defer cancel()
+		if err := shutdown(ctx); err != nil {
+			return err
+		}
+
+		err := <-serveErr
+		if err != nil && !errors.Is(err, ssh.ErrServerClosed) {
+			return err
+		}
+		return nil
+	}
+}
+
 func (a *app) Start() {
-	var err error
 	done := make(chan os.Signal, 1)
 	signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
+	defer signal.Stop(done)
 	log.Printf("Starting SSH server on %s:%d", host, port)
-	go func() {
-		if err = a.ListenAndServe(); err != nil {
-			log.Fatalln(err)
-		}
-	}()
-
-	<-done
-	log.Println("Stopping SSH server")
-	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
-	defer func() { cancel() }()
-	if err := a.Shutdown(ctx); err != nil {
+	if err := runServer(a.ListenAndServe, a.Shutdown, done, 30*time.Second); err != nil {
 		log.Fatalln(err)
 	}
 }

pkgs/chat/main_test.go

@@ -0,0 +1,60 @@
+package main
+
+import (
+	"context"
+	"errors"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/charmbracelet/ssh"
+)
+
+func TestRunServerTreatsShutdownAsCleanStop(t *testing.T) {
+	signals := make(chan os.Signal, 1)
+	signals <- os.Interrupt
+
+	serveRelease := make(chan struct{})
+	shutdownCalled := false
+	err := runServer(
+		func() error {
+			<-serveRelease
+			return ssh.ErrServerClosed
+		},
+		func(context.Context) error {
+			shutdownCalled = true
+			close(serveRelease)
+			return nil
+		},
+		signals,
+		time.Second,
+	)
+
+	if err != nil {
+		t.Fatalf("runServer() error = %v, want nil", err)
+	}
+	if !shutdownCalled {
+		t.Fatal("runServer() did not shut down after receiving a signal")
+	}
+}
+
+func TestRunServerReturnsServeFailure(t *testing.T) {
+	serveErr := errors.New("listen tcp 127.0.0.1:23234: bind: address already in use")
+	shutdownCalled := false
+	err := runServer(
+		func() error { return serveErr },
+		func(context.Context) error {
+			shutdownCalled = true
+			return nil
+		},
+		make(chan os.Signal),
+		time.Second,
+	)
+
+	if !errors.Is(err, serveErr) {
+		t.Fatalf("runServer() error = %v, want %v", err, serveErr)
+	}
+	if shutdownCalled {
+		t.Fatal("runServer() called shutdown after a serve failure")
+	}
+}