DbDataReader left open when cancellation occurs in AdvanceAsync after ExecuteReaderAsync

[drakke]

Description

There's a risk DbDataReader won't be disposed of in GetResultSetAsync method within the src\NHibernate\Async\Loader\Loader.cs class. When a cancellation token is canceled after the DbDataReader is successfully obtained from session.Batcher.ExecuteReaderAsync() but during the AdvanceAsync method call, the reader is not closed.

The issue occurs due to the catch block for OperationCanceledException that throws without closing the reader:

catch (OperationCanceledException) { throw; }
catch (Exception sqle)
{
    ADOExceptionReporter.LogExceptions(sqle);
    session.Batcher.CloseCommand(st, rs);
    throw;
}

Steps to Reproduce

I manually changed the NH source code to force the error to happen. Adding the following right after the DataReader initially gets set:

rs = await (session.Batcher.ExecuteReaderAsync(st, cancellationToken)).ConfigureAwait(false);

// Added code on row 1064 to force OperationCanceledException:
CancellationTokenSource cts = new CancellationTokenSource();
cts.Cancel();
cts.Token.ThrowIfCancellationRequested();

Then I added this test to src\NHibernate.Test\Async\Legacy\SQLLoaderTest.cs:

[Test]
public void ExpiredCancellationTokenThrowsOperationCanceledException()
{
    if (Dialect is Oracle8iDialect)
    {
        return;
    }

    Assert.ThrowsAsync<OperationCanceledException>(() => InnerExpiredCancellationTokenThrowsOperationCanceledException());
}

public async Task InnerExpiredCancellationTokenThrowsOperationCanceledException()
{
    using ISession session = OpenSession();
    using ITransaction txn = session.BeginTransaction();

    IQuery q = session.CreateSQLQuery(@"
        SELECT TOP 100000 a.*, b.* 
        FROM sys.objects a
        CROSS JOIN sys.objects b
        ORDER BY a.name");
    await (q.ListAsync());
}

This leads to an unexpected InvalidOperationException while rolling back the transaction rather than the expected OperationCanceledException.

  Message: 
      Expected: <System.OperationCanceledException>
      But was:  <System.InvalidOperationException: There is already an open DataReader associated with this Command which must be closed first.
       at System.Data.SqlClient.SqlInternalConnectionTds.ValidateConnectionForExecute(SqlCommand command)
       at System.Data.SqlClient.SqlInternalTransaction.Rollback()
       at System.Data.SqlClient.SqlInternalTransaction.Dispose(Boolean disposing)
       at System.Data.SqlClient.SqlInternalTransaction.Dispose()
       at System.Data.SqlClient.SqlTransaction.Dispose(Boolean disposing)
       at NHibernate.Transaction.AdoTransaction.Dispose(Boolean isDisposing) in C:\dev\nhibernate-core\src\NHibernate\Transaction\AdoTransaction.cs:line 383
       at NHibernate.Transaction.AdoTransaction.Dispose() in C:\dev\nhibernate-core\src\NHibernate\Transaction\AdoTransaction.cs:line 352
       at NHibernate.Test.Legacy.SQLLoaderTestAsync.InnerExpiredCancellationTokenThrowsOperationCanceledException() in C:\dev\nhibernate-core\src\NHibernate.Test\Async\Legacy\SQLLoaderTest.cs:line 92
       at NUnit.Framework.Internal.TaskAwaitAdapter.GenericAdapter`1.BlockUntilCompleted()
       at NUnit.Framework.Internal.MessagePumpStrategy.NoMessagePumpStrategy.WaitForCompletion(AwaitAdapter awaiter)
       at NUnit.Framework.Internal.AsyncToSyncAdapter.Await(Func`1 invoke)
       at NUnit.Framework.Assert.ThrowsAsync(IResolveConstraint expression, AsyncTestDelegate code, String message, Object[] args)>

Reproducing without manually canceling the token

It is not consistent, but I was able to reproduce with this test after 1-25 iterations:

public async Task InnerExpiredCancellationTokenThrowsOperationCanceledException()
{
    using var cts = new CancellationTokenSource(100);
    using ISession session = OpenSession();
    using ITransaction txn = session.BeginTransaction();

    IQuery q = session.CreateSQLQuery(@"
        SELECT TOP 100000 a.*, b.* 
        FROM sys.objects a
        CROSS JOIN sys.objects b
        ORDER BY a.name");
    await (q.ListAsync(cts.Token));
}

Possible Resolution

This seems to be generated code, but the goal would be to somehow make sure at least session.Batcher.CloseReader(rs); gets called in case of OperationCanceledException.

catch (OperationCanceledException)
{
    session.Batcher.CloseReader(rs);
    throw; 
}

It might be that session.Batcher.CloseCommand(st, rs); should be used, as in the generic catch block? Calling any of those would ensure the DbDataReader gets closed which would allow a transaction to rollback without encountering InvalidOperationException that hides the original exception.

Related issues

maca88/AsyncGenerator#90
dotnet/SqlClient#26 (Makes this less likely to happen)

[delta-emil]

I ran into this as well. In my case ExecuteReaderAsync did throw an OperationCanceledException. I think it's because the query happened to return many rows streamed in multiple network packages or something (similar to how you sometimes start seeing results way before the query is done in the SQL Management Studio). But the particular point where the OperationCanceledException happens shouldn't matter that much. It can occur at many points along the way, and there is no case where leaving the in the data reader open is okay.

In my case, the logic wants to do a quick update of a record to mark an operation finished unsuccessfully, after the transaction, but was completely confused by the transaction's dispose/roll-back throwing an exception.

---

Here is a test to reproduce it reliably without changing the main source code (useful in order to have a test remain after the fix).
It uses a subclass of the MicrosoftDataSqlClientDriver to return a wrapper of its DbCommand, which triggers the cancellation right after the underlying DbCommand returns the reader.

using System.Data;
using System.Data.Common;
using NHibernate;
using NHibernate.Cfg;
using NHibernate.Cfg.MappingSchema;
using NHibernate.Dialect;
using NHibernate.Driver;
using NHibernate.Mapping.ByCode;
using NHibernate.Tool.hbm2ddl;
using NUnit.Framework;

namespace TestDataReaderOpCancel;

public class Tests
{
    // just an empty MS SQL DB
    private const string ConnectionString = "Server=(local);initial catalog=nhibernate;Integrated Security=SSPI;Encrypt=False";

    private Configuration? _Configuration;
    private bool _SchemaCreated;
    private ISessionFactory? _SessionFactory;

    [SetUp]
    public void SetUp()
    {
        var mapper = new ModelMapper();
        mapper.Class<Person>(ca =>
        {
            ca.Id(x => x.Id, map => map.Generator(Generators.HighLow));
            ca.Property(x => x.NotNullData);
        });
        HbmMapping hbmMapping = mapper.CompileMappingForAllExplicitlyAddedEntities();

        // plug in TestDriver where we can issue the token cancellation at a predictable moment
        //_Configuration = GetConfiguration(hbmMapping, typeof(NHibernate.Driver.MicrosoftDataSqlClientDriver));
        _Configuration = GetConfiguration(hbmMapping, typeof(TestDriver));

        _SessionFactory = _Configuration.BuildSessionFactory();
        new SchemaExport(_Configuration).Create(false, true);
        _SchemaCreated = true;
    }

    private static Configuration GetConfiguration(HbmMapping hbmMapping, Type driverType)
    {
        var configuration = new Configuration();

        configuration.SetProperty("dialect", typeof(MsSql2012Dialect).AssemblyQualifiedName);
        configuration.SetProperty("connection.driver_class", driverType.AssemblyQualifiedName);
        configuration.SetProperty("default_schema", "dbo");
        configuration.SetProperty("connection.connection_string", ConnectionString);

        configuration.AddDeserializedMapping(hbmMapping, null);

        configuration.SetProperty("connection.provider", typeof(NHibernate.Connection.DriverConnectionProvider).AssemblyQualifiedName);
        configuration.SetProperty("connection.isolation", "ReadCommitted");

        configuration.SetProperty("cache.use_second_level_cache", "true");
        configuration.SetProperty("cache.provider_class", typeof(NHibernate.Cache.HashtableCacheProvider).AssemblyQualifiedName);

        configuration.SetProperty("format_sql", "true");

        return configuration;
    }

    [TearDown]
    public void TearDown()
    {
        _SessionFactory?.Dispose();
        if (_SchemaCreated)
        {
            new SchemaExport(_Configuration).Drop(false, true);
        }
    }

    [Test]
    public async Task Test1()
    {
        using (ISession session = _SessionFactory!.OpenSession())
        {
            await session.SaveAsync(new Person { Id = 42, NotNullData = "test" });
        }

        var cancellationTokenSource = new CancellationTokenSource();
        TestDriver.CancellationTokenSource = cancellationTokenSource;

        try
        {
            using (ISession session = _SessionFactory.OpenSession())
            {
                using (ITransaction transaction = session.BeginTransaction())
                {
                    try
                    {
                        await session.QueryOver<Person>().ListAsync(cancellationTokenSource.Token);

                        Assert.Fail("We should have had an exception caused by the cancellation token.");
                    }
                    catch (Exception e)
                    {
                        // Console.WriteLine(e);
                        Assert.That(e, Is.TypeOf<OperationCanceledException>());
                        throw;
                    }

                    // you can call it explicitly, or leave it to the transaction's Dispose
                    // transaction.Rollback();
                }
            }

            Assert.Fail("We should have had an exception caused by the cancellation token.");
        }
        catch (Exception e)
        {
            // The bug is that here we see that Rollback of the transaction failed, instead of still getting the OperationCanceledException
            Assert.That(e, Is.TypeOf<OperationCanceledException>());
        }
    }
}

public class Person
{
    public virtual int Id { get; set; }
    public virtual string NotNullData { get; set; } = "not-null";
}

public class TestDriver : MicrosoftDataSqlClientDriver
{
    public static CancellationTokenSource? CancellationTokenSource { get; set; }

    public override DbCommand CreateCommand()
    {
        return new TestDbCommand(base.CreateCommand());
    }
}

public class TestDbCommand : DbCommand
{
    private DbCommand _underlyingCommand;

    public TestDbCommand(DbCommand underlyingCommand)
    {
        _underlyingCommand = underlyingCommand;
    }

    protected override Task<DbDataReader> ExecuteDbDataReaderAsync(CommandBehavior behavior, CancellationToken cancellationToken)
    {
        var reader = _underlyingCommand.ExecuteReaderAsync(behavior, cancellationToken);

        if (TestDriver.CancellationTokenSource != null)
        {
            TestDriver.CancellationTokenSource.Cancel();
        }

        return reader;
    }

    // the rest simply redirect to the underlying command

    public override string? CommandText
    {
        get => _underlyingCommand.CommandText;
        set => _underlyingCommand.CommandText = value;
    }

    public override int CommandTimeout
    {
        get => _underlyingCommand.CommandTimeout;
        set => _underlyingCommand.CommandTimeout = value;
    }

    public override CommandType CommandType
    {
        get => _underlyingCommand.CommandType;
        set => _underlyingCommand.CommandType = value;
    }

    public override bool DesignTimeVisible
    {
        get => _underlyingCommand.DesignTimeVisible;
        set => _underlyingCommand.DesignTimeVisible = value;
    }

    public override UpdateRowSource UpdatedRowSource
    {
        get => _underlyingCommand.UpdatedRowSource;
        set => _underlyingCommand.UpdatedRowSource = value;
    }

    protected override DbConnection? DbConnection
    {
        get => _underlyingCommand.Connection;
        set => _underlyingCommand.Connection = value;
    }

    protected override DbParameterCollection DbParameterCollection => _underlyingCommand.Parameters;

    protected override DbTransaction? DbTransaction
    {
        get => _underlyingCommand.Transaction;
        set => _underlyingCommand.Transaction = value;
    }

    public override void Cancel() => _underlyingCommand.Cancel();

    public override int ExecuteNonQuery() => _underlyingCommand.ExecuteNonQuery();

    public override object? ExecuteScalar() =>_underlyingCommand.ExecuteScalar();

    public override void Prepare() => _underlyingCommand.Prepare();

    protected override DbParameter CreateDbParameter() => _underlyingCommand.CreateParameter();

    protected override DbDataReader ExecuteDbDataReader(CommandBehavior behavior) => _underlyingCommand.ExecuteReader(behavior);

    public override Task<int> ExecuteNonQueryAsync(CancellationToken cancellationToken) => _underlyingCommand.ExecuteNonQueryAsync(cancellationToken);

    public override Task<object?> ExecuteScalarAsync(CancellationToken cancellationToken) => _underlyingCommand.ExecuteScalarAsync(cancellationToken);

    public override Task PrepareAsync(CancellationToken cancellationToken = default) => _underlyingCommand.PrepareAsync(cancellationToken);

    public override ValueTask DisposeAsync() => _underlyingCommand.DisposeAsync();
}

Also note that an OperationCanceledException at that point is particularly nasty, because even if we could easily change Loader's GetResultSetAsync to have

catch (OperationCanceledException)
{
    session.Batcher.CloseCommand(st, rs);
    throw;
}

that would not resolve the issue. The DbDataReader rs is still null because the exception is thrown before it was returned; and disposing of the command (DbCommand st) doesn't seem to close the data reader. Moreover, the data reader didn't even make it into AbstractBatcher's _readersToClose.

This seems to need handling in DoExecuteReaderAsync itself, or the reference to the open reader gets lost when something like NHybridDataReader.CreateAsync throws.

Alternatively, maybe something more drastic like closing the connection is in order. Maybe that will manage to close the reader we may have lost track of.