dig_blockstore/
pipeline.rs

1//! Async batched write pipeline and canonical range streaming
2//! ([`BLK-006`](../docs/requirements/domains/block_storage/specs/BLK-006.md),
3//! [`BLK-008`](../docs/requirements/domains/block_storage/specs/BLK-008.md)).
4//!
5//! # Scope
6//!
7//! - [`PipelineJob`] — one ingress unit (block + canonical flag + ack channel).
8//! - [`BlockStore::put_pipelined`] — enqueue a block for batched write; returns a
9//!   [`oneshot::Receiver`] that resolves when the worker has committed the batch.
10//! - [`BlockStore::pipeline_write_batch_count`] — [`WriteBatch`] commit counter.
11//! - [`run_write_pipeline`] — background task that drains the mpsc channel into
12//!   one [`WriteBatch`] per flush interval.
13//! - [`flush_pipeline_batch`] — single-batch commit mirroring [`BlockStore::put_block`] semantics.
14//! - [`StreamBlocksInRange`] — readahead-backed iterator over canonical bodies
15//!   for a closed height range.
16
17use std::collections::HashSet;
18use std::sync::atomic::Ordering;
19use std::sync::Arc;
20use std::time::Duration;
21
22use chia_protocol::Bytes32;
23use dig_block::{BlockStatus, L2Block};
24use rocksdb::{ColumnFamily, Direction, IteratorMode, ReadOptions, WriteBatch};
25use tokio::sync::{mpsc, oneshot};
26
27use crate::constants::{CF_BLOCKS, CF_CANONICAL, CF_HEADERS};
28use crate::encoding::{decode_height_key, hash_key, height_key};
29use crate::error::{BlockStoreError, ERR_MUTATION_READ_ONLY};
30use crate::store::{BlockStore, BlockStoreInner};
31use crate::types::BlockRecord;
32
33/// One ingress job for [`run_write_pipeline`]: own the [`L2Block`], canonical flag, and per-block ack channel
34/// ([`BLK-008`](../docs/requirements/domains/block_storage/specs/BLK-008.md) + [`IMPLEMENTATION_ORDER.md`](../docs/requirements/IMPLEMENTATION_ORDER.md) Phase 5).
35///
36/// **Ack semantics:** `Ok(true)` means a **new** row was written to [`CF_BLOCKS`]; `Ok(false)` matches [`BlockStore::put_block`]
37/// idempotency (duplicate hash on disk or duplicate within the same batch).
38pub(crate) type PipelineJob = (
39    L2Block,
40    bool,
41    oneshot::Sender<Result<bool, BlockStoreError>>,
42);
43
44impl BlockStore {
45    /// Build [`ReadOptions`] for sequential [`CF_BLOCKS`] reads inside [`StreamBlocksInRange`] ([`BLK-006`](../docs/requirements/domains/block_storage/specs/BLK-006.md) AC §3).
46    fn blocks_stream_read_options(&self) -> ReadOptions {
47        let mut o = ReadOptions::default();
48        o.set_readahead_size(self.readahead_size);
49        o
50    }
51
52    /// Stream canonical blocks from height `start` through `end` inclusive ([`BLK-006`](../docs/requirements/domains/block_storage/specs/BLK-006.md)).
53    ///
54    /// **Phase 1 — canonical walk:** [`rocksdb::DB::iterator_cf_opt`] over [`CF_CANONICAL`] with
55    /// [`ReadOptions::set_readahead_size`] and iterate bounds
56    /// ([`KEY-002`](../docs/requirements/domains/key_encoding/specs/KEY-002_height_keys.md) big-endian order).
57    ///
58    /// **Phase 2 — lazy bodies:** The returned [`StreamBlocksInRange`] walks the captured `(height, hash)` slice and,
59    /// for each entry, serves [`ShardedBlockCache`](crate::cache::sharded::ShardedBlockCache) hits without RocksDB, or
60    /// [`rocksdb::DB::get_cf_opt`] on [`CF_BLOCKS`] with the same readahead hint (separate [`ReadOptions`]
61    /// instance so canonical and block reads each carry the configured hint).
62    ///
63    /// **Why two phases:** A live RocksDB iterator over [`CF_CANONICAL`] cannot coexist with mutable/immutable borrows
64    /// of `block_cache` / decompressors on every `Iterator::next` without self-referential structs; materializing
65    /// the height→hash list preserves **readahead on the canonical scan** while keeping the public API safe and `'static`-free.
66    ///
67    /// **Errors:** Missing [`CF_BLOCKS`] row for a canonical hash yields [`BlockStoreError::BlockNotFound`] from the stream
68    /// ([`BLK-006`](../docs/requirements/domains/block_storage/specs/BLK-006.md) AC §6). Malformed canonical keys/values map to [`BlockStoreError::Serialization`].
69    ///
70    /// **Empty / inverted range:** If `start > end`, returns an iterator that yields immediately without I/O.
71    pub fn stream_blocks_in_range(
72        &self,
73        start: u64,
74        end: u64,
75    ) -> Result<StreamBlocksInRange<'_>, BlockStoreError> {
76        let cf_blocks = self.cf(CF_BLOCKS)?;
77        if start > end {
78            return Ok(StreamBlocksInRange {
79                store: self,
80                pairs: Vec::new(),
81                idx: 0,
82                read_opts: self.blocks_stream_read_options(),
83                cf_blocks,
84            });
85        }
86        let cf_canon = self.cf(CF_CANONICAL)?;
87        let mut ro_canon = ReadOptions::default();
88        ro_canon.set_readahead_size(self.readahead_size);
89        ro_canon.set_iterate_lower_bound(height_key(start).to_vec());
90        if end < u64::MAX {
91            ro_canon.set_iterate_upper_bound(height_key(end.saturating_add(1)).to_vec());
92        }
93        let iter = self.db.iterator_cf_opt(
94            cf_canon,
95            ro_canon,
96            IteratorMode::From(height_key(start).as_slice(), Direction::Forward),
97        );
98        let mut pairs = Vec::new();
99        for item in iter {
100            let (k, v) = item?;
101            let karr: [u8; 8] = k.as_ref().try_into().map_err(|_| {
102                BlockStoreError::Serialization(
103                    "stream_blocks_in_range: CF_CANONICAL key must be exactly 8 bytes".into(),
104                )
105            })?;
106            let height = decode_height_key(&karr);
107            if height > end {
108                break;
109            }
110            if height < start {
111                continue;
112            }
113            let varr: [u8; 32] = v.as_ref().try_into().map_err(|_| {
114                BlockStoreError::Serialization(
115                    "stream_blocks_in_range: CF_CANONICAL value must be exactly 32 bytes".into(),
116                )
117            })?;
118            pairs.push((height, Bytes32::new(varr)));
119        }
120        Ok(StreamBlocksInRange {
121            store: self,
122            pairs,
123            idx: 0,
124            read_opts: self.blocks_stream_read_options(),
125            cf_blocks,
126        })
127    }
128
129    /// Async batched ingest ([`BLK-008`](../docs/requirements/domains/block_storage/specs/BLK-008.md), [`IMPLEMENTATION_ORDER.md`](../docs/requirements/IMPLEMENTATION_ORDER.md) Phase 5).
130    ///
131    /// **Channel + batching (NORMATIVE §1–3):** Enqueues into a bounded [`mpsc`] queue; a background task accumulates
132    /// up to `pipeline_batch_size` jobs or until `pipeline_flush_ms` elapses,
133    /// then applies **one** [`WriteBatch`] mirroring [`BlockStore::put_block`] semantics.
134    ///
135    /// **Per-block ack ([`IMPLEMENTATION_ORDER.md`](../docs/requirements/IMPLEMENTATION_ORDER.md)):** The returned
136    /// [`oneshot::Receiver`] resolves to the same `Result<bool, BlockStoreError>` shape as [`BlockStore::put_block`]
137    /// (`Ok(true)` inserted, `Ok(false)` duplicate).
138    ///
139    /// **Runtime contract:** The first call lazily spawns [`run_write_pipeline`] via [`tokio::spawn`]; therefore an
140    /// active [`tokio::runtime::Handle`] must exist (integration tests should use `#[tokio::test]`).
141    pub async fn put_pipelined(
142        &self,
143        block: L2Block,
144        canonical: bool,
145    ) -> Result<oneshot::Receiver<Result<bool, BlockStoreError>>, BlockStoreError> {
146        if self.read_only {
147            return Err(BlockStoreError::Serialization(
148                ERR_MUTATION_READ_ONLY.into(),
149            ));
150        }
151        let tx = self.pipeline_sender().await?;
152        let (ack_tx, ack_rx) = oneshot::channel();
153        tx.send((block, canonical, ack_tx))
154            .await
155            .map_err(|_| BlockStoreError::PipelineClosed)?;
156        Ok(ack_rx)
157    }
158
159    /// Count of successful RocksDB [`WriteBatch`] commits executed by the [`BLK-008`](../docs/requirements/domains/block_storage/specs/BLK-008.md) worker.
160    ///
161    /// **Instrumentation:** Used by `tests/blk_008_tests.rs` to prove AC §4 "single `WriteBatch` per flush interval".
162    #[must_use]
163    pub fn pipeline_write_batch_count(&self) -> u64 {
164        self.pipeline_write_batches.load(Ordering::Relaxed) as u64
165    }
166
167    /// Lazily constructs the bounded [`mpsc`] sender and spawns [`run_write_pipeline`].
168    async fn pipeline_sender(&self) -> Result<mpsc::Sender<PipelineJob>, BlockStoreError> {
169        let mut guard = self.pipeline_tx.lock().await;
170        if let Some(tx) = guard.as_ref() {
171            return Ok(tx.clone());
172        }
173        let _handle = tokio::runtime::Handle::try_current().map_err(|_| {
174            BlockStoreError::Serialization(
175                "put_pipelined requires an active Tokio runtime (use #[tokio::test] or Runtime::block_on)"
176                    .into(),
177            )
178        })?;
179        let cap = self.pipeline_channel_capacity;
180        let (tx, rx) = mpsc::channel::<PipelineJob>(cap);
181        let inner = self.inner.clone();
182        let batch = self.pipeline_batch_size;
183        let flush_ms = self.pipeline_flush_ms;
184        tokio::spawn(run_write_pipeline(
185            inner,
186            Arc::new(tokio::sync::Mutex::new(None)),
187            rx,
188            batch,
189            flush_ms,
190        ));
191        *guard = Some(tx.clone());
192        Ok(tx)
193    }
194}
195
196/// Background loop draining [`PipelineJob`] values into batched [`WriteBatch`] commits ([`BLK-008`](../docs/requirements/domains/block_storage/specs/BLK-008.md)).
197///
198/// **Shutdown (AC §8):** Ingress senders live on [`BlockStore::pipeline_tx`], not on [`BlockStoreInner`]. When the last
199/// [`BlockStore`] clone is dropped, the final [`mpsc::Sender`] is released, `rx.recv()` yields `None`, and we
200/// [`flush_pipeline_batch`] any tail buffer before exiting.
201pub(crate) async fn run_write_pipeline(
202    inner: Arc<BlockStoreInner>,
203    _worker_unused_pipeline_tx: Arc<tokio::sync::Mutex<Option<mpsc::Sender<PipelineJob>>>>,
204    mut rx: mpsc::Receiver<PipelineJob>,
205    batch_size: usize,
206    flush_ms: u64,
207) {
208    let store = BlockStore {
209        inner,
210        pipeline_tx: _worker_unused_pipeline_tx,
211    };
212    let mut buf: Vec<PipelineJob> = Vec::with_capacity(batch_size);
213    let tick = Duration::from_millis(flush_ms);
214
215    loop {
216        match rx.recv().await {
217            None => {
218                return;
219            }
220            Some(job) => buf.push(job),
221        }
222        if buf.len() >= batch_size {
223            let _ = flush_pipeline_batch(&store, &mut buf);
224            buf.clear();
225            continue;
226        }
227
228        let mut sleep = Box::pin(tokio::time::sleep(tick));
229        'collect: loop {
230            tokio::select! {
231                biased;
232                maybe = rx.recv() => {
233                    match maybe {
234                        None => {
235                            let _ = flush_pipeline_batch(&store, &mut buf);
236                            return;
237                        }
238                        Some(job) => {
239                            buf.push(job);
240                            if buf.len() >= batch_size {
241                                break 'collect;
242                            }
243                        }
244                    }
245                }
246                _ = &mut sleep, if !buf.is_empty() => {
247                    break 'collect;
248                }
249            }
250        }
251
252        let _ = flush_pipeline_batch(&store, &mut buf);
253        buf.clear();
254    }
255}
256
257/// Applies one RocksDB [`WriteBatch`] for all novel inserts in `jobs`, mirroring [`BlockStore::put_block`].
258///
259/// **Idempotency (AC §5):** Duplicate hashes already on disk **or** repeated within `jobs` are answered with
260/// `Ok(false)` acks and are omitted from the write batch. A completely duplicate batch performs **no** `db.write`.
261///
262/// **Errors:** Build/IO failures notify every still-pending staged [`oneshot`] with [`BlockStoreError::Serialization`]
263/// carrying the diagnostic string, then the function returns `Ok(())` so the worker loop keeps draining (best-effort
264/// bulk ingest semantics; callers observe failure on their ack channel).
265fn flush_pipeline_batch(
266    store: &BlockStore,
267    jobs: &mut Vec<PipelineJob>,
268) -> Result<(), BlockStoreError> {
269    if store.read_only {
270        let pending: Vec<PipelineJob> = std::mem::take(jobs);
271        let msg = ERR_MUTATION_READ_ONLY.to_string();
272        for (_, _, ack) in pending {
273            let _ = ack.send(Err(BlockStoreError::Serialization(msg.clone())));
274        }
275        return Ok(());
276    }
277    let pending: Vec<PipelineJob> = std::mem::take(jobs);
278    let cf_b = match store.cf(CF_BLOCKS) {
279        Ok(c) => c,
280        Err(e) => {
281            let msg = e.to_string();
282            for (_, _, ack) in pending {
283                let _ = ack.send(Err(BlockStoreError::Serialization(format!(
284                    "write pipeline: {msg}"
285                ))));
286            }
287            return Ok(());
288        }
289    };
290    let cf_h = match store.cf(CF_HEADERS) {
291        Ok(c) => c,
292        Err(e) => {
293            let msg = e.to_string();
294            for (_, _, ack) in pending {
295                let _ = ack.send(Err(BlockStoreError::Serialization(format!(
296                    "write pipeline: {msg}"
297                ))));
298            }
299            return Ok(());
300        }
301    };
302    let cf_c = match store.cf(CF_CANONICAL) {
303        Ok(c) => c,
304        Err(e) => {
305            let msg = e.to_string();
306            for (_, _, ack) in pending {
307                let _ = ack.send(Err(BlockStoreError::Serialization(format!(
308                    "write pipeline: {msg}"
309                ))));
310            }
311            return Ok(());
312        }
313    };
314
315    let mut seen: HashSet<Bytes32> = HashSet::new();
316    struct StagedRow {
317        hash: Bytes32,
318        block: L2Block,
319        compressed: Vec<u8>,
320        header_bytes: Vec<u8>,
321        canonical: bool,
322        ack: oneshot::Sender<Result<bool, BlockStoreError>>,
323    }
324    let mut staged: Vec<StagedRow> = Vec::new();
325
326    for (block, canonical, ack) in pending {
327        let hash = block.hash();
328        if !seen.insert(hash) {
329            let _ = ack.send(Ok(false));
330            continue;
331        }
332        let exists = match store.db.get_cf(cf_b, hash_key(&hash).as_slice()) {
333            Ok(o) => o.is_some(),
334            Err(e) => {
335                let _ = ack.send(Err(BlockStoreError::RocksDb(e)));
336                continue;
337            }
338        };
339        if exists {
340            let _ = ack.send(Ok(false));
341            continue;
342        }
343        let compressed = match store.serialize_block(&block) {
344            Ok(b) => b,
345            Err(e) => {
346                let _ = ack.send(Err(e));
347                continue;
348            }
349        };
350        let header_bytes = match BlockStore::serialize_header(&block.header) {
351            Ok(b) => b,
352            Err(e) => {
353                let _ = ack.send(Err(e));
354                continue;
355            }
356        };
357        staged.push(StagedRow {
358            hash,
359            block,
360            compressed,
361            header_bytes,
362            canonical,
363            ack,
364        });
365    }
366
367    let mut wb = WriteBatch::default();
368    for row in &staged {
369        wb.put_cf(
370            cf_b,
371            hash_key(&row.hash).as_slice(),
372            row.compressed.as_slice(),
373        );
374        wb.put_cf(
375            cf_h,
376            hash_key(&row.hash).as_slice(),
377            row.header_bytes.as_slice(),
378        );
379        if row.canonical {
380            wb.put_cf(
381                cf_c,
382                height_key(row.block.height()),
383                hash_key(&row.hash).as_slice(),
384            );
385        }
386    }
387
388    if wb.is_empty() {
389        return Ok(());
390    }
391
392    if let Err(e) = store.db.write(wb) {
393        let msg = format!("write pipeline: rocksdb write failed: {e}");
394        for row in staged {
395            let _ = row
396                .ack
397                .send(Err(BlockStoreError::Serialization(msg.clone())));
398        }
399        return Ok(());
400    }
401
402    for row in &staged {
403        if row.canonical {
404            if let Err(e) = store
405                .canonical_bin
406                .write()
407                .extend_write(row.block.height(), &row.hash)
408            {
409                let msg = format!("write pipeline: canonical.bin mmap update failed: {e}");
410                for row in staged {
411                    let _ = row
412                        .ack
413                        .send(Err(BlockStoreError::Serialization(msg.clone())));
414                }
415                return Ok(());
416            }
417        }
418    }
419
420    store.pipeline_write_batches.fetch_add(1, Ordering::Relaxed);
421
422    for row in staged {
423        let record = BlockRecord::from_header(&row.block.header, BlockStatus::Validated);
424        store.record_cache.lock().insert(row.hash, record);
425        store.block_cache.insert(row.hash, row.block.clone());
426        store
427            .header_cache
428            .insert(row.hash, row.block.header.clone());
429        let ack_res = match store.maybe_train_dictionary() {
430            Ok(()) => Ok(true),
431            Err(e) => Err(e),
432        };
433        let _ = row.ack.send(ack_res);
434    }
435    Ok(())
436}
437
438/// Lazy iterator over canonical block bodies for a closed height range ([`BLK-006`](../docs/requirements/domains/block_storage/specs/BLK-006.md)).
439///
440/// Constructed only via [`BlockStore::stream_blocks_in_range`]. Holds a precomputed `(height, hash)` list from a
441/// readahead-backed scan of [`CF_CANONICAL`], then loads [`CF_BLOCKS`] rows on demand so callers can stop early without
442/// decompressing the remainder ([`BLK-006.md`](../docs/requirements/domains/block_storage/specs/BLK-006.md) implementation notes).
443///
444/// **Invariants:** Heights in `pairs` are strictly ascending (RocksDB canonical ordering). Each successful item matches
445/// the canonical hash at that height; [`L2Block::height`](dig_block::L2Block::height) should equal the stored height
446/// when the database is consistent ([`BLK-001`](../docs/requirements/domains/block_storage/specs/BLK-001.md) write path).
447pub struct StreamBlocksInRange<'a> {
448    store: &'a BlockStore,
449    pairs: Vec<(u64, Bytes32)>,
450    idx: usize,
451    read_opts: ReadOptions,
452    cf_blocks: &'a ColumnFamily,
453}
454
455impl<'a> Iterator for StreamBlocksInRange<'a> {
456    type Item = Result<L2Block, BlockStoreError>;
457
458    fn next(&mut self) -> Option<Self::Item> {
459        if self.idx >= self.pairs.len() {
460            return None;
461        }
462        let (_expected_height, hash) = self.pairs[self.idx];
463        self.idx += 1;
464        if let Some(block) = self.store.block_cache.get_clone(&hash) {
465            return Some(Ok(block));
466        }
467        self.store
468            .cf_blocks_stream_physical_gets
469            .fetch_add(1, Ordering::Relaxed);
470        let raw_opt = match self.store.db.get_cf_opt(
471            self.cf_blocks,
472            hash_key(&hash).as_slice(),
473            &self.read_opts,
474        ) {
475            Ok(o) => o,
476            Err(e) => return Some(Err(e.into())),
477        };
478        let Some(raw) = raw_opt else {
479            return Some(Err(BlockStoreError::BlockNotFound(hash)));
480        };
481        match self.store.deserialize_block(&raw) {
482            Ok(block) => {
483                self.store.block_cache.insert(hash, block.clone());
484                self.store.header_cache.insert(hash, block.header.clone());
485                Some(Ok(block))
486            }
487            Err(e) => Some(Err(e)),
488        }
489    }
490}
dig_blockstore/pipeline.rs

dig_blockstore/
pipeline.rs
