Bump version to 4.1.5-SNAPSHOT; explicitly throw exception in PscTopicUriPartitionSplitReader.fetch()

pom.xml

@@ -7,7 +7,7 @@
     <artifactId>psc-java-oss</artifactId>
     <description>PubSub Client (PSC)</description>
     <url>https://github.com/pinterest/psc</url>
-    <version>4.1.4</version>
+    <version>4.1.5-SNAPSHOT</version>
     <packaging>pom</packaging>
     <name>psc-java-oss</name>
     <modules>

psc-common/pom.xml

@@ -5,7 +5,7 @@
     <parent>
         <artifactId>psc-java-oss</artifactId>
         <groupId>com.pinterest.psc</groupId>
-        <version>4.1.4</version>
+        <version>4.1.5-SNAPSHOT</version>
         <relativePath>../pom.xml</relativePath>
     </parent>
     <modelVersion>4.0.0</modelVersion>

psc-examples/pom.xml

@@ -5,13 +5,13 @@
     <parent>
         <artifactId>psc-java-oss</artifactId>
         <groupId>com.pinterest.psc</groupId>
-        <version>4.1.4</version>
+        <version>4.1.5-SNAPSHOT</version>
         <relativePath>../pom.xml</relativePath>
     </parent>
     <modelVersion>4.0.0</modelVersion>
 
     <artifactId>psc-examples</artifactId>
-    <version>4.1.4</version>
+    <version>4.1.5-SNAPSHOT</version>
     <name>psc-examples</name>
 
     <properties>

psc-flink-logging/pom.xml

@@ -5,13 +5,13 @@
     <parent>
         <artifactId>psc-java-oss</artifactId>
         <groupId>com.pinterest.psc</groupId>
-        <version>4.1.4</version>
+        <version>4.1.5-SNAPSHOT</version>
         <relativePath>../pom.xml</relativePath>
     </parent>
     <modelVersion>4.0.0</modelVersion>
 
     <artifactId>psc-flink-logging</artifactId>
-    <version>4.1.4</version>
+    <version>4.1.5-SNAPSHOT</version>
 
     <dependencies>
         <dependency>

psc-flink/pom.xml

@@ -5,7 +5,7 @@
     <parent>
         <groupId>com.pinterest.psc</groupId>
         <artifactId>psc-java-oss</artifactId>
-        <version>4.1.4</version>
+        <version>4.1.5-SNAPSHOT</version>
         <relativePath>../pom.xml</relativePath>
     </parent>
     <artifactId>psc-flink</artifactId>

psc-flink/src/main/java/com/pinterest/flink/connector/psc/source/reader/PscTopicUriPartitionSplitReader.java

@@ -113,11 +113,21 @@ public RecordsWithSplitIds<PscConsumerMessage<byte[], byte[]>> fetch() throws IO
             // This happens if all assigned partitions are invalid or empty (starting offset >=
             // stopping offset). We just mark empty partitions as finished and return an empty
             // record container, and this consumer will be closed by SplitFetcherManager.
-            PscPartitionSplitRecords recordsBySplits =
-                    new PscPartitionSplitRecords(
-                            PscConsumerMessagesIterable.emptyIterable(), pscSourceReaderMetrics);
-            markEmptySplitsAsFinished(recordsBySplits);
-            return recordsBySplits;
+            if (e.getCause() != null &&
+                    (e.getCause().getClass().equals(IllegalStateException.class) || e.getCause().getClass().equals(WakeupException.class))) {
+                LOG.warn("Caught IllegalStateException or WakeupException in poll(), marking partitions as finished", e);
+                PscPartitionSplitRecords recordsBySplits =
+                        new PscPartitionSplitRecords(
+                                PscConsumerMessagesIterable.emptyIterable(), pscSourceReaderMetrics);
+                markEmptySplitsAsFinished(recordsBySplits);
+                return recordsBySplits;
+            } else {
+                LOG.error("Unrecoverable ConsumerException caught in poll()", e);
+                throw new RuntimeException(e);
+            }
+        } catch (Exception e) {
+            LOG.error("Unrecoverable Exception caught in poll()", e);
+            throw new RuntimeException(e);
         }
         PscPartitionSplitRecords recordsBySplits =
                 new PscPartitionSplitRecords(consumerMessagesIterable, pscSourceReaderMetrics);

psc-flink/src/test/java/com/pinterest/flink/connector/psc/source/reader/PscTopicUriPartitionSplitReaderTest.java

@@ -24,8 +24,10 @@
 import com.pinterest.flink.streaming.connectors.psc.PscTestEnvironmentWithKafkaAsPubSub;
 import com.pinterest.psc.common.TopicUriPartition;
 import com.pinterest.psc.config.PscConfiguration;
+import com.pinterest.psc.consumer.PscConsumer;
 import com.pinterest.psc.consumer.PscConsumerMessage;
 import com.pinterest.psc.exception.ClientException;
+import com.pinterest.psc.exception.consumer.ConsumerException;
 import com.pinterest.psc.exception.consumer.DeserializerException;
 import com.pinterest.psc.exception.startup.ConfigurationException;
 import com.pinterest.psc.serde.ByteArrayDeserializer;
@@ -46,6 +48,7 @@
 import org.apache.flink.runtime.metrics.groups.UnregisteredMetricGroups;
 import org.apache.kafka.clients.consumer.ConsumerConfig;
 import org.apache.kafka.clients.consumer.NoOffsetForPartitionException;
+import org.apache.kafka.common.KafkaException;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Disabled;
@@ -56,7 +59,9 @@
 import org.junit.jupiter.params.provider.NullAndEmptySource;
 import org.junit.jupiter.params.provider.ValueSource;
 
+import java.lang.reflect.Field;
 import java.io.IOException;
+import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
@@ -73,6 +78,9 @@
 import static com.pinterest.flink.connector.psc.testutils.PscSourceTestEnv.NUM_RECORDS_PER_PARTITION;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 /** Unit tests for {@link PscTopicUriPartitionSplitReader}. */
 public class PscTopicUriPartitionSplitReaderTest {
@@ -161,6 +169,44 @@ public void testWakeupThenAssign() throws IOException, ConfigurationException, C
                         new PscTopicUriPartitionSplit(tp, PscTopicUriPartitionSplit.EARLIEST_OFFSET)));
     }
 
+    @Test
+    public void testFetchThrowsOnNonWakeupConsumerException() throws Exception {
+        PscTopicUriPartitionSplitReader reader = createReader();
+        @SuppressWarnings("unchecked")
+        PscConsumer<byte[], byte[]> mockConsumer = (PscConsumer<byte[], byte[]>) mock(PscConsumer.class);
+        when(mockConsumer.poll(any(Duration.class)))
+                .thenThrow(new ConsumerException(new KafkaException("boom")));
+
+        Field consumerField = PscTopicUriPartitionSplitReader.class.getDeclaredField("consumer");
+        consumerField.setAccessible(true);
+        consumerField.set(reader, mockConsumer);
+
+        assertThatThrownBy(reader::fetch)
+                .isInstanceOf(RuntimeException.class)
+                .hasCauseInstanceOf(ConsumerException.class)
+                .hasRootCauseInstanceOf(KafkaException.class);
+    }
+
+    @Test
+    public void testFetchThrowsOnKafkaExceptionWithOomCause() throws Exception {
+        PscTopicUriPartitionSplitReader reader = createReader();
+        @SuppressWarnings("unchecked")
+        PscConsumer<byte[], byte[]> mockConsumer = (PscConsumer<byte[], byte[]>) mock(PscConsumer.class);
+        KafkaException kafkaException =
+                new KafkaException("Failed to load record batch", new OutOfMemoryError("Java heap space"));
+        when(mockConsumer.poll(any(Duration.class)))
+                .thenThrow(new ConsumerException(kafkaException));
+
+        Field consumerField = PscTopicUriPartitionSplitReader.class.getDeclaredField("consumer");
+        consumerField.setAccessible(true);
+        consumerField.set(reader, mockConsumer);
+
+        assertThatThrownBy(reader::fetch)
+                .isInstanceOf(RuntimeException.class)
+                .hasCauseInstanceOf(ConsumerException.class)
+                .hasRootCauseInstanceOf(OutOfMemoryError.class);
+    }
+
     @Test
     public void testNumBytesInCounter() throws Exception {
         final OperatorMetricGroup operatorMetricGroup =

psc-integration-test/pom.xml

@@ -5,7 +5,7 @@
     <parent>
         <artifactId>psc-java-oss</artifactId>
         <groupId>com.pinterest.psc</groupId>
-        <version>4.1.4</version>
+        <version>4.1.5-SNAPSHOT</version>
         <relativePath>../pom.xml</relativePath>
     </parent>
     <modelVersion>4.0.0</modelVersion>

psc-logging/pom.xml

@@ -5,7 +5,7 @@
     <parent>
         <artifactId>psc-java-oss</artifactId>
         <groupId>com.pinterest.psc</groupId>
-        <version>4.1.4</version>
+        <version>4.1.5-SNAPSHOT</version>
         <relativePath>../pom.xml</relativePath>
     </parent>
     <modelVersion>4.0.0</modelVersion>

psc/pom.xml

@@ -5,7 +5,7 @@
     <parent>
         <artifactId>psc-java-oss</artifactId>
         <groupId>com.pinterest.psc</groupId>
-        <version>4.1.4</version>
+        <version>4.1.5-SNAPSHOT</version>
         <relativePath>../pom.xml</relativePath>
     </parent>
     <modelVersion>4.0.0</modelVersion>

psc/src/main/java/com/pinterest/psc/consumer/PscConsumer.java

@@ -702,9 +702,9 @@ public PscConsumerPollMessageIterator<K, V> poll(Duration pollTimeout) throws Co
         acquireAndEnsureOpen();
         try {
             if (messageListener != null) {
-                throw new ConsumerException(ExceptionMessage.MUTUALLY_EXCLUSIVE_APIS("poll()", "MessageListener"));
+                throw new ConsumerException(ExceptionMessage.MUTUALLY_EXCLUSIVE_APIS("poll()", "MessageListener"), new IllegalStateException());
             } else if (!subscribed.get() && !assigned.get()) {
-                throw new ConsumerException(ExceptionMessage.NO_SUBSCRIPTION_ASSIGNMENT("poll()"));
+                throw new ConsumerException(ExceptionMessage.NO_SUBSCRIPTION_ASSIGNMENT("poll()"), new IllegalStateException());
             }
             return internalPoll(pollTimeout);
         } finally {