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

cmake/MujocoOptions.cmake

@@ -99,9 +99,12 @@ if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU" OR (CMAKE_CXX_COMPILER_ID MATCHES "Clang
   if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU")
     # Set -Wimplicit-fallthrough=5 to only allow fallthrough annotation via __attribute__.
     set(EXTRA_COMPILE_OPTIONS ${EXTRA_COMPILE_OPTIONS} -Wimplicit-fallthrough=5
-                              -Wno-maybe-uninitialized
+                              -Wno-maybe-uninitialized -Wno-c++20-extensions
     )
   endif()
+  if(CMAKE_CXX_COMPILER_ID MATCHES "Clang")
+    set(EXTRA_COMPILE_OPTIONS ${EXTRA_COMPILE_OPTIONS} -Wno-c++20-designator)
+  endif()
 endif()
 
 include(MujocoHarden)

doc/APIreference/functions.rst

@@ -3444,7 +3444,7 @@ mju_threadPoolCreate
 
 .. mujoco-include:: mju_threadPoolCreate
 
-Creates a thread pool with the specified number of threads running.
+Create a thread pool with the specified number of threads running.
 
 .. _mju_threadPoolEnqueue:
 
@@ -3453,16 +3453,7 @@ mju_threadPoolEnqueue
 
 .. mujoco-include:: mju_threadPoolEnqueue
 
-Enqueues a task in a thread pool.
+Enqueue a task in a thread pool.
-
-.. _mju_taskJoin:
-
-mju_taskJoin
-~~~~~~~~~~~~
-
-.. mujoco-include:: mju_taskJoin
-
-Waits for a task to complete.
 
 .. _mju_threadPoolDestroy:
 
@@ -3471,5 +3462,23 @@ mju_threadPoolDestroy
 
 .. mujoco-include:: mju_threadPoolDestroy
 
-Destroys a thread pool.
+Destroy a thread pool.
+
+.. _mju_defaultTask:
+
+mju_defaultTask
+~~~~~~~~~~~~~~~
+
+.. mujoco-include:: mju_defaultTask
+
+Initialize an mjTask.
+
+.. _mju_taskJoin:
+
+mju_taskJoin
+~~~~~~~~~~~~
+
+.. mujoco-include:: mju_taskJoin
+
+Wait for a task to complete.
 

doc/includes/references.h

@@ -1443,14 +1443,21 @@ struct mjrContext_ {              // custom OpenGL context
   int     readPixelFormat;        // default color pixel format for mjr_readPixels
 };
 typedef struct mjrContext_ mjrContext;
-struct mjTask_ {
+typedef enum mjtTaskStatus_ {  // status values for mjTask
-  char buffer[24];
+  mjTASK_NEW = 0,              // newly created
-};
+  mjTASK_QUEUED,               // enqueued in a thread pool
-typedef struct mjTask_ mjTask;
+  mjTASK_COMPLETED             // completed execution
+} mjtTaskStatus;
 struct mjThreadPool_ {
-  char buffer[6208];
+  int nworker;  // number of workers in the pool
 };
 typedef struct mjThreadPool_ mjThreadPool;
+struct mjTask_ {        // a task that can be executed by a thread pool.
+  mjfTask func;         // pointer to the function that implements the task
+  void* args;           // arguments to func
+  volatile int status;  // status of the task
+};
+typedef struct mjTask_ mjTask;
 typedef enum mjtButton_ {         // mouse button
   mjBUTTON_NONE = 0,              // no button
   mjBUTTON_LEFT,                  // left button
@@ -2588,9 +2595,8 @@ int mjp_resourceProviderCount(void);
 const mjpResourceProvider* mjp_getResourceProvider(const char* resource_name);
 const mjpResourceProvider* mjp_getResourceProviderAtSlot(int slot);
 mjThreadPool* mju_threadPoolCreate(size_t number_of_threads);
-void mju_threadPoolEnqueue(
+void mju_threadPoolEnqueue(mjThreadPool* thread_pool, mjTask* task);
-ThreadPool* thread_pool, mjTask* task, void*(start_routine)(void*),
-id* args);
-void mju_taskJoin(mjTask* task);
 void mju_threadPoolDestroy(mjThreadPool* thread_pool);
+void mju_defaultTask(mjTask* task);
+void mju_taskJoin(mjTask* task);
 // NOLINTEND

include/mujoco/mjthread.h

@@ -15,33 +15,26 @@
 #ifndef MUJOCO_INCLUDE_MJTHREAD_H_
 #define MUJOCO_INCLUDE_MJTHREAD_H_
 
-// C API for MuJoCo threading
+typedef enum mjtTaskStatus_ {  // status values for mjTask
-#ifdef __cplusplus
+  mjTASK_NEW = 0,              // newly created
-extern "C" {
+  mjTASK_QUEUED,               // enqueued in a thread pool
-#endif
+  mjTASK_COMPLETED             // completed execution
+} mjtTaskStatus;
 
-#include <stddef.h>
+// function pointer type for mjTask
-
+typedef void* (*mjfTask)(void*);
-#include <mujoco/mjexport.h>
-
-// These types are implemented in C++, they're just used as opaque pointers in C
-// to provide type safety for functions.
-struct mjTask_ {
-  char buffer[24];
-};
-typedef struct mjTask_ mjTask;
 
+// An opaque type representing a thread pool.
 struct mjThreadPool_ {
-  char buffer[6208];
+  int nworker;  // number of workers in the pool
 };
 typedef struct mjThreadPool_ mjThreadPool;
 
-typedef void*(*mjStartRoutine_)(void*);
+struct mjTask_ {        // a task that can be executed by a thread pool.
-typedef mjStartRoutine_ mjStartRoutine;
+  mjfTask func;         // pointer to the function that implements the task
-
+  void* args;           // arguments to func
-#ifdef __cplusplus
+  volatile int status;  // status of the task
-}
+};
-#endif
+typedef struct mjTask_ mjTask;
-
 
 #endif  // MUJOCO_INCLUDE_MJTHREAD_H_

include/mujoco/mujoco.h

@@ -1307,19 +1307,21 @@ MJAPI const mjpResourceProvider* mjp_getResourceProviderAtSlot(int slot);
 
 //---------------------- Thread -------------------------------------------------------------------
 
-// Creates a thread pool with the specified number of threads running.
+// Create a thread pool with the specified number of threads running.
 MJAPI mjThreadPool* mju_threadPoolCreate(size_t number_of_threads);
 
-// Enqueues a task in a thread pool.
+// Enqueue a task in a thread pool.
-MJAPI void mju_threadPoolEnqueue(
+MJAPI void mju_threadPoolEnqueue(mjThreadPool* thread_pool, mjTask* task);
-    mjThreadPool* thread_pool, mjTask* task, void*(start_routine)(void*),
-    void* args);
 
-// Waits for a task to complete.
+// Destroy a thread pool.
+MJAPI void mju_threadPoolDestroy(mjThreadPool* thread_pool);
+
+// Initialize an mjTask.
+MJAPI void mju_defaultTask(mjTask* task);
+
+// Wait for a task to complete.
 MJAPI void mju_taskJoin(mjTask* task);
 
-// Destroys a thread pool.
-MJAPI void mju_threadPoolDestroy(mjThreadPool* thread_pool);
 
 #if defined(__cplusplus)
 }

introspect/codegen/generate_enums.py

@@ -56,7 +56,10 @@ class MjEnumVisitor:
       child_kind = child.get('kind')
       if child_kind == 'EnumConstantDecl':
         next_idx = values[-1][1] + 1 if values else 0
-        value = int(child['inner'][0].get('value', next_idx))
+        if 'inner' in child:
+          value = int(child['inner'][0].get('value', next_idx))
+        else:
+          value = next_idx
         values.append((child['name'], value))
     return ast_nodes.EnumDecl(name=name, declname=name, values=dict(values))
 

introspect/enums.py

@@ -642,6 +642,16 @@ ENUMS: Mapping[str, EnumDecl] = dict([
              ('mjFONT_BIG', 2),
          ]),
      )),
+    ('mjtTaskStatus',
+     EnumDecl(
+         name='mjtTaskStatus',
+         declname='enum mjtTaskStatus_',
+         values=dict([
+             ('mjTASK_NEW', 0),
+             ('mjTASK_QUEUED', 1),
+             ('mjTASK_COMPLETED', 2),
+         ]),
+     )),
     ('mjtButton',
      EnumDecl(
          name='mjtButton',

introspect/functions.py

@@ -8269,7 +8269,7 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
                  type=ValueType(name='size_t'),
              ),
          ),
-         doc='Creates a thread pool with the specified number of threads running.',  # pylint: disable=line-too-long
+         doc='Create a thread pool with the specified number of threads running.',  # pylint: disable=line-too-long
      )),
     ('mju_threadPoolEnqueue',
      FunctionDecl(
@@ -8288,32 +8288,8 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
                      inner_type=ValueType(name='mjTask'),
                  ),
              ),
-             FunctionParameterDecl(
-                 name='start_routine',
-                 type=ValueType(name='void *(*)(void *)'),
-             ),
-             FunctionParameterDecl(
-                 name='args',
-                 type=PointerType(
-                     inner_type=ValueType(name='void'),
-                 ),
-             ),
          ),
-         doc='Enqueues a task in a thread pool.',
+         doc='Enqueue a task in a thread pool.',
-     )),
-    ('mju_taskJoin',
-     FunctionDecl(
-         name='mju_taskJoin',
-         return_type=ValueType(name='void'),
-         parameters=(
-             FunctionParameterDecl(
-                 name='task',
-                 type=PointerType(
-                     inner_type=ValueType(name='mjTask'),
-                 ),
-             ),
-         ),
-         doc='Waits for a task to complete.',
      )),
     ('mju_threadPoolDestroy',
      FunctionDecl(
@@ -8327,6 +8303,34 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
                  ),
              ),
          ),
-         doc='Destroys a thread pool.',
+         doc='Destroy a thread pool.',
+     )),
+    ('mju_defaultTask',
+     FunctionDecl(
+         name='mju_defaultTask',
+         return_type=ValueType(name='void'),
+         parameters=(
+             FunctionParameterDecl(
+                 name='task',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjTask'),
+                 ),
+             ),
+         ),
+         doc='Initialize an mjTask.',
+     )),
+    ('mju_taskJoin',
+     FunctionDecl(
+         name='mju_taskJoin',
+         return_type=ValueType(name='void'),
+         parameters=(
+             FunctionParameterDecl(
+                 name='task',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjTask'),
+                 ),
+             ),
+         ),
+         doc='Wait for a task to complete.',
      )),
 ])

introspect/structs.py

@@ -7007,33 +7007,39 @@ STRUCTS: Mapping[str, StructDecl] = dict([
              ),
          ),
      )),
-    ('mjTask',
-     StructDecl(
-         name='mjTask',
-         declname='struct mjTask_',
-         fields=(
-             StructFieldDecl(
-                 name='buffer',
-                 type=ArrayType(
-                     inner_type=ValueType(name='char'),
-                     extents=(24,),
-                 ),
-                 doc='',
-             ),
-         ),
-     )),
     ('mjThreadPool',
      StructDecl(
          name='mjThreadPool',
          declname='struct mjThreadPool_',
          fields=(
              StructFieldDecl(
-                 name='buffer',
+                 name='nworker',
-                 type=ArrayType(
+                 type=ValueType(name='int'),
-                     inner_type=ValueType(name='char'),
+                 doc='number of workers in the pool',
-                     extents=(6208,),
+             ),
+         ),
+     )),
+    ('mjTask',
+     StructDecl(
+         name='mjTask',
+         declname='struct mjTask_',
+         fields=(
+             StructFieldDecl(
+                 name='func',
+                 type=ValueType(name='mjfTask'),
+                 doc='pointer to the function that implements the task',
+             ),
+             StructFieldDecl(
+                 name='args',
+                 type=PointerType(
+                     inner_type=ValueType(name='void'),
                  ),
-                 doc='',
+                 doc='arguments to func',
+             ),
+             StructFieldDecl(
+                 name='status',
+                 type=ValueType(name='int', is_volatile=True),
+                 doc='status of the task',
              ),
          ),
      )),

sample/cmake/SampleOptions.cmake

@@ -99,9 +99,12 @@ if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU" OR (CMAKE_CXX_COMPILER_ID MATCHES "Clang
   if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU")
     # Set -Wimplicit-fallthrough=5 to only allow fallthrough annotation via __attribute__.
     set(EXTRA_COMPILE_OPTIONS ${EXTRA_COMPILE_OPTIONS} -Wimplicit-fallthrough=5
-                              -Wno-maybe-uninitialized
+                              -Wno-maybe-uninitialized -Wno-c++20-extensions
     )
   endif()
+  if(CMAKE_CXX_COMPILER_ID MATCHES "Clang")
+    set(EXTRA_COMPILE_OPTIONS ${EXTRA_COMPILE_OPTIONS} -Wno-c++20-designator)
+  endif()
 endif()
 
 include(MujocoHarden)

simulate/cmake/SimulateOptions.cmake

@@ -99,9 +99,12 @@ if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU" OR (CMAKE_CXX_COMPILER_ID MATCHES "Clang
   if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU")
     # Set -Wimplicit-fallthrough=5 to only allow fallthrough annotation via __attribute__.
     set(EXTRA_COMPILE_OPTIONS ${EXTRA_COMPILE_OPTIONS} -Wimplicit-fallthrough=5
-                              -Wno-maybe-uninitialized
+                              -Wno-maybe-uninitialized -Wno-c++20-extensions
     )
   endif()
+  if(CMAKE_CXX_COMPILER_ID MATCHES "Clang")
+    set(EXTRA_COMPILE_OPTIONS ${EXTRA_COMPILE_OPTIONS} -Wno-c++20-designator)
+  endif()
 endif()
 
 include(MujocoHarden)

src/thread/CMakeLists.txt

@@ -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})

src/thread/task.h

@@ -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_

src/thread/thread_pool.cc

@@ -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 "engine/engine_crossplatform.h"
-#include "thread/task.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 =
+  return new ThreadPoolImpl(number_of_threads);
-      new mujoco::ThreadPool<kMaxThreads>(number_of_threads);
-  return static_cast<mjThreadPool*>(static_cast<void*>(thread_pool));
 }
 
 // start a task in the threadpool
-void mju_threadPoolEnqueue(
+void mju_threadPoolEnqueue(mjThreadPool* thread_pool, mjTask* task) {
-    mjThreadPool* thread_pool, mjTask* task, mjStartRoutine start_routine,
+  auto thread_pool_impl = static_cast<ThreadPoolImpl*>(thread_pool);
-    void* args) {
+  thread_pool_impl->Enqueue(task);
-  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);
 }
 
 // shutdown the threadpool and free the memory
 void mju_threadPoolDestroy(mjThreadPool* thread_pool) {
-  mujoco::ThreadPool<kMaxThreads>* thread_pool_ptr =
+  auto thread_pool_impl = static_cast<ThreadPoolImpl*>(thread_pool);
-      static_cast<mujoco::ThreadPool<kMaxThreads>*>(
+  thread_pool_impl->Shutdown();
-          static_cast<void*>(thread_pool));
+  delete thread_pool_impl;
-  thread_pool_ptr->Shutdown();
-  delete thread_pool_ptr;
 }
-
+}  // namespace mujoco

src/thread/thread_pool.h

@@ -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>
+// Enqueue a task in a thread pool.
-class ThreadPool {
+MJAPI void mju_threadPoolEnqueue(mjThreadPool* thread_pool, mjTask* task);
- 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));
-    }
-  }
 
-  // start a task in the threadpool
+// Destroy a thread pool.
-  void Enqueue(
+MJAPI void mju_threadPoolDestroy(mjThreadPool* thread_pool);
-      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_;
-};
 
+#ifdef __cplusplus
+}  // extern "C"
 }  // namespace mujoco
-
 #endif  // __cplusplus
 
 #endif  // MUJOCO_SRC_THREAD_THREAD_POOL_H_

src/thread/thread_queue.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_

src/thread/thread_task.cc

@@ -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
+namespace mujoco {
-void mju_taskJoin(mjTask* task) {
+void mju_defaultTask(mjTask* task) {
-  mujoco::Task* task_ptr = static_cast<mujoco::Task*>(static_cast<void*>(task));
+  task->func = nullptr;
-  task_ptr->Join();
+  task->args = nullptr;
+  task->status = mjTASK_NEW;
 }
+
+void mju_taskJoin(mjTask* task) {
+  while (GetAtomicTaskStatus(task) != mjTASK_COMPLETED) {
+    std::this_thread::yield();
+  }
+}
+}  // namespace mujoco

src/thread/thread_task.h

@@ -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_

test/thread/CMakeLists.txt

@@ -0,0 +1,19 @@
+# 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
+#
+#     https://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.
+
+mujoco_test(thread_pool_test)
+target_link_libraries(thread_pool_test fixture gmock)
+
+mujoco_test(thread_queue_test)
+target_link_libraries(thread_queue_test fixture gmock)

test/thread/mjthread_test.cc

@@ -1,90 +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.
-
-#include <mujoco/mjthread.h>
-
-#include <atomic>
-
-#include <gtest/gtest.h>
-#include <mujoco/mujoco.h>
-#include "src/thread/task.h"
-#include "src/thread/thread_pool.h"
-
-namespace {
-
-struct TestFunctionArgs_ {
-  int input;
-  // make this atomic to avoid red-herring tsan failures.
-  std::atomic<int> output;
-};
-typedef struct TestFunctionArgs_ TestFunctionArgs;
-
-void* test_function(void* args) {
-  TestFunctionArgs* test_function_args = static_cast<TestFunctionArgs*>(args);
-  if (!test_function_args) {
-    return nullptr;
-  }
-  test_function_args->output = test_function_args->input;
-  return nullptr;
-}
-
-TEST(TestMjThreadPool, EnsureStructClassSizeMatch) {
-  EXPECT_EQ(sizeof(mjTask), sizeof(mujoco::Task));
-  EXPECT_EQ(sizeof(mjThreadPool), sizeof(mujoco::ThreadPool<128>));
-}
-
-TEST(TestMjThreadPool, TestMjThreadPool10Threads) {
-  mjThreadPool* thread_pool = mju_threadPoolCreate(10);
-
-  TestFunctionArgs test_function_args[1000];
-  mjTask tasks[1000];
-  for (int i = 0; i < 1000; ++i) {
-    test_function_args[i].input = i;
-    mju_threadPoolEnqueue(thread_pool, &tasks[i], test_function,
-                          (void*)&test_function_args[i]);
-  }
-
-  for (int i = 0; i < 1000; ++i) {
-    mju_taskJoin(&tasks[i]);
-  }
-
-  for (int i = 0; i < 1000; ++i) {
-    EXPECT_EQ(test_function_args[i].input, test_function_args[i].output);
-  }
-  mju_threadPoolDestroy(thread_pool);
-}
-
-TEST(TestMjThreadPool, TestMjThreadPool100Threads) {
-  mjThreadPool* thread_pool = mju_threadPoolCreate(100);
-
-  TestFunctionArgs test_function_args[1000];
-  mjTask tasks[1000];
-  for (int i = 0; i < 1000; ++i) {
-    test_function_args[i].input = i;
-    mju_threadPoolEnqueue(thread_pool, &tasks[i], test_function,
-                          (void*)&test_function_args[i]);
-  }
-
-  for (int i = 0; i < 1000; ++i) {
-    mju_taskJoin(&tasks[i]);
-  }
-
-  for (int i = 0; i < 1000; ++i) {
-    EXPECT_EQ(test_function_args[i].input, test_function_args[i].output);
-  }
-
-  mju_threadPoolDestroy(thread_pool);
-}
-
-}  // namespace

test/thread/task_test.cc

@@ -1,46 +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.
-
-#include "src/thread/task.h"
-
-#include <gtest/gtest.h>
-
-namespace mujoco {
-namespace {
-
-struct TestFunctionArgs {
-  int input;
-  int output;
-};
-
-void* test_function(void* args) {
-  TestFunctionArgs* test_function_args = (TestFunctionArgs*)args;
-  test_function_args->output = test_function_args->input;
-  return nullptr;
-}
-
-TEST(TestMjThread, TestMjThread) {
-  TestFunctionArgs test_function_args;
-  test_function_args.input = 1;
-  test_function_args.output = 2;
-  Task task;
-  Task::Initialize(
-      &task, test_function, static_cast<void*>(&test_function_args));
-  task.Execute();
-  task.Join();
-  EXPECT_EQ(test_function_args.input, test_function_args.output);
-}
-
-}  // namespace
-}  // namespace mujoco

test/thread/thread_pool_test.cc

@@ -12,8 +12,6 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "src/thread/thread_pool.h"
-
 #include <atomic>
 #include <condition_variable>
 #include <memory>
@@ -21,75 +19,82 @@
 #include <thread>
 
 #include <gtest/gtest.h>
-#include "src/thread/task.h"
+#include <mujoco/mujoco.h>
 
 namespace mujoco {
 namespace {
 
-struct TestFunctionArgs {
+struct TestFunctionArgs_ {
   int input;
   // make this atomic to avoid red-herring tsan failures.
   std::atomic<int> output;
 };
+typedef struct TestFunctionArgs_ TestFunctionArgs;
 
 void* test_function(void* args) {
   TestFunctionArgs* test_function_args = static_cast<TestFunctionArgs*>(args);
+  if (!test_function_args) {
+    return nullptr;
+  }
   test_function_args->output = test_function_args->input;
   return nullptr;
 }
 
 TEST(TestMjThreadPool, TestMjThreadPool10Threads) {
-  ThreadPool<10> thread_pool(10);
+  mjThreadPool* thread_pool = mju_threadPoolCreate(10);
 
   constexpr int kTasks = 1000;
   TestFunctionArgs test_function_args[kTasks];
-  Task tasks[kTasks];
+  mjTask tasks[kTasks];
   for (int i = 0; i < kTasks; ++i) {
     test_function_args[i].input = i;
-    thread_pool.Enqueue(
+    mju_defaultTask(&tasks[i]);
-        &tasks[i], test_function, static_cast<void*>(&test_function_args[i]));
+    tasks[i].func = test_function;
+    tasks[i].args = &test_function_args[i];
+    mju_threadPoolEnqueue(thread_pool, &tasks[i]);
   }
 
   for (int i = 0; i < kTasks; ++i) {
-    tasks[i].Join();
+    mju_taskJoin(&tasks[i]);
   }
 
   for (int i = 0; i < kTasks; ++i) {
     EXPECT_EQ(test_function_args[i].input, test_function_args[i].output);
   }
-
+  mju_threadPoolDestroy(thread_pool);
-  thread_pool.Shutdown();
 }
 
 TEST(TestMjThreadPool, TestMjThreadPool100Threads) {
-  ThreadPool<100> thread_pool(100);
+  mjThreadPool* thread_pool = mju_threadPoolCreate(100);
 
   constexpr int kTasks = 1000;
   TestFunctionArgs test_function_args[kTasks];
-  Task tasks[kTasks];
+  mjTask tasks[kTasks];
   for (int i = 0; i < kTasks; ++i) {
     test_function_args[i].input = i;
-    thread_pool.Enqueue(
+    mju_defaultTask(&tasks[i]);
-        &tasks[i], test_function, static_cast<void*>(&test_function_args[i]));
+    tasks[i].func = test_function;
+    tasks[i].args = &test_function_args[i];
+    mju_threadPoolEnqueue(thread_pool, &tasks[i]);
   }
 
   for (int i = 0; i < kTasks; ++i) {
-    tasks[i].Join();
+    mju_taskJoin(&tasks[i]);
   }
 
   for (int i = 0; i < kTasks; ++i) {
     EXPECT_EQ(test_function_args[i].input, test_function_args[i].output);
   }
 
-  thread_pool.Shutdown();
+  mju_threadPoolDestroy(thread_pool);
 }
 
 TEST(TestMjThreadPool, TestMjThreadPoolManyWriters) {
-  ThreadPool<10> thread_pool(10);
+  mjThreadPool* thread_pool = mju_threadPoolCreate(10);
 
   constexpr int kTasks = 20;
   TestFunctionArgs test_function_args[kTasks];
-  Task tasks[kTasks];
+  mjTask tasks[kTasks];
   std::unique_ptr<std::thread> enqueue_threads[kTasks];
 
   // add tasks to the thread pool from many threads
@@ -98,6 +103,10 @@ TEST(TestMjThreadPool, TestMjThreadPoolManyWriters) {
   bool start = false;
   for (int i = 0; i < kTasks; ++i) {
     test_function_args[i].input = i;
+    mju_defaultTask(&tasks[i]);
+    tasks[i].func = &test_function;
+    tasks[i].args = &test_function_args[i];
+
     enqueue_threads[i] = std::make_unique<std::thread>([&, i] {
       // synchronize all threads adding to the thread_pool at the same time
       {
@@ -105,8 +114,7 @@ TEST(TestMjThreadPool, TestMjThreadPoolManyWriters) {
         start_cv.wait(lock, [&] { return start; });
       }
       // enqueue outside the lock, to get some concurrency
-      thread_pool.Enqueue(
+      mju_threadPoolEnqueue(thread_pool, &tasks[i]);
-          &tasks[i], test_function, static_cast<void*>(&test_function_args[i]));
     });
   }
   {
@@ -120,14 +128,14 @@ TEST(TestMjThreadPool, TestMjThreadPoolManyWriters) {
   }
 
   for (int i = 0; i < kTasks; ++i) {
-    tasks[i].Join();
+    mju_taskJoin(&tasks[i]);
   }
 
   for (int i = 0; i < kTasks; ++i) {
     EXPECT_EQ(test_function_args[i].input, test_function_args[i].output);
   }
 
-  thread_pool.Shutdown();
+  mju_threadPoolDestroy(thread_pool);
 }
 
 }  // namespace

test/thread/thread_queue_test.cc

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "src/thread/lockless_queue.h"
+#include "src/thread/thread_queue.h"
 
 #include <cstddef>
 

unity/Runtime/Bindings/MjBindings.cs

@@ -403,6 +403,11 @@ public enum mjtFont : int{
   mjFONT_SHADOW = 1,
   mjFONT_BIG = 2,
 }
+public enum mjtTaskStatus : int{
+  mjTASK_NEW = 0,
+  mjTASK_QUEUED = 1,
+  mjTASK_COMPLETED = 2,
+}
 public enum mjtButton : int{
   mjBUTTON_NONE = 0,
   mjBUTTON_LEFT = 1,
@@ -2360,16 +2365,6 @@ public unsafe struct mjrContext_ {
   public int readPixelFormat;
 }
 
-[StructLayout(LayoutKind.Sequential)]
-public unsafe struct mjTask_ {
-  public fixed sbyte buffer[24];
-}
-
-[StructLayout(LayoutKind.Sequential)]
-public unsafe struct mjThreadPool_ {
-  public fixed sbyte buffer[6208];
-}
-
 [StructLayout(LayoutKind.Sequential)]
 public unsafe struct mjuiState_ {
   public int nrect;
@@ -3950,14 +3945,5 @@ public static unsafe extern void mjd_subQuat(double* qa, double* qb, double* Da,
 
 [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
 public static unsafe extern void mjd_quatIntegrate(double* vel, double scale, double* Dquat, double* Dvel, double* Dscale);
-
-[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
-public static unsafe extern mjThreadPool_* mju_threadPoolCreate(UIntPtr number_of_threads);
-
-[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
-public static unsafe extern void mju_taskJoin(mjTask_* task);
-
-[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
-public static unsafe extern void mju_threadPoolDestroy(mjThreadPool_* thread_pool);
 }
 }