RavenDB-21084 : HandleClusterDatabaseChanged : if sharded database is…

… disabled we need to remove the ShardedDatabaseContext from ShardedDatabasesCache
aviv86 · Aug 31, 2023 · 1047767 · 1047767
1 parent 58e6295
commit 1047767
Show file tree

Hide file tree

Showing 2 changed files with 88 additions and 22 deletions.
diff --git a/src/Raven.Server/Documents/DatabasesLandlord.cs b/src/Raven.Server/Documents/DatabasesLandlord.cs
@@ -138,6 +138,15 @@ await HandleSpecificClusterDatabaseChanged(
                             var topology = rawRecord.Sharding.Orchestrator.Topology;
                             if (topology.RelevantFor(_serverStore.NodeTag))
                             {
+                                if (rawRecord.IsDisabled)
+                                {
+                                    if (ShardedDatabasesCache.TryGetValue(databaseName, out var shardedDatabaseTask) == false)
+                                        return; // sharded database was already unloaded
+
+                                    UnloadDatabaseInternal(databaseName, shardedDatabaseTask);
+                                    return;
+                                }
+
                                 // we need to update this upon any shard topology change
                                 // and upon migration completion
                                 var databaseContext = GetOrAddShardedDatabaseContext(databaseName, rawRecord);
@@ -301,28 +310,7 @@ private void UnloadDatabase(string databaseName, bool dbRecordIsNull = false)
                 return;
             }
 
-            if (databaseTask.IsCompletedSuccessfully)
-            {
-                UnloadDatabaseInternal(databaseName);
-            }
-            else if (databaseTask.IsCompleted == false)
-            {
-                // This case is when an unload was issued prior to the actual loading of a database.
-                databaseTask.ContinueWith(t =>
-                {
-                    if (databaseTask.IsCompletedSuccessfully)
-                    {
-                        using (_serverStore.ContextPool.AllocateOperationContext(out TransactionOperationContext context))
-                        using (context.OpenReadTransaction())
-                        using (var databaseRecord = _serverStore.Cluster.ReadRawDatabaseRecord(context, databaseName))
-                        {
-                            // unload only if DB is still disabled
-                            if (IsDatabaseDisabled(databaseRecord.Raw))
-                                UnloadDatabaseInternal(databaseName);
-                        }
-                    }
-                });
-            }
+            UnloadDatabaseInternal(databaseName, databaseTask);
         }
 
         public static bool IsDatabaseDisabled(BlittableJsonReaderObject databaseRecord)
@@ -352,6 +340,33 @@ private void UnloadDatabaseInternal(string databaseName, [CallerMemberName] stri
             }
         }
 
+        private void UnloadDatabaseInternal(string databaseName, Task databaseTask)
+        {
+            if (databaseTask.IsCompletedSuccessfully)
+            {
+                UnloadDatabaseInternal(databaseName);
+            }
+            else if (databaseTask.IsCompleted == false)
+            {
+                // This case is when an unload was issued prior to the actual loading of a database.
+                databaseTask.ContinueWith(t =>
+                {
+                    if (databaseTask.IsCompletedSuccessfully)
+                    {
+                        using (_serverStore.ContextPool.AllocateOperationContext(out TransactionOperationContext context))
+                        using (context.OpenReadTransaction())
+                        using (var databaseRecord = _serverStore.Cluster.ReadRawDatabaseRecord(context, databaseName))
+                        {
+                            // unload only if DB is still disabled
+                            if (IsDatabaseDisabled(databaseRecord.Raw))
+                                UnloadDatabaseInternal(databaseName);
+                        }
+                    }
+                });
+            }
+        }
+
+
         public bool ShouldDeleteDatabase(TransactionOperationContext context, string dbName, RawDatabaseRecord rawRecord, bool fromReplication = false)
         {
             var shard = ShardHelper.TryGetShardNumberFromDatabaseName(dbName, out var shardNumber);

diff --git a/test/SlowTests/Sharding/Issues/RavenDB_21084.cs b/test/SlowTests/Sharding/Issues/RavenDB_21084.cs
@@ -0,0 +1,51 @@
+using System.Threading.Tasks;
+using FastTests;
+using Raven.Client.ServerWide.Operations;
+using Raven.Server.Utils;
+using Tests.Infrastructure;
+using Xunit;
+using Xunit.Abstractions;
+
+namespace SlowTests.Sharding.Issues;
+
+public class RavenDB_21084 : RavenTestBase
+{
+    public RavenDB_21084(ITestOutputHelper output) : base(output)
+    {
+    }
+
+    [RavenFact(RavenTestCategory.Sharding)]
+    public async Task DisableShardedDb_ShouldRemoveDbFromShardedDatabasesCache()
+    {
+        using (var store = Sharding.GetDocumentStore())
+        {
+            var dbLandlord = Server.ServerStore.DatabasesLandlord;
+
+            Assert.True(dbLandlord.ShardedDatabasesCache.TryGetValue(store.Database, out _));
+
+            var shardingConfig = await Sharding.GetShardingConfigurationAsync(store);
+            foreach (var shardNumber in shardingConfig.Shards.Keys)
+            {
+                var shard = ShardHelper.ToShardName(store.Database, shardNumber);
+                Assert.True(dbLandlord.DatabasesCache.TryGetValue(shard, out _));
+            }
+
+            var operationResult = await store.Maintenance.Server.SendAsync(new ToggleDatabasesStateOperation(store.Database, disable: true));
+            Assert.True(operationResult.Success);
+            Assert.True(operationResult.Disabled);
+
+            // should remove document databases from DatabasesCache
+            foreach (var shardNumber in shardingConfig.Shards.Keys)
+            {
+                var shard = ShardHelper.ToShardName(store.Database, shardNumber);
+                Assert.False(dbLandlord.DatabasesCache.TryGetValue(shard, out _));
+            }
+
+            // should remove sharded database context from ShardedDatabasesCache
+            Assert.False(dbLandlord.ShardedDatabasesCache.TryGetValue(store.Database, out _));
+
+            await store.Maintenance.Server.SendAsync(new ToggleDatabasesStateOperation(store.Database, disable: false));
+        }
+    }
+
+}