rust/neon
1
0
Fork 0
mirror of https://github.com/neondatabase/neon.git synced 2026-05-26 17:40:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

implement deletion in art tree

Browse Source
This commit is contained in:
Heikki Linnakangas
2025-05-04 23:29:01 +03:00
parent 42df3e5453
commit 884e028a4a
8 changed files with 322 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files

83
Cargo.lock generated
View File

@@ -2595,6 +2595,18 @@ dependencies = [
"wasm-bindgen",
]
[[package]]
name = "getrandom"
version = "0.3.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "73fea8450eea4bac3940448fb7ae50d91f034f941199fcd9d909a5a07aa455f0"
dependencies = [
"cfg-if",
"libc",
"r-efi",
"wasi 0.14.2+wasi-0.2.4",
]
[[package]]
name = "gettid"
version = "0.1.3"
@@ -3817,7 +3829,7 @@ dependencies = [
"procfs",
"prometheus",
"rand 0.8.5",
"rand_distr",
"rand_distr 0.4.3",
"twox-hash",
]
@@ -3906,10 +3918,10 @@ name = "neonart"
version = "0.1.0"
dependencies = [
"crossbeam-utils",
"rand 0.8.5",
"rand 0.9.1",
"rand_distr 0.5.1",
"spin",
"tracing",
"zerocopy 0.8.24",
]
[[package]]
@@ -5328,7 +5340,7 @@ dependencies = [
"postgres_backend",
"pq_proto",
"rand 0.8.5",
"rand_distr",
"rand_distr 0.4.3",
"rcgen",
"redis",
"regex",
@@ -5431,6 +5443,12 @@ dependencies = [
"proc-macro2",
]
[[package]]
name = "r-efi"
version = "5.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "74765f6d916ee2faa39bc8e68e4f3ed8949b48cccdac59983d287a7cb71ce9c5"
[[package]]
name = "rand"
version = "0.7.3"
@@ -5455,6 +5473,16 @@ dependencies = [
"rand_core 0.6.4",
]
[[package]]
name = "rand"
version = "0.9.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9fbfd9d094a40bf3ae768db9361049ace4c0e04a4fd6b359518bd7b73a73dd97"
dependencies = [
"rand_chacha 0.9.0",
"rand_core 0.9.3",
]
[[package]]
name = "rand_chacha"
version = "0.2.2"
@@ -5475,6 +5503,16 @@ dependencies = [
"rand_core 0.6.4",
]
[[package]]
name = "rand_chacha"
version = "0.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb"
dependencies = [
"ppv-lite86",
"rand_core 0.9.3",
]
[[package]]
name = "rand_core"
version = "0.5.1"
@@ -5493,6 +5531,15 @@ dependencies = [
"getrandom 0.2.11",
]
[[package]]
name = "rand_core"
version = "0.9.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "99d9a13982dcf210057a8a78572b2217b667c3beacbf3a0d8b454f6f82837d38"
dependencies = [
"getrandom 0.3.2",
]
[[package]]
name = "rand_distr"
version = "0.4.3"
@@ -5503,6 +5550,16 @@ dependencies = [
"rand 0.8.5",
]
[[package]]
name = "rand_distr"
version = "0.5.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6a8615d50dcf34fa31f7ab52692afec947c4dd0ab803cc87cb3b0b4570ff7463"
dependencies = [
"num-traits",
"rand 0.9.1",
]
[[package]]
name = "rand_hc"
version = "0.2.0"
@@ -8240,6 +8297,15 @@ version = "0.11.0+wasi-snapshot-preview1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423"
[[package]]
name = "wasi"
version = "0.14.2+wasi-0.2.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9683f9a5a998d873c0d21fcbe3c083009670149a8fab228644b8bd36b2c48cb3"
dependencies = [
"wit-bindgen-rt",
]
[[package]]
name = "wasite"
version = "0.1.0"
@@ -8597,6 +8663,15 @@ dependencies = [
"windows-sys 0.48.0",
]
[[package]]
name = "wit-bindgen-rt"
version = "0.39.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6f42320e61fe2cfd34354ecb597f86f413484a798ba44a8ca1165c58d42da6c1"
dependencies = [
"bitflags 2.8.0",
]
[[package]]
name = "workspace_hack"
version = "0.1.0"

5
libs/neonart/Cargo.toml
View File

@@ -9,5 +9,6 @@ crossbeam-utils.workspace = true
spin.workspace = true
tracing.workspace = true
rand.workspace = true # for tests
zerocopy = "0.8"
[dev-dependencies]
rand = "0.9.1"
rand_distr = "0.5.1"

43
libs/neonart/src/algorithm.rs
View File

@@ -60,11 +60,11 @@ pub(crate) fn search<'e, K: Key, V: Value>(
}
}
pub(crate) fn update_fn<'e, K: Key, V: Value, A: ArtAllocator<V>, F>(
pub(crate) fn update_fn<'e, 'g, K: Key, V: Value, A: ArtAllocator<V>, F>(
key: &K,
value_fn: F,
root: RootPtr<V>,
guard: &'e TreeWriteGuard<K, V, A>,
guard: &'g mut TreeWriteGuard<'e, K, V, A>,
) where
F: FnOnce(Option<&V>) -> Option<V>,
{
@@ -84,17 +84,17 @@ pub(crate) fn update_fn<'e, K: Key, V: Value, A: ArtAllocator<V>, F>(
key_bytes,
) {
Ok(()) => break,
Err(ArtError::ConcurrentUpdate) => continue, // retry
Err(ArtError::ConcurrentUpdate) => {
eprintln!("retrying");
continue; // retry
},
Err(ArtError::OutOfMemory) => {
panic!("todo: OOM: try to GC, propagate to caller");
},
Err(ArtError::GarbageQueueFull) => {
if guard.collect_garbage() {
continue;
}
// FIXME: This can happen if someone is holding back the epoch. We should
// wait for the epoch to advance
panic!("todo: GC queue is full and couldn't free up space");
panic!("todo: GC queue is full");
},
}
}
@@ -144,12 +144,12 @@ fn lookup_recurse<'e, V: Value>(
}
// This corresponds to the 'insertOpt' function in the paper
pub(crate) fn update_recurse<'e, K: Key, V: Value, A: ArtAllocator<V>, F>(
pub(crate) fn update_recurse<'e, 'g, K: Key, V: Value, A: ArtAllocator<V>, F>(
key: &[u8],
value_fn: F,
node: NodeRef<'e, V>,
rparent: Option<(ReadLockedNodeRef<V>, u8)>,
guard: &'e TreeWriteGuard<K, V, A>,
guard: &'g mut TreeWriteGuard<'e, K, V, A>,
level: usize,
orig_key: &[u8],
) -> Result<(), ArtError>
@@ -211,7 +211,7 @@ where
match next_node {
ChildOrValue::Value(existing_value_ptr) => {
assert!(key.len() == 1);
let wnode = rnode.upgrade_to_write_lock_or_restart()?;
let mut wnode = rnode.upgrade_to_write_lock_or_restart()?;
// safety: Now that we have acquired the write lock, we have exclusive access to the
// value
@@ -219,7 +219,10 @@ where
if let Some(new_value) = value_fn(Some(vmut)) {
*vmut = new_value;
} else {
// TODO: Treat this as deletion?
// TODO: Shrink the node
// TODO: If the node becomes empty, unlink it from parent
wnode.delete_value(key[0]);
}
wnode.write_unlock();
@@ -320,11 +323,11 @@ fn insert_split_prefix<'e, K: Key, V: Value, A: ArtAllocator<V>>(
// Allocate a node for the new value.
let new_value_node =
allocate_node_for_value(&key[common_prefix_len + 1..], value, guard.allocator)?;
allocate_node_for_value(&key[common_prefix_len + 1..], value, guard.tree_writer.allocator)?;
// Allocate a new internal node with the common prefix
// FIXME: deallocate 'new_value_node' on OOM
let mut prefix_node = node_ref::new_internal(&key[..common_prefix_len], guard.allocator)?;
let mut prefix_node = node_ref::new_internal(&key[..common_prefix_len], guard.tree_writer.allocator)?;
// Add the old node and the new nodes to the new internal node
prefix_node.insert_old_child(old_prefix[common_prefix_len], old_node);
@@ -348,36 +351,34 @@ fn insert_to_node<'e, K: Key, V: Value, A: ArtAllocator<V>>(
if wnode.is_leaf() {
wnode.insert_value(key[0], value);
} else {
let value_child = allocate_node_for_value(&key[1..], value, guard.allocator)?;
let value_child = allocate_node_for_value(&key[1..], value, guard.tree_writer.allocator)?;
wnode.insert_child(key[0], value_child.into_ptr());
}
Ok(())
}
// On entry: 'parent' and 'node' are locked
fn insert_and_grow<'e, K: Key, V: Value, A: ArtAllocator<V>>(
fn insert_and_grow<'e, 'g, K: Key, V: Value, A: ArtAllocator<V>>(
key: &[u8],
value: V,
wnode: &WriteLockedNodeRef<V>,
parent: &mut WriteLockedNodeRef<V>,
parent_key_byte: u8,
guard: &'e TreeWriteGuard<K, V, A>,
guard: &'g mut TreeWriteGuard<'e, K, V, A>,
) -> Result<(), ArtError> {
let mut bigger_node = wnode.grow(guard.allocator)?;
let mut bigger_node = wnode.grow(guard.tree_writer.allocator)?;
if wnode.is_leaf() {
bigger_node.insert_value(key[0], value);
} else {
// FIXME: deallocate 'bigger_node' on OOM
let value_child = allocate_node_for_value(&key[1..], value, guard.allocator)?;
let value_child = allocate_node_for_value(&key[1..], value, guard.tree_writer.allocator)?;
bigger_node.insert_new_child(key[0], value_child);
}
// Replace the pointer in the parent
parent.replace_child(parent_key_byte, bigger_node.into_ptr());
// FIXME: if this errors out, deallocate stuff we already allocated
guard.remember_obsolete_node(wnode.as_ptr())?;
guard.remember_obsolete_node(wnode.as_ptr());
Ok(())
}

85
libs/neonart/src/algorithm/node_ptr.rs
View File

@@ -404,6 +404,19 @@ impl<V: Value> NodePtr<V> {
}
}
pub(crate) fn delete_value(&mut self, key_byte: u8) {
match self.variant_mut() {
NodeVariantMut::Internal4(_)
| NodeVariantMut::Internal16(_)
| NodeVariantMut::Internal48(_)
| NodeVariantMut::Internal256(_) => panic!("delete_value called on internal node"),
NodeVariantMut::Leaf4(n) => n.delete_value(key_byte),
NodeVariantMut::Leaf16(n) => n.delete_value(key_byte),
NodeVariantMut::Leaf48(n) => n.delete_value(key_byte),
NodeVariantMut::Leaf256(n) => n.delete_value(key_byte),
}
}
pub(crate) fn deallocate(self, allocator: &impl ArtAllocator<V>) {
match self.variant() {
NodeVariant::Internal4(_) => allocator.dealloc_node_internal4(self.ptr.cast()),
@@ -766,7 +779,7 @@ impl<V: Value> NodeLeaf4<V> {
}
fn insert_value(&mut self, key_byte: u8, value: V) {
assert!(self.num_values < 16);
assert!(self.num_values < 4);
let idx = self.num_values as usize;
self.child_keys[idx] = key_byte;
@@ -797,6 +810,23 @@ impl<V: Value> NodeLeaf4<V> {
unsafe { ptr.write(init) };
ptr.into()
}
fn delete_value(&mut self, key_byte: u8) {
assert!(self.num_values <= 4);
for i in 0..self.num_values as usize {
if self.child_keys[i] == key_byte {
assert!(self.child_values[i].is_some());
if i < self.num_values as usize - 1 {
self.child_keys[i] = self.child_keys[self.num_values as usize - 1];
self.child_values[i] = std::mem::replace(&mut self.child_values[self.num_values as usize - 1], None);
}
self.num_values -= 1;
return;
}
}
panic!("key to delete not found in leaf4 node");
}
}
impl<V: Value> NodeLeaf16<V> {
@@ -859,6 +889,23 @@ impl<V: Value> NodeLeaf16<V> {
unsafe { ptr.write(init) };
ptr.into()
}
fn delete_value(&mut self, key_byte: u8) {
assert!(self.num_values <= 16);
for i in 0..self.num_values as usize {
if self.child_keys[i as usize] == key_byte {
assert!(self.child_values[i as usize].is_some());
if i < self.num_values as usize - 1 {
self.child_keys[i] = self.child_keys[self.num_values as usize - 1];
self.child_values[i] = std::mem::replace(&mut self.child_values[self.num_values as usize - 1], None);
}
self.num_values -= 1;
return;
}
}
panic!("key to delete not found in leaf16 node");
}
}
impl<V: Value> NodeLeaf48<V> {
@@ -921,6 +968,34 @@ impl<V: Value> NodeLeaf48<V> {
unsafe { ptr.write(init) };
ptr.into()
}
fn delete_value(&mut self, key_byte: u8) {
assert!(self.num_values <= 48);
let idx = self.child_indexes[key_byte as usize];
if idx == INVALID_CHILD_INDEX {
panic!("key to delete not found in leaf48 node");
}
self.child_indexes[key_byte as usize] = INVALID_CHILD_INDEX;
self.num_values -= 1;
if idx < self.num_values {
// Move all existing values with higher indexes down one position
for i in idx as usize ..self.num_values as usize {
self.child_values[i] = std::mem::replace(&mut self.child_values[i + 1], None);
}
// Update all higher indexes
for i in 0..256 {
if self.child_indexes[i] != INVALID_CHILD_INDEX {
if self.child_indexes[i] > idx {
self.child_indexes[i] -= 1;
}
assert!(self.child_indexes[i] < self.num_values);
}
}
}
}
}
impl<V: Value> NodeLeaf256<V> {
@@ -952,6 +1027,14 @@ impl<V: Value> NodeLeaf256<V> {
self.child_values[key_byte as usize] = Some(value);
self.num_values += 1;
}
fn delete_value(&mut self, key_byte: u8) {
if self.child_values[key_byte as usize].is_none() {
panic!("key to delete not found in leaf256 node");
}
self.child_values[key_byte as usize] = None;
self.num_values -= 1;
}
}
impl<V: Value> NodeInternal256<V> {

4
libs/neonart/src/algorithm/node_ref.rs
View File

@@ -152,6 +152,10 @@ impl<'e, V: Value> WriteLockedNodeRef<'e, V> {
self.ptr.insert_value(key_byte, value)
}
pub(crate) fn delete_value(&mut self, key_byte: u8) {
self.ptr.delete_value(key_byte)
}
pub(crate) fn grow<'a, A>(
&self,
allocator: &'a A,

122
libs/neonart/src/lib.rs
View File

@@ -128,6 +128,7 @@ mod epoch;
use algorithm::RootPtr;
use algorithm::node_ptr::NodePtr;
use std::collections::VecDeque;
use std::fmt::Debug;
use std::marker::PhantomData;
use std::ptr::NonNull;
@@ -163,8 +164,6 @@ pub struct Tree<V: Value> {
writer_attached: AtomicBool,
epoch: epoch::EpochShared,
garbage: spin::Mutex<GarbageQueue<V>>,
}
unsafe impl<V: Value + Sync> Sync for Tree<V> {}
@@ -172,43 +171,29 @@ unsafe impl<V: Value + Send> Send for Tree<V> {}
struct GarbageQueueFullError();
struct GarbageQueue<V> {
slots: [(NodePtr<V>, u64); MAX_GARBAGE],
front: usize,
back: usize,
}
struct GarbageQueue<V>(VecDeque<(NodePtr<V>, u64)>);
unsafe impl<V: Value + Sync> Sync for GarbageQueue<V> {}
unsafe impl<V: Value + Send> Send for GarbageQueue<V> {}
impl<V> GarbageQueue<V> {
fn new() -> GarbageQueue<V> {
GarbageQueue {
slots: [const { (NodePtr::null(), 0) }; MAX_GARBAGE],
front: 0,
back: 0,
}
GarbageQueue(VecDeque::with_capacity(MAX_GARBAGE))
}
fn remember_obsolete_node(
&mut self,
ptr: NodePtr<V>,
epoch: u64,
) -> Result<(), GarbageQueueFullError> {
if self.front == self.back.wrapping_add(MAX_GARBAGE) {
return Err(GarbageQueueFullError());
}
self.slots[self.front % MAX_GARBAGE] = (ptr, epoch);
self.front = self.front.wrapping_add(1);
Ok(())
) {
self.0.push_front((ptr, epoch));
}
fn next_obsolete(&mut self, cutoff_epoch: u64) -> Option<NodePtr<V>> {
if self.front == self.back {
return None;
}
let slot = &self.slots[self.back % MAX_GARBAGE];
// FIXME: performing wrapping comparison
if slot.1 < cutoff_epoch {
self.back += 1;
return Some(slot.0);
if let Some(back) = self.0.back() {
if back.1 < cutoff_epoch {
return Some(self.0.pop_back().unwrap().0);
}
}
None
}
@@ -237,6 +222,9 @@ where
epoch_handle: epoch::LocalHandle<'t>,
phantom_key: PhantomData<K>,
/// Obsolete nodes that cannot be recycled until their epoch expires.
garbage: spin::Mutex<GarbageQueue<V>>,
}
/// The backends have a reference to this. It cannot be used to modify the tree
@@ -260,7 +248,6 @@ impl<'a, 't: 'a, K: Key, V: Value, A: ArtAllocator<V>> TreeInitStruct<'t, K, V,
root: algorithm::new_root(allocator),
writer_attached: AtomicBool::new(false),
epoch: epoch::EpochShared::new(),
garbage: spin::Mutex::new(GarbageQueue::new()),
};
unsafe { tree_ptr.write(init) };
@@ -281,6 +268,7 @@ impl<'a, 't: 'a, K: Key, V: Value, A: ArtAllocator<V>> TreeInitStruct<'t, K, V,
allocator: self.allocator,
phantom_key: PhantomData,
epoch_handle: self.tree.epoch.register(),
garbage: spin::Mutex::new(GarbageQueue::new()),
}
}
@@ -294,13 +282,14 @@ impl<'a, 't: 'a, K: Key, V: Value, A: ArtAllocator<V>> TreeInitStruct<'t, K, V,
}
impl<'t, K: Key + Clone, V: Value, A: ArtAllocator<V>> TreeWriteAccess<'t, K, V, A> {
pub fn start_write(&'t self) -> TreeWriteGuard<'t, K, V, A> {
// TODO: grab epoch guard
pub fn start_write<'g>(&'t self) -> TreeWriteGuard<'g, K, V, A>
where 't: 'g
{
TreeWriteGuard {
allocator: self.allocator,
tree: &self.tree,
tree_writer: self,
epoch_pin: self.epoch_handle.pin(),
phantom_key: PhantomData,
created_garbage: false
}
}
@@ -344,49 +333,72 @@ pub struct TreeWriteGuard<'e, K, V, A>
where
K: Key,
V: Value,
A: ArtAllocator<V>,
{
tree: &'e Tree<V>,
allocator: &'e A,
tree_writer: &'e TreeWriteAccess<'e, K, V, A>,
epoch_pin: EpochPin<'e>,
phantom_key: PhantomData<K>,
created_garbage: bool,
}
impl<'t, K: Key, V: Value, A: ArtAllocator<V>> TreeWriteGuard<'t, K, V, A> {
pub fn insert(&mut self, key: &K, value: V) {
/// Get a value
pub fn get(&mut self, key: &K) -> Option<V> {
algorithm::search(key, self.tree_writer.tree.root, &self.epoch_pin)
}
/// Insert a value
pub fn insert(self, key: &K, value: V) {
self.update_with_fn(key, |_| Some(value))
}
pub fn update_with_fn<F>(&mut self, key: &K, value_fn: F)
/// Remove value
pub fn remove(self, key: &K) {
self.update_with_fn(key, |_| None)
}
/// Update key using the given function. All the other modifying operations are based on this.
///
/// The function is passed a reference to the existing value, if any. If the function
/// returns None, the value is removed from the tree (or if there was no existing value,
/// does nothing). If the function returns Some, the existing value is replaced, of if there
/// was no existing value, it is inserted.
pub fn update_with_fn<F>(mut self, key: &K, value_fn: F)
where
F: FnOnce(Option<&V>) -> Option<V>,
{
algorithm::update_fn(key, value_fn, self.tree.root, self)
let result = algorithm::update_fn(key, value_fn, self.tree_writer.tree.root, &mut self);
if self.created_garbage {
let n = self.collect_garbage();
eprintln!("collected {n} obsolete nodes");
}
result
}
pub fn get(&mut self, key: &K) -> Option<V> {
algorithm::search(key, self.tree.root, &self.epoch_pin)
}
fn remember_obsolete_node(&'t self, ptr: NodePtr<V>) -> Result<(), GarbageQueueFullError> {
self.tree
fn remember_obsolete_node(&mut self, ptr: NodePtr<V>) {
self.tree_writer
.garbage
.lock()
.remember_obsolete_node(ptr, self.epoch_pin.epoch)
.remember_obsolete_node(ptr, self.epoch_pin.epoch);
self.created_garbage = true;
}
// returns true if something was free'd up
fn collect_garbage(&'t self) -> bool {
let mut result = false;
self.tree.epoch.advance();
self.tree.epoch.broadcast();
// returns number of nodes recycled
fn collect_garbage(&self) -> usize {
self.tree_writer.tree.epoch.advance();
self.tree_writer.tree.epoch.broadcast();
let cutoff_epoch = self.tree.epoch.get_oldest();
let cutoff_epoch = self.tree_writer.tree.epoch.get_oldest();
let mut garbage_queue = self.tree.garbage.lock();
let mut result = 0;
let mut garbage_queue = self.tree_writer.garbage.lock();
while let Some(ptr) = garbage_queue.next_obsolete(cutoff_epoch) {
ptr.deallocate(self.allocator);
result = true;
ptr.deallocate(self.tree_writer.allocator);
result += 1;
}
result
}
@@ -400,7 +412,7 @@ impl<'t, K: Key, V: Value + Debug> TreeReadGuard<'t, K, V> {
}
impl<'t, K: Key, V: Value + Debug> TreeWriteGuard<'t, K, V, ArtMultiSlabAllocator<'t, V>> {
pub fn get_statistics(&self) -> ArtTreeStatistics {
self.allocator.get_statistics()
self.tree_writer.allocator.get_statistics()
}
}

67
libs/neonart/src/tests.rs
View File

@@ -1,16 +1,19 @@
use std::collections::HashSet;
use std::collections::BTreeMap;
use crate::ArtAllocator;
use crate::ArtMultiSlabAllocator;
use crate::TreeInitStruct;
use crate::{Key, Value};
use rand::seq::SliceRandom;
use rand::thread_rng;
use rand::Rng;
use rand_distr::Zipf;
const TEST_KEY_LEN: usize = 16;
#[derive(Clone, Copy, Debug)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
struct TestKey([u8; TEST_KEY_LEN]);
impl Key for TestKey {
@@ -39,13 +42,10 @@ fn test_inserts<K: Into<TestKey> + Copy>(keys: &[K]) {
let tree_writer = init_struct.attach_writer();
for (idx, k) in keys.iter().enumerate() {
let mut w = tree_writer.start_write();
let w = tree_writer.start_write();
w.insert(&(*k).into(), idx);
eprintln!("INSERTED {:?}", Into::<TestKey>::into(*k));
}
//tree_writer.start_read().dump();
for (idx, k) in keys.iter().enumerate() {
let r = tree_writer.start_read();
let value = r.get(&(*k).into());
@@ -67,7 +67,7 @@ fn dense() {
// Do the same in random orders
for _ in 1..10 {
keys.shuffle(&mut thread_rng());
keys.shuffle(&mut rand::rng());
test_inserts(&keys);
}
}
@@ -90,3 +90,56 @@ fn sparse() {
}
test_inserts(&keys);
}
#[derive(Clone, Copy, Debug)]
struct TestOp(TestKey, Option<usize>);
fn apply_op<A: ArtAllocator<usize>>(op: &TestOp, tree: &crate::TreeWriteAccess<TestKey, usize, A>, shadow: &mut BTreeMap<TestKey, usize>) {
eprintln!("applying op: {op:?}");
// apply the change to the shadow tree first
let shadow_existing = if let Some(v) = op.1 {
shadow.insert(op.0, v)
} else {
shadow.remove(&op.0)
};
// apply to Art tree
let w = tree.start_write();
w.update_with_fn(&op.0, |existing| {
assert_eq!(existing, shadow_existing.as_ref());
return op.1;
});
}
#[test]
fn random_ops() {
const MEM_SIZE: usize = 10000000;
let mut area = Box::new_uninit_slice(MEM_SIZE);
let allocator = ArtMultiSlabAllocator::new(&mut area);
let init_struct = TreeInitStruct::<TestKey, usize, _>::new(allocator);
let tree_writer = init_struct.attach_writer();
let mut shadow: std::collections::BTreeMap<TestKey, usize> = BTreeMap::new();
let distribution = Zipf::new(u128::MAX as f64, 1.1).unwrap();
let mut rng = rand::rng();
for i in 0..100000 {
let key: TestKey = (rng.sample(distribution) as u128).into();
let op = TestOp(
key,
if rng.random_bool(0.75) {
Some(i)
} else {
None
},
);
apply_op(&op, &tree_writer, &mut shadow);
}
}

7
pgxn/neon/communicator/src/integrated_cache.rs
View File

@@ -267,8 +267,7 @@ impl<'t> IntegratedCacheWriteAccess<'t> {
}
pub fn remember_rel_size(&'t self, rel: &RelTag, nblocks: u32) {
let mut w = self.cache_tree.start_write();
let w = self.cache_tree.start_write();
w.insert(
&TreeKey::from(rel),
TreeEntry::Rel(RelEntry {
@@ -286,7 +285,7 @@ impl<'t> IntegratedCacheWriteAccess<'t> {
lw_lsn: Lsn,
) {
if let Some(file_cache) = self.file_cache.as_ref() {
let mut w = self.cache_tree.start_write();
let w = self.cache_tree.start_write();
let key = TreeKey::from((rel, block_number));
@@ -324,7 +323,7 @@ impl<'t> IntegratedCacheWriteAccess<'t> {
/// Forget information about given relation in the cache. (For DROP TABLE and such)
pub fn forget_rel(&'t self, rel: &RelTag) {
// FIXME: not implemented properly. smgrexists() would still return true for this
let mut w = self.cache_tree.start_write();
let w = self.cache_tree.start_write();
w.insert(
&TreeKey::from(rel),
TreeEntry::Rel(RelEntry { nblocks: None }),