diff --git a/src/ci_orchestrator/commands/template.py b/src/ci_orchestrator/commands/template.py index 2fe50ea..a7fb233 100644 --- a/src/ci_orchestrator/commands/template.py +++ b/src/ci_orchestrator/commands/template.py @@ -292,42 +292,61 @@ def _wait_tcp(host: str, port: int, timeout: float) -> bool: return False -def _wait_tcp_down(host: str, port: int, timeout: float, interval: float = 5.0) -> bool: - """Return True once TCP *port* on *host* stops accepting connections (host rebooting). - - Used after issuing a reboot: a pending Windows Update / CBS install can delay - the actual restart for minutes while the WinRM session is still alive, so we - must observe the port go *down* before waiting for it to come back — otherwise - we'd treat the still-up pre-reboot session as "ready". - """ - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - try: - with socket.create_connection((host, port), timeout=5.0): - time.sleep(interval) - except OSError: +def _tcp_open(host: str, port: int, timeout: float = 3.0) -> bool: + """True if *port* on *host* currently accepts a TCP connection.""" + try: + with socket.create_connection((host, port), timeout=timeout): return True - return False + except OSError: + return False + + +def _guest_boot_ticks(transport: WinRmTransport) -> int | None: + """Read the guest's last boot time as FileTime ticks; None if unreachable. + + Used as a reboot signal that is immune to how fast the box restarts — + comparing boot time before/after a reboot is reliable even when the WinRM + port never appears "down" between polls. + """ + try: + result = transport.run( + "(Get-CimInstance Win32_OperatingSystem).LastBootUpTime.ToFileTimeUtc()", + check=False, + ) + except TransportError: + return None + try: + return int(result.stdout.strip()) + except (ValueError, AttributeError): + return None def _reboot_and_wait( transport: WinRmTransport, host: str, port: int, winrm_timeout: int ) -> None: - """Reboot the guest and wait for it to actually go down then come back ready.""" + """Reboot the guest and wait until it comes back with a *newer* boot time. + + A TCP "down" probe can miss a fast reboot entirely, so instead we record the + boot time before the reboot and poll until WinRM answers with a later one. + """ + boot_before = _guest_boot_ticks(transport) with contextlib.suppress(TransportError): transport.run("Restart-Computer -Force", check=False) transport.close() - click.echo("[prepare-win] waiting for VM to go down ...") - if _wait_tcp_down(host, port, timeout=300.0): - click.echo("[prepare-win] VM down — waiting for it to come back ...") - else: - click.echo( - "[prepare-win] [WARN] WinRM still reachable after reboot request; " - "proceeding anyway." - ) - if not _wait_tcp(host, port, float(winrm_timeout)): - raise click.ClickException(f"TCP/{port} on {host} not reachable after reboot.") - _wait_winrm(transport, winrm_timeout) + click.echo("[prepare-win] waiting for guest to reboot (boot-time change) ...") + deadline = time.monotonic() + float(winrm_timeout) + 300.0 + while time.monotonic() < deadline: + time.sleep(5) + if not _tcp_open(host, port): + continue # still down / not back yet + ticks = _guest_boot_ticks(transport) + if ticks is not None and (boot_before is None or ticks > boot_before): + click.echo("[prepare-win] guest rebooted — WinRM back.") + _wait_winrm(transport, winrm_timeout) + return + raise click.ClickException( + f"Guest did not reboot / WinRM not back on {host}:{port} within timeout." + ) def _parse_exit_marker(stdout: str) -> int: diff --git a/tests/python/test_commands_template.py b/tests/python/test_commands_template.py index e5d111f..6ee8d96 100644 --- a/tests/python/test_commands_template.py +++ b/tests/python/test_commands_template.py @@ -363,44 +363,28 @@ def test_wait_tcp_timeout(monkeypatch: pytest.MonkeyPatch) -> None: assert tmpl_mod._wait_tcp("127.0.0.1", 9999, 0.0) is False -def test_wait_tcp_down_detects_down(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr( - tmpl_mod.socket, "create_connection", - lambda *a, **kw: (_ for _ in ()).throw(OSError("down")), - ) - assert tmpl_mod._wait_tcp_down("127.0.0.1", 5986, 5.0) is True - - -def test_wait_tcp_down_still_up(monkeypatch: pytest.MonkeyPatch) -> None: - class _FakeConn: - def __enter__(self) -> _FakeConn: - return self - - def __exit__(self, *_: object) -> None: - pass - - monkeypatch.setattr(tmpl_mod.socket, "create_connection", lambda *a, **kw: _FakeConn()) - monkeypatch.setattr(tmpl_mod.time, "sleep", lambda _: None) - assert tmpl_mod._wait_tcp_down("127.0.0.1", 5986, 0.0) is False - - -def test_reboot_and_wait(monkeypatch: pytest.MonkeyPatch) -> None: +def test_reboot_and_wait_boot_change(monkeypatch: pytest.MonkeyPatch) -> None: calls: list[str] = [] transport = MagicMock() - transport.run.side_effect = lambda *_a, **_k: calls.append("reboot") - transport.close.side_effect = lambda: calls.append("close") - monkeypatch.setattr(tmpl_mod, "_wait_tcp_down", lambda *a, **k: True) - monkeypatch.setattr(tmpl_mod, "_wait_tcp", lambda *a, **k: True) + monkeypatch.setattr(tmpl_mod, "_tcp_open", lambda *a, **k: True) + # boot ticks: before reboot, then a newer value once it's back. + ticks = iter([100, 200]) + monkeypatch.setattr(tmpl_mod, "_guest_boot_ticks", lambda _t: next(ticks)) monkeypatch.setattr(tmpl_mod, "_wait_winrm", lambda *a, **k: calls.append("winrm")) + monkeypatch.setattr(tmpl_mod.time, "sleep", lambda _s: None) tmpl_mod._reboot_and_wait(transport, "10.0.0.5", 5986, 30) - assert calls == ["reboot", "close", "winrm"] + assert "winrm" in calls + transport.close.assert_called_once() -def test_reboot_and_wait_unreachable(monkeypatch: pytest.MonkeyPatch) -> None: +def test_reboot_and_wait_timeout(monkeypatch: pytest.MonkeyPatch) -> None: transport = MagicMock() - monkeypatch.setattr(tmpl_mod, "_wait_tcp_down", lambda *a, **k: True) - monkeypatch.setattr(tmpl_mod, "_wait_tcp", lambda *a, **k: False) - with pytest.raises(click.ClickException, match="not reachable after reboot"): + monkeypatch.setattr(tmpl_mod, "_tcp_open", lambda *a, **k: False) # never comes back + monkeypatch.setattr(tmpl_mod, "_guest_boot_ticks", lambda _t: 100) + monkeypatch.setattr(tmpl_mod.time, "sleep", lambda _s: None) + clock = iter([0.0, 1.0, 1e9]) # deadline calc, one loop check, then past deadline + monkeypatch.setattr(tmpl_mod.time, "monotonic", lambda: next(clock)) + with pytest.raises(click.ClickException, match="did not reboot"): tmpl_mod._reboot_and_wait(transport, "10.0.0.5", 5986, 1) @@ -530,7 +514,7 @@ def _patch_prepare(monkeypatch: pytest.MonkeyPatch, backend: MagicMock, transpor monkeypatch.setattr(tmpl_mod, "WorkstationVmrunBackend", lambda **_: backend) monkeypatch.setattr(tmpl_mod, "WinRmTransport", lambda **_: transport) monkeypatch.setattr(tmpl_mod, "_wait_tcp", lambda *_: True) - monkeypatch.setattr(tmpl_mod, "_wait_tcp_down", lambda *_a, **_k: True) + monkeypatch.setattr(tmpl_mod, "_reboot_and_wait", lambda *_a, **_k: None) monkeypatch.setattr(tmpl_mod, "_wait_winrm", lambda *_: None) monkeypatch.setattr(tmpl_mod, "_check_guest_idle", lambda *_a, **_k: None) monkeypatch.setattr(tmpl_mod.time, "sleep", lambda _: None) @@ -694,7 +678,8 @@ def test_prepare_win_snapshot_exists_recreate( monkeypatch.setattr(tmpl_mod, "WorkstationVmrunBackend", lambda **_: backend) monkeypatch.setattr(tmpl_mod, "WinRmTransport", lambda **_: transport) monkeypatch.setattr(tmpl_mod, "_wait_tcp", lambda *_: True) - monkeypatch.setattr(tmpl_mod, "_wait_tcp_down", lambda *_a, **_k: True) + monkeypatch.setattr(tmpl_mod, "_reboot_and_wait", lambda *_a, **_k: None) + monkeypatch.setattr(tmpl_mod, "_check_guest_idle", lambda *_a, **_k: None) monkeypatch.setattr(tmpl_mod, "_wait_winrm", lambda *_: None) monkeypatch.setattr(tmpl_mod.time, "sleep", lambda _: None) def _sp_run(cmd: list[str], **_: object) -> None: