implement heartbeat command
This commit is contained in:
ars
@@ -0,0 +1,11 @@
|
||||
from nexus_sync.server.app import create_app
|
||||
from nexus_sync.server.config import DEFAULT_COMMAND_POLL_SECONDS, DEFAULT_IDLE_POLL_SECONDS
|
||||
from nexus_sync.server.store import InMemoryStore, Store
|
||||
|
||||
__all__ = [
|
||||
"DEFAULT_COMMAND_POLL_SECONDS",
|
||||
"DEFAULT_IDLE_POLL_SECONDS",
|
||||
"InMemoryStore",
|
||||
"Store",
|
||||
"create_app",
|
||||
]
|
||||
|
||||
@@ -1,5 +1,10 @@
|
||||
import uvicorn
|
||||
|
||||
from nexus_sync.server.app import app
|
||||
|
||||
|
||||
def main() -> None:
|
||||
print("nexus-sync server")
|
||||
uvicorn.run(app, host="0.0.0.0", port=8000)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -0,0 +1,84 @@
|
||||
from datetime import UTC, datetime
|
||||
|
||||
from fastapi import Depends, FastAPI, Header, HTTPException, Request, status
|
||||
from fastapi.exceptions import RequestValidationError
|
||||
from fastapi.responses import JSONResponse
|
||||
|
||||
from nexus_sync.common import HeartbeatRequest, HeartbeatResponse
|
||||
from nexus_sync.server.config import (
|
||||
DEFAULT_COMMAND_POLL_SECONDS,
|
||||
DEFAULT_IDLE_POLL_SECONDS,
|
||||
load_client_tokens,
|
||||
)
|
||||
from nexus_sync.server.store import InMemoryStore, Store
|
||||
|
||||
|
||||
def create_app(
|
||||
*,
|
||||
store: Store | None = None,
|
||||
client_tokens: dict[str, str] | None = None,
|
||||
) -> FastAPI:
|
||||
app = FastAPI(title="nexus-sync", version="0.1.0")
|
||||
app.state.store = store or InMemoryStore()
|
||||
app.state.client_tokens = client_tokens if client_tokens is not None else load_client_tokens()
|
||||
|
||||
@app.exception_handler(RequestValidationError)
|
||||
async def validation_exception_handler(
|
||||
_request: Request,
|
||||
error: RequestValidationError,
|
||||
) -> JSONResponse:
|
||||
return JSONResponse(
|
||||
status_code=status.HTTP_400_BAD_REQUEST, content={"detail": error.errors()}
|
||||
)
|
||||
|
||||
def authorize_token(authorization: str | None = Header(default=None)) -> str:
|
||||
if authorization is None or not authorization.startswith("Bearer "):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail="missing bearer token",
|
||||
)
|
||||
|
||||
token = authorization.removeprefix("Bearer ").strip()
|
||||
owner_client_id = _client_id_for_token(app.state.client_tokens, token)
|
||||
if owner_client_id is None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail="invalid bearer token",
|
||||
)
|
||||
return owner_client_id
|
||||
|
||||
@app.post("/api/v1/client/heartbeat", response_model=HeartbeatResponse)
|
||||
def heartbeat(
|
||||
payload: HeartbeatRequest,
|
||||
token_client_id: str = Depends(authorize_token),
|
||||
) -> HeartbeatResponse:
|
||||
if token_client_id != payload.client_id:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail="token is not allowed for this client_id",
|
||||
)
|
||||
|
||||
now = datetime.now(UTC)
|
||||
app.state.store.upsert_client(payload, now)
|
||||
app.state.store.record_command_result(payload, now)
|
||||
|
||||
command = app.state.store.take_next_command(payload.client_id, now)
|
||||
return HeartbeatResponse(
|
||||
server_time=now,
|
||||
next_poll_after_seconds=(
|
||||
DEFAULT_COMMAND_POLL_SECONDS if command else DEFAULT_IDLE_POLL_SECONDS
|
||||
),
|
||||
command=command,
|
||||
)
|
||||
|
||||
return app
|
||||
|
||||
|
||||
def _client_id_for_token(client_tokens: dict[str, str], token: str) -> str | None:
|
||||
for client_id, expected_token in client_tokens.items():
|
||||
if token == expected_token:
|
||||
return client_id
|
||||
return None
|
||||
|
||||
|
||||
app = create_app()
|
||||
@@ -0,0 +1,16 @@
|
||||
import os
|
||||
|
||||
DEFAULT_IDLE_POLL_SECONDS = 60
|
||||
DEFAULT_COMMAND_POLL_SECONDS = 10
|
||||
|
||||
|
||||
def load_client_tokens() -> dict[str, str]:
|
||||
raw_tokens = os.environ.get("NEXUS_SYNC_CLIENT_TOKENS", "dev-client:dev-token")
|
||||
tokens: dict[str, str] = {}
|
||||
for item in raw_tokens.split(","):
|
||||
if not item.strip():
|
||||
continue
|
||||
client_id, separator, token = item.partition(":")
|
||||
if separator and client_id and token:
|
||||
tokens[client_id] = token
|
||||
return tokens
|
||||
@@ -0,0 +1,143 @@
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from typing import Protocol
|
||||
|
||||
from fastapi import HTTPException, status
|
||||
|
||||
from nexus_sync.common import (
|
||||
ClientRecord,
|
||||
Command,
|
||||
CommandRecord,
|
||||
CommandResultRecord,
|
||||
CommandStatus,
|
||||
HeartbeatRequest,
|
||||
)
|
||||
|
||||
|
||||
class Store(Protocol):
|
||||
def upsert_client(self, heartbeat: HeartbeatRequest, now: datetime) -> ClientRecord:
|
||||
pass
|
||||
|
||||
def record_command_result(self, heartbeat: HeartbeatRequest, now: datetime) -> None:
|
||||
pass
|
||||
|
||||
def take_next_command(self, client_id: str, now: datetime) -> Command | None:
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
class InMemoryStore:
|
||||
clients: dict[str, ClientRecord] = field(default_factory=dict)
|
||||
commands: dict[str, CommandRecord] = field(default_factory=dict)
|
||||
results: list[CommandResultRecord] = field(default_factory=list)
|
||||
|
||||
def upsert_client(self, heartbeat: HeartbeatRequest, now: datetime) -> ClientRecord:
|
||||
existing = self.clients.get(heartbeat.client_id)
|
||||
created_at = existing.created_at if existing else now
|
||||
token_hash = existing.token_hash if existing else None
|
||||
|
||||
record = ClientRecord(
|
||||
id=heartbeat.client_id,
|
||||
hostname=heartbeat.client.hostname,
|
||||
platform=heartbeat.client.platform,
|
||||
version=heartbeat.client.version,
|
||||
created_at=created_at,
|
||||
last_seen_at=now,
|
||||
is_active=True,
|
||||
token_hash=token_hash,
|
||||
)
|
||||
self.clients[heartbeat.client_id] = record
|
||||
return record
|
||||
|
||||
def enqueue_command(self, command: Command, client_id: str, now: datetime) -> CommandRecord:
|
||||
record = CommandRecord(
|
||||
id=command.id,
|
||||
client_id=client_id,
|
||||
kind=command.kind,
|
||||
name=command.name,
|
||||
args=command.args,
|
||||
status=CommandStatus.PENDING,
|
||||
timeout_seconds=command.timeout_seconds,
|
||||
created_at=now,
|
||||
)
|
||||
self.commands[record.id] = record
|
||||
return record
|
||||
|
||||
def record_command_result(self, heartbeat: HeartbeatRequest, now: datetime) -> None:
|
||||
result = heartbeat.last_command_result
|
||||
if result is None:
|
||||
return
|
||||
|
||||
command = self.commands.get(result.command_id)
|
||||
if command is None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_409_CONFLICT,
|
||||
detail="command result references unknown command",
|
||||
)
|
||||
if command.client_id != heartbeat.client_id:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_409_CONFLICT,
|
||||
detail="command result references command for another client",
|
||||
)
|
||||
if command.status in TERMINAL_COMMAND_STATUSES:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_409_CONFLICT,
|
||||
detail="command result references terminal command",
|
||||
)
|
||||
|
||||
finished_at = result.finished_at or now
|
||||
self.commands[command.id] = command.model_copy(
|
||||
update={
|
||||
"status": CommandStatus(result.status.value),
|
||||
"finished_at": finished_at,
|
||||
}
|
||||
)
|
||||
self.results.append(
|
||||
CommandResultRecord(
|
||||
command_id=result.command_id,
|
||||
client_id=heartbeat.client_id,
|
||||
status=result.status,
|
||||
started_at=result.started_at,
|
||||
finished_at=finished_at,
|
||||
return_code=result.return_code,
|
||||
stdout=result.stdout,
|
||||
stderr=result.stderr,
|
||||
received_at=now,
|
||||
)
|
||||
)
|
||||
|
||||
def take_next_command(self, client_id: str, now: datetime) -> Command | None:
|
||||
pending = sorted(
|
||||
(
|
||||
command
|
||||
for command in self.commands.values()
|
||||
if command.client_id == client_id and command.status == CommandStatus.PENDING
|
||||
),
|
||||
key=lambda command: command.created_at,
|
||||
)
|
||||
if not pending:
|
||||
return None
|
||||
|
||||
record = pending[0]
|
||||
self.commands[record.id] = record.model_copy(
|
||||
update={
|
||||
"status": CommandStatus.DELIVERED,
|
||||
"attempts": record.attempts + 1,
|
||||
"delivered_at": now,
|
||||
}
|
||||
)
|
||||
return Command(
|
||||
id=record.id,
|
||||
kind=record.kind,
|
||||
name=record.name,
|
||||
args=record.args,
|
||||
timeout_seconds=record.timeout_seconds,
|
||||
)
|
||||
|
||||
|
||||
TERMINAL_COMMAND_STATUSES = {
|
||||
CommandStatus.SUCCEEDED,
|
||||
CommandStatus.FAILED,
|
||||
CommandStatus.TIMED_OUT,
|
||||
CommandStatus.REJECTED,
|
||||
}
|
||||
Reference in New Issue
Block a user