Merge pull request #21 from Haox-liu/streamLoader_http_timeout

Fix stream loader http hang

Author: DanJing-coder

src/main/assembly/plugin.conf

@@ -54,5 +54,11 @@ secret_key=
 # Whether to generate sql digest for all queries
 enable_compute_all_query_digest=false
 
+# Http connectTimeout ms
+connect_timeout=1000
+
+# Http readTimeout ms
+read_timeout=1000
+
 # Filter conditions when importing audit information
 filter=

src/main/java/com/starrocks/plugin/audit/AuditLoaderPlugin.java

@@ -308,6 +308,8 @@ public static class AuditLoaderConf {
         public static final String QE_SLOW_LOG_MS = "qe_slow_log_ms";
         public static final String MAX_QUEUE_SIZE = "max_queue_size";
         public static final String ENABLE_COMPUTE_ALL_QUERY_DIGEST = "enable_compute_all_query_digest";
+        public static final String CONNECT_TIMEOUT = "connect_timeout";
+        public static final String READ_TIMEOUT = "read_timeout";
 
         public static final String STREAM_LOAD_FILTER = "filter";
 
@@ -325,6 +327,9 @@ public static class AuditLoaderConf {
         public int maxQueueSize = 1000;
 
         public boolean enableComputeAllQueryDigest = false;
+
+        public int connectTimeout = 1000;
+        public int readTimeout = 1000;
         public String streamLoadFilter = "";
 
         public static final String PROP_SECRET_KEY = "secret_key";
@@ -374,6 +379,12 @@ public void init(Map<String, String> properties) throws PluginException {
                 if (properties.containsKey(STREAM_LOAD_FILTER)) {
                     streamLoadFilter = properties.get(STREAM_LOAD_FILTER);
                 }
+                if (properties.containsKey(CONNECT_TIMEOUT)) {
+                    connectTimeout = Integer.parseInt(properties.get(CONNECT_TIMEOUT));
+                }
+                if (properties.containsKey(READ_TIMEOUT)) {
+                    readTimeout = Integer.parseInt(properties.get(READ_TIMEOUT));
+                }
             } catch (Exception e) {
                 throw new PluginException(e.getMessage());
             }

src/main/java/com/starrocks/plugin/audit/StarrocksStreamLoader.java

@@ -44,6 +44,10 @@ public class StarrocksStreamLoader {
     private String loadUrlStr;
     private String authEncoding;
     private String feIdentity;
+
+    private int connectTimeout;
+    private int readTimeout;
+
     private String streamLoadFilter;
 
     public StarrocksStreamLoader(AuditLoaderPlugin.AuditLoaderConf conf) {
@@ -65,6 +69,10 @@ public StarrocksStreamLoader(AuditLoaderPlugin.AuditLoaderConf conf) {
         this.authEncoding = Base64.getEncoder().encodeToString(String.format("%s:%s", user, passwd).getBytes(StandardCharsets.UTF_8));
         // currently, FE identity is FE's IP, so we replace the "." in IP to make it suitable for label
         this.feIdentity = conf.feIdentity.replaceAll("\\.", "_");
+
+        this.connectTimeout = conf.connectTimeout;
+        this.readTimeout = conf.readTimeout;
+      
         this.streamLoadFilter = conf.streamLoadFilter;
     }
 
@@ -89,6 +97,9 @@ private HttpURLConnection getConnection(String urlStr, String label) throws IOEx
         conn.setDoOutput(true);
         conn.setDoInput(true);
 
+        conn.setConnectTimeout(connectTimeout);
+        conn.setReadTimeout(readTimeout);
+
         return conn;
     }