Skip to content

span customizer error handling #187

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Andrew Kent (realark) merged 1 commit into from
Oct 9, 2026
Merged

span customizer error handling #187

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
223 changes: 199 additions & 24 deletions braintrust-sdk/src/main/java/dev/braintrust/trace/BraintrustSpanExporter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,18 +2,32 @@

import dev.braintrust.api.BraintrustOpenApiClient;
import dev.braintrust.config.BraintrustConfig;
import io.opentelemetry.api.common.AttributeKey;
import io.opentelemetry.api.common.Attributes;
import io.opentelemetry.api.trace.SpanContext;
import io.opentelemetry.api.trace.SpanKind;
import io.opentelemetry.api.trace.TraceFlags;
import io.opentelemetry.api.trace.TraceState;
import io.opentelemetry.exporter.otlp.http.trace.OtlpHttpSpanExporter;
import io.opentelemetry.sdk.common.CompletableResultCode;
import io.opentelemetry.sdk.common.InstrumentationLibraryInfo;
import io.opentelemetry.sdk.common.InstrumentationScopeInfo;
import io.opentelemetry.sdk.resources.Resource;
import io.opentelemetry.sdk.trace.data.DelegatingSpanData;
import io.opentelemetry.sdk.trace.data.EventData;
import io.opentelemetry.sdk.trace.data.LinkData;
import io.opentelemetry.sdk.trace.data.SpanData;
import io.opentelemetry.sdk.trace.data.StatusData;
import io.opentelemetry.sdk.trace.export.SpanExporter;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.UnaryOperator;
import java.util.stream.Collectors;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;

/**
Expand All @@ -22,6 +36,9 @@
*/
@Slf4j
class BraintrustSpanExporter implements SpanExporter {
private static final AtomicBoolean customizerFailureLogged = new AtomicBoolean();
private static final AttributeKey<String> CUSTOMIZER_CONTEXT =
AttributeKey.stringKey("braintrust.context_json");
private final BraintrustConfig config;
private final String tracesEndpoint;
private final Map<String, SpanExporter> exporterCache = new ConcurrentHashMap<>();
Expand Down Expand Up @@ -51,17 +68,11 @@ public CompletableResultCode export(Collection<SpanData> spans) {
}

// Finish customization for the entire batch before any attachment or span upload.
Map<String, List<SpanData>> spansByParent;
try {
var exportSpans = spans.stream();
if (!config.spanCustomizers().isEmpty()) {
exportSpans = exportSpans.map(this::customizeSpan);
}
spansByParent = exportSpans.collect(Collectors.groupingBy(this::getParentFromSpan));
} catch (Exception e) {
log.error("Failed to customize spans for export", e);
return CompletableResultCode.ofFailure();
var exportSpans = spans.stream();
if (!config.spanCustomizers().isEmpty()) {
exportSpans = exportSpans.map(this::customizeSpan);
}
var spansByParent = exportSpans.collect(Collectors.groupingBy(this::getParentFromSpan));
Comment thread
realark marked this conversation as resolved.

// Export each group with the appropriate x-bt-parent header
var results =
Expand All @@ -86,24 +97,188 @@ private SpanData customizeSpan(SpanData span) {
var traceId = span.getTraceId();
var spanId = span.getSpanId();
var parentSpanId = span.getParentSpanId();
var parentTraceId = span.getParentSpanContext().getTraceId();
var startEpochNanos = span.getStartEpochNanos();
var endEpochNanos = span.getEndEpochNanos();
var parent = span.getAttributes().get(BraintrustSpanProcessor.PARENT);
var current = span;
for (var customizer : config.spanCustomizers()) {
current = customizer.onSpanExport(current);
if (current == null) {
throw new IllegalStateException("SpanCustomizer.onSpanExport must not return null");
try {
for (var customizer : config.spanCustomizers()) {
var next = customizer.onSpanExport(current);
// SpanData is immutable: returning the input requires no further work.
if (next == current) {
continue;
}
if (next == null) {
current = null;
break;
}
current = new CustomizedSpanData(next);
if (!traceId.equals(current.getTraceId())
|| !spanId.equals(current.getSpanId())
|| !parentSpanId.equals(current.getParentSpanId())
|| !traceId.equals(current.getSpanContext().getTraceId())
|| !spanId.equals(current.getSpanContext().getSpanId())
|| !parentTraceId.equals(current.getParentSpanContext().getTraceId())
|| !parentSpanId.equals(current.getParentSpanContext().getSpanId())) {
current = null;
break;
}
}
if (!traceId.equals(current.getTraceId())
|| !spanId.equals(current.getSpanId())
|| !parentSpanId.equals(current.getParentSpanId())
|| !traceId.equals(current.getSpanContext().getTraceId())
|| !spanId.equals(current.getSpanContext().getSpanId())
|| !parentSpanId.equals(current.getParentSpanContext().getSpanId())) {
throw new IllegalStateException(
"SpanCustomizer.onSpanExport must not change trace ID, span ID, or parent"
+ " span ID");
if (current != null) {
return current;
}
} catch (Exception ignored) {
// Never log hook exceptions: even their messages and stacks may contain span content.
}
if (customizerFailureLogged.compareAndSet(false, true)) {
log.error(
"Span customization failed; exporting only span identity and duration with"
+ " context.customizer_error=true. Further failures will not be logged.");
}
var attributes =
Attributes.builder().put(CUSTOMIZER_CONTEXT, "{\"customizer_error\":true}");
if (parent != null) {
attributes.put(BraintrustSpanProcessor.PARENT, parent);
}
return new StrippedSpanData(
SpanContext.create(
traceId, spanId, TraceFlags.getDefault(), TraceState.getDefault()),
SpanContext.create(
parentTraceId,
parentSpanId,
TraceFlags.getDefault(),
TraceState.getDefault()),
startEpochNanos,
endEpochNanos,
attributes.build());
}

/** Read replacements inside the failure boundary; downstream code never calls their getters. */
@Getter
private static final class CustomizedSpanData implements SpanData {
private final String traceId;
private final String spanId;
private final String parentSpanId;
private final SpanContext spanContext;
private final SpanContext parentSpanContext;
private final String name;
private final SpanKind kind;
private final StatusData status;
private final long startEpochNanos;
private final long endEpochNanos;
private final Attributes attributes;
private final List<EventData> events;
private final List<LinkData> links;
private final int totalAttributeCount;
private final int totalRecordedEvents;
private final int totalRecordedLinks;
private final Resource resource;
private final InstrumentationScopeInfo instrumentationScopeInfo;

@SuppressWarnings("deprecation")
private final InstrumentationLibraryInfo instrumentationLibraryInfo;

private final boolean ended;

@SuppressWarnings("deprecation")
private CustomizedSpanData(SpanData span) {
traceId = span.getTraceId();
spanId = span.getSpanId();
parentSpanId = span.getParentSpanId();
spanContext = Objects.requireNonNull(span.getSpanContext());
parentSpanContext = Objects.requireNonNull(span.getParentSpanContext());
name = Objects.requireNonNull(span.getName());
kind = Objects.requireNonNull(span.getKind());
var spanStatus = span.getStatus();
status = StatusData.create(spanStatus.getStatusCode(), spanStatus.getDescription());
startEpochNanos = span.getStartEpochNanos();
endEpochNanos = span.getEndEpochNanos();
attributes = Objects.requireNonNull(span.getAttributes());
// Materialize nested event/link access too, rather than retaining lazy lists or data.
var spanEvents = span.getEvents();
events =
spanEvents.isEmpty()
? List.of()
: spanEvents.stream()
.map(
event ->
EventData.create(
event.getEpochNanos(),
event.getName(),
event.getAttributes(),
event.getTotalAttributeCount()))
.toList();
var spanLinks = span.getLinks();
links =
spanLinks.isEmpty()
? List.of()
: spanLinks.stream()
.map(
link ->
LinkData.create(
link.getSpanContext(),
link.getAttributes(),
link.getTotalAttributeCount()))
.toList();
totalAttributeCount = span.getTotalAttributeCount();
totalRecordedEvents = span.getTotalRecordedEvents();
totalRecordedLinks = span.getTotalRecordedLinks();
resource = Objects.requireNonNull(span.getResource());
instrumentationScopeInfo = Objects.requireNonNull(span.getInstrumentationScopeInfo());
instrumentationLibraryInfo =
Objects.requireNonNull(span.getInstrumentationLibraryInfo());
ended = span.hasEnded();
}

@Override
public boolean hasEnded() {
return ended;
}
}

/** Allowlist only: never delegate content access back to the original or customized span. */
@Getter
private static final class StrippedSpanData implements SpanData {
private final SpanContext spanContext;
private final SpanContext parentSpanContext;
private final long startEpochNanos;
private final long endEpochNanos;
private final Attributes attributes;
private final int totalAttributeCount;
private final String name = "";
private final SpanKind kind = SpanKind.INTERNAL;
private final StatusData status = StatusData.unset();
private final List<EventData> events = List.of();
private final List<LinkData> links = List.of();
private final int totalRecordedEvents = 0;
private final int totalRecordedLinks = 0;
private final Resource resource = Resource.empty();
private final InstrumentationScopeInfo instrumentationScopeInfo =
InstrumentationScopeInfo.empty();

@SuppressWarnings("deprecation")
private final InstrumentationLibraryInfo instrumentationLibraryInfo =
InstrumentationLibraryInfo.empty();

private StrippedSpanData(
SpanContext spanContext,
SpanContext parentSpanContext,
long startEpochNanos,
long endEpochNanos,
Attributes attributes) {
this.spanContext = spanContext;
this.parentSpanContext = parentSpanContext;
this.startEpochNanos = startEpochNanos;
this.endEpochNanos = endEpochNanos;
this.attributes = attributes;
this.totalAttributeCount = attributes.size();
}

@Override
public boolean hasEnded() {
return true;
}
return current;
}

private SpanData convertAttachments(SpanData span) {
Expand Down
75 changes: 70 additions & 5 deletions braintrust-sdk/src/test/java/dev/braintrust/trace/AttachmentProcessorTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@
import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.MethodSource;
import org.junit.jupiter.params.provider.ValueSource;

/**
* Integration tests inspecting actual OTLP payloads and attachment requests. Additional span
Expand Down Expand Up @@ -345,13 +346,23 @@ public int getTotalAttributeCount() {
assertNull(attachmentRequests.poll(300, TimeUnit.MILLISECONDS));
}

@Test
void rejectedBatchDoesNotUploadEarlierSpansAttachments() throws Exception {
@ParameterizedTest
@ValueSource(booleans = {false, true})
void failedCustomizationSkipsOnlyAffectedSpansAttachments(boolean lazyFailure)
throws Exception {
configBuilder.addSpanCustomizer(
new SpanCustomizer() {
@Override
public SpanData onSpanExport(SpanData span) {
if (span.getName().equals("reject")) {
if (lazyFailure) {
return new DelegatingSpanData(span) {
@Override
public String getName() {
throw new IllegalStateException("Lazy redaction failed");
}
};
}
throw new IllegalStateException("Redaction failed");
}
return span;
Expand All @@ -371,6 +382,12 @@ public SpanData onSpanExport(SpanData span) {
source.get("test")
.spanBuilder("reject")
.setAttribute("braintrust.parent", "project_name:second")
.setAttribute(
INPUT_JSON,
"{\"url\":\"data:application/pdf;base64," + BASE64_PDF + "\"}")
.setAttribute(
BraintrustSpanProcessor.OUTPUT_JSON,
"{\"url\":\"data:application/pdf;base64," + BASE64_PDF + "\"}")
.startSpan();
first.end();
second.end();
Expand All @@ -381,10 +398,58 @@ public SpanData onSpanExport(SpanData span) {
((ReadableSpan) second).toSpanData()))
.join(5, TimeUnit.SECONDS);
assertTrue(result.isDone());
assertFalse(result.isSuccess());
assertNull(attachmentRequests.poll(300, TimeUnit.MILLISECONDS));
assertTrue(exports.isEmpty());
assertTrue(result.isSuccess());
assertArrayEquals(
java.util.Base64.getDecoder().decode(BASE64_PNG),
uploads.poll(5, TimeUnit.SECONDS));
assertNotNull(completedUploads.poll(5, TimeUnit.SECONDS));
for (int i = 0; i < 2; i++) {
var body = exports.poll(5, TimeUnit.SECONDS);
assertNotNull(body, "Expected an OTLP export");
var exported =
ExportTraceServiceRequest.parseFrom(body)
.getResourceSpans(0)
.getScopeSpans(0)
.getSpans(0);
if (exported.getName().equals("accept")) {
var input =
exported.getAttributesList().stream()
.filter(a -> a.getKey().equals(INPUT_JSON.getKey()))
.findFirst()
.orElseThrow()
.getValue()
.getStringValue();
assertAttachmentRef(
BraintrustJsonMapper.get().readTree(input).get("url"), "image/png");
} else {
assertEquals("", exported.getName());
assertFalse(
exported.getAttributesList().stream()
.anyMatch(
a ->
a.getKey().equals(INPUT_JSON.getKey())
|| a.getKey()
.equals(
BraintrustSpanProcessor
.OUTPUT_JSON
.getKey())));
assertEquals(
"{\"customizer_error\":true}",
exported.getAttributesList().stream()
.filter(a -> a.getKey().equals("braintrust.context_json"))
.findFirst()
.orElseThrow()
.getValue()
.getStringValue());
}
}
}
// Closing the exporter drains uploads, so absence here cannot race a pending upload.
assertTrue(uploads.isEmpty(), "No attachment from the stripped span may be uploaded");
assertEquals(
List.of("/api/apikey/login", "/attachment", "/upload", "/attachment/status"),
List.copyOf(attachmentRequests));
assertTrue(exports.isEmpty());
}

@Test
Expand Down
Loading
Loading