add request id to grpc logging fields

Signed-off-by: Coleen Iona Quadros <coleen.quadros27@gmail.com>
thanos-io · Oct 18, 2023 · ae8e6ed · ae8e6ed
1 parent 6cf226c
commit ae8e6ed
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 2 deletions.
diff --git a/pkg/logging/grpc.go b/pkg/logging/grpc.go
@@ -152,6 +152,7 @@ func NewGRPCOption(configYAML []byte) ([]grpc_logging.Option, []string, error) {
 	logOpts = []grpc_logging.Option{
 		grpc_logging.WithLevels(DefaultCodeToLevelGRPC),
 		grpc_logging.WithFieldsFromContext(GetTraceIDAsField),
+		grpc_logging.WithFieldsFromContext(GetRequestIDAsField),
 	}
 	// If the combination is valid, use them, otherwise return error.
 	reqLogDecision, err := getGRPCLoggingOption(globalStart, globalEnd)

diff --git a/pkg/logging/options.go b/pkg/logging/options.go
@@ -6,15 +6,18 @@ package logging
 import (
 	"context"
 	"fmt"
+	"math/rand"
 	"time"
 
+	"github.com/oklog/ulid"
 	"github.com/opentracing/opentracing-go"
 	"go.opentelemetry.io/otel/trace"
 
 	extflag "github.com/efficientgo/tools/extkingpin"
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
 	grpc_logging "github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/logging"
+	middleware "github.com/thanos-io/thanos/pkg/server/http/middleware"
 	"github.com/thanos-io/thanos/pkg/tracing/migration"
 	"google.golang.org/grpc/codes"
 )
@@ -169,6 +172,22 @@ func GetTraceIDAsField(ctx context.Context) grpc_logging.Fields {
 		return grpc_logging.Fields{"traceID", TraceID}
 	}
 	return nil
+
+}
+
+func GetRequestIDAsField(ctx context.Context) grpc_logging.Fields {
+	reqID, ok := middleware.RequestIDFromContext(ctx)
+	if ok {
+		return grpc_logging.Fields{"requestID", reqID}
+	}
+
+	entropy := ulid.Monotonic(rand.New(rand.NewSource(time.Now().UnixNano())), 0)
+	newReqID := ulid.MustNew(ulid.Timestamp(time.Now()), entropy).String()
+
+	// Insert into context (Note: This updated context isn't being used later, so this might be redundant)
+	ctx = middleware.NewContextWithRequestID(ctx, newReqID)
+
+	return grpc_logging.Fields{"requestID", newReqID}
 }
 
 // TODO: @yashrsharma44 - To be deprecated in the next release.

diff --git a/pkg/server/http/middleware/request_id.go b/pkg/server/http/middleware/request_id.go
@@ -17,7 +17,7 @@ type ctxKey int
 const reqIDKey = ctxKey(0)
 
 // newContextWithRequestID creates a context with a request id.
-func newContextWithRequestID(ctx context.Context, rid string) context.Context {
+func NewContextWithRequestID(ctx context.Context, rid string) context.Context {
 	return context.WithValue(ctx, reqIDKey, rid)
 }
 
@@ -36,7 +36,7 @@ func RequestID(h http.Handler) http.HandlerFunc {
 			reqID := ulid.MustNew(ulid.Timestamp(time.Now()), entropy)
 			r.Header.Set("X-Request-ID", reqID.String())
 		}
-		ctx := newContextWithRequestID(r.Context(), reqID)
+		ctx := NewContextWithRequestID(r.Context(), reqID)
 		h.ServeHTTP(w, r.WithContext(ctx))
 	}
 }