diff --git a/csharp/src/DatabricksConnection.cs b/csharp/src/DatabricksConnection.cs index 5ec2ecbf3..4d361e58c 100644 --- a/csharp/src/DatabricksConnection.cs +++ b/csharp/src/DatabricksConnection.cs @@ -116,6 +116,29 @@ internal class DatabricksConnection : SparkHttpConnection // Telemetry private IConnectionTelemetry _telemetry = NoOpConnectionTelemetry.Instance; + + // Issue #477: long-lived "lifetime" Activity that wraps every per-call activity + // emitted on this connection. Opened at the end of the constructor (after + // ValidateProperties succeeds) so it becomes Activity.Current for the synchronous + // Database.Connect caller — that way OpenAsync().Wait(), ApplyServerSidePropertiesAsync().Wait(), + // and every subsequent call (executes, metadata, dispose) automatically chains + // underneath via Activity.Current. Stopped in Dispose AFTER base.Dispose runs + // DisposeClient so the close-session RPC also lands under this root. + // + // Without this wrapper every top-level driver call is its own root TraceId — the + // MemoryStress baseline in #477 shows 12,958 TraceIds across one run, 95% of which + // are single-span throwaway ReadNextRecordBatchAsync roots, breaking trace-tree + // visualization on fleet dashboards. + // + // STATIC ActivitySource is intentional: the per-instance source held by + // is disposed inside TracingConnection.Dispose + // BEFORE returns, which would silently drop the + // Stop() for our lifetime Activity (no listener notification once the + // source is disposed). A process-wide source survives connection disposal and + // keeps notifying listeners. + private static readonly ActivitySource s_lifetimeActivitySource = new ActivitySource("AdbcDrivers.Databricks"); + private Activity? _lifetimeActivity; + private bool _lifetimeActivityStopped; // Stopwatch covering the connection lifetime; started at construction and used to // measure session-open latency for the CREATE_SESSION telemetry event. The wall-clock // time between construction and HandleOpenSessionResponse encompasses transport setup @@ -157,8 +180,30 @@ internal DatabricksConnection( Lz4BufferPool = lz4BufferPool ?? System.Buffers.ArrayPool.Create(maxArrayLength: 4 * 1024 * 1024, maxArraysPerBucket: 10); ValidateProperties(); + + // Issue #477: open the connection-lifetime Activity now, AFTER ValidateProperties + // succeeds, so a throw from validation doesn't leak a half-open Activity. The + // call site (DatabricksDatabase.Connect) runs synchronously, so starting the + // Activity here makes it Activity.Current for the same async context that + // immediately invokes OpenAsync().Wait() and ApplyServerSidePropertiesAsync().Wait(); + // those per-call activities therefore chain underneath via Activity.Current with + // no extra plumbing. Kind=Internal because this is a self-traced scope, not an RPC. + _lifetimeActivity = s_lifetimeActivitySource.StartActivity( + "DatabricksConnection.Lifetime", + ActivityKind.Internal); } + /// + /// The long-lived opened in the constructor that wraps + /// every per-call activity emitted on this connection. Exposed (internal) so + /// can pass + /// explicitly as the parent context to its own lifetime Activity — that guards + /// against the case where the Statement is constructed on a different async + /// context than where the Connection was opened (so + /// would not see this Activity). + /// + internal Activity? LifetimeActivity => _lifetimeActivity; + private void LogConnectionProperties(Activity? activity) { if (activity == null) return; @@ -1100,6 +1145,15 @@ protected override void Dispose(bool disposing) // This is synchronous because Dispose() cannot be async; we block on // GetAwaiter().GetResult(), which is acceptable in Dispose. DisposeTelemetryAsync().GetAwaiter().GetResult(); + + // Issue #477: stop the connection-lifetime Activity LAST, after + // base.Dispose has already run the TCloseSessionReq RPC (which emits + // HiveServer2Connection.DisposeClient as a child of this Activity) and + // after DELETE_SESSION telemetry has flushed. Stopping it earlier would + // re-orphan DisposeClient back to a root span — exactly the case + // ConcurrentDisposeTests flags in #477. Stop is idempotent across + // repeated Dispose() calls. + StopLifetimeActivity(); } return; } @@ -1107,6 +1161,20 @@ protected override void Dispose(bool disposing) base.Dispose(disposing); } + /// + /// Idempotently stops the connection-lifetime Activity (#477). Called from + /// ; calling it from a finalizer / unmanaged-only path + /// is intentionally a no-op because Activity stop is not async-signal-safe and we + /// already accept that GC'd Activity objects leak (their state is GC-tracked). + /// + private void StopLifetimeActivity() + { + if (_lifetimeActivityStopped) return; + _lifetimeActivityStopped = true; + _lifetimeActivity?.Stop(); + _lifetimeActivity = null; + } + /// /// Emits the DELETE_SESSION telemetry event when a session was previously opened. /// is the measured duration of base.Dispose, diff --git a/csharp/src/DatabricksStatement.cs b/csharp/src/DatabricksStatement.cs index a5313de89..ff7627ec2 100644 --- a/csharp/src/DatabricksStatement.cs +++ b/csharp/src/DatabricksStatement.cs @@ -89,6 +89,28 @@ internal class DatabricksStatement : SparkStatement, IHiveServer2Statement private readonly Stopwatch _statementLifetimeStopwatch = Stopwatch.StartNew(); private bool _closeStatementTelemetryEmitted; + // Issue #477: long-lived "lifetime" Activity that wraps every per-call activity + // emitted on behalf of this statement. Parented explicitly to the connection's + // own lifetime Activity (see DatabricksConnection.LifetimeActivity) so the + // hierarchy is robust even when the Statement is constructed on a different + // async context than where the Connection was opened — in that case + // Activity.Current here would NOT see the connection's lifetime Activity. + // Stopped in Dispose AFTER the statement's CLOSE_STATEMENT telemetry has been + // emitted so any close-side spans (DatabricksCompositeReader.Dispose, the + // CloseStatement RPC) chain underneath this scope. + // + // STATIC ActivitySource is intentional: the per-instance source held by + // is shared with the connection and disposed + // inside TracingConnection.Dispose BEFORE + // returns. Listeners registered against that source stop receiving + // notifications once it's disposed, so a Stop() called from the statement + // Dispose path (which might run alongside or after connection disposal in + // pooled scenarios) would silently drop. A process-wide static source + // survives both lifecycles. + private static readonly ActivitySource s_lifetimeActivitySource = new ActivitySource("AdbcDrivers.Databricks"); + private Activity? _lifetimeActivity; + private bool _lifetimeActivityStopped; + // Populated by DatabricksCompositeReader.Dispose when it issues TCloseOperationReq // to the server. Lets us report the actual close-RPC latency in the CLOSE_STATEMENT // telemetry event emitted at statement Dispose, instead of a meaningless 0. @@ -127,6 +149,21 @@ public DatabricksStatement(DatabricksConnection connection) { SetOption(ApacheParameters.PollTimeMilliseconds, DatabricksConstants.DefaultAsyncExecPollIntervalMs.ToString()); } + + // Issue #477: open the statement-lifetime Activity here. Pass the connection's + // lifetime Activity Context explicitly as the parent so the hierarchy is + // robust against cross-async-context construction — Activity.Current is + // AsyncLocal, so if CreateStatement is invoked on a continuation that didn't + // observe the connection-open context, Activity.Current would NOT see + // ConnectionLifetime. The default ActivityContext (when LifetimeActivity is + // null, e.g. no listener was attached when the connection was opened) is + // "no parent", which still lets the statement-lifetime Activity be its own + // root and parent all per-call activities under one TraceId. + var connectionLifetimeContext = connection.LifetimeActivity?.Context ?? default; + _lifetimeActivity = s_lifetimeActivitySource.StartActivity( + "DatabricksStatement.Lifetime", + ActivityKind.Internal, + connectionLifetimeContext); } private StatementTelemetryContext? CreateTelemetryContext(Telemetry.Proto.Statement.Types.Type statementType) @@ -1346,6 +1383,20 @@ protected override void Dispose(bool disposing) } } base.Dispose(disposing); + + // Issue #477: stop the statement-lifetime Activity LAST so any spans emitted + // during base.Dispose (e.g. the per-statement CloseStatement RPC, the + // composite-reader's CloseOperation, etc.) chain underneath this scope. + // Idempotent across repeated Dispose() calls. + if (disposing) + { + if (!_lifetimeActivityStopped) + { + _lifetimeActivityStopped = true; + _lifetimeActivity?.Stop(); + _lifetimeActivity = null; + } + } } /// diff --git a/csharp/test/E2E/Telemetry/ConnectionStatementLifetimeTraceTests.cs b/csharp/test/E2E/Telemetry/ConnectionStatementLifetimeTraceTests.cs new file mode 100644 index 000000000..b5d3ff9e3 --- /dev/null +++ b/csharp/test/E2E/Telemetry/ConnectionStatementLifetimeTraceTests.cs @@ -0,0 +1,254 @@ +/* +* Copyright (c) 2025 ADBC Drivers Contributors +* +* Licensed under the Apache License, Version 2.0 (the "License"); +* you may not use this file except in compliance with the License. +* You may obtain a copy of the License at +* +* http://www.apache.org/licenses/LICENSE-2.0 +* +* Unless required by applicable law or agreed to in writing, software +* distributed under the License is distributed on an "AS IS" BASIS, +* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +* See the License for the specific language governing permissions and +* limitations under the License. +*/ + +using System; +using System.Collections.Generic; +using System.Diagnostics; +using System.Linq; +using System.Threading.Tasks; +using Apache.Arrow; +using Apache.Arrow.Adbc; +using Apache.Arrow.Adbc.Tests; +using Xunit; +using Xunit.Abstractions; + +namespace AdbcDrivers.Databricks.Tests.E2E.Telemetry +{ + /// + /// E2E regression test for issue #477: every top-level driver call creates a + /// fresh root with a new TraceId, leaving fleet + /// trace-tree visualization broken. The MemoryStress run referenced in the + /// issue shows 12,316 of 12,958 TraceIds (95%) as single-span throwaway + /// ReadNextRecordBatchAsync roots. + /// + /// Fix: introduce a long-lived DatabricksConnection.Lifetime activity + /// that wraps the entire connection, and a long-lived + /// DatabricksStatement.Lifetime activity that wraps each statement. + /// Existing per-call activities (HiveServer2Connection.OpenAsync, + /// HiveServer2Statement.ExecuteStatementAsync, ReadNextRecordBatchAsync, + /// DatabricksCompositeReader.Dispose, HiveServer2Connection.DisposeClient) + /// chain underneath via Activity.Current, so every call within one + /// connection-statement session shares the same TraceId and can be + /// visualized as a tree. + /// + /// This test captures every activity emitted on + /// AdbcDrivers.Databricks (and AdbcDrivers.HiveServer2) during + /// a tiny SELECT-1 workflow and asserts the parent relationships. Before + /// the fix, ConnectionLifetime / StatementLifetime spans do not exist, so + /// the test fails with "no DatabricksConnection.Lifetime span found". + /// + public class ConnectionStatementLifetimeTraceTests : TestBase, IDisposable + { + private readonly List _capturedActivities = new(); + private readonly object _capturedLock = new(); + private readonly ActivityListener _activityListener; + private bool _disposed; + + public ConnectionStatementLifetimeTraceTests(ITestOutputHelper? outputHelper) + : base(outputHelper, new DatabricksTestEnvironment.Factory()) + { + Skip.IfNot(Utils.CanExecuteTestConfig(TestConfigVariable)); + + _activityListener = new ActivityListener + { + // Listen to every driver-level ActivitySource. Each connection + // creates its own per-instance source named after the assembly + // (AdbcDrivers.Databricks / AdbcDrivers.HiveServer2 etc.), so + // we cannot match by exact name — match by prefix instead. + ShouldListenTo = source => + source.Name.StartsWith("AdbcDrivers.", StringComparison.Ordinal), + Sample = (ref ActivityCreationOptions options) => ActivitySamplingResult.AllDataAndRecorded, + ActivityStopped = activity => + { + lock (_capturedLock) + { + _capturedActivities.Add(new CapturedActivity( + SourceName: activity.Source.Name, + OperationName: activity.OperationName, + TraceId: activity.TraceId, + SpanId: activity.SpanId, + ParentSpanId: activity.ParentSpanId)); + } + } + }; + ActivitySource.AddActivityListener(_activityListener); + } + + protected override void Dispose(bool disposing) + { + if (!_disposed) + { + if (disposing) + { + _activityListener.Dispose(); + } + _disposed = true; + } + base.Dispose(disposing); + } + + /// + /// Asserts that opening a connection, creating a statement, executing a + /// trivial query, draining the reader, and disposing the statement and + /// connection produces a single coherent trace tree rooted at + /// DatabricksConnection.Lifetime, with + /// DatabricksStatement.Lifetime as a child and every per-call + /// activity chained underneath. + /// + /// Before the fix (issue #477), no Lifetime spans exist at all, so the + /// first "ConnectionLifetime span found" assertion fails. After the + /// fix, every per-call activity that was previously a fresh root will + /// share the ConnectionLifetime TraceId. + /// + [SkippableFact] + public async Task ConnectionLifetime_ParentsStatementAndCalls_Issue477() + { + lock (_capturedLock) { _capturedActivities.Clear(); } + + var parameters = new Dictionary + { + [DatabricksParameters.Protocol] = "thrift", + }; + + // Drive a minimal but complete connection -> statement -> read -> dispose + // workflow so we capture every per-call activity at least once. + using (var connection = NewConnection(TestConfiguration, parameters)) + { + using var statement = connection.CreateStatement(); + statement.SqlQuery = "SELECT 1 AS val"; + var result = await statement.ExecuteQueryAsync(); + + long rows = 0; + using (var reader = result.Stream!) + { + RecordBatch? batch; + while ((batch = await reader.ReadNextRecordBatchAsync()) != null) + { + rows += batch.Length; + } + } + OutputHelper?.WriteLine($"Drained {rows} rows."); + } + + // Snapshot what the listener captured. This is what fleet dashboards + // would see today. + List activities; + lock (_capturedLock) + { + activities = _capturedActivities.ToList(); + } + + OutputHelper?.WriteLine($"Captured {activities.Count} activities:"); + foreach (var a in activities) + { + OutputHelper?.WriteLine($" source={a.SourceName} op={a.OperationName} trace={a.TraceId} span={a.SpanId} parent={a.ParentSpanId}"); + } + + // ---- Assertion 1: ConnectionLifetime span exists and is a root. ---- + var connectionLifetime = activities + .FirstOrDefault(a => a.OperationName == "DatabricksConnection.Lifetime"); + Assert.True( + connectionLifetime != null, + "Issue #477: expected a long-lived span named 'DatabricksConnection.Lifetime' " + + "to wrap every per-call activity emitted for this connection. " + + $"Found {activities.Count} activities, none with that name. " + + "Without this wrapper, every top-level driver call (OpenAsync, " + + "ExecuteStatementAsync, ReadNextRecordBatchAsync, DisposeClient, ...) " + + "is a fresh TraceId — MemoryStress shows 95% of TraceIds are " + + "single-span throwaway ReadNextRecordBatchAsync roots."); + Assert.Equal(default(ActivitySpanId), connectionLifetime!.ParentSpanId); + + // ---- Assertion 2: StatementLifetime span exists and is a child of ConnectionLifetime. ---- + var statementLifetime = activities + .FirstOrDefault(a => a.OperationName == "DatabricksStatement.Lifetime"); + Assert.True( + statementLifetime != null, + "Issue #477: expected a long-lived span named 'DatabricksStatement.Lifetime' " + + "to wrap every per-call activity emitted for this statement (Execute, " + + "ReadNext, Dispose). Found no such span."); + Assert.Equal(connectionLifetime.SpanId, statementLifetime!.ParentSpanId); + Assert.Equal(connectionLifetime.TraceId, statementLifetime.TraceId); + + // ---- Assertion 3: at least one per-call activity is a child of ConnectionLifetime. ---- + // ApplyServerSidePropertiesAsync runs synchronously inside Database.Connect after + // the constructor returns, so it sits directly under the lifetime. If the lifetime + // isn't on the Activity.Current stack, this never chains. + var connectionChildren = activities + .Where(a => a.ParentSpanId == connectionLifetime.SpanId && + a.OperationName != "DatabricksStatement.Lifetime") + .ToList(); + Assert.True( + connectionChildren.Count > 0, + "Issue #477: expected at least one per-call activity (e.g. " + + "ApplyServerSidePropertiesAsync, CreateStatement, DisposeClient) to be a " + + "child of DatabricksConnection.Lifetime. None were. This means the " + + "Lifetime activity is not on the Activity.Current stack when subsequent " + + "calls fire, defeating the purpose of the wrapper. " + + $"Captured operations: [{string.Join(", ", activities.Select(a => a.OperationName).Distinct())}]."); + + // ---- Assertion 4: at least one per-call activity is a child of StatementLifetime. ---- + // ExecuteStatementAsync or ReadNextRecordBatchAsync are the typical signals — but + // ReadNext lives on the reader, which runs in a different async context than the + // statement ctor. ExecuteStatementAsync is the strongest signal that the + // statement's lifetime activity is properly threading through. + var statementChildren = activities + .Where(a => a.ParentSpanId == statementLifetime.SpanId) + .ToList(); + Assert.True( + statementChildren.Count > 0, + "Issue #477: expected at least one per-call activity (Execute, ReadNext, " + + "GetSchema, ...) to be a child of DatabricksStatement.Lifetime. None were. " + + $"Captured operations under StatementLifetime: [{string.Join(", ", statementChildren.Select(a => a.OperationName))}]; " + + $"All captured operations: [{string.Join(", ", activities.Select(a => a.OperationName).Distinct())}]."); + + // ---- Assertion 5: every per-connection captured activity shares the ConnectionLifetime TraceId. ---- + // This is the headline assertion for the fleet visualization concern in #477. + // Today, every per-call activity has its own TraceId; after the fix every activity + // on the connection-scope source (AdbcDrivers.Databricks) shares one root TraceId. + // + // The FeatureFlagCache source (AdbcDrivers.Databricks.FeatureFlagCache) is + // intentionally NOT under the connection scope — it runs in DatabricksDatabase.Connect + // BEFORE the DatabricksConnection constructor (and therefore before + // ConnectionLifetime opens), so it has its own TraceId. That's fine; it's not part + // of the per-connection visualization that #477 fixes. Exclude it from the + // singleton-TraceId check to avoid coupling this test to that orthogonal source. + var connectionScopeActivities = activities + .Where(a => a.SourceName == "AdbcDrivers.Databricks") + .ToList(); + var distinctTraceIds = connectionScopeActivities.Select(a => a.TraceId).Distinct().ToList(); + Assert.True( + distinctTraceIds.Count == 1, + $"Issue #477: expected a single TraceId across the entire connection " + + $"lifetime (so trace-tree visualization works), got {distinctTraceIds.Count} " + + $"distinct TraceIds on source 'AdbcDrivers.Databricks': [{string.Join(", ", distinctTraceIds)}]. " + + "MemoryStress baseline shows 12,958 TraceIds across one test run — the " + + "Lifetime activities collapse that to one TraceId per connection."); + // Sanity check: ConnectionLifetime's TraceId is the one shared by everyone else. + Assert.Equal(connectionLifetime.TraceId, distinctTraceIds[0]); + } + + /// + /// Captured snapshot of a single Activity. We snapshot at ActivityStopped time + /// because the framework recycles Activity state shortly after Stop. + /// + private sealed record CapturedActivity( + string SourceName, + string OperationName, + ActivityTraceId TraceId, + ActivitySpanId SpanId, + ActivitySpanId ParentSpanId); + } +}