improve error handling on binary propagation

1 files changed, 77 insertions, 1 deletions

diff --git a/kamon-core-tests/src/test/scala/kamon/context/BinaryPropagationSpec.scala b/kamon-core-tests/src/test/scala/kamon/context/BinaryPropagationSpec.scala
index 99e72f59..5681d300 100644
--- a/kamon-core-tests/src/test/scala/kamon/context/BinaryPropagationSpec.scala
+++ b/kamon-core-tests/src/test/scala/kamon/context/BinaryPropagationSpec.scala
@@ -8,6 +8,8 @@ import kamon.context.BinaryPropagation.{ByteStreamReader, ByteStreamWriter}
 import kamon.context.Propagation.{EntryReader, EntryWriter}
 import org.scalatest.{Matchers, OptionValues, WordSpec}
 
+import scala.util.Random
+
 class BinaryPropagationSpec extends WordSpec with Matchers with OptionValues {
 
   "The Binary Context Propagation" should {
@@ -22,6 +24,51 @@ class BinaryPropagationSpec extends WordSpec with Matchers with OptionValues {
       writer.size() shouldBe 0
     }
 
+    "handle malformed data in when reading a context" in {
+      val randomBytes = Array.ofDim[Byte](42)
+      Random.nextBytes(randomBytes)
+
+      val context = binaryPropagation.read(ByteStreamReader.of(randomBytes))
+      context.isEmpty() shouldBe true
+    }
+
+    "handle read failures in an entry reader" in {
+      val context = Context.of(
+        BinaryPropagationSpec.StringKey, "string-value",
+        BinaryPropagationSpec.FailStringKey, "fail-read"
+      )
+      val writer = inspectableByteStreamWriter()
+      binaryPropagation.write(context, writer)
+
+      val rtContext = binaryPropagation.read(ByteStreamReader.of(writer.toByteArray))
+      rtContext.tags shouldBe empty
+      rtContext.get(BinaryPropagationSpec.StringKey) shouldBe "string-value"
+      rtContext.get(BinaryPropagationSpec.FailStringKey) shouldBe null
+    }
+
+    "handle write failures in an entry writer" in {
+      val context = Context.of(
+        BinaryPropagationSpec.StringKey, "string-value",
+        BinaryPropagationSpec.FailStringKey, "fail-write"
+      )
+      val writer = inspectableByteStreamWriter()
+      binaryPropagation.write(context, writer)
+
+      val rtContext = binaryPropagation.read(ByteStreamReader.of(writer.toByteArray))
+      rtContext.tags shouldBe empty
+      rtContext.get(BinaryPropagationSpec.StringKey) shouldBe "string-value"
+      rtContext.get(BinaryPropagationSpec.FailStringKey) shouldBe null
+    }
+
+    "handle write failures in an entry writer when the context is too big" in {
+      val context = Context.of(BinaryPropagationSpec.StringKey, "string-value" * 20)
+      val writer = inspectableByteStreamWriter()
+      binaryPropagation.write(context, writer)
+
+      val rtContext = binaryPropagation.read(ByteStreamReader.of(writer.toByteArray))
+      rtContext shouldBe empty
+    }
+
     "round trip a Context that only has tags" in {
       val context = Context.of(Map("hello" -> "world", "kamon" -> "rulez"))
       val writer = inspectableByteStreamWriter()
@@ -61,9 +108,11 @@ class BinaryPropagationSpec extends WordSpec with Matchers with OptionValues {
   val binaryPropagation = BinaryPropagation.from(
     ConfigFactory.parseString(
       """
-        |
+        |max-outgoing-size = 64
         |entries.incoming.string = "kamon.context.BinaryPropagationSpec$StringEntryCodec"
+        |entries.incoming.failString = "kamon.context.BinaryPropagationSpec$FailStringEntryCodec"
         |entries.outgoing.string = "kamon.context.BinaryPropagationSpec$StringEntryCodec"
+        |entries.outgoing.failString = "kamon.context.BinaryPropagationSpec$FailStringEntryCodec"
         |
       """.stripMargin
     ).withFallback(ConfigFactory.load().getConfig("kamon.propagation")), Kamon)
@@ -76,6 +125,7 @@ class BinaryPropagationSpec extends WordSpec with Matchers with OptionValues {
 object BinaryPropagationSpec {
 
   val StringKey = Context.key[String]("string", null)
+  val FailStringKey = Context.key[String]("failString", null)
   val IntegerKey = Context.key[Int]("integer", 0)
 
   class StringEntryCodec extends EntryReader[ByteStreamReader] with EntryWriter[ByteStreamWriter] {
@@ -95,4 +145,30 @@ object BinaryPropagationSpec {
       }
     }
   }
+
+  class FailStringEntryCodec extends EntryReader[ByteStreamReader] with EntryWriter[ByteStreamWriter] {
+
+    override def read(medium: ByteStreamReader, context: Context): Context = {
+      val valueData = medium.readAll()
+
+      if(valueData.length > 0) {
+        val stringValue = new String(valueData)
+        if(stringValue == "fail-read") {
+          sys.error("The fail string entry reader has triggered")
+        }
+
+        context.withKey(FailStringKey, stringValue)
+      } else context
+    }
+
+    override def write(context: Context, medium: ByteStreamWriter): Unit = {
+      val value = context.get(FailStringKey)
+      if(value != null && value != "fail-write") {
+        medium.write(value.getBytes)
+      } else {
+        medium.write(42) // malformed data on purpose
+        sys.error("The fail string entry writer has triggered")
+      }
+    }
+  }
 }
\ No newline at end of file