From 1a39e87c8cb5de165775d753437aab920b48cea5 Mon Sep 17 00:00:00 2001
From: qing <1776177662@qq.com>
Date: Sun, 26 Jul 2026 21:04:42 +0800
Subject: [PATCH] =?UTF-8?q?=E6=9B=B4=E6=96=B0pymsg?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 dnfs/dnfs.py            | 11 +++--------
 dnfs/log_client.py      | 15 +++------------
 pymsg/pymsg/__init__.py | 32 +++++++++++++++++++++++++++++++-
 3 files changed, 37 insertions(+), 21 deletions(-)

diff --git a/dnfs/dnfs.py b/dnfs/dnfs.py
index 60c063a..1145a17 100755
--- a/dnfs/dnfs.py
+++ b/dnfs/dnfs.py
@@ -23,7 +23,6 @@ from waitress import serve
 import os
 import json
 import base64
-import socket as _socket
 import urllib.request
 import urllib.error
 from typing import Dict, Any
@@ -33,7 +32,7 @@ from log_client import logs
 # pymsg 导入
 import sys as _sys
 _sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
-from pymsg import send as _pmsg_send, recv as _pmsg_recv  # noqa: E402
+from pymsg import MsgClient  # noqa: E402
 
 _DBV2_ADDR = '127.0.0.1'
 _DBV2_PORT = 8005
@@ -44,12 +43,8 @@ def _call_dbv2(payload: dict, timeout: int = 10) -> dict:
     """
     payload = {**payload, "_route": "execute"}
     try:
-        sock = _socket.socket(_socket.AF_INET, _socket.SOCK_STREAM)
-        sock.settimeout(timeout)
-        sock.connect((_DBV2_ADDR, _DBV2_PORT))
-        _pmsg_send(sock, json.dumps(payload))
-        resp_text = _pmsg_recv(sock, timeout=timeout)
-        sock.close()
+        client = MsgClient(_DBV2_ADDR, _DBV2_PORT, "dnfs")
+        resp_text = client.call(json.dumps(payload), timeout=timeout)
         return json.loads(resp_text)
     except Exception as e:
         raise ConnectionError(
diff --git a/dnfs/log_client.py b/dnfs/log_client.py
index 1408c1e..6ee468c 100644
--- a/dnfs/log_client.py
+++ b/dnfs/log_client.py
@@ -19,14 +19,13 @@
 import json
 import os as _os
 import sys as _sys
-import socket as _socket
 import threading
 import queue
 import time
 import atexit
 
 _sys.path.insert(0, _os.path.join(_os.path.dirname(__file__), ".."))
-from pymsg import send as _pmsg_send, recv as _pmsg_recv  # noqa: E402
+from pymsg import MsgClient  # noqa: E402
 
 # ── 配置 ──────────────────────────────────────────────────────────────
 
@@ -135,17 +134,9 @@ class LogClient:
     def _send(self, batch):
         """通过 pymsg 短连接发送一批日志到 dbv2。失败时静默丢弃。"""
         try:
-            sock = _socket.socket(_socket.AF_INET, _socket.SOCK_STREAM)
-            sock.settimeout(_TIMEOUT)
-            sock.connect((_DBV2_ADDR, _DBV2_PORT))
+            client = MsgClient(_DBV2_ADDR, _DBV2_PORT, "dnfs-log")
             payload = {"_route": "log/ingest", "logs": batch}
-            _pmsg_send(sock, json.dumps(payload))
-            # 消费应答
-            try:
-                _pmsg_recv(sock, timeout=_TIMEOUT)
-            except Exception:
-                pass
-            sock.close()
+            client.call(json.dumps(payload), timeout=_TIMEOUT)
         except Exception:
             # fail-safe：日志服务不可用，静默丢弃，不影响业务
             pass
diff --git a/pymsg/pymsg/__init__.py b/pymsg/pymsg/__init__.py
index fb8040a..4a604ef 100644
--- a/pymsg/pymsg/__init__.py
+++ b/pymsg/pymsg/__init__.py
@@ -146,7 +146,7 @@ class MsgClient:
         self._sock = _socket.socket(_socket.AF_INET, _socket.SOCK_STREAM)
         self._sock.connect((self.addr, self.port))
 
-    # ---- 发送单条消息（短连接）--------------------------------------------
+    # ---- 发送单条消息（短连接，不等待回复）--------------------------------
     def send(self, msg: str) -> None:
         """发送一条消息到服务器（使用短连接，pymsg 协议）。"""
         try:
@@ -159,6 +159,36 @@ class MsgClient:
                 f"Could not connect to {self.addr}:{self.port}"
             )
 
+    # ---- 请求-应答（短连接）------------------------------------------------
+    def call(self, msg: str, *, timeout: "float | None" = None) -> str:
+        """发送一条消息并阻塞等待单条回复（短连接：connect → send → recv → close）。
+
+        适合请求-应答模式，如 RPC 调用。不会给消息添加 name 前缀。
+
+        Args:
+            msg:     要发送的原始字符串。
+            timeout: 超时秒数，默认 10 秒。
+
+        Returns:
+            str: 服务器返回的字符串。
+
+        Raises:
+            ConnectionRefusedError: 目标服务不可达。
+            TimeoutError:           等待响应超时。
+            ConnectionError:        连接在收全数据前断开。
+        """
+        _tout = timeout if timeout is not None else 10.0
+        try:
+            with _socket.socket(_socket.AF_INET, _socket.SOCK_STREAM) as s:
+                s.settimeout(_tout)
+                s.connect((self.addr, self.port))
+                send(s, msg)
+                return recv(s, timeout=_tout)
+        except ConnectionRefusedError:
+            raise ConnectionRefusedError(
+                f"Could not connect to {self.addr}:{self.port}"
+            )
+
     # ---- 持续监听（阻塞当前线程，回调模式）-------------------------------
     def listen(self, on_message: "callable") -> None:
         """阻塞当前线程，持续接收服务器消息并回调。