diff --git a/internal/api/init.go b/internal/api/init.go index 4ceabda90..53a0aed7b 100644 --- a/internal/api/init.go +++ b/internal/api/init.go @@ -156,13 +156,20 @@ func Start(ctx context.Context, index int, config *Config) error { signal.Notify(sigs, syscall.SIGTERM) ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + shutdown := func() error { + err := server.Shutdown(ctx) + if err != nil { + return errs.WrapMsg(err, "shutdown err") + } + return nil + } + disetcd.RegisterShutDown(shutdown) defer cancel() select { case <-sigs: program.SIGTERMExit() - err := server.Shutdown(ctx) - if err != nil { - return errs.WrapMsg(err, "shutdown err") + if err := shutdown(); err != nil { + return err } case <-netDone: close(netDone) diff --git a/internal/msggateway/ws_server.go b/internal/msggateway/ws_server.go index 7271c3727..095d1adbc 100644 --- a/internal/msggateway/ws_server.go +++ b/internal/msggateway/ws_server.go @@ -2,12 +2,14 @@ package msggateway import ( "context" + "errors" "fmt" "net/http" "sync" "sync/atomic" "time" + "github.com/openimsdk/open-im-server/v3/pkg/common/discovery/etcd" "github.com/openimsdk/open-im-server/v3/pkg/common/webhook" "github.com/openimsdk/open-im-server/v3/pkg/rpccache" pbAuth "github.com/openimsdk/protocol/auth" @@ -161,20 +163,27 @@ func (ws *WsServer) Run(done chan error) error { http.HandleFunc("/", ws.wsHandler) err := server.ListenAndServe() defer close(netDone) - if err != nil && err != http.ErrServerClosed { + if err != nil && !errors.Is(err, http.ErrServerClosed) { netErr = errs.WrapMsg(err, "ws start err", server.Addr) } }() ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) - defer cancel() - var err error - select { - case err = <-done: + shutDown := func() error { sErr := server.Shutdown(ctx) if sErr != nil { return errs.WrapMsg(sErr, "shutdown err") } close(shutdownDone) + return nil + } + etcd.RegisterShutDown(shutDown) + defer cancel() + var err error + select { + case err = <-done: + if err := shutDown(); err != nil { + return err + } if err != nil { return err } diff --git a/pkg/common/discovery/etcd/config_manager.go b/pkg/common/discovery/etcd/config_manager.go index f0e45c886..46b676ec6 100644 --- a/pkg/common/discovery/etcd/config_manager.go +++ b/pkg/common/discovery/etcd/config_manager.go @@ -18,6 +18,14 @@ const ( ConfigKeyPrefix = "/open-im/config/" ) +var ( + ShutDowns []func() error +) + +func RegisterShutDown(shutDown ...func() error) { + ShutDowns = append(ShutDowns, shutDown...) +} + type ConfigManager struct { client *clientv3.Client watchConfigNames []string @@ -83,6 +91,12 @@ func restartServer(ctx context.Context) error { if runtime.GOOS != "windows" { cmd.SysProcAttr = &syscall.SysProcAttr{} } + log.ZInfo(ctx, "shutdown server") + for _, f := range ShutDowns { + if err = f(); err != nil { + log.ZError(ctx, "shutdown fail", err) + } + } log.ZInfo(ctx, "restart server") err = cmd.Start()