pjfanning commented on code in PR #489: URL: https://github.com/apache/pekko-persistence-r2dbc/pull/489#discussion_r3992438979
########## core/src/main/scala/org/apache/pekko/persistence/r2dbc/journal/R2dbcBatchJournal.scala: ########## @@ -0,0 +1,349 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * license agreements; and to You under the Apache License, version 2.0: + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * This file is part of the Apache Pekko project, which was derived from Akka. + */ + +package org.apache.pekko.persistence.r2dbc.journal + +import java.time.Instant +import scala.collection.immutable +import scala.concurrent.{ ExecutionContext, Future, Promise } +import scala.util.{ Failure, Success, Try } +import com.typesafe.config.Config +import io.r2dbc.spi.R2dbcDataIntegrityViolationException +import org.apache.pekko +import org.apache.pekko.persistence.r2dbc.Dialect.{ Postgres, Yugabyte } +import pekko.Done +import pekko.actor.{ ActorRef, Timers } +import pekko.actor.typed.ActorSystem +import pekko.actor.typed.scaladsl.adapter._ +import pekko.annotation.InternalApi +import pekko.event.Logging +import pekko.persistence.AtomicWrite +import pekko.persistence.Persistence +import pekko.persistence.PersistentRepr +import pekko.persistence.journal.AsyncWriteJournal +import pekko.persistence.journal.Tagged +import pekko.persistence.r2dbc.JournalSettings +import pekko.persistence.r2dbc.internal.InstantFactory +import pekko.persistence.r2dbc.internal.PubSub +import pekko.persistence.r2dbc.journal.JournalDao.SerializedEventMetadata +import pekko.persistence.r2dbc.journal.JournalDao.SerializedJournalRow +import pekko.persistence.typed.PersistenceId +import pekko.serialization.Serialization +import pekko.serialization.SerializationExtension +import pekko.serialization.Serializers +import pekko.stream.scaladsl.Sink + +import scala.concurrent.duration.FiniteDuration +import scala.jdk.DurationConverters.JavaDurationOps + +/** + * INTERNAL API + */ +@InternalApi +private[r2dbc] object R2dbcBatchJournal { + case class WriteFinished(persistenceId: String, done: Future[?]) + private case object Flush + private case object FlushDone + + private final case class WriteRequest( + persistenceId: String, + rows: Seq[SerializedJournalRow], + messages: Seq[AtomicWrite], + promise: Promise[Seq[Try[Unit]]] + ) + + def deserializeRow(serialization: Serialization, row: SerializedJournalRow): PersistentRepr = { + if (row.payload.isEmpty) + throw new IllegalStateException("Expected event payload to be loaded.") + val payload = serialization.deserialize(row.payload.get, row.serId, row.serManifest).get + val repr = PersistentRepr( + payload, + row.seqNr, + row.persistenceId, + writerUuid = row.writerUuid, + manifest = "", // FIXME issue #84 + deleted = false, + sender = ActorRef.noSender) + + val reprWithMeta = row.metadata match { + case None => repr + case Some(meta) => + repr.withMetadata(serialization.deserialize(meta.payload, meta.serId, meta.serManifest).get) + } + reprWithMeta + } + +} + +/** + * INTERNAL API + * + * Opt-in journal plugin (`pekko.persistence.r2dbc.batched-journal`) that coalesces concurrent + * writes from different persistence ids into one multi-row statement, trading up to + * `max-batch-time` of write latency for higher throughput at high concurrency. + * + * Mixing persistence ids in a single statement is only safe because the plugin requires + * `use-app-timestamp = on` and `db-timestamp-monotonic-increasing = on`: in that mode + * [[JournalDao]] does not bind the per-persistence-id previous sequence number subselect, and + * timestamps come from the application clock, which therefore must not move backwards. + * The Postgres and Yugabyte dialects are required because the flush relies on `RETURNING`. + * + * Writes are buffered in a bounded queue (`max-queue-size`); incoming writes are rejected with + * a failed future once the queue is full. Flushed batches are serialized: only one batch is in + * flight at a time, which keeps same-persistence-id writes committed in order without relying on + * replay-time coordination, at the cost of not using spare pool capacity. Concurrent flushing can + * be added later if a single flush saturates. + * + * A batch that fails with a database integrity violation is retried in halves so that only the + * offending persistence ids fail. Infrastructure errors fail the whole batch. A batch can contain + * an arbitrarily large number of rows when callers use `persistAll` or `persistAsync` bursts, the + * same as the default journal; this plugin targets many small concurrent writes. + */ +@InternalApi +private[r2dbc] final class R2dbcBatchJournal(config: Config) extends AsyncWriteJournal with Timers { + import R2dbcBatchJournal.WriteFinished + import R2dbcBatchJournal.deserializeRow + import R2dbcBatchJournal.Flush + import R2dbcBatchJournal.FlushDone + import R2dbcBatchJournal.WriteRequest + + implicit val system: ActorSystem[?] = context.system.toTyped + implicit val ec: ExecutionContext = context.dispatcher + + private val log = Logging(context.system, classOf[R2dbcBatchJournal]) + + private val persistenceExt = Persistence(system) + + private val serialization: Serialization = SerializationExtension(context.system) + private val journalSettings = JournalSettings(config) + + require(journalSettings.dialect == Postgres || journalSettings.dialect == Yugabyte, + "Batching is only supported for Postgres and Yugabyte") + require(journalSettings.useAppTimestamp, "use-app-timestamp must be 'on' when using R2dbcBatchJournal") + require(journalSettings.dbTimestampMonotonicIncreasing, + "db-timestamp-monotonic-increasing must be 'on' when using R2dbcBatchJournal") + + private val maxQueueSize: Int = config.getInt("max-queue-size") + private val maxBatchSize: Int = config.getInt("max-batch-size") + private val maxBatchTime: FiniteDuration = config.getDuration("max-batch-time").toScala + + require(maxQueueSize > 0, "max-queue-size must be at least 1 when using R2dbcBatchJournal") + require(maxBatchSize > 0, "max-batch-size must be at least 1 when using R2dbcBatchJournal") + require(maxBatchSize <= maxQueueSize, "max-batch-size must be less than or equal to `max-queue-size`") + + private val journalDao = JournalDao.fromConfig(journalSettings, config) + + private val pubSub: Option[PubSub] = + if (journalSettings.journalPublishEvents) Some(PubSub(system)) + else None + + // if there are pending writes when an actor restarts we must wait for + // them to complete before we can read the highest sequence number, or we will miss it + private val writesInProgress = new java.util.HashMap[String, Future[?]]() + + private val queue = collection.mutable.ArrayDeque[WriteRequest]() + private var noActiveWrite = true + + // set in postStop; failing queued promises completes their futures, whose callbacks must + // not send WriteFinished to an actor that is already terminating + @volatile private var stopping = false + + private def doFlush(): Unit = { + + val count = math.min(maxBatchSize, queue.size.toLong).toInt + val writeRequests = new Array[WriteRequest](count) + + queue.copyToArray(writeRequests, 0, count) + queue.dropInPlace(count) + + def write(requests: Array[WriteRequest]): Future[Unit] = + journalDao + .writeEvents(immutable.ArraySeq.unsafeWrapArray(requests.view.flatMap(_.rows).toArray)) + .map { _ => + requests.foreach(_.promise.trySuccess(Nil)) + requests.foreach(w => publish(w.messages, Future.successful(w.rows.head.dbTimestamp))) + } + .recoverWith { + case _: R2dbcDataIntegrityViolationException if requests.length > 1 => + val (left, right) = requests.splitAt(requests.length / 2) + write(left).flatMap(_ => write(right)) + case exception => + requests.foreach(_.promise.tryFailure(exception)) + Future.unit + } + + write(writeRequests).onComplete(_ => if (!stopping) self ! FlushDone) + } + + override def receivePluginInternal: Receive = { + case WriteFinished(pid, f) => writesInProgress.remove(pid, f) + case Flush => + if (noActiveWrite && queue.nonEmpty) { + noActiveWrite = false + doFlush() + } + case FlushDone => + noActiveWrite = true + if (queue.size >= maxBatchSize) { + noActiveWrite = false + doFlush() + } else if (queue.nonEmpty && !timers.isTimerActive(Flush)) { + timers.startSingleTimer(Flush, Flush, maxBatchTime) + } + } + + override def asyncWriteMessages(messages: immutable.Seq[AtomicWrite]): Future[immutable.Seq[Try[Unit]]] = { + if (queue.length >= maxQueueSize) + Promise + .failed(new IllegalStateException(s"Unable to accept the request, max-queue-size [$maxQueueSize] reached")) + .future + else { + val promise = Promise[immutable.Seq[Try[Unit]]]() + + def atomicWrite(atomicWrite: AtomicWrite): Unit = { Review Comment: I don't like the code here that seems to collapse a Seq[Unit] to a Unit - I'd prefer to have Seq[Int] collapsed to an summed Int or a Seq[Boolean] collapsed to a Boolean - or something that allows the compiler to check if all the work is completed -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
