1717 */
1818package org .apache .beam .sdk .io .pulsar ;
1919
20- import edu .umd .cs .findbugs .annotations .SuppressFBWarnings ;
2120import java .io .IOException ;
21+ import java .time .Duration ;
2222import java .util .concurrent .TimeUnit ;
2323import org .apache .beam .sdk .coders .Coder ;
2424import org .apache .beam .sdk .io .range .OffsetRange ;
25+ import org .apache .beam .sdk .options .PipelineOptions ;
2526import org .apache .beam .sdk .transforms .DoFn ;
2627import org .apache .beam .sdk .transforms .SerializableFunction ;
2728import org .apache .beam .sdk .transforms .splittabledofn .GrowableOffsetRangeTracker ;
3031import org .apache .beam .sdk .transforms .splittabledofn .WatermarkEstimator ;
3132import org .apache .beam .sdk .transforms .splittabledofn .WatermarkEstimators ;
3233import org .apache .beam .sdk .transforms .windowing .BoundedWindow ;
34+ import org .apache .beam .vendor .guava .v32_1_2_jre .com .google .common .base .MoreObjects ;
35+ import org .apache .beam .vendor .guava .v32_1_2_jre .com .google .common .base .Stopwatch ;
36+ import org .apache .beam .vendor .guava .v32_1_2_jre .com .google .common .base .Strings ;
3337import org .apache .beam .vendor .guava .v32_1_2_jre .com .google .common .base .Supplier ;
3438import org .apache .beam .vendor .guava .v32_1_2_jre .com .google .common .base .Suppliers ;
3539import org .apache .pulsar .client .admin .PulsarAdmin ;
4044import org .apache .pulsar .client .api .PulsarClientException ;
4145import org .apache .pulsar .client .api .Reader ;
4246import org .apache .pulsar .client .api .ReaderBuilder ;
47+ import org .checkerframework .checker .nullness .qual .MonotonicNonNull ;
48+ import org .checkerframework .checker .nullness .qual .Nullable ;
4349import org .joda .time .Instant ;
4450import org .slf4j .Logger ;
4551import org .slf4j .LoggerFactory ;
4652
4753/**
48- * Transform for reading from Apache Pulsar. Support is currently incomplete, and there may be bugs;
49- * see https://github.com/apache/beam/issues/31078 for more info, and comment in that issue if you
50- * run into issues with this IO.
54+ * DoFn for reading from Apache Pulsar based on Pulsar {@link Reader} from the start message id. It
55+ * does not support split or acknowledge message get read.
5156 */
5257@ DoFn .UnboundedPerElement
53- @ SuppressWarnings ({"rawtypes" , "nullness" })
54- @ SuppressFBWarnings (value = "CT_CONSTRUCTOR_THROW" , justification = "Initialization is safe." )
55- public class ReadFromPulsarDoFn extends DoFn <PulsarSourceDescriptor , PulsarMessage > {
58+ @ SuppressWarnings ("nullness" )
59+ public class NaiveReadFromPulsarDoFn <T > extends DoFn <PulsarSourceDescriptor , T > {
5660
57- private static final Logger LOG = LoggerFactory .getLogger (ReadFromPulsarDoFn .class );
58- private SerializableFunction <String , PulsarClient > pulsarClientSerializableFunction ;
59- private PulsarClient client ;
60- private PulsarAdmin admin ;
61- private String clientUrl ;
62- private String adminUrl ;
61+ private static final Logger LOG = LoggerFactory .getLogger (NaiveReadFromPulsarDoFn .class );
62+ private final SerializableFunction <String , PulsarClient > clientFn ;
63+ private final SerializableFunction <String , PulsarAdmin > adminFn ;
64+ private final SerializableFunction <Message <?>, T > outputFn ;
65+ private final java .time .Duration pollingTimeout ;
66+ private transient @ MonotonicNonNull PulsarClient client ;
67+ private transient @ MonotonicNonNull PulsarAdmin admin ;
68+ private @ MonotonicNonNull String clientUrl ;
69+ private @ Nullable final String adminUrl ;
6370
6471 private final SerializableFunction <Message <byte []>, Instant > extractOutputTimestampFn ;
6572
66- public ReadFromPulsarDoFn (PulsarIO .Read transform ) {
67- this .extractOutputTimestampFn = transform .getExtractOutputTimestampFn ();
73+ public NaiveReadFromPulsarDoFn (PulsarIO .Read <T > transform ) {
74+ this .extractOutputTimestampFn =
75+ transform .getTimestampType () == PulsarIO .ReadTimestampType .PUBLISH_TIME
76+ ? record -> new Instant (record .getPublishTime ())
77+ : ignored -> Instant .now ();
78+ this .pollingTimeout = Duration .ofSeconds (transform .getConsumerPollingTimeout ());
79+ this .outputFn = transform .getOutputFn ();
6880 this .clientUrl = transform .getClientUrl ();
6981 this .adminUrl = transform .getAdminUrl ();
70- this .pulsarClientSerializableFunction = transform .getPulsarClient ();
82+ this .clientFn =
83+ MoreObjects .firstNonNull (
84+ transform .getPulsarClient (), PulsarIOUtils .PULSAR_CLIENT_SERIALIZABLE_FUNCTION );
85+ this .adminFn =
86+ MoreObjects .firstNonNull (
87+ transform .getPulsarAdmin (), PulsarIOUtils .PULSAR_ADMIN_SERIALIZABLE_FUNCTION );
88+ admin = null ;
7189 }
7290
73- // Open connection to Pulsar clients
91+ /** Open connection to Pulsar clients. */
7492 @ Setup
7593 public void initPulsarClients () throws Exception {
76- if (this .clientUrl == null ) {
77- this .clientUrl = PulsarIOUtils .SERVICE_URL ;
78- }
79- if (this .adminUrl == null ) {
80- this .adminUrl = PulsarIOUtils .SERVICE_HTTP_URL ;
81- }
82-
83- if (this .client == null ) {
84- this .client = pulsarClientSerializableFunction .apply (this .clientUrl );
85- if (this .client == null ) {
86- this .client = PulsarClient .builder ().serviceUrl (clientUrl ).build ();
94+ if (client == null ) {
95+ if (clientUrl == null ) {
96+ clientUrl = PulsarIOUtils .LOCAL_SERVICE_URL ;
8797 }
98+ client = clientFn .apply (clientUrl );
8899 }
89100
90- if (this .admin == null ) {
91- this .admin =
92- PulsarAdmin .builder ()
93- .serviceHttpUrl (adminUrl )
94- .tlsTrustCertsFilePath (null )
95- .allowTlsInsecureConnection (false )
96- .build ();
101+ // admin is optional
102+ if (this .admin == null && !Strings .isNullOrEmpty (adminUrl )) {
103+ admin = adminFn .apply (adminUrl );
97104 }
98105 }
99106
100- // Close connection to Pulsar clients
107+ /** Close connection to Pulsar clients. */
101108 @ Teardown
102109 public void teardown () throws Exception {
103110 this .client .close ();
104- this .admin .close ();
111+ if (this .admin != null ) {
112+ this .admin .close ();
113+ }
105114 }
106115
107116 @ GetInitialRestriction
@@ -152,25 +161,47 @@ public Coder<OffsetRange> getRestrictionCoder() {
152161 public ProcessContinuation processElement (
153162 @ Element PulsarSourceDescriptor pulsarSourceDescriptor ,
154163 RestrictionTracker <OffsetRange , Long > tracker ,
155- WatermarkEstimator watermarkEstimator ,
156- OutputReceiver <PulsarMessage > output )
164+ WatermarkEstimator < Instant > watermarkEstimator ,
165+ OutputReceiver <T > output )
157166 throws IOException {
158167 long startTimestamp = tracker .currentRestriction ().getFrom ();
159168 String topicDescriptor = pulsarSourceDescriptor .getTopic ();
160169 try (Reader <byte []> reader = newReader (this .client , topicDescriptor )) {
161170 if (startTimestamp > 0 ) {
171+ // seek is inclusive
162172 reader .seek (startTimestamp );
163173 }
164- while (true ) {
165- if (reader .hasReachedEndOfTopic ()) {
166- reader .close ();
167- return ProcessContinuation .stop ();
174+ if (reader .hasReachedEndOfTopic ()) {
175+ // topic has terminated
176+ tracker .tryClaim (Long .MAX_VALUE );
177+ reader .close ();
178+ return ProcessContinuation .stop ();
179+ }
180+ final Stopwatch pollTimer = Stopwatch .createUnstarted ();
181+ Duration remainingTimeout = pollingTimeout ;
182+ while (Duration .ZERO .compareTo (remainingTimeout ) < 0 ) {
183+ pollTimer .reset ().start ();
184+ Message <byte []> message =
185+ reader .readNext ((int ) remainingTimeout .toMillis (), TimeUnit .MILLISECONDS );
186+ final Duration elapsed = pollTimer .elapsed ();
187+ try {
188+ remainingTimeout = remainingTimeout .minus (elapsed );
189+ } catch (ArithmeticException e ) {
190+ remainingTimeout = Duration .ZERO ;
168191 }
169- Message <byte []> message = reader .readNext ();
192+ // No progress when the polling timeout expired.
193+ // Self-checkpoint and move to process the next element.
170194 if (message == null ) {
171195 return ProcessContinuation .resume ();
172196 }
173197 Long currentTimestamp = message .getPublishTime ();
198+ if (currentTimestamp < startTimestamp ) {
199+ LOG .warn (
200+ "Skip late message of publish time {} before startTimestamp {}" ,
201+ currentTimestamp ,
202+ startTimestamp );
203+ continue ;
204+ }
174205 // if tracker.tryclaim() return true, sdf must execute work otherwise
175206 // doFn must exit processElement() without doing any work associated
176207 // or claiming more work
@@ -186,12 +217,21 @@ public ProcessContinuation processElement(
186217 return ProcessContinuation .stop ();
187218 }
188219 }
189- PulsarMessage pulsarMessage =
190- new PulsarMessage (message .getTopicName (), message .getPublishTime (), message );
220+ T messageT = outputFn .apply (message );
191221 Instant outputTimestamp = extractOutputTimestampFn .apply (message );
192- output .outputWithTimestamp (pulsarMessage , outputTimestamp );
222+ output .outputWithTimestamp (messageT , outputTimestamp );
193223 }
194224 }
225+ return ProcessContinuation .resume ();
226+ }
227+
228+ @ SplitRestriction
229+ public void splitRestriction (
230+ @ Restriction OffsetRange restriction ,
231+ OutputReceiver <OffsetRange > receiver ,
232+ PipelineOptions unused ) {
233+ // read based on Reader does not support split
234+ receiver .output (restriction );
195235 }
196236
197237 @ GetInitialWatermarkEstimatorState
@@ -221,27 +261,34 @@ public OffsetRangeTracker restrictionTracker(
221261 private static class PulsarLatestOffsetEstimator
222262 implements GrowableOffsetRangeTracker .RangeEndEstimator {
223263
224- private final Supplier <Message > memoizedBacklog ;
264+ private final @ Nullable Supplier <Message < byte []> > memoizedBacklog ;
225265
226- private PulsarLatestOffsetEstimator (PulsarAdmin admin , String topic ) {
227- this .memoizedBacklog =
228- Suppliers .memoizeWithExpiration (
229- () -> {
230- try {
231- Message <byte []> lastMsg = admin .topics ().examineMessage (topic , "latest" , 1 );
232- return lastMsg ;
233- } catch (PulsarAdminException e ) {
234- throw new RuntimeException (e );
235- }
236- },
237- 1 ,
238- TimeUnit .SECONDS );
266+ private PulsarLatestOffsetEstimator (@ Nullable PulsarAdmin admin , String topic ) {
267+ if (admin != null ) {
268+ this .memoizedBacklog =
269+ Suppliers .memoizeWithExpiration (
270+ () -> {
271+ try {
272+ return admin .topics ().examineMessage (topic , "latest" , 1 );
273+ } catch (PulsarAdminException e ) {
274+ throw new RuntimeException (e );
275+ }
276+ },
277+ 1 ,
278+ TimeUnit .SECONDS );
279+ } else {
280+ memoizedBacklog = null ;
281+ }
239282 }
240283
241284 @ Override
242285 public long estimate () {
243- Message <byte []> msg = memoizedBacklog .get ();
244- return msg .getPublishTime ();
286+ if (memoizedBacklog != null ) {
287+ Message <byte []> msg = memoizedBacklog .get ();
288+ return msg .getPublishTime ();
289+ } else {
290+ return Long .MIN_VALUE ;
291+ }
245292 }
246293 }
247294
0 commit comments