fix potential infinite loop in get_count; patch by yukim reviewed by Tyler Hobbs for CASSANDRA-4833

CHANGES.txt

@@ -5,6 +5,7 @@
    (CASSANDRA-4765)
  * fix wrong leveled compaction progress calculation (CASSANDRA-4807)
  * add a close() method to CRAR to prevent leaking file descriptors (CASSANDRA-4820)
+ * fix potential infinite loop in get_count (CASSANDRA-4833)
 
 1.1.6
  * Wait for writes on synchronous read digest mismatch (CASSANDRA-4792)

src/java/org/apache/cassandra/thrift/CassandraServer.java

@@ -428,25 +428,28 @@ public class CassandraServer implements Cassandra.Iface
                                                    Integer.MAX_VALUE);
         }
 
-        int requestedCount = predicate.slice_range.count;
+        final int requestedCount = predicate.slice_range.count;
+        int remaining = requestedCount;
         int pages = 0;
         while (true)
         {
-            predicate.slice_range.count = Math.min(pageSize, requestedCount);
+            predicate.slice_range.count = Math.min(pageSize, Math.max(2, remaining)); // fetch at least two columns
             columns = get_slice(key, column_parent, predicate, consistency_level);
             if (columns.isEmpty())
                 break;
 
-            ColumnOrSuperColumn firstColumn = columns.get(columns.size() - 1);
             ByteBuffer firstName = getName(columns.get(0));
             int newColumns = pages == 0 || !firstName.equals(predicate.slice_range.start) ? columns.size() : columns.size() - 1;
             totalCount += newColumns;
-            requestedCount -= newColumns;
+            // if we over-counted, just return original limit
+            if (totalCount > requestedCount)
+                return requestedCount;
+            remaining -= newColumns;
             pages++;
             // We're done if either:
             //   - We've querying the number of columns requested by the user
             //   - The last page wasn't full
-            if (requestedCount == 0 || columns.size() < predicate.slice_range.count)
+            if (remaining == 0 || columns.size() < predicate.slice_range.count)
                 break;
             else
                 predicate.slice_range.start = getName(columns.get(columns.size() - 1));

test/unit/org/apache/cassandra/service/CassandraServerTest.java

@@ -21,40 +21,68 @@ package org.apache.cassandra.service;
 import org.junit.Test;
 
 import org.apache.cassandra.SchemaLoader;
+import org.apache.cassandra.Util;
+import org.apache.cassandra.config.Schema;
+import org.apache.cassandra.db.DecoratedKey;
+import org.apache.cassandra.db.RowMutation;
+import org.apache.cassandra.db.filter.QueryPath;
+import org.apache.cassandra.thrift.*;
+import org.apache.cassandra.utils.ByteBufferUtil;
 
 public class CassandraServerTest extends SchemaLoader
 {
+    /**
+     * test get_count() to work correctly with 'count' settings around page size.
+     * (CASSANDRA-4833)
+     */
     @Test
-    public void test_get_column() throws Throwable {
-        /*
-        CassandraServer server = new CassandraServer();
-        server.start();
+    public void test_get_count() throws Exception
+    {
+        Schema.instance.clear(); // Schema are now written on disk and will be reloaded
+        new EmbeddedCassandraService().start();
 
-        try {
-            Column c1 = column("c1", "0", 0L);
-            Column c2 = column("c2", "0", 0L);
-            List<Column> columns = new ArrayList<Column>();
-            columns.add(c1);
-            columns.add(c2);
-            Map<String, List<Column>> cfmap = new HashMap<String, List<Column>>();
-            cfmap.put("Standard1", columns);
-            cfmap.put("Standard2", columns);
-
-            BatchMutation m = new BatchMutation("Keyspace1", "key1", cfmap);
-            server.batch_insert(m, 1);
-
-            Column column;
-            column = server.get_column("Keyspace1", "key1", "Standard1:c2");
-            assert column.value.equals("0");
-
-            column = server.get_column("Keyspace1", "key1", "Standard2:c2");
-            assert column.value.equals("0");
-
-            ArrayList<Column> Columns = server.get_slice_strong("Keyspace1", "key1", "Standard1", -1, -1);
-            assert Columns.size() == 2;
-        } finally {
-            server.shutdown();
+        DecoratedKey key = Util.dk("testkey");
+        for (int i = 0; i < 3050; i++)
+        {
+            RowMutation rm = new RowMutation("Keyspace1", key.key);
+            rm.add(new QueryPath("Standard1", null, ByteBufferUtil.bytes(String.valueOf(i))),
+                          ByteBufferUtil.EMPTY_BYTE_BUFFER,
+                          System.currentTimeMillis());
+            rm.apply();
         }
-        */
+
+        CassandraServer server = new CassandraServer();
+        server.set_keyspace("Keyspace1");
+
+        // same as page size
+        int count = server.get_count(key.key, new ColumnParent("Standard1"), predicateWithCount(1024), ConsistencyLevel.ONE);
+        assert count == 1024 : "expected 1024 but was " + count;
+
+        // 1 above page size
+        count = server.get_count(key.key, new ColumnParent("Standard1"), predicateWithCount(1025), ConsistencyLevel.ONE);
+        assert count == 1025 : "expected 1025 but was " + count;
+
+        // above number of columns
+        count = server.get_count(key.key, new ColumnParent("Standard1"), predicateWithCount(4000), ConsistencyLevel.ONE);
+        assert count == 3050 : "expected 3050 but was " + count;
+
+        // same as number of columns
+        count = server.get_count(key.key, new ColumnParent("Standard1"), predicateWithCount(3050), ConsistencyLevel.ONE);
+        assert count == 3050 : "expected 3050 but was " + count;
+
+        // 1 above number of columns
+        count = server.get_count(key.key, new ColumnParent("Standard1"), predicateWithCount(3051), ConsistencyLevel.ONE);
+        assert count == 3050 : "expected 3050 but was " + count;
+    }
+
+    private SlicePredicate predicateWithCount(int count)
+    {
+        SliceRange range = new SliceRange();
+        range.setStart("".getBytes());
+        range.setFinish("".getBytes());
+        range.setCount(count);
+        SlicePredicate predicate = new SlicePredicate();
+        predicate.setSlice_range(range);
+        return predicate;
     }
 }