Share source session across migration workers

CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs

@@ -4,11 +4,17 @@
 using CassandraMigrationProcessor.Infrastructure;
 using CassandraMigrationProcessor.Models;
 namespace CassandraMigrationProcessor.CassandraDriver;
+
+internal sealed record SourceSessionSettings(
+    string ContactPoint,
+    int Port,
+    string Username,
+    int MaxConnectionsPerHost);
+
 /// <summary>
 /// Creates Cassandra ISession instances for source (Cosmos DB)
 /// and target (OSS Cassandra) clusters.
-/// Delegates AAD token management to TokenRefreshManager and
-/// ARM credential discovery to ArmCredentialDiscovery.
+/// Delegates ARM credential discovery to ArmCredentialDiscovery.
 /// </summary>
 public static class CassandraClientFactory
 {
@@ -35,8 +41,6 @@ public static class CassandraClientFactory
     /// <summary>
     /// Create a session to a Cosmos DB Cassandra API account.
     /// Uses SSL on port 10350 with PlainTextAuthProvider.
-    /// Starts proactive token refresh if the password is a
-    /// JWT/AAD token.
     /// Retries on 429/OverloadedException with backoff.
     /// </summary>
     public static ISession CreateSourceSession(
@@ -45,13 +49,8 @@ public static ISession CreateSourceSession(
         int port,
         string username,
         string password,
-        TokenRefreshManager? tokenRefreshManager = null,
         int maxConnectionsPerHost = 0)
     {
-        // Cache parameters for token refresh reconnection
-        tokenRefreshManager?.CacheSourceConnectionParams(
-            contactPoint, port, username);
-
         // Source always uses SSL (Cosmos DB requires it)
         var builder = CreateBaseBuilder(
             contactPoint, port, username, password,
@@ -68,9 +67,7 @@ public static ISession CreateSourceSession(
         {
             try
             {
-                var session = ConnectCluster(builder);
-                RegisterAadTokenRefresh(session, password, tokenRefreshManager);
-                return session;
+                return ConnectCluster(builder);
             }
             catch (Exception ex) when (
                 ExceptionClassifier.IsTransient(ex)
@@ -91,24 +88,6 @@ public static ISession CreateSourceSession(
         throw new UnreachableException();
     }
 
-    /// <summary>
-    /// When <paramref name="password"/> looks like an AAD/JWT bearer
-    /// token and the caller wired up a <see cref="TokenRefreshManager"/>,
-    /// hand the freshly-connected <paramref name="session"/> off so the
-    /// proactive refresh timer can rotate the bearer before it expires.
-    /// No-op when the password is a static credential or the manager is
-    /// not supplied.
-    /// </summary>
-    private static void RegisterAadTokenRefresh(
-        ISession session,
-        string password,
-        TokenRefreshManager? tokenRefreshManager)
-    {
-        if (!TokenRefreshManager.IsLikelyAadToken(password)) return;
-        tokenRefreshManager?.SetManagedSourceSession(session);
-        tokenRefreshManager?.StartTokenRefreshTimer(password);
-    }
-
     /// <summary>
     /// Create a session to an OSS Apache Cassandra cluster.
     /// Tries SSL first, falls back to plain if SSL fails.
@@ -302,56 +281,61 @@ private static ISession ConnectCluster(Builder builder)
         }
     }
 
-    /// <summary>
-    /// Create source session from a Job's properties.
-    /// If SourceUseAad is true or password is missing (e.g.
-    /// on resume after [JsonIgnore]), fetches a fresh AAD
-    /// token automatically.
-    /// </summary>
-    public static ISession CreateSourceSession(
-        MigrationLog MigrationLog, Job job,
-        TokenRefreshManager? tokenRefreshManager = null)
+    internal static string AcquireAadToken()
+    {
+        var credential = new Azure.Identity.DefaultAzureCredential();
+        return credential.GetToken(
+            new Azure.Core.TokenRequestContext(
+                new[] { "https://cosmos.azure.com/.default" }))
+            .Token;
+    }
+
+    internal static (
+        SourceSessionSettings Settings,
+        string Credential) ResolveSourceSession(
+        Job job,
+        int workerCount = 0)
     {
         if (string.IsNullOrEmpty(job.SourceContactPoint))
             throw new ArgumentException("Source contact point is required", nameof(job));
 
-        string password = job.SourcePassword ?? string.Empty;
-
-        // If password is empty (resume) or AAD is enabled,
-        // fetch a fresh token via managed identity
-        if (string.IsNullOrEmpty(password) || job.SourceUseAad)
+        bool useAad = job.SourceUseAad
+            || string.IsNullOrEmpty(job.SourcePassword);
+        string credential = job.SourcePassword ?? string.Empty;
+        if (useAad)
         {
-            password = tokenRefreshManager?.GetFreshAadToken()
-                ?? TokenRefreshManager.AcquireAadToken();
-            // SECURITY: do NOT write the AAD bearer token back into
-            // job.SourcePassword — even though [JsonIgnore] keeps it
-            // off disk, the Blazor "Update Connection Strings" modal
-            // would echo it into a <input value="…"> and leak the
-            // bearer JWT to the browser DOM. Azure.Identity caches
-            // tokens in-process so re-acquiring per call is free.
+            credential = AcquireAadToken();
+            // Do not write the bearer token back to SourcePassword. The
+            // connection editor would otherwise expose it in the browser DOM.
             job.SourceUseAad = true;
         }
 
-        // For AAD auth, derive username from hostname if
-        // not explicitly provided (account name = first
-        // segment of the contact point FQDN).
         string username = job.SourceUsername ?? string.Empty;
         if (string.IsNullOrWhiteSpace(username)
-            && job.SourceUseAad
-            && !string.IsNullOrEmpty(job.SourceContactPoint))
+            && useAad)
         {
             username = job.SourceContactPoint
                 .Split('.')[0];
         }
 
-        return CreateSourceSession(
-            MigrationLog,
-            job.SourceContactPoint,
-            job.SourcePort,
-            username,
-            password,
-            tokenRefreshManager,
-            maxConnectionsPerHost: ResolveMaxConnectionsPerHost(job.SourceMaxConnectionsPerHost, job.MaxConnectionsPerHost));
+        int maxConnectionsPerHost = ResolveMaxConnectionsPerHost(
+            job.SourceMaxConnectionsPerHost,
+            job.MaxConnectionsPerHost);
+        if (maxConnectionsPerHost == 0 && workerCount > 0)
+        {
+            maxConnectionsPerHost = Math.Clamp(
+                (workerCount + 31) / 32,
+                2,
+                8);
+        }
+
+        return (
+            new SourceSessionSettings(
+                job.SourceContactPoint,
+                job.SourcePort,
+                username,
+                maxConnectionsPerHost),
+            credential);
     }
 
     /// <summary>

CassandraMigrationProcessor/CassandraDriver/JobSessionFactory.cs

@@ -0,0 +1,40 @@
+using Cassandra;
+using CassandraMigrationProcessor.Models;
+using CassandraMigrationProcessor.Infrastructure;
+
+namespace CassandraMigrationProcessor.CassandraDriver;
+
+/// <summary>
+/// Creates worker-owned target sessions for a job while limiting simultaneous
+/// opens to prevent a connection storm during startup.
+/// </summary>
+internal sealed class JobSessionFactory
+{
+    private const int MaxConcurrentSessionCreations = 20;
+
+    private readonly MigrationLog _log;
+    private readonly Job _job;
+    private readonly SemaphoreSlim _creationGate = new(
+        MaxConcurrentSessionCreations,
+        MaxConcurrentSessionCreations);
+
+    public JobSessionFactory(MigrationLog log, Job job)
+    {
+        _log = log ?? throw new ArgumentNullException(nameof(log));
+        _job = job ?? throw new ArgumentNullException(nameof(job));
+    }
+
+    public async Task<ISession> CreateSessionAsync(CancellationToken cancellationToken)
+    {
+        await _creationGate.WaitAsync(cancellationToken).ConfigureAwait(false);
+        try
+        {
+            return await CassandraClientFactory.CreateTargetSessionAsync(
+                _log, _job).ConfigureAwait(false);
+        }
+        finally
+        {
+            _creationGate.Release();
+        }
+    }
+}