Merge pull request #17237 from hepin1989/fix-#13516

+act #13516 add getByteString and getBytes method with a given length

akka-actor-tests/src/test/scala/akka/util/ByteStringSpec.scala

@@ -4,24 +4,20 @@
 
 package akka.util
 
-import java.io.{ ByteArrayInputStream, ObjectInputStream, ObjectOutputStream, ByteArrayOutputStream }
+import java.io.{ ByteArrayInputStream, ByteArrayOutputStream, ObjectInputStream, ObjectOutputStream }
+import java.lang.Double.doubleToRawLongBits
+import java.lang.Float.floatToRawIntBits
+import java.nio.{ ByteBuffer, ByteOrder }
+import java.nio.ByteOrder.{ BIG_ENDIAN, LITTLE_ENDIAN }
 
-import org.scalatest.WordSpec
-import org.scalatest.Matchers
-import org.scalatest.prop.Checkers
-import org.scalacheck.Arbitrary
+import org.apache.commons.codec.binary.Hex.encodeHex
 import org.scalacheck.Arbitrary.arbitrary
-import org.scalacheck.Gen
-
-import org.apache.commons.codec.binary.Hex.{ encodeHex, decodeHex }
+import org.scalacheck.{ Arbitrary, Gen }
+import org.scalatest.{ Matchers, WordSpec }
+import org.scalatest.prop.Checkers
 
 import scala.collection.mutable.Builder
 
-import java.nio.{ ByteBuffer }
-import java.nio.ByteOrder, ByteOrder.{ BIG_ENDIAN, LITTLE_ENDIAN }
-import java.lang.Float.floatToRawIntBits
-import java.lang.Double.doubleToRawLongBits
-
 class ByteStringSpec extends WordSpec with Matchers with Checkers {
 
   def genSimpleByteString(min: Int, max: Int) = for {
@@ -482,6 +478,24 @@ class ByteStringSpec extends WordSpec with Matchers with Checkers {
         }
       }
 
+      "getting Bytes with a given length" in {
+        check {
+          slice: ByteStringSlice ⇒
+            val (bytes, _, _) = slice
+            val input = bytes.iterator
+            (input.getBytes(bytes.length).toSeq == bytes) && input.isEmpty
+        }
+      }
+
+      "getting ByteString with a given length" in {
+        check {
+          slice: ByteStringSlice ⇒
+            val (bytes, _, _) = slice
+            val input = bytes.iterator
+            (input.getByteString(bytes.length) == bytes) && input.isEmpty
+        }
+      }
+
       "getting Bytes, using the InputStream wrapper" in {
         // combining skip and both read methods here for more rigorous testing
         check { slice: ByteStringSlice ⇒

akka-actor/src/main/scala/akka/util/ByteIterator.scala

@@ -6,12 +6,9 @@ package akka.util
 
 import java.nio.{ ByteBuffer, ByteOrder }
 
-import scala.collection.{ LinearSeq, IndexedSeqOptimized }
-import scala.collection.mutable.{ Builder, WrappedArray }
-import scala.collection.immutable.{ IndexedSeq, VectorBuilder }
-import scala.collection.generic.CanBuildFrom
-import scala.collection.mutable.{ ListBuffer }
 import scala.annotation.tailrec
+import scala.collection.LinearSeq
+import scala.collection.mutable.ListBuffer
 import scala.reflect.ClassTag
 
 object ByteIterator {
@@ -203,7 +200,7 @@ object ByteIterator {
     final override def len: Int = iterators.foldLeft(0) { _ + _.len }
 
     final override def length: Int = {
-      var result = len
+      val result = len
       clear()
       result
     }
@@ -556,6 +553,26 @@ abstract class ByteIterator extends BufferedIterator[Byte] {
    */
   def getBytes(xs: Array[Byte], offset: Int, n: Int): this.type
 
+  /**
+   * Get a specific number of Bytes from this iterator. In contrast to
+   * copyToArray, this method will fail if this.len < n.
+   */
+  def getBytes(n: Int): Array[Byte] = {
+    val bytes = new Array[Byte](n)
+    getBytes(bytes, 0, n)
+    bytes
+  }
+
+  /**
+   * Get a ByteString with specific number of Bytes from this iterator. In contrast to
+   * copyToArray, this method will fail if this.len < n.
+   */
+  def getByteString(n: Int): ByteString = {
+    val bs = clone.take(n).toByteString
+    drop(n)
+    bs
+  }
+
   /**
    * Get a number of Shorts from this iterator.
    */