Examples with TraceGroup com.amazon.dataprepper.plugins.prepper.oteltracegroup.model.TraceGroup used on opensource projects

Example 1 with TraceGroup

use of com.amazon.dataprepper.plugins.prepper.oteltracegroup.model.TraceGroup in project data-prepper by opensearch-project.

the class OTelTraceGroupPrepper method searchTraceGroupByTraceIds.

private Map<String, TraceGroup> searchTraceGroupByTraceIds(final Collection<String> traceIds) {
    final Map<String, TraceGroup> traceIdToTraceGroup = new HashMap<>();
    final SearchRequest searchRequest = createSearchRequest(traceIds);
    try {
        final SearchResponse searchResponse = restHighLevelClient.search(searchRequest, RequestOptions.DEFAULT);
        final SearchHit[] searchHits = searchResponse.getHits().getHits();
        Arrays.asList(searchHits).forEach(searchHit -> {
            final Optional<Map.Entry<String, TraceGroup>> optionalStringTraceGroupEntry = fromSearchHitToMapEntry(searchHit);
            optionalStringTraceGroupEntry.ifPresent(entry -> traceIdToTraceGroup.put(entry.getKey(), entry.getValue()));
        });
    } catch (Exception e) {
        // TODO: retry for status code 429 of OpenSearchException?
        LOG.error("Search request for traceGroup failed for traceIds: {} due to {}", traceIds, e.getMessage());
    }
    return traceIdToTraceGroup;
}

Example 2 with TraceGroup

use of com.amazon.dataprepper.plugins.prepper.oteltracegroup.model.TraceGroup 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;
}

Example 3 with TraceGroup

use of com.amazon.dataprepper.plugins.prepper.oteltracegroup.model.TraceGroup in project data-prepper by opensearch-project.

the class OTelTraceGroupPrepper method fromSearchHitToMapEntry.

private Optional<Map.Entry<String, TraceGroup>> fromSearchHitToMapEntry(final SearchHit searchHit) {
    final DocumentField traceIdDocField = searchHit.field(OTelTraceGroupPrepperConfig.TRACE_ID_FIELD);
    final DocumentField traceGroupNameDocField = searchHit.field(TraceGroup.TRACE_GROUP_NAME_FIELD);
    final DocumentField traceGroupEndTimeDocField = searchHit.field(TraceGroup.TRACE_GROUP_END_TIME_FIELD);
    final DocumentField traceGroupDurationInNanosDocField = searchHit.field(TraceGroup.TRACE_GROUP_DURATION_IN_NANOS_FIELD);
    final DocumentField traceGroupStatusCodeDocField = searchHit.field(TraceGroup.TRACE_GROUP_STATUS_CODE_FIELD);
    if (Stream.of(traceIdDocField, traceGroupNameDocField, traceGroupEndTimeDocField, traceGroupDurationInNanosDocField, traceGroupStatusCodeDocField).allMatch(Objects::nonNull)) {
        final String traceId = traceIdDocField.getValue();
        final String traceGroupName = traceGroupNameDocField.getValue();
        // Restore trailing zeros for thousand, e.g. 2020-08-20T05:40:46.0895568Z -> 2020-08-20T05:40:46.089556800Z
        final String traceGroupEndTime = Instant.parse(traceGroupEndTimeDocField.getValue()).toString();
        final Number traceGroupDurationInNanos = traceGroupDurationInNanosDocField.getValue();
        final Number traceGroupStatusCode = traceGroupStatusCodeDocField.getValue();
        return Optional.of(new AbstractMap.SimpleEntry<>(traceId, new TraceGroup(traceGroupName, traceGroupEndTime, traceGroupDurationInNanos.longValue(), traceGroupStatusCode.intValue())));
    }
    return Optional.empty();
}