Merge branch 'develop' into nziebart/support-persistable

atlasdb-api/src/main/java/com/palantir/atlasdb/spi/KeyValueServiceConfig.java

@@ -42,4 +42,5 @@ public interface KeyValueServiceConfig {
     default int defaultGetRangesConcurrency() {
         return Math.min(8, concurrentGetRangesThreadPoolSize() / 2);
     }
+
 }

atlasdb-cassandra-integration-tests/src/test/java/com/palantir/atlasdb/keyvalue/cassandra/CassandraGetCandidateCellsForSweepingTest.java

@@ -16,13 +16,18 @@
 
 package com.palantir.atlasdb.keyvalue.cassandra;
 
+import static org.assertj.core.api.Assertions.assertThat;
+
 import org.junit.ClassRule;
+import org.junit.Test;
 import org.mockito.Mockito;
 import org.slf4j.Logger;
 
 import com.palantir.atlasdb.cassandra.CassandraKeyValueServiceConfigManager;
 import com.palantir.atlasdb.containers.CassandraContainer;
 import com.palantir.atlasdb.containers.Containers;
+import com.palantir.atlasdb.encoding.PtBytes;
+import com.palantir.atlasdb.keyvalue.api.ImmutableCandidateCellForSweeping;
 import com.palantir.atlasdb.keyvalue.api.KeyValueService;
 import com.palantir.atlasdb.keyvalue.impl.AbstractGetCandidateCellsForSweepingTest;
 
@@ -38,4 +43,29 @@ protected KeyValueService createKeyValueService() {
                 CassandraContainer.LEADER_CONFIG,
                 Mockito.mock(Logger.class));
     }
+
+    @Test
+    public void returnCandidateIfPossiblyUncommittedTimestamp() {
+        new TestDataBuilder().put(1, 1, 10L).store();
+        assertThat(getAllCandidates(conservativeRequest(PtBytes.EMPTY_BYTE_ARRAY, 40L, 1)))
+                .containsExactly(ImmutableCandidateCellForSweeping.builder()
+                        .cell(cell(1, 1))
+                        .sortedTimestamps(new long[] { 10L })
+                        .isLatestValueEmpty(false)
+                        .numCellsTsPairsExamined(1)
+                        .build());
+    }
+
+    @Test
+    public void returnCandidateIfTwoCommittedTimestamps() {
+        new TestDataBuilder().put(1, 1, 10L).put(1, 1, 20L).store();
+        assertThat(getAllCandidates(conservativeRequest(PtBytes.EMPTY_BYTE_ARRAY, 40L, 1)))
+                .containsExactly(ImmutableCandidateCellForSweeping.builder()
+                        .cell(cell(1, 1))
+                        .sortedTimestamps(new long[] { 10L, 20L })
+                        .isLatestValueEmpty(false)
+                        .numCellsTsPairsExamined(2)
+                        .build());
+    }
+
 }

atlasdb-cassandra-integration-tests/src/test/java/com/palantir/atlasdb/keyvalue/cassandra/CassandraKeyValueServiceSweepTaskRunnerIntegrationTest.java

@@ -15,29 +15,65 @@
  */
 package com.palantir.atlasdb.keyvalue.cassandra;
 
+import java.util.Arrays;
+
+import org.apache.commons.lang3.RandomStringUtils;
 import org.junit.Assert;
+import org.junit.Assume;
 import org.junit.ClassRule;
 import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
 
+import com.palantir.atlasdb.cassandra.CassandraKeyValueServiceConfig;
 import com.palantir.atlasdb.cassandra.CassandraKeyValueServiceConfigManager;
+import com.palantir.atlasdb.cassandra.ImmutableCassandraKeyValueServiceConfig;
 import com.palantir.atlasdb.containers.CassandraContainer;
 import com.palantir.atlasdb.containers.Containers;
 import com.palantir.atlasdb.keyvalue.api.KeyValueService;
 import com.palantir.atlasdb.keyvalue.api.SweepResults;
 import com.palantir.atlasdb.protos.generated.TableMetadataPersistence;
 import com.palantir.atlasdb.sweep.AbstractSweepTaskRunnerTest;
 
+@RunWith(Parameterized.class)
 public class CassandraKeyValueServiceSweepTaskRunnerIntegrationTest extends AbstractSweepTaskRunnerTest {
     @ClassRule
     public static final Containers CONTAINERS = new Containers(
                 CassandraKeyValueServiceSweepTaskRunnerIntegrationTest.class)
             .with(new CassandraContainer());
 
+    @Parameterized.Parameter
+    public boolean useColumnBatchSize;
+
+    @Parameterized.Parameters(name = "Use column batch size parameter = {0}")
+    public static Iterable<?> parameters() {
+        return Arrays.asList(true, false);
+    }
+
+
     @Override
     protected KeyValueService getKeyValueService() {
+        CassandraKeyValueServiceConfig config = useColumnBatchSize
+                ? ImmutableCassandraKeyValueServiceConfig.copyOf(CassandraContainer.KVS_CONFIG)
+                        .withTimestampsGetterBatchSize(10)
+                : CassandraContainer.KVS_CONFIG;
+
         return CassandraKeyValueServiceImpl.create(
-                CassandraKeyValueServiceConfigManager.createSimpleManager(
-                        CassandraContainer.KVS_CONFIG), CassandraContainer.LEADER_CONFIG);
+                CassandraKeyValueServiceConfigManager.createSimpleManager(config), CassandraContainer.LEADER_CONFIG);
+    }
+
+    @Test
+    public void should_not_oom_when_there_are_many_large_values_to_sweep() {
+        Assume.assumeTrue("should_not_oom test will always fail if column batch size is not set!", useColumnBatchSize);
+
+        createTable(TableMetadataPersistence.SweepStrategy.CONSERVATIVE);
+
+        long numInsertions = 100;
+        insertMultipleValues(numInsertions);
+
+        long sweepTimestamp = numInsertions + 1;
+        SweepResults results = completeSweep(sweepTimestamp);
+        Assert.assertEquals(numInsertions - 1, results.getStaleValuesDeleted());
     }
 
     @Test
@@ -53,4 +89,14 @@ public void should_return_values_for_multiple_columns_when_sweeping() {
         Assert.assertEquals(28, results.getStaleValuesDeleted());
     }
 
+    private void insertMultipleValues(long numInsertions) {
+        for (int ts = 1; ts <= numInsertions; ts++) {
+            System.out.println("putting with ts = " + ts);
+            putIntoDefaultColumn("row", makeLongRandomString(), ts);
+        }
+    }
+
+    private String makeLongRandomString() {
+        return RandomStringUtils.random(1_000_000);
+    }
 }

atlasdb-cassandra/src/main/java/com/palantir/atlasdb/keyvalue/cassandra/CassandraKeyValueServiceImpl.java

@@ -107,19 +107,16 @@
 import com.palantir.atlasdb.keyvalue.cassandra.jmx.CassandraJmxCompaction;
 import com.palantir.atlasdb.keyvalue.cassandra.jmx.CassandraJmxCompactionManager;
 import com.palantir.atlasdb.keyvalue.cassandra.paging.CassandraRangePagingIterable;
-import com.palantir.atlasdb.keyvalue.cassandra.paging.CellPager;
-import com.palantir.atlasdb.keyvalue.cassandra.paging.CellPagerBatchSizingStrategy;
 import com.palantir.atlasdb.keyvalue.cassandra.paging.ColumnGetter;
 import com.palantir.atlasdb.keyvalue.cassandra.paging.CqlColumnGetter;
-import com.palantir.atlasdb.keyvalue.cassandra.paging.RowRangeLoader;
-import com.palantir.atlasdb.keyvalue.cassandra.paging.SingleRowColumnPager;
+import com.palantir.atlasdb.keyvalue.cassandra.paging.RowGetter;
 import com.palantir.atlasdb.keyvalue.cassandra.paging.ThriftColumnGetter;
-import com.palantir.atlasdb.keyvalue.cassandra.sweep.CassandraGetCandidateCellsForSweepingImpl;
 import com.palantir.atlasdb.keyvalue.cassandra.thrift.SlicePredicates;
 import com.palantir.atlasdb.keyvalue.cassandra.thrift.SlicePredicates.Limit;
 import com.palantir.atlasdb.keyvalue.cassandra.thrift.SlicePredicates.Range;
 import com.palantir.atlasdb.keyvalue.impl.AbstractKeyValueService;
 import com.palantir.atlasdb.keyvalue.impl.Cells;
+import com.palantir.atlasdb.keyvalue.impl.GetCandidateCellsForSweepingShim;
 import com.palantir.atlasdb.keyvalue.impl.KeyValueServices;
 import com.palantir.atlasdb.keyvalue.impl.LocalRowColumnRangeIterator;
 import com.palantir.atlasdb.logging.LoggingArgs;
@@ -196,9 +193,7 @@ public void close() {
     protected final CassandraKeyValueServiceConfigManager configManager;
 
     private final Optional<CassandraJmxCompactionManager> compactionManager;
-
-    @VisibleForTesting
-    final CassandraClientPool clientPool;
+    private final CassandraClientPool clientPool;
 
     private SchemaMutationLock schemaMutationLock;
     private final Optional<LeaderConfig> leaderConfig;
@@ -214,8 +209,6 @@ public void close() {
     private final TracingQueryRunner queryRunner;
     private final CassandraTables cassandraTables;
 
-    private final CassandraGetCandidateCellsForSweepingImpl getCandidateCellsForSweepingImpl;
-
     private final InitializingWrapper wrapper = new InitializingWrapper();
 
     public static CassandraKeyValueService create(
@@ -272,11 +265,6 @@ protected CassandraKeyValueServiceImpl(Logger log,
 
         this.queryRunner = new TracingQueryRunner(log, tracingPrefs);
         this.cassandraTables = new CassandraTables(clientPool, configManager);
-
-        SingleRowColumnPager singleRowPager = new SingleRowColumnPager(clientPool, queryRunner);
-        CellPager cellPager = new CellPager(
-                singleRowPager, clientPool, queryRunner, new CellPagerBatchSizingStrategy());
-        this.getCandidateCellsForSweepingImpl = new CassandraGetCandidateCellsForSweepingImpl(cellPager);
     }
 
     @Override
@@ -1452,8 +1440,7 @@ public ClosableIterator<RowResult<Set<Long>>> getRangeOfTimestamps(
     @Override
     public ClosableIterator<List<CandidateCellForSweeping>> getCandidateCellsForSweeping(TableReference tableRef,
             CandidateCellForSweepingRequest request) {
-        return ClosableIterators.wrap(
-                getCandidateCellsForSweepingImpl.getCandidateCellsForSweeping(tableRef, request, deleteConsistency));
+        return new GetCandidateCellsForSweepingShim(this).getCandidateCellsForSweeping(tableRef, request);
     }
 
     private ClosableIterator<RowResult<Set<Long>>> getTimestampsInBatchesWithPageCreator(
@@ -1463,18 +1450,12 @@ private ClosableIterator<RowResult<Set<Long>>> getTimestampsInBatchesWithPageCre
             long timestamp,
             ConsistencyLevel consistency) {
         SlicePredicate predicate = SlicePredicates.create(Range.ALL, Limit.ONE);
-
-        RowRangeLoader rowRangeLoader = new RowRangeLoader(clientPool, queryRunner, consistency, tableRef);
+        RowGetter rowGetter = new RowGetter(clientPool, queryRunner, consistency, tableRef);
 
         CqlExecutor cqlExecutor = new CqlExecutor(clientPool, consistency);
         ColumnGetter columnGetter = new CqlColumnGetter(cqlExecutor, tableRef, columnBatchSize);
 
-        return getRangeWithPageCreator(
-                rowRangeLoader,
-                predicate,
-                columnGetter,
-                rangeRequest,
-                TimestampExtractor::new,
+        return getRangeWithPageCreator(rowGetter, predicate, columnGetter, rangeRequest, TimestampExtractor::new,
                 timestamp);
     }
 
@@ -1493,15 +1474,14 @@ private <T> ClosableIterator<RowResult<T>> getRangeWithPageCreator(
             // each column. note that if no columns are specified, it's a special case that means all columns
             predicate = SlicePredicates.create(Range.ALL, Limit.NO_LIMIT);
         }
-        RowRangeLoader rowRangeLoader = new RowRangeLoader(clientPool, queryRunner, consistency, tableRef);
+        RowGetter rowGetter = new RowGetter(clientPool, queryRunner, consistency, tableRef);
         ColumnGetter columnGetter = new ThriftColumnGetter();
 
-        return getRangeWithPageCreator(rowRangeLoader, predicate, columnGetter, rangeRequest, resultsExtractor,
-                startTs);
+        return getRangeWithPageCreator(rowGetter, predicate, columnGetter, rangeRequest, resultsExtractor, startTs);
     }
 
     private <T> ClosableIterator<RowResult<T>> getRangeWithPageCreator(
-            RowRangeLoader rowRangeLoader,
+            RowGetter rowGetter,
             SlicePredicate slicePredicate,
             ColumnGetter columnGetter,
             RangeRequest rangeRequest,
@@ -1515,7 +1495,7 @@ private <T> ClosableIterator<RowResult<T>> getRangeWithPageCreator(
         }
 
         CassandraRangePagingIterable<T> rowResults = new CassandraRangePagingIterable<>(
-                rowRangeLoader,
+                rowGetter,
                 slicePredicate,
                 columnGetter,
                 rangeRequest,