package httphandler import ( "net/http" "sync" "zero/core/httphandler/internal" "zero/core/httpx" "zero/core/load" "zero/core/logx" "zero/core/stat" ) const serviceType = "api" var ( sheddingStat *load.SheddingStat lock sync.Mutex ) func SheddingHandler(shedder load.Shedder, metrics *stat.Metrics) func(http.Handler) http.Handler { if shedder == nil { return func(next http.Handler) http.Handler { return next } } ensureSheddingStat() return func(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { sheddingStat.IncrementTotal() promise, err := shedder.Allow() if err != nil { metrics.AddDrop() sheddingStat.IncrementDrop() logx.Errorf("[http] dropped, %s - %s - %s", r.RequestURI, httpx.GetRemoteAddr(r), r.UserAgent()) w.WriteHeader(http.StatusServiceUnavailable) return } cw := &internal.WithCodeResponseWriter{Writer: w} defer func() { if cw.Code == http.StatusServiceUnavailable { promise.Fail() } else { sheddingStat.IncrementPass() promise.Pass() } }() next.ServeHTTP(cw, r) }) } } func ensureSheddingStat() { lock.Lock() if sheddingStat == nil { sheddingStat = load.NewSheddingStat(serviceType) } lock.Unlock() }