From 5f10075da3d0892ce5d5b67186a0e7bf16fb880a Mon Sep 17 00:00:00 2001 From: Carlos Alberto Cortez Date: Wed, 6 Dec 2017 01:24:58 -0600 Subject: [PATCH] Add a BINARY propagator for MockTracer. --- .../java/io/opentracing/mock/MockTracer.java | 82 +++++++++++++++++++ .../io/opentracing/mock/MockTracerTest.java | 36 ++++++++ 2 files changed, 118 insertions(+) diff --git a/opentracing-mock/src/main/java/io/opentracing/mock/MockTracer.java b/opentracing-mock/src/main/java/io/opentracing/mock/MockTracer.java index b0546f46..cf9e0ac6 100644 --- a/opentracing-mock/src/main/java/io/opentracing/mock/MockTracer.java +++ b/opentracing-mock/src/main/java/io/opentracing/mock/MockTracer.java @@ -20,10 +20,17 @@ import io.opentracing.References; import io.opentracing.SpanContext; import io.opentracing.Tracer; +import io.opentracing.propagation.Binary; import io.opentracing.propagation.Format; import io.opentracing.propagation.TextMap; import io.opentracing.util.ThreadLocalScopeManager; +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.ObjectInputStream; +import java.io.ObjectOutputStream; +import java.io.IOException; +import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.HashMap; import java.util.List; @@ -112,6 +119,81 @@ public MockSpan.MockContext extract(Format format, C carrier) { } }; + Propagator BINARY = new Propagator() { + static final int BUFFER_SIZE = 128; + + @Override + public void inject(MockSpan.MockContext ctx, Format format, C carrier) { + if (carrier instanceof Binary) { + Binary binary = (Binary) carrier; + ByteArrayOutputStream stream = new ByteArrayOutputStream(); + ObjectOutputStream objStream = null; + try { + objStream = new ObjectOutputStream(stream); + objStream.writeLong(ctx.spanId()); + objStream.writeLong(ctx.traceId()); + + for (Map.Entry entry : ctx.baggageItems()) { + objStream.writeUTF(entry.getKey()); + objStream.writeUTF(entry.getValue()); + } + + objStream.flush(); // *need* to flush ObjectOutputStream. + binary.write(ByteBuffer.wrap(stream.toByteArray())); + + } catch (IOException e) { + throw new IllegalArgumentException("Corrupted data"); + } finally { + if (objStream != null) { + try { objStream.close(); } catch (Exception e2) {} + } + } + } else { + throw new IllegalArgumentException("Unknown carrier"); + } + } + + @Override + public MockSpan.MockContext extract(Format format, C carrier) { + Long traceId = null; + Long spanId = null; + Map baggage = new HashMap<>(); + + if (carrier instanceof Binary) { + Binary binary = (Binary) carrier; + ByteArrayOutputStream outputStream = new ByteArrayOutputStream(); + ObjectInputStream objStream = null; + try { + ByteBuffer buffer = ByteBuffer.allocate(BUFFER_SIZE); + for (int readBytes = 0; (readBytes = binary.read(buffer)) > 0; buffer.rewind()) { + outputStream.write(buffer.array(), 0, readBytes); + } + + objStream = new ObjectInputStream(new ByteArrayInputStream(outputStream.toByteArray())); + spanId = objStream.readLong(); + traceId = objStream.readLong(); + + while (objStream.available() > 0) { + baggage.put(objStream.readUTF(), objStream.readUTF()); + } + } catch (IOException e) { + } finally { + if (objStream != null) { + try { objStream.close(); } catch (Exception e2) {} + } + } + } else { + throw new IllegalArgumentException("Unknown carrier"); + } + + if (traceId != null && spanId != null) { + return new MockSpan.MockContext(traceId, spanId, baggage); + } + + return null; + } + }; + Propagator TEXT_MAP = new Propagator() { public static final String SPAN_ID_KEY = "spanid"; public static final String TRACE_ID_KEY = "traceid"; diff --git a/opentracing-mock/src/test/java/io/opentracing/mock/MockTracerTest.java b/opentracing-mock/src/test/java/io/opentracing/mock/MockTracerTest.java index 42df1bac..f9c828fe 100644 --- a/opentracing-mock/src/test/java/io/opentracing/mock/MockTracerTest.java +++ b/opentracing-mock/src/test/java/io/opentracing/mock/MockTracerTest.java @@ -19,9 +19,12 @@ import io.opentracing.Span; import io.opentracing.SpanContext; import io.opentracing.Tracer; +import io.opentracing.propagation.Adapters; import io.opentracing.propagation.Format; import io.opentracing.propagation.TextMapExtractAdapter; import io.opentracing.propagation.TextMapInjectAdapter; +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -197,6 +200,39 @@ public void testTextMapPropagatorHttpHeaders() { Assert.assertEquals(finishedSpans.get(0).context().spanId(), finishedSpans.get(1).parentId()); } + @Test + public void testBinaryPropagator() { + MockTracer tracer = new MockTracer(MockTracer.Propagator.BINARY); + { + Span parentSpan = tracer.buildSpan("foo") + .start(); + parentSpan.setBaggageItem("foobag", "fooitem"); + parentSpan.finish(); + + ByteArrayOutputStream injectStream = new ByteArrayOutputStream(); + tracer.inject(parentSpan.context(), Format.Builtin.BINARY, + Adapters.injectBinary(injectStream)); + + ByteArrayInputStream extractStream = new ByteArrayInputStream(injectStream.toByteArray()); + SpanContext extract = tracer.extract(Format.Builtin.BINARY, Adapters.extractBinary(extractStream)); + + Span childSpan = tracer.buildSpan("bar") + .asChildOf(extract) + .start(); + childSpan.setBaggageItem("barbag", "baritem"); + childSpan.finish(); + } + List finishedSpans = tracer.finishedSpans(); + + Assert.assertEquals(2, finishedSpans.size()); + Assert.assertEquals(finishedSpans.get(0).context().traceId(), finishedSpans.get(1).context().traceId()); + Assert.assertEquals(finishedSpans.get(0).context().spanId(), finishedSpans.get(1).parentId()); + Assert.assertEquals("fooitem", finishedSpans.get(0).getBaggageItem("foobag")); + Assert.assertNull(finishedSpans.get(0).getBaggageItem("barbag")); + Assert.assertEquals("fooitem", finishedSpans.get(1).getBaggageItem("foobag")); + Assert.assertEquals("baritem", finishedSpans.get(1).getBaggageItem("barbag")); + } + @Test public void testReset() { MockTracer mockTracer = new MockTracer();