[Improvement-18611][Task] Support configurable TCP keepalive on the HTTP task

dolphinscheduler-common/src/main/java/org/apache/dolphinscheduler/common/utils/OkHttpUtils.java

@@ -23,10 +23,13 @@
 import org.apache.http.HttpStatus;
 
 import java.io.IOException;
+import java.net.InetAddress;
+import java.net.Socket;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
 
 import javax.annotation.Nullable;
+import javax.net.SocketFactory;
 
 import lombok.NonNull;
 import okhttp3.HttpUrl;
@@ -36,10 +39,60 @@
 import okhttp3.RequestBody;
 import okhttp3.Response;
 
+import com.google.common.annotations.VisibleForTesting;
+
 public class OkHttpUtils {
 
     private static OkHttpClient CLIENT = new OkHttpClient();
 
+    /**
+     * Socket factory that enables TCP keepalive (SO_KEEPALIVE) on the sockets OkHttp creates.
+     */
+    private static final SocketFactory KEEP_ALIVE_SOCKET_FACTORY = new SocketFactory() {
+
+        @Override
+        public Socket createSocket() throws IOException {
+            return enableTcpKeepAlive(SocketFactory.getDefault().createSocket());
+        }
+
+        @Override
+        public Socket createSocket(String host, int port) throws IOException {
+            return enableTcpKeepAlive(SocketFactory.getDefault().createSocket(host, port));
+        }
+
+        @Override
+        public Socket createSocket(String host, int port, InetAddress localHost, int localPort) throws IOException {
+            return enableTcpKeepAlive(SocketFactory.getDefault().createSocket(host, port, localHost, localPort));
+        }
+
+        @Override
+        public Socket createSocket(InetAddress host, int port) throws IOException {
+            return enableTcpKeepAlive(SocketFactory.getDefault().createSocket(host, port));
+        }
+
+        @Override
+        public Socket createSocket(InetAddress address, int port, InetAddress localAddress,
+                                   int localPort) throws IOException {
+            return enableTcpKeepAlive(
+                    SocketFactory.getDefault().createSocket(address, port, localAddress, localPort));
+        }
+
+        private Socket enableTcpKeepAlive(Socket socket) throws IOException {
+            socket.setKeepAlive(true);
+            return socket;
+        }
+    };
+
+    /**
+     * Base client for requests with TCP keepalive enabled. OkHttp ignores the socket factory when
+     * deciding whether a pooled connection can be reused, so it must not share the connection pool
+     * with {@link #CLIENT}, otherwise a connection created without keepalive could be reused by a
+     * keepalive request and vice versa.
+     */
+    private static final OkHttpClient KEEP_ALIVE_CLIENT = new OkHttpClient().newBuilder()
+            .socketFactory(KEEP_ALIVE_SOCKET_FACTORY)
+            .build();
+
     /**
      * http get request
      * @param connectTimeout connect timeout in milliseconds
@@ -54,7 +107,27 @@ public class OkHttpUtils {
                                               int connectTimeout,
                                               int writeTimeout,
                                               int readTimeout) throws IOException {
-        OkHttpClient client = getHttpClient(connectTimeout, writeTimeout, readTimeout);
+        return get(url, okHttpRequestHeaders, requestParams,
+                connectTimeout, writeTimeout, readTimeout, false);
+    }
+
+    /**
+     * http get request
+     * @param connectTimeout connect timeout in milliseconds
+     * @param writeTimeout write timeout in milliseconds
+     * @param readTimeout read timeout in milliseconds
+     * @param keepAlive whether to enable TCP keepalive (SO_KEEPALIVE) on the socket
+     * @return OkHttpResponse
+     * @throws RuntimeException
+     */
+    public static @NonNull OkHttpResponse get(@NonNull String url,
+                                              @Nullable OkHttpRequestHeaders okHttpRequestHeaders,
+                                              @Nullable Map<String, Object> requestParams,
+                                              int connectTimeout,
+                                              int writeTimeout,
+                                              int readTimeout,
+                                              boolean keepAlive) throws IOException {
+        OkHttpClient client = getHttpClient(connectTimeout, writeTimeout, readTimeout, keepAlive);
         String finalUrl = addUrlParams(requestParams, url);
         Request.Builder requestBuilder = new Request.Builder().url(finalUrl);
         addHeader(okHttpRequestHeaders.getHeaders(), requestBuilder);
@@ -81,7 +154,28 @@ public class OkHttpUtils {
                                                int connectTimeout,
                                                int writeTimeout,
                                                int readTimeout) throws IOException {
-        OkHttpClient client = getHttpClient(connectTimeout, writeTimeout, readTimeout);
+        return post(url, okHttpRequestHeaders, requestParamsMap, requestBodyMap,
+                connectTimeout, writeTimeout, readTimeout, false);
+    }
+
+    /**
+     * http post request
+     * @param connectTimeout connect timeout in milliseconds
+     * @param writeTimeout write timeout in milliseconds
+     * @param readTimeout read timeout in milliseconds
+     * @param keepAlive whether to enable TCP keepalive (SO_KEEPALIVE) on the socket
+     * @return OkHttpResponse
+     * @throws RuntimeException
+     */
+    public static @NonNull OkHttpResponse post(@NonNull String url,
+                                               @Nullable OkHttpRequestHeaders okHttpRequestHeaders,
+                                               @Nullable Map<String, Object> requestParamsMap,
+                                               @Nullable Map<String, Object> requestBodyMap,
+                                               int connectTimeout,
+                                               int writeTimeout,
+                                               int readTimeout,
+                                               boolean keepAlive) throws IOException {
+        OkHttpClient client = getHttpClient(connectTimeout, writeTimeout, readTimeout, keepAlive);
         String finalUrl = addUrlParams(requestParamsMap, url);
         Request.Builder requestBuilder = new Request.Builder().url(finalUrl);
         addHeader(okHttpRequestHeaders.getHeaders(), requestBuilder);
@@ -111,7 +205,27 @@ public class OkHttpUtils {
                                               int connectTimeout,
                                               int writeTimeout,
                                               int readTimeout) throws IOException {
-        OkHttpClient client = getHttpClient(connectTimeout, writeTimeout, readTimeout);
+        return put(url, okHttpRequestHeaders, requestBodyMap,
+                connectTimeout, writeTimeout, readTimeout, false);
+    }
+
+    /**
+     * http put request
+     * @param connectTimeout connect timeout in milliseconds
+     * @param writeTimeout write timeout in milliseconds
+     * @param readTimeout read timeout in milliseconds
+     * @param keepAlive whether to enable TCP keepalive (SO_KEEPALIVE) on the socket
+     * @return OkHttpResponse
+     * @throws RuntimeException
+     */
+    public static @NonNull OkHttpResponse put(@NonNull String url,
+                                              @Nullable OkHttpRequestHeaders okHttpRequestHeaders,
+                                              @Nullable Map<String, Object> requestBodyMap,
+                                              int connectTimeout,
+                                              int writeTimeout,
+                                              int readTimeout,
+                                              boolean keepAlive) throws IOException {
+        OkHttpClient client = getHttpClient(connectTimeout, writeTimeout, readTimeout, keepAlive);
         Request.Builder requestBuilder = new Request.Builder().url(url);
         addHeader(okHttpRequestHeaders.getHeaders(), requestBuilder);
         if (requestBodyMap != null) {
@@ -139,7 +253,25 @@ public class OkHttpUtils {
                                                  int connectTimeout,
                                                  int writeTimeout,
                                                  int readTimeout) throws IOException {
-        OkHttpClient client = getHttpClient(connectTimeout, writeTimeout, readTimeout);
+        return delete(url, okHttpRequestHeaders, connectTimeout, writeTimeout, readTimeout, false);
+    }
+
+    /**
+     * http delete request
+     * @param connectTimeout connect timeout in milliseconds
+     * @param writeTimeout write timeout in milliseconds
+     * @param readTimeout read timeout in milliseconds
+     * @param keepAlive whether to enable TCP keepalive (SO_KEEPALIVE) on the socket
+     * @return OkHttpResponse
+     * @throws RuntimeException
+     */
+    public static @NonNull OkHttpResponse delete(@NonNull String url,
+                                                 @Nullable OkHttpRequestHeaders okHttpRequestHeaders,
+                                                 int connectTimeout,
+                                                 int writeTimeout,
+                                                 int readTimeout,
+                                                 boolean keepAlive) throws IOException {
+        OkHttpClient client = getHttpClient(connectTimeout, writeTimeout, readTimeout, keepAlive);
         Request.Builder requestBuilder = new Request.Builder().url(url);
         addHeader(okHttpRequestHeaders.getHeaders(), requestBuilder);
         requestBuilder = requestBuilder.delete();
@@ -182,10 +314,13 @@ private static String getResponseBody(@NonNull Response response) throws IOExcep
         return response.body().string();
     }
 
-    private static OkHttpClient getHttpClient(int connectTimeout,
-                                              int writeTimeout,
-                                              int readTimeout) {
-        return CLIENT.newBuilder()
+    @VisibleForTesting
+    static OkHttpClient getHttpClient(int connectTimeout,
+                                      int writeTimeout,
+                                      int readTimeout,
+                                      boolean keepAlive) {
+        OkHttpClient baseClient = keepAlive ? KEEP_ALIVE_CLIENT : CLIENT;
+        return baseClient.newBuilder()
                 .connectTimeout(connectTimeout, TimeUnit.MILLISECONDS)
                 .writeTimeout(writeTimeout, TimeUnit.MILLISECONDS)
                 .readTimeout(readTimeout, TimeUnit.MILLISECONDS)

dolphinscheduler-common/src/test/java/org/apache/dolphinscheduler/common/utils/OkHttpUtilsTest.java

@@ -0,0 +1,135 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.dolphinscheduler.common.utils;
+
+import org.apache.dolphinscheduler.common.model.OkHttpRequestHeaders;
+
+import java.io.IOException;
+import java.io.OutputStream;
+import java.net.InetSocketAddress;
+import java.net.Socket;
+import java.nio.charset.StandardCharsets;
+import java.util.HashSet;
+import java.util.List;
+import java.util.concurrent.CopyOnWriteArrayList;
+
+import okhttp3.OkHttpClient;
+
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+import com.sun.net.httpserver.HttpServer;
+
+class OkHttpUtilsTest {
+
+    private static final int TIMEOUT = 10_000;
+
+    @Test
+    void testKeepAliveAndPlainClientsAreSeparated() throws Exception {
+        OkHttpClient plainClient = OkHttpUtils.getHttpClient(TIMEOUT, TIMEOUT, TIMEOUT, false);
+        OkHttpClient keepAliveClient = OkHttpUtils.getHttpClient(TIMEOUT, TIMEOUT, TIMEOUT, true);
+
+        Assertions.assertNotSame(plainClient.connectionPool(), keepAliveClient.connectionPool(),
+                "the two settings must not share a connection pool");
+        try (Socket plainSocket = plainClient.socketFactory().createSocket()) {
+            Assertions.assertFalse(plainSocket.getKeepAlive(), "the plain client must not enable TCP keepalive");
+        }
+        try (Socket keepAliveSocket = keepAliveClient.socketFactory().createSocket()) {
+            Assertions.assertTrue(keepAliveSocket.getKeepAlive(), "the keepalive client must enable TCP keepalive");
+        }
+    }
+
+    @Test
+    void testKeepAliveRequestDoesNotReusePlainConnection() throws Exception {
+        resetConnectionPools();
+        List<Integer> remotePorts = new CopyOnWriteArrayList<>();
+        HttpServer server = startServer(remotePorts);
+        try {
+            String url = urlOf(server);
+            OkHttpUtils.get(url, new OkHttpRequestHeaders(), null, TIMEOUT, TIMEOUT, TIMEOUT, false);
+            assertPooledConnectionCount(false, 1);
+            assertPooledConnectionCount(true, 0);
+
+            OkHttpUtils.get(url, new OkHttpRequestHeaders(), null, TIMEOUT, TIMEOUT, TIMEOUT, true);
+            assertPooledConnectionCount(false, 1);
+            assertPooledConnectionCount(true, 1);
+            assertRequestsUseDistinctConnections(remotePorts);
+        } finally {
+            server.stop(0);
+        }
+    }
+
+    @Test
+    void testPlainRequestDoesNotReuseKeepAliveConnection() throws Exception {
+        resetConnectionPools();
+        List<Integer> remotePorts = new CopyOnWriteArrayList<>();
+        HttpServer server = startServer(remotePorts);
+        try {
+            String url = urlOf(server);
+            OkHttpUtils.get(url, new OkHttpRequestHeaders(), null, TIMEOUT, TIMEOUT, TIMEOUT, true);
+            assertPooledConnectionCount(true, 1);
+            assertPooledConnectionCount(false, 0);
+
+            OkHttpUtils.get(url, new OkHttpRequestHeaders(), null, TIMEOUT, TIMEOUT, TIMEOUT, false);
+            assertPooledConnectionCount(true, 1);
+            assertPooledConnectionCount(false, 1);
+            assertRequestsUseDistinctConnections(remotePorts);
+        } finally {
+            server.stop(0);
+        }
+    }
+
+    private static void assertRequestsUseDistinctConnections(List<Integer> remotePorts) {
+        Assertions.assertEquals(2, remotePorts.size(), "both requests should reach the server");
+        Assertions.assertEquals(2, new HashSet<>(remotePorts).size(),
+                "the two requests reused one TCP connection, remote ports: " + remotePorts);
+    }
+
+    private static void assertPooledConnectionCount(boolean keepAlive, int expected) {
+        int actual = clientOf(keepAlive).connectionPool().connectionCount();
+        Assertions.assertEquals(expected, actual,
+                "unexpected pooled connection count of the keepAlive=" + keepAlive + " client");
+    }
+
+    private static void resetConnectionPools() {
+        clientOf(false).connectionPool().evictAll();
+        clientOf(true).connectionPool().evictAll();
+    }
+
+    private static OkHttpClient clientOf(boolean keepAlive) {
+        return OkHttpUtils.getHttpClient(TIMEOUT, TIMEOUT, TIMEOUT, keepAlive);
+    }
+
+    private static HttpServer startServer(List<Integer> remotePorts) throws IOException {
+        HttpServer server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0);
+        server.createContext("/keep-alive", exchange -> {
+            remotePorts.add(exchange.getRemoteAddress().getPort());
+            byte[] body = "ok".getBytes(StandardCharsets.UTF_8);
+            exchange.sendResponseHeaders(200, body.length);
+            try (OutputStream outputStream = exchange.getResponseBody()) {
+                outputStream.write(body);
+            }
+        });
+        server.start();
+        return server;
+    }
+
+    private static String urlOf(HttpServer server) {
+        return "http://127.0.0.1:" + server.getAddress().getPort() + "/keep-alive";
+    }
+}

dolphinscheduler-task-plugin/dolphinscheduler-task-http/src/main/java/org/apache/dolphinscheduler/plugin/task/http/HttpParameters.java

@@ -56,6 +56,12 @@ public class HttpParameters extends AbstractParameters {
      */
     private int connectTimeout;
 
+    /**
+     * Whether to enable TCP keepalive (SO_KEEPALIVE) on the request socket.
+     * The probe interval is controlled by the worker host OS, not by this flag.
+     */
+    private boolean socketKeepAlive;
+
     @Override
     public boolean checkParameters() {
         return StringUtils.isNotEmpty(url) && httpRequestMethod != null

dolphinscheduler-task-plugin/dolphinscheduler-task-http/src/main/java/org/apache/dolphinscheduler/plugin/task/http/HttpTask.java

@@ -154,7 +154,8 @@ private OkHttpResponse sendGetRequest() {
 
         OkHttpResponse okHttpResponse = OkHttpUtils.get(httpParameters.getUrl(), okHttpRequestHeaders,
                 requestParams, httpParameters.getConnectTimeout(),
-                httpParameters.getConnectTimeout(), httpParameters.getConnectTimeout());
+                httpParameters.getConnectTimeout(), httpParameters.getConnectTimeout(),
+                httpParameters.isSocketKeepAlive());
         addDefaultOutput(JSONUtils.toJsonString(okHttpResponse));
         return okHttpResponse;
     }
@@ -168,7 +169,8 @@ private OkHttpResponse sendPostRequest() {
 
         OkHttpResponse okHttpResponse = OkHttpUtils.post(httpParameters.getUrl(), okHttpRequestHeaders, null,
                 requestBody, httpParameters.getConnectTimeout(),
-                httpParameters.getConnectTimeout(), httpParameters.getConnectTimeout());
+                httpParameters.getConnectTimeout(), httpParameters.getConnectTimeout(),
+                httpParameters.isSocketKeepAlive());
         addDefaultOutput(JSONUtils.toJsonString(okHttpResponse));
         return okHttpResponse;
     }
@@ -182,7 +184,8 @@ private OkHttpResponse sendPutRequest() {
 
         OkHttpResponse okHttpResponse = OkHttpUtils.put(httpParameters.getUrl(), okHttpRequestHeaders,
                 requestBody, httpParameters.getConnectTimeout(),
-                httpParameters.getConnectTimeout(), httpParameters.getConnectTimeout());
+                httpParameters.getConnectTimeout(), httpParameters.getConnectTimeout(),
+                httpParameters.isSocketKeepAlive());
         addDefaultOutput(JSONUtils.toJsonString(okHttpResponse));
         return okHttpResponse;
     }
@@ -195,7 +198,7 @@ private OkHttpResponse sendDeleteRequest() {
 
         OkHttpResponse okHttpResponse = OkHttpUtils.delete(httpParameters.getUrl(), okHttpRequestHeaders,
                 httpParameters.getConnectTimeout(), httpParameters.getConnectTimeout(),
-                httpParameters.getConnectTimeout());
+                httpParameters.getConnectTimeout(), httpParameters.isSocketKeepAlive());
         addDefaultOutput(JSONUtils.toJsonString(okHttpResponse));
         return okHttpResponse;
     }