forked from zeromicro/go-zero
/
sheddinginterceptor.go
53 lines (44 loc) · 1.09 KB
/
sheddinginterceptor.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
package serverinterceptors
import (
"context"
"sync"
"github.com/mikokutou1/go-zero-m/core/load"
"github.com/mikokutou1/go-zero-m/core/stat"
"google.golang.org/grpc"
)
const serviceType = "rpc"
var (
sheddingStat *load.SheddingStat
lock sync.Mutex
)
// UnarySheddingInterceptor returns a func that does load shedding on processing unary requests.
func UnarySheddingInterceptor(shedder load.Shedder, metrics *stat.Metrics) grpc.UnaryServerInterceptor {
ensureSheddingStat()
return func(ctx context.Context, req any, info *grpc.UnaryServerInfo,
handler grpc.UnaryHandler) (val any, err error) {
sheddingStat.IncrementTotal()
var promise load.Promise
promise, err = shedder.Allow()
if err != nil {
metrics.AddDrop()
sheddingStat.IncrementDrop()
return
}
defer func() {
if err == context.DeadlineExceeded {
promise.Fail()
} else {
sheddingStat.IncrementPass()
promise.Pass()
}
}()
return handler(ctx, req)
}
}
func ensureSheddingStat() {
lock.Lock()
if sheddingStat == nil {
sheddingStat = load.NewSheddingStat(serviceType)
}
lock.Unlock()
}