def send_task(conn: socket.socket, q: Queue, stop_e: threading.Event, delimiter: bytes, timeout=None):
print("[SEND][INFO] Started")
conn.settimeout(timeout)
try:
while True:
try:
n = q.get(timeout=timeout)
q.task_done()
except Empty:
if stop_e.is_set():
print(SharedData.bold("[SEND][INFO] Event set!"))
return
else:
try:
tcp_send(n, conn, delimiter)
except socket.timeout:
# really just want to use logging and dump logs in other thread..
print(SharedData.red("[Send][CRIT] Connection Broken!"))
break
except Exception:
print(SharedData.bold("[SEND][CRIT] Stopping SEND!"))
stop_e.set()
raise
async def recv_task(
s_receiver: asyncio.StreamReader,
q: asyncio.Queue,
e: asyncio.Event,
delimiter: bytes,
timeout=None,
):
print("[RECV][INFO] Started")
try:
while True:
try:
data = await tcp_recv(s_receiver, delimiter, timeout)
except asyncio.TimeoutError:
print('[RECV][WARN] TIMEOUT')
if e.is_set():
print(SharedData.bold(f"[RECV][INFO] Event set!"))
return
except asyncio.IncompleteReadError:
print(SharedData.red(f"[RECV][CRIT] Disconnected!"))
e.set()
return
else:
await q.put(data)
except Exception:
print(SharedData.bold("[RECV][CRIT] Stopping SEND!"))
e.set()
raise
async def send_task(
s_sender: asyncio.StreamWriter,
q: asyncio.Queue,
e: asyncio.Event,
delimiter: bytes,
timeout=None,
):
print("[SEND][INFO] Started")
try:
while True:
try:
n = await asyncio.wait_for(q.get(), timeout)
q.task_done()
except asyncio.TimeoutError:
if e.is_set():
print(SharedData.bold("[SEND][INFO] Event set!"))
return
else:
try:
await tcp_send(n, s_sender, delimiter, timeout)
except asyncio.TimeoutError:
# really just want to use logging and dump logs in other thread..
print(SharedData.red("[Send][CRIT] Connection Broken!"))
break
except Exception:
print(SharedData.bold("[SEND][CRIT] Stopping SEND!"))
e.set()
raise
def main():
event = threading.Event()
send_q = Queue()
recv_q = Queue()
server_thread = [
threading.Thread(target=send_thread, args=[send_q, event]),
threading.Thread(target=recv_thread, args=[recv_q, event]),
]
workers = [
threading.Thread(target=worker, args=[i, send_q, recv_q, event])
for i in range(config.WORKERS)
]
# start threads
for w in chain(server_thread, workers):
w.start()
for w in workers:
w.join()
event.set()
print(SharedData.bold("[C][info] All workers stopped."))
# send stop signal to server side RECV
print(SharedData.bold("[C][info] Sending kill signal to server RECV."))
send_q.put(config.END_MARK)
# waiting for SEND / RECV to stop
for t in server_thread:
t.join(timeout=5)
# load pickled result from INIT port
print("[C][Info] Fetching Port data from server.")
data = recv_until_eof(c_sock, config.END_MARK)
used_ports, shut_ports = json.loads(data.decode(config.ENCODING))
c_sock.close()
print("\n[Results]")
print(f"Used Ports : {used_ports}")
print(f"Closed Ports: {shut_ports}")
print(f"Excluded : {config.EXCLUDE}")
print(f"\nAll other ports from 1~{config.PORT_MAX} is open.")
async def recv_task(conn: socket.socket, q: Queue, e: threading.Event, delimiter: bytes, timeout=None):
print("[RECV][INFO] Started")
conn.settimeout(timeout)
try:
while True:
try:
data = tcp_recv(conn, delimiter)
except socket.timeout:
print('[RECV][WARN] TIMEOUT')
if e.is_set():
print(SharedData.bold(f"[RECV][INFO] Event set!"))
return
else:
q.put(data)
except Exception:
print(SharedData.bold("[RECV][CRIT] Stopping SEND!"))
e.set()
raise
async def main():
host, s_recv, s_send = await get_connection()
print(f'Config received from {host}.')
print(raw := await tcp_recv(s_recv))
config_json = json.loads(raw)
config = SimpleNamespace(**config_json)
event = asyncio.Event()
send_q = asyncio.Queue()
recv_q = asyncio.Queue()
delimiter = config.READ_UNTIL.encode(config.ENCODING)
server_task = (
asyncio.create_task(
send_task(s_send, send_q, event, delimiter, config.TIMEOUT)),
asyncio.create_task(
recv_task(s_recv, recv_q, event, delimiter, config.TIMEOUT)),
)
workers = [
asyncio.create_task(
worker(i, host, send_q, recv_q, event, delimiter, config.TIMEOUT))
for i in range(config.WORKERS)
]
# waiting for workers to complete.
for t in workers:
await t
print(SharedData.bold("[C][INFO] All workers stopped."))
if event.is_set(): # if set, then worker crashed and set the alarm!
print("task failed! waiting for server task to complete.")
else:
event.set()
for t in server_task:
try:
await t
except ConnectionResetError:
print("Connection reset!")
continue
# wait until server is closed - preventing RuntimeError.
s_send.close()
await s_send.wait_closed()
async def run_workers(
config,
works: asyncio.Queue,
send: asyncio.Queue,
recv: asyncio.Queue,
in_use: asyncio.Queue,
unreachable: asyncio.Queue,
e: asyncio.Event,
):
"""
Handle worker tasks.
"""
delimiter = config.READ_UNTIL.encode(config.ENCODING)
excl = set(config.EXCLUDE)
workers = [
asyncio.create_task(
worker(
i,
works,
send,
recv,
excl,
in_use,
unreachable,
e,
delimiter,
config.TIMEOUT,
)) for i in range(config.WORKERS)
]
for t in workers: # wait until workers are all complete
await t
print(SharedData.bold("[S][info] All workers stopped."))
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."))