from dask.distributed import Client
import time
from dask.utils import SerializableLock


def dummy(i: int) -> int:
    return i


def constrained(lock: SerializableLock, dummy: int):
    assert not lock.locked()
    with lock:
        time.sleep(0.005)
    assert not lock.locked()


N_TASKS = 1000
lock = SerializableLock()
futures = []

if __name__ == "__main__":
    with Client(resources={"foo": 1}, n_workers=2, threads_per_worker=2) as client:
        for i in range(N_TASKS):
            x = client.submit(dummy, i)
            futures.append(client.submit(constrained, lock, x, resources={"foo": 1}))
        client.gather(futures)