ms2k/remora
SHA256
Watch
0
Code Releases Activity

Implement authorized live event replay

Browse source
This commit is contained in:
ms2k 2026-07-17 23:08:28 +09:00
commit 06a7ac56b7
14 changed files with 2370 additions and 364 deletions
Download patch file Download diff file Expand all files Collapse all files

16
crates/remora-admission/src/service.rs
View file

@ -56,7 +56,7 @@ impl AdmissionService {
fingerprint(&request).map_err(|error| AdmissionError::Invalid(error.to_string()))?;
let subject = principal.id.clone();
let principal = principal.clone();
let config = self.config.clone();
let service = self.clone();
let outcome = execute_mutation(
&self.pool,
MutationScope {
@ -64,12 +64,12 @@ impl AdmissionService {
operation: "job.submit",
key: idempotency_key,
},
"submit_job",
"admission_submit_job",
&[],
request_fingerprint,
move |tx| {
Box::pin(async move {
let decision = decide_submission(tx, &principal, &config, request).await?;
let decision = service.submit_in_tx(tx, &principal, request).await?;
let status = match decision {
AdmissionDecision::Accepted(_) => MutationStatus::Accepted,
AdmissionDecision::Rejected { .. } => MutationStatus::Rejected,
@ -92,6 +92,15 @@ impl AdmissionService {
})
}
pub async fn submit_in_tx(
&self,
tx: &mut PgTx<'_>,
principal: &Principal,
request: JobSubmissionDto,
) -> Result<AdmissionDecision, DurableError> {
decide_submission(tx, principal, &self.config, request).await
}
pub async fn revoke_execution_grant(
&self,
principal: &Principal,
@ -249,6 +258,7 @@ async fn decide_submission(
.map_err(DurableError::effect)?;
}
enqueue(tx, "job.dispatch", json!({ "job_id": job_id })).await?;
enqueue(tx, "job.updated", json!({ "job_id": job_id })).await?;
Ok(AdmissionDecision::Accepted(Box::new(AcceptedJob {
job_id,

5
crates/remora-api/Cargo.toml
View file

@ -6,8 +6,10 @@ rust-version.workspace = true
version.workspace = true
[dependencies]
actix-ws = "=0.4.0"
actix-web = { version = "=4.14.0", features = ["rustls-0_23"] }
actix-tls = "=3.5.0"
futures-util = "=0.3.32"
getrandom = "0.3.4"
remora-admission = { path = "../remora-admission" }
remora-attempt = { path = "../remora-attempt" }
@ -25,6 +27,7 @@ sha2.workspace = true
sqlx.workspace = true
thiserror.workspace = true
time.workspace = true
tokio.workspace = true
[dev-dependencies]
base64 = "=0.22.1"
@ -33,4 +36,4 @@ rcgen = "=0.14.8"
remora-coordination = { path = "../remora-coordination" }
reqwest = { version = "=0.13.4", default-features = false, features = ["form", "json", "rustls"] }
rustls = "=0.23.42"
tokio.workspace = true
tokio-tungstenite = "=0.30.0"

33
crates/remora-api/migrations/0001_live_event_history.sql Normal file
View file

@ -0,0 +1,33 @@
CREATE TABLE live_event_history_state (
singleton BOOLEAN PRIMARY KEY DEFAULT TRUE CHECK (singleton),
floor_outbox_id BIGINT NOT NULL DEFAULT 0 CHECK (floor_outbox_id >= 0),
progress_poll_not_before TIMESTAMPTZ NOT NULL DEFAULT '-infinity'
);
INSERT INTO live_event_history_state (singleton) VALUES (TRUE);
CREATE TABLE live_progress_event_state (
job_id TEXT PRIMARY KEY REFERENCES jobs (job_id),
attempt_id TEXT NOT NULL,
fence BIGINT NOT NULL CHECK (fence BETWEEN 1 AND 9007199254740991),
progress_present BOOLEAN NOT NULL,
completed_bytes BIGINT,
total_bytes BIGINT,
completed_members BIGINT,
total_members BIGINT,
updated_at TIMESTAMPTZ NOT NULL,
CONSTRAINT live_progress_event_shape_check CHECK (
(
progress_present
AND completed_bytes IS NOT NULL
AND completed_members IS NOT NULL
)
OR (
NOT progress_present
AND completed_bytes IS NULL
AND total_bytes IS NULL
AND completed_members IS NULL
AND total_members IS NULL
)
)
);

7
crates/remora-api/src/lib.rs
View file

@ -1,7 +1,12 @@
mod migrate;
mod public;
mod worker;
pub use public::{PublicApiConfig, PublicApiServices, PublicApiState, configure_public_api};
pub use migrate::{migrate, migrator};
pub use public::{
PublicApiConfig, PublicApiServices, PublicApiState, configure_public_api,
prune_live_event_history,
};
pub use worker::{
AuthenticatedWorker, WorkerApiState, authenticate_rustls_worker, configure_worker_api,
};

16
crates/remora-api/src/migrate.rs Normal file
View file

@ -0,0 +1,16 @@
use remora_durable::DurableError;
use sqlx::PgPool;
const MIGRATION_TABLE: &str = "_remora_api_migrations";
pub async fn migrate(pool: &PgPool) -> Result<(), DurableError> {
let mut migrator = sqlx::migrate!("./migrations");
migrator.dangerous_set_table_name(MIGRATION_TABLE);
migrator.run(pool).await.map_err(DurableError::from)
}
pub fn migrator() -> sqlx::migrate::Migrator {
let mut migrator = sqlx::migrate!("./migrations");
migrator.dangerous_set_table_name(MIGRATION_TABLE);
migrator
}

148
crates/remora-api/src/public.rs
View file

@ -5,7 +5,7 @@ use actix_web::{
http::{StatusCode, header},
web,
};
use remora_admission::{AdmissionDecision, AdmissionError, AdmissionService};
use remora_admission::{AdmissionDecision, AdmissionService};
use remora_attempt::{AttemptService, CommittedResultRecord, ProgressSnapshot};
use remora_contracts::{
ApiErrorCodeDto, ApiErrorDto, CredentialMaterialDto, DateTimeDto, JobSubmissionDto,
@ -39,6 +39,8 @@ use serde_json::{Value, json};
use sqlx::{PgPool, types::Json};
use time::{OffsetDateTime, format_description::well_known::Rfc3339};
mod events;
const DEFAULT_LIMIT: u32 = 50;
const MAX_LIMIT: u32 = 200;
@ -46,6 +48,9 @@ const MAX_LIMIT: u32 = 200;
pub struct PublicApiConfig {
pub cursor_secret: Vec<u8>,
pub member_access_ttl: Duration,
pub event_retention: Duration,
pub event_poll_interval: Duration,
pub event_heartbeat_interval: Duration,
}
#[derive(Clone)]
@ -60,6 +65,9 @@ pub struct PublicApiState {
results: ResultInventory,
cursors: CursorCodec,
member_access_ttl: Duration,
event_retention: Duration,
event_poll_interval: Duration,
event_heartbeat_interval: Duration,
}
pub struct PublicApiServices {
@ -82,6 +90,14 @@ impl PublicApiState {
"member access expiry must be positive".into(),
));
}
if config.event_retention.is_zero()
|| config.event_poll_interval.is_zero()
|| config.event_heartbeat_interval.is_zero()
{
return Err(DurableError::Config(
"event retention and intervals must be positive".into(),
));
}
Ok(Self {
results: ResultInventory::new(pool.clone()),
cursors: CursorCodec::new(config.cursor_secret)?,
@ -93,6 +109,9 @@ impl PublicApiState {
deletion: services.deletion,
targets: services.targets,
member_access_ttl: config.member_access_ttl,
event_retention: config.event_retention,
event_poll_interval: config.event_poll_interval,
event_heartbeat_interval: config.event_heartbeat_interval,
})
}
@ -173,10 +192,18 @@ pub fn configure_public_api(config: &mut web::ServiceConfig) {
.route(
"/deletions/{operation_id}",
web::get().to(get_deletion_operation),
),
)
.route("/events", web::get().to(events::connect)),
);
}
pub async fn prune_live_event_history(
pool: &PgPool,
before: OffsetDateTime,
) -> Result<u64, DurableError> {
events::prune_history(pool, before).await
}
#[derive(Deserialize)]
struct PageQuery {
cursor: Option<String>,
@ -198,17 +225,69 @@ async fn submit_job(
) -> Result<HttpResponse, PublicApiError> {
let principal = state.principal(&http).await?;
let key = idempotency_key(&http)?;
let outcome = state
.admission
.submit(&principal, key, request.into_inner())
.await
.map_err(PublicApiError::from_admission)?;
let accepted = match outcome.decision {
AdmissionDecision::Accepted(accepted) => accepted,
AdmissionDecision::Rejected { .. } => return Err(PublicApiError::invalid_request()),
let request = request.into_inner();
let request_fingerprint =
fingerprint(&request).map_err(|_| PublicApiError::invalid_request())?;
let admission = state.admission.clone();
let state_for_projection = state.get_ref().clone();
let principal_for_effect = principal.clone();
let outcome = execute_mutation(
&state.pool,
MutationScope {
subject: &principal.id,
operation: "job.submit",
key,
},
"submit_job",
&[],
request_fingerprint,
move |tx| {
Box::pin(async move {
match admission
.submit_in_tx(tx, &principal_for_effect, request)
.await?
{
AdmissionDecision::Accepted(accepted) => {
let projection = load_authorized_job_in_tx(
tx,
&state_for_projection,
&principal_for_effect,
&accepted.job_id,
false,
)
.await
.map_err(|error| DurableError::Effect(error.to_string()))?;
Ok((MutationStatus::Accepted, projection))
}
AdmissionDecision::Rejected { .. } => Ok((
MutationStatus::Rejected,
public_error_body(ApiErrorCodeDto::InvalidRequest, "invalid request"),
)),
}
})
},
)
.await?;
if outcome.disposition == MutationDisposition::Conflict {
return Err(PublicApiError::idempotency_conflict());
}
if outcome.status == MutationStatus::Accepted {
let job_id = outcome
.response
.get("id")
.and_then(Value::as_str)
.ok_or_else(PublicApiError::internal)?;
let job_id = parse_id(job_id.to_owned())?;
authorize_job(&state, &principal, &job_id, ManagementOperation::Read).await?;
}
let response = if outcome.status == MutationStatus::Rejected {
HttpResponse::BadRequest()
.content_type("application/problem+json")
.json(outcome.response)
} else {
HttpResponse::Accepted().json(outcome.response)
};
let job = load_authorized_job(&state, &principal, &accepted.job_id, false).await?;
Ok(with_idempotency_expiry(HttpResponse::Accepted().json(job)))
Ok(with_idempotency_expiry(response))
}
async fn get_job(
@ -241,6 +320,7 @@ async fn list_jobs(
.map(|token| state.cursors.bind(token, principal.id.as_str(), filter))
.transpose()
.map_err(PublicApiError::from_cursor)?;
let event_cursor = event_cursor(&state, &principal).await?;
let mut rows = load_job_page_rows(
&state.pool,
&principal,
@ -251,7 +331,6 @@ async fn list_jobs(
.await?;
let has_more = rows.len() > limit as usize;
rows.truncate(limit as usize);
let event_cursor = event_cursor(&state, &principal).await?;
let mut items = Vec::with_capacity(rows.len());
for row in &rows {
items.push(job_row_projection(
@ -1073,6 +1152,7 @@ async fn load_authorized_job(
job_id: &OpaqueId,
include_deleted: bool,
) -> Result<Value, PublicApiError> {
let cursor = event_cursor(state, principal).await?;
let row = load_job_row(&state.pool, job_id)
.await?
.ok_or_else(PublicApiError::not_found)?;
@ -1080,7 +1160,6 @@ async fn load_authorized_job(
return Err(PublicApiError::not_found());
}
authorize_row(principal, &row, ManagementOperation::Read)?;
let cursor = event_cursor(state, principal).await?;
let progress = current_progress(state, &row).await?;
job_row_projection(&row, progress, &cursor)
}
@ -1654,11 +1733,16 @@ async fn event_cursor(
principal: &Principal,
) -> Result<String, PublicApiError> {
let high_water = sqlx::query_as::<_, (i64, OffsetDateTime)>(
"SELECT id, created_at FROM outbox_messages ORDER BY id DESC LIMIT 1",
"SELECT GREATEST(
COALESCE((SELECT MAX(id) FROM outbox_messages), 0),
history.floor_outbox_id
),
now()
FROM live_event_history_state history
WHERE history.singleton",
)
.fetch_optional(&state.pool)
.await?
.unwrap_or((0, OffsetDateTime::UNIX_EPOCH));
.fetch_one(&state.pool)
.await?;
encode_event_cursor(state, principal, high_water)
}
@ -1668,11 +1752,16 @@ async fn event_cursor_in_tx(
principal: &Principal,
) -> Result<String, PublicApiError> {
let high_water = sqlx::query_as::<_, (i64, OffsetDateTime)>(
"SELECT id, created_at FROM outbox_messages ORDER BY id DESC LIMIT 1",
"SELECT GREATEST(
COALESCE((SELECT MAX(id) FROM outbox_messages), 0),
history.floor_outbox_id
),
now()
FROM live_event_history_state history
WHERE history.singleton",
)
.fetch_optional(tx.as_mut())
.await?
.unwrap_or((0, OffsetDateTime::UNIX_EPOCH));
.fetch_one(tx.as_mut())
.await?;
encode_event_cursor(state, principal, high_water)
}
@ -1721,6 +1810,7 @@ fn authorization_snapshot(principal: &Principal) -> Value {
})
.collect::<Vec<_>>();
scopes.sort_unstable();
scopes.dedup();
json!({
"administrator": principal.administrator,
"scopes": scopes
@ -1875,15 +1965,6 @@ impl PublicApiError {
}
}
fn from_admission(error: AdmissionError) -> Self {
match error {
AdmissionError::IdempotencyConflict => Self::idempotency_conflict(),
AdmissionError::Unauthorized => Self::forbidden(),
AdmissionError::Invalid(_) => Self::invalid_request(),
_ => Self::unavailable(),
}
}
fn from_credential(error: CredentialError) -> Self {
match error {
CredentialError::NotFound => Self::not_found(),
@ -1988,6 +2069,7 @@ mod tests {
};
let first = principal(false, vec![scope("b"), scope("a")]);
let reordered = principal(false, vec![scope("a"), scope("b")]);
let duplicated = principal(false, vec![scope("a"), scope("b"), scope("a")]);
let expanded = principal(false, vec![scope("a"), scope("b"), scope("c")]);
let administrator = principal(true, vec![scope("a"), scope("b")]);
@ -1995,6 +2077,10 @@ mod tests {
authorization_snapshot(&first),
authorization_snapshot(&reordered)
);
assert_eq!(
authorization_snapshot(&first),
authorization_snapshot(&duplicated)
);
assert_ne!(
authorization_snapshot(&first),
authorization_snapshot(&expanded)

860
crates/remora-api/src/public/events.rs Normal file
View file

@ -0,0 +1,860 @@
use actix_web::{HttpRequest, HttpResponse, web};
use futures_util::StreamExt as _;
use remora_contracts::{
CursorDto, DateTimeDto, DeletionChangedDto, DeletionStateDto, EventAttemptDto, EventFrameDto,
EventProgressDto, HeartbeatDto, JobChangedDto, NullableDto, PositiveIntegerDto,
ResyncReasonDto, ResyncRequiredDto, SafeIntegerDto, SchemaVersionDto,
};
use remora_domain::{JobLifecycle, JobPhase, OpaqueId};
use remora_durable::{DurableError, KeysetCursor};
use remora_identity::{ManagementOperation, Principal};
use serde::Deserialize;
use serde_json::Value;
use sqlx::{FromRow, PgPool};
use time::{OffsetDateTime, format_description::well_known::Rfc3339};
use super::{
JobRow, PublicApiError, PublicApiState, authorization_snapshot, authorize_row,
current_progress, encode_event_cursor, event_cursor, load_job_row, parse_id, request_id,
};
const EVENT_TOPICS: [&str; 2] = ["job.updated", "deletion.updated"];
const EVENT_BATCH: i64 = 100;
#[derive(Deserialize)]
pub(super) struct EventQuery {
cursor: Option<String>,
}
pub(super) async fn connect(
state: web::Data<PublicApiState>,
request: HttpRequest,
body: web::Payload,
query: web::Query<EventQuery>,
) -> Result<HttpResponse, PublicApiError> {
let principal = state.principal(&request).await?;
let now = OffsetDateTime::now_utc();
let retention =
time::Duration::try_from(state.event_retention).map_err(|_| PublicApiError::internal())?;
prune_history(&state.pool, now - retention).await?;
let start = classify_cursor(&state, &principal, query.cursor.as_deref(), now).await?;
let (response, session, messages) =
actix_ws::handle(&request, body).map_err(|_| PublicApiError::invalid_request())?;
let state = state.get_ref().clone();
actix_web::rt::spawn(async move {
match start {
EventStart::Replay(after_id) => {
run_connection(state, principal, after_id, session, messages).await;
}
EventStart::Resync(reason) => {
send_resync(state, principal, reason, session).await;
}
}
});
Ok(response)
}
pub(super) async fn prune_history(
pool: &PgPool,
before: OffsetDateTime,
) -> Result<u64, DurableError> {
let mut tx = pool.begin().await?;
let pruned: Vec<i64> = sqlx::query_scalar(
"DELETE FROM outbox_messages
WHERE topic = ANY($1) AND created_at < $2
RETURNING id",
)
.bind(&EVENT_TOPICS[..])
.bind(before)
.fetch_all(tx.as_mut())
.await?;
if let Some(floor) = pruned.iter().max() {
sqlx::query(
"UPDATE live_event_history_state
SET floor_outbox_id = GREATEST(floor_outbox_id, $1)
WHERE singleton",
)
.bind(floor)
.execute(tx.as_mut())
.await?;
}
tx.commit().await?;
Ok(pruned.len() as u64)
}
enum EventStart {
Replay(i64),
Resync(ResyncReasonDto),
}
async fn classify_cursor(
state: &PublicApiState,
principal: &Principal,
token: Option<&str>,
now: OffsetDateTime,
) -> Result<EventStart, PublicApiError> {
let Some(token) = token else {
return Ok(EventStart::Resync(ResyncReasonDto::CursorAbsent));
};
let after_id =
match decode_event_cursor(&state.cursors, state.event_retention, principal, token, now)? {
Ok(after_id) => after_id,
Err(reason) => return Ok(EventStart::Resync(reason)),
};
let (floor, maximum): (i64, i64) = sqlx::query_as(
"SELECT history.floor_outbox_id,
GREATEST(
COALESCE((SELECT MAX(id) FROM outbox_messages), 0),
history.floor_outbox_id
)
FROM live_event_history_state history
WHERE history.singleton",
)
.fetch_one(&state.pool)
.await?;
if after_id < floor || after_id > maximum {
return Ok(EventStart::Resync(ResyncReasonDto::HistoryUnavailable));
}
Ok(EventStart::Replay(after_id))
}
fn decode_event_cursor(
cursors: &remora_durable::CursorCodec,
retention: std::time::Duration,
principal: &Principal,
token: &str,
now: OffsetDateTime,
) -> Result<Result<i64, ResyncReasonDto>, PublicApiError> {
let cursor = match cursors.decode(token) {
Ok(cursor) => cursor,
Err(_) => return Ok(Err(ResyncReasonDto::CursorInvalid)),
};
let Some(after_id) = event_position(&cursor) else {
return Ok(Err(ResyncReasonDto::CursorInvalid));
};
let expected = remora_durable::CursorCodec::filter_fingerprint(&serde_json::json!({
"resource": "events",
"authorization": authorization_snapshot(principal)
}))?;
if cursor.subject != principal.id.as_str() || cursor.filter_fingerprint != expected {
return Ok(Err(ResyncReasonDto::AuthorizationChanged));
}
let retention = time::Duration::try_from(retention).map_err(|_| PublicApiError::internal())?;
if cursor.high_water.created_at + retention < now {
return Ok(Err(ResyncReasonDto::CursorExpired));
}
Ok(Ok(after_id))
}
fn event_position(cursor: &KeysetCursor) -> Option<i64> {
let high_water = cursor
.high_water
.id
.strip_prefix("event-")?
.parse::<i64>()
.ok()?;
let after = cursor.after.as_ref()?;
let after_id = after.id.strip_prefix("event-")?.parse::<i64>().ok()?;
(high_water == after_id && after.created_at == cursor.high_water.created_at).then_some(after_id)
}
async fn send_resync(
state: PublicApiState,
principal: Principal,
reason: ResyncReasonDto,
mut session: actix_ws::Session,
) {
let Ok(cursor) = event_cursor(&state, &principal).await else {
let _ = session.close(None).await;
return;
};
let Ok(frame) = resync_frame(reason, 1, &cursor, OffsetDateTime::now_utc()) else {
let _ = session.close(None).await;
return;
};
let _ = send_frame(&mut session, &frame).await;
let _ = session.close(None).await;
}
async fn run_connection(
state: PublicApiState,
principal: Principal,
mut after_id: i64,
mut session: actix_ws::Session,
mut messages: actix_ws::MessageStream,
) {
let mut sequence = 0_u64;
let mut poll = tokio::time::interval(state.event_poll_interval);
let mut heartbeat = tokio::time::interval(state.event_heartbeat_interval);
poll.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
heartbeat.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
heartbeat.tick().await;
loop {
tokio::select! {
message = messages.next() => {
match message {
Some(Ok(actix_ws::Message::Ping(bytes))) => {
if session.pong(&bytes).await.is_err() {
break;
}
}
Some(Ok(actix_ws::Message::Pong(_))) => {}
Some(Ok(actix_ws::Message::Close(reason))) => {
let _ = session.close(reason).await;
return;
}
Some(Err(_)) | None => break,
Some(Ok(_)) => break,
}
}
_ = poll.tick() => {
if poll_progress(&state).await.is_err() {
break;
}
match replay_events(
&state,
&principal,
&mut after_id,
&mut sequence,
&mut session,
).await {
Ok(ReplayStatus::Current) => {}
Ok(ReplayStatus::HistoryUnavailable) => {
let now = OffsetDateTime::now_utc();
let Ok(cursor) = event_cursor(&state, &principal).await else {
break;
};
let Ok(frame) = resync_frame(
ResyncReasonDto::HistoryUnavailable,
next_sequence(&mut sequence),
&cursor,
now,
) else {
break;
};
let _ = send_frame(&mut session, &frame).await;
break;
}
Err(()) => break,
}
}
_ = heartbeat.tick() => {
let now = OffsetDateTime::now_utc();
let Ok(cursor) = encode_event_cursor(&state, &principal, (after_id, now)) else {
break;
};
let Ok(frame) = heartbeat_frame(next_sequence(&mut sequence), &cursor, now) else {
break;
};
if send_frame(&mut session, &frame).await.is_err() {
break;
}
}
}
}
let _ = session.close(None).await;
}
#[derive(FromRow)]
struct ReplayRow {
id: i64,
topic: String,
payload: Value,
created_at: OffsetDateTime,
}
#[derive(FromRow)]
struct ReplayBatchRow {
floor_outbox_id: i64,
id: Option<i64>,
topic: Option<String>,
payload: Option<Value>,
created_at: Option<OffsetDateTime>,
}
enum ReplayStatus {
Current,
HistoryUnavailable,
}
async fn replay_events(
state: &PublicApiState,
principal: &Principal,
after_id: &mut i64,
sequence: &mut u64,
session: &mut actix_ws::Session,
) -> Result<ReplayStatus, ()> {
for _ in 0..10 {
let batch = sqlx::query_as::<_, ReplayBatchRow>(
"SELECT history.floor_outbox_id, events.id, events.topic,
events.payload, events.created_at
FROM live_event_history_state history
LEFT JOIN LATERAL (
SELECT id, topic, payload, created_at
FROM outbox_messages
WHERE id > $1 AND topic = ANY($2)
ORDER BY id
LIMIT $3
) events ON TRUE
WHERE history.singleton
ORDER BY events.id",
)
.bind(*after_id)
.bind(&EVENT_TOPICS[..])
.bind(EVENT_BATCH)
.fetch_all(&state.pool)
.await
.map_err(|_| ())?;
let Some(first) = batch.first() else {
return Err(());
};
if *after_id < first.floor_outbox_id {
return Ok(ReplayStatus::HistoryUnavailable);
}
let rows = batch
.into_iter()
.filter_map(|row| {
Some(ReplayRow {
id: row.id?,
topic: row.topic?,
payload: row.payload?,
created_at: row.created_at?,
})
})
.collect::<Vec<_>>();
let count = rows.len();
for row in rows {
*after_id = row.id;
let now = OffsetDateTime::now_utc();
let cursor = encode_event_cursor(state, principal, (row.id, now)).map_err(|_| ())?;
let candidate_sequence = sequence.saturating_add(1);
let frame = match row.topic.as_str() {
"job.updated" => {
project_job_event(state, principal, &row, candidate_sequence, &cursor).await
}
"deletion.updated" => {
project_deletion_event(state, principal, &row, candidate_sequence, &cursor)
.await
}
_ => None,
};
if let Some(frame) = frame {
*sequence = candidate_sequence;
send_frame(session, &frame).await?;
}
}
if count < EVENT_BATCH as usize {
break;
}
}
Ok(ReplayStatus::Current)
}
async fn project_job_event(
state: &PublicApiState,
principal: &Principal,
event: &ReplayRow,
sequence: u64,
cursor: &str,
) -> Option<EventFrameDto> {
let job_id = parse_payload_id(&event.payload, "job_id")?;
let row = load_job_row(&state.pool, &job_id).await.ok()??;
authorize_row(principal, &row, ManagementOperation::Read).ok()?;
let progress = current_progress(state, &row)
.await
.ok()
.flatten()
.map(progress_fingerprint);
job_frame(
&row,
progress.as_ref(),
stable_event_id(event.id),
sequence,
cursor,
event.created_at,
)
.ok()
}
async fn project_deletion_event(
state: &PublicApiState,
principal: &Principal,
event: &ReplayRow,
sequence: u64,
cursor: &str,
) -> Option<EventFrameDto> {
let operation_id = parse_payload_id(&event.payload, "operation_id")?;
let (job_id, state_name): (String, String) =
sqlx::query_as("SELECT job_id, state FROM deletion_operations WHERE operation_id = $1")
.bind(operation_id.as_str())
.fetch_optional(&state.pool)
.await
.ok()??;
let job_id = parse_id(job_id).ok()?;
let row = load_job_row(&state.pool, &job_id).await.ok()??;
authorize_row(principal, &row, ManagementOperation::Read).ok()?;
deletion_frame(
job_id,
operation_id,
parse_deletion_state(&state_name)?,
stable_event_id(event.id),
sequence,
cursor,
event.created_at,
)
.ok()
}
#[derive(Clone, Debug, Eq, PartialEq)]
struct ProgressFingerprint {
attempt_id: OpaqueId,
fence: u64,
completed_bytes: u64,
total_bytes: Option<u64>,
completed_members: u64,
total_members: Option<u64>,
}
fn progress_fingerprint(snapshot: remora_attempt::ProgressSnapshot) -> ProgressFingerprint {
ProgressFingerprint {
attempt_id: snapshot.attempt_id,
fence: snapshot.fence,
completed_bytes: snapshot.progress.completed_bytes,
total_bytes: snapshot.progress.total_bytes,
completed_members: snapshot.progress.completed_members,
total_members: snapshot.progress.total_members,
}
}
async fn poll_progress(state: &PublicApiState) -> Result<(), ()> {
let delay = state.event_poll_interval.as_secs_f64();
let claimed = sqlx::query_scalar::<_, bool>(
"UPDATE live_event_history_state
SET progress_poll_not_before = now() + ($1 * interval '1 second')
WHERE singleton AND progress_poll_not_before <= now()
RETURNING singleton",
)
.bind(delay)
.fetch_optional(&state.pool)
.await
.map_err(|_| ())?
.is_some();
if !claimed {
return Ok(());
}
let mut after_job = None::<String>;
loop {
let job_ids: Vec<String> = sqlx::query_scalar(
"SELECT job_id FROM jobs
WHERE phase IN ('running', 'cancelling')
AND ($1::text IS NULL OR job_id > $1)
ORDER BY job_id
LIMIT 200",
)
.bind(after_job.as_deref())
.fetch_all(&state.pool)
.await
.map_err(|_| ())?;
if job_ids.is_empty() {
break;
}
let count = job_ids.len();
for job_id in job_ids {
after_job = Some(job_id.clone());
let Some(row) = load_job_row(&state.pool, &parse_id(job_id.clone()).map_err(|_| ())?)
.await
.map_err(|_| ())?
else {
continue;
};
let progress = current_progress(state, &row)
.await
.ok()
.flatten()
.map(progress_fingerprint);
record_progress_change(&state.pool, &row, progress.as_ref())
.await
.map_err(|_| ())?;
}
if count < 200 {
break;
}
}
Ok(())
}
async fn record_progress_change(
pool: &PgPool,
row: &JobRow,
progress: Option<&ProgressFingerprint>,
) -> Result<(), PublicApiError> {
let Some(attempt_id) = row.authority_attempt_id.as_deref() else {
return Ok(());
};
let mut tx = pool.begin().await?;
let now = OffsetDateTime::now_utc();
let changed = if let Some(progress) = progress {
sqlx::query_scalar::<_, String>(
"INSERT INTO live_progress_event_state (
job_id, attempt_id, fence, progress_present, completed_bytes,
total_bytes, completed_members, total_members, updated_at
) VALUES ($1, $2, $3, TRUE, $4, $5, $6, $7, $8)
ON CONFLICT (job_id) DO UPDATE
SET attempt_id = EXCLUDED.attempt_id,
fence = EXCLUDED.fence,
progress_present = TRUE,
completed_bytes = EXCLUDED.completed_bytes,
total_bytes = EXCLUDED.total_bytes,
completed_members = EXCLUDED.completed_members,
total_members = EXCLUDED.total_members,
updated_at = EXCLUDED.updated_at
WHERE (
live_progress_event_state.attempt_id,
live_progress_event_state.fence,
live_progress_event_state.progress_present,
live_progress_event_state.completed_bytes,
live_progress_event_state.total_bytes,
live_progress_event_state.completed_members,
live_progress_event_state.total_members
) IS DISTINCT FROM (
EXCLUDED.attempt_id,
EXCLUDED.fence,
TRUE,
EXCLUDED.completed_bytes,
EXCLUDED.total_bytes,
EXCLUDED.completed_members,
EXCLUDED.total_members
)
RETURNING job_id",
)
.bind(&row.job_id)
.bind(progress.attempt_id.as_str())
.bind(i64::try_from(progress.fence).map_err(|_| PublicApiError::internal())?)
.bind(i64::try_from(progress.completed_bytes).map_err(|_| PublicApiError::internal())?)
.bind(
progress
.total_bytes
.map(i64::try_from)
.transpose()
.map_err(|_| PublicApiError::internal())?,
)
.bind(i64::try_from(progress.completed_members).map_err(|_| PublicApiError::internal())?)
.bind(
progress
.total_members
.map(i64::try_from)
.transpose()
.map_err(|_| PublicApiError::internal())?,
)
.bind(now)
.fetch_optional(tx.as_mut())
.await?
.is_some()
} else {
sqlx::query_scalar::<_, String>(
"UPDATE live_progress_event_state
SET attempt_id = $2, fence = $3, progress_present = FALSE,
completed_bytes = NULL, total_bytes = NULL,
completed_members = NULL, total_members = NULL, updated_at = $4
WHERE job_id = $1
AND (
attempt_id IS DISTINCT FROM $2
OR fence IS DISTINCT FROM $3
OR progress_present
OR completed_bytes IS NOT NULL
OR total_bytes IS NOT NULL
OR completed_members IS NOT NULL
OR total_members IS NOT NULL
)
RETURNING job_id",
)
.bind(&row.job_id)
.bind(attempt_id)
.bind(row.current_fence)
.bind(now)
.fetch_optional(tx.as_mut())
.await?
.is_some()
};
if changed {
remora_durable::enqueue(
&mut tx,
"job.updated",
serde_json::json!({ "job_id": row.job_id }),
)
.await?;
}
tx.commit().await?;
Ok(())
}
fn job_frame(
row: &JobRow,
progress: Option<&ProgressFingerprint>,
event_id: OpaqueId,
sequence: u64,
cursor: &str,
occurred_at: OffsetDateTime,
) -> Result<EventFrameDto, PublicApiError> {
let authoritative_attempt = row
.authority_attempt_id
.as_deref()
.map(|attempt_id| {
Ok::<_, PublicApiError>(EventAttemptDto {
attempt_id: parse_id(attempt_id)?,
fence: PositiveIntegerDto::new(
u64::try_from(row.current_fence).map_err(|_| PublicApiError::internal())?,
)
.map_err(|_| PublicApiError::internal())?,
})
})
.transpose()?;
let progress = progress
.filter(|progress| {
row.authority_attempt_id.as_deref() == Some(progress.attempt_id.as_str())
&& u64::try_from(row.current_fence).ok() == Some(progress.fence)
})
.map(event_progress)
.transpose()?;
Ok(EventFrameDto::JobChanged(JobChangedDto {
schema_version: SchemaVersionDto,
event_id,
sequence: positive(sequence)?,
cursor: event_cursor_dto(cursor)?,
occurred_at: event_time(occurred_at)?,
job_id: parse_id(row.job_id.clone())?,
job_revision: positive(
u64::try_from(row.revision).map_err(|_| PublicApiError::internal())?,
)?,
phase: parse_phase(&row.phase)?,
lifecycle: parse_lifecycle(&row.lifecycle)?,
authoritative_attempt: NullableDto(authoritative_attempt),
progress: NullableDto(progress),
}))
}
fn deletion_frame(
job_id: OpaqueId,
operation_id: OpaqueId,
state: DeletionStateDto,
event_id: OpaqueId,
sequence: u64,
cursor: &str,
occurred_at: OffsetDateTime,
) -> Result<EventFrameDto, PublicApiError> {
Ok(EventFrameDto::DeletionChanged(DeletionChangedDto {
schema_version: SchemaVersionDto,
event_id,
sequence: positive(sequence)?,
cursor: event_cursor_dto(cursor)?,
occurred_at: event_time(occurred_at)?,
job_id,
operation_id,
state,
}))
}
fn resync_frame(
reason: ResyncReasonDto,
sequence: u64,
cursor: &str,
occurred_at: OffsetDateTime,
) -> Result<EventFrameDto, PublicApiError> {
Ok(EventFrameDto::ResyncRequired(ResyncRequiredDto {
schema_version: SchemaVersionDto,
event_id: request_id(),
sequence: positive(sequence)?,
cursor: event_cursor_dto(cursor)?,
occurred_at: event_time(occurred_at)?,
reason,
}))
}
fn heartbeat_frame(
sequence: u64,
cursor: &str,
occurred_at: OffsetDateTime,
) -> Result<EventFrameDto, PublicApiError> {
Ok(EventFrameDto::Heartbeat(HeartbeatDto {
schema_version: SchemaVersionDto,
event_id: request_id(),
sequence: positive(sequence)?,
cursor: event_cursor_dto(cursor)?,
occurred_at: event_time(occurred_at)?,
}))
}
fn event_progress(progress: &ProgressFingerprint) -> Result<EventProgressDto, PublicApiError> {
Ok(EventProgressDto {
attempt_id: progress.attempt_id.clone(),
fence: positive(progress.fence)?,
completed_bytes: safe(progress.completed_bytes)?,
total_bytes: progress
.total_bytes
.map(|value| safe(value).map(|value| NullableDto(Some(value))))
.transpose()?,
completed_members: safe(progress.completed_members)?,
total_members: progress
.total_members
.map(|value| safe(value).map(|value| NullableDto(Some(value))))
.transpose()?,
})
}
fn parse_payload_id(payload: &Value, field: &str) -> Option<OpaqueId> {
payload
.get(field)?
.as_str()
.and_then(|value| OpaqueId::parse(value).ok())
}
fn stable_event_id(id: i64) -> OpaqueId {
OpaqueId::parse(format!("event_{id:016}")).expect("positive outbox ids form opaque event ids")
}
fn parse_phase(value: &str) -> Result<JobPhase, PublicApiError> {
match value {
"queued" => Ok(JobPhase::Queued),
"running" => Ok(JobPhase::Running),
"retry_wait" => Ok(JobPhase::RetryWait),
"cancelling" => Ok(JobPhase::Cancelling),
"succeeded" => Ok(JobPhase::Succeeded),
"failed" => Ok(JobPhase::Failed),
"cancelled" => Ok(JobPhase::Cancelled),
_ => Err(PublicApiError::internal()),
}
}
fn parse_lifecycle(value: &str) -> Result<JobLifecycle, PublicApiError> {
match value {
"visible" => Ok(JobLifecycle::Visible),
"soft_deleted" => Ok(JobLifecycle::SoftDeleted),
"hard_deleting" => Ok(JobLifecycle::HardDeleting),
"hard_deleted" => Ok(JobLifecycle::HardDeleted),
_ => Err(PublicApiError::internal()),
}
}
fn parse_deletion_state(value: &str) -> Option<DeletionStateDto> {
match value {
"queued" => Some(DeletionStateDto::Queued),
"running" => Some(DeletionStateDto::Running),
"retry_wait" => Some(DeletionStateDto::RetryWait),
"blocked" => Some(DeletionStateDto::Blocked),
"completed" => Some(DeletionStateDto::Completed),
_ => None,
}
}
fn positive(value: u64) -> Result<PositiveIntegerDto, PublicApiError> {
PositiveIntegerDto::new(value).map_err(|_| PublicApiError::internal())
}
fn safe(value: u64) -> Result<SafeIntegerDto, PublicApiError> {
SafeIntegerDto::new(value).map_err(|_| PublicApiError::internal())
}
fn event_cursor_dto(value: &str) -> Result<CursorDto, PublicApiError> {
CursorDto::parse(value).map_err(|_| PublicApiError::internal())
}
fn event_time(value: OffsetDateTime) -> Result<DateTimeDto, PublicApiError> {
DateTimeDto::parse(
value
.format(&Rfc3339)
.map_err(|_| PublicApiError::internal())?,
)
.map_err(|_| PublicApiError::internal())
}
fn next_sequence(sequence: &mut u64) -> u64 {
*sequence = sequence.saturating_add(1);
*sequence
}
async fn send_frame(session: &mut actix_ws::Session, frame: &EventFrameDto) -> Result<(), ()> {
session
.text(serde_json::to_string(frame).map_err(|_| ())?)
.await
.map_err(|_| ())
}
#[cfg(test)]
mod tests {
use remora_domain::{BucketName, StorageScope, TargetPrefix, TargetStoreId};
use remora_durable::{CursorCodec, KeysetPosition, first_page_cursor};
use super::*;
#[test]
fn event_cursor_rejects_other_subject_and_non_event_positions() {
let codec = CursorCodec::new(b"event-cursor-test-secret").unwrap();
let scope = StorageScope {
store: TargetStoreId::parse("target-main").unwrap(),
bucket: BucketName::parse("target-bucket").unwrap(),
prefix: TargetPrefix::parse("models", std::iter::empty()).unwrap(),
};
let principal = Principal {
id: OpaqueId::parse("principal_one").unwrap(),
issuer: "https://issuer.example".into(),
subject: "one".into(),
administrator: false,
scopes: vec![scope],
};
let fingerprint = CursorCodec::filter_fingerprint(&serde_json::json!({
"resource": "events",
"authorization": authorization_snapshot(&principal)
}))
.unwrap();
let issued = OffsetDateTime::now_utc();
let cursor = first_page_cursor(
"principal_two",
fingerprint,
KeysetPosition {
created_at: issued,
id: "event-12".into(),
},
KeysetPosition {
created_at: issued,
id: "event-12".into(),
},
);
let token = codec.encode(&cursor).unwrap();
assert_eq!(
decode_event_cursor(
&codec,
std::time::Duration::from_secs(60),
&principal,
&token,
issued,
)
.unwrap(),
Err(ResyncReasonDto::AuthorizationChanged)
);
let page = first_page_cursor(
principal.id.as_str(),
fingerprint,
KeysetPosition {
created_at: issued,
id: "job_000001".into(),
},
KeysetPosition {
created_at: issued,
id: "job_000001".into(),
},
);
let token = codec.encode(&page).unwrap();
assert_eq!(
decode_event_cursor(
&codec,
std::time::Duration::from_secs(60),
&principal,
&token,
issued,
)
.unwrap(),
Err(ResyncReasonDto::CursorInvalid)
);
}
}

773
crates/remora-api/tests/live_events.rs Normal file
View file

@ -0,0 +1,773 @@
use std::{net::TcpListener, time::Duration};
use actix_web::{App, HttpServer, web};
use futures_util::{SinkExt as _, StreamExt as _};
use remora_admission::{AdmissionConfig, AdmissionService, migrate as admission_migrate};
use remora_api::{
PublicApiConfig, PublicApiServices, PublicApiState, configure_public_api,
migrate as api_migrate, prune_live_event_history,
};
use remora_attempt::{AttemptConfig, AttemptService, migrate as attempt_migrate};
use remora_contracts::{EventFrameDto, ResyncReasonDto};
use remora_coordination::{
ActivateOptions, CoordinationStore, LeaseOwner, ProgressValue, StoreConfig,
};
use remora_credentials::{
CredentialRegistry, VaultClient, VaultConfig, migrate as credentials_migrate,
};
use remora_deletion::{
DeletionConfig, DeletionService, TargetRegistry, migrate as deletion_migrate,
};
use remora_domain::{BucketName, OpaqueId, TargetPrefix, TargetStoreId};
use remora_durable::{PoolConfig, connect, enqueue, migrate as durable_migrate};
use remora_identity::{DefaultTarget, migrate as identity_migrate};
use remora_result::migrate as result_migrate;
use reqwest::{Client, StatusCode};
use serde_json::{Value, json};
use sqlx::PgPool;
use time::OffsetDateTime;
use tokio_tungstenite::{
MaybeTlsStream, WebSocketStream, connect_async,
tungstenite::{
Message,
client::IntoClientRequest,
http::{HeaderValue, header::AUTHORIZATION},
},
};
mod support;
use support::{mint_provider_token, oidc_provider, required_env, revoke_remora_grants};
type Socket = WebSocketStream<MaybeTlsStream<tokio::net::TcpStream>>;
struct Harness {
pool: PgPool,
coordination: CoordinationStore,
client: Client,
token: String,
http_base: String,
ws_base: String,
server: actix_web::dev::ServerHandle,
}
impl Harness {
async fn new() -> Option<Self> {
if std::env::var("REMORA_LIVE_EVENTS_REQUIRE").is_err() {
return None;
}
let pool = connect(
&required_env("DATABASE_URL"),
PoolConfig {
max_connections: 20,
acquire_timeout: Duration::from_secs(10),
idle_timeout: Some(Duration::from_secs(30)),
},
)
.await
.unwrap();
durable_migrate(&pool).await.unwrap();
identity_migrate(&pool).await.unwrap();
credentials_migrate(&pool).await.unwrap();
admission_migrate(&pool).await.unwrap();
attempt_migrate(&pool).await.unwrap();
result_migrate(&pool).await.unwrap();
deletion_migrate(&pool).await.unwrap();
api_migrate(&pool).await.unwrap();
sqlx::query("TRUNCATE TABLE jobs CASCADE")
.execute(&pool)
.await
.unwrap();
sqlx::query("TRUNCATE TABLE outbox_messages RESTART IDENTITY")
.execute(&pool)
.await
.unwrap();
sqlx::query(
"UPDATE live_event_history_state
SET floor_outbox_id = 0, progress_poll_not_before = '-infinity'
WHERE singleton",
)
.execute(&pool)
.await
.unwrap();
let coordination = CoordinationStore::connect(
&required_env("REMORA_COORDINATION_URL"),
format!(
"{}wp22:{}:",
required_env("REMORA_COORDINATION_NAMESPACE"),
nonce()
),
StoreConfig {
global_limit: 16,
lease_ttl: Duration::from_secs(10),
reserve_ttl: Duration::from_secs(10),
progress_ttl: Duration::from_secs(10),
cancel_ttl: None,
},
)
.await
.unwrap();
let attempts = AttemptService::new(
pool.clone(),
coordination.clone(),
AttemptConfig {
max_attempts: 3,
local_limit: 8,
},
)
.unwrap();
let admission = AdmissionService::new(
pool.clone(),
AdmissionConfig {
default_target: DefaultTarget {
store: TargetStoreId::parse("target-main").unwrap(),
bucket: BucketName::parse("target-bucket").unwrap(),
prefix: TargetPrefix::parse("models", ["_remora"]).unwrap(),
},
reserved_target_segments: vec!["_remora".into()],
policy_version: "live-events-v1".into(),
},
)
.unwrap();
let vault = VaultClient::new(VaultConfig {
address: "http://127.0.0.1:8200".into(),
token: "synthetic-unused-token".into(),
mount: "secret".into(),
path_prefix: "unused".into(),
timeout: Duration::from_secs(1),
})
.unwrap();
let credentials = CredentialRegistry::new(pool.clone(), vault);
let targets = TargetRegistry::default();
let deletion =
DeletionService::new(pool.clone(), targets.clone(), DeletionConfig::default()).unwrap();
let state = PublicApiState::new(
pool.clone(),
PublicApiServices {
identity: oidc_provider(),
admission,
attempts,
credentials,
deletion,
targets,
},
PublicApiConfig {
cursor_secret: b"live-events-fixture-cursor-secret".to_vec(),
member_access_ttl: Duration::from_secs(60),
event_retention: Duration::from_millis(500),
event_poll_interval: Duration::from_millis(250),
event_heartbeat_interval: Duration::from_millis(100),
},
)
.unwrap();
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let address = listener.local_addr().unwrap();
let server = HttpServer::new(move || {
App::new()
.app_data(web::Data::new(state.clone()))
.configure(configure_public_api)
})
.listen(listener)
.unwrap()
.run();
let handle = server.handle();
actix_web::rt::spawn(server);
let token = mint_provider_token().await.unwrap();
Some(Self {
pool,
coordination,
client: Client::new(),
token,
http_base: format!("http://{address}/v1"),
ws_base: format!("ws://{address}/v1/events"),
server: handle,
})
}
async fn submit_job(&self) -> Value {
let response = self
.client
.post(format!("{}/jobs", self.http_base))
.bearer_auth(&self.token)
.header("Idempotency-Key", idempotency_key())
.json(&json!({
"label": "live events",
"source": {
"kind": "http",
"url": "https://example.invalid/live",
"result_path": "payload.bin"
}
}))
.send()
.await
.unwrap();
assert_eq!(response.status(), StatusCode::ACCEPTED);
response.json().await.unwrap()
}
async fn get_job(&self, token: &str, job_id: &str) -> reqwest::Response {
self.client
.get(format!("{}/jobs/{job_id}", self.http_base))
.bearer_auth(token)
.send()
.await
.unwrap()
}
async fn list_jobs(&self, token: &str) -> Value {
let response = self
.client
.get(format!("{}/jobs", self.http_base))
.bearer_auth(token)
.send()
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
response.json().await.unwrap()
}
async fn connect(&self, token: &str, cursor: Option<&str>) -> Socket {
let url = cursor.map_or_else(
|| self.ws_base.clone(),
|cursor| format!("{}?cursor={cursor}", self.ws_base),
);
let mut request = url.into_client_request().unwrap();
request.headers_mut().insert(
AUTHORIZATION,
HeaderValue::from_str(&format!("Bearer {token}")).unwrap(),
);
let (socket, response) = connect_async(request).await.unwrap();
assert_eq!(response.status(), 101);
socket
}
}
impl Drop for Harness {
fn drop(&mut self) {
actix_web::rt::spawn(self.server.stop(true));
}
}
#[actix_web::test]
async fn live_event_replay_authorization_progress_and_resync_converge() {
let Some(harness) = Harness::new().await else {
return;
};
let mut absent = harness.connect(&harness.token, None).await;
assert_resync(next_frame(&mut absent).await, ResyncReasonDto::CursorAbsent);
let submitted = harness.submit_job().await;
let job_id = submitted["id"].as_str().unwrap().to_owned();
let initial_cursor = submitted["event_cursor"].as_str().unwrap().to_owned();
let mut invalid = harness
.connect(&harness.token, Some(&format!("{initial_cursor}x")))
.await;
assert_resync(
next_frame(&mut invalid).await,
ResyncReasonDto::CursorInvalid,
);
update_job_and_enqueue(&harness.pool, &job_id).await;
update_job_and_enqueue(&harness.pool, &job_id).await;
let mut first_replay = harness.connect(&harness.token, Some(&initial_cursor)).await;
let first = next_job_frame(&mut first_replay, &job_id).await;
let second = next_job_frame(&mut first_replay, &job_id).await;
let (first_id, first_sequence) = job_event_identity(&first);
let (second_id, second_sequence) = job_event_identity(&second);
assert_ne!(first_id, second_id);
assert_eq!((first_sequence, second_sequence), (1, 2));
let replay_cursor = frame_cursor(&second).to_owned();
let mut duplicate_replay = harness.connect(&harness.token, Some(&initial_cursor)).await;
let duplicate_first = next_job_frame(&mut duplicate_replay, &job_id).await;
let duplicate_second = next_job_frame(&mut duplicate_replay, &job_id).await;
assert_eq!(job_event_identity(&duplicate_first).0, first_id);
assert_eq!(job_event_identity(&duplicate_second).0, second_id);
let mut caught_up = harness.connect(&harness.token, Some(&replay_cursor)).await;
assert!(matches!(
next_frame(&mut caught_up).await,
EventFrameDto::Heartbeat(_)
));
first_replay.close(None).await.unwrap();
duplicate_replay.close(None).await.unwrap();
caught_up.close(None).await.unwrap();
let fresh = harness
.get_job(&harness.token, &job_id)
.await
.json::<Value>()
.await
.unwrap();
let expiring_cursor = fresh["event_cursor"].as_str().unwrap().to_owned();
tokio::time::sleep(Duration::from_millis(600)).await;
let mut expired = harness
.connect(&harness.token, Some(&expiring_cursor))
.await;
assert_resync(
next_frame(&mut expired).await,
ResyncReasonDto::CursorExpired,
);
let before_gap = harness
.get_job(&harness.token, &job_id)
.await
.json::<Value>()
.await
.unwrap()["event_cursor"]
.as_str()
.unwrap()
.to_owned();
update_job_and_enqueue(&harness.pool, &job_id).await;
assert!(
prune_live_event_history(
&harness.pool,
OffsetDateTime::now_utc() + time::Duration::seconds(1)
)
.await
.unwrap()
> 0
);
let mut unavailable = harness.connect(&harness.token, Some(&before_gap)).await;
assert_resync(
next_frame(&mut unavailable).await,
ResyncReasonDto::HistoryUnavailable,
);
let after_gap = harness
.get_job(&harness.token, &job_id)
.await
.json::<Value>()
.await
.unwrap()["event_cursor"]
.as_str()
.unwrap()
.to_owned();
let mut recovered = harness.connect(&harness.token, Some(&after_gap)).await;
assert!(matches!(
next_frame(&mut recovered).await,
EventFrameDto::Heartbeat(_)
));
recovered.close(None).await.unwrap();
let active_gap_cursor = harness
.get_job(&harness.token, &job_id)
.await
.json::<Value>()
.await
.unwrap()["event_cursor"]
.as_str()
.unwrap()
.to_owned();
let mut active_gap = harness
.connect(&harness.token, Some(&active_gap_cursor))
.await;
assert!(matches!(
next_frame(&mut active_gap).await,
EventFrameDto::Heartbeat(_)
));
update_job_and_enqueue(&harness.pool, &job_id).await;
prune_live_event_history(
&harness.pool,
OffsetDateTime::now_utc() + time::Duration::seconds(1),
)
.await
.unwrap();
assert_resync_reason(
next_resync(&mut active_gap).await,
ResyncReasonDto::HistoryUnavailable,
);
let first_attempt = insert_running_attempt(&harness.pool, &job_id, 1).await;
let first_owner = LeaseOwner::new(first_attempt.as_str(), 1, "worker_live_01").unwrap();
assert!(
harness
.coordination
.activate(&job_id, &first_owner, ActivateOptions::default())
.await
.unwrap()
.ok()
);
let progress_cursor = harness
.get_job(&harness.token, &job_id)
.await
.json::<Value>()
.await
.unwrap()["event_cursor"]
.as_str()
.unwrap()
.to_owned();
let mut progress_socket = harness
.connect(&harness.token, Some(&progress_cursor))
.await;
assert!(
harness
.coordination
.report_progress_value(
&job_id,
&first_owner,
ProgressValue {
completed_bytes: 42,
total_bytes: Some(100),
completed_members: 1,
total_members: Some(2),
},
)
.await
.unwrap()
.ok()
);
let progress = next_job_progress(&mut progress_socket, &job_id, Some(42)).await;
assert_matching_progress(&progress, &first_attempt, 1, 42);
let mut progress_replay = harness
.connect(&harness.token, Some(&progress_cursor))
.await;
let replayed_progress = next_job_progress(&mut progress_replay, &job_id, Some(42)).await;
assert_eq!(
job_event_identity(&replayed_progress).0,
job_event_identity(&progress).0
);
let second_attempt = insert_running_attempt(&harness.pool, &job_id, 2).await;
let stale_suppressed = next_job_progress(&mut progress_socket, &job_id, None).await;
assert_job_authority(&stale_suppressed, &second_attempt, 2);
assert!(
harness
.coordination
.release(&job_id, &first_owner)
.await
.unwrap()
.ok()
);
harness.coordination.clean_terminal(&job_id).await.unwrap();
let second_owner = LeaseOwner::new(second_attempt.as_str(), 2, "worker_live_02").unwrap();
assert!(
harness
.coordination
.activate(&job_id, &second_owner, ActivateOptions::default())
.await
.unwrap()
.ok()
);
assert!(
harness
.coordination
.report_progress(&job_id, &second_owner, 77)
.await
.unwrap()
.ok()
);
let restored_progress = next_job_progress(&mut progress_socket, &job_id, Some(77)).await;
assert_matching_progress(&restored_progress, &second_attempt, 2, 77);
assert!(
harness
.coordination
.release(&job_id, &second_owner)
.await
.unwrap()
.ok()
);
harness.coordination.clean_terminal(&job_id).await.unwrap();
let coordination_lost = next_job_progress(&mut progress_socket, &job_id, None).await;
assert_job_authority(&coordination_lost, &second_attempt, 2);
let snapshot_after_loss: Value = harness
.get_job(&harness.token, &job_id)
.await
.json()
.await
.unwrap();
assert!(snapshot_after_loss["progress"].is_null());
transition_cancelled(&harness.pool, &job_id).await;
let deletion_cursor = harness
.get_job(&harness.token, &job_id)
.await
.json::<Value>()
.await
.unwrap()["event_cursor"]
.as_str()
.unwrap()
.to_owned();
let deletion_response = harness
.client
.post(format!("{}/jobs/{job_id}/hard-delete", harness.http_base))
.bearer_auth(&harness.token)
.header("Idempotency-Key", idempotency_key())
.send()
.await
.unwrap();
assert_eq!(deletion_response.status(), StatusCode::ACCEPTED);
let operation: Value = deletion_response.json().await.unwrap();
let operation_id = operation["id"].as_str().unwrap();
let mut deletion_socket = harness
.connect(&harness.token, Some(&deletion_cursor))
.await;
let deletion = next_deletion_frame(&mut deletion_socket, operation_id).await;
let EventFrameDto::DeletionChanged(deletion) = deletion else {
unreachable!()
};
assert_eq!(deletion.job_id.as_str(), job_id);
assert_eq!(
serde_json::to_value(deletion.state).unwrap(),
json!("completed")
);
let authorized_cursor = harness.list_jobs(&harness.token).await["event_cursor"]
.as_str()
.unwrap()
.to_owned();
revoke_remora_grants().await.unwrap();
let reduced_token = mint_provider_token().await.unwrap();
let mut changed = harness
.connect(&reduced_token, Some(&authorized_cursor))
.await;
assert_resync(
next_frame(&mut changed).await,
ResyncReasonDto::AuthorizationChanged,
);
assert_eq!(
harness.get_job(&reduced_token, &job_id).await.status(),
StatusCode::NOT_FOUND
);
let reduced_snapshot = harness.list_jobs(&reduced_token).await;
assert!(reduced_snapshot["items"].as_array().unwrap().is_empty());
let reduced_cursor = reduced_snapshot["event_cursor"].as_str().unwrap();
let mut filtered = harness.connect(&reduced_token, Some(reduced_cursor)).await;
enqueue_job_event(&harness.pool, &job_id).await;
for _ in 0..2 {
assert!(matches!(
next_frame(&mut filtered).await,
EventFrameDto::Heartbeat(_)
));
}
progress_socket.close(None).await.unwrap();
progress_replay.close(None).await.unwrap();
deletion_socket.close(None).await.unwrap();
filtered.close(None).await.unwrap();
}
async fn update_job_and_enqueue(pool: &PgPool, job_id: &str) {
let mut tx = pool.begin().await.unwrap();
sqlx::query("UPDATE jobs SET revision = revision + 1, updated_at = now() WHERE job_id = $1")
.bind(job_id)
.execute(tx.as_mut())
.await
.unwrap();
enqueue(&mut tx, "job.updated", json!({ "job_id": job_id }))
.await
.unwrap();
tx.commit().await.unwrap();
}
async fn enqueue_job_event(pool: &PgPool, job_id: &str) {
let mut tx = pool.begin().await.unwrap();
enqueue(&mut tx, "job.updated", json!({ "job_id": job_id }))
.await
.unwrap();
tx.commit().await.unwrap();
}
async fn insert_running_attempt(pool: &PgPool, job_id: &str, fence: i64) -> OpaqueId {
let attempt_id = opaque_id("attempt");
sqlx::query(
"INSERT INTO attempts (
attempt_id, job_id, fence, worker_incarnation, state, initial_progress,
authorized_at, executing_at, execution_envelope, staging_origin_attempt_id
) VALUES ($1, $2, $3, $4, 'executing', 0, now(), now(), '{}'::jsonb, $1)",
)
.bind(attempt_id.as_str())
.bind(job_id)
.bind(fence)
.bind(format!("worker_live_{fence:02}"))
.execute(pool)
.await
.unwrap();
let mut tx = pool.begin().await.unwrap();
sqlx::query(
"UPDATE jobs
SET phase = 'running', current_fence = $2, authority_attempt_id = $3,
revision = revision + 1, updated_at = now()
WHERE job_id = $1",
)
.bind(job_id)
.bind(fence)
.bind(attempt_id.as_str())
.execute(tx.as_mut())
.await
.unwrap();
enqueue(&mut tx, "job.updated", json!({ "job_id": job_id }))
.await
.unwrap();
tx.commit().await.unwrap();
attempt_id
}
async fn transition_cancelled(pool: &PgPool, job_id: &str) {
let mut tx = pool.begin().await.unwrap();
sqlx::query(
"UPDATE jobs
SET phase = 'cancelled', cancel_requested = TRUE, authority_attempt_id = NULL,
dispatch_needed = FALSE, terminal_at = now(), revision = revision + 1,
updated_at = now()
WHERE job_id = $1",
)
.bind(job_id)
.execute(tx.as_mut())
.await
.unwrap();
enqueue(&mut tx, "job.updated", json!({ "job_id": job_id }))
.await
.unwrap();
tx.commit().await.unwrap();
}
async fn next_frame(socket: &mut Socket) -> EventFrameDto {
tokio::time::timeout(Duration::from_secs(3), async {
loop {
match socket.next().await {
Some(Ok(Message::Text(text))) => {
return serde_json::from_str(text.as_ref()).unwrap();
}
Some(Ok(Message::Ping(bytes))) => {
socket.send(Message::Pong(bytes)).await.unwrap();
}
Some(Ok(Message::Close(reason))) => {
panic!("websocket closed before a frame: {reason:?}");
}
Some(Ok(_)) => {}
Some(Err(error)) => panic!("websocket error: {error}"),
None => panic!("websocket ended before a frame"),
}
}
})
.await
.expect("event frame timeout")
}
async fn next_job_frame(socket: &mut Socket, job_id: &str) -> EventFrameDto {
loop {
let frame = next_frame(socket).await;
if matches!(&frame, EventFrameDto::JobChanged(event) if event.job_id.as_str() == job_id) {
return frame;
}
}
}
async fn next_job_progress(
socket: &mut Socket,
job_id: &str,
completed_bytes: Option<u64>,
) -> EventFrameDto {
loop {
let frame = next_job_frame(socket, job_id).await;
let EventFrameDto::JobChanged(event) = &frame else {
unreachable!()
};
let observed = event
.progress
.0
.as_ref()
.map(|progress| progress.completed_bytes.get());
if observed == completed_bytes {
return frame;
}
}
}
async fn next_deletion_frame(socket: &mut Socket, operation_id: &str) -> EventFrameDto {
loop {
let frame = next_frame(socket).await;
if matches!(
&frame,
EventFrameDto::DeletionChanged(event)
if event.operation_id.as_str() == operation_id
) {
return frame;
}
}
}
async fn next_resync(socket: &mut Socket) -> EventFrameDto {
loop {
let frame = next_frame(socket).await;
if matches!(frame, EventFrameDto::ResyncRequired(_)) {
return frame;
}
}
}
fn assert_resync(frame: EventFrameDto, expected: ResyncReasonDto) {
let sequence = assert_resync_reason(frame, expected);
assert_eq!(sequence, 1);
}
fn assert_resync_reason(frame: EventFrameDto, expected: ResyncReasonDto) -> u64 {
let EventFrameDto::ResyncRequired(frame) = frame else {
panic!("expected resync_required");
};
assert_eq!(frame.reason, expected);
frame.sequence.get()
}
fn job_event_identity(frame: &EventFrameDto) -> (&str, u64) {
let EventFrameDto::JobChanged(frame) = frame else {
panic!("expected job_changed");
};
(frame.event_id.as_str(), frame.sequence.get())
}
fn frame_cursor(frame: &EventFrameDto) -> &str {
match frame {
EventFrameDto::JobChanged(frame) => frame.cursor.as_str(),
EventFrameDto::DeletionChanged(frame) => frame.cursor.as_str(),
EventFrameDto::ResyncRequired(frame) => frame.cursor.as_str(),
EventFrameDto::Heartbeat(frame) => frame.cursor.as_str(),
}
}
fn assert_matching_progress(
frame: &EventFrameDto,
attempt_id: &OpaqueId,
fence: u64,
completed_bytes: u64,
) {
let EventFrameDto::JobChanged(frame) = frame else {
panic!("expected job_changed");
};
let progress = frame.progress.0.as_ref().expect("matching progress");
assert_eq!(progress.attempt_id, *attempt_id);
assert_eq!(progress.fence.get(), fence);
assert_eq!(progress.completed_bytes.get(), completed_bytes);
}
fn assert_job_authority(frame: &EventFrameDto, attempt_id: &OpaqueId, fence: u64) {
let EventFrameDto::JobChanged(frame) = frame else {
panic!("expected job_changed");
};
assert!(frame.progress.0.is_none());
let authority = frame
.authoritative_attempt
.0
.as_ref()
.expect("authoritative attempt");
assert_eq!(authority.attempt_id, *attempt_id);
assert_eq!(authority.fence.get(), fence);
}
fn opaque_id(prefix: &str) -> OpaqueId {
OpaqueId::parse(format!("{prefix}_{}", nonce())).unwrap()
}
fn idempotency_key() -> String {
format!("key-{}-replay-safe", nonce())
}
fn nonce() -> String {
let mut bytes = [0_u8; 8];
getrandom::fill(&mut bytes).unwrap();
bytes.iter().map(|byte| format!("{byte:02x}")).collect()
}

340
crates/remora-api/tests/public_api.rs
View file

@ -1,13 +1,11 @@
use std::{net::TcpListener, time::Duration};
use actix_web::{App, HttpServer, web};
use base64::{
Engine as _,
engine::general_purpose::{STANDARD, URL_SAFE_NO_PAD},
};
use jsonwebtoken::{Algorithm, EncodingKey, Header, encode};
use remora_admission::{AdmissionConfig, AdmissionService, migrate as admission_migrate};
use remora_api::{PublicApiConfig, PublicApiServices, PublicApiState, configure_public_api};
use remora_api::{
PublicApiConfig, PublicApiServices, PublicApiState, configure_public_api,
migrate as api_migrate,
};
use remora_attempt::{AttemptConfig, AttemptService, migrate as attempt_migrate};
use remora_contracts::{
AttemptPageDto, DeletionOperationDto, JobPageDto, JobResourceDto, ResultDescriptorDto,
@ -24,9 +22,7 @@ use remora_deletion::{
};
use remora_domain::{BucketName, OpaqueId, TargetPrefix, TargetStoreId};
use remora_durable::{PoolConfig, connect, migrate as durable_migrate};
use remora_identity::{
DefaultTarget, OidcConfig, OidcIdentityProvider, migrate as identity_migrate,
};
use remora_identity::{DefaultTarget, migrate as identity_migrate};
use remora_result::migrate as result_migrate;
use remora_target::{
CreateObjectRequest, TargetStore, TargetStoreConfig, attempt_manifest_key_in,
@ -37,6 +33,10 @@ use serde_json::{Value, json};
use sqlx::PgPool;
use time::OffsetDateTime;
mod support;
use support::{mint_provider_token, oidc_provider, required_env, revoke_remora_grants};
struct Harness {
pool: PgPool,
coordination: CoordinationStore,
@ -68,6 +68,7 @@ impl Harness {
attempt_migrate(&pool).await.unwrap();
result_migrate(&pool).await.unwrap();
deletion_migrate(&pool).await.unwrap();
api_migrate(&pool).await.unwrap();
sqlx::query("TRUNCATE TABLE jobs CASCADE")
.execute(&pool)
.await
@ -140,6 +141,9 @@ impl Harness {
PublicApiConfig {
cursor_secret: b"public-api-fixture-cursor-secret".to_vec(),
member_access_ttl: Duration::from_secs(60),
event_retention: Duration::from_secs(60),
event_poll_interval: Duration::from_millis(50),
event_heartbeat_interval: Duration::from_secs(1),
},
)
.unwrap();
@ -737,6 +741,14 @@ async fn public_routes_are_authenticated_replayable_authorized_and_redacted() {
.await
.unwrap();
assert_eq!(lost_scope.status(), StatusCode::NOT_FOUND);
let hidden_submit_replay = harness
.request_with_token(reqwest::Method::POST, "/jobs", &unprivileged_token)
.header("Idempotency-Key", &submit_key)
.json(&submission)
.send()
.await
.unwrap();
assert_eq!(hidden_submit_replay.status(), StatusCode::NOT_FOUND);
let hidden_peer_job = harness
.request_with_token(
reqwest::Method::GET,
@ -957,316 +969,6 @@ async fn remove_mutation_record_failure(pool: &PgPool) {
.unwrap();
}
async fn revoke_remora_grants() -> Result<(), String> {
match required_env("REMORA_OIDC_PROVIDER").as_str() {
"keycloak" => revoke_keycloak_remora_roles().await,
"zitadel" => revoke_zitadel_remora_roles().await,
provider => Err(format!("unsupported OIDC provider {provider}")),
}
}
async fn revoke_keycloak_remora_roles() -> Result<(), String> {
let realm = required_env("REMORA_OIDC_TENANT");
let client_id = required_env("REMORA_OIDC_CLIENT_ID");
let username = required_env("REMORA_OIDC_USERNAME");
let client = Client::new();
let admin_response = client
.post("http://keycloak.remora.test:8080/realms/master/protocol/openid-connect/token")
.form(&[
("grant_type", "password"),
("client_id", "admin-cli"),
("username", "remora-admin"),
("password", "synthetic-keycloak"),
])
.send()
.await
.map_err(|error| error.to_string())?;
let admin_body: Value = admin_response
.json()
.await
.map_err(|error| error.to_string())?;
let admin_token = admin_body["access_token"]
.as_str()
.ok_or_else(|| "keycloak admin token missing".to_owned())?;
let users: Value = client
.get(format!(
"http://keycloak.remora.test:8080/admin/realms/{realm}/users?username={username}&exact=true"
))
.bearer_auth(admin_token)
.send()
.await
.map_err(|error| error.to_string())?
.json()
.await
.map_err(|error| error.to_string())?;
let user_id = users
.as_array()
.and_then(|rows| rows.first())
.and_then(|row| row["id"].as_str())
.ok_or_else(|| "keycloak user missing".to_owned())?;
let clients: Value = client
.get(format!(
"http://keycloak.remora.test:8080/admin/realms/{realm}/clients?clientId={client_id}"
))
.bearer_auth(admin_token)
.send()
.await
.map_err(|error| error.to_string())?
.json()
.await
.map_err(|error| error.to_string())?;
let client_uuid = clients
.as_array()
.and_then(|rows| rows.first())
.and_then(|row| row["id"].as_str())
.ok_or_else(|| "keycloak client missing".to_owned())?;
let roles_url = format!(
"http://keycloak.remora.test:8080/admin/realms/{realm}/users/{user_id}/role-mappings/clients/{client_uuid}"
);
let mapped: Value = client
.get(&roles_url)
.bearer_auth(admin_token)
.send()
.await
.map_err(|error| error.to_string())?
.json()
.await
.map_err(|error| error.to_string())?;
let roles = mapped
.as_array()
.cloned()
.unwrap_or_default()
.into_iter()
.filter(|role| {
role["name"]
.as_str()
.is_some_and(|name| name.starts_with("remora:v1:"))
})
.collect::<Vec<_>>();
let response = client
.delete(roles_url)
.bearer_auth(admin_token)
.json(&roles)
.send()
.await
.map_err(|error| error.to_string())?;
if response.status().is_success() {
Ok(())
} else {
Err(format!(
"keycloak role revocation failed: {}",
response.status()
))
}
}
async fn revoke_zitadel_remora_roles() -> Result<(), String> {
let issuer = required_env("REMORA_OIDC_ISSUER");
let organization = required_env("REMORA_OIDC_TENANT");
let user_id = required_env("REMORA_OIDC_USER_ID");
let project_id = required_env("REMORA_OIDC_PROJECT_ID");
let token = required_env("REMORA_OIDC_MGMT_TOKEN");
let host = issuer
.strip_prefix("http://")
.or_else(|| issuer.strip_prefix("https://"))
.unwrap_or(&issuer);
let client = Client::new();
let grants: Value = client
.post(format!("{issuer}/management/v1/users/grants/_search"))
.header("Host", host)
.bearer_auth(&token)
.header("x-zitadel-orgid", &organization)
.json(&json!({ "queries": [{ "userIdQuery": { "userId": user_id } }] }))
.send()
.await
.map_err(|error| error.to_string())?
.json()
.await
.map_err(|error| error.to_string())?;
let grant_id = grants["result"]
.as_array()
.and_then(|rows| {
rows.iter().find(|row| {
row["projectId"]
.as_str()
.is_some_and(|value| value == project_id)
})
})
.and_then(|row| row["id"].as_str())
.ok_or_else(|| "zitadel project grant missing".to_owned())?;
let response = client
.put(format!(
"{issuer}/management/v1/users/{user_id}/grants/{grant_id}"
))
.header("Host", host)
.bearer_auth(token)
.header("x-zitadel-orgid", organization)
.json(&json!({ "roleKeys": ["test-user"] }))
.send()
.await
.map_err(|error| error.to_string())?;
if response.status().is_success() {
Ok(())
} else {
Err(format!(
"zitadel role revocation failed: {}",
response.status()
))
}
}
fn oidc_provider() -> OidcIdentityProvider {
let provider = required_env("REMORA_OIDC_PROVIDER");
let audience = if provider == "zitadel" {
required_env("REMORA_OIDC_PROJECT_ID")
} else {
required_env("REMORA_OIDC_CLIENT_ID")
};
let mut config = OidcConfig::new(required_env("REMORA_OIDC_ISSUER"), audience).unwrap();
if provider == "zitadel" {
let project_id = required_env("REMORA_OIDC_PROJECT_ID");
config = config
.with_zitadel_project_roles()
.with_project_role_claim(format!("urn:zitadel:iam:org:project:{project_id}:roles"))
.with_resource_access_client(required_env("REMORA_OIDC_CLIENT_ID"));
} else {
config = config.with_resource_access_client(required_env("REMORA_OIDC_CLIENT_ID"));
}
OidcIdentityProvider::new(config).unwrap()
}
async fn mint_provider_token() -> Result<String, String> {
match required_env("REMORA_OIDC_PROVIDER").as_str() {
"keycloak" => mint_keycloak_token().await,
"zitadel" => mint_zitadel_token().await,
provider => Err(format!("unsupported OIDC provider {provider}")),
}
}
async fn mint_keycloak_token() -> Result<String, String> {
let response = Client::new()
.post(format!(
"{}/protocol/openid-connect/token",
required_env("REMORA_OIDC_ISSUER").trim_end_matches('/')
))
.form(&[
("grant_type", "password"),
("client_id", required_env("REMORA_OIDC_CLIENT_ID").as_str()),
("username", required_env("REMORA_OIDC_USERNAME").as_str()),
("password", required_env("REMORA_OIDC_PASSWORD").as_str()),
("scope", "openid"),
])
.send()
.await
.map_err(|error| error.to_string())?;
token_from_response(response).await
}
async fn mint_zitadel_token() -> Result<String, String> {
let issuer = required_env("REMORA_OIDC_ISSUER")
.trim_end_matches('/')
.to_owned();
let token_url = format!("{issuer}/oauth/v2/token");
let encoded_key = required_env("REMORA_OIDC_USER_KEY_B64");
let key_json = URL_SAFE_NO_PAD
.decode(encoded_key.as_bytes())
.or_else(|_| STANDARD.decode(encoded_key.as_bytes()))
.map_err(|error| format!("decode user key: {error}"))?;
let key: Value =
serde_json::from_slice(&key_json).map_err(|error| format!("parse user key: {error}"))?;
let key_id = key
.get("keyId")
.or_else(|| key.get("key_id"))
.and_then(Value::as_str)
.ok_or_else(|| "user key missing keyId".to_owned())?
.to_owned();
let subject = key
.get("userId")
.or_else(|| key.get("userID"))
.and_then(Value::as_str)
.map(str::to_owned)
.unwrap_or_else(|| required_env("REMORA_OIDC_USER_ID"));
let private_key = key["key"]
.as_str()
.ok_or_else(|| "user key missing private key".to_owned())?
.replace("\\n", "\n");
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map_err(|error| error.to_string())?
.as_secs() as i64;
let encoding =
EncodingKey::from_rsa_pem(private_key.as_bytes()).map_err(|error| error.to_string())?;
let project_id = required_env("REMORA_OIDC_PROJECT_ID");
let host = issuer
.strip_prefix("http://")
.or_else(|| issuer.strip_prefix("https://"))
.unwrap_or(&issuer);
let scope = format!(
"openid profile \
urn:zitadel:iam:org:project:id:{project_id}:aud \
urn:zitadel:iam:org:project:roles \
urn:zitadel:iam:org:projects:roles \
urn:zitadel:iam:org:project:{project_id}:roles"
);
let client = Client::new();
let mut last_error = String::new();
for audience in [&token_url, &issuer] {
let mut header = Header::new(Algorithm::RS256);
header.kid = Some(key_id.clone());
let assertion = encode(
&header,
&json!({
"iss": subject,
"sub": subject,
"aud": audience,
"iat": now,
"exp": now + 3600
}),
&encoding,
)
.map_err(|error| error.to_string())?;
let response = client
.post(&token_url)
.header("Host", host)
.form(&[
("grant_type", "urn:ietf:params:oauth:grant-type:jwt-bearer"),
("assertion", assertion.as_str()),
("scope", scope.as_str()),
])
.send()
.await
.map_err(|error| error.to_string())?;
if response.status().is_success() {
return token_from_response(response).await;
}
last_error = format!(
"token endpoint: {} {}",
response.status(),
response.text().await.unwrap_or_default()
);
}
Err(last_error)
}
async fn token_from_response(response: reqwest::Response) -> Result<String, String> {
if !response.status().is_success() {
return Err(format!(
"token endpoint: {} {}",
response.status(),
response.text().await.unwrap_or_default()
));
}
let body: Value = response.json().await.map_err(|error| error.to_string())?;
body.get("access_token")
.and_then(Value::as_str)
.map(str::to_owned)
.ok_or_else(|| "token response missing access_token".into())
}
fn required_env(name: &str) -> String {
std::env::var(name).unwrap_or_else(|_| panic!("{name} is required"))
}
fn target_store() -> TargetStore {
let mut config = TargetStoreConfig::from_env().unwrap();
config.target_store_id = "target-main".into();

318
crates/remora-api/tests/support/mod.rs Normal file
View file

@ -0,0 +1,318 @@
use base64::{
Engine as _,
engine::general_purpose::{STANDARD, URL_SAFE_NO_PAD},
};
use jsonwebtoken::{Algorithm, EncodingKey, Header, encode};
use remora_identity::{OidcConfig, OidcIdentityProvider};
use reqwest::Client;
use serde_json::{Value, json};
pub fn required_env(name: &str) -> String {
std::env::var(name).unwrap_or_else(|_| panic!("{name} is required"))
}
pub fn oidc_provider() -> OidcIdentityProvider {
let provider = required_env("REMORA_OIDC_PROVIDER");
let audience = if provider == "zitadel" {
required_env("REMORA_OIDC_PROJECT_ID")
} else {
required_env("REMORA_OIDC_CLIENT_ID")
};
let mut config = OidcConfig::new(required_env("REMORA_OIDC_ISSUER"), audience).unwrap();
if provider == "zitadel" {
let project_id = required_env("REMORA_OIDC_PROJECT_ID");
config = config
.with_zitadel_project_roles()
.with_project_role_claim(format!("urn:zitadel:iam:org:project:{project_id}:roles"))
.with_resource_access_client(required_env("REMORA_OIDC_CLIENT_ID"));
} else {
config = config.with_resource_access_client(required_env("REMORA_OIDC_CLIENT_ID"));
}
OidcIdentityProvider::new(config).unwrap()
}
pub async fn mint_provider_token() -> Result<String, String> {
match required_env("REMORA_OIDC_PROVIDER").as_str() {
"keycloak" => mint_keycloak_token().await,
"zitadel" => mint_zitadel_token().await,
provider => Err(format!("unsupported OIDC provider {provider}")),
}
}
pub async fn revoke_remora_grants() -> Result<(), String> {
match required_env("REMORA_OIDC_PROVIDER").as_str() {
"keycloak" => revoke_keycloak_remora_roles().await,
"zitadel" => revoke_zitadel_remora_roles().await,
provider => Err(format!("unsupported OIDC provider {provider}")),
}
}
async fn mint_keycloak_token() -> Result<String, String> {
let response = Client::new()
.post(format!(
"{}/protocol/openid-connect/token",
required_env("REMORA_OIDC_ISSUER").trim_end_matches('/')
))
.form(&[
("grant_type", "password"),
("client_id", required_env("REMORA_OIDC_CLIENT_ID").as_str()),
("username", required_env("REMORA_OIDC_USERNAME").as_str()),
("password", required_env("REMORA_OIDC_PASSWORD").as_str()),
("scope", "openid"),
])
.send()
.await
.map_err(|error| error.to_string())?;
token_from_response(response).await
}
async fn mint_zitadel_token() -> Result<String, String> {
let issuer = required_env("REMORA_OIDC_ISSUER")
.trim_end_matches('/')
.to_owned();
let token_url = format!("{issuer}/oauth/v2/token");
let encoded_key = required_env("REMORA_OIDC_USER_KEY_B64");
let key_json = URL_SAFE_NO_PAD
.decode(encoded_key.as_bytes())
.or_else(|_| STANDARD.decode(encoded_key.as_bytes()))
.map_err(|error| format!("decode user key: {error}"))?;
let key: Value =
serde_json::from_slice(&key_json).map_err(|error| format!("parse user key: {error}"))?;
let key_id = key
.get("keyId")
.or_else(|| key.get("key_id"))
.and_then(Value::as_str)
.ok_or_else(|| "user key missing keyId".to_owned())?
.to_owned();
let subject = key
.get("userId")
.or_else(|| key.get("userID"))
.and_then(Value::as_str)
.map(str::to_owned)
.unwrap_or_else(|| required_env("REMORA_OIDC_USER_ID"));
let private_key = key["key"]
.as_str()
.ok_or_else(|| "user key missing private key".to_owned())?
.replace("\\n", "\n");
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map_err(|error| error.to_string())?
.as_secs() as i64;
let encoding =
EncodingKey::from_rsa_pem(private_key.as_bytes()).map_err(|error| error.to_string())?;
let project_id = required_env("REMORA_OIDC_PROJECT_ID");
let host = issuer
.strip_prefix("http://")
.or_else(|| issuer.strip_prefix("https://"))
.unwrap_or(&issuer);
let scope = format!(
"openid profile \
urn:zitadel:iam:org:project:id:{project_id}:aud \
urn:zitadel:iam:org:project:roles \
urn:zitadel:iam:org:projects:roles \
urn:zitadel:iam:org:project:{project_id}:roles"
);
let client = Client::new();
let mut last_error = String::new();
for audience in [&token_url, &issuer] {
let mut header = Header::new(Algorithm::RS256);
header.kid = Some(key_id.clone());
let assertion = encode(
&header,
&json!({
"iss": subject,
"sub": subject,
"aud": audience,
"iat": now,
"exp": now + 3600
}),
&encoding,
)
.map_err(|error| error.to_string())?;
let response = client
.post(&token_url)
.header("Host", host)
.form(&[
("grant_type", "urn:ietf:params:oauth:grant-type:jwt-bearer"),
("assertion", assertion.as_str()),
("scope", scope.as_str()),
])
.send()
.await
.map_err(|error| error.to_string())?;
if response.status().is_success() {
return token_from_response(response).await;
}
last_error = format!(
"token endpoint: {} {}",
response.status(),
response.text().await.unwrap_or_default()
);
}
Err(last_error)
}
async fn token_from_response(response: reqwest::Response) -> Result<String, String> {
if !response.status().is_success() {
return Err(format!(
"token endpoint: {} {}",
response.status(),
response.text().await.unwrap_or_default()
));
}
let body: Value = response.json().await.map_err(|error| error.to_string())?;
body.get("access_token")
.and_then(Value::as_str)
.map(str::to_owned)
.ok_or_else(|| "token response missing access_token".into())
}
async fn revoke_keycloak_remora_roles() -> Result<(), String> {
let realm = required_env("REMORA_OIDC_TENANT");
let client_id = required_env("REMORA_OIDC_CLIENT_ID");
let username = required_env("REMORA_OIDC_USERNAME");
let client = Client::new();
let admin_response = client
.post("http://keycloak.remora.test:8080/realms/master/protocol/openid-connect/token")
.form(&[
("grant_type", "password"),
("client_id", "admin-cli"),
("username", "remora-admin"),
("password", "synthetic-keycloak"),
])
.send()
.await
.map_err(|error| error.to_string())?;
let admin_body: Value = admin_response
.json()
.await
.map_err(|error| error.to_string())?;
let admin_token = admin_body["access_token"]
.as_str()
.ok_or_else(|| "keycloak admin token missing".to_owned())?;
let users: Value = client
.get(format!(
"http://keycloak.remora.test:8080/admin/realms/{realm}/users?username={username}&exact=true"
))
.bearer_auth(admin_token)
.send()
.await
.map_err(|error| error.to_string())?
.json()
.await
.map_err(|error| error.to_string())?;
let user_id = users
.as_array()
.and_then(|rows| rows.first())
.and_then(|row| row["id"].as_str())
.ok_or_else(|| "keycloak user missing".to_owned())?;
let clients: Value = client
.get(format!(
"http://keycloak.remora.test:8080/admin/realms/{realm}/clients?clientId={client_id}"
))
.bearer_auth(admin_token)
.send()
.await
.map_err(|error| error.to_string())?
.json()
.await
.map_err(|error| error.to_string())?;
let client_uuid = clients
.as_array()
.and_then(|rows| rows.first())
.and_then(|row| row["id"].as_str())
.ok_or_else(|| "keycloak client missing".to_owned())?;
let roles_url = format!(
"http://keycloak.remora.test:8080/admin/realms/{realm}/users/{user_id}/role-mappings/clients/{client_uuid}"
);
let mapped: Value = client
.get(&roles_url)
.bearer_auth(admin_token)
.send()
.await
.map_err(|error| error.to_string())?
.json()
.await
.map_err(|error| error.to_string())?;
let roles = mapped
.as_array()
.cloned()
.unwrap_or_default()
.into_iter()
.filter(|role| {
role["name"]
.as_str()
.is_some_and(|name| name.starts_with("remora:v1:"))
})
.collect::<Vec<_>>();
let response = client
.delete(roles_url)
.bearer_auth(admin_token)
.json(&roles)
.send()
.await
.map_err(|error| error.to_string())?;
if response.status().is_success() {
Ok(())
} else {
Err(format!(
"keycloak role revocation failed: {}",
response.status()
))
}
}
async fn revoke_zitadel_remora_roles() -> Result<(), String> {
let issuer = required_env("REMORA_OIDC_ISSUER");
let organization = required_env("REMORA_OIDC_TENANT");
let user_id = required_env("REMORA_OIDC_USER_ID");
let project_id = required_env("REMORA_OIDC_PROJECT_ID");
let token = required_env("REMORA_OIDC_MGMT_TOKEN");
let host = issuer
.strip_prefix("http://")
.or_else(|| issuer.strip_prefix("https://"))
.unwrap_or(&issuer);
let client = Client::new();
let grants: Value = client
.post(format!("{issuer}/management/v1/users/grants/_search"))
.header("Host", host)
.bearer_auth(&token)
.header("x-zitadel-orgid", &organization)
.json(&json!({ "queries": [{ "userIdQuery": { "userId": user_id } }] }))
.send()
.await
.map_err(|error| error.to_string())?
.json()
.await
.map_err(|error| error.to_string())?;
let grant_id = grants["result"]
.as_array()
.and_then(|rows| {
rows.iter().find(|row| {
row["projectId"]
.as_str()
.is_some_and(|value| value == project_id)
})
})
.and_then(|row| row["id"].as_str())
.ok_or_else(|| "zitadel project grant missing".to_owned())?;
let response = client
.put(format!(
"{issuer}/management/v1/users/{user_id}/grants/{grant_id}"
))
.header("Host", host)
.bearer_auth(token)
.header("x-zitadel-orgid", organization)
.json(&json!({ "roleKeys": ["test-user"] }))
.send()
.await
.map_err(|error| error.to_string())?;
if response.status().is_success() {
Ok(())
} else {
Err(format!(
"zitadel role revocation failed: {}",
response.status()
))
}
}

3
crates/remora-attempt/src/mutations.rs
View file

@ -342,6 +342,9 @@ impl AttemptService {
json!({ "job_id": authority.job_id }),
)
.await?;
if retry {
enqueue(tx, "job.updated", json!({ "job_id": authority.job_id })).await?;
}
accepted_phase(if retry {
JobPhase::RetryWait
} else {

9
crates/remora-attempt/src/recovery.rs
View file

@ -668,6 +668,8 @@ impl AttemptService {
.retry_not_before
.is_some_and(|due| due <= OffsetDateTime::now_utc())
{
let mut tx = self.pool.begin().await?;
let now = OffsetDateTime::now_utc();
let changed = sqlx::query(
"UPDATE jobs SET phase = 'queued', dispatch_needed = TRUE,
retry_not_before = NULL, revision = revision + 1, updated_at = $2
@ -675,14 +677,16 @@ impl AttemptService {
AND retry_not_before <= $2",
)
.bind(job_id.as_str())
.bind(OffsetDateTime::now_utc())
.execute(&self.pool)
.bind(now)
.execute(tx.as_mut())
.await?;
if changed.rows_affected() == 1 {
enqueue(&mut tx, "job.updated", json!({ "job_id": job_id })).await?;
report.retries_scheduled += 1;
row.phase = "queued".into();
row.dispatch_needed = true;
}
tx.commit().await?;
}
}
if row.phase == "queued"
@ -1222,6 +1226,7 @@ async fn reap_lost_attempt(
.execute(tx.as_mut())
.await?;
enqueue(tx, "job.retry_wait", json!({ "job_id": job_id })).await?;
enqueue(tx, "job.updated", json!({ "job_id": job_id })).await?;
} else {
sqlx::query(
"UPDATE jobs SET phase = 'failed', authority_attempt_id = NULL,

2
crates/remora-attempt/src/service.rs
View file

@ -501,6 +501,7 @@ impl AttemptService {
));
}
enqueue(tx, "job.dispatch", json!({ "job_id": job_id })).await?;
enqueue(tx, "job.updated", json!({ "job_id": job_id })).await?;
Ok(RecoveryOutcome::Reaped)
})
})
@ -816,6 +817,7 @@ async fn authorize_attempt(
.bind(checkpoint_owner_after_claim)
.execute(tx.as_mut())
.await?;
enqueue(tx, "job.updated", json!({ "job_id": input.job_id })).await?;
Ok(ClaimOutcome::Claimed {
authority: AttemptAuthority {
job_id: input.job_id.clone(),

54
crates/remora-deletion/src/service.rs
View file

@ -183,6 +183,12 @@ impl DeletionService {
json!({ "job_id": job_id, "deletion_operation_id": operation_id }),
)
.await?;
enqueue(
tx,
"deletion.updated",
json!({ "job_id": job_id, "operation_id": operation_id }),
)
.await?;
let operation = load_operation_in_tx(tx, operation_id.as_str()).await?;
Ok(HardDeleteOutcome {
mutation: LifecycleMutation::Changed,
@ -224,6 +230,7 @@ impl DeletionService {
) -> Result<Option<DeleteClaim>, DeletionError> {
let claim_token = new_id('c')?;
let claim_seconds = duration_seconds(self.config.claim_ttl)?;
let mut tx = self.pool.begin().await?;
let row = sqlx::query_as::<_, ClaimRow>(
"WITH candidate AS (
SELECT i.item_id
@ -257,19 +264,29 @@ impl DeletionService {
.bind(operation_id.as_str())
.bind(claim_token.as_str())
.bind(claim_seconds)
.fetch_optional(&self.pool)
.fetch_optional(tx.as_mut())
.await?;
if row.is_some() {
sqlx::query(
let job_id: Option<String> = sqlx::query_scalar(
"UPDATE deletion_operations
SET state = 'running', problem_class = NULL, problem_message = NULL,
updated_at = now()
WHERE operation_id = $1 AND state <> 'completed'",
WHERE operation_id = $1 AND state NOT IN ('running', 'completed')
RETURNING job_id",
)
.bind(operation_id.as_str())
.execute(&self.pool)
.fetch_optional(tx.as_mut())
.await?;
if let Some(job_id) = job_id {
enqueue(
&mut tx,
"deletion.updated",
json!({ "job_id": job_id, "operation_id": operation_id }),
)
.await?;
}
}
tx.commit().await?;
row.map(ClaimRow::into_claim).transpose()
}
@ -377,7 +394,7 @@ impl DeletionService {
.await?;
let completed = finalize_if_empty(
&mut tx,
&OpaqueId::parse(job_id).map_err(durable_effect)?,
&OpaqueId::parse(&job_id).map_err(durable_effect)?,
&claim.operation_id,
now,
)
@ -393,6 +410,12 @@ impl DeletionService {
.bind(now)
.execute(tx.as_mut())
.await?;
enqueue(
&mut tx,
"deletion.updated",
json!({ "job_id": job_id, "operation_id": claim.operation_id }),
)
.await?;
}
tx.commit().await?;
Ok(if completed {
@ -433,6 +456,17 @@ impl DeletionService {
.bind(message)
.execute(tx.as_mut())
.await?;
let job_id: String =
sqlx::query_scalar("SELECT job_id FROM deletion_operations WHERE operation_id = $1")
.bind(claim.operation_id.as_str())
.fetch_one(tx.as_mut())
.await?;
enqueue(
&mut tx,
"deletion.updated",
json!({ "job_id": job_id, "operation_id": claim.operation_id }),
)
.await?;
tx.commit().await?;
Ok(())
}
@ -793,7 +827,7 @@ async fn finalize_if_empty(
if remaining {
return Ok(false);
}
sqlx::query(
let operation_changed = sqlx::query(
"UPDATE deletion_operations
SET state = 'completed', problem_class = NULL, problem_message = NULL,
updated_at = $2, completed_at = $2
@ -803,6 +837,14 @@ async fn finalize_if_empty(
.bind(now)
.execute(tx.as_mut())
.await?;
if operation_changed.rows_affected() == 1 {
enqueue(
tx,
"deletion.updated",
json!({ "job_id": job_id, "operation_id": operation_id }),
)
.await?;
}
let changed = sqlx::query(
"UPDATE jobs
SET lifecycle = 'hard_deleted', revision = revision + 1, updated_at = $3