Reducing SQL Connection Latency for Apps Using Azure AAD Authentication

Challenge: connection latency and token overhead 

Consider a cloud-native application deployed in Azure App Service or Kubernetes (AKS) that needs to query an Azure SQL Database for real-time data. The application uses Azure Active Directory (AAD) for secure authentication, but every time the application establishes a new connection to the database, it requests a new AAD token. In high-traffic environments where thousands of requests are processed per second, this repetitive token issuance introduces latency and performance degradation. This delay becomes particularly problematic for time-sensitive applications where every millisecond counts. Each token request impacts response times and creates unnecessary resource consumption. 

Solution: token caching and expiration management 

To mitigate these delays, we can optimize the authentication process by caching the AAD token and reusing it for the duration of its validity (typically 1 hour to 24 hours). Instead of requesting a new token for every database connection, the token is fetched only when the existing one is near expiration. This approach eliminates the repeated authentication overhead and ensures that the application can maintain seamless connectivity to the database without the performance hit of generating a new token for each request. In addition to reducing latency, this approach reduces the number of HTTP calls made to the Azure Active Directory service, resulting in better resource utilization and lower operational costs. 

Concrete performance gains: optimized SQL client connection

As part of the mitigation, we provide a custom code implementation that uses SqlClient, a supported library, to optimize the connection time. The test was conducted with the S0 database, where using a single process and using connection pooling, we opened a connection, executed the SELECT 1, and closed the connection.

During the testing phase with a connection pooler script running for 96 hours (without the AAD token cache), the following results were observed:

• 10 connections took 1 second, representing 0.866% of total connections.
• 1 connection took 4 seconds, representing 0.0866%.
• 1.144 connections took less than 1 second, representing 99.05% of total connections.
• All executions of SELECT 1 were completed in 0 seconds.

These results demonstrate how caching AAD tokens and reusing them effectively reduced connection overhead and improved performance. None of the connections exceeded 5 seconds in duration, while with the default behavior, connections were reaching 30 seconds and more, depending on the environment complexity.

Step-by-step implementation 

Here’s a step-by-step guide on how to implement this solution using C# and the Microsoft.Data.SqlClient package to optimize SQL database connections: 

1. Obtain and cache a token: Instead of requesting a new AAD token with every connection, we obtain a token once and cache it. This is done by leveraging Azure Managed Identity to authenticate the application, which eliminates the need to repeatedly authenticate with Azure Active Directory for every database connection. In this step, we fetch the token once and store it securely for reuse. 
   
   
2. Renew the token only when it’s near expiry We will refresh the token only when it is nearing expiration or has already expired. The application checks the token’s expiration time before attempting to use it. If the token is still valid, it continues to be reused. If it's close to expiration, a new token is fetched. 
    
3. Reuse a single token across multiple connections: The cached token can be used for multiple database connections during its lifetime. Rather than requesting a new token for each new connection, the application will use the same token across all connections until the token is about to expire. 

Code example: optimized SQL connection management 

Here’s an example of how you can implement token caching in a C# application using Microsoft.Data.SqlClient.

using System;
using System.Data.SqlClient;
using System.Diagnostics;
using System.Threading;
using Azure.Identity;

namespace SqlConnectionOptimization
{
    public class SqlConnectionManager
    {
        private string _accessToken;
        private DateTimeOffset _tokenExpiration;
        private readonly string _connectionString = "Server=tcp:servername.database.windows.net,1433;Initial Catalog=DBName;...";
        private readonly Stopwatch _stopwatch = new Stopwatch();

        public SqlConnectionManager()
        {
            _accessToken = string.Empty;
            _tokenExpiration = DateTimeOffset.UtcNow;
        }

        public void Run()
        {
            while (true)
            {
                // Refresh token if necessary
                if (IsTokenExpired())
                {
                    RefreshToken();
                }

                // Establish connection and perform operations
                using (var connection = CreateConnection())
                {
                    LogExecutionTime("Connected");
                    ExecuteQuery(connection);
                    LogExecutionTime("Query Executed");
                }

                // Simulate some idle time between operations
                Log("Waiting before next operation...");
                Thread.Sleep(1000);
            }
        }

        private bool IsTokenExpired()
        {
            return string.IsNullOrEmpty(_accessToken) || DateTimeOffset.UtcNow.AddMinutes(5) >= _tokenExpiration;
        }

        private void RefreshToken()
        {
            _stopwatch.Start();
            try
            {
                var result = FetchAccessToken();
                _accessToken = result.Token;
                _tokenExpiration = result.Expiration;
                LogExecutionTime("Token Refreshed");
                Log($"Token expires at: {_tokenExpiration}");
            }
            catch (Exception ex)
            {
                Log($"Error fetching token: {ex.Message}");
            }
        }

        private (string Token, DateTimeOffset Expiration) FetchAccessToken()
        {
            var managedIdentityCredential = new ManagedIdentityCredential();
            var tokenRequestContext = new Azure.Core.TokenRequestContext(new[] { "https://database.windows.net/" });
            var accessToken = managedIdentityCredential.GetTokenAsync(tokenRequestContext).Result;

            return (accessToken.Token, accessToken.ExpiresOn.UtcDateTime);
        }

        private SqlConnection CreateConnection()
        {
            var connection = new SqlConnection(_connectionString)
            {
                AccessToken = _accessToken
            };

            int retries = 0;
            while (true)
            {
                try
                {
                    connection.Open();
                    return connection;
                }
                catch (Exception ex)
                {
                    retries++;
                    if (retries > 5)
                    {
                        Log($"Error connecting after multiple retries: {ex.Message}");
                        throw;
                    }

                    Log($"Connection attempt failed. Retrying in {retries} seconds...");
                    Thread.Sleep(retries * 1000);
                }
            }
        }

        private void ExecuteQuery(SqlConnection connection)
        {
            var query = "SELECT 1"; // Simple query, replace with real logic as needed
            int retries = 0;

            while (true)
            {
                try
                {
                    using (var command = new SqlCommand(query, connection))
                    {
                        command.CommandTimeout = 5;  // Adjust timeout for more complex queries
                        command.ExecuteScalar();
                    }
                    return;
                }
                catch (Exception ex)
                {
                    retries++;
                    if (retries > 5)
                    {
                        Log($"Max retries reached for query execution: {ex.Message}");
                        throw;
                    }

                    Log($"Query execution failed. Retrying in {retries} seconds...");
                    Thread.Sleep(retries * 1000);
                }
            }
        }

        private void Log(string message)
        {
            Console.WriteLine($"{DateTime.Now:yyyy-MM-dd HH:mm:ss.fff}: {message}");
        }

        private void LogExecutionTime(string action)
        {
            _stopwatch.Stop();
            var elapsed = _stopwatch.Elapsed;
            Log($"{action} - Elapsed time: {elapsed:hh\\:mm\\:ss\\.fff}");
            _stopwatch.Reset();
        }

        public static void Main(string[] args)
        {
            var manager = new SqlConnectionManager();
            manager.Run();
        }
    }
}

Key points in the code

• Token Expiration Check: The IsTokenExpired() method checks whether the token has expired by comparing it to the current time. We’ve added a 5-minute buffer for token expiration. This can be adjusted based on your needs. 

• Managed Identity Authentication: The application uses Azure Managed Identity to authenticate and fetch the token, ensuring secure and scalable access to Azure SQL Database without requiring client secrets. 

• Retry Logic: In the event of a connection failure or query execution failure, the system retries a set number of times with exponential backoff, making it resilient to transient network or authentication issues. 

Conclusion 

By implementing a token caching and expiration management strategy, applications can dramatically improve the performance and scalability of their database interactions, especially in environments with high request volumes. By leveraging Azure Managed Identity for secure, reusable tokens, you can reduce authentication latency and improve the overall efficiency of your SQL database connections.

This approach can also be adapted to any service using Azure SQL Database and Azure Active Directory for authentication. 

Next steps

• Benchmarking: Test the implementation in your environment to quantify the performance gains. 

• Error Handling: Extend the retry logic and error handling to better handle transient failures, especially in production environments. Resources:
  Introducing Configurable Retry Logic in Microsoft.Data.SqlClient v3.0.0-Preview1
  Configurable retry logic in SqlClient
  Troubleshoot transient connection errors

• Scaling: Consider how this strategy can be applied across multiple services in larger architectures.
  
  
• Consider reading and applying managed identity best practices. Resources: Managed identity best practice recommendations