KAFKA-20777: Make a closed StagedMergeIterator throw InvalidStateStoreException

streams/src/main/java/org/apache/kafka/streams/state/internals/StagedMergeIterator.java

@@ -17,6 +17,7 @@
 package org.apache.kafka.streams.state.internals;
 
 import org.apache.kafka.streams.KeyValue;
+import org.apache.kafka.streams.errors.InvalidStateStoreException;
 import org.apache.kafka.streams.state.KeyValueIterator;
 
 import java.util.Iterator;
@@ -77,7 +78,7 @@ private void advanceBase() {
     @Override
     public boolean hasNext() {
         if (closed) {
-            throw new IllegalStateException("Iterator has already been closed.");
+            throw new InvalidStateStoreException("Iterator has already been closed.");
         }
         if (prefetched != null) {
             return true;

streams/src/test/java/org/apache/kafka/streams/state/internals/StagedMergeIteratorTest.java

@@ -18,6 +18,7 @@
 
 import org.apache.kafka.common.utils.Bytes;
 import org.apache.kafka.streams.KeyValue;
+import org.apache.kafka.streams.errors.InvalidStateStoreException;
 import org.apache.kafka.streams.state.KeyValueIterator;
 
 import org.junit.jupiter.api.Test;
@@ -87,6 +88,23 @@ public void shouldPreferStagingOverBaseForSameKey() {
         }
     }
 
+    @Test
+    public void shouldThrowInvalidStateStoreExceptionOnHasNextAfterClose() {
+        final NavigableMap<Bytes, Optional<byte[]>> staging = new TreeMap<>();
+        staging.put(key("a"), Optional.of(val("staged-a")));
+
+        final KeyValueIterator<Bytes, byte[]> iter =
+            new StagedMergeIterator(staging, new ListIterator(List.of()));
+        iter.close();
+
+        // A closed store iterator must signal via InvalidStateStoreException -- the type
+        // SegmentIterator catches to close open iterators gracefully -- not a bare
+        // IllegalStateException, which would escape that handling.
+        assertThrows(InvalidStateStoreException.class, iter::hasNext);
+        assertThrows(InvalidStateStoreException.class, iter::next);
+        assertThrows(InvalidStateStoreException.class, iter::peekNextKey);
+    }
+
     @Test
     public void shouldSkipTombstones() {
         final NavigableMap<Bytes, Optional<byte[]>> staging = new TreeMap<>();