From 339bf3932034a67245423418708e3e4695d36fe6 Mon Sep 17 00:00:00 2001
From: Jim Ma <majinjing3@gmail.com>
Date: Thu, 12 Aug 2021 14:47:38 +0800
Subject: [PATCH] chore: optimize grpc interceptor code (#536)

Signed-off-by: Jim Ma <majinjing3@gmail.com>
---
 go.mod            |  1 -
 go.sum            |  1 -
 pkg/rpc/client.go | 11 ++++-------
 pkg/rpc/server.go | 11 ++++-------
 4 files changed, 8 insertions(+), 16 deletions(-)

diff --git a/go.mod b/go.mod
index 7ff8672233d..a2eab1150f0 100644
--- a/go.mod
+++ b/go.mod
@@ -30,7 +30,6 @@ require (
 	github.com/golang/protobuf v1.5.2
 	github.com/google/uuid v1.2.0
 	github.com/gorilla/mux v1.7.3
-	github.com/grpc-ecosystem/go-grpc-middleware v1.0.0
 	github.com/jarcoal/httpmock v1.0.8
 	github.com/klauspost/compress v1.13.1 // indirect
 	github.com/kr/text v0.2.0 // indirect
diff --git a/go.sum b/go.sum
index 62ede940b49..b30e224d1b1 100644
--- a/go.sum
+++ b/go.sum
@@ -345,7 +345,6 @@ github.com/gorilla/securecookie v1.1.1/go.mod h1:ra0sb63/xPlUeL+yeDciTfxMRAA+MP+
 github.com/gorilla/sessions v1.2.0 h1:S7P+1Hm5V/AT9cjEcUD5uDaQSX0OE577aCXgoaKpYbQ=
 github.com/gorilla/sessions v1.2.0/go.mod h1:dk2InVEVJ0sfLlnXv9EAgkf6ecYs/i80K/zI+bUmuGM=
 github.com/gorilla/websocket v1.4.2/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
-github.com/grpc-ecosystem/go-grpc-middleware v1.0.0 h1:Iju5GlWwrvL6UBg4zJJt3btmonfrMlCDdsejg4CZE7c=
 github.com/grpc-ecosystem/go-grpc-middleware v1.0.0/go.mod h1:FiyG127CGDf3tlThmgyCl78X/SZQqEOJBCDaAfeWzPs=
 github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0/go.mod h1:8NvIoxWQoOIhqOTXgfV/d3M/q6VIi02HzZEHgUlZvzk=
 github.com/grpc-ecosystem/grpc-gateway v1.9.0/go.mod h1:vNeuVxBJEsws4ogUvrchl83t/GYV9WGTSLVdBhOQFDY=
diff --git a/pkg/rpc/client.go b/pkg/rpc/client.go
index 85d4d9569fa..66e5134f441 100644
--- a/pkg/rpc/client.go
+++ b/pkg/rpc/client.go
@@ -22,7 +22,6 @@ import (
 	"sync"
 	"time"
 
-	grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware"
 	"github.com/pkg/errors"
 	"github.com/serialx/hashring"
 	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
@@ -96,12 +95,10 @@ var defaultClientOpts = []grpc.DialOption{
 		Timeout: 10 * time.Second,
 	}),
 	// TODO make grpc interceptor optional
-	grpc.WithStreamInterceptor(grpc_middleware.ChainStreamClient(
-		otelgrpc.StreamClientInterceptor(),
-		streamClientInterceptor)),
-	grpc.WithUnaryInterceptor(grpc_middleware.ChainUnaryClient(
-		otelgrpc.UnaryClientInterceptor(),
-		unaryClientInterceptor)),
+	grpc.WithChainStreamInterceptor(otelgrpc.StreamClientInterceptor(),
+		streamClientInterceptor),
+	grpc.WithChainUnaryInterceptor(otelgrpc.UnaryClientInterceptor(),
+		unaryClientInterceptor),
 }
 
 type ConnOption interface {
diff --git a/pkg/rpc/server.go b/pkg/rpc/server.go
index bc6622c8158..d0639b3bc5c 100644
--- a/pkg/rpc/server.go
+++ b/pkg/rpc/server.go
@@ -29,7 +29,6 @@ import (
 	"syscall"
 	"time"
 
-	grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware"
 	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/keepalive"
@@ -78,12 +77,10 @@ var serverOpts = []grpc.ServerOption{
 	}),
 	grpc.MaxConcurrentStreams(100),
 	// TODO make grpc interceptor optional
-	grpc.StreamInterceptor(grpc_middleware.ChainStreamServer(
-		otelgrpc.StreamServerInterceptor(),
-		streamServerInterceptor)),
-	grpc.UnaryInterceptor(grpc_middleware.ChainUnaryServer(
-		otelgrpc.UnaryServerInterceptor(),
-		unaryServerInterceptor)),
+	grpc.ChainStreamInterceptor(otelgrpc.StreamServerInterceptor(),
+		streamServerInterceptor),
+	grpc.ChainUnaryInterceptor(otelgrpc.UnaryServerInterceptor(),
+		unaryServerInterceptor),
 }
 
 var sp = struct {