Share source session across migration workers

CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs

@@ -17,6 +17,13 @@ namespace CassandraMigrationProcessor.CassandraDriver;
 /// </summary>
 public interface ISessionFactory
 {
+    /// <summary>
+    /// Whether sessions returned by this factory are owned by the caller.
+    /// Shared job sessions are owned by the migration runner instead.
+    /// </summary>
+    bool CallerOwnsSourceSession => true;
+    bool CallerOwnsTargetSession => true;
+
     /// <summary>Mint a new keyspace-agnostic source-cluster session.</summary>
     ISession CreateSourceSession();
 
@@ -25,6 +32,44 @@ public interface ISessionFactory
     Task<ISession> CreateTargetSessionAsync();
 }
 
+/// <summary>
+/// Exposes the runner's job-wide source session while retaining per-worker
+/// target sessions. Sharing the source avoids the Cosmos metadata connection
+/// storm; independent target sessions preserve the writer capacity required
+/// by high-concurrency bulk jobs.
+/// </summary>
+public sealed class SharedSourceSessionFactory : ISessionFactory
+{
+    private readonly ISession _sourceSession;
+    private readonly ISessionFactory _targetSessionFactory;
+    private readonly SemaphoreSlim _targetSessionCreationGate = new(2, 2);
+
+    public SharedSourceSessionFactory(ISession sourceSession, ISessionFactory targetSessionFactory)
+    {
+        _sourceSession = sourceSession ?? throw new ArgumentNullException(nameof(sourceSession));
+        _targetSessionFactory = targetSessionFactory
+            ?? throw new ArgumentNullException(nameof(targetSessionFactory));
+    }
+
+    public bool CallerOwnsSourceSession => false;
+
+    public ISession CreateSourceSession() => _sourceSession;
+
+    public async Task<ISession> CreateTargetSessionAsync()
+    {
+        await _targetSessionCreationGate.WaitAsync().ConfigureAwait(false);
+        try
+        {
+            return await _targetSessionFactory.CreateTargetSessionAsync()
+                .ConfigureAwait(false);
+        }
+        finally
+        {
+            _targetSessionCreationGate.Release();
+        }
+    }
+}
+
 /// <summary>
 /// Default <see cref="ISessionFactory"/> bound to a single
 /// <see cref="Job"/>. Delegates to <see cref="CassandraClientFactory"/>

CassandraMigrationProcessor/DataTransfer/JobPipeline.cs

@@ -20,7 +20,9 @@ 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, ISessionFactory sessionFactory,
+        JobControl control)
     {
         _log = log;
         _pipelineConfig = pipelineConfig;
@@ -43,7 +45,7 @@ public JobPipeline(MigrationLog log, Job job, PipelineConfig pipelineConfig, Job
 
         Context = new PipelineContext(
             _partitions,
-            new JobSessionFactory(log, job, tokenRefreshManager),
+            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,
+                new SharedSourceSessionFactory(
+                    _sourceSession,
+                    new JobSessionFactory(_log, job, _tokenRefreshManager)),
+                _control);
             _pipeline.Start();
 
             await RunCopyPhaseAsync(job, units, partitioning, cancellationToken);

CassandraMigrationProcessor/DataTransfer/PageReader.cs

@@ -30,6 +30,7 @@ internal class PageReader : IDisposable
     private readonly WorkerLog _log;
     private readonly CancellationToken _ct;
     private readonly ISession _sourceSession;
+    private readonly bool _ownsSourceSession;
     private readonly int _pageSize;
     private readonly int _maxReadRetries;
     private readonly bool _preserveCellTtl;
@@ -59,6 +60,7 @@ private PageReader(WorkerLog log, ISessionFactory sessionFactory, ReaderConfig c
         _preserveCellTtl = config.PreserveCellTtlAndWritetime;
         _useJsonCopy = config.UseJsonCopy;
         _sourceSession = sessionFactory.CreateSourceSession();
+        _ownsSourceSession = sessionFactory.CallerOwnsSourceSession;
     }
 
     public static Task<PageReader> CreateAsync(WorkerLog log,
@@ -68,7 +70,11 @@ public static Task<PageReader> CreateAsync(WorkerLog log,
         return Task.FromResult(new PageReader(log, sessionFactory, config, cancellationToken));
     }
 
-    public void Dispose() => MigrationUtilities.SafeDisposeSession(_sourceSession, "PageReader source session");
+    public void Dispose()
+    {
+        if (_ownsSourceSession)
+            MigrationUtilities.SafeDisposeSession(_sourceSession, "PageReader source session");
+    }
 
     /// <summary>
     /// Lazy, idempotent UDT registration for typed reads. The first typed

CassandraMigrationProcessor/DataTransfer/PageWriter.cs

@@ -29,6 +29,7 @@ internal sealed class PageWriter : IDisposable
     private readonly WorkerLog _log;
     private readonly CancellationToken _ct;
     private readonly ISession _targetSession;
+    private readonly bool _ownsTargetSession;
     private readonly int _maxWriteRetries;
     private readonly ConsistencyLevel _targetWriteConsistencyLevel;
     private readonly bool _preserveCellTtl;
@@ -46,7 +47,7 @@ internal sealed class PageWriter : IDisposable
     /// </summary>
     internal Exception? LastWriteException { get; private set; }
 
-    private PageWriter(WorkerLog log, ISession targetSession,
+    private PageWriter(WorkerLog log, ISession targetSession, bool ownsTargetSession,
         WriterConfig config, CancellationToken cancellationToken)
     {
         _log = log;
@@ -56,15 +57,20 @@ private PageWriter(WorkerLog log, ISession targetSession,
         _preserveCellTtl = config.PreserveCellTtlAndWritetime;
         _useJsonCopy = config.UseJsonCopy;
         _targetSession = targetSession;
+        _ownsTargetSession = ownsTargetSession;
     }
 
     public static async Task<PageWriter> CreateAsync(WorkerLog log, ISessionFactory sessionFactory, WriterConfig config, CancellationToken cancellationToken)
     {
         var targetSession = await sessionFactory.CreateTargetSessionAsync();
-        return new PageWriter(log, targetSession, config, cancellationToken);
+        return new PageWriter(log, targetSession, sessionFactory.CallerOwnsTargetSession, config, cancellationToken);
     }
 
-    public void Dispose() => MigrationUtilities.SafeDisposeSession(_targetSession, "PageWriter target session");
+    public void Dispose()
+    {
+        if (_ownsTargetSession)
+            MigrationUtilities.SafeDisposeSession(_targetSession, "PageWriter target session");
+    }
 
     private Task<IRowWriteStrategy> GetStrategyAsync(Partition partition)
     {

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.