From 4b926bb1bd7bd46f6230aa332a41fe200015108b Mon Sep 17 00:00:00 2001
From: Vasiliy Tolstov <v.tolstov@unistack.org>
Date: Tue, 12 Mar 2024 01:58:03 +0300
Subject: [PATCH] fixup metadata

Signed-off-by: Vasiliy Tolstov <v.tolstov@unistack.org>
---
 requestid.go | 48 ++++++++++++++++++++++++++++++++++--------------
 1 file changed, 34 insertions(+), 14 deletions(-)

diff --git a/requestid.go b/requestid.go
index b6fdfb8..af3682c 100644
--- a/requestid.go
+++ b/requestid.go
@@ -31,31 +31,51 @@ var DefaultMetadataKey = textproto.CanonicalMIMEHeaderKey("x-request-id")
 // DefaultMetadataFunc wil be used if user not provide own func to fill metadata
 var DefaultMetadataFunc = func(ctx context.Context) (context.Context, error) {
 	var xid string
-	var err error
 
-	if _, ok := ctx.Value(XRequestIDKey{}).(string); !ok {
+	cid, cok := ctx.Value(XRequestIDKey{}).(string)
+	if cok && cid != "" {
+		xid = cid
+	}
+
+	imd, iok := metadata.FromIncomingContext(ctx)
+	if !iok || imd == nil {
+		imd = metadata.New(1)
+		ctx = metadata.NewIncomingContext(ctx, imd)
+	}
+
+	omd, ook := metadata.FromOutgoingContext(ctx)
+	if !ook || omd == nil {
+		omd = metadata.New(1)
+		ctx = metadata.NewOutgoingContext(ctx, omd)
+	}
+
+	if xid == "" {
+		var id string
+		if id, iok = imd.Get(DefaultMetadataKey); iok && id != "" {
+			xid = id
+		}
+		if id, ook = omd.Get(DefaultMetadataKey); ook && id != "" {
+			xid = id
+		}
+	}
+
+	if xid == "" {
+		var err error
 		xid, err = id.New()
 		if err != nil {
 			return ctx, err
 		}
+	}
+
+	if !cok {
 		ctx = context.WithValue(ctx, XRequestIDKey{}, xid)
 	}
 
-	imd, ok := metadata.FromIncomingContext(ctx)
-	if !ok {
-		imd = metadata.New(1)
-		imd.Set(DefaultMetadataKey, xid)
-		ctx = metadata.NewIncomingContext(ctx, imd)
-	} else if _, ok = imd.Get(DefaultMetadataKey); !ok {
+	if !iok {
 		imd.Set(DefaultMetadataKey, xid)
 	}
 
-	omd, ok := metadata.FromOutgoingContext(ctx)
-	if !ok {
-		omd = metadata.New(1)
-		omd.Set(DefaultMetadataKey, xid)
-		ctx = metadata.NewOutgoingContext(ctx, omd)
-	} else if _, ok = omd.Get(DefaultMetadataKey); !ok {
+	if !ook {
 		omd.Set(DefaultMetadataKey, xid)
 	}