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

6 files changed, 88 insertions, 21 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 {
 
diff --git a/weed/server/common.go b/weed/server/common.go
index 5dad9d81b..516f8bf1c 100644
--- a/weed/server/common.go
+++ b/weed/server/common.go
@@ -3,9 +3,12 @@ package weed_server
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
+	"github.com/google/uuid"
+	"google.golang.org/grpc/metadata"
 	"io"
 	"io/fs"
 	"mime/multipart"
@@ -421,3 +424,21 @@ func ProcessRangeRequest(r *http.Request, w http.ResponseWriter, totalSize int64
 	}
 	return nil
 }
+
+func requestIDMiddleware(h http.HandlerFunc) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		reqID := r.Header.Get(util.RequestIdHttpHeader)
+		if reqID == "" {
+			reqID = uuid.New().String()
+		}
+
+		ctx := context.WithValue(r.Context(), util.RequestIDKey, reqID)
+		ctx = metadata.NewOutgoingContext(ctx,
+			metadata.New(map[string]string{
+				util.RequestIDKey: reqID,
+			}))
+
+		w.Header().Set(util.RequestIdHttpHeader, reqID)
+		h(w, r.WithContext(ctx))
+	}
+}
diff --git a/weed/server/filer_server.go b/weed/server/filer_server.go
index 090c795fa..57af97d84 100644
--- a/weed/server/filer_server.go
+++ b/weed/server/filer_server.go
@@ -189,13 +189,13 @@ func NewFilerServer(defaultMux, readonlyMux *http.ServeMux, option *FilerOption)
 
 	handleStaticResources(defaultMux)
 	if !option.DisableHttp {
-		defaultMux.HandleFunc("/healthz", fs.filerHealthzHandler)
-		defaultMux.HandleFunc("/", fs.filerGuard.WhiteList(fs.filerHandler))
+		defaultMux.HandleFunc("/healthz", requestIDMiddleware(fs.filerHealthzHandler))
+		defaultMux.HandleFunc("/", fs.filerGuard.WhiteList(requestIDMiddleware(fs.filerHandler)))
 	}
 	if defaultMux != readonlyMux {
 		handleStaticResources(readonlyMux)
-		readonlyMux.HandleFunc("/healthz", fs.filerHealthzHandler)
-		readonlyMux.HandleFunc("/", fs.filerGuard.WhiteList(fs.readonlyFilerHandler))
+		readonlyMux.HandleFunc("/healthz", requestIDMiddleware(fs.filerHealthzHandler))
+		readonlyMux.HandleFunc("/", fs.filerGuard.WhiteList(requestIDMiddleware(fs.readonlyFilerHandler)))
 	}
 
 	existingNodes := fs.filer.ListExistingPeerUpdates(context.Background())
diff --git a/weed/server/master_server.go b/weed/server/master_server.go
index 8621708d2..6569fdbd4 100644
--- a/weed/server/master_server.go
+++ b/weed/server/master_server.go
@@ -134,24 +134,24 @@ func NewMasterServer(r *mux.Router, option *MasterOption, peers map[string]pb.Se
 	ms.guard = security.NewGuard(append(ms.option.WhiteList, whiteList...), signingKey, expiresAfterSec, readSigningKey, readExpiresAfterSec)
 
 	handleStaticResources2(r)
-	r.HandleFunc("/", ms.proxyToLeader(ms.uiStatusHandler))
-	r.HandleFunc("/ui/index.html", ms.uiStatusHandler)
+	r.HandleFunc("/", ms.proxyToLeader(requestIDMiddleware(ms.uiStatusHandler)))
+	r.HandleFunc("/ui/index.html", requestIDMiddleware(ms.uiStatusHandler))
 	if !ms.option.DisableHttp {
-		r.HandleFunc("/dir/assign", ms.proxyToLeader(ms.guard.WhiteList(ms.dirAssignHandler)))
-		r.HandleFunc("/dir/lookup", ms.guard.WhiteList(ms.dirLookupHandler))
-		r.HandleFunc("/dir/status", ms.proxyToLeader(ms.guard.WhiteList(ms.dirStatusHandler)))
-		r.HandleFunc("/col/delete", ms.proxyToLeader(ms.guard.WhiteList(ms.collectionDeleteHandler)))
-		r.HandleFunc("/vol/grow", ms.proxyToLeader(ms.guard.WhiteList(ms.volumeGrowHandler)))
-		r.HandleFunc("/vol/status", ms.proxyToLeader(ms.guard.WhiteList(ms.volumeStatusHandler)))
-		r.HandleFunc("/vol/vacuum", ms.proxyToLeader(ms.guard.WhiteList(ms.volumeVacuumHandler)))
-		r.HandleFunc("/submit", ms.guard.WhiteList(ms.submitFromMasterServerHandler))
-		r.HandleFunc("/collection/info", ms.guard.WhiteList(ms.collectionInfoHandler))
+		r.HandleFunc("/dir/assign", ms.proxyToLeader(ms.guard.WhiteList(requestIDMiddleware(ms.dirAssignHandler))))
+		r.HandleFunc("/dir/lookup", ms.guard.WhiteList(requestIDMiddleware(ms.dirLookupHandler)))
+		r.HandleFunc("/dir/status", ms.proxyToLeader(ms.guard.WhiteList(requestIDMiddleware(ms.dirStatusHandler))))
+		r.HandleFunc("/col/delete", ms.proxyToLeader(ms.guard.WhiteList(requestIDMiddleware(ms.collectionDeleteHandler))))
+		r.HandleFunc("/vol/grow", ms.proxyToLeader(ms.guard.WhiteList(requestIDMiddleware(ms.volumeGrowHandler))))
+		r.HandleFunc("/vol/status", ms.proxyToLeader(ms.guard.WhiteList(requestIDMiddleware(ms.volumeStatusHandler))))
+		r.HandleFunc("/vol/vacuum", ms.proxyToLeader(ms.guard.WhiteList(requestIDMiddleware(ms.volumeVacuumHandler))))
+		r.HandleFunc("/submit", ms.guard.WhiteList(requestIDMiddleware(ms.submitFromMasterServerHandler)))
+		r.HandleFunc("/collection/info", ms.guard.WhiteList(requestIDMiddleware(ms.collectionInfoHandler)))
 		/*
 			r.HandleFunc("/stats/health", ms.guard.WhiteList(statsHealthHandler))
 			r.HandleFunc("/stats/counter", ms.guard.WhiteList(statsCounterHandler))
 			r.HandleFunc("/stats/memory", ms.guard.WhiteList(statsMemoryHandler))
 		*/
-		r.HandleFunc("/{fileId}", ms.redirectHandler)
+		r.HandleFunc("/{fileId}", requestIDMiddleware(ms.redirectHandler))
 	}
 
 	ms.Topo.StartRefreshWritableVolumes(
diff --git a/weed/server/volume_server.go b/weed/server/volume_server.go
index a3f072eb7..5c5ebc49a 100644
--- a/weed/server/volume_server.go
+++ b/weed/server/volume_server.go
@@ -115,22 +115,22 @@ func NewVolumeServer(adminMux, publicMux *http.ServeMux, ip string,
 	vs.guard = security.NewGuard(whiteList, signingKey, expiresAfterSec, readSigningKey, readExpiresAfterSec)
 
 	handleStaticResources(adminMux)
-	adminMux.HandleFunc("/status", vs.statusHandler)
-	adminMux.HandleFunc("/healthz", vs.healthzHandler)
+	adminMux.HandleFunc("/status", requestIDMiddleware(vs.statusHandler))
+	adminMux.HandleFunc("/healthz", requestIDMiddleware(vs.healthzHandler))
 	if signingKey == "" || enableUiAccess {
 		// only expose the volume server details for safe environments
-		adminMux.HandleFunc("/ui/index.html", vs.uiStatusHandler)
+		adminMux.HandleFunc("/ui/index.html", requestIDMiddleware(vs.uiStatusHandler))
 		/*
 			adminMux.HandleFunc("/stats/counter", vs.guard.WhiteList(statsCounterHandler))
 			adminMux.HandleFunc("/stats/memory", vs.guard.WhiteList(statsMemoryHandler))
 			adminMux.HandleFunc("/stats/disk", vs.guard.WhiteList(vs.statsDiskHandler))
 		*/
 	}
-	adminMux.HandleFunc("/", vs.privateStoreHandler)
+	adminMux.HandleFunc("/", requestIDMiddleware(vs.privateStoreHandler))
 	if publicMux != adminMux {
 		// separated admin and public port
 		handleStaticResources(publicMux)
-		publicMux.HandleFunc("/", vs.publicReadOnlyHandler)
+		publicMux.HandleFunc("/", requestIDMiddleware(vs.publicReadOnlyHandler))
 	}
 
 	go vs.heartbeat()
diff --git a/weed/util/request_id.go b/weed/util/request_id.go
new file mode 100644
index 000000000..85ec254dc
--- /dev/null
+++ b/weed/util/request_id.go
@@ -0,0 +1,20 @@
+package util
+
+import "context"
+
+const (
+	RequestIdHttpHeader = "X-Request-ID"
+	RequestIDKey        = "x-request-id"
+)
+
+func GetRequestID(ctx context.Context) string {
+	if ctx == nil {
+		return ""
+	}
+	id, _ := ctx.Value(RequestIDKey).(string)
+	return id
+}
+
+func WithRequestID(ctx context.Context, id string) context.Context {
+	return context.WithValue(ctx, RequestIDKey, id)
+}