-
Notifications
You must be signed in to change notification settings - Fork 8
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add suspend and resume functionality
1 parent
4b9593f
commit 441fd83
Showing
19 changed files
with
890 additions
and
107 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
146 changes: 137 additions & 9 deletions
146
backup-s3/src/main/scala/io/aiven/guardian/kafka/backup/s3/BackupClient.scala
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,40 +1,168 @@ | ||
package io.aiven.guardian.kafka.backup.s3 | ||
|
||
import akka.Done | ||
import akka.NotUsed | ||
import akka.actor.ActorSystem | ||
import akka.stream.alpakka.s3.FailedUploadPart | ||
import akka.stream.alpakka.s3.MultipartUploadResult | ||
import akka.stream.alpakka.s3.Part | ||
import akka.stream.alpakka.s3.S3Attributes | ||
import akka.stream.alpakka.s3.S3Headers | ||
import akka.stream.alpakka.s3.S3Settings | ||
import akka.stream.alpakka.s3.SuccessfulUploadPart | ||
import akka.stream.alpakka.s3.UploadPartResponse | ||
import akka.stream.alpakka.s3.scaladsl.S3 | ||
import akka.stream.scaladsl._ | ||
import akka.util.ByteString | ||
import com.typesafe.scalalogging.StrictLogging | ||
import io.aiven.guardian.kafka.KafkaClientInterface | ||
import io.aiven.guardian.kafka.backup.BackupClientInterface | ||
import io.aiven.guardian.kafka.backup.configs.Backup | ||
import io.aiven.guardian.kafka.s3.configs.{S3 => S3Config} | ||
|
||
import scala.collection.immutable | ||
import scala.concurrent.ExecutionContext | ||
import scala.concurrent.Future | ||
|
||
import java.time.Instant | ||
|
||
final case class CurrentS3State(uploadId: String, parts: Seq[Part]) | ||
|
||
class BackupClient[T <: KafkaClientInterface](maybeS3Settings: Option[S3Settings])(implicit | ||
override val kafkaClientInterface: T, | ||
override val backupConfig: Backup, | ||
override val system: ActorSystem, | ||
s3Config: S3Config, | ||
s3Headers: S3Headers | ||
) extends BackupClientInterface[T] { | ||
) extends BackupClientInterface[T] | ||
with StrictLogging { | ||
|
||
override def empty: () => Future[Option[MultipartUploadResult]] = () => Future.successful(None) | ||
|
||
override type BackupResult = Option[MultipartUploadResult] | ||
|
||
override def backupToStorageSink(key: String): Sink[ByteString, Future[BackupResult]] = { | ||
val base = S3 | ||
.multipartUploadWithHeaders( | ||
s3Config.dataBucket, | ||
key, | ||
s3Headers = s3Headers, | ||
chunkingParallelism = 1 // Parallelism is pointless for this usecase since we are directly streaming from Kafka | ||
override type CurrentState = CurrentS3State | ||
|
||
override def getCurrentUploadState(key: String): Future[Option[CurrentS3State]] = { | ||
implicit val ec: ExecutionContext = system.classicSystem.getDispatcher | ||
|
||
val baseListMultipart = S3.listMultipartUpload(s3Config.dataBucket, None) | ||
|
||
for { | ||
incompleteUploads <- | ||
maybeS3Settings | ||
.fold(baseListMultipart)(s3Settings => baseListMultipart.withAttributes(S3Attributes.settings(s3Settings))) | ||
.runWith(Sink.seq) | ||
keys = incompleteUploads.filter(_.key == key) | ||
result <- if (keys.isEmpty) | ||
Future.successful(None) | ||
else { | ||
val listMultipartUploads = keys match { | ||
case Seq(single) => | ||
logger.info( | ||
s"Found previous uploadId: ${single.uploadId} and bucket: ${s3Config.dataBucket} with key: ${single.key}" | ||
) | ||
single | ||
case rest => | ||
val last = rest.minBy(_.initiated)(Ordering[Instant].reverse) | ||
logger.warn( | ||
s"Found multiple previously cancelled uploads for key: $key and bucket: ${s3Config.dataBucket}, picking uploadId: ${last.uploadId}" | ||
) | ||
last | ||
} | ||
val uploadId = listMultipartUploads.uploadId | ||
val baseList = S3.listParts(s3Config.dataBucket, key, listMultipartUploads.uploadId) | ||
|
||
for { | ||
parts <- maybeS3Settings | ||
.fold(baseList)(s3Settings => baseList.withAttributes(S3Attributes.settings(s3Settings))) | ||
.runWith(Sink.seq) | ||
|
||
finalParts = parts.lastOption match { | ||
case Some(part) if part.size >= akka.stream.alpakka.s3.scaladsl.S3.MinChunkSize => | ||
parts | ||
case _ => | ||
// We drop the last part here since its broken | ||
parts.dropRight(1) | ||
} | ||
} yield Some(CurrentS3State(uploadId, finalParts.map(_.toPart))) | ||
} | ||
} yield result | ||
|
||
} | ||
|
||
private[s3] def failureSink | ||
: Sink[(FailedUploadPart, immutable.Iterable[kafkaClientInterface.CursorContext]), Future[Done]] = Sink | ||
.foreach[(FailedUploadPart, immutable.Iterable[kafkaClientInterface.CursorContext])] { case (failedPart, _) => | ||
logger.warn( | ||
s"Failed to upload a chunk into S3 with bucket: ${failedPart.multipartUpload.bucket}, key: ${failedPart.multipartUpload.key}, uploadId: ${failedPart.multipartUpload.uploadId} and partNumber: ${failedPart.partNumber}", | ||
failedPart.exception | ||
) | ||
.mapMaterializedValue(future => future.map(result => Some(result))(ExecutionContext.parasitic)) | ||
} | ||
|
||
private[s3] def successSink | ||
: Sink[(SuccessfulUploadPart, immutable.Iterable[kafkaClientInterface.CursorContext]), Future[Done]] = | ||
kafkaClientInterface.commitCursor | ||
.contramap[(SuccessfulUploadPart, immutable.Iterable[kafkaClientInterface.CursorContext])] { case (_, cursors) => | ||
kafkaClientInterface.batchCursorContext(cursors) | ||
} | ||
|
||
private[s3] def kafkaBatchSink | ||
: Sink[(UploadPartResponse, immutable.Iterable[kafkaClientInterface.CursorContext]), NotUsed] = { | ||
|
||
val success = Flow[(UploadPartResponse, immutable.Iterable[kafkaClientInterface.CursorContext])] | ||
.collectType[(SuccessfulUploadPart, immutable.Iterable[kafkaClientInterface.CursorContext])] | ||
.wireTap { data => | ||
val (part, _) = data | ||
logger.info( | ||
s"Committing kafka cursor for uploadId:${part.multipartUpload.uploadId} key: ${part.multipartUpload.key} and S3 part: ${part.partNumber}" | ||
) | ||
} | ||
.toMat(successSink)(Keep.none) | ||
|
||
val failure = Flow[(UploadPartResponse, immutable.Iterable[kafkaClientInterface.CursorContext])] | ||
.collectType[(FailedUploadPart, immutable.Iterable[kafkaClientInterface.CursorContext])] | ||
.toMat(failureSink)(Keep.none) | ||
|
||
Sink.combine(success, failure)(Broadcast(_)) | ||
} | ||
|
||
override def backupToStorageSink(key: String, | ||
currentState: Option[CurrentS3State] | ||
): Sink[(ByteString, kafkaClientInterface.CursorContext), Future[BackupResult]] = { | ||
|
||
// Note that chunkingParallelism is pointless for this usecase since we are directly streaming from Kafka. | ||
// Furthermore the real `KafkaClient` implementation uses `CommittableOffsetBatch` which is a global singleton so | ||
// we can't have concurrent updates to this data structure. | ||
|
||
val sink = currentState match { | ||
case Some(state) => | ||
logger.info( | ||
s"Resuming previous upload with uploadId: ${state.uploadId} and bucket: ${s3Config.dataBucket} with key: $key" | ||
) | ||
S3.resumeMultipartUploadWithHeadersAndContext[kafkaClientInterface.CursorContext]( | ||
s3Config.dataBucket, | ||
key, | ||
state.uploadId, | ||
state.parts, | ||
kafkaBatchSink, | ||
s3Headers = s3Headers, | ||
chunkingParallelism = 1 | ||
) | ||
case None => | ||
logger.info( | ||
s"Creating new upload with bucket: ${s3Config.dataBucket} and key: $key" | ||
) | ||
S3.multipartUploadWithHeadersAndContext[kafkaClientInterface.CursorContext]( | ||
s3Config.dataBucket, | ||
key, | ||
kafkaBatchSink, | ||
s3Headers = s3Headers, | ||
chunkingParallelism = 1 | ||
) | ||
} | ||
|
||
val base = sink.mapMaterializedValue(future => future.map(result => Some(result))(ExecutionContext.parasitic)) | ||
maybeS3Settings.fold(base)(s3Settings => base.withAttributes(S3Attributes.settings(s3Settings))) | ||
} | ||
} |
33 changes: 33 additions & 0 deletions
33
backup-s3/src/test/scala/io/aiven/guardian/kafka/backup/s3/BackupClientChunkState.scala
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,33 @@ | ||
package io.aiven.guardian.kafka.backup.s3 | ||
|
||
import akka.Done | ||
import akka.actor.ActorSystem | ||
import akka.stream.alpakka.s3.S3Headers | ||
import akka.stream.alpakka.s3.S3Settings | ||
import akka.stream.alpakka.s3.SuccessfulUploadPart | ||
import akka.stream.scaladsl.Sink | ||
import io.aiven.guardian.kafka.KafkaClientInterface | ||
import io.aiven.guardian.kafka.backup.configs.Backup | ||
import io.aiven.guardian.kafka.s3.configs.{S3 => S3Config} | ||
|
||
import scala.collection.immutable | ||
import scala.concurrent.Future | ||
|
||
import java.util.concurrent.ConcurrentLinkedQueue | ||
|
||
class BackupClientChunkState[T <: KafkaClientInterface](maybeS3Settings: Option[S3Settings])(implicit | ||
override val kafkaClientInterface: T, | ||
override val backupConfig: Backup, | ||
override val system: ActorSystem, | ||
s3Config: S3Config, | ||
s3Headers: S3Headers | ||
) extends BackupClient[T](maybeS3Settings) { | ||
val processedChunks: ConcurrentLinkedQueue[SuccessfulUploadPart] = new ConcurrentLinkedQueue[SuccessfulUploadPart]() | ||
|
||
override val successSink | ||
: Sink[(SuccessfulUploadPart, immutable.Iterable[kafkaClientInterface.CursorContext]), Future[Done]] = | ||
super.successSink.contramap { case (part, value) => | ||
processedChunks.add(part) | ||
(part, value) | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
27 changes: 27 additions & 0 deletions
27
backup-s3/src/test/scala/io/aiven/guardian/kafka/backup/s3/KafkaClientWithKillSwitch.scala
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,27 @@ | ||
package io.aiven.guardian.kafka.backup.s3 | ||
|
||
import akka.actor.ActorSystem | ||
import akka.kafka.CommitterSettings | ||
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( | ||
configureConsumer: Option[ | ||
ConsumerSettings[Array[Byte], Array[Byte]] => ConsumerSettings[Array[Byte], Array[Byte]] | ||
] = None, | ||
configureCommitter: Option[ | ||
CommitterSettings => CommitterSettings | ||
] = None, | ||
killSwitch: SharedKillSwitch | ||
)(implicit system: ActorSystem, kafkaClusterConfig: KafkaCluster) | ||
extends KafkaClient(configureConsumer, configureCommitter) { | ||
override def getSource | ||
: SourceWithContext[ReducedConsumerRecord, ConsumerMessage.CommittableOffset, Consumer.Control] = | ||
super.getSource.via(killSwitch.flow) | ||
} |
11 changes: 6 additions & 5 deletions
11
...p-s3/src/test/scala/io/aiven/guardian/kafka/backup/s3/MockedS3BackupClientInterface.scala
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,24 +1,25 @@ | ||
package io.aiven.guardian.kafka.backup.s3 | ||
|
||
import akka.NotUsed | ||
import akka.actor.ActorSystem | ||
import akka.stream.alpakka.s3.S3Headers | ||
import akka.stream.alpakka.s3.S3Settings | ||
import akka.stream.scaladsl.Source | ||
import io.aiven.guardian.kafka.MockedKafkaClientInterface | ||
import io.aiven.guardian.kafka.backup.configs.Backup | ||
import io.aiven.guardian.kafka.backup.configs.TimeConfiguration | ||
import io.aiven.guardian.kafka.models.ReducedConsumerRecord | ||
import io.aiven.guardian.kafka.s3.configs.{S3 => S3Config} | ||
|
||
import scala.concurrent.duration.FiniteDuration | ||
|
||
class MockedS3BackupClientInterface( | ||
kafkaData: Source[ReducedConsumerRecord, NotUsed], | ||
periodSlice: FiniteDuration, | ||
timeConfiguration: TimeConfiguration, | ||
s3Config: S3Config, | ||
maybeS3Settings: Option[S3Settings] | ||
)(implicit val s3Headers: S3Headers) | ||
)(implicit val s3Headers: S3Headers, system: ActorSystem) | ||
extends BackupClient(maybeS3Settings)(new MockedKafkaClientInterface(kafkaData), | ||
Backup(periodSlice), | ||
Backup(timeConfiguration), | ||
implicitly, | ||
s3Config, | ||
implicitly | ||
) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
8 changes: 3 additions & 5 deletions
8
core-backup/src/main/scala/io/aiven/guardian/kafka/backup/configs/Backup.scala
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,8 +1,6 @@ | ||
package io.aiven.guardian.kafka.backup.configs | ||
|
||
import scala.concurrent.duration.FiniteDuration | ||
|
||
/** @param periodSlice | ||
* The time period for each given slice that stores all of the `ReducedConsumerRecord` | ||
/** @param timeConfiguration | ||
* Determines how the backed up objects/files are segregated depending on a time configuration | ||
*/ | ||
final case class Backup(periodSlice: FiniteDuration) | ||
final case class Backup(timeConfiguration: TimeConfiguration) |
23 changes: 23 additions & 0 deletions
23
core-backup/src/main/scala/io/aiven/guardian/kafka/backup/configs/TimeConfiguration.scala
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,23 @@ | ||
package io.aiven.guardian.kafka.backup.configs | ||
|
||
import scala.concurrent.duration.Duration | ||
|
||
import java.time.temporal.ChronoUnit | ||
|
||
sealed trait TimeConfiguration | ||
|
||
/** Backs up objects/files depending on the timestamp fo the first received Kafka message. Suspending/resuming the | ||
* backup client will always create a new object/file | ||
* @param duration | ||
* The maximum span of time for each object/file, when this duration is exceeded a new file is created | ||
*/ | ||
final case class PeriodFromFirst(duration: Duration) extends TimeConfiguration | ||
|
||
/** Backs up objects/files by collecting received Kafka messages into a single time slice based on a [[ChronoUnit]]. | ||
* When suspending/resuming the backup client, this option will reuse existing objects/keys if they fall into the | ||
* currently configured `chronoUnit`. | ||
* @param chronoUnit | ||
* Timestamps for kafka messages that are contained within the configured [[ChronoUnit]] will be placed into the same | ||
* object/file. | ||
*/ | ||
final case class ChronoUnitSlice(chronoUnit: ChronoUnit) extends TimeConfiguration |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters