From cc387c08d4bba23434c9c825a70eb5ce58677bae Mon Sep 17 00:00:00 2001 From: Andrew Kent Date: Fri, 9 Oct 2026 09:34:30 -0600 Subject: [PATCH] span customizer error handling implement: https://github.com/braintrustdata/braintrust-spec/blob/main/skills/instrumentation-spec/references/features/span-customizers.md#hook-failures --- .../trace/BraintrustSpanExporter.java | 223 ++++++++++-- .../trace/AttachmentProcessorTest.java | 75 +++- .../trace/BraintrustSpanExporterTest.java | 327 ++++++++++++++++-- 3 files changed, 569 insertions(+), 56 deletions(-) diff --git a/braintrust-sdk/src/main/java/dev/braintrust/trace/BraintrustSpanExporter.java b/braintrust-sdk/src/main/java/dev/braintrust/trace/BraintrustSpanExporter.java index 5e5eb28f..66f77689 100644 --- a/braintrust-sdk/src/main/java/dev/braintrust/trace/BraintrustSpanExporter.java +++ b/braintrust-sdk/src/main/java/dev/braintrust/trace/BraintrustSpanExporter.java @@ -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; /** @@ -22,6 +36,9 @@ */ @Slf4j class BraintrustSpanExporter implements SpanExporter { + private static final AtomicBoolean customizerFailureLogged = new AtomicBoolean(); + private static final AttributeKey CUSTOMIZER_CONTEXT = + AttributeKey.stringKey("braintrust.context_json"); private final BraintrustConfig config; private final String tracesEndpoint; private final Map exporterCache = new ConcurrentHashMap<>(); @@ -51,17 +68,11 @@ public CompletableResultCode export(Collection spans) { } // Finish customization for the entire batch before any attachment or span upload. - Map> 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)); // Export each group with the appropriate x-bt-parent header var results = @@ -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 events; + private final List 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 events = List.of(); + private final List 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) { diff --git a/braintrust-sdk/src/test/java/dev/braintrust/trace/AttachmentProcessorTest.java b/braintrust-sdk/src/test/java/dev/braintrust/trace/AttachmentProcessorTest.java index 16b37f1e..004d5adc 100644 --- a/braintrust-sdk/src/test/java/dev/braintrust/trace/AttachmentProcessorTest.java +++ b/braintrust-sdk/src/test/java/dev/braintrust/trace/AttachmentProcessorTest.java @@ -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 @@ -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; @@ -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(); @@ -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 diff --git a/braintrust-sdk/src/test/java/dev/braintrust/trace/BraintrustSpanExporterTest.java b/braintrust-sdk/src/test/java/dev/braintrust/trace/BraintrustSpanExporterTest.java index ec07a5ff..609f851d 100644 --- a/braintrust-sdk/src/test/java/dev/braintrust/trace/BraintrustSpanExporterTest.java +++ b/braintrust-sdk/src/test/java/dev/braintrust/trace/BraintrustSpanExporterTest.java @@ -11,6 +11,8 @@ import io.opentelemetry.api.common.Attributes; import io.opentelemetry.api.trace.Span; import io.opentelemetry.api.trace.SpanContext; +import io.opentelemetry.api.trace.SpanKind; +import io.opentelemetry.api.trace.StatusCode; import io.opentelemetry.api.trace.TraceFlags; import io.opentelemetry.api.trace.TraceState; import io.opentelemetry.context.Context; @@ -18,22 +20,28 @@ import io.opentelemetry.sdk.OpenTelemetrySdk; import io.opentelemetry.sdk.logs.SdkLoggerProvider; import io.opentelemetry.sdk.metrics.SdkMeterProvider; +import io.opentelemetry.sdk.resources.Resource; import io.opentelemetry.sdk.trace.ReadableSpan; import io.opentelemetry.sdk.trace.SdkTracerProvider; 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 java.io.ByteArrayInputStream; import java.io.FileInputStream; import java.net.InetSocketAddress; import java.nio.file.Files; import java.security.*; import java.time.Duration; +import java.util.AbstractList; import java.util.HexFormat; import java.util.List; import java.util.concurrent.BlockingQueue; import java.util.concurrent.CountDownLatch; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; import java.util.zip.GZIPInputStream; import javax.net.ssl.*; import org.junit.jupiter.api.AfterEach; @@ -354,49 +362,235 @@ public SpanData onSpanExport(SpanData span) { assertEquals("project_name:test-project", exported.parentHeader()); } + @Test + void replacementSnapshotPreservesCustomizedContent() throws Exception { + var config = + configBuilder() + .addSpanCustomizer( + new SpanCustomizer() { + @Override + public SpanData onSpanExport(SpanData span) { + return new DelegatingSpanData(span) { + @Override + public String getName() { + return "customized"; + } + + @Override + public SpanKind getKind() { + return SpanKind.CLIENT; + } + + @Override + public StatusData getStatus() { + return StatusData.create( + StatusCode.ERROR, "custom error"); + } + + @Override + public List getEvents() { + return List.of( + EventData.create( + 123, + "custom event", + Attributes.of(KEEP, "event value"), + 3)); + } + + @Override + public int getTotalRecordedEvents() { + return 4; + } + + @Override + public List getLinks() { + return List.of( + LinkData.create( + PARENT_CONTEXT, + Attributes.of(KEEP, "link value"), + 2)); + } + + @Override + public int getTotalRecordedLinks() { + return 5; + } + }; + } + }) + .build(); + var span = exportSpan(config, Attributes.of(KEEP, "span value")).span(); + assertEquals("customized", span.getName()); + assertEquals( + io.opentelemetry.proto.trace.v1.Span.SpanKind.SPAN_KIND_CLIENT, span.getKind()); + assertEquals( + io.opentelemetry.proto.trace.v1.Status.StatusCode.STATUS_CODE_ERROR, + span.getStatus().getCode()); + assertEquals("custom error", span.getStatus().getMessage()); + assertStringAttribute(span, KEEP.getKey(), "span value"); + assertEquals(1, span.getEventsCount()); + assertEquals(3, span.getDroppedEventsCount()); + var event = span.getEvents(0); + assertEquals(123, event.getTimeUnixNano()); + assertEquals("custom event", event.getName()); + assertEquals("event value", event.getAttributes(0).getValue().getStringValue()); + assertEquals(2, event.getDroppedAttributesCount()); + assertEquals(1, span.getLinksCount()); + assertEquals(4, span.getDroppedLinksCount()); + var link = span.getLinks(0); + assertEquals( + PARENT_CONTEXT.getTraceId(), + HexFormat.of().formatHex(link.getTraceId().toByteArray())); + assertEquals( + PARENT_CONTEXT.getSpanId(), + HexFormat.of().formatHex(link.getSpanId().toByteArray())); + assertEquals("link value", link.getAttributes(0).getValue().getStringValue()); + assertEquals(1, link.getDroppedAttributesCount()); + } + @ParameterizedTest @EnumSource(InvalidCustomization.class) - void invalidCustomizerFailsEntireBatchBeforeTransport(InvalidCustomization invalid) - throws Exception { + void invalidCustomizerStripsOnlyAffectedSpan(InvalidCustomization invalid) throws Exception { + var laterCalls = new AtomicInteger(); var config = configBuilder() .addSpanCustomizer( new SpanCustomizer() { @Override public SpanData onSpanExport(SpanData span) { - if (span.getName().equals("invalid")) { - return invalid.customize(span); - } - return withAttributes( - span, + var attributes = span.getAttributes().toBuilder() .put(PARENT, "project_name:customized") - .build()); + .put( + "braintrust.context_json", + "{\"secret\":\"partial\"}") + .put("partial", "must not survive") + .build(); + return new DelegatingSpanData( + withAttributes(span, attributes)) { + @Override + public long getStartEpochNanos() { + return 1; + } + + @Override + public long getEndEpochNanos() { + return 2; + } + }; + } + }) + .addSpanCustomizer( + new SpanCustomizer() { + @Override + public SpanData onSpanExport(SpanData span) { + return span.getName().equals("invalid") + ? invalid.customize(span) + : span; + } + }) + .addSpanCustomizer( + new SpanCustomizer() { + @Override + public SpanData onSpanExport(SpanData span) { + laterCalls.incrementAndGet(); + return span; } }) .build(); - try (var provider = SdkTracerProvider.builder().build(); + try (var provider = + SdkTracerProvider.builder() + .setResource( + Resource.create(Attributes.of(KEEP, "secret resource"))) + .build(); var exporter = new BraintrustSpanExporter(config)) { - var tracer = provider.get("test-tracer"); - var first = tracer.spanBuilder("valid").setParent(parentContext()).startSpan(); - var second = tracer.spanBuilder("invalid").setParent(parentContext()).startSpan(); + var tracer = + provider.tracerBuilder("secret-scope") + .setInstrumentationVersion("secret-version") + .setSchemaUrl("https://secret-schema") + .build(); + var first = + tracer.spanBuilder("valid") + .setParent(parentContext()) + .setAttribute(KEEP, "keep me") + .startSpan(); + var second = + tracer.spanBuilder("invalid") + .setParent(parentContext()) + .setSpanKind(SpanKind.CLIENT) + .addLink(PARENT_CONTEXT, Attributes.of(KEEP, "secret link")) + .setAttribute(PARENT, "experiment_id:original") + .setAttribute("braintrust.input_json", "{\"secret\":true}") + .setAttribute("braintrust.context_json", "{\"secret\":\"original\"}") + .setAttribute(KEEP, "secret attribute") + .startSpan(); + second.addEvent("secret event", Attributes.of(KEEP, "secret event attribute")); + second.setStatus(StatusCode.ERROR, "secret status"); first.end(); second.end(); - - var result = - exporter.export( - List.of( - ((ReadableSpan) first).toSpanData(), - ((ReadableSpan) second).toSpanData())) - .join(10, TimeUnit.SECONDS); - assertTrue(result.isDone(), "Export must finish"); - assertFalse(result.isSuccess(), "Invalid customization must fail the entire batch"); + var original = ((ReadableSpan) second).toSpanData(); + var batch = List.of(((ReadableSpan) first).toSpanData(), original); + + // Each submission runs the hooks again, including previously stripped records. + for (int attempt = 0; attempt < 2; attempt++) { + var result = exporter.export(batch).join(10, TimeUnit.SECONDS); + assertTrue(result.isDone(), "Export must finish"); + assertTrue(result.isSuccess(), "A hook failure must not fail the batch"); + for (int group = 0; group < 2; group++) { + var request = requests.poll(10, TimeUnit.SECONDS); + assertNotNull(request); + var exported = request.decodeSpans(); + assertEquals(1, exported.size()); + var span = exported.get(0); + if (span.getName().equals("valid")) { + assertEquals("project_name:customized", request.parentHeader()); + assertStringAttribute(span, KEEP.getKey(), "keep me"); + } else { + assertEquals("experiment_id:original", request.parentHeader()); + assertEquals("", span.getName()); + assertEquals( + original.getTraceId(), + HexFormat.of().formatHex(span.getTraceId().toByteArray())); + assertEquals( + original.getSpanId(), + HexFormat.of().formatHex(span.getSpanId().toByteArray())); + assertEquals( + original.getParentSpanId(), + HexFormat.of().formatHex(span.getParentSpanId().toByteArray())); + assertEquals(original.getStartEpochNanos(), span.getStartTimeUnixNano()); + assertEquals(original.getEndEpochNanos(), span.getEndTimeUnixNano()); + assertEquals("", span.getTraceState()); + assertEquals( + io.opentelemetry.proto.trace.v1.Span.SpanKind.SPAN_KIND_INTERNAL, + span.getKind()); + assertEquals(0, span.getEventsCount()); + assertEquals(0, span.getLinksCount()); + assertEquals(0, span.getDroppedEventsCount()); + assertEquals(0, span.getDroppedLinksCount()); + assertEquals(0, span.getDroppedAttributesCount()); + assertEquals("", span.getStatus().getMessage()); + assertEquals( + io.opentelemetry.proto.trace.v1.Status.StatusCode.STATUS_CODE_UNSET, + span.getStatus().getCode()); + assertEquals(2, span.getAttributesCount()); + assertStringAttribute(span, PARENT.getKey(), "experiment_id:original"); + assertStringAttribute( + span, "braintrust.context_json", "{\"customizer_error\":true}"); + var resource = request.decode().getResourceSpans(0); + assertEquals(0, resource.getResource().getAttributesCount()); + assertEquals("", resource.getSchemaUrl()); + var scope = resource.getScopeSpans(0); + assertEquals("", scope.getSchemaUrl()); + assertEquals("", scope.getScope().getName()); + assertEquals("", scope.getScope().getVersion()); + assertEquals(0, scope.getScope().getAttributesCount()); + } + } + } + assertEquals(2, laterCalls.get(), "Later hooks run only on the healthy record"); } - assertNull( - requests.poll(300, TimeUnit.MILLISECONDS), - "No span from a failed batch may reach HTTP transport"); - assertTrue(requests.isEmpty(), "No additional export requests are allowed"); + assertTrue(requests.isEmpty()); } private BraintrustConfig.Builder configBuilder() { @@ -492,14 +686,18 @@ private record ExportedSpan( String parentHeader) {} private record CapturedRequest(byte[] body, String contentEncoding, String parentHeader) { - List decodeSpans() throws Exception { + ExportTraceServiceRequest decode() throws Exception { byte[] payload = body; if ("gzip".equalsIgnoreCase(contentEncoding)) { try (var gzip = new GZIPInputStream(new ByteArrayInputStream(body))) { payload = gzip.readAllBytes(); } } - return ExportTraceServiceRequest.parseFrom(payload).getResourceSpansList().stream() + return ExportTraceServiceRequest.parseFrom(payload); + } + + List decodeSpans() throws Exception { + return decode().getResourceSpansList().stream() .flatMap(resource -> resource.getScopeSpansList().stream()) .flatMap(scope -> scope.getSpansList().stream()) .toList(); @@ -509,17 +707,81 @@ List decodeSpans() throws Exception { private enum InvalidCustomization { THROWS, RETURNS_NULL, + ATTRIBUTES_THROWS, + NAME_THROWS, + LAZY_EVENT_LIST_THROWS, + EVENT_DATA_THROWS, TRACE_ID, SPAN_ID, PARENT_SPAN_ID, CONTEXT_TRACE_ID, CONTEXT_SPAN_ID, + CONTEXT_PARENT_TRACE_ID, CONTEXT_PARENT_SPAN_ID; SpanData customize(SpanData span) { return switch (this) { case THROWS -> throw new IllegalStateException("customizer failed"); case RETURNS_NULL -> null; + case ATTRIBUTES_THROWS -> + new DelegatingSpanData(span) { + @Override + public Attributes getAttributes() { + throw new IllegalStateException("sensitive attributes"); + } + }; + case NAME_THROWS -> + new DelegatingSpanData(span) { + @Override + public String getName() { + throw new IllegalStateException("sensitive name"); + } + }; + case LAZY_EVENT_LIST_THROWS -> + new DelegatingSpanData(span) { + @Override + public List getEvents() { + return new AbstractList<>() { + @Override + public EventData get(int index) { + throw new IllegalStateException("sensitive event list"); + } + + @Override + public int size() { + return 1; + } + }; + } + }; + case EVENT_DATA_THROWS -> + new DelegatingSpanData(span) { + @Override + public List getEvents() { + return List.of( + new EventData() { + @Override + public String getName() { + throw new IllegalStateException("sensitive event"); + } + + @Override + public Attributes getAttributes() { + return Attributes.empty(); + } + + @Override + public long getEpochNanos() { + return 123; + } + + @Override + public int getTotalAttributeCount() { + return 0; + } + }); + } + }; case TRACE_ID -> new DelegatingSpanData(span) { @Override @@ -563,6 +825,17 @@ public SpanContext getSpanContext() { span.getSpanContext().getTraceState()); } }; + case CONTEXT_PARENT_TRACE_ID -> + new DelegatingSpanData(span) { + @Override + public SpanContext getParentSpanContext() { + return SpanContext.create( + "ffffffffffffffffffffffffffffffff", + span.getParentSpanId(), + TraceFlags.getSampled(), + TraceState.getDefault()); + } + }; case CONTEXT_PARENT_SPAN_ID -> new DelegatingSpanData(span) { @Override