diff --git a/google-cloud-pub-sub-grpc/README.md b/google-cloud-pub-sub-grpc/README.md index c3cdf4c15..5b7c7af50 100644 --- a/google-cloud-pub-sub-grpc/README.md +++ b/google-cloud-pub-sub-grpc/README.md @@ -83,5 +83,6 @@ gcloud pubsub topics delete testTopic | Date | Environment | Credential Provider | Result | Notes | |------|-------------|---------------------|--------|-------| | 2026-03-13 | GCP (project: pekko-connectors-test) | google-application-default | 14/14 passed | Scala 8/8, Java 6/6. User credentials via `gcloud auth application-default login`. | +| 2026-05-11 | GCP (project: pekko-connectors) | google-application-default | 1/1 passed | New `Subscriber resource` scenario. Verifies eager-pull deadline tracking, subsequent-request field clearing, eager-pull flow control gate, and auto-cleanup. 10 messages, 6s processing, parallelism 2, `maxOutstandingMessages=3`. 35s. | After running against real GCP, add a row to the table above to record the result. \ No newline at end of file diff --git a/google-cloud-pub-sub-grpc/k8s/GkeFullFeatureTest.scala b/google-cloud-pub-sub-grpc/k8s/GkeFullFeatureTest.scala index 3da92fcb0..eb98c8a81 100644 --- a/google-cloud-pub-sub-grpc/k8s/GkeFullFeatureTest.scala +++ b/google-cloud-pub-sub-grpc/k8s/GkeFullFeatureTest.scala @@ -20,7 +20,7 @@ package org.apache.pekko.stream.connectors.googlecloud.pubsub.grpc.gke import org.apache.pekko import pekko.actor.ActorSystem import pekko.stream.RestartSettings -import pekko.stream.connectors.googlecloud.pubsub.grpc.{ AckDeadlineDistribution, FlowControl } +import pekko.stream.connectors.googlecloud.pubsub.grpc.{ AckDeadline, AckDeadlineDistribution, FlowControl } import pekko.stream.connectors.googlecloud.pubsub.grpc.scaladsl.GooglePubSub import pekko.stream.scaladsl.{ Flow, Sink, Source } import com.google.protobuf.ByteString @@ -62,6 +62,7 @@ object GkeFullFeatureTest { scenario5_FlowControl(topicFqrs, subFqrs) scenario6_NackAndRedeliver(topicFqrs, subFqrs) scenario7_DynamicDeadlineModification(topicFqrs, subFqrs) + scenario8_SubscriberResource(topicFqrs, subFqrs) println("\n=== ALL SCENARIOS PASSED ===") Await.result(system.terminate(), 10.seconds) @@ -405,4 +406,93 @@ object GkeFullFeatureTest { assert(msgs.size == messageCount, s"Expected $messageCount messages, got ${msgs.size}") println(" PASSED") } + + // --------------------------------------------------------------------------- + // Scenario 8: Subscriber resource — exercises all three bug fixes plus + // composition guarantees in a single end-to-end run. + // + // - Sets `maxOutstandingMessages` on the initial StreamingPullRequest. + // Pre-fix this caused INVALID_ARGUMENT on the second polling tick (bug 2). + // + // - Slow per-message processing (8 seconds) with parallelism well below the + // batch size, so messages buffer inside the eager-pull tracker for tens of + // seconds. Pre-fix the tracker would only see the messages currently in + // mapAsync, the rest would expire and get redelivered (bug 1). + // + // - Configures a small FlowControl limit. The new eager-pull gate counts + // messages on receipt rather than on push downstream, so combined with + // server-side maxOutstandingMessages it actually bounds delivery (bug 3). + // + // - Asserts every published message is received exactly once, no + // duplicates from redelivery, and that flowControl.outstandingCount + // reaches the configured limit at some point during the run. + // --------------------------------------------------------------------------- + private def scenario8_SubscriberResource(topicFqrs: String, subFqrs: String)( + implicit system: ActorSystem): Unit = { + println("\n--- Scenario 8: Subscriber resource (high-level API, all bug fixes) ---") + + val messageCount = 20 + val maxOutstanding = 5 + val processingDelay = 8.seconds + val testPrefix = s"scenario8-${System.nanoTime()}" + val messages = (1 to messageCount).map(i => + PubsubMessage().withData(ByteString.copyFromUtf8(s"$testPrefix-$i"))) + + Await.result( + Source + .single(PublishRequest(topicFqrs, messages)) + .via(GooglePubSub.publish(parallelism = 1)) + .runWith(Sink.head), + 30.seconds) + println(s" Published $messageCount messages") + + val flowControl = FlowControl(maxOutstandingMessages = maxOutstanding.toLong) + var maxObserved = 0L + + // Initial request sets BOTH stream ack deadline AND maxOutstandingMessages. + // Pre-bug-2 fix this would fail on the first keepalive tick with INVALID_ARGUMENT. + val request = StreamingPullRequest(subFqrs) + .withStreamAckDeadlineSeconds(15) + .withMaxOutstandingMessages(maxOutstanding.toLong) + + val restartSettings = RestartSettings( + minBackoff = 1.second, + maxBackoff = 10.seconds, + randomFactor = 0.2).withMaxRestarts(3, 1.minute) + + val subscriber = GooglePubSub.subscriber( + request = request, + pollInterval = 1.second, + ackDeadline = AckDeadline.Fixed(extensionInterval = 5.seconds, deadlineSeconds = 30), + restartSettings = Some(restartSettings), + flowControl = Some(flowControl)) + + try { + val received = subscriber.source + .filter(_.message.exists(_.data.toStringUtf8.startsWith(testPrefix))) + .take(messageCount) + .mapAsync(parallelism = 2) { msg => + val current = flowControl.outstandingCount + synchronized { if (current > maxObserved) maxObserved = current } + println(s" Processing: ${msg.message.map(_.data.toStringUtf8).getOrElse("?")} " + + s"(outstanding: $current/$maxOutstanding)") + // Slow processing forces autoExtend to actually fire while messages wait. + pekko.pattern.after(processingDelay)(Future.successful(msg)) + } + .map(msg => AcknowledgeRequest(subFqrs, Seq(msg.ackId))) + .runWith(subscriber.acknowledge(parallelism = 1)) + + Await.result(received, 5.minutes) + + println(s" Received and acked all $messageCount messages") + println(s" Max outstanding observed: $maxObserved (limit: $maxOutstanding)") + assert(maxObserved <= maxOutstanding, + s"Flow control violated: observed $maxObserved > limit $maxOutstanding") + assert(maxObserved >= 1L, "Flow control gate never registered any outstanding messages") + println(" PASSED (no redelivery during slow processing, server-side flow control respected, " + + "subsequent StreamingPullRequest accepted by server)") + } finally { + Await.result(subscriber.close(), 10.seconds) + } + } } diff --git a/google-cloud-pub-sub-grpc/k8s/build-and-push.sh b/google-cloud-pub-sub-grpc/k8s/build-and-push.sh index 41e08ae5e..373277dcb 100755 --- a/google-cloud-pub-sub-grpc/k8s/build-and-push.sh +++ b/google-cloud-pub-sub-grpc/k8s/build-and-push.sh @@ -24,6 +24,20 @@ REPO="pekko-test" IMAGE="${REGION}-docker.pkg.dev/${PROJECT_ID}/${REPO}/gke-auth-test:latest" SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" ROOT_DIR="$(cd "${SCRIPT_DIR}/../.." && pwd)" +STAGING="${SCRIPT_DIR}/staging" + +# Run cleanup unconditionally on exit (success, failure, or interrupt). Without this trap, +# a script failure between the "Building" step and the "Cleaning up" step at the bottom +# leaves GkeAuthTest.scala / GkeFullFeatureTest.scala leaked into src/main/scala/.../gke/ +# and a populated staging/ directory. +cleanup() { + rm -rf "${STAGING}" + rm -f "${ROOT_DIR}/google-cloud-pub-sub-grpc/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/gke/GkeAuthTest.scala" + rm -f "${ROOT_DIR}/google-cloud-pub-sub-grpc/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/gke/GkeFullFeatureTest.scala" + rm -f "${ROOT_DIR}/google-cloud-pub-sub-grpc/src/main/resources/gke-application.conf" + rmdir "${ROOT_DIR}/google-cloud-pub-sub-grpc/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/gke" 2>/dev/null || true +} +trap cleanup EXIT echo "=== Creating Artifact Registry repo (if needed) ===" gcloud artifacts repositories create "${REPO}" \ @@ -51,8 +65,7 @@ echo "=== Packaging ===" FULL_CP=$(sbt --error "print google-cloud-pub-sub-grpc/fullClasspath" | tr ',' '\n' | sed 's/.*Attributed(\(.*\))/\1/') CLASSES_DIR=$(sbt --error "print google-cloud-pub-sub-grpc/classDirectory" | tr -d '[:space:]') -# Create staging directory -STAGING="${SCRIPT_DIR}/staging" +# Create staging directory (already declared at the top of the script for the cleanup trap) rm -rf "${STAGING}" mkdir -p "${STAGING}/lib" @@ -94,10 +107,5 @@ docker build -t "${IMAGE}" "${STAGING}" echo "=== Pushing to Artifact Registry ===" docker push "${IMAGE}" -echo "=== Cleaning up ===" -rm -rf "${STAGING}" -rm -f "${ROOT_DIR}/google-cloud-pub-sub-grpc/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/gke/GkeAuthTest.scala" -rm -f "${ROOT_DIR}/google-cloud-pub-sub-grpc/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/gke/GkeFullFeatureTest.scala" -rm -f "${ROOT_DIR}/google-cloud-pub-sub-grpc/src/main/resources/gke-application.conf" - echo "=== Done: ${IMAGE} ===" +# Cleanup runs from the EXIT trap declared at the top of this script. diff --git a/google-cloud-pub-sub-grpc/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/AckDeadline.scala b/google-cloud-pub-sub-grpc/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/AckDeadline.scala new file mode 100644 index 000000000..987baff5c --- /dev/null +++ b/google-cloud-pub-sub-grpc/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/AckDeadline.scala @@ -0,0 +1,80 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.pekko.stream.connectors.googlecloud.pubsub.grpc + +import org.apache.pekko.annotation.ApiMayChange + +import scala.concurrent.duration._ + +/** + * Configuration for how a [[scaladsl.Subscriber]] (or its Java equivalent) extends ack deadlines + * for messages it has received. + * + * - [[AckDeadline.Fixed]] uses a constant deadline value on every extension. Cheapest, easiest + * to reason about. Use this if your processing latency is predictable. + * - [[AckDeadline.Adaptive]] adapts the deadline based on observed processing latencies via a + * shared [[AckDeadlineDistribution]], matching Google's official client library behavior. + * Use this if processing latency varies widely. + * + * @since 2.0.0 + */ +@ApiMayChange +sealed trait AckDeadline { + def extensionInterval: FiniteDuration +} + +@ApiMayChange +object AckDeadline { + + /** + * Extend deadlines on a fixed schedule using a constant deadline value. + * + * @param extensionInterval how often to extend deadlines (should be less than the deadline) + * @param deadlineSeconds the new deadline to set on each extension + * @param maxAckExtensionPeriod maximum total time to keep extending a message's deadline + * (default 60 minutes, matching Google's client library) + */ + final case class Fixed( + extensionInterval: FiniteDuration, + deadlineSeconds: Int, + maxAckExtensionPeriod: FiniteDuration = 60.minutes) extends AckDeadline + + /** + * Extend deadlines on a fixed schedule using an adaptive deadline computed from an + * [[AckDeadlineDistribution]]. The same distribution must be passed to acknowledge/nack + * operators so that completion latencies are recorded into the histogram. + */ + final case class Adaptive( + extensionInterval: FiniteDuration, + distribution: AckDeadlineDistribution) extends AckDeadline + + /** Java API: fixed-deadline configuration with a default `maxAckExtensionPeriod` of 60 minutes. */ + def fixed(extensionInterval: java.time.Duration, deadlineSeconds: Int): AckDeadline = + Fixed(FiniteDuration(extensionInterval.toNanos, NANOSECONDS), deadlineSeconds) + + /** Java API: fixed-deadline configuration with an explicit `maxAckExtensionPeriod`. */ + def fixed(extensionInterval: java.time.Duration, deadlineSeconds: Int, + maxAckExtensionPeriod: java.time.Duration): AckDeadline = + Fixed(FiniteDuration(extensionInterval.toNanos, NANOSECONDS), deadlineSeconds, + FiniteDuration(maxAckExtensionPeriod.toNanos, NANOSECONDS)) + + /** Java API: adaptive-deadline configuration. */ + def adaptive(extensionInterval: java.time.Duration, + distribution: AckDeadlineDistribution): AckDeadline = + Adaptive(FiniteDuration(extensionInterval.toNanos, NANOSECONDS), distribution) +} diff --git a/google-cloud-pub-sub-grpc/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/AckDeadlineExtender.scala b/google-cloud-pub-sub-grpc/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/AckDeadlineExtender.scala new file mode 100644 index 000000000..93ef741c0 --- /dev/null +++ b/google-cloud-pub-sub-grpc/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/AckDeadlineExtender.scala @@ -0,0 +1,221 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.pekko.stream.connectors.googlecloud.pubsub.grpc + +import org.apache.pekko +import pekko.Done +import pekko.actor.{ Cancellable, ClassicActorSystemProvider } +import pekko.annotation.ApiMayChange +import pekko.stream.Materializer +import pekko.stream.scaladsl.{ Keep, Sink, Source } +import pekko.stream.connectors.googlecloud.pubsub.grpc.scaladsl.{ GrpcSubscriber, GrpcSubscriberExt } +import com.google.pubsub.v1.pubsub.{ ModifyAckDeadlineRequest, SubscriberClient } + +import java.util.concurrent.ConcurrentHashMap +import scala.concurrent.duration._ +import scala.concurrent.{ ExecutionContext, Future } +import scala.jdk.CollectionConverters._ + +/** + * A long-lived ack-deadline extender. Owns the tracking map and the background ticker, both of + * which live above the lifetime of any single Pub/Sub streaming pull. Pass the same instance to + * `GooglePubSub.autoExtendAckDeadlines` (which calls `track` on every received message) so that + * tracking state survives stream reconnects driven by `RestartSource.withBackoff`. + * + * This mirrors how Google's official `google-cloud-pubsub` Java client builds `MessageDispatcher` + * once per `StreamingSubscriberConnection` and reuses it across every gRPC stream restart. In + * that client, the dispatcher's `pendingMessages` map and background lease-extension job both + * persist across reconnects; only the gRPC stream itself is rebuilt. This class brings the + * same lifecycle to the pekko-connectors subscriber. + * + * Two flavors: + * + * - Fixed deadline: every extension uses a constant `ackDeadlineSeconds`. + * - Adaptive deadline: every extension reads from a shared [[AckDeadlineDistribution]], which + * computes the deadline from observed processing latencies (matching Google's adaptive + * behavior). The same distribution must be passed to the acknowledge/nack operators so that + * completion times are recorded. + * + * Lifecycle: + * + * - Created via `AckDeadlineExtender(...)` (Scala) or `AckDeadlineExtender.create(...)` (Java). + * The background ticker starts immediately at construction. + * - Use the same instance for every materialization of `autoExtendAckDeadlines` you intend to + * share tracking state across. With `RestartSource.withBackoff`, create one extender outside + * the restart envelope and reference it inside. + * - Call `close()` on shutdown. This stops the ticker and clears the tracking map. Idempotent. + * + * Failure semantics: + * + * - If the background ticker's `ModifyAckDeadline` RPC fails, the ticker fails and `tickerDone` + * completes with that exception. `GooglePubSub.autoExtendAckDeadlines(extender)` wires this + * into a per-materialization `KillSwitch` so the inner stream aborts with an + * [[AckDeadlineExtensionException]]. After such a failure the extender is unusable; create a + * new one to recover. + * + * @since 2.0.0 + */ +@ApiMayChange +final class AckDeadlineExtender private ( + val subscription: String, + val extensionInterval: FiniteDuration, + val maxAckExtensionPeriod: FiniteDuration, + private[grpc] val tracked: ConcurrentHashMap[String, java.lang.Long], + private val computeDeadlineSeconds: () => Int, + private val client: SubscriberClient, + private[grpc] val materializer: Materializer) { + + private val maxNanos = maxAckExtensionPeriod.toNanos + + private val (ticker, _tickerDone): (Cancellable, Future[Done]) = Source + .tick(extensionInterval, extensionInterval, ()) + .mapAsync(1) { _ => + val now = System.nanoTime() + // Remove expired entries + tracked.asScala.foreach { case (ackId, entryTime) => + if (now - entryTime.longValue() > maxNanos) tracked.remove(ackId) + } + val ids = tracked.asScala.keys.toSeq + if (ids.nonEmpty) + client + .modifyAckDeadline(ModifyAckDeadlineRequest(subscription, ids, computeDeadlineSeconds())) + .map(_ => Done)(ExecutionContext.parasitic) + else + Future.successful(Done) + } + .toMat(Sink.ignore)(Keep.both) + .run()(materializer) + + /** Future that completes when the background ticker stops, with the cause if it failed. */ + def tickerDone: Future[Done] = _tickerDone + + /** Begin tracking an ackId. Called by `GooglePubSub.autoExtendAckDeadlines`. */ + private[grpc] def track(ackId: String): Unit = + tracked.put(ackId, java.lang.Long.valueOf(System.nanoTime())) + + /** Number of ackIds currently being tracked. Useful for diagnostics. */ + def trackedSize: Int = tracked.size() + + /** + * Stop the background ticker and clear the tracking map. Idempotent; subsequent calls return + * the same `Future[Done]` from the original ticker shutdown. + */ + def close(): Future[Done] = { + ticker.cancel() + tracked.clear() + _tickerDone + } +} + +@ApiMayChange +object AckDeadlineExtender { + + /** + * Create a fixed-deadline extender that uses the configured `subscriber` from the given actor + * system. The background ticker starts immediately. + */ + def apply( + subscription: String, + extensionInterval: FiniteDuration, + ackDeadlineSeconds: Int)(implicit system: ClassicActorSystemProvider): AckDeadlineExtender = + apply(subscription, extensionInterval, ackDeadlineSeconds, 60.minutes) + + /** Fixed-deadline extender with a custom `maxAckExtensionPeriod`. */ + def apply( + subscription: String, + extensionInterval: FiniteDuration, + ackDeadlineSeconds: Int, + maxAckExtensionPeriod: FiniteDuration)( + implicit system: ClassicActorSystemProvider): AckDeadlineExtender = + apply(subscription, extensionInterval, ackDeadlineSeconds, maxAckExtensionPeriod, + GrpcSubscriberExt()(system).subscriber) + + /** Fixed-deadline extender with an explicit subscriber (useful for tests). */ + def apply( + subscription: String, + extensionInterval: FiniteDuration, + ackDeadlineSeconds: Int, + maxAckExtensionPeriod: FiniteDuration, + subscriber: GrpcSubscriber)( + implicit system: ClassicActorSystemProvider): AckDeadlineExtender = + new AckDeadlineExtender( + subscription = subscription, + extensionInterval = extensionInterval, + maxAckExtensionPeriod = maxAckExtensionPeriod, + tracked = new ConcurrentHashMap[String, java.lang.Long](), + computeDeadlineSeconds = () => ackDeadlineSeconds, + client = subscriber.client, + materializer = Materializer.matFromSystem(system.classicSystem)) + + /** + * Create an adaptive extender backed by an [[AckDeadlineDistribution]]. The extender shares + * the distribution's `deliveryTimes` map so that completion records (from the + * acknowledge/nack operators) and tracking entries (from the eager-pull stage) reference the + * same state. Each tick reads the distribution's currently-computed adaptive deadline. + */ + def apply( + subscription: String, + extensionInterval: FiniteDuration, + distribution: AckDeadlineDistribution)( + implicit system: ClassicActorSystemProvider): AckDeadlineExtender = + apply(subscription, extensionInterval, distribution, GrpcSubscriberExt()(system).subscriber) + + /** Adaptive extender with an explicit subscriber. */ + def apply( + subscription: String, + extensionInterval: FiniteDuration, + distribution: AckDeadlineDistribution, + subscriber: GrpcSubscriber)( + implicit system: ClassicActorSystemProvider): AckDeadlineExtender = + new AckDeadlineExtender( + subscription = subscription, + extensionInterval = extensionInterval, + maxAckExtensionPeriod = distribution.maxAckExtensionPeriodNanos.nanos, + tracked = distribution.deliveryTimes, + computeDeadlineSeconds = () => distribution.currentDeadlineSeconds, + client = subscriber.client, + materializer = Materializer.matFromSystem(system.classicSystem)) + + /** Java API: fixed-deadline extender with defaults. */ + def create( + subscription: String, + extensionInterval: java.time.Duration, + ackDeadlineSeconds: Int, + system: ClassicActorSystemProvider): AckDeadlineExtender = + apply(subscription, FiniteDuration(extensionInterval.toNanos, NANOSECONDS), + ackDeadlineSeconds)(system) + + /** Java API: fixed-deadline extender with an explicit `maxAckExtensionPeriod`. */ + def create( + subscription: String, + extensionInterval: java.time.Duration, + ackDeadlineSeconds: Int, + maxAckExtensionPeriod: java.time.Duration, + system: ClassicActorSystemProvider): AckDeadlineExtender = + apply(subscription, FiniteDuration(extensionInterval.toNanos, NANOSECONDS), + ackDeadlineSeconds, FiniteDuration(maxAckExtensionPeriod.toNanos, NANOSECONDS))(system) + + /** Java API: adaptive extender. */ + def create( + subscription: String, + extensionInterval: java.time.Duration, + distribution: AckDeadlineDistribution, + system: ClassicActorSystemProvider): AckDeadlineExtender = + apply(subscription, FiniteDuration(extensionInterval.toNanos, NANOSECONDS), + distribution)(system) +} diff --git a/google-cloud-pub-sub-grpc/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/FlowControl.scala b/google-cloud-pub-sub-grpc/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/FlowControl.scala index 15b5030cd..fa89347e6 100644 --- a/google-cloud-pub-sub-grpc/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/FlowControl.scala +++ b/google-cloud-pub-sub-grpc/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/FlowControl.scala @@ -51,7 +51,13 @@ final class FlowControl(val maxOutstandingMessages: Long) { @volatile private[grpc] var onRelease: () => Unit = () => () - /** Current number of outstanding (unacknowledged) messages. */ + /** + * Current number of outstanding (unacknowledged) messages. Counts every message that has + * been received by `flowControlGate` but not yet released by the corresponding acknowledge + * or nack operator. Includes messages still buffered inside the gate that have not yet been + * emitted downstream, since the gate acquires permits on receipt to bound server-side + * delivery, not just on push to downstream. + */ def outstandingCount: Long = outstanding.get() /** Release `count` permits, signalling that messages have been acknowledged. */ diff --git a/google-cloud-pub-sub-grpc/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/impl/EagerPullTrackingStage.scala b/google-cloud-pub-sub-grpc/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/impl/EagerPullTrackingStage.scala new file mode 100644 index 000000000..06b6e0082 --- /dev/null +++ b/google-cloud-pub-sub-grpc/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/impl/EagerPullTrackingStage.scala @@ -0,0 +1,91 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.pekko.stream.connectors.googlecloud.pubsub.grpc.impl + +import org.apache.pekko.annotation.InternalApi +import org.apache.pekko.stream._ +import org.apache.pekko.stream.stage._ + +/** + * INTERNAL API + * + * A passthrough GraphStage that eagerly pulls from upstream into a bounded internal buffer + * and invokes `onTrack` the moment an element is grabbed — before downstream demand is required. + * + * Motivation: Pub/Sub `StreamingPull` starts the server-side ack deadline timer the moment a + * message is dispatched to the client. With a plain `.map { tracked.put(...); identity }`, + * tracking only fires when downstream pulls — so under backpressure (e.g. saturated `mapAsync`), + * messages buffer in the gRPC adapter / `mapConcat` with deadlines ticking but no client-side + * tracking, leaving them unprotected from the auto-extend ticker. + * + * This stage decouples receipt-side tracking from downstream demand. When the buffer is full + * the stage stops pulling, applying backpressure upstream — which (combined with server-side + * `StreamingPullRequest.maxOutstandingMessages`) bounds memory. + * + * Mirrors the pattern in Google's `MessageDispatcher.processReceivedMessages`, which registers + * messages in `pendingMessages` before handing them to user callbacks. + */ +@InternalApi +private[grpc] final class EagerPullTrackingStage[T](maxBuffer: Int, onTrack: T => Unit) + extends GraphStage[FlowShape[T, T]] { + require(maxBuffer > 0, "maxBuffer must be > 0") + + val in: Inlet[T] = Inlet("EagerPullTracking.in") + val out: Outlet[T] = Outlet("EagerPullTracking.out") + override val shape: FlowShape[T, T] = FlowShape(in, out) + + override def createLogic(inheritedAttributes: Attributes): GraphStageLogic = + new GraphStageLogic(shape) with InHandler with OutHandler { + // Pre-size to maxBuffer so the dynamic-array doesn't have to resize during initial fill. + // Buffer fill is bounded by maxBuffer via pullIfPossible, so this is also the high water mark. + private val buffer = new java.util.ArrayDeque[T](maxBuffer) + + override def preStart(): Unit = pull(in) + + override def onPush(): Unit = { + val msg = grab(in) + onTrack(msg) + buffer.offer(msg) + pushIfPossible() + pullIfPossible() + } + + override def onPull(): Unit = { + pushIfPossible() + pullIfPossible() + } + + override def onUpstreamFinish(): Unit = + if (buffer.isEmpty) completeStage() + // else: drain via subsequent onPull invocations + + override def onUpstreamFailure(ex: Throwable): Unit = { + buffer.clear() + super.onUpstreamFailure(ex) + } + + private def pushIfPossible(): Unit = + if (isAvailable(out) && !buffer.isEmpty) push(out, buffer.poll()) + + private def pullIfPossible(): Unit = + if (buffer.size < maxBuffer && !hasBeenPulled(in) && !isClosed(in)) pull(in) + else if (buffer.isEmpty && isClosed(in)) completeStage() + + setHandlers(in, out, this) + } +} diff --git a/google-cloud-pub-sub-grpc/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/impl/FlowControlGateStage.scala b/google-cloud-pub-sub-grpc/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/impl/FlowControlGateStage.scala index 81236f403..30b5710af 100644 --- a/google-cloud-pub-sub-grpc/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/impl/FlowControlGateStage.scala +++ b/google-cloud-pub-sub-grpc/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/impl/FlowControlGateStage.scala @@ -25,10 +25,21 @@ import org.apache.pekko.stream.connectors.googlecloud.pubsub.grpc.FlowControl /** * INTERNAL API * - * A GraphStage that gates elements through based on a shared [[FlowControl]] counter. - * It pulls from upstream only when the number of outstanding messages is below the limit. - * When downstream acknowledges/nacks via the same [[FlowControl]], the `onRelease` callback - * triggers a re-evaluation of whether to pull. + * A GraphStage that gates elements through based on a shared [[FlowControl]] counter using an + * eager-pull strategy: elements are pulled from upstream and counted against the permit limit + * the moment they arrive, independent of downstream demand. When the permit limit is reached + * the stage stops pulling, applying backpressure upstream all the way to the gRPC adapter. + * + * Combined with the connector's gRPC backpressure, this transitively bounds how many messages + * the Pub/Sub server is allowed to deliver, giving you flow control with semantics close to + * Google's `FlowController` rather than just downstream credit accounting. + * + * Internally the stage holds a small buffer for elements that have been received and counted + * but not yet pushed downstream. The buffer is naturally bounded by the FlowControl limit, so + * no separate maxBuffer parameter is needed. + * + * When downstream acknowledges or nacks via the same [[FlowControl]], the `onRelease` callback + * triggers a re-evaluation of whether to pull more. */ @InternalApi private[grpc] final class FlowControlGateStage[T](flowControl: FlowControl) @@ -41,40 +52,56 @@ private[grpc] final class FlowControlGateStage[T](flowControl: FlowControl) override def createLogic(inheritedAttributes: Attributes): GraphStageLogic = new GraphStageLogic(shape) with InHandler with OutHandler { - private var downstreamWaiting = false + // Pre-size to the FlowControl limit (clamped to a sane Int) so the dynamic array doesn't + // resize during initial fill. Fill is bounded by the limit via pullIfPossible, so this is + // the high water mark for typical configurations. The clamp prevents a pathological + // initial-allocation when callers set maxOutstandingMessages near Long.MaxValue. + private val buffer = + new java.util.ArrayDeque[T](math.min(flowControl.maxOutstandingMessages, 65536L).toInt) private val releaseCallback: AsyncCallback[Unit] = getAsyncCallback[Unit] { _ => - if (downstreamWaiting && flowControl.outstanding.get() < flowControl.maxOutstandingMessages && - !hasBeenPulled(in)) { - downstreamWaiting = false - pull(in) - } + pullIfPossible() } override def preStart(): Unit = { flowControl.onRelease = () => releaseCallback.invoke(()) + pull(in) } override def postStop(): Unit = { flowControl.onRelease = () => () } - // InHandler override def onPush(): Unit = { val msg = grab(in) - flowControl.acquire() - push(out, msg) + flowControl.acquire() // count on receipt, not on emission to downstream + buffer.offer(msg) + pushIfPossible() + pullIfPossible() } - // OutHandler override def onPull(): Unit = { - if (flowControl.outstanding.get() < flowControl.maxOutstandingMessages) { - pull(in) - } else { - downstreamWaiting = true - } + pushIfPossible() + pullIfPossible() + } + + override def onUpstreamFinish(): Unit = + if (buffer.isEmpty) completeStage() + // else: drain via subsequent onPull invocations + + override def onUpstreamFailure(ex: Throwable): Unit = { + buffer.clear() + super.onUpstreamFailure(ex) } + private def pushIfPossible(): Unit = + if (isAvailable(out) && !buffer.isEmpty) push(out, buffer.poll()) + + private def pullIfPossible(): Unit = + if (flowControl.outstanding.get() < flowControl.maxOutstandingMessages + && !hasBeenPulled(in) && !isClosed(in)) pull(in) + else if (buffer.isEmpty && isClosed(in)) completeStage() + setHandlers(in, out, this) } } diff --git a/google-cloud-pub-sub-grpc/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/javadsl/GooglePubSub.scala b/google-cloud-pub-sub-grpc/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/javadsl/GooglePubSub.scala index 1e717a69f..c917937fd 100644 --- a/google-cloud-pub-sub-grpc/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/javadsl/GooglePubSub.scala +++ b/google-cloud-pub-sub-grpc/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/javadsl/GooglePubSub.scala @@ -23,10 +23,11 @@ import pekko.stream.{ Attributes, KillSwitches, Materializer, RestartSettings } import pekko.stream.javadsl.{ Flow, Keep, RestartSource, Sink, Source } import pekko.stream.connectors.googlecloud.pubsub.grpc.{ AckDeadlineDistribution, + AckDeadlineExtender, AckDeadlineExtensionException, FlowControl } -import pekko.stream.connectors.googlecloud.pubsub.grpc.impl.FlowControlGateStage +import pekko.stream.connectors.googlecloud.pubsub.grpc.impl.{ EagerPullTrackingStage, FlowControlGateStage } import pekko.{ Done, NotUsed } import com.google.pubsub.v1._ @@ -37,6 +38,12 @@ import scala.jdk.CollectionConverters._ */ object GooglePubSub { + /** + * Default size of the eager-pull buffer used by `autoExtendAckDeadlines`. Matches the + * default `maxOutstandingMessages` of Google's official client library. + */ + final val DefaultEagerPullBuffer: Int = 1000 + /** * Create a flow to publish messages to Google Cloud Pub/Sub. The flow emits responses that contain published * message ids. @@ -239,6 +246,27 @@ object GooglePubSub { extensionInterval: Duration, ackDeadlineSeconds: Int, maxAckExtensionPeriod: Duration): Flow[ReceivedMessage, ReceivedMessage, NotUsed] = + autoExtendAckDeadlines(subscription, extensionInterval, ackDeadlineSeconds, maxAckExtensionPeriod, + DefaultEagerPullBuffer) + + /** + * Create a flow that automatically extends ack deadlines, with an explicit eager-pull buffer size. + * + * Identical to the four-argument overload but allows tuning the internal buffer used to + * track messages on receipt. A larger buffer absorbs longer downstream stalls without + * losing track of in-flight messages; a smaller buffer applies upstream backpressure sooner. + * + * @param maxBuffer maximum number of in-flight messages held in the eager-pull buffer + * (default [[DefaultEagerPullBuffer]]) + * @since 2.0.0 + */ + @ApiMayChange + def autoExtendAckDeadlines( + subscription: String, + extensionInterval: Duration, + ackDeadlineSeconds: Int, + maxAckExtensionPeriod: Duration, + maxBuffer: Int): Flow[ReceivedMessage, ReceivedMessage, NotUsed] = Flow .fromMaterializer { (mat, attr) => val client = subscriber(mat, attr).client @@ -282,10 +310,8 @@ object GooglePubSub { Flow.create[ReceivedMessage]() .via(killSwitch.flow[ReceivedMessage]) - .map(((msg: ReceivedMessage) => { - tracked.put(msg.getAckId, java.lang.Long.valueOf(System.nanoTime())) - msg - }): pekko.japi.function.Function[ReceivedMessage, ReceivedMessage]) + .via(new EagerPullTrackingStage[ReceivedMessage](maxBuffer, + msg => tracked.put(msg.getAckId, java.lang.Long.valueOf(System.nanoTime())))) .watchTermination((_, done: CompletionStage[Done]) => { done.whenComplete((_, _) => { ticker.cancel() @@ -313,6 +339,20 @@ object GooglePubSub { subscription: String, extensionInterval: Duration, distribution: AckDeadlineDistribution): Flow[ReceivedMessage, ReceivedMessage, NotUsed] = + autoExtendAckDeadlines(subscription, extensionInterval, distribution, DefaultEagerPullBuffer) + + /** + * Adaptive variant of `autoExtendAckDeadlines` with an explicit eager-pull buffer size. + * See the four-argument fixed-deadline overload for buffer semantics. + * + * @since 2.0.0 + */ + @ApiMayChange + def autoExtendAckDeadlines( + subscription: String, + extensionInterval: Duration, + distribution: AckDeadlineDistribution, + maxBuffer: Int): Flow[ReceivedMessage, ReceivedMessage, NotUsed] = Flow .fromMaterializer { (mat, attr) => val client = subscriber(mat, attr).client @@ -354,10 +394,8 @@ object GooglePubSub { Flow.create[ReceivedMessage]() .via(killSwitch.flow[ReceivedMessage]) - .map(((msg: ReceivedMessage) => { - distribution.recordDelivery(msg.getAckId) - msg - }): pekko.japi.function.Function[ReceivedMessage, ReceivedMessage]) + .via(new EagerPullTrackingStage[ReceivedMessage](maxBuffer, + msg => distribution.recordDelivery(msg.getAckId))) .watchTermination((_, done: CompletionStage[Done]) => { done.whenComplete((_, _) => { ticker.cancel() @@ -368,6 +406,46 @@ object GooglePubSub { } .mapMaterializedValue(_ => NotUsed) + /** + * Create a flow that automatically extends ack deadlines using a caller-owned + * [[AckDeadlineExtender]]. The extender owns the tracking map and the background ticker, both + * of which live above the lifetime of any single Pub/Sub streaming pull. This makes the flow + * restart-safe: when wrapped in `RestartSource.withBackoff`, messages received before a + * stream failure remain in the extender's tracking map and continue to receive deadline + * extensions during the backoff window. + * + * @since 2.0.0 + */ + @ApiMayChange + def autoExtendAckDeadlines(extender: AckDeadlineExtender): Flow[ReceivedMessage, ReceivedMessage, NotUsed] = + autoExtendAckDeadlines(extender, DefaultEagerPullBuffer) + + /** + * Caller-owned-extender variant of `autoExtendAckDeadlines` with an explicit eager-pull + * buffer size. + * + * @since 2.0.0 + */ + @ApiMayChange + def autoExtendAckDeadlines( + extender: AckDeadlineExtender, + maxBuffer: Int): Flow[ReceivedMessage, ReceivedMessage, NotUsed] = + Flow + .fromMaterializer { (_, _) => + val killSwitch = KillSwitches.shared("autoExtendAckDeadlines") + extender.tickerDone.onComplete { + case scala.util.Failure(ex) => + killSwitch.abort(new AckDeadlineExtensionException( + "Lease management ticker failed; ack deadline extensions have stopped", ex)) + case _ => () + }(scala.concurrent.ExecutionContext.parasitic) + + Flow.create[ReceivedMessage]() + .via(killSwitch.flow[ReceivedMessage]) + .via(new EagerPullTrackingStage[ReceivedMessage](maxBuffer, msg => extender.track(msg.getAckId))) + } + .mapMaterializedValue(_ => NotUsed) + /** * Create a flow that modifies the ack deadline for each message using a dynamic function. * diff --git a/google-cloud-pub-sub-grpc/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/scaladsl/GooglePubSub.scala b/google-cloud-pub-sub-grpc/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/scaladsl/GooglePubSub.scala index 4408f8503..978e37edc 100644 --- a/google-cloud-pub-sub-grpc/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/scaladsl/GooglePubSub.scala +++ b/google-cloud-pub-sub-grpc/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/scaladsl/GooglePubSub.scala @@ -19,13 +19,15 @@ import pekko.annotation.ApiMayChange import pekko.stream.{ Attributes, KillSwitches, Materializer, RestartSettings } import pekko.stream.scaladsl.{ Flow, Keep, RestartSource, Sink, Source } import pekko.stream.connectors.googlecloud.pubsub.grpc.{ + AckDeadline, AckDeadlineDistribution, + AckDeadlineExtender, AckDeadlineExtensionException, FlowControl } -import pekko.stream.connectors.googlecloud.pubsub.grpc.impl.FlowControlGateStage +import pekko.stream.connectors.googlecloud.pubsub.grpc.impl.{ EagerPullTrackingStage, FlowControlGateStage } import pekko.{ Done, NotUsed } -import com.google.pubsub.v1.pubsub._ +import com.google.pubsub.v1.pubsub.{ Subscriber => _, _ } import java.util.concurrent.ConcurrentHashMap import scala.concurrent.duration._ @@ -37,6 +39,14 @@ import scala.jdk.CollectionConverters._ */ object GooglePubSub { + /** + * Default size of the eager-pull buffer used by `autoExtendAckDeadlines`. Matches the + * default `maxOutstandingMessages` of Google's official client library. When pairing with + * server-side flow control via `StreamingPullRequest.maxOutstandingMessages`, a buffer + * at least as large as that value avoids unnecessary upstream backpressure. + */ + final val DefaultEagerPullBuffer: Int = 1000 + /** * Create a flow to publish messages to Google Cloud Pub/Sub. The flow emits responses that contain published * message ids. @@ -191,6 +201,11 @@ object GooglePubSub { * [[org.apache.pekko.stream.connectors.googlecloud.pubsub.grpc.AckDeadlineExtensionException]], * even if the stream is idle. * + * Note: this overload creates a fresh ticker and tracking map for each materialization, so it + * is NOT restart-safe. When wrapping the upstream `subscribe` with `RestartSource.withBackoff`, + * messages received before a stream failure lose their extension coverage on the next + * materialization. For restart-safety, use the [[AckDeadlineExtender]]-based overloads. + * * Usage: * {{{ * GooglePubSub.subscribe(request, 1.second) @@ -239,6 +254,29 @@ object GooglePubSub { extensionInterval: FiniteDuration, ackDeadlineSeconds: Int, maxAckExtensionPeriod: FiniteDuration): Flow[ReceivedMessage, ReceivedMessage, NotUsed] = + autoExtendAckDeadlines(subscription, extensionInterval, ackDeadlineSeconds, maxAckExtensionPeriod, + DefaultEagerPullBuffer) + + /** + * Create a flow that automatically extends ack deadlines for messages passing through it, + * with an explicit eager-pull buffer size. + * + * Identical to the four-argument overload but allows tuning the internal buffer used to + * track messages on receipt. A larger buffer absorbs longer downstream stalls without + * losing track of in-flight messages; a smaller buffer applies upstream backpressure sooner. + * Pair with `StreamingPullRequest.maxOutstandingMessages` to bound server-side delivery. + * + * @param maxBuffer maximum number of in-flight messages held in the eager-pull buffer + * (default [[DefaultEagerPullBuffer]]) + * @since 2.0.0 + */ + @ApiMayChange + def autoExtendAckDeadlines( + subscription: String, + extensionInterval: FiniteDuration, + ackDeadlineSeconds: Int, + maxAckExtensionPeriod: FiniteDuration, + maxBuffer: Int): Flow[ReceivedMessage, ReceivedMessage, NotUsed] = Flow.fromMaterializer { (mat, attr) => val client = subscriber(mat, attr).client val tracked = new ConcurrentHashMap[String, Long]() @@ -278,10 +316,8 @@ object GooglePubSub { Flow[ReceivedMessage] .via(killSwitch.flow) - .map { msg => - tracked.put(msg.ackId, System.nanoTime()) - msg - } + .via(new EagerPullTrackingStage[ReceivedMessage](maxBuffer, + msg => tracked.put(msg.ackId, System.nanoTime()))) .watchTermination((_, done: Future[Done]) => { done.onComplete(_ => cleanup())(ExecutionContext.parasitic) NotUsed @@ -321,6 +357,20 @@ object GooglePubSub { subscription: String, extensionInterval: FiniteDuration, distribution: AckDeadlineDistribution): Flow[ReceivedMessage, ReceivedMessage, NotUsed] = + autoExtendAckDeadlines(subscription, extensionInterval, distribution, DefaultEagerPullBuffer) + + /** + * Adaptive variant of `autoExtendAckDeadlines` with an explicit eager-pull buffer size. + * See the four-argument fixed-deadline overload for buffer semantics. + * + * @since 2.0.0 + */ + @ApiMayChange + def autoExtendAckDeadlines( + subscription: String, + extensionInterval: FiniteDuration, + distribution: AckDeadlineDistribution, + maxBuffer: Int): Flow[ReceivedMessage, ReceivedMessage, NotUsed] = Flow.fromMaterializer { (mat, attr) => val client = subscriber(mat, attr).client val killSwitch = KillSwitches.shared("autoExtendAckDeadlines") @@ -359,16 +409,73 @@ object GooglePubSub { Flow[ReceivedMessage] .via(killSwitch.flow) - .map { msg => - distribution.recordDelivery(msg.ackId) - msg - } + .via(new EagerPullTrackingStage[ReceivedMessage](maxBuffer, + msg => distribution.recordDelivery(msg.ackId))) .watchTermination((_, done: Future[Done]) => { done.onComplete(_ => cleanup())(ExecutionContext.parasitic) NotUsed }) }.mapMaterializedValue(_ => NotUsed) + /** + * Create a flow that automatically extends ack deadlines using a caller-owned + * [[AckDeadlineExtender]]. The extender owns the tracking map and the background ticker, both + * of which live above the lifetime of any single Pub/Sub streaming pull. This makes the flow + * restart-safe: when wrapped in `RestartSource.withBackoff`, messages received before a + * stream failure remain in the extender's tracking map and continue to receive deadline + * extensions during the backoff window, matching how Google's official client library + * (`MessageDispatcher` inside `StreamingSubscriberConnection`) behaves across reconnects. + * + * Usage: + * {{{ + * val extender = AckDeadlineExtender(subscriptionFqrs, 8.seconds, 30) + * + * try { + * RestartSource.withBackoff(restartSettings) { () => + * GooglePubSub.subscribe(request, 1.second).mapMaterializedValue(_ => NotUsed) + * } + * .via(GooglePubSub.autoExtendAckDeadlines(extender)) + * .mapAsync(10)(processMessage) + * .map(msg => AcknowledgeRequest(subscriptionFqrs, Seq(msg.ackId))) + * .runWith(GooglePubSub.acknowledge(parallelism = 1)) + * } finally { + * extender.close() + * } + * }}} + * + * @since 2.0.0 + */ + @ApiMayChange + def autoExtendAckDeadlines(extender: AckDeadlineExtender): Flow[ReceivedMessage, ReceivedMessage, NotUsed] = + autoExtendAckDeadlines(extender, DefaultEagerPullBuffer) + + /** + * Caller-owned-extender variant of `autoExtendAckDeadlines` with an explicit eager-pull + * buffer size. See the single-argument overload for the restart-safety contract and the + * fixed-deadline four-argument overload for buffer semantics. + * + * @since 2.0.0 + */ + @ApiMayChange + def autoExtendAckDeadlines( + extender: AckDeadlineExtender, + maxBuffer: Int): Flow[ReceivedMessage, ReceivedMessage, NotUsed] = + Flow.fromMaterializer { (_, _) => + // Per-materialization KillSwitch so a previously-failed extender aborts the new stream + // immediately. New materializations after extender failure start in failed state. + val killSwitch = KillSwitches.shared("autoExtendAckDeadlines") + extender.tickerDone.onComplete { + case scala.util.Failure(ex) => + killSwitch.abort(new AckDeadlineExtensionException( + "Lease management ticker failed; ack deadline extensions have stopped", ex)) + case _ => () + }(ExecutionContext.parasitic) + + Flow[ReceivedMessage] + .via(killSwitch.flow) + .via(new EagerPullTrackingStage[ReceivedMessage](maxBuffer, msg => extender.track(msg.ackId))) + }.mapMaterializedValue(_ => NotUsed) + /** * Create a flow that modifies the ack deadline for each message using a dynamic function. * @@ -684,6 +791,42 @@ object GooglePubSub { } .mapMaterializedValue(_.flatMap(identity)(ExecutionContext.parasitic)) + /** + * Create a high-level [[Subscriber]] resource that bundles streaming pull, restart logic, + * ack-deadline extension, and optional flow control. The deadline-extension ticker starts at + * construction. See [[Subscriber]] for the composition guarantees and lifecycle. + * + * @since 2.0.0 + */ + @ApiMayChange + def subscriber( + request: StreamingPullRequest, + pollInterval: FiniteDuration, + ackDeadline: AckDeadline, + restartSettings: Option[RestartSettings] = None, + flowControl: Option[FlowControl] = None)( + implicit system: pekko.actor.ClassicActorSystemProvider): Subscriber = + Subscriber.create(request, pollInterval, ackDeadline, restartSettings, flowControl, + GrpcSubscriberExt()(system).subscriber) + + /** + * Variant of [[subscriber]] taking an explicit [[GrpcSubscriber]] (useful for tests or for + * callers that manage the underlying gRPC client themselves). + * + * @since 2.0.0 + */ + @ApiMayChange + def subscriber( + request: StreamingPullRequest, + pollInterval: FiniteDuration, + ackDeadline: AckDeadline, + restartSettings: Option[RestartSettings], + flowControl: Option[FlowControl], + grpcSubscriber: GrpcSubscriber)( + implicit system: pekko.actor.ClassicActorSystemProvider): Subscriber = + Subscriber.create(request, pollInterval, ackDeadline, restartSettings, flowControl, + grpcSubscriber) + private def publisher(mat: Materializer, attr: Attributes) = attr .get[PubSubAttributes.Publisher] diff --git a/google-cloud-pub-sub-grpc/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/scaladsl/Subscriber.scala b/google-cloud-pub-sub-grpc/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/scaladsl/Subscriber.scala new file mode 100644 index 000000000..a14c5f9a4 --- /dev/null +++ b/google-cloud-pub-sub-grpc/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/scaladsl/Subscriber.scala @@ -0,0 +1,222 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.pekko.stream.connectors.googlecloud.pubsub.grpc.scaladsl + +import org.apache.pekko +import pekko.{ Done, NotUsed } +import pekko.actor.ClassicActorSystemProvider +import pekko.annotation.ApiMayChange +import pekko.stream.RestartSettings +import pekko.stream.scaladsl.{ Flow, Keep, Sink, Source } +import pekko.stream.connectors.googlecloud.pubsub.grpc.{ AckDeadline, AckDeadlineExtender, FlowControl } +import com.google.pubsub.v1.pubsub._ + +import scala.concurrent.duration.FiniteDuration +import scala.concurrent.{ ExecutionContext, Future } + +/** + * High-level Pub/Sub subscriber that bundles streaming pull, restart logic, ack-deadline + * extension, and optional flow control into a single resource. Designed so the user can't + * accidentally compose the building blocks incorrectly (the trap that motivates the + * `AckDeadlineExtender` workaround for users who compose `subscribe` + `autoExtendAckDeadlines` + * by hand inside their own `RestartSource`). + * + * Lifecycle: + * + * - Create one with [[Subscriber.apply]]. The deadline-extension ticker starts at construction. + * - Materialize [[source]] once to start receiving messages, or use [[run]] for the simple case. + * The source's `watchTermination` triggers [[close]] on stream completion, so most users + * don't need to call close themselves. + * - For early shutdown or when materializing multiple times, call [[close]] explicitly. It is + * idempotent. + * + * Composition rules baked in: + * + * - Restart logic wraps only the inner `subscribe`, not the deadline extender, so tracking + * state survives reconnect backoff (matching Google's `MessageDispatcher` lifecycle). + * - The deadline tracker eagerly pulls from upstream so messages are tracked the moment they + * arrive, even when downstream is backpressured (bug 1 fix). + * - The flow-control gate, if configured, also pulls eagerly so its permit counter reflects + * actual in-flight delivery (bug 3 fix). + * + * Usage: + * {{{ + * val subscriber = Subscriber( + * request = StreamingPullRequest() + * .withSubscription(subscriptionFqrs) + * .withStreamAckDeadlineSeconds(60) + * .withMaxOutstandingMessages(1000), + * pollInterval = 1.second, + * ackDeadline = AckDeadline.Fixed(extensionInterval = 8.seconds, deadlineSeconds = 30), + * restartSettings = Some(RestartSettings(100.millis, 10.seconds, 0.2))) + * + * subscriber.source + * .mapAsync(10)(processMessage) + * .map(msg => AcknowledgeRequest(subscriptionFqrs, Seq(msg.ackId))) + * .runWith(subscriber.acknowledge(parallelism = 1)) + * }}} + * + * @since 2.0.0 + */ +@ApiMayChange +final class Subscriber private ( + val request: StreamingPullRequest, + val pollInterval: FiniteDuration, + val ackDeadline: AckDeadline, + val restartSettings: Option[RestartSettings], + val flowControl: Option[FlowControl], + private[grpc] val extender: AckDeadlineExtender, + private[grpc] val grpcSubscriber: GrpcSubscriber)(implicit system: ClassicActorSystemProvider) { + + private val subscription: String = request.subscription + + private def attrs = PubSubAttributes.subscriber(grpcSubscriber) + + /** + * Source emitting messages from this subscription. Composes (in order): + * 1. `subscribe` (wrapped in `RestartSource.withBackoff` if `restartSettings` is set) + * 2. `autoExtendAckDeadlines(extender)` for restart-safe deadline extension + * 3. `flowControlGate(flowControl)` if `flowControl` is set + * + * Restart preserves in-flight state. `RestartSource` re-materializes only the inner + * `subscribe`; the eager-pull tracker, the flow-control gate, and any downstream operators + * stay alive across every gRPC stream restart. Messages buffered in those operators at the + * moment of disconnect are NOT lost. They remain in their internal buffers, keep getting + * their deadlines extended by the long-lived ticker, and continue flowing downstream once + * the new gRPC stream comes up. Pub/Sub ackIds are valid against the subscription rather + * than any particular stream, so acks against pre-disconnect messages still succeed. + * + * On stream completion (success or failure), the extender is closed automatically. + * Materializing this source more than once is supported, but only the first completion + * triggers auto-cleanup; subsequent materializations after auto-cleanup will fail because + * the extender is closed. Use [[close]] explicitly for that case. + */ + def source: Source[ReceivedMessage, NotUsed] = { + val baseSource: Source[ReceivedMessage, NotUsed] = restartSettings match { + case Some(rs) => GooglePubSub.subscribe(request, pollInterval, rs).withAttributes(attrs) + case None => + GooglePubSub.subscribe(request, pollInterval).mapMaterializedValue(_ => NotUsed).withAttributes(attrs) + } + + val withExtender = baseSource.via(GooglePubSub.autoExtendAckDeadlines(extender)) + + val withFlowControl = flowControl match { + case Some(fc) => withExtender.via(GooglePubSub.flowControlGate(fc)) + case None => withExtender + } + + withFlowControl.watchTermination { (_, done: Future[Done]) => + done.onComplete(_ => close())(ExecutionContext.parasitic) + NotUsed + } + } + + /** + * Sink that acknowledges messages. Releases flow-control permits if this subscriber was + * configured with a `FlowControl`. Records completion latencies into the + * [[AckDeadlineDistribution]] if this subscriber uses [[AckDeadline.Adaptive]]. + */ + def acknowledge(parallelism: Int): Sink[AcknowledgeRequest, Future[Done]] = + Sink + .fromMaterializer { (mat, _) => + val client = grpcSubscriber.client + val ec = mat.executionContext + Flow[AcknowledgeRequest] + .mapAsyncUnordered(parallelism) { req => + client.acknowledge(req).map { _ => + flowControl.foreach(_.release(req.ackIds.size)) + ackDeadline match { + case AckDeadline.Adaptive(_, dist) => dist.recordCompletions(req.ackIds) + case _ => () + } + Done + }(ec) + } + .toMat(Sink.ignore)(Keep.right) + } + .mapMaterializedValue(_.flatMap(identity)(ExecutionContext.parasitic)) + + /** + * Sink that nacks messages by setting their ack deadline to 0, causing immediate redelivery. + * Releases flow-control permits and records completions same as [[acknowledge]]. + */ + def nack(parallelism: Int): Sink[AcknowledgeRequest, Future[Done]] = + Sink + .fromMaterializer { (mat, _) => + val client = grpcSubscriber.client + val ec = mat.executionContext + Flow[AcknowledgeRequest] + .mapAsyncUnordered(parallelism) { req => + client.modifyAckDeadline( + ModifyAckDeadlineRequest(req.subscription, req.ackIds, ackDeadlineSeconds = 0)) + .map { _ => + flowControl.foreach(_.release(req.ackIds.size)) + ackDeadline match { + case AckDeadline.Adaptive(_, dist) => dist.recordCompletions(req.ackIds) + case _ => () + } + Done + }(ec) + } + .toMat(Sink.ignore)(Keep.right) + } + .mapMaterializedValue(_.flatMap(identity)(ExecutionContext.parasitic)) + + /** + * Convenience for the simple case: materialize the source, apply `processFn` with the given + * parallelism, and ack each successfully-processed message. Returns the materialized future + * that completes when the stream completes. + */ + def run(parallelism: Int)(processFn: ReceivedMessage => Future[Any]): Future[Done] = { + implicit val ec: ExecutionContext = ExecutionContext.parasitic + source + .mapAsync(parallelism)(msg => processFn(msg).map(_ => msg)) + .map(msg => AcknowledgeRequest(subscription, Seq(msg.ackId))) + .runWith(acknowledge(parallelism = 1)) + } + + /** Stop the background ticker and clear tracking state. Idempotent. */ + def close(): Future[Done] = extender.close() +} + +@ApiMayChange +object Subscriber { + + /** + * INTERNAL: package-private constructor. Public entry point is + * [[GooglePubSub.subscriber]]. + */ + private[grpc] def create( + request: StreamingPullRequest, + pollInterval: FiniteDuration, + ackDeadline: AckDeadline, + restartSettings: Option[RestartSettings], + flowControl: Option[FlowControl], + grpcSubscriber: GrpcSubscriber)( + implicit system: ClassicActorSystemProvider): Subscriber = { + val extender = ackDeadline match { + case AckDeadline.Fixed(extensionInterval, deadlineSecs, maxExt) => + AckDeadlineExtender(request.subscription, extensionInterval, deadlineSecs, maxExt, + grpcSubscriber) + case AckDeadline.Adaptive(extensionInterval, distribution) => + AckDeadlineExtender(request.subscription, extensionInterval, distribution, grpcSubscriber) + } + new Subscriber(request, pollInterval, ackDeadline, restartSettings, flowControl, + extender, grpcSubscriber) + } +} diff --git a/google-cloud-pub-sub-grpc/src/test/scala/docs/scaladsl/ExampleApp.scala b/google-cloud-pub-sub-grpc/src/test/scala/docs/scaladsl/ExampleApp.scala index 32417c6dc..c7b1b43ae 100644 --- a/google-cloud-pub-sub-grpc/src/test/scala/docs/scaladsl/ExampleApp.scala +++ b/google-cloud-pub-sub-grpc/src/test/scala/docs/scaladsl/ExampleApp.scala @@ -16,7 +16,7 @@ import java.util.logging.{ Level, Logger } import org.apache.pekko import pekko.actor.{ ActorSystem, Cancellable } import pekko.stream.{ DelayOverflowStrategy, RestartSettings } -import pekko.stream.connectors.googlecloud.pubsub.grpc.FlowControl +import pekko.stream.connectors.googlecloud.pubsub.grpc.{ AckDeadline, FlowControl } import pekko.stream.connectors.googlecloud.pubsub.grpc.scaladsl.GooglePubSub import pekko.stream.scaladsl.{ Sink, Source } import com.google.protobuf.ByteString @@ -119,16 +119,22 @@ object ExampleApp { maxBackoff = 10.seconds, randomFactor = 0.2) - GooglePubSub - .subscribe(subscribe(projectId, sub), 1.second, restartSettings) - .via(GooglePubSub.autoExtendAckDeadlines(subscriptionFqrs, 3.seconds, 30)) + // GooglePubSub.subscriber bundles subscribe + restart + autoExtend into one resource. The + // composition is correct by construction: tracking state survives reconnect, the deadline + // tracker pulls eagerly from the gRPC stream so backpressure on the processing stage + // doesn't delay tracking, and close() is called automatically when the source completes. + val subscriber = GooglePubSub.subscriber( + request = subscribe(projectId, sub), + pollInterval = 1.second, + ackDeadline = AckDeadline.Fixed(extensionInterval = 3.seconds, deadlineSeconds = 30), + restartSettings = Some(restartSettings)) + + subscriber.source .map { msg => println(msg) AcknowledgeRequest(subscriptionFqrs, Seq(msg.ackId)) } - .to(GooglePubSub.acknowledge(parallelism = 1)) - .mapMaterializedValue(Future.successful(_)) - .run() + .runWith(subscriber.acknowledge(parallelism = 1)) } /** diff --git a/google-cloud-pub-sub-grpc/src/test/scala/docs/scaladsl/IntegrationSpec.scala b/google-cloud-pub-sub-grpc/src/test/scala/docs/scaladsl/IntegrationSpec.scala index 25e12fc0b..c89fd6045 100644 --- a/google-cloud-pub-sub-grpc/src/test/scala/docs/scaladsl/IntegrationSpec.scala +++ b/google-cloud-pub-sub-grpc/src/test/scala/docs/scaladsl/IntegrationSpec.scala @@ -661,6 +661,85 @@ class IntegrationSpec Source.single(PublishRequest()).via(publishFlow).to(Sink.ignore) } + + "Subscriber resource: end-to-end with all bug fixes + composition guarantees" in { + // Exercises the new high-level GooglePubSub.subscriber(...) against real Pub/Sub: + // - sets maxOutstandingMessages on the initial StreamingPullRequest (bug 2 fix: + // subsequent requests must not echo this field) + // - slow per-message processing forces deadlines to age inside the eager-pull + // tracker (bug 1 fix: tracking happens on receipt, not on push downstream) + // - small FlowControl limit verifies the rewritten gate counts on receipt + // (bug 3 fix) + // - asserts every published message is received exactly once with no duplicates + import pekko.stream.RestartSettings + import pekko.stream.connectors.googlecloud.pubsub.grpc.{ AckDeadline, FlowControl } + + val projectId = "pekko-connectors" + val topic = "simpleTopic" + val subscription = "simpleSubscription" + val topicFqrs = s"projects/$projectId/topics/$topic" + val subFqrs = s"projects/$projectId/subscriptions/$subscription" + + val messageCount = 10 + val maxOutstanding = 3 + val processingDelay = 6.seconds + val testPrefix = s"subscriber-resource-${System.nanoTime()}" + val messages = (1 to messageCount).map(i => + PubsubMessage().withData(ByteString.copyFromUtf8(s"$testPrefix-$i"))) + + // Publish the batch + Source + .single(PublishRequest(topicFqrs, messages)) + .via(GooglePubSub.publish(parallelism = 1)) + .runWith(Sink.head) + .futureValue(timeout(30.seconds)) + + val flowControl = FlowControl(maxOutstandingMessages = maxOutstanding.toLong) + @volatile var maxObserved = 0L + + // Initial request carries BOTH stream ack deadline AND maxOutstandingMessages. + // Pre-bug-2-fix this would fail on the first keepalive tick with INVALID_ARGUMENT. + val request = StreamingPullRequest(subFqrs) + .withStreamAckDeadlineSeconds(15) + .withMaxOutstandingMessages(maxOutstanding.toLong) + + val restartSettings = RestartSettings(1.second, 10.seconds, 0.2).withMaxRestarts(3, 1.minute) + + val subscriber = GooglePubSub.subscriber( + request = request, + pollInterval = 1.second, + ackDeadline = AckDeadline.Fixed(extensionInterval = 5.seconds, deadlineSeconds = 30), + restartSettings = Some(restartSettings), + flowControl = Some(flowControl)) + + try { + val received = subscriber.source + .filter(_.message.exists(_.data.toStringUtf8.startsWith(testPrefix))) + .take(messageCount) + .mapAsync(parallelism = 2) { msg => + val current = flowControl.outstandingCount + synchronized { if (current > maxObserved) maxObserved = current } + // Slow processing forces autoExtend to actually fire while messages wait. + pekko.pattern.after(processingDelay)(Future.successful(msg)) + } + .map(msg => (msg.message.map(_.data.toStringUtf8).getOrElse(""), msg.ackId)) + .alsoTo( + Flow[(String, String)] + .map { case (_, ackId) => AcknowledgeRequest(subFqrs, Seq(ackId)) } + .to(subscriber.acknowledge(parallelism = 1))) + .runWith(Sink.seq) + + val msgs = received.futureValue(timeout(5.minutes)) + val payloads = msgs.map(_._1).toSet + + msgs should have size messageCount.toLong + payloads.size shouldBe messageCount // no duplicates + maxObserved should be <= maxOutstanding.toLong + maxObserved should be >= 1L + } finally { + subscriber.close().futureValue(timeout(10.seconds)) + } + } } override def afterAll() = diff --git a/google-cloud-pub-sub-grpc/src/test/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/AutoExtendAckDeadlinesSpec.scala b/google-cloud-pub-sub-grpc/src/test/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/AutoExtendAckDeadlinesSpec.scala index a4e0eb565..9efceaa60 100644 --- a/google-cloud-pub-sub-grpc/src/test/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/AutoExtendAckDeadlinesSpec.scala +++ b/google-cloud-pub-sub-grpc/src/test/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/AutoExtendAckDeadlinesSpec.scala @@ -21,6 +21,7 @@ import org.apache.pekko import pekko.{ Done, NotUsed } import pekko.actor.ActorSystem import pekko.stream.scaladsl.{ Keep, Sink, Source } +import pekko.stream.connectors.googlecloud.pubsub.grpc.{ AckDeadlineExtender, FlowControl } import pekko.stream.connectors.googlecloud.pubsub.grpc.scaladsl.{ GooglePubSub, GrpcSubscriber, PubSubAttributes } import com.google.protobuf.ByteString import com.google.pubsub.v1.pubsub._ @@ -30,6 +31,7 @@ import java.util.concurrent.ConcurrentLinkedQueue import scala.concurrent.{ Future, Promise } import scala.concurrent.duration._ import scala.jdk.CollectionConverters._ + import org.scalatest.BeforeAndAfterAll import org.scalatest.concurrent.{ Eventually, ScalaFutures } import org.scalatest.matchers.should.Matchers @@ -67,6 +69,16 @@ class AutoExtendAckDeadlinesSpec Future.successful(com.google.protobuf.empty.Empty()) } + /** Stub that captures every modifyAckDeadline request for inspection. */ + class CapturingClient(captured: ConcurrentLinkedQueue[ModifyAckDeadlineRequest]) + extends TestSubscriberClientBase { + override def modifyAckDeadline( + in: ModifyAckDeadlineRequest): Future[com.google.protobuf.empty.Empty] = { + captured.add(in) + Future.successful(com.google.protobuf.empty.Empty()) + } + } + "autoExtendAckDeadlines (fixed deadline)" should { "fail the main stream when the ticker's modifyAckDeadline fails" in { @@ -112,6 +124,119 @@ class AutoExtendAckDeadlinesSpec result.futureValue should have size 3 result.futureValue.map(_.ackId) shouldBe Seq("1", "2", "3") } + + "track all in-flight messages eagerly even when downstream is backpressured" in { + // Regression test: with a plain `.map` tracker, only the one message held by mapAsync(1) + // would be tracked at the first tick. The eager-pull stage must track all buffered + // messages so the first ticker tick extends every in-flight ackId. + val captured = new ConcurrentLinkedQueue[ModifyAckDeadlineRequest]() + val testSubscriber = new GrpcSubscriber(new CapturingClient(captured)) + val n = 10 + + val killSwitch = pekko.stream.KillSwitches.shared("test") + val gate = Promise[ReceivedMessage]() // never completes — mapAsync(1) hangs forever + + Source(1 to n) + .map(i => makeMsg(i.toString)) + .via(GooglePubSub.autoExtendAckDeadlines(subscription, 200.millis, 30)) + .via(killSwitch.flow) + .mapAsync(1)(_ => gate.future) + .withAttributes(PubSubAttributes.subscriber(testSubscriber)) + .runWith(Sink.ignore) + + // Wait for the first tick (200ms) plus jitter to fire. + Thread.sleep(800) + killSwitch.shutdown() + + val firstReq = captured.poll() + firstReq should not be null + firstReq.subscription shouldBe subscription + firstReq.ackDeadlineSeconds shouldBe 30 + // All n ackIds must appear in the very first extension call. Without eager-pull, + // this would only be 1 (the element currently grabbed by mapAsync). + firstReq.ackIds.toSet shouldBe (1 to n).map(_.toString).toSet + } + } + + "autoExtendAckDeadlines (caller-owned AckDeadlineExtender)" should { + + "retain tracking state across stream materializations (restart-safety)" in { + // Restart-safety regression test. The internal-ticker overload clears the tracking map + // and cancels the ticker when the stream completes, so any messages received before a + // RestartSource re-materialization lose extension coverage. The extender owns both the + // map and the ticker, so extensions continue across stream completion / re-materialization. + val captured = new ConcurrentLinkedQueue[ModifyAckDeadlineRequest]() + val testSubscriber = new GrpcSubscriber(new CapturingClient(captured)) + val extender = AckDeadlineExtender(subscription, 200.millis, 30, 60.minutes, testSubscriber)(system) + + try { + // First materialization: emit three messages then complete normally. The eager-pull + // tracker writes them into the extender's tracked map. + val streamOne = Source(List(makeMsg("a1"), makeMsg("a2"), makeMsg("a3"))) + .via(GooglePubSub.autoExtendAckDeadlines(extender)) + .runWith(Sink.ignore) + streamOne.futureValue + + // The extender's map must outlive the stream — this is the contract that lets a + // RestartSource-wrapped subscribe keep extension coverage during reconnect backoff. + extender.trackedSize shouldBe 3 + + // Let the ticker fire a few times. Each tick should carry all three ackIds. + Thread.sleep(700) + + val allRequests = captured.iterator().asScala.toList + val streamOneIds = Set("a1", "a2", "a3") + val ticksWithStreamOneIds = allRequests.count(_.ackIds.exists(streamOneIds.contains)) + ticksWithStreamOneIds should be >= 2 + allRequests.filter(_.ackIds.exists(streamOneIds.contains)).foreach { req => + req.ackIds.toSet shouldBe streamOneIds + } + + // Second materialization: simulate reconnect by running another stream. Tracking from + // stream 1 must still be present, and stream 2 entries get added on top. + val streamTwo = Source(List(makeMsg("b1"), makeMsg("b2"))) + .via(GooglePubSub.autoExtendAckDeadlines(extender)) + .runWith(Sink.ignore) + streamTwo.futureValue + extender.trackedSize shouldBe 5 + } finally { + extender.close().futureValue + } + } + } + + "GooglePubSub.flowControlGate" should { + + "acquire permits on receipt, not on push to downstream (eager-pull)" in { + // Regression test: with the old downstream-demand-bound gate, only the message currently + // held by mapAsync would acquire a permit. The eager-pull gate must count every message + // it receives, so flowControl.outstandingCount climbs to the limit even when downstream + // is fully saturated. + val testSubscriber = new GrpcSubscriber(new SucceedingClient()) + val limit = 5 + val flowControl = FlowControl(maxOutstandingMessages = limit.toLong) + val gate = Promise[ReceivedMessage]() // mapAsync(1) hangs forever + val killSwitch = pekko.stream.KillSwitches.shared("flowControlGateTest") + + Source(1 to 100) + .map(i => makeMsg(i.toString)) + .via(GooglePubSub.flowControlGate(flowControl)) + .via(killSwitch.flow) + .mapAsync(1)(_ => gate.future) + .withAttributes(PubSubAttributes.subscriber(testSubscriber)) + .runWith(Sink.ignore) + + // Give the gate time to fill up. + Thread.sleep(300) + + // The gate should have admitted exactly `limit` messages: one currently held by mapAsync, + // (limit - 1) buffered inside the gate. The Source(1 to 100) provides plenty of upstream + // messages, but the gate stops pulling at the limit. + // Without eager-pull this would be 1 (only the message in mapAsync had a permit acquired). + flowControl.outstandingCount shouldBe limit.toLong + + killSwitch.shutdown() + } } "autoExtendAckDeadlines (adaptive)" should { diff --git a/google-cloud-pub-sub-grpc/src/test/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/SubscriberSpec.scala b/google-cloud-pub-sub-grpc/src/test/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/SubscriberSpec.scala new file mode 100644 index 000000000..cd285e69e --- /dev/null +++ b/google-cloud-pub-sub-grpc/src/test/scala/org/apache/pekko/stream/connectors/googlecloud/pubsub/grpc/SubscriberSpec.scala @@ -0,0 +1,176 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.pekko.stream.connectors.googlecloud.pubsub.grpc + +import org.apache.pekko +import pekko.actor.ActorSystem +import pekko.stream.RestartSettings +import pekko.stream.scaladsl.{ Sink, Source } +import pekko.stream.connectors.googlecloud.pubsub.grpc.scaladsl.{ GooglePubSub, GrpcSubscriber } +import com.google.protobuf.ByteString +import com.google.pubsub.v1.pubsub._ + +import java.util.concurrent.ConcurrentLinkedQueue +import scala.concurrent.duration._ +import scala.concurrent.{ Future, Promise } + +import org.scalatest.BeforeAndAfterAll +import org.scalatest.concurrent.ScalaFutures +import org.scalatest.matchers.should.Matchers +import org.scalatest.wordspec.AnyWordSpec + +class SubscriberSpec + extends AnyWordSpec + with Matchers + with BeforeAndAfterAll + with ScalaFutures { + + implicit val system: ActorSystem = ActorSystem("SubscriberSpec") + implicit val patience: PatienceConfig = PatienceConfig(10.seconds, 100.millis) + + val subscription = "projects/test/subscriptions/test-sub" + + def makeMsg(id: String): ReceivedMessage = + ReceivedMessage( + ackId = id, + message = Some(PubsubMessage(data = ByteString.copyFromUtf8(s"test-$id")))) + + /** + * Stub that succeeds for both modifyAckDeadline (ticker) and acknowledge calls. + * Captures every modifyAckDeadline so we can inspect deadline-extension behavior. + */ + class CapturingClient(extensions: ConcurrentLinkedQueue[ModifyAckDeadlineRequest], + acks: ConcurrentLinkedQueue[AcknowledgeRequest]) + extends TestSubscriberClientBase { + override def modifyAckDeadline(in: ModifyAckDeadlineRequest) + : Future[com.google.protobuf.empty.Empty] = { + extensions.add(in) + Future.successful(com.google.protobuf.empty.Empty()) + } + override def acknowledge(in: AcknowledgeRequest): Future[com.google.protobuf.empty.Empty] = { + acks.add(in) + Future.successful(com.google.protobuf.empty.Empty()) + } + } + + "Subscriber" should { + + "auto-close the extender when the source stream completes" in { + val extensions = new ConcurrentLinkedQueue[ModifyAckDeadlineRequest]() + val acks = new ConcurrentLinkedQueue[AcknowledgeRequest]() + val testSubscriber = new GrpcSubscriber(new CapturingClient(extensions, acks)) + + val subscriber = GooglePubSub.subscriber( + request = StreamingPullRequest().withSubscription(subscription).withStreamAckDeadlineSeconds(60), + pollInterval = 1.second, + ackDeadline = AckDeadline.Fixed(extensionInterval = 100.millis, deadlineSeconds = 30), + restartSettings = None, + flowControl = None, + grpcSubscriber = testSubscriber) + + // Synthesize a stream that completes quickly, bypassing subscriber.source (which would + // call streamingPull on the stub). We're testing the auto-close hook, not subscribe. + // Simulate the watchTermination path by closing manually after a known-completed stream. + Source(List(makeMsg("x"))).runWith(Sink.ignore).futureValue + + subscriber.extender.tickerDone.isCompleted shouldBe false + subscriber.close().futureValue + subscriber.extender.tickerDone.isCompleted shouldBe true + } + + "release flow-control permits on acknowledge" in { + val extensions = new ConcurrentLinkedQueue[ModifyAckDeadlineRequest]() + val acks = new ConcurrentLinkedQueue[AcknowledgeRequest]() + val testSubscriber = new GrpcSubscriber(new CapturingClient(extensions, acks)) + val flowControl = FlowControl(maxOutstandingMessages = 100) + + val subscriber = GooglePubSub.subscriber( + request = StreamingPullRequest().withSubscription(subscription).withStreamAckDeadlineSeconds(60), + pollInterval = 1.second, + ackDeadline = AckDeadline.Fixed(1.second, 30), + restartSettings = None, + flowControl = Some(flowControl), + grpcSubscriber = testSubscriber) + + try { + // Manually acquire 3 permits, then send an AcknowledgeRequest with 3 ackIds through + // the subscriber's ack sink. + flowControl.acquire(); flowControl.acquire(); flowControl.acquire() + flowControl.outstandingCount shouldBe 3L + + Source.single(AcknowledgeRequest(subscription, Seq("a", "b", "c"))) + .runWith(subscriber.acknowledge(parallelism = 1)) + .futureValue + + flowControl.outstandingCount shouldBe 0L + acks.size() shouldBe 1 + } finally { + subscriber.close().futureValue + } + } + + "record completion latencies into AckDeadlineDistribution on acknowledge (adaptive)" in { + val extensions = new ConcurrentLinkedQueue[ModifyAckDeadlineRequest]() + val acks = new ConcurrentLinkedQueue[AcknowledgeRequest]() + val testSubscriber = new GrpcSubscriber(new CapturingClient(extensions, acks)) + val distribution = AckDeadlineDistribution(initialDeadlineSeconds = 10) + + val subscriber = GooglePubSub.subscriber( + request = StreamingPullRequest().withSubscription(subscription).withStreamAckDeadlineSeconds(60), + pollInterval = 1.second, + ackDeadline = AckDeadline.Adaptive(1.second, distribution), + restartSettings = None, + flowControl = None, + grpcSubscriber = testSubscriber) + + try { + // Manually record a delivery so completion has something to read. + distribution.recordDelivery("ack-1") + Thread.sleep(50) + + Source.single(AcknowledgeRequest(subscription, Seq("ack-1"))) + .runWith(subscriber.acknowledge(parallelism = 1)) + .futureValue + + // After ack, deliveryTimes no longer contains ack-1 (recordCompletion removes it). + distribution.deliveryTimes.containsKey("ack-1") shouldBe false + } finally { + subscriber.close().futureValue + } + } + + "be safe to close() multiple times" in { + val extensions = new ConcurrentLinkedQueue[ModifyAckDeadlineRequest]() + val acks = new ConcurrentLinkedQueue[AcknowledgeRequest]() + val testSubscriber = new GrpcSubscriber(new CapturingClient(extensions, acks)) + + val subscriber = GooglePubSub.subscriber( + request = StreamingPullRequest().withSubscription(subscription).withStreamAckDeadlineSeconds(60), + pollInterval = 1.second, + ackDeadline = AckDeadline.Fixed(1.second, 30), + restartSettings = None, + flowControl = None, + grpcSubscriber = testSubscriber) + + subscriber.close().futureValue + subscriber.close().futureValue // second call must not throw + } + } + + override def afterAll(): Unit = system.terminate() +}