Golang RPC连接池实现与优化技巧

type RPCClientPool struct {
    addr     string
    pool     chan *rpc.Client
    maxConn  int
    mu       sync.Mutex
}

func (p *RPCClientPool) Get() (*rpc.Client, error) {
    select {
    case client := <-p.pool:
        return client, nil
    default:
        p.mu.Lock()
        if len(p.pool) < p.maxConn {
            conn, err := net.Dial("tcp", p.addr)
            if err != nil {
                p.mu.Unlock()
                return nil, err
            }
            client := rpc.NewClient(conn)
            p.mu.Unlock()
            return client, nil
        }
        p.mu.Unlock()
        // 阻塞式获取（可设置超时）
        client, ok := <-p.pool
        if !ok {
            return nil, errors.New("pool closed")
        }
        return client, nil
    }
}

func (p *RPCClientPool) Put(client *rpc.Client) {
    select {
    case p.pool <- client:
        // 成功放回
    default:
        // 池满，关闭旧连接
        client.Close()
    }
}

func (p *RPCClientPool) Call(serviceMethod string, args interface{}, reply interface{}) error {
    client, err := p.Get()
    if err != nil {
        return err
    }
    defer func() { p.Put(client) }()
// 设置调用超时（可选）
done := make(chan error, 1)
go func() {
    done <- client.Call(serviceMethod, args, reply)
}()

select {
case err = <-done:
    return err
case <-time.After(3 * time.Second):
    return errors.New("call timeout")
}
}

func (p *RPCClientPool) isValid(client *rpc.Client) bool {
    var reply interface{}
    err := client.Call("Arith.Ping", nil, &reply)
    return err == nil
}

func (p *RPCClientPool) isAlive(client *rpc.Client) bool {
    return client.Active() > 0
}

func NewRPCClientPool(addr string, maxConn int) *RPCClientPool {
    return &RPCClientPool{
        addr:    addr,
        pool:    make(chan *rpc.Client, maxConn),
        maxConn: maxConn,
    }
}
func (p *RPCClientPool) Close() {
p.mu.Lock()
close(p.pool)
for client := range p.pool {
client.Close()
}
p.mu.Unlock()
}