Refactor bulk reads to use CursorAsync (#409)

Author: Arkatufus

src/Akka.Persistence.MongoDb.Hosting/MongoDbJournalOptions.cs

@@ -46,6 +46,14 @@ public MongoDbJournalOptions(bool isDefaultPlugin, string identifier = "mongodb"
     /// Read Transaction
     /// </summary>
     public bool? UseReadTransaction { get; set; }
+    
+    /// <summary>
+    /// Set the batch size for read operations.
+    /// If not set (null), this will use the default MongoDb behavior where all queries
+    /// will have a batch size of 101 on the first page read and then will try to read
+    /// as many documents as possible that fits the 16 MeBi limit.
+    /// </summary>
+    public int? ReadBatchSize { get; set; }
 
     /// <summary>
     /// When true, enables BSON serialization (which breaks features like Akka.Cluster.Sharding, AtLeastOnceDelivery, and so on.)
@@ -82,6 +90,8 @@ protected override StringBuilder Build(StringBuilder sb)
         if(UseReadTransaction is not null)
             sb.AppendLine($"use-read-transaction = {UseReadTransaction.ToHocon()}");
 
+        sb.AppendLine($"read-batch-size = {(ReadBatchSize is not null ? ReadBatchSize.ToHocon() : "off")}");
+
         if(LegacySerialization is not null)
             sb.AppendLine($"legacy-serialization = {LegacySerialization.ToHocon()}");
 

src/Akka.Persistence.MongoDb.Tests/Hosting/MongoDbJournalOptionsSpec.cs

@@ -41,6 +41,7 @@ public void DefaultJournalOptionsTest()
             config.GetString("connection-string").Should().Be(baseConfig.GetString("connection-string"));
             config.GetBoolean("use-write-transaction").Should().Be(baseConfig.GetBoolean("use-write-transaction"));
             config.GetBoolean("use-read-transaction").Should().Be(baseConfig.GetBoolean("use-read-transaction"));
+            config.GetBoolean("read-batch-size").Should().Be(baseConfig.GetBoolean("read-batch-size"));
             config.GetBoolean("auto-initialize").Should().Be(baseConfig.GetBoolean("auto-initialize"));
             config.GetString("plugin-dispatcher").Should().Be(baseConfig.GetString("plugin-dispatcher"));
             config.GetString("collection").Should().Be(baseConfig.GetString("collection"));
@@ -68,6 +69,7 @@ public void CustomIdJournalOptionsTest()
             config.GetString("connection-string").Should().Be(baseConfig.GetString("connection-string"));
             config.GetBoolean("use-write-transaction").Should().Be(baseConfig.GetBoolean("use-write-transaction"));
             config.GetBoolean("use-read-transaction").Should().Be(baseConfig.GetBoolean("use-read-transaction"));
+            config.GetBoolean("read-batch-size").Should().Be(baseConfig.GetBoolean("read-batch-size"));
             config.GetBoolean("auto-initialize").Should().Be(baseConfig.GetBoolean("auto-initialize"));
             config.GetString("plugin-dispatcher").Should().Be(baseConfig.GetString("plugin-dispatcher"));
             config.GetString("collection").Should().Be(baseConfig.GetString("collection"));
@@ -88,6 +90,7 @@ public void JournalOptionsTest()
                 MetadataCollection = "metadataCollection",
                 UseWriteTransaction = true,
                 UseReadTransaction = true,
+                ReadBatchSize = 16,
                 LegacySerialization = true,
                 CallTimeout = TimeSpan.FromHours(2)
             };
@@ -104,6 +107,9 @@ public void JournalOptionsTest()
             config.GetString("metadata-collection").Should().Be(options.MetadataCollection);
             config.GetBoolean("use-write-transaction").Should().Be(options.UseWriteTransaction.Value);
             config.GetBoolean("use-read-transaction").Should().Be(options.UseReadTransaction.Value);
+            config.GetString("read-batch-size").ToLowerInvariant().Should().NotBe("off");
+            config.GetString("read-batch-size").ToLowerInvariant().Should().NotBe("false");
+            config.GetInt("read-batch-size").Should().Be(options.ReadBatchSize.Value);
             config.GetBoolean("legacy-serialization").Should().Be(options.LegacySerialization.Value);
             config.GetTimeSpan("call-timeout").Should().Be(options.CallTimeout.Value);
         }
@@ -121,6 +127,7 @@ public void JournalOptionsTest()
               ""ConnectionString"": ""mongodb://localhost:27017"",
               ""UseWriteTransaction"": ""true"",
               ""UseReadTransaction"": ""true"",
+              ""ReadBatchSize"": 16,
               ""Identifier"": ""custommongodb"",
               ""AutoInitialize"": true,
               ""IsDefaultPlugin"": false,
@@ -143,6 +150,7 @@ public void JournalOptionsIConfigurationBindingTest()
             options.ConnectionString.Should().Be("mongodb://localhost:27017");
             options.UseWriteTransaction.Should().BeTrue();
             options.UseReadTransaction.Should().BeTrue();
+            options.ReadBatchSize.Should().Be(16);
             options.Identifier.Should().Be("custommongodb");
             options.AutoInitialize.Should().BeTrue();
             options.IsDefaultPlugin.Should().BeFalse();

src/Akka.Persistence.MongoDb/Journal/MongoDbJournal.cs

@@ -44,6 +44,8 @@ public class MongoDbJournal : AsyncWriteJournal
         private readonly CancellationTokenSource _pendingCommandsCancellation = new();
 
         private readonly Akka.Serialization.Serialization _serialization;
+        
+        private readonly FindOptions? _findOptions;
 
         public MongoDbJournal() : this(MongoDbPersistence.Get(Context.System).JournalSettings)
         {
@@ -58,6 +60,11 @@ private MongoDbJournal(MongoDbJournalSettings settings)
         {
             _settings = settings;
             _serialization = Context.System.Serialization;
+            if (_settings.ReadBatchSize is not null)
+                _findOptions = new FindOptions
+                {
+                    BatchSize = settings.ReadBatchSize,
+                };
         }
 
         private IMongoDatabase GetMongoDb()
@@ -215,11 +222,11 @@ public override async Task ReplayMessagesAsync(
 
             await MaybeReadWithTransaction(async (session, ct) =>
             {
-                var collections = await journalCollection
-                    .ReplayMessagesQuery(session, persistenceId, fromSequenceNr, toSequenceNr, limitValue)
-                    .ToListAsync(ct);
+                var cursor = await journalCollection
+                    .ReplayMessagesQuery(session, persistenceId, fromSequenceNr, toSequenceNr, limitValue, _findOptions)
+                    .ToCursorAsync(ct);
 
-                collections.ForEach(doc => recoveryCallback(ToPersistenceRepresentation(doc, sender)));
+                await cursor.ForEachAsync(doc => recoveryCallback(ToPersistenceRepresentation(doc, sender)), ct);
             }, unitedCts.Token);
         }
 
@@ -256,8 +263,11 @@ private async Task<long> ReplayTaggedMessagesAsync(ReplayTaggedMessages replay)
 
                 var maxOrderingId = maxSeqNoEntry.Value;
 
-                await journalCollection.MessagesQuery(s, fromSequenceNr, toSequenceNr, maxOrderingId, tag, limitValue)
-                    .ForEachAsync(entry =>
+                var cursor = await journalCollection
+                    .MessagesQuery(s, fromSequenceNr, toSequenceNr, maxOrderingId, tag, limitValue, _findOptions)
+                    .ToCursorAsync(ct);
+                
+                await cursor.ForEachAsync(entry =>
                     {
                         var persistent = ToPersistenceRepresentation(entry, ActorRefs.NoSender);
                         foreach (var adapted in AdaptFromJournal(persistent))
@@ -579,8 +589,11 @@ protected virtual async Task<long> ReplayAllEventsAsync(ReplayAllEvents replay)
 
                 var maxOrderingId = maxSeqNoEntry.Value;
 
-                await journalCollection.MessagesQuery(session, fromSequenceNr, toSequenceNr, maxOrderingId, null, limitValue)
-                    .ForEachAsync(entry =>
+                var cursor = await journalCollection
+                    .MessagesQuery(session, fromSequenceNr, toSequenceNr, maxOrderingId, null, limitValue, _findOptions)
+                    .ToCursorAsync(token);
+                
+                await cursor.ForEachAsync(entry =>
                     {
                         var persistent = ToPersistenceRepresentation(entry, ActorRefs.NoSender);
                         foreach (var adapted in AdaptFromJournal(persistent))

src/Akka.Persistence.MongoDb/Journal/MongoDbJournalQueries.cs

@@ -29,7 +29,8 @@ public static IFindFluent<JournalEntry, JournalEntry> ReplayMessagesQuery(
         string persistenceId,
         long fromSequenceNr,
         long toSequenceNr,
-        int limit)
+        int limit,
+        FindOptions? findOptions)
     {
         var builder = Builders<JournalEntry>.Filter;
         var filter = builder.Eq(x => x.PersistenceId, persistenceId);
@@ -40,7 +41,7 @@ public static IFindFluent<JournalEntry, JournalEntry> ReplayMessagesQuery(
 
         var sort = Builders<JournalEntry>.Sort.Ascending(x => x.SequenceNr);
 
-        return (session is not null ? collection.Find(session, filter) : collection.Find(filter))
+        return (session is not null ? collection.Find(session, filter, findOptions) : collection.Find(filter, findOptions))
             .Sort(sort)
             .Limit(limit);
     }
@@ -85,7 +86,8 @@ public static IFindFluent<JournalEntry, JournalEntry> MessagesQuery(
             long toSequenceNr,
             long maxOrderingId,
             string? tag,
-            int limit
+            int limit,
+            FindOptions? findOptions
         )
     {
         var builder = Builders<JournalEntry>.Filter;
@@ -101,7 +103,7 @@ int limit
 
         var sort = Builders<JournalEntry>.Sort.Ascending(x => x.Ordering);
 
-        return (session is not null ? collection.Find(session, filter) : collection.Find(filter))
+        return (session is not null ? collection.Find(session, filter, findOptions) : collection.Find(filter, findOptions))
             .Sort(sort)
             .Limit(limit);
     }

src/Akka.Persistence.MongoDb/MongoDbSettings.cs

@@ -72,7 +72,9 @@ public class MongoDbJournalSettings : MongoDbSettings
     {
         public const string JournalConfigPath = "akka.persistence.journal.mongodb";
 
-        public string MetadataCollection { get; private set; }
+        public string MetadataCollection { get; }
+        
+        public int? ReadBatchSize { get; }
 
         public MongoDbJournalSettings(Config config) : base(config)
         {
@@ -81,6 +83,9 @@ public MongoDbJournalSettings(Config config) : base(config)
                     "MongoDb journal settings cannot be initialized, because required HOCON section couldn't been found");
 
             MetadataCollection = config.GetString("metadata-collection");
+            var readBatchSize = config.GetString("read-batch-size")?.ToLowerInvariant();
+            if(!string.IsNullOrWhiteSpace(readBatchSize) && readBatchSize != "off" && readBatchSize != "false")
+                ReadBatchSize = int.Parse(readBatchSize);
         }
     }
 

src/Akka.Persistence.MongoDb/reference.conf

@@ -12,6 +12,12 @@
 
 			use-read-transaction = on
 
+      # Set the batch size for read operations.
+      # If set to off or false, this will use the default MongoDb behavior where all queries
+      # will have a batch size of 101 on the first page read and then will try to read
+      # as many documents as possible that fits the 16 MeBi limit.
+      read-batch-size = off
+      
 			# should corresponding journal table's indexes be initialized automatically
 			auto-initialize = on