Kinesis: fix maxBytesPerSecond throttling (#3035)

akka · Nov 22, 2023 · d19f621 · d19f621
1 parent e257852
commit d19f621
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/kinesis/src/main/scala/akka/stream/alpakka/kinesis/scaladsl/KinesisFlow.scala b/kinesis/src/main/scala/akka/stream/alpakka/kinesis/scaladsl/KinesisFlow.scala
@@ -104,8 +104,8 @@ object KinesisFlow {
   )(result: PutRecordsResponse): List[(PutRecordsResultEntry, T)] =
     result.records.asScala.toList.zip(entries).map { case (res, (_, t)) => (res, t) }
 
-  private def getPayloadByteSize[T](record: (PutRecordsRequestEntry, T)): Int = record match {
-    case (request, _) => request.partitionKey.length + request.data.asByteBuffer.position()
+  private[kinesis] def getPayloadByteSize[T](record: (PutRecordsRequestEntry, T)): Int = record match {
+    case (request, _) => request.partitionKey.length + request.data.asByteArrayUnsafe.length
   }
 
   def byPartitionAndData(

diff --git a/kinesis/src/test/scala/akka/stream/alpakka/kinesis/KinesisFlowSpec.scala b/kinesis/src/test/scala/akka/stream/alpakka/kinesis/KinesisFlowSpec.scala
@@ -56,6 +56,15 @@ class KinesisFlowSpec extends AnyWordSpec with Matchers with KinesisMock with Lo
         sinkProbe.expectError(FailurePublishingRecords(requestError))
       }
     }
+
+    "compute payload size" in {
+      val r = PutRecordsRequestEntry
+        .builder()
+        .partitionKey("")
+        .data(SdkBytes.fromByteBuffer(ByteString("data").asByteBuffer))
+        .build()
+      KinesisFlow.getPayloadByteSize((r, "")) shouldBe 4
+    }
   }
 
   "KinesisFlowWithUserContext" must {