Implement crash-safe streaming worker runtime
This commit is contained in:
parent
16cba3c317
commit
cdeef3cc1a
1 changed files with 7448 additions and 194 deletions
198
crates/remora-worker/src/descriptor.rs
Normal file
198
crates/remora-worker/src/descriptor.rs
Normal file
|
|
@ -0,0 +1,198 @@
|
|||
use remora_contracts::{
|
||||
CheckpointDescriptorKindDto, CheckpointDescriptorRefDto, DurableDescriptorKindDto,
|
||||
DurableDescriptorRefDto, PositiveIntegerDto, ResolvedSourceDescriptorKindDto,
|
||||
ResolvedSourceDescriptorRefDto, SafeIntegerDto,
|
||||
};
|
||||
use remora_domain::OpaqueId;
|
||||
use serde::Serialize;
|
||||
use serde_json::Value;
|
||||
use sha2::{Digest, Sha256};
|
||||
use sqlx::PgPool;
|
||||
use time::OffsetDateTime;
|
||||
|
||||
use crate::WorkerError;
|
||||
|
||||
const MIGRATION_TABLE: &str = "_remora_worker_migrations";
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct DescriptorStore {
|
||||
pool: PgPool,
|
||||
}
|
||||
|
||||
impl DescriptorStore {
|
||||
#[must_use]
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
pub async fn store_resolved<T: Serialize>(
|
||||
&self,
|
||||
attempt_id: &OpaqueId,
|
||||
payload: &T,
|
||||
item_count: Option<u64>,
|
||||
) -> Result<ResolvedSourceDescriptorRefDto, WorkerError> {
|
||||
let stored = self
|
||||
.store(attempt_id, "resolved_source", payload, item_count)
|
||||
.await?;
|
||||
Ok(ResolvedSourceDescriptorRefDto {
|
||||
id: stored.id,
|
||||
kind: ResolvedSourceDescriptorKindDto::ResolvedSource,
|
||||
schema_version: stored.schema_version,
|
||||
byte_length: stored.byte_length,
|
||||
item_count: stored.item_count,
|
||||
checksum: stored.checksum,
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn store_checkpoint<T: Serialize>(
|
||||
&self,
|
||||
attempt_id: &OpaqueId,
|
||||
payload: &T,
|
||||
item_count: Option<u64>,
|
||||
) -> Result<CheckpointDescriptorRefDto, WorkerError> {
|
||||
let stored = self
|
||||
.store(attempt_id, "checkpoint", payload, item_count)
|
||||
.await?;
|
||||
Ok(CheckpointDescriptorRefDto {
|
||||
id: stored.id,
|
||||
kind: CheckpointDescriptorKindDto::Checkpoint,
|
||||
schema_version: stored.schema_version,
|
||||
byte_length: stored.byte_length,
|
||||
item_count: stored.item_count,
|
||||
checksum: stored.checksum,
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn load_ref(
|
||||
&self,
|
||||
job_id: &OpaqueId,
|
||||
descriptor: &DurableDescriptorRefDto,
|
||||
) -> Result<Option<Value>, WorkerError> {
|
||||
type Row = (String, i32, Value, i64, Option<i64>, String);
|
||||
let row = sqlx::query_as::<_, Row>(
|
||||
"SELECT d.kind, d.schema_version, d.payload, d.byte_length, d.item_count, d.checksum
|
||||
FROM worker_descriptors d
|
||||
JOIN attempts a ON a.attempt_id = d.attempt_id
|
||||
WHERE d.descriptor_id = $1 AND a.job_id = $2",
|
||||
)
|
||||
.bind(descriptor.id.as_str())
|
||||
.bind(job_id.as_str())
|
||||
.fetch_optional(&self.pool)
|
||||
.await
|
||||
.map_err(|_| WorkerError::ApiUnavailable)?;
|
||||
let Some((kind, schema_version, payload, byte_length, item_count, checksum)) = row else {
|
||||
return Ok(None);
|
||||
};
|
||||
let encoded = canonical_bytes(&payload)?;
|
||||
let expected_kind = match descriptor.kind {
|
||||
DurableDescriptorKindDto::ResolvedSource => "resolved_source",
|
||||
DurableDescriptorKindDto::Checkpoint => "checkpoint",
|
||||
};
|
||||
if kind != expected_kind
|
||||
|| u64::try_from(schema_version).ok() != Some(descriptor.schema_version.get())
|
||||
|| u64::try_from(byte_length).ok() != Some(descriptor.byte_length.get())
|
||||
|| item_count.map(u64::try_from).transpose().ok().flatten()
|
||||
!= descriptor.item_count.map(SafeIntegerDto::get)
|
||||
|| usize::try_from(byte_length).ok() != Some(encoded.len())
|
||||
|| checksum != descriptor.checksum
|
||||
|| digest(&encoded) != checksum
|
||||
{
|
||||
return Err(WorkerError::InvalidResponse);
|
||||
}
|
||||
Ok(Some(payload))
|
||||
}
|
||||
|
||||
async fn store<T: Serialize>(
|
||||
&self,
|
||||
attempt_id: &OpaqueId,
|
||||
kind: &str,
|
||||
payload: &T,
|
||||
item_count: Option<u64>,
|
||||
) -> Result<StoredDescriptor, WorkerError> {
|
||||
let payload = serde_json::to_value(payload).map_err(|_| WorkerError::InvalidRequest)?;
|
||||
let encoded = canonical_bytes(&payload)?;
|
||||
let id = descriptor_id();
|
||||
let checksum = digest(&encoded);
|
||||
let byte_length = SafeIntegerDto::new(
|
||||
u64::try_from(encoded.len()).map_err(|_| WorkerError::InvalidRequest)?,
|
||||
)
|
||||
.map_err(|_| WorkerError::InvalidRequest)?;
|
||||
let item_count = item_count
|
||||
.map(SafeIntegerDto::new)
|
||||
.transpose()
|
||||
.map_err(|_| WorkerError::InvalidRequest)?;
|
||||
sqlx::query(
|
||||
"INSERT INTO worker_descriptors (
|
||||
descriptor_id, attempt_id, kind, schema_version, payload,
|
||||
byte_length, item_count, checksum, created_at
|
||||
) VALUES ($1, $2, $3, 1, $4, $5, $6, $7, $8)",
|
||||
)
|
||||
.bind(id.as_str())
|
||||
.bind(attempt_id.as_str())
|
||||
.bind(kind)
|
||||
.bind(&payload)
|
||||
.bind(i64::try_from(byte_length.get()).map_err(|_| WorkerError::InvalidRequest)?)
|
||||
.bind(
|
||||
item_count
|
||||
.map(SafeIntegerDto::get)
|
||||
.map(i64::try_from)
|
||||
.transpose()
|
||||
.map_err(|_| WorkerError::InvalidRequest)?,
|
||||
)
|
||||
.bind(&checksum)
|
||||
.bind(OffsetDateTime::now_utc())
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_err(|_| WorkerError::ApiUnavailable)?;
|
||||
Ok(StoredDescriptor {
|
||||
id,
|
||||
schema_version: PositiveIntegerDto::new(1).expect("one is positive"),
|
||||
byte_length,
|
||||
item_count,
|
||||
checksum,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn migrate(pool: &PgPool) -> Result<(), sqlx::migrate::MigrateError> {
|
||||
migrator().run(pool).await
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn migrator() -> sqlx::migrate::Migrator {
|
||||
let mut migrator = sqlx::migrate!("./migrations");
|
||||
migrator.dangerous_set_table_name(MIGRATION_TABLE);
|
||||
migrator
|
||||
}
|
||||
|
||||
struct StoredDescriptor {
|
||||
id: OpaqueId,
|
||||
schema_version: PositiveIntegerDto,
|
||||
byte_length: SafeIntegerDto,
|
||||
item_count: Option<SafeIntegerDto>,
|
||||
checksum: String,
|
||||
}
|
||||
|
||||
fn canonical_bytes(value: &Value) -> Result<Vec<u8>, WorkerError> {
|
||||
serde_json_canonicalizer::to_vec(value).map_err(|_| WorkerError::InvalidRequest)
|
||||
}
|
||||
|
||||
fn digest(bytes: &[u8]) -> String {
|
||||
let mut value = String::from("sha256:");
|
||||
for byte in Sha256::digest(bytes) {
|
||||
use std::fmt::Write as _;
|
||||
write!(value, "{byte:02x}").expect("writing to String cannot fail");
|
||||
}
|
||||
value
|
||||
}
|
||||
|
||||
fn descriptor_id() -> OpaqueId {
|
||||
let mut bytes = [0_u8; 16];
|
||||
getrandom::fill(&mut bytes).expect("operating system randomness is required");
|
||||
let mut value = String::from("d");
|
||||
for byte in bytes {
|
||||
use std::fmt::Write as _;
|
||||
write!(value, "{byte:02x}").expect("writing to String cannot fail");
|
||||
}
|
||||
OpaqueId::parse(value).expect("generated descriptor identifier is valid")
|
||||
}
|
||||
Loading…
Reference in a new issue