Add option to verify file checksum of output files (facebook#14433)

Summary:
One of the follow ups from facebook#14103. Users will have the option to verify file checksums for all compaction output files before they are installed. This feature helps prevent corrupted SST files from being added to the LSM tree.

Pull Request resolved: facebook#14433

Test Plan: Unit test added

Reviewed By: archang19

Differential Revision: D95648000

Pulled By: jaykorean

fbshipit-source-id: 512f3c1a7449b96a660865531f3537624c89a9cc

Author: jaykorean

db/compaction/compaction_job.cc

@@ -33,6 +33,7 @@
 #include "db/range_del_aggregator.h"
 #include "db/version_edit.h"
 #include "db/version_set.h"
+#include "file/file_util.h"
 #include "file/filename.h"
 #include "file/read_write_util.h"
 #include "file/sst_file_manager_impl.h"
@@ -903,17 +904,17 @@ Status CompactionJob::VerifyOutputFiles() {
       // use_direct_io_for_flush_and_compaction is true, we will regard this
       // verification as user reads since the goal is to cache it here for
       // further user reads
-      ReadOptions verify_table_read_options(Env::IOActivity::kCompaction);
-      verify_table_read_options.verify_checksums = true;
-      verify_table_read_options.readahead_size =
+      ReadOptions verification_read_options(Env::IOActivity::kCompaction);
+      verification_read_options.verify_checksums = true;
+      verification_read_options.readahead_size =
           file_options_for_read_.compaction_readahead_size;
 
       std::unique_ptr<TableReader> table_reader_guard;
       TableReader* table_reader_ptr = table_reader_guard.get();
-      verify_table_read_options.rate_limiter_priority =
+      verification_read_options.rate_limiter_priority =
           GetRateLimiterPriority();
       InternalIterator* iter = cfd->table_cache()->NewIterator(
-          verify_table_read_options, file_options_, cfd->internal_comparator(),
+          verification_read_options, file_options_, cfd->internal_comparator(),
           output_file.meta,
           /*range_del_agg=*/nullptr, compact_->compaction->mutable_cf_options(),
           /*table_reader_ptr=*/&table_reader_ptr,
@@ -941,12 +942,17 @@ Status CompactionJob::VerifyOutputFiles() {
               !!(verify_output_flags & VerifyOutputFlags::kVerifyBlockChecksum);
           const bool should_verify_iteration =
               !!(verify_output_flags & VerifyOutputFlags::kVerifyIteration);
+          const bool should_verify_file_checksum =
+              !!(verify_output_flags &
+                 VerifyOutputFlags::kVerifyFileChecksum) &&
+              db_options_.file_checksum_gen_factory != nullptr &&
+              output_file.meta.file_checksum != kUnknownFileChecksum;
           if (should_verify_block_checksum) {
             assert(table_reader_ptr != nullptr);
             // If verifying iteration as well, verify meta blocks here only to
             // avoid redundant checks on data blocks
             s = table_reader_ptr->VerifyChecksum(
-                verify_table_read_options, TableReaderCaller::kCompaction,
+                verification_read_options, TableReaderCaller::kCompaction,
                 /*meta_blocks_only=*/should_verify_iteration);
           }
           if (s.ok() && should_verify_iteration) {
@@ -967,6 +973,24 @@ Status CompactionJob::VerifyOutputFiles() {
                   "was computed when written");
             }
           }
+          if (s.ok() && should_verify_file_checksum) {
+            std::string file_checksum;
+            std::string file_checksum_func_name;
+            std::string fname =
+                GetTableFileName(output_file.meta.fd.GetNumber());
+            s = GenerateOneFileChecksum(
+                fs_.get(), fname, db_options_.file_checksum_gen_factory.get(),
+                output_file.meta.file_checksum_func_name, &file_checksum,
+                &file_checksum_func_name,
+                verification_read_options.readahead_size,
+                db_options_.allow_mmap_reads, io_tracer_,
+                db_options_.rate_limiter.get(), verification_read_options,
+                stats_, db_options_.clock, file_options_for_read_);
+            if (s.ok() && file_checksum != output_file.meta.file_checksum) {
+              s = Status::Corruption(
+                  "File checksum mismatch for compaction output file " + fname);
+            }
+          }
         }
       }
 

db/compaction/compaction_service_test.cc

@@ -1113,6 +1113,8 @@ TEST_F(CompactionServiceTest, CorruptedOutputVerifyOutputFlags) {
             VerifyOutputFlags::kVerifyBlockChecksum,
         VerifyOutputFlags::kEnableForRemoteCompaction |
             VerifyOutputFlags::kVerifyIteration,
+        VerifyOutputFlags::kEnableForRemoteCompaction |
+            VerifyOutputFlags::kVerifyFileChecksum,
         VerifyOutputFlags::kVerifyAll}) {
     SCOPED_TRACE(
         "verify_output_flags=" +
@@ -1124,6 +1126,7 @@ TEST_F(CompactionServiceTest, CorruptedOutputVerifyOutputFlags) {
     options.disable_auto_compactions = true;
     options.paranoid_file_checks = false;
     options.verify_output_flags = verify_output_flags;
+    options.file_checksum_gen_factory = GetFileChecksumGenCrc32cFactory();
     ReopenWithCompactionService(&options);
     GenerateTestData();
 
@@ -1159,10 +1162,13 @@ TEST_F(CompactionServiceTest, CorruptedOutputVerifyOutputFlags) {
         !!(verify_output_flags & VerifyOutputFlags::kVerifyBlockChecksum);
     const bool should_verify_iteration =
         !!(verify_output_flags & VerifyOutputFlags::kVerifyIteration);
+    const bool should_verify_file_checksum =
+        !!(verify_output_flags & VerifyOutputFlags::kVerifyFileChecksum);
 
     Status s = db_->CompactRange(CompactRangeOptions(), &start, &end);
     if (is_enabled_for_remote_compaction &&
-        (should_verify_block_checksum || should_verify_iteration)) {
+        (should_verify_block_checksum || should_verify_iteration ||
+         should_verify_file_checksum)) {
       ASSERT_NOK(s);
       ASSERT_TRUE(s.IsCorruption());
     } else {

db/db_compaction_test.cc

@@ -14,11 +14,13 @@
 #include "db/db_test_util.h"
 #include "db/dbformat.h"
 #include "env/mock_env.h"
+#include "file/filename.h"
 #include "port/port.h"
 #include "port/stack_trace.h"
 #include "rocksdb/advanced_options.h"
 #include "rocksdb/concurrent_task_limiter.h"
 #include "rocksdb/experimental.h"
+#include "rocksdb/file_checksum.h"
 #include "rocksdb/iostats_context.h"
 #include "rocksdb/sst_file_writer.h"
 #include "test_util/mock_time_env.h"
@@ -11735,6 +11737,59 @@ TEST_F(DBCompactionTest, RoundRobinCleanCutWithSharedBoundary) {
     ASSERT_EQ(Get("key" + std::to_string(k)), "v4");
   }
 }
+
+TEST_F(DBCompactionTest, VerifyFileChecksumOnCompactionOutput) {
+  Options options = CurrentOptions();
+  options.disable_auto_compactions = true;
+  options.file_checksum_gen_factory = GetFileChecksumGenCrc32cFactory();
+  options.verify_output_flags = VerifyOutputFlags::kVerifyFileChecksum |
+                                VerifyOutputFlags::kEnableForLocalCompaction;
+  DestroyAndReopen(options);
+
+  // Create 2 L0 files to trigger compaction
+  for (int i = 0; i < 10; i++) {
+    ASSERT_OK(Put(Key(i), "value" + std::to_string(i)));
+  }
+  ASSERT_OK(Flush());
+
+  for (int i = 5; i < 15; i++) {
+    ASSERT_OK(Put(Key(i), "value2_" + std::to_string(i)));
+  }
+  ASSERT_OK(Flush());
+
+  // Corrupt output files right before verification
+  ROCKSDB_NAMESPACE::SyncPoint::GetInstance()->SetCallBack(
+      "CompactionJob::Run:BeforeVerify", [&](void* /*arg*/) {
+        // Find and corrupt the newest SST file (compaction output)
+        std::vector<std::string> filenames;
+        ASSERT_OK(env_->GetChildren(dbname_, &filenames));
+        uint64_t max_number = 0;
+        std::string target_fname;
+        for (const auto& f : filenames) {
+          uint64_t number;
+          FileType type;
+          if (ParseFileName(f, &number, &type) && type == kTableFile &&
+              number > max_number) {
+            max_number = number;
+            target_fname = dbname_ + "/" + f;
+          }
+        }
+        ASSERT_FALSE(target_fname.empty());
+        ASSERT_OK(test::CorruptFile(env_, target_fname, 0, 1,
+                                    false /* verifyChecksum */));
+      });
+  SyncPoint::GetInstance()->EnableProcessing();
+
+  Status s = db_->CompactRange(CompactRangeOptions(), nullptr, nullptr);
+  ASSERT_TRUE(s.IsCorruption()) << s.ToString();
+  ASSERT_TRUE(
+      std::strstr(s.getState(), "File checksum mismatch for compaction output"))
+      << s.ToString();
+
+  SyncPoint::GetInstance()->DisableProcessing();
+  SyncPoint::GetInstance()->ClearAllCallBacks();
+}
+
 }  // namespace ROCKSDB_NAMESPACE
 
 int main(int argc, char** argv) {

db_stress_tool/db_stress_common.h

@@ -287,6 +287,7 @@ DECLARE_bool(verification_only);
 DECLARE_string(last_level_temperature);
 DECLARE_string(default_write_temperature);
 DECLARE_string(default_temperature);
+DECLARE_uint32(verify_output_flags);
 DECLARE_bool(paranoid_memory_checks);
 DECLARE_bool(memtable_veirfy_per_key_checksum_on_seek);
 

db_stress_tool/db_stress_gflags.cc

@@ -1536,6 +1536,10 @@ DEFINE_uint32(uncache_aggressiveness,
               "obsolete. 0 = disabled, 1 = minimum, 100 = moderate, 10000 = "
               "normal max");
 
+DEFINE_uint32(verify_output_flags, 0,
+              "Sets CF option verify_output_flags as a uint32_t bitmask. "
+              "See VerifyOutputFlags enum for bit definitions.");
+
 DEFINE_bool(paranoid_memory_checks,
             ROCKSDB_NAMESPACE::Options().paranoid_memory_checks,
             "Sets CF option paranoid_memory_checks.");

db_stress_tool/db_stress_test_base.cc

@@ -4532,6 +4532,8 @@ void InitializeOptionsFromFlags(
   options.memtable_protection_bytes_per_key =
       FLAGS_memtable_protection_bytes_per_key;
   options.block_protection_bytes_per_key = FLAGS_block_protection_bytes_per_key;
+  options.verify_output_flags =
+      static_cast<VerifyOutputFlags>(FLAGS_verify_output_flags);
   options.paranoid_memory_checks = FLAGS_paranoid_memory_checks;
   options.memtable_veirfy_per_key_checksum_on_seek =
       FLAGS_memtable_veirfy_per_key_checksum_on_seek;

include/rocksdb/advanced_options.h

@@ -213,8 +213,7 @@ enum class VerifyOutputFlags : uint32_t {
                                   // by comparing the one inserted into a
                                   // file, and what is read back.
 
-  // TODO - Implement
-  // kVerifyFileChecksum = 1 << 2,   // Verify file-level checksum
+  kVerifyFileChecksum = 1 << 2,  // Verify file-level checksum
 
   // Second set of bits: when to enable verification
   kEnableForLocalCompaction = 1 << 10,   // Enable for local compaction

options/options_settable_test.cc

@@ -694,7 +694,7 @@ TEST_F(OptionsSettableTest, ColumnFamilyOptionsAllFieldsSettable) {
       "memtable_op_scan_flush_trigger=123;"
       "memtable_avg_op_scan_flush_trigger=12;"
       "cf_allow_ingest_behind=1;"
-      "verify_output_flags=2049;",
+      "verify_output_flags=2053;",
       new_options));
 
   ASSERT_NE(new_options->blob_cache.get(), nullptr);

options/options_test.cc

@@ -1752,6 +1752,19 @@ TEST_F(OptionsTest, MutableCFOptions) {
       config_options, cf_opts, {{"paranoid_file_checks", "false"}}, &cf_opts));
   ASSERT_EQ(cf_opts.paranoid_file_checks, false);
 
+  // Test verify_output_flags with kVerifyFileChecksum (bit 2)
+  // 2052 = 4 (kVerifyFileChecksum) | 2048 (kEnableForRemoteCompaction)
+  ASSERT_OK(GetColumnFamilyOptionsFromString(
+      config_options, cf_opts, "verify_output_flags=2052;", &cf_opts));
+  ASSERT_NE(
+      (cf_opts.verify_output_flags & VerifyOutputFlags::kVerifyFileChecksum),
+      VerifyOutputFlags::kVerifyNone);
+  ASSERT_NE((cf_opts.verify_output_flags &
+             VerifyOutputFlags::kEnableForRemoteCompaction),
+            VerifyOutputFlags::kVerifyNone);
+  ASSERT_EQ(
+      (cf_opts.verify_output_flags & VerifyOutputFlags::kVerifyBlockChecksum),
+      VerifyOutputFlags::kVerifyNone);
   // Should replace the factory with the new setting
   ASSERT_OK(GetColumnFamilyOptionsFromMap(
       config_options, cf_opts,

tools/db_crashtest.py

@@ -429,6 +429,13 @@ def apply_random_seed_per_iteration():
     "check_multiget_entity_consistency": lambda: random.choice([0, 0, 0, 1]),
     "use_timed_put_one_in": lambda: random.choice([0] * 7 + [1, 5, 10]),
     "universal_max_read_amp": lambda: random.choice([-1] * 3 + [0, 4, 10]),
+    # verify_output_flags is a bitmask: bits 0-2 are verification types
+    # (block checksum, iteration, file checksum), bits 10-11 are when to
+    # enable (local compaction, remote compaction). 0x407 = all types +
+    # local, 0xC07 = all types + local + remote, 0xFFFFFFFF = all.
+    "verify_output_flags": lambda: random.choice(
+        [0] * 3 + [0x407, 0xC07, 0xFFFFFFFF]
+    ),
     "paranoid_memory_checks": lambda: random.choice([0] * 7 + [1]),
     "memtable_veirfy_per_key_checksum_on_seek": lambda: random.choice([0] * 7 + [1]),
     "allow_unprepared_value": lambda: random.choice([0, 1]),