implement RpcConn and XshardConn compatibility layer

This commit is contained in:
iteye 2026-07-10 11:17:39 +08:00
parent 91950ad909
commit a0a8b7004a
10 changed files with 2378 additions and 0 deletions

View file

@ -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:<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:<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())
}
}

View file

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

View file

@ -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")
)

View file

View file

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

View file

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

View file

@ -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:<port> Printed when listening
PONG_OK id=<hex> 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()

View file

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

View file

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

View file

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