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

This commit is contained in:
Aleksey Kosov
2025-05-21 17:57:39 +03:00
committed by GitHub
parent 140b7a7402
commit 5182d46e22
6 changed files with 88 additions and 21 deletions

View File

@@ -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 {