Share source session across migration workers

CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs

@@ -5,24 +5,73 @@
 namespace CassandraMigrationProcessor.CassandraDriver;
 
 /// <summary>
-/// Per-job session factory. Encapsulates everything required to mint a
-/// new source or target <see cref="ISession"/> (job credentials, logger,
-/// optional token refresh manager) so that consumers — primarily
-/// <see cref="DataTransfer.PageReader"/> and
-/// <see cref="DataTransfer.PageWriter"/> — depend on a single
-/// abstraction instead of being threaded the raw <see cref="Job"/> and
-/// <see cref="TokenRefreshManager"/> separately. This keeps worker-side
-/// classes focused on data movement and makes it possible to swap the
-/// connection wiring in tests without touching their constructors.
+/// Creates worker-owned sessions. The consumer determines the session role;
+/// job-owned shared sessions are passed directly instead of using this factory.
 /// </summary>
 public interface ISessionFactory
 {
-    /// <summary>Mint a new keyspace-agnostic source-cluster session.</summary>
-    ISession CreateSourceSession();
+    /// <summary>Mint a new keyspace-agnostic session.</summary>
+    Task<ISession> CreateSessionAsync(CancellationToken cancellationToken);
+}
+
+/// <summary>
+/// Provides a lease on the current shared session. A rotated session is not
+/// disposed until all operations using its leases have completed.
+/// </summary>
+public interface ISessionProvider
+{
+    SessionLease AcquireSession();
+}
+
+public sealed class SessionLease : IDisposable
+{
+    private Action? _release;
+
+    internal SessionLease(ISession session, Action release)
+    {
+        Session = session;
+        _release = release;
+    }
+
+    public ISession Session { get; }
+
+    public void Dispose()
+        => Interlocked.Exchange(ref _release, null)?.Invoke();
+}
 
-    /// <summary>Mint a new keyspace-agnostic target-cluster session. Async because
-    /// target credential discovery may go through ARM.</summary>
-    Task<ISession> CreateTargetSessionAsync();
+/// <summary>
+/// Limits simultaneous session opens. This prevents high-worker jobs from
+/// creating a connection storm during startup.
+/// </summary>
+public sealed class GatedSessionFactory : ISessionFactory, IDisposable
+{
+    private const int MaxConcurrentSessionCreations = 20;
+
+    private readonly ISessionFactory _inner;
+    private readonly SemaphoreSlim _creationGate = new(
+        MaxConcurrentSessionCreations,
+        MaxConcurrentSessionCreations);
+
+    public GatedSessionFactory(ISessionFactory inner)
+    {
+        _inner = inner ?? throw new ArgumentNullException(nameof(inner));
+    }
+
+    public async Task<ISession> CreateSessionAsync(CancellationToken cancellationToken)
+    {
+        await _creationGate.WaitAsync(cancellationToken).ConfigureAwait(false);
+        try
+        {
+            return await _inner.CreateSessionAsync(cancellationToken)
+                .ConfigureAwait(false);
+        }
+        finally
+        {
+            _creationGate.Release();
+        }
+    }
+
+    public void Dispose() => _creationGate.Dispose();
 }
 
 /// <summary>
@@ -34,18 +83,17 @@ public sealed class JobSessionFactory : ISessionFactory
 {
     private readonly MigrationLog _log;
     private readonly Job _job;
-    private readonly TokenRefreshManager? _tokenRefreshManager;
 
-    public JobSessionFactory(MigrationLog log, Job job, TokenRefreshManager? tokenRefreshManager)
+    public JobSessionFactory(MigrationLog log, Job job)
     {
         _log = log;
         _job = job;
-        _tokenRefreshManager = tokenRefreshManager;
     }
 
-    public ISession CreateSourceSession()
-        => CassandraClientFactory.CreateSourceSession(_log, _job, _tokenRefreshManager);
-
-    public Task<ISession> CreateTargetSessionAsync()
-        => CassandraClientFactory.CreateTargetSessionAsync(_log, _job);
+    public async Task<ISession> CreateSessionAsync(CancellationToken cancellationToken)
+    {
+        cancellationToken.ThrowIfCancellationRequested();
+        return await CassandraClientFactory.CreateTargetSessionAsync(_log, _job)
+            .ConfigureAwait(false);
+    }
 }

CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs

@@ -8,12 +8,22 @@ namespace CassandraMigrationProcessor.CassandraDriver;
 /// Manages AAD token lifecycle and proactive refresh for
 /// Cosmos DB Cassandra API connections.
 /// </summary>
-public class TokenRefreshManager : IDisposable
+public class TokenRefreshManager : ISessionProvider, IDisposable
 {
+    private sealed class ManagedSessionState
+    {
+        public ManagedSessionState(ISession session) => Session = session;
+
+        public ISession Session { get; }
+        public int ActiveLeases { get; set; }
+        public bool Retired { get; set; }
+    }
+
     private Timer? _tokenRefreshTimer;
     private readonly object _refreshLock = new();
-    private ISession? _managedSourceSession;
+    private ManagedSessionState? _managedSourceSession;
     private readonly MigrationLog _log;
+    private bool _disposed;
     private DateTime _tokenExpiresAt = DateTime.MinValue;
     private int _consecutiveRefreshFailures;
     private const int MaxRefreshFailures = 6;
@@ -153,17 +163,16 @@ private void TokenRefreshCallback(object? state)
 
                 // If we have a managed session, recreate it
                 if (_managedSourceSession != null
-                    && !_managedSourceSession.IsDisposed
+                    && !_managedSourceSession.Session.IsDisposed
                     && _lastSourceContactPoint != null)
                 {
-                    var oldSession = _managedSourceSession;
-                    _managedSourceSession = CassandraClientFactory.CreateSourceSession(
+                    var newSession = CassandraClientFactory.CreateSourceSession(
                         _log,
                         _lastSourceContactPoint,
                         _lastSourcePort,
                         _lastSourceUsername ?? string.Empty,
                         freshToken);
-                    MigrationUtilities.SafeDisposeSession(oldSession, "TokenRefresh old session");
+                    SetManagedSourceSession(newSession);
                 }
 
                 // Schedule next refresh
@@ -199,14 +208,73 @@ private void TokenRefreshCallback(object? state)
     /// </summary>
     public void SetManagedSourceSession(ISession session)
     {
-        _managedSourceSession = session;
+        ArgumentNullException.ThrowIfNull(session);
+
+        ISession? sessionToDispose = null;
+        lock (_refreshLock)
+        {
+            ObjectDisposedException.ThrowIf(_disposed, this);
+            if (ReferenceEquals(_managedSourceSession?.Session, session))
+                return;
+
+            var previous = _managedSourceSession;
+            _managedSourceSession = new ManagedSessionState(session);
+            if (previous != null)
+            {
+                previous.Retired = true;
+                if (previous.ActiveLeases == 0)
+                    sessionToDispose = previous.Session;
+            }
+        }
+
+        MigrationUtilities.SafeDisposeSession(
+            sessionToDispose, "TokenRefresh retired source session");
+    }
+
+    public SessionLease AcquireSession()
+    {
+        lock (_refreshLock)
+        {
+            ObjectDisposedException.ThrowIf(_disposed, this);
+            var state = _managedSourceSession
+                ?? throw new InvalidOperationException("The source session has not been initialized.");
+            state.ActiveLeases++;
+            return new SessionLease(state.Session, () => ReleaseSession(state));
+        }
+    }
+
+    private void ReleaseSession(ManagedSessionState state)
+    {
+        ISession? sessionToDispose = null;
+        lock (_refreshLock)
+        {
+            state.ActiveLeases--;
+            if (state.Retired && state.ActiveLeases == 0)
+                sessionToDispose = state.Session;
+        }
+
+        MigrationUtilities.SafeDisposeSession(
+            sessionToDispose, "TokenRefresh retired source session");
     }
 
     public void Dispose()
     {
+        ISession? sessionToDispose = null;
         lock (_refreshLock)
         {
+            if (_disposed) return;
+            _disposed = true;
             StopTokenRefreshTimer();
+            if (_managedSourceSession != null)
+            {
+                _managedSourceSession.Retired = true;
+                if (_managedSourceSession.ActiveLeases == 0)
+                    sessionToDispose = _managedSourceSession.Session;
+                _managedSourceSession = null;
+            }
         }
+
+        MigrationUtilities.SafeDisposeSession(
+            sessionToDispose, "TokenRefresh managed source session");
     }
 }

CassandraMigrationProcessor/DataTransfer/DataCopyWorker.cs

@@ -33,7 +33,7 @@ public async Task RunAsync(PipelineContext ctx)
         Partition? current = null;
         try
         {
-            reader = await PageReader.CreateAsync(_workerLog, ctx.SessionFactory, ctx.ReaderConfig, _ct);
+            reader = await PageReader.CreateAsync(_workerLog, ctx.SourceSessionProvider, ctx.ReaderConfig, _ct);
             writer = await PageWriter.CreateAsync(_workerLog, ctx.SessionFactory, ctx.WriterConfig, _ct);
 
             while (!_ct.IsCancellationRequested
@@ -159,7 +159,6 @@ public async Task RunAsync(PipelineContext ctx)
             // instant any worker exits. Channel completion is driven
             // by the orchestrator (offline) or _cts.Cancel (fatal).
             MigrationUtilities.SafeDispose(writer, "worker PageWriter");
-            MigrationUtilities.SafeDispose(reader, "worker PageReader");
         }
     }
 

CassandraMigrationProcessor/DataTransfer/JobPipeline.cs

@@ -18,13 +18,18 @@ internal sealed class JobPipeline : IDisposable, IAsyncDisposable
     private readonly JobControl _control;
     private readonly WorkerPool _workerPool;
     private readonly PartitionManager _partitions;
+    private readonly ISessionFactory _sessionFactory;
     public PipelineContext Context { get; }
 
-    public JobPipeline(MigrationLog log, Job job, PipelineConfig pipelineConfig, JobPartitioning partitioning, TokenRefreshManager? tokenRefreshManager, JobControl control)
+    public JobPipeline(MigrationLog log, Job job, PipelineConfig pipelineConfig,
+        JobPartitioning partitioning, ISessionProvider sourceSessionProvider,
+        ISessionFactory sessionFactory,
+        JobControl control)
     {
         _log = log;
         _pipelineConfig = pipelineConfig;
         _control = control;
+        _sessionFactory = sessionFactory;
 
         bool enableReplay = job.IsOnline;
         _partitions = new PartitionManager(
@@ -43,7 +48,8 @@ public JobPipeline(MigrationLog log, Job job, PipelineConfig pipelineConfig, Job
 
         Context = new PipelineContext(
             _partitions,
-            new JobSessionFactory(log, job, tokenRefreshManager),
+            sourceSessionProvider,
+            sessionFactory,
             readerConfig,
             writerConfig,
             EnableReplay: enableReplay,
@@ -95,5 +101,7 @@ public async ValueTask DisposeAsync()
         // by JobManager — we never cancel or dispose it here.
         await _partitions.DisposeAsync().ConfigureAwait(false);
         MigrationUtilities.SafeDispose(_workerPool, "JobPipeline WorkerPool");
+        if (_sessionFactory is IDisposable disposableFactory)
+            MigrationUtilities.SafeDispose(disposableFactory, "JobPipeline SessionFactory");
     }
 }

CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs

@@ -34,8 +34,9 @@ public class MigrationJobRunner : IAsyncDisposable
     /// Runner-wide source / target sessions opened once in
     /// <see cref="CreateAsync"/> and reused across wildcard expansion,
     /// schema provisioning, and partition discovery. Disposed in
-    /// <see cref="DisposeAsync"/>. Copy workers mint their own sessions
-    /// via <see cref="ISessionFactory"/> for throughput isolation.
+    /// <see cref="DisposeAsync"/>. Copy workers reuse the thread-safe source
+    /// session to avoid multiplying driver metadata topology/schema handshakes,
+    /// while retaining independent target sessions for write throughput.
     /// For simulated runs the target session is a <see cref="NullSession"/>.
     /// </summary>
     private readonly ISession _sourceSession;
@@ -87,14 +88,14 @@ public static async Task<MigrationJobRunner> CreateAsync(
         try
         {
             source = CassandraClientFactory.CreateSourceSession(log, job, tokenRefreshManager);
+            tokenRefreshManager.SetManagedSourceSession(source);
             target = await CassandraClientFactory.CreateTargetSessionAsync(log, job);
             return new MigrationJobRunner(log, job, pipelineConfig, control, tokenRefreshManager, source, target);
         }
         catch
         {
             MigrationUtilities.SafeDisposeSession(target, "MigrationJobRunner target (CreateAsync rollback)");
-            MigrationUtilities.SafeDisposeSession(source, "MigrationJobRunner source (CreateAsync rollback)");
-            tokenRefreshManager.StopTokenRefreshTimer();
+            tokenRefreshManager.Dispose();
             throw;
         }
     }
@@ -173,7 +174,12 @@ public async Task StartAsync()
             var partitioning = await RunPartitioningPhaseAsync(
                 job, units, cancellationToken);
 
-            _pipeline = new JobPipeline(_log, job, _pipelineConfig, partitioning, _tokenRefreshManager, _control);
+            _pipeline = new JobPipeline(
+                _log, job, _pipelineConfig, partitioning,
+                _tokenRefreshManager,
+                new GatedSessionFactory(
+                    new JobSessionFactory(_log, job)),
+                _control);
             _pipeline.Start();
 
             await RunCopyPhaseAsync(job, units, partitioning, cancellationToken);
@@ -245,7 +251,7 @@ public ValueTask DisposeAsync()
         MigrationUtilities.SafeDispose(_pipeline, "JobPipeline (Dispose)");
         _pipeline = null;
         MigrationUtilities.SafeDisposeSession(_targetSession, "MigrationJobRunner target session");
-        MigrationUtilities.SafeDisposeSession(_sourceSession, "MigrationJobRunner source session");
+        _tokenRefreshManager.Dispose();
         return ValueTask.CompletedTask;
     }