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            .query_opt(
179                "SELECT owner_host_id, owner_epoch, home_region, state_version, state_object, last_request_id \
180                 FROM durable_object_placements WHERE object_id = $1",
181                &[&object.as_str()],
182            )
183            .await
184            .context("load PostgreSQL object placement")?;
185        row.map(|row| placement_from_row(object, &row)).transpose()
186    }
187
188    async fn claim(
189        &self,
190        object: &ActorStorageKey,
191        expected: Option<&ObjectPlacement>,
192        owner: &HostId,
193        home_region: &str,
194    ) -> Result<PlacementClaim> {
195        object.validate()?;
196        validate_region(home_region)?;
197        if expected.is_none() {
198            if let Some(row) = self
199                .database
200                .query_opt(
201                    "INSERT INTO durable_object_placements \
202                     (object_id, owner_host_id, owner_epoch, home_region) \
203                     VALUES ($1, $2, 1, $3) ON CONFLICT DO NOTHING \
204                     RETURNING owner_host_id, owner_epoch, home_region, state_version, state_object, last_request_id",
205                    &[&object.as_str(), &owner.as_str(), &home_region],
206                )
207                .await
208                .context("insert PostgreSQL object placement")?
209            {
210                return Ok(PlacementClaim::Acquired(placement_from_row(object, &row)?));
211            }
212            return self.current_claim(object).await;
213        }
214
215        let expected = expected.expect("checked above");
216        ensure!(
217            expected.object == *object && expected.home_region == home_region,
218            "expected object placement does not match the claim"
219        );
220        if &expected.owner == owner {
221            return self.current_claim(object).await;
222        }
223        let expected_epoch = i64::try_from(expected.owner_epoch)
224            .context("object owner epoch exceeds PostgreSQL BIGINT")?;
225        if let Some(row) = self
226            .database
227            .query_opt(
228                "UPDATE durable_object_placements \
229                 SET owner_host_id = $2, owner_epoch = owner_epoch + 1, updated_at = clock_timestamp() \
230                 WHERE object_id = $1 AND owner_host_id = $3 AND owner_epoch = $4 AND home_region = $5 \
231                 RETURNING owner_host_id, owner_epoch, home_region, state_version, state_object, last_request_id",
232                &[
233                    &object.as_str(),
234                    &owner.as_str(),
235                    &expected.owner.as_str(),
236                    &expected_epoch,
237                    &home_region,
238                ],
239            )
240            .await
241            .context("claim PostgreSQL object placement")?
242        {
243            return Ok(PlacementClaim::Acquired(placement_from_row(object, &row)?));
244        }
245        self.current_claim(object).await
246    }
247
248    async fn commit_state(&self, request: &StateCommitRequest) -> Result<StateCommit> {
249        validate_state_commit(request)?;
250        let expected_epoch = i64::try_from(request.owner_epoch)
251            .context("object owner epoch exceeds PostgreSQL BIGINT")?;
252        let expected_version = i64::try_from(request.expected_version)
253            .context("actor state version exceeds PostgreSQL BIGINT")?;
254        let row = self
255            .database
256            .query_opt(
257                "UPDATE durable_object_placements AS placement \
258                 SET state_version = state_version + 1, state_object = $6, last_request_id = $7, updated_at = clock_timestamp() \
259                 FROM durable_object_host_leases AS lease \
260                 WHERE placement.object_id = $1 \
261                   AND placement.owner_host_id = $2 \
262                   AND placement.owner_epoch = $3 \
263                   AND placement.state_version = $4 \
264                   AND lease.host_id = placement.owner_host_id \
265                   AND lease.session_id = $5 \
266                   AND lease.expires_at_ms > (EXTRACT(EPOCH FROM clock_timestamp()) * 1000)::BIGINT \
267                 RETURNING placement.owner_host_id, placement.owner_epoch, placement.home_region, \
268                           placement.state_version, placement.state_object, placement.last_request_id",
269                &[
270                    &request.object.as_str(),
271                    &request.owner.as_str(),
272                    &expected_epoch,
273                    &expected_version,
274                    &request.session_id,
275                    &request.state_object,
276                    &request.request_id,
277                ],
278            )
279            .await
280            .context("commit PostgreSQL actor state head")?;
281        if let Some(row) = row {
282            return Ok(StateCommit::Committed(placement_from_row(
283                &request.object,
284                &row,
285            )?));
286        }
287        let current = self
288            .get(&request.object)
289            .await?
290            .context("actor placement disappeared during state commit")?;
291        Ok(if is_replayed_commit(&current, request) {
292            StateCommit::Committed(current)
293        } else {
294            StateCommit::Current(current)
295        })
296    }
297}
298
299impl PostgresObjectPlacementStore {
300    async fn current_claim(&self, object: &ActorStorageKey) -> Result<PlacementClaim> {
301        self.get(object)
302            .await?
303            .map(PlacementClaim::Current)
304            .context("object placement disappeared during claim")
305    }
306}
307
308fn placement_from_row(
309    object: &ActorStorageKey,
310    row: &tokio_postgres::Row,
311) -> Result<ObjectPlacement> {
312    Ok(ObjectPlacement {
313        object: object.clone(),
314        owner: HostId::new(row.get::<_, String>(0)),
315        owner_epoch: u64::try_from(row.get::<_, i64>(1))
316            .context("PostgreSQL object owner epoch is negative")?,
317        home_region: row.get(2),
318        state_version: u64::try_from(row.get::<_, i64>(3))
319            .context("PostgreSQL actor state version is negative")?,
320        state_object: row.get(4),
321        last_request_id: row.get(5),
322    })
323}
324
325fn validate_state_commit(request: &StateCommitRequest) -> Result<()> {
326    request.object.validate()?;
327    ensure!(
328        !request.owner.as_str().is_empty(),
329        "state commit owner is empty"
330    );
331    ensure!(
332        !request.session_id.is_empty(),
333        "state commit session is empty"
334    );
335    ensure!(
336        request.owner_epoch > 0,
337        "state commit owner epoch must be positive"
338    );
339    ensure!(
340        !request.state_object.is_empty() && request.state_object.len() <= 1024,
341        "state commit object name is invalid"
342    );
343    ensure!(
344        !request.request_id.is_empty() && request.request_id.len() <= 255,
345        "state commit request ID is invalid"
346    );
347    Ok(())
348}
349
350fn is_replayed_commit(current: &ObjectPlacement, request: &StateCommitRequest) -> bool {
351    current.state_version == request.expected_version.saturating_add(1)
352        && current.state_object.as_deref() == Some(&request.state_object)
353        && current.last_request_id.as_deref() == Some(&request.request_id)
354}
355
356pub fn validate_region(region: &str) -> Result<()> {
357    ensure!(
358        !region.is_empty()
359            && region.len() <= 64
360            && region.bytes().all(|byte| {
361                byte.is_ascii_lowercase()
362                    || byte.is_ascii_digit()
363                    || matches!(byte, b'.' | b'_' | b'-')
364            }),
365        "sandbox region is invalid"
366    );
367    Ok(())
368}
369
370#[cfg(test)]
371mod tests {
372    use super::*;
373
374    #[tokio::test]
375    async fn claims_once_and_increments_epoch_on_transfer() -> Result<()> {
376        let store = LocalObjectPlacementStore::default();
377        let object = ActorStorageKey::new("object.v1.project.Counter.one");
378        let first = match store
379            .claim(&object, None, &HostId::new("host-a"), "us-east")
380            .await?
381        {
382            PlacementClaim::Acquired(placement) => placement,
383            claim => anyhow::bail!("unexpected claim: {claim:?}"),
384        };
385        assert_eq!(first.owner_epoch, 1);
386
387        let second = match store
388            .claim(&object, Some(&first), &HostId::new("host-b"), "us-east")
389            .await?
390        {
391            PlacementClaim::Acquired(placement) => placement,
392            claim => anyhow::bail!("unexpected claim: {claim:?}"),
393        };
394        assert_eq!(second.owner, HostId::new("host-b"));
395        assert_eq!(second.owner_epoch, 2);
396        assert_eq!(second.home_region, "us-east");
397        Ok(())
398    }
399
400    #[tokio::test]
401    async fn stale_claim_observes_the_current_owner() -> Result<()> {
402        let store = LocalObjectPlacementStore::default();
403        let object = ActorStorageKey::new("object.v1.project.Counter.one");
404        let PlacementClaim::Acquired(first) = store
405            .claim(&object, None, &HostId::new("host-a"), "us-east")
406            .await?
407        else {
408            anyhow::bail!("first claim was not acquired")
409        };
410        let PlacementClaim::Acquired(second) = store
411            .claim(&object, Some(&first), &HostId::new("host-b"), "us-east")
412            .await?
413        else {
414            anyhow::bail!("second claim was not acquired")
415        };
416        assert_eq!(
417            store
418                .claim(&object, Some(&first), &HostId::new("host-c"), "us-east",)
419                .await?,
420            PlacementClaim::Current(second)
421        );
422        Ok(())
423    }
424}
little_durable_objects/placement.rs

little_durable_objects/
placement.rs
