From 95fb613bf996c715392f9aa1f491609b6acaeff5 Mon Sep 17 00:00:00 2001 From: Yuki Morishita Date: Mon, 22 Oct 2012 12:16:20 -0500 Subject: [PATCH] fix potential infinite loop in get_count; patch by yukim reviewed by Tyler Hobbs for CASSANDRA-4833 --- CHANGES.txt | 1 + .../cassandra/thrift/CassandraServer.java | 13 +-- .../service/CassandraServerTest.java | 86 ++++++++++++------- 3 files changed, 66 insertions(+), 34 deletions(-) diff --git a/CHANGES.txt b/CHANGES.txt index 8822c3bace..f309ef1b4d 100644 --- a/CHANGES.txt +++ b/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) diff --git a/src/java/org/apache/cassandra/thrift/CassandraServer.java b/src/java/org/apache/cassandra/thrift/CassandraServer.java index 3bf155e1b8..39b57f372f 100644 --- a/src/java/org/apache/cassandra/thrift/CassandraServer.java +++ b/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)); diff --git a/test/unit/org/apache/cassandra/service/CassandraServerTest.java b/test/unit/org/apache/cassandra/service/CassandraServerTest.java index 495d6971b0..48701de6f1 100644 --- a/test/unit/org/apache/cassandra/service/CassandraServerTest.java +++ b/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 columns = new ArrayList(); - columns.add(c1); - columns.add(c2); - Map> cfmap = new HashMap>(); - 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 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; } }