diff --git a/example/worker.py b/example/worker.py index 174f0f9..61295ab 100644 --- a/example/worker.py +++ b/example/worker.py @@ -51,6 +51,14 @@ def double_stream(numbers: Iterator[int]) -> Iterator[int]: yield n * n +@expose +def stream_then_raise(n: int) -> Iterator[int]: + """流式输出中途抛异常,用于复现 _dispatch 里 end/error 消息错位的问题""" + for i in range(n): + yield i + raise ValueError("boom: stream_then_raise 中途失败") + + # ── struct(dataclass / dict)类型 ─────────────────────────────────────────── diff --git a/pool_test.go b/pool_test.go index 63affb7..48e7246 100644 --- a/pool_test.go +++ b/pool_test.go @@ -116,6 +116,41 @@ func TestDefaultTimeout(t *testing.T) { } } +// TestStreamErrorMidwayCorruptsNextCall 复现协议层的消息错位 bug: +// 流式 handler(生成器)中途抛异常时,Python 侧 _dispatch 会先发 end 再发 error +// (见 python/gobridge/__init__.py 的 finally 块),但 Go 侧 invokeStreamOut 读到 +// 第一条终止消息(end)就直接返回、把连接标记健康放回池子——那条本该属于这次调用 +// 的 error 消息还留在 socket 里没被读走。workers=1/maxConns=1 保证只有一条连接, +// 下一次调用必然复用同一条连接,从而读到上一次调用遗留的 error,得到错误的结果。 +// +// 这个测试目前应该失败(复现 bug);修复 _dispatch 后应该转为通过。 +func TestStreamErrorMidwayCorruptsNextCall(t *testing.T) { + pool := newTestPool(t, WithWorkers(1), WithMaxConns(1)) + ctx := context.Background() + + ch, err := Invoke[chan int](ctx, pool, "stream_then_raise", 3) + if err != nil { + t.Fatalf("Invoke stream_then_raise: %v", err) + } + var got []int + for v := range ch { + got = append(got, v) + } + if len(got) != 3 { + t.Fatalf("want 3 items before the raise, got %v", got) + } + + // 连接理论上已经被当作"健康"放回池子;下一次调用应该拿到自己的正常结果, + // 而不是上一次调用中途抛出的 ValueError 遗留消息。 + sum, err := Invoke[int](ctx, pool, "add", 1, 2) + if err != nil { + t.Fatalf("next call on the reused connection got corrupted by the previous call's stray error message: %v", err) + } + if sum != 3 { + t.Fatalf("want 3, got %d", sum) + } +} + // TestStreamTimeoutClosesChannelSilently 验证流式输出模式下,ctx 超时不会通过 error // 返回,而是静默关闭已返回的 channel,需要调用方自行检查 ctx.Err()。 func TestStreamTimeoutClosesChannelSilently(t *testing.T) { diff --git a/python/gobridge/__init__.py b/python/gobridge/__init__.py index 0a935fb..36a6447 100644 --- a/python/gobridge/__init__.py +++ b/python/gobridge/__init__.py @@ -336,11 +336,15 @@ def _dispatch(mux: _ConnMux, msg: dict): result = fn(*args) if inspect.isgenerator(result): - try: - for item in result: - mux.write({"id": msg_id, "type": "chunk", "data": _bytes_encode(item)}) - finally: - mux.write({"id": msg_id, "type": "end"}) + # 注意:end 只能在生成器正常耗尽后发送一次。之前这里用 try/finally 包裹, + # 导致生成器中途抛异常时会先发 end、再被下面的 except 发一条 error—— + # Go 侧读到第一条终止消息(end)就会返回并把连接标记健康放回池子, + # 那条多余的 error 消息留在 socket 里,会被下一个复用该连接的调用错误地 + # 当成自己的响应读走。改成不用 finally,异常直接交给下面统一处理, + # 保证一次调用只产生一条终止消息。 + for item in result: + mux.write({"id": msg_id, "type": "chunk", "data": _bytes_encode(item)}) + mux.write({"id": msg_id, "type": "end"}) else: mux.write({"id": msg_id, "type": "result", "data": _bytes_encode(result)})