Pubsub source v2 cancels streaming pull when blocked

istreeter · istreeter · commit 4cde5376dccd · 2024-10-10T00:42:01.000+01:00
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -7,7 +7,7 @@ on:
 
 jobs:
   test:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-22.04
     steps:
       - uses: actions/checkout@v2
       - uses: coursier/cache-action@v6
@@ -22,7 +22,7 @@ jobs:
 
   publish_docker:
     needs: test
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-22.04
     strategy:
       matrix:
         sbtProject:
diff --git a/modules/gcp/src/main/resources/application.conf b/modules/gcp/src/main/resources/application.conf
@@ -14,6 +14,9 @@
     // V2 defaults
     "durationPerAckExtension": "10 minutes"
     "minRemainingDeadline": 0.1
+    "progressTimeout": "10 seconds"
+    "modackOnProgressTimeout": true
+    "cancelOnProgressTimeout": false
   }
   "output": {
     "bad": ${snowplow.defaults.sinks.pubsub}
diff --git a/modules/gcp/src/main/scala/common-streams-extensions/v2/PubsubCheckpointer.scala b/modules/gcp/src/main/scala/common-streams-extensions/v2/PubsubCheckpointer.scala
@@ -88,16 +88,9 @@ class PubsubCheckpointer[F[_]: Async](
       ackDatas <- refAckIds.modify(m => (m.removedAll(c), c.flatMap(m.get)))
       grouped = ackDatas.groupBy(_.channelAffinity)
       _ <- grouped.toVector.parTraverse_ { case (channelAffinity, ackDatas) =>
-             ackDatas.flatMap(_.ackIds).grouped(1000).toVector.traverse_ { ackIds =>
-               // A nack is just a modack with zero duration
-               Utils
-                 .modAck[F](subscription, stub, ackIds, Duration.Zero, channelAffinity)
-                 .retryingOnTransientGrpcFailures
-                 .recoveringOnGrpcInvalidArgument { s =>
-                   // This can happen if ack IDs have expired before we acked
-                   Logger[F].info(s"Ignoring error from GRPC when acking: ${s.getDescription}")
-                 }
-             }
+             val ackIds = ackDatas.flatMap(_.ackIds)
+             // A nack is just a modack with zero duration
+             Utils.modAck[F](subscription, stub, ackIds, Duration.Zero, channelAffinity)
            }
     } yield ()
 }
diff --git a/modules/gcp/src/main/scala/common-streams-extensions/v2/PubsubSourceConfigV2.scala b/modules/gcp/src/main/scala/common-streams-extensions/v2/PubsubSourceConfigV2.scala
@@ -23,7 +23,10 @@ case class PubsubSourceConfigV2(
   durationPerAckExtension: FiniteDuration,
   minRemainingDeadline: Double,
   gcpUserAgent: GcpUserAgent,
-  maxPullsPerTransportChannel: Int
+  maxPullsPerTransportChannel: Int,
+  progressTimeout: FiniteDuration,
+  modackOnProgressTimeout: Boolean,
+  cancelOnProgressTimeout: Boolean
 )
 
 object PubsubSourceConfigV2 {
diff --git a/modules/gcp/src/main/scala/common-streams-extensions/v2/PubsubSourceV2.scala b/modules/gcp/src/main/scala/common-streams-extensions/v2/PubsubSourceV2.scala
@@ -8,10 +8,11 @@
 package com.snowplowanalytics.snowplow.sources.pubsub.v2
 
 import cats.effect.{Async, Deferred, Ref, Resource, Sync}
-import cats.effect.std.Hotswap
+import cats.effect.std.{Hotswap, Queue, QueueSink}
 import cats.effect.kernel.Unique
+import cats.effect.implicits._
 import cats.implicits._
-import fs2.{Chunk, Stream}
+import fs2.{Chunk, Pipe, Stream}
 import org.typelevel.log4cats.Logger
 import org.typelevel.log4cats.slf4j.Slf4jLogger
 
@@ -32,9 +33,8 @@ import org.threeten.bp.{Duration => ThreetenDuration}
 import com.snowplowanalytics.snowplow.pubsub.GcpUserAgent
 import com.snowplowanalytics.snowplow.sources.SourceAndAck
 import com.snowplowanalytics.snowplow.sources.internal.{Checkpointer, LowLevelEvents, LowLevelSource}
-import com.snowplowanalytics.snowplow.sources.pubsub.v2.PubsubRetryOps.implicits._
 
-import scala.concurrent.duration.{DurationDouble, FiniteDuration}
+import scala.concurrent.duration.{Duration, DurationDouble, FiniteDuration}
 import scala.jdk.CollectionConverters._
 
 import java.util.concurrent.{ExecutorService, Executors, LinkedBlockingQueue}
@@ -75,19 +75,51 @@ object PubsubSourceV2 {
       _ <- Stream.eval(deferredResources.complete(PubsubCheckpointer.Resources(stub, refStates)))
     } yield Stream
       .range(0, parallelPullCount)
-      .map { i =>
-        val actionQueue = new LinkedBlockingQueue[SubscriberAction]()
-        val clientId    = UUID.randomUUID
-        val resource    = initializeStreamingPull(config, stub, actionQueue, i, clientId)
-        Stream.resource(Hotswap(resource)).flatMap { case (hotswap, _) =>
-          Stream
-            .eval(pullFromQueue(config, actionQueue, refStates, hotswap, resource, i))
-            .repeat
-            .concurrently(extendDeadlines(config, stub, refStates, i))
-        }
-      }
+      .map(i => miniPubsubStream(config, stub, refStates, i))
       .parJoinUnbounded
 
+  private def miniPubsubStream[F[_]: Async](
+    config: PubsubSourceConfigV2,
+    stub: SubscriberStub,
+    refStates: Ref[F, Map[Unique.Token, PubsubBatchState]],
+    channelAffinity: Int
+  ): Stream[F, LowLevelEvents[Vector[Unique.Token]]] = {
+    val jQueue   = new LinkedBlockingQueue[SubscriberAction]()
+    val clientId = UUID.randomUUID
+    val resource = initializeStreamingPull[F](config, stub, jQueue, channelAffinity, clientId)
+
+    for {
+      (hotswap, _) <- Stream.resource(Hotswap(resource))
+      fs2Queue <- Stream.eval(Queue.synchronous[F, SubscriberAction])
+      _ <- extendDeadlines(config, stub, refStates, channelAffinity).spawn
+      _ <- Stream.eval(queueToQueue(config, jQueue, fs2Queue, stub, channelAffinity)).repeat.spawn
+      lle <- Stream
+               .fromQueueUnterminated(fs2Queue)
+               .through(toLowLevelEvents(config, refStates, hotswap, resource, channelAffinity))
+    } yield lle
+  }
+
+  private def queueToQueue[F[_]: Async](
+    config: PubsubSourceConfigV2,
+    jQueue: LinkedBlockingQueue[SubscriberAction],
+    fs2Queue: QueueSink[F, SubscriberAction],
+    stub: SubscriberStub,
+    channelAffinity: Int
+  ): F[Unit] =
+    resolveNextAction(jQueue).flatMap {
+      case action @ SubscriberAction.ProcessRecords(records, controller, _) =>
+        val fallback = for {
+          _ <- Logger[F].debug(s"Cancelling Pubsub channel $channelAffinity for not making progress")
+          ackIds = records.map(_.getAckId)
+          _ <- if (config.cancelOnProgressTimeout) Sync[F].delay(controller.cancel()) else Sync[F].unit
+          _ <- if (config.modackOnProgressTimeout) Utils.modAck(config.subscription, stub, ackIds, Duration.Zero, channelAffinity)
+               else fs2Queue.offer(action)
+        } yield ()
+        fs2Queue.offer(action).timeoutTo(config.progressTimeout, fallback)
+      case action: SubscriberAction.SubscriberError =>
+        fs2Queue.offer(action)
+    }
+
   /**
    * Modify ack deadlines if we need more time to process the messages
    *
@@ -126,27 +158,19 @@ object PubsubSourceV2 {
       .evalMap { toExtend =>
         if (toExtend.isEmpty)
           Sync[F].sleep(0.5 * config.minRemainingDeadline * config.durationPerAckExtension)
-        else
-          toExtend.sortBy(_.currentDeadline).flatMap(_.ackIds).grouped(1000).toVector.traverse_ { ackIds =>
-            Utils
-              .modAck[F](config.subscription, stub, ackIds, config.durationPerAckExtension, channelAffinity)
-              .retryingOnTransientGrpcFailures
-              .recoveringOnGrpcInvalidArgument { s =>
-                // This can happen if ack IDs were acked before we modAcked
-                Logger[F].info(s"Ignoring error from GRPC when modifying ack IDs: ${s.getDescription}")
-              }
-          }
+        else {
+          val ackIds = toExtend.sortBy(_.currentDeadline).flatMap(_.ackIds)
+          Utils.modAck[F](config.subscription, stub, ackIds, config.durationPerAckExtension, channelAffinity)
+        }
       }
       .repeat
       .drain
 
   /**
-   * Pulls a SubscriberAction from a queue when one becomes available
+   * Pipe from SubscriberAction to LowLevelEvents TODO: Say what else this does
    *
    * @param config
    *   The source configuration
-   * @param queue
-   *   The queue from which to pull a SubscriberAction
    * @param refStates
    *   A map from tokens to the data held about a batch of messages received from pubsub. This
    *   function must update the state to add new batches.
@@ -159,44 +183,48 @@ object PubsubSourceV2 {
    *   Identifies the GRPC channel (TCP connection) creating these Actions. Each GRPC channel has
    *   its own queue, observer, and puller.
    */
-  private def pullFromQueue[F[_]: Async](
+  private def toLowLevelEvents[F[_]: Async](
     config: PubsubSourceConfigV2,
-    queue: LinkedBlockingQueue[SubscriberAction],
     refStates: Ref[F, Map[Unique.Token, PubsubBatchState]],
     hotswap: Hotswap[F, Unit],
     toSwap: Resource[F, Unit],
     channelAffinity: Int
-  ): F[LowLevelEvents[Vector[Unique.Token]]] = {
-    def go(delayOnSubscriberError: FiniteDuration): F[LowLevelEvents[Vector[Unique.Token]]] =
-      resolveNextAction[F, SubscriberAction](queue).flatMap {
-        case SubscriberAction.ProcessRecords(records, controller, timeReceived) =>
-          val chunk = Chunk.from(records.map(_.getMessage.getData.asReadOnlyByteBuffer()))
-          val (tstampSeconds, tstampNanos) =
-            records.map(r => (r.getMessage.getPublishTime.getSeconds, r.getMessage.getPublishTime.getNanos)).min
-          val ackIds = records.map(_.getAckId)
+  ): Pipe[F, SubscriberAction, LowLevelEvents[Vector[Unique.Token]]] =
+    _.flatMap {
+      case SubscriberAction.ProcessRecords(records, controller, timeReceived) =>
+        val chunk = Chunk.from(records.map(_.getMessage.getData.asReadOnlyByteBuffer()))
+        val (tstampSeconds, tstampNanos) =
+          records.map(r => (r.getMessage.getPublishTime.getSeconds, r.getMessage.getPublishTime.getNanos)).min
+        val ackIds = records.map(_.getAckId)
+        Stream.eval {
           for {
             token <- Unique[F].unique
             currentDeadline = timeReceived.plusMillis(config.durationPerAckExtension.toMillis)
             _ <- refStates.update(_ + (token -> PubsubBatchState(currentDeadline, ackIds, channelAffinity)))
             _ <- Sync[F].delay(controller.request(1))
           } yield LowLevelEvents(chunk, Vector(token), Some(Instant.ofEpochSecond(tstampSeconds, tstampNanos.toLong)))
-        case SubscriberAction.SubscriberError(t) =>
-          if (PubsubRetryOps.isRetryableException(t)) {
-            val nextDelay = (2 * delayOnSubscriberError).min((10 + scala.util.Random.nextDouble()).second)
-            // Log at debug level because retryable errors are very frequent.
-            // In particular, if the pubsub subscription is empty then a streaming pull returns UNAVAILABLE
+        }
+      case SubscriberAction.SubscriberError(t) =>
+        if (PubsubRetryOps.isRetryableException(t)) {
+          // val nextDelay = (2 * delayOnSubscriberError).min((10 + scala.util.Random.nextDouble()).second)
+          // Log at debug level because retryable errors are very frequent.
+          // In particular, if the pubsub subscription is empty then a streaming pull returns UNAVAILABLE
+          Stream.eval {
             Logger[F].debug(s"Retryable error on PubSub channel $channelAffinity: ${t.getMessage}") *>
               hotswap.clear *>
-              Async[F].sleep(delayOnSubscriberError) *>
-              hotswap.swap(toSwap) *>
-              go(nextDelay)
-          } else {
-            Logger[F].error(t)("Exception from PubSub source") *> Sync[F].raiseError(t)
-          }
-      }
-
-    go(delayOnSubscriberError = (1.0 + scala.util.Random.nextDouble()).second)
-  }
+              Async[F].sleep((1.0 + scala.util.Random.nextDouble()).second) *> // TODO expotential backoff
+              hotswap.swap(toSwap)
+          }.drain
+        } else if (t.isInstanceOf[java.util.concurrent.CancellationException]) {
+          Stream.eval {
+            Logger[F].debug("Cancellation exception on PubSub channel") *>
+              hotswap.clear *>
+              hotswap.swap(toSwap)
+          }.drain
+        } else {
+          Stream.eval(Logger[F].error(t)("Exception from PubSub source")) *> Stream.raiseError[F](t)
+        }
+    }
 
   private def resolveNextAction[F[_]: Sync, A](queue: LinkedBlockingQueue[A]): F[A] =
     Sync[F].delay(Option[A](queue.poll)).flatMap {
diff --git a/modules/gcp/src/main/scala/common-streams-extensions/v2/Utils.scala b/modules/gcp/src/main/scala/common-streams-extensions/v2/Utils.scala
@@ -9,34 +9,43 @@ package com.snowplowanalytics.snowplow.sources.pubsub.v2
 
 import cats.effect.{Async, Sync}
 import cats.implicits._
+import org.typelevel.log4cats.Logger
 
 import com.google.api.gax.grpc.GrpcCallContext
 import com.google.cloud.pubsub.v1.stub.SubscriberStub
 import com.google.pubsub.v1.ModifyAckDeadlineRequest
 import com.snowplowanalytics.snowplow.pubsub.FutureInterop
+import com.snowplowanalytics.snowplow.sources.pubsub.v2.PubsubRetryOps.implicits._
 
 import scala.concurrent.duration.FiniteDuration
 import scala.jdk.CollectionConverters._
 
 private object Utils {
 
-  def modAck[F[_]: Async](
+  def modAck[F[_]: Async: Logger](
     subscription: PubsubSourceConfigV2.Subscription,
     stub: SubscriberStub,
     ackIds: Vector[String],
     duration: FiniteDuration,
     channelAffinity: Int
-  ): F[Unit] = {
-    val request = ModifyAckDeadlineRequest.newBuilder
-      .setSubscription(subscription.show)
-      .addAllAckIds(ackIds.asJava)
-      .setAckDeadlineSeconds(duration.toSeconds.toInt)
-      .build
-    val context = GrpcCallContext.createDefault.withChannelAffinity(channelAffinity)
-    for {
-      apiFuture <- Sync[F].delay(stub.modifyAckDeadlineCallable.futureCall(request, context))
-      _ <- FutureInterop.fromFuture(apiFuture)
-    } yield ()
-  }
+  ): F[Unit] =
+    ackIds.grouped(1000).toVector.traverse_ { group =>
+      val request = ModifyAckDeadlineRequest.newBuilder
+        .setSubscription(subscription.show)
+        .addAllAckIds(group.asJava)
+        .setAckDeadlineSeconds(duration.toSeconds.toInt)
+        .build
+      val context = GrpcCallContext.createDefault.withChannelAffinity(channelAffinity)
+      val io = for {
+        apiFuture <- Sync[F].delay(stub.modifyAckDeadlineCallable.futureCall(request, context))
+        _ <- FutureInterop.fromFuture(apiFuture)
+      } yield ()
+
+      io.retryingOnTransientGrpcFailures
+        .recoveringOnGrpcInvalidArgument { s =>
+          // This can happen if ack IDs were acked before we modAcked
+          Logger[F].info(s"Ignoring error from GRPC when modifying ack IDs: ${s.getDescription}")
+        }
+    }
 
 }
