fix tests

Author: yifan-c

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

@@ -24,6 +24,7 @@
 import java.util.Set;
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import com.google.common.annotations.VisibleForTesting;
 import org.slf4j.Logger;
@@ -50,9 +51,9 @@ public class CompressionDictionaryScheduler implements ICompressionDictionarySch
     private final String keyspaceName;
     private final String tableName;
     private final ICompressionDictionaryCache cache;
+    private final AtomicBoolean manualTrainingInProgress = new AtomicBoolean(false);
 
     private volatile ScheduledFuture<?> scheduledRefreshTask;
-    private volatile ScheduledFuture<?> scheduledManualTrainingTask;
     private volatile boolean isEnabled;
 
     public CompressionDictionaryScheduler(String keyspaceName,
@@ -88,7 +89,7 @@ public void scheduleSSTableBasedTraining(ICompressionDictionaryTrainer trainer,
                                              CompressionDictionaryTrainingConfig config,
                                              boolean force)
     {
-        if (scheduledManualTrainingTask != null)
+        if (!manualTrainingInProgress.compareAndSet(false, true))
         {
             throw new IllegalStateException("Training already in progress for table " + keyspaceName + '.' + tableName);
         }
@@ -97,25 +98,16 @@ public void scheduleSSTableBasedTraining(ICompressionDictionaryTrainer trainer,
                     keyspaceName, tableName, sstables.size());
 
         // Run the SSTableSamplingTask asynchronously
-        // Use a dummy scheduled task to track that training is in progress
         SSTableSamplingTask task = new SSTableSamplingTask(sstables, trainer, config, force);
         ScheduledExecutors.nonPeriodicTasks.submit(task);
-
-        // Set a placeholder task so status checks know training is in progress
-        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;
+        manualTrainingInProgress.compareAndSet(true, false);
     }
 
     /**
@@ -161,11 +153,7 @@ public void close()
             scheduledRefreshTask = null;
         }
 
-        if (scheduledManualTrainingTask != null)
-        {
-            scheduledManualTrainingTask.cancel(false);
-            scheduledManualTrainingTask = null;
-        }
+        cancelManualTraining();
     }
 
     /**
@@ -265,8 +253,8 @@ public void run()
     }
 
     @VisibleForTesting
-    ScheduledFuture<?> scheduledManualTrainingTask()
+    boolean isManualTrainingRunning()
     {
-        return scheduledManualTrainingTask;
+        return manualTrainingInProgress.get();
     }
 }

test/resources/nodetool/help/compressiondictionary

@@ -5,7 +5,8 @@ SYNOPSIS
         nodetool [(-h <host> | --host <host>)] [(-p <port> | --port <port>)]
                 [(-pp | --print-port)] [(-pw <password> | --password <password>)]
                 [(-pwf <passwordFilePath> | --password-file <passwordFilePath>)]
-                [(-u <username> | --username <username>)] compressiondictionary <command> [<args>]
+                [(-u <username> | --username <username>)] compressiondictionary
+                <command> [<args>]
 
         nodetool [(-h <host> | --host <host>)] [(-p <port> | --port <port>)]
                 [(-pp | --print-port)] [(-pw <password> | --password <password>)]
@@ -39,5 +40,5 @@ COMMANDS
             Manually trigger compression dictionary training for a table. If no
             SSTables are available, the memtable will be flushed first.
 
-            With -f or --force option, Force the dictionary training even if there
-            are not enough samples
+            With --force option, Force the dictionary training even if there are not
+            enough samples

test/resources/nodetool/help/compressiondictionary$train

@@ -42,4 +42,3 @@ OPTIONS
 
         <table>
             The table name
-

test/resources/nodetool/help/nodetool

@@ -21,6 +21,7 @@ The most commonly used nodetool commands are:
     compact                             Force a (major) compaction on one or more tables or user-defined compaction on given SSTables
     compactionhistory                   Print history of compaction
     compactionstats                     Print statistics on compactions
+    compressiondictionary               Manage compression dictionaries
     consensus_admin                     List and mark ranges as migrating between consensus protocols
     datapaths                           Print all directories where data of tables are stored
     decommission                        Decommission the *node I am connecting to*
@@ -155,7 +156,6 @@ The most commonly used nodetool commands are:
     tablestats                          Print statistics on tables
     toppartitions                       Sample and print the most active partitions
     tpstats                             Print usage statistics of thread pools
-    traincompressiondictionary          Manually trigger compression dictionary training for a table. If no SSTables are available, the memtable will be flushed first.
     truncatehints                       Truncate all hints on the local node, or truncate hints for the endpoint(s) specified.
     updatecidrgroup                     Insert/Update a cidr group
     upgradesstables                     Rewrite sstables (for the requested tables) that are not on the current version (thus upgrading them to said current version)

test/unit/org/apache/cassandra/db/compression/CompressionDictionaryIntegrationTest.java

@@ -19,6 +19,7 @@
 package org.apache.cassandra.db.compression;
 
 import java.util.Collections;
+import java.util.concurrent.TimeUnit;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -142,9 +143,9 @@ public void testResourceCleanupOnClose() throws Exception
         manager.close();
 
         assertThat(manager.trainer()).isNull();
-        assertThat(testDict.selfRef().globalCount())
-        .as("Dictionary's reference count should be 0 after closing manager")
-        .isZero();
+        // Dictionary's reference count should be 0 after closing manager
+        // Dictionary is closed in a separate thread. Wait a bit for the reference count to be updated.
+        spinUntilTrue(() -> testDict.selfRef().globalCount() == 0, 1, TimeUnit.SECONDS);
         assertThat(testDict.rawDictionary())
         .as("The raw dictionary bytes should still be accessible")
         .isNotNull();

test/unit/org/apache/cassandra/db/compression/CompressionDictionarySchedulerTest.java

@@ -36,17 +36,15 @@
 
 public class CompressionDictionarySchedulerTest extends CQLTester
 {
-    private static final String KEYSPACE = "scheduler_test_ks";
-    private static final String TABLE = "scheduler_test_table";
-
     private CompressionDictionaryScheduler scheduler;
     private ICompressionDictionaryCache cache;
 
     @Before
     public void setUp()
     {
         cache = new CompressionDictionaryCache();
-        scheduler = new CompressionDictionaryScheduler(KEYSPACE, TABLE, cache, true);
+        // Disable compaction to make the sstable sampling deterministic; to avoid excluding sstables get compacted away.
+        disableCompaction(KEYSPACE);
     }
 
     @After
@@ -63,6 +61,8 @@ public void testScheduleSSTableBasedTrainingWithNoSSTables()
     {
         String table = createTable("CREATE TABLE %s (id int PRIMARY KEY, data text) " +
                                    "WITH compression = {'class': 'ZstdDictionaryCompressor'}");
+        scheduler = new CompressionDictionaryScheduler(KEYSPACE, table, cache, true);
+
         ColumnFamilyStore cfs = Keyspace.open(keyspace()).getColumnFamilyStore(table);
         CompressionDictionaryManager manager = cfs.compressionDictionaryManager();
 
@@ -71,7 +71,7 @@ public void testScheduleSSTableBasedTrainingWithNoSSTables()
 
         // Should not throw, but task will complete quickly with no SSTables
         scheduler.scheduleSSTableBasedTraining(manager.trainer(), sstables, config, true);
-        spinUntilTrue(() -> scheduler.scheduledManualTrainingTask() == null);
+        spinUntilTrue(() -> !scheduler.isManualTrainingRunning());
         assertThat(manager.getCurrent()).isNull();
     }
 
@@ -80,6 +80,8 @@ public void testScheduleSSTableBasedTrainingWithSSTables()
     {
         String table = createTable("CREATE TABLE %s (id int PRIMARY KEY, data text) " +
                                    "WITH compression = {'class': 'ZstdDictionaryCompressor', 'chunk_length_in_kb': '4'}");
+        scheduler = new CompressionDictionaryScheduler(KEYSPACE, table, cache, true);
+
         ColumnFamilyStore cfs = Keyspace.open(keyspace()).getColumnFamilyStore(table);
         CompressionDictionaryManager manager = cfs.compressionDictionaryManager();
 
@@ -95,7 +97,7 @@ public void testScheduleSSTableBasedTrainingWithSSTables()
         scheduler.scheduleSSTableBasedTraining(manager.trainer(), sstables, config, true);
 
         // Task should be scheduled
-        assertThat((Object) scheduler.scheduledManualTrainingTask()).isNotNull();
+        assertThat(scheduler.isManualTrainingRunning()).isTrue();
         // A dictionary should be trained
         spinUntilTrue(() -> manager.getCurrent() != null);
     }