rename ngin to rest

rest/handler/sheddinghandler.go

@@ -0,0 +1,63 @@
+package handler
+
+import (
+	"net/http"
+	"sync"
+
+	"zero/core/load"
+	"zero/core/logx"
+	"zero/core/stat"
+	"zero/rest/internal"
+	"zero/rest/internal/security"
+)
+
+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, internal.GetRemoteAddr(r), r.UserAgent())
+				w.WriteHeader(http.StatusServiceUnavailable)
+				return
+			}
+
+			cw := &security.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()
+}