TaskContext.h

/*
 * Copyright (C) 2019-2021 HERE Europe B.V.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 * SPDX-License-Identifier: Apache-2.0
 * License-Filename: LICENSE
 */
 
#pragma once
 
#include <atomic>
#include <memory>
#include <mutex>
#include <utility>
 
#include <olp/core/client/ApiError.h>
#include <olp/core/client/CancellationContext.h>
#include <olp/core/client/CancellationToken.h>
#include <olp/core/client/Condition.h>
 
namespace olp {
namespace client {
 
class CORE_API TaskContext {
 public:
  template <typename Exec, typename Callback>
  static TaskContext Create(
      Exec execute_func, Callback callback,
      client::CancellationContext context = client::CancellationContext()) {
    TaskContext task;
    task.SetExecutors(std::move(execute_func), std::move(callback),
                      std::move(context));
    return task;
  }
 
  void Execute() const { impl_->Execute(); }
 
  bool BlockingCancel(
      std::chrono::milliseconds timeout = std::chrono::seconds(60)) const {
    return impl_->BlockingCancel(timeout);
  }
 
  client::CancellationToken CancelToken() const { return impl_->CancelToken(); }
 
  bool operator==(const TaskContext& other) const {
    return impl_ == other.impl_;
  }
 
 protected:
  friend struct TaskContextHash;
 
  TaskContext() = default;
 
  template <typename Exec, typename Callback,
#if defined(__cpp_lib_is_invocable) && __cpp_lib_is_invocable >= 201703
            typename ExecResult =
                std::invoke_result_t<Exec, client::CancellationContext>
#else
            typename ExecResult =
                typename std::result_of<Exec(client::CancellationContext)>::type
#endif
            >
  void SetExecutors(Exec execute_func, Callback callback,
                    client::CancellationContext context) {
    impl_ = std::make_shared<TaskContextImpl<ExecResult>>(
        std::move(execute_func), std::move(callback), std::move(context));
  }
 
  class Impl {
   public:
    virtual ~Impl() = default;
 
    virtual void Execute() = 0;
 
    virtual bool BlockingCancel(std::chrono::milliseconds timeout) = 0;
 
    virtual client::CancellationToken CancelToken() = 0;
  };
 
  template <typename Response>
  class TaskContextImpl : public Impl {
   public:
    using ExecuteFunc = std::function<Response(client::CancellationContext)>;
    using UserCallback = std::function<void(Response)>;
 
    TaskContextImpl(ExecuteFunc execute_func, UserCallback callback,
                    client::CancellationContext context)
        : execute_func_(std::move(execute_func)),
          callback_(std::move(callback)),
          context_(std::move(context)),
          state_{State::PENDING} {}
 
    ~TaskContextImpl() override{};
 
    void Execute() override {
      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 (function && !context_.IsCancelled()) {
        auto response = function(context_);
        // Cancel could occur during the function execution. In that case,
        // ignore the response.
        if (!context_.IsCancelled() ||
            (!response.IsSuccessful() &&
             response.GetError().GetErrorCode() == ErrorCode::RequestTimeout)) {
          user_response = std::move(response);
        }
      }
 
      // Reset the context after the task is finished.
      context_.ExecuteOrCancelled([]() { return CancellationToken(); });
 
      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);
    }
 
    bool BlockingCancel(std::chrono::milliseconds timeout) override {
      if (state_.load() == State::COMPLETED) {
        return true;
      }
 
      // Cancels the operation and waits for the notification.
      if (!context_.IsCancelled()) {
        context_.CancelOperation();
      }
 
      {
        std::lock_guard<std::mutex> lock(mutex_);
        execute_func_ = nullptr;
      }
 
      return condition_.Wait(timeout);
    }
 
    client::CancellationToken CancelToken() override {
      auto context = context_;
      return client::CancellationToken(
          [context]() mutable { context.CancelOperation(); });
    }
 
    enum class State {
      PENDING,
      IN_PROGRESS,
      COMPLETED
    };
 
    std::mutex mutex_;
    ExecuteFunc execute_func_;
    UserCallback callback_;
    client::CancellationContext context_;
    client::Condition condition_;
    std::atomic<State> state_;
  };
 
  std::shared_ptr<Impl> impl_;
};
 
struct CORE_API TaskContextHash {
  size_t operator()(const TaskContext& task_context) const {
    return std::hash<std::shared_ptr<TaskContext::Impl>>()(task_context.impl_);
  }
};
 
}  // namespace client
}  // namespace olp