From d2db80caa28e21d0af8617b8a3d7cd83bbdb0443 Mon Sep 17 00:00:00 2001
From: Changyu Bi <changyubi@meta.com>
Date: Thu, 2 Jan 2025 12:10:31 -0800
Subject: [PATCH] Try to deflake unit test
 `RoundRobinSubcompactionsAgainstPressureToken` (#13254)

Summary:
unit test `RoundRobinSubcompactionsAgainstPressureToken.PressureTokenTest` has been [flaky](https://github.com/facebook/rocksdb/actions/runs/12220443012/job/34088263578?fbclid=IwZXh0bgNhZW0CMTEAAR3Vi0p8xxzU1tSpvaeB0RfP_97nOMiONGyZbhdcnN8IXW4tChNVHN3iIhc_aem_SGy-iqplt0GaEHAel_BGQQ). num_planned_subcompactions can be 1 for two reasons: compactions not having enough input files or that there were not enough bg threads. This PR updates the test to try to trigger a larger compaction for subcompactions, and added a callback to verify compactions have enough input files.

Pull Request resolved: https://github.com/facebook/rocksdb/pull/13254

Test Plan: monitor future failure.

Reviewed By: hx235

Differential Revision: D67764944

Pulled By: cbi42

fbshipit-source-id: 4fc9c0bef76c8bfaa54be4f3d78071e2bebee8aa
---
 db/db_compaction_test.cc | 40 ++++++++++++++++++++++++++++------------
 1 file changed, 28 insertions(+), 12 deletions(-)

diff --git a/db/db_compaction_test.cc b/db/db_compaction_test.cc
index c71cf2e0c1..05cf7c4615 100644
--- a/db/db_compaction_test.cc
+++ b/db/db_compaction_test.cc
@@ -6417,10 +6417,13 @@ TEST_P(RoundRobinSubcompactionsAgainstPressureToken, PressureTokenTest) {
   options.level0_file_num_compaction_trigger = 4;
   options.target_file_size_base = kKeysPerBuffer * 1024;
   options.compaction_pri = CompactionPri::kRoundRobin;
-  // Target size is chosen so that filling the level with `kFilesPerLevel` files
-  // will make it oversized by `kNumSubcompactions` files.
+  // Pick a small target level size so that round-robin compaction will pick
+  // more files to compact and trigger subcompactions. Actual number of
+  // subcompactions will be limited by number of bg threads available.
+  // Cannot be too small to cause compaction pressure. See
+  // GetPendingCompactionBytesForCompactionSpeedup().
   options.max_bytes_for_level_base =
-      (kFilesPerLevel - kNumSubcompactions) * kKeysPerBuffer * 1024;
+      (kFilesPerLevel - 10) * kKeysPerBuffer * 1024;
   options.disable_auto_compactions = true;
   // Setup `kNumSubcompactions` threads but limited subcompactions so
   // that RoundRobin requires extra compactions from reserved threads
@@ -6446,21 +6449,34 @@ TEST_P(RoundRobinSubcompactionsAgainstPressureToken, PressureTokenTest) {
     ASSERT_EQ(kFilesPerLevel, NumTableFilesAtLevel(lvl, 0));
   }
 
+  bool compaction_num_input_file_verified = false;
+  SyncPoint::GetInstance()->SetCallBack(
+      "LevelCompactionPicker::PickCompaction:Return", [&](void* arg) {
+        if (!compaction_num_input_file_verified) {
+          Compaction* compaction = static_cast<Compaction*>(arg);
+          // In Round-Robin, # of subcompactions needed is the # of input files
+          ASSERT_GT(compaction->num_input_files(0), 1);
+          compaction_num_input_file_verified = true;
+        }
+      });
+
   // This is a variable for making sure the following callback is called
   // and the assertions in it are indeed excuted.
   bool num_planned_subcompactions_verified = false;
   ROCKSDB_NAMESPACE::SyncPoint::GetInstance()->SetCallBack(
       "CompactionJob::GenSubcompactionBoundaries:0", [&](void* arg) {
-        uint64_t num_planned_subcompactions = *(static_cast<uint64_t*>(arg));
-        if (grab_pressure_token_) {
-          // `kNumSubcompactions` files are selected for round-robin under auto
-          // compaction. The number of planned subcompaction is restricted by
-          // the limited number of max_background_compactions
-          ASSERT_EQ(num_planned_subcompactions, kNumSubcompactions);
-        } else {
-          ASSERT_EQ(num_planned_subcompactions, 1);
+        if (!num_planned_subcompactions_verified) {
+          uint64_t num_planned_subcompactions = *(static_cast<uint64_t*>(arg));
+          if (grab_pressure_token_) {
+            // `kNumSubcompactions` files are selected for round-robin under
+            // auto compaction. The number of planned subcompaction is
+            // restricted by the limited number of max_background_compactions
+            ASSERT_EQ(num_planned_subcompactions, kNumSubcompactions);
+          } else {
+            ASSERT_EQ(num_planned_subcompactions, 1);
+          }
+          num_planned_subcompactions_verified = true;
         }
-        num_planned_subcompactions_verified = true;
       });
 
   // The following 3 dependencies have to be added to ensure the auto