risingwave_meta/hummock/
compactor_manager.rs

1// Copyright 2022 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;
17use std::time::{Duration, Instant, SystemTime};
18
19use fail::fail_point;
20use parking_lot::RwLock;
21use risingwave_hummock_sdk::compact::statistics_compact_task;
22use risingwave_hummock_sdk::compact_task::CompactTask;
23use risingwave_hummock_sdk::{HummockCompactionTaskId, HummockContextId};
24use risingwave_pb::hummock::subscribe_compaction_event_response::Event as ResponseEvent;
25use risingwave_pb::hummock::{
26    CancelCompactTask, CompactTaskProgress, SubscribeCompactionEventResponse,
27};
28use risingwave_pb::iceberg_compaction::subscribe_iceberg_compaction_event_response::Event as IcebergResponseEvent;
29use risingwave_pb::iceberg_compaction::{
30    CancelCompactTask as IcebergCancelCompactTask, SubscribeIcebergCompactionEventResponse,
31};
32use risingwave_pb::id::IcebergCompactionTaskId;
33use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender};
34
35use crate::MetaResult;
36use crate::hummock::model::ext::compaction_task_model_to_assignment;
37use crate::manager::MetaSrvEnv;
38use crate::model::MetadataModelError;
39
40pub type CompactorManagerRef = Arc<CompactorManager>;
41pub type IcebergCompactorManagerRef = Arc<IcebergCompactorManager>;
42
43pub const TASK_RUN_TOO_LONG: &str = "running too long";
44pub const TASK_NOT_FOUND: &str = "task not found";
45pub const TASK_NORMAL: &str = "task is normal, please wait some time";
46
47type CompactorSubscribeStreamSender = UnboundedSender<MetaResult<SubscribeCompactionEventResponse>>;
48type CompactorSubscribeStreamReceiver =
49    UnboundedReceiver<MetaResult<SubscribeCompactionEventResponse>>;
50
51type IcebergCompactorSubscribeStreamSender =
52    UnboundedSender<MetaResult<SubscribeIcebergCompactionEventResponse>>;
53type IcebergCompactorSubscribeStreamReceiver =
54    UnboundedReceiver<MetaResult<SubscribeIcebergCompactionEventResponse>>;
55
56type CompactorSubscribeResponseEvent = ResponseEvent;
57
58type IcebergCompactorSubscribeResponseEvent = IcebergResponseEvent;
59
60/// Wraps the stream between meta node and compactor node.
61/// Compactor node will re-establish the stream when the previous one fails.
62pub struct Compactor {
63    context_id: HummockContextId,
64    sender: CompactorSubscribeStreamSender,
65}
66
67pub struct IcebergCompactor {
68    context_id: HummockContextId,
69    sender: IcebergCompactorSubscribeStreamSender,
70}
71
72struct TaskHeartbeat {
73    task: CompactTask,
74    num_ssts_sealed: u32,
75    num_ssts_uploaded: u32,
76    num_progress_key: u64,
77    num_pending_read_io: u64,
78    num_pending_write_io: u64,
79    create_time: Instant,
80    expire_at: u64,
81
82    update_at: u64,
83}
84
85impl Compactor {
86    pub fn new(context_id: HummockContextId, sender: CompactorSubscribeStreamSender) -> Self {
87        Self { context_id, sender }
88    }
89
90    pub fn send_event(&self, event: CompactorSubscribeResponseEvent) -> MetaResult<()> {
91        fail_point!("compaction_send_task_fail", |_| Err(anyhow::anyhow!(
92            "compaction_send_task_fail"
93        )
94        .into()));
95
96        self.sender
97            .send(Ok(SubscribeCompactionEventResponse {
98                create_at: SystemTime::now()
99                    .duration_since(SystemTime::UNIX_EPOCH)
100                    .expect("Clock may have gone backwards")
101                    .as_millis() as u64,
102                event: Some(event),
103            }))
104            .map_err(|e| anyhow::anyhow!(e))?;
105
106        Ok(())
107    }
108
109    pub fn cancel_task(&self, task_id: u64) -> MetaResult<()> {
110        self.send_event(ResponseEvent::CancelCompactTask(CancelCompactTask {
111            context_id: self.context_id,
112            task_id,
113        }))
114    }
115
116    pub fn cancel_tasks(&self, task_ids: &Vec<u64>) -> MetaResult<()> {
117        for task_id in task_ids {
118            self.cancel_task(*task_id)?;
119        }
120        Ok(())
121    }
122
123    pub fn context_id(&self) -> HummockContextId {
124        self.context_id
125    }
126}
127
128pub trait CompactorManagerTrait {
129    fn add_compactor(&self, context_id: HummockContextId) -> CompactorSubscribeStreamReceiver;
130    fn remove_compactor(&self, context_id: HummockContextId);
131    fn get_compactor(&self, context_id: HummockContextId) -> Option<Arc<Compactor>>;
132    fn next_compactor(&self) -> Option<Arc<Compactor>>;
133    fn compactor_num(&self) -> usize;
134}
135
136/// `CompactorManagerInner` maintains compactors which can process compact task.
137/// A compact task is tracked in `HummockManager::Compaction` via both `CompactStatus` and
138/// `CompactTaskAssignment`.
139///
140/// A compact task can be in one of these states:
141/// 1. Success: an assigned task is reported as success via `CompactStatus::report_compact_task`.
142///    It's the final state.
143/// 2. Failed: an Failed task is reported as success via `CompactStatus::report_compact_task`.
144///    It's the final state.
145/// 3. Cancelled: a task is reported as cancelled via `CompactStatus::report_compact_task`. It's
146///    the final state.
147pub struct CompactorManagerInner {
148    pub task_expired_seconds: u64,
149    pub heartbeat_expired_seconds: u64,
150    task_heartbeats: HashMap<HummockCompactionTaskId, TaskHeartbeat>,
151
152    /// The outer lock is a `RwLock`, so we should still be able to modify each compactor
153    pub compactor_map: HashMap<HummockContextId, Arc<Compactor>>,
154}
155
156impl CompactorManagerInner {
157    pub async fn with_meta(env: MetaSrvEnv) -> MetaResult<Self> {
158        use risingwave_meta_model::compaction_task;
159        use sea_orm::EntityTrait;
160        // Retrieve the existing task assignments from metastore.
161        let task_assignment: Vec<_> = compaction_task::Entity::find()
162            .all(&env.meta_store_ref().conn)
163            .await
164            .map_err(MetadataModelError::from)?
165            .into_iter()
166            .map(compaction_task_model_to_assignment)
167            .collect();
168        let mut manager = Self {
169            task_expired_seconds: env.opts.compaction_task_max_progress_interval_secs,
170            heartbeat_expired_seconds: env.opts.compaction_task_max_heartbeat_interval_secs,
171            task_heartbeats: Default::default(),
172            compactor_map: Default::default(),
173        };
174        // Initialize heartbeat for existing tasks.
175        task_assignment.into_iter().for_each(|assignment| {
176            manager.initiate_task_heartbeat(assignment.compact_task);
177        });
178        Ok(manager)
179    }
180
181    /// Only used for unit test.
182    pub fn for_test() -> Self {
183        Self {
184            task_expired_seconds: 1,
185            heartbeat_expired_seconds: 1,
186            task_heartbeats: Default::default(),
187            compactor_map: Default::default(),
188        }
189    }
190
191    pub fn next_compactor(&self) -> Option<Arc<Compactor>> {
192        use rand::Rng;
193        if self.compactor_map.is_empty() {
194            return None;
195        }
196
197        let rand_index = rand::rng().random_range(0..self.compactor_map.len());
198        let compactor = self.compactor_map.values().nth(rand_index).unwrap().clone();
199
200        Some(compactor)
201    }
202
203    /// Retrieve a receiver of tasks for the compactor identified by `context_id`. The sender should
204    /// be obtained by calling one of the compactor getters.
205    ///
206    ///  If `add_compactor` is called with the same `context_id` more than once, the only cause
207    /// would be compactor re-subscription, as `context_id` is a monotonically increasing
208    /// sequence.
209    pub fn add_compactor(
210        &mut self,
211        context_id: HummockContextId,
212    ) -> CompactorSubscribeStreamReceiver {
213        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
214
215        self.compactor_map
216            .insert(context_id, Arc::new(Compactor::new(context_id, tx)));
217
218        tracing::info!(context_id = %context_id, "Added compactor session");
219
220        rx
221    }
222
223    /// Used when meta exiting to support graceful shutdown.
224    pub fn abort_all_compactors(&mut self) {
225        while let Some(compactor) = self.next_compactor() {
226            self.remove_compactor(compactor.context_id);
227        }
228    }
229
230    pub fn remove_compactor(&mut self, context_id: HummockContextId) {
231        if self.compactor_map.remove(&context_id).is_some() {
232            tracing::info!(context_id = %context_id, "Removed compactor session")
233        };
234    }
235
236    pub fn get_compactor(&self, context_id: HummockContextId) -> Option<Arc<Compactor>> {
237        self.compactor_map.get(&context_id).cloned()
238    }
239
240    pub fn check_tasks_status(
241        &self,
242        tasks: &[HummockCompactionTaskId],
243        slow_task_duration: Duration,
244    ) -> HashMap<HummockCompactionTaskId, (Duration, &'static str)> {
245        let tasks_ids: HashSet<u64> = HashSet::from_iter(tasks.to_vec());
246        let mut ret = HashMap::default();
247        for TaskHeartbeat {
248            task, create_time, ..
249        } in self.task_heartbeats.values()
250        {
251            if !tasks_ids.contains(&task.task_id) {
252                continue;
253            }
254            let pending_time = create_time.elapsed();
255            if pending_time > slow_task_duration {
256                ret.insert(task.task_id, (pending_time, TASK_RUN_TOO_LONG));
257            } else {
258                ret.insert(task.task_id, (pending_time, TASK_NORMAL));
259            }
260        }
261
262        for task_id in tasks {
263            if !ret.contains_key(task_id) {
264                ret.insert(*task_id, (Duration::from_secs(0), TASK_NOT_FOUND));
265            }
266        }
267        ret
268    }
269
270    pub fn get_heartbeat_expired_tasks(&self) -> Vec<CompactTask> {
271        let heartbeat_expired_ts: u64 = SystemTime::now()
272            .duration_since(SystemTime::UNIX_EPOCH)
273            .expect("Clock may have gone backwards")
274            .as_secs()
275            - self.heartbeat_expired_seconds;
276        Self::get_heartbeat_expired_tasks_impl(&self.task_heartbeats, heartbeat_expired_ts)
277    }
278
279    fn get_heartbeat_expired_tasks_impl(
280        task_heartbeats: &HashMap<HummockCompactionTaskId, TaskHeartbeat>,
281        heartbeat_expired_ts: u64,
282    ) -> Vec<CompactTask> {
283        let mut cancellable_tasks = vec![];
284        const MAX_TASK_DURATION_SEC: u64 = 2700;
285
286        for TaskHeartbeat {
287            expire_at,
288            task,
289            create_time,
290            num_ssts_sealed,
291            num_ssts_uploaded,
292            num_progress_key,
293            num_pending_read_io,
294            num_pending_write_io,
295            update_at,
296        } in task_heartbeats.values()
297        {
298            if *update_at < heartbeat_expired_ts {
299                cancellable_tasks.push(task.clone());
300            }
301
302            let task_duration_too_long = create_time.elapsed().as_secs() > MAX_TASK_DURATION_SEC;
303            if task_duration_too_long {
304                let compact_task_statistics = statistics_compact_task(task);
305                tracing::info!(
306                    "CompactionGroupId {} Task {} duration too long create_time {:?} expire_at {:?} num_ssts_sealed {} num_ssts_uploaded {} num_progress_key {} \
307                        pending_read_io_count {} pending_write_io_count {} target_level {} \
308                        base_level {} target_sub_level_id {} task_type {} compact_task_statistics {:?}",
309                    task.compaction_group_id,
310                    task.task_id,
311                    create_time,
312                    expire_at,
313                    num_ssts_sealed,
314                    num_ssts_uploaded,
315                    num_progress_key,
316                    num_pending_read_io,
317                    num_pending_write_io,
318                    task.target_level,
319                    task.base_level,
320                    task.target_sub_level_id,
321                    task.task_type.as_str_name(),
322                    compact_task_statistics
323                );
324            }
325        }
326        cancellable_tasks
327    }
328
329    pub fn initiate_task_heartbeat(&mut self, task: CompactTask) {
330        let now = SystemTime::now()
331            .duration_since(SystemTime::UNIX_EPOCH)
332            .expect("Clock may have gone backwards")
333            .as_secs();
334        self.task_heartbeats.insert(
335            task.task_id,
336            TaskHeartbeat {
337                task,
338                num_ssts_sealed: 0,
339                num_ssts_uploaded: 0,
340                num_progress_key: 0,
341                num_pending_read_io: 0,
342                num_pending_write_io: 0,
343                create_time: Instant::now(),
344                expire_at: now + self.task_expired_seconds,
345                update_at: now,
346            },
347        );
348    }
349
350    pub fn remove_task_heartbeat(&mut self, task_id: u64) {
351        self.task_heartbeats.remove(&task_id).unwrap();
352    }
353
354    pub fn update_task_heartbeats(
355        &mut self,
356        progress_list: &Vec<CompactTaskProgress>,
357    ) -> Vec<CompactTask> {
358        let now = SystemTime::now()
359            .duration_since(SystemTime::UNIX_EPOCH)
360            .expect("Clock may have gone backwards")
361            .as_secs();
362        let mut cancel_tasks = vec![];
363        for progress in progress_list {
364            if let Some(task_ref) = self.task_heartbeats.get_mut(&progress.task_id) {
365                task_ref.update_at = now;
366
367                if task_ref.num_ssts_sealed < progress.num_ssts_sealed
368                    || task_ref.num_ssts_uploaded < progress.num_ssts_uploaded
369                    || task_ref.num_progress_key < progress.num_progress_key
370                {
371                    // Refresh the expired of the task as it is showing progress.
372                    task_ref.expire_at = now + self.task_expired_seconds;
373                    task_ref.num_ssts_sealed = progress.num_ssts_sealed;
374                    task_ref.num_ssts_uploaded = progress.num_ssts_uploaded;
375                    task_ref.num_progress_key = progress.num_progress_key;
376                }
377                task_ref.num_pending_read_io = progress.num_pending_read_io;
378                task_ref.num_pending_write_io = progress.num_pending_write_io;
379
380                // timeout check
381                if task_ref.expire_at < now {
382                    // cancel
383                    cancel_tasks.push(task_ref.task.clone())
384                }
385            }
386        }
387
388        cancel_tasks
389    }
390
391    pub fn compactor_num(&self) -> usize {
392        self.compactor_map.len()
393    }
394
395    pub fn get_progress(&self) -> Vec<CompactTaskProgress> {
396        self.task_heartbeats
397            .values()
398            .map(|hb| CompactTaskProgress {
399                task_id: hb.task.task_id,
400                num_ssts_sealed: hb.num_ssts_sealed,
401                num_ssts_uploaded: hb.num_ssts_uploaded,
402                num_progress_key: hb.num_progress_key,
403                num_pending_read_io: hb.num_pending_read_io,
404                num_pending_write_io: hb.num_pending_write_io,
405                compaction_group_id: Some(hb.task.compaction_group_id),
406            })
407            .collect()
408    }
409}
410
411pub struct CompactorManager {
412    inner: Arc<RwLock<CompactorManagerInner>>,
413}
414
415impl CompactorManager {
416    pub async fn with_meta(env: MetaSrvEnv) -> MetaResult<Self> {
417        let inner = CompactorManagerInner::with_meta(env).await?;
418
419        Ok(Self {
420            inner: Arc::new(RwLock::new(inner)),
421        })
422    }
423
424    /// Only used for unit test.
425    pub fn for_test() -> Self {
426        let inner = CompactorManagerInner::for_test();
427        Self {
428            inner: Arc::new(RwLock::new(inner)),
429        }
430    }
431
432    pub fn next_compactor(&self) -> Option<Arc<Compactor>> {
433        self.inner.read().next_compactor()
434    }
435
436    pub fn add_compactor(&self, context_id: HummockContextId) -> CompactorSubscribeStreamReceiver {
437        self.inner.write().add_compactor(context_id)
438    }
439
440    pub fn abort_all_compactors(&self) {
441        self.inner.write().abort_all_compactors();
442    }
443
444    pub fn remove_compactor(&self, context_id: HummockContextId) {
445        self.inner.write().remove_compactor(context_id)
446    }
447
448    pub fn get_compactor(&self, context_id: HummockContextId) -> Option<Arc<Compactor>> {
449        self.inner.read().get_compactor(context_id)
450    }
451
452    pub fn check_tasks_status(
453        &self,
454        tasks: &[HummockCompactionTaskId],
455        slow_task_duration: Duration,
456    ) -> HashMap<HummockCompactionTaskId, (Duration, &'static str)> {
457        self.inner
458            .read()
459            .check_tasks_status(tasks, slow_task_duration)
460    }
461
462    pub fn get_heartbeat_expired_tasks(&self) -> Vec<CompactTask> {
463        self.inner.read().get_heartbeat_expired_tasks()
464    }
465
466    pub fn initiate_task_heartbeat(&self, task: CompactTask) {
467        self.inner.write().initiate_task_heartbeat(task);
468    }
469
470    pub fn remove_task_heartbeat(&self, task_id: u64) {
471        self.inner.write().remove_task_heartbeat(task_id);
472    }
473
474    pub fn update_task_heartbeats(
475        &self,
476        progress_list: &Vec<CompactTaskProgress>,
477    ) -> Vec<CompactTask> {
478        self.inner.write().update_task_heartbeats(progress_list)
479    }
480
481    pub fn compactor_num(&self) -> usize {
482        self.inner.read().compactor_num()
483    }
484
485    pub fn get_progress(&self) -> Vec<CompactTaskProgress> {
486        self.inner.read().get_progress()
487    }
488}
489
490impl IcebergCompactor {
491    pub fn new(
492        context_id: HummockContextId,
493        sender: IcebergCompactorSubscribeStreamSender,
494    ) -> Self {
495        Self { context_id, sender }
496    }
497
498    pub fn context_id(&self) -> HummockContextId {
499        self.context_id
500    }
501
502    pub fn send_event(&self, event: IcebergCompactorSubscribeResponseEvent) -> MetaResult<()> {
503        fail_point!("iceberg_compaction_send_task_fail", |_| Err(
504            anyhow::anyhow!("iceberg_compaction_send_task_fail").into()
505        ));
506
507        self.sender
508            .send(Ok(SubscribeIcebergCompactionEventResponse {
509                create_at: SystemTime::now()
510                    .duration_since(SystemTime::UNIX_EPOCH)
511                    .expect("Clock may have gone backwards")
512                    .as_millis() as u64,
513                event: Some(event),
514            }))
515            .map_err(|e| anyhow::anyhow!(e))?;
516
517        Ok(())
518    }
519
520    pub fn cancel_task(&self, task_id: IcebergCompactionTaskId) -> MetaResult<()> {
521        self.send_event(IcebergResponseEvent::CancelCompactTask(
522            IcebergCancelCompactTask { task_id },
523        ))
524    }
525}
526
527pub struct IcebergCompactorManagerInner {
528    pub compactor_map: HashMap<HummockContextId, Arc<IcebergCompactor>>,
529}
530
531pub struct IcebergCompactorManager {
532    inner: Arc<RwLock<IcebergCompactorManagerInner>>,
533}
534
535impl IcebergCompactorManager {
536    pub fn new() -> Self {
537        Self {
538            inner: Arc::new(RwLock::new(IcebergCompactorManagerInner {
539                compactor_map: HashMap::new(),
540            })),
541        }
542    }
543
544    pub fn add_compactor(
545        &self,
546        context_id: HummockContextId,
547    ) -> IcebergCompactorSubscribeStreamReceiver {
548        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
549        self.inner
550            .write()
551            .compactor_map
552            .insert(context_id, Arc::new(IcebergCompactor::new(context_id, tx)));
553        tracing::info!(context_id = %context_id, "Added iceberg compactor session");
554        rx
555    }
556
557    pub fn remove_compactor(&self, context_id: HummockContextId) {
558        if self
559            .inner
560            .write()
561            .compactor_map
562            .remove(&context_id)
563            .is_some()
564        {
565            tracing::info!(context_id = %context_id, "Removed iceberg compactor session");
566        }
567    }
568
569    pub fn get_compactor(&self, context_id: HummockContextId) -> Option<Arc<IcebergCompactor>> {
570        self.inner.read().compactor_map.get(&context_id).cloned()
571    }
572
573    pub fn next_compactor(&self) -> Option<Arc<IcebergCompactor>> {
574        use rand::Rng;
575        let compactor_map = &self.inner.read().compactor_map;
576        if compactor_map.is_empty() {
577            return None;
578        }
579        let rand_index = rand::rng().random_range(0..compactor_map.len());
580        compactor_map.values().nth(rand_index).cloned()
581    }
582
583    pub fn compactor_num(&self) -> usize {
584        self.inner.read().compactor_map.len()
585    }
586}
587
588#[cfg(test)]
589mod tests {
590    use std::sync::Arc;
591    use std::time::Duration;
592
593    use risingwave_hummock_sdk::compaction_group::StaticCompactionGroupId;
594    use risingwave_pb::hummock::CompactTaskProgress;
595    use risingwave_pb::iceberg_compaction::subscribe_iceberg_compaction_event_response::Event as IcebergResponseEvent;
596    use risingwave_rpc_client::HummockMetaClient;
597
598    use crate::hummock::compaction::selector::default_compaction_selector;
599    use crate::hummock::test_utils::{
600        add_ssts, register_table_ids_to_compaction_group, setup_compute_env,
601    };
602    use crate::hummock::{CompactorManager, IcebergCompactorManager, MockHummockMetaClient};
603
604    #[tokio::test]
605    async fn test_compactor_manager() {
606        // Initialize metastore with task assignment.
607        let (env, context_id) = {
608            let (env, hummock_manager, _cluster_manager, worker_id) = setup_compute_env(80).await;
609            let context_id = worker_id as _;
610            let hummock_meta_client: Arc<dyn HummockMetaClient> = Arc::new(
611                MockHummockMetaClient::new(hummock_manager.clone(), context_id),
612            );
613            let compactor_manager = hummock_manager.compactor_manager.clone();
614            register_table_ids_to_compaction_group(
615                hummock_manager.as_ref(),
616                &[1],
617                StaticCompactionGroupId::StateDefault,
618            )
619            .await;
620            let _sst_infos =
621                add_ssts(1, hummock_manager.as_ref(), hummock_meta_client.clone()).await;
622            let _receiver = compactor_manager.add_compactor(context_id);
623            hummock_manager
624                .get_compact_task(
625                    StaticCompactionGroupId::StateDefault,
626                    &mut *default_compaction_selector(),
627                )
628                .await
629                .unwrap()
630                .unwrap();
631            (env, context_id)
632        };
633
634        // Restart. Set task_expired_seconds to 0 only to speed up test.
635        let compactor_manager = CompactorManager::with_meta(env).await.unwrap();
636        // Because task assignment exists.
637        // Because compactor gRPC is not established yet.
638        assert_eq!(compactor_manager.compactor_num(), 0);
639        assert!(compactor_manager.get_compactor(context_id).is_none());
640
641        // Ensure task is expired.
642        tokio::time::sleep(Duration::from_secs(2)).await;
643        let expired = compactor_manager.get_heartbeat_expired_tasks();
644        assert_eq!(expired.len(), 1);
645
646        // Mimic no-op compaction heartbeat
647        assert_eq!(compactor_manager.get_heartbeat_expired_tasks().len(), 1);
648
649        // Mimic compaction heartbeat with invalid task id
650        compactor_manager.update_task_heartbeats(&vec![CompactTaskProgress {
651            task_id: expired[0].task_id + 1,
652            num_ssts_sealed: 1,
653            num_ssts_uploaded: 1,
654            num_progress_key: 100,
655            ..Default::default()
656        }]);
657        assert_eq!(compactor_manager.get_heartbeat_expired_tasks().len(), 1);
658
659        // Mimic effective compaction heartbeat
660        compactor_manager.update_task_heartbeats(&vec![CompactTaskProgress {
661            task_id: expired[0].task_id,
662            num_ssts_sealed: 1,
663            num_ssts_uploaded: 1,
664            num_progress_key: 100,
665            ..Default::default()
666        }]);
667        assert_eq!(compactor_manager.get_heartbeat_expired_tasks().len(), 0);
668
669        // Test add
670        assert_eq!(compactor_manager.compactor_num(), 0);
671        assert!(compactor_manager.get_compactor(context_id).is_none());
672        compactor_manager.add_compactor(context_id);
673        assert_eq!(compactor_manager.compactor_num(), 1);
674        assert_eq!(
675            compactor_manager
676                .get_compactor(context_id)
677                .unwrap()
678                .context_id(),
679            context_id
680        );
681        // Test remove
682        compactor_manager.remove_compactor(context_id);
683        assert_eq!(compactor_manager.compactor_num(), 0);
684        assert!(compactor_manager.get_compactor(context_id).is_none());
685    }
686
687    #[test]
688    fn test_iceberg_compactor_manager() {
689        // Test Add and Remove Iceberg Compactor
690        let iceberg_context_id = 1000.into();
691        let iceberg_compactor_manager = IcebergCompactorManager::new();
692        assert_eq!(iceberg_compactor_manager.compactor_num(), 0);
693        assert!(
694            iceberg_compactor_manager
695                .get_compactor(iceberg_context_id)
696                .is_none()
697        );
698        let mut receiver = iceberg_compactor_manager.add_compactor(iceberg_context_id);
699        assert_eq!(iceberg_compactor_manager.compactor_num(), 1);
700        let compactor = iceberg_compactor_manager
701            .get_compactor(iceberg_context_id)
702            .unwrap();
703        assert_eq!(compactor.context_id(), iceberg_context_id);
704
705        compactor.cancel_task(42.into()).unwrap();
706        let event = receiver.try_recv().unwrap().unwrap().event.unwrap();
707        match event {
708            IcebergResponseEvent::CancelCompactTask(cancel_task) => {
709                assert_eq!(cancel_task.task_id, 42);
710            }
711            other => panic!("unexpected iceberg compactor event: {other:?}"),
712        }
713
714        // Test remove
715        iceberg_compactor_manager.remove_compactor(iceberg_context_id);
716        assert_eq!(iceberg_compactor_manager.compactor_num(), 0);
717        assert!(
718            iceberg_compactor_manager
719                .get_compactor(iceberg_context_id)
720                .is_none()
721        );
722    }
723}
risingwave_meta/hummock/compactor_manager.rs

risingwave_meta/hummock/
compactor_manager.rs
