| /* |
| * 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.transforms; |
| |
| import static org.apache.beam.sdk.transforms.Contextful.Fn.Context.wrapProcessContext; |
| import static org.apache.beam.sdk.transforms.DoFn.ProcessContinuation.resume; |
| import static org.apache.beam.sdk.transforms.DoFn.ProcessContinuation.stop; |
| 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.auto.value.AutoValue; |
| import java.io.IOException; |
| import java.io.InputStream; |
| import java.io.OutputStream; |
| import java.io.Serializable; |
| import java.util.Arrays; |
| import java.util.Collections; |
| import java.util.HashSet; |
| import java.util.List; |
| import java.util.Map; |
| import java.util.Objects; |
| import java.util.Set; |
| import javax.annotation.Nullable; |
| import org.apache.beam.sdk.annotations.Experimental; |
| import org.apache.beam.sdk.coders.AtomicCoder; |
| import org.apache.beam.sdk.coders.CannotProvideCoderException; |
| import org.apache.beam.sdk.coders.Coder; |
| import org.apache.beam.sdk.coders.DurationCoder; |
| import org.apache.beam.sdk.coders.InstantCoder; |
| import org.apache.beam.sdk.coders.KvCoder; |
| import org.apache.beam.sdk.coders.ListCoder; |
| import org.apache.beam.sdk.coders.MapCoder; |
| import org.apache.beam.sdk.coders.NullableCoder; |
| import org.apache.beam.sdk.coders.SnappyCoder; |
| import org.apache.beam.sdk.coders.StructuredCoder; |
| import org.apache.beam.sdk.coders.VarIntCoder; |
| import org.apache.beam.sdk.io.range.OffsetRange; |
| import org.apache.beam.sdk.transforms.Contextful.Fn; |
| import org.apache.beam.sdk.transforms.DoFn.BoundedPerElement; |
| import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement; |
| import org.apache.beam.sdk.transforms.Watch.Growth.PollResult; |
| import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker; |
| import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker; |
| import org.apache.beam.sdk.transforms.windowing.BoundedWindow; |
| import org.apache.beam.sdk.util.VarInt; |
| import org.apache.beam.sdk.values.KV; |
| import org.apache.beam.sdk.values.PCollection; |
| import org.apache.beam.sdk.values.TimestampedValue; |
| import org.apache.beam.sdk.values.TypeDescriptor; |
| import org.apache.beam.sdk.values.TypeDescriptors; |
| import org.apache.beam.sdk.values.TypeDescriptors.TypeVariableExtractor; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.MoreObjects; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Maps; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Ordering; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.hash.Funnel; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.hash.Funnels; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.hash.HashCode; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.hash.Hashing; |
| import org.joda.time.Duration; |
| import org.joda.time.Instant; |
| import org.joda.time.ReadableDuration; |
| import org.slf4j.Logger; |
| import org.slf4j.LoggerFactory; |
| |
| /** |
| * Given a "poll function" that produces a potentially growing set of outputs for an input, this |
| * transform simultaneously continuously watches the growth of output sets of all inputs, until a |
| * per-input termination condition is reached. |
| * |
| * <p>The output is returned as an unbounded {@link PCollection} of {@code KV<InputT, OutputT>}, |
| * where each {@code OutputT} is associated with the {@code InputT} that produced it, and is |
| * assigned with the timestamp that the poll function returned when this output was detected for the |
| * first time. |
| * |
| * <p>Hypothetical usage example for watching new files in a collection of directories, where for |
| * each directory we assume that new files will not appear if the directory contains a file named |
| * ".complete": |
| * |
| * <pre>{@code |
| * PCollection<String> directories = ...; // E.g. Create.of(single directory) |
| * PCollection<KV<String, String>> matches = filepatterns.apply(Watch.<String, String>growthOf( |
| * new PollFn<String, String>() { |
| * public PollResult<String> apply(TimestampedValue<String> input) { |
| * String directory = input.getValue(); |
| * List<TimestampedValue<String>> outputs = new ArrayList<>(); |
| * ... List the directory and get creation times of all files ... |
| * boolean isComplete = ... does a file ".complete" exist in the directory ... |
| * return isComplete ? PollResult.complete(outputs) : PollResult.incomplete(outputs); |
| * } |
| * }) |
| * // Poll each directory every 5 seconds |
| * .withPollInterval(Duration.standardSeconds(5)) |
| * // Stop watching each directory 12 hours after it's seen even if it's incomplete |
| * .withTerminationPerInput(afterTotalOf(Duration.standardHours(12))); |
| * }</pre> |
| * |
| * <p>By default, the watermark for a particular input is computed from a poll result as "earliest |
| * timestamp of new elements in this poll result". It can also be set explicitly via {@link |
| * Growth.PollResult#withWatermark} if the {@link Growth.PollFn} can provide a more optimistic |
| * estimate. |
| * |
| * <p>Note: This transform works only in runners supporting Splittable DoFn: see <a |
| * href="https://beam.apache.org/documentation/runners/capability-matrix/">capability matrix</a>. |
| */ |
| @Experimental(Experimental.Kind.SPLITTABLE_DO_FN) |
| public class Watch { |
| private static final Logger LOG = LoggerFactory.getLogger(Watch.class); |
| |
| /** Watches the growth of the given poll function. See class documentation for more details. */ |
| public static <InputT, OutputT> Growth<InputT, OutputT, OutputT> growthOf( |
| Growth.PollFn<InputT, OutputT> pollFn, Requirements requirements) { |
| return new AutoValue_Watch_Growth.Builder<InputT, OutputT, OutputT>() |
| .setTerminationPerInput(Growth.never()) |
| .setPollFn(Contextful.of(pollFn, requirements)) |
| // use null as a signal that this is the identity function and output coder can be |
| // reused as key coder |
| .setOutputKeyFn(null) |
| .build(); |
| } |
| |
| /** Watches the growth of the given poll function. See class documentation for more details. */ |
| public static <InputT, OutputT> Growth<InputT, OutputT, OutputT> growthOf( |
| Growth.PollFn<InputT, OutputT> pollFn) { |
| return growthOf(pollFn, Requirements.empty()); |
| } |
| |
| /** |
| * Watches the growth of the given poll function, using the given "key function" to deduplicate |
| * outputs. For example, if OutputT is a filename + file size, this can be a function that returns |
| * just the filename, so that if the same file is observed multiple times with different sizes, |
| * only the first observation is emitted. |
| * |
| * <p>By default, this is the identity function, i.e. the output is used as its own key. |
| */ |
| public static <InputT, OutputT, KeyT> Growth<InputT, OutputT, KeyT> growthOf( |
| Contextful<Growth.PollFn<InputT, OutputT>> pollFn, |
| SerializableFunction<OutputT, KeyT> outputKeyFn) { |
| checkArgument(pollFn != null, "pollFn can not be null"); |
| checkArgument(outputKeyFn != null, "outputKeyFn can not be null"); |
| return new AutoValue_Watch_Growth.Builder<InputT, OutputT, KeyT>() |
| .setTerminationPerInput(Watch.Growth.never()) |
| .setPollFn(pollFn) |
| .setOutputKeyFn(outputKeyFn) |
| .build(); |
| } |
| |
| /** Implementation of {@link #growthOf}. */ |
| @AutoValue |
| public abstract static class Growth<InputT, OutputT, KeyT> |
| extends PTransform<PCollection<InputT>, PCollection<KV<InputT, OutputT>>> { |
| /** The result of a single invocation of a {@link PollFn}. */ |
| public static final class PollResult<OutputT> { |
| private final List<TimestampedValue<OutputT>> outputs; |
| // null means unspecified (infer automatically). |
| @Nullable private final Instant watermark; |
| |
| private PollResult(List<TimestampedValue<OutputT>> outputs, @Nullable Instant watermark) { |
| this.outputs = outputs; |
| this.watermark = watermark; |
| } |
| |
| List<TimestampedValue<OutputT>> getOutputs() { |
| return outputs; |
| } |
| |
| @Nullable |
| Instant getWatermark() { |
| return watermark; |
| } |
| |
| /** |
| * Returns a new {@link PollResult} like this one with the provided watermark. The watermark |
| * represents an approximate lower bound on timestamps of future new outputs from the {@link |
| * PollFn}. |
| */ |
| public PollResult<OutputT> withWatermark(Instant watermark) { |
| checkNotNull(watermark, "watermark"); |
| return new PollResult<>(outputs, watermark); |
| } |
| |
| /** Returns a new {@link PollResult} like this one with the provided outputs. */ |
| public PollResult<OutputT> withOutputs(List<TimestampedValue<OutputT>> outputs) { |
| checkNotNull(outputs); |
| return new PollResult<>(outputs, watermark); |
| } |
| |
| @Override |
| public String toString() { |
| return MoreObjects.toStringHelper(this) |
| .add("watermark", watermark) |
| .add("outputs", outputs) |
| .toString(); |
| } |
| |
| /** |
| * Constructs a {@link PollResult} with the given outputs and declares that there will be no |
| * new outputs for the current input. The {@link PollFn} will not be called again for this |
| * input. |
| */ |
| public static <OutputT> PollResult<OutputT> complete( |
| List<TimestampedValue<OutputT>> outputs) { |
| return new PollResult<>(outputs, BoundedWindow.TIMESTAMP_MAX_VALUE); |
| } |
| |
| /** Like {@link #complete(List)}, but assigns the same timestamp to all new outputs. */ |
| public static <OutputT> PollResult<OutputT> complete( |
| Instant timestamp, List<OutputT> outputs) { |
| return new PollResult<>( |
| addTimestamp(timestamp, outputs), BoundedWindow.TIMESTAMP_MAX_VALUE); |
| } |
| |
| /** |
| * Constructs a {@link PollResult} with the given outputs and declares that new outputs might |
| * appear for the current input. By default, {@link Watch} will estimate the watermark for |
| * future new outputs as equal to the earliest of the new outputs from this {@link |
| * PollResult}. To specify a more exact watermark, use {@link #withWatermark(Instant)}. |
| */ |
| public static <OutputT> PollResult<OutputT> incomplete( |
| List<TimestampedValue<OutputT>> outputs) { |
| return new PollResult<>(outputs, null); |
| } |
| |
| /** Like {@link #incomplete(List)}, but assigns the same timestamp to all new outputs. */ |
| public static <OutputT> PollResult<OutputT> incomplete( |
| Instant timestamp, List<OutputT> outputs) { |
| return new PollResult<>(addTimestamp(timestamp, outputs), null); |
| } |
| |
| private static <OutputT> List<TimestampedValue<OutputT>> addTimestamp( |
| Instant timestamp, List<OutputT> outputs) { |
| List<TimestampedValue<OutputT>> res = Lists.newArrayListWithExpectedSize(outputs.size()); |
| for (OutputT output : outputs) { |
| res.add(TimestampedValue.of(output, timestamp)); |
| } |
| return res; |
| } |
| |
| @Override |
| public boolean equals(Object o) { |
| if (this == o) { |
| return true; |
| } |
| if (o == null || getClass() != o.getClass()) { |
| return false; |
| } |
| PollResult<?> that = (PollResult<?>) o; |
| return Objects.equals(outputs, that.outputs) && Objects.equals(watermark, that.watermark); |
| } |
| |
| @Override |
| public int hashCode() { |
| return Objects.hash(outputs, watermark); |
| } |
| } |
| |
| /** |
| * A function that computes the current set of outputs for the given input, in the form of a |
| * {@link PollResult}. |
| */ |
| public abstract static class PollFn<InputT, OutputT> |
| implements Fn<InputT, PollResult<OutputT>> {} |
| |
| /** |
| * A strategy for determining whether it is time to stop polling the current input regardless of |
| * whether its output is complete or not. |
| * |
| * <p>Some built-in termination conditions are {@link #never}, {@link #afterTotalOf} and {@link |
| * #afterTimeSinceNewOutput}. Conditions can be combined using {@link #eitherOf} and {@link |
| * #allOf}. Users can also develop custom termination conditions, for example, one might imagine |
| * a condition that terminates after a given time after the first output appears for the input |
| * (unlike {@link #afterTotalOf} which operates relative to when the input itself arrives). |
| * |
| * <p>A {@link TerminationCondition} is provided to {@link |
| * Growth#withTerminationPerInput(TerminationCondition)} and is used to maintain an independent |
| * state of the termination condition for every input, represented as {@code StateT} which must |
| * be immutable, non-null, and encodable via {@link #getStateCoder()}. |
| * |
| * <p>All functions take the wall-clock timestamp as {@link Instant} for convenience of |
| * unit-testing custom termination conditions. |
| */ |
| public interface TerminationCondition<InputT, StateT> extends Serializable { |
| /** Used to encode the state of this {@link TerminationCondition}. */ |
| Coder<StateT> getStateCoder(); |
| |
| /** |
| * Called by the {@link Watch} transform to create a new independent termination state for a |
| * newly arrived {@code InputT}. |
| */ |
| StateT forNewInput(Instant now, @Nullable InputT input); |
| |
| /** |
| * Called by the {@link Watch} transform to compute a new termination state, in case after |
| * calling the {@link PollFn} for the current input, the {@link PollResult} included a |
| * previously unseen {@code OutputT}. |
| */ |
| StateT onSeenNewOutput(Instant now, StateT state); |
| |
| /** |
| * Called by the {@link Watch} transform to determine whether the given termination state |
| * signals that {@link Watch} should stop calling {@link PollFn} for the current input, |
| * regardless of whether the last {@link PollResult} was complete or incomplete. |
| */ |
| boolean canStopPolling(Instant now, StateT state); |
| |
| /** Creates a human-readable representation of the given state of this condition. */ |
| String toString(StateT state); |
| } |
| |
| /** |
| * Returns a {@link TerminationCondition} that never holds (i.e., poll each input until its |
| * output is complete). |
| */ |
| public static <InputT> Never<InputT> never() { |
| return new Never<>(); |
| } |
| |
| /** |
| * Wraps a given input-independent {@link TerminationCondition} as an equivalent condition with |
| * a given input type, passing {@code null} to the original condition as input. |
| */ |
| public static <InputT, StateT> TerminationCondition<InputT, StateT> ignoreInput( |
| TerminationCondition<?, StateT> condition) { |
| return new IgnoreInput<>(condition); |
| } |
| |
| /** |
| * Returns a {@link TerminationCondition} that holds after the given time has elapsed after the |
| * current input was seen. |
| */ |
| public static <InputT> AfterTotalOf<InputT> afterTotalOf(ReadableDuration timeSinceInput) { |
| return afterTotalOf(SerializableFunctions.<InputT, ReadableDuration>constant(timeSinceInput)); |
| } |
| |
| /** Like {@link #afterTotalOf(ReadableDuration)}, but the duration is input-dependent. */ |
| public static <InputT> AfterTotalOf<InputT> afterTotalOf( |
| SerializableFunction<InputT, ReadableDuration> timeSinceInput) { |
| return new AfterTotalOf<>(timeSinceInput); |
| } |
| |
| /** |
| * Returns a {@link TerminationCondition} that holds after the given time has elapsed after the |
| * last time the {@link PollResult} for the current input contained a previously unseen output. |
| */ |
| public static <InputT> AfterTimeSinceNewOutput<InputT> afterTimeSinceNewOutput( |
| ReadableDuration timeSinceNewOutput) { |
| return afterTimeSinceNewOutput( |
| SerializableFunctions.<InputT, ReadableDuration>constant(timeSinceNewOutput)); |
| } |
| |
| /** |
| * Like {@link #afterTimeSinceNewOutput(ReadableDuration)}, but the duration is input-dependent. |
| */ |
| public static <InputT> AfterTimeSinceNewOutput<InputT> afterTimeSinceNewOutput( |
| SerializableFunction<InputT, ReadableDuration> timeSinceNewOutput) { |
| return new AfterTimeSinceNewOutput<>(timeSinceNewOutput); |
| } |
| |
| /** |
| * Returns a {@link TerminationCondition} that holds when at least one of the given two |
| * conditions holds. |
| */ |
| public static <InputT, FirstStateT, SecondStateT> |
| BinaryCombined<InputT, FirstStateT, SecondStateT> eitherOf( |
| TerminationCondition<InputT, FirstStateT> first, |
| TerminationCondition<InputT, SecondStateT> second) { |
| return new BinaryCombined<>(BinaryCombined.Operation.OR, first, second); |
| } |
| |
| /** |
| * Returns a {@link TerminationCondition} that holds when both of the given two conditions hold. |
| */ |
| public static <InputT, FirstStateT, SecondStateT> |
| BinaryCombined<InputT, FirstStateT, SecondStateT> allOf( |
| TerminationCondition<InputT, FirstStateT> first, |
| TerminationCondition<InputT, SecondStateT> second) { |
| return new BinaryCombined<>(BinaryCombined.Operation.AND, first, second); |
| } |
| |
| // Uses Integer rather than Void for state, because termination state must be non-null. |
| static class Never<InputT> implements TerminationCondition<InputT, Integer> { |
| @Override |
| public Coder<Integer> getStateCoder() { |
| return VarIntCoder.of(); |
| } |
| |
| @Override |
| public Integer forNewInput(Instant now, InputT input) { |
| return 0; |
| } |
| |
| @Override |
| public Integer onSeenNewOutput(Instant now, Integer state) { |
| return state; |
| } |
| |
| @Override |
| public boolean canStopPolling(Instant now, Integer state) { |
| return false; |
| } |
| |
| @Override |
| public String toString(Integer state) { |
| return "Never"; |
| } |
| } |
| |
| static class IgnoreInput<InputT, StateT> implements TerminationCondition<InputT, StateT> { |
| private final TerminationCondition<?, StateT> wrapped; |
| |
| IgnoreInput(TerminationCondition<?, StateT> wrapped) { |
| this.wrapped = wrapped; |
| } |
| |
| @Override |
| public Coder<StateT> getStateCoder() { |
| return wrapped.getStateCoder(); |
| } |
| |
| @Override |
| public StateT forNewInput(Instant now, InputT input) { |
| return wrapped.forNewInput(now, null); |
| } |
| |
| @Override |
| public StateT onSeenNewOutput(Instant now, StateT state) { |
| return wrapped.onSeenNewOutput(now, state); |
| } |
| |
| @Override |
| public boolean canStopPolling(Instant now, StateT state) { |
| return wrapped.canStopPolling(now, state); |
| } |
| |
| @Override |
| public String toString(StateT state) { |
| return wrapped.toString(state); |
| } |
| } |
| |
| static class AfterTotalOf<InputT> |
| implements TerminationCondition< |
| InputT, KV<Instant /* timeStarted */, ReadableDuration /* maxTimeSinceInput */>> { |
| private final SerializableFunction<InputT, ReadableDuration> maxTimeSinceInput; |
| |
| private AfterTotalOf(SerializableFunction<InputT, ReadableDuration> maxTimeSinceInput) { |
| this.maxTimeSinceInput = maxTimeSinceInput; |
| } |
| |
| @Override |
| public Coder<KV<Instant, ReadableDuration>> getStateCoder() { |
| return KvCoder.of(InstantCoder.of(), DurationCoder.of()); |
| } |
| |
| @Override |
| public KV<Instant, ReadableDuration> forNewInput(Instant now, InputT input) { |
| return KV.of(now, maxTimeSinceInput.apply(input)); |
| } |
| |
| @Override |
| public KV<Instant, ReadableDuration> onSeenNewOutput( |
| Instant now, KV<Instant, ReadableDuration> state) { |
| return state; |
| } |
| |
| @Override |
| public boolean canStopPolling(Instant now, KV<Instant, ReadableDuration> state) { |
| return new Duration(state.getKey(), now).isLongerThan(state.getValue()); |
| } |
| |
| @Override |
| public String toString(KV<Instant, ReadableDuration> state) { |
| return "AfterTotalOf{" |
| + "timeStarted=" |
| + state.getKey() |
| + ", maxTimeSinceInput=" |
| + state.getValue() |
| + '}'; |
| } |
| } |
| |
| static class AfterTimeSinceNewOutput<InputT> |
| implements TerminationCondition< |
| InputT, |
| KV<Instant /* timeOfLastNewOutput */, ReadableDuration /* maxTimeSinceNewOutput */>> { |
| private final SerializableFunction<InputT, ReadableDuration> maxTimeSinceNewOutput; |
| |
| private AfterTimeSinceNewOutput( |
| SerializableFunction<InputT, ReadableDuration> maxTimeSinceNewOutput) { |
| this.maxTimeSinceNewOutput = maxTimeSinceNewOutput; |
| } |
| |
| @Override |
| public Coder<KV<Instant, ReadableDuration>> getStateCoder() { |
| return KvCoder.of(NullableCoder.of(InstantCoder.of()), DurationCoder.of()); |
| } |
| |
| @Override |
| public KV<Instant, ReadableDuration> forNewInput(Instant now, InputT input) { |
| return KV.of(null, maxTimeSinceNewOutput.apply(input)); |
| } |
| |
| @Override |
| public KV<Instant, ReadableDuration> onSeenNewOutput( |
| Instant now, KV<Instant, ReadableDuration> state) { |
| return KV.of(now, state.getValue()); |
| } |
| |
| @Override |
| public boolean canStopPolling(Instant now, KV<Instant, ReadableDuration> state) { |
| Instant timeOfLastNewOutput = state.getKey(); |
| ReadableDuration maxTimeSinceNewOutput = state.getValue(); |
| return timeOfLastNewOutput != null |
| && new Duration(timeOfLastNewOutput, now).isLongerThan(maxTimeSinceNewOutput); |
| } |
| |
| @Override |
| public String toString(KV<Instant, ReadableDuration> state) { |
| return "AfterTimeSinceNewOutput{" |
| + "timeOfLastNewOutput=" |
| + state.getKey() |
| + ", maxTimeSinceNewOutput=" |
| + state.getValue() |
| + '}'; |
| } |
| } |
| |
| static class BinaryCombined<InputT, FirstStateT, SecondStateT> |
| implements TerminationCondition<InputT, KV<FirstStateT, SecondStateT>> { |
| private enum Operation { |
| OR, |
| AND |
| } |
| |
| private final Operation operation; |
| private final TerminationCondition<InputT, FirstStateT> first; |
| private final TerminationCondition<InputT, SecondStateT> second; |
| |
| public BinaryCombined( |
| Operation operation, |
| TerminationCondition<InputT, FirstStateT> first, |
| TerminationCondition<InputT, SecondStateT> second) { |
| this.operation = operation; |
| this.first = first; |
| this.second = second; |
| } |
| |
| @Override |
| public Coder<KV<FirstStateT, SecondStateT>> getStateCoder() { |
| return KvCoder.of(first.getStateCoder(), second.getStateCoder()); |
| } |
| |
| @Override |
| public KV<FirstStateT, SecondStateT> forNewInput(Instant now, InputT input) { |
| return KV.of(first.forNewInput(now, input), second.forNewInput(now, input)); |
| } |
| |
| @Override |
| public KV<FirstStateT, SecondStateT> onSeenNewOutput( |
| Instant now, KV<FirstStateT, SecondStateT> state) { |
| return KV.of( |
| first.onSeenNewOutput(now, state.getKey()), |
| second.onSeenNewOutput(now, state.getValue())); |
| } |
| |
| @Override |
| public boolean canStopPolling(Instant now, KV<FirstStateT, SecondStateT> state) { |
| switch (operation) { |
| case OR: |
| return first.canStopPolling(now, state.getKey()) |
| || second.canStopPolling(now, state.getValue()); |
| case AND: |
| return first.canStopPolling(now, state.getKey()) |
| && second.canStopPolling(now, state.getValue()); |
| default: |
| throw new UnsupportedOperationException("Unexpected operation " + operation); |
| } |
| } |
| |
| @Override |
| public String toString(KV<FirstStateT, SecondStateT> state) { |
| return operation |
| + "{first=" |
| + first.toString(state.getKey()) |
| + ", second=" |
| + second.toString(state.getValue()) |
| + '}'; |
| } |
| } |
| |
| abstract Contextful<PollFn<InputT, OutputT>> getPollFn(); |
| |
| @Nullable |
| abstract SerializableFunction<OutputT, KeyT> getOutputKeyFn(); |
| |
| @Nullable |
| abstract Coder<KeyT> getOutputKeyCoder(); |
| |
| @Nullable |
| abstract Duration getPollInterval(); |
| |
| @Nullable |
| abstract TerminationCondition<InputT, ?> getTerminationPerInput(); |
| |
| @Nullable |
| abstract Coder<OutputT> getOutputCoder(); |
| |
| abstract Builder<InputT, OutputT, KeyT> toBuilder(); |
| |
| @AutoValue.Builder |
| abstract static class Builder<InputT, OutputT, KeyT> { |
| abstract Builder<InputT, OutputT, KeyT> setPollFn(Contextful<PollFn<InputT, OutputT>> pollFn); |
| |
| abstract Builder<InputT, OutputT, KeyT> setOutputKeyFn( |
| @Nullable SerializableFunction<OutputT, KeyT> outputKeyFn); |
| |
| abstract Builder<InputT, OutputT, KeyT> setOutputKeyCoder(Coder<KeyT> outputKeyCoder); |
| |
| abstract Builder<InputT, OutputT, KeyT> setTerminationPerInput( |
| TerminationCondition<InputT, ?> terminationPerInput); |
| |
| abstract Builder<InputT, OutputT, KeyT> setPollInterval(Duration pollInterval); |
| |
| abstract Builder<InputT, OutputT, KeyT> setOutputCoder(Coder<OutputT> outputCoder); |
| |
| abstract Growth<InputT, OutputT, KeyT> build(); |
| } |
| |
| /** Specifies the coder for the output key. */ |
| public Growth<InputT, OutputT, KeyT> withOutputKeyCoder(Coder<KeyT> outputKeyCoder) { |
| return toBuilder().setOutputKeyCoder(outputKeyCoder).build(); |
| } |
| |
| /** Specifies a {@link TerminationCondition} that will be independently used for every input. */ |
| public Growth<InputT, OutputT, KeyT> withTerminationPerInput( |
| TerminationCondition<InputT, ?> terminationPerInput) { |
| return toBuilder().setTerminationPerInput(terminationPerInput).build(); |
| } |
| |
| /** |
| * Specifies how long to wait after a call to {@link PollFn} before calling it again (if at all |
| * - according to {@link PollResult} and the {@link TerminationCondition}). |
| */ |
| public Growth<InputT, OutputT, KeyT> withPollInterval(Duration pollInterval) { |
| return toBuilder().setPollInterval(pollInterval).build(); |
| } |
| |
| /** |
| * Specifies a {@link Coder} to use for the outputs. If unspecified, it will be inferred from |
| * the output type of {@link PollFn} whenever possible. |
| * |
| * <p>The coder must be deterministic, because the transform will compare encoded outputs for |
| * deduplication between polling rounds. |
| */ |
| public Growth<InputT, OutputT, KeyT> withOutputCoder(Coder<OutputT> outputCoder) { |
| return toBuilder().setOutputCoder(outputCoder).build(); |
| } |
| |
| @Override |
| public PCollection<KV<InputT, OutputT>> expand(PCollection<InputT> input) { |
| checkNotNull(getPollInterval(), "pollInterval"); |
| checkNotNull(getTerminationPerInput(), "terminationPerInput"); |
| |
| Coder<OutputT> outputCoder = getOutputCoder(); |
| if (outputCoder == null) { |
| // If a coder was not specified explicitly, infer it from the OutputT type parameter |
| // of the PollFn. |
| TypeDescriptor<OutputT> outputT = |
| TypeDescriptors.extractFromTypeParameters( |
| getPollFn().getClosure(), |
| PollFn.class, |
| new TypeVariableExtractor<PollFn<InputT, OutputT>, OutputT>() {}); |
| try { |
| outputCoder = input.getPipeline().getCoderRegistry().getCoder(outputT); |
| } catch (CannotProvideCoderException e) { |
| throw new RuntimeException( |
| "Unable to infer coder for OutputT (" |
| + outputT |
| + "). Specify it explicitly using withOutputCoder()."); |
| } |
| } |
| |
| Coder<KeyT> outputKeyCoder = getOutputKeyCoder(); |
| SerializableFunction<OutputT, KeyT> outputKeyFn = getOutputKeyFn(); |
| if (getOutputKeyFn() == null) { |
| // This by construction can happen only if OutputT == KeyT |
| outputKeyCoder = (Coder) outputCoder; |
| outputKeyFn = (SerializableFunction) SerializableFunctions.identity(); |
| } else { |
| if (outputKeyCoder == null) { |
| // If a coder was not specified explicitly, infer it from the OutputT type parameter |
| // of the output key fn. |
| TypeDescriptor<KeyT> keyT = TypeDescriptors.outputOf(getOutputKeyFn()); |
| try { |
| outputKeyCoder = input.getPipeline().getCoderRegistry().getCoder(keyT); |
| } catch (CannotProvideCoderException e) { |
| throw new RuntimeException( |
| "Unable to infer coder for KeyT (" |
| + keyT |
| + "). Specify it explicitly using withOutputKeyCoder()."); |
| } |
| } |
| try { |
| outputKeyCoder.verifyDeterministic(); |
| } catch (Coder.NonDeterministicException e) { |
| throw new IllegalArgumentException( |
| "Key coder " + outputKeyCoder + " must be deterministic"); |
| } |
| } |
| |
| PCollection<KV<InputT, List<TimestampedValue<OutputT>>>> polledPc = |
| input |
| .apply( |
| ParDo.of(new WatchGrowthFn<>(this, outputCoder, outputKeyFn, outputKeyCoder)) |
| .withSideInputs(getPollFn().getRequirements().getSideInputs())) |
| .setCoder( |
| KvCoder.of( |
| input.getCoder(), |
| ListCoder.of(TimestampedValue.TimestampedValueCoder.of(outputCoder)))); |
| return polledPc |
| .apply(ParDo.of(new PollResultSplitFn<>())) |
| .setCoder(KvCoder.of(input.getCoder(), outputCoder)); |
| } |
| } |
| |
| /** A splittable {@link DoFn} that emits {@link PollResult}s outputs. */ |
| @BoundedPerElement |
| private static class PollResultSplitFn<InputT, OutputT> |
| extends DoFn<KV<InputT, List<TimestampedValue<OutputT>>>, KV<InputT, OutputT>> { |
| |
| @ProcessElement |
| public void processElement(ProcessContext c, RestrictionTracker<OffsetRange, Long> tracker) { |
| long position = tracker.currentRestriction().getFrom(); |
| while (tracker.tryClaim(position)) { |
| TimestampedValue<OutputT> value = c.element().getValue().get((int) position); |
| c.outputWithTimestamp(KV.of(c.element().getKey(), value.getValue()), value.getTimestamp()); |
| c.updateWatermark(value.getTimestamp()); |
| position += 1L; |
| } |
| } |
| |
| @GetInitialRestriction |
| public OffsetRange getInitialRestriction(KV<InputT, List<TimestampedValue<OutputT>>> element) { |
| return new OffsetRange(0, element.getValue().size()); |
| } |
| |
| @NewTracker |
| public OffsetRangeTracker newTracker(OffsetRange restriction) { |
| return restriction.newTracker(); |
| } |
| |
| @GetRestrictionCoder |
| public Coder<OffsetRange> getRestrictionCoder() { |
| return OffsetRange.Coder.of(); |
| } |
| } |
| |
| @UnboundedPerElement |
| private static class WatchGrowthFn<InputT, OutputT, KeyT, TerminationStateT> |
| extends DoFn<InputT, KV<InputT, List<TimestampedValue<OutputT>>>> { |
| private final Watch.Growth<InputT, OutputT, KeyT> spec; |
| private final Coder<OutputT> outputCoder; |
| private final SerializableFunction<OutputT, KeyT> outputKeyFn; |
| private final Coder<KeyT> outputKeyCoder; |
| private final Funnel<OutputT> coderFunnel; |
| |
| private WatchGrowthFn( |
| Growth<InputT, OutputT, KeyT> spec, |
| Coder<OutputT> outputCoder, |
| SerializableFunction<OutputT, KeyT> outputKeyFn, |
| Coder<KeyT> outputKeyCoder) { |
| this.spec = spec; |
| this.outputCoder = outputCoder; |
| this.outputKeyFn = outputKeyFn; |
| this.outputKeyCoder = outputKeyCoder; |
| this.coderFunnel = |
| (from, into) -> { |
| try { |
| // Rather than hashing the output itself, hash the output key. |
| KeyT outputKey = outputKeyFn.apply(from); |
| outputKeyCoder.encode(outputKey, Funnels.asOutputStream(into)); |
| } catch (IOException e) { |
| throw new RuntimeException(e); |
| } |
| }; |
| } |
| |
| @ProcessElement |
| public ProcessContinuation process( |
| ProcessContext c, |
| RestrictionTracker<GrowthState, KV<Growth.PollResult<OutputT>, TerminationStateT>> tracker) |
| throws Exception { |
| |
| GrowthState currentRestriction = tracker.currentRestriction(); |
| if (currentRestriction instanceof NonPollingGrowthState) { |
| Growth.PollResult<OutputT> priorPoll = |
| ((NonPollingGrowthState<OutputT>) currentRestriction).getPending(); |
| if (tracker.tryClaim(KV.of(priorPoll, null))) { |
| if (!priorPoll.getOutputs().isEmpty()) { |
| LOG.info( |
| "{} - re-emitting output of prior poll containing {} results.", |
| c.element(), |
| priorPoll.getOutputs().size()); |
| c.output(KV.of(c.element(), priorPoll.getOutputs())); |
| } |
| if (priorPoll.getWatermark() != null) { |
| c.updateWatermark(priorPoll.getWatermark()); |
| } |
| } |
| return stop(); |
| } |
| |
| // Poll for additional elements. |
| Instant now = Instant.now(); |
| Growth.PollResult<OutputT> res = |
| spec.getPollFn().getClosure().apply(c.element(), wrapProcessContext(c)); |
| |
| PollingGrowthState<TerminationStateT> pollingRestriction = |
| (PollingGrowthState<TerminationStateT>) currentRestriction; |
| // Produce a poll result that only contains never seen before results. |
| Growth.PollResult<OutputT> newResults = |
| computeNeverSeenBeforeResults(pollingRestriction, res); |
| |
| // If we had zero new results, attempt to update the watermark if the poll result |
| // provided a watermark. Otherwise attempt to claim all pending outputs. |
| LOG.info( |
| "{} - current round of polling took {} ms and returned {} results, " |
| + "of which {} were new.", |
| c.element(), |
| new Duration(now, Instant.now()).getMillis(), |
| res.getOutputs().size(), |
| newResults.getOutputs().size()); |
| |
| TerminationStateT terminationState = pollingRestriction.getTerminationState(); |
| if (!newResults.getOutputs().isEmpty()) { |
| terminationState = |
| getTerminationCondition().onSeenNewOutput(Instant.now(), terminationState); |
| } |
| |
| if (!tracker.tryClaim(KV.of(newResults, terminationState))) { |
| LOG.info("{} - will not emit poll result tryClaim failed.", c.element()); |
| return stop(); |
| } |
| |
| if (!newResults.getOutputs().isEmpty()) { |
| c.output(KV.of(c.element(), newResults.getOutputs())); |
| } |
| |
| if (newResults.getWatermark() != null) { |
| c.updateWatermark(newResults.getWatermark()); |
| } |
| |
| Instant currentTime = Instant.now(); |
| if (getTerminationCondition().canStopPolling(currentTime, terminationState)) { |
| LOG.info( |
| "{} - told to stop polling by polling function at {} with termination state {}.", |
| c.element(), |
| currentTime, |
| getTerminationCondition().toString(terminationState)); |
| return stop(); |
| } |
| |
| if (BoundedWindow.TIMESTAMP_MAX_VALUE.equals(newResults.getWatermark())) { |
| LOG.info("{} - will stop polling, reached max timestamp.", c.element()); |
| return stop(); |
| } |
| |
| LOG.info( |
| "{} - will resume polling in {} ms.", c.element(), spec.getPollInterval().getMillis()); |
| return resume().withResumeDelay(spec.getPollInterval()); |
| } |
| |
| private HashCode hash128(OutputT value) { |
| return Hashing.murmur3_128().hashObject(value, coderFunnel); |
| } |
| |
| private Growth.PollResult computeNeverSeenBeforeResults( |
| PollingGrowthState<TerminationStateT> state, Growth.PollResult<OutputT> pollResult) { |
| // Collect results to include as newly pending. Note that the poll result may in theory |
| // contain multiple outputs mapping to the same output key - we need to ignore duplicates |
| // here already. |
| Map<HashCode, TimestampedValue<OutputT>> newPending = Maps.newHashMap(); |
| for (TimestampedValue<OutputT> output : pollResult.getOutputs()) { |
| OutputT value = output.getValue(); |
| HashCode hash = hash128(value); |
| if (state.getCompleted().containsKey(hash) || newPending.containsKey(hash)) { |
| continue; |
| } |
| // TODO (https://issues.apache.org/jira/browse/BEAM-2680): |
| // Consider adding only at most N pending elements and ignoring others, |
| // instead relying on future poll rounds to provide them, in order to avoid |
| // blowing up the state. Combined with garbage collection of PollingGrowthState.completed, |
| // this would make the transform scalable to very large poll results. |
| newPending.put(hash, output); |
| } |
| |
| return pollResult.withOutputs( |
| Ordering.natural() |
| .onResultOf((TimestampedValue<OutputT> value) -> value.getTimestamp()) |
| .sortedCopy(newPending.values())); |
| } |
| |
| private Growth.TerminationCondition<InputT, TerminationStateT> getTerminationCondition() { |
| return (Growth.TerminationCondition<InputT, TerminationStateT>) spec.getTerminationPerInput(); |
| } |
| |
| @GetInitialRestriction |
| public GrowthState getInitialRestriction(InputT element) { |
| return PollingGrowthState.of(getTerminationCondition().forNewInput(Instant.now(), element)); |
| } |
| |
| @NewTracker |
| public GrowthTracker<OutputT, TerminationStateT> newTracker(GrowthState restriction) { |
| return new GrowthTracker<>(restriction, coderFunnel); |
| } |
| |
| @GetRestrictionCoder |
| @SuppressWarnings({"unchecked", "rawtypes"}) |
| public Coder<GrowthState> getRestrictionCoder() { |
| return SnappyCoder.of( |
| GrowthStateCoder.of(outputCoder, (Coder) spec.getTerminationPerInput().getStateCoder())); |
| } |
| } |
| |
| /** A base class for all restrictions related to the {@link Growth} SplittableDoFn. */ |
| abstract static class GrowthState {} |
| |
| /** |
| * Stores the prior pending poll results related to the {@link Growth} SplittableDoFn. Used to |
| * represent the primary restriction during checkpoint which should be replayed if the primary |
| * ever needs to be re-executed. |
| */ |
| @AutoValue |
| @VisibleForTesting |
| abstract static class NonPollingGrowthState<OutputT> extends GrowthState { |
| public static <OutputT> NonPollingGrowthState<OutputT> of(Growth.PollResult<OutputT> pending) { |
| return new AutoValue_Watch_NonPollingGrowthState(pending); |
| } |
| |
| /** |
| * Contains all pending results to output. Checkpointing/splitting moves "pending" outputs to |
| * the completed set. |
| */ |
| public abstract Growth.PollResult<OutputT> getPending(); |
| } |
| |
| /** |
| * A restriction for the {@link Growth} transform which represents a polling state. The |
| * restriction represents an unbounded amount of work until one of the termination conditions of |
| * the {@link Growth} transform are met. |
| */ |
| @AutoValue |
| @VisibleForTesting |
| abstract static class PollingGrowthState<TerminationStateT> extends GrowthState { |
| public static <TerminationStateT> PollingGrowthState<TerminationStateT> of( |
| TerminationStateT terminationState) { |
| return new AutoValue_Watch_PollingGrowthState(ImmutableMap.of(), null, terminationState); |
| } |
| |
| public static <TerminationStateT> PollingGrowthState<TerminationStateT> of( |
| ImmutableMap<HashCode, Instant> completed, |
| Instant pollWatermark, |
| TerminationStateT terminationState) { |
| return new AutoValue_Watch_PollingGrowthState(completed, pollWatermark, terminationState); |
| } |
| |
| // Hashes and timestamps of outputs that have already been output and should be omitted |
| // from future polls. Timestamps are preserved to allow garbage-collecting this state |
| // in the future, e.g. dropping elements from "completed" and from |
| // computeNeverSeenBeforeResults() if their timestamp is more than X behind the watermark. |
| // As of writing, we don't do this, but preserve the information for forward compatibility |
| // in case of pipeline update. TODO: do this. |
| public abstract ImmutableMap<HashCode, Instant> getCompleted(); |
| |
| @Nullable |
| public abstract Instant getPollWatermark(); |
| |
| public abstract TerminationStateT getTerminationState(); |
| } |
| |
| @VisibleForTesting |
| static class GrowthTracker<OutputT, TerminationStateT> |
| extends RestrictionTracker<GrowthState, KV<Growth.PollResult<OutputT>, TerminationStateT>> { |
| |
| static final GrowthState EMPTY_STATE = |
| NonPollingGrowthState.of(new PollResult<>(Collections.emptyList(), null)); |
| |
| // Used to hash values. |
| private final Funnel<OutputT> coderFunnel; |
| |
| // non-null after first successful tryClaim() |
| @Nullable private Growth.PollResult<OutputT> claimedPollResult; |
| @Nullable private TerminationStateT claimedTerminationState; |
| @Nullable private ImmutableMap<HashCode, Instant> claimedHashes; |
| |
| // The restriction describing the entire work to be done by the current ProcessElement call. |
| private GrowthState state; |
| // Whether we should stop claiming poll results. |
| private boolean shouldStop; |
| |
| GrowthTracker(GrowthState state, Funnel<OutputT> coderFunnel) { |
| this.state = state; |
| this.coderFunnel = coderFunnel; |
| this.shouldStop = false; |
| } |
| |
| @Override |
| public GrowthState currentRestriction() { |
| return state; |
| } |
| |
| @Override |
| public GrowthState checkpoint() { |
| // residual should contain exactly the work *not* claimed in the current ProcessElement call - |
| // unclaimed pending outputs or future polling output |
| GrowthState residual; |
| |
| if (claimedPollResult == null) { |
| // If we have yet to claim anything then our residual becomes all the work we were meant |
| // to do and we update our current restriction to be empty. |
| residual = state; |
| state = EMPTY_STATE; |
| } else if (state instanceof NonPollingGrowthState) { |
| // Since we have claimed the prior poll, our residual is empty. |
| residual = EMPTY_STATE; |
| } else { |
| // Since we claimed a poll result, our primary becomes the poll result and |
| // our residual becomes everything we have claimed in the past + the current poll result. |
| |
| PollingGrowthState<TerminationStateT> currentState = |
| (PollingGrowthState<TerminationStateT>) state; |
| ImmutableMap.Builder<HashCode, Instant> newCompleted = ImmutableMap.builder(); |
| newCompleted.putAll(currentState.getCompleted()); |
| newCompleted.putAll(claimedHashes); |
| residual = |
| PollingGrowthState.of( |
| newCompleted.build(), |
| Ordering.natural() |
| .nullsFirst() |
| .max(currentState.getPollWatermark(), claimedPollResult.watermark), |
| claimedTerminationState); |
| state = NonPollingGrowthState.of(claimedPollResult); |
| } |
| |
| shouldStop = true; |
| return residual; |
| } |
| |
| private HashCode hash128(OutputT value) { |
| return Hashing.murmur3_128().hashObject(value, coderFunnel); |
| } |
| |
| @Override |
| public void checkDone() throws IllegalStateException { |
| checkState( |
| shouldStop, "Missing tryClaim()/checkpoint() call. Expected " + "one or the other."); |
| } |
| |
| @Override |
| public boolean tryClaim(KV<Growth.PollResult<OutputT>, TerminationStateT> pollResult) { |
| if (shouldStop) { |
| return false; |
| } |
| |
| ImmutableMap.Builder<HashCode, Instant> newClaimedHashesBuilder = ImmutableMap.builder(); |
| for (TimestampedValue<OutputT> value : pollResult.getKey().getOutputs()) { |
| HashCode hash = hash128(value.getValue()); |
| newClaimedHashesBuilder.put(hash, value.getTimestamp()); |
| } |
| ImmutableMap<HashCode, Instant> newClaimedHashes = newClaimedHashesBuilder.build(); |
| |
| if (state instanceof PollingGrowthState) { |
| // If we have previously claimed one of these hashes then return false. |
| if (!Collections.disjoint( |
| newClaimedHashes.keySet(), ((PollingGrowthState) state).getCompleted().keySet())) { |
| return false; |
| } |
| } else { |
| Set<HashCode> expectedHashesToClaim = new HashSet<>(); |
| for (TimestampedValue<OutputT> value : |
| ((NonPollingGrowthState<OutputT>) state).getPending().getOutputs()) { |
| expectedHashesToClaim.add(hash128(value.getValue())); |
| } |
| // We expect to claim the entire poll result from a NonPollingGrowthState. This is |
| // stricter then currently required and could be relaxed if this tracker supported |
| // splitting a NonPollingGrowthState into two smaller NonPollingGrowthStates. |
| if (!expectedHashesToClaim.equals(newClaimedHashes.keySet())) { |
| return false; |
| } |
| } |
| |
| // Only allow claiming a single poll result at a time. |
| shouldStop = true; |
| claimedPollResult = pollResult.getKey(); |
| claimedTerminationState = pollResult.getValue(); |
| claimedHashes = newClaimedHashes; |
| |
| return true; |
| } |
| |
| @Override |
| public String toString() { |
| return MoreObjects.toStringHelper(this) |
| .add("state", state) |
| .add("pollResult", claimedPollResult) |
| .add("terminationState", claimedTerminationState) |
| .add("shouldStop", shouldStop) |
| .toString(); |
| } |
| } |
| |
| private static class HashCode128Coder extends AtomicCoder<HashCode> { |
| private static final HashCode128Coder INSTANCE = new HashCode128Coder(); |
| |
| public static HashCode128Coder of() { |
| return INSTANCE; |
| } |
| |
| @Override |
| public void encode(HashCode value, OutputStream os) throws IOException { |
| checkArgument( |
| value.bits() == 128, "Expected a 128-bit hash code, but got %s bits", value.bits()); |
| byte[] res = new byte[16]; |
| value.writeBytesTo(res, 0, 16); |
| os.write(res); |
| } |
| |
| @Override |
| public HashCode decode(InputStream is) throws IOException { |
| byte[] res = new byte[16]; |
| int numRead = is.read(res, 0, 16); |
| checkArgument(numRead == 16, "Expected to read 16 bytes, but read %s", numRead); |
| return HashCode.fromBytes(res); |
| } |
| } |
| |
| static class GrowthStateCoder<OutputT, TerminationStateT> extends StructuredCoder<GrowthState> { |
| |
| private static final int POLLING_GROWTH_STATE = 0; |
| private static final int NON_POLLING_GROWTH_STATE = 1; |
| |
| public static <OutputT, TerminationStateT> GrowthStateCoder<OutputT, TerminationStateT> of( |
| Coder<OutputT> outputCoder, Coder<TerminationStateT> terminationStateCoder) { |
| return new GrowthStateCoder<>(outputCoder, terminationStateCoder); |
| } |
| |
| private static final MapCoder<HashCode, Instant> COMPLETED_CODER = |
| MapCoder.of(HashCode128Coder.of(), InstantCoder.of()); |
| private static final Coder<Instant> NULLABLE_INSTANT_CODER = |
| NullableCoder.of(InstantCoder.of()); |
| |
| private final Coder<OutputT> outputCoder; |
| private final Coder<List<TimestampedValue<OutputT>>> timestampedOutputCoder; |
| private final Coder<TerminationStateT> terminationStateCoder; |
| |
| private GrowthStateCoder( |
| Coder<OutputT> outputCoder, Coder<TerminationStateT> terminationStateCoder) { |
| this.outputCoder = outputCoder; |
| this.terminationStateCoder = terminationStateCoder; |
| this.timestampedOutputCoder = |
| ListCoder.of(TimestampedValue.TimestampedValueCoder.of(outputCoder)); |
| } |
| |
| @Override |
| public void encode(GrowthState value, OutputStream os) throws IOException { |
| if (value instanceof PollingGrowthState) { |
| VarInt.encode(POLLING_GROWTH_STATE, os); |
| encodePollingGrowthState((PollingGrowthState<TerminationStateT>) value, os); |
| } else if (value instanceof NonPollingGrowthState) { |
| VarInt.encode(NON_POLLING_GROWTH_STATE, os); |
| encodeNonPollingGrowthState((NonPollingGrowthState<OutputT>) value, os); |
| } else { |
| throw new IOException("Unknown growth state: " + value); |
| } |
| } |
| |
| private void encodePollingGrowthState( |
| PollingGrowthState<TerminationStateT> value, OutputStream os) throws IOException { |
| terminationStateCoder.encode(value.getTerminationState(), os); |
| NULLABLE_INSTANT_CODER.encode(value.getPollWatermark(), os); |
| COMPLETED_CODER.encode(value.getCompleted(), os); |
| } |
| |
| private void encodeNonPollingGrowthState(NonPollingGrowthState<OutputT> value, OutputStream os) |
| throws IOException { |
| NULLABLE_INSTANT_CODER.encode(value.getPending().getWatermark(), os); |
| timestampedOutputCoder.encode(value.getPending().getOutputs(), os); |
| } |
| |
| @Override |
| public GrowthState decode(InputStream is) throws IOException { |
| int type = VarInt.decodeInt(is); |
| switch (type) { |
| case NON_POLLING_GROWTH_STATE: |
| return decodeNonPollingGrowthState(is); |
| case POLLING_GROWTH_STATE: |
| return decodePollingGrowthState(is); |
| default: |
| throw new IOException("Unknown growth state type " + type); |
| } |
| } |
| |
| private GrowthState decodeNonPollingGrowthState(InputStream is) throws IOException { |
| Instant watermark = NULLABLE_INSTANT_CODER.decode(is); |
| List<TimestampedValue<OutputT>> values = timestampedOutputCoder.decode(is); |
| return NonPollingGrowthState.of(new Growth.PollResult<>(values, watermark)); |
| } |
| |
| private GrowthState decodePollingGrowthState(InputStream is) throws IOException { |
| TerminationStateT terminationState = terminationStateCoder.decode(is); |
| Instant watermark = NULLABLE_INSTANT_CODER.decode(is); |
| Map<HashCode, Instant> completed = COMPLETED_CODER.decode(is); |
| return PollingGrowthState.of(ImmutableMap.copyOf(completed), watermark, terminationState); |
| } |
| |
| @Override |
| public List<? extends Coder<?>> getCoderArguments() { |
| return Arrays.asList(outputCoder, terminationStateCoder); |
| } |
| |
| @Override |
| public void verifyDeterministic() throws NonDeterministicException { |
| outputCoder.verifyDeterministic(); |
| } |
| } |
| } |
