Modify memory allocation in MuJoCo to be thread safe:

1) protect mj_arenaAllocBytes with mutexes
2) create shards for each thread in the stack and update mj_stackAllocBytes to allocate memory within each shard for a given thread

PiperOrigin-RevId: 568315726
Change-Id: I0dee6694f2a5200fa4df22ade0e68dfaebf637fc

src/engine/engine_io.c

@@ -34,6 +34,7 @@
 #include "engine/engine_util_errmem.h"
 #include "engine/engine_util_misc.h"
 #include "engine/engine_vfs.h"
+#include "thread/thread_pool.h"
 
 #ifdef ADDRESS_SANITIZER
   #include <sanitizer/asan_interface.h>
@@ -1215,15 +1216,29 @@ mjData* mj_copyData(mjData* dest, const mjModel* m, const mjData* src) {
 }
 
 
+static void maybe_lock_alloc_mutex(mjData* d) {
+  if (d->threadpool != 0) {
+    mju_threadPoolLockAllocMutex((mjThreadPool*)d->threadpool);
+  }
+}
+
+static void maybe_unlock_alloc_mutex(mjData* d) {
+  if (d->threadpool != 0) {
+    mju_threadPoolUnlockAllocMutex((mjThreadPool*)d->threadpool);
+  }
+}
+
 
 // allocate memory from the mjData arena
 void* mj_arenaAllocByte(mjData* d, size_t bytes, size_t alignment) {
+  maybe_lock_alloc_mutex(d);
   size_t misalignment = fastmod(d->parena, alignment);
   size_t padding = misalignment ? alignment - misalignment : 0;
 
   // check size
   size_t bytes_available = d->narena - d->pstack;
   if (mjUNLIKELY(d->parena + padding + bytes > bytes_available)) {
+    maybe_unlock_alloc_mutex(d);
     return NULL;
   }
 
@@ -1240,56 +1255,46 @@ void* mj_arenaAllocByte(mjData* d, size_t bytes, size_t alignment) {
   __msan_allocated_memory(result, bytes);
 #endif
 
+  maybe_unlock_alloc_mutex(d);
   return result;
 }
 
 
-
-// internal: allocate size bytes on the mjData stack
+// internal: allocate size bytes on the provided stack shard
 // declared inline so that modular arithmetic with specific alignments can be optimized out
-static inline void* stackalloc(mjData* d, size_t size, size_t alignment) {
+static inline void* stackallocinternal(mjData* d, mjStackInfo* stack_info, size_t size, size_t alignment) {
   // return NULL if empty
   if (mjUNLIKELY(!size)) {
     return NULL;
   }
 
-  // size of entire arena/stack in bytes
-  size_t stack_size_bytes = d->narena;
-
-  // end of the arena
-  uintptr_t end_of_arena_ptr = (uintptr_t)d->arena + stack_size_bytes;
-
-  // current top of the stack
-  uintptr_t end_ptr = end_of_arena_ptr - d->pstack;
-
   // start of the memory to be allocated to the buffer
-  uintptr_t start_ptr = end_ptr - (size + mjREDZONE);
+  uintptr_t start_ptr = stack_info->top - (size + mjREDZONE);
 
   // align the pointer
   start_ptr -= fastmod(start_ptr, alignment);
 
   // new top of the stack
-  uintptr_t new_pstack_ptr = start_ptr - mjREDZONE;
-  size_t new_pstack = end_of_arena_ptr - new_pstack_ptr;
+  uintptr_t new_top_ptr = start_ptr - mjREDZONE;
 
   // exclude red zone from stack usage statistics
-  size_t current_alloc_usage = end_ptr - new_pstack_ptr - 2 * mjREDZONE;
-  size_t usage = current_alloc_usage + d->pstack;
+  size_t current_alloc_usage = stack_info->top - new_top_ptr - 2 * mjREDZONE;
+  size_t usage = current_alloc_usage + (stack_info->bottom - stack_info->top);
 
   // check size
-  size_t stack_available_bytes = end_ptr - ((uintptr_t)d->arena + d->parena);
-  size_t stack_required_bytes = end_ptr - new_pstack_ptr;
+  size_t stack_available_bytes = stack_info->top - stack_info->limit;
+  size_t stack_required_bytes = stack_info->top - new_top_ptr;
   if (mjUNLIKELY(stack_required_bytes > stack_available_bytes)) {
     mju_error("mj_stackAlloc: insufficient memory: max = %zu, available = %zu, requested = %zu "
               "(ne = %d, nf = %d, nefc = %d, ncon = %d)",
-              stack_size_bytes, stack_available_bytes, stack_required_bytes,
+              stack_info->bottom - stack_info->limit, stack_available_bytes, stack_required_bytes,
               d->ne, d->nf, d->nefc, d->ncon);
   }
 
 #ifdef ADDRESS_SANITIZER
   // actual stack usage (without red zone bytes) is stored in the red zone
-  if (d->pstack) {
-    char* prev_pstack_ptr = (char*)(end_of_arena_ptr - d->pstack);
+  if (stack_info->top != stack_info->bottom) {
+    char* prev_pstack_ptr = (char*)(stack_info->top);
     size_t prev_misalign = (uintptr_t)prev_pstack_ptr % _Alignof(size_t);
     size_t* prev_usage_ptr =
         (size_t*)(prev_pstack_ptr +
@@ -1300,9 +1305,9 @@ static inline void* stackalloc(mjData* d, size_t size, size_t alignment) {
   }
 
   // store new stack usage in the red zone
-  size_t misalign = new_pstack_ptr % _Alignof(size_t);
+  size_t misalign = new_top_ptr % _Alignof(size_t);
   size_t* usage_ptr =
-      (size_t*)(new_pstack_ptr + (misalign ? _Alignof(size_t) - misalign : 0));
+      (size_t*)(new_top_ptr + (misalign ? _Alignof(size_t) - misalign : 0));
   ASAN_UNPOISON_MEMORY_REGION(usage_ptr, sizeof(size_t));
   *usage_ptr = usage;
   ASAN_POISON_MEMORY_REGION(usage_ptr, sizeof(size_t));
@@ -1311,45 +1316,95 @@ static inline void* stackalloc(mjData* d, size_t size, size_t alignment) {
   ASAN_UNPOISON_MEMORY_REGION((void*)start_ptr, size);
 #endif
 
-  // update pstack and max usage statistics
-  d->pstack = new_pstack;
-  d->maxuse_stack = mjMAX(d->maxuse_stack, usage);
-  d->maxuse_arena = mjMAX(d->maxuse_arena, usage + d->parena);
+  // update max usage statistics
+  stack_info->top = new_top_ptr;
+  if (!d->threadpool) {
+    d->maxuse_stack = mjMAX(d->maxuse_stack, usage);
+    d->maxuse_arena = mjMAX(d->maxuse_arena, usage + d->parena);
+  } else {
+    size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
+    d->maxuse_threadstack[thread_id] = mjMAX(d->maxuse_threadstack[thread_id], usage);
+  }
 
   return (void*)start_ptr;
 }
 
 
+static inline mjStackInfo get_stack_info_from_data(mjData* d) {
+  mjStackInfo stack_info;
+  stack_info.bottom = (uintptr_t)d->arena + (uintptr_t)d->narena;
+  stack_info.top = stack_info.bottom - d->pstack;
+  stack_info.limit = (uintptr_t)d->arena + (uintptr_t)d->parena;
+  stack_info.stack_base = d->pbase;
+
+  return stack_info;
+}
+
+
+// internal: allocate size bytes in mjData
+// declared inline so that modular arithmetic with specific alignments can be optimized out
+static inline void* stackalloc(mjData* d, size_t size, size_t alignment) {
+  if (!d->threadpool) {
+    mjStackInfo stack_info = get_stack_info_from_data(d);
+
+    void* result = stackallocinternal(d, &stack_info, size, alignment);
+
+    d->pstack = stack_info.bottom - stack_info.top;
+
+    return result;
+  }
+
+  size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
+  mjStackInfo* stack_info = mju_getStackInfoForThread(d, thread_id);
+  return stackallocinternal(d, stack_info, size, alignment);
+}
+
+
+// mjStackInfo mark stack frame, inline so ASAN errors point to correct code unit
+#ifdef ADDRESS_SANITIZER
+__attribute__((always_inline))
+#endif
+static inline void markstackinternal(mjData* d, mjStackInfo* stack_info) {
+  size_t top_old = stack_info->top;
+  mjStackFrame* s =
+      (mjStackFrame*) stackallocinternal(d, stack_info, sizeof(mjStackFrame), _Alignof(mjStackFrame));
+  s->pbase = stack_info->stack_base;
+  s->pstack = top_old;
+#ifdef ADDRESS_SANITIZER
+  // store the program counter to the caller so that we can compare against mj_freeStack later
+  s->pc = __sanitizer_return_address();
+#endif
+  stack_info->stack_base = (uintptr_t) s;
+}
+
 
 // mjData mark stack frame
 #ifdef ADDRESS_SANITIZER
 __attribute__((noinline))
 #endif
 void mj_markStack(mjData* d) {
-  size_t pstack_old = d->pstack;
-  mjStackFrame* s =
-      (mjStackFrame*) stackalloc(d, sizeof(mjStackFrame), _Alignof(mjStackFrame));
-  s->pbase = d->pbase;
-  s->pstack = pstack_old;
-#ifdef ADDRESS_SANITIZER
-  // store the program counter to the caller so that we can compare against mj_freeStack later
-  s->pc = __sanitizer_return_address();
-#endif
-  d->pbase = d->pstack - mjREDZONE;
-}
-
-
-
-// mjData free stack frame
-#ifdef ADDRESS_SANITIZER
-__attribute__((noinline))
-#endif
-void mj_freeStack(mjData* d) {
-  if (mjUNLIKELY(!d->pbase)) {
+  if (!d->threadpool) {
+    mjStackInfo stack_info = get_stack_info_from_data(d);
+    markstackinternal(d, &stack_info);
+    d->pstack = stack_info.bottom - stack_info.top;
+    d->pbase = stack_info.stack_base;
     return;
   }
 
-  mjStackFrame* s = (mjStackFrame*) ((char*)d->arena + d->narena - d->pbase);
+  size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
+  mjStackInfo* stack_info = mju_getStackInfoForThread(d, thread_id);
+  markstackinternal(d, stack_info);
+}
+
+#ifdef ADDRESS_SANITIZER
+__attribute__((always_inline))
+#endif
+static inline void freestackinternal(mjStackInfo* stack_info) {
+  if (mjUNLIKELY(!stack_info->stack_base)) {
+    return;
+  }
+
+  mjStackFrame* s = (mjStackFrame*) stack_info->stack_base;
 #ifdef ADDRESS_SANITIZER
   // raise an error if caller function name doesn't match the most recent caller of mj_markStack
   if (!_mj_comparePcFuncName(s->pc, __sanitizer_return_address())) {
@@ -1365,15 +1420,34 @@ void mj_freeStack(mjData* d) {
 #endif
 
   // restore pbase and pstack
-  d->pbase = s->pbase;
-  d->pstack = s->pstack;
+  stack_info->stack_base = s->pbase;
+  stack_info->top = s->pstack;
 
   // if running under asan, poison the newly freed memory region
 #ifdef ADDRESS_SANITIZER
-  ASAN_POISON_MEMORY_REGION((char*)d->arena + d->parena, d->narena - d->pstack - d->parena);
+  ASAN_POISON_MEMORY_REGION((char*)stack_info->limit, stack_info->top - stack_info->limit);
 #endif
 }
 
+
+// mjData free stack frame
+#ifdef ADDRESS_SANITIZER
+__attribute__((noinline))
+#endif
+void mj_freeStack(mjData* d) {
+  if (!d->threadpool) {
+    mjStackInfo stack_info = get_stack_info_from_data(d);
+    freestackinternal(&stack_info);
+    d->pstack = stack_info.bottom - stack_info.top;
+    d->pbase = stack_info.stack_base;
+    return;
+  }
+
+  size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
+  mjStackInfo* stack_info = mju_getStackInfoForThread(d, thread_id);
+  freestackinternal(stack_info);
+}
+
 void* mj_stackAllocByte(mjData* d, size_t bytes, size_t alignment) {
   return stackalloc(d, bytes, alignment);
 }

src/thread/thread_pool.cc

@@ -14,15 +14,19 @@
 
 #include "thread/thread_pool.h"
 
+#include <stdint.h>
+
 #include <algorithm>
 #include <atomic>
 #include <cstddef>
 #include <memory>
+#include <mutex>
 #include <thread>
 #include <utility>
 #include <vector>
 
 #include <mujoco/mjthread.h>
+#include <mujoco/mujoco.h>
 #include "engine/engine_crossplatform.h"
 #include "engine/engine_util_errmem.h"
 #include "thread/thread_queue.h"
@@ -32,6 +36,11 @@ namespace mujoco {
 namespace {
 constexpr size_t kThreadPoolQueueSize = 640;
 
+// Each thread being run will be assigned a worker_id.
+//   0: main thread
+//   1->n: workers
+thread_local size_t worker_id = 0;
+
 struct WorkerThread {
   // Shutdown function passed to running threads to ensure clean shutdown.
   static void* ShutdownFunction(void* args) {
@@ -58,11 +67,15 @@ class ThreadPoolImpl : public mjThreadPool {
     // initialize worker threads
     for (int i = 0; i < std::min(num_worker, mjMAXTHREADS); ++i) {
       WorkerThread worker{
-        std::make_unique<std::thread>(ThreadPoolWorker, this)};
+        std::make_unique<std::thread>(ThreadPoolWorker, this, i)};
       workers_.push_back(std::move(worker));
     }
   }
 
+  size_t NumberOfThreads() {
+    return workers_.size();
+  }
+
   // start a task in the threadpool
   void Enqueue(mjTask* task) {
     if (mjUNLIKELY(GetAtomicTaskStatus(task).exchange(mjTASK_QUEUED) !=
@@ -89,11 +102,39 @@ class ThreadPoolImpl : public mjThreadPool {
     }
   }
 
+  // registers a worker ID for a given thread
+  void RegisterWorker(const size_t input_worker_id) {
+    worker_id = input_worker_id;
+  }
+
+  // gets the worker id of the current thread
+  size_t GetWorkerId() {
+    return worker_id;
+  }
+
+  void LockAlloc() {
+    alloc_mutex_.lock();
+  }
+
+  void UnlockAlloc() {
+    alloc_mutex_.unlock();
+  }
+
+  bool IsThreadPoolBound() {
+    return thread_pool_bound_;
+  }
+
+  void BindThreadPool() {
+    thread_pool_bound_ = true;
+  }
+
   ~ThreadPoolImpl() { Shutdown(); }
 
  private:
   // method executed by running threads
-  static void ThreadPoolWorker(ThreadPoolImpl* thread_pool) {
+  static void ThreadPoolWorker(
+      ThreadPoolImpl* thread_pool, const size_t thread_index) {
+    worker_id = thread_index + 1;
     while (!thread_pool->shutdown_) {
       auto task = static_cast<mjTask*>(thread_pool->lockless_queue_.pop());
       task->args = task->func(task->args);
@@ -109,11 +150,128 @@ class ThreadPoolImpl : public mjThreadPool {
 
   // queue of tasks to execute
   mujoco::LocklessQueue<void*, kThreadPoolQueueSize> lockless_queue_;
+
+  // Mutex to protect arena allocations.
+  std::mutex alloc_mutex_;
+
+  // Whether or not a ThreadPool was bound using mju_bindThreadPool.
+  bool thread_pool_bound_ = false;
 };
 
 // create a thread pool
 mjThreadPool* mju_threadPoolCreate(size_t number_of_threads) {
-  return new ThreadPoolImpl(number_of_threads);
+  return reinterpret_cast<mjThreadPool*>(new ThreadPoolImpl(number_of_threads));
+}
+
+// gets the number of shards the stack is currently broken into
+static size_t GetNumberOfShards(mjData* d) {
+  if (!d->threadpool) {
+    return 1;
+  }
+  return mju_threadPoolNumberOfThreads((mjThreadPool*)d->threadpool) + 1;
+}
+
+// returns the stack information for the specified thread's shard
+mjStackInfo* mju_getStackInfoForThread(mjData* d, size_t thread_id) {
+  auto thread_pool = (ThreadPoolImpl*)d->threadpool;
+  if (!thread_pool || !thread_pool->IsThreadPoolBound()) {
+    mju_error("Thread Pool not bound, use mju_bindThreadPool to add an mjThreadPool to mjData");
+  }
+
+  // number of threads running in the threadpool plus the main thread
+  size_t number_of_shards = GetNumberOfShards(d);
+
+  // size of entire arena/stack in bytes
+  size_t total_arena_size_bytes = d->narena;
+
+  // set the shard cursor to the end of the arena
+  uintptr_t end_of_arena_ptr = (uintptr_t)d->arena + total_arena_size_bytes;
+
+  // each thread including the main one will get an equal shard of the stack
+  size_t bytes_per_shard = total_arena_size_bytes / (2 * (number_of_shards));
+
+  // ensure the shard is larger than the cache line
+  size_t misalignment = bytes_per_shard % mju_getDestructiveInterferenceSize();
+
+  if (misalignment != 0) {
+    bytes_per_shard += mju_getDestructiveInterferenceSize() - misalignment;
+  }
+
+  if (bytes_per_shard * number_of_shards > total_arena_size_bytes) {
+    mju_error("Arena is not large enough for %zu shards", number_of_shards);
+  }
+
+  uintptr_t result = (end_of_arena_ptr - (thread_id + 1) * bytes_per_shard);
+
+  // align the end of the shard to be mjStackInfo.
+  misalignment = result % alignof(mjStackInfo);
+  result -= misalignment;
+
+  return (mjStackInfo*) result;
+}
+
+// shards the stack for each thread
+static void  ConfigureMultiThreadedStack(mjData* d) {
+  if (!d->threadpool) {
+    mju_error("No thread pool specified for multithreaded operation");
+  }
+
+  size_t number_of_shards = GetNumberOfShards(d);
+
+  // current top of the stack
+  uintptr_t current_limit = (uintptr_t)d->arena + d->narena - d->pstack;
+
+  // set the shard cursor to the end of the arena
+  uintptr_t begin_shard_cursor_ptr = (uintptr_t)d->arena + d->narena;
+
+  for (size_t shard_index = 0; shard_index < number_of_shards; ++shard_index) {
+    mjStackInfo* end_shard_cursor_ptr = mju_getStackInfoForThread(d, shard_index);
+#ifdef ADDRESS_SANITIZER
+  // unpoison stack info
+  ASAN_UNPOISON_MEMORY_REGION((void*)end_shard_cursor_ptr, sizeof(mjStackInfo));
+#endif
+    // handle the main thread's stack which may already have data in it
+    if (shard_index == 0) {
+      // abort if the current stack is already larger than the portion of the stack
+      // that would be reserved for the main thread
+      if ((uintptr_t)end_shard_cursor_ptr > current_limit) {
+        mju_error("mj_bindThreadPool: sharding stack - existing stack larger than shard size: current_size = %zu, "
+                  "max_size = %zu", current_limit, (uintptr_t) end_shard_cursor_ptr);
+      }
+      end_shard_cursor_ptr->top = current_limit;
+      end_shard_cursor_ptr->stack_base = d->pbase;
+    } else {
+      // all other stacks are empty because threads have not been used yet
+      end_shard_cursor_ptr->top = begin_shard_cursor_ptr;
+      end_shard_cursor_ptr->stack_base = 0;
+    }
+
+    end_shard_cursor_ptr->bottom = begin_shard_cursor_ptr;
+    end_shard_cursor_ptr->limit = (uintptr_t)end_shard_cursor_ptr + sizeof(mjStackInfo);
+    begin_shard_cursor_ptr = (uintptr_t)end_shard_cursor_ptr - 1;
+  }
+}
+
+// adds a thread pool to mjData and configures it for multi-threaded use.
+void mju_bindThreadPool(mjData* d, mjThreadPool* thread_pool) {
+  if (d->threadpool) {
+    mju_error("Thread Pool already bound to mjData");
+  }
+
+  d->threadpool = (uintptr_t) thread_pool;
+  ((ThreadPoolImpl*)thread_pool)->BindThreadPool();
+  ConfigureMultiThreadedStack(d);
+}
+
+// gets the number of running threads in the thread pool.
+size_t mju_threadPoolNumberOfThreads(mjThreadPool* thread_pool) {
+  auto thread_pool_impl = static_cast<ThreadPoolImpl*>(thread_pool);
+  return thread_pool_impl->NumberOfThreads();
+}
+
+size_t mju_threadPoolCurrentWorkerId(mjThreadPool* thread_pool) {
+  auto thread_pool_impl = static_cast<ThreadPoolImpl*>(thread_pool);
+  return thread_pool_impl->GetWorkerId();
 }
 
 // start a task in the threadpool
@@ -128,4 +286,23 @@ void mju_threadPoolDestroy(mjThreadPool* thread_pool) {
   thread_pool_impl->Shutdown();
   delete thread_pool_impl;
 }
+
+// locks the allocation mutex to protect Stack and Arena allocations
+void mju_threadPoolLockAllocMutex(mjThreadPool* thread_pool) {
+  auto thread_pool_impl = static_cast<ThreadPoolImpl*>(thread_pool);
+  thread_pool_impl->LockAlloc();
+}
+
+// unlocks the allocation mutex to protect Stack and Arena allocations
+void mju_threadPoolUnlockAllocMutex(mjThreadPool* thread_pool) {
+  auto thread_pool_impl = static_cast<ThreadPoolImpl*>(thread_pool);
+  thread_pool_impl->UnlockAlloc();
+}
+
+// Get the destructive interference size for the architecture.
+size_t mju_getDestructiveInterferenceSize(void) {
+  // return std::hardware_destructive_interference_size;
+  return 128;
+}
+
 }  // namespace mujoco

src/thread/thread_pool.h

@@ -19,21 +19,64 @@
 
 #include <mujoco/mjexport.h>
 #include <mujoco/mjthread.h>
+#include <mujoco/mujoco.h>
 
 #ifdef __cplusplus
 namespace mujoco {
 extern "C" {
 #endif
 
+// MultiThreaded Stack will be an approximately 50/50 split of the entire buffer, with a little
+// wiggle for alignment and caching concerns. The basic layout is to reuse the existing single
+// threaded markers, and then create shards for each thread to use as its stack.
+// Not to scale.
+// |----------|-----------|-----------|-----------|-----------|----------|-----------|-----------|
+// |Used Arena|Free Arena |Shard1     |Shard1     |Shard1     |Shard0    |Shard0     |Shard0     |
+// |%%%%%%%%%%|           |StackInfo  |Free Stack |Used Stack |StackInfo |Free Stack |Used Stack |
+// |%%%%%%%%%%|           |           |           |%%%%%%%%%%%|          |           |%%%%%%%%%%%|
+// |%%%%%%%%%%|           |           |           |%%%%%%%%%%%|          |           |%%%%%%%%%%%|
+// |----------|-----------|-----------|-----------|-----------|----------|-----------|-----------|
+// d->arena   d->parena   d->pstack   shard1->stack_info      shard1->bottom_of_stack            shard1->bottom_of_stack
+//                        shard1->stack_info                  shard0->stack_info     shard0->current_stack
+//                                    shard1->top_of_stack               shard1->top_of_stack
+//                                                shard1->current_stack
+typedef struct {
+  uintptr_t bottom;          // First memory address available to the stack
+  uintptr_t top;             // Current memory address used by the stack
+  uintptr_t limit;             // Top limit of the stack (note this is smaller than bottom, stack grows down)
+  uintptr_t stack_base;               // Current stack base for mark and free stack
+} mjStackInfo;
+
 // Create a thread pool with the specified number of threads running.
 MJAPI mjThreadPool* mju_threadPoolCreate(size_t number_of_threads);
 
+// Returns the stack information for the specified thread's shard.
+mjStackInfo* mju_getStackInfoForThread(mjData* d, size_t thread_id);
+
+// Adds a thread pool to mjData and configures it for multi-threaded use.
+MJAPI void mju_bindThreadPool(mjData* d, mjThreadPool* thread_pool);
+
+// Gets the number of running threads in the thread pool.
+MJAPI size_t mju_threadPoolNumberOfThreads(mjThreadPool* thread_pool);
+
+// Gets the ID of the current thread being executed
+MJAPI size_t mju_threadPoolCurrentWorkerId(mjThreadPool* thread_pool);
+
 // Enqueue a task in a thread pool.
 MJAPI void mju_threadPoolEnqueue(mjThreadPool* thread_pool, mjTask* task);
 
+// Locks the allocation mutex to protect Arena allocations.
+MJAPI void mju_threadPoolLockAllocMutex(mjThreadPool* thread_pool);
+
+// Unlocks the allocation mutex to protect Arena allocations.
+MJAPI void mju_threadPoolUnlockAllocMutex(mjThreadPool* thread_pool);
+
 // Destroy a thread pool.
 MJAPI void mju_threadPoolDestroy(mjThreadPool* thread_pool);
 
+// Get the destructive interference size for the architecture.
+MJAPI size_t mju_getDestructiveInterferenceSize(void);
+
 #ifdef __cplusplus
 }  // extern "C"
 }  // namespace mujoco

test/engine/engine_io_test.cc

@@ -31,6 +31,7 @@
 #include <mujoco/mjxmacro.h>
 #include <mujoco/mujoco.h>
 #include "src/engine/engine_util_errmem.h"
+#include "src/thread/thread_pool.h"
 #include "test/fixture.h"
 
 namespace mujoco {
@@ -764,6 +765,82 @@ TEST_F(EngineIoTest, CanMarkAndFreeStack) {
   mj_deleteModel(model);
 }
 
+struct TestFunctionArgs_ {
+  mjData* d;
+  int input;
+  int stack_output;
+  int arena_output;
+  size_t output_thread_worker;
+};
+typedef TestFunctionArgs_ TestFunctionArgs;
+
+void* TestFunction(void* args) {
+  TestFunctionArgs* test_args = static_cast<TestFunctionArgs*>(args);
+  test_args->output_thread_worker =
+      mju_threadPoolCurrentWorkerId((mjThreadPool*)test_args->d->threadpool);
+  mj_markStack(test_args->d);
+  int* test_ints = mj_stackAllocInt(test_args->d, 10);
+  test_ints[0] = test_args->input;
+  test_args->stack_output = test_ints[0];
+
+  int* test_arena_ints =
+      (int*)mj_arenaAllocByte(test_args->d, sizeof(int) * 10, _Alignof(int));
+  test_arena_ints[0] = test_args->input;
+  test_args->arena_output = test_arena_ints[0];
+
+
+  mj_freeStack(test_args->d);
+  return nullptr;
+}
+
+TEST_F(EngineIoTest, TestStackShardingForThreads) {
+  constexpr char xml[] = R"(
+  <mujoco>
+    <worldbody>
+    </worldbody>
+  </mujoco>
+  )";
+
+  std::array<char, 1024> error;
+  mjModel* model = LoadModelFromString(xml, error.data(), error.size());
+  ASSERT_THAT(model, NotNull()) << "Failed to load model: " << error.data();
+
+  mjData* data = mj_makeData(model);
+  ASSERT_THAT(data, NotNull());
+  mjThreadPool* thread_pool = mju_threadPoolCreate(10);
+  mju_bindThreadPool(data, thread_pool);
+
+  constexpr int kTasks = 1000;
+  TestFunctionArgs test_function_args[kTasks];
+  mjTask tasks[kTasks];
+  for (int i = 0; i < kTasks; ++i) {
+    test_function_args[i].d = data;
+    test_function_args[i].input = i;
+    mju_defaultTask(&tasks[i]);
+    tasks[i].func = TestFunction;
+    tasks[i].args = &test_function_args[i];
+    mju_threadPoolEnqueue(thread_pool, &tasks[i]);
+  }
+
+  mj_markStack(data);
+  int* test_ints = mj_stackAllocInt(data, 10);
+  test_ints[0] = 1;
+  mj_freeStack(data);
+
+  for (int i = 0; i < kTasks; ++i) {
+    mju_taskJoin(&tasks[i]);
+  }
+
+  for (int i = 0; i < kTasks; ++i) {
+    EXPECT_EQ(test_function_args[i].input, test_function_args[i].stack_output);
+    EXPECT_EQ(test_function_args[i].input, test_function_args[i].arena_output);
+  }
+
+  mj_deleteData(data);
+  mj_deleteModel(model);
+  mju_threadPoolDestroy(thread_pool);
+}
+
 #ifdef ADDRESS_SANITIZER
 void MarkFreeStack(mjData* d, bool free) {
   mj_markStack(d);