Make Context implement Serializable, use Codecs instead of default serialization, only BroadcastKeys get serialized

2 files changed, 85 insertions, 1 deletions

diff --git a/kamon-core-tests/src/test/scala/kamon/context/ContextSerializationSpec.scala b/kamon-core-tests/src/test/scala/kamon/context/ContextSerializationSpec.scala
new file mode 100644
index 00000000..f7e7599f
--- /dev/null
+++ b/kamon-core-tests/src/test/scala/kamon/context/ContextSerializationSpec.scala
@@ -0,0 +1,50 @@
+/* =========================================================================================
+ * Copyright © 2013-2017 the kamon project <http://kamon.io/>
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
+ * except in compliance with the License. You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the
+ * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
+ * either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ * =========================================================================================
+ */
+
+package kamon.context
+
+import java.io.{ByteArrayInputStream, ByteArrayOutputStream, ObjectInputStream, ObjectOutputStream}
+
+import kamon.Kamon
+import kamon.testkit.ContextTesting
+import org.scalatest.{Matchers, OptionValues, WordSpec}
+
+class ContextSerializationSpec extends WordSpec with Matchers with ContextTesting with OptionValues {
+  "the Context is Serializable" should {
+    "empty " in {
+      val bos = new ByteArrayOutputStream()
+      val oos = new ObjectOutputStream(bos)
+      oos.writeObject(Context.Empty)
+
+      val ois = new ObjectInputStream(new ByteArrayInputStream(bos.toByteArray))
+      val ctx = ois.readObject().asInstanceOf[Context]
+      ctx shouldBe Context.Empty
+    }
+
+    "full" in {
+      val sCtx = Context(StringBroadcastKey, Some("disi"))
+      val bos = new ByteArrayOutputStream()
+      val oos = new ObjectOutputStream(bos)
+      oos.writeObject(sCtx)
+
+      val ois = new ObjectInputStream(new ByteArrayInputStream(bos.toByteArray))
+      val rCtx = ois.readObject().asInstanceOf[Context]
+      rCtx shouldBe sCtx
+    }
+
+  }
+
+  val ContextCodec = new Codecs(Kamon.config())
+}
\ No newline at end of file
diff --git a/kamon-core/src/main/scala/kamon/context/Context.scala b/kamon-core/src/main/scala/kamon/context/Context.scala
index 85c80d50..340a7803 100644
--- a/kamon-core/src/main/scala/kamon/context/Context.scala
+++ b/kamon-core/src/main/scala/kamon/context/Context.scala
@@ -15,15 +15,48 @@
 
 package kamon.context
 
-class Context private (private[context] val entries: Map[Key[_], Any]) {
+import java.io._
+import java.nio.ByteBuffer
+
+import kamon.Kamon
+
+class Context private (private[context] val entries: Map[Key[_], Any]) extends scala.Serializable {
   def get[T](key: Key[T]): T =
     entries.getOrElse(key, key.emptyValue).asInstanceOf[T]
 
   def withKey[T](key: Key[T], value: T): Context =
     new Context(entries.updated(key, value))
+
+  var _deserializedEntries: Map[Key[_], Any] = Map.empty
+
+  @throws[IOException]
+  private def writeObject(out: ObjectOutputStream): Unit = out.write(
+    Kamon.contextCodec().Binary.encode(this).array()
+  )
+
+  @throws[IOException]
+  @throws[ClassNotFoundException]
+  private def readObject(in: ObjectInputStream): Unit = {
+    val buf = new Array[Byte](in.available())
+    in.readFully(buf)
+    _deserializedEntries = Kamon.contextCodec().Binary.decode(ByteBuffer.wrap(buf)).entries
+  }
+
+  def readResolve(): AnyRef = new Context(_deserializedEntries)
+
+  override def equals(obj: scala.Any): Boolean = {
+    obj != null &&
+    obj.isInstanceOf[Context] &&
+    obj.asInstanceOf[Context].entries != null &&
+    obj.asInstanceOf[Context].entries == this.entries
+  }
+
+  override def hashCode(): Int = entries.hashCode()
+
 }
 
 object Context {
+
   val Empty = new Context(Map.empty)
 
   def apply(): Context =
@@ -37,6 +70,7 @@ object Context {
 
   def create[T](key: Key[T], value: T): Context =
     apply(key, value)
+
 }