Skip to content

fix(csharp/src/Drivers/Databricks): Make mandatory token exchange blocking and fix concurrent handling #3715

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
CurtHagenlocher merged 3 commits into from
Nov 17, 2025
Merged

fix(csharp/src/Drivers/Databricks): Make mandatory token exchange blocking and fix concurrent handling #3715

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
cc7824b
fix(csharp/src/Drivers/Databricks): Make mandatory token exchange blo…
eric-wang-1990 Nov 15, 2025
8dea1fc
fix(csharp/src/Drivers/Databricks): Fix concurrent token exchange han…
eric-wang-1990 Nov 15, 2025
ccd334b
refactor: Remove unused _tokenBeingExchanged field
eric-wang-1990 Nov 15, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
125 changes: 61 additions & 64 deletions csharp/src/Drivers/Databricks/Auth/MandatoryTokenExchangeDelegatingHandler.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,8 @@ namespace Apache.Arrow.Adbc.Drivers.Databricks.Auth
{
/// <summary>
/// HTTP message handler that performs mandatory token exchange for non-Databricks tokens.
/// Uses a non-blocking approach to exchange tokens in the background.
/// Blocks requests while exchanging tokens to ensure the exchanged token is used.
/// Falls back to the original token if the exchange fails.
/// </summary>
internal class MandatoryTokenExchangeDelegatingHandler : DelegatingHandler
{
Expand All @@ -34,8 +35,7 @@ internal class MandatoryTokenExchangeDelegatingHandler : DelegatingHandler
private readonly ITokenExchangeClient _tokenExchangeClient;
private string? _currentToken;
private string? _lastSeenToken;

protected Task? _pendingTokenTask = null;
private Task? _pendingExchange = null;

/// <summary>
/// Initializes a new instance of the <see cref="MandatoryTokenExchangeDelegatingHandler"/> class.
Expand All @@ -59,18 +59,6 @@ public MandatoryTokenExchangeDelegatingHandler(
/// <returns>True if token exchange is needed, false otherwise.</returns>
private bool NeedsTokenExchange(string bearerToken)
{
// If we already started exchange for this token, no need to check again
if (_lastSeenToken == bearerToken)
{
return false;
}

// If we already have a pending token task, don't start another exchange
if (_pendingTokenTask != null)
{
return false;
}

// If we can't parse the token as JWT, default to use existing token
if (!JwtTokenDecoder.TryGetIssuer(bearerToken, out string issuer))
{
Expand All @@ -86,56 +74,87 @@ private bool NeedsTokenExchange(string bearerToken)
}

/// <summary>
/// Starts token exchange in the background if needed.
/// Performs token exchange if needed.
/// </summary>
/// <param name="bearerToken">The bearer token to potentially exchange.</param>
/// <param name="cancellationToken">A cancellation token.</param>
private void StartTokenExchangeIfNeeded(string bearerToken, CancellationToken cancellationToken)
private async Task PerformTokenExchangeIfNeeded(string bearerToken, CancellationToken cancellationToken)
{
if (_lastSeenToken == bearerToken)
// Check if we need exchange (no lock needed for this check)
bool needsExchange = NeedsTokenExchange(bearerToken);

if (!needsExchange)
{
lock (_tokenLock)
{
_lastSeenToken = bearerToken;
}
return;
}

bool needsExchange;
// Wait for any pending exchange to complete first (could be for a different token)
Task? exchangeToAwait = null;
lock (_tokenLock)
{
needsExchange = NeedsTokenExchange(bearerToken);

_lastSeenToken = bearerToken;
if (_pendingExchange != null)
{
exchangeToAwait = _pendingExchange;
}
}

if (!needsExchange)
if (exchangeToAwait != null)
{
return;
await exchangeToAwait;
}

// Start token exchange in the background
_pendingTokenTask = Task.Run(async () =>
// Now check if we need to exchange our token
lock (_tokenLock)
{
try
// If this token was already processed (by us or another concurrent request)
if (_lastSeenToken == bearerToken)
{
TokenExchangeResponse response = await _tokenExchangeClient.ExchangeTokenAsync(
bearerToken,
_identityFederationClientId,
cancellationToken);

lock (_tokenLock)
{
_currentToken = response.AccessToken;
}
return;
}
catch (Exception ex)

// Start new exchange for our token
_lastSeenToken = bearerToken;
_pendingExchange = DoExchangeAsync(bearerToken, cancellationToken);
exchangeToAwait = _pendingExchange;
}

await exchangeToAwait;
}

/// <summary>
/// Performs the actual token exchange operation.
/// </summary>
/// <param name="bearerToken">The bearer token to exchange.</param>
/// <param name="cancellationToken">A cancellation token.</param>
private async Task DoExchangeAsync(string bearerToken, CancellationToken cancellationToken)
{
try
{
TokenExchangeResponse response = await _tokenExchangeClient.ExchangeTokenAsync(
bearerToken,
_identityFederationClientId,
cancellationToken);

lock (_tokenLock)
{
System.Diagnostics.Debug.WriteLine($"Mandatory token exchange failed: {ex.Message}");
_currentToken = response.AccessToken;
}
}, cancellationToken).ContinueWith(_ =>
}
catch (Exception ex)
{
System.Diagnostics.Debug.WriteLine($"Mandatory token exchange failed: {ex.Message}. Continuing with original token.");
}
finally
{
lock (_tokenLock)
{
_pendingTokenTask = null;
_pendingExchange = null;
}
}, TaskScheduler.Default);
}
}

/// <summary>
Expand All @@ -149,7 +168,7 @@ protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage
string? bearerToken = request.Headers.Authorization?.Parameter;
if (!string.IsNullOrEmpty(bearerToken))
{
StartTokenExchangeIfNeeded(bearerToken!, cancellationToken);
await PerformTokenExchangeIfNeeded(bearerToken!, cancellationToken);

string tokenToUse;
lock (_tokenLock)
Expand All @@ -163,27 +182,5 @@ protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage
return await base.SendAsync(request, cancellationToken);
}

protected override void Dispose(bool disposing)
{
if (disposing)
{
// Wait for any pending token task to complete to avoid leaking tasks
if (_pendingTokenTask != null)
{
try
{
// Try to wait for the task to complete, but don't block indefinitely
_pendingTokenTask.Wait(TimeSpan.FromSeconds(10));
}
catch (Exception ex)
{
// Log any exceptions during disposal
System.Diagnostics.Debug.WriteLine($"Exception during token task cleanup: {ex.Message}");
}
}
}

base.Dispose(disposing);
}
}
}
41 changes: 27 additions & 14 deletions csharp/test/Drivers/Databricks/Unit/Auth/MandatoryTokenExchangeDelegatingHandlerTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
*/

using System;
using System.Linq;
using System.Net;
using System.Net.Http;
using System.Text;
Expand Down Expand Up @@ -125,7 +126,7 @@ public async Task SendAsync_WithDatabricksToken_UsesTokenDirectlyWithoutExchange
}

[Fact]
public async Task SendAsync_WithExternalToken_StartsTokenExchangeInBackground()
public async Task SendAsync_WithExternalToken_BlocksUntilTokenExchangeCompletes()
{
var tokenExchangeDelay = TimeSpan.FromMilliseconds(500);
var handler = new MandatoryTokenExchangeDelegatingHandler(
Expand Down Expand Up @@ -165,23 +166,20 @@ public async Task SendAsync_WithExternalToken_StartsTokenExchangeInBackground()

var httpClient = new HttpClient(handler);

// First request should use original token and start background exchange
// First request should block until token exchange completes, then use exchanged token
var startTime = DateTime.UtcNow;
var response = await httpClient.SendAsync(request);
var requestDuration = DateTime.UtcNow - startTime;

Assert.Equal(expectedResponse, response);
Assert.True(requestDuration < tokenExchangeDelay,
$"Request took {requestDuration.TotalMilliseconds}ms, which is longer than the token exchange delay of {tokenExchangeDelay.TotalMilliseconds}ms");
Assert.True(requestDuration >= tokenExchangeDelay,
$"Request took {requestDuration.TotalMilliseconds}ms, which is shorter than the token exchange delay of {tokenExchangeDelay.TotalMilliseconds}ms. Expected blocking behavior.");

Assert.NotNull(capturedRequest);
Assert.Equal("Bearer", capturedRequest.Headers.Authorization?.Scheme);
Assert.Equal(_externalToken, capturedRequest.Headers.Authorization?.Parameter); // First request uses original token
Assert.Equal(_exchangedToken, capturedRequest.Headers.Authorization?.Parameter); // First request uses exchanged token

// Wait for background task to complete
await Task.Delay(tokenExchangeDelay + TimeSpan.FromMilliseconds(1000));

// Make a second request - this should use the exchanged token
// Make a second request - this should also use the exchanged token (cached)
var request2 = new HttpRequestMessage(HttpMethod.Get, "https://example.com/2");
request2.Headers.Authorization = new System.Net.Http.Headers.AuthenticationHeaderValue("Bearer", _externalToken);
HttpRequestMessage? capturedRequest2 = null;
Expand All @@ -198,8 +196,9 @@ public async Task SendAsync_WithExternalToken_StartsTokenExchangeInBackground()

Assert.NotNull(capturedRequest2);
Assert.Equal("Bearer", capturedRequest2.Headers.Authorization?.Scheme);
Assert.Equal(_exchangedToken, capturedRequest2.Headers.Authorization?.Parameter); // Second request uses exchanged token
Assert.Equal(_exchangedToken, capturedRequest2.Headers.Authorization?.Parameter); // Second request uses cached exchanged token

// Token exchange should only be called once
_mockTokenExchangeClient.Verify(
x => x.ExchangeTokenAsync(_externalToken, _identityFederationClientId, It.IsAny<CancellationToken>()),
Times.Once);
Expand Down Expand Up @@ -396,11 +395,15 @@ public async Task SendAsync_WithConcurrentRequestsSameToken_ExchangesTokenOnlyOn
ExpiryTime = DateTime.UtcNow.AddHours(1)
};

// Add a small delay to token exchange to simulate concurrent access
var exchangeCallCount = 0;
var capturedRequests = new System.Collections.Concurrent.ConcurrentBag<HttpRequestMessage>();

// Add a delay to token exchange to ensure concurrent requests arrive while exchange is in progress
_mockTokenExchangeClient
.Setup(x => x.ExchangeTokenAsync(_externalToken, _identityFederationClientId, It.IsAny<CancellationToken>()))
.Returns(async () =>
{
Interlocked.Increment(ref exchangeCallCount);
await Task.Delay(200);
return tokenExchangeResponse;
});
Expand All @@ -410,6 +413,7 @@ public async Task SendAsync_WithConcurrentRequestsSameToken_ExchangesTokenOnlyOn
"SendAsync",
ItExpr.IsAny<HttpRequestMessage>(),
ItExpr.IsAny<CancellationToken>())
.Callback<HttpRequestMessage, CancellationToken>((req, ct) => capturedRequests.Add(req))
.ReturnsAsync(new HttpResponseMessage(HttpStatusCode.OK));

var httpClient = new HttpClient(handler);
Expand All @@ -424,13 +428,22 @@ public async Task SendAsync_WithConcurrentRequestsSameToken_ExchangesTokenOnlyOn

await Task.WhenAll(tasks);

// Wait for any background token exchange to complete
await Task.Delay(1000);

// Token exchange should only be called once despite concurrent requests
_mockTokenExchangeClient.Verify(
x => x.ExchangeTokenAsync(_externalToken, _identityFederationClientId, It.IsAny<CancellationToken>()),
Times.Once);

Assert.Equal(1, exchangeCallCount);

// All requests should have been sent
Assert.Equal(3, capturedRequests.Count);

// All concurrent requests should use the exchanged token
// (they all wait for the same _pendingExchange task)
foreach (var request in capturedRequests)
{
Assert.Equal(_exchangedToken, request.Headers.Authorization?.Parameter);
}
}

[Fact]
Expand Down
Loading