feat: Add invocation ID and attempt count in x-goog-api-client header

Fixes #8881

apis/Google.Cloud.Storage.V1/Google.Cloud.Storage.V1.Tests/FakeStorageService.cs

@@ -1,4 +1,4 @@
-// Copyright 2022 Google LLC
+// Copyright 2022 Google LLC
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -20,20 +20,27 @@
 
 namespace Google.Cloud.Storage.V1.Tests
 {
-    class FakeStorageService : StorageService
+    internal class FakeStorageService : StorageService
     {
         private readonly ReplayingMessageHandler handler;
 
-        public FakeStorageService() : base(new Initializer
+        public FakeStorageService() : this(new ReplayingMessageHandler())
         {
-            HttpClientFactory = new FakeHttpClientFactory(new ReplayingMessageHandler()),
-            ApplicationName = "Fake",
-            GZipEnabled = false
-        })
+        }
+
+        public FakeStorageService(ReplayingMessageHandler handler) : base(CreateInitializer(handler))
         {
-            handler = (ReplayingMessageHandler)HttpClient.MessageHandler.InnerHandler;
+            this.handler = handler;
         }
 
+        private static Initializer CreateInitializer(ReplayingMessageHandler handler) =>
+            new Initializer
+            {
+                HttpClientFactory = new FakeHttpClientFactory(handler),
+                ApplicationName = "Fake",
+                GZipEnabled = false
+            };
+
         public void ExpectRequest<TResponse>(ClientServiceRequest<TResponse> request, TResponse response)
         {
             var httpRequest = request.CreateRequest();

apis/Google.Cloud.Storage.V1/Google.Cloud.Storage.V1.Tests/StorageClientImplRetryTest.cs

@@ -16,7 +16,9 @@
 using Google.Apis.Storage.v1;
 using Google.Apis.Storage.v1.Data;
 using Google.Apis.Util;
+using Google.Cloud.ClientTesting;
 using System;
+using System.Collections.Generic;
 using System.Linq;
 using System.Net;
 using Xunit;
@@ -430,6 +432,47 @@ public class StorageClientImplRetryTest
 
         #endregion
 
+        #region Invocation ID test cases
+        [Fact]
+        public void InvocationIdAndCountAreSet()
+        {
+            var replayingMessageHandler = new ReplayingMessageHandler(VersionHeaderBuilder.HeaderName);
+            var service = new FakeStorageService(replayingMessageHandler);
+            service.HttpClient.MessageHandler.GoogleApiClientHeader = "test/fake";
+
+            var request = service.Buckets.Get("bucket");
+            service.ExpectRequest(request, HttpStatusCode.BadGateway);
+            service.ExpectRequest(request, HttpStatusCode.BadGateway);
+            service.ExpectRequest(request, new Bucket());
+
+            var client = new StorageClientImpl(service);
+            client.GetBucket("bucket");
+            service.Verify();
+
+            var actualHeaders = replayingMessageHandler.CapturedHeaders;
+            Assert.Equal(3, actualHeaders.Count);
+
+            string invocationIdName = RetryHandler.InvocationIdHeaderPart;
+            string attemptCountName = RetryHandler.AttemptCountHeaderPart;
+
+            var request1Parts = ConvertHeader(actualHeaders[0]);
+            Assert.Equal("fake", request1Parts["test"]);
+            Assert.Equal("1", request1Parts[attemptCountName]);
+            var guid = request1Parts[invocationIdName];
+            // Just validate that it's a real GUID...
+            Guid.Parse(guid);
+
+            var request2Parts = ConvertHeader(actualHeaders[1]);
+            Assert.Equal("2", request2Parts[attemptCountName]);
+            Assert.Equal(guid, request2Parts[invocationIdName]);
+
+            var request3Parts = ConvertHeader(actualHeaders[2]);
+            Assert.Equal("3", request3Parts[attemptCountName]);
+            Assert.Equal(guid, request3Parts[invocationIdName]);
+
+            Dictionary<string, string> ConvertHeader(string header) =>
+                header.Split(' ').ToDictionary(piece => piece.Split('/')[0], piece => piece.Split('/')[1]);
+        }
+        #endregion
     }
 }
-

apis/Google.Cloud.Storage.V1/Google.Cloud.Storage.V1/RetryHandler.cs

@@ -1,4 +1,4 @@
-// Copyright 2022 Google Inc. All Rights Reserved.
+// Copyright 2022 Google Inc. All Rights Reserved.
 // 
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -12,11 +12,16 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+using Google.Api.Gax;
 using Google.Apis.Http;
 using Google.Apis.Storage.v1;
 using System;
+using System.Collections.Generic;
+using System.Globalization;
 using System.Linq;
 using System.Net.Http;
+using System.Net.NetworkInformation;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace Google.Cloud.Storage.V1
@@ -28,6 +33,10 @@ namespace Google.Cloud.Storage.V1
     /// </summary>
     internal sealed class RetryHandler : IHttpUnsuccessfulResponseHandler
     {
+        // For testing
+        internal const string InvocationIdHeaderPart = "gccl-invocation-id";
+        internal const string AttemptCountHeaderPart = "gccl-attempt-count";
+
         private static readonly int[] s_retriableErrorCodes =
         {
                 408, // Request timeout
@@ -37,12 +46,17 @@ internal sealed class RetryHandler : IHttpUnsuccessfulResponseHandler
                 503, // Service unavailable
                 504 // Gateway timeout
         };
-        private static RetryHandler s_instance = new RetryHandler();
+        private static readonly RetryHandler s_instance = new RetryHandler();
 
         private RetryHandler() { }
 
-        internal static void MarkAsRetriable<TResponse>(StorageBaseServiceRequest<TResponse> request) =>
+        internal static void MarkAsRetriable<TResponse>(StorageBaseServiceRequest<TResponse> request)
+        {
+            // Note: we can't use ModifyRequest, as the x-goog-api-client header is added later by ConfigurableMessageHandler.
+            // Additionally, that's only called once, and we may want to record the attempt number as well.
+            request.AddExecuteInterceptor(InvocationIdInterceptor.Instance);
             request.AddUnsuccessfulResponseHandler(s_instance);
+        }
 
         // This function is designed to support asynchrony in case we need to examine the response content, but for now we only need the status code
         internal static Task<bool> IsRetriableResponse(HttpResponseMessage response) => 
@@ -65,5 +79,69 @@ public async Task<bool> HandleResponseAsync(HandleUnsuccessfulResponseArgs args)
             await Task.Delay(delay, args.CancellationToken).ConfigureAwait(false);
             return true;
         }
+
+        /// <summary>
+        /// Interceptor which adds a random invocation ID within the x-goog-api-client header,
+        /// along with an attempt count.
+        /// </summary>
+        private sealed class InvocationIdInterceptor : IHttpExecuteInterceptor
+        {
+            internal static InvocationIdInterceptor Instance { get; } = new InvocationIdInterceptor();
+
+            private const string InvocationIdPrefix = InvocationIdHeaderPart + "/";
+            private const string AttemptCountPrefix = AttemptCountHeaderPart + "/";
+
+            private InvocationIdInterceptor()
+            {
+            }
+
+            public Task InterceptAsync(HttpRequestMessage request, CancellationToken cancellationToken)
+            {
+                // If we don't have the header already, or if there isn't a single value,
+                // that's an odd situation: don't add one with just the invocation ID.
+                if (!request.Headers.TryGetValues(VersionHeaderBuilder.HeaderName, out var values) || values.Count() != 1)
+                {
+                    return Task.CompletedTask;
+                }
+                string value = values.Single();
+                List<string> parts = value.Split(' ').ToList();
+
+                bool gotInvocationId = false;
+                bool gotAttemptCount = false;
+                for (int i = 0; i < parts.Count; i++)
+                {
+                    if (parts[i].StartsWith(InvocationIdPrefix, StringComparison.Ordinal))
+                    {
+                        gotInvocationId = true;
+                    }
+                    else if (parts[i].StartsWith(AttemptCountPrefix, StringComparison.Ordinal))
+                    {
+                        gotAttemptCount = true;
+                        string countText = parts[i].Substring(AttemptCountPrefix.Length);
+                        if (int.TryParse(countText, NumberStyles.None, CultureInfo.InvariantCulture, out int count))
+                        {
+                            count++;
+                            parts[i] = AttemptCountPrefix + count.ToString(CultureInfo.InvariantCulture);
+                        }
+                    }
+                }
+                if (!gotInvocationId)
+                {
+                    parts.Add(InvocationIdPrefix + Guid.NewGuid());
+                }
+                if (!gotAttemptCount)
+                {
+                    // TODO: Check this: should we add it on the first request,
+                    // or only subsequent requests? Design doc is unclear.
+                    parts.Add(AttemptCountPrefix + "1");
+                }
+
+                request.Headers.Remove(VersionHeaderBuilder.HeaderName);
+                request.Headers.Add(VersionHeaderBuilder.HeaderName, string.Join(" ", parts));
+
+
+                return Task.CompletedTask;
+            }
+        }
     }
 }

tools/Google.Cloud.ClientTesting/ReplayingMessageHandler.cs

@@ -1,4 +1,4 @@
-// Copyright 2016 Google Inc. All Rights Reserved.
+// Copyright 2016 Google Inc. All Rights Reserved.
 // 
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -30,13 +30,48 @@ namespace Google.Cloud.ClientTesting
     /// </summary>
     public class ReplayingMessageHandler : HttpMessageHandler
     {
+        private readonly string _headerToCapture;
+
         private readonly Queue<Tuple<Uri, string, HttpResponseMessage>> _requestResponses =
             new Queue<Tuple<Uri, string, HttpResponseMessage>>();
 
+        /// <summary>
+        /// The captured headers, or null if headers are not being captured.
+        /// There is one element per request, with an element value of null if the header is not present for the corresponding request.
+        /// </summary>
+        public List<string> CapturedHeaders { get; }
+
+        /// <summary>
+        /// Creates a handler that doesn't capture any headers
+        /// </summary>
+        public ReplayingMessageHandler()
+        {
+        }
+
+        /// <summary>
+        /// Creates a handler that captures the given header in <see cref="CapturedHeaders"/>,
+        /// once per request.
+        /// </summary>
+        public ReplayingMessageHandler(string header)
+        {
+            _headerToCapture = header;
+            CapturedHeaders = new List<string>();
+        }
+
         protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
         {
             Assert.NotEmpty(_requestResponses);
-
+            if (_headerToCapture is string header)
+            {
+                if (request.Headers.TryGetValues(header, out var values))
+                {
+                    CapturedHeaders.Add(string.Join(",", values));
+                }
+                else
+                {
+                    CapturedHeaders.Add(null);
+                }
+            }
             var requestResponse = _requestResponses.Dequeue();
             Uri expectedRequestUri = requestResponse.Item1;
             string expectedRequestContent = requestResponse.Item2;