devplacepy/tests/unit/services/nodes/dispatcher.py

# retoor <retoor@molodetz.nl>

import base64
import json
import random
import time

import pytest

from devplacepy.database import get_table
from devplacepy.database import nodes as store
from devplacepy.services.nodes.broker import BodyError, JobBroker, clean_headers, decode_body
from devplacepy.services.nodes.circuit import CircuitBreaker
from devplacepy.services.nodes.dispatcher import JOB_HEADERS, NodeDispatcher
from devplacepy.services.nodes.ledger import UsageLedger
from devplacepy.services.nodes.policy import DestinationPolicy
from devplacepy.services.nodes.registry import NodeCounters, NodeRegistry
from devplacepy.services.nodes.scheduler import MEGABYTE, NodeScheduler
from devplacepy.services.nodes.session import NodeSession
from devplacepy.services.nodes.spec import SPEC
from devplacepy.utils import generate_uid
from tests.conftest import run_async
from tests.unit.services.nodes.scheduler import Clock, node, settings

LIMITS = SPEC.limits
TARGET_URL = "https://www.reddit.com/r/python.json"


class Rig:
    def __init__(self):
        self.clock = Clock()
        self.counters = NodeCounters()
        self.registry = NodeRegistry(self.counters)
        self.breaker = CircuitBreaker(clock=self.clock)
        self.scheduler = NodeScheduler(self.registry, self.breaker, self.clock, rng=random.Random(2))
        self.broker = JobBroker()
        self.dispatcher = NodeDispatcher(
            self.registry, self.scheduler, self.broker, self.breaker, UsageLedger(), self.counters,
            clock=self.clock,
        )

    def connect(self, **overrides):
        owner = generate_uid()
        device, _ = store.create_device(owner, "d", "native")
        session = NodeSession(
            node().websocket, {"uid": device["uid"], "user_uid": owner, "kind": "native"}
        )
        session.hello = {"modes": ["fetch"]}
        session.settings = settings(**overrides)
        session.roll_day()
        self.dispatcher.attach(session)
        self.registry.admit(session)
        return session

    def enqueue(self, url=TARGET_URL, caller="c1", kind="user", mode="any"):
        target = DestinationPolicy().check(url, "GET", mode)
        return self.broker.enqueue(kind, caller, target, mode)

    def tick(self):
        return run_async(self.dispatcher.tick())

    def answer(self, session, uid, **fields):
        frame = {"type": "job_result", "job_id": uid, "status": 200, "body_b64": base64.b64encode(b"hello").decode()}
        frame.update(fields)
        run_async(session._handle(frame))


@pytest.fixture
def rig(local_db):
    get_table("node_requests").delete()
    yield Rig()
    get_table("node_requests").delete()


def sent_jobs(session):
    return [frame for frame in session.websocket.sent if frame["type"] == "job"]


def test_a_queued_request_becomes_a_job_frame(rig):
    session = rig.connect()
    uid = rig.enqueue()
    assert rig.tick() == 1
    job = sent_jobs(session)[0]
    assert job == {"type": "job", "job_id": uid, "method": "GET", "url": TARGET_URL, "headers": JOB_HEADERS, "timeout": LIMITS.job_timeout_seconds}
    assert store.get_request(uid)["status"] == "assigned"
    assert store.get_request(uid)["device_uid"] == session.device_uid
    assert uid in session.inflight


def test_the_result_completes_the_request_and_is_metered(rig):
    session = rig.connect()
    uid = rig.enqueue()
    rig.tick()
    rig.answer(session, uid, headers={"Content-Type": "application/json", "Set-Cookie": "x=1"})
    result = rig.broker.read(uid)
    assert result.ok and result.http_status == 200 and result.body == b"hello"
    assert result.headers == {"content-type": "application/json"}
    assert uid not in session.inflight
    today = store.usage_today(session.device_uid)
    assert today["requests_ok"] == 1 and today["bytes_down"] == 5 and today["bytes_up"] > len(TARGET_URL)
    assert store.get_donor(session.user_uid)["requests_ok"] == 1
    assert session.used_today == today["bytes_up"] + today["bytes_down"]
    assert rig.counters.snapshot()["jobs_done"] == 1


def test_no_eligible_node_leaves_the_request_queued(rig):
    uid = rig.enqueue()
    assert rig.tick() == 0
    assert store.get_request(uid)["status"] == "queued"
    rig.connect(allowed_platforms=[])
    assert rig.tick() == 0
    assert store.get_request(uid)["status"] == "queued"


def test_one_request_goes_to_exactly_one_node(rig):
    sessions = [rig.connect() for _ in range(4)]
    rig.enqueue()
    assert rig.tick() == 1
    assert sum(len(sent_jobs(item)) for item in sessions) == 1


def test_a_caller_never_gets_its_own_node(rig):
    session = rig.connect()
    uid = rig.enqueue(caller=session.user_uid)
    assert rig.tick() == 0
    assert store.get_request(uid)["status"] == "queued"
    guest = rig.enqueue(caller=session.user_uid, kind="guest")
    assert rig.tick() == 1 and guest in session.inflight


def test_a_policy_failure_at_dispatch_time_fails_the_request(rig):
    rig.connect()
    uid = store.enqueue_request("user", "c", "https://evil.example/", "GET", "x", "any")
    rig.tick()
    row = store.get_request(uid)
    assert row["status"] == "failed" and row["error"] == "host_not_allowed"
    assert rig.counters.snapshot()["jobs_refused"] == 1


def test_a_job_error_fails_the_request_and_counts_against_the_node(rig):
    session = rig.connect()
    uid = rig.enqueue()
    rig.tick()
    run_async(session._handle({"type": "job_error", "job_id": uid, "code": "timeout"}))
    assert store.get_request(uid)["error"] == "timeout"
    assert store.usage_today(session.device_uid)["requests_failed"] == 1
    assert rig.counters.snapshot()["jobs_failed"] == 1


def test_five_node_failures_open_that_nodes_circuit(rig):
    session = rig.connect()
    for _ in range(LIMITS.breaker_failures):
        uid = rig.enqueue()
        session.host_started.clear()
        rig.scheduler._destination_started.clear()
        rig.tick()
        run_async(session._handle({"type": "job_error", "job_id": uid, "code": "network"}))
    follow = rig.enqueue()
    session.host_started.clear()
    rig.scheduler._destination_started.clear()
    assert rig.tick() == 0
    assert store.get_request(follow)["status"] == "queued"
    rig.clock.now += LIMITS.breaker_open_seconds
    assert rig.tick() == 1


def test_donor_limit_codes_pause_the_node_without_blaming_it(rig):
    session = rig.connect(bandwidth_mb_day=50)
    uid = rig.enqueue()
    rig.tick()
    run_async(session._handle({"type": "job_error", "job_id": uid, "code": "bandwidth_exhausted"}))
    assert session.used_today == 50 * MEGABYTE
    assert rig.breaker.allows(("node", session.device_uid))
    other = rig.connect()
    rig.clock.now += 10
    uid = rig.enqueue()
    assert rig.tick() == 1
    run_async(other._handle({"type": "job_error", "job_id": uid, "code": "outside_hours"}))
    assert other.paused_until > rig.clock.now
    assert rig.breaker.allows(("node", other.device_uid))


def test_a_donor_refusal_of_an_allowed_destination_raises_a_flag(rig):
    session = rig.connect()
    uid = rig.enqueue()
    rig.tick()
    run_async(session._handle({"type": "job_error", "job_id": uid, "code": "host_not_allowed"}))
    kinds = [flag["kind"] for flag in store.list_flags()]
    assert "refusal_mismatch" in kinds
    assert rig.breaker.allows(("node", session.device_uid))


@pytest.mark.parametrize(
    "fields,code",
    [
        ({"status": 99}, "bad_status"),
        ({"status": 600}, "bad_status"),
        ({"body_b64": "***not base64***"}, "bad_base64"),
        ({"body_b64": "abc"}, "bad_base64"),
    ],
)
def test_forged_results_are_integrity_failures(rig, fields, code):
    session = rig.connect()
    uid = rig.enqueue()
    rig.tick()
    rig.answer(session, uid, **fields)
    assert store.get_request(uid)["error"] == code
    assert store.usage_today(session.device_uid)["integrity_failures"] == 1
    assert rig.counters.snapshot()["integrity_failures"] == 1
    assert any(flag["kind"] == "integrity" for flag in store.list_flags())
    assert store.get_donor(session.user_uid)["requests_ok"] == 0


def test_an_oversized_body_is_truncated_and_flagged(rig):
    session = rig.connect()
    uid = rig.enqueue()
    rig.tick()
    big = base64.b64encode(b"x" * (LIMITS.max_body_bytes + 10)).decode()
    rig.answer(session, uid, body_b64=big)
    result = rig.broker.read(uid)
    assert result.truncated and len(result.body) == LIMITS.max_body_bytes
    assert rig.counters.snapshot()["results_truncated"] == 1


def test_a_donor_declared_truncation_is_kept(rig):
    session = rig.connect()
    uid = rig.enqueue()
    rig.tick()
    rig.answer(session, uid, truncated=True)
    assert rig.broker.read(uid).truncated


def test_a_result_for_an_unknown_job_is_a_protocol_error(rig):
    session = rig.connect()
    from devplacepy.services.nodes.frames import FrameError

    with pytest.raises(FrameError):
        rig.answer(session, "never-assigned")


def test_a_job_cannot_be_answered_twice(rig):
    session = rig.connect()
    uid = rig.enqueue()
    rig.tick()
    rig.answer(session, uid)
    from devplacepy.services.nodes.frames import FrameError

    with pytest.raises(FrameError):
        rig.answer(session, uid)
    assert store.get_donor(session.user_uid)["requests_ok"] == 1


def test_rate_limited_and_server_error_results_open_the_destination_circuit(rig):
    for index in range(LIMITS.breaker_failures):
        session = rig.connect()
        uid = rig.enqueue()
        rig.scheduler._destination_started.clear()
        rig.tick()
        rig.answer(session, uid, status=429 if index % 2 else 503)
        assert rig.broker.read(uid).ok
        rig.registry.release(session)
    assert not rig.breaker.allows(("destination", "reddit.com"))
    rig.connect()
    rig.enqueue()
    rig.scheduler._destination_started.clear()
    assert rig.tick() == 0


def test_a_healthy_result_closes_the_destination_circuit(rig):
    rig.breaker.failure(("destination", "reddit.com"))
    session = rig.connect()
    uid = rig.enqueue()
    rig.tick()
    rig.answer(session, uid, status=200)
    assert rig.breaker.state(("destination", "reddit.com")) == "closed"


def test_a_job_that_never_answers_times_out(rig):
    session = rig.connect()
    uid = rig.enqueue()
    rig.tick()
    rig.clock.now += LIMITS.job_timeout_seconds + LIMITS.job_grace_seconds + 1
    assert run_async(rig.dispatcher.sweep_inflight()) == 1
    assert store.get_request(uid)["error"] == "timeout"
    assert session.inflight == {}


def test_a_disconnect_fails_every_inflight_job(rig):
    session = rig.connect(max_concurrent=2)
    uids = [rig.enqueue(), rig.enqueue()]
    for _ in uids:
        session.host_started.clear()
        rig.scheduler._destination_started.clear()
        assert rig.tick() == 1
    assert len(session.inflight) == 2
    assert rig.dispatcher.release_session(session) == 2
    for uid in uids:
        row = store.get_request(uid)
        assert row["status"] == "failed" and row["error"] == "node_disconnected"
    assert session.inflight == {}


def test_release_session_marks_assigned_jobs_as_node_disconnected(rig):
    session = rig.connect()
    uid = rig.enqueue()
    rig.tick()
    assert rig.dispatcher.release_session(session) == 1
    assert store.get_request(uid)["error"] == "node_disconnected"
    assert rig.dispatcher.release_session(session) == 0


def test_a_failed_send_releases_the_job(rig):
    session = rig.connect()

    async def broken(*args, **kwargs):
        raise ConnectionError("gone")

    session.send = broken
    uid = rig.enqueue()
    assert rig.tick() == 0
    assert store.get_request(uid)["error"] == "send_failed"
    assert session.inflight == {}


def test_an_idle_tick_does_no_writes_and_no_work(rig):
    rig.connect()
    started = time.perf_counter()
    for _ in range(50):
        assert rig.tick() == 0
    assert (time.perf_counter() - started) / 50 < 0.02


def test_the_cap_refusing_the_ledger_pauses_the_node(rig):
    session = rig.connect(bandwidth_mb_day=1)
    store.save_node_settings(session.user_uid, {"enabled": True, "bandwidth_mb_day": 1})
    session.used_today = 0
    uid = rig.enqueue()
    rig.tick()
    store.record_job(session.device_uid, session.user_uid, ok=True, bytes_down=MEGABYTE, cap_bytes=MEGABYTE)
    rig.answer(session, uid)
    assert rig.broker.read(uid).ok
    assert session.used_today >= MEGABYTE


def test_every_assigned_job_is_answered_at_most_once_under_a_random_workload(rig):
    rng = random.Random(11)
    sessions = [rig.connect(max_concurrent=2) for _ in range(4)]
    seen = set()
    for _ in range(150):
        rig.enqueue(caller=generate_uid())
        rig.clock.now += 10
        rig.scheduler._destination_started.clear()
        rig.tick()
        for session in sessions:
            for job_id in list(session.inflight):
                assert job_id not in seen
                action = rng.random()
                if action < 0.6:
                    rig.answer(session, job_id)
                    seen.add(job_id)
                elif action < 0.8:
                    run_async(session._handle({"type": "job_error", "job_id": job_id, "code": "network"}))
                    seen.add(job_id)
            assert len(session.inflight) <= session.max_concurrent
    open_now = {job for item in sessions for job in item.inflight}
    assert not (open_now & seen)


def test_body_helpers():
    assert decode_body(base64.b64encode(b"abc").decode()) == (b"abc", False)
    with pytest.raises(BodyError):
        decode_body("!!!")
    assert clean_headers({"A": "1", "set-cookie": "s", "Authorization": "x", 5: "x", "b": 5}) == {"a": "1"}
    assert clean_headers("nope") == {}
    many = {f"h{i}": "v" for i in range(500)}
    assert len(clean_headers(many)) == 64
    assert len(next(iter(clean_headers({"k": "v" * 5000}).values()))) == 1024
    assert json.loads(json.dumps(clean_headers({"A": "1"}))) == {"a": "1"}


def test_a_saturated_destination_is_scanned_once_per_tick_not_once_per_request(rig, monkeypatch):
    rig.connect(allowed_platforms=[])
    for index in range(30):
        rig.enqueue(caller=f"c{index}")
    scans = []
    original = rig.scheduler.candidates
    monkeypatch.setattr(rig.scheduler, "candidates", lambda *a: scans.append(a) or original(*a))
    assert rig.tick() == 0
    assert len(scans) == 1


def test_saturation_does_not_hide_other_destinations(rig):
    session = rig.connect()
    for index in range(3):
        rig.enqueue(caller=f"c{index}")
    rig.scheduler._destination_started["reddit.com"] = rig.clock.now
    assert rig.tick() == 0
    rig.clock.now += 10
    assert rig.tick() == 1
    assert len(session.inflight) == 1


def test_the_owner_exclusion_does_not_saturate_the_destination(rig):
    mine = rig.connect()
    other = rig.connect()
    first = rig.enqueue(caller=mine.user_uid)
    second = rig.enqueue(caller=other.user_uid)
    rig.clock.now += 10
    rig.tick()
    assert first in other.inflight
    rig.clock.now += 10
    rig.scheduler._destination_started.clear()
    other.host_started.clear()
    rig.tick()
    assert second in mine.inflight