From 22e0fc216bfbc2e95af463f4dd4f796067852d67 Mon Sep 17 00:00:00 2001
From: wangchuxiao <wangchuxiao97@outlook.com>
Date: Tue, 9 May 2023 11:35:59 +0800
Subject: [PATCH] grpc mw

---
 pkg/common/mw/rpc_client_interceptor.go | 4 ++--
 pkg/common/mw/rpc_server_interceptor.go | 4 ++--
 pkg/startrpc/start.go                   | 6 ++++--
 3 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/pkg/common/mw/rpc_client_interceptor.go b/pkg/common/mw/rpc_client_interceptor.go
index 0b7f69e84..eed52f003 100644
--- a/pkg/common/mw/rpc_client_interceptor.go
+++ b/pkg/common/mw/rpc_client_interceptor.go
@@ -16,10 +16,10 @@ import (
 )
 
 func GrpcClient() grpc.DialOption {
-	return grpc.WithUnaryInterceptor(rpcClientInterceptor)
+	return grpc.WithUnaryInterceptor(RpcClientInterceptor)
 }
 
-func rpcClientInterceptor(ctx context.Context, method string, req, resp interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) (err error) {
+func RpcClientInterceptor(ctx context.Context, method string, req, resp interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) (err error) {
 	if ctx == nil {
 		return errs.ErrInternalServer.Wrap("call rpc request context is nil")
 	}
diff --git a/pkg/common/mw/rpc_server_interceptor.go b/pkg/common/mw/rpc_server_interceptor.go
index 9da31b716..bece93733 100644
--- a/pkg/common/mw/rpc_server_interceptor.go
+++ b/pkg/common/mw/rpc_server_interceptor.go
@@ -29,7 +29,7 @@ func rpcString(v interface{}) string {
 	return fmt.Sprintf("%+v", v)
 }
 
-func rpcServerInterceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
+func RpcServerInterceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
 	defer func() {
 		if r := recover(); r != nil {
 			log.ZError(ctx, "rpc panic", nil, "FullMethod", info.FullMethod, "type:", fmt.Sprintf("%T", r), "panic:", r)
@@ -144,5 +144,5 @@ func rpcServerInterceptor(ctx context.Context, req interface{}, info *grpc.Unary
 }
 
 func GrpcServer() grpc.ServerOption {
-	return grpc.UnaryInterceptor(rpcServerInterceptor)
+	return grpc.UnaryInterceptor(RpcServerInterceptor)
 }
diff --git a/pkg/startrpc/start.go b/pkg/startrpc/start.go
index a9db548a6..fcfff47c1 100644
--- a/pkg/startrpc/start.go
+++ b/pkg/startrpc/start.go
@@ -40,16 +40,18 @@ func Start(rpcPort int, rpcRegisterName string, prometheusPort int, rpcFn func(c
 	if err != nil {
 		return err
 	}
-	options = append(options, mw.GrpcServer()) // ctx 中间件
+	// ctx 中间件
 	if config.Config.Prometheus.Enable {
 		prome.NewGrpcRequestCounter()
 		prome.NewGrpcRequestFailedCounter()
 		prome.NewGrpcRequestSuccessCounter()
-		unaryInterceptor := mw.InterceptChain(grpcPrometheus.UnaryServerInterceptor, grpcPrometheus.UnaryServerInterceptor)
+		unaryInterceptor := mw.InterceptChain(grpcPrometheus.UnaryServerInterceptor, mw.RpcServerInterceptor)
 		options = append(options, []grpc.ServerOption{
 			grpc.StreamInterceptor(grpcPrometheus.StreamServerInterceptor),
 			grpc.UnaryInterceptor(unaryInterceptor),
 		}...)
+	} else {
+		options = append(options, mw.GrpcServer())
 	}
 	srv := grpc.NewServer(options...)
 	defer srv.GracefulStop()