little_durable_objects/
placement.rs

1#[cfg(test)]
2pub(crate) mod testing;
3
4use anyhow::{Context, Result, ensure};
5use async_trait::async_trait;
6use serde::{Deserialize, Serialize};
7
8use crate::{actor_state::ActorStorageKey, host::HostId, postgres::PostgresDatabase};
9
10#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
11pub struct ObjectPlacement {
12    pub object: ActorStorageKey,
13    pub owner: HostId,
14    pub owner_epoch: u64,
15    pub home_region: String,
16    pub state_version: u64,
17    pub state_object: Option<String>,
18    pub last_request_id: Option<String>,
19}
20
21#[derive(Clone, Debug, PartialEq, Eq)]
22pub enum PlacementClaim {
23    Acquired(ObjectPlacement),
24    Current(ObjectPlacement),
25}
26
27#[derive(Clone, Debug, PartialEq, Eq)]
28pub struct StateCommitRequest {
29    pub object: ActorStorageKey,
30    pub owner: HostId,
31    pub session_id: String,
32    pub owner_epoch: u64,
33    pub expected_version: u64,
34    pub state_object: String,
35    pub request_id: String,
36}
37
38#[derive(Clone, Debug, PartialEq, Eq)]
39pub enum StateCommit {
40    Committed(ObjectPlacement),
41    Current(ObjectPlacement),
42}
43
44#[async_trait]
45pub trait ObjectPlacementStore: Send + Sync {
46    async fn get(&self, object: &ActorStorageKey) -> Result<Option<ObjectPlacement>>;
47
48    async fn claim(
49        &self,
50        object: &ActorStorageKey,
51        expected: Option<&ObjectPlacement>,
52        owner: &HostId,
53        home_region: &str,
54    ) -> Result<PlacementClaim>;
55
56    async fn commit_state(&self, request: &StateCommitRequest) -> Result<StateCommit>;
57}
58
59pub struct PostgresObjectPlacementStore {
60    database: PostgresDatabase,
61}
62
63impl PostgresObjectPlacementStore {
64    pub async fn connect(url: &str) -> Result<Self> {
65        Ok(Self::from_database(PostgresDatabase::connect(url).await?))
66    }
67
68    pub(crate) fn from_database(database: PostgresDatabase) -> Self {
69        Self { database }
70    }
71}
72
73#[async_trait]
74impl ObjectPlacementStore for PostgresObjectPlacementStore {
75    async fn get(&self, object: &ActorStorageKey) -> Result<Option<ObjectPlacement>> {
76        let row = self
77            .database
78            .query_opt(
79                "SELECT owner_host_id, owner_epoch, home_region, state_version, state_object, last_request_id \
80                 FROM durable_object_placements WHERE object_id = $1",
81                &[&object.as_str()],
82            )
83            .await
84            .context("load PostgreSQL object placement")?;
85        row.map(|row| placement_from_row(object, &row)).transpose()
86    }
87
88    async fn claim(
89        &self,
90        object: &ActorStorageKey,
91        expected: Option<&ObjectPlacement>,
92        owner: &HostId,
93        home_region: &str,
94    ) -> Result<PlacementClaim> {
95        object.validate()?;
96        validate_region(home_region)?;
97        if expected.is_none() {
98            if let Some(row) = self
99                .database
100                .query_opt(
101                    "INSERT INTO durable_object_placements \
102                     (object_id, owner_host_id, owner_epoch, home_region) \
103                     VALUES ($1, $2, 1, $3) ON CONFLICT DO NOTHING \
104                     RETURNING owner_host_id, owner_epoch, home_region, state_version, state_object, last_request_id",
105                    &[&object.as_str(), &owner.as_str(), &home_region],
106                )
107                .await
108                .context("insert PostgreSQL object placement")?
109            {
110                return Ok(PlacementClaim::Acquired(placement_from_row(object, &row)?));
111            }
112            return self.current_claim(object).await;
113        }
114
115        let expected = expected.expect("checked above");
116        ensure!(
117            expected.object == *object && expected.home_region == home_region,
118            "expected object placement does not match the claim"
119        );
120        if &expected.owner == owner {
121            return self.current_claim(object).await;
122        }
123        let expected_epoch = i64::try_from(expected.owner_epoch)
124            .context("object owner epoch exceeds PostgreSQL BIGINT")?;
125        if let Some(row) = self
126            .database
127            .query_opt(
128                "UPDATE durable_object_placements \
129                 SET owner_host_id = $2, owner_epoch = owner_epoch + 1, updated_at = clock_timestamp() \
130                 WHERE object_id = $1 AND owner_host_id = $3 AND owner_epoch = $4 AND home_region = $5 \
131                 RETURNING owner_host_id, owner_epoch, home_region, state_version, state_object, last_request_id",
132                &[
133                    &object.as_str(),
134                    &owner.as_str(),
135                    &expected.owner.as_str(),
136                    &expected_epoch,
137                    &home_region,
138                ],
139            )
140            .await
141            .context("claim PostgreSQL object placement")?
142        {
143            return Ok(PlacementClaim::Acquired(placement_from_row(object, &row)?));
144        }
145        self.current_claim(object).await
146    }
147
148    async fn commit_state(&self, request: &StateCommitRequest) -> Result<StateCommit> {
149        validate_state_commit(request)?;
150        let expected_epoch = i64::try_from(request.owner_epoch)
151            .context("object owner epoch exceeds PostgreSQL BIGINT")?;
152        let expected_version = i64::try_from(request.expected_version)
153            .context("actor state version exceeds PostgreSQL BIGINT")?;
154        let row = self
155            .database
156            .query_opt(
157                "UPDATE durable_object_placements AS placement \
158                 SET state_version = state_version + 1, state_object = $6, last_request_id = $7, updated_at = clock_timestamp() \
159                 FROM durable_object_host_leases AS lease \
160                 WHERE placement.object_id = $1 \
161                   AND placement.owner_host_id = $2 \
162                   AND placement.owner_epoch = $3 \
163                   AND placement.state_version = $4 \
164                   AND lease.host_id = placement.owner_host_id \
165                   AND lease.session_id = $5 \
166                   AND lease.expires_at_ms > (EXTRACT(EPOCH FROM clock_timestamp()) * 1000)::BIGINT \
167                 RETURNING placement.owner_host_id, placement.owner_epoch, placement.home_region, \
168                           placement.state_version, placement.state_object, placement.last_request_id",
169                &[
170                    &request.object.as_str(),
171                    &request.owner.as_str(),
172                    &expected_epoch,
173                    &expected_version,
174                    &request.session_id,
175                    &request.state_object,
176                    &request.request_id,
177                ],
178            )
179            .await
180            .context("commit PostgreSQL actor state head")?;
181        if let Some(row) = row {
182            return Ok(StateCommit::Committed(placement_from_row(
183                &request.object,
184                &row,
185            )?));
186        }
187        let current = self
188            .get(&request.object)
189            .await?
190            .context("actor placement disappeared during state commit")?;
191        Ok(if is_replayed_commit(&current, request) {
192            StateCommit::Committed(current)
193        } else {
194            StateCommit::Current(current)
195        })
196    }
197}
198
199impl PostgresObjectPlacementStore {
200    async fn current_claim(&self, object: &ActorStorageKey) -> Result<PlacementClaim> {
201        self.get(object)
202            .await?
203            .map(PlacementClaim::Current)
204            .context("object placement disappeared during claim")
205    }
206}
207
208fn placement_from_row(
209    object: &ActorStorageKey,
210    row: &tokio_postgres::Row,
211) -> Result<ObjectPlacement> {
212    Ok(ObjectPlacement {
213        object: object.clone(),
214        owner: HostId::new(row.get::<_, String>(0)),
215        owner_epoch: u64::try_from(row.get::<_, i64>(1))
216            .context("PostgreSQL object owner epoch is negative")?,
217        home_region: row.get(2),
218        state_version: u64::try_from(row.get::<_, i64>(3))
219            .context("PostgreSQL actor state version is negative")?,
220        state_object: row.get(4),
221        last_request_id: row.get(5),
222    })
223}
224
225fn validate_state_commit(request: &StateCommitRequest) -> Result<()> {
226    request.object.validate()?;
227    ensure!(
228        !request.owner.as_str().is_empty(),
229        "state commit owner is empty"
230    );
231    ensure!(
232        !request.session_id.is_empty(),
233        "state commit session is empty"
234    );
235    ensure!(
236        request.owner_epoch > 0,
237        "state commit owner epoch must be positive"
238    );
239    ensure!(
240        !request.state_object.is_empty() && request.state_object.len() <= 1024,
241        "state commit object name is invalid"
242    );
243    ensure!(
244        !request.request_id.is_empty() && request.request_id.len() <= 255,
245        "state commit request ID is invalid"
246    );
247    Ok(())
248}
249
250fn is_replayed_commit(current: &ObjectPlacement, request: &StateCommitRequest) -> bool {
251    current.state_version == request.expected_version.saturating_add(1)
252        && current.state_object.as_deref() == Some(&request.state_object)
253        && current.last_request_id.as_deref() == Some(&request.request_id)
254}
255
256pub fn validate_region(region: &str) -> Result<()> {
257    ensure!(
258        !region.is_empty()
259            && region.len() <= 64
260            && region.bytes().all(|byte| {
261                byte.is_ascii_lowercase()
262                    || byte.is_ascii_digit()
263                    || matches!(byte, b'.' | b'_' | b'-')
264            }),
265        "sandbox region is invalid"
266    );
267    Ok(())
268}
little_durable_objects/placement.rs

little_durable_objects/
placement.rs
