risingwave_meta/manager/
exactly_once_util.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 risingwave_meta_model::{
16    Epoch, ObjectId, SinkId, SinkSchemachange, object, pending_sink_state,
17};
18use risingwave_pb::stream_plan::PbSinkSchemaChange;
19use sea_orm::{
20    ColumnTrait, ConnectionTrait, DatabaseConnection, EntityTrait, Order, QueryFilter, QueryOrder,
21    QuerySelect, Set, TransactionTrait,
22};
23use thiserror_ext::AsReport;
24
25// Helpers for accessing the `pending_sink_state` system table used by exactly-once sink coordinators
26// (both the generic sink coordinator and the Iceberg pk-index sink coordinator).
27
28async fn sink_object_exists<C: ConnectionTrait>(
29    db: &C,
30    sink_id: SinkId,
31) -> Result<bool, sea_orm::DbErr> {
32    let object_id: Option<ObjectId> = object::Entity::find_by_id(sink_id.as_object_id())
33        .select_only()
34        .column(object::Column::Oid)
35        .into_tuple()
36        .one(db)
37        .await?;
38    Ok(object_id.is_some())
39}
40
41pub async fn persist_pre_commit_metadata(
42    db: &DatabaseConnection,
43    sink_id: SinkId,
44    epoch: u64,
45    commit_metadata: Option<Vec<u8>>,
46    schema_change: Option<&PbSinkSchemaChange>,
47) -> anyhow::Result<()> {
48    fail::fail_point!("iceberg_v3_persist_pre_commit_fail", |_| Err(
49        anyhow::anyhow!("injected: iceberg_v3_persist_pre_commit_fail")
50    ));
51    let schema_change = schema_change.map(Into::into);
52    let m = pending_sink_state::ActiveModel {
53        sink_id: Set(sink_id),
54        epoch: Set(epoch as Epoch),
55        sink_state: Set(pending_sink_state::SinkState::Pending),
56        metadata: Set(commit_metadata),
57        schema_change: Set(schema_change),
58    };
59    match pending_sink_state::Entity::insert(m).exec(db).await {
60        Ok(_) => Ok(()),
61        Err(e) => {
62            // `DROP SINK` currently removes the catalog object before the stop barrier reaches
63            // the sink actor. A commit request already in flight can therefore race with the
64            // cascading deletion of `pending_sink_state`. Once the object is gone there is no
65            // state to recover, so let the request finish instead of failing the actor and the
66            // whole database. Other insert errors must still be surfaced.
67            if !sink_object_exists(db, sink_id).await? {
68                tracing::debug!(
69                    %sink_id,
70                    epoch,
71                    "skip persisting exactly-once metadata for a dropped sink"
72                );
73                return Ok(());
74            }
75            tracing::error!(
76                "Error inserting into exactly once system table: {:?}",
77                e.as_report()
78            );
79            Err(e.into())
80        }
81    }
82}
83
84pub async fn commit_and_prune_epoch(
85    db: &DatabaseConnection,
86    sink_id: SinkId,
87    epoch: u64,
88    prev_epoch: Option<u64>,
89) -> anyhow::Result<()> {
90    fail::fail_point!("iceberg_v3_commit_prune_fail", |_| Err(anyhow::anyhow!(
91        "injected: iceberg_v3_commit_prune_fail"
92    )));
93    let txn = db.begin().await?;
94    let update_result = pending_sink_state::Entity::update(pending_sink_state::ActiveModel {
95        sink_id: Set(sink_id),
96        epoch: Set(epoch as Epoch),
97        sink_state: Set(pending_sink_state::SinkState::Committed),
98        ..Default::default()
99    })
100    .exec(&txn)
101    .await;
102
103    if let Err(e) = update_result {
104        // The row may have been cascade-deleted while an external two-phase commit was in
105        // progress. Roll back before checking through `db`: PostgreSQL leaves a transaction in
106        // an aborted state after a statement error.
107        txn.rollback().await?;
108        if !sink_object_exists(db, sink_id).await? {
109            tracing::debug!(
110                %sink_id,
111                epoch,
112                "skip marking exactly-once metadata committed for a dropped sink"
113            );
114            return Ok(());
115        }
116        return Err(e.into());
117    }
118
119    if let Some(prev_epoch) = prev_epoch {
120        pending_sink_state::Entity::delete_many()
121            .filter(
122                pending_sink_state::Column::SinkId
123                    .eq(sink_id)
124                    .and(pending_sink_state::Column::Epoch.eq(prev_epoch as Epoch)),
125            )
126            .exec(&txn)
127            .await?;
128    }
129
130    match txn.commit().await {
131        Ok(_) => Ok(()),
132        Err(e) => {
133            tracing::error!(
134                "Error marking item to committed exactly once system table: {:?}",
135                e.as_report()
136            );
137            Err(e.into())
138        }
139    }
140}
141
142pub async fn clean_aborted_records(
143    db: &DatabaseConnection,
144    sink_id: SinkId,
145    aborted_epochs: Vec<u64>,
146) -> anyhow::Result<()> {
147    if aborted_epochs.is_empty() {
148        return Ok(());
149    }
150
151    match pending_sink_state::Entity::delete_many()
152        .filter(
153            pending_sink_state::Column::SinkId
154                .eq(sink_id)
155                .and(pending_sink_state::Column::Epoch.is_in(aborted_epochs)),
156        )
157        .exec(db)
158        .await
159    {
160        Ok(_) => Ok(()),
161        Err(e) => {
162            tracing::error!(
163                "Error deleting records from exactly once system table: {:?}",
164                e.as_report()
165            );
166            Err(e.into())
167        }
168    }
169}
170
171type PendingSinkStateRow = (
172    Epoch,
173    pending_sink_state::SinkState,
174    Option<Vec<u8>>,
175    Option<SinkSchemachange>,
176);
177
178pub async fn list_sink_states_ordered_by_epoch(
179    db: &DatabaseConnection,
180    sink_id: SinkId,
181) -> anyhow::Result<
182    Vec<(
183        u64,
184        pending_sink_state::SinkState,
185        Option<Vec<u8>>,
186        Option<PbSinkSchemaChange>,
187    )>,
188> {
189    let rows: Vec<PendingSinkStateRow> = match pending_sink_state::Entity::find()
190        .select_only()
191        .columns([
192            pending_sink_state::Column::Epoch,
193            pending_sink_state::Column::SinkState,
194            pending_sink_state::Column::Metadata,
195            pending_sink_state::Column::SchemaChange,
196        ])
197        .filter(pending_sink_state::Column::SinkId.eq(sink_id))
198        .order_by(pending_sink_state::Column::Epoch, Order::Asc)
199        .into_tuple()
200        .all(db)
201        .await
202    {
203        Ok(rows) => rows,
204        Err(e) => {
205            tracing::error!("Error querying pending sink states: {:?}", e.as_report());
206            return Err(e.into());
207        }
208    };
209
210    Ok(rows
211        .into_iter()
212        .map(|(epoch, state, metadata, schema_change)| {
213            (
214                epoch as u64,
215                state,
216                metadata,
217                schema_change.map(|v| v.to_protobuf()),
218            )
219        })
220        .collect())
221}
222
223#[cfg(test)]
224mod tests {
225    use sea_orm::{ConnectionTrait, Database, DatabaseConnection, DbBackend, Statement};
226
227    use super::{commit_and_prune_epoch, persist_pre_commit_metadata};
228
229    async fn prepare_db() -> DatabaseConnection {
230        let db = Database::connect("sqlite::memory:").await.unwrap();
231        for ddl in [
232            "PRAGMA foreign_keys = ON",
233            "CREATE TABLE object (oid INTEGER PRIMARY KEY)",
234            "CREATE TABLE pending_sink_state (\
235                sink_id INTEGER NOT NULL, \
236                epoch BIGINT NOT NULL, \
237                sink_state STRING NOT NULL, \
238                metadata BLOB, \
239                schema_change BLOB, \
240                PRIMARY KEY (sink_id, epoch), \
241                FOREIGN KEY (sink_id) REFERENCES object(oid) ON DELETE CASCADE\
242            )",
243        ] {
244            db.execute(Statement::from_string(DbBackend::Sqlite, ddl))
245                .await
246                .unwrap();
247        }
248        db
249    }
250
251    async fn insert_object(db: &DatabaseConnection, sink_id: i32) {
252        db.execute(Statement::from_sql_and_values(
253            DbBackend::Sqlite,
254            "INSERT INTO object (oid) VALUES (?)",
255            [sink_id.into()],
256        ))
257        .await
258        .unwrap();
259    }
260
261    #[tokio::test]
262    async fn test_exactly_once_metadata_writes_ignore_dropped_sink() {
263        let db = prepare_db().await;
264        insert_object(&db, 1).await;
265
266        persist_pre_commit_metadata(&db, 1.into(), 100, Some(vec![1, 2, 3]), None)
267            .await
268            .unwrap();
269        db.execute(Statement::from_string(
270            DbBackend::Sqlite,
271            "DELETE FROM object WHERE oid = 1",
272        ))
273        .await
274        .unwrap();
275
276        // The first row was cascade-deleted while the coordinator was committing it, and a new
277        // pre-commit request arrived after the parent object was removed. Both are expected
278        // during `DROP SINK` and must not fail the sink actor.
279        commit_and_prune_epoch(&db, 1.into(), 100, None)
280            .await
281            .unwrap();
282        persist_pre_commit_metadata(&db, 1.into(), 101, Some(vec![4, 5, 6]), None)
283            .await
284            .unwrap();
285
286        let row_count = db
287            .query_one(Statement::from_string(
288                DbBackend::Sqlite,
289                "SELECT COUNT(*) AS count FROM pending_sink_state",
290            ))
291            .await
292            .unwrap()
293            .unwrap()
294            .try_get::<i64>("", "count")
295            .unwrap();
296        assert_eq!(row_count, 0);
297    }
298
299    #[tokio::test]
300    async fn test_missing_state_is_error_for_existing_sink() {
301        let db = prepare_db().await;
302        insert_object(&db, 1).await;
303
304        assert!(
305            commit_and_prune_epoch(&db, 1.into(), 100, None)
306                .await
307                .is_err()
308        );
309
310        persist_pre_commit_metadata(&db, 1.into(), 101, None, None)
311            .await
312            .unwrap();
313        assert!(
314            persist_pre_commit_metadata(&db, 1.into(), 101, None, None)
315                .await
316                .is_err()
317        );
318    }
319}
risingwave_meta/manager/exactly_once_util.rs

risingwave_meta/manager/
exactly_once_util.rs
