From e16d690e777c625f12c2838b6719fa913c159c69 Mon Sep 17 00:00:00 2001
From: Joshua Watt <JPEWhacker@gmail.com>
Date: Wed, 29 May 2024 13:39:33 -0600
Subject: [PATCH] bitbake: hashserv: server: Add support for SO_REUSEPORT

SO_REUSEPORT is a socket option that allows multiple servers to listen
on the same TCP port, and the kernel will automatically load balance the
connections between them. This is particularly helpful for the hash
server since it runs in a single thread. To take advantage of a
multi-core server, multiple servers can be started in parallel with this
option (up to 1 per CPU) and the kernel will load balance between them.

(Bitbake rev: d72d5a7decb489e2af0ebc43cfea0ca3e4353e9b)

Signed-off-by: Joshua Watt <JPEWhacker@gmail.com>
Signed-off-by: Richard Purdie <richard.purdie@linuxfoundation.org>
---
 bitbake/bin/bitbake-hashserv     | 10 +++++++++-
 bitbake/lib/bb/asyncrpc/serv.py  | 34 +++++++++++++++++++++++++-------
 bitbake/lib/hashserv/__init__.py |  6 ++++--
 3 files changed, 40 insertions(+), 10 deletions(-)

diff --git a/bitbake/bin/bitbake-hashserv b/bitbake/bin/bitbake-hashserv
index 4bfb7abfbc..01503736b9 100755
--- a/bitbake/bin/bitbake-hashserv
+++ b/bitbake/bin/bitbake-hashserv
@@ -125,6 +125,11 @@ The following permissions are supported by the server:
         default=os.environ.get("HASHSERVER_ADMIN_PASSWORD", None),
         help="Create default admin user with password ADMIN_PASSWORD ($HASHSERVER_ADMIN_PASSWORD)",
     )
+    parser.add_argument(
+        "--reuseport",
+        action="store_true",
+        help="Enable SO_REUSEPORT, allowing multiple servers to bind to the same port for load balancing",
+    )
 
     args = parser.parse_args()
 
@@ -132,7 +137,9 @@ The following permissions are supported by the server:
 
     level = getattr(logging, args.log.upper(), None)
     if not isinstance(level, int):
-        raise ValueError("Invalid log level: %s (Try ERROR/WARNING/INFO/DEBUG)" % args.log)
+        raise ValueError(
+            "Invalid log level: %s (Try ERROR/WARNING/INFO/DEBUG)" % args.log
+        )
 
     logger.setLevel(level)
     console = logging.StreamHandler()
@@ -155,6 +162,7 @@ The following permissions are supported by the server:
         anon_perms=anon_perms,
         admin_username=args.admin_user,
         admin_password=args.admin_password,
+        reuseport=args.reuseport,
     )
     server.serve_forever()
     return 0
diff --git a/bitbake/lib/bb/asyncrpc/serv.py b/bitbake/lib/bb/asyncrpc/serv.py
index a66117acad..46d54fb511 100644
--- a/bitbake/lib/bb/asyncrpc/serv.py
+++ b/bitbake/lib/bb/asyncrpc/serv.py
@@ -138,14 +138,20 @@ class StreamServer(object):
 
 
 class TCPStreamServer(StreamServer):
-    def __init__(self, host, port, handler, logger):
+    def __init__(self, host, port, handler, logger, *, reuseport=False):
         super().__init__(handler, logger)
         self.host = host
         self.port = port
+        self.reuseport = reuseport
 
     def start(self, loop):
         self.server = loop.run_until_complete(
-            asyncio.start_server(self.handle_stream_client, self.host, self.port)
+            asyncio.start_server(
+                self.handle_stream_client,
+                self.host,
+                self.port,
+                reuse_port=self.reuseport,
+            )
         )
 
         for s in self.server.sockets:
@@ -209,11 +215,12 @@ class UnixStreamServer(StreamServer):
 
 
 class WebsocketsServer(object):
-    def __init__(self, host, port, handler, logger):
+    def __init__(self, host, port, handler, logger, *, reuseport=False):
         self.host = host
         self.port = port
         self.handler = handler
         self.logger = logger
+        self.reuseport = reuseport
 
     def start(self, loop):
         import websockets.server
@@ -224,6 +231,7 @@ class WebsocketsServer(object):
                 self.host,
                 self.port,
                 ping_interval=None,
+                reuse_port=self.reuseport,
             )
         )
 
@@ -262,14 +270,26 @@ class AsyncServer(object):
         self.loop = None
         self.run_tasks = []
 
-    def start_tcp_server(self, host, port):
-        self.server = TCPStreamServer(host, port, self._client_handler, self.logger)
+    def start_tcp_server(self, host, port, *, reuseport=False):
+        self.server = TCPStreamServer(
+            host,
+            port,
+            self._client_handler,
+            self.logger,
+            reuseport=reuseport,
+        )
 
     def start_unix_server(self, path):
         self.server = UnixStreamServer(path, self._client_handler, self.logger)
 
-    def start_websocket_server(self, host, port):
-        self.server = WebsocketsServer(host, port, self._client_handler, self.logger)
+    def start_websocket_server(self, host, port, reuseport=False):
+        self.server = WebsocketsServer(
+            host,
+            port,
+            self._client_handler,
+            self.logger,
+            reuseport=reuseport,
+        )
 
     async def _client_handler(self, socket):
         address = socket.address
diff --git a/bitbake/lib/hashserv/__init__.py b/bitbake/lib/hashserv/__init__.py
index 74367eb6b4..ac891e0174 100644
--- a/bitbake/lib/hashserv/__init__.py
+++ b/bitbake/lib/hashserv/__init__.py
@@ -13,6 +13,7 @@ from bb.asyncrpc.client import parse_address, ADDR_TYPE_UNIX, ADDR_TYPE_WS
 
 User = namedtuple("User", ("username", "permissions"))
 
+
 def create_server(
     addr,
     dbname,
@@ -25,6 +26,7 @@ def create_server(
     anon_perms=None,
     admin_username=None,
     admin_password=None,
+    reuseport=False,
 ):
     def sqlite_engine():
         from .sqlite import DatabaseEngine
@@ -60,9 +62,9 @@ def create_server(
         s.start_unix_server(*a)
     elif typ == ADDR_TYPE_WS:
         url = urlparse(a[0])
-        s.start_websocket_server(url.hostname, url.port)
+        s.start_websocket_server(url.hostname, url.port, reuseport=reuseport)
     else:
-        s.start_tcp_server(*a)
+        s.start_tcp_server(*a, reuseport=reuseport)
 
     return s