diff --git a/PowerThreadPool/Core/PowerThreadPool.Control.cs b/PowerThreadPool/Core/PowerThreadPool.Control.cs index a4697ef..c36b949 100644 --- 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 idSet)) { @@ -310,6 +315,11 @@ public List> Fetch(IEnumerable idList, b if (executeResultBase != null) { resultList.Add(executeResultBase.ToTypedResult()); + + 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 index 9e2d021..0119922 100644 --- 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 idSet)) { diff --git a/PowerThreadPool/Core/PowerThreadPool.Group.cs b/PowerThreadPool/Core/PowerThreadPool.Group.cs index c40b7d8..6f6cf56 100644 --- 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 workIDSet)) - { - work.Group = null; - return workIDSet.Remove(workID); - } + return RemoveWorkFromGroup(groupName, work); + } + + return false; + } + + /// + /// Remove work from group. + /// + /// + /// + /// Returns false if either the work or the group does not exist, or if the work does not belong to the group. + private bool RemoveWorkFromGroup(string groupName, WorkBase work) + { + if (_workGroupDic.TryGetValue(groupName, out ConcurrentSet workIDSet)) + { + work.Group = null; + return workIDSet.Remove(work.ID); } return false; diff --git a/PowerThreadPool/Works/Work.cs b/PowerThreadPool/Works/Work.cs index 30fe5ab..6c485bc 100644 --- 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 Dependents => _workOption.Dependents; internal Work(PowerPool powerPool, string id, Func function, object[] param, WorkOption option) diff --git a/PowerThreadPool/Works/WorkBase.cs b/PowerThreadPool/Works/WorkBase.cs index 61717ad..6cf5dbc 100644 --- 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 Dependents { get; } } }