From ee64fb149c2927450f8755df1ab8b31bc3c24b97 Mon Sep 17 00:00:00 2001 From: Jason Volk Date: Tue, 2 Jul 2024 05:56:10 +0000 Subject: [PATCH] optimize get w/ zero-copy ref handle Signed-off-by: Jason Volk --- src/database/handle.rs | 25 +++++++++++++++++++++++++ src/database/map.rs | 8 ++++---- src/database/mod.rs | 2 ++ src/service/globals/data.rs | 2 +- src/service/globals/migrations.rs | 6 +++--- src/service/rooms/alias/data.rs | 2 +- src/service/rooms/event_handler/mod.rs | 4 ++-- src/service/rooms/timeline/data.rs | 2 +- src/service/rooms/timeline/mod.rs | 7 +++++-- src/service/transaction_ids/data.rs | 2 +- src/service/transaction_ids/mod.rs | 2 +- 11 files changed, 46 insertions(+), 16 deletions(-) create mode 100644 src/database/handle.rs diff --git a/src/database/handle.rs b/src/database/handle.rs new file mode 100644 index 00000000..787a5cea --- /dev/null +++ b/src/database/handle.rs @@ -0,0 +1,25 @@ +use std::ops::Deref; + +use rocksdb::DBPinnableSlice; + +pub struct Handle<'a> { + val: DBPinnableSlice<'a>, +} + +impl<'a> From> for Handle<'a> { + fn from(val: DBPinnableSlice<'a>) -> Self { + Self { + val, + } + } +} + +impl Deref for Handle<'_> { + type Target = [u8]; + + fn deref(&self) -> &Self::Target { &self.val } +} + +impl AsRef<[u8]> for Handle<'_> { + fn as_ref(&self) -> &[u8] { &self.val } +} diff --git a/src/database/map.rs b/src/database/map.rs index afd74080..40a810bc 100644 --- a/src/database/map.rs +++ b/src/database/map.rs @@ -5,7 +5,7 @@ use rocksdb::{ AsColumnFamilyRef, ColumnFamily, Direction, IteratorMode, ReadOptions, WriteBatchWithTransaction, WriteOptions, }; -use crate::{or_else, result, watchers::Watchers, Engine, Iter}; +use crate::{or_else, result, watchers::Watchers, Engine, Handle, Iter}; pub struct Map { name: String, @@ -32,11 +32,11 @@ impl Map { })) } - pub fn get(&self, key: &[u8]) -> Result>> { + pub fn get(&self, key: &[u8]) -> Result>> { let read_options = &self.read_options; - let res = self.db.db.get_cf_opt(&self.cf(), key, read_options); + let res = self.db.db.get_pinned_cf_opt(&self.cf(), key, read_options); - result(res) + Ok(result(res)?.map(Handle::from)) } pub fn multi_get(&self, keys: &[&[u8]]) -> Result>>> { diff --git a/src/database/mod.rs b/src/database/mod.rs index 74e598b6..39566896 100644 --- a/src/database/mod.rs +++ b/src/database/mod.rs @@ -1,6 +1,7 @@ mod cork; mod database; mod engine; +mod handle; mod iter; mod map; pub mod maps; @@ -13,6 +14,7 @@ extern crate rust_rocksdb as rocksdb; pub use database::Database; pub(crate) use engine::Engine; +pub use handle::Handle; pub use iter::Iter; pub use map::Map; pub(crate) use util::{or_else, result}; diff --git a/src/service/globals/data.rs b/src/service/globals/data.rs index ccd50ee9..59df5e7f 100644 --- a/src/service/globals/data.rs +++ b/src/service/globals/data.rs @@ -226,7 +226,7 @@ lasttimelinecount_cache: {lasttimelinecount_cache} / {max_lasttimelinecount_cach self.global.insert(b"keypair", &keypair)?; Ok::<_, Error>(keypair) }, - Ok, + |val| Ok(val.to_vec()), )?; let mut parts = keypair_bytes.splitn(2, |&b| b == 0xFF); diff --git a/src/service/globals/migrations.rs b/src/service/globals/migrations.rs index c05e8df8..4109017d 100644 --- a/src/service/globals/migrations.rs +++ b/src/service/globals/migrations.rs @@ -479,7 +479,7 @@ async fn db_lt_8(db: &Arc, _config: &Config) -> Result<()> { .unwrap() .expect("shortroomid should exist"); - let mut new_key = short_room_id; + let mut new_key = short_room_id.to_vec(); new_key.extend_from_slice(count); Some((new_key, v)) @@ -500,7 +500,7 @@ async fn db_lt_8(db: &Arc, _config: &Config) -> Result<()> { .unwrap() .expect("shortroomid should exist"); - let mut new_value = short_room_id; + let mut new_value = short_room_id.to_vec(); new_value.extend_from_slice(count); Some((k, new_value)) @@ -534,7 +534,7 @@ async fn db_lt_9(db: &Arc, _config: &Config) -> Result<()> { .get(room_id) .unwrap() .expect("shortroomid should exist"); - let mut new_key = short_room_id; + let mut new_key = short_room_id.to_vec(); new_key.extend_from_slice(word); new_key.push(0xFF); new_key.extend_from_slice(pdu_id_count); diff --git a/src/service/rooms/alias/data.rs b/src/service/rooms/alias/data.rs index 04f5ca7f..eb99e543 100644 --- a/src/service/rooms/alias/data.rs +++ b/src/service/rooms/alias/data.rs @@ -39,7 +39,7 @@ impl Data { pub(super) fn remove_alias(&self, alias: &RoomAliasId) -> Result<()> { if let Some(room_id) = self.alias_roomid.get(alias.alias().as_bytes())? { - let mut prefix = room_id; + let mut prefix = room_id.to_vec(); prefix.push(0xFF); for (key, _) in self.aliasid_alias.scan_prefix(prefix) { diff --git a/src/service/rooms/event_handler/mod.rs b/src/service/rooms/event_handler/mod.rs index 0468ab38..ed09e2e6 100644 --- a/src/service/rooms/event_handler/mod.rs +++ b/src/service/rooms/event_handler/mod.rs @@ -83,7 +83,7 @@ impl Service { ) -> Result>> { // 1. Skip the PDU if we already have it as a timeline event if let Some(pdu_id) = services().rooms.timeline.get_pdu_id(event_id)? { - return Ok(Some(pdu_id)); + return Ok(Some(pdu_id.to_vec())); } // 1.1 Check the server is in the room @@ -449,7 +449,7 @@ impl Service { ) -> Result>> { // Skip the PDU if we already have it as a timeline event if let Ok(Some(pduid)) = services().rooms.timeline.get_pdu_id(&incoming_pdu.event_id) { - return Ok(Some(pduid)); + return Ok(Some(pduid.to_vec())); } if services() diff --git a/src/service/rooms/timeline/data.rs b/src/service/rooms/timeline/data.rs index ba74b1aa..6609bde8 100644 --- a/src/service/rooms/timeline/data.rs +++ b/src/service/rooms/timeline/data.rs @@ -93,7 +93,7 @@ impl Data { } /// Returns the pdu's id. - pub(super) fn get_pdu_id(&self, event_id: &EventId) -> Result>> { + pub(super) fn get_pdu_id(&self, event_id: &EventId) -> Result>> { self.eventid_pduid.get(event_id.as_bytes()) } diff --git a/src/service/rooms/timeline/mod.rs b/src/service/rooms/timeline/mod.rs index aba712fa..7aed2c0d 100644 --- a/src/service/rooms/timeline/mod.rs +++ b/src/service/rooms/timeline/mod.rs @@ -159,7 +159,9 @@ impl Service { } /// Returns the pdu's id. - pub fn get_pdu_id(&self, event_id: &EventId) -> Result>> { self.db.get_pdu_id(event_id) } + pub fn get_pdu_id(&self, event_id: &EventId) -> Result>> { + self.db.get_pdu_id(event_id) + } /// Returns the pdu. /// @@ -1185,7 +1187,8 @@ impl Service { // Skip the PDU if we already have it as a timeline event if let Some(pdu_id) = self.get_pdu_id(&event_id)? { - info!("We already know {event_id} at {pdu_id:?}"); + let pdu_id = pdu_id.to_vec(); + debug!("We already know {event_id} at {pdu_id:?}"); return Ok(()); } diff --git a/src/service/transaction_ids/data.rs b/src/service/transaction_ids/data.rs index 668852ca..791b46f0 100644 --- a/src/service/transaction_ids/data.rs +++ b/src/service/transaction_ids/data.rs @@ -31,7 +31,7 @@ impl Data { pub(super) fn existing_txnid( &self, user_id: &UserId, device_id: Option<&DeviceId>, txn_id: &TransactionId, - ) -> Result>> { + ) -> Result>> { let mut key = user_id.as_bytes().to_vec(); key.push(0xFF); key.extend_from_slice(device_id.map(DeviceId::as_bytes).unwrap_or_default()); diff --git a/src/service/transaction_ids/mod.rs b/src/service/transaction_ids/mod.rs index 5b94ef5b..f18bbe64 100644 --- a/src/service/transaction_ids/mod.rs +++ b/src/service/transaction_ids/mod.rs @@ -26,7 +26,7 @@ impl Service { pub fn existing_txnid( &self, user_id: &UserId, device_id: Option<&DeviceId>, txn_id: &TransactionId, - ) -> Result>> { + ) -> Result>> { self.db.existing_txnid(user_id, device_id, txn_id) } }