CSHARP-2207: Large bulk writes not always splitting correctly

Author: rstam

src/MongoDB.Driver.Core/Core/Misc/SizeLimitingBatchableSourceSerializer.cs

@@ -52,8 +52,8 @@ public SizeLimitingBatchableSourceSerializer(
             _itemSerializer = Ensure.IsNotNull(itemSerializer, nameof(itemSerializer));
             _itemElementNameValidator = Ensure.IsNotNull(itemElementNameValidator, nameof(itemElementNameValidator));
             _maxBatchCount = Ensure.IsGreaterThanZero(maxBatchCount, nameof(maxBatchCount));
-            _maxItemSize = Ensure.IsBetween(maxItemSize, 1, int.MaxValue, nameof(maxItemSize));
-            _maxBatchSize = Ensure.IsBetween(maxBatchSize, maxItemSize, int.MaxValue, nameof(maxBatchSize));
+            _maxItemSize = Ensure.IsGreaterThanZero(maxItemSize, nameof(maxItemSize));
+            _maxBatchSize = Ensure.IsGreaterThanZero(maxBatchSize, nameof(maxBatchSize));
         }
 
         // public methods
@@ -88,19 +88,23 @@ public override void Serialize(BsonSerializationContext context, BsonSerializati
                     var item = value.Items[value.Offset + i];
                     _itemSerializer.Serialize(context, args, item);
 
-                    var batchSize = binaryWriter?.Position - startPosition;
-                    if (batchSize > _maxBatchSize)
+                    // always process at least one item
+                    if (i > 0)
                     {
-                        if (i > 0 && value.CanBeSplit)
+                        var batchSize = binaryWriter?.Position - startPosition;
+                        if (batchSize > _maxBatchSize)
                         {
-                            binaryWriter.BaseStream.Position = itemPosition.Value; // remove the last item
-                            binaryWriter.BaseStream.SetLength(itemPosition.Value);
-                            value.SetProcessedCount(i);
-                            return;
-                        }
-                        else
-                        {
-                            throw new ArgumentException("Batch is too large.");
+                            if (value.CanBeSplit)
+                            {
+                                binaryWriter.BaseStream.Position = itemPosition.Value; // remove the last item
+                                binaryWriter.BaseStream.SetLength(itemPosition.Value);
+                                value.SetProcessedCount(i);
+                                return;
+                            }
+                            else
+                            {
+                                throw new ArgumentException("Batch is too large.");
+                            }
                         }
                     }
                 }

src/MongoDB.Driver.Core/Core/Operations/RetryableDeleteCommandOperation.cs

@@ -123,8 +123,8 @@ private IBsonSerializer<BatchableSource<DeleteRequest>> CreateBatchSerializer(Co
             if (attempt == 1)
             {
                 var maxBatchCount = Math.Min(MaxBatchCount ?? int.MaxValue, connectionDescription.MaxBatchCount);
-                var maxItemSize = connectionDescription.MaxDocumentSize;
-                var maxBatchSize = connectionDescription.MaxWireDocumentSize;
+                var maxItemSize = connectionDescription.MaxWireDocumentSize;
+                var maxBatchSize = connectionDescription.MaxDocumentSize;
                 return new SizeLimitingBatchableSourceSerializer<DeleteRequest>(DeleteRequestSerializer.Instance, NoOpElementNameValidator.Instance, maxBatchCount, maxItemSize, maxBatchSize);
             }
             else

src/MongoDB.Driver.Core/Core/Operations/RetryableInsertCommandOperation.cs

@@ -150,7 +150,7 @@ private IBsonSerializer<BatchableSource<TDocument>> CreateBatchSerializer(Connec
             {
                 var maxBatchCount = Math.Min(MaxBatchCount ?? int.MaxValue, connectionDescription.MaxBatchCount);
                 var maxItemSize = connectionDescription.MaxDocumentSize;
-                var maxBatchSize = connectionDescription.MaxWireDocumentSize;
+                var maxBatchSize = connectionDescription.MaxDocumentSize;
                 return new SizeLimitingBatchableSourceSerializer<TDocument>(itemSerializer, elementNameValidator, maxBatchCount, maxItemSize, maxBatchSize);
             }
             else

src/MongoDB.Driver.Core/Core/Operations/RetryableUpdateCommandOperation.cs

@@ -145,7 +145,7 @@ private IBsonSerializer<BatchableSource<UpdateRequest>> CreateBatchSerializer(Co
             {
                 var maxBatchCount = Math.Min(MaxBatchCount ?? int.MaxValue, connectionDescription.MaxBatchCount);
                 var maxItemSize = connectionDescription.MaxWireDocumentSize;
-                var maxBatchSize = connectionDescription.MaxWireDocumentSize;
+                var maxBatchSize = connectionDescription.MaxDocumentSize;
                 return new SizeLimitingBatchableSourceSerializer<UpdateRequest>(UpdateRequestSerializer.Instance, NoOpElementNameValidator.Instance, maxBatchCount, maxItemSize, maxBatchSize);
             }
             else

tests/MongoDB.Driver.Core.TestHelpers/EventCapturer.cs

@@ -15,6 +15,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Linq;
 using MongoDB.Bson;
 using MongoDB.Driver.Core.Events;
 using MongoDB.Driver.Core.Misc;
@@ -64,6 +65,17 @@ public void Clear()
             }
         }
 
+        public List<object> Events
+        {
+            get
+            {
+                lock (_lock)
+                {
+                    return _capturedEvents.ToList();
+                }
+            }
+        }
+
         public object Next()
         {
             lock (_lock)

tests/MongoDB.Driver.Core.Tests/Core/Operations/BulkMixedWriteOperationTests.cs

@@ -17,17 +17,13 @@
 using System.Collections.Generic;
 using System.Linq;
 using System.Threading;
-using System.Threading.Tasks;
 using FluentAssertions;
 using MongoDB.Bson;
 using MongoDB.Bson.Serialization.Serializers;
 using MongoDB.Bson.TestHelpers.XunitExtensions;
 using MongoDB.Driver.Core.Bindings;
-using MongoDB.Driver.Core.Clusters;
-using MongoDB.Driver.Core.Configuration;
 using MongoDB.Driver.Core.Events;
 using MongoDB.Driver.Core.Misc;
-using MongoDB.Driver.Core.TestHelpers;
 using MongoDB.Driver.Core.TestHelpers.XunitExtensions;
 using Xunit;
 
@@ -1257,11 +1253,149 @@ public void Execute_with_update_should_send_session_id_when_supported(
             VerifySessionIdWasSentWhenSupported(subject, "update", async);
         }
 
-        private List<BsonDocument> ReadAllFromCollection(IReadBinding binding)
+        [SkippableTheory]
+        [InlineData(new[] { 1 }, new[] { 1 })]
+        [InlineData(new[] { 1, 1 }, new[] { 2 })]
+        [InlineData(new[] { 8388605, 8388605 }, new[] { 2 })]
+        [InlineData(new[] { 8388605, 8388606 }, new[] { 1, 1 })]
+        [InlineData(new[] { 16777216 }, new[] { 1 })]
+        [InlineData(new[] { 16777216, 1 }, new[] { 1, 1 })]
+        [InlineData(new[] { 16777216, 16777216 }, new[] { 1, 1 })]
+        public void Execute_with_multiple_deletes_should_split_batches_as_expected_when_using_write_commands_via_opquery(int[] requestSizes, int[] expectedBatchCounts)
         {
-            var operation = new FindOperation<BsonDocument>(_collectionNamespace, BsonDocumentSerializer.Instance, _messageEncoderSettings);
-            var cursor = ExecuteOperation(operation, binding, false);
-            return ReadCursorToEnd(cursor);
+            RequireServer.Check().Supports(Feature.WriteCommands);
+            DropCollection();
+
+            using (EventContext.BeginOperation())
+            {
+                var eventCapturer = new EventCapturer().Capture<CommandStartedEvent>(e => e.CommandName == "delete" && e.OperationId == EventContext.OperationId);
+                using (var cluster = CoreTestConfiguration.CreateCluster(b => b.Subscribe(eventCapturer)))
+                using (var session = NoCoreSession.NewHandle())
+                using (var binding = new ReadWriteBindingHandle(new WritableServerBinding(cluster, session.Fork())))
+                {
+                    var requests = requestSizes.Select(size => CreateDeleteRequest(size));
+                    var operation = new BulkMixedWriteOperation(_collectionNamespace, requests, _messageEncoderSettings);
+
+                    var result = ExecuteOperation(operation, binding, async: false);
+
+                    result.RequestCount.Should().Be(requestSizes.Length);
+                    var commandStartedEvents = eventCapturer.Events.OfType<CommandStartedEvent>().ToList();
+                    var actualBatchCounts = commandStartedEvents.Select(e => e.Command["deletes"].AsBsonArray.Count).ToList();
+                    actualBatchCounts.Should().Equal(expectedBatchCounts);
+                }
+            }
+        }
+
+        [SkippableTheory]
+        [InlineData(new[] { 1 }, new[] { 1 })]
+        [InlineData(new[] { 1, 1 }, new[] { 2 })]
+        [InlineData(new[] { 8388605, 8388605 }, new[] { 2 })]
+        [InlineData(new[] { 8388605, 8388606 }, new[] { 1, 1 })]
+        [InlineData(new[] { 16777216 }, new[] { 1 })]
+        [InlineData(new[] { 16777216, 1 }, new[] { 1, 1 })]
+        [InlineData(new[] { 16777216, 16777216 }, new[] { 1, 1 })]
+        public void Execute_with_multiple_inserts_should_split_batches_as_expected_when_using_write_commands_via_opquery(int[] documentSizes, int[] expectedBatchCounts)
+        {
+            RequireServer.Check().Supports(Feature.WriteCommands);
+            DropCollection();
+
+            using (EventContext.BeginOperation())
+            {
+                var eventCapturer = new EventCapturer().Capture<CommandStartedEvent>(e => e.CommandName == "insert" && e.OperationId == EventContext.OperationId);
+                using (var cluster = CoreTestConfiguration.CreateCluster(b => b.Subscribe(eventCapturer)))
+                using (var session = NoCoreSession.NewHandle())
+                using (var binding = new ReadWriteBindingHandle(new WritableServerBinding(cluster, session.Fork())))
+                {
+                    var documents = documentSizes.Select((size, index) => CreateDocument(index + 1, size)).ToList();
+                    var requests = documents.Select(d => new InsertRequest(d)).ToList();
+                    var operation = new BulkMixedWriteOperation(_collectionNamespace, requests, _messageEncoderSettings);
+
+                    var result = ExecuteOperation(operation, binding, async: false);
+
+                    result.InsertedCount.Should().Be(documents.Count);
+                    var commandStartedEvents = eventCapturer.Events.OfType<CommandStartedEvent>().ToList();
+                    var actualBatchCounts = commandStartedEvents.Select(e => e.Command["documents"].AsBsonArray.Count).ToList();
+                    actualBatchCounts.Should().Equal(expectedBatchCounts);
+                }
+            }
+        }
+
+        [SkippableTheory]
+        [InlineData(new[] { 1 }, new[] { 1 })]
+        [InlineData(new[] { 1, 1 }, new[] { 2 })]
+        [InlineData(new[] { 8388605, 8388605 }, new[] { 2 })]
+        [InlineData(new[] { 8388605, 8388606 }, new[] { 1, 1 })]
+        [InlineData(new[] { 16777216 }, new[] { 1 })]
+        [InlineData(new[] { 16777216, 1 }, new[] { 1, 1 })]
+        [InlineData(new[] { 16777216, 16777216 }, new[] { 1, 1 })]
+        public void Execute_with_multiple_updates_should_split_batches_as_expected_when_using_write_commands_via_opquery(int[] requestSizes, int[] expectedBatchCounts)
+        {
+            RequireServer.Check().Supports(Feature.WriteCommands);
+            DropCollection();
+
+            using (EventContext.BeginOperation())
+            {
+                var eventCapturer = new EventCapturer().Capture<CommandStartedEvent>(e => e.CommandName == "update" && e.OperationId == EventContext.OperationId);
+                using (var cluster = CoreTestConfiguration.CreateCluster(b => b.Subscribe(eventCapturer)))
+                using (var session = NoCoreSession.NewHandle())
+                using (var binding = new ReadWriteBindingHandle(new WritableServerBinding(cluster, session.Fork())))
+                {
+                    var requests = requestSizes.Select(size => CreateUpdateRequest(size));
+                    var operation = new BulkMixedWriteOperation(_collectionNamespace, requests, _messageEncoderSettings);
+
+                    var result = ExecuteOperation(operation, binding, async: false);
+
+                    result.RequestCount.Should().Be(requestSizes.Length);
+                    var commandStartedEvents = eventCapturer.Events.OfType<CommandStartedEvent>().ToList();
+                    var actualBatchCounts = commandStartedEvents.Select(e => e.Command["updates"].AsBsonArray.Count).ToList();
+                    actualBatchCounts.Should().Equal(expectedBatchCounts);
+                }
+            }
+        }
+
+        // private methods
+        private DeleteRequest CreateDeleteRequest(int size)
+        {
+            var filter = new BsonDocument("filler", "");
+            var requestDocument = new BsonDocument
+            {
+                { "q", filter },
+                { "limit", 1 }
+            };
+            var fillerSize = size - requestDocument.ToBson().Length;
+            if (fillerSize > 0)
+            {
+                filter["filler"] = new string('x', fillerSize);
+            }
+            return new DeleteRequest(filter);
+        }
+
+        private BsonDocument CreateDocument(int id, int size)
+        {
+            var document = new BsonDocument { { "_id", id }, { "filler", "" } };
+            var fillerSize = size - document.ToBson().Length;
+            if (fillerSize > 0)
+            {
+                document["filler"] = new string('x', fillerSize);
+            }
+            return document;
+        }
+
+        private UpdateRequest CreateUpdateRequest(int size)
+        {
+            var filter = new BsonDocument("filler", "");
+            var update = new BsonDocument("$set", new BsonDocument("x", 1));
+            var requestDocument = new BsonDocument
+            {
+                { "q", filter },
+                { "u", update }
+            };
+            var fillerSize = size - requestDocument.ToBson().Length;
+            if (fillerSize > 0)
+            {
+                filter["filler"] = new string('x', fillerSize);
+            }
+            return new UpdateRequest(UpdateType.Update, filter, update);
         }
 
         private void EnsureTestData()
@@ -1275,5 +1409,12 @@ private void EnsureTestData()
                 BsonDocument.Parse("{_id: 5, x: 2 }"),
                 BsonDocument.Parse("{_id: 6, x: 3 }"));
         }
+        
+        private List<BsonDocument> ReadAllFromCollection(IReadBinding binding)
+        {
+            var operation = new FindOperation<BsonDocument>(_collectionNamespace, BsonDocumentSerializer.Instance, _messageEncoderSettings);
+            var cursor = ExecuteOperation(operation, binding, false);
+            return ReadCursorToEnd(cursor);
+        }
     }
 }