awa_worker/

context.rs

pub use crate::runtime::ProgressState;
use crate::storage::RuntimeStorage;
use awa_model::{AwaError, CallbackConfig, JobRow, QueueStorage};
use sqlx::PgPool;
use std::any::Any;
use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::Duration;

/// Proof that this job registered an external callback in the database.
///
/// The public `id` can be sent to the external system. `#[non_exhaustive]`
/// keeps external callers from constructing this type directly.
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct CallbackGuard {
    pub id: uuid::Uuid,
}

impl CallbackGuard {
    fn new(id: uuid::Uuid) -> Self {
        Self { id }
    }

    /// Return the callback UUID persisted for this job.
    pub fn id(&self) -> uuid::Uuid {
        self.id
    }

    #[cfg(feature = "__python-bridge")]
    #[doc(hidden)]
    pub fn from_bridge_token(id: uuid::Uuid) -> Self {
        Self::new(id)
    }
}

#[doc(hidden)]
pub type CallbackToken = CallbackGuard;

/// Context passed to worker handlers during job execution.
///
/// Provides access to the job metadata, shared state (e.g., service dependencies),
/// callback registration, and structured progress reporting.
pub struct JobContext {
    /// The raw job row from the database.
    pub job: JobRow,
    /// Cancellation flag — set to true when shutdown or deadline is signalled.
    cancelled: Arc<AtomicBool>,
    /// Shared state map for dependency injection.
    state: Arc<HashMap<std::any::TypeId, Box<dyn Any + Send + Sync>>>,
    /// Database pool for callback registration and progress flush.
    pool: PgPool,
    /// Active runtime storage backend.
    storage: RuntimeStorage,
    /// Shared progress buffer — written by handler, read by heartbeat service.
    progress: Arc<std::sync::Mutex<ProgressState>>,
}

impl JobContext {
    pub(crate) fn new(
        job: JobRow,
        cancelled: Arc<AtomicBool>,
        state: Arc<HashMap<std::any::TypeId, Box<dyn Any + Send + Sync>>>,
        pool: PgPool,
        storage: RuntimeStorage,
        progress: Arc<std::sync::Mutex<ProgressState>>,
    ) -> Self {
        Self {
            job,
            cancelled,
            state,
            pool,
            storage,
            progress,
        }
    }

    #[doc(hidden)]
    pub fn new_for_testing(
        job: JobRow,
        cancelled: Arc<AtomicBool>,
        state: Arc<HashMap<std::any::TypeId, Box<dyn Any + Send + Sync>>>,
        pool: PgPool,
        progress: Arc<std::sync::Mutex<ProgressState>>,
    ) -> Self {
        Self::new(
            job,
            cancelled,
            state,
            pool,
            RuntimeStorage::Canonical,
            progress,
        )
    }

    /// Check if this job's execution has been cancelled (shutdown or deadline).
    pub fn is_cancelled(&self) -> bool {
        self.cancelled.load(Ordering::SeqCst)
    }

    /// Clone the shared cancellation flag for language bridges.
    pub fn cancellation_flag(&self) -> Arc<AtomicBool> {
        self.cancelled.clone()
    }

    /// Signal cancellation for this job.
    pub fn cancel(&self) {
        self.cancelled.store(true, Ordering::SeqCst);
    }

    /// Extract a shared state value by type.
    ///
    /// State values are registered via `Client::builder().state(value)`.
    /// Register the concrete type you want to extract:
    ///
    /// ```ignore
    /// // Register with the type you'll extract:
    /// let deps = Arc::new(MyDeps::new());
    /// Client::builder(pool)
    ///     .state(deps.clone())  // stores Arc<MyDeps>
    ///     .build()?;
    ///
    /// // In handler — extract the same type:
    /// let deps = ctx.extract::<Arc<MyDeps>>().unwrap();
    /// ```
    pub fn extract<T: Any + Send + Sync + Clone>(&self) -> Option<T> {
        self.state
            .get(&std::any::TypeId::of::<T>())
            .and_then(|v| v.downcast_ref::<T>())
            .cloned()
    }

    /// Get a reference to the database pool.
    pub fn pool(&self) -> &PgPool {
        &self.pool
    }

    /// Get the active queue-storage backend for this job, if any.
    pub fn queue_storage_store(&self) -> Option<Arc<QueueStorage>> {
        self.storage.queue_storage_store()
    }

    /// Register a callback for this job, writing the callback_id to the database
    /// immediately.
    ///
    /// Call this BEFORE sending the callback_id to the external system to avoid
    /// the race condition where the external system fires before the DB knows
    /// about the callback.
    ///
    /// Returns a `CallbackGuard` whose `id` should be included in the URL or
    /// payload sent to the external system.
    pub async fn register_callback(&self, timeout: Duration) -> Result<CallbackGuard, AwaError> {
        let callback_id = match &self.storage {
            RuntimeStorage::Canonical => {
                awa_model::admin::register_callback(
                    &self.pool,
                    self.job.id,
                    self.job.run_lease,
                    timeout,
                )
                .await?
            }
            RuntimeStorage::QueueStorage(runtime) => {
                runtime
                    .store
                    .register_callback(&self.pool, self.job.id, self.job.run_lease, timeout)
                    .await?
            }
        };
        Ok(CallbackGuard::new(callback_id))
    }

    /// Register a callback with CEL expressions for automatic resolution.
    ///
    /// See [`CallbackConfig`] for expression semantics.
    pub async fn register_callback_with_config(
        &self,
        timeout: Duration,
        config: &CallbackConfig,
    ) -> Result<CallbackGuard, AwaError> {
        let callback_id = match &self.storage {
            RuntimeStorage::Canonical => {
                awa_model::admin::register_callback_with_config(
                    &self.pool,
                    self.job.id,
                    self.job.run_lease,
                    timeout,
                    config,
                )
                .await?
            }
            RuntimeStorage::QueueStorage(runtime) => {
                runtime
                    .store
                    .register_callback_with_config(
                        &self.pool,
                        self.job.id,
                        self.job.run_lease,
                        timeout,
                        config,
                    )
                    .await?
            }
        };
        Ok(CallbackGuard::new(callback_id))
    }

    /// Wait for an external callback to resolve, then resume with the payload.
    ///
    /// This enables sequential callbacks: the handler can register a callback,
    /// wait for it, process the result, register another callback, wait again,
    /// and so on — all within a single handler invocation.
    ///
    /// The handler's async task suspends while waiting. The job transitions to
    /// `waiting_external`. When `resume_external(callback_id, payload)` is
    /// called, the job transitions back to `running` and this method returns
    /// the payload.
    ///
    /// The handler holds its permit (worker slot) during the wait. For very
    /// long waits, consider whether the single-shot `WaitForCallback` pattern
    /// (which releases the permit) is more appropriate.
    ///
    /// ```ignore
    /// let token = ctx.register_callback(Duration::from_secs(3600)).await?;
    /// send_to_external_system(token.id());
    /// let payload = ctx.wait_for_callback(token).await?;
    /// // Handler resumes here with the external system's response
    /// ```
    pub async fn wait_for_callback(
        &self,
        guard: CallbackGuard,
    ) -> Result<serde_json::Value, AwaError> {
        use awa_model::admin::CallbackPollResult;

        let callback_id = guard.id();

        let entered = match &self.storage {
            RuntimeStorage::Canonical => {
                awa_model::admin::enter_callback_wait(
                    &self.pool,
                    self.job.id,
                    self.job.run_lease,
                    callback_id,
                )
                .await?
            }
            RuntimeStorage::QueueStorage(runtime) => {
                runtime
                    .store
                    .enter_callback_wait(&self.pool, self.job.id, self.job.run_lease, callback_id)
                    .await?
            }
        };

        let check_state = || async {
            match &self.storage {
                RuntimeStorage::Canonical => {
                    awa_model::admin::check_callback_state(&self.pool, self.job.id, callback_id)
                        .await
                }
                RuntimeStorage::QueueStorage(runtime) => {
                    runtime
                        .store
                        .check_callback_state(&self.pool, self.job.id, callback_id)
                        .await
                }
            }
        };

        if !entered {
            match check_state().await? {
                CallbackPollResult::Resolved(payload) => return Ok(payload),
                CallbackPollResult::Pending => { /* Already in waiting_external */ }
                CallbackPollResult::Stale {
                    token,
                    current,
                    state,
                } => {
                    return Err(AwaError::Validation(format!(
                        "wait_for_callback: token {token} is stale; current callback is {current} in state {state:?}"
                    )));
                }
                CallbackPollResult::UnexpectedState { token, state } => {
                    return Err(AwaError::Validation(format!(
                        "wait_for_callback: job is not waiting on callback {token}; state={state:?}"
                    )));
                }
                CallbackPollResult::NotFound => {
                    return Err(AwaError::Validation(
                        "job not found during callback wait".into(),
                    ));
                }
            }
        }

        loop {
            if self.is_cancelled() {
                return Err(AwaError::Validation(
                    "job cancelled while waiting for callback".into(),
                ));
            }

            match check_state().await? {
                CallbackPollResult::Resolved(payload) => return Ok(payload),
                CallbackPollResult::Pending => {
                    tokio::time::sleep(Duration::from_millis(200)).await;
                }
                CallbackPollResult::Stale {
                    token,
                    current,
                    state,
                } => {
                    return Err(AwaError::Validation(format!(
                        "wait_for_callback: token {token} is stale; current callback is {current} in state {state:?}"
                    )));
                }
                CallbackPollResult::UnexpectedState { token, state } => {
                    return Err(AwaError::Validation(format!(
                        "job left wait_for_callback unexpectedly for token {token}: state={state:?}"
                    )));
                }
                CallbackPollResult::NotFound => {
                    return Err(AwaError::Validation(
                        "job not found during callback wait".into(),
                    ));
                }
            }
        }
    }

    /// Set structured progress (0-100 with message). Sync — writes to in-memory buffer.
    ///
    /// `percent` is clamped to 0-100. For progress without a message, pass `""`.
    pub fn set_progress(&self, percent: u8, message: &str) {
        let mut guard = self.progress.lock().expect("progress lock poisoned");
        guard.set_progress(percent, Some(message));
    }

    /// Shallow-merge keys into progress.metadata for checkpointing. Sync.
    ///
    /// `updates` must be a JSON object. Top-level keys overwrite; nested objects
    /// are replaced, not deep-merged.
    pub fn update_metadata(&self, updates: serde_json::Value) -> Result<(), AwaError> {
        let obj = updates
            .as_object()
            .ok_or_else(|| AwaError::Validation("update_metadata requires a JSON object".into()))?;

        let mut guard = self.progress.lock().expect("progress lock poisoned");
        if !guard.merge_metadata(obj) {
            return Err(AwaError::Validation(
                "progress.metadata is not a JSON object; cannot merge".into(),
            ));
        }
        Ok(())
    }

    /// Force immediate flush of pending progress to DB. For critical checkpoints.
    ///
    /// Does not return success until the progress has been durably written
    /// or the job is no longer in running state (rescued/cancelled).
    pub async fn flush_progress(&self) -> Result<(), AwaError> {
        let (snapshot, target_generation) = {
            let guard = self.progress.lock().expect("progress lock poisoned");
            match guard.pending_snapshot() {
                Some(pair) => pair,
                None => return Ok(()),
            }
        };

        match &self.storage {
            RuntimeStorage::Canonical => {
                let result = sqlx::query(
                    r#"
                    UPDATE awa.jobs_hot
                    SET progress = $2
                    WHERE id = $1 AND state = 'running' AND run_lease = $3
                    "#,
                )
                .bind(self.job.id)
                .bind(&snapshot)
                .bind(self.job.run_lease)
                .execute(&self.pool)
                .await?;

                if result.rows_affected() == 0 {
                    return Ok(());
                }
            }
            RuntimeStorage::QueueStorage(runtime) => {
                runtime
                    .store
                    .flush_progress(
                        &self.pool,
                        self.job.id,
                        self.job.run_lease,
                        snapshot.clone(),
                    )
                    .await?;
            }
        }

        let mut guard = self.progress.lock().expect("progress lock poisoned");
        guard.ack(target_generation);

        Ok(())
    }

    /// Get a clone of the shared progress state Arc (for Python bridge).
    pub fn progress_buffer(&self) -> Arc<std::sync::Mutex<ProgressState>> {
        self.progress.clone()
    }
}