Fix Azure http client leak in Azure filesystem

Author: findepi

lib/trino-filesystem-azure/pom.xml

@@ -105,6 +105,11 @@
             <artifactId>guice</artifactId>
         </dependency>
 
+        <dependency>
+            <groupId>com.squareup.okhttp3</groupId>
+            <artifactId>okhttp</artifactId>
+        </dependency>
+
         <dependency>
             <groupId>io.airlift</groupId>
             <artifactId>configuration</artifactId>
@@ -135,6 +140,11 @@
             <artifactId>trino-spi</artifactId>
         </dependency>
 
+        <dependency>
+            <groupId>jakarta.annotation</groupId>
+            <artifactId>jakarta.annotation-api</artifactId>
+        </dependency>
+
         <dependency>
             <groupId>jakarta.validation</groupId>
             <artifactId>jakarta.validation-api</artifactId>

lib/trino-filesystem-azure/src/main/java/io/trino/filesystem/azure/AzureFileSystemFactory.java

@@ -14,7 +14,7 @@
 package io.trino.filesystem.azure;
 
 import com.azure.core.http.HttpClient;
-import com.azure.core.http.okhttp.OkHttpAsyncClientProvider;
+import com.azure.core.http.okhttp.OkHttpAsyncHttpClientBuilder;
 import com.azure.core.tracing.opentelemetry.OpenTelemetryTracingOptions;
 import com.azure.core.util.HttpClientOptions;
 import com.azure.core.util.TracingOptions;
@@ -24,6 +24,11 @@
 import io.trino.filesystem.TrinoFileSystem;
 import io.trino.filesystem.TrinoFileSystemFactory;
 import io.trino.spi.security.ConnectorIdentity;
+import jakarta.annotation.PreDestroy;
+import okhttp3.ConnectionPool;
+import okhttp3.OkHttpClient;
+
+import java.util.concurrent.TimeUnit;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static java.util.Objects.requireNonNull;
@@ -37,6 +42,7 @@ public class AzureFileSystemFactory
     private final int maxWriteConcurrency;
     private final DataSize maxSingleUploadSize;
     private final TracingOptions tracingOptions;
+    private final OkHttpClient okHttpClient;
     private final HttpClient httpClient;
 
     @Inject
@@ -65,14 +71,39 @@ public AzureFileSystemFactory(
         this.maxWriteConcurrency = maxWriteConcurrency;
         this.maxSingleUploadSize = requireNonNull(maxSingleUploadSize, "maxSingleUploadSize is null");
         this.tracingOptions = new OpenTelemetryTracingOptions().setOpenTelemetry(openTelemetry);
-        this.httpClient = HttpClient.createDefault((HttpClientOptions) new HttpClientOptions()
-                .setHttpClientProvider(OkHttpAsyncClientProvider.class)
-                .setTracingOptions(tracingOptions));
+
+        okHttpClient = new OkHttpClient.Builder().build();
+        HttpClientOptions clientOptions = new HttpClientOptions();
+        clientOptions.setTracingOptions(tracingOptions);
+        httpClient = createAzureHttpClient(okHttpClient, clientOptions);
+    }
+
+    @PreDestroy
+    public void destroy()
+    {
+        okHttpClient.dispatcher().executorService().shutdownNow();
+        okHttpClient.connectionPool().evictAll();
     }
 
     @Override
     public TrinoFileSystem create(ConnectorIdentity identity)
     {
         return new AzureFileSystem(httpClient, tracingOptions, auth, readBlockSize, writeBlockSize, maxWriteConcurrency, maxSingleUploadSize);
     }
+
+    public static HttpClient createAzureHttpClient(OkHttpClient okHttpClient, HttpClientOptions clientOptions)
+    {
+        Integer poolSize = clientOptions.getMaximumConnectionPoolSize();
+        // By default, OkHttp uses a maximum idle connection count of 5.
+        int maximumConnectionPoolSize = (poolSize != null && poolSize > 0) ? poolSize : 5;
+
+        return new OkHttpAsyncHttpClientBuilder(okHttpClient)
+                .proxy(clientOptions.getProxyOptions())
+                .configuration(clientOptions.getConfiguration())
+                .connectionTimeout(clientOptions.getConnectTimeout())
+                .writeTimeout(clientOptions.getWriteTimeout())
+                .readTimeout(clientOptions.getReadTimeout())
+                .connectionPool(new ConnectionPool(maximumConnectionPoolSize,
+                        clientOptions.getConnectionIdleTimeout().toMillis(), TimeUnit.MILLISECONDS)).build();
+    }
 }

lib/trino-filesystem-azure/src/test/java/io/trino/filesystem/azure/AbstractTestAzureFileSystem.java

@@ -61,6 +61,7 @@ protected enum AccountKind
     private String containerName;
     private Location rootLocation;
     private BlobContainerClient blobContainerClient;
+    private AzureFileSystemFactory fileSystemFactory;
     private TrinoFileSystem fileSystem;
 
     protected void initializeWithAccessKey(String account, String accountKey, AccountKind accountKind)
@@ -100,10 +101,11 @@ private void initialize(String account, AzureAuth azureAuth, AccountKind account
             checkState(!isHierarchicalNamespaceEnabled, "Expected hierarchical namespaces to not be enabled for storage account %s and container %s with account kind %s".formatted(account, containerName, accountKind));
         }
 
-        fileSystem = new AzureFileSystemFactory(
+        fileSystemFactory = new AzureFileSystemFactory(
                 OpenTelemetry.noop(),
                 azureAuth,
-                new AzureFileSystemConfig()).create(ConnectorIdentity.ofUser("test"));
+                new AzureFileSystemConfig());
+        fileSystem = fileSystemFactory.create(ConnectorIdentity.ofUser("test"));
 
         cleanupFiles();
     }
@@ -124,6 +126,10 @@ private boolean isHierarchicalNamespaceEnabled()
     void tearDown()
     {
         azureAuth = null;
+        if (fileSystemFactory != null) {
+            fileSystemFactory.destroy();
+            fileSystemFactory = null;
+        }
         fileSystem = null;
         if (blobContainerClient != null) {
             blobContainerClient.deleteIfExists();

plugin/trino-delta-lake/pom.xml

@@ -249,6 +249,12 @@
             </exclusions>
         </dependency>
 
+        <dependency>
+            <groupId>com.squareup.okhttp3</groupId>
+            <artifactId>okhttp</artifactId>
+            <scope>runtime</scope>
+        </dependency>
+
         <dependency>
             <groupId>io.airlift</groupId>
             <artifactId>log-manager</artifactId>
@@ -261,6 +267,12 @@
             <scope>runtime</scope>
         </dependency>
 
+        <dependency>
+            <groupId>io.trino</groupId>
+            <artifactId>trino-filesystem-azure</artifactId>
+            <scope>runtime</scope>
+        </dependency>
+
         <dependency>
             <groupId>org.jetbrains</groupId>
             <artifactId>annotations</artifactId>

plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeAdlsConnectorSmokeTest.java

@@ -13,6 +13,7 @@
  */
 package io.trino.plugin.deltalake;
 
+import com.azure.core.util.HttpClientOptions;
 import com.azure.storage.blob.BlobContainerClient;
 import com.azure.storage.blob.BlobServiceClient;
 import com.azure.storage.blob.BlobServiceClientBuilder;
@@ -24,6 +25,7 @@
 import com.google.common.reflect.ClassPath;
 import io.trino.plugin.hive.containers.HiveHadoop;
 import io.trino.testing.QueryRunner;
+import okhttp3.OkHttpClient;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.TestInstance;
 
@@ -42,6 +44,7 @@
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
+import static io.trino.filesystem.azure.AzureFileSystemFactory.createAzureHttpClient;
 import static io.trino.plugin.hive.containers.HiveHadoop.HIVE3_IMAGE;
 import static io.trino.testing.TestingProperties.requiredNonEmptySystemProperty;
 import static java.lang.String.format;
@@ -58,25 +61,32 @@ public class TestDeltaLakeAdlsConnectorSmokeTest
     private final String container;
     private final String account;
     private final String accessKey;
-    private final BlobContainerClient azureContainerClient;
     private final String adlsDirectory;
+    private BlobContainerClient azureContainerClient;
 
     public TestDeltaLakeAdlsConnectorSmokeTest()
     {
         this.container = requiredNonEmptySystemProperty("testing.azure-abfs-container");
         this.account = requiredNonEmptySystemProperty("testing.azure-abfs-account");
         this.accessKey = requiredNonEmptySystemProperty("testing.azure-abfs-access-key");
-
-        String connectionString = format("DefaultEndpointsProtocol=https;AccountName=%s;AccountKey=%s;EndpointSuffix=core.windows.net", account, accessKey);
-        BlobServiceClient blobServiceClient = new BlobServiceClientBuilder().connectionString(connectionString).buildClient();
-        this.azureContainerClient = blobServiceClient.getBlobContainerClient(container);
         this.adlsDirectory = format("abfs://%s@%s.dfs.core.windows.net/%s/", container, account, bucketName);
     }
 
     @Override
     protected HiveHadoop createHiveHadoop()
             throws Exception
     {
+        String connectionString = format("DefaultEndpointsProtocol=https;AccountName=%s;AccountKey=%s;EndpointSuffix=core.windows.net", account, accessKey);
+        OkHttpClient okHttpClient = new OkHttpClient.Builder().build();
+        closeAfterClass(() -> {
+            okHttpClient.dispatcher().executorService().shutdownNow();
+            okHttpClient.connectionPool().evictAll();
+        });
+        BlobServiceClient blobServiceClient = new BlobServiceClientBuilder().connectionString(connectionString)
+                .httpClient(createAzureHttpClient(okHttpClient, new HttpClientOptions()))
+                .buildClient();
+        this.azureContainerClient = blobServiceClient.getBlobContainerClient(container);
+
         String abfsSpecificCoreSiteXmlContent = Resources.toString(Resources.getResource("io/trino/plugin/deltalake/hdp3.1-core-site.xml.abfs-template"), UTF_8)
                 .replace("%ABFS_ACCESS_KEY%", accessKey)
                 .replace("%ABFS_ACCOUNT%", account);