Add authenticated internal worker API

crates/remora-attempt/src/service.rs

@@ -14,9 +14,10 @@ use sqlx::{PgPool, types::Json};
 use time::OffsetDateTime;
 
 use crate::{
-    ActivationOutcome, AttemptAuthority, AttemptError, AttemptFaultInjector, AttemptFaultPoint,
-    AttemptState, CheckpointProjection, ClaimOutcome, ClaimRejection, RecoveryOutcome,
-    RenewalOutcome,
+    ActivationMutationOutcome, ActivationOutcome, AttemptAuthority, AttemptError,
+    AttemptFaultInjector, AttemptFaultPoint, AttemptState, CheckpointProjection, ClaimOutcome,
+    ClaimRejection, ExecutionEnvelopeRecord, LeaseProjection, RecoveryOutcome,
+    RenewalMutationOutcome, RenewalOutcome,
 };
 
 struct NoFaults;
@@ -179,6 +180,23 @@ impl AttemptService {
         authority: &AttemptAuthority,
         operation_id: &OpaqueId,
     ) -> Result<ActivationOutcome, AttemptError> {
+        Ok(
+            match self.activate_operation(authority, operation_id).await? {
+                ActivationMutationOutcome::Activated { disposition, .. } => {
+                    ActivationOutcome::Activated(disposition)
+                }
+                ActivationMutationOutcome::PreconditionFailed => {
+                    ActivationOutcome::PreconditionFailed
+                }
+            },
+        )
+    }
+
+    pub async fn activate_operation(
+        &self,
+        authority: &AttemptAuthority,
+        operation_id: &OpaqueId,
+    ) -> Result<ActivationMutationOutcome, AttemptError> {
         validate_authority(authority)?;
         let request_fingerprint = fingerprint(&ActivationRequest {
             job_id: &authority.job_id,
@@ -203,16 +221,17 @@ impl AttemptService {
             request_fingerprint,
             move |tx| {
                 Box::pin(async move {
-                    let accepted =
+                    let lease =
                         activate_attempt(tx, &coordination, &authority_for_effect, faults.as_ref())
                             .await?;
                     Ok((
-                        if accepted {
+                        if lease.is_some() {
                             MutationStatus::Accepted
                         } else {
                             MutationStatus::Rejected
                         },
-                        json!({ "activated": accepted }),
+                        serde_json::to_value(ActivationStored { lease })
+                            .map_err(DurableError::effect)?,
                     ))
                 })
             },
@@ -225,7 +244,14 @@ impl AttemptService {
                     .exact_authority_is_executing_and_live(authority)
                     .await?
                 {
-                    return Ok(ActivationOutcome::Activated(MutationDisposition::First));
+                    let lease = self
+                        .live_lease(authority)
+                        .await?
+                        .ok_or_else(|| AttemptError::Invalid("live lease disappeared".into()))?;
+                    return Ok(ActivationMutationOutcome::Activated {
+                        disposition: MutationDisposition::First,
+                        lease,
+                    });
                 }
                 self.release_exact(authority).await;
                 return Err(error.into());
@@ -237,16 +263,24 @@ impl AttemptService {
         }
         if outcome.status == MutationStatus::Rejected {
             self.release_exact(authority).await;
-            return Ok(ActivationOutcome::PreconditionFailed);
+            return Ok(ActivationMutationOutcome::PreconditionFailed);
         }
         if !self
             .exact_authority_is_executing_and_live(authority)
             .await?
         {
             self.release_exact(authority).await;
-            return Ok(ActivationOutcome::PreconditionFailed);
+            return Ok(ActivationMutationOutcome::PreconditionFailed);
         }
-        Ok(ActivationOutcome::Activated(outcome.disposition))
+        let stored: ActivationStored = serde_json::from_value(outcome.response)
+            .map_err(|error| AttemptError::Invalid(error.to_string()))?;
+        let lease = stored
+            .lease
+            .ok_or_else(|| AttemptError::Invalid("accepted activation omitted its lease".into()))?;
+        Ok(ActivationMutationOutcome::Activated {
+            disposition: outcome.disposition,
+            lease,
+        })
     }
 
     pub async fn renew(
@@ -269,6 +303,103 @@ impl AttemptService {
         Ok(RenewalOutcome::Renewed)
     }
 
+    pub async fn renew_operation(
+        &self,
+        authority: &AttemptAuthority,
+        operation_id: &OpaqueId,
+    ) -> Result<RenewalMutationOutcome, AttemptError> {
+        validate_authority(authority)?;
+        let request_fingerprint = fingerprint(&ActivationRequest {
+            job_id: &authority.job_id,
+            attempt_id: &authority.attempt_id,
+            fence: authority.fence,
+            worker_incarnation: &authority.worker_incarnation,
+        })
+        .map_err(|error| AttemptError::Invalid(error.to_string()))?;
+        let eligibility = authority_binding(authority);
+        let authority = authority.clone();
+        let coordination = self.coordination.clone();
+        let faults = self.faults.clone();
+        let outcome = execute_mutation(
+            &self.pool,
+            MutationScope {
+                subject: &self.operation_subject,
+                operation: "attempt.mutation",
+                key: operation_id.as_str(),
+            },
+            "renew_attempt",
+            &eligibility,
+            request_fingerprint,
+            move |tx| {
+                Box::pin(async move {
+                    let Some(row) = lock_authority(tx, &authority.attempt_id).await? else {
+                        return Ok((
+                            MutationStatus::Rejected,
+                            serde_json::to_value(RenewalMutationOutcome::AuthorityLost {
+                                disposition: None,
+                            })
+                            .map_err(DurableError::effect)?,
+                        ));
+                    };
+                    if row.state != AttemptState::Executing
+                        || !row.matches(
+                            &authority.job_id,
+                            authority.fence,
+                            &authority.worker_incarnation,
+                        )
+                    {
+                        return Ok((
+                            MutationStatus::Rejected,
+                            serde_json::to_value(RenewalMutationOutcome::AuthorityLost {
+                                disposition: None,
+                            })
+                            .map_err(DurableError::effect)?,
+                        ));
+                    }
+                    faults.hit(AttemptFaultPoint::RenewalBeforeLease);
+                    let expected = lease_owner(&authority).map_err(DurableError::effect)?;
+                    let renewed = coordination
+                        .renew(authority.job_id.as_str(), &expected)
+                        .await
+                        .map_err(DurableError::effect)?;
+                    faults.hit(AttemptFaultPoint::RenewalAfterLease);
+                    let live = coordination
+                        .get_live_lease(authority.job_id.as_str())
+                        .await
+                        .map_err(DurableError::effect)?;
+                    let response = match live
+                        .filter(|lease| renewed.ok() && lease.record.owner() == expected)
+                    {
+                        Some(lease) => RenewalMutationOutcome::Renewed {
+                            lease: LeaseProjection {
+                                expires_at_ms: lease.record.expires_at,
+                                observed_at_ms: lease.now_ms,
+                            },
+                            disposition: None,
+                        },
+                        None => RenewalMutationOutcome::AuthorityLost { disposition: None },
+                    };
+                    let status = if matches!(response, RenewalMutationOutcome::Renewed { .. }) {
+                        MutationStatus::Accepted
+                    } else {
+                        MutationStatus::Rejected
+                    };
+                    Ok((
+                        status,
+                        serde_json::to_value(response).map_err(DurableError::effect)?,
+                    ))
+                })
+            },
+        )
+        .await?;
+        if outcome.disposition == MutationDisposition::Conflict {
+            return Err(AttemptError::OperationConflict);
+        }
+        let response: RenewalMutationOutcome = serde_json::from_value(outcome.response)
+            .map_err(|error| AttemptError::Invalid(error.to_string()))?;
+        Ok(response.with_disposition(outcome.disposition))
+    }
+
     pub async fn recover_inactive(
         &self,
         authority: &AttemptAuthority,
@@ -372,6 +503,45 @@ impl AttemptService {
             .map_err(Into::into)
     }
 
+    pub async fn historical_authority(
+        &self,
+        attempt_id: &OpaqueId,
+    ) -> Result<Option<AttemptAuthority>, AttemptError> {
+        load_historical_authority(&self.pool, attempt_id)
+            .await
+            .map_err(Into::into)
+    }
+
+    pub async fn execution_envelope(
+        &self,
+        attempt_id: &OpaqueId,
+    ) -> Result<Option<ExecutionEnvelopeRecord>, AttemptError> {
+        let value = sqlx::query_scalar::<_, Option<Json<ExecutionEnvelopeRecord>>>(
+            "SELECT execution_envelope FROM attempts WHERE attempt_id = $1",
+        )
+        .bind(attempt_id.as_str())
+        .fetch_optional(&self.pool)
+        .await?
+        .flatten();
+        Ok(value.map(|value| value.0))
+    }
+
+    pub async fn live_lease(
+        &self,
+        authority: &AttemptAuthority,
+    ) -> Result<Option<LeaseProjection>, AttemptError> {
+        let expected = lease_owner(authority)?;
+        Ok(self
+            .coordination
+            .get_live_lease(authority.job_id.as_str())
+            .await?
+            .filter(|lease| lease.record.owner() == expected)
+            .map(|lease| LeaseProjection {
+                expires_at_ms: lease.record.expires_at,
+                observed_at_ms: lease.now_ms,
+            }))
+    }
+
     pub fn release_local(&self, job_id: &OpaqueId) -> bool {
         self.local_permits.release_all(job_id.as_str())
     }
@@ -451,6 +621,11 @@ struct ActivationRequest<'a> {
     worker_incarnation: &'a OpaqueId,
 }
 
+#[derive(serde::Deserialize, Serialize)]
+struct ActivationStored {
+    lease: Option<LeaseProjection>,
+}
+
 struct AuthorityRow {
     job_id: OpaqueId,
     fence: u64,
@@ -534,11 +709,14 @@ async fn authorize_attempt(
         _ => 0,
     };
     let now = OffsetDateTime::now_utc();
+    let envelope =
+        load_execution_envelope_record(tx, input.job_id, input.attempt_id, fence, input.worker)
+            .await?;
     sqlx::query(
         "INSERT INTO attempts (
             attempt_id, job_id, fence, worker_incarnation, state,
-            initial_progress, authorized_at, executing_at, reaped_at
-         ) VALUES ($1, $2, $3, $4, 'authorized', $5, $6, NULL, NULL)",
+            initial_progress, authorized_at, executing_at, reaped_at, execution_envelope
+         ) VALUES ($1, $2, $3, $4, 'authorized', $5, $6, NULL, NULL, $7)",
     )
     .bind(input.attempt_id.as_str())
     .bind(input.job_id.as_str())
@@ -546,6 +724,7 @@ async fn authorize_attempt(
     .bind(input.worker.as_str())
     .bind(i64::try_from(initial_progress).map_err(DurableError::effect)?)
     .bind(now)
+    .bind(Json(&envelope))
     .execute(tx.as_mut())
     .await?;
     sqlx::query(
@@ -572,21 +751,79 @@ async fn authorize_attempt(
     })
 }
 
+async fn load_execution_envelope_record(
+    tx: &mut remora_durable::PgTx<'_>,
+    job_id: &OpaqueId,
+    attempt_id: &OpaqueId,
+    fence: u64,
+    worker: &OpaqueId,
+) -> Result<ExecutionEnvelopeRecord, DurableError> {
+    type Row = (
+        Json<serde_json::Value>,
+        String,
+        String,
+        String,
+        Json<Vec<String>>,
+        String,
+        Option<String>,
+        Option<String>,
+        Option<Json<crate::DurableDescriptor>>,
+        Option<Json<CheckpointProjection>>,
+    );
+    let row = sqlx::query_as::<_, Row>(
+        "SELECT j.source_projection, j.sensitive_source_ref, j.target_store,
+                j.target_bucket, j.target_prefix, j.execution_grant_id,
+                g.credential_id, g.credential_version_id, j.resolved_source, j.checkpoint
+         FROM jobs j
+         JOIN execution_grants g ON g.grant_id = j.execution_grant_id
+         WHERE j.job_id = $1",
+    )
+    .bind(job_id.as_str())
+    .fetch_one(tx.as_mut())
+    .await?;
+    let grant_id = OpaqueId::parse(row.5).map_err(DurableError::effect)?;
+    Ok(ExecutionEnvelopeRecord {
+        id: grant_id.clone(),
+        job_id: job_id.clone(),
+        attempt_id: attempt_id.clone(),
+        fence,
+        worker_incarnation: worker.clone(),
+        source: row.0.0,
+        sensitive_source_ref: OpaqueId::parse(row.1).map_err(DurableError::effect)?,
+        credential_id: row
+            .6
+            .map(OpaqueId::parse)
+            .transpose()
+            .map_err(DurableError::effect)?,
+        credential_version_id: row
+            .7
+            .map(OpaqueId::parse)
+            .transpose()
+            .map_err(DurableError::effect)?,
+        resolved_source: row.8.map(|value| value.0),
+        target_store: row.2,
+        target_bucket: row.3,
+        target_prefix: row.4.0,
+        execution_grant_id: grant_id,
+        checkpoint: row.9.map(|value| value.0),
+    })
+}
+
 async fn activate_attempt(
     tx: &mut remora_durable::PgTx<'_>,
     coordination: &CoordinationStore,
     authority: &AttemptAuthority,
     faults: &dyn AttemptFaultInjector,
-) -> Result<bool, DurableError> {
+) -> Result<Option<LeaseProjection>, DurableError> {
     let Some(row) = lock_authority(tx, &authority.attempt_id).await? else {
-        return Ok(false);
+        return Ok(None);
     };
     if !row.matches(
         &authority.job_id,
         authority.fence,
         &authority.worker_incarnation,
     ) {
-        return Ok(false);
+        return Ok(None);
     }
     if row.state == AttemptState::Executing {
         let live = coordination
@@ -595,16 +832,28 @@ async fn activate_attempt(
             .map_err(DurableError::effect)?;
         let expected = lease_owner(authority).map_err(DurableError::effect)?;
         if !live.is_some_and(|lease| lease.record.owner() == expected) {
-            return Ok(false);
+            return Ok(None);
         }
-        return Ok(coordination
+        if !coordination
             .renew(authority.job_id.as_str(), &expected)
             .await
             .map_err(DurableError::effect)?
-            .ok());
+            .ok()
+        {
+            return Ok(None);
+        }
+        return Ok(coordination
+            .get_live_lease(authority.job_id.as_str())
+            .await
+            .map_err(DurableError::effect)?
+            .filter(|lease| lease.record.owner() == expected)
+            .map(|lease| LeaseProjection {
+                expires_at_ms: lease.record.expires_at,
+                observed_at_ms: lease.now_ms,
+            }));
     }
     if row.state != AttemptState::Authorized {
-        return Ok(false);
+        return Ok(None);
     }
     faults.hit(AttemptFaultPoint::ActivationBeforeLease);
     let owner = lease_owner(authority).map_err(DurableError::effect)?;
@@ -620,7 +869,7 @@ async fn activate_attempt(
         .await
         .map_err(DurableError::effect)?;
     if !coordinated.ok() {
-        return Ok(false);
+        return Ok(None);
     }
     faults.hit(AttemptFaultPoint::ActivationAfterLease);
     let live = coordination
@@ -628,9 +877,9 @@ async fn activate_attempt(
         .await
         .map_err(DurableError::effect)?;
     let expected = lease_owner(authority).map_err(DurableError::effect)?;
-    if !live.is_some_and(|lease| lease.record.owner() == expected) {
-        return Ok(false);
-    }
+    let Some(live) = live.filter(|lease| lease.record.owner() == expected) else {
+        return Ok(None);
+    };
     let now = OffsetDateTime::now_utc();
     sqlx::query(
         "UPDATE attempts SET state = 'executing', executing_at = $2
@@ -640,7 +889,10 @@ async fn activate_attempt(
     .bind(now)
     .execute(tx.as_mut())
     .await?;
-    Ok(true)
+    Ok(Some(LeaseProjection {
+        expires_at_ms: live.record.expires_at,
+        observed_at_ms: live.now_ms,
+    }))
 }
 
 async fn lock_authority(
@@ -692,6 +944,30 @@ async fn load_authority(
     .transpose()
 }
 
+async fn load_historical_authority(
+    pool: &PgPool,
+    attempt_id: &OpaqueId,
+) -> Result<Option<AttemptAuthority>, DurableError> {
+    let row = sqlx::query_as::<_, (String, String, i64, String, String)>(
+        "SELECT job_id, attempt_id, fence, worker_incarnation, state
+         FROM attempts WHERE attempt_id = $1",
+    )
+    .bind(attempt_id.as_str())
+    .fetch_optional(pool)
+    .await?;
+    row.map(|(job, attempt, fence, worker, state)| {
+        Ok(AttemptAuthority {
+            job_id: OpaqueId::parse(job).map_err(DurableError::effect)?,
+            attempt_id: OpaqueId::parse(attempt).map_err(DurableError::effect)?,
+            fence: u64::try_from(fence).map_err(DurableError::effect)?,
+            worker_incarnation: OpaqueId::parse(worker).map_err(DurableError::effect)?,
+            state: AttemptState::parse(&state)
+                .ok_or_else(|| DurableError::Config(format!("unknown attempt state '{state}'")))?,
+        })
+    })
+    .transpose()
+}
+
 fn authority_row(row: (String, i64, String, String)) -> Result<AuthorityRow, DurableError> {
     Ok(AuthorityRow {
         job_id: OpaqueId::parse(row.0).map_err(DurableError::effect)?,