diff --git a/README.md b/README.md index 9478752..8a8ae87 100644 --- a/README.md +++ b/README.md @@ -186,6 +186,53 @@ def slow_compute(n: int) -> int: > 限制:长时间不释放 GIL 的 C 扩展(如大规模 numpy 矩阵运算)无法被中断,需等其释放 GIL 后才触发。 +### 默认超时(WithDefaultTimeout) + +`Invoke` 本身不带任何默认超时——如果传入的 `ctx` 没有 deadline(比如直接用 `context.Background()`),且池子被占满(所有 worker 的连接都在处理别的请求),调用会**永久阻塞**,不会自动放弃。 + +`WithDefaultTimeout` 用于兜底这种情况:仅当调用方传入的 `ctx` **未设置** deadline 时才生效,调用方显式设置的 `context.WithTimeout` 优先级更高,不会被覆盖。 + +> 池子被占满时新调用是**排队阻塞等待**,不是立刻失败或被跳过——`example/main.go` 中的 `demoBlocking` 用 `workers=2, maxConns=2`(总容量 4)故意占满连接池, +> 再发起一个不设超时的调用,实测会阻塞约 1.8s(等到某个占位任务释放连接)才返回,而不是瞬间失败。 + +```go +pool, _ := gobridge.NewPool("worker.py", + gobridge.WithDefaultTimeout(500 * time.Millisecond), +) + +// 未设置 deadline,池的默认超时自动生效,500ms 后返回 context.DeadlineExceeded +_, err := gobridge.Invoke[string](context.Background(), pool, "sleep_seconds", 2.0) + +// 显式传入的 deadline 优先,不受默认超时影响 +ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) +defer cancel() +result, err := gobridge.Invoke[string](ctx, pool, "sleep_seconds", 1.0) // 正常返回 +``` + +**超时是否会通过 `error` 返回,取决于调用模式:** + +| 调用模式 | 超时表现 | +|---|---| +| 普通调用 / 流式输入 | `Invoke` 返回非 nil 的 `error`(`context.DeadlineExceeded` / `context.Canceled`) | +| 流式输出 / 双向流 | `Invoke` 建立阶段失败会返回 `error`;**建立成功后**若中途超时,只会静默关闭已返回的 channel,不会有第二次 `error` | + +流式模式下,`for v := range ch` 结束后无法区分"正常读完"还是"被超时打断",需要调用方自行检查传入的 `ctx.Err()`: + +```go +ch, err := gobridge.Invoke[chan int](ctx, pool, "slow_range_gen", 1, 10, 100) +if err != nil { + // 建立阶段失败/超时 +} +for v := range ch { + fmt.Println(v) +} +if ctx.Err() != nil { + // channel 是因为超时/取消提前关闭的,不是正常 yield 完 +} +``` + +完整可运行示例见 [example/main.go](example/main.go) 中的 `demoTimeout`(默认超时 / 显式 deadline 优先级 / 流式超时静默关闭 channel)和 `demoBlocking`(池子占满后阻塞排队)。 + ## Session 亲和路由 默认情况下,每次 `Invoke` 通过轮询分配 worker 进程。当多次调用需要共享同一 Python 进程的状态时,可以使用 Session 或 StickyCtx 将调用固定到同一进程。 @@ -394,6 +441,7 @@ pool, err := gobridge.NewPool("worker.py", gobridge.WithSocketDir("/var/run/myapp"), // socket 文件目录,默认 /tmp gobridge.WithStdout(os.Stdout), // 子进程 stdout,默认 os.Stdout gobridge.WithStderr(os.Stderr), // 子进程 stderr,默认 os.Stderr + gobridge.WithDefaultTimeout(10*time.Second), // Invoke 默认超时,默认不启用 ) // 静默模式:丢弃子进程输出 @@ -414,6 +462,7 @@ pool, err := gobridge.NewPool("worker.py", | `WithSocketDir(dir)` | UDS socket 文件目录 | `"/tmp"` | | `WithStdout(w)` | 子进程标准输出 | `os.Stdout` | | `WithStderr(w)` | 子进程标准错误 | `os.Stderr` | +| `WithDefaultTimeout(d)` | `Invoke` 默认超时,仅在传入的 `ctx` 未设置 deadline 时生效 | 不启用 | ## 使用 uv 管理 Python 环境 diff --git a/client.go b/client.go index 1991466..930207f 100644 --- a/client.go +++ b/client.go @@ -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() diff --git a/example/main.go b/example/main.go index d5056b6..e20108f 100644 --- a/example/main.go +++ b/example/main.go @@ -6,6 +6,8 @@ import ( "log" "path/filepath" "runtime" + "sync" + "time" "git.fsdpf.net/go/gobridge" ) @@ -35,6 +37,8 @@ func main() { demoPool(ctx, pool) demoServer(ctx, script) demoSession(ctx, script) + demoTimeout(script) + demoBlocking(script) } func demoPool(ctx context.Context, pool gobridge.Pool) { @@ -277,6 +281,93 @@ func demoSession(ctx context.Context, script string) { } } +func demoTimeout(script string) { + fmt.Println("\n── 默认超时示例(WithDefaultTimeout)─────────────────────────────") + + // workers=1, maxConns=1:整个池子只有一个连接槽位,便于演示"占满后阻塞/超时" + pool, err := gobridge.NewPool(script, + gobridge.WithWorkers(1), + gobridge.WithMaxConns(1), + gobridge.WithDefaultTimeout(500*time.Millisecond), + ) + if err != nil { + log.Fatal(err) + } + defer pool.Close() + + // ── 示例1:ctx 未设置 deadline,池的默认超时自动生效 ───────────────────── + // sleep_seconds(2) 耗时 2s,远超 500ms 默认超时,Invoke 会在 500ms 后返回 + // context.DeadlineExceeded,而不是永久阻塞。 + _, err = gobridge.Invoke[string](context.Background(), pool, "sleep_seconds", 2.0) + fmt.Println("sleep_seconds(2.0)(无 ctx deadline,池默认超时 500ms)→ err =", err) + + // ── 示例2:调用方显式传入的 deadline 优先级更高,不会被默认超时覆盖 ─────── + ctx2, cancel2 := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel2() + msg, err := gobridge.Invoke[string](ctx2, pool, "sleep_seconds", 1.0) + fmt.Println("sleep_seconds(1.0)(显式 3s ctx)→", msg, err) + + // ── 示例3:流式输出模式下,超时只会静默关闭 channel,不返回 error, + // 需要调用方自行在 range 结束后检查 ctx.Err() 才能区分"正常结束"还是"被打断"。 + ctx3, cancel3 := context.WithTimeout(context.Background(), 350*time.Millisecond) + defer cancel3() + ch, err := gobridge.Invoke[chan int](ctx3, pool, "slow_range_gen", 1, 10, 100) + if err != nil { + log.Fatal(err) + } + fmt.Print("slow_range_gen(1, 10, 100ms)(350ms ctx 超时)=") + for v := range ch { + fmt.Print(" ", v) + } + fmt.Println() + if ctx3.Err() != nil { + fmt.Println("channel 提前关闭,ctx.Err() =", ctx3.Err()) + } +} + +func demoBlocking(script string) { + fmt.Println("\n── 连接池占满后阻塞示例(不设超时)───────────────────────────") + + // workers=2, maxConns=2:总容量只有 2*2=4 个连接槽位,便于演示"占满后新调用会怎样" + pool, err := gobridge.NewPool(script, + gobridge.WithWorkers(2), + gobridge.WithMaxConns(2), + ) + if err != nil { + log.Fatal(err) + } + defer pool.Close() + + const capacity = 4 // 2 workers * 2 maxConns + const taskSeconds = 2.0 // 每个占位任务耗时 2s + + // 并发发起 4 个耗时调用,占满全部连接槽位 + var wg sync.WaitGroup + wg.Add(capacity) + for i := 0; i < capacity; i++ { + go func(n int) { + defer wg.Done() + if _, err := gobridge.Invoke[string](context.Background(), pool, "sleep_seconds", taskSeconds); err != nil { + log.Printf("occupy task %d failed: %v", n, err) + } + }(i) + } + + // 留出时间让 4 个占位任务都真正进入 Python 侧 time.sleep(),占满全部槽位 + time.Sleep(200 * time.Millisecond) + fmt.Printf("池子已被 %d 个任务占满(各耗时 %.0fs),此时发起第 5 个调用(ctx 不设超时)……\n", capacity, taskSeconds) + + // 池子被占满时,新调用既不会立刻失败,也不会被跳过,而是排队阻塞等待连接释放。 + // 这里故意不设超时,用等待耗时证明它确实在阻塞排队,而不是瞬间返回。 + callStart := time.Now() + sum, err := gobridge.Invoke[int](context.Background(), pool, "add", 1, 2) + fmt.Printf("第 5 个调用等待 %v 后返回:add(1,2)=%d, err=%v\n", time.Since(callStart).Round(time.Millisecond), sum, err) + // 预期:等待时长接近"占满起 ~taskSeconds 秒后有槽位释放"的剩余时间, + // 说明调用被真正阻塞排队,直到某个占位任务结束、连接被释放后才继续执行。 + + wg.Wait() +} + func demoServer(ctx context.Context, script string) { fmt.Println("\n── Server 全双工示例 ─────────────────────────────────────────────") diff --git a/example/worker.py b/example/worker.py index 5d3e705..174f0f9 100644 --- a/example/worker.py +++ b/example/worker.py @@ -5,6 +5,7 @@ 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 @@ -120,6 +121,24 @@ def bytes_chunks(data: bytes, size: int): 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 全双工示例 ──────────────────────────────────────────────────────── diff --git a/pool.go b/pool.go index aa84121..6e735f9 100644 --- a/pool.go +++ b/pool.go @@ -9,6 +9,7 @@ import ( "reflect" "sync" "sync/atomic" + "time" ) // poolConfig 是进程池内部配置,通过 Option 函数填充 @@ -23,6 +24,7 @@ type poolConfig struct { stdout io.Writer stderr io.Writer handler any + defaultTimeout time.Duration } // Option 是 NewPool 的函数选项 @@ -74,6 +76,14 @@ func WithStderr(w io.Writer) Option { return func(c *poolConfig) { c.stderr = w } } +// WithDefaultTimeout 设置 Invoke 调用的默认超时时间(默认不启用)。 +// 仅当调用方传入的 ctx 未设置 deadline 时才会生效,用于防止调用方忘记设置超时导致 +// Invoke 无限期阻塞(例如所有 worker 连接被占满,或 Python 侧 handler 阻塞不返回)。 +// 调用方通过 context.WithTimeout 显式设置的超时优先级更高,不会被覆盖。 +func WithDefaultTimeout(d time.Duration) Option { + return func(c *poolConfig) { c.defaultTimeout = d } +} + // WithHandlers 注册 Go handler struct,其所有公开方法自动暴露给 Python 通过 call_go() 调用。 // // type MyService struct{} @@ -93,6 +103,7 @@ type Pool interface { acquire(ctx context.Context) (net.Conn, *worker, error) nextReqID() uint64 callbackDispatch(ctx context.Context, msg Message) (any, string) + defaultTimeout() time.Duration } // ── goHandler ──────────────────────────────────────────────────────────────── @@ -105,15 +116,15 @@ type goHandler struct { hasErr bool } - // ── pool(内部实现)───────────────────────────────────────────────────────── type pool struct { - workers []*worker - idx atomic.Uint64 - reqID atomic.Uint64 - mu sync.RWMutex - handlers map[string]goHandler + workers []*worker + idx atomic.Uint64 + reqID atomic.Uint64 + mu sync.RWMutex + handlers map[string]goHandler + defTimeout time.Duration } // NewPool 创建并启动进程池。 @@ -151,7 +162,7 @@ func NewPool(script string, opts ...Option) (Pool, error) { workers[i] = w } - p := &pool{workers: workers, handlers: make(map[string]goHandler)} + p := &pool{workers: workers, handlers: make(map[string]goHandler), defTimeout: cfg.defaultTimeout} if cfg.handler != nil { p.bindHandlers(cfg.handler) @@ -231,6 +242,10 @@ func (p *pool) nextReqID() uint64 { return p.reqID.Add(1) } +func (p *pool) defaultTimeout() time.Duration { + return p.defTimeout +} + func (p *pool) callbackDispatch(ctx context.Context, msg Message) (any, string) { p.mu.RLock() h, ok := p.handlers[msg.Method] diff --git a/pool_test.go b/pool_test.go new file mode 100644 index 0000000..63affb7 --- /dev/null +++ b/pool_test.go @@ -0,0 +1,143 @@ +package gobridge + +import ( + "context" + "errors" + "io" + "os/exec" + "path/filepath" + "runtime" + "sync" + "testing" + "time" +) + +// testWorkerScript 返回 example/worker.py 的绝对路径,供集成测试复用其中的 +// sleep_seconds / add 等函数;没有 python3 时跳过测试。 +func testWorkerScript(t *testing.T) string { + t.Helper() + if _, err := exec.LookPath("python3"); err != nil { + t.Skip("python3 not found in PATH, skipping integration test") + } + _, file, _, _ := runtime.Caller(0) + return filepath.Join(filepath.Dir(file), "example", "worker.py") +} + +func newTestPool(t *testing.T, opts ...Option) Pool { + t.Helper() + script := testWorkerScript(t) + // 注意:不要用 t.TempDir() 当 socketDir——macOS 的 AF_UNIX sun_path 只有 104 + // 字节长度限制,t.TempDir() 产出的路径很容易超限导致 bind 失败,所以沿用库默认的 /tmp。 + opts = append([]Option{ + WithStdout(io.Discard), + WithStderr(io.Discard), + }, opts...) + pool, err := NewPool(script, opts...) + if err != nil { + t.Fatalf("NewPool: %v", err) + } + t.Cleanup(pool.Close) + return pool +} + +// TestPoolExhaustionBlocks 验证:workers=2、maxConns=2(总容量 4)的池子被占满后, +// 新的 Invoke 调用会排队阻塞等待连接释放,而不是立刻失败或被跳过; +// 一旦有连接释放,新调用能正常拿到连接并执行。 +func TestPoolExhaustionBlocks(t *testing.T) { + pool := newTestPool(t, WithWorkers(2), WithMaxConns(2)) + + const capacity = 4 + const taskSeconds = 5.0 + + var wg sync.WaitGroup + wg.Add(capacity) + for i := 0; i < capacity; i++ { + go func() { + defer wg.Done() + if _, err := Invoke[string](context.Background(), pool, "sleep_seconds", taskSeconds); err != nil { + t.Errorf("occupy task failed: %v", err) + } + }() + } + + // 留出时间让 4 个占位任务都真正 acquire 到连接、进入 Python 侧 time.sleep(),占满全部 4 个槽位 + time.Sleep(300 * time.Millisecond) + + // 池子已被占满:用一个很短的超时验证新调用确实在排队等待连接(而非瞬间失败/跳过) + shortCtx, cancel := context.WithTimeout(context.Background(), 300*time.Millisecond) + defer cancel() + start := time.Now() + _, err := Invoke[string](shortCtx, pool, "add", 1, 2) + elapsed := time.Since(start) + + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("want context.DeadlineExceeded while pool exhausted, got %v", err) + } + if elapsed < 250*time.Millisecond { + t.Fatalf("acquire returned too fast (%v),说明没有真正排队等待连接释放", elapsed) + } + + // 等 4 个占位任务(各耗时 5s)全部结束、连接释放后,新调用应恢复正常 + wg.Wait() + sum, err := Invoke[int](context.Background(), pool, "add", 3, 4) + if err != nil { + t.Fatalf("Invoke after release: %v", err) + } + if sum != 7 { + t.Fatalf("want 7, got %d", sum) + } +} + +// TestDefaultTimeout 验证 WithDefaultTimeout:仅在 ctx 未设置 deadline 时生效, +// 调用方显式传入的 deadline 优先级更高,不会被覆盖。 +func TestDefaultTimeout(t *testing.T) { + pool := newTestPool(t, WithWorkers(1), WithMaxConns(1), WithDefaultTimeout(300*time.Millisecond)) + + // ctx 未设置 deadline:池的默认超时应自动生效,而不是永久阻塞 + start := time.Now() + _, err := Invoke[string](context.Background(), pool, "sleep_seconds", 5.0) + elapsed := time.Since(start) + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("want context.DeadlineExceeded from default timeout, got %v", err) + } + if elapsed > 2*time.Second { + t.Fatalf("default timeout did not kick in promptly, took %v", elapsed) + } + + // 显式传入的 deadline 优先级更高,不应被 300ms 的默认超时覆盖 + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + msg, err := Invoke[string](ctx, pool, "sleep_seconds", 1.0) + if err != nil { + t.Fatalf("explicit deadline should override default timeout: %v", err) + } + if msg != "slept 1s" { + t.Fatalf("want %q, got %q", "slept 1s", msg) + } +} + +// TestStreamTimeoutClosesChannelSilently 验证流式输出模式下,ctx 超时不会通过 error +// 返回,而是静默关闭已返回的 channel,需要调用方自行检查 ctx.Err()。 +func TestStreamTimeoutClosesChannelSilently(t *testing.T) { + pool := newTestPool(t, WithWorkers(1), WithMaxConns(1)) + + ctx, cancel := context.WithTimeout(context.Background(), 350*time.Millisecond) + defer cancel() + + ch, err := Invoke[chan int](ctx, pool, "slow_range_gen", 1, 10, 100) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + + var got []int + for v := range ch { + got = append(got, v) + } + + if len(got) == 0 || len(got) >= 9 { + t.Fatalf("want a partial stream (超时提前打断), got %v", got) + } + if ctx.Err() == nil { + t.Fatalf("want ctx.Err() != nil after timeout-truncated stream") + } +} diff --git a/session.go b/session.go index 62c3f1d..a0013d6 100644 --- a/session.go +++ b/session.go @@ -3,6 +3,7 @@ package gobridge import ( "context" "net" + "time" ) // singleWorkerPool 是 Pool 的包装,所有请求固定路由到同一个 worker 进程。 @@ -49,3 +50,7 @@ func (s *singleWorkerPool) nextReqID() uint64 { func (s *singleWorkerPool) callbackDispatch(ctx context.Context, msg Message) (any, string) { return s.pool.callbackDispatch(ctx, msg) } + +func (s *singleWorkerPool) defaultTimeout() time.Duration { + return s.pool.defaultTimeout() +}