GolangRPC负载均衡客户端实现示例

net/rpc

package main

import (
    "fmt"
    "log"
    "math/rand"
    "net/rpc"
    "sync"
    "time"
)

// ServiceDiscovery 模拟服务发现接口，用于获取服务实例列表
type ServiceDiscovery interface {
    GetServices() []string
    // 实际场景中，这里会有注册、注销、健康检查等机制
}

// StaticServiceDiscovery 静态服务发现，简单示例
type StaticServiceDiscovery struct {
    services []string
}

func NewStaticServiceDiscovery(addrs []string) *StaticServiceDiscovery {
    return &StaticServiceDiscovery{services: addrs}
}

func (s *StaticServiceDiscovery) GetServices() []string {
    return s.services
}

// Balancer 负载均衡器接口
type Balancer interface {
    Select(services []string) (string, error)
}

// RoundRobinBalancer 轮询负载均衡器
type RoundRobinBalancer struct {
    mu    sync.Mutex
    index int
}

func (r *RoundRobinBalancer) Select(services []string) (string, error) {
    if len(services) == 0 {
        return "", fmt.Errorf("no services available")
    }
    r.mu.Lock()
    defer r.mu.Unlock()

    selected := services[r.index%len(services)]
    r.index = (r.index + 1) % len(services)
    return selected, nil
}

// RandomBalancer 随机负载均衡器
type RandomBalancer struct{}

func (r *RandomBalancer) Select(services []string) (string, error) {
    if len(services) == 0 {
        return "", fmt.Errorf("no services available")
    }
    rand.Seed(time.Now().UnixNano()) // 实际应用中，rand.Seed只需初始化一次
    index := rand.Intn(len(services))
    return services[index], nil
}

// MyRPCClient 封装了RPC客户端和负载均衡逻辑
type MyRPCClient struct {
    sd       ServiceDiscovery
    balancer Balancer
    clients  map[string]*rpc.Client // 维护与各个服务端的连接
    mu       sync.RWMutex
}

func NewMyRPCClient(sd ServiceDiscovery, balancer Balancer) *MyRPCClient {
    return &MyRPCClient{
        sd:       sd,
        balancer: balancer,
        clients:  make(map[string]*rpc.Client),
    }
}

// getClient 获取或创建到指定地址的RPC连接
func (m *MyRPCClient) getClient(addr string) (*rpc.Client, error) {
    m.mu.RLock()
    client, ok := m.clients[addr]
    m.mu.RUnlock()

    if ok && client != nil {
        // 可以在这里加一个简单的健康检查，确保连接仍然有效
        // 比如尝试一个轻量级的ping方法，如果失败就关闭并重新连接
        return client, nil
    }

    m.mu.Lock()
    defer m.mu.Unlock()
    // 双重检查，避免重复创建
    client, ok = m.clients[addr]
    if ok && client != nil {
        return client, nil
    }

    // 尝试连接
    newClient, err := rpc.Dial("tcp", addr)
    if err != nil {
        log.Printf("Failed to dial RPC server %s: %v", addr, err)
        return nil, err
    }
    m.clients[addr] = newClient
    log.Printf("Successfully connected to RPC server %s", addr)
    return newClient, nil
}

// Call 是对外暴露的RPC调用方法
func (m *MyRPCClient) Call(serviceMethod string, args interface{}, reply interface{}) error {
    services := m.sd.GetServices()
    if len(services) == 0 {
        return fmt.Errorf("no RPC services registered or available")
    }

    // 尝试多次，处理瞬时连接失败
    const maxRetries = 3
    for i := 0; i < maxRetries; i++ {
        addr, err := m.balancer.Select(services)
        if err != nil {
            return fmt.Errorf("failed to select service: %v", err)
        }

        client, err := m.getClient(addr)
        if err != nil {
            log.Printf("Attempt %d: Could not get client for %s, trying another...", i+1, addr)
            // 如果连接失败，考虑将该地址暂时从可用列表中移除，或者等待服务发现更新
            // 在这个简化示例中，我们只是重试
            time.Sleep(100 * time.Millisecond) // 简单退避
            continue
        }

        err = client.Call(serviceMethod, args, reply)
        if err != nil {
            log.Printf("Attempt %d: RPC call to %s failed: %v, trying another...", i+1, addr, err)
            // RPC调用失败，可能是服务端问题，关闭当前连接并尝试重新获取
            m.mu.Lock()
            if oldClient, ok := m.clients[addr]; ok && oldClient == client { // 确保是同一个client
                oldClient.Close()
                delete(m.clients, addr)
                log.Printf("Closed faulty connection to %s", addr)
            }
            m.mu.Unlock()
            time.Sleep(100 * time.Millisecond) // 简单退避
            continue
        }
        return nil // 调用成功
    }

    return fmt.Errorf("all RPC call attempts failed after %d retries", maxRetries)
}

// Close 关闭所有维护的RPC连接
func (m *MyRPCClient) Close() {
    m.mu.Lock()
    defer m.mu.Unlock()
    for addr, client := range m.clients {
        if client != nil {
            client.Close()
            log.Printf("Closed RPC connection to %s", addr)
        }
    }
    m.clients = make(map[string]*rpc.Client) // 清空
}

// 假设的服务端代码 (仅为测试客户端)
type Args struct {
    A, B int
}

type Reply struct {
    C int
}

type Math struct{}

func (m *Math) Add(args *Args, reply *Reply) error {
    reply.C = args.A + args.B
    log.Printf("Server received Add(%d, %d), returning %d", args.A, args.B, reply.C)
    return nil
}

func startServer(addr string) {
    math := new(Math)
    rpc.Register(math)
    listener, err := net.Listen("tcp", addr)
    if err != nil {
        log.Fatalf("Listen error: %v", err)
    }
    log.Printf("RPC server listening on %s", addr)
    go func() {
        for {
            conn, err := listener.Accept()
            if err != nil {
                log.Printf("Accept error: %v", err)
                continue
            }
            go rpc.ServeConn(conn)
        }
    }()
}

func main() {
    // 启动几个RPC服务端实例
    serverAddrs := []string{":1234", ":1235", ":1236"}
    for _, addr := range serverAddrs {
        startServer(addr)
    }
    time.Sleep(time.Second) // 等待服务器启动

    // 初始化服务发现和负载均衡器
    sd := NewStaticServiceDiscovery(serverAddrs)
    // balancer := &RoundRobinBalancer{}
    balancer := &RandomBalancer{} // 切换不同的负载均衡策略

    client := NewMyRPCClient(sd, balancer)
    defer client.Close()

    // 模拟多次RPC调用
    for i := 0; i < 10; i++ {
        args := &Args{A: i, B: i * 2}
        reply := &Reply{}
        err := client.Call("Math.Add", args, reply)
        if err != nil {
            log.Printf("RPC call failed: %v", err)
        } else {
            log.Printf("RPC call successful: %d + %d = %d", args.A, args.B, reply.C)
        }
        time.Sleep(100 * time.Millisecond)
    }
}

MyRPCClient

ServiceDiscovery

Balancer