ackIds = ackedIds.poll();
+ if (ackIds == null) {
+ return;
+ }
+ numAcked.add(nowMsSinceEpoch, ackIds.size());
+ for (String ackId : ackIds) {
+ inFlight.remove(ackId);
+ }
+ }
+ }
+
+ /**
+ * BLOCKING
+ * Extend deadline for all messages which need it.
+ *
+ * CAUTION: If extensions can't keep up with wallclock then we'll never return.
+ */
+ private void extend() throws IOException {
+ while (true) {
+ long nowMsSinceEpoch = System.currentTimeMillis();
+ List> toBeExtended = new ArrayList<>();
+ for (Map.Entry entry : inFlight.entrySet()) {
+ long deadlineMsSinceEpoch = entry.getValue().getMillis();
+ if (deadlineMsSinceEpoch < nowMsSinceEpoch + ACK_TOO_LATE.getMillis()) {
+ // This message may have already expired, in which case it will (eventually) be
+ // made available on a future pull request.
+ // If the message ends up being commited then a future resend will be ignored
+ // downsteam and acked as usual.
+ numLateDeadlines.add(nowMsSinceEpoch, 1);
+ }
+ if (deadlineMsSinceEpoch > nowMsSinceEpoch + ACK_SAFETY.getMillis()) {
+ // No later messages need extending.
+ break;
+ }
+ toBeExtended.add(entry);
+ if (toBeExtended.size() >= ACK_BATCH_SIZE) {
+ // Enough for one batch.
+ break;
+ }
+ }
+ if (toBeExtended.isEmpty()) {
+ // No messages need to be extended.
+ return;
+ }
+ List extensionAckIds = new ArrayList<>(toBeExtended.size());
+ long newDeadlineMsSinceEpoch = nowMsSinceEpoch + ACK_EXTENSION.getMillis();
+ for (Map.Entry entry : toBeExtended) {
+ inFlight.remove(entry.getKey());
+ inFlight.put(entry.getKey(), new Instant(newDeadlineMsSinceEpoch));
+ extensionAckIds.add(entry.getKey());
+ }
+ // BLOCKs until extended.
+ extendBatch(nowMsSinceEpoch, extensionAckIds);
+ }
+ }
+
+ /**
+ * BLOCKING
+ * Fetch another batch of messages from pub/sub.
+ */
+ private void pull() throws IOException {
+ if (inFlight.size() >= MAX_IN_FLIGHT) {
+ // Wait for checkpoint to be finalized before pulling anymore.
+ // There may be lag while checkpoints are persisted and the finalizeCheckpoint method
+ // is invoked. By limiting the in-flight messages we can ensure we don't end up consuming
+ // messages faster than we can checkpoint them.
+ return;
+ }
+
+ long requestTimeMsSinceEpoch = System.currentTimeMillis();
+ long deadlineMsSinceEpoch = requestTimeMsSinceEpoch + ACK_TIMEOUT.getMillis();
+
+ // Pull the next batch.
+ // BLOCKs until received.
+ Collection receivedMessages =
+ pubsubClient.pull(requestTimeMsSinceEpoch,
+ outer.outer.subscription,
+ PULL_BATCH_SIZE);
+ if (receivedMessages.isEmpty()) {
+ // Nothing available yet. Try again later.
+ return;
+ }
+
+ lastReceivedMsSinceEpoch = requestTimeMsSinceEpoch;
+
+ // Capture the received messages.
+ for (PubsubClient.IncomingMessage incomingMessage : receivedMessages) {
+ notYetRead.add(incomingMessage);
+ notYetReadBytes += incomingMessage.elementBytes.length;
+ inFlight.put(incomingMessage.ackId, new Instant(deadlineMsSinceEpoch));
+ numReceived++;
+ numReceivedRecently.add(requestTimeMsSinceEpoch, 1L);
+ minReceivedTimestampMsSinceEpoch.add(requestTimeMsSinceEpoch,
+ incomingMessage.timestampMsSinceEpoch);
+ maxReceivedTimestampMsSinceEpoch.add(requestTimeMsSinceEpoch,
+ incomingMessage.timestampMsSinceEpoch);
+ minUnreadTimestampMsSinceEpoch.add(requestTimeMsSinceEpoch,
+ incomingMessage.timestampMsSinceEpoch);
+ }
+ }
+
+ /**
+ * Log stats if time to do so.
+ */
+ private void stats() {
+ long nowMsSinceEpoch = System.currentTimeMillis();
+ if (lastLogTimestampMsSinceEpoch < 0) {
+ lastLogTimestampMsSinceEpoch = nowMsSinceEpoch;
+ return;
+ }
+ long deltaMs = nowMsSinceEpoch - lastLogTimestampMsSinceEpoch;
+ if (deltaMs < LOG_PERIOD.getMillis()) {
+ return;
+ }
+
+ String messageSkew = "unknown";
+ long minTimestamp = minReceivedTimestampMsSinceEpoch.get(nowMsSinceEpoch);
+ long maxTimestamp = maxReceivedTimestampMsSinceEpoch.get(nowMsSinceEpoch);
+ if (minTimestamp < Long.MAX_VALUE && maxTimestamp > Long.MIN_VALUE) {
+ messageSkew = (maxTimestamp - minTimestamp) + "ms";
+ }
+
+ String watermarkSkew = "unknown";
+ long minWatermark = minWatermarkMsSinceEpoch.get(nowMsSinceEpoch);
+ long maxWatermark = maxWatermarkMsSinceEpoch.get(nowMsSinceEpoch);
+ if (minWatermark < Long.MAX_VALUE && maxWatermark > Long.MIN_VALUE) {
+ watermarkSkew = (maxWatermark - minWatermark) + "ms";
+ }
+
+ LOG.warn("Pubsub {} has "
+ + "{} received messages, "
+ + "{} current unread messages, "
+ + "{} current unread bytes, "
+ + "{} current in-flight msgs, "
+ + "{} current in-flight checkpoints, "
+ + "{} max in-flight checkpoints, "
+ + "{}B/s recent read, "
+ + "{} recent received, "
+ + "{} recent extended, "
+ + "{} recent late extended, "
+ + "{} recent acked, "
+ + "{} recent nacked, "
+ + "{} recent message timestamp skew, "
+ + "{} recent watermark skew, "
+ + "{} recent late messages, "
+ + "{} last reported watermark",
+ outer.outer.subscription,
+ numReceived,
+ notYetRead.size(),
+ notYetReadBytes,
+ inFlight.size(),
+ numInFlightCheckpoints.get(),
+ maxInFlightCheckpoints,
+ numReadBytes.get(nowMsSinceEpoch) / (SAMPLE_PERIOD.getMillis() / 1000L),
+ numReceivedRecently.get(nowMsSinceEpoch),
+ numExtendedDeadlines.get(nowMsSinceEpoch),
+ numLateDeadlines.get(nowMsSinceEpoch),
+ numAcked.get(nowMsSinceEpoch),
+ numNacked.get(nowMsSinceEpoch),
+ messageSkew,
+ watermarkSkew,
+ numLateMessages.get(nowMsSinceEpoch),
+ new Instant(lastWatermarkMsSinceEpoch));
+
+ lastLogTimestampMsSinceEpoch = nowMsSinceEpoch;
+ }
+
+ @Override
+ public boolean start() throws IOException {
+ return advance();
+ }
+
+ /**
+ * BLOCKING
+ * Return {@literal true} if a pub/sub messaage is available, {@literal false} if
+ * none is available at this time or we are over-subscribed. May BLOCK while extending
+ * acks or fetching available messages. Will not block waiting for messages.
+ */
+ @Override
+ public boolean advance() throws IOException {
+ // Emit stats.
+ stats();
+
+ if (current != null) {
+ // Current is consumed. It can no longer contribute to holding back the watermark.
+ minUnreadTimestampMsSinceEpoch.remove(current.requestTimeMsSinceEpoch);
+ current = null;
+ }
+
+ // Retire state associated with acked messages.
+ retire();
+
+ // Extend all pressing deadlines.
+ // Will BLOCK until done.
+ // If the system is pulling messages only to let them sit in a downsteam queue then
+ // this will have the effect of slowing down the pull rate.
+ // However, if the system is genuinely taking longer to process each message then
+ // the work to extend acks would be better done in the background.
+ extend();
+
+ if (notYetRead.isEmpty()) {
+ // Pull another batch.
+ // Will BLOCK until fetch returns, but will not block until a message is available.
+ pull();
+ }
+
+ // Take one message from queue.
+ current = notYetRead.poll();
+ if (current == null) {
+ // Try again later.
+ return false;
+ }
+ notYetReadBytes -= current.elementBytes.length;
+ Preconditions.checkState(notYetReadBytes >= 0);
+ long nowMsSinceEpoch = System.currentTimeMillis();
+ numReadBytes.add(nowMsSinceEpoch, current.elementBytes.length);
+ minReadTimestampMsSinceEpoch.add(nowMsSinceEpoch, current.timestampMsSinceEpoch);
+ if (current.timestampMsSinceEpoch < lastWatermarkMsSinceEpoch) {
+ numLateMessages.add(nowMsSinceEpoch, 1L);
+ }
+
+ // Current message will be persisted by the next checkpoint so it is now safe to ack.
+ safeToAckIds.add(current.ackId);
+ return true;
+ }
+
+ @Override
+ public T getCurrent() throws NoSuchElementException {
+ if (current == null) {
+ throw new NoSuchElementException();
+ }
+ try {
+ return CoderUtils.decodeFromByteArray(outer.outer.elementCoder, current.elementBytes);
+ } catch (CoderException e) {
+ throw new RuntimeException("Unable to decode element from pub/sub message: ", e);
+ }
+ }
+
+ @Override
+ public Instant getCurrentTimestamp() throws NoSuchElementException {
+ if (current == null) {
+ throw new NoSuchElementException();
+ }
+ return new Instant(current.timestampMsSinceEpoch);
+ }
+
+ @Override
+ public byte[] getCurrentRecordId() throws NoSuchElementException {
+ if (current == null) {
+ throw new NoSuchElementException();
+ }
+ return current.recordId;
+ }
+
+ @Override
+ public void close() throws IOException {
+ pubsubClient.close();
+ pubsubClient = null;
+ }
+
+ @Override
+ public Source getCurrentSource() {
+ return outer;
+ }
+
+ @Override
+ public Instant getWatermark() {
+ // NOTE: We'll allow the watermark to go backwards. The underlying runner is responsible
+ // for aggregating all reported watermarks and ensuring the aggregate is latched.
+ // If we attempt to latch locally then it is possible a temporary starvation of one reader
+ // could cause its estimated watermark to fast forward to current system time. Then when
+ // the reader resumes its watermark would be unable to resume tracking.
+ // By letting the underlying runner latch we avoid any problems due to localized starvation.
+ long nowMsSinceEpoch = System.currentTimeMillis();
+ long readMin = minReadTimestampMsSinceEpoch.get(nowMsSinceEpoch);
+ long unreadMin = minUnreadTimestampMsSinceEpoch.get();
+ if (readMin == Long.MAX_VALUE &&
+ unreadMin == Long.MAX_VALUE &&
+ lastReceivedMsSinceEpoch >= 0 &&
+ nowMsSinceEpoch > lastReceivedMsSinceEpoch + SAMPLE_PERIOD.getMillis()) {
+ // We don't currently have any unread messages pending, we have not had any messages
+ // read for a while, and we have not received any new messages from pub/sub for a while.
+ // Advance watermark to current time.
+ // TODO: Estimate a timestamp lag.
+ lastWatermarkMsSinceEpoch = nowMsSinceEpoch;
+ } else if (minReadTimestampMsSinceEpoch.isSignificant() ||
+ minUnreadTimestampMsSinceEpoch.isSignificant()) {
+ // Take minimum of the timestamps in all unread messages and recently read messages.
+ lastWatermarkMsSinceEpoch = Math.min(readMin, unreadMin);
+ }
+ // else: We're not confident enough to estimate a new watermark. Stick with the old one.
+ minWatermarkMsSinceEpoch.add(nowMsSinceEpoch, lastWatermarkMsSinceEpoch);
+ maxWatermarkMsSinceEpoch.add(nowMsSinceEpoch, lastWatermarkMsSinceEpoch);
+ return new Instant(lastWatermarkMsSinceEpoch);
+ }
+
+ @Override
+ public UnboundedSource.CheckpointMark getCheckpointMark() {
+ int cur = numInFlightCheckpoints.incrementAndGet();
+ maxInFlightCheckpoints = Math.max(maxInFlightCheckpoints, cur);
+ // The checkpoint will either be finalized or we'll rollback to an earlier
+ // checkpoint. Thus we can hand off these ack ids to the checkpoint.
+ List snapshotSafeToAckIds = safeToAckIds;
+ safeToAckIds = new ArrayList<>();
+ List snapshotNotYetReadIds = new ArrayList<>(notYetRead.size());
+ for (PubsubClient.IncomingMessage incomingMessage : notYetRead) {
+ snapshotNotYetReadIds.add(incomingMessage.ackId);
+ }
+ return new Checkpoint<>(this, snapshotSafeToAckIds, snapshotNotYetReadIds);
+ }
+
+ @Override
+ public long getSplitBacklogBytes() {
+ return notYetReadBytes;
+ }
+ }
+
+ // ================================================================================
+ // Source
+ // ================================================================================
+
+ private static class Source extends UnboundedSource> {
+ PubsubUnboundedSource outer;
+
+ public Source(PubsubUnboundedSource outer) {
+ this.outer = outer;
+ }
+
+ @Override
+ public List> generateInitialSplits(
+ int desiredNumSplits, PipelineOptions options) throws Exception {
+ List> result = new ArrayList<>(desiredNumSplits);
+ for (int i = 0; i < desiredNumSplits * SCALE_OUT; i++) {
+ // Since the source is immutable and pub/sub automatically shards we simply
+ // replicate ourselves the requested number of times
+ result.add(this);
+ }
+ return result;
+ }
+
+ @Override
+ public UnboundedReader createReader(
+ PipelineOptions options, @Nullable Checkpoint checkpoint) {
+ PubsubUnboundedSource.Reader reader;
+ try {
+ reader =
+ new PubsubUnboundedSource.Reader<>(options.as(GcpOptions.class), this);
+ } catch (GeneralSecurityException | IOException e) {
+ throw new RuntimeException("Unable to subscribe to " + outer.subscription + ": ", e);
+ }
+ if (checkpoint != null) {
+ // Nack all messages we may have lost.
+ try {
+ checkpoint.setReader(reader);
+ // Will BLOCK until nacked.
+ checkpoint.nackAll();
+ } catch (IOException e) {
+ LOG.error("Pubsub {} cannot have {} lost messages nacked, ignoring: {}",
+ outer.subscription, checkpoint.notYetReadIds.size(), e);
+ }
+ }
+ return reader;
+ }
+
+ @Nullable
+ @Override
+ public Coder> getCheckpointMarkCoder() {
+ @SuppressWarnings("unchecked") CheckpointCoder typedCoder =
+ (CheckpointCoder) CHECKPOINT_CODER;
+ return typedCoder;
+ }
+
+ @Override
+ public Coder getDefaultOutputCoder() {
+ return outer.elementCoder;
+ }
+
+ @Override
+ public void validate() {
+ // Nothing to validate.
+ }
+
+ @Override
+ public boolean requiresDeduping() {
+ // We cannot prevent re-offering already read messages after a restore from checkpoint.
+ return true;
+ }
+ }
+
+ // ================================================================================
+ // StatsFn
+ // ================================================================================
+
+ private static class StatsFn extends DoFn {
+ private final Aggregator elementCounter =
+ createAggregator("elements", new Sum.SumLongFn());
+
+
+ @Override
+ public void processElement(ProcessContext c) throws Exception {
+ elementCounter.addValue(1L);
+ c.output(c.element());
+ }
+ }
+
+ // ================================================================================
+ // PubsubUnboundedSource
+ // ================================================================================
+
+ /**
+ * Subscription to read from.
+ */
+ private final String subscription;
+
+ /**
+ * Coder for elements. Elements are effectively double-encoded: first to a byte array
+ * using this checkpointCoder, then to a base-64 string to conform to pub/sub's payload
+ * conventions.
+ */
+ private final Coder elementCoder;
+
+ /**
+ * Pub/sub metadata field holding timestamp of each element, or {@literal null} if should use
+ * pub/sub message publish timestamp instead.
+ */
+ @Nullable
+ private final String timestampLabel;
+
+ /**
+ * Pub/sub metadata field holding id for each element, or {@literal null} if need to generate
+ * a unique id ourselves.
+ */
+ @Nullable
+ private final String idLabel;
+
+ /**
+ * Construct an unbounded source to consume from the pub/sub {@code subscription}.
+ */
+ public PubsubUnboundedSource(
+ String subscription, Coder elementCoder, @Nullable String timestampLabel,
+ @Nullable String idLabel) {
+ this.subscription = Preconditions.checkNotNull(subscription);
+ this.elementCoder = Preconditions.checkNotNull(elementCoder);
+ this.timestampLabel = timestampLabel;
+ this.idLabel = idLabel;
+ }
+
+ @Override
+ public PCollection apply(PBegin input) {
+ String label = "PubsubSource(" + subscription.replaceFirst(".*/", "") + ")";
+ return input.getPipeline().begin()
+ .apply(Read.from(new Source(this)))
+ .apply(ParDo.named(label).of(new StatsFn()));
+ }
+}
diff --git a/sdks/java/core/src/main/java/com/google/cloud/dataflow/sdk/io/UnboundedSource.java b/sdks/java/core/src/main/java/com/google/cloud/dataflow/sdk/io/UnboundedSource.java
index 67b7f1067adc..096888a5f110 100644
--- a/sdks/java/core/src/main/java/com/google/cloud/dataflow/sdk/io/UnboundedSource.java
+++ b/sdks/java/core/src/main/java/com/google/cloud/dataflow/sdk/io/UnboundedSource.java
@@ -110,8 +110,21 @@ public interface CheckpointMark {
* For example, this could be sending acknowledgement requests to an external
* data source such as Pub/Sub.
*
- *
This may be called from any thread, potentially at the same time as calls to the
- * {@code UnboundedReader} that created it.
+ *
Note that:
+ *
+ * - This finalize method may be called from any thread, potentially at the same time as
+ * calls to the {@link UnboundedReader} it was created from.
+ *
- Checkpoints will be finalized in the same order they are created from the
+ * {@link UnboundedReader}.
+ *
- It is possible for multiple checkpoints created from the same {@link UnboundedReader}
+ * to be 'in-flight' (ie not yet finalized) simultaneously.
+ *
- If this call throws an exception then the entire checkpoint will be abandoned and the
+ * reader restarted from an earlier, successfully-finalized checkpoint.
+ *
- The system ensures that if a checkpoint fails for any reason then no later checkpoint
+ * is allowed to be finalized without the reader first being restarted.
+ *
- It is not safe to assume the {@link UnboundedReader} from which this checkpoint was
+ * created still exists at the time this method is called.
+ *
*/
void finalizeCheckpoint() throws IOException;
}
@@ -205,10 +218,10 @@ public byte[] getCurrentRecordId() throws NoSuchElementException {
/**
* Returns a {@link CheckpointMark} representing the progress of this {@code UnboundedReader}.
*
- * The elements read up until this is called will be processed together as a bundle. Once
- * the result of this processing has been durably committed,
- * {@link CheckpointMark#finalizeCheckpoint} will be called on the {@link CheckpointMark}
- * object.
+ *
All elements read up until this method is called will be processed together as a bundle.
+ * Once the result of processing those elements and the returned checkpoint have been durably
+ * committed, {@link CheckpointMark#finalizeCheckpoint} will be called on the
+ * returned {@link CheckpointMark} object.
*
*
The returned object should not be modified.
*
diff --git a/sdks/java/core/src/main/java/com/google/cloud/dataflow/sdk/util/BucketingFunction.java b/sdks/java/core/src/main/java/com/google/cloud/dataflow/sdk/util/BucketingFunction.java
new file mode 100644
index 000000000000..56bec8db6fdb
--- /dev/null
+++ b/sdks/java/core/src/main/java/com/google/cloud/dataflow/sdk/util/BucketingFunction.java
@@ -0,0 +1,152 @@
+/*
+ * 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 com.google.cloud.dataflow.sdk.util;
+
+import com.google.cloud.dataflow.sdk.transforms.Combine;
+import com.google.common.base.Preconditions;
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * Keep track of the minimum/maximum/sum of a set of timestamped long values.
+ * For efficiency, bucket values by their timestamp.
+ */
+public class BucketingFunction {
+ private static class Bucket {
+ private int numSamples;
+ private long combinedValue;
+
+ public Bucket(BucketingFunction outer) {
+ numSamples = 0;
+ combinedValue = outer.function.identity();
+ }
+
+ public void add(BucketingFunction outer, long value) {
+ combinedValue = outer.function.apply(combinedValue, value);
+ numSamples++;
+ }
+
+ public boolean remove() {
+ numSamples--;
+ Preconditions.checkState(numSamples >= 0);
+ return numSamples == 0;
+ }
+
+ public long get() {
+ return combinedValue;
+ }
+ }
+
+ /**
+ * How large a time interval to fit within each bucket.
+ */
+ private final long bucketWidthMs;
+
+ /**
+ * How many buckets are considered 'significant'?
+ */
+ private final int numSignificantBuckets;
+
+ /**
+ * How many samples are considered 'significant'?
+ */
+ private final int numSignificantSamples;
+
+ /**
+ * Function for combining sample values.
+ */
+ private final Combine.BinaryCombineLongFn function;
+
+ /**
+ * Active buckets.
+ */
+ private final Map buckets;
+
+ public BucketingFunction(
+ long bucketWidthMs,
+ int numSignificantBuckets,
+ int numSignificantSamples,
+ Combine.BinaryCombineLongFn function) {
+ this.bucketWidthMs = bucketWidthMs;
+ this.numSignificantBuckets = numSignificantBuckets;
+ this.numSignificantSamples = numSignificantSamples;
+ this.function = function;
+ this.buckets = new HashMap<>();
+ }
+
+ /**
+ * Which bucket key corresponds to {@code timeMsSinceEpoch}.
+ */
+ private long key(long timeMsSinceEpoch) {
+ return timeMsSinceEpoch - (timeMsSinceEpoch % bucketWidthMs);
+ }
+
+ /**
+ * Add one sample of {@code value} (to bucket) at {@code timeMsSinceEpoch}.
+ */
+ public void add(long timeMsSinceEpoch, long value) {
+ long key = key(timeMsSinceEpoch);
+ Bucket bucket = buckets.get(key);
+ if (bucket == null) {
+ bucket = new Bucket(this);
+ buckets.put(key, bucket);
+ }
+ bucket.add(this, value);
+ }
+
+ /**
+ * Remove one sample (from bucket) at {@code timeMsSinceEpoch}.
+ */
+ public void remove(long timeMsSinceEpoch) {
+ long key = key(timeMsSinceEpoch);
+ Bucket bucket = buckets.get(key);
+ if (bucket == null) {
+ return;
+ }
+ if (bucket.remove()) {
+ buckets.remove(key);
+ }
+ }
+
+ /**
+ * Return the (bucketized) combined value of all samples.
+ */
+ public long get() {
+ long result = function.identity();
+ for (Bucket bucket : buckets.values()) {
+ result = function.apply(result, bucket.get());
+ }
+ return result;
+ }
+
+ /**
+ * Is the current result 'significant'? Ie is it drawn from enough buckets
+ * or from enough samples?
+ */
+ public boolean isSignificant() {
+ if (buckets.size() >= numSignificantBuckets) {
+ return true;
+ }
+ int totalSamples = 0;
+ for (Bucket bucket : buckets.values()) {
+ totalSamples += bucket.numSamples;
+ }
+ return totalSamples >= numSignificantSamples;
+ }
+}
diff --git a/sdks/java/core/src/main/java/com/google/cloud/dataflow/sdk/util/MovingFunction.java b/sdks/java/core/src/main/java/com/google/cloud/dataflow/sdk/util/MovingFunction.java
new file mode 100644
index 000000000000..22827e62512a
--- /dev/null
+++ b/sdks/java/core/src/main/java/com/google/cloud/dataflow/sdk/util/MovingFunction.java
@@ -0,0 +1,152 @@
+/*
+ * 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 com.google.cloud.dataflow.sdk.util;
+
+import com.google.cloud.dataflow.sdk.transforms.Combine;
+import com.google.common.base.Preconditions;
+import java.util.Arrays;
+
+/**
+ * Keep track of the moving minimum/maximum/sum of sampled long values. The minimum/maximum/sum
+ * is over at most the last {@link #samplePeriodMs}, and is updated every
+ * {@link #sampleUpdateMs}.
+ */
+public class MovingFunction {
+ /**
+ * How far back to retain samples, in ms.
+ */
+ private final long samplePeriodMs;
+
+ /**
+ * How frequently to update the moving function, in ms.
+ */
+ private final long sampleUpdateMs;
+
+ /**
+ * How many buckets are considered 'significant'?
+ */
+ private final int numSignificantBuckets;
+
+ /**
+ * How many samples are considered 'significant'?
+ */
+ private final int numSignificantSamples;
+
+ /**
+ * Function for combining sample values.
+ */
+ private final Combine.BinaryCombineLongFn function;
+
+ /**
+ * Minimum/maximum/sum of all values per bucket.
+ */
+ private final long[] buckets;
+
+ /**
+ * How many samples have been added to each bucket.
+ */
+ private final int[] numSamples;
+
+ /**
+ * Time of start of current bucket.
+ */
+ private long currentMsSinceEpoch;
+
+ /**
+ * Index of bucket corresponding to above timestamp, or -1 if no entries.
+ */
+ private int currentIndex;
+
+ public MovingFunction(long samplePeriodMs, long sampleUpdateMs,
+ int numSignificantBuckets, int numSignificantSamples,
+ Combine.BinaryCombineLongFn function) {
+ this.samplePeriodMs = samplePeriodMs;
+ this.sampleUpdateMs = sampleUpdateMs;
+ this.numSignificantBuckets = numSignificantBuckets;
+ this.numSignificantSamples = numSignificantSamples;
+ this.function = function;
+ int n = (int) (samplePeriodMs / sampleUpdateMs);
+ buckets = new long[n];
+ Arrays.fill(buckets, function.identity());
+ numSamples = new int[n];
+ Arrays.fill(numSamples, 0);
+ currentMsSinceEpoch = -1;
+ currentIndex = -1;
+ }
+
+ /**
+ * Flush stale values.
+ */
+ private void flush(long nowMsSinceEpoch) {
+ Preconditions.checkArgument(nowMsSinceEpoch >= 0);
+ if (currentIndex < 0) {
+ currentMsSinceEpoch = nowMsSinceEpoch - (nowMsSinceEpoch % sampleUpdateMs);
+ currentIndex = 0;
+ }
+ Preconditions.checkArgument(nowMsSinceEpoch >= currentMsSinceEpoch);
+ int newBuckets =
+ Math.min((int) ((nowMsSinceEpoch - currentMsSinceEpoch) / sampleUpdateMs),
+ buckets.length);
+ while (newBuckets > 0) {
+ currentIndex = (currentIndex + 1) % buckets.length;
+ buckets[currentIndex] = function.identity();
+ numSamples[currentIndex] = 0;
+ newBuckets--;
+ currentMsSinceEpoch += sampleUpdateMs;
+ }
+ }
+
+ /**
+ * Add {@code value} at {@code nowMsSinceEpoch}.
+ */
+ public void add(long nowMsSinceEpoch, long value) {
+ flush(nowMsSinceEpoch);
+ buckets[currentIndex] = function.apply(buckets[currentIndex], value);
+ numSamples[currentIndex]++;
+ }
+
+ /**
+ * Return the minimum/maximum/sum of all retained values within {@link #samplePeriodMs}
+ * of {@code nowMsSinceEpoch}.
+ */
+ public long get(long nowMsSinceEpoch) {
+ flush(nowMsSinceEpoch);
+ long result = function.identity();
+ for (int i = 0; i < buckets.length; i++) {
+ result = function.apply(result, buckets[i]);
+ }
+ return result;
+ }
+
+ /**
+ * Is the current result 'significant'? Ie is it drawn from enough buckets
+ * or from enough samples?
+ */
+ public boolean isSignificant() {
+ int totalSamples = 0;
+ int activeBuckets = 0;
+ for (int i = 0; i < buckets.length; i++) {
+ totalSamples += numSamples[i];
+ if (numSamples[i] > 0) {
+ activeBuckets++;
+ }
+ }
+ return activeBuckets >= numSignificantBuckets || totalSamples >= numSignificantSamples;
+ }
+}
diff --git a/sdks/java/core/src/test/java/com/google/cloud/dataflow/sdk/util/BucketingFunctionTest.java b/sdks/java/core/src/test/java/com/google/cloud/dataflow/sdk/util/BucketingFunctionTest.java
new file mode 100644
index 000000000000..f5e52e059ec1
--- /dev/null
+++ b/sdks/java/core/src/test/java/com/google/cloud/dataflow/sdk/util/BucketingFunctionTest.java
@@ -0,0 +1,104 @@
+/*
+ * 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 com.google.cloud.dataflow.sdk.util;
+
+import com.google.cloud.dataflow.sdk.transforms.Combine;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/**
+ * Tests {@link BucketingFunction}.
+ */
+@RunWith(JUnit4.class)
+public class BucketingFunctionTest {
+
+ private static final long BUCKET_WIDTH = 10;
+ private static final int SIGNIFICANT_BUCKETS = 2;
+ private static final int SIGNIFICANT_SAMPLES = 10;
+
+ private static final Combine.BinaryCombineLongFn SUM =
+ new Combine.BinaryCombineLongFn() {
+ @Override
+ public long apply(long left, long right) {
+ return left + right;
+ }
+
+ @Override
+ public long identity() {
+ return 0;
+ }
+ };
+
+ private BucketingFunction newFunc() {
+ return new
+ BucketingFunction(BUCKET_WIDTH, SIGNIFICANT_BUCKETS,
+ SIGNIFICANT_SAMPLES, SUM);
+ }
+
+ @Test
+ public void significantSamples() {
+ BucketingFunction f = newFunc();
+ assertFalse(f.isSignificant());
+ for (int i = 0; i < SIGNIFICANT_SAMPLES - 1; i++) {
+ f.add(0, 0);
+ assertFalse(f.isSignificant());
+ }
+ f.add(0, 0);
+ assertTrue(f.isSignificant());
+ }
+
+ @Test
+ public void significantBuckets() {
+ BucketingFunction f = newFunc();
+ assertFalse(f.isSignificant());
+ f.add(0, 0);
+ assertFalse(f.isSignificant());
+ f.add(BUCKET_WIDTH, 0);
+ assertTrue(f.isSignificant());
+ }
+
+ @Test
+ public void sum() {
+ BucketingFunction f = newFunc();
+ for (int i = 0; i < 100; i++) {
+ f.add(i, i);
+ assertEquals(((i + 1) * i) / 2, f.get());
+ }
+ }
+
+ @Test
+ public void movingSum() {
+ BucketingFunction f = newFunc();
+ int lost = 0;
+ for (int i = 0; i < 200; i++) {
+ f.add(i, 1);
+ if (i >= 100) {
+ f.remove(i - 100);
+ if (i % BUCKET_WIDTH == BUCKET_WIDTH - 1) {
+ lost += BUCKET_WIDTH;
+ }
+ }
+ assertEquals(i + 1 - lost, f.get());
+ }
+ }
+}
diff --git a/sdks/java/core/src/test/java/com/google/cloud/dataflow/sdk/util/MovingFunctionTest.java b/sdks/java/core/src/test/java/com/google/cloud/dataflow/sdk/util/MovingFunctionTest.java
new file mode 100644
index 000000000000..2ca8b7cc7c6d
--- /dev/null
+++ b/sdks/java/core/src/test/java/com/google/cloud/dataflow/sdk/util/MovingFunctionTest.java
@@ -0,0 +1,115 @@
+/*
+ * 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 com.google.cloud.dataflow.sdk.util;
+
+import com.google.cloud.dataflow.sdk.transforms.Combine;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/**
+ * Tests {@link MovingFunction}.
+ */
+@RunWith(JUnit4.class)
+public class MovingFunctionTest {
+
+ private static final long SAMPLE_PERIOD = 100;
+ private static final long SAMPLE_UPDATE = 10;
+ private static final int SIGNIFICANT_BUCKETS = 2;
+ private static final int SIGNIFICANT_SAMPLES = 10;
+
+ private static final Combine.BinaryCombineLongFn SUM =
+ new Combine.BinaryCombineLongFn() {
+ @Override
+ public long apply(long left, long right) {
+ return left + right;
+ }
+
+ @Override
+ public long identity() {
+ return 0;
+ }
+ };
+
+ private MovingFunction newFunc() {
+ return new
+ MovingFunction(SAMPLE_PERIOD, SAMPLE_UPDATE, SIGNIFICANT_BUCKETS,
+ SIGNIFICANT_SAMPLES, SUM);
+
+ }
+
+ @Test
+ public void significantSamples() {
+ MovingFunction f = newFunc();
+ assertFalse(f.isSignificant());
+ for (int i = 0; i < SIGNIFICANT_SAMPLES - 1; i++) {
+ f.add(0, 0);
+ assertFalse(f.isSignificant());
+ }
+ f.add(0, 0);
+ assertTrue(f.isSignificant());
+ }
+
+ @Test
+ public void significantBuckets() {
+ MovingFunction f = newFunc();
+ assertFalse(f.isSignificant());
+ f.add(0, 0);
+ assertFalse(f.isSignificant());
+ f.add(SAMPLE_UPDATE, 0);
+ assertTrue(f.isSignificant());
+ }
+
+ @Test
+ public void sum() {
+ MovingFunction f = newFunc();
+ for (int i = 0; i < SAMPLE_PERIOD; i++) {
+ f.add(i, i);
+ assertEquals(((i + 1) * i) / 2, f.get(i));
+ }
+ }
+
+ @Test
+ public void movingSum() {
+ MovingFunction f = newFunc();
+ int lost = 0;
+ for (int i = 0; i < SAMPLE_PERIOD * 2; i++) {
+ f.add(i , 1);
+ if (i >= SAMPLE_PERIOD) {
+ if (i % SAMPLE_UPDATE == 0) {
+ lost += SAMPLE_UPDATE;
+ }
+ }
+ assertEquals(i + 1 - lost, f.get(i));
+ }
+ }
+
+ @Test
+ public void jumpingSum() {
+ MovingFunction f = newFunc();
+ f.add(0, 1);
+ f.add(SAMPLE_PERIOD - 1, 1);
+ assertEquals(2, f.get(SAMPLE_PERIOD - 1));
+ assertEquals(1, f.get(SAMPLE_PERIOD + 3 * SAMPLE_UPDATE));
+ assertEquals(0, f.get(SAMPLE_PERIOD * 2));
+ }
+}