/*
 * 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.runners.direct;

import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.greaterThan;
import static org.hamcrest.Matchers.is;
import static org.hamcrest.Matchers.isA;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertThat;
import static org.junit.Assert.fail;

import com.fasterxml.jackson.annotation.JsonIgnore;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.Serializable;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.Optional;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.Callable;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;
import org.apache.beam.runners.direct.DirectRunner.DirectPipelineResult;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.PipelineResult;
import org.apache.beam.sdk.PipelineResult.State;
import org.apache.beam.sdk.PipelineRunner;
import org.apache.beam.sdk.coders.AtomicCoder;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.CoderException;
import org.apache.beam.sdk.coders.ListCoder;
import org.apache.beam.sdk.coders.SerializableCoder;
import org.apache.beam.sdk.coders.VarIntCoder;
import org.apache.beam.sdk.coders.VarLongCoder;
import org.apache.beam.sdk.io.BoundedSource;
import org.apache.beam.sdk.io.CountingSource;
import org.apache.beam.sdk.io.GenerateSequence;
import org.apache.beam.sdk.io.Read;
import org.apache.beam.sdk.io.UnboundedSource;
import org.apache.beam.sdk.options.Default;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.testing.PAssert;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.transforms.Count;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.Flatten;
import org.apache.beam.sdk.transforms.MapElements;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.SerializableFunction;
import org.apache.beam.sdk.transforms.SimpleFunction;
import org.apache.beam.sdk.transforms.Sum;
import org.apache.beam.sdk.transforms.display.DisplayData;
import org.apache.beam.sdk.transforms.windowing.AfterWatermark;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.GlobalWindows;
import org.apache.beam.sdk.transforms.windowing.Window;
import org.apache.beam.sdk.util.CoderUtils;
import org.apache.beam.sdk.util.IllegalMutationException;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PBegin;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionList;
import org.apache.beam.sdk.values.PDone;
import org.apache.beam.sdk.values.TypeDescriptor;
import org.apache.beam.sdk.values.TypeDescriptors;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.checkerframework.checker.nullness.qual.Nullable;
import org.hamcrest.Matchers;
import org.joda.time.Duration;
import org.joda.time.Instant;
import org.junit.Rule;
import org.junit.Test;
import org.junit.internal.matchers.ThrowableMessageMatcher;
import org.junit.rules.ExpectedException;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;

/** Tests for basic {@link DirectRunner} functionality. */
@RunWith(JUnit4.class)
public class DirectRunnerTest implements Serializable {
  @Rule public transient ExpectedException thrown = ExpectedException.none();

  private Pipeline getPipeline() {
    return getPipeline(true);
  }

  private Pipeline getPipeline(boolean blockOnRun) {
    PipelineOptions opts = PipelineOptionsFactory.create();
    opts.setRunner(DirectRunner.class);
    opts.as(DirectOptions.class).setBlockOnRun(blockOnRun);
    return Pipeline.create(opts);
  }

  @Test
  public void defaultRunnerLoaded() {
    assertThat(
        DirectRunner.class,
        Matchers.<Class<? extends PipelineRunner>>equalTo(
            PipelineOptionsFactory.create().getRunner()));
  }

  @Test
  public void wordCountShouldSucceed() throws Throwable {
    Pipeline p = getPipeline();

    PCollection<KV<String, Long>> counts =
        p.apply(Create.of("foo", "bar", "foo", "baz", "bar", "foo"))
            .apply(
                MapElements.via(
                    new SimpleFunction<String, String>() {
                      @Override
                      public String apply(String input) {
                        return input;
                      }
                    }))
            .apply(Count.perElement());
    PCollection<String> countStrs =
        counts.apply(
            MapElements.via(
                new SimpleFunction<KV<String, Long>, String>() {
                  @Override
                  public String apply(KV<String, Long> input) {
                    return String.format("%s: %s", input.getKey(), input.getValue());
                  }
                }));

    PAssert.that(countStrs).containsInAnyOrder("baz: 1", "bar: 2", "foo: 3");

    DirectPipelineResult result = (DirectPipelineResult) p.run();
    result.waitUntilFinish();
  }

  private static AtomicInteger changed;

  @Test
  public void reusePipelineSucceeds() throws Throwable {
    Pipeline p = getPipeline();

    changed = new AtomicInteger(0);
    PCollection<KV<String, Long>> counts =
        p.apply(Create.of("foo", "bar", "foo", "baz", "bar", "foo"))
            .apply(
                MapElements.via(
                    new SimpleFunction<String, String>() {
                      @Override
                      public String apply(String input) {
                        return input;
                      }
                    }))
            .apply(Count.perElement());
    PCollection<String> countStrs =
        counts.apply(
            MapElements.via(
                new SimpleFunction<KV<String, Long>, String>() {
                  @Override
                  public String apply(KV<String, Long> input) {
                    return String.format("%s: %s", input.getKey(), input.getValue());
                  }
                }));

    counts.apply(
        ParDo.of(
            new DoFn<KV<String, Long>, Void>() {
              @ProcessElement
              public void updateChanged(ProcessContext c) {
                changed.getAndIncrement();
              }
            }));

    PAssert.that(countStrs).containsInAnyOrder("baz: 1", "bar: 2", "foo: 3");

    DirectPipelineResult result = (DirectPipelineResult) p.run();
    result.waitUntilFinish();

    DirectPipelineResult otherResult = (DirectPipelineResult) p.run();
    otherResult.waitUntilFinish();

    assertThat("Each element should have been processed twice", changed.get(), equalTo(6));
  }

  @Test
  public void byteArrayCountShouldSucceed() {
    Pipeline p = getPipeline();

    SerializableFunction<Integer, byte[]> getBytes =
        input -> {
          try {
            return CoderUtils.encodeToByteArray(VarIntCoder.of(), input);
          } catch (CoderException e) {
            fail("Unexpected Coder Exception " + e);
            throw new AssertionError("Unreachable");
          }
        };
    TypeDescriptor<byte[]> td = new TypeDescriptor<byte[]>() {};
    PCollection<byte[]> foos =
        p.apply(Create.of(1, 1, 1, 2, 2, 3)).apply(MapElements.into(td).via(getBytes));
    PCollection<byte[]> msync =
        p.apply(Create.of(1, -2, -8, -16)).apply(MapElements.into(td).via(getBytes));
    PCollection<byte[]> bytes = PCollectionList.of(foos).and(msync).apply(Flatten.pCollections());
    PCollection<KV<byte[], Long>> counts = bytes.apply(Count.perElement());
    PCollection<KV<Integer, Long>> countsBackToString =
        counts.apply(
            MapElements.via(
                new SimpleFunction<KV<byte[], Long>, KV<Integer, Long>>() {
                  @Override
                  public KV<Integer, Long> apply(KV<byte[], Long> input) {
                    try {
                      return KV.of(
                          CoderUtils.decodeFromByteArray(VarIntCoder.of(), input.getKey()),
                          input.getValue());
                    } catch (CoderException e) {
                      fail("Unexpected Coder Exception " + e);
                      throw new AssertionError("Unreachable");
                    }
                  }
                }));

    Map<Integer, Long> expected =
        ImmutableMap.<Integer, Long>builder()
            .put(1, 4L)
            .put(2, 2L)
            .put(3, 1L)
            .put(-2, 1L)
            .put(-8, 1L)
            .put(-16, 1L)
            .build();
    PAssert.thatMap(countsBackToString).isEqualTo(expected);
  }

  @Test
  public void splitsInputs() {
    Pipeline p = getPipeline();
    PCollection<Long> longs = p.apply(Read.from(MustSplitSource.of(CountingSource.upTo(3))));

    PAssert.that(longs).containsInAnyOrder(0L, 1L, 2L);
    p.run();
  }

  @Test
  public void cancelShouldStopPipeline() throws Exception {
    PipelineOptions opts = TestPipeline.testingPipelineOptions();
    opts.as(DirectOptions.class).setBlockOnRun(false);
    opts.setRunner(DirectRunner.class);

    final Pipeline p = Pipeline.create(opts);
    p.apply(GenerateSequence.from(0).withRate(1L, Duration.standardSeconds(1)));

    final BlockingQueue<PipelineResult> resultExchange = new ArrayBlockingQueue<>(1);
    Runnable cancelRunnable =
        () -> {
          try {
            resultExchange.take().cancel();
          } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new IllegalStateException(e);
          } catch (IOException e) {
            throw new IllegalStateException(e);
          }
        };

    Callable<PipelineResult> runPipelineRunnable =
        () -> {
          PipelineResult res = p.run();
          try {
            resultExchange.put(res);
          } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new IllegalStateException(e);
          }
          return res;
        };

    ExecutorService executor = Executors.newCachedThreadPool();
    Future<?> cancelResult = executor.submit(cancelRunnable);
    Future<PipelineResult> result = executor.submit(runPipelineRunnable);

    cancelResult.get();
    // If cancel doesn't work, this will hang forever
    result.get().waitUntilFinish();
  }

  @Test
  public void testWaitUntilFinishTimeout() throws Exception {
    DirectOptions options = PipelineOptionsFactory.as(DirectOptions.class);
    options.setBlockOnRun(false);
    options.setRunner(DirectRunner.class);
    Pipeline p = Pipeline.create(options);
    p.apply(Create.of(1L))
        .apply(
            ParDo.of(
                new DoFn<Long, Long>() {
                  @ProcessElement
                  public void hang(ProcessContext context) throws InterruptedException {
                    // Hangs "forever"
                    Thread.sleep(Long.MAX_VALUE);
                  }
                }));
    PipelineResult result = p.run();
    // The pipeline should never complete;
    assertThat(result.getState(), is(State.RUNNING));
    // Must time out, otherwise this test will never complete
    assertEquals(null, result.waitUntilFinish(Duration.millis(1L)));
    // Ensure multiple calls complete
    assertEquals(null, result.waitUntilFinish(Duration.millis(1L)));
  }

  private static final AtomicLong TEARDOWN_CALL = new AtomicLong(-1);

  @Test
  public void tearsDownFnsBeforeFinishing() {
    TEARDOWN_CALL.set(-1);
    final Pipeline pipeline = getPipeline();
    pipeline
        .apply(Create.of("a"))
        .apply(
            ParDo.of(
                new DoFn<String, String>() {
                  @ProcessElement
                  public void onElement(final ProcessContext ctx) {
                    // no-op
                  }

                  @Teardown
                  public void teardown() {
                    // just to not have a fast execution hiding an issue until we have a shutdown
                    // callback
                    try {
                      Thread.sleep(1000);
                    } catch (final InterruptedException e) {
                      throw new AssertionError(e);
                    }
                    TEARDOWN_CALL.set(System.nanoTime());
                  }
                }));
    final PipelineResult pipelineResult = pipeline.run();
    pipelineResult.waitUntilFinish();

    final long doneTs = System.nanoTime();
    final long tearDownTs = TEARDOWN_CALL.get();
    assertThat(tearDownTs, greaterThan(0L));
    assertThat(doneTs, greaterThan(tearDownTs));
  }

  @Test
  public void transformDisplayDataExceptionShouldFail() {
    DoFn<Integer, Integer> brokenDoFn =
        new DoFn<Integer, Integer>() {
          @ProcessElement
          public void processElement(ProcessContext c) throws Exception {}

          @Override
          public void populateDisplayData(DisplayData.Builder builder) {
            throw new RuntimeException("oh noes!");
          }
        };

    Pipeline p = getPipeline();
    p.apply(Create.of(1, 2, 3)).apply(ParDo.of(brokenDoFn));

    thrown.expectMessage(brokenDoFn.getClass().getName());
    thrown.expectCause(ThrowableMessageMatcher.hasMessage(is("oh noes!")));
    p.run();
  }

  /**
   * Tests that a {@link DoFn} that mutates an output with a good equals() fails in the {@link
   * DirectRunner}.
   */
  @Test
  public void testMutatingOutputThenOutputDoFnError() throws Exception {
    Pipeline pipeline = getPipeline();

    pipeline
        .apply(Create.of(42))
        .apply(
            ParDo.of(
                new DoFn<Integer, List<Integer>>() {
                  @ProcessElement
                  public void processElement(ProcessContext c) {
                    List<Integer> outputList = Arrays.asList(1, 2, 3, 4);
                    c.output(outputList);
                    outputList.set(0, 37);
                    c.output(outputList);
                  }
                }));

    thrown.expect(IllegalMutationException.class);
    thrown.expectMessage("output");
    thrown.expectMessage("must not be mutated");
    pipeline.run();
  }

  /**
   * Tests that a {@link DoFn} that mutates an output with a good equals() fails in the {@link
   * DirectRunner}.
   */
  @Test
  public void testMutatingOutputWithEnforcementDisabledSucceeds() throws Exception {
    PipelineOptions options = PipelineOptionsFactory.create();
    options.setRunner(DirectRunner.class);
    options.as(DirectOptions.class).setEnforceImmutability(false);
    Pipeline pipeline = Pipeline.create(options);

    pipeline
        .apply(Create.of(42))
        .apply(
            ParDo.of(
                new DoFn<Integer, List<Integer>>() {
                  @ProcessElement
                  public void processElement(ProcessContext c) {
                    List<Integer> outputList = Arrays.asList(1, 2, 3, 4);
                    c.output(outputList);
                    outputList.set(0, 37);
                    c.output(outputList);
                  }
                }));

    pipeline.run();
  }

  /**
   * Tests that a {@link DoFn} that mutates an output with a good equals() fails in the {@link
   * DirectRunner}.
   */
  @Test
  public void testMutatingOutputThenTerminateDoFnError() throws Exception {
    Pipeline pipeline = getPipeline();

    pipeline
        .apply(Create.of(42))
        .apply(
            ParDo.of(
                new DoFn<Integer, List<Integer>>() {
                  @ProcessElement
                  public void processElement(ProcessContext c) {
                    List<Integer> outputList = Arrays.asList(1, 2, 3, 4);
                    c.output(outputList);
                    outputList.set(0, 37);
                  }
                }));

    thrown.expect(IllegalMutationException.class);
    thrown.expectMessage("output");
    thrown.expectMessage("must not be mutated");
    pipeline.run();
  }

  /**
   * Tests that a {@link DoFn} that mutates an output with a bad equals() still fails in the {@link
   * DirectRunner}.
   */
  @Test
  public void testMutatingOutputCoderDoFnError() throws Exception {
    Pipeline pipeline = getPipeline();

    pipeline
        .apply(Create.of(42))
        .apply(
            ParDo.of(
                new DoFn<Integer, byte[]>() {
                  @ProcessElement
                  public void processElement(ProcessContext c) {
                    byte[] outputArray = new byte[] {0x1, 0x2, 0x3};
                    c.output(outputArray);
                    outputArray[0] = 0xa;
                    c.output(outputArray);
                  }
                }));

    thrown.expect(IllegalMutationException.class);
    thrown.expectMessage("output");
    thrown.expectMessage("must not be mutated");
    pipeline.run();
  }

  /**
   * Tests that a {@link DoFn} that mutates its input with a good equals() fails in the {@link
   * DirectRunner}.
   */
  @Test
  public void testMutatingInputDoFnError() throws Exception {
    Pipeline pipeline = getPipeline();

    pipeline
        .apply(
            Create.of(Arrays.asList(1, 2, 3), Arrays.asList(4, 5, 6))
                .withCoder(ListCoder.of(VarIntCoder.of())))
        .apply(
            ParDo.of(
                new DoFn<List<Integer>, Integer>() {
                  @ProcessElement
                  public void processElement(ProcessContext c) {
                    List<Integer> inputList = c.element();
                    inputList.set(0, 37);
                    c.output(12);
                  }
                }));

    thrown.expect(IllegalMutationException.class);
    thrown.expectMessage("Input");
    thrown.expectMessage("must not be mutated");
    pipeline.run();
  }

  /**
   * Tests that a {@link DoFn} that mutates an input with a bad equals() still fails in the {@link
   * DirectRunner}.
   */
  @Test
  public void testMutatingInputCoderDoFnError() throws Exception {
    Pipeline pipeline = getPipeline();

    pipeline
        .apply(Create.of(new byte[] {0x1, 0x2, 0x3}, new byte[] {0x4, 0x5, 0x6}))
        .apply(
            ParDo.of(
                new DoFn<byte[], Integer>() {
                  @ProcessElement
                  public void processElement(ProcessContext c) {
                    byte[] inputArray = c.element();
                    inputArray[0] = 0xa;
                    c.output(13);
                  }
                }));

    thrown.expect(IllegalMutationException.class);
    thrown.expectMessage("Input");
    thrown.expectMessage("must not be mutated");
    pipeline.run();
  }

  @Test
  public void testUnencodableOutputElement() throws Exception {
    Pipeline p = getPipeline();
    PCollection<Long> pcollection =
        p.apply(Create.of((Void) null))
            .apply(
                ParDo.of(
                    new DoFn<Void, Long>() {
                      @ProcessElement
                      public void processElement(ProcessContext c) {
                        c.output(null);
                      }
                    }))
            .setCoder(VarLongCoder.of());
    pcollection.apply(
        ParDo.of(
            new DoFn<Long, Long>() {
              @ProcessElement
              public void unreachable(ProcessContext c) {
                fail("Pipeline should fail to encode a null Long in VarLongCoder");
              }
            }));

    thrown.expectCause(isA(CoderException.class));
    thrown.expectMessage("cannot encode a null Long");
    p.run();
  }

  @Test
  public void testUnencodableOutputFromBoundedRead() throws Exception {
    Pipeline p = getPipeline();
    p.apply(GenerateSequence.from(0).to(10)).setCoder(new LongNoDecodeCoder());

    thrown.expectCause(isA(CoderException.class));
    thrown.expectMessage("Cannot decode a long");
    p.run();
  }

  @Test
  public void testUnencodableOutputFromUnboundedRead() {
    Pipeline p = getPipeline();
    p.apply(GenerateSequence.from(0)).setCoder(new LongNoDecodeCoder());

    thrown.expectCause(isA(CoderException.class));
    thrown.expectMessage("Cannot decode a long");
    p.run();
  }

  /**
   * Tests that {@link DirectRunner#fromOptions(PipelineOptions)} drops {@link PipelineOptions}
   * marked with {@link JsonIgnore} fields.
   */
  @Test
  public void testFromOptionsIfIgnoredFieldsGettingDropped() {
    TestSerializationOfOptions options =
        PipelineOptionsFactory.fromArgs(
                "--foo=testValue", "--ignoredField=overridden", "--runner=DirectRunner")
            .as(TestSerializationOfOptions.class);

    assertEquals("testValue", options.getFoo());
    assertEquals("overridden", options.getIgnoredField());
    Pipeline p = Pipeline.create(options);
    PCollection<Integer> pc =
        p.apply(Create.of("1"))
            .apply(
                ParDo.of(
                    new DoFn<String, Integer>() {
                      @ProcessElement
                      public void processElement(ProcessContext c) {
                        TestSerializationOfOptions options =
                            c.getPipelineOptions().as(TestSerializationOfOptions.class);
                        assertEquals("testValue", options.getFoo());
                        assertEquals("not overridden", options.getIgnoredField());
                        c.output(Integer.parseInt(c.element()));
                      }
                    }));
    PAssert.that(pc).containsInAnyOrder(1);
    p.run();
  }

  /**
   * Test running of {@link Pipeline} which has two {@link POutput POutputs} and finishing the first
   * one triggers data being fed into the second one.
   */
  @Test(timeout = 10000)
  public void testTwoPOutputsInPipelineWithCascade() throws InterruptedException {

    StaticQueue<Integer> start = StaticQueue.of("start", VarIntCoder.of());
    StaticQueue<Integer> messages = StaticQueue.of("messages", VarIntCoder.of());

    Pipeline pipeline = getPipeline(false);
    pipeline.begin().apply("outputStartSignal", outputStartTo(start));
    PCollection<Integer> result =
        pipeline
            .apply("processMessages", messages.read())
            .apply(
                Window.<Integer>into(new GlobalWindows())
                    .triggering(AfterWatermark.pastEndOfWindow())
                    .discardingFiredPanes()
                    .withAllowedLateness(Duration.ZERO))
            .apply(Sum.integersGlobally());

    // the result should be 6, after the data will have been written
    PAssert.that(result).containsInAnyOrder(6);

    PipelineResult run = pipeline.run();

    // wait until a message has been written to the start queue
    while (start.take() == null) {}

    // and publish messages
    messages.add(1).add(2).add(3).terminate();

    run.waitUntilFinish();
  }

  private PTransform<PBegin, PDone> outputStartTo(StaticQueue<Integer> queue) {
    return new PTransform<PBegin, PDone>() {
      @Override
      public PDone expand(PBegin input) {
        input
            .apply(Create.of(1))
            .apply(
                MapElements.into(TypeDescriptors.voids())
                    .via(
                        in -> {
                          queue.add(in);
                          return null;
                        }));
        return PDone.in(input.getPipeline());
      }
    };
  }

  /**
   * Options for testing if {@link DirectRunner} drops {@link PipelineOptions} marked with {@link
   * JsonIgnore} fields.
   */
  public interface TestSerializationOfOptions extends PipelineOptions {
    String getFoo();

    void setFoo(String foo);

    @JsonIgnore
    @Default.String("not overridden")
    String getIgnoredField();

    void setIgnoredField(String value);
  }

  private static class LongNoDecodeCoder extends AtomicCoder<Long> {
    @Override
    public void encode(Long value, OutputStream outStream) throws IOException {}

    @Override
    public Long decode(InputStream inStream) throws IOException {
      throw new CoderException("Cannot decode a long");
    }
  }

  private static class MustSplitSource<T> extends BoundedSource<T> {
    public static <T> BoundedSource<T> of(BoundedSource<T> underlying) {
      return new MustSplitSource<>(underlying);
    }

    private final BoundedSource<T> underlying;

    public MustSplitSource(BoundedSource<T> underlying) {
      this.underlying = underlying;
    }

    @Override
    public List<? extends BoundedSource<T>> split(
        long desiredBundleSizeBytes, PipelineOptions options) throws Exception {
      // Must have more than
      checkState(
          desiredBundleSizeBytes < getEstimatedSizeBytes(options),
          "Must split into more than one source");
      return underlying.split(desiredBundleSizeBytes, options);
    }

    @Override
    public long getEstimatedSizeBytes(PipelineOptions options) throws Exception {
      return underlying.getEstimatedSizeBytes(options);
    }

    @Override
    public BoundedReader<T> createReader(PipelineOptions options) throws IOException {
      throw new IllegalStateException(
          "The MustSplitSource cannot create a reader without being split first");
    }

    @Override
    public void validate() {
      underlying.validate();
    }

    @Override
    public Coder<T> getOutputCoder() {
      return underlying.getOutputCoder();
    }
  }

  private static class StaticQueue<T> implements Serializable {

    static class StaticQueueSource<T> extends UnboundedSource<T, StaticQueueSource.Checkpoint<T>> {

      static class Checkpoint<T> implements CheckpointMark, Serializable {

        final T read;

        Checkpoint(T read) {
          this.read = read;
        }

        @Override
        public void finalizeCheckpoint() throws IOException {
          // nop
        }
      }

      final StaticQueue<T> queue;

      StaticQueueSource(StaticQueue<T> queue) {
        this.queue = queue;
      }

      @Override
      public List<? extends UnboundedSource<T, Checkpoint<T>>> split(
          int desiredNumSplits, PipelineOptions options) throws Exception {
        return Arrays.asList(this);
      }

      @Override
      public UnboundedReader<T> createReader(PipelineOptions po, Checkpoint<T> cmt) {
        return new UnboundedReader<T>() {

          T read = cmt == null ? null : cmt.read;
          boolean finished = false;

          @Override
          public boolean start() throws IOException {
            return advance();
          }

          @Override
          public boolean advance() throws IOException {
            try {
              Optional<T> taken = queue.take();
              if (taken.isPresent()) {
                read = taken.get();
                return true;
              }
              finished = true;
              return false;
            } catch (InterruptedException ex) {
              throw new IOException(ex);
            }
          }

          @Override
          public Instant getWatermark() {
            if (finished) {
              return BoundedWindow.TIMESTAMP_MAX_VALUE;
            }
            return BoundedWindow.TIMESTAMP_MIN_VALUE;
          }

          @Override
          public CheckpointMark getCheckpointMark() {
            return new Checkpoint(read);
          }

          @Override
          public UnboundedSource<T, ?> getCurrentSource() {
            return StaticQueueSource.this;
          }

          @Override
          public T getCurrent() throws NoSuchElementException {
            return read;
          }

          @Override
          public Instant getCurrentTimestamp() {
            return getWatermark();
          }

          @Override
          public void close() throws IOException {
            // nop
          }
        };
      }

      @SuppressWarnings("unchecked")
      @Override
      public Coder<Checkpoint<T>> getCheckpointMarkCoder() {
        return (Coder) SerializableCoder.of(Checkpoint.class);
      }

      @Override
      public Coder<T> getOutputCoder() {
        return queue.coder;
      }
    }

    static final Map<String, StaticQueue<?>> QUEUES = new ConcurrentHashMap<>();

    static <T> StaticQueue<T> of(String name, Coder<T> coder) {
      return new StaticQueue<>(name, coder);
    }

    private final String name;
    private final Coder<T> coder;
    private final transient BlockingQueue<Optional<T>> queue = new ArrayBlockingQueue<>(10);

    StaticQueue(String name, Coder<T> coder) {
      this.name = name;
      this.coder = coder;
      Preconditions.checkState(
          QUEUES.put(name, this) == null, "Queue " + name + " already exists.");
    }

    StaticQueue<T> add(T elem) {
      queue.add(Optional.of(elem));
      return this;
    }

    @Nullable
    Optional<T> take() throws InterruptedException {
      return queue.take();
    }

    PTransform<PBegin, PCollection<T>> read() {
      return new PTransform<PBegin, PCollection<T>>() {
        @Override
        public PCollection<T> expand(PBegin input) {
          return input.apply("readFrom:" + name, Read.from(asSource()));
        }
      };
    }

    UnboundedSource<T, ?> asSource() {
      return new StaticQueueSource<>(this);
    }

    void terminate() {
      queue.add(Optional.empty());
    }

    private Object readResolve() {
      return QUEUES.get(name);
    }
  }
}