From 80b0d5226ff3309dcc7ecdefa55f20dab3aaa129 Mon Sep 17 00:00:00 2001
From: Alan Lau <alau@confluent.io>
Date: Mon, 6 Jul 2026 11:45:33 -0400
Subject: [PATCH 1/2] Make a closed StagedMergeIterator throw
 InvalidStateStoreException

---
 .../state/internals/StagedMergeIterator.java   |  3 ++-
 .../internals/StagedMergeIteratorTest.java     | 18 ++++++++++++++++++
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/StagedMergeIterator.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/StagedMergeIterator.java
index 76bd3abc5b28c..0259052daac70 100644
--- a/streams/src/main/java/org/apache/kafka/streams/state/internals/StagedMergeIterator.java
+++ b/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("Store iterator for this transaction buffer has already been closed.");
         }
         if (prefetched != null) {
             return true;
diff --git a/streams/src/test/java/org/apache/kafka/streams/state/internals/StagedMergeIteratorTest.java b/streams/src/test/java/org/apache/kafka/streams/state/internals/StagedMergeIteratorTest.java
index 0d99e8cebce16..84db1756f8a08 100644
--- a/streams/src/test/java/org/apache/kafka/streams/state/internals/StagedMergeIteratorTest.java
+++ b/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<>();

From 0f13d0c8605602c0dbef8a132932ab9179c5c307 Mon Sep 17 00:00:00 2001
From: Alan Lau <alau@confluent.io>
Date: Mon, 13 Jul 2026 13:01:27 -0400
Subject: [PATCH 2/2] address review comments

---
 .../kafka/streams/state/internals/StagedMergeIterator.java      | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/StagedMergeIterator.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/StagedMergeIterator.java
index 0259052daac70..8477ef6dad998 100644
--- a/streams/src/main/java/org/apache/kafka/streams/state/internals/StagedMergeIterator.java
+++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/StagedMergeIterator.java
@@ -78,7 +78,7 @@ private void advanceBase() {
     @Override
     public boolean hasNext() {
         if (closed) {
-            throw new InvalidStateStoreException("Store iterator for this transaction buffer has already been closed.");
+            throw new InvalidStateStoreException("Iterator has already been closed.");
         }
         if (prefetched != null) {
             return true;