PubsubSource actualSource = new PubsubSource(source); PubsubReader reader = actualSource.createReader(options, null); SubscriptionPath createdSubscription = reader.subscription; assertThat(createdSubscription, not(nullValue())); CoderUtils.clone(actualSource.getCheckpointMarkCoder(), checkpoint); assertThat(checkpoint.subscriptionPath, not(nullValue())); assertThat(checkpoint.subscriptionPath, equalTo(deserCheckpoint.subscriptionPath)); PubsubReader readerFromOriginal = actualSource.createReader(options, checkpoint); PubsubReader readerFromDeser = actualSource.createReader(options, deserCheckpoint);
@Test public void noSubscriptionSplitGeneratesSubscription() throws Exception { TopicPath topicPath = PubsubClient.topicPathFromName("my_project", "my_topic"); factory = PubsubTestClient.createFactoryForCreateSubscription(); PubsubUnboundedSource source = new PubsubUnboundedSource( factory, StaticValueProvider.of(PubsubClient.projectPathFromId("my_project")), StaticValueProvider.of(topicPath), null /* subscription */, null /* timestampLabel */, null /* idLabel */, false /* needsAttributes */); assertThat(source.getSubscription(), nullValue()); assertThat(source.getSubscription(), nullValue()); PipelineOptions options = PipelineOptionsFactory.create(); List<PubsubSource> splits = (new PubsubSource(source)).split(3, options); // We have at least one returned split assertThat(splits, hasSize(greaterThan(0))); for (PubsubSource split : splits) { // Each split is equal assertThat(split, equalTo(splits.get(0))); } assertThat(splits.get(0).subscriptionPath, not(nullValue())); }
PubsubReader reader = primSource.createReader(p.getOptions(), null); PubsubTestClient pubsubClient = (PubsubTestClient) reader.getPubsubClient();
PubsubReader reader = primSource.createReader(p.getOptions(), null); CoderUtils.encodeToByteArray(primSource.getCheckpointMarkCoder(), checkpoint); checkpoint = CoderUtils.decodeFromByteArray(primSource.getCheckpointMarkCoder(), checkpointBytes); assertEquals(1, checkpoint.notYetReadIds.size()); assertEquals("ackid_1", checkpoint.notYetReadIds.get(0)); reader = primSource.createReader(p.getOptions(), checkpoint); assertTrue(reader.start()); assertEquals("data_1", data(reader.getCurrent()));
@Test public void timeoutAckExtensions() throws IOException { setupOneMessage(); PubsubReader reader = primSource.createReader(p.getOptions(), null); PubsubTestClient pubsubClient = (PubsubTestClient) reader.getPubsubClient(); // Pull the first message but don't take a checkpoint for it. assertTrue(reader.start()); assertEquals(DATA, data(reader.getCurrent())); // Extend the ack. now.addAndGet(55 * 1000); pubsubClient.advance(); assertFalse(reader.advance()); // Let the ack expire. for (int i = 0; i < 3; i++) { now.addAndGet(25 * 1000); pubsubClient.advance(); assertFalse(reader.advance()); } // Wait for resend. now.addAndGet(25 * 1000); pubsubClient.advance(); // Reread the same message. assertTrue(reader.advance()); assertEquals(DATA, data(reader.getCurrent())); // Now ACK the message. PubsubCheckpoint checkpoint = reader.getCheckpointMark(); checkpoint.finalizeCheckpoint(); reader.close(); }
@Test public void timeoutAckAndRereadOneMessage() throws IOException { setupOneMessage(); PubsubReader reader = primSource.createReader(p.getOptions(), null); PubsubTestClient pubsubClient = (PubsubTestClient) reader.getPubsubClient(); assertTrue(reader.start()); assertEquals(DATA, data(reader.getCurrent())); // Let the ACK deadline for the above expire. now.addAndGet(65 * 1000); pubsubClient.advance(); // We'll now receive the same message again. assertTrue(reader.advance()); assertEquals(DATA, data(reader.getCurrent())); assertFalse(reader.advance()); // Now ACK the message. PubsubCheckpoint checkpoint = reader.getCheckpointMark(); checkpoint.finalizeCheckpoint(); reader.close(); }
@Test public void extendAck() throws IOException { setupOneMessage(); PubsubReader reader = primSource.createReader(p.getOptions(), null); PubsubTestClient pubsubClient = (PubsubTestClient) reader.getPubsubClient(); // Pull the first message but don't take a checkpoint for it. assertTrue(reader.start()); assertEquals(DATA, data(reader.getCurrent())); // Extend the ack now.addAndGet(55 * 1000); pubsubClient.advance(); assertFalse(reader.advance()); // Extend the ack again now.addAndGet(25 * 1000); pubsubClient.advance(); assertFalse(reader.advance()); // Now ACK the message. PubsubCheckpoint checkpoint = reader.getCheckpointMark(); checkpoint.finalizeCheckpoint(); reader.close(); }
@Override public List<PubsubSource> split(int desiredNumSplits, PipelineOptions options) throws Exception { List<PubsubSource> result = new ArrayList<>(desiredNumSplits); PubsubSource splitSource = this; if (subscriptionPath == null) { splitSource = new PubsubSource( outer, StaticValueProvider.of(outer.createRandomSubscription(options))); } for (int i = 0; i < desiredNumSplits * SCALE_OUT; i++) { // Since the source is immutable and Pubsub automatically shards we simply // replicate ourselves the requested number of times result.add(splitSource); } return result; }
private void setupOneMessage(Iterable<IncomingMessage> incoming) { now = new AtomicLong(REQ_TIME); clock = () -> now.get(); factory = PubsubTestClient.createFactoryForPull(clock, SUBSCRIPTION, ACK_TIMEOUT_S, incoming); PubsubUnboundedSource source = new PubsubUnboundedSource( clock, factory, null, null, StaticValueProvider.of(SUBSCRIPTION), TIMESTAMP_ATTRIBUTE, ID_ATTRIBUTE, true /* needsAttributes */); primSource = new PubsubSource(source); }
@Test public void readOneMessage() throws IOException { setupOneMessage(); PubsubReader reader = primSource.createReader(p.getOptions(), null); // Read one message. assertTrue(reader.start()); assertEquals(DATA, data(reader.getCurrent())); assertFalse(reader.advance()); // ACK the message. PubsubCheckpoint checkpoint = reader.getCheckpointMark(); checkpoint.finalizeCheckpoint(); reader.close(); }
@Override public PCollection<PubsubMessage> expand(PBegin input) { return input .getPipeline() .begin() .apply(Read.from(new PubsubSource(this))) .apply( "PubsubUnboundedSource.Stats", ParDo.of( new StatsFn(pubsubFactory, subscription, topic, timestampAttribute, idAttribute))); }
/** Tests that checkpoints finalized after the reader is closed succeed. */ @Test public void closeWithActiveCheckpoints() throws Exception { setupOneMessage(); PubsubReader reader = primSource.createReader(p.getOptions(), null); reader.start(); PubsubCheckpoint checkpoint = reader.getCheckpointMark(); reader.close(); checkpoint.finalizeCheckpoint(); } }
@Test public void checkpointCoderIsSane() { setupOneMessage(ImmutableList.of()); CoderProperties.coderSerializable(primSource.getCheckpointMarkCoder()); // Since we only serialize/deserialize the 'notYetReadIds', and we don't want to make // equals on checkpoints ignore those fields, we'll test serialization and deserialization // of checkpoints in multipleReaders below. }