content: /*
   * 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.db.compaction;
  
  import java.util.Iterator;
  import java.util.Map;
  
  import org.junit.After;
  import org.junit.Before;
  import org.junit.BeforeClass;
  import org.junit.Test;
  
  import org.apache.cassandra.SchemaLoader;
  import org.apache.cassandra.Util;
  import org.apache.cassandra.config.DatabaseDescriptor;
  import org.apache.cassandra.db.ColumnFamilyStore;
  import org.apache.cassandra.db.Directories;
  import org.apache.cassandra.db.DirectoriesTest;
  import org.apache.cassandra.db.Keyspace;
  import org.apache.cassandra.db.RowUpdateBuilder;
  import org.apache.cassandra.db.lifecycle.LifecycleTransaction;
  import org.apache.cassandra.io.sstable.format.SSTableReader;
  import org.apache.cassandra.io.util.File;
  import org.apache.cassandra.schema.KeyspaceParams;
  import org.apache.cassandra.schema.TableMetadataRef;
  import org.apache.cassandra.utils.CloseableIterator;
  import org.apache.cassandra.utils.FBUtilities;
  
  import static org.hamcrest.CoreMatchers.instanceOf;
  import static org.hamcrest.MatcherAssert.assertThat;
  import static org.junit.Assert.assertEquals;
  
  public class PartialCompactionsTest extends SchemaLoader
  {
      static final String KEYSPACE = PartialCompactionsTest.class.getSimpleName();
      static final String TABLE = "testtable";
  
      @BeforeClass
      public static void initSchema()
      {
          SchemaLoader.prepareServer();
          SchemaLoader.createKeyspace(KEYSPACE,
                                      KeyspaceParams.simple(1),
                                      SchemaLoader.standardCFMD(KEYSPACE, TABLE));
          CompactionManager.instance.disableAutoCompaction();
  
          LimitableDataDirectory.applyTo(KEYSPACE, TABLE);
      }
  
      @Before
      public void prepareCFS()
      {
          LimitableDataDirectory.setAvailableSpace(cfStore(), null);
      }
  
      @After
      public void truncateCF()
      {
          cfStore().truncateBlocking();
          LifecycleTransaction.waitForDeletions();
      }
  
      private static ColumnFamilyStore cfStore()
      {
          return Keyspace.open(KEYSPACE).getColumnFamilyStore(TABLE);
      }
  
      @Test
      public void shouldNotResurrectDataFromSSTableExcludedDueToInsufficientSpace()
      {
          // given
          ColumnFamilyStore cfs = cfStore();
          int few = 10, many = 10 * few;
  
          // a large sstable as the oldest
          createDataSSTable(cfs, 0, many);
          // more inserts (to have more than one sstable to compact)
          createDataSSTable(cfs, many, many + few);
          // delete data that's in both of the prior sstables
          createTombstonesSSTable(cfs, many - few / 2, many + few / 2);
  
          // emulate there not being enough space to compact all sstables
          LimitableDataDirectory.setAvailableSpace(cfs, enoughSpaceForAllButTheLargestSSTable(cfs));
  
          // when - run a compaction where all tombstones have timed out
          FBUtilities.waitOnFutures(CompactionManager.instance.submitMaximal(cfs, Integer.MAX_VALUE, false, 0));
  
          // then - the tombstones should not be removed
          assertEquals("live sstables after compaction", 2, cfs.getLiveSSTables().size());
          assertEquals("remaining live rows after compaction", many, liveRows(cfs));
      }
  
      private static long enoughSpaceForAllButTheLargestSSTable(ColumnFamilyStore cfs)
      {
          long totalSize = 1, maxSize = 0;
          for (SSTableReader ssTable : cfs.getLiveSSTables())
          {
              long size = ssTable.onDiskLength();
              if (size > maxSize) maxSize = size;
              totalSize += size;
          }
          return totalSize - maxSize;
      }
  
      private static int liveRows(ColumnFamilyStore cfs)
      {
          return Util.getAll(Util.cmd(cfs, "key1").build()).stream()
                     .map(partition -> count(partition.rowIterator(false)))
                     .reduce(Integer::sum)
                     .orElse(0);
      }
  
      private static int count(Iterator<?> iter)
      {
          try (CloseableIterator<?> unused = iter instanceof CloseableIterator ? (CloseableIterator<?>) iter : null)
          {
              int count = 0;
              for (; iter.hasNext(); iter.next())
              {
                  count++;
              }
              return count;
          }
      }
  
      private static void createDataSSTable(ColumnFamilyStore cfs, int firstKey, int endKey)
      {
          for (int i = firstKey; i < endKey; i++)
          {
              new RowUpdateBuilder(cfs.metadata(), 0, "key1")
              .clustering(String.valueOf(i))
              .add("val", String.valueOf(i))
              .build()
              .applyUnsafe();
          }
          cfs.forceBlockingFlush(ColumnFamilyStore.FlushReason.UNIT_TESTS);
      }
  
      private static void createTombstonesSSTable(ColumnFamilyStore cfs, int firstKey, int endKey)
      {
          for (int i = firstKey; i < endKey; i++)
          {
              RowUpdateBuilder.deleteRow(cfs.metadata(), 1, "key1", String.valueOf(i)).applyUnsafe();
          }
          cfs.forceBlockingFlush(ColumnFamilyStore.FlushReason.UNIT_TESTS);
      }
  
      private static class LimitableDataDirectory extends Directories.DataDirectory
      {
          private Long availableSpace;
  
          LimitableDataDirectory(Directories.DataDirectory dataDirectory)
          {
              super(dataDirectory.location);
          }
  
          @Override
          public long getAvailableSpace()
          {
              if (availableSpace != null)
                  return availableSpace;
              return super.getAvailableSpace();
          }
  
          public static void setAvailableSpace(ColumnFamilyStore cfs, Long availableSpace)
          {
              for (Directories.DataDirectory location : cfs.getDirectories().getWriteableLocations())
              {
                  assertThat("ColumnFamilyStore set up with ability to emulate limited disk space",
                             location, instanceOf(LimitableDataDirectory.class));
                  ((LimitableDataDirectory) location).availableSpace = availableSpace;
              }
          }
  
          public static void applyTo(String ks, String cf)
          {
              Keyspace keyspace = Keyspace.open(ks);
              ColumnFamilyStore store = keyspace.getColumnFamilyStore(cf);
              TableMetadataRef metadata = store.metadata;
              keyspace.dropCf(metadata.id, true);
              ColumnFamilyStore cfs = ColumnFamilyStore.createColumnFamilyStore(keyspace, cf, metadata.get(), wrapDirectoriesOf(store), false, false);
              keyspace.initCfCustom(cfs);
          }
  
          private static Directories wrapDirectoriesOf(ColumnFamilyStore cfs)
          {
              Directories.DataDirectory[] original = cfs.getDirectories().getWriteableLocations();
              Directories.DataDirectory[] wrapped = new Directories.DataDirectory[original.length];
              for (int i = 0; i < wrapped.length; i++)
              {
                  wrapped[i] = new LimitableDataDirectory(original[i]);
              }
              return new Directories(cfs.metadata(), wrapped)
              {
                  @Override
                  public boolean hasDiskSpaceForCompactionsAndStreams(Map<File, Long> expectedNewWriteSizes, Map<File, Long> totalCompactionWriteRemaining)
                  {
                      return hasDiskSpaceForCompactionsAndStreams(expectedNewWriteSizes, totalCompactionWriteRemaining, file -> {
                          for (DataDirectory location : getWriteableLocations())
                          {
                              if (file.toPath().startsWith(location.location.toPath())) {
                                  LimitableDataDirectory directory = (LimitableDataDirectory) location;
                                  if (directory.availableSpace != null)
                                  {
                                      DirectoriesTest.FakeFileStore store = new DirectoriesTest.FakeFileStore();
                                      // reverse the computation in Directories.getAvailableSpaceForCompactions
                                      store.usableSpace = Math.round(directory.availableSpace / DatabaseDescriptor.getMaxSpaceForCompactionsPerDrive()) + DatabaseDescriptor.getMinFreeSpacePerDriveInBytes();
                                      return store;
                                  }
                              }
                          }
                          return Directories.getFileStore(file);
                      });
                  }
              };
          }
      }
  }
  