Golang的RPC如何支持压缩传输 对比Snappy与Gzip的性能影响

net.Conn

rpc.ClientCodec

rpc.ServerCodec

net.Conn

io.ReadWriteCloser

rpc.Client

rpc.Server

net.Conn

compress/gzip

Conn

net.Conn

package main

import (
    "compress/gzip"
    "io"
    "log"
    "net"
    "net/rpc"
    "time"
)

// GzipConn 实现了 net.Conn 接口，并提供 Gzip 压缩/解压缩功能
type GzipConn struct {
    conn       net.Conn
    reader     *gzip.Reader
    writer     *gzip.Writer
    readBuffer []byte // 用于 Read 方法的临时缓冲区
}

// NewGzipConn 创建一个新的 GzipConn
func NewGzipConn(conn net.Conn) (*GzipConn, error) {
    // Gzip reader 需要在第一次读取时被 Reset，或者在创建时就传入底层的 io.Reader
    // 这里我们选择在 Read 方法中动态处理或预先创建
    reader, err := gzip.NewReader(conn) // Gzip reader 会从 conn 读取压缩数据
    if err != nil {
        return nil, err
    }
    return &GzipConn{
        conn:       conn,
        reader:     reader,
        writer:     gzip.NewWriter(conn), // Gzip writer 会向 conn 写入压缩数据
        readBuffer: make([]byte, 4096), // 示例缓冲区大小
    }, nil
}

func (c *GzipConn) Read(b []byte) (n int, err error) {
    // 注意：gzip.Reader 在底层 io.Reader EOF 后可能不会立即返回 EOF，
    // 而是等待读取完所有解压数据。这里简化处理，实际生产环境需要更健壮的EOF和错误处理
    return c.reader.Read(b)
}

func (c *GzipConn) Write(b []byte) (n int, err error) {
    n, err = c.writer.Write(b)
    if err != nil {
        return n, err
    }
    // 刷新 writer 确保数据被写入底层 conn
    return n, c.writer.Flush()
}

func (c *GzipConn) Close() error {
    // 关闭 writer 和 reader，然后关闭底层 conn
    err1 := c.writer.Close()
    err2 := c.reader.Close()
    err3 := c.conn.Close()
    if err1 != nil {
        return err1
    }
    if err2 != nil {
        return err2
    }
    return err3
}

// 以下方法只是简单代理到底层 conn，因为它们与数据流无关
func (c *GzipConn) LocalAddr() net.Addr {
    return c.conn.LocalAddr()
}

func (c *GzipConn) RemoteAddr() net.Addr {
    return c.conn.RemoteAddr()
}

func (c *GzipConn) SetDeadline(t time.Time) error {
    return c.conn.SetDeadline(t)
}

func (c *GzipConn) SetReadDeadline(t time.Time) error {
    return c.conn.SetReadDeadline(t)
}

func (c *GzipConn) SetWriteDeadline(t time.Time) error {
    return c.conn.SetWriteDeadline(t)
}

// RPC Server 端使用 GzipConn
func serveGzipRPC(l net.Listener) {
    for {
        conn, err := l.Accept()
        if err != nil {
            log.Printf("Server accept error: %v", err)
            return
        }
        go func() {
            gzipConn, err := NewGzipConn(conn)
            if err != nil {
                log.Printf("Failed to create GzipConn: %v", err)
                conn.Close()
                return
            }
            rpc.ServeConn(gzipConn)
            gzipConn.Close() // 确保连接关闭
        }()
    }
}

// RPC Client 端使用 GzipConn
func dialGzipRPC(network, address string) (*rpc.Client, error) {
    conn, err := net.Dial(network, address)
    if err != nil {
        return nil, err
    }
    gzipConn, err := NewGzipConn(conn)
    if err != nil {
        conn.Close()
        return nil, err
    }
    return rpc.NewClient(gzipConn), nil
}

github.com/golang/snappy/snappy

package main

import (
    "io"
    "log"
    "net"
    "net/rpc"
    "time"

    "github.com/golang/snappy" // 引入 Snappy 库
)

// SnappyConn 实现了 net.Conn 接口，并提供 Snappy 压缩/解压缩功能
type SnappyConn struct {
    conn   net.Conn
    reader *snappy.Reader
    writer *snappy.Writer
}

// NewSnappyConn 创建一个新的 SnappyConn
func NewSnappyConn(conn net.Conn) *SnappyConn {
    return &SnappyConn{
        conn:   conn,
        reader: snappy.NewReader(conn), // Snappy reader 从 conn 读取压缩数据
        writer: snappy.NewBufferedWriter(conn), // Snappy writer 向 conn 写入压缩数据
    }
}

func (c *SnappyConn) Read(b []byte) (n int, err error) {
    return c.reader.Read(b)
}

func (c *SnappyConn) Write(b []byte) (n int, err error) {
    n, err = c.writer.Write(b)
    if err != nil {
        return n, err
    }
    // Snappy writer 同样需要 Flush
    return n, c.writer.Flush()
}

func (c *SnappyConn) Close() error {
    // 关闭 writer 和 reader，然后关闭底层 conn
    err1 := c.writer.Close()
    err2 := c.conn.Close() // snappy.Reader 没有 Close 方法，直接关闭底层 conn
    if err1 != nil {
        return err1
    }
    return err2
}

// 以下方法只是简单代理到底层 conn
func (c *SnappyConn) LocalAddr() net.Addr {
    return c.conn.LocalAddr()
}

func (c *SnappyConn) RemoteAddr() net.Addr {
    return c.conn.RemoteAddr()
}

func (c *SnappyConn) SetDeadline(t time.Time) error {
    return c.conn.SetDeadline(t)
}

func (c *SnappyConn) SetReadDeadline(t time.Time) error {
    return c.conn.SetReadDeadline(t)
}

func (c *SnappyConn) SetWriteDeadline(t time.Time) error {
    return c.conn.SetWriteDeadline(t)
}

// RPC Server 端使用 SnappyConn
func serveSnappyRPC(l net.Listener) {
    for {
        conn, err := l.Accept()
        if err != nil {
            log.Printf("Server accept error: %v", err)
            return
        }
        go func() {
            snappyConn := NewSnappyConn(conn)
            rpc.ServeConn(snappyConn)
            snappyConn.Close()
        }()
    }
}

// RPC Client 端使用 SnappyConn
func dialSnappyRPC(network, address string) (*rpc.Client, error) {
    conn, err := net.Dial(network, address)
    if err != nil {
        return nil, err
    }
    snappyConn := NewSnappyConn(conn)
    return rpc.NewClient(snappyConn), nil
}

net.Conn

Read

Write