Optimize Session Canceling process

Aeliux · Aeliux · commit 5c086199a49e · 2025-01-25T10:29:38.000+03:30
diff --git a/src/Kryptor.Cli/KeyStoreAnalyze/CrackSession.cs b/src/Kryptor.Cli/KeyStoreAnalyze/CrackSession.cs
@@ -14,8 +14,12 @@ public class CrackSession : Session, ISessionHost
         public bool Found { get; private set; }
         public bool Verbose => false;
 
-        public CrackSession(int maxRunningSessions)
+        CancellationTokenSource CancellationTokenSource;
+
+        public CrackSession(int maxRunningSessions, CancellationTokenSource cancellationTokenSource)
         {
+            CancellationTokenSource = cancellationTokenSource;
+
             Progress = -1;
 
             container = new SessionContainer(this, maxRunningSessions);
@@ -39,7 +43,9 @@ protected override async Task<bool> RunAsync(ISessionHost sessionHost, Cancellat
                 NewSession(session, false, true);
             }
 
-            await container.WaitAll(cancellationToken);
+            await container.WaitAll();
+
+            cancellationToken.ThrowIfCancellationRequested();
 
             return true;
         }
@@ -49,10 +55,7 @@ private void StopTimer()
             CalculationTimer.Stop();
             Found = true;
 
-            foreach (CancellationTokenSource token in container.TokenSources)
-            {
-                token.Cancel();
-            }
+            CancellationTokenSource.Cancel();
         }
 
         public void Start(ClientContext context) => throw new System.NotImplementedException();
@@ -64,5 +67,7 @@ private void StopTimer()
         public void MonitorTask(Task task) => throw new System.NotImplementedException();
 
         public Task<TResponse> OnSessionRequest<TResponse>(ISession session, SessionRequest<TResponse> request, CancellationToken cancellationToken) => throw new System.NotImplementedException();
+
+        public CancellationToken GetCancellationToken() => CancellationTokenSource.Token;
     }
 }
diff --git a/src/Kryptor.Cli/KeyStoreAnalyze/SessionHost.cs b/src/Kryptor.Cli/KeyStoreAnalyze/SessionHost.cs
@@ -22,7 +22,7 @@ public override void Start(ClientContext context)
             KeyStoreLoadSession ksLoadSession = CreateKeyStoreLoadSession(ks);
             NewSession(ksLoadSession);
 
-            CrackSession calcSession = new CrackSession(maxRunningSessions);
+            CrackSession calcSession = new CrackSession(maxRunningSessions, MasterToken);
             ksLoadSession.ContinueWith(calcSession);
             NewSession(calcSession);
 
diff --git a/src/Kryptor.Client/ISessionHost.cs b/src/Kryptor.Client/ISessionHost.cs
@@ -26,6 +26,14 @@ public interface ISessionHost
         /// </param>
         void End(bool cancelled);
 
+        /// <summary>
+        /// Gets the cancellation token of the session host.
+        /// </summary>
+        /// <returns>
+        /// The global cancellation token of this session host.
+        /// </returns>
+        CancellationToken GetCancellationToken();
+
         /// <summary>
         /// Starts a new session in this session host.
         /// </summary>
diff --git a/src/Kryptor.Client/SessionContainer.cs b/src/Kryptor.Client/SessionContainer.cs
@@ -30,19 +30,6 @@ public ISession[] Sessions
         /// </summary>
         public Task[] Tasks => SessionPool.Values.Select(x => x.Task).Concat(TaskPool).Where(x => x != null).ToArray();
 
-        /// <summary>
-        /// Gets an array of all cancellation tokens.
-        /// </summary>
-        public CancellationTokenSource[] TokenSources
-        {
-            get
-            {
-                tokenSources ??= SessionPool.Values.Select(x => x.TokenSource).ToArray();
-
-                return tokenSources;
-            }
-        }
-
         /// <summary>
         /// Gets all session holders.
         /// </summary>
diff --git a/src/Kryptor.Client/SessionHolder.cs b/src/Kryptor.Client/SessionHolder.cs
@@ -20,11 +20,6 @@ public class SessionHolder
         /// </summary>
         public Task Task { get; private set; }
 
-        /// <summary>
-        /// Gets the cancellation token source that controls the running session.
-        /// </summary>
-        public CancellationTokenSource TokenSource { get; set; }
-
         internal bool AutoRemove { get; set; }
 
         /// <summary>
@@ -33,18 +28,14 @@ public class SessionHolder
         /// <param name="session">
         /// The session instance.
         /// </param>
-        /// <param name="tokenSource">
-        /// The cancellation token source that controls the running task.
-        /// </param>
-        public SessionHolder(ISession session, CancellationTokenSource tokenSource)
+        public SessionHolder(ISession session)
         {
             if (session.Status != SessionStatus.NotStarted)
             {
                 throw new ArgumentException("The session is already started.");
             }
 
             Session = session;
-            TokenSource = tokenSource;
         }
 
         /// <summary>
@@ -64,7 +55,7 @@ public Task StartTask(ISessionHost sessionHost, bool throwIfRunning = true)
                 return throwIfRunning ? throw new InvalidOperationException("Session is already started.") : null;
             }
 
-            Task = Session.StartAsync(sessionHost, TokenSource.Token);
+            Task = Session.StartAsync(sessionHost, sessionHost.GetCancellationToken());
             return Task;
         }
     }
diff --git a/src/Kryptor.Client/SessionHost.cs b/src/Kryptor.Client/SessionHost.cs
@@ -23,12 +23,18 @@ public abstract class SessionHost : ISessionHost
         /// <summary>
         /// Initializes all <see cref="SessionHost"/> instances.
         /// </summary>
-        protected SessionHost()
+        /// <param name="cancellationTokenSource">
+        /// The cancellation token source to use in this session host.
+        /// </param>
+        protected SessionHost(CancellationTokenSource cancellationTokenSource = null)
         {
             Container = new SessionContainer(this, Environment.ProcessorCount - 1);
-            MasterToken = new CancellationTokenSource();
+            MasterToken = cancellationTokenSource ?? new CancellationTokenSource();
         }
 
+        /// <inheritdoc/>
+        public CancellationToken GetCancellationToken() => MasterToken.Token;
+
         /// <inheritdoc/>
         public abstract void Start(ClientContext context);
 
@@ -44,8 +50,6 @@ public virtual void End(bool canceled)
 
             if (canceled)
             {
-                Container.Cancel();
-
                 MasterToken.Cancel();
 
                 // Tell to the SessionManager to do whatever needed
diff --git a/src/Kryptor.Client/SessionManager.cs b/src/Kryptor.Client/SessionManager.cs
@@ -50,29 +50,16 @@ public SessionContainer(ISessionHost sessionHost, int maxRunningSessions)
         /// </summary>
         /// <remarks>
         /// This method also waits for sessions that are not started yet.
-        /// If the <paramref name="cancellationToken"/> was cancelled, this method will cancel all existing tokens in this container and waits for all of them to end.
-        /// This method intended to be used by sub-session models. in standard session host based model, the monitoring, cancelling and ending is managed by various parts.
+        /// This method intended to be used by sub-session models. in standard session host based model, the monitoring, canceling and ending is managed by various parts.
         /// But in sub-session models, you could achieve all of these features in this method.
-        /// Using this method in session host is not recommended at all because this method is not intended for that porpose.
+        /// Using this method in session host is not recommended at all because this method is not intended for that purpose.
         /// </remarks>
-        /// <param name="cancellationToken">
-        /// The token to monitor for cancellation requests.
-        /// </param>
-        public async Task WaitAll(CancellationToken cancellationToken)
+        public async Task WaitAll()
         {
             while (!Sessions.All(x => x.Status == SessionStatus.Ended))
             {
-                if (cancellationToken.IsCancellationRequested && !_cancellationRequested)
-                {
-                    Cancel();
-
-                    _cancellationRequested = true;
-                }
-
                 await AsyncCompat.Delay(5);
             }
-
-            cancellationToken.ThrowIfCancellationRequested();
         }
 
         /// <summary>
@@ -83,21 +70,10 @@ public async Task WaitAll(CancellationToken cancellationToken)
         /// A Task must be registered by <see cref="AddMonitoringTask(Task)"/> that waits for not started sessions before calling this method.
         /// This method only intended to be used in <see cref="ISessionHost.End(bool)"/>.
         /// so using it in sub-session models is not recommended as it could be lead to unexpected session exit and losing control of sub-sessions.
-        /// it's highly recommended to use <see cref="WaitAll(CancellationToken)"/> method in that situation.
+        /// it's highly recommended to use <see cref="WaitAll()"/> method in that situation.
         /// </remarks>
         public void WaitForRunningTasks() => Task.WaitAll(Tasks);
 
-        /// <summary>
-        /// Sends cancellation request to all cancellation tokens in this container.
-        /// </summary>
-        public void Cancel()
-        {
-            foreach (CancellationTokenSource token in TokenSources.Where(x => !x.IsCancellationRequested))
-            {
-                token.Cancel();
-            }
-        }
-
         /// <summary>
         /// Sets the provided <paramref name="sessionGroup"/> as the SessionManager's cache system.
         /// </summary>
@@ -194,9 +170,7 @@ public static SessionHolder WrapSession(ISession session, bool autoRemove)
                 throw new ArgumentException("The session is already started.");
             }
 
-            CancellationTokenSource tokenSource = new CancellationTokenSource();
-
-            SessionHolder sessionHolder = new SessionHolder(session, tokenSource)
+            SessionHolder sessionHolder = new SessionHolder(session)
             {
                 AutoRemove = autoRemove
             };
@@ -275,6 +249,7 @@ private void ProcessTasks()
                         if (_sessionHost.Verbose)
                         {
                             Console.WriteLine($"Error starting session: {ex.Message}");
+                            break;
                         }
                     }
                 }
@@ -285,11 +260,11 @@ private void ProcessTasks()
             }
         }
 
-        private static bool SafeIsReady(SessionHolder holder)
+        private bool SafeIsReady(SessionHolder holder)
         {
             try
             {
-                return holder.Session.IsReady(holder.TokenSource.Token);
+                return holder.Session.IsReady(_sessionHost.GetCancellationToken());
             }
             catch
             {

Original file line number	Diff line number	Diff line change
`@@ -20,11 +20,6 @@ public class SessionHolder`
`20`	`20`	`/// </summary>`
`21`	`21`	`public Task Task { get; private set; }`
`22`	`22`
`23`		`- /// <summary>`
`24`		`- /// Gets the cancellation token source that controls the running session.`
`25`		`- /// </summary>`
`26`		`- public CancellationTokenSource TokenSource { get; set; }`
`27`		`-`
`28`	`23`	`internal bool AutoRemove { get; set; }`
`29`	`24`
`30`	`25`	`/// <summary>`
`@@ -33,18 +28,14 @@ public class SessionHolder`
`33`	`28`	`/// <param name="session">`
`34`	`29`	`/// The session instance.`
`35`	`30`	`/// </param>`
`36`		`- /// <param name="tokenSource">`
`37`		`- /// The cancellation token source that controls the running task.`
`38`		`- /// </param>`
`39`		`- public SessionHolder(ISession session, CancellationTokenSource tokenSource)`
	`31`	`+ public SessionHolder(ISession session)`
`40`	`32`	`{`
`41`	`33`	`if (session.Status != SessionStatus.NotStarted)`
`42`	`34`	`{`
`43`	`35`	`throw new ArgumentException("The session is already started.");`
`44`	`36`	`}`
`45`	`37`
`46`	`38`	`Session = session;`
`47`		`- TokenSource = tokenSource;`
`48`	`39`	`}`
`49`	`40`
`50`	`41`	`/// <summary>`
`@@ -64,7 +55,7 @@ public Task StartTask(ISessionHost sessionHost, bool throwIfRunning = true)`
`64`	`55`	`return throwIfRunning ? throw new InvalidOperationException("Session is already started.") : null;`
`65`	`56`	`}`
`66`	`57`
`67`		`- Task = Session.StartAsync(sessionHost, TokenSource.Token);`
	`58`	`+ Task = Session.StartAsync(sessionHost, sessionHost.GetCancellationToken());`
`68`	`59`	`return Task;`
`69`	`60`	`}`
`70`	`61`	`}`
Original file line number	Diff line number	Diff line change
`@@ -50,29 +50,16 @@ public SessionContainer(ISessionHost sessionHost, int maxRunningSessions)`
`50`	`50`	`/// </summary>`
`51`	`51`	`/// <remarks>`
`52`	`52`	`/// This method also waits for sessions that are not started yet.`
`53`		`- /// If the <paramref name="cancellationToken"/> was cancelled, this method will cancel all existing tokens in this container and waits for all of them to end.`
`54`		`- /// This method intended to be used by sub-session models. in standard session host based model, the monitoring, cancelling and ending is managed by various parts.`
	`53`	`+ /// This method intended to be used by sub-session models. in standard session host based model, the monitoring, canceling and ending is managed by various parts.`
`55`	`54`	`/// But in sub-session models, you could achieve all of these features in this method.`
`56`		`- /// Using this method in session host is not recommended at all because this method is not intended for that porpose.`
	`55`	`+ /// Using this method in session host is not recommended at all because this method is not intended for that purpose.`
`57`	`56`	`/// </remarks>`
`58`		`- /// <param name="cancellationToken">`
`59`		`- /// The token to monitor for cancellation requests.`
`60`		`- /// </param>`
`61`		`- public async Task WaitAll(CancellationToken cancellationToken)`
	`57`	`+ public async Task WaitAll()`
`62`	`58`	`{`
`63`	`59`	`while (!Sessions.All(x => x.Status == SessionStatus.Ended))`
`64`	`60`	`{`
`65`		`- if (cancellationToken.IsCancellationRequested && !_cancellationRequested)`
`66`		`- {`
`67`		`- Cancel();`
`68`		`-`
`69`		`- _cancellationRequested = true;`
`70`		`- }`
`71`		`-`
`72`	`61`	`await AsyncCompat.Delay(5);`
`73`	`62`	`}`
`74`		`-`
`75`		`- cancellationToken.ThrowIfCancellationRequested();`
`76`	`63`	`}`
`77`	`64`
`78`	`65`	`/// <summary>`
`@@ -83,21 +70,10 @@ public async Task WaitAll(CancellationToken cancellationToken)`
`83`	`70`	`/// A Task must be registered by <see cref="AddMonitoringTask(Task)"/> that waits for not started sessions before calling this method.`
`84`	`71`	`/// This method only intended to be used in <see cref="ISessionHost.End(bool)"/>.`
`85`	`72`	`/// so using it in sub-session models is not recommended as it could be lead to unexpected session exit and losing control of sub-sessions.`
`86`		`- /// it's highly recommended to use <see cref="WaitAll(CancellationToken)"/> method in that situation.`
	`73`	`+ /// it's highly recommended to use <see cref="WaitAll()"/> method in that situation.`
`87`	`74`	`/// </remarks>`
`88`	`75`	`public void WaitForRunningTasks() => Task.WaitAll(Tasks);`
`89`	`76`
`90`		`- /// <summary>`
`91`		`- /// Sends cancellation request to all cancellation tokens in this container.`
`92`		`- /// </summary>`
`93`		`- public void Cancel()`
`94`		`- {`
`95`		`- foreach (CancellationTokenSource token in TokenSources.Where(x => !x.IsCancellationRequested))`
`96`		`- {`
`97`		`- token.Cancel();`
`98`		`- }`
`99`		`- }`
`100`		`-`
`101`	`77`	`/// <summary>`
`102`	`78`	`/// Sets the provided <paramref name="sessionGroup"/> as the SessionManager's cache system.`
`103`	`79`	`/// </summary>`
`@@ -194,9 +170,7 @@ public static SessionHolder WrapSession(ISession session, bool autoRemove)`
`194`	`170`	`throw new ArgumentException("The session is already started.");`
`195`	`171`	`}`
`196`	`172`
`197`		`- CancellationTokenSource tokenSource = new CancellationTokenSource();`
`198`		`-`
`199`		`- SessionHolder sessionHolder = new SessionHolder(session, tokenSource)`
	`173`	`+ SessionHolder sessionHolder = new SessionHolder(session)`
`200`	`174`	`{`
`201`	`175`	`AutoRemove = autoRemove`
`202`	`176`	`};`
`@@ -275,6 +249,7 @@ private void ProcessTasks()`
`275`	`249`	`if (_sessionHost.Verbose)`
`276`	`250`	`{`
`277`	`251`	`Console.WriteLine($"Error starting session: {ex.Message}");`
	`252`	`+ break;`
`278`	`253`	`}`
`279`	`254`	`}`
`280`	`255`	`}`
`@@ -285,11 +260,11 @@ private void ProcessTasks()`
`285`	`260`	`}`
`286`	`261`	`}`
`287`	`262`
`288`		`- private static bool SafeIsReady(SessionHolder holder)`
	`263`	`+ private bool SafeIsReady(SessionHolder holder)`
`289`	`264`	`{`
`290`	`265`	`try`
`291`	`266`	`{`
`292`		`- return holder.Session.IsReady(holder.TokenSource.Token);`
	`267`	`+ return holder.Session.IsReady(_sessionHost.GetCancellationToken());`
`293`	`268`	`}`
`294`	`269`	`catch`
`295`	`270`	`{`