grpc: add shutdown call for external management (#957)

* grpc: add shutdown call for external management

* go mod: tidy

* glorious: suggestion

* Update engine/engine.go

Co-authored-by: Adrian Gallagher <adrian.gallagher@thrasher.io>

* Update engine/rpcserver.go

Co-authored-by: Adrian Gallagher <adrian.gallagher@thrasher.io>

* Update main.go

Co-authored-by: Adrian Gallagher <adrian.gallagher@thrasher.io>

* Update engine/rpcserver.go

Co-authored-by: Scott <gloriousCode@users.noreply.github.com>

Co-authored-by: Ryan O'Hara-Reid <ryan.oharareid@thrasher.io>
Co-authored-by: Adrian Gallagher <adrian.gallagher@thrasher.io>
Co-authored-by: Scott <gloriousCode@users.noreply.github.com>

engine/engine.go

@@ -48,6 +48,7 @@ type Engine struct {
 	currencyStateManager    *CurrencyStateManager
 	Settings                Settings
 	uptime                  time.Time
+	GRPCShutdownSignal      chan struct{}
 	ServicesWG              sync.WaitGroup
 }
 
@@ -181,6 +182,13 @@ func validateSettings(b *Engine, s *Settings, flagSet FlagSet) {
 
 	flagSet.WithBool("grpc", &b.Settings.EnableGRPC, b.Config.RemoteControl.GRPC.Enabled)
 	flagSet.WithBool("grpcproxy", &b.Settings.EnableGRPCProxy, b.Config.RemoteControl.GRPC.GRPCProxyEnabled)
+
+	flagSet.WithBool("grpcshutdown", &b.Settings.EnableGRPCShutdown, b.Config.RemoteControl.GRPC.GRPCAllowBotShutdown)
+	if b.Settings.EnableGRPCShutdown {
+		b.GRPCShutdownSignal = make(chan struct{})
+		go b.waitForGPRCShutdown()
+	}
+
 	flagSet.WithBool("websocketrpc", &b.Settings.EnableWebsocketRPC, b.Config.RemoteControl.WebsocketRPC.Enabled)
 	flagSet.WithBool("deprecatedrpc", &b.Settings.EnableDeprecatedRPC, b.Config.RemoteControl.DeprecatedRPC.Enabled)
 
@@ -260,6 +268,7 @@ func PrintSettings(s *Settings) {
 	gctlog.Debugf(gctlog.Global, "\t Portfolio manager sleep delay: %v\n", s.PortfolioManagerDelay)
 	gctlog.Debugf(gctlog.Global, "\t Enable gPRC: %v", s.EnableGRPC)
 	gctlog.Debugf(gctlog.Global, "\t Enable gRPC Proxy: %v", s.EnableGRPCProxy)
+	gctlog.Debugf(gctlog.Global, "\t Enable gRPC shutdown of bot instance: %v", s.EnableGRPCShutdown)
 	gctlog.Debugf(gctlog.Global, "\t Enable websocket RPC: %v", s.EnableWebsocketRPC)
 	gctlog.Debugf(gctlog.Global, "\t Enable deprecated RPC: %v", s.EnableDeprecatedRPC)
 	gctlog.Debugf(gctlog.Global, "\t Enable comms relayer: %v", s.EnableCommsRelayer)
@@ -968,3 +977,11 @@ func (bot *Engine) SetDefaultWebsocketDataHandler() error {
 	}
 	return bot.websocketRoutineManager.setWebsocketDataHandler(bot.websocketRoutineManager.websocketDataHandler)
 }
+
+// waitForGPRCShutdown routines waits for a signal from the grpc server to
+// send a shutdown signal.
+func (bot *Engine) waitForGPRCShutdown() {
+	<-bot.GRPCShutdownSignal
+	gctlog.Warnln(gctlog.Global, "Captured gRPC shutdown request.")
+	bot.Settings.Shutdown <- struct{}{}
+}

engine/engine_types.go

@@ -24,6 +24,7 @@ type Settings struct {
 	PortfolioManagerDelay       time.Duration
 	EnableGRPC                  bool
 	EnableGRPCProxy             bool
+	EnableGRPCShutdown          bool
 	EnableWebsocketRPC          bool
 	EnableDeprecatedRPC         bool
 	EnableCommsRelayer          bool
@@ -90,6 +91,9 @@ type Settings struct {
 
 	// Withdraw settings
 	WithdrawCacheSize uint64
+
+	// Main shutdown channel
+	Shutdown chan struct{}
 }
 
 const (

engine/rpcserver.go

@@ -53,22 +53,24 @@ import (
 )
 
 var (
-	errExchangeNotLoaded    = errors.New("exchange is not loaded/doesn't exist")
-	errExchangeNotEnabled   = errors.New("exchange is not enabled")
-	errExchangeBaseNotFound = errors.New("cannot get exchange base")
-	errInvalidArguments     = errors.New("invalid arguments received")
-	errExchangeNameUnset    = errors.New("exchange name unset")
-	errCurrencyPairUnset    = errors.New("currency pair unset")
-	errInvalidTimes         = errors.New("invalid start and end times")
-	errAssetTypeDisabled    = errors.New("asset type is disabled")
-	errAssetTypeUnset       = errors.New("asset type unset")
-	errDispatchSystem       = errors.New("dispatch system offline")
-	errCurrencyNotEnabled   = errors.New("currency not enabled")
-	errCurrencyNotSpecified = errors.New("a currency must be specified")
-	errCurrencyPairInvalid  = errors.New("currency provided is not found in the available pairs list")
-	errNoTrades             = errors.New("no trades returned from supplied params")
-	errNilRequestData       = errors.New("nil request data received, cannot continue")
-	errNoAccountInformation = errors.New("account information does not exist")
+	errExchangeNotLoaded       = errors.New("exchange is not loaded/doesn't exist")
+	errExchangeNotEnabled      = errors.New("exchange is not enabled")
+	errExchangeBaseNotFound    = errors.New("cannot get exchange base")
+	errInvalidArguments        = errors.New("invalid arguments received")
+	errExchangeNameUnset       = errors.New("exchange name unset")
+	errCurrencyPairUnset       = errors.New("currency pair unset")
+	errInvalidTimes            = errors.New("invalid start and end times")
+	errAssetTypeDisabled       = errors.New("asset type is disabled")
+	errAssetTypeUnset          = errors.New("asset type unset")
+	errDispatchSystem          = errors.New("dispatch system offline")
+	errCurrencyNotEnabled      = errors.New("currency not enabled")
+	errCurrencyNotSpecified    = errors.New("a currency must be specified")
+	errCurrencyPairInvalid     = errors.New("currency provided is not found in the available pairs list")
+	errNoTrades                = errors.New("no trades returned from supplied params")
+	errNilRequestData          = errors.New("nil request data received, cannot continue")
+	errNoAccountInformation    = errors.New("account information does not exist")
+	errShutdownNotAllowed      = errors.New("shutting down this bot instance is not allowed via gRPC, please enable by command line flag --grpcshutdown or config.json field grpcAllowBotShutdown")
+	errGRPCShutdownSignalIsNil = errors.New("cannot shutdown, gRPC shutdown channel is nil")
 )
 
 // RPCServer struct
@@ -4561,3 +4563,18 @@ func (s *RPCServer) GetCollateral(ctx context.Context, r *gctrpc.GetCollateralRe
 	}
 	return result, nil
 }
+
+// Shutdown terminates bot session externally
+func (s *RPCServer) Shutdown(_ context.Context, _ *gctrpc.ShutdownRequest) (*gctrpc.ShutdownResponse, error) {
+	if !s.Engine.Settings.EnableGRPCShutdown {
+		return nil, errShutdownNotAllowed
+	}
+
+	if s.Engine.GRPCShutdownSignal == nil {
+		return nil, errGRPCShutdownSignalIsNil
+	}
+
+	s.Engine.GRPCShutdownSignal <- struct{}{}
+	s.Engine.GRPCShutdownSignal = nil
+	return &gctrpc.ShutdownResponse{}, nil
+}

engine/rpcserver_test.go

@@ -2304,3 +2304,24 @@ func TestGetCollateral(t *testing.T) {
 		t.Errorf("received '%v', expected '%v'", err, nil)
 	}
 }
+
+func TestShutdown(t *testing.T) {
+	t.Parallel()
+	s := RPCServer{Engine: &Engine{}}
+	_, err := s.Shutdown(context.Background(), &gctrpc.ShutdownRequest{})
+	if !errors.Is(err, errShutdownNotAllowed) {
+		t.Fatalf("received: '%v' but expected: '%v'", err, errShutdownNotAllowed)
+	}
+
+	s.Engine.Settings.EnableGRPCShutdown = true
+	_, err = s.Shutdown(context.Background(), &gctrpc.ShutdownRequest{})
+	if !errors.Is(err, errGRPCShutdownSignalIsNil) {
+		t.Fatalf("received: '%v' but expected: '%v'", err, errGRPCShutdownSignalIsNil)
+	}
+
+	s.Engine.GRPCShutdownSignal = make(chan struct{}, 1)
+	_, err = s.Shutdown(context.Background(), &gctrpc.ShutdownRequest{})
+	if !errors.Is(err, nil) {
+		t.Fatalf("received: '%v' but expected: '%v'", err, nil)
+	}
+}