diff --git a/backend/Cargo.lock b/backend/Cargo.lock index 59eb6a4c6f..fc89f30f19 100644 --- a/backend/Cargo.lock +++ b/backend/Cargo.lock @@ -15551,6 +15551,7 @@ dependencies = [ "strum", "tokio", "tokio-postgres", + "tokio-util", "tracing", "url", "uuid", diff --git a/backend/Cargo.toml b/backend/Cargo.toml index 1a564f3644..5bf65d2e61 100644 --- a/backend/Cargo.toml +++ b/backend/Cargo.toml @@ -517,7 +517,7 @@ serde_urlencoded = "^0" astral-tokio-tar = "^0.6.2" tempfile = "^3" x509-parser = "^0.16" -tokio-util = { version = "=0.7.17", features = ["io"] } +tokio-util = { version = "=0.7.17", features = ["io", "rt"] } json-pointer = "^0" itertools = "^0.14.0" regex = "^1" diff --git a/backend/windmill-api-workspaces/Cargo.toml b/backend/windmill-api-workspaces/Cargo.toml index 7ca1f692ad..9f05a26f28 100644 --- a/backend/windmill-api-workspaces/Cargo.toml +++ b/backend/windmill-api-workspaces/Cargo.toml @@ -50,6 +50,7 @@ sha2.workspace = true sqlx.workspace = true tokio.workspace = true tokio-postgres.workspace = true +tokio-util.workspace = true tracing.workspace = true url.workspace = true uuid.workspace = true diff --git a/backend/windmill-api-workspaces/src/workspaces.rs b/backend/windmill-api-workspaces/src/workspaces.rs index 121c5989d4..b0dc1534ba 100644 --- a/backend/windmill-api-workspaces/src/workspaces.rs +++ b/backend/windmill-api-workspaces/src/workspaces.rs @@ -9079,7 +9079,7 @@ async fn create_workspace_fork( ))); }; - tokio::spawn(async move { + BACKGROUND_FORKS.spawn(async move { let heartbeat = { let (db, fork_id) = (db.clone(), fork_id.clone()); tokio::spawn(async move { @@ -9139,6 +9139,12 @@ async fn create_workspace_fork( const FORK_HEARTBEAT_SECS: u64 = 10; const FORK_HEARTBEAT_STALE_SECS: u64 = 60; +/// Forks being created in the background. The server waits for them on shutdown: a copy cut short +/// never drops the data table databases it already made, and a retry under the same id is refused +/// while they exist. +pub static BACKGROUND_FORKS: std::sync::LazyLock = + std::sync::LazyLock::new(tokio_util::task::TaskTracker::new); + #[derive(Deserialize)] struct CreateWorkspaceForkQuery { /// Return once the request is validated and create the fork in the background; its progress diff --git a/backend/windmill-api/src/lib.rs b/backend/windmill-api/src/lib.rs index 156c6ca898..556163ce1c 100644 --- a/backend/windmill-api/src/lib.rs +++ b/backend/windmill-api/src/lib.rs @@ -1279,6 +1279,15 @@ pub async fn run_server( }); server.await?; + let forks = &windmill_api_workspaces::workspaces::BACKGROUND_FORKS; + forks.close(); + if !forks.is_empty() { + tracing::info!( + "server off. waiting for {} fork(s) being created", + forks.len() + ); + forks.wait().await; + } #[cfg(feature = "agent_worker_server")] for (i, bg_processor) in agent_workers_bg_processor.into_iter().enumerate() { tracing::info!("server off. shutting down agent worker bg processor {i}"); diff --git a/cli/src/commands/workspace/fork.ts b/cli/src/commands/workspace/fork.ts index 0aca0984ff..1a6486ad5d 100644 --- a/cli/src/commands/workspace/fork.ts +++ b/cli/src/commands/workspace/fork.ts @@ -22,6 +22,34 @@ import { readConfigFile, } from "../../core/conf.ts"; +/** + * Wait for a fork requested with `background: true`. A server without background forks ignores + * the flag and has already created the fork when the request returns; it has no status route. + */ +async function waitForForkCreation(parentWorkspace: string, forkId: string) { + let failedPolls = 0; + while (true) { + await new Promise((resolve) => setTimeout(resolve, 1500)); + let result: Awaited>; + try { + result = await wmill.getForkCreationStatus({ + workspace: parentWorkspace, + forkWorkspaceId: forkId, + }); + failedPolls = 0; + } catch (e) { + if ((e as { status?: number }).status === 404) return; + // A poll lost to a proxy says nothing about the fork, which is still being created. + if (++failedPolls >= 5) throw e; + continue; + } + if (result.status === "completed") return; + if (result.status === "failed") { + throw new Error(result.error ?? "Unknown error"); + } + } +} + async function createWorkspaceFork( opts: GlobalOptions & { createWorkspaceName: string | undefined; @@ -376,8 +404,9 @@ async function createWorkspaceFork( // --- Create the fork workspace --- try { - const result = await wmill.createWorkspaceFork({ + await wmill.createWorkspaceFork({ workspace: workspace.workspaceId, + background: true, requestBody: { id: trueWorkspaceId, name: opts.createWorkspaceName ?? workspaceName ?? trueWorkspaceId, @@ -385,8 +414,9 @@ async function createWorkspaceFork( forked_datatables: forkedDatatables, }, }); + await waitForForkCreation(workspace.workspaceId, trueWorkspaceId); - log.info(colors.green(`✅ ${result}`)); + log.info(colors.green(`✅ Created forked workspace ${trueWorkspaceId}`)); } catch (error) { log.error( colors.red(`Failed to create forked workspace: ${(error as Error).message}`),