optimize modelRT routing structure

main.go

@@ -4,6 +4,10 @@ package main
 import (
 	"context"
 	"flag"
+	"net/http"
+	"os"
+	"os/signal"
+	"syscall"
 	"time"
 
 	"modelRT/alert"
@@ -16,6 +20,7 @@ import (
 	"modelRT/logger"
 	"modelRT/middleware"
 	"modelRT/pool"
+	"modelRT/router"
 
 	realtimedata "modelRT/real-time-data"
 
@@ -64,9 +69,6 @@ func main() {
 		sqlDB.Close()
 	}()
 
-	// init logger
-	logger.InitLoggerInstance(modelRTConfig.LoggerConfig)
-
 	// init alert manager
 	_ = alert.InitAlertEventManager()
 
@@ -120,14 +122,34 @@ func main() {
 	// TODO 暂时屏蔽完成 swagger 启动测试
 	// go realtimedata.RealTimeDataComputer(ctx, nil, []string{}, "")
 
-	engine := gin.Default()
+	engine := gin.New()
-	engine.Use(limiter.Middleware)
+	router.RegisterRoutes(engine)
+	server := http.Server{
+		Addr:    ":8080",
+		Handler: engine,
+	}
 
-	// diagram api
+	// creating a System Signal Receiver
-	engine.GET("/diagram/load", handler.CircuitDiagramLoadHandler)
+	done := make(chan os.Signal, 10)
-	engine.POST("/diagram/create", handler.CircuitDiagramCreateHandler)
+	signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
-	engine.POST("/diagram/update", handler.CircuitDiagramUpdateHandler)
+	go func() {
-	engine.POST("/diagram/delete", handler.CircuitDiagramDeleteHandler)
+		<-done
+		if err := server.Shutdown(context.Background()); err != nil {
+			logger.Error(ctx, "ShutdownServerError", "err", err)
+		}
+	}()
+
+	logger.Info(ctx, "Starting ModelRT server...")
+	err = server.ListenAndServe()
+	if err != nil {
+		if err == http.ErrServerClosed {
+			// the service receives the shutdown signal normally and then closes
+			logger.Info(ctx, "Server closed under request")
+		} else {
+			// abnormal shutdown of service
+			logger.Error(ctx, "Server closed unexpected", "err", err)
+		}
+	}
 
 	// real time data api
 	engine.GET("/ws/rtdatas", handler.RealTimeDataReceivehandler)

middleware/trace.go

@@ -1,6 +1,12 @@
 package middleware
 
 import (
+	"bytes"
+	"io"
+	"strings"
+	"time"
+
+	"modelRT/logger"
 	"modelRT/util"
 
 	"github.com/gin-gonic/gin"
@@ -21,3 +27,71 @@ func StartTrace() gin.HandlerFunc {
 		c.Next()
 	}
 }
+
+type bodyLogWriter struct {
+	gin.ResponseWriter
+	body *bytes.Buffer
+}
+
+// 包装一下 gin.ResponseWriter，通过这种方式拦截写响应
+// 让gin写响应的时候先写到 bodyLogWriter 再写gin.ResponseWriter ，
+// 这样利用中间件里输出访问日志时就能拿到响应了
+// https://stackoverflow.com/questions/38501325/how-to-log-response-body-in-gin
+func (w bodyLogWriter) Write(b []byte) (int, error) {
+	w.body.Write(b)
+	return w.ResponseWriter.Write(b)
+}
+
+func LogAccess() gin.HandlerFunc {
+	return func(c *gin.Context) {
+		// 保存body
+		var reqBody []byte
+		contentType := c.GetHeader("Content-Type")
+		// multipart/form-data 文件上传请求, 不在日志里记录body
+		if !strings.Contains(contentType, "multipart/form-data") {
+			reqBody, _ = io.ReadAll(c.Request.Body)
+			c.Request.Body = io.NopCloser(bytes.NewReader(reqBody))
+
+			// var request map[string]interface{}
+			// if err := c.ShouldBindBodyWith(&request, binding.JSON); err != nil {
+			// 	c.JSON(400, gin.H{"error": err.Error()})
+			// 	return
+			// }
+		}
+		start := time.Now()
+		blw := &bodyLogWriter{body: bytes.NewBufferString(""), ResponseWriter: c.Writer}
+		c.Writer = blw
+
+		accessLog(c, "access_start", time.Since(start), reqBody, nil)
+		defer func() {
+			var responseLogging string
+			if c.Writer.Size() > 10*1024 { // 响应大于10KB 不记录
+				responseLogging = "Response data size is too Large to log"
+			} else {
+				responseLogging = blw.body.String()
+			}
+			accessLog(c, "access_end", time.Since(start), reqBody, responseLogging)
+		}()
+		c.Next()
+		return
+	}
+}
+
+func accessLog(c *gin.Context, accessType string, dur time.Duration, body []byte, dataOut interface{}) {
+	req := c.Request
+	bodyStr := string(body)
+	query := req.URL.RawQuery
+	path := req.URL.Path
+	// TODO: 实现Token认证后再把访问日志里也加上token记录
+	// token := c.Request.Header.Get("token")
+	logger.New(c).Info("AccessLog",
+		"type", accessType,
+		"ip", c.ClientIP(),
+		//"token", token,
+		"method", req.Method,
+		"path", path,
+		"query", query,
+		"body", bodyStr,
+		"output", dataOut,
+		"time(ms)", int64(dur/time.Millisecond))
+}

router/diagram.go

@@ -0,0 +1,17 @@
+package router
+
+import (
+	"modelRT/handler"
+
+	"github.com/gin-gonic/gin"
+)
+
+// RegisterRoutes define func of register diagram routes
+func registerDiagramRoutes(rg *gin.RouterGroup) {
+	g := rg.Group("/diagram/")
+	// TODO add diagram middleware
+	g.GET("load", handler.CircuitDiagramLoadHandler)
+	g.POST("create", handler.CircuitDiagramCreateHandler)
+	g.POST("update", handler.CircuitDiagramUpdateHandler)
+	g.POST("delete", handler.CircuitDiagramDeleteHandler)
+}

router/router.go

@@ -0,0 +1,22 @@
+package router
+
+import (
+	"time"
+
+	"modelRT/middleware"
+
+	"github.com/gin-gonic/gin"
+)
+
+var limiter *middleware.Limiter
+
+func init() {
+	limiter = middleware.NewLimiter(10, 1*time.Minute) // 设置限流器，允许每分钟最多请求10次
+}
+
+func RegisterRoutes(engine *gin.Engine) {
+	// use global middlewares
+	engine.Use(middleware.StartTrace(), limiter.Middleware)
+	routeGroup := engine.Group("")
+	registerDiagramRoutes(routeGroup)
+}