package proto_socket_test import ( "encoding/binary" "errors" "net" "strings" "sync" "testing" "time" "google.golang.org/protobuf/proto" toki "git.toki-labs.com/toki/proto-socket/go" "git.toki-labs.com/toki/proto-socket/go/packets" ) type fakeTransport struct { mu sync.Mutex packets []*packets.PacketBase } func (f *fakeTransport) WritePacket(base *packets.PacketBase) error { f.mu.Lock() defer f.mu.Unlock() f.packets = append(f.packets, base) return nil } func (f *fakeTransport) Close() error { return nil } func (f *fakeTransport) sent() []*packets.PacketBase { f.mu.Lock() defer f.mu.Unlock() return append([]*packets.PacketBase{}, f.packets...) } func testParserMap() toki.ParserMap { return toki.ParserMap{ toki.TypeNameOf(&packets.TestData{}): func(b []byte) (proto.Message, error) { m := &packets.TestData{} return m, proto.Unmarshal(b, m) }, } } func TestSendRequestTypeMismatch(t *testing.T) { transport := &fakeTransport{} communicator := toki.NewCommunicator(transport, testParserMap()) defer communicator.Close() errCh := make(chan error, 1) go func() { _, err := toki.SendRequestTyped[*packets.TestData, *packets.TestData]( communicator, &packets.TestData{Index: 1, Message: "hello"}, time.Second, ) errCh <- err }() var requestNonce int32 for deadline := time.Now().Add(time.Second); time.Now().Before(deadline); { sent := transport.sent() if len(sent) == 1 { requestNonce = sent[0].GetNonce() break } time.Sleep(time.Millisecond) } if requestNonce == 0 { t.Fatal("request packet was not sent") } hb, err := proto.Marshal(&packets.HeartBeat{}) if err != nil { t.Fatal(err) } communicator.OnReceivedData(toki.TypeNameOf(&packets.HeartBeat{}), hb, 0, requestNonce) err = <-errCh if err == nil { t.Fatal("expected response type mismatch error") } } func TestSendRequestTimeout(t *testing.T) { transport := &fakeTransport{} communicator := toki.NewCommunicator(transport, testParserMap()) defer communicator.Close() _, err := toki.SendRequestTyped[*packets.TestData, *packets.TestData]( communicator, &packets.TestData{Index: 1, Message: "no response"}, 25*time.Millisecond, ) if err == nil { t.Fatal("expected request timeout error") } if !strings.Contains(err.Error(), "timeout") { t.Fatalf("expected timeout error, got %v", err) } if len(transport.sent()) != 1 { t.Fatalf("expected one request packet to be sent, got %d", len(transport.sent())) } } func TestListenerAndRequestListenerConflict(t *testing.T) { transport := &fakeTransport{} communicator := toki.NewCommunicator(transport, testParserMap()) defer communicator.Close() toki.AddListenerTyped[*packets.TestData](communicator, func(*packets.TestData) {}) defer func() { if recover() == nil { t.Fatal("expected panic") } }() toki.AddRequestListenerTyped[*packets.TestData, *packets.TestData](communicator, func(req *packets.TestData) (*packets.TestData, error) { return req, nil }) } // TestInboundQueueFIFO: onReceivedData가 여러 메시지를 순서대로 전달한다. func TestInboundQueueFIFO(t *testing.T) { transport := &fakeTransport{} communicator := toki.NewCommunicator(transport, testParserMap()) defer communicator.Close() var mu sync.Mutex var received []int32 toki.AddListenerTyped[*packets.TestData](communicator, func(m *packets.TestData) { mu.Lock() received = append(received, m.GetIndex()) mu.Unlock() }) for i := int32(1); i <= 5; i++ { data, _ := proto.Marshal(&packets.TestData{Index: i}) communicator.OnReceivedData(toki.TypeNameOf(&packets.TestData{}), data, int32(i), 0) } deadline := time.Now().Add(time.Second) for time.Now().Before(deadline) { mu.Lock() n := len(received) mu.Unlock() if n == 5 { break } time.Sleep(time.Millisecond) } mu.Lock() defer mu.Unlock() if len(received) != 5 { t.Fatalf("expected 5 messages, got %d", len(received)) } for i, idx := range received { if idx != int32(i+1) { t.Fatalf("FIFO violated: index %d at position %d", idx, i) } } } // TestSlowRequestHandlerResponseOrder: 느린 request handler라도 자동 응답이 FIFO를 보존한다. func TestSlowRequestHandlerResponseOrder(t *testing.T) { transport := &fakeTransport{} communicator := toki.NewCommunicator(transport, testParserMap()) defer communicator.Close() toki.AddRequestListenerTyped[*packets.TestData, *packets.TestData](communicator, func(req *packets.TestData) (*packets.TestData, error) { // 첫 번째 요청은 100ms 지연 if req.GetIndex() == 1 { time.Sleep(100 * time.Millisecond) } return &packets.TestData{Index: req.GetIndex() + 100, Message: "echo"}, nil }) data1, _ := proto.Marshal(&packets.TestData{Index: 1, Message: "first"}) data2, _ := proto.Marshal(&packets.TestData{Index: 2, Message: "second"}) communicator.OnReceivedData(toki.TypeNameOf(&packets.TestData{}), data1, 10, 0) communicator.OnReceivedData(toki.TypeNameOf(&packets.TestData{}), data2, 20, 0) // 두 응답이 전송될 때까지 대기 deadline := time.Now().Add(2 * time.Second) for time.Now().Before(deadline) { if len(transport.sent()) >= 2 { break } time.Sleep(5 * time.Millisecond) } sent := transport.sent() if len(sent) < 2 { t.Fatalf("expected 2 responses, got %d", len(sent)) } // 첫 번째 응답이 nonce=10에 대응, 두 번째가 nonce=20에 대응 (FIFO 순서) if sent[0].GetResponseNonce() != 10 { t.Errorf("expected first response for nonce 10, got %d", sent[0].GetResponseNonce()) } if sent[1].GetResponseNonce() != 20 { t.Errorf("expected second response for nonce 20, got %d", sent[1].GetResponseNonce()) } } // TestCloseCanelsPendingRequests: Close 시 pending sendRequest가 ErrNotConnected로 완료된다. func TestCloseCancelsPendingRequests(t *testing.T) { transport := &fakeTransport{} communicator := toki.NewCommunicator(transport, testParserMap()) errCh := make(chan error, 1) go func() { _, err := toki.SendRequestTyped[*packets.TestData, *packets.TestData]( communicator, &packets.TestData{Index: 99, Message: "pending"}, 5*time.Second, ) errCh <- err }() // 요청 패킷이 전송될 때까지 대기 deadline := time.Now().Add(time.Second) for time.Now().Before(deadline) { if len(transport.sent()) == 1 { break } time.Sleep(time.Millisecond) } communicator.Close() select { case err := <-errCh: if !errors.Is(err, toki.ErrNotConnected) { t.Fatalf("expected ErrNotConnected, got %v", err) } case <-time.After(time.Second): t.Fatal("timed out waiting for pending request to be cancelled") } } // TestInboundBackpressureBlocksAndResumes: inbound channel이 가득 차면 추가 enqueue goroutine이 block되고, // handler가 release된 후 resume되는지 검증한다. func TestInboundBackpressureBlocksAndResumes(t *testing.T) { transport := &fakeTransport{} communicator := toki.NewCommunicator(transport, testParserMap()) defer communicator.Close() handlerBlock := make(chan struct{}) handlerRelease := make(chan struct{}) var blockOnce sync.Once toki.AddRequestListenerTyped[*packets.TestData, *packets.TestData](communicator, func(req *packets.TestData) (*packets.TestData, error) { blockOnce.Do(func() { close(handlerBlock) }) <-handlerRelease return &packets.TestData{Index: req.GetIndex() + 100, Message: "echo"}, nil }) data, _ := proto.Marshal(&packets.TestData{Index: 1}) // 1번째: 핸들러 내부에서 블록 communicator.EnqueueInbound(toki.TypeNameOf(&packets.TestData{}), data, 1, 0) // 핸들러 시작 대기 select { case <-handlerBlock: case <-time.After(time.Second): t.Fatal("timed out waiting for request handler to start") } // 2 ~ 65번째(총 64개)를 Enqueue하여 큐 버퍼(capacity=64)를 가득 채운다 for i := 2; i <= 65; i++ { communicator.EnqueueInbound(toki.TypeNameOf(&packets.TestData{}), data, int32(i), 0) } // 66번째 Enqueue 시도는 block되어야 함 enqueueDone := make(chan struct{}) go func() { communicator.EnqueueInbound(toki.TypeNameOf(&packets.TestData{}), data, 66, 0) close(enqueueDone) }() // block되어 대기 중인지 50ms 대기하며 확인 select { case <-enqueueDone: t.Fatal("expected EnqueueInbound to block on full queue, but it completed") case <-time.After(50 * time.Millisecond): // OK } // 핸들러 해제 close(handlerRelease) // 이제 66번째 enqueue가 완료되어야 함 select { case <-enqueueDone: // OK case <-time.After(2 * time.Second): t.Fatal("timed out waiting for blocked EnqueueInbound to resume") } } // TestGracefulCloseDrainsInFlightRequestHandler: in-flight handler가 실행 중일 때, // Close()가 호출되어도 남은 request가 정상적으로 drain(완료)되고 그 자동 응답이 기록되는지 검증한다. func TestGracefulCloseDrainsInFlightRequestHandler(t *testing.T) { transport := &fakeTransport{} communicator := toki.NewCommunicator(transport, testParserMap()) handlerBlock := make(chan struct{}) handlerRelease := make(chan struct{}) toki.AddRequestListenerTyped[*packets.TestData, *packets.TestData](communicator, func(req *packets.TestData) (*packets.TestData, error) { close(handlerBlock) <-handlerRelease return &packets.TestData{Index: req.GetIndex() + 100, Message: "graceful_response"}, nil }) data, _ := proto.Marshal(&packets.TestData{Index: 5}) communicator.EnqueueInbound(toki.TypeNameOf(&packets.TestData{}), data, 42, 0) // 핸들러가 실행 시작할 때까지 대기 select { case <-handlerBlock: case <-time.After(time.Second): t.Fatal("timed out waiting for request handler to start") } // 다른 고루틴에서 50ms 후 handlerRelease 하도록 설정 go func() { time.Sleep(50 * time.Millisecond) close(handlerRelease) }() // Close()를 비동기가 아닌 동기로 호출. Close()는 receiveDone이 닫힐 때까지 대기함 err := communicator.Close() if err != nil { t.Fatalf("unexpected error on Close: %v", err) } // Close()가 리턴했다는 것은 receiveLoop가 완전히 종료(drain 완료)되었다는 뜻임 if communicator.IsAlive() { t.Fatal("expected communicator to be inactive after Close") } // 자동 응답이 transport에 잘 들어갔는지 검증 sent := transport.sent() if len(sent) != 1 { t.Fatalf("expected 1 sent packet, got %d", len(sent)) } response := sent[0] if response.GetResponseNonce() != 42 { t.Errorf("expected response nonce 42, got %d", response.GetResponseNonce()) } var resData packets.TestData err = proto.Unmarshal(response.GetData(), &resData) if err != nil { t.Fatalf("failed to unmarshal response data: %v", err) } if resData.GetIndex() != 105 || resData.GetMessage() != "graceful_response" { t.Errorf("unexpected response content: index=%d message=%s", resData.GetIndex(), resData.GetMessage()) } } type fakeGatewayItem struct { seq int typeName string data []byte nonce int32 } func TestWorkerGatewayReorderAndDispatch(t *testing.T) { transport := &fakeTransport{} communicator := toki.NewCommunicator(transport, testParserMap()) defer communicator.Close() var mu sync.Mutex var received []int32 toki.AddListenerTyped[*packets.TestData](communicator, func(m *packets.TestData) { mu.Lock() received = append(received, m.GetIndex()) mu.Unlock() }) // Fake gateway: out-of-order 아이템들을 seq 오름차순으로 정렬 후 coordinator로 밀어넣음 items := []fakeGatewayItem{ {seq: 2, typeName: toki.TypeNameOf(&packets.TestData{}), data: func() []byte { d, _ := proto.Marshal(&packets.TestData{Index: 2}); return d }(), nonce: 2}, {seq: 1, typeName: toki.TypeNameOf(&packets.TestData{}), data: func() []byte { d, _ := proto.Marshal(&packets.TestData{Index: 1}); return d }(), nonce: 1}, } // Reorder by seq for i := 0; i < len(items); i++ { for j := i + 1; j < len(items); j++ { if items[i].seq > items[j].seq { items[i], items[j] = items[j], items[i] } } } // Dispatch to coordinator in reordered sequence for _, item := range items { communicator.OnReceivedData(item.typeName, item.data, item.nonce, 0) } deadline := time.Now().Add(time.Second) for time.Now().Before(deadline) { mu.Lock() n := len(received) mu.Unlock() if n == 2 { break } time.Sleep(time.Millisecond) } mu.Lock() defer mu.Unlock() if len(received) != 2 { t.Fatalf("expected 2 messages, got %d", len(received)) } if received[0] != 1 || received[1] != 2 { t.Fatalf("seq ordering was not preserved, got: %v", received) } } // TestInboundGatewayPreservesDispatchOrder: with the worker-pool gateway // attached, a burst of raw frames decoded in parallel is still dispatched to the // listener in input order, proving the coordinator keeps FIFO ownership. func TestInboundGatewayPreservesDispatchOrder(t *testing.T) { transport := &fakeTransport{} communicator := toki.NewCommunicator(transport, testParserMap()) communicator.EnableInboundGateway(4, 64) defer communicator.Close() var mu sync.Mutex var received []int32 toki.AddListenerTyped[*packets.TestData](communicator, func(m *packets.TestData) { mu.Lock() received = append(received, m.GetIndex()) mu.Unlock() }) const total = 200 typeName := toki.TypeNameOf(&packets.TestData{}) for i := int32(1); i <= total; i++ { data, _ := proto.Marshal(&packets.TestData{Index: i}) frame, err := proto.Marshal(&packets.PacketBase{TypeName: typeName, Nonce: i, Data: data}) if err != nil { t.Fatal(err) } communicator.OnReceivedFrame(frame) } deadline := time.Now().Add(3 * time.Second) for time.Now().Before(deadline) { mu.Lock() n := len(received) mu.Unlock() if n == total { break } time.Sleep(time.Millisecond) } mu.Lock() defer mu.Unlock() if len(received) != total { t.Fatalf("expected %d messages, got %d", total, len(received)) } for i, idx := range received { if idx != int32(i+1) { t.Fatalf("gateway dispatch order violated at position %d: got index %d", i, idx) } } } // frameTCP builds a TCP wire frame: 4-byte big-endian length prefix + PacketBase. func frameTCP(t *testing.T, base *packets.PacketBase) []byte { t.Helper() b, err := proto.Marshal(base) if err != nil { t.Fatal(err) } out := make([]byte, 4+len(b)) binary.BigEndian.PutUint32(out[:4], uint32(len(b))) copy(out[4:], b) return out } // countingGateway is a test InboundGateway that records every Submit and then // forwards the decoded frame to the supplied sink. It lets a transport read-path // test prove the read loop actually routes raw frames through the gateway rather // than decoding inline and bypassing it. type countingGateway struct { onSubmit func(toki.InboundFrame) } func (g *countingGateway) Submit(frame toki.InboundFrame) { g.onSubmit(frame) } func (g *countingGateway) Close() {} // TestTcpReadLoopRoutesThroughGateway: a real TcpClient.readLoop must hand raw // frames to the attached gateway. The counting gateway fails the test if the // read loop decodes inline and bypasses the gateway (the original wiring gap). func TestTcpReadLoopRoutesThroughGateway(t *testing.T) { serverConn, clientConn := net.Pipe() // intervalSec=0 disables heartbeats, so the client never writes to the pipe. client := toki.NewTcpClient(clientConn, 0, 0, testParserMap()) defer client.Close() var mu sync.Mutex var received []int32 submitted := 0 gw := &countingGateway{onSubmit: func(frame toki.InboundFrame) { mu.Lock() submitted++ mu.Unlock() base := &packets.PacketBase{} if err := proto.Unmarshal(frame.Bytes, base); err != nil { return } client.OnReceivedData(base.GetTypeName(), base.GetData(), base.GetNonce(), base.GetResponseNonce()) }} client.AttachInboundGateway(gw) toki.AddListenerTyped[*packets.TestData](&client.Communicator, func(m *packets.TestData) { mu.Lock() received = append(received, m.GetIndex()) mu.Unlock() }) const total = 20 typeName := toki.TypeNameOf(&packets.TestData{}) frames := make([][]byte, 0, total) for i := int32(1); i <= total; i++ { data, err := proto.Marshal(&packets.TestData{Index: i}) if err != nil { t.Fatal(err) } frames = append(frames, frameTCP(t, &packets.PacketBase{TypeName: typeName, Nonce: i, Data: data})) } go func() { for _, f := range frames { if _, err := serverConn.Write(f); err != nil { return } } }() deadline := time.Now().Add(3 * time.Second) for time.Now().Before(deadline) { mu.Lock() n := len(received) mu.Unlock() if n == total { break } time.Sleep(time.Millisecond) } mu.Lock() defer mu.Unlock() if submitted != total { t.Fatalf("read loop bypassed gateway: submitted=%d, want %d", submitted, total) } if len(received) != total { t.Fatalf("expected %d messages, got %d", total, len(received)) } for i, idx := range received { if idx != int32(i+1) { t.Fatalf("dispatch order violated at position %d: got index %d", i, idx) } } } // TestTcpReadLoopWorkerGatewayPreservesOrder: a real TcpClient.readLoop with the // production worker-pool gateway (EnableInboundGateway) still dispatches frames // to the listener in input order despite parallel decode. func TestTcpReadLoopWorkerGatewayPreservesOrder(t *testing.T) { serverConn, clientConn := net.Pipe() client := toki.NewTcpClient(clientConn, 0, 0, testParserMap()) defer client.Close() client.EnableInboundGateway(4, 64) var mu sync.Mutex var received []int32 toki.AddListenerTyped[*packets.TestData](&client.Communicator, func(m *packets.TestData) { mu.Lock() received = append(received, m.GetIndex()) mu.Unlock() }) const total = 100 typeName := toki.TypeNameOf(&packets.TestData{}) frames := make([][]byte, 0, total) for i := int32(1); i <= total; i++ { data, err := proto.Marshal(&packets.TestData{Index: i}) if err != nil { t.Fatal(err) } frames = append(frames, frameTCP(t, &packets.PacketBase{TypeName: typeName, Nonce: i, Data: data})) } go func() { for _, f := range frames { if _, err := serverConn.Write(f); err != nil { return } } }() deadline := time.Now().Add(3 * time.Second) for time.Now().Before(deadline) { mu.Lock() n := len(received) mu.Unlock() if n == total { break } time.Sleep(time.Millisecond) } mu.Lock() defer mu.Unlock() if len(received) != total { t.Fatalf("expected %d messages, got %d", total, len(received)) } for i, idx := range received { if idx != int32(i+1) { t.Fatalf("dispatch order violated at position %d: got index %d", i, idx) } } } // TestTcpReadLoopParseErrorDisconnects: with no gateway, a malformed frame in // the real read loop must still trigger the parse-error disconnect. func TestTcpReadLoopParseErrorDisconnects(t *testing.T) { serverConn, clientConn := net.Pipe() client := toki.NewTcpClient(clientConn, 0, 0, testParserMap()) defer client.Close() // length=3 prefix followed by bytes that are not valid protobuf wire format. bad := []byte{0, 0, 0, 3, 0xff, 0xff, 0xff} go func() { _, _ = serverConn.Write(bad) }() deadline := time.Now().Add(2 * time.Second) for time.Now().Before(deadline) { if r := client.DisconnectInfo().Reason; r != toki.DisconnectReasonUnknown { break } time.Sleep(time.Millisecond) } info := client.DisconnectInfo() if info.Reason != toki.DisconnectReasonTCPPacketParse { t.Fatalf("expected %s disconnect, got %q", toki.DisconnectReasonTCPPacketParse, info.Reason) } } // TestTcpReadLoopGatewayParseErrorDisconnects: with the gateway attached, a // malformed frame decoded off the read loop must still surface as the same // parse-error disconnect via the gateway error path. func TestTcpReadLoopGatewayParseErrorDisconnects(t *testing.T) { serverConn, clientConn := net.Pipe() client := toki.NewTcpClient(clientConn, 0, 0, testParserMap()) defer client.Close() client.EnableInboundGateway(2, 16) // length=3 prefix followed by bytes that are not valid protobuf wire format. bad := []byte{0, 0, 0, 3, 0xff, 0xff, 0xff} go func() { _, _ = serverConn.Write(bad) }() deadline := time.Now().Add(2 * time.Second) for time.Now().Before(deadline) { if r := client.DisconnectInfo().Reason; r != toki.DisconnectReasonUnknown { break } time.Sleep(time.Millisecond) } info := client.DisconnectInfo() if info.Reason != toki.DisconnectReasonTCPPacketParse { t.Fatalf("expected %s disconnect via gateway error path, got %q", toki.DisconnectReasonTCPPacketParse, info.Reason) } } func TestAliasCollisionOnInit(t *testing.T) { defer func() { if recover() == nil { t.Fatal("expected panic on alias collision") } }() badParserMap := toki.ParserMap{ "a.b.TestData": func(b []byte) (proto.Message, error) { return &packets.TestData{}, nil }, "x.y.TestData": func(b []byte) (proto.Message, error) { return &packets.TestData{}, nil }, } transport := &fakeTransport{} _ = toki.NewCommunicator(transport, badParserMap) } func TestLegacyAliasRouting(t *testing.T) { transport := &fakeTransport{} communicator := toki.NewCommunicator(transport, testParserMap()) defer communicator.Close() // 1. Test Listener routing with legacy simple name var listenerCalled bool var receivedMsg proto.Message toki.AddListenerTyped[*packets.TestData](communicator, func(m *packets.TestData) { listenerCalled = true receivedMsg = m }) testMsgBytes, _ := proto.Marshal(&packets.TestData{Index: 42, Message: "test"}) communicator.OnReceivedData("TestData", testMsgBytes, 1, 0) deadline := time.Now().Add(time.Second) for time.Now().Before(deadline) { if listenerCalled { break } time.Sleep(time.Millisecond) } if !listenerCalled { t.Fatal("listener was not called with legacy simple name") } if receivedMsg.(*packets.TestData).Index != 42 { t.Fatalf("expected Index 42, got %d", receivedMsg.(*packets.TestData).Index) } // 2. Test Pending request matching with legacy simple name errCh := make(chan error, 1) var responseMsg proto.Message go func() { res, err := toki.SendRequestTyped[*packets.TestData, *packets.TestData]( communicator, &packets.TestData{Index: 1, Message: "req"}, time.Second, ) responseMsg = res errCh <- err }() var requestNonce int32 for deadline = time.Now().Add(time.Second); time.Now().Before(deadline); { sent := transport.sent() if len(sent) == 1 { requestNonce = sent[0].GetNonce() break } time.Sleep(time.Millisecond) } if requestNonce == 0 { t.Fatal("request packet was not sent") } communicator.OnReceivedData("TestData", testMsgBytes, 0, requestNonce) err := <-errCh if err != nil { t.Fatalf("expected successful response matching, got error: %v", err) } if responseMsg.(*packets.TestData).Index != 42 { t.Fatalf("expected response Index 42, got %d", responseMsg.(*packets.TestData).Index) } } func TestRequestHandlerAliasSimpleNameInbound(t *testing.T) { transport := &fakeTransport{} communicator := toki.NewCommunicator(transport, testParserMap()) defer communicator.Close() toki.AddRequestListenerTyped[*packets.TestData, *packets.TestData](communicator, func(req *packets.TestData) (*packets.TestData, error) { return &packets.TestData{Index: req.GetIndex() + 1, Message: "echo"}, nil }) reqBytes, _ := proto.Marshal(&packets.TestData{Index: 7, Message: "req"}) communicator.OnReceivedData("TestData", reqBytes, 42, 0) deadline := time.Now().Add(time.Second) for time.Now().Before(deadline) { if len(transport.sent()) >= 1 { break } time.Sleep(time.Millisecond) } sent := transport.sent() if len(sent) != 1 { t.Fatalf("expected 1 response packet, got %d", len(sent)) } if sent[0].GetResponseNonce() != 42 { t.Errorf("expected responseNonce 42, got %d", sent[0].GetResponseNonce()) } decoded := &packets.TestData{} if err := proto.Unmarshal(sent[0].GetData(), decoded); err != nil { t.Fatalf("failed to decode response payload: %v", err) } if decoded.GetIndex() != 8 { t.Errorf("expected response index 8, got %d", decoded.GetIndex()) } }