Pubsub source can use streaming pull

The PubsubSourceConfig now takes an option `streamingPull` which can be
set to true to use a GRPC streaming pull to fetch from pubsub.

Streaming pull is better for apps with low latency requirements.

Author: istreeter

modules/pubsub/src/main/resources/reference.conf

@@ -6,6 +6,13 @@ snowplow.defaults: {
       minRemainingAckDeadline: 0.1
       maxMessagesPerPull: 1000
       debounceRequests: "100 millis"
+      streamingPull: true
+      retries: {
+        transientErrors: {
+          delay: "100 millis"
+          attempts: 10
+        }
+      }
     }
   }
 
@@ -15,6 +22,12 @@ snowplow.defaults: {
       requestByteThreshold: 1000000
       # Equal to 9 MB. Pubsub message size limit 10 MB however, we further reduce it to 9 MB to be on the safe side.
       maxRecordSize: 9000000
+      retries: {
+        transientErrors: {
+          delay: "100 millis"
+          attempts: 10
+        }
+      }
     }
   }
 

modules/pubsub/src/main/scala/com/snowplowanalytics/snowplow/streams/pubsub/PubsubRetryOps.scala

@@ -7,25 +7,60 @@
  */
 package com.snowplowanalytics.snowplow.streams.pubsub
 
+import cats.MonadError
 import cats.implicits._
-import cats.effect.Async
+import cats.effect.{Async, Poll}
 import com.google.api.gax.rpc.{ApiException, StatusCode}
 import io.grpc.Status
 import org.typelevel.log4cats.Logger
-import retry.RetryPolicies
+import retry.{RetryPolicies, Sleep}
 import retry.implicits._
 
-import scala.concurrent.duration.DurationDouble
+import scala.concurrent.duration.FiniteDuration
 
 private[pubsub] object PubsubRetryOps {
 
   object implicits {
     implicit class Ops[F[_], A](val f: F[A]) extends AnyVal {
 
-      def retryingOnTransientGrpcFailures(implicit F: Async[F], L: Logger[F]): F[A] =
+      /**
+       * Retries an effect when an exception matches known retryable GRPC errors
+       *
+       * The retry policy is full jitter, with capped number of attempts
+       */
+      def retryingOnTransientGrpcFailures(retryDelay: FiniteDuration, retryAttempts: Int)(implicit F: Async[F], L: Logger[F]): F[A] =
+        retryingOnTransientGrpcFailuresImpl(retryDelay, retryAttempts)
+
+      /**
+       * Retries an effect when an exception matches known retryable GRPC errors
+       *
+       * The retry policy is full jitter, with capped number of attempts
+       *
+       * This variant takes a `Poll[F]` as a parameter. This is helpful when the effect is run
+       * inside `Sync[F].uncancelable` but we don't mind if the effect is cancelled in between retry
+       * attempts.
+       */
+      def retryingOnTransientGrpcFailuresWithCancelableSleep(
+        retryDelay: FiniteDuration,
+        retryAttempts: Int,
+        poll: Poll[F]
+      )(implicit F: Async[F],
+        L: Logger[F]
+      ): F[A] = {
+        implicit val s: Sleep[F] = (delay: FiniteDuration) => poll(Async[F].sleep(delay))
+        retryingOnTransientGrpcFailuresImpl(retryDelay, retryAttempts)
+      }
+
+      private def retryingOnTransientGrpcFailuresImpl(
+        retryDelay: FiniteDuration,
+        retryAttempts: Int
+      )(implicit F: MonadError[F, Throwable],
+        L: Logger[F],
+        s: Sleep[F]
+      ): F[A] =
         f.retryingOnSomeErrors(
           isWorthRetrying = { e => isRetryableException(e).pure[F] },
-          policy          = RetryPolicies.fullJitter(1.second),
+          policy          = RetryPolicies.fullJitter[F](retryDelay).join(RetryPolicies.limitRetries(retryAttempts - 1)),
           onError = { case (t, _) =>
             Logger[F].info(t)(s"Pubsub retryable GRPC error will be retried: ${t.getMessage}")
           }

modules/pubsub/src/main/scala/com/snowplowanalytics/snowplow/streams/pubsub/PubsubSinkConfig.scala

@@ -10,14 +10,23 @@ package com.snowplowanalytics.snowplow.streams.pubsub
 import cats.Id
 import io.circe.Decoder
 import io.circe.generic.semiauto._
+import io.circe.config.syntax._
+
+import scala.concurrent.duration.FiniteDuration
 
 case class PubsubSinkConfigM[M[_]](
   topic: M[PubsubSinkConfig.Topic],
   batchSize: Int,
-  requestByteThreshold: Int
+  requestByteThreshold: Int,
+  retries: PubsubSinkConfig.Retries
 )
 
 object PubsubSinkConfig {
+
+  case class TransientErrorRetrying(delay: FiniteDuration, attempts: Int)
+
+  case class Retries(transientErrors: TransientErrorRetrying)
+
   case class Topic(projectId: String, topicId: String)
 }
 
@@ -34,13 +43,20 @@ object PubsubSinkConfigM {
           Left("Expected format: projects/<project>/topics/<topic>")
       }
 
-  implicit def decoder: Decoder[PubsubSinkConfig] = deriveDecoder[PubsubSinkConfig]
+  implicit def decoder: Decoder[PubsubSinkConfig] = {
+    implicit val transientErrorDecoder = deriveDecoder[TransientErrorRetrying]
+    implicit val retriesDecoder        = deriveDecoder[Retries]
+    deriveDecoder[PubsubSinkConfig]
+  }
 
-  implicit def optionalDecoder: Decoder[Option[PubsubSinkConfig]] =
+  implicit def optionalDecoder: Decoder[Option[PubsubSinkConfig]] = {
+    implicit val transientErrorDecoder = deriveDecoder[TransientErrorRetrying]
+    implicit val retriesDecoder        = deriveDecoder[Retries]
     deriveDecoder[PubsubSinkConfigM[Option]].map {
-      case PubsubSinkConfigM(Some(t), a, b) =>
-        Some(PubsubSinkConfigM[Id](t, a, b))
+      case PubsubSinkConfigM(Some(t), a, b, c) =>
+        Some(PubsubSinkConfigM[Id](t, a, b, c))
       case _ =>
         None
     }
+  }
 }

modules/pubsub/src/main/scala/com/snowplowanalytics/snowplow/streams/pubsub/PubsubSourceConfig.scala

@@ -24,7 +24,7 @@ import scala.concurrent.duration.FiniteDuration
  *   Controls how many RPCs may be opened concurrently from the client to the PubSub server. The
  *   maximum number of RPCs is equal to this factor multiplied by the number of available cpu cores.
  *   Increasing this factor can increase the rate of events provided by the Source to the
- *   application.
+ *   application. Only used when `streamingPull` is false.
  * @param durationPerAckExtension
  *   Ack deadlines are extended for this duration. For common-streams apps this should be set
  *   slightly larger than the maximum time we expect between app receiving the message and acking
@@ -41,18 +41,32 @@ import scala.concurrent.duration.FiniteDuration
  *   Adds an artifical delay between consecutive requests to pubsub for more messages. Under some
  *   circumstances, this was found to slightly alleviate a problem in which pubsub might re-deliver
  *   the same messages multiple times.
+ * @param streamingPull
+ *   Controls whether to use "Streaming Pull" (true) or "Unary Pull" (false) GRPC method to fetch
+ *   events. Use "Unary Pull" (false) for apps which have long delays in between acks (e.g. Lake
+ *   Loader), to avoid a problem in which PubSub occasionally re-delivers the same messages, causing
+ *   downstream duplicates. Use "Streaming Pull" (true) for apps that ack often, and which have
+ *   lower latency requirements (e.g. Enrich).
+ * @param retries
+ *   Configuration for retrying transient errors
  */
 case class PubsubSourceConfig(
   subscription: PubsubSourceConfig.Subscription,
   parallelPullFactor: BigDecimal,
   durationPerAckExtension: FiniteDuration,
   minRemainingAckDeadline: BigDecimal,
   maxMessagesPerPull: Int,
-  debounceRequests: FiniteDuration
+  debounceRequests: FiniteDuration,
+  streamingPull: Boolean,
+  retries: PubsubSourceConfig.Retries
 )
 
 object PubsubSourceConfig {
 
+  case class TransientErrorRetrying(delay: FiniteDuration, attempts: Int)
+
+  case class Retries(transientErrors: TransientErrorRetrying)
+
   case class Subscription(projectId: String, subscriptionId: String)
 
   object Subscription {
@@ -71,5 +85,9 @@ object PubsubSourceConfig {
           Left("Expected format: projects/<project>/subscriptions/<subscription>")
       }
 
-  implicit def decoder: Decoder[PubsubSourceConfig] = deriveDecoder[PubsubSourceConfig]
+  implicit def decoder: Decoder[PubsubSourceConfig] = {
+    implicit val transientErrorDecoder = deriveDecoder[TransientErrorRetrying]
+    implicit val retriesDecoder        = deriveDecoder[Retries]
+    deriveDecoder[PubsubSourceConfig]
+  }
 }

modules/pubsub/src/main/scala/com/snowplowanalytics/snowplow/streams/pubsub/sink/PubsubSink.scala

@@ -80,7 +80,7 @@ private[pubsub] object PubsubSink {
           apiFuture <- Sync[F].delay(stub.publishCallable.futureCall(request))
           _ <- FutureInterop.fromFuture_(apiFuture)
         } yield ()
-        io.retryingOnTransientGrpcFailures
+        io.retryingOnTransientGrpcFailures(config.retries.transientErrors.delay, config.retries.transientErrors.attempts)
       }
 
   /**

modules/pubsub/src/main/scala/com/snowplowanalytics/snowplow/streams/pubsub/source/PubsubCheckpointer.scala

@@ -35,7 +35,7 @@ import com.snowplowanalytics.snowplow.streams.pubsub.PubsubRetryOps.implicits._
  *   state.
  */
 private class PubsubCheckpointer[F[_]: Async](
-  subscription: PubsubSourceConfig.Subscription,
+  config: PubsubSourceConfig,
   stub: SubscriberStub,
   refAckIds: Ref[F, Map[Unique.Token, PubsubBatchState]]
 ) extends Checkpointer[F, Vector[Unique.Token]] {
@@ -57,12 +57,13 @@ private class PubsubCheckpointer[F[_]: Async](
     for {
       ackDatas <- refAckIds.modify(m => (m -- c, c.flatMap(m.get)))
       _ <- ackDatas.flatMap(_.ackIds).grouped(1000).toVector.traverse_ { ackIds =>
-             val request = AcknowledgeRequest.newBuilder.setSubscription(subscription.show).addAllAckIds(ackIds.asJava).build
+             val request = AcknowledgeRequest.newBuilder.setSubscription(config.subscription.show).addAllAckIds(ackIds.asJava).build
              val attempt = for {
                apiFuture <- Sync[F].delay(stub.acknowledgeCallable.futureCall(request))
                _ <- FutureInterop.fromFuture[F, com.google.protobuf.Empty](apiFuture)
              } yield ()
-             attempt.retryingOnTransientGrpcFailures
+             attempt
+               .retryingOnTransientGrpcFailures(config.retries.transientErrors.delay, config.retries.transientErrors.attempts)
                .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}")
@@ -81,6 +82,6 @@ private class PubsubCheckpointer[F[_]: Async](
       ackDatas <- refAckIds.modify(m => (m -- c, c.flatMap(m.get)))
       ackIds = ackDatas.flatMap(_.ackIds)
       // A nack is just a modack with zero duration
-      _ <- Utils.modAck[F](subscription, stub, ackIds, Duration.Zero)
+      _ <- Utils.modAck[F](config, stub, ackIds, Duration.Zero)
     } yield ()
 }