twitter
diff --git a/‎CHANGELOG.rst‎
Lines changed: 3 additions & 0 deletions b/‎CHANGELOG.rst‎
Lines changed: 3 additions & 0 deletions
diff --git a/‎util-core/src/main/scala/com/twitter/io/Reader.scala‎
Lines changed: 43 additions & 0 deletions b/‎util-core/src/main/scala/com/twitter/io/Reader.scala‎
Lines changed: 43 additions & 0 deletions
diff --git a/‎util-core/src/test/scala/com/twitter/io/ReaderTest.scala‎
Lines changed: 68 additions & 9 deletions b/‎util-core/src/test/scala/com/twitter/io/ReaderTest.scala‎
Lines changed: 68 additions & 9 deletions
@@ -19,6 +19,9 @@ New Features
 * util-core: Introducing `Reader.chunked` that chunks the output of a given reader.
   ``PHAB_ID=D206676``
 
+* util-core: Added Reader#framed for consuming data framed by a user supplied function.
+  ``PHAB_ID=D212396``
+
 * util-security: Add `NullSslSession` related objects for use with non-existent
   `SSLSession`s.  ``PHAB_ID=D201421``
 
 
@@ -57,6 +57,7 @@ object Reader {
     }
 
     def read(n: Int): Future[Option[Buf]] = synchronized {
+      // flatMap to `this` to prevent allocating
       if (state.isEmpty) r.read(Int.MaxValue).flatMap(this)
       else {
         val result = state.slice(0, chunkSize)
@@ -68,6 +69,40 @@ object Reader {
     def discard(): Unit = r.discard()
   }
 
+  // see Reader.framed
+  private final class Framed(r: Reader[Buf], framer: Buf => Seq[Buf])
+    extends Reader[Buf] with (Option[Buf] => Future[Option[Buf]]) {
+
+    private[this] var frames: Seq[Buf] = Seq.empty
+
+    // we only enter here when `frames` is empty.
+    def apply(in: Option[Buf]): Future[Option[Buf]] = synchronized {
+      in match {
+        case Some(data) =>
+          frames = framer(data)
+          read(Int.MaxValue)
+        case None =>
+          Future.None
+      }
+    }
+
+    def read(n: Int): Future[Option[Buf]] = synchronized {
+      frames match {
+        case nextFrame :: rst =>
+          frames = rst
+          Future.value(Some(nextFrame))
+        case _ =>
+          // flatMap to `this` to prevent allocating
+          r.read(Int.MaxValue).flatMap(this)
+      }
+    }
+
+    def discard(): Unit = synchronized {
+      frames = Seq.empty
+      r.discard()
+    }
+  }
+
   /**
    * Read the entire bytestream presented by `r`.
    */
@@ -258,4 +293,12 @@ object Reader {
    * }}}
    */
   def copy(r: Reader[Buf], w: Writer[Buf]): Future[Unit] = copy(r, w, Writer.BufferSize)
+
+  /**
+   * Wraps a [[ Reader[Buf] ]] and emits frames as decided by `framer`.
+   *
+   * @note The returned `Reader` may not be thread safe depending on the behavior
+   *       of the framer.
+   */
+  def framed(r: Reader[Buf], framer: Buf => Seq[Buf]): Reader[Buf] = new Framed(r, framer)
 }
@@ -2,14 +2,16 @@ package com.twitter.io
 
 import com.twitter.concurrent.AsyncStream
 import com.twitter.conversions.time._
+import com.twitter.conversions.storage._
 import com.twitter.util.{Await, Awaitable, Future, Promise}
 import java.io.{ByteArrayInputStream, ByteArrayOutputStream}
 import java.util.concurrent.atomic.AtomicBoolean
 import org.mockito.Mockito._
-import org.scalacheck.Gen
+import org.scalacheck.{Arbitrary, Gen}
 import org.scalatest.concurrent.{Eventually, IntegrationPatience}
 import org.scalatest.prop.GeneratorDrivenPropertyChecks
 import org.scalatest.{FunSuite, Matchers}
+import scala.annotation.tailrec
 
 class ReaderTest
     extends FunSuite
@@ -81,16 +83,16 @@ class ReaderTest
     } yield (s, i)
 
     forAll(stringAndChunk) { case (s, i) =>
-      val r = Reader.chunked(Reader.fromBuf(Buf.Utf8(s)), i)
+        val r = Reader.chunked(Reader.fromBuf(Buf.Utf8(s)), i)
 
-      def readLoop(): Unit = await(r.read(Int.MaxValue)) match {
-        case Some(b) =>
-          assert(b.length <= i)
-          readLoop()
-        case None => ()
-      }
+        def readLoop(): Unit = await(r.read(Int.MaxValue)) match {
+          case Some(b) =>
+            assert(b.length <= i)
+            readLoop()
+          case None => ()
+        }
 
-      readLoop()
+        readLoop()
     }
   }
 
@@ -202,4 +204,61 @@ class ReaderTest
     assert(tailEvaluated.get())
   }
 
+  test("Reader.framed reads framed data") {
+    val getByteArrays: Gen[Seq[Buf]] = Gen.listOf(
+      for {
+        // limit arrays to a few kilobytes, otherwise we may generate a very large amount of data
+        numBytes <- Gen.choose(0.bytes.inBytes, 2.kilobytes.inBytes)
+        bytes <- Gen.containerOfN[Array, Byte](numBytes.toInt, Arbitrary.arbitrary[Byte])
+      } yield Buf.ByteArray.Owned(bytes)
+    )
+
+    forAll(getByteArrays) { buffers: Seq[Buf] =>
+      val buffersWithLength = buffers.map(buf => Buf.U32BE(buf.length).concat(buf))
+
+      val r = Reader.framed(BufReader(Buf(buffersWithLength)), new ReaderTest.U32BEFramer())
+
+      // read all of the frames
+      buffers.foreach { buf =>
+        assert(await(r.read(Int.MaxValue)).contains(buf))
+      }
+
+      // make sure the reader signals EOF
+      assert(await(r.read(Int.MaxValue)).isEmpty)
+    }
+  }
+
+  test("Reader.framed reads empty frames") {
+    val r = Reader.framed(BufReader(Buf.U32BE(0)), new ReaderTest.U32BEFramer())
+    assert(await(r.read(Int.MaxValue)).contains(Buf.Empty))
+    assert(await(r.read(Int.MaxValue)).isEmpty)
+  }
+
+}
+
+object ReaderTest {
+
+  /**
+   * Used to test Reader.framed, extract fields in terms of
+   * frames, signified by a 32-bit BE value preceding
+   * each frame.
+   */
+  private class U32BEFramer() extends (Buf => Seq[Buf]) {
+    var state: Buf = Buf.Empty
+
+    @tailrec
+    private def loop(acc: Seq[Buf], buf: Buf): Seq[Buf] = {
+      buf match {
+        case Buf.U32BE(l, d) if d.length >= l =>
+          loop(acc :+ d.slice(0, l), d.slice(l, d.length))
+        case _ =>
+          state = buf
+          acc
+      }
+    }
+
+    def apply(buf: Buf): Seq[Buf] = synchronized {
+      loop(Seq.empty, state concat buf)
+    }
+  }
 }