Split SymDB payload when too large (#7838)

jpbempel · web-flow · commit b2332408b074 · 2024-10-28T18:19:03.000+01:00
SymDB payload cannot be larger than 50MB otherwise datadog agent
will generate an error. When serializing the upload request we are now
verifying that we are below this limit, otherwise we are splitting
the payload first by uploading by jar scope, but if a jar scope is
still large we are splitting by class scopes. If down to one class
the payload is still too large we are dropping the upload.
diff --git a/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/sink/SymbolSink.java b/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/sink/SymbolSink.java
@@ -1,6 +1,7 @@
 package com.datadog.debugger.sink;
 
 import com.datadog.debugger.symbol.Scope;
+import com.datadog.debugger.symbol.ScopeType;
 import com.datadog.debugger.symbol.ServiceVersion;
 import com.datadog.debugger.uploader.BatchUploader;
 import com.datadog.debugger.util.MoshiHelper;
@@ -9,6 +10,8 @@
 import datadog.trace.util.TagsHelper;
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.ArrayBlockingQueue;
 import java.util.concurrent.BlockingQueue;
@@ -29,24 +32,30 @@ public class SymbolSink {
           + "\"service\": \"%s\",%n"
           + "\"runtimeId\": \"%s\"%n"
           + "}";
+  static final int MAX_SYMDB_UPLOAD_SIZE = 50 * 1024 * 1024;
 
   private final String serviceName;
   private final String env;
   private final String version;
   private final BatchUploader symbolUploader;
+  private final int maxPayloadSize;
   private final BatchUploader.MultiPartContent event;
   private final BlockingQueue<Scope> scopes = new ArrayBlockingQueue<>(CAPACITY);
   private final Stats stats = new Stats();
 
   public SymbolSink(Config config) {
-    this(config, new BatchUploader(config, config.getFinalDebuggerSymDBUrl(), RETRY_POLICY));
+    this(
+        config,
+        new BatchUploader(config, config.getFinalDebuggerSymDBUrl(), RETRY_POLICY),
+        MAX_SYMDB_UPLOAD_SIZE);
   }
 
-  SymbolSink(Config config, BatchUploader symbolUploader) {
+  SymbolSink(Config config, BatchUploader symbolUploader, int maxPayloadSize) {
     this.serviceName = TagsHelper.sanitize(config.getServiceName());
     this.env = TagsHelper.sanitize(config.getEnv());
     this.version = TagsHelper.sanitize(config.getVersion());
     this.symbolUploader = symbolUploader;
+    this.maxPayloadSize = maxPayloadSize;
     byte[] eventContent =
         String.format(
                 EVENT_FORMAT, TagsHelper.sanitize(config.getServiceName()), config.getRuntimeId())
@@ -87,15 +96,119 @@ public void flush() {
     String json =
         SERVICE_VERSION_ADAPTER.toJson(
             new ServiceVersion(serviceName, env, version, "JAVA", scopesToSerialize));
-    LOGGER.debug("Sending {} jar scopes size={}", scopesToSerialize.size(), json.length());
-    updateStats(scopesToSerialize, json);
+    if (json.length() > maxPayloadSize) {
+      LOGGER.debug(
+          "Upload split is required for {} scopes: {}/{}",
+          scopesToSerialize.size(),
+          json.length(),
+          maxPayloadSize);
+      splitAndSend(scopesToSerialize);
+    } else {
+      LOGGER.debug("Sending {} jar scopes size={}", scopesToSerialize.size(), json.length());
+      updateStats(scopesToSerialize, json);
+      LOGGER.debug("SymbolSink stats: {}", stats);
+      symbolUploader.uploadAsMultipart(
+          "",
+          event,
+          new BatchUploader.MultiPartContent(
+              json.getBytes(StandardCharsets.UTF_8), "file", "file.json"));
+    }
+  }
+
+  /*
+   * Try to split recursively the scopes to send in smaller chunks
+   * first try by splitting the jar scopes, then by splitting the class scopes
+   * stopped the recursion when the scopes are small enough to be sent or <2 classes in one scope
+   */
+  private void splitAndSend(List<Scope> scopesToSerialize) {
+    if (scopesToSerialize.size() > 1) {
+      // try to split by jar scopes: one scope per request
+      if (scopesToSerialize.size() < BatchUploader.MAX_ENQUEUED_REQUESTS) {
+        for (Scope scope : scopesToSerialize) {
+          String json =
+              SERVICE_VERSION_ADAPTER.toJson(
+                  new ServiceVersion(
+                      serviceName, env, version, "JAVA", Collections.singletonList(scope)));
+          if (json.length() > maxPayloadSize) {
+            // this jar scope is still too big, split it by classes
+            LOGGER.debug(
+                "Upload split is required for jar scope {}: {}/{}",
+                scope.getName(),
+                json.length(),
+                maxPayloadSize);
+            splitAndSend(Collections.singletonList(scope));
+            continue;
+          }
+          LOGGER.debug("Sending {} jar scope size={}", scope.getName(), json.length());
+          doUpload(Collections.singletonList(scope), json);
+        }
+      } else {
+        // split the list of jar scope in 2 list jar scopes with half of the jar scopes
+        int half = scopesToSerialize.size() / 2;
+        List<Scope> firstHalf = scopesToSerialize.subList(0, half);
+        List<Scope> secondHalf = scopesToSerialize.subList(half, scopesToSerialize.size());
+        LOGGER.debug("split jar scope list in 2: {} and {}", firstHalf.size(), secondHalf.size());
+        String jsonFirstHalf =
+            SERVICE_VERSION_ADAPTER.toJson(
+                new ServiceVersion(serviceName, env, version, "JAVA", firstHalf));
+        if (jsonFirstHalf.length() > maxPayloadSize) {
+          LOGGER.warn(
+              "Cannot split jar scope list in 2, first half is too big: {}",
+              jsonFirstHalf.length());
+          return;
+        }
+        doUpload(firstHalf, jsonFirstHalf);
+        String jsonSecondHalf =
+            SERVICE_VERSION_ADAPTER.toJson(
+                new ServiceVersion(serviceName, env, version, "JAVA", secondHalf));
+        if (jsonSecondHalf.length() > maxPayloadSize) {
+          LOGGER.warn(
+              "Cannot split jar scope list in 2, second half is too big: {}",
+              jsonSecondHalf.length());
+          return;
+        }
+        doUpload(secondHalf, jsonSecondHalf);
+      }
+    } else {
+      Scope jarScope = scopesToSerialize.get(0);
+      if (jarScope.getScopes() == null) {
+        LOGGER.debug("No class scopes to send for jar scope {}", jarScope.getName());
+        return;
+      }
+      if (jarScope.getScopes().size() < 2) {
+        LOGGER.warn(
+            "Cannot split jar scope with less than 2 classes scope: {}", jarScope.getName());
+        return;
+      }
+      // split the jar scope in 2 jar scopes with half of the class scopes
+      int half = jarScope.getScopes().size() / 2;
+      List<Scope> firstHalf = jarScope.getScopes().subList(0, half);
+      List<Scope> secondHalf = jarScope.getScopes().subList(half, jarScope.getScopes().size());
+      LOGGER.debug(
+          "split jar scope {} in 2 jar scopes: {} and {}",
+          jarScope.getName(),
+          firstHalf.size(),
+          secondHalf.size());
+      splitAndSend(
+          Arrays.asList(
+              createJarScope(jarScope.getName(), firstHalf),
+              createJarScope(jarScope.getName(), secondHalf)));
+    }
+  }
+
+  private void doUpload(List<Scope> scopes, String json) {
+    updateStats(scopes, json);
     symbolUploader.uploadAsMultipart(
         "",
         event,
         new BatchUploader.MultiPartContent(
             json.getBytes(StandardCharsets.UTF_8), "file", "file.json"));
   }
 
+  private static Scope createJarScope(String jarName, List<Scope> classScopes) {
+    return Scope.builder(ScopeType.JAR, jarName, 0, 0).name(jarName).scopes(classScopes).build();
+  }
+
   private void updateStats(List<Scope> scopesToSerialize, String json) {
     for (Scope scope : scopesToSerialize) {
       stats.updateStats(scope.getScopes() != null ? scope.getScopes().size() : 0, json.length());
diff --git a/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/uploader/BatchUploader.java b/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/uploader/BatchUploader.java
@@ -74,7 +74,7 @@ public RetryPolicy(int maxFailures) {
   private static final String HEADER_DD_ENTITY_ID = "Datadog-Entity-ID";
   static final String HEADER_DD_API_KEY = "DD-API-KEY";
   static final int MAX_RUNNING_REQUESTS = 10;
-  static final int MAX_ENQUEUED_REQUESTS = 20;
+  public static final int MAX_ENQUEUED_REQUESTS = 20;
   static final int TERMINATION_TIMEOUT = 5;
 
   private final String containerId;
diff --git a/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/sink/SymbolSinkTest.java b/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/sink/SymbolSinkTest.java
@@ -1,5 +1,7 @@
 package com.datadog.debugger.sink;
 
+import static com.datadog.debugger.sink.SymbolSink.MAX_SYMDB_UPLOAD_SIZE;
+import static java.util.Collections.singletonList;
 import static org.junit.jupiter.api.Assertions.*;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
@@ -10,7 +12,9 @@
 import datadog.trace.api.Config;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
+import org.jetbrains.annotations.NotNull;
 import org.junit.jupiter.api.Test;
 
 class SymbolSinkTest {
@@ -20,7 +24,7 @@ public void testSimpleFlush() {
     SymbolUploaderMock symbolUploaderMock = new SymbolUploaderMock();
     Config config = mock(Config.class);
     when(config.getServiceName()).thenReturn("service1");
-    SymbolSink symbolSink = new SymbolSink(config, symbolUploaderMock);
+    SymbolSink symbolSink = new SymbolSink(config, symbolUploaderMock, MAX_SYMDB_UPLOAD_SIZE);
     symbolSink.addScope(Scope.builder(ScopeType.JAR, null, 0, 0).build());
     symbolSink.flush();
     assertEquals(2, symbolUploaderMock.multiPartContents.size());
@@ -43,17 +47,13 @@ public void testMultiScopeFlush() {
     SymbolUploaderMock symbolUploaderMock = new SymbolUploaderMock();
     Config config = mock(Config.class);
     when(config.getServiceName()).thenReturn("service1");
-    SymbolSink symbolSink = new SymbolSink(config, symbolUploaderMock);
+    SymbolSink symbolSink = new SymbolSink(config, symbolUploaderMock, MAX_SYMDB_UPLOAD_SIZE);
     symbolSink.addScope(Scope.builder(ScopeType.JAR, "jar1.jar", 0, 0).build());
     symbolSink.addScope(Scope.builder(ScopeType.JAR, "jar2.jar", 0, 0).build());
     symbolSink.flush();
     // only 1 request because we are batching the scopes
     assertEquals(2, symbolUploaderMock.multiPartContents.size());
-    BatchUploader.MultiPartContent eventContent = symbolUploaderMock.multiPartContents.get(0);
-    assertEquals("event", eventContent.getPartName());
-    BatchUploader.MultiPartContent symbolContent = symbolUploaderMock.multiPartContents.get(1);
-    assertEquals("file", symbolContent.getPartName());
-    String strContent = new String(symbolContent.getContent());
+    String strContent = assertMultipartContent(symbolUploaderMock, 0);
     assertTrue(strContent.contains("\"source_file\":\"jar1.jar\""));
     assertTrue(strContent.contains("\"source_file\":\"jar2.jar\""));
   }
@@ -63,7 +63,7 @@ public void testQueueFull() {
     SymbolUploaderMock symbolUploaderMock = new SymbolUploaderMock();
     Config config = mock(Config.class);
     when(config.getServiceName()).thenReturn("service1");
-    SymbolSink symbolSink = new SymbolSink(config, symbolUploaderMock);
+    SymbolSink symbolSink = new SymbolSink(config, symbolUploaderMock, MAX_SYMDB_UPLOAD_SIZE);
     for (int i = 0; i < SymbolSink.CAPACITY; i++) {
       symbolSink.addScope(Scope.builder(ScopeType.JAR, "jar1.jar", 0, 0).build());
     }
@@ -81,6 +81,131 @@ public void testQueueFull() {
             .contains("\"source_file\":\"jar2.jar\""));
   }
 
+  @Test
+  public void splitByJarScopes() {
+    SymbolUploaderMock symbolUploaderMock = new SymbolUploaderMock();
+    Config config = mock(Config.class);
+    when(config.getServiceName()).thenReturn("service1");
+    SymbolSink symbolSink = new SymbolSink(config, symbolUploaderMock, 1024);
+    final int NUM_JAR_SCOPES = 10;
+    for (int i = 0; i < NUM_JAR_SCOPES; i++) {
+      symbolSink.addScope(
+          Scope.builder(ScopeType.JAR, "jar" + i + ".jar", 0, 0)
+              .scopes(singletonList(Scope.builder(ScopeType.CLASS, "class" + i, 0, 0).build()))
+              .build());
+    }
+    symbolSink.flush();
+    // split upload request per jar scope
+    assertEquals(NUM_JAR_SCOPES * 2, symbolUploaderMock.multiPartContents.size());
+    for (int i = 0; i < NUM_JAR_SCOPES * 2; i += 2) {
+      String strContent = assertMultipartContent(symbolUploaderMock, i);
+      assertTrue(strContent.contains("\"source_file\":\"jar" + (i / 2) + ".jar\""));
+    }
+  }
+
+  @Test
+  public void splitTootManyJarScopes() {
+    SymbolUploaderMock symbolUploaderMock = new SymbolUploaderMock();
+    Config config = mock(Config.class);
+    when(config.getServiceName()).thenReturn("service1");
+    SymbolSink symbolSink = new SymbolSink(config, symbolUploaderMock, 2048);
+    final int NUM_JAR_SCOPES = 21;
+    for (int i = 0; i < NUM_JAR_SCOPES; i++) {
+      symbolSink.addScope(
+          Scope.builder(ScopeType.JAR, "jar" + i + ".jar", 0, 0)
+              .scopes(singletonList(Scope.builder(ScopeType.CLASS, "class" + i, 0, 0).build()))
+              .build());
+    }
+    symbolSink.flush();
+    // split upload request per half jar scopes
+    assertEquals(2 * 2, symbolUploaderMock.multiPartContents.size());
+    String strContent1 = assertMultipartContent(symbolUploaderMock, 0);
+    assertTrue(strContent1.contains("\"source_file\":\"jar0.jar\""));
+    assertTrue(strContent1.contains("\"source_file\":\"jar9.jar\""));
+    String strContent2 = assertMultipartContent(symbolUploaderMock, 2);
+    assertTrue(strContent2.contains("\"source_file\":\"jar10.jar\""));
+    assertTrue(strContent2.contains("\"source_file\":\"jar20.jar\""));
+  }
+
+  @Test
+  public void splitByClassScopes() {
+    SymbolUploaderMock symbolUploaderMock = new SymbolUploaderMock();
+    Config config = mock(Config.class);
+    when(config.getServiceName()).thenReturn("service1");
+    SymbolSink symbolSink = new SymbolSink(config, symbolUploaderMock, 1024);
+    final int NUM_CLASS_SCOPES = 10;
+    List<Scope> classScopes = new ArrayList<>();
+    for (int i = 0; i < NUM_CLASS_SCOPES; i++) {
+      classScopes.add(
+          Scope.builder(ScopeType.CLASS, "class" + i, 0, 0)
+              .scopes(
+                  Collections.singletonList(
+                      Scope.builder(ScopeType.METHOD, "class" + i, 0, 0)
+                          .name("method" + i)
+                          .build()))
+              .build());
+    }
+    symbolSink.addScope(
+        Scope.builder(ScopeType.JAR, "jar1.jar", 0, 0)
+            .name("jar1.jar")
+            .scopes(classScopes)
+            .build());
+    symbolSink.flush();
+    // split upload request per jar scope
+    final int EXPECTED_REQUESTS = 4;
+    assertEquals(EXPECTED_REQUESTS * 2, symbolUploaderMock.multiPartContents.size());
+    List<List<String>> expectedSourceFiles =
+        Arrays.asList(
+            Arrays.asList("class0", "class1"),
+            Arrays.asList("class2", "class3", "class4"),
+            Arrays.asList("class5", "class6"),
+            Arrays.asList("class7", "class8", "class9"));
+    for (int i = 0; i < EXPECTED_REQUESTS * 2; i += 2) {
+      String strContent = assertMultipartContent(symbolUploaderMock, i);
+      for (String sourceFile : expectedSourceFiles.get(i / 2)) {
+        assertTrue(strContent.contains("\"source_file\":\"" + sourceFile + "\""));
+      }
+    }
+  }
+
+  @Test
+  public void splitByClassScopesImpossible() {
+    SymbolUploaderMock symbolUploaderMock = new SymbolUploaderMock();
+    Config config = mock(Config.class);
+    when(config.getServiceName()).thenReturn("service1");
+    SymbolSink symbolSink = new SymbolSink(config, symbolUploaderMock, 1);
+    final int NUM_CLASS_SCOPES = 10;
+    List<Scope> classScopes = new ArrayList<>();
+    for (int i = 0; i < NUM_CLASS_SCOPES; i++) {
+      classScopes.add(
+          Scope.builder(ScopeType.CLASS, "class" + i, 0, 0)
+              .scopes(
+                  Collections.singletonList(
+                      Scope.builder(ScopeType.METHOD, "class" + i, 0, 0)
+                          .name("method" + i)
+                          .build()))
+              .build());
+    }
+    symbolSink.addScope(
+        Scope.builder(ScopeType.JAR, "jar1.jar", 0, 0)
+            .name("jar1.jar")
+            .scopes(classScopes)
+            .build());
+    symbolSink.flush();
+    // no request to upload because we cannot split the jar scope
+    assertTrue(symbolUploaderMock.multiPartContents.isEmpty());
+  }
+
+  @NotNull
+  private static String assertMultipartContent(SymbolUploaderMock symbolUploaderMock, int index) {
+    BatchUploader.MultiPartContent eventContent = symbolUploaderMock.multiPartContents.get(index);
+    assertEquals("event", eventContent.getPartName());
+    BatchUploader.MultiPartContent symbolContent =
+        symbolUploaderMock.multiPartContents.get(index + 1);
+    assertEquals("file", symbolContent.getPartName());
+    return new String(symbolContent.getContent());
+  }
+
   static class SymbolUploaderMock extends BatchUploader {
     final List<MultiPartContent> multiPartContents = new ArrayList<>();
 
