raft/rpc.go

package raft

import (
	"context"
	"fmt"
	"io"
	"net"
	"sync"
	"time"
)

// Transport defines the interface for RPC communication
type Transport interface {
	// Start starts the transport
	Start() error

	// Stop stops the transport
	Stop() error

	// RequestVote sends a RequestVote RPC to the target node
	RequestVote(ctx context.Context, target string, args *RequestVoteArgs) (*RequestVoteReply, error)

	// AppendEntries sends an AppendEntries RPC to the target node
	AppendEntries(ctx context.Context, target string, args *AppendEntriesArgs) (*AppendEntriesReply, error)

	// InstallSnapshot sends an InstallSnapshot RPC to the target node
	InstallSnapshot(ctx context.Context, target string, args *InstallSnapshotArgs) (*InstallSnapshotReply, error)

	// ForwardPropose forwards a propose request to the leader
	ForwardPropose(ctx context.Context, target string, args *ProposeArgs) (*ProposeReply, error)

	// ForwardAddNode forwards an AddNode request to the leader
	ForwardAddNode(ctx context.Context, target string, args *AddNodeArgs) (*AddNodeReply, error)

	// ForwardRemoveNode forwards a RemoveNode request to the leader
	ForwardRemoveNode(ctx context.Context, target string, args *RemoveNodeArgs) (*RemoveNodeReply, error)

	// TimeoutNow sends a TimeoutNow RPC for leadership transfer
	TimeoutNow(ctx context.Context, target string, args *TimeoutNowArgs) (*TimeoutNowReply, error)

	// ReadIndex sends a ReadIndex RPC for linearizable reads
	ReadIndex(ctx context.Context, target string, args *ReadIndexArgs) (*ReadIndexReply, error)

	// ForwardGet sends a Get RPC for remote KV reads
	ForwardGet(ctx context.Context, target string, args *GetArgs) (*GetReply, error)

	// SetRPCHandler sets the handler for incoming RPCs
	SetRPCHandler(handler RPCHandler)
}

// RPCHandler handles incoming RPCs
type RPCHandler interface {
	HandleRequestVote(args *RequestVoteArgs) *RequestVoteReply
	HandleAppendEntries(args *AppendEntriesArgs) *AppendEntriesReply
	HandleInstallSnapshot(args *InstallSnapshotArgs) *InstallSnapshotReply
	HandlePropose(args *ProposeArgs) *ProposeReply
	HandleAddNode(args *AddNodeArgs) *AddNodeReply
	HandleRemoveNode(args *RemoveNodeArgs) *RemoveNodeReply
	HandleTimeoutNow(args *TimeoutNowArgs) *TimeoutNowReply
	HandleReadIndex(args *ReadIndexArgs) *ReadIndexReply
	HandleGet(args *GetArgs) *GetReply
}

// TCPTransport implements Transport using raw TCP with binary protocol
// This is more efficient than HTTP for high-frequency RPCs
type TCPTransport struct {
	mu sync.RWMutex

	localAddr  string
	handler    RPCHandler
	logger     Logger
	listener   net.Listener
	shutdownCh chan struct{}

	// Connection pool
	connPool map[string]chan net.Conn
	poolSize int
}

// NewTCPTransport creates a new TCP transport
func NewTCPTransport(localAddr string, poolSize int, logger Logger) *TCPTransport {
	if logger == nil {
		logger = &NoopLogger{}
	}
	if poolSize <= 0 {
		poolSize = 5
	}

	return &TCPTransport{
		localAddr:  localAddr,
		logger:     logger,
		shutdownCh: make(chan struct{}),
		connPool:   make(map[string]chan net.Conn),
		poolSize:   poolSize,
	}
}

// SetRPCHandler sets the handler for incoming RPCs
func (t *TCPTransport) SetRPCHandler(handler RPCHandler) {
	t.mu.Lock()
	defer t.mu.Unlock()
	t.handler = handler
}

// Start starts the TCP server
func (t *TCPTransport) Start() error {
	var err error
	t.listener, err = net.Listen("tcp", t.localAddr)
	if err != nil {
		return fmt.Errorf("failed to listen on %s: %w", t.localAddr, err)
	}

	go t.acceptLoop()

	t.logger.Info("TCP Transport started on %s", t.localAddr)
	return nil
}

// acceptLoop accepts incoming connections
func (t *TCPTransport) acceptLoop() {
	for {
		select {
		case <-t.shutdownCh:
			return
		default:
		}

		conn, err := t.listener.Accept()
		if err != nil {
			select {
			case <-t.shutdownCh:
				return
			default:
				t.logger.Error("Accept error: %v", err)
				continue
			}
		}

		go t.handleConnection(conn)
	}
}

// handleConnection handles an incoming connection
func (t *TCPTransport) handleConnection(conn net.Conn) {
	defer conn.Close()

	for {
		select {
		case <-t.shutdownCh:
			return
		default:
		}

		// Set read deadline
		conn.SetReadDeadline(time.Now().Add(30 * time.Second))

		// Read message type (1 byte)
		typeBuf := make([]byte, 1)
		if _, err := io.ReadFull(conn, typeBuf); err != nil {
			if err != io.EOF {
				t.logger.Debug("Read type error: %v", err)
			}
			return
		}

		// Read message length (4 bytes)
		lenBuf := make([]byte, 4)
		if _, err := io.ReadFull(conn, lenBuf); err != nil {
			t.logger.Debug("Read length error: %v", err)
			return
		}

		length := uint32(lenBuf[0])<<24 | uint32(lenBuf[1])<<16 | uint32(lenBuf[2])<<8 | uint32(lenBuf[3])
		if length > 10*1024*1024 { // 10MB limit
			t.logger.Error("Message too large: %d", length)
			return
		}

		// Read message body
		body := make([]byte, length)
		if _, err := io.ReadFull(conn, body); err != nil {
			t.logger.Debug("Read body error: %v", err)
			return
		}

		// Handle message
		var response []byte
		var err error

		t.mu.RLock()
		handler := t.handler
		t.mu.RUnlock()

		if handler == nil {
			t.logger.Error("No handler registered")
			return
		}

		switch RPCType(typeBuf[0]) {
		case RPCRequestVote:
			var args RequestVoteArgs
			if err := DefaultCodec.Unmarshal(body, &args); err != nil {
				t.logger.Error("Unmarshal RequestVote error: %v", err)
				return
			}
			reply := handler.HandleRequestVote(&args)
			response, err = DefaultCodec.Marshal(reply)

		case RPCAppendEntries:
			var args AppendEntriesArgs
			if err := DefaultCodec.Unmarshal(body, &args); err != nil {
				t.logger.Error("Unmarshal AppendEntries error: %v", err)
				return
			}
			reply := handler.HandleAppendEntries(&args)
			response, err = DefaultCodec.Marshal(reply)

		case RPCInstallSnapshot:
			var args InstallSnapshotArgs
			if err := DefaultCodec.Unmarshal(body, &args); err != nil {
				t.logger.Error("Unmarshal InstallSnapshot error: %v", err)
				return
			}
			reply := handler.HandleInstallSnapshot(&args)
			response, err = DefaultCodec.Marshal(reply)

		case RPCPropose:
			var args ProposeArgs
			if err := DefaultCodec.Unmarshal(body, &args); err != nil {
				t.logger.Error("Unmarshal Propose error: %v", err)
				return
			}
			reply := handler.HandlePropose(&args)
			response, err = DefaultCodec.Marshal(reply)

		case RPCAddNode:
			var args AddNodeArgs
			if err := DefaultCodec.Unmarshal(body, &args); err != nil {
				t.logger.Error("Unmarshal AddNode error: %v", err)
				return
			}
			reply := handler.HandleAddNode(&args)
			response, err = DefaultCodec.Marshal(reply)

		case RPCRemoveNode:
			var args RemoveNodeArgs
			if err := DefaultCodec.Unmarshal(body, &args); err != nil {
				t.logger.Error("Unmarshal RemoveNode error: %v", err)
				return
			}
			reply := handler.HandleRemoveNode(&args)
			response, err = DefaultCodec.Marshal(reply)

		case RPCTimeoutNow:
			var args TimeoutNowArgs
			if err := DefaultCodec.Unmarshal(body, &args); err != nil {
				t.logger.Error("Unmarshal TimeoutNow error: %v", err)
				return
			}
			reply := handler.HandleTimeoutNow(&args)
			response, err = DefaultCodec.Marshal(reply)

		case RPCReadIndex:
			var args ReadIndexArgs
			if err := DefaultCodec.Unmarshal(body, &args); err != nil {
				t.logger.Error("Unmarshal ReadIndex error: %v", err)
				return
			}
			reply := handler.HandleReadIndex(&args)
			response, err = DefaultCodec.Marshal(reply)

		case RPCGet:
			var args GetArgs
			if err := DefaultCodec.Unmarshal(body, &args); err != nil {
				t.logger.Error("Unmarshal Get error: %v", err)
				return
			}
			reply := handler.HandleGet(&args)
			response, err = DefaultCodec.Marshal(reply)

		default:
			t.logger.Error("Unknown RPC type: %d", typeBuf[0])
			return
		}

		if err != nil {
			t.logger.Error("Marshal response error: %v", err)
			return
		}

		// Write response
		conn.SetWriteDeadline(time.Now().Add(5 * time.Second))

		respLen := make([]byte, 4)
		respLen[0] = byte(len(response) >> 24)
		respLen[1] = byte(len(response) >> 16)
		respLen[2] = byte(len(response) >> 8)
		respLen[3] = byte(len(response))

		if _, err := conn.Write(respLen); err != nil {
			t.logger.Debug("Write response length error: %v", err)
			return
		}
		if _, err := conn.Write(response); err != nil {
			t.logger.Debug("Write response error: %v", err)
			return
		}
	}
}

// Stop stops the TCP server
func (t *TCPTransport) Stop() error {
	close(t.shutdownCh)

	// Close all pooled connections
	t.mu.Lock()
	for _, pool := range t.connPool {
		close(pool)
		for conn := range pool {
			conn.Close()
		}
	}
	t.connPool = make(map[string]chan net.Conn)
	t.mu.Unlock()

	if t.listener != nil {
		return t.listener.Close()
	}
	return nil
}

// getConn gets a connection from the pool or creates a new one
func (t *TCPTransport) getConn(target string) (net.Conn, error) {
	t.mu.Lock()
	pool, ok := t.connPool[target]
	if !ok {
		pool = make(chan net.Conn, t.poolSize)
		t.connPool[target] = pool
	}
	t.mu.Unlock()

	select {
	case conn := <-pool:
		return conn, nil
	default:
		return net.DialTimeout("tcp", target, 5*time.Second)
	}
}

// putConn returns a connection to the pool
func (t *TCPTransport) putConn(target string, conn net.Conn) {
	t.mu.RLock()
	pool, ok := t.connPool[target]
	t.mu.RUnlock()

	if !ok {
		conn.Close()
		return
	}

	select {
	case pool <- conn:
	default:
		conn.Close()
	}
}

// sendTCPRPC sends an RPC over TCP
func (t *TCPTransport) sendTCPRPC(ctx context.Context, target string, rpcType RPCType, args interface{}, reply interface{}) error {
	conn, err := t.getConn(target)
	if err != nil {
		return fmt.Errorf("failed to get connection: %w", err)
	}

	data, err := DefaultCodec.Marshal(args)
	if err != nil {
		conn.Close()
		return fmt.Errorf("failed to marshal request: %w", err)
	}

	// Set deadline from context
	deadline, ok := ctx.Deadline()
	if !ok {
		deadline = time.Now().Add(5 * time.Second)
	}
	conn.SetDeadline(deadline)

	// Write message: [type(1)][length(4)][body]
	header := make([]byte, 5)
	header[0] = byte(rpcType)
	header[1] = byte(len(data) >> 24)
	header[2] = byte(len(data) >> 16)
	header[3] = byte(len(data) >> 8)
	header[4] = byte(len(data))

	if _, err := conn.Write(header); err != nil {
		conn.Close()
		return fmt.Errorf("failed to write header: %w", err)
	}
	if _, err := conn.Write(data); err != nil {
		conn.Close()
		return fmt.Errorf("failed to write body: %w", err)
	}

	// Read response length
	lenBuf := make([]byte, 4)
	if _, err := io.ReadFull(conn, lenBuf); err != nil {
		conn.Close()
		return fmt.Errorf("failed to read response length: %w", err)
	}

	length := uint32(lenBuf[0])<<24 | uint32(lenBuf[1])<<16 | uint32(lenBuf[2])<<8 | uint32(lenBuf[3])
	if length > 10*1024*1024 {
		conn.Close()
		return fmt.Errorf("response too large: %d", length)
	}

	// Read response body
	respBody := make([]byte, length)
	if _, err := io.ReadFull(conn, respBody); err != nil {
		conn.Close()
		return fmt.Errorf("failed to read response body: %w", err)
	}

	if err := DefaultCodec.Unmarshal(respBody, reply); err != nil {
		conn.Close()
		return fmt.Errorf("failed to unmarshal response: %w", err)
	}

	// Return connection to pool
	t.putConn(target, conn)
	return nil
}

// RequestVote sends a RequestVote RPC
func (t *TCPTransport) RequestVote(ctx context.Context, target string, args *RequestVoteArgs) (*RequestVoteReply, error) {
	var reply RequestVoteReply
	err := t.sendTCPRPC(ctx, target, RPCRequestVote, args, &reply)
	if err != nil {
		return nil, err
	}
	return &reply, nil
}

// AppendEntries sends an AppendEntries RPC
func (t *TCPTransport) AppendEntries(ctx context.Context, target string, args *AppendEntriesArgs) (*AppendEntriesReply, error) {
	var reply AppendEntriesReply
	err := t.sendTCPRPC(ctx, target, RPCAppendEntries, args, &reply)
	if err != nil {
		return nil, err
	}
	return &reply, nil
}

// InstallSnapshot sends an InstallSnapshot RPC
func (t *TCPTransport) InstallSnapshot(ctx context.Context, target string, args *InstallSnapshotArgs) (*InstallSnapshotReply, error) {
	var reply InstallSnapshotReply
	err := t.sendTCPRPC(ctx, target, RPCInstallSnapshot, args, &reply)
	if err != nil {
		return nil, err
	}
	return &reply, nil
}

// ForwardPropose forwards a propose request to the leader
func (t *TCPTransport) ForwardPropose(ctx context.Context, target string, args *ProposeArgs) (*ProposeReply, error) {
	var reply ProposeReply
	err := t.sendTCPRPC(ctx, target, RPCPropose, args, &reply)
	if err != nil {
		return nil, err
	}
	return &reply, nil
}

// ForwardAddNode forwards an AddNode request to the leader
func (t *TCPTransport) ForwardAddNode(ctx context.Context, target string, args *AddNodeArgs) (*AddNodeReply, error) {
	var reply AddNodeReply
	err := t.sendTCPRPC(ctx, target, RPCAddNode, args, &reply)
	if err != nil {
		return nil, err
	}
	return &reply, nil
}

// ForwardRemoveNode forwards a RemoveNode request to the leader
func (t *TCPTransport) ForwardRemoveNode(ctx context.Context, target string, args *RemoveNodeArgs) (*RemoveNodeReply, error) {
	var reply RemoveNodeReply
	err := t.sendTCPRPC(ctx, target, RPCRemoveNode, args, &reply)
	if err != nil {
		return nil, err
	}
	return &reply, nil
}

// TimeoutNow sends a TimeoutNow RPC for leadership transfer
func (t *TCPTransport) TimeoutNow(ctx context.Context, target string, args *TimeoutNowArgs) (*TimeoutNowReply, error) {
	var reply TimeoutNowReply
	err := t.sendTCPRPC(ctx, target, RPCTimeoutNow, args, &reply)
	if err != nil {
		return nil, err
	}
	return &reply, nil
}

// ReadIndex sends a ReadIndex RPC for linearizable reads
func (t *TCPTransport) ReadIndex(ctx context.Context, target string, args *ReadIndexArgs) (*ReadIndexReply, error) {
	var reply ReadIndexReply
	err := t.sendTCPRPC(ctx, target, RPCReadIndex, args, &reply)
	if err != nil {
		return nil, err
	}
	return &reply, nil
}

// ForwardGet sends a Get RPC for remote KV reads
func (t *TCPTransport) ForwardGet(ctx context.Context, target string, args *GetArgs) (*GetReply, error) {
	var reply GetReply
	err := t.sendTCPRPC(ctx, target, RPCGet, args, &reply)
	if err != nil {
		return nil, err
	}
	return &reply, nil
}