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.stress;
  
  import java.io.File;
  import java.io.IOError;
  import java.net.URI;
  import java.util.ArrayList;
  import java.util.Collections;
  import java.util.List;
  import java.util.Map;
  import java.util.Random;
  import java.util.Set;
  import java.util.concurrent.CountDownLatch;
  import java.util.concurrent.ExecutorService;
  import java.util.concurrent.Executors;
  import java.util.concurrent.TimeUnit;
  
  import com.google.common.collect.Lists;
  import com.google.common.util.concurrent.Uninterruptibles;
  
  import org.apache.cassandra.config.DatabaseDescriptor;
  import org.apache.cassandra.cql3.statements.schema.CreateTableStatement;
  import org.apache.cassandra.db.ColumnFamilyStore;
  import org.apache.cassandra.db.Directories;
  import org.apache.cassandra.db.Keyspace;
  import org.apache.cassandra.db.commitlog.CommitLog;
  import org.apache.cassandra.db.compaction.CompactionManager;
  import org.apache.cassandra.db.lifecycle.LifecycleTransaction;
  import org.apache.cassandra.dht.IPartitioner;
  import org.apache.cassandra.dht.Token;
  import org.apache.cassandra.io.sstable.Component;
  import org.apache.cassandra.io.sstable.Descriptor;
  import org.apache.cassandra.io.sstable.StressCQLSSTableWriter;
  import org.apache.cassandra.io.sstable.format.SSTableFormat.Components;
  import org.apache.cassandra.io.sstable.format.SSTableReader;
  import org.apache.cassandra.io.util.FileUtils;
  import org.apache.cassandra.locator.InetAddressAndPort;
  import org.apache.cassandra.stress.generate.PartitionGenerator;
  import org.apache.cassandra.stress.generate.SeedManager;
  import org.apache.cassandra.stress.operations.userdefined.SchemaInsert;
  import org.apache.cassandra.stress.settings.StressSettings;
  import org.apache.cassandra.tcm.ClusterMetadata;
  import org.apache.cassandra.tcm.ClusterMetadataService;
  import org.apache.cassandra.tools.nodetool.CompactionStats;
  import org.apache.cassandra.tools.nodetool.formatter.TableBuilder;
  import org.apache.cassandra.utils.FBUtilities;
  import org.apache.cassandra.utils.JVMStabilityInspector;
  import org.apache.cassandra.utils.concurrent.Future;
  import picocli.CommandLine;
  import picocli.CommandLine.Command;
  import picocli.CommandLine.Option;
  
  /**
   * Tool that allows fast route to loading data for arbitrary schemas to disk
   * and compacting them.
   */
  public abstract class CompactionStress implements Runnable
  {
      @Option(names = { "-p", "--profile" }, description = "Path to stress yaml file", required = true)
      String profile;
  
      @Option(names = { "-d", "--datadir" }, description = "Data directory (can be used many times to specify multiple data dirs)", required = true)
      List<String> dataDirs;
  
      @Option(names = { "-v", "--vnodes" }, description = "number of local tokens to generate (default 256)")
      Integer numTokens = 256;
  
      static
      {
          DatabaseDescriptor.daemonInitialization();
          CommitLog.instance.start();
      }
  
      List<File> getDataDirectories()
      {
          List<File> dataDirectories = new ArrayList<>(dataDirs.size());
          for (String dataDir : dataDirs)
          {
              File outputDir = new File(dataDir);
  
              if (!outputDir.exists())
              {
                  System.err.println("Invalid output dir (missing): " + outputDir);
                  System.exit(1);
              }
  
              if (!outputDir.isDirectory())
              {
                  System.err.println("Invalid output dir (not a directory): " + outputDir);
                  System.exit(2);
              }
  
              if (!outputDir.canWrite())
              {
                  System.err.println("Invalid output dir (no write permissions): " + outputDir);
                  System.exit(3);
              }
  
              dataDirectories.add(outputDir);
          }
  
          return dataDirectories;
      }
  
      ColumnFamilyStore initCf(StressProfile stressProfile, boolean loadSSTables)
      {
          generateTokens(stressProfile.seedStr, numTokens);
  
          CreateTableStatement.Raw createStatement = stressProfile.getCreateStatement();
          List<File> dataDirectories = getDataDirectories();
  
          ColumnFamilyStore cfs = StressCQLSSTableWriter.Builder.createOfflineTable(createStatement, Collections.EMPTY_LIST, dataDirectories);
  
          if (loadSSTables)
          {
              Directories.SSTableLister lister = cfs.getDirectories().sstableLister(Directories.OnTxnErr.IGNORE).skipTemporary(true);
              List<SSTableReader> sstables = new ArrayList<>();
  
              //Offline open sstables
              for (Map.Entry<Descriptor, Set<Component>> entry : lister.list().entrySet())
              {
                  Set<Component> components = entry.getValue();
                  if (!components.contains(Components.DATA))
                      continue;
  
                  try
                  {
                      SSTableReader sstable = SSTableReader.openNoValidation(entry.getKey(), components, cfs);
                      sstables.add(sstable);
                  }
                  catch (Exception e)
                  {
                      JVMStabilityInspector.inspectThrowable(e);
                      System.err.println(String.format("Error Loading %s: %s", entry.getKey(), e.getMessage()));
                  }
              }
  
              cfs.disableAutoCompaction();
  
              // We want to add the SSTables without firing their indexing by any eventual unsupported 2i
              if (cfs.indexManager.hasIndexes())
                  throw new IllegalStateException("CompactionStress does not support secondary indexes");
  
              //Register with cfs
              cfs.addSSTables(sstables);
          }
  
          return cfs;
      }
  
      StressProfile getStressProfile()
      {
          try
          {
              File yamlFile = new File(profile);
              return StressProfile.load(yamlFile.exists() ? yamlFile.toURI() : URI.create(profile));
          }
          catch ( IOError e)
          {
              e.printStackTrace();
              System.err.print("Invalid profile URI : " + profile);
              System.exit(4);
          }
  
          return null;
      }
  
      /**
       * Populate tokenMetadata consistently across runs.
       *
       * We need consistency to write and compact the same data offline
       * in the case of a range aware sstable writer.
       */
      private void generateTokens(String seed, Integer numTokens)
      {
          Random random = new Random(seed.hashCode());
  
          IPartitioner p = ClusterMetadata.current().tokenMap.partitioner();
         // tokenMetadata.clearUnsafe();
          for (int i = 1; i <= numTokens; i++)
          {
              InetAddressAndPort addr = FBUtilities.getBroadcastAddressAndPort();
              List<Token> tokens = Lists.newArrayListWithCapacity(numTokens);
              for (int j = 0; j < numTokens; ++j)
                  tokens.add(p.getRandomToken(random));
  
  //            tokenMetadata.updateNormalTokens(tokens, addr);
          }
      }
  
      @Command(name = "compact", description = "Compact data in directory")
      public static class Compaction extends CompactionStress
      {
  
          @Option(names = { "-m", "--maximal" }, description = "Force maximal compaction (default true)")
          Boolean maximal = false;
  
          @Option(names = { "-t", "--threads" }, description = "Number of compactor threads to use for bg compactions (default 4)")
          Integer threads = 4;
  
          public void run()
          {
              ClusterMetadataService.initializeForTools(true);
              //Setup
              CompactionManager.instance.setMaximumCompactorThreads(threads);
              CompactionManager.instance.setCoreCompactorThreads(threads);
              CompactionManager.instance.setRateInBytes(0);
  
              StressProfile stressProfile = getStressProfile();
              ColumnFamilyStore cfs = initCf(stressProfile, true);
              cfs.getCompactionStrategyManager().compactionLogger.enable();
  
              List<Future<?>> futures = new ArrayList<>(threads);
              if (maximal)
              {
                  futures = CompactionManager.instance.submitMaximal(cfs, FBUtilities.nowInSeconds(), false, Integer.MAX_VALUE);
              }
              else
              {
                  cfs.enableAutoCompaction();
                  cfs.getCompactionStrategyManager().enable();
                  for (int i = 0; i < threads; i++)
                      futures.addAll(CompactionManager.instance.submitBackground(cfs));
              }
  
              long working;
              //Report compaction stats while working
              while ((working = futures.stream().filter(f -> !f.isDone()).count()) > 0 || CompactionManager.instance.getActiveCompactions() > 0 || (!maximal && cfs.getCompactionStrategyManager().getEstimatedRemainingTasks() > 0))
              {
                  //Re-up any bg jobs
                  if (!maximal)
                  {
                      for (long i = working; i < threads; i++)
                          futures.addAll(CompactionManager.instance.submitBackground(cfs));
                  }
  
                  reportCompactionStats();
                  Uninterruptibles.sleepUninterruptibly(10, TimeUnit.SECONDS);
              }
  
              System.out.println("Finished! Shutting down...");
              CompactionManager.instance.forceShutdown();
  
              //Wait for cleanup to finish before forcing
              Uninterruptibles.sleepUninterruptibly(1, TimeUnit.SECONDS);
              LifecycleTransaction.removeUnfinishedLeftovers(cfs);
          }
      }
  
      void reportCompactionStats()
      {
          System.out.println("========");
          System.out.println(String.format("Pending compactions: %d\n", CompactionManager.instance.getPendingTasks()));
          CompactionStats.reportCompactionTable(CompactionManager.instance.getCompactions(), 0, true, System.out, new TableBuilder());
      }
  
  
      @Command(name = "write", description = "write data directly to disk")
      public static class DataWriter extends CompactionStress
      {
          private static double BYTES_IN_GIB = 1024 * 1014 * 1024;
  
          @Option(names = { "-g", "--gbsize" }, description = "Total GB size on disk you wish to write", required = true)
          Integer totalSizeGiB;
  
          @Option(names = { "-t", "--threads" }, description = "Number of sstable writer threads (default 2)")
          Integer threads = 2;
  
          @Option(names = { "-c", "--partition-count" }, description = "Number of partitions to loop over (default 1000000)")
          Integer partitions = 1000000;
  
          @Option(names = { "-b", "--buffer-size-mb" }, description = "Buffer in MiB writes before writing new sstable (default 128)")
          Integer bufferSize = 128;
  
          @Option(names = { "-r", "--range-aware" }, description = "Splits the local ranges in number of data directories and makes sure we never write the same token in two different directories (default true)")
          Boolean makeRangeAware = true;
  
          public void run()
          {
              ClusterMetadataService.initializeForTools(true);
              Keyspace.setInitialized();
              StressProfile stressProfile = getStressProfile();
              ColumnFamilyStore cfs = initCf(stressProfile, false);
              Directories directories = cfs.getDirectories();
  
              StressSettings settings = StressSettings.parse(new String[]{ "write", "-pop seq=1.." + partitions });
              SeedManager seedManager = new SeedManager(settings);
              PartitionGenerator generator = stressProfile.getOfflineGenerator();
              WorkManager workManager = new WorkManager.FixedWorkManager(Long.MAX_VALUE);
  
              ExecutorService executorService = Executors.newFixedThreadPool(threads);
              CountDownLatch finished = new CountDownLatch(threads);
  
              for (int i = 0; i < threads; i++)
              {
                  //Every thread needs it's own writer
                  final SchemaInsert insert = stressProfile.getOfflineInsert(null, generator, seedManager, settings);
                  final StressCQLSSTableWriter tableWriter = insert.createWriter(cfs, bufferSize, makeRangeAware);
                  executorService.submit(() -> {
                      try
                      {
                          insert.runOffline(tableWriter, workManager);
                      }
                      catch (Exception e)
                      {
                          e.printStackTrace();
                      }
                      finally
                      {
                          FileUtils.closeQuietly(tableWriter);
                          finished.countDown();
                      }
                  });
              }
  
              double currentSizeGiB;
              while ((currentSizeGiB = directories.getRawDirectoriesSize() / BYTES_IN_GIB) < totalSizeGiB)
              {
                  if (finished.getCount() == 0)
                      break;
  
                  System.out.println(String.format("Written %.2fGB of %dGB", currentSizeGiB, totalSizeGiB));
  
                  Uninterruptibles.sleepUninterruptibly(3, TimeUnit.SECONDS);
              }
  
              workManager.stop();
              Uninterruptibles.awaitUninterruptibly(finished);
  
              currentSizeGiB = directories.getRawDirectoriesSize() / BYTES_IN_GIB;
              System.out.println(String.format("Finished writing %.2fGB", currentSizeGiB));
          }
      }
  
      @Command(name = "compaction-stress",
               description = "benchmark for compaction",
               subcommands = { CommandLine.HelpCommand.class, CompactionStress.Compaction.class, CompactionStress.DataWriter.class })
      public static class CompactionStressCommand implements Runnable
      {
          @Override
          public void run()
          {
              CommandLine.usage(this, System.out);
          }
      }
  
      public static void main(String[] args)
      {
          CommandLine commandLine = new CommandLine(CompactionStressCommand.class);
          commandLine.setExecutionExceptionHandler((ex, c, arg) -> {
              ex.printStackTrace(System.out);
              return 6;
          });
          System.exit(commandLine.execute(args));
      }
  }
  
  
  