juarezr / spark-streaming-google-pubsub / 60

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

Build # 60

Build Type

push

github

Committed by

web-flow

Commit Message

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

Coverage Stats

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

39.13

/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.CredentialsProvider;
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.common.util.concurrent.MoreExecutors;
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);

  private final RetryPolicy retryPolicy;
  private PubSubConnectionErrors connectionErrors;
  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.connectionErrors = new PubSubConnectionErrors(this.config.subscriptionPath());

    this.emulator = this.config.emulatorHost().map(PubSubEmulator::new).orElse(null);
    try {
      CredentialsProvider authProvider = this.credentialsProvider.providerFor(this.config);
      if (this.emulator == null) {
        connectionErrors.verifySubscriptionAccessible(authProvider);
      }
      applySeekIfNeeded(authProvider);

      startSubscriberAsyncWith(authProvider);

      this.connectionErrors.throwIfFatalConnection();
    } catch (Exception e) {
      this.close();
      if (e instanceof IOException) {
        throw (IOException) e;
      }
      throw new IOException("Unable to start Pub/Sub subscriber", e);
    }
  }

  private void startSubscriberAsyncWith(CredentialsProvider authProvider) {

    Subscriber.Builder builder =
        Subscriber.newBuilder(this.config.subscriptionPath(), new PubSubMessageReceiver());
    builder.setFlowControlSettings(flowControl());

    long ackExtendSecs = config.effectiveAckDeadline(null).getSeconds();
    long ackExtendMax = Math.max(60L, ackExtendSecs);
    org.threeten.bp.Duration ackExtend = org.threeten.bp.Duration.ofSeconds(ackExtendMax);
    org.threeten.bp.Duration ackMaxPeriod = org.threeten.bp.Duration.ofMinutes(60);
    builder.setMaxAckExtensionPeriod(ackMaxPeriod);
    builder.setMaxDurationPerAckExtension(ackExtend);

    if (this.emulator != null) {
      this.emulator.configureSubscriber(builder);
    } else {
      builder.setCredentialsProvider(authProvider);
    }
    this.subscriber = builder.build();
    this.subscriber.addListener(connectionErrors, MoreExecutors.directExecutor());
    this.subscriber.startAsync().awaitRunning();
  }

  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(CredentialsProvider authProvider) 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(authProvider);
      }
      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);
      }
    }
    if (this.connectionErrors != null) {
      try {
        this.connectionErrors.throwIfFatalConnection();
      } catch (IOException e) {
        throw new IllegalStateException(e.getMessage(), 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);
  }

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

juarezr / spark-streaming-google-pubsub / 60

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