package org.rx.socks;
import lombok.extern.slf4j.Slf4j;
import org.rx.core.LogWriter;
import org.rx.core.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.core.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();
}
}
package org.rx.socks;
import org.rx.core.Disposable;
import org.rx.core.LogWriter;
import static org.rx.core.Contract.isNull;
public abstract class Traceable extends Disposable {
private LogWriter tracer;
public LogWriter getTracer() {
return tracer;
}
public synchronized void setTracer(LogWriter tracer) {
this.tracer = isNull(tracer, new LogWriter());
}
}
package org.rx.socks;
import lombok.extern.slf4j.Slf4j;
import org.rx.beans.$;
import org.rx.beans.Tuple;
import org.rx.util.BufferSegment;
import org.rx.util.BytesSegment;
import org.rx.core.*;
import org.rx.core.AsyncTask;
import org.rx.io.MemoryStream;
import java.io.IOException;
import java.net.*;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import static org.rx.beans.$.$;
import static org.rx.core.Contract.isNull;
import static org.rx.core.Contract.require;
@Slf4j
public class DirectSocket extends Traceable implements AutoCloseable {
@FunctionalInterface
public interface SocketSupplier {
Tuple<AutoCloseable, Socket> get(MemoryStream pack);
}
private static class ClientItem {
private final DirectSocket owner;
private final BufferSegment segment;
public final NetworkStream stream;
public final AutoCloseable toSock;
public final NetworkStream toStream;
public ClientItem(Socket client, DirectSocket owner) {
this.owner = owner;
segment = new BufferSegment(Contract.config.getDefaultBufferSize(), 2);
try {
stream = new NetworkStream(client, segment.alloc());
if (owner.directAddress != null) {
SocketPool.PooledSocket pooledSocket = App.retry(owner.connectRetryCount,
p -> SocketPool.Pool.borrowSocket(p.directAddress), owner);
toSock = pooledSocket;
toStream = new NetworkStream(pooledSocket.socket, segment.alloc(), false);
return;
}
if (owner.directSupplier != null) {
MemoryStream firstPack = new MemoryStream(32, true);
BytesSegment buffer = stream.getSegment();
int read;
while ((read = stream.readSegment()) > 0) {
System.out.println("----:" + Bytes.toString(buffer.array, buffer.offset, read));
firstPack.write(buffer.array, buffer.offset, read);
Tuple<AutoCloseable, Socket> toSocks;
if ((toSocks = owner.directSupplier.get(firstPack)) != null) {
toSock = toSocks.left;
firstPack.writeTo(toStream = new NetworkStream(toSocks.right, segment.alloc(), false));
return;
}
}
log.info("DirectSocket ClientState directSupplier read: {}
content: {}", read,
Bytes.toString(firstPack.toArray(), 0, firstPack.getLength()));
}
} catch (IOException ex) {
throw new SocketException((InetSocketAddress) client.getLocalSocketAddress(), ex);
}
throw new SocketException((InetSocketAddress) client.getLocalSocketAddress(),
"DirectSocket directSupplier error");
}
public void closeSocket() {
owner.getTracer().info("client close socket[%s->%s]..", Sockets.getId(stream.getSocket(), false),
Sockets.getId(stream.getSocket(), true));
owner.clients.remove(this);
stream.close();
}
public void closeToSocket(boolean pooling) {
owner.getTracer().info("client %s socket[%s->%s]..", pooling ? "pooling" : "close",
Sockets.getId(toStream.getSocket(), false), Sockets.getId(toStream.getSocket(), true));
if (pooling) {
try {
toSock.close();
} catch (Exception ex) {
ex.printStackTrace();
}
} else {
Sockets.close(toStream.getSocket());
}
}
}
public static final SocketSupplier HttpSupplier = pack -> {
String line = Bytes.readLine(pack.getBuffer());
if (line == null) {
return null;
}
InetSocketAddress authority;
try {
authority = Sockets.parseEndpoint(
new URL(line.split(" ")[1])
.getAuthority());
} catch (MalformedURLException ex) {
throw SystemException.wrap(ex);
}
SocketPool.PooledSocket pooledSocket = App.retry(2,
p -> SocketPool.Pool.borrowSocket(p),
authority);
return Tuple.of(pooledSocket, pooledSocket.socket);
};
private static final int DefaultBacklog = 128;
private static final int DefaultConnectRetryCount = 4;
private final ServerSocket server;
private final List<ClientItem> clients;
private volatile int connectRetryCount;
private InetSocketAddress directAddress;
private SocketSupplier directSupplier;
@Override
public boolean isClosed() {
return !(!super.isClosed() && !server.isClosed());
}
public InetSocketAddress getLocalAddress() {
return (InetSocketAddress) server.getLocalSocketAddress();
}
public NQuery<Tuple<Socket, Socket>> getClients() {
return NQuery.of(clients).select(p -> Tuple.of(p.stream.getSocket(), p.toStream.getSocket()));
}
public int getConnectRetryCount() {
return connectRetryCount;
}
public void setConnectRetryCount(int connectRetryCount) {
if (connectRetryCount <= 0) {
connectRetryCount = 1;
}
this.connectRetryCount = connectRetryCount;
}
public DirectSocket(int listenPort, InetSocketAddress directAddr) {
this(new InetSocketAddress(Sockets.AnyAddress, listenPort), directAddr, null);
}
public DirectSocket(InetSocketAddress listenAddr, InetSocketAddress directAddr, SocketSupplier directSupplier) {
require(listenAddr);
require(this, directAddr != null || directSupplier != null);
try {
server = new ServerSocket();
server.setReuseAddress(true);
server.bind(listenAddr, DefaultBacklog);
} catch (IOException ex) {
throw new SocketException(listenAddr, ex);
}
directAddress = directAddr;
this.directSupplier = directSupplier;
clients = Collections.synchronizedList(new ArrayList<>());
connectRetryCount = DefaultConnectRetryCount;
String taskName = String.format("DirectSocket[%s->%s]", listenAddr, isNull(directAddress, "autoAddress"));
LogWriter tracer = new LogWriter();
tracer.setPrefix(taskName + " ");
setTracer(tracer);
AsyncTask.TaskFactory.run(() -> {
getTracer().info("start..");
while (!isClosed()) {
try {
ClientItem client = new ClientItem(server.accept(), this);
clients.add(client);
onReceive(client, taskName);
} catch (IOException ex) {
log.error(taskName, ex);
}
}
close();
}, taskName);
}
@Override
protected void freeObjects() {
try {
for (ClientItem client : NQuery.of(clients)) {
client.closeSocket();
}
clients.clear();
server.close();
} catch (IOException ex) {
log.error("DirectSocket close", ex);
}
getTracer().info("stop..");
}
private void onReceive(ClientItem client, String taskName) {
AsyncTask.TaskFactory.run(() -> {
try {
int recv = client.stream.directTo(client.toStream, (p1, p2) -> {
getTracer().info("sent %s bytes from %s to %s..", p2,
Sockets.getId(client.stream.getSocket(), true),
Sockets.getId(client.toStream.getSocket(), false));
return true;
});
getTracer().info("socket[%s->%s] closing with %s", Sockets.getId(client.stream.getSocket(), false),
Sockets.getId(client.stream.getSocket(), true), recv);
} catch (SystemException ex) {
$<java.net.SocketException> out = $();
if (ex.tryGet(out, java.net.SocketException.class)) {
if (out.v.getMessage().contains("Socket closed")) {
//ignore
log.debug("DirectTo ignore socket closed");
return;
}
}
throw ex;
} finally {
client.closeSocket();
}
}, String.format("%s[networkStream]", taskName));
AsyncTask.TaskFactory.run(() -> {
int recv = NetworkStream.StreamEOF;
try {
recv = client.toStream.directTo(client.stream, (p1, p2) -> {
getTracer().info("recv %s bytes from %s to %s..", p2,
Sockets.getId(client.toStream.getSocket(), false),
Sockets.getId(client.stream.getSocket(), true));
return true;
});
getTracer().info("socket[%s->%s] closing with %s", Sockets.getId(client.toStream.getSocket(), false),
Sockets.getId(client.toStream.getSocket(), true), recv);
} catch (SystemException ex) {
$<java.net.SocketException> out = $();
if (ex.tryGet(out, java.net.SocketException.class)) {
if (out.v.getMessage().contains("Socket closed")) {
//ignore
log.debug("DirectTo ignore socket closed");
return;
}
}
throw ex;
} finally {
client.closeToSocket(recv == NetworkStream.CannotWrite);
}
}, String.format("%s[toNetworkStream]", taskName));
}
}
package org.rx.socks;
import lombok.extern.slf4j.Slf4j;
import org.rx.util.BytesSegment;
import org.rx.io.IOStream;
import java.io.IOException;
import java.net.Socket;
import static org.rx.core.Contract.require;
import static org.rx.socks.Sockets.shutdown;
@Slf4j
public final class NetworkStream extends IOStream {
@FunctionalInterface
public interface DirectPredicate {
boolean test(BytesSegment buffer, int count);
}
public static final int SocketEOF = 0;
public static final int StreamEOF = -1;
public static final int CannotWrite = -2;
private final boolean ownsSocket;
private final Socket socket;
private final BytesSegment segment;
public boolean isConnected() {
return !isClosed() && !socket.isClosed() && socket.isConnected();
}
@Override
public boolean canRead() {
return super.canRead() && checkSocket(socket, false);
}
@Override
public boolean canWrite() {
return super.canWrite() && checkSocket(socket, true);
}
private static boolean checkSocket(Socket sock, boolean isWrite) {
return !sock.isClosed() && sock.isConnected() && !(isWrite ? sock.isOutputShutdown() : sock.isInputShutdown());
}
public Socket getSocket() {
return socket;
}
public BytesSegment getSegment() {
return segment;
}
public NetworkStream(Socket socket, BytesSegment segment) throws IOException {
this(socket, segment, true);
}
public NetworkStream(Socket socket, BytesSegment segment, boolean ownsSocket) throws IOException {
super(socket.getInputStream(), socket.getOutputStream());
this.ownsSocket = ownsSocket;
this.socket = socket;
this.segment = segment;
}
@Override
protected void freeObjects() {
try {
log.info("NetworkStream freeObjects ownsSocket={} socket[{}][closed={}]", ownsSocket,
Sockets.getId(socket, false), socket.isClosed());
if (ownsSocket) {
//super.freeObjects(); Ignore this!!
Sockets.close(socket, 1);
}
} finally {
segment.close();
}
}
int readSegment() {
return read(segment.array, segment.offset, segment.count);
}
void writeSegment(int count) {
write(segment.array, segment.offset, count);
}
public int directTo(NetworkStream to, DirectPredicate onEach) {
checkNotClosed();
require(to);
int recv = StreamEOF;
while (canRead() && (recv = read(segment.array, segment.offset, segment.count)) >= -1) {
if (recv <= 0) {
if (ownsSocket) {
log.debug("DirectTo read {} flag and shutdown send", recv);
shutdown(socket, 1);
}
break;
}
if (!to.canWrite()) {
log.debug("DirectTo read {} bytes and can't write", recv);
recv = CannotWrite;
break;
}
to.write(segment.array, segment.offset, recv);
if (onEach != null && !onEach.test(segment, recv)) {
recv = StreamEOF;
break;
}
}
if (to.canWrite()) {
to.flush();
}
return recv;
}
}
package org.rx.socks;
import org.rx.core.SystemException;
import java.net.InetSocketAddress;
public class SocketException extends SystemException {
private InetSocketAddress localAddress;
public InetSocketAddress getLocalAddress() {
return localAddress;
}
public SocketException(InetSocketAddress localAddress, Exception ex) {
super(ex);
this.localAddress = localAddress;
}
public SocketException(InetSocketAddress localAddress, String msg) {
super(msg);
this.localAddress = localAddress;
}
}
package org.rx.socks;
import java.io.IOException;
import java.net.*;
import io.netty.bootstrap.Bootstrap;
import io.netty.buffer.Unpooled;
import io.netty.channel.Channel;
import io.netty.channel.ChannelFutureListener;
import io.netty.channel.EventLoopGroup;
import io.netty.channel.ServerChannel;
import io.netty.channel.epoll.Epoll;
import io.netty.channel.epoll.EpollEventLoopGroup;
import io.netty.channel.epoll.EpollServerSocketChannel;
import io.netty.channel.epoll.EpollSocketChannel;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.nio.NioServerSocketChannel;
import io.netty.channel.socket.nio.NioSocketChannel;
import org.rx.core.Strings;
import org.rx.core.SystemException;
import org.rx.core.WeakCache;
import org.springframework.util.CollectionUtils;
import java.util.List;
import java.util.Properties;
import java.util.function.Function;
import static org.rx.core.Contract.require;
import static org.rx.core.Contract.values;
public final class Sockets {
public static final InetAddress LocalAddress, AnyAddress;
static {
LocalAddress = InetAddress.getLoopbackAddress();
try {
AnyAddress = InetAddress.getByName("0.0.0.0");
} catch (Exception ex) {
throw SystemException.wrap(ex);
}
}
public InetAddress[] getAddresses(String host) {
return (InetAddress[]) WeakCache.getOrStore("Sockets.getAddresses", values(host), p -> {
try {
return InetAddress.getAllByName(host);
} catch (UnknownHostException ex) {
throw SystemException.wrap(ex);
}
});
}
public static InetSocketAddress getAnyEndpoint(int port) {
return new InetSocketAddress(AnyAddress, port);
}
public static InetSocketAddress parseEndpoint(String endpoint) {
require(endpoint);
String[] arr = Strings.split(endpoint, ":", 2);
return new InetSocketAddress(arr[0], Integer.parseInt(arr[1]));
}
public static void writeAndFlush(Channel channel, Object... packs) {
require(channel);
channel.eventLoop().execute(() -> {
for (Object pack : packs) {
channel.write(pack);
}
channel.flush();
});
}
public static EventLoopGroup bossEventLoop() {
return eventLoopGroup(1);
}
public static EventLoopGroup workEventLoop() {
return eventLoopGroup(0);
}
public static EventLoopGroup eventLoopGroup(int threadAmount) {
return Epoll.isAvailable() ? new EpollEventLoopGroup(threadAmount) : new NioEventLoopGroup(threadAmount); //NioEventLoopGroup(0, TaskFactory.getExecutor());
}
public static Bootstrap bootstrap() {
return bootstrap(getChannelClass());
}
public static Bootstrap bootstrap(Class<? extends Channel> channelClass) {
require(channelClass);
return new Bootstrap().group(channelClass.getName().startsWith("Epoll") ? new EpollEventLoopGroup() : new NioEventLoopGroup()).channel(channelClass);
}
public static Bootstrap bootstrap(Channel channel) {
require(channel);
return new Bootstrap().group(channel.eventLoop()).channel(channel.getClass());
}
public static Class<? extends ServerChannel> getServerChannelClass() {
return Epoll.isAvailable() ? EpollServerSocketChannel.class : NioServerSocketChannel.class;
}
public static Class<? extends Channel> getChannelClass() {
return Epoll.isAvailable() ? EpollSocketChannel.class : NioSocketChannel.class;
}
public static void closeOnFlushed(Channel channel) {
require(channel);
if (!channel.isActive()) {
return;
}
channel.writeAndFlush(Unpooled.EMPTY_BUFFER).addListener(ChannelFutureListener.CLOSE);
}
public static void close(Socket socket) {
close(socket, 1 | 2);
}
public static void close(Socket socket, int flags) {
require(socket);
if (!socket.isClosed()) {
shutdown(socket, flags);
try {
socket.setSoLinger(true, 2);
socket.close();
} catch (IOException ex) {
throw SystemException.wrap(ex);
}
}
}
/**
* @param socket
* @param flags Send=1, Receive=2
*/
public static void shutdown(Socket socket, int flags) {
require(socket);
if (!socket.isClosed() && socket.isConnected()) {
try {
if ((flags & 1) == 1 && !socket.isOutputShutdown()) {
socket.shutdownOutput();
}
if ((flags & 2) == 2 && !socket.isInputShutdown()) {
socket.shutdownInput();
}
} catch (IOException ex) {
throw SystemException.wrap(ex);
}
}
}
public static String getId(Socket sock, boolean isRemote) {
require(sock);
InetSocketAddress addr = (InetSocketAddress) (isRemote ? sock.getRemoteSocketAddress()
: sock.getLocalSocketAddress());
return addr.getHostString() + ":" + addr.getPort();
}
public static <T> T httpProxyInvoke(String proxyAddr, Function<String, T> func) {
setHttpProxy(proxyAddr);
try {
return func.apply(proxyAddr);
} finally {
clearHttpProxy();
}
}
public static void setHttpProxy(String proxyAddr) {
setHttpProxy(proxyAddr, null, null, null);
}
public static void setHttpProxy(String proxyAddr, List<String> nonProxyHosts, String userName, String password) {
InetSocketAddress ipe = parseEndpoint(proxyAddr);
Properties prop = System.getProperties();
prop.setProperty("http.proxyHost", ipe.getAddress().getHostAddress());
prop.setProperty("http.proxyPort", String.valueOf(ipe.getPort()));
prop.setProperty("https.proxyHost", ipe.getAddress().getHostAddress());
prop.setProperty("https.proxyPort", String.valueOf(ipe.getPort()));
if (!CollectionUtils.isEmpty(nonProxyHosts)) {
//如"localhost|192.168.0.*"
prop.setProperty("http.nonProxyHosts", String.join("|", nonProxyHosts));
}
if (userName != null && password != null) {
Authenticator.setDefault(new UserAuthenticator(userName, password));
}
}
public static void clearHttpProxy() {
System.clearProperty("http.proxyHost");
System.clearProperty("http.proxyPort");
System.clearProperty("https.proxyHost");
System.clearProperty("https.proxyPort");
System.clearProperty("http.nonProxyHosts");
}
static class UserAuthenticator extends Authenticator {
private String userName;
private String password;
public UserAuthenticator(String userName, String password) {
this.userName = userName;
this.password = password;
}
protected PasswordAuthentication getPasswordAuthentication() {
return new PasswordAuthentication(userName, password.toCharArray());
}
}
}
