def worker(id_: int, send, recv, event: threading.Event):
q: Queue
send: Queue
recv: Queue
while not event.is_set():
# announce server that the worker is ready.
print(f"[CS{id_:2}][Info] Worker {id_:2} READY.")
send.put(id_)
try:
p = recv.get()
except Empty:
continue
recv.task_done()
print(f"[CS{id_:2}][Info] Worker {id_:2} received {p}.")
if p > config.PORT_MAX:
print(SharedData.cyan(f"[CS{id_:2}][Info] Stop Signal received!"))
break
print(f"[CS{id_:2}][Info] Connecting Port {p}.")
child_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
child_sock.settimeout(config.TIMEOUT)
try:
child_sock.connect((host, p))
except socket.timeout:
print(SharedData.red(f"[CS{id_:2}][Info] Port {p} Timeout."))
except OSError:
print(SharedData.red(f"[CS{id_:2}][Warn] Port {p} in use."))
else:
print(SharedData.green(f"[CS{id_:2}][Info] Port {p} is open."))
child_sock.close()
async def worker(id_: int, host, send, recv, event, delimiter, timeout=None):
q: asyncio.Queue
send: asyncio.Queue
recv: asyncio.Queue
event: asyncio.Event
try:
# if one thread crashes, will trigger event and gradually stop all threads.
while not event.is_set():
# announce server that the worker is ready.
print(f"[CS{id_:2}][INFO] Worker {id_:2} READY.")
await send.put(id_)
try:
p = await asyncio.wait_for(recv.get(), timeout=timeout)
p = int(p)
recv.task_done()
except asyncio.TimeoutError:
print(
SharedData.red(
f"[CS{id_:2}][WARN] Worker {id_:2} timeout fetching from Queue."
))
continue
except ValueError:
print(
SharedData.cyan(
f"[CS{id_:2}][INFO] Stop Signal received!"))
break
print(f"[CS{id_:2}][INFO] Connecting Port {p}.")
try:
child_recv, child_send = await asyncio.wait_for(
asyncio.open_connection(host, p), timeout)
except asyncio.TimeoutError:
print(
SharedData.purple(f"[CS{id_:2}][INFO] Port {p} timeout."))
except OSError:
print(
SharedData.red(
f"[CS{id_:2}][WARN] Port {p} connection refused."))
else:
try:
print(await tcp_recv(child_recv,
delimiter,
timeout=timeout))
except asyncio.TimeoutError:
print(
SharedData.purple(
f"[CS{id_:2}][INFO] Port {p} timeout."))
except asyncio.IncompleteReadError:
print(
SharedData.red(
f"[CS{id_:2}][WARN] Port {p} disconnected!"))
else:
print(f"[CS{id_:2}][INFO] Port {p} open.")
print(
SharedData.green(
f"[CS{id_:2}][INFO] Port {p} is available."))
finally:
child_send.close()
await child_send.wait_closed()
except Exception:
# trigger event to stop all threads.
print(SharedData.red(f"[CS{id_:2}][CRIT] Exception Event set!."))
event.set()
raise
print(SharedData.bold(f"[CS{id_:2}][INFO] Task Finished."))
async def worker(
id_,
task_q: asyncio.Queue,
send: asyncio.Queue,
recv: asyncio.Queue,
exclude: set,
used: asyncio.Queue,
unreachable: asyncio.Queue,
event: asyncio.Event,
delimiter: bytes,
timeout=None,
):
async def worker_handler(
reader: asyncio.StreamReader,
writer: asyncio.StreamWriter,
port: int,
handle_finished: asyncio.Event,
):
print(SharedData.green(f"[SS{id_:2}][INFO] --- IN HANDLER ---"))
await tcp_send(p, writer, delimiter, timeout=timeout)
print(SharedData.green(f"[SS{id_:2}][INFO] Port {port} is open."))
writer.close()
await writer.wait_closed()
handle_finished.set()
try:
while not task_q.empty() and not event.is_set():
# receive worker announcement.
try:
worker_id = await asyncio.wait_for(recv.get(), timeout)
recv.task_done()
except asyncio.TimeoutError:
print(SharedData.red(f"[SS{id_:2}][Warn] Timeout."))
continue
print(f"[SS{id_:2}][INFO] Worker {worker_id} available.")
# get next work.
print(f"[SS{id_:2}][INFO] Getting new port.")
# if timeout getting port, either task is empty or just coroutine delayed.
try:
p: int = await asyncio.wait_for(task_q.get(), timeout)
task_q.task_done()
except asyncio.TimeoutError:
if task_q.empty():
break
else:
await recv.put(worker_id)
continue
# put back in and run again.
# check if port is in blacklist.
if p in exclude:
print(SharedData.cyan(f"[SS{id_:2}][INFO] Skipping Port {p}."))
continue
print(f"[SS{id_:2}][INFO] Sending port {p} to client.")
await send.put(p)
handle_ev = asyncio.Event()
print(f"[SS{id_:2}][INFO] Trying to serve port {p}.")
try:
# child_sock = await asyncio.wait_for(asyncio.start_server(
# lambda r, w: worker_handler(r, w, p, handle_ev), port=p), TIMEOUT_FACTOR)
child_sock = await asyncio.start_server(
lambda r, w: worker_handler(r, w, p, handle_ev), port=p)
# except asyncio.TimeoutError:
# # not sure why start_server gets timeout.
# # maybe I need to control number of task so opening server don't hang.
# print(SharedData.red(f"[SS{id_:2}][Warn] Port {p} timeout while opening."))
# await unreachable.put(p)
except AssertionError:
print(
SharedData.red(
f"[SS{id_:2}][INFO] Port {p} assertion failed!"))
await unreachable.put(p)
except OSError:
print(SharedData.red(f"[SS{id_:2}][Warn] Port {p} in use."))
await used.put(p)
else:
try:
await child_sock.start_serving()
await asyncio.wait_for(handle_ev.wait(), timeout)
except asyncio.TimeoutError:
print(
SharedData.red(f"[SS{id_:2}][Warn] Port {p} timeout."))
await unreachable.put(p)
finally:
child_sock.close()
await child_sock.wait_closed()
# Send end signal to client.
# first worker catching this signal will go offline.
print(SharedData.cyan(f"[SS{id_:2}][INFO] Done. Sending stop signal."))
await send.put("DONE"
) # causing int type-error on client side workers.
except Exception:
# trigger event to stop all threads.
print(SharedData.red(f"[SS{id_:2}][CRIT] Exception Event set!."))
event.set()
raise
if event.is_set():
print(SharedData.bold(f"[SS{id_:2}][WARN] Task Finished by event."))
else:
print(SharedData.bold(f"[SS{id_:2}][INFO] Task Finished."))