|
| 1 | +package agent |
| 2 | + |
| 3 | +import ( |
| 4 | + "context" |
| 5 | + "net" |
| 6 | + "testing" |
| 7 | + "time" |
| 8 | + |
| 9 | + agentlocal "github.com/cuihairu/croupier/internal/platform/agentlocal" |
| 10 | + transportcore "github.com/cuihairu/croupier/internal/transport" |
| 11 | + tcptr "github.com/cuihairu/croupier/internal/transport/tcp" |
| 12 | + sdkv1 "github.com/cuihairu/croupier/pkg/pb/croupier/sdk/v1" |
| 13 | + "github.com/cuihairu/croupier/pkg/protocol" |
| 14 | + "github.com/stretchr/testify/assert" |
| 15 | + "github.com/stretchr/testify/require" |
| 16 | + "google.golang.org/protobuf/proto" |
| 17 | +) |
| 18 | + |
| 19 | +// providerCallDeadline：metadata timeout_ms clamp 与默认值语义。 |
| 20 | +func TestProviderCallDeadlineSemantics(t *testing.T) { |
| 21 | + h := NewLocalHandler(nil, "/tmp", "agent-1", nil) |
| 22 | + |
| 23 | + // 默认 15s（对齐 Server 派发层，替代旧硬编码 10s） |
| 24 | + assert.Equal(t, 15*time.Second, h.providerCallDeadline(nil)) |
| 25 | + assert.Equal(t, 15*time.Second, h.providerCallDeadline(map[string]string{"timeout_ms": "garbage"})) |
| 26 | + |
| 27 | + // 请求声明更小值 → 生效 |
| 28 | + assert.Equal(t, 3*time.Second, h.providerCallDeadline(map[string]string{"timeout_ms": "3000"})) |
| 29 | + // 低于 1s 提到 1s；高于配置默认 → clamp 回默认 |
| 30 | + assert.Equal(t, time.Second, h.providerCallDeadline(map[string]string{"timeout_ms": "50"})) |
| 31 | + assert.Equal(t, 15*time.Second, h.providerCallDeadline(map[string]string{"timeout_ms": "999999"})) |
| 32 | + |
| 33 | + // 配置自定义默认：请求声明可小于配置，不可超过 |
| 34 | + h.SetProviderCallTimeout(2 * time.Second) |
| 35 | + assert.Equal(t, 2*time.Second, h.providerCallDeadline(nil)) |
| 36 | + assert.Equal(t, 1500*time.Millisecond, h.providerCallDeadline(map[string]string{"timeout_ms": "1500"})) |
| 37 | + assert.Equal(t, 2*time.Second, h.providerCallDeadline(map[string]string{"timeout_ms": "30000"})) |
| 38 | + |
| 39 | + // 非法配置回落默认；超上限截断 |
| 40 | + h.SetProviderCallTimeout(-1) |
| 41 | + assert.Equal(t, 15*time.Second, h.providerCallDeadline(nil)) |
| 42 | + h.SetProviderCallTimeout(5 * time.Minute) |
| 43 | + assert.Equal(t, 60*time.Second, h.providerCallDeadline(nil)) |
| 44 | +} |
| 45 | + |
| 46 | +// 端到端：慢 provider 在 metadata 声明的预算内未完成 → 调用方拿到超时错误。 |
| 47 | +func TestHandleInvokeRespectsMetadataTimeout(t *testing.T) { |
| 48 | + listener, err := NewTCPLocalListener(&TCPLocalListenerConfig{Address: "127.0.0.1:0"}, nil, nil) |
| 49 | + require.NoError(t, err) |
| 50 | + t.Cleanup(func() { require.NoError(t, listener.Close()) }) |
| 51 | + |
| 52 | + ctx, cancel := context.WithCancel(context.Background()) |
| 53 | + t.Cleanup(cancel) |
| 54 | + go func() { _ = listener.Serve(ctx) }() |
| 55 | + |
| 56 | + conn, err := net.Dial("tcp", listener.Addr()) |
| 57 | + require.NoError(t, err) |
| 58 | + provider := tcptr.NewMuxConn(conn, nil, transportcore.HandlerFunc( |
| 59 | + func(_ context.Context, msgID uint32, _ uint32, body []byte) ([]byte, error) { |
| 60 | + if msgID == protocol.MsgInvokeRequest { |
| 61 | + req := &sdkv1.InvokeRequest{} |
| 62 | + require.NoError(t, proto.Unmarshal(body, req)) |
| 63 | + // 慢于声明的 1200ms 预算 |
| 64 | + time.Sleep(2500 * time.Millisecond) |
| 65 | + return proto.Marshal(&sdkv1.InvokeResponse{Payload: []byte(`{}`)}) |
| 66 | + } |
| 67 | + return nil, nil |
| 68 | + })) |
| 69 | + go func() { _ = provider.Run(ctx) }() |
| 70 | + t.Cleanup(func() { _ = provider.Close() }) |
| 71 | + |
| 72 | + connectBody, err := proto.Marshal(&sdkv1.ProviderConnectRequest{ |
| 73 | + ServiceId: "slow-provider", |
| 74 | + Version: "1.0.0", |
| 75 | + Functions: []*sdkv1.ProviderFunctionDescriptor{{Id: "slow.fn", Version: "1.0.0"}}, |
| 76 | + }) |
| 77 | + require.NoError(t, err) |
| 78 | + _, _, err = provider.Call(ctx, protocol.MsgProviderConnectRequest, connectBody) |
| 79 | + require.NoError(t, err) |
| 80 | + |
| 81 | + session, ok := listener.SessionStore().GetByServiceID("slow-provider") |
| 82 | + require.True(t, ok) |
| 83 | + require.NotNil(t, session.Conn()) |
| 84 | + |
| 85 | + // 生产装配形态：session 连接回调把函数注册进 local store（与 |
| 86 | + // app.go SetOnConnect 相同），pickInstance 才能路由到该函数。 |
| 87 | + store := agentlocal.NewLocalStore() |
| 88 | + listener.SetOnConnect(func(sess *ProviderSession) { |
| 89 | + store.Register(sess.SessionID, sess.ServiceID, sess.Conn().RemoteAddr(), sess.Version, sess.Functions, nil) |
| 90 | + }) |
| 91 | + // 重新握手触发 onConnect（上面的注册发生在 SetOnConnect 之前） |
| 92 | + conn2, err := net.Dial("tcp", listener.Addr()) |
| 93 | + require.NoError(t, err) |
| 94 | + provider2 := tcptr.NewMuxConn(conn2, nil, transportcore.HandlerFunc( |
| 95 | + func(_ context.Context, msgID uint32, _ uint32, _ []byte) ([]byte, error) { |
| 96 | + if msgID == protocol.MsgInvokeRequest { |
| 97 | + time.Sleep(2500 * time.Millisecond) |
| 98 | + return proto.Marshal(&sdkv1.InvokeResponse{Payload: []byte(`{}`)}) |
| 99 | + } |
| 100 | + return nil, nil |
| 101 | + })) |
| 102 | + go func() { _ = provider2.Run(ctx) }() |
| 103 | + t.Cleanup(func() { _ = provider2.Close() }) |
| 104 | + _, _, err = provider2.Call(ctx, protocol.MsgProviderConnectRequest, connectBody) |
| 105 | + require.NoError(t, err) |
| 106 | + |
| 107 | + h := NewLocalHandler(store, "/tmp", "agent-1", nil) |
| 108 | + h.SetProviderSessionStore(listener.SessionStore()) |
| 109 | + |
| 110 | + // 1) 声明 1200ms 预算，provider 睡 2.5s → 必须在预算内超时 |
| 111 | + invokeBody, err := proto.Marshal(&sdkv1.InvokeRequest{ |
| 112 | + FunctionId: "slow.fn", |
| 113 | + Payload: []byte(`{}`), |
| 114 | + Metadata: map[string]string{"timeout_ms": "1200"}, |
| 115 | + }) |
| 116 | + require.NoError(t, err) |
| 117 | + |
| 118 | + start := time.Now() |
| 119 | + _, err = h.Handle(ctx, protocol.MsgInvokeRequest, 1, invokeBody) |
| 120 | + require.Error(t, err, "调用必须超时") |
| 121 | + elapsed := time.Since(start) |
| 122 | + assert.Less(t, elapsed, 2300*time.Millisecond, "应在声明的 ~1.2s 预算内失败") |
| 123 | + assert.GreaterOrEqual(t, elapsed, 1100*time.Millisecond, "不应早于预算误伤") |
| 124 | + |
| 125 | + // 2) 不声明预算 → 走默认 15s，provider 2.5s 正常返回成功 |
| 126 | + invokeBody2, err := proto.Marshal(&sdkv1.InvokeRequest{ |
| 127 | + FunctionId: "slow.fn", |
| 128 | + Payload: []byte(`{}`), |
| 129 | + }) |
| 130 | + require.NoError(t, err) |
| 131 | + resp, err := h.Handle(ctx, protocol.MsgInvokeRequest, 2, invokeBody2) |
| 132 | + require.NoError(t, err, "默认预算下慢 provider 应成功") |
| 133 | + assert.NotEmpty(t, resp) |
| 134 | +} |
0 commit comments