This is an automated email from the ASF dual-hosted git repository.
bhulette pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 0943d85 [BEAM-14129] Restructure PubsubLiteIO Read side to produce
smaller bundles (#17125)
0943d85 is described below
commit 0943d85113e77505f2a1291d7b256b9a4b1e5dec
Author: dpcollins-google <40498610+dpcollins-goo...@users.noreply.github.com>
AuthorDate: Wed Mar 23 11:48:53 2022 -0400
[BEAM-14129] Restructure PubsubLiteIO Read side to produce smaller bundles
(#17125)
* [BEAM-14129] Restructure PubsubLiteIO Read side to produce smaller bundles
* [BEAM-14129] Restructure PubsubLiteIO Read side to produce smaller bundles
* fixes
* fixes
* fixes
* fixes
* fixes
* fixes
---
.../org/apache/beam/gradle/BeamModulePlugin.groovy | 8 +-
.../internal/MemoryBufferedSubscriber.java | 55 ++++++
.../internal/MemoryBufferedSubscriberImpl.java | 173 ++++++++++++++++++
...erverPublisherCache.java => MemoryLimiter.java} | 21 ++-
.../gcp/pubsublite/internal/MemoryLimiterImpl.java | 92 ++++++++++
.../internal/OffsetByteRangeTracker.java | 104 ++---------
.../internal/PerServerPublisherCache.java | 7 +-
...herCache.java => PerServerSubscriberCache.java} | 14 +-
.../internal/PerSubscriptionPartitionSdf.java | 21 ++-
.../io/gcp/pubsublite/internal/PubsubLiteSink.java | 4 +-
.../{PublisherCache.java => ServiceCache.java} | 54 +++---
.../pubsublite/internal/SubscribeTransform.java | 50 +++---
.../pubsublite/internal/SubscriberAssembler.java | 1 +
.../SubscriptionPartitionProcessorImpl.java | 180 +++++++------------
.../internal/MemoryBufferedSubscriberImplTest.java | 200 +++++++++++++++++++++
.../internal/OffsetByteRangeTrackerTest.java | 37 +---
.../SubscriptionPartitionProcessorImplTest.java | 180 ++++++++-----------
17 files changed, 786 insertions(+), 415 deletions(-)
diff --git
a/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy
b/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy
index eb48d48..6b57085 100644
--- a/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy
+++ b/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy
@@ -573,11 +573,11 @@ class BeamModulePlugin implements Plugin<Project> {
google_cloud_datastore_v1_proto_client :
"com.google.cloud.datastore:datastore-v1-proto-client:2.1.3",
google_cloud_firestore :
"com.google.cloud:google-cloud-firestore", //
google_cloud_platform_libraries_bom sets version
google_cloud_pubsub :
"com.google.cloud:google-cloud-pubsub", // google_cloud_platform_libraries_bom
sets version
- google_cloud_pubsublite :
"com.google.cloud:google-cloud-pubsublite", //
google_cloud_platform_libraries_bom sets version
+ google_cloud_pubsublite :
"com.google.cloud:google-cloud-pubsublite:1.5.0", // TODO(dpcollins-google):
Let google_cloud_platform_libraries_bom set version once high enough
// The GCP Libraries BOM dashboard shows the versions set by the BOM:
//
https://storage.googleapis.com/cloud-opensource-java-dashboard/com.google.cloud/libraries-bom/24.3.0/artifact_details.html
// Update libraries-bom version on
sdks/java/container/license_scripts/dep_urls_java.yaml
- google_cloud_platform_libraries_bom :
"com.google.cloud:libraries-bom:24.4.0",
+ google_cloud_platform_libraries_bom :
"com.google.cloud:libraries-bom:24.4.0", // TODO(updater): Remove the
pubsublite pin to 1.5.0 if the BOM includes a higher version
google_cloud_spanner :
"com.google.cloud:google-cloud-spanner", // google_cloud_platform_libraries_bom
sets version
google_cloud_spanner_test :
"com.google.cloud:google-cloud-spanner:$google_cloud_spanner_version:tests",
google_code_gson :
"com.google.code.gson:gson:$google_code_gson_version",
@@ -599,7 +599,7 @@ class BeamModulePlugin implements Plugin<Project> {
grpc_core : "io.grpc:grpc-core", //
google_cloud_platform_libraries_bom sets version
grpc_google_cloud_firestore_v1 :
"com.google.api.grpc:grpc-google-cloud-firestore-v1", //
google_cloud_platform_libraries_bom sets version
grpc_google_cloud_pubsub_v1 :
"com.google.api.grpc:grpc-google-cloud-pubsub-v1", //
google_cloud_platform_libraries_bom sets version
- grpc_google_cloud_pubsublite_v1 :
"com.google.api.grpc:grpc-google-cloud-pubsublite-v1", //
google_cloud_platform_libraries_bom sets version
+ grpc_google_cloud_pubsublite_v1 :
"com.google.api.grpc:grpc-google-cloud-pubsublite-v1:1.5.0", //
TODO(dpcollins-google): Let google_cloud_platform_libraries_bom set version
once high enough
grpc_google_common_protos :
"com.google.api.grpc:grpc-google-common-protos", //
google_cloud_platform_libraries_bom sets version
grpc_grpclb : "io.grpc:grpc-grpclb",
// google_cloud_platform_libraries_bom sets version
grpc_protobuf : "io.grpc:grpc-protobuf",
// google_cloud_platform_libraries_bom sets version
@@ -662,7 +662,7 @@ class BeamModulePlugin implements Plugin<Project> {
proto_google_cloud_datastore_v1 :
"com.google.api.grpc:proto-google-cloud-datastore-v1", //
google_cloud_platform_libraries_bom sets version
proto_google_cloud_firestore_v1 :
"com.google.api.grpc:proto-google-cloud-firestore-v1", //
google_cloud_platform_libraries_bom sets version
proto_google_cloud_pubsub_v1 :
"com.google.api.grpc:proto-google-cloud-pubsub-v1", //
google_cloud_platform_libraries_bom sets version
- proto_google_cloud_pubsublite_v1 :
"com.google.api.grpc:proto-google-cloud-pubsublite-v1", //
google_cloud_platform_libraries_bom sets version
+ proto_google_cloud_pubsublite_v1 :
"com.google.api.grpc:proto-google-cloud-pubsublite-v1:1.5.0", //
TODO(dpcollins-google): Let google_cloud_platform_libraries_bom set version
once high enough
proto_google_cloud_spanner_v1 :
"com.google.api.grpc:proto-google-cloud-spanner-v1", //
google_cloud_platform_libraries_bom sets version
proto_google_cloud_spanner_admin_database_v1:
"com.google.api.grpc:proto-google-cloud-spanner-admin-database-v1", //
google_cloud_platform_libraries_bom sets version
proto_google_common_protos :
"com.google.api.grpc:proto-google-common-protos", //
google_cloud_platform_libraries_bom sets version
diff --git
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/MemoryBufferedSubscriber.java
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/MemoryBufferedSubscriber.java
new file mode 100644
index 0000000..27b6d83
--- /dev/null
+++
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/MemoryBufferedSubscriber.java
@@ -0,0 +1,55 @@
+/*
+ * 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.beam.sdk.io.gcp.pubsublite.internal;
+
+import com.google.api.core.ApiFuture;
+import com.google.api.core.ApiService;
+import com.google.api.gax.rpc.ApiException;
+import com.google.cloud.pubsublite.Offset;
+import com.google.cloud.pubsublite.proto.SequencedMessage;
+import java.util.Optional;
+
+interface MemoryBufferedSubscriber extends ApiService {
+
+ /**
+ * Get the current fetch offset of this subscriber. This offset will be less
than or equal to all
+ * future messages returned by this object.
+ */
+ Offset fetchOffset();
+
+ /**
+ * Notify this subscriber that all messages that have been removed with
`pop` should no longer be
+ * counted against the memory budget.
+ *
+ * <p>Acquire a new memory buffer and allow any bytes which are now
available to be sent to this
+ * job.
+ */
+ void rebuffer() throws ApiException;
+
+ /** Return the head message from the buffer if it exists, or empty()
otherwise. */
+ Optional<SequencedMessage> peek();
+
+ /** Remove the head message from the buffer. Throws if no messages exist in
the buffer. */
+ void pop();
+
+ /**
+ * Return a future which will be satisfied when data is likely available or
the subscriber has
+ * failed.
+ */
+ ApiFuture<Void> onData();
+}
diff --git
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/MemoryBufferedSubscriberImpl.java
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/MemoryBufferedSubscriberImpl.java
new file mode 100644
index 0000000..d836797
--- /dev/null
+++
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/MemoryBufferedSubscriberImpl.java
@@ -0,0 +1,173 @@
+/*
+ * 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.beam.sdk.io.gcp.pubsublite.internal;
+
+import com.google.api.core.ApiFuture;
+import com.google.api.core.ApiFutures;
+import com.google.api.core.SettableApiFuture;
+import com.google.api.gax.rpc.ApiException;
+import com.google.cloud.pubsublite.Offset;
+import com.google.cloud.pubsublite.Partition;
+import com.google.cloud.pubsublite.internal.CheckedApiException;
+import com.google.cloud.pubsublite.internal.ProxyService;
+import com.google.cloud.pubsublite.internal.wire.Subscriber;
+import com.google.cloud.pubsublite.proto.FlowControlRequest;
+import com.google.cloud.pubsublite.proto.SequencedMessage;
+import java.util.ArrayDeque;
+import java.util.List;
+import java.util.Optional;
+import java.util.Queue;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import org.apache.beam.sdk.io.gcp.pubsublite.internal.MemoryLimiter.Block;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+class MemoryBufferedSubscriberImpl extends ProxyService implements
MemoryBufferedSubscriber {
+ private static final Logger LOG =
LoggerFactory.getLogger(MemoryBufferedSubscriberImpl.class);
+
+ private final Partition partition;
+ private final MemoryLimiter limiter;
+ private final Subscriber subscriber;
+ private long targetMemory;
+ private Offset fetchOffset;
+ private Block memBlock;
+
+ private long bytesOutstandingToServer = 0;
+ private long bytesOutstanding = 0;
+ private final Queue<SequencedMessage> messages = new ArrayDeque<>();
+ private SettableApiFuture<Void> newData = SettableApiFuture.create();
+ private boolean shutdown = false;
+
+ // onReceive will not be called inline as subscriber is not started.
+ // addServices is intended to be called from the constructor.
+ @SuppressWarnings({"methodref.receiver.bound.invalid",
"method.invocation.invalid"})
+ public MemoryBufferedSubscriberImpl(
+ Partition partition,
+ Offset startOffset,
+ MemoryLimiter limiter,
+ Function<Consumer<List<SequencedMessage>>, Subscriber>
subscriberFactory) {
+ this.partition = partition;
+ this.fetchOffset = startOffset;
+ this.limiter = limiter;
+ this.targetMemory = limiter.maxBlockSize();
+ this.subscriber = subscriberFactory.apply(this::onReceive);
+ addServices(this.subscriber);
+ memBlock = limiter.claim(targetMemory);
+ }
+
+ @Override
+ protected synchronized void start() throws CheckedApiException {
+ bytesOutstandingToServer += memBlock.claimed();
+ bytesOutstanding += memBlock.claimed();
+ subscriber.allowFlow(
+ FlowControlRequest.newBuilder()
+ .setAllowedBytes(memBlock.claimed())
+ .setAllowedMessages(Long.MAX_VALUE)
+ .build());
+ }
+
+ @Override
+ protected synchronized void stop() {
+ if (shutdown) {
+ return;
+ }
+ shutdown = true;
+ newData.set(null);
+ memBlock.close();
+ }
+
+ @Override
+ protected synchronized void handlePermanentError(CheckedApiException e) {
+ stop();
+ }
+
+ private synchronized void onReceive(List<SequencedMessage> batch) {
+ if (shutdown) {
+ return;
+ }
+ for (SequencedMessage message : batch) {
+ bytesOutstandingToServer -= message.getSizeBytes();
+ }
+ messages.addAll(batch);
+ newData.set(null);
+ newData = SettableApiFuture.create();
+ }
+
+ @Override
+ public synchronized Offset fetchOffset() {
+ return fetchOffset;
+ }
+
+ @Override
+ public synchronized void rebuffer() throws ApiException {
+ if (shutdown) {
+ return;
+ }
+ if (bytesOutstandingToServer < (targetMemory / 3)) {
+ // Server is delivering lots of data, increase the target so that it is
not throttled.
+ targetMemory = Math.min(limiter.maxBlockSize(), targetMemory * 2);
+ } else if (bytesOutstandingToServer > (2 * targetMemory / 3)) {
+ // Server is delivering little data, decrease the target so that memory
can be used for other
+ // users of the limiter.
+ targetMemory = Math.max(limiter.minBlockSize(), targetMemory / 2);
+ }
+ long claimTarget = Math.max(bytesOutstanding, targetMemory);
+ memBlock.close();
+ memBlock = limiter.claim(claimTarget);
+ long toAllow = Math.max(memBlock.claimed() - bytesOutstanding, 0);
+ if (toAllow > 0) {
+ bytesOutstanding += toAllow;
+ bytesOutstandingToServer += toAllow;
+ try {
+
subscriber.allowFlow(FlowControlRequest.newBuilder().setAllowedBytes(toAllow).build());
+ } catch (CheckedApiException e) {
+ throw e.underlying;
+ }
+ } else {
+ LOG.debug(
+ "Not claiming memory: partition {} outstanding {} to server {}
target {} claimed {} messages {}",
+ partition,
+ bytesOutstanding,
+ bytesOutstandingToServer,
+ targetMemory,
+ memBlock.claimed(),
+ messages.size());
+ }
+ }
+
+ @Override
+ public synchronized Optional<SequencedMessage> peek() {
+ return Optional.ofNullable(messages.peek());
+ }
+
+ @Override
+ public synchronized void pop() {
+ SequencedMessage message = messages.remove();
+ bytesOutstanding -= message.getSizeBytes();
+ fetchOffset = Offset.of(message.getCursor().getOffset() + 1);
+ }
+
+ @Override
+ public synchronized ApiFuture<Void> onData() {
+ if (shutdown || !messages.isEmpty()) {
+ return ApiFutures.immediateFuture(null);
+ }
+ return newData;
+ }
+}
diff --git
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/PerServerPublisherCache.java
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/MemoryLimiter.java
similarity index 65%
copy from
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/PerServerPublisherCache.java
copy to
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/MemoryLimiter.java
index 07953cb..02cee19 100644
---
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/PerServerPublisherCache.java
+++
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/MemoryLimiter.java
@@ -18,17 +18,20 @@
package org.apache.beam.sdk.io.gcp.pubsublite.internal;
/**
- * A shared cache per-worker instance of Pub/Sub Lite publishers.
- *
- * <p>Pub/Sub Lite publishers connect to all available partitions: it would be
a pessimization for
- * all instances of the PubsubLiteSink to do this.
+ * A class which tracks blocks of memory which have been given out, and tries
to limit total memory
+ * size.
*/
-final class PerServerPublisherCache {
- private PerServerPublisherCache() {}
+interface MemoryLimiter {
+ Block claim(long toAcquire);
+
+ long minBlockSize();
+
+ long maxBlockSize();
- static final PublisherCache PUBLISHER_CACHE = new PublisherCache();
+ interface Block extends AutoCloseable {
+ long claimed();
- static {
- Runtime.getRuntime().addShutdownHook(new Thread(PUBLISHER_CACHE::close));
+ @Override
+ void close();
}
}
diff --git
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/MemoryLimiterImpl.java
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/MemoryLimiterImpl.java
new file mode 100644
index 0000000..3e54cc5
--- /dev/null
+++
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/MemoryLimiterImpl.java
@@ -0,0 +1,92 @@
+/*
+ * 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.beam.sdk.io.gcp.pubsublite.internal;
+
+import static
org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import javax.annotation.concurrent.GuardedBy;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+class MemoryLimiterImpl implements MemoryLimiter {
+ private static final Logger LOG =
LoggerFactory.getLogger(MemoryLimiterImpl.class);
+ private final long minBlockSize;
+ private final long maxBlockSize;
+ private final long maxAvailable;
+
+ @GuardedBy("this")
+ private long available;
+
+ public MemoryLimiterImpl(long minBlockSize, long maxBlockSize, long
maxAvailable) {
+ this.minBlockSize = minBlockSize;
+ this.maxBlockSize = maxBlockSize;
+ this.maxAvailable = maxAvailable;
+ this.available = maxAvailable;
+ }
+
+ @Override
+ public synchronized Block claim(long toAcquire) {
+ toAcquire = Math.max(Math.min(toAcquire, available / 2), minBlockSize);
+ available -= toAcquire;
+ return new Block(toAcquire);
+ }
+
+ @Override
+ public long minBlockSize() {
+ return minBlockSize;
+ }
+
+ @Override
+ public long maxBlockSize() {
+ return maxBlockSize;
+ }
+
+ private synchronized void release(long toRelease) {
+ available += toRelease;
+ checkState(available <= maxAvailable);
+ }
+
+ public class Block implements MemoryLimiter.Block {
+ public final long claimed;
+ private boolean released = false;
+
+ private Block(long claimed) {
+ this.claimed = claimed;
+ }
+
+ @Override
+ public long claimed() {
+ return claimed;
+ }
+
+ @Override
+ public void close() {
+ checkState(!released);
+ released = true;
+ release(claimed);
+ }
+
+ @Override
+ public void finalize() {
+ if (!released) {
+ LOG.error("Failed to release memory block- likely SDF implementation
error.");
+ close();
+ }
+ }
+ }
+}
diff --git
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/OffsetByteRangeTracker.java
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/OffsetByteRangeTracker.java
index 44e6246..f4ac39d 100644
---
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/OffsetByteRangeTracker.java
+++
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/OffsetByteRangeTracker.java
@@ -17,45 +17,28 @@
*/
package org.apache.beam.sdk.io.gcp.pubsublite.internal;
+import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull;
import static
org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
-import static
org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkNotNull;
import static
org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
import com.google.cloud.pubsublite.Offset;
import com.google.cloud.pubsublite.proto.ComputeMessageStatsResponse;
-import java.util.concurrent.TimeUnit;
import javax.annotation.Nullable;
import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
import org.apache.beam.sdk.transforms.splittabledofn.SplitResult;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Stopwatch;
-import org.joda.time.Duration;
/**
* OffsetByteRangeTracker is an unbounded restriction tracker for Pub/Sub lite
partitions that
* tracks offsets for checkpointing and bytes for progress.
- *
- * <p>Any valid instance of an OffsetByteRangeTracker tracks one of exactly
two types of ranges: -
- * Unbounded ranges whose last offset is Long.MAX_VALUE - Completed ranges
that are either empty
- * (From == To) or fully claimed (lastClaimed == To - 1)
- *
- * <p>Also prevents splitting until minTrackingTime has passed or
minBytesReceived have been
- * received. IMPORTANT: minTrackingTime must be strictly smaller than the SDF
read timeout when it
- * would return ProcessContinuation.resume().
*/
class OffsetByteRangeTracker extends TrackerWithProgress {
private final TopicBacklogReader unownedBacklogReader;
- private final Duration minTrackingTime;
- private final long minBytesReceived;
- private final Stopwatch stopwatch;
- private OffsetByteRange range;
+ private final OffsetRangeTracker rangeTracker;
+ private long bytes;
private @Nullable Long lastClaimed;
- public OffsetByteRangeTracker(
- OffsetByteRange range,
- TopicBacklogReader unownedBacklogReader,
- Stopwatch stopwatch,
- Duration minTrackingTime,
- long minBytesReceived) {
+ public OffsetByteRangeTracker(OffsetByteRange range, TopicBacklogReader
unownedBacklogReader) {
checkArgument(
range.getRange().getTo() == Long.MAX_VALUE,
"May only construct OffsetByteRangeTracker with an unbounded range
with no progress.");
@@ -63,10 +46,8 @@ class OffsetByteRangeTracker extends TrackerWithProgress {
range.getByteCount() == 0L,
"May only construct OffsetByteRangeTracker with an unbounded range
with no progress.");
this.unownedBacklogReader = unownedBacklogReader;
- this.minTrackingTime = minTrackingTime;
- this.minBytesReceived = minBytesReceived;
- this.stopwatch = stopwatch.reset().start();
- this.range = range;
+ this.rangeTracker = new OffsetRangeTracker(range.getRange());
+ this.bytes = range.getByteCount();
}
@Override
@@ -76,34 +57,17 @@ class OffsetByteRangeTracker extends TrackerWithProgress {
@Override
public boolean tryClaim(OffsetByteProgress position) {
- long toClaim = position.lastOffset().value();
- checkArgument(
- lastClaimed == null || toClaim > lastClaimed,
- "Trying to claim offset %s while last attempted was %s",
- position.lastOffset().value(),
- lastClaimed);
- checkArgument(
- toClaim >= range.getRange().getFrom(),
- "Trying to claim offset %s before start of the range %s",
- toClaim,
- range);
- // split() has already been called, truncating this range. No more offsets
may be claimed.
- if (range.getRange().getTo() != Long.MAX_VALUE) {
- boolean isRangeEmpty = range.getRange().getTo() ==
range.getRange().getFrom();
- boolean isValidClosedRange = nextOffset() == range.getRange().getTo();
- checkState(
- isRangeEmpty || isValidClosedRange,
- "Violated class precondition: offset range improperly split. Please
report a beam bug.");
+ if (!rangeTracker.tryClaim(position.lastOffset().value())) {
return false;
}
- lastClaimed = toClaim;
- range = OffsetByteRange.of(range.getRange(), range.getByteCount() +
position.batchBytes());
+ lastClaimed = position.lastOffset().value();
+ bytes += position.batchBytes();
return true;
}
@Override
public OffsetByteRange currentRestriction() {
- return range;
+ return OffsetByteRange.of(rangeTracker.currentRestriction(), bytes);
}
private long nextOffset() {
@@ -111,62 +75,30 @@ class OffsetByteRangeTracker extends TrackerWithProgress {
return lastClaimed == null ? currentRestriction().getRange().getFrom() :
lastClaimed + 1;
}
- /**
- * Whether the tracker has received enough data/been running for enough time
that it can
- * checkpoint and be confident it can get sufficient throughput.
- */
- private boolean receivedEnough() {
- Duration duration =
Duration.millis(stopwatch.elapsed(TimeUnit.MILLISECONDS));
- if (duration.isLongerThan(minTrackingTime)) {
- return true;
- }
- if (currentRestriction().getByteCount() >= minBytesReceived) {
- return true;
- }
- return false;
- }
-
@Override
public @Nullable SplitResult<OffsetByteRange> trySplit(double
fractionOfRemainder) {
// Cannot split a bounded range. This should already be completely claimed.
- if (range.getRange().getTo() != Long.MAX_VALUE) {
+ if (rangeTracker.currentRestriction().getTo() != Long.MAX_VALUE) {
return null;
}
- if (!receivedEnough()) {
+ @Nullable SplitResult<OffsetRange> ranges =
rangeTracker.trySplit(fractionOfRemainder);
+ if (ranges == null) {
return null;
}
- range =
- OffsetByteRange.of(
- new OffsetRange(currentRestriction().getRange().getFrom(),
nextOffset()),
- range.getByteCount());
+
checkArgument(rangeTracker.currentRestriction().equals(ranges.getPrimary()));
return SplitResult.of(
- this.range, OffsetByteRange.of(new OffsetRange(nextOffset(),
Long.MAX_VALUE), 0));
+ currentRestriction(),
OffsetByteRange.of(checkArgumentNotNull(ranges.getResidual())));
}
@Override
- @SuppressWarnings("unboxing.of.nullable")
public void checkDone() throws IllegalStateException {
- if (range.getRange().getFrom() == range.getRange().getTo()) {
- return;
- }
- checkState(
- lastClaimed != null,
- "Last attempted offset should not be null. No work was claimed in
non-empty range %s.",
- range);
- long lastClaimedNotNull = checkNotNull(lastClaimed);
- checkState(
- lastClaimedNotNull >= range.getRange().getTo() - 1,
- "Last attempted offset was %s in range %s, claiming work in [%s, %s)
was not attempted",
- lastClaimedNotNull,
- range,
- lastClaimedNotNull + 1,
- range.getRange().getTo());
+ rangeTracker.checkDone();
}
@Override
public Progress getProgress() {
ComputeMessageStatsResponse stats =
this.unownedBacklogReader.computeMessageStats(Offset.of(nextOffset()));
- return Progress.from(range.getByteCount(), stats.getMessageBytes());
+ return Progress.from(bytes, stats.getMessageBytes());
}
}
diff --git
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/PerServerPublisherCache.java
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/PerServerPublisherCache.java
index 07953cb..ee76169 100644
---
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/PerServerPublisherCache.java
+++
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/PerServerPublisherCache.java
@@ -17,6 +17,10 @@
*/
package org.apache.beam.sdk.io.gcp.pubsublite.internal;
+import com.google.cloud.pubsublite.MessageMetadata;
+import com.google.cloud.pubsublite.internal.Publisher;
+import org.apache.beam.sdk.io.gcp.pubsublite.PublisherOptions;
+
/**
* A shared cache per-worker instance of Pub/Sub Lite publishers.
*
@@ -26,7 +30,8 @@ package org.apache.beam.sdk.io.gcp.pubsublite.internal;
final class PerServerPublisherCache {
private PerServerPublisherCache() {}
- static final PublisherCache PUBLISHER_CACHE = new PublisherCache();
+ static final ServiceCache<PublisherOptions, Publisher<MessageMetadata>>
PUBLISHER_CACHE =
+ new ServiceCache<>();
static {
Runtime.getRuntime().addShutdownHook(new Thread(PUBLISHER_CACHE::close));
diff --git
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/PerServerPublisherCache.java
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/PerServerSubscriberCache.java
similarity index 66%
copy from
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/PerServerPublisherCache.java
copy to
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/PerServerSubscriberCache.java
index 07953cb..130ecc8 100644
---
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/PerServerPublisherCache.java
+++
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/PerServerSubscriberCache.java
@@ -18,17 +18,17 @@
package org.apache.beam.sdk.io.gcp.pubsublite.internal;
/**
- * A shared cache per-worker instance of Pub/Sub Lite publishers.
+ * A shared cache per-worker instance of Pub/Sub Lite subscribers.
*
- * <p>Pub/Sub Lite publishers connect to all available partitions: it would be
a pessimization for
- * all instances of the PubsubLiteSink to do this.
+ * <p>This allows subscribers to buffer in the background while processElement
is not running.
*/
-final class PerServerPublisherCache {
- private PerServerPublisherCache() {}
+final class PerServerSubscriberCache {
+ private PerServerSubscriberCache() {}
- static final PublisherCache PUBLISHER_CACHE = new PublisherCache();
+ static final ServiceCache<SubscriptionPartition, MemoryBufferedSubscriber>
CACHE =
+ new ServiceCache<>();
static {
- Runtime.getRuntime().addShutdownHook(new Thread(PUBLISHER_CACHE::close));
+ Runtime.getRuntime().addShutdownHook(new Thread(CACHE::close));
}
}
diff --git
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/PerSubscriptionPartitionSdf.java
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/PerSubscriptionPartitionSdf.java
index d387cf1..0996449 100644
---
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/PerSubscriptionPartitionSdf.java
+++
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/PerSubscriptionPartitionSdf.java
@@ -28,8 +28,11 @@ import
org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
import
org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
import org.joda.time.Duration;
import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
class PerSubscriptionPartitionSdf extends DoFn<SubscriptionPartition,
SequencedMessage> {
+ private static final Logger LOG =
LoggerFactory.getLogger(PerSubscriptionPartitionSdf.class);
private final Duration maxSleepTime;
private final ManagedBacklogReaderFactory backlogReaderFactory;
private final SubscriptionPartitionProcessorFactory processorFactory;
@@ -60,6 +63,13 @@ class PerSubscriptionPartitionSdf extends
DoFn<SubscriptionPartition, SequencedM
backlogReaderFactory.close();
}
+ /**
+ * The initial watermark state is not allowed to return less than the
element's input timestamp.
+ *
+ * <p>The polling logic for identifying new partitions will export all
preexisting partitions with
+ * very old (EPOCH) initial watermarks, and any new partitions with a recent
watermark likely to
+ * be before all messages that could exist on that partition given the
polling delay.
+ */
@GetInitialWatermarkEstimatorState
public Instant getInitialWatermarkState(@Timestamp Instant elementTimestamp)
{
return elementTimestamp;
@@ -76,20 +86,25 @@ class PerSubscriptionPartitionSdf extends
DoFn<SubscriptionPartition, SequencedM
@Element SubscriptionPartition subscriptionPartition,
OutputReceiver<SequencedMessage> receiver)
throws Exception {
+ LOG.debug("Starting process for {} at {}", subscriptionPartition,
Instant.now());
SubscriptionPartitionProcessor processor =
processorFactory.newProcessor(subscriptionPartition, tracker,
receiver);
ProcessContinuation result = processor.runFor(maxSleepTime);
+ LOG.debug("Starting commit for {} at {}", subscriptionPartition,
Instant.now());
+ // TODO(dpcollins-google): Move commits to a bundle finalizer for drain
correctness
processor
.lastClaimed()
.ifPresent(
- lastClaimedOffset -> {
- Offset commitOffset = Offset.of(lastClaimedOffset.value() + 1);
+ lastClaimed -> {
try {
-
committerFactory.apply(subscriptionPartition).commitOffset(commitOffset);
+ committerFactory
+ .apply(subscriptionPartition)
+ .commitOffset(Offset.of(lastClaimed.value() + 1));
} catch (Exception e) {
throw ExtractStatus.toCanonical(e).underlying;
}
});
+ LOG.debug("Finishing process for {} at {}", subscriptionPartition,
Instant.now());
return result;
}
diff --git
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/PubsubLiteSink.java
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/PubsubLiteSink.java
index f370919..e23a66d 100644
---
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/PubsubLiteSink.java
+++
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/PubsubLiteSink.java
@@ -53,7 +53,9 @@ public class PubsubLiteSink extends DoFn<PubSubMessage, Void>
{
private final Publisher<MessageMetadata> publisher;
RunState(PublisherOptions options) {
- publisher = PerServerPublisherCache.PUBLISHER_CACHE.get(options);
+ publisher =
+ PerServerPublisherCache.PUBLISHER_CACHE.get(
+ options, () -> new PublisherAssembler(options).newPublisher());
}
void publish(PubSubMessage message) {
diff --git
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/PublisherCache.java
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/ServiceCache.java
similarity index 55%
rename from
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/PublisherCache.java
rename to
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/ServiceCache.java
index ac85ba9..5e124c0 100644
---
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/PublisherCache.java
+++
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/ServiceCache.java
@@ -17,59 +17,63 @@
*/
package org.apache.beam.sdk.io.gcp.pubsublite.internal;
+import com.google.api.core.ApiService;
import com.google.api.core.ApiService.Listener;
import com.google.api.core.ApiService.State;
import com.google.api.gax.rpc.ApiException;
-import com.google.cloud.pubsublite.MessageMetadata;
-import com.google.cloud.pubsublite.internal.Publisher;
+import com.google.cloud.pubsublite.internal.wire.ApiServiceUtils;
import com.google.cloud.pubsublite.internal.wire.SystemExecutors;
import com.google.errorprone.annotations.concurrent.GuardedBy;
import java.util.HashMap;
-import org.apache.beam.sdk.io.gcp.pubsublite.PublisherOptions;
+import java.util.function.Supplier;
import
org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
-/** A map of working publishers by PublisherOptions. */
-class PublisherCache implements AutoCloseable {
- private final Logger logger = LoggerFactory.getLogger(PublisherCache.class);
+/** A map of working ApiServices by identifying key. The key must be hashable.
*/
+class ServiceCache<K, V extends ApiService> implements AutoCloseable {
+ private final Logger logger = LoggerFactory.getLogger(ServiceCache.class);
@GuardedBy("this")
- private final HashMap<PublisherOptions, Publisher<MessageMetadata>>
livePublishers =
- new HashMap<>();
+ private final HashMap<K, V> liveMap = new HashMap<>();
- private synchronized void evict(PublisherOptions options) {
- livePublishers.remove(options);
+ private synchronized void evict(K key, V service) {
+ liveMap.remove(key, service);
}
- synchronized Publisher<MessageMetadata> get(PublisherOptions options) throws
ApiException {
- Publisher<MessageMetadata> publisher = livePublishers.get(options);
- if (publisher != null) {
- return publisher;
+ synchronized V get(K key, Supplier<V> factory) throws ApiException {
+ V service = liveMap.get(key);
+ if (service != null) {
+ return service;
}
- publisher = new PublisherAssembler(options).newPublisher();
- livePublishers.put(options, publisher);
- publisher.addListener(
+ V newService = factory.get();
+ liveMap.put(key, newService);
+ newService.addListener(
new Listener() {
@Override
public void failed(State s, Throwable t) {
- logger.warn("Publisher failed.", t);
- evict(options);
+ logger.warn(newService.getClass().getSimpleName() + " failed.", t);
+ evict(key, newService);
+ }
+
+ @Override
+ public void terminated(State from) {
+ evict(key, newService);
}
},
SystemExecutors.getFuturesExecutor());
- publisher.startAsync().awaitRunning();
- return publisher;
+ newService.startAsync().awaitRunning();
+ return newService;
}
@VisibleForTesting
- synchronized void set(PublisherOptions options, Publisher<MessageMetadata>
toCache) {
- livePublishers.put(options, toCache);
+ synchronized void set(K key, V service) {
+ liveMap.put(key, service);
}
@Override
public synchronized void close() {
- livePublishers.forEach((options, publisher) -> publisher.stopAsync());
- livePublishers.clear();
+ ApiServiceUtils.blockingShutdown(liveMap.values());
+ liveMap.clear();
}
}
diff --git
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/SubscribeTransform.java
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/SubscribeTransform.java
index 4b572b6..97c2899 100644
---
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/SubscribeTransform.java
+++
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/SubscribeTransform.java
@@ -30,6 +30,7 @@ import com.google.cloud.pubsublite.internal.wire.Subscriber;
import com.google.cloud.pubsublite.proto.SequencedMessage;
import java.util.List;
import java.util.function.Consumer;
+import java.util.function.Supplier;
import java.util.stream.Collectors;
import org.apache.beam.sdk.io.gcp.pubsublite.SubscriberOptions;
import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
@@ -38,12 +39,16 @@ import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
import org.apache.beam.sdk.values.PBegin;
import org.apache.beam.sdk.values.PCollection;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Stopwatch;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.math.LongMath;
import org.joda.time.Duration;
public class SubscribeTransform extends PTransform<PBegin,
PCollection<SequencedMessage>> {
private static final long MEBIBYTE = 1L << 20;
+ private static final long SOFT_MEMORY_LIMIT = 512 * MEBIBYTE;
+ private static final long MIN_PER_PARTITION_MEMORY = 4 * MEBIBYTE;
+ private static final long MAX_PER_PARTITION_MEMORY = 100 * MEBIBYTE;
+
+ private static final MemoryLimiter LIMITER =
+ new MemoryLimiterImpl(MIN_PER_PARTITION_MEMORY,
MAX_PER_PARTITION_MEMORY, SOFT_MEMORY_LIMIT);
private final SubscriberOptions options;
@@ -74,40 +79,37 @@ public class SubscribeTransform extends PTransform<PBegin,
PCollection<Sequenced
private SubscriptionPartitionProcessor newPartitionProcessor(
SubscriptionPartition subscriptionPartition,
RestrictionTracker<OffsetByteRange, OffsetByteProgress> tracker,
- OutputReceiver<SequencedMessage> receiver)
- throws ApiException {
- checkSubscription(subscriptionPartition);
+ OutputReceiver<SequencedMessage> receiver) {
+ Supplier<MemoryBufferedSubscriber> newSubscriber =
+ () ->
+ newBufferedSubscriber(
+ subscriptionPartition,
+ Offset.of(tracker.currentRestriction().getRange().getFrom()));
return new SubscriptionPartitionProcessorImpl(
+ subscriptionPartition,
tracker,
receiver,
- consumer ->
- newSubscriber(
- subscriptionPartition.partition(),
- Offset.of(tracker.currentRestriction().getRange().getFrom()),
- consumer),
- options.flowControlSettings());
+ () -> PerServerSubscriberCache.CACHE.get(subscriptionPartition,
newSubscriber));
}
- private TopicBacklogReader newBacklogReader(SubscriptionPartition
subscriptionPartition) {
+ private MemoryBufferedSubscriber newBufferedSubscriber(
+ SubscriptionPartition subscriptionPartition, Offset startOffset) throws
ApiException {
checkSubscription(subscriptionPartition);
- return new SubscriberAssembler(options,
subscriptionPartition.partition()).getBacklogReader();
+ return new MemoryBufferedSubscriberImpl(
+ subscriptionPartition.partition(),
+ startOffset,
+ LIMITER,
+ consumer -> newSubscriber(subscriptionPartition.partition(),
startOffset, consumer));
}
- private long calculateMinWindowBytes() {
- long minFromFlowControl =
-
LongMath.saturatedMultiply(options.flowControlSettings().bytesOutstanding(),
10);
- // Dataflow will not accept outputs larger than 1 GiB. Cap the maximum at
750 MiB to avoid this.
- return Math.min(minFromFlowControl, 750 * MEBIBYTE);
+ private TopicBacklogReader newBacklogReader(SubscriptionPartition
subscriptionPartition) {
+ checkSubscription(subscriptionPartition);
+ return new SubscriberAssembler(options,
subscriptionPartition.partition()).getBacklogReader();
}
private TrackerWithProgress newRestrictionTracker(
TopicBacklogReader backlogReader, OffsetByteRange initial) {
- return new OffsetByteRangeTracker(
- initial,
- backlogReader,
- Stopwatch.createUnstarted(),
- options.minBundleTimeout(),
- calculateMinWindowBytes());
+ return new OffsetByteRangeTracker(initial, backlogReader);
}
private InitialOffsetReader newInitialOffsetReader(SubscriptionPartition
subscriptionPartition) {
diff --git
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/SubscriberAssembler.java
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/SubscriberAssembler.java
index 7825ab8..d65375e 100644
---
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/SubscriberAssembler.java
+++
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/SubscriberAssembler.java
@@ -114,6 +114,7 @@ class SubscriberAssembler {
.setMessageConsumer(consumer)
.setSubscriptionPath(options.subscriptionPath())
.setPartition(partition)
+ .setRetryStreamRaces(false)
.setStreamFactory(
responseStream -> {
ApiCallContext context =
diff --git
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/SubscriptionPartitionProcessorImpl.java
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/SubscriptionPartitionProcessorImpl.java
index 8b2eff9..818ccf7 100644
---
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/SubscriptionPartitionProcessorImpl.java
+++
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/SubscriptionPartitionProcessorImpl.java
@@ -17,161 +17,111 @@
*/
package org.apache.beam.sdk.io.gcp.pubsublite.internal;
-import static
com.google.cloud.pubsublite.internal.wire.ApiServiceUtils.blockingShutdown;
+import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull;
-import com.google.api.core.ApiService.Listener;
-import com.google.api.core.ApiService.State;
import com.google.cloud.pubsublite.Offset;
-import com.google.cloud.pubsublite.cloudpubsub.FlowControlSettings;
-import com.google.cloud.pubsublite.internal.CheckedApiException;
-import com.google.cloud.pubsublite.internal.ExtractStatus;
-import com.google.cloud.pubsublite.internal.wire.Subscriber;
-import com.google.cloud.pubsublite.internal.wire.SystemExecutors;
-import com.google.cloud.pubsublite.proto.FlowControlRequest;
import com.google.cloud.pubsublite.proto.SequencedMessage;
import com.google.protobuf.util.Timestamps;
-import java.util.List;
import java.util.Optional;
-import java.util.concurrent.SynchronousQueue;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
-import java.util.function.Consumer;
-import java.util.function.Function;
-import javax.annotation.Nullable;
+import java.util.concurrent.TimeoutException;
+import java.util.function.Supplier;
import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
import org.apache.beam.sdk.transforms.DoFn.ProcessContinuation;
import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
-import
org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
-import
org.apache.beam.vendor.guava.v26_0_jre.com.google.common.util.concurrent.SettableFuture;
import org.joda.time.Duration;
import org.joda.time.Instant;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
-class SubscriptionPartitionProcessorImpl extends Listener
- implements SubscriptionPartitionProcessor, AutoCloseable {
+class SubscriptionPartitionProcessorImpl implements
SubscriptionPartitionProcessor {
private static final Logger LOG =
LoggerFactory.getLogger(SubscriptionPartitionProcessorImpl.class);
+ private final SubscriptionPartition subscriptionPartition;
private final RestrictionTracker<OffsetByteRange, OffsetByteProgress>
tracker;
private final OutputReceiver<SequencedMessage> receiver;
- private final Subscriber subscriber;
- private final SettableFuture<Void> completionFuture =
SettableFuture.create();
- // Queue to transfer messages from subscriber callback to runFor downcall.
- private final SynchronousQueue<List<SequencedMessage>> transfer = new
SynchronousQueue<>();
- private final FlowControlSettings flowControlSettings;
+ private final MemoryBufferedSubscriber subscriber;
private Optional<Offset> lastClaimedOffset = Optional.empty();
- @SuppressWarnings("methodref.receiver.bound.invalid")
+ // getReadySubscriber doesn't reference the subscriber member.
+ @SuppressWarnings("method.invocation.invalid")
SubscriptionPartitionProcessorImpl(
+ SubscriptionPartition subscriptionPartition,
RestrictionTracker<OffsetByteRange, OffsetByteProgress> tracker,
OutputReceiver<SequencedMessage> receiver,
- Function<Consumer<List<SequencedMessage>>, Subscriber> subscriberFactory,
- FlowControlSettings flowControlSettings) {
+ Supplier<MemoryBufferedSubscriber> subscriberFactory) {
+ this.subscriptionPartition = subscriptionPartition;
this.tracker = tracker;
this.receiver = receiver;
- this.subscriber = subscriberFactory.apply(this::onSubscriberMessages);
- this.flowControlSettings = flowControlSettings;
- }
-
- @Override
- public void failed(State from, Throwable failure) {
- completionFuture.setException(ExtractStatus.toCanonical(failure));
- }
-
- private void onSubscriberMessages(List<SequencedMessage> messages) {
- try {
- while (!completionFuture.isDone()) {
- if (transfer.offer(messages, 10, TimeUnit.MILLISECONDS)) {
- return;
- }
- }
- } catch (Throwable t) {
- throw ExtractStatus.toCanonical(t).underlying;
- }
- }
-
- @SuppressWarnings("argument.type.incompatible")
- private void start() {
- this.subscriber.addListener(this, SystemExecutors.getFuturesExecutor());
- this.subscriber.startAsync();
- this.subscriber.awaitRunning();
- try {
- this.subscriber.allowFlow(
- FlowControlRequest.newBuilder()
- .setAllowedBytes(flowControlSettings.bytesOutstanding())
- .setAllowedMessages(flowControlSettings.messagesOutstanding())
- .build());
- } catch (Throwable t) {
- throw ExtractStatus.toCanonical(t).underlying;
- }
- }
-
- private void handleMessages(List<SequencedMessage> messages) {
- if (completionFuture.isDone()) {
- return;
- }
- Offset lastOffset =
Offset.of(Iterables.getLast(messages).getCursor().getOffset());
- long byteSize =
messages.stream().mapToLong(SequencedMessage::getSizeBytes).sum();
- if (tracker.tryClaim(OffsetByteProgress.of(lastOffset, byteSize))) {
- lastClaimedOffset = Optional.of(lastOffset);
- messages.forEach(
- message ->
- receiver.outputWithTimestamp(
- message, new
Instant(Timestamps.toMillis(message.getPublishTime()))));
- try {
- subscriber.allowFlow(
- FlowControlRequest.newBuilder()
- .setAllowedBytes(byteSize)
- .setAllowedMessages(messages.size())
- .build());
- } catch (CheckedApiException e) {
- completionFuture.setException(e);
- }
- } else {
- completionFuture.set(null);
- }
+ this.subscriber = getReadySubscriber(subscriberFactory);
}
@Override
@SuppressWarnings("argument.type.incompatible")
public ProcessContinuation runFor(Duration duration) {
- Instant deadline = Instant.now().plus(duration);
- start();
- try (SubscriptionPartitionProcessorImpl closeThis = this) {
- while (!completionFuture.isDone() && deadline.isAfterNow()) {
- @Nullable List<SequencedMessage> messages = transfer.poll(10,
TimeUnit.MILLISECONDS);
- if (messages != null) {
- handleMessages(messages);
+ Instant maxReadTime = Instant.now().plus(duration);
+ while (subscriber.isRunning()) {
+ // Read any available data.
+ for (Optional<SequencedMessage> next = subscriber.peek();
+ next.isPresent();
+ next = subscriber.peek()) {
+ SequencedMessage message = next.get();
+ Offset messageOffset = Offset.of(message.getCursor().getOffset());
+ if (tracker.tryClaim(OffsetByteProgress.of(messageOffset,
message.getSizeBytes()))) {
+ subscriber.pop();
+ lastClaimedOffset = Optional.of(messageOffset);
+ receiver.outputWithTimestamp(
+ message, new
Instant(Timestamps.toMillis(message.getPublishTime())));
+ } else {
+ // Our claim failed, return stop()
+ return ProcessContinuation.stop();
}
}
- } catch (Throwable t) {
- throw ExtractStatus.toCanonical(t).underlying;
- }
- // Determine return code after shutdown.
- if (completionFuture.isDone()) {
- // Call get() to ensure there is no exception.
+ // Try waiting for new data.
try {
- completionFuture.get();
- } catch (Throwable t) {
- throw ExtractStatus.toCanonical(t).underlying;
+ Duration readTime = new Duration(Instant.now(), maxReadTime);
+ Future<Void> onData = subscriber.onData();
+ checkArgumentNotNull(onData);
+ onData.get(readTime.getMillis(), TimeUnit.MILLISECONDS);
+ } catch (TimeoutException e) {
+ // Read timed out without us being cut off, yield to the runtime.
+ return ProcessContinuation.resume();
+ } catch (InterruptedException | ExecutionException e2) {
+ // We should never be interrupted by beam, and onData should never
return an error.
+ throw new RuntimeException(e2);
}
- // CompletionFuture set with null when tryClaim returned false.
- return ProcessContinuation.stop();
}
+ // Subscriber is no longer running, it has likely failed. Yield to the
runtime to retry reading
+ // with a new subscriber.
return ProcessContinuation.resume();
}
@Override
- public void close() {
- try {
- blockingShutdown(subscriber);
- } catch (Throwable t) {
- // Don't propagate errors on subscriber shutdown.
- LOG.info("Error on subscriber shutdown.", t);
- }
- }
-
- @Override
public Optional<Offset> lastClaimed() {
return lastClaimedOffset;
}
+
+ private MemoryBufferedSubscriber getReadySubscriber(
+ Supplier<MemoryBufferedSubscriber> getOrCreate) {
+ Offset startOffset =
Offset.of(tracker.currentRestriction().getRange().getFrom());
+ while (true) {
+ MemoryBufferedSubscriber subscriber = getOrCreate.get();
+ Offset fetchOffset = subscriber.fetchOffset();
+ if (startOffset.equals(fetchOffset)) {
+ subscriber.rebuffer(); // TODO(dpcollins-google): Move this to a
bundle finalizer
+ return subscriber;
+ }
+ LOG.info(
+ "Discarding subscriber due to mismatch, this should be rare. {},
start: {} fetch: {}",
+ subscriptionPartition,
+ startOffset,
+ fetchOffset);
+ try {
+ subscriber.stopAsync().awaitTerminated();
+ } catch (Exception ignored) {
+ }
+ }
+ }
}
diff --git
a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/MemoryBufferedSubscriberImplTest.java
b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/MemoryBufferedSubscriberImplTest.java
new file mode 100644
index 0000000..31eae3c
--- /dev/null
+++
b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/MemoryBufferedSubscriberImplTest.java
@@ -0,0 +1,200 @@
+/*
+ * 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.beam.sdk.io.gcp.pubsublite.internal;
+
+import static
com.google.cloud.pubsublite.internal.testing.UnitTestExamples.example;
+import static
org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkNotNull;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.reset;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.MockitoAnnotations.initMocks;
+
+import com.google.cloud.pubsublite.Offset;
+import com.google.cloud.pubsublite.Partition;
+import com.google.cloud.pubsublite.internal.testing.FakeApiService;
+import com.google.cloud.pubsublite.internal.wire.Subscriber;
+import com.google.cloud.pubsublite.proto.Cursor;
+import com.google.cloud.pubsublite.proto.SequencedMessage;
+import java.util.List;
+import java.util.concurrent.Future;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import
org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.Timeout;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.Mock;
+import org.mockito.Spy;
+
+@RunWith(JUnit4.class)
+public class MemoryBufferedSubscriberImplTest {
+ private static final long MAX_MEMORY = 1024;
+
+ @Rule public Timeout globalTimeout = Timeout.seconds(30);
+
+ abstract static class FakeSubscriber extends FakeApiService implements
Subscriber {}
+
+ @Spy FakeSubscriber subscriber;
+ @Mock Function<Consumer<List<SequencedMessage>>, Subscriber>
subscriberFactory;
+ @Mock MemoryLimiter limiter;
+ @Mock MemoryLimiter.Block block;
+
+ MemoryBufferedSubscriber bufferedSubscriber;
+
+ Consumer<List<SequencedMessage>> consumer;
+
+ private static SequencedMessage messageWithSize(long size) {
+ return SequencedMessage.newBuilder().setSizeBytes(size).build();
+ }
+
+ @Before
+ public void setUp() {
+ initMocks(this);
+ doAnswer(
+ args -> {
+ consumer = args.getArgument(0);
+ return subscriber;
+ })
+ .when(subscriberFactory)
+ .apply(any());
+ doReturn(1L).when(limiter).minBlockSize();
+ doReturn(MAX_MEMORY).when(limiter).maxBlockSize();
+ checkNotNull(block);
+ checkNotNull(limiter);
+ doReturn(block).when(limiter).claim(anyLong());
+ doReturn(MAX_MEMORY).when(block).claimed();
+ bufferedSubscriber =
+ new MemoryBufferedSubscriberImpl(
+ example(Partition.class), example(Offset.class), limiter,
subscriberFactory);
+ checkNotNull(consumer);
+ bufferedSubscriber.startAsync().awaitRunning();
+ verify(subscriber).startAsync();
+ assertTrue(subscriber.isRunning());
+ }
+
+ @Test
+ public void underlyingFailureFails() {
+ subscriber.fail(new RuntimeException("bad"));
+ assertThrows(Exception.class, subscriber::awaitTerminated);
+ }
+
+ @Test
+ public void rebufferReducesToOutstandingWhenLittleData() {
+ consumer.accept(ImmutableList.of(messageWithSize(MAX_MEMORY / 4)));
+ bufferedSubscriber.pop();
+ bufferedSubscriber.rebuffer();
+ verify(block).close();
+ verify(limiter).claim(3 * MAX_MEMORY / 4);
+ }
+
+ @Test
+ public void rebufferCannotGoBelowMin() {
+ long minBlock = MAX_MEMORY * 4 / 5;
+ doReturn(minBlock).when(limiter).minBlockSize();
+ for (int i = 0; i < 1000; ++i) {
+ // Rebuffer many times with no data to bring down the target value
+ bufferedSubscriber.rebuffer();
+ }
+ reset(limiter);
+ doReturn(minBlock).when(limiter).minBlockSize();
+ doReturn(block).when(limiter).claim(anyLong());
+ // Deliver enough data that 3 * minBlock / 4 is outstanding, buffer is
allowed to and will
+ // shrink except it is limited by min block size.
+ consumer.accept(ImmutableList.of(messageWithSize(2 * MAX_MEMORY / 5)));
+ bufferedSubscriber.pop();
+ bufferedSubscriber.rebuffer();
+ verify(limiter).claim(minBlock);
+ }
+
+ @Test
+ public void rebufferStaysSameOnHalfDelivered() {
+ consumer.accept(ImmutableList.of(messageWithSize(MAX_MEMORY / 4)));
+ bufferedSubscriber.pop();
+ bufferedSubscriber.rebuffer();
+ verify(limiter).claim(3 * MAX_MEMORY / 4);
+ consumer.accept(ImmutableList.of(messageWithSize(3 * MAX_MEMORY / 8)));
+ bufferedSubscriber.pop();
+ bufferedSubscriber.rebuffer();
+ verify(limiter).claim(3 * MAX_MEMORY / 4);
+ }
+
+ @Test
+ public void rebufferGrowsOnMoreDelivered() {
+ consumer.accept(ImmutableList.of(messageWithSize(MAX_MEMORY / 4)));
+ bufferedSubscriber.pop();
+ bufferedSubscriber.rebuffer();
+ verify(limiter).claim(3 * MAX_MEMORY / 4);
+ consumer.accept(
+ ImmutableList.of(messageWithSize(MAX_MEMORY / 2),
messageWithSize(MAX_MEMORY / 8)));
+ bufferedSubscriber.rebuffer();
+ verify(limiter, times(2)).claim(MAX_MEMORY); // once in setup
+ }
+
+ @Test
+ public void onDataSatisfiedOnShutdown() throws Exception {
+ Future<Void> onData = bufferedSubscriber.onData();
+ bufferedSubscriber.stopAsync();
+ onData.get();
+ }
+
+ @Test
+ public void onDataSatisfiedOnError() throws Exception {
+ Future<Void> onData = bufferedSubscriber.onData();
+ subscriber.fail(new RuntimeException("bad"));
+ onData.get();
+ }
+
+ @Test
+ public void onDataSatisfiedOnData() throws Exception {
+ SequencedMessage message1 =
+ SequencedMessage.newBuilder()
+
.setCursor(Cursor.newBuilder().setOffset(example(Offset.class).value() + 10))
+ .setSizeBytes(1)
+ .build();
+ SequencedMessage message2 =
+ SequencedMessage.newBuilder()
+
.setCursor(Cursor.newBuilder().setOffset(example(Offset.class).value() + 20))
+ .setSizeBytes(1)
+ .build();
+ Future<Void> onData = bufferedSubscriber.onData();
+ assertFalse(onData.isDone());
+ consumer.accept(ImmutableList.of(message1, message2));
+ onData.get();
+ assertTrue(bufferedSubscriber.onData().isDone()); // Still messages,
onData is satisfied.
+ assertEquals(bufferedSubscriber.fetchOffset(), example(Offset.class));
+ assertEquals(bufferedSubscriber.peek().get(), message1);
+ bufferedSubscriber.pop();
+ assertEquals(bufferedSubscriber.fetchOffset(),
Offset.of(message1.getCursor().getOffset() + 1));
+ assertTrue(bufferedSubscriber.onData().isDone()); // Still messages,
onData is satisfied.
+ assertEquals(bufferedSubscriber.peek().get(), message2);
+ bufferedSubscriber.pop();
+ assertEquals(bufferedSubscriber.fetchOffset(),
Offset.of(message2.getCursor().getOffset() + 1));
+ assertFalse(bufferedSubscriber.onData().isDone());
+ }
+}
diff --git
a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/OffsetByteRangeTrackerTest.java
b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/OffsetByteRangeTrackerTest.java
index f7d206d..79585ab 100644
---
a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/OffsetByteRangeTrackerTest.java
+++
b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/OffsetByteRangeTrackerTest.java
@@ -19,6 +19,7 @@ package org.apache.beam.sdk.io.gcp.pubsublite.internal;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertThrows;
import static org.junit.Assert.assertTrue;
@@ -34,9 +35,7 @@ import
com.google.cloud.pubsublite.proto.ComputeMessageStatsResponse;
import org.apache.beam.sdk.io.range.OffsetRange;
import
org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.Progress;
import org.apache.beam.sdk.transforms.splittabledofn.SplitResult;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Stopwatch;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Ticker;
-import org.joda.time.Duration;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
@@ -58,13 +57,7 @@ public class OffsetByteRangeTrackerTest {
public void setUp() {
initMocks(this);
when(ticker.read()).thenReturn(0L);
- tracker =
- new OffsetByteRangeTracker(
- OffsetByteRange.of(RANGE, 0),
- unownedBacklogReader,
- Stopwatch.createUnstarted(ticker),
- Duration.millis(500),
- MIN_BYTES);
+ tracker = new OffsetByteRangeTracker(OffsetByteRange.of(RANGE, 0),
unownedBacklogReader);
}
@Test
@@ -136,31 +129,7 @@ public class OffsetByteRangeTrackerTest {
@Test
public void cannotClaimSplitRange() {
assertTrue(tracker.tryClaim(OffsetByteProgress.of(Offset.of(1_000),
MIN_BYTES)));
- assertTrue(tracker.trySplit(IGNORED_FRACTION) != null);
+ assertNotNull(tracker.trySplit(IGNORED_FRACTION));
assertFalse(tracker.tryClaim(OffsetByteProgress.of(Offset.of(1_001),
MIN_BYTES)));
}
-
- @Test
- public void cannotSplitNotEnoughBytesOrTime() {
- assertTrue(tracker.tryClaim(OffsetByteProgress.of(Offset.of(1_000),
MIN_BYTES - 2)));
- assertTrue(tracker.tryClaim(OffsetByteProgress.of(Offset.of(1_001), 1)));
- when(ticker.read()).thenReturn(100_000_000L);
- assertTrue(tracker.trySplit(IGNORED_FRACTION) == null);
- }
-
- @Test
- public void canSplitTimeOnly() {
- assertTrue(tracker.tryClaim(OffsetByteProgress.of(Offset.of(1_000),
MIN_BYTES - 2)));
- assertTrue(tracker.tryClaim(OffsetByteProgress.of(Offset.of(1_001), 1)));
- when(ticker.read()).thenReturn(1_000_000_000L);
- assertTrue(tracker.trySplit(IGNORED_FRACTION) != null);
- }
-
- @Test
- public void canSplitBytesOnly() {
- assertTrue(tracker.tryClaim(OffsetByteProgress.of(Offset.of(1_000),
MIN_BYTES - 2)));
- assertTrue(tracker.tryClaim(OffsetByteProgress.of(Offset.of(1_001), 2)));
- when(ticker.read()).thenReturn(100_000_000L);
- assertTrue(tracker.trySplit(IGNORED_FRACTION) != null);
- }
}
diff --git
a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/SubscriptionPartitionProcessorImplTest.java
b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/SubscriptionPartitionProcessorImplTest.java
index 5977560..679861b 100644
---
a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/SubscriptionPartitionProcessorImplTest.java
+++
b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/SubscriptionPartitionProcessorImplTest.java
@@ -18,40 +18,37 @@
package org.apache.beam.sdk.io.gcp.pubsublite.internal;
import static
com.google.cloud.pubsublite.internal.testing.UnitTestExamples.example;
-import static
org.apache.beam.sdk.io.gcp.pubsublite.SubscriberOptions.DEFAULT_FLOW_CONTROL;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertThrows;
import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.inOrder;
+import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
import static org.mockito.MockitoAnnotations.initMocks;
+import com.google.api.core.ApiFutures;
import com.google.api.core.SettableApiFuture;
import com.google.api.gax.rpc.ApiException;
import com.google.api.gax.rpc.StatusCode.Code;
import com.google.cloud.pubsublite.Offset;
+import com.google.cloud.pubsublite.Partition;
+import com.google.cloud.pubsublite.SubscriptionPath;
import com.google.cloud.pubsublite.internal.CheckedApiException;
import com.google.cloud.pubsublite.internal.testing.FakeApiService;
-import com.google.cloud.pubsublite.internal.wire.Subscriber;
-import com.google.cloud.pubsublite.internal.wire.SystemExecutors;
import com.google.cloud.pubsublite.proto.Cursor;
-import com.google.cloud.pubsublite.proto.FlowControlRequest;
import com.google.cloud.pubsublite.proto.SequencedMessage;
import com.google.protobuf.util.Timestamps;
-import java.util.List;
-import java.util.function.Consumer;
-import java.util.function.Function;
+import java.util.Optional;
+import java.util.function.Supplier;
import org.apache.beam.sdk.io.range.OffsetRange;
import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
import org.apache.beam.sdk.transforms.DoFn.ProcessContinuation;
import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
-import
org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.joda.time.Duration;
import org.joda.time.Instant;
import org.junit.Before;
@@ -63,24 +60,23 @@ import org.junit.runners.JUnit4;
import org.mockito.InOrder;
import org.mockito.Mock;
import org.mockito.Spy;
-import org.mockito.stubbing.Answer;
@RunWith(JUnit4.class)
@SuppressWarnings("initialization.fields.uninitialized")
public class SubscriptionPartitionProcessorImplTest {
+ private static final SubscriptionPartition PARTITION =
+ SubscriptionPartition.of(example(SubscriptionPath.class),
example(Partition.class));
+
@Spy RestrictionTracker<OffsetByteRange, OffsetByteProgress> tracker;
@Mock OutputReceiver<SequencedMessage> receiver;
- @Mock Function<Consumer<List<SequencedMessage>>, Subscriber>
subscriberFactory;
+ @Mock Supplier<MemoryBufferedSubscriber> subscriberFactory;
@Rule public Timeout globalTimeout = Timeout.seconds(30);
- abstract static class FakeSubscriber extends FakeApiService implements
Subscriber {}
+ abstract static class FakeSubscriber extends FakeApiService implements
MemoryBufferedSubscriber {}
@Spy FakeSubscriber subscriber;
- Consumer<List<SequencedMessage>> leakedConsumer;
- SubscriptionPartitionProcessor processor;
-
private static SequencedMessage messageWithOffset(long offset) {
return SequencedMessage.newBuilder()
.setCursor(Cursor.newBuilder().setOffset(offset))
@@ -96,134 +92,106 @@ public class SubscriptionPartitionProcessorImplTest {
@Before
public void setUp() {
initMocks(this);
- when(subscriberFactory.apply(any()))
- .then(
- args -> {
- leakedConsumer = args.getArgument(0);
- return subscriber;
- });
- processor =
- new SubscriptionPartitionProcessorImpl(
- tracker, receiver, subscriberFactory, DEFAULT_FLOW_CONTROL);
- assertNotNull(leakedConsumer);
+ when(subscriberFactory.get()).thenReturn(subscriber);
+ when(tracker.currentRestriction()).thenReturn(initialRange());
+ doReturn(true).when(subscriber).isRunning();
+ doReturn(example(Offset.class)).when(subscriber).fetchOffset();
+ doReturn(SettableApiFuture.create()).when(subscriber).onData();
+ }
+
+ private SubscriptionPartitionProcessor newProcessor() {
+ return new SubscriptionPartitionProcessorImpl(PARTITION, tracker,
receiver, subscriberFactory);
}
@Test
- public void lifecycle() throws Exception {
- when(tracker.currentRestriction()).thenReturn(initialRange());
+ public void lifecycle() {
+ SubscriptionPartitionProcessor processor = newProcessor();
assertEquals(ProcessContinuation.resume(),
processor.runFor(Duration.millis(10)));
- InOrder order = inOrder(subscriber);
- order.verify(subscriber).startAsync();
- order.verify(subscriber).awaitRunning();
- order
- .verify(subscriber)
- .allowFlow(
- FlowControlRequest.newBuilder()
- .setAllowedBytes(DEFAULT_FLOW_CONTROL.bytesOutstanding())
- .setAllowedMessages(DEFAULT_FLOW_CONTROL.messagesOutstanding())
- .build());
- order.verify(subscriber).stopAsync();
- order.verify(subscriber).awaitTerminated();
+ InOrder order = inOrder(subscriberFactory, subscriber);
+ order.verify(subscriberFactory).get();
+ order.verify(subscriber).fetchOffset();
+ order.verify(subscriber).rebuffer();
}
@Test
- public void lifecycleFlowControlThrows() throws Exception {
- when(tracker.currentRestriction()).thenReturn(initialRange());
- doThrow(new
CheckedApiException(Code.OUT_OF_RANGE)).when(subscriber).allowFlow(any());
- assertThrows(ApiException.class, () -> processor.runFor(Duration.ZERO));
+ public void lifecycleOffsetMismatch() {
+ MemoryBufferedSubscriber badSubscriber = spy(FakeSubscriber.class);
+ doReturn(Offset.of(example(Offset.class).value() +
1)).when(badSubscriber).fetchOffset();
+ doThrow(new
RuntimeException("Ignored")).when(badSubscriber).awaitTerminated();
+ doReturn(badSubscriber, subscriber).when(subscriberFactory).get();
+ SubscriptionPartitionProcessor processor = newProcessor();
+ assertEquals(ProcessContinuation.resume(),
processor.runFor(Duration.millis(10)));
+ InOrder order = inOrder(subscriberFactory, badSubscriber, subscriber);
+ order.verify(subscriberFactory).get();
+ order.verify(badSubscriber).fetchOffset();
+ order.verify(badSubscriber).stopAsync();
+ order.verify(badSubscriber).awaitTerminated();
+ order.verify(subscriberFactory).get();
+ order.verify(subscriber).fetchOffset();
+ order.verify(subscriber).rebuffer();
}
@Test
- public void subscriberFailureFails() throws Exception {
- when(tracker.currentRestriction()).thenReturn(initialRange());
- doAnswer(
- (Answer<Void>)
- args -> {
- subscriber.fail(new CheckedApiException(Code.OUT_OF_RANGE));
- return null;
- })
- .when(subscriber)
- .awaitRunning();
- ApiException e =
- assertThrows(
- // Longer wait is needed due to listener asynchrony, but should
never wait this long.
- ApiException.class, () ->
processor.runFor(Duration.standardMinutes(2)));
- assertEquals(Code.OUT_OF_RANGE, e.getStatusCode().getCode());
+ public void lifecycleRebufferThrows() throws Exception {
+ doThrow(new
CheckedApiException(Code.OUT_OF_RANGE).underlying).when(subscriber).rebuffer();
+ assertThrows(ApiException.class, this::newProcessor);
}
@Test
- public void allowFlowFailureFails() throws Exception {
- when(tracker.currentRestriction()).thenReturn(initialRange());
- when(tracker.tryClaim(any())).thenReturn(true);
- doThrow(new
CheckedApiException(Code.OUT_OF_RANGE)).when(subscriber).allowFlow(any());
- SystemExecutors.getFuturesExecutor()
- .execute(() ->
leakedConsumer.accept(ImmutableList.of(messageWithOffset(1))));
- ApiException e =
- assertThrows(ApiException.class, () ->
processor.runFor(Duration.standardHours(10)));
- assertEquals(Code.OUT_OF_RANGE, e.getStatusCode().getCode());
+ public void subscriberFailureReturnsResume() throws Exception {
+ SubscriptionPartitionProcessor processor = newProcessor();
+ doReturn(ApiFutures.immediateFuture(null)).when(subscriber).onData();
+ doReturn(false).when(subscriber).isRunning();
+ assertEquals(ProcessContinuation.resume(),
processor.runFor(Duration.standardHours(1)));
}
@Test
public void timeoutReturnsResume() {
+ SubscriptionPartitionProcessor processor = newProcessor();
assertEquals(ProcessContinuation.resume(),
processor.runFor(Duration.millis(10)));
assertFalse(processor.lastClaimed().isPresent());
}
@Test
- public void failedClaimCausesStop() throws Exception {
+ public void failedClaimCausesStop() {
+ SubscriptionPartitionProcessor processor = newProcessor();
+
when(tracker.tryClaim(any())).thenReturn(false);
- SettableApiFuture<Void> runDone = SettableApiFuture.create();
- SystemExecutors.getFuturesExecutor()
- .execute(
- () -> {
- assertEquals(
- ProcessContinuation.stop(),
processor.runFor(Duration.standardHours(10)));
- runDone.set(null);
- });
- leakedConsumer.accept(ImmutableList.of(messageWithOffset(1)));
- runDone.get();
+ doReturn(ApiFutures.immediateFuture(null)).when(subscriber).onData();
+ doReturn(Optional.of(messageWithOffset(1))).when(subscriber).peek();
+
+ assertEquals(ProcessContinuation.stop(),
processor.runFor(Duration.standardHours(10)));
verify(tracker, times(1)).tryClaim(any());
+ verify(subscriber, times(0)).pop();
assertFalse(processor.lastClaimed().isPresent());
- // Future calls to process don't try to claim.
- leakedConsumer.accept(ImmutableList.of(messageWithOffset(2)));
- verify(tracker, times(1)).tryClaim(any());
}
@Test
- public void successfulClaimThenTimeout() throws Exception {
- when(tracker.tryClaim(any())).thenReturn(true);
- SettableApiFuture<Void> runDone = SettableApiFuture.create();
- SystemExecutors.getFuturesExecutor()
- .execute(
- () -> {
- assertEquals(
- ProcessContinuation.resume(),
processor.runFor(Duration.standardSeconds(3)));
- runDone.set(null);
- });
+ public void successfulClaimThenTimeout() {
+ doReturn(true).when(tracker).tryClaim(any());
+ doReturn(ApiFutures.immediateFuture(null), SettableApiFuture.create())
+ .when(subscriber)
+ .onData();
SequencedMessage message1 = messageWithOffset(1);
SequencedMessage message3 = messageWithOffset(3);
- leakedConsumer.accept(ImmutableList.of(message1, message3));
- runDone.get();
- InOrder order = inOrder(tracker, receiver, subscriber);
- order
- .verify(tracker)
- .tryClaim(
- OffsetByteProgress.of(Offset.of(3), message1.getSizeBytes() +
message3.getSizeBytes()));
+ doReturn(Optional.of(message1), Optional.of(message3), Optional.empty())
+ .when(subscriber)
+ .peek();
+
+ SubscriptionPartitionProcessor processor = newProcessor();
+ assertEquals(ProcessContinuation.resume(),
processor.runFor(Duration.standardSeconds(3)));
+
+ InOrder order = inOrder(tracker, receiver);
+ order.verify(tracker).tryClaim(OffsetByteProgress.of(Offset.of(1),
message1.getSizeBytes()));
order
.verify(receiver)
.outputWithTimestamp(message1, new
Instant(Timestamps.toMillis(message1.getPublishTime())));
+ order.verify(tracker).tryClaim(OffsetByteProgress.of(Offset.of(3),
message3.getSizeBytes()));
order
.verify(receiver)
.outputWithTimestamp(message3, new
Instant(Timestamps.toMillis(message3.getPublishTime())));
- order
- .verify(subscriber)
- .allowFlow(
- FlowControlRequest.newBuilder()
- .setAllowedMessages(2)
- .setAllowedBytes(message1.getSizeBytes() +
message3.getSizeBytes())
- .build());
assertEquals(processor.lastClaimed().get(), Offset.of(3));
}
}
