private static void advanceOnce(UnboundedReader<?> reader, boolean isStarted) throws IOException { if (!isStarted && reader.start()) { return; } while (!reader.advance()) { // very rarely will there be more than one attempts. // In case of a bug we might end up looping forever, and test will fail with a timeout. // Avoid hard cpu spinning in case of a test failure. try { Thread.sleep(1); } catch (InterruptedException e) { throw new RuntimeException(e); } } }
@Override public Object readCurrent() { try { if (!isStarted) { isStarted = true; isCurrentAvailable = reader.start(); } else { isCurrentAvailable = reader.advance(); } } catch (final Exception e) { throw new RuntimeException(e); } if (isCurrentAvailable) { final O elem = reader.getCurrent(); return WindowedValue.timestampedValueInGlobalWindow(elem, reader.getCurrentTimestamp()); } else { throw new NoSuchElementException(); } }
@Override public void emitTuples() { try { if (!available) { available = reader.advance(); } if (available) { OutputT data = reader.getCurrent(); Instant timestamp = reader.getCurrentTimestamp(); available = reader.advance(); if (traceTuples) { LOG.debug("\nemitting '{}' timestamp {}\n", data, timestamp); } output.emit( DataTuple.of( WindowedValue.of(data, timestamp, GlobalWindow.INSTANCE, PaneInfo.NO_FIRING))); } } catch (Exception e) { Throwables.propagateIfPossible(e); throw new RuntimeException(e); } } }
} while (numElements < ARBITRARY_MAX_ELEMENTS && reader.advance()); Instant watermark = reader.getWatermark();
@ProcessElement public void process( @Element Shard<T> shard, OutputReceiver<ValueWithRecordId<T>> out, PipelineOptions options) throws Exception { Instant endTime = shard.getMaxReadTime() == null ? null : Instant.now().plus(shard.getMaxReadTime()); if (shard.getMaxNumRecords() <= 0 || (shard.getMaxReadTime() != null && shard.getMaxReadTime().getMillis() == 0)) { return; } try (UnboundedSource.UnboundedReader<T> reader = SerializableUtils.clone(shard.getSource()).createReader(options, null)) { for (long i = 0L; i < shard.getMaxNumRecords(); ++i) { boolean available = (i == 0) ? reader.start() : reader.advance(); if (!available && !advanceWithBackoff(reader, endTime)) { break; } out.outputWithTimestamp( new ValueWithRecordId<T>(reader.getCurrent(), reader.getCurrentRecordId()), reader.getCurrentTimestamp()); } reader.getCheckpointMark().finalizeCheckpoint(); } }
@Test public void testUnboundedSourceCheckpointMark() throws Exception { UnboundedSource<Long, CounterMark> source = CountingSource.unboundedWithTimestampFn(new ValueAsTimestampFn()); UnboundedReader<Long> reader = source.createReader(null, null); final long numToSkip = 3; assertTrue(reader.start()); // Advance the source numToSkip elements and manually save state. for (long l = 0; l < numToSkip; ++l) { reader.advance(); } // Confirm that we get the expected element in sequence before checkpointing. assertEquals(numToSkip, (long) reader.getCurrent()); assertEquals(numToSkip, reader.getCurrentTimestamp().getMillis()); // Checkpoint and restart, and confirm that the source continues correctly. CounterMark mark = CoderUtils.clone(source.getCheckpointMarkCoder(), (CounterMark) reader.getCheckpointMark()); reader = source.createReader(null, mark); assertTrue(reader.start()); // Confirm that we get the next element in sequence. assertEquals(numToSkip + 1, (long) reader.getCurrent()); assertEquals(numToSkip + 1, reader.getCurrentTimestamp().getMillis()); } }
private boolean advanceWithBackoff(UnboundedReader<T> reader, Instant endTime) throws IOException { // Try reading from the source with exponential backoff BackOff backoff = BACKOFF_FACTORY.backoff(); long nextSleep = backoff.nextBackOffMillis(); while (true) { if (nextSleep == BackOff.STOP || (endTime != null && Instant.now().isAfter(endTime))) { return false; } if (reader.advance()) { return true; } Uninterruptibles.sleepUninterruptibly(nextSleep, TimeUnit.MILLISECONDS); nextSleep = backoff.nextBackOffMillis(); } } }
private boolean startReader( UnboundedReader<OutputT> reader, UnboundedSourceShard<OutputT, CheckpointMarkT> shard) throws IOException { if (shard.getExistingReader() == null) { return reader.start(); } else { return shard.getExistingReader().advance(); } }