Branch data     Line data    Source code

       1                 :            : /*
       2                 :            :  * Copyright (C) 2026 by Thun Lu. All rights reserved.
       3                 :            :  * Author: Thun Lu <thun.lu@zohomail.cn>
       4                 :            :  * Repo:   https://github.com/thun-res/vlink
       5                 :            :  *  _    __   __      _           __
       6                 :            :  * | |  / /  / /     (_) ____    / /__
       7                 :            :  * | | / /  / /     / / / __ \  / //_/
       8                 :            :  * | |/ /  / /___  / / / / / / / ,<
       9                 :            :  * |___/  /_____/ /_/ /_/ /_/ /_/|_|
      10                 :            :  *
      11                 :            :  * Licensed under the Apache License, Version 2.0 (the "License");
      12                 :            :  * you may not use this file except in compliance with the License.
      13                 :            :  * You may obtain a copy of the License at
      14                 :            :  *
      15                 :            :  *     http://www.apache.org/licenses/LICENSE-2.0
      16                 :            :  *
      17                 :            :  * Unless required by applicable law or agreed to in writing, software
      18                 :            :  * distributed under the License is distributed on an "AS IS" BASIS,
      19                 :            :  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
      20                 :            :  * See the License for the specific language governing permissions and
      21                 :            :  * limitations under the License.
      22                 :            :  */
      23                 :            : 
      24                 :            : #include "./base/task_handle.h"
      25                 :            : 
      26                 :            : #include <chrono>
      27                 :            : #include <memory>
      28                 :            : #include <mutex>
      29                 :            : #include <utility>
      30                 :            : 
      31                 :            : #include "./base/condition_variable.h"
      32                 :            : #include "./base/memory_resource.h"
      33                 :            : 
      34                 :            : namespace vlink {
      35                 :            : 
      36                 :        117 : [[maybe_unused]] static bool is_terminal(TaskExecutionState state) noexcept {
      37   [ +  +  +  + ]:         85 :   return state == TaskExecutionState::kCompleted || state == TaskExecutionState::kCancelled ||
      38   [ +  +  +  +  :        202 :          state == TaskExecutionState::kDropped || state == TaskExecutionState::kRejected ||
                   +  + ]
      39                 :        117 :          state == TaskExecutionState::kFailed;
      40                 :            : }
      41                 :            : 
      42                 :            : template <typename StateT>
      43                 :         56 : [[maybe_unused]] static void release_parent_registration(StateT& state,
      44                 :            :                                                          CancellationRegistration& registration) noexcept {
      45                 :         56 :   registration = std::move(state.parent_registration);
      46                 :         56 : }
      47                 :            : 
      48                 :            : // TaskHandle::State
      49                 :            : struct TaskHandle::State final {
      50                 :            :   mutable std::mutex mtx;
      51                 :            :   ConditionVariable cv;
      52                 :            :   TaskExecutionState state{TaskExecutionState::kInvalid};
      53                 :            :   CancellationSource cancellation_source;
      54                 :            :   CancellationRegistration parent_registration;
      55                 :            : };
      56                 :            : 
      57                 :            : // TrackedTask
      58                 :            : class TrackedTask final {
      59                 :            :  public:
      60                 :         51 :   TrackedTask(TaskHandle handle, MoveFunction<void()>&& callback) noexcept
      61                 :         51 :       : handle_(std::move(handle)), callback_(std::move(callback)) {}
      62                 :            : 
      63                 :         51 :   TrackedTask(TrackedTask&& other) noexcept
      64                 :         51 :       : handle_(std::move(other.handle_)), callback_(std::move(other.callback_)), armed_(other.armed_) {
      65                 :         51 :     other.armed_ = false;
      66                 :         51 :   }
      67                 :            : 
      68                 :            :   TrackedTask& operator=(TrackedTask&& other) noexcept {
      69                 :            :     if VUNLIKELY (this == &other) {
      70                 :            :       return *this;
      71                 :            :     }
      72                 :            : 
      73                 :            :     if VUNLIKELY (armed_) {
      74                 :            :       TaskHandle::drop_task_if_queued(handle_);
      75                 :            :     }
      76                 :            : 
      77                 :            :     handle_ = std::move(other.handle_);
      78                 :            :     callback_ = std::move(other.callback_);
      79                 :            :     armed_ = other.armed_;
      80                 :            : 
      81                 :            :     other.armed_ = false;
      82                 :            : 
      83                 :            :     return *this;
      84                 :            :   }
      85                 :            : 
      86                 :        102 :   ~TrackedTask() {
      87         [ +  + ]:        102 :     if VUNLIKELY (armed_) {
      88                 :         16 :       TaskHandle::drop_task_if_queued(handle_);
      89                 :            :     }
      90                 :        102 :   }
      91                 :            : 
      92                 :         35 :   void operator()() {
      93                 :         35 :     armed_ = false;
      94                 :            : 
      95         [ +  + ]:         35 :     if VUNLIKELY (!TaskHandle::begin_task_execution(handle_)) {
      96                 :          4 :       return;
      97                 :            :     }
      98                 :            : 
      99                 :            :     try {
     100         [ +  - ]:         31 :       if VLIKELY (callback_) {
     101         [ +  + ]:         31 :         callback_();
     102                 :            :       }
     103                 :            : 
     104         [ +  - ]:         27 :       TaskHandle::complete_task_execution(handle_);
     105                 :          4 :     } catch (...) {
     106         [ +  - ]:          4 :       TaskHandle::fail_task_execution(handle_);
     107                 :          4 :     }
     108                 :            :   }
     109                 :            : 
     110                 :            :  private:
     111                 :            :   TaskHandle handle_;
     112                 :            :   MoveFunction<void()> callback_;
     113                 :            :   bool armed_{true};
     114                 :            : 
     115                 :            :   VLINK_DISALLOW_COPY_AND_ASSIGN(TrackedTask)
     116                 :            : };
     117                 :            : 
     118                 :            : // TaskHandle
     119                 :          6 : TaskHandle::TaskHandle() noexcept = default;
     120                 :            : 
     121                 :        270 : TaskHandle::~TaskHandle() = default;
     122                 :            : 
     123                 :         53 : TaskHandle::TaskHandle(const TaskHandle&) noexcept = default;
     124                 :            : 
     125                 :          2 : TaskHandle& TaskHandle::operator=(const TaskHandle&) noexcept = default;
     126                 :            : 
     127                 :        155 : TaskHandle::TaskHandle(TaskHandle&&) noexcept = default;
     128                 :            : 
     129                 :          2 : TaskHandle& TaskHandle::operator=(TaskHandle&&) noexcept = default;
     130                 :            : 
     131                 :         12 : bool TaskHandle::valid() const noexcept { return state_ != nullptr; }
     132                 :            : 
     133                 :        328 : TaskExecutionState TaskHandle::state() const noexcept {
     134         [ +  + ]:        328 :   if VUNLIKELY (!state_) {
     135                 :          7 :     return TaskExecutionState::kInvalid;
     136                 :            :   }
     137                 :            : 
     138                 :        321 :   std::lock_guard lock(state_->mtx);
     139                 :            : 
     140                 :        321 :   return state_->state;
     141                 :        321 : }
     142                 :            : 
     143                 :         28 : bool TaskHandle::is_done() const noexcept { return is_terminal(state()); }
     144                 :            : 
     145                 :          2 : CancellationToken TaskHandle::cancellation_token() const noexcept {
     146         [ +  + ]:          2 :   if VUNLIKELY (!state_) {
     147                 :          1 :     return {};
     148                 :            :   }
     149                 :            : 
     150                 :          1 :   return state_->cancellation_source.token();
     151                 :            : }
     152                 :            : 
     153         [ +  - ]:          9 : bool TaskHandle::cancel() const { return request_cancel(state_); }
     154                 :            : 
     155                 :         46 : bool TaskHandle::wait(int timeout_ms) const {
     156         [ +  + ]:         46 :   if VUNLIKELY (!state_) {
     157                 :          4 :     return false;
     158                 :            :   }
     159                 :            : 
     160         [ +  - ]:         42 :   std::unique_lock lock(state_->mtx);
     161                 :            : 
     162                 :         65 :   auto predicate = [this] { return is_terminal(state_->state); };
     163                 :            : 
     164         [ +  + ]:         42 :   if (timeout_ms < 0) {
     165                 :          2 :     state_->cv.wait(lock, predicate);
     166                 :            : 
     167                 :          2 :     return true;
     168                 :            :   }
     169                 :            : 
     170                 :         40 :   return state_->cv.wait_for(lock, std::chrono::milliseconds(timeout_ms), predicate);
     171                 :         42 : }
     172                 :            : 
     173                 :         56 : TaskHandle::TaskHandle(std::shared_ptr<State> state) noexcept : state_(std::move(state)) {}
     174                 :            : 
     175                 :         56 : TaskHandle TaskHandle::make_task_handle(const CancellationToken& parent_token) {
     176         [ +  - ]:         56 :   auto state = MemoryResource::make_shared<TaskHandle::State>();
     177                 :            : 
     178         [ +  + ]:         56 :   if VLIKELY (parent_token.valid()) {
     179                 :          6 :     std::weak_ptr<TaskHandle::State> weak_state = state;
     180                 :            : 
     181                 :          6 :     auto registration = parent_token.register_callback([weak_state]() {
     182                 :          6 :       auto locked = weak_state.lock();
     183                 :            : 
     184         [ +  - ]:          6 :       if VLIKELY (locked) {
     185         [ +  - ]:          6 :         (void)TaskHandle::request_cancel(std::move(locked));
     186                 :            :       }
     187   [ +  -  +  - ]:         18 :     });
     188                 :            : 
     189                 :            :     {
     190         [ +  - ]:          6 :       std::lock_guard lock(state->mtx);
     191                 :            : 
     192                 :          6 :       state->parent_registration = std::move(registration);
     193                 :          6 :     }
     194                 :            : 
     195         [ +  + ]:          6 :     if VUNLIKELY (parent_token.is_cancellation_requested()) {
     196         [ +  - ]:          3 :       (void)TaskHandle::request_cancel(state);
     197                 :            :     }
     198                 :          6 :   }
     199                 :            : 
     200                 :        112 :   return TaskHandle{std::move(state)};
     201                 :         56 : }
     202                 :            : 
     203                 :         51 : MoveFunction<void()> TaskHandle::make_tracked_task(TaskHandle handle, MoveFunction<void()>&& callback) {
     204         [ +  - ]:        102 :   return TrackedTask{std::move(handle), std::move(callback)};
     205                 :            : }
     206                 :            : 
     207                 :         56 : void TaskHandle::mark_task_queued(const TaskHandle& handle) {
     208         [ -  + ]:         56 :   if VUNLIKELY (!handle.state_) {
     209                 :          3 :     return;
     210                 :            :   }
     211                 :            : 
     212                 :         56 :   CancellationRegistration parent_registration;
     213                 :            : 
     214                 :            :   {
     215         [ +  - ]:         56 :     std::lock_guard lock(handle.state_->mtx);
     216                 :            : 
     217         [ +  + ]:         56 :     if VUNLIKELY (handle.state_->state != TaskExecutionState::kInvalid) {
     218                 :          3 :       return;
     219                 :            :     }
     220                 :            : 
     221         [ -  + ]:         53 :     if VUNLIKELY (handle.state_->cancellation_source.is_cancellation_requested()) {
     222                 :          0 :       handle.state_->state = TaskExecutionState::kCancelled;
     223                 :            : 
     224                 :          0 :       release_parent_registration(*handle.state_, parent_registration);
     225                 :            :     } else {
     226                 :         53 :       handle.state_->state = TaskExecutionState::kQueued;
     227                 :            :     }
     228         [ +  + ]:         56 :   }
     229                 :            : 
     230                 :         53 :   handle.state_->cv.notify_all();
     231         [ +  + ]:         56 : }
     232                 :            : 
     233                 :          9 : void TaskHandle::mark_task_rejected(const TaskHandle& handle) {
     234         [ -  + ]:          9 :   if VUNLIKELY (!handle.state_) {
     235                 :          0 :     return;
     236                 :            :   }
     237                 :            : 
     238                 :          9 :   CancellationRegistration parent_registration;
     239                 :            : 
     240                 :            :   {
     241         [ +  - ]:          9 :     std::lock_guard lock(handle.state_->mtx);
     242                 :            : 
     243         [ +  - ]:          9 :     if VLIKELY (!is_terminal(handle.state_->state)) {
     244                 :          9 :       handle.state_->state = TaskExecutionState::kRejected;
     245                 :            : 
     246                 :          9 :       release_parent_registration(*handle.state_, parent_registration);
     247                 :            :     }
     248                 :          9 :   }
     249                 :            : 
     250                 :          9 :   handle.state_->cv.notify_all();
     251                 :          9 : }
     252                 :            : 
     253                 :         35 : bool TaskHandle::begin_task_execution(const TaskHandle& handle) {
     254         [ -  + ]:         35 :   if VUNLIKELY (!handle.state_) {
     255                 :          0 :     return false;
     256                 :            :   }
     257                 :            : 
     258         [ +  - ]:         35 :   std::lock_guard lock(handle.state_->mtx);
     259                 :            : 
     260         [ +  + ]:         35 :   if VUNLIKELY (handle.state_->state != TaskExecutionState::kQueued) {
     261                 :          4 :     return false;
     262                 :            :   }
     263                 :            : 
     264         [ -  + ]:         31 :   if VUNLIKELY (handle.state_->cancellation_source.is_cancellation_requested()) {
     265                 :          0 :     handle.state_->state = TaskExecutionState::kCancelled;
     266                 :            : 
     267                 :          0 :     handle.state_->cv.notify_all();
     268                 :            : 
     269                 :          0 :     return false;
     270                 :            :   }
     271                 :            : 
     272                 :         31 :   handle.state_->state = TaskExecutionState::kRunning;
     273                 :            : 
     274                 :         31 :   handle.state_->cv.notify_all();
     275                 :            : 
     276                 :         31 :   return true;
     277                 :         35 : }
     278                 :            : 
     279                 :         27 : void TaskHandle::complete_task_execution(const TaskHandle& handle) {
     280         [ -  + ]:         27 :   if VUNLIKELY (!handle.state_) {
     281                 :          0 :     return;
     282                 :            :   }
     283                 :            : 
     284                 :         27 :   CancellationRegistration parent_registration;
     285                 :            : 
     286                 :            :   {
     287         [ +  - ]:         27 :     std::lock_guard lock(handle.state_->mtx);
     288                 :            : 
     289         [ +  - ]:         27 :     if VLIKELY (handle.state_->state == TaskExecutionState::kRunning) {
     290                 :         27 :       handle.state_->state = TaskExecutionState::kCompleted;
     291                 :            : 
     292                 :         27 :       release_parent_registration(*handle.state_, parent_registration);
     293                 :            :     }
     294                 :         27 :   }
     295                 :            : 
     296                 :         27 :   handle.state_->cv.notify_all();
     297                 :         27 : }
     298                 :            : 
     299                 :          4 : void TaskHandle::fail_task_execution(const TaskHandle& handle) {
     300         [ -  + ]:          4 :   if VUNLIKELY (!handle.state_) {
     301                 :          0 :     return;
     302                 :            :   }
     303                 :            : 
     304                 :          4 :   CancellationRegistration parent_registration;
     305                 :            : 
     306                 :            :   {
     307         [ +  - ]:          4 :     std::lock_guard lock(handle.state_->mtx);
     308                 :            : 
     309         [ +  - ]:          4 :     if VLIKELY (handle.state_->state == TaskExecutionState::kRunning) {
     310                 :          4 :       handle.state_->state = TaskExecutionState::kFailed;
     311                 :            : 
     312                 :          4 :       release_parent_registration(*handle.state_, parent_registration);
     313                 :            :     }
     314                 :          4 :   }
     315                 :            : 
     316                 :          4 :   handle.state_->cv.notify_all();
     317                 :          4 : }
     318                 :            : 
     319                 :         16 : void TaskHandle::drop_task_if_queued(const TaskHandle& handle) {
     320         [ -  + ]:         16 :   if VUNLIKELY (!handle.state_) {
     321                 :          0 :     return;
     322                 :            :   }
     323                 :            : 
     324                 :         16 :   CancellationRegistration parent_registration;
     325                 :            : 
     326                 :            :   {
     327         [ +  - ]:         16 :     std::lock_guard lock(handle.state_->mtx);
     328                 :            : 
     329         [ +  + ]:         16 :     if VLIKELY (handle.state_->state == TaskExecutionState::kQueued) {
     330                 :          7 :       handle.state_->state = TaskExecutionState::kDropped;
     331                 :            : 
     332                 :          7 :       release_parent_registration(*handle.state_, parent_registration);
     333                 :            :     }
     334                 :         16 :   }
     335                 :            : 
     336                 :         16 :   handle.state_->cv.notify_all();
     337                 :         16 : }
     338                 :            : 
     339                 :         18 : bool TaskHandle::request_cancel(std::shared_ptr<State> state) {
     340         [ +  + ]:         18 :   if VUNLIKELY (!state) {
     341                 :          3 :     return false;
     342                 :            :   }
     343                 :            : 
     344                 :         15 :   bool changed = false;
     345                 :            : 
     346                 :         15 :   CancellationRegistration parent_registration;
     347                 :            : 
     348                 :            :   {
     349         [ +  - ]:         15 :     std::lock_guard lock(state->mtx);
     350                 :            : 
     351         [ +  + ]:         15 :     if VUNLIKELY (is_terminal(state->state)) {
     352                 :          5 :       return false;
     353                 :            :     }
     354                 :            : 
     355   [ +  +  +  +  :         10 :     if VLIKELY (state->state == TaskExecutionState::kInvalid || state->state == TaskExecutionState::kQueued) {
                   +  + ]
     356                 :          9 :       state->state = TaskExecutionState::kCancelled;
     357                 :            : 
     358                 :          9 :       release_parent_registration(*state, parent_registration);
     359                 :            : 
     360                 :          9 :       changed = true;
     361                 :            :     }
     362         [ +  + ]:         15 :   }
     363                 :            : 
     364         [ +  + ]:         10 :   if VLIKELY (changed) {
     365                 :          9 :     state->cv.notify_all();
     366                 :            :   }
     367                 :            : 
     368   [ +  -  -  +  :         10 :   changed = state->cancellation_source.request_cancel() || changed;
                   -  - ]
     369                 :            : 
     370                 :         10 :   return changed;
     371                 :         15 : }
     372                 :            : 
     373                 :            : }  // namespace vlink