From 709a1bc6dbf55ddebb38db4c9fd9b0e09fb6357f Mon Sep 17 00:00:00 2001 From: Nitesh Vijay Date: Tue, 18 Aug 2026 11:40:31 +0530 Subject: [PATCH 01/32] perf: share source session across workers Reuse the job-wide thread-safe source session while preserving per-worker target sessions. Gate target session creation to two concurrent opens to avoid connection storms during high-worker startup. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 078ac7b9-aa7f-44a1-9053-44c346cec5e1 --- .../CassandraDriver/ISessionFactory.cs | 45 +++++++++++++++++++ .../DataTransfer/JobPipeline.cs | 6 ++- .../DataTransfer/MigrationJobRunner.cs | 12 +++-- .../DataTransfer/PageReader.cs | 8 +++- .../DataTransfer/PageWriter.cs | 12 +++-- README.md | 2 +- 6 files changed, 75 insertions(+), 10 deletions(-) diff --git a/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs b/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs index 6068a87..6711bdf 100644 --- a/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs +++ b/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs @@ -17,6 +17,13 @@ namespace CassandraMigrationProcessor.CassandraDriver; /// public interface ISessionFactory { + /// + /// Whether sessions returned by this factory are owned by the caller. + /// Shared job sessions are owned by the migration runner instead. + /// + bool CallerOwnsSourceSession => true; + bool CallerOwnsTargetSession => true; + /// Mint a new keyspace-agnostic source-cluster session. ISession CreateSourceSession(); @@ -25,6 +32,44 @@ public interface ISessionFactory Task CreateTargetSessionAsync(); } +/// +/// 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. +/// +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 CreateTargetSessionAsync() + { + await _targetSessionCreationGate.WaitAsync().ConfigureAwait(false); + try + { + return await _targetSessionFactory.CreateTargetSessionAsync() + .ConfigureAwait(false); + } + finally + { + _targetSessionCreationGate.Release(); + } + } +} + /// /// Default bound to a single /// . Delegates to diff --git a/CassandraMigrationProcessor/DataTransfer/JobPipeline.cs b/CassandraMigrationProcessor/DataTransfer/JobPipeline.cs index dd48feb..590f909 100644 --- a/CassandraMigrationProcessor/DataTransfer/JobPipeline.cs +++ b/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, diff --git a/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs b/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs index ccd2778..b4076f8 100644 --- a/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs +++ b/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs @@ -34,8 +34,9 @@ public class MigrationJobRunner : IAsyncDisposable /// Runner-wide source / target sessions opened once in /// and reused across wildcard expansion, /// schema provisioning, and partition discovery. Disposed in - /// . Copy workers mint their own sessions - /// via for throughput isolation. + /// . 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 . /// 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); diff --git a/CassandraMigrationProcessor/DataTransfer/PageReader.cs b/CassandraMigrationProcessor/DataTransfer/PageReader.cs index 243e1f0..0d1ca40 100644 --- a/CassandraMigrationProcessor/DataTransfer/PageReader.cs +++ b/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 CreateAsync(WorkerLog log, @@ -68,7 +70,11 @@ public static Task 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"); + } /// /// Lazy, idempotent UDT registration for typed reads. The first typed diff --git a/CassandraMigrationProcessor/DataTransfer/PageWriter.cs b/CassandraMigrationProcessor/DataTransfer/PageWriter.cs index 86c2981..cb81d6a 100644 --- a/CassandraMigrationProcessor/DataTransfer/PageWriter.cs +++ b/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 /// 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 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 GetStrategyAsync(Partition partition) { diff --git a/README.md b/README.md index 291d8e5..fbc3924 100644 --- a/README.md +++ b/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. From f852240a837d1a32a72817c3c754e6c30cd0e83a Mon Sep 17 00:00:00 2001 From: Nitesh Vijay Date: Tue, 18 Aug 2026 14:55:08 +0530 Subject: [PATCH 02/32] refactor: make session ownership explicit Pass the runner-owned source session directly to readers and restrict the factory abstraction to worker-owned target sessions. This removes ownership booleans and makes disposal responsibilities structural. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 078ac7b9-aa7f-44a1-9053-44c346cec5e1 --- .../CassandraDriver/ISessionFactory.cs | 96 ------------------- .../CassandraDriver/ITargetSessionFactory.cs | 67 +++++++++++++ .../DataTransfer/DataCopyWorker.cs | 5 +- .../DataTransfer/JobPipeline.cs | 7 +- .../DataTransfer/MigrationJobRunner.cs | 6 +- .../DataTransfer/PageReader.cs | 18 +--- .../DataTransfer/PageWriter.cs | 13 +-- .../DataTransfer/PipelineContext.cs | 8 +- .../Models/TableCopySpec.cs | 4 +- 9 files changed, 93 insertions(+), 131 deletions(-) delete mode 100644 CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs create mode 100644 CassandraMigrationProcessor/CassandraDriver/ITargetSessionFactory.cs diff --git a/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs b/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs deleted file mode 100644 index 6711bdf..0000000 --- a/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs +++ /dev/null @@ -1,96 +0,0 @@ -using Cassandra; -using CassandraMigrationProcessor.Models; -using CassandraMigrationProcessor.Infrastructure; - -namespace CassandraMigrationProcessor.CassandraDriver; - -/// -/// Per-job session factory. Encapsulates everything required to mint a -/// new source or target (job credentials, logger, -/// optional token refresh manager) so that consumers — primarily -/// and -/// — depend on a single -/// abstraction instead of being threaded the raw and -/// 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. -/// -public interface ISessionFactory -{ - /// - /// Whether sessions returned by this factory are owned by the caller. - /// Shared job sessions are owned by the migration runner instead. - /// - bool CallerOwnsSourceSession => true; - bool CallerOwnsTargetSession => true; - - /// Mint a new keyspace-agnostic source-cluster session. - ISession CreateSourceSession(); - - /// Mint a new keyspace-agnostic target-cluster session. Async because - /// target credential discovery may go through ARM. - Task CreateTargetSessionAsync(); -} - -/// -/// 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. -/// -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 CreateTargetSessionAsync() - { - await _targetSessionCreationGate.WaitAsync().ConfigureAwait(false); - try - { - return await _targetSessionFactory.CreateTargetSessionAsync() - .ConfigureAwait(false); - } - finally - { - _targetSessionCreationGate.Release(); - } - } -} - -/// -/// Default bound to a single -/// . Delegates to -/// so the connection-construction policy stays in one place. -/// -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) - { - _log = log; - _job = job; - _tokenRefreshManager = tokenRefreshManager; - } - - public ISession CreateSourceSession() - => CassandraClientFactory.CreateSourceSession(_log, _job, _tokenRefreshManager); - - public Task CreateTargetSessionAsync() - => CassandraClientFactory.CreateTargetSessionAsync(_log, _job); -} diff --git a/CassandraMigrationProcessor/CassandraDriver/ITargetSessionFactory.cs b/CassandraMigrationProcessor/CassandraDriver/ITargetSessionFactory.cs new file mode 100644 index 0000000..c7f38b3 --- /dev/null +++ b/CassandraMigrationProcessor/CassandraDriver/ITargetSessionFactory.cs @@ -0,0 +1,67 @@ +using Cassandra; +using CassandraMigrationProcessor.Models; +using CassandraMigrationProcessor.Infrastructure; + +namespace CassandraMigrationProcessor.CassandraDriver; + +/// +/// Creates worker-owned target sessions. Source sessions are job-owned and +/// passed directly to readers, so their lifetime cannot be confused with the +/// per-worker target-session lifetime. +/// +public interface ITargetSessionFactory +{ + /// Mint a new keyspace-agnostic target-cluster session. Async because + /// target credential discovery may go through ARM. + Task CreateTargetSessionAsync(); +} + +/// +/// Limits simultaneous target-session opens while retaining one target +/// session per worker. This prevents high-worker jobs from creating a +/// connection storm during startup. +/// +public sealed class GatedTargetSessionFactory : ITargetSessionFactory +{ + private readonly ITargetSessionFactory _inner; + private readonly SemaphoreSlim _creationGate = new(2, 2); + + public GatedTargetSessionFactory(ITargetSessionFactory inner) + { + _inner = inner ?? throw new ArgumentNullException(nameof(inner)); + } + + public async Task CreateTargetSessionAsync() + { + await _creationGate.WaitAsync().ConfigureAwait(false); + try + { + return await _inner.CreateTargetSessionAsync() + .ConfigureAwait(false); + } + finally + { + _creationGate.Release(); + } + } +} + +/// +/// Default bound to a single +/// . Delegates to +/// so the connection-construction policy stays in one place. +/// +public sealed class JobTargetSessionFactory : ITargetSessionFactory +{ + private readonly MigrationLog _log; + private readonly Job _job; + + public JobTargetSessionFactory(MigrationLog log, Job job) + { + _log = log; + _job = job; + } + + public Task CreateTargetSessionAsync() + => CassandraClientFactory.CreateTargetSessionAsync(_log, _job); +} diff --git a/CassandraMigrationProcessor/DataTransfer/DataCopyWorker.cs b/CassandraMigrationProcessor/DataTransfer/DataCopyWorker.cs index 743a2de..0e760fa 100644 --- a/CassandraMigrationProcessor/DataTransfer/DataCopyWorker.cs +++ b/CassandraMigrationProcessor/DataTransfer/DataCopyWorker.cs @@ -33,8 +33,8 @@ public async Task RunAsync(PipelineContext ctx) Partition? current = null; try { - reader = await PageReader.CreateAsync(_workerLog, ctx.SessionFactory, ctx.ReaderConfig, _ct); - writer = await PageWriter.CreateAsync(_workerLog, ctx.SessionFactory, ctx.WriterConfig, _ct); + reader = await PageReader.CreateAsync(_workerLog, ctx.SourceSession, ctx.ReaderConfig, _ct); + writer = await PageWriter.CreateAsync(_workerLog, ctx.TargetSessionFactory, ctx.WriterConfig, _ct); while (!_ct.IsCancellationRequested && !ctx.Control.IsFatal) @@ -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"); } } diff --git a/CassandraMigrationProcessor/DataTransfer/JobPipeline.cs b/CassandraMigrationProcessor/DataTransfer/JobPipeline.cs index 590f909..7c50db5 100644 --- a/CassandraMigrationProcessor/DataTransfer/JobPipeline.cs +++ b/CassandraMigrationProcessor/DataTransfer/JobPipeline.cs @@ -1,3 +1,4 @@ +using Cassandra; using CassandraMigrationProcessor.CassandraDriver; using CassandraMigrationProcessor.Infrastructure; using CassandraMigrationProcessor.Models; @@ -21,7 +22,8 @@ internal sealed class JobPipeline : IDisposable, IAsyncDisposable public PipelineContext Context { get; } public JobPipeline(MigrationLog log, Job job, PipelineConfig pipelineConfig, - JobPartitioning partitioning, ISessionFactory sessionFactory, + JobPartitioning partitioning, ISession sourceSession, + ITargetSessionFactory targetSessionFactory, JobControl control) { _log = log; @@ -45,7 +47,8 @@ public JobPipeline(MigrationLog log, Job job, PipelineConfig pipelineConfig, Context = new PipelineContext( _partitions, - sessionFactory, + sourceSession, + targetSessionFactory, readerConfig, writerConfig, EnableReplay: enableReplay, diff --git a/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs b/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs index b4076f8..4de20a8 100644 --- a/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs +++ b/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs @@ -176,9 +176,9 @@ public async Task StartAsync() _pipeline = new JobPipeline( _log, job, _pipelineConfig, partitioning, - new SharedSourceSessionFactory( - _sourceSession, - new JobSessionFactory(_log, job, _tokenRefreshManager)), + _sourceSession, + new GatedTargetSessionFactory( + new JobTargetSessionFactory(_log, job)), _control); _pipeline.Start(); diff --git a/CassandraMigrationProcessor/DataTransfer/PageReader.cs b/CassandraMigrationProcessor/DataTransfer/PageReader.cs index 0d1ca40..f912246 100644 --- a/CassandraMigrationProcessor/DataTransfer/PageReader.cs +++ b/CassandraMigrationProcessor/DataTransfer/PageReader.cs @@ -25,12 +25,11 @@ 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. /// -internal class PageReader : IDisposable +internal class PageReader { 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; @@ -51,7 +50,7 @@ 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; @@ -59,21 +58,14 @@ private PageReader(WorkerLog log, ISessionFactory sessionFactory, ReaderConfig c _maxReadRetries = config.MaxReadRetries; _preserveCellTtl = config.PreserveCellTtlAndWritetime; _useJsonCopy = config.UseJsonCopy; - _sourceSession = sessionFactory.CreateSourceSession(); - _ownsSourceSession = sessionFactory.CallerOwnsSourceSession; + _sourceSession = sourceSession; } public static Task CreateAsync(WorkerLog log, - ISessionFactory sessionFactory, ReaderConfig config, + ISession sourceSession, ReaderConfig config, CancellationToken cancellationToken) { - return Task.FromResult(new PageReader(log, sessionFactory, config, cancellationToken)); - } - - public void Dispose() - { - if (_ownsSourceSession) - MigrationUtilities.SafeDisposeSession(_sourceSession, "PageReader source session"); + return Task.FromResult(new PageReader(log, sourceSession, config, cancellationToken)); } /// diff --git a/CassandraMigrationProcessor/DataTransfer/PageWriter.cs b/CassandraMigrationProcessor/DataTransfer/PageWriter.cs index cb81d6a..9ae6031 100644 --- a/CassandraMigrationProcessor/DataTransfer/PageWriter.cs +++ b/CassandraMigrationProcessor/DataTransfer/PageWriter.cs @@ -29,7 +29,6 @@ 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; @@ -47,7 +46,7 @@ internal sealed class PageWriter : IDisposable /// internal Exception? LastWriteException { get; private set; } - private PageWriter(WorkerLog log, ISession targetSession, bool ownsTargetSession, + private PageWriter(WorkerLog log, ISession targetSession, WriterConfig config, CancellationToken cancellationToken) { _log = log; @@ -57,20 +56,16 @@ private PageWriter(WorkerLog log, ISession targetSession, bool ownsTargetSession _preserveCellTtl = config.PreserveCellTtlAndWritetime; _useJsonCopy = config.UseJsonCopy; _targetSession = targetSession; - _ownsTargetSession = ownsTargetSession; } - public static async Task CreateAsync(WorkerLog log, ISessionFactory sessionFactory, WriterConfig config, CancellationToken cancellationToken) + public static async Task CreateAsync(WorkerLog log, ITargetSessionFactory sessionFactory, WriterConfig config, CancellationToken cancellationToken) { var targetSession = await sessionFactory.CreateTargetSessionAsync(); - return new PageWriter(log, targetSession, sessionFactory.CallerOwnsTargetSession, config, cancellationToken); + return new PageWriter(log, targetSession, config, cancellationToken); } public void Dispose() - { - if (_ownsTargetSession) - MigrationUtilities.SafeDisposeSession(_targetSession, "PageWriter target session"); - } + => MigrationUtilities.SafeDisposeSession(_targetSession, "PageWriter target session"); private Task GetStrategyAsync(Partition partition) { diff --git a/CassandraMigrationProcessor/DataTransfer/PipelineContext.cs b/CassandraMigrationProcessor/DataTransfer/PipelineContext.cs index 652df55..01aa7df 100644 --- a/CassandraMigrationProcessor/DataTransfer/PipelineContext.cs +++ b/CassandraMigrationProcessor/DataTransfer/PipelineContext.cs @@ -1,3 +1,4 @@ +using Cassandra; using CassandraMigrationProcessor.CassandraDriver; namespace CassandraMigrationProcessor.DataTransfer; @@ -6,15 +7,16 @@ namespace CassandraMigrationProcessor.DataTransfer; /// Shared (job-wide) state passed to every worker. Holds the /// 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, target-session factory, +/// reader / writer tunables, the replay configuration knobs, and /// the unified (cancellation + first-fault). /// Per-table state is resolved through /// pass-through accessors. /// internal record PipelineContext( PartitionManager Partitions, - ISessionFactory SessionFactory, + ISession SourceSession, + ITargetSessionFactory TargetSessionFactory, ReaderConfig ReaderConfig, WriterConfig WriterConfig, bool EnableReplay, diff --git a/CassandraMigrationProcessor/Models/TableCopySpec.cs b/CassandraMigrationProcessor/Models/TableCopySpec.cs index f06aacf..3cbc5c4 100644 --- a/CassandraMigrationProcessor/Models/TableCopySpec.cs +++ b/CassandraMigrationProcessor/Models/TableCopySpec.cs @@ -3,8 +3,8 @@ namespace CassandraMigrationProcessor.Models; /// /// 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 -/// ISessionFactory. +/// here — readers use the job-wide source session and writers open +/// worker-owned sessions through ITargetSessionFactory. /// public record TableCopySpec( string KeyspaceName, From a63bb4515ef8dcecd4b4a33fa2ad9ad11fbd6193 Mon Sep 17 00:00:00 2001 From: Nitesh Vijay Date: Tue, 18 Aug 2026 17:08:19 +0530 Subject: [PATCH 03/32] refactor: retain session factory name Keep the established ISessionFactory name while preserving explicit source and target ownership. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 078ac7b9-aa7f-44a1-9053-44c346cec5e1 --- .../{ITargetSessionFactory.cs => ISessionFactory.cs} | 12 ++++++------ .../DataTransfer/JobPipeline.cs | 2 +- .../DataTransfer/PageWriter.cs | 2 +- .../DataTransfer/PipelineContext.cs | 2 +- CassandraMigrationProcessor/Models/TableCopySpec.cs | 2 +- 5 files changed, 10 insertions(+), 10 deletions(-) rename CassandraMigrationProcessor/CassandraDriver/{ITargetSessionFactory.cs => ISessionFactory.cs} (83%) diff --git a/CassandraMigrationProcessor/CassandraDriver/ITargetSessionFactory.cs b/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs similarity index 83% rename from CassandraMigrationProcessor/CassandraDriver/ITargetSessionFactory.cs rename to CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs index c7f38b3..ef2d538 100644 --- a/CassandraMigrationProcessor/CassandraDriver/ITargetSessionFactory.cs +++ b/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs @@ -9,7 +9,7 @@ namespace CassandraMigrationProcessor.CassandraDriver; /// passed directly to readers, so their lifetime cannot be confused with the /// per-worker target-session lifetime. /// -public interface ITargetSessionFactory +public interface ISessionFactory { /// Mint a new keyspace-agnostic target-cluster session. Async because /// target credential discovery may go through ARM. @@ -21,12 +21,12 @@ public interface ITargetSessionFactory /// session per worker. This prevents high-worker jobs from creating a /// connection storm during startup. /// -public sealed class GatedTargetSessionFactory : ITargetSessionFactory +public sealed class GatedTargetSessionFactory : ISessionFactory { - private readonly ITargetSessionFactory _inner; + private readonly ISessionFactory _inner; private readonly SemaphoreSlim _creationGate = new(2, 2); - public GatedTargetSessionFactory(ITargetSessionFactory inner) + public GatedTargetSessionFactory(ISessionFactory inner) { _inner = inner ?? throw new ArgumentNullException(nameof(inner)); } @@ -47,11 +47,11 @@ public async Task CreateTargetSessionAsync() } /// -/// Default bound to a single +/// Default bound to a single /// . Delegates to /// so the connection-construction policy stays in one place. /// -public sealed class JobTargetSessionFactory : ITargetSessionFactory +public sealed class JobTargetSessionFactory : ISessionFactory { private readonly MigrationLog _log; private readonly Job _job; diff --git a/CassandraMigrationProcessor/DataTransfer/JobPipeline.cs b/CassandraMigrationProcessor/DataTransfer/JobPipeline.cs index 7c50db5..1d1666a 100644 --- a/CassandraMigrationProcessor/DataTransfer/JobPipeline.cs +++ b/CassandraMigrationProcessor/DataTransfer/JobPipeline.cs @@ -23,7 +23,7 @@ internal sealed class JobPipeline : IDisposable, IAsyncDisposable public JobPipeline(MigrationLog log, Job job, PipelineConfig pipelineConfig, JobPartitioning partitioning, ISession sourceSession, - ITargetSessionFactory targetSessionFactory, + ISessionFactory targetSessionFactory, JobControl control) { _log = log; diff --git a/CassandraMigrationProcessor/DataTransfer/PageWriter.cs b/CassandraMigrationProcessor/DataTransfer/PageWriter.cs index 9ae6031..5d4a215 100644 --- a/CassandraMigrationProcessor/DataTransfer/PageWriter.cs +++ b/CassandraMigrationProcessor/DataTransfer/PageWriter.cs @@ -58,7 +58,7 @@ private PageWriter(WorkerLog log, ISession targetSession, _targetSession = targetSession; } - public static async Task CreateAsync(WorkerLog log, ITargetSessionFactory sessionFactory, WriterConfig config, CancellationToken cancellationToken) + public static async Task CreateAsync(WorkerLog log, ISessionFactory sessionFactory, WriterConfig config, CancellationToken cancellationToken) { var targetSession = await sessionFactory.CreateTargetSessionAsync(); return new PageWriter(log, targetSession, config, cancellationToken); diff --git a/CassandraMigrationProcessor/DataTransfer/PipelineContext.cs b/CassandraMigrationProcessor/DataTransfer/PipelineContext.cs index 01aa7df..8e3dc78 100644 --- a/CassandraMigrationProcessor/DataTransfer/PipelineContext.cs +++ b/CassandraMigrationProcessor/DataTransfer/PipelineContext.cs @@ -16,7 +16,7 @@ namespace CassandraMigrationProcessor.DataTransfer; internal record PipelineContext( PartitionManager Partitions, ISession SourceSession, - ITargetSessionFactory TargetSessionFactory, + ISessionFactory TargetSessionFactory, ReaderConfig ReaderConfig, WriterConfig WriterConfig, bool EnableReplay, diff --git a/CassandraMigrationProcessor/Models/TableCopySpec.cs b/CassandraMigrationProcessor/Models/TableCopySpec.cs index 3cbc5c4..5d3ec94 100644 --- a/CassandraMigrationProcessor/Models/TableCopySpec.cs +++ b/CassandraMigrationProcessor/Models/TableCopySpec.cs @@ -4,7 +4,7 @@ namespace CassandraMigrationProcessor.Models; /// Immutable description of a single table copy. Identifies the source /// and target keyspace/table; runtime sessions are not threaded through /// here — readers use the job-wide source session and writers open -/// worker-owned sessions through ITargetSessionFactory. +/// worker-owned sessions through ISessionFactory. /// public record TableCopySpec( string KeyspaceName, From 8ea0f3759e623a89c2896b201fffbc7acdcde311 Mon Sep 17 00:00:00 2001 From: Nitesh Vijay Date: Tue, 18 Aug 2026 17:15:41 +0530 Subject: [PATCH 04/32] refactor: generalize session factory naming Use generic factory and method names while keeping the factory responsible for worker-created sessions at its call site. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 078ac7b9-aa7f-44a1-9053-44c346cec5e1 --- .../CassandraDriver/ISessionFactory.cs | 29 +++++++++---------- .../DataTransfer/DataCopyWorker.cs | 2 +- .../DataTransfer/JobPipeline.cs | 4 +-- .../DataTransfer/MigrationJobRunner.cs | 4 +-- .../DataTransfer/PageWriter.cs | 2 +- .../DataTransfer/PipelineContext.cs | 4 +-- 6 files changed, 21 insertions(+), 24 deletions(-) diff --git a/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs b/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs index ef2d538..04244f3 100644 --- a/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs +++ b/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs @@ -5,38 +5,35 @@ namespace CassandraMigrationProcessor.CassandraDriver; /// -/// Creates worker-owned target sessions. Source sessions are job-owned and -/// passed directly to readers, so their lifetime cannot be confused with the -/// per-worker target-session lifetime. +/// Creates worker-owned sessions. The consumer determines the session role; +/// job-owned shared sessions are passed directly instead of using this factory. /// public interface ISessionFactory { - /// Mint a new keyspace-agnostic target-cluster session. Async because - /// target credential discovery may go through ARM. - Task CreateTargetSessionAsync(); + /// Mint a new keyspace-agnostic session. + Task CreateSessionAsync(); } /// -/// Limits simultaneous target-session opens while retaining one target -/// session per worker. This prevents high-worker jobs from creating a -/// connection storm during startup. +/// Limits simultaneous session opens. This prevents high-worker jobs from +/// creating a connection storm during startup. /// -public sealed class GatedTargetSessionFactory : ISessionFactory +public sealed class GatedSessionFactory : ISessionFactory { private readonly ISessionFactory _inner; private readonly SemaphoreSlim _creationGate = new(2, 2); - public GatedTargetSessionFactory(ISessionFactory inner) + public GatedSessionFactory(ISessionFactory inner) { _inner = inner ?? throw new ArgumentNullException(nameof(inner)); } - public async Task CreateTargetSessionAsync() + public async Task CreateSessionAsync() { await _creationGate.WaitAsync().ConfigureAwait(false); try { - return await _inner.CreateTargetSessionAsync() + return await _inner.CreateSessionAsync() .ConfigureAwait(false); } finally @@ -51,17 +48,17 @@ public async Task CreateTargetSessionAsync() /// . Delegates to /// so the connection-construction policy stays in one place. /// -public sealed class JobTargetSessionFactory : ISessionFactory +public sealed class JobSessionFactory : ISessionFactory { private readonly MigrationLog _log; private readonly Job _job; - public JobTargetSessionFactory(MigrationLog log, Job job) + public JobSessionFactory(MigrationLog log, Job job) { _log = log; _job = job; } - public Task CreateTargetSessionAsync() + public Task CreateSessionAsync() => CassandraClientFactory.CreateTargetSessionAsync(_log, _job); } diff --git a/CassandraMigrationProcessor/DataTransfer/DataCopyWorker.cs b/CassandraMigrationProcessor/DataTransfer/DataCopyWorker.cs index 0e760fa..58252c3 100644 --- a/CassandraMigrationProcessor/DataTransfer/DataCopyWorker.cs +++ b/CassandraMigrationProcessor/DataTransfer/DataCopyWorker.cs @@ -34,7 +34,7 @@ public async Task RunAsync(PipelineContext ctx) try { reader = await PageReader.CreateAsync(_workerLog, ctx.SourceSession, ctx.ReaderConfig, _ct); - writer = await PageWriter.CreateAsync(_workerLog, ctx.TargetSessionFactory, ctx.WriterConfig, _ct); + writer = await PageWriter.CreateAsync(_workerLog, ctx.SessionFactory, ctx.WriterConfig, _ct); while (!_ct.IsCancellationRequested && !ctx.Control.IsFatal) diff --git a/CassandraMigrationProcessor/DataTransfer/JobPipeline.cs b/CassandraMigrationProcessor/DataTransfer/JobPipeline.cs index 1d1666a..87f6f74 100644 --- a/CassandraMigrationProcessor/DataTransfer/JobPipeline.cs +++ b/CassandraMigrationProcessor/DataTransfer/JobPipeline.cs @@ -23,7 +23,7 @@ internal sealed class JobPipeline : IDisposable, IAsyncDisposable public JobPipeline(MigrationLog log, Job job, PipelineConfig pipelineConfig, JobPartitioning partitioning, ISession sourceSession, - ISessionFactory targetSessionFactory, + ISessionFactory sessionFactory, JobControl control) { _log = log; @@ -48,7 +48,7 @@ public JobPipeline(MigrationLog log, Job job, PipelineConfig pipelineConfig, Context = new PipelineContext( _partitions, sourceSession, - targetSessionFactory, + sessionFactory, readerConfig, writerConfig, EnableReplay: enableReplay, diff --git a/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs b/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs index 4de20a8..9e8241f 100644 --- a/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs +++ b/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs @@ -177,8 +177,8 @@ public async Task StartAsync() _pipeline = new JobPipeline( _log, job, _pipelineConfig, partitioning, _sourceSession, - new GatedTargetSessionFactory( - new JobTargetSessionFactory(_log, job)), + new GatedSessionFactory( + new JobSessionFactory(_log, job)), _control); _pipeline.Start(); diff --git a/CassandraMigrationProcessor/DataTransfer/PageWriter.cs b/CassandraMigrationProcessor/DataTransfer/PageWriter.cs index 5d4a215..19fc641 100644 --- a/CassandraMigrationProcessor/DataTransfer/PageWriter.cs +++ b/CassandraMigrationProcessor/DataTransfer/PageWriter.cs @@ -60,7 +60,7 @@ private PageWriter(WorkerLog log, ISession targetSession, public static async Task 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); } diff --git a/CassandraMigrationProcessor/DataTransfer/PipelineContext.cs b/CassandraMigrationProcessor/DataTransfer/PipelineContext.cs index 8e3dc78..94cf6a6 100644 --- a/CassandraMigrationProcessor/DataTransfer/PipelineContext.cs +++ b/CassandraMigrationProcessor/DataTransfer/PipelineContext.cs @@ -7,7 +7,7 @@ namespace CassandraMigrationProcessor.DataTransfer; /// Shared (job-wide) state passed to every worker. Holds the /// that all tables seed into and /// every worker pulls from (and which owns the cooldown scheduler for -/// delayed recycles), the shared source session, target-session factory, +/// delayed recycles), the shared source session, worker session factory, /// reader / writer tunables, the replay configuration knobs, and /// the unified (cancellation + first-fault). /// Per-table state is resolved through @@ -16,7 +16,7 @@ namespace CassandraMigrationProcessor.DataTransfer; internal record PipelineContext( PartitionManager Partitions, ISession SourceSession, - ISessionFactory TargetSessionFactory, + ISessionFactory SessionFactory, ReaderConfig ReaderConfig, WriterConfig WriterConfig, bool EnableReplay, From 9dcc1d332cb83f5a456970abf3a3a7a9f5bf5cb0 Mon Sep 17 00:00:00 2001 From: Nitesh Vijay Date: Tue, 18 Aug 2026 17:19:18 +0530 Subject: [PATCH 05/32] perf: allow twenty concurrent session opens Increase the startup session-creation gate from two to twenty concurrent opens. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 078ac7b9-aa7f-44a1-9053-44c346cec5e1 --- .../CassandraDriver/ISessionFactory.cs | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs b/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs index 04244f3..19f3db2 100644 --- a/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs +++ b/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs @@ -20,8 +20,12 @@ public interface ISessionFactory /// public sealed class GatedSessionFactory : ISessionFactory { + private const int MaxConcurrentSessionCreations = 20; + private readonly ISessionFactory _inner; - private readonly SemaphoreSlim _creationGate = new(2, 2); + private readonly SemaphoreSlim _creationGate = new( + MaxConcurrentSessionCreations, + MaxConcurrentSessionCreations); public GatedSessionFactory(ISessionFactory inner) { From 341d8c748ad7a69aec1b9fb1a3cd10cd6a30ea88 Mon Sep 17 00:00:00 2001 From: Nitesh Vijay Date: Tue, 18 Aug 2026 17:32:59 +0530 Subject: [PATCH 06/32] fix: preserve source sessions during token rotation Lease the current managed source session for each reader operation so AAD refresh can rotate sessions without disposing one still in use. Make worker session creation cancellable and dispose the creation gate with the pipeline. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 078ac7b9-aa7f-44a1-9053-44c346cec5e1 --- .../CassandraDriver/ISessionFactory.cs | 45 ++++++++-- .../CassandraDriver/TokenRefreshManager.cs | 82 +++++++++++++++++-- .../DataTransfer/DataCopyWorker.cs | 2 +- .../DataTransfer/JobPipeline.cs | 9 +- .../DataTransfer/MigrationJobRunner.cs | 8 +- .../DataTransfer/PageReader.cs | 33 ++++---- .../DataTransfer/PageWriter.cs | 2 +- .../DataTransfer/PipelineContext.cs | 3 +- 8 files changed, 145 insertions(+), 39 deletions(-) diff --git a/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs b/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs index 19f3db2..85bc28d 100644 --- a/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs +++ b/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs @@ -11,14 +11,39 @@ namespace CassandraMigrationProcessor.CassandraDriver; public interface ISessionFactory { /// Mint a new keyspace-agnostic session. - Task CreateSessionAsync(); + Task CreateSessionAsync(CancellationToken cancellationToken); +} + +/// +/// Provides a lease on the current shared session. A rotated session is not +/// disposed until all operations using its leases have completed. +/// +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(); } /// /// Limits simultaneous session opens. This prevents high-worker jobs from /// creating a connection storm during startup. /// -public sealed class GatedSessionFactory : ISessionFactory +public sealed class GatedSessionFactory : ISessionFactory, IDisposable { private const int MaxConcurrentSessionCreations = 20; @@ -32,12 +57,12 @@ public GatedSessionFactory(ISessionFactory inner) _inner = inner ?? throw new ArgumentNullException(nameof(inner)); } - public async Task CreateSessionAsync() + public async Task CreateSessionAsync(CancellationToken cancellationToken) { - await _creationGate.WaitAsync().ConfigureAwait(false); + await _creationGate.WaitAsync(cancellationToken).ConfigureAwait(false); try { - return await _inner.CreateSessionAsync() + return await _inner.CreateSessionAsync(cancellationToken) .ConfigureAwait(false); } finally @@ -45,6 +70,8 @@ public async Task CreateSessionAsync() _creationGate.Release(); } } + + public void Dispose() => _creationGate.Dispose(); } /// @@ -63,6 +90,10 @@ public JobSessionFactory(MigrationLog log, Job job) _job = job; } - public Task CreateSessionAsync() - => CassandraClientFactory.CreateTargetSessionAsync(_log, _job); + public async Task CreateSessionAsync(CancellationToken cancellationToken) + { + cancellationToken.ThrowIfCancellationRequested(); + return await CassandraClientFactory.CreateTargetSessionAsync(_log, _job) + .ConfigureAwait(false); + } } diff --git a/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs b/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs index a9133e8..d4c401c 100644 --- a/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs +++ b/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs @@ -8,12 +8,22 @@ namespace CassandraMigrationProcessor.CassandraDriver; /// Manages AAD token lifecycle and proactive refresh for /// Cosmos DB Cassandra API connections. /// -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) /// 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"); } } diff --git a/CassandraMigrationProcessor/DataTransfer/DataCopyWorker.cs b/CassandraMigrationProcessor/DataTransfer/DataCopyWorker.cs index 58252c3..ccdeded 100644 --- a/CassandraMigrationProcessor/DataTransfer/DataCopyWorker.cs +++ b/CassandraMigrationProcessor/DataTransfer/DataCopyWorker.cs @@ -33,7 +33,7 @@ public async Task RunAsync(PipelineContext ctx) Partition? current = null; try { - reader = await PageReader.CreateAsync(_workerLog, ctx.SourceSession, 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 diff --git a/CassandraMigrationProcessor/DataTransfer/JobPipeline.cs b/CassandraMigrationProcessor/DataTransfer/JobPipeline.cs index 87f6f74..2b84dde 100644 --- a/CassandraMigrationProcessor/DataTransfer/JobPipeline.cs +++ b/CassandraMigrationProcessor/DataTransfer/JobPipeline.cs @@ -1,4 +1,3 @@ -using Cassandra; using CassandraMigrationProcessor.CassandraDriver; using CassandraMigrationProcessor.Infrastructure; using CassandraMigrationProcessor.Models; @@ -19,16 +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, ISession sourceSession, + JobPartitioning partitioning, ISessionProvider sourceSessionProvider, ISessionFactory sessionFactory, JobControl control) { _log = log; _pipelineConfig = pipelineConfig; _control = control; + _sessionFactory = sessionFactory; bool enableReplay = job.IsOnline; _partitions = new PartitionManager( @@ -47,7 +48,7 @@ public JobPipeline(MigrationLog log, Job job, PipelineConfig pipelineConfig, Context = new PipelineContext( _partitions, - sourceSession, + sourceSessionProvider, sessionFactory, readerConfig, writerConfig, @@ -100,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"); } } diff --git a/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs b/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs index 9e8241f..2ba3947 100644 --- a/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs +++ b/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs @@ -88,14 +88,14 @@ public static async Task 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; } } @@ -176,7 +176,7 @@ public async Task StartAsync() _pipeline = new JobPipeline( _log, job, _pipelineConfig, partitioning, - _sourceSession, + _tokenRefreshManager, new GatedSessionFactory( new JobSessionFactory(_log, job)), _control); @@ -251,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; } diff --git a/CassandraMigrationProcessor/DataTransfer/PageReader.cs b/CassandraMigrationProcessor/DataTransfer/PageReader.cs index f912246..0d1fa0e 100644 --- a/CassandraMigrationProcessor/DataTransfer/PageReader.cs +++ b/CassandraMigrationProcessor/DataTransfer/PageReader.cs @@ -29,12 +29,12 @@ internal class PageReader { private readonly WorkerLog _log; private readonly CancellationToken _ct; - private readonly ISession _sourceSession; + private readonly ISessionProvider _sourceSessionProvider; private readonly int _pageSize; private readonly int _maxReadRetries; private readonly bool _preserveCellTtl; private readonly bool _useJsonCopy; - private readonly ConcurrentDictionary _udtRegistrations = new(); + private readonly ConcurrentDictionary<(ISession Session, string Keyspace), Task> _udtRegistrations = new(); /// /// Most recent transient exception observed during retry-exhausted @@ -50,7 +50,7 @@ internal class PageReader // hints parking a worker for minutes. private const int MaxRetryDelayMs = 30_000; - private PageReader(WorkerLog log, ISession sourceSession, ReaderConfig config, CancellationToken cancellationToken) + private PageReader(WorkerLog log, ISessionProvider sourceSessionProvider, ReaderConfig config, CancellationToken cancellationToken) { _log = log; _ct = cancellationToken; @@ -58,14 +58,15 @@ private PageReader(WorkerLog log, ISession sourceSession, ReaderConfig config, C _maxReadRetries = config.MaxReadRetries; _preserveCellTtl = config.PreserveCellTtlAndWritetime; _useJsonCopy = config.UseJsonCopy; - _sourceSession = sourceSession; + _sourceSessionProvider = sourceSessionProvider + ?? throw new ArgumentNullException(nameof(sourceSessionProvider)); } public static Task CreateAsync(WorkerLog log, - ISession sourceSession, ReaderConfig config, + ISessionProvider sourceSessionProvider, ReaderConfig config, CancellationToken cancellationToken) { - return Task.FromResult(new PageReader(log, sourceSession, config, cancellationToken)); + return Task.FromResult(new PageReader(log, sourceSessionProvider, config, cancellationToken)); } /// @@ -73,27 +74,30 @@ public static Task CreateAsync(WorkerLog log, /// table registers every UDT in the keyspace because this reader can /// subsequently process other tables that reference different UDTs. /// - private Task EnsureUdtsRegisteredAsync(Partition partition) + private async Task EnsureUdtsRegisteredAsync(Partition partition) { // JSON read path bypasses CLR-side UDT decoding entirely. if (!partition.Table.IsCounterTable && _useJsonCopy) - return Task.CompletedTask; + return; - return _udtRegistrations.GetOrAdd(partition.Table.Spec.KeyspaceName, async ks => + using var lease = _sourceSessionProvider.AcquireSession(); + var sourceSession = lease.Session; + var keyspace = partition.Table.Spec.KeyspaceName; + await _udtRegistrations.GetOrAdd((sourceSession, keyspace), async key => { try { - var allUdts = await SchemaManager.GetUserDefinedTypesAsync(_sourceSession, ks); - await DynamicUdtRegistrar.RegisterAsync(_sourceSession, ks, allUdts); + var allUdts = await SchemaManager.GetUserDefinedTypesAsync(key.Session, key.Keyspace); + await DynamicUdtRegistrar.RegisterAsync(key.Session, key.Keyspace, allUdts); } catch (Exception ex) { // Do NOT swallow: UDT mapping is required for correct // row decoding. Surface as fatal. - _log.WriteLine($"FATAL: UDT mapping registration on source failed for {ks}: {ex.Message}", LogType.Error); + _log.WriteLine($"FATAL: UDT mapping registration on source failed for {key.Keyspace}: {ex.Message}", LogType.Error); throw; } - }); + }).ConfigureAwait(false); } /// @@ -214,8 +218,9 @@ internal record ReadResult( // re-queue this partition via cooldown — LastPagingState is // intact and will retry the same page once the source stops // throttling. + using var lease = _sourceSessionProvider.AcquireSession(); var resultSet = await RetryExecutor.ExecuteOrDefaultAsync( - operation: _ => _sourceSession.ExecuteAsync(stmt).WaitAsync(_ct), + operation: _ => lease.Session.ExecuteAsync(stmt).WaitAsync(_ct), maxAttempts: _maxReadRetries, shouldRetry: ExceptionClassifier.IsTransient, delayFor: (ex, attempt) => TimeSpan.FromMilliseconds( diff --git a/CassandraMigrationProcessor/DataTransfer/PageWriter.cs b/CassandraMigrationProcessor/DataTransfer/PageWriter.cs index 19fc641..08a7ed0 100644 --- a/CassandraMigrationProcessor/DataTransfer/PageWriter.cs +++ b/CassandraMigrationProcessor/DataTransfer/PageWriter.cs @@ -60,7 +60,7 @@ private PageWriter(WorkerLog log, ISession targetSession, public static async Task CreateAsync(WorkerLog log, ISessionFactory sessionFactory, WriterConfig config, CancellationToken cancellationToken) { - var targetSession = await sessionFactory.CreateSessionAsync(); + var targetSession = await sessionFactory.CreateSessionAsync(cancellationToken); return new PageWriter(log, targetSession, config, cancellationToken); } diff --git a/CassandraMigrationProcessor/DataTransfer/PipelineContext.cs b/CassandraMigrationProcessor/DataTransfer/PipelineContext.cs index 94cf6a6..ad2d63f 100644 --- a/CassandraMigrationProcessor/DataTransfer/PipelineContext.cs +++ b/CassandraMigrationProcessor/DataTransfer/PipelineContext.cs @@ -1,4 +1,3 @@ -using Cassandra; using CassandraMigrationProcessor.CassandraDriver; namespace CassandraMigrationProcessor.DataTransfer; @@ -15,7 +14,7 @@ namespace CassandraMigrationProcessor.DataTransfer; /// internal record PipelineContext( PartitionManager Partitions, - ISession SourceSession, + ISessionProvider SourceSessionProvider, ISessionFactory SessionFactory, ReaderConfig ReaderConfig, WriterConfig WriterConfig, From 1f5c1e96e13607418e6b58bc515a2e21adc3c4eb Mon Sep 17 00:00:00 2001 From: Nitesh Vijay Date: Tue, 18 Aug 2026 17:43:34 +0530 Subject: [PATCH 07/32] refactor: defer disposal of rotated sessions Keep retired source sessions alive for a bounded two-minute grace period and resolve the current session for every read retry attempt, allowing retries to move to a refreshed AAD session without lease tracking. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 078ac7b9-aa7f-44a1-9053-44c346cec5e1 --- .../CassandraDriver/ISessionFactory.cs | 23 +----- .../CassandraDriver/TokenRefreshManager.cs | 82 +++++++++---------- .../DataTransfer/PageReader.cs | 8 +- 3 files changed, 47 insertions(+), 66 deletions(-) diff --git a/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs b/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs index 85bc28d..9c8a06f 100644 --- a/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs +++ b/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs @@ -15,28 +15,13 @@ public interface ISessionFactory } /// -/// Provides a lease on the current shared session. A rotated session is not -/// disposed until all operations using its leases have completed. +/// Resolves the current shared session. Implementations may rotate the +/// underlying session while keeping retired instances alive for in-flight +/// operations. /// 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(); + ISession GetSession(); } /// diff --git a/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs b/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs index d4c401c..430ab83 100644 --- a/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs +++ b/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs @@ -10,23 +10,21 @@ namespace CassandraMigrationProcessor.CassandraDriver; /// 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 ManagedSessionState? _managedSourceSession; + private ISession? _managedSourceSession; + private readonly HashSet _retiredSourceSessions = + new(ReferenceEqualityComparer.Instance); private readonly MigrationLog _log; private bool _disposed; private DateTime _tokenExpiresAt = DateTime.MinValue; private int _consecutiveRefreshFailures; private const int MaxRefreshFailures = 6; + // A read attempt is capped at 60 seconds. Keep a rotated session alive + // for twice that time so an in-flight attempt can finish, while the next + // retry resolves and uses the newly refreshed session. + private static readonly TimeSpan RetiredSessionDisposalDelay = + TimeSpan.FromMinutes(2); private string? _lastSourceContactPoint; private int _lastSourcePort; @@ -163,7 +161,7 @@ private void TokenRefreshCallback(object? state) // If we have a managed session, recreate it if (_managedSourceSession != null - && !_managedSourceSession.Session.IsDisposed + && !_managedSourceSession.IsDisposed && _lastSourceContactPoint != null) { var newSession = CassandraClientFactory.CreateSourceSession( @@ -210,71 +208,69 @@ public void SetManagedSourceSession(ISession session) { ArgumentNullException.ThrowIfNull(session); - ISession? sessionToDispose = null; + ISession? retiredSession = null; lock (_refreshLock) { ObjectDisposedException.ThrowIf(_disposed, this); - if (ReferenceEquals(_managedSourceSession?.Session, session)) + if (ReferenceEquals(_managedSourceSession, session)) return; - var previous = _managedSourceSession; - _managedSourceSession = new ManagedSessionState(session); - if (previous != null) - { - previous.Retired = true; - if (previous.ActiveLeases == 0) - sessionToDispose = previous.Session; - } + retiredSession = _managedSourceSession; + _managedSourceSession = session; + if (retiredSession != null) + _retiredSourceSessions.Add(retiredSession); } - MigrationUtilities.SafeDisposeSession( - sessionToDispose, "TokenRefresh retired source session"); + if (retiredSession != null) + _ = DisposeRetiredSessionAfterDelayAsync(retiredSession); } - public SessionLease AcquireSession() + public ISession GetSession() { lock (_refreshLock) { ObjectDisposedException.ThrowIf(_disposed, this); - var state = _managedSourceSession + return _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) + private async Task DisposeRetiredSessionAfterDelayAsync(ISession session) { - ISession? sessionToDispose = null; + await Task.Delay(RetiredSessionDisposalDelay).ConfigureAwait(false); + + bool shouldDispose; lock (_refreshLock) { - state.ActiveLeases--; - if (state.Retired && state.ActiveLeases == 0) - sessionToDispose = state.Session; + shouldDispose = _retiredSourceSessions.Remove(session); } - MigrationUtilities.SafeDisposeSession( - sessionToDispose, "TokenRefresh retired source session"); + if (shouldDispose) + { + MigrationUtilities.SafeDisposeSession( + session, "TokenRefresh deferred source session"); + } } public void Dispose() { - ISession? sessionToDispose = null; + List sessionsToDispose; lock (_refreshLock) { if (_disposed) return; _disposed = true; StopTokenRefreshTimer(); + sessionsToDispose = _retiredSourceSessions.ToList(); + _retiredSourceSessions.Clear(); if (_managedSourceSession != null) - { - _managedSourceSession.Retired = true; - if (_managedSourceSession.ActiveLeases == 0) - sessionToDispose = _managedSourceSession.Session; - _managedSourceSession = null; - } + sessionsToDispose.Add(_managedSourceSession); + _managedSourceSession = null; } - MigrationUtilities.SafeDisposeSession( - sessionToDispose, "TokenRefresh managed source session"); + foreach (var session in sessionsToDispose) + { + MigrationUtilities.SafeDisposeSession( + session, "TokenRefresh managed source session"); + } } } diff --git a/CassandraMigrationProcessor/DataTransfer/PageReader.cs b/CassandraMigrationProcessor/DataTransfer/PageReader.cs index 0d1fa0e..6d1435e 100644 --- a/CassandraMigrationProcessor/DataTransfer/PageReader.cs +++ b/CassandraMigrationProcessor/DataTransfer/PageReader.cs @@ -80,8 +80,7 @@ private async Task EnsureUdtsRegisteredAsync(Partition partition) if (!partition.Table.IsCounterTable && _useJsonCopy) return; - using var lease = _sourceSessionProvider.AcquireSession(); - var sourceSession = lease.Session; + var sourceSession = _sourceSessionProvider.GetSession(); var keyspace = partition.Table.Spec.KeyspaceName; await _udtRegistrations.GetOrAdd((sourceSession, keyspace), async key => { @@ -218,9 +217,10 @@ internal record ReadResult( // re-queue this partition via cooldown — LastPagingState is // intact and will retry the same page once the source stops // throttling. - using var lease = _sourceSessionProvider.AcquireSession(); var resultSet = await RetryExecutor.ExecuteOrDefaultAsync( - operation: _ => lease.Session.ExecuteAsync(stmt).WaitAsync(_ct), + operation: _ => _sourceSessionProvider.GetSession() + .ExecuteAsync(stmt) + .WaitAsync(_ct), maxAttempts: _maxReadRetries, shouldRetry: ExceptionClassifier.IsTransient, delayFor: (ex, attempt) => TimeSpan.FromMilliseconds( From 7a30d26f4a8d2099fd19901d4f9755f407944973 Mon Sep 17 00:00:00 2001 From: Nitesh Vijay Date: Tue, 18 Aug 2026 17:44:28 +0530 Subject: [PATCH 08/32] refactor: extend session disposal grace period Keep rotated source sessions alive for ten minutes before disposal while retries continue resolving the latest session. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 078ac7b9-aa7f-44a1-9053-44c346cec5e1 --- .../CassandraDriver/TokenRefreshManager.cs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs b/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs index 430ab83..334c00e 100644 --- a/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs +++ b/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs @@ -21,10 +21,10 @@ public class TokenRefreshManager : ISessionProvider, IDisposable private int _consecutiveRefreshFailures; private const int MaxRefreshFailures = 6; // A read attempt is capped at 60 seconds. Keep a rotated session alive - // for twice that time so an in-flight attempt can finish, while the next - // retry resolves and uses the newly refreshed session. + // for a generous bounded grace period so in-flight operations can finish, + // while each retry resolves and uses the newly refreshed session. private static readonly TimeSpan RetiredSessionDisposalDelay = - TimeSpan.FromMinutes(2); + TimeSpan.FromMinutes(10); private string? _lastSourceContactPoint; private int _lastSourcePort; From 9da94498ab0782cd8cc25da7dae7a27aea917b7c Mon Sep 17 00:00:00 2001 From: Nitesh Vijay Date: Tue, 18 Aug 2026 17:49:32 +0530 Subject: [PATCH 09/32] refactor: separate session rotation from token refresh Move current and deferred session ownership into RotatingSessionProvider. TokenRefreshManager now handles token scheduling only and registers refreshed sessions with the provider. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 078ac7b9-aa7f-44a1-9053-44c346cec5e1 --- .../CassandraDriver/CassandraClientFactory.cs | 2 +- .../CassandraDriver/ISessionFactory.cs | 10 -- .../RotatingSessionProvider.cs | 103 ++++++++++++++++++ .../CassandraDriver/TokenRefreshManager.cs | 87 ++------------- .../DataTransfer/MigrationJobRunner.cs | 16 ++- 5 files changed, 128 insertions(+), 90 deletions(-) create mode 100644 CassandraMigrationProcessor/CassandraDriver/RotatingSessionProvider.cs diff --git a/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs b/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs index 6932b6f..d29d46f 100644 --- a/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs +++ b/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs @@ -105,7 +105,7 @@ private static void RegisterAadTokenRefresh( TokenRefreshManager? tokenRefreshManager) { if (!TokenRefreshManager.IsLikelyAadToken(password)) return; - tokenRefreshManager?.SetManagedSourceSession(session); + tokenRefreshManager?.RegisterSourceSession(session); tokenRefreshManager?.StartTokenRefreshTimer(password); } diff --git a/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs b/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs index 9c8a06f..0b625e6 100644 --- a/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs +++ b/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs @@ -14,16 +14,6 @@ public interface ISessionFactory Task CreateSessionAsync(CancellationToken cancellationToken); } -/// -/// Resolves the current shared session. Implementations may rotate the -/// underlying session while keeping retired instances alive for in-flight -/// operations. -/// -public interface ISessionProvider -{ - ISession GetSession(); -} - /// /// Limits simultaneous session opens. This prevents high-worker jobs from /// creating a connection storm during startup. diff --git a/CassandraMigrationProcessor/CassandraDriver/RotatingSessionProvider.cs b/CassandraMigrationProcessor/CassandraDriver/RotatingSessionProvider.cs new file mode 100644 index 0000000..f19df90 --- /dev/null +++ b/CassandraMigrationProcessor/CassandraDriver/RotatingSessionProvider.cs @@ -0,0 +1,103 @@ +using Cassandra; +using CassandraMigrationProcessor.Infrastructure; + +namespace CassandraMigrationProcessor.CassandraDriver; + +/// +/// Resolves the current shared session while retaining rotated sessions for a +/// bounded grace period so in-flight operations can complete. +/// +public interface ISessionProvider +{ + ISession GetSession(); +} + +public sealed class RotatingSessionProvider : ISessionProvider, IDisposable +{ + private static readonly TimeSpan RetiredSessionDisposalDelay = + TimeSpan.FromMinutes(10); + + private readonly object _sync = new(); + private readonly HashSet _retiredSessions = + new(ReferenceEqualityComparer.Instance); + private ISession? _currentSession; + private bool _disposed; + + public ISession GetSession() + { + lock (_sync) + { + ObjectDisposedException.ThrowIf(_disposed, this); + return _currentSession + ?? throw new InvalidOperationException("The session provider has not been initialized."); + } + } + + internal bool TryGetSession(out ISession? session) + { + lock (_sync) + { + session = _currentSession; + return !_disposed && session != null; + } + } + + internal void SetSession(ISession session) + { + ArgumentNullException.ThrowIfNull(session); + + ISession? retiredSession; + lock (_sync) + { + ObjectDisposedException.ThrowIf(_disposed, this); + if (ReferenceEquals(_currentSession, session)) + return; + + retiredSession = _currentSession; + _currentSession = session; + if (retiredSession != null) + _retiredSessions.Add(retiredSession); + } + + if (retiredSession != null) + _ = DisposeRetiredSessionAfterDelayAsync(retiredSession); + } + + private async Task DisposeRetiredSessionAfterDelayAsync(ISession session) + { + await Task.Delay(RetiredSessionDisposalDelay).ConfigureAwait(false); + + bool shouldDispose; + lock (_sync) + { + shouldDispose = _retiredSessions.Remove(session); + } + + if (shouldDispose) + { + MigrationUtilities.SafeDisposeSession( + session, "Deferred rotated session"); + } + } + + public void Dispose() + { + List sessionsToDispose; + lock (_sync) + { + if (_disposed) return; + _disposed = true; + sessionsToDispose = _retiredSessions.ToList(); + _retiredSessions.Clear(); + if (_currentSession != null) + sessionsToDispose.Add(_currentSession); + _currentSession = null; + } + + foreach (var session in sessionsToDispose) + { + MigrationUtilities.SafeDisposeSession( + session, "Rotating session provider"); + } + } +} diff --git a/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs b/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs index 334c00e..995c82a 100644 --- a/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs +++ b/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs @@ -8,31 +8,27 @@ namespace CassandraMigrationProcessor.CassandraDriver; /// Manages AAD token lifecycle and proactive refresh for /// Cosmos DB Cassandra API connections. /// -public class TokenRefreshManager : ISessionProvider, IDisposable +public class TokenRefreshManager : IDisposable { private Timer? _tokenRefreshTimer; private readonly object _refreshLock = new(); - private ISession? _managedSourceSession; - private readonly HashSet _retiredSourceSessions = - new(ReferenceEqualityComparer.Instance); + private readonly RotatingSessionProvider _sourceSessions; private readonly MigrationLog _log; - private bool _disposed; private DateTime _tokenExpiresAt = DateTime.MinValue; private int _consecutiveRefreshFailures; private const int MaxRefreshFailures = 6; - // A read attempt is capped at 60 seconds. Keep a rotated session alive - // for a generous bounded grace period so in-flight operations can finish, - // while each retry resolves and uses the newly refreshed session. - private static readonly TimeSpan RetiredSessionDisposalDelay = - TimeSpan.FromMinutes(10); private string? _lastSourceContactPoint; private int _lastSourcePort; private string? _lastSourceUsername; - public TokenRefreshManager(MigrationLog log) + public TokenRefreshManager( + MigrationLog log, + RotatingSessionProvider sourceSessions) { _log = log; + _sourceSessions = sourceSessions + ?? throw new ArgumentNullException(nameof(sourceSessions)); } /// @@ -160,8 +156,8 @@ private void TokenRefreshCallback(object? state) string freshToken = GetFreshAadToken(); // If we have a managed session, recreate it - if (_managedSourceSession != null - && !_managedSourceSession.IsDisposed + if (_sourceSessions.TryGetSession(out var currentSession) + && !currentSession!.IsDisposed && _lastSourceContactPoint != null) { var newSession = CassandraClientFactory.CreateSourceSession( @@ -170,7 +166,7 @@ private void TokenRefreshCallback(object? state) _lastSourcePort, _lastSourceUsername ?? string.Empty, freshToken); - SetManagedSourceSession(newSession); + _sourceSessions.SetSession(newSession); } // Schedule next refresh @@ -204,73 +200,14 @@ private void TokenRefreshCallback(object? state) /// Set the managed source session so the token refresh /// timer can reconnect it proactively. /// - public void SetManagedSourceSession(ISession session) - { - ArgumentNullException.ThrowIfNull(session); - - ISession? retiredSession = null; - lock (_refreshLock) - { - ObjectDisposedException.ThrowIf(_disposed, this); - if (ReferenceEquals(_managedSourceSession, session)) - return; - - retiredSession = _managedSourceSession; - _managedSourceSession = session; - if (retiredSession != null) - _retiredSourceSessions.Add(retiredSession); - } - - if (retiredSession != null) - _ = DisposeRetiredSessionAfterDelayAsync(retiredSession); - } - - public ISession GetSession() - { - lock (_refreshLock) - { - ObjectDisposedException.ThrowIf(_disposed, this); - return _managedSourceSession - ?? throw new InvalidOperationException("The source session has not been initialized."); - } - } - - private async Task DisposeRetiredSessionAfterDelayAsync(ISession session) - { - await Task.Delay(RetiredSessionDisposalDelay).ConfigureAwait(false); - - bool shouldDispose; - lock (_refreshLock) - { - shouldDispose = _retiredSourceSessions.Remove(session); - } - - if (shouldDispose) - { - MigrationUtilities.SafeDisposeSession( - session, "TokenRefresh deferred source session"); - } - } + public void RegisterSourceSession(ISession session) + => _sourceSessions.SetSession(session); public void Dispose() { - List sessionsToDispose; lock (_refreshLock) { - if (_disposed) return; - _disposed = true; StopTokenRefreshTimer(); - sessionsToDispose = _retiredSourceSessions.ToList(); - _retiredSourceSessions.Clear(); - if (_managedSourceSession != null) - sessionsToDispose.Add(_managedSourceSession); - _managedSourceSession = null; - } - - foreach (var session in sessionsToDispose) - { - MigrationUtilities.SafeDisposeSession( - session, "TokenRefresh managed source session"); } } } diff --git a/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs b/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs index 2ba3947..795a91f 100644 --- a/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs +++ b/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs @@ -20,6 +20,7 @@ public class MigrationJobRunner : IAsyncDisposable private readonly PipelineConfig _pipelineConfig; private readonly JobControl _control; private readonly TokenRefreshManager _tokenRefreshManager; + private readonly RotatingSessionProvider _sourceSessions; private int _consecutiveAuthErrors; // Last auth exception observed by HandleMigrationUnitError; // attached as inner when the consecutive-auth threshold trips so @@ -55,6 +56,7 @@ private MigrationJobRunner( PipelineConfig pipelineConfig, JobControl control, TokenRefreshManager tokenRefreshManager, + RotatingSessionProvider sourceSessions, ISession sourceSession, ISession targetSession) { @@ -63,6 +65,7 @@ private MigrationJobRunner( _pipelineConfig = pipelineConfig; _control = control; _tokenRefreshManager = tokenRefreshManager; + _sourceSessions = sourceSessions; _sourceSession = sourceSession; _targetSession = targetSession; } @@ -82,20 +85,24 @@ public static async Task CreateAsync( ArgumentNullException.ThrowIfNull(control); var pipelineConfig = PipelineConfig.Resolve(job, config); - var tokenRefreshManager = new TokenRefreshManager(log); + var sourceSessions = new RotatingSessionProvider(); + var tokenRefreshManager = new TokenRefreshManager(log, sourceSessions); ISession? source = null; ISession? target = null; try { source = CassandraClientFactory.CreateSourceSession(log, job, tokenRefreshManager); - tokenRefreshManager.SetManagedSourceSession(source); + sourceSessions.SetSession(source); target = await CassandraClientFactory.CreateTargetSessionAsync(log, job); - return new MigrationJobRunner(log, job, pipelineConfig, control, tokenRefreshManager, source, target); + return new MigrationJobRunner( + log, job, pipelineConfig, control, tokenRefreshManager, + sourceSessions, source, target); } catch { MigrationUtilities.SafeDisposeSession(target, "MigrationJobRunner target (CreateAsync rollback)"); tokenRefreshManager.Dispose(); + sourceSessions.Dispose(); throw; } } @@ -176,7 +183,7 @@ public async Task StartAsync() _pipeline = new JobPipeline( _log, job, _pipelineConfig, partitioning, - _tokenRefreshManager, + _sourceSessions, new GatedSessionFactory( new JobSessionFactory(_log, job)), _control); @@ -252,6 +259,7 @@ public ValueTask DisposeAsync() _pipeline = null; MigrationUtilities.SafeDisposeSession(_targetSession, "MigrationJobRunner target session"); _tokenRefreshManager.Dispose(); + _sourceSessions.Dispose(); return ValueTask.CompletedTask; } From 25f2cc753d7405c656a58a54bda3405483139d6f Mon Sep 17 00:00:00 2001 From: Nitesh Vijay Date: Tue, 18 Aug 2026 17:53:50 +0530 Subject: [PATCH 10/32] refactor: let session provider refresh itself Give RotatingSessionProvider a credential-aware session factory and replace external session assignment with Initialize and Refresh operations. TokenRefreshManager now supplies only refreshed credentials. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 078ac7b9-aa7f-44a1-9053-44c346cec5e1 --- .../CassandraDriver/CassandraClientFactory.cs | 40 ++++++++++------ .../RotatingSessionProvider.cs | 47 +++++++++++++------ .../CassandraDriver/TokenRefreshManager.cs | 38 +-------------- .../DataTransfer/MigrationJobRunner.cs | 6 ++- 4 files changed, 64 insertions(+), 67 deletions(-) diff --git a/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs b/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs index d29d46f..1c2ec24 100644 --- a/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs +++ b/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs @@ -48,10 +48,6 @@ public static ISession CreateSourceSession( 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, @@ -69,7 +65,7 @@ public static ISession CreateSourceSession( try { var session = ConnectCluster(builder); - RegisterAadTokenRefresh(session, password, tokenRefreshManager); + RegisterAadTokenRefresh(password, tokenRefreshManager); return session; } catch (Exception ex) when ( @@ -94,18 +90,16 @@ public static ISession CreateSourceSession( /// /// When looks like an AAD/JWT bearer /// token and the caller wired up a , - /// hand the freshly-connected off so the - /// proactive refresh timer can rotate the bearer before it expires. + /// start the proactive refresh timer so the bearer is rotated before it + /// expires. /// No-op when the password is a static credential or the manager is /// not supplied. /// private static void RegisterAadTokenRefresh( - ISession session, string password, TokenRefreshManager? tokenRefreshManager) { if (!TokenRefreshManager.IsLikelyAadToken(password)) return; - tokenRefreshManager?.RegisterSourceSession(session); tokenRefreshManager?.StartTokenRefreshTimer(password); } @@ -332,9 +326,27 @@ public static ISession CreateSourceSession( job.SourceUseAad = true; } - // For AAD auth, derive username from hostname if - // not explicitly provided (account name = first - // segment of the contact point FQDN). + return CreateSourceSessionWithCredential( + MigrationLog, job, password, tokenRefreshManager); + } + + internal static ISession CreateSourceSessionWithCredential( + MigrationLog migrationLog, + Job job, + string credential) + { + return CreateSourceSessionWithCredential( + migrationLog, job, credential, tokenRefreshManager: null); + } + + private static ISession CreateSourceSessionWithCredential( + MigrationLog migrationLog, + Job job, + string credential, + TokenRefreshManager? tokenRefreshManager) + { + // 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 @@ -345,11 +357,11 @@ public static ISession CreateSourceSession( } return CreateSourceSession( - MigrationLog, + migrationLog, job.SourceContactPoint, job.SourcePort, username, - password, + credential, tokenRefreshManager, maxConnectionsPerHost: ResolveMaxConnectionsPerHost(job.SourceMaxConnectionsPerHost, job.MaxConnectionsPerHost)); } diff --git a/CassandraMigrationProcessor/CassandraDriver/RotatingSessionProvider.cs b/CassandraMigrationProcessor/CassandraDriver/RotatingSessionProvider.cs index f19df90..651e9b8 100644 --- a/CassandraMigrationProcessor/CassandraDriver/RotatingSessionProvider.cs +++ b/CassandraMigrationProcessor/CassandraDriver/RotatingSessionProvider.cs @@ -18,11 +18,18 @@ public sealed class RotatingSessionProvider : ISessionProvider, IDisposable TimeSpan.FromMinutes(10); private readonly object _sync = new(); + private readonly Func _sessionFactory; private readonly HashSet _retiredSessions = new(ReferenceEqualityComparer.Instance); private ISession? _currentSession; private bool _disposed; + public RotatingSessionProvider(Func sessionFactory) + { + _sessionFactory = sessionFactory + ?? throw new ArgumentNullException(nameof(sessionFactory)); + } + public ISession GetSession() { lock (_sync) @@ -33,34 +40,46 @@ public ISession GetSession() } } - internal bool TryGetSession(out ISession? session) + public void Initialize(ISession session) { + ArgumentNullException.ThrowIfNull(session); + lock (_sync) { - session = _currentSession; - return !_disposed && session != null; + ObjectDisposedException.ThrowIf(_disposed, this); + if (_currentSession != null) + throw new InvalidOperationException("The session provider is already initialized."); + _currentSession = session; } } - internal void SetSession(ISession session) + public void Refresh(string credential) { - ArgumentNullException.ThrowIfNull(session); + ArgumentException.ThrowIfNullOrWhiteSpace(credential); + var session = _sessionFactory(credential); ISession? retiredSession; - lock (_sync) + try { - ObjectDisposedException.ThrowIf(_disposed, this); - if (ReferenceEquals(_currentSession, session)) - return; + lock (_sync) + { + ObjectDisposedException.ThrowIf(_disposed, this); + if (_currentSession == null) + throw new InvalidOperationException("The session provider has not been initialized."); - retiredSession = _currentSession; - _currentSession = session; - if (retiredSession != null) + retiredSession = _currentSession; + _currentSession = session; _retiredSessions.Add(retiredSession); + } + } + catch + { + MigrationUtilities.SafeDisposeSession( + session, "Unpublished refreshed session"); + throw; } - if (retiredSession != null) - _ = DisposeRetiredSessionAfterDelayAsync(retiredSession); + _ = DisposeRetiredSessionAfterDelayAsync(retiredSession); } private async Task DisposeRetiredSessionAfterDelayAsync(ISession session) diff --git a/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs b/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs index 995c82a..9076550 100644 --- a/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs +++ b/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs @@ -1,4 +1,3 @@ -using Cassandra; using System.IdentityModel.Tokens.Jwt; using CassandraMigrationProcessor.Infrastructure; using CassandraMigrationProcessor.Models; @@ -18,10 +17,6 @@ public class TokenRefreshManager : IDisposable private int _consecutiveRefreshFailures; private const int MaxRefreshFailures = 6; - private string? _lastSourceContactPoint; - private int _lastSourcePort; - private string? _lastSourceUsername; - public TokenRefreshManager( MigrationLog log, RotatingSessionProvider sourceSessions) @@ -31,18 +26,6 @@ public TokenRefreshManager( ?? throw new ArgumentNullException(nameof(sourceSessions)); } - /// - /// Cache source connection parameters so the token refresh - /// timer can reconnect with a fresh token. - /// - internal void CacheSourceConnectionParams( - string contactPoint, int port, string username) - { - _lastSourceContactPoint = contactPoint; - _lastSourcePort = port; - _lastSourceUsername = username; - } - /// /// Detect if a password looks like an AAD/JWT token /// (very long base64-ish string). @@ -155,19 +138,7 @@ private void TokenRefreshCallback(object? state) { 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); - } + _sourceSessions.Refresh(freshToken); // Schedule next refresh _consecutiveRefreshFailures = 0; @@ -196,13 +167,6 @@ private void TokenRefreshCallback(object? state) } } - /// - /// Set the managed source session so the token refresh - /// timer can reconnect it proactively. - /// - public void RegisterSourceSession(ISession session) - => _sourceSessions.SetSession(session); - public void Dispose() { lock (_refreshLock) diff --git a/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs b/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs index 795a91f..fae15af 100644 --- a/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs +++ b/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs @@ -85,14 +85,16 @@ public static async Task CreateAsync( ArgumentNullException.ThrowIfNull(control); var pipelineConfig = PipelineConfig.Resolve(job, config); - var sourceSessions = new RotatingSessionProvider(); + var sourceSessions = new RotatingSessionProvider( + credential => CassandraClientFactory.CreateSourceSessionWithCredential( + log, job, credential)); var tokenRefreshManager = new TokenRefreshManager(log, sourceSessions); ISession? source = null; ISession? target = null; try { source = CassandraClientFactory.CreateSourceSession(log, job, tokenRefreshManager); - sourceSessions.SetSession(source); + sourceSessions.Initialize(source); target = await CassandraClientFactory.CreateTargetSessionAsync(log, job); return new MigrationJobRunner( log, job, pipelineConfig, control, tokenRefreshManager, From fe1118a31ce7e5f24944635b85ff482732310f3e Mon Sep 17 00:00:00 2001 From: Nitesh Vijay Date: Tue, 18 Aug 2026 18:19:58 +0530 Subject: [PATCH 11/32] refactor: create initial session in provider Resolve only the initial source credential in MigrationJobRunner and let RotatingSessionProvider create and own both initial and refreshed sessions through its session factory. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 078ac7b9-aa7f-44a1-9053-44c346cec5e1 --- .../CassandraDriver/CassandraClientFactory.cs | 28 +++++++++++-------- .../RotatingSessionProvider.cs | 25 ++++++++++++----- .../DataTransfer/MigrationJobRunner.cs | 7 +++-- 3 files changed, 39 insertions(+), 21 deletions(-) diff --git a/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs b/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs index 1c2ec24..4958f19 100644 --- a/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs +++ b/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs @@ -305,29 +305,33 @@ private static ISession ConnectCluster(Builder builder) public static ISession CreateSourceSession( MigrationLog MigrationLog, Job job, TokenRefreshManager? tokenRefreshManager = null) + { + string credential = ResolveSourceCredential(job, tokenRefreshManager); + + return CreateSourceSessionWithCredential( + MigrationLog, job, credential, tokenRefreshManager); + } + + internal static string ResolveSourceCredential( + Job job, + TokenRefreshManager? tokenRefreshManager = null) { 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) + string credential = job.SourcePassword ?? string.Empty; + if (string.IsNullOrEmpty(credential) || job.SourceUseAad) { - password = tokenRefreshManager?.GetFreshAadToken() + credential = 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 and leak the - // bearer JWT to the browser DOM. Azure.Identity caches - // tokens in-process so re-acquiring per call is free. + // would echo it into an and leak the + // bearer JWT to the browser DOM. job.SourceUseAad = true; } - - return CreateSourceSessionWithCredential( - MigrationLog, job, password, tokenRefreshManager); + return credential; } internal static ISession CreateSourceSessionWithCredential( diff --git a/CassandraMigrationProcessor/CassandraDriver/RotatingSessionProvider.cs b/CassandraMigrationProcessor/CassandraDriver/RotatingSessionProvider.cs index 651e9b8..a46d088 100644 --- a/CassandraMigrationProcessor/CassandraDriver/RotatingSessionProvider.cs +++ b/CassandraMigrationProcessor/CassandraDriver/RotatingSessionProvider.cs @@ -40,16 +40,27 @@ public ISession GetSession() } } - public void Initialize(ISession session) + public ISession Initialize(string credential) { - ArgumentNullException.ThrowIfNull(session); + ArgumentException.ThrowIfNullOrWhiteSpace(credential); - lock (_sync) + var session = _sessionFactory(credential); + try { - ObjectDisposedException.ThrowIf(_disposed, this); - if (_currentSession != null) - throw new InvalidOperationException("The session provider is already initialized."); - _currentSession = session; + lock (_sync) + { + ObjectDisposedException.ThrowIf(_disposed, this); + if (_currentSession != null) + throw new InvalidOperationException("The session provider is already initialized."); + _currentSession = session; + } + return session; + } + catch + { + MigrationUtilities.SafeDisposeSession( + session, "Unpublished initial session"); + throw; } } diff --git a/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs b/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs index fae15af..3a47434 100644 --- a/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs +++ b/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs @@ -93,8 +93,11 @@ public static async Task CreateAsync( ISession? target = null; try { - source = CassandraClientFactory.CreateSourceSession(log, job, tokenRefreshManager); - sourceSessions.Initialize(source); + string sourceCredential = CassandraClientFactory.ResolveSourceCredential( + job, tokenRefreshManager); + source = sourceSessions.Initialize(sourceCredential); + if (TokenRefreshManager.IsLikelyAadToken(sourceCredential)) + tokenRefreshManager.StartTokenRefreshTimer(sourceCredential); target = await CassandraClientFactory.CreateTargetSessionAsync(log, job); return new MigrationJobRunner( log, job, pipelineConfig, control, tokenRefreshManager, From 87d700048cee7bf76dcd46b62126e9477475d63d Mon Sep 17 00:00:00 2001 From: Nitesh Vijay Date: Tue, 18 Aug 2026 18:59:05 +0530 Subject: [PATCH 12/32] refactor: capture immutable source settings Resolve source endpoint, username, port, and pool sizing once into SourceSessionSettings. The rotating provider factory now captures those settings and receives only the changing credential. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 078ac7b9-aa7f-44a1-9053-44c346cec5e1 --- .../CassandraDriver/CassandraClientFactory.cs | 47 ++++++++++++++----- .../DataTransfer/MigrationJobRunner.cs | 15 +++--- 2 files changed, 43 insertions(+), 19 deletions(-) diff --git a/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs b/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs index 4958f19..c504f2e 100644 --- a/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs +++ b/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs @@ -4,6 +4,13 @@ using CassandraMigrationProcessor.Infrastructure; using CassandraMigrationProcessor.Models; namespace CassandraMigrationProcessor.CassandraDriver; + +internal sealed record SourceSessionSettings( + string ContactPoint, + int Port, + string Username, + int MaxConnectionsPerHost); + /// /// Creates Cassandra ISession instances for source (Cosmos DB) /// and target (OSS Cassandra) clusters. @@ -307,9 +314,10 @@ public static ISession CreateSourceSession( TokenRefreshManager? tokenRefreshManager = null) { string credential = ResolveSourceCredential(job, tokenRefreshManager); + var settings = ResolveSourceSessionSettings(job); return CreateSourceSessionWithCredential( - MigrationLog, job, credential, tokenRefreshManager); + MigrationLog, settings, credential, tokenRefreshManager); } internal static string ResolveSourceCredential( @@ -336,38 +344,51 @@ internal static string ResolveSourceCredential( internal static ISession CreateSourceSessionWithCredential( MigrationLog migrationLog, - Job job, + SourceSessionSettings settings, string credential) { return CreateSourceSessionWithCredential( - migrationLog, job, credential, tokenRefreshManager: null); + migrationLog, settings, credential, tokenRefreshManager: null); } private static ISession CreateSourceSessionWithCredential( MigrationLog migrationLog, - Job job, + SourceSessionSettings settings, string credential, TokenRefreshManager? tokenRefreshManager) { - // For AAD auth, derive username from hostname if not explicitly - // provided (account name = first segment of the contact point FQDN). + return CreateSourceSession( + migrationLog, + settings.ContactPoint, + settings.Port, + settings.Username, + credential, + tokenRefreshManager, + settings.MaxConnectionsPerHost); + } + + internal static SourceSessionSettings ResolveSourceSessionSettings(Job job) + { + if (string.IsNullOrEmpty(job.SourceContactPoint)) + throw new ArgumentException("Source contact point is required", nameof(job)); + + bool useAad = job.SourceUseAad + || string.IsNullOrEmpty(job.SourcePassword); 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, + return new SourceSessionSettings( job.SourceContactPoint, job.SourcePort, username, - credential, - tokenRefreshManager, - maxConnectionsPerHost: ResolveMaxConnectionsPerHost(job.SourceMaxConnectionsPerHost, job.MaxConnectionsPerHost)); + ResolveMaxConnectionsPerHost( + job.SourceMaxConnectionsPerHost, + job.MaxConnectionsPerHost)); } /// diff --git a/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs b/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs index 3a47434..1ff79ba 100644 --- a/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs +++ b/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs @@ -85,14 +85,17 @@ public static async Task CreateAsync( ArgumentNullException.ThrowIfNull(control); var pipelineConfig = PipelineConfig.Resolve(job, config); - var sourceSessions = new RotatingSessionProvider( - credential => CassandraClientFactory.CreateSourceSessionWithCredential( - log, job, credential)); - var tokenRefreshManager = new TokenRefreshManager(log, sourceSessions); + RotatingSessionProvider? sourceSessions = null; + TokenRefreshManager? tokenRefreshManager = null; ISession? source = null; ISession? target = null; try { + var sourceSettings = CassandraClientFactory.ResolveSourceSessionSettings(job); + sourceSessions = new RotatingSessionProvider( + credential => CassandraClientFactory.CreateSourceSessionWithCredential( + log, sourceSettings, credential)); + tokenRefreshManager = new TokenRefreshManager(log, sourceSessions); string sourceCredential = CassandraClientFactory.ResolveSourceCredential( job, tokenRefreshManager); source = sourceSessions.Initialize(sourceCredential); @@ -106,8 +109,8 @@ public static async Task CreateAsync( catch { MigrationUtilities.SafeDisposeSession(target, "MigrationJobRunner target (CreateAsync rollback)"); - tokenRefreshManager.Dispose(); - sourceSessions.Dispose(); + tokenRefreshManager?.Dispose(); + sourceSessions?.Dispose(); throw; } } From 7aade0e242b71ac8bb1940ed77ebf25af0696c71 Mon Sep 17 00:00:00 2001 From: Nitesh Vijay Date: Tue, 18 Aug 2026 19:05:02 +0530 Subject: [PATCH 13/32] refactor: use credential session factory object Replace the session-creation delegate with ICredentialSessionFactory and a concrete SourceSessionFactory that owns immutable source settings. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 078ac7b9-aa7f-44a1-9053-44c346cec5e1 --- .../CassandraDriver/CassandraClientFactory.cs | 6 ---- .../RotatingSessionProvider.cs | 13 ++++++--- .../CassandraDriver/SourceSessionFactory.cs | 28 +++++++++++++++++++ .../DataTransfer/MigrationJobRunner.cs | 3 +- 4 files changed, 38 insertions(+), 12 deletions(-) create mode 100644 CassandraMigrationProcessor/CassandraDriver/SourceSessionFactory.cs diff --git a/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs b/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs index c504f2e..4a0f8c8 100644 --- a/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs +++ b/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs @@ -5,12 +5,6 @@ using CassandraMigrationProcessor.Models; namespace CassandraMigrationProcessor.CassandraDriver; -internal sealed record SourceSessionSettings( - string ContactPoint, - int Port, - string Username, - int MaxConnectionsPerHost); - /// /// Creates Cassandra ISession instances for source (Cosmos DB) /// and target (OSS Cassandra) clusters. diff --git a/CassandraMigrationProcessor/CassandraDriver/RotatingSessionProvider.cs b/CassandraMigrationProcessor/CassandraDriver/RotatingSessionProvider.cs index a46d088..952071f 100644 --- a/CassandraMigrationProcessor/CassandraDriver/RotatingSessionProvider.cs +++ b/CassandraMigrationProcessor/CassandraDriver/RotatingSessionProvider.cs @@ -12,19 +12,24 @@ public interface ISessionProvider ISession GetSession(); } +public interface ICredentialSessionFactory +{ + ISession CreateSession(string credential); +} + public sealed class RotatingSessionProvider : ISessionProvider, IDisposable { private static readonly TimeSpan RetiredSessionDisposalDelay = TimeSpan.FromMinutes(10); private readonly object _sync = new(); - private readonly Func _sessionFactory; + private readonly ICredentialSessionFactory _sessionFactory; private readonly HashSet _retiredSessions = new(ReferenceEqualityComparer.Instance); private ISession? _currentSession; private bool _disposed; - public RotatingSessionProvider(Func sessionFactory) + public RotatingSessionProvider(ICredentialSessionFactory sessionFactory) { _sessionFactory = sessionFactory ?? throw new ArgumentNullException(nameof(sessionFactory)); @@ -44,7 +49,7 @@ public ISession Initialize(string credential) { ArgumentException.ThrowIfNullOrWhiteSpace(credential); - var session = _sessionFactory(credential); + var session = _sessionFactory.CreateSession(credential); try { lock (_sync) @@ -68,7 +73,7 @@ public void Refresh(string credential) { ArgumentException.ThrowIfNullOrWhiteSpace(credential); - var session = _sessionFactory(credential); + var session = _sessionFactory.CreateSession(credential); ISession? retiredSession; try { diff --git a/CassandraMigrationProcessor/CassandraDriver/SourceSessionFactory.cs b/CassandraMigrationProcessor/CassandraDriver/SourceSessionFactory.cs new file mode 100644 index 0000000..c7868de --- /dev/null +++ b/CassandraMigrationProcessor/CassandraDriver/SourceSessionFactory.cs @@ -0,0 +1,28 @@ +using Cassandra; +using CassandraMigrationProcessor.Infrastructure; + +namespace CassandraMigrationProcessor.CassandraDriver; + +internal sealed record SourceSessionSettings( + string ContactPoint, + int Port, + string Username, + int MaxConnectionsPerHost); + +internal sealed class SourceSessionFactory : ICredentialSessionFactory +{ + private readonly MigrationLog _log; + private readonly SourceSessionSettings _settings; + + public SourceSessionFactory( + MigrationLog log, + SourceSessionSettings settings) + { + _log = log; + _settings = settings; + } + + public ISession CreateSession(string credential) + => CassandraClientFactory.CreateSourceSessionWithCredential( + _log, _settings, credential); +} diff --git a/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs b/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs index 1ff79ba..512e94e 100644 --- a/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs +++ b/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs @@ -93,8 +93,7 @@ public static async Task CreateAsync( { var sourceSettings = CassandraClientFactory.ResolveSourceSessionSettings(job); sourceSessions = new RotatingSessionProvider( - credential => CassandraClientFactory.CreateSourceSessionWithCredential( - log, sourceSettings, credential)); + new SourceSessionFactory(log, sourceSettings)); tokenRefreshManager = new TokenRefreshManager(log, sourceSessions); string sourceCredential = CassandraClientFactory.ResolveSourceCredential( job, tokenRefreshManager); From d5d95a93def68084a7293d5ecc3f738a3a522b16 Mon Sep 17 00:00:00 2001 From: Nitesh Vijay Date: Wed, 19 Aug 2026 12:37:17 +0530 Subject: [PATCH 14/32] perf: share source UDT registration cache Move source UDT registration from each PageReader into one job-wide cache keyed by session and keyspace. Use Lazy to guarantee concurrent workers execute registration once. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 078ac7b9-aa7f-44a1-9053-44c346cec5e1 --- .../DataTransfer/DataCopyWorker.cs | 7 ++- .../DataTransfer/JobPipeline.cs | 1 + .../DataTransfer/PageReader.cs | 40 +++++++-------- .../DataTransfer/PipelineContext.cs | 1 + .../SourceUdtRegistrationCache.cs | 51 +++++++++++++++++++ 5 files changed, 79 insertions(+), 21 deletions(-) create mode 100644 CassandraMigrationProcessor/DataTransfer/SourceUdtRegistrationCache.cs diff --git a/CassandraMigrationProcessor/DataTransfer/DataCopyWorker.cs b/CassandraMigrationProcessor/DataTransfer/DataCopyWorker.cs index ccdeded..cdd4dd6 100644 --- a/CassandraMigrationProcessor/DataTransfer/DataCopyWorker.cs +++ b/CassandraMigrationProcessor/DataTransfer/DataCopyWorker.cs @@ -33,7 +33,12 @@ public async Task RunAsync(PipelineContext ctx) Partition? current = null; try { - reader = await PageReader.CreateAsync(_workerLog, ctx.SourceSessionProvider, ctx.ReaderConfig, _ct); + reader = await PageReader.CreateAsync( + _workerLog, + ctx.SourceSessionProvider, + ctx.SourceUdtRegistrations, + ctx.ReaderConfig, + _ct); writer = await PageWriter.CreateAsync(_workerLog, ctx.SessionFactory, ctx.WriterConfig, _ct); while (!_ct.IsCancellationRequested diff --git a/CassandraMigrationProcessor/DataTransfer/JobPipeline.cs b/CassandraMigrationProcessor/DataTransfer/JobPipeline.cs index 2b84dde..bfd4900 100644 --- a/CassandraMigrationProcessor/DataTransfer/JobPipeline.cs +++ b/CassandraMigrationProcessor/DataTransfer/JobPipeline.cs @@ -49,6 +49,7 @@ public JobPipeline(MigrationLog log, Job job, PipelineConfig pipelineConfig, Context = new PipelineContext( _partitions, sourceSessionProvider, + new SourceUdtRegistrationCache(), sessionFactory, readerConfig, writerConfig, diff --git a/CassandraMigrationProcessor/DataTransfer/PageReader.cs b/CassandraMigrationProcessor/DataTransfer/PageReader.cs index 6d1435e..07a9e17 100644 --- a/CassandraMigrationProcessor/DataTransfer/PageReader.cs +++ b/CassandraMigrationProcessor/DataTransfer/PageReader.cs @@ -2,7 +2,6 @@ using CassandraMigrationProcessor.Infrastructure; using CassandraMigrationProcessor.CassandraDriver; using CassandraMigrationProcessor.Models; -using System.Collections.Concurrent; using System.Diagnostics; namespace CassandraMigrationProcessor.DataTransfer; @@ -30,11 +29,11 @@ internal class PageReader private readonly WorkerLog _log; private readonly CancellationToken _ct; private readonly ISessionProvider _sourceSessionProvider; + private readonly SourceUdtRegistrationCache _sourceUdtRegistrations; private readonly int _pageSize; private readonly int _maxReadRetries; private readonly bool _preserveCellTtl; private readonly bool _useJsonCopy; - private readonly ConcurrentDictionary<(ISession Session, string Keyspace), Task> _udtRegistrations = new(); /// /// Most recent transient exception observed during retry-exhausted @@ -50,7 +49,12 @@ internal class PageReader // hints parking a worker for minutes. private const int MaxRetryDelayMs = 30_000; - private PageReader(WorkerLog log, ISessionProvider sourceSessionProvider, ReaderConfig config, CancellationToken cancellationToken) + private PageReader( + WorkerLog log, + ISessionProvider sourceSessionProvider, + SourceUdtRegistrationCache sourceUdtRegistrations, + ReaderConfig config, + CancellationToken cancellationToken) { _log = log; _ct = cancellationToken; @@ -60,13 +64,22 @@ private PageReader(WorkerLog log, ISessionProvider sourceSessionProvider, Reader _useJsonCopy = config.UseJsonCopy; _sourceSessionProvider = sourceSessionProvider ?? throw new ArgumentNullException(nameof(sourceSessionProvider)); + _sourceUdtRegistrations = sourceUdtRegistrations + ?? throw new ArgumentNullException(nameof(sourceUdtRegistrations)); } public static Task CreateAsync(WorkerLog log, - ISessionProvider sourceSessionProvider, ReaderConfig config, + ISessionProvider sourceSessionProvider, + SourceUdtRegistrationCache sourceUdtRegistrations, + ReaderConfig config, CancellationToken cancellationToken) { - return Task.FromResult(new PageReader(log, sourceSessionProvider, config, cancellationToken)); + return Task.FromResult(new PageReader( + log, + sourceSessionProvider, + sourceUdtRegistrations, + config, + cancellationToken)); } /// @@ -82,21 +95,8 @@ private async Task EnsureUdtsRegisteredAsync(Partition partition) var sourceSession = _sourceSessionProvider.GetSession(); var keyspace = partition.Table.Spec.KeyspaceName; - await _udtRegistrations.GetOrAdd((sourceSession, keyspace), async key => - { - try - { - var allUdts = await SchemaManager.GetUserDefinedTypesAsync(key.Session, key.Keyspace); - await DynamicUdtRegistrar.RegisterAsync(key.Session, key.Keyspace, allUdts); - } - catch (Exception ex) - { - // Do NOT swallow: UDT mapping is required for correct - // row decoding. Surface as fatal. - _log.WriteLine($"FATAL: UDT mapping registration on source failed for {key.Keyspace}: {ex.Message}", LogType.Error); - throw; - } - }).ConfigureAwait(false); + await _sourceUdtRegistrations.EnsureRegisteredAsync( + sourceSession, keyspace, _log).ConfigureAwait(false); } /// diff --git a/CassandraMigrationProcessor/DataTransfer/PipelineContext.cs b/CassandraMigrationProcessor/DataTransfer/PipelineContext.cs index ad2d63f..75ef9c3 100644 --- a/CassandraMigrationProcessor/DataTransfer/PipelineContext.cs +++ b/CassandraMigrationProcessor/DataTransfer/PipelineContext.cs @@ -15,6 +15,7 @@ namespace CassandraMigrationProcessor.DataTransfer; internal record PipelineContext( PartitionManager Partitions, ISessionProvider SourceSessionProvider, + SourceUdtRegistrationCache SourceUdtRegistrations, ISessionFactory SessionFactory, ReaderConfig ReaderConfig, WriterConfig WriterConfig, diff --git a/CassandraMigrationProcessor/DataTransfer/SourceUdtRegistrationCache.cs b/CassandraMigrationProcessor/DataTransfer/SourceUdtRegistrationCache.cs new file mode 100644 index 0000000..eee3cc1 --- /dev/null +++ b/CassandraMigrationProcessor/DataTransfer/SourceUdtRegistrationCache.cs @@ -0,0 +1,51 @@ +using Cassandra; +using CassandraMigrationProcessor.CassandraDriver; +using CassandraMigrationProcessor.Models; +using System.Collections.Concurrent; + +namespace CassandraMigrationProcessor.DataTransfer; + +/// +/// Coordinates source UDT registration across every reader in a job. Each +/// session/keyspace pair is registered once; refreshed sessions get their own +/// registration entry. +/// +internal sealed class SourceUdtRegistrationCache +{ + private readonly ConcurrentDictionary<(ISession Session, string Keyspace), Lazy> + _registrations = new(); + + public Task EnsureRegisteredAsync( + ISession session, + string keyspace, + WorkerLog log) + { + return _registrations.GetOrAdd( + (session, keyspace), + key => new Lazy( + () => RegisterAsync(key.Session, key.Keyspace, log), + LazyThreadSafetyMode.ExecutionAndPublication)) + .Value; + } + + private static async Task RegisterAsync( + ISession session, + string keyspace, + WorkerLog log) + { + try + { + var allUdts = await SchemaManager.GetUserDefinedTypesAsync( + session, keyspace); + await DynamicUdtRegistrar.RegisterAsync( + session, keyspace, allUdts); + } + catch (Exception ex) + { + log.WriteLine( + $"FATAL: UDT mapping registration on source failed for {keyspace}: {ex.Message}", + LogType.Error); + throw; + } + } +} From f03fb5ce64cf9cb5094bc7a9ed31062aeb93d479 Mon Sep 17 00:00:00 2001 From: Nitesh Vijay Date: Wed, 19 Aug 2026 12:51:04 +0530 Subject: [PATCH 15/32] refactor: consolidate source session wrapper Merge source-session rotation and job-wide UDT registration into SourceSessionWrapper. Each read retry resolves a session and registers UDT mappings on that exact session before query execution. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 078ac7b9-aa7f-44a1-9053-44c346cec5e1 --- .../CassandraDriver/SourceSessionFactory.cs | 5 ++ ...ionProvider.cs => SourceSessionWrapper.cs} | 54 +++++++++++++------ .../CassandraDriver/TokenRefreshManager.cs | 4 +- .../DataTransfer/DataCopyWorker.cs | 3 +- .../DataTransfer/JobPipeline.cs | 5 +- .../DataTransfer/MigrationJobRunner.cs | 8 +-- .../DataTransfer/PageReader.cs | 51 ++++++------------ .../DataTransfer/PipelineContext.cs | 3 +- .../SourceUdtRegistrationCache.cs | 51 ------------------ 9 files changed, 68 insertions(+), 116 deletions(-) rename CassandraMigrationProcessor/CassandraDriver/{RotatingSessionProvider.cs => SourceSessionWrapper.cs} (71%) delete mode 100644 CassandraMigrationProcessor/DataTransfer/SourceUdtRegistrationCache.cs diff --git a/CassandraMigrationProcessor/CassandraDriver/SourceSessionFactory.cs b/CassandraMigrationProcessor/CassandraDriver/SourceSessionFactory.cs index c7868de..d8fb0e7 100644 --- a/CassandraMigrationProcessor/CassandraDriver/SourceSessionFactory.cs +++ b/CassandraMigrationProcessor/CassandraDriver/SourceSessionFactory.cs @@ -9,6 +9,11 @@ internal sealed record SourceSessionSettings( string Username, int MaxConnectionsPerHost); +public interface ICredentialSessionFactory +{ + ISession CreateSession(string credential); +} + internal sealed class SourceSessionFactory : ICredentialSessionFactory { private readonly MigrationLog _log; diff --git a/CassandraMigrationProcessor/CassandraDriver/RotatingSessionProvider.cs b/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs similarity index 71% rename from CassandraMigrationProcessor/CassandraDriver/RotatingSessionProvider.cs rename to CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs index 952071f..bf6caeb 100644 --- a/CassandraMigrationProcessor/CassandraDriver/RotatingSessionProvider.cs +++ b/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs @@ -1,23 +1,15 @@ using Cassandra; using CassandraMigrationProcessor.Infrastructure; +using System.Collections.Concurrent; namespace CassandraMigrationProcessor.CassandraDriver; /// -/// Resolves the current shared session while retaining rotated sessions for a -/// bounded grace period so in-flight operations can complete. +/// Owns the shared source-session lifecycle and session-scoped UDT mappings. +/// Rotated sessions remain available for a bounded grace period so in-flight +/// operations can complete. /// -public interface ISessionProvider -{ - ISession GetSession(); -} - -public interface ICredentialSessionFactory -{ - ISession CreateSession(string credential); -} - -public sealed class RotatingSessionProvider : ISessionProvider, IDisposable +public sealed class SourceSessionWrapper : IDisposable { private static readonly TimeSpan RetiredSessionDisposalDelay = TimeSpan.FromMinutes(10); @@ -26,16 +18,36 @@ public sealed class RotatingSessionProvider : ISessionProvider, IDisposable private readonly ICredentialSessionFactory _sessionFactory; private readonly HashSet _retiredSessions = new(ReferenceEqualityComparer.Instance); + private readonly ConcurrentDictionary<(ISession Session, string Keyspace), Lazy> + _udtRegistrations = new(); private ISession? _currentSession; private bool _disposed; - public RotatingSessionProvider(ICredentialSessionFactory sessionFactory) + public SourceSessionWrapper(ICredentialSessionFactory sessionFactory) { _sessionFactory = sessionFactory ?? throw new ArgumentNullException(nameof(sessionFactory)); } - public ISession GetSession() + public async Task GetSessionForReadAsync( + string keyspace, + bool registerUdts) + { + var session = GetSession(); + if (registerUdts) + { + await _udtRegistrations.GetOrAdd( + (session, keyspace), + key => new Lazy( + () => RegisterUdtsAsync(key.Session, key.Keyspace), + LazyThreadSafetyMode.ExecutionAndPublication)) + .Value + .ConfigureAwait(false); + } + return session; + } + + private ISession GetSession() { lock (_sync) { @@ -45,6 +57,16 @@ public ISession GetSession() } } + private static async Task RegisterUdtsAsync( + ISession session, + string keyspace) + { + var allUdts = await SchemaManager.GetUserDefinedTypesAsync( + session, keyspace); + await DynamicUdtRegistrar.RegisterAsync( + session, keyspace, allUdts); + } + public ISession Initialize(string credential) { ArgumentException.ThrowIfNullOrWhiteSpace(credential); @@ -132,7 +154,7 @@ public void Dispose() foreach (var session in sessionsToDispose) { MigrationUtilities.SafeDisposeSession( - session, "Rotating session provider"); + session, "Source session wrapper"); } } } diff --git a/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs b/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs index 9076550..5fa7eb1 100644 --- a/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs +++ b/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs @@ -11,7 +11,7 @@ public class TokenRefreshManager : IDisposable { private Timer? _tokenRefreshTimer; private readonly object _refreshLock = new(); - private readonly RotatingSessionProvider _sourceSessions; + private readonly SourceSessionWrapper _sourceSessions; private readonly MigrationLog _log; private DateTime _tokenExpiresAt = DateTime.MinValue; private int _consecutiveRefreshFailures; @@ -19,7 +19,7 @@ public class TokenRefreshManager : IDisposable public TokenRefreshManager( MigrationLog log, - RotatingSessionProvider sourceSessions) + SourceSessionWrapper sourceSessions) { _log = log; _sourceSessions = sourceSessions diff --git a/CassandraMigrationProcessor/DataTransfer/DataCopyWorker.cs b/CassandraMigrationProcessor/DataTransfer/DataCopyWorker.cs index cdd4dd6..708fe8c 100644 --- a/CassandraMigrationProcessor/DataTransfer/DataCopyWorker.cs +++ b/CassandraMigrationProcessor/DataTransfer/DataCopyWorker.cs @@ -35,8 +35,7 @@ public async Task RunAsync(PipelineContext ctx) { reader = await PageReader.CreateAsync( _workerLog, - ctx.SourceSessionProvider, - ctx.SourceUdtRegistrations, + ctx.SourceSession, ctx.ReaderConfig, _ct); writer = await PageWriter.CreateAsync(_workerLog, ctx.SessionFactory, ctx.WriterConfig, _ct); diff --git a/CassandraMigrationProcessor/DataTransfer/JobPipeline.cs b/CassandraMigrationProcessor/DataTransfer/JobPipeline.cs index bfd4900..f90aa26 100644 --- a/CassandraMigrationProcessor/DataTransfer/JobPipeline.cs +++ b/CassandraMigrationProcessor/DataTransfer/JobPipeline.cs @@ -22,7 +22,7 @@ internal sealed class JobPipeline : IDisposable, IAsyncDisposable public PipelineContext Context { get; } public JobPipeline(MigrationLog log, Job job, PipelineConfig pipelineConfig, - JobPartitioning partitioning, ISessionProvider sourceSessionProvider, + JobPartitioning partitioning, SourceSessionWrapper sourceSession, ISessionFactory sessionFactory, JobControl control) { @@ -48,8 +48,7 @@ public JobPipeline(MigrationLog log, Job job, PipelineConfig pipelineConfig, Context = new PipelineContext( _partitions, - sourceSessionProvider, - new SourceUdtRegistrationCache(), + sourceSession, sessionFactory, readerConfig, writerConfig, diff --git a/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs b/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs index 512e94e..4e6157b 100644 --- a/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs +++ b/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs @@ -20,7 +20,7 @@ public class MigrationJobRunner : IAsyncDisposable private readonly PipelineConfig _pipelineConfig; private readonly JobControl _control; private readonly TokenRefreshManager _tokenRefreshManager; - private readonly RotatingSessionProvider _sourceSessions; + private readonly SourceSessionWrapper _sourceSessions; private int _consecutiveAuthErrors; // Last auth exception observed by HandleMigrationUnitError; // attached as inner when the consecutive-auth threshold trips so @@ -56,7 +56,7 @@ private MigrationJobRunner( PipelineConfig pipelineConfig, JobControl control, TokenRefreshManager tokenRefreshManager, - RotatingSessionProvider sourceSessions, + SourceSessionWrapper sourceSessions, ISession sourceSession, ISession targetSession) { @@ -85,14 +85,14 @@ public static async Task CreateAsync( ArgumentNullException.ThrowIfNull(control); var pipelineConfig = PipelineConfig.Resolve(job, config); - RotatingSessionProvider? sourceSessions = null; + SourceSessionWrapper? sourceSessions = null; TokenRefreshManager? tokenRefreshManager = null; ISession? source = null; ISession? target = null; try { var sourceSettings = CassandraClientFactory.ResolveSourceSessionSettings(job); - sourceSessions = new RotatingSessionProvider( + sourceSessions = new SourceSessionWrapper( new SourceSessionFactory(log, sourceSettings)); tokenRefreshManager = new TokenRefreshManager(log, sourceSessions); string sourceCredential = CassandraClientFactory.ResolveSourceCredential( diff --git a/CassandraMigrationProcessor/DataTransfer/PageReader.cs b/CassandraMigrationProcessor/DataTransfer/PageReader.cs index 07a9e17..63c2166 100644 --- a/CassandraMigrationProcessor/DataTransfer/PageReader.cs +++ b/CassandraMigrationProcessor/DataTransfer/PageReader.cs @@ -28,8 +28,7 @@ internal class PageReader { private readonly WorkerLog _log; private readonly CancellationToken _ct; - private readonly ISessionProvider _sourceSessionProvider; - private readonly SourceUdtRegistrationCache _sourceUdtRegistrations; + private readonly SourceSessionWrapper _sourceSession; private readonly int _pageSize; private readonly int _maxReadRetries; private readonly bool _preserveCellTtl; @@ -51,8 +50,7 @@ internal class PageReader private PageReader( WorkerLog log, - ISessionProvider sourceSessionProvider, - SourceUdtRegistrationCache sourceUdtRegistrations, + SourceSessionWrapper sourceSession, ReaderConfig config, CancellationToken cancellationToken) { @@ -62,43 +60,22 @@ private PageReader( _maxReadRetries = config.MaxReadRetries; _preserveCellTtl = config.PreserveCellTtlAndWritetime; _useJsonCopy = config.UseJsonCopy; - _sourceSessionProvider = sourceSessionProvider - ?? throw new ArgumentNullException(nameof(sourceSessionProvider)); - _sourceUdtRegistrations = sourceUdtRegistrations - ?? throw new ArgumentNullException(nameof(sourceUdtRegistrations)); + _sourceSession = sourceSession + ?? throw new ArgumentNullException(nameof(sourceSession)); } public static Task CreateAsync(WorkerLog log, - ISessionProvider sourceSessionProvider, - SourceUdtRegistrationCache sourceUdtRegistrations, + SourceSessionWrapper sourceSession, ReaderConfig config, CancellationToken cancellationToken) { return Task.FromResult(new PageReader( log, - sourceSessionProvider, - sourceUdtRegistrations, + sourceSession, config, cancellationToken)); } - /// - /// Lazy, idempotent UDT registration for typed reads. The first typed - /// table registers every UDT in the keyspace because this reader can - /// subsequently process other tables that reference different UDTs. - /// - private async Task EnsureUdtsRegisteredAsync(Partition partition) - { - // JSON read path bypasses CLR-side UDT decoding entirely. - if (!partition.Table.IsCounterTable && _useJsonCopy) - return; - - var sourceSession = _sourceSessionProvider.GetSession(); - var keyspace = partition.Table.Spec.KeyspaceName; - await _sourceUdtRegistrations.EnsureRegisteredAsync( - sourceSession, keyspace, _log).ConfigureAwait(false); - } - /// /// One page of source rows together with the chunk and per-row /// CDC metadata (writetime + TTL expiry). Exactly one of @@ -137,8 +114,6 @@ internal record ReadResult( /// private async Task ReadJsonPageAsync(Partition partition) { - await EnsureUdtsRegisteredAsync(partition); - var stopwatch = Stopwatch.StartNew(); var (resultSet, elapsed) = await ExecutePageAsync(partition, useJson: true, stopwatch); if (resultSet == null) return null; @@ -176,8 +151,6 @@ internal record ReadResult( /// private async Task ReadTypedPageAsync(Partition partition) { - await EnsureUdtsRegisteredAsync(partition); - var stopwatch = Stopwatch.StartNew(); var (resultSet, elapsed) = await ExecutePageAsync(partition, useJson: false, stopwatch); if (resultSet == null) return null; @@ -218,9 +191,15 @@ internal record ReadResult( // intact and will retry the same page once the source stops // throttling. var resultSet = await RetryExecutor.ExecuteOrDefaultAsync( - operation: _ => _sourceSessionProvider.GetSession() - .ExecuteAsync(stmt) - .WaitAsync(_ct), + operation: async _ => + { + var sourceSession = await _sourceSession.GetSessionForReadAsync( + partition.Table.Spec.KeyspaceName, + registerUdts: !useJson).ConfigureAwait(false); + return await sourceSession.ExecuteAsync(stmt) + .WaitAsync(_ct) + .ConfigureAwait(false); + }, maxAttempts: _maxReadRetries, shouldRetry: ExceptionClassifier.IsTransient, delayFor: (ex, attempt) => TimeSpan.FromMilliseconds( diff --git a/CassandraMigrationProcessor/DataTransfer/PipelineContext.cs b/CassandraMigrationProcessor/DataTransfer/PipelineContext.cs index 75ef9c3..7fea3b9 100644 --- a/CassandraMigrationProcessor/DataTransfer/PipelineContext.cs +++ b/CassandraMigrationProcessor/DataTransfer/PipelineContext.cs @@ -14,8 +14,7 @@ namespace CassandraMigrationProcessor.DataTransfer; /// internal record PipelineContext( PartitionManager Partitions, - ISessionProvider SourceSessionProvider, - SourceUdtRegistrationCache SourceUdtRegistrations, + SourceSessionWrapper SourceSession, ISessionFactory SessionFactory, ReaderConfig ReaderConfig, WriterConfig WriterConfig, diff --git a/CassandraMigrationProcessor/DataTransfer/SourceUdtRegistrationCache.cs b/CassandraMigrationProcessor/DataTransfer/SourceUdtRegistrationCache.cs deleted file mode 100644 index eee3cc1..0000000 --- a/CassandraMigrationProcessor/DataTransfer/SourceUdtRegistrationCache.cs +++ /dev/null @@ -1,51 +0,0 @@ -using Cassandra; -using CassandraMigrationProcessor.CassandraDriver; -using CassandraMigrationProcessor.Models; -using System.Collections.Concurrent; - -namespace CassandraMigrationProcessor.DataTransfer; - -/// -/// Coordinates source UDT registration across every reader in a job. Each -/// session/keyspace pair is registered once; refreshed sessions get their own -/// registration entry. -/// -internal sealed class SourceUdtRegistrationCache -{ - private readonly ConcurrentDictionary<(ISession Session, string Keyspace), Lazy> - _registrations = new(); - - public Task EnsureRegisteredAsync( - ISession session, - string keyspace, - WorkerLog log) - { - return _registrations.GetOrAdd( - (session, keyspace), - key => new Lazy( - () => RegisterAsync(key.Session, key.Keyspace, log), - LazyThreadSafetyMode.ExecutionAndPublication)) - .Value; - } - - private static async Task RegisterAsync( - ISession session, - string keyspace, - WorkerLog log) - { - try - { - var allUdts = await SchemaManager.GetUserDefinedTypesAsync( - session, keyspace); - await DynamicUdtRegistrar.RegisterAsync( - session, keyspace, allUdts); - } - catch (Exception ex) - { - log.WriteLine( - $"FATAL: UDT mapping registration on source failed for {keyspace}: {ex.Message}", - LogType.Error); - throw; - } - } -} From 7df21ac55482cb0e62c54aac92ad27d128204c3c Mon Sep 17 00:00:00 2001 From: Nitesh Vijay Date: Wed, 19 Aug 2026 12:56:18 +0530 Subject: [PATCH 16/32] refactor: expose typed source session API Replace the registration flag with explicit GetSession and GetTypedSessionAsync methods. PageReader supplies a keyspace only for typed reads; UDT registration remains encapsulated by SourceSessionWrapper. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 078ac7b9-aa7f-44a1-9053-44c346cec5e1 --- .../CassandraDriver/SourceSessionWrapper.cs | 28 +++++++++---------- .../DataTransfer/PageReader.cs | 7 +++-- 2 files changed, 17 insertions(+), 18 deletions(-) diff --git a/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs b/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs index bf6caeb..33d0db9 100644 --- a/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs +++ b/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs @@ -29,25 +29,23 @@ public SourceSessionWrapper(ICredentialSessionFactory sessionFactory) ?? throw new ArgumentNullException(nameof(sessionFactory)); } - public async Task GetSessionForReadAsync( - string keyspace, - bool registerUdts) + public ISession GetSession() + => GetCurrentSession(); + + public async Task GetTypedSessionAsync(string keyspace) { - var session = GetSession(); - if (registerUdts) - { - await _udtRegistrations.GetOrAdd( - (session, keyspace), - key => new Lazy( - () => RegisterUdtsAsync(key.Session, key.Keyspace), - LazyThreadSafetyMode.ExecutionAndPublication)) - .Value - .ConfigureAwait(false); - } + var session = GetCurrentSession(); + await _udtRegistrations.GetOrAdd( + (session, keyspace), + key => new Lazy( + () => RegisterUdtsAsync(key.Session, key.Keyspace), + LazyThreadSafetyMode.ExecutionAndPublication)) + .Value + .ConfigureAwait(false); return session; } - private ISession GetSession() + private ISession GetCurrentSession() { lock (_sync) { diff --git a/CassandraMigrationProcessor/DataTransfer/PageReader.cs b/CassandraMigrationProcessor/DataTransfer/PageReader.cs index 63c2166..c310d26 100644 --- a/CassandraMigrationProcessor/DataTransfer/PageReader.cs +++ b/CassandraMigrationProcessor/DataTransfer/PageReader.cs @@ -193,9 +193,10 @@ internal record ReadResult( var resultSet = await RetryExecutor.ExecuteOrDefaultAsync( operation: async _ => { - var sourceSession = await _sourceSession.GetSessionForReadAsync( - partition.Table.Spec.KeyspaceName, - registerUdts: !useJson).ConfigureAwait(false); + var sourceSession = useJson + ? _sourceSession.GetSession() + : await _sourceSession.GetTypedSessionAsync( + partition.Table.Spec.KeyspaceName).ConfigureAwait(false); return await sourceSession.ExecuteAsync(stmt) .WaitAsync(_ct) .ConfigureAwait(false); From c04db77db0ce65cf6bc6591ac9a7ffddeafb0f9a Mon Sep 17 00:00:00 2001 From: Nitesh Vijay Date: Wed, 19 Aug 2026 13:44:58 +0530 Subject: [PATCH 17/32] fix: harden shared session lifecycle Avoid disposing the startup gate while workers may still release it, make UDT registration failures recoverable and diagnosable, resolve all runner source operations through the rotating wrapper, prune retired-session mappings, prevent token timer resurrection, and size the shared source pool from worker count. Keep the target session creation gate at 20. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 078ac7b9-aa7f-44a1-9053-44c346cec5e1 --- .../CassandraDriver/CassandraClientFactory.cs | 19 +++++++-- .../CassandraDriver/ISessionFactory.cs | 4 +- .../CassandraDriver/SourceSessionWrapper.cs | 40 ++++++++++++++++--- .../CassandraDriver/TokenRefreshManager.cs | 5 +++ .../DataTransfer/JobPipeline.cs | 4 -- .../DataTransfer/MigrationJobRunner.cs | 35 ++++++++-------- .../DataTransfer/PageReader.cs | 3 +- 7 files changed, 75 insertions(+), 35 deletions(-) diff --git a/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs b/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs index 4a0f8c8..d28a413 100644 --- a/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs +++ b/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs @@ -361,7 +361,9 @@ private static ISession CreateSourceSessionWithCredential( settings.MaxConnectionsPerHost); } - internal static SourceSessionSettings ResolveSourceSessionSettings(Job job) + internal static SourceSessionSettings ResolveSourceSessionSettings( + Job job, + int workerCount = 0) { if (string.IsNullOrEmpty(job.SourceContactPoint)) throw new ArgumentException("Source contact point is required", nameof(job)); @@ -376,13 +378,22 @@ internal static SourceSessionSettings ResolveSourceSessionSettings(Job job) .Split('.')[0]; } + 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, - ResolveMaxConnectionsPerHost( - job.SourceMaxConnectionsPerHost, - job.MaxConnectionsPerHost)); + maxConnectionsPerHost); } /// diff --git a/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs b/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs index 0b625e6..a2619c3 100644 --- a/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs +++ b/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs @@ -18,7 +18,7 @@ public interface ISessionFactory /// Limits simultaneous session opens. This prevents high-worker jobs from /// creating a connection storm during startup. /// -public sealed class GatedSessionFactory : ISessionFactory, IDisposable +public sealed class GatedSessionFactory : ISessionFactory { private const int MaxConcurrentSessionCreations = 20; @@ -45,8 +45,6 @@ public async Task CreateSessionAsync(CancellationToken cancellationTok _creationGate.Release(); } } - - public void Dispose() => _creationGate.Dispose(); } /// diff --git a/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs b/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs index 33d0db9..6d20cfb 100644 --- a/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs +++ b/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs @@ -4,6 +4,14 @@ namespace CassandraMigrationProcessor.CassandraDriver; +internal sealed class SourceUdtRegistrationException : Exception +{ + public SourceUdtRegistrationException(string keyspace, Exception innerException) + : base($"UDT mapping registration failed for source keyspace '{keyspace}'.", innerException) + { + } +} + /// /// Owns the shared source-session lifecycle and session-scoped UDT mappings. /// Rotated sessions remain available for a bounded grace period so in-flight @@ -35,13 +43,24 @@ public ISession GetSession() public async Task GetTypedSessionAsync(string keyspace) { var session = GetCurrentSession(); - await _udtRegistrations.GetOrAdd( - (session, keyspace), + var key = (Session: session, Keyspace: keyspace); + var registration = _udtRegistrations.GetOrAdd( + key, key => new Lazy( () => RegisterUdtsAsync(key.Session, key.Keyspace), - LazyThreadSafetyMode.ExecutionAndPublication)) - .Value - .ConfigureAwait(false); + LazyThreadSafetyMode.ExecutionAndPublication)); + try + { + await registration.Value.ConfigureAwait(false); + } + catch (Exception ex) + { + ((ICollection>>) + _udtRegistrations).Remove(new KeyValuePair< + (ISession Session, string Keyspace), Lazy>( + key, registration)); + throw new SourceUdtRegistrationException(keyspace, ex); + } return session; } @@ -130,11 +149,21 @@ private async Task DisposeRetiredSessionAfterDelayAsync(ISession session) if (shouldDispose) { + RemoveUdtRegistrations(session); MigrationUtilities.SafeDisposeSession( session, "Deferred rotated session"); } } + private void RemoveUdtRegistrations(ISession session) + { + foreach (var key in _udtRegistrations.Keys) + { + if (ReferenceEquals(key.Session, session)) + _udtRegistrations.TryRemove(key, out _); + } + } + public void Dispose() { List sessionsToDispose; @@ -147,6 +176,7 @@ public void Dispose() if (_currentSession != null) sessionsToDispose.Add(_currentSession); _currentSession = null; + _udtRegistrations.Clear(); } foreach (var session in sessionsToDispose) diff --git a/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs b/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs index 5fa7eb1..fa6d986 100644 --- a/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs +++ b/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs @@ -13,6 +13,7 @@ public class TokenRefreshManager : IDisposable private readonly object _refreshLock = new(); private readonly SourceSessionWrapper _sourceSessions; private readonly MigrationLog _log; + private bool _disposed; private DateTime _tokenExpiresAt = DateTime.MinValue; private int _consecutiveRefreshFailures; private const int MaxRefreshFailures = 6; @@ -98,6 +99,7 @@ public void StartTokenRefreshTimer( { lock (_refreshLock) { + if (_disposed) return; StopTokenRefreshTimer(); DateTime expiry = GetTokenExpiry(currentToken); @@ -134,6 +136,7 @@ private void TokenRefreshCallback(object? state) { lock (_refreshLock) { + if (_disposed) return; try { string freshToken = GetFreshAadToken(); @@ -171,6 +174,8 @@ public void Dispose() { lock (_refreshLock) { + if (_disposed) return; + _disposed = true; StopTokenRefreshTimer(); } } diff --git a/CassandraMigrationProcessor/DataTransfer/JobPipeline.cs b/CassandraMigrationProcessor/DataTransfer/JobPipeline.cs index f90aa26..d789b9e 100644 --- a/CassandraMigrationProcessor/DataTransfer/JobPipeline.cs +++ b/CassandraMigrationProcessor/DataTransfer/JobPipeline.cs @@ -18,7 +18,6 @@ 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, @@ -29,7 +28,6 @@ public JobPipeline(MigrationLog log, Job job, PipelineConfig pipelineConfig, _log = log; _pipelineConfig = pipelineConfig; _control = control; - _sessionFactory = sessionFactory; bool enableReplay = job.IsOnline; _partitions = new PartitionManager( @@ -101,7 +99,5 @@ 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"); } } diff --git a/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs b/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs index 4e6157b..66a1856 100644 --- a/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs +++ b/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs @@ -32,15 +32,13 @@ public class MigrationJobRunner : IAsyncDisposable private JobPipeline? _pipeline; /// - /// Runner-wide source / target sessions opened once in - /// and reused across wildcard expansion, - /// schema provisioning, and partition discovery. Disposed in - /// . Copy workers reuse the thread-safe source - /// session to avoid multiplying driver metadata topology/schema handshakes, - /// while retaining independent target sessions for write throughput. + /// Runner-wide target session opened once in . + /// Source operations resolve the current session through + /// so AAD rotation is honored throughout + /// wildcard expansion, schema provisioning, partition discovery, and copy. + /// Copy workers retain independent target sessions for write throughput. /// For simulated runs the target session is a . /// - private readonly ISession _sourceSession; private readonly ISession _targetSession; /// @@ -57,7 +55,6 @@ private MigrationJobRunner( JobControl control, TokenRefreshManager tokenRefreshManager, SourceSessionWrapper sourceSessions, - ISession sourceSession, ISession targetSession) { _log = log; @@ -66,7 +63,6 @@ private MigrationJobRunner( _control = control; _tokenRefreshManager = tokenRefreshManager; _sourceSessions = sourceSessions; - _sourceSession = sourceSession; _targetSession = targetSession; } @@ -87,23 +83,23 @@ public static async Task CreateAsync( var pipelineConfig = PipelineConfig.Resolve(job, config); SourceSessionWrapper? sourceSessions = null; TokenRefreshManager? tokenRefreshManager = null; - ISession? source = null; ISession? target = null; try { - var sourceSettings = CassandraClientFactory.ResolveSourceSessionSettings(job); + var sourceSettings = CassandraClientFactory.ResolveSourceSessionSettings( + job, pipelineConfig.WorkerCount); sourceSessions = new SourceSessionWrapper( new SourceSessionFactory(log, sourceSettings)); tokenRefreshManager = new TokenRefreshManager(log, sourceSessions); string sourceCredential = CassandraClientFactory.ResolveSourceCredential( job, tokenRefreshManager); - source = sourceSessions.Initialize(sourceCredential); + sourceSessions.Initialize(sourceCredential); if (TokenRefreshManager.IsLikelyAadToken(sourceCredential)) tokenRefreshManager.StartTokenRefreshTimer(sourceCredential); target = await CassandraClientFactory.CreateTargetSessionAsync(log, job); return new MigrationJobRunner( log, job, pipelineConfig, control, tokenRefreshManager, - sourceSessions, source, target); + sourceSessions, target); } catch { @@ -376,7 +372,7 @@ private async Task RunSchemaPhaseAsync( .Distinct(StringComparer.Ordinal) .ToList(); await SchemaManager.WarnAboutUnreplicatedSchemaAsync( - _sourceSession, inScopeKeyspaces, _log); + _sourceSessions.GetSession(), inScopeKeyspaces, _log); } catch (Exception ex) { @@ -506,7 +502,9 @@ await Parallel.ForEachAsync(units, options, async (mu, token) => _log.WriteLine($"[Partitioning] Discovering partitions for {mu.KeyspaceName}.{mu.TableName}", LogType.Info); try { - await DiscoverUnitPartitioningAsync(job, mu, _sourceSession, partitioner, chunks, collectLock); + await DiscoverUnitPartitioningAsync( + job, mu, _sourceSessions.GetSession(), + partitioner, chunks, collectLock); } catch (OperationCanceledException) { @@ -788,7 +786,7 @@ await _targetSession.ExecuteAsync(new SimpleStatement( } bool existed = await SchemaManager.TableExistsAsync(_targetSession, mu.KeyspaceName, mu.TableName); - await SchemaManager.SyncSchemaAsync(_sourceSession, _targetSession, + await SchemaManager.SyncSchemaAsync(_sourceSessions.GetSession(), _targetSession, mu.KeyspaceName, mu.TableName, mu.KeyspaceName, mu.TableName, _log); if (!existed) _log.WriteLine($"Created target table {mu.KeyspaceName}.{mu.TableName}", LogType.Info); @@ -935,11 +933,12 @@ void AddExpandedUnit(string keyspaceName, string tableName) => try { - var tables = await CassandraQueries.ListTablesAsync(_sourceSession, keyspace); + var sourceSession = _sourceSessions.GetSession(); + var tables = await CassandraQueries.ListTablesAsync(sourceSession, keyspace); foreach (var tableName in tables) { cancellationToken.ThrowIfCancellationRequested(); - if (await IsTableAccessibleAsync(_sourceSession, keyspace, tableName, cancellationToken)) + if (await IsTableAccessibleAsync(sourceSession, keyspace, tableName, cancellationToken)) { AddExpandedUnit(keyspace, tableName); } diff --git a/CassandraMigrationProcessor/DataTransfer/PageReader.cs b/CassandraMigrationProcessor/DataTransfer/PageReader.cs index c310d26..b6e78f5 100644 --- a/CassandraMigrationProcessor/DataTransfer/PageReader.cs +++ b/CassandraMigrationProcessor/DataTransfer/PageReader.cs @@ -202,7 +202,8 @@ internal record ReadResult( .ConfigureAwait(false); }, maxAttempts: _maxReadRetries, - shouldRetry: ExceptionClassifier.IsTransient, + shouldRetry: ex => ex is not SourceUdtRegistrationException + && ExceptionClassifier.IsTransient(ex), delayFor: (ex, attempt) => TimeSpan.FromMilliseconds( Math.Min(ExceptionClassifier.GetRetryDelayMs(ex, attempt), MaxRetryDelayMs)), onRetry: (ex, attempt) => From 9850625c4ac891ad739f0ea78cec37b5bde94e5e Mon Sep 17 00:00:00 2001 From: Nitesh Vijay Date: Wed, 19 Aug 2026 14:15:47 +0530 Subject: [PATCH 18/32] style: expand method implementations Replace expression-bodied method implementations introduced by this PR with explicit method bodies. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 078ac7b9-aa7f-44a1-9053-44c346cec5e1 --- .../CassandraDriver/SourceSessionFactory.cs | 4 +++- .../CassandraDriver/SourceSessionWrapper.cs | 4 +++- CassandraMigrationProcessor/DataTransfer/PageWriter.cs | 6 +++++- 3 files changed, 11 insertions(+), 3 deletions(-) diff --git a/CassandraMigrationProcessor/CassandraDriver/SourceSessionFactory.cs b/CassandraMigrationProcessor/CassandraDriver/SourceSessionFactory.cs index d8fb0e7..3aaecdf 100644 --- a/CassandraMigrationProcessor/CassandraDriver/SourceSessionFactory.cs +++ b/CassandraMigrationProcessor/CassandraDriver/SourceSessionFactory.cs @@ -28,6 +28,8 @@ public SourceSessionFactory( } public ISession CreateSession(string credential) - => CassandraClientFactory.CreateSourceSessionWithCredential( + { + return CassandraClientFactory.CreateSourceSessionWithCredential( _log, _settings, credential); + } } diff --git a/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs b/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs index 6d20cfb..6f2d152 100644 --- a/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs +++ b/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs @@ -38,7 +38,9 @@ public SourceSessionWrapper(ICredentialSessionFactory sessionFactory) } public ISession GetSession() - => GetCurrentSession(); + { + return GetCurrentSession(); + } public async Task GetTypedSessionAsync(string keyspace) { diff --git a/CassandraMigrationProcessor/DataTransfer/PageWriter.cs b/CassandraMigrationProcessor/DataTransfer/PageWriter.cs index 08a7ed0..756c5a8 100644 --- a/CassandraMigrationProcessor/DataTransfer/PageWriter.cs +++ b/CassandraMigrationProcessor/DataTransfer/PageWriter.cs @@ -65,7 +65,11 @@ public static async Task CreateAsync(WorkerLog log, ISessionFactory } public void Dispose() - => MigrationUtilities.SafeDisposeSession(_targetSession, "PageWriter target session"); + { + MigrationUtilities.SafeDisposeSession( + _targetSession, + "PageWriter target session"); + } private Task GetStrategyAsync(Partition partition) { From 2a6af55c0a9d1108fd3c6a8d646e659da3237e8c Mon Sep 17 00:00:00 2001 From: Nitesh Vijay Date: Wed, 19 Aug 2026 14:20:09 +0530 Subject: [PATCH 19/32] Encapsulate token refresh in source sessions Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 078ac7b9-aa7f-44a1-9053-44c346cec5e1 --- .../CassandraDriver/SourceSessionWrapper.cs | 18 ++++++++++++++++-- .../CassandraDriver/TokenRefreshManager.cs | 10 +++++----- .../DataTransfer/MigrationJobRunner.cs | 18 ++++-------------- 3 files changed, 25 insertions(+), 21 deletions(-) diff --git a/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs b/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs index 6f2d152..df6fef4 100644 --- a/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs +++ b/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs @@ -24,6 +24,7 @@ public sealed class SourceSessionWrapper : IDisposable private readonly object _sync = new(); private readonly ICredentialSessionFactory _sessionFactory; + private readonly TokenRefreshManager _tokenRefreshManager; private readonly HashSet _retiredSessions = new(ReferenceEqualityComparer.Instance); private readonly ConcurrentDictionary<(ISession Session, string Keyspace), Lazy> @@ -31,10 +32,13 @@ public sealed class SourceSessionWrapper : IDisposable private ISession? _currentSession; private bool _disposed; - public SourceSessionWrapper(ICredentialSessionFactory sessionFactory) + public SourceSessionWrapper( + MigrationLog log, + ICredentialSessionFactory sessionFactory) { _sessionFactory = sessionFactory ?? throw new ArgumentNullException(nameof(sessionFactory)); + _tokenRefreshManager = new TokenRefreshManager(log, Refresh); } public ISession GetSession() @@ -100,7 +104,6 @@ public ISession Initialize(string credential) throw new InvalidOperationException("The session provider is already initialized."); _currentSession = session; } - return session; } catch { @@ -108,6 +111,10 @@ public ISession Initialize(string credential) session, "Unpublished initial session"); throw; } + + if (TokenRefreshManager.IsLikelyAadToken(credential)) + _tokenRefreshManager.StartTokenRefreshTimer(credential); + return session; } public void Refresh(string credential) @@ -166,8 +173,15 @@ private void RemoveUdtRegistrations(ISession session) } } + public void StopTokenRefresh() + { + _tokenRefreshManager.StopTokenRefreshTimer(); + } + public void Dispose() { + _tokenRefreshManager.Dispose(); + List sessionsToDispose; lock (_sync) { diff --git a/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs b/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs index fa6d986..658e25c 100644 --- a/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs +++ b/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs @@ -11,7 +11,7 @@ public class TokenRefreshManager : IDisposable { private Timer? _tokenRefreshTimer; private readonly object _refreshLock = new(); - private readonly SourceSessionWrapper _sourceSessions; + private readonly Action _refreshSession; private readonly MigrationLog _log; private bool _disposed; private DateTime _tokenExpiresAt = DateTime.MinValue; @@ -20,11 +20,11 @@ public class TokenRefreshManager : IDisposable public TokenRefreshManager( MigrationLog log, - SourceSessionWrapper sourceSessions) + Action refreshSession) { _log = log; - _sourceSessions = sourceSessions - ?? throw new ArgumentNullException(nameof(sourceSessions)); + _refreshSession = refreshSession + ?? throw new ArgumentNullException(nameof(refreshSession)); } /// @@ -141,7 +141,7 @@ private void TokenRefreshCallback(object? state) { string freshToken = GetFreshAadToken(); - _sourceSessions.Refresh(freshToken); + _refreshSession(freshToken); // Schedule next refresh _consecutiveRefreshFailures = 0; diff --git a/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs b/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs index 66a1856..4e8b023 100644 --- a/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs +++ b/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs @@ -19,7 +19,6 @@ public class MigrationJobRunner : IAsyncDisposable private readonly Job _job; private readonly PipelineConfig _pipelineConfig; private readonly JobControl _control; - private readonly TokenRefreshManager _tokenRefreshManager; private readonly SourceSessionWrapper _sourceSessions; private int _consecutiveAuthErrors; // Last auth exception observed by HandleMigrationUnitError; @@ -53,7 +52,6 @@ private MigrationJobRunner( Job job, PipelineConfig pipelineConfig, JobControl control, - TokenRefreshManager tokenRefreshManager, SourceSessionWrapper sourceSessions, ISession targetSession) { @@ -61,7 +59,6 @@ private MigrationJobRunner( _job = job; _pipelineConfig = pipelineConfig; _control = control; - _tokenRefreshManager = tokenRefreshManager; _sourceSessions = sourceSessions; _targetSession = targetSession; } @@ -82,29 +79,24 @@ public static async Task CreateAsync( var pipelineConfig = PipelineConfig.Resolve(job, config); SourceSessionWrapper? sourceSessions = null; - TokenRefreshManager? tokenRefreshManager = null; ISession? target = null; try { var sourceSettings = CassandraClientFactory.ResolveSourceSessionSettings( job, pipelineConfig.WorkerCount); sourceSessions = new SourceSessionWrapper( + log, new SourceSessionFactory(log, sourceSettings)); - tokenRefreshManager = new TokenRefreshManager(log, sourceSessions); string sourceCredential = CassandraClientFactory.ResolveSourceCredential( - job, tokenRefreshManager); + job); sourceSessions.Initialize(sourceCredential); - if (TokenRefreshManager.IsLikelyAadToken(sourceCredential)) - tokenRefreshManager.StartTokenRefreshTimer(sourceCredential); target = await CassandraClientFactory.CreateTargetSessionAsync(log, job); return new MigrationJobRunner( - log, job, pipelineConfig, control, tokenRefreshManager, - sourceSessions, target); + log, job, pipelineConfig, control, sourceSessions, target); } catch { MigrationUtilities.SafeDisposeSession(target, "MigrationJobRunner target (CreateAsync rollback)"); - tokenRefreshManager?.Dispose(); sourceSessions?.Dispose(); throw; } @@ -257,11 +249,9 @@ public async Task StartAsync() /// public ValueTask DisposeAsync() { - _tokenRefreshManager.StopTokenRefreshTimer(); MigrationUtilities.SafeDispose(_pipeline, "JobPipeline (Dispose)"); _pipeline = null; MigrationUtilities.SafeDisposeSession(_targetSession, "MigrationJobRunner target session"); - _tokenRefreshManager.Dispose(); _sourceSessions.Dispose(); return ValueTask.CompletedTask; } @@ -873,7 +863,7 @@ public void Stop() // without waiting for the outer Task to observe the cancel. MigrationUtilities.SafeDispose(_pipeline, "JobPipeline (Stop)"); _pipeline = null; - _tokenRefreshManager.StopTokenRefreshTimer(); + _sourceSessions.StopTokenRefresh(); } /// From 6000e4777dfddae8801ed230c01c79c5344c0cbe Mon Sep 17 00:00:00 2001 From: Nitesh Vijay Date: Wed, 19 Aug 2026 14:33:13 +0530 Subject: [PATCH 20/32] Inline source token refresh lifecycle Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 078ac7b9-aa7f-44a1-9053-44c346cec5e1 --- .../CassandraDriver/CassandraClientFactory.cs | 58 ++---- .../CassandraDriver/SourceSessionWrapper.cs | 135 ++++++++++++- .../CassandraDriver/TokenRefreshManager.cs | 182 ------------------ 3 files changed, 144 insertions(+), 231 deletions(-) delete mode 100644 CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs diff --git a/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs b/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs index d28a413..9f540fa 100644 --- a/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs +++ b/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs @@ -8,8 +8,7 @@ namespace CassandraMigrationProcessor.CassandraDriver; /// /// 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. /// public static class CassandraClientFactory { @@ -36,8 +35,6 @@ public static class CassandraClientFactory /// /// 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. /// public static ISession CreateSourceSession( @@ -46,7 +43,6 @@ public static ISession CreateSourceSession( int port, string username, string password, - TokenRefreshManager? tokenRefreshManager = null, int maxConnectionsPerHost = 0) { // Source always uses SSL (Cosmos DB requires it) @@ -65,9 +61,7 @@ public static ISession CreateSourceSession( { try { - var session = ConnectCluster(builder); - RegisterAadTokenRefresh(password, tokenRefreshManager); - return session; + return ConnectCluster(builder); } catch (Exception ex) when ( ExceptionClassifier.IsTransient(ex) @@ -88,22 +82,6 @@ public static ISession CreateSourceSession( throw new UnreachableException(); } - /// - /// When looks like an AAD/JWT bearer - /// token and the caller wired up a , - /// start the proactive refresh timer so the bearer is rotated before it - /// expires. - /// No-op when the password is a static credential or the manager is - /// not supplied. - /// - private static void RegisterAadTokenRefresh( - string password, - TokenRefreshManager? tokenRefreshManager) - { - if (!TokenRefreshManager.IsLikelyAadToken(password)) return; - tokenRefreshManager?.StartTokenRefreshTimer(password); - } - /// /// Create a session to an OSS Apache Cassandra cluster. /// Tries SSL first, falls back to plain if SSL fails. @@ -304,19 +282,16 @@ private static ISession ConnectCluster(Builder builder) /// token automatically. /// public static ISession CreateSourceSession( - MigrationLog MigrationLog, Job job, - TokenRefreshManager? tokenRefreshManager = null) + MigrationLog MigrationLog, Job job) { - string credential = ResolveSourceCredential(job, tokenRefreshManager); + string credential = ResolveSourceCredential(job); var settings = ResolveSourceSessionSettings(job); return CreateSourceSessionWithCredential( - MigrationLog, settings, credential, tokenRefreshManager); + MigrationLog, settings, credential); } - internal static string ResolveSourceCredential( - Job job, - TokenRefreshManager? tokenRefreshManager = null) + internal static string ResolveSourceCredential(Job job) { if (string.IsNullOrEmpty(job.SourceContactPoint)) throw new ArgumentException("Source contact point is required", nameof(job)); @@ -324,8 +299,7 @@ internal static string ResolveSourceCredential( string credential = job.SourcePassword ?? string.Empty; if (string.IsNullOrEmpty(credential) || job.SourceUseAad) { - credential = tokenRefreshManager?.GetFreshAadToken() - ?? TokenRefreshManager.AcquireAadToken(); + credential = 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 @@ -336,20 +310,19 @@ internal static string ResolveSourceCredential( return credential; } - internal static ISession CreateSourceSessionWithCredential( - MigrationLog migrationLog, - SourceSessionSettings settings, - string credential) + internal static string AcquireAadToken() { - return CreateSourceSessionWithCredential( - migrationLog, settings, credential, tokenRefreshManager: null); + var credential = new Azure.Identity.DefaultAzureCredential(); + return credential.GetToken( + new Azure.Core.TokenRequestContext( + new[] { "https://cosmos.azure.com/.default" })) + .Token; } - private static ISession CreateSourceSessionWithCredential( + internal static ISession CreateSourceSessionWithCredential( MigrationLog migrationLog, SourceSessionSettings settings, - string credential, - TokenRefreshManager? tokenRefreshManager) + string credential) { return CreateSourceSession( migrationLog, @@ -357,7 +330,6 @@ private static ISession CreateSourceSessionWithCredential( settings.Port, settings.Username, credential, - tokenRefreshManager, settings.MaxConnectionsPerHost); } diff --git a/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs b/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs index df6fef4..bb4c23c 100644 --- a/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs +++ b/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs @@ -1,6 +1,8 @@ using Cassandra; using CassandraMigrationProcessor.Infrastructure; +using CassandraMigrationProcessor.Models; using System.Collections.Concurrent; +using System.IdentityModel.Tokens.Jwt; namespace CassandraMigrationProcessor.CassandraDriver; @@ -21,24 +23,31 @@ public sealed class SourceSessionWrapper : IDisposable { private static readonly TimeSpan RetiredSessionDisposalDelay = TimeSpan.FromMinutes(10); + private const int MaxRefreshFailures = 6; private readonly object _sync = new(); + private readonly object _refreshLock = new(); + private readonly MigrationLog _log; private readonly ICredentialSessionFactory _sessionFactory; - private readonly TokenRefreshManager _tokenRefreshManager; private readonly HashSet _retiredSessions = new(ReferenceEqualityComparer.Instance); private readonly ConcurrentDictionary<(ISession Session, string Keyspace), Lazy> _udtRegistrations = new(); private ISession? _currentSession; + private Timer? _tokenRefreshTimer; + private DateTime _tokenExpiresAt = DateTime.MinValue; + private int _consecutiveRefreshFailures; + private bool _tokenRefreshEnabled; + private bool _tokenRefreshDisposed; private bool _disposed; public SourceSessionWrapper( MigrationLog log, ICredentialSessionFactory sessionFactory) { + _log = log ?? throw new ArgumentNullException(nameof(log)); _sessionFactory = sessionFactory ?? throw new ArgumentNullException(nameof(sessionFactory)); - _tokenRefreshManager = new TokenRefreshManager(log, Refresh); } public ISession GetSession() @@ -112,8 +121,8 @@ public ISession Initialize(string credential) throw; } - if (TokenRefreshManager.IsLikelyAadToken(credential)) - _tokenRefreshManager.StartTokenRefreshTimer(credential); + if (IsLikelyAadToken(credential)) + StartTokenRefresh(credential); return session; } @@ -173,14 +182,128 @@ private void RemoveUdtRegistrations(ISession session) } } + private static bool IsLikelyAadToken(string? credential) + { + return credential != null && credential.Length > 200; + } + + private static DateTime GetTokenExpiry(string token) + { + try + { + var handler = new JwtSecurityTokenHandler(); + if (handler.CanReadToken(token)) + { + var jwt = handler.ReadJwtToken(token); + return jwt.ValidTo; + } + } + catch (Exception ex) + { + Console.WriteLine( + $"[Warning] Failed to read AAD token expiry: {ex.Message}"); + } + + return DateTime.MaxValue; + } + + private void StartTokenRefresh(string currentToken) + { + lock (_refreshLock) + { + if (_tokenRefreshDisposed) return; + _tokenRefreshEnabled = true; + ScheduleTokenRefresh(currentToken); + } + } + + private void ScheduleTokenRefresh(string currentToken) + { + _tokenRefreshTimer?.Dispose(); + + DateTime expiry = GetTokenExpiry(currentToken); + if (expiry == DateTime.MaxValue) + expiry = DateTime.UtcNow.AddMinutes(50); + + _tokenExpiresAt = expiry; + + TimeSpan delay = expiry - DateTime.UtcNow + - TimeSpan.FromMinutes(5); + if (delay < TimeSpan.FromMinutes(1)) + delay = TimeSpan.FromMinutes(1); + + _tokenRefreshTimer = new Timer( + RefreshTokenCallback, null, + delay, Timeout.InfiniteTimeSpan); + } + + private void RefreshTokenCallback(object? state) + { + lock (_refreshLock) + { + if (_tokenRefreshDisposed || !_tokenRefreshEnabled) return; + + try + { + string freshToken = CassandraClientFactory.AcquireAadToken(); + Refresh(freshToken); + + _consecutiveRefreshFailures = 0; + ScheduleTokenRefresh(freshToken); + } + catch (Exception ex) + { + _consecutiveRefreshFailures++; + int seconds = Math.Min( + 300, + 30 * (1 << Math.Min( + _consecutiveRefreshFailures - 1, 4))); + bool tokenAlreadyExpired = + DateTime.UtcNow >= _tokenExpiresAt; + LogType severity = + _consecutiveRefreshFailures >= MaxRefreshFailures + || tokenAlreadyExpired + ? LogType.Error + : LogType.Warning; + string message = + $"Token refresh failed (attempt {_consecutiveRefreshFailures}, " + + $"retrying in {seconds}s, tokenExpiresAt={_tokenExpiresAt:O}): " + + ex.Message; + Console.WriteLine($"[{severity}] {message}"); + _log.WriteLine(message, severity); + + _tokenRefreshTimer?.Dispose(); + if (_tokenRefreshEnabled && !_tokenRefreshDisposed) + { + _tokenRefreshTimer = new Timer( + RefreshTokenCallback, null, + TimeSpan.FromSeconds(seconds), + Timeout.InfiniteTimeSpan); + } + } + } + } + public void StopTokenRefresh() { - _tokenRefreshManager.StopTokenRefreshTimer(); + lock (_refreshLock) + { + _tokenRefreshEnabled = false; + _tokenRefreshTimer?.Dispose(); + _tokenRefreshTimer = null; + } } public void Dispose() { - _tokenRefreshManager.Dispose(); + lock (_refreshLock) + { + if (_tokenRefreshDisposed) return; + _tokenRefreshDisposed = true; + _tokenRefreshEnabled = false; + _tokenRefreshTimer?.Dispose(); + _tokenRefreshTimer = null; + } List sessionsToDispose; lock (_sync) diff --git a/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs b/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs deleted file mode 100644 index 658e25c..0000000 --- a/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs +++ /dev/null @@ -1,182 +0,0 @@ -using System.IdentityModel.Tokens.Jwt; -using CassandraMigrationProcessor.Infrastructure; -using CassandraMigrationProcessor.Models; - -namespace CassandraMigrationProcessor.CassandraDriver; -/// -/// Manages AAD token lifecycle and proactive refresh for -/// Cosmos DB Cassandra API connections. -/// -public class TokenRefreshManager : IDisposable -{ - private Timer? _tokenRefreshTimer; - private readonly object _refreshLock = new(); - private readonly Action _refreshSession; - private readonly MigrationLog _log; - private bool _disposed; - private DateTime _tokenExpiresAt = DateTime.MinValue; - private int _consecutiveRefreshFailures; - private const int MaxRefreshFailures = 6; - - public TokenRefreshManager( - MigrationLog log, - Action refreshSession) - { - _log = log; - _refreshSession = refreshSession - ?? throw new ArgumentNullException(nameof(refreshSession)); - } - - /// - /// Detect if a password looks like an AAD/JWT token - /// (very long base64-ish string). - /// - public static bool IsLikelyAadToken(string? password) - { - return password != null && password.Length > 200; - } - - /// - /// Acquire a fresh AAD token for Cosmos DB Cassandra - /// without tracking expiry state. Use for one-shot - /// sessions that do not need proactive refresh. - /// - public static string AcquireAadToken() - { - return AcquireTokenInternal().Token; - } - - /// - /// Generate a fresh AAD token for Cosmos DB Cassandra. - /// Uses DefaultAzureCredential (Managed Identity in - /// App Service, Azure CLI locally). - /// - public string GetFreshAadToken() - { - var tokenResult = AcquireTokenInternal(); - _tokenExpiresAt = tokenResult.ExpiresOn.UtcDateTime; - return tokenResult.Token; - } - - private static Azure.Core.AccessToken AcquireTokenInternal() - { - var credential = new Azure.Identity.DefaultAzureCredential(); - return credential.GetToken( - new Azure.Core.TokenRequestContext( - new[] { "https://cosmos.azure.com/.default" })); - } - - /// - /// Parse the "exp" claim from a JWT to determine when - /// it expires. Returns DateTime.MaxValue if parsing fails. - /// - public static DateTime GetTokenExpiry(string token) - { - try - { - var handler = new JwtSecurityTokenHandler(); - if (handler.CanReadToken(token)) - { - var jwt = handler.ReadJwtToken(token); - return jwt.ValidTo; - } - } - catch (Exception ex) - { - Console.WriteLine($"[WARN] GetTokenExpiry failed: {ex.Message}"); - } - return DateTime.MaxValue; - } - - /// - /// Start the proactive token refresh timer. Schedules - /// a refresh 5 minutes before the token expires. - /// If the token can't be parsed, defaults to refreshing - /// every 50 minutes (tokens typically live 60-75 min). - /// - public void StartTokenRefreshTimer( - string currentToken) - { - lock (_refreshLock) - { - if (_disposed) return; - StopTokenRefreshTimer(); - - DateTime expiry = GetTokenExpiry(currentToken); - if (expiry == DateTime.MaxValue) - { - // Can't parse — refresh every 50 minutes - expiry = DateTime.UtcNow.AddMinutes(50); - } - - _tokenExpiresAt = expiry; - - // Refresh 5 minutes before expiry, minimum 1 min - TimeSpan delay = expiry - DateTime.UtcNow - - TimeSpan.FromMinutes(5); - if (delay < TimeSpan.FromMinutes(1)) - delay = TimeSpan.FromMinutes(1); - - _tokenRefreshTimer = new Timer( - TokenRefreshCallback, null, - delay, Timeout.InfiniteTimeSpan); - } - } - - /// - /// Stop the proactive token refresh timer. - /// - public void StopTokenRefreshTimer() - { - _tokenRefreshTimer?.Dispose(); - _tokenRefreshTimer = null; - } - - private void TokenRefreshCallback(object? state) - { - lock (_refreshLock) - { - if (_disposed) return; - try - { - string freshToken = GetFreshAadToken(); - - _refreshSession(freshToken); - - // Schedule next refresh - _consecutiveRefreshFailures = 0; - StartTokenRefreshTimer(freshToken); - } - catch (Exception ex) - { - _consecutiveRefreshFailures++; - // Exponential backoff capped at 5 min: - // 1: 30s 2: 1m 3: 2m 4: 4m 5+: 5m - int seconds = Math.Min(300, 30 * (1 << Math.Min(_consecutiveRefreshFailures - 1, 4))); - bool tokenAlreadyExpired = DateTime.UtcNow >= _tokenExpiresAt; - LogType severity = (_consecutiveRefreshFailures >= MaxRefreshFailures || tokenAlreadyExpired) - ? LogType.Error - : LogType.Warning; - string msg = $"Token refresh failed (attempt {_consecutiveRefreshFailures}, " + - $"retrying in {seconds}s, tokenExpiresAt={_tokenExpiresAt:O}): {ex.Message}"; - Console.WriteLine($"[{severity}] {msg}"); - _log?.WriteLine(msg, severity); - StopTokenRefreshTimer(); - _tokenRefreshTimer = new Timer( - TokenRefreshCallback, null, - TimeSpan.FromSeconds(seconds), - Timeout.InfiniteTimeSpan); - } - } - } - - public void Dispose() - { - lock (_refreshLock) - { - if (_disposed) return; - _disposed = true; - StopTokenRefreshTimer(); - } - } -} From d8bd8f521914d5309a68841d43b0e37491c6d29c Mon Sep 17 00:00:00 2001 From: Nitesh Vijay Date: Wed, 19 Aug 2026 14:41:45 +0530 Subject: [PATCH 21/32] Centralize retriable operation execution Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 078ac7b9-aa7f-44a1-9053-44c346cec5e1 --- .../CassandraDriver/ArmCredentialDiscovery.cs | 154 ++++++++------- .../CassandraDriver/CassandraClientFactory.cs | 38 ++-- .../CassandraDriver/SchemaManager.cs | 46 ++--- .../DataTransfer/MigrationJobRunner.cs | 19 +- .../DataTransfer/PageReader.cs | 17 +- .../DataTransfer/RowWriteRetry.cs | 61 +++--- .../Infrastructure/RetryExecutor.cs | 184 +++++++++--------- .../Infrastructure/RetryPolicy.cs | 86 ++++++-- 8 files changed, 332 insertions(+), 273 deletions(-) diff --git a/CassandraMigrationProcessor/CassandraDriver/ArmCredentialDiscovery.cs b/CassandraMigrationProcessor/CassandraDriver/ArmCredentialDiscovery.cs index 4a9cf79..2ae6103 100644 --- a/CassandraMigrationProcessor/CassandraDriver/ArmCredentialDiscovery.cs +++ b/CassandraMigrationProcessor/CassandraDriver/ArmCredentialDiscovery.cs @@ -1,3 +1,4 @@ +using CassandraMigrationProcessor.Infrastructure; using System.Net; using System.Net.Http.Headers; using System.Text.Json; @@ -28,6 +29,17 @@ internal class ArmCredentialResult private const int ThrottleRetries = 3; + private sealed class ArmThrottleException : Exception + { + public TimeSpan RetryAfter { get; } + + public ArmThrottleException(TimeSpan retryAfter) + : base("ARM request was throttled.") + { + RetryAfter = retryAfter; + } + } + /// Azure Instance Metadata Service — well-known endpoint (docs.microsoft.com/azure/virtual-machines/instance-metadata-service) private const string ImdsEndpoint = "http://169.254.169.254/metadata/instance"; @@ -42,78 +54,86 @@ internal class ArmCredentialResult private static async Task SendArmRequestAsync( Func buildRequest, string context) { - for (int attempt = 1; attempt <= ThrottleRetries; attempt++) - { - using var req = buildRequest(); - var resp = await _armHttpClient.SendAsync(req); + var retryPolicy = RetryPolicy.Create( + ThrottleRetries, + exception => exception is ArmThrottleException, + (exception, _) => + ((ArmThrottleException)exception).RetryAfter); + return await RetryExecutor.ExecuteOrDefaultAsync( + async (attempt, _) => + { + using var request = buildRequest(); + var response = await _armHttpClient.SendAsync(request); - if (resp.IsSuccessStatusCode) - return resp; + if (response.IsSuccessStatusCode) + return response; - switch (resp.StatusCode) - { - case HttpStatusCode.NotFound: - Console.WriteLine( - $"[INFO] ARM ({context}): 404 — no matching " + - $"resource in this subscription."); - resp.Dispose(); - return null; - - case HttpStatusCode.Unauthorized: - resp.Dispose(); - throw new InvalidOperationException( - $"ARM ({context}) returned 401 Unauthorized. " + - $"The current identity's token was rejected. " + - $"Re-acquire credentials and try again."); - - case HttpStatusCode.Forbidden: - resp.Dispose(); - throw new InvalidOperationException( - $"ARM ({context}) returned 403 Forbidden. " + - $"The caller lacks the RBAC role required " + - $"(typically 'Cosmos DB Account Reader Role' " + - $"or 'DocumentDB Account Contributor')."); - - case HttpStatusCode.TooManyRequests: - // Retry-After comes in two RFC 7231 §7.1.3 shapes - // that are mutually exclusive on the wire: - // "Retry-After: 30" -> Delta - // "Retry-After: Wed, 21 Oct" -> Date - // ARM normally uses Delta but is allowed to send - // Date; we previously silently fell through to - // 2*attempt seconds on the Date form and pounded - // a still-throttled endpoint. - var ra = resp.Headers.RetryAfter; - TimeSpan retryAfter = ra?.Delta - ?? (ra?.Date is { } d - ? d - DateTimeOffset.UtcNow - : (TimeSpan?)null) - ?? TimeSpan.FromSeconds(2 * attempt); - if (retryAfter < TimeSpan.Zero) - retryAfter = TimeSpan.FromSeconds(2 * attempt); - Console.WriteLine( - $"[WARN] ARM ({context}): 429 throttle — " + - $"sleeping {retryAfter.TotalSeconds:F1}s " + - $"(attempt {attempt}/{ThrottleRetries})."); - resp.Dispose(); - if (attempt == ThrottleRetries) return null; - await Task.Delay(retryAfter); - continue; + switch (response.StatusCode) + { + case HttpStatusCode.NotFound: + Console.WriteLine( + $"[INFO] ARM ({context}): 404 — no matching " + + $"resource in this subscription."); + response.Dispose(); + return null!; + + case HttpStatusCode.Unauthorized: + response.Dispose(); + throw new InvalidOperationException( + $"ARM ({context}) returned 401 Unauthorized. " + + $"The current identity's token was rejected. " + + $"Re-acquire credentials and try again."); - default: - var code = (int)resp.StatusCode; - resp.Dispose(); - if (code >= 500) + case HttpStatusCode.Forbidden: + response.Dispose(); + throw new InvalidOperationException( + $"ARM ({context}) returned 403 Forbidden. " + + $"The caller lacks the RBAC role required " + + $"(typically 'Cosmos DB Account Reader Role' " + + $"or 'DocumentDB Account Contributor')."); + + case HttpStatusCode.TooManyRequests: + var retryAfter = ResolveRetryAfter( + response.Headers.RetryAfter, + TimeSpan.FromSeconds(2 * attempt)); + response.Dispose(); + throw new ArmThrottleException(retryAfter); + + default: + var code = (int)response.StatusCode; + var statusCode = response.StatusCode; + response.Dispose(); + if (code >= 500) + throw new InvalidOperationException( + $"ARM ({context}) returned {code} " + + $"({statusCode}) — service outage. " + + $"Retry later."); throw new InvalidOperationException( $"ARM ({context}) returned {code} " + - $"({resp.StatusCode}) — service outage. " + - $"Retry later."); - throw new InvalidOperationException( - $"ARM ({context}) returned {code} " + - $"({resp.StatusCode})."); - } - } - return null; + $"({statusCode})."); + } + }, + retryPolicy, + (exception, attempt) => + { + var throttled = (ArmThrottleException)exception; + Console.WriteLine( + $"[WARN] ARM ({context}): 429 throttle — " + + $"sleeping {throttled.RetryAfter.TotalSeconds:F1}s " + + $"(attempt {attempt}/{ThrottleRetries})."); + }); + } + + private static TimeSpan ResolveRetryAfter( + RetryConditionHeaderValue? retryAfter, + TimeSpan fallback) + { + TimeSpan delay = retryAfter?.Delta + ?? (retryAfter?.Date is { } date + ? date - DateTimeOffset.UtcNow + : (TimeSpan?)null) + ?? fallback; + return delay < TimeSpan.Zero ? fallback : delay; } /// diff --git a/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs b/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs index 9f540fa..20152f3 100644 --- a/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs +++ b/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs @@ -1,5 +1,4 @@ using Cassandra; -using System.Diagnostics; using System.Security.Authentication; using CassandraMigrationProcessor.Infrastructure; using CassandraMigrationProcessor.Models; @@ -50,36 +49,21 @@ public static ISession CreateSourceSession( contactPoint, port, username, password, useSsl: true, maxConnectionsPerHost); - // Single connect+register success path. The loop covers all - // attempts; the `when (attempt < MaxRetries)` filter swallows - // transient failures only on attempts 1..MaxRetries-1, so on - // the final attempt any exception — transient or not — - // propagates out unhandled, matching the original "Final - // attempt — let exception propagate" semantics. const int MaxRetries = 5; - for (int attempt = 1; attempt <= MaxRetries; attempt++) - { - try - { - return ConnectCluster(builder); - } - catch (Exception ex) when ( - ExceptionClassifier.IsTransient(ex) - && attempt < MaxRetries) + var retryPolicy = RetryPolicy.Create( + MaxRetries, + ExceptionClassifier.IsTransient, + (exception, attempt) => TimeSpan.FromMilliseconds( + ExceptionClassifier.GetRetryDelayMs(exception, attempt))); + return RetryExecutor.Execute( + _ => ConnectCluster(builder), + retryPolicy, + (exception, attempt) => { - int delayMs = ExceptionClassifier.GetRetryDelayMs(ex, attempt); MigrationLog.WriteLine( - $"Source connect retry " + - $"{attempt}: {ex.Message}", + $"Source connect retry {attempt}: {exception.Message}", LogType.Warning); - Thread.Sleep(delayMs); - } - } - - // Unreachable: the loop either returns on success or rethrows - // on the final attempt (the `when` filter is false when - // attempt == MaxRetries). - throw new UnreachableException(); + }); } /// diff --git a/CassandraMigrationProcessor/CassandraDriver/SchemaManager.cs b/CassandraMigrationProcessor/CassandraDriver/SchemaManager.cs index dc608a6..5274e32 100644 --- a/CassandraMigrationProcessor/CassandraDriver/SchemaManager.cs +++ b/CassandraMigrationProcessor/CassandraDriver/SchemaManager.cs @@ -721,34 +721,30 @@ public static async Task TableExistsAsync(ISession session, string keyspac if (!tables.Contains(table, StringComparer.OrdinalIgnoreCase)) return false; - for (int attempt = 1; attempt <= ThrottleMaxRetries; attempt++) + var retryPolicy = RetryPolicy.Create( + ThrottleMaxRetries, + ExceptionClassifier.IsThrottle, + (_, attempt) => TimeSpan.FromSeconds( + Math.Min(attempt * 3, 30))); + try { - try - { - var probe = new SimpleStatement( - $"SELECT * FROM \"{keyspace}\".\"{table}\" LIMIT 1"); - probe.SetPageSize(1); - probe.SetAutoPage(false); - probe.SetReadTimeoutMillis(ProbeTimeoutMs); - await session.ExecuteAsync(probe); - return true; - } - catch (Exception ex) - { - if (ExceptionClassifier.IsThrottle(ex) && attempt < ThrottleMaxRetries) + return await RetryExecutor.ExecuteAsync( + async (_, _) => { - int delaySec = Math.Min(attempt * 3, 30); - await Task.Delay(delaySec * 1000); - continue; - } - - if (ExceptionClassifier.IsNotFound(ex)) - return false; - - throw; - } + var probe = new SimpleStatement( + $"SELECT * FROM \"{keyspace}\".\"{table}\" LIMIT 1"); + probe.SetPageSize(1); + probe.SetAutoPage(false); + probe.SetReadTimeoutMillis(ProbeTimeoutMs); + await session.ExecuteAsync(probe); + return true; + }, + retryPolicy); + } + catch (Exception ex) when (ExceptionClassifier.IsNotFound(ex)) + { + return false; } - return false; } /// diff --git a/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs b/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs index 4e8b023..9a376e4 100644 --- a/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs +++ b/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs @@ -617,14 +617,15 @@ private async Task RunOfflineFinalizeAsync(Job job) private Task ProcessWithRetryAsync(Job job, TableMigration mu, JobPartitioning partitioning, CancellationToken token) { return RetryExecutor.ExecuteAsync( - operation: async _ => + operation: async (_, _) => { await ProcessMigrationUnitAsync(job, mu, partitioning, token); return 0; }, - maxAttempts: MigrationDefaults.MaxTableRetries, - shouldRetry: ExceptionClassifier.IsTransient, - delayFor: (ex, attempt) => RetryPolicy.FromException(ex, attempt), + policy: RetryPolicy.Create( + MigrationDefaults.MaxTableRetries, + ExceptionClassifier.IsTransient, + RetryPolicy.FromException), onRetry: (ex, attempt) => _log.WriteLine( $"Table retry {attempt} for {mu.KeyspaceName}.{mu.TableName}: {ex.Message}", LogType.Warning), @@ -957,7 +958,7 @@ private async Task IsTableAccessibleAsync( try { return await RetryExecutor.ExecuteAsync( - operation: _ => + operation: (_, _) => { var probe = new SimpleStatement( $"SELECT * FROM \"{keyspace}\".\"{tableName}\" WHERE COSMOS_CHANGEFEED_FROM_START() = true"); @@ -967,9 +968,11 @@ private async Task IsTableAccessibleAsync( session.Execute(probe); return Task.FromResult(true); }, - maxAttempts: 10, - shouldRetry: ExceptionClassifier.IsThrottle, - delayFor: (_, attempt) => TimeSpan.FromSeconds(Math.Min(attempt * 3, 30)), + policy: RetryPolicy.Create( + 10, + ExceptionClassifier.IsThrottle, + (_, attempt) => TimeSpan.FromSeconds( + Math.Min(attempt * 3, 30))), cancellationToken: cancellationToken); } catch (OperationCanceledException) diff --git a/CassandraMigrationProcessor/DataTransfer/PageReader.cs b/CassandraMigrationProcessor/DataTransfer/PageReader.cs index b6e78f5..70ae0e2 100644 --- a/CassandraMigrationProcessor/DataTransfer/PageReader.cs +++ b/CassandraMigrationProcessor/DataTransfer/PageReader.cs @@ -191,7 +191,7 @@ internal record ReadResult( // intact and will retry the same page once the source stops // throttling. var resultSet = await RetryExecutor.ExecuteOrDefaultAsync( - operation: async _ => + operation: async (_, _) => { var sourceSession = useJson ? _sourceSession.GetSession() @@ -201,12 +201,15 @@ internal record ReadResult( .WaitAsync(_ct) .ConfigureAwait(false); }, - maxAttempts: _maxReadRetries, - shouldRetry: ex => ex is not SourceUdtRegistrationException - && ExceptionClassifier.IsTransient(ex), - delayFor: (ex, attempt) => TimeSpan.FromMilliseconds( - Math.Min(ExceptionClassifier.GetRetryDelayMs(ex, attempt), MaxRetryDelayMs)), - onRetry: (ex, attempt) => + policy: RetryPolicy.Create( + _maxReadRetries, + ex => ex is not SourceUdtRegistrationException + && ExceptionClassifier.IsTransient(ex), + (ex, attempt) => TimeSpan.FromMilliseconds( + Math.Min( + ExceptionClassifier.GetRetryDelayMs(ex, attempt), + MaxRetryDelayMs))), + onFailure: (ex, attempt) => { LastRetryExhaustionException = ex; _log.WriteLine( diff --git a/CassandraMigrationProcessor/DataTransfer/RowWriteRetry.cs b/CassandraMigrationProcessor/DataTransfer/RowWriteRetry.cs index 77d6e28..cfa5bef 100644 --- a/CassandraMigrationProcessor/DataTransfer/RowWriteRetry.cs +++ b/CassandraMigrationProcessor/DataTransfer/RowWriteRetry.cs @@ -95,41 +95,42 @@ public static async Task ExecuteRowGroupsAsync( string rowKind, CancellationToken cancellationToken) { - for (int n = 1; n <= policy.MaxAttempts; n++) + int attempts = 0; + try { - cancellationToken.ThrowIfCancellationRequested(); - var start = Stopwatch.GetTimestamp(); - try - { - await attempt(); - long elapsed = (Stopwatch.GetTimestamp() - start) * 1000 / Stopwatch.Frequency; - return (WriteOutcome.Success, elapsed, null); - } - catch (OperationCanceledException) - { - throw; - } - catch (Exception ex) - { - if (ExceptionClassifier.IsFatal(ex)) - { - log.WriteLine($"FATAL {rowKind}: {ex.GetType().Name}: {ex.Message}", - LogType.Error); - return (WriteOutcome.Fatal, 0, ex); - } - - if (ExceptionClassifier.IsTransient(ex) && n < policy.MaxAttempts) + long elapsed = await RetryExecutor.ExecuteAsync( + async (attemptNumber, _) => { - await Task.Delay(policy.DelayBeforeRetry(n), cancellationToken); - continue; - } - - log.WriteLine($"{rowKind} FAILED after {n} attempt(s): {ex.GetType().Name}: {ex.Message}", + attempts = attemptNumber; + var start = Stopwatch.GetTimestamp(); + await attempt().ConfigureAwait(false); + return (Stopwatch.GetTimestamp() - start) + * 1000 / Stopwatch.Frequency; + }, + policy, + cancellationToken: cancellationToken); + return (WriteOutcome.Success, elapsed, null); + } + catch (OperationCanceledException) + { + throw; + } + catch (Exception ex) + { + if (ExceptionClassifier.IsFatal(ex)) + { + log.WriteLine( + $"FATAL {rowKind}: {ex.GetType().Name}: {ex.Message}", LogType.Error); - return (WriteOutcome.Failed, 0, ex); + return (WriteOutcome.Fatal, 0, ex); } + + log.WriteLine( + $"{rowKind} FAILED after {attempts} attempt(s): " + + $"{ex.GetType().Name}: {ex.Message}", + LogType.Error); + return (WriteOutcome.Failed, 0, ex); } - return (WriteOutcome.Failed, 0, null); } private static void ApplyToCounters( diff --git a/CassandraMigrationProcessor/Infrastructure/RetryExecutor.cs b/CassandraMigrationProcessor/Infrastructure/RetryExecutor.cs index 7c45aa3..fe2b814 100644 --- a/CassandraMigrationProcessor/Infrastructure/RetryExecutor.cs +++ b/CassandraMigrationProcessor/Infrastructure/RetryExecutor.cs @@ -1,146 +1,138 @@ namespace CassandraMigrationProcessor.Infrastructure; /// -/// Shared transient-fault retry helper. Wraps an async operation with -/// linear backoff retry on transient Cassandra exceptions (timeouts, -/// throttles, transport errors as classified by -/// ). Caller-agnostic — -/// nothing here is schema- or query-specific; see callers in -/// SchemaManager, PageReader, etc. +/// Executes synchronous and asynchronous operations using caller-provided +/// retry policies. Operation-specific exception classification, delay, and +/// logging remain outside the executor. /// internal static class RetryExecutor { - /// - /// Execute an async operation with retry on transient errors. - /// Delay between attempts is taken from - /// (which honours - /// server RetryAfterMs hints, applies exponential backoff - /// with jitter, and caps the per-sleep ceiling). The supplied - /// cancellation token is honoured both during the operation and - /// during the backoff sleep, so Stop observes promptly instead - /// of waiting for the next retry timer to fire. - /// public static async Task ExecuteAsync( + Func> operation, + RetryPolicy policy, + Action? onRetry = null, + Action? onFailure = null, + CancellationToken cancellationToken = default) + { + var result = await ExecuteCoreAsync( + operation, policy, onRetry, onFailure, cancellationToken) + .ConfigureAwait(false); + if (result.Succeeded) + return result.Value!; + + System.Runtime.ExceptionServices.ExceptionDispatchInfo + .Capture(result.Exception!) + .Throw(); + throw new System.Diagnostics.UnreachableException(); + } + + public static Task ExecuteAsync( Func> operation, - int maxRetries = MigrationDefaults.TransientRetryMaxAttempts, - int baseDelayMs = MigrationDefaults.TransientRetryBaseDelayMs, CancellationToken cancellationToken = default) { - Exception? lastException = null; - for (int attempt = 1; attempt <= maxRetries; attempt++) - { - cancellationToken.ThrowIfCancellationRequested(); - try - { - return await operation().ConfigureAwait(false); - } - catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) - { - throw; - } - catch (Exception ex) when (attempt < maxRetries - && ExceptionClassifier.IsTransient(ex)) - { - lastException = ex; - var delay = Math.Max( - ExceptionClassifier.GetRetryDelayMs(ex, attempt), - attempt * baseDelayMs); - await Task.Delay(delay, cancellationToken) - .ConfigureAwait(false); - } - } - throw lastException ?? new TimeoutException("Operation timed out after all retries"); + return ExecuteAsync( + (_, _) => operation(), + RetryPolicy.Transient(), + cancellationToken: cancellationToken); } - /// - /// Non-generic overload for fire-and-forget operations. - /// public static Task ExecuteAsync( Func operation, - int maxRetries = MigrationDefaults.TransientRetryMaxAttempts, - int baseDelayMs = MigrationDefaults.TransientRetryBaseDelayMs, CancellationToken cancellationToken = default) { - return ExecuteAsync( - async () => { await operation().ConfigureAwait(false); return 0; }, - maxRetries, baseDelayMs, cancellationToken); + return ExecuteAsync( + async (_, _) => + { + await operation().ConfigureAwait(false); + return true; + }, + RetryPolicy.Transient(), + cancellationToken: cancellationToken); } - /// - /// Generic overload with a caller-supplied retry predicate and - /// custom backoff function — used by sites that need to retry on a - /// narrower set than - /// (e.g. throttle-only retries for table-accessibility probes) or - /// need a different backoff curve. The cancellation token is - /// honoured both during the operation and during the sleep so Stop - /// observes promptly. - /// - public static async Task ExecuteAsync( - Func> operation, - int maxAttempts, - Predicate shouldRetry, - Func delayFor, - Action? onRetry = null, - CancellationToken cancellationToken = default) + public static T Execute( + Func operation, + RetryPolicy policy, + Action? onRetry = null) { - Exception? lastException = null; - for (int attempt = 1; attempt <= maxAttempts; attempt++) + for (int attempt = 1; ; attempt++) { - cancellationToken.ThrowIfCancellationRequested(); try { - return await operation(attempt).ConfigureAwait(false); + return operation(attempt); } - catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) + catch (Exception ex) when ( + attempt < policy.MaxAttempts + && policy.ShouldRetry(ex)) { - throw; - } - catch (Exception ex) when (attempt < maxAttempts && shouldRetry(ex)) - { - lastException = ex; onRetry?.Invoke(ex, attempt); - await Task.Delay(delayFor(ex, attempt), cancellationToken) - .ConfigureAwait(false); + Thread.Sleep(policy.DelayBeforeRetry(ex, attempt)); } } - throw lastException ?? new TimeoutException("Operation timed out after all retries"); } - /// - /// Variant that returns default(T) instead of throwing when - /// every attempt fails on a retryable exception. Used by PageReader - /// where read-retry exhaustion must surface as "no page; re-queue - /// via cooldown" rather than as a thrown error. Non-retryable - /// exceptions still propagate. - /// public static async Task ExecuteOrDefaultAsync( - Func> operation, - int maxAttempts, - Predicate shouldRetry, - Func delayFor, + Func> operation, + RetryPolicy policy, Action? onRetry = null, + Action? onFailure = null, CancellationToken cancellationToken = default) where T : class { - for (int attempt = 1; attempt <= maxAttempts; attempt++) + var result = await ExecuteCoreAsync( + operation, policy, onRetry, onFailure, cancellationToken) + .ConfigureAwait(false); + return result.Succeeded ? result.Value : null; + } + + private static async Task> ExecuteCoreAsync( + Func> operation, + RetryPolicy policy, + Action? onRetry, + Action? onFailure, + CancellationToken cancellationToken) + { + for (int attempt = 1; ; attempt++) { cancellationToken.ThrowIfCancellationRequested(); try { - return await operation(attempt).ConfigureAwait(false); + var value = await operation(attempt, cancellationToken) + .ConfigureAwait(false); + return RetryResult.Success(value); } catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) { throw; } - catch (Exception ex) when (shouldRetry(ex)) + catch (Exception ex) when (policy.ShouldRetry(ex)) { + onFailure?.Invoke(ex, attempt); + if (attempt >= policy.MaxAttempts) + return RetryResult.Failure(ex); + onRetry?.Invoke(ex, attempt); - if (attempt >= maxAttempts) return null; - await Task.Delay(delayFor(ex, attempt), cancellationToken) + await Task.Delay( + policy.DelayBeforeRetry(ex, attempt), + cancellationToken) .ConfigureAwait(false); } } - return null; + } + + private readonly record struct RetryResult( + bool Succeeded, + T? Value, + Exception? Exception) + { + public static RetryResult Success(T value) + { + return new RetryResult(true, value, null); + } + + public static RetryResult Failure(Exception exception) + { + return new RetryResult(false, default, exception); + } } } diff --git a/CassandraMigrationProcessor/Infrastructure/RetryPolicy.cs b/CassandraMigrationProcessor/Infrastructure/RetryPolicy.cs index 04795b5..13ea349 100644 --- a/CassandraMigrationProcessor/Infrastructure/RetryPolicy.cs +++ b/CassandraMigrationProcessor/Infrastructure/RetryPolicy.cs @@ -16,14 +16,26 @@ namespace CassandraMigrationProcessor.Infrastructure; public sealed class RetryPolicy { public int MaxAttempts { get; } - private readonly Func _delayFor; + private readonly Predicate _shouldRetry; + private readonly Func _delayFor; - private RetryPolicy(int maxAttempts, Func delayFor) + private RetryPolicy( + int maxAttempts, + Predicate shouldRetry, + Func delayFor) { if (maxAttempts < 1) throw new ArgumentOutOfRangeException(nameof(maxAttempts), "Must be >= 1."); MaxAttempts = maxAttempts; - _delayFor = delayFor; + _shouldRetry = shouldRetry + ?? throw new ArgumentNullException(nameof(shouldRetry)); + _delayFor = delayFor + ?? throw new ArgumentNullException(nameof(delayFor)); + } + + public bool ShouldRetry(Exception exception) + { + return _shouldRetry(exception); } /// @@ -31,26 +43,71 @@ private RetryPolicy(int maxAttempts, Func delayFor) /// (1-based, so the delay returned for attempt=1 is the wait /// between the first failed attempt and the second attempt). /// - public TimeSpan DelayBeforeRetry(int attempt) => _delayFor(attempt); + public TimeSpan DelayBeforeRetry(Exception exception, int attempt) + { + return _delayFor(exception, attempt); + } + + public static RetryPolicy Create( + int maxAttempts, + Predicate shouldRetry, + Func delayFor) + { + return new RetryPolicy(maxAttempts, shouldRetry, delayFor); + } /// /// Linear backoff: baseDelay × attempt. Matches the historical /// behaviour of the bulk-copy retry loop. /// - public static RetryPolicy Linear(int maxAttempts, TimeSpan baseDelay) - => new(maxAttempts, attempt => TimeSpan.FromMilliseconds(baseDelay.TotalMilliseconds * attempt)); + public static RetryPolicy Linear( + int maxAttempts, + TimeSpan baseDelay, + Predicate? shouldRetry = null) + { + return new RetryPolicy( + maxAttempts, + shouldRetry ?? ExceptionClassifier.IsTransient, + (_, attempt) => TimeSpan.FromMilliseconds( + baseDelay.TotalMilliseconds * attempt)); + } /// /// Capped exponential backoff: baseDelay × 2^(attempt-1), /// clamped to . Suitable when a retry storm is /// likely (target throttling, e.g. 429s). /// - public static RetryPolicy Exponential(int maxAttempts, TimeSpan baseDelay, TimeSpan cap) - => new(maxAttempts, attempt => - { - var ms = baseDelay.TotalMilliseconds * Math.Pow(2, attempt - 1); - return ms >= cap.TotalMilliseconds ? cap : TimeSpan.FromMilliseconds(ms); - }); + public static RetryPolicy Exponential( + int maxAttempts, + TimeSpan baseDelay, + TimeSpan cap, + Predicate? shouldRetry = null) + { + return new RetryPolicy( + maxAttempts, + shouldRetry ?? ExceptionClassifier.IsTransient, + (_, attempt) => + { + var ms = baseDelay.TotalMilliseconds + * Math.Pow(2, attempt - 1); + return ms >= cap.TotalMilliseconds + ? cap + : TimeSpan.FromMilliseconds(ms); + }); + } + + public static RetryPolicy Transient( + int maxAttempts = MigrationDefaults.TransientRetryMaxAttempts, + int minimumDelayMs = MigrationDefaults.TransientRetryBaseDelayMs) + { + return new RetryPolicy( + maxAttempts, + ExceptionClassifier.IsTransient, + (exception, attempt) => TimeSpan.FromMilliseconds( + Math.Max( + ExceptionClassifier.GetRetryDelayMs(exception, attempt), + attempt * minimumDelayMs))); + } /// /// Single-shot delay computed from a server-hinted exception via @@ -60,5 +117,8 @@ public static RetryPolicy Exponential(int maxAttempts, TimeSpan baseDelay, TimeS /// Returned as a ready for Task.Delay. /// public static TimeSpan FromException(Exception ex, int attempt) - => TimeSpan.FromMilliseconds(ExceptionClassifier.GetRetryDelayMs(ex, attempt)); + { + return TimeSpan.FromMilliseconds( + ExceptionClassifier.GetRetryDelayMs(ex, attempt)); + } } From 974705c04711a5ad395c63476ebd421bf2bd59bd Mon Sep 17 00:00:00 2001 From: Nitesh Vijay Date: Wed, 19 Aug 2026 16:56:18 +0530 Subject: [PATCH 22/32] Simplify shared session architecture Remove the redundant retry rewrite and collapse source and target session factory layers while preserving shared-source rotation, UDT registration, and gated worker target creation. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 078ac7b9-aa7f-44a1-9053-44c346cec5e1 --- .../CassandraDriver/ArmCredentialDiscovery.cs | 154 +++++++-------- .../CassandraDriver/CassandraClientFactory.cs | 118 +++++------ .../CassandraDriver/ISessionFactory.cs | 72 ------- .../CassandraDriver/JobSessionFactory.cs | 40 ++++ .../CassandraDriver/SchemaManager.cs | 46 +++-- .../CassandraDriver/SourceSessionFactory.cs | 35 ---- .../CassandraDriver/SourceSessionWrapper.cs | 94 +++++---- .../DataTransfer/DataCopyWorker.cs | 2 +- .../DataTransfer/JobPipeline.cs | 3 +- .../DataTransfer/MigrationJobRunner.cs | 41 ++-- .../DataTransfer/PageReader.cs | 34 +--- .../DataTransfer/PageWriter.cs | 2 +- .../DataTransfer/PipelineContext.cs | 2 +- .../DataTransfer/RowWriteRetry.cs | 61 +++--- .../Infrastructure/RetryExecutor.cs | 184 +++++++++--------- .../Infrastructure/RetryPolicy.cs | 86 ++------ .../Models/TableCopySpec.cs | 2 +- 17 files changed, 413 insertions(+), 563 deletions(-) delete mode 100644 CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs create mode 100644 CassandraMigrationProcessor/CassandraDriver/JobSessionFactory.cs delete mode 100644 CassandraMigrationProcessor/CassandraDriver/SourceSessionFactory.cs diff --git a/CassandraMigrationProcessor/CassandraDriver/ArmCredentialDiscovery.cs b/CassandraMigrationProcessor/CassandraDriver/ArmCredentialDiscovery.cs index 2ae6103..4a9cf79 100644 --- a/CassandraMigrationProcessor/CassandraDriver/ArmCredentialDiscovery.cs +++ b/CassandraMigrationProcessor/CassandraDriver/ArmCredentialDiscovery.cs @@ -1,4 +1,3 @@ -using CassandraMigrationProcessor.Infrastructure; using System.Net; using System.Net.Http.Headers; using System.Text.Json; @@ -29,17 +28,6 @@ internal class ArmCredentialResult private const int ThrottleRetries = 3; - private sealed class ArmThrottleException : Exception - { - public TimeSpan RetryAfter { get; } - - public ArmThrottleException(TimeSpan retryAfter) - : base("ARM request was throttled.") - { - RetryAfter = retryAfter; - } - } - /// Azure Instance Metadata Service — well-known endpoint (docs.microsoft.com/azure/virtual-machines/instance-metadata-service) private const string ImdsEndpoint = "http://169.254.169.254/metadata/instance"; @@ -54,86 +42,78 @@ public ArmThrottleException(TimeSpan retryAfter) private static async Task SendArmRequestAsync( Func buildRequest, string context) { - var retryPolicy = RetryPolicy.Create( - ThrottleRetries, - exception => exception is ArmThrottleException, - (exception, _) => - ((ArmThrottleException)exception).RetryAfter); - return await RetryExecutor.ExecuteOrDefaultAsync( - async (attempt, _) => - { - using var request = buildRequest(); - var response = await _armHttpClient.SendAsync(request); + for (int attempt = 1; attempt <= ThrottleRetries; attempt++) + { + using var req = buildRequest(); + var resp = await _armHttpClient.SendAsync(req); - if (response.IsSuccessStatusCode) - return response; + if (resp.IsSuccessStatusCode) + return resp; - switch (response.StatusCode) - { - case HttpStatusCode.NotFound: - Console.WriteLine( - $"[INFO] ARM ({context}): 404 — no matching " + - $"resource in this subscription."); - response.Dispose(); - return null!; - - case HttpStatusCode.Unauthorized: - response.Dispose(); - throw new InvalidOperationException( - $"ARM ({context}) returned 401 Unauthorized. " + - $"The current identity's token was rejected. " + - $"Re-acquire credentials and try again."); + switch (resp.StatusCode) + { + case HttpStatusCode.NotFound: + Console.WriteLine( + $"[INFO] ARM ({context}): 404 — no matching " + + $"resource in this subscription."); + resp.Dispose(); + return null; + + case HttpStatusCode.Unauthorized: + resp.Dispose(); + throw new InvalidOperationException( + $"ARM ({context}) returned 401 Unauthorized. " + + $"The current identity's token was rejected. " + + $"Re-acquire credentials and try again."); + + case HttpStatusCode.Forbidden: + resp.Dispose(); + throw new InvalidOperationException( + $"ARM ({context}) returned 403 Forbidden. " + + $"The caller lacks the RBAC role required " + + $"(typically 'Cosmos DB Account Reader Role' " + + $"or 'DocumentDB Account Contributor')."); + + case HttpStatusCode.TooManyRequests: + // Retry-After comes in two RFC 7231 §7.1.3 shapes + // that are mutually exclusive on the wire: + // "Retry-After: 30" -> Delta + // "Retry-After: Wed, 21 Oct" -> Date + // ARM normally uses Delta but is allowed to send + // Date; we previously silently fell through to + // 2*attempt seconds on the Date form and pounded + // a still-throttled endpoint. + var ra = resp.Headers.RetryAfter; + TimeSpan retryAfter = ra?.Delta + ?? (ra?.Date is { } d + ? d - DateTimeOffset.UtcNow + : (TimeSpan?)null) + ?? TimeSpan.FromSeconds(2 * attempt); + if (retryAfter < TimeSpan.Zero) + retryAfter = TimeSpan.FromSeconds(2 * attempt); + Console.WriteLine( + $"[WARN] ARM ({context}): 429 throttle — " + + $"sleeping {retryAfter.TotalSeconds:F1}s " + + $"(attempt {attempt}/{ThrottleRetries})."); + resp.Dispose(); + if (attempt == ThrottleRetries) return null; + await Task.Delay(retryAfter); + continue; - case HttpStatusCode.Forbidden: - response.Dispose(); - throw new InvalidOperationException( - $"ARM ({context}) returned 403 Forbidden. " + - $"The caller lacks the RBAC role required " + - $"(typically 'Cosmos DB Account Reader Role' " + - $"or 'DocumentDB Account Contributor')."); - - case HttpStatusCode.TooManyRequests: - var retryAfter = ResolveRetryAfter( - response.Headers.RetryAfter, - TimeSpan.FromSeconds(2 * attempt)); - response.Dispose(); - throw new ArmThrottleException(retryAfter); - - default: - var code = (int)response.StatusCode; - var statusCode = response.StatusCode; - response.Dispose(); - if (code >= 500) - throw new InvalidOperationException( - $"ARM ({context}) returned {code} " + - $"({statusCode}) — service outage. " + - $"Retry later."); + default: + var code = (int)resp.StatusCode; + resp.Dispose(); + if (code >= 500) throw new InvalidOperationException( $"ARM ({context}) returned {code} " + - $"({statusCode})."); - } - }, - retryPolicy, - (exception, attempt) => - { - var throttled = (ArmThrottleException)exception; - Console.WriteLine( - $"[WARN] ARM ({context}): 429 throttle — " + - $"sleeping {throttled.RetryAfter.TotalSeconds:F1}s " + - $"(attempt {attempt}/{ThrottleRetries})."); - }); - } - - private static TimeSpan ResolveRetryAfter( - RetryConditionHeaderValue? retryAfter, - TimeSpan fallback) - { - TimeSpan delay = retryAfter?.Delta - ?? (retryAfter?.Date is { } date - ? date - DateTimeOffset.UtcNow - : (TimeSpan?)null) - ?? fallback; - return delay < TimeSpan.Zero ? fallback : delay; + $"({resp.StatusCode}) — service outage. " + + $"Retry later."); + throw new InvalidOperationException( + $"ARM ({context}) returned {code} " + + $"({resp.StatusCode})."); + } + } + return null; } /// diff --git a/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs b/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs index 20152f3..614dfef 100644 --- a/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs +++ b/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs @@ -1,9 +1,16 @@ using Cassandra; +using System.Diagnostics; using System.Security.Authentication; using CassandraMigrationProcessor.Infrastructure; using CassandraMigrationProcessor.Models; namespace CassandraMigrationProcessor.CassandraDriver; +internal sealed record SourceSessionSettings( + string ContactPoint, + int Port, + string Username, + int MaxConnectionsPerHost); + /// /// Creates Cassandra ISession instances for source (Cosmos DB) /// and target (OSS Cassandra) clusters. @@ -49,21 +56,36 @@ public static ISession CreateSourceSession( contactPoint, port, username, password, useSsl: true, maxConnectionsPerHost); + // Single connect+register success path. The loop covers all + // attempts; the `when (attempt < MaxRetries)` filter swallows + // transient failures only on attempts 1..MaxRetries-1, so on + // the final attempt any exception — transient or not — + // propagates out unhandled, matching the original "Final + // attempt — let exception propagate" semantics. const int MaxRetries = 5; - var retryPolicy = RetryPolicy.Create( - MaxRetries, - ExceptionClassifier.IsTransient, - (exception, attempt) => TimeSpan.FromMilliseconds( - ExceptionClassifier.GetRetryDelayMs(exception, attempt))); - return RetryExecutor.Execute( - _ => ConnectCluster(builder), - retryPolicy, - (exception, attempt) => + for (int attempt = 1; attempt <= MaxRetries; attempt++) + { + try + { + return ConnectCluster(builder); + } + catch (Exception ex) when ( + ExceptionClassifier.IsTransient(ex) + && attempt < MaxRetries) { + int delayMs = ExceptionClassifier.GetRetryDelayMs(ex, attempt); MigrationLog.WriteLine( - $"Source connect retry {attempt}: {exception.Message}", + $"Source connect retry " + + $"{attempt}: {ex.Message}", LogType.Warning); - }); + Thread.Sleep(delayMs); + } + } + + // Unreachable: the loop either returns on success or rethrows + // on the final attempt (the `when` filter is false when + // attempt == MaxRetries). + throw new UnreachableException(); } /// @@ -259,41 +281,6 @@ private static ISession ConnectCluster(Builder builder) } } - /// - /// 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. - /// - public static ISession CreateSourceSession( - MigrationLog MigrationLog, Job job) - { - string credential = ResolveSourceCredential(job); - var settings = ResolveSourceSessionSettings(job); - - return CreateSourceSessionWithCredential( - MigrationLog, settings, credential); - } - - internal static string ResolveSourceCredential(Job job) - { - if (string.IsNullOrEmpty(job.SourceContactPoint)) - throw new ArgumentException("Source contact point is required", nameof(job)); - - string credential = job.SourcePassword ?? string.Empty; - if (string.IsNullOrEmpty(credential) || job.SourceUseAad) - { - credential = 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 an and leak the - // bearer JWT to the browser DOM. - job.SourceUseAad = true; - } - return credential; - } - internal static string AcquireAadToken() { var credential = new Azure.Identity.DefaultAzureCredential(); @@ -303,21 +290,9 @@ internal static string AcquireAadToken() .Token; } - internal static ISession CreateSourceSessionWithCredential( - MigrationLog migrationLog, - SourceSessionSettings settings, - string credential) - { - return CreateSourceSession( - migrationLog, - settings.ContactPoint, - settings.Port, - settings.Username, - credential, - settings.MaxConnectionsPerHost); - } - - internal static SourceSessionSettings ResolveSourceSessionSettings( + internal static ( + SourceSessionSettings Settings, + string Credential) ResolveSourceSession( Job job, int workerCount = 0) { @@ -326,6 +301,15 @@ internal static SourceSessionSettings ResolveSourceSessionSettings( bool useAad = job.SourceUseAad || string.IsNullOrEmpty(job.SourcePassword); + string credential = job.SourcePassword ?? string.Empty; + if (useAad) + { + 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; + } + string username = job.SourceUsername ?? string.Empty; if (string.IsNullOrWhiteSpace(username) && useAad) @@ -345,11 +329,13 @@ internal static SourceSessionSettings ResolveSourceSessionSettings( 8); } - return new SourceSessionSettings( - job.SourceContactPoint, - job.SourcePort, - username, - maxConnectionsPerHost); + return ( + new SourceSessionSettings( + job.SourceContactPoint, + job.SourcePort, + username, + maxConnectionsPerHost), + credential); } /// diff --git a/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs b/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs deleted file mode 100644 index a2619c3..0000000 --- a/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs +++ /dev/null @@ -1,72 +0,0 @@ -using Cassandra; -using CassandraMigrationProcessor.Models; -using CassandraMigrationProcessor.Infrastructure; - -namespace CassandraMigrationProcessor.CassandraDriver; - -/// -/// Creates worker-owned sessions. The consumer determines the session role; -/// job-owned shared sessions are passed directly instead of using this factory. -/// -public interface ISessionFactory -{ - /// Mint a new keyspace-agnostic session. - Task CreateSessionAsync(CancellationToken cancellationToken); -} - -/// -/// Limits simultaneous session opens. This prevents high-worker jobs from -/// creating a connection storm during startup. -/// -public sealed class GatedSessionFactory : ISessionFactory -{ - 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 CreateSessionAsync(CancellationToken cancellationToken) - { - await _creationGate.WaitAsync(cancellationToken).ConfigureAwait(false); - try - { - return await _inner.CreateSessionAsync(cancellationToken) - .ConfigureAwait(false); - } - finally - { - _creationGate.Release(); - } - } -} - -/// -/// Default bound to a single -/// . Delegates to -/// so the connection-construction policy stays in one place. -/// -public sealed class JobSessionFactory : ISessionFactory -{ - private readonly MigrationLog _log; - private readonly Job _job; - - public JobSessionFactory(MigrationLog log, Job job) - { - _log = log; - _job = job; - } - - public async Task CreateSessionAsync(CancellationToken cancellationToken) - { - cancellationToken.ThrowIfCancellationRequested(); - return await CassandraClientFactory.CreateTargetSessionAsync(_log, _job) - .ConfigureAwait(false); - } -} diff --git a/CassandraMigrationProcessor/CassandraDriver/JobSessionFactory.cs b/CassandraMigrationProcessor/CassandraDriver/JobSessionFactory.cs new file mode 100644 index 0000000..f41f5e3 --- /dev/null +++ b/CassandraMigrationProcessor/CassandraDriver/JobSessionFactory.cs @@ -0,0 +1,40 @@ +using Cassandra; +using CassandraMigrationProcessor.Models; +using CassandraMigrationProcessor.Infrastructure; + +namespace CassandraMigrationProcessor.CassandraDriver; + +/// +/// Creates worker-owned target sessions for a job while limiting simultaneous +/// opens to prevent a connection storm during startup. +/// +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 CreateSessionAsync(CancellationToken cancellationToken) + { + await _creationGate.WaitAsync(cancellationToken).ConfigureAwait(false); + try + { + return await CassandraClientFactory.CreateTargetSessionAsync( + _log, _job).ConfigureAwait(false); + } + finally + { + _creationGate.Release(); + } + } +} diff --git a/CassandraMigrationProcessor/CassandraDriver/SchemaManager.cs b/CassandraMigrationProcessor/CassandraDriver/SchemaManager.cs index 5274e32..dc608a6 100644 --- a/CassandraMigrationProcessor/CassandraDriver/SchemaManager.cs +++ b/CassandraMigrationProcessor/CassandraDriver/SchemaManager.cs @@ -721,30 +721,34 @@ public static async Task TableExistsAsync(ISession session, string keyspac if (!tables.Contains(table, StringComparer.OrdinalIgnoreCase)) return false; - var retryPolicy = RetryPolicy.Create( - ThrottleMaxRetries, - ExceptionClassifier.IsThrottle, - (_, attempt) => TimeSpan.FromSeconds( - Math.Min(attempt * 3, 30))); - try + for (int attempt = 1; attempt <= ThrottleMaxRetries; attempt++) { - return await RetryExecutor.ExecuteAsync( - async (_, _) => + try + { + var probe = new SimpleStatement( + $"SELECT * FROM \"{keyspace}\".\"{table}\" LIMIT 1"); + probe.SetPageSize(1); + probe.SetAutoPage(false); + probe.SetReadTimeoutMillis(ProbeTimeoutMs); + await session.ExecuteAsync(probe); + return true; + } + catch (Exception ex) + { + if (ExceptionClassifier.IsThrottle(ex) && attempt < ThrottleMaxRetries) { - var probe = new SimpleStatement( - $"SELECT * FROM \"{keyspace}\".\"{table}\" LIMIT 1"); - probe.SetPageSize(1); - probe.SetAutoPage(false); - probe.SetReadTimeoutMillis(ProbeTimeoutMs); - await session.ExecuteAsync(probe); - return true; - }, - retryPolicy); - } - catch (Exception ex) when (ExceptionClassifier.IsNotFound(ex)) - { - return false; + int delaySec = Math.Min(attempt * 3, 30); + await Task.Delay(delaySec * 1000); + continue; + } + + if (ExceptionClassifier.IsNotFound(ex)) + return false; + + throw; + } } + return false; } /// diff --git a/CassandraMigrationProcessor/CassandraDriver/SourceSessionFactory.cs b/CassandraMigrationProcessor/CassandraDriver/SourceSessionFactory.cs deleted file mode 100644 index 3aaecdf..0000000 --- a/CassandraMigrationProcessor/CassandraDriver/SourceSessionFactory.cs +++ /dev/null @@ -1,35 +0,0 @@ -using Cassandra; -using CassandraMigrationProcessor.Infrastructure; - -namespace CassandraMigrationProcessor.CassandraDriver; - -internal sealed record SourceSessionSettings( - string ContactPoint, - int Port, - string Username, - int MaxConnectionsPerHost); - -public interface ICredentialSessionFactory -{ - ISession CreateSession(string credential); -} - -internal sealed class SourceSessionFactory : ICredentialSessionFactory -{ - private readonly MigrationLog _log; - private readonly SourceSessionSettings _settings; - - public SourceSessionFactory( - MigrationLog log, - SourceSessionSettings settings) - { - _log = log; - _settings = settings; - } - - public ISession CreateSession(string credential) - { - return CassandraClientFactory.CreateSourceSessionWithCredential( - _log, _settings, credential); - } -} diff --git a/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs b/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs index bb4c23c..9e04246 100644 --- a/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs +++ b/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs @@ -19,7 +19,7 @@ public SourceUdtRegistrationException(string keyspace, Exception innerException) /// Rotated sessions remain available for a bounded grace period so in-flight /// operations can complete. /// -public sealed class SourceSessionWrapper : IDisposable +internal sealed class SourceSessionWrapper : IDisposable { private static readonly TimeSpan RetiredSessionDisposalDelay = TimeSpan.FromMinutes(10); @@ -28,7 +28,7 @@ public sealed class SourceSessionWrapper : IDisposable private readonly object _sync = new(); private readonly object _refreshLock = new(); private readonly MigrationLog _log; - private readonly ICredentialSessionFactory _sessionFactory; + private readonly SourceSessionSettings _settings; private readonly HashSet _retiredSessions = new(ReferenceEqualityComparer.Instance); private readonly ConcurrentDictionary<(ISession Session, string Keyspace), Lazy> @@ -41,23 +41,48 @@ public sealed class SourceSessionWrapper : IDisposable private bool _tokenRefreshDisposed; private bool _disposed; - public SourceSessionWrapper( + private SourceSessionWrapper( MigrationLog log, - ICredentialSessionFactory sessionFactory) + SourceSessionSettings settings) { _log = log ?? throw new ArgumentNullException(nameof(log)); - _sessionFactory = sessionFactory - ?? throw new ArgumentNullException(nameof(sessionFactory)); + _settings = settings + ?? throw new ArgumentNullException(nameof(settings)); + } + + public static SourceSessionWrapper Create( + MigrationLog log, + Job job, + int workerCount) + { + var source = CassandraClientFactory.ResolveSourceSession( + job, workerCount); + var wrapper = new SourceSessionWrapper(log, source.Settings); + try + { + wrapper.Initialize(source.Credential); + return wrapper; + } + catch + { + wrapper.Dispose(); + throw; + } } public ISession GetSession() { - return GetCurrentSession(); + lock (_sync) + { + ObjectDisposedException.ThrowIf(_disposed, this); + return _currentSession + ?? throw new InvalidOperationException("The session provider has not been initialized."); + } } public async Task GetTypedSessionAsync(string keyspace) { - var session = GetCurrentSession(); + var session = GetSession(); var key = (Session: session, Keyspace: keyspace); var registration = _udtRegistrations.GetOrAdd( key, @@ -79,16 +104,6 @@ public async Task GetTypedSessionAsync(string keyspace) return session; } - private ISession GetCurrentSession() - { - lock (_sync) - { - ObjectDisposedException.ThrowIf(_disposed, this); - return _currentSession - ?? throw new InvalidOperationException("The session provider has not been initialized."); - } - } - private static async Task RegisterUdtsAsync( ISession session, string keyspace) @@ -99,11 +114,11 @@ await DynamicUdtRegistrar.RegisterAsync( session, keyspace, allUdts); } - public ISession Initialize(string credential) + public void Initialize(string credential) { ArgumentException.ThrowIfNullOrWhiteSpace(credential); - var session = _sessionFactory.CreateSession(credential); + var session = CreateSession(credential); try { lock (_sync) @@ -121,16 +136,15 @@ public ISession Initialize(string credential) throw; } - if (IsLikelyAadToken(credential)) + if (credential.Length > 200) StartTokenRefresh(credential); - return session; } - public void Refresh(string credential) + private void Refresh(string credential) { ArgumentException.ThrowIfNullOrWhiteSpace(credential); - var session = _sessionFactory.CreateSession(credential); + var session = CreateSession(credential); ISession? retiredSession; try { @@ -155,6 +169,17 @@ public void Refresh(string credential) _ = DisposeRetiredSessionAfterDelayAsync(retiredSession); } + private ISession CreateSession(string credential) + { + return CassandraClientFactory.CreateSourceSession( + _log, + _settings.ContactPoint, + _settings.Port, + _settings.Username, + credential, + _settings.MaxConnectionsPerHost); + } + private async Task DisposeRetiredSessionAfterDelayAsync(ISession session) { await Task.Delay(RetiredSessionDisposalDelay).ConfigureAwait(false); @@ -182,11 +207,6 @@ private void RemoveUdtRegistrations(ISession session) } } - private static bool IsLikelyAadToken(string? credential) - { - return credential != null && credential.Length > 200; - } - private static DateTime GetTokenExpiry(string token) { try @@ -219,7 +239,7 @@ private void StartTokenRefresh(string currentToken) private void ScheduleTokenRefresh(string currentToken) { - _tokenRefreshTimer?.Dispose(); + StopTokenRefreshCore(); DateTime expiry = GetTokenExpiry(currentToken); if (expiry == DateTime.MaxValue) @@ -272,7 +292,7 @@ private void RefreshTokenCallback(object? state) Console.WriteLine($"[{severity}] {message}"); _log.WriteLine(message, severity); - _tokenRefreshTimer?.Dispose(); + StopTokenRefreshCore(); if (_tokenRefreshEnabled && !_tokenRefreshDisposed) { _tokenRefreshTimer = new Timer( @@ -289,11 +309,16 @@ public void StopTokenRefresh() lock (_refreshLock) { _tokenRefreshEnabled = false; - _tokenRefreshTimer?.Dispose(); - _tokenRefreshTimer = null; + StopTokenRefreshCore(); } } + private void StopTokenRefreshCore() + { + _tokenRefreshTimer?.Dispose(); + _tokenRefreshTimer = null; + } + public void Dispose() { lock (_refreshLock) @@ -301,8 +326,7 @@ public void Dispose() if (_tokenRefreshDisposed) return; _tokenRefreshDisposed = true; _tokenRefreshEnabled = false; - _tokenRefreshTimer?.Dispose(); - _tokenRefreshTimer = null; + StopTokenRefreshCore(); } List sessionsToDispose; diff --git a/CassandraMigrationProcessor/DataTransfer/DataCopyWorker.cs b/CassandraMigrationProcessor/DataTransfer/DataCopyWorker.cs index 708fe8c..557a098 100644 --- a/CassandraMigrationProcessor/DataTransfer/DataCopyWorker.cs +++ b/CassandraMigrationProcessor/DataTransfer/DataCopyWorker.cs @@ -33,7 +33,7 @@ public async Task RunAsync(PipelineContext ctx) Partition? current = null; try { - reader = await PageReader.CreateAsync( + reader = new PageReader( _workerLog, ctx.SourceSession, ctx.ReaderConfig, diff --git a/CassandraMigrationProcessor/DataTransfer/JobPipeline.cs b/CassandraMigrationProcessor/DataTransfer/JobPipeline.cs index d789b9e..0bbd668 100644 --- a/CassandraMigrationProcessor/DataTransfer/JobPipeline.cs +++ b/CassandraMigrationProcessor/DataTransfer/JobPipeline.cs @@ -22,7 +22,6 @@ internal sealed class JobPipeline : IDisposable, IAsyncDisposable public JobPipeline(MigrationLog log, Job job, PipelineConfig pipelineConfig, JobPartitioning partitioning, SourceSessionWrapper sourceSession, - ISessionFactory sessionFactory, JobControl control) { _log = log; @@ -47,7 +46,7 @@ public JobPipeline(MigrationLog log, Job job, PipelineConfig pipelineConfig, Context = new PipelineContext( _partitions, sourceSession, - sessionFactory, + new JobSessionFactory(log, job), readerConfig, writerConfig, EnableReplay: enableReplay, diff --git a/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs b/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs index 9a376e4..eabba58 100644 --- a/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs +++ b/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs @@ -82,14 +82,8 @@ public static async Task CreateAsync( ISession? target = null; try { - var sourceSettings = CassandraClientFactory.ResolveSourceSessionSettings( - job, pipelineConfig.WorkerCount); - sourceSessions = new SourceSessionWrapper( - log, - new SourceSessionFactory(log, sourceSettings)); - string sourceCredential = CassandraClientFactory.ResolveSourceCredential( - job); - sourceSessions.Initialize(sourceCredential); + sourceSessions = SourceSessionWrapper.Create( + log, job, pipelineConfig.WorkerCount); target = await CassandraClientFactory.CreateTargetSessionAsync(log, job); return new MigrationJobRunner( log, job, pipelineConfig, control, sourceSessions, target); @@ -179,8 +173,6 @@ public async Task StartAsync() _pipeline = new JobPipeline( _log, job, _pipelineConfig, partitioning, _sourceSessions, - new GatedSessionFactory( - new JobSessionFactory(_log, job)), _control); _pipeline.Start(); @@ -617,15 +609,14 @@ private async Task RunOfflineFinalizeAsync(Job job) private Task ProcessWithRetryAsync(Job job, TableMigration mu, JobPartitioning partitioning, CancellationToken token) { return RetryExecutor.ExecuteAsync( - operation: async (_, _) => + operation: async _ => { await ProcessMigrationUnitAsync(job, mu, partitioning, token); return 0; }, - policy: RetryPolicy.Create( - MigrationDefaults.MaxTableRetries, - ExceptionClassifier.IsTransient, - RetryPolicy.FromException), + maxAttempts: MigrationDefaults.MaxTableRetries, + shouldRetry: ExceptionClassifier.IsTransient, + delayFor: (ex, attempt) => RetryPolicy.FromException(ex, attempt), onRetry: (ex, attempt) => _log.WriteLine( $"Table retry {attempt} for {mu.KeyspaceName}.{mu.TableName}: {ex.Message}", LogType.Warning), @@ -924,12 +915,16 @@ void AddExpandedUnit(string keyspaceName, string tableName) => try { - var sourceSession = _sourceSessions.GetSession(); - var tables = await CassandraQueries.ListTablesAsync(sourceSession, keyspace); + var tables = await CassandraQueries.ListTablesAsync( + _sourceSessions.GetSession(), keyspace); foreach (var tableName in tables) { cancellationToken.ThrowIfCancellationRequested(); - if (await IsTableAccessibleAsync(sourceSession, keyspace, tableName, cancellationToken)) + if (await IsTableAccessibleAsync( + _sourceSessions.GetSession(), + keyspace, + tableName, + cancellationToken)) { AddExpandedUnit(keyspace, tableName); } @@ -958,7 +953,7 @@ private async Task IsTableAccessibleAsync( try { return await RetryExecutor.ExecuteAsync( - operation: (_, _) => + operation: _ => { var probe = new SimpleStatement( $"SELECT * FROM \"{keyspace}\".\"{tableName}\" WHERE COSMOS_CHANGEFEED_FROM_START() = true"); @@ -968,11 +963,9 @@ private async Task IsTableAccessibleAsync( session.Execute(probe); return Task.FromResult(true); }, - policy: RetryPolicy.Create( - 10, - ExceptionClassifier.IsThrottle, - (_, attempt) => TimeSpan.FromSeconds( - Math.Min(attempt * 3, 30))), + maxAttempts: 10, + shouldRetry: ExceptionClassifier.IsThrottle, + delayFor: (_, attempt) => TimeSpan.FromSeconds(Math.Min(attempt * 3, 30)), cancellationToken: cancellationToken); } catch (OperationCanceledException) diff --git a/CassandraMigrationProcessor/DataTransfer/PageReader.cs b/CassandraMigrationProcessor/DataTransfer/PageReader.cs index 70ae0e2..f69281e 100644 --- a/CassandraMigrationProcessor/DataTransfer/PageReader.cs +++ b/CassandraMigrationProcessor/DataTransfer/PageReader.cs @@ -21,8 +21,7 @@ internal record ReaderConfig(int PageSize, int MaxReadRetries, bool PreserveCell /// source session is keyspace-agnostic; per-table state (columns, /// identifiers, UDT registrations) is resolved from /// at read time. UDT registration is -/// cached per keyspace so the first partition for each table pays the -/// cost and subsequent partitions reuse it. +/// cached job-wide per physical session and keyspace. /// internal class PageReader { @@ -48,7 +47,7 @@ internal class PageReader // hints parking a worker for minutes. private const int MaxRetryDelayMs = 30_000; - private PageReader( + public PageReader( WorkerLog log, SourceSessionWrapper sourceSession, ReaderConfig config, @@ -64,18 +63,6 @@ private PageReader( ?? throw new ArgumentNullException(nameof(sourceSession)); } - public static Task CreateAsync(WorkerLog log, - SourceSessionWrapper sourceSession, - ReaderConfig config, - CancellationToken cancellationToken) - { - return Task.FromResult(new PageReader( - log, - sourceSession, - config, - cancellationToken)); - } - /// /// One page of source rows together with the chunk and per-row /// CDC metadata (writetime + TTL expiry). Exactly one of @@ -191,7 +178,7 @@ internal record ReadResult( // intact and will retry the same page once the source stops // throttling. var resultSet = await RetryExecutor.ExecuteOrDefaultAsync( - operation: async (_, _) => + operation: async _ => { var sourceSession = useJson ? _sourceSession.GetSession() @@ -201,15 +188,12 @@ internal record ReadResult( .WaitAsync(_ct) .ConfigureAwait(false); }, - policy: RetryPolicy.Create( - _maxReadRetries, - ex => ex is not SourceUdtRegistrationException - && ExceptionClassifier.IsTransient(ex), - (ex, attempt) => TimeSpan.FromMilliseconds( - Math.Min( - ExceptionClassifier.GetRetryDelayMs(ex, attempt), - MaxRetryDelayMs))), - onFailure: (ex, attempt) => + maxAttempts: _maxReadRetries, + shouldRetry: ex => ex is not SourceUdtRegistrationException + && ExceptionClassifier.IsTransient(ex), + delayFor: (ex, attempt) => TimeSpan.FromMilliseconds( + Math.Min(ExceptionClassifier.GetRetryDelayMs(ex, attempt), MaxRetryDelayMs)), + onRetry: (ex, attempt) => { LastRetryExhaustionException = ex; _log.WriteLine( diff --git a/CassandraMigrationProcessor/DataTransfer/PageWriter.cs b/CassandraMigrationProcessor/DataTransfer/PageWriter.cs index 756c5a8..8fb5ad9 100644 --- a/CassandraMigrationProcessor/DataTransfer/PageWriter.cs +++ b/CassandraMigrationProcessor/DataTransfer/PageWriter.cs @@ -58,7 +58,7 @@ private PageWriter(WorkerLog log, ISession targetSession, _targetSession = targetSession; } - public static async Task CreateAsync(WorkerLog log, ISessionFactory sessionFactory, WriterConfig config, CancellationToken cancellationToken) + public static async Task CreateAsync(WorkerLog log, JobSessionFactory sessionFactory, WriterConfig config, CancellationToken cancellationToken) { var targetSession = await sessionFactory.CreateSessionAsync(cancellationToken); return new PageWriter(log, targetSession, config, cancellationToken); diff --git a/CassandraMigrationProcessor/DataTransfer/PipelineContext.cs b/CassandraMigrationProcessor/DataTransfer/PipelineContext.cs index 7fea3b9..2412d38 100644 --- a/CassandraMigrationProcessor/DataTransfer/PipelineContext.cs +++ b/CassandraMigrationProcessor/DataTransfer/PipelineContext.cs @@ -15,7 +15,7 @@ namespace CassandraMigrationProcessor.DataTransfer; internal record PipelineContext( PartitionManager Partitions, SourceSessionWrapper SourceSession, - ISessionFactory SessionFactory, + JobSessionFactory SessionFactory, ReaderConfig ReaderConfig, WriterConfig WriterConfig, bool EnableReplay, diff --git a/CassandraMigrationProcessor/DataTransfer/RowWriteRetry.cs b/CassandraMigrationProcessor/DataTransfer/RowWriteRetry.cs index cfa5bef..77d6e28 100644 --- a/CassandraMigrationProcessor/DataTransfer/RowWriteRetry.cs +++ b/CassandraMigrationProcessor/DataTransfer/RowWriteRetry.cs @@ -95,42 +95,41 @@ public static async Task ExecuteRowGroupsAsync( string rowKind, CancellationToken cancellationToken) { - int attempts = 0; - try + for (int n = 1; n <= policy.MaxAttempts; n++) { - long elapsed = await RetryExecutor.ExecuteAsync( - async (attemptNumber, _) => - { - attempts = attemptNumber; - var start = Stopwatch.GetTimestamp(); - await attempt().ConfigureAwait(false); - return (Stopwatch.GetTimestamp() - start) - * 1000 / Stopwatch.Frequency; - }, - policy, - cancellationToken: cancellationToken); - return (WriteOutcome.Success, elapsed, null); - } - catch (OperationCanceledException) - { - throw; - } - catch (Exception ex) - { - if (ExceptionClassifier.IsFatal(ex)) + cancellationToken.ThrowIfCancellationRequested(); + var start = Stopwatch.GetTimestamp(); + try { - log.WriteLine( - $"FATAL {rowKind}: {ex.GetType().Name}: {ex.Message}", - LogType.Error); - return (WriteOutcome.Fatal, 0, ex); + await attempt(); + long elapsed = (Stopwatch.GetTimestamp() - start) * 1000 / Stopwatch.Frequency; + return (WriteOutcome.Success, elapsed, null); + } + catch (OperationCanceledException) + { + throw; } + catch (Exception ex) + { + if (ExceptionClassifier.IsFatal(ex)) + { + log.WriteLine($"FATAL {rowKind}: {ex.GetType().Name}: {ex.Message}", + LogType.Error); + return (WriteOutcome.Fatal, 0, ex); + } - log.WriteLine( - $"{rowKind} FAILED after {attempts} attempt(s): " + - $"{ex.GetType().Name}: {ex.Message}", - LogType.Error); - return (WriteOutcome.Failed, 0, ex); + if (ExceptionClassifier.IsTransient(ex) && n < policy.MaxAttempts) + { + await Task.Delay(policy.DelayBeforeRetry(n), cancellationToken); + continue; + } + + log.WriteLine($"{rowKind} FAILED after {n} attempt(s): {ex.GetType().Name}: {ex.Message}", + LogType.Error); + return (WriteOutcome.Failed, 0, ex); + } } + return (WriteOutcome.Failed, 0, null); } private static void ApplyToCounters( diff --git a/CassandraMigrationProcessor/Infrastructure/RetryExecutor.cs b/CassandraMigrationProcessor/Infrastructure/RetryExecutor.cs index fe2b814..7c45aa3 100644 --- a/CassandraMigrationProcessor/Infrastructure/RetryExecutor.cs +++ b/CassandraMigrationProcessor/Infrastructure/RetryExecutor.cs @@ -1,138 +1,146 @@ namespace CassandraMigrationProcessor.Infrastructure; /// -/// Executes synchronous and asynchronous operations using caller-provided -/// retry policies. Operation-specific exception classification, delay, and -/// logging remain outside the executor. +/// Shared transient-fault retry helper. Wraps an async operation with +/// linear backoff retry on transient Cassandra exceptions (timeouts, +/// throttles, transport errors as classified by +/// ). Caller-agnostic — +/// nothing here is schema- or query-specific; see callers in +/// SchemaManager, PageReader, etc. /// internal static class RetryExecutor { + /// + /// Execute an async operation with retry on transient errors. + /// Delay between attempts is taken from + /// (which honours + /// server RetryAfterMs hints, applies exponential backoff + /// with jitter, and caps the per-sleep ceiling). The supplied + /// cancellation token is honoured both during the operation and + /// during the backoff sleep, so Stop observes promptly instead + /// of waiting for the next retry timer to fire. + /// public static async Task ExecuteAsync( - Func> operation, - RetryPolicy policy, - Action? onRetry = null, - Action? onFailure = null, - CancellationToken cancellationToken = default) - { - var result = await ExecuteCoreAsync( - operation, policy, onRetry, onFailure, cancellationToken) - .ConfigureAwait(false); - if (result.Succeeded) - return result.Value!; - - System.Runtime.ExceptionServices.ExceptionDispatchInfo - .Capture(result.Exception!) - .Throw(); - throw new System.Diagnostics.UnreachableException(); - } - - public static Task ExecuteAsync( Func> operation, + int maxRetries = MigrationDefaults.TransientRetryMaxAttempts, + int baseDelayMs = MigrationDefaults.TransientRetryBaseDelayMs, CancellationToken cancellationToken = default) { - return ExecuteAsync( - (_, _) => operation(), - RetryPolicy.Transient(), - cancellationToken: cancellationToken); + Exception? lastException = null; + for (int attempt = 1; attempt <= maxRetries; attempt++) + { + cancellationToken.ThrowIfCancellationRequested(); + try + { + return await operation().ConfigureAwait(false); + } + catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) + { + throw; + } + catch (Exception ex) when (attempt < maxRetries + && ExceptionClassifier.IsTransient(ex)) + { + lastException = ex; + var delay = Math.Max( + ExceptionClassifier.GetRetryDelayMs(ex, attempt), + attempt * baseDelayMs); + await Task.Delay(delay, cancellationToken) + .ConfigureAwait(false); + } + } + throw lastException ?? new TimeoutException("Operation timed out after all retries"); } + /// + /// Non-generic overload for fire-and-forget operations. + /// public static Task ExecuteAsync( Func operation, + int maxRetries = MigrationDefaults.TransientRetryMaxAttempts, + int baseDelayMs = MigrationDefaults.TransientRetryBaseDelayMs, CancellationToken cancellationToken = default) { - return ExecuteAsync( - async (_, _) => - { - await operation().ConfigureAwait(false); - return true; - }, - RetryPolicy.Transient(), - cancellationToken: cancellationToken); + return ExecuteAsync( + async () => { await operation().ConfigureAwait(false); return 0; }, + maxRetries, baseDelayMs, cancellationToken); } - public static T Execute( - Func operation, - RetryPolicy policy, - Action? onRetry = null) + /// + /// Generic overload with a caller-supplied retry predicate and + /// custom backoff function — used by sites that need to retry on a + /// narrower set than + /// (e.g. throttle-only retries for table-accessibility probes) or + /// need a different backoff curve. The cancellation token is + /// honoured both during the operation and during the sleep so Stop + /// observes promptly. + /// + public static async Task ExecuteAsync( + Func> operation, + int maxAttempts, + Predicate shouldRetry, + Func delayFor, + Action? onRetry = null, + CancellationToken cancellationToken = default) { - for (int attempt = 1; ; attempt++) + Exception? lastException = null; + for (int attempt = 1; attempt <= maxAttempts; attempt++) { + cancellationToken.ThrowIfCancellationRequested(); try { - return operation(attempt); + return await operation(attempt).ConfigureAwait(false); } - catch (Exception ex) when ( - attempt < policy.MaxAttempts - && policy.ShouldRetry(ex)) + catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) { + throw; + } + catch (Exception ex) when (attempt < maxAttempts && shouldRetry(ex)) + { + lastException = ex; onRetry?.Invoke(ex, attempt); - Thread.Sleep(policy.DelayBeforeRetry(ex, attempt)); + await Task.Delay(delayFor(ex, attempt), cancellationToken) + .ConfigureAwait(false); } } + throw lastException ?? new TimeoutException("Operation timed out after all retries"); } + /// + /// Variant that returns default(T) instead of throwing when + /// every attempt fails on a retryable exception. Used by PageReader + /// where read-retry exhaustion must surface as "no page; re-queue + /// via cooldown" rather than as a thrown error. Non-retryable + /// exceptions still propagate. + /// public static async Task ExecuteOrDefaultAsync( - Func> operation, - RetryPolicy policy, + Func> operation, + int maxAttempts, + Predicate shouldRetry, + Func delayFor, Action? onRetry = null, - Action? onFailure = null, CancellationToken cancellationToken = default) where T : class { - var result = await ExecuteCoreAsync( - operation, policy, onRetry, onFailure, cancellationToken) - .ConfigureAwait(false); - return result.Succeeded ? result.Value : null; - } - - private static async Task> ExecuteCoreAsync( - Func> operation, - RetryPolicy policy, - Action? onRetry, - Action? onFailure, - CancellationToken cancellationToken) - { - for (int attempt = 1; ; attempt++) + for (int attempt = 1; attempt <= maxAttempts; attempt++) { cancellationToken.ThrowIfCancellationRequested(); try { - var value = await operation(attempt, cancellationToken) - .ConfigureAwait(false); - return RetryResult.Success(value); + return await operation(attempt).ConfigureAwait(false); } catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) { throw; } - catch (Exception ex) when (policy.ShouldRetry(ex)) + catch (Exception ex) when (shouldRetry(ex)) { - onFailure?.Invoke(ex, attempt); - if (attempt >= policy.MaxAttempts) - return RetryResult.Failure(ex); - onRetry?.Invoke(ex, attempt); - await Task.Delay( - policy.DelayBeforeRetry(ex, attempt), - cancellationToken) + if (attempt >= maxAttempts) return null; + await Task.Delay(delayFor(ex, attempt), cancellationToken) .ConfigureAwait(false); } } - } - - private readonly record struct RetryResult( - bool Succeeded, - T? Value, - Exception? Exception) - { - public static RetryResult Success(T value) - { - return new RetryResult(true, value, null); - } - - public static RetryResult Failure(Exception exception) - { - return new RetryResult(false, default, exception); - } + return null; } } diff --git a/CassandraMigrationProcessor/Infrastructure/RetryPolicy.cs b/CassandraMigrationProcessor/Infrastructure/RetryPolicy.cs index 13ea349..04795b5 100644 --- a/CassandraMigrationProcessor/Infrastructure/RetryPolicy.cs +++ b/CassandraMigrationProcessor/Infrastructure/RetryPolicy.cs @@ -16,26 +16,14 @@ namespace CassandraMigrationProcessor.Infrastructure; public sealed class RetryPolicy { public int MaxAttempts { get; } - private readonly Predicate _shouldRetry; - private readonly Func _delayFor; + private readonly Func _delayFor; - private RetryPolicy( - int maxAttempts, - Predicate shouldRetry, - Func delayFor) + private RetryPolicy(int maxAttempts, Func delayFor) { if (maxAttempts < 1) throw new ArgumentOutOfRangeException(nameof(maxAttempts), "Must be >= 1."); MaxAttempts = maxAttempts; - _shouldRetry = shouldRetry - ?? throw new ArgumentNullException(nameof(shouldRetry)); - _delayFor = delayFor - ?? throw new ArgumentNullException(nameof(delayFor)); - } - - public bool ShouldRetry(Exception exception) - { - return _shouldRetry(exception); + _delayFor = delayFor; } /// @@ -43,71 +31,26 @@ public bool ShouldRetry(Exception exception) /// (1-based, so the delay returned for attempt=1 is the wait /// between the first failed attempt and the second attempt). /// - public TimeSpan DelayBeforeRetry(Exception exception, int attempt) - { - return _delayFor(exception, attempt); - } - - public static RetryPolicy Create( - int maxAttempts, - Predicate shouldRetry, - Func delayFor) - { - return new RetryPolicy(maxAttempts, shouldRetry, delayFor); - } + public TimeSpan DelayBeforeRetry(int attempt) => _delayFor(attempt); /// /// Linear backoff: baseDelay × attempt. Matches the historical /// behaviour of the bulk-copy retry loop. /// - public static RetryPolicy Linear( - int maxAttempts, - TimeSpan baseDelay, - Predicate? shouldRetry = null) - { - return new RetryPolicy( - maxAttempts, - shouldRetry ?? ExceptionClassifier.IsTransient, - (_, attempt) => TimeSpan.FromMilliseconds( - baseDelay.TotalMilliseconds * attempt)); - } + public static RetryPolicy Linear(int maxAttempts, TimeSpan baseDelay) + => new(maxAttempts, attempt => TimeSpan.FromMilliseconds(baseDelay.TotalMilliseconds * attempt)); /// /// Capped exponential backoff: baseDelay × 2^(attempt-1), /// clamped to . Suitable when a retry storm is /// likely (target throttling, e.g. 429s). /// - public static RetryPolicy Exponential( - int maxAttempts, - TimeSpan baseDelay, - TimeSpan cap, - Predicate? shouldRetry = null) - { - return new RetryPolicy( - maxAttempts, - shouldRetry ?? ExceptionClassifier.IsTransient, - (_, attempt) => - { - var ms = baseDelay.TotalMilliseconds - * Math.Pow(2, attempt - 1); - return ms >= cap.TotalMilliseconds - ? cap - : TimeSpan.FromMilliseconds(ms); - }); - } - - public static RetryPolicy Transient( - int maxAttempts = MigrationDefaults.TransientRetryMaxAttempts, - int minimumDelayMs = MigrationDefaults.TransientRetryBaseDelayMs) - { - return new RetryPolicy( - maxAttempts, - ExceptionClassifier.IsTransient, - (exception, attempt) => TimeSpan.FromMilliseconds( - Math.Max( - ExceptionClassifier.GetRetryDelayMs(exception, attempt), - attempt * minimumDelayMs))); - } + public static RetryPolicy Exponential(int maxAttempts, TimeSpan baseDelay, TimeSpan cap) + => new(maxAttempts, attempt => + { + var ms = baseDelay.TotalMilliseconds * Math.Pow(2, attempt - 1); + return ms >= cap.TotalMilliseconds ? cap : TimeSpan.FromMilliseconds(ms); + }); /// /// Single-shot delay computed from a server-hinted exception via @@ -117,8 +60,5 @@ public static RetryPolicy Transient( /// Returned as a ready for Task.Delay. /// public static TimeSpan FromException(Exception ex, int attempt) - { - return TimeSpan.FromMilliseconds( - ExceptionClassifier.GetRetryDelayMs(ex, attempt)); - } + => TimeSpan.FromMilliseconds(ExceptionClassifier.GetRetryDelayMs(ex, attempt)); } diff --git a/CassandraMigrationProcessor/Models/TableCopySpec.cs b/CassandraMigrationProcessor/Models/TableCopySpec.cs index 5d3ec94..704f5e0 100644 --- a/CassandraMigrationProcessor/Models/TableCopySpec.cs +++ b/CassandraMigrationProcessor/Models/TableCopySpec.cs @@ -4,7 +4,7 @@ namespace CassandraMigrationProcessor.Models; /// Immutable description of a single table copy. Identifies the source /// and target keyspace/table; runtime sessions are not threaded through /// here — readers use the job-wide source session and writers open -/// worker-owned sessions through ISessionFactory. +/// worker-owned sessions through JobSessionFactory. /// public record TableCopySpec( string KeyspaceName, From b96e1b3f430e7e8dbcb5125cf3a03169b32501bc Mon Sep 17 00:00:00 2001 From: Nitesh Vijay Date: Wed, 19 Aug 2026 17:21:32 +0530 Subject: [PATCH 23/32] Group source session lifecycle state Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 078ac7b9-aa7f-44a1-9053-44c346cec5e1 --- .../CassandraDriver/SourceSessionWrapper.cs | 186 ++++++++---------- .../DataTransfer/MigrationJobRunner.cs | 3 +- 2 files changed, 78 insertions(+), 111 deletions(-) diff --git a/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs b/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs index 9e04246..8855b2f 100644 --- a/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs +++ b/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs @@ -25,57 +25,39 @@ internal sealed class SourceSessionWrapper : IDisposable TimeSpan.FromMinutes(10); private const int MaxRefreshFailures = 6; - private readonly object _sync = new(); - private readonly object _refreshLock = new(); private readonly MigrationLog _log; private readonly SourceSessionSettings _settings; - private readonly HashSet _retiredSessions = - new(ReferenceEqualityComparer.Instance); - private readonly ConcurrentDictionary<(ISession Session, string Keyspace), Lazy> - _udtRegistrations = new(); - private ISession? _currentSession; - private Timer? _tokenRefreshTimer; - private DateTime _tokenExpiresAt = DateTime.MinValue; - private int _consecutiveRefreshFailures; - private bool _tokenRefreshEnabled; - private bool _tokenRefreshDisposed; - private bool _disposed; - - private SourceSessionWrapper( - MigrationLog log, - SourceSessionSettings settings) - { - _log = log ?? throw new ArgumentNullException(nameof(log)); - _settings = settings - ?? throw new ArgumentNullException(nameof(settings)); - } + private readonly SessionState _sessions = new(); + private readonly TokenRefreshState _tokenRefresh = new(); - public static SourceSessionWrapper Create( + public SourceSessionWrapper( MigrationLog log, Job job, int workerCount) { + _log = log ?? throw new ArgumentNullException(nameof(log)); var source = CassandraClientFactory.ResolveSourceSession( job, workerCount); - var wrapper = new SourceSessionWrapper(log, source.Settings); + _settings = source.Settings; try { - wrapper.Initialize(source.Credential); - return wrapper; + _sessions.Current = CreateSession(source.Credential); + if (source.Credential.Length > 200) + StartTokenRefresh(source.Credential); } catch { - wrapper.Dispose(); + Dispose(); throw; } } public ISession GetSession() { - lock (_sync) + lock (_sessions.Sync) { - ObjectDisposedException.ThrowIf(_disposed, this); - return _currentSession + ObjectDisposedException.ThrowIf(_sessions.Disposed, this); + return _sessions.Current ?? throw new InvalidOperationException("The session provider has not been initialized."); } } @@ -84,7 +66,7 @@ public async Task GetTypedSessionAsync(string keyspace) { var session = GetSession(); var key = (Session: session, Keyspace: keyspace); - var registration = _udtRegistrations.GetOrAdd( + var registration = _sessions.UdtRegistrations.GetOrAdd( key, key => new Lazy( () => RegisterUdtsAsync(key.Session, key.Keyspace), @@ -96,7 +78,7 @@ public async Task GetTypedSessionAsync(string keyspace) catch (Exception ex) { ((ICollection>>) - _udtRegistrations).Remove(new KeyValuePair< + _sessions.UdtRegistrations).Remove(new KeyValuePair< (ISession Session, string Keyspace), Lazy>( key, registration)); throw new SourceUdtRegistrationException(keyspace, ex); @@ -114,32 +96,6 @@ await DynamicUdtRegistrar.RegisterAsync( session, keyspace, allUdts); } - public void Initialize(string credential) - { - ArgumentException.ThrowIfNullOrWhiteSpace(credential); - - var session = CreateSession(credential); - try - { - lock (_sync) - { - ObjectDisposedException.ThrowIf(_disposed, this); - if (_currentSession != null) - throw new InvalidOperationException("The session provider is already initialized."); - _currentSession = session; - } - } - catch - { - MigrationUtilities.SafeDisposeSession( - session, "Unpublished initial session"); - throw; - } - - if (credential.Length > 200) - StartTokenRefresh(credential); - } - private void Refresh(string credential) { ArgumentException.ThrowIfNullOrWhiteSpace(credential); @@ -148,15 +104,15 @@ private void Refresh(string credential) ISession? retiredSession; try { - lock (_sync) + lock (_sessions.Sync) { - ObjectDisposedException.ThrowIf(_disposed, this); - if (_currentSession == null) + ObjectDisposedException.ThrowIf(_sessions.Disposed, this); + if (_sessions.Current == null) throw new InvalidOperationException("The session provider has not been initialized."); - retiredSession = _currentSession; - _currentSession = session; - _retiredSessions.Add(retiredSession); + retiredSession = _sessions.Current; + _sessions.Current = session; + _sessions.Retired.Add(retiredSession); } } catch @@ -185,9 +141,9 @@ private async Task DisposeRetiredSessionAfterDelayAsync(ISession session) await Task.Delay(RetiredSessionDisposalDelay).ConfigureAwait(false); bool shouldDispose; - lock (_sync) + lock (_sessions.Sync) { - shouldDispose = _retiredSessions.Remove(session); + shouldDispose = _sessions.Retired.Remove(session); } if (shouldDispose) @@ -200,10 +156,10 @@ private async Task DisposeRetiredSessionAfterDelayAsync(ISession session) private void RemoveUdtRegistrations(ISession session) { - foreach (var key in _udtRegistrations.Keys) + foreach (var key in _sessions.UdtRegistrations.Keys) { if (ReferenceEquals(key.Session, session)) - _udtRegistrations.TryRemove(key, out _); + _sessions.UdtRegistrations.TryRemove(key, out _); } } @@ -229,10 +185,9 @@ private static DateTime GetTokenExpiry(string token) private void StartTokenRefresh(string currentToken) { - lock (_refreshLock) + lock (_tokenRefresh.Sync) { - if (_tokenRefreshDisposed) return; - _tokenRefreshEnabled = true; + if (_tokenRefresh.Disposed) return; ScheduleTokenRefresh(currentToken); } } @@ -245,57 +200,57 @@ private void ScheduleTokenRefresh(string currentToken) if (expiry == DateTime.MaxValue) expiry = DateTime.UtcNow.AddMinutes(50); - _tokenExpiresAt = expiry; + _tokenRefresh.ExpiresAt = expiry; TimeSpan delay = expiry - DateTime.UtcNow - TimeSpan.FromMinutes(5); if (delay < TimeSpan.FromMinutes(1)) delay = TimeSpan.FromMinutes(1); - _tokenRefreshTimer = new Timer( + _tokenRefresh.Timer = new Timer( RefreshTokenCallback, null, delay, Timeout.InfiniteTimeSpan); } private void RefreshTokenCallback(object? state) { - lock (_refreshLock) + lock (_tokenRefresh.Sync) { - if (_tokenRefreshDisposed || !_tokenRefreshEnabled) return; + if (_tokenRefresh.Disposed || _tokenRefresh.Timer == null) return; try { string freshToken = CassandraClientFactory.AcquireAadToken(); Refresh(freshToken); - _consecutiveRefreshFailures = 0; + _tokenRefresh.ConsecutiveFailures = 0; ScheduleTokenRefresh(freshToken); } catch (Exception ex) { - _consecutiveRefreshFailures++; + _tokenRefresh.ConsecutiveFailures++; int seconds = Math.Min( 300, 30 * (1 << Math.Min( - _consecutiveRefreshFailures - 1, 4))); + _tokenRefresh.ConsecutiveFailures - 1, 4))); bool tokenAlreadyExpired = - DateTime.UtcNow >= _tokenExpiresAt; + DateTime.UtcNow >= _tokenRefresh.ExpiresAt; LogType severity = - _consecutiveRefreshFailures >= MaxRefreshFailures + _tokenRefresh.ConsecutiveFailures >= MaxRefreshFailures || tokenAlreadyExpired ? LogType.Error : LogType.Warning; string message = - $"Token refresh failed (attempt {_consecutiveRefreshFailures}, " + - $"retrying in {seconds}s, tokenExpiresAt={_tokenExpiresAt:O}): " + + $"Token refresh failed (attempt {_tokenRefresh.ConsecutiveFailures}, " + + $"retrying in {seconds}s, tokenExpiresAt={_tokenRefresh.ExpiresAt:O}): " + ex.Message; Console.WriteLine($"[{severity}] {message}"); _log.WriteLine(message, severity); StopTokenRefreshCore(); - if (_tokenRefreshEnabled && !_tokenRefreshDisposed) + if (!_tokenRefresh.Disposed) { - _tokenRefreshTimer = new Timer( + _tokenRefresh.Timer = new Timer( RefreshTokenCallback, null, TimeSpan.FromSeconds(seconds), Timeout.InfiniteTimeSpan); @@ -304,42 +259,34 @@ private void RefreshTokenCallback(object? state) } } - public void StopTokenRefresh() - { - lock (_refreshLock) - { - _tokenRefreshEnabled = false; - StopTokenRefreshCore(); - } - } - private void StopTokenRefreshCore() { - _tokenRefreshTimer?.Dispose(); - _tokenRefreshTimer = null; + _tokenRefresh.Timer?.Dispose(); + _tokenRefresh.Timer = null; } public void Dispose() { - lock (_refreshLock) + lock (_tokenRefresh.Sync) { - if (_tokenRefreshDisposed) return; - _tokenRefreshDisposed = true; - _tokenRefreshEnabled = false; - StopTokenRefreshCore(); + if (!_tokenRefresh.Disposed) + { + _tokenRefresh.Disposed = true; + StopTokenRefreshCore(); + } } List sessionsToDispose; - lock (_sync) + lock (_sessions.Sync) { - if (_disposed) return; - _disposed = true; - sessionsToDispose = _retiredSessions.ToList(); - _retiredSessions.Clear(); - if (_currentSession != null) - sessionsToDispose.Add(_currentSession); - _currentSession = null; - _udtRegistrations.Clear(); + if (_sessions.Disposed) return; + _sessions.Disposed = true; + sessionsToDispose = _sessions.Retired.ToList(); + _sessions.Retired.Clear(); + if (_sessions.Current != null) + sessionsToDispose.Add(_sessions.Current); + _sessions.Current = null; + _sessions.UdtRegistrations.Clear(); } foreach (var session in sessionsToDispose) @@ -348,4 +295,25 @@ public void Dispose() session, "Source session wrapper"); } } + + private sealed class SessionState + { + public object Sync { get; } = new(); + public HashSet Retired { get; } = + new(ReferenceEqualityComparer.Instance); + public ConcurrentDictionary< + (ISession Session, string Keyspace), + Lazy> UdtRegistrations { get; } = new(); + public ISession? Current { get; set; } + public bool Disposed { get; set; } + } + + private sealed class TokenRefreshState + { + public object Sync { get; } = new(); + public Timer? Timer { get; set; } + public DateTime ExpiresAt { get; set; } = DateTime.MinValue; + public int ConsecutiveFailures { get; set; } + public bool Disposed { get; set; } + } } diff --git a/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs b/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs index eabba58..c7719ab 100644 --- a/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs +++ b/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs @@ -82,7 +82,7 @@ public static async Task CreateAsync( ISession? target = null; try { - sourceSessions = SourceSessionWrapper.Create( + sourceSessions = new SourceSessionWrapper( log, job, pipelineConfig.WorkerCount); target = await CassandraClientFactory.CreateTargetSessionAsync(log, job); return new MigrationJobRunner( @@ -855,7 +855,6 @@ public void Stop() // without waiting for the outer Task to observe the cancel. MigrationUtilities.SafeDispose(_pipeline, "JobPipeline (Stop)"); _pipeline = null; - _sourceSessions.StopTokenRefresh(); } /// From 3971db535d39f9ac203711acf6a58b4a1756f7c3 Mon Sep 17 00:00:00 2001 From: Nitesh Vijay Date: Wed, 19 Aug 2026 17:26:47 +0530 Subject: [PATCH 24/32] Flatten source session wrapper state Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 078ac7b9-aa7f-44a1-9053-44c346cec5e1 --- .../CassandraDriver/SourceSessionWrapper.cs | 143 +++++++----------- 1 file changed, 55 insertions(+), 88 deletions(-) diff --git a/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs b/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs index 8855b2f..7596524 100644 --- a/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs +++ b/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs @@ -25,10 +25,18 @@ internal sealed class SourceSessionWrapper : IDisposable TimeSpan.FromMinutes(10); private const int MaxRefreshFailures = 6; + private readonly object _sync = new(); + private readonly object _refreshLock = new(); private readonly MigrationLog _log; private readonly SourceSessionSettings _settings; - private readonly SessionState _sessions = new(); - private readonly TokenRefreshState _tokenRefresh = new(); + private readonly HashSet _retiredSessions = + new(ReferenceEqualityComparer.Instance); + private readonly ConcurrentDictionary<(ISession Session, string Keyspace), Lazy> + _udtRegistrations = new(); + private ISession? _currentSession; + private Timer? _tokenRefreshTimer; + private DateTime _tokenExpiresAt = DateTime.MinValue; + private int _consecutiveRefreshFailures; public SourceSessionWrapper( MigrationLog log, @@ -39,26 +47,17 @@ public SourceSessionWrapper( var source = CassandraClientFactory.ResolveSourceSession( job, workerCount); _settings = source.Settings; - try - { - _sessions.Current = CreateSession(source.Credential); - if (source.Credential.Length > 200) - StartTokenRefresh(source.Credential); - } - catch - { - Dispose(); - throw; - } + _currentSession = CreateSession(source.Credential); + if (source.Credential.Length > 200) + StartTokenRefresh(source.Credential); } public ISession GetSession() { - lock (_sessions.Sync) + lock (_sync) { - ObjectDisposedException.ThrowIf(_sessions.Disposed, this); - return _sessions.Current - ?? throw new InvalidOperationException("The session provider has not been initialized."); + return _currentSession + ?? throw new ObjectDisposedException(nameof(SourceSessionWrapper)); } } @@ -66,7 +65,7 @@ public async Task GetTypedSessionAsync(string keyspace) { var session = GetSession(); var key = (Session: session, Keyspace: keyspace); - var registration = _sessions.UdtRegistrations.GetOrAdd( + var registration = _udtRegistrations.GetOrAdd( key, key => new Lazy( () => RegisterUdtsAsync(key.Session, key.Keyspace), @@ -78,7 +77,7 @@ public async Task GetTypedSessionAsync(string keyspace) catch (Exception ex) { ((ICollection>>) - _sessions.UdtRegistrations).Remove(new KeyValuePair< + _udtRegistrations).Remove(new KeyValuePair< (ISession Session, string Keyspace), Lazy>( key, registration)); throw new SourceUdtRegistrationException(keyspace, ex); @@ -104,15 +103,14 @@ private void Refresh(string credential) ISession? retiredSession; try { - lock (_sessions.Sync) + lock (_sync) { - ObjectDisposedException.ThrowIf(_sessions.Disposed, this); - if (_sessions.Current == null) - throw new InvalidOperationException("The session provider has not been initialized."); + if (_currentSession == null) + throw new ObjectDisposedException(nameof(SourceSessionWrapper)); - retiredSession = _sessions.Current; - _sessions.Current = session; - _sessions.Retired.Add(retiredSession); + retiredSession = _currentSession; + _currentSession = session; + _retiredSessions.Add(retiredSession); } } catch @@ -141,9 +139,9 @@ private async Task DisposeRetiredSessionAfterDelayAsync(ISession session) await Task.Delay(RetiredSessionDisposalDelay).ConfigureAwait(false); bool shouldDispose; - lock (_sessions.Sync) + lock (_sync) { - shouldDispose = _sessions.Retired.Remove(session); + shouldDispose = _retiredSessions.Remove(session); } if (shouldDispose) @@ -156,10 +154,10 @@ private async Task DisposeRetiredSessionAfterDelayAsync(ISession session) private void RemoveUdtRegistrations(ISession session) { - foreach (var key in _sessions.UdtRegistrations.Keys) + foreach (var key in _udtRegistrations.Keys) { if (ReferenceEquals(key.Session, session)) - _sessions.UdtRegistrations.TryRemove(key, out _); + _udtRegistrations.TryRemove(key, out _); } } @@ -185,9 +183,8 @@ private static DateTime GetTokenExpiry(string token) private void StartTokenRefresh(string currentToken) { - lock (_tokenRefresh.Sync) + lock (_refreshLock) { - if (_tokenRefresh.Disposed) return; ScheduleTokenRefresh(currentToken); } } @@ -200,93 +197,84 @@ private void ScheduleTokenRefresh(string currentToken) if (expiry == DateTime.MaxValue) expiry = DateTime.UtcNow.AddMinutes(50); - _tokenRefresh.ExpiresAt = expiry; + _tokenExpiresAt = expiry; TimeSpan delay = expiry - DateTime.UtcNow - TimeSpan.FromMinutes(5); if (delay < TimeSpan.FromMinutes(1)) delay = TimeSpan.FromMinutes(1); - _tokenRefresh.Timer = new Timer( + _tokenRefreshTimer = new Timer( RefreshTokenCallback, null, delay, Timeout.InfiniteTimeSpan); } private void RefreshTokenCallback(object? state) { - lock (_tokenRefresh.Sync) + lock (_refreshLock) { - if (_tokenRefresh.Disposed || _tokenRefresh.Timer == null) return; + if (_tokenRefreshTimer == null) return; try { string freshToken = CassandraClientFactory.AcquireAadToken(); Refresh(freshToken); - _tokenRefresh.ConsecutiveFailures = 0; + _consecutiveRefreshFailures = 0; ScheduleTokenRefresh(freshToken); } catch (Exception ex) { - _tokenRefresh.ConsecutiveFailures++; + _consecutiveRefreshFailures++; int seconds = Math.Min( 300, 30 * (1 << Math.Min( - _tokenRefresh.ConsecutiveFailures - 1, 4))); + _consecutiveRefreshFailures - 1, 4))); bool tokenAlreadyExpired = - DateTime.UtcNow >= _tokenRefresh.ExpiresAt; + DateTime.UtcNow >= _tokenExpiresAt; LogType severity = - _tokenRefresh.ConsecutiveFailures >= MaxRefreshFailures + _consecutiveRefreshFailures >= MaxRefreshFailures || tokenAlreadyExpired ? LogType.Error : LogType.Warning; string message = - $"Token refresh failed (attempt {_tokenRefresh.ConsecutiveFailures}, " + - $"retrying in {seconds}s, tokenExpiresAt={_tokenRefresh.ExpiresAt:O}): " + + $"Token refresh failed (attempt {_consecutiveRefreshFailures}, " + + $"retrying in {seconds}s, tokenExpiresAt={_tokenExpiresAt:O}): " + ex.Message; Console.WriteLine($"[{severity}] {message}"); _log.WriteLine(message, severity); StopTokenRefreshCore(); - if (!_tokenRefresh.Disposed) - { - _tokenRefresh.Timer = new Timer( - RefreshTokenCallback, null, - TimeSpan.FromSeconds(seconds), - Timeout.InfiniteTimeSpan); - } + _tokenRefreshTimer = new Timer( + RefreshTokenCallback, null, + TimeSpan.FromSeconds(seconds), + Timeout.InfiniteTimeSpan); } } } private void StopTokenRefreshCore() { - _tokenRefresh.Timer?.Dispose(); - _tokenRefresh.Timer = null; + _tokenRefreshTimer?.Dispose(); + _tokenRefreshTimer = null; } public void Dispose() { - lock (_tokenRefresh.Sync) + lock (_refreshLock) { - if (!_tokenRefresh.Disposed) - { - _tokenRefresh.Disposed = true; - StopTokenRefreshCore(); - } + StopTokenRefreshCore(); } List sessionsToDispose; - lock (_sessions.Sync) + lock (_sync) { - if (_sessions.Disposed) return; - _sessions.Disposed = true; - sessionsToDispose = _sessions.Retired.ToList(); - _sessions.Retired.Clear(); - if (_sessions.Current != null) - sessionsToDispose.Add(_sessions.Current); - _sessions.Current = null; - _sessions.UdtRegistrations.Clear(); + sessionsToDispose = _retiredSessions.ToList(); + _retiredSessions.Clear(); + if (_currentSession != null) + sessionsToDispose.Add(_currentSession); + _currentSession = null; + _udtRegistrations.Clear(); } foreach (var session in sessionsToDispose) @@ -295,25 +283,4 @@ public void Dispose() session, "Source session wrapper"); } } - - private sealed class SessionState - { - public object Sync { get; } = new(); - public HashSet Retired { get; } = - new(ReferenceEqualityComparer.Instance); - public ConcurrentDictionary< - (ISession Session, string Keyspace), - Lazy> UdtRegistrations { get; } = new(); - public ISession? Current { get; set; } - public bool Disposed { get; set; } - } - - private sealed class TokenRefreshState - { - public object Sync { get; } = new(); - public Timer? Timer { get; set; } - public DateTime ExpiresAt { get; set; } = DateTime.MinValue; - public int ConsecutiveFailures { get; set; } - public bool Disposed { get; set; } - } } From 3cdff2953e0de80fc681e4c6023f4a827da0f529 Mon Sep 17 00:00:00 2001 From: Nitesh Vijay Date: Wed, 19 Aug 2026 17:30:33 +0530 Subject: [PATCH 25/32] Use lock-free current session reads Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 078ac7b9-aa7f-44a1-9053-44c346cec5e1 --- .../CassandraDriver/SourceSessionWrapper.cs | 62 +++++++++---------- 1 file changed, 30 insertions(+), 32 deletions(-) diff --git a/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs b/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs index 7596524..8d5090d 100644 --- a/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs +++ b/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs @@ -33,10 +33,11 @@ internal sealed class SourceSessionWrapper : IDisposable new(ReferenceEqualityComparer.Instance); private readonly ConcurrentDictionary<(ISession Session, string Keyspace), Lazy> _udtRegistrations = new(); - private ISession? _currentSession; + private ISession _currentSession; private Timer? _tokenRefreshTimer; private DateTime _tokenExpiresAt = DateTime.MinValue; private int _consecutiveRefreshFailures; + private int _disposed; public SourceSessionWrapper( MigrationLog log, @@ -49,16 +50,15 @@ public SourceSessionWrapper( _settings = source.Settings; _currentSession = CreateSession(source.Credential); if (source.Credential.Length > 200) - StartTokenRefresh(source.Credential); + ScheduleTokenRefresh(source.Credential); } public ISession GetSession() { - lock (_sync) - { - return _currentSession - ?? throw new ObjectDisposedException(nameof(SourceSessionWrapper)); - } + ObjectDisposedException.ThrowIf( + Volatile.Read(ref _disposed) != 0, + this); + return Volatile.Read(ref _currentSession); } public async Task GetTypedSessionAsync(string keyspace) @@ -105,11 +105,12 @@ private void Refresh(string credential) { lock (_sync) { - if (_currentSession == null) - throw new ObjectDisposedException(nameof(SourceSessionWrapper)); + ObjectDisposedException.ThrowIf( + Volatile.Read(ref _disposed) != 0, + this); retiredSession = _currentSession; - _currentSession = session; + Volatile.Write(ref _currentSession, session); _retiredSessions.Add(retiredSession); } } @@ -181,14 +182,6 @@ private static DateTime GetTokenExpiry(string token) return DateTime.MaxValue; } - private void StartTokenRefresh(string currentToken) - { - lock (_refreshLock) - { - ScheduleTokenRefresh(currentToken); - } - } - private void ScheduleTokenRefresh(string currentToken) { StopTokenRefreshCore(); @@ -213,7 +206,9 @@ private void RefreshTokenCallback(object? state) { lock (_refreshLock) { - if (_tokenRefreshTimer == null) return; + if (Volatile.Read(ref _disposed) != 0 + || _tokenRefreshTimer == null) + return; try { @@ -245,10 +240,13 @@ private void RefreshTokenCallback(object? state) _log.WriteLine(message, severity); StopTokenRefreshCore(); - _tokenRefreshTimer = new Timer( - RefreshTokenCallback, null, - TimeSpan.FromSeconds(seconds), - Timeout.InfiniteTimeSpan); + if (Volatile.Read(ref _disposed) == 0) + { + _tokenRefreshTimer = new Timer( + RefreshTokenCallback, null, + TimeSpan.FromSeconds(seconds), + Timeout.InfiniteTimeSpan); + } } } } @@ -261,20 +259,20 @@ private void StopTokenRefreshCore() public void Dispose() { + List sessionsToDispose; lock (_refreshLock) { + if (Interlocked.Exchange(ref _disposed, 1) != 0) + return; StopTokenRefreshCore(); - } - List sessionsToDispose; - lock (_sync) - { - sessionsToDispose = _retiredSessions.ToList(); - _retiredSessions.Clear(); - if (_currentSession != null) + lock (_sync) + { + sessionsToDispose = _retiredSessions.ToList(); + _retiredSessions.Clear(); sessionsToDispose.Add(_currentSession); - _currentSession = null; - _udtRegistrations.Clear(); + _udtRegistrations.Clear(); + } } foreach (var session in sessionsToDispose) From ca6a58b13835965648b722375b95b04f7fd96a31 Mon Sep 17 00:00:00 2001 From: Nitesh Vijay Date: Thu, 20 Aug 2026 13:57:24 +0530 Subject: [PATCH 26/32] Move AAD token ownership to source wrapper Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 078ac7b9-aa7f-44a1-9053-44c346cec5e1 --- .../CassandraDriver/CassandraClientFactory.cs | 34 ++++-------------- .../CassandraDriver/SourceSessionWrapper.cs | 35 +++++++++++++++---- 2 files changed, 35 insertions(+), 34 deletions(-) diff --git a/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs b/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs index 614dfef..0ca6091 100644 --- a/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs +++ b/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs @@ -281,18 +281,7 @@ private static ISession ConnectCluster(Builder builder) } } - 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( + internal static SourceSessionSettings ResolveSourceSessionSettings( Job job, int workerCount = 0) { @@ -301,15 +290,6 @@ internal static ( bool useAad = job.SourceUseAad || string.IsNullOrEmpty(job.SourcePassword); - string credential = job.SourcePassword ?? string.Empty; - if (useAad) - { - 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; - } - string username = job.SourceUsername ?? string.Empty; if (string.IsNullOrWhiteSpace(username) && useAad) @@ -329,13 +309,11 @@ internal static ( 8); } - return ( - new SourceSessionSettings( - job.SourceContactPoint, - job.SourcePort, - username, - maxConnectionsPerHost), - credential); + return new SourceSessionSettings( + job.SourceContactPoint, + job.SourcePort, + username, + maxConnectionsPerHost); } /// diff --git a/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs b/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs index 8d5090d..d8cc21a 100644 --- a/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs +++ b/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs @@ -45,12 +45,12 @@ public SourceSessionWrapper( int workerCount) { _log = log ?? throw new ArgumentNullException(nameof(log)); - var source = CassandraClientFactory.ResolveSourceSession( + _settings = CassandraClientFactory.ResolveSourceSessionSettings( job, workerCount); - _settings = source.Settings; - _currentSession = CreateSession(source.Credential); - if (source.Credential.Length > 200) - ScheduleTokenRefresh(source.Credential); + string credential = ResolveCredential(job); + _currentSession = CreateSession(credential); + if (job.SourceUseAad) + ScheduleTokenRefresh(credential); } public ISession GetSession() @@ -182,6 +182,29 @@ private static DateTime GetTokenExpiry(string token) return DateTime.MaxValue; } + private static string ResolveCredential(Job job) + { + string credential = job.SourcePassword ?? string.Empty; + if (string.IsNullOrEmpty(credential) || job.SourceUseAad) + { + 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; + } + + return credential; + } + + private static string AcquireAadToken() + { + var credential = new Azure.Identity.DefaultAzureCredential(); + return credential.GetToken( + new Azure.Core.TokenRequestContext( + new[] { "https://cosmos.azure.com/.default" })) + .Token; + } + private void ScheduleTokenRefresh(string currentToken) { StopTokenRefreshCore(); @@ -212,7 +235,7 @@ private void RefreshTokenCallback(object? state) try { - string freshToken = CassandraClientFactory.AcquireAadToken(); + string freshToken = AcquireAadToken(); Refresh(freshToken); _consecutiveRefreshFailures = 0; From 17403a4ea29c361367968ed84bda6322a29e735a Mon Sep 17 00:00:00 2001 From: Nitesh Vijay Date: Thu, 20 Aug 2026 14:09:16 +0530 Subject: [PATCH 27/32] Require explicit source AAD mode Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 078ac7b9-aa7f-44a1-9053-44c346cec5e1 --- .../CassandraDriver/CassandraClientFactory.cs | 4 +--- .../CassandraDriver/SourceSessionWrapper.cs | 12 +++--------- 2 files changed, 4 insertions(+), 12 deletions(-) diff --git a/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs b/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs index 0ca6091..be47ef1 100644 --- a/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs +++ b/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs @@ -288,11 +288,9 @@ internal static SourceSessionSettings ResolveSourceSessionSettings( if (string.IsNullOrEmpty(job.SourceContactPoint)) throw new ArgumentException("Source contact point is required", nameof(job)); - bool useAad = job.SourceUseAad - || string.IsNullOrEmpty(job.SourcePassword); string username = job.SourceUsername ?? string.Empty; if (string.IsNullOrWhiteSpace(username) - && useAad) + && job.SourceUseAad) { username = job.SourceContactPoint .Split('.')[0]; diff --git a/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs b/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs index d8cc21a..684b49c 100644 --- a/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs +++ b/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs @@ -184,16 +184,10 @@ private static DateTime GetTokenExpiry(string token) private static string ResolveCredential(Job job) { - string credential = job.SourcePassword ?? string.Empty; - if (string.IsNullOrEmpty(credential) || job.SourceUseAad) - { - 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; - } + if (job.SourceUseAad) + return AcquireAadToken(); - return credential; + return job.SourcePassword ?? string.Empty; } private static string AcquireAadToken() From e484d1355a3a32b4b6537cf7f1951c227746cb9c Mon Sep 17 00:00:00 2001 From: Nitesh Vijay Date: Thu, 20 Aug 2026 14:12:55 +0530 Subject: [PATCH 28/32] Inline source session settings Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 078ac7b9-aa7f-44a1-9053-44c346cec5e1 --- .../CassandraDriver/CassandraClientFactory.cs | 39 --------- .../CassandraDriver/SourceSessionWrapper.cs | 83 +++++++++++-------- 2 files changed, 47 insertions(+), 75 deletions(-) diff --git a/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs b/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs index be47ef1..19ab2f0 100644 --- a/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs +++ b/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs @@ -5,12 +5,6 @@ using CassandraMigrationProcessor.Models; namespace CassandraMigrationProcessor.CassandraDriver; -internal sealed record SourceSessionSettings( - string ContactPoint, - int Port, - string Username, - int MaxConnectionsPerHost); - /// /// Creates Cassandra ISession instances for source (Cosmos DB) /// and target (OSS Cassandra) clusters. @@ -281,39 +275,6 @@ private static ISession ConnectCluster(Builder builder) } } - internal static SourceSessionSettings ResolveSourceSessionSettings( - Job job, - int workerCount = 0) - { - if (string.IsNullOrEmpty(job.SourceContactPoint)) - throw new ArgumentException("Source contact point is required", nameof(job)); - - string username = job.SourceUsername ?? string.Empty; - if (string.IsNullOrWhiteSpace(username) - && job.SourceUseAad) - { - username = job.SourceContactPoint - .Split('.')[0]; - } - - 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); - } - /// /// Per-side connection pool sizing. The per-side override /// ( / diff --git a/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs b/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs index 684b49c..7bbd535 100644 --- a/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs +++ b/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs @@ -25,8 +25,7 @@ internal sealed class SourceSessionWrapper : IDisposable TimeSpan.FromMinutes(10); private const int MaxRefreshFailures = 6; - private readonly object _sync = new(); - private readonly object _refreshLock = new(); + private readonly object _lifecycleLock = new(); private readonly MigrationLog _log; private readonly SourceSessionSettings _settings; private readonly HashSet _retiredSessions = @@ -45,8 +44,36 @@ public SourceSessionWrapper( int workerCount) { _log = log ?? throw new ArgumentNullException(nameof(log)); - _settings = CassandraClientFactory.ResolveSourceSessionSettings( - job, workerCount); + ArgumentNullException.ThrowIfNull(job); + if (string.IsNullOrEmpty(job.SourceContactPoint)) + throw new ArgumentException( + "Source contact point is required", + nameof(job)); + + string username = job.SourceUsername ?? string.Empty; + if (string.IsNullOrWhiteSpace(username) + && job.SourceUseAad) + { + username = job.SourceContactPoint.Split('.')[0]; + } + + int maxConnectionsPerHost = + CassandraClientFactory.ResolveMaxConnectionsPerHost( + job.SourceMaxConnectionsPerHost, + job.MaxConnectionsPerHost); + if (maxConnectionsPerHost == 0 && workerCount > 0) + { + maxConnectionsPerHost = Math.Clamp( + (workerCount + 31) / 32, + 2, + 8); + } + + _settings = new SourceSessionSettings( + job.SourceContactPoint, + job.SourcePort, + username, + maxConnectionsPerHost); string credential = ResolveCredential(job); _currentSession = CreateSession(credential); if (job.SourceUseAad) @@ -97,29 +124,10 @@ await DynamicUdtRegistrar.RegisterAsync( private void Refresh(string credential) { - ArgumentException.ThrowIfNullOrWhiteSpace(credential); - var session = CreateSession(credential); - ISession? retiredSession; - try - { - lock (_sync) - { - ObjectDisposedException.ThrowIf( - Volatile.Read(ref _disposed) != 0, - this); - - retiredSession = _currentSession; - Volatile.Write(ref _currentSession, session); - _retiredSessions.Add(retiredSession); - } - } - catch - { - MigrationUtilities.SafeDisposeSession( - session, "Unpublished refreshed session"); - throw; - } + var retiredSession = _currentSession; + Volatile.Write(ref _currentSession, session); + _retiredSessions.Add(retiredSession); _ = DisposeRetiredSessionAfterDelayAsync(retiredSession); } @@ -140,7 +148,7 @@ private async Task DisposeRetiredSessionAfterDelayAsync(ISession session) await Task.Delay(RetiredSessionDisposalDelay).ConfigureAwait(false); bool shouldDispose; - lock (_sync) + lock (_lifecycleLock) { shouldDispose = _retiredSessions.Remove(session); } @@ -221,7 +229,7 @@ private void ScheduleTokenRefresh(string currentToken) private void RefreshTokenCallback(object? state) { - lock (_refreshLock) + lock (_lifecycleLock) { if (Volatile.Read(ref _disposed) != 0 || _tokenRefreshTimer == null) @@ -277,19 +285,16 @@ private void StopTokenRefreshCore() public void Dispose() { List sessionsToDispose; - lock (_refreshLock) + lock (_lifecycleLock) { if (Interlocked.Exchange(ref _disposed, 1) != 0) return; StopTokenRefreshCore(); - lock (_sync) - { - sessionsToDispose = _retiredSessions.ToList(); - _retiredSessions.Clear(); - sessionsToDispose.Add(_currentSession); - _udtRegistrations.Clear(); - } + sessionsToDispose = _retiredSessions.ToList(); + _retiredSessions.Clear(); + sessionsToDispose.Add(_currentSession); + _udtRegistrations.Clear(); } foreach (var session in sessionsToDispose) @@ -298,4 +303,10 @@ public void Dispose() session, "Source session wrapper"); } } + + private sealed record SourceSessionSettings( + string ContactPoint, + int Port, + string Username, + int MaxConnectionsPerHost); } From e830545690a73f5c193cde74da2e007af6547649 Mon Sep 17 00:00:00 2001 From: Nitesh Vijay Date: Fri, 21 Aug 2026 11:32:45 +0530 Subject: [PATCH 29/32] Fail jobs on hidden operational errors Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 078ac7b9-aa7f-44a1-9053-44c346cec5e1 --- .../CassandraDriver/ArmCredentialDiscovery.cs | 13 ++- .../CassandraDriver/CassandraClientFactory.cs | 7 +- .../CassandraDriver/SchemaManager.cs | 25 +++--- .../CassandraDriver/SourceSessionWrapper.cs | 85 +++++++------------ .../Context/JobStore.cs | 45 +++++----- .../Context/JsonStore.cs | 15 ++-- .../Context/MigrationJobContext.cs | 13 ++- .../Context/UnitStore.cs | 74 +++++++--------- .../DataTransfer/MigrationJobRunner.cs | 29 ++++--- .../Infrastructure/TableDiscovery.cs | 10 ++- .../Infrastructure/TableMigrationMapper.cs | 53 ++++++------ .../Models/TableMigration.cs | 2 +- .../Persistence/DiskPersistence.cs | 48 +++++------ .../Service/JobManager.cs | 11 +-- 14 files changed, 197 insertions(+), 233 deletions(-) diff --git a/CassandraMigrationProcessor/CassandraDriver/ArmCredentialDiscovery.cs b/CassandraMigrationProcessor/CassandraDriver/ArmCredentialDiscovery.cs index 4a9cf79..3e2e462 100644 --- a/CassandraMigrationProcessor/CassandraDriver/ArmCredentialDiscovery.cs +++ b/CassandraMigrationProcessor/CassandraDriver/ArmCredentialDiscovery.cs @@ -96,7 +96,10 @@ internal class ArmCredentialResult $"sleeping {retryAfter.TotalSeconds:F1}s " + $"(attempt {attempt}/{ThrottleRetries})."); resp.Dispose(); - if (attempt == ThrottleRetries) return null; + if (attempt == ThrottleRetries) + throw new InvalidOperationException( + $"ARM ({context}) remained throttled after " + + $"{ThrottleRetries} attempts."); await Task.Delay(retryAfter); continue; @@ -265,7 +268,9 @@ internal static async Task DiscoverTargetCredentialsViaArm( } catch (Exception ex) { - Console.Error.WriteLine($"ARM discovery: {ex.Message}"); + throw new InvalidOperationException( + "ARM target credential discovery failed.", + ex); } return null; } @@ -351,7 +356,9 @@ internal static async Task DiscoverTargetCredentialsViaArm( } catch (Exception ex) { - Console.Error.WriteLine($"ARM discovery: {ex.Message}"); + throw new InvalidOperationException( + "ARM target credential discovery failed.", + ex); } return null; } diff --git a/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs b/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs index 19ab2f0..0127398 100644 --- a/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs +++ b/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs @@ -328,7 +328,12 @@ public static async Task CreateTargetSessionAsync( } catch (Exception ex) { - MigrationLog?.WriteLine($"ARM credential discovery failed: {ex.Message}", LogType.Debug); + MigrationLog?.WriteLine( + $"ARM target credential discovery failed: {ex.Message}", + LogType.Error); + throw new InvalidOperationException( + "ARM target credential discovery failed.", + ex); } } diff --git a/CassandraMigrationProcessor/CassandraDriver/SchemaManager.cs b/CassandraMigrationProcessor/CassandraDriver/SchemaManager.cs index dc608a6..f49e235 100644 --- a/CassandraMigrationProcessor/CassandraDriver/SchemaManager.cs +++ b/CassandraMigrationProcessor/CassandraDriver/SchemaManager.cs @@ -662,10 +662,11 @@ private static async Task { log?.WriteLine( $"Could not read replication for source keyspace " + - $"\"{sourceKeyspace}\" ({ex.GetType().Name}: {ex.Message}); " + - $"falling back to SimpleStrategy default.", - LogType.Warning); - return new KeyspaceReplicationInfo(null, null, null); + $"\"{sourceKeyspace}\" ({ex.GetType().Name}: {ex.Message}).", + LogType.Error); + throw new InvalidOperationException( + $"Failed to read replication for source keyspace '{sourceKeyspace}'.", + ex); } } @@ -697,7 +698,7 @@ private static async Task> GetTargetDataCentersAsync(ISession se if (!string.IsNullOrWhiteSpace(dc)) dcs.Add(dc); } } - catch + catch (InvalidQueryException) { // Targets that do not expose system.local/system.peers // (or reject the query) fall through to single-DC @@ -1153,10 +1154,12 @@ private static async Task catch (Exception ex) { log?.WriteLine( - $"[Schema] {keyspace}.{table}: failed to read source table options ({ex.GetType().Name}: {ex.Message}); " + - $"target table will use distribution defaults for TTL / gc_grace / compaction / compression / caching.", - LogType.Warning); - return new ForwardableTableOptions(string.Empty, Array.Empty()); + $"[Schema] {keyspace}.{table}: failed to read source table options " + + $"({ex.GetType().Name}: {ex.Message}).", + LogType.Error); + throw new InvalidOperationException( + $"Failed to read source table options for '{keyspace}.{table}'.", + ex); } } @@ -1189,7 +1192,7 @@ private static T TryGet(Row row, string column, T fallback) var v = row.GetValue(column); return v is null ? fallback : v; } - catch { return fallback; } + catch (ArgumentException) { return fallback; } } private static bool RowHasNonEmptyMap(Row row, string column) @@ -1199,6 +1202,6 @@ private static bool RowHasNonEmptyMap(Row row, string column) var map = row.GetValue>(column); return map != null && map.Count > 0; } - catch { return false; } + catch (ArgumentException) { return false; } } } diff --git a/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs b/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs index 7bbd535..181a91f 100644 --- a/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs +++ b/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs @@ -23,10 +23,10 @@ internal sealed class SourceSessionWrapper : IDisposable { private static readonly TimeSpan RetiredSessionDisposalDelay = TimeSpan.FromMinutes(10); - private const int MaxRefreshFailures = 6; private readonly object _lifecycleLock = new(); private readonly MigrationLog _log; + private readonly Action _reportFatalFailure; private readonly SourceSessionSettings _settings; private readonly HashSet _retiredSessions = new(ReferenceEqualityComparer.Instance); @@ -34,16 +34,17 @@ internal sealed class SourceSessionWrapper : IDisposable _udtRegistrations = new(); private ISession _currentSession; private Timer? _tokenRefreshTimer; - private DateTime _tokenExpiresAt = DateTime.MinValue; - private int _consecutiveRefreshFailures; private int _disposed; public SourceSessionWrapper( MigrationLog log, Job job, - int workerCount) + int workerCount, + Action reportFatalFailure) { _log = log ?? throw new ArgumentNullException(nameof(log)); + _reportFatalFailure = reportFatalFailure + ?? throw new ArgumentNullException(nameof(reportFatalFailure)); ArgumentNullException.ThrowIfNull(job); if (string.IsNullOrEmpty(job.SourceContactPoint)) throw new ArgumentException( @@ -77,7 +78,7 @@ public SourceSessionWrapper( string credential = ResolveCredential(job); _currentSession = CreateSession(credential); if (job.SourceUseAad) - ScheduleTokenRefresh(credential); + ScheduleTokenRefresh(GetTokenExpiry(credential)); } public ISession GetSession() @@ -172,22 +173,21 @@ private void RemoveUdtRegistrations(ISession session) private static DateTime GetTokenExpiry(string token) { - try - { - var handler = new JwtSecurityTokenHandler(); - if (handler.CanReadToken(token)) - { - var jwt = handler.ReadJwtToken(token); - return jwt.ValidTo; - } - } - catch (Exception ex) - { - Console.WriteLine( - $"[Warning] Failed to read AAD token expiry: {ex.Message}"); - } + if (string.IsNullOrWhiteSpace(token)) + throw new InvalidOperationException( + "AAD token acquisition returned an empty token."); - return DateTime.MaxValue; + var handler = new JwtSecurityTokenHandler(); + if (!handler.CanReadToken(token)) + throw new InvalidOperationException( + "AAD token acquisition returned a token that is not a readable JWT."); + + var expiry = handler.ReadJwtToken(token).ValidTo; + if (expiry == DateTime.MinValue) + throw new InvalidOperationException( + "AAD token does not contain a valid expiration time."); + + return expiry; } private static string ResolveCredential(Job job) @@ -207,16 +207,10 @@ private static string AcquireAadToken() .Token; } - private void ScheduleTokenRefresh(string currentToken) + private void ScheduleTokenRefresh(DateTime expiry) { StopTokenRefreshCore(); - DateTime expiry = GetTokenExpiry(currentToken); - if (expiry == DateTime.MaxValue) - expiry = DateTime.UtcNow.AddMinutes(50); - - _tokenExpiresAt = expiry; - TimeSpan delay = expiry - DateTime.UtcNow - TimeSpan.FromMinutes(5); if (delay < TimeSpan.FromMinutes(1)) @@ -229,6 +223,7 @@ private void ScheduleTokenRefresh(string currentToken) private void RefreshTokenCallback(object? state) { + Exception? fatalFailure = null; lock (_lifecycleLock) { if (Volatile.Read(ref _disposed) != 0 @@ -238,42 +233,22 @@ private void RefreshTokenCallback(object? state) try { string freshToken = AcquireAadToken(); + DateTime expiry = GetTokenExpiry(freshToken); Refresh(freshToken); - - _consecutiveRefreshFailures = 0; - ScheduleTokenRefresh(freshToken); + ScheduleTokenRefresh(expiry); } catch (Exception ex) { - _consecutiveRefreshFailures++; - int seconds = Math.Min( - 300, - 30 * (1 << Math.Min( - _consecutiveRefreshFailures - 1, 4))); - bool tokenAlreadyExpired = - DateTime.UtcNow >= _tokenExpiresAt; - LogType severity = - _consecutiveRefreshFailures >= MaxRefreshFailures - || tokenAlreadyExpired - ? LogType.Error - : LogType.Warning; string message = - $"Token refresh failed (attempt {_consecutiveRefreshFailures}, " + - $"retrying in {seconds}s, tokenExpiresAt={_tokenExpiresAt:O}): " + - ex.Message; - Console.WriteLine($"[{severity}] {message}"); - _log.WriteLine(message, severity); - + $"AAD token refresh failed. Aborting migration job: {ex.Message}"; + _log.WriteLine(message, LogType.Error); StopTokenRefreshCore(); - if (Volatile.Read(ref _disposed) == 0) - { - _tokenRefreshTimer = new Timer( - RefreshTokenCallback, null, - TimeSpan.FromSeconds(seconds), - Timeout.InfiniteTimeSpan); - } + fatalFailure = new InvalidOperationException(message, ex); } } + + if (fatalFailure != null) + _reportFatalFailure(fatalFailure); } private void StopTokenRefreshCore() diff --git a/CassandraMigrationProcessor/Context/JobStore.cs b/CassandraMigrationProcessor/Context/JobStore.cs index d61613a..ea9bd74 100644 --- a/CassandraMigrationProcessor/Context/JobStore.cs +++ b/CassandraMigrationProcessor/Context/JobStore.cs @@ -1,5 +1,4 @@ using System.Collections.Concurrent; -using CassandraMigrationProcessor.Infrastructure; using CassandraMigrationProcessor.Models; namespace CassandraMigrationProcessor.Context; @@ -71,15 +70,12 @@ private static void SerializeAndPersist(Job job) if (_jobs.TryGetValue(jobId, out var cached)) return cached; - return MigrationUtilities.SafeExecute(() => - { - var loadedObject = JsonStore.Read( - GetJobDefinitionPath(jobId)); - if (loadedObject == null) - return null; - _jobs[jobId] = loadedObject; - return loadedObject; - }, (Job?)null, $"LoadJob({jobId})"); + var loadedObject = JsonStore.Read( + GetJobDefinitionPath(jobId)); + if (loadedObject == null) + return null; + _jobs[jobId] = loadedObject; + return loadedObject; } /// Retrieves a job by ID, preferring the active in-memory job if it matches. @@ -102,27 +98,24 @@ public static List GetAllJobs(List ids) /// Persists a job to disk and updates the in-memory cache. public static bool SaveJob(Job job) { - return MigrationUtilities.SafeExecute(() => + lock (_writeJobLock) { - lock (_writeJobLock) + SerializeAndPersist(job); + _jobs[job.Id] = job; + if (!string.IsNullOrEmpty( + MigrationJobContext.Instance + .ActiveMigrationJobId) + && job.Id + == MigrationJobContext.Instance + .ActiveMigrationJobId) { - SerializeAndPersist(job); - _jobs[job.Id] = job; - if (!string.IsNullOrEmpty( - MigrationJobContext.Instance - .ActiveMigrationJobId) - && job.Id - == MigrationJobContext.Instance - .ActiveMigrationJobId) + lock (_cacheLock) { - lock (_cacheLock) - { - _cachedActiveJob = job; - } + _cachedActiveJob = job; } } - return true; - }, false, "SaveJob"); + } + return true; } internal static void PersistActiveJobUnderLock() diff --git a/CassandraMigrationProcessor/Context/JsonStore.cs b/CassandraMigrationProcessor/Context/JsonStore.cs index b3916a2..1bede0c 100644 --- a/CassandraMigrationProcessor/Context/JsonStore.cs +++ b/CassandraMigrationProcessor/Context/JsonStore.cs @@ -33,17 +33,22 @@ internal static class JsonStore /// /// Serializes to JSON and writes it to - /// . Returns the underlying - /// result, or false - /// when the store is unavailable. + /// . Persistence failures are propagated so + /// callers cannot report a successful checkpoint or state transition + /// that was never durably written. /// internal static bool Write( string path, T value, bool indented = true) { var store = Store; - if (store == null) return false; + if (store == null) + throw new InvalidOperationException( + "Document store is not initialized."); var json = JsonConvert.SerializeObject( value, indented ? Formatting.Indented : Formatting.None); - return store.Write(path, json); + if (!store.Write(path, json)) + throw new IOException( + $"Document store failed to write '{path}'."); + return true; } } diff --git a/CassandraMigrationProcessor/Context/MigrationJobContext.cs b/CassandraMigrationProcessor/Context/MigrationJobContext.cs index 377cd60..9f986bd 100644 --- a/CassandraMigrationProcessor/Context/MigrationJobContext.cs +++ b/CassandraMigrationProcessor/Context/MigrationJobContext.cs @@ -321,16 +321,13 @@ private JobListLoadAttempt TryLoadJobListOnce(string path) public bool SaveJobList() { - return MigrationUtilities.SafeExecute(() => + if (JobIndex != null) { - if (JobIndex != null) + lock (_writeJobListLock) { - lock (_writeJobListLock) - { - JsonStore.Write(JobStore.JobRegistryPath, JobIndex); - } + JsonStore.Write(JobStore.JobRegistryPath, JobIndex); } - return true; - }, false, "SaveJobList"); + } + return true; } } diff --git a/CassandraMigrationProcessor/Context/UnitStore.cs b/CassandraMigrationProcessor/Context/UnitStore.cs index 8a785f6..461e3a1 100644 --- a/CassandraMigrationProcessor/Context/UnitStore.cs +++ b/CassandraMigrationProcessor/Context/UnitStore.cs @@ -30,35 +30,32 @@ public static TableMigration GetUnit( public static bool SaveUnit( TableMigration mu, bool updateParent) { - return MigrationUtilities.SafeExecute(() => - { - if (mu == null) return false; + ArgumentNullException.ThrowIfNull(mu); - if (mu.ParentJob == null && MigrationJobContext.Instance.CurrentlyActiveJob != null) - mu.ParentJob = - MigrationJobContext.Instance.CurrentlyActiveJob; + if (mu.ParentJob == null && MigrationJobContext.Instance.CurrentlyActiveJob != null) + mu.ParentJob = + MigrationJobContext.Instance.CurrentlyActiveJob; - if (mu.ParentJob != null && updateParent) - TableMigrationMapper.UpdateParentJob(mu); + if (mu.ParentJob != null && updateParent) + TableMigrationMapper.UpdateParentJob(mu); - lock (_writeMULock) - { - JsonStore.Write( - JobStore.GetUnitDocumentPath(mu.JobId, mu.Id), mu); - } + lock (_writeMULock) + { + JsonStore.Write( + JobStore.GetUnitDocumentPath(mu.JobId, mu.Id), mu); + } - if (MigrationJobContext.Instance.CurrentlyActiveJob != null - && updateParent) - { - JobStore.PersistActiveJobUnderLock(); - } + if (MigrationJobContext.Instance.CurrentlyActiveJob != null + && updateParent) + { + JobStore.PersistActiveJobUnderLock(); + } - if (MigrationJobContext.Instance.MigrationUnitsCache != null) - MigrationJobContext.Instance.MigrationUnitsCache - .UpdateMigrationUnit(mu); + if (MigrationJobContext.Instance.MigrationUnitsCache != null) + MigrationJobContext.Instance.MigrationUnitsCache + .UpdateMigrationUnit(mu); - return true; - }, false, "SaveUnit"); + return true; } /// Removes a migration unit from its parent job and deletes it from storage. @@ -67,33 +64,28 @@ public static bool RemoveUnit(TableMigrationSummary unit) if (unit == null || unit.ParentJob == null) return false; - return MigrationUtilities.SafeExecute(() => - { - var job = unit.ParentJob; - var index = job.Tables - .FindIndex(mu => mu.Id == unit.Id); - if (index == -1) return false; + var job = unit.ParentJob; + var index = job.Tables + .FindIndex(mu => mu.Id == unit.Id); + if (index == -1) return false; - job.Tables.RemoveAt(index); + job.Tables.RemoveAt(index); - if (!MigrationJobContext.Instance.SaveMigrationJob(job)) - return false; + MigrationJobContext.Instance.SaveMigrationJob(job); - var filePath = JobStore.GetUnitDocumentPath(unit.JobId, unit.Id); - MigrationJobContext.Instance.Store.Delete(filePath); + var filePath = JobStore.GetUnitDocumentPath(unit.JobId, unit.Id); + if (!MigrationJobContext.Instance.Store.Delete(filePath)) + throw new IOException( + $"Failed to delete migration unit '{filePath}'."); - return true; - }, false, "RemoveUnit"); + return true; } public static TableMigration GetFromStorage( string jobId, string unitId) { - return MigrationUtilities.SafeExecute(() => - { - return JsonStore.Read( - JobStore.GetUnitDocumentPath(jobId, unitId)); - }, (TableMigration)null, $"GetFromStorage({jobId}, {unitId})"); + return JsonStore.Read( + JobStore.GetUnitDocumentPath(jobId, unitId)); } public static List GetMigrationUnitsToMigrate( diff --git a/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs b/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs index c7719ab..6fcc951 100644 --- a/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs +++ b/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs @@ -83,7 +83,10 @@ public static async Task CreateAsync( try { sourceSessions = new SourceSessionWrapper( - log, job, pipelineConfig.WorkerCount); + log, + job, + pipelineConfig.WorkerCount, + control.ReportFault); target = await CassandraClientFactory.CreateTargetSessionAsync(log, job); return new MigrationJobRunner( log, job, pipelineConfig, control, sourceSessions, target); @@ -893,17 +896,16 @@ void AddExpandedUnit(string keyspaceName, string tableName) => } catch (ArgumentException ex) { - _log.WriteLine( - $"Skipping invalid namespace entry '{fullName}': {ex.Message}", - LogType.Warning); - continue; + throw new ArgumentException( + $"Invalid namespace entry '{fullName}'.", + nameof(job.Namespaces), + ex); } if (string.IsNullOrEmpty(keyspace) || string.IsNullOrEmpty(table)) { - _log.WriteLine( - $"Skipping namespace entry '{fullName}' — empty keyspace or table after parsing.", - LogType.Warning); - continue; + throw new ArgumentException( + $"Namespace entry '{fullName}' contains an empty keyspace or table.", + nameof(job.Namespaces)); } if (table != "*") @@ -935,7 +937,9 @@ void AddExpandedUnit(string keyspaceName, string tableName) => } catch (Exception ex) { - _log.WriteLine($"Failed to discover tables in keyspace {keyspace}: {ex.Message}", LogType.Error); + throw new InvalidOperationException( + $"Failed to discover tables in source keyspace '{keyspace}'.", + ex); } } @@ -973,8 +977,9 @@ private async Task IsTableAccessibleAsync( } catch (Exception vex) { - _log.WriteLine($"Skipping {keyspace}.{tableName}: {vex.Message}", LogType.Warning); - return false; + throw new InvalidOperationException( + $"Source table accessibility check failed for '{keyspace}.{tableName}'.", + vex); } } } diff --git a/CassandraMigrationProcessor/Infrastructure/TableDiscovery.cs b/CassandraMigrationProcessor/Infrastructure/TableDiscovery.cs index 31f9cae..558bb2f 100644 --- a/CassandraMigrationProcessor/Infrastructure/TableDiscovery.cs +++ b/CassandraMigrationProcessor/Infrastructure/TableDiscovery.cs @@ -22,7 +22,7 @@ public static class TableDiscovery private static List? TryDeserializeJson(string input, string context) { try { return JsonConvert.DeserializeObject>(input); } - catch (Exception ex) + catch (JsonException ex) { Console.WriteLine($"[WARN] {context}: {ex.Message}"); return null; @@ -130,9 +130,12 @@ public static async Task> // Tolerates the '*' wildcard sentinel; see ParseNamespaceEntries. (keyspace, table) = CqlIdentifier.SplitNamespaceEntry(fullName); } - catch (ArgumentException) + catch (ArgumentException ex) { - continue; // skip malformed entries + throw new ArgumentException( + $"Invalid namespace entry '{fullName}'.", + nameof(namespacesToMigrate), + ex); } if (!unitsToAdd.Any(x => @@ -189,4 +192,3 @@ public static Tuple ValidateNamespaceFormat( return Tuple.Create(true, normalizedOutput, string.Empty); } } - diff --git a/CassandraMigrationProcessor/Infrastructure/TableMigrationMapper.cs b/CassandraMigrationProcessor/Infrastructure/TableMigrationMapper.cs index c9b6a27..72e97e4 100644 --- a/CassandraMigrationProcessor/Infrastructure/TableMigrationMapper.cs +++ b/CassandraMigrationProcessor/Infrastructure/TableMigrationMapper.cs @@ -13,38 +13,35 @@ public static class TableMigrationMapper public static bool UpdateParentJob(TableMigration unit) { - if (unit.ParentJob == null) return false; + if (unit.ParentJob == null) + throw new InvalidOperationException( + $"Migration unit '{unit.KeyspaceName}.{unit.TableName}' has no parent job."); - try + lock (_updateParentLock) { - lock (_updateParentLock) - { - var index = unit.ParentJob.Tables - .FindIndex(mu => mu.Id == unit.Id); - if (index == -1) return false; + var index = unit.ParentJob.Tables + .FindIndex(mu => mu.Id == unit.Id); + if (index == -1) + throw new InvalidOperationException( + $"Migration unit '{unit.KeyspaceName}.{unit.TableName}' is missing from its parent job."); - var target = unit.ParentJob.Tables[index]; - // Flush-and-reset the per-batch accumulator at the - // explicit sync boundary, then surface the unit via - // ToSummary. Only overwrite the sticky "last flushed - // batch" when this flush actually drained fresh - // activity (flushed > 0); idle ticks preserve the - // previous sticky value so the dashboard does not zero - // the column between UI renders while replay is - // actively applying rows. - long flushed = Interlocked.Exchange( - ref unit._changeFeedUpdatesInLastBatch, 0); - if (flushed > 0) - Interlocked.Exchange( - ref unit._changeFeedLastFlushedBatch, flushed); - ToSummary(unit, target); - } - return true; - } - catch - { - return false; + var target = unit.ParentJob.Tables[index]; + // Flush-and-reset the per-batch accumulator at the + // explicit sync boundary, then surface the unit via + // ToSummary. Only overwrite the sticky "last flushed + // batch" when this flush actually drained fresh + // activity (flushed > 0); idle ticks preserve the + // previous sticky value so the dashboard does not zero + // the column between UI renders while replay is + // actively applying rows. + long flushed = Interlocked.Exchange( + ref unit._changeFeedUpdatesInLastBatch, 0); + if (flushed > 0) + Interlocked.Exchange( + ref unit._changeFeedLastFlushedBatch, flushed); + ToSummary(unit, target); } + return true; } public static TableMigrationSummary ToSummary( diff --git a/CassandraMigrationProcessor/Models/TableMigration.cs b/CassandraMigrationProcessor/Models/TableMigration.cs index 0a21618..afa3330 100644 --- a/CassandraMigrationProcessor/Models/TableMigration.cs +++ b/CassandraMigrationProcessor/Models/TableMigration.cs @@ -294,7 +294,7 @@ private void AssignIfZero(JObject source, string field, Action assign) { return node.ToObject(); } - catch (Exception ex) + catch (JsonException ex) { readFailed = true; Console.Error.WriteLine( diff --git a/CassandraMigrationProcessor/Persistence/DiskPersistence.cs b/CassandraMigrationProcessor/Persistence/DiskPersistence.cs index 2fd2e0b..0d7977f 100644 --- a/CassandraMigrationProcessor/Persistence/DiskPersistence.cs +++ b/CassandraMigrationProcessor/Persistence/DiskPersistence.cs @@ -159,9 +159,7 @@ public bool Write(string id, string jsonContent) if (string.IsNullOrWhiteSpace(jsonContent)) throw new ArgumentException("JSON content cannot be null or empty", nameof(jsonContent)); - return MigrationUtilities.SafeExecute( - () => FileSystem.WriteAllText(GetFilePath(id), jsonContent), - false, $"Write({id})"); + return FileSystem.WriteAllText(GetFilePath(id), jsonContent); } public string? Read(string id) @@ -169,9 +167,10 @@ public bool Write(string id, string jsonContent) _ = Logs(); RequireJsonId(id, nameof(id)); - return MigrationUtilities.SafeExecute( - () => FileSystem.ReadAllText(GetFilePath(id)), - null, $"Read({id})"); + var path = GetFilePath(id); + return FileSystem.Exists(path) + ? FileSystem.ReadAllText(path) + : null; } public bool Exists(string id) @@ -181,9 +180,7 @@ public bool Exists(string id) if (!id.EndsWith(FILE_EXTENSION)) throw new ArgumentException($"ID must end with {FILE_EXTENSION} extension", nameof(id)); - return MigrationUtilities.SafeExecute( - () => FileSystem.Exists(GetFilePath(id)), - false, $"Exists({id})"); + return FileSystem.Exists(GetFilePath(id)); } /// @@ -196,32 +193,27 @@ public bool Delete(string id) if (string.IsNullOrWhiteSpace(id)) throw new ArgumentException("ID cannot be null or empty", nameof(id)); - return MigrationUtilities.SafeExecute(() => + if (id.EndsWith(FILE_EXTENSION)) { - if (id.EndsWith(FILE_EXTENSION)) - { - var filePath = GetFilePath(id); - if (!FileSystem.Exists(filePath)) return false; - FileSystem.DeleteIfExists(filePath); - return true; - } - return FileSystem.DeleteDirectory(GetDirectoryPath(id), recursive: true); - }, false, $"Delete({id})"); + var filePath = GetFilePath(id); + if (!FileSystem.Exists(filePath)) return false; + FileSystem.DeleteIfExists(filePath); + return true; + } + return FileSystem.DeleteDirectory(GetDirectoryPath(id), recursive: true); } public List ListIds() { _ = Logs(); - return MigrationUtilities.SafeExecute(() => - { - var files = FileSystem.ListFiles(_storagePath, "*" + FILE_EXTENSION, recursive: true); - return files - .Select(f => Path.GetRelativePath(_storagePath, f) - .Replace('/', '\\') - .Replace(Path.DirectorySeparatorChar, '\\')) - .ToList(); - }, new List(), "ListIds"); + var files = FileSystem.ListFiles( + _storagePath, "*" + FILE_EXTENSION, recursive: true); + return files + .Select(f => Path.GetRelativePath(_storagePath, f) + .Replace('/', '\\') + .Replace(Path.DirectorySeparatorChar, '\\')) + .ToList(); } // --- Log operations delegated to LogPersistence --- diff --git a/CassandraMigrationWebApp/Service/JobManager.cs b/CassandraMigrationWebApp/Service/JobManager.cs index 577c804..ce2ca38 100644 --- a/CassandraMigrationWebApp/Service/JobManager.cs +++ b/CassandraMigrationWebApp/Service/JobManager.cs @@ -371,16 +371,7 @@ public Task StartMigration(Job job, string sourceConnectionString, string target foreach (var staleJob in staleRunningJobs) { staleJob.Status = JobStatus.Pending; - try - { - _context.SaveMigrationJob(staleJob); - } - catch (Exception ex) - { - _log.WriteLine( - $"Failed to clear stale Running status for job {staleJob.Id}: {ex.Message}", - LogType.Warning); - } + _context.SaveMigrationJob(staleJob); } _context.ActiveMigrationJobId = job.Id; From 12b484956464e39d9b4ed1bbf339006dfca55ef6 Mon Sep 17 00:00:00 2001 From: Nitesh Vijay Date: Fri, 21 Aug 2026 11:45:42 +0530 Subject: [PATCH 30/32] Log AAD source session rotation Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 078ac7b9-aa7f-44a1-9053-44c346cec5e1 --- .../CassandraDriver/SourceSessionWrapper.cs | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs b/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs index 181a91f..ffa0d96 100644 --- a/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs +++ b/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs @@ -159,6 +159,9 @@ private async Task DisposeRetiredSessionAfterDelayAsync(ISession session) RemoveUdtRegistrations(session); MigrationUtilities.SafeDisposeSession( session, "Deferred rotated session"); + _log.WriteLine( + "Retired AAD source session disposed after the rotation grace period.", + LogType.Info); } } @@ -219,6 +222,10 @@ private void ScheduleTokenRefresh(DateTime expiry) _tokenRefreshTimer = new Timer( RefreshTokenCallback, null, delay, Timeout.InfiniteTimeSpan); + _log.WriteLine( + $"AAD source token refresh scheduled for " + + $"{DateTime.UtcNow.Add(delay):O}; token expires {expiry:O}.", + LogType.Info); } private void RefreshTokenCallback(object? state) @@ -236,6 +243,9 @@ private void RefreshTokenCallback(object? state) DateTime expiry = GetTokenExpiry(freshToken); Refresh(freshToken); ScheduleTokenRefresh(expiry); + _log.WriteLine( + "AAD source session refreshed successfully.", + LogType.Info); } catch (Exception ex) { From 074af7d5e85969f0ebc43122d2d8a39125d8b72b Mon Sep 17 00:00:00 2001 From: Nitesh Vijay Date: Fri, 21 Aug 2026 11:53:01 +0530 Subject: [PATCH 31/32] Allow controlled AAD refresh timing Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 078ac7b9-aa7f-44a1-9053-44c346cec5e1 --- .../CassandraDriver/SourceSessionWrapper.cs | 25 ++++++++++++++++++- 1 file changed, 24 insertions(+), 1 deletion(-) diff --git a/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs b/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs index ffa0d96..e873445 100644 --- a/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs +++ b/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs @@ -23,6 +23,10 @@ internal sealed class SourceSessionWrapper : IDisposable { private static readonly TimeSpan RetiredSessionDisposalDelay = TimeSpan.FromMinutes(10); + private static readonly TimeSpan DefaultTokenRefreshLeadTime = + TimeSpan.FromMinutes(5); + private const string TokenRefreshLeadMinutesSetting = + "CMT_AAD_TOKEN_REFRESH_LEAD_MINUTES"; private readonly object _lifecycleLock = new(); private readonly MigrationLog _log; @@ -214,8 +218,9 @@ private void ScheduleTokenRefresh(DateTime expiry) { StopTokenRefreshCore(); + TimeSpan refreshLeadTime = ResolveTokenRefreshLeadTime(); TimeSpan delay = expiry - DateTime.UtcNow - - TimeSpan.FromMinutes(5); + - refreshLeadTime; if (delay < TimeSpan.FromMinutes(1)) delay = TimeSpan.FromMinutes(1); @@ -228,6 +233,24 @@ private void ScheduleTokenRefresh(DateTime expiry) LogType.Info); } + private static TimeSpan ResolveTokenRefreshLeadTime() + { + string? configured = + Environment.GetEnvironmentVariable( + TokenRefreshLeadMinutesSetting); + if (string.IsNullOrWhiteSpace(configured)) + return DefaultTokenRefreshLeadTime; + + if (!int.TryParse(configured, out int minutes) + || minutes <= 0) + { + throw new InvalidOperationException( + $"{TokenRefreshLeadMinutesSetting} must be a positive integer."); + } + + return TimeSpan.FromMinutes(minutes); + } + private void RefreshTokenCallback(object? state) { Exception? fatalFailure = null; From ae409bd2b0a917451aa97d551014ea85d06ddec2 Mon Sep 17 00:00:00 2001 From: Nitesh Vijay Date: Fri, 21 Aug 2026 12:06:43 +0530 Subject: [PATCH 32/32] Support accelerated AAD refresh validation Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 078ac7b9-aa7f-44a1-9053-44c346cec5e1 --- .../CassandraDriver/SourceSessionWrapper.cs | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs b/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs index e873445..0664eff 100644 --- a/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs +++ b/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs @@ -218,11 +218,15 @@ private void ScheduleTokenRefresh(DateTime expiry) { StopTokenRefreshCore(); - TimeSpan refreshLeadTime = ResolveTokenRefreshLeadTime(); + TimeSpan refreshLeadTime = + ResolveTokenRefreshLeadTime(out bool isConfigured); TimeSpan delay = expiry - DateTime.UtcNow - refreshLeadTime; - if (delay < TimeSpan.FromMinutes(1)) - delay = TimeSpan.FromMinutes(1); + TimeSpan minimumDelay = isConfigured + ? TimeSpan.FromSeconds(10) + : TimeSpan.FromMinutes(1); + if (delay < minimumDelay) + delay = minimumDelay; _tokenRefreshTimer = new Timer( RefreshTokenCallback, null, @@ -233,13 +237,17 @@ private void ScheduleTokenRefresh(DateTime expiry) LogType.Info); } - private static TimeSpan ResolveTokenRefreshLeadTime() + private static TimeSpan ResolveTokenRefreshLeadTime( + out bool isConfigured) { string? configured = Environment.GetEnvironmentVariable( TokenRefreshLeadMinutesSetting); if (string.IsNullOrWhiteSpace(configured)) + { + isConfigured = false; return DefaultTokenRefreshLeadTime; + } if (!int.TryParse(configured, out int minutes) || minutes <= 0) @@ -248,6 +256,7 @@ private static TimeSpan ResolveTokenRefreshLeadTime() $"{TokenRefreshLeadMinutesSetting} must be a positive integer."); } + isConfigured = true; return TimeSpan.FromMinutes(minutes); }