fix(transactions): scope storage fault injection

src/Orleans.Transactions.TestKit.Base/FaultInjection/ControlledInjection/FaultInjectionAzureTableTransactionStateStorage.cs

@@ -42,9 +42,12 @@ public async Task<string> Store(
             long? abortAfter
         )
         {
-            faultInjector.BeforeStore();
+            var transactionIds = ControlledTransactionFaultInjectorExtensions.GetTransactionIds(
+                metadata,
+                statesToPrepare);
+            faultInjector.BeforeStore(transactionIds);
             var result = await this.stateStorage.Store(expectedETag, metadata, statesToPrepare, commitUpTo, abortAfter);
-            faultInjector.AfterStore();
+            faultInjector.AfterStore(transactionIds);
             return result;
         }
     }

src/Orleans.Transactions.TestKit.Base/FaultInjection/ControlledInjection/FaultInjectionDynamoDBTransactionStateStorage.cs

@@ -55,9 +55,12 @@ public async Task<string> Store(
         long? abortAfter
     )
     {
-        faultInjector.BeforeStore();
+        var transactionIds = ControlledTransactionFaultInjectorExtensions.GetTransactionIds(
+            metadata,
+            statesToPrepare);
+        faultInjector.BeforeStore(transactionIds);
         var result = await this.stateStorage.Store(expectedETag, metadata, statesToPrepare, commitUpTo, abortAfter);
-        faultInjector.AfterStore();
+        faultInjector.AfterStore(transactionIds);
         return result;
     }
 }

src/Orleans.Transactions.TestKit.Base/FaultInjection/ControlledInjection/FaultInjectionTransactionReource.cs

@@ -35,12 +35,7 @@ public async Task<TransactionalStatus> PrepareAndCommit(Guid transactionId, Acce
                 out var injectionType);
             if (injectBeforeStore)
             {
-                if (injectionType == FaultInjectionType.ExceptionBeforeStore)
-                    this.faultInjector.InjectBeforeStore = true;
-                if (injectionType == FaultInjectionType.ExceptionAfterStore)
-                    this.faultInjector.InjectAfterStore = true;
-                if (injectionType == FaultInjectionType.GenericExceptionAfterStore)
-                    this.faultInjector.InjectGenericAfterStore = true;
+                this.faultInjector.Arm(transactionId, injectionType, requireTransactionMatch: true);
                 LogInformationInjectedFaultBeforePrepareAndCommit(this.logger, context.GrainInstance, transactionId, injectionType);
                 FaultInjectionDiagnosticEvents.Emit(new(
                     this.context.GrainId,
@@ -235,12 +230,7 @@ public async Task Confirm(Guid transactionId, DateTime timeStamp)
                 out var injectionType);
             if (injectBeforeStore)
             {
-                if (injectionType == FaultInjectionType.ExceptionBeforeStore)
-                    this.faultInjector.InjectBeforeStore = true;
-                if (injectionType == FaultInjectionType.ExceptionAfterStore)
-                    this.faultInjector.InjectAfterStore = true;
-                if (injectionType == FaultInjectionType.GenericExceptionAfterStore)
-                    this.faultInjector.InjectGenericAfterStore = true;
+                this.faultInjector.Arm(transactionId, injectionType);
                 LogInformationInjectedFaultBeforeConfirm(this.logger, context.GrainInstance, transactionId, injectionType);
                 FaultInjectionDiagnosticEvents.Emit(new(
                     this.context.GrainId,
@@ -273,12 +263,7 @@ public async Task Prepare(Guid transactionId, AccessCounter accessCount, DateTim
                 out var injectionType);
             if (injectBeforeStore)
             {
-                if (injectionType == FaultInjectionType.ExceptionBeforeStore)
-                    this.faultInjector.InjectBeforeStore = true;
-                if (injectionType == FaultInjectionType.ExceptionAfterStore)
-                    this.faultInjector.InjectAfterStore = true;
-                if (injectionType == FaultInjectionType.GenericExceptionAfterStore)
-                    this.faultInjector.InjectGenericAfterStore = true;
+                this.faultInjector.Arm(transactionId, injectionType);
                 LogInformationInjectedFaultBeforePrepare(this.logger, this.context.GrainInstance, transactionId, injectionType);
                 FaultInjectionDiagnosticEvents.Emit(new(
                     this.context.GrainId,

src/Orleans.Transactions.TestKit.Base/FaultInjection/ControlledInjection/IControlledFaultInjector.cs

@@ -1,9 +1,91 @@
-namespace Orleans.Transactions.TestKit
+using System;
+using System.Collections.Generic;
+using System.Collections.Immutable;
+using Orleans.Transactions.Abstractions;
+
+namespace Orleans.Transactions.TestKit
 {
     public interface IControlledTransactionFaultInjector : ITransactionFaultInjector
     {
         bool InjectBeforeStore { get; set; }
         bool InjectAfterStore { get; set; }
         bool InjectGenericAfterStore { get; set; }
     }
+
+    internal interface ITransactionScopedFaultInjector
+    {
+        void Arm(Guid transactionId, FaultInjectionType injectionType, bool requireTransactionMatch);
+        void BeforeStore(ImmutableArray<Guid> transactionIds);
+        void AfterStore(ImmutableArray<Guid> transactionIds);
+    }
+
+    internal static class ControlledTransactionFaultInjectorExtensions
+    {
+        public static void Arm(
+            this IControlledTransactionFaultInjector faultInjector,
+            Guid transactionId,
+            FaultInjectionType injectionType,
+            bool requireTransactionMatch = false)
+        {
+            if (faultInjector is ITransactionScopedFaultInjector scopedFaultInjector)
+            {
+                scopedFaultInjector.Arm(transactionId, injectionType, requireTransactionMatch);
+                return;
+            }
+
+            faultInjector.InjectBeforeStore = injectionType == FaultInjectionType.ExceptionBeforeStore;
+            faultInjector.InjectAfterStore = injectionType == FaultInjectionType.ExceptionAfterStore;
+            faultInjector.InjectGenericAfterStore = injectionType == FaultInjectionType.GenericExceptionAfterStore;
+        }
+
+        public static ImmutableArray<Guid> GetTransactionIds<TState>(
+            TransactionalStateMetaData metadata,
+            List<PendingTransactionState<TState>>? statesToPrepare)
+            where TState : class, new()
+        {
+            var result = ImmutableArray.CreateBuilder<Guid>(
+                metadata.CommitRecords.Count + (statesToPrepare?.Count ?? 0));
+            result.AddRange(metadata.CommitRecords.Keys);
+            if (statesToPrepare is not null)
+            {
+                foreach (var state in statesToPrepare)
+                {
+                    if (Guid.TryParse(state.TransactionId, out var transactionId))
+                    {
+                        result.Add(transactionId);
+                    }
+                }
+            }
+
+            return result.MoveToImmutable();
+        }
+
+        public static void BeforeStore(
+            this ITransactionFaultInjector faultInjector,
+            ImmutableArray<Guid> transactionIds)
+        {
+            if (faultInjector is ITransactionScopedFaultInjector scopedFaultInjector)
+            {
+                scopedFaultInjector.BeforeStore(transactionIds);
+            }
+            else
+            {
+                faultInjector.BeforeStore();
+            }
+        }
+
+        public static void AfterStore(
+            this ITransactionFaultInjector faultInjector,
+            ImmutableArray<Guid> transactionIds)
+        {
+            if (faultInjector is ITransactionScopedFaultInjector scopedFaultInjector)
+            {
+                scopedFaultInjector.AfterStore(transactionIds);
+            }
+            else
+            {
+                faultInjector.AfterStore();
+            }
+        }
+    }
 }

src/Orleans.Transactions.TestKit.Base/FaultInjection/ControlledInjection/SimpleAzureStorageExceptionInjector.cs

@@ -1,13 +1,16 @@
 using System;
+using System.Collections.Immutable;
 using System.ComponentModel;
 using System.Runtime.Serialization;
 using Azure;
 using Microsoft.Extensions.Logging;
 
 namespace Orleans.Transactions.TestKit
 {
-    public partial class SimpleAzureStorageExceptionInjector : IControlledTransactionFaultInjector
+    public partial class SimpleAzureStorageExceptionInjector : IControlledTransactionFaultInjector, ITransactionScopedFaultInjector
     {
+        private readonly object lockObj = new();
+        private Guid? targetTransactionId;
         public bool InjectBeforeStore { get; set; }
         public bool InjectAfterStore { get; set; }
         public bool InjectGenericAfterStore { get; set; }
@@ -22,37 +25,87 @@ public SimpleAzureStorageExceptionInjector(ILogger<SimpleAzureStorageExceptionIn
 
         public void AfterStore()
         {
-            if (InjectAfterStore)
+            this.AfterStore(default);
+        }
+
+        void ITransactionScopedFaultInjector.Arm(
+            Guid transactionId,
+            FaultInjectionType injectionType,
+            bool requireTransactionMatch)
+        {
+            lock (this.lockObj)
             {
-                InjectAfterStore = false;
-                this.injectionAfterStoreCounter++;
-                var message = $"Storage exception thrown after store, thrown total {injectionAfterStoreCounter}";
-                LogInformationMessage(this.logger, message);
-                throw new SimpleAzureStorageException(message);
+                this.targetTransactionId = requireTransactionMatch ? transactionId : null;
+                this.InjectBeforeStore = injectionType == FaultInjectionType.ExceptionBeforeStore;
+                this.InjectAfterStore = injectionType == FaultInjectionType.ExceptionAfterStore;
+                this.InjectGenericAfterStore = injectionType == FaultInjectionType.GenericExceptionAfterStore;
             }
+        }
+
+        void ITransactionScopedFaultInjector.BeforeStore(ImmutableArray<Guid> transactionIds)
+            => this.BeforeStore(transactionIds);
 
-            if (InjectGenericAfterStore)
+        void ITransactionScopedFaultInjector.AfterStore(ImmutableArray<Guid> transactionIds)
+            => this.AfterStore(transactionIds);
+
+        private void AfterStore(ImmutableArray<Guid> transactionIds)
+        {
+            lock (this.lockObj)
             {
-                InjectGenericAfterStore = false;
-                this.genericInjectionAfterStoreCounter++;
-                var message = $"Generic storage exception thrown after store, thrown total {genericInjectionAfterStoreCounter}";
-                LogInformationMessage(this.logger, message);
-                throw new InvalidOperationException(message);
+                if (!this.IsTargetStore(transactionIds))
+                {
+                    return;
+                }
+
+                if (this.InjectAfterStore)
+                {
+                    this.InjectAfterStore = false;
+                    this.targetTransactionId = null;
+                    this.injectionAfterStoreCounter++;
+                    var message = $"Storage exception thrown after store, thrown total {injectionAfterStoreCounter}";
+                    LogInformationMessage(this.logger, message);
+                    throw new SimpleAzureStorageException(message);
+                }
+
+                if (this.InjectGenericAfterStore)
+                {
+                    this.InjectGenericAfterStore = false;
+                    this.targetTransactionId = null;
+                    this.genericInjectionAfterStoreCounter++;
+                    var message = $"Generic storage exception thrown after store, thrown total {genericInjectionAfterStoreCounter}";
+                    LogInformationMessage(this.logger, message);
+                    throw new InvalidOperationException(message);
+                }
             }
         }
 
         public void BeforeStore()
         {
-            if (InjectBeforeStore)
+            this.BeforeStore(default);
+        }
+
+        private void BeforeStore(ImmutableArray<Guid> transactionIds)
+        {
+            lock (this.lockObj)
             {
-                InjectBeforeStore = false;
+                if (!this.IsTargetStore(transactionIds) || !this.InjectBeforeStore)
+                {
+                    return;
+                }
+
+                this.InjectBeforeStore = false;
+                this.targetTransactionId = null;
                 this.injectionBeforeStoreCounter++;
                 var message = $"Storage exception thrown before store. Thrown total {injectionBeforeStoreCounter}";
                 LogInformationMessage(this.logger, message);
                 throw new SimpleAzureStorageException(message);
             }
         }
 
+        private bool IsTargetStore(ImmutableArray<Guid> transactionIds)
+            => this.targetTransactionId is not { } targetTransactionId
+                || (!transactionIds.IsDefaultOrEmpty && transactionIds.IndexOf(targetTransactionId) >= 0);
+
         [LoggerMessage(
             Level = LogLevel.Information,
             Message = "{Message}"

test/Transactions/Orleans.Transactions.Tests/FaultInjectionControlTests.cs

@@ -1,3 +1,6 @@
+using System;
+using System.Collections.Immutable;
+using Microsoft.Extensions.Logging.Abstractions;
 using Orleans.Transactions.TestKit;
 using Xunit;
 
@@ -42,4 +45,66 @@ public void NonMatchingPhaseDoesNotConsumeFault()
         Assert.Equal(TransactionFaultInjectPhase.BeforePrepareAndCommit, control.FaultInjectionPhase);
         Assert.Equal(FaultInjectionType.ExceptionAfterStore, control.FaultInjectionType);
     }
+
+    [Theory]
+    [InlineData(FaultInjectionType.ExceptionBeforeStore)]
+    [InlineData(FaultInjectionType.ExceptionAfterStore)]
+    [InlineData(FaultInjectionType.GenericExceptionAfterStore)]
+    public void ScopedStorageFaultIgnoresUnrelatedBatches(FaultInjectionType injectionType)
+    {
+        var targetTransactionId = Guid.NewGuid();
+        var injector = new SimpleAzureStorageExceptionInjector(
+            NullLogger<SimpleAzureStorageExceptionInjector>.Instance);
+        var scopedInjector = Assert.IsAssignableFrom<ITransactionScopedFaultInjector>(injector);
+        scopedInjector.Arm(targetTransactionId, injectionType, requireTransactionMatch: true);
+
+        var unrelatedTransactionIds = ImmutableArray.Create(Guid.NewGuid());
+        scopedInjector.BeforeStore(unrelatedTransactionIds);
+        scopedInjector.AfterStore(unrelatedTransactionIds);
+
+        var targetTransactionIds = ImmutableArray.Create(targetTransactionId);
+        AssertFaultInjected(scopedInjector, injectionType, targetTransactionIds);
+
+        scopedInjector.BeforeStore(targetTransactionIds);
+        scopedInjector.AfterStore(targetTransactionIds);
+    }
+
+    [Theory]
+    [InlineData(FaultInjectionType.ExceptionBeforeStore)]
+    [InlineData(FaultInjectionType.ExceptionAfterStore)]
+    [InlineData(FaultInjectionType.GenericExceptionAfterStore)]
+    public void UnscopedStorageFaultInjectsIntoUnidentifiedBatch(FaultInjectionType injectionType)
+    {
+        var injector = new SimpleAzureStorageExceptionInjector(
+            NullLogger<SimpleAzureStorageExceptionInjector>.Instance);
+        var scopedInjector = Assert.IsAssignableFrom<ITransactionScopedFaultInjector>(injector);
+        scopedInjector.Arm(Guid.NewGuid(), injectionType, requireTransactionMatch: false);
+
+        AssertFaultInjected(scopedInjector, injectionType, ImmutableArray<Guid>.Empty);
+
+        scopedInjector.BeforeStore(ImmutableArray<Guid>.Empty);
+        scopedInjector.AfterStore(ImmutableArray<Guid>.Empty);
+    }
+
+    private static void AssertFaultInjected(
+        ITransactionScopedFaultInjector faultInjector,
+        FaultInjectionType injectionType,
+        ImmutableArray<Guid> transactionIds)
+    {
+        if (injectionType == FaultInjectionType.ExceptionBeforeStore)
+        {
+            Assert.Throws<SimpleAzureStorageException>(() => faultInjector.BeforeStore(transactionIds));
+            return;
+        }
+
+        faultInjector.BeforeStore(transactionIds);
+        if (injectionType == FaultInjectionType.ExceptionAfterStore)
+        {
+            Assert.Throws<SimpleAzureStorageException>(() => faultInjector.AfterStore(transactionIds));
+        }
+        else
+        {
+            Assert.Throws<InvalidOperationException>(() => faultInjector.AfterStore(transactionIds));
+        }
+    }
 }