[Internal] [FaultInjection] ThinClient: Fixes payload of injected errors for requests going to ThinClient Proxy

Microsoft.Azure.Cosmos/FaultInjection/src/implementation/FaultInjectionRuleProcessor.cs

@@ -189,7 +189,8 @@ await BackoffRetryUtility<IEnumerable<string>>.ExecuteAsync(
                     result.GetDelay(),
                     result.GetSuppressServiceRequests(),
                     result.GetInjectionRate(),
-                    this.applicationContext));
+                    this.applicationContext, 
+                    this.globalEndpointManager));
         }
 
         private async Task<IFaultInjectionRuleInternal> GetEffectiveConnectionErrorRule(FaultInjectionRule rule)

Microsoft.Azure.Cosmos/src/DocumentClient.cs

@@ -1088,7 +1088,8 @@ private async Task<bool> GetInitializationTaskAsync(IStoreClientFactory storeCli
                 httpClient: this.httpClient,
                 globalPartitionEndpointManager: this.PartitionKeyRangeLocation,
                 isThinClientEnabled: this.isThinClientEnabled,
-                userAgentContainer: this.ConnectionPolicy.UserAgentContainer);
+                userAgentContainer: this.ConnectionPolicy.UserAgentContainer,
+                this.chaosInterceptor);
 
             this.GatewayStoreModel = gatewayStoreModel;
 

Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs

@@ -17,7 +17,8 @@ namespace Microsoft.Azure.Cosmos
     using Microsoft.Azure.Cosmos.Routing;
     using Microsoft.Azure.Cosmos.Tracing;
     using Microsoft.Azure.Documents;
-    using Microsoft.Azure.Documents.Collections;
+    using Microsoft.Azure.Documents.Collections;
+    using Microsoft.Azure.Documents.FaultInjection;
     using Newtonsoft.Json;
 
     // Marking it as non-sealed in order to unit test it using Moq framework
@@ -26,7 +27,8 @@ internal class GatewayStoreModel : IStoreModelExtension, IDisposable
         private static readonly string sessionConsistencyAsString = ConsistencyLevel.Session.ToString();
         private readonly GlobalPartitionEndpointManager globalPartitionEndpointManager;
         private readonly ISessionContainer sessionContainer;
-        private readonly DocumentClientEventSource eventSource;
+        private readonly DocumentClientEventSource eventSource;
+        private readonly IChaosInterceptor chaosInterceptor;
 
         internal readonly GlobalEndpointManager endpointManager;
         internal readonly ConsistencyLevel defaultConsistencyLevel;
@@ -48,13 +50,15 @@ public GatewayStoreModel(
              CosmosHttpClient httpClient,
              GlobalPartitionEndpointManager globalPartitionEndpointManager,
              bool isThinClientEnabled,
-             UserAgentContainer userAgentContainer = null)
+             UserAgentContainer userAgentContainer = null,
+             IChaosInterceptor chaosInterceptor = null)
         {
             this.endpointManager = endpointManager;
             this.sessionContainer = sessionContainer;
             this.defaultConsistencyLevel = defaultConsistencyLevel;
             this.eventSource = eventSource;
-            this.globalPartitionEndpointManager = globalPartitionEndpointManager;
+            this.globalPartitionEndpointManager = globalPartitionEndpointManager;
+            this.chaosInterceptor = chaosInterceptor;
             this.gatewayStoreClient = new GatewayStoreClient(
                 httpClient,
                 this.eventSource,
@@ -68,7 +72,8 @@ public GatewayStoreModel(
                     userAgentContainer,
                     this.eventSource,
                     globalPartitionEndpointManager,
-                    serializerSettings);
+                    serializerSettings,
+                    this.chaosInterceptor);
             }
 
             this.globalPartitionEndpointManager.SetBackgroundConnectionPeriodicRefreshTask(

Microsoft.Azure.Cosmos/src/HttpClient/CosmosHttpClientCore.cs

@@ -359,7 +359,6 @@ private async Task<HttpResponseMessage> SendHttpHelperAsync(
                     {
                         if (this.chaosInterceptor != null && documentServiceRequest != null)
                         {
-                            this.SetFaultInjectionHeader(documentServiceRequest, requestMessage);
                             (bool hasFault, HttpResponseMessage fiResponseMessage) = await this.InjectFaultsAsync(cancellationTokenSource, documentServiceRequest, requestMessage);
                             if (hasFault)
                             {
@@ -374,7 +373,6 @@ private async Task<HttpResponseMessage> SendHttpHelperAsync(
 
                         if (this.chaosInterceptor != null && documentServiceRequest != null)
                         {
-                            this.SetFaultInjectionHeader(documentServiceRequest, requestMessage);
                             CancellationToken fiToken = cancellationTokenSource.Token;
                             fiToken.ThrowIfCancellationRequested();
                             await this.chaosInterceptor.OnAfterHttpSendAsync(documentServiceRequest, fiToken);
@@ -469,11 +467,6 @@ private async Task<HttpResponseMessage> SendHttpHelperAsync(
             }
         }
 
-        private void SetFaultInjectionHeader(DocumentServiceRequest documentServiceRequest, HttpRequestMessage requestMessage)
-        {
-            documentServiceRequest.Headers.Set("FAULTINJECTION_GW_URI", requestMessage.RequestUri.ToString());
-        }
-
         private async Task<(bool, HttpResponseMessage)> InjectFaultsAsync(
             CancellationTokenSource cancellationTokenSource, 
             DocumentServiceRequest documentServiceRequest, 

Microsoft.Azure.Cosmos/src/ThinClientStoreClient.cs

@@ -14,6 +14,7 @@ namespace Microsoft.Azure.Cosmos
     using Microsoft.Azure.Cosmos.Core.Trace;
     using Microsoft.Azure.Cosmos.Routing;
     using Microsoft.Azure.Documents;
+    using Microsoft.Azure.Documents.FaultInjection;
     using Newtonsoft.Json;
     using static Microsoft.Azure.Cosmos.ThinClientTransportSerializer;
 
@@ -26,13 +27,15 @@ internal class ThinClientStoreClient : GatewayStoreClient
         private readonly GlobalPartitionEndpointManager globalPartitionEndpointManager;
         private readonly ObjectPool<BufferProviderWrapper> bufferProviderWrapperPool;
         private readonly UserAgentContainer userAgentContainer;
+        private readonly IChaosInterceptor chaosInterceptor;
 
         public ThinClientStoreClient(
             CosmosHttpClient httpClient,
             UserAgentContainer userAgentContainer,
             ICommunicationEventSource eventSource,
             GlobalPartitionEndpointManager globalPartitionEndpointManager,
-            JsonSerializerSettings serializerSettings = null)
+            JsonSerializerSettings serializerSettings = null,
+            IChaosInterceptor chaosInterceptor = null)
             : base(httpClient,
                   eventSource,
                   globalPartitionEndpointManager,
@@ -43,6 +46,7 @@ public ThinClientStoreClient(
             this.userAgentContainer = userAgentContainer
                 ?? throw new ArgumentNullException(nameof(userAgentContainer),
                 "UserAgentContainer cannot be null when initializing ThinClientStoreClient.");
+            this.chaosInterceptor = chaosInterceptor;
         }
 
         public override async Task<DocumentServiceResponse> InvokeAsync(
@@ -63,6 +67,19 @@ public override async Task<DocumentServiceResponse> InvokeAsync(
                 clientCollectionCache,
                 cancellationToken))
             {
+                if (this.chaosInterceptor != null)
+                {
+                    request.Headers.Set("FAULTINJECTION_IS_PROXY", "true");
+                    (bool hasFault, HttpResponseMessage fiResponseMessage) = await this.chaosInterceptor.OnHttpRequestCallAsync(request, cancellationToken);
+                    if (hasFault)
+                    {
+                        DefaultTrace.TraceInformation("Chaos interceptor injected fault for request: {0}", request);
+                        fiResponseMessage.RequestMessage = responseMessage.RequestMessage;
+                        request.Headers.Remove("FAULTINJECTION_IS_PROXY");
+                        return await ThinClientStoreClient.ParseResponseAsync(fiResponseMessage, request.SerializerSettings ?? base.SerializerSettings, request);
+                    }
+                }
+
                 HttpResponseMessage proxyResponse = await ThinClientTransportSerializer.ConvertProxyResponseAsync(responseMessage);
                 return await ThinClientStoreClient.ParseResponseAsync(proxyResponse, request.SerializerSettings ?? base.SerializerSettings, request);
             }