private async ValueTask <TResult> InternalExecuteAndPropagateCancellationAsync <TState, TResult>(
TState state,
Func <TState, CancellationToken, ValueTask <TResult> > executeAsync,
CancellationToken cancellationToken,
bool isConnectionMonitoringQuery)
{
Invariant.Require(cancellationToken.CanBeCanceled);
using var _ = await this.AcquireConnectionLockIfNeeded(isConnectionMonitoringQuery).ConfigureAwait(false);
// Note: for now we cannot pass cancellationToken to PrepareAsync() because this will break on Postgres which
// is the only db we currently support that needs Prepare currently. See https://github.com/npgsql/npgsql/issues/4209
await this.PrepareIfNeededAsync(CancellationToken.None).ConfigureAwait(false);
try
{
return(await executeAsync(state, cancellationToken).ConfigureAwait(false));
}
catch (Exception ex)
// Canceled SQL operations throw SqlException/InvalidOperationException instead of OCE.
// That means that downstream operations end up faulted instead of canceled. We
// wrap with OCE here to correctly propagate cancellation
when(cancellationToken.IsCancellationRequested && this._connection.IsCommandCancellationException(ex))
{
throw new OperationCanceledException(
"Command was canceled",
ex,
cancellationToken
);
}
}
private bool StartMonitorWorkerIfNeededNoLock()
{
Invariant.Require(this._state != State.Disposed);
// never monitor external connections
if (this._isExternallyOwnedConnection)
{
return(false);
}
// If we're in the active state, we already have a worker. If we're not in the idle
// state, we're not supposed to be running
if (this._state != State.Idle)
{
return(false);
}
// skip if there's nothing to do
if (this._keepaliveCadence.IsInfinite && !this.HasRegisteredMonitoringHandlesNoLock)
{
return(false);
}
this._monitorStateChangedTokenSource = new CancellationTokenSource();
// Set up the task as a continuation on the previous task to avoid concurrency in the case where the previous
// one is spinning down. If we change states in rapid succession we could end up with multiple tasks queued up
// but this shouldn't matter since when the active one ultimately stops all the others will follow in rapid succession
this._monitoringWorkerTask = this._monitoringWorkerTask
.ContinueWith((_, state) => ((ConnectionMonitor)state).MonitorWorkerLoop(), state: this)
.Unwrap();
this._state = State.Active;
return(true);
}
private async ValueTask <TResult> InternalExecuteAndPropagateCancellationAsync <TState, TResult>(
TState state,
Func <TState, CancellationToken, ValueTask <TResult> > executeAsync,
CancellationToken cancellationToken,
bool isConnectionMonitoringQuery)
{
Invariant.Require(cancellationToken.CanBeCanceled);
using var _ = await this.AcquireConnectionLockIfNeeded(isConnectionMonitoringQuery).ConfigureAwait(false);
await this.PrepareIfNeededAsync(cancellationToken).ConfigureAwait(false);
try
{
return(await executeAsync(state, cancellationToken).ConfigureAwait(false));
}
catch (Exception ex)
// Canceled SQL operations throw SqlException/InvalidOperationException instead of OCE.
// That means that downstream operations end up faulted instead of canceled. We
// wrap with OCE here to correctly propagate cancellation
when(cancellationToken.IsCancellationRequested && this._connection.IsCommandCancellationException(ex))
{
throw new OperationCanceledException(
"Command was canceled",
ex,
cancellationToken
);
}
}
public override bool IsCommandCancellationException(Exception exception)
{
const int CanceledNumber = 0;
// fast path using default SqlClient
if (exception is SqlException sqlException && sqlException.Number == CanceledNumber)
{
return(true);
}
var exceptionType = exception.GetType();
// since SqlException is sealed (as of 2020-01-26)
if (exceptionType.ToString() == "System.Data.SqlClient.SqlException")
{
var numberProperty = exceptionType
.GetProperty(nameof(SqlException.Number), BindingFlags.Public | BindingFlags.Instance);
Invariant.Require(numberProperty != null);
if (numberProperty != null)
{
return(Equals(numberProperty.GetValue(exception), CanceledNumber));
}
}
// this shows up when you call DbCommand.Cancel()
return(exception is InvalidOperationException);
}
private SqlApplicationLock(Mode mode, bool isUpgrade = false)
{
Invariant.Require(!isUpgrade || mode == Mode.Exclusive);
this._mode = mode;
this._isUpgrade = isUpgrade;
}
private void CloseOrCancelMonitoringHandleRegistrationsNoLock(bool isCancel)
{
Invariant.Require(this._state == State.AutoStopped || this._state == State.Stopped || this._state == State.Disposed);
if (this._monitoringHandleRegistrations == null)
{
return;
}
foreach (var kvp in this._monitoringHandleRegistrations)
{
var cancellationTokenSource = kvp.Value;
if (isCancel)
{
// cancel in a background thread in case we have hangs or errors
Task.Run(() =>
{
try { cancellationTokenSource.Cancel(); }
finally { cancellationTokenSource.Dispose(); }
});
}
else
{
cancellationTokenSource.Dispose();
}
}
this._monitoringHandleRegistrations.Clear();
}
private void OnConnectionStateChanged(object sender, StateChangeEventArgs args)
{
if (args.OriginalState == ConnectionState.Open && args.CurrentState != ConnectionState.Open)
{
lock (this.Lock)
{
if (this._state == State.Idle || this._state == State.Active)
{
this._state = State.AutoStopped;
this.CloseOrCancelMonitoringHandleRegistrationsNoLock(isCancel: true);
}
Invariant.Require(!this.HasRegisteredMonitoringHandlesNoLock);
}
}
else if (args.OriginalState != ConnectionState.Open && args.CurrentState == ConnectionState.Open)
{
lock (this.Lock)
{
if (this._state == State.AutoStopped)
{
this.StartNoLock();
}
}
}
}
public async ValueTask <Result> TryAcquireAsync <TLockCookie>(
string name,
TimeoutValue timeout,
IDbSynchronizationStrategy <TLockCookie> strategy,
TimeoutValue keepaliveCadence,
CancellationToken cancellationToken,
bool opportunistic)
where TLockCookie : class
{
using var mutextHandle = await this._mutex.TryAcquireAsync(opportunistic?TimeSpan.Zero : Timeout.InfiniteTimeSpan, cancellationToken).ConfigureAwait(false);
if (mutextHandle == null)
{
// mutex wasn't free, so just give up
Invariant.Require(opportunistic);
// The current lock is busy so we allow retry but on a different lock instance. We can't safely dispose
// since we never acquired the mutex so we can't check _heldLocks
return(new Result(MultiplexedConnectionLockRetry.Retry, canSafelyDispose: false));
}
try
{
if (this._heldLocksToKeepaliveCadences.ContainsKey(name))
{
// we won't try to hold the same lock twice on one connection. At some point, we could
// support this case in-memory using a counter for each multiply-held lock name and being careful
// with modes
return(this.GetFailureResultNoLock(isAlreadyHeld: true, opportunistic, timeout));
}
if (!this._connection.CanExecuteQueries)
{
await this._connection.OpenAsync(cancellationToken).ConfigureAwait(false);
}
var lockCookie = await strategy.TryAcquireAsync(this._connection, name, opportunistic?TimeSpan.Zero : timeout, cancellationToken).ConfigureAwait(false);
if (lockCookie != null)
{
var handle = new ManagedFinalizationDistributedLockHandle(new Handle <TLockCookie>(this, strategy, name, lockCookie));
this._heldLocksToKeepaliveCadences.Add(name, keepaliveCadence);
if (!keepaliveCadence.IsInfinite)
{
this.SetKeepaliveCadenceNoLock();
}
return(new Result(handle));
}
// we failed to acquire the lock, so we should retry if we were being opportunistic and artificially
// shortened the timeout
return(this.GetFailureResultNoLock(isAlreadyHeld: false, opportunistic, timeout));
}
finally
{
await this.CloseConnectionIfNeededNoLockAsync().ConfigureAwait(false);
}
}
public void Start()
{
Invariant.Require(!this._isExternallyOwnedConnection);
lock (this.Lock)
{
Invariant.Require(this._state == State.Stopped);
this.StartNoLock();
}
}
public static bool HasSufficientSuccesses(int successCount, int databaseCount)
{
// a majority is required
var threshold = (databaseCount / 2) + 1;
// While in theory this should return true if we have more than enough, we never expect this to be
// called except with just enough or not enough due to how we've implemented our approaches.
Invariant.Require(successCount <= threshold);
return(successCount >= threshold);
}
public override int CountActiveSessions(string applicationName)
{
Invariant.Require(applicationName.Length <= this.MaxApplicationNameLength);
using var connection = new OracleConnection(DefaultConnectionString);
connection.Open();
using var command = connection.CreateCommand();
command.CommandText = "SELECT COUNT(*) FROM v$session WHERE client_info = :applicationName AND status != 'KILLED'";
command.Parameters.Add("applicationName", applicationName);
return((int)(decimal)command.ExecuteScalar() !);
}
public override async Task SleepAsync(TimeSpan sleepTime, CancellationToken cancellationToken, Func <DatabaseCommand, CancellationToken, ValueTask <int> > executor)
{
Invariant.Require(sleepTime >= TimeSpan.Zero && sleepTime < TimeSpan.FromDays(1));
using var command = this.CreateCommand();
command.SetCommandText(@"WAITFOR DELAY @delay");
command.AddParameter("delay", sleepTime.ToString(@"hh\:mm\:ss\.fff"), DbType.AnsiStringFixedLength);
command.SetTimeout(sleepTime);
await executor(command, cancellationToken).ConfigureAwait(false);
}
public override int CountActiveSessions(string applicationName)
{
Invariant.Require(applicationName.Length <= this.MaxApplicationNameLength);
using var connection = new NpgsqlConnection(DefaultConnectionString);
connection.Open();
using var command = connection.CreateCommand();
command.CommandText = "SELECT COUNT(*)::int FROM pg_stat_activity WHERE application_name = @applicationName";
command.Parameters.AddWithValue("applicationName", applicationName);
return((int)command.ExecuteScalar() !);
}
public override int CountActiveSessions(string applicationName)
{
Invariant.Require(applicationName.Length <= this.MaxApplicationNameLength);
using var connection = new Microsoft.Data.SqlClient.SqlConnection(DefaultConnectionString);
connection.Open();
using var command = connection.CreateCommand();
command.CommandText = $@"SELECT COUNT(*) FROM sys.dm_exec_sessions WHERE program_name = @applicationName";
command.Parameters.AddWithValue("applicationName", applicationName);
return((int)command.ExecuteScalar());
}
public static bool HasTooManyFailuresOrFaults(int failureOrFaultCount, int databaseCount)
{
// For an odd number of databases, we need a majority to make success impossible. For an
// even number, however, getting to 50% failures/faults is sufficient to rule out getting
// a majority of successes.
var threshold = (databaseCount / 2) + (databaseCount % 2);
// While in theory this should return true if we have more than enough, we never expect this to be
// called except with just enough or not enough due to how we've implemented our approaches.
Invariant.Require(failureOrFaultCount <= threshold);
return(failureOrFaultCount >= threshold);
}
public ConnectionMonitor(DatabaseConnection connection)
{
this._weakConnection = new WeakReference <DatabaseConnection>(connection);
this._isExternallyOwnedConnection = connection.IsExernallyOwned;
// stopped not autostopped here so that the statechange handler will not cause a start
this._state = connection.CanExecuteQueries ? State.Idle : State.Stopped;
Invariant.Require(this._state == State.Stopped || this._isExternallyOwnedConnection);
if (connection.InnerConnection is DbConnection dbConnection)
{
dbConnection.StateChange += this._stateChangedHandler = this.OnConnectionStateChanged;
}
}
private async ValueTask StopOrDisposeAsync(bool isDispose)
{
Task?task;
lock (this.Lock)
{
if (isDispose)
{
this._state = State.Disposed;
}
else
{
Invariant.Require(!this._isExternallyOwnedConnection);
Invariant.Require(this._state != State.Disposed);
this._state = State.Stopped;
}
// If we have any registered monitoring handles, clear them out.
// We don't cancel them since if the helper was stopped that indicates
// proper disposal rather than loss of the connection
this.CloseOrCancelMonitoringHandleRegistrationsNoLock(isCancel: false);
task = this._monitoringWorkerTask;
// Note: synchronous cancel here should be safe because we've already set
// the state to disposed above which the monitoring loop will check if it
// takes over the Cancel() thread.
this._monitorStateChangedTokenSource?.Cancel();
// unsubscribe from state change tracking
if (this._stateChangedHandler != null &&
this._weakConnection.TryGetTarget(out var connection))
{
((DbConnection)connection.InnerConnection).StateChange -= this._stateChangedHandler;
}
}
if (task != null)
{
if (SyncViaAsync.IsSynchronous)
{
task.GetAwaiter().GetResult();
}
else
{
await task.ConfigureAwait(false);
}
}
}
public void Dispose()
{
var strategy = Interlocked.Exchange(ref this._strategy, null);
if (strategy != null)
{
Invariant.Require(strategy._preparedForHandleLost);
try { strategy._killHandleAction?.Invoke(); }
finally
{
strategy._killHandleAction = null;
strategy._preparedForHandleLost = false;
}
}
}
// note: we could have this return an IAsyncDisposable which would allow you to close the transaction
// without closing the connection. However, we don't currently have any use-cases for that
public async ValueTask BeginTransactionAsync()
{
Invariant.Require(!this.HasTransaction);
using var _ = await this.ConnectionMonitor.AcquireConnectionLockAsync(CancellationToken.None).ConfigureAwait(false);
this._transaction =
#if NETSTANDARD2_1
!SyncViaAsync.IsSynchronous && this.InnerConnection is DbConnection dbConnection
? await dbConnection.BeginTransactionAsync().ConfigureAwait(false)
:
#elif NETSTANDARD2_0 || NET461
#else
ERROR
#endif
this.InnerConnection.BeginTransaction();
}
private async ValueTask DisposeOrCloseAsync(bool isDispose)
{
Invariant.Require(isDispose || !this.IsExernallyOwned);
try
{
await(isDispose ? this.ConnectionMonitor.DisposeAsync() : this.ConnectionMonitor.StopAsync()).ConfigureAwait(false);
}
finally
{
if (!this.IsExernallyOwned)
{
try { await this.DisposeTransactionAsync(isClosingOrDisposingConnection : true).ConfigureAwait(false); }
finally
{
#if NETSTANDARD2_1
if (!SyncViaAsync.IsSynchronous && this.InnerConnection is DbConnection dbConnection)
{
await(isDispose ? dbConnection.DisposeAsync() : dbConnection.CloseAsync().AsValueTask()).ConfigureAwait(false);
}
else
{
SyncDisposeConnection();
}
#elif NETSTANDARD2_0 || NET461
SyncDisposeConnection();
#else
ERROR
#endif
}
}
}
void SyncDisposeConnection()
{
if (isDispose)
{
this.InnerConnection.Dispose();
}
else
{
this.InnerConnection.Close();
}
}
}
public void SetKeepaliveCadence(TimeoutValue keepaliveCadence)
{
Invariant.Require(!this._isExternallyOwnedConnection);
lock (this.Lock)
{
Invariant.Require(this._state != State.Disposed);
var originalKeepaliveCadence = this._keepaliveCadence;
this._keepaliveCadence = keepaliveCadence;
if (!this.StartMonitorWorkerIfNeededNoLock() &&
this._state == State.Active &&
!this.HasRegisteredMonitoringHandlesNoLock &&
keepaliveCadence.CompareTo(originalKeepaliveCadence) < 0)
{
// If we get here, then we already have an active worker performing
// keepalive on a longer cadence. Since that worker is likely asleep,
// we fire state changed to wake it up
this.FireStateChangedNoLock();
}
}
}
private async Task <bool> WaitForAcquireAsync(IReadOnlyDictionary <IDatabase, Task <bool> > tryAcquireTasks)
{
using var timeout = new TimeoutTask(this._primitive.AcquireTimeout, this._cancellationToken);
var incompleteTasks = new HashSet <Task>(tryAcquireTasks.Values)
{
timeout.Task
};
var successCount = 0;
var failCount = 0;
var faultCount = 0;
while (true)
{
var completed = await Task.WhenAny(incompleteTasks).ConfigureAwait(false);
if (completed == timeout.Task)
{
await completed.ConfigureAwait(false); // propagates cancellation
return(false); // true timeout
}
if (completed.Status == TaskStatus.RanToCompletion)
{
var result = await((Task <bool>)completed).ConfigureAwait(false);
if (result)
{
++successCount;
if (RedLockHelper.HasSufficientSuccesses(successCount, this._databases.Count))
{
return(true);
}
}
else
{
++failCount;
if (RedLockHelper.HasTooManyFailuresOrFaults(failCount, this._databases.Count))
{
return(false);
}
}
}
else // faulted or canceled
{
// if we get too many faults, the lock is not possible to acquire, so we should throw
++faultCount;
if (RedLockHelper.HasTooManyFailuresOrFaults(faultCount, this._databases.Count))
{
var faultingTasks = tryAcquireTasks.Values.Where(t => t.IsCanceled || t.IsFaulted)
.ToArray();
if (faultingTasks.Length == 1)
{
await faultingTasks[0].ConfigureAwait(false); // propagate the error
}
throw new AggregateException(faultingTasks.Select(t => t.Exception ?? new TaskCanceledException(t).As <Exception>()))
.Flatten();
}
++failCount;
if (RedLockHelper.HasTooManyFailuresOrFaults(failCount, this._databases.Count))
{
return(false);
}
}
incompleteTasks.Remove(completed);
Invariant.Require(incompleteTasks.Count > 1, "should be more than just timeout left");
}
}
public async Task <bool?> TryExtendAsync()
{
Invariant.Require(!SyncViaAsync.IsSynchronous, "should only be called from a background renewal thread which is async");
var incompleteTasks = new HashSet <Task>();
foreach (var kvp in this._tryAcquireOrRenewTasks.ToArray())
{
if (kvp.Value.IsCompleted)
{
incompleteTasks.Add(
this._tryAcquireOrRenewTasks[kvp.Key] = Helpers.SafeCreateTask(
state => state.primitive.TryExtendAsync(state.database),
(primitive: this._primitive, database: kvp.Key)
)
);
}
else
{
// if the previous acquire/renew is still going, just keep waiting for that
incompleteTasks.Add(kvp.Value);
}
}
// For extension we use the same timeout as acquire. This ensures the same min validity time which should be
// sufficient to keep extending
using var timeout = new TimeoutTask(this._primitive.AcquireTimeout, this._cancellationToken);
incompleteTasks.Add(timeout.Task);
var databaseCount = this._tryAcquireOrRenewTasks.Count;
var successCount = 0;
var failCount = 0;
while (true)
{
var completed = await Task.WhenAny(incompleteTasks).ConfigureAwait(false);
if (completed == timeout.Task)
{
await completed.ConfigureAwait(false); // propagate cancellation
return(null); // inconclusive
}
if (completed.Status == TaskStatus.RanToCompletion && ((Task <bool>)completed).Result)
{
++successCount;
if (RedLockHelper.HasSufficientSuccesses(successCount, databaseCount))
{
return(true);
}
}
else
{
// note that we treat faulted and failed the same in extend. There's no reason to throw, since
// this is just called by the extend loop. While in theory a fault could indicate some kind of post-success
// failure, most likely it means the db is unreachable and so it is safest to consider it a failure
++failCount;
if (RedLockHelper.HasTooManyFailuresOrFaults(failCount, databaseCount))
{
return(false);
}
}
incompleteTasks.Remove(completed);
}
}
public override IDisposable?PrepareForHandleLost()
{
Invariant.Require(!this._preparedForHandleLost);
this._preparedForHandleLost = true;
return(new HandleLostScope(this));
}
public ValueTask DisposeAsync()
{
Invariant.Require(this._heldLocksToKeepaliveCadences.Count == 0);
return(this._connection.DisposeAsync());
}
public override void PerformAdditionalCleanupForHandleAbandonment()
{
Invariant.Require(this._preparedForHandleAbandonment);
Thread.Sleep(TimeSpan.FromSeconds(.5));
}
public async ValueTask <IDistributedSynchronizationHandle?> TryAcquireAsync <TLockCookie>(
TimeoutValue timeout,
IDbSynchronizationStrategy <TLockCookie> strategy,
CancellationToken cancellationToken,
IDistributedSynchronizationHandle?contextHandle)
where TLockCookie : class
{
IDistributedSynchronizationHandle?result = null;
IAsyncDisposable?connectionResource = null;
try
{
DatabaseConnection connection;
bool transactionScoped;
if (contextHandle != null)
{
connection = GetContextHandleConnection <TLockCookie>(contextHandle);
transactionScoped = false;
}
else
{
connectionResource = connection = this._connectionFactory();
if (connection.IsExernallyOwned)
{
Invariant.Require(!this._scopeToOwnedTransaction);
if (!connection.CanExecuteQueries)
{
throw new InvalidOperationException("The connection and/or transaction are disposed or closed");
}
transactionScoped = false;
}
else
{
await connection.OpenAsync(cancellationToken).ConfigureAwait(false);
if (this._scopeToOwnedTransaction)
{
await connection.BeginTransactionAsync().ConfigureAwait(false);
}
transactionScoped = this._scopeToOwnedTransaction;
}
}
var lockCookie = await strategy.TryAcquireAsync(connection, this._name, timeout, cancellationToken).ConfigureAwait(false);
if (lockCookie != null)
{
result = new Handle <TLockCookie>(connection, strategy, this._name, lockCookie, transactionScoped, connectionResource);
if (!this._keepaliveCadence.IsInfinite)
{
connection.SetKeepaliveCadence(this._keepaliveCadence);
}
}
}
finally
{
// if we fail to acquire or throw, make sure to clean up the connection
if (result == null && connectionResource != null)
{
await connectionResource.DisposeAsync().ConfigureAwait(false);
}
}
return(result);
}
public Pool(TimeoutValue maxAge)
{
Invariant.Require(!maxAge.IsInfinite);
this._maxAge = maxAge.TimeSpan;
}
