diff --git a/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/base/HttpServer.groovy b/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/base/HttpServer.groovy deleted file mode 100644 index 0e06968c28c..00000000000 --- a/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/base/HttpServer.groovy +++ /dev/null @@ -1,12 +0,0 @@ -package datadog.trace.agent.test.base - -import java.util.concurrent.TimeoutException - -interface HttpServer { - - void start() throws TimeoutException - - void stop() - - URI address() -} diff --git a/dd-java-agent/testing/src/main/java/datadog/trace/agent/test/base/HttpServer.java b/dd-java-agent/testing/src/main/java/datadog/trace/agent/test/base/HttpServer.java new file mode 100644 index 00000000000..620a135ac9a --- /dev/null +++ b/dd-java-agent/testing/src/main/java/datadog/trace/agent/test/base/HttpServer.java @@ -0,0 +1,13 @@ +package datadog.trace.agent.test.base; + +import java.net.URI; +import java.util.concurrent.TimeoutException; + +public interface HttpServer { + + void start() throws TimeoutException; + + void stop(); + + URI address(); +} diff --git a/dd-java-agent/testing/src/main/java/datadog/trace/agent/test/server/http/JavaTestHttpServer.java b/dd-java-agent/testing/src/main/java/datadog/trace/agent/test/server/http/JavaTestHttpServer.java new file mode 100644 index 00000000000..1ff635fba5e --- /dev/null +++ b/dd-java-agent/testing/src/main/java/datadog/trace/agent/test/server/http/JavaTestHttpServer.java @@ -0,0 +1,628 @@ +package datadog.trace.agent.test.server.http; + +import static datadog.trace.agent.test.server.http.HttpServletRequestExtractAdapter.GETTER; +import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.extractContextAndGetSpanContext; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; + +import datadog.trace.agent.test.base.HttpServer; +import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; +import datadog.trace.bootstrap.instrumentation.api.Tags; +import de.thetaphi.forbiddenapis.SuppressForbidden; +import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; +import java.io.ByteArrayOutputStream; +import java.io.File; +import java.io.FileOutputStream; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.net.URI; +import java.net.URISyntaxException; +import java.net.URL; +import java.nio.charset.StandardCharsets; +import java.security.KeyManagementException; +import java.security.NoSuchAlgorithmException; +import java.security.cert.X509Certificate; +import java.util.ArrayList; +import java.util.Enumeration; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; +import javax.net.ssl.HostnameVerifier; +import javax.net.ssl.SSLContext; +import javax.net.ssl.TrustManager; +import javax.net.ssl.X509TrustManager; +import javax.servlet.ServletException; +import javax.servlet.http.HttpServletRequest; +import javax.servlet.http.HttpServletResponse; +import org.eclipse.jetty.http.HttpMethod; +import org.eclipse.jetty.http.HttpVersion; +import org.eclipse.jetty.server.Handler; +import org.eclipse.jetty.server.HttpConfiguration; +import org.eclipse.jetty.server.HttpConnectionFactory; +import org.eclipse.jetty.server.Request; +import org.eclipse.jetty.server.SecureRequestCustomizer; +import org.eclipse.jetty.server.Server; +import org.eclipse.jetty.server.ServerConnector; +import org.eclipse.jetty.server.SslConnectionFactory; +import org.eclipse.jetty.server.handler.AbstractHandler; +import org.eclipse.jetty.server.handler.HandlerList; +import org.eclipse.jetty.util.ssl.SslContextFactory; +import org.eclipse.jetty.util.thread.QueuedThreadPool; + +@SuppressFBWarnings({"IS2_INCONSISTENT_SYNC", "PA_PUBLIC_PRIMITIVE_ATTRIBUTE"}) +public class JavaTestHttpServer implements AutoCloseable { + + @FunctionalInterface + public interface RequestHandler { + void handle(HandlerApi api) throws Exception; + } + + private final Server internalServer; + private HandlersSpec handlers; + private Consumer customizer = s -> {}; + + public String keystorePath; + private URI address; + private URI secureAddress; + private final AtomicReference last = new AtomicReference<>(); + + public final SSLContext sslContext; + + private final X509TrustManager trustManager = + new X509TrustManager() { + @Override + public X509Certificate[] getAcceptedIssuers() { + return new X509Certificate[0]; + } + + @Override + public void checkClientTrusted(X509Certificate[] certificate, String str) {} + + @Override + public void checkServerTrusted(X509Certificate[] certificate, String str) {} + }; + + private final HostnameVerifier hostnameVerifier = + (hostname, session) -> "localhost".equals(hostname); + + public static JavaTestHttpServer httpServer(Consumer spec) { + JavaTestHttpServer server = new JavaTestHttpServer(); + spec.accept(server); + server.start(); + return server; + } + + private JavaTestHttpServer() { + // In some versions, Jetty requires max threads > than some arbitrary calculated value. + // The calculated value can be high in CI. There is no easy way to override the configuration + // in a version-neutral way. + internalServer = new Server(new QueuedThreadPool(400)); + try { + sslContext = SSLContext.getInstance("TLSv1.2"); + sslContext.init(null, new TrustManager[] {trustManager}, null); + } catch (NoSuchAlgorithmException | KeyManagementException e) { + throw new IllegalStateException(e); + } + } + + @SuppressForbidden + public JavaTestHttpServer start() { + if (internalServer.isStarted()) { + return this; + } + synchronized (this) { + if (!internalServer.isRunning()) { + if (handlers == null) { + throw new IllegalStateException("handlers must be defined"); + } + HandlerList handlerList = new HandlerList(); + handlerList.setHandlers(handlers.configured.toArray(new Handler[0])); + internalServer.setHandler(handlerList); + + HttpConfiguration httpConfiguration = new HttpConfiguration(); + + // HTTP + ServerConnector http = + new ServerConnector(internalServer, new HttpConnectionFactory(httpConfiguration)); + http.setHost("localhost"); + http.setPort(0); + internalServer.addConnector(http); + + // HTTPS + SslContextFactory sslContextFactory = new SslContextFactory(); + keystorePath = + extractKeystoreToDisk(JavaTestHttpServer.class.getResource("datadog.jks")).getPath(); + sslContextFactory.setKeyStorePath(keystorePath); + sslContextFactory.setKeyStorePassword("datadog"); + HttpConfiguration httpsConfiguration = new HttpConfiguration(httpConfiguration); + httpsConfiguration.addCustomizer(new SecureRequestCustomizer()); + ServerConnector https = + new ServerConnector( + internalServer, + new SslConnectionFactory(sslContextFactory, HttpVersion.HTTP_1_1.asString()), + new HttpConnectionFactory(httpsConfiguration)); + https.setHost("localhost"); + https.setPort(0); + internalServer.addConnector(https); + + customizer.accept(internalServer); + try { + internalServer.start(); + } catch (Exception e) { + throw new IllegalStateException(e); + } + // set after starting, otherwise two callbacks get added. + internalServer.setStopAtShutdown(true); + + try { + address = new URI("http://localhost:" + http.getLocalPort()); + secureAddress = new URI("https://localhost:" + https.getLocalPort()); + } catch (URISyntaxException e) { + throw new IllegalStateException(e); + } + } + } + long startTime = System.nanoTime(); + long rem = TimeUnit.SECONDS.toMillis(5); + while (!internalServer.isStarted()) { + if (rem <= 0) { + throw new RuntimeException( + new TimeoutException( + "Failed to start server " + this + " on port " + address.getPort())); + } + try { + Thread.sleep(Math.min(rem, 100)); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + long endTime = System.nanoTime(); + rem -= TimeUnit.NANOSECONDS.toMillis(endTime - startTime); + startTime = endTime; + } + System.out.println("Started server " + this + " on " + address + " and " + secureAddress); + return this; + } + + private File extractKeystoreToDisk(URL internalFile) { + try (InputStream inputStream = internalFile.openStream()) { + File tempFile = File.createTempFile("datadog", ".jks"); + tempFile.deleteOnExit(); + try (OutputStream out = new FileOutputStream(tempFile)) { + byte[] buffer = new byte[1024]; + int len; + while ((len = inputStream.read(buffer)) != -1) { + out.write(buffer, 0, len); + } + } + return tempFile; + } catch (IOException e) { + throw new IllegalStateException(e); + } + } + + @SuppressForbidden + public JavaTestHttpServer stop() { + System.out.println("Stopping server " + this + " on " + address + " and " + secureAddress); + try { + internalServer.stop(); + } catch (Exception e) { + throw new IllegalStateException(e); + } + return this; + } + + @Override + public void close() { + stop(); + } + + public URI getAddress() { + return address; + } + + public URI getSecureAddress() { + return secureAddress; + } + + public X509TrustManager getTrustManager() { + return trustManager; + } + + public HostnameVerifier getHostnameVerifier() { + return hostnameVerifier; + } + + public HandlerApi.RequestApi getLastRequest() { + return last.get(); + } + + public void customizer(Consumer spec) { + this.customizer = spec; + } + + public void handlers(Consumer spec) { + if (handlers != null) { + throw new IllegalStateException("handlers already defined"); + } + handlers = new HandlersSpec(); + spec.accept(handlers); + } + + public HttpServer asHttpServer() { + return new HttpServerAdapter(this, false); + } + + public HttpServer asHttpServer(boolean secure) { + return new HttpServerAdapter(this, secure); + } + + public final class HandlersSpec { + final List configured = new ArrayList<>(); + + public void get(String path, RequestHandler spec) { + if (path == null) { + throw new IllegalArgumentException("path must not be null"); + } + configured.add(new PathHandler(HttpMethod.GET, path, spec)); + } + + public void post(String path, RequestHandler spec) { + if (path == null) { + throw new IllegalArgumentException("path must not be null"); + } + configured.add(new PathHandler(HttpMethod.POST, path, spec)); + } + + public void put(String path, RequestHandler spec) { + if (path == null) { + throw new IllegalArgumentException("path must not be null"); + } + configured.add(new PathHandler(HttpMethod.PUT, path, spec)); + } + + public void connect(RequestHandler spec) { + configured.add(new MethodHandler(HttpMethod.CONNECT, spec)); + } + + public void prefix(String path, RequestHandler spec) { + configured.add(new PrefixHandler(path, spec)); + } + + public void all(RequestHandler spec) { + configured.add(new AllHandler(spec)); + } + } + + private class AllHandler extends AbstractHandler { + final RequestHandler spec; + + AllHandler(RequestHandler spec) { + this.spec = spec; + } + + @Override + public void handle( + String target, + Request baseRequest, + HttpServletRequest request, + HttpServletResponse response) + throws IOException, ServletException { + send(baseRequest, response); + } + + void send(Request baseRequest, HttpServletResponse response) { + HandlerApi api = new HandlerApi(baseRequest, response); + last.set(api.getRequest()); + try { + spec.handle(api); + } catch (Exception e) { + try { + api.getResponse().status(500).send(e.getMessage()); + } catch (Exception ignored) { + // ignore + } + e.printStackTrace(); + } + } + } + + private class MethodHandler extends AllHandler { + private final String method; + + MethodHandler(HttpMethod method, RequestHandler spec) { + super(spec); + this.method = method.name(); + } + + @Override + public void handle( + String target, + Request baseRequest, + HttpServletRequest request, + HttpServletResponse response) + throws IOException, ServletException { + if (request.getMethod().equalsIgnoreCase(method)) { + super.handle(target, baseRequest, request, response); + } + } + } + + private class PathHandler extends MethodHandler { + private final String path; + + PathHandler(HttpMethod method, String path, RequestHandler spec) { + super(method, spec); + this.path = path.startsWith("/") ? path : "/" + path; + } + + @Override + public void handle( + String target, + Request baseRequest, + HttpServletRequest request, + HttpServletResponse response) + throws IOException, ServletException { + if (path.equals(target)) { + super.handle(target, baseRequest, request, response); + } + } + } + + private class PrefixHandler extends AllHandler { + private final String prefix; + + PrefixHandler(String prefix, RequestHandler spec) { + super(spec); + this.prefix = prefix.startsWith("/") ? prefix : "/" + prefix; + } + + @Override + public void handle( + String target, + Request baseRequest, + HttpServletRequest request, + HttpServletResponse response) + throws IOException, ServletException { + if (target.startsWith(prefix)) { + super.handle(target, baseRequest, request, response); + } + } + } + + public static class HandlerApi { + private final RequestApi req; + private final HttpServletResponse resp; + + private HandlerApi(Request request, HttpServletResponse response) { + this.req = new RequestApi(request); + this.resp = response; + } + + public RequestApi getRequest() { + return req; + } + + public ResponseApi getResponse() { + return new ResponseApi(); + } + + public void redirect(String uri) throws IOException { + resp.sendRedirect(uri); + req.orig.setHandled(true); + } + + public void handleDistributedRequest() { + boolean isDDServer = true; + String header = req.getHeader("is-dd-server"); + if (header != null) { + isDDServer = Boolean.parseBoolean(header); + } + if (isDDServer) { + AgentSpanContext extractedContext = extractContextAndGetSpanContext(req.orig, GETTER); + if (extractedContext != null) { + startSpan("test", "test-http-server", extractedContext) + .setTag("path", req.getPath()) + .setTag(Tags.SPAN_KIND, Tags.SPAN_KIND_SERVER) + .finish(); + } else { + startSpan("test", "test-http-server") + .setTag("path", req.getPath()) + .setTag(Tags.SPAN_KIND, Tags.SPAN_KIND_SERVER) + .finish(); + } + } + } + + public static class RequestApi { + final Request orig; + final String path; + final Headers headers; + final int contentLength; + final String contentType; + final byte[] body; + + RequestApi(Request req) { + this.orig = req; + this.path = req.getPathInfo(); + this.headers = new Headers(req); + this.contentLength = req.getContentLength(); + String ct = req.getContentType(); + this.contentType = ct; + try (InputStream is = req.getInputStream()) { + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + byte[] buffer = new byte[1024]; + int read; + while ((read = is.read(buffer)) != -1) { + baos.write(buffer, 0, read); + } + this.body = baos.toByteArray(); + } catch (IOException e) { + throw new IllegalStateException(e); + } + } + + public String getPath() { + return path; + } + + public int getContentLength() { + return contentLength; + } + + public String getContentType() { + if (contentType == null) { + return null; + } + int idx = contentType.indexOf(';'); + return idx >= 0 ? contentType.substring(0, idx) : contentType; + } + + public Headers getHeaders() { + return headers; + } + + public String getHeader(String header) { + return headers.get(header); + } + + public byte[] getBody() { + return body; + } + + public String getText() { + return new String(body); + } + + public String getParameter(String parameter) { + return orig.getParameter(parameter); + } + } + + public class ResponseApi { + private static final String DEFAULT_TYPE = "text/plain;charset=utf-8"; + private int status = 200; + private final Map headers = new HashMap<>(); + + public ResponseApi status(int status) { + this.status = status; + return this; + } + + public ResponseApi addHeader(String headerName, String headerValue) { + this.headers.put(headerName, headerValue); + return this; + } + + public void send() { + sendWithType(DEFAULT_TYPE); + } + + public void sendWithType(String contentType) { + if (contentType == null) { + throw new IllegalArgumentException("contentType must not be null"); + } + if (req.orig.isHandled()) { + throw new IllegalStateException("response already handled"); + } + req.orig.setContentType(contentType); + resp.setStatus(status); + for (Map.Entry e : headers.entrySet()) { + resp.addHeader(e.getKey(), e.getValue()); + } + req.orig.setHandled(true); + } + + public void send(String body) { + sendWithType(DEFAULT_TYPE, body); + } + + public void sendWithType(String contentType, String body) { + if (body == null) { + throw new IllegalArgumentException("body must not be null"); + } + sendWithType(contentType); + byte[] bytes = body.getBytes(StandardCharsets.UTF_8); + resp.setContentLength(bytes.length); + try { + resp.getWriter().print(body); + } catch (IOException e) { + throw new IllegalStateException(e); + } + } + + public void send(byte[] body) { + sendWithType(DEFAULT_TYPE, body); + } + + public void sendWithType(String contentType, byte[] body) { + if (body == null) { + throw new IllegalArgumentException("body must not be null"); + } + sendWithType(contentType); + resp.setContentLength(body.length); + try { + resp.getOutputStream().write(body); + } catch (IOException e) { + throw new IllegalStateException(e); + } + } + } + } + + public static class Headers { + private final Map headers = new HashMap<>(); + + private Headers(Request request) { + Enumeration names = request.getHeaderNames(); + if (names != null) { + while (names.hasMoreElements()) { + String name = names.nextElement(); + headers.put(name, request.getHeader(name)); + } + } + } + + public String get(String header) { + return headers.get(header); + } + } + + public static class HttpServerAdapter implements HttpServer { + final JavaTestHttpServer server; + final boolean secure; + URI address; + + public HttpServerAdapter(JavaTestHttpServer server, boolean secure) { + this.server = server; + this.secure = secure; + } + + @Override + public void start() throws TimeoutException { + server.start(); + address = secure ? server.secureAddress : server.address; + if (!address.getPath().endsWith("/")) { + try { + address = new URI(address.toString() + "/"); + } catch (URISyntaxException e) { + throw new IllegalStateException(e); + } + } + } + + @Override + public void stop() { + server.stop(); + } + + @Override + public URI address() { + return address; + } + } +} diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/datastreams/CheckpointerTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/datastreams/CheckpointerTest.groovy deleted file mode 100644 index b41a8d3ad13..00000000000 --- a/dd-trace-core/src/test/groovy/datadog/trace/core/datastreams/CheckpointerTest.groovy +++ /dev/null @@ -1,53 +0,0 @@ -package datadog.trace.core.datastreams - - -import datadog.trace.api.experimental.DataStreamsContextCarrier -import datadog.trace.bootstrap.instrumentation.api.AgentTracer -import datadog.trace.core.test.DDCoreSpecification - -import static datadog.trace.api.config.GeneralConfig.DATA_STREAMS_ENABLED - -class CheckpointerTest extends DDCoreSpecification { - void 'test setting produce & consume checkpoint'() { - setup: - // Enable DSM - injectSysConfig(DATA_STREAMS_ENABLED, 'true') - // Create a test tracer - def tracer = tracerBuilder().build() - AgentTracer.forceRegister(tracer) - // Get the test checkpointer - def checkpointer = tracer.getDataStreamsCheckpointer() - // Declare the carrier to test injected data - def carrier = new CustomContextCarrier() - // Start and activate a span - def span = tracer.buildSpan('test', 'dsm-checkpoint').start() - def scope = tracer.activateSpan(span) - - when: - // Trigger produce checkpoint - checkpointer.setProduceCheckpoint('kafka', 'testTopic', carrier) - checkpointer.setConsumeCheckpoint('kafka', 'testTopic', carrier) - // Clean up span - scope.close() - span.finish() - - then: - carrier.entries().any { entry -> entry.getKey() == "dd-pathway-ctx-base64" } - span.context().pathwayContext.hash != 0 - } - - class CustomContextCarrier implements DataStreamsContextCarrier { - - private Map data = new HashMap<>() - - @Override - Set> entries() { - return data.entrySet() - } - - @Override - void set(String key, String value) { - data.put(key, value) - } - } -} diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/datastreams/DataStreamsTransactionExtractorsTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/datastreams/DataStreamsTransactionExtractorsTest.groovy deleted file mode 100644 index be3bdc9d003..00000000000 --- a/dd-trace-core/src/test/groovy/datadog/trace/core/datastreams/DataStreamsTransactionExtractorsTest.groovy +++ /dev/null @@ -1,23 +0,0 @@ -package datadog.trace.core.datastreams - -import datadog.trace.api.datastreams.DataStreamsTransactionExtractor -import datadog.trace.core.test.DDCoreSpecification - -class DataStreamsTransactionExtractorsTest extends DDCoreSpecification { - def "Deserialize from json"() { - when: - def list = DataStreamsTransactionExtractors.deserialize("""[ - {"name": "extractor", "type": "HTTP_OUT_HEADERS", "value": "transaction_id"}, - {"name": "second_extractor", "type": "HTTP_IN_HEADERS", "value": "transaction_id"} - ]""") - def extractors = list.getExtractors() - then: - extractors.size() == 2 - extractors[0].getName() == "extractor" - extractors[0].getType() == DataStreamsTransactionExtractor.Type.HTTP_OUT_HEADERS - extractors[0].getValue() == "transaction_id" - extractors[1].getName() == "second_extractor" - extractors[1].getType() == DataStreamsTransactionExtractor.Type.HTTP_IN_HEADERS - extractors[1].getValue() == "transaction_id" - } -} diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/datastreams/DataStreamsWritingTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/datastreams/DataStreamsWritingTest.groovy deleted file mode 100644 index f47904f522b..00000000000 --- a/dd-trace-core/src/test/groovy/datadog/trace/core/datastreams/DataStreamsWritingTest.groovy +++ /dev/null @@ -1,537 +0,0 @@ -package datadog.trace.core.datastreams - -import datadog.communication.ddagent.DDAgentFeaturesDiscovery -import datadog.communication.ddagent.SharedCommunicationObjects -import datadog.communication.http.OkHttpUtils -import datadog.trace.api.Config -import datadog.trace.api.ProcessTags -import datadog.trace.api.TraceConfig -import datadog.trace.api.WellKnownTags -import datadog.trace.api.datastreams.DataStreamsTags -import datadog.trace.api.time.ControllableTimeSource -import datadog.trace.api.datastreams.StatsPoint -import datadog.trace.core.DDTraceCoreInfo -import datadog.trace.core.test.DDCoreSpecification -import okhttp3.HttpUrl -import okio.BufferedSource -import okio.GzipSource -import okio.Okio -import org.msgpack.core.MessagePack -import org.msgpack.core.MessageUnpacker -import spock.lang.AutoCleanup -import spock.lang.Shared -import spock.util.concurrent.PollingConditions - -import static datadog.trace.agent.test.server.http.TestHttpServer.httpServer -import static datadog.trace.api.config.GeneralConfig.EXPERIMENTAL_PROPAGATE_PROCESS_TAGS_ENABLED -import static java.util.concurrent.TimeUnit.SECONDS - -/** - * This test class exists because a real integration test is not possible. see DataStreamsIntegrationTest - */ -class DataStreamsWritingTest extends DDCoreSpecification { - @Shared - List requestBodies - - @AutoCleanup - @Shared - def server = httpServer { - handlers { - post(DDAgentFeaturesDiscovery.V01_DATASTREAMS_ENDPOINT) { - owner.owner.owner.requestBodies.add(request.body) - response.status(200).send() - } - } - } - - static final DEFAULT_BUCKET_DURATION_NANOS = Config.get().getDataStreamsBucketDurationNanoseconds() - def setup() { - requestBodies = [] - } - - def "Service overrides split buckets"() { - given: - def conditions = new PollingConditions(timeout: 2) - - def testOkhttpClient = OkHttpUtils.buildHttpClient(HttpUrl.get(server.address), 5000L) - - def features = Stub(DDAgentFeaturesDiscovery) { - supportsDataStreams() >> true - } - - def wellKnownTags = new WellKnownTags("runtimeid", "hostname", "test", Config.get().getServiceName(), "version", "java") - - def fakeConfig = Stub(Config) { - getAgentUrl() >> server.address.toString() - getWellKnownTags() >> wellKnownTags - getPrimaryTag() >> "region-1" - } - - def sharedCommObjects = new SharedCommunicationObjects() - sharedCommObjects.featuresDiscovery = features - sharedCommObjects.agentHttpClient = testOkhttpClient - sharedCommObjects.createRemaining(fakeConfig) - - def timeSource = new ControllableTimeSource() - - def traceConfig = Mock(TraceConfig) { - isDataStreamsEnabled() >> true - } - def serviceNameOverride = "service-name-override" - - when: - def dataStreams = new DefaultDataStreamsMonitoring(fakeConfig, sharedCommObjects, timeSource, { traceConfig }) - dataStreams.start() - dataStreams.setThreadServiceName(serviceNameOverride) - dataStreams.add(new StatsPoint(DataStreamsTags.create(null, null), 9, 0, 10, timeSource.currentTimeNanos, 0, 0, 0, serviceNameOverride)) - dataStreams.trackBacklog(DataStreamsTags.createWithPartition("kafka_produce", "testTopic", "1", null, null), 130) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS) - // force flush - dataStreams.report() - dataStreams.close() - dataStreams.clearThreadServiceName() - then: - conditions.eventually { - assert requestBodies.size() == 1 - } - GzipSource gzipSource = new GzipSource(Okio.source(new ByteArrayInputStream(requestBodies[0]))) - - BufferedSource bufferedSource = Okio.buffer(gzipSource) - MessageUnpacker unpacker = MessagePack.newDefaultUnpacker(bufferedSource.inputStream()) - - assert unpacker.unpackMapHeader() == 9 - assert unpacker.unpackString() == "Env" - assert unpacker.unpackString() == "test" - assert unpacker.unpackString() == "Service" - assert unpacker.unpackString() == serviceNameOverride - } - - def "Write bucket to mock server with process tags enabled #processTagsEnabled"() { - setup: - injectSysConfig(EXPERIMENTAL_PROPAGATE_PROCESS_TAGS_ENABLED, "$processTagsEnabled") - ProcessTags.reset() - - def conditions = new PollingConditions(timeout: 2) - - def testOkhttpClient = OkHttpUtils.buildHttpClient(HttpUrl.get(server.address), 5000L) - - def features = Stub(DDAgentFeaturesDiscovery) { - supportsDataStreams() >> true - } - - def wellKnownTags = new WellKnownTags("runtimeid", "hostname", "test", Config.get().getServiceName(), "version", "java") - - def fakeConfig = Stub(Config) { - getAgentUrl() >> server.address.toString() - getWellKnownTags() >> wellKnownTags - getPrimaryTag() >> "region-1" - } - - def sharedCommObjects = new SharedCommunicationObjects() - sharedCommObjects.featuresDiscovery = features - sharedCommObjects.agentHttpClient = testOkhttpClient - sharedCommObjects.createRemaining(fakeConfig) - - def timeSource = new ControllableTimeSource() - - def traceConfig = Mock(TraceConfig) { - isDataStreamsEnabled() >> true - } - - when: - def dataStreams = new DefaultDataStreamsMonitoring(fakeConfig, sharedCommObjects, timeSource, { traceConfig }) - dataStreams.start() - dataStreams.add(new StatsPoint(DataStreamsTags.create(null, null), 9, 0, 10, timeSource.currentTimeNanos, 0, 0, 0, null)) - dataStreams.add(new StatsPoint(DataStreamsTags.create("testType", DataStreamsTags.Direction.INBOUND, "testTopic", "testGroup", null), 1, 2, 5, timeSource.currentTimeNanos, 0, 0, 0, null)) - dataStreams.trackBacklog(DataStreamsTags.createWithPartition("kafka_produce", "testTopic", "1", null, null), 100) - dataStreams.trackBacklog(DataStreamsTags.createWithPartition("kafka_produce", "testTopic", "1", null, null), 130) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS - 100l) - dataStreams.add(new StatsPoint(DataStreamsTags.create("testType", DataStreamsTags.Direction.INBOUND, "testTopic", "testGroup", null), 1, 2, 5, timeSource.currentTimeNanos, SECONDS.toNanos(10), SECONDS.toNanos(10), 10, null)) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.add(new StatsPoint(DataStreamsTags.create("testType", DataStreamsTags.Direction.INBOUND, "testTopic", "testGroup", null), 1, 2, 5, timeSource.currentTimeNanos, SECONDS.toNanos(5), SECONDS.toNanos(5), 5, null)) - dataStreams.add(new StatsPoint(DataStreamsTags.create("testType", DataStreamsTags.Direction.INBOUND, "testTopic2", "testGroup", null), 3, 4, 6, timeSource.currentTimeNanos, SECONDS.toNanos(2), 0, 2, null)) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.close() - - then: - conditions.eventually { - assert requestBodies.size() == 1 - } - - validateMessage(requestBodies[0], processTagsEnabled) - - cleanup: - injectSysConfig(EXPERIMENTAL_PROPAGATE_PROCESS_TAGS_ENABLED, "true") - ProcessTags.reset() - - where: - processTagsEnabled << [true, false] - } - - def "Write Kafka configs to mock server"() { - given: - def conditions = new PollingConditions(timeout: 2) - - def testOkhttpClient = OkHttpUtils.buildHttpClient(HttpUrl.get(server.address), 5000L) - - def features = Stub(DDAgentFeaturesDiscovery) { - supportsDataStreams() >> true - } - - def wellKnownTags = new WellKnownTags("runtimeid", "hostname", "test", Config.get().getServiceName(), "version", "java") - - def fakeConfig = Stub(Config) { - getAgentUrl() >> server.address.toString() - getWellKnownTags() >> wellKnownTags - getPrimaryTag() >> "region-1" - } - - def sharedCommObjects = new SharedCommunicationObjects() - sharedCommObjects.featuresDiscovery = features - sharedCommObjects.agentHttpClient = testOkhttpClient - sharedCommObjects.createRemaining(fakeConfig) - - def timeSource = new ControllableTimeSource() - - def traceConfig = Mock(TraceConfig) { - isDataStreamsEnabled() >> true - } - - when: - def dataStreams = new DefaultDataStreamsMonitoring(fakeConfig, sharedCommObjects, timeSource, { traceConfig }) - dataStreams.start() - - // Report a producer and consumer config - dataStreams.reportKafkaConfig("kafka_producer", "", "", ["bootstrap.servers": "localhost:9092", "acks": "all", "linger.ms": "5"]) - dataStreams.reportKafkaConfig("kafka_consumer", "", "test-group", ["bootstrap.servers": "localhost:9092", "group.id": "test-group", "auto.offset.reset": "earliest"]) - - // Also add a stats point so the bucket is not empty of stats - dataStreams.add(new StatsPoint(DataStreamsTags.create(null, null), 9, 0, 10, timeSource.currentTimeNanos, 0, 0, 0, null)) - - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.close() - - then: - conditions.eventually { - assert requestBodies.size() == 1 - } - - validateKafkaConfigMessage(requestBodies[0]) - } - - def "Duplicate Kafka configs are each serialized in the payload"() { - given: - def conditions = new PollingConditions(timeout: 2) - - def testOkhttpClient = OkHttpUtils.buildHttpClient(HttpUrl.get(server.address), 5000L) - - def features = Stub(DDAgentFeaturesDiscovery) { - supportsDataStreams() >> true - } - - def wellKnownTags = new WellKnownTags("runtimeid", "hostname", "test", Config.get().getServiceName(), "version", "java") - - def fakeConfig = Stub(Config) { - getAgentUrl() >> server.address.toString() - getWellKnownTags() >> wellKnownTags - getPrimaryTag() >> "region-1" - } - - def sharedCommObjects = new SharedCommunicationObjects() - sharedCommObjects.featuresDiscovery = features - sharedCommObjects.agentHttpClient = testOkhttpClient - sharedCommObjects.createRemaining(fakeConfig) - - def timeSource = new ControllableTimeSource() - - def traceConfig = Mock(TraceConfig) { - isDataStreamsEnabled() >> true - } - - when: - def dataStreams = new DefaultDataStreamsMonitoring(fakeConfig, sharedCommObjects, timeSource, { traceConfig }) - dataStreams.start() - - // Report the same producer config twice — both should be serialized - dataStreams.reportKafkaConfig("kafka_producer", "", "", ["bootstrap.servers": "localhost:9092", "acks": "all"]) - dataStreams.reportKafkaConfig("kafka_producer", "", "", ["bootstrap.servers": "localhost:9092", "acks": "all"]) - - // Also add a stats point so the bucket has content - dataStreams.add(new StatsPoint(DataStreamsTags.create(null, null), 9, 0, 10, timeSource.currentTimeNanos, 0, 0, 0, null)) - - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.close() - - then: - conditions.eventually { - assert requestBodies.size() == 1 - } - - validateDuplicateKafkaConfigMessage(requestBodies[0]) - } - - def validateKafkaConfigMessage(byte[] message) { - GzipSource gzipSource = new GzipSource(Okio.source(new ByteArrayInputStream(message))) - BufferedSource bufferedSource = Okio.buffer(gzipSource) - MessageUnpacker unpacker = MessagePack.newDefaultUnpacker(bufferedSource.inputStream()) - - // Outer map (same structure as other payloads) - def outerMapSize = unpacker.unpackMapHeader() - // Skip to Stats array - boolean foundStats = false - for (int i = 0; i < outerMapSize; i++) { - def key = unpacker.unpackString() - if (key == "Stats") { - foundStats = true - def numBuckets = unpacker.unpackArrayHeader() - assert numBuckets >= 1 - - // Parse first bucket - def bucketMapSize = unpacker.unpackMapHeader() - boolean foundConfigs = false - for (int j = 0; j < bucketMapSize; j++) { - def bucketKey = unpacker.unpackString() - if (bucketKey == "Configs") { - foundConfigs = true - def numConfigs = unpacker.unpackArrayHeader() - assert numConfigs == 2 - - // Collect configs in a map keyed by type - Map> configsByType = [:] - numConfigs.times { - assert unpacker.unpackMapHeader() == 4 - assert unpacker.unpackString() == "Type" - def type = unpacker.unpackString() - assert unpacker.unpackString() == "KafkaClusterId" - unpacker.unpackString() // skip cluster id value - assert unpacker.unpackString() == "ConsumerGroup" - unpacker.unpackString() // skip consumer group value - assert unpacker.unpackString() == "Config" - def configSize = unpacker.unpackMapHeader() - Map configEntries = [:] - configSize.times { - def ck = unpacker.unpackString() - def cv = unpacker.unpackString() - configEntries[ck] = cv - } - configsByType[type] = configEntries - } - - // Verify producer config - assert configsByType.containsKey("kafka_producer") - assert configsByType["kafka_producer"]["bootstrap.servers"] == "localhost:9092" - assert configsByType["kafka_producer"]["acks"] == "all" - assert configsByType["kafka_producer"]["linger.ms"] == "5" - - // Verify consumer config - assert configsByType.containsKey("kafka_consumer") - assert configsByType["kafka_consumer"]["bootstrap.servers"] == "localhost:9092" - assert configsByType["kafka_consumer"]["group.id"] == "test-group" - assert configsByType["kafka_consumer"]["auto.offset.reset"] == "earliest" - } else { - unpacker.skipValue() - } - } - assert foundConfigs : "Configs field not found in bucket" - - // Skip remaining buckets - for (int b = 1; b < numBuckets; b++) { - unpacker.skipValue() - } - } else { - unpacker.skipValue() - } - } - assert foundStats : "Stats field not found in payload" - return true - } - - def validateDuplicateKafkaConfigMessage(byte[] message) { - GzipSource gzipSource = new GzipSource(Okio.source(new ByteArrayInputStream(message))) - BufferedSource bufferedSource = Okio.buffer(gzipSource) - MessageUnpacker unpacker = MessagePack.newDefaultUnpacker(bufferedSource.inputStream()) - - def outerMapSize = unpacker.unpackMapHeader() - boolean foundStats = false - for (int i = 0; i < outerMapSize; i++) { - def key = unpacker.unpackString() - if (key == "Stats") { - foundStats = true - def numBuckets = unpacker.unpackArrayHeader() - assert numBuckets >= 1 - - // Parse first bucket - def bucketMapSize = unpacker.unpackMapHeader() - boolean foundConfigs = false - for (int j = 0; j < bucketMapSize; j++) { - def bucketKey = unpacker.unpackString() - if (bucketKey == "Configs") { - foundConfigs = true - def numConfigs = unpacker.unpackArrayHeader() - // Both configs should be present (no deduplication) - assert numConfigs == 2 - - numConfigs.times { - assert unpacker.unpackMapHeader() == 4 - assert unpacker.unpackString() == "Type" - assert unpacker.unpackString() == "kafka_producer" - assert unpacker.unpackString() == "KafkaClusterId" - unpacker.unpackString() // skip cluster id value - assert unpacker.unpackString() == "ConsumerGroup" - unpacker.unpackString() // skip consumer group value - assert unpacker.unpackString() == "Config" - def configSize = unpacker.unpackMapHeader() - Map configEntries = [:] - configSize.times { - def ck = unpacker.unpackString() - def cv = unpacker.unpackString() - configEntries[ck] = cv - } - assert configEntries["bootstrap.servers"] == "localhost:9092" - assert configEntries["acks"] == "all" - } - } else { - unpacker.skipValue() - } - } - assert foundConfigs : "Configs field not found in bucket" - - for (int b = 1; b < numBuckets; b++) { - unpacker.skipValue() - } - } else { - unpacker.skipValue() - } - } - assert foundStats : "Stats field not found in payload" - return true - } - - def validateMessage(byte[] message, boolean processTagsEnabled) { - GzipSource gzipSource = new GzipSource(Okio.source(new ByteArrayInputStream(message))) - - BufferedSource bufferedSource = Okio.buffer(gzipSource) - MessageUnpacker unpacker = MessagePack.newDefaultUnpacker(bufferedSource.inputStream()) - - assert unpacker.unpackMapHeader() == 8 + (processTagsEnabled ? 1 : 0) - assert unpacker.unpackString() == "Env" - assert unpacker.unpackString() == "test" - assert unpacker.unpackString() == "Service" - assert unpacker.unpackString() == Config.get().getServiceName() - assert unpacker.unpackString() == "Lang" - assert unpacker.unpackString() == "java" - assert unpacker.unpackString() == "PrimaryTag" - assert unpacker.unpackString() == "region-1" - assert unpacker.unpackString() == "TracerVersion" - assert unpacker.unpackString() == DDTraceCoreInfo.VERSION - assert unpacker.unpackString() == "Version" - assert unpacker.unpackString() == "version" - assert unpacker.unpackString() == "Stats" - assert unpacker.unpackArrayHeader() == 2 // 2 time buckets - - // FIRST BUCKET - assert unpacker.unpackMapHeader() == 4 - assert unpacker.unpackString() == "Start" - unpacker.skipValue() - assert unpacker.unpackString() == "Duration" - assert unpacker.unpackLong() == DEFAULT_BUCKET_DURATION_NANOS - assert unpacker.unpackString() == "Stats" - assert unpacker.unpackArrayHeader() == 2 // 2 groups in first bucket - - Set availableSizes = [5, 6] // we don't know the order the groups will be reported - 2.times { - int mapHeaderSize = unpacker.unpackMapHeader() - assert availableSizes.remove(mapHeaderSize) - if (mapHeaderSize == 5) { - // empty topic group - assert unpacker.unpackString() == "PathwayLatency" - unpacker.skipValue() - assert unpacker.unpackString() == "EdgeLatency" - unpacker.skipValue() - assert unpacker.unpackString() == "PayloadSize" - unpacker.skipValue() - assert unpacker.unpackString() == "Hash" - assert unpacker.unpackLong() == 9 - assert unpacker.unpackString() == "ParentHash" - assert unpacker.unpackLong() == 0 - } else { - //other group - assert unpacker.unpackString() == "PathwayLatency" - unpacker.skipValue() - assert unpacker.unpackString() == "EdgeLatency" - unpacker.skipValue() - assert unpacker.unpackString() == "PayloadSize" - unpacker.skipValue() - assert unpacker.unpackString() == "Hash" - assert unpacker.unpackLong() == 1 - assert unpacker.unpackString() == "ParentHash" - assert unpacker.unpackLong() == 2 - assert unpacker.unpackString() == "EdgeTags" - assert unpacker.unpackArrayHeader() == 4 - assert unpacker.unpackString() == "direction:in" - assert unpacker.unpackString() == "topic:testTopic" - assert unpacker.unpackString() == "type:testType" - assert unpacker.unpackString() == "group:testGroup" - } - } - - // Kafka stats - assert unpacker.unpackString() == "Backlogs" - assert unpacker.unpackArrayHeader() == 1 - assert unpacker.unpackMapHeader() == 2 - assert unpacker.unpackString() == "Tags" - assert unpacker.unpackArrayHeader() == 3 - assert unpacker.unpackString() == "topic:testTopic" - assert unpacker.unpackString() == "type:kafka_produce" - assert unpacker.unpackString() == "partition:1" - assert unpacker.unpackString() == "Value" - assert unpacker.unpackLong() == 130 - - // SECOND BUCKET - assert unpacker.unpackMapHeader() == 3 - assert unpacker.unpackString() == "Start" - unpacker.skipValue() - assert unpacker.unpackString() == "Duration" - assert unpacker.unpackLong() == DEFAULT_BUCKET_DURATION_NANOS - assert unpacker.unpackString() == "Stats" - assert unpacker.unpackArrayHeader() == 2 // 2 groups in second bucket - - Set availableHashes = [1L, 3L] // we don't know the order the groups will be reported - 2.times { - assert unpacker.unpackMapHeader() == 6 - assert unpacker.unpackString() == "PathwayLatency" - unpacker.skipValue() - assert unpacker.unpackString() == "EdgeLatency" - unpacker.skipValue() - assert unpacker.unpackString() == "PayloadSize" - unpacker.skipValue() - assert unpacker.unpackString() == "Hash" - def hash = unpacker.unpackLong() - assert availableHashes.remove(hash) - assert unpacker.unpackString() == "ParentHash" - assert unpacker.unpackLong() == (hash == 1 ? 2 : 4) - assert unpacker.unpackString() == "EdgeTags" - assert unpacker.unpackArrayHeader() == 4 - assert unpacker.unpackString() == "direction:in" - assert unpacker.unpackString() == (hash == 1 ? "topic:testTopic" : "topic:testTopic2") - assert unpacker.unpackString() == "type:testType" - assert unpacker.unpackString() == "group:testGroup" - } - - assert unpacker.unpackString() == "ProductMask" - assert unpacker.unpackLong() == 1 - - def processTags = ProcessTags.getTagsAsStringList() - assert unpacker.hasNext() == (processTags != null) - if (processTags != null) { - assert unpacker.unpackString() == "ProcessTags" - assert unpacker.unpackArrayHeader() == processTags.size() - processTags.each { - assert unpacker.unpackString() == it - } - } - - return true - } -} - diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/datastreams/DefaultDataStreamsMonitoringTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/datastreams/DefaultDataStreamsMonitoringTest.groovy deleted file mode 100644 index 217dd721d08..00000000000 --- a/dd-trace-core/src/test/groovy/datadog/trace/core/datastreams/DefaultDataStreamsMonitoringTest.groovy +++ /dev/null @@ -1,1513 +0,0 @@ -package datadog.trace.core.datastreams - -import datadog.communication.ddagent.DDAgentFeaturesDiscovery -import datadog.trace.api.Config -import datadog.trace.api.TraceConfig -import datadog.trace.api.datastreams.DataStreamsTags -import datadog.trace.api.datastreams.KafkaConfigReport -import datadog.trace.api.datastreams.SchemaRegistryUsage -import datadog.trace.api.datastreams.StatsPoint -import datadog.trace.api.experimental.DataStreamsContextCarrier -import datadog.trace.api.time.ControllableTimeSource -import datadog.trace.common.metrics.EventListener -import datadog.trace.common.metrics.Sink -import datadog.trace.core.test.DDCoreSpecification -import spock.util.concurrent.PollingConditions - -import java.util.concurrent.TimeUnit -import java.util.function.BiConsumer - -import static DefaultDataStreamsMonitoring.FEATURE_CHECK_INTERVAL_NANOS -import static java.util.concurrent.TimeUnit.SECONDS - -class DefaultDataStreamsMonitoringTest extends DDCoreSpecification { - static final DEFAULT_BUCKET_DURATION_NANOS = Config.get().getDataStreamsBucketDurationNanoseconds() - - def "No payloads written if data streams not supported or not enabled"() { - given: - def conditions = new PollingConditions(timeout: 1) - def features = Stub(DDAgentFeaturesDiscovery) { - supportsDataStreams() >> enabledAtAgent - } - def timeSource = new ControllableTimeSource() - def payloadWriter = Mock(DatastreamsPayloadWriter) - def sink = Mock(Sink) - - def traceConfig = Mock(TraceConfig) { - isDataStreamsEnabled() >> enabledInConfig - } - - when: - def dataStreams = new DefaultDataStreamsMonitoring(sink, features, timeSource, { traceConfig }, payloadWriter, DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.start() - dataStreams.add(new StatsPoint(DataStreamsTags.create("testType", null, "testTopic", "testGroup", null), 0, 0, 0, timeSource.currentTimeNanos, 0, 0, 0, null)) - dataStreams.report() - - then: - conditions.eventually { - assert dataStreams.inbox.isEmpty() - assert dataStreams.thread.state != Thread.State.RUNNABLE - } - 0 * payloadWriter.writePayload(_) - - cleanup: - dataStreams.close() - - where: - enabledAtAgent | enabledInConfig - false | true - true | false - false | false - } - - def "Schema sampler samples with correct weights"() { - given: - def features = Stub(DDAgentFeaturesDiscovery) { - supportsDataStreams() >> true - } - def timeSource = new ControllableTimeSource() - timeSource.set(1e12 as long) - def payloadWriter = Mock(DatastreamsPayloadWriter) - def sink = Mock(Sink) - def traceConfig = Mock(TraceConfig) { - isDataStreamsEnabled() >> true - } - - when: - def dataStreams = new DefaultDataStreamsMonitoring(sink, features, timeSource, { traceConfig }, payloadWriter, DEFAULT_BUCKET_DURATION_NANOS) - - then: - // the first received schema is sampled, with a weight of one. - dataStreams.canSampleSchema("schema1") - dataStreams.trySampleSchema("schema1") == 1 - // the sampling is done by topic, so a schema on a different topic will also be sampled at once, also with a weight of one. - dataStreams.canSampleSchema("schema2") - dataStreams.trySampleSchema("schema2") == 1 - // no time has passed from the last sampling, so the same schema is not sampled again (two times in a row). - !dataStreams.canSampleSchema("schema1") - !dataStreams.canSampleSchema("schema1") - timeSource.advance(30*1e9 as long) - // now, 30 seconds have passed, so the schema is sampled again, with a weight of 3 (so it includes the two times the schema was not sampled). - dataStreams.canSampleSchema("schema1") - dataStreams.trySampleSchema("schema1") == 3 - } - - def "Context carrier adapter test"() { - given: - def carrier = new CustomContextCarrier() - def keyName = "keyName" - def keyValue = "keyValue" - def extracted = "" - - when: - DataStreamsContextCarrierAdapter.INSTANCE.set(carrier, keyName, keyValue) - DataStreamsContextCarrierAdapter.INSTANCE.forEachKeyValue(carrier, new BiConsumer() { - @Override - void accept(String key, String value) { - if (key == keyName) { - extracted = value - } - } - }) - then: - extracted == keyValue - } - - def "Write group after a delay"() { - given: - def conditions = new PollingConditions(timeout: 1) - def features = Stub(DDAgentFeaturesDiscovery) { - supportsDataStreams() >> true - } - def timeSource = new ControllableTimeSource() - def sink = Mock(Sink) - def payloadWriter = new CapturingPayloadWriter() - - def traceConfig = Mock(TraceConfig) { - isDataStreamsEnabled() >> true - } - - when: - def dataStreams = new DefaultDataStreamsMonitoring(sink, features, timeSource, { traceConfig }, payloadWriter, DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.start() - def tg = DataStreamsTags.create("testType", null, "testTopic", "testGroup", null) - dataStreams.add(new StatsPoint(tg, 1, 2, 3, timeSource.currentTimeNanos, 0, 0, 0, null)) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.report() - - then: - conditions.eventually { - assert dataStreams.inbox.isEmpty() - assert dataStreams.thread.state != Thread.State.RUNNABLE - assert payloadWriter.buckets.size() == 1 - } - - with(payloadWriter.buckets.get(0)) { - groups.size() == 1 - - with(groups.iterator().next()) { - tags.type == "type:testType" - tags.group == "group:testGroup" - tags.topic == "topic:testTopic" - tags.nonNullSize() == 3 - hash == 1 - parentHash == 2 - } - } - - cleanup: - payloadWriter.close() - dataStreams.close() - } - - // This test relies on automatic reporting instead of manually calling report - def "SLOW Write group after a delay"() { - given: - def conditions = new PollingConditions(timeout: 1) - def features = Stub(DDAgentFeaturesDiscovery) { - supportsDataStreams() >> true - } - def timeSource = new ControllableTimeSource() - def sink = Mock(Sink) - def payloadWriter = new CapturingPayloadWriter() - def bucketDuration = TimeUnit.MILLISECONDS.toNanos(200) - - def traceConfig = Mock(TraceConfig) { - isDataStreamsEnabled() >> true - } - - when: - def dataStreams = new DefaultDataStreamsMonitoring(sink, features, timeSource, { traceConfig }, payloadWriter, bucketDuration) - dataStreams.start() - def tg = DataStreamsTags.create("testType", null, "testTopic", "testGroup", null) - dataStreams.add(new StatsPoint(tg, 1, 2, 3, timeSource.currentTimeNanos, 0, 0, 0, null)) - timeSource.advance(bucketDuration) - - then: - conditions.eventually { - assert dataStreams.inbox.isEmpty() - assert dataStreams.thread.state != Thread.State.RUNNABLE - assert payloadWriter.buckets.size() == 1 - } - - with(payloadWriter.buckets.get(0)) { - groups.size() == 1 - - with(groups.iterator().next()) { - tags.type == "type:testType" - tags.group == "group:testGroup" - tags.topic == "topic:testTopic" - tags.nonNullSize() == 3 - hash == 1 - parentHash == 2 - } - } - - cleanup: - payloadWriter.close() - dataStreams.close() - } - - def "Groups for current bucket are not reported"() { - given: - def conditions = new PollingConditions(timeout: 1) - def features = Stub(DDAgentFeaturesDiscovery) { - supportsDataStreams() >> true - } - def timeSource = new ControllableTimeSource() - def sink = Mock(Sink) - def payloadWriter = new CapturingPayloadWriter() - - def traceConfig = Mock(TraceConfig) { - isDataStreamsEnabled() >> true - } - - when: - def dataStreams = new DefaultDataStreamsMonitoring(sink, features, timeSource, { traceConfig }, payloadWriter, DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.start() - def tg = DataStreamsTags.create("testType", null, "testTopic", "testGroup", null) - dataStreams.add(new StatsPoint(tg, 1, 2, 3, timeSource.currentTimeNanos, 0, 0, 0, null)) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.add(new StatsPoint(tg, 3, 4, 3, timeSource.currentTimeNanos, 0, 0, 0, null)) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS - 100l) - dataStreams.report() - - then: - conditions.eventually { - assert dataStreams.inbox.isEmpty() - assert dataStreams.thread.state != Thread.State.RUNNABLE - assert payloadWriter.buckets.size() == 1 - } - - with(payloadWriter.buckets.get(0)) { - groups.size() == 1 - - with(groups.iterator().next()) { - tags.type == "type:testType" - tags.group == "group:testGroup" - tags.topic == "topic:testTopic" - tags.nonNullSize() == 3 - hash == 1 - parentHash == 2 - } - } - - cleanup: - payloadWriter.close() - dataStreams.close() - } - - def "All groups written in close"() { - given: - def conditions = new PollingConditions(timeout: 1) - def features = Stub(DDAgentFeaturesDiscovery) { - supportsDataStreams() >> true - } - def timeSource = new ControllableTimeSource() - def sink = Mock(Sink) - def payloadWriter = new CapturingPayloadWriter() - - def traceConfig = Mock(TraceConfig) { - isDataStreamsEnabled() >> true - } - - when: - def dataStreams = new DefaultDataStreamsMonitoring(sink, features, timeSource, { traceConfig }, payloadWriter, DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.start() - def tg = DataStreamsTags.create("testType", null, "testTopic", "testGroup", null) - def tg2 = DataStreamsTags.create("testType", null, "testTopic2", "testGroup", null) - dataStreams.add(new StatsPoint(tg, 1, 2, 5, timeSource.currentTimeNanos, 0, 0, 0, null)) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.add(new StatsPoint(tg2, 3, 4, 6, timeSource.currentTimeNanos, 0, 0, 0, null)) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS - 100l) - dataStreams.close() - - then: - conditions.eventually { - assert dataStreams.inbox.isEmpty() - assert dataStreams.thread.state != Thread.State.RUNNABLE - assert payloadWriter.buckets.size() == 2 - } - - with(payloadWriter.buckets.get(0)) { - groups.size() == 1 - - with(groups.iterator().next()) { - tags.type == "type:testType" - tags.group == "group:testGroup" - tags.topic == "topic:testTopic" - tags.nonNullSize() == 3 - hash == 1 - parentHash == 2 - } - } - - with(payloadWriter.buckets.get(1)) { - groups.size() == 1 - - with(groups.iterator().next()) { - tags.type == "type:testType" - tags.group == "group:testGroup" - tags.topic == "topic:testTopic2" - tags.nonNullSize() == 3 - hash == 3 - parentHash == 4 - } - } - - cleanup: - payloadWriter.close() - } - - def "Kafka offsets are tracked"() { - given: - def conditions = new PollingConditions(timeout: 1) - def features = Stub(DDAgentFeaturesDiscovery) { - supportsDataStreams() >> true - } - def timeSource = new ControllableTimeSource() - def sink = Mock(Sink) - def payloadWriter = new CapturingPayloadWriter() - - def traceConfig = Mock(TraceConfig) { - isDataStreamsEnabled() >> true - } - - when: - def dataStreams = new DefaultDataStreamsMonitoring(sink, features, timeSource, { traceConfig }, payloadWriter, DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.start() - dataStreams.trackBacklog(DataStreamsTags.createWithPartition("kafka_commit", "testTopic", "2", null, "testGroup"), 23) - dataStreams.trackBacklog(DataStreamsTags.createWithPartition("kafka_commit", "testTopic", "2", null, "testGroup"), 24) - dataStreams.trackBacklog(DataStreamsTags.createWithPartition("kafka_produce", "testTopic", "2", null, null), 23) - dataStreams.trackBacklog(DataStreamsTags.createWithPartition("kafka_produce", "testTopic2", "2", null, null), 23) - dataStreams.trackBacklog(DataStreamsTags.createWithPartition("kafka_produce", "testTopic", "2", null, null), 45) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.report() - - then: - conditions.eventually { - assert dataStreams.inbox.isEmpty() - assert dataStreams.thread.state != Thread.State.RUNNABLE - assert payloadWriter.buckets.size() == 1 - } - - with(payloadWriter.buckets.get(0)) { - backlogs.size() == 3 - def list = backlogs.sort({ it.key.toString() }) - with(list[0]) { - it.key == DataStreamsTags.createWithPartition("kafka_commit", "testTopic", "2", null, "testGroup") - it.value == 24 - } - with(list[1]) { - it.key == DataStreamsTags.createWithPartition("kafka_produce", "testTopic", "2", null, null) - it.value == 45 - } - with(list[2]) { - it.key == DataStreamsTags.createWithPartition("kafka_produce", "testTopic2", "2", null, null) - it.value == 23 - } - } - - cleanup: - payloadWriter.close() - dataStreams.close() - } - - def "Groups from multiple buckets are reported"() { - given: - def conditions = new PollingConditions(timeout: 1) - def features = Stub(DDAgentFeaturesDiscovery) { - supportsDataStreams() >> true - } - def timeSource = new ControllableTimeSource() - def sink = Mock(Sink) - def payloadWriter = new CapturingPayloadWriter() - - def traceConfig = Mock(TraceConfig) { - isDataStreamsEnabled() >> true - } - - when: - def dataStreams = new DefaultDataStreamsMonitoring(sink, features, timeSource, { traceConfig }, payloadWriter, DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.start() - def tg = DataStreamsTags.create("testType", null, "testTopic", "testGroup", null) - dataStreams.add(new StatsPoint(tg, 1, 2, 5, timeSource.currentTimeNanos, 0, 0, 0, null)) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS*10) - def tg2 = DataStreamsTags.create("testType", null, "testTopic2", "testGroup", null) - dataStreams.add(new StatsPoint(tg2, 3, 4, 6, timeSource.currentTimeNanos, 0, 0, 0, null)) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.report() - - then: - conditions.eventually { - assert dataStreams.inbox.isEmpty() - assert dataStreams.thread.state != Thread.State.RUNNABLE - assert payloadWriter.buckets.size() == 2 - } - - with(payloadWriter.buckets.get(0)) { - groups.size() == 1 - groups - with(groups.iterator().next()) { - tags.nonNullSize() == 3 - tags.getType() == "type:testType" - tags.getGroup() == "group:testGroup" - tags.getTopic() == "topic:testTopic" - hash == 1 - parentHash == 2 - } - } - - with(payloadWriter.buckets.get(1)) { - groups.size() == 1 - - with(groups.iterator().next()) { - tags.getType() == "type:testType" - tags.getGroup() == "group:testGroup" - tags.getTopic() == "topic:testTopic2" - tags.nonNullSize() == 3 - hash == 3 - parentHash == 4 - } - } - - cleanup: - payloadWriter.close() - dataStreams.close() - } - - def "Multiple points are correctly grouped in multiple buckets"() { - given: - def conditions = new PollingConditions(timeout: 1) - def features = Stub(DDAgentFeaturesDiscovery) { - supportsDataStreams() >> true - } - def timeSource = new ControllableTimeSource() - def sink = Mock(Sink) - def payloadWriter = new CapturingPayloadWriter() - - def traceConfig = Mock(TraceConfig) { - isDataStreamsEnabled() >> true - } - - when: - def tg = DataStreamsTags.create("testType", null, "testTopic", "testGroup", null) - def dataStreams = new DefaultDataStreamsMonitoring(sink, features, timeSource, { traceConfig }, payloadWriter, DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.start() - dataStreams.add(new StatsPoint(tg, 1, 2, 1, timeSource.currentTimeNanos, 0, 0, 0, null)) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS - 100l) - dataStreams.add(new StatsPoint(tg, 1, 2, 1, timeSource.currentTimeNanos, SECONDS.toNanos(10), SECONDS.toNanos(10), 10, null)) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.add(new StatsPoint(tg, 1, 2,1, timeSource.currentTimeNanos, SECONDS.toNanos(5), SECONDS.toNanos(5), 5, null)) - dataStreams.add(new StatsPoint(tg, 3, 4, 5, timeSource.currentTimeNanos, SECONDS.toNanos(2), 0, 0, null)) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.report() - - then: - conditions.eventually { - assert dataStreams.inbox.isEmpty() - assert dataStreams.thread.state != Thread.State.RUNNABLE - assert payloadWriter.buckets.size() == 2 - } - - with(payloadWriter.buckets.get(0)) { - groups.size() == 1 - - with(groups.iterator().next()) { - tags.type == "type:testType" - tags.group == "group:testGroup" - tags.topic == "topic:testTopic" - tags.nonNullSize() == 3 - hash == 1 - parentHash == 2 - Math.abs((pathwayLatency.getMaxValue()-10)/10) < 0.01 - } - } - - with(payloadWriter.buckets.get(1)) { - groups.size() == 2 - - List sortedGroups = new ArrayList<>(groups) - sortedGroups.sort({ it.hash }) - - with(sortedGroups[0]) { - hash == 1 - parentHash == 2 - tags.type == "type:testType" - tags.group == "group:testGroup" - tags.topic == "topic:testTopic" - tags.nonNullSize() == 3 - Math.abs((pathwayLatency.getMaxValue()-5)/5) < 0.01 - } - - with(sortedGroups[1]) { - hash == 3 - parentHash == 4 - tags.type == "type:testType" - tags.group == "group:testGroup" - tags.topic == "topic:testTopic" - tags.nonNullSize() == 3 - Math.abs((pathwayLatency.getMaxValue()-2)/2) < 0.01 - } - } - - cleanup: - payloadWriter.close() - dataStreams.close() - } - - def "feature upgrade"() { - given: - def conditions = new PollingConditions(timeout: 1) - boolean supportsDataStreaming = false - def features = Mock(DDAgentFeaturesDiscovery) { - supportsDataStreams() >> { return supportsDataStreaming } - } - def timeSource = new ControllableTimeSource() - def sink = Mock(Sink) - def payloadWriter = new CapturingPayloadWriter() - - def traceConfig = Mock(TraceConfig) { - isDataStreamsEnabled() >> true - } - - when: "reporting points when data streams is not supported" - def dataStreams = new DefaultDataStreamsMonitoring(sink, features, timeSource, { traceConfig }, payloadWriter, DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.start() - def tg = DataStreamsTags.create("testType", null, "testTopic", "testGroup", null) - dataStreams.add(new StatsPoint(tg, 1, 2, 3, timeSource.currentTimeNanos, 0, 0, 0, null)) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.report() - - then: "no buckets are reported" - conditions.eventually { - assert dataStreams.inbox.isEmpty() - assert dataStreams.thread.state != Thread.State.RUNNABLE - } - - payloadWriter.buckets.isEmpty() - - when: "report called multiple times without advancing past check interval" - dataStreams.report() - dataStreams.report() - dataStreams.report() - - then: "features are not rechecked" - conditions.eventually { - assert dataStreams.inbox.isEmpty() - assert dataStreams.thread.state != Thread.State.RUNNABLE - } - 0 * features.discover() - payloadWriter.buckets.isEmpty() - - when: "submitting points after an upgrade" - supportsDataStreaming = true - timeSource.advance(FEATURE_CHECK_INTERVAL_NANOS) - dataStreams.report() - - dataStreams.add(new StatsPoint(tg, 1, 2, 3, timeSource.currentTimeNanos, 0, 0, 0, null)) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.report() - - then: "points are now reported" - conditions.eventually { - assert dataStreams.inbox.isEmpty() - assert payloadWriter.buckets.size() == 1 - } - - with(payloadWriter.buckets.get(0)) { - groups.size() == 1 - - with(groups.iterator().next()) { - tags.type == "type:testType" - tags.group == "group:testGroup" - tags.topic == "topic:testTopic" - tags.nonNullSize() == 3 - hash == 1 - parentHash == 2 - } - } - - cleanup: - payloadWriter.close() - dataStreams.close() - } - - def "feature downgrade then upgrade"() { - given: - def conditions = new PollingConditions(timeout: 1) - boolean supportsDataStreaming = true - def features = Stub(DDAgentFeaturesDiscovery) { - supportsDataStreams() >> { return supportsDataStreaming } - } - def timeSource = new ControllableTimeSource() - def sink = Mock(Sink) - def payloadWriter = new CapturingPayloadWriter() - - def traceConfig = Mock(TraceConfig) { - isDataStreamsEnabled() >> true - } - - when: "reporting points after a downgrade" - def dataStreams = new DefaultDataStreamsMonitoring(sink, features, timeSource, { traceConfig }, payloadWriter, DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.start() - supportsDataStreaming = false - dataStreams.onEvent(EventListener.EventType.DOWNGRADED, "") - def tg = DataStreamsTags.create("testType", null, "testTopic", "testGroup", null) - dataStreams.add(new StatsPoint(tg, 1, 2, 3, timeSource.currentTimeNanos, 0, 0, 0, null)) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.report() - - then: "no buckets are reported" - conditions.eventually { - assert dataStreams.inbox.isEmpty() - assert dataStreams.thread.state != Thread.State.RUNNABLE - } - - payloadWriter.buckets.isEmpty() - - when: "submitting points after an upgrade" - supportsDataStreaming = true - timeSource.advance(FEATURE_CHECK_INTERVAL_NANOS) - dataStreams.report() - - dataStreams.add(new StatsPoint(tg, 1, 2, 3, timeSource.currentTimeNanos, 0, 0, 0, null)) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.report() - - then: "points are now reported" - conditions.eventually { - assert dataStreams.inbox.isEmpty() - assert payloadWriter.buckets.size() == 1 - } - - with(payloadWriter.buckets.get(0)) { - groups.size() == 1 - - with(groups.iterator().next()) { - tags.type == "type:testType" - tags.group == "group:testGroup" - tags.topic == "topic:testTopic" - tags.nonNullSize() == 3 - hash == 1 - parentHash == 2 - } - } - - cleanup: - payloadWriter.close() - dataStreams.close() - } - - def "dynamic config enable and disable"() { - given: - def conditions = new PollingConditions(timeout: 1) - boolean supportsDataStreaming = true - def features = Mock(DDAgentFeaturesDiscovery) { - supportsDataStreams() >> { return supportsDataStreaming } - } - def timeSource = new ControllableTimeSource() - def sink = Mock(Sink) - def payloadWriter = new CapturingPayloadWriter() - - boolean dsmEnabled = false - def traceConfig = Mock(TraceConfig) { - isDataStreamsEnabled() >> { return dsmEnabled } - } - - when: "reporting points when data streams is not enabled" - def tg = DataStreamsTags.create("testType", null, "testTopic", "testGroup", null) - def dataStreams = new DefaultDataStreamsMonitoring(sink, features, timeSource, { traceConfig }, payloadWriter, DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.start() - dataStreams.add(new StatsPoint(tg, 1, 2, 3, timeSource.currentTimeNanos, 0, 0, 0, null)) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.report() - - then: "no buckets are reported" - conditions.eventually { - assert dataStreams.inbox.isEmpty() - assert dataStreams.thread.state != Thread.State.RUNNABLE - } - - payloadWriter.buckets.isEmpty() - - when: "submitting points after dynamically enabled" - dsmEnabled = true - dataStreams.report() - - dataStreams.add(new StatsPoint(tg, 1, 2, 3, timeSource.currentTimeNanos, 0, 0, 0, null)) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.report() - - then: "points are now reported" - conditions.eventually { - assert dataStreams.inbox.isEmpty() - assert payloadWriter.buckets.size() == 1 - } - - with(payloadWriter.buckets.get(0)) { - groups.size() == 1 - - with(groups.iterator().next()) { - tags.type == "type:testType" - tags.group == "group:testGroup" - tags.topic == "topic:testTopic" - tags.nonNullSize() == 3 - hash == 1 - parentHash == 2 - } - } - - when: "disabling data streams dynamically" - dsmEnabled = false - dataStreams.report() - - then: "inbox is processed" - conditions.eventually { - assert dataStreams.inbox.isEmpty() - } - - when: "submitting points after being disabled" - payloadWriter.buckets.clear() - - dataStreams.add(new StatsPoint(tg, 1, 2, 3, timeSource.currentTimeNanos, 0, 0, 0, null)) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.report() - - then: "points are no longer reported" - conditions.eventually { - assert dataStreams.inbox.isEmpty() - assert payloadWriter.buckets.isEmpty() - } - - cleanup: - payloadWriter.close() - dataStreams.close() - } - - def "feature and dynamic config upgrade interactions"() { - given: - def conditions = new PollingConditions(timeout: 1) - boolean supportsDataStreaming = false - def features = Mock(DDAgentFeaturesDiscovery) { - supportsDataStreams() >> { return supportsDataStreaming } - } - def timeSource = new ControllableTimeSource() - def sink = Mock(Sink) - def payloadWriter = new CapturingPayloadWriter() - - boolean dsmEnabled = false - def traceConfig = Mock(TraceConfig) { - isDataStreamsEnabled() >> { return dsmEnabled } - } - - when: "reporting points when data streams is not supported" - def dataStreams = new DefaultDataStreamsMonitoring(sink, features, timeSource, { traceConfig }, payloadWriter, DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.start() - def tg = DataStreamsTags.create("testType", null, "testTopic", "testGroup", null) - dataStreams.add(new StatsPoint(tg, 1, 2, 3, timeSource.currentTimeNanos, 0, 0, 0, null)) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.report() - - then: "no buckets are reported" - conditions.eventually { - assert dataStreams.inbox.isEmpty() - assert dataStreams.thread.state != Thread.State.RUNNABLE - } - - payloadWriter.buckets.isEmpty() - - when: "submitting points after an upgrade with dsm disabled" - supportsDataStreaming = true - timeSource.advance(FEATURE_CHECK_INTERVAL_NANOS) - dataStreams.report() - - dataStreams.add(new StatsPoint(tg, 1, 2, 3, timeSource.currentTimeNanos, 0, 0, 0, null)) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.report() - - then: "points are not reported" - conditions.eventually { - assert dataStreams.inbox.isEmpty() - assert payloadWriter.buckets.isEmpty() - } - - when: "dsm is enabled dynamically" - dsmEnabled = true - dataStreams.report() - - dataStreams.add(new StatsPoint(tg, 1, 2, 3, timeSource.currentTimeNanos, 0, 0, 0, null)) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.report() - - then: "points are now reported" - conditions.eventually { - assert dataStreams.inbox.isEmpty() - assert payloadWriter.buckets.size() == 1 - } - - with(payloadWriter.buckets.get(0)) { - groups.size() == 1 - - with(groups.iterator().next()) { - tags.type == "type:testType" - tags.group == "group:testGroup" - tags.topic == "topic:testTopic" - tags.nonNullSize() == 3 - hash == 1 - parentHash == 2 - } - } - - cleanup: - payloadWriter.close() - dataStreams.close() - } - - def "more feature and dynamic config upgrade interactions"() { - given: - def conditions = new PollingConditions(timeout: 1) - boolean supportsDataStreaming = false - def features = Mock(DDAgentFeaturesDiscovery) { - supportsDataStreams() >> { return supportsDataStreaming } - } - def timeSource = new ControllableTimeSource() - def sink = Mock(Sink) - def payloadWriter = new CapturingPayloadWriter() - - boolean dsmEnabled = false - def traceConfig = Mock(TraceConfig) { - isDataStreamsEnabled() >> { return dsmEnabled } - } - - when: "reporting points when data streams is not supported" - def dataStreams = new DefaultDataStreamsMonitoring(sink, features, timeSource, { traceConfig }, payloadWriter, DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.start() - def tg = DataStreamsTags.create("testType", null, "testTopic", "testGroup", null) - dataStreams.add(new StatsPoint(tg, 1, 2, 3, timeSource.currentTimeNanos, 0, 0, 0, null)) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.report() - - then: "no buckets are reported" - conditions.eventually { - assert dataStreams.inbox.isEmpty() - assert dataStreams.thread.state != Thread.State.RUNNABLE - } - - payloadWriter.buckets.isEmpty() - - when: "enabling dsm when not supported by agent" - dsmEnabled = true - dataStreams.report() - - dataStreams.add(new StatsPoint(tg, 1, 2, 3, timeSource.currentTimeNanos, 0, 0, 0, null)) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.report() - - then: "points are not reported" - conditions.eventually { - assert dataStreams.inbox.isEmpty() - assert payloadWriter.buckets.isEmpty() - } - - cleanup: - payloadWriter.close() - dataStreams.close() - } - - def "Schema registry usages are aggregated by operation"() { - given: - def conditions = new PollingConditions(timeout: 2) - def features = Stub(DDAgentFeaturesDiscovery) { - supportsDataStreams() >> true - } - def timeSource = new ControllableTimeSource() - def payloadWriter = new CapturingPayloadWriter() - def sink = Mock(Sink) - def traceConfig = Mock(TraceConfig) { - isDataStreamsEnabled() >> true - } - - when: - def dataStreams = new DefaultDataStreamsMonitoring(sink, features, timeSource, { traceConfig }, payloadWriter, DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.start() - - // Record serialize and deserialize operations - dataStreams.reportSchemaRegistryUsage("test-topic", "test-cluster", 123, true, false, "serialize") - dataStreams.reportSchemaRegistryUsage("test-topic", "test-cluster", 123, true, false, "serialize") // duplicate serialize - dataStreams.reportSchemaRegistryUsage("test-topic", "test-cluster", 123, true, false, "deserialize") - dataStreams.reportSchemaRegistryUsage("test-topic", "test-cluster", 456, true, true, "serialize") // different schema/key - - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.report() - - then: - conditions.eventually { - assert dataStreams.inbox.isEmpty() - assert dataStreams.thread.state != Thread.State.RUNNABLE - assert payloadWriter.buckets.size() == 1 - } - - with(payloadWriter.buckets.get(0)) { - schemaRegistryUsages.size() == 3 // 3 unique combinations - - // Find serialize operation for schema 123 (should have count 2) - def serializeUsage = schemaRegistryUsages.find { e -> - e.key.schemaId == 123 && e.key.operation == "serialize" && !e.key.isKey - } - serializeUsage != null - serializeUsage.value == 2L // Aggregated 2 serialize operations - - // Find deserialize operation for schema 123 (should have count 1) - def deserializeUsage = schemaRegistryUsages.find { e -> - e.key.schemaId == 123 && e.key.operation == "deserialize" && !e.key.isKey - } - deserializeUsage != null - deserializeUsage.value == 1L - - // Find serialize operation for schema 456 with isKey=true (should have count 1) - def keySerializeUsage = schemaRegistryUsages.find { e -> - e.key.schemaId == 456 && e.key.operation == "serialize" && e.key.isKey - } - keySerializeUsage != null - keySerializeUsage.value == 1L - } - - cleanup: - payloadWriter.close() - dataStreams.close() - } - - def "SchemaKey equals and hashCode work correctly"() { - given: - def key1 = new StatsBucket.SchemaKey("topic1", "cluster1", 123, true, false, "serialize") - def key2 = new StatsBucket.SchemaKey("topic1", "cluster1", 123, true, false, "serialize") - def key3 = new StatsBucket.SchemaKey("topic2", "cluster1", 123, true, false, "serialize") // different topic - def key4 = new StatsBucket.SchemaKey("topic1", "cluster2", 123, true, false, "serialize") // different cluster - def key5 = new StatsBucket.SchemaKey("topic1", "cluster1", 456, true, false, "serialize") // different schema - def key6 = new StatsBucket.SchemaKey("topic1", "cluster1", 123, false, false, "serialize") // different success - def key7 = new StatsBucket.SchemaKey("topic1", "cluster1", 123, true, true, "serialize") // different isKey - def key8 = new StatsBucket.SchemaKey("topic1", "cluster1", 123, true, false, "deserialize") // different operation - - expect: - // Reflexive - key1.equals(key1) - key1.hashCode() == key1.hashCode() - - // Symmetric - key1.equals(key2) - key2.equals(key1) - key1.hashCode() == key2.hashCode() - - // Different topic - !key1.equals(key3) - !key3.equals(key1) - - // Different cluster - !key1.equals(key4) - !key4.equals(key1) - - // Different schema ID - !key1.equals(key5) - !key5.equals(key1) - - // Different success - !key1.equals(key6) - !key6.equals(key1) - - // Different isKey - !key1.equals(key7) - !key7.equals(key1) - - // Different operation - !key1.equals(key8) - !key8.equals(key1) - - // Null check - !key1.equals(null) - - // Different class - !key1.equals("not a schema key") - } - - def "StatsBucket aggregates schema registry usages correctly"() { - given: - def bucket = new StatsBucket(1000L, 10000L) - def usage1 = new SchemaRegistryUsage("topic1", "cluster1", 123, true, false, "serialize", 1000L, null) - def usage2 = new SchemaRegistryUsage("topic1", "cluster1", 123, true, false, "serialize", 2000L, null) - def usage3 = new SchemaRegistryUsage("topic1", "cluster1", 123, true, false, "deserialize", 3000L, null) - - when: - bucket.addSchemaRegistryUsage(usage1) - bucket.addSchemaRegistryUsage(usage2) // should increment count for same key - bucket.addSchemaRegistryUsage(usage3) // different operation, new key - - def usages = bucket.getSchemaRegistryUsages() - def usageMap = usages.collectEntries { [(it.key): it.value] } - - then: - usages.size() == 2 - - // Check serialize count - def serializeKey = new StatsBucket.SchemaKey("topic1", "cluster1", 123, true, false, "serialize") - usageMap[serializeKey] == 2L - - // Check deserialize count - def deserializeKey = new StatsBucket.SchemaKey("topic1", "cluster1", 123, true, false, "deserialize") - usageMap[deserializeKey] == 1L - - // Check that different operations create different keys - serializeKey != deserializeKey - } - - def "Kafka producer config is reported in bucket"() { - given: - def conditions = new PollingConditions(timeout: 1) - def features = Stub(DDAgentFeaturesDiscovery) { - supportsDataStreams() >> true - } - def timeSource = new ControllableTimeSource() - def sink = Mock(Sink) - def payloadWriter = new CapturingPayloadWriter() - - def traceConfig = Mock(TraceConfig) { - isDataStreamsEnabled() >> true - } - - when: - def dataStreams = new DefaultDataStreamsMonitoring(sink, features, timeSource, { traceConfig }, payloadWriter, DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.start() - dataStreams.reportKafkaConfig("kafka_producer", "", "", ["bootstrap.servers": "localhost:9092", "acks": "all"]) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.report() - - then: - conditions.eventually { - assert dataStreams.inbox.isEmpty() - assert dataStreams.thread.state != Thread.State.RUNNABLE - assert payloadWriter.buckets.size() == 1 - } - - with(payloadWriter.buckets.get(0)) { - kafkaConfigs.size() == 1 - with(kafkaConfigs.get(0)) { - type == "kafka_producer" - config["bootstrap.servers"] == "localhost:9092" - config["acks"] == "all" - } - } - - cleanup: - payloadWriter.close() - dataStreams.close() - } - - def "Kafka consumer config is reported in bucket"() { - given: - def conditions = new PollingConditions(timeout: 1) - def features = Stub(DDAgentFeaturesDiscovery) { - supportsDataStreams() >> true - } - def timeSource = new ControllableTimeSource() - def sink = Mock(Sink) - def payloadWriter = new CapturingPayloadWriter() - - def traceConfig = Mock(TraceConfig) { - isDataStreamsEnabled() >> true - } - - when: - def dataStreams = new DefaultDataStreamsMonitoring(sink, features, timeSource, { traceConfig }, payloadWriter, DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.start() - dataStreams.reportKafkaConfig("kafka_consumer", "", "test-group", ["bootstrap.servers": "localhost:9092", "group.id": "test-group", "auto.offset.reset": "earliest"]) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.report() - - then: - conditions.eventually { - assert dataStreams.inbox.isEmpty() - assert dataStreams.thread.state != Thread.State.RUNNABLE - assert payloadWriter.buckets.size() == 1 - } - - with(payloadWriter.buckets.get(0)) { - kafkaConfigs.size() == 1 - with(kafkaConfigs.get(0)) { - type == "kafka_consumer" - config["bootstrap.servers"] == "localhost:9092" - config["group.id"] == "test-group" - config["auto.offset.reset"] == "earliest" - } - } - - cleanup: - payloadWriter.close() - dataStreams.close() - } - - def "Duplicate Kafka configs are each reported in the bucket"() { - given: - def conditions = new PollingConditions(timeout: 1) - def features = Stub(DDAgentFeaturesDiscovery) { - supportsDataStreams() >> true - } - def timeSource = new ControllableTimeSource() - def sink = Mock(Sink) - def payloadWriter = new CapturingPayloadWriter() - - def traceConfig = Mock(TraceConfig) { - isDataStreamsEnabled() >> true - } - - when: "reporting the same config twice" - def dataStreams = new DefaultDataStreamsMonitoring(sink, features, timeSource, { traceConfig }, payloadWriter, DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.start() - def config1 = ["bootstrap.servers": "localhost:9092", "acks": "all"] - def config2 = ["bootstrap.servers": "localhost:9092", "acks": "all"] - dataStreams.reportKafkaConfig("kafka_producer", "", "", config1) - dataStreams.reportKafkaConfig("kafka_producer", "", "", config2) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.report() - - then: "both configs are reported in the bucket" - conditions.eventually { - assert dataStreams.inbox.isEmpty() - assert dataStreams.thread.state != Thread.State.RUNNABLE - assert payloadWriter.buckets.size() == 1 - } - - with(payloadWriter.buckets.get(0)) { - kafkaConfigs.size() == 2 - kafkaConfigs.every { it.type == "kafka_producer" } - kafkaConfigs.every { it.config["bootstrap.servers"] == "localhost:9092" } - kafkaConfigs.every { it.config["acks"] == "all" } - } - - cleanup: - payloadWriter.close() - dataStreams.close() - } - - def "Kafka configs reported in separate buckets appear in each bucket"() { - given: - def conditions = new PollingConditions(timeout: 1) - def features = Stub(DDAgentFeaturesDiscovery) { - supportsDataStreams() >> true - } - def timeSource = new ControllableTimeSource() - def sink = Mock(Sink) - def payloadWriter = new CapturingPayloadWriter() - - def traceConfig = Mock(TraceConfig) { - isDataStreamsEnabled() >> true - } - - when: "reporting a config in the first bucket" - def dataStreams = new DefaultDataStreamsMonitoring(sink, features, timeSource, { traceConfig }, payloadWriter, DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.start() - def config = ["bootstrap.servers": "localhost:9092", "acks": "all"] - dataStreams.reportKafkaConfig("kafka_producer", "", "", config) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.report() - - then: "first bucket has the config" - conditions.eventually { - assert dataStreams.inbox.isEmpty() - assert dataStreams.thread.state != Thread.State.RUNNABLE - assert payloadWriter.buckets.size() == 1 - } - - with(payloadWriter.buckets.get(0)) { - kafkaConfigs.size() == 1 - } - - when: "reporting the same config again in a new bucket" - payloadWriter.buckets.clear() - dataStreams.reportKafkaConfig("kafka_producer", "", "", config) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.report() - - then: "second bucket also has the config" - conditions.eventually { - assert dataStreams.inbox.isEmpty() - assert dataStreams.thread.state != Thread.State.RUNNABLE - assert payloadWriter.buckets.size() == 1 - } - - with(payloadWriter.buckets.get(0)) { - kafkaConfigs.size() == 1 - with(kafkaConfigs.get(0)) { - type == "kafka_producer" - config["bootstrap.servers"] == "localhost:9092" - config["acks"] == "all" - } - } - - cleanup: - payloadWriter.close() - dataStreams.close() - } - - def "Different Kafka configs are both reported"() { - given: - def conditions = new PollingConditions(timeout: 1) - def features = Stub(DDAgentFeaturesDiscovery) { - supportsDataStreams() >> true - } - def timeSource = new ControllableTimeSource() - def sink = Mock(Sink) - def payloadWriter = new CapturingPayloadWriter() - - def traceConfig = Mock(TraceConfig) { - isDataStreamsEnabled() >> true - } - - when: "reporting producer and consumer configs" - def dataStreams = new DefaultDataStreamsMonitoring(sink, features, timeSource, { traceConfig }, payloadWriter, DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.start() - dataStreams.reportKafkaConfig("kafka_producer", "", "", ["bootstrap.servers": "localhost:9092", "acks": "all"]) - dataStreams.reportKafkaConfig("kafka_consumer", "", "my-group", ["bootstrap.servers": "localhost:9092", "group.id": "my-group"]) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.report() - - then: "both configs are reported" - conditions.eventually { - assert dataStreams.inbox.isEmpty() - assert dataStreams.thread.state != Thread.State.RUNNABLE - assert payloadWriter.buckets.size() == 1 - } - - with(payloadWriter.buckets.get(0)) { - kafkaConfigs.size() == 2 - - def producerConfig = kafkaConfigs.find { it.type == "kafka_producer" } - producerConfig != null - producerConfig.config["bootstrap.servers"] == "localhost:9092" - producerConfig.config["acks"] == "all" - - def consumerConfig = kafkaConfigs.find { it.type == "kafka_consumer" } - consumerConfig != null - consumerConfig.config["bootstrap.servers"] == "localhost:9092" - consumerConfig.config["group.id"] == "my-group" - } - - cleanup: - payloadWriter.close() - dataStreams.close() - } - - def "Kafka configs with different values for same type are not deduplicated"() { - given: - def conditions = new PollingConditions(timeout: 1) - def features = Stub(DDAgentFeaturesDiscovery) { - supportsDataStreams() >> true - } - def timeSource = new ControllableTimeSource() - def sink = Mock(Sink) - def payloadWriter = new CapturingPayloadWriter() - - def traceConfig = Mock(TraceConfig) { - isDataStreamsEnabled() >> true - } - - when: "reporting two producer configs with different settings" - def dataStreams = new DefaultDataStreamsMonitoring(sink, features, timeSource, { traceConfig }, payloadWriter, DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.start() - dataStreams.reportKafkaConfig("kafka_producer", "", "", ["bootstrap.servers": "localhost:9092", "acks": "all"]) - dataStreams.reportKafkaConfig("kafka_producer", "", "", ["bootstrap.servers": "localhost:9093", "acks": "1"]) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.report() - - then: "both configs are reported because they have different values" - conditions.eventually { - assert dataStreams.inbox.isEmpty() - assert dataStreams.thread.state != Thread.State.RUNNABLE - assert payloadWriter.buckets.size() == 1 - } - - with(payloadWriter.buckets.get(0)) { - kafkaConfigs.size() == 2 - } - - cleanup: - payloadWriter.close() - dataStreams.close() - } - - def "Kafka configs are reported alongside stats points"() { - given: - def conditions = new PollingConditions(timeout: 1) - def features = Stub(DDAgentFeaturesDiscovery) { - supportsDataStreams() >> true - } - def timeSource = new ControllableTimeSource() - def sink = Mock(Sink) - def payloadWriter = new CapturingPayloadWriter() - - def traceConfig = Mock(TraceConfig) { - isDataStreamsEnabled() >> true - } - - when: "reporting both stats points and kafka configs" - def dataStreams = new DefaultDataStreamsMonitoring(sink, features, timeSource, { traceConfig }, payloadWriter, DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.start() - def tg = DataStreamsTags.create("testType", null, "testTopic", "testGroup", null) - dataStreams.add(new StatsPoint(tg, 1, 2, 3, timeSource.currentTimeNanos, 0, 0, 0, null)) - dataStreams.reportKafkaConfig("kafka_producer", "", "", ["bootstrap.servers": "localhost:9092"]) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.report() - - then: "bucket contains both stats groups and kafka configs" - conditions.eventually { - assert dataStreams.inbox.isEmpty() - assert dataStreams.thread.state != Thread.State.RUNNABLE - assert payloadWriter.buckets.size() == 1 - } - - with(payloadWriter.buckets.get(0)) { - groups.size() == 1 - kafkaConfigs.size() == 1 - - with(groups.iterator().next()) { - tags.type == "type:testType" - hash == 1 - parentHash == 2 - } - - with(kafkaConfigs.get(0)) { - type == "kafka_producer" - config["bootstrap.servers"] == "localhost:9092" - } - } - - cleanup: - payloadWriter.close() - dataStreams.close() - } - - def "Kafka configs not reported when DSM is disabled"() { - given: - def conditions = new PollingConditions(timeout: 1) - def features = Stub(DDAgentFeaturesDiscovery) { - supportsDataStreams() >> enabledAtAgent - } - def timeSource = new ControllableTimeSource() - def payloadWriter = new CapturingPayloadWriter() - def sink = Mock(Sink) - - def traceConfig = Mock(TraceConfig) { - isDataStreamsEnabled() >> enabledInConfig - } - - when: - def dataStreams = new DefaultDataStreamsMonitoring(sink, features, timeSource, { traceConfig }, payloadWriter, DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.start() - dataStreams.reportKafkaConfig("kafka_producer", "", "", ["bootstrap.servers": "localhost:9092"]) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.report() - - then: - conditions.eventually { - assert dataStreams.inbox.isEmpty() - assert dataStreams.thread.state != Thread.State.RUNNABLE - } - payloadWriter.buckets.isEmpty() - - cleanup: - payloadWriter.close() - dataStreams.close() - - where: - enabledAtAgent | enabledInConfig - false | true - true | false - false | false - } - - def "Kafka configs flushed on close"() { - given: - def conditions = new PollingConditions(timeout: 1) - def features = Stub(DDAgentFeaturesDiscovery) { - supportsDataStreams() >> true - } - def timeSource = new ControllableTimeSource() - def sink = Mock(Sink) - def payloadWriter = new CapturingPayloadWriter() - - def traceConfig = Mock(TraceConfig) { - isDataStreamsEnabled() >> true - } - - when: - def dataStreams = new DefaultDataStreamsMonitoring(sink, features, timeSource, { traceConfig }, payloadWriter, DEFAULT_BUCKET_DURATION_NANOS) - dataStreams.start() - dataStreams.reportKafkaConfig("kafka_producer", "", "", ["bootstrap.servers": "localhost:9092"]) - timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS - 100l) - dataStreams.close() - - then: "configs in the current bucket are flushed on close" - conditions.eventually { - assert dataStreams.inbox.isEmpty() - assert dataStreams.thread.state != Thread.State.RUNNABLE - assert payloadWriter.buckets.size() == 1 - } - - with(payloadWriter.buckets.get(0)) { - kafkaConfigs.size() == 1 - with(kafkaConfigs.get(0)) { - type == "kafka_producer" - config["bootstrap.servers"] == "localhost:9092" - } - } - - cleanup: - payloadWriter.close() - } - - def "KafkaConfigReport equals and hashCode work correctly"() { - given: - def config1 = new KafkaConfigReport("kafka_producer", "", "", ["bootstrap.servers": "localhost:9092", "acks": "all"], 1000L, null) - def config2 = new KafkaConfigReport("kafka_producer", "", "", ["bootstrap.servers": "localhost:9092", "acks": "all"], 2000L, null) - def config3 = new KafkaConfigReport("kafka_consumer", "", "", ["bootstrap.servers": "localhost:9092", "acks": "all"], 1000L, null) - def config4 = new KafkaConfigReport("kafka_producer", "", "", ["bootstrap.servers": "localhost:9093"], 1000L, null) - def config5 = new KafkaConfigReport("kafka_producer", "", "", ["bootstrap.servers": "localhost:9092", "acks": "all"], 1000L, "other-service") - - expect: - // Reflexive - config1.equals(config1) - config1.hashCode() == config1.hashCode() - - // Same type and config, different timestamp -- equals (timestamp is NOT part of equals) - config1.equals(config2) - config2.equals(config1) - config1.hashCode() == config2.hashCode() - - // Same type and config, different serviceNameOverride -- equals (serviceNameOverride is NOT part of equals) - config1.equals(config5) - config5.equals(config1) - config1.hashCode() == config5.hashCode() - - // Different type - !config1.equals(config3) - !config3.equals(config1) - - // Different config values - !config1.equals(config4) - !config4.equals(config1) - - // Null check - !config1.equals(null) - - // Different class - !config1.equals("not a config report") - } - - def "StatsBucket stores Kafka configs"() { - given: - def bucket = new StatsBucket(1000L, 10000L) - def config1 = new KafkaConfigReport("kafka_producer", "", "", ["acks": "all"], 1000L, null) - def config2 = new KafkaConfigReport("kafka_consumer", "", "test", ["group.id": "test"], 2000L, null) - - when: - bucket.addKafkaConfig(config1) - bucket.addKafkaConfig(config2) - - then: - bucket.getKafkaConfigs().size() == 2 - bucket.getKafkaConfigs().get(0).type == "kafka_producer" - bucket.getKafkaConfigs().get(0).config["acks"] == "all" - bucket.getKafkaConfigs().get(1).type == "kafka_consumer" - bucket.getKafkaConfigs().get(1).config["group.id"] == "test" - } -} - -class CapturingPayloadWriter implements DatastreamsPayloadWriter { - boolean accepting = true - List buckets = new ArrayList<>() - - void writePayload(Collection payload, String serviceNameOverride) { - if (accepting) { - buckets.addAll(payload) - } - } - - void close() { - // Stop accepting new buckets so any late submissions by the reporting thread aren't seen - accepting = false - } -} - -class CustomContextCarrier implements DataStreamsContextCarrier { - - private Map data = new HashMap<>() - - @Override - Set> entries() { - return data.entrySet() - } - - @Override - void set(String key, String value) { - data.put(key, value) - } -} diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/datastreams/DefaultPathwayContextTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/datastreams/DefaultPathwayContextTest.groovy deleted file mode 100644 index 43f75f5052b..00000000000 --- a/dd-trace-core/src/test/groovy/datadog/trace/core/datastreams/DefaultPathwayContextTest.groovy +++ /dev/null @@ -1,653 +0,0 @@ -package datadog.trace.core.datastreams - -import datadog.communication.ddagent.DDAgentFeaturesDiscovery -import datadog.trace.api.BaseHash -import datadog.trace.api.Config -import datadog.trace.api.DDTraceId -import datadog.trace.api.TagMap -import datadog.trace.api.TraceConfig -import datadog.trace.api.datastreams.DataStreamsTags -import datadog.trace.api.datastreams.StatsPoint -import datadog.trace.api.time.ControllableTimeSource -import datadog.trace.bootstrap.instrumentation.api.AgentPropagation -import datadog.trace.bootstrap.instrumentation.api.AgentSpan -import datadog.trace.bootstrap.instrumentation.api.AgentTracer -import datadog.trace.common.metrics.Sink -import datadog.trace.core.propagation.ExtractedContext -import datadog.trace.core.test.DDCoreSpecification -import java.util.function.Consumer -import static datadog.context.Context.root -import static datadog.trace.api.TracePropagationStyle.DATADOG -import static datadog.trace.api.datastreams.DataStreamsContext.create -import static datadog.trace.api.datastreams.DataStreamsContext.fromTags -import static datadog.trace.api.datastreams.PathwayContext.PROPAGATION_KEY_BASE64 -import static java.util.concurrent.TimeUnit.MILLISECONDS - -class DefaultPathwayContextTest extends DDCoreSpecification { - long baseHash = 12 - - static final DEFAULT_BUCKET_DURATION_NANOS = Config.get().getDataStreamsBucketDurationNanoseconds() - def pointConsumer = new Consumer() { - List points = [] - - @Override - void accept(StatsPoint point) { - points.add(point) - } - } - - void verifyFirstPoint(StatsPoint point) { - assert point.parentHash == 0 - assert point.pathwayLatencyNano == 0 - assert point.edgeLatencyNano == 0 - assert point.payloadSizeBytes == 0 - } - - def "First Set checkpoint starts the context."() { - given: - def timeSource = new ControllableTimeSource() - def context = new DefaultPathwayContext(timeSource, null) - - when: - timeSource.advance(50) - context.setCheckpoint(fromTags(DataStreamsTags.create("internal", null)), pointConsumer) - - then: - context.isStarted() - pointConsumer.points.size() == 1 - verifyFirstPoint(pointConsumer.points[0]) - } - - def "Checkpoint generated"() { - given: - def timeSource = new ControllableTimeSource() - def context = new DefaultPathwayContext(timeSource, null) - - when: - timeSource.advance(50) - context.setCheckpoint(fromTags(DataStreamsTags.create("internal", DataStreamsTags.Direction.INBOUND)), pointConsumer) - timeSource.advance(25) - def tags = DataStreamsTags.create("kafka", DataStreamsTags.Direction.OUTBOUND, "topic", "group", null) - context.setCheckpoint(fromTags(tags), pointConsumer) - - then: - context.isStarted() - pointConsumer.points.size() == 2 - verifyFirstPoint(pointConsumer.points[0]) - with(pointConsumer.points[1]) { - tags.group == "group:group" - tags.topic == "topic:topic" - tags.type == "type:kafka" - tags.getDirection() == "direction:out" - tags.nonNullSize() == 4 - parentHash == pointConsumer.points[0].hash - hash != 0 - pathwayLatencyNano == 25 - edgeLatencyNano == 25 - } - } - - def "Checkpoint with payload size"() { - given: - def timeSource = new ControllableTimeSource() - def context = new DefaultPathwayContext(timeSource, null) - - when: - timeSource.advance(25) - context.setCheckpoint( - create(DataStreamsTags.create("kafka", null, "topic", "group", null), 0, 72), - pointConsumer) - - then: - context.isStarted() - pointConsumer.points.size() == 1 - with(pointConsumer.points[0]) { - tags.getGroup() == "group:group" - tags.getTopic() == "topic:topic" - tags.getType() == "type:kafka" - tags.nonNullSize() == 3 - hash != 0 - payloadSizeBytes == 72 - } - } - - def "Multiple checkpoints generated"() { - given: - def timeSource = new ControllableTimeSource() - def context = new DefaultPathwayContext(timeSource, null) - - when: - timeSource.advance(50) - context.setCheckpoint(fromTags(DataStreamsTags.create("kafka", DataStreamsTags.Direction.OUTBOUND)), pointConsumer) - timeSource.advance(25) - def tg = DataStreamsTags.create("kafka", DataStreamsTags.Direction.INBOUND, "topic", "group", null) - context.setCheckpoint(fromTags(tg), pointConsumer) - timeSource.advance(30) - context.setCheckpoint(fromTags(tg), pointConsumer) - - then: - context.isStarted() - pointConsumer.points.size() == 3 - verifyFirstPoint(pointConsumer.points[0]) - with(pointConsumer.points[1]) { - tags.nonNullSize() == 4 - tags.direction == "direction:in" - tags.group == "group:group" - tags.topic == "topic:topic" - tags.type == "type:kafka" - parentHash == pointConsumer.points[0].hash - hash != 0 - pathwayLatencyNano == 25 - edgeLatencyNano == 25 - } - with(pointConsumer.points[2]) { - tags.nonNullSize() == 4 - tags.direction == "direction:in" - tags.group == "group:group" - tags.topic == "topic:topic" - tags.type == "type:kafka" - // this point should have the first point as parent, - // as the loop protection will reset the parent if two identical - // points (same hash for tag values) are about to form a hierarchy - parentHash == pointConsumer.points[0].hash - hash != 0 - pathwayLatencyNano == 55 - edgeLatencyNano == 30 - } - } - - def "Exception thrown when trying to encode unstarted context"() { - given: - def timeSource = new ControllableTimeSource() - def context = new DefaultPathwayContext(timeSource, null) - - when: - context.encode() - - then: - thrown(IllegalStateException) - } - - def "Set checkpoint with dataset tags"() { - given: - def timeSource = new ControllableTimeSource() - def context = new DefaultPathwayContext(timeSource, null) - - when: - timeSource.advance(MILLISECONDS.toNanos(50)) - context.setCheckpoint(fromTags(DataStreamsTags.createWithDataset("s3", DataStreamsTags.Direction.INBOUND, null, "my_object.csv", "my_bucket")), pointConsumer) - def encoded = context.encode() - timeSource.advance(MILLISECONDS.toNanos(2)) - def decodedContext = DefaultPathwayContext.decode(timeSource, null, encoded) - timeSource.advance(MILLISECONDS.toNanos(25)) - def tg = DataStreamsTags.createWithDataset("s3", DataStreamsTags.Direction.OUTBOUND, null, "my_object.csv", "my_bucket") - context.setCheckpoint(fromTags(tg), pointConsumer) - - then: - decodedContext.isStarted() - pointConsumer.points.size() == 2 - - // all points should have datasetHash, which is not equal to hash or 0 - for (def i = 0; i < pointConsumer.points.size(); i++){ - pointConsumer.points[i].aggregationHash != pointConsumer.points[i].hash - pointConsumer.points[i].aggregationHash != 0 - } - } - - def "Encoding and decoding (base64) a context"() { - // Timesource needs to be advanced in milliseconds because encoding truncates to millis - given: - def timeSource = new ControllableTimeSource() - def context = new DefaultPathwayContext(timeSource, null) - - when: - timeSource.advance(MILLISECONDS.toNanos(50)) - context.setCheckpoint(fromTags(DataStreamsTags.create("internal", DataStreamsTags.Direction.INBOUND)), pointConsumer) - def encoded = context.encode() - timeSource.advance(MILLISECONDS.toNanos(2)) - def decodedContext = DefaultPathwayContext.decode(timeSource, null, encoded) - timeSource.advance(MILLISECONDS.toNanos(25)) - context.setCheckpoint(fromTags(DataStreamsTags.create("kafka", null, "topic", "group", null)), pointConsumer) - - then: - decodedContext.isStarted() - pointConsumer.points.size() == 2 - - with(pointConsumer.points[1]) { - tags.nonNullSize() == 3 - tags.getGroup() == "group:group" - tags.getType() == "type:kafka" - tags.getTopic() == "topic:topic" - parentHash == pointConsumer.points[0].hash - hash != 0 - pathwayLatencyNano == MILLISECONDS.toNanos(27) - edgeLatencyNano == MILLISECONDS.toNanos(27) - } - } - - def "Set checkpoint with timestamp"() { - given: - def timeSource = new ControllableTimeSource() - def context = new DefaultPathwayContext(timeSource, null) - def timeFromQueue = timeSource.getCurrentTimeMillis() - 200 - when: - context.setCheckpoint(create(DataStreamsTags.create("internal", null), timeFromQueue, 0), pointConsumer) - then: - context.isStarted() - pointConsumer.points.size() == 1 - with(pointConsumer.points[0]) { - tags.getType() == "type:internal" - tags.nonNullSize() == 1 - parentHash == 0 - hash != 0 - pathwayLatencyNano == MILLISECONDS.toNanos(200) - edgeLatencyNano == MILLISECONDS.toNanos(200) - } - } - - def "Encoding and decoding (base64) with contexts and checkpoints"() { - // Timesource needs to be advanced in milliseconds because encoding truncates to millis - given: - def timeSource = new ControllableTimeSource() - def context = new DefaultPathwayContext(timeSource, null) - - when: - timeSource.advance(MILLISECONDS.toNanos(50)) - context.setCheckpoint(fromTags(DataStreamsTags.create("internal", DataStreamsTags.Direction.INBOUND)), pointConsumer) - - def encoded = context.encode() - timeSource.advance(MILLISECONDS.toNanos(1)) - def decodedContext = DefaultPathwayContext.decode(timeSource, null, encoded) - timeSource.advance(MILLISECONDS.toNanos(25)) - context.setCheckpoint(fromTags(DataStreamsTags.create("kafka", DataStreamsTags.Direction.OUTBOUND, "topic", "group", null)), pointConsumer) - - then: - decodedContext.isStarted() - pointConsumer.points.size() == 2 - with(pointConsumer.points[1]) { - tags.group == "group:group" - tags.topic == "topic:topic" - tags.type == "type:kafka" - tags.direction == "direction:out" - tags.nonNullSize() == 4 - parentHash == pointConsumer.points[0].hash - hash != 0 - pathwayLatencyNano == MILLISECONDS.toNanos(26) - edgeLatencyNano == MILLISECONDS.toNanos(26) - } - - when: - def secondEncode = decodedContext.encode() - timeSource.advance(MILLISECONDS.toNanos(2)) - def secondDecode = DefaultPathwayContext.decode(timeSource, null, secondEncode) - timeSource.advance(MILLISECONDS.toNanos(30)) - context.setCheckpoint(fromTags(DataStreamsTags.create("kafka", DataStreamsTags.Direction.INBOUND, "topicB", "group", null)), pointConsumer) - - then: - secondDecode.isStarted() - pointConsumer.points.size() == 3 - with(pointConsumer.points[2]) { - tags.group == "group:group" - tags.topic == "topic:topicB" - tags.type == "type:kafka" - tags.direction == "direction:in" - tags.nonNullSize() == 4 - parentHash == pointConsumer.points[1].hash - hash != 0 - pathwayLatencyNano == MILLISECONDS.toNanos(58) - edgeLatencyNano == MILLISECONDS.toNanos(32) - } - } - - def "Encoding and decoding (base64) with injects and extracts"() { - // Timesource needs to be advanced in milliseconds because encoding truncates to millis - given: - def timeSource = new ControllableTimeSource() - def context = new DefaultPathwayContext(timeSource, null) - def contextVisitor = new Base64MapContextVisitor() - - when: - timeSource.advance(MILLISECONDS.toNanos(50)) - context.setCheckpoint(fromTags(DataStreamsTags.create("internal", DataStreamsTags.Direction.INBOUND)), pointConsumer) - - def encoded = context.encode() - Map carrier = [(PROPAGATION_KEY_BASE64): encoded, "someotherkey": "someothervalue"] - timeSource.advance(MILLISECONDS.toNanos(1)) - def decodedContext = DefaultPathwayContext.extract(carrier, contextVisitor, timeSource, null) - timeSource.advance(MILLISECONDS.toNanos(25)) - context.setCheckpoint(fromTags(DataStreamsTags.create("kafka", DataStreamsTags.Direction.OUTBOUND, "topic", "group", null)), pointConsumer) - - then: - decodedContext.isStarted() - pointConsumer.points.size() == 2 - with(pointConsumer.points[1]) { - tags.nonNullSize() == 4 - tags.group == "group:group" - tags.topic == "topic:topic" - tags.type == "type:kafka" - tags.direction == "direction:out" - parentHash == pointConsumer.points[0].hash - hash != 0 - pathwayLatencyNano == MILLISECONDS.toNanos(26) - edgeLatencyNano == MILLISECONDS.toNanos(26) - } - - when: - def secondEncode = decodedContext.encode() - carrier = [(PROPAGATION_KEY_BASE64): secondEncode] - timeSource.advance(MILLISECONDS.toNanos(2)) - def secondDecode = DefaultPathwayContext.extract(carrier, contextVisitor, timeSource, null) - timeSource.advance(MILLISECONDS.toNanos(30)) - context.setCheckpoint(fromTags(DataStreamsTags.create("kafka", DataStreamsTags.Direction.INBOUND, "topicB", "group", null)), pointConsumer) - - then: - secondDecode.isStarted() - pointConsumer.points.size() == 3 - with(pointConsumer.points[2]) { - tags.nonNullSize() == 4 - tags.group == "group:group" - tags.topic == "topic:topicB" - tags.type == "type:kafka" - tags.direction == "direction:in" - parentHash == pointConsumer.points[1].hash - hash != 0 - pathwayLatencyNano == MILLISECONDS.toNanos(58) - edgeLatencyNano == MILLISECONDS.toNanos(32) - } - } - - def "Encoding and decoding (SQS-formatted) with injects and extracts"() { - // Timesource needs to be advanced in milliseconds because encoding truncates to millis - given: - def timeSource = new ControllableTimeSource() - def context = new DefaultPathwayContext(timeSource, null) - def contextVisitor = new Base64MapContextVisitor() - - when: - timeSource.advance(MILLISECONDS.toNanos(50)) - context.setCheckpoint(fromTags(DataStreamsTags.create("internal", DataStreamsTags.Direction.INBOUND)), pointConsumer) - - def encoded = context.encode() - Map carrier = [(PROPAGATION_KEY_BASE64): encoded, "someotherkey": "someothervalue"] - timeSource.advance(MILLISECONDS.toNanos(1)) - def decodedContext = DefaultPathwayContext.extract(carrier, contextVisitor, timeSource, null) - timeSource.advance(MILLISECONDS.toNanos(25)) - context.setCheckpoint(fromTags(DataStreamsTags.create("sqs", DataStreamsTags.Direction.OUTBOUND, "topic", null, null)), pointConsumer) - - then: - decodedContext.isStarted() - pointConsumer.points.size() == 2 - with(pointConsumer.points[1]) { - tags.direction == "direction:out" - tags.topic == "topic:topic" - tags.type == "type:sqs" - tags.nonNullSize() == 3 - parentHash == pointConsumer.points[0].hash - hash != 0 - pathwayLatencyNano == MILLISECONDS.toNanos(26) - edgeLatencyNano == MILLISECONDS.toNanos(26) - } - - when: - def secondEncode = decodedContext.encode() - carrier = [(PROPAGATION_KEY_BASE64): secondEncode] - timeSource.advance(MILLISECONDS.toNanos(2)) - def secondDecode = DefaultPathwayContext.extract(carrier, contextVisitor, timeSource, null) - timeSource.advance(MILLISECONDS.toNanos(30)) - context.setCheckpoint(fromTags(DataStreamsTags.create("sqs", DataStreamsTags.Direction.INBOUND, "topicB", null, null)), pointConsumer) - - then: - secondDecode.isStarted() - pointConsumer.points.size() == 3 - with(pointConsumer.points[2]) { - tags.type == "type:sqs" - tags.topic == "topic:topicB" - tags.nonNullSize() == 3 - parentHash == pointConsumer.points[1].hash - hash != 0 - pathwayLatencyNano == MILLISECONDS.toNanos(58) - edgeLatencyNano == MILLISECONDS.toNanos(32) - } - } - - def "Empty tags not set"() { - given: - def timeSource = new ControllableTimeSource() - def context = new DefaultPathwayContext(timeSource, null) - - when: - timeSource.advance(50) - context.setCheckpoint(fromTags(DataStreamsTags.create("internal", DataStreamsTags.Direction.INBOUND)), pointConsumer) - timeSource.advance(25) - context.setCheckpoint(fromTags(DataStreamsTags.create("type", DataStreamsTags.Direction.OUTBOUND, "topic", "group", null)), pointConsumer) - timeSource.advance(25) - context.setCheckpoint(fromTags(DataStreamsTags.create(null, null)), pointConsumer) - - then: - context.isStarted() - pointConsumer.points.size() == 3 - verifyFirstPoint(pointConsumer.points[0]) - with(pointConsumer.points[1]) { - tags.type == "type:type" - tags.topic == "topic:topic" - tags.group == "group:group" - tags.direction == "direction:out" - tags.nonNullSize() == 4 - parentHash == pointConsumer.points[0].hash - hash != 0 - pathwayLatencyNano == 25 - edgeLatencyNano == 25 - } - with(pointConsumer.points[2]) { - tags.nonNullSize() == 0 - parentHash == pointConsumer.points[1].hash - hash != 0 - pathwayLatencyNano == 50 - edgeLatencyNano == 25 - } - } - - def "Check context extractor decorator behavior"() { - given: - def sink = Mock(Sink) - def features = Stub(DDAgentFeaturesDiscovery) { - supportsDataStreams() >> true - } - def timeSource = new ControllableTimeSource() - def payloadWriter = Mock(DatastreamsPayloadWriter) - - def globalTraceConfig = Mock(TraceConfig) { - isDataStreamsEnabled() >> { return dynamicConfigEnabled } - } - - def tracerApi = Mock(AgentTracer.TracerAPI) { - captureTraceConfig() >> globalTraceConfig - } - AgentTracer.TracerAPI originalTracer = AgentTracer.get() - AgentTracer.forceRegister(tracerApi) - - def dataStreams = new DefaultDataStreamsMonitoring(sink, features, timeSource, { globalTraceConfig }, payloadWriter, DEFAULT_BUCKET_DURATION_NANOS) - - BaseHash.updateBaseHash(baseHash) - def context = new DefaultPathwayContext(timeSource, null) - timeSource.advance(MILLISECONDS.toNanos(50)) - context.setCheckpoint(fromTags(DataStreamsTags.create("internal", DataStreamsTags.Direction.INBOUND)), pointConsumer) - def encoded = context.encode() - Map carrier = [ - (PROPAGATION_KEY_BASE64): encoded, - "someotherkey": "someothervalue" - ] - def contextVisitor = new Base64MapContextVisitor() - def propagator = dataStreams.propagator() - - when: - def extractedContext = propagator.extract(root(), carrier, contextVisitor) - def extractedSpan = AgentSpan.fromContext(extractedContext) - - then: - encoded == "L+lDG/Pa9hRkZA==" - !dynamicConfigEnabled || extractedSpan != null - if (dynamicConfigEnabled) { - def extracted = extractedSpan.context() - assert extracted != null - assert extracted.pathwayContext != null - assert extracted.pathwayContext.isStarted() - } - - cleanup: - AgentTracer.forceRegister(originalTracer) - - where: - dynamicConfigEnabled << [true, false] - } - - def "Check context extractor decorator behavior when trace data is null"() { - given: - def sink = Mock(Sink) - def features = Stub(DDAgentFeaturesDiscovery) { - supportsDataStreams() >> true - } - def timeSource = new ControllableTimeSource() - def payloadWriter = Mock(DatastreamsPayloadWriter) - - def globalTraceConfig = Mock(TraceConfig) { - isDataStreamsEnabled() >> { return globalDsmEnabled } - } - - def tracerApi = Mock(AgentTracer.TracerAPI) { - captureTraceConfig() >> globalTraceConfig - } - AgentTracer.TracerAPI originalTracer = AgentTracer.get() - AgentTracer.forceRegister(tracerApi) - - def dataStreams = new DefaultDataStreamsMonitoring(sink, features, timeSource, { globalTraceConfig }, payloadWriter, DEFAULT_BUCKET_DURATION_NANOS) - - BaseHash.updateBaseHash(baseHash) - def context = new DefaultPathwayContext(timeSource, null) - timeSource.advance(MILLISECONDS.toNanos(50)) - context.setCheckpoint(fromTags(DataStreamsTags.create("internal", DataStreamsTags.Direction.INBOUND)), pointConsumer) - def encoded = context.encode() - - Map carrier = [(PROPAGATION_KEY_BASE64): encoded, "someotherkey": "someothervalue"] - def contextVisitor = new Base64MapContextVisitor() - def propagator = dataStreams.propagator() - - when: - def extractedContext = propagator.extract(root(), carrier, contextVisitor) - def extractedSpan = AgentSpan.fromContext(extractedContext) - - then: - encoded == "L+lDG/Pa9hRkZA==" - if (globalDsmEnabled) { - extractedSpan != null - def extracted = extractedSpan.context() - extracted != null - extracted.pathwayContext != null - extracted.pathwayContext.isStarted() - } else { - extractedSpan == null - } - - cleanup: - AgentTracer.forceRegister(originalTracer) - - where: - globalDsmEnabled << [true, false] - } - - def "Check context extractor decorator behavior when local trace config is null"() { - given: - def sink = Mock(Sink) - def features = Stub(DDAgentFeaturesDiscovery) { - supportsDataStreams() >> true - } - def timeSource = new ControllableTimeSource() - def payloadWriter = Mock(DatastreamsPayloadWriter) - - def globalTraceConfig = Mock(TraceConfig) { - isDataStreamsEnabled() >> { return globalDsmEnabled } - } - - def tracerApi = Mock(AgentTracer.TracerAPI) { - captureTraceConfig() >> globalTraceConfig - } - AgentTracer.TracerAPI originalTracer = AgentTracer.get() - AgentTracer.forceRegister(tracerApi) - - def dataStreams = new DefaultDataStreamsMonitoring(sink, features, timeSource, { globalTraceConfig }, payloadWriter, DEFAULT_BUCKET_DURATION_NANOS) - - BaseHash.updateBaseHash(baseHash) - def context = new DefaultPathwayContext(timeSource, null) - timeSource.advance(MILLISECONDS.toNanos(50)) - context.setCheckpoint(fromTags(DataStreamsTags.create("internal", DataStreamsTags.Direction.INBOUND)), pointConsumer) - def encoded = context.encode() - Map carrier = [(PROPAGATION_KEY_BASE64): encoded, "someotherkey": "someothervalue"] - def contextVisitor = new Base64MapContextVisitor() - def spanContext = new ExtractedContext(DDTraceId.ONE, 1, 0, null, 0, - null, (TagMap)null, null, null, globalTraceConfig, DATADOG) - def baseContext = AgentSpan.fromSpanContext(spanContext).storeInto(root()) - def propagator = dataStreams.propagator() - - when: - def extractedContext = propagator.extract(baseContext, carrier, contextVisitor) - def extractedSpan = AgentSpan.fromContext(extractedContext) - - then: - extractedSpan != null - - when: - def extracted = extractedSpan.context() - - then: - extracted != null - encoded == "L+lDG/Pa9hRkZA==" - if (globalDsmEnabled) { - extracted.pathwayContext != null - extracted.pathwayContext.isStarted() - } else { - extracted.pathwayContext == null - } - - cleanup: - AgentTracer.forceRegister(originalTracer) - - where: - globalDsmEnabled << [true, false] - } - - def "Check context extractor decorator behavior when trace data and dsm data are null"() { - given: - def sink = Mock(Sink) - def features = Stub(DDAgentFeaturesDiscovery) { - supportsDataStreams() >> true - } - def timeSource = new ControllableTimeSource() - def payloadWriter = Mock(DatastreamsPayloadWriter) - - def traceConfig = Mock(TraceConfig) { - isDataStreamsEnabled() >> true - } - - def dataStreams = new DefaultDataStreamsMonitoring(sink, features, timeSource, { traceConfig }, payloadWriter, DEFAULT_BUCKET_DURATION_NANOS) - - Map carrier = ["someotherkey": "someothervalue"] - def contextVisitor = new Base64MapContextVisitor() - def propagator = dataStreams.propagator() - - when: - def extractedContext = propagator.extract(root(), carrier, contextVisitor) - def extractedSpan = AgentSpan.fromContext(extractedContext) - - then: - extractedSpan == null - } - - class Base64MapContextVisitor implements AgentPropagation.ContextVisitor> { - @Override - void forEachKey(Map carrier, AgentPropagation.KeyClassifier classifier) { - for (Map.Entry entry : carrier.entrySet()) { - classifier.accept(entry.key, entry.value) - } - } - } -} diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/datastreams/SchemaBuilderTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/datastreams/SchemaBuilderTest.groovy deleted file mode 100644 index b18e56dcf9c..00000000000 --- a/dd-trace-core/src/test/groovy/datadog/trace/core/datastreams/SchemaBuilderTest.groovy +++ /dev/null @@ -1,44 +0,0 @@ -package datadog.trace.core.datastreams - -import datadog.trace.bootstrap.instrumentation.api.Schema -import datadog.trace.bootstrap.instrumentation.api.SchemaIterator -import datadog.trace.core.test.DDCoreSpecification - -class SchemaBuilderTest extends DDCoreSpecification { - - class Iterator implements SchemaIterator{ - - @Override - void iterateOverSchema(datadog.trace.bootstrap.instrumentation.api.SchemaBuilder builder) { - HashMap extension = new HashMap(1) - extension.put("x-test-extension-1", "hello") - extension.put("x-test-extension-2", "world") - builder.addProperty("person", "name", false, "string", "name of the person", null, null, null, null) - builder.addProperty("person", "phone_numbers", true, "string", null, null, null, null, null) - builder.addProperty("person", "person_name", false, "string", null, null, null, null, null) - builder.addProperty("person", "address", false, "object", null, "#/components/schemas/address", null, null, null) - builder.addProperty("address", "zip", false, "number", null, null, "int", null, null) - builder.addProperty("address", "street", false, "string", null, null, null, null, extension) - } - } - - def "schema is converted correctly to JSON"() { - given: - SchemaBuilder builder = new SchemaBuilder(new Iterator()) - - when: - boolean shouldExtractPerson = builder.shouldExtractSchema("person", 0) - boolean shouldExtractAddress = builder.shouldExtractSchema("address", 1) - boolean shouldExtractPerson2 = builder.shouldExtractSchema("person", 0) - boolean shouldExtractTooDeep = builder.shouldExtractSchema("city", 11) - Schema schema = builder.build() - - then: - "{\"components\":{\"schemas\":{\"person\":{\"properties\":{\"name\":{\"description\":\"name of the person\",\"type\":\"string\"},\"phone_numbers\":{\"items\":{\"type\":\"string\"},\"type\":\"array\"},\"person_name\":{\"type\":\"string\"},\"address\":{\"\$ref\":\"#/components/schemas/address\",\"type\":\"object\"}},\"type\":\"object\"},\"address\":{\"properties\":{\"zip\":{\"format\":\"int\",\"type\":\"number\"},\"street\":{\"extensions\":{\"x-test-extension-1\":\"hello\",\"x-test-extension-2\":\"world\"},\"type\":\"string\"}},\"type\":\"object\"}}},\"openapi\":\"3.0.0\"}" == schema.definition - "16548065305426330543" == schema.id - shouldExtractPerson - shouldExtractAddress - !shouldExtractPerson2 - !shouldExtractTooDeep - } -} diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/datastreams/SchemaSamplerTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/datastreams/SchemaSamplerTest.groovy deleted file mode 100644 index 41615e95542..00000000000 --- a/dd-trace-core/src/test/groovy/datadog/trace/core/datastreams/SchemaSamplerTest.groovy +++ /dev/null @@ -1,30 +0,0 @@ -package datadog.trace.core.datastreams - -import datadog.trace.core.test.DDCoreSpecification - -class SchemaSamplerTest extends DDCoreSpecification { - - def "schema sampler samples with correct weights"() { - given: - long currentTimeMillis = 100000 - SchemaSampler sampler = new SchemaSampler() - - when: - boolean canSample1 = sampler.canSample(currentTimeMillis) - int weight1 = sampler.trySample(currentTimeMillis) - boolean canSample2= sampler.canSample(currentTimeMillis + 1000) - boolean canSample3 = sampler.canSample(currentTimeMillis + 2000) - boolean canSample4 = sampler.canSample(currentTimeMillis + 30000) - int weight4 = sampler.trySample(currentTimeMillis + 30000) - boolean canSample5 = sampler.canSample(currentTimeMillis + 30001) - - then: - canSample1 - weight1 == 1 - !canSample2 - !canSample3 - canSample4 - weight4 == 3 - !canSample5 - } -} diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/datastreams/TransactionContainerTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/datastreams/TransactionContainerTest.groovy deleted file mode 100644 index a7593acb19e..00000000000 --- a/dd-trace-core/src/test/groovy/datadog/trace/core/datastreams/TransactionContainerTest.groovy +++ /dev/null @@ -1,49 +0,0 @@ -package datadog.trace.core.datastreams - -import datadog.trace.api.datastreams.TransactionInfo -import datadog.trace.core.test.DDCoreSpecification - -class TransactionContainerTest extends DDCoreSpecification { - def "test with no resize"() { - given: - TransactionInfo.resetCache() - def container = new TransactionContainer(1024) - container.add(new TransactionInfo("1", 1, "1")) - container.add(new TransactionInfo("2", 2, "2")) - def data = container.getData() - - expect: - data.size() == 22 - data == new byte[] { - 1, 0, 0, 0, 0, 0, 0, 0, 1, 1, 49, 2, 0, 0, 0, 0, 0, 0, 0, 2, 1, 50 - } - } - - def "test with with resize"() { - given: - TransactionInfo.resetCache() - def container = new TransactionContainer(10) - container.add(new TransactionInfo("1", 1, "1")) - container.add(new TransactionInfo("2", 2, "2")) - def data = container.getData() - - expect: - data.size() == 22 - data == new byte[] { - 1, 0, 0, 0, 0, 0, 0, 0, 1, 1, 49, 2, 0, 0, 0, 0, 0, 0, 0, 2, 1, 50 - } - } - - def "test checkpoint map"() { - given: - TransactionInfo.resetCache() - new TransactionInfo("1", 1, "1") - new TransactionInfo("2", 2, "2") - def data = TransactionInfo.getCheckpointIdCacheBytes() - expect: - data.size() == 6 - data == new byte[] { - 1, 1, 49, 2, 1, 50 - } - } -} diff --git a/dd-trace-core/src/test/java/datadog/trace/api/datastreams/TransactionInfoTestBridge.java b/dd-trace-core/src/test/java/datadog/trace/api/datastreams/TransactionInfoTestBridge.java new file mode 100644 index 00000000000..9cd65594f25 --- /dev/null +++ b/dd-trace-core/src/test/java/datadog/trace/api/datastreams/TransactionInfoTestBridge.java @@ -0,0 +1,11 @@ +package datadog.trace.api.datastreams; + +/** + * Bridge class to allow tests to access package-private method exposed by the {@code + * TransactionInfo} + */ +public class TransactionInfoTestBridge { + public static void resetCache() { + TransactionInfo.resetCache(); + } +} diff --git a/dd-trace-core/src/test/java/datadog/trace/core/datastreams/CheckpointerTest.java b/dd-trace-core/src/test/java/datadog/trace/core/datastreams/CheckpointerTest.java new file mode 100644 index 00000000000..26e4a157aba --- /dev/null +++ b/dd-trace-core/src/test/java/datadog/trace/core/datastreams/CheckpointerTest.java @@ -0,0 +1,64 @@ +package datadog.trace.core.datastreams; + +import static datadog.trace.api.config.GeneralConfig.DATA_STREAMS_ENABLED; +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import datadog.trace.api.experimental.DataStreamsCheckpointer; +import datadog.trace.api.experimental.DataStreamsContextCarrier; +import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.AgentTracer; +import datadog.trace.core.CoreTracer; +import datadog.trace.core.DDCoreJavaSpecification; +import datadog.trace.core.DDSpan; +import datadog.trace.junit.utils.config.WithConfig; +import java.util.HashMap; +import java.util.Map; +import java.util.Set; +import org.junit.jupiter.api.Test; + +// Enable DSM +@WithConfig(key = DATA_STREAMS_ENABLED, value = "true") +public class CheckpointerTest extends DDCoreJavaSpecification { + @Test + void testSettingProduceAndConsumeCheckpoint() { + // Create a test tracer + CoreTracer tracer = tracerBuilder().build(); + AgentTracer.forceRegister(tracer); + // Get the test checkpointer + DataStreamsCheckpointer checkpointer = tracer.getDataStreamsCheckpointer(); + // Declare the carrier to test injected data + CustomContextCarrier carrier = new CustomContextCarrier(); + // Start and activate a span + AgentSpan span = tracer.buildSpan("test", "dsm-checkpoint").start(); + AgentScope scope = tracer.activateSpan(span); + + // Trigger produce checkpoint + checkpointer.setProduceCheckpoint("kafka", "testTopic", carrier); + checkpointer.setConsumeCheckpoint("kafka", "testTopic", carrier); + // Clean up span + scope.close(); + span.finish(); + + boolean hasPathwayCtxBase64 = + carrier.entries().stream() + .anyMatch(entry -> "dd-pathway-ctx-base64".equals(entry.getKey())); + assertTrue(hasPathwayCtxBase64); + assertNotEquals(0L, ((DDSpan) span).context().getPathwayContext().getHash()); + } + + static class CustomContextCarrier implements DataStreamsContextCarrier { + private final Map data = new HashMap<>(); + + @Override + public Set> entries() { + return data.entrySet(); + } + + @Override + public void set(String key, String value) { + data.put(key, value); + } + } +} diff --git a/dd-trace-core/src/test/java/datadog/trace/core/datastreams/DataStreamsTransactionExtractorsTest.java b/dd-trace-core/src/test/java/datadog/trace/core/datastreams/DataStreamsTransactionExtractorsTest.java new file mode 100644 index 00000000000..7565f19afee --- /dev/null +++ b/dd-trace-core/src/test/java/datadog/trace/core/datastreams/DataStreamsTransactionExtractorsTest.java @@ -0,0 +1,31 @@ +package datadog.trace.core.datastreams; + +import static datadog.trace.api.datastreams.DataStreamsTransactionExtractor.Type.HTTP_IN_HEADERS; +import static datadog.trace.api.datastreams.DataStreamsTransactionExtractor.Type.HTTP_OUT_HEADERS; +import static org.junit.jupiter.api.Assertions.assertEquals; + +import datadog.trace.api.datastreams.DataStreamsTransactionExtractor; +import datadog.trace.core.DDCoreJavaSpecification; +import java.util.List; +import org.junit.jupiter.api.Test; + +public class DataStreamsTransactionExtractorsTest extends DDCoreJavaSpecification { + @Test + void deserializeFromJson() { + DataStreamsTransactionExtractors list = + DataStreamsTransactionExtractors.deserialize( + "[\n" + + " {\"name\": \"extractor\", \"type\": \"HTTP_OUT_HEADERS\", \"value\": \"transaction_id\"},\n" + + " {\"name\": \"second_extractor\", \"type\": \"HTTP_IN_HEADERS\", \"value\": \"transaction_id\"}\n" + + "]"); + List extractors = list.getExtractors(); + + assertEquals(2, extractors.size()); + assertEquals("extractor", extractors.get(0).getName()); + assertEquals(HTTP_OUT_HEADERS, extractors.get(0).getType()); + assertEquals("transaction_id", extractors.get(0).getValue()); + assertEquals("second_extractor", extractors.get(1).getName()); + assertEquals(HTTP_IN_HEADERS, extractors.get(1).getType()); + assertEquals("transaction_id", extractors.get(1).getValue()); + } +} diff --git a/dd-trace-core/src/test/java/datadog/trace/core/datastreams/DataStreamsWritingTest.java b/dd-trace-core/src/test/java/datadog/trace/core/datastreams/DataStreamsWritingTest.java new file mode 100644 index 00000000000..52ad1a11447 --- /dev/null +++ b/dd-trace-core/src/test/java/datadog/trace/core/datastreams/DataStreamsWritingTest.java @@ -0,0 +1,649 @@ +package datadog.trace.core.datastreams; + +import static datadog.trace.api.config.GeneralConfig.EXPERIMENTAL_PROPAGATE_PROCESS_TAGS_ENABLED; +import static java.util.concurrent.TimeUnit.SECONDS; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import datadog.communication.ddagent.DDAgentFeaturesDiscovery; +import datadog.communication.ddagent.SharedCommunicationObjects; +import datadog.communication.http.OkHttpUtils; +import datadog.trace.agent.test.server.http.JavaTestHttpServer; +import datadog.trace.api.Config; +import datadog.trace.api.ProcessTags; +import datadog.trace.api.TraceConfig; +import datadog.trace.api.WellKnownTags; +import datadog.trace.api.datastreams.DataStreamsTags; +import datadog.trace.api.datastreams.StatsPoint; +import datadog.trace.api.time.ControllableTimeSource; +import datadog.trace.core.DDCoreJavaSpecification; +import datadog.trace.core.DDTraceCoreInfo; +import datadog.trace.junit.utils.config.WithConfigExtension; +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.CopyOnWriteArrayList; +import okhttp3.HttpUrl; +import okhttp3.OkHttpClient; +import okio.BufferedSource; +import okio.GzipSource; +import okio.Okio; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import org.msgpack.core.MessagePack; +import org.msgpack.core.MessageUnpacker; + +/** + * This test class exists because a real integration test is not possible. see + * DataStreamsIntegrationTest + */ +public class DataStreamsWritingTest extends DDCoreJavaSpecification { + + private static long defaultBucketDurationNanos; + + private static JavaTestHttpServer server; + private static HttpUrl serverAddress; + private static List requestBodies; + + @BeforeAll + static void startServer() { + defaultBucketDurationNanos = Config.get().getDataStreamsBucketDurationNanoseconds(); + requestBodies = new CopyOnWriteArrayList<>(); + server = + JavaTestHttpServer.httpServer( + s -> + s.handlers( + h -> + h.post( + DDAgentFeaturesDiscovery.V01_DATASTREAMS_ENDPOINT, + api -> { + requestBodies.add(api.getRequest().getBody()); + api.getResponse().status(200).send(); + }))); + serverAddress = HttpUrl.get(server.getAddress()); + } + + @AfterAll + static void stopServer() { + if (server != null) { + server.close(); + } + } + + @BeforeEach + void resetRequestBodies() { + requestBodies.clear(); + } + + private static void awaitOneRequestBody() throws InterruptedException { + long deadline = System.currentTimeMillis() + 2000; + while (requestBodies.size() < 1 && System.currentTimeMillis() < deadline) { + Thread.sleep(20); + } + assertEquals(1, requestBodies.size()); + } + + @Test + void serviceOverridesSplitBuckets() throws InterruptedException, IOException { + WellKnownTags wellKnownTags = + new WellKnownTags( + "runtimeid", "hostname", "test", Config.get().getServiceName(), "version", "java"); + Config fakeConfig = mock(Config.class); + when(fakeConfig.getAgentUrl()).thenReturn(serverAddress.toString()); + when(fakeConfig.getWellKnownTags()).thenReturn(wellKnownTags); + when(fakeConfig.getPrimaryTag()).thenReturn("region-1"); + + OkHttpClient testOkhttpClient = OkHttpUtils.buildHttpClient(serverAddress, 5000L); + DDAgentFeaturesDiscovery features = mock(DDAgentFeaturesDiscovery.class); + when(features.supportsDataStreams()).thenReturn(true); + SharedCommunicationObjects sharedCommObjects = new SharedCommunicationObjects(); + sharedCommObjects.setFeaturesDiscovery(features); + sharedCommObjects.agentHttpClient = testOkhttpClient; + sharedCommObjects.createRemaining(fakeConfig); + + ControllableTimeSource timeSource = new ControllableTimeSource(); + TraceConfig traceConfig = mock(TraceConfig.class); + when(traceConfig.isDataStreamsEnabled()).thenReturn(true); + String serviceNameOverride = "service-name-override"; + + DefaultDataStreamsMonitoring dataStreams = + new DefaultDataStreamsMonitoring( + fakeConfig, sharedCommObjects, timeSource, () -> traceConfig); + dataStreams.start(); + dataStreams.setThreadServiceName(serviceNameOverride); + dataStreams.add( + new StatsPoint( + DataStreamsTags.create(null, null), + 9, + 0, + 10, + timeSource.getCurrentTimeNanos(), + 0, + 0, + 0, + serviceNameOverride)); + dataStreams.trackBacklog( + DataStreamsTags.createWithPartition("kafka_produce", "testTopic", "1", null, null), 130); + timeSource.advance(defaultBucketDurationNanos); + // force flush + dataStreams.report(); + dataStreams.close(); + dataStreams.clearThreadServiceName(); + + awaitOneRequestBody(); + GzipSource gzipSource = + new GzipSource(Okio.source(new ByteArrayInputStream(requestBodies.get(0)))); + BufferedSource bufferedSource = Okio.buffer(gzipSource); + MessageUnpacker unpacker = MessagePack.newDefaultUnpacker(bufferedSource.inputStream()); + + assertEquals(9, unpacker.unpackMapHeader()); + assertEquals("Env", unpacker.unpackString()); + assertEquals("test", unpacker.unpackString()); + assertEquals("Service", unpacker.unpackString()); + assertEquals(serviceNameOverride, unpacker.unpackString()); + } + + @ParameterizedTest(name = "Write bucket to mock server with process tags enabled {0}") + @ValueSource(booleans = {true, false}) + void writeBucketToMockServer(boolean processTagsEnabled) + throws InterruptedException, IOException { + WithConfigExtension.injectSysConfig( + EXPERIMENTAL_PROPAGATE_PROCESS_TAGS_ENABLED, Boolean.toString(processTagsEnabled)); + ProcessTags.reset(Config.get()); + + WellKnownTags wellKnownTags = + new WellKnownTags( + "runtimeid", "hostname", "test", Config.get().getServiceName(), "version", "java"); + Config fakeConfig = mock(Config.class); + when(fakeConfig.getAgentUrl()).thenReturn(serverAddress.toString()); + when(fakeConfig.getWellKnownTags()).thenReturn(wellKnownTags); + when(fakeConfig.getPrimaryTag()).thenReturn("region-1"); + + OkHttpClient testOkhttpClient = OkHttpUtils.buildHttpClient(serverAddress, 5000L); + DDAgentFeaturesDiscovery features = mock(DDAgentFeaturesDiscovery.class); + when(features.supportsDataStreams()).thenReturn(true); + SharedCommunicationObjects sharedCommObjects = new SharedCommunicationObjects(); + sharedCommObjects.setFeaturesDiscovery(features); + sharedCommObjects.agentHttpClient = testOkhttpClient; + sharedCommObjects.createRemaining(fakeConfig); + + ControllableTimeSource timeSource = new ControllableTimeSource(); + TraceConfig traceConfig = mock(TraceConfig.class); + when(traceConfig.isDataStreamsEnabled()).thenReturn(true); + + DefaultDataStreamsMonitoring dataStreams = + new DefaultDataStreamsMonitoring( + fakeConfig, sharedCommObjects, timeSource, () -> traceConfig); + try { + dataStreams.start(); + dataStreams.add( + new StatsPoint( + DataStreamsTags.create(null, null), + 9, + 0, + 10, + timeSource.getCurrentTimeNanos(), + 0, + 0, + 0, + null)); + dataStreams.add( + new StatsPoint( + DataStreamsTags.create( + "testType", DataStreamsTags.Direction.INBOUND, "testTopic", "testGroup", null), + 1, + 2, + 5, + timeSource.getCurrentTimeNanos(), + 0, + 0, + 0, + null)); + dataStreams.trackBacklog( + DataStreamsTags.createWithPartition("kafka_produce", "testTopic", "1", null, null), 100); + dataStreams.trackBacklog( + DataStreamsTags.createWithPartition("kafka_produce", "testTopic", "1", null, null), 130); + timeSource.advance(defaultBucketDurationNanos - 100L); + dataStreams.add( + new StatsPoint( + DataStreamsTags.create( + "testType", DataStreamsTags.Direction.INBOUND, "testTopic", "testGroup", null), + 1, + 2, + 5, + timeSource.getCurrentTimeNanos(), + SECONDS.toNanos(10), + SECONDS.toNanos(10), + 10, + null)); + timeSource.advance(defaultBucketDurationNanos); + dataStreams.add( + new StatsPoint( + DataStreamsTags.create( + "testType", DataStreamsTags.Direction.INBOUND, "testTopic", "testGroup", null), + 1, + 2, + 5, + timeSource.getCurrentTimeNanos(), + SECONDS.toNanos(5), + SECONDS.toNanos(5), + 5, + null)); + dataStreams.add( + new StatsPoint( + DataStreamsTags.create( + "testType", DataStreamsTags.Direction.INBOUND, "testTopic2", "testGroup", null), + 3, + 4, + 6, + timeSource.getCurrentTimeNanos(), + SECONDS.toNanos(2), + 0, + 2, + null)); + timeSource.advance(defaultBucketDurationNanos); + dataStreams.close(); + + awaitOneRequestBody(); + validateMessage(requestBodies.get(0), processTagsEnabled); + } finally { + // cleanup + WithConfigExtension.injectSysConfig(EXPERIMENTAL_PROPAGATE_PROCESS_TAGS_ENABLED, "true"); + ProcessTags.reset(Config.get()); + } + } + + @Test + void writeKafkaConfigsToMockServer() throws InterruptedException, IOException { + WellKnownTags wellKnownTags = + new WellKnownTags( + "runtimeid", "hostname", "test", Config.get().getServiceName(), "version", "java"); + Config fakeConfig = mock(Config.class); + when(fakeConfig.getAgentUrl()).thenReturn(serverAddress.toString()); + when(fakeConfig.getWellKnownTags()).thenReturn(wellKnownTags); + when(fakeConfig.getPrimaryTag()).thenReturn("region-1"); + + OkHttpClient testOkhttpClient = OkHttpUtils.buildHttpClient(serverAddress, 5000L); + DDAgentFeaturesDiscovery features = mock(DDAgentFeaturesDiscovery.class); + when(features.supportsDataStreams()).thenReturn(true); + SharedCommunicationObjects sharedCommObjects = new SharedCommunicationObjects(); + sharedCommObjects.setFeaturesDiscovery(features); + sharedCommObjects.agentHttpClient = testOkhttpClient; + sharedCommObjects.createRemaining(fakeConfig); + + ControllableTimeSource timeSource = new ControllableTimeSource(); + TraceConfig traceConfig = mock(TraceConfig.class); + when(traceConfig.isDataStreamsEnabled()).thenReturn(true); + + DefaultDataStreamsMonitoring dataStreams = + new DefaultDataStreamsMonitoring( + fakeConfig, sharedCommObjects, timeSource, () -> traceConfig); + dataStreams.start(); + + // Report a producer and consumer config + Map producerConfig = new HashMap<>(); + producerConfig.put("bootstrap.servers", "localhost:9092"); + producerConfig.put("acks", "all"); + producerConfig.put("linger.ms", "5"); + dataStreams.reportKafkaConfig("kafka_producer", "", "", producerConfig); + + Map consumerConfig = new HashMap<>(); + consumerConfig.put("bootstrap.servers", "localhost:9092"); + consumerConfig.put("group.id", "test-group"); + consumerConfig.put("auto.offset.reset", "earliest"); + dataStreams.reportKafkaConfig("kafka_consumer", "", "test-group", consumerConfig); + + // Also add a stats point so the bucket is not empty of stats + dataStreams.add( + new StatsPoint( + DataStreamsTags.create(null, null), + 9, + 0, + 10, + timeSource.getCurrentTimeNanos(), + 0, + 0, + 0, + null)); + + timeSource.advance(defaultBucketDurationNanos); + dataStreams.close(); + + awaitOneRequestBody(); + validateKafkaConfigMessage(requestBodies.get(0)); + } + + @Test + void duplicateKafkaConfigsAreEachSerializedInPayload() throws InterruptedException, IOException { + WellKnownTags wellKnownTags = + new WellKnownTags( + "runtimeid", "hostname", "test", Config.get().getServiceName(), "version", "java"); + Config fakeConfig = mock(Config.class); + when(fakeConfig.getAgentUrl()).thenReturn(serverAddress.toString()); + when(fakeConfig.getWellKnownTags()).thenReturn(wellKnownTags); + when(fakeConfig.getPrimaryTag()).thenReturn("region-1"); + + OkHttpClient testOkhttpClient = OkHttpUtils.buildHttpClient(serverAddress, 5000L); + DDAgentFeaturesDiscovery features = mock(DDAgentFeaturesDiscovery.class); + when(features.supportsDataStreams()).thenReturn(true); + SharedCommunicationObjects sharedCommObjects = new SharedCommunicationObjects(); + sharedCommObjects.setFeaturesDiscovery(features); + sharedCommObjects.agentHttpClient = testOkhttpClient; + sharedCommObjects.createRemaining(fakeConfig); + + ControllableTimeSource timeSource = new ControllableTimeSource(); + TraceConfig traceConfig = mock(TraceConfig.class); + when(traceConfig.isDataStreamsEnabled()).thenReturn(true); + + DefaultDataStreamsMonitoring dataStreams = + new DefaultDataStreamsMonitoring( + fakeConfig, sharedCommObjects, timeSource, () -> traceConfig); + dataStreams.start(); + + // Report the same producer config twice — both should be serialized + Map producerConfig = new HashMap<>(); + producerConfig.put("bootstrap.servers", "localhost:9092"); + producerConfig.put("acks", "all"); + dataStreams.reportKafkaConfig("kafka_producer", "", "", producerConfig); + dataStreams.reportKafkaConfig("kafka_producer", "", "", producerConfig); + + // Also add a stats point so the bucket has content + dataStreams.add( + new StatsPoint( + DataStreamsTags.create(null, null), + 9, + 0, + 10, + timeSource.getCurrentTimeNanos(), + 0, + 0, + 0, + null)); + + timeSource.advance(defaultBucketDurationNanos); + dataStreams.close(); + + awaitOneRequestBody(); + validateDuplicateKafkaConfigMessage(requestBodies.get(0)); + } + + private void validateKafkaConfigMessage(byte[] message) throws IOException { + GzipSource gzipSource = new GzipSource(Okio.source(new ByteArrayInputStream(message))); + BufferedSource bufferedSource = Okio.buffer(gzipSource); + MessageUnpacker unpacker = MessagePack.newDefaultUnpacker(bufferedSource.inputStream()); + + // Outer map (same structure as other payloads) + int outerMapSize = unpacker.unpackMapHeader(); + // Skip to Stats array + boolean foundStats = false; + for (int i = 0; i < outerMapSize; i++) { + String key = unpacker.unpackString(); + if ("Stats".equals(key)) { + foundStats = true; + int numBuckets = unpacker.unpackArrayHeader(); + assertTrue(numBuckets >= 1); + + // Parse first bucket + int bucketMapSize = unpacker.unpackMapHeader(); + boolean foundConfigs = false; + for (int j = 0; j < bucketMapSize; j++) { + String bucketKey = unpacker.unpackString(); + if ("Configs".equals(bucketKey)) { + foundConfigs = true; + int numConfigs = unpacker.unpackArrayHeader(); + assertEquals(2, numConfigs); + + // Collect configs in a map keyed by type + Map> configsByType = new HashMap<>(); + for (int n = 0; n < numConfigs; n++) { + assertEquals(4, unpacker.unpackMapHeader()); + assertEquals("Type", unpacker.unpackString()); + String type = unpacker.unpackString(); + assertEquals("KafkaClusterId", unpacker.unpackString()); + unpacker.unpackString(); // skip cluster id value + assertEquals("ConsumerGroup", unpacker.unpackString()); + unpacker.unpackString(); // skip consumer group value + assertEquals("Config", unpacker.unpackString()); + int configSize = unpacker.unpackMapHeader(); + Map configEntries = new HashMap<>(); + for (int c = 0; c < configSize; c++) { + String ck = unpacker.unpackString(); + String cv = unpacker.unpackString(); + configEntries.put(ck, cv); + } + configsByType.put(type, configEntries); + } + + // Verify producer config + assertTrue(configsByType.containsKey("kafka_producer")); + assertEquals( + "localhost:9092", configsByType.get("kafka_producer").get("bootstrap.servers")); + assertEquals("all", configsByType.get("kafka_producer").get("acks")); + assertEquals("5", configsByType.get("kafka_producer").get("linger.ms")); + + // Verify consumer config + assertTrue(configsByType.containsKey("kafka_consumer")); + assertEquals( + "localhost:9092", configsByType.get("kafka_consumer").get("bootstrap.servers")); + assertEquals("test-group", configsByType.get("kafka_consumer").get("group.id")); + assertEquals("earliest", configsByType.get("kafka_consumer").get("auto.offset.reset")); + } else { + unpacker.skipValue(); + } + } + assertTrue(foundConfigs, "Configs field not found in bucket"); + + // Skip remaining buckets + for (int b = 1; b < numBuckets; b++) { + unpacker.skipValue(); + } + } else { + unpacker.skipValue(); + } + } + assertTrue(foundStats, "Stats field not found in payload"); + } + + private void validateDuplicateKafkaConfigMessage(byte[] message) throws IOException { + GzipSource gzipSource = new GzipSource(Okio.source(new ByteArrayInputStream(message))); + BufferedSource bufferedSource = Okio.buffer(gzipSource); + MessageUnpacker unpacker = MessagePack.newDefaultUnpacker(bufferedSource.inputStream()); + + int outerMapSize = unpacker.unpackMapHeader(); + boolean foundStats = false; + for (int i = 0; i < outerMapSize; i++) { + String key = unpacker.unpackString(); + if ("Stats".equals(key)) { + foundStats = true; + int numBuckets = unpacker.unpackArrayHeader(); + assertTrue(numBuckets >= 1); + + // Parse first bucket + int bucketMapSize = unpacker.unpackMapHeader(); + boolean foundConfigs = false; + for (int j = 0; j < bucketMapSize; j++) { + String bucketKey = unpacker.unpackString(); + if ("Configs".equals(bucketKey)) { + foundConfigs = true; + int numConfigs = unpacker.unpackArrayHeader(); + // Both configs should be present (no deduplication) + assertEquals(2, numConfigs); + + for (int n = 0; n < numConfigs; n++) { + assertEquals(4, unpacker.unpackMapHeader()); + assertEquals("Type", unpacker.unpackString()); + assertEquals("kafka_producer", unpacker.unpackString()); + assertEquals("KafkaClusterId", unpacker.unpackString()); + unpacker.unpackString(); // skip cluster id value + assertEquals("ConsumerGroup", unpacker.unpackString()); + unpacker.unpackString(); // skip consumer group value + assertEquals("Config", unpacker.unpackString()); + int configSize = unpacker.unpackMapHeader(); + Map configEntries = new HashMap<>(); + for (int c = 0; c < configSize; c++) { + String ck = unpacker.unpackString(); + String cv = unpacker.unpackString(); + configEntries.put(ck, cv); + } + assertEquals("localhost:9092", configEntries.get("bootstrap.servers")); + assertEquals("all", configEntries.get("acks")); + } + } else { + unpacker.skipValue(); + } + } + assertTrue(foundConfigs, "Configs field not found in bucket"); + + for (int b = 1; b < numBuckets; b++) { + unpacker.skipValue(); + } + } else { + unpacker.skipValue(); + } + } + assertTrue(foundStats, "Stats field not found in payload"); + } + + private void validateMessage(byte[] message, boolean processTagsEnabled) throws IOException { + GzipSource gzipSource = new GzipSource(Okio.source(new ByteArrayInputStream(message))); + BufferedSource bufferedSource = Okio.buffer(gzipSource); + MessageUnpacker unpacker = MessagePack.newDefaultUnpacker(bufferedSource.inputStream()); + + assertEquals(8 + (processTagsEnabled ? 1 : 0), unpacker.unpackMapHeader()); + assertEquals("Env", unpacker.unpackString()); + assertEquals("test", unpacker.unpackString()); + assertEquals("Service", unpacker.unpackString()); + assertEquals(Config.get().getServiceName(), unpacker.unpackString()); + assertEquals("Lang", unpacker.unpackString()); + assertEquals("java", unpacker.unpackString()); + assertEquals("PrimaryTag", unpacker.unpackString()); + assertEquals("region-1", unpacker.unpackString()); + assertEquals("TracerVersion", unpacker.unpackString()); + assertEquals(DDTraceCoreInfo.VERSION, unpacker.unpackString()); + assertEquals("Version", unpacker.unpackString()); + assertEquals("version", unpacker.unpackString()); + assertEquals("Stats", unpacker.unpackString()); + assertEquals(2, unpacker.unpackArrayHeader()); // 2 time buckets + + // FIRST BUCKET + assertEquals(4, unpacker.unpackMapHeader()); + assertEquals("Start", unpacker.unpackString()); + unpacker.skipValue(); + assertEquals("Duration", unpacker.unpackString()); + assertEquals(defaultBucketDurationNanos, unpacker.unpackLong()); + assertEquals("Stats", unpacker.unpackString()); + assertEquals(2, unpacker.unpackArrayHeader()); // 2 groups in first bucket + + // we don't know the order the groups will be reported + Set availableSizes = new HashSet<>(); + availableSizes.add(5); + availableSizes.add(6); + for (int g = 0; g < 2; g++) { + int mapHeaderSize = unpacker.unpackMapHeader(); + assertTrue(availableSizes.remove(mapHeaderSize)); + if (mapHeaderSize == 5) { + // empty topic group + assertEquals("PathwayLatency", unpacker.unpackString()); + unpacker.skipValue(); + assertEquals("EdgeLatency", unpacker.unpackString()); + unpacker.skipValue(); + assertEquals("PayloadSize", unpacker.unpackString()); + unpacker.skipValue(); + assertEquals("Hash", unpacker.unpackString()); + assertEquals(9L, unpacker.unpackLong()); + assertEquals("ParentHash", unpacker.unpackString()); + assertEquals(0L, unpacker.unpackLong()); + } else { + // other group + assertEquals("PathwayLatency", unpacker.unpackString()); + unpacker.skipValue(); + assertEquals("EdgeLatency", unpacker.unpackString()); + unpacker.skipValue(); + assertEquals("PayloadSize", unpacker.unpackString()); + unpacker.skipValue(); + assertEquals("Hash", unpacker.unpackString()); + assertEquals(1L, unpacker.unpackLong()); + assertEquals("ParentHash", unpacker.unpackString()); + assertEquals(2L, unpacker.unpackLong()); + assertEquals("EdgeTags", unpacker.unpackString()); + assertEquals(4, unpacker.unpackArrayHeader()); + assertEquals("direction:in", unpacker.unpackString()); + assertEquals("topic:testTopic", unpacker.unpackString()); + assertEquals("type:testType", unpacker.unpackString()); + assertEquals("group:testGroup", unpacker.unpackString()); + } + } + + // Kafka stats + assertEquals("Backlogs", unpacker.unpackString()); + assertEquals(1, unpacker.unpackArrayHeader()); + assertEquals(2, unpacker.unpackMapHeader()); + assertEquals("Tags", unpacker.unpackString()); + assertEquals(3, unpacker.unpackArrayHeader()); + assertEquals("topic:testTopic", unpacker.unpackString()); + assertEquals("type:kafka_produce", unpacker.unpackString()); + assertEquals("partition:1", unpacker.unpackString()); + assertEquals("Value", unpacker.unpackString()); + assertEquals(130L, unpacker.unpackLong()); + + // SECOND BUCKET + assertEquals(3, unpacker.unpackMapHeader()); + assertEquals("Start", unpacker.unpackString()); + unpacker.skipValue(); + assertEquals("Duration", unpacker.unpackString()); + assertEquals(defaultBucketDurationNanos, unpacker.unpackLong()); + assertEquals("Stats", unpacker.unpackString()); + assertEquals(2, unpacker.unpackArrayHeader()); // 2 groups in second bucket + + // we don't know the order the groups will be reported + Set availableHashes = new HashSet<>(); + availableHashes.add(1L); + availableHashes.add(3L); + for (int g = 0; g < 2; g++) { + assertEquals(6, unpacker.unpackMapHeader()); + assertEquals("PathwayLatency", unpacker.unpackString()); + unpacker.skipValue(); + assertEquals("EdgeLatency", unpacker.unpackString()); + unpacker.skipValue(); + assertEquals("PayloadSize", unpacker.unpackString()); + unpacker.skipValue(); + assertEquals("Hash", unpacker.unpackString()); + long hash = unpacker.unpackLong(); + assertTrue(availableHashes.remove(hash)); + assertEquals("ParentHash", unpacker.unpackString()); + assertEquals(hash == 1L ? 2L : 4L, unpacker.unpackLong()); + assertEquals("EdgeTags", unpacker.unpackString()); + assertEquals(4, unpacker.unpackArrayHeader()); + assertEquals("direction:in", unpacker.unpackString()); + assertEquals(hash == 1L ? "topic:testTopic" : "topic:testTopic2", unpacker.unpackString()); + assertEquals("type:testType", unpacker.unpackString()); + assertEquals("group:testGroup", unpacker.unpackString()); + } + + assertEquals("ProductMask", unpacker.unpackString()); + assertEquals(1L, unpacker.unpackLong()); + + List processTags = ProcessTags.getTagsAsStringList(); + if (processTags == null) { + assertFalse(unpacker.hasNext()); + } else { + assertTrue(unpacker.hasNext()); + assertEquals("ProcessTags", unpacker.unpackString()); + assertEquals(processTags.size(), unpacker.unpackArrayHeader()); + for (String tag : processTags) { + assertEquals(tag, unpacker.unpackString()); + } + } + } +} diff --git a/dd-trace-core/src/test/java/datadog/trace/core/datastreams/DefaultDataStreamsMonitoringTest.java b/dd-trace-core/src/test/java/datadog/trace/core/datastreams/DefaultDataStreamsMonitoringTest.java new file mode 100644 index 00000000000..c1c29026c49 --- /dev/null +++ b/dd-trace-core/src/test/java/datadog/trace/core/datastreams/DefaultDataStreamsMonitoringTest.java @@ -0,0 +1,1575 @@ +package datadog.trace.core.datastreams; + +import static datadog.trace.core.datastreams.DefaultDataStreamsMonitoring.FEATURE_CHECK_INTERVAL_NANOS; +import static datadog.trace.core.datastreams.DefaultDataStreamsMonitoringTestBridge.getThreadState; +import static datadog.trace.core.datastreams.DefaultDataStreamsMonitoringTestBridge.isInboxEmpty; +import static java.util.concurrent.TimeUnit.SECONDS; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.RETURNS_SMART_NULLS; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import datadog.communication.ddagent.DDAgentFeaturesDiscovery; +import datadog.metrics.api.Histograms; +import datadog.metrics.impl.DDSketchHistograms; +import datadog.trace.api.Config; +import datadog.trace.api.TraceConfig; +import datadog.trace.api.datastreams.DataStreamsTags; +import datadog.trace.api.datastreams.KafkaConfigReport; +import datadog.trace.api.datastreams.SchemaRegistryUsage; +import datadog.trace.api.datastreams.StatsPoint; +import datadog.trace.api.experimental.DataStreamsContextCarrier; +import datadog.trace.api.time.ControllableTimeSource; +import datadog.trace.common.metrics.EventListener; +import datadog.trace.common.metrics.Sink; +import datadog.trace.core.DDCoreJavaSpecification; +import java.util.ArrayList; +import java.util.Collection; +import java.util.Comparator; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Map.Entry; +import java.util.Set; +import java.util.concurrent.TimeUnit; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.tabletest.junit.TableTest; + +public class DefaultDataStreamsMonitoringTest extends DDCoreJavaSpecification { + + private static final long DEFAULT_BUCKET_DURATION_NANOS = + Config.get().getDataStreamsBucketDurationNanoseconds(); + + @BeforeAll + static void registerHistograms() { + Histograms.register(DDSketchHistograms.FACTORY); + } + + private static void awaitIdle(DefaultDataStreamsMonitoring dataStreams) + throws InterruptedException { + long deadline = System.currentTimeMillis() + 1000; + while (System.currentTimeMillis() < deadline) { + if (isInboxEmpty(dataStreams) && getThreadState(dataStreams) != Thread.State.RUNNABLE) { + return; + } + Thread.sleep(10); + } + assertTrue(isInboxEmpty(dataStreams)); + assertNotEquals(Thread.State.RUNNABLE, getThreadState(dataStreams)); + } + + private static void awaitBuckets( + DefaultDataStreamsMonitoring dataStreams, CapturingPayloadWriter writer, int expectedBuckets) + throws InterruptedException { + long deadline = System.currentTimeMillis() + 1000; + while (System.currentTimeMillis() < deadline) { + if (writer.buckets.size() == expectedBuckets + && isInboxEmpty(dataStreams) + && getThreadState(dataStreams) != Thread.State.RUNNABLE) { + return; + } + Thread.sleep(10); + } + assertTrue(isInboxEmpty(dataStreams)); + assertNotEquals(Thread.State.RUNNABLE, getThreadState(dataStreams)); + assertEquals(expectedBuckets, writer.buckets.size()); + } + + private static DDAgentFeaturesDiscovery stubFeatures(boolean supportsDataStreams) { + DDAgentFeaturesDiscovery features = mock(DDAgentFeaturesDiscovery.class, RETURNS_SMART_NULLS); + when(features.supportsDataStreams()).thenReturn(supportsDataStreams); + return features; + } + + private static TraceConfig stubTraceConfig(boolean dataStreamsEnabled) { + TraceConfig traceConfig = mock(TraceConfig.class, RETURNS_SMART_NULLS); + when(traceConfig.isDataStreamsEnabled()).thenReturn(dataStreamsEnabled); + return traceConfig; + } + + @TableTest({ + "scenario | enabledAtAgent | enabledInConfig", + "agent off, config on | false | true ", + "agent on, config off | true | false ", + "both off | false | false " + }) + void noPayloadsWrittenIfDataStreamsNotSupportedOrNotEnabled( + boolean enabledAtAgent, boolean enabledInConfig) throws InterruptedException { + DDAgentFeaturesDiscovery features = stubFeatures(enabledAtAgent); + ControllableTimeSource timeSource = new ControllableTimeSource(); + DatastreamsPayloadWriter payloadWriter = mock(DatastreamsPayloadWriter.class); + Sink sink = mock(Sink.class); + TraceConfig traceConfig = stubTraceConfig(enabledInConfig); + + DefaultDataStreamsMonitoring dataStreams = + new DefaultDataStreamsMonitoring( + sink, + features, + timeSource, + () -> traceConfig, + payloadWriter, + DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.start(); + dataStreams.add( + new StatsPoint( + DataStreamsTags.create("testType", null, "testTopic", "testGroup", null), + 0, + 0, + 0, + timeSource.getCurrentTimeNanos(), + 0, + 0, + 0, + null)); + dataStreams.report(); + + awaitIdle(dataStreams); + verify(payloadWriter, org.mockito.Mockito.never()).writePayload(any(), any()); + + // cleanup + dataStreams.close(); + } + + @Test + void schemaSamplerSamplesWithCorrectWeights() { + DDAgentFeaturesDiscovery features = stubFeatures(true); + ControllableTimeSource timeSource = new ControllableTimeSource(); + timeSource.set(1000000000000L); + DatastreamsPayloadWriter payloadWriter = mock(DatastreamsPayloadWriter.class); + Sink sink = mock(Sink.class); + TraceConfig traceConfig = stubTraceConfig(true); + + DefaultDataStreamsMonitoring dataStreams = + new DefaultDataStreamsMonitoring( + sink, + features, + timeSource, + () -> traceConfig, + payloadWriter, + DEFAULT_BUCKET_DURATION_NANOS); + + // the first received schema is sampled, with a weight of one. + assertTrue(dataStreams.canSampleSchema("schema1")); + assertEquals(1, dataStreams.trySampleSchema("schema1")); + // the sampling is done by topic, so a schema on a different topic will also be sampled at once, + // also with a weight of one. + assertTrue(dataStreams.canSampleSchema("schema2")); + assertEquals(1, dataStreams.trySampleSchema("schema2")); + // no time has passed from the last sampling, so the same schema is not sampled again (two times + // in a row). + assertFalse(dataStreams.canSampleSchema("schema1")); + assertFalse(dataStreams.canSampleSchema("schema1")); + timeSource.advance((long) (30 * 1e9)); + // now, 30 seconds have passed, so the schema is sampled again, with a weight of 3 (so it + // includes the two times the schema was not sampled). + assertTrue(dataStreams.canSampleSchema("schema1")); + assertEquals(3, dataStreams.trySampleSchema("schema1")); + } + + @Test + void contextCarrierAdapterTest() { + CustomContextCarrier carrier = new CustomContextCarrier(); + String keyName = "keyName"; + String keyValue = "keyValue"; + String[] extracted = new String[] {""}; + + DataStreamsContextCarrierAdapter.INSTANCE.set(carrier, keyName, keyValue); + DataStreamsContextCarrierAdapter.INSTANCE.forEachKeyValue( + carrier, + (key, value) -> { + if (keyName.equals(key)) { + extracted[0] = value; + } + }); + + assertEquals(keyValue, extracted[0]); + } + + @Test + void writeGroupAfterADelay() throws InterruptedException { + DDAgentFeaturesDiscovery features = stubFeatures(true); + ControllableTimeSource timeSource = new ControllableTimeSource(); + Sink sink = mock(Sink.class); + CapturingPayloadWriter payloadWriter = new CapturingPayloadWriter(); + TraceConfig traceConfig = stubTraceConfig(true); + + DefaultDataStreamsMonitoring dataStreams = + new DefaultDataStreamsMonitoring( + sink, + features, + timeSource, + () -> traceConfig, + payloadWriter, + DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.start(); + DataStreamsTags tags = DataStreamsTags.create("testType", null, "testTopic", "testGroup", null); + dataStreams.add(new StatsPoint(tags, 1, 2, 3, timeSource.getCurrentTimeNanos(), 0, 0, 0, null)); + timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.report(); + + awaitBuckets(dataStreams, payloadWriter, 1); + + StatsBucket bucket = payloadWriter.buckets.get(0); + assertEquals(1, bucket.getGroups().size()); + StatsGroup group = bucket.getGroups().iterator().next(); + assertEquals("type:testType", group.getTags().getType()); + assertEquals("group:testGroup", group.getTags().getGroup()); + assertEquals("topic:testTopic", group.getTags().getTopic()); + assertEquals(3, group.getTags().nonNullSize()); + assertEquals(1L, group.getHash()); + assertEquals(2L, group.getParentHash()); + + // cleanup + payloadWriter.close(); + dataStreams.close(); + } + + // This test relies on automatic reporting instead of manually calling report + @Test + void slowWriteGroupAfterADelay() throws InterruptedException { + DDAgentFeaturesDiscovery features = stubFeatures(true); + ControllableTimeSource timeSource = new ControllableTimeSource(); + Sink sink = mock(Sink.class); + CapturingPayloadWriter payloadWriter = new CapturingPayloadWriter(); + long bucketDuration = TimeUnit.MILLISECONDS.toNanos(200); + TraceConfig traceConfig = stubTraceConfig(true); + + DefaultDataStreamsMonitoring dataStreams = + new DefaultDataStreamsMonitoring( + sink, features, timeSource, () -> traceConfig, payloadWriter, bucketDuration); + dataStreams.start(); + DataStreamsTags tags = DataStreamsTags.create("testType", null, "testTopic", "testGroup", null); + dataStreams.add(new StatsPoint(tags, 1, 2, 3, timeSource.getCurrentTimeNanos(), 0, 0, 0, null)); + timeSource.advance(bucketDuration); + + awaitBuckets(dataStreams, payloadWriter, 1); + + StatsBucket bucket = payloadWriter.buckets.get(0); + assertEquals(1, bucket.getGroups().size()); + StatsGroup group = bucket.getGroups().iterator().next(); + assertEquals("type:testType", group.getTags().getType()); + assertEquals("group:testGroup", group.getTags().getGroup()); + assertEquals("topic:testTopic", group.getTags().getTopic()); + assertEquals(3, group.getTags().nonNullSize()); + assertEquals(1L, group.getHash()); + assertEquals(2L, group.getParentHash()); + + // cleanup + payloadWriter.close(); + dataStreams.close(); + } + + @Test + void groupsForCurrentBucketAreNotReported() throws InterruptedException { + DDAgentFeaturesDiscovery features = stubFeatures(true); + ControllableTimeSource timeSource = new ControllableTimeSource(); + Sink sink = mock(Sink.class); + CapturingPayloadWriter payloadWriter = new CapturingPayloadWriter(); + TraceConfig traceConfig = stubTraceConfig(true); + + DefaultDataStreamsMonitoring dataStreams = + new DefaultDataStreamsMonitoring( + sink, + features, + timeSource, + () -> traceConfig, + payloadWriter, + DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.start(); + DataStreamsTags tags = DataStreamsTags.create("testType", null, "testTopic", "testGroup", null); + dataStreams.add(new StatsPoint(tags, 1, 2, 3, timeSource.getCurrentTimeNanos(), 0, 0, 0, null)); + timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.add(new StatsPoint(tags, 3, 4, 3, timeSource.getCurrentTimeNanos(), 0, 0, 0, null)); + timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS - 100L); + dataStreams.report(); + + awaitBuckets(dataStreams, payloadWriter, 1); + + StatsBucket bucket = payloadWriter.buckets.get(0); + assertEquals(1, bucket.getGroups().size()); + StatsGroup group = bucket.getGroups().iterator().next(); + assertEquals("type:testType", group.getTags().getType()); + assertEquals("group:testGroup", group.getTags().getGroup()); + assertEquals("topic:testTopic", group.getTags().getTopic()); + assertEquals(3, group.getTags().nonNullSize()); + assertEquals(1L, group.getHash()); + assertEquals(2L, group.getParentHash()); + + // cleanup + payloadWriter.close(); + dataStreams.close(); + } + + @Test + void allGroupsWrittenInClose() throws InterruptedException { + DDAgentFeaturesDiscovery features = stubFeatures(true); + ControllableTimeSource timeSource = new ControllableTimeSource(); + Sink sink = mock(Sink.class); + CapturingPayloadWriter payloadWriter = new CapturingPayloadWriter(); + TraceConfig traceConfig = stubTraceConfig(true); + + DefaultDataStreamsMonitoring dataStreams = + new DefaultDataStreamsMonitoring( + sink, + features, + timeSource, + () -> traceConfig, + payloadWriter, + DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.start(); + DataStreamsTags tags1 = + DataStreamsTags.create("testType", null, "testTopic", "testGroup", null); + DataStreamsTags tags2 = + DataStreamsTags.create("testType", null, "testTopic2", "testGroup", null); + dataStreams.add( + new StatsPoint(tags1, 1, 2, 5, timeSource.getCurrentTimeNanos(), 0, 0, 0, null)); + timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.add( + new StatsPoint(tags2, 3, 4, 6, timeSource.getCurrentTimeNanos(), 0, 0, 0, null)); + timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS - 100L); + dataStreams.close(); + + awaitBuckets(dataStreams, payloadWriter, 2); + + StatsBucket bucket0 = payloadWriter.buckets.get(0); + assertEquals(1, bucket0.getGroups().size()); + StatsGroup group0 = bucket0.getGroups().iterator().next(); + assertEquals("type:testType", group0.getTags().getType()); + assertEquals("group:testGroup", group0.getTags().getGroup()); + assertEquals("topic:testTopic", group0.getTags().getTopic()); + assertEquals(3, group0.getTags().nonNullSize()); + assertEquals(1L, group0.getHash()); + assertEquals(2L, group0.getParentHash()); + + StatsBucket bucket1 = payloadWriter.buckets.get(1); + assertEquals(1, bucket1.getGroups().size()); + StatsGroup group1 = bucket1.getGroups().iterator().next(); + assertEquals("type:testType", group1.getTags().getType()); + assertEquals("group:testGroup", group1.getTags().getGroup()); + assertEquals("topic:testTopic2", group1.getTags().getTopic()); + assertEquals(3, group1.getTags().nonNullSize()); + assertEquals(3L, group1.getHash()); + assertEquals(4L, group1.getParentHash()); + + // cleanup + payloadWriter.close(); + } + + @Test + void kafkaOffsetsAreTracked() throws InterruptedException { + DDAgentFeaturesDiscovery features = stubFeatures(true); + ControllableTimeSource timeSource = new ControllableTimeSource(); + Sink sink = mock(Sink.class); + CapturingPayloadWriter payloadWriter = new CapturingPayloadWriter(); + TraceConfig traceConfig = stubTraceConfig(true); + + DefaultDataStreamsMonitoring dataStreams = + new DefaultDataStreamsMonitoring( + sink, + features, + timeSource, + () -> traceConfig, + payloadWriter, + DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.start(); + dataStreams.trackBacklog( + DataStreamsTags.createWithPartition("kafka_commit", "testTopic", "2", null, "testGroup"), + 23); + dataStreams.trackBacklog( + DataStreamsTags.createWithPartition("kafka_commit", "testTopic", "2", null, "testGroup"), + 24); + dataStreams.trackBacklog( + DataStreamsTags.createWithPartition("kafka_produce", "testTopic", "2", null, null), 23); + dataStreams.trackBacklog( + DataStreamsTags.createWithPartition("kafka_produce", "testTopic2", "2", null, null), 23); + dataStreams.trackBacklog( + DataStreamsTags.createWithPartition("kafka_produce", "testTopic", "2", null, null), 45); + timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.report(); + + awaitBuckets(dataStreams, payloadWriter, 1); + + StatsBucket bucket = payloadWriter.buckets.get(0); + Collection> backlogs = bucket.getBacklogs(); + assertEquals(3, backlogs.size()); + List> sortedBacklogs = new ArrayList<>(backlogs); + sortedBacklogs.sort(Comparator.comparing(e -> e.getKey().toString())); + assertEquals( + DataStreamsTags.createWithPartition("kafka_commit", "testTopic", "2", null, "testGroup"), + sortedBacklogs.get(0).getKey()); + assertEquals(24L, sortedBacklogs.get(0).getValue()); + assertEquals( + DataStreamsTags.createWithPartition("kafka_produce", "testTopic", "2", null, null), + sortedBacklogs.get(1).getKey()); + assertEquals(45L, sortedBacklogs.get(1).getValue()); + assertEquals( + DataStreamsTags.createWithPartition("kafka_produce", "testTopic2", "2", null, null), + sortedBacklogs.get(2).getKey()); + assertEquals(23L, sortedBacklogs.get(2).getValue()); + + // cleanup + payloadWriter.close(); + dataStreams.close(); + } + + @Test + void groupsFromMultipleBucketsAreReported() throws InterruptedException { + DDAgentFeaturesDiscovery features = stubFeatures(true); + ControllableTimeSource timeSource = new ControllableTimeSource(); + Sink sink = mock(Sink.class); + CapturingPayloadWriter payloadWriter = new CapturingPayloadWriter(); + TraceConfig traceConfig = stubTraceConfig(true); + + DefaultDataStreamsMonitoring dataStreams = + new DefaultDataStreamsMonitoring( + sink, + features, + timeSource, + () -> traceConfig, + payloadWriter, + DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.start(); + DataStreamsTags tags1 = + DataStreamsTags.create("testType", null, "testTopic", "testGroup", null); + dataStreams.add( + new StatsPoint(tags1, 1, 2, 5, timeSource.getCurrentTimeNanos(), 0, 0, 0, null)); + timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS * 10); + DataStreamsTags tags2 = + DataStreamsTags.create("testType", null, "testTopic2", "testGroup", null); + dataStreams.add( + new StatsPoint(tags2, 3, 4, 6, timeSource.getCurrentTimeNanos(), 0, 0, 0, null)); + timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.report(); + + awaitBuckets(dataStreams, payloadWriter, 2); + + StatsBucket bucket0 = payloadWriter.buckets.get(0); + assertEquals(1, bucket0.getGroups().size()); + StatsGroup group0 = bucket0.getGroups().iterator().next(); + assertEquals(3, group0.getTags().nonNullSize()); + assertEquals("type:testType", group0.getTags().getType()); + assertEquals("group:testGroup", group0.getTags().getGroup()); + assertEquals("topic:testTopic", group0.getTags().getTopic()); + assertEquals(1L, group0.getHash()); + assertEquals(2L, group0.getParentHash()); + + StatsBucket bucket1 = payloadWriter.buckets.get(1); + assertEquals(1, bucket1.getGroups().size()); + StatsGroup group1 = bucket1.getGroups().iterator().next(); + assertEquals("type:testType", group1.getTags().getType()); + assertEquals("group:testGroup", group1.getTags().getGroup()); + assertEquals("topic:testTopic2", group1.getTags().getTopic()); + assertEquals(3, group1.getTags().nonNullSize()); + assertEquals(3L, group1.getHash()); + assertEquals(4L, group1.getParentHash()); + + // cleanup + payloadWriter.close(); + dataStreams.close(); + } + + @Test + void multiplePointsAreCorrectlyGroupedInMultipleBuckets() throws InterruptedException { + DDAgentFeaturesDiscovery features = stubFeatures(true); + ControllableTimeSource timeSource = new ControllableTimeSource(); + Sink sink = mock(Sink.class); + CapturingPayloadWriter payloadWriter = new CapturingPayloadWriter(); + TraceConfig traceConfig = stubTraceConfig(true); + + DataStreamsTags tags = DataStreamsTags.create("testType", null, "testTopic", "testGroup", null); + DefaultDataStreamsMonitoring dataStreams = + new DefaultDataStreamsMonitoring( + sink, + features, + timeSource, + () -> traceConfig, + payloadWriter, + DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.start(); + dataStreams.add(new StatsPoint(tags, 1, 2, 1, timeSource.getCurrentTimeNanos(), 0, 0, 0, null)); + timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS - 100L); + dataStreams.add( + new StatsPoint( + tags, + 1, + 2, + 1, + timeSource.getCurrentTimeNanos(), + SECONDS.toNanos(10), + SECONDS.toNanos(10), + 10, + null)); + timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.add( + new StatsPoint( + tags, + 1, + 2, + 1, + timeSource.getCurrentTimeNanos(), + SECONDS.toNanos(5), + SECONDS.toNanos(5), + 5, + null)); + dataStreams.add( + new StatsPoint( + tags, 3, 4, 5, timeSource.getCurrentTimeNanos(), SECONDS.toNanos(2), 0, 0, null)); + timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.report(); + + awaitBuckets(dataStreams, payloadWriter, 2); + + StatsBucket bucket0 = payloadWriter.buckets.get(0); + assertEquals(1, bucket0.getGroups().size()); + StatsGroup group0 = bucket0.getGroups().iterator().next(); + assertEquals("type:testType", group0.getTags().getType()); + assertEquals("group:testGroup", group0.getTags().getGroup()); + assertEquals("topic:testTopic", group0.getTags().getTopic()); + assertEquals(3, group0.getTags().nonNullSize()); + assertEquals(1L, group0.getHash()); + assertEquals(2L, group0.getParentHash()); + assertTrue(Math.abs((group0.getPathwayLatency().getMaxValue() - 10) / 10) < 0.01); + + StatsBucket bucket1 = payloadWriter.buckets.get(1); + assertEquals(2, bucket1.getGroups().size()); + List sortedGroups = new ArrayList<>(bucket1.getGroups()); + sortedGroups.sort(Comparator.comparingLong(StatsGroup::getHash)); + + StatsGroup sortedGroup0 = sortedGroups.get(0); + assertEquals(1L, sortedGroup0.getHash()); + assertEquals(2L, sortedGroup0.getParentHash()); + assertEquals("type:testType", sortedGroup0.getTags().getType()); + assertEquals("group:testGroup", sortedGroup0.getTags().getGroup()); + assertEquals("topic:testTopic", sortedGroup0.getTags().getTopic()); + assertEquals(3, sortedGroup0.getTags().nonNullSize()); + assertTrue(Math.abs((sortedGroup0.getPathwayLatency().getMaxValue() - 5) / 5) < 0.01); + + StatsGroup sortedGroup1 = sortedGroups.get(1); + assertEquals(3L, sortedGroup1.getHash()); + assertEquals(4L, sortedGroup1.getParentHash()); + assertEquals("type:testType", sortedGroup1.getTags().getType()); + assertEquals("group:testGroup", sortedGroup1.getTags().getGroup()); + assertEquals("topic:testTopic", sortedGroup1.getTags().getTopic()); + assertEquals(3, sortedGroup1.getTags().nonNullSize()); + assertTrue(Math.abs((sortedGroup1.getPathwayLatency().getMaxValue() - 2) / 2) < 0.01); + + // cleanup + payloadWriter.close(); + dataStreams.close(); + } + + @Test + void featureUpgrade() throws InterruptedException { + boolean[] supportsDataStreaming = new boolean[] {false}; + DDAgentFeaturesDiscovery features = mock(DDAgentFeaturesDiscovery.class, RETURNS_SMART_NULLS); + when(features.supportsDataStreams()).thenAnswer(invocation -> supportsDataStreaming[0]); + + ControllableTimeSource timeSource = new ControllableTimeSource(); + Sink sink = mock(Sink.class); + CapturingPayloadWriter payloadWriter = new CapturingPayloadWriter(); + TraceConfig traceConfig = stubTraceConfig(true); + + // reporting points when data streams is not supported + DefaultDataStreamsMonitoring dataStreams = + new DefaultDataStreamsMonitoring( + sink, + features, + timeSource, + () -> traceConfig, + payloadWriter, + DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.start(); + DataStreamsTags tags = DataStreamsTags.create("testType", null, "testTopic", "testGroup", null); + dataStreams.add(new StatsPoint(tags, 1, 2, 3, timeSource.getCurrentTimeNanos(), 0, 0, 0, null)); + timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.report(); + + // no buckets are reported + awaitIdle(dataStreams); + assertTrue(payloadWriter.buckets.isEmpty()); + + // report called multiple times without advancing past check interval + dataStreams.report(); + dataStreams.report(); + dataStreams.report(); + + // features are not rechecked + awaitIdle(dataStreams); + verify(features, org.mockito.Mockito.never()).discover(); + assertTrue(payloadWriter.buckets.isEmpty()); + + // submitting points after an upgrade + supportsDataStreaming[0] = true; + timeSource.advance(FEATURE_CHECK_INTERVAL_NANOS); + dataStreams.report(); + + dataStreams.add(new StatsPoint(tags, 1, 2, 3, timeSource.getCurrentTimeNanos(), 0, 0, 0, null)); + timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.report(); + + // points are now reported + awaitBuckets(dataStreams, payloadWriter, 1); + + StatsBucket bucket = payloadWriter.buckets.get(0); + assertEquals(1, bucket.getGroups().size()); + StatsGroup group = bucket.getGroups().iterator().next(); + assertEquals("type:testType", group.getTags().getType()); + assertEquals("group:testGroup", group.getTags().getGroup()); + assertEquals("topic:testTopic", group.getTags().getTopic()); + assertEquals(3, group.getTags().nonNullSize()); + assertEquals(1L, group.getHash()); + assertEquals(2L, group.getParentHash()); + + // cleanup + payloadWriter.close(); + dataStreams.close(); + } + + @Test + void featureDowngradeThenUpgrade() throws InterruptedException { + boolean[] supportsDataStreaming = new boolean[] {true}; + DDAgentFeaturesDiscovery features = mock(DDAgentFeaturesDiscovery.class, RETURNS_SMART_NULLS); + when(features.supportsDataStreams()).thenAnswer(invocation -> supportsDataStreaming[0]); + + ControllableTimeSource timeSource = new ControllableTimeSource(); + Sink sink = mock(Sink.class); + CapturingPayloadWriter payloadWriter = new CapturingPayloadWriter(); + TraceConfig traceConfig = stubTraceConfig(true); + + // reporting points after a downgrade + DefaultDataStreamsMonitoring dataStreams = + new DefaultDataStreamsMonitoring( + sink, + features, + timeSource, + () -> traceConfig, + payloadWriter, + DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.start(); + supportsDataStreaming[0] = false; + dataStreams.onEvent(EventListener.EventType.DOWNGRADED, ""); + DataStreamsTags tags = DataStreamsTags.create("testType", null, "testTopic", "testGroup", null); + dataStreams.add(new StatsPoint(tags, 1, 2, 3, timeSource.getCurrentTimeNanos(), 0, 0, 0, null)); + timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.report(); + + // no buckets are reported + awaitIdle(dataStreams); + assertTrue(payloadWriter.buckets.isEmpty()); + + // submitting points after an upgrade + supportsDataStreaming[0] = true; + timeSource.advance(FEATURE_CHECK_INTERVAL_NANOS); + dataStreams.report(); + + dataStreams.add(new StatsPoint(tags, 1, 2, 3, timeSource.getCurrentTimeNanos(), 0, 0, 0, null)); + timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.report(); + + // points are now reported + awaitBuckets(dataStreams, payloadWriter, 1); + + StatsBucket bucket = payloadWriter.buckets.get(0); + assertEquals(1, bucket.getGroups().size()); + StatsGroup group = bucket.getGroups().iterator().next(); + assertEquals("type:testType", group.getTags().getType()); + assertEquals("group:testGroup", group.getTags().getGroup()); + assertEquals("topic:testTopic", group.getTags().getTopic()); + assertEquals(3, group.getTags().nonNullSize()); + assertEquals(1L, group.getHash()); + assertEquals(2L, group.getParentHash()); + + // cleanup + payloadWriter.close(); + dataStreams.close(); + } + + @Test + void dynamicConfigEnableAndDisable() throws InterruptedException { + boolean[] supportsDataStreaming = new boolean[] {true}; + DDAgentFeaturesDiscovery features = mock(DDAgentFeaturesDiscovery.class, RETURNS_SMART_NULLS); + when(features.supportsDataStreams()).thenAnswer(invocation -> supportsDataStreaming[0]); + + ControllableTimeSource timeSource = new ControllableTimeSource(); + Sink sink = mock(Sink.class); + CapturingPayloadWriter payloadWriter = new CapturingPayloadWriter(); + boolean[] dsmEnabled = new boolean[] {false}; + TraceConfig traceConfig = mock(TraceConfig.class, RETURNS_SMART_NULLS); + when(traceConfig.isDataStreamsEnabled()).thenAnswer(invocation -> dsmEnabled[0]); + + // reporting points when data streams is not enabled + DataStreamsTags tags = DataStreamsTags.create("testType", null, "testTopic", "testGroup", null); + DefaultDataStreamsMonitoring dataStreams = + new DefaultDataStreamsMonitoring( + sink, + features, + timeSource, + () -> traceConfig, + payloadWriter, + DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.start(); + dataStreams.add(new StatsPoint(tags, 1, 2, 3, timeSource.getCurrentTimeNanos(), 0, 0, 0, null)); + timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.report(); + + // no buckets are reported + awaitIdle(dataStreams); + assertTrue(payloadWriter.buckets.isEmpty()); + + // submitting points after dynamically enabled + dsmEnabled[0] = true; + dataStreams.report(); + + dataStreams.add(new StatsPoint(tags, 1, 2, 3, timeSource.getCurrentTimeNanos(), 0, 0, 0, null)); + timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.report(); + + // points are now reported + awaitBuckets(dataStreams, payloadWriter, 1); + + StatsBucket bucket = payloadWriter.buckets.get(0); + assertEquals(1, bucket.getGroups().size()); + StatsGroup group = bucket.getGroups().iterator().next(); + assertEquals("type:testType", group.getTags().getType()); + assertEquals("group:testGroup", group.getTags().getGroup()); + assertEquals("topic:testTopic", group.getTags().getTopic()); + assertEquals(3, group.getTags().nonNullSize()); + assertEquals(1L, group.getHash()); + assertEquals(2L, group.getParentHash()); + + // disabling data streams dynamically + dsmEnabled[0] = false; + dataStreams.report(); + + // inbox is processed + awaitIdle(dataStreams); + + // submitting points after being disabled + payloadWriter.buckets.clear(); + + dataStreams.add(new StatsPoint(tags, 1, 2, 3, timeSource.getCurrentTimeNanos(), 0, 0, 0, null)); + timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.report(); + + // points are no longer reported + awaitIdle(dataStreams); + assertTrue(payloadWriter.buckets.isEmpty()); + + // cleanup + payloadWriter.close(); + dataStreams.close(); + } + + @Test + void featureAndDynamicConfigUpgradeInteractions() throws InterruptedException { + boolean[] supportsDataStreaming = new boolean[] {false}; + DDAgentFeaturesDiscovery features = mock(DDAgentFeaturesDiscovery.class, RETURNS_SMART_NULLS); + when(features.supportsDataStreams()).thenAnswer(invocation -> supportsDataStreaming[0]); + + ControllableTimeSource timeSource = new ControllableTimeSource(); + Sink sink = mock(Sink.class); + CapturingPayloadWriter payloadWriter = new CapturingPayloadWriter(); + boolean[] dsmEnabled = new boolean[] {false}; + TraceConfig traceConfig = mock(TraceConfig.class, RETURNS_SMART_NULLS); + when(traceConfig.isDataStreamsEnabled()).thenAnswer(invocation -> dsmEnabled[0]); + + // reporting points when data streams is not supported + DefaultDataStreamsMonitoring dataStreams = + new DefaultDataStreamsMonitoring( + sink, + features, + timeSource, + () -> traceConfig, + payloadWriter, + DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.start(); + DataStreamsTags tags = DataStreamsTags.create("testType", null, "testTopic", "testGroup", null); + dataStreams.add(new StatsPoint(tags, 1, 2, 3, timeSource.getCurrentTimeNanos(), 0, 0, 0, null)); + timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.report(); + + // no buckets are reported + awaitIdle(dataStreams); + assertTrue(payloadWriter.buckets.isEmpty()); + + // submitting points after an upgrade with dsm disabled + supportsDataStreaming[0] = true; + timeSource.advance(FEATURE_CHECK_INTERVAL_NANOS); + dataStreams.report(); + + dataStreams.add(new StatsPoint(tags, 1, 2, 3, timeSource.getCurrentTimeNanos(), 0, 0, 0, null)); + timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.report(); + + // points are not reported + awaitIdle(dataStreams); + assertTrue(payloadWriter.buckets.isEmpty()); + + // dsm is enabled dynamically + dsmEnabled[0] = true; + dataStreams.report(); + + dataStreams.add(new StatsPoint(tags, 1, 2, 3, timeSource.getCurrentTimeNanos(), 0, 0, 0, null)); + timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.report(); + + // points are now reported + awaitBuckets(dataStreams, payloadWriter, 1); + + StatsBucket bucket = payloadWriter.buckets.get(0); + assertEquals(1, bucket.getGroups().size()); + StatsGroup group = bucket.getGroups().iterator().next(); + assertEquals("type:testType", group.getTags().getType()); + assertEquals("group:testGroup", group.getTags().getGroup()); + assertEquals("topic:testTopic", group.getTags().getTopic()); + assertEquals(3, group.getTags().nonNullSize()); + assertEquals(1L, group.getHash()); + assertEquals(2L, group.getParentHash()); + + // cleanup + payloadWriter.close(); + dataStreams.close(); + } + + @Test + void moreFeatureAndDynamicConfigUpgradeInteractions() throws InterruptedException { + boolean[] supportsDataStreaming = new boolean[] {false}; + DDAgentFeaturesDiscovery features = mock(DDAgentFeaturesDiscovery.class, RETURNS_SMART_NULLS); + when(features.supportsDataStreams()).thenAnswer(invocation -> supportsDataStreaming[0]); + + ControllableTimeSource timeSource = new ControllableTimeSource(); + Sink sink = mock(Sink.class); + CapturingPayloadWriter payloadWriter = new CapturingPayloadWriter(); + boolean[] dsmEnabled = new boolean[] {false}; + TraceConfig traceConfig = mock(TraceConfig.class, RETURNS_SMART_NULLS); + when(traceConfig.isDataStreamsEnabled()).thenAnswer(invocation -> dsmEnabled[0]); + + // reporting points when data streams is not supported + DefaultDataStreamsMonitoring dataStreams = + new DefaultDataStreamsMonitoring( + sink, + features, + timeSource, + () -> traceConfig, + payloadWriter, + DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.start(); + DataStreamsTags tags = DataStreamsTags.create("testType", null, "testTopic", "testGroup", null); + dataStreams.add(new StatsPoint(tags, 1, 2, 3, timeSource.getCurrentTimeNanos(), 0, 0, 0, null)); + timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.report(); + + // no buckets are reported + awaitIdle(dataStreams); + assertTrue(payloadWriter.buckets.isEmpty()); + + // enabling dsm when not supported by agent + dsmEnabled[0] = true; + dataStreams.report(); + + dataStreams.add(new StatsPoint(tags, 1, 2, 3, timeSource.getCurrentTimeNanos(), 0, 0, 0, null)); + timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.report(); + + // points are not reported + awaitIdle(dataStreams); + assertTrue(payloadWriter.buckets.isEmpty()); + + // cleanup + payloadWriter.close(); + dataStreams.close(); + } + + @Test + void schemaRegistryUsagesAreAggregatedByOperation() throws InterruptedException { + DDAgentFeaturesDiscovery features = stubFeatures(true); + ControllableTimeSource timeSource = new ControllableTimeSource(); + CapturingPayloadWriter payloadWriter = new CapturingPayloadWriter(); + Sink sink = mock(Sink.class); + TraceConfig traceConfig = stubTraceConfig(true); + + DefaultDataStreamsMonitoring dataStreams = + new DefaultDataStreamsMonitoring( + sink, + features, + timeSource, + () -> traceConfig, + payloadWriter, + DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.start(); + + // Record serialize and deserialize operations + dataStreams.reportSchemaRegistryUsage( + "test-topic", "test-cluster", 123, true, false, "serialize"); + // duplicate serialize + dataStreams.reportSchemaRegistryUsage( + "test-topic", "test-cluster", 123, true, false, "serialize"); + dataStreams.reportSchemaRegistryUsage( + "test-topic", "test-cluster", 123, true, false, "deserialize"); + // different schema/key + dataStreams.reportSchemaRegistryUsage( + "test-topic", "test-cluster", 456, true, true, "serialize"); + + timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.report(); + + awaitBuckets(dataStreams, payloadWriter, 1); + + StatsBucket bucket = payloadWriter.buckets.get(0); + Collection> usages = bucket.getSchemaRegistryUsages(); + assertEquals(3, usages.size()); // 3 unique combinations + + // Find serialize operation for schema 123 (should have count 2) + Entry serializeUsage = findUsage(usages, 123, "serialize", false); + assertNotNull(serializeUsage); + assertEquals(2L, serializeUsage.getValue()); // Aggregated 2 serialize operations + + // Find deserialize operation for schema 123 (should have count 1) + Entry deserializeUsage = + findUsage(usages, 123, "deserialize", false); + assertNotNull(deserializeUsage); + assertEquals(1L, deserializeUsage.getValue()); + + // Find serialize operation for schema 456 with isKey=true (should have count 1) + Entry keySerializeUsage = + findUsage(usages, 456, "serialize", true); + assertNotNull(keySerializeUsage); + assertEquals(1L, keySerializeUsage.getValue()); + + // cleanup + payloadWriter.close(); + dataStreams.close(); + } + + private static Entry findUsage( + Collection> usages, + int schemaId, + String operation, + boolean isKey) { + for (Entry entry : usages) { + StatsBucket.SchemaKey key = entry.getKey(); + if (key.getSchemaId() == schemaId + && operation.equals(key.getOperation()) + && key.isKey() == isKey) { + return entry; + } + } + return null; + } + + @Test + void schemaKeyEqualsAndHashCodeWorkCorrectly() { + StatsBucket.SchemaKey key1 = + new StatsBucket.SchemaKey("topic1", "cluster1", 123, true, false, "serialize"); + StatsBucket.SchemaKey key2 = + new StatsBucket.SchemaKey("topic1", "cluster1", 123, true, false, "serialize"); + // different topic + StatsBucket.SchemaKey key3 = + new StatsBucket.SchemaKey("topic2", "cluster1", 123, true, false, "serialize"); + // different cluster + StatsBucket.SchemaKey key4 = + new StatsBucket.SchemaKey("topic1", "cluster2", 123, true, false, "serialize"); + // different schema + StatsBucket.SchemaKey key5 = + new StatsBucket.SchemaKey("topic1", "cluster1", 456, true, false, "serialize"); + // different success + StatsBucket.SchemaKey key6 = + new StatsBucket.SchemaKey("topic1", "cluster1", 123, false, false, "serialize"); + // different isKey + StatsBucket.SchemaKey key7 = + new StatsBucket.SchemaKey("topic1", "cluster1", 123, true, true, "serialize"); + // different operation + StatsBucket.SchemaKey key8 = + new StatsBucket.SchemaKey("topic1", "cluster1", 123, true, false, "deserialize"); + + // Reflexive + assertEquals(key1, key1); + assertEquals(key1.hashCode(), key1.hashCode()); + + // Symmetric + assertEquals(key1, key2); + assertEquals(key2, key1); + assertEquals(key1.hashCode(), key2.hashCode()); + + // Different topic + assertNotEquals(key1, key3); + assertNotEquals(key3, key1); + + // Different cluster + assertNotEquals(key1, key4); + assertNotEquals(key4, key1); + + // Different schema ID + assertNotEquals(key1, key5); + assertNotEquals(key5, key1); + + // Different success + assertNotEquals(key1, key6); + assertNotEquals(key6, key1); + + // Different isKey + assertNotEquals(key1, key7); + assertNotEquals(key7, key1); + + // Different operation + assertNotEquals(key1, key8); + assertNotEquals(key8, key1); + + // Null check + assertNotEquals(null, key1); + + // Different class + assertNotEquals("not a schema key", key1); + } + + @Test + void statsBucketAggregatesSchemaRegistryUsagesCorrectly() { + StatsBucket bucket = new StatsBucket(1000L, 10000L); + SchemaRegistryUsage usage1 = + new SchemaRegistryUsage("topic1", "cluster1", 123, true, false, "serialize", 1000L, null); + SchemaRegistryUsage usage2 = + new SchemaRegistryUsage("topic1", "cluster1", 123, true, false, "serialize", 2000L, null); + SchemaRegistryUsage usage3 = + new SchemaRegistryUsage("topic1", "cluster1", 123, true, false, "deserialize", 3000L, null); + + bucket.addSchemaRegistryUsage(usage1); + // should increment count for same key + bucket.addSchemaRegistryUsage(usage2); + // different operation, new key + bucket.addSchemaRegistryUsage(usage3); + + Collection> usages = bucket.getSchemaRegistryUsages(); + Map usageMap = new HashMap<>(); + for (Entry entry : usages) { + usageMap.put(entry.getKey(), entry.getValue()); + } + + assertEquals(2, usages.size()); + + // Check serialize count + StatsBucket.SchemaKey serializeKey = + new StatsBucket.SchemaKey("topic1", "cluster1", 123, true, false, "serialize"); + assertEquals(2L, usageMap.get(serializeKey)); + + // Check deserialize count + StatsBucket.SchemaKey deserializeKey = + new StatsBucket.SchemaKey("topic1", "cluster1", 123, true, false, "deserialize"); + assertEquals(1L, usageMap.get(deserializeKey)); + + // Check that different operations create different keys + assertNotEquals(serializeKey, deserializeKey); + } + + @Test + void kafkaProducerConfigIsReportedInBucket() throws InterruptedException { + DDAgentFeaturesDiscovery features = stubFeatures(true); + ControllableTimeSource timeSource = new ControllableTimeSource(); + Sink sink = mock(Sink.class); + CapturingPayloadWriter payloadWriter = new CapturingPayloadWriter(); + TraceConfig traceConfig = stubTraceConfig(true); + + DefaultDataStreamsMonitoring dataStreams = + new DefaultDataStreamsMonitoring( + sink, + features, + timeSource, + () -> traceConfig, + payloadWriter, + DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.start(); + Map kafkaConfig = new HashMap<>(); + kafkaConfig.put("bootstrap.servers", "localhost:9092"); + kafkaConfig.put("acks", "all"); + dataStreams.reportKafkaConfig("kafka_producer", "", "", kafkaConfig); + timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.report(); + + awaitBuckets(dataStreams, payloadWriter, 1); + + StatsBucket bucket = payloadWriter.buckets.get(0); + List kafkaConfigs = bucket.getKafkaConfigs(); + assertEquals(1, kafkaConfigs.size()); + KafkaConfigReport report = kafkaConfigs.get(0); + assertEquals("kafka_producer", report.getType()); + assertEquals("localhost:9092", report.getConfig().get("bootstrap.servers")); + assertEquals("all", report.getConfig().get("acks")); + + // cleanup + payloadWriter.close(); + dataStreams.close(); + } + + @Test + void kafkaConsumerConfigIsReportedInBucket() throws InterruptedException { + DDAgentFeaturesDiscovery features = stubFeatures(true); + ControllableTimeSource timeSource = new ControllableTimeSource(); + Sink sink = mock(Sink.class); + CapturingPayloadWriter payloadWriter = new CapturingPayloadWriter(); + TraceConfig traceConfig = stubTraceConfig(true); + + DefaultDataStreamsMonitoring dataStreams = + new DefaultDataStreamsMonitoring( + sink, + features, + timeSource, + () -> traceConfig, + payloadWriter, + DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.start(); + Map kafkaConfig = new HashMap<>(); + kafkaConfig.put("bootstrap.servers", "localhost:9092"); + kafkaConfig.put("group.id", "test-group"); + kafkaConfig.put("auto.offset.reset", "earliest"); + dataStreams.reportKafkaConfig("kafka_consumer", "", "test-group", kafkaConfig); + timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.report(); + + awaitBuckets(dataStreams, payloadWriter, 1); + + StatsBucket bucket = payloadWriter.buckets.get(0); + List kafkaConfigs = bucket.getKafkaConfigs(); + assertEquals(1, kafkaConfigs.size()); + KafkaConfigReport report = kafkaConfigs.get(0); + assertEquals("kafka_consumer", report.getType()); + assertEquals("localhost:9092", report.getConfig().get("bootstrap.servers")); + assertEquals("test-group", report.getConfig().get("group.id")); + assertEquals("earliest", report.getConfig().get("auto.offset.reset")); + + // cleanup + payloadWriter.close(); + dataStreams.close(); + } + + @Test + void duplicateKafkaConfigsAreEachReportedInTheBucket() throws InterruptedException { + DDAgentFeaturesDiscovery features = stubFeatures(true); + ControllableTimeSource timeSource = new ControllableTimeSource(); + Sink sink = mock(Sink.class); + CapturingPayloadWriter payloadWriter = new CapturingPayloadWriter(); + TraceConfig traceConfig = stubTraceConfig(true); + + // reporting the same config twice + DefaultDataStreamsMonitoring dataStreams = + new DefaultDataStreamsMonitoring( + sink, + features, + timeSource, + () -> traceConfig, + payloadWriter, + DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.start(); + Map config1 = new HashMap<>(); + config1.put("bootstrap.servers", "localhost:9092"); + config1.put("acks", "all"); + Map config2 = new HashMap<>(); + config2.put("bootstrap.servers", "localhost:9092"); + config2.put("acks", "all"); + dataStreams.reportKafkaConfig("kafka_producer", "", "", config1); + dataStreams.reportKafkaConfig("kafka_producer", "", "", config2); + timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.report(); + + // both configs are reported in the bucket + awaitBuckets(dataStreams, payloadWriter, 1); + + StatsBucket bucket = payloadWriter.buckets.get(0); + List kafkaConfigs = bucket.getKafkaConfigs(); + assertEquals(2, kafkaConfigs.size()); + for (KafkaConfigReport report : kafkaConfigs) { + assertEquals("kafka_producer", report.getType()); + assertEquals("localhost:9092", report.getConfig().get("bootstrap.servers")); + assertEquals("all", report.getConfig().get("acks")); + } + + // cleanup + payloadWriter.close(); + dataStreams.close(); + } + + @Test + void kafkaConfigsReportedInSeparateBucketsAppearInEachBucket() throws InterruptedException { + DDAgentFeaturesDiscovery features = stubFeatures(true); + ControllableTimeSource timeSource = new ControllableTimeSource(); + Sink sink = mock(Sink.class); + CapturingPayloadWriter payloadWriter = new CapturingPayloadWriter(); + TraceConfig traceConfig = stubTraceConfig(true); + + // reporting a config in the first bucket + DefaultDataStreamsMonitoring dataStreams = + new DefaultDataStreamsMonitoring( + sink, + features, + timeSource, + () -> traceConfig, + payloadWriter, + DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.start(); + Map kafkaConfig = new HashMap<>(); + kafkaConfig.put("bootstrap.servers", "localhost:9092"); + kafkaConfig.put("acks", "all"); + dataStreams.reportKafkaConfig("kafka_producer", "", "", kafkaConfig); + timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.report(); + + // first bucket has the config + awaitBuckets(dataStreams, payloadWriter, 1); + assertEquals(1, payloadWriter.buckets.get(0).getKafkaConfigs().size()); + + // reporting the same config again in a new bucket + payloadWriter.buckets.clear(); + dataStreams.reportKafkaConfig("kafka_producer", "", "", kafkaConfig); + timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.report(); + + // second bucket also has the config + awaitBuckets(dataStreams, payloadWriter, 1); + + StatsBucket bucket = payloadWriter.buckets.get(0); + List kafkaConfigs = bucket.getKafkaConfigs(); + assertEquals(1, kafkaConfigs.size()); + KafkaConfigReport report = kafkaConfigs.get(0); + assertEquals("kafka_producer", report.getType()); + assertEquals("localhost:9092", report.getConfig().get("bootstrap.servers")); + assertEquals("all", report.getConfig().get("acks")); + + // cleanup + payloadWriter.close(); + dataStreams.close(); + } + + @Test + void differentKafkaConfigsAreBothReported() throws InterruptedException { + DDAgentFeaturesDiscovery features = stubFeatures(true); + ControllableTimeSource timeSource = new ControllableTimeSource(); + Sink sink = mock(Sink.class); + CapturingPayloadWriter payloadWriter = new CapturingPayloadWriter(); + TraceConfig traceConfig = stubTraceConfig(true); + + // reporting producer and consumer configs + DefaultDataStreamsMonitoring dataStreams = + new DefaultDataStreamsMonitoring( + sink, + features, + timeSource, + () -> traceConfig, + payloadWriter, + DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.start(); + Map producerConfig = new HashMap<>(); + producerConfig.put("bootstrap.servers", "localhost:9092"); + producerConfig.put("acks", "all"); + dataStreams.reportKafkaConfig("kafka_producer", "", "", producerConfig); + Map consumerConfig = new HashMap<>(); + consumerConfig.put("bootstrap.servers", "localhost:9092"); + consumerConfig.put("group.id", "my-group"); + dataStreams.reportKafkaConfig("kafka_consumer", "", "my-group", consumerConfig); + timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.report(); + + // both configs are reported + awaitBuckets(dataStreams, payloadWriter, 1); + + StatsBucket bucket = payloadWriter.buckets.get(0); + List kafkaConfigs = bucket.getKafkaConfigs(); + assertEquals(2, kafkaConfigs.size()); + + KafkaConfigReport producerReport = findReport(kafkaConfigs, "kafka_producer"); + assertNotNull(producerReport); + assertEquals("localhost:9092", producerReport.getConfig().get("bootstrap.servers")); + assertEquals("all", producerReport.getConfig().get("acks")); + + KafkaConfigReport consumerReport = findReport(kafkaConfigs, "kafka_consumer"); + assertNotNull(consumerReport); + assertEquals("localhost:9092", consumerReport.getConfig().get("bootstrap.servers")); + assertEquals("my-group", consumerReport.getConfig().get("group.id")); + + // cleanup + payloadWriter.close(); + dataStreams.close(); + } + + private static KafkaConfigReport findReport(List configs, String type) { + for (KafkaConfigReport report : configs) { + if (type.equals(report.getType())) { + return report; + } + } + return null; + } + + @Test + void kafkaConfigsWithDifferentValuesForSameTypeAreNotDeduplicated() throws InterruptedException { + DDAgentFeaturesDiscovery features = stubFeatures(true); + ControllableTimeSource timeSource = new ControllableTimeSource(); + Sink sink = mock(Sink.class); + CapturingPayloadWriter payloadWriter = new CapturingPayloadWriter(); + TraceConfig traceConfig = stubTraceConfig(true); + + // reporting two producer configs with different settings + DefaultDataStreamsMonitoring dataStreams = + new DefaultDataStreamsMonitoring( + sink, + features, + timeSource, + () -> traceConfig, + payloadWriter, + DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.start(); + Map config1 = new HashMap<>(); + config1.put("bootstrap.servers", "localhost:9092"); + config1.put("acks", "all"); + dataStreams.reportKafkaConfig("kafka_producer", "", "", config1); + Map config2 = new HashMap<>(); + config2.put("bootstrap.servers", "localhost:9093"); + config2.put("acks", "1"); + dataStreams.reportKafkaConfig("kafka_producer", "", "", config2); + timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.report(); + + // both configs are reported because they have different values + awaitBuckets(dataStreams, payloadWriter, 1); + assertEquals(2, payloadWriter.buckets.get(0).getKafkaConfigs().size()); + + // cleanup + payloadWriter.close(); + dataStreams.close(); + } + + @Test + void kafkaConfigsAreReportedAlongsideStatsPoints() throws InterruptedException { + DDAgentFeaturesDiscovery features = stubFeatures(true); + ControllableTimeSource timeSource = new ControllableTimeSource(); + Sink sink = mock(Sink.class); + CapturingPayloadWriter payloadWriter = new CapturingPayloadWriter(); + TraceConfig traceConfig = stubTraceConfig(true); + + // reporting both stats points and kafka configs + DefaultDataStreamsMonitoring dataStreams = + new DefaultDataStreamsMonitoring( + sink, + features, + timeSource, + () -> traceConfig, + payloadWriter, + DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.start(); + DataStreamsTags tags = DataStreamsTags.create("testType", null, "testTopic", "testGroup", null); + dataStreams.add(new StatsPoint(tags, 1, 2, 3, timeSource.getCurrentTimeNanos(), 0, 0, 0, null)); + Map kafkaConfig = new HashMap<>(); + kafkaConfig.put("bootstrap.servers", "localhost:9092"); + dataStreams.reportKafkaConfig("kafka_producer", "", "", kafkaConfig); + timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.report(); + + // bucket contains both stats groups and kafka configs + awaitBuckets(dataStreams, payloadWriter, 1); + + StatsBucket bucket = payloadWriter.buckets.get(0); + assertEquals(1, bucket.getGroups().size()); + List kafkaConfigs = bucket.getKafkaConfigs(); + assertEquals(1, kafkaConfigs.size()); + + StatsGroup group = bucket.getGroups().iterator().next(); + assertEquals("type:testType", group.getTags().getType()); + assertEquals(1L, group.getHash()); + assertEquals(2L, group.getParentHash()); + + KafkaConfigReport report = kafkaConfigs.get(0); + assertEquals("kafka_producer", report.getType()); + assertEquals("localhost:9092", report.getConfig().get("bootstrap.servers")); + + // cleanup + payloadWriter.close(); + dataStreams.close(); + } + + @TableTest({ + "scenario | enabledAtAgent | enabledInConfig", + "agent off, config on | false | true ", + "agent on, config off | true | false ", + "both off | false | false " + }) + void kafkaConfigsNotReportedWhenDSMIsDisabled(boolean enabledAtAgent, boolean enabledInConfig) + throws InterruptedException { + DDAgentFeaturesDiscovery features = stubFeatures(enabledAtAgent); + ControllableTimeSource timeSource = new ControllableTimeSource(); + CapturingPayloadWriter payloadWriter = new CapturingPayloadWriter(); + Sink sink = mock(Sink.class); + TraceConfig traceConfig = stubTraceConfig(enabledInConfig); + + DefaultDataStreamsMonitoring dataStreams = + new DefaultDataStreamsMonitoring( + sink, + features, + timeSource, + () -> traceConfig, + payloadWriter, + DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.start(); + Map kafkaConfig = new HashMap<>(); + kafkaConfig.put("bootstrap.servers", "localhost:9092"); + dataStreams.reportKafkaConfig("kafka_producer", "", "", kafkaConfig); + timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.report(); + + awaitIdle(dataStreams); + assertTrue(payloadWriter.buckets.isEmpty()); + + // cleanup + payloadWriter.close(); + dataStreams.close(); + } + + @Test + void kafkaConfigsFlushedOnClose() throws InterruptedException { + DDAgentFeaturesDiscovery features = stubFeatures(true); + ControllableTimeSource timeSource = new ControllableTimeSource(); + Sink sink = mock(Sink.class); + CapturingPayloadWriter payloadWriter = new CapturingPayloadWriter(); + TraceConfig traceConfig = stubTraceConfig(true); + + DefaultDataStreamsMonitoring dataStreams = + new DefaultDataStreamsMonitoring( + sink, + features, + timeSource, + () -> traceConfig, + payloadWriter, + DEFAULT_BUCKET_DURATION_NANOS); + dataStreams.start(); + Map kafkaConfig = new HashMap<>(); + kafkaConfig.put("bootstrap.servers", "localhost:9092"); + dataStreams.reportKafkaConfig("kafka_producer", "", "", kafkaConfig); + timeSource.advance(DEFAULT_BUCKET_DURATION_NANOS - 100L); + dataStreams.close(); + + // configs in the current bucket are flushed on close + awaitBuckets(dataStreams, payloadWriter, 1); + + StatsBucket bucket = payloadWriter.buckets.get(0); + List kafkaConfigs = bucket.getKafkaConfigs(); + assertEquals(1, kafkaConfigs.size()); + KafkaConfigReport report = kafkaConfigs.get(0); + assertEquals("kafka_producer", report.getType()); + assertEquals("localhost:9092", report.getConfig().get("bootstrap.servers")); + + // cleanup + payloadWriter.close(); + } + + @Test + void kafkaConfigReportEqualsAndHashCodeWorkCorrectly() { + Map sameConfig = new HashMap<>(); + sameConfig.put("bootstrap.servers", "localhost:9092"); + sameConfig.put("acks", "all"); + Map sameConfig2 = new HashMap<>(); + sameConfig2.put("bootstrap.servers", "localhost:9092"); + sameConfig2.put("acks", "all"); + Map differentConfig = new HashMap<>(); + differentConfig.put("bootstrap.servers", "localhost:9093"); + + KafkaConfigReport config1 = + new KafkaConfigReport("kafka_producer", "", "", sameConfig, 1000L, null); + // different timestamp + KafkaConfigReport config2 = + new KafkaConfigReport("kafka_producer", "", "", sameConfig2, 2000L, null); + // different type + KafkaConfigReport config3 = + new KafkaConfigReport("kafka_consumer", "", "", sameConfig, 1000L, null); + // different config values + KafkaConfigReport config4 = + new KafkaConfigReport("kafka_producer", "", "", differentConfig, 1000L, null); + // different serviceNameOverride + KafkaConfigReport config5 = + new KafkaConfigReport("kafka_producer", "", "", sameConfig, 1000L, "other-service"); + + // Reflexive + assertEquals(config1, config1); + assertEquals(config1.hashCode(), config1.hashCode()); + + // Same type and config, different timestamp -- equals (timestamp is NOT part of equals) + assertEquals(config1, config2); + assertEquals(config2, config1); + assertEquals(config1.hashCode(), config2.hashCode()); + + // Same type and config, different serviceNameOverride -- equals (serviceNameOverride is NOT + // part of equals) + assertEquals(config1, config5); + assertEquals(config5, config1); + assertEquals(config1.hashCode(), config5.hashCode()); + + // Different type + assertNotEquals(config1, config3); + assertNotEquals(config3, config1); + + // Different config values + assertNotEquals(config1, config4); + assertNotEquals(config4, config1); + + // Null check + assertNotEquals(null, config1); + + // Different class + assertNotEquals("not a config report", config1); + } + + @Test + void statsBucketStoresKafkaConfigs() { + StatsBucket bucket = new StatsBucket(1000L, 10000L); + Map producerCfg = new HashMap<>(); + producerCfg.put("acks", "all"); + KafkaConfigReport config1 = + new KafkaConfigReport("kafka_producer", "", "", producerCfg, 1000L, null); + Map consumerCfg = new HashMap<>(); + consumerCfg.put("group.id", "test"); + KafkaConfigReport config2 = + new KafkaConfigReport("kafka_consumer", "", "test", consumerCfg, 2000L, null); + + bucket.addKafkaConfig(config1); + bucket.addKafkaConfig(config2); + + assertEquals(2, bucket.getKafkaConfigs().size()); + assertEquals("kafka_producer", bucket.getKafkaConfigs().get(0).getType()); + assertEquals("all", bucket.getKafkaConfigs().get(0).getConfig().get("acks")); + assertEquals("kafka_consumer", bucket.getKafkaConfigs().get(1).getType()); + assertEquals("test", bucket.getKafkaConfigs().get(1).getConfig().get("group.id")); + } + + static class CapturingPayloadWriter implements DatastreamsPayloadWriter { + boolean accepting = true; + final List buckets = new ArrayList<>(); + + @Override + public void writePayload(Collection payload, String serviceNameOverride) { + if (accepting) { + buckets.addAll(payload); + } + } + + void close() { + // Stop accepting new buckets so any late submissions by the reporting thread aren't seen + accepting = false; + } + } + + static class CustomContextCarrier implements DataStreamsContextCarrier { + private final Map data = new HashMap<>(); + + @Override + public Set> entries() { + return data.entrySet(); + } + + @Override + public void set(String key, String value) { + data.put(key, value); + } + } +} diff --git a/dd-trace-core/src/test/java/datadog/trace/core/datastreams/DefaultDataStreamsMonitoringTestBridge.java b/dd-trace-core/src/test/java/datadog/trace/core/datastreams/DefaultDataStreamsMonitoringTestBridge.java new file mode 100644 index 00000000000..37c204757fd --- /dev/null +++ b/dd-trace-core/src/test/java/datadog/trace/core/datastreams/DefaultDataStreamsMonitoringTestBridge.java @@ -0,0 +1,40 @@ +package datadog.trace.core.datastreams; + +import java.lang.reflect.Field; +import org.jctools.queues.MessagePassingQueue; + +final class DefaultDataStreamsMonitoringTestBridge { + private static final Field INBOX_FIELD; + private static final Field THREAD_FIELD; + + static { + try { + INBOX_FIELD = DefaultDataStreamsMonitoring.class.getDeclaredField("inbox"); + INBOX_FIELD.setAccessible(true); + THREAD_FIELD = DefaultDataStreamsMonitoring.class.getDeclaredField("thread"); + THREAD_FIELD.setAccessible(true); + } catch (NoSuchFieldException e) { + throw new IllegalStateException(e); + } + } + + private DefaultDataStreamsMonitoringTestBridge() {} + + static boolean isInboxEmpty(DefaultDataStreamsMonitoring monitoring) { + try { + MessagePassingQueue inbox = (MessagePassingQueue) INBOX_FIELD.get(monitoring); + return inbox.isEmpty(); + } catch (IllegalAccessException e) { + throw new IllegalStateException(e); + } + } + + static Thread.State getThreadState(DefaultDataStreamsMonitoring monitoring) { + try { + Thread thread = (Thread) THREAD_FIELD.get(monitoring); + return thread.getState(); + } catch (IllegalAccessException e) { + throw new IllegalStateException(e); + } + } +} diff --git a/dd-trace-core/src/test/java/datadog/trace/core/datastreams/DefaultPathwayContextTest.java b/dd-trace-core/src/test/java/datadog/trace/core/datastreams/DefaultPathwayContextTest.java new file mode 100644 index 00000000000..013a73e2686 --- /dev/null +++ b/dd-trace-core/src/test/java/datadog/trace/core/datastreams/DefaultPathwayContextTest.java @@ -0,0 +1,735 @@ +package datadog.trace.core.datastreams; + +import static datadog.context.Context.root; +import static datadog.trace.api.TracePropagationStyle.DATADOG; +import static datadog.trace.api.datastreams.DataStreamsContext.create; +import static datadog.trace.api.datastreams.DataStreamsContext.fromTags; +import static datadog.trace.api.datastreams.PathwayContext.PROPAGATION_KEY_BASE64; +import static java.util.concurrent.TimeUnit.MILLISECONDS; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.RETURNS_SMART_NULLS; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import datadog.communication.ddagent.DDAgentFeaturesDiscovery; +import datadog.context.Context; +import datadog.context.propagation.Propagator; +import datadog.trace.api.BaseHash; +import datadog.trace.api.Config; +import datadog.trace.api.DDTraceId; +import datadog.trace.api.TagMap; +import datadog.trace.api.TraceConfig; +import datadog.trace.api.datastreams.DataStreamsTags; +import datadog.trace.api.datastreams.StatsPoint; +import datadog.trace.api.time.ControllableTimeSource; +import datadog.trace.bootstrap.instrumentation.api.AgentPropagation; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.AgentTracer; +import datadog.trace.common.metrics.Sink; +import datadog.trace.core.DDCoreJavaSpecification; +import datadog.trace.core.propagation.ExtractedContext; +import java.io.IOException; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.function.Consumer; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import org.tabletest.junit.TableTest; + +public class DefaultPathwayContextTest extends DDCoreJavaSpecification { + private static final long DEFAULT_BUCKET_DURATION_NANOS = + Config.get().getDataStreamsBucketDurationNanoseconds(); + private static final long BASE_HASH = 12L; + + private CapturingPointConsumer pointConsumer; + + @BeforeEach + void setupConsumer() { + pointConsumer = new CapturingPointConsumer(); + } + + private static void verifyFirstPoint(StatsPoint point) { + assertEquals(0L, point.getParentHash()); + assertEquals(0L, point.getPathwayLatencyNano()); + assertEquals(0L, point.getEdgeLatencyNano()); + assertEquals(0L, point.getPayloadSizeBytes()); + } + + @Test + void firstSetCheckpointStartsTheContext() { + ControllableTimeSource timeSource = new ControllableTimeSource(); + DefaultPathwayContext context = new DefaultPathwayContext(timeSource, null); + + timeSource.advance(50); + context.setCheckpoint(fromTags(DataStreamsTags.create("internal", null)), pointConsumer); + + assertTrue(context.isStarted()); + assertEquals(1, pointConsumer.points.size()); + verifyFirstPoint(pointConsumer.points.get(0)); + } + + @Test + void checkpointGenerated() { + ControllableTimeSource timeSource = new ControllableTimeSource(); + DefaultPathwayContext context = new DefaultPathwayContext(timeSource, null); + + timeSource.advance(50); + context.setCheckpoint( + fromTags(DataStreamsTags.create("internal", DataStreamsTags.Direction.INBOUND)), + pointConsumer); + timeSource.advance(25); + DataStreamsTags tags = + DataStreamsTags.create("kafka", DataStreamsTags.Direction.OUTBOUND, "topic", "group", null); + context.setCheckpoint(fromTags(tags), pointConsumer); + + assertTrue(context.isStarted()); + assertEquals(2, pointConsumer.points.size()); + verifyFirstPoint(pointConsumer.points.get(0)); + StatsPoint second = pointConsumer.points.get(1); + assertEquals("group:group", second.getTags().getGroup()); + assertEquals("topic:topic", second.getTags().getTopic()); + assertEquals("type:kafka", second.getTags().getType()); + assertEquals("direction:out", second.getTags().getDirection()); + assertEquals(4, second.getTags().nonNullSize()); + assertEquals(pointConsumer.points.get(0).getHash(), second.getParentHash()); + assertNotEquals(0L, second.getHash()); + assertEquals(25L, second.getPathwayLatencyNano()); + assertEquals(25L, second.getEdgeLatencyNano()); + } + + @Test + void checkpointWithPayloadSize() { + ControllableTimeSource timeSource = new ControllableTimeSource(); + DefaultPathwayContext context = new DefaultPathwayContext(timeSource, null); + + timeSource.advance(25); + context.setCheckpoint( + create(DataStreamsTags.create("kafka", null, "topic", "group", null), 0, 72), + pointConsumer); + + assertTrue(context.isStarted()); + assertEquals(1, pointConsumer.points.size()); + StatsPoint first = pointConsumer.points.get(0); + assertEquals("group:group", first.getTags().getGroup()); + assertEquals("topic:topic", first.getTags().getTopic()); + assertEquals("type:kafka", first.getTags().getType()); + assertEquals(3, first.getTags().nonNullSize()); + assertNotEquals(0L, first.getHash()); + assertEquals(72L, first.getPayloadSizeBytes()); + } + + @Test + void multipleCheckpointsGenerated() { + ControllableTimeSource timeSource = new ControllableTimeSource(); + DefaultPathwayContext context = new DefaultPathwayContext(timeSource, null); + + timeSource.advance(50); + context.setCheckpoint( + fromTags(DataStreamsTags.create("kafka", DataStreamsTags.Direction.OUTBOUND)), + pointConsumer); + timeSource.advance(25); + DataStreamsTags tags = + DataStreamsTags.create("kafka", DataStreamsTags.Direction.INBOUND, "topic", "group", null); + context.setCheckpoint(fromTags(tags), pointConsumer); + timeSource.advance(30); + context.setCheckpoint(fromTags(tags), pointConsumer); + + assertTrue(context.isStarted()); + assertEquals(3, pointConsumer.points.size()); + verifyFirstPoint(pointConsumer.points.get(0)); + StatsPoint second = pointConsumer.points.get(1); + assertEquals(4, second.getTags().nonNullSize()); + assertEquals("direction:in", second.getTags().getDirection()); + assertEquals("group:group", second.getTags().getGroup()); + assertEquals("topic:topic", second.getTags().getTopic()); + assertEquals("type:kafka", second.getTags().getType()); + assertEquals(pointConsumer.points.get(0).getHash(), second.getParentHash()); + assertNotEquals(0L, second.getHash()); + assertEquals(25L, second.getPathwayLatencyNano()); + assertEquals(25L, second.getEdgeLatencyNano()); + StatsPoint third = pointConsumer.points.get(2); + assertEquals(4, third.getTags().nonNullSize()); + assertEquals("direction:in", third.getTags().getDirection()); + assertEquals("group:group", third.getTags().getGroup()); + assertEquals("topic:topic", third.getTags().getTopic()); + assertEquals("type:kafka", third.getTags().getType()); + // this point should have the first point as parent, + // as the loop protection will reset the parent if two identical + // points (same hash for tag values) are about to form a hierarchy + assertEquals(pointConsumer.points.get(0).getHash(), third.getParentHash()); + assertNotEquals(0L, third.getHash()); + assertEquals(55L, third.getPathwayLatencyNano()); + assertEquals(30L, third.getEdgeLatencyNano()); + } + + @Test + void exceptionThrownWhenTryingToEncodeUnstartedContext() { + ControllableTimeSource timeSource = new ControllableTimeSource(); + DefaultPathwayContext context = new DefaultPathwayContext(timeSource, null); + + assertThrows(IllegalStateException.class, context::encode); + } + + @Test + void setCheckpointWithDatasetTags() throws IOException { + ControllableTimeSource timeSource = new ControllableTimeSource(); + DefaultPathwayContext context = new DefaultPathwayContext(timeSource, null); + + timeSource.advance(MILLISECONDS.toNanos(50)); + context.setCheckpoint( + fromTags( + DataStreamsTags.createWithDataset( + "s3", DataStreamsTags.Direction.INBOUND, null, "my_object.csv", "my_bucket")), + pointConsumer); + String encoded = context.encode(); + timeSource.advance(MILLISECONDS.toNanos(2)); + DefaultPathwayContext decodedContext = DefaultPathwayContext.decode(timeSource, null, encoded); + timeSource.advance(MILLISECONDS.toNanos(25)); + DataStreamsTags tags = + DataStreamsTags.createWithDataset( + "s3", DataStreamsTags.Direction.OUTBOUND, null, "my_object.csv", "my_bucket"); + context.setCheckpoint(fromTags(tags), pointConsumer); + + assertTrue(decodedContext.isStarted()); + assertEquals(2, pointConsumer.points.size()); + + // all points should have datasetHash, which is not equal to hash or 0 + for (StatsPoint point : pointConsumer.points) { + assertNotEquals(point.getHash(), point.getAggregationHash()); + assertNotEquals(0L, point.getAggregationHash()); + } + } + + @Test + void encodingAndDecodingBase64AContext() throws IOException { + // Timesource needs to be advanced in milliseconds because encoding truncates to millis + ControllableTimeSource timeSource = new ControllableTimeSource(); + DefaultPathwayContext context = new DefaultPathwayContext(timeSource, null); + + timeSource.advance(MILLISECONDS.toNanos(50)); + context.setCheckpoint( + fromTags(DataStreamsTags.create("internal", DataStreamsTags.Direction.INBOUND)), + pointConsumer); + String encoded = context.encode(); + timeSource.advance(MILLISECONDS.toNanos(2)); + DefaultPathwayContext decodedContext = DefaultPathwayContext.decode(timeSource, null, encoded); + timeSource.advance(MILLISECONDS.toNanos(25)); + context.setCheckpoint( + fromTags(DataStreamsTags.create("kafka", null, "topic", "group", null)), pointConsumer); + + assertTrue(decodedContext.isStarted()); + assertEquals(2, pointConsumer.points.size()); + + StatsPoint second = pointConsumer.points.get(1); + assertEquals(3, second.getTags().nonNullSize()); + assertEquals("group:group", second.getTags().getGroup()); + assertEquals("type:kafka", second.getTags().getType()); + assertEquals("topic:topic", second.getTags().getTopic()); + assertEquals(pointConsumer.points.get(0).getHash(), second.getParentHash()); + assertNotEquals(0L, second.getHash()); + assertEquals(MILLISECONDS.toNanos(27), second.getPathwayLatencyNano()); + assertEquals(MILLISECONDS.toNanos(27), second.getEdgeLatencyNano()); + } + + @Test + void setCheckpointWithTimestamp() { + ControllableTimeSource timeSource = new ControllableTimeSource(); + DefaultPathwayContext context = new DefaultPathwayContext(timeSource, null); + long timeFromQueue = timeSource.getCurrentTimeMillis() - 200; + + context.setCheckpoint( + create(DataStreamsTags.create("internal", null), timeFromQueue, 0), pointConsumer); + + assertTrue(context.isStarted()); + assertEquals(1, pointConsumer.points.size()); + StatsPoint first = pointConsumer.points.get(0); + assertEquals("type:internal", first.getTags().getType()); + assertEquals(1, first.getTags().nonNullSize()); + assertEquals(0L, first.getParentHash()); + assertNotEquals(0L, first.getHash()); + assertEquals(MILLISECONDS.toNanos(200), first.getPathwayLatencyNano()); + assertEquals(MILLISECONDS.toNanos(200), first.getEdgeLatencyNano()); + } + + @Test + void encodingAndDecodingBase64WithContextsAndCheckpoints() throws IOException { + // Timesource needs to be advanced in milliseconds because encoding truncates to millis + ControllableTimeSource timeSource = new ControllableTimeSource(); + DefaultPathwayContext context = new DefaultPathwayContext(timeSource, null); + + timeSource.advance(MILLISECONDS.toNanos(50)); + context.setCheckpoint( + fromTags(DataStreamsTags.create("internal", DataStreamsTags.Direction.INBOUND)), + pointConsumer); + + String encoded = context.encode(); + timeSource.advance(MILLISECONDS.toNanos(1)); + DefaultPathwayContext decodedContext = DefaultPathwayContext.decode(timeSource, null, encoded); + timeSource.advance(MILLISECONDS.toNanos(25)); + context.setCheckpoint( + fromTags( + DataStreamsTags.create( + "kafka", DataStreamsTags.Direction.OUTBOUND, "topic", "group", null)), + pointConsumer); + + assertTrue(decodedContext.isStarted()); + assertEquals(2, pointConsumer.points.size()); + StatsPoint second = pointConsumer.points.get(1); + assertEquals("group:group", second.getTags().getGroup()); + assertEquals("topic:topic", second.getTags().getTopic()); + assertEquals("type:kafka", second.getTags().getType()); + assertEquals("direction:out", second.getTags().getDirection()); + assertEquals(4, second.getTags().nonNullSize()); + assertEquals(pointConsumer.points.get(0).getHash(), second.getParentHash()); + assertNotEquals(0L, second.getHash()); + assertEquals(MILLISECONDS.toNanos(26), second.getPathwayLatencyNano()); + assertEquals(MILLISECONDS.toNanos(26), second.getEdgeLatencyNano()); + + String secondEncode = decodedContext.encode(); + timeSource.advance(MILLISECONDS.toNanos(2)); + DefaultPathwayContext secondDecode = + DefaultPathwayContext.decode(timeSource, null, secondEncode); + timeSource.advance(MILLISECONDS.toNanos(30)); + context.setCheckpoint( + fromTags( + DataStreamsTags.create( + "kafka", DataStreamsTags.Direction.INBOUND, "topicB", "group", null)), + pointConsumer); + + assertTrue(secondDecode.isStarted()); + assertEquals(3, pointConsumer.points.size()); + StatsPoint third = pointConsumer.points.get(2); + assertEquals("group:group", third.getTags().getGroup()); + assertEquals("topic:topicB", third.getTags().getTopic()); + assertEquals("type:kafka", third.getTags().getType()); + assertEquals("direction:in", third.getTags().getDirection()); + assertEquals(4, third.getTags().nonNullSize()); + assertEquals(pointConsumer.points.get(1).getHash(), third.getParentHash()); + assertNotEquals(0L, third.getHash()); + assertEquals(MILLISECONDS.toNanos(58), third.getPathwayLatencyNano()); + assertEquals(MILLISECONDS.toNanos(32), third.getEdgeLatencyNano()); + } + + @Test + void encodingAndDecodingBase64WithInjectsAndExtracts() throws IOException { + // Timesource needs to be advanced in milliseconds because encoding truncates to millis + ControllableTimeSource timeSource = new ControllableTimeSource(); + DefaultPathwayContext context = new DefaultPathwayContext(timeSource, null); + Base64MapContextVisitor contextVisitor = new Base64MapContextVisitor(); + + timeSource.advance(MILLISECONDS.toNanos(50)); + context.setCheckpoint( + fromTags(DataStreamsTags.create("internal", DataStreamsTags.Direction.INBOUND)), + pointConsumer); + + String encoded = context.encode(); + Map carrier = new HashMap<>(); + carrier.put(PROPAGATION_KEY_BASE64, encoded); + carrier.put("someotherkey", "someothervalue"); + timeSource.advance(MILLISECONDS.toNanos(1)); + DefaultPathwayContext decodedContext = + DefaultPathwayContext.extract(carrier, contextVisitor, timeSource, null); + timeSource.advance(MILLISECONDS.toNanos(25)); + context.setCheckpoint( + fromTags( + DataStreamsTags.create( + "kafka", DataStreamsTags.Direction.OUTBOUND, "topic", "group", null)), + pointConsumer); + + assertTrue(decodedContext.isStarted()); + assertEquals(2, pointConsumer.points.size()); + StatsPoint second = pointConsumer.points.get(1); + assertEquals(4, second.getTags().nonNullSize()); + assertEquals("group:group", second.getTags().getGroup()); + assertEquals("topic:topic", second.getTags().getTopic()); + assertEquals("type:kafka", second.getTags().getType()); + assertEquals("direction:out", second.getTags().getDirection()); + assertEquals(pointConsumer.points.get(0).getHash(), second.getParentHash()); + assertNotEquals(0L, second.getHash()); + assertEquals(MILLISECONDS.toNanos(26), second.getPathwayLatencyNano()); + assertEquals(MILLISECONDS.toNanos(26), second.getEdgeLatencyNano()); + + String secondEncode = decodedContext.encode(); + carrier = new HashMap<>(); + carrier.put(PROPAGATION_KEY_BASE64, secondEncode); + timeSource.advance(MILLISECONDS.toNanos(2)); + DefaultPathwayContext secondDecode = + DefaultPathwayContext.extract(carrier, contextVisitor, timeSource, null); + timeSource.advance(MILLISECONDS.toNanos(30)); + context.setCheckpoint( + fromTags( + DataStreamsTags.create( + "kafka", DataStreamsTags.Direction.INBOUND, "topicB", "group", null)), + pointConsumer); + + assertTrue(secondDecode.isStarted()); + assertEquals(3, pointConsumer.points.size()); + StatsPoint third = pointConsumer.points.get(2); + assertEquals(4, third.getTags().nonNullSize()); + assertEquals("group:group", third.getTags().getGroup()); + assertEquals("topic:topicB", third.getTags().getTopic()); + assertEquals("type:kafka", third.getTags().getType()); + assertEquals("direction:in", third.getTags().getDirection()); + assertEquals(pointConsumer.points.get(1).getHash(), third.getParentHash()); + assertNotEquals(0L, third.getHash()); + assertEquals(MILLISECONDS.toNanos(58), third.getPathwayLatencyNano()); + assertEquals(MILLISECONDS.toNanos(32), third.getEdgeLatencyNano()); + } + + @Test + void encodingAndDecodingSqsFormattedWithInjectsAndExtracts() throws IOException { + // Timesource needs to be advanced in milliseconds because encoding truncates to millis + ControllableTimeSource timeSource = new ControllableTimeSource(); + DefaultPathwayContext context = new DefaultPathwayContext(timeSource, null); + Base64MapContextVisitor contextVisitor = new Base64MapContextVisitor(); + + timeSource.advance(MILLISECONDS.toNanos(50)); + context.setCheckpoint( + fromTags(DataStreamsTags.create("internal", DataStreamsTags.Direction.INBOUND)), + pointConsumer); + + String encoded = context.encode(); + Map carrier = new HashMap<>(); + carrier.put(PROPAGATION_KEY_BASE64, encoded); + carrier.put("someotherkey", "someothervalue"); + timeSource.advance(MILLISECONDS.toNanos(1)); + DefaultPathwayContext decodedContext = + DefaultPathwayContext.extract(carrier, contextVisitor, timeSource, null); + timeSource.advance(MILLISECONDS.toNanos(25)); + context.setCheckpoint( + fromTags( + DataStreamsTags.create("sqs", DataStreamsTags.Direction.OUTBOUND, "topic", null, null)), + pointConsumer); + + assertTrue(decodedContext.isStarted()); + assertEquals(2, pointConsumer.points.size()); + StatsPoint second = pointConsumer.points.get(1); + assertEquals("direction:out", second.getTags().getDirection()); + assertEquals("topic:topic", second.getTags().getTopic()); + assertEquals("type:sqs", second.getTags().getType()); + assertEquals(3, second.getTags().nonNullSize()); + assertEquals(pointConsumer.points.get(0).getHash(), second.getParentHash()); + assertNotEquals(0L, second.getHash()); + assertEquals(MILLISECONDS.toNanos(26), second.getPathwayLatencyNano()); + assertEquals(MILLISECONDS.toNanos(26), second.getEdgeLatencyNano()); + + String secondEncode = decodedContext.encode(); + carrier = new HashMap<>(); + carrier.put(PROPAGATION_KEY_BASE64, secondEncode); + timeSource.advance(MILLISECONDS.toNanos(2)); + DefaultPathwayContext secondDecode = + DefaultPathwayContext.extract(carrier, contextVisitor, timeSource, null); + timeSource.advance(MILLISECONDS.toNanos(30)); + context.setCheckpoint( + fromTags( + DataStreamsTags.create("sqs", DataStreamsTags.Direction.INBOUND, "topicB", null, null)), + pointConsumer); + + assertTrue(secondDecode.isStarted()); + assertEquals(3, pointConsumer.points.size()); + StatsPoint third = pointConsumer.points.get(2); + assertEquals("type:sqs", third.getTags().getType()); + assertEquals("topic:topicB", third.getTags().getTopic()); + assertEquals(3, third.getTags().nonNullSize()); + assertEquals(pointConsumer.points.get(1).getHash(), third.getParentHash()); + assertNotEquals(0L, third.getHash()); + assertEquals(MILLISECONDS.toNanos(58), third.getPathwayLatencyNano()); + assertEquals(MILLISECONDS.toNanos(32), third.getEdgeLatencyNano()); + } + + @Test + void emptyTagsNotSet() { + ControllableTimeSource timeSource = new ControllableTimeSource(); + DefaultPathwayContext context = new DefaultPathwayContext(timeSource, null); + + timeSource.advance(50); + context.setCheckpoint( + fromTags(DataStreamsTags.create("internal", DataStreamsTags.Direction.INBOUND)), + pointConsumer); + timeSource.advance(25); + context.setCheckpoint( + fromTags( + DataStreamsTags.create( + "type", DataStreamsTags.Direction.OUTBOUND, "topic", "group", null)), + pointConsumer); + timeSource.advance(25); + context.setCheckpoint(fromTags(DataStreamsTags.create(null, null)), pointConsumer); + + assertTrue(context.isStarted()); + assertEquals(3, pointConsumer.points.size()); + verifyFirstPoint(pointConsumer.points.get(0)); + StatsPoint second = pointConsumer.points.get(1); + assertEquals("type:type", second.getTags().getType()); + assertEquals("topic:topic", second.getTags().getTopic()); + assertEquals("group:group", second.getTags().getGroup()); + assertEquals("direction:out", second.getTags().getDirection()); + assertEquals(4, second.getTags().nonNullSize()); + assertEquals(pointConsumer.points.get(0).getHash(), second.getParentHash()); + assertNotEquals(0L, second.getHash()); + assertEquals(25L, second.getPathwayLatencyNano()); + assertEquals(25L, second.getEdgeLatencyNano()); + StatsPoint third = pointConsumer.points.get(2); + assertEquals(0, third.getTags().nonNullSize()); + assertEquals(pointConsumer.points.get(1).getHash(), third.getParentHash()); + assertNotEquals(0L, third.getHash()); + assertEquals(50L, third.getPathwayLatencyNano()); + assertEquals(25L, third.getEdgeLatencyNano()); + } + + @TableTest({ + "scenario | dynamicConfigEnabled", + "enabled | true ", + "disabled | false " + }) + void checkContextExtractorDecoratorBehavior(boolean dynamicConfigEnabled) throws IOException { + Sink sink = mock(Sink.class); + DDAgentFeaturesDiscovery features = mock(DDAgentFeaturesDiscovery.class, RETURNS_SMART_NULLS); + when(features.supportsDataStreams()).thenReturn(true); + ControllableTimeSource timeSource = new ControllableTimeSource(); + DatastreamsPayloadWriter payloadWriter = mock(DatastreamsPayloadWriter.class); + + TraceConfig globalTraceConfig = mock(TraceConfig.class, RETURNS_SMART_NULLS); + when(globalTraceConfig.isDataStreamsEnabled()).thenReturn(dynamicConfigEnabled); + + AgentTracer.TracerAPI tracerApi = mock(AgentTracer.TracerAPI.class, RETURNS_SMART_NULLS); + when(tracerApi.captureTraceConfig()).thenReturn(globalTraceConfig); + AgentTracer.TracerAPI originalTracer = AgentTracer.get(); + AgentTracer.forceRegister(tracerApi); + + try { + DefaultDataStreamsMonitoring dataStreams = + new DefaultDataStreamsMonitoring( + sink, + features, + timeSource, + () -> globalTraceConfig, + payloadWriter, + DEFAULT_BUCKET_DURATION_NANOS); + + BaseHash.updateBaseHash(BASE_HASH); + DefaultPathwayContext context = new DefaultPathwayContext(timeSource, null); + timeSource.advance(MILLISECONDS.toNanos(50)); + context.setCheckpoint( + fromTags(DataStreamsTags.create("internal", DataStreamsTags.Direction.INBOUND)), + pointConsumer); + String encoded = context.encode(); + Map carrier = new HashMap<>(); + carrier.put(PROPAGATION_KEY_BASE64, encoded); + carrier.put("someotherkey", "someothervalue"); + Base64MapContextVisitor contextVisitor = new Base64MapContextVisitor(); + Propagator propagator = dataStreams.propagator(); + + Context extractedContext = propagator.extract(root(), carrier, contextVisitor); + AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + + assertEquals("L+lDG/Pa9hRkZA==", encoded); + if (dynamicConfigEnabled) { + assertNotNull(extractedSpan); + assertNotNull(extractedSpan.context()); + assertNotNull(extractedSpan.context().getPathwayContext()); + assertTrue(extractedSpan.context().getPathwayContext().isStarted()); + } + } finally { + // cleanup + AgentTracer.forceRegister(originalTracer); + } + } + + @TableTest({ + "scenario | globalDsmEnabled", + "enabled | true ", + "disabled | false " + }) + void checkContextExtractorDecoratorBehaviorWhenTraceDataIsNull(boolean globalDsmEnabled) + throws IOException { + Sink sink = mock(Sink.class); + DDAgentFeaturesDiscovery features = mock(DDAgentFeaturesDiscovery.class, RETURNS_SMART_NULLS); + when(features.supportsDataStreams()).thenReturn(true); + ControllableTimeSource timeSource = new ControllableTimeSource(); + DatastreamsPayloadWriter payloadWriter = mock(DatastreamsPayloadWriter.class); + + TraceConfig globalTraceConfig = mock(TraceConfig.class, RETURNS_SMART_NULLS); + when(globalTraceConfig.isDataStreamsEnabled()).thenReturn(globalDsmEnabled); + + AgentTracer.TracerAPI tracerApi = mock(AgentTracer.TracerAPI.class, RETURNS_SMART_NULLS); + when(tracerApi.captureTraceConfig()).thenReturn(globalTraceConfig); + AgentTracer.TracerAPI originalTracer = AgentTracer.get(); + AgentTracer.forceRegister(tracerApi); + + try { + DefaultDataStreamsMonitoring dataStreams = + new DefaultDataStreamsMonitoring( + sink, + features, + timeSource, + () -> globalTraceConfig, + payloadWriter, + DEFAULT_BUCKET_DURATION_NANOS); + + BaseHash.updateBaseHash(BASE_HASH); + DefaultPathwayContext context = new DefaultPathwayContext(timeSource, null); + timeSource.advance(MILLISECONDS.toNanos(50)); + context.setCheckpoint( + fromTags(DataStreamsTags.create("internal", DataStreamsTags.Direction.INBOUND)), + pointConsumer); + String encoded = context.encode(); + + Map carrier = new HashMap<>(); + carrier.put(PROPAGATION_KEY_BASE64, encoded); + carrier.put("someotherkey", "someothervalue"); + Base64MapContextVisitor contextVisitor = new Base64MapContextVisitor(); + Propagator propagator = dataStreams.propagator(); + + Context extractedContext = propagator.extract(root(), carrier, contextVisitor); + AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + + assertEquals("L+lDG/Pa9hRkZA==", encoded); + if (globalDsmEnabled) { + assertNotNull(extractedSpan); + assertNotNull(extractedSpan.context()); + assertNotNull(extractedSpan.context().getPathwayContext()); + assertTrue(extractedSpan.context().getPathwayContext().isStarted()); + } else { + assertNull(extractedSpan); + } + } finally { + // cleanup + AgentTracer.forceRegister(originalTracer); + } + } + + @ParameterizedTest + @ValueSource(booleans = {true, false}) + void checkContextExtractorDecoratorBehaviorWhenLocalTraceConfigIsNull(boolean globalDsmEnabled) + throws IOException { + Sink sink = mock(Sink.class); + DDAgentFeaturesDiscovery features = mock(DDAgentFeaturesDiscovery.class, RETURNS_SMART_NULLS); + when(features.supportsDataStreams()).thenReturn(true); + ControllableTimeSource timeSource = new ControllableTimeSource(); + DatastreamsPayloadWriter payloadWriter = mock(DatastreamsPayloadWriter.class); + + TraceConfig globalTraceConfig = mock(TraceConfig.class, RETURNS_SMART_NULLS); + when(globalTraceConfig.isDataStreamsEnabled()).thenReturn(globalDsmEnabled); + + AgentTracer.TracerAPI tracerApi = mock(AgentTracer.TracerAPI.class, RETURNS_SMART_NULLS); + when(tracerApi.captureTraceConfig()).thenReturn(globalTraceConfig); + AgentTracer.TracerAPI originalTracer = AgentTracer.get(); + AgentTracer.forceRegister(tracerApi); + + try { + DefaultDataStreamsMonitoring dataStreams = + new DefaultDataStreamsMonitoring( + sink, + features, + timeSource, + () -> globalTraceConfig, + payloadWriter, + DEFAULT_BUCKET_DURATION_NANOS); + + BaseHash.updateBaseHash(BASE_HASH); + DefaultPathwayContext context = new DefaultPathwayContext(timeSource, null); + timeSource.advance(MILLISECONDS.toNanos(50)); + context.setCheckpoint( + fromTags(DataStreamsTags.create("internal", DataStreamsTags.Direction.INBOUND)), + pointConsumer); + String encoded = context.encode(); + Map carrier = new HashMap<>(); + carrier.put(PROPAGATION_KEY_BASE64, encoded); + carrier.put("someotherkey", "someothervalue"); + Base64MapContextVisitor contextVisitor = new Base64MapContextVisitor(); + ExtractedContext spanContext = + new ExtractedContext( + DDTraceId.ONE, + 1, + 0, + null, + 0, + null, + (TagMap) null, + null, + null, + globalTraceConfig, + DATADOG); + Context baseContext = AgentSpan.fromSpanContext(spanContext).storeInto(root()); + Propagator propagator = dataStreams.propagator(); + + Context extractedContext = propagator.extract(baseContext, carrier, contextVisitor); + AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + + assertNotNull(extractedSpan); + + Object extracted = extractedSpan.context(); + + assertNotNull(extracted); + assertEquals("L+lDG/Pa9hRkZA==", encoded); + if (globalDsmEnabled) { + assertNotNull(extractedSpan.context().getPathwayContext()); + assertTrue(extractedSpan.context().getPathwayContext().isStarted()); + } else { + assertNull(extractedSpan.context().getPathwayContext()); + } + } finally { + // cleanup + AgentTracer.forceRegister(originalTracer); + } + } + + @Test + void checkContextExtractorDecoratorBehaviorWhenTraceDataAndDsmDataAreNull() { + Sink sink = mock(Sink.class); + DDAgentFeaturesDiscovery features = mock(DDAgentFeaturesDiscovery.class, RETURNS_SMART_NULLS); + when(features.supportsDataStreams()).thenReturn(true); + ControllableTimeSource timeSource = new ControllableTimeSource(); + DatastreamsPayloadWriter payloadWriter = mock(DatastreamsPayloadWriter.class); + + TraceConfig traceConfig = mock(TraceConfig.class, RETURNS_SMART_NULLS); + when(traceConfig.isDataStreamsEnabled()).thenReturn(true); + + DefaultDataStreamsMonitoring dataStreams = + new DefaultDataStreamsMonitoring( + sink, + features, + timeSource, + () -> traceConfig, + payloadWriter, + DEFAULT_BUCKET_DURATION_NANOS); + + Map carrier = new HashMap<>(); + carrier.put("someotherkey", "someothervalue"); + Base64MapContextVisitor contextVisitor = new Base64MapContextVisitor(); + Propagator propagator = dataStreams.propagator(); + + Context extractedContext = propagator.extract(root(), carrier, contextVisitor); + AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + + assertNull(extractedSpan); + } + + static class CapturingPointConsumer implements Consumer { + final List points = new ArrayList<>(); + + @Override + public void accept(StatsPoint point) { + points.add(point); + } + } + + static class Base64MapContextVisitor + implements AgentPropagation.ContextVisitor> { + @Override + public void forEachKey(Map carrier, AgentPropagation.KeyClassifier classifier) { + for (Map.Entry entry : carrier.entrySet()) { + classifier.accept(entry.getKey(), entry.getValue()); + } + } + } +} diff --git a/dd-trace-core/src/test/java/datadog/trace/core/datastreams/SchemaBuilderTest.java b/dd-trace-core/src/test/java/datadog/trace/core/datastreams/SchemaBuilderTest.java new file mode 100644 index 00000000000..51f000363d2 --- /dev/null +++ b/dd-trace-core/src/test/java/datadog/trace/core/datastreams/SchemaBuilderTest.java @@ -0,0 +1,60 @@ +package datadog.trace.core.datastreams; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import datadog.trace.bootstrap.instrumentation.api.Schema; +import datadog.trace.bootstrap.instrumentation.api.SchemaIterator; +import datadog.trace.core.DDCoreJavaSpecification; +import java.util.HashMap; +import org.junit.jupiter.api.Test; + +public class SchemaBuilderTest extends DDCoreJavaSpecification { + + static class Iterator implements SchemaIterator { + @Override + public void iterateOverSchema( + datadog.trace.bootstrap.instrumentation.api.SchemaBuilder builder) { + HashMap extension = new HashMap<>(1); + extension.put("x-test-extension-1", "hello"); + extension.put("x-test-extension-2", "world"); + builder.addProperty( + "person", "name", false, "string", "name of the person", null, null, null, null); + builder.addProperty("person", "phone_numbers", true, "string", null, null, null, null, null); + builder.addProperty("person", "person_name", false, "string", null, null, null, null, null); + builder.addProperty( + "person", + "address", + false, + "object", + null, + "#/components/schemas/address", + null, + null, + null); + builder.addProperty("address", "zip", false, "number", null, null, "int", null, null); + builder.addProperty("address", "street", false, "string", null, null, null, null, extension); + } + } + + @Test + void schemaIsConvertedCorrectlyToJson() { + SchemaBuilder builder = new SchemaBuilder(new Iterator()); + + boolean shouldExtractPerson = builder.shouldExtractSchema("person", 0); + boolean shouldExtractAddress = builder.shouldExtractSchema("address", 1); + boolean shouldExtractPerson2 = builder.shouldExtractSchema("person", 0); + boolean shouldExtractTooDeep = builder.shouldExtractSchema("city", 11); + Schema schema = builder.build(); + + String expectedDefinition = + "{\"components\":{\"schemas\":{\"person\":{\"properties\":{\"name\":{\"description\":\"name of the person\",\"type\":\"string\"},\"phone_numbers\":{\"items\":{\"type\":\"string\"},\"type\":\"array\"},\"person_name\":{\"type\":\"string\"},\"address\":{\"$ref\":\"#/components/schemas/address\",\"type\":\"object\"}},\"type\":\"object\"},\"address\":{\"properties\":{\"zip\":{\"format\":\"int\",\"type\":\"number\"},\"street\":{\"extensions\":{\"x-test-extension-1\":\"hello\",\"x-test-extension-2\":\"world\"},\"type\":\"string\"}},\"type\":\"object\"}}},\"openapi\":\"3.0.0\"}"; + assertEquals(expectedDefinition, schema.definition); + assertEquals("16548065305426330543", schema.id); + assertTrue(shouldExtractPerson); + assertTrue(shouldExtractAddress); + assertFalse(shouldExtractPerson2); + assertFalse(shouldExtractTooDeep); + } +} diff --git a/dd-trace-core/src/test/java/datadog/trace/core/datastreams/SchemaSamplerTest.java b/dd-trace-core/src/test/java/datadog/trace/core/datastreams/SchemaSamplerTest.java new file mode 100644 index 00000000000..42479eb85cc --- /dev/null +++ b/dd-trace-core/src/test/java/datadog/trace/core/datastreams/SchemaSamplerTest.java @@ -0,0 +1,33 @@ +package datadog.trace.core.datastreams; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import datadog.trace.core.DDCoreJavaSpecification; +import org.junit.jupiter.api.Test; + +public class SchemaSamplerTest extends DDCoreJavaSpecification { + + @Test + void schemaSamplerSamplesWithCorrectWeights() { + long currentTimeMillis = 100000; + SchemaSampler sampler = new SchemaSampler(); + + boolean canSample1 = sampler.canSample(currentTimeMillis); + int weight1 = sampler.trySample(currentTimeMillis); + boolean canSample2 = sampler.canSample(currentTimeMillis + 1000); + boolean canSample3 = sampler.canSample(currentTimeMillis + 2000); + boolean canSample4 = sampler.canSample(currentTimeMillis + 30000); + int weight4 = sampler.trySample(currentTimeMillis + 30000); + boolean canSample5 = sampler.canSample(currentTimeMillis + 30001); + + assertTrue(canSample1); + assertEquals(1, weight1); + assertFalse(canSample2); + assertFalse(canSample3); + assertTrue(canSample4); + assertEquals(3, weight4); + assertFalse(canSample5); + } +} diff --git a/dd-trace-core/src/test/java/datadog/trace/core/datastreams/TransactionContainerTest.java b/dd-trace-core/src/test/java/datadog/trace/core/datastreams/TransactionContainerTest.java new file mode 100644 index 00000000000..5f37096515e --- /dev/null +++ b/dd-trace-core/src/test/java/datadog/trace/core/datastreams/TransactionContainerTest.java @@ -0,0 +1,50 @@ +package datadog.trace.core.datastreams; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; + +import datadog.trace.api.datastreams.TransactionInfo; +import datadog.trace.api.datastreams.TransactionInfoTestBridge; +import datadog.trace.core.DDCoreJavaSpecification; +import org.junit.jupiter.api.Test; + +public class TransactionContainerTest extends DDCoreJavaSpecification { + + private static final byte[] EXPECTED_CONTAINER_DATA = + new byte[] {1, 0, 0, 0, 0, 0, 0, 0, 1, 1, 49, 2, 0, 0, 0, 0, 0, 0, 0, 2, 1, 50}; + + @Test + void testWithNoResize() { + TransactionInfoTestBridge.resetCache(); + TransactionContainer container = new TransactionContainer(1024); + container.add(new TransactionInfo("1", 1, "1")); + container.add(new TransactionInfo("2", 2, "2")); + byte[] data = container.getData(); + + assertEquals(22, data.length); + assertArrayEquals(EXPECTED_CONTAINER_DATA, data); + } + + @Test + void testWithResize() { + TransactionInfoTestBridge.resetCache(); + TransactionContainer container = new TransactionContainer(10); + container.add(new TransactionInfo("1", 1, "1")); + container.add(new TransactionInfo("2", 2, "2")); + byte[] data = container.getData(); + + assertEquals(22, data.length); + assertArrayEquals(EXPECTED_CONTAINER_DATA, data); + } + + @Test + void testCheckpointMap() { + TransactionInfoTestBridge.resetCache(); + new TransactionInfo("1", 1, "1"); + new TransactionInfo("2", 2, "2"); + byte[] data = TransactionInfo.getCheckpointIdCacheBytes(); + + assertEquals(6, data.length); + assertArrayEquals(new byte[] {1, 1, 49, 2, 1, 50}, data); + } +}