Use killswitch's in RealS3BackupClientSpec and properly terminate JSON

Author: mdedetrich

backup-s3/src/test/scala/io/aiven/guardian/kafka/backup/s3/BackupClientSpec.scala

@@ -49,7 +49,9 @@ trait BackupClientSpec
     with StrictLogging {
 
   implicit val ec: ExecutionContext            = system.dispatcher
-  implicit val defaultPatience: PatienceConfig = PatienceConfig(90 seconds, 100 millis)
+  implicit val defaultPatience: PatienceConfig = PatienceConfig(5 minutes, 100 millis)
+  implicit override val generatorDrivenConfig: PropertyCheckConfiguration =
+    PropertyCheckConfiguration(minSuccessful = 1, minSize = 1)
 
   val ThrottleElements: Int          = 100
   val ThrottleAmount: FiniteDuration = 1 millis
@@ -172,7 +174,7 @@ trait BackupClientSpec
             })
           keysWithRecords <- Future.sequence(keysWithSource.map { case (key, source) =>
                                source
-                                 .via(CirceStreamSupport.decode[List[ReducedConsumerRecord]])
+                                 .via(CirceStreamSupport.decode[List[Option[ReducedConsumerRecord]]])
                                  .toMat(Sink.collection)(Keep.right)
                                  .run()
                                  .map(list => (key, list.flatten))
@@ -182,7 +184,9 @@ trait BackupClientSpec
                      OffsetDateTime.parse(date).toEpochSecond
                    }(Ordering[Long].reverse)
           flattened = sorted.flatMap { case (_, records) => records }
-        } yield flattened
+        } yield flattened.collect { case Some(reducedConsumerRecord) =>
+          reducedConsumerRecord
+        }
         val observed = calculatedFuture.futureValue
 
         kafkaDataWithTimePeriod.data.containsSlice(observed) mustEqual true

backup-s3/src/test/scala/io/aiven/guardian/kafka/backup/s3/KafkaClientWithKillSwitch.scala

@@ -0,0 +1,21 @@
+package io.aiven.guardian.kafka.backup.s3
+
+import akka.actor.ActorSystem
+import akka.kafka.ConsumerMessage
+import akka.kafka.ConsumerSettings
+import akka.kafka.scaladsl.Consumer
+import akka.stream.SharedKillSwitch
+import akka.stream.scaladsl.SourceWithContext
+import io.aiven.guardian.kafka.KafkaClient
+import io.aiven.guardian.kafka.configs.KafkaCluster
+import io.aiven.guardian.kafka.models.ReducedConsumerRecord
+
+class KafkaClientWithKillSwitch(
+    configure: Option[ConsumerSettings[Array[Byte], Array[Byte]] => ConsumerSettings[Array[Byte], Array[Byte]]] = None,
+    killSwitch: SharedKillSwitch
+)(implicit system: ActorSystem, kafkaClusterConfig: KafkaCluster)
+    extends KafkaClient(configure) {
+  override def getSource
+      : SourceWithContext[ReducedConsumerRecord, ConsumerMessage.CommittableOffset, Consumer.Control] =
+    super.getSource.via(killSwitch.flow)
+}

backup-s3/src/test/scala/io/aiven/guardian/kafka/backup/s3/RealS3BackupClientSpec.scala

@@ -3,14 +3,16 @@ package io.aiven.guardian.kafka.backup.s3
 import akka.actor.ActorSystem
 import akka.kafka.ProducerSettings
 import akka.kafka.scaladsl.Producer
+import akka.stream.KillSwitches
+import akka.stream.SharedKillSwitch
 import akka.stream.alpakka.s3.S3Settings
 import akka.stream.alpakka.s3.scaladsl.S3
 import akka.stream.scaladsl.Sink
 import akka.stream.scaladsl.Source
+import akka.util.ByteString
 import io.aiven.guardian.akka.AnyPropTestKit
 import io.aiven.guardian.kafka.Generators.KafkaDataInChunksWithTimePeriod
 import io.aiven.guardian.kafka.Generators.kafkaDataWithMinSizeGen
-import io.aiven.guardian.kafka.KafkaClient
 import io.aiven.guardian.kafka.KafkaClusterTest
 import io.aiven.guardian.kafka.Utils._
 import io.aiven.guardian.kafka.backup.configs.Backup
@@ -26,6 +28,7 @@ import org.apache.kafka.clients.producer.ProducerRecord
 import org.apache.kafka.common.serialization.ByteArraySerializer
 import org.mdedetrich.akka.stream.support.CirceStreamSupport
 
+import scala.concurrent.Future
 import scala.concurrent.duration._
 import scala.jdk.CollectionConverters._
 import scala.jdk.FutureConverters._
@@ -47,12 +50,26 @@ class RealS3BackupClientSpec
   override lazy val bucketPrefix: Option[String]          = Some("guardian-")
   override lazy val enableCleanup: Option[FiniteDuration] = Some(5 seconds)
 
+  case object TerminationException extends Exception("termination-exception")
+
   def reducedConsumerRecordsToJson(reducedConsumerRecords: List[ReducedConsumerRecord]): Array[Byte] = {
     import io.aiven.guardian.kafka.codecs.Circe._
     import io.circe.syntax._
     reducedConsumerRecords.asJson.noSpaces.getBytes
   }
 
+  def createKafkaClient(
+      killSwitch: SharedKillSwitch
+  )(implicit kafkaClusterConfig: KafkaCluster): KafkaClientWithKillSwitch =
+    new KafkaClientWithKillSwitch(
+      configure = Some(
+        _.withBootstrapServers(
+          container.bootstrapServers
+        ).withGroupId("test-group")
+      ),
+      killSwitch
+    )
+
   property("suspend/resume works correctly") {
     forAll(kafkaDataWithMinSizeGen(S3.MinChunkSize, 5, reducedConsumerRecordsToJson),
            s3ConfigGen(useVirtualDotHost, bucketPrefix)
@@ -72,18 +89,15 @@ class RealS3BackupClientSpec
         ProducerSettings(system, new ByteArraySerializer, new ByteArraySerializer)
           .withBootstrapServers(container.bootstrapServers)
 
-      val backupClient = new BackupClient(Some(s3Settings))(new KafkaClient(
-                                                              configure = Some(
-                                                                _.withBootstrapServers(
-                                                                  container.bootstrapServers
-                                                                ).withGroupId("test-group")
-                                                              )
-                                                            ),
-                                                            implicitly,
-                                                            implicitly,
-                                                            implicitly,
-                                                            implicitly
-      )
+      val firstKillSwitch = KillSwitches.shared("first-kill-switch")
+
+      val backupClient =
+        new BackupClient(Some(s3Settings))(createKafkaClient(firstKillSwitch),
+                                           implicitly,
+                                           implicitly,
+                                           implicitly,
+                                           implicitly
+        )
 
       val asProducerRecord = data.map { reducedConsumerRecord =>
         val keyAsByteArray   = Base64.getDecoder.decode(reducedConsumerRecord.key)
@@ -110,45 +124,57 @@ class RealS3BackupClientSpec
       val calculatedFuture = for {
         _ <- createTopics.all().toCompletableFuture.asScala
         _ <- createBucket(s3Config.dataBucket)
-        firstControl = backupClient.backup.run()
-        _            = baseSource.runWith(Producer.plainSink(producerSettings))
-        _ <- akka.pattern.after(15 seconds)(firstControl.stop())
-        secondBackupClient = new BackupClient(Some(s3Settings))(new KafkaClient(
-                                                                  configure = Some(
-                                                                    _.withBootstrapServers(
-                                                                      container.bootstrapServers
-                                                                    ).withGroupId("test-group")
-                                                                  )
-                                                                ),
-                                                                implicitly,
-                                                                implicitly,
-                                                                implicitly,
-                                                                implicitly
-                             )
-        secondControl = secondBackupClient.backup.run()
-        _ <- akka.pattern.after(15 seconds)(secondControl.stop())
+        _ = backupClient.backup.run()
+        _ = baseSource.runWith(Producer.plainSink(producerSettings))
+        _ <- akka.pattern.after(15 seconds)(Future {
+               firstKillSwitch.abort(TerminationException)
+             })
+        secondKillSwitch = KillSwitches.shared("second-kill-switch")
+        secondBackupClient =
+          new BackupClient(Some(s3Settings))(createKafkaClient(secondKillSwitch),
+                                             implicitly,
+                                             implicitly,
+                                             implicitly,
+                                             implicitly
+          )
+        _ <- akka.pattern.after(10 seconds)(Future(secondBackupClient.backup.run()))
+        _ <- akka.pattern.after(30 seconds)(Future {
+               secondKillSwitch.abort(TerminationException)
+             })
         // Need to manually merge the bucket since we haven't gone over time window
-        incomplete <- S3.listMultipartUpload(s3Config.dataBucket, None).withAttributes(s3Attrs).runWith(Sink.seq)
+        incomplete <- akka.pattern.after(10 seconds)(
+                        S3.listMultipartUpload(s3Config.dataBucket, None).withAttributes(s3Attrs).runWith(Sink.seq)
+                      )
 
         grouped = incomplete.groupBy(_.key)
 
         (key, uploadIds) = grouped.head
-        uploadId         = uploadIds.head
+        upload           = uploadIds.head
 
-        parts <- S3.listParts(s3Config.dataBucket, key, uploadId.uploadId).runWith(Sink.seq)
+        parts <- S3.listParts(s3Config.dataBucket, key, upload.uploadId).runWith(Sink.seq)
 
-        _ <- S3.completeMultipartUpload(s3Config.dataBucket, key, uploadId.uploadId, parts.map(_.toPart))
+        _ <- Source
+               .single(ByteString("null]"))
+               .runWith(
+                 S3.resumeMultipartUpload(s3Config.dataBucket, key, upload.uploadId, parts.map(_.toPart))
+                   .withAttributes(s3Attrs)
+               )
 
         downloaded <-
-          akka.pattern.after(5.seconds)(
-            S3.download(s3Config.dataBucket, key).withAttributes(s3Attrs).runWith(Sink.head).flatMap {
-              case Some((downloadSource, _)) =>
-                downloadSource.via(CirceStreamSupport.decode[List[ReducedConsumerRecord]]).runWith(Sink.seq)
-              case None => throw new Exception(s"Expected object in bucket ${s3Config.dataBucket} with key $key")
-            }
+          akka.pattern.after(10 seconds)(
+            S3.download(s3Config.dataBucket, key)
+              .withAttributes(s3Attrs)
+              .runWith(Sink.head)
+              .flatMap {
+                case Some((downloadSource, _)) =>
+                  downloadSource.via(CirceStreamSupport.decode[List[Option[ReducedConsumerRecord]]]).runWith(Sink.seq)
+                case None => throw new Exception(s"Expected object in bucket ${s3Config.dataBucket} with key $key")
+              }
           )
 
-      } yield downloaded.flatten
+      } yield downloaded.toList.flatten.collect { case Some(reducedConsumerRecord) =>
+        reducedConsumerRecord
+      }
 
       val downloaded = calculatedFuture.futureValue
 
@@ -161,7 +187,7 @@ class RealS3BackupClientSpec
         }
         .toMap
 
-      val inputAsKey = kafkaDataInChunksWithTimePeriod.data.flatten
+      val inputAsKey = data
         .groupBy(_.key)
         .view
         .mapValues { reducedConsumerRecords =>

core/src/main/scala/io/aiven/guardian/kafka/KafkaClient.scala

@@ -53,7 +53,7 @@ class KafkaClient(
   /** @return
     *   A `SourceWithContext` that returns a Kafka Stream which automatically handles committing of cursors
     */
-  override val getSource: SourceWithContext[ReducedConsumerRecord, CommittableOffset, Consumer.Control] =
+  override def getSource: SourceWithContext[ReducedConsumerRecord, CommittableOffset, Consumer.Control] =
     Consumer
       .sourceWithOffsetContext(consumerSettings, subscriptions)
       .map(consumerRecord =>
@@ -72,7 +72,7 @@ class KafkaClient(
   /** @return
     *   A `Sink` that allows you to commit a `CursorContext` to Kafka to signify you have processed a message
     */
-  override val commitCursor: Sink[CommittableOffsetBatch, Future[Done]] = Committer.sink(committerSettings)
+  override def commitCursor: Sink[CommittableOffsetBatch, Future[Done]] = Committer.sink(committerSettings)
 
   /** How to batch an immutable iterable of `CursorContext` into a `BatchedCursorContext`
     * @param cursors