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
import asyncio
from concurrent.futures import ThreadPoolExecutor

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"


def run_in_thread(coro):
    with ThreadPoolExecutor(max_workers=1) as pool:
        return pool.submit(asyncio.run, coro).result()


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,
        )
        self.dispatcher.redundancy_percent = staticmethod(lambda: 0)

    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", "asn_class": "residential"}
        )
        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


def test_the_claim_records_the_donor_address_class(rig):
    session = rig.connect()
    session.asn_class = "mobile"
    uid = rig.enqueue()
    rig.tick()
    assert store.get_request(uid)["device_asn_class"] == "mobile"


def test_outcomes_are_observed_per_class(rig):
    session = rig.connect()
    uid = rig.enqueue()
    rig.tick()
    rig.answer(session, uid)
    other = rig.enqueue(url="https://www.reddit.com/r/two.json")
    rig.clock.now += 100
    rig.tick()
    rig.answer(session, other, status=403)
    third = rig.enqueue(url="https://www.reddit.com/r/three.json")
    rig.clock.now += 100
    rig.tick()
    rig.answer(session, third, status=404)
    seen = rig.counters.outcomes()
    assert seen["outcome:reddit:residential:ok"] == 2 and seen["outcome:reddit:residential:blocked"] == 1


def test_a_node_failure_is_observed_but_a_donor_limit_is_not(rig):
    session = rig.connect()
    uid = rig.enqueue()
    rig.tick()
    run_async(session._handle({"type": "job_error", "job_id": uid, "code": "fetch_failed"}))
    assert rig.counters.outcomes().get("outcome:reddit:residential:failed") == 1
    rig.clock.now += 100
    session.paused_until = 0
    second = rig.enqueue(url="https://www.reddit.com/r/two.json")
    rig.tick()
    run_async(session._handle({"type": "job_error", "job_id": second, "code": "outside_hours"}))
    assert rig.counters.outcomes().get("outcome:reddit:residential:failed") == 1


def test_on_battery_pauses_the_donor_without_blaming_it(rig):
    session = rig.connect()
    uid = rig.enqueue()
    rig.tick()
    run_async(session._handle({"type": "job_error", "job_id": uid, "code": "on_battery"}))
    assert session.paused_until > rig.clock()
    assert not rig.breaker.failures or rig.breaker.allows(("node", session.device_uid))
    assert store.get_request(uid)["error"] == "on_battery"


class SampleAll:
    def __init__(self, rig):
        rig.dispatcher.redundancy_percent = staticmethod(lambda: 100)


def mismatch_flags(device_uid=None):
    return [
        flag
        for flag in store.list_flags()
        if flag["kind"] == "redundancy_mismatch" and (device_uid is None or flag["device_uid"] == device_uid)
    ]


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


def test_a_sampled_fetch_result_is_cross_checked_on_a_different_donor(rig):
    SampleAll(rig)
    first, second = rig.connect(), rig.connect()
    uid = rig.enqueue()
    rig.tick()
    holder = first if uid in sent_to(first) else second
    other = second if holder is first else first
    rig.answer(holder, uid, headers={"content-type": "application/json"}, body_b64=base64.b64encode(b'{"kind":"Listing","data":{}}').decode())
    shadows = list(get_table("node_requests").find(shadow_of=uid))
    assert len(shadows) == 1 and shadows[0]["exclude_device"] == holder.device_uid
    assert shadows[0]["caller_kind"] == "shadow"
    rig.clock.now += 100
    assert rig.tick() == 1
    assert sent_to(other) and shadows[0]["uid"] in sent_to(other)
    rig.answer(other, shadows[0]["uid"], headers={"content-type": "application/json"}, body_b64=base64.b64encode(b'{"kind":"Listing","data":{}}').decode())
    counts = rig.counters.snapshot()
    assert counts["redundancy_checked"] == 1 and counts["redundancy_matched"] == 1
    assert not mismatch_flags(other.device_uid)


def test_a_forged_second_opinion_raises_a_flag_and_no_ban(rig):
    SampleAll(rig)
    first, second = rig.connect(), rig.connect()
    uid = rig.enqueue()
    rig.tick()
    holder = first if uid in sent_to(first) else second
    other = second if holder is first else first
    rig.answer(holder, uid, headers={"content-type": "application/json"}, body_b64=base64.b64encode(b'{"kind":"Listing","data":{}}').decode())
    shadow = list(get_table("node_requests").find(shadow_of=uid))[0]["uid"]
    rig.clock.now += 100
    rig.tick()
    rig.answer(other, shadow, headers={"content-type": "application/json"}, body_b64=base64.b64encode(b'{"buy":"now"}').decode())
    assert rig.counters.snapshot()["redundancy_mismatches"] == 1
    flags = mismatch_flags(other.device_uid)
    assert len(flags) == 1
    assert effective_status_of(other) == "active"


def effective_status_of(session):
    return store.effective_status(store.get_device(session.device_uid))


def test_a_shadow_never_spends_the_caller_quota_and_is_not_cross_checked_again(rig):
    SampleAll(rig)
    first, second = rig.connect(), rig.connect()
    uid = rig.enqueue()
    rig.tick()
    holder = first if uid in sent_to(first) else second
    other = second if holder is first else first
    rig.answer(holder, uid)
    shadow = list(get_table("node_requests").find(shadow_of=uid))[0]["uid"]
    rig.clock.now += 100
    rig.tick()
    rig.answer(other, shadow)
    assert get_table("node_requests").count(shadow_of={"!=": ""}) == 1


def test_sampling_is_deterministic_and_close_to_the_requested_rate():
    uids = [generate_uid() for _ in range(10000)]
    chosen = [uid for uid in uids if NodeDispatcher.sampled(uid, 5)]
    assert chosen == [uid for uid in uids if NodeDispatcher.sampled(uid, 5)]
    assert 350 <= len(chosen) <= 650
    assert not any(NodeDispatcher.sampled(uid, 0) for uid in uids[:500])
    assert all(NodeDispatcher.sampled(uid, 100) for uid in uids[:500])


def test_the_percent_is_clamped_to_the_spec_ceiling(local_db):
    from devplacepy.database import set_setting

    set_setting("nodes_redundancy_percent", "99")
    try:
        assert NodeDispatcher.redundancy_percent() == LIMITS.redundancy_max_percent
        set_setting("nodes_redundancy_percent", "-4")
        assert NodeDispatcher.redundancy_percent() == 0
    finally:
        set_setting("nodes_redundancy_percent", str(LIMITS.redundancy_percent))


def test_a_missing_original_is_inconclusive_not_a_flag(rig):
    session = rig.connect()
    rig.dispatcher._compare_shadow(session, type("J", (), {"shadow_of": "gone", "destination": "reddit.com"})(), 200, {}, b"x")
    assert rig.counters.snapshot()["redundancy_inconclusive"] == 1
    assert not mismatch_flags(session.device_uid)


def test_a_waiter_in_the_dispatcher_process_wakes_without_polling(rig):
    import asyncio

    from devplacepy.services.nodes.wake import HUB

    async def scenario():
        HUB.attach()
        target = DestinationPolicy().check(TARGET_URL)
        uid = rig.broker.enqueue("user", "c1", target, "any")
        store.claim_request(uid, "d1")
        loop = asyncio.get_running_loop()
        loop.call_later(0.02, lambda: rig.broker.complete(uid, 200, {}, b"ok", False))
        started = time.monotonic()
        result = await rig.broker.wait(uid, 5)
        return result, time.monotonic() - started

    try:
        result, elapsed = run_in_thread(scenario())
    finally:
        HUB.enabled = True
    assert result.ok and result.body == b"ok"
    assert elapsed < LIMITS.poll_interval_ms / 1000


def test_a_waiter_without_the_hub_still_finishes_by_polling(rig):
    import asyncio

    from devplacepy.services.nodes.wake import HUB

    async def scenario():
        target = DestinationPolicy().check(TARGET_URL)
        uid = rig.broker.enqueue("user", "c2", target, "any")
        store.claim_request(uid, "d1")
        asyncio.get_running_loop().call_later(0.02, lambda: rig.broker.complete(uid, 200, {}, b"ok", False))
        return await rig.broker.wait(uid, 5)

    HUB._loop = None
    assert run_in_thread(scenario()).ok


def test_enqueue_signals_the_dispatcher(rig):
    import asyncio

    from devplacepy.services.nodes.wake import HUB

    async def scenario():
        HUB.attach()
        target = DestinationPolicy().check(TARGET_URL)
        woke = asyncio.ensure_future(HUB.wait_queued(5))
        await asyncio.sleep(0)
        rig.broker.enqueue("user", "c3", target, "any")
        return await asyncio.wait_for(woke, 1)

    assert run_in_thread(scenario()) is True