putShortVolatile is not volatile in InMemoryTrie

This commit fixes several issues found using the Cassandra
edge case explorer skill for the BTI feature.

patch by Francisco Guerrero; reviewed by Aleksey Yeschenko for CASSANDRA-21353

CHANGES.txt

@@ -1,4 +1,5 @@
 5.0.9
+ * putShortVolatile is not volatile in InMemoryTrie (CASSANDRA-21353)
  * Fix RequestFailureReason serializer and nits in a few others (CASSANDRA-21437)
  * Remove golang dependency in gen-doc and replace with python implementation (CASSANDRA-21432)
  * Use estimated compressed size for tables to check if there is enough free space for a compaction (CASSANDRA-21245)

src/java/org/apache/cassandra/db/tries/InMemoryTrie.java

@@ -128,7 +128,7 @@ public class InMemoryTrie<T> extends InMemoryReadTrie<T>
 
     final void putShortVolatile(int pos, short value)
     {
-        getChunk(pos).putShort(inChunkPointer(pos), value);
+        getChunk(pos).putShortVolatile(inChunkPointer(pos), value);
     }
 
     final void putByte(int pos, byte value)

src/java/org/apache/cassandra/io/sstable/AbstractSSTableIterator.java

@@ -97,6 +97,7 @@ public abstract class AbstractSSTableIterator<RIE extends AbstractRowIndexEntry>
         }
         else
         {
+            Reader reader = null;
             boolean shouldCloseFile = file == null;
             try
             {
@@ -119,15 +120,18 @@ public abstract class AbstractSSTableIterator<RIE extends AbstractRowIndexEntry>
 
                     // Note that this needs to be called after file != null and after the partitionDeletion has been set, but before readStaticRow
                     // (since it uses it) so we can't move that up (but we'll be able to simplify as soon as we drop support for the old file format).
-                    this.reader = createReader(indexEntry, file, shouldCloseFile);
+                    reader = createReader(indexEntry, file, shouldCloseFile);
                     this.staticRow = readStaticRow(sstable, file, helper, columns.fetchedColumns().statics);
                 }
                 else
                 {
                     this.partitionLevelDeletion = indexEntry.deletionTime();
                     this.staticRow = Rows.EMPTY_STATIC_ROW;
-                    this.reader = createReader(indexEntry, file, shouldCloseFile);
+                    reader = createReader(indexEntry, file, shouldCloseFile);
                 }
+
+                this.reader = reader;
+
                 if (!partitionLevelDeletion.validate())
                     UnfilteredValidation.handleInvalid(metadata(), key, sstable, "partitionLevelDeletion="+partitionLevelDeletion.toString());
 
@@ -137,9 +141,25 @@ public abstract class AbstractSSTableIterator<RIE extends AbstractRowIndexEntry>
                 if (reader == null && file != null && shouldCloseFile)
                     file.close();
             }
-            catch (IOException e)
+            catch (CorruptSSTableException | IOException e)
             {
                 sstable.markSuspect();
+
+                if (reader != null)
+                {
+                    try
+                    {
+                        reader.close();
+                        // reader will close the file internally, so there's no
+                        // need to close it in the next block
+                        shouldCloseFile = false;
+                    }
+                    catch (IOException suppressed)
+                    {
+                        e.addSuppressed(suppressed);
+                    }
+                }
+
                 String filePath = file.getPath();
                 if (shouldCloseFile)
                 {
@@ -152,6 +172,8 @@ public abstract class AbstractSSTableIterator<RIE extends AbstractRowIndexEntry>
                         e.addSuppressed(suppressed);
                     }
                 }
+                if (e instanceof CorruptSSTableException)
+                    throw (CorruptSSTableException)e;
                 throw new CorruptSSTableException(e, filePath);
             }
         }

src/java/org/apache/cassandra/io/sstable/format/bti/BtiTableWriter.java

@@ -329,12 +329,15 @@ public class BtiTableWriter extends SortedTableWriter<BtiFormatPartitionWriter,
         {
             checkState(!dataWriterOpened, "Data writer has been already opened.");
 
-            return DataComponent.buildWriter(descriptor,
+            SequentialWriter sequentialWriter = DataComponent.buildWriter(descriptor,
-                                             getTableMetadataRef().getLocal(),
+                    getTableMetadataRef().getLocal(),
-                                             getIOOptions().writerOptions,
+                    getIOOptions().writerOptions,
-                                             getMetadataCollector(),
+                    getMetadataCollector(),
-                                             ensuringInBuildInternalContext(operationType),
+                    ensuringInBuildInternalContext(operationType),
-                                             getIOOptions().flushCompression);
+                    getIOOptions().flushCompression);
+            dataWriterOpened = true;
+
+            return sequentialWriter;
         }
 
         @Override

src/java/org/apache/cassandra/io/tries/ReverseValueIterator.java

@@ -61,7 +61,16 @@ public class ReverseValueIterator<Concrete extends ReverseValueIterator<Concrete
     {
         super(source, root);
         limit = null;
-        initializeNoRightBound(root, NOT_AT_LIMIT, false);
+
+        try
+        {
+            initializeNoRightBound(root, NOT_AT_LIMIT, false);
+        }
+        catch (Throwable t)
+        {
+            super.close();
+            throw t;
+        }
     }
 
     /**
@@ -80,10 +89,18 @@ public class ReverseValueIterator<Concrete extends ReverseValueIterator<Concrete
         super(source, root);
         limit = start != null ? start.asComparableBytes(BYTE_COMPARABLE_VERSION) : null;
 
-        if (end != null)
+        try
-            initializeWithRightBound(root, end.asComparableBytes(BYTE_COMPARABLE_VERSION), admitPrefix, limit != null);
+        {
-        else
+            if (end != null)
-            initializeNoRightBound(root, limit != null ? limit.next() : NOT_AT_LIMIT, admitPrefix);
+                initializeWithRightBound(root, end.asComparableBytes(BYTE_COMPARABLE_VERSION), admitPrefix, limit != null);
+            else
+                initializeNoRightBound(root, limit != null ? limit.next() : NOT_AT_LIMIT, admitPrefix);
+        }
+        catch (Throwable t)
+        {
+            super.close();
+            throw t;
+        }
     }
 
     void initializeWithRightBound(long root, ByteSource endStream, boolean admitPrefix, boolean hasLimit)

src/java/org/apache/cassandra/io/tries/ValueIterator.java

@@ -72,7 +72,16 @@ public class ValueIterator<CONCRETE extends ValueIterator<CONCRETE>> extends Wal
         super(source, root);
         limit = null;
         collector = collecting ? new TransitionBytesCollector() : null;
-        initializeNoLeftBound(root, 256);
+
+        try
+        {
+            initializeNoLeftBound(root, 256);
+        }
+        catch (Throwable t)
+        {
+            super.close();
+            throw t;
+        }
     }
 
     protected ValueIterator(Rebufferer source, long root, ByteComparable start, ByteComparable end, boolean admitPrefix)
@@ -97,10 +106,18 @@ public class ValueIterator<CONCRETE extends ValueIterator<CONCRETE>> extends Wal
         limit = end != null ? end.asComparableBytes(BYTE_COMPARABLE_VERSION) : null;
         collector = collecting ? new TransitionBytesCollector() : null;
 
-        if (start != null)
+        try
-            initializeWithLeftBound(root, start.asComparableBytes(BYTE_COMPARABLE_VERSION), admitPrefix, limit != null);
+        {
-        else
+            if (start != null)
-            initializeNoLeftBound(root, limit != null ? limit.next() : 256);
+                initializeWithLeftBound(root, start.asComparableBytes(BYTE_COMPARABLE_VERSION), admitPrefix, limit != null);
+            else
+                initializeNoLeftBound(root, limit != null ? limit.next() : 256);
+        }
+        catch (Throwable t)
+        {
+            super.close();
+            throw t;
+        }
     }
 
     private void initializeWithLeftBound(long root, ByteSource startStream, boolean admitPrefix, boolean atLimit)
@@ -110,76 +127,60 @@ public class ValueIterator<CONCRETE extends ValueIterator<CONCRETE>> extends Wal
         int limitByte;
         long payloadedNode = -1;
 
-        try
+        // Follow start position while we still have a prefix, stacking path and saving prefixes.
+        go(root);
+        while (true)
         {
-            // Follow start position while we still have a prefix, stacking path and saving prefixes.
+            int s = startStream.next();
-            go(root);
+            childIndex = search(s);
-            while (true)
+
+            // For a separator trie the latest payload met along the prefix is a potential match for start
+            if (admitPrefix)
             {
-                int s = startStream.next();
+                if (childIndex == 0 || childIndex == -1)
-                childIndex = search(s);
-
-                // For a separator trie the latest payload met along the prefix is a potential match for start
-                if (admitPrefix)
                 {
-                    if (childIndex == 0 || childIndex == -1)
+                    if (hasPayload())
-                    {
+                        payloadedNode = position;
-                        if (hasPayload())
-                            payloadedNode = position;
-                    }
-                    else
-                    {
-                        payloadedNode = -1;
-                    }
                 }
-
+                else
-                limitByte = 256;
-                if (atLimit)
                 {
-                    limitByte = limit.next();
+                    payloadedNode = -1;
-                    if (s < limitByte)
-                        atLimit = false;
                 }
-                if (childIndex < 0)
-                    break;
-
-                prev = new IterationPosition(position, childIndex, limitByte, prev);
-                go(transition(childIndex)); // child index is positive, transition must exist
             }
 
-            childIndex = -1 - childIndex - 1;
+            limitByte = 256;
-            stack = new IterationPosition(position, childIndex, limitByte, prev);
+            if (atLimit)
+            {
+                limitByte = limit.next();
+                if (s < limitByte)
+                    atLimit = false;
+            }
+            if (childIndex < 0)
+                break;
 
-            // Advancing now gives us first match if we didn't find one already.
+            prev = new IterationPosition(position, childIndex, limitByte, prev);
-            if (payloadedNode != -1)
+            go(transition(childIndex)); // child index is positive, transition must exist
-                next = payloadedNode;
-            else
-                next = advanceNode();
-        }
-        catch (Throwable t)
-        {
-            super.close();
-            throw t;
         }
+
+        childIndex = -1 - childIndex - 1;
+        stack = new IterationPosition(position, childIndex, limitByte, prev);
+
+        // Advancing now gives us first match if we didn't find one already.
+        if (payloadedNode != -1)
+            next = payloadedNode;
+        else
+            next = advanceNode();
     }
 
     private void initializeNoLeftBound(long root, int limitByte)
     {
         stack = new IterationPosition(root, -1, limitByte, null);
 
-        try
+        go(root);
-        {
+        if (hasPayload())
-            go(root);
+            next = root;
-            if (hasPayload())
+        else
-                next = root;
+            next = advanceNode();
-            else
-                next = advanceNode();
-        }
-        catch (Throwable t)
-        {
-            super.close();
-            throw t;
-        }
     }
 
     /**

src/java/org/apache/cassandra/io/tries/Walker.java

@@ -44,7 +44,7 @@ import org.apache.lucene.util.ArrayUtil;
 public class Walker<CONCRETE extends Walker<CONCRETE>> implements AutoCloseable
 {
     /** Value used to indicate a branch (e.g. lesser/greaterBranch) does not exist. */
-    public static int NONE = TrieNode.NONE;
+    public static final int NONE = TrieNode.NONE;
 
     private final Rebufferer source;
     protected final long root;

test/unit/org/apache/cassandra/io/tries/ReverseValueIteratorLeakTest.java

@@ -0,0 +1,112 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.cassandra.io.tries;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+import org.junit.Test;
+
+import org.apache.cassandra.io.util.DataOutputBuffer;
+import org.apache.cassandra.io.util.Rebufferer;
+import org.apache.cassandra.utils.bytecomparable.ByteComparable;
+import org.apache.cassandra.utils.bytecomparable.ByteSource;
+
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+
+/**
+ * Reproducer: ReverseValueIterator leaks Rebufferer on constructor exception
+ *
+ * <p>If initializeWithRightBound() or initializeNoRightBound() throws after the Walker
+ * super-constructor acquires a Rebufferer, the Rebufferer is never closed. Compare with
+ * ValueIterator which wraps initialization in try-catch with super.close().
+ *
+ * Expected: Rebufferer.closeReader() is called when initialization throws.
+ * Actual (buggy): Rebufferer leaks — closeReader() never called.
+ * Failure criterion: closeReader() not called after constructor exception.
+ */
+@SuppressWarnings({"unchecked", "RedundantSuppression"})
+public class ReverseValueIteratorLeakTest extends AbstractTrieTestBase
+{
+    // ORACLE: Rebufferer.closeReader() must be called if constructor throws
+    @Test
+    public void testRebuffererClosedOnInitWithRightBoundFailure() throws IOException
+    {
+        // HARNESS: build a valid trie, wrap its Rebufferer with close tracking,
+        //          then trigger a failure during initializeWithRightBound
+        DataOutputBuffer buf = new DataOutputBufferPaged();
+        IncrementalTrieWriter<Integer> builder = newTrieWriter(serializer, buf);
+        builder.add(source("aaa"), 1);
+        builder.add(source("bbb"), 2);
+        long root = builder.complete();
+
+        AtomicBoolean closeReaderCalled = new AtomicBoolean(false);
+        Rebufferer trackingSource = new TrackingRebufferer(buf.asNewBuffer(), closeReaderCalled);
+
+        // TRIGGER: provide a ByteComparable whose ByteSource throws during iteration,
+        //          causing initializeWithRightBound to fail after Walker acquires the Rebufferer
+        ByteComparable throwingEnd = v -> new ByteSource()
+        {
+            private int calls = 0;
+
+            @Override
+            public int next()
+            {
+                if (++calls > 1)
+                    throw new RuntimeException("simulated failure during trie traversal");
+                return 'b';
+            }
+        };
+
+        try
+        {
+            new ReverseValueIterator<>(trackingSource, root, source("aaa"), throwingEnd, false);
+            fail("Constructor should have thrown");
+        }
+        catch (RuntimeException e)
+        {
+            assertTrue("Expected simulated failure", e.getMessage().contains("simulated failure"));
+        }
+
+        // ORACLE: Rebufferer must have been closed despite the constructor failure
+        assertTrue("Rebufferer.closeReader() must be called when ReverseValueIterator " +
+                   "constructor fails during initializeWithRightBound — without the fix, " +
+                   "the Rebufferer leaks",
+                   closeReaderCalled.get());
+    }
+
+    private static class TrackingRebufferer extends ByteBufRebufferer
+    {
+        private final AtomicBoolean closeReaderCalled;
+
+        TrackingRebufferer(ByteBuffer buffer, AtomicBoolean closeReaderCalled)
+        {
+            super(buffer);
+            this.closeReaderCalled = closeReaderCalled;
+        }
+
+        @Override
+        public void closeReader()
+        {
+            closeReaderCalled.set(true);
+        }
+    }
+}