risingwave_meta/manager/iceberg_pk_index_sink/
manager.rs

1// Copyright 2026 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;
16use std::sync::Arc;
17
18use anyhow::anyhow;
19use parking_lot::RwLock;
20use risingwave_connector::sink::catalog::SinkId;
21use risingwave_connector::sink::iceberg::IcebergConfig;
22use risingwave_pb::stream_service::barrier_complete_response::IcebergPkIndexSinkMetadata as PbIcebergPkIndexSinkMetadata;
23use sea_orm::DatabaseConnection;
24use tokio::sync::Mutex;
25use tracing::warn;
26
27use super::coordinator::IcebergPkIndexSinkCoordinator;
28
29type CoordinatorRef = Arc<Mutex<IcebergPkIndexSinkCoordinator>>;
30
31/// Manager for the Iceberg pk-index sink path, cheap to clone.
32#[derive(Clone)]
33pub struct IcebergPkIndexSinkManager {
34    inner: Arc<ManagerInner>,
35}
36
37struct ManagerInner {
38    db: DatabaseConnection,
39    /// Read on every pre-commit/commit (only to clone out the `CoordinatorRef`, never held across an await);
40    /// written only by register/unregister/reset, which are rare control-plane events.
41    coordinators: RwLock<HashMap<SinkId, CoordinatorRef>>,
42}
43
44impl IcebergPkIndexSinkManager {
45    pub fn new(db: DatabaseConnection) -> Self {
46        IcebergPkIndexSinkManager {
47            inner: Arc::new(ManagerInner {
48                db,
49                coordinators: RwLock::new(HashMap::new()),
50            }),
51        }
52    }
53
54    /// Register an Iceberg pk-index sink so its commit coordinator is ready to receive epoch reports. Builds and
55    /// fully initializes the coordinator (loading the iceberg table and draining any recovered pending
56    /// commits) BEFORE inserting it, so a successful return means the sink is ready to serve. Idempotent:
57    /// registering the same `sink_id` replaces the existing coordinator.
58    pub async fn register_sink(
59        &self,
60        sink_id: SinkId,
61        iceberg_config: IcebergConfig,
62    ) -> anyhow::Result<()> {
63        // Initialize (load + recover + drain) outside the map lock; this is the slow, fallible part.
64        let coordinator =
65            IcebergPkIndexSinkCoordinator::init(sink_id, iceberg_config, self.inner.db.clone())
66                .await?;
67
68        let prev = self
69            .inner
70            .coordinators
71            .write()
72            .insert(sink_id, Arc::new(Mutex::new(coordinator)));
73        if prev.is_some() {
74            // Replacing an existing coordinator. Any in-flight commit on the old one keeps it alive via its
75            // own `Arc` until it finishes; the snapshot_id idempotency check guards against double-commit.
76            warn!(%sink_id, "iceberg pk-index sink coordinator re-registered; replacing previous instance");
77        }
78        Ok(())
79    }
80
81    /// Pre-commit phase for one epoch: persist the merged report under `pending_sink_state` (no iceberg IO).
82    /// The barrier-complete path awaits this BEFORE issuing hummock `commit_epoch`.
83    pub async fn pre_commit_epoch(
84        &self,
85        sink_id: SinkId,
86        prev_epoch: u64,
87        reports: Vec<PbIcebergPkIndexSinkMetadata>,
88    ) -> anyhow::Result<()> {
89        let coordinator = self.coordinator(sink_id)?;
90        coordinator
91            .lock()
92            .await
93            .pre_commit(prev_epoch, reports)
94            .await
95    }
96
97    /// Commit phase for one epoch: run an iceberg `overwrite_files` transaction and mark its pending row
98    /// committed. The barrier-complete path awaits this AFTER hummock `commit_epoch`.
99    pub async fn commit_epoch(&self, sink_id: SinkId) -> anyhow::Result<()> {
100        let coordinator = self.coordinator(sink_id)?;
101        coordinator.lock().await.commit().await
102    }
103
104    /// Unregister the given `sink_id`(s)' coordinator(s) (e.g. at DROP SINK time). Unregistering an unknown
105    /// `sink_id` is a no-op.
106    pub fn unregister_sinks(&self, sink_ids: Vec<SinkId>) {
107        let mut coordinators = self.inner.coordinators.write();
108        for sink_id in sink_ids {
109            coordinators.remove(&sink_id);
110        }
111    }
112
113    /// Drop every coordinator. Used at recovery time.
114    pub fn reset(&self) {
115        self.inner.coordinators.write().clear();
116    }
117
118    fn coordinator(&self, sink_id: SinkId) -> anyhow::Result<CoordinatorRef> {
119        self.inner
120            .coordinators
121            .read()
122            .get(&sink_id)
123            .cloned()
124            .ok_or_else(|| {
125                anyhow!(
126                    "iceberg pk-index sink coordinator for sink {} is not registered",
127                    sink_id
128                )
129            })
130    }
131}
risingwave_meta/manager/iceberg_pk_index_sink/manager.rs

risingwave_meta/manager/iceberg_pk_index_sink/
manager.rs
