SephamoreSlim lock to prevent simultaneous fetchResults + poll status

Author: toddmeng-db

csharp/src/Drivers/Databricks/CloudFetch/CloudFetchResultFetcher.cs

@@ -21,6 +21,7 @@
 using System.Diagnostics;
 using System.Threading;
 using System.Threading.Tasks;
+using Apache.Arrow.Adbc.Drivers.Databricks;
 using Apache.Hive.Service.Rpc.Thrift;
 
 namespace Apache.Arrow.Adbc.Drivers.Apache.Databricks.CloudFetch
@@ -30,7 +31,7 @@ namespace Apache.Arrow.Adbc.Drivers.Apache.Databricks.CloudFetch
     /// </summary>
     internal sealed class CloudFetchResultFetcher : ICloudFetchResultFetcher
     {
-        private readonly IHiveServer2Statement _statement;
+        private readonly DatabricksStatement _statement;
         private readonly ICloudFetchMemoryBufferManager _memoryManager;
         private readonly BlockingCollection<IDownloadResult> _downloadQueue;
         private long _startOffset;
@@ -49,7 +50,7 @@ internal sealed class CloudFetchResultFetcher : ICloudFetchResultFetcher
         /// <param name="downloadQueue">The queue to add download tasks to.</param>
         /// <param name="prefetchCount">The number of result chunks to prefetch.</param>
         public CloudFetchResultFetcher(
-            IHiveServer2Statement statement,
+            DatabricksStatement statement,
             ICloudFetchMemoryBufferManager memoryManager,
             BlockingCollection<IDownloadResult> downloadQueue,
             long batchSize)
@@ -197,7 +198,8 @@ private async Task FetchNextResultBatchAsync(CancellationToken cancellationToken
             TFetchResultsResp response;
             try
             {
-                response = await _statement.Client.FetchResults(request, cancellationToken).ConfigureAwait(false);
+                // Use thread-safe method to fetch results
+                response = await _statement.FetchResultsAsync(request, cancellationToken).ConfigureAwait(false);
             }
             catch (Exception ex)
             {

csharp/src/Drivers/Databricks/DatabricksOperationStatusPoller.cs

@@ -29,13 +29,13 @@ namespace Apache.Arrow.Adbc.Drivers.Databricks
     /// </summary>
     internal class DatabricksOperationStatusPoller : IDisposable
     {
-        private readonly IHiveServer2Statement _statement;
+        private readonly DatabricksStatement _statement;
         private readonly int _heartbeatIntervalSeconds;
         // internal cancellation token source - won't affect the external token
         private CancellationTokenSource? _internalCts;
         private Task? _operationStatusPollingTask;
 
-        public DatabricksOperationStatusPoller(IHiveServer2Statement statement, int heartbeatIntervalSeconds = DatabricksConstants.DefaultOperationStatusPollingIntervalSeconds)
+        public DatabricksOperationStatusPoller(DatabricksStatement statement, int heartbeatIntervalSeconds = DatabricksConstants.DefaultOperationStatusPollingIntervalSeconds)
         {
             _statement = statement ?? throw new ArgumentNullException(nameof(statement));
             _heartbeatIntervalSeconds = heartbeatIntervalSeconds;
@@ -70,7 +70,8 @@ private async Task PollOperationStatus(CancellationToken cancellationToken)
                     if (operationHandle == null) break;
 
                     var request = new TGetOperationStatusReq(operationHandle);
-                    var response = await _statement.Client.GetOperationStatus(request, cancellationToken);
+                    
+                    var response = await _statement.GetOperationStatusAsync(request, cancellationToken);
                     await Task.Delay(TimeSpan.FromSeconds(_heartbeatIntervalSeconds), cancellationToken);
 
                     // end the heartbeat if the command has terminated

csharp/src/Drivers/Databricks/DatabricksReader.cs

@@ -78,7 +78,7 @@ public DatabricksReader(DatabricksStatement statement, Schema schema, bool isLz4
                 }
 
                 TFetchResultsReq request = new TFetchResultsReq(this.statement.OperationHandle!, TFetchOrientation.FETCH_NEXT, this.statement.BatchSize);
-                TFetchResultsResp response = await this.statement.Connection.Client!.FetchResults(request, cancellationToken);
+                TFetchResultsResp response = await this.statement.FetchResultsAsync(request, cancellationToken);
 
                 // Make sure we get the arrowBatches
                 this.batches = response.Results.ArrowBatches;

csharp/src/Drivers/Databricks/DatabricksStatement.cs

@@ -16,6 +16,8 @@
 */
 
 using System;
+using System.Threading;
+using System.Threading.Tasks;
 using Apache.Arrow.Adbc.Drivers.Apache;
 using Apache.Arrow.Adbc.Drivers.Apache.Spark;
 using Apache.Arrow.Adbc.Drivers.Apache.Databricks.CloudFetch;
@@ -26,11 +28,14 @@ namespace Apache.Arrow.Adbc.Drivers.Databricks
     /// <summary>
     /// Databricks-specific implementation of <see cref="AdbcStatement"/>
     /// </summary>
-    internal class DatabricksStatement : SparkStatement, IHiveServer2Statement
+    internal class DatabricksStatement : SparkStatement, IHiveServer2Statement, IDisposable
     {
         private bool useCloudFetch;
         private bool canDecompressLz4;
         private long maxBytesPerFile;
+        
+        // Semaphore lock to ensure that polling and fetching results do not both use transport at the same time
+        private readonly SemaphoreSlim _clientSemaphore = new SemaphoreSlim(1, 1);
 
         public DatabricksStatement(DatabricksConnection connection)
             : base(connection)
@@ -70,6 +75,58 @@ public TSparkDirectResults? DirectResults
         // Cast the Client to IAsync for CloudFetch compatibility
         TCLIService.IAsync IHiveServer2Statement.Client => Connection.Client;
 
+        /// <summary>
+        /// Executes a client operation in a thread-safe manner.
+        /// </summary>
+        /// <typeparam name="TResult">The type of the result.</typeparam>
+        /// <param name="operation">The operation to execute.</param>
+        /// <param name="cancellationToken">The cancellation token.</param>
+        /// <returns>A task representing the asynchronous operation with the result.</returns>
+        public async Task<TResult> ExecuteClientOperationAsync<TResult>(
+            Func<TCLIService.IAsync, CancellationToken, Task<TResult>> operation,
+            CancellationToken cancellationToken)
+        {
+            await _clientSemaphore.WaitAsync(cancellationToken).ConfigureAwait(false);
+            try
+            {
+                return await operation(Connection.Client, cancellationToken).ConfigureAwait(false);
+            }
+            finally
+            {
+                _clientSemaphore.Release();
+            }
+        }
+
+        /// <summary>
+        /// Gets the operation status in a thread-safe manner.
+        /// </summary>
+        /// <param name="request">The get operation status request.</param>
+        /// <param name="cancellationToken">The cancellation token.</param>
+        /// <returns>A task representing the asynchronous operation with the result.</returns>
+        public Task<TGetOperationStatusResp> GetOperationStatusAsync(
+            TGetOperationStatusReq request,
+            CancellationToken cancellationToken)
+        {
+            return ExecuteClientOperationAsync(
+                (client, token) => client.GetOperationStatus(request, token),
+                cancellationToken);
+        }
+
+        /// <summary>
+        /// Fetches results in a thread-safe manner.
+        /// </summary>
+        /// <param name="request">The fetch results request.</param>
+        /// <param name="cancellationToken">The cancellation token.</param>
+        /// <returns>A task representing the asynchronous operation with the result.</returns>
+        public Task<TFetchResultsResp> FetchResultsAsync(
+            TFetchResultsReq request,
+            CancellationToken cancellationToken)
+        {
+            return ExecuteClientOperationAsync(
+                (client, token) => client.FetchResults(request, token),
+                cancellationToken);
+        }
+
         public override void SetOption(string key, string value)
         {
             switch (key)
@@ -151,5 +208,18 @@ internal void SetMaxBytesPerFile(long maxBytesPerFile)
         {
             this.maxBytesPerFile = maxBytesPerFile;
         }
+        
+        /// <summary>
+        /// Disposes the resources used by the statement.
+        /// </summary>
+        protected override void Dispose(bool disposing)
+        {
+            if (disposing)
+            {
+                _clientSemaphore.Dispose();
+            }
+            
+            base.Dispose(disposing);
+        }
     }
 }