server: improve chained interceptors performance (#4524)

server.go

@@ -1115,22 +1115,24 @@ func chainUnaryServerInterceptors(s *Server) {
 	} else if len(interceptors) == 1 {
 		chainedInt = interceptors[0]
 	} else {
-		chainedInt = func(ctx context.Context, req interface{}, info *UnaryServerInfo, handler UnaryHandler) (interface{}, error) {
+		chainedInt = chainUnaryInterceptors(interceptors)
-			return interceptors[0](ctx, req, info, getChainUnaryHandler(interceptors, 0, info, handler))
-		}
 	}
 
 	s.opts.unaryInt = chainedInt
 }
 
-// getChainUnaryHandler recursively generate the chained UnaryHandler
+func chainUnaryInterceptors(interceptors []UnaryServerInterceptor) UnaryServerInterceptor {
-func getChainUnaryHandler(interceptors []UnaryServerInterceptor, curr int, info *UnaryServerInfo, finalHandler UnaryHandler) UnaryHandler {
+	return func(ctx context.Context, req interface{}, info *UnaryServerInfo, handler UnaryHandler) (interface{}, error) {
-	if curr == len(interceptors)-1 {
+		var i int
-		return finalHandler
+		var next UnaryHandler
-	}
+		next = func(ctx context.Context, req interface{}) (interface{}, error) {
-
+			if i == len(interceptors)-1 {
-	return func(ctx context.Context, req interface{}) (interface{}, error) {
+				return interceptors[i](ctx, req, info, handler)
-		return interceptors[curr+1](ctx, req, info, getChainUnaryHandler(interceptors, curr+1, info, finalHandler))
+			}
+			i++
+			return interceptors[i-1](ctx, req, info, next)
+		}
+		return next(ctx, req)
 	}
 }
 
@@ -1398,22 +1400,24 @@ func chainStreamServerInterceptors(s *Server) {
 	} else if len(interceptors) == 1 {
 		chainedInt = interceptors[0]
 	} else {
-		chainedInt = func(srv interface{}, ss ServerStream, info *StreamServerInfo, handler StreamHandler) error {
+		chainedInt = chainStreamInterceptors(interceptors)
-			return interceptors[0](srv, ss, info, getChainStreamHandler(interceptors, 0, info, handler))
-		}
 	}
 
 	s.opts.streamInt = chainedInt
 }
 
-// getChainStreamHandler recursively generate the chained StreamHandler
+func chainStreamInterceptors(interceptors []StreamServerInterceptor) StreamServerInterceptor {
-func getChainStreamHandler(interceptors []StreamServerInterceptor, curr int, info *StreamServerInfo, finalHandler StreamHandler) StreamHandler {
+	return func(srv interface{}, ss ServerStream, info *StreamServerInfo, handler StreamHandler) error {
-	if curr == len(interceptors)-1 {
+		var i int
-		return finalHandler
+		var next StreamHandler
-	}
+		next = func(srv interface{}, ss ServerStream) error {
-
+			if i == len(interceptors)-1 {
-	return func(srv interface{}, ss ServerStream) error {
+				return interceptors[i](srv, ss, info, handler)
-		return interceptors[curr+1](srv, ss, info, getChainStreamHandler(interceptors, curr+1, info, finalHandler))
+			}
+			i++
+			return interceptors[i-1](srv, ss, info, next)
+		}
+		return next(srv, ss)
 	}
 }
 

server_test.go

@@ -22,6 +22,7 @@ import (
 	"context"
 	"net"
 	"reflect"
+	"strconv"
 	"strings"
 	"testing"
 	"time"
@@ -130,3 +131,59 @@ func (s) TestStreamContext(t *testing.T) {
 		t.Fatalf("GetStreamFromContext(%v) = %v, %t, want: %v, true", ctx, stream, ok, expectedStream)
 	}
 }
+
+func BenchmarkChainUnaryInterceptor(b *testing.B) {
+	for _, n := range []int{1, 3, 5, 10} {
+		n := n
+		b.Run(strconv.Itoa(n), func(b *testing.B) {
+			interceptors := make([]UnaryServerInterceptor, 0, n)
+			for i := 0; i < n; i++ {
+				interceptors = append(interceptors, func(
+					ctx context.Context, req interface{}, info *UnaryServerInfo, handler UnaryHandler,
+				) (interface{}, error) {
+					return handler(ctx, req)
+				})
+			}
+
+			s := NewServer(ChainUnaryInterceptor(interceptors...))
+			b.ReportAllocs()
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				if _, err := s.opts.unaryInt(context.Background(), nil, nil,
+					func(ctx context.Context, req interface{}) (interface{}, error) {
+						return nil, nil
+					},
+				); err != nil {
+					b.Fatal(err)
+				}
+			}
+		})
+	}
+}
+
+func BenchmarkChainStreamInterceptor(b *testing.B) {
+	for _, n := range []int{1, 3, 5, 10} {
+		n := n
+		b.Run(strconv.Itoa(n), func(b *testing.B) {
+			interceptors := make([]StreamServerInterceptor, 0, n)
+			for i := 0; i < n; i++ {
+				interceptors = append(interceptors, func(
+					srv interface{}, ss ServerStream, info *StreamServerInfo, handler StreamHandler,
+				) error {
+					return handler(srv, ss)
+				})
+			}
+
+			s := NewServer(ChainStreamInterceptor(interceptors...))
+			b.ReportAllocs()
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				if err := s.opts.streamInt(nil, nil, nil, func(srv interface{}, stream ServerStream) error {
+					return nil
+				}); err != nil {
+					b.Fatal(err)
+				}
+			}
+		})
+	}
+}