leaderless_log_compactor_server.py

from __future__ import annotations

import argparse
import json
import os
import random
import signal
import sys
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from http import HTTPStatus
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Any, Callable, Protocol
from urllib.parse import urlsplit

import oxia.ex

from .leaderless_log_common import TopicPartition
from .leaderless_log_compactor import (
    CompactorClaim,
    LeaderlessLogCompactor,
    PartitionNotInitialized,
)
from .leaderless_log_metrics import LeaderlessLogMetrics
from .s3_billing import (
    S3BucketUsageRefreshResult,
    S3BucketUsageRefresher,
    build_bucket_usage_refresher_from_env,
)


class CompactorLike(Protocol):
    def acquire_partition_claim(
        self,
        topic: str,
        partition: int,
        *,
        compactor_id: str,
    ) -> CompactorClaim | None: ...

    def release_partition_claim(self, claim: CompactorClaim) -> None: ...

    def compact_once(
        self,
        topic: str,
        partition: int,
        *,
        max_offsets: int | None = None,
    ) -> object | None: ...

    def discover_partitions(self) -> set[TopicPartition]: ...

    def get_partition_status(self, topic: str, partition: int) -> dict[str, Any]: ...

    def close(self) -> None: ...


@dataclass(frozen=True)
class CompactorServiceConfig:
    host: str
    port: int
    compactor_id: str
    targets: tuple[TopicPartition, ...]
    workers: int
    discovery_interval_ms: int
    idle_sleep_ms: int
    error_sleep_ms: int
    idle_jitter_ms: int
    error_jitter_ms: int
    max_offsets_per_run: int | None = None

    def __post_init__(self) -> None:
        if self.host == "":
            raise ValueError("host must be non-empty")
        if self.port <= 0:
            raise ValueError("port must be a positive integer")
        if self.compactor_id == "":
            raise ValueError("compactor_id must be non-empty")
        if self.workers <= 0:
            raise ValueError("workers must be a positive integer")
        if self.discovery_interval_ms <= 0:
            raise ValueError("discovery_interval_ms must be a positive integer")
        if self.idle_sleep_ms <= 0:
            raise ValueError("idle_sleep_ms must be a positive integer")
        if self.error_sleep_ms <= 0:
            raise ValueError("error_sleep_ms must be a positive integer")
        if self.idle_jitter_ms < 0:
            raise ValueError("idle_jitter_ms must be >= 0")
        if self.error_jitter_ms < 0:
            raise ValueError("error_jitter_ms must be >= 0")
        if self.max_offsets_per_run is not None and self.max_offsets_per_run <= 0:
            raise ValueError("max_offsets_per_run must be a positive integer")


class CompactorService:
    def __init__(
        self,
        config: CompactorServiceConfig,
        compactor_factory: Callable[[], CompactorLike],
        *,
        metrics: LeaderlessLogMetrics | None = None,
        random_source: random.Random | None = None,
    ) -> None:
        self._config = config
        self._compactor_factory = compactor_factory
        self._metrics = metrics or LeaderlessLogMetrics()
        self._random = random_source or random.Random()
        self._stop_event = threading.Event()
        self._thread_local = threading.local()
        self._compactors: list[CompactorLike] = []
        self._compactors_lock = threading.Lock()
        self._state_lock = threading.Lock()
        self._configured_targets = set(config.targets)
        self._known_targets = set(config.targets)
        self._discovered_targets: set[TopicPartition] = set()
        self._inflight: set[TopicPartition] = set()
        self._next_discovery_at_ms = 0
        self._executor = ThreadPoolExecutor(
            max_workers=config.workers,
            thread_name_prefix="compactor-worker",
        )
        self._scheduler_thread = threading.Thread(
            target=self._run_scheduler_loop,
            name="compactor-scheduler",
            daemon=True,
        )
        self._metrics.configure_compactor(
            compactor_id=config.compactor_id,
            host=config.host,
            port=config.port,
            workers=config.workers,
            discovery_interval_ms=config.discovery_interval_ms,
        )
        self._metrics.set_compactor_target_counts(
            configured_targets=len(self._configured_targets),
            discovered_targets=0,
            known_targets=len(self._known_targets),
        )
        self._metrics.set_compactor_inflight(0)

    def start(self) -> None:
        if self._scheduler_thread.is_alive():
            return
        self._scheduler_thread.start()

    def close(self) -> None:
        self._metrics.set_compactor_status("stopping")
        self._stop_event.set()
        if self._scheduler_thread.is_alive():
            self._scheduler_thread.join(timeout=5)
        self._executor.shutdown(wait=True, cancel_futures=False)
        self._close_all_compactors()

    def run_discovery_pass(self) -> set[TopicPartition]:
        compactor = self._get_thread_compactor()
        discovered_at_ms = _now_ms()
        started = time.perf_counter()
        try:
            discovered = compactor.discover_partitions()
        except oxia.ex.SessionNotFound:
            self._discard_thread_compactor(compactor)
            raise
        additions: set[TopicPartition] = set()
        with self._state_lock:
            for tp in discovered:
                if tp in self._known_targets:
                    continue
                self._known_targets.add(tp)
                if tp not in self._configured_targets:
                    self._discovered_targets.add(tp)
                additions.add(tp)
            configured_targets = len(self._configured_targets)
            discovered_targets = len(self._discovered_targets)
            known_targets = len(self._known_targets)
        self._metrics.set_compactor_target_counts(
            configured_targets=configured_targets,
            discovered_targets=discovered_targets,
            known_targets=known_targets,
        )
        self._metrics.observe_compactor_discovery(
            discovered_at_ms=discovered_at_ms,
            duration_seconds=time.perf_counter() - started,
            additions=len(additions),
        )
        self._next_discovery_at_ms = discovered_at_ms + self._config.discovery_interval_ms
        return additions

    def known_targets(self) -> set[TopicPartition]:
        with self._state_lock:
            return set(self._known_targets)

    def _run_scheduler_loop(self) -> None:
        while not self._stop_event.is_set():
            now_ms = _now_ms()
            if now_ms >= self._next_discovery_at_ms:
                try:
                    self.run_discovery_pass()
                except Exception as exc:
                    print(
                        f"compactor discovery failed: {exc}",
                        file=sys.stderr,
                    )
                    self._sleep(self._config.error_sleep_ms, self._config.error_jitter_ms)
                    continue

            started_at_ms = _now_ms()
            started = time.perf_counter()
            submitted = self._submit_scheduler_pass()
            self._metrics.observe_compactor_scheduler_pass(
                started_at_ms=started_at_ms,
                duration_seconds=time.perf_counter() - started,
            )

            if self._stop_event.is_set():
                break
            if submitted > 0:
                self._sleep(self._config.idle_sleep_ms, self._config.idle_jitter_ms)
                continue
            if self._inflight_count() > 0:
                self._stop_event.wait(0.05)
                continue
            self._sleep(self._config.idle_sleep_ms, self._config.idle_jitter_ms)

    def _submit_scheduler_pass(self) -> int:
        targets = list(self.known_targets())
        if not targets:
            return 0
        self._random.shuffle(targets)
        submitted = 0
        for tp in targets:
            if self._stop_event.is_set():
                break
            if not self._try_submit_target(tp):
                if self._inflight_count() >= self._config.workers:
                    break
                continue
            submitted += 1
        return submitted

    def _try_submit_target(self, tp: TopicPartition) -> bool:
        with self._state_lock:
            if tp in self._inflight:
                return False
            if len(self._inflight) >= self._config.workers:
                return False
            self._inflight.add(tp)
            inflight = len(self._inflight)
        self._metrics.set_compactor_inflight(inflight)
        try:
            self._executor.submit(self._run_partition_attempt, tp)
        except Exception:
            self._finish_inflight(tp)
            raise
        return True

    def _run_partition_attempt(self, tp: TopicPartition) -> None:
        compactor = self._get_thread_compactor()
        started = time.perf_counter()
        claim_started: float | None = None
        claim: CompactorClaim | None = None
        claim_attempted = False
        claim_acquired = False
        result = "error"
        error: Exception | None = None
        try:
            claim_attempted = True
            claim = compactor.acquire_partition_claim(
                tp.topic,
                tp.partition,
                compactor_id=self._config.compactor_id,
            )
            if claim is None:
                result = "claim_busy"
                return
            claim_acquired = True
            claim_started = time.perf_counter()
            had_pending_compaction = self._has_pending_compaction(compactor, tp)
            outcome = compactor.compact_once(
                tp.topic,
                tp.partition,
                max_offsets=self._config.max_offsets_per_run,
            )
            if outcome is None:
                result = "no_candidate"
            elif had_pending_compaction:
                result = "recovered"
            else:
                result = "success"
        except Exception as exc:
            error = exc
            result = "error"
            if isinstance(exc, oxia.ex.SessionNotFound):
                self._discard_thread_compactor(compactor)
        finally:
            claim_hold_duration = None
            if claim is not None:
                if claim_started is not None:
                    claim_hold_duration = time.perf_counter() - claim_started
                try:
                    compactor.release_partition_claim(claim)
                except Exception as exc:
                    if error is None:
                        error = exc
                        result = "error"
            self._metrics.observe_compactor_attempt(
                result=result,
                duration_seconds=time.perf_counter() - started,
                claim_attempted=claim_attempted,
                claim_acquired=claim_acquired,
                claim_hold_duration_seconds=claim_hold_duration,
            )
            self._finish_inflight(tp)
            if error is not None:
                print(
                    f"compactor attempt failed for {tp.topic}[{tp.partition}]: "
                    f"{type(error).__name__}: {error!r}",
                    file=sys.stderr,
                )

    @staticmethod
    def _has_pending_compaction(compactor: CompactorLike, tp: TopicPartition) -> bool:
        try:
            status = compactor.get_partition_status(tp.topic, tp.partition)
        except PartitionNotInitialized:
            return False
        return status.get("pending_compaction") is not None

    def _finish_inflight(self, tp: TopicPartition) -> None:
        with self._state_lock:
            self._inflight.discard(tp)
            inflight = len(self._inflight)
        self._metrics.set_compactor_inflight(inflight)

    def _inflight_count(self) -> int:
        with self._state_lock:
            return len(self._inflight)

    def _sleep(self, base_ms: int, jitter_ms: int) -> None:
        delay_ms = base_ms
        if jitter_ms > 0:
            delay_ms += self._random.randint(0, jitter_ms)
        self._stop_event.wait(delay_ms / 1000.0)

    def _get_thread_compactor(self) -> CompactorLike:
        compactor = getattr(self._thread_local, "compactor", None)
        if compactor is not None:
            return compactor
        compactor = self._compactor_factory()
        with self._compactors_lock:
            self._compactors.append(compactor)
        self._thread_local.compactor = compactor
        return compactor

    def _discard_thread_compactor(self, compactor: CompactorLike) -> None:
        if getattr(self._thread_local, "compactor", None) is compactor:
            try:
                delattr(self._thread_local, "compactor")
            except AttributeError:
                pass
        with self._compactors_lock:
            try:
                self._compactors.remove(compactor)
            except ValueError:
                pass
        try:
            compactor.close()
        except Exception:
            pass

    def _close_all_compactors(self) -> None:
        with self._compactors_lock:
            compactors = list(self._compactors)
            self._compactors.clear()
        for compactor in compactors:
            compactor.close()


class CompactorHTTPServer(ThreadingHTTPServer):
    allow_reuse_address = True
    daemon_threads = True

    def __init__(
        self,
        server_address: tuple[str, int],
        service: CompactorService,
        *,
        s3_bucket_usage_refresher: S3BucketUsageRefresher,
        metrics: LeaderlessLogMetrics | None = None,
    ) -> None:
        super().__init__(server_address, CompactorRequestHandler)
        self.service = service
        self.metrics = metrics or service._metrics
        self.s3_bucket_usage_refresher = s3_bucket_usage_refresher


class CompactorRequestHandler(BaseHTTPRequestHandler):
    server_version = "LeaderlessLogCompactor/0.1"

    def do_GET(self) -> None:
        path = self._request_path()
        started = time.perf_counter()
        status = HTTPStatus.INTERNAL_SERVER_ERROR
        try:
            if path == "/health":
                status = HTTPStatus.OK
                self._write_json(status, self.server.metrics.health_snapshot())  # type: ignore[attr-defined]
                return
            if path == "/metrics":
                status = HTTPStatus.OK
                self._write_metrics(status)
                return
            if path == "/metrics/prometheus":
                status = HTTPStatus.OK
                self._write_metrics_prometheus(status)
                return
            status = HTTPStatus.NOT_FOUND
            self._write_json(status, {"error": "not found"})
        finally:
            self.server.metrics.observe_http_request(  # type: ignore[attr-defined]
                method="GET",
                path=path,
                status=int(status),
                duration_seconds=time.perf_counter() - started,
            )

    def do_POST(self) -> None:
        started = time.perf_counter()
        status = HTTPStatus.NOT_FOUND
        try:
            self._write_json(status, {"error": "not found"})
        finally:
            self.server.metrics.observe_http_request(  # type: ignore[attr-defined]
                method="POST",
                path=self._request_path(),
                status=int(status),
                duration_seconds=time.perf_counter() - started,
            )

    def _request_path(self) -> str:
        return urlsplit(self.path).path

    def _write_json(self, status: HTTPStatus, payload: dict[str, Any]) -> None:
        body = json.dumps(payload, separators=(",", ":")).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def _write_metrics(self, status: HTTPStatus) -> None:
        self.server.s3_bucket_usage_refresher.schedule_refresh_if_stale()  # type: ignore[attr-defined]
        body = self.server.metrics.render_latest()  # type: ignore[attr-defined]
        self.send_response(status)
        self.send_header("Content-Type", LeaderlessLogMetrics.CONTENT_TYPE)
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def _write_metrics_prometheus(self, status: HTTPStatus) -> None:
        self.server.s3_bucket_usage_refresher.schedule_refresh_if_stale()  # type: ignore[attr-defined]
        body = self.server.metrics.render_prometheus()  # type: ignore[attr-defined]
        self.send_response(status)
        self.send_header("Content-Type", LeaderlessLogMetrics.PROMETHEUS_CONTENT_TYPE)
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)


def parse_targets(raw: str | None) -> tuple[TopicPartition, ...]:
    if raw is None or raw.strip() == "":
        return ()
    targets: list[TopicPartition] = []
    seen: set[TopicPartition] = set()
    for item in raw.split(","):
        token = item.strip()
        if token == "":
            continue
        if ":" not in token:
            raise ValueError(f"invalid compactor target: {token}")
        topic, partition_text = token.rsplit(":", 1)
        if topic == "":
            raise ValueError(f"invalid compactor target: {token}")
        try:
            partition = int(partition_text)
        except ValueError as exc:
            raise ValueError(f"invalid compactor target: {token}") from exc
        if partition < 0:
            raise ValueError(f"invalid compactor target: {token}")
        tp = TopicPartition(topic=topic, partition=partition)
        if tp in seen:
            continue
        seen.add(tp)
        targets.append(tp)
    return tuple(targets)


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Long-running HTTP compactor for the leaderless log.",
    )
    parser.add_argument("--host")
    parser.add_argument("--port", type=int)
    parser.add_argument("--compactor-id")
    parser.add_argument("--targets")
    parser.add_argument("--workers", type=int)
    parser.add_argument("--discovery-interval-ms", type=int)
    parser.add_argument("--idle-sleep-ms", type=int)
    parser.add_argument("--error-sleep-ms", type=int)
    parser.add_argument("--idle-jitter-ms", type=int)
    parser.add_argument("--error-jitter-ms", type=int)
    parser.add_argument("--max-offsets-per-run", type=int)
    return parser.parse_args(argv)


def build_config_from_args(args: argparse.Namespace) -> CompactorServiceConfig:
    host = args.host if args.host is not None else os.environ.get("LLOG_COMPACTOR_BIND_HOST", "127.0.0.1")
    port = args.port if args.port is not None else int(os.environ.get("LLOG_COMPACTOR_PORT", "8081"))
    compactor_id = (
        args.compactor_id
        or os.environ.get("LLOG_COMPACTOR_ID")
        or _default_compactor_id(host, port)
    )
    targets = parse_targets(args.targets if args.targets is not None else os.environ.get("LLOG_COMPACTOR_TARGETS"))
    workers = args.workers if args.workers is not None else int(os.environ.get("LLOG_COMPACTOR_WORKERS", "1"))
    discovery_interval_ms = args.discovery_interval_ms if args.discovery_interval_ms is not None else int(
        os.environ.get("LLOG_COMPACTOR_DISCOVERY_INTERVAL_MS", "30000")
    )
    idle_sleep_ms = args.idle_sleep_ms if args.idle_sleep_ms is not None else int(
        os.environ.get("LLOG_COMPACTOR_IDLE_SLEEP_MS", "1000")
    )
    error_sleep_ms = args.error_sleep_ms if args.error_sleep_ms is not None else int(
        os.environ.get("LLOG_COMPACTOR_ERROR_SLEEP_MS", "2000")
    )
    idle_jitter_ms = args.idle_jitter_ms if args.idle_jitter_ms is not None else int(
        os.environ.get("LLOG_COMPACTOR_IDLE_JITTER_MS", "250")
    )
    error_jitter_ms = args.error_jitter_ms if args.error_jitter_ms is not None else int(
        os.environ.get("LLOG_COMPACTOR_ERROR_JITTER_MS", str(idle_jitter_ms))
    )
    max_offsets_per_run = args.max_offsets_per_run
    if max_offsets_per_run is None:
        env_value = os.environ.get("LLOG_COMPACTOR_MAX_OFFSETS_PER_RUN")
        if env_value not in (None, ""):
            max_offsets_per_run = int(env_value)
    return CompactorServiceConfig(
        host=host,
        port=port,
        compactor_id=compactor_id,
        targets=targets,
        workers=workers,
        discovery_interval_ms=discovery_interval_ms,
        idle_sleep_ms=idle_sleep_ms,
        error_sleep_ms=error_sleep_ms,
        idle_jitter_ms=idle_jitter_ms,
        error_jitter_ms=error_jitter_ms,
        max_offsets_per_run=max_offsets_per_run,
    )


def build_compactor_factory(
    *,
    metrics: LeaderlessLogMetrics,
    client_identifier_base: str | None,
) -> Callable[[], LeaderlessLogCompactor]:
    def factory() -> LeaderlessLogCompactor:
        client_identifier = client_identifier_base
        if client_identifier is not None:
            client_identifier = f"{client_identifier}:{threading.current_thread().name}"
        return LeaderlessLogCompactor.from_env(
            metrics=metrics,
            client_identifier=client_identifier,
        )

    return factory


def _observe_s3_bucket_usage_refresh(
    metrics: LeaderlessLogMetrics,
    result: S3BucketUsageRefreshResult,
) -> None:
    if not result.refreshed or result.refreshed_at_ms is None or result.success is None:
        return
    snapshot = result.snapshot
    metrics.observe_s3_bucket_usage_refresh(
        refreshed_at_ms=result.refreshed_at_ms,
        success=result.success,
        object_count=None if snapshot is None else snapshot.object_count,
        stored_bytes=None if snapshot is None else snapshot.stored_bytes,
        error=result.error,
    )


def _default_compactor_id(host: str, port: int) -> str:
    return f"{host}:{port}"


def _now_ms() -> int:
    return int(time.time() * 1000)


def _env_bool(name: str, *, default: bool = False) -> bool:
    raw = os.environ.get(name)
    if raw is None or raw == "":
        return default
    return raw.strip().lower() in {"1", "true", "yes", "on"}


def main() -> int:
    args = parse_args()
    config = build_config_from_args(args)
    metrics = LeaderlessLogMetrics(
        enable_compaction_partition_metrics=_env_bool(
            "LLOG_COMPACTION_PARTITION_METRICS"
        )
    )
    s3_bucket_usage_refresher = build_bucket_usage_refresher_from_env(
        result_callback=lambda result: _observe_s3_bucket_usage_refresh(metrics, result),
        operation_observer=lambda operation, result, duration_seconds: metrics.observe_s3_operation(
            operation=operation,
            result=result,
            duration_seconds=duration_seconds,
        ),
    )
    metrics.set_s3_bucket_usage_refresh_interval_seconds(
        s3_bucket_usage_refresher.refresh_interval_seconds
    )
    s3_bucket_usage_refresher.refresh_now()
    client_identifier_base = os.environ.get("LLOG_CLIENT_IDENTIFIER") or config.compactor_id
    service = CompactorService(
        config,
        build_compactor_factory(
            metrics=metrics,
            client_identifier_base=client_identifier_base,
        ),
        metrics=metrics,
    )
    server = CompactorHTTPServer(
        (config.host, config.port),
        service,
        metrics=metrics,
        s3_bucket_usage_refresher=s3_bucket_usage_refresher,
    )
    shutdown_requested = threading.Event()

    def _request_shutdown(signum: int, _frame: object | None) -> None:
        del signum
        if shutdown_requested.is_set():
            return
        shutdown_requested.set()
        metrics.set_compactor_status("stopping")
        threading.Thread(target=server.shutdown, name="compactor-shutdown", daemon=True).start()

    signal.signal(signal.SIGINT, _request_shutdown)
    signal.signal(signal.SIGTERM, _request_shutdown)

    service.start()
    print(
        f"listening on http://{config.host}:{config.port} compactor_id={config.compactor_id}"
    )
    try:
        server.serve_forever()
    finally:
        server.server_close()
        service.close()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())