From a0a8b7004a6deff0e06e8515c2b55b4609270526 Mon Sep 17 00:00:00 2001 From: iteye Date: Fri, 10 Jul 2026 11:17:39 +0800 Subject: [PATCH] implement RpcConn and XshardConn compatibility layer --- qkc/cluster/slave/compat_test.go | 319 ++++++++ qkc/cluster/slave/connection.go | 544 +++++++++++++ qkc/cluster/slave/errors.go | 17 + .../slave/testdata/pyproto/__init__.py | 0 qkc/cluster/slave/testdata/pyproto/frame.py | 35 + .../slave/testdata/pyproto/messages.py | 73 ++ qkc/cluster/slave/testdata/pyproto/peer.py | 127 ++++ qkc/cluster/slave/xshard_conn.go | 230 ++++++ qkc/cluster/slave/xshard_pool.go | 321 ++++++++ qkc/cluster/slave/xshard_test.go | 712 ++++++++++++++++++ 10 files changed, 2378 insertions(+) create mode 100644 qkc/cluster/slave/compat_test.go create mode 100644 qkc/cluster/slave/connection.go create mode 100644 qkc/cluster/slave/errors.go create mode 100644 qkc/cluster/slave/testdata/pyproto/__init__.py create mode 100644 qkc/cluster/slave/testdata/pyproto/frame.py create mode 100644 qkc/cluster/slave/testdata/pyproto/messages.py create mode 100644 qkc/cluster/slave/testdata/pyproto/peer.py create mode 100644 qkc/cluster/slave/xshard_conn.go create mode 100644 qkc/cluster/slave/xshard_pool.go create mode 100644 qkc/cluster/slave/xshard_test.go diff --git a/qkc/cluster/slave/compat_test.go b/qkc/cluster/slave/compat_test.go new file mode 100644 index 0000000000..6907d4e094 --- /dev/null +++ b/qkc/cluster/slave/compat_test.go @@ -0,0 +1,319 @@ +// Copyright 2026-2027, QuarkChain. + +package slave + +import ( + "bufio" + "context" + "fmt" + "net" + "os" + "os/exec" + "path/filepath" + "runtime" + "strings" + "testing" + "time" + + "github.com/ethereum/go-ethereum/log" +) + +// startPythonPeer starts a Python protocol peer subprocess and returns the +// TCP port and a cleanup function. The peer listens on a random port (port=0) +// and prints "PORT:" to stdout when ready. +func startPythonPeer(t *testing.T, extraArgs ...string) (int, func()) { + t.Helper() + + _, filename, _, ok := runtime.Caller(0) + if !ok { + t.Fatal("cannot get caller path") + } + pyScript := filepath.Join(filepath.Dir(filename), "testdata", "pyproto", "peer.py") + + if _, err := exec.LookPath("python3"); err != nil { + t.Skip("python3 not found in PATH") + } + if _, err := os.Stat(pyScript); err != nil { + t.Skipf("peer.py not found at %s", pyScript) + } + + args := []string{pyScript, "--port", "0", "--id", "py", "--shards", "1"} + args = append(args, extraArgs...) + + cmd := exec.Command("python3", args...) + stdout, err := cmd.StdoutPipe() + if err != nil { + t.Fatalf("stdout pipe: %v", err) + } + cmd.Stderr = os.Stderr + + if err := cmd.Start(); err != nil { + t.Fatalf("start python peer: %v", err) + } + + // Read PORT: line from stdout. + portCh := make(chan int, 1) + errCh := make(chan error, 1) + go func() { + scanner := bufio.NewScanner(stdout) + for scanner.Scan() { + line := scanner.Text() + if strings.HasPrefix(line, "PORT:") { + var port int + if _, err := fmt.Sscanf(line, "PORT:%d", &port); err == nil { + portCh <- port + return + } + } + } + errCh <- scanner.Err() + }() + + var port int + select { + case port = <-portCh: + case err := <-errCh: + cmd.Process.Kill() + cmd.Wait() + t.Fatalf("read port from python peer: %v", err) + case <-time.After(5 * time.Second): + cmd.Process.Kill() + cmd.Wait() + t.Fatal("timeout waiting for python peer port") + } + + cleanup := func() { + cmd.Process.Kill() + cmd.Wait() + } + + return port, cleanup +} + +// dialPythonPeer starts a Python peer, dials its TCP port, wraps the +// connection in an XshardConn, and starts it. Returns the XshardConn and a +// cleanup function. +func dialPythonPeer(t *testing.T, extraArgs ...string) (*XshardConn, func()) { + t.Helper() + + port, cleanupPy := startPythonPeer(t, extraArgs...) + + conn, err := net.Dial("tcp", fmt.Sprintf("127.0.0.1:%d", port)) + if err != nil { + cleanupPy() + t.Fatalf("dial python peer: %v", err) + } + + xc := NewXshardConnFromConn(conn, 0, []byte("go"), []uint32{1}, log.New()) + xc.Start() + + cleanup := func() { + xc.Close() + conn.Close() + cleanupPy() + } + + return xc, cleanup +} + +// --------------------------------------------------------------------------- +// Test: Python → Go PING/PONG +// +// Validates: Python SlaveConnection.send_ping() initiator behavior. +// Python sends PING, Go XshardConn.handlePing() records identity and replies +// PONG. Tests that Go correctly receives and responds to a Python-initiated +// PING/PONG exchange. +// --------------------------------------------------------------------------- +func TestPythonCompat_PingPong_PythonToGo(t *testing.T) { + xc, cleanup := dialPythonPeer(t, "--send-ping") + defer cleanup() + + // Wait for Go side to receive PING from Python. + // Python peer sends PING immediately after accept. + if !xc.WaitUntilPingReceived() { + t.Fatal("Go did not receive PING from Python peer") + } + + // Verify Go recorded Python's identity from the PING. + if got := string(xc.RemoteID()); got != "py" { + t.Fatalf("RemoteID: got %q, want %q", got, "py") + } + shards := xc.RemoteFullShardIDList() + if len(shards) != 1 || shards[0] != 1 { + t.Fatalf("RemoteFullShardIDList: got %v, want [1]", shards) + } +} + +// --------------------------------------------------------------------------- +// Test: Go → Python PING/PONG +// +// Validates: Go XshardConn.SendPing() outbound PING/PONG exchange. +// Go sends PING, Python SlaveConnection.handle_ping() records identity and +// replies PONG. Tests that Go's SendPing() correctly parses Python's PONG +// response. +// --------------------------------------------------------------------------- +func TestPythonCompat_PingPong_GoToPython(t *testing.T) { + xc, cleanup := dialPythonPeer(t) + defer cleanup() + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + id, shardList, err := xc.SendPing(ctx) + if err != nil { + t.Fatalf("SendPing: %v", err) + } + + if string(id) != "py" { + t.Fatalf("SendPing returned id %q, want %q", string(id), "py") + } + if len(shardList) != 1 || shardList[0] != 1 { + t.Fatalf("SendPing returned shardList %v, want [1]", shardList) + } +} + +// --------------------------------------------------------------------------- +// Test: RPC request/response matching +// +// Validates: Python's echo-RPC behavior (opcode → opcode+1, same rpc_id, +// same payload). Verifies that Go's RPC ID generation, pending map lifecycle, +// and response matching work correctly when communicating with a Python peer. +// --------------------------------------------------------------------------- +func TestPythonCompat_RPCRequestResponse(t *testing.T) { + xc, cleanup := dialPythonPeer(t) + defer cleanup() + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + // Send a request with opcode=0x10. Python echoes back opcode=0x11. + payload := []byte("hello-rpc") + resp, err := xc.SendRPC(ctx, 0x10, payload) + if err != nil { + t.Fatalf("SendRPC: %v", err) + } + + if resp.Opcode != 0x11 { + t.Fatalf("response opcode: got 0x%02x, want 0x11", resp.Opcode) + } + if string(resp.Payload) != string(payload) { + t.Fatalf("response payload: got %q, want %q", string(resp.Payload), string(payload)) + } + + // Send a second RPC with a different payload to verify sequential RPCs. + payload2 := []byte("second-rpc") + resp2, err := xc.SendRPC(ctx, 0x10, payload2) + if err != nil { + t.Fatalf("second SendRPC: %v", err) + } + + if resp2.Opcode != 0x11 { + t.Fatalf("second response opcode: got 0x%02x, want 0x11", resp2.Opcode) + } + if string(resp2.Payload) != string(payload2) { + t.Fatalf("second response payload: got %q, want %q", string(resp2.Payload), string(payload2)) + } + + // Verify RPC IDs are unique (each response matches its own request). + if resp.RPCID == resp2.RPCID { + t.Fatal("RPC IDs should be unique") + } +} + +// --------------------------------------------------------------------------- +// Test: Connection close propagation +// +// Validates: Python's SlaveConnection.close() behavior. +// When the Python peer disconnects, Go's readLoop must detect the TCP close +// and call Close(). After close, any RPC must fail with ErrConnectionClosed. +// +// Note: Testing mid-flight RPC wakeup is non-deterministic because Python +// echoes the response before the process is killed. This test verifies the +// deterministic post-close behavior instead. +// --------------------------------------------------------------------------- +func TestPythonCompat_ConnectionClosePropagation(t *testing.T) { + port, cleanupPy := startPythonPeer(t) + + conn, err := net.Dial("tcp", fmt.Sprintf("127.0.0.1:%d", port)) + if err != nil { + cleanupPy() + t.Fatalf("dial: %v", err) + } + + xc := NewXshardConnFromConn(conn, 0, []byte("go"), []uint32{1}, log.New()) + xc.Start() + defer xc.Close() + + // Kill the Python peer — this closes the TCP connection from the other end. + cleanupPy() + + // Wait for Go to detect the connection close. + select { + case <-xc.WaitUntilClosed(): + case <-time.After(5 * time.Second): + t.Fatal("Go did not detect connection close within 5 seconds") + } + + if !xc.IsClosed() { + t.Fatal("XshardConn should be closed after Python disconnect") + } + + // Any RPC after close should fail with ErrConnectionClosed. + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + _, err = xc.SendRPC(ctx, 0x01, []byte("test")) + if err != ErrConnectionClosed { + t.Fatalf("expected ErrConnectionClosed after close, got %v", err) + } +} + +// --------------------------------------------------------------------------- +// Test: Pool reconnect after Remove +// +// Validates: Python's SlaveConnectionManager.connect_to_slave() reconnection +// behavior. After a connection is removed from the pool and the slave ID is +// cleaned up, a new connection to a peer with the same identity must be +// accepted. Tests the XshardPool.Remove() → slaveIDs cleanup → reconnection +// invariant. +// --------------------------------------------------------------------------- +func TestPythonCompat_PoolReconnect(t *testing.T) { + pool := NewXshardPool(log.New()) + defer pool.Close() + + // --- First connection --- + xc1, cleanup1 := dialPythonPeer(t) + defer cleanup1() + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + if err := pool.VerifyAndAdd(ctx, 1, xc1, []byte("py"), []uint32{1}); err != nil { + t.Fatalf("first VerifyAndAdd: %v", err) + } + if pool.OutboundSize() != 1 { + t.Fatalf("pool size after add: got %d, want 1", pool.OutboundSize()) + } + + // Remove and verify the pool is empty. + pool.Remove(1, xc1) + if pool.OutboundSize() != 0 { + t.Fatalf("pool size after remove: got %d, want 0", pool.OutboundSize()) + } + + // Clean up the first peer before starting the second. + cleanup1() + + // --- Second connection (same identity, should be accepted) --- + xc2, cleanup2 := dialPythonPeer(t) + defer cleanup2() + + ctx2, cancel2 := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel2() + + if err := pool.VerifyAndAdd(ctx2, 1, xc2, []byte("py"), []uint32{1}); err != nil { + t.Fatalf("second VerifyAndAdd (reconnect) failed: %v", err) + } + if pool.OutboundSize() != 1 { + t.Fatalf("pool size after reconnect: got %d, want 1", pool.OutboundSize()) + } +} diff --git a/qkc/cluster/slave/connection.go b/qkc/cluster/slave/connection.go new file mode 100644 index 0000000000..fdb0c8b18b --- /dev/null +++ b/qkc/cluster/slave/connection.go @@ -0,0 +1,544 @@ +// Copyright 2026-2027, QuarkChain. + +package slave + +import ( + "bufio" + "context" + "fmt" + "io" + "net" + "sync" + "sync/atomic" + + "github.com/ethereum/go-ethereum/log" + "github.com/ethereum/go-ethereum/qkc/cluster/wire" + "github.com/ethereum/go-ethereum/qkc/serialize" +) + +// serializeBytes serializes a wire message using the qkc/serialize package. +func serializeBytes(v any) ([]byte, error) { + return serialize.SerializeToBytes(v) +} + +// deserializeBytes deserializes a wire message from payload bytes. +func deserializeBytes(p []byte, v any) error { + return serialize.Deserialize(serialize.NewByteBuffer(p), v) +} + +// TypedHandler processes a deserialized request and returns a deserialized +// response. The framework handles payload serialization/deserialization. +type TypedHandler func(req any) (resp any, err error) + +// OpSerializer describes how to deserialize a request and serialize a response +// for a specific opcode. It mirrors Python's op_ser_map entries. +type OpSerializer struct { + NewRequest func() any + Deserialize func([]byte, any) error + Serialize func(any) ([]byte, error) + ResponseOpCode byte // optional: if non-zero, used as response opcode +} + +// OpSerializerFor creates an OpSerializer for wire types R (request) and S (response). +func OpSerializerFor[R, S any]() *OpSerializer { + return &OpSerializer{ + NewRequest: func() any { return new(R) }, + Deserialize: func(p []byte, v any) error { + return deserializeBytes(p, v) + }, + Serialize: func(v any) ([]byte, error) { + return serializeBytes(v) + }, + } +} + +// ConnectionState mirrors Python's protocol.ConnectionState. +type ConnectionState int32 + +const ( + ConnectionStateConnecting ConnectionState = iota + ConnectionStateActive + ConnectionStateClosed +) + +// ── transport: pure I/O layer ───────────────────────────────────────────────── + +// transport wraps a net.Conn with metadata-aware frame read/write. +// writeMu serializes writes because bufio.Writer is not goroutine-safe and +// both SendRPC (any goroutine) and readLoop handler goroutines write frames. +type transport struct { + conn net.Conn + r *bufio.Reader + w *bufio.Writer + + writeMu sync.Mutex + + readFrameFn func(io.Reader) (*wire.Frame, error) + writeFrameFn func(io.Writer, *wire.Frame) error + + remoteAddr string +} + +func newTransport( + conn net.Conn, + readFrame func(io.Reader) (*wire.Frame, error), + writeFrame func(io.Writer, *wire.Frame) error, +) *transport { + return &transport{ + conn: conn, + r: bufio.NewReader(conn), + w: bufio.NewWriter(conn), + readFrameFn: readFrame, + writeFrameFn: writeFrame, + remoteAddr: conn.RemoteAddr().String(), + } +} + +func (t *transport) readFrame() (*wire.Frame, error) { + return t.readFrameFn(t.r) +} + +func (t *transport) writeFrame(f *wire.Frame) error { + t.writeMu.Lock() + defer t.writeMu.Unlock() + + if err := t.writeFrameFn(t.w, f); err != nil { + return fmt.Errorf("write frame: %w", err) + } + if err := t.w.Flush(); err != nil { + return fmt.Errorf("flush: %w", err) + } + return nil +} + +func (t *transport) close() error { + return t.conn.Close() +} + +func (t *transport) RemoteAddr() string { + return t.remoteAddr +} + +// ── rpcConn: RPC protocol engine ───────────────────────────────────────────── + +// rpcResult is the value delivered over a pending RPC response channel. +type rpcResult struct { + frame *wire.Frame + err error +} + +// rpcConn is the shared RPC engine used by XshardConn (and later MasterConn). +// It handles lifecycle, handler/serializer registration, readLoop dispatch, +// RPC request/response matching, and monotonic RPC ID validation. +// +// The forwarder hook is an extension point for MasterConn to route peer traffic +// to PeerShardConn. For XshardConn it remains nil. +// +// Lock ordering (must be maintained to avoid deadlocks): +// +// closeMu → pendingMu (SendRPCMeta, Close) +// closeMu → stateMu (Close) +// +// pendingMu and stateMu are never held together; readLoop only holds pendingMu. +type rpcConn struct { + *transport + + stateMu sync.Mutex + state ConnectionState + activeChan chan struct{} + closedChan chan struct{} + + errChan chan error + startOnce sync.Once + + handlersMu sync.RWMutex + typedHandlers map[byte]TypedHandler + + serializersMu sync.RWMutex + serializers map[byte]*OpSerializer + + nonRPCOps map[byte]struct{} + + pendingMu sync.Mutex + pending map[uint64]chan rpcResult + + nextRPCID uint64 + + // peerRPCID tracks the most recent inbound RPC ID for monotonic validation. + // Initialized to -1 (like Python) so the first valid rpc_id must be >= 1. + peerRPCID int64 + peerRPCIDMu sync.Mutex + + // validateRPCID is called by readLoop for every RPC request frame. + // Default: simple global monotonic validation. + // MasterConn replaces with per-peer tracking. + validateRPCID func(clusterPeerID uint64, rpcID uint64) bool + + forwarder func(*wire.Frame) bool + forwarderMu sync.RWMutex + + closeMu sync.Mutex + closed bool + + log log.Logger +} + +func newRPCConn( + conn net.Conn, + readFrame func(io.Reader) (*wire.Frame, error), + writeFrame func(io.Writer, *wire.Frame) error, + logger log.Logger, +) *rpcConn { + if logger == nil { + logger = log.Root() + } + rc := &rpcConn{ + transport: newTransport(conn, readFrame, writeFrame), + typedHandlers: make(map[byte]TypedHandler), + serializers: make(map[byte]*OpSerializer), + pending: make(map[uint64]chan rpcResult), + peerRPCID: -1, + nonRPCOps: make(map[byte]struct{}), + state: ConnectionStateConnecting, + activeChan: make(chan struct{}), + closedChan: make(chan struct{}), + errChan: make(chan error, 1), + log: logger, + } + rc.validateRPCID = rc.defaultValidateRPCID + return rc +} + +// Start transitions the connection to ACTIVE and launches the read loop. +func (c *rpcConn) Start() { + c.startOnce.Do(func() { + c.stateMu.Lock() + c.state = ConnectionStateActive + close(c.activeChan) + c.stateMu.Unlock() + go c.readLoop() + }) +} + +// Close closes the connection and wakes all pending RPCs. +func (c *rpcConn) Close() error { + c.closeMu.Lock() + if c.closed { + c.closeMu.Unlock() + return nil + } + c.closed = true + c.closeMu.Unlock() + + c.stateMu.Lock() + if c.state != ConnectionStateClosed { + c.state = ConnectionStateClosed + close(c.closedChan) + // Wake up any goroutines waiting on WaitUntilActive(). + // Matches Python's finally block in active_and_loop_forever that sets active_event. + select { + case <-c.activeChan: + // Already closed (Start was called) + default: + close(c.activeChan) + } + } + c.stateMu.Unlock() + + c.pendingMu.Lock() + for rpcID, ch := range c.pending { + select { + case ch <- rpcResult{err: ErrConnectionClosed}: + default: + } + delete(c.pending, rpcID) + } + c.pendingMu.Unlock() + + return c.transport.close() +} + +func (c *rpcConn) defaultValidateRPCID(clusterPeerID uint64, rpcID uint64) bool { + c.peerRPCIDMu.Lock() + defer c.peerRPCIDMu.Unlock() + if int64(rpcID) <= c.peerRPCID { + return false + } + c.peerRPCID = int64(rpcID) + return true +} + +// RegisterTypedHandlers registers opcode handlers. Nil handlers panic. +func (c *rpcConn) RegisterTypedHandlers(handlers map[byte]TypedHandler) { + c.handlersMu.Lock() + defer c.handlersMu.Unlock() + for opcode, handler := range handlers { + if handler == nil { + panic("handler must not be nil") + } + c.typedHandlers[opcode] = handler + } +} + +// RegisterOpSerializers registers opcode serializers. +func (c *rpcConn) RegisterOpSerializers(serializers map[byte]*OpSerializer) { + c.serializersMu.Lock() + defer c.serializersMu.Unlock() + for opcode, ser := range serializers { + if ser == nil { + panic("serializer must not be nil") + } + c.serializers[opcode] = ser + } +} + +// RegisterNonRPCOps marks opcodes as non-RPC (fire-and-forget), meaning they +// must have rpc_id == 0. +func (c *rpcConn) RegisterNonRPCOps(ops []byte) { + c.handlersMu.Lock() + defer c.handlersMu.Unlock() + for _, op := range ops { + c.nonRPCOps[op] = struct{}{} + } +} + +// SetForwarder installs a raw-frame forwarder hook. If it returns true the +// frame is consumed and readLoop continues without dispatching it. +func (c *rpcConn) SetForwarder(f func(*wire.Frame) bool) { + c.forwarderMu.Lock() + defer c.forwarderMu.Unlock() + c.forwarder = f +} + +// SendRPC sends a request with zero metadata and waits for the response. +// For connections that need metadata (e.g. MasterConn with 12-byte +// ClusterMetadata), use SendRPCMeta directly. +func (c *rpcConn) SendRPC(ctx context.Context, opcode byte, payload []byte) (*wire.Frame, error) { + return c.SendRPCMeta(ctx, opcode, payload, wire.ClusterMetadata{}) +} + +// SendRPCMeta sends a request with the given metadata and waits for the response. +// XshardConn uses zero metadata (0-byte wire format). +// MasterConn uses ClusterMetadata{Branch, ClusterPeerID} (12-byte wire format). +func (c *rpcConn) SendRPCMeta(ctx context.Context, opcode byte, payload []byte, meta wire.ClusterMetadata) (*wire.Frame, error) { + c.stateMu.Lock() + state := c.state + c.stateMu.Unlock() + + switch state { + case ConnectionStateClosed: + return nil, ErrConnectionClosed + case ConnectionStateConnecting: + return nil, ErrNotActive + } + + c.closeMu.Lock() + if c.closed { + c.closeMu.Unlock() + return nil, ErrConnectionClosed + } + + rpcID := atomic.AddUint64(&c.nextRPCID, 1) + respChan := make(chan rpcResult, 1) + c.pendingMu.Lock() + c.pending[rpcID] = respChan + c.pendingMu.Unlock() + c.closeMu.Unlock() + + defer func() { + c.pendingMu.Lock() + delete(c.pending, rpcID) + c.pendingMu.Unlock() + }() + + frame := &wire.Frame{ + Meta: meta, + Opcode: opcode, + RPCID: rpcID, + Payload: payload, + } + if err := c.transport.writeFrame(frame); err != nil { + return nil, err + } + + select { + case res := <-respChan: + if res.err != nil { + return nil, res.err + } + if res.frame == nil { + return nil, ErrConnectionClosed + } + return res.frame, nil + case <-ctx.Done(): + return nil, fmt.Errorf("rpc timeout: %w", ctx.Err()) + } +} + +// ── Read loop ───────────────────────────────────────────────────────────────── + +// readLoop reads frames until a fatal error, then closes the connection. +// Follows Python's protocol validation rules strictly. +func (c *rpcConn) readLoop() { + defer c.Close() + + for { + frame, err := c.transport.readFrame() + if err != nil { + select { + case c.errChan <- err: + default: + } + return + } + + // Forwarder hook (extension point for MasterConn). + c.forwarderMu.RLock() + fwd := c.forwarder + c.forwarderMu.RUnlock() + if fwd != nil && fwd(frame) { + continue + } + + c.handlersMu.RLock() + handler, isRequest := c.typedHandlers[frame.Opcode] + _, isNonRPC := c.nonRPCOps[frame.Opcode] + c.handlersMu.RUnlock() + + c.serializersMu.RLock() + ser := c.serializers[frame.Opcode] + c.serializersMu.RUnlock() + + // No handler: could be a pending RPC response or unsupported opcode. + if !isRequest { + if frame.RPCID != 0 { + c.pendingMu.Lock() + if ch, ok := c.pending[frame.RPCID]; ok { + delete(c.pending, frame.RPCID) + c.pendingMu.Unlock() + select { + case ch <- rpcResult{frame: frame}: + default: + c.log.Warn("response channel full", "rpcid", frame.RPCID) + } + continue + } + c.pendingMu.Unlock() + // INTENTIONAL DEVIATION FROM PYTHON: Python closes connection on + // unexpected RPC response (rpc_id not in rpc_future_map). Go keeps + // connection open and logs error. This is more robust for distributed + // systems where late/duplicate responses are normal after timeout. + // If strict Python compatibility is needed, change to: return + c.log.Error("unexpected rpc response (rpc_id not in pending map)", + "rpcid", frame.RPCID, "opcode", frame.Opcode) + continue + } + c.log.Warn("unsupported opcode", "opcode", frame.Opcode) + return + } + + if ser == nil { + c.log.Warn("handler without serializer", "opcode", frame.Opcode) + return + } + + if isNonRPC && frame.RPCID != 0 { + c.log.Warn("non-rpc command with non-zero rpc_id", "opcode", frame.Opcode, "rpcid", frame.RPCID) + return + } + + if !isNonRPC { + if !c.validateRPCID(frame.Meta.ClusterPeerID, frame.RPCID) { + c.log.Warn("incorrect rpc request id sequence", "rpcid", frame.RPCID) + return + } + } + + go c.dispatch(frame, handler, ser) + } +} + +func (c *rpcConn) dispatch(frame *wire.Frame, handler TypedHandler, ser *OpSerializer) { + defer func() { + if r := recover(); r != nil { + c.log.Error("handler panic", "opcode", frame.Opcode, "panic", r) + c.Close() + } + }() + + req := ser.NewRequest() + if err := ser.Deserialize(frame.Payload, req); err != nil { + c.log.Error("deserialize failed", "opcode", frame.Opcode, "err", err) + c.Close() + return + } + + resp, err := handler(req) + if err != nil { + // NOTE: All handler errors close the connection. This matches Python's + // close_with_error pattern and is intentional for protocol safety. + // The QuarkChain cluster protocol treats handler errors as fatal because + // there's no error response mechanism — the only way to signal failure + // is to close the connection. If recoverable errors are needed in the + // future, the protocol would need to be extended with error responses. + c.log.Error("handler error", "opcode", frame.Opcode, "err", err) + c.Close() + return + } + + if frame.RPCID == 0 { + return // non-RPC: no response + } + + respPayload, err := ser.Serialize(resp) + if err != nil { + c.log.Error("serialize response failed", "opcode", frame.Opcode, "err", err) + c.Close() + return + } + respOp := frame.Opcode + 1 + if ser.ResponseOpCode != 0 { + respOp = ser.ResponseOpCode + } + respFrame := &wire.Frame{ + Meta: frame.Meta, + Opcode: respOp, + RPCID: frame.RPCID, + Payload: respPayload, + } + if err := c.transport.writeFrame(respFrame); err != nil { + c.log.Error("write response failed", "opcode", respFrame.Opcode, "err", err) + c.Close() + } +} + +// ── Query helpers ───────────────────────────────────────────────────────────── + +func (c *rpcConn) Error() <-chan error { return c.errChan } +func (c *rpcConn) RemoteAddr() string { return c.transport.RemoteAddr() } +func (c *rpcConn) WaitUntilActive() <-chan struct{} { return c.activeChan } +func (c *rpcConn) WaitUntilClosed() <-chan struct{} { return c.closedChan } + +func (c *rpcConn) State() ConnectionState { + c.stateMu.Lock() + defer c.stateMu.Unlock() + return c.state +} + +func (c *rpcConn) IsActive() bool { + c.stateMu.Lock() + defer c.stateMu.Unlock() + return c.state == ConnectionStateActive +} + +func (c *rpcConn) IsClosed() bool { + c.stateMu.Lock() + defer c.stateMu.Unlock() + return c.state == ConnectionStateClosed +} + +func (c *rpcConn) Closed() bool { + c.closeMu.Lock() + defer c.closeMu.Unlock() + return c.closed +} diff --git a/qkc/cluster/slave/errors.go b/qkc/cluster/slave/errors.go new file mode 100644 index 0000000000..d6bce48a85 --- /dev/null +++ b/qkc/cluster/slave/errors.go @@ -0,0 +1,17 @@ +// Copyright 2026-2027, QuarkChain. + +package slave + +import ( + "errors" +) + +var ( + // ErrConnectionClosed is returned when an operation is attempted on a + // connection that has already been closed. + ErrConnectionClosed = errors.New("connection closed") + + // ErrNotActive is returned when an RPC is attempted on a connection that + // has not been started (state != ACTIVE). + ErrNotActive = errors.New("connection not active") +) diff --git a/qkc/cluster/slave/testdata/pyproto/__init__.py b/qkc/cluster/slave/testdata/pyproto/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/qkc/cluster/slave/testdata/pyproto/frame.py b/qkc/cluster/slave/testdata/pyproto/frame.py new file mode 100644 index 0000000000..8582eaea86 --- /dev/null +++ b/qkc/cluster/slave/testdata/pyproto/frame.py @@ -0,0 +1,35 @@ +"""Frame read/write for slave-to-slave protocol (0-byte metadata). + +Wire format: [4B payload_len][1B opcode][8B rpc_id][payload] + +This matches Go's qkc/cluster/wire ReadFrameNoMeta/WriteFrameNoMeta. +""" +import struct + + +def read_frame(conn): + """Read one frame from conn. Returns (opcode, rpc_id, payload) or None on EOF.""" + header = conn.recv(13) # 4 (payload_len) + 1 (opcode) + 8 (rpc_id) + if not header: + return None + if len(header) < 13: + raise ConnectionError("truncated frame header") + + payload_len = struct.unpack('>I', header[0:4])[0] + opcode = header[4] + rpc_id = struct.unpack('>Q', header[5:13])[0] + + payload = b'' + while len(payload) < payload_len: + chunk = conn.recv(payload_len - len(payload)) + if not chunk: + raise ConnectionError("truncated frame payload") + payload += chunk + + return (opcode, rpc_id, payload) + + +def write_frame(conn, opcode, rpc_id, payload): + """Write one frame to conn.""" + header = struct.pack('>I', len(payload)) + bytes([opcode]) + struct.pack('>Q', rpc_id) + conn.sendall(header + payload) \ No newline at end of file diff --git a/qkc/cluster/slave/testdata/pyproto/messages.py b/qkc/cluster/slave/testdata/pyproto/messages.py new file mode 100644 index 0000000000..3767f7634c --- /dev/null +++ b/qkc/cluster/slave/testdata/pyproto/messages.py @@ -0,0 +1,73 @@ +"""Message serialization matching Go's qkc/serialize + qkc/cluster/wire messages. + +Conventions (from Go's serialize package): + - []byte: 4-byte big-endian length prefix + raw bytes + - []uint32: 4-byte big-endian count prefix + big-endian uint32 values + - *RootBlock with ser:"nil": 0x00 = nil +""" +import struct + + +def serialize_ping_request(id_bytes, full_shard_id_list): + """Serialize PingRequest matching Go's PingRequest + serialize. + + Fields: + ID: []byte (4B len + raw) + FullShardIDList: []uint32 (4B count + uint32[]) + RootTip: *RootBlock (nil marker 0x00) + """ + data = b'' + data += struct.pack('>I', len(id_bytes)) + id_bytes + data += struct.pack('>I', len(full_shard_id_list)) + for shard_id in full_shard_id_list: + data += struct.pack('>I', shard_id) + data += b'\x00' # RootTip: nil + return data + + +def serialize_pong_response(id_bytes, full_shard_id_list): + """Serialize PongResponse matching Go's PongResponse + serialize. + + Fields: + ID: []byte (4B len + raw) + FullShardIDList: []uint32 (4B count + uint32[]) + """ + data = b'' + data += struct.pack('>I', len(id_bytes)) + id_bytes + data += struct.pack('>I', len(full_shard_id_list)) + for shard_id in full_shard_id_list: + data += struct.pack('>I', shard_id) + return data + + +def parse_ping_request(data): + """Parse PingRequest payload. Returns (id, full_shard_id_list).""" + offset = 0 + id_len = struct.unpack('>I', data[offset:offset + 4])[0] + offset += 4 + id_bytes = data[offset:offset + id_len] + offset += id_len + count = struct.unpack('>I', data[offset:offset + 4])[0] + offset += 4 + shard_list = [] + for _ in range(count): + shard_list.append(struct.unpack('>I', data[offset:offset + 4])[0]) + offset += 4 + # Skip RootTip nil marker (1 byte) + return (id_bytes, shard_list) + + +def parse_pong_response(data): + """Parse PongResponse payload. Returns (id, full_shard_id_list).""" + offset = 0 + id_len = struct.unpack('>I', data[offset:offset + 4])[0] + offset += 4 + id_bytes = data[offset:offset + id_len] + offset += id_len + count = struct.unpack('>I', data[offset:offset + 4])[0] + offset += 4 + shard_list = [] + for _ in range(count): + shard_list.append(struct.unpack('>I', data[offset:offset + 4])[0]) + offset += 4 + return (id_bytes, shard_list) \ No newline at end of file diff --git a/qkc/cluster/slave/testdata/pyproto/peer.py b/qkc/cluster/slave/testdata/pyproto/peer.py new file mode 100644 index 0000000000..ff07324b67 --- /dev/null +++ b/qkc/cluster/slave/testdata/pyproto/peer.py @@ -0,0 +1,127 @@ +#!/usr/bin/env python3 +"""Minimal SlaveConnection protocol peer for Go compatibility tests. + +This peer implements only the slave-to-slave protocol that XshardConn needs: + - Frame read/write (0-byte metadata, matching ReadFrameNoMeta/WriteFrameNoMeta) + - PING/PONG identity exchange (ClusterOp 0x81/0x82) + - Echo RPC (opcode → opcode+1, same rpc_id, same payload) + +Usage: + python3 peer.py --port 0 --id "py" --shards "1,2" [--send-ping] + + --port TCP port to listen on (0 = random, actual port printed to stdout) + --id Peer identity (string, encoded as UTF-8 bytes) + --shards Comma-separated list of full shard IDs, e.g. "1,2" + --send-ping Send PING immediately after connect, wait for PONG, then enter read loop + +Output: + PORT: Printed when listening + PONG_OK id= Printed when --send-ping PONG is received + PING_RECEIVED ... Printed when PING is received from peer + DISCONNECTED Printed when connection closes + +Behavior: + - Listens on TCP, accepts one connection + - If --send-ping: sends PING (rpc_id=1), waits for PONG, prints PONG_OK + - Read loop: + PING(0x81) → record peer identity, reply PONG(0x82) + any opcode → reply opcode+1, same rpc_id, same payload + - On disconnect: exits +""" +import argparse +import socket +import struct +import sys + +from frame import read_frame, write_frame +from messages import ( + serialize_ping_request, + serialize_pong_response, + parse_ping_request, + parse_pong_response, +) + +CLUSTER_OP_PING = 0x81 +CLUSTER_OP_PONG = 0x82 + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--port', type=int, required=True) + parser.add_argument('--id', type=str, required=True) + parser.add_argument('--shards', type=str, required=True) + parser.add_argument('--send-ping', action='store_true') + args = parser.parse_args() + + peer_id = args.id.encode('utf-8') + shard_list = [int(s) for s in args.shards.split(',')] + + server = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + server.bind(('127.0.0.1', args.port)) + server.listen(1) + + actual_port = server.getsockname()[1] + print(f"PORT:{actual_port}", flush=True) + + conn, addr = server.accept() + + try: + if args.send_ping: + _do_send_ping(conn, peer_id, shard_list) + + _read_loop(conn, peer_id, shard_list) + + except (ConnectionError, BrokenPipeError, OSError): + pass + finally: + conn.close() + server.close() + print("DISCONNECTED", flush=True) + + +def _do_send_ping(conn, peer_id, shard_list): + """Send PING (rpc_id=1), wait for PONG, validate and print result.""" + ping_payload = serialize_ping_request(peer_id, shard_list) + write_frame(conn, CLUSTER_OP_PING, 1, ping_payload) + + frame = read_frame(conn) + if frame is None: + print("ERROR: no pong received", flush=True) + sys.exit(1) + + opcode, rpc_id, payload = frame + if opcode != CLUSTER_OP_PONG: + print(f"ERROR: expected PONG(0x{CLUSTER_OP_PONG:02x}), got 0x{opcode:02x}", flush=True) + sys.exit(1) + if rpc_id != 1: + print(f"ERROR: expected rpc_id 1, got {rpc_id}", flush=True) + sys.exit(1) + + peer_id_recv, _ = parse_pong_response(payload) + print(f"PONG_OK id={peer_id_recv.hex()}", flush=True) + + +def _read_loop(conn, peer_id, shard_list): + """Read frames, handle PING or echo RPC, until disconnect.""" + while True: + frame = read_frame(conn) + if frame is None: + break + + opcode, rpc_id, payload = frame + + if opcode == CLUSTER_OP_PING: + peer_id_recv, peer_shards = parse_ping_request(payload) + shard_str = ",".join(str(s) for s in peer_shards) + print(f"PING_RECEIVED id={peer_id_recv.hex()} shards={shard_str}", flush=True) + + pong_payload = serialize_pong_response(peer_id, shard_list) + write_frame(conn, CLUSTER_OP_PONG, rpc_id, pong_payload) + else: + # Echo RPC: opcode+1, same rpc_id, same payload + write_frame(conn, opcode + 1, rpc_id, payload) + + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/qkc/cluster/slave/xshard_conn.go b/qkc/cluster/slave/xshard_conn.go new file mode 100644 index 0000000000..174e55a91b --- /dev/null +++ b/qkc/cluster/slave/xshard_conn.go @@ -0,0 +1,230 @@ +// Copyright 2026-2027, QuarkChain. + +package slave + +import ( + "context" + "fmt" + "io" + "net" + "sync" + "time" + + "github.com/ethereum/go-ethereum/log" + "github.com/ethereum/go-ethereum/qkc/cluster/wire" +) + +const defaultDialTimeout = 10 * time.Second + +// XshardConn is a direct TCP connection to another slave node for cross-shard +// traffic. It uses 0-byte metadata (slave↔slave mode) and corresponds to Python's +// SlaveConnection. +// +// Architecture: +// +// XshardConn embeds *rpcConn embeds *transport +// +// No forwarder — all frames are dispatched locally. RPC ID validation is +// global monotonic (the default in rpcConn). +type XshardConn struct { + *rpcConn + + // local identity of this slave, used in PONG responses. + localID []byte + localFullShardIDList []uint32 + + // peer identity state, protected by its own mutex (not rpcConn.closeMu). + stateMu sync.Mutex + remoteID []byte + remoteFullShardIDList []uint32 + pingReceived chan struct{} + pingOnce sync.Once +} + +// NewXshardConn dials another slave and returns an XshardConn. +// Call RegisterHandlers then Start before using the connection. +// maxPayloadSize controls frame payload size limit; 0 disables the limit. +// localID and localFullShardIDList identify this slave and are used in PONG responses. +func NewXshardConn(addr string, maxPayloadSize uint32, localID []byte, localFullShardIDList []uint32, logger log.Logger) (*XshardConn, error) { + conn, err := net.DialTimeout("tcp", addr, defaultDialTimeout) + if err != nil { + return nil, fmt.Errorf("dial xshard slave %s: %w", addr, err) + } + return newXshardConn(conn, maxPayloadSize, localID, localFullShardIDList, logger), nil +} + +// NewXshardConnFromConn wraps an accepted net.Conn as an XshardConn. +// maxPayloadSize controls frame payload size limit; 0 disables the limit. +// localID and localFullShardIDList identify this slave and are used in PONG responses. +func NewXshardConnFromConn(conn net.Conn, maxPayloadSize uint32, localID []byte, localFullShardIDList []uint32, logger log.Logger) *XshardConn { + return newXshardConn(conn, maxPayloadSize, localID, localFullShardIDList, logger) +} + +func newXshardConn(conn net.Conn, maxPayloadSize uint32, localID []byte, localFullShardIDList []uint32, logger log.Logger) *XshardConn { + readFrame := func(r io.Reader) (*wire.Frame, error) { + return wire.ReadFrameNoMeta(r, maxPayloadSize) + } + xc := &XshardConn{ + rpcConn: newRPCConn(conn, readFrame, wire.WriteFrameNoMeta, logger), + localID: append([]byte(nil), localID...), + localFullShardIDList: append([]uint32(nil), localFullShardIDList...), + pingReceived: make(chan struct{}), + } + + // Register serializers for all opcodes that SlaveConnection understands. + // This matches Python's SLAVE_OP_SERIALIZER_MAP. + xc.rpcConn.RegisterOpSerializers(map[byte]*OpSerializer{ + byte(wire.ClusterOpPing): OpSerializerFor[wire.PingRequest, wire.PongResponse](), + byte(wire.ClusterOpAddXshardTxListRequest): OpSerializerFor[wire.AddXshardTxListRequest, wire.AddXshardTxListResponse](), + byte(wire.ClusterOpBatchAddXshardTxListRequest): OpSerializerFor[wire.BatchAddXshardTxListRequest, wire.BatchAddXshardTxListResponse](), + }) + + // PING is handled internally by SlaveConnection in Python; register the + // built-in handler immediately so it works even if the caller never calls + // RegisterHandlers. + xc.rpcConn.RegisterTypedHandlers(map[byte]TypedHandler{ + byte(wire.ClusterOpPing): xc.handlePing, + }) + + return xc +} + +// handlePing is the built-in PING handler. It records peer identity, validates +// the shard list, and returns a PONG with this slave's identity. +func (x *XshardConn) handlePing(req any) (any, error) { + ping := req.(*wire.PingRequest) + + // Record peer identity (only on first ping, matches Python's "if not self.id") + x.stateMu.Lock() + if len(x.remoteID) == 0 { + x.remoteID = append([]byte(nil), ping.ID...) + x.remoteFullShardIDList = append([]uint32(nil), ping.FullShardIDList...) + } + // Check stored shard list (matches Python's self.full_shard_id_list check) + storedShardList := x.remoteFullShardIDList + x.stateMu.Unlock() + + if len(storedShardList) == 0 { + // Returning error causes rpcConn to close connection (Python's close_with_error) + return nil, fmt.Errorf("empty shard list from slave %s", ping.ID) + } + + // Signal ping received AFTER check passes (matches Python's ping_received_event.set()) + if !x.rpcConn.Closed() { + x.pingOnce.Do(func() { close(x.pingReceived) }) + } + + return &wire.PongResponse{ + ID: append([]byte(nil), x.localID...), + FullShardIDList: append([]uint32(nil), x.localFullShardIDList...), + }, nil +} + +// RegisterHandlers registers user-provided opcode handlers. PING is always +// handled internally (see handlePing). If the user registers a PING handler, +// it is wrapped so that peer identity recording and empty-shard-list validation +// still happen first; the user's returned response object is then sent as the +// PONG body. +func (x *XshardConn) RegisterHandlers(handlers map[byte]TypedHandler) { + wrapped := make(map[byte]TypedHandler, len(handlers)) + for opcode, handler := range handlers { + if opcode != byte(wire.ClusterOpPing) { + wrapped[opcode] = handler + } + } + + if userPingHandler, ok := handlers[byte(wire.ClusterOpPing)]; ok { + wrapped[byte(wire.ClusterOpPing)] = func(req any) (any, error) { + ping := req.(*wire.PingRequest) + + // Record peer identity (only on first ping) + x.stateMu.Lock() + if len(x.remoteID) == 0 { + x.remoteID = append([]byte(nil), ping.ID...) + x.remoteFullShardIDList = append([]uint32(nil), ping.FullShardIDList...) + } + // Check stored shard list + storedShardList := x.remoteFullShardIDList + x.stateMu.Unlock() + + if len(storedShardList) == 0 { + return nil, fmt.Errorf("empty shard list from slave %s", ping.ID) + } + + // Signal ping received AFTER check passes + if !x.rpcConn.Closed() { + x.pingOnce.Do(func() { close(x.pingReceived) }) + } + + return userPingHandler(req) + } + } + + x.rpcConn.RegisterTypedHandlers(wrapped) +} + +// RemoteID returns the peer's slave ID, populated after the first PING. +func (x *XshardConn) RemoteID() []byte { + x.stateMu.Lock() + defer x.stateMu.Unlock() + return append([]byte(nil), x.remoteID...) +} + +// RemoteFullShardIDList returns the peer's full shard ID list, populated after +// the first PING. +func (x *XshardConn) RemoteFullShardIDList() []uint32 { + x.stateMu.Lock() + defer x.stateMu.Unlock() + return append([]uint32(nil), x.remoteFullShardIDList...) +} + +// WaitUntilPingReceived blocks until the first PING is received or the +// connection is closed. It returns true if the connection is still alive. +func (x *XshardConn) WaitUntilPingReceived() bool { + select { + case <-x.pingReceived: + return !x.rpcConn.Closed() + case <-x.rpcConn.Error(): + return false + } +} + +// SendPing sends a PING request and waits for PONG response. It returns the +// peer's id and full_shard_id_list from the PONG response. +// This is the outbound half of the slave-to-slave identity exchange, +// corresponding to Python's SlaveConnection.send_ping(). +// The connection must have been started (Start() called). +func (x *XshardConn) SendPing(ctx context.Context) (id []byte, shardList []uint32, err error) { + payload, err := serializeBytes(&wire.PingRequest{ + ID: x.localID, + FullShardIDList: x.localFullShardIDList, + RootTip: nil, // slave-to-slave: no root tip required + }) + if err != nil { + return nil, nil, fmt.Errorf("serialize ping: %w", err) + } + + frame, err := x.rpcConn.SendRPC(ctx, byte(wire.ClusterOpPing), payload) + if err != nil { + return nil, nil, fmt.Errorf("send ping: %w", err) + } + + var pong wire.PongResponse + if err := deserializeBytes(frame.Payload, &pong); err != nil { + return nil, nil, fmt.Errorf("deserialize pong: %w", err) + } + + return pong.ID, pong.FullShardIDList, nil +} + +// SendXshardTxList sends an AddXshardTxListRequest via RPC and returns the response. +// Python's ADD_XSHARD_TX_LIST_REQUEST is an RPC (in SLAVE_OP_RPC_MAP), not fire-and-forget. +func (x *XshardConn) SendXshardTxList(ctx context.Context, payload []byte) (*wire.Frame, error) { + return x.rpcConn.SendRPC(ctx, byte(wire.ClusterOpAddXshardTxListRequest), payload) +} + +// SendBatchXshardTxList sends a BatchAddXshardTxListRequest via RPC and returns the response. +// Python's BATCH_ADD_XSHARD_TX_LIST_REQUEST is an RPC (in SLAVE_OP_RPC_MAP). +func (x *XshardConn) SendBatchXshardTxList(ctx context.Context, payload []byte) (*wire.Frame, error) { + return x.rpcConn.SendRPC(ctx, byte(wire.ClusterOpBatchAddXshardTxListRequest), payload) +} diff --git a/qkc/cluster/slave/xshard_pool.go b/qkc/cluster/slave/xshard_pool.go new file mode 100644 index 0000000000..9e342b6f0d --- /dev/null +++ b/qkc/cluster/slave/xshard_pool.go @@ -0,0 +1,321 @@ +// Copyright 2026-2027, QuarkChain. + +package slave + +import ( + "bytes" + "context" + "fmt" + "sync" + + "github.com/ethereum/go-ethereum/log" + "github.com/ethereum/go-ethereum/qkc/cluster/wire" +) + +// XshardPool manages direct slave-to-slave xshard connections, indexed by full +// shard ID. It corresponds to Python's SlaveConnectionManager. +type XshardPool struct { + mu sync.RWMutex + conns map[uint32][]*XshardConn + inbound []*XshardConn + slaveIDs map[string]bool // Tracks slave IDs to prevent duplicate connections + closed bool + log log.Logger +} + +// NewXshardPool creates a new, empty connection pool. +func NewXshardPool(logger log.Logger) *XshardPool { + return &XshardPool{ + conns: make(map[uint32][]*XshardConn), + slaveIDs: make(map[string]bool), + log: logger, + } +} + +// Add adds a connection to the pool for the given full shard ID. +// If the pool is already closed, the connection is closed immediately. +// If the slave ID is already tracked, the connection is closed and a warning is logged. +func (p *XshardPool) Add(fullShardID uint32, conn *XshardConn) { + p.mu.Lock() + if p.closed { + p.mu.Unlock() + conn.Close() + p.log.Warn("xshard pool closed, closing outbound conn immediately", "remote", conn.RemoteAddr()) + return + } + + // Check for duplicate slave ID (matches Python's slave_ids deduplication) + remoteID := string(conn.RemoteID()) + if remoteID != "" && p.slaveIDs[remoteID] { + p.mu.Unlock() + conn.Close() + p.log.Warn("duplicate slave connection rejected", "slave_id", remoteID, "full_shard_id", fullShardID) + return + } + + // Track the slave ID + if remoteID != "" { + p.slaveIDs[remoteID] = true + } + + p.conns[fullShardID] = append(p.conns[fullShardID], conn) + p.mu.Unlock() + p.log.Info("added xshard connection", "full_shard_id", fullShardID, "remote", conn.RemoteAddr()) +} + +// VerifyAndAdd performs PING-based identity verification on an outbound +// connection before adding it to the pool. It matches Python's +// SlaveConnectionManager.connect_to_slave(). +// +// The connection must already have been started (Start() called). +// On verification failure the connection is closed. +func (p *XshardPool) VerifyAndAdd(ctx context.Context, fullShardID uint32, conn *XshardConn, expectedID []byte, expectedShardList []uint32) error { + id, shardList, err := conn.SendPing(ctx) + if err != nil { + conn.Close() + return fmt.Errorf("ping failed for %s: %w", conn.RemoteAddr(), err) + } + if !bytes.Equal(id, expectedID) { + conn.Close() + return fmt.Errorf("slave id mismatch for %s: expected %x, got %x", conn.RemoteAddr(), expectedID, id) + } + if len(shardList) != len(expectedShardList) { + conn.Close() + return fmt.Errorf("shard list length mismatch for %s: expected %d, got %d", conn.RemoteAddr(), len(expectedShardList), len(shardList)) + } + for i := range shardList { + if shardList[i] != expectedShardList[i] { + conn.Close() + return fmt.Errorf("shard list mismatch for %s: expected %v, got %v", conn.RemoteAddr(), expectedShardList, shardList) + } + } + p.Add(fullShardID, conn) + return nil +} + +// Get returns a snapshot of connections for the given full shard ID. +func (p *XshardPool) Get(fullShardID uint32) []*XshardConn { + p.mu.RLock() + conns := p.conns[fullShardID] + result := make([]*XshardConn, len(conns)) + copy(result, conns) + p.mu.RUnlock() + return result +} + +// Remove removes a specific connection from the pool. It also cleans up the +// slave ID tracking so the same slave can reconnect later. +func (p *XshardPool) Remove(fullShardID uint32, conn *XshardConn) { + p.mu.Lock() + defer p.mu.Unlock() + + conns := p.conns[fullShardID] + for i, c := range conns { + if c == conn { + copy(conns[i:], conns[i+1:]) + conns[len(conns)-1] = nil + p.conns[fullShardID] = conns[:len(conns)-1] + if len(p.conns[fullShardID]) == 0 { + delete(p.conns, fullShardID) + } + if remoteID := string(conn.RemoteID()); remoteID != "" { + delete(p.slaveIDs, remoteID) + } + p.log.Info("removed xshard connection", "full_shard_id", fullShardID, "remote", conn.RemoteAddr()) + return + } + } +} + +// RemoveTarget removes and closes all connections for a full shard ID. +func (p *XshardPool) RemoveTarget(fullShardID uint32) { + p.mu.Lock() + conns := p.conns[fullShardID] + delete(p.conns, fullShardID) + for _, conn := range conns { + if remoteID := string(conn.RemoteID()); remoteID != "" { + delete(p.slaveIDs, remoteID) + } + } + p.mu.Unlock() + + for _, conn := range conns { + conn.Close() + } + p.log.Info("removed all xshard connections to shard", "full_shard_id", fullShardID) +} + +// SendXshardTx broadcasts xshard transactions to all active connections for the +// target shard via RPC. Returns the first successful response or an error if no +// connection exists or all connections fail. +// +// This matches Python's broadcast_xshard_tx_list behavior: sends to ALL connections +// concurrently and checks that all responses have error_code == 0. +func (p *XshardPool) SendXshardTx(ctx context.Context, fullShardID uint32, payload []byte) (*wire.Frame, error) { + conns := p.Get(fullShardID) + if len(conns) == 0 { + return nil, fmt.Errorf("no xshard connection to full shard %d", fullShardID) + } + + // Filter active connections + var activeConns []*XshardConn + for _, conn := range conns { + if conn.IsActive() && !conn.Closed() { + activeConns = append(activeConns, conn) + } + } + + if len(activeConns) == 0 { + return nil, fmt.Errorf("no live xshard connection to full shard %d", fullShardID) + } + + // Broadcast to all active connections concurrently (matches Python's asyncio.gather) + type result struct { + resp *wire.Frame + err error + } + results := make([]result, len(activeConns)) + var wg sync.WaitGroup + + for i, conn := range activeConns { + wg.Add(1) + go func(idx int, c *XshardConn) { + defer wg.Done() + resp, err := c.SendXshardTxList(ctx, payload) + results[idx] = result{resp: resp, err: err} + }(i, conn) + } + wg.Wait() + + // Check all responses (matches Python's check(all([response.error_code == 0 ...]))) + var firstErr error + var firstResp *wire.Frame + for _, r := range results { + if r.err != nil { + if firstErr == nil { + firstErr = r.err + } + continue + } + if firstResp == nil { + firstResp = r.resp + } + } + + if firstErr != nil { + return nil, firstErr + } + return firstResp, nil +} + +// TrackInbound registers an already-started inbound connection for lifecycle +// management. The pool will close it when Close is called. +// +// TrackInbound only handles lifecycle (close-on-shutdown). Use WatchAndIndex +// to additionally wait for identity exchange and index by shard for routing. +func (p *XshardPool) TrackInbound(conn *XshardConn) { + p.mu.Lock() + if p.closed { + p.mu.Unlock() + conn.Close() + p.log.Warn("xshard pool closed, closing inbound conn immediately", "remote", conn.RemoteAddr()) + return + } + p.inbound = append(p.inbound, conn) + p.mu.Unlock() + p.log.Info("tracked inbound xshard connection", "remote", conn.RemoteAddr()) +} + +// WatchAndIndex waits for the inbound connection to complete PING-based identity +// exchange, then indexes it by all remote shard IDs for routing purposes. +// It also registers the slave ID for deduplication. +// +// Returns false if the connection closes before identity exchange completes. +// The connection should already be tracked via TrackInbound before calling this. +func (p *XshardPool) WatchAndIndex(conn *XshardConn) bool { + if !conn.WaitUntilPingReceived() { + p.log.Warn("inbound xshard connection closed before ping", "remote", conn.RemoteAddr()) + return false + } + + remoteID := conn.RemoteID() + shardList := conn.RemoteFullShardIDList() + + p.mu.Lock() + if p.closed { + p.mu.Unlock() + conn.Close() + return false + } + + // Register slave ID for deduplication + if len(remoteID) > 0 { + p.slaveIDs[string(remoteID)] = true + } + + // Index by remote shard IDs for routing + for _, shardID := range shardList { + p.conns[shardID] = append(p.conns[shardID], conn) + } + p.mu.Unlock() + + p.log.Info("indexed inbound xshard connection", "remote_id", string(remoteID), "shards", shardList) + return true +} + +// Close closes all connections in the pool and prevents new additions. +func (p *XshardPool) Close() { + p.mu.Lock() + if p.closed { + p.mu.Unlock() + return + } + p.closed = true + + var allConns []*XshardConn + for _, conns := range p.conns { + allConns = append(allConns, conns...) + } + allConns = append(allConns, p.inbound...) + + p.conns = nil + p.inbound = nil + p.slaveIDs = nil + p.mu.Unlock() + + for _, conn := range allConns { + conn.Close() + } + p.log.Info("xshard pool closed", "connections", len(allConns)) +} + +// OutboundSize returns the number of outbound connections (indexed by shard ID). +func (p *XshardPool) OutboundSize() int { + p.mu.RLock() + defer p.mu.RUnlock() + + total := 0 + for _, conns := range p.conns { + total += len(conns) + } + return total +} + +// InboundSize returns the number of tracked inbound connections. +func (p *XshardPool) InboundSize() int { + p.mu.RLock() + defer p.mu.RUnlock() + return len(p.inbound) +} + +// Targets returns all full shard IDs that have outbound connections. +func (p *XshardPool) Targets() []uint32 { + p.mu.RLock() + defer p.mu.RUnlock() + + targets := make([]uint32, 0, len(p.conns)) + for id := range p.conns { + targets = append(targets, id) + } + return targets +} diff --git a/qkc/cluster/slave/xshard_test.go b/qkc/cluster/slave/xshard_test.go new file mode 100644 index 0000000000..1eb07985d3 --- /dev/null +++ b/qkc/cluster/slave/xshard_test.go @@ -0,0 +1,712 @@ +// Copyright 2026-2027, QuarkChain. + +package slave + +import ( + "context" + "fmt" + "net" + "sync" + "testing" + "time" + + "github.com/ethereum/go-ethereum/log" + "github.com/ethereum/go-ethereum/qkc/cluster/wire" + "github.com/ethereum/go-ethereum/qkc/serialize" +) + +// writeRawFrame writes a raw frame directly to the underlying TCP connection, +// bypassing the connection's frame writer. Used to craft malformed/invalid frames +// for protocol-validation tests. +func writeRawFrame(t *testing.T, conn net.Conn, frame *wire.Frame) { + t.Helper() + if err := wire.WriteFrameNoMeta(conn, frame); err != nil { + t.Fatalf("write raw frame: %v", err) + } +} + +// newTestConnPair creates a pair of XshardConns connected over a local TCP +// socket. The caller is responsible for calling cleanup. +func newTestConnPair(t *testing.T) (client, server *XshardConn, cleanup func()) { + t.Helper() + return newTestConnPairWithIdentity(t, []byte("client-slave"), []uint32{0x00010001}, []byte("server-slave"), []uint32{0x00030004}) +} + +func newTestConnPairWithIdentity(t *testing.T, clientID []byte, clientShards []uint32, serverID []byte, serverShards []uint32) (client, server *XshardConn, cleanup func()) { + t.Helper() + + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("listen: %v", err) + } + + var serverConn net.Conn + var acceptErr error + accepted := make(chan struct{}) + go func() { + defer close(accepted) + serverConn, acceptErr = ln.Accept() + ln.Close() + }() + + clientConn, err := net.Dial("tcp", ln.Addr().String()) + if err != nil { + t.Fatalf("dial: %v", err) + } + <-accepted + if acceptErr != nil { + t.Fatalf("accept: %v", acceptErr) + } + + logger := log.New() + client = NewXshardConnFromConn(clientConn, 0, clientID, clientShards, logger) // 0 = no limit (matches Python) + server = NewXshardConnFromConn(serverConn, 0, serverID, serverShards, logger) + cleanup = func() { + client.Close() + server.Close() + } + return +} + +// TestXshardConn_DefaultPingHandler verifies that PING is handled internally +// even when the server does not register a PING handler. The server still +// records peer identity and returns a PONG with its own identity. +func TestXshardConn_DefaultPingHandler(t *testing.T) { + clientID := []byte("client-slave") + clientShards := []uint32{0x00010001} + serverID := []byte("server-slave") + serverShards := []uint32{0x00030004} + + client, server, cleanup := newTestConnPairWithIdentity(t, clientID, clientShards, serverID, serverShards) + defer cleanup() + + // Server does NOT register any handler; PING should be handled internally. + server.Start() + client.Start() + + pingPayload, err := serialize.SerializeToBytes(&wire.PingRequest{ + ID: clientID, + FullShardIDList: clientShards, + RootTip: nil, + }) + if err != nil { + t.Fatalf("serialize ping: %v", err) + } + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + resp, err := client.SendRPC(ctx, byte(wire.ClusterOpPing), pingPayload) + if err != nil { + t.Fatalf("send ping rpc: %v", err) + } + if resp.Opcode != byte(wire.ClusterOpPong) { + t.Fatalf("expected opcode 0x%x, got 0x%x", wire.ClusterOpPong, resp.Opcode) + } + + var pong wire.PongResponse + if err := serialize.Deserialize(serialize.NewByteBuffer(resp.Payload), &pong); err != nil { + t.Fatalf("deserialize pong: %v", err) + } + if string(pong.ID) != string(serverID) { + t.Fatalf("pong id mismatch: got %s, expected %s", pong.ID, serverID) + } + if len(pong.FullShardIDList) != len(serverShards) { + t.Fatalf("pong shard list mismatch: got %v", pong.FullShardIDList) + } + + if !server.WaitUntilPingReceived() { + t.Fatal("server did not receive ping") + } + if string(server.RemoteID()) != string(clientID) { + t.Fatalf("server remote id mismatch: got %s", server.RemoteID()) + } +} + +func TestXshardConn_RPCRoundTrip(t *testing.T) { + client, server, cleanup := newTestConnPair(t) + defer cleanup() + + clientID := []byte("client-slave") + clientShards := []uint32{0x00010001, 0x00010002} + serverID := []byte("server-slave") + serverShards := []uint32{0x00030004} + + server.RegisterHandlers(map[byte]TypedHandler{ + byte(wire.ClusterOpPing): func(req any) (any, error) { + _ = req.(*wire.PingRequest) + return &wire.PongResponse{ + ID: serverID, + FullShardIDList: serverShards, + }, nil + }, + }) + server.Start() + client.Start() + + pingPayload, err := serialize.SerializeToBytes(&wire.PingRequest{ + ID: clientID, + FullShardIDList: clientShards, + RootTip: nil, // OK for SlaveConnection (master doesn't use it) + }) + if err != nil { + t.Fatalf("serialize ping: %v", err) + } + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + resp, err := client.SendRPC(ctx, byte(wire.ClusterOpPing), pingPayload) + if err != nil { + t.Fatalf("send ping rpc: %v", err) + } + if resp.Opcode != byte(wire.ClusterOpPong) { + t.Fatalf("expected opcode 0x%x, got 0x%x", wire.ClusterOpPong, resp.Opcode) + } + + var pong wire.PongResponse + if err := serialize.Deserialize(serialize.NewByteBuffer(resp.Payload), &pong); err != nil { + t.Fatalf("deserialize pong: %v", err) + } + if string(pong.ID) != string(serverID) { + t.Fatalf("pong id mismatch: got %s", pong.ID) + } + + if !server.WaitUntilPingReceived() { + t.Fatal("server did not receive ping") + } + if string(server.RemoteID()) != string(clientID) { + t.Fatalf("server remote id mismatch: got %s", server.RemoteID()) + } + if len(server.RemoteFullShardIDList()) != len(clientShards) { + t.Fatalf("server remote shard list mismatch: got %v", server.RemoteFullShardIDList()) + } +} + +// TestXshardConn_RejectEmptyShardList verifies that empty shard list causes +// connection close (Python's close_with_error behavior). The peer ID is still +// recorded before closing, matching Python's handle_ping. +func TestXshardConn_RejectEmptyShardList(t *testing.T) { + client, server, cleanup := newTestConnPair(t) + defer cleanup() + + server.RegisterHandlers(map[byte]TypedHandler{ + byte(wire.ClusterOpPing): func(req any) (any, error) { + // This handler won't be called because wrapper rejects empty shard list first. + t.Fatal("user handler should not be called for empty shard list") + return nil, nil + }, + }) + server.Start() + client.Start() + + pingPayload, err := serialize.SerializeToBytes(&wire.PingRequest{ + ID: []byte("bad-slave"), + FullShardIDList: []uint32{}, // empty list + RootTip: nil, + }) + if err != nil { + t.Fatalf("serialize ping: %v", err) + } + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + // Python: empty shard list causes close_with_error (connection close, no response). + _, err = client.SendRPC(ctx, byte(wire.ClusterOpPing), pingPayload) + if err == nil { + t.Fatal("expected error due to connection close, got nil") + } + // The error should be connection closed (readLoop returns after handler error). + if err != ErrConnectionClosed { + t.Logf("got error: %v (expected ErrConnectionClosed or timeout)", err) + } + + // Python: id is recorded BEFORE close_with_error is called. + // The wrapper records the id first, then checks shard list. + if string(server.RemoteID()) != "bad-slave" { + t.Fatalf("expected remote ID 'bad-slave', got %v", server.RemoteID()) + } +} + +// TestXshardConn_UnsupportedOpcodeClosesConnection verifies that unsupported +// opcode causes connection close (Python's close_with_error behavior). +func TestXshardConn_UnsupportedOpcodeClosesConnection(t *testing.T) { + client, server, cleanup := newTestConnPair(t) + defer cleanup() + + server.Start() + client.Start() + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + // Send a request for an opcode that has no handler. + _, err := client.SendRPC(ctx, byte(wire.ClusterOpAddRootBlockRequest), []byte("payload")) + if err == nil { + t.Fatal("expected error due to connection close, got nil") + } + // Connection should be closed by server due to unsupported opcode. + if err != ErrConnectionClosed { + t.Logf("got error: %v (expected ErrConnectionClosed or timeout)", err) + } +} + +// TestXshardConn_HandlerErrorClosesConnection verifies that handler error +// causes connection close (Python's close_with_error behavior). +func TestXshardConn_HandlerErrorClosesConnection(t *testing.T) { + client, server, cleanup := newTestConnPair(t) + defer cleanup() + + server.RegisterHandlers(map[byte]TypedHandler{ + byte(wire.ClusterOpAddRootBlockRequest): func(req any) (any, error) { + _ = req + return nil, fmt.Errorf("intentional error") //nolint:govet // test error + }, + }) + server.Start() + client.Start() + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + _, err := client.SendRPC(ctx, byte(wire.ClusterOpAddRootBlockRequest), []byte("payload")) + if err == nil { + t.Fatal("expected error due to connection close, got nil") + } +} + +// TestXshardConn_HandlerPanicClosesConnection verifies that handler panic +// causes connection close (Python's close_with_error behavior). +func TestXshardConn_HandlerPanicClosesConnection(t *testing.T) { + client, server, cleanup := newTestConnPair(t) + defer cleanup() + + server.RegisterHandlers(map[byte]TypedHandler{ + byte(wire.ClusterOpAddRootBlockRequest): func(req any) (any, error) { + _ = req + panic("intentional panic") + }, + }) + server.Start() + client.Start() + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + _, err := client.SendRPC(ctx, byte(wire.ClusterOpAddRootBlockRequest), []byte("payload")) + if err == nil { + t.Fatal("expected error due to connection close, got nil") + } +} + +// TestXshardConn_CloseWakesPendingRPC verifies that Close wakes all pending RPCs. +// Uses a sync channel instead of time.Sleep for reliable testing. +func TestXshardConn_CloseWakesPendingRPC(t *testing.T) { + client, _, cleanup := newTestConnPair(t) + defer cleanup() + + // Server intentionally left unstarted so it never replies. + client.Start() + + var wg sync.WaitGroup + wg.Add(1) + errChan := make(chan error, 1) + go func() { + wg.Done() // Signal that goroutine is ready + _, err := client.SendRPC(context.Background(), byte(wire.ClusterOpPing), []byte("ping")) + errChan <- err + }() + + wg.Wait() // Wait for goroutine to start (reliable synchronization) + client.Close() + + select { + case err := <-errChan: + if err != ErrConnectionClosed { + t.Fatalf("expected ErrConnectionClosed, got %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("pending RPC was not woken by Close") + } +} + +// TestXshardConn_SendXshardTxList verifies RPC mode for AddXshardTxListRequest. +// The handler must return a proper response (AddXshardTxListResponse). +func TestXshardConn_SendXshardTxList(t *testing.T) { + client, server, cleanup := newTestConnPair(t) + defer cleanup() + + server.RegisterHandlers(map[byte]TypedHandler{ + byte(wire.ClusterOpAddXshardTxListRequest): func(req any) (any, error) { + _ = req.(*wire.AddXshardTxListRequest) + // Return success response (Python: AddXshardTxListResponse(error_code=0)) + return &wire.AddXshardTxListResponse{ErrorCode: 0}, nil + }, + }) + server.Start() + client.Start() + + txList := wire.RawBytes([]byte("tx-list")) + req := &wire.AddXshardTxListRequest{ + Branch: 0x00010001, + MinorBlockHash: [32]byte{1, 2, 3}, + TxList: &txList, + } + payload, err := serialize.SerializeToBytes(req) + if err != nil { + t.Fatalf("serialize request: %v", err) + } + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + resp, err := client.SendXshardTxList(ctx, payload) + if err != nil { + t.Fatalf("send xshard tx list: %v", err) + } + if resp.Opcode != byte(wire.ClusterOpAddXshardTxListResponse) { + t.Fatalf("unexpected response opcode 0x%x", resp.Opcode) + } + + var xshardResp wire.AddXshardTxListResponse + if err := serialize.Deserialize(serialize.NewByteBuffer(resp.Payload), &xshardResp); err != nil { + t.Fatalf("deserialize response: %v", err) + } + if xshardResp.ErrorCode != 0 { + t.Fatalf("expected error_code 0, got %d", xshardResp.ErrorCode) + } +} + +// TestXshardConn_SendBatchXshardTxList verifies RPC mode for BatchAddXshardTxListRequest. +func TestXshardConn_SendBatchXshardTxList(t *testing.T) { + client, server, cleanup := newTestConnPair(t) + defer cleanup() + + server.RegisterHandlers(map[byte]TypedHandler{ + byte(wire.ClusterOpBatchAddXshardTxListRequest): func(req any) (any, error) { + _ = req.(*wire.BatchAddXshardTxListRequest) + return &wire.BatchAddXshardTxListResponse{ErrorCode: 0}, nil + }, + }) + server.Start() + client.Start() + + txList := wire.RawBytes([]byte("tx1")) + req := &wire.BatchAddXshardTxListRequest{ + AddXshardTxListRequestList: []wire.AddXshardTxListRequest{ + {Branch: 0x00010001, MinorBlockHash: [32]byte{1, 2, 3}, TxList: &txList}, + }, + } + payload, err := serialize.SerializeToBytes(req) + if err != nil { + t.Fatalf("serialize request: %v", err) + } + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + resp, err := client.SendBatchXshardTxList(ctx, payload) + if err != nil { + t.Fatalf("send batch xshard tx list: %v", err) + } + if resp.Opcode != byte(wire.ClusterOpBatchAddXshardTxListResponse) { + t.Fatalf("unexpected response opcode 0x%x", resp.Opcode) + } + + var batchResp wire.BatchAddXshardTxListResponse + if err := serialize.Deserialize(serialize.NewByteBuffer(resp.Payload), &batchResp); err != nil { + t.Fatalf("deserialize response: %v", err) + } + if batchResp.ErrorCode != 0 { + t.Fatalf("expected error_code 0, got %d", batchResp.ErrorCode) + } +} + +func TestXshardPool_AddGetRemove(t *testing.T) { + pool := NewXshardPool(log.New()) + defer pool.Close() + + // Use stub connections that are never started. + _, conn1, cleanup1 := newTestConnPair(t) + defer cleanup1() + _, conn2, cleanup2 := newTestConnPair(t) + defer cleanup2() + + pool.Add(0x00010001, conn1) + pool.Add(0x00010001, conn2) + pool.Add(0x00020001, conn1) + + if got := pool.OutboundSize(); got != 3 { + t.Fatalf("expected pool outbound size 3, got %d", got) + } + + conns := pool.Get(0x00010001) + if len(conns) != 2 { + t.Fatalf("expected 2 conns for shard 0x00010001, got %d", len(conns)) + } + + pool.Remove(0x00010001, conn1) + if got := pool.OutboundSize(); got != 2 { + t.Fatalf("expected pool outbound size 2 after remove, got %d", got) + } + conns = pool.Get(0x00010001) + if len(conns) != 1 || conns[0] != conn2 { + t.Fatalf("expected only conn2 for shard 0x00010001") + } + + targets := pool.Targets() + if len(targets) != 2 { + t.Fatalf("expected 2 targets, got %d", len(targets)) + } +} + +func TestXshardPool_RemoveTargetClosesConnections(t *testing.T) { + pool := NewXshardPool(log.New()) + defer pool.Close() + + _, conn, cleanup := newTestConnPair(t) + defer cleanup() + + conn.Start() + pool.Add(0x00010001, conn) + pool.RemoveTarget(0x00010001) + + if pool.OutboundSize() != 0 { + t.Fatalf("expected pool outbound size 0, got %d", pool.OutboundSize()) + } + + // A closed connection rejects further RPCs. + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + _, err := conn.SendRPC(ctx, byte(wire.ClusterOpPing), []byte("ping")) + if err != ErrConnectionClosed { + t.Fatalf("expected ErrConnectionClosed, got %v", err) + } +} + +func TestXshardPool_TrackInboundClose(t *testing.T) { + pool := NewXshardPool(log.New()) + + _, conn, cleanup := newTestConnPair(t) + defer cleanup() + + conn.Start() + pool.TrackInbound(conn) + pool.Close() + + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + _, err := conn.SendRPC(ctx, byte(wire.ClusterOpPing), []byte("ping")) + if err != ErrConnectionClosed { + t.Fatalf("expected ErrConnectionClosed after pool close, got %v", err) + } +} + +func TestXshardPool_SendXshardTxNoConnection(t *testing.T) { + pool := NewXshardPool(log.New()) + defer pool.Close() + + ctx := context.Background() + _, err := pool.SendXshardTx(ctx, 0x00010001, []byte("tx")) + if err == nil { + t.Fatal("expected error when no connection exists") + } +} + +func TestXshardPool_ClosedPoolRejectsAdd(t *testing.T) { + pool := NewXshardPool(log.New()) + pool.Close() + + _, conn, cleanup := newTestConnPair(t) + defer cleanup() + + conn.Start() + pool.Add(0x00010001, conn) + + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + _, err := conn.SendRPC(ctx, byte(wire.ClusterOpPing), []byte("ping")) + if err != ErrConnectionClosed { + t.Fatalf("expected ErrConnectionClosed, got %v", err) + } +} + +// TestXshardConn_RPCIDMonotonic verifies RPC ID monotonic validation. +// Sending a duplicate RPC ID causes the server to close the connection. +func TestXshardConn_RPCIDMonotonic(t *testing.T) { + client, server, cleanup := newTestConnPair(t) + defer cleanup() + + server.RegisterHandlers(map[byte]TypedHandler{ + byte(wire.ClusterOpPing): func(req any) (any, error) { + _ = req.(*wire.PingRequest) + return &wire.PongResponse{ + ID: []byte("server"), + FullShardIDList: []uint32{0x00030004}, + }, nil + }, + }) + server.Start() + client.Start() + + pingPayload, _ := serialize.SerializeToBytes(&wire.PingRequest{ + ID: []byte("client"), + FullShardIDList: []uint32{0x00010001}, + }) + + // Manually send two PING frames with the same RPC ID (=1). + writeRawFrame(t, client.conn, &wire.Frame{ + Opcode: byte(wire.ClusterOpPing), + RPCID: 1, + Payload: pingPayload, + }) + writeRawFrame(t, client.conn, &wire.Frame{ + Opcode: byte(wire.ClusterOpPing), + RPCID: 1, // duplicate rpc_id: should trigger close + Payload: pingPayload, + }) + + // Wait for server to close the connection. + select { + case <-server.WaitUntilClosed(): + case <-time.After(2 * time.Second): + t.Fatal("server did not close connection after duplicate rpc_id") + } + + if !server.IsClosed() { + t.Fatal("server should be closed") + } +} + +// TestXshardConn_RPCIDDecreasing verifies that a decreasing RPC ID closes the connection. +func TestXshardConn_RPCIDDecreasing(t *testing.T) { + client, server, cleanup := newTestConnPair(t) + defer cleanup() + + server.RegisterHandlers(map[byte]TypedHandler{ + byte(wire.ClusterOpPing): func(req any) (any, error) { + _ = req.(*wire.PingRequest) + return &wire.PongResponse{ + ID: []byte("server"), + FullShardIDList: []uint32{0x00030004}, + }, nil + }, + }) + server.Start() + client.Start() + + pingPayload, _ := serialize.SerializeToBytes(&wire.PingRequest{ + ID: []byte("client"), + FullShardIDList: []uint32{0x00010001}, + }) + + // Send rpc_id=2 then rpc_id=1 (decreasing). + writeRawFrame(t, client.conn, &wire.Frame{ + Opcode: byte(wire.ClusterOpPing), + RPCID: 2, + Payload: pingPayload, + }) + writeRawFrame(t, client.conn, &wire.Frame{ + Opcode: byte(wire.ClusterOpPing), + RPCID: 1, // decreasing rpc_id: should trigger close + Payload: pingPayload, + }) + + select { + case <-server.WaitUntilClosed(): + case <-time.After(2 * time.Second): + t.Fatal("server did not close connection after decreasing rpc_id") + } +} + +// TestXshardConn_MultipleRPCs verifies multiple sequential RPCs work correctly. +func TestXshardConn_MultipleRPCs(t *testing.T) { + client, server, cleanup := newTestConnPair(t) + defer cleanup() + + callCount := 0 + server.RegisterHandlers(map[byte]TypedHandler{ + byte(wire.ClusterOpPing): func(req any) (any, error) { + _ = req.(*wire.PingRequest) + callCount++ + return &wire.PongResponse{ + ID: []byte("server"), + FullShardIDList: []uint32{0x00010001}, + }, nil + }, + }) + server.Start() + client.Start() + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + pingPayload, _ := serialize.SerializeToBytes(&wire.PingRequest{ + ID: []byte("client"), + FullShardIDList: []uint32{0x00010001}, + }) + + // Send multiple RPCs in sequence. + for i := 0; i < 5; i++ { + _, err := client.SendRPC(ctx, byte(wire.ClusterOpPing), pingPayload) + if err != nil { + t.Fatalf("rpc %d failed: %v", i+1, err) + } + } + + if callCount != 5 { + t.Fatalf("expected 5 handler calls, got %d", callCount) + } +} + +// TestXshardConn_RecordPingOnlyOnce verifies that recordPing only updates +// on first PING (matches Python's handle_ping behavior). +func TestXshardConn_RecordPingOnlyOnce(t *testing.T) { + client, server, cleanup := newTestConnPair(t) + defer cleanup() + + server.RegisterHandlers(map[byte]TypedHandler{ + byte(wire.ClusterOpPing): func(req any) (any, error) { + _ = req.(*wire.PingRequest) + return &wire.PongResponse{ + ID: []byte("server"), + FullShardIDList: []uint32{0x00010001}, + }, nil + }, + }) + server.Start() + client.Start() + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + // First PING with one shard list. + ping1, _ := serialize.SerializeToBytes(&wire.PingRequest{ + ID: []byte("client1"), + FullShardIDList: []uint32{0x00010001, 0x00010002}, + }) + _, err := client.SendRPC(ctx, byte(wire.ClusterOpPing), ping1) + if err != nil { + t.Fatalf("first ping failed: %v", err) + } + + firstID := server.RemoteID() + firstShards := server.RemoteFullShardIDList() + + // Second PING with different shard list (should not overwrite). + ping2, _ := serialize.SerializeToBytes(&wire.PingRequest{ + ID: []byte("client2"), + FullShardIDList: []uint32{0x00030004}, + }) + _, err = client.SendRPC(ctx, byte(wire.ClusterOpPing), ping2) + if err != nil { + t.Fatalf("second ping failed: %v", err) + } + + // RemoteID and RemoteFullShardIDList should NOT have changed. + if string(server.RemoteID()) != string(firstID) { + t.Fatalf("remote ID changed: got %s, expected %s", server.RemoteID(), firstID) + } + if len(server.RemoteFullShardIDList()) != len(firstShards) { + t.Fatalf("remote shard list changed: got %v, expected %v", server.RemoteFullShardIDList(), firstShards) + } +}