From a644e094426a159585569a3f867636642a661bad Mon Sep 17 00:00:00 2001 From: Jakob Cornell Date: Sat, 28 Jan 2023 18:33:50 -0600 Subject: [PATCH] Set socket option to enable Strikebot traffic partitioning --- strikebot/strikebot/setup.cfg | 2 +- strikebot/strikebot/src/strikebot/__main__.py | 21 +++++++++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/strikebot/strikebot/setup.cfg b/strikebot/strikebot/setup.cfg index 573a48a..8ae4a37 100644 --- a/strikebot/strikebot/setup.cfg +++ b/strikebot/strikebot/setup.cfg @@ -1,6 +1,6 @@ [metadata] name = strikebot -version = 0.0.8 +version = 0.0.9 [options] package_dir = diff --git a/strikebot/strikebot/src/strikebot/__main__.py b/strikebot/strikebot/src/strikebot/__main__.py index fc80e23..4b4d12a 100644 --- a/strikebot/strikebot/src/strikebot/__main__.py +++ b/strikebot/strikebot/src/strikebot/__main__.py @@ -10,15 +10,18 @@ import argparse import configparser import datetime as dt import logging +import socket import ssl from trio import open_memory_channel, open_nursery, open_signal_receiver +from trio._socket import _SocketType try: from trio.lowlevel import current_root_task, Task except ImportError: from trio.hazmat import current_root_task, Task +from trio.socket import socket as orig_socket import trio_asyncio import triopg @@ -106,6 +109,24 @@ async def _signal_handler_impl(): _print_task_info() +# Patch Trio to set SO_PRIORITY = 1 on sockets. This is for everything except database traffic (i.e. WebSockets +# connections and Reddit HTTP connections). triopg happens to be the only of the three that doesn't use the Trio socket +# wrappers, so we can get away with patching Trio itself. +# This is fragile, so TODO: look into changes to asks/anyio/trio-websocket APIs to allow setting socket options. +def socket_patch(*args) -> _SocketType: + sock = orig_socket(*args) + sock.setsockopt(socket.SOL_SOCKET, socket.SO_PRIORITY, 1) + return sock + +# asks (via anyio) calls socket directly (module member access) +import trio.socket as mod +mod.socket = socket_patch + +# trio-websocket calls socket via open_tcp_stream in this helper module which has a direct import of the function +import trio._highlevel_open_tcp_stream as mod +mod.socket = socket_patch + + ap = argparse.ArgumentParser(__package__) ap.add_argument("config_path") args = ap.parse_args() -- 2.30.2