risingwave_meta/serving/
mod.rs

1// Copyright 2023 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};
16use std::sync::Arc;
17
18use parking_lot::RwLock;
19use risingwave_common::bitmap::Bitmap;
20use risingwave_common::hash::WorkerSlotMapping;
21use risingwave_common::vnode_mapping::vnode_placement::place_vnode;
22use risingwave_meta_model::{TableId, WorkerId};
23use risingwave_pb::common::{WorkerNode, WorkerType};
24use risingwave_pb::meta::serving_table_vnode_mappings::PbServingTableVnodeMapping;
25use risingwave_pb::meta::subscribe_response::{Info, Operation};
26use risingwave_pb::meta::table_fragments::fragment::FragmentDistributionType;
27use risingwave_pb::meta::{
28    FragmentWorkerSlotMapping, FragmentWorkerSlotMappings, PbServingTableVnodeMappings,
29    PbTableRefillRuntimeConfig,
30};
31use tokio::sync::oneshot::Sender;
32use tokio::task::JoinHandle;
33
34use crate::MetaResult;
35use crate::controller::fragment::FragmentServingInfo;
36use crate::controller::session_params::SessionParamsControllerRef;
37use crate::manager::{LocalNotification, MetadataManager, NotificationManagerRef, WorkerKey};
38use crate::model::FragmentId;
39
40pub type ServingVnodeMappingRef = Arc<ServingVnodeMapping>;
41
42#[derive(Default)]
43pub struct ServingVnodeMapping {
44    serving_vnode_mappings: RwLock<HashMap<FragmentId, WorkerSlotMapping>>,
45}
46
47impl ServingVnodeMapping {
48    pub fn all(&self) -> HashMap<FragmentId, WorkerSlotMapping> {
49        self.serving_vnode_mappings.read().clone()
50    }
51
52    pub(crate) fn table_vnode_mappings_by_worker(
53        &self,
54        worker_ids: impl IntoIterator<Item = WorkerId>,
55        fragment_serving_infos: &HashMap<FragmentId, FragmentServingInfo>,
56    ) -> HashMap<WorkerId, HashMap<TableId, Bitmap>> {
57        let mut result = worker_ids
58            .into_iter()
59            .map(|worker_id| (worker_id, HashMap::new()))
60            .collect::<HashMap<_, _>>();
61        let serving_vnode_mappings = self.all();
62
63        for (fragment_id, info) in fragment_serving_infos {
64            let Some(table_id) = info.result_table_id else {
65                continue;
66            };
67            let Some(mapping) = serving_vnode_mappings.get(fragment_id) else {
68                continue;
69            };
70
71            for (worker_slot_id, bitmap) in mapping.to_bitmaps() {
72                let Some(table_mappings) = result.get_mut(&worker_slot_id.worker_id()) else {
73                    continue;
74                };
75                table_mappings
76                    .entry(table_id)
77                    .and_modify(|current| *current |= &bitmap)
78                    .or_insert(bitmap);
79            }
80        }
81        result
82    }
83
84    /// Upsert mapping for given fragments according to the latest `workers`.
85    /// Returns (successful updates, failed updates).
86    pub fn upsert(
87        &self,
88        fragment_serving_infos: &HashMap<FragmentId, FragmentServingInfo>,
89        workers: &[WorkerNode],
90        max_serving_parallelism: Option<u64>,
91    ) -> (HashMap<FragmentId, WorkerSlotMapping>, HashSet<FragmentId>) {
92        let mut serving_vnode_mappings = self.serving_vnode_mappings.write();
93        Self::upsert_locked(
94            &mut serving_vnode_mappings,
95            fragment_serving_infos,
96            workers,
97            max_serving_parallelism,
98        )
99    }
100
101    /// Rebuild mappings from a complete catalog snapshot.
102    ///
103    /// Unlike [`Self::upsert`], this removes mappings for fragments absent from the snapshot.
104    pub(crate) fn reconcile(
105        &self,
106        fragment_serving_infos: &HashMap<FragmentId, FragmentServingInfo>,
107        workers: &[WorkerNode],
108        max_serving_parallelism: Option<u64>,
109    ) -> (HashMap<FragmentId, WorkerSlotMapping>, HashSet<FragmentId>) {
110        let mut serving_vnode_mappings = self.serving_vnode_mappings.write();
111        serving_vnode_mappings
112            .retain(|fragment_id, _| fragment_serving_infos.contains_key(fragment_id));
113        Self::upsert_locked(
114            &mut serving_vnode_mappings,
115            fragment_serving_infos,
116            workers,
117            max_serving_parallelism,
118        )
119    }
120
121    fn upsert_locked(
122        serving_vnode_mappings: &mut HashMap<FragmentId, WorkerSlotMapping>,
123        fragment_serving_infos: &HashMap<FragmentId, FragmentServingInfo>,
124        workers: &[WorkerNode],
125        max_serving_parallelism: Option<u64>,
126    ) -> (HashMap<FragmentId, WorkerSlotMapping>, HashSet<FragmentId>) {
127        let mut upserted: HashMap<FragmentId, WorkerSlotMapping> = HashMap::default();
128        let mut failed: HashSet<FragmentId> = HashSet::default();
129        for (fragment_id, info) in fragment_serving_infos {
130            let new_mapping = {
131                let old_mapping = serving_vnode_mappings.get(fragment_id);
132                let max_parallelism = match info.distribution_type {
133                    FragmentDistributionType::Unspecified => unreachable!(),
134                    FragmentDistributionType::Single => Some(1),
135                    FragmentDistributionType::Hash => None,
136                }
137                .or_else(|| max_serving_parallelism.map(|p| p as usize));
138                place_vnode(old_mapping, workers, max_parallelism, info.vnode_count)
139            };
140            match new_mapping {
141                None => {
142                    serving_vnode_mappings.remove(fragment_id);
143                    failed.insert(*fragment_id);
144                }
145                Some(mapping) => {
146                    serving_vnode_mappings.insert(*fragment_id, mapping.clone());
147                    upserted.insert(*fragment_id, mapping);
148                }
149            }
150        }
151        (upserted, failed)
152    }
153
154    fn remove(&self, fragment_ids: &[FragmentId]) {
155        let mut mappings = self.serving_vnode_mappings.write();
156        for fragment_id in fragment_ids {
157            mappings.remove(fragment_id);
158        }
159    }
160}
161
162pub(crate) fn to_fragment_worker_slot_mapping(
163    mappings: &HashMap<FragmentId, WorkerSlotMapping>,
164) -> Vec<FragmentWorkerSlotMapping> {
165    mappings
166        .iter()
167        .map(|(&fragment_id, mapping)| FragmentWorkerSlotMapping {
168            fragment_id,
169            mapping: Some(mapping.to_protobuf()),
170        })
171        .collect()
172}
173
174pub(crate) fn to_deleted_fragment_worker_slot_mapping(
175    fragment_ids: impl Iterator<Item = FragmentId>,
176) -> Vec<FragmentWorkerSlotMapping> {
177    fragment_ids
178        .map(|fragment_id| FragmentWorkerSlotMapping {
179            fragment_id,
180            mapping: None,
181        })
182        .collect()
183}
184
185pub async fn on_meta_start(
186    notification_manager: NotificationManagerRef,
187    metadata_manager: &MetadataManager,
188    serving_vnode_mapping: ServingVnodeMappingRef,
189    max_serving_parallelism: Option<u64>,
190) {
191    let (serving_compute_nodes, fragment_serving_infos) = fetch_serving_infos(metadata_manager)
192        .await
193        .expect("fail to fetch serving infos");
194    let (mappings, failed) = serving_vnode_mapping.reconcile(
195        &fragment_serving_infos,
196        &serving_compute_nodes,
197        max_serving_parallelism,
198    );
199    tracing::debug!(
200        "Initialize serving vnode mapping snapshot for fragments {:?}.",
201        mappings.keys()
202    );
203    if !failed.is_empty() {
204        tracing::warn!(
205            "Fail to update serving vnode mapping for fragments {:?}.",
206            failed
207        );
208    }
209    notification_manager.notify_frontend_without_version(
210        Operation::Snapshot,
211        Info::ServingWorkerSlotMappings(FragmentWorkerSlotMappings {
212            mappings: to_fragment_worker_slot_mapping(&mappings),
213        }),
214    );
215}
216
217pub(crate) async fn fetch_serving_infos(
218    metadata_manager: &MetadataManager,
219) -> MetaResult<(Vec<WorkerNode>, HashMap<FragmentId, FragmentServingInfo>)> {
220    let fragment_serving_infos = metadata_manager
221        .catalog_controller
222        .fragment_serving_infos()
223        .await?;
224    let serving_compute_nodes = metadata_manager
225        .cluster_controller
226        .list_active_serving_workers()
227        .await?;
228    Ok((
229        serving_compute_nodes,
230        fragment_serving_infos
231            .into_iter()
232            .map(|(fragment_id, info)| (fragment_id as FragmentId, info))
233            .collect(),
234    ))
235}
236
237pub(crate) async fn sync_serving_table_vnode_mappings_to_hummock(
238    notification_manager: &NotificationManagerRef,
239    serving_vnode_mapping: &ServingVnodeMappingRef,
240    workers: &[WorkerNode],
241    fragment_serving_infos: &HashMap<FragmentId, FragmentServingInfo>,
242) {
243    let mut table_vnode_mappings = serving_vnode_mapping.table_vnode_mappings_by_worker(
244        workers.iter().map(|worker| worker.id),
245        fragment_serving_infos,
246    );
247    for worker in workers {
248        let Some(host) = worker.host.clone() else {
249            tracing::warn!(worker_id = %worker.id, "serving worker host not found");
250            continue;
251        };
252        let table_vnode_mapping = table_vnode_mappings
253            .remove(&worker.id)
254            .expect("requested worker must have a table vnode mapping");
255        let config = PbTableRefillRuntimeConfig {
256            serving_table_vnode_mappings: Some(to_pb_serving_table_vnode_mappings(
257                &table_vnode_mapping,
258            )),
259            ..Default::default()
260        };
261        notification_manager.notify_hummock_targeted_update(
262            WorkerKey(host),
263            Info::TableRefillRuntimeConfig(config),
264        );
265    }
266}
267
268pub fn start_serving_vnode_mapping_worker(
269    notification_manager: NotificationManagerRef,
270    metadata_manager: MetadataManager,
271    serving_vnode_mapping: ServingVnodeMappingRef,
272    session_params: SessionParamsControllerRef,
273) -> (JoinHandle<()>, Sender<()>) {
274    let (local_notification_tx, mut local_notification_rx) = tokio::sync::mpsc::unbounded_channel();
275    let (shutdown_tx, mut shutdown_rx) = tokio::sync::oneshot::channel();
276    notification_manager.insert_local_sender(local_notification_tx);
277    let join_handle = tokio::spawn(async move {
278        let reset = || async {
279            let (workers, fragment_serving_infos) = fetch_serving_infos(&metadata_manager)
280                .await
281                .expect("fail to fetch serving infos");
282            let max_serving_parallelism = session_params
283                .get_params()
284                .await
285                .batch_parallelism()
286                .map(|p| p.get());
287            let (mappings, failed) = serving_vnode_mapping.reconcile(
288                &fragment_serving_infos,
289                &workers,
290                max_serving_parallelism,
291            );
292            tracing::debug!(
293                "Update serving vnode mapping snapshot for fragments {:?}.",
294                mappings.keys()
295            );
296            if !failed.is_empty() {
297                tracing::warn!(
298                    "Fail to update serving vnode mapping for fragments {:?}.",
299                    failed
300                );
301            }
302            notification_manager.notify_frontend_without_version(
303                Operation::Snapshot,
304                Info::ServingWorkerSlotMappings(FragmentWorkerSlotMappings {
305                    mappings: to_fragment_worker_slot_mapping(&mappings),
306                }),
307            );
308            sync_serving_table_vnode_mappings_to_hummock(
309                &notification_manager,
310                &serving_vnode_mapping,
311                &workers,
312                &fragment_serving_infos,
313            )
314            .await;
315        };
316        loop {
317            tokio::select! {
318                notification = local_notification_rx.recv() => {
319                    match notification {
320                        Some(notification) => {
321                            match notification {
322                                LocalNotification::WorkerNodeActivated(w) | LocalNotification::WorkerNodeDeleted(w) =>  {
323                                    if w.r#type() != WorkerType::ComputeNode || !w.property.as_ref().is_some_and(|p| p.is_serving) {
324                                        continue;
325                                    }
326                                    reset().await;
327                                }
328                                LocalNotification::BatchParallelismChange => {
329                                    reset().await;
330                                }
331                                LocalNotification::ServingFragmentMappingsUpsert(fragment_ids) => {
332                                    if fragment_ids.is_empty() {
333                                        continue;
334                                    }
335                                    let (workers, fragment_serving_infos) = fetch_serving_infos(&metadata_manager)
336                                        .await
337                                        .expect("fail to fetch serving infos");
338                                    let filtered_fragment_serving_infos = fragment_ids.iter().filter_map(|frag_id| {
339                                        match fragment_serving_infos.get(frag_id) {
340                                            Some(info) => Some((*frag_id, info.clone())),
341                                            None => {
342                                                tracing::warn!(fragment_id = %frag_id, "fragment serving info not found");
343                                                None
344                                            }
345                                        }
346                                    }).collect();
347                                    let max_serving_parallelism = session_params
348                                        .get_params()
349                                        .await
350                                        .batch_parallelism()
351                                        .map(|p|p.get());
352                                    let (upserted, failed) = serving_vnode_mapping.upsert(&filtered_fragment_serving_infos, &workers, max_serving_parallelism);
353                                    if !upserted.is_empty() {
354                                        tracing::debug!("Update serving vnode mapping for fragments {:?}.", upserted.keys());
355                                        notification_manager.notify_frontend_without_version(Operation::Update, Info::ServingWorkerSlotMappings(FragmentWorkerSlotMappings{ mappings: to_fragment_worker_slot_mapping(&upserted) }));
356                                    }
357                                    if !failed.is_empty() {
358                                        tracing::warn!("Fail to update serving vnode mapping for fragments {:?}.", failed);
359                                        notification_manager.notify_frontend_without_version(Operation::Delete, Info::ServingWorkerSlotMappings(FragmentWorkerSlotMappings{ mappings: to_deleted_fragment_worker_slot_mapping(failed.iter().cloned())}));
360                                    }
361                                    sync_serving_table_vnode_mappings_to_hummock(
362                                        &notification_manager,
363                                        &serving_vnode_mapping,
364                                        &workers,
365                                        &fragment_serving_infos,
366                                    )
367                                    .await;
368                                }
369                                LocalNotification::ServingFragmentMappingsDelete(fragment_ids) => {
370                                    if fragment_ids.is_empty() {
371                                        continue;
372                                    }
373
374                                    tracing::debug!("Delete serving vnode mapping for fragments {:?}.", fragment_ids);
375
376                                    serving_vnode_mapping.remove(&fragment_ids);
377                                    notification_manager.notify_frontend_without_version(Operation::Delete, Info::ServingWorkerSlotMappings(FragmentWorkerSlotMappings{ mappings: to_deleted_fragment_worker_slot_mapping(fragment_ids.iter().cloned()) }));
378                                    let (workers, fragment_serving_infos) = fetch_serving_infos(&metadata_manager)
379                                        .await
380                                        .expect("fail to fetch serving infos");
381                                    sync_serving_table_vnode_mappings_to_hummock(
382                                        &notification_manager,
383                                        &serving_vnode_mapping,
384                                        &workers,
385                                        &fragment_serving_infos,
386                                    )
387                                    .await;
388                                }
389                                _ => {}
390                            }
391                        }
392                        None => {
393                            return;
394                        }
395                    }
396                }
397                _ = &mut shutdown_rx => {
398                    return;
399                }
400            }
401        }
402    });
403    (join_handle, shutdown_tx)
404}
405
406pub(crate) fn to_pb_serving_table_vnode_mappings(
407    table_vnode_mapping: &HashMap<TableId, Bitmap>,
408) -> PbServingTableVnodeMappings {
409    PbServingTableVnodeMappings {
410        mappings: table_vnode_mapping
411            .iter()
412            .map(|(table_id, bitmap)| PbServingTableVnodeMapping {
413                table_id: table_id.as_raw_id(),
414                bitmap: Some(bitmap.to_protobuf()),
415            })
416            .collect(),
417    }
418}
419
420#[cfg(test)]
421mod tests {
422    use risingwave_common::hash::{VirtualNode, WorkerSlotId};
423    use risingwave_pb::common::{HostAddress, worker_node};
424    use risingwave_pb::meta::subscribe_response::Info;
425    use risingwave_pb::meta::{SubscribeResponse, SubscribeType};
426    use tokio::sync::mpsc;
427
428    use super::*;
429    use crate::controller::SqlMetaStore;
430    use crate::manager::NotificationManager;
431
432    fn serving_worker(id: u32) -> WorkerNode {
433        WorkerNode {
434            id: id.into(),
435            r#type: WorkerType::ComputeNode.into(),
436            host: Some(HostAddress {
437                host: format!("host{}", id),
438                port: id as i32,
439            }),
440            state: worker_node::State::Running as i32,
441            property: Some(worker_node::Property {
442                is_serving: true,
443                parallelism: 1,
444                ..Default::default()
445            }),
446            ..Default::default()
447        }
448    }
449
450    #[test]
451    fn test_table_vnode_mappings_merge_worker_slots() {
452        let worker1 = WorkerId::new(1);
453        let worker2 = WorkerId::new(2);
454        let worker3 = WorkerId::new(3);
455        let slot1 = WorkerSlotId::new(worker1, 0);
456        let slot2 = WorkerSlotId::new(worker1, 1);
457        let slot3 = WorkerSlotId::new(worker2, 0);
458        let fragment_id = FragmentId::new(233);
459        let table_id = TableId::new(234);
460        let mapping = WorkerSlotMapping::new_uniform(
461            [slot1, slot2, slot3].into_iter(),
462            VirtualNode::COUNT_FOR_TEST,
463        );
464        let slot_bitmaps = mapping.to_bitmaps();
465        let serving_vnode_mapping = ServingVnodeMapping {
466            serving_vnode_mappings: RwLock::new(HashMap::from([(fragment_id, mapping)])),
467        };
468        let fragment_serving_infos = HashMap::from([(
469            fragment_id,
470            FragmentServingInfo {
471                result_table_id: Some(table_id),
472                distribution_type: FragmentDistributionType::Hash,
473                vnode_count: VirtualNode::COUNT_FOR_TEST,
474            },
475        )]);
476
477        let mappings = serving_vnode_mapping
478            .table_vnode_mappings_by_worker([worker1, worker2, worker3], &fragment_serving_infos);
479        let mut worker1_bitmap = slot_bitmaps[&slot1].clone();
480        worker1_bitmap |= &slot_bitmaps[&slot2];
481        assert_eq!(mappings[&worker1][&table_id], worker1_bitmap);
482        assert_eq!(mappings[&worker2][&table_id], slot_bitmaps[&slot3]);
483        assert!(mappings[&worker3].is_empty());
484    }
485
486    #[tokio::test]
487    async fn test_sync_serving_table_vnode_mappings_to_hummock_targets_only_result_tables() {
488        let notification_manager =
489            Arc::new(NotificationManager::new(SqlMetaStore::for_test().await).await);
490        let worker1 = serving_worker(1);
491        let worker2 = serving_worker(2);
492        let workers = vec![worker1.clone(), worker2.clone()];
493        let worker_key1 = WorkerKey(worker1.host.clone().unwrap());
494        let worker_key2 = WorkerKey(worker2.host.clone().unwrap());
495        let (tx1, mut rx1) = mpsc::unbounded_channel();
496        let (tx2, mut rx2) = mpsc::unbounded_channel();
497        notification_manager.insert_sender(SubscribeType::Hummock, worker_key1, tx1);
498        notification_manager.insert_sender(SubscribeType::Hummock, worker_key2, tx2);
499
500        let table_id = TableId::new(233);
501        let fragment_id = FragmentId::new(234);
502        let private_fragment_id = FragmentId::new(235);
503        let serving_vnode_mapping = Arc::new(ServingVnodeMapping::default());
504        let fragment_serving_infos = HashMap::from([
505            (
506                fragment_id,
507                FragmentServingInfo {
508                    result_table_id: Some(table_id),
509                    distribution_type: FragmentDistributionType::Hash,
510                    vnode_count: VirtualNode::COUNT_FOR_TEST,
511                },
512            ),
513            (
514                private_fragment_id,
515                FragmentServingInfo {
516                    result_table_id: None,
517                    distribution_type: FragmentDistributionType::Hash,
518                    vnode_count: VirtualNode::COUNT_FOR_TEST,
519                },
520            ),
521        ]);
522        let (upserted, failed) =
523            serving_vnode_mapping.upsert(&fragment_serving_infos, &workers, None);
524        assert_eq!(upserted.len(), 2);
525        assert!(failed.is_empty());
526        let expected_bitmaps = upserted[&fragment_id]
527            .to_bitmaps()
528            .into_iter()
529            .map(|(worker_slot_id, bitmap)| (worker_slot_id.worker_id(), bitmap))
530            .collect::<HashMap<_, _>>();
531        // Give the non-result fragment a distinct mapping. If it were incorrectly
532        // associated with `table_id`, it would expand the table's ownership to all vnodes.
533        serving_vnode_mapping.serving_vnode_mappings.write().insert(
534            private_fragment_id,
535            WorkerSlotMapping::new_single(WorkerSlotId::new(worker1.id, 0)),
536        );
537
538        sync_serving_table_vnode_mappings_to_hummock(
539            &notification_manager,
540            &serving_vnode_mapping,
541            &workers,
542            &fragment_serving_infos,
543        )
544        .await;
545
546        let response1 = rx1.recv().await.unwrap().unwrap();
547        let response2 = rx2.recv().await.unwrap().unwrap();
548        assert!(rx1.try_recv().is_err());
549        assert!(rx2.try_recv().is_err());
550
551        let serving_bitmap = |response: SubscribeResponse| {
552            let Some(Info::TableRefillRuntimeConfig(config)) = response.info else {
553                panic!("expect table refill runtime config");
554            };
555            let mappings = config.serving_table_vnode_mappings.unwrap().mappings;
556            assert_eq!(mappings.len(), 1);
557            assert_eq!(mappings[0].table_id, table_id.as_raw_id());
558            Bitmap::from(mappings[0].bitmap.clone().unwrap())
559        };
560        let bitmap1 = serving_bitmap(response1);
561        let bitmap2 = serving_bitmap(response2);
562        assert_eq!(bitmap1, expected_bitmaps[&worker1.id]);
563        assert_eq!(bitmap2, expected_bitmaps[&worker2.id]);
564        assert_eq!(
565            bitmap1.count_ones() + bitmap2.count_ones(),
566            VirtualNode::COUNT_FOR_TEST
567        );
568    }
569
570    fn hash_serving_info() -> FragmentServingInfo {
571        FragmentServingInfo {
572            result_table_id: None,
573            distribution_type: FragmentDistributionType::Hash,
574            vnode_count: VirtualNode::COUNT_FOR_TEST,
575        }
576    }
577
578    #[test]
579    fn test_reconcile_exactly_matches_full_snapshot_and_removes_failed_placements() {
580        let mapping = ServingVnodeMapping::default();
581        let worker = serving_worker(1);
582        let stale_fragment = FragmentId::new(1);
583        let retained_fragment = FragmentId::new(2);
584        let added_fragment = FragmentId::new(3);
585
586        let initial_snapshot = HashMap::from([
587            (stale_fragment, hash_serving_info()),
588            (retained_fragment, hash_serving_info()),
589        ]);
590        mapping.upsert(&initial_snapshot, std::slice::from_ref(&worker), None);
591        assert!(mapping.all().contains_key(&stale_fragment));
592        let retained_mapping = mapping.all()[&retained_fragment].clone();
593
594        let current_snapshot = HashMap::from([
595            (retained_fragment, hash_serving_info()),
596            (added_fragment, hash_serving_info()),
597        ]);
598        let (reconciled, failed) = mapping.reconcile(&current_snapshot, &[worker], None);
599        assert!(failed.is_empty());
600        assert_eq!(reconciled.len(), 2);
601        assert!(reconciled.contains_key(&retained_fragment));
602        assert!(reconciled.contains_key(&added_fragment));
603
604        let mappings = mapping.all();
605        assert_eq!(mappings.len(), 2);
606        assert!(mappings.contains_key(&retained_fragment));
607        assert!(mappings.contains_key(&added_fragment));
608        assert_eq!(mappings[&retained_fragment], retained_mapping);
609
610        let final_snapshot = HashMap::from([(retained_fragment, hash_serving_info())]);
611        let (reconciled, failed) = mapping.reconcile(&final_snapshot, &[], None);
612        assert!(reconciled.is_empty());
613        assert_eq!(failed, HashSet::from([retained_fragment]));
614        assert!(mapping.all().is_empty());
615    }
616
617    #[test]
618    fn test_upsert_keeps_fragments_absent_from_incremental_update() {
619        let mapping = ServingVnodeMapping::default();
620        let worker = serving_worker(1);
621        let first_fragment = FragmentId::new(1);
622        let second_fragment = FragmentId::new(2);
623
624        let first_update = HashMap::from([(first_fragment, hash_serving_info())]);
625        mapping.upsert(&first_update, std::slice::from_ref(&worker), None);
626        let second_update = HashMap::from([(second_fragment, hash_serving_info())]);
627        mapping.upsert(&second_update, &[worker], None);
628
629        let mappings = mapping.all();
630        assert!(mappings.contains_key(&first_fragment));
631        assert!(mappings.contains_key(&second_fragment));
632    }
633}
risingwave_meta/serving/mod.rs

risingwave_meta/serving/
mod.rs
