#![cfg(feature = "ctx-gates")]
use std::alloc::{GlobalAlloc, Layout, System};
use std::sync::atomic::{AtomicUsize, Ordering};
use armdb::format_migration::plan::{TransformBuffers, TransformOutcome, build_plan};
use armdb::{CollectionMeta, FormatAction, FormatMigration};
use armour_core::GetType;
struct CountingAlloc;
static ALLOCS: AtomicUsize = AtomicUsize::new(0);
static LIVE_BYTES: AtomicUsize = AtomicUsize::new(0);
static PEAK_BYTES: AtomicUsize = AtomicUsize::new(0);
fn add_live(bytes: usize) {
let live = LIVE_BYTES.fetch_add(bytes, Ordering::Relaxed) + bytes;
PEAK_BYTES.fetch_max(live, Ordering::Relaxed);
}
unsafe impl GlobalAlloc for CountingAlloc {
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
ALLOCS.fetch_add(1, Ordering::Relaxed);
let ptr = unsafe { System.alloc(layout) };
if !ptr.is_null() {
add_live(layout.size());
}
ptr
}
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
LIVE_BYTES.fetch_sub(layout.size(), Ordering::Relaxed);
unsafe { System.dealloc(ptr, layout) }
}
unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
ALLOCS.fetch_add(1, Ordering::Relaxed);
let ptr = unsafe { System.alloc_zeroed(layout) };
if !ptr.is_null() {
add_live(layout.size());
}
ptr
}
unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, size: usize) -> *mut u8 {
ALLOCS.fetch_add(1, Ordering::Relaxed);
let new_ptr = unsafe { System.realloc(ptr, layout, size) };
if !new_ptr.is_null() {
if size >= layout.size() {
add_live(size - layout.size());
} else {
LIVE_BYTES.fetch_sub(layout.size() - size, Ordering::Relaxed);
}
}
new_ptr
}
}
#[global_allocator]
static GLOBAL: CountingAlloc = CountingAlloc;
#[derive(GetType)]
struct V;
fn keep(_: &[u8], _: &[u8], _: &mut Vec<u8>) -> armdb::DbResult<FormatAction> {
Ok(FormatAction::Keep)
}
const KEEP: FormatMigration = FormatMigration::raw(1, keep);
impl CollectionMeta for V {
type SelfId = [u8; 8];
const NAME: &'static str = "format_migration_alloc";
const VERSION: u16 = 2;
const FORMAT_MIGRATIONS: &'static [FormatMigration] = &[KEEP];
}
#[derive(GetType)]
struct WriteV;
fn write(_: &[u8], _: &[u8], out: &mut Vec<u8>) -> armdb::DbResult<FormatAction> {
out.extend_from_slice(b"fixed migration output payload");
Ok(FormatAction::Write)
}
const WRITE: FormatMigration = FormatMigration::raw(1, write);
impl CollectionMeta for WriteV {
type SelfId = [u8; 8];
const NAME: &'static str = "format_migration_alloc_write";
const VERSION: u16 = 2;
const FORMAT_MIGRATIONS: &'static [FormatMigration] = &[WRITE];
}
#[test]
fn format_migration_alloc_keep_transform_is_steady_state_allocation_free() {
let plan = build_plan::<[u8; 8], V>(V::NAME, Some((1, V::TYPE.h())), &[])
.unwrap()
.unwrap();
let mut buffers = TransformBuffers::default();
buffers.a.reserve_exact(256);
buffers.b.reserve_exact(256);
let input = b"fixed migration input";
let before = ALLOCS.load(Ordering::Relaxed);
for _ in 0..10_000 {
let outcome = plan
.transform_from_version(1, b"key", input, &mut buffers)
.unwrap();
assert!(matches!(&outcome, TransformOutcome::Keep(_)));
assert_eq!(outcome.as_bytes(), Some(&input[..]));
}
assert_eq!(ALLOCS.load(Ordering::Relaxed) - before, 0);
}
#[test]
fn format_migration_alloc_write_has_a_fixed_warmed_allocation_count() {
let plan = build_plan::<[u8; 8], WriteV>(WriteV::NAME, Some((1, WriteV::TYPE.h())), &[])
.unwrap()
.unwrap();
let mut buffers = TransformBuffers::default();
buffers.a.reserve_exact(256);
buffers.b.reserve_exact(256);
let _ = plan
.transform_from_version(1, b"key", b"input", &mut buffers)
.unwrap();
let count_for = |rows: usize, buffers: &mut TransformBuffers| {
let before = ALLOCS.load(Ordering::Relaxed);
for _ in 0..rows {
let outcome = plan
.transform_from_version(1, b"key", b"input", buffers)
.unwrap();
assert!(matches!(outcome, TransformOutcome::Write(_)));
}
ALLOCS.load(Ordering::Relaxed) - before
};
let one = count_for(1, &mut buffers);
let many = count_for(10_000, &mut buffers);
assert_eq!(
many, one,
"warmed Write allocations must not scale with rows"
);
}
#[derive(Clone, Debug, PartialEq, rapira::Rapira, GetType)]
struct LargeV1 {
payload: Vec<u8>,
}
impl CollectionMeta for LargeV1 {
type SelfId = [u8; 8];
const NAME: &'static str = "format_migration_bounded_shard_values";
const VERSION: u16 = 1;
}
#[derive(Clone, Debug, PartialEq, rapira::Rapira, GetType)]
struct LargeV2 {
payload: Vec<u8>,
}
fn rewrite_large(_key: &[u8], old: &[u8], out: &mut Vec<u8>) -> armdb::DbResult<FormatAction> {
out.extend_from_slice(old);
Ok(FormatAction::Write)
}
const REWRITE_LARGE: FormatMigration = FormatMigration::raw(1, rewrite_large);
impl CollectionMeta for LargeV2 {
type SelfId = [u8; 8];
const NAME: &'static str = LargeV1::NAME;
const VERSION: u16 = 2;
const FORMAT_MIGRATIONS: &'static [FormatMigration] = &[REWRITE_LARGE];
}
#[test]
fn format_migration_alloc_rewrite_does_not_retain_whole_shard_values() {
use armdb::armour::Db;
use armdb::{Config, NoHook, RapiraCodec};
const ROWS: u64 = 24;
const VALUE_BYTES: usize = 256 * 1024;
const MAX_TRANSIENT_BYTES: usize = 4 * 1024 * 1024;
let dir = tempfile::tempdir().expect("tempdir");
let mut config = Config::test();
config.shard_count = 1;
config.max_file_size = 2 * 1024 * 1024;
{
let db = Db::open_test(dir.path()).expect("open source db");
let tree = db
.open_var_typed_tree::<LargeV1, RapiraCodec, _>(config.clone(), NoHook, &[])
.expect("open source tree");
let payload = vec![0x5a; VALUE_BYTES];
for key in 0..ROWS {
tree.put(
&key.to_be_bytes(),
&LargeV1 {
payload: payload.clone(),
},
)
.expect("seed large row");
}
db.close().expect("close source db");
}
let db = Db::open_test(dir.path()).expect("reopen migration db");
let baseline = LIVE_BYTES.load(Ordering::Relaxed);
PEAK_BYTES.store(baseline, Ordering::Relaxed);
db.migrate_format::<LargeV2>(&config)
.expect("migrate large rows");
let transient = PEAK_BYTES.load(Ordering::Relaxed) - baseline;
assert!(
transient < MAX_TRANSIENT_BYTES,
"migration retained {transient} transient bytes for a {}-byte shard; values must be streamed",
ROWS as usize * VALUE_BYTES
);
}