juarezr / spark-streaming-google-pubsub / 57

Committed 01 Oct 2026 07:02AM UTC coverage: 66.187% (-2.3%) from 68.457%

Build # 57

Build Type

push

github

Committed by

web-flow

Commit Message

fix: update AutoBatchTime logic for receiveTime estimation with write-first approach (#20)

- fix: update AutoBatchTime logic for receiveTime estimation with
write-first approach
- feat: add processingTime configuration to PubSubConfig to hint batch
interval on AutoBatchTime
- fix: remove unneded cap on writeBudget for AutoBatchTime

Coverage Stats

689 of 1169 branches covered (58.94%)

Branch coverage included in aggregate %.

232 of 354 new or added lines in 12 files covered. (65.54%)

15 existing lines in 1 file now uncovered.

1517 of 2164 relevant lines covered (70.1%)

0.7 hits per line

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

39.39

/src/main/java/io/github/juarezr/spark/pubsub/structured/PubSubClient.java

package io.github.juarezr.spark.pubsub.structured;

import com.google.api.core.ApiService;
import com.google.api.gax.batching.FlowControlSettings;
import com.google.api.gax.batching.FlowController;
import com.google.api.gax.core.FixedCredentialsProvider;
import com.google.cloud.pubsub.v1.AckReplyConsumer;
import com.google.cloud.pubsub.v1.MessageReceiver;
import com.google.cloud.pubsub.v1.Subscriber;
import com.google.cloud.pubsub.v1.SubscriptionAdminClient;
import com.google.cloud.pubsub.v1.SubscriptionAdminSettings;
import com.google.protobuf.Timestamp;
import com.google.pubsub.v1.PubsubMessage;
import com.google.pubsub.v1.SeekRequest;
import io.github.juarezr.spark.pubsub.config.PubSubConfig;
import io.github.juarezr.spark.pubsub.config.SeekMode;
import java.io.Closeable;
import java.io.IOException;
import java.io.Serializable;
import java.time.Duration;
import java.time.Instant;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Warm {@link Subscriber} (StreamingPull) plus a bounded queue. The receiver never acks; Spark
 * commit calls {@link #acknowledge}.
 */
final class PubSubClient implements Closeable, Serializable {

  private static final long serialVersionUID = -3861530486L;
  private static final Logger LOG = LoggerFactory.getLogger(PubSubClient.class);

  /** Pub/Sub Acknowledge / ModifyAckDeadline cap. */
  static final int ACK_CHUNK = 1000;

  private final RetryPolicy retryPolicy;
  private transient PubSubEmulator emulator;

  private final PubSubConfig config;
  private final PubSubCredentialsProvider credentialsProvider;
  private transient Subscriber subscriber;
  private transient LinkedBlockingQueue<HeldMessage> queue;
  private transient ConcurrentHashMap<String, AckReplyConsumer> consumers;
  private transient AtomicLong outstandingBytes;
  private transient AtomicLong queuedBytes;
  private transient AtomicBoolean stopped;
  private transient boolean seekApplied;

  /** Caps the next {@link #poll(Duration)} so overflow stays queued for the following batch. */
  private transient int nextPollMax = Integer.MAX_VALUE;

  private final AtomicLong lastSeenNewestPublishMillis = new AtomicLong(Long.MIN_VALUE);

  final class PubSubMessageReceiver implements MessageReceiver {

    @Override
    public void receiveMessage(PubsubMessage message, AckReplyConsumer consumer) {
      if (stopped != null && stopped.get()) {
        consumer.nack();
        return;
      }
      PulledMessage pulled = toPulled(message);
      long cap = config.batchSize() > 0 ? config.batchSize() : PubSubConfig.DEFAULT_BATCH_SIZE;
      while (queuedBytes != null && queuedBytes.get() >= cap && stopped != null && !stopped.get()) {
        try {
          TimeUnit.MILLISECONDS.sleep(10);
        } catch (InterruptedException e) {
          Thread.currentThread().interrupt();
          consumer.nack();
          return;
        }
      }
      if (stopped != null && stopped.get()) {
        consumer.nack();
        return;
      }
      consumers.put(pulled.ackId(), consumer);
      queuedBytes.addAndGet(pulled.data().length);
      outstandingBytes.addAndGet(pulled.data().length);
      try {
        queue.put(new HeldMessage(pulled, consumer));
      } catch (InterruptedException e) {
        Thread.currentThread().interrupt();
        consumers.remove(pulled.ackId());
        queuedBytes.addAndGet(-pulled.data().length);
        outstandingBytes.addAndGet(-pulled.data().length);
        consumer.nack();
      }
    }

    private PulledMessage toPulled(PubsubMessage message) {
      byte[] data = message.getData().toByteArray();
      long publishMillis =
          message.getPublishTime().getSeconds() * 1000L
              + message.getPublishTime().getNanos() / 1_000_000L;
      String messageId = message.getMessageId();
      return new PulledMessage(
          messageId,
          data,
          message.getAttributesMap(),
          publishMillis,
          message.getOrderingKey(),
          messageId);
    }
  }

  static final class HeldMessage {
    final PulledMessage message;
    final AckReplyConsumer consumer;

    HeldMessage(PulledMessage message, AckReplyConsumer consumer) {
      this.message = message;
      this.consumer = consumer;
    }
  }

  PubSubClient(PubSubConfig config) {
    this(
        config,
        new PubSubCredentialsProvider(config.credentialsFile().orElse(null)),
        RetryPolicy.defaults(config.effectiveMaxRetryTime()));
  }

  PubSubClient(
      PubSubConfig config, PubSubCredentialsProvider credentialsProvider, RetryPolicy retryPolicy) {
    this.config = config;
    this.credentialsProvider = credentialsProvider;
    this.retryPolicy = retryPolicy;
    this.retryPolicy.lastPublishTime(this.lastSeenNewestPublishMillis::get);
  }

  synchronized void start() throws IOException {
    if (this.subscriber != null) {
      return;
    }
    this.outstandingBytes = new AtomicLong(0);
    this.queuedBytes = new AtomicLong(0);
    this.queue = new LinkedBlockingQueue<>();
    this.consumers = new ConcurrentHashMap<>();
    this.stopped = new AtomicBoolean(false);

    this.emulator = this.config.emulatorHost().map(PubSubEmulator::new).orElse(null);
    try {
      applySeekIfNeeded();
      Subscriber.Builder builder =
          Subscriber.newBuilder(this.config.subscriptionPath(), new PubSubMessageReceiver());
      builder.setFlowControlSettings(flowControl());
      org.threeten.bp.Duration ackExtend =
          org.threeten.bp.Duration.ofSeconds(
              Math.max(60L, config.effectiveAckDeadline(null).getSeconds()));
      builder.setMaxAckExtensionPeriod(org.threeten.bp.Duration.ofMinutes(60));
      builder.setMaxDurationPerAckExtension(ackExtend);
      if (this.emulator != null) {
        this.emulator.configureSubscriber(builder);
      } else {
        builder.setCredentialsProvider(
            FixedCredentialsProvider.create(this.credentialsProvider.getCredentials()));
      }
      this.subscriber = builder.build();
      this.subscriber.startAsync().awaitRunning();
    } catch (Exception e) {
      this.close();
      if (e instanceof IOException) {
        throw (IOException) e;
      }
      throw new IOException("Unable to start Pub/Sub subscriber", e);
    }
  }

  private FlowControlSettings flowControl() {
    long bytes = config.batchSize() > 0 ? config.batchSize() : PubSubConfig.DEFAULT_BATCH_SIZE;
    long outstandingBytes = bytes > Long.MAX_VALUE / 2 ? Long.MAX_VALUE : bytes * 2;
    FlowControlSettings.Builder flow =
        FlowControlSettings.newBuilder()
            .setMaxOutstandingRequestBytes(outstandingBytes)
            .setLimitExceededBehavior(FlowController.LimitExceededBehavior.Block);
    if (config.batchCount() > 0) {
      long count = config.batchCount();
      long elements = count > Long.MAX_VALUE / 2 ? Long.MAX_VALUE : count * 2;
      flow.setMaxOutstandingElementCount(elements);
    }
    return flow.build();
  }

  private void applySeekIfNeeded() throws IOException {
    if (seekApplied || this.config.seekMode() == SeekMode.NONE) {
      return;
    }
    final SeekRequest.Builder seek = getSeekRequestBuilder();
    final SubscriptionAdminSettings.Builder builder = SubscriptionAdminSettings.newBuilder();
    try (PubSubEmulator adminEmulator =
        this.config.emulatorHost().map(PubSubEmulator::new).orElse(null)) {
      if (adminEmulator != null) {
        adminEmulator.configureSubscriptionAdmin(builder);
      } else {
        builder.setCredentialsProvider(
            FixedCredentialsProvider.create(this.credentialsProvider.getCredentials()));
      }
      try (SubscriptionAdminClient admin = SubscriptionAdminClient.create(builder.build())) {
        LOG.warn(
            "CONNECTION: Applying seek={} on subscription {} (explicit rewind requested)",
            this.config.seekMode(),
            this.config.subscriptionPath());
        admin.seek(seek.build());
      }
    }
    seekApplied = true;
  }

  private SeekRequest.Builder getSeekRequestBuilder() {
    final SeekRequest.Builder seek =
        SeekRequest.newBuilder().setSubscription(this.config.subscriptionPath());

    switch (this.config.seekMode()) {
      case BEGINNING:
        seek.setTime(Timestamp.newBuilder().setSeconds(0).setNanos(0).build());
        break;
      case TIMESTAMP:
        final Instant seekInstant = this.config.seekTimeAsInstant();
        final Timestamp seekTime =
            Timestamp.newBuilder()
                .setSeconds(seekInstant.getEpochSecond())
                .setNanos(seekInstant.getNano())
                .build();
        seek.setTime(seekTime);
        break;
      case SNAPSHOT:
        seek.setSnapshot(this.config.seekSnapshot().orElseThrow());
        break;
      case NONE:
      default:
    }
    return seek;
  }

  private void ensureStarted() {
    if (this.subscriber == null) {
      try {
        start();
      } catch (IOException e) {
        throw new IllegalStateException("CONNECTION: Failed to start Pub/Sub client", e);
      }
    }
  }

  @Override
  public synchronized void close() {
    if (stopped != null) {
      stopped.set(true);
    }
    if (this.subscriber != null) {
      try {
        ApiService service = this.subscriber.stopAsync();
        service.awaitTerminated(5, TimeUnit.SECONDS);
      } catch (Exception e) {
        LOG.warn("CONNECTION: Error shutting down subscriber", e);
      } finally {
        this.subscriber = null;
      }
    }
    nackQueued();
    if (this.emulator != null) {
      this.emulator.close();
      this.emulator = null;
    }
  }

  private void nackQueued() {
    if (queue == null) {
      return;
    }
    HeldMessage held;
    while ((held = queue.poll()) != null) {
      try {
        held.consumer.nack();
      } catch (RuntimeException e) {
        LOG.warn("CONNECTION: Failed to nack queued message on close: {}", e.getMessage());
      }
    }
    if (consumers != null) {
      consumers.clear();
    }
    if (queuedBytes != null) {
      queuedBytes.set(0);
    }
  }

  /**
   * Upper bound for the next {@link #poll(Duration)}. Messages beyond the cap stay in the queue.
   * {@code maxMessages <= 0} means no cap.
   */
  void limitNextPoll(int maxMessages) {
    this.nextPollMax = maxMessages <= 0 ? Integer.MAX_VALUE : maxMessages;
  }

  /**
   * Wait up to {@code timeout} for messages, up to {@link #limitNextPoll(int)}. When capped, keeps
   * waiting within the same timeout until the cap is reached or the deadline passes. When uncapped,
   * waits only for the first message, then drains the queue without blocking. Does not ack.
   */
  List<PulledMessage> poll(Duration timeout) {
    ensureStarted();
    int max = this.nextPollMax;
    this.nextPollMax = Integer.MAX_VALUE;
    List<PulledMessage> messages = new ArrayList<>();
    if (max <= 0) {
      return messages;
    }
    long deadlineNanos = System.nanoTime() + Math.max(1L, timeout.toNanos());
    try {
      long waitNanos = deadlineNanos - System.nanoTime();
      if (waitNanos <= 0L) {
        return messages;
      }
      HeldMessage first = queue.poll(waitNanos, TimeUnit.NANOSECONDS);
      if (first == null) {
        return messages;
      }
      take(first, messages);
      if (max == Integer.MAX_VALUE) {
        HeldMessage next;
        while ((next = queue.poll()) != null) {
          take(next, messages);
        }
      } else {
        while (messages.size() < max) {
          waitNanos = deadlineNanos - System.nanoTime();
          if (waitNanos <= 0L) {
            break;
          }
          HeldMessage next = queue.poll(waitNanos, TimeUnit.NANOSECONDS);
          if (next == null) {
            break;
          }
          take(next, messages);
        }
      }
    } catch (InterruptedException e) {
      Thread.currentThread().interrupt();
    }
    final PublishTimeWindow window = PublishTimeWindow.of(messages);
    if (window != null) {
      this.lastSeenNewestPublishMillis.set(window.newestMillis());
    }
    return messages;
  }

  private void take(HeldMessage held, List<PulledMessage> messages) {
    messages.add(held.message);
    if (queuedBytes != null) {
      queuedBytes.addAndGet(-held.message.data().length);
    }
  }

  void acknowledge(List<String> ackIds) {
    reply(ackIds, true);
  }

  void nack(List<String> ackIds) {
    reply(ackIds, false);
  }

  /** Subscriber renews leases; kept for {@link AckLeaseWatchdog}. */
  void extendAckDeadline(List<String> ackIds, int deadlineSeconds) {}

  private void reply(List<String> ackIds, boolean ack) {
    if (ackIds == null || ackIds.isEmpty()) {
      return;
    }
    if (consumers == null) {
      return;
    }
    for (String id : ackIds) {
      AckReplyConsumer consumer = consumers.remove(id);
      if (consumer == null) {
        continue;
      }
      try {
        if (ack) {
          consumer.ack();
        } else {
          consumer.nack();
        }
      } catch (RuntimeException e) {
        LOG.warn("ACK: Failed to {} {}: {}", ack ? "ack" : "nack", id, e.getMessage());
      }
    }
  }

  void releaseMessages(List<PulledMessage> messages) {
    if (this.outstandingBytes != null) {
      long bytes = PulledMessage.payloadBytes(messages);
      if (bytes > 0) {
        this.outstandingBytes.addAndGet(-bytes);
      }
    }
  }

  void resetOutstandingBytes() {
    if (this.outstandingBytes != null) {
      this.outstandingBytes.set(0);
    }
  }

  long outstandingBytes() {
    return this.outstandingBytes == null ? 0L : this.outstandingBytes.get();
  }

  long retryAttempts() {
    return retryPolicy.retryAttempts();
  }

  long lastSeenNewestPublishMillis() {
    return lastSeenNewestPublishMillis.get();
  }

  public static String asString(final Duration duration) {
    return duration.toString().substring(2).replaceAll("(\\d[HMS])(?!$)", "$1").toLowerCase();
  }
}

1	package io.github.juarezr.spark.pubsub.structured;
2
3	import com.google.api.core.ApiService;
4	import com.google.api.gax.batching.FlowControlSettings;
5	import com.google.api.gax.batching.FlowController;
6	import com.google.api.gax.core.FixedCredentialsProvider;
7	import com.google.cloud.pubsub.v1.AckReplyConsumer;
8	import com.google.cloud.pubsub.v1.MessageReceiver;
9	import com.google.cloud.pubsub.v1.Subscriber;
10	import com.google.cloud.pubsub.v1.SubscriptionAdminClient;
11	import com.google.cloud.pubsub.v1.SubscriptionAdminSettings;
12	import com.google.protobuf.Timestamp;
13	import com.google.pubsub.v1.PubsubMessage;
14	import com.google.pubsub.v1.SeekRequest;
15	import io.github.juarezr.spark.pubsub.config.PubSubConfig;
16	import io.github.juarezr.spark.pubsub.config.SeekMode;
17	import java.io.Closeable;
18	import java.io.IOException;
19	import java.io.Serializable;
20	import java.time.Duration;
21	import java.time.Instant;
22	import java.util.ArrayList;
23	import java.util.List;
24	import java.util.concurrent.ConcurrentHashMap;
25	import java.util.concurrent.LinkedBlockingQueue;
26	import java.util.concurrent.TimeUnit;
27	import java.util.concurrent.atomic.AtomicBoolean;
28	import java.util.concurrent.atomic.AtomicLong;
29	import org.slf4j.Logger;
30	import org.slf4j.LoggerFactory;
31
32	/**
33	* Warm {@link Subscriber} (StreamingPull) plus a bounded queue. The receiver never acks; Spark
34	* commit calls {@link #acknowledge}.
35	*/
36	final class PubSubClient implements Closeable, Serializable {
37
38	private static final long serialVersionUID = -3861530486L;
39	private static final Logger LOG = LoggerFactory.getLogger(PubSubClient.class);	1 ✔
40
41	/** Pub/Sub Acknowledge / ModifyAckDeadline cap. */
42	static final int ACK_CHUNK = 1000;
43
44	private final RetryPolicy retryPolicy;
45	private transient PubSubEmulator emulator;
46
47	private final PubSubConfig config;
48	private final PubSubCredentialsProvider credentialsProvider;
49	private transient Subscriber subscriber;
50	private transient LinkedBlockingQueue<HeldMessage> queue;
51	private transient ConcurrentHashMap<String, AckReplyConsumer> consumers;
52	private transient AtomicLong outstandingBytes;
53	private transient AtomicLong queuedBytes;
54	private transient AtomicBoolean stopped;
55	private transient boolean seekApplied;
56
57	/** Caps the next {@link #poll(Duration)} so overflow stays queued for the following batch. */
58	private transient int nextPollMax = Integer.MAX_VALUE;	1 ✔
59
60	private final AtomicLong lastSeenNewestPublishMillis = new AtomicLong(Long.MIN_VALUE);	1 ✔
61
62	final class PubSubMessageReceiver implements MessageReceiver {	×
63
64	@Override
65	public void receiveMessage(PubsubMessage message, AckReplyConsumer consumer) {
66	if (stopped != null && stopped.get()) {	×
67	consumer.nack();	×
68	return;	×
69	}
70	PulledMessage pulled = toPulled(message);	×
71	long cap = config.batchSize() > 0 ? config.batchSize() : PubSubConfig.DEFAULT_BATCH_SIZE;	×
72	while (queuedBytes != null && queuedBytes.get() >= cap && stopped != null && !stopped.get()) {	×
73	try {
74	TimeUnit.MILLISECONDS.sleep(10);	×
75	} catch (InterruptedException e) {	×
76	Thread.currentThread().interrupt();	×
77	consumer.nack();	×
78	return;	×
79	}	×
80	}
81	if (stopped != null && stopped.get()) {	×
82	consumer.nack();	×
83	return;	×
84	}
85	consumers.put(pulled.ackId(), consumer);	×
86	queuedBytes.addAndGet(pulled.data().length);	×
87	outstandingBytes.addAndGet(pulled.data().length);	×
88	try {
89	queue.put(new HeldMessage(pulled, consumer));	×
90	} catch (InterruptedException e) {	×
91	Thread.currentThread().interrupt();	×
92	consumers.remove(pulled.ackId());	×
93	queuedBytes.addAndGet(-pulled.data().length);	×
94	outstandingBytes.addAndGet(-pulled.data().length);	×
95	consumer.nack();	×
96	}	×
97	}	×
98
99	private PulledMessage toPulled(PubsubMessage message) {
100	byte[] data = message.getData().toByteArray();	×
101	long publishMillis =	×
102	message.getPublishTime().getSeconds() * 1000L	×
103	+ message.getPublishTime().getNanos() / 1_000_000L;	×
104	String messageId = message.getMessageId();	×
105	return new PulledMessage(	×
106	messageId,
107	data,
108	message.getAttributesMap(),	×
109	publishMillis,
110	message.getOrderingKey(),	×
111	messageId);
112	}
113	}
114
115	static final class HeldMessage {
116	final PulledMessage message;
117	final AckReplyConsumer consumer;
118
119	HeldMessage(PulledMessage message, AckReplyConsumer consumer) {	1 ✔
120	this.message = message;	1 ✔
121	this.consumer = consumer;	1 ✔
122	}	1 ✔
123	}
124
125	PubSubClient(PubSubConfig config) {
126	this(	1 ✔
127	config,
128	new PubSubCredentialsProvider(config.credentialsFile().orElse(null)),	1 ✔
129	RetryPolicy.defaults(config.effectiveMaxRetryTime()));	1 ✔
130	}	1 ✔
131
132	PubSubClient(
133	PubSubConfig config, PubSubCredentialsProvider credentialsProvider, RetryPolicy retryPolicy) {	1 ✔
134	this.config = config;	1 ✔
135	this.credentialsProvider = credentialsProvider;	1 ✔
136	this.retryPolicy = retryPolicy;	1 ✔
137	this.retryPolicy.lastPublishTime(this.lastSeenNewestPublishMillis::get);	1 ✔
138	}	1 ✔
139
140	synchronized void start() throws IOException {
141	if (this.subscriber != null) {	×
142	return;	×
143	}
144	this.outstandingBytes = new AtomicLong(0);	×
145	this.queuedBytes = new AtomicLong(0);	×
146	this.queue = new LinkedBlockingQueue<>();	×
147	this.consumers = new ConcurrentHashMap<>();	×
148	this.stopped = new AtomicBoolean(false);	×
149
150	this.emulator = this.config.emulatorHost().map(PubSubEmulator::new).orElse(null);	×
151	try {
152	applySeekIfNeeded();	×
153	Subscriber.Builder builder =	×
154	Subscriber.newBuilder(this.config.subscriptionPath(), new PubSubMessageReceiver());	×
155	builder.setFlowControlSettings(flowControl());	×
156	org.threeten.bp.Duration ackExtend =	×
157	org.threeten.bp.Duration.ofSeconds(	×
158	Math.max(60L, config.effectiveAckDeadline(null).getSeconds()));	×
159	builder.setMaxAckExtensionPeriod(org.threeten.bp.Duration.ofMinutes(60));	×
160	builder.setMaxDurationPerAckExtension(ackExtend);	×
161	if (this.emulator != null) {	×
162	this.emulator.configureSubscriber(builder);	×
163	} else {
164	builder.setCredentialsProvider(	×
165	FixedCredentialsProvider.create(this.credentialsProvider.getCredentials()));	×
166	}
167	this.subscriber = builder.build();	×
168	this.subscriber.startAsync().awaitRunning();	×
169	} catch (Exception e) {	×
170	this.close();	×
171	if (e instanceof IOException) {	×
172	throw (IOException) e;	×
173	}
174	throw new IOException("Unable to start Pub/Sub subscriber", e);	×
175	}	×
176	}	×
177
178	private FlowControlSettings flowControl() {
179	long bytes = config.batchSize() > 0 ? config.batchSize() : PubSubConfig.DEFAULT_BATCH_SIZE;	×
180	long outstandingBytes = bytes > Long.MAX_VALUE / 2 ? Long.MAX_VALUE : bytes * 2;	×
181	FlowControlSettings.Builder flow =
182	FlowControlSettings.newBuilder()	×
183	.setMaxOutstandingRequestBytes(outstandingBytes)	×
184	.setLimitExceededBehavior(FlowController.LimitExceededBehavior.Block);	×
185	if (config.batchCount() > 0) {	×
186	long count = config.batchCount();	×
187	long elements = count > Long.MAX_VALUE / 2 ? Long.MAX_VALUE : count * 2;	×
188	flow.setMaxOutstandingElementCount(elements);	×
189	}
190	return flow.build();	×
191	}
192
193	private void applySeekIfNeeded() throws IOException {
194	if (seekApplied \|\| this.config.seekMode() == SeekMode.NONE) {	×
195	return;	×
196	}
197	final SeekRequest.Builder seek = getSeekRequestBuilder();	×
198	final SubscriptionAdminSettings.Builder builder = SubscriptionAdminSettings.newBuilder();	×
199	try (PubSubEmulator adminEmulator =	×
200	this.config.emulatorHost().map(PubSubEmulator::new).orElse(null)) {	×
201	if (adminEmulator != null) {	×
202	adminEmulator.configureSubscriptionAdmin(builder);	×
203	} else {
204	builder.setCredentialsProvider(	×
205	FixedCredentialsProvider.create(this.credentialsProvider.getCredentials()));	×
206	}
207	try (SubscriptionAdminClient admin = SubscriptionAdminClient.create(builder.build())) {	×
208	LOG.warn(	×
209	"CONNECTION: Applying seek={} on subscription {} (explicit rewind requested)",
210	this.config.seekMode(),	×
211	this.config.subscriptionPath());	×
212	admin.seek(seek.build());	×
213	}
214	}
215	seekApplied = true;	×
216	}	×
217
218	private SeekRequest.Builder getSeekRequestBuilder() {
219	final SeekRequest.Builder seek =
220	SeekRequest.newBuilder().setSubscription(this.config.subscriptionPath());	×
221
222	switch (this.config.seekMode()) {	×
223	case BEGINNING:
224	seek.setTime(Timestamp.newBuilder().setSeconds(0).setNanos(0).build());	×
225	break;	×
226	case TIMESTAMP:
NEW 227	final Instant seekInstant = this.config.seekTimeAsInstant();	×
228	final Timestamp seekTime =
229	Timestamp.newBuilder()	×
NEW 230	.setSeconds(seekInstant.getEpochSecond())	×
NEW 231	.setNanos(seekInstant.getNano())	×
232	.build();	×
233	seek.setTime(seekTime);	×
234	break;	×
235	case SNAPSHOT:
236	seek.setSnapshot(this.config.seekSnapshot().orElseThrow());	×
237	break;	×
238	case NONE:
239	default:
240	}
241	return seek;	×
242	}
243
244	private void ensureStarted() {
245	if (this.subscriber == null) {	1 !
246	try {
247	start();	×
248	} catch (IOException e) {	×
249	throw new IllegalStateException("CONNECTION: Failed to start Pub/Sub client", e);	×
250	}	×
251	}
252	}	1 ✔
253
254	@Override
255	public synchronized void close() {
256	if (stopped != null) {	1 !
257	stopped.set(true);	1 ✔
258	}
259	if (this.subscriber != null) {	1 !
260	try {
261	ApiService service = this.subscriber.stopAsync();	×
262	service.awaitTerminated(5, TimeUnit.SECONDS);	×
263	} catch (Exception e) {	×
264	LOG.warn("CONNECTION: Error shutting down subscriber", e);	×
265	} finally {
266	this.subscriber = null;	×
267	}
268	}
269	nackQueued();	1 ✔
270	if (this.emulator != null) {	1 !
271	this.emulator.close();	×
272	this.emulator = null;	×
273	}
274	}	1 ✔
275
276	private void nackQueued() {
277	if (queue == null) {	1 !
278	return;	×
279	}
280	HeldMessage held;
281	while ((held = queue.poll()) != null) {	1 ✔
282	try {
283	held.consumer.nack();	1 ✔
284	} catch (RuntimeException e) {	×
285	LOG.warn("CONNECTION: Failed to nack queued message on close: {}", e.getMessage());	×
286	}	1 ✔
287	}
288	if (consumers != null) {	1 !
289	consumers.clear();	1 ✔
290	}
291	if (queuedBytes != null) {	1 !
292	queuedBytes.set(0);	1 ✔
293	}
294	}	1 ✔
295
296	/**
297	* Upper bound for the next {@link #poll(Duration)}. Messages beyond the cap stay in the queue.
298	* {@code maxMessages <= 0} means no cap.
299	*/
300	void limitNextPoll(int maxMessages) {
301	this.nextPollMax = maxMessages <= 0 ? Integer.MAX_VALUE : maxMessages;	1 !
302	}	1 ✔
303
304	/**
305	* Wait up to {@code timeout} for messages, up to {@link #limitNextPoll(int)}. When capped, keeps
306	* waiting within the same timeout until the cap is reached or the deadline passes. When uncapped,
307	* waits only for the first message, then drains the queue without blocking. Does not ack.
308	*/
309	List<PulledMessage> poll(Duration timeout) {
310	ensureStarted();	1 ✔
311	int max = this.nextPollMax;	1 ✔
312	this.nextPollMax = Integer.MAX_VALUE;	1 ✔
313	List<PulledMessage> messages = new ArrayList<>();	1 ✔
314	if (max <= 0) {	1 !
315	return messages;	×
316	}
317	long deadlineNanos = System.nanoTime() + Math.max(1L, timeout.toNanos());	1 ✔
318	try {
319	long waitNanos = deadlineNanos - System.nanoTime();	1 ✔
320	if (waitNanos <= 0L) {	1 !
NEW 321	return messages;	×
322	}
323	HeldMessage first = queue.poll(waitNanos, TimeUnit.NANOSECONDS);	1 ✔
324	if (first == null) {	1 ✔
325	return messages;	1 ✔
326	}
327	take(first, messages);	1 ✔
328	if (max == Integer.MAX_VALUE) {	1 !
329	HeldMessage next;
NEW 330	while ((next = queue.poll()) != null) {	×
NEW 331	take(next, messages);	×
332	}
NEW 333	} else {	×
334	while (messages.size() < max) {	1 ✔
335	waitNanos = deadlineNanos - System.nanoTime();	1 ✔
336	if (waitNanos <= 0L) {	1 !
NEW 337	break;	×
338	}
339	HeldMessage next = queue.poll(waitNanos, TimeUnit.NANOSECONDS);	1 ✔
340	if (next == null) {	1 !
NEW 341	break;	×
342	}
343	take(next, messages);	1 ✔
344	}	1 ✔
345	}
346	} catch (InterruptedException e) {	×
347	Thread.currentThread().interrupt();	×
348	}	1 ✔
349	final PublishTimeWindow window = PublishTimeWindow.of(messages);	1 ✔
350	if (window != null) {	1 !
351	this.lastSeenNewestPublishMillis.set(window.newestMillis());	1 ✔
352	}
353	return messages;	1 ✔
354	}
355
356	private void take(HeldMessage held, List<PulledMessage> messages) {
357	messages.add(held.message);	1 ✔
358	if (queuedBytes != null) {	1 !
359	queuedBytes.addAndGet(-held.message.data().length);	1 ✔
360	}
361	}	1 ✔
362
363	void acknowledge(List<String> ackIds) {
364	reply(ackIds, true);	1 ✔
365	}	1 ✔
366
367	void nack(List<String> ackIds) {
368	reply(ackIds, false);	1 ✔
369	}	1 ✔
370
371	/** Subscriber renews leases; kept for {@link AckLeaseWatchdog}. */
372	void extendAckDeadline(List<String> ackIds, int deadlineSeconds) {}	×
373
374	private void reply(List<String> ackIds, boolean ack) {
375	if (ackIds == null \|\| ackIds.isEmpty()) {	1 ✔
376	return;	1 ✔
377	}
378	if (consumers == null) {	1 !
379	return;	×
380	}
381	for (String id : ackIds) {	1 ✔
382	AckReplyConsumer consumer = consumers.remove(id);	1 ✔
383	if (consumer == null) {	1 ✔
384	continue;	1 ✔
385	}
386	try {
387	if (ack) {	1 ✔
388	consumer.ack();	1 ✔
389	} else {
390	consumer.nack();	1 ✔
391	}
392	} catch (RuntimeException e) {	×
393	LOG.warn("ACK: Failed to {} {}: {}", ack ? "ack" : "nack", id, e.getMessage());	×
394	}	1 ✔
395	}	1 ✔
396	}	1 ✔
397
398	void releaseMessages(List<PulledMessage> messages) {
399	if (this.outstandingBytes != null) {	1 !
400	long bytes = PulledMessage.payloadBytes(messages);	1 ✔
401	if (bytes > 0) {	1 ✔
402	this.outstandingBytes.addAndGet(-bytes);	1 ✔
403	}
404	}
405	}	1 ✔
406
407	void resetOutstandingBytes() {
408	if (this.outstandingBytes != null) {	1 !
409	this.outstandingBytes.set(0);	1 ✔
410	}
411	}	1 ✔
412
413	long outstandingBytes() {
414	return this.outstandingBytes == null ? 0L : this.outstandingBytes.get();	1 !
415	}
416
417	long retryAttempts() {
418	return retryPolicy.retryAttempts();	×
419	}
420
421	long lastSeenNewestPublishMillis() {
422	return lastSeenNewestPublishMillis.get();	1 ✔
423	}
424
425	public static String asString(final Duration duration) {
426	return duration.toString().substring(2).replaceAll("(\\d[HMS])(?!$)", "$1").toLowerCase();	1 ✔
427	}
428	}

juarezr / spark-streaming-google-pubsub / 57

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

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