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

src/thread/thread_queue.h

@@ -0,0 +1,152 @@
+// 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_LOCKLESS_QUEUE_H_
+#define MUJOCO_SRC_THREAD_LOCKLESS_QUEUE_H_
+
+#include <atomic>
+#include <climits>
+#include <cstddef>
+#include <thread>
+
+namespace mujoco {
+
+// A Lockless Queue allows for sending information quickly between different
+// threads. This is a Multi-Producer Multi-Consumer Lockless Queue allowing for
+// multiple threads to be adding items to the queue while multiple threads are
+// consuming items from the queue. Internally it uses a Ring Buffer for storage
+// so it will not grow as items are added. Push will block if the Queue is full
+// and Pop will block if it is empty.
+//
+// For a basic overview of this category of structures:
+// https://www.linuxjournal.com/content/lock-free-multi-producer-multi-consumer-queue-ring-buffer
+template  <typename T, size_t buffer_capacity>
+class LocklessQueue {
+ public:
+  bool full() const {
+    return full_internal(
+        convert_to_index(read_cursor_), convert_to_index(write_cursor_));
+  }
+
+  bool empty() const {
+    return maximum_read_cursor_ == read_cursor_;
+  }
+
+  // Push an element into the queue.
+  void push(const T& input) {
+    // Reserve a slot in the queue
+    size_t current_write_cursor;
+    size_t dummy_current_write_cursor;
+    size_t next_write_cursor;
+    size_t current_write_index;
+    size_t current_read_index;
+    do {
+      // Check if the queue is full.
+      do {
+        current_write_cursor = write_cursor_.load();
+        current_write_index = convert_to_index(current_write_cursor);
+        next_write_cursor = get_next_cursor(current_write_cursor);
+
+        current_read_index = convert_to_index(read_cursor_.load());
+      } while (full_internal(current_read_index, current_write_index));
+
+      // Once it's not full, attempt to grab a slot to write.
+      dummy_current_write_cursor = current_write_cursor;
+    } while (!write_cursor_.compare_exchange_weak(
+                 dummy_current_write_cursor, next_write_cursor));
+
+    // Write the entry.
+    buffer_[current_write_index].store(input);
+
+    // Increment maximum read cursor. Note here it has to wait if the compare
+    // and exchange fails as another thread might not have completed its write.
+    do {
+      dummy_current_write_cursor = current_write_cursor;
+    } while (!maximum_read_cursor_.compare_exchange_weak(
+                 dummy_current_write_cursor, next_write_cursor));
+  }
+
+  // Pop an element from the queue.
+  T pop() {
+    size_t current_read_cursor;
+    size_t dummy_current_read_cursor;
+    size_t current_read_index;
+    size_t next_read_cursor;
+    size_t current_maximum_read_cursor;
+    size_t current_maximum_read_index;
+    bool empty = false;
+    T result;
+    do {
+      // Wait until the queue has an element
+      do {
+        if (empty) {
+          std::this_thread::yield();
+        }
+        current_read_cursor = read_cursor_.load();
+        current_maximum_read_cursor = maximum_read_cursor_.load();
+
+        current_read_index = convert_to_index(current_read_cursor);
+        current_maximum_read_index = convert_to_index(
+            current_maximum_read_cursor);
+
+        empty = empty_internal(
+            current_read_index, current_maximum_read_index);
+      } while (empty);
+
+      next_read_cursor = get_next_cursor(current_read_cursor);
+
+      // Attempt to grab the element, if unsuccessful then wait for the next
+      // element to arrive.
+      result = buffer_[current_read_index].load();
+      dummy_current_read_cursor = current_read_cursor;
+    } while (!read_cursor_.compare_exchange_weak(
+                 dummy_current_read_cursor, next_read_cursor));
+
+    return result;
+  }
+
+ private:
+  size_t convert_to_index(size_t input) const {
+    return input % internal_buffer_capacity_;
+  }
+
+  size_t get_next_cursor(size_t input) const {
+    return (input + 1) % cursor_max_;
+  }
+
+  size_t get_next_index(size_t input) const {
+    return convert_to_index(get_next_cursor(input));
+  }
+
+  bool full_internal(size_t read_index, size_t write_index) const {
+    return get_next_index(write_index) == read_index;
+  }
+
+  bool empty_internal(size_t read_index, size_t write_index) const {
+    return read_index == write_index;
+  }
+
+  const size_t internal_buffer_capacity_ = buffer_capacity + 1;
+  const size_t cursor_max_ = UINT_MAX - (UINT_MAX % internal_buffer_capacity_);
+
+  std::atomic<size_t> read_cursor_ = 0;
+  std::atomic<size_t> write_cursor_ = 0;
+  std::atomic<size_t> maximum_read_cursor_ = 0;
+
+  std::atomic<T> buffer_[(buffer_capacity + 1)];
+};
+
+}  // namespace mujoco
+
+#endif  // MUJOCO_SRC_THREAD_LOCKLESS_QUEUE_H_