diff --git a/server/serve_stop_test.go b/server/serve_stop_test.go new file mode 100644 index 0000000000..8e990b125a --- /dev/null +++ b/server/serve_stop_test.go @@ -0,0 +1,112 @@ +package server + +import ( + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestStopBeforeServe verifies that Stop() is safe to call before Serve() +// and does not panic. +func TestStopBeforeServe(t *testing.T) { + srv, err := NewServer() + require.NoError(t, err) + // Should not panic when server has not been started + assert.NotPanics(t, func() { + srv.Stop() + }) +} + +// TestStopIdempotent verifies that calling Stop() multiple times is safe. +func TestStopIdempotent(t *testing.T) { + srv, err := NewServer() + require.NoError(t, err) + + // Start Serve in a goroutine; it will block on stopCh. + serveErrCh := make(chan error, 1) + go func() { + serveErrCh <- srv.Serve() + }() + + // Give Serve() time to reach the blocking phase. + time.Sleep(200 * time.Millisecond) + + // First Stop should succeed. + assert.NotPanics(t, func() { + srv.Stop() + }) + + // Wait for Serve to return. + select { + case err := <-serveErrCh: + assert.NoError(t, err) + case <-time.After(2 * time.Second): + t.Fatal("Serve() did not return after Stop()") + } + + // Second Stop should be safe (no-op). + assert.NotPanics(t, func() { + srv.Stop() + }) +} + +// TestServeStopGracefulShutdown verifies that Stop() causes Serve() to return gracefully. +func TestServeStopGracefulShutdown(t *testing.T) { + srv, err := NewServer() + require.NoError(t, err) + + serveErrCh := make(chan error, 1) + go func() { + serveErrCh <- srv.Serve() + }() + + // Allow Serve() to proceed past initialization and reach the blocking phase. + time.Sleep(200 * time.Millisecond) + + // Trigger shutdown. + srv.Stop() + + // Serve should return without error. + select { + case err := <-serveErrCh: + assert.NoError(t, err) + case <-time.After(2 * time.Second): + t.Fatal("Serve() did not return after Stop()") + } +} + +// TestServeAfterStop verifies that Serve() can be called again after Stop(). +func TestServeAfterStop(t *testing.T) { + srv, err := NewServer() + require.NoError(t, err) + + // First cycle + serveErrCh1 := make(chan error, 1) + go func() { + serveErrCh1 <- srv.Serve() + }() + time.Sleep(200 * time.Millisecond) + srv.Stop() + select { + case err := <-serveErrCh1: + assert.NoError(t, err) + case <-time.After(2 * time.Second): + t.Fatal("first Serve() did not return after Stop()") + } + + // Second cycle + serveErrCh2 := make(chan error, 1) + go func() { + serveErrCh2 <- srv.Serve() + }() + time.Sleep(200 * time.Millisecond) + srv.Stop() + select { + case err := <-serveErrCh2: + assert.NoError(t, err) + case <-time.After(2 * time.Second): + t.Fatal("second Serve() did not return after Stop()") + } +} diff --git a/server/server.go b/server/server.go index cff53f39f1..8a65af606d 100644 --- a/server/server.go +++ b/server/server.go @@ -57,6 +57,9 @@ type Server struct { interfaceNameServices map[string]*ServiceOptions // indicate whether the server is already started serve bool + // stopCh is used to signal graceful shutdown of Serve(). + // Closing this channel causes Serve() to return instead of blocking forever with select{}. + stopCh chan struct{} } // ServiceInfo Deprecated: common.ServiceInfo type alias, just for compatible with old generate pb.go file @@ -323,6 +326,13 @@ func (s *Server) Serve() error { } // prevent multiple calls to Serve s.serve = true + // re-create stopCh if it was closed by a previous Stop(), + // allowing Serve() to be called again after graceful shutdown. + select { + case <-s.stopCh: + s.stopCh = make(chan struct{}) + default: + } // release lock in case causing deadlock s.mu.Unlock() @@ -356,7 +366,24 @@ func (s *Server) Serve() error { probe.SetStartupComplete(true) probe.SetReady(true) - select {} + // Block until Stop() is called or the stopCh is closed, + // enabling graceful shutdown instead of an unrecoverable hard spin. + <-s.stopCh + return nil +} + +// Stop signals the server to shut down gracefully. +// It is safe to call multiple times; only the first call has effect. +// After Stop, Serve() will return and the server can be started again +// by calling Serve() (if desired). +func (s *Server) Stop() { + s.mu.Lock() + defer s.mu.Unlock() + if !s.serve { + return + } + s.serve = false + close(s.stopCh) } // In order to expose internal services @@ -463,6 +490,7 @@ func NewServer(opts ...ServerOption) (*Server, error) { cfg: newSrvOpts, svcOptsMap: make(map[string]*ServiceOptions), interfaceNameServices: make(map[string]*ServiceOptions), + stopCh: make(chan struct{}), } return srv, nil }