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 6932b6f..0127398 100644 --- a/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs +++ b/CassandraMigrationProcessor/CassandraDriver/CassandraClientFactory.cs @@ -4,11 +4,11 @@ using CassandraMigrationProcessor.Infrastructure; using CassandraMigrationProcessor.Models; 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 { @@ -35,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( @@ -45,13 +43,8 @@ public static ISession CreateSourceSession( int port, string username, string password, - 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, @@ -68,9 +61,7 @@ public static ISession CreateSourceSession( { try { - var session = ConnectCluster(builder); - RegisterAadTokenRefresh(session, password, tokenRefreshManager); - return session; + return ConnectCluster(builder); } catch (Exception ex) when ( ExceptionClassifier.IsTransient(ex) @@ -91,24 +82,6 @@ public static ISession CreateSourceSession( throw new UnreachableException(); } - /// - /// 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. - /// 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?.SetManagedSourceSession(session); - tokenRefreshManager?.StartTokenRefreshTimer(password); - } - /// /// Create a session to an OSS Apache Cassandra cluster. /// Tries SSL first, falls back to plain if SSL fails. @@ -302,58 +275,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, - 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) - { - password = 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. - job.SourceUseAad = true; - } - - // 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 - && !string.IsNullOrEmpty(job.SourceContactPoint)) - { - username = job.SourceContactPoint - .Split('.')[0]; - } - - return CreateSourceSession( - MigrationLog, - job.SourceContactPoint, - job.SourcePort, - username, - password, - tokenRefreshManager, - maxConnectionsPerHost: ResolveMaxConnectionsPerHost(job.SourceMaxConnectionsPerHost, job.MaxConnectionsPerHost)); - } - /// /// Per-side connection pool sizing. The per-side override /// ( / @@ -407,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/ISessionFactory.cs b/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs deleted file mode 100644 index 6068a87..0000000 --- a/CassandraMigrationProcessor/CassandraDriver/ISessionFactory.cs +++ /dev/null @@ -1,51 +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 -{ - /// 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(); -} - -/// -/// 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/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 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 new file mode 100644 index 0000000..0664eff --- /dev/null +++ b/CassandraMigrationProcessor/CassandraDriver/SourceSessionWrapper.cs @@ -0,0 +1,329 @@ +using Cassandra; +using CassandraMigrationProcessor.Infrastructure; +using CassandraMigrationProcessor.Models; +using System.Collections.Concurrent; +using System.IdentityModel.Tokens.Jwt; + +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 +/// operations can complete. +/// +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; + private readonly Action _reportFatalFailure; + 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 int _disposed; + + public SourceSessionWrapper( + MigrationLog log, + Job job, + 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( + "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) + ScheduleTokenRefresh(GetTokenExpiry(credential)); + } + + public ISession GetSession() + { + ObjectDisposedException.ThrowIf( + Volatile.Read(ref _disposed) != 0, + this); + return Volatile.Read(ref _currentSession); + } + + public async Task GetTypedSessionAsync(string keyspace) + { + var session = GetSession(); + var key = (Session: session, Keyspace: keyspace); + var registration = _udtRegistrations.GetOrAdd( + key, + key => new Lazy( + () => RegisterUdtsAsync(key.Session, key.Keyspace), + 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; + } + + private static async Task RegisterUdtsAsync( + ISession session, + string keyspace) + { + var allUdts = await SchemaManager.GetUserDefinedTypesAsync( + session, keyspace); + await DynamicUdtRegistrar.RegisterAsync( + session, keyspace, allUdts); + } + + private void Refresh(string credential) + { + var session = CreateSession(credential); + var retiredSession = _currentSession; + Volatile.Write(ref _currentSession, session); + _retiredSessions.Add(retiredSession); + + _ = 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); + + bool shouldDispose; + lock (_lifecycleLock) + { + shouldDispose = _retiredSessions.Remove(session); + } + + if (shouldDispose) + { + RemoveUdtRegistrations(session); + MigrationUtilities.SafeDisposeSession( + session, "Deferred rotated session"); + _log.WriteLine( + "Retired AAD source session disposed after the rotation grace period.", + LogType.Info); + } + } + + private void RemoveUdtRegistrations(ISession session) + { + foreach (var key in _udtRegistrations.Keys) + { + if (ReferenceEquals(key.Session, session)) + _udtRegistrations.TryRemove(key, out _); + } + } + + private static DateTime GetTokenExpiry(string token) + { + if (string.IsNullOrWhiteSpace(token)) + throw new InvalidOperationException( + "AAD token acquisition returned an empty token."); + + 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) + { + if (job.SourceUseAad) + return AcquireAadToken(); + + return job.SourcePassword ?? string.Empty; + } + + 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(DateTime expiry) + { + StopTokenRefreshCore(); + + TimeSpan refreshLeadTime = + ResolveTokenRefreshLeadTime(out bool isConfigured); + TimeSpan delay = expiry - DateTime.UtcNow + - refreshLeadTime; + TimeSpan minimumDelay = isConfigured + ? TimeSpan.FromSeconds(10) + : TimeSpan.FromMinutes(1); + if (delay < minimumDelay) + delay = minimumDelay; + + _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 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) + { + throw new InvalidOperationException( + $"{TokenRefreshLeadMinutesSetting} must be a positive integer."); + } + + isConfigured = true; + return TimeSpan.FromMinutes(minutes); + } + + private void RefreshTokenCallback(object? state) + { + Exception? fatalFailure = null; + lock (_lifecycleLock) + { + if (Volatile.Read(ref _disposed) != 0 + || _tokenRefreshTimer == null) + return; + + try + { + string freshToken = AcquireAadToken(); + DateTime expiry = GetTokenExpiry(freshToken); + Refresh(freshToken); + ScheduleTokenRefresh(expiry); + _log.WriteLine( + "AAD source session refreshed successfully.", + LogType.Info); + } + catch (Exception ex) + { + string message = + $"AAD token refresh failed. Aborting migration job: {ex.Message}"; + _log.WriteLine(message, LogType.Error); + StopTokenRefreshCore(); + fatalFailure = new InvalidOperationException(message, ex); + } + } + + if (fatalFailure != null) + _reportFatalFailure(fatalFailure); + } + + private void StopTokenRefreshCore() + { + _tokenRefreshTimer?.Dispose(); + _tokenRefreshTimer = null; + } + + public void Dispose() + { + List sessionsToDispose; + lock (_lifecycleLock) + { + if (Interlocked.Exchange(ref _disposed, 1) != 0) + return; + StopTokenRefreshCore(); + + sessionsToDispose = _retiredSessions.ToList(); + _retiredSessions.Clear(); + sessionsToDispose.Add(_currentSession); + _udtRegistrations.Clear(); + } + + foreach (var session in sessionsToDispose) + { + MigrationUtilities.SafeDisposeSession( + session, "Source session wrapper"); + } + } + + private sealed record SourceSessionSettings( + string ContactPoint, + int Port, + string Username, + int MaxConnectionsPerHost); +} diff --git a/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs b/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs deleted file mode 100644 index a9133e8..0000000 --- a/CassandraMigrationProcessor/CassandraDriver/TokenRefreshManager.cs +++ /dev/null @@ -1,212 +0,0 @@ -using Cassandra; -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 ISession? _managedSourceSession; - private readonly MigrationLog _log; - private DateTime _tokenExpiresAt = DateTime.MinValue; - private int _consecutiveRefreshFailures; - private const int MaxRefreshFailures = 6; - - private string? _lastSourceContactPoint; - private int _lastSourcePort; - private string? _lastSourceUsername; - - public TokenRefreshManager(MigrationLog log) - { - _log = log; - } - - /// - /// 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). - /// - 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) - { - 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) - { - try - { - string freshToken = GetFreshAadToken(); - - // If we have a managed session, recreate it - if (_managedSourceSession != null - && !_managedSourceSession.IsDisposed - && _lastSourceContactPoint != null) - { - var oldSession = _managedSourceSession; - _managedSourceSession = CassandraClientFactory.CreateSourceSession( - _log, - _lastSourceContactPoint, - _lastSourcePort, - _lastSourceUsername ?? string.Empty, - freshToken); - MigrationUtilities.SafeDisposeSession(oldSession, "TokenRefresh old session"); - } - - // 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); - } - } - } - - /// - /// Set the managed source session so the token refresh - /// timer can reconnect it proactively. - /// - public void SetManagedSourceSession(ISession session) - { - _managedSourceSession = session; - } - - public void Dispose() - { - lock (_refreshLock) - { - StopTokenRefreshTimer(); - } - } -} 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/DataCopyWorker.cs b/CassandraMigrationProcessor/DataTransfer/DataCopyWorker.cs index 743a2de..557a098 100644 --- a/CassandraMigrationProcessor/DataTransfer/DataCopyWorker.cs +++ b/CassandraMigrationProcessor/DataTransfer/DataCopyWorker.cs @@ -33,7 +33,11 @@ public async Task RunAsync(PipelineContext ctx) Partition? current = null; try { - reader = await PageReader.CreateAsync(_workerLog, ctx.SessionFactory, ctx.ReaderConfig, _ct); + reader = new PageReader( + _workerLog, + ctx.SourceSession, + ctx.ReaderConfig, + _ct); writer = await PageWriter.CreateAsync(_workerLog, ctx.SessionFactory, ctx.WriterConfig, _ct); while (!_ct.IsCancellationRequested @@ -159,7 +163,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 dd48feb..0bbd668 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, SourceSessionWrapper sourceSession, + JobControl control) { _log = log; _pipelineConfig = pipelineConfig; @@ -43,7 +45,8 @@ public JobPipeline(MigrationLog log, Job job, PipelineConfig pipelineConfig, Job Context = new PipelineContext( _partitions, - new JobSessionFactory(log, job, tokenRefreshManager), + sourceSession, + new JobSessionFactory(log, job), readerConfig, writerConfig, EnableReplay: enableReplay, diff --git a/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs b/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs index ccd2778..6fcc951 100644 --- a/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs +++ b/CassandraMigrationProcessor/DataTransfer/MigrationJobRunner.cs @@ -19,7 +19,7 @@ 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; // attached as inner when the consecutive-auth threshold trips so @@ -31,14 +31,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 mint their own sessions - /// via for throughput isolation. + /// 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; /// @@ -53,16 +52,14 @@ private MigrationJobRunner( Job job, PipelineConfig pipelineConfig, JobControl control, - TokenRefreshManager tokenRefreshManager, - ISession sourceSession, + SourceSessionWrapper sourceSessions, ISession targetSession) { _log = log; _job = job; _pipelineConfig = pipelineConfig; _control = control; - _tokenRefreshManager = tokenRefreshManager; - _sourceSession = sourceSession; + _sourceSessions = sourceSessions; _targetSession = targetSession; } @@ -81,20 +78,23 @@ public static async Task CreateAsync( ArgumentNullException.ThrowIfNull(control); var pipelineConfig = PipelineConfig.Resolve(job, config); - var tokenRefreshManager = new TokenRefreshManager(log); - ISession? source = null; + SourceSessionWrapper? sourceSessions = null; ISession? target = null; try { - source = CassandraClientFactory.CreateSourceSession(log, job, tokenRefreshManager); + sourceSessions = new SourceSessionWrapper( + log, + job, + pipelineConfig.WorkerCount, + control.ReportFault); target = await CassandraClientFactory.CreateTargetSessionAsync(log, job); - return new MigrationJobRunner(log, job, pipelineConfig, control, tokenRefreshManager, source, target); + return new MigrationJobRunner( + log, job, pipelineConfig, control, sourceSessions, target); } catch { MigrationUtilities.SafeDisposeSession(target, "MigrationJobRunner target (CreateAsync rollback)"); - MigrationUtilities.SafeDisposeSession(source, "MigrationJobRunner source (CreateAsync rollback)"); - tokenRefreshManager.StopTokenRefreshTimer(); + sourceSessions?.Dispose(); throw; } } @@ -173,7 +173,10 @@ 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, + _sourceSessions, + _control); _pipeline.Start(); await RunCopyPhaseAsync(job, units, partitioning, cancellationToken); @@ -241,11 +244,10 @@ public async Task StartAsync() /// public ValueTask DisposeAsync() { - _tokenRefreshManager.StopTokenRefreshTimer(); MigrationUtilities.SafeDispose(_pipeline, "JobPipeline (Dispose)"); _pipeline = null; MigrationUtilities.SafeDisposeSession(_targetSession, "MigrationJobRunner target session"); - MigrationUtilities.SafeDisposeSession(_sourceSession, "MigrationJobRunner source session"); + _sourceSessions.Dispose(); return ValueTask.CompletedTask; } @@ -355,7 +357,7 @@ private async Task RunSchemaPhaseAsync( .Distinct(StringComparer.Ordinal) .ToList(); await SchemaManager.WarnAboutUnreplicatedSchemaAsync( - _sourceSession, inScopeKeyspaces, _log); + _sourceSessions.GetSession(), inScopeKeyspaces, _log); } catch (Exception ex) { @@ -485,7 +487,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) { @@ -767,7 +771,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); @@ -854,7 +858,6 @@ public void Stop() // without waiting for the outer Task to observe the cancel. MigrationUtilities.SafeDispose(_pipeline, "JobPipeline (Stop)"); _pipeline = null; - _tokenRefreshManager.StopTokenRefreshTimer(); } /// @@ -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 != "*") @@ -914,11 +916,16 @@ void AddExpandedUnit(string keyspaceName, string tableName) => try { - 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); } @@ -930,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); } } @@ -968,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/DataTransfer/PageReader.cs b/CassandraMigrationProcessor/DataTransfer/PageReader.cs index 243e1f0..f69281e 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; @@ -22,19 +21,17 @@ 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 : IDisposable +internal class PageReader { private readonly WorkerLog _log; private readonly CancellationToken _ct; - private readonly ISession _sourceSession; + private readonly SourceSessionWrapper _sourceSession; private readonly int _pageSize; private readonly int _maxReadRetries; private readonly bool _preserveCellTtl; private readonly bool _useJsonCopy; - private readonly ConcurrentDictionary _udtRegistrations = new(); /// /// Most recent transient exception observed during retry-exhausted @@ -50,7 +47,11 @@ 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) + public PageReader( + WorkerLog log, + SourceSessionWrapper sourceSession, + ReaderConfig config, + CancellationToken cancellationToken) { _log = log; _ct = cancellationToken; @@ -58,44 +59,8 @@ private PageReader(WorkerLog log, ISessionFactory sessionFactory, ReaderConfig c _maxReadRetries = config.MaxReadRetries; _preserveCellTtl = config.PreserveCellTtlAndWritetime; _useJsonCopy = config.UseJsonCopy; - _sourceSession = sessionFactory.CreateSourceSession(); - } - - public static Task CreateAsync(WorkerLog log, - ISessionFactory sessionFactory, ReaderConfig config, - CancellationToken cancellationToken) - { - return Task.FromResult(new PageReader(log, sessionFactory, config, cancellationToken)); - } - - public void Dispose() => MigrationUtilities.SafeDisposeSession(_sourceSession, "PageReader source session"); - - /// - /// 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 Task EnsureUdtsRegisteredAsync(Partition partition) - { - // JSON read path bypasses CLR-side UDT decoding entirely. - if (!partition.Table.IsCounterTable && _useJsonCopy) - return Task.CompletedTask; - - return _udtRegistrations.GetOrAdd(partition.Table.Spec.KeyspaceName, async ks => - { - try - { - var allUdts = await SchemaManager.GetUserDefinedTypesAsync(_sourceSession, ks); - await DynamicUdtRegistrar.RegisterAsync(_sourceSession, ks, 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); - throw; - } - }); + _sourceSession = sourceSession + ?? throw new ArgumentNullException(nameof(sourceSession)); } /// @@ -136,8 +101,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; @@ -175,8 +138,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; @@ -217,9 +178,19 @@ internal record ReadResult( // intact and will retry the same page once the source stops // throttling. var resultSet = await RetryExecutor.ExecuteOrDefaultAsync( - operation: _ => _sourceSession.ExecuteAsync(stmt).WaitAsync(_ct), + operation: async _ => + { + var sourceSession = useJson + ? _sourceSession.GetSession() + : await _sourceSession.GetTypedSessionAsync( + partition.Table.Spec.KeyspaceName).ConfigureAwait(false); + return await sourceSession.ExecuteAsync(stmt) + .WaitAsync(_ct) + .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) => diff --git a/CassandraMigrationProcessor/DataTransfer/PageWriter.cs b/CassandraMigrationProcessor/DataTransfer/PageWriter.cs index 86c2981..8fb5ad9 100644 --- a/CassandraMigrationProcessor/DataTransfer/PageWriter.cs +++ b/CassandraMigrationProcessor/DataTransfer/PageWriter.cs @@ -58,13 +58,18 @@ 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.CreateTargetSessionAsync(); + var targetSession = await sessionFactory.CreateSessionAsync(cancellationToken); return new PageWriter(log, targetSession, config, cancellationToken); } - public void Dispose() => MigrationUtilities.SafeDisposeSession(_targetSession, "PageWriter target session"); + public void Dispose() + { + MigrationUtilities.SafeDisposeSession( + _targetSession, + "PageWriter target session"); + } private Task GetStrategyAsync(Partition partition) { diff --git a/CassandraMigrationProcessor/DataTransfer/PipelineContext.cs b/CassandraMigrationProcessor/DataTransfer/PipelineContext.cs index 652df55..2412d38 100644 --- a/CassandraMigrationProcessor/DataTransfer/PipelineContext.cs +++ b/CassandraMigrationProcessor/DataTransfer/PipelineContext.cs @@ -6,15 +6,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, worker 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, + SourceSessionWrapper SourceSession, + JobSessionFactory SessionFactory, ReaderConfig ReaderConfig, WriterConfig WriterConfig, bool EnableReplay, 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/TableCopySpec.cs b/CassandraMigrationProcessor/Models/TableCopySpec.cs index f06aacf..704f5e0 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 JobSessionFactory. /// public record TableCopySpec( string KeyspaceName, 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; 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.