Fix range queries on early-open BTI files

patch by Branimir Lambov; reviewed by Stefan Miklosovic for CASSANDRA-20976

CHANGES.txt

@@ -1,4 +1,5 @@
 5.0.6
+ * Fix range queries on early-open BTI files (CASSANDRA-20976)
  * Avoid re-initializing underlying iterator in LazilyInitializedUnfilteredRowIterator after closing (CASSANDRA-20972)
  * Flush SAI segment builder when current SSTable writer is switched (CASSANDRA-20752)
  * Throw RTE instead of FSError when RTE is thrown from FileUtis.write in TOCComponent (CASSANDRA-20917)

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

@@ -262,8 +262,8 @@ public class BtiTableWriter extends SortedTableWriter<BtiFormatPartitionWriter,
         PartitionIndex completedPartitionIndex()
         {
             complete();
-            rowIndexFHBuilder.withLengthOverride(0);
-            partitionIndexFHBuilder.withLengthOverride(0);
+            rowIndexFHBuilder.withLengthOverride(NO_LENGTH_OVERRIDE);
+            partitionIndexFHBuilder.withLengthOverride(NO_LENGTH_OVERRIDE);
             try
             {
                 return PartitionIndex.load(partitionIndexFHBuilder, metadata.getLocal().partitioner, false);

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

@@ -88,7 +88,7 @@ public class PartitionIndex implements SharedCloseable
         this.root = trieRoot;
     }
 
-    private PartitionIndex(PartitionIndex src)
+    protected PartitionIndex(PartitionIndex src)
     {
         this(src.fh, src.root, src.keyCount, src.first, src.last);
     }
@@ -439,7 +439,7 @@ public class PartitionIndex implements SharedCloseable
         }
         catch (Throwable t)
         {
-            logger.warn("Failed to dump trie to {} due to exception {}", fileName, t);
+            logger.warn("Failed to dump trie to {} due to exception", fileName, t);
         }
     }
 

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

@@ -116,7 +116,7 @@ class PartitionIndexBuilder implements AutoCloseable
         }
         finally
         {
-            fhBuilder.withLengthOverride(-1);
+            fhBuilder.withLengthOverride(FileHandle.Builder.NO_LENGTH_OVERRIDE);
         }
 
     }

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

@@ -42,6 +42,19 @@ class PartitionIndexEarly extends PartitionIndex
         this.tail = tail;
     }
 
+    protected PartitionIndexEarly(PartitionIndexEarly partitionIndexEarly)
+    {
+        super(partitionIndexEarly);
+        this.cutoff = partitionIndexEarly.cutoff;
+        this.tail = partitionIndexEarly.tail;
+    }
+
+    @Override
+    public PartitionIndex sharedCopy()
+    {
+        return new PartitionIndexEarly(this);
+    }
+
     @Override
     protected Rebufferer instantiateRebufferer()
     {

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

@@ -83,7 +83,7 @@ class PartitionIterator extends PartitionIndex.IndexPosIterator implements KeyRe
             partitionIterator.advance();
             return partitionIterator;
         }
-        catch (IOException | RuntimeException ex)
+        catch (Throwable ex)
         {
             if (partitionIterator != null)
             {

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

@@ -75,7 +75,7 @@ public class Walker<CONCRETE extends Walker<CONCRETE>> implements AutoCloseable
             bh = source.rebuffer(root);
             buf = bh.buffer();
         }
-        catch (RuntimeException ex)
+        catch (Throwable ex)
         {
             if (bh != null) bh.release();
             source.closeReader();

test/unit/org/apache/cassandra/cql3/EarlyOpenCompactionTest.java

@@ -0,0 +1,225 @@
+/*
+ * 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.cql3;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Random;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import org.junit.After;
+import org.junit.Assert;
+import org.junit.Test;
+
+import org.apache.cassandra.config.DatabaseDescriptor;
+import org.apache.cassandra.db.ColumnFamilyStore;
+import org.hamcrest.Matchers;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+public class EarlyOpenCompactionTest extends CQLTester
+{
+    private static final int NUM_PARTITIONS = 1000;
+    private static final int NUM_ROWS_PER_PARTITION = 100;
+    private static final int VALUE_SIZE = 1000; // ~1KB per row
+    private static final int VERIFICATION_THREADS = 4;
+
+    private final AtomicBoolean stopVerification = new AtomicBoolean(false);
+    private final AtomicInteger verificationErrors = new AtomicInteger(0);
+    private final Random random = new Random();
+    private ExecutorService executor;
+
+    @After
+    public void cleanupAfter() throws Throwable
+    {
+        stopVerification.set(true);
+        if (executor != null)
+        {
+            executor.shutdownNow();
+            executor.awaitTermination(1, TimeUnit.MINUTES);
+        }
+        DatabaseDescriptor.setSSTablePreemptiveOpenIntervalInMiB(50);
+    }
+
+    @Test
+    public void testEarlyOpenDuringCompaction() throws Throwable
+    {
+        // Create a table with a simple schema
+        createTable("CREATE TABLE %s (" +
+                   "pk int, " +
+                   "ck int, " +
+                   "data text, " +
+                   "PRIMARY KEY (pk, ck)" +
+                   ")");
+
+        ColumnFamilyStore cfs = getCurrentColumnFamilyStore();
+        disableCompaction();
+        
+        // Insert data to create multiple SSTables
+        System.out.println("Inserting test data...");
+        for (int i = 0; i < NUM_PARTITIONS; i++)
+        {
+            for (int j = 0; j < NUM_ROWS_PER_PARTITION; j++)
+            {
+                String value = randomString(VALUE_SIZE);
+                execute("INSERT INTO %s (pk, ck, data) VALUES (?, ?, ?)", i, j, value);
+            }
+            
+            // Flush from time to time to get 10 sstables
+            if (i > 0 && i % Math.max(1, NUM_PARTITIONS / 10) == 0)
+            {
+                flush();
+            }
+        }
+        
+        // Final flush to ensure all data is written
+        flush();
+        
+        // Verify we have multiple SSTables
+        int sstableCount = cfs.getLiveSSTables().size();
+        assertTrue("Expected multiple SSTables, got: " + sstableCount, sstableCount > 1);
+        
+        // Start verification threads
+        System.out.println("Starting verification threads...");
+        executor = Executors.newFixedThreadPool(VERIFICATION_THREADS);
+        List<Future<?>> futures = new ArrayList<>();
+        
+        for (int i = 0; i < VERIFICATION_THREADS; i++)
+        {
+            futures.add(executor.submit(new VerificationTask()));
+        }
+        
+        // Wait a bit to ensure verification is running
+        Thread.sleep(1000);
+        
+        // Set early open interval to 1MiB to trigger early open during compaction
+        System.out.println("Setting early open interval to 1MiB...");
+        DatabaseDescriptor.setSSTablePreemptiveOpenIntervalInMiB(1);
+        // Slow down compaction to give the verifier time to fail.
+        DatabaseDescriptor.setCompactionThroughputMebibytesPerSec(10);
+        
+        // Trigger compaction and await its completion
+        System.out.println("Starting compaction...");
+        cfs.enableAutoCompaction(true);
+
+        // Let verification run for a while during and after compaction
+        System.out.println("Verifying data during and after compaction...");
+        Thread.sleep(1000);
+        
+        // Stop verification
+        stopVerification.set(true);
+        
+        // Wait for verification to complete
+        for (Future<?> future : futures)
+        {
+            try
+            {
+                future.get(10, TimeUnit.SECONDS);
+            }
+            catch (Exception e)
+            {
+                System.err.println("Verification task failed: " + e);
+                e.printStackTrace();
+            }
+        }
+        
+        // Verify no errors occurred during verification
+        int errors = verificationErrors.get();
+        assertEquals("Found " + errors + " verification errors. Check logs for details.", 0, errors);
+        
+        System.out.println("Test completed successfully");
+    }
+    
+    private class VerificationTask implements Runnable
+    {
+        @Override
+        public void run()
+        {
+            try
+            {
+                Random localRandom = new Random(Thread.currentThread().getId());
+                
+                while (!stopVerification.get() && !Thread.currentThread().isInterrupted())
+                {
+                    // Randomly choose between point query and partition range query
+                    if (localRandom.nextBoolean())
+                    {
+                        // Point query
+                        int pk = localRandom.nextInt(NUM_PARTITIONS * 110 / 100); // 10% chance outside
+                        int ck = localRandom.nextInt(NUM_ROWS_PER_PARTITION * 110 / 100); // 10% chance outside
+                        
+                        try
+                        {
+                            Assert.assertEquals(pk < NUM_PARTITIONS && ck < NUM_ROWS_PER_PARTITION ? 1 : 0,
+                                                execute("SELECT data FROM %s WHERE pk = ? AND ck = ?", pk, ck).size());
+                        }
+                        catch (Throwable t)
+                        {
+                            verificationErrors.incrementAndGet();
+                            System.err.println("Point query failed for pk=" + pk + ", ck=" + ck + ": " + t);
+                            t.printStackTrace();
+                        }
+                    }
+                    else
+                    {
+                        // Partition range query
+                        int pk = localRandom.nextInt(NUM_PARTITIONS);
+                        
+                        try
+                        {
+                            Assert.assertThat(execute("SELECT data FROM %s WHERE token(pk) <= token(?) AND token(pk) >= token(?)", pk, pk).size(),
+                                              Matchers.greaterThanOrEqualTo(NUM_ROWS_PER_PARTITION));
+                        }
+                        catch (Throwable t)
+                        {
+                            verificationErrors.incrementAndGet();
+                            System.err.println("Range query failed for pk in (" + pk + ", " + (pk + 1) + ", " + (pk + 2) + "): " + t);
+                            t.printStackTrace();
+                        }
+                    }
+                    
+                    // Add a small delay to prevent overwhelming the system
+                    Thread.yield();
+                }
+            }
+            catch (Throwable t)
+            {
+                verificationErrors.incrementAndGet();
+                System.err.println("Verification task failed: " + t);
+                t.printStackTrace();
+            }
+        }
+    }
+    
+    private String randomString(int length)
+    {
+        StringBuilder sb = new StringBuilder(length);
+        for (int i = 0; i < length; i++)
+        {
+            sb.append((char)('a' + random.nextInt(26)));
+        }
+        return sb.toString();
+    }
+}