Reimplement Pub/Sub Lite's I/O using UnboundedSource.

SDF is not well supported by the default mode of dataflow, and UnboundedSource is supported by all modes.
apache · Aug 8, 2022 · 6ccf85d · 6ccf85d
1 parent 2f08b0e
commit 6ccf85d
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 9 deletions.
diff --git a/...orm/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/UnboundedReaderImpl.java b/...orm/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/UnboundedReaderImpl.java
@@ -102,15 +102,18 @@ public boolean advance() throws IOException {
  throw new IOException("Subscriber failed: ", subscriber.failureCause());
  }
  if (advanced) {
- SequencedMessage peeked = subscriber.peek().get();
- Offset nextOffset = Offset.of(peeked.getCursor().getOffset() + 1);
- checkState(nextOffset.value() > fetchOffset.value());
- fetchOffset = nextOffset;
- lastMessageTimestamp = Optional.of(getTimestamp(peeked));
  subscriber.pop();
  }
- advanced = subscriber.peek().isPresent();
- return advanced;
+ Optional<SequencedMessage> next = subscriber.peek();
+ advanced = next.isPresent();
+ if (!advanced) {
+ return false;
+ }
+ Offset nextOffset = Offset.of(next.get().getCursor().getOffset() + 1);
+ checkState(nextOffset.value() > fetchOffset.value());
+ fetchOffset = nextOffset;
+ lastMessageTimestamp = Optional.of(getTimestamp(next.get()));
+ return true;
  }
 
  @Override

diff --git a/...src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/UnboundedReaderImplTest.java b/...src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/UnboundedReaderImplTest.java
@@ -105,12 +105,14 @@ public void setUp() {
 
  @Test
  public void startAdvances() throws Exception {
- SequencedMessage message = messageWith("abc", 2, Instant.now());
+ Instant ts = Instant.now();
+ SequencedMessage message = messageWith("abc", 2, ts);
  doReturn(Optional.of(message)).when(subscriber).peek();
  assertTrue(reader.start());
  verify(subscriber).startAsync();
  verify(subscriber).awaitRunning(1, TimeUnit.MINUTES);
  assertEquals(reader.getCurrent(), message);
+ assertEquals(reader.getWatermark(), ts);
  }
 
  @Test
@@ -120,6 +122,7 @@ public void startAdvancesNoMessage() throws Exception {
  verify(subscriber).startAsync();
  verify(subscriber).awaitRunning(1, TimeUnit.MINUTES);
  assertThrows(NoSuchElementException.class, reader::getCurrent);
+ assertEquals(BoundedWindow.TIMESTAMP_MIN_VALUE, reader.getWatermark());
  }
 
  @Test
@@ -142,7 +145,7 @@ public void advanceWithPreviousValue() throws Exception {
  assertTrue(reader.advance());
  assertEquals(reader.getCurrent(), message1);
  assertEquals(reader.getCurrentTimestamp(), ts1);
- assertEquals(reader.getWatermark(), BoundedWindow.TIMESTAMP_MIN_VALUE);
+ assertEquals(reader.getWatermark(), ts1);
  doAnswer(
  (Answer<Void>)
  args -> {