Skip to content

Commit 251d73b

Browse files
vonzshikvonzshik
authored
Add tracing for physical connection open (#6091)
Closes #4136
1 parent cf9d243 commit 251d73b

7 files changed

Lines changed: 228 additions & 18 deletions

File tree

‎src/Npgsql/Internal/NpgsqlConnector.cs‎

Lines changed: 24 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -485,9 +485,18 @@ internal async Task Open(NpgsqlTimeout timeout, bool async, CancellationToken ca
485485
LogMessages.OpeningPhysicalConnection(ConnectionLogger, Host, Port, Database, UserFacingConnectionString);
486486
var startOpenTimestamp = Stopwatch.GetTimestamp();
487487

488+
Activity? activity = null;
489+
488490
try
489491
{
490-
await OpenCore(this, Settings.SslMode, timeout, async, cancellationToken).ConfigureAwait(false);
492+
var username = await GetUsernameAsync(async, cancellationToken).ConfigureAwait(false);
493+
494+
activity = NpgsqlActivitySource.ConnectionOpen(this);
495+
496+
await OpenCore(this, username, Settings.SslMode, timeout, async, cancellationToken).ConfigureAwait(false);
497+
498+
if (activity is not null)
499+
NpgsqlActivitySource.Enrich(activity, this);
491500

492501
await DataSource.Bootstrap(this, timeout, forceReload: false, async, cancellationToken).ConfigureAwait(false);
493502

@@ -510,6 +519,8 @@ internal async Task Open(NpgsqlTimeout timeout, bool async, CancellationToken ca
510519
// It is intentionally not awaited and will run as long as the connector is alive.
511520
// The CommandsInFlightWriter channel is completed in Cleanup, which should cause this task
512521
// to complete.
522+
// Make sure we do not flow AsyncLocals like Activity.Current
523+
using var __ = ExecutionContext.SuppressFlow();
513524
_ = Task.Run(MultiplexingReadLoop, CancellationToken.None)
514525
.ContinueWith(t =>
515526
{
@@ -540,7 +551,7 @@ internal async Task Open(NpgsqlTimeout timeout, bool async, CancellationToken ca
540551
{
541552
if (async)
542553
await DataSource.ConnectionInitializerAsync(tempConnection).ConfigureAwait(false);
543-
else if (!async)
554+
else
544555
DataSource.ConnectionInitializer(tempConnection);
545556
}
546557
finally
@@ -553,26 +564,31 @@ internal async Task Open(NpgsqlTimeout timeout, bool async, CancellationToken ca
553564
}
554565
}
555566

567+
if (activity is not null)
568+
NpgsqlActivitySource.CommandStop(activity);
569+
556570
LogMessages.OpenedPhysicalConnection(
557-
ConnectionLogger, Host, Port, Database, UserFacingConnectionString, (long)Stopwatch.GetElapsedTime(startOpenTimestamp).TotalMilliseconds, Id);
571+
ConnectionLogger, Host, Port, Database, UserFacingConnectionString,
572+
(long)Stopwatch.GetElapsedTime(startOpenTimestamp).TotalMilliseconds, Id);
558573
}
559574
catch (Exception e)
560575
{
576+
if (activity is not null)
577+
NpgsqlActivitySource.SetException(activity, e);
561578
Break(e);
562579
throw;
563580
}
564581

565582
static async Task OpenCore(
566583
NpgsqlConnector conn,
584+
string username,
567585
SslMode sslMode,
568586
NpgsqlTimeout timeout,
569587
bool async,
570588
CancellationToken cancellationToken)
571589
{
572590
await conn.RawOpen(sslMode, timeout, async, cancellationToken).ConfigureAwait(false);
573591

574-
var username = await conn.GetUsernameAsync(async, cancellationToken).ConfigureAwait(false);
575-
576592
timeout.CheckAndApply(conn);
577593
conn.WriteStartupMessage(username);
578594
await conn.Flush(async, cancellationToken).ConfigureAwait(false);
@@ -595,6 +611,7 @@ static async Task OpenCore(
595611
// If Allow was specified and we failed (without SSL), retry with SSL
596612
await OpenCore(
597613
conn,
614+
username,
598615
sslMode == SslMode.Prefer ? SslMode.Disable : SslMode.Require,
599616
timeout,
600617
async,
@@ -754,6 +771,8 @@ async Task RawOpen(SslMode sslMode, NpgsqlTimeout timeout, bool async, Cancellat
754771
else
755772
Connect(timeout);
756773

774+
ConnectionLogger.LogTrace("Socket connected to {Host}:{Port}", Host, Port);
775+
757776
_baseStream = new NetworkStream(_socket, true);
758777
_stream = _baseStream;
759778

@@ -810,8 +829,6 @@ async Task RawOpen(SslMode sslMode, NpgsqlTimeout timeout, bool async, Cancellat
810829
if (ReadBuffer.ReadBytesLeft > 0)
811830
throw new NpgsqlException("Additional unencrypted data received after SSL negotiation - this should never happen, and may be an indication of a man-in-the-middle attack.");
812831
}
813-
814-
ConnectionLogger.LogTrace("Socket connected to {Host}:{Port}", Host, Port);
815832
}
816833
catch
817834
{

‎src/Npgsql/MultiplexingDataSource.cs‎

Lines changed: 21 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,8 @@ internal MultiplexingDataSource(
5656
_connectionLogger = dataSourceConfig.LoggingConfiguration.ConnectionLogger;
5757
_commandLogger = dataSourceConfig.LoggingConfiguration.CommandLogger;
5858

59+
// Make sure we do not flow AsyncLocals like Activity.Current
60+
using var _ = ExecutionContext.SuppressFlow();
5961
_multiplexWriteLoop = Task.Run(MultiplexingWriteLoop, CancellationToken.None)
6062
.ContinueWith(t =>
6163
{
@@ -106,15 +108,28 @@ async Task MultiplexingWriteLoop()
106108
break;
107109
}
108110

109-
connector = await OpenNewConnector(
110-
command.InternalConnection!,
111-
new NpgsqlTimeout(TimeSpan.FromSeconds(Settings.Timeout)),
112-
async: true,
113-
CancellationToken.None).ConfigureAwait(false);
111+
// At no point should we ever have an activity here
112+
Debug.Assert(Activity.Current is null);
113+
// Set current activity as the one from the command
114+
// So child activities from physical open are bound to it
115+
Activity.Current = command.CurrentActivity;
116+
117+
try
118+
{
119+
connector = await OpenNewConnector(
120+
command.InternalConnection!,
121+
new NpgsqlTimeout(TimeSpan.FromSeconds(Settings.Timeout)),
122+
async: true,
123+
CancellationToken.None).ConfigureAwait(false);
124+
}
125+
finally
126+
{
127+
Activity.Current = null;
128+
}
114129

115130
if (connector != null)
116131
{
117-
// Managed to created a new connector
132+
// Managed to create a new connector
118133
connector.Connection = null;
119134

120135
// See increment under over-capacity mode below

‎src/Npgsql/NpgsqlActivitySource.cs‎

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ namespace Npgsql;
99

1010
static class NpgsqlActivitySource
1111
{
12-
static readonly ActivitySource Source = new("Npgsql", "0.1.0");
12+
static readonly ActivitySource Source = new("Npgsql", "0.2.0");
1313

1414
internal static bool IsEnabled => Source.HasListeners();
1515

@@ -61,6 +61,22 @@ static class NpgsqlActivitySource
6161
return activity;
6262
}
6363

64+
internal static Activity? ConnectionOpen(NpgsqlConnector connector)
65+
{
66+
if (!connector.DataSource.Configuration.TracingOptions.EnablePhysicalOpenTracing)
67+
return null;
68+
69+
var dbName = connector.Settings.Database ?? connector.InferredUserName;
70+
var activity = Source.StartActivity(dbName, ActivityKind.Client);
71+
if (activity is not { IsAllDataRequested: true })
72+
return activity;
73+
74+
activity.SetTag("db.system", "postgresql");
75+
activity.SetTag("db.connection_string", connector.UserFacingConnectionString);
76+
77+
return activity;
78+
}
79+
6480
internal static void Enrich(Activity activity, NpgsqlConnector connector)
6581
{
6682
if (!activity.IsAllDataRequested)

‎src/Npgsql/NpgsqlCommand.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@ public class NpgsqlCommand : DbCommand, ICloneable, IComponent
5151

5252
internal List<NpgsqlBatchCommand> InternalBatchCommands { get; }
5353

54-
Activity? CurrentActivity;
54+
internal Activity? CurrentActivity { get; private set; }
5555

5656
/// <summary>
5757
/// Returns details about each statement that this command has executed.

‎src/Npgsql/NpgsqlTracingOptionsBuilder.cs‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ public sealed class NpgsqlTracingOptionsBuilder
1515
Func<NpgsqlCommand, string?>? _commandSpanNameProvider;
1616
Func<NpgsqlBatch, string?>? _batchSpanNameProvider;
1717
bool _enableFirstResponseEvent = true;
18+
bool _enablePhysicalOpenTracing = true;
1819

1920
internal NpgsqlTracingOptionsBuilder()
2021
{
@@ -88,6 +89,16 @@ public NpgsqlTracingOptionsBuilder EnableFirstResponseEvent(bool enable = true)
8889
return this;
8990
}
9091

92+
/// <summary>
93+
/// Gets or sets a value indicating whether to trace physical connection open.
94+
/// Default is true to preserve existing behavior.
95+
/// </summary>
96+
public NpgsqlTracingOptionsBuilder EnablePhysicalOpenTracing(bool enable = true)
97+
{
98+
_enablePhysicalOpenTracing = enable;
99+
return this;
100+
}
101+
91102
internal NpgsqlTracingOptions Build() => new()
92103
{
93104
CommandFilter = _commandFilter,
@@ -96,7 +107,8 @@ public NpgsqlTracingOptionsBuilder EnableFirstResponseEvent(bool enable = true)
96107
BatchEnrichmentCallback = _batchEnrichmentCallback,
97108
CommandSpanNameProvider = _commandSpanNameProvider,
98109
BatchSpanNameProvider = _batchSpanNameProvider,
99-
EnableFirstResponseEvent = _enableFirstResponseEvent
110+
EnableFirstResponseEvent = _enableFirstResponseEvent,
111+
EnablePhysicalOpenTracing = _enablePhysicalOpenTracing
100112
};
101113
}
102114

@@ -109,4 +121,5 @@ sealed class NpgsqlTracingOptions
109121
internal Func<NpgsqlCommand, string?>? CommandSpanNameProvider { get; init; }
110122
internal Func<NpgsqlBatch, string?>? BatchSpanNameProvider { get; init; }
111123
internal bool EnableFirstResponseEvent { get; init; }
124+
internal bool EnablePhysicalOpenTracing { get; init; }
112125
}

‎src/Npgsql/PublicAPI.Unshipped.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ Npgsql.NpgsqlTracingOptionsBuilder.ConfigureCommandEnrichmentCallback(System.Act
3838
Npgsql.NpgsqlTracingOptionsBuilder.ConfigureCommandFilter(System.Func<Npgsql.NpgsqlCommand!, bool>? commandFilter) -> Npgsql.NpgsqlTracingOptionsBuilder!
3939
Npgsql.NpgsqlTracingOptionsBuilder.ConfigureCommandSpanNameProvider(System.Func<Npgsql.NpgsqlCommand!, string?>? commandSpanNameProvider) -> Npgsql.NpgsqlTracingOptionsBuilder!
4040
Npgsql.NpgsqlTracingOptionsBuilder.EnableFirstResponseEvent(bool enable = true) -> Npgsql.NpgsqlTracingOptionsBuilder!
41+
Npgsql.NpgsqlTracingOptionsBuilder.EnablePhysicalOpenTracing(bool enable = true) -> Npgsql.NpgsqlTracingOptionsBuilder!
4142
Npgsql.NpgsqlTypeLoadingOptionsBuilder
4243
Npgsql.NpgsqlTypeLoadingOptionsBuilder.EnableTableCompositesLoading(bool enable = true) -> Npgsql.NpgsqlTypeLoadingOptionsBuilder!
4344
Npgsql.NpgsqlTypeLoadingOptionsBuilder.EnableTypeLoading(bool enable = true) -> Npgsql.NpgsqlTypeLoadingOptionsBuilder!

0 commit comments

Comments
 (0)