Use Modify instead of Get / Update for Ref Updates (#134)

* Use Modify instead of Get / Update for Ref Updates * Parallel PutRecords
etspaceman · Jul 6, 2021 · 10d17a1 · 10d17a1
1 parent 169dd64
commit 10d17a1
Show file tree

Hide file tree

Showing 21 changed files with 244 additions and 218 deletions.
diff --git a/src/fun/scala/kinesis/mock/PutRecordTests.scala b/src/fun/scala/kinesis/mock/PutRecordTests.scala
@@ -18,7 +18,7 @@ class PutRecordTests extends munit.CatsEffectSuite with AwsFunctionalTests {
     for {
       recordRequests <- IO(
         putRecordRequestArb.arbitrary
-          .take(5)
+          .take(20)
           .toVector
           .map(_.copy(streamName = resources.streamName))
           .map(x =>
@@ -34,7 +34,7 @@ class PutRecordTests extends munit.CatsEffectSuite with AwsFunctionalTests {
               .build()
           )
       )
-      _ <- recordRequests.traverse(x =>
+      _ <- recordRequests.parTraverse(x =>
         resources.kinesisClient.putRecord(x).toIO
       )
       shards <- resources.kinesisClient
@@ -68,7 +68,7 @@ class PutRecordTests extends munit.CatsEffectSuite with AwsFunctionalTests {
       )
       res = gets.flatMap(_.records().asScala.toVector)
     } yield assert(
-      res.length == 5 && res.forall(rec =>
+      res.length == 20 && res.forall(rec =>
         recordRequests.exists(req =>
           req.data.asByteArray.sameElements(rec.data.asByteArray)
             && req.partitionKey == rec.partitionKey

diff --git a/src/fun/scala/kinesis/mock/PutRecordsTests.scala b/src/fun/scala/kinesis/mock/PutRecordsTests.scala
@@ -16,27 +16,29 @@ class PutRecordsTests extends munit.CatsEffectSuite with AwsFunctionalTests {
 
   fixture.test("It should put records") { resources =>
     for {
-      req <- IO(
-        PutRecordsRequest
-          .builder()
-          .records(
-            putRecordsRequestEntryArb.arbitrary
-              .take(5)
-              .toVector
-              .map(x =>
-                PutRecordsRequestEntry
-                  .builder()
-                  .data(SdkBytes.fromByteArray(x.data))
-                  .partitionKey(x.partitionKey)
-                  .maybeTransform(x.explicitHashKey)(_.explicitHashKey(_))
-                  .build()
-              )
-              .asJava
-          )
-          .streamName(resources.streamName.streamName)
-          .build()
+      reqs <- IO(
+        List.fill(10)(
+          PutRecordsRequest
+            .builder()
+            .records(
+              putRecordsRequestEntryArb.arbitrary
+                .take(5)
+                .toVector
+                .map(x =>
+                  PutRecordsRequestEntry
+                    .builder()
+                    .data(SdkBytes.fromByteArray(x.data))
+                    .partitionKey(x.partitionKey)
+                    .maybeTransform(x.explicitHashKey)(_.explicitHashKey(_))
+                    .build()
+                )
+                .asJava
+            )
+            .streamName(resources.streamName.streamName)
+            .build()
+        )
       )
-      _ <- resources.kinesisClient.putRecords(req).toIO
+      _ <- reqs.parTraverse(req => resources.kinesisClient.putRecords(req).toIO)
       shards <- resources.kinesisClient
         .listShards(
           ListShardsRequest
@@ -67,14 +69,15 @@ class PutRecordsTests extends munit.CatsEffectSuite with AwsFunctionalTests {
           .toIO
       )
       res = gets.flatMap(_.records().asScala.toVector)
+      records = reqs.flatMap(_.records.asScala.toVector)
     } yield assert(
-      res.length == 5 && res.forall(rec =>
-        req.records.asScala.toVector.exists(req =>
+      res.length == 50 && res.forall(rec =>
+        records.exists(req =>
           req.data.asByteArray.sameElements(rec.data.asByteArray)
             && req.partitionKey == rec.partitionKey
         )
       ),
-      s"$res\n${req.records()}"
+      s"$res\n$records"
     )
   }
 }
diff --git a/src/main/scala/kinesis/mock/api/AddTagsToStreamRequest.scala b/src/main/scala/kinesis/mock/api/AddTagsToStreamRequest.scala
@@ -8,6 +8,7 @@ import cats.syntax.all._
 import io.circe
 
 import kinesis.mock.models._
+import kinesis.mock.syntax.either._
 import kinesis.mock.validations.CommonValidations
 
 // https://docs.aws.amazon.com/kinesis/latest/APIReference/API_AddTagsToStream.html
@@ -19,7 +20,7 @@ final case class AddTagsToStreamRequest(
 ) {
   def addTagsToStream(
       streamsRef: Ref[IO, Streams]
-  ): IO[Response[Unit]] = streamsRef.get.flatMap { streams =>
+  ): IO[Response[Unit]] = streamsRef.modify { streams =>
     CommonValidations
       .validateStreamName(streamName)
       .flatMap(_ =>
@@ -62,11 +63,10 @@ final case class AddTagsToStreamRequest(
             ).mapN((_, _, _, _, _) => stream)
           )
       )
-      .traverse(stream =>
-        streamsRef.update(x =>
-          x.updateStream(stream.copy(tags = stream.tags |+| tags))
-        )
+      .map(stream =>
+        (streams.updateStream(stream.copy(tags = stream.tags |+| tags)), ())
       )
+      .sequenceWithDefault(streams)
   }
 }
 

diff --git a/src/main/scala/kinesis/mock/api/CreateStreamRequest.scala b/src/main/scala/kinesis/mock/api/CreateStreamRequest.scala
@@ -2,12 +2,13 @@ package kinesis.mock
 package api
 
 import cats.Eq
+import cats.effect.IO
 import cats.effect.concurrent.Ref
-import cats.effect.{Concurrent, IO}
 import cats.syntax.all._
 import io.circe
 
 import kinesis.mock.models._
+import kinesis.mock.syntax.either._
 import kinesis.mock.validations.CommonValidations
 
 final case class CreateStreamRequest(shardCount: Int, streamName: StreamName) {
@@ -16,8 +17,8 @@ final case class CreateStreamRequest(shardCount: Int, streamName: StreamName) {
       shardLimit: Int,
       awsRegion: AwsRegion,
       awsAccountId: AwsAccountId
-  )(implicit C: Concurrent[IO]): IO[Response[Unit]] =
-    streamsRef.get.flatMap { streams =>
+  ): IO[Response[Unit]] =
+    streamsRef.modify { streams =>
       (
         CommonValidations.validateStreamName(streamName),
         if (streams.streams.contains(streamName))
@@ -36,16 +37,14 @@ final case class CreateStreamRequest(shardCount: Int, streamName: StreamName) {
           ).asLeft
         else Right(()),
         CommonValidations.validateShardLimit(shardCount, streams, shardLimit)
-      ).traverseN { (_, _, _, _, _) =>
+      ).mapN { (_, _, _, _, _) =>
         val newStream =
           StreamData.create(shardCount, streamName, awsRegion, awsAccountId)
-        for {
-          res <- streamsRef
-            .update(x =>
-              x.copy(streams = x.streams + (streamName -> newStream))
-            )
-        } yield res
-      }
+        (
+          streams.copy(streams = streams.streams + (streamName -> newStream)),
+          ()
+        )
+      }.sequenceWithDefault(streams)
     }
 }
 

diff --git a/src/main/scala/kinesis/mock/api/DecreaseStreamRetentionPeriodRequest.scala b/src/main/scala/kinesis/mock/api/DecreaseStreamRetentionPeriodRequest.scala
@@ -10,6 +10,7 @@ import cats.syntax.all._
 import io.circe
 
 import kinesis.mock.models._
+import kinesis.mock.syntax.either._
 import kinesis.mock.validations.CommonValidations
 
 // https://docs.aws.amazon.com/kinesis/latest/APIReference/API_DecreaseStreamRetention.html
@@ -19,7 +20,7 @@ final case class DecreaseStreamRetentionPeriodRequest(
 ) {
   def decreaseStreamRetention(
       streamsRef: Ref[IO, Streams]
-  ): IO[Response[Unit]] = streamsRef.get.flatMap { streams =>
+  ): IO[Response[Unit]] = streamsRef.modify { streams =>
     CommonValidations
       .validateStreamName(streamName)
       .flatMap(_ =>
@@ -38,13 +39,15 @@ final case class DecreaseStreamRetentionPeriodRequest(
             ).mapN((_, _, _) => stream)
           )
       )
-      .traverse(stream =>
-        streamsRef.update(x =>
-          x.updateStream(
+      .map(stream =>
+        (
+          streams.updateStream(
             stream.copy(retentionPeriod = retentionPeriodHours.hours)
-          )
+          ),
+          ()
         )
       )
+      .sequenceWithDefault(streams)
   }
 }
 

diff --git a/src/main/scala/kinesis/mock/api/DeleteStreamRequest.scala b/src/main/scala/kinesis/mock/api/DeleteStreamRequest.scala
@@ -8,6 +8,7 @@ import cats.syntax.all._
 import io.circe
 
 import kinesis.mock.models._
+import kinesis.mock.syntax.either._
 import kinesis.mock.validations.CommonValidations
 
 final case class DeleteStreamRequest(
@@ -17,7 +18,7 @@ final case class DeleteStreamRequest(
   def deleteStream(
       streamsRef: Ref[IO, Streams]
   ): IO[Response[Unit]] =
-    streamsRef.get.flatMap { streams =>
+    streamsRef.modify { streams =>
       CommonValidations
         .validateStreamName(streamName)
         .flatMap(_ =>
@@ -37,7 +38,7 @@ final case class DeleteStreamRequest(
               ).mapN((_, _) => stream)
             )
         )
-        .traverse { stream =>
+        .map { stream =>
           val deletingStream = Map(
             streamName -> stream.copy(
               shards = Map.empty,
@@ -47,16 +48,14 @@ final case class DeleteStreamRequest(
               consumers = Map.empty
             )
           )
-
-          for {
-            res <- streamsRef
-              .update(streams =>
-                streams.copy(
-                  streams = streams.streams ++ deletingStream
-                )
-              )
-          } yield res
+          (
+            streams.copy(
+              streams = streams.streams ++ deletingStream
+            ),
+            ()
+          )
         }
+        .sequenceWithDefault(streams)
     }
 }
 

diff --git a/src/main/scala/kinesis/mock/api/DeregisterStreamConsumerRequest.scala b/src/main/scala/kinesis/mock/api/DeregisterStreamConsumerRequest.scala
@@ -8,6 +8,7 @@ import cats.syntax.all._
 import io.circe
 
 import kinesis.mock.models._
+import kinesis.mock.syntax.either._
 import kinesis.mock.validations.CommonValidations
 
 // https://docs.aws.amazon.com/kinesis/latest/APIReference/API_DeregisterStreamConsumer.html
@@ -17,26 +18,26 @@ final case class DeregisterStreamConsumerRequest(
     streamArn: Option[String]
 ) {
   private def deregister(
-      streamsRef: Ref[IO, Streams],
+      streams: Streams,
       consumer: Consumer,
       stream: StreamData
-  ): IO[Consumer] = {
-    val newConsumer = consumer.copy(consumerStatus = ConsumerStatus.DELETING)
+  ): (Streams, Consumer) = {
+    val newConsumer =
+      consumer.copy(consumerStatus = ConsumerStatus.DELETING)
 
-    streamsRef
-      .update(x =>
-        x.updateStream(
-          stream.copy(consumers =
-            stream.consumers + (consumer.consumerName -> newConsumer)
-          )
+    (
+      streams.updateStream(
+        stream.copy(consumers =
+          stream.consumers + (consumer.consumerName -> newConsumer)
         )
-      )
-      .as(newConsumer)
+      ),
+      newConsumer
+    )
   }
 
   def deregisterStreamConsumer(
       streamsRef: Ref[IO, Streams]
-  ): IO[Response[Consumer]] = streamsRef.get.flatMap { streams =>
+  ): IO[Response[Consumer]] = streamsRef.modify { streams =>
     (consumerArn, consumerName, streamArn) match {
       case (Some(cArn), _, _) =>
         CommonValidations
@@ -50,9 +51,10 @@ final case class DeregisterStreamConsumerRequest(
                 s"Consumer $consumerName is not in an ACTIVE state"
               ).asLeft
           }
-          .traverse { case (consumer, stream) =>
-            deregister(streamsRef, consumer, stream)
+          .map { case (consumer, stream) =>
+            deregister(streams, consumer, stream)
           }
+          .sequenceWithDefault(streams)
       case (None, Some(cName), Some(sArn)) =>
         CommonValidations
           .findStreamByArn(sArn, streams)
@@ -68,11 +70,13 @@ final case class DeregisterStreamConsumerRequest(
 
             }
           }
-          .traverse { case (consumer, stream) =>
-            deregister(streamsRef, consumer, stream)
+          .map { case (consumer, stream) =>
+            deregister(streams, consumer, stream)
           }
+          .sequenceWithDefault(streams)
       case _ =>
-        IO(
+        (
+          streams,
           InvalidArgumentException(
             "ConsumerArn or both ConsumerName and StreamARN are required for this request."
           ).asLeft

diff --git a/src/main/scala/kinesis/mock/api/DisableEnhancedMonitoringRequest.scala b/src/main/scala/kinesis/mock/api/DisableEnhancedMonitoringRequest.scala
@@ -4,10 +4,10 @@ package api
 import cats.Eq
 import cats.effect.IO
 import cats.effect.concurrent.Ref
-import cats.syntax.all._
 import io.circe
 
 import kinesis.mock.models._
+import kinesis.mock.syntax.either._
 import kinesis.mock.validations.CommonValidations
 
 // https://docs.aws.amazon.com/kinesis/latest/APIReference/API_DisableEnhancedMonitoring.html
@@ -18,33 +18,31 @@ final case class DisableEnhancedMonitoringRequest(
   def disableEnhancedMonitoring(
       streamsRef: Ref[IO, Streams]
   ): IO[Response[DisableEnhancedMonitoringResponse]] =
-    streamsRef.get.flatMap { streams =>
+    streamsRef.modify { streams =>
       CommonValidations
         .validateStreamName(streamName)
         .flatMap(_ => CommonValidations.findStream(streamName, streams))
-        .traverse { stream =>
+        .map { stream =>
           val current =
             stream.enhancedMonitoring.flatMap(_.shardLevelMetrics)
           val desired =
             if (shardLevelMetrics.contains(ShardLevelMetric.ALL))
               Vector.empty
             else current.diff(shardLevelMetrics)
 
-          streamsRef
-            .update(x =>
-              x.updateStream(
-                stream
-                  .copy(enhancedMonitoring = Vector(ShardLevelMetrics(desired)))
-              )
-            )
-            .as(
-              DisableEnhancedMonitoringResponse(
-                current,
-                desired,
-                streamName
-              )
+          (
+            streams.updateStream(
+              stream
+                .copy(enhancedMonitoring = Vector(ShardLevelMetrics(desired)))
+            ),
+            DisableEnhancedMonitoringResponse(
+              current,
+              desired,
+              streamName
             )
+          )
         }
+        .sequenceWithDefault(streams)
     }
 }