Skip to content

CSHARP-5887: Simplify retryable read and writes #1882

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
papafe wants to merge 5 commits into
base: main
Choose a base branch
from
Open

CSHARP-5887: Simplify retryable read and writes #1882

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
db819bb
Merge
papafe Feb 26, 2026
0fc39ef
Small fix
papafe Feb 26, 2026
aeceafc
Small fix
papafe Feb 26, 2026
843cfdf
Rename
papafe Mar 4, 2026
5446592
Corrected typo
papafe Mar 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 7 additions & 1 deletion src/MongoDB.Driver/Core/Misc/BatchableSource.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,12 +88,18 @@ public BatchableSource(IReadOnlyList<T> items, bool canBeSplit = false)
/// <param name="count">The count.</param>
/// <param name="canBeSplit">if set to <c>true</c> the batch can be split.</param>
public BatchableSource(IReadOnlyList<T> items, int offset, int count, bool canBeSplit)
: this(items, offset, count, 0, canBeSplit)
{

}

internal BatchableSource(IReadOnlyList<T> items, int offset, int count, int processedCount, bool canBeSplit)
{
_items = Ensure.IsNotNull(items, nameof(items));
_offset = Ensure.IsBetween(offset, 0, items.Count, nameof(offset));
_count = Ensure.IsBetween(count, 0, items.Count - offset, nameof(count));
_processedCount = Ensure.IsBetween(processedCount, 0, count, nameof(processedCount));
_canBeSplit = canBeSplit;
_processedCount = 0;
}

// public properties
Expand Down
16 changes: 10 additions & 6 deletions src/MongoDB.Driver/Core/Operations/AggregateOperation.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@

namespace MongoDB.Driver.Core.Operations
{
internal sealed class AggregateOperation<TResult> : IReadOperation<IAsyncCursor<TResult>>, IExecutableInRetryableReadContext<IAsyncCursor<TResult>>
internal sealed class AggregateOperation<TResult> : IReadOperation<IAsyncCursor<TResult>>, IExecutableInRetryableReadContext<IAsyncCursor<TResult>>, ICommandCreator
{
// fields
private bool? _allowDiskUse;
Expand Down Expand Up @@ -279,7 +279,7 @@ public IAsyncCursor<TResult> Execute(OperationContext operationContext, IReadBin
Ensure.IsNotNull(binding, nameof(binding));

using (BeginOperation())
using (var context = RetryableReadContext.Create(operationContext, binding, _retryRequested))
using (var context = new RetryableReadContext(binding, _retryRequested))
{
return Execute(operationContext, context);
}
Expand Down Expand Up @@ -308,7 +308,7 @@ public async Task<IAsyncCursor<TResult>> ExecuteAsync(OperationContext operation
Ensure.IsNotNull(binding, nameof(binding));

using (BeginOperation())
using (var context = await RetryableReadContext.CreateAsync(operationContext, binding, _retryRequested).ConfigureAwait(false))
using (var context = new RetryableReadContext(binding, _retryRequested))
{
return await ExecuteAsync(operationContext, context).ConfigureAwait(false);
}
Expand All @@ -331,7 +331,7 @@ public async Task<IAsyncCursor<TResult>> ExecuteAsync(OperationContext operation
}
}

internal BsonDocument CreateCommand(OperationContext operationContext, ICoreSession session, ConnectionDescription connectionDescription)
public BsonDocument CreateCommand(OperationContext operationContext, ICoreSession session, ConnectionDescription connectionDescription)
{
var readConcern = ReadConcernHelper.GetReadConcernForCommand(session, connectionDescription, _readConcern);
var command = new BsonDocument
Expand Down Expand Up @@ -362,9 +362,13 @@ internal BsonDocument CreateCommand(OperationContext operationContext, ICoreSess
private ReadCommandOperation<AggregateResult> CreateOperation(OperationContext operationContext, RetryableReadContext context)
{
var databaseNamespace = _collectionNamespace == null ? _databaseNamespace : _collectionNamespace.DatabaseNamespace;
var command = CreateCommand(operationContext, context.Binding.Session, context.Channel.ConnectionDescription);
var serializer = new AggregateResultDeserializer(_resultSerializer);
return new ReadCommandOperation<AggregateResult>(databaseNamespace, command, serializer, MessageEncoderSettings, OperationName)
return new ReadCommandOperation<AggregateResult>(
databaseNamespace,
this,
serializer,
MessageEncoderSettings,
OperationName)
{
RetryRequested = _retryRequested // might be overridden by retryable read context
};
Expand Down
4 changes: 2 additions & 2 deletions src/MongoDB.Driver/Core/Operations/BulkMixedWriteOperation.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -145,7 +145,7 @@ public WriteConcern WriteConcern
public BulkWriteOperationResult Execute(OperationContext operationContext, IWriteBinding binding)
{
using (BeginOperation())
using (var context = RetryableWriteContext.Create(operationContext, binding, IsOperationRetryable()))
using (var context = new RetryableWriteContext(binding, IsOperationRetryable()))
{
EnsureHintIsSupportedIfAnyRequestHasHint();
var helper = new BatchHelper(_requests, _isOrdered, _writeConcern);
Expand All @@ -160,7 +160,7 @@ public BulkWriteOperationResult Execute(OperationContext operationContext, IWrit
public async Task<BulkWriteOperationResult> ExecuteAsync(OperationContext operationContext, IWriteBinding binding)
{
using (BeginOperation())
using (var context = await RetryableWriteContext.CreateAsync(operationContext, binding, IsOperationRetryable()).ConfigureAwait(false))
using (var context = new RetryableWriteContext(binding, IsOperationRetryable()))
{
EnsureHintIsSupportedIfAnyRequestHasHint();
var helper = new BatchHelper(_requests, _isOrdered, _writeConcern);
Expand Down
6 changes: 3 additions & 3 deletions src/MongoDB.Driver/Core/Operations/BulkUnmixedWriteOperationBase.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -131,7 +131,7 @@ public BulkWriteOperationResult Execute(OperationContext operationContext, Retry
public BulkWriteOperationResult Execute(OperationContext operationContext, IWriteBinding binding)
{
using (BeginOperation())
using (var context = RetryableWriteContext.Create(operationContext, binding, IsOperationRetryable()))
using (var context = new RetryableWriteContext(binding, IsOperationRetryable()))
{
return Execute(operationContext, context);
}
Expand All @@ -147,7 +147,7 @@ public Task<BulkWriteOperationResult> ExecuteAsync(OperationContext operationCon
public async Task<BulkWriteOperationResult> ExecuteAsync(OperationContext operationContext, IWriteBinding binding)
{
using (BeginOperation())
using (var context = await RetryableWriteContext.CreateAsync(operationContext, binding, IsOperationRetryable()).ConfigureAwait(false))
using (var context = new RetryableWriteContext(binding, IsOperationRetryable()))
{
return await ExecuteAsync(operationContext, context).ConfigureAwait(false);
}
Expand Down Expand Up @@ -286,7 +286,7 @@ public BulkWriteOperationResult CreateFinalResultOrThrow(IChannelHandle channel)
{
var combiner = new BulkWriteBatchResultCombiner(_batchResults, _writeConcern.IsAcknowledged);
var remainingRequests = _requests.GetUnprocessedItems();
return combiner.CreateResultOrThrowIfHasErrors(channel.ConnectionDescription.ConnectionId, remainingRequests);
return combiner.CreateResultOrThrowIfHasErrors(channel?.ConnectionDescription.ConnectionId, remainingRequests);
}

// private methods
Expand Down
8 changes: 4 additions & 4 deletions src/MongoDB.Driver/Core/Operations/ChangeStreamOperation.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -266,7 +266,7 @@ public IChangeStreamCursor<TResult> Execute(OperationContext operationContext, I
IAsyncCursor<RawBsonDocument> cursor;
ICursorBatchInfo cursorBatchInfo;
BsonTimestamp initialOperationTime;
using (var context = RetryableReadContext.Create(operationContext, binding, _retryRequested))
using (var context = new RetryableReadContext(binding, _retryRequested))
{
cursor = ExecuteAggregateOperation(operationContext, context);
cursorBatchInfo = (ICursorBatchInfo)cursor;
Expand Down Expand Up @@ -301,7 +301,7 @@ public async Task<IChangeStreamCursor<TResult>> ExecuteAsync(OperationContext op
IAsyncCursor<RawBsonDocument> cursor;
ICursorBatchInfo cursorBatchInfo;
BsonTimestamp initialOperationTime;
using (var context = await RetryableReadContext.CreateAsync(operationContext, binding, _retryRequested).ConfigureAwait(false))
using (var context = new RetryableReadContext(binding, _retryRequested))
{
cursor = await ExecuteAggregateOperationAsync(operationContext, context).ConfigureAwait(false);
cursorBatchInfo = (ICursorBatchInfo)cursor;
Expand All @@ -326,7 +326,7 @@ public async Task<IChangeStreamCursor<TResult>> ExecuteAsync(OperationContext op
/// <inheritdoc />
public IAsyncCursor<RawBsonDocument> Resume(OperationContext operationContext, IReadBinding binding)
{
using (var context = RetryableReadContext.Create(operationContext, binding, retryRequested: false))
using (var context = new RetryableReadContext(binding, retryRequested: false))
{
return ExecuteAggregateOperation(operationContext, context);
}
Expand All @@ -335,7 +335,7 @@ public IAsyncCursor<RawBsonDocument> Resume(OperationContext operationContext, I
/// <inheritdoc />
public async Task<IAsyncCursor<RawBsonDocument>> ResumeAsync(OperationContext operationContext, IReadBinding binding)
{
using (var context = await RetryableReadContext.CreateAsync(operationContext, binding, retryRequested: false).ConfigureAwait(false))
using (var context = new RetryableReadContext(binding, retryRequested: false))
{
return await ExecuteAggregateOperationAsync(operationContext, context).ConfigureAwait(false);
}
Expand Down
8 changes: 4 additions & 4 deletions src/MongoDB.Driver/Core/Operations/ClientBulkWriteOperation.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -96,7 +96,7 @@ protected override IEnumerable<BatchableCommandMessageSection> CreateCommandPayl
var bulkWriteResults = new BulkWriteRawResult();
while (true)
{
using var context = RetryableWriteContext.Create(operationContext, binding, GetEffectiveRetryRequested());
using var context = new RetryableWriteContext(binding, GetEffectiveRetryRequested());
BsonDocument serverResponse = null;
try
{
Expand All @@ -112,7 +112,7 @@ protected override IEnumerable<BatchableCommandMessageSection> CreateCommandPayl
bulkWriteResults.TopLevelException = commandException;
serverResponse = commandException.Result;
}
catch (Exception exception)
catch (Exception exception) when (context.Channel is not null)
{
bulkWriteResults.TopLevelException = exception;
}
Expand Down Expand Up @@ -154,7 +154,7 @@ protected override IEnumerable<BatchableCommandMessageSection> CreateCommandPayl
var bulkWriteResults = new BulkWriteRawResult();
while (true)
{
using var context = RetryableWriteContext.Create(operationContext, binding, GetEffectiveRetryRequested());
using var context = new RetryableWriteContext(binding, GetEffectiveRetryRequested());
BsonDocument serverResponse = null;
try
{
Expand All @@ -170,7 +170,7 @@ protected override IEnumerable<BatchableCommandMessageSection> CreateCommandPayl
bulkWriteResults.TopLevelException = commandException;
serverResponse = commandException.Result;
}
catch (Exception exception)
catch (Exception exception) when (context.Channel is not null)
{
bulkWriteResults.TopLevelException = exception;
}
Expand Down
11 changes: 10 additions & 1 deletion src/MongoDB.Driver/Core/Operations/CommandOperationBase.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ protected CommandOperationBase(
MessageEncoderSettings messageEncoderSettings)
{
_databaseNamespace = Ensure.IsNotNull(databaseNamespace, nameof(databaseNamespace));
_command = Ensure.IsNotNull(command, nameof(command));
_command = command; //can be null
Copy link
Contributor Author

@papafe papafe Feb 17, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is used so that we can modify the command once the ReadCommandOperation has been created.
This is done because some operations need to have the operationContext to properly create the command, but now the operationContext is not available until later.
I'm not a fan of this implementation, to be honest.
We got two other possibilities in my opinion:

  1. We make ReadCommandOperation equivalent to WriteCommandOperation, so it does not concern with retryability. Then we make a new class, RetryableReadCommandOperation that has retryability and uses ReadCommandOperation inside
  2. Instead of having the possibility of setting the command, we can add a method to modify it when we got the operation context.

I think nr1 is more desirable, but requires further much more refactoring.

_resultSerializer = Ensure.IsNotNull(resultSerializer, nameof(resultSerializer));
_messageEncoderSettings = messageEncoderSettings;
}
Expand Down Expand Up @@ -146,6 +146,15 @@ protected async Task<TCommandResult> ExecuteProtocolAsync(
}
}

/// <summary>
/// Sets the command to be executed. This is used by derived classes that build commands dynamically.
/// </summary>
/// <param name="command">The command.</param>
protected void SetCommand(BsonDocument command)
{
_command = Ensure.IsNotNull(command, nameof(command));
}

private BsonDocument GetEffectiveAdditionalOptions()
{
if (_additionalOptions == null && _comment == null)
Expand Down
14 changes: 9 additions & 5 deletions src/MongoDB.Driver/Core/Operations/CountOperation.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@

namespace MongoDB.Driver.Core.Operations
{
internal sealed class CountOperation : IReadOperation<long>, IExecutableInRetryableReadContext<long>
internal sealed class CountOperation : IReadOperation<long>, IExecutableInRetryableReadContext<long>, ICommandCreator
{
private Collation _collation;
private readonly CollectionNamespace _collectionNamespace;
Expand Down Expand Up @@ -133,7 +133,7 @@ public long Execute(OperationContext operationContext, IReadBinding binding)
Ensure.IsNotNull(binding, nameof(binding));

using (BeginOperation())
using (var context = RetryableReadContext.Create(operationContext, binding, _retryRequested))
using (var context = new RetryableReadContext(binding, _retryRequested))
{
return Execute(operationContext, context);
}
Expand All @@ -151,7 +151,7 @@ public async Task<long> ExecuteAsync(OperationContext operationContext, IReadBin
Ensure.IsNotNull(binding, nameof(binding));

using (BeginOperation())
using (var context = await RetryableReadContext.CreateAsync(operationContext, binding, _retryRequested).ConfigureAwait(false))
using (var context = new RetryableReadContext(binding, _retryRequested))
{
return await ExecuteAsync(operationContext, context).ConfigureAwait(false);
}
Expand All @@ -168,8 +168,12 @@ public async Task<long> ExecuteAsync(OperationContext operationContext, Retryabl

private ReadCommandOperation<BsonDocument> CreateOperation(OperationContext operationContext, RetryableReadContext context)
{
var command = CreateCommand(operationContext, context.Binding.Session, context.Channel.ConnectionDescription);
return new ReadCommandOperation<BsonDocument>(_collectionNamespace.DatabaseNamespace, command, BsonDocumentSerializer.Instance, _messageEncoderSettings, OperationName)
return new ReadCommandOperation<BsonDocument>(
_collectionNamespace.DatabaseNamespace,
this,
BsonDocumentSerializer.Instance,
_messageEncoderSettings,
OperationName)
{
RetryRequested = _retryRequested // might be overridden by retryable read context
};
Expand Down
16 changes: 10 additions & 6 deletions src/MongoDB.Driver/Core/Operations/DistinctOperation.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@

namespace MongoDB.Driver.Core.Operations
{
internal sealed class DistinctOperation<TValue> : IReadOperation<IAsyncCursor<TValue>>
internal sealed class DistinctOperation<TValue> : IReadOperation<IAsyncCursor<TValue>>, ICommandCreator
{
private Collation _collation;
private CollectionNamespace _collectionNamespace;
Expand Down Expand Up @@ -111,7 +111,7 @@ public IAsyncCursor<TValue> Execute(OperationContext operationContext, IReadBind
Ensure.IsNotNull(binding, nameof(binding));

using (BeginOperation())
using (var context = RetryableReadContext.Create(operationContext, binding, _retryRequested))
using (var context = new RetryableReadContext(binding, _retryRequested))
{
var operation = CreateOperation(operationContext, context);
var result = operation.Execute(operationContext, context);
Expand All @@ -127,7 +127,7 @@ public async Task<IAsyncCursor<TValue>> ExecuteAsync(OperationContext operationC
Ensure.IsNotNull(binding, nameof(binding));

using (BeginOperation())
using (var context = await RetryableReadContext.CreateAsync(operationContext, binding, _retryRequested).ConfigureAwait(false))
using (var context = new RetryableReadContext(binding, _retryRequested))
{
var operation = CreateOperation(operationContext, context);
var result = await operation.ExecuteAsync(operationContext, context).ConfigureAwait(false);
Expand All @@ -153,14 +153,18 @@ public BsonDocument CreateCommand(OperationContext operationContext, ICoreSessio
};
}

private EventContext.OperationNameDisposer BeginOperation() => EventContext.BeginOperation(OperationName);
private EventContext.OperationIdDisposer BeginOperation() => EventContext.BeginOperation(null, OperationName);
Copy link
Contributor Author

@papafe papafe Feb 26, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is something strange. If we use this overloads, then the name of the command gets set properly.
This happens because before this PR, we were connecting to the server when creating the RetryableContext, and at that time the command name was set properly. Now this is done inside the retryable operation executor, and by that time (for example) ReadCommandOperation.Execute is called and that contains BeginOperation(null, null), that puts the command name to null. This means that the command name in ClusterSelectingServerEvent is null.
We need to understand if this is the way we want to go, and so do it for other read operations as well.


private ReadCommandOperation<DistinctResult> CreateOperation(OperationContext operationContext, RetryableReadContext context)
{
var command = CreateCommand(operationContext, context.Binding.Session, context.Channel.ConnectionDescription);
var serializer = new DistinctResultDeserializer(_valueSerializer);

return new ReadCommandOperation<DistinctResult>(_collectionNamespace.DatabaseNamespace, command, serializer, _messageEncoderSettings, OperationName)
return new ReadCommandOperation<DistinctResult>(
_collectionNamespace.DatabaseNamespace,
this,
serializer,
_messageEncoderSettings,
OperationName)
{
RetryRequested = _retryRequested // might be overridden by retryable read context
};
Expand Down
8 changes: 4 additions & 4 deletions src/MongoDB.Driver/Core/Operations/EstimatedDocumentCountOperation.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
/* Copyright 2021-present MongoDB Inc.
/* Copyright 2021-present MongoDB Inc.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand Down Expand Up @@ -73,7 +73,7 @@ public long Execute(OperationContext operationContext, IReadBinding binding)
Ensure.IsNotNull(binding, nameof(binding));

using (BeginOperation())
using (var context = RetryableReadContext.Create(operationContext, binding, _retryRequested))
using (var context = new RetryableReadContext(binding, _retryRequested))
{
var operation = CreateCountOperation();

Expand All @@ -86,15 +86,15 @@ public async Task<long> ExecuteAsync(OperationContext operationContext, IReadBin
Ensure.IsNotNull(binding, nameof(binding));

using (BeginOperation())
using (var context = RetryableReadContext.Create(operationContext, binding, _retryRequested))
using (var context = new RetryableReadContext(binding, _retryRequested))
{
var operation = CreateCountOperation();

return await operation.ExecuteAsync(operationContext, context).ConfigureAwait(false);
}
}

private EventContext.OperationNameDisposer BeginOperation() => EventContext.BeginOperation(OperationName);
private EventContext.OperationIdDisposer BeginOperation() => EventContext.BeginOperation(null, OperationName);

private IExecutableInRetryableReadContext<long> CreateCountOperation()
{
Expand Down
Loading