Add PerKeyPlacement support (#13459)

Summary:
This PR adds support for PerKeyPlacement in Remote Compaction.

The `seqno_to_time_mapping` is already available from the table properties of the input files. `preserve_internal_time_seconds` and `preclude_last_level_data_seconds` are directly read from the OPTIONS file upon db open in the remote worker. The necessary changes include:

- Add `is_penultimate_level_output` and `file_temperature` to the `CompactionServiceOutputFile`
- When building the output for the remote compaction, get the outputs for penultimate level and last level separately, serialize them with the two additional information added in this PR.
- When deserializing the result from the primary, SubcompactionState's `GetOutputs()` now takes `is_penultimate_level`. This allows us to determine which level to place the output file.
- Include stats from `compaction_stats.penultimate_level_stats` in the remote compaction result

# To Follow up
- Stats to be fixed. Stats are not being populated correctly for PerKeyPlacement even for non-remote compactions.
- Clean up / Reconcile the "penultimate" naming by replacing with "proximal"

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

Test Plan:
Updated the unit test

```
./compaction_service_test
```

Reviewed By: pdillinger

Differential Revision: D71007211

Pulled By: jaykorean

fbshipit-source-id: f926e56df17239875d849d46b8b940f8cd5f1825

db/compaction/compaction_job.h

@@ -431,6 +431,9 @@ struct CompactionServiceOutputFile {
   bool marked_for_compaction;
   UniqueId64x2 unique_id{};
   TableProperties table_properties;
+  // TODO: clean up the rest of the "penultimate" naming in the codebase
+  bool is_proximal_level_output;  // == is_penultimate_level_output
+  Temperature file_temperature;
 
   CompactionServiceOutputFile() = default;
   CompactionServiceOutputFile(
@@ -440,7 +443,8 @@ struct CompactionServiceOutputFile {
       uint64_t _epoch_number, const std::string& _file_checksum,
       const std::string& _file_checksum_func_name, uint64_t _paranoid_hash,
       bool _marked_for_compaction, UniqueId64x2 _unique_id,
-      const TableProperties& _table_properties)
+      const TableProperties& _table_properties, bool _is_proximal_level_output,
+      Temperature _file_temperature)
       : file_name(name),
         smallest_seqno(smallest),
         largest_seqno(largest),
@@ -454,7 +458,9 @@ struct CompactionServiceOutputFile {
         paranoid_hash(_paranoid_hash),
         marked_for_compaction(_marked_for_compaction),
         unique_id(std::move(_unique_id)),
-        table_properties(_table_properties) {}
+        table_properties(_table_properties),
+        is_proximal_level_output(_is_proximal_level_output),
+        file_temperature(_file_temperature) {}
 };
 
 // CompactionServiceResult contains the compaction result from a different db

db/compaction/compaction_job_test.cc

@@ -1682,7 +1682,8 @@ TEST_F(CompactionJobTest, ResultSerialization) {
         file_checksum /* file_checksum */,
         file_checksum_func_name /* file_checksum_func_name */,
         rnd64.Uniform(UINT64_MAX) /* paranoid_hash */,
-        rnd.OneIn(2) /* marked_for_compaction */, id /* unique_id */, tp);
+        rnd.OneIn(2) /* marked_for_compaction */, id /* unique_id */, tp,
+        false /* is_proximal_level_output */, Temperature::kHot);
   }
   result.output_level = rnd.Uniform(10);
   result.output_path = rnd.RandomString(rnd.Uniform(kStrMaxLen));
@@ -1736,6 +1737,8 @@ TEST_F(CompactionJobTest, ResultSerialization) {
     ASSERT_EQ(deserialized_tmp.output_files[0].file_checksum, file_checksum);
     ASSERT_EQ(deserialized_tmp.output_files[0].file_checksum_func_name,
               file_checksum_func_name);
+    ASSERT_EQ(deserialized_tmp.output_files[0].file_temperature,
+              Temperature::kHot);
   }
 
   // Test unknown field

db/compaction/compaction_outputs.h

@@ -30,13 +30,16 @@ class CompactionOutputs {
   // compaction output file
   struct Output {
     Output(FileMetaData&& _meta, const InternalKeyComparator& _icmp,
-           bool _enable_hash, bool _finished, uint64_t precalculated_hash)
+           bool _enable_hash, bool _finished, uint64_t precalculated_hash,
+           bool _is_penultimate_level)
         : meta(std::move(_meta)),
           validator(_icmp, _enable_hash, precalculated_hash),
-          finished(_finished) {}
+          finished(_finished),
+          is_penultimate_level(_is_penultimate_level) {}
     FileMetaData meta;
     OutputValidator validator;
     bool finished;
+    bool is_penultimate_level;
     std::shared_ptr<const TableProperties> table_properties;
   };
 
@@ -52,7 +55,7 @@ class CompactionOutputs {
                  bool enable_hash, bool finished = false,
                  uint64_t precalculated_hash = 0) {
     outputs_.emplace_back(std::move(meta), icmp, enable_hash, finished,
-                          precalculated_hash);
+                          precalculated_hash, is_penultimate_level_);
   }
 
   // Set new table builder for the current output

db/compaction/compaction_service_job.cc

@@ -239,12 +239,15 @@ CompactionJob::ProcessKeyValueCompactionWithCompactionService(
     meta.file_checksum_func_name = file.file_checksum_func_name;
     meta.marked_for_compaction = file.marked_for_compaction;
     meta.unique_id = file.unique_id;
+    meta.temperature = file.file_temperature;
 
     auto cfd = compaction->column_family_data();
-    sub_compact->Current().AddOutput(std::move(meta),
-                                     cfd->internal_comparator(), false, true,
-                                     file.paranoid_hash);
-    sub_compact->Current().UpdateTableProperties(file.table_properties);
+    CompactionOutputs* compaction_outputs =
+        sub_compact->Outputs(file.is_proximal_level_output);
+    assert(compaction_outputs);
+    compaction_outputs->AddOutput(std::move(meta), cfd->internal_comparator(),
+                                  false, true, file.paranoid_hash);
+    compaction_outputs->UpdateTableProperties(file.table_properties);
   }
   sub_compact->compaction_job_stats = compaction_result.stats;
   sub_compact->Current().SetNumOutputRecords(
@@ -273,14 +276,12 @@ void CompactionServiceCompactionJob::RecordCompactionIOStats() {
 
 void CompactionServiceCompactionJob::UpdateCompactionJobStats(
     const InternalStats::CompactionStats& stats) const {
-  compaction_job_stats_->elapsed_micros = stats.micros;
-
   // output information only in remote compaction
-  compaction_job_stats_->total_output_bytes = stats.bytes_written;
-  compaction_job_stats_->total_output_bytes_blob = stats.bytes_written_blob;
-  compaction_job_stats_->num_output_records = stats.num_output_records;
-  compaction_job_stats_->num_output_files = stats.num_output_files;
-  compaction_job_stats_->num_output_files_blob = stats.num_output_files_blob;
+  compaction_job_stats_->total_output_bytes += stats.bytes_written;
+  compaction_job_stats_->total_output_bytes_blob += stats.bytes_written_blob;
+  compaction_job_stats_->num_output_records += stats.num_output_records;
+  compaction_job_stats_->num_output_files += stats.num_output_files;
+  compaction_job_stats_->num_output_files_blob += stats.num_output_files_blob;
 }
 
 CompactionServiceCompactionJob::CompactionServiceCompactionJob(
@@ -344,15 +345,15 @@ Status CompactionServiceCompactionJob::Run() {
 
   ProcessKeyValueCompaction(sub_compact);
 
-  compaction_stats_.stats.micros =
+  compaction_job_stats_->elapsed_micros =
       db_options_.clock->NowMicros() - start_micros;
-  compaction_stats_.stats.cpu_micros =
+  compaction_job_stats_->cpu_micros =
       sub_compact->compaction_job_stats.cpu_micros;
 
   RecordTimeToHistogram(stats_, COMPACTION_TIME,
-                        compaction_stats_.stats.micros);
+                        compaction_job_stats_->elapsed_micros);
   RecordTimeToHistogram(stats_, COMPACTION_CPU_TIME,
-                        compaction_stats_.stats.cpu_micros);
+                        compaction_job_stats_->cpu_micros);
 
   Status status = sub_compact->status;
   IOStatus io_s = sub_compact->io_status;
@@ -390,6 +391,9 @@ Status CompactionServiceCompactionJob::Run() {
   // 2. Update the Output information in the Compaction Job Stats with
   // aggregated Internal Compaction Stats.
   UpdateCompactionJobStats(compaction_stats_.stats);
+  if (compaction_stats_.has_penultimate_level_output) {
+    UpdateCompactionJobStats(compaction_stats_.penultimate_level_stats);
+  }
 
   // 3. Set fields that are not propagated as part of aggregations above
   compaction_result_->stats.is_manual_compaction = c->is_manual_compaction();
@@ -413,7 +417,8 @@ Status CompactionServiceCompactionJob::Run() {
           meta.file_creation_time, meta.epoch_number, meta.file_checksum,
           meta.file_checksum_func_name, output_file.validator.GetHash(),
           meta.marked_for_compaction, meta.unique_id,
-          *output_file.table_properties);
+          *output_file.table_properties, output_file.is_penultimate_level,
+          meta.temperature);
     }
   }
 
@@ -585,7 +590,16 @@ static std::unordered_map<std::string, OptionTypeInfo>
             const auto this_one = static_cast<const TableProperties*>(addr1);
             const auto that_one = static_cast<const TableProperties*>(addr2);
             return this_one->AreEqual(opts, that_one, mismatch);
-          }}}};
+          }}},
+        {"is_proximal_level_output",
+         {offsetof(struct CompactionServiceOutputFile,
+                   is_proximal_level_output),
+          OptionType::kBoolean, OptionVerificationType::kNormal,
+          OptionTypeFlags::kNone}},
+        {"file_temperature",
+         {offsetof(struct CompactionServiceOutputFile, file_temperature),
+          OptionType::kTemperature, OptionVerificationType::kNormal,
+          OptionTypeFlags::kNone}}};
 
 static std::unordered_map<std::string, OptionTypeInfo>
     compaction_job_stats_type_info = {

db/compaction/compaction_service_test.cc

@@ -1188,34 +1188,34 @@ TEST_F(CompactionServiceTest, PrecludeLastLevel) {
 
   for (int i = 0; i < kNumTrigger; i++) {
     for (int j = 0; j < kNumKeys; j++) {
-      // FIXME: need to assign outputs to levels to allow overlapping ranges:
-      // ASSERT_OK(Put(Key(j * kNumTrigger + i), "v" + std::to_string(i)));
-      // instead of this (too easy):
-      ASSERT_OK(Put(Key(i * kNumKeys + j), "v" + std::to_string(i)));
+      ASSERT_OK(Put(Key(j * kNumTrigger + i), "v" + std::to_string(i)));
     }
     ASSERT_OK(Flush());
   }
   ASSERT_OK(dbfull()->TEST_WaitForCompact());
 
   // Data split between penultimate (kUnknown) and last (kCold) levels
-  // FIXME: need to assign outputs to levels to get this:
-  // ASSERT_EQ("0,0,0,0,0,1,1", FilesPerLevel());
-  // ASSERT_GT(GetSstSizeHelper(Temperature::kUnknown), 0);
-  // ASSERT_GT(GetSstSizeHelper(Temperature::kCold), 0);
-  // instead of this (WRONG but currently expected):
-  ASSERT_EQ("0,0,0,0,0,0,2", FilesPerLevel());
-  // Check manifest temperatures
+  ASSERT_EQ("0,0,0,0,0,1,1", FilesPerLevel());
   ASSERT_GT(GetSstSizeHelper(Temperature::kUnknown), 0);
-  ASSERT_EQ(GetSstSizeHelper(Temperature::kCold), 0);
+  ASSERT_GT(GetSstSizeHelper(Temperature::kCold), 0);
+
   // TODO: Check FileSystem temperatures with FileTemperatureTestFS
 
   for (int i = 0; i < kNumTrigger; i++) {
     for (int j = 0; j < kNumKeys; j++) {
-      // FIXME
-      // ASSERT_EQ(Get(Key(j * kNumTrigger + i)), "v" + std::to_string(i));
-      ASSERT_EQ(Get(Key(i * kNumKeys + j)), "v" + std::to_string(i));
+      ASSERT_EQ(Get(Key(j * kNumTrigger + i)), "v" + std::to_string(i));
     }
   }
+
+  // Verify Output Stats
+  auto my_cs = GetCompactionService();
+  CompactionServiceResult result;
+  my_cs->GetResult(&result);
+  ASSERT_OK(result.status);
+  ASSERT_GT(result.stats.cpu_micros, 0);
+  ASSERT_GT(result.stats.elapsed_micros, 0);
+  ASSERT_EQ(result.stats.num_output_records, kNumTrigger * kNumKeys);
+  ASSERT_EQ(result.stats.num_output_files, 2);
 }
 
 TEST_F(CompactionServiceTest, ConcurrentCompaction) {

db/compaction/subcompaction_state.h

@@ -169,6 +169,15 @@ class SubcompactionState {
     return *current_outputs_;
   }
 
+  CompactionOutputs* Outputs(bool is_penultimate_level) {
+    assert(compaction);
+    if (is_penultimate_level) {
+      assert(compaction->SupportsPerKeyPlacement());
+      return &penultimate_level_outputs_;
+    }
+    return &compaction_outputs_;
+  }
+
   CompactionRangeDelAggregator* RangeDelAgg() const {
     return range_del_agg_.get();
   }

unreleased_history/new_features/per_key_placement_remote_compaction.md

@@ -0,0 +1 @@
+Added per-key-placement feature in Remote Compaction