risingwave_meta/barrier/
utils.rs

1// Copyright 2024 RisingWave Labs
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::collections::{BTreeMap, HashMap, HashSet, VecDeque};
16
17use itertools::Itertools;
18use risingwave_common::catalog::TableId;
19use risingwave_common::util::epoch::EpochPair;
20use risingwave_common::util::stream_graph_visitor::visit_stream_node_cont;
21use risingwave_hummock_sdk::sstable_info::SstableInfo;
22use risingwave_hummock_sdk::table_stats::from_prost_table_stats_map;
23use risingwave_hummock_sdk::table_watermark::{
24    TableWatermarks, merge_multiple_new_table_watermarks,
25};
26use risingwave_hummock_sdk::vector_index::{VectorIndexAdd, VectorIndexDelta};
27use risingwave_hummock_sdk::{HummockSstableObjectId, LocalSstableInfo};
28use risingwave_meta_model::WorkerId;
29use risingwave_pb::catalog::PbTable;
30use risingwave_pb::catalog::table::PbTableType;
31use risingwave_pb::hummock::vector_index_delta::PbVectorIndexInit;
32use risingwave_pb::stream_plan::stream_node::NodeBody;
33use risingwave_pb::stream_service::BarrierCompleteResponse;
34use risingwave_pb::stream_service::barrier_complete_response::IcebergPkIndexSinkMetadata;
35
36use crate::barrier::CreateStreamingJobCommandInfo;
37use crate::barrier::command::PostCollectCommand;
38use crate::barrier::complete_task::CompleteBarrierTask;
39use crate::barrier::partial_graph::PartialGraphBarrierInfo;
40use crate::hummock::NewTableFragmentInfo;
41
42#[expect(clippy::type_complexity)]
43pub(super) fn collect_resp_info(
44    resps: Vec<BarrierCompleteResponse>,
45) -> (
46    HashMap<HummockSstableObjectId, WorkerId>,
47    Vec<LocalSstableInfo>,
48    HashMap<TableId, TableWatermarks>,
49    Vec<SstableInfo>,
50    HashMap<TableId, Vec<VectorIndexAdd>>,
51    HashSet<TableId>,
52    Vec<IcebergPkIndexSinkMetadata>,
53) {
54    let mut sst_to_worker: HashMap<HummockSstableObjectId, _> = HashMap::new();
55    let mut synced_ssts: Vec<LocalSstableInfo> = vec![];
56    let mut table_watermarks = Vec::with_capacity(resps.len());
57    let mut old_value_ssts = Vec::with_capacity(resps.len());
58    let mut vector_index_adds = HashMap::new();
59    let mut truncate_tables: HashSet<TableId> = HashSet::new();
60    let mut iceberg_pk_index_sink_metadata = Vec::new();
61
62    for resp in resps {
63        let ssts_iter = resp.synced_sstables.into_iter().map(|local_sst| {
64            let sst_info = local_sst.sst.expect("field not None");
65            sst_to_worker.insert(sst_info.object_id, resp.worker_id);
66            LocalSstableInfo::new(
67                sst_info.into(),
68                from_prost_table_stats_map(local_sst.table_stats_map),
69                local_sst.created_at,
70            )
71        });
72        synced_ssts.extend(ssts_iter);
73        table_watermarks.push(resp.table_watermarks);
74        old_value_ssts.extend(resp.old_value_sstables.into_iter().map(|s| s.into()));
75        for (table_id, vector_index_add) in resp.vector_index_adds {
76            vector_index_adds
77                .try_insert(
78                    table_id,
79                    vector_index_add
80                        .adds
81                        .into_iter()
82                        .map(VectorIndexAdd::from)
83                        .collect(),
84                )
85                .expect("non-duplicate");
86        }
87        truncate_tables.extend(resp.truncate_tables);
88        iceberg_pk_index_sink_metadata.extend(resp.iceberg_pk_index_sink_metadata);
89    }
90
91    (
92        sst_to_worker,
93        synced_ssts,
94        merge_multiple_new_table_watermarks(
95            table_watermarks
96                .into_iter()
97                .map(|watermarks| {
98                    watermarks
99                        .into_iter()
100                        .map(|(table_id, watermarks)| {
101                            (table_id, TableWatermarks::from(&watermarks))
102                        })
103                        .collect()
104                })
105                .collect_vec(),
106        ),
107        old_value_ssts,
108        vector_index_adds,
109        truncate_tables,
110        iceberg_pk_index_sink_metadata,
111    )
112}
113
114pub(super) fn collect_new_vector_index_info(
115    info: &CreateStreamingJobCommandInfo,
116) -> Option<&PbTable> {
117    let mut vector_index_table = None;
118    {
119        for fragment in info.stream_job_fragments.fragments.values() {
120            visit_stream_node_cont(&fragment.nodes, |node| {
121                match node.node_body.as_ref().unwrap() {
122                    NodeBody::VectorIndexWrite(vector_index_write) => {
123                        let index_table = vector_index_write.table.as_ref().unwrap();
124                        assert_eq!(index_table.table_type, PbTableType::VectorIndex as i32);
125                        vector_index_table = Some(index_table);
126                        false
127                    }
128                    _ => true,
129                }
130            })
131        }
132        vector_index_table
133    }
134}
135
136pub(super) fn collect_independent_job_commit_epoch_info(
137    task: &mut CompleteBarrierTask,
138    epoch: u64,
139    resps: Vec<BarrierCompleteResponse>,
140    barrier_info: &PartialGraphBarrierInfo,
141) {
142    let (
143        sst_to_context,
144        sstables,
145        new_table_watermarks,
146        old_value_sst,
147        vector_index_adds,
148        truncate_tables,
149        iceberg_pk_index_sink_metadata,
150    ) = collect_resp_info(resps);
151    assert!(old_value_sst.is_empty());
152    let commit_info = &mut task.commit_info;
153    commit_info.sst_to_context.extend(sst_to_context);
154    commit_info.sstables.extend(sstables);
155    commit_info
156        .new_table_watermarks
157        .extend(new_table_watermarks);
158    for (table_id, vector_index_adds) in vector_index_adds {
159        commit_info
160            .vector_index_delta
161            .try_insert(table_id, VectorIndexDelta::Adds(vector_index_adds))
162            .expect("non-duplicate");
163    }
164    commit_info.truncate_tables.extend(truncate_tables);
165    barrier_info
166        .table_ids_to_commit
167        .iter()
168        .for_each(|table_id| {
169            commit_info
170                .tables_to_commit
171                .try_insert(*table_id, epoch)
172                .expect("non duplicate");
173        });
174    if let PostCollectCommand::CreateStreamingJob { info, .. } = &barrier_info.post_collect_command
175    {
176        commit_info
177            .new_table_fragment_infos
178            .push(NewTableFragmentInfo {
179                table_ids: barrier_info.table_ids_to_commit.clone(),
180            });
181        if let Some(index_table) = collect_new_vector_index_info(info) {
182            commit_info
183                .vector_index_delta
184                .try_insert(
185                    index_table.id,
186                    VectorIndexDelta::Init(PbVectorIndexInit {
187                        info: Some(index_table.vector_index_info.unwrap()),
188                    }),
189                )
190                .expect("non-duplicate");
191        }
192    };
193    task.iceberg_pk_index_sink_metadata
194        .extend(iceberg_pk_index_sink_metadata);
195}
196
197pub(super) type NodeToCollect = HashSet<WorkerId>;
198pub(super) fn is_valid_after_worker_err(
199    node_to_collect: &NodeToCollect,
200    worker_id: WorkerId,
201) -> bool {
202    !node_to_collect.contains(&worker_id)
203}
204
205#[derive(Debug)]
206struct InflightBarrierNode<K, Item, Info> {
207    epoch: EpochPair,
208    to_collect: HashSet<K>,
209    collected: HashMap<K, Item>,
210    info: Info,
211}
212
213#[derive(Debug)]
214struct CollectedBarrierNode<K, Item, Info> {
215    epoch: EpochPair,
216    collected: HashMap<K, Item>,
217    info: Info,
218}
219
220#[derive(Debug)]
221pub(super) struct BarrierItemCollector<K, Item, Info> {
222    /// `prev_epoch` -> barrier
223    inflight_barriers: BTreeMap<u64, InflightBarrierNode<K, Item, Info>>,
224    /// newer epoch at the back. `push_back` and `pop_front`
225    collected_barriers: VecDeque<CollectedBarrierNode<K, Item, Info>>,
226    /// Whether successive barriers are expected to be epoch-contiguous (the `curr` epoch of one
227    /// barrier equals the `prev` epoch of the next). This holds for the main per-barrier collector,
228    /// where every injected barrier is enqueued, but NOT for collectors that are only fed a sparse
229    /// subset of barriers (e.g. the finishing-jobs collector, enqueued only when a snapshot-backfill
230    /// job finishes). When `false`, only monotonicity and non-duplication are checked.
231    require_contiguous: bool,
232}
233
234impl<K: std::fmt::Debug + Eq + std::hash::Hash, Item: std::fmt::Debug, Info: std::fmt::Debug>
235    BarrierItemCollector<K, Item, Info>
236{
237    pub(super) fn new(require_contiguous: bool) -> Self {
238        Self {
239            inflight_barriers: Default::default(),
240            collected_barriers: Default::default(),
241            require_contiguous,
242        }
243    }
244
245    pub(super) fn is_empty(&self) -> bool {
246        self.inflight_barriers.is_empty() && self.collected_barriers.is_empty()
247    }
248
249    pub(super) fn inflight_barrier_num(&self) -> usize {
250        self.inflight_barriers.len()
251    }
252
253    pub(super) fn collected_barrier_num(&self) -> usize {
254        self.collected_barriers.len()
255    }
256
257    pub(super) fn enqueue(&mut self, epoch: EpochPair, to_collect: HashSet<K>, info: Info) {
258        assert!(!to_collect.is_empty());
259        if let Some((last_prev_epoch, last_barrier)) = self.inflight_barriers.last_key_value() {
260            if self.require_contiguous {
261                assert_eq!(last_barrier.epoch.curr, epoch.prev);
262            }
263            assert!(*last_prev_epoch < epoch.prev);
264        }
265        self.inflight_barriers
266            .try_insert(
267                epoch.prev,
268                InflightBarrierNode {
269                    epoch,
270                    to_collect,
271                    collected: Default::default(),
272                    info,
273                },
274            )
275            .expect("non-duplicated");
276    }
277
278    pub(super) fn collect(&mut self, prev_epoch: u64, key: K, item: Item) {
279        let inflight_barrier = self
280            .inflight_barriers
281            .get_mut(&prev_epoch)
282            .expect("should exist");
283        assert!(inflight_barrier.to_collect.remove(&key));
284        inflight_barrier
285            .collected
286            .try_insert(key, item)
287            .expect("non-duplicate");
288    }
289
290    pub(super) fn barrier_collected(&mut self) -> Option<(EpochPair, &Info)> {
291        if let Some(entry) = self.inflight_barriers.first_entry()
292            && entry.get().to_collect.is_empty()
293        {
294            let InflightBarrierNode {
295                epoch,
296                collected,
297                info,
298                ..
299            } = entry.remove();
300            if self.require_contiguous
301                && let Some(prev_barrier) = self.collected_barriers.back()
302            {
303                assert_eq!(prev_barrier.epoch.curr, epoch.prev);
304            }
305            self.collected_barriers.push_back(CollectedBarrierNode {
306                epoch,
307                collected,
308                info,
309            });
310            Some((
311                epoch,
312                &self.collected_barriers.back().expect("non-empty").info,
313            ))
314        } else {
315            None
316        }
317    }
318
319    pub(super) fn advance_collected(&mut self) {
320        while self.barrier_collected().is_some() {}
321    }
322
323    pub(super) fn first_inflight_epoch(&self) -> Option<EpochPair> {
324        self.inflight_barriers
325            .first_key_value()
326            .map(|(_, barrier)| barrier.epoch)
327    }
328
329    pub(super) fn last_collected(&self) -> Option<(EpochPair, &HashMap<K, Item>, &Info)> {
330        self.collected_barriers
331            .back()
332            .map(|barrier| (barrier.epoch, &barrier.collected, &barrier.info))
333    }
334
335    pub(super) fn iter_infos(&self) -> impl Iterator<Item = &Info> + '_ {
336        // Preserve epoch order: collected barriers are older and kept from old to new,
337        // followed by inflight barriers ordered by `prev_epoch` in the BTreeMap.
338        self.collected_barriers
339            .iter()
340            .map(|barrier| &barrier.info)
341            .chain(self.inflight_barriers.values().map(|barrier| &barrier.info))
342    }
343
344    pub(super) fn into_infos(self) -> impl Iterator<Item = Info> {
345        self.inflight_barriers
346            .into_values()
347            .map(|barrier| barrier.info)
348            .chain(
349                self.collected_barriers
350                    .into_iter()
351                    .map(|barrier| barrier.info),
352            )
353    }
354
355    pub(super) fn iter_to_collect(&self) -> impl Iterator<Item = &HashSet<K>> + '_ {
356        self.inflight_barriers
357            .values()
358            .map(|barrier| &barrier.to_collect)
359    }
360
361    pub(super) fn take_collected_if(
362        &mut self,
363        cond: impl FnOnce(EpochPair) -> bool,
364    ) -> Option<(EpochPair, HashMap<K, Item>, Info)> {
365        self.collected_barriers
366            .pop_front_if(move |barrier| cond(barrier.epoch))
367            .map(|barrier| (barrier.epoch, barrier.collected, barrier.info))
368    }
369}
370
371#[cfg(test)]
372mod tests {
373    use std::collections::HashSet;
374
375    use risingwave_common::util::epoch::EpochPair;
376
377    use super::BarrierItemCollector;
378
379    /// The finishing-jobs collector is fed only on the sparse subset of barriers where a
380    /// snapshot-backfill job finishes, so enqueued epochs are not contiguous. Enqueuing a barrier
381    /// whose `prev` does not match the previous barrier's `curr` must be tolerated.
382    #[test]
383    fn non_contiguous_enqueue_allowed_when_not_required() {
384        let mut collector = BarrierItemCollector::<u32, (), ()>::new(false);
385        collector.enqueue(EpochPair::new(200, 100), HashSet::from([1]), ());
386        // gap: prev (500) != previous curr (200)
387        collector.enqueue(EpochPair::new(600, 500), HashSet::from([2]), ());
388        assert_eq!(collector.inflight_barrier_num(), 2);
389    }
390
391    /// The main per-barrier collector enqueues every barrier, so contiguity is a real invariant and
392    /// must still be asserted.
393    #[test]
394    #[should_panic]
395    fn non_contiguous_enqueue_rejected_when_required() {
396        let mut collector = BarrierItemCollector::<u32, (), ()>::new(true);
397        collector.enqueue(EpochPair::new(200, 100), HashSet::from([1]), ());
398        collector.enqueue(EpochPair::new(600, 500), HashSet::from([2]), ());
399    }
400}
risingwave_meta/barrier/utils.rs

risingwave_meta/barrier/
utils.rs
