Remove write-based sampling from nodetool traincompressiondictionary

Author: yifan-c

src/java/org/apache/cassandra/db/compression/CompressionDictionaryManager.java

@@ -19,7 +19,6 @@
 package org.apache.cassandra.db.compression;
 
 import java.nio.ByteBuffer;
-import java.util.Map;
 import java.util.Set;
 import javax.annotation.Nullable;
 
@@ -113,8 +112,6 @@ public synchronized void maybeReloadFromSchema(CompressionParams newParams)
 
             if (needsNewTrainer)
             {
-                // The manual training should be cancelled if a new trainer is needed
-                scheduler.cancelManualTraining();
                 // Close existing trainer and create a new one
                 if (trainer != null)
                 {
@@ -201,7 +198,7 @@ public void onNewDictionaryAvailable(CompressionDictionary.DictId dictionaryId)
     }
 
     @Override
-    public synchronized void train(Map<String, String> options)
+    public synchronized void train()
     {
         // Validate table supports dictionary compression
         if (!isEnabled)
@@ -214,32 +211,26 @@ public synchronized void train(Map<String, String> options)
             throw new IllegalStateException("Dictionary trainer is not available for table " + keyspaceName + '.' + tableName);
         }
 
-        // Parse and validate training options
-        ManualTrainingOptions trainingOptions = ManualTrainingOptions.fromStringMap(options);
-
-        if (trainingOptions.useExistingSSTables())
+        // SSTable-based training: sample from existing SSTables
+        Set<SSTableReader> sstables = columnFamilyStore.getLiveSSTables();
+        if (sstables.isEmpty())
         {
-            // SSTable-based training: sample from existing SSTables
-            Set<SSTableReader> sstables = columnFamilyStore.getLiveSSTables();
+            logger.info("No SSTables available for training in table {}.{}, flushing memtable first",
+                       keyspaceName, tableName);
+            columnFamilyStore.forceBlockingFlush(ColumnFamilyStore.FlushReason.USER_FORCED);
+            sstables = columnFamilyStore.getLiveSSTables();
+
             if (sstables.isEmpty())
             {
-                throw new IllegalStateException("No SSTables available for training in table " + keyspaceName + '.' + tableName);
+                throw new IllegalStateException("No SSTables available for training in table " + keyspaceName + '.' + tableName + " after flush");
             }
-
-            logger.info("Starting SSTable-based training for {}.{} with {} SSTables",
-                       keyspaceName, tableName, sstables.size());
-
-            trainer.start(true);
-            scheduler.scheduleSSTableBasedTraining(trainingOptions, trainer, sstables, createTrainingConfig());
         }
-        else
-        {
-            // Write-based training: sample from new writes
-            logger.info("Starting write-based training for {}.{}", keyspaceName, tableName);
 
-            trainer.start(true);
-            scheduler.scheduleManualTraining(trainingOptions, trainer);
-        }
+        logger.info("Starting SSTable-based training for {}.{} with {} SSTables",
+                   keyspaceName, tableName, sstables.size());
+
+        trainer.start(true);
+        scheduler.scheduleSSTableBasedTraining(trainer, sstables, createTrainingConfig());
     }
 
     @Override
@@ -253,17 +244,6 @@ public String getTrainingStatus()
         return dictionaryTrainer.getTrainingStatus().toString();
     }
 
-    @Override
-    public void updateSamplingRate(int samplingRate)
-    {
-        ICompressionDictionaryTrainer dictionaryTrainer = trainer;
-        if (dictionaryTrainer == null)
-        {
-            throw new IllegalArgumentException("Dictionary trainer is not available for table " + keyspaceName + '.' + tableName);
-        }
-        dictionaryTrainer.updateSamplingRate(samplingRate);
-    }
-
     @Override
     public long getSampleCount()
     {

src/java/org/apache/cassandra/db/compression/CompressionDictionaryManagerMBean.java

@@ -18,20 +18,19 @@
 
 package org.apache.cassandra.db.compression;
 
-import java.util.Map;
-
 public interface CompressionDictionaryManagerMBean
 {
     String MBEAN_NAME = "org.apache.cassandra.db.compression:type=CompressionDictionaryManager";
 
     /**
-     * Starts sampling and training for this table.
-     * 
-     * @param options options for the training process (currently unused, reserved for future extensions)
+     * Starts training from existing SSTables for this table.
+     * Samples chunks from all live SSTables and trains a compression dictionary.
+     * If no SSTables are available, automatically flushes the memtable first.
+     *
      * @throws UnsupportedOperationException if table doesn't support dictionary compression
-     * @throws IllegalStateException if training is already in progress for this table
+     * @throws IllegalStateException if training is already in progress for this table or no SSTables available after flush
      */
-    void train(Map<String, String> options);
+    void train();
 
     /**
      * Gets the current training status for this table.
@@ -41,15 +40,6 @@ public interface CompressionDictionaryManagerMBean
      */
     String getTrainingStatus();
 
-    /**
-     * Updates the sampling rate for the trainer.
-     *
-     * @param samplingRate the new sampling rate. For exmaple, 1 = sample every time (100%);
-     *                     2 = expect sample 1/2 of data (50%), n = expect sample 1/n of data
-     * @throws IllegalArgumentException if sampling rate is invalid or trainer is not available
-     */
-    void updateSamplingRate(int samplingRate);
-
     /**
      * Gets the number of samples collected so far during training.
      *

src/java/org/apache/cassandra/db/compression/CompressionDictionaryScheduler.java

@@ -31,10 +31,8 @@
 
 import org.apache.cassandra.concurrent.ScheduledExecutors;
 import org.apache.cassandra.config.DatabaseDescriptor;
-import org.apache.cassandra.db.compression.ICompressionDictionaryTrainer.TrainingStatus;
 import org.apache.cassandra.io.sstable.format.SSTableReader;
 import org.apache.cassandra.schema.SystemDistributedKeyspace;
-import org.apache.cassandra.utils.Clock;
 import org.apache.cassandra.utils.concurrent.Ref;
 
 /**
@@ -85,41 +83,7 @@ public void scheduleRefreshTask()
     }
 
     @Override
-    public void scheduleManualTraining(ManualTrainingOptions options, ICompressionDictionaryTrainer trainer)
-    {
-        if (scheduledManualTrainingTask != null)
-        {
-            throw new IllegalStateException("Training already in progress for table " + keyspaceName + '.' + tableName);
-        }
-
-        int maxSamplingDurationSeconds = options.getMaxSamplingDurationSeconds();
-
-        logger.info("Starting manual dictionary training for {}.{} with max sampling duration: {} seconds",
-                    keyspaceName, tableName, maxSamplingDurationSeconds);
-
-        long deadlineMillis = Clock.Global.currentTimeMillis() + TimeUnit.SECONDS.toMillis(maxSamplingDurationSeconds);
-
-        ManualTrainingTask task = new ManualTrainingTask(deadlineMillis, trainer);
-
-        // Check every second whether it gets enough samples and completes training
-        scheduledManualTrainingTask = ScheduledExecutors.scheduledTasks
-                                      .scheduleWithFixedDelay(task, 1, 1, TimeUnit.SECONDS);
-    }
-
-    @Override
-    public void cancelManualTraining()
-    {
-        ScheduledFuture<?> future = scheduledManualTrainingTask;
-        if (future != null)
-        {
-            future.cancel(false);
-        }
-        scheduledManualTrainingTask = null;
-    }
-
-    @Override
-    public void scheduleSSTableBasedTraining(ManualTrainingOptions options,
-                                             ICompressionDictionaryTrainer trainer,
+    public void scheduleSSTableBasedTraining(ICompressionDictionaryTrainer trainer,
                                              Set<SSTableReader> sstables,
                                              CompressionDictionaryTrainingConfig config)
     {
@@ -140,6 +104,19 @@ public void scheduleSSTableBasedTraining(ManualTrainingOptions options,
         scheduledManualTrainingTask = ScheduledExecutors.scheduledTasks.schedule(() -> {}, 1, TimeUnit.HOURS);
     }
 
+    /**
+     * Cancels the in-progress manual training task.
+     */
+    private void cancelManualTraining()
+    {
+        ScheduledFuture<?> future = scheduledManualTrainingTask;
+        if (future != null)
+        {
+            future.cancel(false);
+        }
+        scheduledManualTrainingTask = null;
+    }
+
     /**
      * Sets the enabled state of the scheduler. When disabled, refresh tasks will not execute.
      *
@@ -190,55 +167,8 @@ public void close()
         }
     }
 
-    private class ManualTrainingTask implements Runnable
-    {
-        private final long deadlineMillis;
-        private final ICompressionDictionaryTrainer trainer;
-        private boolean isTraining = false;
-
-        private ManualTrainingTask(long deadlineMillis, ICompressionDictionaryTrainer trainer)
-        {
-            this.deadlineMillis = deadlineMillis;
-            this.trainer = trainer;
-        }
-
-        @Override
-        public void run()
-        {
-            if (trainer.getTrainingStatus() == TrainingStatus.NOT_STARTED)
-            {
-                logger.warn("Trainer is not started. Stop training dictionary for table {}.{}", keyspaceName, tableName);
-                cancelManualTraining();
-                return;
-            }
-
-            long now = Clock.Global.currentTimeMillis();
-            // Force training if there are not enough samples, but we have hit the max sampling duration
-            boolean reachedDeadline = now >= deadlineMillis;
-            if (!isTraining && (trainer.isReady() || reachedDeadline))
-            {
-                // Set isTraining to only enter the branch once
-                isTraining = true;
-                trainer.trainDictionaryAsync(reachedDeadline)
-                       .addCallback((dictionary, throwable) -> {
-                           cancelManualTraining();
-                           if (throwable != null)
-                           {
-                               logger.error("Manual dictionary training failed for {}.{}: {}",
-                                            keyspaceName, tableName, throwable.getMessage());
-                           }
-                           else
-                           {
-                               logger.info("Manual dictionary training completed for {}.{}", keyspaceName, tableName);
-                           }
-                       });
-            }
-        }
-    }
-
     /**
      * Task that samples chunks from existing SSTables and triggers training.
-     * Similar to ManualTrainingTask but reads from disk instead of waiting for writes.
      * Acquires references to SSTables to prevent them from being deleted during sampling.
      */
     private class SSTableSamplingTask implements Runnable
@@ -270,7 +200,7 @@ private SSTableSamplingTask(Set<SSTableReader> sstables,
                 else
                 {
                     logger.debug("Couldn't acquire reference to SSTable {}. It may have been removed.",
-                                sstable.descriptor);
+                                 sstable.descriptor);
                 }
             }
 

src/java/org/apache/cassandra/db/compression/ICompressionDictionaryScheduler.java

@@ -37,34 +37,18 @@ public interface ICompressionDictionaryScheduler extends AutoCloseable
      */
     void scheduleRefreshTask();
 
-    /**
-     * Schedules manual training with the specified options.
-     *
-     * @param options parsed and validated training options
-     * @param trainer the trainer to use
-     * @throws IllegalStateException if training is already in progress
-     */
-    void scheduleManualTraining(ManualTrainingOptions options, ICompressionDictionaryTrainer trainer);
-
     /**
      * Schedules SSTable-based training that samples from existing SSTables.
      *
-     * @param options parsed and validated training options
      * @param trainer the trainer to use
      * @param sstables the set of SSTables to sample from
      * @param config the training configuration
      * @throws IllegalStateException if training is already in progress
      */
-    void scheduleSSTableBasedTraining(ManualTrainingOptions options,
-                                      ICompressionDictionaryTrainer trainer,
+    void scheduleSSTableBasedTraining(ICompressionDictionaryTrainer trainer,
                                       Set<SSTableReader> sstables,
                                       CompressionDictionaryTrainingConfig config);
 
-    /**
-     * Cancel the in-progress manual training
-     */
-    void cancelManualTraining();
-
     /**
      * Sets the enabled state of the scheduler. When disabled, refresh tasks will not execute.
      *