#![cfg(all(
feature = "typed-tree",
feature = "armour",
feature = "rpc",
feature = "rapira-codec",
feature = "var-collections"
))]
use std::ops::Bound;
use std::sync::Arc;
use armdb::armour::{Db, RpcHandler};
use armdb::{CollectionMeta, Config, DbError, FixedConfig, NoHook, RapiraCodec};
use armour_core::GetType;
use armour_rpc::protocol::WriteKey;
use rapira::Rapira;
use tempfile::tempdir;
use zerocopy::IntoBytes;
#[derive(Clone, Debug, PartialEq, Rapira, GetType)]
struct MapTypedItem {
value: u64,
}
impl CollectionMeta for MapTypedItem {
type SelfId = [u8; 8];
const NAME: &'static str = "rpc_map_typed";
const VERSION: u16 = 1;
}
#[derive(Clone, Debug, PartialEq, Rapira, GetType)]
struct MapVarItem {
value: u64,
}
impl CollectionMeta for MapVarItem {
type SelfId = [u8; 8];
const NAME: &'static str = "rpc_map_var";
const VERSION: u16 = 1;
}
#[derive(
Clone,
Copy,
Debug,
PartialEq,
GetType,
zerocopy::FromBytes,
zerocopy::IntoBytes,
zerocopy::Immutable,
)]
#[repr(C)]
struct MapZeroBc(u64);
impl CollectionMeta for MapZeroBc {
type SelfId = [u8; 8];
const NAME: &'static str = "rpc_map_zero_bc";
const VERSION: u16 = 1;
}
#[derive(
Clone,
Copy,
Debug,
PartialEq,
GetType,
zerocopy::FromBytes,
zerocopy::IntoBytes,
zerocopy::Immutable,
)]
#[repr(C)]
struct MapZeroFx(u64);
impl CollectionMeta for MapZeroFx {
type SelfId = [u8; 8];
const NAME: &'static str = "rpc_map_zero_fx";
const VERSION: u16 = 1;
}
#[derive(Clone, Debug, PartialEq, Rapira, GetType)]
struct TreeTypedItem {
value: u64,
}
impl CollectionMeta for TreeTypedItem {
type SelfId = [u8; 8];
const NAME: &'static str = "rpc_map_tree_parity";
const VERSION: u16 = 1;
}
fn key(id: u64) -> Vec<u8> {
id.to_be_bytes().to_vec()
}
fn get_handler(db: &Db, name: &str) -> Arc<dyn RpcHandler> {
let hash = xxhash_rust::xxh3::xxh3_64(name.as_bytes());
db.build_tree_map()
.get(&hash)
.cloned()
.unwrap_or_else(|| panic!("handler for '{name}' not registered"))
}
fn encode_typed(id: u64) -> Vec<u8> {
let codec = RapiraCodec;
let mut buf = Vec::new();
armdb::Codec::encode_to(&codec, &MapTypedItem { value: id }, &mut buf).unwrap();
buf
}
fn encode_var(id: u64) -> Vec<u8> {
let codec = RapiraCodec;
let mut buf = Vec::new();
armdb::Codec::encode_to(&codec, &MapVarItem { value: id }, &mut buf).unwrap();
buf
}
fn encode_zero(id: u64) -> Vec<u8> {
MapZeroBc(id).as_bytes().to_vec()
}
fn encode_zero_fx(id: u64) -> Vec<u8> {
MapZeroFx(id).as_bytes().to_vec()
}
fn iterable_config(reversed: bool) -> Config {
Config::balanced()
.shard_count(2)
.hints(true)
.iterable(true)
.reversed(reversed)
.build()
}
fn fixed_iterable_config(reversed: bool) -> FixedConfig {
FixedConfig {
shard_count: 2,
grow_step: 32,
iterable: true,
reversed,
..FixedConfig::test()
}
}
fn insert_rows(
handler: &dyn RpcHandler,
encode: impl Fn(u64) -> Vec<u8>,
) -> Vec<(Vec<u8>, Vec<u8>)> {
(1_u64..=5)
.map(|id| {
let value = encode(id);
handler
.write(WriteKey::Provided(key(id)), None, value.clone())
.unwrap();
(key(id), value)
})
.collect()
}
fn assert_scan_contract(
handler: &dyn RpcHandler,
rows_ascending: &[(Vec<u8>, Vec<u8>)],
reversed: bool,
) {
let expected: Vec<_> = if reversed {
rows_ascending.iter().cloned().rev().collect()
} else {
rows_ascending.to_vec()
};
assert_eq!(
handler
.range(Bound::Unbounded, Bound::Unbounded, 0)
.unwrap(),
expected,
);
assert_eq!(
handler
.range(Bound::Included(key(2)), Bound::Excluded(key(5)), 0)
.unwrap()
.into_iter()
.map(|(key, _)| key)
.collect::<Vec<_>>(),
if reversed {
vec![key(4), key(3), key(2)]
} else {
vec![key(2), key(3), key(4)]
},
);
assert_eq!(
handler
.range(Bound::Excluded(key(2)), Bound::Included(key(4)), 0)
.unwrap()
.into_iter()
.map(|(key, _)| key)
.collect::<Vec<_>>(),
if reversed {
vec![key(4), key(3)]
} else {
vec![key(3), key(4)]
},
);
assert_eq!(
handler
.range(Bound::Unbounded, Bound::Included(key(3)), 2)
.unwrap()
.len(),
2,
);
assert_eq!(
handler
.range_keys(Bound::Included(key(2)), Bound::Excluded(key(5)), 0)
.unwrap(),
if reversed {
vec![key(4), key(3), key(2)]
} else {
vec![key(2), key(3), key(4)]
},
);
assert_eq!(
handler.first().unwrap().map(|(key, _)| key),
Some(if reversed { key(5) } else { key(1) }),
);
assert_eq!(
handler.last().unwrap().map(|(key, _)| key),
Some(if reversed { key(1) } else { key(5) }),
);
}
fn assert_keyset_pagination(handler: &dyn RpcHandler, reversed: bool) {
let mut cursor: Option<Vec<u8>> = None;
let mut keys = Vec::new();
loop {
let (start, end) = match (cursor.as_ref(), reversed) {
(None, _) => (Bound::Unbounded, Bound::Unbounded),
(Some(key), false) => (Bound::Excluded(key.clone()), Bound::Unbounded),
(Some(key), true) => (Bound::Unbounded, Bound::Excluded(key.clone())),
};
let page = handler.range_keys(start, end, 2).unwrap();
if page.is_empty() {
break;
}
cursor = page.last().cloned();
keys.extend(page);
}
assert_eq!(
keys,
if reversed {
(1_u64..=5).rev().map(key).collect::<Vec<_>>()
} else {
(1_u64..=5).map(key).collect::<Vec<_>>()
},
);
}
fn open_typed_map(db: &Db, reversed: bool) -> Arc<armdb::typed_map!(MapTypedItem, RapiraCodec)> {
db.open_typed_map::<MapTypedItem, RapiraCodec, _>(iterable_config(reversed), NoHook, &[])
.expect("open typed map")
}
fn open_var_map(db: &Db, reversed: bool) -> Arc<armdb::var_typed_map!(MapVarItem, RapiraCodec)> {
db.open_var_typed_map::<MapVarItem, RapiraCodec, _>(iterable_config(reversed), NoHook, &[])
.expect("open var typed map")
}
fn open_zero_bc_map(db: &Db, reversed: bool) -> Arc<armdb::zero_map!(MapZeroBc)> {
db.open_zero_map::<MapZeroBc, { std::mem::size_of::<MapZeroBc>() }, _>(
iterable_config(reversed),
NoHook,
&[],
)
.expect("open zero map bitcask")
}
fn open_zero_fx_map(db: &Db, reversed: bool) {
db.open_zero_map_fixed::<MapZeroFx, { std::mem::size_of::<MapZeroFx>() }, _>(
fixed_iterable_config(reversed),
NoHook,
&[],
)
.expect("open zero map fixed");
}
#[test]
fn rpc_map_typed_ascending_contract() {
let dir = tempdir().expect("tmp");
let db = Db::open(dir.path()).expect("open");
let _map = open_typed_map(&db, false);
let h = get_handler(&db, MapTypedItem::NAME);
let rows = insert_rows(&*h, encode_typed);
assert_scan_contract(&*h, &rows, false);
assert_keyset_pagination(&*h, false);
}
#[test]
fn rpc_map_typed_descending_contract() {
let dir = tempdir().expect("tmp");
let db = Db::open(dir.path()).expect("open");
let _map = open_typed_map(&db, true);
let h = get_handler(&db, MapTypedItem::NAME);
let rows = insert_rows(&*h, encode_typed);
assert_scan_contract(&*h, &rows, true);
assert_keyset_pagination(&*h, true);
}
#[test]
fn rpc_map_zero_bitcask_ascending_contract() {
let dir = tempdir().expect("tmp");
let db = Db::open(dir.path()).expect("open");
let _map = open_zero_bc_map(&db, false);
let h = get_handler(&db, MapZeroBc::NAME);
let rows = insert_rows(&*h, encode_zero);
assert_scan_contract(&*h, &rows, false);
assert_keyset_pagination(&*h, false);
}
#[test]
fn rpc_map_zero_bitcask_descending_contract() {
let dir = tempdir().expect("tmp");
let db = Db::open(dir.path()).expect("open");
let _map = open_zero_bc_map(&db, true);
let h = get_handler(&db, MapZeroBc::NAME);
let rows = insert_rows(&*h, encode_zero);
assert_scan_contract(&*h, &rows, true);
assert_keyset_pagination(&*h, true);
}
#[test]
fn rpc_map_zero_fixed_ascending_contract() {
let dir = tempdir().expect("tmp");
let db = Db::open(dir.path()).expect("open");
open_zero_fx_map(&db, false);
let h = get_handler(&db, MapZeroFx::NAME);
let rows = insert_rows(&*h, encode_zero_fx);
assert_scan_contract(&*h, &rows, false);
assert_keyset_pagination(&*h, false);
}
#[test]
fn rpc_map_zero_fixed_descending_contract() {
let dir = tempdir().expect("tmp");
let db = Db::open(dir.path()).expect("open");
open_zero_fx_map(&db, true);
let h = get_handler(&db, MapZeroFx::NAME);
let rows = insert_rows(&*h, encode_zero_fx);
assert_scan_contract(&*h, &rows, true);
assert_keyset_pagination(&*h, true);
}
#[test]
fn rpc_map_var_typed_ascending_contract() {
let dir = tempdir().expect("tmp");
let db = Db::open(dir.path()).expect("open");
let _map = open_var_map(&db, false);
let h = get_handler(&db, MapVarItem::NAME);
let rows = insert_rows(&*h, encode_var);
assert_scan_contract(&*h, &rows, false);
assert_keyset_pagination(&*h, false);
}
#[test]
fn rpc_map_var_typed_descending_contract() {
let dir = tempdir().expect("tmp");
let db = Db::open(dir.path()).expect("open");
let _map = open_var_map(&db, true);
let h = get_handler(&db, MapVarItem::NAME);
let rows = insert_rows(&*h, encode_var);
assert_scan_contract(&*h, &rows, true);
assert_keyset_pagination(&*h, true);
}
fn assert_non_iterable_scan_ops(handler: &dyn RpcHandler) {
let bad_bound = Bound::Included(vec![1, 2, 3]);
assert!(matches!(handler.first(), Err(DbError::NotImplemented)));
assert!(matches!(handler.last(), Err(DbError::NotImplemented)));
assert!(matches!(
handler.range(bad_bound.clone(), Bound::Unbounded, 0),
Err(DbError::NotImplemented)
));
assert!(matches!(
handler.range_keys(bad_bound, Bound::Unbounded, 0),
Err(DbError::NotImplemented)
));
}
#[test]
fn rpc_map_non_iterable_rejects_all_scan_ops_before_bounds() {
let dir = tempdir().expect("tmp");
let db = Db::open(dir.path()).expect("open");
let _typed = db
.open_typed_map::<MapTypedItem, RapiraCodec, _>(Config::balanced().build(), NoHook, &[])
.expect("typed map");
let typed_h = get_handler(&db, MapTypedItem::NAME);
assert_non_iterable_scan_ops(&*typed_h);
typed_h
.write(WriteKey::Provided(key(1)), None, encode_typed(1))
.unwrap();
assert!(typed_h.get(&key(1)).unwrap().is_some());
assert!(typed_h.contains(&key(1)).unwrap());
assert_eq!(typed_h.count(false).unwrap(), 1);
let _var = db
.open_var_typed_map::<MapVarItem, RapiraCodec, _>(Config::balanced().build(), NoHook, &[])
.expect("var map");
let var_h = get_handler(&db, MapVarItem::NAME);
assert_non_iterable_scan_ops(&*var_h);
var_h
.write(WriteKey::Provided(key(1)), None, encode_var(1))
.unwrap();
assert!(var_h.get(&key(1)).unwrap().is_some());
assert!(var_h.contains(&key(1)).unwrap());
assert_eq!(var_h.count(false).unwrap(), 1);
let _zero_bc = db
.open_zero_map::<MapZeroBc, { std::mem::size_of::<MapZeroBc>() }, _>(
Config::balanced().build(),
NoHook,
&[],
)
.expect("zero bc map");
let zero_bc_h = get_handler(&db, MapZeroBc::NAME);
assert_non_iterable_scan_ops(&*zero_bc_h);
zero_bc_h
.write(WriteKey::Provided(key(1)), None, encode_zero(1))
.unwrap();
assert!(zero_bc_h.get(&key(1)).unwrap().is_some());
assert!(zero_bc_h.contains(&key(1)).unwrap());
assert_eq!(zero_bc_h.count(false).unwrap(), 1);
let _zero_fx = db
.open_zero_map_fixed::<MapZeroFx, { std::mem::size_of::<MapZeroFx>() }, _>(
FixedConfig::test(),
NoHook,
&[],
)
.expect("zero fx map");
let zero_fx_h = get_handler(&db, MapZeroFx::NAME);
assert_non_iterable_scan_ops(&*zero_fx_h);
zero_fx_h
.write(WriteKey::Provided(key(1)), None, encode_zero_fx(1))
.unwrap();
assert!(zero_fx_h.get(&key(1)).unwrap().is_some());
assert!(zero_fx_h.contains(&key(1)).unwrap());
assert_eq!(zero_fx_h.count(false).unwrap(), 1);
}
#[test]
fn rpc_map_empty_iterable_endpoints_are_none() {
let dir = tempdir().expect("tmp");
let db = Db::open(dir.path()).expect("open");
let _typed = open_typed_map(&db, false);
let typed_h = get_handler(&db, MapTypedItem::NAME);
assert!(matches!(typed_h.first(), Ok(None)));
assert!(matches!(typed_h.last(), Ok(None)));
let _var = open_var_map(&db, false);
let var_h = get_handler(&db, MapVarItem::NAME);
assert!(matches!(var_h.first(), Ok(None)));
assert!(matches!(var_h.last(), Ok(None)));
let _zero_bc = open_zero_bc_map(&db, false);
let zero_bc_h = get_handler(&db, MapZeroBc::NAME);
assert!(matches!(zero_bc_h.first(), Ok(None)));
assert!(matches!(zero_bc_h.last(), Ok(None)));
open_zero_fx_map(&db, false);
let zero_fx_h = get_handler(&db, MapZeroFx::NAME);
assert!(matches!(zero_fx_h.first(), Ok(None)));
assert!(matches!(zero_fx_h.last(), Ok(None)));
}
#[test]
fn rpc_map_typed_matches_typed_tree_bytes() {
for reversed in [false, true] {
let dir = tempdir().expect("tmp");
let db = Db::open(dir.path()).expect("open");
let cfg = iterable_config(reversed);
let _tree = db
.open_typed_tree::<TreeTypedItem, RapiraCodec, _>(cfg.clone(), NoHook, &[])
.expect("tree");
let _map = db
.open_typed_map::<MapTypedItem, RapiraCodec, _>(cfg, NoHook, &[])
.expect("map");
let tree_h = get_handler(&db, TreeTypedItem::NAME);
let map_h = get_handler(&db, MapTypedItem::NAME);
for id in 1_u64..=5 {
let value = encode_typed(id);
tree_h
.write(WriteKey::Provided(key(id)), None, value.clone())
.unwrap();
map_h
.write(WriteKey::Provided(key(id)), None, value)
.unwrap();
}
assert_eq!(
map_h.range(Bound::Unbounded, Bound::Unbounded, 0).unwrap(),
tree_h.range(Bound::Unbounded, Bound::Unbounded, 0).unwrap(),
);
assert_eq!(
map_h
.range_keys(Bound::Included(key(2)), Bound::Excluded(key(5)), 0)
.unwrap(),
tree_h
.range_keys(Bound::Included(key(2)), Bound::Excluded(key(5)), 0)
.unwrap(),
);
}
}
#[test]
fn rpc_map_handler_scannability_matches_collection() {
#[derive(Clone, Debug, PartialEq, Rapira, GetType)]
struct NonIterableItem {
value: u64,
}
impl CollectionMeta for NonIterableItem {
type SelfId = [u8; 8];
const NAME: &'static str = "rpc_map_non_iterable";
const VERSION: u16 = 1;
}
let dir = tempdir().expect("tmp");
let db = Db::open(dir.path()).expect("open");
let _tree = db
.open_typed_tree::<TreeTypedItem, RapiraCodec, _>(iterable_config(false), NoHook, &[])
.expect("tree");
assert!(get_handler(&db, TreeTypedItem::NAME).scannable());
let _iterable_map = open_typed_map(&db, false);
assert!(get_handler(&db, MapTypedItem::NAME).scannable());
let _non_iterable = db
.open_typed_map::<NonIterableItem, RapiraCodec, _>(Config::balanced().build(), NoHook, &[])
.expect("non-iterable map");
assert!(!get_handler(&db, NonIterableItem::NAME).scannable());
}
#[test]
fn rpc_map_invalid_bound_on_iterable_is_client_error() {
let dir = tempdir().expect("tmp");
let db = Db::open(dir.path()).expect("open");
let _map = open_typed_map(&db, false);
let h = get_handler(&db, MapTypedItem::NAME);
let bad = Bound::Included(vec![1, 2, 3]);
let err = h
.range(bad.clone(), Bound::Unbounded, 0)
.expect_err("invalid bound");
assert!(matches!(err, DbError::Client("invalid bound key bytes")));
assert_eq!(err.status_code(), 400);
let err = h
.range_keys(bad, Bound::Unbounded, 0)
.expect_err("invalid bound");
assert!(matches!(err, DbError::Client("invalid bound key bytes")));
assert_eq!(err.status_code(), 400);
}