From 9f41192bffbf4ec9c8c557dc2940fa2c71b92baa Mon Sep 17 00:00:00 2001
From: Leonard Lyubich <leonard@nspcc.ru>
Date: Mon, 18 Jan 2021 11:56:14 +0300
Subject: [PATCH] [#323] cmd/node: Compose functions for closing components in
 app

Closing callback can be registered in app through onShutdown method.

Signed-off-by: Leonard Lyubich <leonard@nspcc.ru>
---
 cmd/neofs-node/config.go  |  2 ++
 cmd/neofs-node/control.go |  4 ++++
 cmd/neofs-node/grpc.go    | 16 ++++++++++++++++
 cmd/neofs-node/main.go    | 11 +++++++----
 4 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/cmd/neofs-node/config.go b/cmd/neofs-node/config.go
index edb2cce4..62407cea 100644
--- a/cmd/neofs-node/config.go
+++ b/cmd/neofs-node/config.go
@@ -187,6 +187,8 @@ type cfg struct {
 	netStatus *atomic.Int32
 
 	healthStatus *atomic.Int32
+
+	closers []func()
 }
 
 type cfgGRPC struct {
diff --git a/cmd/neofs-node/control.go b/cmd/neofs-node/control.go
index ed9a6989..47700068 100644
--- a/cmd/neofs-node/control.go
+++ b/cmd/neofs-node/control.go
@@ -62,6 +62,10 @@ func initControlService(c *cfg) {
 		c.cfgControlService.server = grpc.NewServer()
 	}
 
+	c.onShutdown(func() {
+		stopGRPC("NeoFS Control API", c.cfgControlService.server, c.log)
+	})
+
 	control.RegisterControlServiceServer(c.cfgControlService.server, ctlSvc)
 
 	c.workers = append(c.workers, newWorkerFromFunc(func(ctx context.Context) {
diff --git a/cmd/neofs-node/grpc.go b/cmd/neofs-node/grpc.go
index f757871b..aaed1aca 100644
--- a/cmd/neofs-node/grpc.go
+++ b/cmd/neofs-node/grpc.go
@@ -4,6 +4,8 @@ import (
 	"fmt"
 	"net"
 
+	"github.com/nspcc-dev/neofs-node/pkg/util/logger"
+	"go.uber.org/zap"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/reflection"
 )
@@ -17,6 +19,10 @@ func initGRPC(c *cfg) {
 	c.cfgGRPC.server = grpc.NewServer(
 		grpc.MaxSendMsgSize(c.viper.GetInt(cfgMaxMsgSize)),
 	)
+
+	c.onShutdown(func() {
+		stopGRPC("NeoFS Public API", c.cfgGRPC.server, c.log)
+	})
 }
 
 func serveGRPC(c *cfg) {
@@ -38,3 +44,13 @@ func serveGRPC(c *cfg) {
 		}
 	}()
 }
+
+func stopGRPC(name string, s *grpc.Server, l *logger.Logger) {
+	l = l.With(zap.String("name", name))
+
+	l.Info("stopping gRPC server...")
+
+	s.GracefulStop()
+
+	l.Info("gRPC server stopped successfully")
+}
diff --git a/cmd/neofs-node/main.go b/cmd/neofs-node/main.go
index 4ff55f63..249193ac 100644
--- a/cmd/neofs-node/main.go
+++ b/cmd/neofs-node/main.go
@@ -78,12 +78,15 @@ func wait(c *cfg) {
 }
 
 func shutdown(c *cfg) {
-	c.cfgGRPC.server.GracefulStop()
-	c.cfgControlService.server.GracefulStop()
-
-	c.log.Info("gRPC server stopped")
+	for _, closer := range c.closers {
+		closer()
+	}
 
 	goOffline(c)
 
 	c.wg.Wait()
 }
+
+func (c *cfg) onShutdown(f func()) {
+	c.closers = append(c.closers, f)
+}