From 89b2782049ea491aa6611cb8da17e30d32c1933c Mon Sep 17 00:00:00 2001 From: euzu Date: Sun, 28 Dec 2025 13:34:53 +0100 Subject: [PATCH] BPlusTree space management for smaller value types --- backend/src/api/endpoints/websocket_api.rs | 4 +- backend/src/repository/bplustree.rs | 305 ++++++++++++--------- 2 files changed, 178 insertions(+), 131 deletions(-) diff --git a/backend/src/api/endpoints/websocket_api.rs b/backend/src/api/endpoints/websocket_api.rs index 3ec3d34da..f9f25df94 100644 --- a/backend/src/api/endpoints/websocket_api.rs +++ b/backend/src/api/endpoints/websocket_api.rs @@ -305,7 +305,7 @@ async fn handle_socket(mut socket: WebSocket, app_state: Arc, auth_req maybe_msg = socket.recv() => { if let Some(msg) = maybe_msg { if let Err(e) = handle_incoming_message(msg, &mut socket, &mut handler, &app_state, auth_required, secret_key.as_ref()).await { - error!("WebSocket message handling error: {e}"); + trace!("WebSocket message handling error: {e}"); break; } } else { @@ -315,7 +315,7 @@ async fn handle_socket(mut socket: WebSocket, app_state: Arc, auth_req Ok(event) = event_rx.recv() => { if let Err(e) = handle_event_message(&mut socket, event, &handler).await { - error!("Failed to send ws event: {e}"); + trace!("Failed to send ws event: {e}"); break; } } diff --git a/backend/src/repository/bplustree.rs b/backend/src/repository/bplustree.rs index c4024fc04..7d1b16d99 100644 --- a/backend/src/repository/bplustree.rs +++ b/backend/src/repository/bplustree.rs @@ -301,83 +301,71 @@ where fn serialize_to_block( &self, file: &mut W, - buffer: &mut Vec, // must be length BLOCK_SIZE + buffer: &mut Vec, offset: u64, ) -> io::Result { - // on-disk layout with minimal allocations. - let mut current_offset = offset; - - // Set node type (no need to zero entire buffer upfront) - let buffer_slice = &mut buffer[..]; - buffer_slice[0] = u8::from(self.is_leaf); - let mut write_pos = FLAG_SIZE; - - // ---- Write keys (length + bytes) into the first block ---- let keys_encoded = binary_serialize(&self.keys)?; - let keys_len = keys_encoded.len(); - buffer_slice[write_pos..write_pos + LEN_SIZE] - .copy_from_slice(&u32::try_from(keys_len).map_err(to_io_error)?.to_le_bytes()); - write_pos += LEN_SIZE; - // NOTE: By design of the legacy layout, keys are expected to fit into the first block. - buffer_slice[write_pos..write_pos + keys_len].copy_from_slice(&keys_encoded); - write_pos += keys_len; - drop(keys_encoded); - + let keys_len = u32::try_from(keys_encoded.len()).map_err(to_io_error)?; + if self.is_leaf { - // ---- Leaf nodes: write value info (offset + length) ---- - let info_encoded = binary_serialize(&self.value_info)?; - let info_len = info_encoded.len(); - - // CRITICAL CHECK: Ensure metadata fits in block - if write_pos + LEN_SIZE + info_len > BLOCK_SIZE { - return Err(io::Error::other(format!("Leaf node overflow: keys ({keys_len}) + value_info ({info_len}) exceeds block size"))); - } - - buffer_slice[write_pos..write_pos + LEN_SIZE] - .copy_from_slice(&u32::try_from(info_len).map_err(to_io_error)?.to_le_bytes()); - write_pos += LEN_SIZE; - buffer_slice[write_pos..write_pos + info_len].copy_from_slice(&info_encoded); - write_pos += info_len; - - // Zero unused portion and write first block - if write_pos < BLOCK_SIZE { - buffer_slice[write_pos..BLOCK_SIZE].fill(0u8); - } - file.seek(SeekFrom::Start(offset))?; - file.write_all(&buffer_slice[..BLOCK_SIZE])?; - current_offset += BLOCK_SIZE as u64; + let info_encoded = binary_serialize(&self.value_info)?; + let info_len = u32::try_from(info_encoded.len()).map_err(to_io_error)?; + + let content_size = FLAG_SIZE + LEN_SIZE + keys_encoded.len() + LEN_SIZE + info_encoded.len(); + let blocks = content_size.div_ceil(BLOCK_SIZE); + + file.seek(SeekFrom::Start(offset))?; + + let mut data = Vec::with_capacity(blocks * BLOCK_SIZE); + data.push(1u8); + data.extend_from_slice(&keys_len.to_le_bytes()); + data.extend_from_slice(&keys_encoded); + data.extend_from_slice(&info_len.to_le_bytes()); + data.extend_from_slice(&info_encoded); + + let pad_len = (blocks * BLOCK_SIZE) - data.len(); + if pad_len > 0 { + data.extend(std::iter::repeat(0).take(pad_len)); + } + + file.write_all(&data)?; + + Ok(offset + (blocks as u64 * BLOCK_SIZE as u64)) } else { - // Internal node: write breadth-first style pointers - let pointer_offset_within_first_block = offset + write_pos as u64; - - // Zero unused portion and write first block - if write_pos < BLOCK_SIZE { - buffer_slice[write_pos..BLOCK_SIZE].fill(0u8); - } - file.seek(SeekFrom::Start(offset))?; - file.write_all(&buffer_slice[..BLOCK_SIZE])?; - current_offset += BLOCK_SIZE as u64; - - let mut pointer_vec: Vec = Vec::with_capacity(self.children.len()); - for child in &self.children { - pointer_vec.push(current_offset); - current_offset = child.serialize_to_block(file, buffer, current_offset)?; - } - - let pointer_encoded = binary_serialize(&pointer_vec)?; - let pointer_len = u32::try_from(pointer_encoded.len()).map_err(to_io_error)?; - - // CRITICAL CHECK: Ensure pointers fit in block - if write_pos + LEN_SIZE + pointer_encoded.len() > BLOCK_SIZE { - return Err(io::Error::other(format!("Internal node overflow during recursion: keys ({}) + pointers ({}) exceeds block size", keys_len, pointer_encoded.len()))); - } - - file.seek(SeekFrom::Start(pointer_offset_within_first_block))?; - file.write_all(&pointer_len.to_le_bytes())?; - file.write_all(&pointer_encoded)?; + let ptr_count = self.children.len(); + let ptr_encoded_size = 8 + 8 * ptr_count; + + let content_size = FLAG_SIZE + LEN_SIZE + keys_encoded.len() + LEN_SIZE + ptr_encoded_size; + let blocks_needed = content_size.div_ceil(BLOCK_SIZE); + + let parent_start = offset; + let mut current_offset = parent_start + (blocks_needed as u64 * BLOCK_SIZE as u64); + + let mut pointers = Vec::with_capacity(ptr_count); + for child in &self.children { + pointers.push(current_offset); + current_offset = child.serialize_to_block(file, buffer, current_offset)?; + } + + let pointers_encoded = binary_serialize(&pointers)?; + let pointers_len = u32::try_from(pointers_encoded.len()).map_err(to_io_error)?; + + file.seek(SeekFrom::Start(parent_start))?; + let mut data = Vec::with_capacity(blocks_needed * BLOCK_SIZE); + data.push(0u8); + data.extend_from_slice(&keys_len.to_le_bytes()); + data.extend_from_slice(&keys_encoded); + data.extend_from_slice(&pointers_len.to_le_bytes()); + data.extend_from_slice(&pointers_encoded); + + let pad_len = (blocks_needed * BLOCK_SIZE) - data.len(); + if pad_len > 0 { + data.extend(std::iter::repeat(0).take(pad_len)); + } + file.write_all(&data)?; + + Ok(current_offset) } - - Ok(current_offset) } /// Serialize the tree in breadth-first order for better disk locality @@ -642,68 +630,71 @@ where fn deserialize_from_block( file: &mut R, - buffer: &mut [u8], + buffer: &mut Vec, offset: u64, nested: bool, ) -> io::Result<(Self, Option>)> { - // Read the full first block into buffer (always aligned on-disk) file.seek(SeekFrom::Start(offset))?; - file.read_exact(buffer)?; - - let slice = &buffer[..]; - - // Node type - let is_leaf = slice[0] == 1u8; - let mut read_pos = FLAG_SIZE; - - // ---- Keys ---- - let keys_length = u32_from_bytes(&slice[read_pos..read_pos + LEN_SIZE])? as usize; + + let header_required = FLAG_SIZE + LEN_SIZE; + if buffer.len() < header_required { + buffer.resize(header_required, 0); + } + + file.read_exact(&mut buffer[0..header_required])?; + + let is_leaf = buffer[0] != 0; + let keys_len = u32_from_bytes(&buffer[FLAG_SIZE..FLAG_SIZE + LEN_SIZE])? as usize; + + let min_required = header_required + keys_len + LEN_SIZE; + if buffer.len() < min_required { + buffer.resize(min_required, 0); + } + + file.read_exact(&mut buffer[header_required..min_required])?; + + let mut read_pos = FLAG_SIZE + LEN_SIZE; + let keys: Vec = binary_deserialize(&buffer[read_pos..read_pos + keys_len])?; + read_pos += keys_len; + + let payload_len = u32_from_bytes(&buffer[read_pos..read_pos + LEN_SIZE])? as usize; read_pos += LEN_SIZE; - let keys: Vec = binary_deserialize(&slice[read_pos..read_pos + keys_length])?; - read_pos += keys_length; - - // ---- Value info (offset, length) for leaf nodes ---- - let (value_info, values): (Vec, Vec) = if is_leaf { - // Read value_info - let info_length = u32_from_bytes(&slice[read_pos..read_pos + LEN_SIZE])? as usize; - read_pos += LEN_SIZE; - let info: Vec = binary_deserialize(&slice[read_pos..read_pos + info_length])?; - - // Values are loaded on-demand when nested=true - if nested { - let mut vals = Vec::with_capacity(info.len()); - for value_info in &info { - let value = Self::load_value_from_info(file, value_info)?; - vals.push(value); - } - (info, vals) - } else { - (info, Vec::new()) - } + + let total_required = min_required + payload_len; + if buffer.len() < total_required { + buffer.resize(total_required, 0); + } + + file.read_exact(&mut buffer[min_required..total_required])?; + + let (value_info, values, children, children_pointer) = if is_leaf { + let info: Vec = binary_deserialize(&buffer[read_pos..read_pos + payload_len])?; + let vals = if nested { + let mut v = Vec::with_capacity(info.len()); + for i in &info { + v.push(Self::load_value_from_info(file, i)?); + } + v + } else { + Vec::new() + }; + (info, vals, Vec::new(), None) } else { - (Vec::new(), Vec::new()) + let pointers: Vec = binary_deserialize(&buffer[read_pos..read_pos + payload_len])?; + let nodes = if nested { + let mut n = Vec::with_capacity(pointers.len()); + let mut child_buf = Vec::with_capacity(BLOCK_SIZE); + for &ptr in &pointers { + let (child, _) = Self::deserialize_from_block(file, &mut child_buf, ptr, nested)?; + n.push(child); + } + n + } else { + Vec::new() + }; + (Vec::new(), Vec::new(), nodes, Some(pointers)) }; - - // ---- Pointers for internal nodes ---- - let (children, children_pointer): (Vec, Option>) = if is_leaf { - (Vec::new(), None) - } else { - let pointers_length = u32_from_bytes(&slice[read_pos..read_pos + LEN_SIZE])? as usize; - read_pos += LEN_SIZE; - let pointers: Vec = binary_deserialize(&slice[read_pos..read_pos + pointers_length])?; - if nested { - let mut nodes = Vec::with_capacity(pointers.len()); - let mut child_buffer = vec![0u8; BLOCK_SIZE]; - for &ptr in &pointers { - let (child, _) = Self::deserialize_from_block(file, &mut child_buffer, ptr, nested)?; - nodes.push(child); - } - (nodes, None) - } else { - (Vec::new(), Some(pointers)) - } - }; - + Ok((Self { keys, children, is_leaf, value_info, values }, children_pointer)) } @@ -886,7 +877,7 @@ where root: BPlusTreeNode::::new(true), inner_order, leaf_order, - dirty: true, + dirty: true, // an empty tree is stored! } } @@ -1013,7 +1004,7 @@ where fn query_tree( file: &mut R, - buffer: &mut [u8], + buffer: &mut Vec, cache: &mut IndexMap>, key: &K, start_offset: u64, @@ -1082,7 +1073,7 @@ where fn query_tree_le( file: &mut R, - buffer: &mut [u8], + buffer: &mut Vec, cache: &mut IndexMap>, key: &K, start_offset: u64, @@ -2770,5 +2761,61 @@ mod tests { Ok(()) } + + #[test] + fn test_large_keys_multiblock_node() -> io::Result<()> { + let tempdir = tempfile::tempdir()?; + let filepath = tempdir.path().join("tree_multiblock.bin"); + + let mut tree = BPlusTree::::new(); + + // 5 keys of 2000 bytes each. Total ~10KB keys. + // Should span ~3 blocks (4KB each). + for i in 0..5 { + let key = format!("{:04}{}", i, "a".repeat(2000)); + tree.insert(key, i); + } + + tree.store(&filepath)?; + drop(tree); + + let loaded = BPlusTree::::load(&filepath)?; + for i in 0..5 { + let key = format!("{:04}{}", i, "a".repeat(2000)); + assert_eq!(loaded.query(&key), Some(i).as_ref()); + } + Ok(()) + } + + #[test] + fn test_upsert_multiblock_node() -> io::Result<()> { + let tempdir = tempfile::tempdir()?; + let filepath = tempdir.path().join("tree_multiblock_upsert.bin"); + + let mut tree = BPlusTree::::new(); + tree.store(&filepath)?; + drop(tree); + + let mut updater = BPlusTreeUpdate::::try_new(&filepath)?; + + // Upsert large keys + let mut batch = Vec::new(); + let keys: Vec = (0..5).map(|i| format!("{:04}{}", i, "b".repeat(2000))).collect(); + let vals: Vec = (0..5).collect(); + + for i in 0..5 { + batch.push((&keys[i], &vals[i])); + } + + updater.upsert_batch(&batch)?; + drop(updater); + + let mut query = BPlusTreeQuery::::try_new(&filepath)?; + for i in 0..5 { + let val = query.query(&keys[i]).unwrap(); + assert_eq!(val, vals[i]); + } + Ok(()) + } }