Optimizing performance for Pulsar function archive download (apache#4082

) * Optimizing performance for download * fixing update and create * fix upload * remove commented out code * remove blank lines * fix error messages * fix for sources and sinks * fix getting exception * fix auth headers * cleaning up * fix print messages
JamieCressey · Apr 20, 2019 · c3e8a33 · c3e8a33
1 parent e5fbb89
commit c3e8a33
Show file tree

Hide file tree

Showing 9 changed files with 242 additions and 85 deletions.
diff --git a/pulsar-client-admin/src/main/java/org/apache/pulsar/client/admin/PulsarAdmin.java b/pulsar-client-admin/src/main/java/org/apache/pulsar/client/admin/PulsarAdmin.java
@@ -51,6 +51,7 @@
 import org.apache.pulsar.client.api.PulsarClientException;
 import org.apache.pulsar.client.impl.auth.AuthenticationDisabled;
 import org.apache.pulsar.client.impl.conf.ClientConfigurationData;
+import org.asynchttpclient.AsyncHttpClient;
 import org.glassfish.jersey.client.ClientConfig;
 import org.glassfish.jersey.client.ClientProperties;
 import org.glassfish.jersey.jackson.JacksonFeature;
@@ -81,6 +82,7 @@ public class PulsarAdmin implements Closeable {
     private final ResourceQuotas resourceQuotas;
     private final ClientConfigurationData clientConfigData;
     private final Client client;
+    private final AsyncHttpClient httpAsyncClient;
     private final String serviceUrl;
     private final Lookup lookups;
     private final Functions functions;
@@ -146,11 +148,13 @@ public PulsarAdmin(String serviceUrl,
             auth.start();
         }
 
+        AsyncHttpConnectorProvider asyncConnectorProvider = new AsyncHttpConnectorProvider(clientConfigData);
+
         ClientConfig httpConfig = new ClientConfig();
         httpConfig.property(ClientProperties.FOLLOW_REDIRECTS, true);
         httpConfig.property(ClientProperties.ASYNC_THREADPOOL_SIZE, 8);
         httpConfig.register(MultiPartFeature.class);
-        httpConfig.connectorProvider(new AsyncHttpConnectorProvider(clientConfigData));
+        httpConfig.connectorProvider(asyncConnectorProvider);
 
         ClientBuilder clientBuilder = ClientBuilder.newBuilder()
                 .withConfig(httpConfig)
@@ -165,6 +169,10 @@ public PulsarAdmin(String serviceUrl,
         this.serviceUrl = serviceUrl;
         root = client.target(serviceUrl);
 
+        this.httpAsyncClient = asyncConnectorProvider.getConnector(
+                Math.toIntExact(TimeUnit.SECONDS.toMillis(this.connectTimeout)),
+                Math.toIntExact(TimeUnit.SECONDS.toMillis(this.readTimeout))).getHttpClient();
+
         this.clusters = new ClustersImpl(root, auth);
         this.brokers = new BrokersImpl(root, auth);
         this.brokerStats = new BrokerStatsImpl(root, auth);
@@ -175,9 +183,9 @@ public PulsarAdmin(String serviceUrl,
         this.nonPersistentTopics = new NonPersistentTopicsImpl(root, auth);
         this.resourceQuotas = new ResourceQuotasImpl(root, auth);
         this.lookups = new LookupImpl(root, auth, useTls);
-        this.functions = new FunctionsImpl(root, auth);
-        this.source = new SourceImpl(root, auth);
-        this.sink = new SinkImpl(root, auth);
+        this.functions = new FunctionsImpl(root, auth, httpAsyncClient);
+        this.source = new SourceImpl(root, auth, httpAsyncClient);
+        this.sink = new SinkImpl(root, auth, httpAsyncClient);
         this.worker = new WorkerImpl(root, auth);
         this.schemas = new SchemasImpl(root, auth);
         this.bookies = new BookiesImpl(root, auth);
@@ -387,6 +395,11 @@ public void close() {
             LOG.error("Failed to close the authentication service", e);
         }
         client.close();
-    }
 
+        try {
+            httpAsyncClient.close();
+        } catch (IOException e) {
+           LOG.error("Failed to close http async client", e);
+        }
+    }
 }
diff --git a/pulsar-client-admin/src/main/java/org/apache/pulsar/client/admin/PulsarAdminException.java b/pulsar-client-admin/src/main/java/org/apache/pulsar/client/admin/PulsarAdminException.java
@@ -23,26 +23,34 @@
 import javax.ws.rs.WebApplicationException;
 import javax.ws.rs.core.MediaType;
 
+import lombok.extern.slf4j.Slf4j;
 import org.apache.pulsar.common.policies.data.ErrorData;
+import org.apache.pulsar.common.util.ObjectMapperFactory;
 
 
 @SuppressWarnings("serial")
+@Slf4j
 public class PulsarAdminException extends Exception {
     private static final int DEFAULT_STATUS_CODE = 500;
 
     private final String httpError;
     private final int statusCode;
 
     private static String getReasonFromServer(WebApplicationException e) {
-        if (MediaType.APPLICATION_JSON.equals(e.getResponse().getHeaderString("Content-Type"))) {
+        try {
+            return e.getResponse().readEntity(ErrorData.class).reason;
+        } catch (Exception ex) {
             try {
-                return e.getResponse().readEntity(ErrorData.class).reason;
-            } catch (Exception ex) {
-                // could not parse output to ErrorData class
-                return e.getMessage();
+                return ObjectMapperFactory.getThreadLocal().readValue(e.getResponse().getEntity().toString(), ErrorData.class).reason;
+            } catch (Exception ex1) {
+                try {
+                    return ObjectMapperFactory.getThreadLocal().readValue(e.getMessage(), ErrorData.class).reason;
+                } catch (Exception ex2) {
+                    // could not parse output to ErrorData class
+                    return e.getMessage();
+                }
             }
         }
-        return e.getMessage();
     }
 
     public PulsarAdminException(ClientErrorException e) {

diff --git a/pulsar-client-admin/src/main/java/org/apache/pulsar/client/admin/internal/BaseResource.java b/pulsar-client-admin/src/main/java/org/apache/pulsar/client/admin/internal/BaseResource.java
@@ -37,7 +37,6 @@
 import org.apache.pulsar.client.admin.PulsarAdminException.ConflictException;
 import org.apache.pulsar.client.admin.PulsarAdminException.ConnectException;
 import org.apache.pulsar.client.admin.PulsarAdminException.GettingAuthenticationDataException;
-import org.apache.pulsar.client.admin.PulsarAdminException.HttpErrorException;
 import org.apache.pulsar.client.admin.PulsarAdminException.NotAllowedException;
 import org.apache.pulsar.client.admin.PulsarAdminException.NotAuthorizedException;
 import org.apache.pulsar.client.admin.PulsarAdminException.NotFoundException;
@@ -52,7 +51,7 @@
 public abstract class BaseResource {
     private static final Logger log = LoggerFactory.getLogger(BaseResource.class);
 
-    private final Authentication auth;
+    protected final Authentication auth;
 
     protected BaseResource(Authentication auth) {
         this.auth = auth;
@@ -200,5 +199,4 @@ public WebApplicationException getApiException(Response response) {
             throw new WebApplicationException(response);
         }
     }
-
 }
diff --git a/...client-admin/src/main/java/org/apache/pulsar/client/admin/internal/ComponentResource.java b/...client-admin/src/main/java/org/apache/pulsar/client/admin/internal/ComponentResource.java
@@ -0,0 +1,47 @@
+/**
+ * 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.pulsar.client.admin.internal;
+
+import org.apache.pulsar.client.admin.PulsarAdminException;
+import org.apache.pulsar.client.api.Authentication;
+import org.asynchttpclient.RequestBuilder;
+
+import java.util.Map;
+
+public class ComponentResource extends BaseResource {
+
+    protected ComponentResource(Authentication auth) {
+        super(auth);
+    }
+
+    public RequestBuilder addAuthHeaders(RequestBuilder requestBuilder) throws PulsarAdminException {
+
+        try {
+            if (auth != null && auth.getAuthData().hasDataForHttp()) {
+                for (Map.Entry<String, String> header : auth.getAuthData().getHttpHeaders()) {
+                    requestBuilder.addHeader(header.getKey(), header.getValue());
+                }
+            }
+
+            return requestBuilder;
+        } catch (Throwable t) {
+            throw new PulsarAdminException.GettingAuthenticationDataException(t);
+        }
+    }
+}
diff --git a/pulsar-client-admin/src/main/java/org/apache/pulsar/client/admin/internal/FunctionsImpl.java b/pulsar-client-admin/src/main/java/org/apache/pulsar/client/admin/internal/FunctionsImpl.java
@@ -20,6 +20,7 @@
 
 import com.google.gson.Gson;
 import com.google.gson.reflect.TypeToken;
+import io.netty.handler.codec.http.HttpHeaders;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.pulsar.client.admin.Functions;
@@ -32,6 +33,14 @@
 import org.apache.pulsar.common.policies.data.ErrorData;
 import org.apache.pulsar.common.policies.data.FunctionStats;
 import org.apache.pulsar.common.policies.data.FunctionStatus;
+import org.apache.pulsar.common.util.ObjectMapperFactory;
+import org.asynchttpclient.AsyncHandler;
+import org.asynchttpclient.AsyncHttpClient;
+import org.asynchttpclient.HttpResponseBodyPart;
+import org.asynchttpclient.HttpResponseStatus;
+import org.asynchttpclient.RequestBuilder;
+import org.asynchttpclient.request.body.multipart.FilePart;
+import org.asynchttpclient.request.body.multipart.StringPart;
 import org.glassfish.jersey.media.multipart.FormDataBodyPart;
 import org.glassfish.jersey.media.multipart.FormDataMultiPart;
 import org.glassfish.jersey.media.multipart.file.FileDataBodyPart;
@@ -42,21 +51,27 @@
 import javax.ws.rs.core.MediaType;
 import javax.ws.rs.core.Response;
 import java.io.File;
-import java.io.IOException;
-import java.io.InputStream;
-import java.nio.file.StandardCopyOption;
+import java.io.FileOutputStream;
+import java.nio.channels.FileChannel;
 import java.util.List;
 import java.util.Set;
+import java.util.concurrent.Future;
 import java.util.stream.Collectors;
 
+import static org.asynchttpclient.Dsl.get;
+import static org.asynchttpclient.Dsl.post;
+import static org.asynchttpclient.Dsl.put;
+
 @Slf4j
-public class FunctionsImpl extends BaseResource implements Functions {
+public class FunctionsImpl extends ComponentResource implements Functions {
 
     private final WebTarget functions;
+    private final AsyncHttpClient asyncHttpClient;
 
-    public FunctionsImpl(WebTarget web, Authentication auth) {
+    public FunctionsImpl(WebTarget web, Authentication auth, AsyncHttpClient asyncHttpClient) {
         super(auth);
         this.functions = web.path("/admin/v3/functions");
+        this.asyncHttpClient = asyncHttpClient;
     }
 
     @Override
@@ -145,18 +160,19 @@ public FunctionStats getFunctionStats(String tenant, String namespace, String fu
     @Override
     public void createFunction(FunctionConfig functionConfig, String fileName) throws PulsarAdminException {
         try {
-            final FormDataMultiPart mp = new FormDataMultiPart();
+            RequestBuilder builder = post(functions.path(functionConfig.getTenant()).path(functionConfig.getNamespace()).path(functionConfig.getName()).getUri().toASCIIString())
+                    .addBodyPart(new StringPart("functionConfig", ObjectMapperFactory.getThreadLocal().writeValueAsString(functionConfig), MediaType.APPLICATION_JSON));
 
             if (fileName != null && !fileName.startsWith("builtin://")) {
                 // If the function code is built in, we don't need to submit here
-                mp.bodyPart(new FileDataBodyPart("data", new File(fileName), MediaType.APPLICATION_OCTET_STREAM_TYPE));
+               builder.addBodyPart(new FilePart("data", new File(fileName), MediaType.APPLICATION_OCTET_STREAM));
+            }
+            org.asynchttpclient.Response response = asyncHttpClient.executeRequest(addAuthHeaders(builder).build()).get();
+
+            if (response.getStatusCode() < 200 || response.getStatusCode() >= 300) {
+                throw getApiException(Response.status(response.getStatusCode()).entity(response.getResponseBody()).build());
             }
 
-            mp.bodyPart(new FormDataBodyPart("functionConfig",
-                new Gson().toJson(functionConfig),
-                MediaType.APPLICATION_JSON_TYPE));
-            request(functions.path(functionConfig.getTenant()).path(functionConfig.getNamespace()).path(functionConfig.getName()))
-                    .post(Entity.entity(mp, MediaType.MULTIPART_FORM_DATA), ErrorData.class);
         } catch (Exception e) {
             throw getApiException(e);
         }
@@ -192,18 +208,18 @@ public void deleteFunction(String cluster, String namespace, String function) th
     @Override
     public void updateFunction(FunctionConfig functionConfig, String fileName) throws PulsarAdminException {
         try {
-            final FormDataMultiPart mp = new FormDataMultiPart();
+            RequestBuilder builder = put(functions.path(functionConfig.getTenant()).path(functionConfig.getNamespace()).path(functionConfig.getName()).getUri().toASCIIString())
+                    .addBodyPart(new StringPart("functionConfig", ObjectMapperFactory.getThreadLocal().writeValueAsString(functionConfig), MediaType.APPLICATION_JSON));
 
             if (fileName != null && !fileName.startsWith("builtin://")) {
                 // If the function code is built in, we don't need to submit here
-                mp.bodyPart(new FileDataBodyPart("data", new File(fileName), MediaType.APPLICATION_OCTET_STREAM_TYPE));
+                builder.addBodyPart(new FilePart("data", new File(fileName), MediaType.APPLICATION_OCTET_STREAM));
             }
+            org.asynchttpclient.Response response = asyncHttpClient.executeRequest(addAuthHeaders(builder).build()).get();
 
-            mp.bodyPart(new FormDataBodyPart("functionConfig",
-                new Gson().toJson(functionConfig),
-                MediaType.APPLICATION_JSON_TYPE));
-            request(functions.path(functionConfig.getTenant()).path(functionConfig.getNamespace()).path(functionConfig.getName()))
-                    .put(Entity.entity(mp, MediaType.MULTIPART_FORM_DATA), ErrorData.class);
+            if (response.getStatusCode() < 200 || response.getStatusCode() >= 300) {
+                throw getApiException(Response.status(response.getStatusCode()).entity(response.getResponseBody()).build());
+            }
         } catch (Exception e) {
             throw getApiException(e);
         }
@@ -314,29 +330,73 @@ public void startFunction(String tenant, String namespace, String functionName)
     @Override
     public void uploadFunction(String sourceFile, String path) throws PulsarAdminException {
         try {
-            final FormDataMultiPart mp = new FormDataMultiPart();
-
-            mp.bodyPart(new FileDataBodyPart("data", new File(sourceFile), MediaType.APPLICATION_OCTET_STREAM_TYPE));
+            RequestBuilder builder = post(functions.path("upload").getUri().toASCIIString())
+                    .addBodyPart(new FilePart("data", new File(sourceFile), MediaType.APPLICATION_OCTET_STREAM))
+                    .addBodyPart(new StringPart("path", path, MediaType.TEXT_PLAIN));
 
-            mp.bodyPart(new FormDataBodyPart("path", path, MediaType.TEXT_PLAIN_TYPE));
-            request(functions.path("upload"))
-                    .post(Entity.entity(mp, MediaType.MULTIPART_FORM_DATA), ErrorData.class);
+            org.asynchttpclient.Response response = asyncHttpClient.executeRequest(addAuthHeaders(builder).build()).get();
+            if (response.getStatusCode() < 200 || response.getStatusCode() >= 300) {
+                throw getApiException(Response.status(response.getStatusCode()).entity(response.getResponseBody()).build());
+            }
         } catch (Exception e) {
             throw getApiException(e);
         }
     }
 
     @Override
     public void downloadFunction(String destinationPath, String path) throws PulsarAdminException {
+
+        HttpResponseStatus status;
         try {
-            InputStream response = request(functions.path("download")
-                    .queryParam("path", path)).get(InputStream.class);
-            if (response != null) {
-                File targetFile = new File(destinationPath);
-                java.nio.file.Files.copy(
-                        response,
-                        targetFile.toPath(),
-                        StandardCopyOption.REPLACE_EXISTING);
+            File file = new File(destinationPath);
+            if (!file.exists()) {
+                file.createNewFile();
+            }
+            FileChannel os = new FileOutputStream(new File(destinationPath)).getChannel();
+            WebTarget target = functions.path("download").queryParam("path", path);
+
+            RequestBuilder builder = get(target.getUri().toASCIIString());
+
+            Future<HttpResponseStatus> whenStatusCode
+                    = asyncHttpClient.executeRequest(addAuthHeaders(builder).build(), new AsyncHandler<HttpResponseStatus>() {
+                private HttpResponseStatus status;
+
+                @Override
+                public State onStatusReceived(HttpResponseStatus responseStatus) throws Exception {
+                    status = responseStatus;
+                    if (status.getStatusCode() != Response.Status.OK.getStatusCode()) {
+                        return State.ABORT;
+                    }
+                    return State.CONTINUE;
+                }
+
+                @Override
+                public State onHeadersReceived(HttpHeaders headers) throws Exception {
+                    return State.CONTINUE;
+                }
+
+                @Override
+                public State onBodyPartReceived(HttpResponseBodyPart bodyPart) throws Exception {
+
+                    os.write(bodyPart.getBodyByteBuffer());
+                    return State.CONTINUE;
+                }
+
+                @Override
+                public HttpResponseStatus onCompleted() throws Exception {
+                    return status;
+                }
+
+                @Override
+                public void onThrowable(Throwable t) {
+                }
+            });
+
+            status = whenStatusCode.get();
+            os.close();
+
+            if (status.getStatusCode() < 200 || status.getStatusCode() >= 300) {
+                throw getApiException(Response.status(status.getStatusCode()).entity(status.getStatusText()).build());
             }
         } catch (Exception e) {
             throw getApiException(e);