fix: grace shutdown for gw (#1478)

Signed-off-by: rfyiamcool <rfyiamcool@163.com>

pkg/common/cmd/msg_gateway.go

@@ -17,12 +17,12 @@ package cmd
 import (
 	"log"
 
-	"github.com/openimsdk/open-im-server/v3/internal/msggateway"
-	v3config "github.com/openimsdk/open-im-server/v3/pkg/common/config"
-
 	"github.com/spf13/cobra"
 
 	"github.com/OpenIMSDK/protocol/constant"
+
+	"github.com/openimsdk/open-im-server/v3/internal/msggateway"
+	v3config "github.com/openimsdk/open-im-server/v3/pkg/common/config"
 )
 
 type MsgGatewayCmd struct {
@@ -60,14 +60,19 @@ func (m *MsgGatewayCmd) Exec() error {
 	m.addRunE()
 	return m.Execute()
 }
+
 func (m *MsgGatewayCmd) GetPortFromConfig(portType string) int {
-	if portType == constant.FlagWsPort {
+	switch portType {
+	case constant.FlagWsPort:
 		return v3config.Config.LongConnSvr.OpenImWsPort[0]
-	} else if portType == constant.FlagPort {
+
+	case constant.FlagPort:
 		return v3config.Config.LongConnSvr.OpenImMessageGatewayPort[0]
-	} else if portType == constant.FlagPrometheusPort {
+
+	case constant.FlagPrometheusPort:
 		return v3config.Config.Prometheus.MessageGatewayPrometheusPort[0]
-	} else {
+
+	default:
 		return 0
 	}
 }

pkg/common/startrpc/start.go

@@ -15,14 +15,21 @@
 package startrpc
 
 import (
+	"errors"
 	"fmt"
 	"log"
 	"net"
 	"net/http"
+	"os"
+	"os/signal"
 	"strconv"
+	"sync"
+	"syscall"
+	"time"
 
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
+	"golang.org/x/sync/errgroup"
 
 	"github.com/openimsdk/open-im-server/v3/pkg/common/config"
 	"github.com/openimsdk/open-im-server/v3/pkg/common/prommetrics"
@@ -56,31 +63,37 @@ func Start(
 	if err != nil {
 		return err
 	}
+
 	defer listener.Close()
 	client, err := kdisc.NewDiscoveryRegister(config.Config.Envs.Discovery)
 	if err != nil {
 		return utils.Wrap1(err)
 	}
+
 	defer client.Close()
 	client.AddOption(mw.GrpcClient(), grpc.WithTransportCredentials(insecure.NewCredentials()))
 	registerIP, err := network.GetRpcRegisterIP(config.Config.Rpc.RegisterIP)
 	if err != nil {
 		return err
 	}
+
 	var reg *prometheus.Registry
 	var metric *grpcprometheus.ServerMetrics
-	// ctx 中间件
 	if config.Config.Prometheus.Enable {
-		//////////////////////////
 		cusMetrics := prommetrics.GetGrpcCusMetrics(rpcRegisterName)
-		reg, metric, err = prommetrics.NewGrpcPromObj(cusMetrics)
+		reg, metric, _ = prommetrics.NewGrpcPromObj(cusMetrics)
 		options = append(options, mw.GrpcServer(), grpc.StreamInterceptor(metric.StreamServerInterceptor()),
 			grpc.UnaryInterceptor(metric.UnaryServerInterceptor()))
 	} else {
 		options = append(options, mw.GrpcServer())
 	}
+
 	srv := grpc.NewServer(options...)
-	defer srv.GracefulStop()
+	once := sync.Once{}
+	defer func() {
+		once.Do(srv.GracefulStop)
+	}()
+
 	err = rpcFn(client, srv)
 	if err != nil {
 		return utils.Wrap1(err)
@@ -94,7 +107,10 @@ func Start(
 	if err != nil {
 		return utils.Wrap1(err)
 	}
-	go func() {
+
+	var wg errgroup.Group
+
+	wg.Go(func() error {
 		if config.Config.Prometheus.Enable && prometheusPort != 0 {
 			metric.InitializeMetrics(srv)
 			// Create a HTTP server for prometheus.
@@ -103,7 +119,34 @@ func Start(
 				log.Fatal("Unable to start a http server.")
 			}
 		}
+		return nil
+	})
+
+	wg.Go(func() error {
+		return utils.Wrap1(srv.Serve(listener))
+	})
+
+	sigs := make(chan os.Signal, 1)
+	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
+	<-sigs
+
+	var (
+		done = make(chan struct{}, 1)
+		gerr error
+	)
+
+	go func() {
+		once.Do(srv.GracefulStop)
+		gerr = wg.Wait()
+		close(done)
 	}()
 
-	return utils.Wrap1(srv.Serve(listener))
+	select {
+	case <-done:
+		return gerr
+
+	case <-time.After(15 * time.Second):
+		return utils.Wrap1(errors.New("timeout exit"))
+	}
+
 }