use std::{convert::AsRef, fmt::Debug, io::Write, sync::Arc}; use arrayvec::ArrayVec; use conduit::{err, implement, utils::IterStream, Err, Result}; use futures::{future, Future, FutureExt, Stream, StreamExt}; use rocksdb::DBPinnableSlice; use serde::Serialize; use crate::{ keyval::KeyBuf, ser, util::{is_incomplete, map_err, or_else}, Handle, }; type RocksdbResult<'a> = Result>, rocksdb::Error>; /// Fetch a value from the database into cache, returning a reference-handle /// asynchronously. The key is serialized into an allocated buffer to perform /// the query. #[implement(super::Map)] #[inline] pub fn qry(self: &Arc, key: &K) -> impl Future>> + Send where K: Serialize + ?Sized + Debug, { let mut buf = KeyBuf::new(); self.bqry(key, &mut buf) } /// Fetch a value from the database into cache, returning a reference-handle /// asynchronously. The key is serialized into a fixed-sized buffer to perform /// the query. The maximum size is supplied as const generic parameter. #[implement(super::Map)] #[inline] pub fn aqry(self: &Arc, key: &K) -> impl Future>> + Send where K: Serialize + ?Sized + Debug, { let mut buf = ArrayVec::::new(); self.bqry(key, &mut buf) } /// Fetch a value from the database into cache, returning a reference-handle /// asynchronously. The key is serialized into a user-supplied Writer. #[implement(super::Map)] #[tracing::instrument(skip(self, buf), level = "trace")] pub fn bqry(self: &Arc, key: &K, buf: &mut B) -> impl Future>> + Send where K: Serialize + ?Sized + Debug, B: Write + AsRef<[u8]>, { let key = ser::serialize(buf, key).expect("failed to serialize query key"); self.get(key) } #[implement(super::Map)] #[tracing::instrument(skip(self, keys), fields(%self), level = "trace")] pub fn get_batch<'a, I, K>(self: &'a Arc, keys: I) -> impl Stream>> + Send + 'a where I: Iterator + ExactSizeIterator + Debug + Send + 'a, K: AsRef<[u8]> + Debug + Send + ?Sized + Sync + 'a, { keys.stream() .map(move |key| self.get(key)) .buffered(self.db.server.config.db_pool_workers.saturating_mul(2)) } /// Fetch a value from the database into cache, returning a reference-handle /// asynchronously. The key is referenced directly to perform the query. #[implement(super::Map)] #[tracing::instrument(skip(self, key), fields(%self), level = "trace")] pub fn get(self: &Arc, key: &K) -> impl Future>> + Send where K: AsRef<[u8]> + Debug + ?Sized, { use crate::pool::{Cmd, Get}; let cached = self.get_cached(key); if matches!(cached, Err(_) | Ok(Some(_))) { return future::ready(cached.map(|res| res.expect("Option is Some"))).boxed(); } debug_assert!(matches!(cached, Ok(None)), "expected status Incomplete"); let cmd = Cmd::Get(Get { map: self.clone(), key: key.as_ref().into(), res: None, }); self.db.pool.execute(cmd).boxed() } #[implement(super::Map)] #[tracing::instrument(skip(self, keys), name = "batch_blocking", level = "trace")] pub(crate) fn get_batch_blocking<'a, I, K>(&self, keys: I) -> impl Iterator>> + Send where I: Iterator + ExactSizeIterator + Debug + Send, K: AsRef<[u8]> + Debug + Send + ?Sized + Sync + 'a, { // Optimization can be `true` if key vector is pre-sorted **by the column // comparator**. const SORTED: bool = false; let read_options = &self.read_options; self.db .db .batched_multi_get_cf_opt(&self.cf(), keys, SORTED, read_options) .into_iter() .map(into_result_handle) } /// Fetch a value from the database into cache, returning a reference-handle. /// The key is referenced directly to perform the query. This is a thread- /// blocking call. #[implement(super::Map)] #[tracing::instrument(skip(self, key), name = "blocking", level = "trace")] pub fn get_blocking(&self, key: &K) -> Result> where K: AsRef<[u8]> + ?Sized, { let res = self .db .db .get_pinned_cf_opt(&self.cf(), key, &self.read_options); into_result_handle(res) } /// Fetch a value from the cache without I/O. #[implement(super::Map)] #[tracing::instrument(skip(self, key), name = "cache", level = "trace")] pub(crate) fn get_cached(&self, key: &K) -> Result>> where K: AsRef<[u8]> + Debug + ?Sized, { let res = self .db .db .get_pinned_cf_opt(&self.cf(), key, &self.cache_read_options); match res { // cache hit; not found Ok(None) => Err!(Request(NotFound("Not found in database"))), // cache hit; value found Ok(Some(res)) => Ok(Some(Handle::from(res))), // cache miss; unknown Err(e) if is_incomplete(&e) => Ok(None), // some other error occurred Err(e) => or_else(e), } } fn into_result_handle(result: RocksdbResult<'_>) -> Result> { result .map_err(map_err)? .map(Handle::from) .ok_or(err!(Request(NotFound("Not found in database")))) }