From 3f931adbd56d9756c6cbc188a4198c2fd8d3739b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=A9=99=E5=AD=90=E6=98=82?=
 <44974561+czasg@users.noreply.github.com>
Date: Thu, 25 Jul 2024 17:10:49 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E6=94=AF=E6=8C=81fork?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 pywss/__init__.py | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/pywss/__init__.py b/pywss/__init__.py
index bdfeed1..3c6e7cf 100644
--- a/pywss/__init__.py
+++ b/pywss/__init__.py
@@ -1,6 +1,7 @@
 # coding: utf-8
 import os
 import re
+import sys
 import json
 import time
 import gzip
@@ -757,6 +758,25 @@ def watchdog(self, interval: float = float(os.environ.get("PYWSS_WATCHDOG_INTERV
             time.sleep(interval)
         log.warning("exit")
 
+    def fork_processes(self, num_processes: int):
+        if sys.platform == "win32":
+            raise Exception("fork not available on windows")
+        pids = {}
+        num_processes = num_processes if num_processes > 0 else os.cpu_count()
+        for i in range(num_processes):
+            pid = os.fork()
+            if pid == 0:  # child process
+                return pid
+            else:
+                pids[pid] = i
+        # main process
+        while pids:
+            pid, status = os.wait()
+            if pid not in pids:
+                continue
+            pids.pop(pid)
+        return None
+
     def run(
             self,
             host: str = "0.0.0.0",
@@ -767,6 +787,8 @@ def run(
             thread_pool_size: int = int(os.environ.get("PYWSS_THREAD_POOL_SIZE", min(30, (os.cpu_count() or 1) * 5))),
             thread_pool_idle_time: int = int(os.environ.get("PYWSS_THREAD_POOL_IDLE_TIME", 300)),
             watch: bool = os.environ.get("PYWSS_WATCHDOG_ENABLE", "false").lower() == "true",
+            fork: bool = False,
+            num_processes: int = 0,
     ) -> None:
         # build app with [route:handler]
         self.build()
@@ -775,8 +797,15 @@ def run(
         # rigister signal closing
         for sig in (signal.SIGTERM, signal.SIGINT, signal.SIGILL):
             signal.signal(sig, lambda *args: self.close())
+        # use os.fork
+        if fork:
+            pid = self.fork_processes(num_processes)
+            if pid is None:
+                return
         # socket
         with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+            if fork:  # child
+                sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
             sock.bind((host, port))
             sock.listen(select_size)
             # queue of threading pool