fix: 增加 listenAndServe 启动 和 shutdown 关闭

internal/example/main.go

@@ -1,6 +1,10 @@
 package main
 
 import (
+	"context"
+	"log"
+	"time"
+
 	"github.com/tangx/rum-gonic/internal/example/apis"
 	"github.com/tangx/rum-gonic/internal/example/apis/index"
 	"github.com/tangx/rum-gonic/internal/example/injector/redis"
@@ -31,7 +35,39 @@ func main() {
 	// r.Static("/userindex", "userindex")
 	r.Static("/userindex", "/data/gopath/src/github.com/tangx/rum-gonic/userindex")
 
+	/* 启动方式 */
+	// 1. 普通方式启动
+	normalRun(r)
+
+	// 2. 启动并控制退出
+	listenAndServe(r)
+}
+
+func normalRun(r *rum.Engine) {
 	if err := r.Run(); err != nil {
 		panic(err)
 	}
 }
+
+func listenAndServe(r *rum.Engine) {
+
+	go func() {
+		if err := r.ListenAndServe(); err != nil {
+			panic(err)
+		}
+	}()
+
+	ctx := context.Background()
+	ctx, stop := context.WithTimeout(ctx, 10*time.Second)
+	defer stop()
+
+	time.Sleep(30 * time.Second)
+
+	if err := r.Shutdown(ctx); err != nil {
+		log.Println("强制关闭 engine")
+	}
+
+	log.Println("rum 已经退出")
+	// 一分钟后关闭
+	time.Sleep(1 * time.Minute)
+}

rum/engine.go

@@ -1,13 +1,19 @@
 package rum
 
 import (
+	"context"
+	"errors"
+	"net/http"
+
 	"github.com/gin-gonic/gin"
 )
 
 type Engine struct {
 	engine *gin.Engine
 
 	*RouterGroup
+
+	srv *http.Server
 }
 
 // Default 默认 gin.Engine
@@ -30,14 +36,6 @@ func New(r *gin.Engine) *Engine {
 	return e
 }
 
-// Run 监听 tcp 端口
-func (e *Engine) Run(addrs ...string) error {
-
-	e.register()
-
-	return e.engine.Run(addrs...)
-}
-
 func (e *Engine) register() {
 	root := &RouterGroup{
 		ginRG: &e.engine.RouterGroup,
@@ -54,3 +52,46 @@ func (e *Engine) WithContextCompose(injectors ...ContextInjectorFunc) *Engine {
 
 	return e
 }
+
+// Run 监听 tcp 端口
+func (e *Engine) Run(addrs ...string) error {
+
+	e.register()
+
+	return e.engine.Run(addrs...)
+}
+
+// ListenAndServe 启动服务， 可以配合 Shutdown(ctx) 自定义退出
+func (e *Engine) ListenAndServe(addrs ...string) error {
+	e.register()
+
+	if e.srv == nil {
+		e.srv = &http.Server{
+			Addr:    addr(addrs...),
+			Handler: e.engine,
+		}
+	}
+
+	err := e.srv.ListenAndServe()
+	if err != nil && err != http.ErrServerClosed {
+		return err
+	}
+	return nil
+}
+
+// Shutdown 用户自定义退出。 如果不是使用 ListenAndServe 启动则不能使用。
+func (e *Engine) Shutdown(ctx context.Context) error {
+	if e.srv == nil {
+		return errors.New(NotRunWithListenAndServe)
+	}
+
+	return e.srv.Shutdown(ctx)
+}
+
+func addr(addrs ...string) string {
+	if len(addrs) == 0 {
+		return ":8080"
+	}
+
+	return addrs[1]
+}

rum/errors.go

@@ -0,0 +1,5 @@
+package rum
+
+var (
+	NotRunWithListenAndServe = "rum not run with ListenAndServe method"
+)