scheduler.hpp

// Copyright 2024 Cellumation GmbH.
//
// 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.
#pragma once
 
#include <algorithm>
#include <condition_variable>
#include <deque>
#include <functional>
#include <list>
#include <memory>
#include <mutex>
#include <utility>
 
#include <rclcpp/callback_group.hpp>
#include "global_event_id_provider.hpp"
 
namespace rclcpp
{
namespace executors
{
namespace cbg_executor
{
struct Worker
{
  std::mutex mutex;
  std::condition_variable condition_variable;
  std::atomic<bool> wakeup = false;
 
  void prepare_block()
  {
    wakeup = false;
  }
 
  void block()
  {
    std::unique_lock lk(mutex);
    condition_variable.wait(lk, [this]() -> bool {
        return wakeup;
    });
  }
 
  void block_for(std::chrono::nanoseconds timeout)
  {
    std::unique_lock lk(mutex);
    condition_variable.wait_for(lk, timeout, [this]() -> bool {
        return wakeup;
    });
  }
 
  void unblock()
  {
    {
      std::unique_lock lk(mutex);
      wakeup = true;
    }
    condition_variable.notify_one();
  }
 
  /*
   * This function shall be used if there is a risk that the worker
   * might get deleted after wakeup. This function is supposed to
   * be used in the shutdown case.
   * This function is less performant as the unblock function, as
   * the notofied thread may be tempoarlily blocked again until the
   * mutex is released.
   */
  void unblock_thread_safe()
  {
    std::unique_lock lk(mutex);
    wakeup = true;
    condition_variable.notify_one();
  }
};
 
struct WorkerQueue
{
  std::mutex workers_mutex;
  std::deque<Worker *> workers;
 
  bool release_workers = false;
 
  Worker * pop_blocked_worker_thread()
  {
    if(workers.empty()) {
      // no threads available
      return nullptr;
    }
    Worker * worker = workers.front();
    workers.pop_front();
    return worker;
  }
 
  bool push_blocked_worker_thread(Worker * worker)
  {
    if(release_workers) {
      return false;
    }
    workers.push_front(worker);
    return true;
  }
 
  void remove_worker_thread(Worker * worker)
  {
    auto it = std::find(workers.begin(), workers.end(), worker);
    if(it != workers.end()) {
      workers.erase(it);
    }
  }
 
  std::deque<Worker *> release_all_worker_threads()
  {
    std::deque<Worker *> cpy;
    cpy.swap(workers);
    release_workers = true;
    return cpy;
  }
};
 
class CBGScheduler
{
public:
  struct WaitableWithEventType
  {
    rclcpp::Waitable::WeakPtr waitable;
    int internal_event_type;
 
    bool expired() const
    {
      return waitable.expired();
    }
  };
 
  struct CallbackEventType
  {
    explicit CallbackEventType(std::function<void()> callback)
    : callback(std::move(callback))
    {
    }
 
    std::function<void()> callback;
 
    bool expired() const
    {
      return false;
    }
  };
 
  struct CallbackGroupHandle
  {
    explicit CallbackGroupHandle(CBGScheduler & scheduler, CallbackGroupType type)
    : scheduler(scheduler), type(type)
    {
    }
 
    CallbackGroupHandle(const CallbackGroupHandle &) = delete;
    CallbackGroupHandle(CallbackGroupHandle &&) = delete;
 
    virtual ~CallbackGroupHandle() = default;
 
    CallbackGroupHandle & operator=(const CallbackGroupHandle &) = delete;
    CallbackGroupHandle & operator=(CallbackGroupHandle &&) = delete;
 
    virtual std::function<void(size_t)> get_ready_callback_for_entity(
      const rclcpp::SubscriptionBase::WeakPtr & entity) = 0;
    virtual std::function<void(std::function<void()> executed_callback)>
    get_ready_callback_for_entity(const rclcpp::TimerBase::WeakPtr & entity) = 0;
    virtual std::function<void(size_t)> get_ready_callback_for_entity(
      const rclcpp::ClientBase::WeakPtr & entity) = 0;
    virtual std::function<void(size_t)> get_ready_callback_for_entity(
      const rclcpp::ServiceBase::WeakPtr & entity) = 0;
    virtual std::function<void(size_t,
      int)> get_ready_callback_for_entity(const rclcpp::Waitable::WeakPtr & entity) = 0;
    virtual std::function<void(size_t)> get_ready_callback_for_entity(
      const CallbackEventType & entity) = 0;
 
    void mark_as_executed()
    {
      {
        std::lock_guard l(ready_mutex);
        not_ready = false;
 
        if(!has_ready_entities()) {
          idle = true;
          return;
        }
      }
      // inform scheduler that we have more work
      scheduler.callback_group_ready(this, false);
    }
 
    CallbackGroupType get_type() {return type;}
 
    bool is_ready();
 
    // true if this cbg is inside the scheduler's queue
    bool in_queue = false;
 
protected:
    CBGScheduler & scheduler;
 
    virtual bool has_ready_entities() const = 0;
 
    template<typename add_fun>
    void add_ready_entity(const add_fun & fun)
    {
      {
        std::lock_guard l(ready_mutex);
 
        fun();
 
        if(not_ready || !idle) {
          return;
        }
 
        idle = false;
      }
 
      // If we reached this point, we were idle and now have work,
      // therefore we need to move this callback group into the list
      // of ready callback groups.
      scheduler.callback_group_ready(this, true);
    }
 
    void mark_as_skipped()
    {
      if(!has_ready_entities()) {
        idle = true;
      }
    }
 
    void mark_as_executing()
    {
      if (type != CallbackGroupType::Reentrant) {
        not_ready = true;
      }
    }
 
    std::mutex ready_mutex;
 
private:
    // will be set if cbg is mutual exclusive and something is executing
    bool not_ready = false;
 
    // true, if nothing is beeing executed, and there are no pending events
    bool idle = true;
 
    // type of the underlying callback group
    CallbackGroupType type;
  };
 
  struct ExecutableEntity
  {
    // if called executes the entity
    std::function<void()> execute_function;
    // The callback group associated with the entity. Can be nullptr.
    CallbackGroupHandle *callback_handle = nullptr;
  };
 
  explicit CBGScheduler(std::function<void ()> sync_function)
  : sync_function(sync_function) {}
  CBGScheduler(const CBGScheduler &) = delete;
  CBGScheduler(CBGScheduler &&) = delete;
  virtual ~CBGScheduler() = default;
 
  CBGScheduler & operator=(const CBGScheduler &) = delete;
  CBGScheduler & operator=(CBGScheduler &&) = delete;
 
  CallbackGroupHandle * add_callback_group(const rclcpp::CallbackGroup::SharedPtr & callback_group)
  {
    auto uPtr = get_handle_for_callback_group(callback_group);
    CallbackGroupHandle * ret = uPtr.get();
 
    std::lock_guard lk(ready_callback_groups_mutex);
 
    callback_groups.push_back(std::move(uPtr));
    return ret;
  }
 
  void remove_callback_group(const CallbackGroupHandle *callback_handle)
  {
    std::lock_guard lk(ready_callback_groups_mutex);
 
    auto cbg_it = std::find(ready_callback_groups.begin(),
          ready_callback_groups.end(), callback_handle);
    if (cbg_it != ready_callback_groups.end()) {
      ready_callback_groups.erase(cbg_it);
    }
 
    callback_groups.remove_if([&callback_handle] (const auto & e) {
        return e.get() == callback_handle;
    });
  }
 
  void callback_group_ready(CallbackGroupHandle *handle, bool callback_group_was_idle)
  {
    {
      std::lock_guard l(ready_callback_groups_mutex);
 
      if (!handle->in_queue) {
        ready_callback_groups.push_back(handle);
        handle->in_queue = true;
      }
    }
 
    if(callback_group_was_idle) {
      unblock_one_worker_thread();
    }
  }
 
  struct ExecutableEntityWithInfo
  {
    std::optional<ExecutableEntity> entity;
    bool moreEntitiesReady{};
  };
 
  ExecutableEntityWithInfo get_next_ready_entity()
  {
    std::lock_guard l(ready_callback_groups_mutex);
    worker_checking_for_work = false;
 
    if(needs_sync) {
      needs_sync = false;
      return ExecutableEntityWithInfo{.entity =
          ExecutableEntity{.execute_function = sync_function, .callback_handle = nullptr},
        .moreEntitiesReady = false};
    }
 
    return get_next_ready_entity_intern();
  }
 
  ExecutableEntityWithInfo get_next_ready_entity(
    GlobalEventIdProvider::MonotonicId max_id)
  {
    std::lock_guard l(ready_callback_groups_mutex);
    worker_checking_for_work = false;
 
    if(needs_sync) {
      needs_sync = false;
      return ExecutableEntityWithInfo{.entity =
          ExecutableEntity{.execute_function = sync_function, .callback_handle = nullptr},
        .moreEntitiesReady = false};
    }
 
    return get_next_ready_entity_intern(max_id);
  }
 
  virtual ExecutableEntityWithInfo get_next_ready_entity_intern() = 0;
  virtual ExecutableEntityWithInfo get_next_ready_entity_intern(
    GlobalEventIdProvider::MonotonicId max_id) = 0;
 
  void trigger_sync()
  {
    bool wake_worker = false;
    {
      std::lock_guard l(ready_callback_groups_mutex);
      if(!needs_sync) {
        wake_worker = true;
      }
      needs_sync = true;
    }
    if(wake_worker) {
      unblock_one_worker_thread();
    }
  }
 
  void mark_entity_as_executed(const ExecutableEntity & e)
  {
    if(e.callback_handle != nullptr) {
      e.callback_handle->mark_as_executed();
    }
  }
 
  void suppress_thread_wakeup()
  {
    std::lock_guard l(ready_callback_groups_mutex);
    worker_checking_for_work = true;
  }
 
  void unblock_one_worker_thread()
  {
    Worker * worker = nullptr;
    {
      std::scoped_lock l(ready_callback_groups_mutex, worker_queue.workers_mutex);
      if(worker_checking_for_work) {
        return;
      }
      worker = worker_queue.pop_blocked_worker_thread();
      if(worker == nullptr) {
        return;
      }
      worker_checking_for_work = true;
    }
    worker->unblock();
  }
 
  void block_worker_thread(Worker * worker)
  {
    if(!prepare_and_enqueue_worker(worker)) {
      return;
    }
    worker->block();
  }
 
  void block_worker_thread_for(Worker * worker, std::chrono::nanoseconds timeout)
  {
    if(!prepare_and_enqueue_worker(worker)) {
      return;
    }
    worker->block_for(timeout);
    // on timeout the worker is still in the queue, remove it, as the
    // worker may not be valid any more after this call
    std::lock_guard lk(worker_queue.workers_mutex);
    worker_queue.remove_worker_thread(worker);
  }
 
  void release_all_worker_threads()
  {
    std::deque<Worker *> workers;
    {
      std::lock_guard lk(worker_queue.workers_mutex);
      workers = worker_queue.release_all_worker_threads();
    }
    for(Worker * worker : workers) {
      worker->unblock_thread_safe();
    }
  }
 
protected:
  virtual std::unique_ptr<CallbackGroupHandle> get_handle_for_callback_group(
    const rclcpp::CallbackGroup::SharedPtr & callback_group) = 0;
 
  bool prepare_and_enqueue_worker(Worker * worker)
  {
    worker->prepare_block();
    std::scoped_lock l(ready_callback_groups_mutex, worker_queue.workers_mutex);
    if(needs_sync || !ready_callback_groups.empty()) {
      return false;
    }
    return worker_queue.push_blocked_worker_thread(worker);
  }
 
  // sync function, will be triggered if the executor needs
  // resync. E.g. if entities / cbg or nodes were added / removed
  std::function<void ()> sync_function;
 
  bool needs_sync = false;
 
  std::mutex ready_callback_groups_mutex;
  std::deque<CallbackGroupHandle *> ready_callback_groups;
 
  bool worker_checking_for_work = false;
 
  WorkerQueue worker_queue;
 
  std::list<std::unique_ptr<CallbackGroupHandle>> callback_groups;
};
}  // namespace cbg_executor
}  // namespace executors
}  // namespace rclcpp