Skip to content

Commit 3a025fa

Browse files
rojiroji
committed
Implement batching error barrier control
Closes #4205
1 parent 58e3a4c commit 3a025fa

9 files changed

Lines changed: 421 additions & 43 deletions

File tree

‎src/Npgsql/Internal/NpgsqlConnector.cs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1358,7 +1358,8 @@ internal ValueTask<IBackendMessage> ReadMessage(bool async, DataRowLoadingMode d
13581358
// an RFQ. Instead, the server closes the connection immediately
13591359
throw error;
13601360
}
1361-
else if (PostgresErrorCodes.IsCriticalFailure(error, clusterError: false))
1361+
1362+
if (PostgresErrorCodes.IsCriticalFailure(error, clusterError: false))
13621363
{
13631364
// Consider the connection dead
13641365
throw connector.Break(error);

‎src/Npgsql/NpgsqlBatch.cs‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,32 @@ protected override DbTransaction? DbTransaction
5454
set => Transaction = (NpgsqlTransaction?)value;
5555
}
5656

57+
/// <summary>
58+
/// Controls whether to place error barriers between all batch commands within this batch. Default to <see langword="false" />.
59+
/// </summary>
60+
/// <remarks>
61+
/// <para>
62+
/// By default, any exception in a command causes later commands in the batch to be skipped, and earlier commands to be rolled back.
63+
/// Enabling error barriers ensures that errors do not affect other commands in the batch.
64+
/// </para>
65+
/// <para>
66+
/// Note that if the batch is executed within an explicit transaction, the first error places the transaction in a failed state,
67+
/// causing all later commands to fail in any case. As a result, this option is useful mainly when there is no explicit transaction.
68+
/// </para>
69+
/// <para>
70+
/// At the PostgreSQL wire protocol level, this corresponds to inserting a Sync message between each command, rather than grouping
71+
/// all the batch's commands behind a single terminating Sync.
72+
/// </para>
73+
/// <para>
74+
/// To control error barriers on a command-by-command basis, see <see cref="NpgsqlBatchCommand.AppendErrorBarrier" />.
75+
/// </para>
76+
/// </remarks>
77+
public bool EnableErrorBarriers
78+
{
79+
get => Command.EnableErrorBarriers;
80+
set => Command.EnableErrorBarriers = value;
81+
}
82+
5783
/// <summary>
5884
/// Marks all of the batch's result columns as either known or unknown.
5985
/// Unknown results column are requested them from PostgreSQL in text format, and Npgsql makes no

‎src/Npgsql/NpgsqlBatchCommand.cs‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,31 @@ public override string CommandText
3232
/// <inheritdoc cref="DbBatchCommand.Parameters"/>
3333
public new NpgsqlParameterCollection Parameters { get; } = new();
3434

35+
/// <summary>
36+
/// Appends an error barrier after this batch command. Defaults to the value of <see cref="NpgsqlBatch.EnableErrorBarriers" /> on the
37+
/// batch.
38+
/// </summary>
39+
/// <remarks>
40+
/// <para>
41+
/// By default, any exception in a command causes later commands in the batch to be skipped, and earlier commands to be rolled back.
42+
/// Appending an error barrier ensures that errors from this command (or previous ones) won't cause later commands to be skipped,
43+
/// and that errors from later commands won't cause this command (or previous ones) to be rolled back).
44+
/// </para>
45+
/// <para>
46+
/// Note that if the batch is executed within an explicit transaction, the first error places the transaction in a failed state,
47+
/// causing all later commands to fail in any case. As a result, this option is useful mainly when there is no explicit transaction.
48+
/// </para>
49+
/// <para>
50+
/// At the PostgreSQL wire protocol level, this corresponds to inserting a Sync message after this command, rather than grouping
51+
/// all the batch's commands behind a single terminating Sync.
52+
/// </para>
53+
/// <para>
54+
/// Controlling error barriers on a command-by-command basis is an advanced feature, consider enabling error barriers for the entire
55+
/// batch via <see cref="NpgsqlBatch.EnableErrorBarriers" />.
56+
/// </para>
57+
/// </remarks>
58+
public bool? AppendErrorBarrier { get; set; }
59+
3560
/// <summary>
3661
/// The number of rows affected or retrieved.
3762
/// </summary>

‎src/Npgsql/NpgsqlCommand.cs‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,8 @@ public class NpgsqlCommand : DbCommand, ICloneable, IComponent
7676
internal static readonly bool EnableSqlRewriting;
7777
#endif
7878

79+
internal bool EnableErrorBarriers { get; set; }
80+
7981
static readonly List<NpgsqlParameter> EmptyParameters = new();
8082

8183
static readonly SingleThreadSynchronizationContext SingleThreadSynchronizationContext = new("NpgsqlRemainingAsyncSendWorker");
@@ -1000,11 +1002,18 @@ await connector.WriteBind(
10001002

10011003
await connector.WriteExecute(0, async, cancellationToken);
10021004

1005+
if (batchCommand.AppendErrorBarrier ?? EnableErrorBarriers)
1006+
await connector.WriteSync(async, cancellationToken);
1007+
10031008
if (pStatement != null)
10041009
pStatement.LastUsed = DateTime.UtcNow;
10051010
}
10061011

1007-
await connector.WriteSync(async, cancellationToken);
1012+
if (InternalBatchCommands.Count == 0 ||
1013+
!EnableErrorBarriers && InternalBatchCommands[InternalBatchCommands.Count - 1].AppendErrorBarrier != true)
1014+
{
1015+
await connector.WriteSync(async, cancellationToken);
1016+
}
10081017

10091018
if (flush)
10101019
await connector.Flush(async, cancellationToken);

‎src/Npgsql/NpgsqlDataReader.cs‎

Lines changed: 113 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
using System.IO;
1010
using System.Linq;
1111
using System.Runtime.CompilerServices;
12+
using System.Runtime.ExceptionServices;
1213
using System.Text;
1314
using System.Threading;
1415
using System.Threading.Tasks;
@@ -283,9 +284,24 @@ async Task<bool> Read(bool async, CancellationToken cancellationToken = default)
283284
throw new ArgumentOutOfRangeException();
284285
}
285286

286-
var msg2 = await ReadMessage(async);
287-
ProcessMessage(msg2);
288-
return msg2.Code == BackendMessageCode.DataRow;
287+
var msg = await ReadMessage(async);
288+
289+
switch (msg.Code)
290+
{
291+
case BackendMessageCode.DataRow:
292+
ProcessMessage(msg);
293+
return true;
294+
295+
case BackendMessageCode.CommandComplete:
296+
case BackendMessageCode.EmptyQueryResponse:
297+
ProcessMessage(msg);
298+
if (_statements[StatementIndex].AppendErrorBarrier ?? Command.EnableErrorBarriers)
299+
Expect<ReadyForQueryMessage>(await Connector.ReadMessage(async), Connector);
300+
return false;
301+
302+
default:
303+
throw Connector.UnexpectedMessageReceived(msg.Code);
304+
}
289305
}
290306
catch
291307
{
@@ -335,10 +351,11 @@ public override bool NextResult() => (_isSchemaOnly ? NextResultSchemaOnly(false
335351
/// <returns>A task representing the asynchronous operation.</returns>
336352
public override Task<bool> NextResultAsync(CancellationToken cancellationToken)
337353
{
338-
using (NoSynchronizationContextScope.Enter())
339-
return _isSchemaOnly
340-
? NextResultSchemaOnly(async: true, cancellationToken: cancellationToken)
341-
: NextResult(async: true, cancellationToken: cancellationToken);
354+
using var _ = NoSynchronizationContextScope.Enter();
355+
356+
return _isSchemaOnly
357+
? NextResultSchemaOnly(async: true, cancellationToken: cancellationToken)
358+
: NextResult(async: true, cancellationToken: cancellationToken);
342359
}
343360

344361
/// <summary>
@@ -370,7 +387,12 @@ async Task<bool> NextResult(bool async, bool isConsuming = false, CancellationTo
370387
case BackendMessageCode.CommandComplete:
371388
case BackendMessageCode.EmptyQueryResponse:
372389
ProcessMessage(completedMsg);
390+
391+
if (_statements[StatementIndex].AppendErrorBarrier ?? Command.EnableErrorBarriers)
392+
Expect<ReadyForQueryMessage>(await Connector.ReadMessage(async), Connector);
393+
373394
break;
395+
374396
default:
375397
continue;
376398
}
@@ -472,6 +494,10 @@ async Task<bool> NextResult(bool async, bool isConsuming = false, CancellationTo
472494
}
473495

474496
ProcessMessage(msg);
497+
498+
if (statement.AppendErrorBarrier ?? Command.EnableErrorBarriers)
499+
Expect<ReadyForQueryMessage>(await Connector.ReadMessage(async), Connector);
500+
475501
continue;
476502
}
477503

@@ -494,30 +520,34 @@ async Task<bool> NextResult(bool async, bool isConsuming = false, CancellationTo
494520
switch (msg.Code)
495521
{
496522
case BackendMessageCode.DataRow:
523+
return true;
497524
case BackendMessageCode.CommandComplete:
498-
break;
525+
if (statement.AppendErrorBarrier ?? Command.EnableErrorBarriers)
526+
Expect<ReadyForQueryMessage>(await Connector.ReadMessage(async), Connector);
527+
return true;
499528
default:
500529
throw Connector.UnexpectedMessageReceived(msg.Code);
501530
}
502-
503-
return true;
504531
}
505532

506533
// There are no more queries, we're done. Read the RFQ.
507-
ProcessMessage(Expect<ReadyForQueryMessage>(await Connector.ReadMessage(async), Connector));
534+
if (_statements.Count == 0 ||
535+
!Command.EnableErrorBarriers && _statements[_statements.Count - 1].AppendErrorBarrier != true)
536+
{
537+
Expect<ReadyForQueryMessage>(await Connector.ReadMessage(async), Connector);
538+
}
539+
State = ReaderState.Consumed;
508540
RowDescription = null;
509541
return false;
510542
}
511543
catch (Exception e)
512544
{
513-
State = ReaderState.Consumed;
514-
515545
// Reference the triggering statement from the exception
516546
if (e is PostgresException postgresException && StatementIndex >= 0 && StatementIndex < _statements.Count)
517547
{
518548
postgresException.BatchCommand = _statements[StatementIndex];
519549

520-
// Prevent the command or batch from by recycled (by the connection) when it's disposed. This is important since
550+
// Prevent the command or batch from being recycled (by the connection) when it's disposed. This is important since
521551
// the exception is very likely to escape the using statement of the command, and by that time some other user may
522552
// already be using the recycled instance.
523553
if (!Command.IsWrappedByBatch)
@@ -526,9 +556,8 @@ async Task<bool> NextResult(bool async, bool isConsuming = false, CancellationTo
526556
}
527557
}
528558

529-
// An error means all subsequent statements were skipped by PostgreSQL.
530-
// If any of them were being prepared, we need to update our bookkeeping to put
531-
// them back in unprepared state.
559+
// For the statement that errored, if it was being prepared we need to update our bookkeeping to put them back in unprepared
560+
// state.
532561
for (; StatementIndex < _statements.Count; StatementIndex++)
533562
{
534563
var statement = _statements[StatementIndex];
@@ -537,6 +566,34 @@ async Task<bool> NextResult(bool async, bool isConsuming = false, CancellationTo
537566
statement.IsPreparing = false;
538567
statement.PreparedStatement!.AbortPrepare();
539568
}
569+
570+
if (statement.AppendErrorBarrier ?? Command.EnableErrorBarriers)
571+
break;
572+
}
573+
574+
// In normal, non-isolated batching, we've consumed the result set and are done.
575+
// However, if an isolated command was present after the error, we now have to consume the rest of the result set.
576+
// Note that Consume calls NextResult (this method) recursively, the isConsuming flag tells us we're in this mode.
577+
if (StatementIndex == _statements.Count)
578+
{
579+
State = ReaderState.Consumed;
580+
}
581+
else if (!isConsuming)
582+
{
583+
switch (State)
584+
{
585+
case ReaderState.Consumed:
586+
case ReaderState.Closed:
587+
case ReaderState.Disposed:
588+
// The exception may have caused the connector to break (e.g. I/O), and so the reader is already closed.
589+
break;
590+
default:
591+
// We provide Consume with the first exception which we've just caught.
592+
// If it encounters other exceptions while consuming the rest of the result set, it will raise an AggregateException,
593+
// otherwise it will rethrow this first exception.
594+
await Consume(async, firstException: e);
595+
break;
596+
}
540597
}
541598

542599
throw;
@@ -672,8 +729,9 @@ async Task<bool> NextResultSchemaOnly(bool async, bool isConsuming = false, Canc
672729
// There are no more queries, we're done. Read to the RFQ.
673730
if (!_statements.All(s => s.IsPrepared))
674731
{
675-
ProcessMessage(Expect<ReadyForQueryMessage>(await Connector.ReadMessage(async), Connector));
732+
Expect<ReadyForQueryMessage>(await Connector.ReadMessage(async), Connector);
676733
RowDescription = null;
734+
State = ReaderState.Consumed;
677735
}
678736

679737
return false;
@@ -748,10 +806,6 @@ internal void ProcessMessage(IBackendMessage msg)
748806
State = ReaderState.BetweenResults;
749807
return;
750808

751-
case BackendMessageCode.ReadyForQuery:
752-
State = ReaderState.Consumed;
753-
return;
754-
755809
default:
756810
throw new Exception("Received unexpected backend message of type " + msg.Code);
757811
}
@@ -901,14 +955,44 @@ public override int FieldCount
901955
/// Consumes all result sets for this reader, leaving the connector ready for sending and processing further
902956
/// queries
903957
/// </summary>
904-
async Task Consume(bool async)
958+
async Task Consume(bool async, Exception? firstException = null)
905959
{
906-
// Skip over the other result sets. Note that this does tally records affected
907-
// from CommandComplete messages, and properly sets state for auto-prepared statements
908-
if (_isSchemaOnly)
909-
while (await NextResultSchemaOnly(async, isConsuming: true)) {}
910-
else
911-
while (await NextResult(async, isConsuming: true)) {}
960+
var exceptions = firstException is null ? null : new List<Exception> { firstException };
961+
962+
// Skip over the other result sets. Note that this does tally records affected from CommandComplete messages, and properly sets
963+
// state for auto-prepared statements
964+
while (true)
965+
{
966+
try
967+
{
968+
if (!(_isSchemaOnly
969+
? await NextResultSchemaOnly(async, isConsuming: true)
970+
: await NextResult(async, isConsuming: true)))
971+
{
972+
break;
973+
}
974+
}
975+
catch (Exception e)
976+
{
977+
exceptions ??= new();
978+
exceptions.Add(e);
979+
}
980+
}
981+
982+
Debug.Assert(exceptions?.Count != 0);
983+
984+
switch (exceptions?.Count)
985+
{
986+
case null:
987+
return;
988+
case 1:
989+
ExceptionDispatchInfo.Capture(exceptions[0]).Throw();
990+
return;
991+
default:
992+
throw new NpgsqlException(
993+
"Multiple exceptions occurred when consuming the result set",
994+
new AggregateException(exceptions));
995+
}
912996
}
913997

914998
/// <summary>

‎src/Npgsql/PostgresMinimalDatabaseInfo.cs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ class PostgresMinimalDatabaseInfoFactory : INpgsqlDatabaseInfoFactory
2020

2121
class PostgresMinimalDatabaseInfo : PostgresDatabaseInfo
2222
{
23-
static PostgresType[]? TypesWithMultiranges, TypesWithoutMultiranges;
23+
static PostgresType[]? _typesWithMultiranges, _typesWithoutMultiranges;
2424

2525
static PostgresType[] CreateTypes(bool withMultiranges)
2626
=> typeof(NpgsqlDbType).GetFields()
@@ -50,8 +50,8 @@ static PostgresType[] CreateTypes(bool withMultiranges)
5050

5151
protected override IEnumerable<PostgresType> GetTypes()
5252
=> SupportsMultirangeTypes
53-
? TypesWithMultiranges ??= CreateTypes(withMultiranges: true)
54-
: TypesWithoutMultiranges ??= CreateTypes(withMultiranges: false);
53+
? _typesWithMultiranges ??= CreateTypes(withMultiranges: true)
54+
: _typesWithoutMultiranges ??= CreateTypes(withMultiranges: false);
5555

5656
internal PostgresMinimalDatabaseInfo(NpgsqlConnector conn)
5757
: base(conn)

‎src/Npgsql/PublicAPI.Unshipped.txt‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,8 @@
11
#nullable enable
2+
Npgsql.NpgsqlBatch.EnableErrorBarriers.get -> bool
3+
Npgsql.NpgsqlBatch.EnableErrorBarriers.set -> void
4+
Npgsql.NpgsqlBatchCommand.AppendErrorBarrier.get -> bool?
5+
Npgsql.NpgsqlBatchCommand.AppendErrorBarrier.set -> void
26
Npgsql.NpgsqlLoggingConfiguration
37
static Npgsql.NpgsqlLoggingConfiguration.InitializeLogging(Microsoft.Extensions.Logging.ILoggerFactory! loggerFactory, bool parameterLoggingEnabled = false) -> void
48
*REMOVED*Npgsql.NpgsqlConnection.Settings.get -> Npgsql.NpgsqlConnectionStringBuilder!

‎src/Npgsql/Util/PGUtil.cs‎

Lines changed: 4 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -27,14 +27,10 @@ static Statics()
2727

2828
[MethodImpl(MethodImplOptions.AggressiveInlining)]
2929
internal static T Expect<T>(IBackendMessage msg, NpgsqlConnector connector)
30-
{
31-
if (msg is T asT)
32-
return asT;
33-
34-
throw connector.Break(
35-
new NpgsqlException($"Received backend message {msg.Code} while expecting {typeof(T).Name}. " +
36-
"Please file a bug."));
37-
}
30+
=> msg is T t
31+
? t
32+
: throw connector.Break(
33+
new NpgsqlException($"Received backend message {msg.Code} while expecting {typeof(T).Name}. Please file a bug."));
3834

3935
internal static DeferDisposable Defer(Action action) => new(action);
4036
internal static DeferDisposable<T> Defer<T>(Action<T> action, T arg) => new(action, arg);

0 commit comments

Comments
 (0)