Enable ZSTD compression for profiling (#8862)

Author: Matt

dd-java-agent/agent-profiling/profiling-controller/src/main/java/com/datadog/profiling/controller/ProfilerSettingsSupport.java

@@ -123,10 +123,12 @@ protected ProfilerSettingsSupport(
         configProvider.getInteger(
             ProfilingConfig.PROFILING_UPLOAD_TIMEOUT,
             ProfilingConfig.PROFILING_UPLOAD_TIMEOUT_DEFAULT);
+    // First try the new debug upload compression property, and fall back to the deprecated one
     uploadCompression =
         configProvider.getString(
-            ProfilingConfig.PROFILING_UPLOAD_COMPRESSION,
-            ProfilingConfig.PROFILING_UPLOAD_COMPRESSION_DEFAULT);
+            ProfilingConfig.PROFILING_DEBUG_UPLOAD_COMPRESSION,
+            ProfilingConfig.PROFILING_DEBUG_UPLOAD_COMPRESSION_DEFAULT,
+            ProfilingConfig.PROFILING_UPLOAD_COMPRESSION);
     allocationProfilingEnabled =
         configProvider.getBoolean(
             ProfilingConfig.PROFILING_ALLOCATION_ENABLED,

dd-java-agent/agent-profiling/profiling-uploader/build.gradle

@@ -31,6 +31,7 @@ dependencies {
 
   implementation libs.okhttp
   implementation libs.lz4
+  implementation group: 'io.airlift', name: 'aircompressor', version: '2.0.2'
 
   testImplementation libs.bundles.junit5
   testImplementation project(':dd-java-agent:agent-profiling:profiling-testing')

dd-java-agent/agent-profiling/profiling-uploader/src/main/java/com/datadog/profiling/uploader/CompressingRequestBody.java

@@ -2,6 +2,7 @@
 
 import datadog.trace.api.Platform;
 import datadog.trace.api.profiling.RecordingInputStream;
+import io.airlift.compress.zstd.ZstdOutputStream;
 import java.io.BufferedOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
@@ -86,6 +87,7 @@ interface RetryBackoff {
   private static final int[] LZ4_MAGIC = new int[] {0x04, 0x22, 0x4D, 0x18};
   private static final int ZIP_MAGIC[] = new int[] {80, 75, 3, 4};
   private static final int GZ_MAGIC[] = new int[] {31, 139};
+  private static final int ZSTD_MAGIC[] = new int[] {0x28, 0xB5, 0x2F, 0xFD};
 
   private final InputStreamSupplier inputStreamSupplier;
   private final OutputStreamMappingFunction outputStreamMapper;
@@ -269,7 +271,24 @@ public void close() throws IOException {
    */
   static boolean isCompressed(@Nonnull final InputStream is) throws IOException {
     checkMarkSupported(is);
-    return isGzip(is) || isLz4(is) || isZip(is);
+    return isGzip(is) || isLz4(is) || isZip(is) || isZstd(is);
+  }
+
+  /**
+   * Check whether the stream represents ZSTD data
+   *
+   * @param is input stream; must support {@linkplain InputStream#mark(int)}
+   * @return {@literal true} if the stream represents ZSTD data
+   * @throws IOException
+   */
+  static boolean isZstd(@Nonnull final InputStream is) throws IOException {
+    checkMarkSupported(is);
+    is.mark(ZSTD_MAGIC.length);
+    try {
+      return hasMagic(is, ZSTD_MAGIC);
+    } finally {
+      is.reset();
+    }
   }
 
   /**
@@ -331,12 +350,11 @@ private static void checkMarkSupported(@Nonnull final InputStream is) throws IOE
 
   private static OutputStreamMappingFunction getOutputStreamMapper(
       @Nonnull CompressionType compressionType) {
-    // currently only gzip and off are supported
-    // this needs to be updated once more compression types are added
-    compressionType =
-        (Platform.isNativeImage() && compressionType != CompressionType.OFF
-            ? CompressionType.GZIP
-            : compressionType);
+    // Handle native image compatibility
+    if (Platform.isNativeImage() && compressionType != CompressionType.OFF) {
+      compressionType = CompressionType.GZIP;
+    }
+
     switch (compressionType) {
       case GZIP:
         {
@@ -346,6 +364,10 @@ private static OutputStreamMappingFunction getOutputStreamMapper(
         {
           return out -> out;
         }
+      case ZSTD:
+        {
+          return CompressingRequestBody::toZstdStream;
+        }
       case ON:
       case LZ4:
       default:
@@ -355,6 +377,12 @@ private static OutputStreamMappingFunction getOutputStreamMapper(
     }
   }
 
+  private static OutputStream toZstdStream(@Nonnull OutputStream os) throws IOException {
+    // Default compression level is 3 which provides a good balance between performance and
+    // compression ratio
+    return new ZstdOutputStream(os);
+  }
+
   private static OutputStream toLz4Stream(@Nonnull OutputStream os) throws IOException {
     return new LZ4FrameOutputStream(
         os,

dd-java-agent/agent-profiling/profiling-uploader/src/main/java/com/datadog/profiling/uploader/CompressionType.java

@@ -12,7 +12,9 @@ enum CompressionType {
   /** Lower compression ratio with less CPU overhead * */
   LZ4,
   /** Better compression ratio for the price of higher CPU usage * */
-  GZIP;
+  GZIP,
+  /** High compression ratio with reasonable CPU usage * */
+  ZSTD;
 
   private static final Logger log = LoggerFactory.getLogger(CompressionType.class);
 
@@ -30,8 +32,10 @@ static CompressionType of(String type) {
         return LZ4;
       case "gzip":
         return GZIP;
+      case "zstd":
+        return ZSTD;
       default:
-        log.warn("Unrecognizable compression type: {}. Defaulting to 'on'.", type);
+        log.warn("Unrecognizable compression type: {}. Defaulting to 'lz4'.", type);
         return ON;
     }
   }

dd-java-agent/agent-profiling/profiling-uploader/src/test/java/com/datadog/profiling/uploader/CompressingRequestBodyTest.java

@@ -6,6 +6,8 @@
 import static org.mockito.Mockito.when;
 
 import datadog.trace.api.profiling.RecordingInputStream;
+import io.airlift.compress.zstd.ZstdInputStream;
+import io.airlift.compress.zstd.ZstdOutputStream;
 import java.io.BufferedInputStream;
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
@@ -186,6 +188,15 @@ void writeTo(CompressionType compressionType) throws IOException {
           assertEquals(compressed.length, instance.getWrittenBytes());
           break;
         }
+      case ZSTD:
+        {
+          assertTrue(CompressingRequestBody.isZstd(compressedStream));
+          byte[] uncompressed = IOUtils.toByteArray(new ZstdInputStream(compressedStream));
+          assertArrayEquals(recordingData, uncompressed);
+          assertEquals(recordingData.length, instance.getReadBytes());
+          assertEquals(compressed.length, instance.getWrittenBytes());
+          break;
+        }
     }
   }
 
@@ -210,6 +221,11 @@ void writeToRecompression(CompressionType targetType) throws IOException {
             compressedStream = new GZIPOutputStream(baos);
             break;
           }
+        case ZSTD:
+          {
+            compressedStream = new ZstdOutputStream(baos);
+            break;
+          }
       }
       assertNotNull(compressedStream);
 

dd-java-agent/build.gradle

@@ -306,7 +306,7 @@ tasks.register('checkAgentJarSize').configure {
   doLast {
     // Arbitrary limit to prevent unintentional increases to the agent jar size
     // Raise or lower as required
-    assert shadowJar.archiveFile.get().getAsFile().length() <= 31 * 1024 * 1024
+    assert shadowJar.archiveFile.get().getAsFile().length() <= 32 * 1024 * 1024
   }
 
   dependsOn "shadowJar"

dd-trace-api/src/main/java/datadog/trace/api/config/ProfilingConfig.java

@@ -33,8 +33,20 @@ public final class ProfilingConfig {
       "profiling.jfr-template-override-file";
   public static final String PROFILING_UPLOAD_TIMEOUT = "profiling.upload.timeout";
   public static final int PROFILING_UPLOAD_TIMEOUT_DEFAULT = 30;
+  /**
+   * @deprecated Use {@link #PROFILING_DEBUG_UPLOAD_COMPRESSION} instead. This will be removed in a
+   *     future release.
+   */
+  @Deprecated
   public static final String PROFILING_UPLOAD_COMPRESSION = "profiling.upload.compression";
-  public static final String PROFILING_UPLOAD_COMPRESSION_DEFAULT = "on";
+  /**
+   * Default compression value. Supported values are: - "on": equivalent to "lz4", will later be
+   * "zstd" - "off": disables compression - "lz4": uses LZ4 compression (fast with moderate
+   * compression ratio) - "gzip": uses GZIP compression (higher compression ratio but slower) -
+   * "zstd": uses ZSTD compression (high compression ratio with reasonable performance)
+   */
+  public static final String PROFILING_DEBUG_UPLOAD_COMPRESSION_DEFAULT = "lz4";
+
   public static final String PROFILING_PROXY_HOST = "profiling.proxy.host";
   public static final String PROFILING_PROXY_PORT = "profiling.proxy.port";
   public static final int PROFILING_PROXY_PORT_DEFAULT = 8080;
@@ -192,6 +204,14 @@ public final class ProfilingConfig {
 
   public static final String PROFILING_DEBUG_DUMP_PATH = "profiling.debug.dump_path";
   public static final String PROFILING_DEBUG_JFR_DISABLED = "profiling.debug.jfr.disabled";
+  /**
+   * Configuration for profile upload compression. Supported values are: - "on": equivalent to "lz4"
+   * - "off": disables compression - "lz4": uses LZ4 compression (fast with moderate compression
+   * ratio) - "gzip": uses GZIP compression (higher compression ratio but slower) - "zstd": uses
+   * ZSTD compression (high compression ratio with reasonable performance)
+   */
+  public static final String PROFILING_DEBUG_UPLOAD_COMPRESSION =
+      "profiling.debug.upload.compression";
 
   public static final String PROFILING_CONTEXT_ATTRIBUTES = "profiling.context.attributes";
 

internal-api/src/main/java/datadog/trace/api/Config.java

@@ -1259,7 +1259,7 @@ PROFILING_DATADOG_PROFILER_ENABLED, isDatadogProfilerSafeInCurrentEnvironment())
         configProvider.getInteger(PROFILING_UPLOAD_TIMEOUT, PROFILING_UPLOAD_TIMEOUT_DEFAULT);
     profilingUploadCompression =
         configProvider.getString(
-            PROFILING_UPLOAD_COMPRESSION, PROFILING_UPLOAD_COMPRESSION_DEFAULT);
+            PROFILING_UPLOAD_COMPRESSION, PROFILING_DEBUG_UPLOAD_COMPRESSION_DEFAULT);
     profilingProxyHost = configProvider.getString(PROFILING_PROXY_HOST);
     profilingProxyPort =
         configProvider.getInteger(PROFILING_PROXY_PORT, PROFILING_PROXY_PORT_DEFAULT);