TraceStorageTopology implements Supplier<Topology> { @Override public Topology get() { StreamsBuilder builder = new StreamsBuilder(); if (traceSearchEnabled || traceByIdQueryEnabled) { builder .addStateStore(Stores.keyValueStoreBuilder( Stores.persistentKeyValueStore(TRACES_STORE_NAME), Serdes.String(), spansSerde).withLoggingDisabled()) .addStateStore(Stores.keyValueStoreBuilder( Stores.persistentKeyValueStore(SPAN_IDS_BY_TS_STORE_NAME), Serdes.Long(), spanIdsSerde).withLoggingDisabled()); KStream<String, List<Span>> spansStream = builder .stream(spansTopic, Consumed.with(Serdes.String(), spansSerde)); spansStream.process(() -> new Processor<String, List<Span>>() { KeyValueStore<String, List<Span>> tracesStore; KeyValueStore<Long, Set<String>> spanIdsByTsStore; @SuppressWarnings("unchecked") @Override public void init(ProcessorContext context) { tracesStore = (KeyValueStore<String, List<Span>>) context.getStateStore(TRACES_STORE_NAME); spanIdsByTsStore = (KeyValueStore<Long, Set<String>>) context.getStateStore(SPAN_IDS_BY_TS_STORE_NAME); context.schedule( traceTtlCheckInterval, PunctuationType.STREAM_TIME, timestamp -> { if (traceTtl.toMillis() > 0 && tracesStore.approximateNumEntries() > minTracesStored) { long from = 0L; long to = timestamp - traceTtl.toMillis(); try (final KeyValueIterator<Long, Set<String>> range = spanIdsByTsStore.range(from, MILLISECONDS.toMicros(to))) { range.forEachRemaining(record -> { spanIdsByTsStore.delete(record.key); for (String traceId : record.value) { tracesStore.delete(traceId); } }); LOG.debug( "Traces deletion emitted at {}, approx. number of traces stored {}", Instant.ofEpochMilli(to).atZone(ZoneId.systemDefault()), tracesStore.approximateNumEntries()); } } }); } @Override public void process(String traceId, List<Span> spans) { if (!spans.isEmpty()) { List<Span> currentSpans = tracesStore.get(traceId); if (currentSpans == null) { currentSpans = new ArrayList<>(); } else { brokenTracesTotal.increment(); } currentSpans.addAll(spans); tracesStore.put(traceId, currentSpans); long timestamp = spans.get(0).timestamp(); Set<String> currentSpanIds = spanIdsByTsStore.get(timestamp); if (currentSpanIds == null) currentSpanIds = new LinkedHashSet<>(); currentSpanIds.add(traceId); spanIdsByTsStore.put(timestamp, currentSpanIds); } } @Override public void close() { } }, TRACES_STORE_NAME, SPAN_IDS_BY_TS_STORE_NAME); if (traceSearchEnabled) { builder .addStateStore(Stores.keyValueStoreBuilder( Stores.inMemoryKeyValueStore(SERVICE_NAMES_STORE_NAME), Serdes.String(), Serdes.String())) .addStateStore(Stores.keyValueStoreBuilder( Stores.inMemoryKeyValueStore(SPAN_NAMES_STORE_NAME), Serdes.String(), namesSerde)) .addStateStore(Stores.keyValueStoreBuilder( Stores.inMemoryKeyValueStore(REMOTE_SERVICE_NAMES_STORE_NAME), Serdes.String(), namesSerde)) .addStateStore(Stores.keyValueStoreBuilder( Stores.persistentKeyValueStore(AUTOCOMPLETE_TAGS_STORE_NAME), Serdes.String(), namesSerde)); spansStream.process(() -> new Processor<String, List<Span>>() { KeyValueStore<String, String> serviceNameStore; KeyValueStore<String, Set<String>> spanNamesStore; KeyValueStore<String, Set<String>> remoteServiceNamesStore; KeyValueStore<String, Set<String>> autocompleteTagsStore; @SuppressWarnings("unchecked") @Override public void init(ProcessorContext context) { serviceNameStore = (KeyValueStore<String, String>) context.getStateStore(SERVICE_NAMES_STORE_NAME); spanNamesStore = (KeyValueStore<String, Set<String>>) context.getStateStore(SPAN_NAMES_STORE_NAME); remoteServiceNamesStore = (KeyValueStore<String, Set<String>>) context.getStateStore( REMOTE_SERVICE_NAMES_STORE_NAME); autocompleteTagsStore = (KeyValueStore<String, Set<String>>) context.getStateStore( AUTOCOMPLETE_TAGS_STORE_NAME); } @Override public void process(String traceId, List<Span> spans) { for (Span span : spans) { if (span.localServiceName() != null) { serviceNameStore.putIfAbsent(span.localServiceName(), span.localServiceName()); if (span.name() != null) { Set<String> spanNames = spanNamesStore.get(span.localServiceName()); if (spanNames == null) spanNames = new LinkedHashSet<>(); if (!spanNames.contains(span.name())) { spanNames.add(span.name()); spanNamesStore.put(span.localServiceName(), spanNames); } } if (span.remoteServiceName() != null) { Set<String> remoteServiceNames = remoteServiceNamesStore.get(span.localServiceName()); if (remoteServiceNames == null) remoteServiceNames = new LinkedHashSet<>(); if (!remoteServiceNames.contains(span.remoteServiceName())) { remoteServiceNames.add(span.remoteServiceName()); remoteServiceNamesStore.put(span.localServiceName(), remoteServiceNames); } } } if (!span.tags().isEmpty()) { autoCompleteKeys.forEach(tagKey -> { String value = span.tags().get(tagKey); if (value != null) { Set<String> values = autocompleteTagsStore.get(tagKey); if (values == null) values = new LinkedHashSet<>(); if (!values.contains(value)) { values.add(value); autocompleteTagsStore.put(tagKey, values); } } }); } } } @Override public void close() { } }, SERVICE_NAMES_STORE_NAME, SPAN_NAMES_STORE_NAME, REMOTE_SERVICE_NAMES_STORE_NAME, AUTOCOMPLETE_TAGS_STORE_NAME); } } return builder.build(); } TraceStorageTopology(
      String spansTopic,
      List<String> autoCompleteKeys,
      Duration traceTtl,
      Duration traceTtlCheckInterval,
      long minTracesStored,
      boolean traceByIdQueryEnabled,
      boolean traceSearchEnabled
  ); @Override Topology get(); static final String TRACES_STORE_NAME; static final String SPAN_IDS_BY_TS_STORE_NAME; static final String SERVICE_NAMES_STORE_NAME; static final String SPAN_NAMES_STORE_NAME; static final String REMOTE_SERVICE_NAMES_STORE_NAME; static final String AUTOCOMPLETE_TAGS_STORE_NAME; }
@Test void should_doNothing_whenAllDisabled() { Duration traceTtl = Duration.ofMillis(5); Duration traceTtlCheckInterval = Duration.ofMinutes(1); List<String> autocompleteKeys = Collections.singletonList("environment"); Topology topology = new TraceStorageTopology( spansTopic, autocompleteKeys, traceTtl, traceTtlCheckInterval, 0, false, false).get(); TopologyDescription description = topology.describe(); assertThat(description.subtopologies()).hasSize(0); TopologyTestDriver testDriver = new TopologyTestDriver(topology, props); testDriver.close(); }
@Test void should_persistSpans_and_onlyQueryTraces_whenEnabled() { Duration traceTtl = Duration.ofMillis(5); Duration traceTtlCheckInterval = Duration.ofMinutes(1); List<String> autocompleteKeys = Collections.singletonList("environment"); SpansSerde spansSerde = new SpansSerde(); Topology topology = new TraceStorageTopology( spansTopic, autocompleteKeys, traceTtl, traceTtlCheckInterval, 0, true, false).get(); TopologyDescription description = topology.describe(); assertThat(description.subtopologies()).hasSize(1); TopologyTestDriver testDriver = new TopologyTestDriver(topology, props); ConsumerRecordFactory<String, List<Span>> factory = new ConsumerRecordFactory<>(spansTopic, new StringSerializer(), spansSerde.serializer()); Span a = Span.newBuilder().traceId("a").id("a").name("op_a").kind(Span.Kind.CLIENT) .localEndpoint(Endpoint.newBuilder().serviceName("svc_a").build()) .timestamp(10000L).duration(11L) .putTag("environment", "dev") .build(); Span b = Span.newBuilder().traceId("a").id("b").name("op_b").kind(Span.Kind.SERVER) .localEndpoint(Endpoint.newBuilder().serviceName("svc_b").build()) .timestamp(10000L).duration(10L) .build(); Span c = Span.newBuilder().traceId("c").id("c").name("op_a").kind(Span.Kind.CLIENT) .localEndpoint(Endpoint.newBuilder().serviceName("svc_a").build()) .timestamp(10000L).duration(11L) .putTag("environment", "dev") .build(); List<Span> spans = Arrays.asList(a, b, c); testDriver.pipeInput(factory.create(spansTopic, a.traceId(), spans, 10L)); KeyValueStore<String, List<Span>> traces = testDriver.getKeyValueStore(TRACES_STORE_NAME); assertThat(traces.get(a.traceId())).containsExactlyElementsOf(spans); KeyValueStore<Long, Set<String>> spanIdsByTs = testDriver.getKeyValueStore(SPAN_IDS_BY_TS_STORE_NAME); KeyValueIterator<Long, Set<String>> ids = spanIdsByTs.all(); assertThat(ids).hasNext(); assertThat(ids.next().value).containsExactly(a.traceId()); KeyValueStore<String, String> serviceNames = testDriver.getKeyValueStore(SERVICE_NAMES_STORE_NAME); assertThat(serviceNames).isNull(); KeyValueStore<String, Set<String>> spanNames = testDriver.getKeyValueStore(SPAN_NAMES_STORE_NAME); assertThat(spanNames).isNull(); KeyValueStore<String, Set<String>> autocompleteTags = testDriver.getKeyValueStore(AUTOCOMPLETE_TAGS_STORE_NAME); assertThat(autocompleteTags).isNull(); testDriver.close(); spansSerde.close(); }
@Test void should_persistSpans_and_searchQueryTraces_whenAllEnabled() { Duration traceTtl = Duration.ofMillis(5); Duration traceTtlCheckInterval = Duration.ofMinutes(1); List<String> autocompleteKeys = Collections.singletonList("environment"); SpansSerde spansSerde = new SpansSerde(); Topology topology = new TraceStorageTopology( spansTopic, autocompleteKeys, traceTtl, traceTtlCheckInterval, 0, true, true).get(); TopologyDescription description = topology.describe(); assertThat(description.subtopologies()).hasSize(1); TopologyTestDriver testDriver = new TopologyTestDriver(topology, props); ConsumerRecordFactory<String, List<Span>> factory = new ConsumerRecordFactory<>(spansTopic, new StringSerializer(), spansSerde.serializer()); Span a = Span.newBuilder().traceId("a").id("a").name("op_a").kind(Span.Kind.CLIENT) .localEndpoint(Endpoint.newBuilder().serviceName("svc_a").build()) .timestamp(10000L).duration(11L) .putTag("environment", "dev") .build(); Span b = Span.newBuilder().traceId("a").id("b").name("op_b").kind(Span.Kind.SERVER) .localEndpoint(Endpoint.newBuilder().serviceName("svc_b").build()) .timestamp(10000L).duration(10L) .build(); Span c = Span.newBuilder().traceId("c").id("c").name("op_a").kind(Span.Kind.CLIENT) .localEndpoint(Endpoint.newBuilder().serviceName("svc_a").build()) .timestamp(10000L).duration(11L) .putTag("environment", "dev") .build(); List<Span> spans = Arrays.asList(a, b, c); testDriver.pipeInput(factory.create(spansTopic, a.traceId(), spans, 10L)); KeyValueStore<String, List<Span>> traces = testDriver.getKeyValueStore(TRACES_STORE_NAME); assertThat(traces.get(a.traceId())).containsExactlyElementsOf(spans); KeyValueStore<Long, Set<String>> spanIdsByTs = testDriver.getKeyValueStore(SPAN_IDS_BY_TS_STORE_NAME); KeyValueIterator<Long, Set<String>> ids = spanIdsByTs.all(); assertThat(ids).hasNext(); assertThat(ids.next().value).containsExactly(a.traceId()); KeyValueStore<String, String> serviceNames = testDriver.getKeyValueStore(SERVICE_NAMES_STORE_NAME); List<String> serviceNameList = new ArrayList<>(); serviceNames.all().forEachRemaining(serviceName -> serviceNameList.add(serviceName.value)); assertThat(serviceNameList).hasSize(2); assertThat(serviceNames.get("svc_a")).isEqualTo("svc_a"); assertThat(serviceNames.get("svc_b")).isEqualTo("svc_b"); KeyValueStore<String, Set<String>> spanNames = testDriver.getKeyValueStore(SPAN_NAMES_STORE_NAME); assertThat(spanNames.get("svc_a")).containsExactly("op_a"); assertThat(spanNames.get("svc_b")).containsExactly("op_b"); KeyValueStore<String, Set<String>> autocompleteTags = testDriver.getKeyValueStore(AUTOCOMPLETE_TAGS_STORE_NAME); assertThat(autocompleteTags.get("environment")).containsExactly("dev"); Span d = Span.newBuilder() .traceId("d") .id("d") .timestamp( MILLISECONDS.toMicros(traceTtlCheckInterval.toMillis()) + MILLISECONDS.toMicros(20)) .build(); testDriver.pipeInput( factory.create(spansTopic, d.traceId(), Collections.singletonList(d), traceTtlCheckInterval.plusMillis(1).toMillis())); assertThat(traces.get(a.traceId())).isNull(); testDriver.close(); spansSerde.close(); }