Skip to main content

risingwave_meta/hummock/manager/
transaction.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};
16use std::ops::{Deref, DerefMut};
17use std::sync::Arc;
18
19use parking_lot::Mutex;
20use risingwave_common::catalog::TableId;
21use risingwave_hummock_sdk::change_log::{EpochNewChangeLog, TableChangeLog};
22use risingwave_hummock_sdk::compaction_group::StaticCompactionGroupId;
23use risingwave_hummock_sdk::sstable_info::SstableInfo;
24use risingwave_hummock_sdk::table_watermark::TableWatermarks;
25use risingwave_hummock_sdk::vector_index::VectorIndexDelta;
26use risingwave_hummock_sdk::version::{GroupDelta, HummockVersion, HummockVersionDelta};
27use risingwave_hummock_sdk::{
28    CompactionGroupId, FrontendHummockVersionDelta, HummockSstableId, HummockVersionId,
29};
30use risingwave_pb::hummock::{
31    CompatibilityVersion, GroupConstruct, HummockVersionDeltas, HummockVersionStats,
32    StateTableInfoDelta,
33};
34use risingwave_pb::meta::subscribe_response::{Info, Operation};
35use sea_orm::{ConnectionTrait, EntityTrait};
36
37use super::TableCommittedEpochNotifiers;
38use crate::hummock::model::CompactionGroup;
39use crate::hummock::model::ext::to_table_change_log_meta_store_model;
40use crate::manager::{MetaOpts, NotificationManager};
41use crate::model::{
42    InMemValTransaction, MetadataModelResult, Transactional, ValTransaction, VarTransaction,
43};
44use crate::rpc::metrics::MetaMetrics;
45
46fn trigger_delta_log_stats(metrics: &MetaMetrics, total_number: usize) {
47    metrics.delta_log_count.set(total_number as _);
48}
49
50#[derive(Default)]
51struct TableChangeLogTransactionDelta {
52    updates: Vec<HashMap<TableId, EpochNewChangeLog>>,
53    delete_all: HashSet<TableId>,
54}
55
56pub(super) fn trigger_version_stat(metrics: &MetaMetrics, current_version: &HummockVersion) {
57    metrics
58        .version_size
59        .set(current_version.estimated_encode_len() as i64);
60    metrics
61        .current_version_id
62        .set(current_version.id.as_i64_id());
63}
64
65pub(super) struct HummockVersionTransaction<'a> {
66    orig_version: &'a mut Arc<HummockVersion>,
67    orig_deltas: &'a mut BTreeMap<HummockVersionId, HummockVersionDelta>,
68    orig_table_change_log: &'a mut HashMap<TableId, TableChangeLog>,
69    notification_manager: &'a NotificationManager,
70    table_committed_epoch_notifiers: Option<&'a Mutex<TableCommittedEpochNotifiers>>,
71    meta_metrics: &'a MetaMetrics,
72    version_stat_tx: &'a tokio::sync::mpsc::UnboundedSender<Arc<HummockVersion>>,
73
74    pre_applied_version: Option<(
75        HummockVersion,
76        Vec<HummockVersionDelta>,
77        TableChangeLogTransactionDelta,
78    )>,
79    disable_apply_to_txn: bool,
80    opts: &'a MetaOpts,
81}
82
83impl<'a> HummockVersionTransaction<'a> {
84    pub(super) fn new(
85        version: &'a mut Arc<HummockVersion>,
86        deltas: &'a mut BTreeMap<HummockVersionId, HummockVersionDelta>,
87        table_change_log: &'a mut HashMap<TableId, TableChangeLog>,
88        notification_manager: &'a NotificationManager,
89        table_committed_epoch_notifiers: Option<&'a Mutex<TableCommittedEpochNotifiers>>,
90        meta_metrics: &'a MetaMetrics,
91        opts: &'a MetaOpts,
92        version_stat_tx: &'a tokio::sync::mpsc::UnboundedSender<Arc<HummockVersion>>,
93    ) -> Self {
94        Self {
95            orig_version: version,
96            orig_deltas: deltas,
97            orig_table_change_log: table_change_log,
98            pre_applied_version: None,
99            disable_apply_to_txn: false,
100            notification_manager,
101            table_committed_epoch_notifiers,
102            meta_metrics,
103            opts,
104            version_stat_tx,
105        }
106    }
107
108    pub(super) fn disable_apply_to_txn(&mut self) {
109        assert!(
110            self.pre_applied_version.is_none(),
111            "should only call disable at the beginning of txn"
112        );
113        self.disable_apply_to_txn = true;
114    }
115
116    pub(super) fn latest_version(&self) -> &HummockVersion {
117        if let Some((version, _, _)) = &self.pre_applied_version {
118            version
119        } else {
120            self.orig_version.as_ref()
121        }
122    }
123
124    pub(super) fn new_delta<'b>(&'b mut self) -> SingleDeltaTransaction<'a, 'b> {
125        let delta = self.latest_version().version_delta_after();
126        SingleDeltaTransaction {
127            version_txn: self,
128            delta: Some((delta, HashMap::new())),
129        }
130    }
131
132    fn pre_apply(
133        &mut self,
134        delta: HummockVersionDelta,
135        change_log_updates: HashMap<TableId, EpochNewChangeLog>,
136    ) {
137        let (version, deltas, table_change_log_delta) =
138            self.pre_applied_version.get_or_insert_with(|| {
139                (
140                    self.orig_version.as_ref().clone(),
141                    Vec::with_capacity(1),
142                    TableChangeLogTransactionDelta::default(),
143                )
144            });
145        let changed_table_info = version.apply_version_delta(&delta);
146        // The in-memory change logs are only updated after the metastore transaction succeeds, so
147        // complete deletion is derived from the original state. A deletion that becomes eligible
148        // after multiple deltas in one transaction may be deferred to the next transaction.
149        let delete_all = HummockVersion::collect_gc_change_log_delta(
150            self.orig_table_change_log.keys(),
151            &change_log_updates,
152            &delta.removed_table_ids,
153            &delta.state_table_info_delta,
154            &changed_table_info,
155        );
156        table_change_log_delta.delete_all.extend(delete_all);
157        if !change_log_updates.is_empty() {
158            table_change_log_delta.updates.push(change_log_updates);
159        }
160        deltas.push(delta);
161    }
162
163    /// Returns a duplicate delta, used by time travel.
164    pub(super) fn pre_commit_epoch(
165        &mut self,
166        tables_to_commit: &HashMap<TableId, u64>,
167        new_compaction_groups: Vec<CompactionGroup>,
168        group_id_to_sub_levels: BTreeMap<CompactionGroupId, Vec<Vec<SstableInfo>>>,
169        new_table_ids: &HashMap<TableId, CompactionGroupId>,
170        new_table_watermarks: HashMap<TableId, TableWatermarks>,
171        change_log_delta: HashMap<TableId, EpochNewChangeLog>,
172        vector_index_delta: HashMap<TableId, VectorIndexDelta>,
173        group_id_to_truncate_tables: HashMap<CompactionGroupId, HashSet<TableId>>,
174    ) -> HummockVersionDelta {
175        let mut new_version_delta = self.new_delta();
176        new_version_delta.new_table_watermarks = new_table_watermarks;
177        new_version_delta.set_change_log_delta(change_log_delta);
178        new_version_delta.vector_index_delta = vector_index_delta;
179
180        for compaction_group in &new_compaction_groups {
181            let group_deltas = &mut new_version_delta
182                .group_deltas
183                .entry(compaction_group.group_id())
184                .or_default()
185                .group_deltas;
186
187            #[expect(deprecated)]
188            group_deltas.push(GroupDelta::GroupConstruct(Box::new(GroupConstruct {
189                group_config: Some(compaction_group.compaction_config().as_ref().clone()),
190                group_id: compaction_group.group_id(),
191                parent_group_id: StaticCompactionGroupId::NewCompactionGroup as CompactionGroupId,
192                new_sst_start_id: HummockSstableId::default(), // No need to set it when `NewCompactionGroup`
193                table_ids: vec![],
194                version: CompatibilityVersion::LATEST as _,
195                split_key: None,
196            })));
197        }
198
199        // Append SSTs to a new version.
200        for (compaction_group_id, sub_levels) in group_id_to_sub_levels {
201            let group_deltas = &mut new_version_delta
202                .group_deltas
203                .entry(compaction_group_id)
204                .or_default()
205                .group_deltas;
206
207            for sub_level in sub_levels {
208                group_deltas.push(GroupDelta::NewL0SubLevel(sub_level));
209            }
210        }
211
212        for (compaction_group_id, table_ids) in group_id_to_truncate_tables {
213            let group_deltas = &mut new_version_delta
214                .group_deltas
215                .entry(compaction_group_id)
216                .or_default()
217                .group_deltas;
218
219            group_deltas.push(GroupDelta::PruneTableIdsFromSsts(
220                table_ids.into_iter().collect(),
221            ));
222        }
223
224        // update state table info
225        new_version_delta.with_latest_version(|version, delta| {
226            for (table_id, cg_id) in new_table_ids {
227                assert!(
228                    !version.state_table_info.info().contains_key(table_id),
229                    "newly added table exists previously: {:?}",
230                    table_id
231                );
232                let committed_epoch = *tables_to_commit.get(table_id).expect("newly added table must exist in tables_to_commit");
233                delta.state_table_info_delta.insert(
234                    *table_id,
235                    StateTableInfoDelta {
236                        committed_epoch,
237                        compaction_group_id: *cg_id,
238                    },
239                );
240            }
241
242            for (table_id, committed_epoch) in tables_to_commit {
243                if new_table_ids.contains_key(table_id) {
244                    continue;
245                }
246                let info = version.state_table_info.info().get(table_id).unwrap_or_else(|| {
247                    panic!("tables_to_commit {:?} contains table_id {} that is not newly added but not exists previously", tables_to_commit, table_id);
248                });
249                assert!(delta
250                    .state_table_info_delta
251                    .insert(
252                        *table_id,
253                        StateTableInfoDelta {
254                            committed_epoch: *committed_epoch,
255                            compaction_group_id: info.compaction_group_id,
256                        }
257                    )
258                    .is_none());
259            }
260        });
261
262        let time_travel_delta = (*new_version_delta).clone();
263        new_version_delta.pre_apply();
264        time_travel_delta
265    }
266}
267
268impl InMemValTransaction for HummockVersionTransaction<'_> {
269    fn commit(self) {
270        if let Some((version, deltas, table_change_log_delta)) = self.pre_applied_version {
271            *self.orig_version = Arc::new(version);
272            for change_log_delta in table_change_log_delta.updates {
273                HummockVersion::apply_change_log_delta(
274                    self.orig_table_change_log,
275                    &change_log_delta,
276                );
277            }
278            self.orig_table_change_log
279                .retain(|table_id, _| !table_change_log_delta.delete_all.contains(table_id));
280
281            if !self.disable_apply_to_txn {
282                let pb_deltas = deltas.iter().map(|delta| delta.to_protobuf()).collect();
283                self.notification_manager.notify_hummock_without_version(
284                    Operation::Add,
285                    Info::HummockVersionDeltas(risingwave_pb::hummock::HummockVersionDeltas {
286                        version_deltas: pb_deltas,
287                    }),
288                );
289                self.notification_manager.notify_frontend_without_version(
290                    Operation::Update,
291                    Info::HummockVersionDeltas(HummockVersionDeltas {
292                        version_deltas: deltas
293                            .iter()
294                            .map(|delta| {
295                                FrontendHummockVersionDelta::from_delta(delta).to_protobuf()
296                            })
297                            .collect(),
298                    }),
299                );
300                if let Some(table_committed_epoch_notifiers) = self.table_committed_epoch_notifiers
301                {
302                    table_committed_epoch_notifiers
303                        .lock()
304                        .notify_deltas(&deltas);
305                }
306            }
307
308            for delta in deltas {
309                assert!(self.orig_deltas.insert(delta.id, delta.clone()).is_none());
310            }
311
312            trigger_delta_log_stats(self.meta_metrics, self.orig_deltas.len());
313            let _ = self.version_stat_tx.send(self.orig_version.clone());
314        }
315    }
316}
317
318impl<TXN> ValTransaction<TXN> for HummockVersionTransaction<'_>
319where
320    TXN: ConnectionTrait,
321    HummockVersionDelta: Transactional<TXN>,
322    HummockVersionStats: Transactional<TXN>,
323{
324    async fn apply_to_txn(&self, txn: &mut TXN) -> MetadataModelResult<()> {
325        if self.disable_apply_to_txn {
326            return Ok(());
327        }
328        if let Some((_, deltas, table_change_log_delta)) = &self.pre_applied_version {
329            // These upsert_in_transaction can be batched. However, we know len(deltas) is always 1 currently.
330            for delta in deltas {
331                delta.upsert_in_transaction(txn).await?;
332            }
333
334            let insert_batch_size = self.opts.table_change_log_insert_batch_size as usize;
335            use futures::stream::{self, StreamExt};
336            use sea_orm::{ColumnTrait, Condition, QueryFilter};
337            let insert_iter = table_change_log_delta
338                .updates
339                .iter()
340                .flat_map(|updates| updates.iter())
341                .map(|(table_id, new_log)| (*table_id, new_log));
342            let mut stream = stream::iter(insert_iter).chunks(insert_batch_size);
343            while let Some(change_log_batch) = stream.next().await {
344                let insert_many = change_log_batch
345                    .into_iter()
346                    .map(|(table_id, change_log)| {
347                        to_table_change_log_meta_store_model(table_id, change_log)
348                    })
349                    .collect::<Vec<_>>();
350                risingwave_meta_model::hummock_table_change_log::Entity::insert_many(insert_many)
351                    .on_empty_do_nothing()
352                    .exec(txn)
353                    .await?;
354            }
355
356            let delete_batch_size = self.opts.table_change_log_delete_batch_size as usize;
357            let delete_iter = table_change_log_delta.delete_all.iter().copied();
358
359            let mut stream = stream::iter(delete_iter).chunks(delete_batch_size);
360            while let Some(change_log_batch) = stream.next().await {
361                let mut condition = Condition::any();
362                for table_id in change_log_batch {
363                    let table_condition = Condition::all().add(
364                        risingwave_meta_model::hummock_table_change_log::Column::TableId
365                            .eq(table_id),
366                    );
367                    condition = condition.add(table_condition);
368                }
369                risingwave_meta_model::hummock_table_change_log::Entity::delete_many()
370                    .filter(condition)
371                    .exec(txn)
372                    .await?;
373            }
374        }
375        Ok(())
376    }
377}
378
379pub(super) struct SingleDeltaTransaction<'a, 'b> {
380    version_txn: &'b mut HummockVersionTransaction<'a>,
381    delta: Option<(HummockVersionDelta, HashMap<TableId, EpochNewChangeLog>)>,
382}
383
384impl SingleDeltaTransaction<'_, '_> {
385    pub(super) fn latest_version(&self) -> &HummockVersion {
386        self.version_txn.latest_version()
387    }
388
389    fn set_change_log_delta(&mut self, change_log_delta: HashMap<TableId, EpochNewChangeLog>) {
390        self.delta.as_mut().expect("should exist").1 = change_log_delta;
391    }
392
393    pub(super) fn pre_apply(mut self) {
394        let (delta, change_log_delta) = self.delta.take().unwrap();
395        self.version_txn.pre_apply(delta, change_log_delta);
396    }
397
398    pub(super) fn with_latest_version(
399        &mut self,
400        f: impl FnOnce(&HummockVersion, &mut HummockVersionDelta),
401    ) {
402        f(
403            self.version_txn.latest_version(),
404            &mut self.delta.as_mut().expect("should exist").0,
405        )
406    }
407}
408
409impl Deref for SingleDeltaTransaction<'_, '_> {
410    type Target = HummockVersionDelta;
411
412    fn deref(&self) -> &Self::Target {
413        &self.delta.as_ref().expect("should exist").0
414    }
415}
416
417impl DerefMut for SingleDeltaTransaction<'_, '_> {
418    fn deref_mut(&mut self) -> &mut Self::Target {
419        &mut self.delta.as_mut().expect("should exist").0
420    }
421}
422
423impl Drop for SingleDeltaTransaction<'_, '_> {
424    fn drop(&mut self) {
425        if let Some((delta, change_log_delta)) = self.delta.take() {
426            self.version_txn.pre_apply(delta, change_log_delta);
427        }
428    }
429}
430
431pub(super) struct HummockVersionStatsTransaction<'a> {
432    stats: VarTransaction<'a, HummockVersionStats>,
433    notification_manager: &'a NotificationManager,
434}
435
436impl<'a> HummockVersionStatsTransaction<'a> {
437    pub(super) fn new(
438        stats: &'a mut HummockVersionStats,
439        notification_manager: &'a NotificationManager,
440    ) -> Self {
441        Self {
442            stats: VarTransaction::new(stats),
443            notification_manager,
444        }
445    }
446}
447
448impl InMemValTransaction for HummockVersionStatsTransaction<'_> {
449    fn commit(self) {
450        if self.stats.has_new_value() {
451            let stats = self.stats.clone();
452            self.stats.commit();
453            self.notification_manager
454                .notify_frontend_without_version(Operation::Update, Info::HummockStats(stats));
455        }
456    }
457}
458
459impl<TXN> ValTransaction<TXN> for HummockVersionStatsTransaction<'_>
460where
461    TXN: ConnectionTrait,
462    HummockVersionStats: Transactional<TXN>,
463{
464    async fn apply_to_txn(&self, txn: &mut TXN) -> MetadataModelResult<()> {
465        self.stats.apply_to_txn(txn).await
466    }
467}
468
469impl Deref for HummockVersionStatsTransaction<'_> {
470    type Target = HummockVersionStats;
471
472    fn deref(&self) -> &Self::Target {
473        self.stats.deref()
474    }
475}
476
477impl DerefMut for HummockVersionStatsTransaction<'_> {
478    fn deref_mut(&mut self) -> &mut Self::Target {
479        self.stats.deref_mut()
480    }
481}
482
483#[cfg(test)]
484mod tests {
485    use itertools::Itertools;
486    use risingwave_hummock_sdk::change_log::EpochNewChangeLog;
487    use risingwave_pb::hummock::StateTableInfo;
488
489    use super::*;
490
491    fn new_change_log(checkpoint_epoch: u64) -> EpochNewChangeLog {
492        EpochNewChangeLog {
493            new_value: vec![],
494            old_value: vec![],
495            non_checkpoint_epochs: vec![],
496            checkpoint_epoch,
497        }
498    }
499
500    #[test]
501    fn test_apply_change_log_delta() {
502        let table_id = TableId::new(1);
503        let mut table_change_logs = HashMap::from([(
504            table_id,
505            TableChangeLog::new([new_change_log(1), new_change_log(2)]),
506        )]);
507        HummockVersion::apply_change_log_delta(
508            &mut table_change_logs,
509            &HashMap::from([(table_id, new_change_log(3))]),
510        );
511
512        assert_eq!(
513            table_change_logs[&table_id].epochs().collect_vec(),
514            vec![1, 2, 3]
515        );
516    }
517
518    #[test]
519    fn test_collect_gc_change_log_delta() {
520        let table_id = TableId::new(1);
521        let removed_table_id = TableId::new(2);
522        let current_table_ids = HashSet::from([table_id, removed_table_id]);
523        let state_table_info_delta = HashMap::from([(
524            table_id,
525            StateTableInfoDelta {
526                committed_epoch: 2,
527                compaction_group_id: 1.into(),
528            },
529        )]);
530        let changed_table_info = HashMap::from([(
531            table_id,
532            Some(StateTableInfo {
533                committed_epoch: 1,
534                compaction_group_id: 1.into(),
535            }),
536        )]);
537
538        assert_eq!(
539            HummockVersion::collect_gc_change_log_delta(
540                current_table_ids.iter(),
541                &HashMap::<TableId, EpochNewChangeLog>::new(),
542                &HashSet::from([removed_table_id]),
543                &state_table_info_delta,
544                &changed_table_info,
545            ),
546            HashSet::from([table_id, removed_table_id])
547        );
548        assert_eq!(
549            HummockVersion::collect_gc_change_log_delta(
550                current_table_ids.iter(),
551                &HashMap::from([(table_id, new_change_log(2))]),
552                &HashSet::new(),
553                &state_table_info_delta,
554                &changed_table_info,
555            ),
556            HashSet::new()
557        );
558    }
559}