1use 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 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 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(), table_ids: vec![],
194 version: CompatibilityVersion::LATEST as _,
195 split_key: None,
196 })));
197 }
198
199 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 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 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}