what
ee5e5b96af
流式输出/双向流的 handler（Python 生成器）如果执行过程中抛异常， Invoke[chan T] 本身的 err 只描述"调用有没有发起成功"，跟这个异常 无关（永远是 nil），channel 只会静默提前关闭，调用方原本完全无法 感知。新增 WithStreamErrors(ctx) 返回一个包过的 ctx 和一个查询函数 streamErr，opt-in 之后可以查到具体错误。 错误记录挂在 WithStreamErrors 返回的 ctx 的对象图里（context.WithValue）， 不是全局表——调用方不再引用 ctx/channel 时会被 GC 自然回收，不需要 任何显式清理逻辑，也不依赖 ctx.Done()，即使用 context.Background() 也能正常释放；ctx 之后被别的 context.With*（包括 StickyCtx）再包一层 也不影响查询。 同时补充完整的自动化测试覆盖 example/main.go 里演示过的所有功能： 四种调用模式 × int/struct/slice/[]byte 的组合（client_test.go）、 WithHandlers/call_go 全双工（handlers_test.go）、NewSession 隔离性 与 StickyCtx 路由（session_test.go），之前这些只能靠人肉跑 go run 看输出，现在都有真实断言。
252 lines
7.6 KiB
Python
252 lines
7.6 KiB
Python
import sys
|
||
import os
|
||
|
||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "python"))
|
||
|
||
import dataclasses
|
||
import threading
|
||
import time
|
||
from typing import Iterator
|
||
|
||
from gobridge import expose, call_go, run, worker_id, worker_count
|
||
|
||
# ── worker_id / worker_count ──────────────────────────────────────────────────
|
||
# 只有 worker 0 才执行一次性初始化（如监听端口、建立长连接等），
|
||
# 其余 worker 跳过，避免端口冲突 / 重复连接。
|
||
print(f"[worker {worker_id}/{worker_count}] started", flush=True)
|
||
if worker_id == 0:
|
||
|
||
def _init_shared_resource():
|
||
# 示例：此处可启动 WebSocket 客户端、监听 TCP 端口等
|
||
print(f"[worker {worker_id}] shared resource initialized", flush=True)
|
||
|
||
threading.Thread(target=_init_shared_resource, daemon=True).start()
|
||
|
||
|
||
# ── 基础类型 ─────────────────────────────────────────────────────────────────
|
||
|
||
|
||
@expose
|
||
def add(a: int, b: int) -> int:
|
||
return a + b
|
||
|
||
|
||
@expose
|
||
def get_env(name: str) -> str:
|
||
"""读取当前进程环境变量，取不到时返回空字符串；用于验证 WithEnv 注入是否生效"""
|
||
return os.environ.get(name, "")
|
||
|
||
|
||
@expose
|
||
def count_threads() -> int:
|
||
"""返回当前 Python 进程存活线程数，用于检测阻塞在 chunk_q.get() 上的线程是否泄漏"""
|
||
return threading.active_count()
|
||
|
||
|
||
@expose
|
||
def range_gen(start: int, stop: int) -> Iterator[int]:
|
||
"""流式输出：对应 Go 侧 Invoke[chan int]"""
|
||
for i in range(start, stop):
|
||
yield i
|
||
|
||
|
||
@expose
|
||
def sum_stream(numbers: Iterator[int]) -> int:
|
||
"""流式输入：对应 Go 侧传入 chan int 参数"""
|
||
return sum(numbers)
|
||
|
||
|
||
@expose
|
||
def double_stream(numbers: Iterator[int]) -> Iterator[int]:
|
||
"""双向流：输入每个数，yield 其平方"""
|
||
for n in numbers:
|
||
yield n * n
|
||
|
||
|
||
@expose
|
||
def stream_then_raise(n: int) -> Iterator[int]:
|
||
"""流式输出执行过程中抛异常（还没 yield 完就失败），
|
||
用于复现 _dispatch 里 end/error 消息错位的问题"""
|
||
for i in range(n):
|
||
if i == n - 1:
|
||
raise ValueError(f"boom: stream_then_raise 在第 {i} 个元素时失败")
|
||
yield i
|
||
|
||
|
||
# ── struct（dataclass / dict）类型 ───────────────────────────────────────────
|
||
|
||
|
||
@dataclasses.dataclass
|
||
class User:
|
||
id: int
|
||
name: str
|
||
score: float
|
||
level: str = ""
|
||
|
||
|
||
@expose
|
||
def get_user(uid: int) -> dict:
|
||
"""普通调用：返回一个 struct（Go 对应 User）"""
|
||
return {"id": uid, "name": f"user_{uid}", "score": uid * 1.5}
|
||
|
||
|
||
@expose
|
||
def total_score(users: list) -> float:
|
||
"""slice 输入：接收 []User，返回总分"""
|
||
return sum(u["score"] for u in users)
|
||
|
||
|
||
@expose
|
||
def enrich_users(users: list) -> list:
|
||
"""slice 输入输出：为每个 user 追加 level 字段"""
|
||
result = []
|
||
for u in users:
|
||
u = dict(u)
|
||
u["level"] = "gold" if u["score"] >= 10 else "silver"
|
||
result.append(u)
|
||
return result
|
||
|
||
|
||
@expose
|
||
def gen_users(count: int) -> Iterator[dict]:
|
||
"""流式输出 struct：yield 多个 User，对应 Go 侧 Invoke[chan User]"""
|
||
for i in range(1, count + 1):
|
||
yield {"id": i, "name": f"user_{i}", "score": float(i * 3)}
|
||
|
||
|
||
@expose
|
||
def process_users(users: Iterator[dict]) -> Iterator[dict]:
|
||
"""双向流 struct：输入流式 User，yield 处理后的 User"""
|
||
for u in users:
|
||
yield {"id": u["id"], "name": u["name"].upper(), "score": u["score"] * 2}
|
||
|
||
|
||
# ── []byte / bytes 示例 ──────────────────────────────────────────────────────
|
||
|
||
|
||
@expose
|
||
def bytes_reverse(data: bytes) -> bytes:
|
||
"""接收 []byte，返回翻转后的 []byte"""
|
||
return data[::-1]
|
||
|
||
|
||
@expose
|
||
def bytes_concat(a: bytes, b: bytes) -> bytes:
|
||
"""接收两个 []byte 参数，返回拼接结果"""
|
||
return a + b
|
||
|
||
|
||
@expose
|
||
def bytes_chunks(data: bytes, size: int):
|
||
"""流式输出：将 []byte 按 size 切分，逐块 yield（对应 Go Invoke[chan []byte]）"""
|
||
for i in range(0, len(data), size):
|
||
yield data[i:i + size]
|
||
|
||
|
||
# ── 超时示例 ─────────────────────────────────────────────────────────────────
|
||
|
||
|
||
@expose
|
||
def sleep_seconds(n: float) -> str:
|
||
"""模拟一次耗时阻塞调用（不检查 ctx 取消），用于演示 WithDefaultTimeout"""
|
||
time.sleep(n)
|
||
return f"slept {n}s"
|
||
|
||
|
||
@expose
|
||
def slow_range_gen(start: int, stop: int, delay_ms: int) -> Iterator[int]:
|
||
"""流式输出，每个元素之间人为延迟，用于演示 ctx 超时会提前关闭 channel"""
|
||
for i in range(start, stop):
|
||
time.sleep(delay_ms / 1000)
|
||
yield i
|
||
|
||
|
||
# ── Server 全双工示例 ────────────────────────────────────────────────────────
|
||
|
||
|
||
@expose
|
||
def compute_with_go_mul(a: int, b: int) -> int:
|
||
"""示例1：call_go[int] 指定返回类型"""
|
||
return call_go[int]("Multiply", a, b)
|
||
|
||
|
||
@expose
|
||
def squared_with_log(n: int) -> Iterator[int]:
|
||
"""示例2：流式输出，每次 yield 前 call_go("Log") 回调 Go"""
|
||
for i in range(1, n + 1):
|
||
call_go("Log", f"yielding {i}² = {i * i}")
|
||
yield i * i
|
||
|
||
|
||
@expose
|
||
def to_upper(s: str) -> str:
|
||
"""辅助方法：被 Go 的 EnrichName handler 内部调用"""
|
||
return s.upper()
|
||
|
||
|
||
@expose
|
||
def full_chain(name: str) -> str:
|
||
"""示例3：Go→Python→Go→Python 四层链路
|
||
|
||
full_chain("world")
|
||
→ call_go[str]("EnrichName", "world") # Python 调 Go
|
||
→ Invoke[string](ctx, serv, "to_upper", "world") # Go 再调 Python
|
||
← "WORLD"
|
||
← "Hello, WORLD!"
|
||
← "Hello, WORLD!"
|
||
"""
|
||
return call_go[str]("EnrichName", name)
|
||
|
||
|
||
@expose
|
||
def get_user_via_go(uid: int) -> dict:
|
||
"""示例4：call_go[User] 自动将 Go 返回的 dict 构造为 dataclass 实例"""
|
||
user = call_go[User]("MakeUser", uid) # Go 返回 {"id":..,"name":..,"score":..}
|
||
user.level = "gold" if user.score >= 10 else "silver"
|
||
return dataclasses.asdict(user)
|
||
|
||
|
||
# ── Session 亲和示例 ──────────────────────────────────────────────────────────
|
||
# _sessions 保存每个 session 的状态，key 由调用方提供
|
||
_sessions: dict = {}
|
||
|
||
# _global_counter 是进程级全局变量，同一 worker 的所有 session 共享
|
||
_global_counter: int = 0
|
||
|
||
|
||
@expose
|
||
def global_increment(delta: int) -> str:
|
||
global _global_counter
|
||
_global_counter += delta
|
||
return f"[worker {worker_id}] counter = {_global_counter}"
|
||
|
||
|
||
@expose
|
||
def global_get() -> str:
|
||
return f"[worker {worker_id}] counter = {_global_counter}"
|
||
|
||
|
||
@expose
|
||
def session_init(session_id: str, value: int) -> str:
|
||
_sessions[session_id] = {"value": value, "steps": []}
|
||
return f"[worker {worker_id}] session {session_id} init with {value}"
|
||
|
||
|
||
@expose
|
||
def session_step(session_id: str, delta: int) -> int:
|
||
s = _sessions[session_id]
|
||
s["value"] += delta
|
||
s["steps"].append(delta)
|
||
return s["value"]
|
||
|
||
|
||
@expose
|
||
def session_result(session_id: str) -> dict:
|
||
return _sessions.pop(session_id)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
run()
|
||
print("worker_id", worker_id)
|
||
print("worker_count", worker_count)
|