diff --git a/sdks/java/io/solace/build.gradle b/sdks/java/io/solace/build.gradle index ef0d49891f08..5436aa4e7bac 100644 --- a/sdks/java/io/solace/build.gradle +++ b/sdks/java/io/solace/build.gradle @@ -27,9 +27,12 @@ description = "Apache Beam :: SDKs :: Java :: IO :: Solace" ext.summary = """IO to read and write to Solace destinations (queues and topics).""" dependencies { + implementation platform(library.java.opentelemetry_bom) implementation enforcedPlatform(library.java.google_cloud_platform_libraries_bom) implementation library.java.vendored_guava_32_1_2_jre implementation project(path: ":sdks:java:core", configuration: "shadow") + implementation library.java.opentelemetry_api + implementation library.java.opentelemetry_context implementation library.java.slf4j_api implementation library.java.joda_time implementation library.java.solace @@ -51,6 +54,8 @@ dependencies { implementation library.java.jackson_databind testImplementation library.java.junit + testImplementation library.java.opentelemetry_sdk + testImplementation library.java.opentelemetry_sdk_testing testImplementation project(path: ":sdks:java:io:common") testImplementation project(path: ":sdks:java:testing:test-utils") testImplementation project(path: ":sdks:java:core", configuration: "shadowTest") diff --git a/sdks/java/io/solace/src/main/java/org/apache/beam/sdk/io/solace/SolaceIO.java b/sdks/java/io/solace/src/main/java/org/apache/beam/sdk/io/solace/SolaceIO.java index 57f65ae62651..610e8481211f 100644 --- a/sdks/java/io/solace/src/main/java/org/apache/beam/sdk/io/solace/SolaceIO.java +++ b/sdks/java/io/solace/src/main/java/org/apache/beam/sdk/io/solace/SolaceIO.java @@ -28,7 +28,18 @@ import com.solacesystems.jcsmp.JCSMPFactory; import com.solacesystems.jcsmp.Queue; import com.solacesystems.jcsmp.Topic; +import io.opentelemetry.api.OpenTelemetry; +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.SpanKind; +import io.opentelemetry.api.trace.Tracer; +import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator; +import io.opentelemetry.context.Context; +import io.opentelemetry.context.Scope; +import io.opentelemetry.context.propagation.TextMapGetter; +import io.opentelemetry.context.propagation.TextMapSetter; import java.io.IOException; +import java.util.HashMap; +import java.util.Map; import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.annotations.Internal; import org.apache.beam.sdk.coders.CannotProvideCoderException; @@ -47,7 +58,9 @@ import org.apache.beam.sdk.io.solace.write.UnboundedBatchedSolaceWriter; import org.apache.beam.sdk.io.solace.write.UnboundedStreamingSolaceWriter; import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.options.SdkHarnessOptions; import org.apache.beam.sdk.schemas.NoSuchSchemaException; +import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.MapElements; import org.apache.beam.sdk.transforms.PTransform; import org.apache.beam.sdk.transforms.ParDo; @@ -423,6 +436,7 @@ public class SolaceIO { Duration.standardSeconds(30); private static final Duration DEFAULT_ACK_DEADLINE = Duration.standardSeconds(30); public static final boolean DEFAULT_NACK_ON_TIMEOUT = false; + public static final boolean DEFAULT_ENABLE_OPENTELEMETRY_TRACING = false; public static final int DEFAULT_WRITER_NUM_SHARDS = 20; public static final int DEFAULT_WRITER_CLIENTS_PER_WORKER = 4; public static final Boolean DEFAULT_WRITER_PUBLISH_LATENCY_METRICS = false; @@ -472,7 +486,8 @@ public static Read read() { .setDeduplicateRecords(DEFAULT_DEDUPLICATE_RECORDS) .setWatermarkIdleDurationThreshold(DEFAULT_WATERMARK_IDLE_DURATION_THRESHOLD) .setAckDeadline(DEFAULT_ACK_DEADLINE) - .setNackOnTimeout(DEFAULT_NACK_ON_TIMEOUT)); + .setNackOnTimeout(DEFAULT_NACK_ON_TIMEOUT) + .setEnableOpenTelemetryTracing(DEFAULT_ENABLE_OPENTELEMETRY_TRACING)); } /** @@ -503,7 +518,8 @@ public static Read read( .setDeduplicateRecords(DEFAULT_DEDUPLICATE_RECORDS) .setWatermarkIdleDurationThreshold(DEFAULT_WATERMARK_IDLE_DURATION_THRESHOLD) .setAckDeadline(DEFAULT_ACK_DEADLINE) - .setNackOnTimeout(DEFAULT_NACK_ON_TIMEOUT)); + .setNackOnTimeout(DEFAULT_NACK_ON_TIMEOUT) + .setEnableOpenTelemetryTracing(DEFAULT_ENABLE_OPENTELEMETRY_TRACING)); } /** @@ -706,6 +722,16 @@ public Read withSessionServiceFactory(SessionServiceFactory sessionServiceFac return this; } + /** + * Optional, default: false. Enables OpenTelemetry distributed tracing by extracting the W3C + * trace context from incoming {@link Solace.Record#getUserProperties()} and creating a consumer + * span around downstream processing. + */ + public Read withEnableOpenTelemetryTracing() { + configurationBuilder.setEnableOpenTelemetryTracing(true); + return this; + } + @AutoValue abstract static class Configuration { @@ -733,9 +759,12 @@ abstract static class Configuration { abstract boolean getNackOnTimeout(); + abstract boolean getEnableOpenTelemetryTracing(); + public static Builder builder() { Builder builder = - new org.apache.beam.sdk.io.solace.AutoValue_SolaceIO_Read_Configuration.Builder(); + new org.apache.beam.sdk.io.solace.AutoValue_SolaceIO_Read_Configuration.Builder() + .setEnableOpenTelemetryTracing(DEFAULT_ENABLE_OPENTELEMETRY_TRACING); return builder; } @@ -767,6 +796,8 @@ abstract Builder setParseFn( abstract Builder setNackOnTimeout(boolean nackOnTimeout); + abstract Builder setEnableOpenTelemetryTracing(boolean enableOpenTelemetryTracing); + abstract Configuration build(); } } @@ -793,20 +824,36 @@ public PCollection expand(PBegin input) { Coder coder = inferCoder(input.getPipeline(), configuration.getTypeDescriptor()); - return input.apply( - org.apache.beam.sdk.io.Read.from( - new UnboundedSolaceSource<>( - initializedQueue, - sempClientFactory, - sessionServiceFactory, - configuration.getMaxNumConnections(), - configuration.getDeduplicateRecords(), - coder, - configuration.getTimestampFn(), - configuration.getWatermarkIdleDurationThreshold(), - configuration.getParseFn(), - configuration.getAckDeadline(), - configuration.getNackOnTimeout()))); + PCollection output = + input.apply( + org.apache.beam.sdk.io.Read.from( + new UnboundedSolaceSource<>( + initializedQueue, + sempClientFactory, + sessionServiceFactory, + configuration.getMaxNumConnections(), + configuration.getDeduplicateRecords(), + coder, + configuration.getTimestampFn(), + configuration.getWatermarkIdleDurationThreshold(), + configuration.getParseFn(), + configuration.getAckDeadline(), + configuration.getNackOnTimeout()))); + + if (configuration.getEnableOpenTelemetryTracing()) { + checkArgument( + TypeDescriptor.of(Solace.Record.class).isSupertypeOf(configuration.getTypeDescriptor()), + "SolaceIO.Read: withEnableOpenTelemetryTracing() requires Solace.Record output " + + "(SolaceIO.read()); user properties are not available after a custom parseFn."); + output = + output + .apply( + "Extract OpenTelemetry context from Header", + ParDo.of(new OpenTelemetryHeaderConsumer<>())) + .setCoder(coder); + } + + return output; } @VisibleForTesting @@ -871,6 +918,118 @@ private Queue initializeQueueForTopicIfNeeded( } } + @VisibleForTesting + static class OpenTelemetryHeaderConsumer extends DoFn { + private transient @Nullable Tracer tracer; + + @Setup + public void setup(PipelineOptions options) { + OpenTelemetry openTelemetry = options.as(SdkHarnessOptions.class).getOpenTelemetry(); + tracer = openTelemetry.getTracer("SolaceIO"); + } + + private static final TextMapGetter HEADER_GETTER = + new TextMapGetter() { + @Override + public Iterable keys(Record carrier) { + return carrier.getUserProperties().keySet(); + } + + @Override + public @Nullable String get(@Nullable Record carrier, String key) { + if (carrier == null) { + return null; + } + Solace.UserPropertyValue direct = carrier.getUserProperties().get(key); + if (direct != null && direct.getString() != null) { + return direct.getString(); + } + for (Map.Entry entry : + carrier.getUserProperties().entrySet()) { + if (entry.getKey().equalsIgnoreCase(key)) { + return entry.getValue().getString(); + } + } + return null; + } + }; + + @ProcessElement + public void processElement(@Element T element, OutputReceiver receiver) { + Tracer currentTracer = this.tracer; + if (currentTracer == null) { + receiver.output(element); + return; + } + + Context context = Context.current(); + if (element instanceof Record) { + context = + W3CTraceContextPropagator.getInstance() + .extract(context, (Record) element, HEADER_GETTER); + } + Span span = + currentTracer + .spanBuilder("SolaceIO.Read") + .setSpanKind(SpanKind.CONSUMER) + .setParent(context) + .startSpan(); + try (Scope scope = span.makeCurrent()) { + receiver.output(element); + } finally { + span.end(); + } + } + } + + @VisibleForTesting + static class OpenTelemetryHeaderPropagator extends DoFn { + private transient @Nullable Tracer tracer; + + @Setup + public void setup(PipelineOptions options) { + OpenTelemetry openTelemetry = options.as(SdkHarnessOptions.class).getOpenTelemetry(); + tracer = openTelemetry.getTracer("SolaceIO"); + } + + private static final TextMapSetter HEADER_SETTER = + new TextMapSetter() { + @Override + public void set(@Nullable Record carrier, String key, String value) { + if (carrier != null && value != null) { + carrier.getUserProperties().put(key, Solace.UserPropertyValue.of(value)); + } + } + }; + + @ProcessElement + public void processElement(@Element Record element, OutputReceiver receiver) { + Tracer currentTracer = this.tracer; + if (currentTracer == null) { + receiver.output(element); + return; + } + + Span span = + currentTracer + .spanBuilder("SolaceIO.Write") + .setSpanKind(SpanKind.PRODUCER) + .setParent(Context.current()) + .startSpan(); + try (Scope scope = span.makeCurrent()) { + Record updatedElement = + element.toBuilder() + .setUserProperties(new HashMap<>(element.getUserProperties())) + .build(); + W3CTraceContextPropagator.getInstance() + .inject(Context.current(), updatedElement, HEADER_SETTER); + receiver.output(updatedElement); + } finally { + span.end(); + } + } + } + public enum SubmissionMode { HIGHER_THROUGHPUT, LOWER_LATENCY, @@ -1061,6 +1220,15 @@ public Write withErrorHandler(ErrorHandler errorHandler) { return toBuilder().setErrorHandler(errorHandler).build(); } + /** + * Optional, default: false. Enables OpenTelemetry distributed tracing by creating a producer + * span and injecting the W3C trace context into outgoing {@link + * Solace.Record#getUserProperties()}. + */ + public Write withEnableOpenTelemetryTracing() { + return toBuilder().setEnableOpenTelemetryTracing(true).build(); + } + abstract int getNumShards(); abstract int getNumberOfClientsPerWorker(); @@ -1081,6 +1249,8 @@ public Write withErrorHandler(ErrorHandler errorHandler) { abstract @Nullable ErrorHandler getErrorHandler(); + abstract boolean getEnableOpenTelemetryTracing(); + static Builder builder() { return new AutoValue_SolaceIO_Write.Builder() .setDeliveryMode(DEFAULT_WRITER_DELIVERY_MODE) @@ -1088,7 +1258,8 @@ static Builder builder() { .setNumberOfClientsPerWorker(DEFAULT_WRITER_CLIENTS_PER_WORKER) .setPublishLatencyMetrics(DEFAULT_WRITER_PUBLISH_LATENCY_METRICS) .setDispatchMode(DEFAULT_WRITER_SUBMISSION_MODE) - .setWriterType(DEFAULT_WRITER_TYPE); + .setWriterType(DEFAULT_WRITER_TYPE) + .setEnableOpenTelemetryTracing(DEFAULT_ENABLE_OPENTELEMETRY_TRACING); } abstract Builder toBuilder(); @@ -1115,6 +1286,8 @@ abstract static class Builder { abstract Builder setErrorHandler(ErrorHandler errorHandler); + abstract Builder setEnableOpenTelemetryTracing(boolean enableOpenTelemetryTracing); + abstract Write build(); } @@ -1146,6 +1319,12 @@ public SolaceOutput expand(PCollection input) { MapElements.into(TypeDescriptor.of(Solace.Record.class)) .via(checkNotNull(getFormatFunction()))); + if (getEnableOpenTelemetryTracing()) { + records = + records.apply( + "Propagate OpenTelemetry Tracing", ParDo.of(new OpenTelemetryHeaderPropagator())); + } + PCollection withGlobalWindow = records.apply("Global window", Window.into(new GlobalWindows())); diff --git a/sdks/java/io/solace/src/main/java/org/apache/beam/sdk/io/solace/data/Solace.java b/sdks/java/io/solace/src/main/java/org/apache/beam/sdk/io/solace/data/Solace.java index b3345dad6e0c..08e20b8761ca 100644 --- a/sdks/java/io/solace/src/main/java/org/apache/beam/sdk/io/solace/data/Solace.java +++ b/sdks/java/io/solace/src/main/java/org/apache/beam/sdk/io/solace/data/Solace.java @@ -435,6 +435,8 @@ public static Builder builder() { .setUserProperties(Collections.emptyMap()); } + public abstract Builder toBuilder(); + @AutoValue.Builder public abstract static class Builder { public abstract Builder setMessageId(String messageId); diff --git a/sdks/java/io/solace/src/test/java/org/apache/beam/sdk/io/solace/MockProducer.java b/sdks/java/io/solace/src/test/java/org/apache/beam/sdk/io/solace/MockProducer.java index a1712633535b..f90e1a14d522 100644 --- a/sdks/java/io/solace/src/test/java/org/apache/beam/sdk/io/solace/MockProducer.java +++ b/sdks/java/io/solace/src/test/java/org/apache/beam/sdk/io/solace/MockProducer.java @@ -21,6 +21,8 @@ import com.solacesystems.jcsmp.Destination; import com.solacesystems.jcsmp.JCSMPException; import java.time.Instant; +import java.util.ArrayList; +import java.util.Collections; import java.util.List; import org.apache.beam.sdk.io.solace.broker.MessageProducer; import org.apache.beam.sdk.io.solace.broker.PublishResultHandler; @@ -29,6 +31,19 @@ import org.apache.beam.sdk.transforms.SerializableFunction; public abstract class MockProducer implements MessageProducer { + private static final List PUBLISHED_RECORDS = + Collections.synchronizedList(new ArrayList<>()); + + public static List getPublishedRecords() { + synchronized (PUBLISHED_RECORDS) { + return new ArrayList<>(PUBLISHED_RECORDS); + } + } + + public static void clearPublishedRecords() { + PUBLISHED_RECORDS.clear(); + } + final PublishResultHandler handler; public MockProducer(PublishResultHandler handler) { @@ -67,6 +82,7 @@ public void publishSingleMessage( Destination topicOrQueue, boolean useCorrelationKeyLatency, DeliveryMode deliveryMode) { + PUBLISHED_RECORDS.add(msg); if (useCorrelationKeyLatency) { handler.responseReceivedEx( Solace.PublishResult.builder() diff --git a/sdks/java/io/solace/src/test/java/org/apache/beam/sdk/io/solace/SolaceIOReadTest.java b/sdks/java/io/solace/src/test/java/org/apache/beam/sdk/io/solace/SolaceIOReadTest.java index 29e82142f94f..637ec2f9c21c 100644 --- a/sdks/java/io/solace/src/test/java/org/apache/beam/sdk/io/solace/SolaceIOReadTest.java +++ b/sdks/java/io/solace/src/test/java/org/apache/beam/sdk/io/solace/SolaceIOReadTest.java @@ -24,15 +24,28 @@ import com.solacesystems.jcsmp.BytesXMLMessage; import com.solacesystems.jcsmp.Destination; +import com.solacesystems.jcsmp.JCSMPFactory; import com.solacesystems.jcsmp.Queue; +import com.solacesystems.jcsmp.SDTException; +import com.solacesystems.jcsmp.SDTMap; import com.solacesystems.jcsmp.Topic; import com.solacesystems.jcsmp.impl.ReplicationGroupMessageIdImpl; +import io.opentelemetry.api.GlobalOpenTelemetry; +import io.opentelemetry.api.trace.SpanKind; +import io.opentelemetry.sdk.OpenTelemetrySdk; +import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; +import io.opentelemetry.sdk.trace.SdkTracerProvider; +import io.opentelemetry.sdk.trace.data.SpanData; +import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; +import io.opentelemetry.sdk.trace.samplers.Sampler; import java.io.IOException; import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.List; +import java.util.Set; import java.util.UUID; import java.util.concurrent.atomic.AtomicInteger; +import java.util.stream.Collectors; import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.io.UnboundedSource.CheckpointMark; import org.apache.beam.sdk.io.UnboundedSource.UnboundedReader; @@ -47,6 +60,7 @@ import org.apache.beam.sdk.io.solace.read.SolaceCheckpointMark; import org.apache.beam.sdk.io.solace.read.UnboundedSolaceSource; import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.options.SdkHarnessOptions; import org.apache.beam.sdk.testing.CoderProperties; import org.apache.beam.sdk.testing.PAssert; import org.apache.beam.sdk.testing.TestPipeline; @@ -57,6 +71,7 @@ import org.apache.beam.sdk.values.TypeDescriptor; import org.apache.beam.sdk.values.TypeDescriptors; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; import org.joda.time.Instant; import org.junit.Rule; import org.junit.Test; @@ -696,4 +711,140 @@ public void testWithNackOnTimeout() { .withNackOnTimeout(false); assertFalse(readWithoutNack.configurationBuilder.build().getNackOnTimeout()); } + + @Test + public void testWithEnableOpenTelemetryTracingConfiguration() { + Read defaultRead = + getDefaultRead() + .withSessionServiceFactory( + MockSessionServiceFactory.builder() + .sessionServiceType(SessionServiceType.EMPTY) + .build()); + assertFalse(defaultRead.configurationBuilder.build().getEnableOpenTelemetryTracing()); + + Read readWithTracing = + getDefaultRead() + .withSessionServiceFactory( + MockSessionServiceFactory.builder() + .sessionServiceType(SessionServiceType.EMPTY) + .build()) + .withEnableOpenTelemetryTracing(); + assertTrue(readWithTracing.configurationBuilder.build().getEnableOpenTelemetryTracing()); + } + + @Test + public void testReadWithOpenTelemetryTracingRejectsCustomType() { + Read read = + SolaceIO.read( + TypeDescriptor.of(SimpleRecord.class), + input -> + new SimpleRecord( + new String(input.getBytes(), StandardCharsets.UTF_8), + input.getApplicationMessageId()), + input -> Instant.ofEpochMilli(1708100477061L)) + .from(Solace.Queue.fromName("queue")) + .withSempClientFactory(MockSempClientFactory.getDefaultMock()) + .withSessionServiceFactory( + MockSessionServiceFactory.builder() + .sessionServiceType(SessionServiceType.EMPTY) + .build()) + .withEnableOpenTelemetryTracing(); + + assertThrows( + IllegalArgumentException.class, () -> read.expand(PBegin.in(TestPipeline.create()))); + } + + @Test + public void testReadWithOpenTelemetryTracing() { + String traceId = "4bf92f3577b34da6a3ce929d0e0e4736"; + String parentSpanId1 = "00f067aa0ba902b7"; + String parentSpanId2 = "11f067aa0ba902b8"; + String traceparent1 = "00-" + traceId + "-" + parentSpanId1 + "-01"; + String traceparent2 = "00-" + traceId + "-" + parentSpanId2 + "-01"; + String tracestate1 = "congo=t61rcWkgMzE"; + + InMemorySpanExporter spanExporter = InMemorySpanExporter.create(); + SdkTracerProvider tracerProvider = + SdkTracerProvider.builder() + .setSampler(Sampler.alwaysOn()) + .addSpanProcessor(SimpleSpanProcessor.create(spanExporter)) + .build(); + GlobalOpenTelemetry.resetForTest(); + OpenTelemetrySdk openTelemetry = + OpenTelemetrySdk.builder().setTracerProvider(tracerProvider).buildAndRegisterGlobal(); + pipeline.getOptions().as(SdkHarnessOptions.class).setOpenTelemetry(openTelemetry); + + try { + SerializableFunction recordFn = + index -> { + BytesXMLMessage msg0 = SolaceDataUtils.getBytesXmlMessage("payload_test0", "450"); + SDTMap props0 = JCSMPFactory.onlyInstance().createMap(); + try { + props0.putString("traceparent", traceparent1); + props0.putString("tracestate", tracestate1); + } catch (SDTException e) { + throw new RuntimeException(e); + } + msg0.setProperties(props0); + + BytesXMLMessage msg1 = SolaceDataUtils.getBytesXmlMessage("payload_test1", "451"); + SDTMap props1 = JCSMPFactory.onlyInstance().createMap(); + try { + props1.putString("TraceParent", traceparent2); + } catch (SDTException e) { + throw new RuntimeException(e); + } + msg1.setProperties(props1); + + BytesXMLMessage msg2 = SolaceDataUtils.getBytesXmlMessage("payload_test2", "452"); + + List messages = ImmutableList.of(msg0, msg1, msg2); + return getOrNull(index, messages); + }; + + SessionServiceFactory fakeSessionServiceFactory = + MockSessionServiceFactory.builder().recordFn(recordFn).minMessagesReceived(3).build(); + + PCollection messageIds = + pipeline + .apply( + "Read from Solace", + getDefaultRead() + .withSessionServiceFactory(fakeSessionServiceFactory) + .withEnableOpenTelemetryTracing()) + .apply(MapElements.into(TypeDescriptors.strings()).via(Record::getMessageId)); + + PAssert.that(messageIds).containsInAnyOrder("450", "451", "452"); + pipeline.run(); + + List spans = spanExporter.getFinishedSpanItems(); + assertEquals(3, spans.size()); + for (SpanData span : spans) { + assertEquals("SolaceIO.Read", span.getName()); + assertEquals(SpanKind.CONSUMER, span.getKind()); + } + + Set parentSpanIds = + spans.stream().map(SpanData::getParentSpanId).collect(Collectors.toSet()); + assertTrue(parentSpanIds.containsAll(ImmutableSet.of(parentSpanId1, parentSpanId2))); + + SpanData span1 = + spans.stream() + .filter(s -> parentSpanId1.equals(s.getParentSpanId())) + .findFirst() + .orElseThrow(AssertionError::new); + assertEquals(traceId, span1.getTraceId()); + assertEquals("t61rcWkgMzE", span1.getParentSpanContext().getTraceState().get("congo")); + + SpanData span2 = + spans.stream() + .filter(s -> parentSpanId2.equals(s.getParentSpanId())) + .findFirst() + .orElseThrow(AssertionError::new); + assertEquals(traceId, span2.getTraceId()); + } finally { + tracerProvider.close(); + GlobalOpenTelemetry.resetForTest(); + } + } } diff --git a/sdks/java/io/solace/src/test/java/org/apache/beam/sdk/io/solace/SolaceIOWriteTest.java b/sdks/java/io/solace/src/test/java/org/apache/beam/sdk/io/solace/SolaceIOWriteTest.java index c67458b2c98b..66a262b179a6 100644 --- a/sdks/java/io/solace/src/test/java/org/apache/beam/sdk/io/solace/SolaceIOWriteTest.java +++ b/sdks/java/io/solace/src/test/java/org/apache/beam/sdk/io/solace/SolaceIOWriteTest.java @@ -18,11 +18,26 @@ package org.apache.beam.sdk.io.solace; import static org.apache.beam.sdk.values.TypeDescriptors.strings; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; +import com.solacesystems.jcsmp.BytesXMLMessage; import com.solacesystems.jcsmp.DeliveryMode; +import io.opentelemetry.api.GlobalOpenTelemetry; +import io.opentelemetry.api.trace.SpanKind; +import io.opentelemetry.sdk.OpenTelemetrySdk; +import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; +import io.opentelemetry.sdk.trace.SdkTracerProvider; +import io.opentelemetry.sdk.trace.data.SpanData; +import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; +import io.opentelemetry.sdk.trace.samplers.Sampler; import java.nio.charset.StandardCharsets; import java.util.List; +import java.util.Map; import java.util.Objects; +import java.util.Set; import java.util.stream.Collectors; import java.util.stream.Stream; import org.apache.beam.sdk.Pipeline; @@ -36,10 +51,13 @@ import org.apache.beam.sdk.io.solace.data.Solace.Record; import org.apache.beam.sdk.io.solace.data.SolaceDataUtils; import org.apache.beam.sdk.io.solace.write.SolaceOutput; +import org.apache.beam.sdk.options.SdkHarnessOptions; import org.apache.beam.sdk.testing.PAssert; import org.apache.beam.sdk.testing.TestPipeline; import org.apache.beam.sdk.testing.TestStream; +import org.apache.beam.sdk.transforms.Create; import org.apache.beam.sdk.transforms.MapElements; +import org.apache.beam.sdk.transforms.ParDo; import org.apache.beam.sdk.transforms.errorhandling.BadRecord; import org.apache.beam.sdk.transforms.errorhandling.ErrorHandler; import org.apache.beam.sdk.transforms.errorhandling.ErrorHandlingTestUtils.ErrorSinkTransform; @@ -47,6 +65,7 @@ import org.apache.beam.sdk.values.PCollection; import org.apache.beam.sdk.values.TypeDescriptor; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; import org.joda.time.Duration; import org.joda.time.Instant; import org.junit.Rule; @@ -373,4 +392,198 @@ public void testWriteWithExceptionRecords() throws Exception { .isEqualTo((long) payloads.size()); pipeline.run(); } + + @Test + public void testWithEnableOpenTelemetryTracingConfiguration() { + SolaceIO.Write defaultWrite = SolaceIO.write(); + assertFalse(defaultWrite.getEnableOpenTelemetryTracing()); + + SolaceIO.Write writeWithTracing = SolaceIO.write().withEnableOpenTelemetryTracing(); + assertTrue(writeWithTracing.getEnableOpenTelemetryTracing()); + } + + @Test + public void testWriteWithOpenTelemetryTracingStreaming() throws Exception { + runWriteWithOpenTelemetryTracingTest(WriterType.STREAMING); + } + + @Test + public void testWriteWithOpenTelemetryTracingBatched() throws Exception { + runWriteWithOpenTelemetryTracingTest(WriterType.BATCHED); + } + + private void runWriteWithOpenTelemetryTracingTest(WriterType writerType) throws Exception { + InMemorySpanExporter spanExporter = InMemorySpanExporter.create(); + SdkTracerProvider tracerProvider = + SdkTracerProvider.builder() + .setSampler(Sampler.alwaysOn()) + .addSpanProcessor(SimpleSpanProcessor.create(spanExporter)) + .build(); + GlobalOpenTelemetry.resetForTest(); + OpenTelemetrySdk openTelemetry = + OpenTelemetrySdk.builder().setTracerProvider(tracerProvider).buildAndRegisterGlobal(); + pipeline.getOptions().as(SdkHarnessOptions.class).setOpenTelemetry(openTelemetry); + MockProducer.clearPublishedRecords(); + + try { + ErrorHandler> errorHandler = + pipeline.registerBadRecordErrorHandler(new ErrorSinkTransform()); + + PCollection records = getRecords(pipeline); + SolaceOutput output = + records.apply( + "Write to Solace", + SolaceIO.write() + .to(Solace.Queue.fromName("queue")) + .withSubmissionMode(SubmissionMode.LOWER_LATENCY) + .withWriterType(writerType) + .withDeliveryMode(DeliveryMode.PERSISTENT) + .withSessionServiceFactory( + MockSessionServiceFactory.builder() + .mode(SubmissionMode.LOWER_LATENCY) + .sessionServiceType(SessionServiceType.WITH_SUCCEEDING_PRODUCER) + .build()) + .withErrorHandler(errorHandler) + .withEnableOpenTelemetryTracing()); + + PCollection ids = getIdsPCollection(output); + PAssert.that(ids).containsInAnyOrder(keys); + errorHandler.close(); + PAssert.that(errorHandler.getOutput()).empty(); + pipeline.run(); + + List spans = spanExporter.getFinishedSpanItems(); + assertEquals(keys.size(), spans.size()); + for (SpanData span : spans) { + assertEquals("SolaceIO.Write", span.getName()); + assertEquals(SpanKind.PRODUCER, span.getKind()); + } + + Set expectedTraceparents = + spans.stream() + .map( + s -> + String.format( + "00-%s-%s-%s", + s.getTraceId(), + s.getSpanId(), + s.getSpanContext().getTraceFlags().asHex())) + .collect(Collectors.toSet()); + + List publishedRecords = MockProducer.getPublishedRecords(); + assertEquals(keys.size(), publishedRecords.size()); + for (Record publishedRecord : publishedRecords) { + Solace.UserPropertyValue traceparentProp = + publishedRecord.getUserProperties().get("traceparent"); + assertNotNull(traceparentProp); + String traceparent = traceparentProp.getString(); + assertNotNull(traceparent); + assertTrue( + "Unexpected traceparent: " + traceparent + ", expected one of: " + expectedTraceparents, + expectedTraceparents.contains(traceparent)); + + BytesXMLMessage jcsmpMsg = Solace.SolaceRecordMapper.toMessage(publishedRecord); + assertNotNull(jcsmpMsg.getProperties()); + assertEquals(traceparent, jcsmpMsg.getProperties().getString("traceparent")); + } + } finally { + MockProducer.clearPublishedRecords(); + tracerProvider.close(); + GlobalOpenTelemetry.resetForTest(); + } + } + + @Test + public void testOpenTelemetryTraceContextPropagationConsumerToProducer() { + String traceId = "4bf92f3577b34da6a3ce929d0e0e4736"; + String upstreamSpanId = "00f067aa0ba902b7"; + String incomingTraceparent = "00-" + traceId + "-" + upstreamSpanId + "-01"; + String incomingTracestate = "congo=t61rcWkgMzE"; + + InMemorySpanExporter spanExporter = InMemorySpanExporter.create(); + SdkTracerProvider tracerProvider = + SdkTracerProvider.builder() + .setSampler(Sampler.alwaysOn()) + .addSpanProcessor(SimpleSpanProcessor.create(spanExporter)) + .build(); + GlobalOpenTelemetry.resetForTest(); + OpenTelemetrySdk openTelemetry = + OpenTelemetrySdk.builder().setTracerProvider(tracerProvider).buildAndRegisterGlobal(); + pipeline.getOptions().as(SdkHarnessOptions.class).setOpenTelemetry(openTelemetry); + + try { + Record inputRecord = + Record.builder() + .setMessageId("id0") + .setPayload("payload_test0".getBytes(StandardCharsets.UTF_8)) + .setUserProperties( + ImmutableMap.of( + "traceparent", Solace.UserPropertyValue.of(incomingTraceparent), + "tracestate", Solace.UserPropertyValue.of(incomingTracestate), + "customKey", Solace.UserPropertyValue.of("customVal"))) + .build(); + + PCollection propagated = + pipeline + .apply(Create.of(inputRecord)) + .apply( + "Extract OpenTelemetry context", + ParDo.of(new SolaceIO.OpenTelemetryHeaderConsumer<>())) + .apply( + "Propagate OpenTelemetry context", + ParDo.of(new SolaceIO.OpenTelemetryHeaderPropagator())); + + PCollection> outputProps = + propagated.apply( + MapElements.into(new TypeDescriptor>() {}) + .via( + r -> + ImmutableMap.of( + "traceparent", + Objects.requireNonNull( + r.getUserProperties().get("traceparent").getString()), + "tracestate", + Objects.requireNonNull( + r.getUserProperties().get("tracestate").getString()), + "customKey", + Objects.requireNonNull( + r.getUserProperties().get("customKey").getString())))); + + PAssert.thatSingleton(outputProps) + .satisfies( + props -> { + assertEquals("customVal", props.get("customKey")); + assertEquals(incomingTracestate, props.get("tracestate")); + assertTrue(props.get("traceparent").startsWith("00-" + traceId + "-")); + return null; + }); + + pipeline.run(); + + List spans = spanExporter.getFinishedSpanItems(); + assertEquals(2, spans.size()); + + SpanData readSpan = + spans.stream() + .filter(s -> "SolaceIO.Read".equals(s.getName())) + .findFirst() + .orElseThrow(AssertionError::new); + SpanData writeSpan = + spans.stream() + .filter(s -> "SolaceIO.Write".equals(s.getName())) + .findFirst() + .orElseThrow(AssertionError::new); + + assertEquals(SpanKind.CONSUMER, readSpan.getKind()); + assertEquals(traceId, readSpan.getTraceId()); + assertEquals(upstreamSpanId, readSpan.getParentSpanId()); + + assertEquals(SpanKind.PRODUCER, writeSpan.getKind()); + assertEquals(traceId, writeSpan.getTraceId()); + assertEquals(readSpan.getSpanId(), writeSpan.getParentSpanId()); + } finally { + tracerProvider.close(); + GlobalOpenTelemetry.resetForTest(); + } + } } diff --git a/sdks/java/io/solace/src/test/java/org/apache/beam/sdk/io/solace/data/SolaceDataUtils.java b/sdks/java/io/solace/src/test/java/org/apache/beam/sdk/io/solace/data/SolaceDataUtils.java index 5e0b8182cf58..2d4c4455eeaf 100644 --- a/sdks/java/io/solace/src/test/java/org/apache/beam/sdk/io/solace/data/SolaceDataUtils.java +++ b/sdks/java/io/solace/src/test/java/org/apache/beam/sdk/io/solace/data/SolaceDataUtils.java @@ -168,6 +168,7 @@ private static BytesXMLMessage getBytesXmlMessageInternal( ? replicationGroupMessageId : DEFAULT_REPLICATION_GROUP_ID; return new BytesXMLMessage() { + private @Nullable SDTMap properties = null; @Override public byte[] getBytes() { @@ -367,7 +368,7 @@ public int getPriority() { @Override public SDTMap getProperties() { - return null; + return properties; } @Override @@ -633,7 +634,9 @@ public void setMessageType(MessageType arg0) {} public void setPriority(int arg0) {} @Override - public void setProperties(SDTMap arg0) {} + public void setProperties(SDTMap arg0) { + this.properties = arg0; + } @Override public void setQueueNameLocation(int arg0, int arg1) {}