Fixes to Cryo support

typelevel · Jun 1, 2024 · 081b6c3 · 081b6c3
1 parent d6712bc
commit 081b6c3
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 8 deletions.
diff --git a/core/shared/src/main/scala/fabric/cryo/ByteBufferPool.scala b/core/shared/src/main/scala/fabric/cryo/ByteBufferPool.scala
@@ -50,7 +50,7 @@ object ByteBufferPool extends ObjectPool[ByteBuffer] {
     }
 
   override protected def resetForPool(bb: ByteBuffer): Option[ByteBuffer] = {
-    bb.rewind()
+    bb.clear()
     Some(bb)
   }
 

diff --git a/core/shared/src/main/scala/fabric/cryo/Cryo.scala b/core/shared/src/main/scala/fabric/cryo/Cryo.scala
@@ -66,8 +66,9 @@ object Cryo {
       }
     case Str(s, _) =>
       bb.put(identifiers.Str)
-      bb.putInt(s.length)
-      bb.put(s.getBytes("UTF-8"))
+      val bytes = s.getBytes("UTF-8")
+      bb.putInt(bytes.length)
+      bb.put(bytes)
       ()
     case NumInt(l, _) =>
       bb.put(identifiers.NumInt)
@@ -113,8 +114,9 @@ object Cryo {
         bytes
       }
     catch {
-      case _: BufferOverflowException if ByteBufferPool.AutoResize =>
+      case t: BufferOverflowException if ByteBufferPool.AutoResize =>
         val size = bytes(json)
+        if (size < ByteBufferPool.ByteBufferSize) throw t
         println(
           s"*** WARNING: ByteBufferPool overflowed (Tried: $size, Actual: ${ByteBufferPool.ByteBufferSize}). Resizing the pool and trying again."
         )
@@ -149,9 +151,14 @@ object Cryo {
     case v => throw new UnsupportedOperationException(s"Unsupported: $v")
   }
 
-  def thaw(bytes: Array[Byte]): Json = ByteBufferPool.use { bb =>
-    bb.put(bytes)
-    bb.flip()
-    thaw(bb)
+  def thaw(bytes: Array[Byte]): Json = {
+    if (ByteBufferPool.ByteBufferSize < bytes.length) {
+      ByteBufferPool.resizePool(bytes.length)
+    }
+    ByteBufferPool.use { bb =>
+      bb.put(bytes)
+      bb.flip()
+      thaw(bb)
+    }
   }
 }
diff --git a/core/shared/src/test/scala/spec/CryoSpec.scala b/core/shared/src/test/scala/spec/CryoSpec.scala
@@ -64,6 +64,26 @@ class CryoSpec extends AnyWordSpec with Matchers {
       val thawed = Cryo.thaw(bytes)
       json should be(thawed)
     }
+    "freeze and that a multi-type obj using the pool" in {
+      val json = obj(
+        "string" -> "Hello, World",
+        "integer" -> 42,
+        "bigDecimal" -> BigDecimal(123),
+        "array" -> arr(
+          1,
+          "two",
+          BigDecimal(3)
+        ),
+        "bool" -> true,
+        "null" -> Null
+      )
+      val bytes = Cryo.freeze(json)
+      val thawed = Cryo.thaw(bytes)
+      json should be(thawed)
+    }
+    "freeze and thaw null" in {
+      Cryo.thaw(Cryo.freeze(Null)) should be(Null)
+    }
     "dispose the pool and overflow the new pool to verify resizing" in {
       ByteBufferPool.dispose()
       ByteBufferPool.ByteBufferSize = 10