java 简单 SocketPool

Posted 2022-06-17 googler

package org.rx.socks;

import lombok.extern.slf4j.Slf4j;
import org.rx.common.LogWriter;
import org.rx.common.NQuery;
import org.rx.beans.DateTime;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.Socket;

import java.util.Map;
import java.util.Timer;
import java.util.TimerTask;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedDeque;

import static org.rx.common.Contract.require;

@Slf4j
public final class SocketPool extends Traceable implements AutoCloseable 
    public static final class PooledSocket implements AutoCloseable 
        private final SocketPool owner;
        private DateTime lastActive;
        public final Socket socket;

        public boolean isConnected() 
            return !owner.isClosed() && !socket.isClosed() && socket.isConnected();
        

        public DateTime getLastActive() 
            return lastActive;
        

        public void setLastActive(DateTime lastActive) 
            this.lastActive = lastActive;
        

        private PooledSocket(SocketPool owner, Socket socket) 
            this.owner = owner;
            this.socket = socket;
            lastActive = DateTime.utcNow();
        

        @Override
        public void close() 
            owner.returnSocket(this);
        
    

    public static final SocketPool Pool = new SocketPool();
    private static final int DefaultConnectTimeout = 30000;
    private static final int DefaultMaxIdleMillis = 120000;
    private static final int DefaultMaxSocketsCount = 64;
    private final ConcurrentHashMap<InetSocketAddress, ConcurrentLinkedDeque<PooledSocket>> pool;
    private volatile int connectTimeout;
    private volatile int maxIdleMillis;
    private volatile int maxSocketsCount;
    private final Timer timer;
    private volatile boolean isTimerRun;

    public int getConnectTimeout() 
        return connectTimeout;
    

    public void setConnectTimeout(int connectTimeout) 
        this.connectTimeout = connectTimeout;
    

    public int getMaxIdleMillis() 
        return maxIdleMillis;
    

    public void setMaxIdleMillis(int maxIdleMillis) 
        if (maxIdleMillis <= 0) 
            maxIdleMillis = DefaultMaxIdleMillis;
        
        this.maxIdleMillis = maxIdleMillis;
    

    public int getMaxSocketsCount() 
        return maxSocketsCount;
    

    public void setMaxSocketsCount(int maxSocketsCount) 
        if (maxSocketsCount < 0) 
            maxSocketsCount = 0;
        
        this.maxSocketsCount = maxSocketsCount;
    

    private SocketPool() 
        pool = new ConcurrentHashMap<>();
        connectTimeout = DefaultConnectTimeout;
        maxIdleMillis = DefaultMaxIdleMillis;
        maxSocketsCount = DefaultMaxSocketsCount;
        String n = "SocketPool";
        timer = new Timer(n, true);
        LogWriter tracer = new LogWriter();
        tracer.setPrefix(n + " ");
        tracer.info("started..");
        setTracer(tracer);
    

    @Override
    protected void freeObjects() 
        clear();
    

    private void runTimer() 
        if (isTimerRun) 
            return;
        
        synchronized (timer) 
            if (isTimerRun) 
                return;
            

            long period = 90000;
            timer.schedule(new TimerTask() 
                @Override
                public void run() 
                    clearIdleSockets();
                
            , period, period);
            isTimerRun = true;
        
        getTracer().info("runTimer..");
    

    private void clearIdleSockets() 
        for (Map.Entry<InetSocketAddress, ConcurrentLinkedDeque<PooledSocket>> entry : NQuery.of(pool.entrySet())) 
            ConcurrentLinkedDeque<PooledSocket> sockets = entry.getValue();
            if (sockets == null) 
                continue;
            

            for (PooledSocket socket : NQuery.of(sockets)) 
                if (!socket.isConnected()
                        || DateTime.utcNow().subtract(socket.getLastActive()).getTotalMilliseconds() >= maxIdleMillis) 
                    sockets.remove(socket);
                    getTracer().info("clear idle socket[local=%s, remote=%s]..",
                            Sockets.getId(socket.socket, false), Sockets.getId(socket.socket, true));
                
            
            if (sockets.isEmpty()) 
                pool.remove(entry.getKey());
            
        
        if (pool.size() == 0) 
            stopTimer();
        
    

    private void stopTimer() 
        synchronized (timer) 
            timer.cancel();
            timer.purge();
            isTimerRun = false;
        
        getTracer().info("stopTimer..");
    

    private ConcurrentLinkedDeque<PooledSocket> getSockets(InetSocketAddress remoteAddr) 
        ConcurrentLinkedDeque<PooledSocket> sockets = pool.get(remoteAddr);
        if (sockets == null) 
            pool.put(remoteAddr, sockets = new ConcurrentLinkedDeque<>());
            runTimer();
        
        return sockets;
    

    public PooledSocket borrowSocket(InetSocketAddress remoteAddr) 
        checkNotClosed();
        require(remoteAddr);

        boolean isExisted = true;
        ConcurrentLinkedDeque<PooledSocket> sockets = getSockets(remoteAddr);
        PooledSocket pooledSocket;
        if ((pooledSocket = sockets.pollFirst()) == null) 
            Socket sock = new Socket();
            try 
                sock.connect(remoteAddr, connectTimeout);
             catch (IOException ex) 
                throw new SocketException(remoteAddr, ex);
            
            pooledSocket = new PooledSocket(this, sock);
            isExisted = false;
        
        if (!pooledSocket.isConnected()) 
            if (isExisted) 
                sockets.remove(pooledSocket);
            
            return borrowSocket(remoteAddr);
        
        Socket sock = pooledSocket.socket;
        getTracer().info("borrow %s socket[local=%s, remote=%s]..", isExisted ? "existed" : "new",
                Sockets.getId(sock, false), Sockets.getId(sock, true));
        return pooledSocket;
    

    public void returnSocket(PooledSocket pooledSocket) 
        checkNotClosed();
        require(pooledSocket);

        String action = "return";
        try 
            if (!pooledSocket.isConnected()) 
                action = "discard closed";
                return;
            
            pooledSocket.setLastActive(DateTime.utcNow());
            ConcurrentLinkedDeque<PooledSocket> sockets = getSockets(
                    (InetSocketAddress) pooledSocket.socket.getRemoteSocketAddress());
            if (sockets.size() >= maxSocketsCount || sockets.contains(pooledSocket)) 
                action = "discard contains";
                return;
            

            sockets.addFirst(pooledSocket);
         finally 
            Socket sock = pooledSocket.socket;
            getTracer().info("%s socket[local=%s, remote=%s]..", action, Sockets.getId(sock, false),
                    Sockets.getId(sock, true));
        
    

    public void clear() 
        checkNotClosed();

        for (Socket socket : NQuery.of(pool.values()).selectMany(p -> p).select(p -> p.socket)) 
            try 
                getTracer().info("clear socket[local=%s, remote=%s]..", Sockets.getId(socket, false),
                        Sockets.getId(socket, true));
                Sockets.close(socket);
             catch (Exception ex) 
                log.error("SocketPool clear", ex);
            
        
        pool.clear();