Added middleware for processing request_id grpc and http requests (#6805)

1 files changed, 26 insertions, 0 deletions

diff --git a/weed/pb/grpc_client_server.go b/weed/pb/grpc_client_server.go
index 777dfb402..ba0d7d0cc 100644
--- a/weed/pb/grpc_client_server.go
+++ b/weed/pb/grpc_client_server.go
@@ -3,6 +3,7 @@ package pb
 import (
 	"context"
 	"fmt"
+	"github.com/google/uuid"
 	"google.golang.org/grpc/metadata"
 	"math/rand/v2"
 	"net/http"
@@ -58,6 +59,7 @@ func NewGrpcServer(opts ...grpc.ServerOption) *grpc.Server {
 		}),
 		grpc.MaxRecvMsgSize(Max_Message_Size),
 		grpc.MaxSendMsgSize(Max_Message_Size),
+		grpc.UnaryInterceptor(requestIDUnaryInterceptor()),
 	)
 	for _, opt := range opts {
 		if opt != nil {
@@ -118,6 +120,30 @@ func getOrCreateConnection(address string, waitForReady bool, opts ...grpc.DialO
 	return vgc, nil
 }
 
+func requestIDUnaryInterceptor() grpc.UnaryServerInterceptor {
+	return func(
+		ctx context.Context,
+		req interface{},
+		info *grpc.UnaryServerInfo,
+		handler grpc.UnaryHandler,
+	) (interface{}, error) {
+		md, _ := metadata.FromIncomingContext(ctx)
+		idList := md.Get(util.RequestIDKey)
+		var reqID string
+		if len(idList) > 0 {
+			reqID = idList[0]
+		}
+		if reqID == "" {
+			reqID = uuid.New().String()
+		}
+
+		ctx = util.WithRequestID(ctx, reqID)
+		grpc.SetTrailer(ctx, metadata.Pairs(util.RequestIDKey, reqID))
+
+		return handler(ctx, req)
+	}
+}
+
 // WithGrpcClient In streamingMode, always use a fresh connection. Otherwise, try to reuse an existing connection.
 func WithGrpcClient(streamingMode bool, signature int32, fn func(*grpc.ClientConn) error, address string, waitForReady bool, opts ...grpc.DialOption) error {