Set socket option to enable Strikebot traffic partitioning

diff --git a/strikebot/strikebot/setup.cfg b/strikebot/strikebot/setup.cfg
index 573a48ad99f97e890a4c3f943073c5a023e39491..8ae4a37bdcbcc4ef793cd62ffef21c8b42f20ac8 100644 (file)
--- 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 fc80e23947a04006283212d84d9ce763e58da1ca..4b4d12af84ee21eb1e10a5c510ee166e244ab25c 100644 (file)
--- 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()