基于room的多租户WebSocket连接管理实现

code-server-like/app/sockets/terminal_service.py

@@ -125,26 +125,23 @@ def read_and_forward_pty_output(fd=None, pid=None, room_id=None, namespace=None)
 class VSCLikeNameSpace(Namespace):
     def on_connect(self):
         """new client connected"""
-        # 确保terminal_config存在于session中
-        if 'terminal_config' not in session:
-            session['terminal_config'] = TERM_INIT_CONFIG.copy()
-            session.modified = True
+        # 为每个客户端生成唯一的会话ID和房间ID
+        session_id = str(uuid.uuid4())
+        room_id = f"terminal_{session_id}"
         
-        terminal_config = session['terminal_config']
+        # 加入房间
+        join_room(room_id)
         
-        # 检查是否已经有运行中的子进程
-        if terminal_config.get('child_pid'):
-            try:
-                # 验证进程是否真的存在
-                child_process = psutil.Process(terminal_config['child_pid'])
-                if child_process.status() in ('running', 'sleeping'):
-                    current_app.logger.debug("Already running child process: {}".format(terminal_config['child_pid']))
-                    return
-            except psutil.NoSuchProcess:
-                # 进程不存在，清理配置
-                terminal_config.pop('child_pid', None)
-                terminal_config.pop('fd', None)
-                session.modified = True
+        # 确保terminal_config存在于session中，每个客户端有独立配置
+        terminal_config = {
+            **TERM_INIT_CONFIG.copy(),
+            'session_id': session_id,
+            'room_id': room_id
+        }
+        session['terminal_config'] = terminal_config
+        session.modified = True
+        
+        current_app.logger.debug(f"Client connected, session_id: {session_id}, room_id: {room_id}")
         
         # create child process attached to a pty we can read from and write to
         (child_pid, fd) = pty.fork()
@@ -155,7 +152,6 @@ class VSCLikeNameSpace(Namespace):
             # of this subprocess
             try:
                 # 获取终端配置
-                terminal_config = session.get('terminal_config', {})
                 term_type = terminal_config.get('term_type')
                 
                 if not term_type:
@@ -164,7 +160,7 @@ class VSCLikeNameSpace(Namespace):
                 
                 path = TERM_INIT_CONFIG.get('client_path', {}).get(term_type, None)
                 if not path or not os.path.exists(path):
-                    print("Can't locate {} binary at {}, exit".format(term_type, path))
+                    print(f"Can't locate {term_type} binary at {path}, exit")
                     os._exit(1)
                 
                 # 获取连接参数
@@ -185,29 +181,29 @@ class VSCLikeNameSpace(Namespace):
                 elif term_type == 'ssh':
                     # 使用ssh连接
                     ssh_port = port if port else 22  # 默认端口22
-                    os.execl(path, 'ssh', '-p', str(ssh_port), '{}@{}'.format(username, domain))
+                    os.execl(path, 'ssh', '-p', str(ssh_port), f'{username}@{domain}')
                 else:
-                    print("Wrong term type {}, exit".format(term_type))
+                    print(f"Wrong term type {term_type}, exit")
                     os._exit(1)
             except Exception as e:
-                print("Error in child process: {}", str(e))
+                print(f"Error in child process: {e}")
                 os._exit(1)
         else:
-            # 更新会话配置
+            # 更新会话配置，保存到session中
             terminal_config['fd'] = fd
             terminal_config['child_pid'] = child_pid
-            terminal_config['room_id'] = rooms()[0]
+            session['terminal_config'] = terminal_config
             session.modified = True
             
             # 设置初始窗口大小
             set_winsize(fd, 50, 50)
             
             # 记录日志
-            current_app.logger.debug("child pid = {}".format(child_pid))
-            current_app.logger.debug("rooms of this session = {}".format(rooms()))
+            current_app.logger.debug(f"Child process created: {child_pid} for room: {room_id}")
+            current_app.logger.debug(f"Client joined room: {room_id}")
             
-            # 启动后台任务读取pty输出
-            socketio.start_background_task(read_and_forward_pty_output, fd, child_pid, rooms()[0], self)
+            # 启动后台任务读取pty输出，传入正确的room_id
+            socketio.start_background_task(read_and_forward_pty_output, fd, child_pid, room_id, self)
             current_app.logger.debug("background task running")
 
     def on_pty_input(self, data):
@@ -323,13 +319,23 @@ class VSCLikeNameSpace(Namespace):
         terminal_config = session.get('terminal_config', {})
         child_pid = terminal_config.get('child_pid')
         fd = terminal_config.get('fd')
+        room_id = terminal_config.get('room_id')
+        session_id = terminal_config.get('session_id')
+        
+        current_app.logger.debug(f"Client disconnecting, session_id: {session_id}, room_id: {room_id}")
+        
+        # 退出房间
+        if room_id:
+            leave_room(room_id)
+            current_app.logger.debug(f"Client left room: {room_id}")
         
         # 先关闭文件描述符，避免其他操作尝试使用无效的文件描述符
         if fd:
             try:
                 os.close(fd)
-            except Exception:
-                pass
+                current_app.logger.debug(f"Closed file descriptor: {fd}")
+            except Exception as e:
+                current_app.logger.debug(f"Error closing file descriptor: {e}")
         
         if child_pid:
             try:
@@ -340,23 +346,28 @@ class VSCLikeNameSpace(Namespace):
                     child_process.terminate()
                     # Wait for the process to terminate and collect its exit status
                     child_process.wait(timeout=2)
-                    current_app.logger.debug('user left the pty alone, terminated and waited')
-            except psutil.NoSuchProcess as err:
+                    current_app.logger.debug(f'User left pty alone, terminated and waited: {child_pid}')
+            except psutil.NoSuchProcess:
                 # Process already terminated, try to wait anyway to clean up any zombie
                 try:
                     os.waitpid(child_pid, os.WNOHANG)
-                except Exception:
-                    pass
+                    current_app.logger.debug(f'Cleaned up zombie process: {child_pid}')
+                except Exception as e:
+                    current_app.logger.debug(f"Error cleaning up zombie process: {e}")
             except psutil.TimeoutExpired:
                 # If process didn't terminate in time, kill it forcefully
                 try:
                     child_process.kill()
                     child_process.wait(timeout=1)
-                except Exception:
-                    pass
+                    current_app.logger.debug(f'Forcefully killed process: {child_pid}')
+                except Exception as e:
+                    current_app.logger.debug(f"Error forcefully killing process: {e}")
             except Exception as err:
                 current_app.logger.error(f'Error terminating process: {err}')
         
-        # Reset session config
-        session['terminal_config'] = TERM_INIT_CONFIG
+        # Clear session config for this client
+        session.pop('terminal_config', None)
+        session.modified = True
+        
+        current_app.logger.debug(f"Client disconnected, session_id: {session_id}, room_id: {room_id}")
         current_app.logger.debug('Client disconnected')