ferntree 0.6.0

//! # Concurrency Tests for Ferntree B+ Tree
//!
//! This module contains multi-threaded tests to verify the correctness
//! of the concurrent B+ tree implementation under various contention scenarios.
//!
//! ## Test Categories
//!
//! - Basic concurrent tests: Lower contention scenarios
//! - Stress tests: Higher contention scenarios to validate behavior under load

use ferntree::{OptimisticRead, Tree};
use rand::prelude::*;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::thread;
use std::time::Duration;

// ===========================================================================
// Basic Concurrent Insert Tests
// ===========================================================================

#[test]
fn concurrent_insert_disjoint_ranges() {
	let tree = Arc::new(Tree::<i32, i32>::new());
	let num_threads = 4;
	let entries_per_thread = 100;

	let handles: Vec<_> = (0..num_threads)
		.map(|t| {
			let tree = Arc::clone(&tree);
			thread::spawn(move || {
				for i in 0..entries_per_thread {
					let key = t * entries_per_thread + i;
					tree.insert(key, key * 10);
				}
			})
		})
		.collect();

	for h in handles {
		h.join().unwrap();
	}

	// Single-threaded invariant check after all concurrent operations complete
	tree.assert_invariants();
	assert_eq!(tree.len(), (num_threads * entries_per_thread) as usize);

	// Verify all entries
	for t in 0..num_threads {
		for i in 0..entries_per_thread {
			let key = t * entries_per_thread + i;
			assert_eq!(tree.lookup(&key, |v| *v), Some(key * 10), "Missing key {}", key);
		}
	}
}

#[test]
fn concurrent_insert_same_keys() {
	let tree = Arc::new(Tree::<i32, i32>::new());
	let num_threads = 4;
	let iterations = 100;

	// All threads repeatedly insert the same small set of keys
	let handles: Vec<_> = (0..num_threads)
		.map(|t| {
			let tree = Arc::clone(&tree);
			thread::spawn(move || {
				for i in 0..iterations {
					let key = i % 10; // Only 10 unique keys
					tree.insert(key, t); // Value is thread ID
				}
			})
		})
		.collect();

	for h in handles {
		h.join().unwrap();
	}

	// Single-threaded invariant check after all concurrent operations complete
	tree.assert_invariants();

	// Should have exactly 10 entries
	assert_eq!(tree.len(), 10);

	// Each key should have a valid thread ID as value
	for key in 0..10 {
		let value = tree.lookup(&key, |v| *v).expect("Key should exist");
		assert!(value < num_threads, "Invalid value {} for key {}", value, key);
	}
}

// ===========================================================================
// Basic Concurrent Lookup Tests
// ===========================================================================

#[test]
fn many_concurrent_readers() {
	let tree = Arc::new(Tree::<i32, i32>::new());
	let num_readers = 4;
	let entries = 100;

	// Pre-insert data
	for i in 0..entries {
		tree.insert(i, i * 10);
	}

	tree.assert_invariants();

	let handles: Vec<_> = (0..num_readers)
		.map(|_| {
			let tree = Arc::clone(&tree);
			thread::spawn(move || {
				// Each reader does a full scan
				let mut iter = tree.raw_iter();
				iter.seek_to_first();

				let mut count = 0;
				while let Some((k, v)) = iter.next() {
					assert_eq!(*v, *k * 10);
					count += 1;
				}
				count
			})
		})
		.collect();

	for h in handles {
		let count = h.join().unwrap();
		assert_eq!(count, entries);
	}

	// Final invariant check
	tree.assert_invariants();
}

// ===========================================================================
// Basic Mixed Operation Tests
// ===========================================================================

#[test]
fn concurrent_insert_and_lookup() {
	let tree = Arc::new(Tree::<i32, i32>::new());
	let entries = 100;

	// Pre-insert some data
	for i in 0..entries {
		tree.insert(i, i * 10);
	}

	tree.assert_invariants();

	let tree_writer = Arc::clone(&tree);
	let tree_reader = Arc::clone(&tree);

	let writer = thread::spawn(move || {
		for i in entries..(entries + 50) {
			tree_writer.insert(i, i * 10);
		}
	});

	let reader = thread::spawn(move || {
		let mut found = 0;
		for i in 0..entries {
			if tree_reader.lookup(&i, |v| *v).is_some() {
				found += 1;
			}
		}
		found
	});

	writer.join().unwrap();
	let found = reader.join().unwrap();

	// Final invariant check after all concurrent operations
	tree.assert_invariants();

	// Reader should have found entries
	assert!(found > 0);
	assert!(tree.len() >= entries as usize);
}

// ===========================================================================
// Concurrent Iteration Tests
// ===========================================================================

#[test]
fn iterate_while_inserting() {
	let tree = Arc::new(Tree::<i32, i32>::new());

	// Pre-insert some data
	for i in 0..50 {
		tree.insert(i, i);
	}

	tree.assert_invariants();

	let tree_writer = Arc::clone(&tree);
	let tree_reader = Arc::clone(&tree);

	let writer = thread::spawn(move || {
		for i in 50..75 {
			tree_writer.insert(i, i);
		}
	});

	let reader = thread::spawn(move || {
		let mut iter = tree_reader.raw_iter();
		iter.seek_to_first();

		let mut prev = -1i32;
		let mut count = 0;
		while let Some((k, _)) = iter.next() {
			// Keys should be in sorted order
			assert!(*k > prev, "Order violation: {} not > {}", *k, prev);
			prev = *k;
			count += 1;
		}
		count
	});

	writer.join().unwrap();
	let count = reader.join().unwrap();

	// Final invariant check
	tree.assert_invariants();

	// Reader should have seen a consistent snapshot
	assert!(count > 0);
}

// ===========================================================================
// Split Under Contention Tests
// ===========================================================================

#[test]
fn concurrent_inserts_cause_splits() {
	let tree = Arc::new(Tree::<i32, i32>::new());
	let num_threads = 2;
	let entries_per_thread = 50;

	// Insert enough entries concurrently to cause splits
	let handles: Vec<_> = (0..num_threads)
		.map(|t| {
			let tree = Arc::clone(&tree);
			thread::spawn(move || {
				for i in 0..entries_per_thread {
					tree.insert(t * entries_per_thread + i, i);
				}
			})
		})
		.collect();

	for h in handles {
		h.join().unwrap();
	}

	// Final invariant check after concurrent splits
	tree.assert_invariants();

	// Tree should have grown
	assert!(tree.height() > 1, "Tree should have split");
	assert_eq!(tree.len(), (num_threads * entries_per_thread) as usize);

	// Verify all entries
	for t in 0..num_threads {
		for i in 0..entries_per_thread {
			let key = t * entries_per_thread + i;
			assert_eq!(tree.lookup(&key, |v| *v), Some(i), "Missing key {}", key);
		}
	}
}

// ===========================================================================
// Basic Remove Tests
// ===========================================================================

#[test]
fn concurrent_removes() {
	let tree = Arc::new(Tree::<i32, i32>::new());
	let entries = 100;

	// Pre-insert entries
	for i in 0..entries {
		tree.insert(i, i);
	}

	tree.assert_invariants();

	let num_threads = 2;
	let entries_per_thread = entries / num_threads;

	// Concurrently remove entries
	let handles: Vec<_> = (0..num_threads)
		.map(|t| {
			let tree = Arc::clone(&tree);
			thread::spawn(move || {
				for i in 0..entries_per_thread {
					let key = t * entries_per_thread + i;
					tree.remove(&key);
				}
			})
		})
		.collect();

	for h in handles {
		h.join().unwrap();
	}

	// Final invariant check after concurrent removes
	tree.assert_invariants();

	// Tree should be empty
	assert!(tree.is_empty());
}

// ===========================================================================
// Stress Tests
// ===========================================================================

/// Higher contention stress test - validates behavior under heavy contention
#[test]
fn stress_concurrent_insert_overlapping_ranges() {
	let tree = Arc::new(Tree::<i32, i32>::new());
	let num_threads = 8;
	let entries_per_thread = 1000;
	let key_range = 500; // High contention - all threads insert into 0..500

	let handles: Vec<_> = (0..num_threads)
		.map(|t| {
			let tree = Arc::clone(&tree);
			thread::spawn(move || {
				let mut rng = rand::rng();
				for _ in 0..entries_per_thread {
					let key: i32 = rng.random_range(0..key_range);
					tree.insert(key, t);
				}
			})
		})
		.collect();

	for h in handles {
		h.join().unwrap();
	}

	// Final invariant check
	tree.assert_invariants();

	// Tree should have at most key_range entries
	assert!(tree.len() <= key_range as usize);
}

/// Mixed operations stress test
#[test]
fn stress_concurrent_mixed_operations() {
	let tree = Arc::new(Tree::<i32, i32>::new());
	let num_threads = 8;
	let operations_per_thread = 1000;
	let key_range = 500;

	// Pre-insert some data
	for i in 0..key_range {
		tree.insert(i, i);
	}

	tree.assert_invariants();

	let handles: Vec<_> = (0..num_threads)
		.map(|t| {
			let tree = Arc::clone(&tree);
			thread::spawn(move || {
				let mut rng = rand::rng();
				for _ in 0..operations_per_thread {
					let key: i32 = rng.random_range(0..key_range);
					let op: u8 = rng.random_range(0..3);

					match op {
						0 => {
							tree.insert(key, t);
						}
						1 => {
							tree.remove(&key);
						}
						2 => {
							tree.lookup(&key, |v| *v);
						}
						_ => unreachable!(),
					}
				}
			})
		})
		.collect();

	for h in handles {
		h.join().unwrap();
	}

	// Final invariant check
	tree.assert_invariants();

	// Verify tree is still functional
	let len = tree.len();
	assert!(len <= key_range as usize);
}

/// High contention single key stress test
#[test]
fn stress_high_contention_single_key() {
	let tree = Arc::new(Tree::<i32, i32>::new());
	let num_threads = 8;
	let iterations = 1000;

	tree.insert(42, 0);

	let handles: Vec<_> = (0..num_threads)
		.map(|t| {
			let tree = Arc::clone(&tree);
			thread::spawn(move || {
				for i in 0..iterations {
					if i % 2 == 0 {
						tree.insert(42, t);
					} else {
						tree.lookup(&42, |v| *v);
					}
				}
			})
		})
		.collect();

	for h in handles {
		h.join().unwrap();
	}

	// Final invariant check
	tree.assert_invariants();

	// Key should still exist
	assert!(tree.lookup(&42, |v| *v).is_some());
	assert_eq!(tree.len(), 1);
}

/// Sustained mixed operations stress test
#[test]
fn stress_sustained_mixed_operations() {
	let tree = Arc::new(Tree::<i32, i32>::new());
	let num_threads = 4;
	let duration_ms = 500;

	let running = Arc::new(AtomicUsize::new(1));

	let handles: Vec<_> = (0..num_threads)
		.map(|t| {
			let tree = Arc::clone(&tree);
			let running = Arc::clone(&running);
			thread::spawn(move || {
				let mut rng = rand::rng();
				let mut ops = 0u64;

				while running.load(Ordering::Relaxed) == 1 {
					let key: i32 = rng.random_range(0..1000);
					let op: u8 = rng.random_range(0..10);

					match op {
						0..=3 => {
							tree.insert(key, t);
						}
						4..=5 => {
							tree.remove(&key);
						}
						6..=9 => {
							tree.lookup(&key, |v| *v);
						}
						_ => unreachable!(),
					}
					ops += 1;
				}

				ops
			})
		})
		.collect();

	// Let it run for the specified duration
	thread::sleep(Duration::from_millis(duration_ms));
	running.store(0, Ordering::Relaxed);

	let total_ops: u64 = handles.into_iter().map(|h| h.join().unwrap()).sum();

	// Final invariant check
	tree.assert_invariants();

	// Should have performed many operations
	assert!(total_ops > 100, "Only {} operations performed", total_ops);
}

/// Large scale concurrent inserts stress test
#[test]
fn stress_large_scale_concurrent_inserts() {
	let tree = Arc::new(Tree::<i32, i32>::new());
	let num_threads = 8;
	let entries_per_thread = 1000;

	let handles: Vec<_> = (0..num_threads)
		.map(|t| {
			let tree = Arc::clone(&tree);
			thread::spawn(move || {
				for i in 0..entries_per_thread {
					let key = t * entries_per_thread + i;
					tree.insert(key, key * 10);
				}
			})
		})
		.collect();

	for h in handles {
		h.join().unwrap();
	}

	// Final invariant check
	tree.assert_invariants();

	assert_eq!(tree.len(), (num_threads * entries_per_thread) as usize);
}

/// Producer-consumer pattern stress test
#[test]
fn stress_producer_consumer() {
	let tree = Arc::new(Tree::<i32, i32>::new());
	let num_producers = 4;
	let num_consumers = 4;
	let entries_per_producer = 500;

	let produced = Arc::new(AtomicUsize::new(0));

	// Producers insert entries
	let producer_handles: Vec<_> = (0..num_producers)
		.map(|p| {
			let tree = Arc::clone(&tree);
			let produced = Arc::clone(&produced);
			thread::spawn(move || {
				for i in 0..entries_per_producer {
					let key = p * entries_per_producer + i;
					tree.insert(key, key * 10);
					produced.fetch_add(1, Ordering::Relaxed);
				}
			})
		})
		.collect();

	// Consumers lookup entries
	let consumer_handles: Vec<_> = (0..num_consumers)
		.map(|_| {
			let tree = Arc::clone(&tree);
			let produced = Arc::clone(&produced);
			thread::spawn(move || {
				let mut rng = rand::rng();
				let mut found = 0u64;
				let total_entries = num_producers * entries_per_producer;

				while produced.load(Ordering::Relaxed) < total_entries as usize {
					let key: i32 = rng.random_range(0..total_entries);
					if tree.lookup(&key, |v| *v).is_some() {
						found += 1;
					}
				}

				found
			})
		})
		.collect();

	for h in producer_handles {
		h.join().unwrap();
	}

	let total_found: u64 = consumer_handles.into_iter().map(|h| h.join().unwrap()).sum();

	// Final invariant check
	tree.assert_invariants();

	// Consumers should have found some entries
	assert!(total_found > 0);

	// All entries should be present
	assert_eq!(tree.len(), (num_producers * entries_per_producer) as usize);
}

// ===========================================================================
// Regression: concurrent lookup of values containing interior pointers
// ===========================================================================
//
// Regression for https://github.com/surrealdb/ferntree/issues/4
//
// Before the fix, `Tree::lookup` ran the user closure under a purely
// optimistic guard. For value types with interior pointers (Vec, String,
// `bytes::Bytes`, etc.) a torn read of the value's length/tag bytes during a
// concurrent mutation triggers UB inside the value type's own methods. The
// fix is to hold a shared lock on the leaf for the duration of the closure.

#[test]
fn concurrent_lookup_with_interior_pointer_values() {
	// Tiny key space so almost every commit lands on the same leaf as a
	// concurrent one — this is what the original repro relies on.
	const NUM_KEYS: u32 = 16;
	const WRITERS: usize = 12;
	const READERS: usize = 12;
	const OPS_PER_THREAD: usize = 2_000;

	// `Vec<u64>` has the same "interior pointer" hazard the original
	// reproducer targeted: a torn read of len/ptr/cap during concurrent
	// `push` / `truncate` triggers UB inside `Vec`'s own methods. The fix
	// is to hold a shared lock on the leaf across the lookup closure.
	type Versions = Vec<u64>;

	let tree: Arc<Tree<u32, Versions>> = Arc::new(Tree::new());

	// Pre-populate every key so writers exercise the same leaf-edit path as
	// the original consumer (seek_exact -> mutate the Vec in place).
	for k in 0..NUM_KEYS {
		let sv: Versions = vec![0];
		tree.insert(k, sv);
	}

	let mut handles = Vec::new();

	for w in 0..WRITERS {
		let tree = Arc::clone(&tree);
		handles.push(thread::spawn(move || {
			let mut rng = StdRng::seed_from_u64(0xA110C + w as u64);
			for i in 0..OPS_PER_THREAD {
				let key = rng.random_range(0..NUM_KEYS);
				let mut iter = tree.raw_iter_mut();
				if iter.seek_exact(&key) {
					let (_, versions) = iter.next().expect("seek_exact returned true");
					// Mix of grow and shrink so the Vec triggers
					// reallocations repeatedly.
					if versions.len() > 4 && (i & 1) == 0 {
						versions.truncate(2);
					} else {
						versions.push(i as u64);
					}
				}
			}
		}));
	}

	for r in 0..READERS {
		let tree = Arc::clone(&tree);
		handles.push(thread::spawn(move || {
			let mut rng = StdRng::seed_from_u64(0xBEE5 + r as u64);
			for _ in 0..OPS_PER_THREAD {
				let key = rng.random_range(0..NUM_KEYS);
				// Touch the Vec in a way that forces method dispatch
				// (sums the slice). Before the fix, a torn read of the
				// ptr/len during a concurrent push/truncate would read
				// freed heap memory.
				let _ = tree.lookup(&key, |v| v.iter().copied().sum::<u64>());
			}
		}));
	}

	for h in handles {
		h.join().expect("worker thread panicked");
	}

	tree.assert_invariants();
}

// ===========================================================================
// Optimistic-read fast-path concurrent stress (V- and K-deferred drops)
// ===========================================================================
//
// These tests previously lived here (out of `cargo miri test --lib`)
// because Miri's data-race detector flagged the protocol's
// unsynchronised non-atomic V/K reads racing with non-atomic writes.
// The atomic-mirror migration (see [`ferntree::atomic_slot`]) replaces
// the underlying `ptr::read` / `mem::replace` with `Acquire` /
// `Release` atomic loads and stores, so the tests now live in
// `src/lib.rs::tests` and run cleanly under
// `cargo miri test --all-features --lib`. See
// `epoch_deferred_drop_optimistic_reader_vs_defer_writer` and
// `k_deferred_drop_optimistic_reader_vs_defer_writer` there.

/// Mock refcounted blob acting as `V`. Cheap clone via `Arc`. Opts into
/// `EPOCH_DEFERRED_DROP = true` so writes route the displaced V through
/// the epoch GC, keeping interior pointers alive across the reader's
/// snapshot/use window.
#[derive(Clone)]
#[allow(dead_code)] // retained for any future stress tests this file gains
struct RefcountedBlob(Arc<Vec<u8>>);

// SAFETY: `RefcountedBlob` wraps `Arc<Vec<u8>>` which is `Send + Sync`. A
// bitwise snapshot followed by a successful version recheck yields a valid
// `Arc`. `EPOCH_DEFERRED_DROP = true`, combined with using `insert_defer`
// / `remove_defer` exclusively for writes, keeps the underlying `Vec`'s
// buffer alive across the reader's snapshot/use window.
unsafe impl OptimisticRead for RefcountedBlob {
	const EPOCH_DEFERRED_DROP: bool = true;
	type Slot = ferntree::atomic_slot::BoxedSlot<Self>;
}

#[cfg(any())]
#[test]
fn epoch_deferred_drop_optimistic_reader_vs_defer_writer() {
	let tree: Arc<Tree<i32, RefcountedBlob>> = Arc::new(Tree::new());
	let stop = Arc::new(AtomicBool::new(false));

	for i in 0..200 {
		tree.insert_defer(i, RefcountedBlob(Arc::new(vec![i as u8; 32])));
	}

	let mut handles = Vec::new();
	for _ in 0..4 {
		let tree = Arc::clone(&tree);
		let stop = Arc::clone(&stop);
		handles.push(thread::spawn(move || {
			while !stop.load(Ordering::Relaxed) {
				for k in 0..200 {
					if let Some(blob) = tree.lookup_optimistic(&k, |v| v.clone()) {
						// Touch the buffer so the optimiser keeps the clone
						// alive past the lookup — if the buffer were freed
						// behind our back this would UB / segfault under
						// ASan/TSan.
						let s: usize = blob.0.iter().map(|&b| b as usize).sum();
						std::hint::black_box(s);
					}
				}
			}
		}));
	}

	let writer = {
		let tree = Arc::clone(&tree);
		let stop = Arc::clone(&stop);
		thread::spawn(move || {
			for round in 0..50 {
				for k in 0..200 {
					let v = RefcountedBlob(Arc::new(vec![(k + round) as u8; 32]));
					tree.insert_defer(k, v);
				}
			}
			stop.store(true, Ordering::Relaxed);
		})
	};

	writer.join().unwrap();
	for h in handles {
		h.join().unwrap();
	}
}

/// Stress for the K-deferred-drop extension: K is a refcounted type with
/// `EPOCH_DEFERRED_DROP = true`, and writers alternate `insert_defer` /
/// `remove_defer` so leaf K is actually dropped (`remove_defer` defers
/// both K and V drops via the epoch GC). Readers `lookup_optimistic`
/// hammer the descent which snapshots K's interior pointer.
#[cfg(any())]
#[test]
fn k_deferred_drop_optimistic_reader_vs_defer_writer() {
	const KEYS: u8 = 64;
	let mk_key = |i: u8| RcKey(Arc::new(vec![i; 8]));

	let tree: Arc<Tree<RcKey, u64>> = Arc::new(Tree::new());
	let stop = Arc::new(AtomicBool::new(false));

	for i in 0..KEYS {
		tree.insert_defer(mk_key(i), i as u64);
	}

	let mut handles = Vec::new();
	for _ in 0..4 {
		let tree = Arc::clone(&tree);
		let stop = Arc::clone(&stop);
		handles.push(thread::spawn(move || {
			while !stop.load(Ordering::Relaxed) {
				for i in 0..KEYS {
					let k = mk_key(i);
					let _ = tree.lookup_optimistic(&k, |v| *v);
				}
			}
		}));
	}

	let writer = {
		let tree = Arc::clone(&tree);
		let stop = Arc::clone(&stop);
		thread::spawn(move || {
			for round in 0..50u64 {
				for i in 0..KEYS {
					let k = mk_key(i);
					if round % 2 == 0 {
						tree.remove_defer(&k);
					} else {
						tree.insert_defer(k, round * 100 + i as u64);
					}
				}
			}
			stop.store(true, Ordering::Relaxed);
		})
	};

	writer.join().unwrap();
	for h in handles {
		h.join().unwrap();
	}
}