[PECO-2009] Improvements to thrift-server client (#552)

apache http client maintains a pool of connections to mitigate TCP-handshake for every request. however, server can unilaterally decide to close the connection to avoid resource drainage. a problem with classical network IO model (used by apache client 4.x) is that the socket only reacts to incoming events when it is blocked on an operation. the client, therefore can miss the event of server closing the connection. this creates a need to remove expired/idle connections from the pool on client side.

there are broadly two ways of doing so:

close such connections manually in different code paths (currently being done using the method)

create a background thread. from apache http client 4.5 onwards, they have included [IdleConnectionEvictor](https://www.javadoc.io/doc/org.apache.httpcomponents/httpclient/4.5.5/org/apache/http/impl/client/IdleConnectionEvictor.html) in library so that users don't have to write code for the background thread. we can use this to avoid manually calling the method

summarising the improvements:

- use byte array input stream for receiving data structure
- use IdleConnectionEvictor to close expired/idle connections

Author: jayantsing-db

src/main/java/com/databricks/jdbc/api/impl/arrow/RemoteChunkProvider.java

@@ -131,7 +131,6 @@ public ArrowResultChunk getChunk() throws DatabricksSQLException {
       return null;
     }
     ArrowResultChunk chunk = chunkIndexToChunksMap.get(currentChunkIndex);
-    httpClient.closeExpiredAndIdleConnections();
     synchronized (chunk) {
       try {
         while (!isDownloadComplete(chunk.getStatus())) {
@@ -188,7 +187,6 @@ public void close() {
     this.isClosed = true;
     this.chunkDownloaderExecutorService.shutdownNow();
     this.chunkIndexToChunksMap.values().forEach(ArrowResultChunk::releaseChunk);
-    httpClient.closeExpiredAndIdleConnections();
   }
 
   /** Release the memory for previous chunk since it is already consumed */

src/main/java/com/databricks/jdbc/dbclient/IDatabricksHttpClient.java

@@ -14,6 +14,4 @@ public interface IDatabricksHttpClient {
    * @return http response
    */
   CloseableHttpResponse execute(HttpUriRequest request) throws DatabricksHttpException;
-
-  void closeExpiredAndIdleConnections();
 }

src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpClient.java

@@ -33,6 +33,7 @@
 import org.apache.http.conn.ssl.NoopHostnameVerifier;
 import org.apache.http.impl.client.CloseableHttpClient;
 import org.apache.http.impl.client.HttpClientBuilder;
+import org.apache.http.impl.client.IdleConnectionEvictor;
 import org.apache.http.impl.conn.DefaultSchemePortResolver;
 import org.apache.http.impl.conn.PoolingHttpClientConnectionManager;
 import org.apache.http.ssl.SSLContextBuilder;
@@ -52,15 +53,17 @@ public class DatabricksHttpClient implements IDatabricksHttpClient {
   private static PoolingHttpClientConnectionManager connectionManager;
   private final CloseableHttpClient httpClient;
   protected static int idleHttpConnectionExpiry;
-  private CloseableHttpClient httpDisabledSSLClient;
   private DatabricksHttpRetryHandler retryHandler;
+  private IdleConnectionEvictor idleConnectionEvictor;
 
   private DatabricksHttpClient(IDatabricksConnectionContext connectionContext) {
     initializeConnectionManager(connectionContext);
     httpClient = makeClosableHttpClient(connectionContext);
-    httpDisabledSSLClient = makeClosableDisabledSslHttpClient();
     idleHttpConnectionExpiry = connectionContext.getIdleHttpConnectionExpiry();
     retryHandler = new DatabricksHttpRetryHandler(connectionContext);
+    idleConnectionEvictor =
+        new IdleConnectionEvictor(connectionManager, idleHttpConnectionExpiry, TimeUnit.SECONDS);
+    idleConnectionEvictor.start();
   }
 
   @VisibleForTesting
@@ -219,30 +222,6 @@ public CloseableHttpResponse execute(HttpUriRequest request) throws DatabricksHt
     return null;
   }
 
-  public CloseableHttpResponse executeWithoutCertVerification(HttpUriRequest request)
-      throws DatabricksHttpException {
-    LOGGER.debug(
-        String.format("Executing HTTP request [{%s}]", RequestSanitizer.sanitizeRequest(request)));
-    try {
-      return httpDisabledSSLClient.execute(request);
-    } catch (Exception e) {
-      throwHttpException(e, request, LogLevel.DEBUG);
-    }
-    return null;
-  }
-
-  @Override
-  public void closeExpiredAndIdleConnections() {
-    if (connectionManager != null) {
-      synchronized (connectionManager) {
-        LOGGER.debug(
-            String.format("connection pool stats: {%s}", connectionManager.getTotalStats()));
-        connectionManager.closeExpiredConnections();
-        connectionManager.closeIdleConnections(idleHttpConnectionExpiry, TimeUnit.SECONDS);
-      }
-    }
-  }
-
   static String getUserAgent() {
     String sdkUserAgent = UserAgent.asString();
     // Split the string into parts
@@ -270,6 +249,9 @@ public static synchronized void removeInstance(IDatabricksConnectionContext cont
     DatabricksHttpClient instance = instances.remove(contextKey);
     if (instance != null) {
       try {
+        if (instance.idleConnectionEvictor != null) {
+          instance.idleConnectionEvictor.shutdown();
+        }
         instance.httpClient.close();
       } catch (IOException e) {
         LOGGER.debug(String.format("Caught error while closing http client. Error %s", e));

src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksHttpTTransport.java

@@ -6,15 +6,17 @@
 import com.databricks.jdbc.log.JdbcLogger;
 import com.databricks.jdbc.log.JdbcLoggerFactory;
 import com.google.common.annotations.VisibleForTesting;
+import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
-import java.io.InputStream;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
+import org.apache.http.HttpEntity;
 import org.apache.http.client.methods.CloseableHttpResponse;
 import org.apache.http.client.methods.HttpPost;
 import org.apache.http.entity.ByteArrayEntity;
+import org.apache.http.util.EntityUtils;
 import org.apache.thrift.TConfiguration;
 import org.apache.thrift.transport.TTransport;
 import org.apache.thrift.transport.TTransportException;
@@ -23,20 +25,19 @@ public class DatabricksHttpTTransport extends TTransport {
 
   private static final JdbcLogger LOGGER =
       JdbcLoggerFactory.getLogger(DatabricksHttpTTransport.class);
+  private static final Map<String, String> DEFAULT_HEADERS =
+      Collections.unmodifiableMap(getDefaultHeaders());
   private final DatabricksHttpClient httpClient;
   private final String url;
   private Map<String, String> customHeaders = Collections.emptyMap();
   private final ByteArrayOutputStream requestBuffer;
-  private InputStream inputStream = null;
-  private CloseableHttpResponse response = null;
-  private static final Map<String, String> DEFAULT_HEADERS =
-      Collections.unmodifiableMap(getDefaultHeaders());
+  private ByteArrayInputStream responseBuffer;
 
   public DatabricksHttpTTransport(DatabricksHttpClient httpClient, String url) {
     this.httpClient = httpClient;
-    this.httpClient.closeExpiredAndIdleConnections();
     this.url = url;
     this.requestBuffer = new ByteArrayOutputStream();
+    this.responseBuffer = null;
   }
 
   @Override
@@ -51,65 +52,20 @@ public void open() throws TTransportException {
   }
 
   @Override
-  public void close() {
-    this.httpClient.closeExpiredAndIdleConnections();
-    if (inputStream != null) {
-      try {
-        inputStream.close();
-      } catch (IOException e) {
-        LOGGER.error(
-            String.format("Failed to close inputStream with error {%s}. Skipping the close.", e));
-      }
-      inputStream = null;
-    }
-    if (response != null) {
-      try {
-        response.close();
-      } catch (IOException e) {
-        LOGGER.error(String.format("Failed to close response with error {%s}", e.toString()));
-      }
-      response = null;
-    }
-  }
-
-  public void setCustomHeaders(Map<String, String> headers) {
-    if (headers != null) {
-      customHeaders = new HashMap<>(headers);
-    } else {
-      customHeaders = Collections.emptyMap();
-    }
-  }
-
-  @VisibleForTesting
-  Map<String, String> getCustomHeaders() {
-    return customHeaders;
-  }
-
-  @VisibleForTesting
-  InputStream getInputStream() {
-    return inputStream;
-  }
-
-  @VisibleForTesting
-  void setInputStream(InputStream inputStream) {
-    this.inputStream = inputStream;
-  }
+  public void close() {}
 
   @Override
   public int read(byte[] buf, int off, int len) throws TTransportException {
-    if (inputStream == null) {
-      throw new TTransportException("Response buffer is empty, no request.");
+    if (responseBuffer == null) {
+      LOGGER.error("Response buffer is empty, no response.");
+      throw new TTransportException("Response buffer is empty, no response.");
     }
-    try {
-      int ret = inputStream.read(buf, off, len);
-      if (ret == -1) {
-        throw new TTransportException("No more data available.");
-      }
-      return ret;
-    } catch (IOException e) {
-      LOGGER.error(String.format("Failed to read inputStream with error {%s}", e.toString()));
-      throw new TTransportException(e);
+    int numBytes = responseBuffer.read(buf, off, len);
+    if (numBytes == -1) {
+      LOGGER.error("No data available to read.");
+      throw new TTransportException("No more data available.");
     }
+    return numBytes;
   }
 
   @Override
@@ -119,32 +75,34 @@ public void write(byte[] buf, int off, int len) {
 
   @Override
   public void flush() throws TTransportException {
-    try {
-      HttpPost request = new HttpPost(this.url);
-      DEFAULT_HEADERS.forEach(request::addHeader);
-      if (customHeaders != null) {
-        customHeaders.forEach(request::addHeader);
-      }
-      request.setEntity(new ByteArrayEntity(requestBuffer.toByteArray()));
-      response = httpClient.execute(request);
+    HttpPost request = new HttpPost(this.url);
+    DEFAULT_HEADERS.forEach(request::addHeader);
+
+    if (customHeaders != null) {
+      customHeaders.forEach(request::addHeader);
+    }
+
+    // Set the request entity
+    request.setEntity(new ByteArrayEntity(requestBuffer.toByteArray()));
+
+    // Execute the request and handle the response
+    try (CloseableHttpResponse response = httpClient.execute(request)) {
       ValidationUtil.checkHTTPError(response);
-      inputStream = response.getEntity().getContent();
-      requestBuffer.reset();
-    } catch (DatabricksHttpException | IOException e) {
-      Throwable cause = e;
-      while (cause != null) {
-        if (cause instanceof DatabricksHttpException) {
-          throw new TTransportException(
-              TTransportException.UNKNOWN, "Failed to flush data to server: " + cause.getMessage());
-        }
-        cause = cause.getCause();
-      }
-      httpClient.closeExpiredAndIdleConnections();
 
+      // Read the response
+      HttpEntity entity = response.getEntity();
+      if (entity != null) {
+        byte[] responseBytes = EntityUtils.toByteArray(entity);
+        responseBuffer = new ByteArrayInputStream(responseBytes);
+      }
+    } catch (DatabricksHttpException | IOException e) {
       String errorMessage = "Failed to flush data to server: " + e.getMessage();
-      LOGGER.error(errorMessage);
+      LOGGER.error(e, errorMessage);
       throw new TTransportException(TTransportException.UNKNOWN, errorMessage);
     }
+
+    // Reset the request buffer
+    requestBuffer.reset();
   }
 
   @Override
@@ -158,10 +116,28 @@ public void updateKnownMessageSize(long size) throws TTransportException {}
   @Override
   public void checkReadBytesAvailable(long numBytes) throws TTransportException {}
 
+  public void setCustomHeaders(Map<String, String> headers) {
+    if (headers != null) {
+      customHeaders = new HashMap<>(headers);
+    } else {
+      customHeaders = Collections.emptyMap();
+    }
+  }
+
   private static Map<String, String> getDefaultHeaders() {
     Map<String, String> headers = new HashMap<>();
     headers.put("Content-Type", "application/x-thrift");
     headers.put("Accept", "application/x-thrift");
     return headers;
   }
+
+  @VisibleForTesting
+  Map<String, String> getCustomHeaders() {
+    return customHeaders;
+  }
+
+  @VisibleForTesting
+  void setResponseBuffer(ByteArrayInputStream responseBuffer) {
+    this.responseBuffer = responseBuffer;
+  }
 }