Skip to content

Commit

Permalink
Add nested Thrift object support to SpanProtocol.
Browse files Browse the repository at this point in the history
  • Loading branch information
finlaycurran committed Feb 27, 2024
1 parent 346ef18 commit cc33a7b
Showing 1 changed file with 25 additions and 4 deletions.
29 changes: 25 additions & 4 deletions src/main/java/io/opentracing/thrift/SpanProtocol.java
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,11 @@ public class SpanProtocol extends TProtocolDecorator {
private final ClientSpanDecorator spanDecorator;
static final short SPAN_FIELD_ID = 3333; // Magic number
private boolean oneWay;
private boolean injected;

/**
* Field level tracking to ensure that we only include span data once in the case of nested Thrift objects.
*/
private int level;

/**
* Encloses the specified protocol. Take tracer from GlobalTracer
Expand Down Expand Up @@ -98,13 +102,15 @@ public SpanProtocol(TProtocol protocol, Tracer tracer, SpanHolder spanHolder,

@Override
public void writeMessageBegin(TMessage tMessage) throws TException {
// Always reset the level at message begin. The last field stop will insert the span at this level.
level = 0;

Span span = tracer.buildSpan(tMessage.name)
.withTag(Tags.SPAN_KIND.getKey(), Tags.SPAN_KIND_CLIENT)
.start();
spanHolder.setSpan(span);

oneWay = tMessage.type == TMessageType.ONEWAY;
injected = false;

spanDecorator.decorate(span, tMessage);
super.writeMessageBegin(tMessage);
Expand All @@ -125,7 +131,7 @@ public void writeMessageEnd() throws TException {

@Override
public void writeFieldStop() throws TException {
if (!injected) {
if (level == 0) {
Span span = spanHolder.getSpan();
if (span != null) {
Map<String, String> map = new HashMap<>();
Expand All @@ -139,13 +145,28 @@ public void writeFieldStop() throws TException {
}
super.writeMapEnd();
super.writeFieldEnd();
injected = true;
}
}

super.writeFieldStop();
}

@Override
public void writeFieldBegin(TField tField) throws TException
{
// Begin increments nested level
level++;
super.writeFieldBegin(tField);
}

@Override
public void writeFieldEnd() throws TException
{
// End decrements nested level
level--;
super.writeFieldEnd();
}

@Override
public TMessage readMessageBegin() throws TException {
try {
Expand Down

0 comments on commit cc33a7b

Please sign in to comment.