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 org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.*;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.CopyOnWriteArraySet;
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 Dong Liu
 */
public class RpcClientChannelGroup implements AutoCloseable {

    private volatile List channels;
    private final AtomicInteger roundIdx = new AtomicInteger(0);
    private ReadWriteLock lock = new ReentrantReadWriteLock();

    private static Logger logger = LoggerFactory.getLogger(RpcClientChannelGroup.class);

    public RpcClientChannelGroup(RpcClientChannel... clientChannels) {
        this.channels = Collections.unmodifiableList(Arrays.asList(clientChannels));
    }

    public RpcClientChannelGroup(List channels) {
        List channelList = new ArrayList<>(channels);
        this.channels = Collections.unmodifiableList(channelList);
    }

    /**
     * select one rpc channel. null if no one exists.
     * the channel returned may has been closed
     */
    public RpcClientChannel select() {
        List localChannels = channels;
        if (localChannels.size() == 0) {
            return null;
        }

        int idx = roundIdx.getAndIncrement();
        idx = idx % localChannels.size();
        if (idx < 0) {
            idx = idx + localChannels.size();
        }
        return localChannels.get(idx);
    }

    /**
     * add one channel
     *
     * @param rpcClientChannel
     * @return true if success, false if already have channel with same server info
     */
    public boolean add(RpcClientChannel rpcClientChannel) {
        ServerInfo serverInfo = rpcClientChannel.getServerInfo();
        lock.writeLock().lock();
        try {
            for (RpcClientChannel channel : channels) {
                ServerInfo si = channel.getServerInfo();
                if (sameServer(serverInfo, si)) {
                    return false;
                }
            }
            List newChannels = new ArrayList<>(channels.size() + 1);
            newChannels.addAll(channels);
            newChannels.add(rpcClientChannel);
            this.channels = Collections.unmodifiableList(newChannels);
        } finally {
            lock.writeLock().unlock();
        }
        return true;
    }

    /**
     * remove and close one channel
     *
     * @param serverInfo
     * @return true if removed; false if not exists
     */
    public boolean remove(ServerInfo serverInfo) {
        RpcClientChannel removed = null;
        lock.writeLock().lock();
        try {
            List newChannels = new ArrayList<>(channels.size());
            for (RpcClientChannel channel : channels) {
                ServerInfo si = channel.getServerInfo();
                if (sameServer(si, serverInfo)) {
                    removed = channel;
                } else {
                    newChannels.add(removed);
                }
            }
            if (removed != null) {
                this.channels = Collections.unmodifiableList(newChannels);
            }
        } finally {
            lock.writeLock().unlock();
        }
        if (removed != null) {
            try {
                removed.close();
            } catch (RuntimeException e) {
                logger.error("", e);
            }
            return true;
        } else {
            return false;
        }
    }

    @Override
    public void close() {
        List channelsCopy;
        lock.writeLock().lock();
        try {
            channelsCopy = new ArrayList<>(this.channels);
            this.channels = Collections.emptyList();
        } finally {
            lock.writeLock().unlock();
        }
        for (RpcClientChannel channel : channelsCopy) {
            channel.close();
        }
    }

    private boolean sameServer(ServerInfo s1, ServerInfo s2) {
        return s1.getHost().equals(s2.getHost()) && s1.getPort() == s2.getPort();
    }
}