Fix ThriftSender max span size check

Signed-off-by: Yeonghoon Park <me@yhpark.io>
jaegertracing · Nov 30, 2019 · 5f51594 · 5f51594
1 parent 46798bd
commit 5f51594
Show file tree

Hide file tree

Showing 3 changed files with 71 additions and 15 deletions.
diff --git a/jaeger-thrift/src/main/java/io/jaegertracing/thrift/internal/senders/ThriftSender.java b/jaeger-thrift/src/main/java/io/jaegertracing/thrift/internal/senders/ThriftSender.java
@@ -54,9 +54,9 @@ public int append(JaegerSpan span) throws SenderException {
 
     io.jaegertracing.thriftjava.Span thriftSpan = JaegerThriftSpanConverter.convertSpan(span);
     int spanSize = calculateSpanSize(thriftSpan);
-    if (spanSize > getMaxSpanBytes()) {
+    if ((processBytesSize + spanSize) > getMaxSpanBytes()) {
       throw new SenderException(String.format("ThriftSender received a span that was too large, size = %d, max = %d",
-          spanSize, getMaxSpanBytes()), null, 1);
+          (processBytesSize + spanSize), getMaxSpanBytes()), null, 1);
     }
 
     byteBufferSize += spanSize;

diff --git a/jaeger-thrift/src/test/java/io/jaegertracing/thrift/internal/senders/ThriftSenderTest.java b/jaeger-thrift/src/test/java/io/jaegertracing/thrift/internal/senders/ThriftSenderTest.java
@@ -49,6 +49,21 @@ public void send(Process process, List<Span> spans) throws SenderException {
     sender.calculateSpanSize(null);
   }
 
+  @Test(expected = SenderException.class)
+  public void appendFail() throws Exception {
+    int size = 0;
+    ThriftSender sender = new ThriftSender(ProtocolType.Compact, 0) {
+      @Override
+      public void send(Process process, List<Span> spans) throws SenderException {
+        throw new SenderException("", null, spans.size());
+      }
+    };
+
+    JaegerTracer tracer = new JaegerTracer.Builder("failure").build();
+    sender.append(tracer.buildSpan("flush-fail").start());
+    sender.flush();
+  }
+
   @Test(expected = SenderException.class)
   public void flushFail() throws Exception {
     ThriftSender sender = new ThriftSender(ProtocolType.Compact, 0) {

diff --git a/jaeger-thrift/src/test/java/io/jaegertracing/thrift/internal/senders/UdpSenderTest.java b/jaeger-thrift/src/test/java/io/jaegertracing/thrift/internal/senders/UdpSenderTest.java
@@ -84,21 +84,29 @@ public void tearDown() throws Exception {
     reporter.close();
   }
 
-  @Test(expected = SenderException.class)
-  public void testAppendSpanTooLarge() throws Exception {
+  /**
+   * create a mock span with target size
+   */
+  private JaegerSpan buildSpanWithSize(int targetSpanSize) throws Exception {
     JaegerSpan jaegerSpan = tracer.buildSpan("raza").start();
-    String msg = "";
-    for (int i = 0; i < 10001; i++) {
-      msg += ".";
-      jaegerSpan.log(msg);
-    }
+    jaegerSpan.setTag("mock", "");
+    io.jaegertracing.thriftjava.Span emptySpan = JaegerThriftSpanConverter.convertSpan(jaegerSpan);
 
-    try {
-      sender.append(jaegerSpan);
-    } catch (SenderException e) {
-      assertEquals(e.getDroppedSpanCount(), 1);
-      throw e;
-    }
+    int emptySpanSize = sender.getSize(emptySpan);
+    int freePacketSizeLeft = targetSpanSize - emptySpanSize;
+
+    // each "0" takes 1 byte in UTF-8
+    // size header seems to take 1 byte for every 2^7 string length
+    int lengthBytes = (31 - Integer.numberOfLeadingZeros(freePacketSizeLeft)) / 7;
+    jaegerSpan.setTag("mock", String.join("", Collections.nCopies(freePacketSizeLeft - lengthBytes, "0")));
+
+    io.jaegertracing.thriftjava.Span span = JaegerThriftSpanConverter.convertSpan(jaegerSpan);
+    int spanSize = sender.getSize(span);
+
+    // sanity check
+    assertEquals(targetSpanSize, spanSize);
+
+    return jaegerSpan;
   }
 
   @Test
@@ -131,6 +139,39 @@ public void testAppend() throws Exception {
     assertEquals(expectedNumSpans, result);
   }
 
+  @Test
+  public void testAppendMaxSize() throws Exception {
+    Process process = new Process(tracer.getServiceName())
+            .setTags(JaegerThriftSpanConverter.buildTags(tracer.tags()));
+
+    int processSize = sender.getSize(process);
+
+    JaegerSpan jaegerSpan = buildSpanWithSize(maxPacketSize - UdpSender.EMIT_BATCH_OVERHEAD - processSize);
+
+    sender.append(jaegerSpan);
+
+    // add a span that overflows the limit to hit the last branch
+    int result = sender.append(jaegerSpan);
+
+    assertEquals(1, result);
+  }
+
+  @Test(expected = SenderException.class)
+  public void testAppendSpanTooLarge() throws Exception {
+    Process process = new Process(tracer.getServiceName())
+            .setTags(JaegerThriftSpanConverter.buildTags(tracer.tags()));
+    int processSize = sender.getSize(process);
+
+    JaegerSpan jaegerSpan = buildSpanWithSize(maxPacketSize - UdpSender.EMIT_BATCH_OVERHEAD - processSize + 1);
+
+    try {
+      sender.append(jaegerSpan);
+    } catch (SenderException e) {
+      assertEquals(e.getDroppedSpanCount(), 1);
+      throw e;
+    }
+  }
+
   @Test
   public void testFlushSendsSpan() throws Exception {
     int timeout = 50; // in milliseconds