diff --git a/Cargo.toml b/Cargo.toml index 8a6dd368..96faa868 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,6 +20,11 @@ path = "examples/append_compact_purge.rs" name = "raft-engine-fork" path = "examples/fork.rs" +[[example]] +name = "json-codec" +path = "examples/append_with_json_codec.rs" +required-features = ["serde-json"] + [[test]] name = "failpoints" path = "tests/failpoints/mod.rs" @@ -59,6 +64,8 @@ rayon = "1.5" rhai = { version = "1.7", features = ["sync"], optional = true } scopeguard = "1.1" serde = { version = "1.0", features = ["derive"] } +serde_json = { version = "1.0", optional = true } +bincode = { version = "1.3", optional = true } serde_repr = "0.1" strum = { version = "0.26.2", features = ["derive"] } thiserror = "1.0" @@ -84,10 +91,13 @@ internals = [] nightly = ["prometheus/nightly"] failpoints = ["fail/failpoints"] scripting = ["rhai"] +serde-bincode = ["dep:bincode"] +serde-json = ["dep:serde_json"] swap = ["nightly", "memmap2"] std_fs = [] nightly_group = ["nightly", "swap"] +serde_group = ["serde-bincode", "serde-json"] [patch.crates-io] raft-proto = { git = "https://github.com/tikv/raft-rs", branch = "master" } diff --git a/Makefile b/Makefile index da01ab1b..9a4ba9f1 100644 --- a/Makefile +++ b/Makefile @@ -45,18 +45,18 @@ CLIPPY_WHITELIST += -A clippy::bool_assert_comparison ## Run clippy. clippy: ifdef WITH_NIGHTLY_FEATURES - cargo ${TOOLCHAIN_ARGS} clippy --all --features nightly_group,failpoints --all-targets -- -D clippy::all ${CLIPPY_WHITELIST} + cargo ${TOOLCHAIN_ARGS} clippy --all --features nightly_group,serde_group,failpoints --all-targets -- -D clippy::all ${CLIPPY_WHITELIST} else - cargo ${TOOLCHAIN_ARGS} clippy --all --features failpoints --all-targets -- -D clippy::all ${CLIPPY_WHITELIST} + cargo ${TOOLCHAIN_ARGS} clippy --all --features serde_group,failpoints --all-targets -- -D clippy::all ${CLIPPY_WHITELIST} endif ## Run tests. test: ifdef WITH_NIGHTLY_FEATURES - cargo ${TOOLCHAIN_ARGS} test --all --features nightly_group ${EXTRA_CARGO_ARGS} -- --nocapture + cargo ${TOOLCHAIN_ARGS} test --all --features nightly_group,serde_group ${EXTRA_CARGO_ARGS} -- --nocapture cargo ${TOOLCHAIN_ARGS} test --test failpoints --features nightly_group,failpoints ${EXTRA_CARGO_ARGS} -- --test-threads 1 --nocapture else - cargo ${TOOLCHAIN_ARGS} test --all ${EXTRA_CARGO_ARGS} -- --nocapture + cargo ${TOOLCHAIN_ARGS} test --all --features serde_group ${EXTRA_CARGO_ARGS} -- --nocapture cargo ${TOOLCHAIN_ARGS} test --test failpoints --features failpoints ${EXTRA_CARGO_ARGS} -- --test-threads 1 --nocapture endif @@ -66,9 +66,9 @@ test_matrix: $(error Must run test matrix with nightly features. Please reset WITH_STABLE_TOOLCHAIN.) else test_matrix: test - cargo ${TOOLCHAIN_ARGS} test --all ${EXTRA_CARGO_ARGS} -- --nocapture + cargo ${TOOLCHAIN_ARGS} test --all --features serde_group ${EXTRA_CARGO_ARGS} -- --nocapture cargo ${TOOLCHAIN_ARGS} test --test failpoints --features failpoints ${EXTRA_CARGO_ARGS} -- --test-threads 1 --nocapture - cargo ${TOOLCHAIN_ARGS} test --all --features nightly_group,std_fs ${EXTRA_CARGO_ARGS} -- --nocapture + cargo ${TOOLCHAIN_ARGS} test --all --features nightly_group,std_fs,serde_group ${EXTRA_CARGO_ARGS} -- --nocapture cargo ${TOOLCHAIN_ARGS} test --test failpoints --features nightly_group,std_fs,failpoints ${EXTRA_CARGO_ARGS} -- --test-threads 1 --nocapture endif diff --git a/examples/append_with_json_codec.rs b/examples/append_with_json_codec.rs new file mode 100644 index 00000000..f5fa6614 --- /dev/null +++ b/examples/append_with_json_codec.rs @@ -0,0 +1,147 @@ +// Copyright (c) 2017-present, PingCAP, Inc. Licensed under Apache-2.0. + +use raft_engine::{Config, Engine, JsonCodec, LogBatch, MessageExt, ReadableSize}; +use rand::thread_rng; +use rand_distr::{Distribution, Normal}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +struct JsonEntry { + index: u64, + term: u64, + payload: String, +} + +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +struct RegionState { + last_index: u64, + last_round: u64, +} + +struct JsonEntryExt; + +impl MessageExt for JsonEntryExt { + type Entry = JsonEntry; + + fn index(e: &Self::Entry) -> u64 { + e.index + } +} + +const DATA_DIR: &str = "append_with_json_codec"; +const ROUNDS: u64 = 256; +const WRITES_PER_ROUND: u64 = 512; +const COMPACT_OFFSET: u64 = 32; + +fn main() { + env_logger::init(); + + let config = Config { + dir: DATA_DIR.to_owned(), + // Small thresholds so that a short run actually rotates files, purges + // and exercises the rewrite path. + target_file_size: ReadableSize::kb(128), + purge_threshold: ReadableSize::mb(1), + batch_compression_threshold: ReadableSize::kb(0), + ..Default::default() + }; + let engine = Engine::open(config).expect("Open raft engine"); + let recovered = engine.raft_groups(); + if recovered.is_empty() { + println!("[EXAMPLE] starting from an empty {DATA_DIR}/"); + } else { + println!( + "[EXAMPLE] recovered {} raft groups written by a previous run", + recovered.len() + ); + } + + let mut rand_regions = Normal::new(8.0, 4.0) + .unwrap() + .sample_iter(thread_rng()) + .map(|x: f64| (x as u64) % 16); + let mut rand_compacts = Normal::new(COMPACT_OFFSET as f64, 16.0) + .unwrap() + .sample_iter(thread_rng()) + .map(|x: f64| x as u64); + + let mut batch = LogBatch::with_capacity(256); + let payload = "x".repeat(1024); + + for round in 1..=ROUNDS { + for _ in 0..WRITES_PER_ROUND { + let region = rand_regions.next().unwrap(); + let mut state = engine + .get_value::(region, b"state") + .unwrap() + .unwrap_or_default(); + + state.last_index += 1; // manually update the state + state.last_round = round; + + let entry = JsonEntry { + index: state.last_index, + term: round, + payload: payload.clone(), + }; + batch + .add_entries_with::(region, &[entry]) + .unwrap(); + batch + .put_value::(region, b"state".to_vec(), &state) + .unwrap(); + engine.write(&mut batch, false).unwrap(); + + if state.last_index % COMPACT_OFFSET == 0 { + let rand_compact_offset = rand_compacts.next().unwrap(); + if state.last_index > rand_compact_offset { + let compact_to = state.last_index - rand_compact_offset; + engine.compact_to(region, compact_to); + } + } + } + + for region in engine.purge_expired_files().unwrap() { + let state = engine + .get_value::(region, b"state") + .unwrap() + .unwrap(); + let compact_to = state.last_index.saturating_sub(7); + engine.compact_to(region, compact_to); + println!("[EXAMPLE] round {round}: force compact {region} to {compact_to}"); + } + } + engine.sync().unwrap(); + + // Read everything back and check it survived the whole cycle. + let mut regions = engine.raft_groups(); + regions.sort_unstable(); + for ®ion in ®ions { + let (first, last) = match (engine.first_index(region), engine.last_index(region)) { + (Some(f), Some(l)) => (f, l), + _ => continue, + }; + let mut entries = Vec::new(); + engine + .fetch_entries_to_with::( + region, + first, + last + 1, + None, + &mut entries, + ) + .unwrap(); + assert_eq!(entries.len() as u64, last - first + 1); + for (offset, e) in entries.iter().enumerate() { + assert_eq!(e.index, first + offset as u64); + assert_eq!(e.payload, payload); + } + let state = engine + .get_value::(region, b"state") + .unwrap() + .unwrap(); + assert_eq!(state.last_index, last); + println!("[EXAMPLE] region {region}: entries [{first}, {last}] verified"); + } + println!("[EXAMPLE] done, data left in {DATA_DIR}/"); +} diff --git a/src/engine.rs b/src/engine.rs index 1f1d55fd..323673d6 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -8,7 +8,7 @@ use std::thread::{Builder as ThreadBuilder, JoinHandle}; use std::time::{Duration, Instant}; use log::{error, info}; -use protobuf::{Message, parse_from_bytes}; +use protobuf::Message; use crate::config::{Config, RecoveryMode}; use crate::consistency::ConsistencyChecker; @@ -21,6 +21,7 @@ use crate::memtable::{EntryIndex, MemTableRecoverContextFactory, MemTables}; use crate::metrics::*; use crate::pipe_log::{FileBlockHandle, FileId, LogQueue, PipeLog}; use crate::purge::{PurgeHook, PurgeManager}; +use crate::value_codec::{ProtobufCodec, ValueCodec}; use crate::write_barrier::{WriteBarrier, Writer}; use crate::{Error, GlobalStats, Result, perf_context}; @@ -257,10 +258,14 @@ where } pub fn get_message(&self, region_id: u64, key: &[u8]) -> Result> { + self.get_value::(region_id, key) + } + + pub fn get_value>(&self, region_id: u64, key: &[u8]) -> Result> { let _t = StopWatch::new(&*ENGINE_READ_MESSAGE_DURATION_HISTOGRAM); if let Some(memtable) = self.memtables.get(region_id) { if let Some(value) = memtable.read().get(key) { - return Ok(Some(parse_from_bytes(&value)?)); + return Ok(Some(C::decode(&value)?)); } } Ok(None) @@ -282,14 +287,29 @@ where start_key: Option<&[u8]>, end_key: Option<&[u8]>, reverse: bool, - mut callback: C, + callback: C, ) -> Result<()> where S: Message, C: FnMut(&[u8], S) -> bool, + { + self.scan_values::(region_id, start_key, end_key, reverse, callback) + } + + pub fn scan_values( + &self, + region_id: u64, + start_key: Option<&[u8]>, + end_key: Option<&[u8]>, + reverse: bool, + mut callback: C, + ) -> Result<()> + where + VC: ValueCodec, + C: FnMut(&[u8], S) -> bool, { self.scan_raw_messages(region_id, start_key, end_key, reverse, move |k, raw_v| { - if let Ok(v) = parse_from_bytes(raw_v) { + if let Ok(v) = VC::decode(raw_v) { callback(k, v) } else { true @@ -319,16 +339,24 @@ where Ok(()) } - pub fn get_entry( - &self, - region_id: u64, - log_idx: u64, - ) -> Result> { + pub fn get_entry(&self, region_id: u64, log_idx: u64) -> Result> + where + ProtobufCodec: ValueCodec, + { + self.get_entry_with::(region_id, log_idx) + } + + /// Reads one log entry, decoded with `C`. + pub fn get_entry_with(&self, region_id: u64, log_idx: u64) -> Result> + where + C: ValueCodec, + M: MessageExt, + { let _t = StopWatch::new(&*ENGINE_READ_ENTRY_DURATION_HISTOGRAM); if let Some(memtable) = self.memtables.get(region_id) { if let Some(idx) = memtable.read().get_entry(log_idx) { ENGINE_READ_ENTRY_COUNT_HISTOGRAM.observe(1.0); - return Ok(Some(read_entry_from_file::( + return Ok(Some(read_entry_from_file::( self.pipe_log.as_ref(), &idx, )?)); @@ -343,7 +371,10 @@ where self.purge_manager.purge_expired_files() } - /// Returns count of fetched entries. + /// Returns count of fetched entries, decoded with [`ProtobufCodec`]. + /// + /// For entries stored with another codec, use + /// [`Engine::fetch_entries_to_with`]. pub fn fetch_entries_to( &self, region_id: u64, @@ -351,7 +382,26 @@ where end: u64, max_size: Option, vec: &mut Vec, - ) -> Result { + ) -> Result + where + ProtobufCodec: ValueCodec, + { + self.fetch_entries_to_with::(region_id, begin, end, max_size, vec) + } + + /// Returns count of fetched entries, decoded with `C`. + pub fn fetch_entries_to_with( + &self, + region_id: u64, + begin: u64, + end: u64, + max_size: Option, + vec: &mut Vec, + ) -> Result + where + C: ValueCodec, + M: MessageExt, + { let _t = StopWatch::new(&*ENGINE_READ_ENTRY_DURATION_HISTOGRAM); if let Some(memtable) = self.memtables.get(region_id) { let mut ents_idx: Vec = Vec::with_capacity((end - begin) as usize); @@ -360,7 +410,7 @@ where .fetch_entries_to(begin, end, max_size, &mut ents_idx)?; for i in ents_idx.iter() { vec.push({ - match read_entry_from_file::(self.pipe_log.as_ref(), i) { + match read_entry_from_file::(self.pipe_log.as_ref(), i) { Ok(entry) => entry, Err(e) => { // The index is not found in the file, it means the entry is already @@ -373,7 +423,7 @@ where // scenario where the index could become stale while // being concurrently updated by the `rewrite` operation. if let Some(idx) = immutable.get_entry(i.index) { - read_entry_from_file::(self.pipe_log.as_ref(), &idx)? + read_entry_from_file::(self.pipe_log.as_ref(), &idx)? } else { return Err(e); } @@ -619,9 +669,10 @@ thread_local! { static BLOCK_CACHE: BlockCache = BlockCache::new(); } -pub(crate) fn read_entry_from_file(pipe_log: &P, idx: &EntryIndex) -> Result +pub(crate) fn read_entry_from_file(pipe_log: &P, idx: &EntryIndex) -> Result where - M: MessageExt, + C: ValueCodec, + M: MessageExt, P: PipeLog, { BLOCK_CACHE.with(|cache| { @@ -635,11 +686,21 @@ where )?, ); } - let e = parse_from_bytes( + let e = C::decode( &cache.block.borrow() [idx.entry_offset as usize..(idx.entry_offset + idx.entry_len) as usize], )?; - assert_eq!(M::index(&e), idx.index); + // A mismatching index almost always means the data was written with a + // different value codec (or is corrupted). Report it instead of + // panicking in the read path. + if M::index(&e) != idx.index { + return Err(Error::Corruption(format!( + "log entry index mismatch: decoded {}, expected {}. The data may have \ + been written with a different value codec", + M::index(&e), + idx.index, + ))); + } Ok(e) }) } @@ -835,6 +896,101 @@ pub(crate) mod tests { } } + #[test] + fn test_entries_with_non_protobuf_codec() { + use crate::value_codec::test_codecs::{RawCodec, RawEntry, RawExt}; + let dir = tempfile::Builder::new() + .prefix("test_entries_with_non_protobuf_codec") + .tempdir() + .unwrap(); + let cfg = Config { + dir: dir.path().to_str().unwrap().to_owned(), + target_file_size: ReadableSize::kb(4), + batch_compression_threshold: ReadableSize::kb(1), + ..Default::default() + }; + let engine = RaftLogEngine::open(cfg).unwrap(); + let entries: Vec = (1..21).map(|i| RawEntry::new(i, 512)).collect(); + for rid in 1..=3 { + let mut batch = LogBatch::default(); + batch + .add_entries_with::(rid, &entries) + .unwrap(); + batch + .put_value::(rid, b"state".to_vec(), &entries[19]) + .unwrap(); + engine.write(&mut batch, true).unwrap(); + } + let check = |engine: &RaftLogEngine| { + for rid in 1..=3 { + assert_eq!(engine.first_index(rid), Some(1)); + assert_eq!(engine.last_index(rid), Some(20)); + for e in &entries { + assert_eq!( + engine + .get_entry_with::(rid, e.index) + .unwrap() + .as_ref(), + Some(e) + ); + } + let mut fetched = Vec::new(); + let n = engine + .fetch_entries_to_with::(rid, 1, 21, None, &mut fetched) + .unwrap(); + assert_eq!(n, entries.len()); + assert_eq!(fetched, entries); + assert_eq!( + engine + .get_value::(rid, b"state") + .unwrap() + .as_ref(), + Some(&entries[19]) + ); + } + }; + check(&engine); + // ... and again after recovery. + let engine = engine.reopen(); + check(&engine); + } + #[test] + fn test_read_entry_reports_index_mismatch() { + use crate::value_codec::test_codecs::{RawCodec, RawEntry, RawExt, RawExtWrongIndex}; + let dir = tempfile::Builder::new() + .prefix("test_read_entry_reports_index_mismatch") + .tempdir() + .unwrap(); + let cfg = Config { + dir: dir.path().to_str().unwrap().to_owned(), + ..Default::default() + }; + let engine = RaftLogEngine::open(cfg).unwrap(); + let entries: Vec = (1..6).map(|i| RawEntry::new(i, 32)).collect(); + let mut batch = LogBatch::default(); + batch + .add_entries_with::(1, &entries) + .unwrap(); + engine.write(&mut batch, true).unwrap(); + match engine.get_entry_with::(1, 3) { + Err(Error::Corruption(msg)) => { + assert!(msg.contains("index mismatch"), "unexpected message: {msg}"); + assert!( + msg.contains("value codec"), + "message should hint at the codec: {msg}" + ); + } + other => panic!("expected Error::Corruption, got {other:?}"), + } + assert_eq!( + engine + .get_entry_with::(1, 3) + .unwrap() + .as_ref(), + Some(&entries[2]) + ); + } + #[test] fn test_clean_raft_group() { fn run_steps(steps: &[Option<(u64, u64)>]) { diff --git a/src/lib.rs b/src/lib.rs index 5ca303d6..281169e3 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -57,6 +57,7 @@ mod swappy_allocator; #[cfg(test)] mod test_util; mod util; +mod value_codec; mod write_barrier; pub mod env; @@ -68,6 +69,11 @@ pub use log_batch::{Command, LogBatch, MessageExt}; pub use metrics::{PerfContext, get_perf_context, set_perf_context, take_perf_context}; pub use pipe_log::Version; pub use util::ReadableSize; +#[cfg(feature = "serde-bincode")] +pub use value_codec::BincodeCodec; +#[cfg(feature = "serde-json")] +pub use value_codec::JsonCodec; +pub use value_codec::{ProtobufCodec, ValueCodec}; #[cfg(feature = "internals")] pub mod internals { diff --git a/src/log_batch.rs b/src/log_batch.rs index 3555bdfd..90f94551 100644 --- a/src/log_batch.rs +++ b/src/log_batch.rs @@ -17,6 +17,7 @@ use crate::memtable::EntryIndex; use crate::metrics::StopWatch; use crate::pipe_log::{FileBlockHandle, FileId, LogFileContext, ReactiveBytes}; use crate::util::{crc32, lz4}; +use crate::value_codec::{ProtobufCodec, ValueCodec}; use crate::{Error, Result, perf_context}; pub(crate) const LOG_BATCH_HEADER_LEN: usize = 16; @@ -34,10 +35,14 @@ const MAX_LOG_BATCH_BUFFER_CAP: usize = 8 * 1024 * 1024; // 2GiB, The maximum content length accepted by lz4 compression. const MAX_LOG_ENTRIES_SIZE_PER_BATCH: usize = i32::MAX as usize; -/// `MessageExt` trait allows for probing log index from a specific type of -/// protobuf messages. -pub trait MessageExt: Send + Sync { - type Entry: Message + Clone + PartialEq; +/// `MessageExt` describes a type of log entry: how to probe its Raft log index, +/// and -- through the codec parameter `C` -- how it is serialized onto disk. +pub trait MessageExt: Send + Sync +where + C: ValueCodec, +{ + /// The log entry type. + type Entry: Clone + PartialEq; fn index(e: &Self::Entry) -> u64; } @@ -477,7 +482,16 @@ impl LogItemBatch { } pub fn put_message(&mut self, region_id: u64, key: Vec, s: &S) -> Result<()> { - self.put(region_id, key, s.write_to_bytes()?); + self.put_value::(region_id, key, s) + } + + pub fn put_value, S>( + &mut self, + region_id: u64, + key: Vec, + v: &S, + ) -> Result<()> { + self.put(region_id, key, C::encode_to_vec(v)?); Ok(()) } @@ -643,12 +657,20 @@ impl LogBatch { Ok(()) } - /// Adds some protobuf log entries into the log batch. - pub fn add_entries( - &mut self, - region_id: u64, - entries: &[M::Entry], - ) -> Result<()> { + /// Adds log entries to the batch using [`ProtobufCodec`]. + pub fn add_entries(&mut self, region_id: u64, entries: &[M::Entry]) -> Result<()> + where + ProtobufCodec: ValueCodec, + { + self.add_entries_with::(region_id, entries) + } + + /// Adds some log entries into the log batch, encoded with `C`. + pub fn add_entries_with(&mut self, region_id: u64, entries: &[M::Entry]) -> Result<()> + where + C: ValueCodec, + M: MessageExt, + { debug_assert!(self.buf_state == BufState::Open); if entries.is_empty() { return Ok(()); @@ -663,7 +685,14 @@ impl LogBatch { })(); for e in entries { let buf_offset = self.buf.len(); - e.write_to_vec(&mut self.buf)?; + if let Err(err) = C::encode_to(e, &mut self.buf) { + // Leave the batch in a usable state before bailing out, + // otherwise every later `debug_assert!(buf_state == Open)` + // would trip. + self.buf.truncate(old_buf_len); + self.buf_state = BufState::Open; + return Err(err); + } if self.buf.len() > max_entries_size + LOG_BATCH_HEADER_LEN { self.buf.truncate(old_buf_len); self.buf_state = BufState::Open; @@ -728,13 +757,23 @@ impl LogBatch { /// Adds a protobuf key value pair into the log batch. pub fn put_message(&mut self, region_id: u64, key: Vec, s: &S) -> Result<()> { + self.put_value::(region_id, key, s) + } + + /// Adds a key value pair into the log batch, encoding the value with `C`. + pub fn put_value, S>( + &mut self, + region_id: u64, + key: Vec, + v: &S, + ) -> Result<()> { if crate::is_internal_key(&key, None) { return Err(Error::InvalidArgument(format!( "key prefix `{:?}` reserved for internal use", crate::INTERNAL_KEY_PREFIX ))); } - self.item_batch.put_message(region_id, key, s) + self.item_batch.put_value::(region_id, key, v) } /// Adds a key value pair into the log batch. @@ -1116,7 +1155,6 @@ mod tests { use super::*; use crate::pipe_log::{LogQueue, Version}; use crate::test_util::{catch_unwind_silent, generate_entries, generate_entry_indexes_opt}; - use protobuf::parse_from_bytes; use raft::eraftpb::Entry; use strum::IntoEnumIterator; @@ -1124,14 +1162,17 @@ mod tests { buf: &[u8], entry_indexes: &[EntryIndex], _encoded: bool, - ) -> Vec { + ) -> Vec + where + ProtobufCodec: ValueCodec, + { let mut entries = Vec::with_capacity(entry_indexes.len()); for ei in entry_indexes { let block = LogBatch::decode_entries_block(buf, ei.entries.unwrap(), ei.compression_type) .unwrap(); entries.push( - parse_from_bytes( + ProtobufCodec::decode( &block[ei.entry_offset as usize..(ei.entry_offset + ei.entry_len) as usize], ) .unwrap(), @@ -1579,6 +1620,103 @@ mod tests { .unwrap_err(), Error::InvalidArgument(_) )); + // The codec-generic entry point must apply the same guard. + assert!(matches!( + batch + .put_value::( + 0, + crate::make_internal_key(ATOMIC_GROUP_KEY), + &Entry::new() + ) + .unwrap_err(), + Error::InvalidArgument(_) + )); + } + + #[test] + fn test_put_value_matches_put_message_bytes() { + // `put_value::` must produce exactly the bytes the old + // `put_message` produced. + let mut e = Entry::new(); + e.set_index(3); + e.set_term(5); + e.set_data(vec![b'z'; 128].into()); + + let mut a = LogBatch::default(); + a.put_message(1, b"k".to_vec(), &e).unwrap(); + let mut b = LogBatch::default(); + b.put_value::(1, b"k".to_vec(), &e) + .unwrap(); + + let extract = |batch: &LogBatch| -> Vec { + match &batch.item_batch.items[0].content { + LogItemContent::Kv(kv) => kv.value.clone().unwrap(), + _ => unreachable!(), + } + }; + assert_eq!(extract(&a), extract(&b)); + assert_eq!(extract(&a), e.write_to_bytes().unwrap()); + } + + #[test] + fn test_add_entries_encodes_with_codec() { + // `add_entries` must lay entries out back-to-back in the shared buffer, + // exactly as the old `write_to_vec` loop did. + let entries = generate_entries(1, 6, Some(&[b'q'; 32])); + let mut batch = LogBatch::default(); + batch.add_entries::(7, &entries).unwrap(); + + let mut expected = Vec::new(); + for e in &entries { + e.write_to_vec(&mut expected).unwrap(); + } + assert_eq!(&batch.buf[LOG_BATCH_HEADER_LEN..], &expected[..]); + } + #[test] + fn test_add_entries_restores_state_on_codec_error() { + use crate::value_codec::test_codecs::{FAILING_INDEX, FailExt, FailingCodec, RawEntry}; + let mut batch = LogBatch::default(); + let first = generate_entries(1, 4, Some(&[b'p'; 24])); + batch.add_entries::(7, &first).unwrap(); + let buf_len_before = batch.buf.len(); + let items_before = batch.item_batch.items.len(); + assert!(buf_len_before > LOG_BATCH_HEADER_LEN); + let mixed = [ + RawEntry::new(1, 8), + RawEntry::new(2, 8), + RawEntry::new(FAILING_INDEX, 8), + ]; + assert!( + batch + .add_entries_with::(1, &mixed) + .is_err() + ); + assert_eq!(batch.buf_state, BufState::Open); + assert_eq!( + batch.buf.len(), + buf_len_before, + "rollback ate the wrong range" + ); + assert_eq!(batch.item_batch.items.len(), items_before); + let second = generate_entries(10, 15, Some(&[b'q'; 32])); + batch.add_entries::(8, &second).unwrap(); + let mut decoded = Vec::new(); + for item in &batch.item_batch.items { + if let LogItemContent::EntryIndexes(eis) = &item.content { + for ei in &eis.0 { + let s = LOG_BATCH_HEADER_LEN + ei.entry_offset as usize; + let e = s + ei.entry_len as usize; + decoded.push(ProtobufCodec::decode(&batch.buf[s..e]).ok()); + } + } + } + let expected: Vec> = first + .iter() + .chain(second.iter()) + .cloned() + .map(Some) + .collect(); + assert_eq!(decoded, expected, "entry offsets desynced from the buffer"); } #[test] diff --git a/src/value_codec.rs b/src/value_codec.rs new file mode 100644 index 00000000..fc10daed --- /dev/null +++ b/src/value_codec.rs @@ -0,0 +1,304 @@ +// Copyright (c) 2017-present, PingCAP, Inc. Licensed under Apache-2.0. + +//! Pluggable value codecs. +//! +//! The built-in [`ProtobufCodec`] reproduces the byte layout used by Raft +//! Engine 0.4 and earlier, so adopting this abstraction is on-disk compatible. +//! [`BincodeCodec`] and [`JsonCodec`] let entries be plain `serde` types +//! instead of protobuf messages. + +use crate::Result; + +/// Describes how values of type `T` are encoded into and decoded from the log. +/// +/// # Example +/// +/// ``` +/// use raft_engine::{Result, ValueCodec}; +/// +/// struct RawBytesCodec; +/// +/// impl ValueCodec> for RawBytesCodec { +/// fn encode_to(v: &Vec, buf: &mut Vec) -> Result<()> { +/// buf.extend_from_slice(v); +/// Ok(()) +/// } +/// +/// fn decode(bytes: &[u8]) -> Result> { +/// Ok(bytes.to_owned()) +/// } +/// } +/// ``` +pub trait ValueCodec { + /// Appends the encoded form of `v` to `buf`. + /// + /// Implementations must only append; bytes already in `buf` belong to + /// other records in the same [`LogBatch`](crate::LogBatch). + fn encode_to(v: &T, buf: &mut Vec) -> Result<()>; + + /// Decodes a value from a complete encoded byte slice. + fn decode(bytes: &[u8]) -> Result; + + /// Encodes `v` into a freshly allocated buffer. + /// + /// The default implementation starts from an empty `Vec`. Codecs that can + /// cheaply compute the encoded size should override this to pre-allocate. + #[inline] + fn encode_to_vec(v: &T) -> Result> { + let mut buf = Vec::new(); + Self::encode_to(v, &mut buf)?; + Ok(buf) + } +} + +/// The legacy codec, backed by `rust-protobuf` 2.x. +/// +/// Byte-for-byte identical to the encoding used by Raft Engine before value +/// codecs were introduced. This is what you want unless you are creating a +/// brand new data directory. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct ProtobufCodec; + +impl ValueCodec for ProtobufCodec { + #[inline] + fn encode_to(v: &T, buf: &mut Vec) -> Result<()> { + v.write_to_vec(buf)?; + Ok(()) + } + + #[inline] + fn decode(bytes: &[u8]) -> Result { + Ok(protobuf::parse_from_bytes(bytes)?) + } + + /// PERF: overrides the provided default, which would start from an empty + /// `Vec` and grow it. `Message::write_to_bytes` pre-sizes the buffer using + /// `compute_size()`, exactly reproducing the pre-codec `put_message` path. + #[inline] + fn encode_to_vec(v: &T) -> Result> { + Ok(v.write_to_bytes()?) + } +} + +/// A `serde` codec backed by [`bincode`], available under the `serde-bincode` +/// feature. +#[cfg(feature = "serde-bincode")] +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct BincodeCodec; + +#[cfg(feature = "serde-bincode")] +impl ValueCodec for BincodeCodec +where + T: serde::Serialize + serde::de::DeserializeOwned, +{ + #[inline] + fn encode_to(v: &T, buf: &mut Vec) -> Result<()> { + // NOTE: `serialize_into` appends through `io::Write`, it does not + // clobber the existing contents of `buf`. + bincode::serialize_into(&mut *buf, v).map_err(|e| box_err!("bincode encode: {}", e)) + } + + #[inline] + fn decode(bytes: &[u8]) -> Result { + bincode::deserialize(bytes).map_err(|e| box_err!("bincode decode: {}", e)) + } +} + +/// A `serde` codec backed by [`serde_json`], available under the `serde-json` +/// feature. +#[cfg(feature = "serde-json")] +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct JsonCodec; + +#[cfg(feature = "serde-json")] +impl ValueCodec for JsonCodec +where + T: serde::Serialize + serde::de::DeserializeOwned, +{ + #[inline] + fn encode_to(v: &T, buf: &mut Vec) -> Result<()> { + serde_json::to_writer(&mut *buf, v).map_err(|e| box_err!("json encode: {}", e)) + } + + #[inline] + fn decode(bytes: &[u8]) -> Result { + serde_json::from_slice(bytes).map_err(|e| box_err!("json decode: {}", e)) + } +} + +#[cfg(test)] +pub(crate) mod test_codecs { + use super::*; + use crate::log_batch::MessageExt; + #[derive(Clone, Debug, Default, PartialEq)] + pub(crate) struct RawEntry { + pub index: u64, + pub data: Vec, + } + impl RawEntry { + pub(crate) fn new(index: u64, data_len: usize) -> Self { + RawEntry { + index, + data: vec![(index % 251) as u8; data_len], + } + } + } + pub(crate) struct RawCodec; + impl ValueCodec for RawCodec { + fn encode_to(v: &RawEntry, buf: &mut Vec) -> Result<()> { + buf.extend_from_slice(&v.index.to_le_bytes()); + buf.extend_from_slice(&v.data); + Ok(()) + } + fn decode(bytes: &[u8]) -> Result { + if bytes.len() < 8 { + return Err(box_err!("raw codec: input too short")); + } + Ok(RawEntry { + index: u64::from_le_bytes(bytes[..8].try_into().unwrap()), + data: bytes[8..].to_owned(), + }) + } + } + + pub(crate) struct FailingCodec; + + impl ValueCodec for FailingCodec { + fn encode_to(v: &RawEntry, buf: &mut Vec) -> Result<()> { + if v.index == FAILING_INDEX { + buf.extend_from_slice(b"PARTIAL-GARBAGE"); + return Err(box_err!("codec blew up")); + } + RawCodec::encode_to(v, buf) + } + + fn decode(bytes: &[u8]) -> Result { + RawCodec::decode(bytes) + } + } + pub(crate) const FAILING_INDEX: u64 = u64::MAX; + + pub(crate) struct RawExt; + + impl MessageExt for RawExt { + type Entry = RawEntry; + + fn index(e: &Self::Entry) -> u64 { + e.index + } + } + + pub(crate) struct RawExtWrongIndex; + + impl MessageExt for RawExtWrongIndex { + type Entry = RawEntry; + fn index(e: &Self::Entry) -> u64 { + e.index + 1 + } + } + + pub(crate) struct FailExt; + + impl MessageExt for FailExt { + type Entry = RawEntry; + fn index(e: &Self::Entry) -> u64 { + e.index + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use protobuf::Message; + use raft::eraftpb::Entry; + + fn entry(index: u64, data_len: usize) -> Entry { + let mut e = Entry::new(); + e.set_index(index); + e.set_term(7); + e.set_data(vec![b'x'; data_len].into()); + e + } + + #[test] + fn test_protobuf_codec_matches_legacy_encoding() { + for len in [0, 1, 7, 128, 4096] { + let e = entry(len as u64 + 1, len); + let mut legacy_appended = Vec::new(); + e.write_to_vec(&mut legacy_appended).unwrap(); + let mut via_codec = Vec::new(); + ProtobufCodec::encode_to(&e, &mut via_codec).unwrap(); + assert_eq!(legacy_appended, via_codec, "len={len}"); + assert_eq!( + e.write_to_bytes().unwrap(), + ProtobufCodec::encode_to_vec(&e).unwrap(), + "len={len}" + ); + } + } + + /// `encode_to` must append, never clobber: `LogBatch` packs every entry of + /// a batch into one shared buffer. + #[test] + fn test_encode_to_appends() { + let prefix = b"already here".to_vec(); + + let mut buf = prefix.clone(); + let e = entry(1, 32); + ProtobufCodec::encode_to(&e, &mut buf).unwrap(); + assert_eq!(&buf[..prefix.len()], &prefix[..]); + let decoded: Entry = ProtobufCodec::decode(&buf[prefix.len()..]).unwrap(); + assert_eq!(decoded, e); + + #[cfg(feature = "serde-bincode")] + { + let mut buf = prefix.clone(); + let v = vec![1u32, 2, 3]; + BincodeCodec::encode_to(&v, &mut buf).unwrap(); + assert_eq!(&buf[..prefix.len()], &prefix[..]); + let decoded: Vec = BincodeCodec::decode(&buf[prefix.len()..]).unwrap(); + assert_eq!(decoded, v); + } + } + + #[test] + fn test_protobuf_codec_roundtrip() { + let e = entry(42, 256); + let bytes = ProtobufCodec::encode_to_vec(&e).unwrap(); + let decoded: Entry = ProtobufCodec::decode(&bytes).unwrap(); + assert_eq!(decoded, e); + } + + #[cfg(feature = "serde-bincode")] + #[test] + fn test_bincode_codec_roundtrip() { + use serde::{Deserialize, Serialize}; + + #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] + struct Payload { + index: u64, + name: String, + blob: Vec, + } + let v = Payload { + index: 9, + name: "hello".to_owned(), + blob: vec![3; 100], + }; + let bytes = BincodeCodec::encode_to_vec(&v).unwrap(); + let decoded: Payload = BincodeCodec::decode(&bytes).unwrap(); + assert_eq!(decoded, v); + assert!(>::decode(&[0xff, 0x01]).is_err()); + } + + #[cfg(feature = "serde-json")] + #[test] + fn test_json_codec_roundtrip() { + let v = vec!["a".to_owned(), "b".to_owned()]; + let bytes = JsonCodec::encode_to_vec(&v).unwrap(); + assert_eq!(bytes, br#"["a","b"]"#.to_vec()); + let decoded: Vec = JsonCodec::decode(&bytes).unwrap(); + assert_eq!(decoded, v); + } +}