Skip to content

Commit

Permalink
Add option to disable concurrency detection
Browse files Browse the repository at this point in the history
  • Loading branch information
@roji
roji committed Feb 14, 2021
1 parent ede9db2 commit 46a33c2
Show file tree
Hide file tree
Showing 22 changed files with 549 additions and 256 deletions.
101 changes: 60 additions & 41 deletions ...e.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.QueryingEnumerable.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ private sealed class QueryingEnumerable<T> : IEnumerable<T>, IAsyncEnumerable<T>
private readonly string _partitionKey;
private readonly IDiagnosticsLogger<DbLoggerCategory.Query> _queryLogger;
private readonly bool _standAloneStateManager;
private readonly bool _concurrencyDetectionEnabled;

public QueryingEnumerable(
CosmosQueryContext cosmosQueryContext,
Expand All @@ -44,7 +45,8 @@ public QueryingEnumerable(
Func<CosmosQueryContext, JObject, T> shaper,
Type contextType,
string partitionKeyFromExtension,
bool standAloneStateManager)
bool standAloneStateManager,
bool concurrencyDetectionEnabled)
{
_cosmosQueryContext = cosmosQueryContext;
_sqlExpressionFactory = sqlExpressionFactory;
Expand All @@ -54,6 +56,7 @@ public QueryingEnumerable(
_contextType = contextType;
_queryLogger = cosmosQueryContext.QueryLogger;
_standAloneStateManager = standAloneStateManager;
_concurrencyDetectionEnabled = concurrencyDetectionEnabled;

var partitionKey = selectExpression.GetPartitionKey(cosmosQueryContext.ParameterValues);
if (partitionKey != null && partitionKeyFromExtension != null && partitionKeyFromExtension != partitionKey)
Expand Down Expand Up @@ -113,6 +116,7 @@ private sealed class Enumerator : IEnumerator<T>
private readonly string _partitionKey;
private readonly IDiagnosticsLogger<DbLoggerCategory.Query> _queryLogger;
private readonly bool _standAloneStateManager;
private readonly IConcurrencyDetector _concurrencyDetector;

private IEnumerator<JObject> _enumerator;

Expand All @@ -126,6 +130,10 @@ public Enumerator(QueryingEnumerable<T> queryingEnumerable)
_partitionKey = queryingEnumerable._partitionKey;
_queryLogger = queryingEnumerable._queryLogger;
_standAloneStateManager = queryingEnumerable._standAloneStateManager;

_concurrencyDetector = queryingEnumerable._concurrencyDetectionEnabled
? _cosmosQueryContext.ConcurrencyDetector
: null;
}

public T Current { get; private set; }
Expand All @@ -135,41 +143,44 @@ object IEnumerator.Current

public bool MoveNext()
{
_concurrencyDetector?.EnterCriticalSection();

try
{
using (_cosmosQueryContext.ConcurrencyDetector.EnterCriticalSection())
if (_enumerator == null)
{
if (_enumerator == null)
{
var sqlQuery = _queryingEnumerable.GenerateQuery();
var sqlQuery = _queryingEnumerable.GenerateQuery();

EntityFrameworkEventSource.Log.QueryExecuting();
EntityFrameworkEventSource.Log.QueryExecuting();

_enumerator = _cosmosQueryContext.CosmosClient
.ExecuteSqlQuery(
_selectExpression.Container,
_partitionKey,
sqlQuery)
.GetEnumerator();
_cosmosQueryContext.InitializeStateManager(_standAloneStateManager);
}
_enumerator = _cosmosQueryContext.CosmosClient
.ExecuteSqlQuery(
_selectExpression.Container,
_partitionKey,
sqlQuery)
.GetEnumerator();
_cosmosQueryContext.InitializeStateManager(_standAloneStateManager);
}

var hasNext = _enumerator.MoveNext();
var hasNext = _enumerator.MoveNext();

Current
= hasNext
? _shaper(_cosmosQueryContext, _enumerator.Current)
: default;
Current
= hasNext
? _shaper(_cosmosQueryContext, _enumerator.Current)
: default;

return hasNext;
}
return hasNext;
}
catch (Exception exception)
{
_queryLogger.QueryIterationFailed(_contextType, exception);

throw;
}
finally
{
_concurrencyDetector?.ExitCriticalSection();
}
}

public void Dispose()
Expand All @@ -193,6 +204,7 @@ private sealed class AsyncEnumerator : IAsyncEnumerator<T>
private readonly IDiagnosticsLogger<DbLoggerCategory.Query> _queryLogger;
private readonly bool _standAloneStateManager;
private readonly CancellationToken _cancellationToken;
private readonly IConcurrencyDetector _concurrencyDetector;

private IAsyncEnumerator<JObject> _enumerator;

Expand All @@ -207,47 +219,54 @@ public AsyncEnumerator(QueryingEnumerable<T> queryingEnumerable, CancellationTok
_queryLogger = queryingEnumerable._queryLogger;
_standAloneStateManager = queryingEnumerable._standAloneStateManager;
_cancellationToken = cancellationToken;

_concurrencyDetector = queryingEnumerable._concurrencyDetectionEnabled
? _cosmosQueryContext.ConcurrencyDetector
: null;
}

public T Current { get; private set; }

public async ValueTask<bool> MoveNextAsync()
{
_concurrencyDetector?.EnterCriticalSection();

try
{
using (_cosmosQueryContext.ConcurrencyDetector.EnterCriticalSection())
if (_enumerator == null)
{
if (_enumerator == null)
{
var sqlQuery = _queryingEnumerable.GenerateQuery();
var sqlQuery = _queryingEnumerable.GenerateQuery();

EntityFrameworkEventSource.Log.QueryExecuting();
EntityFrameworkEventSource.Log.QueryExecuting();

_enumerator = _cosmosQueryContext.CosmosClient
.ExecuteSqlQueryAsync(
_selectExpression.Container,
_partitionKey,
sqlQuery)
.GetAsyncEnumerator(_cancellationToken);
_cosmosQueryContext.InitializeStateManager(_standAloneStateManager);
}
_enumerator = _cosmosQueryContext.CosmosClient
.ExecuteSqlQueryAsync(
_selectExpression.Container,
_partitionKey,
sqlQuery)
.GetAsyncEnumerator(_cancellationToken);
_cosmosQueryContext.InitializeStateManager(_standAloneStateManager);
}

var hasNext = await _enumerator.MoveNextAsync().ConfigureAwait(false);
var hasNext = await _enumerator.MoveNextAsync().ConfigureAwait(false);

Current
= hasNext
? _shaper(_cosmosQueryContext, _enumerator.Current)
: default;
Current
= hasNext
? _shaper(_cosmosQueryContext, _enumerator.Current)
: default;

return hasNext;
}
return hasNext;
}
catch (Exception exception)
{
_queryLogger.QueryIterationFailed(_contextType, exception);

throw;
}
finally
{
_concurrencyDetector?.ExitCriticalSection();
}
}

public ValueTask DisposeAsync()
Expand Down
109 changes: 66 additions & 43 deletions .../Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.ReadItemQueryingEnumerable.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,20 +35,23 @@ private sealed class ReadItemQueryingEnumerable<T> : IEnumerable<T>, IAsyncEnume
private readonly Type _contextType;
private readonly IDiagnosticsLogger<DbLoggerCategory.Query> _queryLogger;
private readonly bool _standAloneStateManager;
private readonly bool _concurrencyDetectionEnabled;

public ReadItemQueryingEnumerable(
CosmosQueryContext cosmosQueryContext,
ReadItemExpression readItemExpression,
Func<CosmosQueryContext, JObject, T> shaper,
Type contextType,
bool standAloneStateManager)
bool standAloneStateManager,
bool concurrencyDetectionEnabled)
{
_cosmosQueryContext = cosmosQueryContext;
_readItemExpression = readItemExpression;
_shaper = shaper;
_contextType = contextType;
_queryLogger = _cosmosQueryContext.QueryLogger;
_standAloneStateManager = standAloneStateManager;
_concurrencyDetectionEnabled = concurrencyDetectionEnabled;
}

public IAsyncEnumerator<T> GetAsyncEnumerator(CancellationToken cancellationToken = default)
Expand Down Expand Up @@ -173,6 +176,7 @@ private sealed class Enumerator : IEnumerator<T>, IAsyncEnumerator<T>
private readonly Type _contextType;
private readonly IDiagnosticsLogger<DbLoggerCategory.Query> _queryLogger;
private readonly bool _standAloneStateManager;
private readonly bool _concurrencyDetectionEnabled;
private readonly ReadItemQueryingEnumerable<T> _readItemEnumerable;
private readonly CancellationToken _cancellationToken;

Expand All @@ -187,6 +191,7 @@ public Enumerator(ReadItemQueryingEnumerable<T> readItemEnumerable, Cancellation
_contextType = readItemEnumerable._contextType;
_queryLogger = readItemEnumerable._queryLogger;
_standAloneStateManager = readItemEnumerable._standAloneStateManager;
_concurrencyDetectionEnabled = readItemEnumerable._concurrencyDetectionEnabled;
_readItemEnumerable = readItemEnumerable;
_cancellationToken = cancellationToken;
}
Expand All @@ -198,82 +203,100 @@ object IEnumerator.Current

public bool MoveNext()
{
if (_concurrencyDetectionEnabled)
{
_cosmosQueryContext.ConcurrencyDetector.EnterCriticalSection();
}

try
{
using (_cosmosQueryContext.ConcurrencyDetector.EnterCriticalSection())
if (!_hasExecuted)
{
if (!_hasExecuted)
if (!_readItemEnumerable.TryGetResourceId(out var resourceId))
{
if (!_readItemEnumerable.TryGetResourceId(out var resourceId))
{
throw new InvalidOperationException(CosmosStrings.ResourceIdMissing);
}
throw new InvalidOperationException(CosmosStrings.ResourceIdMissing);
}

if (!_readItemEnumerable.TryGetPartitionId(out var partitionKey))
{
throw new InvalidOperationException(CosmosStrings.PartitionKeyMissing);
}
if (!_readItemEnumerable.TryGetPartitionId(out var partitionKey))
{
throw new InvalidOperationException(CosmosStrings.PartitionKeyMissing);
}

EntityFrameworkEventSource.Log.QueryExecuting();
EntityFrameworkEventSource.Log.QueryExecuting();

_item = _cosmosQueryContext.CosmosClient.ExecuteReadItem(
_readItemExpression.Container,
partitionKey,
resourceId);
_item = _cosmosQueryContext.CosmosClient.ExecuteReadItem(
_readItemExpression.Container,
partitionKey,
resourceId);

return ShapeResult();
}

return false;
return ShapeResult();
}

return false;
}
catch (Exception exception)
{
_queryLogger.QueryIterationFailed(_contextType, exception);

throw;
}
finally
{
if (_concurrencyDetectionEnabled)
{
_cosmosQueryContext.ConcurrencyDetector.ExitCriticalSection();
}
}
}

public async ValueTask<bool> MoveNextAsync()
{
if (_concurrencyDetectionEnabled)
{
_cosmosQueryContext.ConcurrencyDetector.EnterCriticalSection();
}

try
{
using (_cosmosQueryContext.ConcurrencyDetector.EnterCriticalSection())
if (!_hasExecuted)
{
if (!_hasExecuted)
if (!_readItemEnumerable.TryGetResourceId(out var resourceId))
{
throw new InvalidOperationException(CosmosStrings.ResourceIdMissing);
}

if (!_readItemEnumerable.TryGetPartitionId(out var partitionKey))
{
if (!_readItemEnumerable.TryGetResourceId(out var resourceId))
{
throw new InvalidOperationException(CosmosStrings.ResourceIdMissing);
}

if (!_readItemEnumerable.TryGetPartitionId(out var partitionKey))
{
throw new InvalidOperationException(CosmosStrings.PartitionKeyMissing);
}

EntityFrameworkEventSource.Log.QueryExecuting();

_item = await _cosmosQueryContext.CosmosClient.ExecuteReadItemAsync(
_readItemExpression.Container,
partitionKey,
resourceId,
_cancellationToken)
.ConfigureAwait(false);

return ShapeResult();
throw new InvalidOperationException(CosmosStrings.PartitionKeyMissing);
}

return false;
EntityFrameworkEventSource.Log.QueryExecuting();

_item = await _cosmosQueryContext.CosmosClient.ExecuteReadItemAsync(
_readItemExpression.Container,
partitionKey,
resourceId,
_cancellationToken)
.ConfigureAwait(false);

return ShapeResult();
}

return false;
}
catch (Exception exception)
{
_queryLogger.QueryIterationFailed(_contextType, exception);

throw;
}
finally
{
if (_concurrencyDetectionEnabled)
{
_cosmosQueryContext.ConcurrencyDetector.ExitCriticalSection();
}
}
}

public void Dispose()
Expand Down
Loading

0 comments on commit 46a33c2

Please sign in to comment.