Merge pull request #57 from ZjzMisaka/fix-unable-to-fetch-results-by-…

…group [*] fix: unable to fetch results by group
ZjzMisaka · Sep 22, 2024 · 4fe11d9 · 4fe11d9
2 parents 8af2f38 + 8477567
commit 4fe11d9
Show file tree

Hide file tree

Showing 5 changed files with 39 additions and 9 deletions.
diff --git a/PowerThreadPool/Core/PowerThreadPool.Control.cs b/PowerThreadPool/Core/PowerThreadPool.Control.cs
@@ -47,8 +47,13 @@ public void StopIfRequested()
             }
             else if (work != null)
             {
-                _aliveWorkDic.TryRemove(work.ID, out _);
-                if (work.Group != null)
+                // 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)
+                {
+                    _aliveWorkDic.TryRemove(work.ID, out _);
+                }
+                if (work.Group != null && !work.ShouldStoreResult)
                 {
                     if (_workGroupDic.TryGetValue(work.Group, out ConcurrentSet<string> idSet))
                     {
@@ -310,6 +315,11 @@ public List<ExecuteResult<TResult>> Fetch<TResult>(IEnumerable<string> idList, b
                     if (executeResultBase != null)
                     {
                         resultList.Add(executeResultBase.ToTypedResult<TResult>());
+
+                        if (removeAfterFetch && _aliveWorkDic.TryRemove(id, out WorkBase work))
+                        {
+                            RemoveWorkFromGroup(work.Group, work);
+                        }
                     }
                     else
                     {

diff --git a/PowerThreadPool/Core/PowerThreadPool.Events.cs b/PowerThreadPool/Core/PowerThreadPool.Events.cs
@@ -143,8 +143,13 @@ internal void WorkCallbackEnd(WorkBase work, Status status)
                 CallbackEnd.Invoke(work.ID);
             }
 
-            _aliveWorkDic.TryRemove(work.ID, out _);
-            if (work.Group != null)
+            // 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)
+            {
+                _aliveWorkDic.TryRemove(work.ID, out _);
+            }
+            if (work.Group != null && !work.ShouldStoreResult)
             {
                 if (_workGroupDic.TryGetValue(work.Group, out ConcurrentSet<string> idSet))
                 {

diff --git a/PowerThreadPool/Core/PowerThreadPool.Group.cs b/PowerThreadPool/Core/PowerThreadPool.Group.cs
@@ -74,11 +74,24 @@ public bool RemoveWorkFromGroup(string groupName, string workID)
         {
             if (_aliveWorkDic.TryGetValue(workID, out WorkBase work))
             {
-                if (_workGroupDic.TryGetValue(groupName, out ConcurrentSet<string> workIDSet))
-                {
-                    work.Group = null;
-                    return workIDSet.Remove(workID);
-                }
+                return RemoveWorkFromGroup(groupName, work);
+            }
+
+            return false;
+        }
+
+        /// <summary>
+        /// Remove work from group.
+        /// </summary>
+        /// <param name="groupName"></param>
+        /// <param name="work"></param>
+        /// <returns>Returns false if either the work or the group does not exist, or if the work does not belong to the group.</returns>
+        private bool RemoveWorkFromGroup(string groupName, WorkBase work)
+        {
+            if (_workGroupDic.TryGetValue(groupName, out ConcurrentSet<string> workIDSet))
+            {
+                work.Group = null;
+                return workIDSet.Remove(work.ID);
             }
 
             return false;

diff --git a/PowerThreadPool/Works/Work.cs b/PowerThreadPool/Works/Work.cs
@@ -34,6 +34,7 @@ internal override string Group
         internal override TimeoutOption WorkTimeoutOption => _workOption.TimeoutOption;
         internal override RetryOption RetryOption => _workOption.RetryOption;
         internal override bool LongRunning => _workOption.LongRunning;
+        internal override bool ShouldStoreResult => _workOption.ShouldStoreResult;
         internal override ConcurrentSet<string> Dependents => _workOption.Dependents;
 
         internal Work(PowerPool powerPool, string id, Func<object[], TResult> function, object[] param, WorkOption<TResult> option)

diff --git a/PowerThreadPool/Works/WorkBase.cs b/PowerThreadPool/Works/WorkBase.cs
@@ -52,6 +52,7 @@ internal bool IsPausing
         internal abstract TimeoutOption WorkTimeoutOption { get; }
         internal abstract RetryOption RetryOption { get; }
         internal abstract bool LongRunning { get; }
+        internal abstract bool ShouldStoreResult { get; }
         internal abstract ConcurrentSet<string> Dependents { get; }
     }
 }