Share source session across migration workers

[Nitesh Vijay (niteshvijay1995)]

Summary

• use a focused SourceSessionWrapper for source session creation, rotation, deferred disposal, and session-scoped UDT registration
• resolve immutable source settings once and pass a concrete credential-aware session factory to the wrapper
• keep TokenRefreshManager focused on acquiring fresh tokens and triggering wrapper refresh
• resolve every source operation through the current wrapper session so AAD rotation is honored throughout the job
• remove faulted UDT registration entries and surface registration failures distinctly from read throttling
• prune UDT mappings when retired sessions are disposed and prevent token-refresh timer resurrection after shutdown
• retain per-worker target sessions, allow up to 20 concurrent session opens, and keep the gate alive for in-flight worker startup
• derive shared-source connection-pool sizing from worker count when the operator has not configured it

Validation

• dotnet build CassandraMigration.sln --nologo --verbosity minimal
• 40,000,000-row, 100-table sustained migration completed with exact target parity and zero failures using the earlier two-open gate
• exact v3.4 baseline: 7,439.4s at 5.376k rows/s
• shared-source design: 6,730.7s at 5.943k rows/s (10.5% faster)
• source connections reduced from approximately 130 to 2
• no metadata 429, BusyPool, NoHostAvailable, or fatal errors in the sustained run
• the current 20-open gate and subsequent lifecycle hardening are build-validated; they have not yet been rerun through the 40M workload

Fixes #57

[Copilot]

Pull request overview

Updates the migration worker session lifecycle to reduce Cosmos DB metadata request bursts by sharing a single source ISession across copy workers, while keeping per-worker target sessions to maintain write throughput.

Changes:

• Introduces SharedSourceSessionFactory to reuse the runner-wide source session across workers and throttle concurrent target-session creation.
• Tracks session ownership in PageReader/PageWriter to avoid disposing shared sessions.
• Wires the new session factory through MigrationJobRunner → JobPipeline → PipelineContext.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
README.md	Updates feature description to reflect shared source sessions + independent target sessions.
CassandraMigrationProcessor/DataTransfer/PageWriter.cs	Adds target-session ownership tracking to avoid disposing shared sessions.
CassandraMigrationProcessor/DataTransfer/PageReader.cs	Adds source-session ownership tracking to avoid disposing shared sessions.
CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs	Builds the pipeline with a shared-source session factory.
CassandraMigrationProcessor/DataTransfer/JobPipeline.cs	Accepts an ISessionFactory instead of constructing one internally.
CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs	Adds ownership flags and implements SharedSourceSessionFactory with a target-session creation gate.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Pull request overview

Copilot reviewed 10 out of 10 changed files in this pull request and generated no new comments.

Suppressed comments (3)

CassandraMigrationProcessor/DataTransfer/PageReader.cs:57

• PageReader now accepts a shared source session; add an explicit null-check so a miswired caller fails fast with a clear exception rather than an NRE later during reads.

    private PageReader(WorkerLog log, ISession sourceSession, ReaderConfig config, CancellationToken cancellationToken)
    {
        _log = log;
        _ct = cancellationToken;
        _pageSize = config.PageSize;

CassandraMigrationProcessor/DataTransfer/PageReader.cs:69

• Because the source session is now shared across all workers, per-worker UDT registration can cause redundant concurrent system_schema.types queries and repeated UDT registration work (especially for counter tables / non-JSON copy). Consider moving the UDT-registration cache to a job-wide location (e.g., PipelineContext) so keyspace UDT registration happens once per job+session instead of once per worker.

    public static Task<PageReader> CreateAsync(WorkerLog log,
        ISession sourceSession, ReaderConfig config,
        CancellationToken cancellationToken)
    {
        return Task.FromResult(new PageReader(log, sourceSession, config, cancellationToken));
    }

CassandraMigrationProcessor/CassandraDriver/ITargetSessionFactory.cs:40

• GatedTargetSessionFactory waits on the semaphore without observing cancellation. If a job is paused/stopped/faulted during startup, workers queued behind the gate may still open sessions before noticing cancellation, delaying shutdown and potentially creating unnecessary connections. Consider adding a CancellationToken parameter to ITargetSessionFactory.CreateTargetSessionAsync and passing it through to SemaphoreSlim.WaitAsync(ct) and the underlying CassandraClientFactory calls.

    public async Task<ISession> CreateTargetSessionAsync()
    {
        await _creationGate.WaitAsync().ConfigureAwait(false);
        try
        {
            return await _inner.CreateTargetSessionAsync()
                .ConfigureAwait(false);

[Copilot]

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated no new comments.

Suppressed comments (2)

CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs:32

• GatedTargetSessionFactory owns a SemaphoreSlim, but the factory isn’t disposable and JobPipeline.DisposeAsync() doesn’t attempt to dispose the target-session factory. While SemaphoreSlim only allocates a wait handle lazily, disposing it is still the standard way to avoid potential resource leaks if AvailableWaitHandle is ever used in the future.

Consider implementing IDisposable on GatedTargetSessionFactory (and any other factory that holds disposable resources) and disposing it when the pipeline is torn down.

public sealed class GatedTargetSessionFactory : ISessionFactory
{
    private readonly ISessionFactory _inner;
    private readonly SemaphoreSlim _creationGate = new(2, 2);

    public GatedTargetSessionFactory(ISessionFactory inner)
    {
        _inner = inner ?? throw new ArgumentNullException(nameof(inner));
    }

CassandraMigrationProcessor/DataTransfer/PageReader.cs:62

• Now that PageReader takes a shared job-wide source session, UDT registration work is still cached per reader instance (each worker has its own ConcurrentDictionary). That means each worker that touches typed reads can independently query system_schema.types and call UserDefinedTypes.Define(...) for the same keyspace, which partially undercuts the goal of reducing metadata/schema pressure.

Consider moving the UDT-registration cache to a job-wide scope (e.g., PipelineContext) keyed by keyspace (and bound to the specific source session) so each keyspace’s UDT registration runs once per job, regardless of worker count.

    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 = sourceSession;
    }

[Copilot]

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated no new comments.

Suppressed comments (3)

CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs:15

• The interface docs imply the factory can create sessions for multiple roles, but the only implementation in this PR (JobSessionFactory) always creates target sessions and the only consumer (PageWriter) uses it for target writes. Updating the XML docs to explicitly describe this as a target-session factory will reduce future confusion/misuse.

/// <summary>
/// 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 session.</summary>
    Task<ISession> CreateSessionAsync();
}

CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs:24

• The gate size 2 is a hard-coded magic number. Defining it as a named constant makes the intent clearer and avoids having to hunt for the value if it ever needs to become configurable.

    private const int MaxConcurrentSessionCreations = 20;

CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs:37

• WaitAsync() is called without a CancellationToken, so worker shutdown/cancellation can be delayed while waiting for the gate (and then while opening the session). If you want session creation to respond promptly to job cancellation, consider threading the job CancellationToken through ISessionFactory (e.g., CreateSessionAsync(CancellationToken)) and using it in WaitAsync(ct) (and any downstream session-open work).

    {
        _inner = inner ?? throw new ArgumentNullException(nameof(inner));
    }

    public async Task<ISession> CreateSessionAsync()
    {
        await _creationGate.WaitAsync().ConfigureAwait(false);

[Copilot]

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated 1 comment.

Suppressed comments (2)

CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs:15

• ISessionFactory.CreateSessionAsync is named as a generic session creator, but the only provided implementation (JobSessionFactory) always returns a target session (CassandraClientFactory.CreateTargetSessionAsync). This mismatch makes it easy for future consumers to accidentally use it for source sessions or assume role-agnostic semantics. Consider renaming to something explicit (e.g., ITargetSessionFactory/CreateTargetSessionAsync) or updating the documentation to clearly state it always creates target sessions.

public interface ISessionFactory
{
    /// <summary>Mint a new keyspace-agnostic session.</summary>
    Task<ISession> CreateSessionAsync();
}

CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs:39

• GatedSessionFactory.CreateSessionAsync waits on the semaphore without any cancellation mechanism. During shutdown/cancellation, worker startup can remain blocked on WaitAsync even though callers (e.g., PageWriter.CreateAsync) already have a CancellationToken available. Consider extending the factory API to accept a CancellationToken and using WaitAsync(token) so session creation can be cancelled promptly.

    public async Task<ISession> CreateSessionAsync()
    {
        await _creationGate.WaitAsync().ConfigureAwait(false);
        try
        {

[Copilot]

Pull request overview

Copilot reviewed 10 out of 10 changed files in this pull request and generated no new comments.

Suppressed comments (2)

CassandraMigrationProcessor/DataTransfer/PageReader.cs:28

• The PageReader summary comment says UDT registration is cached per keyspace, but the implementation now caches registrations per (session, keyspace) (see _udtRegistrations key). Updating the comment would avoid misleading readers about behavior during source-session rotation.

/// cached per keyspace so the first partition for each table pays the
/// cost and subsequent partitions reuse it.
/// </summary>
internal class PageReader

CassandraMigrationProcessor/DataTransfer/PageReader.cs:86

• _udtRegistrations is keyed by (ISession, keyspace). When TokenRefreshManager rotates the managed session, each worker will accumulate dictionary entries for every prior session, retaining references to disposed sessions and growing unbounded over long online runs. Consider pruning entries for non-current sessions when acquiring the current session.

        using var lease = _sourceSessionProvider.AcquireSession();
        var sourceSession = lease.Session;
        var keyspace = partition.Table.Spec.KeyspaceName;
        await _udtRegistrations.GetOrAdd((sourceSession, keyspace), async key =>
        {

[Copilot]

Pull request overview

Copilot reviewed 12 out of 12 changed files in this pull request and generated no new comments.

Suppressed comments (4)

CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs:173

• TokenRefreshCallback performs potentially long-running work (token acquisition + CassandraClientFactory.CreateSourceSession, which includes retry delays) while holding _refreshLock. GetSession() also takes _refreshLock, and PageReader now calls GetSession() on every retry attempt, so a slow refresh can block all reads (and may amplify latency / trigger timeouts). Consider reducing the lock scope: compute the new session outside the lock, then swap it in under the lock (and dispose/retire the old one after).

            try
            {
                string freshToken = GetFreshAadToken();

                // If we have a managed session, recreate it
                if (_sourceSessions.TryGetSession(out var currentSession)
                    && !currentSession!.IsDisposed
                    && _lastSourceContactPoint != null)
                {
                    var newSession = CassandraClientFactory.CreateSourceSession(
                        _log,
                        _lastSourceContactPoint,
                        _lastSourcePort,
                        _lastSourceUsername ?? string.Empty,
                        freshToken);
                    _sourceSessions.SetSession(newSession);
                }

                // Schedule next refresh
                _consecutiveRefreshFailures = 0;

CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs:246

• DisposeRetiredSessionAfterDelayAsync schedules an untracked Task.Delay(10 min) per rotation and the task captures the TokenRefreshManager instance. After job shutdown, Dispose() disposes sessions immediately, but these delayed tasks still keep the manager rooted until they complete, which can accumulate in a long-running service with frequent jobs/refreshes. Consider using a CancellationTokenSource canceled in Dispose() (and pass it to Task.Delay) so pending delayed-disposal tasks can exit promptly on shutdown.
  CassandraMigrationProcessor/DataTransfer/PageReader.cs:37
• _udtRegistrations uses the default equality for (ISession, string) keys. Since ISession equality is not guaranteed to be reference-based, different rotated sessions could be treated as the same key, causing UDT registration to be skipped for a new session (or, conversely, duplicated unexpectedly). Using explicit reference equality for the ISession part makes the cache semantics match the intent of “per session + keyspace” registration.

    private readonly ConcurrentDictionary<(ISession Session, string Keyspace), Task> _udtRegistrations = new();

CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs:181

• JobPipeline now gives workers an ISessionProvider that can rotate/dispose the source session, but the runner continues to use the captured _sourceSession for long-running phases (notably partition discovery). If an AAD refresh rotates the session mid-phase, _sourceSession can become a “retired” session and be disposed after the grace period, causing runner-side schema/partition queries to start failing with disposed/auth errors. Consider using the same ISessionProvider for runner operations as well (resolve the session per operation or per retry), or otherwise ensure the runner never uses a session instance that can be retired while the run is active.

                $"Migrating {units.Count} tables with {_pipelineConfig.WorkerCount} shared workers");

            await RunSchemaPhaseAsync(job, units, cancellationToken);

            var partitioning = await RunPartitioningPhaseAsync(

[Copilot]

Pull request overview

Copilot reviewed 13 out of 13 changed files in this pull request and generated no new comments.

Suppressed comments (3)

CassandraMigrationProcessor/CassandraDriver/JobSessionFactory.cs:40

• After making JobSessionFactory disposable, add a Dispose implementation to dispose the _creationGate semaphore so the owning pipeline can release resources deterministically.

        finally
        {
            _creationGate.Release();
        }
    }
}

CassandraMigrationProcessor/CassandraDriver/JobSessionFactory.cs:11

• JobSessionFactory holds a SemaphoreSlim but does not implement IDisposable. If multiple jobs/pipelines are created over the lifetime of the service, this makes it easy to leak resources (especially if the semaphore’s wait handle is ever materialized). Consider making the factory disposable so the owning pipeline can tear it down cleanly.

This issue also appears on line 35 of the same file.

internal sealed class JobSessionFactory

CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs:102

• The UDT-registration failure path removes the cached registration via the ICollection<KeyValuePair<...>>.Remove(...) API, which only removes if both key and value match. Using ConcurrentDictionary.TryRemove(key, out _) is simpler and guarantees the faulted entry is cleared even if another thread has already replaced the value for the same key.

        catch (Exception ex)
        {
            ((ICollection<KeyValuePair<(ISession Session, string Keyspace), Lazy<Task>>>)
                _udtRegistrations).Remove(new KeyValuePair<
                    (ISession Session, string Keyspace), Lazy<Task>>(
                    key, registration));
            throw new SourceUdtRegistrationException(keyspace, ex);
        }

[Copilot]

Pull request overview

Copilot reviewed 13 out of 13 changed files in this pull request and generated no new comments.

Suppressed comments (3)

CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs:46

• Token refresh is started based solely on source.Credential.Length > 200. This can incorrectly treat a long static password as an AAD token and start the refresh timer, which will then attempt AcquireAadToken() and rotate credentials unexpectedly. Since ResolveSourceSession already sets job.SourceUseAad = true when AAD is in use, key the refresh behavior off that flag instead of credential length.

            _sessions.Current = CreateSession(source.Credential);
            if (source.Credential.Length > 200)
                StartTokenRefresh(source.Credential);

CassandraMigrationProcessor/CassandraDriver/JobSessionFactory.cs:20

• JobSessionFactory owns a SemaphoreSlim (_creationGate) but never disposes it. Since this factory is created per job/pipeline, repeated runs can accumulate these resources in a long-running process. Implement IDisposable and dispose the gate when the pipeline/job shuts down.

internal sealed class JobSessionFactory
{
    private const int MaxConcurrentSessionCreations = 20;

    private readonly MigrationLog _log;
    private readonly Job _job;
    private readonly SemaphoreSlim _creationGate = new(
        MaxConcurrentSessionCreations,
        MaxConcurrentSessionCreations);

CassandraMigrationProcessor/DataTransfer/JobPipeline.cs:50

• A new per-job JobSessionFactory is created here, but JobPipeline never disposes it (and it currently holds a SemaphoreSlim). If JobSessionFactory becomes IDisposable (recommended), JobPipeline should retain a reference and dispose it during shutdown after workers have exited/cancelled; otherwise long-running services running multiple jobs can leak these factories.

        Context = new PipelineContext(
            _partitions,
            sourceSession,
            new JobSessionFactory(log, job),
            readerConfig,

[Copilot]

Pull request overview

Copilot reviewed 13 out of 13 changed files in this pull request and generated no new comments.

Suppressed comments (3)

CassandraMigrationProcessor/CassandraDriver/JobSessionFactory.cs:31

• JobSessionFactory should expose Dispose() to dispose the internal SemaphoreSlim gate. Without this, the gate may keep OS resources alive beyond the job lifetime (especially if its WaitHandle is ever created).

    public async Task<ISession> CreateSessionAsync(CancellationToken cancellationToken)
    {
        await _creationGate.WaitAsync(cancellationToken).ConfigureAwait(false);
        try
        {

CassandraMigrationProcessor/CassandraDriver/JobSessionFactory.cs:11

• JobSessionFactory owns a SemaphoreSlim (_creationGate). SemaphoreSlim is IDisposable and should be disposed when the job/pipeline ends to avoid leaking wait handles/resources in long-running hosts. Implement IDisposable on this type so the owning pipeline can dispose it deterministically.

This issue also appears on line 27 of the same file.

internal sealed class JobSessionFactory

CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs:82

• The failed UDT-registration entry removal is more complex than necessary and relies on ICollection<KeyValuePair<...>>.Remove with a constructed KeyValuePair. ConcurrentDictionary already provides TryRemove by key, which is clearer and avoids any subtle KVP equality/value-instance coupling.

            ((ICollection<KeyValuePair<(ISession Session, string Keyspace), Lazy<Task>>>)
                _udtRegistrations).Remove(new KeyValuePair<
                    (ISession Session, string Keyspace), Lazy<Task>>(
                    key, registration));

[Copilot]

Pull request overview

Copilot reviewed 13 out of 13 changed files in this pull request and generated no new comments.

Suppressed comments (2)

CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs:53

• Token refresh is scheduled based on credential length (>200), which can misclassify long static passwords as AAD/JWT tokens and cause the refresh timer to rotate the source session using managed identity unexpectedly. Since ResolveSourceSession already sets job.SourceUseAad when AAD is in use, schedule refresh based on that flag instead of length.

        _currentSession = CreateSession(source.Credential);
        if (source.Credential.Length > 200)
            ScheduleTokenRefresh(source.Credential);

CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs:83

• The failure-path cleanup for _udtRegistrations uses an explicit ICollection<KeyValuePair<...>> cast plus Remove(KeyValuePair) which is hard to read and easy to get wrong. You can keep the same “remove only if still mapped to this Lazy” semantics with a TryGetValue + ReferenceEquals check and then TryRemove.

            ((ICollection<KeyValuePair<(ISession Session, string Keyspace), Lazy<Task>>>)
                _udtRegistrations).Remove(new KeyValuePair<
                    (ISession Session, string Keyspace), Lazy<Task>>(
                    key, registration));
            throw new SourceUdtRegistrationException(keyspace, ex);

[Copilot]

Pull request overview

Copilot reviewed 13 out of 13 changed files in this pull request and generated no new comments.

Suppressed comments (2)

CassandraMigrationProcessor/DataTransfer/JobPipeline.cs:50

• JobPipeline creates a per-job JobSessionFactory, but DisposeAsync() only disposes _partitions and _workerPool. Since JobSessionFactory owns the session-open concurrency gate, it should be disposed with the pipeline to ensure any underlying resources are released and to make shutdown behavior explicit.

        Context = new PipelineContext(
            _partitions,
            sourceSession,
            new JobSessionFactory(log, job),
            readerConfig,

CassandraMigrationProcessor/CassandraDriver/JobSessionFactory.cs:19

• JobSessionFactory owns a SemaphoreSlim gate (_creationGate) but the type is not disposable, so the gate can’t be deterministically cleaned up when a job finishes. Consider implementing IDisposable (or IAsyncDisposable) on JobSessionFactory and disposing _creationGate, then disposing the factory from JobPipeline.

internal sealed class JobSessionFactory
{
    private const int MaxConcurrentSessionCreations = 20;

    private readonly MigrationLog _log;
    private readonly Job _job;
    private readonly SemaphoreSlim _creationGate = new(
        MaxConcurrentSessionCreations,
        MaxConcurrentSessionCreations);

[Copilot]

Pull request overview

Copilot reviewed 13 out of 13 changed files in this pull request and generated no new comments.

Suppressed comments (3)

CassandraMigrationProcessor/DataTransfer/JobPipeline.cs:51

• JobSessionFactory owns a SemaphoreSlim (_creationGate) which implements IDisposable, but JobPipeline constructs the factory and never disposes it. In a long-running web app that runs many jobs, this can leak wait handles over time. Consider making JobSessionFactory disposable and disposing it from JobPipeline.DisposeAsync (or holding the gate in a shared, app-wide singleton instead of per-job).

        Context = new PipelineContext(
            _partitions,
            sourceSession,
            new JobSessionFactory(log, job),
            readerConfig,
            writerConfig,

CassandraMigrationProcessor/CassandraDriver/JobSessionFactory.cs:20

• JobSessionFactory creates a SemaphoreSlim (_creationGate) but never disposes it. Since a JobSessionFactory is created per job, this can leak wait handles across repeated migrations in a long-running process. Implement IDisposable and dispose _creationGate, and ensure the owning component (JobPipeline) calls Dispose/DisposeAsync during shutdown.

internal sealed class JobSessionFactory
{
    private const int MaxConcurrentSessionCreations = 20;

    private readonly MigrationLog _log;
    private readonly Job _job;
    private readonly SemaphoreSlim _creationGate = new(
        MaxConcurrentSessionCreations,
        MaxConcurrentSessionCreations);

CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs:82

• The failure-path cleanup for a faulted UDT registration uses an ICollection<KeyValuePair<...>> cast and Remove with a constructed KeyValuePair. This is harder to read and easier to get subtly wrong than using the native ConcurrentDictionary API. Prefer TryRemove(key, out _) to clearly express the intent to remove the cached registration so the next attempt can retry.

        if (job.SourceUseAad)
            ScheduleTokenRefresh(credential);
    }

    public ISession GetSession()

[Copilot]

Pull request overview

Copilot reviewed 13 out of 13 changed files in this pull request and generated 1 comment.

Suppressed comments (1)

CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs:110

• The catch-path removal of a failed UDT registration uses an unusual ICollection<KeyValuePair<...>>.Remove(...) pattern. Using TryRemove(key, out _) is simpler and avoids depending on KeyValuePair/value equality semantics, while still ensuring subsequent calls can retry registration after a failure.

            ((ICollection<KeyValuePair<(ISession Session, string Keyspace), Lazy<Task>>>)
                _udtRegistrations).Remove(new KeyValuePair<
                    (ISession Session, string Keyspace), Lazy<Task>>(
                    key, registration));
            throw new SourceUdtRegistrationException(keyspace, ex);