fix x-ydb-trace-id

internal/conn/conn.go

@@ -329,7 +329,7 @@ func (c *conn) Invoke(
 	c.touchLastUsage()
 	defer c.touchLastUsage()
 
-	traceID, err := newTraceID()
+	ctx, traceID, err := meta.TraceID(ctx)
 	if err != nil {
 		return xerrors.WithStackTrace(err)
 	}
@@ -422,7 +422,7 @@ func (c *conn) NewStream(
 	c.touchLastUsage()
 	defer c.touchLastUsage()
 
-	traceID, err := newTraceID()
+	ctx, traceID, err := meta.TraceID(ctx)
 	if err != nil {
 		return nil, xerrors.WithStackTrace(err)
 	}

internal/meta/context.go

@@ -8,7 +8,17 @@ import (
 
 // WithTraceID returns a copy of parent context with traceID
 func WithTraceID(ctx context.Context, traceID string) context.Context {
-	return metadata.AppendToOutgoingContext(ctx, HeaderTraceID, traceID)
+	if md, has := metadata.FromOutgoingContext(ctx); !has || len(md[HeaderTraceID]) == 0 {
+		return metadata.AppendToOutgoingContext(ctx, HeaderTraceID, traceID)
+	}
+	return ctx
+}
+
+func traceID(ctx context.Context) (string, bool) {
+	if md, has := metadata.FromOutgoingContext(ctx); has && len(md[HeaderTraceID]) > 0 {
+		return md[HeaderTraceID][0], true
+	}
+	return "", false
 }
 
 // WithUserAgent returns a copy of parent context with custom user-agent info

internal/meta/trace_id.go

@@ -0,0 +1,30 @@
+package meta
+
+import (
+	"context"
+
+	"github.com/google/uuid"
+	"google.golang.org/grpc/metadata"
+
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
+)
+
+type newTraceIDOpts struct {
+	newRandom func() (uuid.UUID, error)
+}
+
+func TraceID(ctx context.Context, opts ...func(opts *newTraceIDOpts)) (context.Context, string, error) {
+	if id, has := traceID(ctx); has {
+		return ctx, id, nil
+	}
+	options := newTraceIDOpts{newRandom: uuid.NewRandom}
+	for _, opt := range opts {
+		opt(&options)
+	}
+	uuid, err := options.newRandom()
+	if err != nil {
+		return ctx, "", xerrors.WithStackTrace(err)
+	}
+	id := uuid.String()
+	return metadata.AppendToOutgoingContext(ctx, HeaderTraceID, id), id, nil
+}

internal/meta/trace_id_test.go

@@ -0,0 +1,57 @@
+package meta
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc/metadata"
+)
+
+func TestTraceID(t *testing.T) {
+	t.Run("TraceID from rand", func(t *testing.T) {
+		ctx, id, err := TraceID(
+			context.Background(),
+			func(opts *newTraceIDOpts) {
+				opts.newRandom = func() (uuid.UUID, error) {
+					return uuid.UUID{}, nil
+				}
+			},
+		)
+		require.NoError(t, err)
+		require.Equal(t, "00000000-0000-0000-0000-000000000000", id)
+		md, has := metadata.FromOutgoingContext(ctx)
+		require.True(t, has)
+		require.Equal(t, 1, len(md[HeaderTraceID]))
+		require.Equal(t, id, md[HeaderTraceID][0])
+	})
+	t.Run("TraceID from rand failed", func(t *testing.T) {
+		_, _, err := TraceID(
+			context.Background(),
+			func(opts *newTraceIDOpts) {
+				opts.newRandom = func() (uuid.UUID, error) {
+					return uuid.UUID{}, errors.New("")
+				}
+			},
+		)
+		require.Error(t, err)
+	})
+	t.Run("TraceID from outgoing metadata", func(t *testing.T) {
+		ctx, id, err := TraceID(
+			WithTraceID(context.Background(), "{test}"),
+			func(opts *newTraceIDOpts) {
+				opts.newRandom = func() (uuid.UUID, error) {
+					return uuid.UUID{}, errors.New("")
+				}
+			},
+		)
+		require.NoError(t, err)
+		require.Equal(t, "{test}", id)
+		md, has := metadata.FromOutgoingContext(ctx)
+		require.True(t, has)
+		require.Equal(t, 1, len(md[HeaderTraceID]))
+		require.Equal(t, id, md[HeaderTraceID][0])
+	})
+}