Validate GrainId Type and Key on reminder creation

wassim-k · wassim-k · commit c8d86e86d8dc · 2025-09-23T06:30:10.000+10:00
diff --git a/Orleans.Providers.MongoDB/Membership/Store/Multiple/MultipleMembershipCollection.cs b/Orleans.Providers.MongoDB/Membership/Store/Multiple/MultipleMembershipCollection.cs
@@ -49,7 +49,7 @@ public async Task<bool> UpsertRow(string deploymentId, MembershipEntry entry, st
             using var session = await Client.StartSessionAsync();
             return await session.WithTransactionAsync(async (sessionHandle, ct) =>
             {
-                var hasUpsertedTable = await tableVersionCollection.UpsertAsync(session, tableVersion, deploymentId);
+                var hasUpsertedTable = await tableVersionCollection.UpsertAsync(sessionHandle, tableVersion, deploymentId);
 
                 if (!hasUpsertedTable)
                 {
@@ -106,10 +106,10 @@ public async Task<MembershipTableData> ReadAll(string deploymentId)
             using var session = await Client.StartSessionAsync();
             return await session.WithTransactionAsync(async (sessionHandle, ct) =>
             {
-                var tableVersion = tableVersionCollection.GetTableVersionAsync(deploymentId);
+                var tableVersion = tableVersionCollection.GetTableVersionAsync(sessionHandle, deploymentId);
 
                 var entries =
-                    Collection.Find(x => x.DeploymentId == deploymentId)
+                    Collection.Find(sessionHandle, x => x.DeploymentId == deploymentId)
                         .ToListAsync(cancellationToken: ct);
 
                 await Task.WhenAll(tableVersion, entries);
@@ -123,12 +123,12 @@ public async Task<MembershipTableData> ReadRow(string deploymentId, SiloAddress
             using var session = await Client.StartSessionAsync();
             return await session.WithTransactionAsync(async (sessionHandle, ct) =>
             {
-                var tableVersion = tableVersionCollection.GetTableVersionAsync(deploymentId);
+                var tableVersion = tableVersionCollection.GetTableVersionAsync(sessionHandle, deploymentId);
 
                 var id = ReturnId(deploymentId, address);
 
                 var entries =
-                    Collection.Find(x => x.Id == id)
+                    Collection.Find(sessionHandle, x => x.Id == id)
                         .ToListAsync(cancellationToken: ct);
 
                 await Task.WhenAll(tableVersion, entries);
@@ -151,12 +151,19 @@ public Task CleanupDefunctSiloEntries(string deploymentId, DateTimeOffset before
             return Collection.DeleteManyAsync(x => x.DeploymentId == deploymentId && x.Status != (int)SiloStatus.Active && x.Timestamp < beforeUtc);
         }
 
-        public Task DeleteMembershipTableEntries(string deploymentId)
+        public async Task DeleteMembershipTableEntries(string deploymentId)
         {
-            return Task.WhenAll(
-                Collection.DeleteManyAsync(x => x.DeploymentId == deploymentId),
-                tableVersionCollection.DeleteAsync(deploymentId)
-            );
+            using var session = await Client.StartSessionAsync();
+
+            await session.WithTransactionAsync(async (sessionHandle, ct) =>
+            {
+                await Task.WhenAll(
+                    Collection.DeleteManyAsync(sessionHandle, x => x.DeploymentId == deploymentId),
+                    tableVersionCollection.DeleteAsync(sessionHandle, deploymentId)
+                );
+
+                return true;
+            });
         }
 
         private static MembershipTableData ReturnMembershipTableData(IEnumerable<MongoMembershipDocument> membershipList, TableVersion tableVersion)
diff --git a/Orleans.Providers.MongoDB/Membership/Store/Multiple/TableVersionCollection.cs b/Orleans.Providers.MongoDB/Membership/Store/Multiple/TableVersionCollection.cs
@@ -26,14 +26,14 @@ protected override string CollectionName()
             return $"{collectionPrefix}OrleansMembershipV3_TableVersion";
         }
 
-        public Task DeleteAsync(string deploymentId)
+        public Task DeleteAsync(IClientSessionHandle session, string deploymentId)
         {
-            return Collection.DeleteOneAsync(x => x.DeploymentId == deploymentId);
+            return Collection.DeleteOneAsync(session, x => x.DeploymentId == deploymentId);
         }
 
-        public async Task<TableVersion> GetTableVersionAsync(string deploymentId)
+        public async Task<TableVersion> GetTableVersionAsync(IClientSessionHandle session, string deploymentId)
         {
-            var deployment = await Collection.Find(x => x.DeploymentId == deploymentId).FirstOrDefaultAsync();
+            var deployment = await Collection.Find(session, x => x.DeploymentId == deploymentId).FirstOrDefaultAsync();
 
             if (deployment == null)
             {
@@ -50,7 +50,7 @@ public async Task<bool> UpsertAsync(IClientSessionHandle session, TableVersion t
             try
             {
                 await Collection.ReplaceOneAsync(session,
-                    x => x.DeploymentId == deploymentId && x.VersionEtag == tableVersion.VersionEtag, 
+                    x => x.DeploymentId == deploymentId && x.VersionEtag == tableVersion.VersionEtag,
                     update,
                     UpsertReplace);
 
diff --git a/Orleans.Providers.MongoDB/Orleans.Providers.MongoDB.csproj b/Orleans.Providers.MongoDB/Orleans.Providers.MongoDB.csproj
@@ -1,11 +1,11 @@
-﻿<Project Sdk="Microsoft.NET.Sdk">
+<Project Sdk="Microsoft.NET.Sdk">
 
   <PropertyGroup>
     <TargetFrameworks>net8.0</TargetFrameworks>
     <Authors>laredoza,sebastianstehle,wassim-k</Authors>
     <Company>Orleans.Providers.MongoDB</Company>
     <Copyright>MIT</Copyright>
-    <Description>A MongoDb implementation of the Orleans Providers. This includes custering (IMembershipTable and IGatewayListProvider), reminders (IReminderTable) and storage providers (IGrainStorage).</Description>
+    <Description>A MongoDb implementation of the Orleans Providers. This includes clustering (IMembershipTable and IGatewayListProvider), reminders (IReminderTable) and storage providers (IGrainStorage).</Description>
     <EmbedUntrackedSources>true</EmbedUntrackedSources>
     <GeneratePackageOnBuild>false</GeneratePackageOnBuild>
     <IncludeSymbols>true</IncludeSymbols>
@@ -18,12 +18,12 @@
   </PropertyGroup>
 
   <ItemGroup>
-    <PackageReference Include="Microsoft.Orleans.Reminders" Version="9.1.2" />
+    <PackageReference Include="Microsoft.Orleans.Reminders" Version="9.0.1" />
     <PackageReference Include="Microsoft.SourceLink.GitHub" Version="8.0.0">
       <PrivateAssets>all</PrivateAssets>
       <IncludeAssets>runtime; build; native; contentfiles; analyzers; buildtransitive</IncludeAssets>
     </PackageReference>
-    <PackageReference Include="MongoDB.Driver" Version="3.2.1" />
+    <PackageReference Include="MongoDB.Driver" Version="3.0.0" />
   </ItemGroup>
 
   <ItemGroup>
diff --git a/Orleans.Providers.MongoDB/Reminders/Store/MongoReminderCollection.cs b/Orleans.Providers.MongoDB/Reminders/Store/MongoReminderCollection.cs
@@ -200,6 +200,14 @@ await Collection.ReplaceOneAsync(x => x.Id == id,
 
         private static string ReturnId(string serviceId, GrainId grainId, string reminderName)
         {
+            var grainType = grainId.Type.ToString();
+            var grainKey = grainId.Key.ToString();
+
+            if (grainType is null || grainKey is null)
+            {
+                throw new ArgumentNullException(nameof(grainId));
+            }
+
             return $"{serviceId}_{grainId}_{reminderName}";
         }
     }
diff --git a/Orleans.Providers.MongoDB/StorageProviders/MongoGrainStorage.cs b/Orleans.Providers.MongoDB/StorageProviders/MongoGrainStorage.cs
@@ -1,6 +1,5 @@
 ﻿using System;
 using System.Collections.Concurrent;
-using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.Logging;
 using MongoDB.Driver;
@@ -12,14 +11,13 @@
 
 namespace Orleans.Providers.MongoDB.StorageProviders
 {
-    public class MongoGrainStorage : IGrainStorage, ILifecycleParticipant<ISiloLifecycle>
+    public class MongoGrainStorage : IGrainStorage
     {
         private readonly ConcurrentDictionary<string, MongoGrainStorageCollection> collections = new ConcurrentDictionary<string, MongoGrainStorageCollection>();
         private readonly MongoDBGrainStorageOptions options;
         private readonly IMongoClient mongoClient;
         private readonly ILogger<MongoGrainStorage> logger;
         private readonly IGrainStateSerializer serializer;
-        private IMongoDatabase database;
 
         public MongoGrainStorage(
             IMongoClientFactory mongoClientFactory,
@@ -32,21 +30,6 @@ public MongoGrainStorage(
             this.serializer = options.GrainStateSerializer;
         }
 
-        public void Participate(ISiloLifecycle lifecycle)
-        {
-            lifecycle.Subscribe<MongoGrainStorage>(ServiceLifecycleStage.ApplicationServices, Init);
-        }
-
-        private Task Init(CancellationToken ct)
-        {
-            return DoAndLog(nameof(Init), () =>
-            {
-                database = mongoClient.GetDatabase(options.DatabaseName);
-
-                return Task.CompletedTask;
-            });
-        }
-
         public Task ReadStateAsync<T>(string stateName, GrainId grainId, IGrainState<T> grainState)
         {
             return DoAndLog(nameof(ReadStateAsync), () =>
diff --git a/Test/GrainInterfaces/Orleans.Providers.MongoDB.Test.GrainInterfaces.csproj b/Test/GrainInterfaces/Orleans.Providers.MongoDB.Test.GrainInterfaces.csproj
@@ -6,9 +6,9 @@
   </PropertyGroup>
 
   <ItemGroup>
-    <PackageReference Include="Microsoft.Orleans.Sdk" Version="9.1.2" />
-    <PackageReference Include="Microsoft.Orleans.Core.Abstractions" Version="9.1.2" />
-    <PackageReference Include="Microsoft.Orleans.Reminders" Version="9.1.2" />
+    <PackageReference Include="Microsoft.Orleans.Sdk" Version="9.0.1" />
+    <PackageReference Include="Microsoft.Orleans.Core.Abstractions" Version="9.0.1" />
+    <PackageReference Include="Microsoft.Orleans.Reminders" Version="9.0.1" />
   </ItemGroup>
 
 </Project>
diff --git a/Test/Grains/Orleans.Providers.MongoDB.Test.Grains.csproj b/Test/Grains/Orleans.Providers.MongoDB.Test.Grains.csproj
@@ -6,9 +6,9 @@
   </PropertyGroup>
 
   <ItemGroup>
-    <PackageReference Include="Microsoft.Orleans.Sdk" Version="9.1.2" />
-    <PackageReference Include="Microsoft.Orleans.Core" Version="9.1.2" />
-    <PackageReference Include="Microsoft.Orleans.Streaming" Version="9.1.2" />
+    <PackageReference Include="Microsoft.Orleans.Sdk" Version="9.0.1" />
+    <PackageReference Include="Microsoft.Orleans.Core" Version="9.0.1" />
+    <PackageReference Include="Microsoft.Orleans.Streaming" Version="9.0.1" />
   </ItemGroup>
 
   <ItemGroup>
diff --git a/Test/Host/Orleans.Providers.MongoDB.Test.Host.csproj b/Test/Host/Orleans.Providers.MongoDB.Test.Host.csproj
@@ -1,4 +1,4 @@
-﻿<Project Sdk="Microsoft.NET.Sdk">
+<Project Sdk="Microsoft.NET.Sdk">
   
   <PropertyGroup>
     <OutputType>Exe</OutputType>
@@ -8,7 +8,10 @@
   <ItemGroup>
     <PackageReference Include="Microsoft.Extensions.Logging.Console" Version="9.0.2" />
     <PackageReference Include="Microsoft.Extensions.Hosting" Version="9.0.2" />
-    <PackageReference Include="MongoSandbox8" Version="1.0.1" />
+    <PackageReference Include="MongoSandbox.Core" Version="2.0.0" />
+    <PackageReference Include="MongoSandbox8.runtime.linux-x64" Version="2.0.0" Condition="$([MSBuild]::IsOSPlatform('Linux'))" />
+    <PackageReference Include="MongoSandbox8.runtime.osx-arm64" Version="2.0.0" Condition="$([MSBuild]::IsOSPlatform('OSX'))" />
+    <PackageReference Include="MongoSandbox8.runtime.win-x64" Version="2.0.0" Condition="$([MSBuild]::IsOSPlatform('Windows'))" />
   </ItemGroup>
 
   <ItemGroup>
diff --git a/Test/Host/Program.cs b/Test/Host/Program.cs
@@ -25,7 +25,7 @@ public static async Task Main(string[] args)
         {
             var createShardKey = false;
 
-            using var mongoRunner = MongoRunner.Run(new MongoRunnerOptions { KillMongoProcessesWhenCurrentProcessExits = true });
+            using var mongoRunner = MongoRunner.Run();
 
             Console.WriteLine("MongoDB ConnectionString: {0}", mongoRunner.ConnectionString);
 
diff --git a/UnitTest/Orleans.Providers.MongoDB.UnitTest.csproj b/UnitTest/Orleans.Providers.MongoDB.UnitTest.csproj
@@ -1,4 +1,4 @@
-﻿<Project Sdk="Microsoft.NET.Sdk">
+<Project Sdk="Microsoft.NET.Sdk">
   <PropertyGroup>
     <TargetFramework>net8.0</TargetFramework>
     <AutoGenerateBindingRedirects>true</AutoGenerateBindingRedirects>
@@ -9,8 +9,11 @@
   </PropertyGroup>
 
   <ItemGroup>
-    <PackageReference Include="Microsoft.NET.Test.Sdk" Version="17.13.0" />
-    <PackageReference Include="MongoSandbox8" Version="1.0.1" />
+    <PackageReference Include="Microsoft.NET.Test.Sdk" Version="17.14.1" />
+    <PackageReference Include="MongoSandbox.Core" Version="2.0.0" />
+    <PackageReference Include="MongoSandbox8.runtime.linux-x64" Version="2.0.0" Condition="$([MSBuild]::IsOSPlatform('Linux'))" />
+    <PackageReference Include="MongoSandbox8.runtime.osx-arm64" Version="2.0.0" Condition="$([MSBuild]::IsOSPlatform('OSX'))" />
+    <PackageReference Include="MongoSandbox8.runtime.win-x64" Version="2.0.0" Condition="$([MSBuild]::IsOSPlatform('Windows'))" />
     <PackageReference Include="xunit" Version="2.9.3" />
   </ItemGroup>
 

Original file line number	Diff line number	Diff line change
`@@ -26,14 +26,14 @@ protected override string CollectionName()`
`26`	`26`	`return $"{collectionPrefix}OrleansMembershipV3_TableVersion";`
`27`	`27`	`}`
`28`	`28`
`29`		`- public Task DeleteAsync(string deploymentId)`
	`29`	`+ public Task DeleteAsync(IClientSessionHandle session, string deploymentId)`
`30`	`30`	`{`
`31`		`- return Collection.DeleteOneAsync(x => x.DeploymentId == deploymentId);`
	`31`	`+ return Collection.DeleteOneAsync(session, x => x.DeploymentId == deploymentId);`
`32`	`32`	`}`
`33`	`33`
`34`		`- public async Task<TableVersion> GetTableVersionAsync(string deploymentId)`
	`34`	`+ public async Task<TableVersion> GetTableVersionAsync(IClientSessionHandle session, string deploymentId)`
`35`	`35`	`{`
`36`		`- var deployment = await Collection.Find(x => x.DeploymentId == deploymentId).FirstOrDefaultAsync();`
	`36`	`+ var deployment = await Collection.Find(session, x => x.DeploymentId == deploymentId).FirstOrDefaultAsync();`
`37`	`37`
`38`	`38`	`if (deployment == null)`
`39`	`39`	`{`
`@@ -50,7 +50,7 @@ public async Task<bool> UpsertAsync(IClientSessionHandle session, TableVersion t`
`50`	`50`	`try`
`51`	`51`	`{`
`52`	`52`	`await Collection.ReplaceOneAsync(session,`
`53`		`- x => x.DeploymentId == deploymentId && x.VersionEtag == tableVersion.VersionEtag,`
	`53`	`+ x => x.DeploymentId == deploymentId && x.VersionEtag == tableVersion.VersionEtag,`
`54`	`54`	`update,`
`55`	`55`	`UpsertReplace);`
`56`	`56`
Original file line number	Diff line number	Diff line change
`@@ -200,6 +200,14 @@ await Collection.ReplaceOneAsync(x => x.Id == id,`
`200`	`200`
`201`	`201`	`private static string ReturnId(string serviceId, GrainId grainId, string reminderName)`
`202`	`202`	`{`
	`203`	`+ var grainType = grainId.Type.ToString();`
	`204`	`+ var grainKey = grainId.Key.ToString();`
	`205`	`+`
	`206`	`+ if (grainType is null \|\| grainKey is null)`
	`207`	`+ {`
	`208`	`+ throw new ArgumentNullException(nameof(grainId));`
	`209`	`+ }`
	`210`	`+`
`203`	`211`	`return $"{serviceId}_{grainId}_{reminderName}";`
`204`	`212`	`}`
`205`	`213`	`}`
Original file line number	Diff line number	Diff line change
`@@ -25,7 +25,7 @@ public static async Task Main(string[] args)`
`25`	`25`	`{`
`26`	`26`	`var createShardKey = false;`
`27`	`27`
`28`		`- using var mongoRunner = MongoRunner.Run(new MongoRunnerOptions { KillMongoProcessesWhenCurrentProcessExits = true });`
	`28`	`+ using var mongoRunner = MongoRunner.Run();`
`29`	`29`
`30`	`30`	`Console.WriteLine("MongoDB ConnectionString: {0}", mongoRunner.ConnectionString);`
`31`	`31`