gobridge/pool_test.go

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

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

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

	baseline, err := Invoke[int](ctx, pool, "count_threads")
	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")
	}
}