Pubsub source v2 re-sends streaming request when blocked

snowplow-incubator · Oct 11, 2024 · ad1b4ea · ad1b4ea
1 parent e8ea8ea
commit ad1b4ea
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 31 deletions.
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
@@ -24,10 +24,7 @@ case class PubsubSourceConfigV2(
   minRemainingDeadline: Double,
   gcpUserAgent: GcpUserAgent,
   maxPullsPerTransportChannel: Int,
-  progressTimeout: FiniteDuration,
-  modackOnProgressTimeout: Boolean,
-  cancelOnProgressTimeout: Boolean,
-  consistentClientId: Boolean
+  progressTimeout: FiniteDuration
 )
 
 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
@@ -34,7 +34,7 @@ import com.snowplowanalytics.snowplow.pubsub.GcpUserAgent
 import com.snowplowanalytics.snowplow.sources.SourceAndAck
 import com.snowplowanalytics.snowplow.sources.internal.{Checkpointer, LowLevelEvents, LowLevelSource}
 
-import scala.concurrent.duration.{Duration, DurationDouble, FiniteDuration}
+import scala.concurrent.duration.{DurationDouble, FiniteDuration}
 import scala.jdk.CollectionConverters._
 
 import java.util.concurrent.{ExecutorService, Executors, LinkedBlockingQueue}
@@ -92,7 +92,7 @@ object PubsubSourceV2 {
       (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
+      _ <- Stream.eval(queueToQueue(config, jQueue, fs2Queue, channelAffinity, hotswap)).repeat.spawn
       lle <- Stream
                .fromQueueUnterminated(fs2Queue)
                .through(toLowLevelEvents(config, refStates, hotswap, resource, channelAffinity))
@@ -103,26 +103,18 @@ object PubsubSourceV2 {
     config: PubsubSourceConfigV2,
     jQueue: LinkedBlockingQueue[SubscriberAction],
     fs2Queue: QueueSink[F, SubscriberAction],
-    stub: SubscriberStub,
-    channelAffinity: Int
+    channelAffinity: Int,
+    hotswap: Hotswap[F, KeepAlive[F]]
   ): F[Unit] =
     resolveNextAction(jQueue).flatMap {
-      case action @ SubscriberAction.ProcessRecords(records, controller, _) =>
-        val fallback = if (config.modackOnProgressTimeout) {
-          val ackIds = records.map(_.getAckId)
-          if (config.cancelOnProgressTimeout)
-            Logger[F].debug(s"Cancelling Pubsub channel $channelAffinity for not making progress") *>
-              Sync[F].delay(controller.cancel()) *> Utils.modAck(config.subscription, stub, ackIds, Duration.Zero, channelAffinity)
-          else
-            Logger[F].debug(s"Nacking on Pubsub channel $channelAffinity for not making progress") *>
-              Sync[F].delay(controller.request(1)) *> Utils.modAck(config.subscription, stub, ackIds, Duration.Zero, channelAffinity)
-        } else {
-          if (config.cancelOnProgressTimeout)
-            Logger[F].debug(s"Cancelling Pubsub channel $channelAffinity for not making progress") *>
-              Sync[F].delay(controller.cancel()) *> fs2Queue.offer(action)
-          else
-            fs2Queue.offer(action)
-        }
+      case action @ SubscriberAction.ProcessRecords(_, _, _) =>
+        def fallback: F[Unit] =
+          hotswap.get.use {
+            case Some(keepAlive) =>
+              Logger[F].debug(s"Sending keepalie for channel $channelAffinity") *>
+                keepAlive.keepAlive
+            case None => Sync[F].unit
+          } >> fs2Queue.offer(action).timeoutTo(config.progressTimeout, fallback)
         fs2Queue.offer(action).timeoutTo(config.progressTimeout, fallback)
       case action: SubscriberAction.SubscriberError =>
         fs2Queue.offer(action)
@@ -194,8 +186,8 @@ object PubsubSourceV2 {
   private def toLowLevelEvents[F[_]: Async](
     config: PubsubSourceConfigV2,
     refStates: Ref[F, Map[Unique.Token, PubsubBatchState]],
-    hotswap: Hotswap[F, Unit],
-    toSwap: Resource[F, Unit],
+    hotswap: Hotswap[F, KeepAlive[F]],
+    toSwap: Resource[F, KeepAlive[F]],
     channelAffinity: Int
   ): Pipe[F, SubscriberAction, LowLevelEvents[Vector[Unique.Token]]] =
     _.flatMap {
@@ -282,7 +274,7 @@ object PubsubSourceV2 {
     actionQueue: LinkedBlockingQueue[SubscriberAction],
     channelAffinity: Int,
     clientId: UUID
-  ): Resource[F, Unit] = {
+  ): Resource[F, KeepAlive[F]] = {
 
     val observer = new ResponseObserver[StreamingPullResponse] {
       var controller: StreamController = _
@@ -314,7 +306,7 @@ object PubsubSourceV2 {
     val request = StreamingPullRequest.newBuilder
       .setSubscription(config.subscription.show)
       .setStreamAckDeadlineSeconds(config.durationPerAckExtension.toSeconds.toInt)
-      .setClientId(if (config.consistentClientId) clientId.toString else UUID.randomUUID.toString)
+      .setClientId(clientId.toString)
       .setMaxOutstandingMessages(0)
       .setMaxOutstandingBytes(0)
       .build
@@ -323,11 +315,18 @@ object PubsubSourceV2 {
       .make(Sync[F].delay(subStub.streamingPullCallable.splitCall(observer, context))) { stream =>
         Sync[F].delay(stream.closeSendWithError(Status.CANCELLED.asException))
       }
-      .evalMap { stream =>
-        Sync[F].delay(stream.send(request))
+      .map { stream =>
+        new KeepAlive[F] {
+          def keepAlive: F[Unit] =
+            Sync[F].delay(stream.send(request))
+        }
       }
-      .void
+      .evalTap(_.keepAlive)
+
+  }
 
+  trait KeepAlive[F[_]] {
+    def keepAlive: F[Unit]
   }
 
   private def executorResource[F[_]: Sync, E <: ExecutorService](make: F[E]): Resource[F, E] =