|
52 | 52 | import datadog.trace.relocate.api.IOLogger; |
53 | 53 | import datadog.trace.util.PidHelper; |
54 | 54 | import delight.fileupload.FileUpload; |
| 55 | +import io.airlift.compress.zstd.ZstdInputStream; |
55 | 56 | import java.io.BufferedInputStream; |
56 | 57 | import java.io.ByteArrayInputStream; |
57 | 58 | import java.io.ByteArrayOutputStream; |
@@ -337,7 +338,7 @@ public void testZippedInput() throws Exception { |
337 | 338 | } |
338 | 339 |
|
339 | 340 | @ParameterizedTest |
340 | | - @ValueSource(strings = {"on", "lz4", "gzip", "off", "invalid"}) |
| 341 | + @ValueSource(strings = {"on", "lz4", "gzip", "zstd", "off", "invalid"}) |
341 | 342 | public void testCompression(final String compression) throws Exception { |
342 | 343 | when(config.getApiKey()).thenReturn(null); |
343 | 344 | when(config.getProfilingUploadCompression()).thenReturn(compression); |
@@ -372,6 +373,8 @@ public void testCompression(final String compression) throws Exception { |
372 | 373 | byte[] uploadedBytes = rawJfr.get(); |
373 | 374 | if (compression.equals("gzip")) { |
374 | 375 | uploadedBytes = unGzip(uploadedBytes); |
| 376 | + } else if (compression.equals("zstd")) { |
| 377 | + uploadedBytes = unZstd(uploadedBytes); |
375 | 378 | } else if (compression.equals("on") |
376 | 379 | || compression.equals("lz4") |
377 | 380 | || compression.equals("invalid")) { |
@@ -895,6 +898,13 @@ private static byte[] unLz4(final byte[] compressed) throws IOException { |
895 | 898 | return result.toByteArray(); |
896 | 899 | } |
897 | 900 |
|
| 901 | + private static byte[] unZstd(final byte[] compressed) throws IOException { |
| 902 | + final InputStream stream = new ZstdInputStream(new ByteArrayInputStream(compressed)); |
| 903 | + final ByteArrayOutputStream result = new ByteArrayOutputStream(); |
| 904 | + ByteStreams.copy(stream, result); |
| 905 | + return result.toByteArray(); |
| 906 | + } |
| 907 | + |
898 | 908 | private void uploadAndWait(final RecordingType recordingType, final RecordingData data) |
899 | 909 | throws InterruptedException { |
900 | 910 | final CountDownLatch latch = new CountDownLatch(1); |
|
0 commit comments