Onboard ingestion onto sync stack (Azure#33146)

* Onboard ingestion to sync stack

Author: Liudmila Molkova

sdk/monitor/azure-monitor-ingestion/README.md

@@ -128,7 +128,7 @@ System.out.println("Logs uploaded successfully");
 ### Upload custom logs with max concurrency
 
 If the in input logs collection is too large, the client will split the input into multiple smaller requests. These 
-requests are sent serially, by default, but by configuring the max concurrency in `UploadLogsOptions`, these requests
+requests are sent serially, by default, but by configuring the max concurrency in `LogsUploadOptions`, these requests
 can be concurrently sent to the service as shown in the example below.
 
 ```java readme-sample-uploadLogsWithMaxConcurrency

sdk/monitor/azure-monitor-ingestion/pom.xml

@@ -48,6 +48,7 @@
         <javaModulesSurefireArgLine>
             --add-exports com.azure.core/com.azure.core.implementation.jackson=ALL-UNNAMED
             --add-exports com.azure.core/com.azure.core.implementation.util=ALL-UNNAMED
+            --add-exports com.azure.monitor.ingestion/com.azure.monitor.ingestion.implementation=ALL-UNNAMED
             --add-opens com.azure.monitor.ingestion/com.azure.monitor.ingestion=ALL-UNNAMED
         </javaModulesSurefireArgLine>
     </properties>

sdk/monitor/azure-monitor-ingestion/src/main/java/com/azure/monitor/ingestion/LogsIngestionAsyncClient.java

@@ -15,37 +15,27 @@
 import com.azure.core.http.rest.Response;
 import com.azure.core.util.BinaryData;
 import com.azure.core.util.Context;
-import com.azure.core.util.logging.ClientLogger;
-import com.azure.core.util.serializer.JsonSerializer;
-import com.azure.core.util.serializer.JsonSerializerProviders;
-import com.azure.core.util.serializer.ObjectSerializer;
+import com.azure.monitor.ingestion.implementation.Batcher;
 import com.azure.monitor.ingestion.implementation.IngestionUsingDataCollectionRulesAsyncClient;
 import com.azure.monitor.ingestion.implementation.LogsIngestionRequest;
 import com.azure.monitor.ingestion.implementation.UploadLogsResponseHolder;
 import com.azure.monitor.ingestion.models.LogsUploadError;
 import com.azure.monitor.ingestion.models.LogsUploadException;
 import com.azure.monitor.ingestion.models.LogsUploadOptions;
-import com.fasterxml.jackson.core.JsonFactory;
-import com.fasterxml.jackson.core.JsonGenerator;
-import reactor.core.publisher.Flux;
-import reactor.core.publisher.FluxSink;
 import reactor.core.publisher.Mono;
 import reactor.core.publisher.SynchronousSink;
 
-import java.io.ByteArrayOutputStream;
-import java.io.IOException;
-import java.io.UncheckedIOException;
-import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.Collections;
-import java.util.Iterator;
 import java.util.List;
 import java.util.Objects;
 import java.util.function.Consumer;
-import java.util.stream.Collectors;
-import java.util.zip.GZIPOutputStream;
 
 import static com.azure.core.util.FluxUtil.withContext;
+import static com.azure.monitor.ingestion.implementation.Utils.CONTENT_ENCODING;
+import static com.azure.monitor.ingestion.implementation.Utils.GZIP;
+import static com.azure.monitor.ingestion.implementation.Utils.getConcurrency;
+import static com.azure.monitor.ingestion.implementation.Utils.gzipRequest;
 
 /**
  * The asynchronous client for uploading logs to Azure Monitor.
@@ -62,12 +52,6 @@
  */
 @ServiceClient(isAsync = true, builder = LogsIngestionClientBuilder.class)
 public final class LogsIngestionAsyncClient {
-    private static final ClientLogger LOGGER = new ClientLogger(LogsIngestionAsyncClient.class);
-    private static final String CONTENT_ENCODING = "Content-Encoding";
-    private static final long MAX_REQUEST_PAYLOAD_SIZE = 1024 * 1024; // 1 MB
-    private static final String GZIP = "gzip";
-    private static final JsonSerializer DEFAULT_SERIALIZER = JsonSerializerProviders.createInstance(true);
-
     private final IngestionUsingDataCollectionRulesAsyncClient service;
 
     LogsIngestionAsyncClient(IngestionUsingDataCollectionRulesAsyncClient service) {
@@ -201,71 +185,15 @@ Mono<Void> upload(String ruleId, String streamName,
      */
     private Mono<Void> splitAndUpload(String ruleId, String streamName, Iterable<Object> logs,
                                       LogsUploadOptions options, Context context) {
-        // set concurrency to 1 as default
-        int concurrency = 1;
-        ObjectSerializer objectSerializer = DEFAULT_SERIALIZER;
-        if (options != null) {
-            if (options.getObjectSerializer() != null) {
-                objectSerializer = options.getObjectSerializer();
-            }
-            if (options.getMaxConcurrency() != null) {
-                concurrency = options.getMaxConcurrency();
-            }
-        }
 
-        ObjectSerializer serializer = objectSerializer;
-        Iterator<Object> iterator = logs.iterator();
+        int concurrency = getConcurrency(options);
 
-        return Flux.<LogsIngestionRequest>create(emitter -> createHttpRequest(serializer, iterator, emitter))
-                .flatMapSequential(request -> uploadToService(ruleId, streamName, context, request), concurrency)
-                .<LogsUploadException>handle((responseHolder, sink) -> processResponse(options, responseHolder, sink))
-                .collectList()
-                .handle((result, sink) -> processExceptions(result, sink));
-    }
-
-    /**
-     * This method converts raw {@link Object} logs into serialized and gzipped byte array that forms the request
-     * body of the service call.
-     * @param serializer The serializer used to serialize logs.
-     * @param iterator The input logs interator.
-     * @param emitter The Flux emitter to which the HTTP request is emitted.
-     */
-    private void createHttpRequest(ObjectSerializer serializer, Iterator<Object> iterator, FluxSink<LogsIngestionRequest> emitter) {
-        try {
-            long currentBatchSize = 0;
-            ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
-            JsonGenerator generator = JsonFactory.builder().build().createGenerator(byteArrayOutputStream);
-            generator.writeStartArray();
-            List<String> serializedLogs = new ArrayList<>();
-            List<Object> originalLogsRequest = new ArrayList<>();
-            while (iterator.hasNext()) {
-                Object currentLog = iterator.next();
-                byte[] bytes = serializer.serializeToBytes(currentLog);
-                int currentLogSize = bytes.length;
-                currentBatchSize += currentLogSize;
-                if (currentBatchSize > MAX_REQUEST_PAYLOAD_SIZE) {
-                    writeLogsAndCloseJsonGenerator(generator, serializedLogs);
-                    byte[] zippedRequestBody = gzipRequest(byteArrayOutputStream.toByteArray());
-                    emitter.next(new LogsIngestionRequest(originalLogsRequest, zippedRequestBody));
-                    byteArrayOutputStream = new ByteArrayOutputStream();
-                    generator = JsonFactory.builder().build().createGenerator(byteArrayOutputStream);
-                    generator.writeStartArray();
-                    currentBatchSize = currentLogSize;
-                    originalLogsRequest = new ArrayList<>();
-                    serializedLogs.clear();
-                }
-                serializedLogs.add(new String(bytes, StandardCharsets.UTF_8));
-                originalLogsRequest.add(currentLog);
-            }
-            if (currentBatchSize > 0) {
-                writeLogsAndCloseJsonGenerator(generator, serializedLogs);
-                byte[] zippedRequestBody = gzipRequest(byteArrayOutputStream.toByteArray());
-                emitter.next(new LogsIngestionRequest(originalLogsRequest, zippedRequestBody));
-            }
-            emitter.complete();
-        } catch (IOException e) {
-            emitter.error(e);
-        }
+        return new Batcher(options, logs)
+            .toFlux()
+            .flatMapSequential(request -> uploadToService(ruleId, streamName, context, request), concurrency)
+            .<LogsUploadException>handle((responseHolder, sink) -> processResponse(options, responseHolder, sink))
+            .collectList()
+            .handle((result, sink) -> processExceptions(result, sink));
     }
 
     private void processExceptions(List<LogsUploadException> result, SynchronousSink<Void> sink) {
@@ -311,28 +239,4 @@ private Mono<UploadLogsResponseHolder> uploadToService(String ruleId, String str
                 .onErrorResume(HttpResponseException.class,
                         ex -> Mono.fromSupplier(() -> new UploadLogsResponseHolder(request, ex)));
     }
-
-    private void writeLogsAndCloseJsonGenerator(JsonGenerator generator, List<String> serializedLogs) throws IOException {
-        generator.writeRaw(serializedLogs.stream()
-                .collect(Collectors.joining(",")));
-        generator.writeEndArray();
-        generator.close();
-    }
-
-    /**
-     * Gzips the input byte array.
-     * @param bytes The input byte array.
-     * @return gzipped byte array.
-     */
-    private byte[] gzipRequest(byte[] bytes) {
-        // This should be moved to azure-core and should be enabled when the client library requests for gzipping the
-        // request body content.
-        ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
-        try (GZIPOutputStream zip = new GZIPOutputStream(byteArrayOutputStream)) {
-            zip.write(bytes);
-        } catch (IOException exception) {
-            throw LOGGER.logExceptionAsError(new UncheckedIOException(exception));
-        }
-        return byteArrayOutputStream.toByteArray();
-    }
 }

sdk/monitor/azure-monitor-ingestion/src/main/java/com/azure/monitor/ingestion/LogsIngestionClient.java

@@ -10,12 +10,35 @@
 import com.azure.core.exception.HttpResponseException;
 import com.azure.core.exception.ResourceModifiedException;
 import com.azure.core.exception.ResourceNotFoundException;
+import com.azure.core.http.HttpHeader;
 import com.azure.core.http.rest.RequestOptions;
 import com.azure.core.http.rest.Response;
 import com.azure.core.util.BinaryData;
 import com.azure.core.util.Context;
+import com.azure.core.util.logging.ClientLogger;
+import com.azure.monitor.ingestion.implementation.Batcher;
+import com.azure.monitor.ingestion.implementation.IngestionUsingDataCollectionRulesClient;
+import com.azure.monitor.ingestion.implementation.LogsIngestionRequest;
+import com.azure.monitor.ingestion.implementation.UploadLogsResponseHolder;
+import com.azure.monitor.ingestion.models.LogsUploadError;
+import com.azure.monitor.ingestion.models.LogsUploadException;
 import com.azure.monitor.ingestion.models.LogsUploadOptions;
 
+import java.util.List;
+import java.util.Objects;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.function.Consumer;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+import static com.azure.monitor.ingestion.implementation.Utils.CONTENT_ENCODING;
+import static com.azure.monitor.ingestion.implementation.Utils.GZIP;
+import static com.azure.monitor.ingestion.implementation.Utils.createThreadPool;
+import static com.azure.monitor.ingestion.implementation.Utils.getConcurrency;
+import static com.azure.monitor.ingestion.implementation.Utils.gzipRequest;
+import static com.azure.monitor.ingestion.implementation.Utils.registerShutdownHook;
+
 /**
  * The synchronous client for uploading logs to Azure Monitor.
  *
@@ -30,12 +53,20 @@
  * <!-- end com.azure.monitor.ingestion.LogsIngestionClient.instantiation -->
  */
 @ServiceClient(builder = LogsIngestionClientBuilder.class)
-public final class LogsIngestionClient {
+public final class LogsIngestionClient implements AutoCloseable {
+    private static final ClientLogger LOGGER = new ClientLogger(LogsIngestionClient.class);
+    private static final String HTTP_REST_PROXY_SYNC_PROXY_ENABLE = "com.azure.core.http.restproxy.syncproxy.enable";
+    private static final Context ENABLE_SYNC_CONTEXT = new Context(HTTP_REST_PROXY_SYNC_PROXY_ENABLE, true);
+    private final IngestionUsingDataCollectionRulesClient client;
 
-    private final LogsIngestionAsyncClient asyncClient;
+    // dynamic thread pool that scales up and down on demand.
+    private final ExecutorService threadPool;
+    private final Thread shutdownHook;
 
-    LogsIngestionClient(LogsIngestionAsyncClient asyncClient) {
-        this.asyncClient = asyncClient;
+    LogsIngestionClient(IngestionUsingDataCollectionRulesClient client) {
+        this.client = client;
+        this.threadPool = createThreadPool();
+        this.shutdownHook = registerShutdownHook(this.threadPool, 5);
     }
 
     /**
@@ -61,7 +92,7 @@ public final class LogsIngestionClient {
      */
     @ServiceMethod(returns = ReturnType.SINGLE)
     public void upload(String ruleId, String streamName, Iterable<Object> logs) {
-        asyncClient.upload(ruleId, streamName, logs).block();
+        upload(ruleId, streamName, logs, null);
     }
 
     /**
@@ -90,7 +121,7 @@ public void upload(String ruleId, String streamName, Iterable<Object> logs) {
     @ServiceMethod(returns = ReturnType.SINGLE)
     public void upload(String ruleId, String streamName,
                                    Iterable<Object> logs, LogsUploadOptions options) {
-        asyncClient.upload(ruleId, streamName, logs, options, Context.NONE).block();
+        upload(ruleId, streamName, logs, options, Context.NONE);
     }
 
     /**
@@ -111,7 +142,64 @@ public void upload(String ruleId, String streamName,
     @ServiceMethod(returns = ReturnType.SINGLE)
     public void upload(String ruleId, String streamName,
                        Iterable<Object> logs, LogsUploadOptions options, Context context) {
-        asyncClient.upload(ruleId, streamName, logs, options, context).block();
+        Objects.requireNonNull(ruleId, "'ruleId' cannot be null.");
+        Objects.requireNonNull(streamName, "'streamName' cannot be null.");
+        Objects.requireNonNull(logs, "'logs' cannot be null.");
+
+        context = enableSync(context);
+
+        Consumer<LogsUploadError> uploadLogsErrorConsumer = options == null ? null : options.getLogsUploadErrorConsumer();
+
+        RequestOptions requestOptions = new RequestOptions();
+        requestOptions.addHeader(CONTENT_ENCODING, GZIP);
+        requestOptions.setContext(context);
+
+        Stream<UploadLogsResponseHolder> responses = new Batcher(options, logs)
+            .toStream()
+            .map(r -> uploadToService(ruleId, streamName, requestOptions, r));
+
+        responses = submit(responses, getConcurrency(options))
+            .filter(response -> response.getException() != null);
+
+        if (uploadLogsErrorConsumer != null) {
+            responses.forEach(response -> uploadLogsErrorConsumer.accept(new LogsUploadError(response.getException(), response.getRequest().getLogs())));
+            return;
+        }
+
+        final int[] failedLogCount = new int[1];
+        List<HttpResponseException> exceptions = responses
+            .map(response -> {
+                failedLogCount[0] += response.getRequest().getLogs().size();
+                return response.getException();
+            })
+            .collect(Collectors.toList());
+
+        if (exceptions.size() > 0) {
+            throw LOGGER.logExceptionAsError(new LogsUploadException(exceptions, failedLogCount[0]));
+        }
+    }
+
+    private Stream<UploadLogsResponseHolder> submit(Stream<UploadLogsResponseHolder> responseStream, int concurrency) {
+        if (concurrency == 1) {
+            return responseStream;
+        }
+
+        try {
+            return threadPool.submit(() -> responseStream).get();
+        } catch (InterruptedException | ExecutionException e) {
+            throw LOGGER.logExceptionAsError(new RuntimeException(e));
+        }
+    }
+
+    private UploadLogsResponseHolder uploadToService(String ruleId, String streamName, RequestOptions requestOptions, LogsIngestionRequest request) {
+        HttpResponseException exception = null;
+        try {
+            client.uploadWithResponse(ruleId, streamName, BinaryData.fromBytes(request.getRequestBody()), requestOptions);
+        } catch (HttpResponseException ex) {
+            exception = ex;
+        }
+
+        return new UploadLogsResponseHolder(request, exception);
     }
 
     /**
@@ -147,6 +235,37 @@ public void upload(String ruleId, String streamName,
     @ServiceMethod(returns = ReturnType.SINGLE)
     public Response<Void> uploadWithResponse(
             String ruleId, String streamName, BinaryData logs, RequestOptions requestOptions) {
-        return asyncClient.uploadWithResponse(ruleId, streamName, logs, requestOptions).block();
+        Objects.requireNonNull(ruleId, "'ruleId' cannot be null.");
+        Objects.requireNonNull(streamName, "'streamName' cannot be null.");
+        Objects.requireNonNull(logs, "'logs' cannot be null.");
+
+        if (requestOptions == null) {
+            requestOptions = new RequestOptions();
+        }
+
+        requestOptions.setContext(enableSync(requestOptions.getContext()));
+        requestOptions.addRequestCallback(request -> {
+            HttpHeader httpHeader = request.getHeaders().get(CONTENT_ENCODING);
+            if (httpHeader == null) {
+                BinaryData gzippedRequest = BinaryData.fromBytes(gzipRequest(logs.toBytes()));
+                request.setBody(gzippedRequest);
+                request.setHeader(CONTENT_ENCODING, GZIP);
+            }
+        });
+        return client.uploadWithResponse(ruleId, streamName, logs, requestOptions);
+    }
+
+    private static Context enableSync(Context context) {
+        if (context == null || context == Context.NONE) {
+            return ENABLE_SYNC_CONTEXT;
+        }
+
+        return context.addData(HTTP_REST_PROXY_SYNC_PROXY_ENABLE, true);
+    }
+
+    @Override
+    public void close() {
+        threadPool.shutdown();
+        Runtime.getRuntime().removeShutdownHook(shutdownHook);
     }
 }

sdk/monitor/azure-monitor-ingestion/src/main/java/com/azure/monitor/ingestion/LogsIngestionClientBuilder.java

@@ -189,7 +189,13 @@ public LogsIngestionClientBuilder serviceVersion(LogsIngestionServiceVersion ser
      * @return A synchronous {@link LogsIngestionClient}.
      */
     public LogsIngestionClient buildClient() {
-        return new LogsIngestionClient(buildAsyncClient());
+        if (endpoint == null) {
+            throw LOGGER.logExceptionAsError(new IllegalStateException("endpoint is required to build the client."));
+        }
+        if (tokenCredential == null) {
+            throw LOGGER.logExceptionAsError(new IllegalStateException("credential is required to build the client."));
+        }
+        return new LogsIngestionClient(innerLogBuilder.buildClient());
     }
 
     /**