mito2/
region_write_ctx.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::mem;
16use std::sync::Arc;
17use std::sync::atomic::{AtomicU64, Ordering};
18
19use api::v1::{BulkWalEntry, Mutation, OpType, Rows, WalEntry, WriteHint};
20use futures::stream::{FuturesUnordered, StreamExt};
21use snafu::ResultExt;
22use store_api::logstore::LogStore;
23use store_api::logstore::provider::Provider;
24use store_api::storage::{RegionId, SequenceNumber};
25
26use crate::error::{Error, Result, WriteGroupSnafu};
27use crate::memtable::KeyValues;
28use crate::memtable::bulk::part::BulkPart;
29use crate::metrics;
30use crate::region::version::{VersionControlData, VersionControlRef, VersionRef};
31use crate::request::OptionOutputTx;
32use crate::wal::{EntryId, WalWriter};
33
34/// Notifier to notify write result on drop.
35struct WriteNotify {
36    /// Error to send to the waiter.
37    err: Option<Arc<Error>>,
38    /// Sender to send write result to the waiter for this mutation.
39    sender: OptionOutputTx,
40    /// Number of rows to be written.
41    num_rows: usize,
42}
43
44impl WriteNotify {
45    /// Creates a new notify from the `sender`.
46    fn new(sender: OptionOutputTx, num_rows: usize) -> WriteNotify {
47        WriteNotify {
48            err: None,
49            sender,
50            num_rows,
51        }
52    }
53
54    /// Send result to the waiter.
55    fn notify_result(&mut self) {
56        if let Some(err) = &self.err {
57            // Try to send the error to waiters.
58            self.sender
59                .send_mut(Err(err.clone()).context(WriteGroupSnafu));
60        } else {
61            // Send success result.
62            self.sender.send_mut(Ok(self.num_rows));
63        }
64    }
65}
66
67impl Drop for WriteNotify {
68    fn drop(&mut self) {
69        self.notify_result();
70    }
71}
72
73/// Context to keep region metadata and buffer write requests.
74pub(crate) struct RegionWriteCtx {
75    /// Id of region to write.
76    region_id: RegionId,
77    /// Version of the region while creating the context.
78    version: VersionRef,
79    /// VersionControl of the region.
80    version_control: VersionControlRef,
81    /// Next sequence number to write.
82    ///
83    /// The context assigns a unique sequence number for each row.
84    next_sequence: SequenceNumber,
85    /// Next entry id of WAL to write.
86    next_entry_id: EntryId,
87    /// Valid WAL entry to write.
88    ///
89    /// We keep [WalEntry] instead of mutations to avoid taking mutations
90    /// out of the context to construct the wal entry when we write to the wal.
91    wal_entry: WalEntry,
92    /// Wal options of the region being written to.
93    provider: Provider,
94    /// Notifiers to send write results to waiters.
95    ///
96    /// The i-th notify is for i-th mutation.
97    notifiers: Vec<WriteNotify>,
98    /// Notifiers for bulk requests.
99    bulk_notifiers: Vec<WriteNotify>,
100    /// Pending bulk write requests
101    pub(crate) bulk_parts: Vec<BulkPart>,
102    /// The write operation is failed and we should not write to the mutable memtable.
103    failed: bool,
104
105    // Metrics:
106    /// Rows to put.
107    pub(crate) put_num: usize,
108    /// Rows to delete.
109    pub(crate) delete_num: usize,
110    /// The total bytes written to the region.
111    pub(crate) written_bytes: Option<Arc<AtomicU64>>,
112}
113
114impl RegionWriteCtx {
115    /// Returns an empty context.
116    pub(crate) fn new(
117        region_id: RegionId,
118        version_control: &VersionControlRef,
119        provider: Provider,
120        written_bytes: Option<Arc<AtomicU64>>,
121    ) -> RegionWriteCtx {
122        let VersionControlData {
123            version,
124            committed_sequence,
125            last_entry_id,
126            ..
127        } = version_control.current();
128
129        RegionWriteCtx {
130            region_id,
131            version,
132            version_control: version_control.clone(),
133            next_sequence: committed_sequence + 1,
134            next_entry_id: last_entry_id + 1,
135            wal_entry: WalEntry::default(),
136            provider,
137            notifiers: Vec::new(),
138            bulk_notifiers: vec![],
139            failed: false,
140            put_num: 0,
141            delete_num: 0,
142            bulk_parts: vec![],
143            written_bytes,
144        }
145    }
146
147    /// Push mutation to the context.
148    /// This method adopts the sequence number in parameters if present.
149    pub(crate) fn push_mutation(
150        &mut self,
151        op_type: i32,
152        rows: Option<Rows>,
153        write_hint: Option<WriteHint>,
154        tx: OptionOutputTx,
155        sequence: Option<SequenceNumber>,
156    ) {
157        if let Some(sequence) = sequence {
158            self.next_sequence = sequence;
159        }
160        let num_rows = rows.as_ref().map(|rows| rows.rows.len()).unwrap_or(0);
161        self.wal_entry.mutations.push(Mutation {
162            op_type,
163            sequence: self.next_sequence,
164            rows,
165            write_hint,
166        });
167
168        let notify = WriteNotify::new(tx, num_rows);
169        // Notifiers are 1:1 map to mutations.
170        self.notifiers.push(notify);
171
172        // Increase sequence number.
173        self.next_sequence += num_rows as u64;
174
175        // Update metrics.
176        match OpType::try_from(op_type) {
177            Ok(OpType::Delete) => self.delete_num += num_rows,
178            Ok(OpType::Put) => self.put_num += num_rows,
179            Err(_) => (),
180        }
181    }
182
183    /// Encode and add WAL entry to the writer.
184    pub(crate) fn add_wal_entry<S: LogStore>(
185        &mut self,
186        wal_writer: &mut WalWriter<S>,
187    ) -> Result<()> {
188        wal_writer.add_entry(
189            self.region_id,
190            self.next_entry_id,
191            &self.wal_entry,
192            &self.provider,
193        )?;
194        self.next_entry_id += 1;
195        Ok(())
196    }
197
198    pub(crate) fn version(&self) -> &VersionRef {
199        &self.version
200    }
201
202    #[cfg(test)]
203    pub(crate) fn version_control(&self) -> &VersionControlRef {
204        &self.version_control
205    }
206
207    /// Returns whether writes in this context should skip WAL.
208    pub(crate) fn skip_wal(&self) -> bool {
209        self.provider == Provider::Noop || self.version.options.skip_wal
210    }
211
212    /// Sets error and marks all write operations are failed.
213    pub(crate) fn set_error(&mut self, err: Arc<Error>) {
214        // Set error for all notifiers.
215        for notify in &mut self.notifiers {
216            notify.err = Some(err.clone());
217        }
218        for notify in &mut self.bulk_notifiers {
219            notify.err = Some(err.clone());
220        }
221
222        // Fail the whole write operation.
223        self.failed = true;
224    }
225
226    /// Returns whether the write operation is already marked as failed.
227    pub(crate) fn is_failed(&self) -> bool {
228        self.failed
229    }
230
231    /// Updates next entry id.
232    pub(crate) fn set_next_entry_id(&mut self, next_entry_id: EntryId) {
233        self.next_entry_id = next_entry_id
234    }
235
236    /// Returns the next entry id to write.
237    #[cfg(test)]
238    pub(crate) fn next_entry_id(&self) -> EntryId {
239        self.next_entry_id
240    }
241
242    /// Consumes mutations and writes them into mutable memtable.
243    pub(crate) async fn write_memtable(&mut self) {
244        debug_assert_eq!(self.notifiers.len(), self.wal_entry.mutations.len());
245
246        if self.failed {
247            return;
248        }
249
250        let mutable_memtable = self.version.memtables.mutable.clone();
251        let prev_memory_usage = if self.written_bytes.is_some() {
252            Some(mutable_memtable.memory_usage())
253        } else {
254            None
255        };
256
257        let mutations = mem::take(&mut self.wal_entry.mutations)
258            .into_iter()
259            .enumerate()
260            .filter_map(|(i, mutation)| {
261                let kvs = KeyValues::new(&self.version.metadata, mutation)?;
262                Some((i, kvs))
263            })
264            .collect::<Vec<_>>();
265
266        if mutations.len() == 1 {
267            if let Err(err) = mutable_memtable.write(&mutations[0].1) {
268                self.notifiers[mutations[0].0].err = Some(Arc::new(err));
269            }
270        } else {
271            let mut tasks = FuturesUnordered::new();
272            for (i, kvs) in mutations {
273                let mutable = mutable_memtable.clone();
274                // use tokio runtime to schedule tasks.
275                tasks.push(common_runtime::spawn_blocking_global(move || {
276                    (i, mutable.write(&kvs))
277                }));
278            }
279
280            while let Some(result) = tasks.next().await {
281                // first unwrap the result from `spawn` above
282                let (i, result) = result.unwrap();
283                if let Err(err) = result {
284                    self.notifiers[i].err = Some(Arc::new(err));
285                }
286            }
287        }
288
289        if let Some(written_bytes) = &self.written_bytes {
290            let new_memory_usage = mutable_memtable.memory_usage();
291            let bytes = new_memory_usage.saturating_sub(prev_memory_usage.unwrap_or_default());
292            written_bytes.fetch_add(bytes as u64, Ordering::Relaxed);
293        }
294    }
295
296    pub(crate) fn push_bulk(
297        &mut self,
298        sender: OptionOutputTx,
299        mut bulk: BulkPart,
300        sequence: Option<SequenceNumber>,
301    ) -> bool {
302        if let Some(sequence) = sequence {
303            self.next_sequence = sequence;
304        }
305        bulk.sequence = self.next_sequence;
306        let entry = match BulkWalEntry::try_from(&bulk) {
307            Ok(entry) => entry,
308            Err(e) => {
309                sender.send(Err(e));
310                return false;
311            }
312        };
313
314        self.bulk_notifiers
315            .push(WriteNotify::new(sender, bulk.num_rows()));
316
317        // Add bulk wal entry
318        self.wal_entry.bulk_entries.push(entry);
319        self.next_sequence += bulk.num_rows() as u64;
320        self.bulk_parts.push(bulk);
321        true
322    }
323
324    pub(crate) async fn write_bulk(&mut self) {
325        if self.failed || self.bulk_parts.is_empty() {
326            return;
327        }
328        #[cfg(test)]
329        test_hooks::pause_before_bulk_install(self.region_id, &self.version_control).await;
330        let _timer = metrics::REGION_WORKER_HANDLE_WRITE_ELAPSED
331            .with_label_values(&["write_bulk"])
332            .start_timer();
333
334        let mutable_memtable = &self.version.memtables.mutable;
335        let prev_memory_usage = if self.written_bytes.is_some() {
336            Some(mutable_memtable.memory_usage())
337        } else {
338            None
339        };
340
341        if self.bulk_parts.len() == 1 {
342            let part = self.bulk_parts.swap_remove(0);
343            let num_rows = part.num_rows();
344            if let Err(e) = self.version.memtables.mutable.write_bulk(part) {
345                self.bulk_notifiers[0].err = Some(Arc::new(e));
346            } else {
347                self.put_num += num_rows;
348            }
349            return;
350        }
351
352        let mut tasks = FuturesUnordered::new();
353        for (i, part) in self.bulk_parts.drain(..).enumerate() {
354            let mutable = mutable_memtable.clone();
355            tasks.push(common_runtime::spawn_blocking_global(move || {
356                let num_rows = part.num_rows();
357                (i, mutable.write_bulk(part), num_rows)
358            }));
359        }
360        while let Some(result) = tasks.next().await {
361            // first unwrap the result from `spawn` above
362            let (i, result, num_rows) = result.unwrap();
363            if let Err(err) = result {
364                self.bulk_notifiers[i].err = Some(Arc::new(err));
365            } else {
366                self.put_num += num_rows;
367            }
368        }
369
370        if let Some(written_bytes) = &self.written_bytes {
371            let new_memory_usage = mutable_memtable.memory_usage();
372            let bytes = new_memory_usage.saturating_sub(prev_memory_usage.unwrap_or_default());
373            written_bytes.fetch_add(bytes as u64, Ordering::Relaxed);
374        }
375    }
376
377    /// Publishes the assigned sequences and entry id to the region's committed
378    /// watermark. Only call after both [`write_memtable`](Self::write_memtable)
379    /// and [`write_bulk`](Self::write_bulk) have completed; a failed context
380    /// must not publish.
381    pub(crate) fn publish_sequence_and_entry_id(&self) {
382        if self.failed {
383            return;
384        }
385        self.version_control
386            .set_sequence_and_entry_id(self.next_sequence - 1, self.next_entry_id - 1);
387    }
388}
389
390/// Test-only hooks to make write ordering races deterministic.
391#[cfg(test)]
392pub(crate) mod test_hooks {
393    use std::sync::Mutex;
394    use std::sync::atomic::{AtomicU64, Ordering};
395
396    use store_api::storage::RegionId;
397    use tokio::sync::watch;
398
399    use crate::region::version::VersionControlRef;
400
401    /// Channels of an armed bulk-install barrier; dropping the senders (by
402    /// disarming) unblocks writes paused on it.
403    struct ActiveBarrier {
404        id: u64,
405        /// Only bulk writes for this region on this version control (Arc
406        /// identity) pause at the barrier.
407        target_region_id: RegionId,
408        target_version_control: VersionControlRef,
409        reached: watch::Sender<bool>,
410        release: watch::Sender<bool>,
411    }
412
413    static ACTIVE_BARRIER: Mutex<Option<ActiveBarrier>> = Mutex::new(None);
414    static NEXT_BARRIER_ID: AtomicU64 = AtomicU64::new(1);
415
416    fn lock_active_barrier() -> std::sync::MutexGuard<'static, Option<ActiveBarrier>> {
417        // Never let a poisoned mutex (e.g. a panic in another test while
418        // holding the lock) hang or break unrelated tests.
419        ACTIVE_BARRIER
420            .lock()
421            .unwrap_or_else(|poisoned| poisoned.into_inner())
422    }
423
424    /// RAII guard: releasing (or dropping) it unblocks a paused write and
425    /// disarms the barrier.
426    pub(crate) struct BulkInstallBarrier {
427        id: u64,
428        reached_rx: watch::Receiver<bool>,
429        release_tx: watch::Sender<bool>,
430        released: bool,
431    }
432
433    impl BulkInstallBarrier {
434        pub(crate) async fn wait_until_reached(&mut self) {
435            if !*self.reached_rx.borrow() {
436                let _ = self.reached_rx.wait_for(|reached| *reached).await;
437            }
438        }
439
440        pub(crate) fn release(&mut self) {
441            if self.released {
442                return;
443            }
444            self.released = true;
445            let _ = self.release_tx.send(true);
446            disarm_barrier(self.id);
447        }
448    }
449
450    impl Drop for BulkInstallBarrier {
451        fn drop(&mut self) {
452            self.release();
453        }
454    }
455
456    /// Arms the bulk-install barrier and returns the owning guard; any
457    /// previously armed barrier is replaced.
458    pub(crate) fn arm_bulk_install_barrier(
459        target_region_id: RegionId,
460        target_version_control: VersionControlRef,
461    ) -> BulkInstallBarrier {
462        let (reached_tx, reached_rx) = watch::channel(false);
463        let (release_tx, _release_rx) = watch::channel(false);
464        let id = NEXT_BARRIER_ID.fetch_add(1, Ordering::Relaxed);
465        let mut active = lock_active_barrier();
466        *active = Some(ActiveBarrier {
467            id,
468            target_region_id,
469            target_version_control,
470            reached: reached_tx,
471            release: release_tx.clone(),
472        });
473        BulkInstallBarrier {
474            id,
475            reached_rx,
476            release_tx,
477            released: false,
478        }
479    }
480
481    fn disarm_barrier(id: u64) {
482        let mut active = lock_active_barrier();
483        if active.as_ref().is_some_and(|barrier| barrier.id == id) {
484            *active = None;
485        }
486    }
487
488    /// Pauses a bulk write before installing its parts until the barrier is
489    /// released or disarmed.
490    pub(crate) async fn pause_before_bulk_install(
491        region_id: RegionId,
492        version_control: &VersionControlRef,
493    ) {
494        let (reached_tx, release_rx) = {
495            let active = lock_active_barrier();
496            match active.as_ref() {
497                Some(barrier)
498                    if barrier.target_region_id == region_id
499                        && std::sync::Arc::ptr_eq(
500                            &barrier.target_version_control,
501                            version_control,
502                        ) =>
503                {
504                    (barrier.reached.clone(), barrier.release.subscribe())
505                }
506                _ => return,
507            }
508        };
509        let _ = reached_tx.send(true);
510        let mut release_rx = release_rx;
511        if !*release_rx.borrow() {
512            // The sender is dropped when the barrier is disarmed, which makes
513            // `wait_for` return an error instead of hanging forever.
514            let _ = release_rx.wait_for(|released| *released).await;
515        }
516    }
517}
518
519#[cfg(test)]
520mod tests {
521    use std::sync::Arc;
522
523    use common_recordbatch::DfRecordBatch;
524    use datatypes::arrow::array::{ArrayRef, TimestampMillisecondArray};
525    use datatypes::arrow::datatypes::{DataType, Field, Schema};
526    use store_api::logstore::provider::Provider;
527    use tokio::sync::oneshot;
528
529    use super::*;
530    use crate::error::UnexpectedSnafu;
531    use crate::memtable::bulk::part::BulkPart;
532    use crate::test_util::version_util::VersionControlBuilder;
533
534    #[test]
535    fn test_set_error_marks_bulk_notifiers_failed() {
536        let builder = VersionControlBuilder::new();
537        let region_id = builder.region_id();
538        let version_control = Arc::new(builder.build());
539        let mut ctx =
540            RegionWriteCtx::new(region_id, &version_control, Provider::noop_provider(), None);
541        let (tx, rx) = oneshot::channel();
542
543        assert!(ctx.push_bulk(OptionOutputTx::from(tx), new_bulk_part(), None));
544        ctx.set_error(Arc::new(
545            UnexpectedSnafu {
546                reason: "wal failed".to_string(),
547            }
548            .build(),
549        ));
550        drop(ctx);
551
552        let result = rx.blocking_recv().unwrap();
553        assert!(result.is_err(), "bulk notifier should report WAL error");
554    }
555
556    fn new_bulk_part() -> BulkPart {
557        let schema = Arc::new(Schema::new(vec![Field::new(
558            "ts",
559            DataType::Timestamp(datatypes::arrow::datatypes::TimeUnit::Millisecond, None),
560            false,
561        )]));
562        let arrays = vec![Arc::new(TimestampMillisecondArray::from(vec![1, 2])) as ArrayRef];
563        let batch = DfRecordBatch::try_new(schema, arrays).unwrap();
564
565        BulkPart {
566            batch,
567            max_timestamp: 2,
568            min_timestamp: 1,
569            sequence: 0,
570            timestamp_index: 0,
571            raw_data: None,
572        }
573    }
574}
mito2/region_write_ctx.rs

mito2/
region_write_ctx.rs
