save freshly generated requestID to ctx

Author: qrort

internal/audit/audit_event.go

@@ -153,10 +153,11 @@ func GRPCCallAuditEvent(
 	err error,
 ) *GRPCCallEvent {
 	s, r := getStatus(inProgress, err)
+	requestID, _ := grpcinfo.GetRequestID(ctx)
 	return &GRPCCallEvent{
 		GenericAuditFields: GenericAuditFields{
 			ID:             uuid.New().String(),
-			IdempotencyKey: grpcinfo.GetRequestID(ctx),
+			IdempotencyKey: requestID,
 			TraceID:        formatTraceID(grpcinfo.GetTraceID(ctx)),
 			Service:        "ydbcp",
 			SpecVersion:    "1.0",

internal/audit/audit_event_test.go

@@ -15,6 +15,7 @@ import (
 	"os"
 	"testing"
 	"time"
+	"ydbcp/internal/server/grpcinfo"
 	"ydbcp/internal/types"
 	pb "ydbcp/pkg/proto/ydbcp/v1alpha1"
 )
@@ -382,3 +383,20 @@ func (m *mockAddr) Network() string {
 func (m *mockAddr) String() string {
 	return m.address
 }
+
+func TestWithGRPCInfo(t *testing.T) {
+	ctx := context.Background()
+	ctx = grpcinfo.WithGRPCInfo(ctx)
+	SetAuditFieldsForRequest(
+		ctx, &AuditFields{
+			ContainerID: "container-1",
+			Database:    "db-1",
+		},
+	)
+
+	requestID, _ := grpcinfo.GetRequestID(ctx)
+	fields := GetAuditFieldsForRequest(requestID)
+	require.NotNil(t, fields)
+	require.Equal(t, "container-1", fields.ContainerID)
+	require.Equal(t, "db-1", fields.Database)
+}

internal/audit/audit_interceptor.go

@@ -22,7 +22,8 @@ var (
 )
 
 func SetAuditFieldsForRequest(ctx context.Context, fields *AuditFields) {
-	containerStore.Store(grpcinfo.GetRequestID(ctx), fields)
+	requestID, _ := grpcinfo.GetRequestID(ctx)
+	containerStore.Store(requestID, fields)
 }
 
 func GetAuditFieldsForRequest(requestID string) *AuditFields {
@@ -42,8 +43,6 @@ func NewAuditGRPCInterceptor(provider auth.AuthProvider) grpc.UnaryServerInterce
 		ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler,
 	) (interface{}, error) {
 		ctx = grpcinfo.WithGRPCInfo(ctx)
-		requestID := grpcinfo.GetRequestID(ctx)
-		ctx = grpcinfo.SetRequestID(ctx, requestID)
 		subject, _ := authHelper.Authenticate(ctx, provider)
 		token, _ := authHelper.GetMaskedToken(ctx, provider)
 		pm, ok := req.(proto.Message)
@@ -55,6 +54,7 @@ func NewAuditGRPCInterceptor(provider auth.AuthProvider) grpc.UnaryServerInterce
 			)
 		}
 		response, grpcErr := handler(ctx, req)
+		requestID, _ := grpcinfo.GetRequestID(ctx)
 		fields := GetAuditFieldsForRequest(requestID)
 		defer ClearAuditFieldsForRequest(requestID)
 		ReportGRPCCallEnd(ctx, info.FullMethod, subject, fields.ContainerID, fields.Database, token, grpcErr)

internal/server/grpcinfo/grpcinfo.go

@@ -27,17 +27,17 @@ func SetRequestID(ctx context.Context, id string) context.Context {
 	return context.WithValue(ctx, ctxKeyRequestID{}, id)
 }
 
-func GetRequestID(ctx context.Context) string {
+func GetRequestID(ctx context.Context) (string, bool) {
 	if id, ok := ctx.Value(ctxKeyRequestID{}).(string); ok {
-		return id
+		return id, false
 	}
 	for _, key := range []string{"RequestID", "RequestId", "request-id", "request_id"} {
 		val := getFromCtx(ctx, key)
 		if val != nil {
-			return *val
+			return *val, false
 		}
 	}
-	return uuid.New().String()
+	return uuid.New().String(), true
 }
 
 func GetTraceID(ctx context.Context) *string {
@@ -76,12 +76,15 @@ func WithGRPCInfo(ctx context.Context) context.Context {
 	if method, ok := grpc.Method(ctx); ok {
 		ctx = xlog.With(ctx, zap.String("GRPCMethod", method))
 	}
-	requestID := GetRequestID(ctx)
+	requestID, newID := GetRequestID(ctx)
 	ctx = xlog.With(ctx, zap.String("RequestID", requestID))
 	err := grpc.SendHeader(ctx, metadata.Pairs("X-Request-ID", requestID))
 	if err != nil {
 		xlog.Error(ctx, "failed to set X-Request-ID header", zap.Error(err))
 	}
 	xlog.Debug(ctx, "New grpc request")
+	if newID {
+		ctx = SetRequestID(ctx, requestID)
+	}
 	return ctx
 }

internal/server/grpcinfo/grpcinfo_test.go

@@ -2,10 +2,11 @@ package grpcinfo
 
 import (
 	"context"
-	"testing"
-
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"google.golang.org/grpc/metadata"
+	"testing"
+	"ydbcp/internal/util/xlog"
 )
 
 func TestGetRemoteAddressChain(t *testing.T) {
@@ -44,3 +45,31 @@ func TestGetTraceID(t *testing.T) {
 	ctx = context.WithValue(ctx, "trace_id", traceID)
 	assert.Equal(t, traceID, *GetTraceID(ctx))
 }
+
+func TestRequestID(t *testing.T) {
+	ctx := context.Background()
+
+	id, generated := GetRequestID(ctx)
+	require.True(t, generated)
+
+	ctx = SetRequestID(ctx, id)
+
+	id2, generated := GetRequestID(ctx)
+	require.False(t, generated)
+	require.Equal(t, id, id2)
+}
+
+func TestWithGRPCInfo(t *testing.T) {
+	logger, err := xlog.SetupLogging("DEBUG")
+	require.NoError(t, err)
+	xlog.SetInternalLogger(logger)
+
+	ctx := context.Background()
+	ctx = WithGRPCInfo(ctx)
+
+	id, generated := GetRequestID(ctx)
+	require.False(t, generated)
+	id2, generated := GetRequestID(ctx)
+	require.False(t, generated)
+	require.Equal(t, id, id2)
+}