train Walk and tackle the memory leak issue.

communication/server.py

@@ -1,4 +1,5 @@
 import logging
+import os
 import socket
 import time
 from select import select
@@ -15,6 +16,11 @@ class Server:
         self.__socket: socket.socket = self._create_socket()
         self.__send_buff = []
         self.__rcv_buffer_size = 1024
+        self.__rcv_buffer_default_size = 1024
+        self.__max_msg_size = 1048576
+        self.__shrink_threshold = 8192
+        self.__shrink_after_msgs = 200
+        self.__small_msg_streak = 0
         self.__rcv_buffer = bytearray(self.__rcv_buffer_size)
 
     def _create_socket(self) -> socket.socket:
@@ -105,6 +111,10 @@ class Server:
 
             msg_size = int.from_bytes(self.__rcv_buffer[:4], byteorder="big", signed=False)
 
+            # Guard against corrupted frame lengths that would trigger huge allocations.
+            if msg_size <= 0 or msg_size > self.__max_msg_size:
+                raise ConnectionResetError
+
             if msg_size > self.__rcv_buffer_size:
                 self.__rcv_buffer_size = msg_size
                 self.__rcv_buffer = bytearray(self.__rcv_buffer_size)
@@ -120,6 +130,15 @@ class Server:
                 message=self.__rcv_buffer[:msg_size].decode()
             )
 
+            if msg_size <= self.__shrink_threshold and self.__rcv_buffer_size > self.__rcv_buffer_default_size:
+                self.__small_msg_streak += 1
+                if self.__small_msg_streak >= self.__shrink_after_msgs:
+                    self.__rcv_buffer_size = self.__rcv_buffer_default_size
+                    self.__rcv_buffer = bytearray(self.__rcv_buffer_size)
+                    self.__small_msg_streak = 0
+            else:
+                self.__small_msg_streak = 0
+
             # 如果socket没有更多数据就退出
             if len(select([self.__socket], [], [], 0.0)[0]) == 0:
                 break