Add missing OperationCanceledException in async code where required

Tools/packages.config

@@ -7,7 +7,7 @@
   <package id="NUnit.Extension.NUnitV2ResultWriter" version="3.6.0" targetFramework="net461" />
   <package id="NUnit.Extension.TeamCityEventListener" version="1.0.2" targetFramework="net461" />
   <package id="NUnit.Extension.VSProjectLoader" version="3.6.0" targetFramework="net461" />
-  <package id="CSharpAsyncGenerator.CommandLine" version="0.8.2.1" targetFramework="net461" />
+  <package id="CSharpAsyncGenerator.CommandLine" version="0.8.2.2" targetFramework="net461" />
   <package id="vswhere" version="2.1.4" targetFramework="net461" />
   <package id="gitreleasemanager" version="0.7.0" targetFramework="net461" />
 </packages>

src/NHibernate.Test/Async/Cascade/Circle/MultiPathCircleCascadeTest.cs

@@ -155,6 +155,7 @@ public async Task MergeEntityWithNonNullableEntityNullAsync()
 					await (session.MergeAsync(route, cancellationToken));
 					Assert.Fail("should have thrown an exception");
 				}
+				catch (OperationCanceledException) { throw; }
 				catch (Exception ex)
 				{
 					Assert.That(ex, Is.TypeOf(typeof(PropertyValueException)));

src/NHibernate.Test/Async/DebugConnectionProvider.cs

@@ -30,6 +30,7 @@ public override async Task<DbConnection> GetConnectionAsync(CancellationToken ca
 				connections.TryAdd(connection, 0);
 				return connection;
 			}
+			catch (OperationCanceledException) { throw; }
 			catch (Exception e)
 			{
 				throw new HibernateException("Could not open connection to: " + ConnectionString, e);

src/NHibernate.Test/Async/NHSpecificTest/NH1553/MsSQL/SnapshotIsolationUpdateConflictTest.cs

@@ -77,16 +77,17 @@ private void SavePerson(Person p)
 		/// Tests, that NHibernate detects the update conflict and returns a StaleObjectStateException as expected.
 		/// This behaviour is part of the standard NHibernate feature set.
 		/// </summary>
+		/// <param name="cancellationToken">A cancellation token that can be used to cancel the work</param>
 		[Test]
-		public async Task UpdateConflictDetectedByNHAsync()
+		public async Task UpdateConflictDetectedByNHAsync(CancellationToken cancellationToken = default(CancellationToken))
 		{
-			Person p1 = await (LoadPersonAsync());
-			Person p2 = await (LoadPersonAsync());
+			Person p1 = await (LoadPersonAsync(cancellationToken));
+			Person p2 = await (LoadPersonAsync(cancellationToken));
 
 			p1.IdentificationNumber++;
 			p2.IdentificationNumber += 2;
 
-			await (SavePersonAsync(p1));
+			await (SavePersonAsync(p1, cancellationToken));
 			Assert.That(p1.Version, Is.EqualTo(person.Version + 1));
 
 			var expectedException = Sfi.Settings.IsBatchVersionedDataEnabled
@@ -95,7 +96,7 @@ public async Task UpdateConflictDetectedByNHAsync()
 				        .And.Property("EntityName").EqualTo(typeof(Person).FullName)
 				        .And.Property("Identifier").EqualTo(p2.Id);
 
-			Assert.That(() => SavePersonAsync(p2), expectedException);
+			Assert.That(() => SavePersonAsync(p2, cancellationToken), expectedException);
 		}
 
 		/// <summary>

src/NHibernate/Async/AdoNet/AbstractBatcher.cs

@@ -123,6 +123,7 @@ public async Task<int> ExecuteNonQueryAsync(DbCommand cmd, CancellationToken can
 			{
 				return await (cmd.ExecuteNonQueryAsync(cancellationToken)).ConfigureAwait(false);
 			}
+			catch (OperationCanceledException) { throw; }
 			catch (Exception e)
 			{
 				e.Data["actual-sql-query"] = cmd.CommandText;
@@ -150,6 +151,7 @@ public virtual async Task<DbDataReader> ExecuteReaderAsync(DbCommand cmd, Cancel
 			{
 				reader = await (cmd.ExecuteReaderAsync(cancellationToken)).ConfigureAwait(false);
 			}
+			catch (OperationCanceledException) { throw; }
 			catch (Exception e)
 			{
 				e.Data["actual-sql-query"] = cmd.CommandText;

src/NHibernate/Async/Collection/Generic/PersistentGenericIdentifierBag.cs

@@ -193,6 +193,7 @@ public override async Task PreInsertAsync(ICollectionPersister persister, Cancel
 					}
 				}
 			}
+			catch (OperationCanceledException) { throw; }
 			catch (Exception sqle)
 			{
 				throw new ADOException("Could not generate idbag row id.", sqle);

src/NHibernate/Async/Connection/DriverConnectionProvider.cs

@@ -39,6 +39,7 @@ public override async Task<DbConnection> GetConnectionAsync(CancellationToken ca
 				conn.ConnectionString = ConnectionString;
 				await (conn.OpenAsync(cancellationToken)).ConfigureAwait(false);
 			}
+			catch (OperationCanceledException) { throw; }
 			catch (Exception)
 			{
 				conn.Dispose();

src/NHibernate/Async/Context/ThreadLocalSessionContext.cs

@@ -38,13 +38,15 @@ private static async Task CleanupAnyOrphanedSessionAsync(ISessionFactory factory
 						{
 							await (orphan.Transaction.RollbackAsync(cancellationToken)).ConfigureAwait(false);
 						}
+						catch (OperationCanceledException) { throw; }
 						catch (Exception ex)
 						{
 							log.Debug(ex, "Unable to rollback transaction for orphaned session");
 						}
 					}
 					orphan.Close();
 				}
+				catch (OperationCanceledException) { throw; }
 				catch (Exception ex)
 				{
 					log.Debug(ex, "Unable to close orphaned session");

src/NHibernate/Async/Dialect/Lock/SelectLockingStrategy.cs

@@ -64,6 +64,7 @@ public async Task LockAsync(object id, object version, object obj, ISessionImple
 					session.Batcher.CloseCommand(st, rs);
 				}
 			}
+			catch (OperationCanceledException) { throw; }
 			catch (HibernateException)
 			{
 				// Do not call Convert on HibernateExceptions

src/NHibernate/Async/Driver/NDataReader.cs

@@ -60,6 +60,7 @@ public static async Task<NDataReader> CreateAsync(DbDataReader reader, bool isMi
 
 				dataReader.results = resultList.ToArray();
 			}
+			catch (OperationCanceledException) { throw; }
 			catch (Exception e)
 			{
 				throw new ADOException("There was a problem converting an DbDataReader to NDataReader", e);

src/NHibernate/Async/Engine/Query/NativeSQLQueryPlan.cs

@@ -96,6 +96,7 @@ public async Task<int> PerformExecuteUpdateAsync(QueryParameters queryParameters
 					}
 				}
 			}
+			catch (OperationCanceledException) { throw; }
 			catch (HibernateException)
 			{
 				throw;

src/NHibernate/Async/Hql/Ast/ANTLR/Exec/AbstractStatementExecutor.cs

@@ -114,6 +114,7 @@ protected virtual async Task DropTemporaryTableIfNecessaryAsync(IQueryable persi
 					ps = await (session.Batcher.PrepareCommandAsync(CommandType.Text, commandText, Array.Empty<SqlType>(), cancellationToken)).ConfigureAwait(false);
 					await (session.Batcher.ExecuteNonQueryAsync(ps, cancellationToken)).ConfigureAwait(false);
 				}
+				catch (OperationCanceledException) { throw; }
 				catch (Exception t)
 				{
 					log.Warn(t, "unable to cleanup temporary id table after use [{0}]", t);
@@ -150,6 +151,7 @@ public async Task DoWorkAsync(DbConnection connection, DbTransaction transaction
 					await (stmnt.ExecuteNonQueryAsync(cancellationToken)).ConfigureAwait(false);
 					session.Factory.Settings.SqlStatementLogger.LogCommand(stmnt, FormatStyle.Ddl);
 				}
+				catch (OperationCanceledException) { throw; }
 				catch (Exception t)
 				{
 					log.Debug(t, "unable to create temporary id table [{0}]", t.Message);
@@ -186,6 +188,7 @@ public async Task DoWorkAsync(DbConnection connection, DbTransaction transaction
 					await (stmnt.ExecuteNonQueryAsync(cancellationToken)).ConfigureAwait(false);
 					session.Factory.Settings.SqlStatementLogger.LogCommand(stmnt, FormatStyle.Ddl);
 				}
+				catch (OperationCanceledException) { throw; }
 				catch (Exception t)
 				{
 					log.Warn("unable to drop temporary id table after use [{0}]", t.Message);

src/NHibernate/Async/Id/Enhanced/TableGenerator.cs

@@ -101,6 +101,7 @@ public override async Task<object> DoWorkInCurrentTransactionAsync(ISessionImple
 						result = Convert.ToInt64(selectedValue);
 					}
 				}
+				catch (OperationCanceledException) { throw; }
 				catch (Exception ex)
 				{
 					log.Error(ex, "Unable to read or initialize hi value in {0}", TableName);
@@ -124,6 +125,7 @@ public override async Task<object> DoWorkInCurrentTransactionAsync(ISessionImple
 						updatedRows = await (updateCmd.ExecuteNonQueryAsync(cancellationToken)).ConfigureAwait(false);
 					}
 				}
+				catch (OperationCanceledException) { throw; }
 				catch (Exception ex)
 				{
 					log.Error(ex, "Unable to update hi value in {0}", TableName);

src/NHibernate/Async/Id/Enhanced/TableStructure.cs

@@ -55,6 +55,7 @@ public override async Task<object> DoWorkInCurrentTransactionAsync(ISessionImple
 					}
 					result = Convert.ToInt64(selectedValue);
 				}
+				catch (OperationCanceledException) { throw; }
 				catch (Exception sqle)
 				{
 					Log.Error(sqle, "could not read a hi value");
@@ -76,6 +77,7 @@ public override async Task<object> DoWorkInCurrentTransactionAsync(ISessionImple
 						updatedRows = await (updateCmd.ExecuteNonQueryAsync(cancellationToken)).ConfigureAwait(false);
 					}
 				}
+				catch (OperationCanceledException) { throw; }
 				catch (Exception sqle)
 				{
 					Log.Error(sqle, "could not update hi value in: {0}", _tableName);

src/NHibernate/Async/Id/IdentifierGeneratorFactory.cs

@@ -72,6 +72,7 @@ public static async Task<object> GetAsync(DbDataReader rs, IType type, ISessionI
 			{
 				return await (type.NullSafeGetAsync(rs, rs.GetName(0), session, null, cancellationToken)).ConfigureAwait(false);
 			}
+			catch (OperationCanceledException) { throw; }
 			catch (Exception e)
 			{
 				throw new IdentifierGenerationException("could not retrieve identifier value", e);

src/NHibernate/Async/Id/NativeGuidGenerator.cs

@@ -56,6 +56,7 @@ public async Task<object> GenerateAsync(ISessionImplementor session, object obj,
 					session.Batcher.CloseCommand(st, reader);
 				}
 			}
+			catch (OperationCanceledException) { throw; }
 			catch (Exception sqle)
 			{
 				throw ADOExceptionHelper.Convert(session.Factory.SQLExceptionConverter, sqle, "could not retrieve GUID", sql);

src/NHibernate/Async/Id/TableGenerator.cs

@@ -87,6 +87,7 @@ public override async Task<object> DoWorkInCurrentTransactionAsync(ISessionImple
 					}
 					result = Convert.ToInt64(columnType.Get(rs, 0, session));
 				}
+				catch (OperationCanceledException) { throw; }
 				catch (Exception e)
 				{
 					log.Error(e, "could not read a hi value");
@@ -114,6 +115,7 @@ public override async Task<object> DoWorkInCurrentTransactionAsync(ISessionImple
 
 					rows = await (ups.ExecuteNonQueryAsync(cancellationToken)).ConfigureAwait(false);
 				}
+				catch (OperationCanceledException) { throw; }
 				catch (Exception e)
 				{
 					log.Error(e, "could not update hi value in: {0}", tableName);

src/NHibernate/Async/Impl/MultiCriteriaImpl.cs

@@ -205,6 +205,7 @@ private async Task GetResultsFromDatabaseAsync(IList results, CancellationToken
 					}
 				}
 			}
+			catch (OperationCanceledException) { throw; }
 			catch (Exception sqle)
 			{
 				log.Error(sqle, "Failed to execute multi criteria: [{0}]", resultSetsCommand.Sql);

src/NHibernate/Async/Impl/MultiQueryImpl.cs

@@ -173,6 +173,7 @@ protected async Task<List<object>> DoListAsync(CancellationToken cancellationTok
 					}
 				}
 			}
+			catch (OperationCanceledException) { throw; }
 			catch (Exception sqle)
 			{
 				log.Error(sqle, "Failed to execute multi query: [{0}]", resultSetsCommand.Sql);

src/NHibernate/Async/Impl/SessionImpl.cs

@@ -268,6 +268,7 @@ private async Task ListAsync(IQueryExpression queryExpression, QueryParameters q
 						await (plan.PerformListAsync(queryParameters, this, results, cancellationToken)).ConfigureAwait(false);
 						success = true;
 					}
+					catch (OperationCanceledException) { throw; }
 					catch (HibernateException)
 					{
 						// Do not call Convert on HibernateExceptions
@@ -1065,6 +1066,7 @@ private async Task FilterAsync(object collection, string filter, QueryParameters
 						await (plan.PerformListAsync(queryParameters, this, results, cancellationToken)).ConfigureAwait(false);
 						success = true;
 					}
+					catch (OperationCanceledException) { throw; }
 					catch (HibernateException)
 					{
 						// Do not call Convert on HibernateExceptions
@@ -1155,6 +1157,7 @@ public override async Task ListAsync(CriteriaImpl criteria, IList results, Cance
 						}
 						success = true;
 					}
+					catch (OperationCanceledException) { throw; }
 					catch (HibernateException)
 					{
 						// Do not call Convert on HibernateExceptions

src/NHibernate/Async/Impl/StatelessSessionImpl.cs

@@ -104,6 +104,7 @@ public override async Task ListAsync(IQueryExpression queryExpression, QueryPara
 					await (plan.PerformListAsync(queryParameters, this, results, cancellationToken)).ConfigureAwait(false);
 					success = true;
 				}
+				catch (OperationCanceledException) { throw; }
 				catch (HibernateException)
 				{
 					// Do not call Convert on HibernateExceptions
@@ -145,6 +146,7 @@ public override async Task ListAsync(CriteriaImpl criteria, IList results, Cance
 					}
 					success = true;
 				}
+				catch (OperationCanceledException) { throw; }
 				catch (HibernateException)
 				{
 					// Do not call Convert on HibernateExceptions

src/NHibernate/Async/Loader/Loader.cs

@@ -116,6 +116,7 @@ protected async Task<object> LoadSingleRowAsync(DbDataReader resultSet, ISession
 					await (GetRowFromResultSetAsync(resultSet, session, queryParameters, GetLockModes(queryParameters.LockModes), null,
 										hydratedObjects, new EntityKey[entitySpan], returnProxies, cancellationToken)).ConfigureAwait(false);
 			}
+			catch (OperationCanceledException) { throw; }
 			catch (HibernateException)
 			{
 				throw; // Don't call Convert on HibernateExceptions
@@ -298,6 +299,7 @@ private async Task<IList> DoQueryAsync(ISessionImplementor session, QueryParamet
 						Log.Debug("done processing result set ({0} rows)", count);
 					}
 				}
+				catch (OperationCanceledException) { throw; }
 				catch (Exception e)
 				{
 					e.Data["actual-sql-query"] = st.CommandText;
@@ -812,6 +814,7 @@ protected internal virtual async Task<DbCommand> PrepareQueryCommandAsync(QueryP
 				driver.RemoveUnusedCommandParameters(command, sqlString);
 				driver.ExpandQueryParameters(command, sqlString, sqlCommand.ParameterTypes);
 			}
+			catch (OperationCanceledException) { throw; }
 			catch (HibernateException)
 			{
 				session.Batcher.CloseCommand(command, null);
@@ -893,6 +896,7 @@ protected async Task<DbDataReader> GetResultSetAsync(
 				}
 				return rs;
 			}
+			catch (OperationCanceledException) { throw; }
 			catch (Exception sqle)
 			{
 				ADOExceptionReporter.LogExceptions(sqle);
@@ -922,6 +926,7 @@ protected async Task<IList> LoadEntityAsync(ISessionImplementor session, object
 										optionalIdentifier);
 				result = await (DoQueryAndInitializeNonLazyCollectionsAsync(session, qp, false, cancellationToken)).ConfigureAwait(false);
 			}
+			catch (OperationCanceledException) { throw; }
 			catch (HibernateException)
 			{
 				throw;
@@ -955,6 +960,7 @@ protected async Task<IList> LoadEntityAsync(ISessionImplementor session, object
 														   new QueryParameters(new IType[] { keyType, indexType },
 																			   new object[] { key, index }), false, cancellationToken)).ConfigureAwait(false);
 			}
+			catch (OperationCanceledException) { throw; }
 			catch (Exception sqle)
 			{
 				throw ADOExceptionHelper.Convert(_factory.SQLExceptionConverter, sqle, "could not collection element by index",
@@ -989,6 +995,7 @@ protected internal async Task<IList> LoadEntityBatchAsync(ISessionImplementor se
 														   new QueryParameters(types, ids, optionalObject, optionalEntityName,
 																			   optionalId), false, cancellationToken)).ConfigureAwait(false);
 			}
+			catch (OperationCanceledException) { throw; }
 			catch (HibernateException)
 			{
 				throw;
@@ -1021,6 +1028,7 @@ public async Task LoadCollectionAsync(ISessionImplementor session, object id, IT
 			{
 				await (DoQueryAndInitializeNonLazyCollectionsAsync(session, new QueryParameters(new IType[] { type }, ids, ids), true, cancellationToken)).ConfigureAwait(false);
 			}
+			catch (OperationCanceledException) { throw; }
 			catch (HibernateException)
 			{
 				// Do not call Convert on HibernateExceptions
@@ -1053,6 +1061,7 @@ public async Task LoadCollectionBatchAsync(ISessionImplementor session, object[]
 			{
 				await (DoQueryAndInitializeNonLazyCollectionsAsync(session, new QueryParameters(idTypes, ids, ids), true, cancellationToken)).ConfigureAwait(false);
 			}
+			catch (OperationCanceledException) { throw; }
 			catch (HibernateException)
 			{
 				// Do not call Convert on HibernateExceptions
@@ -1082,6 +1091,7 @@ protected async Task LoadCollectionSubselectAsync(ISessionImplementor session, o
 													   new QueryParameters(parameterTypes, parameterValues, namedParameters, ids),
 													   true, cancellationToken)).ConfigureAwait(false);
 			}
+			catch (OperationCanceledException) { throw; }
 			catch (HibernateException)
 			{
 				// Do not call Convert on HibernateExceptions
@@ -1283,6 +1293,7 @@ protected async Task<IList> DoListAsync(ISessionImplementor session, QueryParame
 			{
 				result = await (DoQueryAndInitializeNonLazyCollectionsAsync(session, queryParameters, true, forcedResultTransformer, cancellationToken)).ConfigureAwait(false);
 			}
+			catch (OperationCanceledException) { throw; }
 			catch (HibernateException)
 			{
 				// Do not call Convert on HibernateExceptions

src/NHibernate/Async/Persister/Collection/AbstractCollectionPersister.cs

@@ -217,6 +217,7 @@ public async Task RemoveAsync(object id, ISessionImplementor session, Cancellati
 							expectation.VerifyOutcomeNonBatched(await (session.Batcher.ExecuteNonQueryAsync(st, cancellationToken)).ConfigureAwait(false), st);
 						}
 					}
+					catch (OperationCanceledException) { throw; }
 					catch (Exception e)
 					{
 						if (useBatch)
@@ -376,6 +377,7 @@ public async Task DeleteRowsAsync(IPersistentCollection collection, object id, I
 							}
 							count++;
 						}
+						catch (OperationCanceledException) { throw; }
 						catch (Exception e)
 						{
 							if (useBatch)
@@ -517,6 +519,7 @@ protected async Task<object> PerformInsertAsync(object ownerId, IPersistentColle
 					expectation.VerifyOutcomeNonBatched(await (session.Batcher.ExecuteNonQueryAsync(st, cancellationToken)).ConfigureAwait(false), st);
 				}
 			}
+			catch (OperationCanceledException) { throw; }
 			catch (Exception e)
 			{
 				if (useBatch)

src/NHibernate/Async/Persister/Collection/BasicCollectionPersister.cs

@@ -98,6 +98,7 @@ protected override async Task<int> DoUpdateRowsAsync(object id, IPersistentColle
 								expectation.VerifyOutcomeNonBatched(await (session.Batcher.ExecuteNonQueryAsync(st, cancellationToken)).ConfigureAwait(false), st);
 							}
 						}
+						catch (OperationCanceledException) { throw; }
 						catch (Exception e)
 						{
 							if (useBatch)