Preemptive cancellation of TaskContext jobs

- Enable pre-execution cancellation in TaskContext by registering a
  cancellation token in constructor that invokes PreExecuteCancel()
- Clear cancellation token in ~TaskContextImpl to prevent dangling refs
- Fix CancellationContext::CancelOperation() to swap out the token
  before calling Cancel(), avoiding re-entrancy issues
- Update BlockingCancel test to cover pre-execution cancellation

Relates-To: HERESDK-12253
Signed-off-by: Mykhailo Diachenko <ext-mykhailo.z.diachenko@here.com>

Author: diachenko-mischa

olp-cpp-sdk-core/include/olp/core/client/CancellationContext.inl

@@ -19,6 +19,7 @@
 
 #pragma once
 
+#include <utility>
 namespace olp {
 namespace client {
 
@@ -58,9 +59,10 @@ inline void CancellationContext::CancelOperation() {
     return;
   }
 
-  impl_->sub_operation_cancel_token_.Cancel();
-  impl_->sub_operation_cancel_token_ = CancellationToken();
+  auto token = CancellationToken();
+  std::swap(token, impl_->sub_operation_cancel_token_);
   impl_->is_cancelled_ = true;
+  token.Cancel();
 }
 
 inline bool CancellationContext::IsCancelled() const {

olp-cpp-sdk-core/include/olp/core/client/TaskContext.h

@@ -193,9 +193,19 @@ class CORE_API TaskContext {
         : execute_func_(std::move(execute_func)),
           callback_(std::move(callback)),
           context_(std::move(context)),
-          state_{State::PENDING} {}
+          state_{State::PENDING} {
+      context_.ExecuteOrCancelled(
+          [&]() -> CancellationToken {
+            return CancellationToken([&] { PreExecuteCancel(); });
+          },
+          []() {});
+    }
 
-    ~TaskContextImpl() override{};
+    ~TaskContextImpl() override {
+      // Overwrite the token to prevent dangling reference to this.
+      context_.ExecuteOrCancelled(
+          []() -> CancellationToken { return CancellationToken(); }, []() {});
+    };
 
     /**
      * @brief Checks for the cancellation, executes the task, and calls
@@ -249,6 +259,40 @@ class CORE_API TaskContext {
       state_.store(State::COMPLETED);
     }
 
+    void PreExecuteCancel() {
+      State expected_state = State::PENDING;
+
+      if (!state_.compare_exchange_strong(expected_state, State::IN_PROGRESS)) {
+        return;
+      }
+
+      // Moving the user callback and function guarantee that they are
+      // executed exactly once
+      ExecuteFunc function = nullptr;
+      UserCallback callback = nullptr;
+
+      {
+        std::lock_guard<std::mutex> lock(mutex_);
+        function = std::move(execute_func_);
+        callback = std::move(callback_);
+      }
+
+      Response user_response =
+          client::ApiError(client::ErrorCode::Cancelled, "Cancelled");
+
+      if (callback) {
+        callback(std::move(user_response));
+      }
+
+      // Resources need to be released before the notification, else lambas
+      // would have captured resources like network or `TaskScheduler`.
+      function = nullptr;
+      callback = nullptr;
+
+      condition_.Notify();
+      state_.store(State::COMPLETED);
+    }
+
     /**
      * @brief Cancels the operation and waits for the notification.
      *

olp-cpp-sdk-core/tests/client/TaskContextTest.cpp

@@ -136,19 +136,29 @@ TEST(TaskContextTest, BlockingCancel) {
 
   Callback callback = [&](Response r) { response = std::move(r); };
 
-  TaskContext context = TaskContext::Create(func, callback);
+  {
+    SCOPED_TRACE("Pre-exec cancellation");
+    TaskContext context = TaskContext::Create(func, callback);
+    EXPECT_TRUE(context.BlockingCancel(std::chrono::seconds(0)));
+    EXPECT_FALSE(response.IsSuccessful());
+    EXPECT_EQ(response.GetError().GetErrorCode(), ErrorCode::Cancelled);
+  }
 
-  EXPECT_FALSE(context.BlockingCancel(std::chrono::seconds(0)));
+  {
+    SCOPED_TRACE("Pre-exec cancellation race");
+    response = Response{};
+    TaskContext context = TaskContext::Create(func, callback);
 
-  std::thread cancel_thread([&]() { EXPECT_TRUE(context.BlockingCancel()); });
+    std::thread cancel_thread([&]() { EXPECT_TRUE(context.BlockingCancel()); });
 
-  std::thread execute_thread([&]() { context.Execute(); });
+    std::thread execute_thread([&]() { context.Execute(); });
 
-  execute_thread.join();
-  cancel_thread.join();
+    execute_thread.join();
+    cancel_thread.join();
 
-  EXPECT_FALSE(response.IsSuccessful());
-  EXPECT_EQ(response.GetError().GetErrorCode(), ErrorCode::Cancelled);
+    EXPECT_FALSE(response.IsSuccessful());
+    EXPECT_EQ(response.GetError().GetErrorCode(), ErrorCode::Cancelled);
+  }
 }
 
 TEST(TaskContextTest, BlockingCancelIsWaiting) {

olp-cpp-sdk-dataservice-read/tests/VersionedLayerClientImplTest.cpp

@@ -17,6 +17,7 @@
  * License-Filename: LICENSE
  */
 
+#include <future>
 #include <thread>
 
 #include <gtest/gtest.h>
@@ -1138,15 +1139,19 @@ TEST(VersionedLayerClientTest, PrefetchPartitionsCancel) {
                                         settings);
   {
     SCOPED_TRACE("Cancel request");
-    std::promise<void> block_promise;
-    auto block_future = block_promise.get_future();
-    settings.task_scheduler->ScheduleTask(
-        [&block_future]() { block_future.get(); });
+    std::promise<void> block_task;
+    std::promise<void> block_main;
+    auto block_task_future = block_task.get_future();
+    settings.task_scheduler->ScheduleTask([&block_task_future, &block_main]() {
+      block_main.set_value();
+      block_task_future.get();
+    });
     auto cancellable = client.PrefetchPartitions(request, nullptr);
 
     // cancel the request and unblock queue
     cancellable.GetCancellationToken().Cancel();
-    block_promise.set_value();
+    block_main.get_future().wait();
+    block_task.set_value();
     auto future = cancellable.GetFuture();
 
     ASSERT_EQ(future.wait_for(kTimeout), std::future_status::ready);

olp-cpp-sdk-dataservice-read/tests/VolatileLayerClientImplTest.cpp

@@ -357,12 +357,10 @@ TEST(VolatileLayerClientImplTest, GetDataCancelOnClientDestroy) {
     read::DataResponse data_response;
     {
       // Client owns the task scheduler
-      auto caller_thread_id = std::this_thread::get_id();
       read::VolatileLayerClientImpl client(kHrn, kLayerId, std::move(settings));
       client.GetData(read::DataRequest().WithPartitionId(kPartitionId),
                      [&](read::DataResponse response) {
                        data_response = std::move(response);
-                       EXPECT_NE(caller_thread_id, std::this_thread::get_id());
                      });
     }
 
@@ -1013,8 +1011,7 @@ TEST(VolatileLayerClientImplTest, PrefetchTilesCancelOnClientDestroy) {
 
     read::PrefetchTilesResponse response;
     {
-      // Client owns the task scheduler
-      auto caller_thread_id = std::this_thread::get_id();
+      // Client owns the task schedule
       read::VolatileLayerClientImpl client(kHrn, kLayerId, std::move(settings));
       std::vector<olp::geo::TileKey> tile_keys = {
           olp::geo::TileKey::FromHereTile(kTileId)};
@@ -1023,11 +1020,10 @@ TEST(VolatileLayerClientImplTest, PrefetchTilesCancelOnClientDestroy) {
                          .WithMinLevel(11)
                          .WithMaxLevel(12);
 
-      client.PrefetchTiles(
-          request, [&](read::PrefetchTilesResponse prefetch_response) {
-            response = std::move(prefetch_response);
-            EXPECT_NE(caller_thread_id, std::this_thread::get_id());
-          });
+      client.PrefetchTiles(request,
+                           [&](read::PrefetchTilesResponse prefetch_response) {
+                             response = std::move(prefetch_response);
+                           });
     }
 
     // Callback must be called during client destructor.

tests/integration/olp-cpp-sdk-dataservice-read/VersionedLayerClientPrefetchPartitionsTest.cpp

@@ -301,15 +301,20 @@ TEST_F(VersionedLayerClientPrefetchPartitionsTest, PrefetchPartitionsCancel) {
   const auto request =
       read::PrefetchPartitionsRequest().WithPartitionIds(partitions);
 
-  std::promise<void> block_promise;
-  auto block_future = block_promise.get_future();
+  std::promise<void> block_task_promise;
+  std::promise<void> block_main_promise;
+  auto block_future = block_task_promise.get_future();
   settings_.task_scheduler->ScheduleTask(
-      [&block_future]() { block_future.get(); });
+      [&block_future, &block_main_promise]() {
+        block_main_promise.set_value();
+        block_future.get();
+      });
   auto cancellable = client.PrefetchPartitions(request);
 
   // cancel the request and unblock queue
   cancellable.GetCancellationToken().Cancel();
-  block_promise.set_value();
+  block_main_promise.get_future().get();
+  block_task_promise.set_value();
   auto future = cancellable.GetFuture();
 
   ASSERT_EQ(future.wait_for(kTimeout), std::future_status::ready);