net.dongliu.prettypb.rpc.client.RpcClientChannelGroup Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of prettypb-rpc Show documentation

proto rpc libs, compatible with proto-rpc-pro

There is a newer version: 0.3.5

package net.dongliu.prettypb.rpc.client;

import net.dongliu.prettypb.rpc.protocol.ServerInfo;

import java.util.Collection;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.ReadWriteLock;
import java.util.concurrent.locks.ReentrantReadWriteLock;

/**
 * a group of rpc channel for one rpc client
 *
 * @author dongliu
 */
public class RpcClientChannelGroup implements AutoCloseable {

    private volatile Map channelMap = new HashMap<>();
    private final AtomicInteger roundIdx = new AtomicInteger(0);
    private ReadWriteLock lock = new ReentrantReadWriteLock();

    /**
     * select one rpc channel. null if no one exists.
     * the channel returned may has been closed
     */
    public RpcClientChannel selectChannel() {
        Collection channels;
        lock.readLock().lock();
        try {
            channels = channelMap.values();
        } finally {
            lock.readLock().unlock();
        }
        if (channels.isEmpty()) {
            return null;
        }
        RpcClientChannel[] channelArray = channels.toArray(new RpcClientChannel[channels.size()]);
        int idx = roundIdx.getAndIncrement();
        return channelArray[idx % channelArray.length];
    }

    @Override
    public synchronized void close() {
        Collection channels;
        lock.writeLock().lock();
        try {
            channels = channelMap.values();
            channelMap.clear();
        } finally {
            lock.writeLock().unlock();
        }
        for (RpcClientChannel channel : channels) {
            channel.close();
        }
    }

    public void add(ServerInfo serverInfo, RpcClientChannel rpcClientChannel) {
        lock.writeLock().lock();
        try {
            this.channelMap.put(getServerKey(serverInfo), rpcClientChannel);
        } finally {
            lock.writeLock().unlock();
        }
    }

    public synchronized void remove(ServerInfo serverInfo) {
        RpcClientChannel channel;
        lock.writeLock().lock();
        try {
            channel = channelMap.remove(getServerKey(serverInfo));
        } finally {
            lock.writeLock().unlock();
        }
        if (channel != null) {
            channel.close();
        }
    }

    private String getServerKey(ServerInfo serverInfo) {
        return serverInfo.getHost() + ":" + serverInfo.getPort();
    }
}