Fix some BTree concurrency bugs

Author: komozoi

include/fs/BTree.h

@@ -24,6 +24,8 @@
 #include "stdint.h"
 #include "stddef.h"
 #include "string.h"
+#include <mutex>
+#include <shared_mutex>
 
 
 /**
@@ -85,6 +87,7 @@ class BTreeBase {
 	 * @return true if the key is already present, false if the key was not previously present and has been inserted
 	 */
 	bool insert(const T& val) {
+		//std::unique_lock<std::shared_mutex> lock(treeMutex);
 		btree_node_t<T, N> target;
 		btree_node_header_t<N>& header = target.header;
 		uint64_t offset;
@@ -126,6 +129,7 @@ class BTreeBase {
 	 * @return true if the key was already present and overwritten, false if the key was not previously present and has been newly inserted
 	 */
 	bool overwrite(const T& val) {
+		//std::unique_lock<std::shared_mutex> lock(treeMutex);
 		btree_node_t<T, N> target;
 		uint64_t offset;
 
@@ -171,6 +175,7 @@ class BTreeBase {
 	 * @return true if the value was found, false if not
 	 */
 	bool find(T& val) {
+		//std::shared_lock<std::shared_mutex> lock(treeMutex);
 		btree_node_t<T, N> target;
 		uint64_t offset;
 
@@ -190,6 +195,7 @@ class BTreeBase {
 	 * @return true if a match was found, false if not
 	 */
 	bool findNext(T& val) {
+		//std::shared_lock<std::shared_mutex> lock(treeMutex);
 		btree_node_t<T, N> node = getRootNode();
 		uint64_t offset = getRootOffset();
 		bool found = false;
@@ -243,6 +249,7 @@ class BTreeBase {
 	 * @return true if the value was found and removed, false if not
 	 */
 	bool remove(T& val) {
+		//std::unique_lock<std::shared_mutex> lock(treeMutex);
 		btree_node_t<T, N> node;
 		uint64_t offset;
 
@@ -277,6 +284,14 @@ class BTreeBase {
 	 */
 	explicit BTreeBase(int (*compare) (const T&, const T&)) : compare(compare) {}
 
+	// The mutex is non-copyable/non-movable but instances of BTreeBase are
+	// embedded in other containers (e.g. TreeMap) that rely on copy/move.
+	// Each instance gets its own fresh mutex.
+	BTreeBase(const BTreeBase& other) : compare(other.compare) {}
+	BTreeBase(BTreeBase&& other) noexcept : compare(other.compare) {}
+	BTreeBase& operator=(const BTreeBase& other) { compare = other.compare; return *this; }
+	BTreeBase& operator=(BTreeBase&& other) noexcept { compare = other.compare; return *this; }
+
 	/**
 	 * Finds which node is a good candidate for adding the given value.  The outputted node will always be at the
 	 * bottom of the tree, unless the value is already present.  If the value is present in the tree, then this will
@@ -524,6 +539,8 @@ class BTreeBase {
 	virtual void overwriteNodeHeader(uint64_t offset, const btree_node_t<T, N>& node) = 0;
 
 	int (*compare) (const T&, const T&); /**< Comparison function for elements of type T. */
+
+	mutable std::shared_mutex mutex;
 };
 
 
@@ -575,9 +592,7 @@ class BTree: public BTreeBase<T, N> {
 	 */
 	void initialize() {
 		btree_node_t<T, N> root{{0, 0, 0, -1, 0, {}}, {}};
-
-		file.seek(rootOffset);
-		file.write(root);
+		file.pwrite(root, rootOffset);
 	}
 
 protected:
@@ -589,9 +604,7 @@ class BTree: public BTreeBase<T, N> {
 	 */
 	btree_node_t<T, N> getNode(uint64_t offset) const override {
 		btree_node_t<T, N> resultat;
-
-		file.seek(offset, SEEK_SET);
-		file.read(resultat);
+		file.pread(resultat, (off_t)offset);
 		return resultat;
 	}
 
@@ -618,7 +631,7 @@ class BTree: public BTreeBase<T, N> {
 	 */
 	uint64_t addNode(const btree_node_t<T, N>& node) override {
 		uint64_t offset = file.seekToEndWithPadding(8);
-		file.write(node);
+		file.pwrite(node, offset);
 		return offset;
 	}
 
@@ -628,8 +641,7 @@ class BTree: public BTreeBase<T, N> {
 	 * @param node Node data.
 	 */
 	void overwriteNode(uint64_t offset, const btree_node_t<T, N>& node) override {
-		file.seek(offset);
-		file.write(node);
+		file.pwrite(node, (off_t)offset);
 	}
 
 	/**
@@ -638,8 +650,7 @@ class BTree: public BTreeBase<T, N> {
 	 * @param node Node data containing the header.
 	 */
 	void overwriteNodeHeader(uint64_t offset, const btree_node_t<T, N>& node) override {
-		file.seek(offset);
-		file.write(&node, sizeof(node.header));
+		file.pwrite(&node, sizeof(node.header), (off_t)offset);
 	}
 
 	off_t rootOffset; /**< Offset of the root node in the file. */

src/tests/test_FreeSpaceFile.cpp

@@ -18,6 +18,9 @@
 
 #include <string>
 #include <cstdint>
+#include <thread>
+#include <vector>
+#include <atomic>
 #include "universaltime.h"
 #include "gtest/gtest.h"
 #include "fs/FdHandle.h"
@@ -279,6 +282,48 @@ TEST(FreeSpaceFileTest, ReinsertAfterAllocate) {
 	if (!HasFailure()) unlink(test_file);
 }
 
+// Regression test: BTree is documented as threadsafe, and FreeSpaceFile is
+// accessed concurrently (e.g. from a ThreadPool by CatalogFile).  Allocating
+// and freeing regions from multiple threads must not corrupt the BTree's
+// on-disk nodes or crash inside the comparator.
+TEST(FreeSpaceFileTest, ConcurrentAllocateAndFree) {
+	std::string path = makeTempFileName("freespace_concurrent");
+	const char* test_file = path.c_str();
+
+	unlink(test_file);
+	FdHandle handle = FdHandle::open(test_file, O_RDWR | O_CREAT, 0660);
+	FreeSpaceFile fs(std::move(handle));
+
+	const int numThreads = 8;
+	const int opsPerThread = 200;
+
+	std::atomic<bool> failed(false);
+	std::vector<std::thread> threads;
+	threads.reserve(numThreads);
+
+	for (int t = 0; t < numThreads; t++) {
+		threads.emplace_back([&fs, &failed, t, opsPerThread]() {
+			for (int i = 0; i < opsPerThread; i++) {
+				uint32_t size = 64 + ((uint32_t)((t * 131) + (i * 17)) & 0x3FF);
+				off_t off = fs.getFreeRegion(size);
+				if (off < 0) {
+					failed.store(true);
+					return;
+				}
+				fs.markFreeRegion(off, size);
+			}
+		});
+	}
+
+	for (std::thread& th : threads)
+		th.join();
+
+	EXPECT_FALSE(failed.load());
+
+	if (!HasFailure()) unlink(test_file);
+}
+
+
 TEST(FreeSpaceFileTest, NeverAllocatesBeforeHeader) {
 	std::string path = makeTempFileName("freespace_header");
 	const char* test_file = path.c_str();