Implement crash-safe streaming worker runtime

crates/remora-object-source/src/config.rs

@@ -12,6 +12,7 @@ pub struct ObjectSourceConfig {
     region: String,
     access_key_id: String,
     secret_access_key: String,
+    session_token: Option<String>,
 }
 
 impl ObjectSourceConfig {
@@ -21,6 +22,24 @@ impl ObjectSourceConfig {
         region: impl Into<String>,
         access_key_id: impl Into<String>,
         secret_access_key: impl Into<String>,
+    ) -> Result<Self, ObjectSourceError> {
+        Self::new_with_session_token(
+            endpoint,
+            bucket,
+            region,
+            access_key_id,
+            secret_access_key,
+            None,
+        )
+    }
+
+    pub fn new_with_session_token(
+        endpoint: &str,
+        bucket: impl Into<String>,
+        region: impl Into<String>,
+        access_key_id: impl Into<String>,
+        secret_access_key: impl Into<String>,
+        session_token: Option<String>,
     ) -> Result<Self, ObjectSourceError> {
         let endpoint =
             CanonicalUrl::parse(endpoint).map_err(|_| ObjectSourceError::InvalidRequest)?;
@@ -28,7 +47,11 @@ impl ObjectSourceConfig {
         let region = region.into();
         let access_key_id = access_key_id.into();
         let secret_access_key = secret_access_key.into();
-        if bucket.is_empty() || access_key_id.is_empty() || secret_access_key.is_empty() {
+        if bucket.is_empty()
+            || access_key_id.is_empty()
+            || secret_access_key.is_empty()
+            || session_token.as_ref().is_some_and(String::is_empty)
+        {
             return Err(ObjectSourceError::InvalidRequest);
         }
         Ok(Self {
@@ -37,6 +60,7 @@ impl ObjectSourceConfig {
             region,
             access_key_id,
             secret_access_key,
+            session_token,
         })
     }
 
@@ -71,6 +95,10 @@ impl ObjectSourceConfig {
     pub(crate) fn secret_access_key(&self) -> &str {
         &self.secret_access_key
     }
+
+    pub(crate) fn session_token(&self) -> Option<&str> {
+        self.session_token.as_deref()
+    }
 }
 
 impl fmt::Debug for ObjectSourceConfig {

crates/remora-object-source/src/identity.rs

@@ -51,7 +51,7 @@ impl ObjectSourceIdentity {
     ) -> Result<Self, ObjectSourceError> {
         let bucket = bucket.into();
         let key = key.into();
-        let version_id = version_id.filter(|value| !value.is_empty() && value != "null");
+        let version_id = version_id.filter(|value| !value.is_empty());
         if bucket.is_empty()
             || bucket.chars().any(char::is_control)
             || key.is_empty()

crates/remora-object-source/src/source.rs

@@ -14,6 +14,7 @@ use crate::{
 pub struct ObjectSourceRequest {
     key: String,
     result_path: MemberPath,
+    version_id: Option<String>,
 }
 
 impl ObjectSourceRequest {
@@ -26,8 +27,21 @@ impl ObjectSourceRequest {
             key,
             result_path: MemberPath::parse(result_path)
                 .map_err(|_| ObjectSourceError::InvalidRequest)?,
+            version_id: None,
         })
     }
+
+    pub fn with_version_id(
+        mut self,
+        version_id: impl Into<String>,
+    ) -> Result<Self, ObjectSourceError> {
+        let version_id = version_id.into();
+        if version_id.is_empty() || version_id.chars().any(char::is_control) {
+            return Err(ObjectSourceError::InvalidRequest);
+        }
+        self.version_id = Some(version_id);
+        Ok(self)
+    }
 }
 
 impl fmt::Debug for ObjectSourceRequest {
@@ -112,7 +126,7 @@ pub struct ObjectSource {
 impl ObjectSource {
     pub fn connect(config: ObjectSourceConfig) -> Result<Self, ObjectSourceError> {
         opendal::install_default();
-        let builder = S3::default()
+        let mut builder = S3::default()
             .bucket(config.bucket())
             .endpoint(config.endpoint().as_url().as_str())
             .region(config.region())
@@ -120,6 +134,9 @@ impl ObjectSource {
             .secret_access_key(config.secret_access_key())
             .disable_config_load()
             .disable_ec2_metadata();
+        if let Some(token) = config.session_token() {
+            builder = builder.session_token(token);
+        }
         let operator = Operator::new(builder).map_err(map_open_error)?;
         Ok(Self { config, operator })
     }
@@ -128,17 +145,24 @@ impl ObjectSource {
         &self,
         request: ObjectSourceRequest,
     ) -> Result<OpenedObjectSource, ObjectSourceError> {
-        let metadata = self
-            .operator
-            .stat(&request.key)
-            .await
-            .map_err(map_open_error)?;
+        let mut stat = self.operator.stat_with(&request.key);
+        if let Some(version_id) = &request.version_id {
+            stat = stat.version(version_id);
+        }
+        let metadata = stat.await.map_err(map_open_error)?;
+        if request.version_id.as_deref().is_some_and(|requested| {
+            normalized_version(metadata.version()).as_deref() != Some(requested)
+        }) {
+            return Err(ObjectSourceError::SourceChanged);
+        }
         let identity = ObjectSourceIdentity::from_parts(
             self.config.endpoint().clone(),
             self.config.bucket().to_owned(),
             request.key,
             metadata.content_length(),
-            normalized_version(metadata.version()),
+            request
+                .version_id
+                .or_else(|| normalized_version(metadata.version())),
             metadata.etag().map(ObjectValidator::parse).transpose()?,
         )?;
         let body = self.open_body(&identity, 0, false).await?;
@@ -255,7 +279,5 @@ impl ObjectSource {
 }
 
 fn normalized_version(version: Option<&str>) -> Option<String> {
-    version
-        .filter(|value| !value.is_empty() && *value != "null")
-        .map(str::to_owned)
+    version.filter(|value| !value.is_empty()).map(str::to_owned)
 }