little_durable_objects/
placement.rs

1use std::{collections::HashMap, sync::Mutex};
2
3use anyhow::{Context, Result, ensure};
4use async_trait::async_trait;
5use serde::{Deserialize, Serialize};
6
7use crate::{actor_state::ActorStorageKey, host::HostId, postgres::PostgresDatabase};
8
9#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
10pub struct ObjectPlacement {
11    pub object: ActorStorageKey,
12    pub owner: HostId,
13    pub owner_epoch: u64,
14    pub home_region: String,
15    pub state_version: u64,
16    pub state_object: Option<String>,
17    pub last_request_id: Option<String>,
18}
19
20#[derive(Clone, Debug, PartialEq, Eq)]
21pub enum PlacementClaim {
22    Acquired(ObjectPlacement),
23    Current(ObjectPlacement),
24}
25
26#[derive(Clone, Debug, PartialEq, Eq)]
27pub struct StateCommitRequest {
28    pub object: ActorStorageKey,
29    pub owner: HostId,
30    pub session_id: String,
31    pub owner_epoch: u64,
32    pub expected_version: u64,
33    pub state_object: String,
34    pub request_id: String,
35}
36
37#[derive(Clone, Debug, PartialEq, Eq)]
38pub enum StateCommit {
39    Committed(ObjectPlacement),
40    Current(ObjectPlacement),
41}
42
43#[async_trait]
44pub trait ObjectPlacementStore: Send + Sync {
45    async fn get(&self, object: &ActorStorageKey) -> Result<Option<ObjectPlacement>>;
46
47    async fn claim(
48        &self,
49        object: &ActorStorageKey,
50        expected: Option<&ObjectPlacement>,
51        owner: &HostId,
52        home_region: &str,
53    ) -> Result<PlacementClaim>;
54
55    async fn commit_state(&self, request: &StateCommitRequest) -> Result<StateCommit>;
56}
57
58#[derive(Default)]
59pub struct LocalObjectPlacementStore {
60    placements: Mutex<HashMap<ActorStorageKey, ObjectPlacement>>,
61}
62
63#[async_trait]
64impl ObjectPlacementStore for LocalObjectPlacementStore {
65    async fn get(&self, object: &ActorStorageKey) -> Result<Option<ObjectPlacement>> {
66        Ok(self
67            .placements
68            .lock()
69            .map_err(|_| anyhow::anyhow!("object placement lock poisoned"))?
70            .get(object)
71            .cloned())
72    }
73
74    async fn claim(
75        &self,
76        object: &ActorStorageKey,
77        expected: Option<&ObjectPlacement>,
78        owner: &HostId,
79        home_region: &str,
80    ) -> Result<PlacementClaim> {
81        validate_region(home_region)?;
82        let mut placements = self
83            .placements
84            .lock()
85            .map_err(|_| anyhow::anyhow!("object placement lock poisoned"))?;
86        match placements.get(object) {
87            None if expected.is_none() => {
88                let placement = ObjectPlacement {
89                    object: object.clone(),
90                    owner: owner.clone(),
91                    owner_epoch: 1,
92                    home_region: home_region.to_owned(),
93                    state_version: 0,
94                    state_object: None,
95                    last_request_id: None,
96                };
97                placements.insert(object.clone(), placement.clone());
98                Ok(PlacementClaim::Acquired(placement))
99            }
100            Some(current) if expected == Some(current) => {
101                ensure!(
102                    current.home_region == home_region,
103                    "object home region cannot change"
104                );
105                if &current.owner == owner {
106                    return Ok(PlacementClaim::Current(current.clone()));
107                }
108                let placement = ObjectPlacement {
109                    object: object.clone(),
110                    owner: owner.clone(),
111                    owner_epoch: current
112                        .owner_epoch
113                        .checked_add(1)
114                        .context("object owner epoch overflow")?,
115                    home_region: home_region.to_owned(),
116                    state_version: current.state_version,
117                    state_object: current.state_object.clone(),
118                    last_request_id: current.last_request_id.clone(),
119                };
120                placements.insert(object.clone(), placement.clone());
121                Ok(PlacementClaim::Acquired(placement))
122            }
123            Some(current) => Ok(PlacementClaim::Current(current.clone())),
124            None => anyhow::bail!("expected object placement no longer exists"),
125        }
126    }
127
128    async fn commit_state(&self, request: &StateCommitRequest) -> Result<StateCommit> {
129        validate_state_commit(request)?;
130        let mut placements = self
131            .placements
132            .lock()
133            .map_err(|_| anyhow::anyhow!("object placement lock poisoned"))?;
134        let current = placements
135            .get(&request.object)
136            .cloned()
137            .context("actor placement does not exist")?;
138        if is_replayed_commit(&current, request) {
139            return Ok(StateCommit::Committed(current));
140        }
141        if current.owner != request.owner
142            || current.owner_epoch != request.owner_epoch
143            || current.state_version != request.expected_version
144        {
145            return Ok(StateCommit::Current(current));
146        }
147        let mut committed = current;
148        committed.state_version = committed
149            .state_version
150            .checked_add(1)
151            .context("actor state version overflow")?;
152        committed.state_object = Some(request.state_object.clone());
153        committed.last_request_id = Some(request.request_id.clone());
154        placements.insert(request.object.clone(), committed.clone());
155        Ok(StateCommit::Committed(committed))
156    }
157}
158
159pub struct PostgresObjectPlacementStore {
160    database: PostgresDatabase,
161}
162
163impl PostgresObjectPlacementStore {
164    pub async fn connect(url: &str) -> Result<Self> {
165        Ok(Self::from_database(PostgresDatabase::connect(url).await?))
166    }
167
168    pub(crate) fn from_database(database: PostgresDatabase) -> Self {
169        Self { database }
170    }
171}
172
173#[async_trait]
174impl ObjectPlacementStore for PostgresObjectPlacementStore {
175    async fn get(&self, object: &ActorStorageKey) -> Result<Option<ObjectPlacement>> {
176        let row = self
177            .database
178            .client()
179            .query_opt(
180                "SELECT owner_host_id, owner_epoch, home_region, state_version, state_object, last_request_id \
181                 FROM durable_object_placements WHERE object_id = $1",
182                &[&object.as_str()],
183            )
184            .await
185            .context("load PostgreSQL object placement")?;
186        row.map(|row| placement_from_row(object, &row)).transpose()
187    }
188
189    async fn claim(
190        &self,
191        object: &ActorStorageKey,
192        expected: Option<&ObjectPlacement>,
193        owner: &HostId,
194        home_region: &str,
195    ) -> Result<PlacementClaim> {
196        object.validate()?;
197        validate_region(home_region)?;
198        if expected.is_none() {
199            if let Some(row) = self
200                .database
201                .client()
202                .query_opt(
203                    "INSERT INTO durable_object_placements \
204                     (object_id, owner_host_id, owner_epoch, home_region) \
205                     VALUES ($1, $2, 1, $3) ON CONFLICT DO NOTHING \
206                     RETURNING owner_host_id, owner_epoch, home_region, state_version, state_object, last_request_id",
207                    &[&object.as_str(), &owner.as_str(), &home_region],
208                )
209                .await
210                .context("insert PostgreSQL object placement")?
211            {
212                return Ok(PlacementClaim::Acquired(placement_from_row(object, &row)?));
213            }
214            return self.current_claim(object).await;
215        }
216
217        let expected = expected.expect("checked above");
218        ensure!(
219            expected.object == *object && expected.home_region == home_region,
220            "expected object placement does not match the claim"
221        );
222        if &expected.owner == owner {
223            return self.current_claim(object).await;
224        }
225        let expected_epoch = i64::try_from(expected.owner_epoch)
226            .context("object owner epoch exceeds PostgreSQL BIGINT")?;
227        if let Some(row) = self
228            .database
229            .client()
230            .query_opt(
231                "UPDATE durable_object_placements \
232                 SET owner_host_id = $2, owner_epoch = owner_epoch + 1, updated_at = clock_timestamp() \
233                 WHERE object_id = $1 AND owner_host_id = $3 AND owner_epoch = $4 AND home_region = $5 \
234                 RETURNING owner_host_id, owner_epoch, home_region, state_version, state_object, last_request_id",
235                &[
236                    &object.as_str(),
237                    &owner.as_str(),
238                    &expected.owner.as_str(),
239                    &expected_epoch,
240                    &home_region,
241                ],
242            )
243            .await
244            .context("claim PostgreSQL object placement")?
245        {
246            return Ok(PlacementClaim::Acquired(placement_from_row(object, &row)?));
247        }
248        self.current_claim(object).await
249    }
250
251    async fn commit_state(&self, request: &StateCommitRequest) -> Result<StateCommit> {
252        validate_state_commit(request)?;
253        let expected_epoch = i64::try_from(request.owner_epoch)
254            .context("object owner epoch exceeds PostgreSQL BIGINT")?;
255        let expected_version = i64::try_from(request.expected_version)
256            .context("actor state version exceeds PostgreSQL BIGINT")?;
257        let row = self
258            .database
259            .client()
260            .query_opt(
261                "UPDATE durable_object_placements AS placement \
262                 SET state_version = state_version + 1, state_object = $6, last_request_id = $7, updated_at = clock_timestamp() \
263                 FROM durable_object_host_leases AS lease \
264                 WHERE placement.object_id = $1 \
265                   AND placement.owner_host_id = $2 \
266                   AND placement.owner_epoch = $3 \
267                   AND placement.state_version = $4 \
268                   AND lease.host_id = placement.owner_host_id \
269                   AND lease.session_id = $5 \
270                   AND lease.expires_at_ms > (EXTRACT(EPOCH FROM clock_timestamp()) * 1000)::BIGINT \
271                 RETURNING placement.owner_host_id, placement.owner_epoch, placement.home_region, \
272                           placement.state_version, placement.state_object, placement.last_request_id",
273                &[
274                    &request.object.as_str(),
275                    &request.owner.as_str(),
276                    &expected_epoch,
277                    &expected_version,
278                    &request.session_id,
279                    &request.state_object,
280                    &request.request_id,
281                ],
282            )
283            .await
284            .context("commit PostgreSQL actor state head")?;
285        if let Some(row) = row {
286            return Ok(StateCommit::Committed(placement_from_row(
287                &request.object,
288                &row,
289            )?));
290        }
291        let current = self
292            .get(&request.object)
293            .await?
294            .context("actor placement disappeared during state commit")?;
295        Ok(if is_replayed_commit(&current, request) {
296            StateCommit::Committed(current)
297        } else {
298            StateCommit::Current(current)
299        })
300    }
301}
302
303impl PostgresObjectPlacementStore {
304    async fn current_claim(&self, object: &ActorStorageKey) -> Result<PlacementClaim> {
305        self.get(object)
306            .await?
307            .map(PlacementClaim::Current)
308            .context("object placement disappeared during claim")
309    }
310}
311
312fn placement_from_row(
313    object: &ActorStorageKey,
314    row: &tokio_postgres::Row,
315) -> Result<ObjectPlacement> {
316    Ok(ObjectPlacement {
317        object: object.clone(),
318        owner: HostId::new(row.get::<_, String>(0)),
319        owner_epoch: u64::try_from(row.get::<_, i64>(1))
320            .context("PostgreSQL object owner epoch is negative")?,
321        home_region: row.get(2),
322        state_version: u64::try_from(row.get::<_, i64>(3))
323            .context("PostgreSQL actor state version is negative")?,
324        state_object: row.get(4),
325        last_request_id: row.get(5),
326    })
327}
328
329fn validate_state_commit(request: &StateCommitRequest) -> Result<()> {
330    request.object.validate()?;
331    ensure!(
332        !request.owner.as_str().is_empty(),
333        "state commit owner is empty"
334    );
335    ensure!(
336        !request.session_id.is_empty(),
337        "state commit session is empty"
338    );
339    ensure!(
340        request.owner_epoch > 0,
341        "state commit owner epoch must be positive"
342    );
343    ensure!(
344        !request.state_object.is_empty() && request.state_object.len() <= 1024,
345        "state commit object name is invalid"
346    );
347    ensure!(
348        !request.request_id.is_empty() && request.request_id.len() <= 255,
349        "state commit request ID is invalid"
350    );
351    Ok(())
352}
353
354fn is_replayed_commit(current: &ObjectPlacement, request: &StateCommitRequest) -> bool {
355    current.state_version == request.expected_version.saturating_add(1)
356        && current.state_object.as_deref() == Some(&request.state_object)
357        && current.last_request_id.as_deref() == Some(&request.request_id)
358}
359
360pub fn validate_region(region: &str) -> Result<()> {
361    ensure!(
362        !region.is_empty()
363            && region.len() <= 64
364            && region.bytes().all(|byte| {
365                byte.is_ascii_lowercase()
366                    || byte.is_ascii_digit()
367                    || matches!(byte, b'.' | b'_' | b'-')
368            }),
369        "sandbox region is invalid"
370    );
371    Ok(())
372}
373
374#[cfg(test)]
375mod tests {
376    use super::*;
377
378    #[tokio::test]
379    async fn claims_once_and_increments_epoch_on_transfer() -> Result<()> {
380        let store = LocalObjectPlacementStore::default();
381        let object = ActorStorageKey::new("object.v1.project.Counter.one");
382        let first = match store
383            .claim(&object, None, &HostId::new("host-a"), "us-east")
384            .await?
385        {
386            PlacementClaim::Acquired(placement) => placement,
387            claim => anyhow::bail!("unexpected claim: {claim:?}"),
388        };
389        assert_eq!(first.owner_epoch, 1);
390
391        let second = match store
392            .claim(&object, Some(&first), &HostId::new("host-b"), "us-east")
393            .await?
394        {
395            PlacementClaim::Acquired(placement) => placement,
396            claim => anyhow::bail!("unexpected claim: {claim:?}"),
397        };
398        assert_eq!(second.owner, HostId::new("host-b"));
399        assert_eq!(second.owner_epoch, 2);
400        assert_eq!(second.home_region, "us-east");
401        Ok(())
402    }
403
404    #[tokio::test]
405    async fn stale_claim_observes_the_current_owner() -> Result<()> {
406        let store = LocalObjectPlacementStore::default();
407        let object = ActorStorageKey::new("object.v1.project.Counter.one");
408        let PlacementClaim::Acquired(first) = store
409            .claim(&object, None, &HostId::new("host-a"), "us-east")
410            .await?
411        else {
412            anyhow::bail!("first claim was not acquired")
413        };
414        let PlacementClaim::Acquired(second) = store
415            .claim(&object, Some(&first), &HostId::new("host-b"), "us-east")
416            .await?
417        else {
418            anyhow::bail!("second claim was not acquired")
419        };
420        assert_eq!(
421            store
422                .claim(&object, Some(&first), &HostId::new("host-c"), "us-east",)
423                .await?,
424            PlacementClaim::Current(second)
425        );
426        Ok(())
427    }
428}
little_durable_objects/placement.rs

little_durable_objects/
placement.rs
