juarezr / spark-streaming-google-pubsub / 54

Committed 29 Sep 2026 02:55AM UTC coverage: 68.457%. First build

Build # 54

Build Type

push

github

Committed by

web-flow

Commit Message

feat: add code coverage support with JaCoCo and Coveralls (#19)

- Introduced new Maven profiles for coverage reporting and Coveralls
integration.
- Updated README and documentation to include coverage instructions and
usage.
- Enhanced CI workflow to run coverage tests and publish reports to
Coveralls when applicable.
- Added integration tests for PubSubClient, validating message
publishing and acknowledgment against the Pub/Sub emulator

Coverage Stats

647 of 1061 branches covered (60.98%)

Branch coverage included in aggregate %.

1 of 38 new or added lines in 2 files covered. (2.63%)

1430 of 1973 relevant lines covered (72.48%)

0.72 hits per line

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

38.44

/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 instant = this.config.seekTimeAsInstant();
        final Timestamp seekTime =
            Timestamp.newBuilder()
                .setSeconds(instant.getEpochSecond())
                .setNanos(instant.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 the first message, then drain what is already queued up to
   * {@link #limitNextPoll(int)}. 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;
    }
    try {
      HeldMessage first = queue.poll(Math.max(1L, timeout.toNanos()), TimeUnit.NANOSECONDS);
      if (first == null) {
        return messages;
      }
      take(first, messages);
      HeldMessage next;
      while (messages.size() < max && (next = queue.poll()) != null) {
        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
NEW 62	final class PubSubMessageReceiver implements MessageReceiver {	×
63
64	@Override
65	public void receiveMessage(PubsubMessage message, AckReplyConsumer consumer) {
NEW 66	if (stopped != null && stopped.get()) {	×
NEW 67	consumer.nack();	×
NEW 68	return;	×
69	}
NEW 70	PulledMessage pulled = toPulled(message);	×
NEW 71	long cap = config.batchSize() > 0 ? config.batchSize() : PubSubConfig.DEFAULT_BATCH_SIZE;	×
NEW 72	while (queuedBytes != null && queuedBytes.get() >= cap && stopped != null && !stopped.get()) {	×
73	try {
NEW 74	TimeUnit.MILLISECONDS.sleep(10);	×
NEW 75	} catch (InterruptedException e) {	×
NEW 76	Thread.currentThread().interrupt();	×
NEW 77	consumer.nack();	×
NEW 78	return;	×
NEW 79	}	×
80	}
NEW 81	if (stopped != null && stopped.get()) {	×
NEW 82	consumer.nack();	×
NEW 83	return;	×
84	}
NEW 85	consumers.put(pulled.ackId(), consumer);	×
NEW 86	queuedBytes.addAndGet(pulled.data().length);	×
NEW 87	outstandingBytes.addAndGet(pulled.data().length);	×
88	try {
NEW 89	queue.put(new HeldMessage(pulled, consumer));	×
NEW 90	} catch (InterruptedException e) {	×
NEW 91	Thread.currentThread().interrupt();	×
NEW 92	consumers.remove(pulled.ackId());	×
NEW 93	queuedBytes.addAndGet(-pulled.data().length);	×
NEW 94	outstandingBytes.addAndGet(-pulled.data().length);	×
NEW 95	consumer.nack();	×
NEW 96	}	×
NEW 97	}	×
98
99	private PulledMessage toPulled(PubsubMessage message) {
NEW 100	byte[] data = message.getData().toByteArray();	×
NEW 101	long publishMillis =	×
NEW 102	message.getPublishTime().getSeconds() * 1000L	×
NEW 103	+ message.getPublishTime().getNanos() / 1_000_000L;	×
NEW 104	String messageId = message.getMessageId();	×
NEW 105	return new PulledMessage(	×
106	messageId,
107	data,
NEW 108	message.getAttributesMap(),	×
109	publishMillis,
NEW 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 =	×
NEW 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:
227	final Instant instant = this.config.seekTimeAsInstant();	×
228	final Timestamp seekTime =
229	Timestamp.newBuilder()	×
230	.setSeconds(instant.getEpochSecond())	×
231	.setNanos(instant.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 the first message, then drain what is already queued up to
306	* {@link #limitNextPoll(int)}. Does not ack.
307	*/
308	List<PulledMessage> poll(Duration timeout) {
309	ensureStarted();	1 ✔
310	int max = this.nextPollMax;	1 ✔
311	this.nextPollMax = Integer.MAX_VALUE;	1 ✔
312	List<PulledMessage> messages = new ArrayList<>();	1 ✔
313	if (max <= 0) {	1 !
314	return messages;	×
315	}
316	try {
317	HeldMessage first = queue.poll(Math.max(1L, timeout.toNanos()), TimeUnit.NANOSECONDS);	1 ✔
318	if (first == null) {	1 ✔
319	return messages;	1 ✔
320	}
321	take(first, messages);	1 ✔
322	HeldMessage next;
323	while (messages.size() < max && (next = queue.poll()) != null) {	1 !
324	take(next, messages);	1 ✔
325	}
326	} catch (InterruptedException e) {	×
327	Thread.currentThread().interrupt();	×
328	}	1 ✔
329	final PublishTimeWindow window = PublishTimeWindow.of(messages);	1 ✔
330	if (window != null) {	1 !
331	this.lastSeenNewestPublishMillis.set(window.newestMillis());	1 ✔
332	}
333	return messages;	1 ✔
334	}
335
336	private void take(HeldMessage held, List<PulledMessage> messages) {
337	messages.add(held.message);	1 ✔
338	if (queuedBytes != null) {	1 !
339	queuedBytes.addAndGet(-held.message.data().length);	1 ✔
340	}
341	}	1 ✔
342
343	void acknowledge(List<String> ackIds) {
344	reply(ackIds, true);	1 ✔
345	}	1 ✔
346
347	void nack(List<String> ackIds) {
348	reply(ackIds, false);	1 ✔
349	}	1 ✔
350
351	/** Subscriber renews leases; kept for {@link AckLeaseWatchdog}. */
352	void extendAckDeadline(List<String> ackIds, int deadlineSeconds) {}	×
353
354	private void reply(List<String> ackIds, boolean ack) {
355	if (ackIds == null \|\| ackIds.isEmpty()) {	1 ✔
356	return;	1 ✔
357	}
358	if (consumers == null) {	1 !
359	return;	×
360	}
361	for (String id : ackIds) {	1 ✔
362	AckReplyConsumer consumer = consumers.remove(id);	1 ✔
363	if (consumer == null) {	1 ✔
364	continue;	1 ✔
365	}
366	try {
367	if (ack) {	1 ✔
368	consumer.ack();	1 ✔
369	} else {
370	consumer.nack();	1 ✔
371	}
372	} catch (RuntimeException e) {	×
373	LOG.warn("ACK: Failed to {} {}: {}", ack ? "ack" : "nack", id, e.getMessage());	×
374	}	1 ✔
375	}	1 ✔
376	}	1 ✔
377
378	void releaseMessages(List<PulledMessage> messages) {
379	if (this.outstandingBytes != null) {	1 !
380	long bytes = PulledMessage.payloadBytes(messages);	1 ✔
381	if (bytes > 0) {	1 ✔
382	this.outstandingBytes.addAndGet(-bytes);	1 ✔
383	}
384	}
385	}	1 ✔
386
387	void resetOutstandingBytes() {
388	if (this.outstandingBytes != null) {	1 !
389	this.outstandingBytes.set(0);	1 ✔
390	}
391	}	1 ✔
392
393	long outstandingBytes() {
394	return this.outstandingBytes == null ? 0L : this.outstandingBytes.get();	1 !
395	}
396
397	long retryAttempts() {
398	return retryPolicy.retryAttempts();	×
399	}
400
401	long lastSeenNewestPublishMillis() {
402	return lastSeenNewestPublishMillis.get();	1 ✔
403	}
404
405	public static String asString(final Duration duration) {
406	return duration.toString().substring(2).replaceAll("(\\d[HMS])(?!$)", "$1").toLowerCase();	1 ✔
407	}
408	}

juarezr / spark-streaming-google-pubsub / 54

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