Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
45 changes: 45 additions & 0 deletions python/packages/jumpstarter/jumpstarter/exporter/exporter_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -442,6 +442,51 @@ def _make_exporter_for_report_status():
return exporter


class TestBeforeLeaseHookRaceGuard:
async def test_new_lease_after_before_hook_race_recovery(self):
"""After recovering from the beforeLease hook race condition
(lease expired during hook), a new lease must be accepted and
processed normally."""
from jumpstarter.config.exporter import HookConfigV1Alpha1, HookInstanceConfigV1Alpha1
from jumpstarter.exporter.hooks import HookExecutor

hook_config = HookConfigV1Alpha1(
before_lease=HookInstanceConfigV1Alpha1(script="echo setup", timeout=10),
)
hook_executor = HookExecutor(config=hook_config)

lease_ctx_1 = make_lease_context(lease_name="expired-lease")
lease_ctx_1.lease_ended.set()

statuses = []

async def track_status(status, message=""):
statuses.append(status)

exporter = make_exporter(lease_ctx_1, hook_executor)
exporter._report_status = AsyncMock(side_effect=track_status)

await hook_executor.run_before_lease_hook(
lease_ctx_1, exporter._report_status, exporter.stop, exporter._request_lease_release
)

assert lease_ctx_1.before_lease_hook.is_set()
await exporter._cleanup_after_lease(lease_ctx_1)
assert ExporterStatus.AVAILABLE in statuses

lease_ctx_2 = make_lease_context(lease_name="new-lease")
exporter._lease_context = lease_ctx_2

statuses.clear()
await hook_executor.run_before_lease_hook(
lease_ctx_2, exporter._report_status, exporter.stop, exporter._request_lease_release
)

assert ExporterStatus.LEASE_READY in statuses, (
f"New lease must reach LEASE_READY when lease is still active. Statuses: {statuses}"
)


class TestReportStatusGrpcErrorHandling:
async def test_unimplemented_grpc_error_logs_warning(self, caplog):
"""When ReportStatus returns UNIMPLEMENTED, a warning is logged
Expand Down
7 changes: 7 additions & 0 deletions python/packages/jumpstarter/jumpstarter/exporter/hooks.py
Original file line number Diff line number Diff line change
Expand Up @@ -615,6 +615,13 @@ async def run_before_lease_hook(
LogSource.BEFORE_LEASE_HOOK,
)

if lease_scope.lease_ended.is_set():
logger.info(
"Lease %s ended during beforeLease hook, skipping LEASE_READY transition",
lease_scope.lease_name,
)
return

if warning:
msg = f"{HOOK_WARNING_PREFIX}beforeLease hook warning: {warning}"
else:
Expand Down
96 changes: 96 additions & 0 deletions python/packages/jumpstarter/jumpstarter/exporter/hooks_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -1178,3 +1178,99 @@ async def mock_report_status(status, msg):
assert msg.startswith(HOOK_WARNING_PREFIX), (
f"Expected AVAILABLE message to start with '{HOOK_WARNING_PREFIX}', got: '{msg}'"
)


class TestBeforeLeaseHookLeaseEndedGuard:
"""Tests for the race condition where beforeLease hook completes after
the lease has already expired. When lease_ended is set, the hook must
NOT set status to LEASE_READY, preventing the exporter from being
stuck in LEASE_READY permanently."""

async def test_run_before_lease_hook_skips_lease_ready_when_lease_ended(self, lease_scope) -> None:
"""When the lease has already ended by the time the beforeLease hook
completes, status must NOT be set to LEASE_READY."""
hook_config = HookConfigV1Alpha1(
before_lease=HookInstanceConfigV1Alpha1(script="echo setup", timeout=10),
)
executor = HookExecutor(config=hook_config)

lease_scope.lease_ended.set()

status_calls = []

async def mock_report_status(status, msg):
status_calls.append((status, msg))

mock_shutdown = MagicMock()

await executor.run_before_lease_hook(
lease_scope,
mock_report_status,
mock_shutdown,
)

lease_ready_calls = [s for s, _ in status_calls if s == ExporterStatus.LEASE_READY]
assert len(lease_ready_calls) == 0, (
f"LEASE_READY must NOT be set when lease has already ended, got: {status_calls}"
)

hook_started_calls = [s for s, _ in status_calls if s == ExporterStatus.BEFORE_LEASE_HOOK]
assert len(hook_started_calls) == 1, (
f"BEFORE_LEASE_HOOK must be reported (hook must run) even when lease has ended, got: {status_calls}"
)

async def test_run_before_lease_hook_sets_event_even_when_lease_ended(self, lease_scope) -> None:
"""The before_lease_hook event must always be set (via the finally block)
even when the lease has ended, to unblock downstream waiters."""
hook_config = HookConfigV1Alpha1(
before_lease=HookInstanceConfigV1Alpha1(script="echo setup", timeout=10),
)
executor = HookExecutor(config=hook_config)

lease_scope.lease_ended.set()

mock_report_status = AsyncMock()
mock_shutdown = MagicMock()

await executor.run_before_lease_hook(
lease_scope,
mock_report_status,
mock_shutdown,
)

assert lease_scope.before_lease_hook.is_set(), (
"before_lease_hook event must be set even when lease has ended"
)

async def test_run_before_lease_hook_warn_skips_lease_ready_when_lease_ended(self, lease_scope) -> None:
"""When hook fails with on_failure=warn and the lease has already ended,
LEASE_READY must still be skipped."""
hook_config = HookConfigV1Alpha1(
before_lease=HookInstanceConfigV1Alpha1(script="exit 1", timeout=10, on_failure="warn"),
)
executor = HookExecutor(config=hook_config)

lease_scope.lease_ended.set()

status_calls = []

async def mock_report_status(status, msg):
status_calls.append((status, msg))

mock_shutdown = MagicMock()

await executor.run_before_lease_hook(
lease_scope,
mock_report_status,
mock_shutdown,
)

lease_ready_calls = [s for s, _ in status_calls if s == ExporterStatus.LEASE_READY]
assert len(lease_ready_calls) == 0, (
f"LEASE_READY must NOT be set when lease has ended (even with warn), got: {status_calls}"
)

hook_started_calls = [s for s, _ in status_calls if s == ExporterStatus.BEFORE_LEASE_HOOK]
assert len(hook_started_calls) == 1, (
f"BEFORE_LEASE_HOOK must be reported (hook must run) even when lease has ended, got: {status_calls}"
)
Loading