diff --git a/dd-trace-core/src/main/java/datadog/trace/common/metrics/Sink.java b/dd-trace-core/src/main/java/datadog/trace/common/metrics/Sink.java index f8261339926..eeeaf64d2b0 100644 --- a/dd-trace-core/src/main/java/datadog/trace/common/metrics/Sink.java +++ b/dd-trace-core/src/main/java/datadog/trace/common/metrics/Sink.java @@ -1,15 +1,8 @@ package datadog.trace.common.metrics; import datadog.communication.serialization.ChunkedBufferConsumer; -import java.nio.ByteBuffer; -import java.util.Collections; public interface Sink extends ChunkedBufferConsumer { - /** Accepts a payload held in a single buffer. */ - default void accept(int messageCount, ByteBuffer buffer) { - accept(messageCount, Collections.singletonList(buffer)); - } - void register(EventListener listener); } diff --git a/dd-trace-core/src/main/java/datadog/trace/core/datastreams/MsgPackDatastreamsPayloadWriter.java b/dd-trace-core/src/main/java/datadog/trace/core/datastreams/MsgPackDatastreamsPayloadWriter.java index f94210309d9..19616a33f67 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/datastreams/MsgPackDatastreamsPayloadWriter.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/datastreams/MsgPackDatastreamsPayloadWriter.java @@ -2,7 +2,7 @@ import static java.nio.charset.StandardCharsets.ISO_8859_1; -import datadog.communication.serialization.GrowableBuffer; +import datadog.communication.serialization.ChunkedStreamingBuffer; import datadog.communication.serialization.Writable; import datadog.communication.serialization.WritableFormatter; import datadog.communication.serialization.msgpack.MsgPackWriter; @@ -14,6 +14,7 @@ import datadog.trace.api.datastreams.TransactionInfo; import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString; import datadog.trace.common.metrics.Sink; +import java.nio.BufferOverflowException; import java.util.Collection; import java.util.List; import java.util.Map; @@ -62,18 +63,27 @@ public class MsgPackDatastreamsPayloadWriter implements DatastreamsPayloadWriter private static final byte[] CONFIG_MEMBER_PROTOCOL = "MemberProtocol".getBytes(ISO_8859_1); private static final byte[] CONFIG_ENTRIES = "Config".getBytes(ISO_8859_1); - private static final int INITIAL_CAPACITY = 512 * 1024; + private static final int MAX_PAYLOAD_SIZE = 5 << 20; private final WritableFormatter writer; private final Sink sink; - private final GrowableBuffer buffer; + private final ChunkedStreamingBuffer buffer; private final WellKnownTags wellKnownTags; private final byte[] tracerVersionValue; private final byte[] primaryTagValue; public MsgPackDatastreamsPayloadWriter( Sink sink, WellKnownTags wellKnownTags, String tracerVersion, String primaryTag) { - buffer = new GrowableBuffer(INITIAL_CAPACITY); + this(sink, wellKnownTags, tracerVersion, primaryTag, MAX_PAYLOAD_SIZE); + } + + MsgPackDatastreamsPayloadWriter( + Sink sink, + WellKnownTags wellKnownTags, + String tracerVersion, + String primaryTag, + int maxPayloadSize) { + buffer = new ChunkedStreamingBuffer(maxPayloadSize, sink); writer = new MsgPackWriter(buffer); this.sink = sink; this.wellKnownTags = wellKnownTags; @@ -108,6 +118,20 @@ public long getProductsMask() { @Override public void writePayload(Collection data, String serviceNameOverride) { + try { + write(data, serviceNameOverride); + } catch (BufferOverflowException ignored) { + // the buffer discards the rest of the payload, and mark() below reports it + } + if (buffer.mark()) { + buffer.flush(); + } else { + log.debug("dropping a data streams payload larger than {} bytes", buffer.capacity()); + buffer.reset(); + } + } + + private void write(Collection data, String serviceNameOverride) { final List processTags = ProcessTags.getTagsAsUTF8ByteStringList(); final boolean hasProcessTags = processTags != null; writer.startMap(8 + (hasProcessTags ? 1 : 0)); @@ -199,10 +223,6 @@ public void writePayload(Collection data, String serviceNameOverrid writer.startArray(processTags.size()); processTags.forEach(writer::writeUTF8); } - - buffer.mark(); - sink.accept(buffer.messageCount(), buffer.slice()); - buffer.reset(); } private void writeBucket(StatsBucket bucket, Writable packer) { diff --git a/dd-trace-core/src/test/java/datadog/trace/core/datastreams/MsgPackDatastreamsPayloadWriterTest.java b/dd-trace-core/src/test/java/datadog/trace/core/datastreams/MsgPackDatastreamsPayloadWriterTest.java new file mode 100644 index 00000000000..648b77c07d3 --- /dev/null +++ b/dd-trace-core/src/test/java/datadog/trace/core/datastreams/MsgPackDatastreamsPayloadWriterTest.java @@ -0,0 +1,59 @@ +package datadog.trace.core.datastreams; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import datadog.trace.api.WellKnownTags; +import datadog.trace.common.metrics.EventListener; +import datadog.trace.common.metrics.Sink; +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import org.junit.jupiter.api.Test; + +class MsgPackDatastreamsPayloadWriterTest { + + private static final WellKnownTags TAGS = + new WellKnownTags("runtimeid", "hostname", "env", "service", "version", "language"); + + @Test + void payloadOverTheCapIsDroppedWithoutThrowing() { + int size = payloadSize(); + + CapturingSink tooSmall = new CapturingSink(); + new MsgPackDatastreamsPayloadWriter(tooSmall, TAGS, "1.0", null, size - 1) + .writePayload(Collections.emptyList(), null); + assertTrue(tooSmall.payloads.isEmpty()); + + CapturingSink exact = new CapturingSink(); + new MsgPackDatastreamsPayloadWriter(exact, TAGS, "1.0", null, size) + .writePayload(Collections.emptyList(), null); + assertEquals(Collections.singletonList(size), exact.payloads); + } + + private static int payloadSize() { + CapturingSink sink = new CapturingSink(); + new MsgPackDatastreamsPayloadWriter(sink, TAGS, "1.0", null) + .writePayload(Collections.emptyList(), null); + assertEquals(1, sink.payloads.size()); + return sink.payloads.get(0); + } + + /** Records the size of each payload it receives. */ + private static final class CapturingSink implements Sink { + final List payloads = new ArrayList<>(); + + @Override + public void accept(int messageCount, List chunks) { + int size = 0; + for (ByteBuffer chunk : chunks) { + size += chunk.remaining(); + } + payloads.add(size); + } + + @Override + public void register(EventListener listener) {} + } +}