diff --git a/src/Npgsql/NpgsqlMultiHostDataSource.cs b/src/Npgsql/NpgsqlMultiHostDataSource.cs index 4e6f42eeae..aa4d4f8ebc 100644 --- a/src/Npgsql/NpgsqlMultiHostDataSource.cs +++ b/src/Npgsql/NpgsqlMultiHostDataSource.cs @@ -278,6 +278,12 @@ static bool IsOnline(DatabaseState state, TargetSessionAttributes preferredType) return connector; } + catch (OperationCanceledException oce) when (cancellationToken.IsCancellationRequested && oce.CancellationToken == cancellationToken) + { + if (connector is not null) + pool.Return(connector); + throw; + } catch (Exception ex) { exceptions.Add(ex); diff --git a/test/Npgsql.Tests/MultipleHostsTests.cs b/test/Npgsql.Tests/MultipleHostsTests.cs index 398c0520ff..8085ca8c6c 100644 --- a/test/Npgsql.Tests/MultipleHostsTests.cs +++ b/test/Npgsql.Tests/MultipleHostsTests.cs @@ -1139,13 +1139,39 @@ public async Task OpenConnection_when_canceled_throws_TaskCanceledException() { var builder = new NpgsqlDataSourceBuilder(ConnectionString); await using var dataSource = builder.BuildMultiHost(); - using var cts = new CancellationTokenSource(); - cts.Cancel(); + var cancellationToken = new CancellationToken(true); var ex = Assert.ThrowsAsync(async () => { - await using var connection = await dataSource.OpenConnectionAsync(cts.Token); + await using var connection = await dataSource.OpenConnectionAsync(cancellationToken); }); - Assert.That(ex.CancellationToken, Is.EqualTo(cts.Token)); + Assert.That(ex.CancellationToken, Is.EqualTo(cancellationToken)); + } + + [Test] + public async Task OpenConnection_when_canceled_during_TryGet_throws_OperationCanceledException() + { + await using var primary1 = PgPostmasterMock.Start(state: Primary); + await using var primary2 = PgPostmasterMock.Start(state: Primary); + + var connectionString = new NpgsqlConnectionStringBuilder($"Host={primary1.Host}:{primary1.Port},{primary2.Host}:{primary2.Port}") + { + ServerCompatibilityMode = ServerCompatibilityMode.NoTypeLoading, + MaxPoolSize = 1 + }.ToString(); + await using var dataSource = new NpgsqlDataSourceBuilder(connectionString).BuildMultiHost(); + + // Exhaust the pool so that TryGetIdleOrNew returns null and we fall through to TryGet + await using var conn1 = await dataSource.OpenConnectionAsync(TargetSessionAttributes.Primary); + await using var conn2 = await dataSource.OpenConnectionAsync(TargetSessionAttributes.Primary); + + var cancellationToken = new CancellationToken(true); + + var ex = Assert.ThrowsAsync(async () => + { + await using var conn3 = await dataSource.OpenConnectionAsync(TargetSessionAttributes.Primary, cancellationToken); + }); + + Assert.That(ex.CancellationToken, Is.EqualTo(cancellationToken)); } [Test, IssueLink("https://github.com/npgsql/npgsql/issues/4181")]