fix: Fixes to server shutdown stream

server/cmd/main.go

@@ -2,9 +2,13 @@ package main
 
 import (
 	"context"
+	"errors"
 	"log"
+	"net/http"
 	"os"
+	"os/signal"
 	"strings"
+	"syscall"
 	"time"
 
 	"gitea.hostxtra.co.uk/mrhid6/vantage/server/internal/api"
@@ -141,13 +145,15 @@ func serve() {
 	}
 	log.Printf("message bus ready as node %s", bus.NodeID())
 
-	ctx := context.Background()
+	// Cancelled on SIGTERM/SIGINT. Everything below that takes a context — the
+	// housekeeping jobs, the leader lock — stops when the pod is asked to.
+	ctx, shutdown := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
+	defer shutdown()
 
-	go func() {
-		if err := grpcserver.StartGRPC(9090); err != nil {
-			log.Fatalf("gRPC server error: %v", err)
-		}
-	}()
+	stopGRPC, err := grpcserver.StartGRPC(9090)
+	if err != nil {
+		log.Fatalf("gRPC server error: %v", err)
+	}
 
 	// Everything below runs on exactly one replica at a time.
 	//
@@ -183,12 +189,37 @@ func serve() {
 	r.Use(corsMiddleware())
 	api.RegisterRoutes(r)
 
-	log.Println("REST server listening on :8080")
-	if err := r.Run(":8080"); err != nil {
-		log.Fatalf("REST server error: %v", err)
+	srv := &http.Server{Addr: ":8080", Handler: r}
+	go func() {
+		log.Println("REST server listening on :8080")
+		if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
+			log.Fatalf("REST server error: %v", err)
+		}
+	}()
+
+	<-ctx.Done()
+	log.Println("shutdown signal received")
+
+	// gRPC first, and this ordering is the point of the whole exercise. Stopping
+	// it runs each CommandStream handler's deferred release, which clears that
+	// agent's presence claim; until that happens another replica will keep
+	// dispatching commands to this process. Draining HTTP first would leave the
+	// claims held for the length of the drain.
+	stopGRPC()
+
+	drainCtx, cancelDrain := context.WithTimeout(context.Background(), httpDrainTimeout)
+	defer cancelDrain()
+	if err := srv.Shutdown(drainCtx); err != nil {
+		log.Printf("REST server shutdown: %v", err)
 	}
+	log.Println("shutdown complete")
 }
 
+// How long in-flight REST requests are given to finish. Console tunnels are
+// long-lived WebSockets that will not end on their own, so this is a ceiling
+// rather than a target; the relays behind them are already gone by this point.
+const httpDrainTimeout = 10 * time.Second
+
 func corsMiddleware() gin.HandlerFunc {
 	return func(c *gin.Context) {
 		c.Header("Access-Control-Allow-Origin", "*")

server/internal/grpc/server.go

@@ -228,10 +228,19 @@ func (s *vantageServer) CommandStream(stream pb.Vantage_CommandStreamServer) err
 	}
 }
 
-func StartGRPC(port int) error {
+// StartGRPC serves the agent API until stop is called.
+//
+// It returns a stop function rather than serving forever because an abrupt exit
+// is not a neutral act here: every CommandStream handler holds an agent's
+// presence claim, released by a deferred call that a killed process never runs.
+// The claim then outlives its owner for the remainder of its 30s TTL, during
+// which dispatch believes the agent is reachable, publishes to a channel with
+// no subscriber, and fails as "agent offline" — a pod that has already exited
+// still answering for an agent it can no longer reach.
+func StartGRPC(port int) (stop func(), err error) {
 	lis, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
 	if err != nil {
-		return fmt.Errorf("failed to listen: %w", err)
+		return nil, fmt.Errorf("failed to listen: %w", err)
 	}
 
 	s := grpc.NewServer(
@@ -248,6 +257,39 @@ func StartGRPC(port int) error {
 	)
 	pb.RegisterVantageServer(s, &vantageServer{})
 
-	log.Printf("gRPC server listening on :%d", port)
-	return s.Serve(lis)
+	go func() {
+		log.Printf("gRPC server listening on :%d", port)
+		if err := s.Serve(lis); err != nil {
+			log.Fatalf("gRPC server error: %v", err)
+		}
+	}()
+
+	// GracefulStop sends GOAWAY and waits for the handlers to return, which is
+	// what runs those deferred releases and, on the agent's side, ends the
+	// stream with a clean error it reconnects from immediately rather than
+	// waiting out a TCP timeout.
+	//
+	// It is bounded: an idle CommandStream returns as soon as its context is
+	// cancelled, but a console relay mid-transfer would otherwise hold the
+	// process past the pod's grace period and earn a SIGKILL — which is the
+	// abrupt exit this exists to avoid.
+	return func() {
+		done := make(chan struct{})
+		go func() {
+			s.GracefulStop()
+			close(done)
+		}()
+		select {
+		case <-done:
+			log.Println("gRPC server stopped gracefully")
+		case <-time.After(grpcStopTimeout):
+			log.Printf("gRPC server did not stop within %s, forcing", grpcStopTimeout)
+			s.Stop()
+		}
+	}, nil
 }
+
+// How long GracefulStop is given before outstanding streams are cut. Comfortably
+// inside the chart's termination grace period, so the forced stop below still
+// leaves time for the HTTP server to drain.
+const grpcStopTimeout = 10 * time.Second