risingwave_meta/barrier/
backfill_order_control.rs

1// Copyright 2025 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::{HashMap, HashSet};
16
17use risingwave_common::catalog::{FragmentTypeFlag, TableId};
18pub use risingwave_common::id::ActorId;
19
20use super::rpc::build_locality_fragment_state_table_mapping;
21use crate::controller::fragment::InflightFragmentInfo;
22use crate::model::FragmentId;
23use crate::stream::ExtendedFragmentBackfillOrder;
24
25#[derive(Clone, Debug, Default)]
26pub struct BackfillNode {
27    fragment_id: FragmentId,
28    /// How many more actors need to finish,
29    /// before this fragment can finish backfilling.
30    remaining_actors: HashSet<ActorId>,
31    /// How many more dependencies need to finish,
32    /// before this fragment can be backfilled.
33    remaining_dependencies: HashSet<FragmentId>,
34    children: Vec<FragmentId>,
35}
36
37/// Actor done                   -> update `fragment_id` state
38/// Operator done                -> update downstream operator dependency
39/// Operator's dependencies done -> queue operator for backfill
40#[derive(Clone, Debug, Default)]
41pub struct BackfillOrderState {
42    // The order plan.
43    current_backfill_nodes: HashMap<FragmentId, BackfillNode>,
44    // Remaining nodes to finish
45    remaining_backfill_nodes: HashMap<FragmentId, BackfillNode>,
46    // The mapping between actors and fragment_ids
47    actor_to_fragment_id: HashMap<ActorId, FragmentId>,
48    // The mapping between fragment_ids and table_ids of locality provider fragments
49    locality_fragment_state_table_mapping: HashMap<FragmentId, Vec<TableId>>,
50}
51
52/// Get nodes with some dependencies.
53/// These should initially be paused until their dependencies are done.
54pub fn get_nodes_with_backfill_dependencies(
55    backfill_orders: &ExtendedFragmentBackfillOrder,
56) -> HashSet<FragmentId> {
57    backfill_orders.values().flatten().copied().collect()
58}
59
60// constructor
61impl BackfillOrderState {
62    pub fn new(
63        backfill_orders: &ExtendedFragmentBackfillOrder,
64        fragment_infos: &HashMap<FragmentId, InflightFragmentInfo>,
65        locality_fragment_state_table_mapping: HashMap<FragmentId, Vec<TableId>>,
66    ) -> Self {
67        tracing::debug!(?backfill_orders, "initialize backfill order state");
68        let actor_to_fragment_id: HashMap<ActorId, FragmentId> = fragment_infos
69            .iter()
70            .flat_map(|(fragment_id, fragment)| {
71                fragment
72                    .actors
73                    .keys()
74                    .map(|actor_id| (*actor_id, *fragment_id))
75            })
76            .collect();
77
78        let mut backfill_nodes: HashMap<FragmentId, BackfillNode> = HashMap::new();
79
80        for fragment in fragment_infos.values() {
81            if fragment.fragment_type_mask.contains_any([
82                FragmentTypeFlag::StreamScan,
83                FragmentTypeFlag::SourceScan,
84                FragmentTypeFlag::LocalityProvider,
85            ]) {
86                let fragment_id = fragment.fragment_id;
87                backfill_nodes.insert(
88                    fragment_id,
89                    BackfillNode {
90                        fragment_id,
91                        remaining_actors: fragment.actors.keys().copied().collect(),
92                        remaining_dependencies: Default::default(),
93                        children: backfill_orders
94                            .get(&fragment_id)
95                            .cloned()
96                            .unwrap_or_else(Vec::new),
97                    },
98                );
99            }
100        }
101
102        for (fragment_id, children) in backfill_orders.iter() {
103            for child in children {
104                let child_node = backfill_nodes.get_mut(child).unwrap();
105                child_node.remaining_dependencies.insert(*fragment_id);
106            }
107        }
108
109        let mut current_backfill_nodes = HashMap::new();
110        let mut remaining_backfill_nodes = HashMap::new();
111        for (fragment_id, node) in backfill_nodes {
112            if node.remaining_dependencies.is_empty() {
113                current_backfill_nodes.insert(fragment_id, node);
114            } else {
115                remaining_backfill_nodes.insert(fragment_id, node);
116            }
117        }
118
119        Self {
120            current_backfill_nodes,
121            remaining_backfill_nodes,
122            actor_to_fragment_id,
123            locality_fragment_state_table_mapping,
124        }
125    }
126
127    pub fn recover_from_fragment_infos(
128        backfill_orders: &ExtendedFragmentBackfillOrder,
129        fragment_infos: &HashMap<FragmentId, InflightFragmentInfo>,
130    ) -> Self {
131        tracing::debug!(
132            ?backfill_orders,
133            "initialize backfill order state from recovery"
134        );
135        let actor_to_fragment_id = fragment_infos
136            .iter()
137            .flat_map(|(fragment_id, fragment)| {
138                fragment
139                    .actors
140                    .keys()
141                    .map(|actor_id| (*actor_id, *fragment_id))
142            })
143            .collect();
144
145        let locality_fragment_state_table_mapping =
146            build_locality_fragment_state_table_mapping(fragment_infos);
147
148        let mut backfill_nodes: HashMap<FragmentId, BackfillNode> = HashMap::new();
149
150        for (fragment_id, fragment) in fragment_infos {
151            if fragment.fragment_type_mask.contains_any([
152                FragmentTypeFlag::StreamScan,
153                FragmentTypeFlag::SourceScan,
154                FragmentTypeFlag::LocalityProvider,
155            ]) {
156                backfill_nodes.insert(
157                    *fragment_id,
158                    BackfillNode {
159                        fragment_id: *fragment_id,
160                        remaining_actors: fragment.actors.keys().copied().collect(),
161                        remaining_dependencies: Default::default(),
162                        children: backfill_orders
163                            .get(fragment_id)
164                            .cloned()
165                            .unwrap_or_else(Vec::new),
166                    },
167                );
168            }
169        }
170
171        for (fragment_id, children) in backfill_orders.iter() {
172            for child in children {
173                let child_node = backfill_nodes.get_mut(child).unwrap();
174                child_node.remaining_dependencies.insert(*fragment_id);
175            }
176        }
177
178        let mut current_backfill_nodes = HashMap::new();
179        let mut remaining_backfill_nodes = HashMap::new();
180        for (fragment_id, node) in backfill_nodes {
181            if node.remaining_dependencies.is_empty() {
182                current_backfill_nodes.insert(fragment_id, node);
183            } else {
184                remaining_backfill_nodes.insert(fragment_id, node);
185            }
186        }
187
188        Self {
189            current_backfill_nodes,
190            remaining_backfill_nodes,
191            actor_to_fragment_id,
192            locality_fragment_state_table_mapping,
193        }
194    }
195}
196
197// state transitions
198impl BackfillOrderState {
199    pub fn finish_actor(&mut self, actor_id: ActorId) -> Vec<FragmentId> {
200        let Some(fragment_id) = self.actor_to_fragment_id.get(&actor_id) else {
201            tracing::error!(%actor_id, "fragment not found for actor");
202            return vec![];
203        };
204        // NOTE(kwannoel):
205        // Backfill order are specified by the user, for instance:
206        // t1->t2 means that t1 must be backfilled before t2.
207        // However, each snapshot executor may finish ahead of time if there's no data to backfill.
208        // For instance, if t2 has no data to backfill,
209        // and t1 has a lot of data to backfill,
210        // t2's scan operator might finish immediately,
211        // and t2 will finish before t1.
212        // In such cases, we should directly update it in remaining backfill nodes instead,
213        // so we should track whether a node finished in order.
214        let (node, is_in_order) = match self.current_backfill_nodes.get_mut(fragment_id) {
215            Some(node) => (node, true),
216            None => {
217                let Some(node) = self.remaining_backfill_nodes.get_mut(fragment_id) else {
218                    tracing::error!(
219                        %fragment_id,
220                        %actor_id,
221                        "fragment not found in current_backfill_nodes or remaining_backfill_nodes"
222                    );
223                    return vec![];
224                };
225                (node, false)
226            }
227        };
228
229        assert!(node.remaining_actors.remove(&actor_id), "missing actor");
230        tracing::debug!(
231            %actor_id,
232            remaining_actors = node.remaining_actors.len(),
233            %fragment_id,
234            "finish_backfilling_actor"
235        );
236        if node.remaining_actors.is_empty() && is_in_order {
237            self.finish_fragment(*fragment_id)
238        } else {
239            vec![]
240        }
241    }
242
243    pub fn finish_fragment(&mut self, fragment_id: FragmentId) -> Vec<FragmentId> {
244        let mut newly_scheduled = vec![];
245        // Decrease the remaining_dependency_count of the children.
246        // If the remaining_dependency_count is 0, add the child to the current_backfill_nodes.
247        if let Some(node) = self.current_backfill_nodes.remove(&fragment_id) {
248            for child_id in &node.children {
249                let newly_scheduled_child_finished = {
250                    let child = self.remaining_backfill_nodes.get_mut(child_id).unwrap();
251                    assert!(
252                        child.remaining_dependencies.remove(&fragment_id),
253                        "missing dependency"
254                    );
255                    if child.remaining_dependencies.is_empty() {
256                        let should_schedule = !child.remaining_actors.is_empty();
257                        if should_schedule {
258                            tracing::debug!(fragment_id = ?child_id, "schedule next backfill node");
259                        }
260                        self.current_backfill_nodes
261                            .insert(child.fragment_id, child.clone());
262                        if should_schedule {
263                            newly_scheduled.push(child.fragment_id);
264                        }
265                        child.remaining_actors.is_empty()
266                    } else {
267                        false
268                    }
269                };
270                if newly_scheduled_child_finished {
271                    newly_scheduled.extend(self.finish_fragment(*child_id));
272                }
273            }
274        } else {
275            tracing::error!(%fragment_id, "fragment not found in current_backfill_nodes");
276            return vec![];
277        }
278        newly_scheduled
279    }
280
281    pub fn current_backfill_node_fragment_ids(&self) -> Vec<FragmentId> {
282        self.current_backfill_nodes.keys().copied().collect()
283    }
284
285    pub fn get_locality_fragment_state_table_mapping(&self) -> &HashMap<FragmentId, Vec<TableId>> {
286        &self.locality_fragment_state_table_mapping
287    }
288
289    /// Refresh actor mapping after reschedule and return newly scheduled fragments.
290    pub fn refresh_actors(
291        &mut self,
292        fragment_actors: &HashMap<FragmentId, HashSet<ActorId>>,
293    ) -> Vec<FragmentId> {
294        self.actor_to_fragment_id = fragment_actors
295            .iter()
296            .flat_map(|(fragment_id, actors)| {
297                actors.iter().map(|actor_id| (*actor_id, *fragment_id))
298            })
299            .collect();
300
301        for node in self
302            .current_backfill_nodes
303            .values_mut()
304            .chain(self.remaining_backfill_nodes.values_mut())
305        {
306            if let Some(actors) = fragment_actors.get(&node.fragment_id) {
307                node.remaining_actors = actors.iter().copied().collect();
308            } else {
309                node.remaining_actors.clear();
310            }
311        }
312
313        let finished_fragments: Vec<_> = self
314            .current_backfill_nodes
315            .iter()
316            .filter(|(_, node)| node.remaining_actors.is_empty())
317            .map(|(fragment_id, _)| *fragment_id)
318            .collect();
319
320        let mut newly_scheduled = vec![];
321        for fragment_id in finished_fragments {
322            newly_scheduled.extend(self.finish_fragment(fragment_id));
323        }
324        newly_scheduled
325    }
326}
risingwave_meta/barrier/backfill_order_control.rs

risingwave_meta/barrier/
backfill_order_control.rs
