Merge pull request #66 from ZjzMisaka/add-optional-pre-stop-action-pa…

…rameter-to-StopIfRequested-method [+] Add optional pre-stop action parameter to StopIfRequested method
ZjzMisaka · Oct 16, 2024 · 3fa86bb · 3fa86bb
2 parents d292742 + 632a577
commit 3fa86bb
Show file tree

Hide file tree

Showing 2 changed files with 149 additions and 1 deletion.
diff --git a/PowerThreadPool/Core/PowerThreadPool.Control.cs b/PowerThreadPool/Core/PowerThreadPool.Control.cs
@@ -35,19 +35,31 @@ public void PauseIfRequested()
         /// To exit the logic, the function will throw a PowerThreadPool.Exceptions.WorkStopException. Do not catch it. 
         /// If you do not want to exit the logic in this way (for example, if you have some unmanaged resources that need to be released before exiting), it is recommended to use CheckIfRequestedStop. 
         /// </summary>
-        public void StopIfRequested()
+        /// <param name="beforeStop">
+        /// An optional function that is executed before the stop process.
+        /// Return false to prevent stopping.
+        /// </param>
+        public void StopIfRequested(Func<bool> beforeStop = null)
         {
             WorkBase work = null;
             bool res = CheckIfRequestedStopAndGetWork(ref work);
 
             if (!res)
             {
+                if (beforeStop != null && !beforeStop())
+                {
+                    return;
+                }
                 _aliveWorkDic.Clear();
                 _workGroupDic.Clear();
                 throw new WorkStopException();
             }
             else if (work != null)
             {
+                if (beforeStop != null && !beforeStop())
+                {
+                    return;
+                }
                 // If the result needs to be stored, there is a possibility of fetching the result through Group.
                 // Therefore, Work should not be removed from _aliveWorkDic and _workGroupDic for the time being
                 if (work.Group == null || !work.ShouldStoreResult)

diff --git a/UnitTest/ControlTest.cs b/UnitTest/ControlTest.cs
@@ -1042,6 +1042,142 @@ public async void TestStopByIDUseCheckIfRequestedStop()
             Assert.Equal(id, resID);
         }
 
+        [Fact]
+        public async void TestStopByIDDoBeforeStop()
+        {
+            PowerPool powerPool = new PowerPool();
+            List<long> logList = new List<long>();
+
+            object lockObj = new object();
+
+            string id = null;
+            string resID = null;
+            powerPool.WorkStarted += (s, e) =>
+            {
+                powerPool.Stop(e.ID);
+            };
+
+            id = powerPool.QueueWorkItem(() =>
+            {
+                long start = GetNowSs();
+                while (GetNowSs() - start <= 1000)
+                {
+                    powerPool.StopIfRequested(() => true);
+                    Thread.Sleep(1);
+                }
+            }, (res) =>
+            {
+                resID = res.Status == Status.Stopped ? "Stopped" + res.ID : "Ended" + res.ID;
+            });
+
+            await powerPool.WaitAsync(id);
+            await powerPool.WaitAsync();
+
+            Assert.Equal("Stopped" + id, resID);
+        }
+
+        [Fact]
+        public void TestStopAllDoBeforeStop()
+        {
+            PowerPool powerPool = new PowerPool();
+            List<long> logList = new List<long>();
+
+            object lockObj = new object();
+
+            string id = null;
+            string resID = null;
+            powerPool.WorkStarted += (s, e) =>
+            {
+                powerPool.Stop();
+            };
+
+            id = powerPool.QueueWorkItem(() =>
+            {
+                long start = GetNowSs();
+                while (GetNowSs() - start <= 1000)
+                {
+                    powerPool.StopIfRequested(() => true);
+                    Thread.Sleep(1);
+                }
+            }, (res) =>
+            {
+                resID = res.Status == Status.Stopped ? "Stopped" + res.ID : "Ended" + res.ID;
+            });
+
+            powerPool.Wait(id);
+            powerPool.Wait();
+
+            Assert.Equal("Stopped" + id, resID);
+        }
+
+        [Fact]
+        public async void TestStopByIDDoBeforeStopReturnFalse()
+        {
+            PowerPool powerPool = new PowerPool();
+            List<long> logList = new List<long>();
+
+            object lockObj = new object();
+
+            string id = null;
+            string resID = null;
+            powerPool.WorkStarted += (s, e) =>
+            {
+                powerPool.Stop(e.ID);
+            };
+
+            id = powerPool.QueueWorkItem(() =>
+            {
+                long start = GetNowSs();
+                while (GetNowSs() - start <= 1000)
+                {
+                    powerPool.StopIfRequested(() => false);
+                    Thread.Sleep(1);
+                }
+            }, (res) =>
+            {
+                resID = res.Status == Status.Stopped ? "Stopped" + res.ID : "Ended" + res.ID;
+            });
+
+            await powerPool.WaitAsync(id);
+            await powerPool.WaitAsync();
+
+            Assert.Equal("Ended" + id, resID);
+        }
+
+        [Fact]
+        public void TestStopAllDoBeforeStopReturnFalse()
+        {
+            PowerPool powerPool = new PowerPool();
+            List<long> logList = new List<long>();
+
+            object lockObj = new object();
+
+            string id = null;
+            string resID = null;
+            powerPool.WorkStarted += (s, e) =>
+            {
+                powerPool.Stop();
+            };
+
+            id = powerPool.QueueWorkItem(() =>
+            {
+                long start = GetNowSs();
+                while (GetNowSs() - start <= 1000)
+                {
+                    powerPool.StopIfRequested(() => false);
+                    Thread.Sleep(1);
+                }
+            }, (res) =>
+            {
+                resID = res.Status == Status.Stopped ? "Stopped" + res.ID : "Ended" + res.ID;
+            });
+
+            powerPool.Wait(id);
+            powerPool.Wait();
+
+            Assert.Equal("Ended" + id, resID);
+        }
+
         [Fact]
         public void TestCancelByID()
         {