spiral-sdk 0.1.10

//! Task polling and collection workflows over generated resource operations.

use std::time::{Duration, Instant};

use spiral_contract::call::CallError;
use uuid::Uuid;

use super::generated::tasks::AuthenticatedClient;
pub use crate::api::tasks::*;
use crate::ids::ProjectId;
/// Authenticated Task operations generated from the service contract.
pub type HttpTaskClient = AuthenticatedClient;

/// Polling policy for waiting until a Task finishes.
#[derive(Clone, Copy, Debug)]
pub struct TaskWaitOptions {
    /// Initial delay between status reads, doubling to at most ten seconds.
    pub poll_interval: Duration,
    /// Optional client deadline; reaching it does not cancel the Task.
    pub timeout: Option<Duration>,
}

impl Default for TaskWaitOptions {
    fn default() -> Self {
        Self {
            poll_interval: Duration::from_secs(1),
            timeout: None,
        }
    }
}

/// A Task wait deadline expired or a status read failed.
#[derive(Debug, thiserror::Error)]
pub enum TaskWaitError {
    /// The Task may still be running after the client stops polling.
    #[error("timed out before Task {task_id} finished")]
    Timeout { task_id: Uuid },
    /// The generated status request failed.
    #[error(transparent)]
    Service(#[from] CallError<TasksError>),
}

impl TaskStatus {
    /// Whether the Task has finished and will not run another Attempt.
    pub fn is_terminal(&self) -> bool {
        matches!(self, Self::Succeeded | Self::Failed | Self::Cancelled)
    }
}

impl AuthenticatedClient {
    /// Poll until the Task finishes, tolerating up to four consecutive transient failures.
    pub async fn wait(
        &self,
        task_id: Uuid,
        options: TaskWaitOptions,
    ) -> Result<Task, TaskWaitError> {
        let started = Instant::now();
        let mut interval = options.poll_interval;
        let mut failures = 0;
        loop {
            match self.read_task(task_id).await {
                Ok(details) => {
                    failures = 0;
                    if details.task.status.is_terminal() {
                        return Ok(details.task);
                    }
                }
                Err(error) => {
                    failures += 1;
                    if !transient(&error) || failures >= 5 {
                        return Err(error.into());
                    }
                }
            }
            if options
                .timeout
                .is_some_and(|limit| started.elapsed() >= limit)
            {
                return Err(TaskWaitError::Timeout { task_id });
            }
            tokio::time::sleep(interval).await;
            interval = interval.saturating_mul(2).min(Duration::from_secs(10));
        }
    }

    /// Collect every Worker pool page using the server's opaque continuation tokens.
    pub async fn list_all_worker_pools(
        &self,
        project_id: &ProjectId,
    ) -> Result<ListWorkerPoolsResponse, CallError<TasksError>> {
        let mut request = ListWorkerPoolsRequest {
            page_size: Some(1000),
            page_token: None,
            worker_pool_id: None,
        };
        let mut response = self
            .list_worker_pools(project_id.to_string(), request.clone())
            .await?;
        while let Some(token) = response.next_page_token.take() {
            request.page_token = Some(token);
            let page = self
                .list_worker_pools(project_id.to_string(), request.clone())
                .await?;
            response.worker_pools.extend(page.worker_pools);
            response.stats.extend(page.stats);
            response.next_page_token = page.next_page_token;
        }
        Ok(response)
    }

    /// Apply one lifecycle action to an existing Worker pool.
    pub async fn worker_pool_action(
        &self,
        project_id: &ProjectId,
        worker_pool_id: Uuid,
        action: WorkerPoolAction,
    ) -> Result<WorkerPoolResponse, CallError<TasksError>> {
        match action {
            WorkerPoolAction::Activate => {
                self.activate_worker_pool(project_id.to_string(), worker_pool_id)
                    .await
            }
            WorkerPoolAction::Pause => {
                self.pause_worker_pool(project_id.to_string(), worker_pool_id)
                    .await
            }
            WorkerPoolAction::Retire => {
                self.retire_worker_pool(project_id.to_string(), worker_pool_id)
                    .await
            }
        }
    }
}

/// A lifecycle action for an existing Worker pool.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum WorkerPoolAction {
    /// Activate the pool.
    Activate,
    /// Pause the pool.
    Pause,
    /// Retire the pool.
    Retire,
}

fn transient(error: &CallError<TasksError>) -> bool {
    match error {
        CallError::Service(failure) => failure.status() >= 500,
        CallError::UnknownRemote(failure) => failure.status() >= 500,
        CallError::Transport(_) | CallError::Internal { .. } => true,
        CallError::Protocol(_) => false,
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;
    use std::sync::atomic::{AtomicUsize, Ordering};

    use futures::{FutureExt as _, TryStreamExt as _};
    use serde_json::{from_value, json};
    use wiremock::matchers::{header, method, path, query_param};
    use wiremock::{Mock, MockServer, ResponseTemplate};

    use super::*;
    use crate::auth::{Token, TokenSource};

    #[tokio::test]
    async fn generated_task_pages_preserve_filters_and_refresh_credentials() {
        let server = MockServer::start().await;
        let resource = Uuid::new_v4();
        for (token, cursor) in [("token-0", Some("next /&")), ("token-1", None)] {
            Mock::given(method("GET"))
                .and(path("/deployment/api/tasks/projects/project/tasks"))
                .and(header("authorization", format!("Bearer {token}")))
                .and(query_param("resourceId", resource.to_string()))
                .and(query_param("resourceRole", "target"))
                .and(query_param("status", "failed"))
                .and(query_param("pageSize", "10"))
                .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                    "tasks": [], "targetResources": {}, "nextPageToken": cursor,
                })))
                .expect(1)
                .mount(&server)
                .await;
        }
        let calls = Arc::new(AtomicUsize::new(0));
        let provider = TokenSource::new(move || {
            let token = Token::from(format!("token-{}", calls.fetch_add(1, Ordering::SeqCst)));
            async move { Ok(Some(token)) }.boxed()
        });
        let client = HttpTaskClient::new(&format!("{}/deployment", server.uri()))
            .unwrap()
            .with_token_source(provider);
        let request: ListTasksRequest = from_value(json!({
            "resourceId": resource, "resourceRole": "target", "status": "failed", "pageSize": 10,
            "pageToken": "resume /&",
        }))
        .unwrap();
        let listing = client.list_tasks_list("project".into(), request);
        assert!(server.received_requests().await.unwrap().is_empty());
        let tasks: Vec<_> = listing.stream().try_collect().await.unwrap();
        assert!(tasks.is_empty());
        let requests = server.received_requests().await.unwrap();
        assert_eq!(
            requests[0]
                .url
                .query_pairs()
                .find(|(name, _)| name == "pageToken")
                .unwrap()
                .1,
            "resume /&"
        );
        assert_eq!(
            requests[1]
                .url
                .query_pairs()
                .find(|(name, _)| name == "pageToken")
                .unwrap()
                .1,
            "next /&"
        );
    }
}