Implementing a Zero-Copy Message Broker in Go Using Linux splice(2) and TCP Zero-Copy — code snippets

snippet-1.go

package broker

import (
	"fmt"
	"sync"
	"golang.org/x/sys/unix"
)

// Pipe represents a pair of kernel pipe file descriptors used for splicing.
type Pipe struct {
	ReadFd  int
	WriteFd int
}

// Close releases the pipe descriptors back to the OS.
func (p *Pipe) Close() {
	_ = unix.Close(p.ReadFd)
	_ = unix.Close(p.WriteFd)
}

// PipePool manages a pool of pre-allocated Linux pipes to avoid syscall overhead during active routing.
type PipePool struct {
	pool sync.Pool
}

// NewPipePool initializes a pool of pipes.
func NewPipePool() *PipePool {
	return &PipePool{
		pool: sync.Pool{
			New: func() interface{} {
				var fds [2]int
				// Create the pipe with O_NONBLOCK and O_CLOEXEC to integrate cleanly with Go's netpoller
				err := unix.Pipe2(fds[:], unix.O_NONBLOCK|unix.O_CLOEXEC)
				if err != nil {
					panic(fmt.Sprintf("failed to allocate kernel pipe: %v", err))
				}
				
				// Maximize pipe capacity (default is typically 64KB, increase to 1MB for BDP tuning)
				// /proc/sys/fs/pipe-max-size defines the system-wide limit for non-root users
				_, _, _ = unix.Syscall(
					unix.SYS_FCNTL,
					uintptr(fds[0]),
					uintptr(unix.F_SETPIPE_SZ),
					uintptr(1024*1024), // 1MB pipe capacity
				)

				return &Pipe{
					ReadFd:  fds[0],
					WriteFd: fds[1],
				}
			},
		},
	}
}

// Get retrieves an active Pipe from the pool.
func (p *PipePool) Get() *Pipe {
	return p.pool.Get().(*Pipe)
}

// Put returns a Pipe back to the pool.
func (p *PipePool) Put(pipe *Pipe) {
	p.pool.Put(pipe)
}

snippet-2.go

package broker

import (
	"errors"
	"net"
	"syscall"
)

// getRawFd extracts the underlying system file descriptor from a TCP connection.
func getRawFd(conn *net.TCPConn) (syscall.RawConn, error) {
	if conn == nil {
		return nil, errors.New("connection is nil")
	}
	
	rawConn, err := conn.SyscallConn()
	if err != nil {
		return nil, err
	}
	
	return rawConn, nil
}

snippet-3.go

package broker

import (
	"io"
	"syscall"
	"golang.org/x/sys/unix"
)

// SpliceTransfer handles zero-copy routing between two TCP connections.
func SpliceTransfer(srcRaw, dstRaw syscall.RawConn, pipe *Pipe, bytesToTransfer int64) (int64, error) {
	var totalSpliced int64 = 0

	// We must loop until all requested bytes are moved or an error occurs.
	for totalSpliced < bytesToTransfer {
		remaining := bytesToTransfer - totalSpliced
		chunkSize := remaining
		if chunkSize > 1024*1024 { // Cap chunks to match our 1MB pipe capacity
			chunkSize = 1024 * 1024
		}

		var splicedIn int
		var spliceInErr error

		// Step 1: Splice from the Source Socket into the Pipe Write FD
		err := srcRaw.Read(func(srcFd uintptr) bool {
			// SPLICE_F_NONBLOCK ensures the call doesn't block the OS thread
			// SPLICE_F_MOVE hints to the kernel that we want to move memory pages instead of copying
			n, err := unix.Splice(
				int(srcFd),
				nil,
				pipe.WriteFd,
				nil,
				int(chunkSize),
				unix.SPLICE_F_NONBLOCK|unix.SPLICE_F_MOVE,
			)
			if err != nil {
				if err == unix.EAGAIN || err == unix.EWOULDBLOCK {
					// Socket is not ready yet; yield to netpoller
					return false
				}
				spliceInErr = err
				return true
			}
			if n == 0 {
				spliceInErr = io.EOF
				return true
			}
			splicedIn = int(n)
			return true
		})

		if err != nil {
			return totalSpliced, err
		}
		if spliceInErr != nil {
			return totalSpliced, spliceInErr
		}

		// Step 2: Splice from the Pipe Read FD into the Destination Socket
		var splicedOut int
		var spliceOutErr error
		bytesInPipe := splicedIn

		for bytesInPipe > 0 {
			err = dstRaw.Write(func(dstFd uintptr) bool {
				n, err := unix.Splice(
					pipe.ReadFd,
					nil,
					int(dstFd),
					nil,
					bytesInPipe,
					unix.SPLICE_F_NONBLOCK|unix.SPLICE_F_MOVE,
				)
				if err != nil {
					if err == unix.EAGAIN || err == unix.EWOULDBLOCK {
						// Destination socket queue is full; yield to netpoller
						return false
					}
					spliceOutErr = err
					return true
				}
				if n == 0 {
					spliceOutErr = io.ErrUnexpectedEOF
					return true
				}
				splicedOut = int(n)
				return true
			})

			if err != nil {
				return totalSpliced, err
			}
			if spliceOutErr != nil {
				return totalSpliced, spliceOutErr
			}

			bytesInPipe -= splicedOut
			totalSpliced += int64(splicedOut)
		}
	}

	return totalSpliced, nil
}

snippet-4.go

package broker

import (
	"encoding/binary"
	"errors"
	"io"
	"net"
)

const HeaderSize = 12

// FrameHeader defines our simple broker wire protocol:
// Offset 0-3: Magic Bytes (0xDEADBEEF)
// Offset 4-7: Topic ID (uint32)
// Offset 8-11: Payload Length (uint32)
type FrameHeader struct {
	TopicID       uint32
	PayloadLength uint32
}

// ReadHeader parses the fixed-size packet header from the reader.
func ReadHeader(r io.Reader) (FrameHeader, error) {
	var buf [HeaderSize]byte
	_, err := io.ReadFull(r, buf[:])
	if err != nil {
		return FrameHeader{}, err
	}

	magic := binary.BigEndian.Uint32(buf[0:4])
	if magic != 0xDEADBEEF {
		return FrameHeader{}, errors.New("invalid protocol magic byte sequence")
	}

	return FrameHeader{
		TopicID:       binary.BigEndian.Uint32(buf[4:8]),
		PayloadLength: binary.BigEndian.Uint32(buf[8:12]),
	}, nil
}

// RouteMessage coordinates reading the header and splicing the payload to the subscriber.
func RouteMessage(pubConn, subConn *net.TCPConn, pipePool *PipePool) (int64, error) {
	// Parse the frame header in user space
	header, err := ReadHeader(pubConn)
	if err != nil {
		return 0, err
	}

	if header.PayloadLength == 0 {
		return 0, nil // Nothing to splice
	}

	// Extract raw system connections
	pubRaw, err := getRawFd(pubConn)
	if err != nil {
		return 0, err
	}

	subRaw, err := getRawFd(subConn)
	if err != nil {
		return 0, err
	}

	// Retrieve a pipe from our pool
	pipe := pipePool.Get()
	defer pipePool.Put(pipe)

	// Perform the zero-copy splice transfer
	n, err := SpliceTransfer(pubRaw, subRaw, pipe, int64(header.PayloadLength))
	return n, err
}

snippet-5.go

package broker

import (
	"errors"
	"syscall"
	"unsafe"
	"golang.org/x/sys/unix"
)

// EnableTCPZeroCopy configures a connection's socket to use MSG_ZEROCOPY for transmissions.
func EnableTCPZeroCopy(rawConn syscall.RawConn) error {
	var sockOptErr error
	err := rawConn.Control(func(fd uintptr) {
		sockOptErr = unix.SetsockoptInt(int(fd), unix.SOL_SOCKET, unix.SO_ZEROCOPY, 1)
	})
	if err != nil {
		return err
	}
	return sockOptErr
}

// TxCompletion represents the range of sequence numbers confirmed by the network card.
type TxCompletion struct {
	Lowest  uint32
	Highest uint32
}

// ReadTxCompletions reads transmission notifications from the socket error queue.
// This is required to know when pinned user-space buffers can be safely reused.
func ReadTxCompletions(rawConn syscall.RawConn) (TxCompletion, error) {
	var comp TxCompletion
	var recvErr error

	err := rawConn.Control(func(fd uintptr) {
		// Use a 512-byte buffer for the control message (cmsg) containing the sock_extended_err struct
		cmsgBuf := make([]byte, 512)
		
		// Set up msghdr struct to receive from the MSG_ERRQUEUE
		var msg unix.Msghdr
		msg.Control = &cmsgBuf[0]
		msg.Controllen = uint64(len(cmsgBuf))

		// Recvmsg blocks or returns EAGAIN if no errors are pending in the queue
		_, _, _, _, err := unix.Recvmsg(int(fd), nil, cmsgBuf, unix.MSG_ERRQUEUE|unix.MSG_DONTWAIT)
		if err != nil {
			recvErr = err
			return
		}

		// Parse the control messages (cmsgs) to locate the zero-copy notification
		cmsgs, err := unix.ParseSocketControlMessage(cmsgBuf[:msg.Controllen])
		if err != nil {
			recvErr = err
			return
		}

		for _, cmsg := range cmsgs {
			if cmsg.Header.Level == unix.SOL_IP && cmsg.Header.Type == unix.IP_RECVERR {
				// Access the extended error structure
				serr := (*unix.SockExtendedErr)(unsafe.Pointer(&cmsg.Data[0]))
				if serr.Origin == unix.SO_EE_ORIGIN_ZEROCOPY {
					// Check for successful transmission confirmation
					if serr.Errno == uint32(syscall.SO_EE_CODE_ZEROCOPY_COPIED) {
						// Code 1 (COPIED) means the kernel had to fall back to a copy
						// This can happen if the payload was too small or memory was constrained
					}
					
					// info holds the sequence number range confirmed by the NIC
					comp.Lowest = serr.Info
					comp.Highest = serr.Info // Single notification can confirm a range of packets
					return
				}
			}
		}
		recvErr = errors.New("no zero-copy completion found in error queue")
	})

	if err != nil {
		return TxCompletion{}, err
	}
	return comp, recvErr
}