gobridge/pool_test.go

package gobridge

import (
	"context"
	"errors"
	"io"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"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
}

// stableThreadCount 轮询 count_threads()，直到连续两次读数一致再返回，
// 用于绕开 pool 刚创建时 Python 侧连接处理线程还没必然全部起稳的竞态。
func stableThreadCount(t *testing.T, ctx context.Context, pool Pool) (int, error) {
	t.Helper()
	prev := -1
	deadline := time.Now().Add(3 * time.Second)
	for time.Now().Before(deadline) {
		n, err := Invoke[int](ctx, pool, "count_threads")
		if err != nil {
			return 0, err
		}
		if n == prev {
			return n, nil
		}
		prev = n
		time.Sleep(50 * time.Millisecond)
	}
	return prev, nil
}

// 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)
	}
}

// TestWithEnv 验证 WithEnv 设置的环境变量真的被传进了子进程，
// 而不是被静默忽略（回归：曾经上层 framework 的配置解析漏了 env 字段，
// 导致 yaml 里配的 env 从没真正生效过）。
func TestWithEnv(t *testing.T) {
	pool := newTestPool(t, WithWorkers(1), WithMaxConns(1),
		WithEnv("GOBRIDGE_TEST_VAR=hello-from-go"))

	got, err := Invoke[string](context.Background(), pool, "get_env", "GOBRIDGE_TEST_VAR")
	if err != nil {
		t.Fatalf("Invoke get_env: %v", err)
	}
	if got != "hello-from-go" {
		t.Fatalf("want %q, got %q", "hello-from-go", got)
	}

	// 没有通过 WithEnv 设置的变量应该读不到，避免测试因为宿主机碰巧存在
	// 同名环境变量而误判通过。
	got, err = Invoke[string](context.Background(), pool, "get_env", "GOBRIDGE_TEST_VAR_UNSET")
	if err != nil {
		t.Fatalf("Invoke get_env (unset): %v", err)
	}
	if got != "" {
		t.Fatalf("want empty string for unset var, got %q", got)
	}
}

// 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)
	}
}

// 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) != 2 {
		t.Fatalf("want 2 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)
	}
}

// TestStreamError 验证 WithStreamErrors：opt-in 之后，流式调用执行过程中因为 Python
// 侧真实抛出的异常提前结束时，能通过它返回的 streamErr 函数查到具体错误（不需要再传
// ctx）；正常跑完的流式调用查询结果是 nil。
//
// 这个方案没有全局登记表——错误记录挂在 WithStreamErrors 返回的那个 ctx 的对象图里，
// 调用方不再引用这个 ctx（和对应的 ch）时，整条链会被 GC 自然回收，不需要任何显式
// 清理逻辑，也不依赖 ctx.Done()，context.Background() 一样能正常释放。
func TestStreamError(t *testing.T) {
	pool := newTestPool(t, WithWorkers(1), WithMaxConns(2))

	t.Run("opt-in 后能查到 Python 侧真实抛出的异常", func(t *testing.T) {
		ctx, streamErr := WithStreamErrors(context.Background())
		// stream_then_raise 是 example/worker.py 里真实的 Python 函数：
		// 在还没 yield 完 n 个数之前，执行过程中就 raise ValueError("boom: ...")，
		// 不是模拟出来的错误，也不是等全部 yield 完才失败。
		ch, err := Invoke[chan int](ctx, pool, "stream_then_raise", 3)
		if err != nil {
			t.Fatalf("Invoke: %v", err)
		}
		var got []int
		for v := range ch {
			got = append(got, v)
		}
		if len(got) != 2 {
			t.Fatalf("want 2 items before the raise, got %v", got)
		}
		err = streamErr(ch)
		if err == nil {
			t.Fatal("want non-nil error after Python 侧执行过程中抛异常")
		}
		if !strings.Contains(err.Error(), "boom") {
			t.Fatalf("want error containing %q, got %v", "boom", err)
		}
		t.Logf("streamErr 查到的错误: %v", err)
	})

	t.Run("正常跑完不会有错误", func(t *testing.T) {
		ctx, streamErr := WithStreamErrors(context.Background())
		ch, err := Invoke[chan int](ctx, pool, "range_gen", 1, 4)
		if err != nil {
			t.Fatalf("Invoke: %v", err)
		}
		var got []int
		for v := range ch {
			got = append(got, v)
		}
		if len(got) != 3 {
			t.Fatalf("want 3 items, got %v", got)
		}
		if err := streamErr(ch); err != nil {
			t.Fatalf("want nil，got %v", err)
		}
	})

	t.Run("再包一层 ctx（比如加超时）不会丢失 WithStreamErrors 记录的信息", func(t *testing.T) {
		ctx, streamErr := WithStreamErrors(context.Background())
		ctx, cancel := context.WithTimeout(ctx, 5*time.Second) // WithStreamErrors 之后又包了一层
		defer cancel()
		ch, err := Invoke[chan int](ctx, pool, "stream_then_raise", 3)
		if err != nil {
			t.Fatalf("Invoke: %v", err)
		}
		for range ch {
		}
		if err := streamErr(ch); err == nil {
			t.Fatal("再包一层 WithTimeout 之后应该还能查到 WithStreamErrors 记录的错误")
		}
	})
}

// TestStreamOutIgnoresDefaultTimeout 验证 WithDefaultTimeout 不会套用到流式输出/
// 双向流调用上：流式聊天这类响应可能持续很久，如果套用全局默认超时会在正常输出
// 过程中把它腰斩。没有显式传入 deadline 时，流式调用应该完全不受池子默认超时影响，
// 只受调用方显式设置的 ctx 约束。
func TestStreamOutIgnoresDefaultTimeout(t *testing.T) {
	pool := newTestPool(t, WithWorkers(1), WithMaxConns(1), WithDefaultTimeout(200*time.Millisecond))

	// slow_range_gen(1, 10, 100) 总耗时 ~900ms，远超 200ms 的默认超时。
	// 用 context.Background()（无 deadline）发起调用：如果默认超时被错误地套用到
	// 流式输出上，channel 会在 ~200ms 提前关闭，读不满 9 个元素。
	ch, err := Invoke[chan int](context.Background(), 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) != 9 {
		t.Fatalf("want 9 items (WithDefaultTimeout 不应影响流式输出), got %v", got)
	}
}

// TestStreamInputCancelDoesNotLeakThread 验证流式输入模式下，handler 线程阻塞在
// _ChunkIter.__next__ → chunk_q.get() 等待下一个输入块时，如果 ctx 取消导致连接被
// Go 关闭，Python 侧能否真正把这个线程唤醒退出，而不是永久卡住。
//
// _ConnMux._reader 在连接关闭时只处理了 call_q（推 None）和 _active_tids（注入
// InterruptedError），完全没有触碰 chunk_q；而 handler 线程当前正阻塞在
// chunk_q.get() 里——问题是 PyThreadState_SetAsyncExc 能不能真正打断一个阻塞在
// queue.Queue.get()（无 timeout）上的线程。如果不能，这个线程会永久泄漏：既不会
// 收到任何新数据，也不会因为异常注入而退出，因为 chunk_q 以后也不会再有任何写入。
func TestStreamInputCancelDoesNotLeakThread(t *testing.T) {
	pool := newTestPool(t, WithWorkers(1), WithMaxConns(2))
	ctx := context.Background()

	// 刚创建的 pool，Python 那边给每条预建连接起的 _handle_conn/_reader 线程
	// 还没必然全部起稳（server.accept() 是异步接受的），第一次量出来的线程数
	// 可能偏低导致误判。轮询到连续两次读数一致再当作稳定的 baseline。
	baseline, err := stableThreadCount(t, ctx, pool)
	if err != nil {
		t.Fatalf("Invoke count_threads (baseline): %v", err)
	}

	// inputCh 故意不发送任何数据也不关闭，让 sum_stream 的 handler 线程卡在
	// _ChunkIter.__next__ → chunk_q.get() 上等待第一个 chunk。
	inputCh := make(chan int)
	shortCtx, cancel := context.WithTimeout(ctx, 300*time.Millisecond)
	defer cancel()
	_, err = Invoke[int](shortCtx, pool, "sum_stream", inputCh)
	if !errors.Is(err, context.DeadlineExceeded) {
		t.Fatalf("want context.DeadlineExceeded, got %v", err)
	}

	// 给 Python 侧一点时间处理连接关闭 / InterruptedError 注入
	time.Sleep(500 * time.Millisecond)

	after, err := Invoke[int](ctx, pool, "count_threads")
	if err != nil {
		t.Fatalf("Invoke count_threads (after): %v", err)
	}

	t.Logf("thread count: baseline=%d after=%d", baseline, after)
	if after > baseline {
		t.Fatalf("线程泄漏：baseline=%d after=%d，卡在 chunk_q.get() 上的 handler 线程没有被 InterruptedError 唤醒退出", baseline, after)
	}
}

// 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")
	}
}