feat(csharp/src/Drivers/Apache): tracing - work-in-progress

csharp/src/Apache.Arrow.Adbc/Apache.Arrow.Adbc.csproj

@@ -7,6 +7,8 @@
   </PropertyGroup>
   <ItemGroup>
     <PackageReference Include="Apache.Arrow" Version="18.0.0" />
+    <PackageReference Include="System.Diagnostics.DiagnosticSource" Version="8.0.1" />
+    <PackageReference Include="System.Text.Json" Version="8.0.5" />
   </ItemGroup>
   <ItemGroup Condition="$([MSBuild]::IsTargetFrameworkCompatible($(TargetFramework), 'net6.0'))">
     <Compile Remove="C\NativeLibrary.cs" />

csharp/src/Apache.Arrow.Adbc/C/CAdbcDriverImporter.cs

@@ -23,6 +23,7 @@
 using System.Text;
 using System.Threading;
 using Apache.Arrow.Adbc.Extensions;
+using Apache.Arrow.Adbc.Tracing;
 using Apache.Arrow.C;
 using Apache.Arrow.Ipc;
 
@@ -399,7 +400,7 @@ public unsafe override AdbcConnection Connect(IReadOnlyDictionary<string, string
 
                         caller.Call(Driver.ConnectionInit, ref nativeConnection, ref _nativeDatabase);
 
-                        result = new ImportedAdbcConnection(_driver, nativeConnection);
+                        result = new ImportedAdbcConnection(_driver, nativeConnection, options);
                     }
                 }
                 finally
@@ -462,16 +463,18 @@ private unsafe void Dispose(bool disposing)
         /// <summary>
         /// A native implementation of <see cref="AdbcConnection"/>
         /// </summary>
-        internal sealed class ImportedAdbcConnection : AdbcConnection
+        internal sealed class ImportedAdbcConnection : TracingConnection
         {
+            private static readonly string s_tracingBaseName = typeof(ImportedAdbcConnection).FullName!;
             private readonly ImportedAdbcDriver _driver;
             private CAdbcConnection _nativeConnection;
             private bool _disposed;
             private bool? _autoCommit;
             private IsolationLevel? _isolationLevel;
             private bool? _readOnly;
 
-            internal ImportedAdbcConnection(ImportedAdbcDriver driver, CAdbcConnection nativeConnection)
+            internal ImportedAdbcConnection(ImportedAdbcDriver driver, CAdbcConnection nativeConnection, IReadOnlyDictionary<string, string>? properties)
+                : base(properties)
             {
                 _driver = driver.AddReference();
                 _nativeConnection = nativeConnection;
@@ -521,6 +524,8 @@ public override bool ReadOnly
                 }
             }
 
+            public override string TracingBaseName => s_tracingBaseName;
+
             public unsafe override AdbcStatement CreateStatement()
             {
                 CAdbcStatement nativeStatement = new CAdbcStatement();
@@ -540,7 +545,7 @@ public unsafe override AdbcStatement CreateStatement()
                                 (connection, &nativeStatement, &caller._error));
                         }
 
-                        result = new ImportedAdbcStatement(_driver, nativeStatement);
+                        result = new ImportedAdbcStatement(_driver, nativeStatement, ActivitySource);
                     }
                 }
                 finally
@@ -713,7 +718,7 @@ public override void Dispose()
                 GC.SuppressFinalize(this);
             }
 
-            private unsafe void Dispose(bool disposing)
+            protected override unsafe void Dispose(bool disposing)
             {
                 if (!_disposed)
                 {
@@ -745,14 +750,16 @@ private unsafe void Dispose(bool disposing)
         /// <summary>
         /// A native implementation of <see cref="AdbcStatement"/>
         /// </summary>
-        internal sealed class ImportedAdbcStatement : AdbcStatement
+        internal sealed class ImportedAdbcStatement : TracingStatement
         {
+            private static readonly string s_tracingBaseName = typeof(ImportedAdbcStatement).FullName!;
             private ImportedAdbcDriver _driver;
             private CAdbcStatement _nativeStatement;
             private byte[]? _substraitPlan;
             private bool _disposed;
 
-            internal ImportedAdbcStatement(ImportedAdbcDriver driver, CAdbcStatement nativeStatement)
+            internal ImportedAdbcStatement(ImportedAdbcDriver driver, CAdbcStatement nativeStatement, ActivitySource? activitySource)
+                : base(activitySource)
             {
                 _driver = driver.AddReference();
                 _nativeStatement = nativeStatement;
@@ -804,6 +811,8 @@ public unsafe override byte[]? SubstraitPlan
                 }
             }
 
+            public override string TracingBaseName => s_tracingBaseName;
+
             public unsafe override void Bind(RecordBatch batch, Schema schema)
             {
                 using (CallHelper caller = new CallHelper())

csharp/src/Apache.Arrow.Adbc/Tracing/ITracingObject.cs

@@ -0,0 +1,46 @@
+/*
+* Licensed to the Apache Software Foundation (ASF) under one or more
+* contributor license agreements.  See the NOTICE file distributed with
+* this work for additional information regarding copyright ownership.
+* The ASF licenses this file to You 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.Diagnostics;
+
+namespace Apache.Arrow.Adbc.Tracing
+{
+    /// <summary>
+    /// Provides interface for a tracing object.
+    /// </summary>
+    public interface ITracingObject
+    {
+        /// <summary>
+        /// Gets the <see cref="System.Diagnostics.ActivitySource"/> for the current object. The value may be null if the containing object
+        /// did not set this property - typically because tracing was not enabled.
+        /// </summary>
+        ActivitySource? ActivitySource { get; }
+
+        /// <summary>
+        /// Starts a new activity with a given name. 
+        /// </summary>
+        /// <param name="methodName"></param>
+        /// <returns>Returns a new <see cref="System.Diagnostics.Activity"/> instance if there is an active listener for the <see cref="ActivitySource"/>.
+        /// It returns <c>null</c>, otherwise.</returns>
+        Activity? StartActivity(string methodName);
+
+        /// <summary>
+        /// Gets the base name of the tracing object. Typically this is the full name of the class that is tracing.
+        /// </summary>
+        string TracingBaseName { get; }
+    }
+}

csharp/src/Apache.Arrow.Adbc/Tracing/TracingArrowArrayStream.cs

@@ -0,0 +1,37 @@
+/*
+* Licensed to the Apache Software Foundation (ASF) under one or more
+* contributor license agreements.  See the NOTICE file distributed with
+* this work for additional information regarding copyright ownership.
+* The ASF licenses this file to You 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.Diagnostics;
+using System.Threading;
+using System.Threading.Tasks;
+using Apache.Arrow.Ipc;
+
+namespace Apache.Arrow.Adbc.Tracing
+{
+    /// <summary>
+    /// Provides tracing capability for objects that need to implement the <see cref="IArrowArrayStream"/> interface.
+    /// </summary>
+    /// <param name="activitySource"></param>
+    public abstract class TracingArrowArrayStream(ActivitySource? activitySource) : TracingObject(activitySource), IArrowArrayStream
+    {
+        public abstract ValueTask<RecordBatch?> ReadNextRecordBatchAsync(CancellationToken cancellationToken = default);
+
+        public abstract void Dispose();
+
+        public abstract Schema Schema { get; }
+    }
+}

csharp/src/Apache.Arrow.Adbc/Tracing/TracingConnection.cs

@@ -0,0 +1,171 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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.Concurrent;
+using System.Collections.Generic;
+using System.Diagnostics;
+using System.IO;
+using System.Reflection;
+using System.Text.Json;
+using System.Threading.Tasks;
+
+namespace Apache.Arrow.Adbc.Tracing
+{
+    /// <summary>
+    /// Provides an <see cref="AdbcConnection"> which can instrument tracing activity.
+    /// </summary>
+    public abstract class TracingConnection : AdbcConnection, ITracingObject
+    {
+        private bool _disposed = false;
+        internal const string ProductVersionDefault = "1.0.0";
+        private static readonly string s_activitySourceName = Assembly.GetExecutingAssembly().GetName().Name!;
+        private static readonly string s_assemblyVersion = GetProductVersion();
+        private static readonly ConcurrentDictionary<string, ActivityListener> s_listeners = new();
+        private readonly ConcurrentQueue<Activity> _activityQueue = new();
+        private readonly IReadOnlyDictionary<string, string>? _options;
+        private DirectoryInfo? _traceDirectory;
+
+        protected TracingConnection(IReadOnlyDictionary<string, string>? options = default)
+        {
+            _options = options ?? new Dictionary<string, string>();
+            EnsureTracing();
+        }
+
+        protected TracingConnection(bool isTracingEnabled, string traceLocation, int traceMaxFileSizeKb, int traceMaxFiles)
+        {
+            var options = new Dictionary<string, string>();
+            options[TracingOptions.Connection.Trace] = isTracingEnabled.ToString();
+            options[TracingOptions.Connection.TraceLocation] = traceLocation;
+            options[TracingOptions.Connection.TraceFileMaxSizeKb] = traceMaxFileSizeKb.ToString();
+            options[TracingOptions.Connection.TraceFileMaxFiles] = traceMaxFiles.ToString();
+            _options = options;
+            EnsureTracing();
+        }
+
+        public ActivitySource? ActivitySource { get; private set; }
+
+        protected static string GetProductVersion()
+        {
+            FileVersionInfo fileVersionInfo = FileVersionInfo.GetVersionInfo(Assembly.GetExecutingAssembly().Location);
+            return fileVersionInfo.ProductVersion ?? ProductVersionDefault;
+        }
+
+        private void EnsureTracing()
+        {
+            if (true || _options.TryGetValue(TracingOptions.Connection.Trace, out string? traceOption) && bool.TryParse(traceOption, out bool traceEnabled))
+            {
+
+                // TODO: Handle exceptions
+                if (_options?.TryGetValue(TracingOptions.Connection.TraceLocation, out string? traceLocation) != true || !Directory.Exists(traceLocation))
+                {
+                    string? traceLocationDefault = Environment.GetFolderPath(Environment.SpecialFolder.UserProfile);
+                    _traceDirectory = new DirectoryInfo(traceLocationDefault);
+                }
+                else
+                {
+                    // TODO: If not exist, try to create
+                    _traceDirectory = new DirectoryInfo(traceLocation);
+                }
+                // TODO: Check if folder is writable
+
+
+                // TODO: Determine the best handling of listener lifetimes.
+                // Key of listeners collection should be ouput file location
+                ActivityListener listener = s_listeners.GetOrAdd(s_activitySourceName + "." + _traceDirectory.FullName, (_) => new()
+                {
+                    ShouldListenTo = (source) => source.Name == s_activitySourceName,
+                    Sample = (ref ActivityCreationOptions<ActivityContext> options) => ActivitySamplingResult.AllDataAndRecorded,
+                    ActivityStarted = OnActivityStarted,
+                    ActivityStopped = OnActivityStopped,
+                    SampleUsingParentId = (ref ActivityCreationOptions<string> options) => ActivitySamplingResult.AllDataAndRecorded,
+                });
+                // This is a singleton add, if the lister is the same.
+                ActivitySource.AddActivityListener(listener);
+                // THis is an new instance and needs to be disposed later
+                ActivitySource = new(s_activitySourceName, s_assemblyVersion);
+            }
+        }
+
+        private void OnActivityStarted(Activity activity)
+        {
+            _activityQueue.Enqueue(activity);
+            // Intentionally avoid await.
+            DequeueAndWrite("started")
+                .ContinueWith(t => Console.WriteLine(t.Exception), TaskContinuationOptions.OnlyOnFaulted);
+        }
+
+        private void OnActivityStopped(Activity activity)
+        {
+            _activityQueue.Enqueue(activity);
+            // Intentionally avoid await.
+            DequeueAndWrite("stopped")
+                .ContinueWith(t => Console.WriteLine(t.Exception), TaskContinuationOptions.OnlyOnFaulted);
+        }
+
+        // TODO: Encapsulate this separately
+        private Task DequeueAndWrite(string state)
+        {
+            if (_activityQueue.TryDequeue(out Activity? activity))
+            {
+                if (activity != null)
+                {
+                    try
+                    {
+                        string json = JsonSerializer.Serialize(new { State = state, Activity = activity });
+                        Console.WriteLine(json);
+                    }
+                    catch (NotSupportedException ex)
+                    {
+                        Console.WriteLine(ex.Message);
+                    }
+                }
+            }
+
+            return Task.CompletedTask;
+        }
+
+        public Activity? StartActivity(string methodName)
+        {
+            return StartActivity(ActivitySource, TracingBaseName, methodName);
+        }
+
+        public abstract string TracingBaseName { get; }
+
+        protected internal static Activity? StartActivity(ActivitySource? activitySource, string typeName, string methodName) =>
+            activitySource?.StartActivity(typeName + "." + methodName);
+
+        protected virtual void Dispose(bool disposing)
+        {
+            if (!_disposed)
+            {
+                if (disposing)
+                {
+                    ActivitySource?.Dispose();
+                }
+                _disposed = true;
+            }
+        }
+
+        public override void Dispose()
+        {
+            Dispose(true);
+            base.Dispose();
+            GC.SuppressFinalize(this);
+        }
+    }
+}

csharp/src/Apache.Arrow.Adbc/Tracing/TracingObject.cs

@@ -0,0 +1,37 @@
+/*
+* Licensed to the Apache Software Foundation (ASF) under one or more
+* contributor license agreements.  See the NOTICE file distributed with
+* this work for additional information regarding copyright ownership.
+* The ASF licenses this file to You 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.Diagnostics;
+
+namespace Apache.Arrow.Adbc.Tracing
+{
+    /// <summary>
+    /// Provides an implementation of the <see cref="ITracingObject"/> interface.
+    /// </summary>
+    /// <param name="activitySource">The <see cref="System.Diagnostics.ActivitySource"/> to trace on.</param>
+    public abstract class TracingObject(ActivitySource? activitySource) : ITracingObject
+    {
+        public ActivitySource? ActivitySource { get; private set; } = activitySource;
+
+        public Activity? StartActivity(string methodName)
+        {
+            return TracingConnection.StartActivity(ActivitySource, TracingBaseName, methodName);
+        }
+
+        public abstract string TracingBaseName { get; }
+    }
+}