Skip to content

O
opkg
提交 d57375cc authored 作者: mooncake's avatar mooncake
浏览文件
操作

update

上级 54ae8fdb
全部展开 显示空白字符变更
内嵌 并排
正在显示 包含 3662 行增加 和 8 行删除
gin/middleware/auth.go 0 → 100644
浏览文件 @ d57375cc
package middleware
import (
"github.com/gin-gonic/gin"
"gitlab.wanzhuangkj.com/tush/opkg/gin/response"
"gitlab.wanzhuangkj.com/tush/opkg/pkg/errcode"
"gitlab.wanzhuangkj.com/tush/opkg/pkg/jwt"
"gitlab.wanzhuangkj.com/tush/opkg/pkg/logger"
"gitlab.wanzhuangkj.com/tush/opkg/pkg/sf"
ctxutils "gitlab.wanzhuangkj.com/tush/opkg/xutils/ctxutils"
)
const (
// HeaderAuthorizationKey http header authorization key
HeaderAuthorizationKey = "Authorization"
)
type jwtOptions struct {
isSwitchHTTPCode bool
verify VerifyFn // verify function, only use in Auth
}
// JwtOption set the jwt options.
type JwtOption func(*jwtOptions)
func (o *jwtOptions) apply(opts ...JwtOption) {
for _, opt := range opts {
opt(o)
}
}
func defaultJwtOptions() *jwtOptions {
return &jwtOptions{
isSwitchHTTPCode: false,
verify: nil,
}
}
// WithSwitchHTTPCode switch to http code
func WithSwitchHTTPCode() JwtOption {
return func(o *jwtOptions) {
o.isSwitchHTTPCode = true
}
}
// WithVerify set verify function
func WithVerify(verify VerifyFn) JwtOption {
return func(o *jwtOptions) {
o.verify = verify
}
}
func responseUnauthorized(c *gin.Context, isSwitchHTTPCode bool) {
if isSwitchHTTPCode {
response.Out(c, errcode.Unauthorized)
} else {
response.ErrorE(c, errcode.Unauthorized)
}
}
// -------------------------------------------------------------------------------------------
// VerifyFn verify function, tokenTail10 is the last 10 characters of the token.
type VerifyFn func(claims *jwt.Claims, tokenTail10 string, c *gin.Context) error
// Auth authorization
func Auth(opts ...JwtOption) gin.HandlerFunc {
o := defaultJwtOptions()
o.apply(opts...)
return func(c *gin.Context) {
authorization := c.GetHeader(HeaderAuthorizationKey)
if authorization == "" || len(authorization) < 10 {
responseUnauthorized(c, o.isSwitchHTTPCode)
c.Abort()
return
}
token := authorization[7:] // remove Bearer prefix
claims, err := jwt.ParseToken(token)
if err != nil {
logger.Error("parseToken error", logger.Err(err), ctxutils.GinTraceIDField(c))
responseUnauthorized(c, o.isSwitchHTTPCode)
c.Abort()
return
}
if o.verify != nil {
tokenTail10 := token[len(token)-10:]
if err = o.verify(claims, tokenTail10, c); err != nil {
logger.Error("verify token error", logger.Err(err), logger.Any("claims", claims), ctxutils.GinTraceIDField(c))
responseUnauthorized(c, o.isSwitchHTTPCode)
c.Abort()
return
}
} else {
xsfID, _ := sf.ParseString(claims.UID)
c.Set(ctxutils.KeyUID, xsfID)
c.Set(ctxutils.KeyUName, claims.Name)
}
c.Next()
}
}
func URLFields() gin.HandlerFunc {
return func(c *gin.Context) {
params := c.Request.URL.Query()
for key, values := range params {
if len(values) > 0 {
c.Set(key, values[0])
}
}
}
}
type KV = map[string]interface{}
func HeaderFields(keys []string) gin.HandlerFunc {
return func(c *gin.Context) {
kv := KV{}
for _, key := range keys {
c.Set(key, c.Request.Header.Get(key))
kv[key] = c.Request.Header.Get(key)
}
logger.Info("[middleware][headerFields]header kvs", logger.Any("kv", kv), ctxutils.GinTraceIDField(c))
}
}
// -------------------------------------------------------------------------------------------
// VerifyCustomFn verify custom function, tokenTail10 is the last 10 characters of the token.
type VerifyCustomFn func(claims *jwt.CustomClaims, tokenTail10 string, c *gin.Context) error
// AuthCustom custom authentication
func AuthCustom(verify VerifyCustomFn, opts ...JwtOption) gin.HandlerFunc {
o := defaultJwtOptions()
o.apply(opts...)
return func(c *gin.Context) {
authorization := c.GetHeader(HeaderAuthorizationKey)
if authorization == "" || len(authorization) < 10 {
responseUnauthorized(c, o.isSwitchHTTPCode)
c.Abort()
return
}
token := authorization[7:] // remove Bearer prefix
claims, err := jwt.ParseCustomToken(token)
if err != nil {
logger.Error("parseToken error", logger.Err(err), ctxutils.GinTraceIDField(c))
responseUnauthorized(c, o.isSwitchHTTPCode)
c.Abort()
return
}
tokenTail10 := token[len(token)-10:]
if err = verify(claims, tokenTail10, c); err != nil {
logger.Error("verify token error", logger.Err(err), logger.Any("fields", claims.Fields), ctxutils.GinTraceIDField(c))
responseUnauthorized(c, o.isSwitchHTTPCode)
c.Abort()
return
}
c.Next()
}
}
gin/middleware/breaker.go 0 → 100644
浏览文件 @ d57375cc
package middleware
import (
"net/http"
"github.com/gin-gonic/gin"
"gitlab.wanzhuangkj.com/tush/opkg/gin/response"
"gitlab.wanzhuangkj.com/tush/opkg/pkg/container/group"
"gitlab.wanzhuangkj.com/tush/opkg/pkg/shield/circuitbreaker"
)
// ErrNotAllowed error not allowed.
var ErrNotAllowed = circuitbreaker.ErrNotAllowed
// CircuitBreakerOption set the circuit breaker circuitBreakerOptions.
type CircuitBreakerOption func(*circuitBreakerOptions)
type circuitBreakerOptions struct {
group *group.Group
// http code for circuit breaker, default already includes 500 and 503
validCodes map[int]struct{}
// degrade func
degradeHandler func(c *gin.Context)
}
func defaultCircuitBreakerOptions() *circuitBreakerOptions {
return &circuitBreakerOptions{
group: group.NewGroup(func() interface{} {
return circuitbreaker.NewBreaker()
}),
validCodes: map[int]struct{}{
http.StatusInternalServerError: {},
http.StatusServiceUnavailable: {},
},
}
}
func (o *circuitBreakerOptions) apply(opts ...CircuitBreakerOption) {
for _, opt := range opts {
opt(o)
}
}
// WithGroup with circuit breaker group.
// NOTE: implements generics circuitbreaker.CircuitBreaker
func WithGroup(g *group.Group) CircuitBreakerOption {
return func(o *circuitBreakerOptions) {
if g != nil {
o.group = g
}
}
}
// WithValidCode http code to mark failed
func WithValidCode(code ...int) CircuitBreakerOption {
return func(o *circuitBreakerOptions) {
for _, c := range code {
o.validCodes[c] = struct{}{}
}
}
}
// WithDegradeHandler set degrade handler function
func WithDegradeHandler(handler func(c *gin.Context)) CircuitBreakerOption {
return func(o *circuitBreakerOptions) {
o.degradeHandler = handler
}
}
// CircuitBreaker a circuit breaker middleware
func CircuitBreaker(opts ...CircuitBreakerOption) gin.HandlerFunc {
o := defaultCircuitBreakerOptions()
o.apply(opts...)
return func(c *gin.Context) {
breaker := o.group.Get(c.FullPath()).(circuitbreaker.CircuitBreaker)
if err := breaker.Allow(); err != nil {
// NOTE: when client reject request locally, keep adding counter let the drop ratio higher.
breaker.MarkFailed()
if o.degradeHandler != nil {
o.degradeHandler(c)
} else {
response.Output(c, http.StatusServiceUnavailable, err.Error())
}
c.Abort()
return
}
c.Next()
code := c.Writer.Status()
// NOTE: need to check internal and service unavailable error
_, isHit := o.validCodes[code]
if isHit {
breaker.MarkFailed()
} else {
breaker.MarkSuccess()
}
}
}
gin/middleware/cors.go 0 → 100644
浏览文件 @ d57375cc
package middleware
import (
"time"
"github.com/gin-contrib/cors"
"github.com/gin-gonic/gin"
)
type corsOptions struct {
allowOrigins []string
}
type CorsOption func(*corsOptions)
func defaultCorsOptions() *corsOptions {
return &corsOptions{
allowOrigins: []string{"*"},
}
}
func (o *corsOptions) apply(opts ...CorsOption) {
for _, opt := range opts {
opt(o)
}
}
func WithCorsOption(allowOrigins []string) CorsOption {
return func(o *corsOptions) {
o.allowOrigins = allowOrigins
}
}
// Cors cross domain
func Cors(opts ...CorsOption) gin.HandlerFunc {
o := defaultCorsOptions()
o.apply(opts...)
return cors.New(
cors.Config{
AllowOrigins: o.allowOrigins,
AllowMethods: []string{"GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS"},
AllowHeaders: []string{"Origin", "Authorization", "Content-Type", "Accept"},
ExposeHeaders: []string{"Content-Length", "text/plain", "Authorization", "Content-Type"},
AllowCredentials: true,
MaxAge: 12 * time.Hour,
},
)
}
// func CorsQT() gin.HandlerFunc {
// return cors.New(
// cors.Config{
// AllowOrigins: []string{"http://wanzhuangkj.com", "https://wanzhuangkj.com", "http://hzqitu.com", "https://hzqitu.com"},
// AllowMethods: []string{"GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS"},
// AllowHeaders: []string{"Origin", "Authorization", "Content-Type", "Accept"},
// ExposeHeaders: []string{"Content-Length", "text/plain", "Authorization", "Content-Type"},
// AllowCredentials: true,
// MaxAge: 12 * time.Hour,
// },
// )
// }
gin/middleware/logging.go 0 → 100644
浏览文件 @ d57375cc
package middleware
import (
"bytes"
"io"
"net/http"
"strings"
"time"
"gitlab.wanzhuangkj.com/tush/opkg/pkg/httpcli/entity"
ctxutils "gitlab.wanzhuangkj.com/tush/opkg/xutils/ctxutils"
"gitlab.wanzhuangkj.com/tush/opkg/pkg/ips"
"github.com/gin-gonic/gin"
"go.uber.org/zap"
)
var (
// Print body max length
defaultMaxLength = 300
defaultMediaMaxLength = 300
// default zap log
defaultLogger, _ = zap.NewProduction()
// Ignore route list
defaultIgnoreRoutes = map[string]struct{}{
"/ping": {},
"/pong": {},
"/health": {},
}
emptyBody = []byte("")
contentMark = []byte(" ...... ")
)
// Option set the gin logger options.
type Option func(*options)
func defaultOptions() *options {
return &options{
maxLength: defaultMaxLength,
mediaMaxLength: defaultMediaMaxLength,
log: defaultLogger,
ignoreRoutes: defaultIgnoreRoutes,
traceIDFrom: 0,
}
}
type options struct {
maxLength int
mediaMaxLength int
appName string
log *zap.Logger
ignoreRoutes map[string]struct{}
traceIDFrom int // 0: ignore, 1: from context, 2: from header
}
func (o *options) apply(opts ...Option) {
for _, opt := range opts {
opt(o)
}
}
// WithMaxLen logger content max length
func WithMaxLen(maxLen int) Option {
return func(o *options) {
if maxLen < len(contentMark) {
panic("maxLen should be greater than or equal to 8")
}
o.maxLength = maxLen
}
}
func WithMediaMaxLen(maxLen int) Option {
return func(o *options) {
if maxLen < len(contentMark) {
panic("mediaMaxLength should be greater than or equal to 8")
}
o.mediaMaxLength = maxLen
}
}
// WithLog set log
func WithLog(log *zap.Logger) Option {
return func(o *options) {
if log != nil {
o.log = log
}
}
}
func WithAppName(appName string) Option {
return func(o *options) {
o.appName = appName
}
}
// WithIgnoreRoutes no logger content routes
func WithIgnoreRoutes(routes ...string) Option {
return func(o *options) {
for _, route := range routes {
o.ignoreRoutes[route] = struct{}{}
}
}
}
// WithTraceIDFromContext name is field in context, default value is request_id
func WithTraceIDFromContext() Option {
return func(o *options) {
o.traceIDFrom = 1
}
}
// WithTraceIDFromHeader name is field in header, default value is X-Request-Id
func WithTraceIDFromHeader() Option {
return func(o *options) {
o.traceIDFrom = 2
}
}
// ------------------------------------------------------------------------------------------
type bodyLogWriter struct {
gin.ResponseWriter
body *bytes.Buffer
}
func (w bodyLogWriter) Write(b []byte) (int, error) {
w.body.Write(b)
return w.ResponseWriter.Write(b)
}
// If there is sensitive information in the body, you can use WithIgnoreRoutes set the route to ignore logging
func getResponseBody(buf *bytes.Buffer, maxLen int) []byte {
l := buf.Len()
if l == 0 {
return []byte("")
} else if l > maxLen {
l = maxLen
}
body := make([]byte, l)
n, _ := buf.Read(body)
if n == 0 {
return emptyBody
} else if n < maxLen {
if isLastByteNewline(body) {
return body[:n-1]
}
return body[:n]
}
return append(body[:maxLen-len(contentMark)], contentMark...)
}
func isLastByteNewline(data []byte) bool {
if len(data) == 0 {
return false
}
return data[len(data)-1] == '\n'
}
// If there is sensitive information in the body, you can use WithIgnoreRoutes set the route to ignore logging
func getRequestBody(buf *bytes.Buffer, maxLen int) []byte {
l := buf.Len()
if l == 0 {
return []byte("")
} else if l < maxLen {
return buf.Bytes()
}
body := make([]byte, maxLen)
copy(body, buf.Bytes())
return append(body[:maxLen-len(contentMark)], contentMark...)
}
// Logging print request and response info
func Logging(opts ...Option) gin.HandlerFunc {
o := defaultOptions()
o.apply(opts...)
return func(c *gin.Context) {
start := time.Now()
// ignore printing of the specified route
if _, ok := o.ignoreRoutes[c.Request.URL.Path]; ok {
c.Next()
return
}
// print input information before processing
buf := &bytes.Buffer{}
_, _ = buf.ReadFrom(c.Request.Body)
fields := []zap.Field{
zap.String("method", c.Request.Method),
zap.String("url", c.Request.URL.String()),
zap.Any("headers", c.Request.Header),
}
if c.Request.Method == http.MethodPost || c.Request.Method == http.MethodPut || c.Request.Method == http.MethodPatch || c.Request.Method == http.MethodDelete {
fields = append(fields,
zap.Int("size", buf.Len()),
zap.ByteString("body", getRequestBody(buf, o.maxLength)),
)
}
entity.SetCopyReq(c, buf)
c.Set(ctxutils.KeyAppName, o.appName)
reqID := ""
if o.traceIDFrom == 1 {
if v, isExist := c.Get(ctxutils.ContextTraceIDKey); isExist {
if requestID, ok := v.(string); ok {
reqID = requestID
fields = append(fields, zap.String(ctxutils.ContextTraceIDKey, reqID))
}
}
} else if o.traceIDFrom == 2 {
reqID = c.Request.Header.Get(ctxutils.HeaderXRequestIDKey)
fields = append(fields, zap.String(ctxutils.ContextTraceIDKey, reqID))
}
o.log.Info("[middleware][logging]<<<<<<<<<req", fields...)
c.Request.Body = io.NopCloser(buf)
//replace writer
newWriter := &bodyLogWriter{body: &bytes.Buffer{}, ResponseWriter: c.Writer}
c.Writer = newWriter
ip := ips.GetClientIP(c)
c.Set(ctxutils.KeyClientIP, ip)
// processing requests
c.Next()
rspBodyMax := c.GetInt(ctxutils.KeyRspBodyMax)
contentType := c.Writer.Header().Get("Content-Type")
isMedia := false
if strings.Contains(contentType, "image") {
isMedia = true
}
//print return message after processing
fields = []zap.Field{
zap.Int("code", c.Writer.Status()),
zap.String("cost", time.Since(start).String()),
zap.String("method", c.Request.Method),
zap.String("url", c.Request.URL.Path),
zap.String("client_ip", ip),
zap.Int("size", newWriter.body.Len()),
zap.Any("header", c.Writer.Header()),
}
if isMedia {
fields = append(fields, zap.ByteString("body", getResponseBody(newWriter.body, o.mediaMaxLength)))
} else {
if rspBodyMax > 0 {
fields = append(fields, zap.ByteString("body", getResponseBody(newWriter.body, rspBodyMax)))
} else {
fields = append(fields, zap.ByteString("body", getResponseBody(newWriter.body, o.maxLength)))
}
}
if reqID != "" {
fields = append(fields, zap.String(ctxutils.ContextTraceIDKey, reqID))
}
o.log.Info("[middleware][logging]>>>>>>>>>rsp", fields...)
}
}
// SimpleLog print response info
func SimpleLog(opts ...Option) gin.HandlerFunc {
o := defaultOptions()
o.apply(opts...)
return func(c *gin.Context) {
start := time.Now()
// ignore printing of the specified route
if _, ok := o.ignoreRoutes[c.Request.URL.Path]; ok {
c.Next()
return
}
reqID := ""
switch o.traceIDFrom {
case 1:
if v, isExist := c.Get(ctxutils.ContextTraceIDKey); isExist {
if requestID, ok := v.(string); ok {
reqID = requestID
}
}
case 2:
reqID = c.Request.Header.Get(ctxutils.HeaderXRequestIDKey)
}
// processing requests
c.Next()
// print return message after processing
fields := []zap.Field{
zap.Int("code", c.Writer.Status()),
zap.String("method", c.Request.Method),
zap.String("url", c.Request.URL.String()),
zap.String("cost", time.Since(start).String()),
zap.Int("size", c.Writer.Size()),
}
if reqID != "" {
fields = append(fields, zap.String(ctxutils.ContextTraceIDKey, reqID))
}
o.log.Info("[middleware][logging]]>>>>>>>>>rsp", fields...)
}
}
gin/middleware/metrics/gin_grafana.json 0 → 100644
浏览文件 @ d57375cc
差异被折叠。
gin/middleware/metrics/metrics.go 0 → 100644
浏览文件 @ d57375cc
// Package metrics is gin metrics library, collect five metrics, "uptime", "http_request_count_total",
// "http_request_duration_seconds", "http_request_size_bytes", "http_response_size_bytes".
package metrics
import (
"net/http"
"strconv"
"time"
"github.com/gin-gonic/gin"
"github.com/prometheus/client_golang/prometheus"
"github.com/prometheus/client_golang/prometheus/promhttp"
)
var (
namespace = "gin"
labels = []string{"status", "path", "method"}
uptime = prometheus.NewCounterVec(
prometheus.CounterOpts{
Namespace: namespace,
Name: "uptime",
Help: "HTTP service uptime, updated every minute",
}, nil,
)
reqCount = prometheus.NewCounterVec(
prometheus.CounterOpts{
Namespace: namespace,
Name: "http_request_count_total",
Help: "Total number of HTTP requests made.",
}, labels,
)
reqDuration = prometheus.NewHistogramVec(
prometheus.HistogramOpts{
Namespace: namespace,
Name: "http_request_duration_seconds",
Help: "HTTP request latencies in seconds.",
}, labels,
)
reqSizeBytes = prometheus.NewSummaryVec(
prometheus.SummaryOpts{
Namespace: namespace,
Name: "http_request_size_bytes",
Help: "HTTP request sizes in bytes.",
}, labels,
)
respSizeBytes = prometheus.NewSummaryVec(
prometheus.SummaryOpts{
Namespace: namespace,
Name: "http_response_size_bytes",
Help: "HTTP response sizes in bytes.",
}, labels,
)
)
// init registers the prometheus metrics
func initPrometheus() {
prometheus.MustRegister(uptime, reqCount, reqDuration, reqSizeBytes, respSizeBytes)
go recordUptime()
}
// recordUptime increases service uptime per 1 minute.
func recordUptime() {
for range time.Tick(time.Minute) {
uptime.WithLabelValues().Inc()
}
}
// calcRequestSize returns the size of request object.
func calcRequestSize(r *http.Request) float64 {
size := 0
if r.URL != nil {
size = len(r.URL.String())
}
size += len(r.Method)
size += len(r.Proto)
for name, values := range r.Header {
size += len(name)
for _, value := range values {
size += len(value)
}
}
size += len(r.Host)
// r.Form and r.MultipartForm are assumed to be included in r.URL.
if r.ContentLength != -1 {
size += int(r.ContentLength)
}
return float64(size)
}
// ------------------------------------------------------------------------------------------
// metricsHandler wrappers the standard http.Handler to gin.HandlerFunc
func metricsHandler() gin.HandlerFunc {
return func(c *gin.Context) {
handler := promhttp.Handler()
handler.ServeHTTP(c.Writer, c.Request)
}
}
// Metrics returns a gin.HandlerFunc for exporting some Web metrics
func Metrics(r *gin.Engine, opts ...Option) gin.HandlerFunc {
o := defaultOptions()
o.apply(opts...)
// init prometheus
initPrometheus()
r.GET(o.metricsPath, metricsHandler())
return func(c *gin.Context) {
start := time.Now()
c.Next()
ok := o.isIgnoreCodeStatus(c.Writer.Status()) ||
o.isIgnorePath(c.Request.URL.Path) ||
o.checkIgnoreMethod(c.Request.Method)
if ok {
return
}
// no response content will return -1
respSize := c.Writer.Size()
if respSize < 0 {
respSize = 0
}
lvs := []string{strconv.Itoa(c.Writer.Status()), c.Request.URL.Path, c.Request.Method}
reqCount.WithLabelValues(lvs...).Inc()
reqDuration.WithLabelValues(lvs...).Observe(time.Since(start).Seconds())
reqSizeBytes.WithLabelValues(lvs...).Observe(calcRequestSize(c.Request))
respSizeBytes.WithLabelValues(lvs...).Observe(float64(respSize))
}
}
gin/middleware/metrics/option.go 0 → 100644
浏览文件 @ d57375cc
package metrics
import (
"strings"
)
// Option set the metrics options.
type Option func(*options)
type options struct {
metricsPath string
ignoreStatusCodes map[int]struct{}
ignoreRequestPaths map[string]struct{}
ignoreRequestMethods map[string]struct{}
}
// defaultOptions default value
func defaultOptions() *options {
return &options{
metricsPath: "/metrics",
ignoreStatusCodes: nil,
ignoreRequestPaths: nil,
ignoreRequestMethods: nil,
}
}
func (o *options) apply(opts ...Option) {
for _, opt := range opts {
opt(o)
}
}
// WithMetricsPath set metrics path
func WithMetricsPath(metricsPath string) Option {
return func(o *options) {
o.metricsPath = metricsPath
}
}
// WithIgnoreStatusCodes ignore status codes
func WithIgnoreStatusCodes(statusCodes ...int) Option {
return func(o *options) {
codeMaps := make(map[int]struct{}, len(statusCodes))
for _, code := range statusCodes {
codeMaps[code] = struct{}{}
}
o.ignoreStatusCodes = codeMaps
}
}
// WithIgnoreRequestPaths ignore request paths
func WithIgnoreRequestPaths(paths ...string) Option {
return func(o *options) {
pathMaps := make(map[string]struct{}, len(paths))
for _, path := range paths {
pathMaps[path] = struct{}{}
}
o.ignoreRequestPaths = pathMaps
}
}
// WithIgnoreRequestMethods ignore request methods
func WithIgnoreRequestMethods(methods ...string) Option {
return func(o *options) {
methodMaps := make(map[string]struct{}, len(methods))
for _, method := range methods {
methodMaps[strings.ToUpper(method)] = struct{}{}
}
o.ignoreRequestMethods = methodMaps
}
}
func (o *options) isIgnoreCodeStatus(statusCode int) bool {
if o.ignoreStatusCodes == nil {
return false
}
_, ok := o.ignoreStatusCodes[statusCode]
return ok
}
func (o *options) isIgnorePath(path string) bool {
if o.ignoreRequestPaths == nil {
return false
}
_, ok := o.ignoreRequestPaths[path]
return ok
}
func (o *options) checkIgnoreMethod(method string) bool {
if o.ignoreRequestMethods == nil {
return false
}
_, ok := o.ignoreRequestMethods[strings.ToUpper(method)]
return ok
}
gin/middleware/metrics_wz/metrics.go 0 → 100644
浏览文件 @ d57375cc
package metrics
import (
"os"
"time"
"github.com/gin-gonic/gin"
"github.com/prometheus/client_golang/prometheus"
"github.com/prometheus/client_golang/prometheus/promhttp"
"github.com/spf13/cast"
)
var (
// namespace = "xmall"
namespace = ""
// uptime = prometheus.NewCounterVec(
// prometheus.CounterOpts{
// Namespace: namespace,
// Name: "uptime",
// Help: "HTTP service uptime, updated every minute",
// }, nil,
// )
reqCount = prometheus.NewCounterVec(
prometheus.CounterOpts{
Namespace: namespace,
Name: "client_code_total",
Help: "Total number of HTTP requests made.",
}, []string{"kind", "operation", "code"},
)
reqDuration = prometheus.NewHistogramVec(prometheus.HistogramOpts{
Namespace: namespace,
Name: "server_requests_duration_sec",
Help: "Server requests duratio(sec).",
Buckets: []float64{0.005, 0.01, 0.025, 0.05, 0.1, 0.250, 0.5, 1},
}, []string{"kind", "operation"})
)
func init() {
if ns := os.Getenv("CONTAINER_NAMESPACE"); ns != "" {
namespace = ns
}
}
// init registers the prometheus metrics
func initPrometheus() {
// prometheus.MustRegister(uptime, reqCount, reqDuration)
prometheus.MustRegister(reqCount, reqDuration)
}
// recordUptime increases service uptime per 1 minute.
// func recordUptime() {
// for range time.Tick(time.Minute) {
// uptime.WithLabelValues().Inc()
// }
// }
// metricsHandler wrappers the standard http.Handler to gin.HandlerFunc
func metricsHandler() gin.HandlerFunc {
return func(c *gin.Context) {
handler := promhttp.Handler()
handler.ServeHTTP(c.Writer, c.Request)
}
}
// Metrics returns a gin.HandlerFunc for exporting some Web metrics
func Metrics(r *gin.Engine, opts ...Option) gin.HandlerFunc {
o := defaultOptions()
o.apply(opts...)
// init prometheus
initPrometheus()
r.GET(o.metricsPath, metricsHandler())
return func(c *gin.Context) {
start := time.Now()
c.Next()
ok := o.isIgnoreCodeStatus(c.Writer.Status()) ||
o.isIgnorePath(c.Request.URL.Path) ||
o.checkIgnoreMethod(c.Request.Method)
if ok {
return
}
reqCount.With(prometheus.Labels{
"kind": "http",
"operation": c.Request.Method + "." + c.Request.URL.Path,
"code": cast.ToString(c.Writer.Status()),
}).Inc()
reqDuration.With(prometheus.Labels{
"kind": "http",
"operation": c.Request.Method + "." + c.Request.URL.Path,
}).Observe(time.Since(start).Seconds())
}
}
gin/middleware/metrics_wz/option.go 0 → 100644
浏览文件 @ d57375cc
package metrics
import (
"strings"
)
// Option set the metrics options.
type Option func(*options)
type options struct {
metricsPath string
ignoreStatusCodes map[int]struct{}
ignoreRequestPaths map[string]struct{}
ignoreRequestMethods map[string]struct{}
}
// defaultOptions default value
func defaultOptions() *options {
return &options{
metricsPath: "/metrics",
ignoreStatusCodes: nil,
ignoreRequestPaths: nil,
ignoreRequestMethods: nil,
}
}
func (o *options) apply(opts ...Option) {
for _, opt := range opts {
opt(o)
}
}
// WithMetricsPath set metrics path
func WithMetricsPath(metricsPath string) Option {
return func(o *options) {
o.metricsPath = metricsPath
}
}
// WithIgnoreStatusCodes ignore status codes
func WithIgnoreStatusCodes(statusCodes ...int) Option {
return func(o *options) {
codeMaps := make(map[int]struct{}, len(statusCodes))
for _, code := range statusCodes {
codeMaps[code] = struct{}{}
}
o.ignoreStatusCodes = codeMaps
}
}
// WithIgnoreRequestPaths ignore request paths
func WithIgnoreRequestPaths(paths ...string) Option {
return func(o *options) {
pathMaps := make(map[string]struct{}, len(paths))
for _, path := range paths {
pathMaps[path] = struct{}{}
}
o.ignoreRequestPaths = pathMaps
}
}
// WithIgnoreRequestMethods ignore request methods
func WithIgnoreRequestMethods(methods ...string) Option {
return func(o *options) {
methodMaps := make(map[string]struct{}, len(methods))
for _, method := range methods {
methodMaps[strings.ToUpper(method)] = struct{}{}
}
o.ignoreRequestMethods = methodMaps
}
}
func (o *options) isIgnoreCodeStatus(statusCode int) bool {
if o.ignoreStatusCodes == nil {
return false
}
_, ok := o.ignoreStatusCodes[statusCode]
return ok
}
func (o *options) isIgnorePath(path string) bool {
if o.ignoreRequestPaths == nil {
return false
}
_, ok := o.ignoreRequestPaths[path]
return ok
}
func (o *options) checkIgnoreMethod(method string) bool {
if o.ignoreRequestMethods == nil {
return false
}
_, ok := o.ignoreRequestMethods[strings.ToUpper(method)]
return ok
}
gin/middleware/recovery.go 0 → 100644
浏览文件 @ d57375cc
package middleware
import (
"net/http"
ctxUtil "gitlab.wanzhuangkj.com/tush/opkg/xutils/ctxutils"
"github.com/gin-gonic/gin"
"go.uber.org/zap"
)
func CustomRecoveryWithLogger(logger *zap.Logger) gin.HandlerFunc {
return gin.CustomRecovery(func(c *gin.Context, err interface{}) {
logger.Error("Recovered from panic",
zap.String("path", c.Request.URL.Path),
zap.String("method", c.Request.Method),
zap.String("client_ip", c.ClientIP()),
zap.Any("error", err),
zap.Stack("stack"),
ctxUtil.GinTraceIDField(c),
)
c.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{
"code": 500,
"message": "服务器内部错误",
})
})
}
gin/middleware/requstid.go 0 → 100644
浏览文件 @ d57375cc
package middleware
import (
"context"
"net/http"
"time"
"github.com/gin-gonic/gin"
"gitlab.wanzhuangkj.com/tush/opkg/pkg/sf"
ctxutils "gitlab.wanzhuangkj.com/tush/opkg/xutils/ctxutils"
"go.uber.org/zap"
)
// RequestIDOption set the request id options.
type RequestIDOption func(*requestIDOptions)
type requestIDOptions struct {
contextRequestIDKey string
headerXRequestIDKey string
appName string
}
func defaultRequestIDOptions() *requestIDOptions {
return &requestIDOptions{
contextRequestIDKey: ctxutils.ContextTraceIDKey,
headerXRequestIDKey: ctxutils.HeaderXRequestIDKey,
}
}
func (o *requestIDOptions) apply(opts ...RequestIDOption) {
for _, opt := range opts {
opt(o)
}
}
func (o *requestIDOptions) setRequestIDKey() {
if o.contextRequestIDKey != ctxutils.ContextTraceIDKey {
ctxutils.ContextTraceIDKey = o.contextRequestIDKey
}
if o.headerXRequestIDKey != ctxutils.HeaderXRequestIDKey {
ctxutils.HeaderXRequestIDKey = o.headerXRequestIDKey
}
}
// Withctx.ContextRequestIDKey set context request id key, minimum length of 4
func WithContextRequestIDKey(key string) RequestIDOption {
return func(o *requestIDOptions) {
if len(key) < 4 {
return
}
o.contextRequestIDKey = key
}
}
// WithHeaderRequestIDKey set header request id key, minimum length of 4
func WithHeaderRequestIDKey(key string) RequestIDOption {
return func(o *requestIDOptions) {
if len(key) < 4 {
return
}
o.headerXRequestIDKey = key
}
}
// CtxKeyString for context.WithValue key type
type CtxKeyString string
// RequestIDKey request_id
var RequestIDKey = CtxKeyString(ctxutils.ContextTraceIDKey)
// -------------------------------------------------------------------------------------------
// RequestID is an interceptor that injects a 'request id' into the context and request/response header of each request.
func RequestID(opts ...RequestIDOption) gin.HandlerFunc {
// customized request id key
o := defaultRequestIDOptions()
o.apply(opts...)
o.setRequestIDKey()
return func(c *gin.Context) {
requestID := c.Request.Header.Get(ctxutils.HeaderXRequestIDKey)
// Create request id
if requestID == "" {
requestID = sf.GenerateID().String()
c.Request.Header.Set(ctxutils.HeaderXRequestIDKey, requestID)
}
st := time.Now()
// Expose it for use in the application
c.Set(ctxutils.ContextTraceIDKey, requestID)
c.Set(ctxutils.KeyApiStartTime, st)
// Set X-Request-Id header
c.Writer.Header().Set(ctxutils.HeaderXRequestIDKey, requestID)
c.Next()
}
}
// HeaderRequestID get request id from the header
func HeaderRequestID(c *gin.Context) string {
return c.Request.Header.Get(ctxutils.HeaderXRequestIDKey)
}
// HeaderRequestIDField get request id field from header
func HeaderRequestIDField(c *gin.Context) zap.Field {
return zap.String(ctxutils.HeaderXRequestIDKey, HeaderRequestID(c))
}
// -------------------------------------------------------------------------------------------
// RequestHeaderKey request header key
var RequestHeaderKey = "request_header_key"
// AdaptCtx adapt context, if ctx is gin.Context, return gin.Context and context of the transformation
func AdaptCtx(ctx context.Context) (*gin.Context, context.Context) {
c, ok := ctx.(*gin.Context)
if ok {
ctx = ctxutils.WrapCtx(c)
}
return c, ctx
}
// GetFromCtx get value from context
func GetFromCtx(ctx context.Context, key string) interface{} {
return ctx.Value(key)
}
// GetFromHeader get value from header
func GetFromHeader(ctx context.Context, key string) string {
header, ok := ctx.Value(RequestHeaderKey).(http.Header)
if !ok {
return ""
}
return header.Get(key)
}
// GetFromHeaders get values from header
func GetFromHeaders(ctx context.Context, key string) []string {
header, ok := ctx.Value(RequestHeaderKey).(http.Header)
if !ok {
return []string{}
}
return header.Values(key)
}
gin/middleware/tracing.go 0 → 100644
浏览文件 @ d57375cc
package middleware
import (
"fmt"
"github.com/gin-gonic/gin"
otelcontrib "go.opentelemetry.io/contrib"
"go.opentelemetry.io/otel"
"go.opentelemetry.io/otel/attribute"
"go.opentelemetry.io/otel/propagation"
semconv "go.opentelemetry.io/otel/semconv/v1.12.0"
oteltrace "go.opentelemetry.io/otel/trace"
)
const (
tracerKey = "otel-tracer"
tracerName = "otelgin"
)
type traceConfig struct {
TracerProvider oteltrace.TracerProvider
Propagators propagation.TextMapPropagator
}
// TraceOption specifies instrumentation configuration options.
type TraceOption func(*traceConfig)
// WithPropagators specifies propagators to use for extracting
// information from the HTTP requests. If none are specified, global
// ones will be used.
func WithPropagators(propagators propagation.TextMapPropagator) TraceOption {
return func(cfg *traceConfig) {
cfg.Propagators = propagators
}
}
// WithTracerProvider specifies a tracer provider to use for creating a tracer.
// If none is specified, the global provider is used.
func WithTracerProvider(provider oteltrace.TracerProvider) TraceOption {
return func(cfg *traceConfig) {
cfg.TracerProvider = provider
}
}
// Tracing returns interceptor that will trace incoming requests.
// The service parameter should describe the name of the (virtual)
// server handling the request.
func Tracing(serviceName string, opts ...TraceOption) gin.HandlerFunc {
cfg := traceConfig{}
for _, opt := range opts {
opt(&cfg)
}
if cfg.TracerProvider == nil {
cfg.TracerProvider = otel.GetTracerProvider()
}
tracer := cfg.TracerProvider.Tracer(
tracerName,
oteltrace.WithInstrumentationVersion(otelcontrib.Version()),
)
if cfg.Propagators == nil {
cfg.Propagators = otel.GetTextMapPropagator()
}
return func(c *gin.Context) {
c.Set(tracerKey, tracer)
savedCtx := c.Request.Context()
defer func() {
c.Request = c.Request.WithContext(savedCtx)
}()
ctx := cfg.Propagators.Extract(savedCtx, propagation.HeaderCarrier(c.Request.Header))
route := c.FullPath()
tOpts := []oteltrace.SpanStartOption{
oteltrace.WithAttributes(semconv.NetAttributesFromHTTPRequest("tcp", c.Request)...),
oteltrace.WithAttributes(semconv.EndUserAttributesFromHTTPRequest(c.Request)...),
oteltrace.WithAttributes(semconv.HTTPServerAttributesFromHTTPRequest(serviceName, route, c.Request)...),
oteltrace.WithSpanKind(oteltrace.SpanKindServer),
}
spanName := route
if spanName == "" {
spanName = fmt.Sprintf("HTTP %s route not found", c.Request.Method)
}
ctx, span := tracer.Start(ctx, spanName, tOpts...)
defer span.End()
// pass the span through the request context
c.Request = c.Request.WithContext(ctx)
// serve the request to the next interceptor
c.Next()
status := c.Writer.Status()
attrs := semconv.HTTPAttributesFromHTTPStatusCode(status)
spanStatus, spanMessage := semconv.SpanStatusFromHTTPStatusCode(status)
span.SetAttributes(attrs...)
span.SetStatus(spanStatus, spanMessage)
if len(c.Errors) > 0 {
span.SetAttributes(attribute.String("gin.errors", c.Errors.String()))
}
}
}
gin/response/response.go 0 → 100644
浏览文件 @ d57375cc
// Package response provides wrapper gin returns json data in the same format.
package response
import (
"bytes"
"encoding/json"
"net/http"
"strconv"
"time"
"gitlab.wanzhuangkj.com/tush/opkg/pkg/httpcli/entity"
"gitlab.wanzhuangkj.com/tush/opkg/pkg/sgorm/query"
"gitlab.wanzhuangkj.com/tush/opkg/xutils/ctxutils"
"github.com/gin-gonic/gin"
"gitlab.wanzhuangkj.com/tush/opkg/pkg/errcode"
"gitlab.wanzhuangkj.com/tush/opkg/pkg/logger"
"gitlab.wanzhuangkj.com/tush/opkg/pkg/xerrors/xerror"
)
const CustomErrorCode = 0
type BasePageReply struct {
BaseListReply
Total int64 `json:"total" comment:"total" encomment:"total" form:"total" example:"2000"` // 总行数
PageSize int64 `json:"pageSize" comment:"pageSize" encomment:"pageSize" form:"pageSize" example:"1"` // 页大小
PageIndex int64 `json:"pageIndex" comment:"pageIndex" encomment:"pageIndex" form:"pageIndex" example:"1"` // 当前页码
PageTotal int64 `json:"pageTotal" comment:"pageTotal" encomment:"pageTotal" form:"pageTotal" example:"99"` // 总页数
}
type BasePageReply2 struct {
BaseListReply2
Total int64 `json:"total" comment:"total" encomment:"total" form:"total" example:"2000"` // 总行数
PageSize int64 `json:"pageSize" comment:"pageSize" encomment:"pageSize" form:"pageSize" example:"1"` // 页大小
PageIndex int64 `json:"pageIndex" comment:"pageIndex" encomment:"pageIndex" form:"pageIndex" example:"1"` // 当前页码
PageTotal int64 `json:"pageTotal" comment:"pageTotal" encomment:"pageTotal" form:"pageTotal" example:"99"` // 总页数
}
type BaseListReply struct {
List []any `json:"list" comment:"list" encomment:"list" form:"list"` //
}
type BaseListReply2 struct {
List any `json:"list" comment:"list" encomment:"list" form:"list"` //
}
type BaseReply struct {
Result
}
// Result output data format
type Result struct {
Code int `json:"code" form:"code" example:"1"`
Msg string `json:"message" form:"message" example:"OK"`
Data interface{} `json:"data"`
TimeStamp int64 `json:"timestamp" form:"data" example:"1718850053"`
TraceID string `json:"traceID" form:"traceID" example:"7f0000016673920576f0659c620086b0"`
}
func newResp(c *gin.Context, code int, msg string, data interface{}) *Result {
resp := &Result{
Code: code,
Msg: msg,
TimeStamp: time.Now().Unix(),
TraceID: ctxutils.GenerateTraceID(c),
}
// ensure that the data field is not nil on return, note that it is not nil when resp.data=[]interface {}, it is serialized to null
//if data == nil {
// resp.Data = &struct{}{}
//} else {
resp.Data = data
//}
// resp.Data = data
return resp
}
var jsonContentType = []string{"application/json; charset=utf-8"}
func writeContentType(w http.ResponseWriter, value []string) {
header := w.Header()
if val := header["Content-Type"]; len(val) == 0 {
header["Content-Type"] = value
}
}
func writeJSON(c *gin.Context, code int, res interface{}) {
c.Writer.WriteHeader(code)
writeContentType(c.Writer, jsonContentType)
if res != nil {
resBytes, _ := json.Marshal(res)
buf := bytes.NewBuffer(resBytes)
entity.SetCopyRsp(c, buf)
_, _ = c.Writer.Write(resBytes)
}
//err := json.NewEncoder(c.Writer).Encode(res)
//if err != nil {
// fmt.Printf("json encode error, err = %s\n", err.Error())
//}
}
func respJSONWithStatusCode(c *gin.Context, code int, msg string, data ...interface{}) {
var firstData interface{}
if len(data) > 0 {
firstData = data[0]
}
resp := newResp(c, code, msg, firstData)
writeJSON(c, code, resp)
}
// Output return standard HTTP status codes and message, parameter code is HTTP status code
func Output(c *gin.Context, code int, data ...interface{}) {
switch code {
case http.StatusOK:
respJSONWithStatusCode(c, http.StatusOK, errcode.Success.Msg(), data...)
case http.StatusBadRequest:
respJSONWithStatusCode(c, http.StatusBadRequest, errcode.InvalidParams.Msg(), data...)
case http.StatusUnauthorized:
respJSONWithStatusCode(c, http.StatusUnauthorized, errcode.Unauthorized.Msg(), data...)
case http.StatusForbidden:
respJSONWithStatusCode(c, http.StatusForbidden, errcode.Forbidden.Msg(), data...)
case http.StatusNotFound:
respJSONWithStatusCode(c, http.StatusNotFound, errcode.NotFound.Msg(), data...)
case http.StatusRequestTimeout:
respJSONWithStatusCode(c, http.StatusRequestTimeout, errcode.Timeout.Msg(), data...)
case http.StatusConflict:
respJSONWithStatusCode(c, http.StatusConflict, errcode.Conflict.Msg(), data...)
case http.StatusInternalServerError:
respJSONWithStatusCode(c, http.StatusInternalServerError, errcode.InternalServerError.Msg(), data...)
case http.StatusTooManyRequests:
respJSONWithStatusCode(c, http.StatusTooManyRequests, errcode.LimitExceed.Msg(), data...)
case http.StatusServiceUnavailable:
respJSONWithStatusCode(c, http.StatusServiceUnavailable, errcode.ServiceUnavailable.Msg(), data...)
default:
respJSONWithStatusCode(c, code, http.StatusText(code), data...)
}
}
// Out returns the standard HTTP status code and message, parameter err is errcode.Error
func Out(c *gin.Context, err *errcode.Error, data ...interface{}) {
code := err.ToHTTPCode()
switch code {
case http.StatusOK:
respJSONWithStatusCode(c, http.StatusOK, errcode.Success.Msg(), data...)
case http.StatusInternalServerError:
respJSONWithStatusCode(c, http.StatusInternalServerError, err.Msg(), data...)
case http.StatusBadRequest:
respJSONWithStatusCode(c, http.StatusBadRequest, err.Msg(), data...)
case http.StatusUnauthorized:
respJSONWithStatusCode(c, http.StatusUnauthorized, err.Msg(), data...)
case http.StatusForbidden:
respJSONWithStatusCode(c, http.StatusForbidden, err.Msg(), data...)
case http.StatusNotFound:
respJSONWithStatusCode(c, http.StatusNotFound, err.Msg(), data...)
case http.StatusRequestTimeout:
respJSONWithStatusCode(c, http.StatusRequestTimeout, err.Msg(), data...)
case http.StatusConflict:
respJSONWithStatusCode(c, http.StatusConflict, err.Msg(), data...)
case http.StatusTooManyRequests:
respJSONWithStatusCode(c, http.StatusTooManyRequests, err.Msg(), data...)
case http.StatusServiceUnavailable:
respJSONWithStatusCode(c, http.StatusServiceUnavailable, err.Msg(), data...)
default:
respJSONWithStatusCode(c, http.StatusNotExtended, err.Msg(), data...)
}
}
// status code flat 200, custom error codes in data.code
func respJSONWith200(c *gin.Context, code int, msg string, data ...interface{}) {
c.Writer.Header().Set(ctxutils.HeaderXTimestampKey, strconv.FormatInt(time.Now().Unix(), 10))
if len(data) > 0 {
writeJSON(c, http.StatusOK, newResp(c, code, msg, data[0]))
return
}
resp := newResp(c, code, msg, nil)
writeJSON(c, http.StatusOK, resp)
}
// Success return success
func Success(c *gin.Context, data ...interface{}) {
c.Set(ctxutils.KeyRspCode, 1)
respJSONWith200(c, errcode.Success.Code(), errcode.Success.Msg(), data...)
}
// SuccessWithPage return success
//
// func SuccessWithPage[T any](c *gin.Context, list []*T, total int64) {
// currPage := 1
// pageSize := 10
// if req, ok := c.Get("reqbody"); ok {
// if pagination, ok := req.(query.IPagination); ok {
// currPage = pagination.GetPageIndex()
// pageSize = pagination.GetPageSize()
// }
// }
// totalPage := total / int64(pageSize)
// if total%int64(pageSize) > 0 {
// totalPage = totalPage + 1
// }
// c.Set(ctxutils.KeyRspCode, 1)
// if list == nil {
// list = []*T{}
// }
// respJSONWith200(c, errcode.Success.Code(), errcode.Success.Msg(), gin.H{
// "list": list, "totalCount": total, "currPage": currPage, "pageSize": pageSize, "totalPage": totalPage,
// })
// }
//
// SuccessWithPage return success
func SuccessWithPage[T any](c *gin.Context, list []*T, total int64) {
pageIndex := 1
pageSize := 10
if req, ok := c.Get(ctxutils.KeyPagination); ok {
if pagination, ok := req.(query.Pagination); ok {
pageIndex = pagination.PageIndex
if pagination.PageSize > 0 {
pageSize = pagination.PageSize
}
}
}
if pageIndex <= 0 {
pageIndex = 1
}
if pageSize <= 0 {
pageSize = 10
}
pageTotal := total / int64(pageSize)
if total%int64(pageSize) > 0 {
pageTotal = pageTotal + 1
}
c.Set(ctxutils.KeyRspCode, 1)
if list == nil {
list = []*T{}
}
respJSONWith200(c, errcode.Success.Code(), errcode.Success.Msg(), gin.H{
"list": list, "total": total, "pageIndex": pageIndex, "pageTotal": pageTotal, "pageSize": pageSize,
})
}
// SuccessWithList return success
func SuccessWithList[T any](c *gin.Context, list []*T) {
c.Set(ctxutils.KeyRspCode, 1)
if list == nil {
list = []*T{}
}
respJSONWith200(c, errcode.Success.Code(), errcode.Success.Msg(), gin.H{"list": list})
}
// ErrorE return error
func ErrorE(c *gin.Context, err *errcode.Error, data ...interface{}) {
c.Set(ctxutils.KeyRspCode, 0)
respJSONWith200(c, err.Code(), err.Msg(), data...)
}
func Error(c *gin.Context, err error) {
c.Set(ctxutils.KeyRspCode, 0)
msg := ""
if err != nil {
msg = err.Error()
if e, ok := err.(*xerror.Error); ok {
//if !errcode.IsSysDefinedError(xerr.Code()) {
// logger.Error(err.Error(), logger.Err(err), ctxutils.GinTraceIDField(c))
//}
if e.Code() != errcode.InvalidParams.Code() {
logger.Error(err.Error(), logger.Err(err), ctxutils.GinTraceIDField(c))
}
} else {
logger.Error(err.Error(), logger.Err(err), ctxutils.GinTraceIDField(c))
}
}
respJSONWith200(c, CustomErrorCode, msg)
}
go.mod
浏览文件 @ d57375cc
......@@ -5,16 +5,25 @@ go 1.24.3
require (
github.com/aliyun/aliyun-oss-go-sdk v3.0.2+incompatible
github.com/bwmarrin/snowflake v0.3.0
github.com/gin-contrib/cors v1.7.6
github.com/gin-gonic/gin v1.11.0
github.com/go-redsync/redsync/v4 v4.14.0
github.com/gocarina/gocsv v0.0.0-20240520201108-78e41c74b4b1
github.com/golang-jwt/jwt/v5 v5.3.0
github.com/huandu/xstrings v1.5.0
github.com/jinzhu/copier v0.4.0
github.com/mojocn/base64Captcha v1.3.8
github.com/pkg/errors v0.9.1
github.com/prometheus/client_golang v1.19.1
github.com/redis/go-redis/v9 v9.16.0
github.com/shirou/gopsutil/v3 v3.24.5
github.com/spf13/cast v1.10.0
github.com/stretchr/testify v1.11.1
github.com/uptrace/opentelemetry-go-extra/otelgorm v0.3.2
gitlab.wanzhuangkj.com/tush/xpkg v1.2.30
go.opentelemetry.io/contrib v1.38.0
go.opentelemetry.io/otel v1.37.0
go.opentelemetry.io/otel/trace v1.37.0
go.uber.org/zap v1.27.0
golang.org/x/sync v0.17.0
google.golang.org/grpc v1.76.0
......@@ -30,21 +39,23 @@ require (
require (
filippo.io/edwards25519 v1.1.0 // indirect
github.com/beorn7/perks v1.0.1 // indirect
github.com/bytedance/sonic v1.14.0 // indirect
github.com/bytedance/sonic/loader v0.3.0 // indirect
github.com/cespare/xxhash/v2 v2.3.0 // indirect
github.com/cloudwego/base64x v0.1.6 // indirect
github.com/davecgh/go-spew v1.1.1 // indirect
github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
github.com/gabriel-vasile/mimetype v1.4.8 // indirect
github.com/gabriel-vasile/mimetype v1.4.9 // indirect
github.com/gin-contrib/sse v1.1.0 // indirect
github.com/go-logr/logr v1.4.3 // indirect
github.com/go-logr/stdr v1.2.2 // indirect
github.com/go-ole/go-ole v1.2.6 // indirect
github.com/go-playground/locales v0.14.1 // indirect
github.com/go-playground/universal-translator v0.18.1 // indirect
github.com/go-playground/validator/v10 v10.27.0 // indirect
github.com/go-sql-driver/mysql v1.8.1 // indirect
github.com/goccy/go-json v0.10.2 // indirect
github.com/goccy/go-json v0.10.5 // indirect
github.com/goccy/go-yaml v1.18.0 // indirect
github.com/golang/freetype v0.0.0-20170609003504-e2365dfdc4a0 // indirect
github.com/hashicorp/errwrap v1.1.0 // indirect
......@@ -58,21 +69,28 @@ require (
github.com/json-iterator/go v1.1.12 // indirect
github.com/klauspost/cpuid/v2 v2.3.0 // indirect
github.com/leodido/go-urn v1.4.0 // indirect
github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 // indirect
github.com/mattn/go-isatty v0.0.20 // indirect
github.com/mattn/go-sqlite3 v1.14.22 // indirect
github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421 // indirect
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
github.com/modern-go/reflect2 v1.0.2 // indirect
github.com/pelletier/go-toml/v2 v2.2.4 // indirect
github.com/pmezard/go-difflib v1.0.0 // indirect
github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c // indirect
github.com/prometheus/client_model v0.6.0 // indirect
github.com/prometheus/common v0.48.0 // indirect
github.com/prometheus/procfs v0.12.0 // indirect
github.com/quic-go/qpack v0.5.1 // indirect
github.com/quic-go/quic-go v0.54.0 // indirect
github.com/shoenig/go-m1cpu v0.1.6 // indirect
github.com/tklauser/go-sysconf v0.3.12 // indirect
github.com/tklauser/numcpus v0.6.1 // indirect
github.com/twitchyliquid64/golang-asm v0.15.1 // indirect
github.com/ugorji/go/codec v1.3.0 // indirect
github.com/uptrace/opentelemetry-go-extra/otelsql v0.3.2 // indirect
github.com/yusufpapurcu/wmi v1.2.4 // indirect
go.opentelemetry.io/auto/sdk v1.1.0 // indirect
go.opentelemetry.io/otel v1.37.0 // indirect
go.opentelemetry.io/otel/metric v1.37.0 // indirect
go.opentelemetry.io/otel/trace v1.37.0 // indirect
go.uber.org/mock v0.5.0 // indirect
go.uber.org/multierr v1.10.0 // indirect
golang.org/x/arch v0.20.0 // indirect
......
go.sum
浏览文件 @ d57375cc
差异被折叠。
pkg/container/group/example_test.go 0 → 100644
浏览文件 @ d57375cc
package group
import "fmt"
type Counter struct {
Value int
}
func (c *Counter) Incr() {
c.Value++
}
func ExampleGroup_Get() {
group := NewGroup(func() interface{} {
fmt.Println("Only Once")
return &Counter{}
})
// Create a new Counter
group.Get("pass").(*Counter).Incr()
// Get the created Counter again.
group.Get("pass").(*Counter).Incr()
// Output:
// Only Once
}
func ExampleGroup_Reset() {
group := NewGroup(func() interface{} {
return &Counter{}
})
// Reset the new function and clear all created objects.
group.Reset(func() interface{} {
fmt.Println("reset")
return &Counter{}
})
// Create a new Counter
group.Get("pass").(*Counter).Incr()
// Output:reset
}
pkg/container/group/group.go 0 → 100644
浏览文件 @ d57375cc
// Package group provides a sample lazy load container.
// The group only creating a new object not until the object is needed by user.
// And it will cache all the objects to reduce the creation of object.
package group
import "sync"
// Group is a lazy load container.
type Group struct {
new func() interface{}
vals map[string]interface{}
sync.RWMutex
}
// NewGroup news a group container.
func NewGroup(fn func() interface{}) *Group {
if fn == nil {
panic("container.group: can't assign a nil to the new function")
}
return &Group{
new: fn,
vals: make(map[string]interface{}),
}
}
// Get gets the object by the given key.
func (g *Group) Get(key string) interface{} {
g.RLock()
v, ok := g.vals[key]
if ok {
g.RUnlock()
return v
}
g.RUnlock()
// slow path for group don`t have specified key value
g.Lock()
defer g.Unlock()
v, ok = g.vals[key]
if ok {
return v
}
v = g.new()
g.vals[key] = v
return v
}
// Reset resets the new function and deletes all existing objects.
func (g *Group) Reset(fn func() interface{}) {
if fn == nil {
panic("container.group: can't assign a nil to the new function")
}
g.Lock()
g.new = fn
g.Unlock()
g.Clear()
}
// Clear deletes all objects.
func (g *Group) Clear() {
g.Lock()
g.vals = make(map[string]interface{})
g.Unlock()
}
pkg/container/group/group_test.go 0 → 100644
浏览文件 @ d57375cc
package group
import (
"reflect"
"testing"
)
func TestGroupGet(t *testing.T) {
count := 0
g := NewGroup(func() interface{} {
count++
return count
})
v := g.Get("key_0")
if !reflect.DeepEqual(v.(int), 1) {
t.Errorf("expect 1, actual %v", v)
}
v = g.Get("key_1")
if !reflect.DeepEqual(v.(int), 2) {
t.Errorf("expect 2, actual %v", v)
}
v = g.Get("key_0")
if !reflect.DeepEqual(v.(int), 1) {
t.Errorf("expect 1, actual %v", v)
}
if !reflect.DeepEqual(count, 2) {
t.Errorf("expect count 2, actual %v", count)
}
}
func TestGroupReset(t *testing.T) {
g := NewGroup(func() interface{} {
return 1
})
g.Get("key")
call := false
g.Reset(func() interface{} {
call = true
return 1
})
length := 0
for range g.vals {
length++
}
if !reflect.DeepEqual(length, 0) {
t.Errorf("expect length 0, actual %v", length)
}
g.Get("key")
if !reflect.DeepEqual(call, true) {
t.Errorf("expect call true, actual %v", call)
}
}
func TestGroupClear(t *testing.T) {
g := NewGroup(func() interface{} {
return 1
})
g.Get("key")
length := 0
for range g.vals {
length++
}
if !reflect.DeepEqual(length, 1) {
t.Errorf("expect length 1, actual %v", length)
}
g.Clear()
length = 0
for range g.vals {
length++
}
if !reflect.DeepEqual(length, 0) {
t.Errorf("expect length 0, actual %v", length)
}
}
pkg/jwt/jwt.go 0 → 100644
浏览文件 @ d57375cc
// Package jwt is token generation and validation.
package jwt
import (
"time"
"github.com/golang-jwt/jwt/v5"
)
// ErrTokenExpired expired
var ErrTokenExpired = jwt.ErrTokenExpired
var opt *options
// Init initialize jwt
func Init(opts ...Option) {
o := defaultOptions()
o.apply(opts...)
opt = o
}
// Claims standard claims, include uid, name, and RegisteredClaims
type Claims struct {
UID string `json:"uid"`
Name string `json:"name"`
jwt.RegisteredClaims
}
// GenerateToken generate token by uid and name, use universal Claims
func GenerateToken(uid string, name string) (string, time.Time, error) {
expireAt := time.Now()
if opt == nil {
return "", expireAt, errInit
}
expireAt = time.Now().Add(opt.expire)
claims := Claims{
uid,
name,
jwt.RegisteredClaims{
ExpiresAt: jwt.NewNumericDate(expireAt),
IssuedAt: jwt.NewNumericDate(time.Now()),
Issuer: opt.issuer,
},
}
token := jwt.NewWithClaims(opt.signingMethod, claims)
tk, err := token.SignedString(opt.signingKey)
if err != nil {
return tk, expireAt, err
}
return tk, expireAt, nil
}
// ParseToken parse token, return universal Claims
func ParseToken(tokenString string) (*Claims, error) {
if opt == nil {
return nil, errInit
}
token, err := jwt.ParseWithClaims(tokenString, &Claims{}, func(token *jwt.Token) (interface{}, error) {
return opt.signingKey, nil
})
if err != nil {
return nil, err
}
if claims, ok := token.Claims.(*Claims); ok && token.Valid {
return claims, nil
}
return nil, errSignature
}
// RefreshToken refresh token
func RefreshToken(tokenString string) (string, error) {
claims, err := ParseToken(tokenString)
if err != nil {
return "", err
}
claims.RegisteredClaims.ExpiresAt = jwt.NewNumericDate(time.Now().Add(opt.expire))
claims.RegisteredClaims.IssuedAt = jwt.NewNumericDate(time.Now())
token := jwt.NewWithClaims(opt.signingMethod, claims)
return token.SignedString(opt.signingKey)
}
// -------------------------------------------------------------------------------------------
// KV map type
type KV = map[string]interface{}
// CustomClaims custom fields claims
type CustomClaims struct {
Fields KV `json:"fields"`
jwt.RegisteredClaims
}
// Get custom field value by key, if not found, return false
func (c *CustomClaims) Get(key string) (val interface{}, isExist bool) {
if c.Fields == nil {
return nil, false
}
val, isExist = c.Fields[key]
return val, isExist
}
// GetString custom field value by key, if not found, return false
func (c *CustomClaims) GetString(key string) (string, bool) {
val, isExist := c.Get(key)
if isExist {
str, ok := val.(string)
return str, ok
}
return "", false
}
// GetInt custom field value by key, if not found, return false
func (c *CustomClaims) GetInt(key string) (int, bool) {
val, isExist := c.Get(key)
if isExist {
if v, ok := val.(float64); ok {
return int(v), true
}
if v, ok := val.(int); ok {
return v, true
}
}
return 0, false
}
// GetUint64 custom field value by key, if not found, return false
func (c *CustomClaims) GetUint64(key string) (uint64, bool) {
val, isExist := c.Get(key)
if isExist {
if v, ok := val.(float64); ok {
return uint64(v), true
}
if v, ok := val.(uint64); ok {
return v, true
}
}
return 0, false
}
// GenerateCustomToken generate token by custom fields, use CustomClaims
func GenerateCustomToken(kv map[string]interface{}) (string, error) {
if opt == nil {
return "", errInit
}
claims := CustomClaims{
kv,
jwt.RegisteredClaims{
ExpiresAt: jwt.NewNumericDate(time.Now().Add(opt.expire)),
IssuedAt: jwt.NewNumericDate(time.Now()),
Issuer: opt.issuer,
},
}
token := jwt.NewWithClaims(opt.signingMethod, claims)
return token.SignedString(opt.signingKey)
}
// ParseCustomToken parse token, return CustomClaims
func ParseCustomToken(tokenString string) (*CustomClaims, error) {
if opt == nil {
return nil, errInit
}
cc := &CustomClaims{}
token, err := jwt.ParseWithClaims(tokenString, cc, func(token *jwt.Token) (interface{}, error) {
return opt.signingKey, nil
})
if err != nil {
return nil, err
}
if !token.Valid {
return nil, errSignature
}
return cc, nil
}
// RefreshCustomToken refresh custom token
func RefreshCustomToken(tokenString string) (string, error) {
claims, err := ParseCustomToken(tokenString)
if err != nil {
return "", err
}
claims.RegisteredClaims.ExpiresAt = jwt.NewNumericDate(time.Now().Add(opt.expire))
claims.RegisteredClaims.IssuedAt = jwt.NewNumericDate(time.Now())
token := jwt.NewWithClaims(opt.signingMethod, claims)
return token.SignedString(opt.signingKey)
}
pkg/jwt/jwt_test.go 0 → 100644
浏览文件 @ d57375cc
package jwt
import (
"errors"
"fmt"
"testing"
"time"
"github.com/stretchr/testify/assert"
)
func TestGenerateToken(t *testing.T) {
opt = nil
token, _, err := GenerateToken("123", "")
assert.Error(t, err)
Init()
token, _, err = GenerateToken("123", "")
assert.NoError(t, err)
v, err := ParseToken(token)
if err != nil {
t.Fatal(err)
}
t.Log(v)
time.Sleep(time.Second)
newToken, err := RefreshToken(token)
if err != nil {
t.Fatal(err)
}
t.Log(token, newToken)
}
func TestParseToken(t *testing.T) {
opt = nil
v, err := ParseToken("token")
assert.Error(t, err)
uid := "123"
name := "admin"
Init(
WithSigningKey("123456"),
WithExpire(time.Second),
WithSigningMethod(HS512),
)
// success
token, _, err := GenerateToken(uid, name)
if err != nil {
t.Fatal(err)
}
fmt.Println(token)
v, err = ParseToken(token)
if err != nil {
t.Fatal(err)
}
fmt.Println(v)
// invalid token format
token2 := "xxx.xxx.xxx"
v, err = ParseToken(token2)
assert.Error(t, err)
// signature failure
token3 := token + "xxx"
v, err = ParseToken(token3)
assert.Error(t, err)
// token has expired
token, _, err = GenerateToken(uid, name)
if err != nil {
t.Fatal(err)
}
time.Sleep(time.Second * 2)
v, err = ParseToken(token)
assert.True(t, errors.Is(err, ErrTokenExpired))
}
func TestParse(t *testing.T) {
Init(
WithSigningKey("6m2KywHIPEaWWIxffGpp8Dfl3z5SjE5WeeDStc1T64klnpoAqHrHN01vneteCAGE"),
WithExpire(time.Second),
WithSigningMethod(HS256),
)
token := "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJmaWVsZHMiOnsiY29tcGFueUlEIjoiNTIiLCJ1VHlwZSI6MiwidWlkIjoiMTMiLCJ1bmFtZSI6IuW8oOS4iSJ9LCJleHAiOjE3NTc2ODI2NzksImlhdCI6MTc1NzY2MTA3OX0.WqtvxF4XW52bC4mglFOg61i4D9mi7Bkb1ZZ00F4PB6k"
claims, err := ParseCustomToken(token)
if err != nil {
t.Error(err)
}
fmt.Println(claims)
}
func TestGenerateCustomToken(t *testing.T) {
var (
id uint64 = 20
name string = "admin"
age int = 10
fields = KV{"id": id, "name": name, "age": age}
)
Init()
token, err := GenerateCustomToken(fields)
assert.NoError(t, err)
claims, err := ParseCustomToken(token)
assert.NoError(t, err)
idValue, _ := claims.GetUint64("id")
assert.Equal(t, idValue, fields["id"])
nameValue, _ := claims.GetString("name")
assert.Equal(t, nameValue, fields["name"])
ageValue, _ := claims.GetInt("age")
assert.Equal(t, ageValue, fields["age"])
_, ok := claims.Get("foo")
assert.Equal(t, ok, false)
claims.Fields = nil
foo, _ := claims.Get("foo")
assert.Nil(t, foo)
time.Sleep(time.Second)
newToken, err := RefreshCustomToken(token)
if err != nil {
t.Fatal(err)
}
t.Log(token, newToken)
}
func TestParseCustomToken(t *testing.T) {
fields := KV{"id": 123, "foo": "bar"}
opt = nil
v, err := ParseCustomToken("token")
assert.Error(t, err)
Init(
WithSigningKey("123456"),
WithExpire(time.Second),
WithSigningMethod(HS512),
)
// success
token, err := GenerateCustomToken(fields)
if err != nil {
t.Fatal(err)
}
fmt.Println(token)
v, err = ParseCustomToken(token)
if err != nil {
t.Fatal(err)
}
fmt.Println(v)
// invalid token format
token2 := "xxx.xxx.xxx"
v, err = ParseCustomToken(token2)
assert.Error(t, err)
// signature failure
token3 := token + "xxx"
v, err = ParseCustomToken(token3)
assert.Error(t, err)
// token has expired
token, err = GenerateCustomToken(fields)
if err != nil {
t.Fatal(err)
}
time.Sleep(time.Second * 2)
v, err = ParseCustomToken(token)
assert.True(t, errors.Is(err, ErrTokenExpired))
}
pkg/jwt/option.go 0 → 100644
浏览文件 @ d57375cc
package jwt
import (
"errors"
"time"
"github.com/golang-jwt/jwt/v5"
)
var (
// HS256 Method
HS256 = jwt.SigningMethodHS256
// HS384 Method
HS384 = jwt.SigningMethodHS384
// HS512 Method
HS512 = jwt.SigningMethodHS512
)
var (
defaultSigningKey = []byte("zaq12wsxmko0") // default key
defaultSigningMethod = HS256 // default HS256
defaultExpire = 24 * time.Hour // default expiration
defaultIssuer = ""
)
type options struct {
signingKey []byte
expire time.Duration
issuer string
signingMethod *jwt.SigningMethodHMAC
}
func defaultOptions() *options {
return &options{
signingKey: defaultSigningKey,
signingMethod: defaultSigningMethod,
expire: defaultExpire,
issuer: defaultIssuer,
}
}
// Option set the jwt options.
type Option func(*options)
func (o *options) apply(opts ...Option) {
for _, opt := range opts {
opt(o)
}
}
// WithSigningKey set signing key value
func WithSigningKey(key string) Option {
return func(o *options) {
o.signingKey = []byte(key)
}
}
// WithSigningMethod set signing method value
func WithSigningMethod(sm *jwt.SigningMethodHMAC) Option {
return func(o *options) {
o.signingMethod = sm
}
}
// WithExpire set expire value
func WithExpire(d time.Duration) Option {
return func(o *options) {
o.expire = d
}
}
// WithIssuer set issuer value
func WithIssuer(issuer string) Option {
return func(o *options) {
o.issuer = issuer
}
}
var (
errSignature = errors.New("signature failure")
errInit = errors.New("not yet initialized jwt, usage 'jwt.Init()'")
)
pkg/jwt/option_test.go 0 → 100644
浏览文件 @ d57375cc
package jwt
import (
"testing"
"time"
"github.com/golang-jwt/jwt/v5"
"github.com/stretchr/testify/assert"
)
func TestInit(t *testing.T) {
Init(WithSigningKey("foo"))
}
func TestWithExpire(t *testing.T) {
testData := time.Second * 3
opt := WithExpire(testData)
o := new(options)
o.apply(opt)
assert.Equal(t, testData, o.expire)
}
func TestWithIssuer(t *testing.T) {
testData := "issuer"
opt := WithIssuer(testData)
o := new(options)
o.apply(opt)
assert.Equal(t, testData, o.issuer)
}
func TestWithSigningKey(t *testing.T) {
testData := "key"
opt := WithSigningKey(testData)
o := new(options)
o.apply(opt)
assert.Equal(t, testData, string(o.signingKey))
}
func TestWithSigningMethod(t *testing.T) {
testData := jwt.SigningMethodHS384
opt := WithSigningMethod(testData)
o := new(options)
o.apply(opt)
assert.Equal(t, testData, o.signingMethod)
}
func Test_defaultOptions(t *testing.T) {
o := defaultOptions()
assert.NotNil(t, o)
}
func Test_options_apply(t *testing.T) {
testData := "key"
opt := WithSigningKey(testData)
o := new(options)
o.apply(opt)
assert.Equal(t, testData, string(o.signingKey))
}
pkg/logger/logger.go
浏览文件 @ d57375cc
......@@ -58,7 +58,7 @@ func Init(cfg Config) (err error) {
if err != nil {
return err
}
Info("[olog] initialized")
Info("[logger] initialized")
return nil
}
......@@ -95,10 +95,10 @@ func InitWithOptions(opts ...Option) (*zap.Logger, error) {
if err != nil {
panic(err)
}
str = fmt.Sprintf("[olog] config is output to 'terminal', format=%s, level=%s", encoding, levelName)
str = fmt.Sprintf("[logger] config is output to 'terminal', format=%s, level=%s", encoding, levelName)
} else {
zapLog = log2File(encoding, levelName, o.fileConfig)
str = fmt.Sprintf("[olog] config is output to 'file', format=%s, level=%s, file=%s", encoding, levelName, o.fileConfig.filename)
str = fmt.Sprintf("[logger] config is output to 'file', format=%s, level=%s, file=%s", encoding, levelName, o.fileConfig.filename)
}
if len(o.hooks) > 0 {
......
pkg/shield/circuitbreaker/README.md 0 → 100644
浏览文件 @ d57375cc
## circuitbreaker
Circuit Breaker for web middleware and rpc interceptor.
<br>
### Example of use
**gin circuit breaker middleware**
```go
import "gitlab.wanzhuangkj.com/tush/xpkg/pkg/shield/circuitbreaker"
// CircuitBreaker a circuit breaker middleware
func CircuitBreaker(opts ...CircuitBreakerOption) gin.HandlerFunc {
o := defaultCircuitBreakerOptions()
o.apply(opts...)
return func(c *gin.Context) {
breaker := o.group.Get(c.FullPath()).(circuitbreaker.CircuitBreaker)
if err := breaker.Allow(); err != nil {
// NOTE: when client reject request locally, keep adding counter let the drop ratio higher.
breaker.MarkFailed()
response.Output(c, http.StatusServiceUnavailable, err.Error())
c.Abort()
return
}
c.Next()
code := c.Writer.Status()
// NOTE: need to check internal and service unavailable error, e.g. http.StatusInternalServerError
_, isHit := o.validCodes[code]
if isHit {
breaker.MarkFailed()
} else {
breaker.MarkSuccess()
}
}
}
```
<br>
**rpc server circuit breaker interceptor**
```go
import "gitlab.wanzhuangkj.com/tush/xpkg/pkg/shield/circuitbreaker"
// UnaryServerCircuitBreaker server-side unary circuit breaker interceptor
func UnaryServerCircuitBreaker(opts ...CircuitBreakerOption) grpc.UnaryServerInterceptor {
o := defaultCircuitBreakerOptions()
o.apply(opts...)
return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
breaker := o.group.Get(info.FullMethod).(circuitbreaker.CircuitBreaker)
if err := breaker.Allow(); err != nil {
// NOTE: when client reject request locally, keep adding let the drop ratio higher.
breaker.MarkFailed()
return nil, errcode.StatusServiceUnavailable.ToRPCErr(err.Error())
}
reply, err := handler(ctx, req)
if err != nil {
// NOTE: need to check internal and service unavailable error
s, ok := status.FromError(err)
_, isHit := o.validCodes[s.Code()]
if ok && isHit {
breaker.MarkFailed()
} else {
breaker.MarkSuccess()
}
}
return reply, err
}
}
```
pkg/shield/circuitbreaker/circuitbreaker.go 0 → 100644
浏览文件 @ d57375cc
// Package circuitbreaker is an adaptive circuit breaker library, support for use in gin middleware and grpc interceptors.
package circuitbreaker
import (
"errors"
)
// ErrNotAllowed error not allowed.
var ErrNotAllowed = errors.New("circuitbreaker: not allowed for circuit open")
// CircuitBreaker is a circuit breaker.
type CircuitBreaker interface {
Allow() error
MarkSuccess()
MarkFailed()
}
pkg/shield/circuitbreaker/circuitbreaker_test.go 0 → 100644
浏览文件 @ d57375cc
package circuitbreaker_test
import (
"testing"
"gitlab.wanzhuangkj.com/tush/xpkg/pkg/shield/circuitbreaker"
)
// This is an example of using a circuit breaker Do() when return nil.
func TestCircuitBreaker(t *testing.T) {
b := circuitbreaker.NewBreaker()
for i := 0; i < 1000; i++ {
b.MarkSuccess()
}
for i := 0; i < 100; i++ {
b.MarkFailed()
}
err := b.Allow()
t.Log(err)
// Output: err=<nil>
}
pkg/shield/circuitbreaker/sre.go 0 → 100644
浏览文件 @ d57375cc
package circuitbreaker
import (
"math"
"math/rand"
"sync"
"sync/atomic"
"time"
"gitlab.wanzhuangkj.com/tush/opkg/pkg/shield/window"
)
// Option is sre breaker option function.
type Option func(*options)
const (
// StateOpen when circuit breaker open, request not allowed, after sleep
// some duration, allow one single request for testing the health, if ok
// then state reset to closed, if not continue the step.
StateOpen int32 = iota
// StateClosed when circuit breaker closed, request allowed, the breaker
// calc the succeed ratio, if request num greater request setting and
// ratio lower than the setting ratio, then reset state to open.
StateClosed
)
var (
_ CircuitBreaker = &Breaker{}
)
// options is a breaker options.
type options struct {
success float64
request int64
bucket int
window time.Duration
}
// WithSuccess with the K = 1 / Success value of sre breaker, default success is 0.5
// Reducing the K will make adaptive throttling behave more aggressively,
// Increasing the K will make adaptive throttling behave less aggressively.
func WithSuccess(s float64) Option {
return func(c *options) {
c.success = s
}
}
// WithRequest with the minimum number of requests allowed.
func WithRequest(r int64) Option {
return func(c *options) {
c.request = r
}
}
// WithWindow with the duration size of the statistical window.
func WithWindow(d time.Duration) Option {
return func(c *options) {
c.window = d
}
}
// WithBucket set the bucket number in a window duration.
func WithBucket(b int) Option {
return func(c *options) {
c.bucket = b
}
}
// Breaker is a sre CircuitBreaker pattern.
type Breaker struct {
stat window.RollingCounter
r *rand.Rand
// rand.New(...) returns a non thread safe object
randLock sync.Mutex
// Reducing the k will make adaptive throttling behave more aggressively,
// Increasing the k will make adaptive throttling behave less aggressively.
k float64
request int64
state int32
}
// NewBreaker return a sreBresker with options
func NewBreaker(opts ...Option) CircuitBreaker {
opt := options{
success: 0.6,
request: 100,
bucket: 10,
window: 3 * time.Second,
}
for _, o := range opts {
o(&opt)
}
counterOpts := window.RollingCounterOpts{
Size: opt.bucket,
BucketDuration: time.Duration(int64(opt.window) / int64(opt.bucket)),
}
stat := window.NewRollingCounter(counterOpts)
return &Breaker{
stat: stat,
r: rand.New(rand.NewSource(time.Now().UnixNano())),
request: opt.request,
k: 1 / opt.success,
state: StateClosed,
}
}
func (b *Breaker) summary() (success int64, total int64) {
b.stat.Reduce(func(iterator window.Iterator) float64 {
for iterator.Next() {
bucket := iterator.Bucket()
total += bucket.Count
for _, p := range bucket.Points {
success += int64(p)
}
}
return 0
})
return //nolint
}
// Allow request if error returns nil.
func (b *Breaker) Allow() error {
// The number of requests accepted by the backend
accepts, total := b.summary()
// The number of requests attempted by the application layer(at the client, on top of the adaptive throttling system)
requests := b.k * float64(accepts)
// check overflow requests = K * accepts
if total < b.request || float64(total) < requests {
atomic.CompareAndSwapInt32(&b.state, StateOpen, StateClosed)
return nil
}
atomic.CompareAndSwapInt32(&b.state, StateClosed, StateOpen)
dr := math.Max(0, (float64(total)-requests)/float64(total+1))
drop := b.trueOnProba(dr)
if drop {
return ErrNotAllowed
}
return nil
}
// MarkSuccess mark request is success.
func (b *Breaker) MarkSuccess() {
b.stat.Add(1)
}
// MarkFailed mark request is failed.
func (b *Breaker) MarkFailed() {
// NOTE: when client reject request locally, keep adding counter let the drop ratio higher.
b.stat.Add(0)
}
func (b *Breaker) trueOnProba(proba float64) (truth bool) {
b.randLock.Lock()
truth = b.r.Float64() < proba
b.randLock.Unlock()
return truth
}
pkg/shield/cpu/cgroup.go 0 → 100644
浏览文件 @ d57375cc
package cpu
import (
"bufio"
"fmt"
"io"
"os"
"path"
"strconv"
"strings"
)
const cgroupRootDir = "/sys/fs/cgroup"
// cgroup Linux cgroup
type cgroup struct {
cgroupSet map[string]string
}
// CPUCFSQuotaUs cpu.cfs_quota_us
func (c *cgroup) CPUCFSQuotaUs() (int64, error) {
data, err := readFile(path.Join(c.cgroupSet["cpu"], "cpu.cfs_quota_us"))
if err != nil {
return 0, err
}
return strconv.ParseInt(data, 10, 64)
}
// CPUCFSPeriodUs cpu.cfs_period_us
func (c *cgroup) CPUCFSPeriodUs() (uint64, error) {
data, err := readFile(path.Join(c.cgroupSet["cpu"], "cpu.cfs_period_us"))
if err != nil {
return 0, err
}
return parseUint(data)
}
// CPUAcctUsage cpuacct.usage
func (c *cgroup) CPUAcctUsage() (uint64, error) {
data, err := readFile(path.Join(c.cgroupSet["cpuacct"], "cpuacct.usage"))
if err != nil {
return 0, err
}
return parseUint(data)
}
// CPUAcctUsagePerCPU cpuacct.usage_percpu
func (c *cgroup) CPUAcctUsagePerCPU() ([]uint64, error) {
data, err := readFile(path.Join(c.cgroupSet["cpuacct"], "cpuacct.usage_percpu"))
if err != nil {
return nil, err
}
var usage []uint64
for _, v := range strings.Fields(data) {
var u uint64
if u, err = parseUint(v); err != nil {
return nil, err
}
// fix possible_cpu:https://www.ibm.com/support/knowledgecenter/en/linuxonibm/com.ibm.linux.z.lgdd/lgdd_r_posscpusparm.html
if u != 0 {
usage = append(usage, u)
}
}
return usage, nil
}
// CPUSetCPUs cpuset.cpus
func (c *cgroup) CPUSetCPUs() ([]uint64, error) {
data, err := readFile(path.Join(c.cgroupSet["cpuset"], "cpuset.cpus"))
if err != nil {
return nil, err
}
cpus, err := ParseUintList(data)
if err != nil {
return nil, err
}
sets := make([]uint64, 0)
for k := range cpus {
sets = append(sets, uint64(k))
}
return sets, nil
}
// CurrentcGroup get current process cgroup
func currentcGroup() (*cgroup, error) {
pid := os.Getpid()
cgroupFile := fmt.Sprintf("/proc/%d/cgroup", pid)
cgroupSet := make(map[string]string)
fp, err := os.Open(cgroupFile)
if err != nil {
return nil, err
}
defer fp.Close() //nolint
buf := bufio.NewReader(fp)
for {
line, err := buf.ReadString('\n')
if err != nil {
if err == io.EOF {
break
}
return nil, err
}
col := strings.Split(strings.TrimSpace(line), ":")
if len(col) != 3 {
return nil, fmt.Errorf("invalid cgroup format %s", line)
}
dir := col[2]
// When dir is not equal to /, it must be in docker
if dir != "/" {
cgroupSet[col[1]] = path.Join(cgroupRootDir, col[1])
if strings.Contains(col[1], ",") {
for _, k := range strings.Split(col[1], ",") {
cgroupSet[k] = path.Join(cgroupRootDir, k)
}
}
} else {
cgroupSet[col[1]] = path.Join(cgroupRootDir, col[1], col[2])
if strings.Contains(col[1], ",") {
for _, k := range strings.Split(col[1], ",") {
cgroupSet[k] = path.Join(cgroupRootDir, k, col[2])
}
}
}
}
return &cgroup{cgroupSet: cgroupSet}, nil
}
pkg/shield/cpu/cgroup_cpu.go 0 → 100644
浏览文件 @ d57375cc
package cpu
import (
"bufio"
"errors"
"os"
"strconv"
"strings"
pscpu "github.com/shirou/gopsutil/v3/cpu"
)
type cgroupCPU struct {
frequency uint64
quota float64
cores uint64
preSystem uint64
preTotal uint64
}
func newCgroupCPU() (cpu *cgroupCPU, err error) {
cores, err := pscpu.Counts(true)
if err != nil || cores == 0 {
var cpus []uint64
cpus, err = perCPUUsage()
if err != nil {
return nil, err
}
cores = len(cpus)
}
sets, err := cpuSets()
if err != nil {
return nil, err
}
quota := float64(len(sets))
cq, err := cpuQuota()
if err == nil && cq != -1 {
var period uint64
if period, err = cpuPeriod(); err != nil {
return nil, err
}
limit := float64(cq) / float64(period)
if limit < quota {
quota = limit
}
}
maxFreq := cpuMaxFreq()
preSystem, err := systemCPUUsage()
if err != nil {
return nil, err
}
preTotal, err := totalCPUUsage()
if err != nil {
return nil, err
}
cpu = &cgroupCPU{
frequency: maxFreq,
quota: quota,
cores: uint64(cores),
preSystem: preSystem,
preTotal: preTotal,
}
return cpu, nil
}
func (cpu *cgroupCPU) Usage() (u uint64, err error) {
var (
total uint64
system uint64
)
total, err = totalCPUUsage()
if err != nil {
return 0, err
}
system, err = systemCPUUsage()
if err != nil {
return 0, err
}
if system != cpu.preSystem {
u = uint64(float64((total-cpu.preTotal)*cpu.cores*1e3) / (float64(system-cpu.preSystem) * cpu.quota))
}
cpu.preSystem = system
cpu.preTotal = total
return u, err
}
func (cpu *cgroupCPU) Info() Info {
return Info{
Frequency: cpu.frequency,
Quota: cpu.quota,
}
}
const nanoSecondsPerSecond = 1e9
// ErrNoCFSLimit is no quota limit
var ErrNoCFSLimit = errors.New("no quota limit")
var clockTicksPerSecond = uint64(getClockTicks())
// systemCPUUsage returns the host system's cpu usage in
// nanoseconds. An error is returned if the format of the underlying
// file does not match.
//
// Uses /proc/stat defined by POSIX. Looks for the cpu
// statistics line and then sums up the first seven fields
// provided. See man 5 proc for details on specific field
// information.
func systemCPUUsage() (usage uint64, err error) {
var (
line string
f *os.File
)
if f, err = os.Open("/proc/stat"); err != nil {
return usage, err
}
bufReader := bufio.NewReaderSize(nil, 128)
defer func() {
bufReader.Reset(nil)
_ = f.Close()
}()
bufReader.Reset(f)
for err == nil {
if line, err = bufReader.ReadString('\n'); err != nil {
return usage, err
}
parts := strings.Fields(line)
if parts[0] == "cpu" {
if len(parts) < 8 {
err = errors.New("bad format of cpu stats")
return usage, err
}
var totalClockTicks uint64
for _, i := range parts[1:8] {
var v uint64
if v, err = strconv.ParseUint(i, 10, 64); err != nil {
return usage, err
}
totalClockTicks += v
}
usage = (totalClockTicks * nanoSecondsPerSecond) / clockTicksPerSecond
return usage, nil
}
}
err = errors.New("bad stats format")
return usage, err
}
func totalCPUUsage() (usage uint64, err error) {
var cg *cgroup
if cg, err = currentcGroup(); err != nil {
return 0, err
}
return cg.CPUAcctUsage()
}
func perCPUUsage() (usage []uint64, err error) {
var cg *cgroup
if cg, err = currentcGroup(); err != nil {
return nil, err
}
return cg.CPUAcctUsagePerCPU()
}
func cpuSets() (sets []uint64, err error) {
var cg *cgroup
if cg, err = currentcGroup(); err != nil {
return nil, err
}
return cg.CPUSetCPUs()
}
func cpuQuota() (quota int64, err error) {
var cg *cgroup
if cg, err = currentcGroup(); err != nil {
return 0, err
}
return cg.CPUCFSQuotaUs()
}
func cpuPeriod() (peroid uint64, err error) {
var cg *cgroup
if cg, err = currentcGroup(); err != nil {
return 0, err
}
return cg.CPUCFSPeriodUs()
}
func cpuFreq() uint64 {
lines, err := readLines("/proc/cpuinfo")
if err != nil {
return 0
}
for _, line := range lines {
fields := strings.Split(line, ":")
if len(fields) < 2 {
continue
}
key := strings.TrimSpace(fields[0])
value := strings.TrimSpace(fields[1])
if key == "cpu MHz" || key == "clock" {
// treat this as the fallback value, thus we ignore error
if t, err := strconv.ParseFloat(strings.Replace(value, "MHz", "", 1), 64); err == nil {
return uint64(t * 1000.0 * 1000.0)
}
}
}
return 0
}
func cpuMaxFreq() uint64 {
feq := cpuFreq()
data, err := readFile("/sys/devices/system/cpu/cpu0/cpufreq/cpuinfo_max_freq")
if err != nil {
return feq
}
// override the max freq from /proc/cpuinfo
cfeq, err := parseUint(data)
if err == nil {
feq = cfeq
}
return feq
}
// GetClockTicks get the OS's ticks per second
func getClockTicks() int {
// TODO figure out a better alternative for platforms where we're missing cgo
//
// TODO Windows. This could be implemented using Win32 QueryPerformanceFrequency().
// https://msdn.microsoft.com/en-us/library/windows/desktop/ms644905(v=vs.85).aspx
//
// An example of its usage can be found here.
// https://msdn.microsoft.com/en-us/library/windows/desktop/dn553408(v=vs.85).aspx
return 100
}
pkg/shield/cpu/cgroup_test.go 0 → 100644
浏览文件 @ d57375cc
package cpu
import "testing"
func TestCgroup(t *testing.T) {
cgroup := &cgroup{cgroupSet: map[string]string{}}
us, err := cgroup.CPUCFSQuotaUs()
t.Log(us, err)
pus, err := cgroup.CPUCFSPeriodUs()
t.Log(pus, err)
usage, err := cgroup.CPUAcctUsage()
t.Log(usage, err)
pUsage, err := cgroup.CPUAcctUsagePerCPU()
t.Log(pUsage, err)
cpus, err := cgroup.CPUSetCPUs()
t.Log(cpus, err)
cgr, err := currentcGroup()
t.Log(cgr, err)
}
func TestCgroupCPUObj(t *testing.T) {
cgCPU := &cgroupCPU{
frequency: 3000,
quota: 12,
cores: 12,
preSystem: 10,
preTotal: 10,
}
usage, err := cgCPU.Usage()
t.Log(usage, err)
info := cgCPU.Info()
t.Log(info)
}
func TestCgroupCPU(t *testing.T) {
usage, err := systemCPUUsage()
t.Log(usage, err)
usage, err = totalCPUUsage()
t.Log(usage, err)
usages, err := perCPUUsage()
t.Log(usages, err)
sets, err := cpuSets()
t.Log(sets, err)
quota, err := cpuQuota()
t.Log(quota, err)
period, err := cpuPeriod()
t.Log(period, err)
freq := cpuFreq()
t.Log(freq, err)
maxFreq := cpuMaxFreq()
t.Log(maxFreq, err)
}
pkg/shield/cpu/psutil_cpu.go 0 → 100644
浏览文件 @ d57375cc
// Package cpu is a library that calculates cpu and memory usage.
package cpu
import (
"time"
"github.com/shirou/gopsutil/v3/cpu"
)
type psutilCPU struct {
interval time.Duration
}
func newPsutilCPU(interval time.Duration) (*psutilCPU, error) {
psCPU := &psutilCPU{interval: interval}
_, err := psCPU.Usage()
if err != nil {
return nil, err
}
return psCPU, nil
}
func (ps *psutilCPU) Usage() (uint64, error) {
var u uint64
percents, err := cpu.Percent(ps.interval, false)
if err == nil {
u = uint64(percents[0] * 10)
}
return u, err
}
func (ps *psutilCPU) Info() Info {
stats, err := cpu.Info()
if err != nil {
return Info{}
}
cores, err := cpu.Counts(true)
if err != nil {
return Info{}
}
return Info{
Frequency: uint64(stats[0].Mhz),
Quota: float64(cores),
}
}
pkg/shield/cpu/stat.go 0 → 100644
浏览文件 @ d57375cc
package cpu
import (
"fmt"
"sync/atomic"
"time"
)
const (
interval time.Duration = time.Millisecond * 500
)
var (
stats CPU
usage uint64
)
// CPU is cpu stat usage.
type CPU interface {
Usage() (u uint64, e error)
Info() Info
}
func init() {
var err error
stats, err = newCgroupCPU()
if err != nil {
errStr := err.Error()
stats, err = newPsutilCPU(interval)
if err != nil {
errStr += " | " + err.Error()
fmt.Printf(`
[ERROR] cgroup cpu init and psutil cpu init are all failed, %s.
the dependency library https://github.com/shirou/gopsutil does not support getting this CPU information.
[NOTE] After using the project code generated by xmall, please set "enableStat", "enableLimit" and "enableCircuitBreaker"
in configs/xxx.yml to false, which will not affect the normal use of the application.
`, errStr)
return
}
}
go func() {
ticker := time.NewTicker(interval)
defer ticker.Stop()
for {
<-ticker.C
u, err := stats.Usage()
if err == nil && u != 0 {
atomic.StoreUint64(&usage, u)
}
}
}()
}
// Stat cpu stat.
type Stat struct {
Usage uint64 // cpu use ratio.
}
// Info cpu info.
type Info struct {
Frequency uint64
Quota float64
}
// ReadStat read cpu stat.
func ReadStat(stat *Stat) {
stat.Usage = atomic.LoadUint64(&usage)
}
// GetInfo get cpu info.
func GetInfo() Info {
return stats.Info()
}
pkg/shield/cpu/stat_test.go 0 → 100644
浏览文件 @ d57375cc
package cpu
import (
"testing"
"time"
)
func TestStat(t *testing.T) {
time.Sleep(time.Second * 2)
var s Stat
var i Info
ReadStat(&s)
i = GetInfo()
t.Log(s.Usage, i.Frequency, i.Quota)
}
pkg/shield/cpu/utils.go 0 → 100644
浏览文件 @ d57375cc
package cpu
import (
"bufio"
"fmt"
"os"
"strconv"
"strings"
)
func readFile(path string) (string, error) {
contents, err := os.ReadFile(path)
if err != nil {
return "", err
}
return strings.TrimSpace(string(contents)), nil
}
func parseUint(s string) (uint64, error) {
v, err := strconv.ParseUint(s, 10, 64)
if err != nil {
intValue, intErr := strconv.ParseInt(s, 10, 64)
// 1. Handle negative values greater than MinInt64 (and)
// 2. Handle negative values lesser than MinInt64
if intErr == nil && intValue < 0 {
return 0, nil
} else if intErr != nil &&
intErr.(*strconv.NumError).Err == strconv.ErrRange &&
intValue < 0 {
return 0, nil
}
return 0, err
}
return v, nil
}
// ParseUintList parses and validates the specified string as the value
// found in some cgroup file (e.g. cpuset.cpus, cpuset.mems), which could be
// one of the formats below. Note that duplicates are actually allowed in the
// input string. It returns a map[int]bool with available elements from val
// set to true.
// Supported formats:
// 7
// 1-6
// 0,3-4,7,8-10
// 0-0,0,1-7
// 03,1-3 <- this is gonna get parsed as [1,2,3]
// 3,2,1
// 0-2,3,1
func ParseUintList(val string) (map[int]bool, error) {
if val == "" {
return map[int]bool{}, nil
}
availableInts := make(map[int]bool)
split := strings.Split(val, ",")
errInvalidFormat := fmt.Errorf("os/stat: invalid format: %s", val)
for _, r := range split {
if !strings.Contains(r, "-") {
v, err := strconv.Atoi(r)
if err != nil {
return nil, errInvalidFormat
}
availableInts[v] = true
} else {
split := strings.SplitN(r, "-", 2)
min, err := strconv.Atoi(split[0])
if err != nil {
return nil, errInvalidFormat
}
max, err := strconv.Atoi(split[1])
if err != nil {
return nil, errInvalidFormat
}
if max < min {
return nil, errInvalidFormat
}
for i := min; i <= max; i++ {
availableInts[i] = true
}
}
}
return availableInts, nil
}
// ReadLines reads contents from a file and splits them by new lines.
// A convenience wrapper to ReadLinesOffsetN(filename, 0, -1).
func readLines(filename string) ([]string, error) {
return readLinesOffsetN(filename, 0, -1)
}
// ReadLinesOffsetN reads contents from file and splits them by new line.
// The offset tells at which line number to start.
// The count determines the number of lines to read (starting from offset):
//
// n >= 0: at most n lines
// n < 0: whole file
func readLinesOffsetN(filename string, offset uint, n int) ([]string, error) {
f, err := os.Open(filename)
if err != nil {
return []string{""}, err
}
defer f.Close() //nolint
var ret []string
r := bufio.NewReader(f)
for i := 0; i < n+int(offset) || n < 0; i++ {
line, err := r.ReadString('\n')
if err != nil {
break
}
if i < int(offset) {
continue
}
ret = append(ret, strings.Trim(line, "\n"))
}
return ret, nil
}
pkg/shield/cpu/utils_test.go 0 → 100644
浏览文件 @ d57375cc
package cpu
import (
"testing"
"github.com/stretchr/testify/assert"
)
func Testutils(t *testing.T) {
_, err := readFile("notfound")
assert.Error(t, err)
_, err = readFile("utils.go")
assert.NoError(t, err)
_, err = parseUint("-100")
assert.Nil(t, err)
_, err = parseUint("abc")
assert.Error(t, err)
_, err = parseUint("100")
_, err = ParseUintList("")
assert.Nil(t, err)
_, err = readLines("notfound")
assert.Error(t, err)
_, err = readLinesOffsetN("test.data", 0, -1)
assert.NoError(t, err)
}
func TestParseUintList(t *testing.T) {
testData := []string{
"",
"7",
"1-6",
"0,3-4,7,8-10",
"0-0,0,1-7",
"03,1-3",
"3,2,1",
"0-2,3,1",
}
for _, val := range testData {
_, err := ParseUintList(val)
assert.NoError(t, err)
}
}
pkg/shield/ratelimit/bbr.go 0 → 100644
浏览文件 @ d57375cc
// Package ratelimit is an adaptive rate limit library, support for use in gin middleware and grpc interceptors.
package ratelimit
import (
"math"
"runtime"
"sync/atomic"
"time"
"gitlab.wanzhuangkj.com/tush/xpkg/pkg/shield/cpu"
"gitlab.wanzhuangkj.com/tush/xpkg/pkg/shield/window"
)
var (
gCPU int64
decay = 0.95
_ Limiter = &BBR{}
)
type (
cpuGetter func() int64
// Option function for bbr limiter
Option func(*options)
)
func init() {
go cpuproc()
}
// cpu = cpuᵗ⁻¹ * decay + cpuᵗ * (1 - decay)
func cpuproc() {
ticker := time.NewTicker(time.Millisecond * 500) // same to cpu sample rate
defer func() {
ticker.Stop()
if err := recover(); err != nil {
go cpuproc()
}
}()
// EMA algorithm: https://blog.csdn.net/m0_38106113/article/details/81542863
for range ticker.C {
stat := &cpu.Stat{}
cpu.ReadStat(stat)
stat.Usage = min(stat.Usage, 1000)
prevCPU := atomic.LoadInt64(&gCPU)
curCPU := int64(float64(prevCPU)*decay + float64(stat.Usage)*(1.0-decay))
atomic.StoreInt64(&gCPU, curCPU)
}
}
func min(l, r uint64) uint64 {
if l < r {
return l
}
return r
}
// Stat contains the metrics snapshot of bbr.
type Stat struct {
CPU int64
InFlight int64
MaxInFlight int64
MinRt int64
MaxPass int64
}
// counterCache is used to cache maxPASS and minRt result.
// Value of current bucket is not counted in real time.
// Cache time is equal to a bucket duration.
type counterCache struct {
val int64
time time.Time
}
// options of bbr limiter.
type options struct {
// WindowSize defines time duration per window
Window time.Duration
// BucketNum defines bucket number for each window
Bucket int
// CPUThreshold
CPUThreshold int64
// CPUQuota
CPUQuota float64
}
// WithWindow with window size.
func WithWindow(d time.Duration) Option {
return func(o *options) {
o.Window = d
}
}
// WithBucket with bucket ize.
func WithBucket(b int) Option {
return func(o *options) {
o.Bucket = b
}
}
// WithCPUThreshold with cpu threshold;
func WithCPUThreshold(threshold int64) Option {
return func(o *options) {
o.CPUThreshold = threshold
}
}
// WithCPUQuota with real cpu quota(if it can not collect from process correct);
func WithCPUQuota(quota float64) Option {
return func(o *options) {
o.CPUQuota = quota
}
}
// BBR implements bbr-like limiter.
// It is inspired by sentinel.
// https://github.com/alibaba/Sentinel/wiki/%E7%B3%BB%E7%BB%9F%E8%87%AA%E9%80%82%E5%BA%94%E9%99%90%E6%B5%81
type BBR struct {
cpu cpuGetter
passStat window.RollingCounter
rtStat window.RollingCounter
inFlight int64
bucketPerSecond int64
bucketDuration time.Duration
// prevDropTime defines previous start drop since initTime
prevDropTime atomic.Value
maxPASSCache atomic.Value
minRtCache atomic.Value
opts options
}
// NewLimiter returns a bbr limiter
func NewLimiter(opts ...Option) *BBR {
opt := options{
Window: time.Second * 10,
Bucket: 100,
CPUThreshold: 800,
}
for _, o := range opts {
o(&opt)
}
bucketDuration := opt.Window / time.Duration(opt.Bucket)
passStat := window.NewRollingCounter(window.RollingCounterOpts{Size: opt.Bucket, BucketDuration: bucketDuration})
rtStat := window.NewRollingCounter(window.RollingCounterOpts{Size: opt.Bucket, BucketDuration: bucketDuration})
limiter := &BBR{
opts: opt,
passStat: passStat,
rtStat: rtStat,
bucketDuration: bucketDuration,
bucketPerSecond: int64(time.Second / bucketDuration),
cpu: func() int64 { return atomic.LoadInt64(&gCPU) },
}
if opt.CPUQuota != 0 {
// if cpuQuota is set, use new cpuGetter,Calculate the real CPU value based on the number of CPUs and Quota.
limiter.cpu = func() int64 {
return int64(float64(atomic.LoadInt64(&gCPU)) * float64(runtime.NumCPU()) / opt.CPUQuota)
}
}
return limiter
}
func (l *BBR) maxPASS() int64 {
passCache := l.maxPASSCache.Load()
if passCache != nil {
ps := passCache.(*counterCache)
if l.timespan(ps.time) < 1 {
return ps.val
}
}
rawMaxPass := int64(l.passStat.Reduce(func(iterator window.Iterator) float64 {
var result = 1.0
for i := 1; iterator.Next() && i < l.opts.Bucket; i++ {
bucket := iterator.Bucket()
count := 0.0
for _, p := range bucket.Points {
count += p
}
result = math.Max(result, count)
}
return result
}))
l.maxPASSCache.Store(&counterCache{
val: rawMaxPass,
time: time.Now(),
})
return rawMaxPass
}
// timespan returns the passed bucket count
// since lastTime, if it is one bucket duration earlier than
// the last recorded time, it will return the BucketNum.
func (l *BBR) timespan(lastTime time.Time) int {
v := int(time.Since(lastTime) / l.bucketDuration)
if v > -1 {
return v
}
return l.opts.Bucket
}
func (l *BBR) minRT() int64 {
rtCache := l.minRtCache.Load()
if rtCache != nil {
rc := rtCache.(*counterCache)
if l.timespan(rc.time) < 1 {
return rc.val
}
}
rawMinRT := int64(math.Ceil(l.rtStat.Reduce(func(iterator window.Iterator) float64 {
var result = math.MaxFloat64
for i := 1; iterator.Next() && i < l.opts.Bucket; i++ {
bucket := iterator.Bucket()
if len(bucket.Points) == 0 {
continue
}
total := 0.0
for _, p := range bucket.Points {
total += p
}
avg := total / float64(bucket.Count)
result = math.Min(result, avg)
}
return result
})))
if rawMinRT <= 0 {
rawMinRT = 1
}
l.minRtCache.Store(&counterCache{
val: rawMinRT,
time: time.Now(),
})
return rawMinRT
}
func (l *BBR) maxInFlight() int64 {
return int64(math.Floor(float64(l.maxPASS()*l.minRT()*l.bucketPerSecond)/1000.0) + 0.5)
}
func (l *BBR) shouldDrop() bool {
now := time.Duration(time.Now().UnixNano())
if l.cpu() < l.opts.CPUThreshold {
// current cpu payload below the threshold
prevDropTime, _ := l.prevDropTime.Load().(time.Duration)
if prevDropTime == 0 {
// haven't start drop,
// accept current request
return false
}
if now-prevDropTime <= time.Second {
// just start drop one second ago, check current inflight count
inFlight := atomic.LoadInt64(&l.inFlight)
return inFlight > 1 && inFlight > l.maxInFlight()
}
l.prevDropTime.Store(time.Duration(0))
return false
}
// current cpu payload exceeds the threshold
inFlight := atomic.LoadInt64(&l.inFlight)
drop := inFlight > 1 && inFlight > l.maxInFlight()
if drop {
prevDrop, _ := l.prevDropTime.Load().(time.Duration)
if prevDrop != 0 {
// already started drop, return directly
return drop
}
// store start drop time
l.prevDropTime.Store(now)
}
return drop
}
// Stat tasks a snapshot of the bbr limiter.
func (l *BBR) Stat() Stat {
return Stat{
CPU: l.cpu(),
MinRt: l.minRT(),
MaxPass: l.maxPASS(),
MaxInFlight: l.maxInFlight(),
InFlight: atomic.LoadInt64(&l.inFlight),
}
}
// Allow checks all inbound traffic.
// Once overload is detected, it raises limit.ErrLimitExceed error.
func (l *BBR) Allow() (DoneFunc, error) {
if l.shouldDrop() {
return nil, ErrLimitExceed
}
atomic.AddInt64(&l.inFlight, 1)
start := time.Now().UnixNano()
ms := float64(time.Millisecond)
return func(DoneInfo) {
rt := int64(math.Ceil(float64(time.Now().UnixNano()-start)) / ms) //nolint
l.rtStat.Add(rt)
atomic.AddInt64(&l.inFlight, -1)
l.passStat.Add(1)
}, nil
}
pkg/shield/ratelimit/bbr_test.go 0 → 100644
浏览文件 @ d57375cc
差异被折叠。
pkg/shield/ratelimit/ratelimit.go 0 → 100644
浏览文件 @ d57375cc
差异被折叠。
pkg/shield/window/counter.go 0 → 100644
浏览文件 @ d57375cc
差异被折叠。
pkg/shield/window/counter_test.go 0 → 100644
浏览文件 @ d57375cc
差异被折叠。
pkg/shield/window/iterator.go 0 → 100644
浏览文件 @ d57375cc
差异被折叠。
pkg/shield/window/policy.go 0 → 100644
浏览文件 @ d57375cc
差异被折叠。
pkg/shield/window/policy_test.go 0 → 100644
浏览文件 @ d57375cc
差异被折叠。
pkg/shield/window/reduce.go 0 → 100644
浏览文件 @ d57375cc
差异被折叠。
pkg/shield/window/reduce_test.go 0 → 100644
浏览文件 @ d57375cc
差异被折叠。
pkg/shield/window/window.go 0 → 100644
浏览文件 @ d57375cc
差异被折叠。
pkg/shield/window/window_test.go 0 → 100644
浏览文件 @ d57375cc
差异被折叠。
Markdown 格式
0% 或
您添加了 0 人 到此讨论。请谨慎行事。
请先完成此评论的编辑！
请 注册 或者 后发表评论