Examples with Record com.amazon.dataprepper.model.record.Record used on opensource projects

Example 41 with Record

use of com.amazon.dataprepper.model.record.Record in project data-prepper by opensearch-project.

the class OTelTraceRawPrepperTest method testExportRequestFlushByParentSpan.

@Test
public void testExportRequestFlushByParentSpan() throws IOException {
    final ExportTraceServiceRequest exportTraceServiceRequest = buildExportTraceServiceRequestFromJsonFile(TEST_REQUEST_TWO_FULL_TRACE_GROUP_JSON_FILE);
    final List<Record<String>> processedRecords = (List<Record<String>>) oTelTraceRawPrepper.doExecute(Collections.singletonList(new Record<>(exportTraceServiceRequest)));
    Assertions.assertThat(processedRecords.size()).isEqualTo(6);
    Assertions.assertThat(getMissingTraceGroupFieldsSpanCount(processedRecords)).isEqualTo(0);
}

Example 42 with Record

use of com.amazon.dataprepper.model.record.Record in project data-prepper by opensearch-project.

the class OTelTraceRawPrepperTest method testExportRequestFlushByMixedMultiThread.

@Test
public void testExportRequestFlushByMixedMultiThread() throws IOException, InterruptedException, ExecutionException {
    final ExportTraceServiceRequest exportTraceServiceRequest1 = buildExportTraceServiceRequestFromJsonFile(TEST_REQUEST_ONE_FULL_TRACE_GROUP_JSON_FILE);
    final ExportTraceServiceRequest exportTraceServiceRequest2 = buildExportTraceServiceRequestFromJsonFile(TEST_REQUEST_ONE_TRACE_GROUP_MISSING_ROOT_JSON_FILE);
    List<Record<String>> processedRecords = new ArrayList<>();
    List<Future<Collection<Record<String>>>> futures = new ArrayList<>();
    futures.addAll(submitExportTraceServiceRequests(Collections.singletonList(exportTraceServiceRequest1)));
    futures.addAll(submitExportTraceServiceRequests(Collections.singletonList(exportTraceServiceRequest2)));
    for (Future<Collection<Record<String>>> future : futures) {
        processedRecords.addAll(future.get());
    }
    await().atMost(2 * TEST_TRACE_FLUSH_INTERVAL, TimeUnit.SECONDS).untilAsserted(() -> {
        List<Future<Collection<Record<String>>>> futureList = submitExportTraceServiceRequests(Collections.emptyList());
        for (Future<Collection<Record<String>>> future : futureList) {
            processedRecords.addAll(future.get());
        }
        Assertions.assertThat(processedRecords.size()).isEqualTo(5);
        Assertions.assertThat(getMissingTraceGroupFieldsSpanCount(processedRecords)).isEqualTo(2);
    });
}

Example 43 with Record

use of com.amazon.dataprepper.model.record.Record in project data-prepper by opensearch-project.

the class OTelTraceRawPrepperTest method testPrepareForShutdown.

@Test
public void testPrepareForShutdown() throws Exception {
    // Assert no records in memory
    assertTrue(oTelTraceRawPrepper.isReadyForShutdown());
    // Add records to memory/queue
    final ExportTraceServiceRequest exportTraceServiceRequest = buildExportTraceServiceRequestFromJsonFile(TEST_REQUEST_TWO_TRACE_GROUP_MISSING_ROOTS_JSON_FILE);
    oTelTraceRawPrepper.doExecute(Collections.singletonList(new Record<>(exportTraceServiceRequest)));
    // Assert records exist in memory
    assertFalse(oTelTraceRawPrepper.isReadyForShutdown());
    // Force records to be flushed
    oTelTraceRawPrepper.prepareForShutdown();
    oTelTraceRawPrepper.doExecute(Collections.emptyList());
    // Assert records have been flushed
    assertTrue(oTelTraceRawPrepper.isReadyForShutdown());
}

Example 44 with Record

use of com.amazon.dataprepper.model.record.Record in project data-prepper by opensearch-project.

the class OTelTraceRawPrepperTest method testExportRequestFlushByGC.

@Test
public void testExportRequestFlushByGC() throws IOException {
    final ExportTraceServiceRequest exportTraceServiceRequest = buildExportTraceServiceRequestFromJsonFile(TEST_REQUEST_TWO_TRACE_GROUP_MISSING_ROOTS_JSON_FILE);
    oTelTraceRawPrepper.doExecute(Collections.singletonList(new Record<>(exportTraceServiceRequest)));
    await().atMost(2 * TEST_TRACE_FLUSH_INTERVAL, TimeUnit.SECONDS).untilAsserted(() -> {
        final List<Record<String>> processedRecords = (List<Record<String>>) oTelTraceRawPrepper.doExecute(Collections.emptyList());
        Assertions.assertThat(processedRecords.size()).isEqualTo(4);
        Assertions.assertThat(getMissingTraceGroupFieldsSpanCount(processedRecords)).isEqualTo(4);
    });
}

Example 45 with Record

use of com.amazon.dataprepper.model.record.Record in project data-prepper by opensearch-project.

the class OTelTraceGroupPrepper method doExecute.

@Override
public Collection<Record<String>> doExecute(final Collection<Record<String>> rawSpanStringRecords) {
    final List<Record<String>> recordsOut = new LinkedList<>();
    final Map<Record<String>, Map<String, Object>> recordMissingTraceGroupToRawSpanMap = new HashMap<>();
    final Set<String> traceIdsToLookUp = new HashSet<>();
    for (Record<String> record : rawSpanStringRecords) {
        try {
            final Map<String, Object> rawSpanMap = OBJECT_MAPPER.readValue(record.getData(), MAP_TYPE_REFERENCE);
            final String traceGroupName = (String) rawSpanMap.get(TraceGroup.TRACE_GROUP_NAME_FIELD);
            final String traceId = (String) rawSpanMap.get(OTelTraceGroupPrepperConfig.TRACE_ID_FIELD);
            if (Strings.isNullOrEmpty(traceGroupName)) {
                traceIdsToLookUp.add(traceId);
                recordMissingTraceGroupToRawSpanMap.put(record, rawSpanMap);
                recordsInMissingTraceGroupCounter.increment();
            } else {
                recordsOut.add(record);
            }
        } catch (JsonProcessingException e) {
            LOG.error("Failed to parse the record: [{}]", record.getData());
        }
    }
    final Map<String, TraceGroup> traceIdToTraceGroup = searchTraceGroupByTraceIds(traceIdsToLookUp);
    for (final Map.Entry<Record<String>, Map<String, Object>> entry : recordMissingTraceGroupToRawSpanMap.entrySet()) {
        final Record<String> record = entry.getKey();
        final Map<String, Object> rawSpanMap = entry.getValue();
        final String traceId = (String) rawSpanMap.get(OTelTraceGroupPrepperConfig.TRACE_ID_FIELD);
        final TraceGroup traceGroup = traceIdToTraceGroup.get(traceId);
        if (traceGroup != null) {
            try {
                Map<String, Object> traceGroupMap = OBJECT_MAPPER.convertValue(traceGroup, MAP_TYPE_REFERENCE);
                rawSpanMap.putAll(traceGroupMap);
                final String newData = OBJECT_MAPPER.writeValueAsString(rawSpanMap);
                recordsOut.add(new Record<>(newData, record.getMetadata()));
                recordsOutFixedTraceGroupCounter.increment();
            } catch (Exception e) {
                recordsOut.add(record);
                recordsOutMissingTraceGroupCounter.increment();
                LOG.error("Failed to process the raw span: [{}]", record.getData(), e);
            }
        } else {
            recordsOut.add(record);
            recordsOutMissingTraceGroupCounter.increment();
            final String spanId = (String) rawSpanMap.get(OTelTraceGroupPrepperConfig.SPAN_ID_FIELD);
            LOG.warn("Failed to find traceGroup for spanId: {} due to traceGroup missing for traceId: {}", spanId, traceId);
        }
    }
    return recordsOut;
}