From cd68f7886b7d8f5000ef672a8c7edcfc191960de Mon Sep 17 00:00:00 2001 From: Michael Marshall Date: Wed, 17 Dec 2025 15:43:56 -0600 Subject: [PATCH] CNDB-16336: Handle failure in IndexSearcher::toMetaSortedIterator, close rowIdIterator --- .../index/sai/disk/v1/IndexSearcher.java | 40 +++++++----- .../index/sai/cql/VectorTypeTest.java | 64 +++++++++++++++++++ 2 files changed, 88 insertions(+), 16 deletions(-) diff --git a/src/java/org/apache/cassandra/index/sai/disk/v1/IndexSearcher.java b/src/java/org/apache/cassandra/index/sai/disk/v1/IndexSearcher.java index 16d14bdd8ac3..38cef5334bcf 100644 --- a/src/java/org/apache/cassandra/index/sai/disk/v1/IndexSearcher.java +++ b/src/java/org/apache/cassandra/index/sai/disk/v1/IndexSearcher.java @@ -171,24 +171,32 @@ protected KeyRangeIterator toPrimaryKeyIterator(PostingList postingList, QueryCo protected CloseableIterator toMetaSortedIterator(CloseableIterator rowIdIterator, QueryContext queryContext) throws IOException { - if (rowIdIterator == null || !rowIdIterator.hasNext()) + try + { + if (rowIdIterator == null || !rowIdIterator.hasNext()) + { + FileUtils.closeQuietly(rowIdIterator); + return CloseableIterator.emptyIterator(); + } + + IndexSearcherContext searcherContext = new IndexSearcherContext(metadata.minKey, + metadata.maxKey, + metadata.minSSTableRowId, + metadata.maxSSTableRowId, + metadata.segmentRowIdOffset, + queryContext, + null); + var pkm = primaryKeyMapFactory.newPerSSTablePrimaryKeyMap(); + return new RowIdToPrimaryKeyWithSortKeyIterator(indexContext, + pkm.getSSTableId(), + rowIdIterator, + pkm, + searcherContext); + } + catch (Throwable t) { FileUtils.closeQuietly(rowIdIterator); - return CloseableIterator.emptyIterator(); + throw t; } - - IndexSearcherContext searcherContext = new IndexSearcherContext(metadata.minKey, - metadata.maxKey, - metadata.minSSTableRowId, - metadata.maxSSTableRowId, - metadata.segmentRowIdOffset, - queryContext, - null); - var pkm = primaryKeyMapFactory.newPerSSTablePrimaryKeyMap(); - return new RowIdToPrimaryKeyWithSortKeyIterator(indexContext, - pkm.getSSTableId(), - rowIdIterator, - pkm, - searcherContext); } } diff --git a/test/unit/org/apache/cassandra/index/sai/cql/VectorTypeTest.java b/test/unit/org/apache/cassandra/index/sai/cql/VectorTypeTest.java index bf7cc8789132..d4fa63a85ba8 100644 --- a/test/unit/org/apache/cassandra/index/sai/cql/VectorTypeTest.java +++ b/test/unit/org/apache/cassandra/index/sai/cql/VectorTypeTest.java @@ -36,11 +36,13 @@ import org.apache.cassandra.index.sai.StorageAttachedIndex; import org.apache.cassandra.index.sai.disk.v1.IndexWriterConfig; import org.apache.cassandra.index.sai.disk.v1.SegmentBuilder; +import org.apache.cassandra.index.sai.disk.vector.AutoResumingNodeScoreIterator; import org.apache.cassandra.index.sai.disk.vector.CassandraOnHeapGraph; import org.apache.cassandra.index.sai.disk.vector.VectorSourceModel; import org.apache.cassandra.index.sai.plan.QueryController; import org.apache.cassandra.inject.ActionBuilder; import org.apache.cassandra.inject.Expression; +import org.apache.cassandra.inject.Injection; import org.apache.cassandra.inject.Injections; import org.apache.cassandra.inject.InvokePointBuilder; import org.apache.cassandra.service.ClientState; @@ -1057,4 +1059,66 @@ public void testUpdateRowScoreToWorsePositionButIncludeInBatch() // query with ANN only assertRows(execute("SELECT c FROM %s ORDER BY r ANN OF [0.1, 0.1] LIMIT 10"), row(2), row(1)); } + + @Test + public void testRowIdIteratorClosedOnHasNextFailure() throws Throwable + { + createTable("CREATE TABLE %s (pk int, vec vector, PRIMARY KEY(pk))"); + createIndex("CREATE CUSTOM INDEX ON %s(vec) USING 'StorageAttachedIndex'"); + + // Track if the rowIdIterator's close method is called + Injections.Counter closeCounter = Injections.newCounter("rowIdIteratorCloseCounter") + .add(InvokePointBuilder.newInvokePoint() + .onClass(AutoResumingNodeScoreIterator.class) + .onMethod("close")) + .build(); + + // Inject failure at hasNext in toMetaSortedIterator + Injection hasNextFailure = Injections.newCustom("fail_on_hasNext") + .add(InvokePointBuilder.newInvokePoint() + .onClass(AutoResumingNodeScoreIterator.class) + .onMethod("computeNext") + .atEntry()) + .add(ActionBuilder.newActionBuilder() + .actions() + .doThrow(RuntimeException.class, Expression.quote("Injected hasNext failure!"))) + .build(); + + try + { + Injections.inject(closeCounter); + Injections.inject(hasNextFailure); + + // Insert data + execute("INSERT INTO %s (pk, vec) VALUES (1, [1.0, 1.0])"); + execute("INSERT INTO %s (pk, vec) VALUES (2, [2.0, 2.0])"); + flush(); + + // Reset counter before the test + closeCounter.reset(); + + // Enable the failure injection + hasNextFailure.enable(); + + // Execute query that will trigger toMetaSortedIterator and fail at hasNext + assertThatThrownBy(() -> executeInternal("SELECT pk FROM %s ORDER BY vec ANN OF [1.5, 1.5] LIMIT 2")) + .hasMessageContaining("Injected hasNext failure!"); + + // Verify that close was called on the rowIdIterator despite the failure + // The close should be called in the catch block of toMetaSortedIterator (line 198) + assertThat(closeCounter.get()).as("rowIdIterator should be closed when hasNext fails") + .isGreaterThan(0); + + // Remove failure and confirm we can still query + hasNextFailure.disable(); + + // Confrm subsequent queries succeed because we close the iterator and release the graph searcher + execute("SELECT pk FROM %s ORDER BY vec ANN OF [1.5, 1.5] LIMIT 2"); + } + finally + { + hasNextFailure.disable(); + closeCounter.disable(); + } + } }