From cc33a7b93d2103efed3d11266a58701ffd25b6a5 Mon Sep 17 00:00:00 2001
From: Finlay Curran <finlay.curran@spond.com>
Date: Tue, 27 Feb 2024 15:21:27 +0100
Subject: [PATCH] Add nested Thrift object support to SpanProtocol.

Modification taken from: https://github.com/jspofford/java-thrift
Also see: https://github.com/opentracing-contrib/java-thrift/issues/19
---
 .../io/opentracing/thrift/SpanProtocol.java   | 29 ++++++++++++++++---
 1 file changed, 25 insertions(+), 4 deletions(-)

diff --git a/src/main/java/io/opentracing/thrift/SpanProtocol.java b/src/main/java/io/opentracing/thrift/SpanProtocol.java
index d441601..116553f 100644
--- a/src/main/java/io/opentracing/thrift/SpanProtocol.java
+++ b/src/main/java/io/opentracing/thrift/SpanProtocol.java
@@ -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
@@ -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);
@@ -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<>();
@@ -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 {