from __future__ import annotations

import asyncio
import gc
import pickle
from collections import defaultdict
from collections.abc import Iterator

import pytest
from tlz import first

from dask.sizeof import sizeof

import distributed.profile as profile
from distributed import Nanny, Worker, wait
from distributed.protocol.serialize import Serialize
from distributed.scheduler import TaskState as SchedulerTaskState
from distributed.utils import recursive_to_dict
from distributed.utils_test import (
    NO_AMM,
    _LockedCommPool,
    assert_story,
    freeze_data_fetching,
    gen_cluster,
    inc,
    wait_for_state,
)
from distributed.worker_state_machine import (
    AcquireReplicasEvent,
    AddKeysMsg,
    ComputeTaskEvent,
    Execute,
    ExecuteFailureEvent,
    ExecuteSuccessEvent,
    FreeKeysEvent,
    GatherDep,
    GatherDepFailureEvent,
    GatherDepSuccessEvent,
    Instruction,
    InvalidTaskState,
    InvalidTransition,
    PauseEvent,
    RecommendationsConflict,
    RefreshWhoHasEvent,
    ReleaseWorkerDataMsg,
    RemoveReplicasEvent,
    RescheduleEvent,
    RescheduleMsg,
    SecedeEvent,
    SerializedTask,
    StateMachineEvent,
    TaskErredMsg,
    TaskState,
    TransitionCounterMaxExceeded,
    UnpauseEvent,
    UpdateDataEvent,
    merge_recs_instructions,
)


def test_instruction_match():
    i = ReleaseWorkerDataMsg(key="x", stimulus_id="s1")
    assert i == ReleaseWorkerDataMsg(key="x", stimulus_id="s1")
    assert i != ReleaseWorkerDataMsg(key="y", stimulus_id="s1")
    assert i != ReleaseWorkerDataMsg(key="x", stimulus_id="s2")
    assert i != RescheduleMsg(key="x", stimulus_id="s1")

    assert i == ReleaseWorkerDataMsg.match(key="x")
    assert i == ReleaseWorkerDataMsg.match(stimulus_id="s1")
    assert i != ReleaseWorkerDataMsg.match(key="y")
    assert i != ReleaseWorkerDataMsg.match(stimulus_id="s2")
    assert i != RescheduleMsg.match(key="x")


def test_TaskState_tracking(cleanup):
    gc.collect()
    x = TaskState("x")
    assert len(TaskState._instances) == 1
    assert first(TaskState._instances) == x
    del x
    assert len(TaskState._instances) == 0


def test_TaskState_get_nbytes():
    assert TaskState("x", nbytes=123).get_nbytes() == 123
    # Default to distributed.scheduler.default-data-size
    assert TaskState("y").get_nbytes() == 1024


def test_TaskState_eq():
    """Test that TaskState objects are hashable and that two identical objects compare
    as different. See comment in TaskState.__hash__ for why.
    """
    a = TaskState("x")
    b = TaskState("x")
    assert a != b
    s = {a, b}
    assert len(s) == 2


def test_TaskState__to_dict():
    """Tasks that are listed as dependencies or dependents of other tasks are dumped as
    a short repr and always appear in full directly under Worker.state.tasks.
    Uninteresting fields are omitted.
    """
    x = TaskState("x", state="memory", done=True)
    y = TaskState("y", priority=(0,), dependencies={x})
    x.dependents.add(y)
    actual = recursive_to_dict([x, y])
    assert actual == [
        {
            "key": "x",
            "state": "memory",
            "done": True,
            "dependents": ["<TaskState 'y' released>"],
        },
        {
            "key": "y",
            "state": "released",
            "dependencies": ["<TaskState 'x' memory>"],
            "priority": [0],
        },
    ]


def test_TaskState_repr():
    ts = TaskState("x")
    assert str(ts) == "<TaskState 'x' released>"
    ts.state = "cancelled"
    ts.previous = "flight"
    assert str(ts) == "<TaskState 'x' cancelled(flight)>"
    ts.state = "resumed"
    ts.next = "waiting"
    assert str(ts) == "<TaskState 'x' resumed(flight->waiting)>"


def test_WorkerState__to_dict(ws):
    ws.handle_stimulus(
        AcquireReplicasEvent(
            who_has={"x": ["127.0.0.1:1235"]}, nbytes={"x": 123}, stimulus_id="s1"
        )
    )
    ws.handle_stimulus(
        UpdateDataEvent(data={"y": object()}, report=False, stimulus_id="s2")
    )

    actual = recursive_to_dict(ws)
    # Remove timestamps
    for ev in actual["log"]:
        del ev[-1]
    for stim in actual["stimulus_log"]:
        del stim["handled"]

    expect = {
        "address": "127.0.0.1:1",
        "busy_workers": [],
        "constrained": [],
        "data": {"y": None},
        "data_needed": {},
        "executing": [],
        "has_what": {"127.0.0.1:1235": ["x"]},
        "in_flight_tasks": ["x"],
        "in_flight_workers": {"127.0.0.1:1235": ["x"]},
        "log": [
            ["x", "ensure-task-exists", "released", "s1"],
            ["x", "released", "fetch", "fetch", {}, "s1"],
            ["gather-dependencies", "127.0.0.1:1235", ["x"], "s1"],
            ["x", "fetch", "flight", "flight", {}, "s1"],
            ["y", "put-in-memory", "s2"],
            ["y", "receive-from-scatter", "s2"],
        ],
        "long_running": [],
        "missing_dep_flight": [],
        "nthreads": 1,
        "ready": [],
        "running": True,
        "stimulus_log": [
            {
                "cls": "AcquireReplicasEvent",
                "stimulus_id": "s1",
                "who_has": {"x": ["127.0.0.1:1235"]},
                "nbytes": {"x": 123},
            },
            {
                "cls": "UpdateDataEvent",
                "data": {"y": None},
                "report": False,
                "stimulus_id": "s2",
            },
        ],
        "tasks": {
            "x": {
                "coming_from": "127.0.0.1:1235",
                "key": "x",
                "nbytes": 123,
                "priority": [1],
                "state": "flight",
                "who_has": ["127.0.0.1:1235"],
            },
            "y": {
                "key": "y",
                "nbytes": sizeof(object()),
                "state": "memory",
            },
        },
        "transition_counter": 2,
    }
    assert actual == expect


def test_WorkerState_pickle(ws):
    """Test pickle round-trip.

    Big caveat
    ----------
    WorkerState, on its own, can be serialized with pickle; it doesn't need cloudpickle.
    A WorkerState extracted from a Worker might, as data contents may only be
    serializable with cloudpickle. Some objects created externally and not designed
    for network transfer - namely, the SpillBuffer - may not be serializable at all.
    """
    ws.handle_stimulus(
        AcquireReplicasEvent(
            who_has={"x": ["127.0.0.1:1235"]}, nbytes={"x": 123}, stimulus_id="s1"
        )
    )
    ws.handle_stimulus(UpdateDataEvent(data={"y": 123}, report=False, stimulus_id="s"))
    ws2 = pickle.loads(pickle.dumps(ws))
    assert ws2.tasks.keys() == {"x", "y"}
    assert ws2.data == {"y": 123}


@pytest.mark.parametrize(
    "cls,kwargs",
    [
        (
            InvalidTransition,
            dict(key="x", start="released", finish="waiting", story=[]),
        ),
        (
            TransitionCounterMaxExceeded,
            dict(key="x", start="released", finish="waiting", story=[]),
        ),
        (InvalidTaskState, dict(key="x", state="released", story=[])),
    ],
)
@pytest.mark.parametrize("positional", [False, True])
def test_pickle_exceptions(cls, kwargs, positional):
    if positional:
        e = cls(*kwargs.values())
    else:
        e = cls(**kwargs)
    e2 = pickle.loads(pickle.dumps(e))
    assert type(e2) is type(e)
    for k, v in kwargs.items():
        assert getattr(e2, k) == v


def traverse_subclasses(cls: type) -> Iterator[type]:
    yield cls
    for subcls in cls.__subclasses__():
        yield from traverse_subclasses(subcls)


@pytest.mark.parametrize(
    "cls",
    [
        *traverse_subclasses(Instruction),
        *traverse_subclasses(StateMachineEvent),
    ],
)
def test_slots(cls):
    params = [
        k
        for k in dir(cls)
        if not k.startswith("_")
        and k not in ("op", "handled")
        and not callable(getattr(cls, k))
    ]
    inst = cls(**dict.fromkeys(params))
    assert not hasattr(inst, "__dict__")


def test_sendmsg_to_dict():
    # Arbitrary sample class
    smsg = ReleaseWorkerDataMsg(key="x", stimulus_id="test")
    assert smsg.to_dict() == {
        "op": "release-worker-data",
        "key": "x",
        "stimulus_id": "test",
    }


def test_merge_recs_instructions():
    x = TaskState("x")
    y = TaskState("y")
    instr1 = RescheduleMsg(key="foo", stimulus_id="test")
    instr2 = RescheduleMsg(key="bar", stimulus_id="test")
    assert merge_recs_instructions(
        ({x: "memory"}, [instr1]),
        ({y: "released"}, [instr2]),
    ) == (
        {x: "memory", y: "released"},
        [instr1, instr2],
    )

    # Identical recommendations are silently ignored; incompatible ones raise
    assert merge_recs_instructions(({x: "memory"}, []), ({x: "memory"}, [])) == (
        {x: "memory"},
        [],
    )
    with pytest.raises(RecommendationsConflict):
        merge_recs_instructions(({x: "memory"}, []), ({x: "released"}, []))


def test_event_to_dict_with_annotations():
    """Test recursive_to_dict(ev), where ev is a subclass of StateMachineEvent that
    defines its own annotations
    """
    ev = RescheduleEvent(stimulus_id="test", key="x")
    ev2 = ev.to_loggable(handled=11.22)
    assert ev2 == ev
    d = recursive_to_dict(ev2)
    assert d == {
        "cls": "RescheduleEvent",
        "stimulus_id": "test",
        "handled": 11.22,
        "key": "x",
    }
    ev3 = StateMachineEvent.from_dict(d)
    assert ev3 == ev


def test_event_to_dict_without_annotations():
    """Test recursive_to_dict(ev), where ev is a subclass of StateMachineEvent that
    does not define its own annotations
    """
    ev = PauseEvent(stimulus_id="test")
    ev2 = ev.to_loggable(handled=11.22)
    assert ev2 == ev
    d = recursive_to_dict(ev2)
    assert d == {
        "cls": "PauseEvent",
        "stimulus_id": "test",
        "handled": 11.22,
    }
    ev3 = StateMachineEvent.from_dict(d)
    assert ev3 == ev


def test_computetask_to_dict():
    """The potentially very large ComputeTaskEvent.run_spec is not stored in the log"""
    ev = ComputeTaskEvent(
        key="x",
        who_has={"y": ["w1"]},
        nbytes={"y": 123},
        priority=(0,),
        duration=123.45,
        run_spec=None,
        resource_restrictions={},
        actor=False,
        annotations={},
        stimulus_id="test",
        function=b"blob",
        args=b"blob",
        kwargs=None,
    )
    assert ev.run_spec == SerializedTask(function=b"blob", args=b"blob")
    ev2 = ev.to_loggable(handled=11.22)
    assert ev2.handled == 11.22
    assert ev2.run_spec == SerializedTask(task=None)
    assert ev.run_spec == SerializedTask(function=b"blob", args=b"blob")
    d = recursive_to_dict(ev2)
    assert d == {
        "cls": "ComputeTaskEvent",
        "key": "x",
        "who_has": {"y": ["w1"]},
        "nbytes": {"y": 123},
        "priority": [0],
        "run_spec": [None, None, None, None],
        "duration": 123.45,
        "resource_restrictions": {},
        "actor": False,
        "annotations": {},
        "stimulus_id": "test",
        "handled": 11.22,
        "function": None,
        "args": None,
        "kwargs": None,
    }
    ev3 = StateMachineEvent.from_dict(d)
    assert isinstance(ev3, ComputeTaskEvent)
    assert ev3.run_spec == SerializedTask(task=None)
    assert ev3.priority == (0,)  # List is automatically converted back to tuple


def test_computetask_dummy():
    ev = ComputeTaskEvent.dummy("x", stimulus_id="s")
    assert ev == ComputeTaskEvent(
        key="x",
        who_has={},
        nbytes={},
        priority=(0,),
        duration=1.0,
        run_spec=None,
        resource_restrictions={},
        actor=False,
        annotations={},
        stimulus_id="s",
        function=None,
        args=None,
        kwargs=None,
    )

    # nbytes is generated from who_has if omitted
    ev2 = ComputeTaskEvent.dummy("x", who_has={"y": "127.0.0.1:2"}, stimulus_id="s")
    assert ev2.nbytes == {"y": 1}


def test_updatedata_to_dict():
    """The potentially very large UpdateDataEvent.data is not stored in the log"""
    ev = UpdateDataEvent(
        data={"x": "foo", "y": "bar"},
        report=True,
        stimulus_id="test",
    )
    ev2 = ev.to_loggable(handled=11.22)
    assert ev2.handled == 11.22
    assert ev2.data == {"x": None, "y": None}
    d = recursive_to_dict(ev2)
    assert d == {
        "cls": "UpdateDataEvent",
        "data": {"x": None, "y": None},
        "report": True,
        "stimulus_id": "test",
        "handled": 11.22,
    }
    ev3 = StateMachineEvent.from_dict(d)
    assert isinstance(ev3, UpdateDataEvent)
    assert ev3.data == {"x": None, "y": None}


def test_executesuccess_to_dict():
    """The potentially very large ExecuteSuccessEvent.value is not stored in the log"""
    ev = ExecuteSuccessEvent(
        stimulus_id="test",
        key="x",
        value=123,
        start=123.4,
        stop=456.7,
        nbytes=890,
        type=int,
    )
    ev2 = ev.to_loggable(handled=11.22)
    assert ev2.value is None
    assert ev.value == 123
    d = recursive_to_dict(ev2)
    assert d == {
        "cls": "ExecuteSuccessEvent",
        "stimulus_id": "test",
        "handled": 11.22,
        "key": "x",
        "value": None,
        "nbytes": 890,
        "start": 123.4,
        "stop": 456.7,
        "type": "<class 'int'>",
    }
    ev3 = StateMachineEvent.from_dict(d)
    assert isinstance(ev3, ExecuteSuccessEvent)
    assert ev3.stimulus_id == "test"
    assert ev3.handled == 11.22
    assert ev3.key == "x"
    assert ev3.value is None
    assert ev3.start == 123.4
    assert ev3.stop == 456.7
    assert ev3.nbytes == 890
    assert ev3.type is None


def test_executesuccess_dummy():
    ev = ExecuteSuccessEvent.dummy("x", stimulus_id="s")
    assert ev == ExecuteSuccessEvent(
        key="x",
        value=None,
        start=0.0,
        stop=1.0,
        nbytes=1,
        type=None,
        stimulus_id="s",
    )

    ev2 = ExecuteSuccessEvent.dummy("x", 123, stimulus_id="s")
    assert ev2.value == 123


def test_executefailure_to_dict():
    ev = ExecuteFailureEvent(
        stimulus_id="test",
        key="x",
        start=123.4,
        stop=456.7,
        exception=Serialize(ValueError("foo")),
        traceback=Serialize("lose me"),
        exception_text="exc text",
        traceback_text="tb text",
    )
    ev2 = ev.to_loggable(handled=11.22)
    assert ev2 == ev
    d = recursive_to_dict(ev2)
    assert d == {
        "cls": "ExecuteFailureEvent",
        "stimulus_id": "test",
        "handled": 11.22,
        "key": "x",
        "start": 123.4,
        "stop": 456.7,
        "exception": "<Serialize: foo>",
        "traceback": "<Serialize: lose me>",
        "exception_text": "exc text",
        "traceback_text": "tb text",
    }
    ev3 = StateMachineEvent.from_dict(d)
    assert isinstance(ev3, ExecuteFailureEvent)
    assert ev3.stimulus_id == "test"
    assert ev3.handled == 11.22
    assert ev3.key == "x"
    assert ev3.start == 123.4
    assert ev3.stop == 456.7
    assert isinstance(ev3.exception, Serialize)
    assert isinstance(ev3.exception.data, Exception)
    assert ev3.traceback is None
    assert ev3.exception_text == "exc text"
    assert ev3.traceback_text == "tb text"


def test_executefailure_dummy():
    ev = ExecuteFailureEvent.dummy("x", stimulus_id="s")
    assert ev == ExecuteFailureEvent(
        key="x",
        start=None,
        stop=None,
        exception=Serialize(None),
        traceback=None,
        exception_text="",
        traceback_text="",
        stimulus_id="s",
    )


@gen_cluster(client=True)
async def test_fetch_to_compute(c, s, a, b):
    with freeze_data_fetching(b):
        f1 = c.submit(inc, 1, workers=[a.address], key="f1", allow_other_workers=True)
        f2 = c.submit(inc, f1, workers=[b.address], key="f2")
        await wait_for_state(f1.key, "fetch", b)
        await a.close()

    await f2

    assert_story(
        b.state.log,
        # FIXME: This log should be replaced with a StateMachineEvent log
        [
            (f2.key, "compute-task", "released"),
            # This is a "please fetch" request. We don't have anything like
            # this, yet. We don't see the request-dep signal in here because we
            # do not wait for the key to be actually scheduled
            (f1.key, "ensure-task-exists", "released"),
            # After the worker failed, we're instructed to forget f2 before
            # something new comes in
            ("free-keys", (f2.key,)),
            (f1.key, "compute-task", "released"),
            (f1.key, "put-in-memory"),
            (f2.key, "compute-task", "released"),
        ],
    )


@gen_cluster(client=True)
async def test_fetch_via_amm_to_compute(c, s, a, b):
    with freeze_data_fetching(b):
        f1 = c.submit(inc, 1, workers=[a.address], key="f1", allow_other_workers=True)
        await f1
        s.request_acquire_replicas(b.address, [f1.key], stimulus_id="test")
        await wait_for_state(f1.key, "fetch", b)
        await a.close()

    await f1

    assert_story(
        b.state.log,
        # FIXME: This log should be replaced with a StateMachineEvent log
        [
            (f1.key, "ensure-task-exists", "released"),
            (f1.key, "released", "fetch", "fetch", {}),
            (f1.key, "compute-task", "fetch"),
            (f1.key, "put-in-memory"),
        ],
    )


@pytest.mark.parametrize("as_deps", [False, True])
@gen_cluster(
    client=True,
    nthreads=[("", 1)] * 3,
    config=NO_AMM,
)
async def test_lose_replica_during_fetch(c, s, w1, w2, w3, as_deps):
    """
    as_deps=True
        0. task x is a dependency of y1 and y2
        1. scheduler calls handle_compute("y1", who_has={"x": [w2, w3]}) on w1
        2. x transitions released -> fetch
        3. the network stack is busy, so x does not transition to flight yet.
        4. scheduler calls handle_compute("y2", who_has={"x": [w3]}) on w1
        5. when x finally reaches the top of the data_needed heap, w1 will not try
           contacting w2

    as_deps=False
        1. scheduler calls handle_acquire_replicas(who_has={"x": [w2, w3]}) on w1
        2. x transitions released -> fetch
        3. the network stack is busy, so x does not transition to flight yet.
        4. scheduler calls handle_acquire_replicas(who_has={"x": [w3]}) on w1
        5. when x finally reaches the top of the data_needed heap, w1 will not try
           contacting w2
    """
    x = (await c.scatter({"x": 1}, workers=[w2.address, w3.address], broadcast=True))[
        "x"
    ]

    # Make sure find_missing is not involved
    w1.periodic_callbacks["find-missing"].stop()

    with freeze_data_fetching(w1, jump_start=True):
        if as_deps:
            y1 = c.submit(inc, x, key="y1", workers=[w1.address])
        else:
            s.request_acquire_replicas(w1.address, ["x"], stimulus_id="test")

        await wait_for_state("x", "fetch", w1)
        assert w1.state.tasks["x"].who_has == {w2.address, w3.address}

        assert len(s.tasks["x"].who_has) == 2
        await w2.close()
        while len(s.tasks["x"].who_has) > 1:
            await asyncio.sleep(0.01)

        if as_deps:
            y2 = c.submit(inc, x, key="y2", workers=[w1.address])
        else:
            s.request_acquire_replicas(w1.address, ["x"], stimulus_id="test")

        while w1.state.tasks["x"].who_has != {w3.address}:
            await asyncio.sleep(0.01)

    await wait_for_state("x", "memory", w1)
    assert_story(
        w1.state.story("request-dep"),
        [("request-dep", w3.address, {"x"})],
        # This tests that there has been no attempt to contact w2.
        # If the assumption being tested breaks, this will fail 50% of the times.
        strict=True,
    )


@gen_cluster(client=True, nthreads=[("", 1)] * 2)
async def test_fetch_to_missing_on_busy(c, s, a, b):
    """
    1. task x is a dependency of y
    2. scheduler calls handle_compute("y", who_has={"x": [b]}) on a
    3. x transitions released -> fetch -> flight; a connects to b
    4. b responds it's busy. x transitions flight -> fetch
    5. The busy state triggers an RPC call to Scheduler.who_has
    6. the scheduler responds {"x": []}, because w1 in the meantime has lost the key.
    7. x is transitioned fetch -> missing
    """
    # Note: submit and scatter are different. If you lose all workers holding the
    # replicas of a scattered key, the scheduler forgets the task, which in turn would
    # trigger a free-keys response to request-refresh-who-has.
    x = c.submit(inc, 1, key="x", workers=[b.address])
    await x

    b.transfer_outgoing_count_limit = 0
    # Crucially, unlike with `c.submit(inc, x, workers=[a.address])`, the scheduler
    # doesn't keep track of acquire-replicas requests, so it won't proactively inform a
    # when we call remove_worker later on
    s.request_acquire_replicas(a.address, ["x"], stimulus_id="test")

    # state will flip-flop between fetch and flight every 150ms, which is the retry
    # period for busy workers.
    await wait_for_state("x", "fetch", a)
    assert b.address in a.state.busy_workers

    # Sever connection between b and s, but not between b and a.
    # If a tries fetching from b after this, b will keep responding {status: busy}.
    b.periodic_callbacks["heartbeat"].stop()
    await s.remove_worker(b.address, close=False, stimulus_id="test")

    await wait_for_state("x", "missing", a)

    assert_story(
        a.state.story("x"),
        [
            ("x", "ensure-task-exists", "released"),
            ("x", "released", "fetch", "fetch", {}),
            ("gather-dependencies", b.address, {"x"}),
            ("x", "fetch", "flight", "flight", {}),
            ("request-dep", b.address, {"x"}),
            ("busy-gather", b.address, {"x"}),
            ("x", "flight", "fetch", "fetch", {}),
            ("x", "fetch", "missing", "missing", {}),
        ],
        # There may be a round of find_missing() after this.
        # Due to timings, there also may be multiple attempts to connect from a to b.
        strict=False,
    )


def test_new_replica_while_all_workers_in_flight(ws):
    """A task is stuck in 'fetch' state because all workers that hold a replica are in
    flight. While in this state, a new replica appears on a different worker and the
    scheduler informs the waiting worker through a new acquire-replicas or
    compute-task op.

    In real life, this will typically happen when the Active Memory Manager replicates a
    key to multiple workers and some workers are much faster than others to acquire it,
    due to unrelated tasks being in flight, so 2 seconds later the AMM reiterates the
    request, passing a larger who_has.

    Test that, when this happens, the task is immediately acquired from the new worker,
    without waiting for the original replica holders to get out of flight.
    """
    ws2 = "127.0.0.1:2"
    ws3 = "127.0.0.1:3"
    instructions = ws.handle_stimulus(
        AcquireReplicasEvent(
            who_has={"x": [ws2]},
            nbytes={"x": 1},
            stimulus_id="s1",
        ),
        AcquireReplicasEvent(
            who_has={"y": [ws2]},
            nbytes={"y": 1},
            stimulus_id="s2",
        ),
        AcquireReplicasEvent(
            who_has={"y": [ws2, ws3]},
            nbytes={"y": 1},
            stimulus_id="s3",
        ),
    )
    assert instructions == [
        GatherDep(
            worker=ws2,
            to_gather={"x"},
            total_nbytes=1,
            stimulus_id="s1",
        ),
        GatherDep(
            worker=ws3,
            to_gather={"y"},
            total_nbytes=1,
            stimulus_id="s3",
        ),
    ]
    assert ws.tasks["x"].state == "flight"
    assert ws.tasks["y"].state == "flight"


@gen_cluster(client=True)
async def test_cancelled_while_in_flight(c, s, a, b):
    event = asyncio.Event()
    a.rpc = _LockedCommPool(a.rpc, write_event=event)

    x = c.submit(inc, 1, key="x", workers=[b.address])
    y = c.submit(inc, x, key="y", workers=[a.address])
    await wait_for_state("x", "flight", a)
    y.release()
    await wait_for_state("x", "cancelled", a)

    # Let the comm from b to a return the result
    event.set()
    # upon reception, x transitions cancelled->forgotten
    while a.state.tasks:
        await asyncio.sleep(0.01)


@gen_cluster(client=True, config=NO_AMM)
async def test_in_memory_while_in_flight(c, s, a, b):
    """
    1. A client scatters x to a
    2. The scheduler does not know about scattered keys until the three-way round-trip
       between client, worker, and scheduler has been completed (see Scheduler.scatter)
    3. In the middle of that handshake, a client (not necessarily the same client) calls
       ``{op: compute-task, key: x}`` on b and then
       ``{op: compute-task, key: y, who_has: {x: [b]}`` on a, which triggers a
       gather_dep call to copy x key from b to a.
    4. while x is in flight from b to a, the scatter finishes, which triggers
       update_data, which in turn transitions x from flight to memory.
    5. later on, gather_dep finishes, but the key is already in memory.
    """
    event = asyncio.Event()
    a.rpc = _LockedCommPool(a.rpc, write_event=event)

    x = c.submit(inc, 1, key="x", workers=[b.address])
    y = c.submit(inc, x, key="y", workers=[a.address])
    await wait_for_state("x", "flight", a)
    a.update_data({"x": 3})
    await wait_for_state("x", "memory", a)

    # Let the comm from b to a return the result
    event.set()
    assert await y == 4  # Data in flight from b has been discarded


@gen_cluster(client=True)
async def test_forget_data_needed(c, s, a, b):
    """
    1. A task transitions to fetch and is added to data_needed
    2. _ensure_communicating runs, but the network is saturated so the task is not
       popped from data_needed
    3. Task is forgotten
    4. Task is recreated from scratch and transitioned to fetch again
    5. BUG: at the moment of writing this test, adding to data_needed silently did
       nothing, because it still contained the forgotten task, which is a different
       TaskState instance which will be no longer updated.
    6. _ensure_communicating runs. It pops the forgotten task and discards it.
    7. We now have a task stuck in fetch state.
    """
    x = c.submit(inc, 1, key="x", workers=[a.address])
    with freeze_data_fetching(b):
        y = c.submit(inc, x, key="y", workers=[b.address])
        await wait_for_state("x", "fetch", b)
        x.release()
        y.release()
        while s.tasks or a.state.tasks or b.state.tasks:
            await asyncio.sleep(0.01)

    x = c.submit(inc, 2, key="x", workers=[a.address])
    y = c.submit(inc, x, key="y", workers=[b.address])
    assert await y == 4


@gen_cluster(client=True, nthreads=[("", 1)] * 3)
async def test_missing_handle_compute_dependency(c, s, w1, w2, w3):
    """Test that it is OK for a dependency to be in state missing if a dependent is
    asked to be computed
    """
    w3.periodic_callbacks["find-missing"].stop()

    f1 = c.submit(inc, 1, key="f1", workers=[w1.address])
    f2 = c.submit(inc, 2, key="f2", workers=[w1.address])
    await wait_for_state(f1.key, "memory", w1)

    w3.handle_stimulus(
        AcquireReplicasEvent(
            who_has={f1.key: [w2.address]}, nbytes={f1.key: 1}, stimulus_id="acquire"
        )
    )
    await wait_for_state(f1.key, "missing", w3)

    f3 = c.submit(sum, [f1, f2], key="f3", workers=[w3.address])

    await f3


@gen_cluster(client=True, nthreads=[("", 1)] * 3)
async def test_missing_to_waiting(c, s, w1, w2, w3):
    w3.periodic_callbacks["find-missing"].stop()

    f1 = c.submit(inc, 1, key="f1", workers=[w1.address], allow_other_workers=True)
    await wait_for_state(f1.key, "memory", w1)

    w3.handle_stimulus(
        AcquireReplicasEvent(
            who_has={f1.key: [w2.address]}, nbytes={f1.key: 1}, stimulus_id="acquire"
        )
    )
    await wait_for_state(f1.key, "missing", w3)

    await w2.close()
    await w1.close()

    await f1


@gen_cluster(client=True, Worker=Nanny)
async def test_task_state_instance_are_garbage_collected(c, s, a, b):
    futs = c.map(inc, range(10))
    red = c.submit(sum, futs)
    f1 = c.submit(inc, red, pure=False)
    f2 = c.submit(inc, red, pure=False)

    async def check(dask_worker):
        while dask_worker.tasks:
            await asyncio.sleep(0.01)
        with profile.lock:
            gc.collect()
        assert not TaskState._instances

    await c.gather([f2, f1])
    del futs, red, f1, f2
    await c.run(check)

    async def check(dask_scheduler):
        while dask_scheduler.tasks:
            await asyncio.sleep(0.01)
        with profile.lock:
            gc.collect()
        assert not SchedulerTaskState._instances

    await c.run_on_scheduler(check)


@gen_cluster(client=True, nthreads=[("", 1)] * 3)
async def test_fetch_to_missing_on_refresh_who_has(c, s, w1, w2, w3):
    """
    1. Two tasks, x and y, are only available on a busy worker.
       The worker sends request-refresh-who-has to the scheduler.
    2. The scheduler responds that x has become missing, while y has gained an
       additional replica
    3. The handler for RefreshWhoHasEvent empties x.who_has and recommends a transition
       to missing.
    4. Before the recommendation can be implemented, the same event invokes
       _ensure_communicating to let y to transition to flight. This in turn pops x from
       data_needed - but x has an empty who_has, which is an exceptional situation.
    5. The transition fetch->missing is executed, but x is no longer in
       data_needed - another exceptional situation.
    """
    x = c.submit(inc, 1, key="x", workers=[w1.address])
    y = c.submit(inc, 2, key="y", workers=[w1.address])
    await wait([x, y])
    w1.transfer_outgoing_count_limit = 0
    s.request_acquire_replicas(w3.address, ["x", "y"], stimulus_id="test1")

    # The tasks will now flip-flop between fetch and flight every 150ms
    # (see Worker.retry_busy_worker_later)
    await wait_for_state("x", "fetch", w3)
    await wait_for_state("y", "fetch", w3)
    assert w1.address in w3.state.busy_workers
    # w3 sent {op: request-refresh-who-has, keys: [x, y]}
    # There also may have been enough time for a refresh-who-has message to come back,
    # which reiterated what w3 already knew:
    # {op: refresh-who-has, who_has={x: [w1.address], y: [w1.address]}}

    # Let's instead simulate that, while request-refresh-who-has was in transit,
    # w2 gained a replica of y and w1 closed down.
    # When request-refresh-who-has lands, the scheduler will respond:
    # {op: refresh-who-has, who_has={x: [], y: [w2.address]}}
    w3.handle_stimulus(
        RefreshWhoHasEvent(who_has={"x": [], "y": [w2.address]}, stimulus_id="test2")
    )
    assert w3.state.tasks["x"].state == "missing"
    assert w3.state.tasks["y"].state == "flight"
    assert w3.state.tasks["y"].who_has == {w2.address}


@gen_cluster(client=True, nthreads=[("", 1)])
async def test_fetch_to_missing_on_network_failure(c, s, a):
    """
    1. Two tasks, x and y, are respectively in flight and fetch state from the same
       worker, which holds the only replica of both.
    2. gather_dep for x returns GatherDepNetworkFailureEvent
    3. The event empties has_what, x.who_has, and y.who_has.
    4. The same event invokes _ensure_communicating, which pops y from data_needed
       - but y has an empty who_has, which is an exceptional situation.
       _ensure_communicating recommends a transition to missing for x.
    5. The fetch->missing transition is executed, but y is no longer in data_needed -
       another exceptional situation.
    """
    block_get_data = asyncio.Event()

    class BlockedBreakingWorker(Worker):
        async def get_data(self, comm, *args, **kwargs):
            await block_get_data.wait()
            raise OSError("fake error")

    async with BlockedBreakingWorker(s.address) as b:
        x = c.submit(inc, 1, key="x", workers=[b.address])
        y = c.submit(inc, 2, key="y", workers=[b.address])
        await wait([x, y])
        s.request_acquire_replicas(a.address, ["x"], stimulus_id="test_x")
        await wait_for_state("x", "flight", a)
        s.request_acquire_replicas(a.address, ["y"], stimulus_id="test_y")
        await wait_for_state("y", "fetch", a)

        block_get_data.set()

        await wait_for_state("x", "missing", a)
        await wait_for_state("y", "missing", a)


@gen_cluster()
async def test_deprecated_worker_attributes(s, a, b):
    n = a.state.generation
    msg = (
        "The `Worker.generation` attribute has been moved to "
        "`Worker.state.generation`"
    )
    with pytest.warns(FutureWarning, match=msg):
        assert a.generation == n
    with pytest.warns(FutureWarning, match=msg):
        a.generation -= 1
        assert a.generation == n - 1
    assert a.state.generation == n - 1

    # Old and new names differ
    msg = (
        "The `Worker.in_flight_tasks` attribute has been moved to "
        "`Worker.state.in_flight_tasks_count`"
    )
    with pytest.warns(FutureWarning, match=msg):
        assert a.in_flight_tasks == 0

    with pytest.warns(FutureWarning, match="attribute has been removed"):
        assert a.data_needed == set()
    with pytest.warns(FutureWarning, match="attribute has been removed"):
        assert a.waiting_for_data_count == 0


@pytest.mark.parametrize("n_remote_workers", [1, 2])
@pytest.mark.parametrize(
    "nbytes,n_in_flight_per_worker",
    [
        (int(10e6), 3),
        (int(20e6), 2),
        (int(30e6), 1),
        (int(60e6), 1),
    ],
)
def test_aggregate_gather_deps(ws, nbytes, n_in_flight_per_worker, n_remote_workers):
    ws.transfer_message_bytes_limit = int(50e6)
    wss = [f"127.0.0.1:{2 + i}" for i in range(n_remote_workers)]
    who_has = {f"x{i}": [wss[i // 3]] for i in range(3 * n_remote_workers)}
    instructions = ws.handle_stimulus(
        AcquireReplicasEvent(
            who_has=who_has,
            nbytes={task: nbytes for task in who_has.keys()},
            stimulus_id="s1",
        )
    )
    assert instructions == [
        GatherDep.match(worker=remote, stimulus_id="s1") for remote in wss
    ]
    assert all(
        len(instruction.to_gather) == n_in_flight_per_worker
        for instruction in instructions
    )
    assert len(ws.in_flight_tasks) == n_in_flight_per_worker * n_remote_workers
    assert (
        ws.transfer_incoming_bytes == nbytes * n_in_flight_per_worker * n_remote_workers
    )
    assert ws.transfer_incoming_count == n_remote_workers
    assert ws.transfer_incoming_count_total == n_remote_workers


def test_gather_priority(ws):
    """Test that tasks are fetched in the following order:

    1. by task priority
    2. in case of tie, from local workers first
    3. in case of tie, from the worker with the most tasks queued
    4. in case of tie, from a random worker (which is actually deterministic).
    """
    ws.transfer_incoming_count_limit = 4

    instructions = ws.handle_stimulus(
        PauseEvent(stimulus_id="pause"),
        # Note: tasks fetched by acquire-replicas always have priority=(1, )
        AcquireReplicasEvent(
            who_has={
                # Remote + local
                "x1": ["127.0.0.2:1", "127.0.0.1:2"],
                # Remote. After getting x11 from .1, .2  will have less tasks than .3
                "x2": ["127.0.0.2:1"],
                "x3": ["127.0.0.3:1"],
                "x4": ["127.0.0.3:1"],
                # It will be a random choice between .2, .4, .5, .6, and .7
                "x5": ["127.0.0.4:1"],
                "x6": ["127.0.0.5:1"],
                "x7": ["127.0.0.6:1"],
                # This will be fetched first because it's on the same worker as y
                "x8": ["127.0.0.7:1"],
            },
            # Substantial nbytes prevents transfer_incoming_count_limit to be
            # overridden by transfer_incoming_bytes_throttle_threshold,
            # but it's less than transfer_message_bytes_limit
            nbytes={f"x{i}": 4 * 2**20 for i in range(1, 9)},
            stimulus_id="compute1",
        ),
        # A higher-priority task, even if scheduled later, is fetched first
        ComputeTaskEvent.dummy(
            key="z",
            who_has={"y": ["127.0.0.7:1"]},
            priority=(0,),
            stimulus_id="compute2",
        ),
        UnpauseEvent(stimulus_id="unpause"),
    )

    assert instructions == [
        # Highest-priority task first. Lower priority tasks from the same worker are
        # shoved into the same instruction (up to 50MB worth)
        GatherDep(
            stimulus_id="unpause",
            worker="127.0.0.7:1",
            to_gather={"y", "x8"},
            total_nbytes=1 + 4 * 2**20,
        ),
        # Followed by local workers
        GatherDep(
            stimulus_id="unpause",
            worker="127.0.0.1:2",
            to_gather={"x1"},
            total_nbytes=4 * 2**20,
        ),
        # Followed by remote workers with the most tasks
        GatherDep(
            stimulus_id="unpause",
            worker="127.0.0.3:1",
            to_gather={"x3", "x4"},
            total_nbytes=8 * 2**20,
        ),
        # Followed by other remote workers, randomly.
        # Determinism is guaranteed by a statically-seeded random number generator.
        # FIXME It would have not been deterministic if we instead of multiple keys we
        #       had used a single key with multiple workers, because sets
        #       (like TaskState.who_has) change order at every interpreter restart.
        GatherDep(
            stimulus_id="unpause",
            worker="127.0.0.4:1",
            to_gather={"x5"},
            total_nbytes=4 * 2**20,
        ),
    ]
    expected_bytes = 1 + 4 * 2**20 + 4 * 2**20 + 8 * 2**20 + 4 * 2**20
    assert ws.transfer_incoming_bytes == expected_bytes
    assert ws.transfer_incoming_count == 4
    assert ws.transfer_incoming_count_total == 4


@pytest.mark.parametrize("state", ["executing", "long-running"])
def test_task_acquires_resources(ws, state):
    ws.available_resources = {"R": 1}
    ws.total_resources = {"R": 1}

    ws.handle_stimulus(
        ComputeTaskEvent.dummy(
            key="x", resource_restrictions={"R": 1}, stimulus_id="compute"
        )
    )
    if state == "long-running":
        ws.handle_stimulus(
            SecedeEvent(key="x", compute_duration=1.0, stimulus_id="secede")
        )
    assert ws.tasks["x"].state == state
    assert ws.available_resources == {"R": 0}


@pytest.mark.parametrize(
    "done_ev_cls", [ExecuteSuccessEvent, ExecuteFailureEvent, RescheduleEvent]
)
def test_task_releases_resources(ws_with_running_task, done_ev_cls):
    ws = ws_with_running_task
    assert ws.available_resources == {"R": 0}

    ws.handle_stimulus(done_ev_cls.dummy("x", stimulus_id="success"))
    assert ws.available_resources == {"R": 1}


def test_task_with_dependencies_acquires_resources(ws):
    ws.available_resources = {"R": 1}
    ws.total_resources = {"R": 1}
    ws2 = "127.0.0.1:2"
    ws.handle_stimulus(
        ComputeTaskEvent.dummy(
            "y", who_has={"x": [ws2]}, resource_restrictions={"R": 1}, stimulus_id="s1"
        )
    )
    assert ws.tasks["x"].state == "flight"
    assert ws.tasks["y"].state == "waiting"
    assert ws.available_resources == {"R": 1}

    instructions = ws.handle_stimulus(
        GatherDepSuccessEvent(
            worker=ws2, data={"x": 123}, total_nbytes=8, stimulus_id="s2"
        )
    )
    assert instructions == [
        AddKeysMsg(keys=["x"], stimulus_id="s2"),
        Execute(key="y", stimulus_id="s2"),
    ]
    assert ws.tasks["y"].state == "executing"
    assert ws.available_resources == {"R": 0}


@pytest.mark.parametrize(
    "done_ev_cls", [ExecuteSuccessEvent, ExecuteFailureEvent, RescheduleEvent]
)
def test_resumed_task_releases_resources(ws_with_running_task, done_ev_cls):
    ws = ws_with_running_task
    assert ws.available_resources == {"R": 0}
    ws2 = "127.0.0.1:2"

    ws.handle_stimulus(FreeKeysEvent("cancel", ["x"]))
    assert ws.tasks["x"].state == "cancelled"
    assert ws.available_resources == {"R": 0}

    instructions = ws.handle_stimulus(
        ComputeTaskEvent.dummy("y", who_has={"x": [ws2]}, stimulus_id="compute")
    )
    assert not instructions
    assert ws.tasks["x"].state == "resumed"
    assert ws.available_resources == {"R": 0}

    ws.handle_stimulus(done_ev_cls.dummy("x", stimulus_id="s2"))
    assert ws.available_resources == {"R": 1}


@gen_cluster()
async def test_clean_log(s, a, b):
    """Test that brand new workers start with a clean log"""
    assert not a.state.log
    assert not a.state.stimulus_log


def test_running_task_in_all_running_tasks(ws_with_running_task):
    ws = ws_with_running_task
    ws2 = "127.0.0.1:2"
    ts = ws.tasks["x"]
    assert ts in ws.all_running_tasks

    ws.handle_stimulus(FreeKeysEvent(keys=["x"], stimulus_id="s1"))
    assert ts.state == "cancelled"
    assert ts in ws.all_running_tasks

    ws.handle_stimulus(
        ComputeTaskEvent.dummy("y", who_has={"x": [ws2]}, stimulus_id="s2")
    )
    assert ts.state == "resumed"
    assert ts in ws.all_running_tasks


@pytest.mark.parametrize(
    "done_ev_cls", [ExecuteSuccessEvent, ExecuteFailureEvent, RescheduleEvent]
)
def test_done_task_not_in_all_running_tasks(ws_with_running_task, done_ev_cls):
    ws = ws_with_running_task
    ts = ws.tasks["x"]
    assert ts in ws.all_running_tasks

    ws.handle_stimulus(done_ev_cls.dummy("x", stimulus_id="s1"))
    assert ts not in ws.all_running_tasks


@pytest.mark.parametrize(
    "done_ev_cls", [ExecuteSuccessEvent, ExecuteFailureEvent, RescheduleEvent]
)
def test_done_resumed_task_not_in_all_running_tasks(ws_with_running_task, done_ev_cls):
    ws = ws_with_running_task
    ws2 = "127.0.0.1:2"

    ws.handle_stimulus(
        FreeKeysEvent(keys=["x"], stimulus_id="s1"),
        ComputeTaskEvent.dummy("y", who_has={"x": [ws2]}, stimulus_id="s2"),
        done_ev_cls.dummy("x", stimulus_id="s3"),
    )
    ts = ws.tasks["x"]
    assert ts not in ws.all_running_tasks


@pytest.mark.xfail(reason="https://github.com/dask/distributed/issues/6705")
def test_gather_dep_failure(ws):
    """Simulate a task failing to unpickle when it reaches the destination worker after
    a flight.

    See also test_worker_memory.py::test_workerstate_fail_to_pickle_flight,
    where the task instead is gathered successfully, but fails to spill.
    """
    ws2 = "127.0.0.1:2"
    instructions = ws.handle_stimulus(
        ComputeTaskEvent.dummy("y", who_has={"x": [ws2]}, stimulus_id="s1"),
        GatherDepFailureEvent.from_exception(
            Exception(), worker=ws2, total_nbytes=1, stimulus_id="s2"
        ),
    )
    assert instructions == [
        GatherDep(worker=ws2, to_gather={"x"}, total_nbytes=1, stimulus_id="s1"),
        TaskErredMsg.match(key="x", stimulus_id="s2"),
    ]
    assert ws.tasks["x"].state == "error"
    assert ws.tasks["y"].state == "waiting"  # Not ready
    assert ws.transfer_incoming_bytes == 0
    assert ws.transfer_incoming_count == 0
    assert ws.transfer_incoming_count_total == 1


def test_transfer_incoming_metrics(ws):
    assert ws.transfer_incoming_bytes == 0
    assert ws.transfer_incoming_count == 0
    assert ws.transfer_incoming_count_total == 0

    ws2 = "127.0.0.1:2"
    ws.handle_stimulus(
        ComputeTaskEvent.dummy(
            "b", who_has={"a": [ws2]}, nbytes={"a": 7}, stimulus_id="s1"
        )
    )
    assert ws.transfer_incoming_bytes == 7
    assert ws.transfer_incoming_count == 1
    assert ws.transfer_incoming_count_total == 1

    ws.handle_stimulus(
        GatherDepSuccessEvent(
            worker=ws2, data={"a": 123}, total_nbytes=7, stimulus_id="s2"
        )
    )
    assert ws.transfer_incoming_bytes == 0
    assert ws.transfer_incoming_count == 0
    assert ws.transfer_incoming_count_total == 1

    ws.handle_stimulus(
        ComputeTaskEvent.dummy(
            "e",
            who_has={"c": [ws2], "d": [ws2]},
            nbytes={"c": 11, "d": 13},
            stimulus_id="s2",
        )
    )
    assert ws.transfer_incoming_bytes == 24
    assert ws.transfer_incoming_count == 1
    assert ws.transfer_incoming_count_total == 2

    ws.handle_stimulus(
        GatherDepSuccessEvent(
            worker=ws2, data={"c": 123, "d": 234}, total_nbytes=24, stimulus_id="s3"
        )
    )
    assert ws.transfer_incoming_bytes == 0
    assert ws.transfer_incoming_count == 0
    assert ws.transfer_incoming_count_total == 2

    ws3 = "127.0.0.1:3"
    ws.handle_stimulus(
        ComputeTaskEvent.dummy(
            "h",
            who_has={"f": [ws2], "g": [ws3]},
            nbytes={"f": 17, "g": 19},
            stimulus_id="s4",
        )
    )
    assert ws.transfer_incoming_bytes == 36
    assert ws.transfer_incoming_count == 2
    assert ws.transfer_incoming_count_total == 4

    ws.handle_stimulus(
        GatherDepSuccessEvent(
            worker=ws3, data={"g": 345}, total_nbytes=19, stimulus_id="s5"
        )
    )
    assert ws.transfer_incoming_bytes == 17
    assert ws.transfer_incoming_count == 1
    assert ws.transfer_incoming_count_total == 4

    ws.handle_stimulus(
        GatherDepSuccessEvent(
            worker=ws2, data={"g": 456}, total_nbytes=17, stimulus_id="s6"
        )
    )
    assert ws.transfer_incoming_bytes == 0
    assert ws.transfer_incoming_count == 0
    assert ws.transfer_incoming_count_total == 4


def test_throttling_does_not_affect_first_transfer(ws):
    ws.transfer_incoming_count_limit = 100
    ws.transfer_incoming_bytes_limit = 100
    ws.transfer_incoming_bytes_throttle_threshold = 1
    ws.transfer_message_bytes_limit = 100
    ws2 = "127.0.0.1:2"
    ws.handle_stimulus(
        ComputeTaskEvent.dummy(
            "c",
            who_has={"a": [ws2]},
            nbytes={"a": 200},
            stimulus_id="s1",
        )
    )
    assert ws.tasks["a"].state == "flight"


def test_message_target_does_not_affect_first_transfer_on_different_worker(ws):
    ws.transfer_incoming_count_limit = 100
    ws.transfer_incoming_bytes_limit = 600
    ws.transfer_message_bytes_limit = 100
    ws.transfer_incoming_bytes_throttle_threshold = 1
    ws2 = "127.0.0.1:2"
    ws3 = "127.0.0.1:3"
    ws.handle_stimulus(
        ComputeTaskEvent.dummy(
            "c",
            who_has={"a": [ws2], "b": [ws3]},
            nbytes={"a": 200, "b": 200},
            stimulus_id="s1",
        )
    )
    assert ws.tasks["a"].state == "flight"
    assert ws.tasks["b"].state == "flight"


def test_throttle_incoming_transfers_on_count_limit(ws):
    ws.transfer_incoming_count_limit = 1
    ws.transfer_incoming_bytes_limit = 100_000
    ws.transfer_incoming_bytes_throttle_threshold = 1
    ws2 = "127.0.0.1:2"
    ws3 = "127.0.0.1:3"
    who_has = {"a": [ws2], "b": [ws3]}
    ws.handle_stimulus(
        ComputeTaskEvent.dummy(
            "c",
            who_has=who_has,
            nbytes={"a": 100, "b": 100},
            stimulus_id="s1",
        )
    )
    tasks_by_state = defaultdict(list)
    for ts in ws.tasks.values():
        tasks_by_state[ts.state].append(ts)
    assert len(tasks_by_state["flight"]) == 1
    assert len(tasks_by_state["fetch"]) == 1
    assert ws.transfer_incoming_bytes == 100

    in_flight_task = tasks_by_state["flight"][0]
    ws.handle_stimulus(
        GatherDepSuccessEvent(
            worker=who_has[in_flight_task.key][0],
            data={in_flight_task.key: 123},
            total_nbytes=100,
            stimulus_id="s2",
        )
    )
    assert tasks_by_state["flight"][0].state == "memory"
    assert tasks_by_state["fetch"][0].state == "flight"
    assert ws.transfer_incoming_bytes == 100


def test_throttling_incoming_transfer_on_transfer_bytes_same_worker(ws):
    ws.transfer_incoming_count_limit = 100
    ws.transfer_incoming_bytes_limit = 250
    ws.transfer_incoming_bytes_throttle_threshold = 1
    ws2 = "127.0.0.1:2"
    ws.handle_stimulus(
        ComputeTaskEvent.dummy(
            "d",
            who_has={"a": [ws2], "b": [ws2], "c": [ws2]},
            nbytes={"a": 100, "b": 100, "c": 100},
            stimulus_id="s1",
        )
    )
    tasks_by_state = defaultdict(list)
    for ts in ws.tasks.values():
        tasks_by_state[ts.state].append(ts)
    assert ws.transfer_incoming_bytes == 200
    assert len(tasks_by_state["flight"]) == 2
    assert len(tasks_by_state["fetch"]) == 1

    ws.handle_stimulus(
        GatherDepSuccessEvent(
            worker=ws2,
            data={ts.key: 123 for ts in tasks_by_state["flight"]},
            total_nbytes=200,
            stimulus_id="s2",
        )
    )
    assert all(ts.state == "memory" for ts in tasks_by_state["flight"])
    assert all(ts.state == "flight" for ts in tasks_by_state["fetch"])


def test_throttling_incoming_transfer_on_transfer_bytes_different_workers(ws):
    ws.transfer_incoming_count_limit = 100
    ws.transfer_incoming_bytes_limit = 150
    ws.transfer_incoming_bytes_throttle_threshold = 1
    ws2 = "127.0.0.1:2"
    ws3 = "127.0.0.1:3"
    who_has = {"a": [ws2], "b": [ws3]}
    ws.handle_stimulus(
        ComputeTaskEvent.dummy(
            "c",
            who_has=who_has,
            nbytes={"a": 100, "b": 100},
            stimulus_id="s1",
        )
    )
    tasks_by_state = defaultdict(list)
    for ts in ws.tasks.values():
        tasks_by_state[ts.state].append(ts)
    assert ws.transfer_incoming_bytes == 100
    assert len(tasks_by_state["flight"]) == 1
    assert len(tasks_by_state["fetch"]) == 1

    in_flight_task = tasks_by_state["flight"][0]
    ws.handle_stimulus(
        GatherDepSuccessEvent(
            worker=who_has[in_flight_task.key][0],
            data={in_flight_task.key: 123},
            total_nbytes=100,
            stimulus_id="s2",
        )
    )
    assert tasks_by_state["flight"][0].state == "memory"
    assert tasks_by_state["fetch"][0].state == "flight"


def test_do_not_throttle_connections_while_below_threshold(ws):
    ws.transfer_incoming_count_limit = 1
    ws.transfer_incoming_bytes_limit = 200
    ws.transfer_incoming_bytes_throttle_threshold = 50
    ws2 = "127.0.0.1:2"
    ws3 = "127.0.0.1:3"
    ws4 = "127.0.0.1:4"
    ws.handle_stimulus(
        ComputeTaskEvent.dummy(
            "b",
            who_has={"a": [ws2]},
            nbytes={"a": 1},
            stimulus_id="s1",
        )
    )
    assert ws.tasks["a"].state == "flight"

    ws.handle_stimulus(
        ComputeTaskEvent.dummy(
            "d",
            who_has={"c": [ws3]},
            nbytes={"c": 1},
            stimulus_id="s2",
        )
    )
    assert ws.tasks["c"].state == "flight"

    ws.handle_stimulus(
        ComputeTaskEvent.dummy(
            "f",
            who_has={"e": [ws4]},
            nbytes={"e": 100},
            stimulus_id="s3",
        )
    )
    assert ws.tasks["e"].state == "flight"
    assert ws.transfer_incoming_bytes == 102


def test_throttle_on_transfer_bytes_regardless_of_threshold(ws):
    ws.transfer_incoming_count_limit = 1
    ws.transfer_incoming_bytes_limit = 100
    ws.transfer_incoming_bytes_throttle_threshold = 50
    ws2 = "127.0.0.1:2"
    ws3 = "127.0.0.1:3"
    ws.handle_stimulus(
        ComputeTaskEvent.dummy(
            "b",
            who_has={"a": [ws2]},
            nbytes={"a": 1},
            stimulus_id="s1",
        )
    )
    assert ws.tasks["a"].state == "flight"

    ws.handle_stimulus(
        ComputeTaskEvent.dummy(
            "d",
            who_has={"c": [ws3]},
            nbytes={"c": 100},
            stimulus_id="s2",
        )
    )
    assert ws.tasks["c"].state == "fetch"
    assert ws.transfer_incoming_bytes == 1


def test_worker_nbytes(ws_with_running_task):
    ws = ws_with_running_task
    ws2 = "127.0.0.1:2"
    assert ws.nbytes == 0

    # executing->memory
    ws.handle_stimulus(ExecuteSuccessEvent.dummy("x", nbytes=12, stimulus_id="s1"))
    assert ws.nbytes == 12

    # flight->memory
    ws.handle_stimulus(
        AcquireReplicasEvent(who_has={"y": [ws2]}, nbytes={"y": 13}, stimulus_id="s2")
    )
    assert ws.nbytes == 12
    ws.handle_stimulus(
        GatherDepSuccessEvent(
            worker=ws2,
            data={"y": "foo"},
            total_nbytes=13,
            stimulus_id="s3",
        )
    )
    assert ws.nbytes == 12 + 13

    # released -> memory (scatter)
    ws.handle_stimulus(
        UpdateDataEvent(data={"z": "bar"}, report=False, stimulus_id="s3")
    )
    assert ws.nbytes == 12 + 13 + sizeof("bar")

    # actors
    ws.handle_stimulus(
        ComputeTaskEvent.dummy("w", actor=True, stimulus_id="s4"),
        ExecuteSuccessEvent.dummy("w", nbytes=14, stimulus_id="s5"),
    )
    assert ws.nbytes == 12 + 13 + sizeof("bar") + 14

    # memory -> released by FreeKeysEvent
    ws.handle_stimulus(FreeKeysEvent(keys=["z"], stimulus_id="s6"))
    assert ws.nbytes == 12 + 13 + 14

    # memory -> released by RemoveReplicasEvent
    ws.handle_stimulus(RemoveReplicasEvent(keys=["x", "y", "w"], stimulus_id="s7"))
    assert ws.nbytes == 0


def test_fetch_count(ws):
    ws.transfer_incoming_count_limit = 0
    ws2 = "127.0.0.1:2"
    ws3 = "127.0.0.1:3"
    assert ws.fetch_count == 0
    # Saturate comms
    # released->fetch->flight
    ws.handle_stimulus(
        AcquireReplicasEvent(who_has={"a": [ws2]}, nbytes={"a": 1}, stimulus_id="s1"),
        AcquireReplicasEvent(
            who_has={"b": [ws2, ws3]}, nbytes={"b": 1}, stimulus_id="s2"
        ),
    )
    assert ws.tasks["b"].coming_from == ws3
    assert ws.fetch_count == 0

    # released->fetch
    # d is in two data_needed heaps
    ws.handle_stimulus(
        AcquireReplicasEvent(
            who_has={"c": [ws2], "d": [ws2, ws3]},
            nbytes={"c": 1, "d": 1},
            stimulus_id="s3",
        )
    )
    assert ws.fetch_count == 2

    # fetch->released
    ws.handle_stimulus(FreeKeysEvent(keys={"c", "d"}, stimulus_id="s4"))
    assert ws.fetch_count == 0

    # flight->missing
    ws.handle_stimulus(
        GatherDepSuccessEvent(worker=ws2, data={}, total_nbytes=0, stimulus_id="s5")
    )
    assert ws.tasks["a"].state == "missing"
    print(ws.tasks)
    assert ws.fetch_count == 0
    assert len(ws.missing_dep_flight) == 1

    # flight->fetch
    ws.handle_stimulus(
        ComputeTaskEvent.dummy(
            "clog", who_has={"clog_dep": [ws2]}, priority=(-1,), stimulus_id="s6"
        ),
        GatherDepSuccessEvent(worker=ws3, data={}, total_nbytes=0, stimulus_id="s7"),
    )
    assert ws.tasks["b"].state == "fetch"
    assert ws.fetch_count == 1
    assert len(ws.missing_dep_flight) == 1


def test_task_counts(ws):
    assert ws.task_counts == {
        "constrained": 0,
        "executing": 0,
        "fetch": 0,
        "flight": 0,
        "long-running": 0,
        "memory": 0,
        "missing": 0,
        "other": 0,
        "ready": 0,
        "waiting": 0,
    }


def test_task_counts_with_actors(ws):
    ws.handle_stimulus(ComputeTaskEvent.dummy("x", actor=True, stimulus_id="s1"))
    assert ws.actors == {"x": None}
    assert ws.task_counts == {
        "constrained": 0,
        "executing": 1,
        "fetch": 0,
        "flight": 0,
        "long-running": 0,
        "memory": 0,
        "missing": 0,
        "other": 0,
        "ready": 0,
        "waiting": 0,
    }
    ws.handle_stimulus(ExecuteSuccessEvent.dummy("x", value=123, stimulus_id="s2"))
    assert ws.actors == {"x": 123}
    assert ws.task_counts == {
        "constrained": 0,
        "executing": 0,
        "fetch": 0,
        "flight": 0,
        "long-running": 0,
        "memory": 1,
        "missing": 0,
        "other": 0,
        "ready": 0,
        "waiting": 0,
    }