329 lines
12 KiB
Python
329 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import importlib
|
|
import threading
|
|
import uuid
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from concurrent.futures import TimeoutError as FutureTimeoutError
|
|
from datetime import UTC, datetime
|
|
|
|
import pytest
|
|
from sqlalchemy import event, select
|
|
from sqlalchemy.engine import Engine
|
|
from sqlalchemy.orm import Session
|
|
|
|
pytest_plugins = ("tests.test_managed_identities_postgres",)
|
|
|
|
from modelforge_api.domain.agent_protocol import AgentMetadata, HeartbeatRequest
|
|
from modelforge_api.domain.node_lifecycle import NodeLifecycleChange, NodeLifecyclePreviewRequest
|
|
from modelforge_api.persistence.models import ComputeNode, NodeCredential
|
|
from modelforge_api.services.node_agent import (
|
|
AgentAuthorizationError,
|
|
NodeAgentService,
|
|
NodeAuthenticationEvidence,
|
|
secret_hash,
|
|
)
|
|
from modelforge_api.services.node_decommission import (
|
|
NodeDecommissionError,
|
|
NodeDecommissionService,
|
|
)
|
|
from modelforge_api.settings import Settings
|
|
from tests.test_managed_identities_postgres import IdentityDatabase
|
|
|
|
|
|
def test_agent_request_waits_for_disable_and_rechecks_locked_node(
|
|
identity_database: IdentityDatabase,
|
|
) -> None:
|
|
"""PostgreSQL proof that a pre-body auth result cannot outlive a lifecycle disable."""
|
|
importlib.import_module("alembic.command").upgrade(identity_database.migration(), "head")
|
|
engine = identity_database.runtime
|
|
node_id = uuid.uuid4()
|
|
credential_id = uuid.uuid4()
|
|
raw = f"mfnode_{credential_id}_lifecycle-race-secret"
|
|
with Session(engine) as session:
|
|
node = ComputeNode(
|
|
id=node_id,
|
|
key=f"postgres-race-{node_id}",
|
|
hostname="postgres-race-node",
|
|
display_name="PostgreSQL lifecycle race",
|
|
enabled=True,
|
|
status="active",
|
|
liveness_state="offline",
|
|
inventory={"source": "test"},
|
|
observation_source="remote_agent",
|
|
)
|
|
session.add(node)
|
|
session.add(
|
|
NodeCredential(
|
|
id=credential_id,
|
|
compute_node_id=node_id,
|
|
secret_hash=secret_hash(raw),
|
|
scope="node.publish",
|
|
)
|
|
)
|
|
session.commit()
|
|
|
|
preview = NodeDecommissionService(session).preview_enabled_change(
|
|
node_id, NodeLifecyclePreviewRequest(target_enabled=False)
|
|
)
|
|
change = NodeLifecycleChange(
|
|
target_enabled=False,
|
|
expected_generation=preview.node_generation,
|
|
dependency_digest=preview.dependency_digest,
|
|
operator="postgres-race-test",
|
|
reason="Prove that disable serializes with agent authentication.",
|
|
confirmation=preview.persisted_identity,
|
|
)
|
|
|
|
request_session = Session(engine, expire_on_commit=False)
|
|
agent_service = NodeAgentService(request_session, Settings(_env_file=None))
|
|
(_credential, stale_node), evidence = agent_service.authenticate_with_evidence(
|
|
f"Bearer {raw}"
|
|
)
|
|
request_session.commit()
|
|
assert stale_node.enabled is True
|
|
|
|
lifecycle_has_locks = threading.Event()
|
|
allow_lifecycle_to_commit = threading.Event()
|
|
agent_lock_attempted = threading.Event()
|
|
|
|
with Session(engine) as lifecycle_session:
|
|
lifecycle_service = NodeDecommissionService(lifecycle_session)
|
|
lock_dependencies = lifecycle_service._lock_dependency_tables
|
|
|
|
def hold_lifecycle_locks() -> None:
|
|
lock_dependencies()
|
|
lifecycle_has_locks.set()
|
|
if not allow_lifecycle_to_commit.wait(timeout=20):
|
|
raise RuntimeError("timed out holding lifecycle locks in race test")
|
|
|
|
lifecycle_service._lock_dependency_tables = hold_lifecycle_locks # type: ignore[method-assign]
|
|
|
|
@event.listens_for(engine, "before_cursor_execute")
|
|
def observe_agent_lock_attempt(
|
|
_connection: object,
|
|
_cursor: object,
|
|
statement: str,
|
|
_parameters: object,
|
|
_context: object,
|
|
_executemany: bool,
|
|
) -> None:
|
|
normalized = statement.lower()
|
|
if (
|
|
threading.current_thread().name.startswith("agent-request")
|
|
and "compute_nodes" in normalized
|
|
and "for update" in normalized
|
|
):
|
|
agent_lock_attempted.set()
|
|
|
|
try:
|
|
with ThreadPoolExecutor(max_workers=1, thread_name_prefix="lifecycle") as lifecycle_pool:
|
|
lifecycle_future = lifecycle_pool.submit(
|
|
lifecycle_service.apply_enabled_change, node_id, change
|
|
)
|
|
assert lifecycle_has_locks.wait(timeout=10), "lifecycle never acquired its locks"
|
|
with ThreadPoolExecutor(max_workers=1, thread_name_prefix="agent-request") as agent_pool:
|
|
agent_future = agent_pool.submit(
|
|
_publish_authenticated_heartbeat,
|
|
agent_service,
|
|
evidence,
|
|
raw,
|
|
stale_node.key,
|
|
)
|
|
assert agent_lock_attempted.wait(timeout=10), "agent never requested its node lock"
|
|
with pytest.raises(FutureTimeoutError):
|
|
agent_future.result(timeout=0.25)
|
|
|
|
allow_lifecycle_to_commit.set()
|
|
lifecycle_future.result(timeout=30)
|
|
assert agent_future.result(timeout=30) == "rejected"
|
|
finally:
|
|
allow_lifecycle_to_commit.set()
|
|
event.remove(engine, "before_cursor_execute", observe_agent_lock_attempt)
|
|
request_session.close()
|
|
|
|
with Session(engine) as session:
|
|
persisted = session.scalar(select(ComputeNode).where(ComputeNode.id == node_id))
|
|
assert persisted is not None
|
|
assert persisted.enabled is False
|
|
assert persisted.liveness_state == "disabled"
|
|
assert persisted.generation == change.expected_generation + 1
|
|
|
|
|
|
def test_heartbeat_wins_node_lock_and_invalidates_lifecycle_preview(
|
|
identity_database: IdentityDatabase,
|
|
) -> None:
|
|
"""A heartbeat serialized before disable advances the CAS generation."""
|
|
importlib.import_module("alembic.command").upgrade(identity_database.migration(), "head")
|
|
engine = identity_database.runtime
|
|
node_id = uuid.uuid4()
|
|
credential_id = uuid.uuid4()
|
|
raw = f"mfnode_{credential_id}_heartbeat-first-race-secret"
|
|
with Session(engine) as session:
|
|
node = ComputeNode(
|
|
id=node_id,
|
|
key=f"postgres-heartbeat-race-{node_id}",
|
|
hostname="postgres-heartbeat-race-node",
|
|
display_name="PostgreSQL heartbeat race",
|
|
enabled=True,
|
|
status="active",
|
|
liveness_state="offline",
|
|
inventory={"source": "test"},
|
|
observation_source="remote_agent",
|
|
)
|
|
session.add(node)
|
|
session.add(
|
|
NodeCredential(
|
|
id=credential_id,
|
|
compute_node_id=node_id,
|
|
secret_hash=secret_hash(raw),
|
|
scope="node.publish",
|
|
)
|
|
)
|
|
session.commit()
|
|
lifecycle = NodeDecommissionService(session)
|
|
preview = lifecycle.preview_enabled_change(
|
|
node_id, NodeLifecyclePreviewRequest(target_enabled=False)
|
|
)
|
|
change = NodeLifecycleChange(
|
|
target_enabled=False,
|
|
expected_generation=preview.node_generation,
|
|
dependency_digest=preview.dependency_digest,
|
|
operator="postgres-race-test",
|
|
reason="Prove that fresh heartbeat evidence invalidates this preview.",
|
|
confirmation=preview.persisted_identity,
|
|
)
|
|
|
|
agent_session = Session(engine, expire_on_commit=False)
|
|
agent_service = NodeAgentService(agent_session, Settings(_env_file=None))
|
|
(_credential, stale_node), evidence = agent_service.authenticate_with_evidence(
|
|
f"Bearer {raw}"
|
|
)
|
|
agent_session.commit()
|
|
heartbeat_has_node_lock = threading.Event()
|
|
allow_heartbeat_to_commit = threading.Event()
|
|
lifecycle_lock_attempted = threading.Event()
|
|
original_commit = agent_session.commit
|
|
|
|
def hold_heartbeat_commit() -> None:
|
|
heartbeat_has_node_lock.set()
|
|
if not allow_heartbeat_to_commit.wait(timeout=20):
|
|
raise RuntimeError("timed out holding the heartbeat node lock")
|
|
original_commit()
|
|
|
|
agent_session.commit = hold_heartbeat_commit # type: ignore[method-assign]
|
|
|
|
@event.listens_for(engine, "before_cursor_execute")
|
|
def observe_lifecycle_lock_attempt(
|
|
_connection: object,
|
|
_cursor: object,
|
|
statement: str,
|
|
_parameters: object,
|
|
_context: object,
|
|
_executemany: bool,
|
|
) -> None:
|
|
normalized = statement.lower()
|
|
if (
|
|
threading.current_thread().name.startswith("lifecycle-after-heartbeat")
|
|
and "compute_nodes" in normalized
|
|
and "for update" in normalized
|
|
):
|
|
lifecycle_lock_attempted.set()
|
|
|
|
request = HeartbeatRequest(
|
|
identity_key=stale_node.key,
|
|
metadata=AgentMetadata(
|
|
agent_version="1.3.0-rc.1",
|
|
protocol_version=1,
|
|
supported_capabilities=["hardware.inventory"],
|
|
started_at=datetime.now(UTC),
|
|
),
|
|
observed_at=datetime.now(UTC),
|
|
)
|
|
|
|
try:
|
|
with ThreadPoolExecutor(max_workers=1, thread_name_prefix="heartbeat") as agent_pool:
|
|
heartbeat_future = agent_pool.submit(
|
|
_publish_and_hold_authenticated_heartbeat,
|
|
agent_service,
|
|
evidence,
|
|
raw,
|
|
request,
|
|
)
|
|
assert heartbeat_has_node_lock.wait(timeout=10), "heartbeat never acquired node lock"
|
|
with ThreadPoolExecutor(
|
|
max_workers=1, thread_name_prefix="lifecycle-after-heartbeat"
|
|
) as lifecycle_pool:
|
|
lifecycle_future = lifecycle_pool.submit(
|
|
_apply_stale_lifecycle_change, engine, node_id, change
|
|
)
|
|
assert lifecycle_lock_attempted.wait(timeout=10), "lifecycle never requested node lock"
|
|
with pytest.raises(FutureTimeoutError):
|
|
lifecycle_future.result(timeout=0.25)
|
|
allow_heartbeat_to_commit.set()
|
|
heartbeat_future.result(timeout=30)
|
|
assert lifecycle_future.result(timeout=30) == "node_generation_conflict"
|
|
finally:
|
|
allow_heartbeat_to_commit.set()
|
|
event.remove(engine, "before_cursor_execute", observe_lifecycle_lock_attempt)
|
|
agent_session.close()
|
|
|
|
with Session(engine) as session:
|
|
persisted = session.scalar(select(ComputeNode).where(ComputeNode.id == node_id))
|
|
assert persisted is not None
|
|
assert persisted.enabled is True
|
|
assert persisted.liveness_state == "online"
|
|
assert persisted.last_heartbeat_at is not None
|
|
assert persisted.generation == change.expected_generation + 1
|
|
|
|
|
|
def _publish_and_hold_authenticated_heartbeat(
|
|
service: NodeAgentService,
|
|
evidence: NodeAuthenticationEvidence,
|
|
raw_credential: str,
|
|
request: HeartbeatRequest,
|
|
) -> None:
|
|
_credential, node = service.reuse_authentication(evidence, f"Bearer {raw_credential}")
|
|
service.heartbeat(node, request)
|
|
|
|
|
|
def _apply_stale_lifecycle_change(
|
|
engine: Engine, node_id: uuid.UUID, change: NodeLifecycleChange
|
|
) -> str:
|
|
with Session(engine) as session:
|
|
try:
|
|
NodeDecommissionService(session).apply_enabled_change(node_id, change)
|
|
except NodeDecommissionError as exc:
|
|
if exc.code != "node_generation_conflict":
|
|
raise
|
|
return exc.code
|
|
return "applied"
|
|
|
|
|
|
def _publish_authenticated_heartbeat(
|
|
service: NodeAgentService,
|
|
evidence: NodeAuthenticationEvidence,
|
|
raw_credential: str,
|
|
identity_key: str,
|
|
) -> str:
|
|
try:
|
|
_credential, node = service.reuse_authentication(
|
|
evidence, f"Bearer {raw_credential}"
|
|
)
|
|
service.heartbeat(
|
|
node,
|
|
HeartbeatRequest(
|
|
identity_key=identity_key,
|
|
metadata=AgentMetadata(
|
|
agent_version="1.3.0-rc.1",
|
|
protocol_version=1,
|
|
supported_capabilities=["hardware.inventory"],
|
|
started_at=datetime.now(UTC),
|
|
),
|
|
observed_at=datetime.now(UTC),
|
|
),
|
|
)
|
|
except AgentAuthorizationError:
|
|
return "rejected"
|
|
return "published"
|