fix: grace shutdown for gw (#1478)

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

internal/msggateway/init.go

@@ -19,6 +19,7 @@ import (
 	"time"
 
 	"github.com/OpenIMSDK/tools/utils"
+	"golang.org/x/sync/errgroup"
 
 	"github.com/openimsdk/open-im-server/v3/pkg/common/config"
 )
@@ -43,12 +44,22 @@ func RunWsAndServer(rpcPort, wsPort, prometheusPort int) error {
 	if err != nil {
 		return err
 	}
+
 	hubServer := NewServer(rpcPort, prometheusPort, longServer)
-	go func() {
-		err := hubServer.Start()
+
+	wg := errgroup.Group{}
+	wg.Go(func() error {
+		err = hubServer.Start()
 		if err != nil {
-			panic(utils.Wrap1(err))
+			return utils.Wrap1(err)
 		}
-	}()
-	return hubServer.LongConnServer.Run()
+		return err
+	})
+
+	wg.Go(func() error {
+		return hubServer.LongConnServer.Run()
+	})
+
+	err = wg.Wait()
+	return err
 }

internal/msggateway/n_ws_server.go

@@ -18,9 +18,12 @@ import (
 	"context"
 	"errors"
 	"net/http"
+	"os"
+	"os/signal"
 	"strconv"
 	"sync"
 	"sync/atomic"
+	"syscall"
 	"time"
 
 	"github.com/go-playground/validator/v10"
@@ -156,10 +159,22 @@ func NewWsServer(opts ...Option) (*WsServer, error) {
 }
 
 func (ws *WsServer) Run() error {
-	var client *Client
-	go func() {
+	var (
+		client *Client
+		wg     errgroup.Group
+
+		sigs = make(chan os.Signal, 1)
+		done = make(chan struct{}, 1)
+	)
+
+	server := http.Server{Addr: ":" + utils.IntToString(ws.port), Handler: nil}
+
+	wg.Go(func() error {
 		for {
 			select {
+			case <-done:
+				return nil
+
 			case client = <-ws.registerChan:
 				ws.registerClient(client)
 			case client = <-ws.unregisterChan:
@@ -168,10 +183,34 @@ func (ws *WsServer) Run() error {
 				ws.multiTerminalLoginChecker(onlineInfo.clientOK, onlineInfo.oldClients, onlineInfo.newClient)
 			}
 		}
+	})
+
+	wg.Go(func() error {
+		http.HandleFunc("/", ws.wsHandler)
+		return server.ListenAndServe()
+	})
+
+	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
+	<-sigs
+
+	go func() {
+		ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
+		defer cancel()
+
+		// graceful exit operation for server
+		_ = server.Shutdown(ctx)
+		_ = wg.Wait()
+		close(done)
 	}()
-	http.HandleFunc("/", ws.wsHandler)
-	// http.HandleFunc("/metrics", func(w http.ResponseWriter, r *http.Request) {})
-	return http.ListenAndServe(":"+utils.IntToString(ws.port), nil) // Start listening
+
+	select {
+	case <-done:
+		return nil
+
+	case <-time.After(15 * time.Second):
+		return utils.Wrap1(errors.New("timeout exit"))
+	}
+
 }
 
 var concurrentRequest = 3