content: /*
   * 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.flink.translation.wrappers.streaming.io.source.unbounded;
  
  import static org.apache.beam.runners.flink.translation.wrappers.streaming.io.source.unbounded.FlinkUnboundedSourceReader.PENDING_BYTES_METRIC_NAME;
  import static org.junit.Assert.assertEquals;
  import static org.junit.Assert.assertFalse;
  import static org.junit.Assert.assertNotNull;
  import static org.junit.Assert.assertNull;
  import static org.junit.Assert.assertTrue;
  
  import java.util.ArrayList;
  import java.util.List;
  import java.util.Map;
  import java.util.Random;
  import java.util.concurrent.CompletableFuture;
  import java.util.concurrent.ScheduledExecutorService;
  import java.util.concurrent.atomic.AtomicReference;
  import java.util.function.Function;
  import org.apache.beam.runners.flink.FlinkPipelineOptions;
  import org.apache.beam.runners.flink.translation.wrappers.streaming.io.TestCountingSource;
  import org.apache.beam.runners.flink.translation.wrappers.streaming.io.source.EmptyUnboundedSource;
  import org.apache.beam.runners.flink.translation.wrappers.streaming.io.source.FlinkSourceReaderTestBase;
  import org.apache.beam.runners.flink.translation.wrappers.streaming.io.source.FlinkSourceSplit;
  import org.apache.beam.runners.flink.translation.wrappers.streaming.io.source.SourceTestMetrics.TestMetricGroup;
  import org.apache.beam.sdk.io.Source;
  import org.apache.beam.sdk.options.PipelineOptions;
  import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
  import org.apache.beam.sdk.values.KV;
  import org.apache.beam.sdk.values.ValueWithRecordId;
  import org.apache.beam.sdk.values.WindowedValue;
  import org.apache.flink.api.common.eventtime.Watermark;
  import org.apache.flink.api.connector.source.ReaderOutput;
  import org.apache.flink.api.connector.source.SourceOutput;
  import org.apache.flink.api.connector.source.SourceReader;
  import org.apache.flink.api.connector.source.SourceReaderContext;
  import org.apache.flink.core.io.InputStatus;
  import org.apache.flink.core.testutils.ManuallyTriggeredScheduledExecutorService;
  import org.apache.flink.metrics.Gauge;
  import org.checkerframework.checker.nullness.qual.Nullable;
  import org.joda.time.Instant;
  import org.junit.Test;
  
  /** Unite tests for {@link FlinkUnboundedSourceReader}. */
  public class FlinkUnboundedSourceReaderTest
      extends FlinkSourceReaderTestBase<WindowedValue<ValueWithRecordId<KV<Integer, Integer>>>> {
  
    @Test
    public void testSnapshotStateAndRestore() throws Exception {
      final int numSplits = 2;
      final int numRecordsPerSplit = 10;
  
      List<FlinkSourceSplit<KV<Integer, Integer>>> splits =
          createSplits(numSplits, numRecordsPerSplit, 0);
      RecordsValidatingOutput validatingOutput = new RecordsValidatingOutput(splits);
      List<FlinkSourceSplit<KV<Integer, Integer>>> snapshot;
  
      // Create a reader, take a snapshot.
      try (SourceReader<
              WindowedValue<ValueWithRecordId<KV<Integer, Integer>>>,
              FlinkSourceSplit<KV<Integer, Integer>>>
          reader = createReader()) {
        pollAndValidate(reader, splits, validatingOutput, numSplits * numRecordsPerSplit / 2);
        snapshot = reader.snapshotState(0L);
      }
  
      // Create another reader, add the snapshot splits back.
      try (SourceReader<
              WindowedValue<ValueWithRecordId<KV<Integer, Integer>>>,
              FlinkSourceSplit<KV<Integer, Integer>>>
          reader = createReader()) {
        pollAndValidate(reader, snapshot, validatingOutput, Integer.MAX_VALUE);
      }
    }
  
    /**
     * This is a concurrency correctness test. It verifies that the main thread is always waken up by
     * the alarm runner executed in the executor thread.
     */
    @Test(timeout = 60000L)
    public void testIsAvailableAlwaysWakenUp() throws Exception {
      final int numFuturesRequired = 1_000;
      List<CompletableFuture<Void>> futures = new ArrayList<>();
      AtomicReference<Exception> exceptionRef = new AtomicReference<>();
  
      List<FlinkSourceSplit<KV<Integer, Integer>>> splits = new ArrayList<>();
      splits.add(new FlinkSourceSplit<>(0, new DummySource(Integer.MAX_VALUE)));
      RecordsValidatingOutput validatingOutput = new RecordsValidatingOutput(splits);
      ManuallyTriggeredScheduledExecutorService executor =
          new ManuallyTriggeredScheduledExecutorService();
  
      try (SourceReader<
              WindowedValue<ValueWithRecordId<KV<Integer, Integer>>>,
              FlinkSourceSplit<KV<Integer, Integer>>>
          reader = createReader(executor, Long.MAX_VALUE)) {
        reader.start();
        reader.addSplits(splits);
  
        Thread mainThread =
            new Thread(
                () -> {
                  try {
                    while (futures.size() < numFuturesRequired) {
                      // This poll will return NOTHING_AVAILABLE after each record emission.
                      if (reader.pollNext(validatingOutput) == InputStatus.NOTHING_AVAILABLE) {
                        CompletableFuture<Void> future = reader.isAvailable();
                        future.get();
                        futures.add(future);
                      }
                    }
                  } catch (Exception e) {
                    if (!exceptionRef.compareAndSet(null, e)) {
                      exceptionRef.get().addSuppressed(e);
                    }
                  }
                },
                "MainThread");
  
        Thread executorThread =
            new Thread(
                () -> {
                  while (futures.size() < numFuturesRequired) {
                    executor.triggerScheduledTasks();
                  }
                },
                "ExecutorThread");
  
        mainThread.start();
        executorThread.start();
        executorThread.join();
      }
    }
  
    @Test
    public void testIsAvailableOnSplitChangeWhenNoDataAvailableForAliveReaders() throws Exception {
      List<FlinkSourceSplit<KV<Integer, Integer>>> splits1 = new ArrayList<>();
      List<FlinkSourceSplit<KV<Integer, Integer>>> splits2 = new ArrayList<>();
      splits1.add(new FlinkSourceSplit<>(0, new DummySource(0)));
      splits2.add(new FlinkSourceSplit<>(1, new DummySource(0)));
      RecordsValidatingOutput validatingOutput = new RecordsValidatingOutput(splits1);
      ManuallyTriggeredScheduledExecutorService executor =
          new ManuallyTriggeredScheduledExecutorService();
  
      try (SourceReader<
              WindowedValue<ValueWithRecordId<KV<Integer, Integer>>>,
              FlinkSourceSplit<KV<Integer, Integer>>>
          reader = createReader(executor, Long.MAX_VALUE)) {
        reader.start();
        reader.addSplits(splits1);
  
        assertEquals(
            "The reader should have nothing available",
            InputStatus.NOTHING_AVAILABLE,
            reader.pollNext(validatingOutput));
  
        CompletableFuture<Void> future1 = reader.isAvailable();
        assertFalse("Future1 should be uncompleted without live split.", future1.isDone());
  
        reader.addSplits(splits2);
        assertTrue("Future1 should be completed upon addition of new splits.", future1.isDone());
  
        CompletableFuture<Void> future2 = reader.isAvailable();
        assertFalse("Future2 should be uncompleted without live split.", future2.isDone());
  
        reader.notifyNoMoreSplits();
        assertTrue("Future2 should be completed upon NoMoreSplitsNotification.", future2.isDone());
      }
    }
  
    @Test
    public void testWatermark() throws Exception {
      ManuallyTriggeredScheduledExecutorService executor =
          new ManuallyTriggeredScheduledExecutorService();
      try (FlinkUnboundedSourceReader<KV<Integer, Integer>> reader =
          (FlinkUnboundedSourceReader<KV<Integer, Integer>>) createReader(executor, -1L)) {
        List<FlinkSourceSplit<KV<Integer, Integer>>> splits = createSplits(2, 10, 0);
        RecordsValidatingOutput validatingOutput = new RecordsValidatingOutput(splits);
  
        reader.start();
        reader.addSplits(splits);
  
        // Poll 3 records from split 0 and 2 records from split 1.
        for (int i = 0; i < 5; i++) {
          reader.pollNext(validatingOutput);
        }
  
        Map<String, TestSourceOutput> sourceOutputs = validatingOutput.createdSourceOutputs();
        assertEquals("There should be 2 source outputs created.", 2, sourceOutputs.size());
        assertNull(sourceOutputs.get("0").watermark());
        assertNull(sourceOutputs.get("1").watermark());
  
        // Trigger the periodic task marking the watermark emission flag.
        executor.triggerScheduledTasks();
        // Poll one more time to actually emit the watermark. Getting record value 2 from split_1.
        reader.pollNext(validatingOutput);
  
        assertEquals(3, sourceOutputs.get("0").watermark().getTimestamp());
        assertEquals(2, sourceOutputs.get("1").watermark().getTimestamp());
  
        // Poll one more time to ensure no additional watermark is emitted. Getting record value 3
        // from split_0.
        reader.pollNext(validatingOutput);
        assertEquals(3, sourceOutputs.get("0").watermark().getTimestamp());
        assertEquals(2, sourceOutputs.get("1").watermark().getTimestamp());
  
        // Trigger the task to mark the watermark emission flag again.
        executor.triggerScheduledTasks();
        // Poll to actually emit the watermark. Getting (split_1 -> 3).
        reader.pollNext(validatingOutput);
  
        assertEquals(4, sourceOutputs.get("0").watermark().getTimestamp());
        assertEquals(3, sourceOutputs.get("1").watermark().getTimestamp());
      }
    }
  
    @Test
    public void testWatermarkOnEmptySource() throws Exception {
      ManuallyTriggeredScheduledExecutorService executor =
          new ManuallyTriggeredScheduledExecutorService();
      AtomicReference<Instant> watermark = new AtomicReference<>(BoundedWindow.TIMESTAMP_MIN_VALUE);
      ReaderOutput<WindowedValue<ValueWithRecordId<KV<Integer, Integer>>>> output =
          new ReaderOutput<WindowedValue<ValueWithRecordId<KV<Integer, Integer>>>>() {
            @Override
            public void collect(WindowedValue<ValueWithRecordId<KV<Integer, Integer>>> unused) {}
  
            @Override
            public void collect(
                WindowedValue<ValueWithRecordId<KV<Integer, Integer>>> unused, long l) {}
  
            @Override
            public void emitWatermark(Watermark w) {
              watermark.compareAndSet(
                  BoundedWindow.TIMESTAMP_MIN_VALUE, Instant.ofEpochMilli(w.getTimestamp()));
            }
  
            @Override
            public void markIdle() {}
  
            @Override
            public SourceOutput<WindowedValue<ValueWithRecordId<KV<Integer, Integer>>>>
                createOutputForSplit(String s) {
              return this;
            }
  
            @Override
            public void releaseOutputForSplit(String s) {}
  
            @Override
            public void markActive() {}
          };
      Instant now = Instant.now();
      try (FlinkUnboundedSourceReader<KV<Integer, Integer>> reader =
          (FlinkUnboundedSourceReader<KV<Integer, Integer>>) createReader(executor, -1L)) {
        List<FlinkSourceSplit<KV<Integer, Integer>>> splits = createEmptySplits(2);
        reader.start();
        reader.addSplits(splits);
        reader.notifyNoMoreSplits();
  
        for (int i = 0; i < 4; i++) {
          assertEquals(InputStatus.NOTHING_AVAILABLE, reader.pollNext(output));
        }
  
        // move first reader to 'now'
        ((EmptyUnboundedSource<KV<Integer, Integer>>) splits.get(0).getBeamSplitSource())
            .setWatermark(now);
        // force trigger timeout
        executor.triggerScheduledTasks();
        for (int i = 0; i < 4; i++) {
          assertEquals(InputStatus.NOTHING_AVAILABLE, reader.pollNext(output));
        }
  
        // check we have emitted watermark
        assertEquals(now, watermark.get());
  
        // move first reader to end of time
        ((EmptyUnboundedSource<KV<Integer, Integer>>) splits.get(0).getBeamSplitSource())
            .setWatermark(BoundedWindow.TIMESTAMP_MAX_VALUE);
  
        for (int i = 0; i < 4; i++) {
          assertEquals(InputStatus.NOTHING_AVAILABLE, reader.pollNext(output));
        }
  
        // move the second reader to end of time
        ((EmptyUnboundedSource<KV<Integer, Integer>>) splits.get(1).getBeamSplitSource())
            .setWatermark(BoundedWindow.TIMESTAMP_MAX_VALUE);
  
        assertEquals(InputStatus.END_OF_INPUT, reader.pollNext(output));
      }
    }
  
    @Test
    public void testWatermarkOnNoSplits() throws Exception {
      ManuallyTriggeredScheduledExecutorService executor =
          new ManuallyTriggeredScheduledExecutorService();
      try (FlinkUnboundedSourceReader<KV<Integer, Integer>> reader =
          (FlinkUnboundedSourceReader<KV<Integer, Integer>>) createReader(executor, -1L)) {
        reader.start();
        reader.notifyNoMoreSplits();
        assertEquals(InputStatus.END_OF_INPUT, reader.pollNext(null));
      }
    }
  
    @Test
    public void testPendingBytesMetric() throws Exception {
      ManuallyTriggeredScheduledExecutorService executor =
          new ManuallyTriggeredScheduledExecutorService();
      TestMetricGroup testMetricGroup = new TestMetricGroup();
      try (SourceReader<
              WindowedValue<ValueWithRecordId<KV<Integer, Integer>>>,
              FlinkSourceSplit<KV<Integer, Integer>>>
          reader = createReader(executor, 0L, null, testMetricGroup)) {
        reader.start();
  
        List<FlinkSourceSplit<KV<Integer, Integer>>> splits = createSplits(2, 10, 0);
        reader.addSplits(splits);
        RecordsValidatingOutput validatingOutput = new RecordsValidatingOutput(splits);
  
        // Need to poll once to create all the readers.
        reader.pollNext(validatingOutput);
        Gauge<Long> pendingBytesGauge =
            (Gauge<Long>) testMetricGroup.registeredGauge.get(PENDING_BYTES_METRIC_NAME);
        assertNotNull(pendingBytesGauge);
        // The TestCountingSource.CountingSourceReader always return 7L as backlog bytes. Because we
        // have 2 splits,
        // the expected value is the magic number 14 here.
        assertEquals(14L, pendingBytesGauge.getValue().longValue());
      }
    }
  
    @Test
    public void testCheckMarksFinalized() throws Exception {
  
      final int numSplits = 2;
      final int numRecordsPerSplit = 10;
  
      List<FlinkSourceSplit<KV<Integer, Integer>>> splits =
          createSplits(numSplits, numRecordsPerSplit, 0);
      RecordsValidatingOutput validatingOutput = new RecordsValidatingOutput(splits);
      // Create a reader, take a snapshot.
      try (SourceReader<
              WindowedValue<ValueWithRecordId<KV<Integer, Integer>>>,
              FlinkSourceSplit<KV<Integer, Integer>>>
          reader = createReader()) {
        List<Integer> finalizeTracker = new ArrayList<>();
        TestCountingSource.setFinalizeTracker(finalizeTracker);
        pollAndValidate(reader, splits, validatingOutput, numSplits * numRecordsPerSplit / 2);
        assertTrue(finalizeTracker.isEmpty());
        reader.snapshotState(0L);
        // notifyCheckpointComplete is normally called by the SourceOperator
        reader.notifyCheckpointComplete(0L);
        // every split should be finalized
        assertEquals(numSplits, finalizeTracker.size());
        pollAndValidate(reader, splits, validatingOutput, numSplits);
        // no notifyCheckpointComplete here, assume the checkpoint failed
        reader.snapshotState(1L);
        pollAndValidate(reader, splits, validatingOutput, numSplits);
        reader.snapshotState(2L);
        reader.notifyCheckpointComplete(2L);
        // 2 * numSplits more should be finalized
        assertEquals(3 * numSplits, finalizeTracker.size());
      }
    }
  
    // --------------- private helper classes -----------------
    /** A source whose advance() method only returns true occasionally. */
    private static class DummySource extends TestCountingSource {
  
      public DummySource(int numMessagesPerShard) {
        super(numMessagesPerShard);
      }
  
      @Override
      public CountingSourceReader createReader(
          PipelineOptions options, @Nullable CounterMark checkpointMark) {
        CountingSourceReader reader = new DummySourceReader();
        createdReaders().add(reader);
        return reader;
      }
  
      private class DummySourceReader extends TestCountingSource.CountingSourceReader {
        private final Random random = new Random();
  
        public DummySourceReader() {
          super(0);
        }
  
        @Override
        public boolean advance() {
          // Return true once every three times advance is invoked.
          if (random.nextInt(3) == 0) {
            return super.advance();
          } else {
            return false;
          }
        }
      }
    }
  
    // --------------- abstract methods impl ------------------
    @Override
    protected KV<Integer, Integer> getKVPairs(
        WindowedValue<ValueWithRecordId<KV<Integer, Integer>>> record) {
      return record.getValue().getValue();
    }
  
    @Override
    protected FlinkUnboundedSourceReader<KV<Integer, Integer>> createReader(
        ScheduledExecutorService executor,
        long idleTimeoutMs,
        Function<WindowedValue<ValueWithRecordId<KV<Integer, Integer>>>, Long> timestampExtractor,
        TestMetricGroup metricGroup) {
      FlinkPipelineOptions pipelineOptions = FlinkPipelineOptions.defaults();
      pipelineOptions.setShutdownSourcesAfterIdleMs(idleTimeoutMs);
      pipelineOptions.setAutoWatermarkInterval(10L);
      SourceReaderContext mockContext = createSourceReaderContext(metricGroup);
      if (executor != null) {
        return new FlinkUnboundedSourceReader<>(
            "FlinkUnboundedReader", mockContext, pipelineOptions, executor, timestampExtractor);
      } else {
        return new FlinkUnboundedSourceReader<>(
            "FlinkUnboundedReader", mockContext, pipelineOptions, timestampExtractor);
      }
    }
  
    @Override
    protected Source<KV<Integer, Integer>> createBeamSource(int splitIndex, int numRecordsPerSplit) {
      return new TestCountingSource(numRecordsPerSplit)
          .withShardNumber(splitIndex)
          .withFixedNumSplits(1);
    }
  }
  