diff --git a/config.go b/config.go index 2fc81c9..02bb0d2 100644 --- a/config.go +++ b/config.go @@ -171,6 +171,11 @@ type Config struct { // the server is going to use to listen for (and issue) shutdown // commands on. ShutdownInterceptor signal.Interceptor + + // RequestShutdown is a call-back function that can be called in order + // to indicate that pool has received a critical error and needs to shut + // down. + RequestShutdown func() } // DebugConfig is a set of debug options used for development and testing only. @@ -217,6 +222,7 @@ func DefaultConfig() Config { // "no value set". BatchVersion: -1, }, + RequestShutdown: func() {}, } } diff --git a/rpcserver.go b/rpcserver.go index c43cf1e..4d2e303 100644 --- a/rpcserver.go +++ b/rpcserver.go @@ -259,7 +259,7 @@ func (s *rpcServer) serverHandler(blockChan chan int32, blockErrChan chan error) if err != nil && !errors.Is(err, order.ErrMismatchErr) { rpcLog.Errorf("Error handling server message: "+ "%v", err) - s.server.cfg.ShutdownInterceptor.RequestShutdown() + s.server.cfg.RequestShutdown() } case err := <-s.auctioneer.StreamErrChan: @@ -286,7 +286,7 @@ func (s *rpcServer) serverHandler(blockChan chan int32, blockErrChan chan error) if err != nil { rpcLog.Errorf("Unable to receive block "+ "notification: %v", err) - s.server.cfg.ShutdownInterceptor.RequestShutdown() + s.server.cfg.RequestShutdown() } // In case the server is shutting down. @@ -2648,7 +2648,7 @@ func (s *rpcServer) StopDaemon(_ context.Context, _ *poolrpc.StopDaemonRequest) (*poolrpc.StopDaemonResponse, error) { rpcLog.Infof("Stop requested through RPC, gracefully shutting down") - s.server.cfg.ShutdownInterceptor.RequestShutdown() + s.server.cfg.RequestShutdown() return &poolrpc.StopDaemonResponse{}, nil } diff --git a/run.go b/run.go index 884c470..36b46a9 100644 --- a/run.go +++ b/run.go @@ -18,6 +18,7 @@ func Run(cfg *Config) error { if err != nil { return err } + cfg.RequestShutdown = cfg.ShutdownInterceptor.RequestShutdown logWriter = build.NewRotatingLogWriter() SetupLoggers(logWriter, cfg.ShutdownInterceptor)