laminar_core/checkpoint_decision/
allocation.rs

1//! Monotonic checkpoint ID reservation and allocation.
2
3use super::{
4    Bytes, CheckpointDecisionStore, DecisionError, DecisionStoreUpdateMode, DeploymentIdentity,
5    ObjectStore, PutMode, PutOptions, PutPayload, UpdateVersion, VersionedCheckpointIdHead,
6    DEPLOYMENT_IDENTITY_MAX_BYTES, DEPLOYMENT_IDENTITY_VERSION, LOCAL_RESERVATION_SIZE,
7};
8
9impl CheckpointDecisionStore {
10    fn consume_local_reservation(&self, minimum: u64) -> Option<u64> {
11        let mut reservation = self.local_reservation.lock();
12        let checkpoint_id = reservation.next_id?.max(minimum);
13        if checkpoint_id > reservation.end {
14            reservation.next_id = None;
15            return None;
16        }
17        reservation.next_id = checkpoint_id
18            .checked_add(1)
19            .filter(|next| *next <= reservation.end);
20        Some(checkpoint_id)
21    }
22
23    fn install_local_reservation(&self, first: u64, end: u64) {
24        let mut reservation = self.local_reservation.lock();
25        reservation.end = end;
26        reservation.next_id = first.checked_add(1).filter(|next| *next <= end);
27    }
28
29    async fn allocate_local_checkpoint_id_at_least(
30        &self,
31        deployment_id: &str,
32        minimum: u64,
33    ) -> Result<u64, DecisionError> {
34        if let Some(checkpoint_id) = self.consume_local_reservation(minimum) {
35            return Ok(checkpoint_id);
36        }
37        let lock = self.local_metadata_rmw_lock.as_ref().ok_or_else(|| {
38            DecisionError::Conflict(
39                "local decision store is missing its namespace write lock".into(),
40            )
41        })?;
42        let _guard = lock.lock().await;
43        let current = self
44            .read_checkpoint_id_head(deployment_id)
45            .await?
46            .ok_or_else(|| {
47                DecisionError::Conflict(format!(
48                    "checkpoint ID authority for deployment {deployment_id} disappeared"
49                ))
50            })?;
51        let checkpoint_id = current
52            .head
53            .checkpoint_id
54            .checked_add(1)
55            .map(|next| next.max(minimum))
56            .ok_or_else(|| DecisionError::Conflict("checkpoint ID space exhausted u64".into()))?;
57        let reservation_end = checkpoint_id.saturating_add(LOCAL_RESERVATION_SIZE - 1);
58        let head = DeploymentIdentity {
59            version: DEPLOYMENT_IDENTITY_VERSION,
60            id: deployment_id.to_owned(),
61            allocator_mode: self.update_mode,
62            checkpoint_id: reservation_end,
63            allocation_id: uuid::Uuid::now_v7().to_string(),
64        };
65        let payload = Self::encode_control_record(
66            "deployment identity",
67            &head,
68            DEPLOYMENT_IDENTITY_MAX_BYTES,
69        )?;
70        let result = self
71            .store
72            .put_opts(
73                &Self::deployment_identity_path(),
74                PutPayload::from(payload),
75                PutOptions {
76                    mode: PutMode::Overwrite,
77                    ..PutOptions::default()
78                },
79            )
80            .await;
81        match result {
82            Ok(put_result) => {
83                self.cache_checkpoint_id_head(Some(VersionedCheckpointIdHead {
84                    head,
85                    update_version: put_result.into(),
86                }));
87                self.install_local_reservation(checkpoint_id, reservation_end);
88                Ok(checkpoint_id)
89            }
90            Err(error) => {
91                let observed = self.read_checkpoint_id_head(deployment_id).await?;
92                if observed.as_ref().is_some_and(|value| value.head == head) {
93                    self.cache_checkpoint_id_head(observed);
94                    self.install_local_reservation(checkpoint_id, reservation_end);
95                    Ok(checkpoint_id)
96                } else {
97                    self.cache_checkpoint_id_head(observed);
98                    Err(DecisionError::Io(error.to_string()))
99                }
100            }
101        }
102    }
103
104    async fn read_checkpoint_id_head(
105        &self,
106        deployment_id: &str,
107    ) -> Result<Option<VersionedCheckpointIdHead>, DecisionError> {
108        let observed = self.read_deployment_identity().await?;
109        if let Some(observed) = observed.as_ref() {
110            if observed.head.id != deployment_id {
111                return Err(DecisionError::Conflict(format!(
112                    "checkpoint ID head belongs to deployment {}, current deployment is {deployment_id}",
113                    observed.head.id
114                )));
115            }
116        }
117        Ok(observed)
118    }
119
120    pub(super) fn cache_checkpoint_id_head(&self, head: Option<VersionedCheckpointIdHead>) {
121        *self.checkpoint_id_head.lock() = head;
122    }
123
124    fn validate_checkpoint_id_head_progress(
125        prior: Option<&VersionedCheckpointIdHead>,
126        observed: Option<&VersionedCheckpointIdHead>,
127    ) -> Result<(), DecisionError> {
128        match (prior, observed) {
129            (Some(prior), None) => Err(DecisionError::Conflict(format!(
130                "checkpoint ID head disappeared after durable ID {}",
131                prior.head.checkpoint_id
132            ))),
133            (Some(prior), Some(observed))
134                if observed.head.checkpoint_id < prior.head.checkpoint_id =>
135            {
136                Err(DecisionError::Conflict(format!(
137                    "checkpoint ID head regressed from {} to {}",
138                    prior.head.checkpoint_id, observed.head.checkpoint_id
139                )))
140            }
141            (Some(prior), Some(observed))
142                if observed.head.checkpoint_id == prior.head.checkpoint_id
143                    && observed.head != prior.head =>
144            {
145                Err(DecisionError::Conflict(format!(
146                    "checkpoint ID {} changed allocation identity without advancing",
147                    prior.head.checkpoint_id
148                )))
149            }
150            _ => Ok(()),
151        }
152    }
153
154    async fn allocate_shared_checkpoint_id_at_least(
155        &self,
156        deployment_id: &str,
157        minimum: u64,
158    ) -> Result<u64, DecisionError> {
159        let allocation_id = uuid::Uuid::now_v7().to_string();
160        let mut current = self.checkpoint_id_head.lock().clone();
161        if current.is_none() {
162            current = self.read_checkpoint_id_head(deployment_id).await?;
163            if current.is_none() {
164                return Err(DecisionError::Conflict(format!(
165                    "checkpoint ID authority for deployment {deployment_id} disappeared"
166                )));
167            }
168            self.cache_checkpoint_id_head(current.clone());
169        }
170
171        loop {
172            let versioned = current.as_ref().ok_or_else(|| {
173                DecisionError::Conflict(format!(
174                    "checkpoint ID authority for deployment {deployment_id} disappeared"
175                ))
176            })?;
177            let checkpoint_id = versioned
178                .head
179                .checkpoint_id
180                .checked_add(1)
181                .map(|next| next.max(minimum))
182                .ok_or_else(|| {
183                    DecisionError::Conflict("checkpoint ID space exhausted u64".to_owned())
184                })?;
185            let head = DeploymentIdentity {
186                version: DEPLOYMENT_IDENTITY_VERSION,
187                id: deployment_id.to_owned(),
188                allocator_mode: self.update_mode,
189                checkpoint_id,
190                allocation_id: allocation_id.clone(),
191            };
192            let payload = serde_json::to_vec(&head)
193                .map(Bytes::from)
194                .map_err(|error| DecisionError::Conflict(error.to_string()))?;
195            let mode = PutMode::Update(versioned.update_version.clone());
196            let result = self
197                .store
198                .put_opts(
199                    &Self::deployment_identity_path(),
200                    PutPayload::from(payload),
201                    PutOptions {
202                        mode,
203                        ..PutOptions::default()
204                    },
205                )
206                .await;
207            match result {
208                Ok(put_result) => {
209                    let update_version: UpdateVersion = put_result.into();
210                    if update_version.e_tag.is_none() && update_version.version.is_none() {
211                        // The create/update itself is authoritative, but this response cannot
212                        // safely seed the next CAS. Force a metadata read on the next allocation.
213                        self.cache_checkpoint_id_head(None);
214                    } else {
215                        self.cache_checkpoint_id_head(Some(VersionedCheckpointIdHead {
216                            head,
217                            update_version,
218                        }));
219                    }
220                    return Ok(checkpoint_id);
221                }
222                Err(
223                    object_store::Error::Precondition { .. }
224                    | object_store::Error::AlreadyExists { .. }
225                    | object_store::Error::NotFound { .. },
226                ) => {
227                    let observed = self.read_checkpoint_id_head(deployment_id).await?;
228                    Self::validate_checkpoint_id_head_progress(
229                        current.as_ref(),
230                        observed.as_ref(),
231                    )?;
232                    current = observed;
233                    self.cache_checkpoint_id_head(current.clone());
234                    tokio::task::yield_now().await;
235                }
236                Err(error) => {
237                    let observed = self.read_checkpoint_id_head(deployment_id).await?;
238                    Self::validate_checkpoint_id_head_progress(
239                        current.as_ref(),
240                        observed.as_ref(),
241                    )?;
242                    if observed.as_ref().is_some_and(|value| value.head == head) {
243                        self.cache_checkpoint_id_head(observed);
244                        return Ok(checkpoint_id);
245                    }
246                    if observed
247                        .as_ref()
248                        .is_some_and(|value| value.head.checkpoint_id >= checkpoint_id)
249                    {
250                        current = observed;
251                        self.cache_checkpoint_id_head(current.clone());
252                        tokio::task::yield_now().await;
253                        continue;
254                    }
255                    self.cache_checkpoint_id_head(observed);
256                    return Err(DecisionError::Io(error.to_string()));
257                }
258            }
259        }
260    }
261
262    /// Allocate the next globally ordered checkpoint ID at or above `minimum`.
263    ///
264    /// Shared stores advance exactly one ID with native compare-and-swap. A certified local single
265    /// writer reserves a durable range in the deployment singleton, consumes it in memory, and
266    /// burns any unused suffix after restart.
267    ///
268    /// # Errors
269    /// Object-store I/O, malformed or foreign durable state, a shared store without conditional
270    /// update support, or exhaustion of the `u64` ID space.
271    pub async fn allocate_checkpoint_id_at_least(
272        &self,
273        minimum: u64,
274    ) -> Result<u64, DecisionError> {
275        if minimum == 0 {
276            return Err(DecisionError::Conflict(
277                "minimum checkpoint ID must be nonzero".to_owned(),
278            ));
279        }
280
281        let deployment_id = self.load_or_create_deployment_id().await?;
282        let _guard = self.metadata_write_lock.lock().await;
283        match self.update_mode {
284            DecisionStoreUpdateMode::NativeCas => {
285                self.allocate_shared_checkpoint_id_at_least(&deployment_id, minimum)
286                    .await
287            }
288            DecisionStoreUpdateMode::LocalSingleWriter => {
289                self.allocate_local_checkpoint_id_at_least(&deployment_id, minimum)
290                    .await
291            }
292        }
293    }
294
295    #[cfg(test)]
296    pub(super) async fn checkpoint_id_reservation_high_watermark(
297        &self,
298    ) -> Result<u64, DecisionError> {
299        let deployment_id = self.load_or_create_deployment_id().await?;
300        self.read_checkpoint_id_head(&deployment_id)
301            .await?
302            .map(|head| head.head.checkpoint_id)
303            .ok_or_else(|| {
304                DecisionError::Conflict(format!(
305                    "checkpoint ID authority for deployment {deployment_id} disappeared"
306                ))
307            })
308    }
309
310    #[cfg(test)]
311    /// Allocate from the default floor in unit tests.
312    ///
313    /// # Errors
314    ///
315    /// Returns an error when the durable allocator is unavailable or invalid.
316    pub async fn allocate_checkpoint_id(&self) -> Result<u64, DecisionError> {
317        self.allocate_checkpoint_id_at_least(1).await
318    }
319}
laminar_core/checkpoint_decision/allocation.rs

laminar_core/checkpoint_decision/
allocation.rs
