main: bound GracefulStop so a hung stream cannot cost the access-state save

BUILD.bazel

@@ -205,7 +205,14 @@ oci_push(
 
 go_test(
     name = "go_default_test",
-    srcs = ["main_test.go"],
+    srcs = [
+        "main_stop_test.go",
+        "main_test.go",
+    ],
     embed = [":go_default_library"],
-    deps = ["//config:go_default_library"],
+    deps = [
+        "//config:go_default_library",
+        "@org_golang_google_grpc//:go_default_library",
+        "@org_golang_google_grpc//credentials/insecure:go_default_library",
+    ],
 )

main.go

@@ -133,7 +133,7 @@ func run(ctx *cli.Context) error {
 			if !grpcSem.TryAcquire(1) {
 				if grpcServer != nil {
 					log.Println("Stopping gRPC server")
-					grpcServer.GracefulStop()
+					stopGrpcServer(grpcServer, grpcGracefulStopTimeout)
 					log.Println("gRPC server stopped")
 				}
 			}
@@ -337,6 +337,30 @@ func run(ctx *cli.Context) error {
 	return err
 }
 
+// grpcGracefulStopTimeout bounds GracefulStop. A single stream whose client
+// has gone away without closing it would otherwise hold the drain until
+// systemd's TimeoutStopSec fires and SIGKILLs the process before the index's
+// access state is saved.
+const grpcGracefulStopTimeout = 10 * time.Second
+
+// stopGrpcServer drains in-flight RPCs for at most timeout, then closes the
+// remaining connections. Stop unblocks a pending GracefulStop, so the drain
+// goroutine always returns.
+func stopGrpcServer(s *grpc.Server, timeout time.Duration) {
+	drained := make(chan struct{})
+	go func() {
+		s.GracefulStop()
+		close(drained)
+	}()
+	select {
+	case <-drained:
+	case <-time.After(timeout):
+		log.Printf("gRPC server did not drain within %s, closing remaining connections", timeout)
+		s.Stop()
+		<-drained
+	}
+}
+
 // validateMultiBackendTrust refuses to start a multi-backend (s3_proxy
 // backends map) node whose companion trust toggles are off. The backends map
 // only makes sense on an L1 that routes per-tenant traffic, and that routing

main_stop_test.go

@@ -0,0 +1,86 @@
+package main
+
+import (
+	"context"
+	"net"
+	"testing"
+	"time"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+)
+
+// A client that opens a stream and never finishes it keeps GracefulStop
+// waiting forever; stopGrpcServer must still return once the bound expires.
+func TestStopGrpcServerBoundsAHungStream(t *testing.T) {
+	lis, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	streamOpened := make(chan struct{})
+	desc := grpc.StreamDesc{StreamName: "Hang", ServerStreams: true, ClientStreams: true}
+	svc := &grpc.ServiceDesc{
+		ServiceName: "test.Hang",
+		HandlerType: (*any)(nil),
+		Streams: []grpc.StreamDesc{{
+			StreamName:    desc.StreamName,
+			ServerStreams: true,
+			ClientStreams: true,
+			Handler: func(_ any, ss grpc.ServerStream) error {
+				close(streamOpened)
+				<-ss.Context().Done()
+				return ss.Context().Err()
+			},
+		}},
+	}
+	s := grpc.NewServer()
+	s.RegisterService(svc, struct{}{})
+	go func() { _ = s.Serve(lis) }()
+
+	conn, err := grpc.NewClient(lis.Addr().String(), grpc.WithTransportCredentials(insecure.NewCredentials()))
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer conn.Close()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	if _, err := conn.NewStream(ctx, &desc, "/test.Hang/Hang"); err != nil {
+		t.Fatal(err)
+	}
+	select {
+	case <-streamOpened:
+	case <-time.After(5 * time.Second):
+		t.Fatal("stream never reached the handler")
+	}
+
+	start := time.Now()
+	done := make(chan struct{})
+	go func() {
+		stopGrpcServer(s, 200*time.Millisecond)
+		close(done)
+	}()
+	select {
+	case <-done:
+	case <-time.After(5 * time.Second):
+		t.Fatal("stopGrpcServer did not return with a hung stream open")
+	}
+	if elapsed := time.Since(start); elapsed < 200*time.Millisecond {
+		t.Fatalf("stopGrpcServer returned after %s, before the drain bound", elapsed)
+	}
+}
+
+func TestStopGrpcServerReturnsEarlyWhenIdle(t *testing.T) {
+	lis, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	s := grpc.NewServer()
+	go func() { _ = s.Serve(lis) }()
+
+	start := time.Now()
+	stopGrpcServer(s, 10*time.Second)
+	if elapsed := time.Since(start); elapsed > 2*time.Second {
+		t.Fatalf("idle server took %s to stop", elapsed)
+	}
+}