fix: refactoring code of graceful exits (#1885)

* fix: plant a layer

* fix: print chanal

* fix: print sigs

* fix: print the sigs

* fix: reconstruct exit gracefully

* fix: fix the timeout

* fix: fix the netDone

* fix: fix the process exit

* fix: refactor the elegant startup code

* fix: fix the Signal.Notify

* fix: fix the code

* fix: remove not used header import.

* Update init.go

* fix: fix the InitConfig error

* fix: fix branch name

* fix: fix the signal value

* fix: replace the signal with SIGTERM

* fix: fix the script

* fix: fix the unsolve error

* fix: return the SIGTERM received,shutting down

* fix: fix the tranfer exit error

* fix: fix the error

* fix: replace the SIGnal

* fix: del the error return in tranfer

* fix: fix SIGTERM error

* fix: del the unreachalbe code

* fix: fix the make stop print  error

---------

Co-authored-by: OpenIM-Gordon <46924906+FGadvancer@users.noreply.github.com>
Co-authored-by: Xinwei Xiong (cubxxw) <3293172751nss@gmail.com>

internal/msggateway/init.go

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

internal/msggateway/n_ws_server.go

@@ -20,12 +20,9 @@ import (
 	"errors"
 	"fmt"
 	"net/http"
-	"os"
-	"os/signal"
 	"strconv"
 	"sync"
 	"sync/atomic"
-	"syscall"
 	"time"
 
 	"github.com/OpenIMSDK/tools/apiresp"
@@ -49,7 +46,7 @@ import (
 )
 
 type LongConnServer interface {
-	Run() error
+	Run(done chan error) error
 	wsHandler(w http.ResponseWriter, r *http.Request)
 	GetUserAllCons(userID string) ([]*Client, bool)
 	GetUserPlatformCons(userID string, platform int) ([]*Client, bool, bool)
@@ -169,23 +166,20 @@ func NewWsServer(opts ...Option) (*WsServer, error) {
 	}, nil
 }
 
-func (ws *WsServer) Run() error {
+func (ws *WsServer) Run(done chan error) error {
 	var (
-		client *Client
-		wg     errgroup.Group
-
-		sigs = make(chan os.Signal, 1)
-		done = make(chan struct{}, 1)
+		client       *Client
+		netErr       error
+		shutdownDone = make(chan struct{}, 1)
 	)
 
 	server := http.Server{Addr: ":" + utils.IntToString(ws.port), Handler: nil}
 
-	wg.Go(func() error {
+	go func() {
 		for {
 			select {
-			case <-done:
-				return nil
-
+			case <-shutdownDone:
+				return
 			case client = <-ws.registerChan:
 				ws.registerClient(client)
 			case client = <-ws.unregisterChan:
@@ -194,33 +188,32 @@ 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)
 	}()
-
+	netDone := make(chan struct{}, 1)
+	go func() {
+		http.HandleFunc("/", ws.wsHandler)
+		err := server.ListenAndServe()
+		if err != nil && err != http.ErrServerClosed {
+			netErr = errs.Wrap(err, "ws start err", server.Addr)
+			close(netDone)
+		}
+	}()
+	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
+	defer cancel()
+	var err error
 	select {
-	case <-done:
-		return nil
-
-	case <-time.After(15 * time.Second):
-		return utils.Wrap1(errors.New("timeout exit"))
+	case err = <-done:
+		sErr := server.Shutdown(ctx)
+		if sErr != nil {
+			return errs.Wrap(sErr, "shutdown err")
+		}
+		close(shutdownDone)
+		if err != nil {
+			return err
+		}
+	case <-netDone:
 	}
+	return netErr
 
 }
 

internal/msgtransfer/init.go

@@ -18,23 +18,11 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"net/http"
-	"os"
-	"os/signal"
-	"sync"
-	"syscall"
-	"time"
-
 	"github.com/OpenIMSDK/tools/errs"
-
 	"github.com/OpenIMSDK/tools/log"
-	"github.com/OpenIMSDK/tools/mw"
-	"github.com/prometheus/client_golang/prometheus"
-	"github.com/prometheus/client_golang/prometheus/collectors"
-	"github.com/prometheus/client_golang/prometheus/promhttp"
-	"google.golang.org/grpc"
-	"google.golang.org/grpc/credentials/insecure"
+	util "github.com/openimsdk/open-im-server/v3/pkg/util/genutil"
 
+	"github.com/OpenIMSDK/tools/mw"
 	"github.com/openimsdk/open-im-server/v3/pkg/common/config"
 	"github.com/openimsdk/open-im-server/v3/pkg/common/db/cache"
 	"github.com/openimsdk/open-im-server/v3/pkg/common/db/controller"
@@ -42,12 +30,22 @@ import (
 	kdisc "github.com/openimsdk/open-im-server/v3/pkg/common/discoveryregister"
 	"github.com/openimsdk/open-im-server/v3/pkg/common/prommetrics"
 	"github.com/openimsdk/open-im-server/v3/pkg/rpcclient"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/collectors"
+	"github.com/prometheus/client_golang/prometheus/promhttp"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+	"net/http"
+	"os"
+	"os/signal"
+	"syscall"
 )
 
 type MsgTransfer struct {
 	historyCH      *OnlineHistoryRedisConsumerHandler // 这个消费者聚合消息, 订阅的topic：ws2ms_chat, 修改通知发往msg_to_modify topic, 消息存入redis后Incr Redis, 再发消息到ms2pschat topic推送， 发消息到msg_to_mongo topic持久化
 	historyMongoCH *OnlineHistoryMongoConsumerHandler // mongoDB批量插入, 成功后删除redis中消息，以及处理删除通知消息删除的 订阅的topic: msg_to_mongo
-	// modifyCH       *ModifyMsgConsumerHandler          // 负责消费修改消息通知的consumer, 订阅的topic: msg_to_modify
+	ctx            context.Context
+	cancel         context.CancelFunc
 }
 
 func StartTransfer(prometheusPort int) error {
@@ -65,10 +63,6 @@ func StartTransfer(prometheusPort int) error {
 		return err
 	}
 	client, err := kdisc.NewDiscoveryRegister(config.Config.Envs.Discovery)
-	/*
-		client, err := openkeeper.NewClient(config.Config.Zookeeper.ZkAddr, config.Config.Zookeeper.Schema,
-			openkeeper.WithFreq(time.Hour), openkeeper.WithRoundRobin(), openkeeper.WithUserNameAndPassword(config.Config.Zookeeper.Username,
-				config.Config.Zookeeper.Password), openkeeper.WithTimeout(10), openkeeper.WithLogger(log.NewZkLogger()))*/
 	if err != nil {
 		return err
 	}
@@ -109,27 +103,22 @@ func NewMsgTransfer(msgDatabase controller.CommonMsgDatabase, conversationRpcCli
 }
 
 func (m *MsgTransfer) Start(prometheusPort int) error {
-	ctx := context.Background()
 	fmt.Println("start msg transfer", "prometheusPort:", prometheusPort)
 	if prometheusPort <= 0 {
 		return errs.Wrap(errors.New("prometheusPort not correct"))
 	}
+	m.ctx, m.cancel = context.WithCancel(context.Background())
 
-	var wg sync.WaitGroup
+	var (
+		netDone = make(chan struct{}, 1)
+		netErr  error
+	)
 
-	wg.Add(1)
-	go func() {
-		defer wg.Done()
-
-		m.historyCH.historyConsumerGroup.RegisterHandleAndConsumer(ctx, m.historyCH)
-	}()
-
-	wg.Add(1)
-	go func() {
-		defer wg.Done()
-
-		m.historyMongoCH.historyConsumerGroup.RegisterHandleAndConsumer(ctx, m.historyMongoCH)
-	}()
+	onError := func(ctx context.Context, err error, errInfo string) {
+		log.ZWarn(ctx, errInfo, err)
+	}
+	go m.historyCH.historyConsumerGroup.RegisterHandleAndConsumer(m.ctx, m.historyCH, onError)
+	go m.historyMongoCH.historyConsumerGroup.RegisterHandleAndConsumer(m.ctx, m.historyMongoCH, onError)
 
 	if config.Config.Prometheus.Enable {
 		go func() {
@@ -141,30 +130,28 @@ func (m *MsgTransfer) Start(prometheusPort int) error {
 			http.Handle("/metrics", promhttp.HandlerFor(proreg, promhttp.HandlerOpts{Registry: proreg}))
 			err := http.ListenAndServe(fmt.Sprintf(":%d", prometheusPort), nil)
 			if err != nil && err != http.ErrServerClosed {
-				panic(err)
+				netErr = errs.Wrap(err, fmt.Sprintf("prometheus start err: %d", prometheusPort))
+				netDone <- struct{}{}
 			}
 		}()
 	}
 
 	sigs := make(chan os.Signal, 1)
-	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
-	<-sigs
-
-	// graceful close kafka client.
-	go m.historyCH.historyConsumerGroup.Close()
-	go m.historyMongoCH.historyConsumerGroup.Close()
-
-	done := make(chan struct{}, 1)
-	go func() {
-		wg.Wait()
-		close(done)
-	}()
-
+	signal.Notify(sigs, syscall.SIGTERM)
 	select {
-	case <-done:
-		log.ZInfo(context.Background(), "msgtrasfer exit successfully")
-	case <-time.After(15 * time.Second):
-		log.ZError(context.Background(), "msgtransfer force to exit, timeout 15s", nil)
+	case <-sigs:
+		util.SIGUSR1Exit()
+		// graceful close kafka client.
+		m.cancel()
+		m.historyCH.historyConsumerGroup.Close()
+		m.historyMongoCH.historyConsumerGroup.Close()
+
+	case <-netDone:
+		m.cancel()
+		m.historyCH.historyConsumerGroup.Close()
+		m.historyMongoCH.historyConsumerGroup.Close()
+		close(netDone)
+		return netErr
 	}
 
 	return nil

internal/push/callback.go

@@ -130,7 +130,7 @@ func callbackBeforeSuperGroupOnlinePush(
 	if err := http.CallBackPostReturn(ctx, config.Config.Callback.CallbackUrl, req, resp, config.Config.Callback.CallbackBeforeSuperGroupOnlinePush); err != nil {
 		return err
 	}
-	return nil
+
 	if len(resp.UserIDs) != 0 {
 		*pushToUserIDs = resp.UserIDs
 	}

internal/push/consumer_init.go

@@ -14,7 +14,10 @@
 
 package push
 
-import "context"
+import (
+	"context"
+	"github.com/OpenIMSDK/tools/log"
+)
 
 type Consumer struct {
 	pushCh       ConsumerHandler
@@ -32,6 +35,9 @@ func NewConsumer(pusher *Pusher) (*Consumer, error) {
 }
 
 func (c *Consumer) Start() {
+	onError := func(ctx context.Context, err error, errInfo string) {
+		log.ZWarn(ctx, errInfo, err)
+	}
+	go c.pushCh.pushConsumerGroup.RegisterHandleAndConsumer(context.Background(), &c.pushCh, onError)
 
-	go c.pushCh.pushConsumerGroup.RegisterHandleAndConsumer(context.Background(), &c.pushCh)
 }

internal/push/push_rpc_server.go

@@ -16,8 +16,6 @@ package push
 
 import (
 	"context"
-	"sync"
-
 	"github.com/OpenIMSDK/tools/utils"
 
 	"google.golang.org/grpc"
@@ -58,23 +56,18 @@ func Start(client discoveryregistry.SvcDiscoveryRegistry, server *grpc.Server) e
 		&groupRpcClient,
 		&msgRpcClient,
 	)
-	var wg sync.WaitGroup
-	wg.Add(2)
-	go func() {
-		defer wg.Done()
-		pbpush.RegisterPushMsgServiceServer(server, &pushServer{
-			pusher: pusher,
-		})
-	}()
+
+	pbpush.RegisterPushMsgServiceServer(server, &pushServer{
+		pusher: pusher,
+	})
+
 	consumer, err := NewConsumer(pusher)
 	if err != nil {
 		return err
 	}
-	go func() {
-		defer wg.Done()
-		consumer.Start()
-	}()
-	wg.Wait()
+
+	consumer.Start()
+
 	return nil
 }
 

internal/tools/cron_task.go

@@ -64,7 +64,7 @@ func StartTask() error {
 	crontab.Start()
 
 	sigs := make(chan os.Signal, 1)
-	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
+	signal.Notify(sigs, syscall.SIGTERM)
 	<-sigs
 
 	// stop crontab, Wait for the running task to exit.