Add method for persistent state facets (#105)

Co-authored-by: Alex McAuliffe <[email protected]>

src/OrleansTestKit/OrleansTestKit.csproj

@@ -30,7 +30,7 @@
   </PropertyGroup>
 
   <ItemGroup>
-    <PackageReference Include="Microsoft.Orleans.OrleansRuntime" Version="3.1.7" />
+    <PackageReference Include="Microsoft.Orleans.OrleansRuntime" Version="3.4.3" />
     <PackageReference Include="Microsoft.SourceLink.GitHub" Version="1.1.1">
       <IncludeAssets>runtime; build; native; contentfiles; analyzers</IncludeAssets>
       <PrivateAssets>all</PrivateAssets>

src/OrleansTestKit/Storage/StorageExtensions.cs

@@ -1,29 +1,111 @@
 using System;
 using System.ComponentModel;
+using System.Reflection;
+using System.Threading.Tasks;
+using Moq;
+using Orleans.Core;
+using Orleans.Runtime;
+using Orleans.Storage;
 using Orleans.TestKit.Storage;
 
 namespace Orleans.TestKit
 {
     public static class StorageExtensions
     {
-        public static TState State<TState>(this TestKitSilo silo) where TState : class, new()
+        public static TState State<TGrain, TState>(this TestKitSilo silo)
+            where TGrain : Grain<TState>
         {
             if (silo == null)
             {
                 throw new ArgumentNullException(nameof(silo));
             }
 
-            return silo.StorageManager.GetStorage<TState>().State;
+            return silo.StorageManager.GetGrainStorage<TGrain, TState>().State;
         }
 
-        public static TestStorageStats StorageStats(this TestKitSilo silo)
+        public static TestStorageStats StorageStats<TGrain, TState>(this TestKitSilo silo)
+            where TGrain : Grain<TState>
         {
             if (silo == null)
             {
                 throw new ArgumentNullException(nameof(silo));
             }
 
-            return silo.StorageManager.StorageStats;
+            return silo.StorageManager.GetStorageStats<TGrain, TState>();
+        }
+
+        public static IStorage<T> AddGrainState<TGrain, T>(
+            this TestKitSilo silo,
+            T state = default)
+            where TGrain : Grain<T>
+        {
+            if (silo == null)
+            {
+                throw new ArgumentNullException(nameof(silo));
+            }
+
+            var storage = silo.StorageManager.GetGrainStorage<TGrain, T>();
+            storage.State = state ?? Activator.CreateInstance<T>();
+            return storage;
+        }
+
+        /// <summary>
+        /// Add persistent state to the silo for a given type. If a state is provided that will be the loaded state otherwise a new <typeparamref name="T"/>
+        /// </summary>
+        /// <remarks>
+        /// If neither StateName or StorageName are provided then we resolve just based on type, otherwise we try state name and optionally a storage name
+        /// </remarks>
+        /// <typeparam name="T">The type of data in the state</typeparam>
+        /// <param name="silo">The silo to add the state to</param>
+        /// <param name="stateName">The state name on the persistent state parameter</param>
+        /// <param name="storageName">The storage name on the persistent state parameter</param>
+        /// <param name="state">The state to set as default if any</param>
+        /// <returns>The persistent state</returns>
+        public static IPersistentState<T> AddPersistentState<T>(
+            this TestKitSilo silo,
+            string stateName,
+            string storageName = default,
+            T state = default)
+        {
+            return silo.AddPersistentStateStorage(
+                stateName,
+                storageName,
+                new TestStorage<T>(state ?? Activator.CreateInstance<T>()));
+        }
+
+        /// <summary>
+        /// Add persistent state to the silo for a given type. If a storage is provided that will provide the loaded state otherwise a new <typeparamref name="T"/>
+        /// </summary>
+        /// <remarks>
+        /// If neither StateName or StorageName are provided then we resolve just based on type, otherwise we try state name and optionally a storage name
+        /// </remarks>
+        /// <typeparam name="T">The type of data in the state</typeparam>
+        /// <param name="silo">The silo to add the state to</param>
+        /// <param name="stateName">The state name on the persistent state parameter</param>
+        /// <param name="storageName">The storage name on the persistent state parameter</param>
+        /// <param name="storage">The storage to use, if null a default implementation will be created</param>
+        /// <returns>The persistent state</returns>
+        public static IPersistentState<T> AddPersistentStateStorage<T>(
+            this TestKitSilo silo,
+            string stateName,
+            string storageName = default,
+            IStorage<T> storage = default)
+        {
+            var normalizedStorage = storage ?? new TestStorage<T>(Activator.CreateInstance<T>());
+            var normalizedStorageName = storageName ?? "Default";
+
+            if (silo == null)
+            {
+                throw new ArgumentNullException(nameof(silo));
+            }
+
+            if (string.IsNullOrWhiteSpace(stateName))
+            {
+                throw new ArgumentException("A state name must be provided", nameof(stateName));
+            }
+
+            silo.StorageManager.AddStorage(storage, stateName);
+            return silo.StorageManager.StateAttributeFactoryMapper.AddPersistentState(normalizedStorage, stateName, normalizedStorageName);
         }
     }
 }

src/OrleansTestKit/Storage/StorageManager.cs

@@ -1,5 +1,6 @@
 using System;
-using System.Diagnostics.CodeAnalysis;
+using System.Collections.Generic;
+using System.Linq;
 using Orleans.Core;
 
 namespace Orleans.TestKit.Storage
@@ -8,34 +9,63 @@ public sealed class StorageManager
     {
         private readonly TestKitOptions _options;
 
-        private object _storage;
+        private readonly Dictionary<string, object> _storages = new();
 
-        public StorageManager(TestKitOptions options) =>
+        public StorageManager(TestKitOptions options)
+        {
             _options = options ?? throw new ArgumentNullException(nameof(options));
+            StateAttributeFactoryMapper = new TestPersistentStateAttributeToFactoryMapper(this);
+        }
+
+        internal TestPersistentStateAttributeToFactoryMapper StateAttributeFactoryMapper { get; }
+
+        public IStorage<TState> GetGrainStorage<TGrain, TState>() where TGrain : Grain<TState>
+            => GetStorage<TState>(typeof(TGrain).FullName);
 
-        public IStorage<TState> GetStorage<TState>()
+        public IStorage<TState> GetStorage<TState>(string stateName)
         {
-            if (_storage == null)
+            if (string.IsNullOrWhiteSpace(stateName))
+            {
+                foreach (var kvp in _storages)
+                {
+                    if (kvp.Value is TestStorage<TState> typedStorage)
+                    {
+                        return typedStorage;
+                    }
+                }
+
+                throw new InvalidOperationException($"Unable to find any storage with type '{typeof(TState).FullName}'");
+            }
+
+            if (_storages.TryGetValue(stateName, out var storage) is false)
             {
-                _storage = _options.StorageFactory?.Invoke(typeof(TState)) ?? new TestStorage<TState>();
+                storage = _storages[stateName] = _options.StorageFactory?.Invoke(typeof(TState)) ?? new TestStorage<TState>();
             }
 
-            return _storage as IStorage<TState>;
+            return storage as IStorage<TState>;
         }
 
-        public TestStorageStats StorageStats
+        public TestStorageStats GetStorageStats<TGrain, TState>() where TGrain : Grain<TState>
+            => GetStorageStats(typeof(TGrain).FullName);
+
+        public TestStorageStats GetStorageStats(string stateName)
         {
-            get
+            var normalisedStateName = stateName ?? "Default";
+
+            if (_storages.TryGetValue(normalisedStateName, out var storage))
             {
-                //There should only be one state in here since there is only 1 grain under test
-                var stats = _storage as IStorageStats;
+                var stats = storage as IStorageStats;
                 return stats?.Stats;
             }
+
+            return null;
         }
 
-        [Obsolete("Use StorageStats property")]
-        [SuppressMessage("Design", "CA1024:Use properties where appropriate", Justification = "Keeping for backwards compatibility.")]
-        public TestStorageStats GetStorageStats() =>
-            StorageStats;
+        internal void AddStorage<TState>(IStorage<TState> storage, string stateName = default)
+        {
+            var normalisedStateName = stateName ?? "Default";
+
+            _storages[normalisedStateName] = storage;
+        }
     }
 }

src/OrleansTestKit/Storage/TestPersistentStateAttributeToFactoryMapper.cs

@@ -0,0 +1,121 @@
+using System;
+using System.Collections.Generic;
+using System.Reflection;
+using System.Threading.Tasks;
+using Orleans.Core;
+using Orleans.Runtime;
+
+namespace Orleans.TestKit.Storage
+{
+    public sealed class TestPersistentStateAttributeToFactoryMapper : IAttributeToFactoryMapper<PersistentStateAttribute>
+    {
+        private readonly StorageManager storageManager;
+        private readonly Dictionary<Type, Dictionary<(string StateName, string StorageName), object>>
+            registeredStorage = new();
+        private readonly MethodInfo AddEmptyStateMethod = typeof(TestPersistentStateAttributeToFactoryMapper)
+            .GetMethod(nameof(TestPersistentStateAttributeToFactoryMapper.AddEmptyState), BindingFlags.Instance | BindingFlags.NonPublic);
+
+        public TestPersistentStateAttributeToFactoryMapper(StorageManager storageManager) => this.storageManager = storageManager;
+
+        public IPersistentState<T> AddPersistentState<T>(
+            IStorage<T> storage,
+            string stateName,
+            string storageName)
+        {
+            if (storage is null)
+            {
+                throw new ArgumentNullException(nameof(storage));
+            }
+
+            if (string.IsNullOrWhiteSpace(stateName))
+            {
+                throw new ArgumentException($"'{nameof(stateName)}' cannot be null or whitespace.", nameof(stateName));
+            }
+
+            if (string.IsNullOrWhiteSpace(storageName))
+            {
+                throw new ArgumentException($"'{nameof(storageName)}' cannot be null or whitespace.", nameof(storageName));
+            }
+
+            var dict = registeredStorage.TryGetValue(typeof(T), out var typeStateRegistry)
+                ? typeStateRegistry
+                : registeredStorage[typeof(T)] = new Dictionary<(string StateName, string StorageName), object>(1);
+
+            var fake = new PersistentStateFake<T>(storage);
+            dict[(stateName, storageName)] = fake;
+
+            return fake;
+        }
+
+        public Factory<IGrainActivationContext, object> GetFactory(ParameterInfo parameter, PersistentStateAttribute metadata)
+        {
+            if (parameter is null)
+            {
+                throw new ArgumentNullException(nameof(parameter));
+            }
+
+            if (metadata is null)
+            {
+                throw new ArgumentNullException(nameof(metadata));
+            }
+
+            if (parameter.ParameterType.IsGenericType &&
+                parameter.ParameterType.GetGenericTypeDefinition() != typeof(IPersistentState<>))
+            {
+                throw new InvalidOperationException($"No persistent state for the parameter '{parameter.Name}'");
+            }
+
+            var parameterType = parameter.ParameterType.GenericTypeArguments[0];
+            var stateName = metadata.StateName;
+            var storageName = metadata.StorageName ?? "Default";
+
+            if (registeredStorage.TryGetValue(parameterType, out var typeStateRegistry))
+            {
+                // If we must have the state and the storage name so lookup by both
+                if (typeStateRegistry.TryGetValue((metadata.StateName, metadata.StorageName), out var persistentState))
+                {
+                    return _ => persistentState;
+                }
+            }
+
+            var state = AddEmptyStateMethod.MakeGenericMethod(parameterType).Invoke(
+                this,
+                new[] { stateName, storageName });
+
+            return _ => state;
+        }
+
+        private IPersistentState<TState> AddEmptyState<TState>(string stateName, string storageName)
+        {
+            var storage = storageManager.GetStorage<TState>(stateName);
+            return AddPersistentState(storage, stateName, storageName);
+        }
+    }
+
+    internal class PersistentStateFake<TState> : IPersistentState<TState>, IStorageStats
+    {
+        private readonly IStorage<TState> _storage;
+
+        public PersistentStateFake(IStorage<TState> storage) => _storage = storage;
+
+        public TState State
+        {
+            get => _storage.State;
+            set => _storage.State = value;
+        }
+
+        public string Etag => _storage.Etag;
+
+        public bool RecordExists => _storage.RecordExists;
+
+        public TestStorageStats Stats => _storage is IStorageStats statsStorage
+            ? statsStorage.Stats
+            : null;
+
+        public Task ClearStateAsync() => _storage.ClearStateAsync();
+
+        public Task ReadStateAsync() => _storage.ReadStateAsync();
+
+        public Task WriteStateAsync() => _storage.WriteStateAsync();
+    }
+}

src/OrleansTestKit/Storage/TestStorage.cs

@@ -8,10 +8,14 @@ internal class TestStorage<TState> :
         IStorageStats,
         IStorage<TState>
     {
-        public TestStorage()
+        public TestStorage() : this(CreateState())
+        {
+        }
+
+        public TestStorage(TState state)
         {
             Stats = new TestStorageStats() { Reads = -1 };
-            InitializeState();
+            State = state;
         }
 
         public string Etag => throw new System.NotImplementedException();
@@ -24,7 +28,7 @@ public TestStorage()
 
         public Task ClearStateAsync()
         {
-            InitializeState();
+            State = CreateState();
             Stats.Clears++;
             RecordExists = false;
             return Task.CompletedTask;
@@ -43,15 +47,15 @@ public Task WriteStateAsync()
             return Task.CompletedTask;
         }
 
-        private void InitializeState()
+        private static TState CreateState()
         {
             if (!typeof(TState).IsValueType && typeof(TState).GetConstructor(Type.EmptyTypes) == null)
             {
                 throw new NotSupportedException(
                     $"No parameterless constructor defined for {typeof(TState).Name}. This is currently not supported");
             }
 
-            State = Activator.CreateInstance<TState>();
+            return Activator.CreateInstance<TState>();
         }
     }
 }

src/OrleansTestKit/TestGrainRuntime.cs

@@ -57,7 +57,7 @@ public void DelayDeactivation(Grain grain, TimeSpan timeSpan)
             Mock.Object.DelayDeactivation(grain, timeSpan);
         }
 
-        public IStorage<TGrainState> GetStorage<TGrainState>(Grain grain) =>
-            _storageManager.GetStorage<TGrainState>();
+        public IStorage<TGrainState> GetStorage<TGrainState>(Grain grain)
+            => _storageManager.GetStorage<TGrainState>(grain.GetType().FullName);
     }
 }

src/OrleansTestKit/TestKitSilo.cs

@@ -111,6 +111,9 @@ private async Task<T> CreateGrainAsync<T>(IGrainIdentity identity)
                 throw new Exception("A grain has already been created in this silo. Only 1 grain per test silo should ever be created. Add grain probes for supporting grains.");
             }
 
+            // Add state attribute mapping for storage facets
+            this.AddService<IAttributeToFactoryMapper<PersistentStateAttribute>>(StorageManager.StateAttributeFactoryMapper);
+
             _isGrainCreated = true;
             Grain grain;
             var grainContext = new TestGrainActivationContext