risingwave_storage/hummock/
object_id_manager.rs

1// Copyright 2025 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::VecDeque;
16use std::ops::DerefMut;
17use std::sync::Arc;
18
19use parking_lot::Mutex;
20use risingwave_hummock_sdk::{HummockRawObjectId, HummockSstableObjectId, ObjectIdRange};
21use risingwave_pb::hummock::GetNewObjectIdsRequest;
22use risingwave_rpc_client::{GrpcCompactorProxyClient, HummockMetaClient};
23use sync_point::sync_point;
24use thiserror_ext::AsReport;
25use tokio::sync::oneshot;
26
27use crate::hummock::{HummockError, HummockResult};
28pub type ObjectIdManagerRef = Arc<ObjectIdManager>;
29
30#[async_trait::async_trait]
31pub trait GetObjectId: Send + Sync {
32    async fn get_new_sst_object_id(&self) -> HummockResult<HummockSstableObjectId>;
33}
34
35/// Caches SST object ids fetched from meta.
36pub struct ObjectIdManager {
37    // Lock order: `wait_queue` before `available_object_ids`.
38    wait_queue: Mutex<Option<Vec<oneshot::Sender<bool>>>>,
39    available_object_ids: Mutex<ObjectIdRange>,
40    remote_fetch_number: u32,
41    hummock_meta_client: Arc<dyn HummockMetaClient>,
42}
43
44impl ObjectIdManager {
45    pub fn new(hummock_meta_client: Arc<dyn HummockMetaClient>, remote_fetch_number: u32) -> Self {
46        Self {
47            wait_queue: Default::default(),
48            available_object_ids: Mutex::new(ObjectIdRange::new(u64::MIN, u64::MIN)),
49            remote_fetch_number,
50            hummock_meta_client,
51        }
52    }
53
54    pub async fn get_new_object_id<I: From<HummockRawObjectId>>(&self) -> HummockResult<I> {
55        self.map_next_object_id(|available_object_ids| available_object_ids.get_next_object_id())
56            .await
57            .map(Into::into)
58    }
59
60    /// Executes `f` with next SST id.
61    /// May fetch new SST ids via RPC.
62    async fn map_next_object_id<F>(&self, f: F) -> HummockResult<HummockRawObjectId>
63    where
64        F: Fn(&mut ObjectIdRange) -> Option<HummockRawObjectId>,
65    {
66        loop {
67            // 1. Try to get
68            if let Some(new_id) = f(self.available_object_ids.lock().deref_mut()) {
69                return Ok(new_id);
70            }
71            // 2. Otherwise either fetch new ids, or wait for previous fetch if any.
72            let waiter = {
73                let mut guard = self.wait_queue.lock();
74                if let Some(new_id) = f(self.available_object_ids.lock().deref_mut()) {
75                    return Ok(new_id);
76                }
77                let wait_queue = guard.deref_mut();
78                if let Some(wait_queue) = wait_queue {
79                    let (tx, rx) = oneshot::channel();
80                    wait_queue.push(tx);
81                    Some(rx)
82                } else {
83                    *wait_queue = Some(vec![]);
84                    None
85                }
86            };
87            if let Some(waiter) = waiter {
88                // Wait for previous fetch
89                sync_point!("MAP_NEXT_SST_OBJECT_ID.AS_FOLLOWER");
90                let _ = waiter.await;
91                continue;
92            }
93            // Fetch new ids.
94            // Release the fetch role and wake followers even if the leader is cancelled.
95            let mut notify_guard = scopeguard::guard(false, |success| self.notify_waiters(success));
96            sync_point!("MAP_NEXT_SST_OBJECT_ID.AS_LEADER");
97            sync_point!("MAP_NEXT_SST_OBJECT_ID.BEFORE_FETCH");
98            async move {
99                let new_sst_ids = self
100                    .hummock_meta_client
101                    .get_new_object_ids(self.remote_fetch_number)
102                    .await
103                    .map_err(HummockError::meta_error)?;
104                sync_point!("MAP_NEXT_SST_OBJECT_ID.AFTER_FETCH");
105                sync_point!("MAP_NEXT_SST_OBJECT_ID.BEFORE_FILL_CACHE");
106                // Update local cache.
107                let result = {
108                    let mut guard = self.available_object_ids.lock();
109                    let available_object_ids = guard.deref_mut();
110                    if new_sst_ids.start_id < available_object_ids.end_id {
111                        Err(HummockError::meta_error(format!(
112                            "SST id moves backwards. new {} < old {}",
113                            new_sst_ids.start_id, available_object_ids.end_id
114                        )))
115                    } else {
116                        *available_object_ids = new_sst_ids;
117                        Ok(())
118                    }
119                };
120                *notify_guard = result.is_ok();
121                result
122            }
123            .await?;
124        }
125    }
126
127    fn notify_waiters(&self, success: bool) {
128        let mut guard = self.wait_queue.lock();
129        let wait_queue = guard.deref_mut().take().unwrap();
130        for notify in wait_queue {
131            let _ = notify.send(success);
132        }
133    }
134}
135
136#[async_trait::async_trait]
137impl GetObjectId for ObjectIdManager {
138    /// Returns a new SST id.
139    /// The id is guaranteed to be monotonic increasing.
140    async fn get_new_sst_object_id(&self) -> HummockResult<HummockSstableObjectId> {
141        self.get_new_object_id().await
142    }
143}
144
145struct SharedComapctorObjectIdManagerCore {
146    output_object_ids: VecDeque<HummockSstableObjectId>,
147    client: Option<GrpcCompactorProxyClient>,
148    sstable_id_remote_fetch_number: u32,
149}
150impl SharedComapctorObjectIdManagerCore {
151    pub fn new(
152        output_object_ids: VecDeque<HummockSstableObjectId>,
153        client: GrpcCompactorProxyClient,
154        sstable_id_remote_fetch_number: u32,
155    ) -> Self {
156        Self {
157            output_object_ids,
158            client: Some(client),
159            sstable_id_remote_fetch_number,
160        }
161    }
162
163    pub fn for_test(output_object_ids: VecDeque<u64>) -> Self {
164        Self {
165            output_object_ids: output_object_ids.into_iter().map(|x| x.into()).collect(),
166            client: None,
167            sstable_id_remote_fetch_number: 0,
168        }
169    }
170}
171/// `SharedComapctorObjectIdManager` is used to get output sst id for serverless compaction.
172pub struct SharedComapctorObjectIdManager {
173    core: tokio::sync::Mutex<SharedComapctorObjectIdManagerCore>,
174}
175
176impl SharedComapctorObjectIdManager {
177    pub fn new(
178        output_object_ids: VecDeque<HummockSstableObjectId>,
179        client: GrpcCompactorProxyClient,
180        sstable_id_remote_fetch_number: u32,
181    ) -> Arc<Self> {
182        Arc::new(Self {
183            core: tokio::sync::Mutex::new(SharedComapctorObjectIdManagerCore::new(
184                output_object_ids,
185                client,
186                sstable_id_remote_fetch_number,
187            )),
188        })
189    }
190
191    pub fn for_test(output_object_ids: VecDeque<u64>) -> Arc<Self> {
192        Arc::new(Self {
193            core: tokio::sync::Mutex::new(SharedComapctorObjectIdManagerCore::for_test(
194                output_object_ids,
195            )),
196        })
197    }
198}
199
200#[async_trait::async_trait]
201impl GetObjectId for SharedComapctorObjectIdManager {
202    async fn get_new_sst_object_id(&self) -> HummockResult<HummockSstableObjectId> {
203        let mut guard = self.core.lock().await;
204        let core = guard.deref_mut();
205
206        if let Some(first_element) = core.output_object_ids.pop_front() {
207            Ok(first_element)
208        } else {
209            tracing::warn!(
210                "The pre-allocated object ids are used up, and new object id are obtained through RPC."
211            );
212            let request = GetNewObjectIdsRequest {
213                number: core.sstable_id_remote_fetch_number,
214            };
215            match core
216                .client
217                .as_mut()
218                .expect("GrpcCompactorProxyClient is None")
219                .get_new_sst_ids(request)
220                .await
221            {
222                Ok(response) => {
223                    let resp = response.into_inner();
224                    let start_id = resp.start_id;
225                    core.output_object_ids.extend(
226                        ((start_id + 1)..resp.end_id).map(Into::<HummockSstableObjectId>::into),
227                    );
228                    Ok(start_id.into())
229                }
230                Err(e) => Err(HummockError::other(format!(
231                    "Fail to get new sst id: {}",
232                    e.as_report()
233                ))),
234            }
235        }
236    }
237}
risingwave_storage/hummock/object_id_manager.rs

risingwave_storage/hummock/
object_id_manager.rs
