feat: 添加 WithDefaultTimeout 默认超时配置
Invoke 之前完全依赖调用方传入的 ctx 控制超时,池子被占满或 Python 侧 handler 阻塞时,未设置 deadline 的调用会永久阻塞且不报错。新增 WithDefaultTimeout 选项,仅在 ctx 未设置 deadline 时兜底生效,调用方 显式设置的超时优先级更高。 同时补充 example 中的阻塞/超时演示(demoTimeout、demoBlocking)和 pool_test.go 集成测试,覆盖池占满排队、默认超时、显式 deadline 优先级、 流式输出超时后 channel 静默关闭等场景。
This commit is contained in:
@@ -12,6 +12,19 @@ import (
|
||||
"sync"
|
||||
)
|
||||
|
||||
// applyDefaultTimeout 在 ctx 未设置 deadline 时,套用 pool 配置的默认超时(WithDefaultTimeout)。
|
||||
// 调用方显式设置的 deadline 优先级更高,不会被覆盖;未配置默认超时时返回原 ctx 和 no-op cancel。
|
||||
func applyDefaultTimeout(ctx context.Context, pool Pool) (context.Context, context.CancelFunc) {
|
||||
if _, ok := ctx.Deadline(); ok {
|
||||
return ctx, func() {}
|
||||
}
|
||||
d := pool.defaultTimeout()
|
||||
if d <= 0 {
|
||||
return ctx, func() {}
|
||||
}
|
||||
return context.WithTimeout(ctx, d)
|
||||
}
|
||||
|
||||
// Invoke 调用 Python 暴露的函数,支持四种模式:
|
||||
//
|
||||
// 普通调用: Invoke[int](ctx, pool, "Add", 3, 4)
|
||||
@@ -119,6 +132,9 @@ func readResult(ctx context.Context, conn net.Conn, pool Pool, write func(Messag
|
||||
func invokeRegular[R any](ctx context.Context, pool Pool, method string, args ...any) (R, error) {
|
||||
var zero R
|
||||
|
||||
ctx, cancel := applyDefaultTimeout(ctx, pool)
|
||||
defer cancel()
|
||||
|
||||
argsJSON, err := json.Marshal(args)
|
||||
if err != nil {
|
||||
return zero, fmt.Errorf("marshal args: %w", err)
|
||||
@@ -167,13 +183,17 @@ func invokeRegular[R any](ctx context.Context, pool Pool, method string, args ..
|
||||
func invokeStreamOut[R any](ctx context.Context, pool Pool, method string, rt reflect.Type, args ...any) (R, error) {
|
||||
var zero R
|
||||
|
||||
ctx, cancel := applyDefaultTimeout(ctx, pool)
|
||||
|
||||
argsJSON, err := json.Marshal(args)
|
||||
if err != nil {
|
||||
cancel()
|
||||
return zero, fmt.Errorf("marshal args: %w", err)
|
||||
}
|
||||
|
||||
conn, w, err := pool.acquire(ctx)
|
||||
if err != nil {
|
||||
cancel()
|
||||
return zero, err
|
||||
}
|
||||
|
||||
@@ -185,6 +205,7 @@ func invokeStreamOut[R any](ctx context.Context, pool Pool, method string, rt re
|
||||
Args: argsJSON,
|
||||
}); err != nil {
|
||||
w.release(conn, false)
|
||||
cancel()
|
||||
return zero, contextErr(ctx, fmt.Errorf("write call: %w", err))
|
||||
}
|
||||
|
||||
@@ -198,6 +219,7 @@ func invokeStreamOut[R any](ctx context.Context, pool Pool, method string, rt re
|
||||
stop()
|
||||
ch.Close()
|
||||
w.release(conn, ctx.Err() == nil)
|
||||
cancel()
|
||||
}()
|
||||
for {
|
||||
msg, err := readResult(ctx, conn, pool, write)
|
||||
@@ -220,6 +242,9 @@ func invokeStreamOut[R any](ctx context.Context, pool Pool, method string, rt re
|
||||
func invokeStreamIn[R any](ctx context.Context, pool Pool, method string, streamArgIdx int, streamCh reflect.Value, args ...any) (R, error) {
|
||||
var zero R
|
||||
|
||||
ctx, cancel := applyDefaultTimeout(ctx, pool)
|
||||
defer cancel()
|
||||
|
||||
jsonArgs := make([]any, len(args))
|
||||
copy(jsonArgs, args)
|
||||
jsonArgs[streamArgIdx] = nil
|
||||
@@ -306,17 +331,21 @@ func invokeStreamIn[R any](ctx context.Context, pool Pool, method string, stream
|
||||
func invokeStreamBoth[R any](ctx context.Context, pool Pool, method string, streamArgIdx int, streamCh reflect.Value, rt reflect.Type, args ...any) (R, error) {
|
||||
var zero R
|
||||
|
||||
ctx, cancel := applyDefaultTimeout(ctx, pool)
|
||||
|
||||
jsonArgs := make([]any, len(args))
|
||||
copy(jsonArgs, args)
|
||||
jsonArgs[streamArgIdx] = nil
|
||||
|
||||
argsJSON, err := json.Marshal(jsonArgs)
|
||||
if err != nil {
|
||||
cancel()
|
||||
return zero, fmt.Errorf("marshal args: %w", err)
|
||||
}
|
||||
|
||||
conn, w, err := pool.acquire(ctx)
|
||||
if err != nil {
|
||||
cancel()
|
||||
return zero, err
|
||||
}
|
||||
|
||||
@@ -330,6 +359,7 @@ func invokeStreamBoth[R any](ctx context.Context, pool Pool, method string, stre
|
||||
StreamArgIdx: streamArgIdx,
|
||||
}); err != nil {
|
||||
w.release(conn, false)
|
||||
cancel()
|
||||
return zero, contextErr(ctx, fmt.Errorf("write call: %w", err))
|
||||
}
|
||||
|
||||
@@ -338,8 +368,15 @@ func invokeStreamBoth[R any](ctx context.Context, pool Pool, method string, stre
|
||||
var mu sync.Mutex
|
||||
write := func(msg Message) { mu.Lock(); writeMsg(conn, msg); mu.Unlock() } //nolint
|
||||
|
||||
// 两个 goroutine 共用同一个 ctx,defaultTimeout 产生的 cancel 必须等两者都结束才能调用,
|
||||
// 否则先完成的一方会提前取消掉另一方仍在进行的操作。
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(2)
|
||||
go func() { wg.Wait(); cancel() }()
|
||||
|
||||
// 写入 goroutine:输入 channel → Python chunks
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for {
|
||||
val, ok, cancelled := chanRecv(ctx, streamCh)
|
||||
if cancelled || !ok {
|
||||
@@ -361,6 +398,7 @@ func invokeStreamBoth[R any](ctx context.Context, pool Pool, method string, stre
|
||||
|
||||
// 读取 goroutine:Python chunks → 输出 channel,内联处理 callback
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
stop := watchCtx(ctx, conn, id, write)
|
||||
defer func() {
|
||||
stop()
|
||||
|
||||
Reference in New Issue
Block a user