ms2k/remora
SHA256
Watch
0
Code Releases Activity

Implement crash-safe streaming worker runtime

Browse source
This commit is contained in:
ms2k 2026-07-17 18:51:44 +09:00
commit cdeef3cc1a
1 changed files with 7448 additions and 194 deletions
Download patch file Download diff file Expand all files Collapse all files

162
crates/remora-result/src/inventory.rs
View file

@ -3,7 +3,7 @@ use remora_domain::{MemberPath, OpaqueId};
use remora_stream::TargetBody;
use remora_target::{
ContentSha256, CreateObjectRequest, ObjectIdentity, StoredObject, TargetStore,
validate_attempt_key,
attempt_manifest_key_in, attempt_object_prefix, validate_attempt_key,
};
use serde::Serialize;
use sqlx::PgPool;
@ -26,10 +26,25 @@ impl ResultInventory {
}
pub async fn create(&self, attempt_id: &OpaqueId) -> Result<(), ResultError> {
let result = sqlx::query("INSERT INTO result_inventories (attempt_id) VALUES ($1)")
.bind(attempt_id.as_str())
.execute(&self.pool)
.await;
self.create_bound(attempt_id, attempt_id, "").await
}
pub async fn create_bound(
&self,
attempt_id: &OpaqueId,
staging_origin_attempt_id: &OpaqueId,
target_prefix: &str,
) -> Result<(), ResultError> {
let result = sqlx::query(
"INSERT INTO result_inventories (
attempt_id, staging_origin_attempt_id, target_prefix
) VALUES ($1, $2, $3)",
)
.bind(attempt_id.as_str())
.bind(staging_origin_attempt_id.as_str())
.bind(target_prefix)
.execute(&self.pool)
.await;
match result {
Ok(_) => Ok(()),
Err(error) if is_unique_violation(&error, "result_inventories_pkey") => {
@ -39,14 +54,100 @@ impl ResultInventory {
}
}
pub async fn open_or_create(&self, attempt_id: &OpaqueId) -> Result<(), ResultError> {
self.open_or_create_bound(attempt_id, attempt_id, "").await
}
pub async fn open_or_create_bound(
&self,
attempt_id: &OpaqueId,
staging_origin_attempt_id: &OpaqueId,
target_prefix: &str,
) -> Result<(), ResultError> {
match self
.create_bound(attempt_id, staging_origin_attempt_id, target_prefix)
.await
{
Ok(()) => Ok(()),
Err(ResultError::InventoryExists) => {
let binding = sqlx::query_as::<_, (String, Option<String>, Option<String>)>(
"SELECT state, staging_origin_attempt_id, target_prefix
FROM result_inventories WHERE attempt_id = $1",
)
.bind(attempt_id.as_str())
.fetch_one(&self.pool)
.await?;
if binding.0 == "open"
&& binding.1.as_deref() == Some(staging_origin_attempt_id.as_str())
&& binding.2.as_deref() == Some(target_prefix)
{
Ok(())
} else {
Err(ResultError::InventoryNotOpen)
}
}
Err(error) => Err(error),
}
}
pub async fn verify_member_path(
&self,
attempt_id: &OpaqueId,
path: &str,
target: &TargetStore,
) -> Result<Option<ResultMember>, ResultError> {
let member = sqlx::query_as::<_, MemberRow>(
"SELECT member_id, path, size, object_store, object_bucket, object_key,
object_version_id, object_sha256
FROM result_inventory_members
WHERE attempt_id = $1 AND path = $2",
)
.bind(attempt_id.as_str())
.bind(path)
.fetch_optional(&self.pool)
.await?
.map(MemberRow::into_member)
.transpose()?;
let Some(member) = member else {
return Ok(None);
};
let observed = target.verify_exact(&member.exact_identity()).await?;
if observed.identity != member.exact_identity()
|| observed.size != member.size()
|| observed.content_sha256.as_ref() != Some(member.content_sha256())
{
return Err(ResultError::Integrity);
}
Ok(Some(member))
}
pub async fn add_member(
&self,
attempt_id: &OpaqueId,
member: &ResultMember,
) -> Result<(), ResultError> {
validate_member_identity(attempt_id, member)?;
let size = i64::try_from(member.size()).map_err(|_| ResultError::SizeOverflow)?;
let mut tx = self.pool.begin().await?;
let (staging_origin, target_prefix) =
sqlx::query_as::<_, (Option<String>, Option<String>)>(
"SELECT staging_origin_attempt_id, target_prefix
FROM result_inventories WHERE attempt_id = $1",
)
.bind(attempt_id.as_str())
.fetch_one(tx.as_mut())
.await?;
let staging_origin = staging_origin
.ok_or(ResultError::InvalidObjectIdentity)
.and_then(|value| {
OpaqueId::parse(value).map_err(|_| ResultError::InvalidObjectIdentity)
})?;
validate_member_identity(
target_prefix
.as_deref()
.ok_or(ResultError::InvalidObjectIdentity)?,
&staging_origin,
member,
)?;
let state: Option<String> = sqlx::query_scalar(
"SELECT state FROM result_inventories WHERE attempt_id = $1 FOR UPDATE",
)
@ -158,7 +259,12 @@ impl ResultInventory {
}
self.validate_members(attempt_id, target).await?;
let (member_count, total_bytes) = self.inventory_totals(attempt_id).await?;
let request = CreateObjectRequest::for_attempt_manifest(attempt_id, manifest_id);
let (staging_origin, target_prefix) = self.inventory_binding(attempt_id).await?;
let request = CreateObjectRequest::for_target_attempt_manifest(
&target_prefix,
&staging_origin,
manifest_id,
);
let mut body = target.open_create_only(&request).await?;
body.write_all(&manifest_header(attempt_id))
.await
@ -217,7 +323,9 @@ impl ResultInventory {
return Err(ResultError::StaleAttempt);
}
require_supported_version(input.schema_version)?;
let expected_key = remora_target::attempt_manifest_key(attempt_id, &input.manifest_id);
let (staging_origin, target_prefix) = self.inventory_binding(attempt_id).await?;
let expected_key =
attempt_manifest_key_in(&target_prefix, &staging_origin, &input.manifest_id);
if input.manifest.identity.key != expected_key {
return Err(ResultError::InvalidObjectIdentity);
}
@ -291,12 +399,26 @@ impl ResultInventory {
attempt_id: &OpaqueId,
target: &TargetStore,
) -> Result<(), ResultError> {
let (staging_origin, target_prefix) =
sqlx::query_as::<_, (Option<String>, Option<String>)>(
"SELECT staging_origin_attempt_id, target_prefix
FROM result_inventories WHERE attempt_id = $1",
)
.bind(attempt_id.as_str())
.fetch_one(&self.pool)
.await?;
let staging_origin = staging_origin
.ok_or(ResultError::InvalidObjectIdentity)
.and_then(|value| {
OpaqueId::parse(value).map_err(|_| ResultError::InvalidObjectIdentity)
})?;
let target_prefix = target_prefix.ok_or(ResultError::InvalidObjectIdentity)?;
let mut rows = sqlx::query_as::<_, MemberRow>(MEMBER_SELECT)
.bind(attempt_id.as_str())
.fetch(&self.pool);
while let Some(row) = rows.try_next().await? {
let member = row.into_member()?;
validate_member_identity(attempt_id, &member)?;
validate_member_identity(&target_prefix, &staging_origin, &member)?;
let observed = target.verify_exact(&member.exact_identity()).await?;
if observed.identity != member.exact_identity()
|| observed.size != member.size()
@ -318,6 +440,25 @@ impl ResultInventory {
Ok((count as u64, bytes as u64))
}
async fn inventory_binding(
&self,
attempt_id: &OpaqueId,
) -> Result<(OpaqueId, String), ResultError> {
let (origin, prefix) = sqlx::query_as::<_, (Option<String>, Option<String>)>(
"SELECT staging_origin_attempt_id, target_prefix
FROM result_inventories WHERE attempt_id = $1",
)
.bind(attempt_id.as_str())
.fetch_one(&self.pool)
.await?;
let origin = origin
.ok_or(ResultError::InvalidObjectIdentity)
.and_then(|value| {
OpaqueId::parse(value).map_err(|_| ResultError::InvalidObjectIdentity)
})?;
Ok((origin, prefix.ok_or(ResultError::InvalidObjectIdentity)?))
}
async fn persist_sealed(&self, input: &PublicationInput) -> Result<(), ResultError> {
let result = sqlx::query(
"UPDATE result_inventories SET
@ -462,11 +603,12 @@ impl InventoryRow {
}
fn validate_member_identity(
target_prefix: &str,
attempt_id: &OpaqueId,
member: &ResultMember,
) -> Result<(), ResultError> {
validate_attempt_key(member.object().key()).map_err(|_| ResultError::InvalidObjectIdentity)?;
let prefix = format!("_remora/a/{}/o/", attempt_id.as_str());
let prefix = attempt_object_prefix(target_prefix, attempt_id);
if !member.object().key().starts_with(&prefix) {
return Err(ResultError::StaleAttempt);
}