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

* grpc: add shutdown call for external management * go mod: tidy * glorious: suggestion * Update engine/engine.go Co-authored-by: Adrian Gallagher <[email protected]> * Update engine/rpcserver.go Co-authored-by: Adrian Gallagher <[email protected]> * Update main.go Co-authored-by: Adrian Gallagher <[email protected]> * Update engine/rpcserver.go Co-authored-by: Scott <[email protected]> Co-authored-by: Ryan O'Hara-Reid <[email protected]> Co-authored-by: Adrian Gallagher <[email protected]> Co-authored-by: Scott <[email protected]>
cranktakular · May 20, 2022 · ee9c35d · ee9c35d
1 parent c492e66
commit ee9c35d
Show file tree

Hide file tree

Showing 13 changed files with 1,189 additions and 848 deletions.
diff --git a/cmd/gctcli/commands.go b/cmd/gctcli/commands.go
@@ -5037,3 +5037,26 @@ func getCollateral(c *cli.Context) error {
 	jsonOutput(result)
 	return nil
 }
+
+var shutdownCommand = &cli.Command{
+	Name:   "shutdown",
+	Usage:  "shuts down bot instance",
+	Action: shutdown,
+}
+
+func shutdown(c *cli.Context) error {
+	conn, cancel, err := setupClient(c)
+	if err != nil {
+		return err
+	}
+	defer closeConn(conn, cancel)
+
+	client := gctrpc.NewGoCryptoTraderServiceClient(conn)
+	result, err := client.Shutdown(c.Context, &gctrpc.ShutdownRequest{})
+	if err != nil {
+		return err
+	}
+
+	jsonOutput(result)
+	return nil
+}
diff --git a/cmd/gctcli/main.go b/cmd/gctcli/main.go
@@ -202,6 +202,7 @@ func main() {
 		currencyStateManagementCommand,
 		getFuturesPositionsCommand,
 		getCollateralCommand,
+		shutdownCommand,
 	}
 
 	ctx, cancel := context.WithCancel(context.Background())

diff --git a/config/config_types.go b/config/config_types.go
@@ -190,6 +190,7 @@ type GRPCConfig struct {
 	ListenAddress          string `json:"listenAddress"`
 	GRPCProxyEnabled       bool   `json:"grpcProxyEnabled"`
 	GRPCProxyListenAddress string `json:"grpcProxyListenAddress"`
+	GRPCAllowBotShutdown   bool   `json:"grpcAllowBotShutdown"`
 	TimeInNanoSeconds      bool   `json:"timeInNanoSeconds"`
 }
 

diff --git a/engine/engine.go b/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{}{}
+}
diff --git a/engine/engine_types.go b/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 (

diff --git a/engine/rpcserver.go b/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
+}
diff --git a/engine/rpcserver_test.go b/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)
+	}
+}