diff --git a/src/mito2/benches/bench_wal_encode.rs b/src/mito2/benches/bench_wal_encode.rs index 15e93c712f..62150d7752 100644 --- a/src/mito2/benches/bench_wal_encode.rs +++ b/src/mito2/benches/bench_wal_encode.rs @@ -26,7 +26,7 @@ use std::hint::black_box; use api::v1::helper::{field_column_schema, tag_column_schema, time_index_column_schema}; use api::v1::{ColumnDataType, Mutation, OpType, Row, Rows, Value, WalEntry, value}; use criterion::{Criterion, Throughput, criterion_group, criterion_main}; -use mito2::wal::encoder::WalEntryEncoder; +use mito2::wal::encoder::{WalEntryEncoder, WalEntryMask}; use prost::Message; /// Builds a `Rows` with 3 string tags + timestamp + 3 float fields, mirroring a @@ -106,9 +106,22 @@ fn bench_wal_encode(c: &mut Criterion) { group.bench_function("wal_entry_encoder", |b| { let mut encoder = WalEntryEncoder::new(); - b.iter(|| black_box(encoder.encode_to_vec(&entry))); + b.iter(|| black_box(encoder.encode_to_vec(&entry, &WalEntryMask::default()))); }); + // Fixed input and encoder, changing only the WAL selection. Mask setup stays + // outside the timed loop; payloads are never cloned by the production path. + for skipped_count in [0, 1, 2, 4] { + let mut mask = WalEntryMask::default(); + for index in 0..entry.mutations.len() { + mask.push_mutation(index, index >= skipped_count); + } + group.bench_function(format!("mask_skip_{skipped_count}_of_4"), |b| { + let mut encoder = WalEntryEncoder::new(); + b.iter(|| black_box(encoder.encode_to_vec(black_box(&entry), black_box(&mask)))); + }); + } + group.finish(); } diff --git a/src/mito2/src/engine/basic_test.rs b/src/mito2/src/engine/basic_test.rs index 0de989f4f9..04c045010b 100644 --- a/src/mito2/src/engine/basic_test.rs +++ b/src/mito2/src/engine/basic_test.rs @@ -1321,3 +1321,341 @@ async fn test_all_index_metas_list_all_types_with_format(flat_format: bool, expe assert_eq!(expect_format, debug_format); } + +#[tokio::test] +async fn test_request_skip_wal_recovery() { + // Identical workload, changing only the request-level WAL policy. + for flat_format in [false, true] { + for skip_wal in [false, true] { + let mut env = TestEnv::new().await; + let engine = env + .create_engine(MitoConfig { + default_flat_format: flat_format, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let table_dir = request.table_dir.clone(); + let schema = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + let affected = engine + .handle_request( + region_id, + RegionRequest::Put(RegionPutRequest { + rows: Rows { + schema, + rows: build_rows_for_key("a", 0, 4, 0), + }, + hint: None, + partition_expr_version: None, + skip_wal, + }), + ) + .await + .unwrap(); + assert_eq!(affected.affected_rows, 4); + let current = engine + .get_region(region_id) + .unwrap() + .version_control + .current(); + assert_eq!(current.committed_sequence, 4); + assert_eq!(current.last_entry_id, u64::from(!skip_wal)); + let stream = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + let before = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!(before.iter().map(|b| b.num_rows()).sum::(), 4); + + reopen_region(&engine, region_id, table_dir, false, HashMap::new()).await; + let stream = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + let after = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!( + after.iter().map(|b| b.num_rows()).sum::(), + if skip_wal { 0 } else { 4 } + ); + } + } +} + +#[tokio::test] +async fn test_request_skip_wal_flush_watermarks() { + for flat_format in [false, true] { + let mut env = TestEnv::new().await; + let engine = env + .create_engine(MitoConfig { + default_flat_format: flat_format, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let schema = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + // Establish a nonzero flushed baseline, skip one write, then resume WAL. + for (round, skip_wal) in [false, true, false].into_iter().enumerate() { + let region = engine.get_region(region_id).unwrap(); + let before = region.version_control.current(); + let flushed_entry_id = engine + .region_statistic(region_id) + .unwrap() + .manifest + .data_flushed_entry_id(); + let affected = engine + .handle_request( + region_id, + RegionRequest::Put(RegionPutRequest { + rows: Rows { + schema: schema.clone(), + rows: build_rows_for_key("a", 0, 4, 0), + }, + hint: None, + partition_expr_version: None, + skip_wal, + }), + ) + .await + .unwrap(); + assert_eq!(affected.affected_rows, 4); + let after_write = region.version_control.current(); + let expected_entry_id = before.last_entry_id + u64::from(!skip_wal); + assert_eq!(after_write.last_entry_id, expected_entry_id); + assert_eq!( + after_write.committed_sequence, + before.committed_sequence + 4 + ); + assert_eq!(after_write.version.flushed_entry_id, flushed_entry_id); + assert_eq!( + engine + .region_statistic(region_id) + .unwrap() + .manifest + .data_flushed_entry_id(), + flushed_entry_id + ); + assert_eq!( + after_write.version.flushed_sequence, + before.version.flushed_sequence + ); + + flush_region(&engine, region_id, None).await; + let after_flush = region.version_control.current(); + assert_eq!(after_flush.last_entry_id, expected_entry_id); + assert_eq!(after_flush.version.flushed_sequence, (round as u64 + 1) * 4); + assert_eq!( + engine + .region_statistic(region_id) + .unwrap() + .manifest + .data_flushed_entry_id(), + expected_entry_id + ); + if skip_wal { + assert_eq!(expected_entry_id, flushed_entry_id); + } else { + assert!(expected_entry_id > flushed_entry_id); + } + } + } +} + +/// Returns `(committed_sequence, last_entry_id)` from one region version snapshot. +/// Returns `None` if the region is not open. +fn region_write_watermarks(engine: &MitoEngine, region_id: RegionId) -> Option<(u64, u64)> { + let region = engine.find_region(region_id)?; + let current = region.version_control.current(); + Some((current.committed_sequence, current.last_entry_id)) +} + +#[tokio::test] +async fn test_request_skip_wal_mixed_batch_recovery() { + use crate::region_write_ctx::RegionWriteCtx; + use crate::request::OptionOutputTx; + use crate::test_util::LogStoreImpl; + use crate::wal::Wal; + + for flat_format in [false, true] { + for skip_wal in [false, true] { + for flush_before_reopen in [false, true] { + let mut env = TestEnv::new().await; + let engine = env + .create_engine(MitoConfig { + default_flat_format: flat_format, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let table_dir = request.table_dir.clone(); + let schema = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + let region = engine.get_region(region_id).unwrap(); + let LogStoreImpl::RaftEngine(store) = env.get_log_store().unwrap() else { + panic!("expected the default local WAL"); + }; + let wal = Wal::new(store); + // Assemble one real worker write context deterministically, instead + // of relying on concurrently submitted requests landing in one batch. + let mut ctx = RegionWriteCtx::new( + region_id, + ®ion.version_control, + region.provider.clone(), + None, + ); + let mut receivers = Vec::with_capacity(4); + for (index, skip) in [false, skip_wal, false, skip_wal].into_iter().enumerate() { + let (tx, rx) = tokio::sync::oneshot::channel(); + receivers.push(rx); + ctx.push_mutation( + api::v1::OpType::Put as i32, + Some(Rows { + schema: schema.clone(), + rows: build_rows_for_key("a", index * 2, index * 2 + 2, index * 2), + }), + None, + OptionOutputTx::from(tx), + None, + skip, + ); + } + let mut writer = wal.writer(); + ctx.add_wal_entry(&mut writer).unwrap(); + let response = writer.write_to_wal().await.unwrap(); + assert_eq!(response.last_entry_ids.get(®ion_id), Some(&1)); + ctx.write_memtable().await; + ctx.publish_sequence_and_entry_id(); + drop(ctx); + for rx in receivers { + assert_eq!(rx.await.unwrap().unwrap(), 2); + } + assert_eq!(region_write_watermarks(&engine, region_id), Some((8, 1))); + assert_eq!( + request_skip_wal_timestamps(&engine, region_id).await, + (0..8).map(|i| i * 1000).collect::>() + ); + + // Inspect persisted WAL mutations, not just an encoder or counter. + let mut reader = wal.wal_entry_reader(®ion.provider, region_id, None); + let entries = reader + .read(®ion.provider, 1) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(entries.len(), 1); + assert_eq!(entries[0].0, 1); + assert_eq!( + entries[0] + .1 + .mutations + .iter() + .map(|m| m.sequence) + .collect::>(), + if skip_wal { + vec![1, 5] + } else { + vec![1, 3, 5, 7] + } + ); + assert!(entries[0].1.bulk_entries.is_empty()); + drop(region); + if flush_before_reopen { + flush_region(&engine, region_id, None).await; + let current = engine + .get_region(region_id) + .unwrap() + .version_control + .current(); + assert_eq!( + ( + current.version.flushed_sequence, + current.version.flushed_entry_id + ), + (8, 1) + ); + } + + // A no-flush close discards memtables. Only WAL-backed rows recover + // unless an explicit flush has already persisted all requests. + reopen_region(&engine, region_id, table_dir, true, HashMap::new()).await; + let loses_skipped_rows = skip_wal && !flush_before_reopen; + let mut expected_timestamps = if loses_skipped_rows { + vec![0, 1000, 4000, 5000] + } else { + (0..8).map(|i| i * 1000).collect::>() + }; + assert_eq!( + request_skip_wal_timestamps(&engine, region_id).await, + expected_timestamps + ); + let recovered_sequence = if loses_skipped_rows { 6 } else { 8 }; + assert_eq!( + region_write_watermarks(&engine, region_id), + Some((recovered_sequence, 1)) + ); + + // A subsequent default request still writes WAL, even after a + // trailing skipped request or a flush with sequence/entry-id gaps. + put_rows( + &engine, + region_id, + Rows { + schema, + rows: build_rows_for_key("a", 8, 9, 8), + }, + ) + .await; + assert_eq!( + region_write_watermarks(&engine, region_id), + Some((recovered_sequence + 1, 2)) + ); + expected_timestamps.push(8000); + assert_eq!( + request_skip_wal_timestamps(&engine, region_id).await, + expected_timestamps + ); + } + } + } +} + +async fn request_skip_wal_timestamps(engine: &MitoEngine, region_id: RegionId) -> Vec { + use datatypes::arrow::array::TimestampMillisecondArray; + + let stream = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + let mut timestamps = batches + .iter() + .flat_map(|batch| { + batch + .df_record_batch() + .column(2) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect::>(); + timestamps.sort_unstable(); + timestamps +} diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index ebabbbd295..c71b270054 100644 --- a/src/mito2/src/region/opener.rs +++ b/src/mito2/src/region/opener.rs @@ -953,6 +953,7 @@ where OptionOutputTx::none(), // We should respect the sequence in WAL during replay. Some(mutation.sequence), + false, ); } diff --git a/src/mito2/src/region_write_ctx.rs b/src/mito2/src/region_write_ctx.rs index 400f7db1fc..39b16fa99a 100644 --- a/src/mito2/src/region_write_ctx.rs +++ b/src/mito2/src/region_write_ctx.rs @@ -29,6 +29,7 @@ use crate::memtable::bulk::part::BulkPart; use crate::metrics; use crate::region::version::{VersionControlData, VersionControlRef, VersionRef}; use crate::request::OptionOutputTx; +use crate::wal::encoder::WalEntryMask; use crate::wal::{EntryId, WalWriter}; /// Notifier to notify write result on drop. @@ -89,6 +90,8 @@ pub(crate) struct RegionWriteCtx { /// We keep [WalEntry] instead of mutations to avoid taking mutations /// out of the context to construct the wal entry when we write to the wal. wal_entry: WalEntry, + /// In-memory WAL selection; the entry itself retains every request in order. + wal_mask: WalEntryMask, /// Wal options of the region being written to. provider: Provider, /// Notifiers to send write results to waiters. @@ -133,6 +136,7 @@ impl RegionWriteCtx { next_sequence: committed_sequence + 1, next_entry_id: last_entry_id + 1, wal_entry: WalEntry::default(), + wal_mask: WalEntryMask::default(), provider, notifiers: Vec::new(), bulk_notifiers: vec![], @@ -153,21 +157,24 @@ impl RegionWriteCtx { write_hint: Option, tx: OptionOutputTx, sequence: Option, + skip_wal: bool, ) { if let Some(sequence) = sequence { self.next_sequence = sequence; } let num_rows = rows.as_ref().map(|rows| rows.rows.len()).unwrap_or(0); - self.wal_entry.mutations.push(Mutation { + let mutation = Mutation { op_type, sequence: self.next_sequence, rows, write_hint, - }); + }; - let notify = WriteNotify::new(tx, num_rows); - // Notifiers are 1:1 map to mutations. - self.notifiers.push(notify); + self.wal_mask + .push_mutation(self.wal_entry.mutations.len(), !skip_wal); + self.wal_entry.mutations.push(mutation); + // Notifiers are a 1:1 map to the complete, unfiltered mutation list. + self.notifiers.push(WriteNotify::new(tx, num_rows)); // Increase sequence number. self.next_sequence += num_rows as u64; @@ -185,10 +192,11 @@ impl RegionWriteCtx { &mut self, wal_writer: &mut WalWriter, ) -> Result<()> { - wal_writer.add_entry( + wal_writer.add_entry_with_mask( self.region_id, self.next_entry_id, &self.wal_entry, + &self.wal_mask, &self.provider, )?; self.next_entry_id += 1; @@ -206,7 +214,9 @@ impl RegionWriteCtx { /// Returns whether writes in this context should skip WAL. pub(crate) fn skip_wal(&self) -> bool { - self.provider == Provider::Noop || self.version.options.skip_wal + self.provider == Provider::Noop + || self.version.options.skip_wal + || !self.wal_mask.has_entries(&self.wal_entry) } /// Sets error and marks all write operations are failed. @@ -523,6 +533,7 @@ mod tests { use common_recordbatch::DfRecordBatch; use datatypes::arrow::array::{ArrayRef, TimestampMillisecondArray}; use datatypes::arrow::datatypes::{DataType, Field, Schema}; + use prost::Message; use store_api::logstore::provider::Provider; use tokio::sync::oneshot; @@ -531,6 +542,145 @@ mod tests { use crate::memtable::bulk::part::BulkPart; use crate::test_util::version_util::VersionControlBuilder; + #[test] + fn test_request_skip_wal_preserves_sequences_and_other_writes() { + // Ablate only the request flag: the workload and sequence allocation stay identical. + for skip_wal in [false, true] { + let builder = VersionControlBuilder::new(); + let region_id = builder.region_id(); + let version_control = Arc::new(builder.build()); + let mut ctx = RegionWriteCtx::new( + region_id, + &version_control, + Provider::raft_engine_provider(region_id.as_u64()), + None, + ); + for (op_type, skip) in [ + (OpType::Put, skip_wal), + (OpType::Put, false), + (OpType::Delete, false), + ] { + ctx.push_mutation( + op_type as i32, + Some(Rows { + schema: vec![], + rows: vec![api::v1::Row::default(); 2], + }), + None, + OptionOutputTx::none(), + None, + skip, + ); + } + assert!(ctx.push_bulk(OptionOutputTx::none(), new_bulk_part(), None)); + assert!(!ctx.skip_wal()); + assert_eq!(ctx.next_sequence, 9); + assert_eq!(ctx.wal_entry.bulk_entries.len(), 1); + assert_eq!(ctx.bulk_parts[0].sequence, 7); + let sequences: Vec<_> = ctx + .wal_mask + .mutations(&ctx.wal_entry.mutations) + .map(|m| m.sequence) + .collect(); + assert_eq!(sequences, if skip_wal { vec![3, 5] } else { vec![1, 3, 5] }); + // Check the actual wire bytes, not only the mask's bookkeeping. + let encoded = crate::wal::encoder::WalEntryEncoder::new() + .encode_to_vec(&ctx.wal_entry, &ctx.wal_mask); + let decoded = WalEntry::decode(encoded.as_slice()).unwrap(); + assert_eq!( + decoded + .mutations + .iter() + .map(|m| m.sequence) + .collect::>(), + sequences + ); + assert_eq!(decoded.bulk_entries, ctx.wal_entry.bulk_entries); + assert_eq!(ctx.wal_entry.mutations.len(), 3); + assert_eq!(ctx.wal_entry.mutations[0].sequence, 1); + assert_eq!( + ctx.wal_entry.mutations.last().unwrap().op_type, + OpType::Delete as i32 + ); + } + } + + #[test] + fn test_internal_delete_respects_skip_wal_flag() { + for skip_wal in [false, true] { + let builder = VersionControlBuilder::new(); + let region_id = builder.region_id(); + let version_control = Arc::new(builder.build()); + let mut ctx = RegionWriteCtx::new( + region_id, + &version_control, + Provider::raft_engine_provider(region_id.as_u64()), + None, + ); + ctx.push_mutation( + OpType::Delete as i32, + Some(Rows { + schema: vec![], + rows: vec![api::v1::Row::default(); 2], + }), + None, + OptionOutputTx::none(), + None, + skip_wal, + ); + assert_eq!(ctx.skip_wal(), skip_wal); + assert_eq!(ctx.next_sequence, 3); + assert_eq!(ctx.delete_num, 2); + assert_eq!(ctx.wal_entry.mutations.len(), 1); + let encoded = crate::wal::encoder::WalEntryEncoder::new() + .encode_to_vec(&ctx.wal_entry, &ctx.wal_mask); + let decoded = WalEntry::decode(encoded.as_slice()).unwrap(); + if skip_wal { + assert!(decoded.mutations.is_empty()); + } else { + assert_eq!(decoded, ctx.wal_entry); + } + } + } + + #[test] + fn test_all_request_skip_wal_keeps_entry_id_and_propagates_errors() { + let builder = VersionControlBuilder::new(); + let region_id = builder.region_id(); + let version_control = Arc::new(builder.build()); + let mut ctx = RegionWriteCtx::new( + region_id, + &version_control, + Provider::raft_engine_provider(region_id.as_u64()), + None, + ); + let (tx, rx) = oneshot::channel(); + ctx.push_mutation( + OpType::Put as i32, + Some(Rows { + schema: vec![], + rows: vec![api::v1::Row::default(); 2], + }), + None, + OptionOutputTx::from(tx), + None, + true, + ); + assert!(ctx.skip_wal()); + assert_eq!(ctx.wal_entry.mutations.len(), 1); + assert_eq!(ctx.wal_mask.mutations(&ctx.wal_entry.mutations).count(), 0); + assert_eq!(ctx.next_entry_id(), 1); + assert_eq!(ctx.next_sequence, 3); + ctx.set_error(Arc::new( + UnexpectedSnafu { + reason: "wal failed".to_string(), + } + .build(), + )); + drop(ctx); + assert!(rx.blocking_recv().unwrap().is_err()); + } + #[test] fn test_set_error_marks_bulk_notifiers_failed() { let builder = VersionControlBuilder::new(); diff --git a/src/mito2/src/request.rs b/src/mito2/src/request.rs index 5f354c4603..0d918e542b 100644 --- a/src/mito2/src/request.rs +++ b/src/mito2/src/request.rs @@ -145,6 +145,12 @@ impl WriteRequest { }) } + /// Sets the request-level WAL policy. + pub fn with_skip_wal(mut self, skip_wal: bool) -> Self { + self.skip_wal = skip_wal; + self + } + /// Sets the write hint. pub fn with_hint(mut self, hint: Option) -> Self { self.hint = hint; @@ -675,6 +681,7 @@ impl WorkerRequest { let mut write_request = WriteRequest::new(region_id, OpType::Put, v.rows, region_metadata.clone())? .with_hint(v.hint) + .with_skip_wal(v.skip_wal) .with_partition_expr_version(v.partition_expr_version); if write_request.primary_key_encoding() == PrimaryKeyEncoding::Dense && let Some(region_metadata) = ®ion_metadata @@ -1876,6 +1883,25 @@ mod tests { ); } + #[test] + fn test_delete_request_defaults_to_writing_wal() { + let (request, _receiver) = WorkerRequest::try_from_region_request( + RegionId::new(1, 1), + RegionRequest::Delete(store_api::region_request::RegionDeleteRequest { + rows: Rows::default(), + hint: None, + partition_expr_version: None, + }), + None, + ) + .unwrap(); + let WorkerRequest::Write(request) = request else { + panic!("expected a write request"); + }; + assert_eq!(request.request.op_type, OpType::Delete); + assert!(!request.request.skip_wal); + } + #[test] fn test_write_request_metadata() { let rows = Rows { diff --git a/src/mito2/src/wal.rs b/src/mito2/src/wal.rs index 316cbbd612..99612e5e08 100644 --- a/src/mito2/src/wal.rs +++ b/src/mito2/src/wal.rs @@ -26,7 +26,7 @@ use std::sync::Arc; use api::v1::WalEntry; use common_error::ext::BoxedError; use common_telemetry::debug; -use encoder::WalEntryEncoder; +use encoder::{WalEntryEncoder, WalEntryMask}; use entry_reader::NoopEntryReader; use futures::future::BoxFuture; use futures::stream::BoxStream; @@ -203,6 +203,24 @@ impl WalWriter { entry_id: EntryId, wal_entry: &WalEntry, provider: &Provider, + ) -> Result<()> { + self.add_entry_with_mask( + region_id, + entry_id, + wal_entry, + &WalEntryMask::default(), + provider, + ) + } + + /// Adds only the entries selected by an in-memory WAL mask. + pub fn add_entry_with_mask( + &mut self, + region_id: RegionId, + entry_id: EntryId, + wal_entry: &WalEntry, + mask: &WalEntryMask, + provider: &Provider, ) -> Result<()> { // Gets or inserts with a newly built provider. let provider = self @@ -210,7 +228,7 @@ impl WalWriter { .entry(region_id) .or_insert_with(|| provider.clone()); - let data = self.encoder.encode_to_vec(wal_entry); + let data = self.encoder.encode_to_vec(wal_entry, mask); let entry = self .store .entry(data, entry_id, region_id, provider) diff --git a/src/mito2/src/wal/encoder.rs b/src/mito2/src/wal/encoder.rs index 5f8b492950..a36cca50c5 100644 --- a/src/mito2/src/wal/encoder.rs +++ b/src/mito2/src/wal/encoder.rs @@ -55,6 +55,8 @@ //! Leaf messages are delegated to prost, so changes to them need no update here. use api::v1::{Mutation, Row, Rows, Value, WalEntry}; +use common_base::BitVec; +use itertools::Either; use prost::Message; use prost::encoding::{WireType, encode_key, encode_varint, encoded_len_varint, key_len}; @@ -76,6 +78,58 @@ fn msg_field_len(tag: u32, body_len: usize) -> usize { key_len(tag) + encoded_len_varint(body_len as u64) + body_len } +/// An in-memory selection of mutations to persist. Bulk entries are always selected. +/// +/// No bitmap is allocated until the first excluded mutation. Once allocated, +/// each bit corresponds to a mutation in the original entry: true includes it +/// in WAL, false skips it. The mask is never persisted. +#[derive(Debug, Default)] +pub struct WalEntryMask { + selected_mutations: Option, +} + +impl WalEntryMask { + /// Appends the WAL selection for the mutation at `index` in arrival order. + pub fn push_mutation(&mut self, index: usize, write_wal: bool) { + if let Some(selected) = &mut self.selected_mutations { + debug_assert_eq!(selected.len(), index); + selected.push(write_wal); + } else if !write_wal { + let mut selected = BitVec::with_capacity(index + 1); + selected.resize(index, true); + selected.push(false); + self.selected_mutations = Some(selected); + } + } + + /// Returns whether any data remains to write to WAL after applying the mask. + /// Bulk entries are always included, regardless of the mutation mask. + pub fn has_entries(&self, entry: &WalEntry) -> bool { + !entry.bulk_entries.is_empty() + || match &self.selected_mutations { + Some(selected) => { + debug_assert_eq!(selected.len(), entry.mutations.len()); + selected.any() + } + None => !entry.mutations.is_empty(), + } + } + + /// Borrows selected mutations without cloning payloads or reordering entries. + pub fn mutations<'a>( + &'a self, + mutations: &'a [Mutation], + ) -> impl Iterator { + match &self.selected_mutations { + Some(selected) => { + debug_assert_eq!(selected.len(), mutations.len()); + Either::Left(selected.iter_ones().map(|index| &mutations[index])) + } + None => Either::Right(mutations.iter()), + } + } +} + /// A reusable encoder that caches message body sizes between its size pass and /// its encode pass. #[derive(Default)] @@ -90,13 +144,13 @@ impl WalEntryEncoder { } /// Encodes `entry` to a new `Vec`, byte-for-byte identical to - /// `entry.encode_to_vec()`. - pub fn encode_to_vec(&mut self, entry: &WalEntry) -> Vec { + /// encoding the selected entries with prost. The default mask selects all entries. + pub fn encode_to_vec(&mut self, entry: &WalEntry, mask: &WalEntryMask) -> Vec { self.sizes.clear(); - let body_len = self.size_entry(entry); + let body_len = self.size_entry(entry, mask); let mut buf = Vec::with_capacity(body_len); let mut cursor = 0; - self.encode_entry(entry, &mut buf, &mut cursor); + self.encode_entry(entry, mask, &mut buf, &mut cursor); // Invariants of the two-pass design. Kept as `debug_assert` to avoid any // overhead on the hot write path; correctness is covered by the // byte-for-byte equality tests against prost. @@ -121,7 +175,7 @@ impl WalEntryEncoder { /// Returns the body length of the `WalEntry` (no length delimiter; it is /// the root). Pushes cached slots for all nested message nodes. - fn size_entry(&mut self, entry: &WalEntry) -> usize { + fn size_entry(&mut self, entry: &WalEntry, mask: &WalEntryMask) -> usize { // Exhaustive destructure (no `..`): adding a field to `WalEntry` in // greptime-proto makes this fail to compile, forcing this encoder to be // updated rather than silently dropping the new field from the WAL. @@ -130,7 +184,7 @@ impl WalEntryEncoder { bulk_entries, } = entry; let mut body = 0; - for m in mutations { + for m in mask.mutations(mutations) { let mb = self.size_mutation(m); body += msg_field_len(MUTATION_TAG, mb); } @@ -234,13 +288,19 @@ impl WalEntryEncoder { // `next_size`, consuming that child's slot in pre-order) and writes the // key + length delimiter; the callee then writes only the body. - fn encode_entry(&self, entry: &WalEntry, buf: &mut Vec, cursor: &mut usize) { + fn encode_entry( + &self, + entry: &WalEntry, + mask: &WalEntryMask, + buf: &mut Vec, + cursor: &mut usize, + ) { // Exhaustive destructure: see note in `size_entry`. let WalEntry { mutations, bulk_entries, } = entry; - for m in mutations { + for m in mask.mutations(mutations) { let mb = self.next_size(cursor); encode_key(MUTATION_TAG, WireType::LengthDelimited, buf); encode_varint(mb as u64, buf); @@ -369,7 +429,7 @@ mod tests { fn assert_byte_identical(entry: &WalEntry) { let expected = entry.encode_to_vec(); - let actual = WalEntryEncoder::new().encode_to_vec(entry); + let actual = WalEntryEncoder::new().encode_to_vec(entry, &WalEntryMask::default()); assert_eq!( expected, actual, @@ -379,6 +439,145 @@ mod tests { ); } + #[test] + fn test_mask_matches_prost_for_every_mutation_subset() { + let entry = WalEntry { + mutations: (0..4) + .map(|i| Mutation { + op_type: if i == 2 { OpType::Delete } else { OpType::Put } as i32, + sequence: 1 + i * 3, + rows: Some(sample_rows(3, i % 2 == 0)), + write_hint: Some(WriteHint { + primary_key_encoding: 1, + }), + }) + .collect(), + bulk_entries: vec![BulkWalEntry { + sequence: 13, + max_ts: 100, + min_ts: 10, + timestamp_index: 3, + body: None, + }], + }; + // Reuse one encoder across every mask to catch stale cached-size slots. + let mut encoder = WalEntryEncoder::new(); + for bits in 0..16 { + let mut mask = WalEntryMask::default(); + for index in 0..4 { + mask.push_mutation(index, bits & (1 << index) == 0); + } + let expected = WalEntry { + mutations: entry + .mutations + .iter() + .enumerate() + .filter(|(index, _)| bits & (1 << index) == 0) + .map(|(_, mutation)| mutation.clone()) + .collect(), + bulk_entries: entry.bulk_entries.clone(), + }; + let actual = encoder.encode_to_vec(&entry, &mask); + assert_eq!(actual, expected.encode_to_vec(), "mask {bits:04b}"); + assert_eq!(WalEntry::decode(actual.as_slice()).unwrap(), expected); + assert!(mask.has_entries(&entry), "bulk is never masked out"); + assert_eq!(entry.mutations.len(), 4, "mask must not consume the input"); + // Even when every ordinary mutation is skipped, bulk still consumes + // exactly one cached size slot and is encoded normally. + if bits == 15 { + assert_eq!(encoder.sizes.len(), 1); + } + } + } + + #[test] + fn test_mask_growth_across_byte_boundaries() { + let mut encoder = WalEntryEncoder::new(); + for count in [17, 65] { + let entry = WalEntry { + mutations: (0..count) + .map(|index| Mutation { + op_type: OpType::Put as i32, + sequence: index as u64 + 1, + rows: Some(sample_rows(1, false)), + write_hint: None, + }) + .collect(), + bulk_entries: vec![], + }; + let mut all_selected = WalEntryMask::default(); + for index in 0..count { + all_selected.push_mutation(index, true); + } + assert!(all_selected.selected_mutations.is_none()); + assert_eq!( + encoder.encode_to_vec(&entry, &all_selected), + entry.encode_to_vec() + ); + + for skip_last in [false, true] { + let expected_bits: Vec<_> = (0..count) + .map(|index| { + if index == count - 1 { + !skip_last + } else { + ![0, 8, 16].contains(&index) + } + }) + .collect(); + let mut mask = WalEntryMask::default(); + for (index, &write_wal) in expected_bits.iter().enumerate() { + mask.push_mutation(index, write_wal); + assert_eq!(mask.selected_mutations.as_ref().unwrap().len(), index + 1); + } + let selected = mask.selected_mutations.as_ref().unwrap(); + for (index, &expected) in expected_bits.iter().enumerate() { + assert_eq!( + selected[index], expected, + "count={count}, index={index}, skip_last={skip_last}" + ); + } + let expected = WalEntry { + mutations: entry + .mutations + .iter() + .zip(&expected_bits) + .filter(|(_, selected)| **selected) + .map(|(mutation, _)| mutation.clone()) + .collect(), + bulk_entries: vec![], + }; + assert!(mask.has_entries(&entry)); + let encoded = encoder.encode_to_vec(&entry, &mask); + assert_eq!(encoded, expected.encode_to_vec()); + assert_eq!(WalEntry::decode(encoded.as_slice()).unwrap(), expected); + } + } + } + + #[test] + fn test_all_skip_mask_does_not_size_or_allocate_output() { + let entry = WalEntry { + mutations: vec![Mutation { + op_type: OpType::Put as i32, + sequence: 1, + rows: Some(sample_rows(100, true)), + write_hint: None, + }], + bulk_entries: vec![], + }; + let mut mask = WalEntryMask::default(); + assert!(mask.selected_mutations.is_none()); + assert!(mask.has_entries(&entry)); + mask.push_mutation(0, false); + assert!(!mask.has_entries(&entry)); + let mut encoder = WalEntryEncoder::new(); + let encoded = encoder.encode_to_vec(&entry, &mask); + assert_eq!(encoded, WalEntry::default().encode_to_vec()); + assert_eq!(encoded.capacity(), 0); + assert_eq!(encoder.sizes.capacity(), 0); + } + #[test] fn test_matches_prost_basic() { let entry = WalEntry { @@ -451,7 +650,10 @@ mod tests { }], bulk_entries: vec![], }; - assert_eq!(entry.encode_to_vec(), enc.encode_to_vec(&entry)); + assert_eq!( + entry.encode_to_vec(), + enc.encode_to_vec(&entry, &WalEntryMask::default()) + ); } } diff --git a/src/mito2/src/worker/handle_write.rs b/src/mito2/src/worker/handle_write.rs index 9694075624..1a22e04a36 100644 --- a/src/mito2/src/worker/handle_write.rs +++ b/src/mito2/src/worker/handle_write.rs @@ -451,6 +451,7 @@ impl RegionWorkerLoop { sender_req.request.hint, sender_req.sender, None, + sender_req.request.skip_wal, ); } } @@ -613,14 +614,21 @@ async fn write_wal( region_ctxs: &mut HashMap, ) -> bool { let mut wal_writer = wal.writer(); + let mut has_wal_entries = false; for region_ctx in region_ctxs.values_mut() { if region_ctx.skip_wal() { continue; } if let Err(e) = region_ctx.add_wal_entry(&mut wal_writer).map_err(Arc::new) { region_ctx.set_error(e); + } else { + has_wal_entries = true; } } + // All-skipped batches should not touch the log store, even with an empty append. + if !has_wal_entries { + return true; + } match wal_writer.write_to_wal().await.map_err(Arc::new) { Ok(response) => { for (region_id, region_ctx) in region_ctxs.iter_mut() { @@ -880,6 +888,7 @@ mod tests { fn new_region_ctx( region_id: RegionId, + skip_wal: bool, ) -> (RegionWriteCtx, oneshot::Receiver>) { let version_control = Arc::new(VersionControlBuilder::new().build()); let mut ctx = RegionWriteCtx::new( @@ -908,10 +917,38 @@ mod tests { None, OptionOutputTx::from(tx), None, + skip_wal, ); (ctx, rx) } + #[tokio::test] + async fn test_request_skip_wal_does_not_append_empty_batch() { + // Only change the request flag. A failing log store demonstrates that + // the all-skipped path never invokes append_batch, including empty appends. + for skip_wal in [false, true] { + let region_id = RegionId::new(1, 1); + let wal = Wal::new(Arc::new(MockLogStore { + fail_append: true, + ..Default::default() + })); + let (ctx, rx) = new_region_ctx(region_id, skip_wal); + let version_control = ctx.version_control().clone(); + let mut contexts = HashMap::from([(region_id, ctx)]); + assert_eq!(write_wal(&wal, &mut contexts).await, skip_wal); + if skip_wal { + let ctx = contexts.get_mut(®ion_id).unwrap(); + assert_eq!(ctx.next_entry_id(), 1); + ctx.write_memtable().await; + ctx.publish_sequence_and_entry_id(); + assert_eq!(version_control.committed_sequence(), 1); + assert_eq!(version_control.current().last_entry_id, 0); + } + drop(contexts); + assert_eq!(rx.await.unwrap().is_ok(), skip_wal); + } + } + #[tokio::test] async fn test_write_wal_skips_region_failed_to_build_entry() { let failing_region = RegionId::new(1, 1); @@ -922,10 +959,10 @@ mod tests { })); let mut region_ctxs = HashMap::new(); - let (ctx, failing_rx) = new_region_ctx(failing_region); + let (ctx, failing_rx) = new_region_ctx(failing_region, false); let failing_committed_sequence = ctx.version_control().committed_sequence(); region_ctxs.insert(failing_region, ctx); - let (ctx, ok_rx) = new_region_ctx(ok_region); + let (ctx, ok_rx) = new_region_ctx(ok_region, false); let ok_committed_sequence = ctx.version_control().committed_sequence(); region_ctxs.insert(ok_region, ctx); let entry_id = region_ctxs[&ok_region].next_entry_id(); @@ -1026,7 +1063,7 @@ mod tests { })); let mut region_ctxs = HashMap::new(); - let (ctx, rx) = new_region_ctx(failing_region); + let (ctx, rx) = new_region_ctx(failing_region, false); region_ctxs.insert(failing_region, ctx); // Writing an empty batch to the WAL succeeds, the failed region must not panic @@ -1047,7 +1084,7 @@ mod tests { })); let mut region_ctxs = HashMap::new(); - let (ctx, rx) = new_region_ctx(region_id); + let (ctx, rx) = new_region_ctx(region_id, false); region_ctxs.insert(region_id, ctx); assert!(!write_wal(&wal, &mut region_ctxs).await); diff --git a/src/store-api/src/region_request.rs b/src/store-api/src/region_request.rs index d3690a6281..00025f54a4 100644 --- a/src/store-api/src/region_request.rs +++ b/src/store-api/src/region_request.rs @@ -217,9 +217,9 @@ fn make_region_puts(inserts: InsertRequests) -> Result