feat(Spanner.V1): Add Spanner request ID header

robertvoinescu-work · robertvoinescu-work · commit 626fec3c229e · 2025-11-11T20:27:45.000Z
Adds the  class to generate a unique ID for each RPC. This also modifies the  to inject this ID into the  header for all Spanner RPCs. Includes comprehensive tests for ID generation and header injection.
diff --git a/apis/Google.Cloud.Spanner.V1/Google.Cloud.Spanner.V1.Tests/SpannerRequestIdTests.cs b/apis/Google.Cloud.Spanner.V1/Google.Cloud.Spanner.V1.Tests/SpannerRequestIdTests.cs
@@ -0,0 +1,228 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+using Grpc.Core;
+using Google.Api.Gax.Grpc;
+using System;
+using System.Linq;
+using System.Threading.Tasks;
+using Xunit;
+using System.Collections.Generic;
+
+namespace Google.Cloud.Spanner.V1.Tests;
+
+public class SpannerRequestIdTests
+{
+    private const string SampleDatabaseName = "projects/proj/instances/inst/databases/db";
+    private const string SampleSessionName = "projects/proj/instances/inst/databases/db/sessions/sess";
+
+    // Tests for SpannerRequestId class itself
+    [Fact]
+    public void SpannerRequestId_Format()
+    {
+        var expectedClientId = 123;
+        var expectedRequestSequence = 456;
+        var requestId = new SpannerRequestId(expectedClientId, expectedRequestSequence);
+
+        // The expected format is 6 parts that break down as such:
+        // {VersionId}.{s_processId}.{ClientId}.{ChannelId}.{RequestSequence}.{AttemptNum}
+        var idString = requestId.IncrementAttempt();
+        var parts = idString.Split('.');
+
+        Assert.Equal(6, parts.Length);
+        var versionId = parts[0];
+        var processIdString = parts[1];
+        var clientId = parts[2];
+        var channelId = parts[3];
+        var requestSequence = parts[4];
+        var attemptNum = parts[5];
+
+        Assert.Equal("1", versionId);
+        Assert.Equal(expectedClientId.ToString(), clientId);
+        Assert.Equal("0", channelId);
+        Assert.Equal(expectedRequestSequence.ToString(), requestSequence);
+        Assert.Equal("1", attemptNum);
+        Assert.True(ulong.TryParse(processIdString, out _));
+    }
+
+    [Fact]
+    public void SpannerRequestId_IncrementAttemptNum()
+    {
+        var requestId = new SpannerRequestId(1, 1);
+        for (int expectedAttemptNum = 1; expectedAttemptNum < 5; expectedAttemptNum++)
+        {
+            var parts = requestId.IncrementAttempt().Split('.');
+            Assert.Equal(6, parts.Length);
+            var actualAttemptNum = parts[5];
+            Assert.Equal($"{expectedAttemptNum}", actualAttemptNum);
+        }
+    }
+
+    [Fact]
+    public void SpannerRequestId_ProcessIdIsStatic()
+    {
+        var requestId1 = new SpannerRequestId(1, 1);
+        var requestId2 = new SpannerRequestId(2, 2);
+
+        // The process ID is part of the ToString() output, so we can compare the generated strings
+        // and extract the process ID.
+        string requestId1String = requestId1.IncrementAttempt();
+        string requestId2String = requestId2.IncrementAttempt();
+
+        string processId1 = requestId1String.Split('.')[1];
+        string processId2 = requestId2String.Split('.')[1];
+
+        Assert.Equal(processId1, processId2);
+    }
+
+    // Tests for header injection
+    public static TheoryData<Action<SpannerClient>> SpannerClientActions { get; } = new TheoryData<Action<SpannerClient>>
+    {
+        client => client.CreateSession(new CreateSessionRequest { Database = SampleDatabaseName }),
+        client => client.BatchCreateSessions(new BatchCreateSessionsRequest { Database = SampleDatabaseName }),
+        client => client.GetSession(new GetSessionRequest { Name = SampleSessionName }),
+        client => client.ListSessions(new ListSessionsRequest { Database = SampleDatabaseName }).AsRawResponses().First(),
+        client => client.DeleteSession(new DeleteSessionRequest { Name = SampleSessionName }),
+        client => client.ExecuteSql(new ExecuteSqlRequest { Session = SampleSessionName }),
+        client => client.ExecuteBatchDml(new ExecuteBatchDmlRequest { Session = SampleSessionName }),
+        client => client.Read(new ReadRequest { Session = SampleSessionName }),
+        client => client.BeginTransaction(new BeginTransactionRequest { Session = SampleSessionName }),
+        client => client.Commit(new CommitRequest { Session = SampleSessionName }),
+        client => client.Rollback(new RollbackRequest { Session = SampleSessionName }),
+        client => client.PartitionQuery(new PartitionQueryRequest { Session = SampleSessionName }),
+        client => client.PartitionRead(new PartitionReadRequest { Session = SampleSessionName })
+    };
+
+    [Theory]
+    [MemberData(nameof(SpannerClientActions))]
+    public void SetsHeaderOnRpcCalls(Action<SpannerClient> action)
+    {
+        var invoker = new RetryFakeCallInvoker();
+        var client = new SpannerClientBuilder { CallInvoker = invoker }.Build();
+        action(client);
+        Metadata.Entry entry = Assert.Single(invoker.CapturedMetadata[0], e => e.Key == SpannerRequestId.HeaderKey);
+        Assert.NotNull(entry.Value);
+    }
+
+    [Fact]
+    public void IncrementsRequestIdOnRetry()
+    {
+        var invoker = new RetryFakeCallInvoker(failCount: 1);
+        var settings = new SpannerSettings
+        {
+            // Configure the CreateSession call to retry on Unavailable errors.
+            CallSettings = CallSettings.FromRetry(RetrySettings.FromExponentialBackoff(
+                maxAttempts: 3,
+                initialBackoff: TimeSpan.FromMilliseconds(1),
+                maxBackoff: TimeSpan.FromMilliseconds(1),
+                backoffMultiplier: 1.0,
+                retryFilter: RetrySettings.FilterForStatusCodes(StatusCode.Unavailable)))
+        };
+        var client = new SpannerClientBuilder { CallInvoker = invoker, Settings = settings }.Build();
+
+        client.CreateSession(new CreateSessionRequest { Database = SampleDatabaseName });
+        client.CreateSession(new CreateSessionRequest { Database = SampleDatabaseName });
+        client.CreateSession(new CreateSessionRequest { Database = SampleDatabaseName });
+
+        // Assert that the invoker was called four times for the three client calls.
+        // The first call should have failed the first time and succeeded on retry.
+        Assert.Equal(4, invoker.CapturedMetadata.Count);
+
+        var requestIds = invoker.CapturedMetadata
+            .Select(m => m.Single(e => e.Key == SpannerRequestId.HeaderKey).Value)
+            .Select(id => new SpannerRequestIdParts(id))
+            .ToList();
+
+        Assert.Equal((1, 1), (requestIds[0].RequestSequence, requestIds[0].AttemptNum));
+        Assert.Equal((1, 2), (requestIds[1].RequestSequence, requestIds[1].AttemptNum));
+        Assert.Equal((2, 1), (requestIds[2].RequestSequence, requestIds[2].AttemptNum));
+        Assert.Equal((3, 1), (requestIds[3].RequestSequence, requestIds[3].AttemptNum));
+    }
+
+    /// <summary>
+    /// A fake call invoker that fails the first time it's called with an "Unavailable" status,
+    /// and succeeds on subsequent calls. It captures the metadata for every call.
+    /// </summary>
+    private class RetryFakeCallInvoker : CallInvoker
+    {
+        private int _callCount = 0;
+        private readonly int _failCount;
+        public List<Metadata> CapturedMetadata { get; } = new List<Metadata>();
+
+        public RetryFakeCallInvoker(int failCount = 0) => _failCount = failCount;
+
+        public override TResponse BlockingUnaryCall<TRequest, TResponse>(Method<TRequest, TResponse> method, string host, CallOptions options, TRequest request)
+        {
+            CapturedMetadata.Add(options.Headers);
+            _callCount++;
+
+            if (_callCount <= _failCount)
+            {
+                // Fail the first time with a retryable error.
+                throw new RpcException(new Status(StatusCode.Unavailable, "Transient error"));
+            }
+
+            // Succeed on the second attempt.
+            return (TResponse)Activator.CreateInstance(typeof(TResponse));
+        }
+
+        public override AsyncClientStreamingCall<TRequest, TResponse> AsyncClientStreamingCall<TRequest, TResponse>(Method<TRequest, TResponse> method, string host, CallOptions options) =>
+            throw new NotImplementedException();
+
+        public override AsyncDuplexStreamingCall<TRequest, TResponse> AsyncDuplexStreamingCall<TRequest, TResponse>(Method<TRequest, TResponse> method, string host, CallOptions options) =>
+            throw new NotImplementedException();
+
+        public override AsyncServerStreamingCall<TResponse> AsyncServerStreamingCall<TRequest, TResponse>(Method<TRequest, TResponse> method, string host, CallOptions options, TRequest request) =>
+            throw new NotImplementedException();
+
+        public override AsyncUnaryCall<TResponse> AsyncUnaryCall<TRequest, TResponse>(Method<TRequest, TResponse> method, string host, CallOptions options, TRequest request)
+        {
+            CapturedMetadata.Add(options.Headers);
+            _callCount++;
+
+            if (_callCount <= _failCount)
+            {
+                // Fail the first time with a retryable error.
+                throw new RpcException(new Status(StatusCode.Unavailable, "Transient error"));
+            }
+
+            // Succeed on the second attempt.
+            var response = (TResponse)Activator.CreateInstance(typeof(TResponse));
+            return new AsyncUnaryCall<TResponse>(Task.FromResult(response), Task.FromResult(new Metadata()), () => Status.DefaultSuccess, () => new Metadata(), () => { });
+        }
+    }
+
+    private struct SpannerRequestIdParts
+    {
+        public int VersionId { get; }
+        public ulong ProcessId { get; }
+        public int ClientId { get; }
+        public int ChannelId { get; }
+        public int RequestSequence { get; }
+        public int AttemptNum { get; }
+
+        public SpannerRequestIdParts(string requestId)
+        {
+            var parts = requestId.Split('.');
+            Assert.Equal(6, parts.Length);
+
+            VersionId = int.Parse(parts[0]);
+            ProcessId = ulong.Parse(parts[1]);
+            ClientId = int.Parse(parts[2]);
+            ChannelId = int.Parse(parts[3]);
+            RequestSequence = int.Parse(parts[4]);
+            AttemptNum = int.Parse(parts[5]);
+        }
+    }
+}
diff --git a/apis/Google.Cloud.Spanner.V1/Google.Cloud.Spanner.V1/SpannerClientPartial.cs b/apis/Google.Cloud.Spanner.V1/Google.Cloud.Spanner.V1/SpannerClientPartial.cs
@@ -14,6 +14,7 @@
 
 using Google.Api.Gax.Grpc;
 using Google.Cloud.Spanner.Common.V1;
+using System.Threading;
 
 namespace Google.Cloud.Spanner.V1
 {
@@ -48,6 +49,25 @@ internal void MaybeApplyRouteToLeaderHeader(ref CallSettings settings)
 
     public partial class SpannerClientImpl
     {
+        /// <summary>
+        /// A counter incremented on the instation of each new Spanner client
+        /// </summary>
+        private static int s_clientIdCounter;
+
+        /// <summary>
+        /// Creates the request ID associated with each RPC.
+        /// </summary>
+        private SpannerRequestId.Generator _requestIdGenerator;
+
+        /// <summary>
+        /// The process ID, assigned to each outgoing request in the x-goog-spanner-request-id header.
+        /// </summary>
+        public string ProcessId
+        {
+            get => SpannerRequestId.s_processId;
+            set => SpannerRequestId.s_processId = value;
+        }
+
         /// <summary>
         /// The name of the header used for efficiently routing requests.
         /// </summary>
@@ -64,78 +84,100 @@ public partial class SpannerClientImpl
         partial void OnConstruction(Spanner.SpannerClient grpcClient, SpannerSettings effectiveSettings, ClientHelper clientHelper)
         {
             Settings = effectiveSettings;
+            _requestIdGenerator = new SpannerRequestId.Generator(Interlocked.Increment(ref s_clientIdCounter));
         }
 
         partial void Modify_CreateSessionRequest(ref CreateSessionRequest request, ref CallSettings settings)
         {
             ApplyResourcePrefixHeaderFromDatabase(ref settings, request.Database);
             MaybeApplyRouteToLeaderHeader(ref settings);
+            ApplyRequestIdToHeader(ref settings);
         }
 
         partial void Modify_BatchCreateSessionsRequest(ref BatchCreateSessionsRequest request, ref CallSettings settings)
         {
             ApplyResourcePrefixHeaderFromDatabase(ref settings, request.Database);
             MaybeApplyRouteToLeaderHeader(ref settings);
+            ApplyRequestIdToHeader(ref settings);
         }
 
         partial void Modify_GetSessionRequest(ref GetSessionRequest request, ref CallSettings settings)
         {
             ApplyResourcePrefixHeaderFromSession(ref settings, request.Name);
             MaybeApplyRouteToLeaderHeader(ref settings);
+            ApplyRequestIdToHeader(ref settings);
         }
 
-        partial void Modify_ListSessionsRequest(ref ListSessionsRequest request, ref CallSettings settings) =>
+        partial void Modify_ListSessionsRequest(ref ListSessionsRequest request, ref CallSettings settings)
+        {
             // This operation is never routed to leader so we don't call MaybeApplyRouteToLeaderHeader.
             ApplyResourcePrefixHeaderFromDatabase(ref settings, request.Database);
+            ApplyRequestIdToHeader(ref settings);
+        }
 
-        partial void Modify_DeleteSessionRequest(ref DeleteSessionRequest request, ref CallSettings settings) =>
+        partial void Modify_DeleteSessionRequest(ref DeleteSessionRequest request, ref CallSettings settings)
+        {
             // This operation is never routed to leader so we don't call MaybeApplyRouteToLeaderHeader.
             ApplyResourcePrefixHeaderFromSession(ref settings, request.Name);
+            ApplyRequestIdToHeader(ref settings);
+        }
 
-        partial void Modify_ExecuteSqlRequest(ref ExecuteSqlRequest request, ref CallSettings settings) =>
+        partial void Modify_ExecuteSqlRequest(ref ExecuteSqlRequest request, ref CallSettings settings)
+        {
             // This operations is routed to leader only if the transaction it uses is of a certain type.
             // We don't have that information here so the leader routing header needs to be applied elsewhere.
             ApplyResourcePrefixHeaderFromSession(ref settings, request.Session);
+            ApplyRequestIdToHeader(ref settings);
+        }
 
         partial void Modify_ExecuteBatchDmlRequest(ref ExecuteBatchDmlRequest request, ref CallSettings settings)
         {
             ApplyResourcePrefixHeaderFromSession(ref settings, request.Session);
             MaybeApplyRouteToLeaderHeader(ref settings);
+            ApplyRequestIdToHeader(ref settings);
         }
 
-        partial void Modify_ReadRequest(ref ReadRequest request, ref CallSettings settings) =>
+        partial void Modify_ReadRequest(ref ReadRequest request, ref CallSettings settings)
+        {
             // This operations is routed to leader only if the transaction it uses is of a certain type.
             // We don't have that information here so the leader routing header needs to be applied elsewhere.
             ApplyResourcePrefixHeaderFromSession(ref settings, request.Session);
+            ApplyRequestIdToHeader(ref settings);
+        }
 
         partial void Modify_BeginTransactionRequest(ref BeginTransactionRequest request, ref CallSettings settings)
         {
             ApplyResourcePrefixHeaderFromSession(ref settings, request.Session);
             MaybeApplyRouteToLeaderHeader(ref settings, request.Options?.ModeCase ?? TransactionOptions.ModeOneofCase.None);
+            ApplyRequestIdToHeader(ref settings);
         }
 
         partial void Modify_CommitRequest(ref CommitRequest request, ref CallSettings settings)
         {
             ApplyResourcePrefixHeaderFromSession(ref settings, request.Session);
             MaybeApplyRouteToLeaderHeader(ref settings);
+            ApplyRequestIdToHeader(ref settings);
         }
 
         partial void Modify_RollbackRequest(ref RollbackRequest request, ref CallSettings settings)
         {
             ApplyResourcePrefixHeaderFromSession(ref settings, request.Session);
             MaybeApplyRouteToLeaderHeader(ref settings);
+            ApplyRequestIdToHeader(ref settings);
         }
 
         partial void Modify_PartitionQueryRequest(ref PartitionQueryRequest request, ref CallSettings settings)
         {
             ApplyResourcePrefixHeaderFromSession(ref settings, request.Session);
             MaybeApplyRouteToLeaderHeader(ref settings);
+            ApplyRequestIdToHeader(ref settings);
         }
 
         partial void Modify_PartitionReadRequest(ref PartitionReadRequest request, ref CallSettings settings)
         {
             ApplyResourcePrefixHeaderFromSession(ref settings, request.Session);
             MaybeApplyRouteToLeaderHeader(ref settings);
+            ApplyRequestIdToHeader(ref settings);
         }
 
         internal static void ApplyResourcePrefixHeaderFromDatabase(ref CallSettings settings, string resource)
@@ -167,5 +209,15 @@ internal static void ApplyResourcePrefixHeaderFromSession(ref CallSettings setti
                 settings = settings.WithHeader(ResourcePrefixHeader, database.ToString());
             }
         }
+
+        internal void ApplyRequestIdToHeader(ref CallSettings settings)
+        {
+            var requestId = _requestIdGenerator.Generate();
+
+            // The header mutation increments the attempt each time the header is populated which happens right before a
+            // call. This allows us to increment the attempt num and assign a unique request id in the case of retries
+            var newSettings = CallSettings.FromHeaderMutation(metadata => metadata.Add(SpannerRequestId.HeaderKey, requestId.IncrementAttempt()));
+            settings = settings.MergedWith(newSettings);
+        }
     }
 }
diff --git a/apis/Google.Cloud.Spanner.V1/Google.Cloud.Spanner.V1/SpannerRequestId.cs b/apis/Google.Cloud.Spanner.V1/Google.Cloud.Spanner.V1/SpannerRequestId.cs
