Share source session across migration workers

CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs

@@ -5,24 +5,42 @@
 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();
+}
+
+/// <summary>
+/// Limits simultaneous session opens. This prevents high-worker jobs from
+/// creating a connection storm during startup.
+/// </summary>
+public sealed class GatedSessionFactory : ISessionFactory
+{
+    private readonly ISessionFactory _inner;
+    private readonly SemaphoreSlim _creationGate = new(2, 2);
 
-    /// <summary>Mint a new keyspace-agnostic target-cluster session. Async because
-    /// target credential discovery may go through ARM.</summary>
-    Task<ISession> CreateTargetSessionAsync();
+    public GatedSessionFactory(ISessionFactory inner)
+    {
+        _inner = inner ?? throw new ArgumentNullException(nameof(inner));
+    }
+
+    public async Task<ISession> CreateSessionAsync()
+    {
+        await _creationGate.WaitAsync().ConfigureAwait(false);
+        try
+        {
+            return await _inner.CreateSessionAsync()
+                .ConfigureAwait(false);
+        }
+        finally
+        {
+            _creationGate.Release();
+        }
+    }
 }
 
 /// <summary>
@@ -34,18 +52,13 @@ 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()
+    public Task<ISession> CreateSessionAsync()
         => CassandraClientFactory.CreateTargetSessionAsync(_log, _job);
 }

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.SourceSession, 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

@@ -1,3 +1,4 @@
+using Cassandra;
 using CassandraMigrationProcessor.CassandraDriver;
 using CassandraMigrationProcessor.Infrastructure;
 using CassandraMigrationProcessor.Models;
@@ -20,7 +21,10 @@ internal sealed class JobPipeline : IDisposable, IAsyncDisposable
     private readonly PartitionManager _partitions;
     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, ISession sourceSession,
+        ISessionFactory sessionFactory,
+        JobControl control)
     {
         _log = log;
         _pipelineConfig = pipelineConfig;
@@ -43,7 +47,8 @@ public JobPipeline(MigrationLog log, Job job, PipelineConfig pipelineConfig, Job
 
         Context = new PipelineContext(
             _partitions,
-            new JobSessionFactory(log, job, tokenRefreshManager),
+            sourceSession,
+            sessionFactory,
             readerConfig,
             writerConfig,
             EnableReplay: enableReplay,

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;
@@ -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,
+                _sourceSession,
+                new GatedSessionFactory(
+                    new JobSessionFactory(_log, job)),
+                _control);
             _pipeline.Start();
 
             await RunCopyPhaseAsync(job, units, partitioning, cancellationToken);

CassandraMigrationProcessor/DataTransfer/PageReader.cs

@@ -25,7 +25,7 @@ internal record ReaderConfig(int PageSize, int MaxReadRetries, bool PreserveCell
 /// cached per keyspace so the first partition for each table pays the
 /// cost and subsequent partitions reuse it.
 /// </summary>
-internal class PageReader : IDisposable
+internal class PageReader
 {
     private readonly WorkerLog _log;
     private readonly CancellationToken _ct;
@@ -50,26 +50,24 @@ internal class PageReader : IDisposable
     // hints parking a worker for minutes.
     private const int MaxRetryDelayMs = 30_000;
 
-    private PageReader(WorkerLog log, ISessionFactory sessionFactory, ReaderConfig config, CancellationToken cancellationToken)
+    private PageReader(WorkerLog log, ISession sourceSession, ReaderConfig config, CancellationToken cancellationToken)
     {
         _log = log;
         _ct = cancellationToken;
         _pageSize = config.PageSize;
         _maxReadRetries = config.MaxReadRetries;
         _preserveCellTtl = config.PreserveCellTtlAndWritetime;
         _useJsonCopy = config.UseJsonCopy;
-        _sourceSession = sessionFactory.CreateSourceSession();
+        _sourceSession = sourceSession;
     }
 
     public static Task<PageReader> CreateAsync(WorkerLog log,
-        ISessionFactory sessionFactory, ReaderConfig config,
+        ISession sourceSession, ReaderConfig config,
         CancellationToken cancellationToken)
     {
-        return Task.FromResult(new PageReader(log, sessionFactory, config, cancellationToken));
+        return Task.FromResult(new PageReader(log, sourceSession, config, cancellationToken));
     }
 
-    public void Dispose() => MigrationUtilities.SafeDisposeSession(_sourceSession, "PageReader source session");
-
     /// <summary>
     /// Lazy, idempotent UDT registration for typed reads. The first typed
     /// table registers every UDT in the keyspace because this reader can

CassandraMigrationProcessor/DataTransfer/PageWriter.cs

@@ -60,11 +60,12 @@ private PageWriter(WorkerLog log, ISession targetSession,
 
     public static async Task<PageWriter> CreateAsync(WorkerLog log, ISessionFactory sessionFactory, WriterConfig config, CancellationToken cancellationToken)
     {
-        var targetSession = await sessionFactory.CreateTargetSessionAsync();
+        var targetSession = await sessionFactory.CreateSessionAsync();
         return new PageWriter(log, targetSession, config, cancellationToken);
     }
 
-    public void Dispose() => MigrationUtilities.SafeDisposeSession(_targetSession, "PageWriter target session");
+    public void Dispose()
+        => MigrationUtilities.SafeDisposeSession(_targetSession, "PageWriter target session");
 
     private Task<IRowWriteStrategy> GetStrategyAsync(Partition partition)
     {

CassandraMigrationProcessor/DataTransfer/PipelineContext.cs

@@ -1,3 +1,4 @@
+using Cassandra;
 using CassandraMigrationProcessor.CassandraDriver;
 
 namespace CassandraMigrationProcessor.DataTransfer;
@@ -6,14 +7,15 @@ namespace CassandraMigrationProcessor.DataTransfer;
 /// Shared (job-wide) state passed to every worker. Holds the
 /// <see cref="DataTransfer.PartitionManager"/> that all tables seed into and
 /// every worker pulls from (and which owns the cooldown scheduler for
-/// delayed recycles), the connection capability used by readers and
-/// writers, reader / writer tunables, the replay configuration knobs, and
+/// delayed recycles), the shared source session, worker session factory,
+/// reader / writer tunables, the replay configuration knobs, and
 /// the unified <see cref="JobControl"/> (cancellation + first-fault).
 /// Per-table state is resolved through <see cref="Partition"/>
 /// pass-through accessors.
 /// </summary>
 internal record PipelineContext(
     PartitionManager Partitions,
+    ISession SourceSession,
     ISessionFactory SessionFactory,
     ReaderConfig ReaderConfig,
     WriterConfig WriterConfig,

CassandraMigrationProcessor/Models/TableCopySpec.cs

@@ -3,8 +3,8 @@ namespace CassandraMigrationProcessor.Models;
 /// <summary>
 /// Immutable description of a single table copy. Identifies the source
 /// and target keyspace/table; runtime sessions are not threaded through
-/// here — readers and writers open sessions via the job-wide
-/// <c>ISessionFactory</c>.
+/// here — readers use the job-wide source session and writers open
+/// worker-owned sessions through <c>ISessionFactory</c>.
 /// </summary>
 public record TableCopySpec(
     string KeyspaceName,

README.md

@@ -6,7 +6,7 @@ A .NET 9 Blazor Server web app that migrates data from **Azure Cosmos DB for Apa
 
 - **Schema auto-sync** — discovers source keyspaces, tables, columns, clustering keys, and static columns; generates DDL on the target automatically.
 - **Feed-range partitioned bulk copy** — splits each table into token-range chunks and copies them with configurable parallelism.
-- **Per-worker Cassandra sessions** — each parallel worker maintains its own driver session for maximum throughput.
+- **Efficient Cassandra sessions** — workers share the thread-safe source session while retaining independent target sessions for write throughput.
 - **Checkpoint-based pause / resume** — stop at any time and continue later with zero data loss; state is persisted to disk.
 - **Online mode (change feed)** — after the initial bulk copy, replays Cosmos DB change feed events to keep the target in sync until cutover.
 - **Wildcard table selection** — migrate all tables in a keyspace with `keyspace.*` or pick individual tables.