initial import

core/rpc/serverinterceptors/crashinterceptor.go

@@ -0,0 +1,43 @@
+package serverinterceptors
+
+import (
+	"context"
+	"runtime/debug"
+
+	"zero/core/logx"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+func StreamCrashInterceptor(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo,
+	handler grpc.StreamHandler) (err error) {
+	defer handleCrash(func(r interface{}) {
+		err = toPanicError(r)
+	})
+
+	return handler(srv, stream)
+}
+
+func UnaryCrashInterceptor() grpc.UnaryServerInterceptor {
+	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo,
+		handler grpc.UnaryHandler) (resp interface{}, err error) {
+		defer handleCrash(func(r interface{}) {
+			err = toPanicError(r)
+		})
+
+		return handler(ctx, req)
+	}
+}
+
+func handleCrash(handler func(interface{})) {
+	if r := recover(); r != nil {
+		handler(r)
+	}
+}
+
+func toPanicError(r interface{}) error {
+	logx.Errorf("%+v %s", r, debug.Stack())
+	return status.Errorf(codes.Internal, "panic: %v", r)
+}

core/rpc/serverinterceptors/prommetricinterceptor.go

@@ -0,0 +1,45 @@
+package serverinterceptors
+
+import (
+	"context"
+	"strconv"
+	"time"
+
+	"zero/core/metric"
+	"zero/core/timex"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/status"
+)
+
+const serverNamespace = "rpc_server"
+
+var (
+	metricServerReqDur = metric.NewHistogramVec(&metric.HistogramVecOpts{
+		Namespace: serverNamespace,
+		Subsystem: "requests",
+		Name:      "duration_ms",
+		Help:      "rpc server requests duration(ms).",
+		Labels:    []string{"method"},
+		Buckets:   []float64{5, 10, 25, 50, 100, 250, 500, 1000},
+	})
+
+	metricServerReqCodeTotal = metric.NewCounterVec(&metric.CounterVecOpts{
+		Namespace: serverNamespace,
+		Subsystem: "requests",
+		Name:      "code_total",
+		Help:      "rpc server requests code count.",
+		Labels:    []string{"method", "code"},
+	})
+)
+
+func UnaryPromMetricInterceptor() grpc.UnaryServerInterceptor {
+	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
+		startTime := timex.Now()
+		resp, err := handler(ctx, req)
+		metricServerReqDur.Observe(int64(timex.Since(startTime)/time.Millisecond), info.FullMethod)
+		metricServerReqCodeTotal.Inc(info.FullMethod, strconv.Itoa(int(status.Code(err))))
+		return resp, err
+	}
+
+}

core/rpc/serverinterceptors/sheddinginterceptor.go

@@ -0,0 +1,53 @@
+package serverinterceptors
+
+import (
+	"context"
+	"sync"
+
+	"zero/core/load"
+	"zero/core/stat"
+
+	"google.golang.org/grpc"
+)
+
+const serviceType = "rpc"
+
+var (
+	sheddingStat *load.SheddingStat
+	lock         sync.Mutex
+)
+
+func UnarySheddingInterceptor(shedder load.Shedder, metrics *stat.Metrics) grpc.UnaryServerInterceptor {
+	ensureSheddingStat()
+
+	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo,
+		handler grpc.UnaryHandler) (val interface{}, 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()
+}

core/rpc/serverinterceptors/statinterceptor.go

@@ -0,0 +1,52 @@
+package serverinterceptors
+
+import (
+	"context"
+	"encoding/json"
+	"time"
+
+	"zero/core/logx"
+	"zero/core/stat"
+	"zero/core/timex"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/peer"
+)
+
+const serverSlowThreshold = time.Millisecond * 500
+
+func UnaryStatInterceptor(metrics *stat.Metrics) grpc.UnaryServerInterceptor {
+	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo,
+		handler grpc.UnaryHandler) (resp interface{}, err error) {
+		defer handleCrash(func(r interface{}) {
+			err = toPanicError(r)
+		})
+
+		startTime := timex.Now()
+		defer func() {
+			duration := timex.Since(startTime)
+			metrics.Add(stat.Task{
+				Duration: duration,
+			})
+			logDuration(ctx, info.FullMethod, req, duration)
+		}()
+
+		return handler(ctx, req)
+	}
+}
+
+func logDuration(ctx context.Context, method string, req interface{}, duration time.Duration) {
+	var addr string
+	client, ok := peer.FromContext(ctx)
+	if ok {
+		addr = client.Addr.String()
+	}
+	content, err := json.Marshal(req)
+	if err != nil {
+		logx.Errorf("%s - %s", addr, err.Error())
+	} else if duration > serverSlowThreshold {
+		logx.WithDuration(duration).Slowf("[RPC] slowcall - %s - %s - %s", addr, method, string(content))
+	} else {
+		logx.WithDuration(duration).Infof("%s - %s - %s", addr, method, string(content))
+	}
+}

core/rpc/serverinterceptors/timeoutinterceptor.go

@@ -0,0 +1,19 @@
+package serverinterceptors
+
+import (
+	"context"
+	"time"
+
+	"zero/core/contextx"
+
+	"google.golang.org/grpc"
+)
+
+func UnaryTimeoutInterceptor(timeout time.Duration) grpc.UnaryServerInterceptor {
+	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo,
+		handler grpc.UnaryHandler) (resp interface{}, err error) {
+		ctx, cancel := contextx.ShrinkDeadline(ctx, timeout)
+		defer cancel()
+		return handler(ctx, req)
+	}
+}

core/rpc/serverinterceptors/tracinginterceptor.go

@@ -0,0 +1,29 @@
+package serverinterceptors
+
+import (
+	"context"
+
+	"zero/core/trace"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
+)
+
+func UnaryTracingInterceptor(serviceName string) grpc.UnaryServerInterceptor {
+	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo,
+		handler grpc.UnaryHandler) (resp interface{}, err error) {
+		md, ok := metadata.FromIncomingContext(ctx)
+		if !ok {
+			return handler(ctx, req)
+		}
+
+		carrier, err := trace.Extract(trace.GrpcFormat, md)
+		if err != nil {
+			return handler(ctx, req)
+		}
+
+		ctx, span := trace.StartServerSpan(ctx, carrier, serviceName, info.FullMethod)
+		defer span.Finish()
+		return handler(ctx, req)
+	}
+}