• Home
  • Features
  • Pricing
  • Docs
  • Announcements
  • Sign In

juarezr / spark-streaming-google-pubsub / 60

02 Oct 2026 07:10AM UTC coverage: 70.532% (+4.3%) from 66.187%
60

push

github

web-flow
fix: implement connection error handling to halt on non-recoverable configuration errors (#21)

- fix: implement connection error handling to halt on non-recoverable
configuration errors
- fix: validation logic to ignore AUTO placeholders for receiveTime and
processingTime
- test: add parameterized tests for schema and metadata modes
- test: add unit tests for dynamic and mixed schema modes
- refactor: removed dead code from the StreamingPull migration leftovers

764 of 1196 branches covered (63.88%)

Branch coverage included in aggregate %.

59 of 115 new or added lines in 9 files covered. (51.3%)

2 existing lines in 2 files now uncovered.

1608 of 2167 relevant lines covered (74.2%)

0.74 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

80.98
/src/main/java/io/github/juarezr/spark/pubsub/structured/PubSubMicroBatchStream.java
1
package io.github.juarezr.spark.pubsub.structured;
2

3
import io.github.juarezr.spark.pubsub.config.AckMode;
4
import io.github.juarezr.spark.pubsub.config.GatherMode;
5
import io.github.juarezr.spark.pubsub.config.PubSubConfig;
6
import java.io.IOException;
7
import java.time.Duration;
8
import java.time.Instant;
9
import java.util.ArrayList;
10
import java.util.List;
11
import java.util.Map;
12
import java.util.Optional;
13
import java.util.concurrent.ConcurrentHashMap;
14
import java.util.concurrent.ConcurrentMap;
15
import java.util.concurrent.atomic.AtomicInteger;
16
import java.util.concurrent.atomic.AtomicLong;
17
import org.apache.spark.sql.connector.read.InputPartition;
18
import org.apache.spark.sql.connector.read.PartitionReaderFactory;
19
import org.apache.spark.sql.connector.read.streaming.MicroBatchStream;
20
import org.apache.spark.sql.connector.read.streaming.Offset;
21
import org.apache.spark.sql.connector.read.streaming.ReadLimit;
22
import org.apache.spark.sql.connector.read.streaming.ReportsSourceMetrics;
23
import org.apache.spark.sql.connector.read.streaming.SupportsTriggerAvailableNow;
24
import org.apache.spark.sql.types.StructType;
25
import org.slf4j.Logger;
26
import org.slf4j.LoggerFactory;
27

28
/**
29
 * Spark micro-batch source that polls a warm streaming-pull {@link PubSubClient} queue (not unary
30
 * Pub/Sub Pull RPCs). Progress is tracked with synthetic offsets; the subscription cursor remains
31
 * the durable source of truth across process restarts (no rewind unless configured).
32
 */
33
final class PubSubMicroBatchStream
34
    implements MicroBatchStream, ReportsSourceMetrics, SupportsTriggerAvailableNow {
35
  private static final Logger LOG = LoggerFactory.getLogger(PubSubMicroBatchStream.class);
1 ✔
36

37
  /**
38
   * Pub/Sub has no log-end offset. AvailableNow treats this many consecutive empty polls as idle
39
   * and returns {@code null} so Spark stops.
40
   */
41
  static final int AVAILABLE_NOW_IDLE_POLLS = 3;
42

43
  static final Duration RECEIVE_IDLE = Duration.ofSeconds(1);
1 ✔
44

45
  private final PubSubConfig config;
46
  private final StructType readSchema;
47
  private final PubSubClient client;
48
  private final AckCoordinator ackCoordinator;
49
  private final AtomicLong nextBatchId = new AtomicLong(0);
1 ✔
50
  private final AtomicInteger lastGatherMessageCount = new AtomicInteger(0);
1 ✔
51
  private final AtomicLong lastGatherPayloadBytes = new AtomicLong(0);
1 ✔
52
  private volatile Long lastGatherNewestMessageAgeMs;
53
  private final AtomicLong lastReportedRetryAttempts = new AtomicLong(0);
1 ✔
54
  private final int numPartitions;
55
  private volatile PubSubOffset lastProduced;
56
  private volatile PubSubOffset currentOffset = PubSubOffset.empty(-1L);
1 ✔
57
  private volatile boolean firstBatchLogged;
58
  private volatile Long lastGatheredBatchId;
59
  private volatile PublishTimeWindow lastGatheredWindow;
60
  private final ConcurrentMap<Long, List<PulledMessage>> messagesByBatch =
1 ✔
61
      new ConcurrentHashMap<>();
62
  private final AutoBatchTime autoBatchTime;
63
  private volatile boolean cancelled;
64
  private volatile boolean availableNow;
65
  private volatile boolean limitReached;
66

67
  PubSubMicroBatchStream(PubSubConfig config, int numPartitions) {
68
    this(config, numPartitions, PubSubSchema.tableSchema(config, null));
×
69
  }
×
70

71
  PubSubMicroBatchStream(PubSubConfig config, int numPartitions, StructType readSchema) {
72
    this(config, numPartitions, new PubSubClient(config), true, readSchema);
×
73
  }
×
74

75
  PubSubMicroBatchStream(
76
      PubSubConfig config, int numPartitions, PubSubClient client, boolean startClient) {
77
    this(config, numPartitions, client, startClient, PubSubSchema.tableSchema(config, null));
1 ✔
78
  }
1 ✔
79

80
  PubSubMicroBatchStream(
81
      PubSubConfig config,
82
      int numPartitions,
83
      PubSubClient client,
84
      boolean startClient,
85
      StructType readSchema) {
1 ✔
86
    this.config = config;
1 ✔
87
    this.readSchema = readSchema == null ? PubSubSchema.tableSchema(config, null) : readSchema;
1 !
88
    this.client = client;
1 ✔
89
    this.ackCoordinator = new AckCoordinator(config.ackMode());
1 ✔
90
    this.autoBatchTime = AutoBatchTime.create(config);
1 ✔
91
    this.numPartitions = Math.max(1, numPartitions);
1 ✔
92
    if (!startClient) {
1 !
93
      return;
1 ✔
94
    }
95
    try {
96
      this.client.start();
×
97
    } catch (IOException e) {
×
98
      throw new IllegalStateException("Unable to start Pub/Sub client", e);
×
99
    }
×
100
  }
×
101

102
  @Override
103
  public void prepareForTriggerAvailableNow() {
104
    this.availableNow = true;
1 ✔
105
    if (config.limitTime().isPresent()) {
1 ✔
106
      LOG.info(
1 ✔
107
          "BATCH: AvailableNow recovery drain; limitTime={} exclusive; stop at bound or {}"
108
              + " consecutive empty polls",
109
          config.limitTime().get(),
1 ✔
110
          AVAILABLE_NOW_IDLE_POLLS);
1 ✔
111
    } else {
112
      LOG.info(
1 ✔
113
          "BATCH: AvailableNow drain enabled; stop after {} consecutive empty polls"
114
              + " (Pub/Sub has no log-end offset; use seek=snapshot for recovery)",
115
          AVAILABLE_NOW_IDLE_POLLS);
1 ✔
116
    }
117
  }
1 ✔
118

119
  @Override
120
  public ReadLimit getDefaultReadLimit() {
121
    if (config.batchCount() > 0) {
1 !
122
      return ReadLimit.maxRows(config.batchCount());
1 ✔
123
    }
124
    return ReadLimit.allAvailable();
×
125
  }
126

127
  @Override
128
  public Offset reportLatestOffset() {
129
    return lastProduced != null ? lastProduced : currentOffset;
1 !
130
  }
131

132
  @Override
133
  public Offset latestOffset() {
134
    return latestOffset(null, AdmissionLimits.from(config, null), false);
1 ✔
135
  }
136

137
  @Override
138
  public Offset latestOffset(Offset startOffset, ReadLimit limit) {
139
    return latestOffset(startOffset, AdmissionLimits.from(config, limit), true);
1 ✔
140
  }
141

142
  /**
143
   * Spark 3.5 commits batch N-1 only when constructing batch N. Returning the same {@code
144
   * lastProduced} after Spark has already accepted it as {@code startOffset} makes {@code
145
   * constructNextBatch} see no new data, so {@link #commit} never runs and the query idles with the
146
   * first batch still uncommitted.
147
   */
148
  private Offset latestOffset(Offset startOffset, AdmissionLimits limits, boolean nullIfEmpty) {
149

150
    if (config.limitTime().isPresent() && !this.availableNow) {
1 ✔
151
      throw new IllegalStateException(
1 ✔
152
          "limitTime requires Trigger.AvailableNow(); ProcessingTime cannot stop at the bound");
153
    }
154
    if (lastProduced != null && !startConsumed(startOffset, lastProduced)) {
1 ✔
155
      return lastProduced;
1 ✔
156
    }
157
    if (this.limitReached) {
1 ✔
158
      resetMetrics();
1 ✔
159
      return nullIfEmpty ? null : this.currentOffset;
1 !
160
    }
161
    if (!availableNow && autoBatchTime.shouldProbe()) {
1 ✔
162
      resetMetrics();
1 ✔
163
      return nullIfEmpty ? null : currentOffset;
1 ✔
164
    }
165
    List<PulledMessage> pulled = gatherMessages(limits);
1 ✔
166
    if (pulled.isEmpty()) {
1 ✔
167
      resetMetrics();
1 ✔
168
      return nullIfEmpty ? null : currentOffset;
1 ✔
169
    }
170
    long batchId = nextBatchId.getAndIncrement();
1 ✔
171
    PubSubOffset offset = produceNextOffset(pulled, batchId);
1 ✔
172
    this.lastProduced = offset;
1 ✔
173
    updateMetrics(pulled, batchId);
1 ✔
174
    retainGatheredWindow(batchId, pulled);
1 ✔
175
    return offset;
1 ✔
176
  }
177

178
  private PubSubOffset produceNextOffset(List<PulledMessage> pulled, long batchId) {
179
    PubSubOffset offset = new PubSubOffset(batchId);
1 ✔
180
    String batchKey = Long.toString(batchId);
1 ✔
181
    messagesByBatch.put(batchId, pulled);
1 ✔
182
    try {
183
      if (config.ackMode() == AckMode.AFTER_COMMIT) {
1 ✔
184
        ackCoordinator.registerBatch(batchKey, pulled);
1 ✔
185
      } else if (config.gatherMode() == GatherMode.IMMEDIATE) {
1 ✔
186
        ackCoordinator.registerBatch(batchKey, pulled);
1 ✔
187
        ackCoordinator.onPulled(client, batchKey);
1 ✔
188
      }
189
    } catch (RuntimeException e) {
×
190
      ackCoordinator.abort(client, batchKey);
×
191
      messagesByBatch.remove(batchId);
×
192
      throw e;
×
193
    }
1 ✔
194
    return offset;
1 ✔
195
  }
196

197
  private void resetMetrics() {
198
    updateMetrics(null, -1L);
1 ✔
199
  }
1 ✔
200

201
  private void updateMetrics(List<PulledMessage> pulled, long batchId) {
202
    if (pulled == null) {
1 ✔
203
      this.lastGatherMessageCount.set(0);
1 ✔
204
      this.lastGatherPayloadBytes.set(0);
1 ✔
205
      this.lastGatherNewestMessageAgeMs = null;
1 ✔
206
    } else {
207
      final int pulledSize = pulled.size();
1 ✔
208
      final long pulledBytes = PulledMessage.payloadBytes(pulled);
1 ✔
209
      this.lastGatherMessageCount.set(pulledSize);
1 ✔
210
      this.lastGatherPayloadBytes.set(pulledBytes);
1 ✔
211
      long now = System.currentTimeMillis();
1 ✔
212
      this.lastGatherNewestMessageAgeMs = PubSubSourceMetrics.newestMessageAgeMs(pulled, now);
1 ✔
213
      LOG.debug(
1 ✔
214
          "BATCH: micro-batch batchId={} messages={} bytes={}", batchId, pulledSize, pulledBytes);
1 ✔
215
    }
216
  }
1 ✔
217

218
  private List<PulledMessage> gatherMessages(AdmissionLimits gatherLimits) {
219
    AdmissionLimits limits =
220
        availableNow
1 ✔
221
            ? gatherLimits.withDrainUntilIdle()
1 ✔
222
            : gatherLimits.withWaitTime(this.autoBatchTime.receiveWindow(gatherLimits.waitTime()));
1 ✔
223

224
    long gatherStarted = System.nanoTime();
1 ✔
225
    List<PulledMessage> pulled =
226
        limits.drainUntilIdle()
1 ✔
227
            ? gatherUntilIdleOrMax(limits)
1 ✔
228
            : limits.singlePoll()
1 ✔
229
                ? gatherMessagesFromSinglePoll(limits)
1 ✔
230
                : gatherMessagesUntilDeadline(limits);
1 ✔
231
    long gatherFinished = System.nanoTime();
1 ✔
232

233
    this.autoBatchTime.onGatherFinished(gatherFinished - gatherStarted, pulled.size());
1 ✔
234
    return pulled;
1 ✔
235
  }
236

237
  private List<PulledMessage> gatherMessagesFromSinglePoll(AdmissionLimits limits) {
238

239
    List<PulledMessage> messages = new ArrayList<>();
1 ✔
240
    if (gatherAborted()) {
1 !
241
      return messages;
×
242
    }
243
    if (limits.remainingRows(0) <= 0) {
1 !
244
      return messages;
×
245
    }
246
    List<PulledMessage> pulled = poll(limits, messages.size(), RECEIVE_IDLE);
1 ✔
247
    return retainPolled(messages, pulled, limits, 0L).messages;
1 ✔
248
  }
249

250
  private List<PulledMessage> gatherMessagesUntilDeadline(AdmissionLimits limits) {
251

252
    List<PulledMessage> messages = new ArrayList<>();
1 ✔
253
    Duration wait = limits.waitTime();
1 ✔
254
    final long waitNanos = wait == null ? RECEIVE_IDLE.toNanos() : wait.toNanos();
1 !
255
    final long deadlineNanos = System.nanoTime() + waitNanos;
1 ✔
256
    long payloadBytes = 0L;
1 ✔
257
    boolean debounced = false;
1 ✔
258
    try {
259
      while (System.nanoTime() < deadlineNanos) {
1 ✔
260
        if (gatherAborted()) {
1 !
NEW
261
          return abortGatherNackAndRelease(messages, null);
×
262
        }
263
        if (limits.reachedMax(messages.size(), payloadBytes)
1 !
264
            || limits.remainingRows(messages.size()) <= 0) {
1 !
265
          break;
×
266
        }
267
        Duration slice = remainingWait(deadlineNanos);
1 ✔
268
        List<PulledMessage> pulled = poll(limits, messages.size(), slice);
1 ✔
269
        if (gatherAborted()) {
1 ✔
270
          return abortGatherNackAndRelease(messages, pulled);
1 ✔
271
        }
272
        if (!pulled.isEmpty()) {
1 ✔
273
          RetainResult retained = retainPolled(messages, pulled, limits, payloadBytes);
1 ✔
274
          payloadBytes = retained.payloadBytes;
1 ✔
275
          messages = retained.messages;
1 ✔
276
        } else if (config.gatherMode() == GatherMode.BATCH && messages.isEmpty() && !debounced) {
1 !
277
          debounced = true;
1 ✔
278
          List<PulledMessage> pulled2 =
1 ✔
279
              poll(limits, messages.size(), min(RECEIVE_IDLE, remainingWait(deadlineNanos)));
1 ✔
280
          if (gatherAborted()) {
1 !
NEW
281
            return abortGatherNackAndRelease(messages, pulled2);
×
282
          }
283
          if (pulled2.isEmpty()) {
1 !
284
            break;
1 ✔
285
          }
286
          RetainResult retained = retainPolled(messages, pulled2, limits, payloadBytes);
×
287
          payloadBytes = retained.payloadBytes;
×
288
          messages = retained.messages;
×
289
        } else if (config.gatherMode() == GatherMode.BATCH || limits.minRowsMet(messages.size())) {
1 !
290
          break;
×
291
        }
292
        if (limits.reachedMax(messages.size(), payloadBytes)) {
1 ✔
293
          break;
1 ✔
294
        }
295
        if (config.gatherMode() == GatherMode.IMMEDIATE && limits.minRowsMet(messages.size())) {
1 ✔
296
          break;
1 ✔
297
        }
298
      }
1 ✔
299
    } catch (RuntimeException e) {
×
NEW
300
      abortGatherNackAndRelease(messages, null);
×
301
      throw e;
×
302
    }
1 ✔
303
    return messages;
1 ✔
304
  }
305

306
  private List<PulledMessage> poll(AdmissionLimits limits, int already, Duration timeout) {
307
    client.limitNextPoll(limits.remainingRows(already));
1 ✔
308
    return client.poll(timeout);
1 ✔
309
  }
310

311
  private Duration remainingWait(long deadlineNanos) {
312
    return Duration.ofNanos(Math.max(1L, deadlineNanos - System.nanoTime()));
1 ✔
313
  }
314

315
  private static final class RetainResult {
316
    final List<PulledMessage> messages;
317
    final long payloadBytes;
318

319
    RetainResult(List<PulledMessage> messages, long payloadBytes) {
1 ✔
320
      this.messages = messages;
1 ✔
321
      this.payloadBytes = payloadBytes;
1 ✔
322
    }
1 ✔
323
  }
324

325
  private RetainResult retainPolled(
326
      List<PulledMessage> messages,
327
      List<PulledMessage> pulled,
328
      AdmissionLimits limits,
329
      long payloadBytes) {
330

331
    List<PulledMessage> overflow = new ArrayList<>();
1 ✔
332
    for (PulledMessage message : pulled) {
1 ✔
333
      if (limits.reachedMax(messages.size(), payloadBytes)
1 ✔
334
          || limits.remainingRows(messages.size()) <= 0) {
1 !
335
        overflow.add(message);
1 ✔
336
        continue;
1 ✔
337
      }
338
      messages.add(message);
1 ✔
339
      payloadBytes += message.data().length;
1 ✔
340
    }
1 ✔
341
    if (!overflow.isEmpty()) {
1 ✔
342
      if (config.ackMode() == AckMode.AFTER_COMMIT) {
1 !
343
        try {
344
          client.nack(PulledMessage.ackIds(overflow));
1 ✔
345
        } catch (RuntimeException e) {
×
346
          LOG.warn("BATCH: Failed to nack overflow messages: {}", e.getMessage());
×
347
        } finally {
348
          client.releaseMessages(overflow);
1 ✔
349
        }
350
      }
351
    }
352
    List<PulledMessage> kept =
353
        overflow.isEmpty() ? pulled : pulled.subList(0, pulled.size() - overflow.size());
1 ✔
354
    if (config.ackMode() == AckMode.EARLY && !kept.isEmpty()) {
1 !
355
      client.acknowledge(PulledMessage.ackIds(kept));
1 ✔
356
      client.releaseMessages(kept);
1 ✔
357
    }
358
    return new RetainResult(messages, payloadBytes);
1 ✔
359
  }
360

361
  /**
362
   * AvailableNow receive: keep polling until {@code batchCount}/{@code batchSize} or {@link
363
   * #AVAILABLE_NOW_IDLE_POLLS} consecutive empty polls.
364
   */
365
  private List<PulledMessage> gatherUntilIdleOrMax(AdmissionLimits limits) {
366
    List<PulledMessage> messages = new ArrayList<>();
1 ✔
367
    long payloadBytes = 0L;
1 ✔
368
    int emptyStreak = 0;
1 ✔
369
    try {
370
      while (true) {
371
        if (gatherAborted()) {
1 !
NEW
372
          return abortGatherNackAndRelease(messages, null);
×
373
        }
374
        if (limits.reachedMax(messages.size(), payloadBytes)
1 ✔
375
            || limits.remainingRows(messages.size()) <= 0) {
1 !
376
          break;
×
377
        }
378
        List<PulledMessage> pulled = poll(limits, messages.size(), RECEIVE_IDLE);
1 ✔
379
        if (gatherAborted()) {
1 !
NEW
380
          return abortGatherNackAndRelease(messages, pulled);
×
381
        }
382
        if (pulled.isEmpty()) {
1 ✔
383
          emptyStreak++;
1 ✔
384
          if (emptyStreak >= AVAILABLE_NOW_IDLE_POLLS) {
1 ✔
385
            LOG.info(
1 ✔
386
                "BATCH: AvailableNow idle after {} consecutive empty polls; received={}",
387
                emptyStreak,
1 ✔
388
                messages.size());
1 ✔
389
            break;
1 ✔
390
          }
391
          continue;
392
        }
393
        BoundSplit split = splitAtLimit(pulled);
1 ✔
394
        if (!split.overLimit.isEmpty()) {
1 ✔
395
          this.limitReached = true;
1 ✔
396
          nackOverLimit(split.overLimit);
1 ✔
397
          LOG.info(
1 ✔
398
              "BATCH: AvailableNow reached limitTime={}; inRange={} overLimit={}",
399
              config.limitTime().orElse("-"),
1 ✔
400
              split.inRange.size(),
1 ✔
401
              split.overLimit.size());
1 ✔
402

403
          if (!split.inRange.isEmpty()) {
1 ✔
404
            RetainResult retained = retainPolled(messages, split.inRange, limits, payloadBytes);
1 ✔
405
            payloadBytes = retained.payloadBytes;
1 ✔
406
            messages = retained.messages;
1 ✔
407
          }
1 ✔
408
          break;
409
        }
410
        emptyStreak = 0;
1 ✔
411
        RetainResult retained = retainPolled(messages, split.inRange, limits, payloadBytes);
1 ✔
412
        payloadBytes = retained.payloadBytes;
1 ✔
413
        messages = retained.messages;
1 ✔
414
      }
1 ✔
415
    } catch (RuntimeException e) {
×
NEW
416
      abortGatherNackAndRelease(messages, null);
×
417
      throw e;
×
418
    }
1 ✔
419
    return messages;
1 ✔
420
  }
421

422
  private static final class BoundSplit {
423
    final List<PulledMessage> inRange;
424
    final List<PulledMessage> overLimit;
425

426
    BoundSplit(List<PulledMessage> inRange, List<PulledMessage> overLimit) {
1 ✔
427
      this.inRange = inRange;
1 ✔
428
      this.overLimit = overLimit;
1 ✔
429
    }
1 ✔
430
  }
431

432
  private BoundSplit splitAtLimit(List<PulledMessage> pulled) {
433
    Optional<Instant> limit = config.limitTimeAsInstant();
1 ✔
434
    if (limit.isEmpty()) {
1 ✔
435
      return new BoundSplit(pulled, List.of());
1 ✔
436
    }
437
    long boundMillis = limit.get().toEpochMilli();
1 ✔
438
    List<PulledMessage> inRange = new ArrayList<>();
1 ✔
439
    List<PulledMessage> overLimit = new ArrayList<>();
1 ✔
440
    for (PulledMessage message : pulled) {
1 ✔
441
      if (message.publishTimeMillis() >= boundMillis) {
1 ✔
442
        overLimit.add(message);
1 ✔
443
      } else {
444
        inRange.add(message);
1 ✔
445
      }
446
    }
1 ✔
447
    return new BoundSplit(inRange, overLimit);
1 ✔
448
  }
449

450
  private void nackOverLimit(List<PulledMessage> overLimit) {
451
    if (overLimit.isEmpty()) {
1 !
452
      return;
×
453
    }
454
    try {
455
      client.nack(PulledMessage.ackIds(overLimit));
1 ✔
456
    } catch (RuntimeException e) {
×
457
      LOG.warn(
×
458
          "BATCH: Failed to nack {} over-limit messages: {}", overLimit.size(), e.getMessage());
×
459
    } finally {
460
      client.releaseMessages(overLimit);
1 ✔
461
    }
462
  }
1 ✔
463

464
  private boolean gatherAborted() {
465
    return cancelled || Thread.currentThread().isInterrupted();
1 !
466
  }
467

468
  private List<PulledMessage> abortGatherNackAndRelease(
469
      List<PulledMessage> messages, List<PulledMessage> pulled) {
470

471
    if (pulled != null) {
1 !
472
      messages.addAll(pulled);
1 ✔
473
    }
474
    if (config.ackMode() == AckMode.AFTER_COMMIT && !messages.isEmpty()) {
1 !
475
      try {
476
        client.nack(PulledMessage.ackIds(messages));
1 ✔
477
      } catch (RuntimeException nackError) {
×
478
        LOG.warn("BATCH: Failed to nack messages after gather abort: {}", nackError.getMessage());
×
479
      } finally {
480
        client.releaseMessages(messages);
1 ✔
481
      }
482
    }
483
    return List.of();
1 ✔
484
  }
485

486
  private static Duration min(Duration left, Duration right) {
487
    return left.compareTo(right) <= 0 ? left : right;
1 !
488
  }
489

490
  /** True when Spark already persisted {@code produced} as the start of the next range. */
491
  static boolean startConsumed(Offset startOffset, PubSubOffset produced) {
492
    if (startOffset == null || produced == null) {
1 !
493
      return false;
1 ✔
494
    }
495
    final long startBatchId =
496
        startOffset instanceof PubSubOffset
1 !
497
            ? ((PubSubOffset) startOffset).batchId()
1 ✔
498
            : PubSubOffset.fromJson(startOffset.json()).batchId();
1 ✔
499
    return startBatchId >= produced.batchId();
1 ✔
500
  }
501

502
  @Override
503
  public Map<String, String> metrics(final Optional<Offset> latestConsumedOffset) {
504
    PubSubOffset produced = lastProduced;
1 ✔
505
    final Long producedBatchId = produced == null ? null : produced.batchId();
1 ✔
506
    final long retryTotal = client.retryAttempts();
1 ✔
507
    final long reportedTotal = lastReportedRetryAttempts.getAndSet(retryTotal);
1 ✔
508
    final long retryThisBatch =
1 ✔
509
        PubSubSourceMetrics.retryAttemptsThisBatch(retryTotal, reportedTotal);
1 ✔
510
    return PubSubSourceMetrics.snapshot(
1 ✔
511
        lastGatherMessageCount.get(),
1 ✔
512
        lastGatherPayloadBytes.get(),
1 ✔
513
        lastGatherNewestMessageAgeMs,
514
        client.outstandingBytes(),
1 ✔
515
        producedBatchId,
516
        latestConsumedOffset,
517
        retryThisBatch,
518
        retryTotal);
519
  }
520

521
  @Override
522
  public InputPartition[] planInputPartitions(Offset start, Offset end) {
523
    PubSubOffset endOffset = (PubSubOffset) end;
1 ✔
524
    List<PulledMessage> messages = messagesByBatch.getOrDefault(endOffset.batchId(), List.of());
1 ✔
525
    if (messages.isEmpty()) {
1 !
NEW
526
      int gathered = lastGatherMessageCount.get();
×
NEW
527
      if (gathered > 0) {
×
UNCOV
528
        LOG.warn(
×
529
            "BATCH: planInputPartitions batchId={} has 0 messages but lastGatherMessageCount={}",
530
            endOffset.batchId(),
×
NEW
531
            gathered);
×
532
      }
533
      return new InputPartition[] {new PubSubInputPartition(messages)};
×
534
    }
535
    int parts = Math.min(numPartitions, messages.size());
1 ✔
536
    List<List<PulledMessage>> slices = new ArrayList<>(parts);
1 ✔
537
    for (int i = 0; i < parts; i++) {
1 ✔
538
      slices.add(new ArrayList<>());
1 ✔
539
    }
540
    for (int i = 0; i < messages.size(); i++) {
1 ✔
541
      slices.get(i % parts).add(messages.get(i));
1 ✔
542
    }
543
    InputPartition[] partitions = new InputPartition[parts];
1 ✔
544
    for (int i = 0; i < parts; i++) {
1 ✔
545
      partitions[i] = new PubSubInputPartition(slices.get(i));
1 ✔
546
    }
547
    return partitions;
1 ✔
548
  }
549

550
  @Override
551
  public PartitionReaderFactory createReaderFactory() {
552
    return new PubSubPartitionReaderFactory(readSchema);
×
553
  }
554

555
  @Override
556
  public Offset initialOffset() {
557
    return currentOffset;
1 ✔
558
  }
559

560
  @Override
561
  public Offset deserializeOffset(String json) {
562
    PubSubOffset offset = PubSubOffset.fromJson(json);
×
563
    nextBatchId.updateAndGet(v -> Math.max(v, offset.batchId() + 1));
×
564
    currentOffset = offset;
×
565
    return offset;
×
566
  }
567

568
  @Override
569
  public void commit(Offset end) {
570
    PubSubOffset endOffset = (PubSubOffset) end;
1 ✔
571
    String batchKey = Long.toString(endOffset.batchId());
1 ✔
572
    if (config.ackMode() == AckMode.AFTER_COMMIT) {
1 ✔
573
      List<PulledMessage> messages = messagesByBatch.getOrDefault(endOffset.batchId(), List.of());
1 ✔
574
      List<String> ackIds = PulledMessage.ackIds(messages);
1 ✔
575
      try {
576
        if (!ackIds.isEmpty()) {
1 !
577
          try {
578
            client.acknowledge(ackIds);
1 ✔
579
            LOG.debug(
1 ✔
580
                "BATCH: Committed (acked) {} messages for batch {}",
581
                ackIds.size(),
1 ✔
582
                endOffset.batchId());
1 ✔
583
          } catch (RuntimeException e) {
×
584
            try {
585
              client.nack(ackIds);
×
586
            } catch (RuntimeException nackError) {
×
587
              LOG.warn(
×
588
                  "BATCH: Failed to nack {} messages after ack failure for batch {}",
589
                  ackIds.size(),
×
590
                  endOffset.batchId(),
×
591
                  nackError.getMessage());
×
592
            }
×
593
            throw e;
×
594
          }
1 ✔
595
        }
596
      } finally {
597
        if (!ackIds.isEmpty()) {
1 !
598
          client.releaseMessages(messages);
1 ✔
599
        }
600
        ackCoordinator.discard(batchKey);
1 ✔
601
        finishBatch(endOffset.batchId());
1 ✔
602
      }
603
      currentOffset = endOffset;
1 ✔
604
    } else {
1 ✔
605
      ackCoordinator.commit(client, batchKey);
1 ✔
606
      finishBatch(endOffset.batchId());
1 ✔
607
      currentOffset = endOffset;
1 ✔
608
    }
609
    LOG.debug("BATCH: Committed offset batchId={}", endOffset.batchId());
1 ✔
610
    autoBatchTime.onCommit();
1 ✔
611
  }
1 ✔
612

613
  @Override
614
  public void stop() {
615
    cancelled = true;
1 ✔
616
    try {
617
      int uncommitted = 0;
1 ✔
618
      if (lastProduced != null && config.ackMode() == AckMode.AFTER_COMMIT) {
1 !
619
        String batchKey = Long.toString(lastProduced.batchId());
1 ✔
620
        List<PulledMessage> messages = ackCoordinator.messagesForBatch(batchKey);
1 ✔
621
        if (!messages.isEmpty()) {
1 !
622
          uncommitted = messages.size();
1 ✔
623
          ackCoordinator.abort(client, batchKey);
1 ✔
624
          messagesByBatch.remove(lastProduced.batchId());
1 ✔
625
        }
626
      }
627
      logStop(uncommitted);
1 ✔
628
    } finally {
629
      ackCoordinator.clear();
1 ✔
630
      messagesByBatch.clear();
1 ✔
631
      client.resetOutstandingBytes();
1 ✔
632
      client.close();
1 ✔
633
    }
634
  }
1 ✔
635

636
  private void retainGatheredWindow(long batchId, List<PulledMessage> pulled) {
637
    PublishTimeWindow window = PublishTimeWindow.of(pulled);
1 ✔
638
    lastGatheredBatchId = batchId;
1 ✔
639
    lastGatheredWindow = window;
1 ✔
640
    if (firstBatchLogged || window == null) {
1 !
641
      return;
1 ✔
642
    }
643
    this.firstBatchLogged = true;
1 ✔
644
    LOG.info(window.formatStats(batchId, "BATCH: First batch:"));
1 ✔
645
  }
1 ✔
646

647
  private void logStop(int uncommitted) {
648
    PublishTimeWindow window = lastGatheredWindow;
1 ✔
649
    if (window == null) {
1 !
650
      LOG.info("BATCH: Stopping stream; no messages gathered");
×
651
      return;
×
652
    }
653
    final long batchId = this.lastGatheredBatchId == null ? -1L : this.lastGatheredBatchId;
1 !
654
    final String prefix =
655
        uncommitted <= 0
1 !
656
            ? "BATCH: Stopping stream; 0 uncommitted messages found; Last batch is:"
×
657
            : String.format(
1 ✔
658
                "BATCH: Stopping stream; %d uncommitted messages will redeliver if not committed;"
659
                    + " Last batch is:",
660
                uncommitted);
1 ✔
661
    LOG.info(window.formatStats(batchId, prefix));
1 ✔
662
  }
1 ✔
663

664
  PublishTimeWindow lastGatheredWindow() {
665
    return lastGatheredWindow;
1 ✔
666
  }
667

668
  boolean firstBatchLogged() {
669
    return firstBatchLogged;
1 ✔
670
  }
671

672
  Long lastGatheredBatchId() {
673
    return lastGatheredBatchId;
1 ✔
674
  }
675

676
  boolean availableNow() {
677
    return availableNow;
1 ✔
678
  }
679

680
  boolean limitReached() {
681
    return limitReached;
1 ✔
682
  }
683

684
  @Override
685
  public String toString() {
686
    return PubSubConfig.SHORT_NAME + ":" + config.subscriptionPath();
1 ✔
687
  }
688

689
  private void finishBatch(long batchId) {
690
    messagesByBatch.remove(batchId);
1 ✔
691
    if (lastProduced != null && lastProduced.batchId() == batchId) {
1 !
692
      lastProduced = null;
1 ✔
693
    }
694
  }
1 ✔
695
}
STATUS · Troubleshooting · Open an Issue · Sales · Support · CAREERS · ENTERPRISE · START FREE TRIAL · SCHEDULE DEMO
ANNOUNCEMENTS · TWITTER · TOS & SLA · Supported CI Services · What's a CI service? · Automated Testing

© 2026 Coveralls, Inc