Increase export throughput (#2204)

* Revert "Allow concurrent exports (#2181)"

This reverts commit b2c60e7.

* Pending exports rewrite

* Add log message

Author: trask

agent/agent-tooling/src/main/java/com/microsoft/applicationinsights/agent/internal/telemetry/BatchSpanProcessor.java

@@ -28,8 +28,10 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Queue;
+import java.util.Set;
 import java.util.concurrent.ArrayBlockingQueue;
 import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -62,6 +64,7 @@ public static BatchSpanProcessorBuilder builder(TelemetryChannel spanExporter) {
       int maxQueueSize,
       int maxExportBatchSize,
       long exporterTimeoutNanos,
+      int maxPendingExports,
       String queueName) {
     MpscArrayQueue<TelemetryItem> queue = new MpscArrayQueue<>(maxQueueSize);
     this.worker =
@@ -70,6 +73,7 @@ public static BatchSpanProcessorBuilder builder(TelemetryChannel spanExporter) {
             scheduleDelayNanos,
             maxExportBatchSize,
             exporterTimeoutNanos,
+            maxPendingExports,
             queue,
             queue.capacity(),
             queueName);
@@ -100,6 +104,7 @@ private static final class Worker implements Runnable {
     private final long scheduleDelayNanos;
     private final int maxExportBatchSize;
     private final long exporterTimeoutNanos;
+    private final int maxPendingExports;
 
     private long nextExportTime;
 
@@ -118,21 +123,29 @@ private static final class Worker implements Runnable {
     private volatile boolean continueWork = true;
     private final ArrayList<TelemetryItem> batch;
 
+    private final Set<CompletableResultCode> pendingExports =
+        Collections.newSetFromMap(new ConcurrentHashMap<>());
+
     private static final OperationLogger queuingSpanLogger =
         new OperationLogger(BatchSpanProcessor.class, "Queuing span");
 
+    private static final OperationLogger addAsyncExport =
+        new OperationLogger(BatchSpanProcessor.class, "Add async export");
+
     private Worker(
         TelemetryChannel spanExporter,
         long scheduleDelayNanos,
         int maxExportBatchSize,
         long exporterTimeoutNanos,
+        int maxPendingExports,
         Queue<TelemetryItem> queue,
         int queueCapacity,
         String queueName) {
       this.spanExporter = spanExporter;
       this.scheduleDelayNanos = scheduleDelayNanos;
       this.maxExportBatchSize = maxExportBatchSize;
       this.exporterTimeoutNanos = exporterTimeoutNanos;
+      this.maxPendingExports = maxPendingExports;
       this.queue = queue;
       this.queueCapacity = queueCapacity;
       this.queueName = queueName;
@@ -206,8 +219,12 @@ private void flush() {
         }
       }
       exportCurrentBatch();
-      flushRequested.get().succeed();
-      flushRequested.set(null);
+      CompletableResultCode.ofAll(pendingExports).join(exporterTimeoutNanos, TimeUnit.NANOSECONDS);
+      CompletableResultCode flushResult = flushRequested.get();
+      if (flushResult != null) {
+        flushResult.succeed();
+        flushRequested.set(null);
+      }
     }
 
     private void updateNextExportTime() {
@@ -232,28 +249,6 @@ private CompletableResultCode shutdown() {
     }
 
     private CompletableResultCode forceFlush() {
-      CompletableResultCode overallResult = new CompletableResultCode();
-      CompletableResultCode workerResult = forceFlushWorker();
-      workerResult.whenComplete(
-          () -> {
-            if (!workerResult.isSuccess()) {
-              overallResult.fail();
-              return;
-            }
-            CompletableResultCode exporterResult = spanExporter.flush();
-            exporterResult.whenComplete(
-                () -> {
-                  if (exporterResult.isSuccess()) {
-                    overallResult.succeed();
-                  } else {
-                    overallResult.fail();
-                  }
-                });
-          });
-      return overallResult;
-    }
-
-    private CompletableResultCode forceFlushWorker() {
       CompletableResultCode flushResult = new CompletableResultCode();
       // we set the atomic here to trigger the worker loop to do a flush of the entire queue.
       if (flushRequested.compareAndSet(null, flushResult)) {
@@ -274,7 +269,20 @@ private void exportCurrentBatch() {
       try {
         // batching, retry, logging, and writing to disk on failure occur downstream
         CompletableResultCode result = spanExporter.send(Collections.unmodifiableList(batch));
-        result.join(exporterTimeoutNanos, TimeUnit.NANOSECONDS);
+        if (pendingExports.size() < maxPendingExports - 1) {
+          addAsyncExport.recordSuccess();
+          pendingExports.add(result);
+          result.whenComplete(
+              () -> {
+                pendingExports.remove(result);
+              });
+        } else {
+          addAsyncExport.recordFailure(
+              "Max number of concurrent exports "
+                  + maxPendingExports
+                  + " has been hit, may see some export throttling due to this");
+          result.join(exporterTimeoutNanos, TimeUnit.NANOSECONDS);
+        }
       } finally {
         batch.clear();
       }

agent/agent-tooling/src/main/java/com/microsoft/applicationinsights/agent/internal/telemetry/BatchSpanProcessorBuilder.java

@@ -34,12 +34,14 @@ final class BatchSpanProcessorBuilder {
   private static final int DEFAULT_MAX_QUEUE_SIZE = 2048;
   private static final int DEFAULT_MAX_EXPORT_BATCH_SIZE = 512;
   private static final int DEFAULT_EXPORT_TIMEOUT_MILLIS = 30_000;
+  private static final int DEFAULT_MAX_PENDING_EXPORTS = 1;
 
   private final TelemetryChannel spanExporter;
   private long scheduleDelayNanos = TimeUnit.MILLISECONDS.toNanos(DEFAULT_SCHEDULE_DELAY_MILLIS);
   private int maxQueueSize = DEFAULT_MAX_QUEUE_SIZE;
   private int maxExportBatchSize = DEFAULT_MAX_EXPORT_BATCH_SIZE;
   private long exporterTimeoutNanos = TimeUnit.MILLISECONDS.toNanos(DEFAULT_EXPORT_TIMEOUT_MILLIS);
+  private int maxPendingExports = DEFAULT_MAX_PENDING_EXPORTS;
 
   BatchSpanProcessorBuilder(TelemetryChannel spanExporter) {
     this.spanExporter = requireNonNull(spanExporter, "spanExporter");
@@ -122,6 +124,31 @@ public BatchSpanProcessorBuilder setMaxExportBatchSize(int maxExportBatchSize) {
     return this;
   }
 
+  /**
+   * The maximum number of exports that can be pending at any time.
+   *
+   * <p>The {@link BatchSpanProcessor}'s single worker thread will keep processing as many batches
+   * as it can without blocking on the {@link io.opentelemetry.sdk.common.CompletableResultCode}s
+   * that are returned from the {@code spanExporter}, but it will limit the total number of pending
+   * exports in flight to this number.
+   *
+   * <p>Default value is {@code 1}.
+   *
+   * @param maxPendingExports the maximum number of exports that can be pending at any time.
+   * @return this.
+   * @see BatchSpanProcessorBuilder#DEFAULT_MAX_PENDING_EXPORTS
+   */
+  public BatchSpanProcessorBuilder setMaxPendingExports(int maxPendingExports) {
+    checkArgument(maxPendingExports > 0, "maxPendingExports must be positive.");
+    this.maxPendingExports = maxPendingExports;
+    return this;
+  }
+
+  // Visible for testing
+  int getMaxPendingExports() {
+    return maxPendingExports;
+  }
+
   /**
    * Returns a new {@link io.opentelemetry.sdk.trace.export.BatchSpanProcessor} that batches, then
    * converts spans to proto and forwards them to the given {@code spanExporter}.
@@ -136,6 +163,7 @@ public BatchSpanProcessor build(String queueName) {
         maxQueueSize,
         maxExportBatchSize,
         exporterTimeoutNanos,
+        maxPendingExports,
         queueName);
   }
 }