Add Support for passing state between Before and After state operatio…

…ns (#22)

.devcontainer/Dockerfile

@@ -0,0 +1,16 @@
+# See here for image contents: https://github.com/microsoft/vscode-dev-containers/tree/v0.245.2/containers/dotnet/.devcontainer/base.Dockerfile
+
+# [Choice] .NET version: 6.0, 3.1, 6.0-bullseye, 3.1-bullseye, 6.0-focal, 3.1-focal
+ARG VARIANT="6.0-bullseye-slim"
+FROM mcr.microsoft.com/vscode/devcontainers/dotnet:0-${VARIANT}
+
+# [Choice] Node.js version: none, lts/*, 18, 16, 14
+ARG NODE_VERSION="none"
+RUN if [ "${NODE_VERSION}" != "none" ]; then su vscode -c "umask 0002 && . /usr/local/share/nvm/nvm.sh && nvm install ${NODE_VERSION} 2>&1"; fi
+
+# [Optional] Uncomment this section to install additional OS packages.
+# RUN apt-get update && export DEBIAN_FRONTEND=noninteractive \
+#     && apt-get -y install --no-install-recommends <your-package-list-here>
+
+# [Optional] Uncomment this line to install global node packages.
+# RUN su vscode -c "source /usr/local/share/nvm/nvm.sh && npm install -g <your-package-here>" 2>&1

.devcontainer/devcontainer.json

@@ -0,0 +1,60 @@
+// For format details, see https://aka.ms/devcontainer.json. For config options, see the README at:
+// https://github.com/microsoft/vscode-dev-containers/tree/v0.245.2/containers/dotnet
+{
+	"name": "C# (.NET)",
+	"build": {
+		"dockerfile": "Dockerfile",
+		"args": {
+			// Update 'VARIANT' to pick a .NET Core version: 3.1, 6.0
+			// Append -bullseye or -focal to pin to an OS version.
+			"VARIANT": "6.0-bullseye",
+			// Options
+			"NODE_VERSION": "none"
+		}
+	},
+	// Configure tool-specific properties.
+	"customizations": {
+		// Configure properties specific to VS Code.
+		"vscode": {
+			// Add the IDs of extensions you want installed when the container is created.
+			"extensions": [
+				"ms-dotnettools.csharp"
+			]
+		}
+	},
+	// Use 'forwardPorts' to make a list of ports inside the container available locally.
+	// "forwardPorts": [5000, 5001],
+	// [Optional] To reuse of your local HTTPS dev cert:
+	//
+	// 1. Export it locally using this command:
+	//    * Windows PowerShell:
+	//        dotnet dev-certs https --trust; dotnet dev-certs https -ep "$env:USERPROFILE/.aspnet/https/aspnetapp.pfx" -p "SecurePwdGoesHere"
+	//    * macOS/Linux terminal:
+	//        dotnet dev-certs https --trust; dotnet dev-certs https -ep "${HOME}/.aspnet/https/aspnetapp.pfx" -p "SecurePwdGoesHere"
+	// 
+	// 2. Uncomment these 'remoteEnv' lines:
+	//    "remoteEnv": {
+	// 	      "ASPNETCORE_Kestrel__Certificates__Default__Password": "SecurePwdGoesHere",
+	//        "ASPNETCORE_Kestrel__Certificates__Default__Path": "/home/vscode/.aspnet/https/aspnetapp.pfx",
+	//    },
+	//
+	// 3. Do one of the following depending on your scenario:
+	//    * When using GitHub Codespaces and/or Remote - Containers:
+	//      1. Start the container
+	//      2. Drag ~/.aspnet/https/aspnetapp.pfx into the root of the file explorer
+	//      3. Open a terminal in VS Code and run "mkdir -p /home/vscode/.aspnet/https && mv aspnetapp.pfx /home/vscode/.aspnet/https"
+	//
+	//    * If only using Remote - Containers with a local container, uncomment this line instead:
+	//      "mounts": [ "source=${env:HOME}${env:USERPROFILE}/.aspnet/https,target=/home/vscode/.aspnet/https,type=bind" ],
+	// Use 'postCreateCommand' to run commands after the container is created.
+	// "postCreateCommand": "dotnet restore",
+	// Comment out to connect as root instead. More info: https://aka.ms/vscode-remote/containers/non-root.
+	"remoteUser": "vscode",
+	"features": {
+		"ghcr.io/devcontainers/features/dotnet:1": {
+			"version": "latest"
+		},
+		"git": "latest",
+		"git-lfs": "latest"
+	}
+}

.vscode/launch.json

@@ -0,0 +1,15 @@
+{
+    "configurations": [
+    {
+        "name": ".NET Core Launch (console)",
+        "type": "coreclr",
+        "request": "launch",
+        "preLaunchTask": "build",
+        "program": "${workspaceFolder}/Source/Sample/bin/Debug/net6.0/Sample.dll",
+        "args": [],
+        "cwd": "${workspaceFolder}",
+        "stopAtEntry": false,
+        "console": "internalConsole"
+    }
+    ]
+}

.vscode/tasks.json

@@ -0,0 +1,24 @@
+{
+    // See https://go.microsoft.com/fwlink/?LinkId=733558
+    // for the documentation about the tasks.json format
+    "version": "2.0.0",
+    "tasks": [
+        {
+            "label": "build",
+            "command": "dotnet",
+            "type": "shell",
+            "args": [
+                "build",
+                // Ask dotnet build to generate full paths for file names.
+                "/property:GenerateFullPaths=true",
+                // Do not generate summary otherwise it leads to duplicate errors in Problems panel
+                "/consoleloggerparameters:NoSummary"
+            ],
+            "group": "build",
+            "presentation": {
+                "reveal": "silent"
+            },
+            "problemMatcher": "$msCompile"
+        }
+    ]
+}

Source/Orleans.StorageProviderInterceptors/Infrastructure/NamedStorageInterceptorFactory.cs

@@ -292,30 +292,33 @@ public TState State
         /// <inheritdoc/>
         public async Task ClearStateAsync()
         {
-            if (!await this.options.OnBeforeClearStateAsync(this.context, this))
+            var (preventOperation, state) = await this.options.OnBeforeClearStateAsync(this.context, this);
+            if (!preventOperation)
             {
                 await this.storage.ClearStateAsync();
-                await this.options.OnAfterClearStateAsync(this.context, this);
+                await this.options.OnAfterClearStateAsync(this.context, this, state);
             }
         }
 
         /// <inheritdoc/>
         public async Task WriteStateAsync()
         {
-            if (!await this.options.OnBeforeWriteStateFunc(this.context, this))
+            var (preventOperation, state) = await this.options.OnBeforeWriteStateFunc(this.context, this);
+            if (!preventOperation)
             {
                 await this.storage.WriteStateAsync();
-                await this.options.OnAfterWriteStateFunc(this.context, this);
+                await this.options.OnAfterWriteStateFunc(this.context, this, state);
             }
         }
 
         /// <inheritdoc/>
         public async Task ReadStateAsync()
         {
-            if (!await this.options.OnBeforeReadStateAsync(this.context, this))
+            var (preventOperation, state) = await this.options.OnBeforeReadStateAsync(this.context, this);
+            if (!preventOperation)
             {
                 await this.storage.ReadStateAsync();
-                await this.options.OnAfterReadStateFunc.Invoke(this.context, this);
+                await this.options.OnAfterReadStateFunc.Invoke(this.context, this, state);
             }
         }
 

Source/Orleans.StorageProviderInterceptors/Infrastructure/StorageInterceptorOptions.cs

@@ -10,30 +10,30 @@ public class StorageInterceptorOptions<TState>
     /// <summary>
     /// Called before a ClearStateAsync(); return false to prevent writing.
     /// </summary>
-    public Func<IGrainActivationContext, IPersistentState<TState>, ValueTask<bool>> OnBeforeClearStateAsync { get; set; } = (_, _) => ValueTask.FromResult(false);
+    public Func<IGrainActivationContext, IPersistentState<TState>, ValueTask<(bool PreventOperation, object? SharedState)>> OnBeforeClearStateAsync { get; set; } = (_, _) => ValueTask.FromResult((false, (object?)null));
 
     /// <summary>
     /// Called after ClearStateAsync();
     /// </summary>
-    public Func<IGrainActivationContext, IPersistentState<TState>, ValueTask> OnAfterClearStateAsync { get; set; } = (_, _) => ValueTask.CompletedTask;
+    public Func<IGrainActivationContext, IPersistentState<TState>, object?, ValueTask> OnAfterClearStateAsync { get; set; } = (_, _, _) => ValueTask.CompletedTask;
 
     /// <summary>
     /// Called before ReadStateAsync(); return false to prevent Reading
     /// </summary>
-    public Func<IGrainActivationContext, IPersistentState<TState>, ValueTask<bool>> OnBeforeReadStateAsync { get; set; } = (_, _) => ValueTask.FromResult(false);
+    public Func<IGrainActivationContext, IPersistentState<TState>, ValueTask<(bool PreventOperation, object? SharedState)>> OnBeforeReadStateAsync { get; set; } = (_, _) => ValueTask.FromResult((false, (object?)null));
 
     /// <summary>
     /// Called after ReadStateAsync()
     /// </summary>
-    public Func<IGrainActivationContext, IPersistentState<TState>, ValueTask> OnAfterReadStateFunc { get; set; } = (_, _) => ValueTask.CompletedTask;
+    public Func<IGrainActivationContext, IPersistentState<TState>, object?, ValueTask> OnAfterReadStateFunc { get; set; } = (_, _, _) => ValueTask.CompletedTask;
 
     /// <summary>
     /// Called before WriteStateAsync(); return false to prevent writing
     /// </summary>
-    public Func<IGrainActivationContext, IPersistentState<TState>, ValueTask<bool>> OnBeforeWriteStateFunc { get; set; } = (_, _) => ValueTask.FromResult(true);
+    public Func<IGrainActivationContext, IPersistentState<TState>, ValueTask<(bool PreventOperation, object? SharedState)>> OnBeforeWriteStateFunc { get; set; } = (_, _) => ValueTask.FromResult((false, (object?)null));
 
     /// <summary>
     /// Called after WriteStateAsync()
     /// </summary>
-    public Func<IGrainActivationContext, IPersistentState<TState>, ValueTask> OnAfterWriteStateFunc { get; set; } = (_, _) => ValueTask.CompletedTask;
+    public Func<IGrainActivationContext, IPersistentState<TState>, object?, ValueTask> OnAfterWriteStateFunc { get; set; } = (_, _, _) => ValueTask.CompletedTask;
 }

Source/Sample/Program.cs

@@ -20,51 +20,74 @@
         .AddStorageInterceptors()
         .UseGenericStorageInterceptor<Dictionary<string, string>>("SecretsStorage", "secretsState", c =>
         {
-            c.OnBeforeReadStateAsync = (grainActivationContext, currentState) =>
-            {
-                Console.WriteLine($"OnBeforeReadState: {grainActivationContext.GrainIdentity.IdentityString}: Count Is {currentState.State.Count}");
-                return ValueTask.FromResult(true);
-            };
-
-            c.OnAfterReadStateFunc = (grainActivationContext, currentState) =>
-            {
-                Console.WriteLine($"OnAfterReadState: {grainActivationContext.GrainIdentity.IdentityString}: Count Is {currentState.State.Count}");
-
-                // Do a deep copy
-                //  Dictionary<string, string>? stateToModify = JsonConvert.DeserializeObject<Dictionary<string, string>>(JsonConvert.SerializeObject(currentState));
-
-                foreach (var (key, value) in currentState.State)
-                {
-                    Console.WriteLine($"Intercepted: {key}: {value}");
-
-                    // Decrypt the data
-                    currentState.State[key] = currentState.State[key].Replace('3', 'e');
-                }
-                return ValueTask.CompletedTask;
-            };
 
             c.OnBeforeWriteStateFunc = (grainActivationContext, currentState) =>
                 {
+                    var unencryptedValues = new Dictionary<string, string>(currentState.State.Count);
                     Console.WriteLine($"OnBeforeWriteState: {grainActivationContext.GrainIdentity.IdentityString}: Count Is {currentState.State.Count}");
                     foreach (var (key, value) in currentState.State)
                     {
                         Console.WriteLine($"Intercepted: {key}: {value}");
 
+                        // Save the original state to return to the grain
+                        unencryptedValues.Add(key, value);
+
                         // Encrypt the data
                         currentState.State[key] = currentState.State[key].Replace('e', '3');
                     }
-                    return ValueTask.FromResult(true);
+                    return ValueTask.FromResult((false, (object?)unencryptedValues));
                 };
 
-            c.OnAfterWriteStateFunc = (grainActivationContext, currentState) =>
+            c.OnAfterWriteStateFunc = (grainActivationContext, currentState, sharedState) =>
             {
+                var unencryptedValues = (Dictionary<string, string>)sharedState!;
                 Console.WriteLine($"OnAfterWriteState: {grainActivationContext.GrainIdentity.IdentityString}: Count Is {currentState.State.Count}");
                 foreach (var (key, value) in currentState.State)
                 {
                     Console.WriteLine($"What was actually persisted: {key}: {value}");
+
+                    currentState.State[key] = unencryptedValues[key];
+                    Console.WriteLine($"What will be returned to grain: {key}: {unencryptedValues[key]}");
                 }
                 return ValueTask.CompletedTask;
             };
+
+            c.OnBeforeReadStateAsync = (grainActivationContext, currentState) =>
+            {
+                Console.WriteLine($"OnBeforeReadState: {grainActivationContext.GrainIdentity.IdentityString}: Count Is {currentState.State.Count}");
+
+                var unencryptedValues = new Dictionary<string, string>(currentState.State.Count);
+                foreach (var (key, value) in currentState.State)
+                {
+                    Console.WriteLine($"Unencrypted Values: {key}: {value}");
+
+                    // Save the original state to return to the grain
+                    unencryptedValues.Add(key, value);
+                }
+                return ValueTask.FromResult((false, (object?)unencryptedValues));
+            };
+
+            c.OnAfterReadStateFunc = (grainActivationContext, currentState, sharedState) =>
+            {
+                var unencryptedValues = (Dictionary<string, string>)sharedState!;
+                if (!currentState.RecordExists)
+                {
+                    return ValueTask.CompletedTask;
+                }
+
+                var list = sharedState as List<string>;
+                Console.WriteLine($"OnAfterReadState: {grainActivationContext.GrainIdentity.IdentityString}: Count Is {currentState.State.Count}");
+
+                foreach (var (key, value) in currentState.State)
+                {
+                    Console.WriteLine($"Encrypted Values: {key}: {value}");
+
+                    // Decrypt the data
+                    currentState.State[key] = currentState.State[key].Replace('3', 'e');
+                }
+                return ValueTask.CompletedTask;
+            };
+
         }))
     .Build();
 
@@ -81,12 +104,16 @@
 var secretStore = grainFactory.GetGrain<ISecretStorageGrain>("friend");
 
 // Call the grain and print the result to the console
-await secretStore.AddSecret("Password", "Now you See the secrets and now they are seen as safe!");
+await secretStore.AddOrUpdateSecret("Password", "Now you See the secrets and now they are seen as safe!");
 
 var result = await secretStore.GetSecret("Password");
 
-Console.WriteLine("\n\n{0}\n\n", result);
+Console.WriteLine($"Original Value: {result}");
 
+await secretStore.AddOrUpdateSecret("Password", "Seeeeeeeeeecrets");
+result = await secretStore.GetSecret("Password");
+
+Console.WriteLine($"Updated Value: {result}");
 Console.ReadLine();
 
 await host.StopAsync();
@@ -98,18 +125,19 @@ internal class SecretStorageGrain : Grain, ISecretStorageGrain
         private readonly IPersistentState<Dictionary<string, string>> secrets;
 
         public SecretStorageGrain([StorageInterceptor("SecretsStorage", "secretsState")] IPersistentState<Dictionary<string, string>> state) => this.secrets = state;
-        public async Task AddSecret(string name, string value)
+        public async Task AddOrUpdateSecret(string name, string value)
         {
-            this.secrets.State.Add(name, value);
+            this.secrets.State[name] = value;
             await this.secrets.WriteStateAsync();
         }
 
+
         public Task<string> GetSecret(string name) => Task.FromResult(this.secrets.State[name]);
     }
 
     internal interface ISecretStorageGrain : IGrainWithStringKey
     {
-        Task AddSecret(string name, string value);
+        Task AddOrUpdateSecret(string name, string value);
         Task<string> GetSecret(string name);
     }
 }

Tests/Orleans.StorageProviderInterceptors.Test/Class1Test.cs

@@ -5,8 +5,5 @@ namespace Orleans.StorageProviderInterceptors.Test;
 public class Class1Test
 {
     [Fact]
-    public void Given_When_Then() =>
-        //var class1 = new Class1();
-
-        Assert.True(true);
+    public void Given_When_Then() => Assert.True(true);
 }