rename rpcx to zrpc

zrpc/internal/serverinterceptors/authinterceptor.go

@@ -0,0 +1,30 @@
+package serverinterceptors
+
+import (
+	"context"
+
+	"github.com/tal-tech/go-zero/zrpc/internal/auth"
+	"google.golang.org/grpc"
+)
+
+func StreamAuthorizeInterceptor(authenticator *auth.Authenticator) grpc.StreamServerInterceptor {
+	return func(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo,
+		handler grpc.StreamHandler) error {
+		if err := authenticator.Authenticate(stream.Context()); err != nil {
+			return err
+		}
+
+		return handler(srv, stream)
+	}
+}
+
+func UnaryAuthorizeInterceptor(authenticator *auth.Authenticator) grpc.UnaryServerInterceptor {
+	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo,
+		handler grpc.UnaryHandler) (interface{}, error) {
+		if err := authenticator.Authenticate(ctx); err != nil {
+			return nil, err
+		}
+
+		return handler(ctx, req)
+	}
+}

zrpc/internal/serverinterceptors/authinterceptor_test.go

@@ -0,0 +1,200 @@
+package serverinterceptors
+
+import (
+	"context"
+	"testing"
+
+	"github.com/alicebob/miniredis"
+	"github.com/stretchr/testify/assert"
+	"github.com/tal-tech/go-zero/core/stores/redis"
+	"github.com/tal-tech/go-zero/zrpc/internal/auth"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
+)
+
+func TestStreamAuthorizeInterceptor(t *testing.T) {
+	tests := []struct {
+		name     string
+		app      string
+		token    string
+		strict   bool
+		hasError bool
+	}{
+		{
+			name:     "strict=false",
+			strict:   false,
+			hasError: false,
+		},
+		{
+			name:     "strict=true",
+			strict:   true,
+			hasError: true,
+		},
+		{
+			name:     "strict=true,with token",
+			app:      "foo",
+			token:    "bar",
+			strict:   true,
+			hasError: false,
+		},
+		{
+			name:     "strict=true,with error token",
+			app:      "foo",
+			token:    "error",
+			strict:   true,
+			hasError: true,
+		},
+	}
+
+	r := miniredis.NewMiniRedis()
+	assert.Nil(t, r.Start())
+	defer r.Close()
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			store := redis.NewRedis(r.Addr(), redis.NodeType)
+			if len(test.app) > 0 {
+				assert.Nil(t, store.Hset("apps", test.app, test.token))
+				defer store.Hdel("apps", test.app)
+			}
+
+			authenticator, err := auth.NewAuthenticator(store, "apps", test.strict)
+			assert.Nil(t, err)
+			interceptor := StreamAuthorizeInterceptor(authenticator)
+			md := metadata.New(map[string]string{
+				"app":   "foo",
+				"token": "bar",
+			})
+			ctx := metadata.NewIncomingContext(context.Background(), md)
+			stream := mockedStream{ctx: ctx}
+			err = interceptor(nil, stream, nil, func(srv interface{}, stream grpc.ServerStream) error {
+				return nil
+			})
+			if test.hasError {
+				assert.NotNil(t, err)
+			} else {
+				assert.Nil(t, err)
+			}
+		})
+	}
+}
+
+func TestUnaryAuthorizeInterceptor(t *testing.T) {
+	tests := []struct {
+		name     string
+		app      string
+		token    string
+		strict   bool
+		hasError bool
+	}{
+		{
+			name:     "strict=false",
+			strict:   false,
+			hasError: false,
+		},
+		{
+			name:     "strict=true",
+			strict:   true,
+			hasError: true,
+		},
+		{
+			name:     "strict=true,with token",
+			app:      "foo",
+			token:    "bar",
+			strict:   true,
+			hasError: false,
+		},
+		{
+			name:     "strict=true,with error token",
+			app:      "foo",
+			token:    "error",
+			strict:   true,
+			hasError: true,
+		},
+	}
+
+	r := miniredis.NewMiniRedis()
+	assert.Nil(t, r.Start())
+	defer r.Close()
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			store := redis.NewRedis(r.Addr(), redis.NodeType)
+			if len(test.app) > 0 {
+				assert.Nil(t, store.Hset("apps", test.app, test.token))
+				defer store.Hdel("apps", test.app)
+			}
+
+			authenticator, err := auth.NewAuthenticator(store, "apps", test.strict)
+			assert.Nil(t, err)
+			interceptor := UnaryAuthorizeInterceptor(authenticator)
+			md := metadata.New(map[string]string{
+				"app":   "foo",
+				"token": "bar",
+			})
+			ctx := metadata.NewIncomingContext(context.Background(), md)
+			_, err = interceptor(ctx, nil, nil,
+				func(ctx context.Context, req interface{}) (interface{}, error) {
+					return nil, nil
+				})
+			if test.hasError {
+				assert.NotNil(t, err)
+			} else {
+				assert.Nil(t, err)
+			}
+			if test.strict {
+				_, err = interceptor(context.Background(), nil, nil,
+					func(ctx context.Context, req interface{}) (interface{}, error) {
+						return nil, nil
+					})
+				assert.NotNil(t, err)
+
+				var md metadata.MD
+				ctx := metadata.NewIncomingContext(context.Background(), md)
+				_, err = interceptor(ctx, nil, nil,
+					func(ctx context.Context, req interface{}) (interface{}, error) {
+						return nil, nil
+					})
+				assert.NotNil(t, err)
+
+				md = metadata.New(map[string]string{
+					"app":   "",
+					"token": "",
+				})
+				ctx = metadata.NewIncomingContext(context.Background(), md)
+				_, err = interceptor(ctx, nil, nil,
+					func(ctx context.Context, req interface{}) (interface{}, error) {
+						return nil, nil
+					})
+				assert.NotNil(t, err)
+			}
+		})
+	}
+}
+
+type mockedStream struct {
+	ctx context.Context
+}
+
+func (m mockedStream) SetHeader(md metadata.MD) error {
+	return nil
+}
+
+func (m mockedStream) SendHeader(md metadata.MD) error {
+	return nil
+}
+
+func (m mockedStream) SetTrailer(md metadata.MD) {
+}
+
+func (m mockedStream) Context() context.Context {
+	return m.ctx
+}
+
+func (m mockedStream) SendMsg(v interface{}) error {
+	return nil
+}
+
+func (m mockedStream) RecvMsg(v interface{}) error {
+	return nil
+}

zrpc/internal/serverinterceptors/crashinterceptor.go

@@ -0,0 +1,42 @@
+package serverinterceptors
+
+import (
+	"context"
+	"runtime/debug"
+
+	"github.com/tal-tech/go-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)
+}

zrpc/internal/serverinterceptors/crashinterceptor_test.go

@@ -0,0 +1,31 @@
+package serverinterceptors
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/tal-tech/go-zero/core/logx"
+	"google.golang.org/grpc"
+)
+
+func init() {
+	logx.Disable()
+}
+
+func TestStreamCrashInterceptor(t *testing.T) {
+	err := StreamCrashInterceptor(nil, nil, nil, func(
+		srv interface{}, stream grpc.ServerStream) error {
+		panic("mock panic")
+	})
+	assert.NotNil(t, err)
+}
+
+func TestUnaryCrashInterceptor(t *testing.T) {
+	interceptor := UnaryCrashInterceptor()
+	_, err := interceptor(context.Background(), nil, nil,
+		func(ctx context.Context, req interface{}) (interface{}, error) {
+			panic("mock panic")
+		})
+	assert.NotNil(t, err)
+}

zrpc/internal/serverinterceptors/prommetricinterceptor.go

@@ -0,0 +1,44 @@
+package serverinterceptors
+
+import (
+	"context"
+	"strconv"
+	"time"
+
+	"github.com/tal-tech/go-zero/core/metric"
+	"github.com/tal-tech/go-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
+	}
+}

zrpc/internal/serverinterceptors/prommetricinterceptor_test.go

@@ -0,0 +1,19 @@
+package serverinterceptors
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"google.golang.org/grpc"
+)
+
+func TestUnaryPromMetricInterceptor(t *testing.T) {
+	interceptor := UnaryPromMetricInterceptor()
+	_, err := interceptor(context.Background(), nil, &grpc.UnaryServerInfo{
+		FullMethod: "/",
+	}, func(ctx context.Context, req interface{}) (interface{}, error) {
+		return nil, nil
+	})
+	assert.Nil(t, err)
+}

zrpc/internal/serverinterceptors/sheddinginterceptor.go

@@ -0,0 +1,52 @@
+package serverinterceptors
+
+import (
+	"context"
+	"sync"
+
+	"github.com/tal-tech/go-zero/core/load"
+	"github.com/tal-tech/go-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()
+}

zrpc/internal/serverinterceptors/sheddinginterceptor_test.go

@@ -0,0 +1,77 @@
+package serverinterceptors
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/tal-tech/go-zero/core/load"
+	"github.com/tal-tech/go-zero/core/stat"
+	"google.golang.org/grpc"
+)
+
+func TestUnarySheddingInterceptor(t *testing.T) {
+	tests := []struct {
+		name      string
+		allow     bool
+		handleErr error
+		expect    error
+	}{
+		{
+			name:      "allow",
+			allow:     true,
+			handleErr: nil,
+			expect:    nil,
+		},
+		{
+			name:      "allow",
+			allow:     true,
+			handleErr: context.DeadlineExceeded,
+			expect:    context.DeadlineExceeded,
+		},
+		{
+			name:      "reject",
+			allow:     false,
+			handleErr: nil,
+			expect:    load.ErrServiceOverloaded,
+		},
+	}
+
+	for _, test := range tests {
+		test := test
+		t.Run(test.name, func(t *testing.T) {
+			t.Parallel()
+
+			shedder := mockedShedder{allow: test.allow}
+			metrics := stat.NewMetrics("mock")
+			interceptor := UnarySheddingInterceptor(shedder, metrics)
+			_, err := interceptor(context.Background(), nil, &grpc.UnaryServerInfo{
+				FullMethod: "/",
+			}, func(ctx context.Context, req interface{}) (interface{}, error) {
+				return nil, test.handleErr
+			})
+			assert.Equal(t, test.expect, err)
+		})
+	}
+}
+
+type mockedShedder struct {
+	allow bool
+}
+
+func (m mockedShedder) Allow() (load.Promise, error) {
+	if m.allow {
+		return mockedPromise{}, nil
+	} else {
+		return nil, load.ErrServiceOverloaded
+	}
+}
+
+type mockedPromise struct {
+}
+
+func (m mockedPromise) Pass() {
+}
+
+func (m mockedPromise) Fail() {
+}

zrpc/internal/serverinterceptors/statinterceptor.go

@@ -0,0 +1,51 @@
+package serverinterceptors
+
+import (
+	"context"
+	"encoding/json"
+	"time"
+
+	"github.com/tal-tech/go-zero/core/logx"
+	"github.com/tal-tech/go-zero/core/stat"
+	"github.com/tal-tech/go-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))
+	}
+}

zrpc/internal/serverinterceptors/statinterceptor_test.go

@@ -0,0 +1,32 @@
+package serverinterceptors
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/tal-tech/go-zero/core/stat"
+	"google.golang.org/grpc"
+)
+
+func TestUnaryStatInterceptor(t *testing.T) {
+	metrics := stat.NewMetrics("mock")
+	interceptor := UnaryStatInterceptor(metrics)
+	_, err := interceptor(context.Background(), nil, &grpc.UnaryServerInfo{
+		FullMethod: "/",
+	}, func(ctx context.Context, req interface{}) (interface{}, error) {
+		return nil, nil
+	})
+	assert.Nil(t, err)
+}
+
+func TestUnaryStatInterceptor_crash(t *testing.T) {
+	metrics := stat.NewMetrics("mock")
+	interceptor := UnaryStatInterceptor(metrics)
+	_, err := interceptor(context.Background(), nil, &grpc.UnaryServerInfo{
+		FullMethod: "/",
+	}, func(ctx context.Context, req interface{}) (interface{}, error) {
+		panic("error")
+	})
+	assert.NotNil(t, err)
+}

zrpc/internal/serverinterceptors/timeoutinterceptor.go

@@ -0,0 +1,18 @@
+package serverinterceptors
+
+import (
+	"context"
+	"time"
+
+	"github.com/tal-tech/go-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)
+	}
+}

zrpc/internal/serverinterceptors/timeoutinterceptor_test.go

@@ -0,0 +1,41 @@
+package serverinterceptors
+
+import (
+	"context"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"google.golang.org/grpc"
+)
+
+func TestUnaryTimeoutInterceptor(t *testing.T) {
+	interceptor := UnaryTimeoutInterceptor(time.Millisecond * 10)
+	_, err := interceptor(context.Background(), nil, &grpc.UnaryServerInfo{
+		FullMethod: "/",
+	}, func(ctx context.Context, req interface{}) (interface{}, error) {
+		return nil, nil
+	})
+	assert.Nil(t, err)
+}
+
+func TestUnaryTimeoutInterceptor_timeout(t *testing.T) {
+	const timeout = time.Millisecond * 10
+	interceptor := UnaryTimeoutInterceptor(timeout)
+	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
+	defer cancel()
+	var wg sync.WaitGroup
+	wg.Add(1)
+	_, err := interceptor(ctx, nil, &grpc.UnaryServerInfo{
+		FullMethod: "/",
+	}, func(ctx context.Context, req interface{}) (interface{}, error) {
+		defer wg.Done()
+		tm, ok := ctx.Deadline()
+		assert.True(t, ok)
+		assert.True(t, tm.Before(time.Now().Add(timeout+time.Millisecond)))
+		return nil, nil
+	})
+	wg.Wait()
+	assert.Nil(t, err)
+}

zrpc/internal/serverinterceptors/tracinginterceptor.go

@@ -0,0 +1,28 @@
+package serverinterceptors
+
+import (
+	"context"
+
+	"github.com/tal-tech/go-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)
+	}
+}

zrpc/internal/serverinterceptors/tracinginterceptor_test.go

@@ -0,0 +1,48 @@
+package serverinterceptors
+
+import (
+	"context"
+	"sync"
+	"sync/atomic"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/tal-tech/go-zero/core/trace/tracespec"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
+)
+
+func TestUnaryTracingInterceptor(t *testing.T) {
+	interceptor := UnaryTracingInterceptor("foo")
+	var run int32
+	var wg sync.WaitGroup
+	wg.Add(1)
+	_, err := interceptor(context.Background(), nil, &grpc.UnaryServerInfo{
+		FullMethod: "/",
+	}, func(ctx context.Context, req interface{}) (interface{}, error) {
+		defer wg.Done()
+		atomic.AddInt32(&run, 1)
+		return nil, nil
+	})
+	wg.Wait()
+	assert.Nil(t, err)
+	assert.Equal(t, int32(1), atomic.LoadInt32(&run))
+}
+
+func TestUnaryTracingInterceptor_GrpcFormat(t *testing.T) {
+	interceptor := UnaryTracingInterceptor("foo")
+	var wg sync.WaitGroup
+	wg.Add(1)
+	var md metadata.MD
+	ctx := metadata.NewIncomingContext(context.Background(), md)
+	_, err := interceptor(ctx, nil, &grpc.UnaryServerInfo{
+		FullMethod: "/",
+	}, func(ctx context.Context, req interface{}) (interface{}, error) {
+		defer wg.Done()
+		assert.True(t, len(ctx.Value(tracespec.TracingKey).(tracespec.Trace).TraceId()) > 0)
+		assert.True(t, len(ctx.Value(tracespec.TracingKey).(tracespec.Trace).SpanId()) > 0)
+		return nil, nil
+	})
+	wg.Wait()
+	assert.Nil(t, err)
+}