core.go 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468
  1. package mux
  2. import (
  3. "errors"
  4. "fmt"
  5. "golang.org/x/time/rate"
  6. "net/http"
  7. "net/url"
  8. "runtime/debug"
  9. "time"
  10. "git.bvbej.com/bvbej/base-golang/pkg/color"
  11. "git.bvbej.com/bvbej/base-golang/pkg/env"
  12. "git.bvbej.com/bvbej/base-golang/pkg/errno"
  13. "git.bvbej.com/bvbej/base-golang/pkg/trace"
  14. "github.com/gin-contrib/pprof"
  15. "github.com/gin-gonic/gin"
  16. "github.com/prometheus/client_golang/prometheus/promhttp"
  17. cors "github.com/rs/cors/wrapper/gin"
  18. "go.uber.org/multierr"
  19. "go.uber.org/zap"
  20. )
  21. const (
  22. _EverySize = 1e3
  23. _MaxBurstSize = _EverySize * 60
  24. )
  25. var limiter = rate.NewLimiter(rate.Limit(_EverySize), _MaxBurstSize)
  26. type Option func(*option)
  27. type option struct {
  28. enableCors bool
  29. enableRate bool
  30. enablePProf bool
  31. enablePrometheus bool
  32. enableOpenBrowser string
  33. staticDirs []string
  34. panicNotify OnPanicNotify
  35. recordMetrics RecordMetrics
  36. }
  37. type Failure struct {
  38. ResultCode int `json:"result_code"` // 业务码
  39. ResultInfo string `json:"result_info"` // 描述信息
  40. }
  41. type Success struct {
  42. ResultCode int `json:"result_code"` // 业务码
  43. ResultData any `json:"result_data"` //返回数据
  44. }
  45. /******************************************************************************/
  46. // OnPanicNotify 发生panic时通知用
  47. type OnPanicNotify func(ctx Context, err any, stackInfo string)
  48. // RecordMetrics 记录prometheus指标用
  49. // 如果使用AliasForRecordMetrics配置了别名,uri将被替换为别名。
  50. type RecordMetrics func(method, uri string, success bool, costSeconds float64)
  51. // DisableTrace 禁用追踪链
  52. func DisableTrace(ctx Context) {
  53. ctx.disableTrace()
  54. }
  55. // WithPanicNotify 设置panic时的通知回调
  56. func WithPanicNotify(notify OnPanicNotify) Option {
  57. return func(opt *option) {
  58. opt.panicNotify = notify
  59. fmt.Println(color.Green("* [register panic notify]"))
  60. }
  61. }
  62. // WithRecordMetrics 设置记录prometheus记录指标回调
  63. func WithRecordMetrics(record RecordMetrics) Option {
  64. return func(opt *option) {
  65. opt.recordMetrics = record
  66. }
  67. }
  68. // WithEnableCors 开启CORS
  69. func WithEnableCors() Option {
  70. return func(opt *option) {
  71. opt.enableCors = true
  72. fmt.Println(color.Green("* [register cors]"))
  73. }
  74. }
  75. // WithEnableRate 开启限流
  76. func WithEnableRate() Option {
  77. return func(opt *option) {
  78. opt.enableRate = true
  79. fmt.Println(color.Green("* [register rate]"))
  80. }
  81. }
  82. // WithStaticDir 设置静态文件目录
  83. func WithStaticDir(dirs []string) Option {
  84. return func(opt *option) {
  85. opt.staticDirs = dirs
  86. fmt.Println(color.Green("* [register rate]"))
  87. }
  88. }
  89. // AliasForRecordMetrics 对请求uri起个别名,用于prometheus记录指标。
  90. // 如:Get /user/:username 这样的uri,因为username会有非常多的情况,这样记录prometheus指标会非常的不有好。
  91. func AliasForRecordMetrics(path string) HandlerFunc {
  92. return func(ctx Context) {
  93. ctx.setAlias(path)
  94. }
  95. }
  96. /******************************************************************************/
  97. // RouterGroup 包装gin的RouterGroup
  98. type RouterGroup interface {
  99. Group(string, ...HandlerFunc) RouterGroup
  100. IRoutes
  101. }
  102. var _ IRoutes = (*router)(nil)
  103. // IRoutes 包装gin的IRoutes
  104. type IRoutes interface {
  105. Any(string, ...HandlerFunc)
  106. GET(string, ...HandlerFunc)
  107. POST(string, ...HandlerFunc)
  108. DELETE(string, ...HandlerFunc)
  109. PATCH(string, ...HandlerFunc)
  110. PUT(string, ...HandlerFunc)
  111. OPTIONS(string, ...HandlerFunc)
  112. HEAD(string, ...HandlerFunc)
  113. }
  114. type router struct {
  115. group *gin.RouterGroup
  116. }
  117. func (r *router) Group(relativePath string, handlers ...HandlerFunc) RouterGroup {
  118. group := r.group.Group(relativePath, wrapHandlers(handlers...)...)
  119. return &router{group: group}
  120. }
  121. func (r *router) Any(relativePath string, handlers ...HandlerFunc) {
  122. r.group.Any(relativePath, wrapHandlers(handlers...)...)
  123. }
  124. func (r *router) GET(relativePath string, handlers ...HandlerFunc) {
  125. r.group.GET(relativePath, wrapHandlers(handlers...)...)
  126. }
  127. func (r *router) POST(relativePath string, handlers ...HandlerFunc) {
  128. r.group.POST(relativePath, wrapHandlers(handlers...)...)
  129. }
  130. func (r *router) DELETE(relativePath string, handlers ...HandlerFunc) {
  131. r.group.DELETE(relativePath, wrapHandlers(handlers...)...)
  132. }
  133. func (r *router) PATCH(relativePath string, handlers ...HandlerFunc) {
  134. r.group.PATCH(relativePath, wrapHandlers(handlers...)...)
  135. }
  136. func (r *router) PUT(relativePath string, handlers ...HandlerFunc) {
  137. r.group.PUT(relativePath, wrapHandlers(handlers...)...)
  138. }
  139. func (r *router) OPTIONS(relativePath string, handlers ...HandlerFunc) {
  140. r.group.OPTIONS(relativePath, wrapHandlers(handlers...)...)
  141. }
  142. func (r *router) HEAD(relativePath string, handlers ...HandlerFunc) {
  143. r.group.HEAD(relativePath, wrapHandlers(handlers...)...)
  144. }
  145. func wrapHandlers(handlers ...HandlerFunc) []gin.HandlerFunc {
  146. list := make([]gin.HandlerFunc, len(handlers))
  147. for i, handler := range handlers {
  148. fn := handler
  149. list[i] = func(c *gin.Context) {
  150. ctx := newContext(c)
  151. defer releaseContext(ctx)
  152. fn(ctx)
  153. }
  154. }
  155. return list
  156. }
  157. /******************************************************************************/
  158. var _ Mux = (*mux)(nil)
  159. type Mux interface {
  160. http.Handler
  161. Group(relativePath string, handlers ...HandlerFunc) RouterGroup
  162. Routes() gin.RoutesInfo
  163. HandlerFunc(relativePath string, handlerFunc gin.HandlerFunc)
  164. }
  165. type mux struct {
  166. engine *gin.Engine
  167. }
  168. func (m *mux) ServeHTTP(w http.ResponseWriter, req *http.Request) {
  169. m.engine.ServeHTTP(w, req)
  170. }
  171. func (m *mux) Group(relativePath string, handlers ...HandlerFunc) RouterGroup {
  172. return &router{
  173. group: m.engine.Group(relativePath, wrapHandlers(handlers...)...),
  174. }
  175. }
  176. func (m *mux) Routes() gin.RoutesInfo {
  177. return m.engine.Routes()
  178. }
  179. func (m *mux) HandlerFunc(relativePath string, handlerFunc gin.HandlerFunc) {
  180. m.engine.GET(relativePath, handlerFunc)
  181. }
  182. func New(logger *zap.Logger, options ...Option) (Mux, error) {
  183. if logger == nil {
  184. return nil, errors.New("logger required")
  185. }
  186. gin.SetMode(gin.ReleaseMode)
  187. gin.DisableBindValidation()
  188. newMux := &mux{
  189. engine: gin.New(),
  190. }
  191. fmt.Println(color.Green(fmt.Sprintf("* [register env %s]", env.Active().Value())))
  192. // withoutLogPaths 这些请求,默认不记录日志
  193. withoutTracePaths := map[string]bool{
  194. "/metrics": true,
  195. "/favicon.ico": true,
  196. "/system/health": true,
  197. }
  198. opt := new(option)
  199. for _, f := range options {
  200. f(opt)
  201. }
  202. if opt.enablePProf {
  203. pprof.Register(newMux.engine)
  204. fmt.Println(color.Green("* [register pprof]"))
  205. }
  206. if opt.enablePrometheus {
  207. newMux.engine.GET("/metrics", gin.WrapH(promhttp.Handler()))
  208. fmt.Println(color.Green("* [register prometheus]"))
  209. }
  210. if opt.enableCors {
  211. newMux.engine.Use(cors.AllowAll())
  212. }
  213. if opt.staticDirs != nil {
  214. for _, dir := range opt.staticDirs {
  215. newMux.engine.StaticFS(dir, gin.Dir(dir, false))
  216. }
  217. }
  218. // recover两次,防止处理时发生panic,尤其是在OnPanicNotify中。
  219. newMux.engine.Use(func(ctx *gin.Context) {
  220. defer func() {
  221. if err := recover(); err != nil {
  222. logger.Error("got panic", zap.String("panic", fmt.Sprintf("%+v", err)), zap.String("stack", string(debug.Stack())))
  223. }
  224. }()
  225. ctx.Next()
  226. })
  227. newMux.engine.Use(func(ctx *gin.Context) {
  228. ts := time.Now()
  229. newCtx := newContext(ctx)
  230. defer releaseContext(newCtx)
  231. newCtx.init()
  232. newCtx.setLogger(logger)
  233. if !withoutTracePaths[ctx.Request.URL.Path] {
  234. if traceId := newCtx.GetHeader(trace.Header); traceId != "" {
  235. newCtx.setTrace(trace.New(traceId))
  236. } else {
  237. newCtx.setTrace(trace.New(""))
  238. }
  239. }
  240. defer func() {
  241. if err := recover(); err != nil {
  242. stackInfo := string(debug.Stack())
  243. logger.Error("got panic", zap.String("panic", fmt.Sprintf("%+v", err)), zap.String("stack", stackInfo))
  244. newCtx.AbortWithError(errno.NewError(
  245. http.StatusInternalServerError,
  246. http.StatusInternalServerError,
  247. http.StatusText(http.StatusInternalServerError)),
  248. )
  249. if notify := opt.panicNotify; notify != nil {
  250. notify(newCtx, err, stackInfo)
  251. }
  252. }
  253. if ctx.Writer.Status() == http.StatusNotFound {
  254. return
  255. }
  256. var (
  257. response any
  258. businessCode int
  259. businessCodeMsg string
  260. abortErr error
  261. graphResponse any
  262. )
  263. if ctx.IsAborted() {
  264. for i := range ctx.Errors { // gin error
  265. multierr.AppendInto(&abortErr, ctx.Errors[i])
  266. }
  267. if err := newCtx.abortError(); err != nil { // customer err
  268. multierr.AppendInto(&abortErr, err.GetErr())
  269. response = err
  270. businessCode = err.GetBusinessCode()
  271. businessCodeMsg = err.GetMsg()
  272. if x := newCtx.Trace(); x != nil {
  273. newCtx.SetHeader(trace.Header, x.ID())
  274. }
  275. ctx.JSON(err.GetHttpCode(), &Failure{
  276. ResultCode: businessCode,
  277. ResultInfo: businessCodeMsg,
  278. })
  279. }
  280. } else {
  281. response = newCtx.getPayload()
  282. if response != nil {
  283. if x := newCtx.Trace(); x != nil {
  284. newCtx.SetHeader(trace.Header, x.ID())
  285. }
  286. ctx.JSON(http.StatusOK, response)
  287. }
  288. }
  289. graphResponse = newCtx.getGraphPayload()
  290. if opt.recordMetrics != nil {
  291. uri := newCtx.Path()
  292. if alias := newCtx.Alias(); alias != "" {
  293. uri = alias
  294. }
  295. opt.recordMetrics(
  296. newCtx.Method(),
  297. uri,
  298. !ctx.IsAborted() && ctx.Writer.Status() == http.StatusOK,
  299. time.Since(ts).Seconds(),
  300. )
  301. }
  302. var t *trace.Trace
  303. if x := newCtx.Trace(); x != nil {
  304. t = x.(*trace.Trace)
  305. } else {
  306. return
  307. }
  308. decodedURL, _ := url.QueryUnescape(ctx.Request.URL.RequestURI())
  309. t.WithRequest(&trace.Request{
  310. TTL: "un-limit",
  311. Method: ctx.Request.Method,
  312. DecodedURL: decodedURL,
  313. Header: ctx.Request.Header,
  314. Body: string(newCtx.RawData()),
  315. })
  316. var responseBody any
  317. if response != nil {
  318. responseBody = response
  319. }
  320. if graphResponse != nil {
  321. responseBody = graphResponse
  322. }
  323. t.WithResponse(&trace.Response{
  324. Header: ctx.Writer.Header(),
  325. HttpCode: ctx.Writer.Status(),
  326. HttpCodeMsg: http.StatusText(ctx.Writer.Status()),
  327. BusinessCode: businessCode,
  328. BusinessCodeMsg: businessCodeMsg,
  329. Body: responseBody,
  330. CostSeconds: time.Since(ts).Seconds(),
  331. })
  332. t.Success = !ctx.IsAborted() && ctx.Writer.Status() == http.StatusOK
  333. t.CostSeconds = time.Since(ts).Seconds()
  334. logger.Info("core-interceptor",
  335. zap.Any("method", ctx.Request.Method),
  336. zap.Any("path", decodedURL),
  337. zap.Any("http_code", ctx.Writer.Status()),
  338. zap.Any("business_code", businessCode),
  339. zap.Any("success", t.Success),
  340. zap.Any("cost_seconds", t.CostSeconds),
  341. zap.Any("trace_id", t.Identifier),
  342. zap.Any("trace_info", t),
  343. zap.Error(abortErr),
  344. )
  345. }()
  346. ctx.Next()
  347. })
  348. if opt.enableRate {
  349. newMux.engine.Use(func(ctx *gin.Context) {
  350. newCtx := newContext(ctx)
  351. defer releaseContext(newCtx)
  352. if !limiter.Allow() {
  353. newCtx.AbortWithError(errno.NewError(
  354. http.StatusTooManyRequests,
  355. http.StatusTooManyRequests,
  356. http.StatusText(http.StatusTooManyRequests)),
  357. )
  358. return
  359. }
  360. ctx.Next()
  361. })
  362. }
  363. newMux.engine.NoMethod(wrapHandlers(DisableTrace)...)
  364. newMux.engine.NoRoute(wrapHandlers(DisableTrace)...)
  365. system := newMux.Group("/system")
  366. {
  367. // 健康检查
  368. system.GET("/health", func(ctx Context) {
  369. resp := &struct {
  370. Timestamp time.Time `json:"timestamp"`
  371. Environment string `json:"environment"`
  372. Host string `json:"host"`
  373. Status string `json:"status"`
  374. }{
  375. Timestamp: time.Now(),
  376. Environment: env.Active().Value(),
  377. Host: ctx.Host(),
  378. Status: "ok",
  379. }
  380. ctx.Payload(resp)
  381. })
  382. }
  383. return newMux, nil
  384. }