CSHARP-4900: Uploading a duplicate file larger than the original one causes errors in downloading the original file (#1244)

Author: sanych-sun

src/MongoDB.Driver.GridFS/GridFSForwardOnlyUploadStream.cs

@@ -23,7 +23,6 @@
 using MongoDB.Bson.Serialization;
 using MongoDB.Driver.Core.Bindings;
 using MongoDB.Driver.Core.Operations;
-using MongoDB.Shared;
 
 namespace MongoDB.Driver.GridFS
 {
@@ -374,6 +373,32 @@ protected override void Dispose(bool disposing)
             base.Dispose(disposing);
         }
 
+        private void ExecuteOrSetAbortedOnException(Action action)
+        {
+            try
+            {
+                action();
+            }
+            catch
+            {
+                _aborted = true;
+                throw;
+            }
+        }
+
+        private async Task ExecuteOrSetAbortedOnExceptionAsync(Func<Task> action)
+        {
+            try
+            {
+                await action().ConfigureAwait(false);
+            }
+            catch
+            {
+                _aborted = true;
+                throw;
+            }
+        }
+
         private IMongoCollection<BsonDocument> GetChunksCollection()
         {
             return GetCollection("chunks");
@@ -476,34 +501,38 @@ private void TruncateFinalChunk()
         }
 
         private void WriteBatch(CancellationToken cancellationToken)
-        {
-            var chunksCollection = GetChunksCollection();
-            var chunkDocuments = CreateWriteBatchChunkDocuments();
-            chunksCollection.InsertMany(chunkDocuments, cancellationToken: cancellationToken);
-            _batch.Clear();
-        }
+            => ExecuteOrSetAbortedOnException(() =>
+            {
+                var chunksCollection = GetChunksCollection();
+                var chunkDocuments = CreateWriteBatchChunkDocuments();
+                chunksCollection.InsertMany(chunkDocuments, cancellationToken: cancellationToken);
+                _batch.Clear();
+            });
 
-        private async Task WriteBatchAsync(CancellationToken cancellationToken)
-        {
-            var chunksCollection = GetChunksCollection();
-            var chunkDocuments = CreateWriteBatchChunkDocuments();
-            await chunksCollection.InsertManyAsync(chunkDocuments, cancellationToken: cancellationToken).ConfigureAwait(false);
-            _batch.Clear();
-        }
+        private Task WriteBatchAsync(CancellationToken cancellationToken)
+            => ExecuteOrSetAbortedOnExceptionAsync(async () =>
+            {
+                var chunksCollection = GetChunksCollection();
+                var chunkDocuments = CreateWriteBatchChunkDocuments();
+                await chunksCollection.InsertManyAsync(chunkDocuments, cancellationToken: cancellationToken).ConfigureAwait(false);
+                _batch.Clear();
+            });
 
         private void WriteFilesCollectionDocument(CancellationToken cancellationToken)
-        {
-            var filesCollection = GetFilesCollection();
-            var filesCollectionDocument = CreateFilesCollectionDocument();
-            filesCollection.InsertOne(filesCollectionDocument, cancellationToken: cancellationToken);
-        }
+            => ExecuteOrSetAbortedOnException(() =>
+            {
+                var filesCollection = GetFilesCollection();
+                var filesCollectionDocument = CreateFilesCollectionDocument();
+                filesCollection.InsertOne(filesCollectionDocument, cancellationToken: cancellationToken);
+            });
 
-        private async Task WriteFilesCollectionDocumentAsync(CancellationToken cancellationToken)
-        {
-            var filesCollection = GetFilesCollection();
-            var filesCollectionDocument = CreateFilesCollectionDocument();
-            await filesCollection.InsertOneAsync(filesCollectionDocument, cancellationToken: cancellationToken).ConfigureAwait(false);
-        }
+        private Task WriteFilesCollectionDocumentAsync(CancellationToken cancellationToken)
+            => ExecuteOrSetAbortedOnExceptionAsync(() =>
+            {
+                var filesCollection = GetFilesCollection();
+                var filesCollectionDocument = CreateFilesCollectionDocument();
+                return filesCollection.InsertOneAsync(filesCollectionDocument, cancellationToken: cancellationToken);
+            });
 
         private void WriteFinalBatch(CancellationToken cancellationToken)
         {

tests/MongoDB.Driver.GridFS.Tests/GridFSUploadStreamTests.cs

@@ -15,8 +15,11 @@
 
 using System;
 using System.IO;
+using System.Linq;
 using System.Threading;
+using System.Threading.Tasks;
 using FluentAssertions;
+using MongoDB.Bson;
 using MongoDB.TestHelpers.XunitExtensions;
 using MongoDB.Driver.Tests;
 using Xunit;
@@ -71,6 +74,43 @@ public void Flush_should_not_throw(
             action.ShouldNotThrow();
         }
 
+        [Theory]
+        [ParameterAttributeData]
+        public async Task Upload_of_duplicate_file_should_not_invalidate_existing_data(
+            [Values(false, true)] bool async)
+        {
+            var content1 = Enumerable.Repeat((byte)1, 20).ToArray();
+            var content2 = Enumerable.Repeat((byte)2, 100).ToArray();
+
+            var fileId = ObjectId.GenerateNewId();
+            var fileName = "filename";
+            var uploadOptions = new GridFSUploadOptions
+            {
+                ChunkSizeBytes = 10,
+                BatchSize = 5,
+            };
+
+            var bucket = CreateBucket();
+            if (async)
+            {
+                await bucket.UploadFromBytesAsync(fileId, fileName, content1, uploadOptions);
+            }
+            else
+            {
+                bucket.UploadFromBytes(fileId, fileName, content1, uploadOptions);
+            }
+
+            var exception = async
+                ? await Record.ExceptionAsync(() =>  bucket.UploadFromBytesAsync(fileId, fileName, content2, uploadOptions))
+                : Record.Exception(() => bucket.UploadFromBytes(fileId, fileName, content2, uploadOptions));
+            exception.Should().BeOfType<MongoBulkWriteException<BsonDocument>>();
+
+            var uploadedContent = async
+                ? await bucket.DownloadAsBytesAsync(fileId)
+                : bucket.DownloadAsBytes(fileId);
+            uploadedContent.Should().Equal(content1);
+        }
+
         // private methods
         private IGridFSBucket CreateBucket()
         {