Refactor thread pool implementation.
- Make mjTask non-opaque and remove C++ Task class. - Make C++ thread pool a subclass of a skeletal mjThreadPool C struct. - Make the mjTask status enum more consistent with the rest of MuJoCo. - Change mju_threadPoolEnqueue to take just the mjTask. Users must now prepare the mjTask by assigning the function pointer and argument into the struct. - Rename files in thread/ to be more consistent with the rest of MuJoCo. - Run threading tests in CMake. - Allow use of C++20 designated initializers. Otherwise the functionality remains identical. PiperOrigin-RevId: 564374675 Change-Id: I37c9894566bc39faf217e5aa97a4e2713a70e467
This commit is contained in:
Saran Tunyasuvunakool
committed by
Copybara-Service
Copybara-Service
parent
0ba10cc4f0
commit
78183e60e1
@@ -13,11 +13,11 @@
|
||||
# limitations under the License.
|
||||
|
||||
set(MUJOCO_THREAD_SRCS
|
||||
lockless_queue.h
|
||||
task.cc
|
||||
task.h
|
||||
thread_pool.cc
|
||||
thread_pool.h
|
||||
thread_queue.h
|
||||
thread_task.cc
|
||||
thread_task.h
|
||||
)
|
||||
|
||||
target_sources(mujoco PRIVATE ${MUJOCO_THREAD_SRCS})
|
||||
|
||||
@@ -1,70 +0,0 @@
|
||||
// Copyright 2023 DeepMind Technologies Limited
|
||||
//
|
||||
// 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.
|
||||
// IWYU pragma: private, include "third_party/mujoco/include/mujoco.h"
|
||||
// IWYU pragma: friend "third_party/(py/)?mujoco/.*"
|
||||
|
||||
#ifndef MUJOCO_SRC_THREAD_TASK_H_
|
||||
#define MUJOCO_SRC_THREAD_TASK_H_
|
||||
|
||||
#ifdef __cplusplus
|
||||
|
||||
#include <atomic>
|
||||
#include <new>
|
||||
#include <thread>
|
||||
|
||||
namespace mujoco {
|
||||
|
||||
class Task {
|
||||
public:
|
||||
using FunctionPtr = void* (*)(void*);
|
||||
enum Status {
|
||||
QUEUED,
|
||||
COMPLETE,
|
||||
};
|
||||
|
||||
static void Initialize(
|
||||
Task* task,
|
||||
FunctionPtr start_routine,
|
||||
void* args) {
|
||||
// instantiate a task at the pointer passed in
|
||||
new(task) Task();
|
||||
task->start_routine_ = start_routine;
|
||||
task->args_ = args;
|
||||
task->status_ = Status::QUEUED;
|
||||
}
|
||||
|
||||
void Execute() {
|
||||
args_ = start_routine_(args_);
|
||||
status_ = Status::COMPLETE;
|
||||
}
|
||||
|
||||
void Join() {
|
||||
while (status_ != Status::COMPLETE) {
|
||||
std::this_thread::yield();
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
FunctionPtr start_routine_;
|
||||
|
||||
void* args_;
|
||||
|
||||
std::atomic<Status> status_ = Status::QUEUED;
|
||||
};
|
||||
|
||||
} // namespace mujoco
|
||||
|
||||
#endif // __cplusplus
|
||||
|
||||
#endif // MUJOCO_SRC_THREAD_TASK_H_
|
||||
+96
-21
@@ -14,39 +14,114 @@
|
||||
|
||||
#include "thread/thread_pool.h"
|
||||
|
||||
#include <atomic>
|
||||
#include <cstddef>
|
||||
#include <memory>
|
||||
#include <thread>
|
||||
#include <vector>
|
||||
|
||||
#include <mujoco/mjthread.h>
|
||||
#include <mujoco/mujoco.h>
|
||||
#include "thread/task.h"
|
||||
#include "engine/engine_crossplatform.h"
|
||||
#include "engine/engine_util_errmem.h"
|
||||
#include "thread/thread_queue.h"
|
||||
#include "thread/thread_task.h"
|
||||
|
||||
static constexpr size_t kMaxThreads = 128;
|
||||
namespace mujoco {
|
||||
namespace {
|
||||
constexpr size_t kThreadPoolQueueSize = 640;
|
||||
|
||||
struct WorkerThread {
|
||||
// Shutdown function passed to running threads to ensure clean shutdown.
|
||||
static void* ShutdownFunction(void* args) {
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
// Thread for the worker.
|
||||
std::unique_ptr<std::thread> thread_;
|
||||
|
||||
// An mjTask for shutting down this worker.
|
||||
mjTask shutdown_task_ {
|
||||
.func = &ShutdownFunction,
|
||||
.args = nullptr,
|
||||
};
|
||||
};
|
||||
} // namespace
|
||||
|
||||
// Concrete C++ class definition for mjThreadPool.
|
||||
// (The public mjThreadPool C struct is an opaque one.)
|
||||
class ThreadPoolImpl : public mjThreadPool {
|
||||
public:
|
||||
ThreadPoolImpl(int num_worker) : mjThreadPool{.nworker = num_worker} {
|
||||
// initialize worker threads
|
||||
for (int i = 0; i < num_worker; ++i) {
|
||||
workers_.push_back(
|
||||
{std::make_unique<std::thread>(ThreadPoolWorker, this)});
|
||||
}
|
||||
}
|
||||
|
||||
// start a task in the threadpool
|
||||
void Enqueue(mjTask* task) {
|
||||
if (mjUNLIKELY(GetAtomicTaskStatus(task).exchange(mjTASK_QUEUED) !=
|
||||
mjTASK_NEW)) {
|
||||
mjERROR("task->status is not mjTASK_NEW");
|
||||
}
|
||||
lockless_queue_.push(task);
|
||||
}
|
||||
|
||||
// shutdown the threadpool
|
||||
void Shutdown() {
|
||||
if (shutdown_) {
|
||||
return;
|
||||
}
|
||||
|
||||
shutdown_ = true;
|
||||
std::vector<mjTask> shutdown_tasks(workers_.size());
|
||||
for (auto& worker : workers_) {
|
||||
Enqueue(&worker.shutdown_task_);
|
||||
}
|
||||
|
||||
for (auto& worker : workers_) {
|
||||
worker.thread_->join();
|
||||
}
|
||||
}
|
||||
|
||||
~ThreadPoolImpl() { Shutdown(); }
|
||||
|
||||
private:
|
||||
// method executed by running threads
|
||||
static void ThreadPoolWorker(ThreadPoolImpl* thread_pool) {
|
||||
while (!thread_pool->shutdown_) {
|
||||
auto task = static_cast<mjTask*>(thread_pool->lockless_queue_.pop());
|
||||
task->args = task->func(task->args);
|
||||
GetAtomicTaskStatus(task).store(mjTASK_COMPLETED);
|
||||
}
|
||||
}
|
||||
|
||||
// indicates whether the thread pool is being shut down
|
||||
std::atomic<bool> shutdown_ = false;
|
||||
|
||||
// OS threads that are running in this pool
|
||||
std::vector<WorkerThread> workers_;
|
||||
|
||||
// queue of tasks to execute
|
||||
mujoco::LocklessQueue<void*, kThreadPoolQueueSize> lockless_queue_;
|
||||
};
|
||||
|
||||
// create a thread pool
|
||||
mjThreadPool* mju_threadPoolCreate(size_t number_of_threads) {
|
||||
mujoco::ThreadPool<kMaxThreads>* thread_pool =
|
||||
new mujoco::ThreadPool<kMaxThreads>(number_of_threads);
|
||||
return static_cast<mjThreadPool*>(static_cast<void*>(thread_pool));
|
||||
return new ThreadPoolImpl(number_of_threads);
|
||||
}
|
||||
|
||||
// start a task in the threadpool
|
||||
void mju_threadPoolEnqueue(
|
||||
mjThreadPool* thread_pool, mjTask* task, mjStartRoutine start_routine,
|
||||
void* args) {
|
||||
mujoco::ThreadPool<kMaxThreads>* thread_pool_ptr =
|
||||
static_cast<mujoco::ThreadPool<kMaxThreads>*>(
|
||||
static_cast<void*>(thread_pool));
|
||||
thread_pool_ptr->Enqueue(
|
||||
static_cast<mujoco::Task*>(static_cast<void*>(task)), start_routine,
|
||||
args);
|
||||
void mju_threadPoolEnqueue(mjThreadPool* thread_pool, mjTask* task) {
|
||||
auto thread_pool_impl = static_cast<ThreadPoolImpl*>(thread_pool);
|
||||
thread_pool_impl->Enqueue(task);
|
||||
}
|
||||
|
||||
// shutdown the threadpool and free the memory
|
||||
void mju_threadPoolDestroy(mjThreadPool* thread_pool) {
|
||||
mujoco::ThreadPool<kMaxThreads>* thread_pool_ptr =
|
||||
static_cast<mujoco::ThreadPool<kMaxThreads>*>(
|
||||
static_cast<void*>(thread_pool));
|
||||
thread_pool_ptr->Shutdown();
|
||||
delete thread_pool_ptr;
|
||||
auto thread_pool_impl = static_cast<ThreadPoolImpl*>(thread_pool);
|
||||
thread_pool_impl->Shutdown();
|
||||
delete thread_pool_impl;
|
||||
}
|
||||
|
||||
} // namespace mujoco
|
||||
|
||||
+15
-75
@@ -11,92 +11,32 @@
|
||||
// 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.
|
||||
// IWYU pragma: private, include "third_party/mujoco/include/mujoco.h"
|
||||
// IWYU pragma: friend "third_party/(py/)?mujoco/.*"
|
||||
|
||||
#ifndef MUJOCO_SRC_THREAD_THREAD_POOL_H_
|
||||
#define MUJOCO_SRC_THREAD_THREAD_POOL_H_
|
||||
|
||||
#include <stddef.h>
|
||||
|
||||
#include <mujoco/mjexport.h>
|
||||
#include <mujoco/mjthread.h>
|
||||
|
||||
#ifdef __cplusplus
|
||||
|
||||
#include <atomic>
|
||||
#include <cstddef>
|
||||
#include <thread>
|
||||
|
||||
#include "thread/lockless_queue.h"
|
||||
#include "thread/task.h"
|
||||
|
||||
namespace mujoco {
|
||||
extern "C" {
|
||||
#endif
|
||||
|
||||
static constexpr size_t kThreadPoolQueueSize = 640;
|
||||
// Create a thread pool with the specified number of threads running.
|
||||
MJAPI mjThreadPool* mju_threadPoolCreate(size_t number_of_threads);
|
||||
|
||||
template <size_t max_number_of_threads>
|
||||
class ThreadPool {
|
||||
public:
|
||||
ThreadPool(size_t number_of_threads)
|
||||
: number_of_threads_(number_of_threads) {
|
||||
for (int i = 0; i < number_of_threads_; ++i) {
|
||||
threads_[i] = std::thread(ThreadPoolWorker, static_cast<void*>(this));
|
||||
}
|
||||
}
|
||||
// Enqueue a task in a thread pool.
|
||||
MJAPI void mju_threadPoolEnqueue(mjThreadPool* thread_pool, mjTask* task);
|
||||
|
||||
// start a task in the threadpool
|
||||
void Enqueue(
|
||||
Task* task, Task::FunctionPtr start_routine, void* args) {
|
||||
Task::Initialize(task, start_routine, args);
|
||||
lockless_queue_.push(static_cast<void*>(task));
|
||||
}
|
||||
|
||||
// shutdown the threadpool
|
||||
void Shutdown() {
|
||||
if (shutdown_) {
|
||||
return;
|
||||
}
|
||||
|
||||
shutdown_ = true;
|
||||
Task shutdown_tasks[max_number_of_threads];
|
||||
for (int i = 0; i < number_of_threads_; ++i) {
|
||||
Enqueue(&shutdown_tasks[i], ShutdownFunction, nullptr);
|
||||
}
|
||||
|
||||
for (int i = 0; i < number_of_threads_; ++i) {
|
||||
threads_[i].join();
|
||||
}
|
||||
}
|
||||
|
||||
~ThreadPool() { Shutdown(); }
|
||||
|
||||
private:
|
||||
// method executed by running threads
|
||||
static void ThreadPoolWorker(void* arg) {
|
||||
ThreadPool<max_number_of_threads>* thread_pool =
|
||||
static_cast<ThreadPool<max_number_of_threads>*>(arg);
|
||||
while (!thread_pool->shutdown_) {
|
||||
Task* task = static_cast<Task*>(thread_pool->lockless_queue_.pop());
|
||||
task->Execute();
|
||||
}
|
||||
}
|
||||
|
||||
// shutdown function passed to running threads to ensure cleans shutdown
|
||||
static void* ShutdownFunction(void* args) {
|
||||
return NULL;
|
||||
}
|
||||
|
||||
// is the thread pool is being shut down
|
||||
std::atomic<bool> shutdown_ = false;
|
||||
|
||||
// actual number of running threads in the threadpool
|
||||
const size_t number_of_threads_;
|
||||
|
||||
// OS threads that are running in this pool
|
||||
std::thread threads_[max_number_of_threads];
|
||||
|
||||
// queue of tasks to execute
|
||||
LocklessQueue<void*, kThreadPoolQueueSize> lockless_queue_;
|
||||
};
|
||||
// Destroy a thread pool.
|
||||
MJAPI void mju_threadPoolDestroy(mjThreadPool* thread_pool);
|
||||
|
||||
#ifdef __cplusplus
|
||||
} // extern "C"
|
||||
} // namespace mujoco
|
||||
|
||||
#endif // __cplusplus
|
||||
|
||||
#endif // MUJOCO_SRC_THREAD_THREAD_POOL_H_
|
||||
|
||||
@@ -11,8 +11,6 @@
|
||||
// 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.
|
||||
// IWYU pragma: private, include "third_party/mujoco/include/mujoco.h"
|
||||
// IWYU pragma: friend "third_party/(py/)?mujoco/.*"
|
||||
|
||||
#ifndef MUJOCO_SRC_THREAD_LOCKLESS_QUEUE_H_
|
||||
#define MUJOCO_SRC_THREAD_LOCKLESS_QUEUE_H_
|
||||
@@ -12,13 +12,22 @@
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#include "thread/task.h"
|
||||
#include "thread/thread_task.h"
|
||||
|
||||
#include <thread>
|
||||
|
||||
#include <mujoco/mjthread.h>
|
||||
#include <mujoco/mujoco.h>
|
||||
|
||||
// waits for a task to complete
|
||||
void mju_taskJoin(mjTask* task) {
|
||||
mujoco::Task* task_ptr = static_cast<mujoco::Task*>(static_cast<void*>(task));
|
||||
task_ptr->Join();
|
||||
namespace mujoco {
|
||||
void mju_defaultTask(mjTask* task) {
|
||||
task->func = nullptr;
|
||||
task->args = nullptr;
|
||||
task->status = mjTASK_NEW;
|
||||
}
|
||||
|
||||
void mju_taskJoin(mjTask* task) {
|
||||
while (GetAtomicTaskStatus(task) != mjTASK_COMPLETED) {
|
||||
std::this_thread::yield();
|
||||
}
|
||||
}
|
||||
} // namespace mujoco
|
||||
@@ -0,0 +1,50 @@
|
||||
// Copyright 2023 DeepMind Technologies Limited
|
||||
//
|
||||
// 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.
|
||||
|
||||
#ifndef MUJOCO_SRC_THREAD_THREAD_TASK_H_
|
||||
#define MUJOCO_SRC_THREAD_THREAD_TASK_H_
|
||||
|
||||
#include <atomic>
|
||||
#include <new>
|
||||
#include <type_traits>
|
||||
|
||||
#include <mujoco/mjexport.h>
|
||||
#include <mujoco/mjthread.h>
|
||||
|
||||
#ifdef __cplusplus
|
||||
namespace mujoco {
|
||||
extern "C" {
|
||||
#endif
|
||||
|
||||
// Initialize an mjTask.
|
||||
MJAPI void mju_defaultTask(mjTask* task);
|
||||
|
||||
// Wait for a task to complete.
|
||||
MJAPI void mju_taskJoin(mjTask* task);
|
||||
|
||||
#ifdef __cplusplus
|
||||
} // extern "C"
|
||||
|
||||
using TaskStatus = std::remove_volatile_t<decltype(mjTask::status)>;
|
||||
inline std::atomic<TaskStatus>& GetAtomicTaskStatus(mjTask* task) {
|
||||
static_assert(sizeof(std::atomic<TaskStatus>) == sizeof(TaskStatus));
|
||||
static_assert(alignof(std::atomic<TaskStatus>) == alignof(TaskStatus));
|
||||
static_assert(std::atomic<TaskStatus>::is_always_lock_free);
|
||||
return *std::launder(reinterpret_cast<std::atomic<TaskStatus>*>(
|
||||
const_cast<TaskStatus*>(&task->status)));
|
||||
}
|
||||
} // namespace mujoco
|
||||
#endif // __cplusplus
|
||||
|
||||
#endif // MUJOCO_SRC_THREAD_THREAD_TASK_H_
|
||||
Reference in New Issue
Block a user