diff --git a/BUILD.bazel b/BUILD.bazel index b0b6dfa..0d21bcc 100644 --- a/BUILD.bazel +++ b/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", + ], ) diff --git a/main.go b/main.go index b5827e7..c84b808 100644 --- a/main.go +++ b/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 diff --git a/main_stop_test.go b/main_stop_test.go new file mode 100644 index 0000000..9113d9c --- /dev/null +++ b/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) + } +}