feat: 支持自定义线程池

Author: meethigher

src/main/java/top/meethigher/TCPReverseProxy.java

@@ -9,9 +9,9 @@
 import java.net.InetAddress;
 import java.net.ServerSocket;
 import java.net.Socket;
-import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.concurrent.ForkJoinPool;
 
 public class TCPReverseProxy {
 
@@ -29,24 +29,35 @@ public class TCPReverseProxy {
 
     private ServerSocket serverSocket;
 
-    private final ExecutorService executor;
+    private final ExecutorService bossExecutor = Executors.newFixedThreadPool(1);
+
+    private final ExecutorService workerExecutor;
+
+
+    public TCPReverseProxy(String sourceHost, int sourcePort, String targetHost, int targetPort, ExecutorService workerExecutor) {
+        this.sourceHost = sourceHost;
+        this.sourcePort = sourcePort;
+        this.targetHost = targetHost;
+        this.targetPort = targetPort;
+        this.workerExecutor = workerExecutor;
+    }
 
     public TCPReverseProxy(String sourceHost, int sourcePort, String targetHost, int targetPort) {
         this.sourceHost = sourceHost;
         this.sourcePort = sourcePort;
         this.targetHost = targetHost;
         this.targetPort = targetPort;
-        executor = Executors.newFixedThreadPool(1);
+        workerExecutor = ForkJoinPool.commonPool();
     }
 
     public void start() {
-        executor.submit(() -> {
+        bossExecutor.submit(() -> {
             try {
                 serverSocket = new ServerSocket(sourcePort, 0, InetAddress.getByName(sourceHost));
                 log.info("proxy server started {}:{} <--> {}:{}", sourceHost, sourcePort, targetHost, targetPort);
                 while (!serverSocket.isClosed()) {
                     Socket sourceSocket = serverSocket.accept();
-                    CompletableFuture.runAsync(() -> {
+                    workerExecutor.execute(() -> {
                         try {
                             /**
                              * 将sourceSocket获取的内容，写入到targetSocket
@@ -55,7 +66,7 @@ public void start() {
                             Socket targetSocket = new Socket(targetHost, targetPort);
                             log.info("{} connected, proxy to {}", sourceSocket.getRemoteSocketAddress().toString(), targetSocket.getRemoteSocketAddress().toString());
                             // 监听源端主动传入的消息写给目标端
-                            CompletableFuture.runAsync(() -> {
+                            workerExecutor.execute(() -> {
                                 // isClosed只能监听本地连接状态。若远端关闭或者网络问题，无法监听到。
                                 if (sourceSocket.isClosed() || targetSocket.isClosed()) {
                                     return;
@@ -80,7 +91,7 @@ public void start() {
                                 }
                             });
                             // 监听目标端主动写回的消息写回源端
-                            CompletableFuture.runAsync(() -> {
+                            workerExecutor.execute(() -> {
                                 if (sourceSocket.isClosed() || targetSocket.isClosed()) {
                                     return;
                                 }
@@ -122,7 +133,7 @@ public void stop() {
             } catch (Exception ignore) {
             }
         }
-        executor.shutdown();
+        bossExecutor.shutdown();
         log.info("proxy server stoped {}:{} <--> {}:{}", sourceHost, sourcePort, targetHost, targetPort);
     }
 

src/test/java/top/meethigher/TCPReverseProxyTest.java

@@ -2,33 +2,37 @@
 
 import org.junit.Test;
 
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.ThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
 
 public class TCPReverseProxyTest {
 
 
     @Test
     public void testSSH() throws Exception {
-        TCPReverseProxy localhost = new TCPReverseProxy("localhost", 11, "10.0.0.9", 22);
+        TCPReverseProxy localhost = new TCPReverseProxy("127.0.0.1", 11, "10.0.0.9", 22);
         localhost.start();
-        TimeUnit.HOURS.sleep(1);
+        TimeUnit.MINUTES.sleep(1);
         localhost.stop();
     }
 
 
     @Test
     public void testPostgreSQL() throws Exception {
-        TCPReverseProxy localhost = new TCPReverseProxy("localhost", 22, "10.0.0.9", 5432);
+        TCPReverseProxy localhost = new TCPReverseProxy("0.0.0.0", 22,
+                "10.0.0.9", 5432,
+                new ThreadPoolExecutor(50, 50, 1, TimeUnit.MINUTES, new LinkedBlockingQueue<>()));
         localhost.start();
-        TimeUnit.HOURS.sleep(1);
+        TimeUnit.MINUTES.sleep(1);
         localhost.stop();
     }
 
     @Test
     public void testFTP() throws Exception {
-        TCPReverseProxy localhost = new TCPReverseProxy("localhost", 33, "localhost", 66);
+        TCPReverseProxy localhost = new TCPReverseProxy("127.0.0.1", 33, "127.0.0.1", 66);
         localhost.start();
-        TimeUnit.HOURS.sleep(1);
+        TimeUnit.MINUTES.sleep(1);
         localhost.stop();
     }
 }