Skip to content

Fix wake recovery after snapshot restore timeout #87

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
fcasal wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Fix wake recovery after snapshot restore timeout #87

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c1757c8
Fix wake recovery after snapshot restore timeout
fcasal Sep 8, 2026
65c9436
Preserve Tailscale SSH configuration when resuming wake
fcasal Sep 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -178,6 +178,9 @@ dropkit destroy my-droplet
3. Snapshots are tagged with `owner:<username>` and `size:<size-slug>` for tracking
4. After waking, you're prompted to delete the snapshot (default: yes)

Wake waits up to 15 minutes by default. Use `dropkit wake my-droplet --timeout 1800`
to wait up to 30 minutes.

**Note:** Snapshots are billed at $0.06/GB/month, which is typically much cheaper than keeping a droplet running.

### Cloud-Init Customization
Expand Down
15 changes: 10 additions & 5 deletions dropkit/api.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,10 @@ def __init__(self, message: str, status_code: int | None = None):
super().__init__(message)


class DropletWaitTimeoutError(DigitalOceanAPIError):
"""The local wait expired; droplet creation may still be in progress."""


PROTECTED_TAGS = {"owner", "firewall"}


Expand Down Expand Up @@ -405,12 +409,13 @@ def wait_for_droplet_active(

Raises:
ValueError: If droplet_id is not positive
DigitalOceanAPIError: If timeout is reached or droplet enters error state
DropletWaitTimeoutError: If the local wait expires
DigitalOceanAPIError: If the droplet enters an error state or an API request fails
"""
import time

self._validate_positive_int(droplet_id, "droplet_id")
start_time = time.time()
start_time = time.monotonic()

while True:
droplet = self.get_droplet(droplet_id)
Expand All @@ -423,9 +428,9 @@ def wait_for_droplet_active(
f"Droplet entered error state: {droplet.get('name', droplet_id)}"
)

elapsed = time.time() - start_time
if elapsed > timeout:
raise DigitalOceanAPIError(
elapsed = time.monotonic() - start_time
if elapsed >= timeout:
raise DropletWaitTimeoutError(
f"Timeout waiting for droplet to become active (waited {elapsed:.0f}s)"
)

Expand Down
138 changes: 97 additions & 41 deletions dropkit/main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@
from rich.prompt import Confirm, Prompt
from rich.table import Table

from dropkit.api import DigitalOceanAPI, DigitalOceanAPIError
from dropkit.api import DigitalOceanAPI, DigitalOceanAPIError, DropletWaitTimeoutError
from dropkit.cloudinit import render_cloud_init
from dropkit.config import DEFAULT_IMAGE, DEFAULT_REGION, DEFAULT_SIZE, Config, DropkitConfig
from dropkit.lock import requires_lock
Expand Down Expand Up @@ -4036,17 +4036,21 @@ def wake(
droplet_name: str = typer.Argument(
..., autocompletion=complete_snapshot_name, help="Name of the hibernated droplet to restore"
),
timeout: int = typer.Option(
900, "--timeout", min=1, help="Seconds to wait for restoration (default: 15 minutes)"
),
no_tailscale: bool = typer.Option(False, "--no-tailscale", help="Skip Tailscale VPN re-setup"),
):
"""
Wake a hibernated droplet (restore from snapshot).

This will create a new droplet from the hibernated snapshot.
Run again after a timeout to resume setup of the same restored droplet.
After successful restoration, you'll be prompted to delete the snapshot.

If the original droplet had Tailscale lockdown enabled, this command will
re-setup Tailscale after the droplet becomes active. Use --no-tailscale to
skip this and keep public SSH access.
skip this. When resuming, existing Tailscale SSH access is preserved.

Use 'dropkit destroy <name>' to delete a hibernated snapshot without restoring.
"""
Expand All @@ -4064,10 +4068,6 @@ def wake(

# Check if a droplet with this name already exists
existing_droplet, _ = find_user_droplet(api, droplet_name)
if existing_droplet:
console.print(f"[red]Error: A droplet named '{droplet_name}' already exists.[/red]")
console.print("[dim]Destroy or rename the existing droplet first.[/dim]")
raise typer.Exit(1)

# Find the hibernated snapshot
snapshot_name = get_snapshot_name(droplet_name)
Expand All @@ -4079,6 +4079,13 @@ def wake(
if not snapshot:
console.print(f"[red]Error: No hibernated snapshot found for '{droplet_name}'[/red]")
console.print(f"[dim]Expected snapshot name: {snapshot_name}[/dim]")
if existing_droplet:
console.print(
f"[dim]The droplet already exists. Check it with: dropkit info {droplet_name}[/dim]"
)
console.print(
f"[dim]To configure SSH, run: dropkit config-ssh {droplet_name}[/dim]"
)
raise typer.Exit(1)

snapshot_id_str = snapshot.get("id")
Expand All @@ -4087,6 +4094,23 @@ def wake(
raise typer.Exit(1)
snapshot_id = int(snapshot_id_str) # API returns string, convert to int

if existing_droplet:
# A matching name alone is insufficient: only resume a restore of this snapshot.
source_image_id = existing_droplet.get("image", {}).get("id")
if str(source_image_id) != str(snapshot_id):
console.print(
f"[red]Error: A droplet named '{droplet_name}' already exists, "
"but was not restored from this snapshot.[/red]"
)
console.print(f"[dim]Inspect it with: dropkit info {droplet_name}[/dim]")
raise typer.Exit(1)
status = existing_droplet.get("status")
if status not in {"new", "active"}:
console.print(f"[red]Cannot resume wake: droplet status is '{status}'.[/red]")
if status == "off":
console.print(f"[dim]Power it on first: dropkit on {droplet_name}[/dim]")
raise typer.Exit(1)

# Get snapshot details
size_gb = snapshot.get("size_gigabytes", 0)
regions = snapshot.get("regions", [])
Expand Down Expand Up @@ -4120,33 +4144,60 @@ def wake(
)
console.print()

# Create droplet from snapshot
console.print(f"[dim]Creating droplet '{droplet_name}' from snapshot...[/dim]")

# Build tags for new droplet
tags_list = build_droplet_tags(username, list(config.defaults.extra_tags))

droplet = api.create_droplet_from_snapshot(
name=droplet_name,
region=original_region,
size=original_size,
snapshot_id=snapshot_id,
tags=tags_list,
ssh_keys=config.cloudinit.ssh_key_ids,
)
if existing_droplet:
droplet = existing_droplet
console.print(f"[dim]Resuming wake for existing droplet '{droplet_name}'...[/dim]")
else:
# Create droplet from snapshot
console.print(f"[dim]Creating droplet '{droplet_name}' from snapshot...[/dim]")

# Build tags for new droplet
tags_list = build_droplet_tags(username, list(config.defaults.extra_tags))

droplet = api.create_droplet_from_snapshot(
name=droplet_name,
region=original_region,
size=original_size,
snapshot_id=snapshot_id,
tags=tags_list,
ssh_keys=config.cloudinit.ssh_key_ids,
)

droplet_id = droplet.get("id")
if not droplet_id:
console.print("[red]Error: Failed to get droplet ID from API response[/red]")
raise typer.Exit(1)

console.print(f"[green]✓[/green] Droplet created (ID: [cyan]{droplet_id}[/cyan])")
console.print(f"[green]✓[/green] Droplet ID: [cyan]{droplet_id}[/cyan]")

# Wait for droplet to become active
console.print("[dim]Waiting for droplet to become active...[/dim]")

with console.status("[cyan]Waiting...[/cyan]"):
active_droplet = api.wait_for_droplet_active(droplet_id)
try:
with console.status("[cyan]Waiting...[/cyan]"):
active_droplet = api.wait_for_droplet_active(droplet_id, timeout=timeout)
except DropletWaitTimeoutError:
console.print(
f"[yellow]Stopped waiting after {timeout}s. DigitalOcean may still be "
f"restoring droplet '{droplet_name}' (ID: {droplet_id}).[/yellow]"
)
console.print("[dim]The droplet and snapshot have been kept.[/dim]")
console.print(f"[dim]Check status: dropkit info {droplet_name}[/dim]")
console.print(
f"[dim]Resume setup: dropkit wake {droplet_name} --timeout {timeout}"
f"{' --no-tailscale' if no_tailscale else ''}[/dim]"
)
raise typer.Exit(1)
except DigitalOceanAPIError:
console.print(
f"[yellow]Could not check readiness of droplet {droplet_id}. "
"The droplet and snapshot have been kept.[/yellow]"
)
console.print(
f"[dim]Check status with dropkit info {droplet_name}, "
"then rerun wake to resume setup.[/dim]"
)
raise

# Get IP address
networks = active_droplet.get("networks", {})
Expand All @@ -4162,34 +4213,42 @@ def wake(
console.print(f"[green]✓[/green] Droplet is active (IP: [cyan]{ip_address}[/cyan])")
else:
console.print("[green]✓[/green] Droplet is active")
console.print("[yellow]⚠[/yellow] Could not determine IP address")
console.print("[red]Could not determine public IP address; snapshot kept.[/red]")
console.print(f"[dim]Retry setup: dropkit wake {droplet_name}[/dim]")
raise typer.Exit(1)

# Add SSH config entry
if ip_address and config.ssh.auto_update:
try:
console.print("[dim]Configuring SSH...[/dim]")
ssh_hostname = get_ssh_hostname(droplet_name)
add_ssh_host(
config_path=config.ssh.config_path,
host_name=ssh_hostname,
hostname=ip_address,
user=username,
identity_file=config.ssh.identity_file,
)
console.print("[green]✓[/green] SSH config updated")
if existing_droplet and is_droplet_tailscale_locked(config, droplet_name):
# A previous wake may have blocked public SSH before being interrupted.
console.print("[green]✓[/green] Preserved Tailscale SSH config")
else:
add_ssh_host(
config_path=config.ssh.config_path,
host_name=ssh_hostname,
hostname=ip_address,
user=username,
identity_file=config.ssh.identity_file,
)
console.print("[green]✓[/green] SSH config updated")
except Exception as e:
console.print(f"[yellow]⚠[/yellow] Could not update SSH config: {e}")
console.print(f"[red]Could not update SSH config: {e}[/red]")
console.print(
f"[dim]Snapshot kept. Fix the SSH configuration error, "
f"then retry: dropkit wake {droplet_name}[/dim]"
)
raise typer.Exit(1)

# Handle Tailscale re-setup if the original droplet had Tailscale lockdown
if was_tailscale_locked and ip_address:
ssh_hostname = get_ssh_hostname(droplet_name)
if no_tailscale:
console.print()
console.print("[yellow]⚠[/yellow] Original droplet had Tailscale lockdown enabled.")
console.print(
"[dim]Skipping Tailscale setup (--no-tailscale). "
"Public SSH access available.[/dim]"
)
console.print("[dim]Skipping Tailscale setup (--no-tailscale).[/dim]")
console.print(
f"[dim]Enable Tailscale later with: "
f"[cyan]dropkit enable-tailscale {droplet_name}[/cyan][/dim]"
Expand All @@ -4207,10 +4266,7 @@ def wake(
tailscale_ip = setup_tailscale(ssh_hostname, username, config)

if not tailscale_ip:
console.print(
"[yellow]⚠[/yellow] Tailscale setup incomplete. "
"Public SSH access remains available."
)
console.print("[yellow]⚠[/yellow] Tailscale setup incomplete.")
console.print(
f"[dim]Complete setup later with: "
f"[cyan]dropkit enable-tailscale {droplet_name}[/cyan][/dim]"
Expand Down
Loading