Respect the DynamoDB scan page size

fs2-aws-dynamodb/src/main/scala/fs2/aws/dynamodb/StreamScan.scala

@@ -16,8 +16,7 @@ trait StreamScan[F[_]] {
   /** Scans Dynamodb into the FS2 stream
     *
     *  @param scanRequest AWS SDK2 DynamoDB scan request
-    *  @param pageSize the page size of the scan, defines how many
-    *                  records will be loaded to the memory
+    *  @param pageSize the maximum number of records requested per scan page; must be positive
     *  @return an fs2 Stream that emits RAW DynamoDB items, terminates, once scan exhausted
     */
   def scanDynamoDB(
@@ -33,24 +32,24 @@ object StreamScan {
           scanRequest: ScanRequest,
           pageSize: Int
       ): Stream[F, Chunk[JMap[String, AttributeValue]]] =
-        for {
+        if (pageSize <= 0) Stream.raiseError[F](new IllegalArgumentException("pageSize must be greater than 0"))
+        else for {
           dispatcher <- Stream.resource(Dispatcher.parallel[F])
           queue      <- Stream.eval(Queue.bounded[F, Option[Chunk[JMap[String, AttributeValue]]]](1))
           sub        <- Stream.eval(Ref[F].of[Option[Subscription]](None))
           error      <- Stream.eval(Deferred[F, Throwable])
           _          <- Stream.eval(
             ddb
-              .scanPaginator(scanRequest)
+              .scanPaginator(scanRequest.toBuilder().limit(pageSize).build())
               .map { publisher =>
-                // subscribe to the paginator, every time we request to deliver next pageSize items from the DDB table
-                // we use FS2 Queue as bounded buffer with size 1, this way we implement back pressure, not allowing
-                // paginator exhaust memory
+                // Limit each DynamoDB response to pageSize items and request one response at a time.
+                // The bounded FS2 queue provides back pressure while a page is consumed.
                 publisher.subscribe(new Subscriber[ScanResponse] {
 
                   override def onSubscribe(s: Subscription): Unit =
                     dispatcher
                       .unsafeRunSync(
-                        sub.set(s.some) >> Async[F].delay(s.request(pageSize.toLong))
+                        sub.set(s.some) >> Async[F].delay(s.request(1))
                       )
 
                   override def onNext(t: ScanResponse): Unit =
@@ -59,7 +58,7 @@ object StreamScan {
                         for {
                           _ <- queue.offer(Chunk(t.items().asScala.toList*).some)
                           s <- sub.get
-                          _ <- Async[F].delay(s.map(_.request(pageSize.toLong)))
+                          _ <- Async[F].delay(s.map(_.request(1)))
                         } yield ()
                       )
 
@@ -71,10 +70,10 @@ object StreamScan {
                 })
               }
           )
-          stream <- Stream.fromQueueNoneTerminated(queue) ++ Stream.eval(error.tryGet).flatMap {
+          stream <- (Stream.fromQueueNoneTerminated(queue) ++ Stream.eval(error.tryGet).flatMap {
             case Some(value) => Stream.raiseError[F](value)
             case None        => Stream.empty
-          }
+          }).onFinalize(sub.get.flatMap(_.traverse_(s => Async[F].delay(s.cancel()))))
         } yield stream
 
     }

fs2-aws-dynamodb/src/test/scala/fs2/aws/dynamodb/StreamScanCancellationSuite.scala

@@ -0,0 +1,52 @@
+package fs2.aws.dynamodb
+
+import cats.effect.IO
+import io.laserdisc.pure.dynamodb.tagless.DynamoDbAsyncClientOp
+import munit.CatsEffectSuite
+import org.mockito.ArgumentMatchers.any
+import org.mockito.Mockito.{doAnswer, mock, verify, when}
+import org.reactivestreams.{Subscriber, Subscription}
+import software.amazon.awssdk.services.dynamodb.model.{ScanRequest, ScanResponse}
+import software.amazon.awssdk.services.dynamodb.paginators.ScanPublisher
+
+import java.util.concurrent.atomic.AtomicBoolean
+
+class StreamScanCancellationSuite extends CatsEffectSuite {
+  test("invalid page sizes fail through the stream error channel") {
+    val ddb = mock(classOf[DynamoDbAsyncClientOp[IO]])
+    StreamScan[IO](ddb)
+      .scanDynamoDB(ScanRequest.builder().build(), 0)
+      .compile
+      .drain
+      .attempt
+      .map(result => assert(result.left.exists(_.isInstanceOf[IllegalArgumentException])))
+  }
+
+  test("early termination cancels the publisher subscription") {
+    val ddb          = mock(classOf[DynamoDbAsyncClientOp[IO]])
+    val publisher    = mock(classOf[ScanPublisher])
+    val subscription = mock(classOf[Subscription])
+    val emitted      = new AtomicBoolean(false)
+    when(ddb.scanPaginator(any(classOf[ScanRequest]))).thenReturn(IO.pure(publisher))
+    doAnswer { invocation =>
+      val subscriber = invocation.getArgument[Subscriber[ScanResponse]](0)
+      doAnswer { _ =>
+        if (emitted.compareAndSet(false, true)) {
+          val thread = new Thread(() => subscriber.onNext(ScanResponse.builder().build()))
+          thread.setDaemon(true)
+          thread.start()
+        }
+        null
+      }.when(subscription).request(1L)
+      subscriber.onSubscribe(subscription)
+      null
+    }.when(publisher).subscribe(any(classOf[Subscriber[ScanResponse]]))
+
+    StreamScan[IO](ddb)
+      .scanDynamoDB(ScanRequest.builder().build(), 3)
+      .take(1)
+      .compile
+      .drain
+      .map(_ => verify(subscription).cancel())
+  }
+}

fs2-aws-dynamodb/src/test/scala/fs2/aws/dynamodb/StreamScanSpec.scala

@@ -40,17 +40,19 @@ class StreamScanSpec extends AnyWordSpec with Matchers with ScalaFutures {
                 )
                 .build()
             )
-            scanned <- StreamScan[IO](ddb)
+            chunks <- StreamScan[IO](ddb)
               .scanDynamoDB(ScanRequest.builder().tableName(tableName).build(), 3)
-              .unchunks
               .compile
-              .toList
-          } yield scanned.map(_.get("name").s()) should contain theSameElementsAs List(
-            "Dmytro",
-            "Barry",
-            "Ryan",
-            "Vlad"
-          )
+              .toVector
+          } yield {
+            chunks.forall(_.size <= 3) shouldBe true
+            chunks.flatMap(_.toList).map(_.get("name").s()) should contain theSameElementsAs List(
+              "Dmytro",
+              "Barry",
+              "Ryan",
+              "Vlad"
+            )
+          }
         }
         .unsafeToFuture()
         .futureValue

fs2-aws-examples/src/main/scala/fs2/aws/examples/DynamoParallelScan.scala

@@ -29,7 +29,6 @@ class DDBReader[F[_]](ss: StreamScan[F]) {
         .segment(segment)
         .totalSegments(totalSegments)
         .tableName(ddbName)
-        .limit(100)
         .build(),
       pageSize
     )