Deflake ThreadStatus related unit tests (#12858)

Summary:
Unit tests `DBTest.ThreadStatusFlush` and `DBTestWithParam.ThreadStatusSingleCompaction` have been flaky and fail with error message
```
[ RUN      ] DBTest.ThreadStatusFlush
op_count: 0, expected_count 1
thread id: 718113, thread status: , cf_name
thread id: 718114, thread status: , cf_name pikachu
/__w/rocksdb/rocksdb/db/db_test.cc:4817: Failure
Value of: VerifyOperationCount(env_, ThreadStatus::OP_FLUSH, 1)
  Actual: false
Expected: true
[  FAILED  ] DBTest.ThreadStatusFlush (106 ms)

[ RUN      ] DBTestWithParam/DBTestWithParam.ThreadStatusSingleCompaction/0
db/db_test.cc:4673: Failure
Expected equality of these values:
  op_count
    Which is: 0
  expected_count
    Which is: 1
[  FAILED  ] DBTestWithParam/DBTestWithParam.ThreadStatusSingleCompaction/0, where GetParam() = (1, false)
```

One cause for this is that before flush/compaction finishes, we will go through `~WritableFileWriter()`, either for WAL or SST file, and temporarily set thread_operation to UNKNOWN. This UNKNOWN thread operation seem to be there for some stress test verification. This PR fixes these tests by setting the IOActivity in ~WritableFileWriter() for debug build.

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

Test Plan: monitor future test failure.

Reviewed By: hx235

Differential Revision: D59691564

Pulled By: cbi42

fbshipit-source-id: 3f96998bba9d42aba50d1830c2b51bef2dd6705f

2 files changed, 34 insertions, 93 deletions

diff --git a/db_stress_tool/db_stress_env_wrapper.h b/db_stress_tool/db_stress_env_wrapper.h
index ab4a074fc..5ea9e8b6e 100644
--- a/db_stress_tool/db_stress_env_wrapper.h
+++ b/db_stress_tool/db_stress_env_wrapper.h
@@ -13,6 +13,20 @@
 #include "monitoring/thread_status_util.h"
 
 namespace ROCKSDB_NAMESPACE {
+namespace {
+void CheckIOActivity(const IOOptions& options) {
+#ifndef NDEBUG
+  const ThreadStatus::OperationType thread_op =
+      ThreadStatusUtil::GetThreadOperation();
+  Env::IOActivity io_activity =
+      ThreadStatusUtil::TEST_GetExpectedIOActivity(thread_op);
+  assert(io_activity == Env::IOActivity::kUnknown ||
+         io_activity == options.io_activity);
+#else
+  (void)options;
+#endif
+}
+}  // namespace
 class DbStressRandomAccessFileWrapper : public FSRandomAccessFileOwnerWrapper {
  public:
   explicit DbStressRandomAccessFileWrapper(
@@ -83,144 +97,67 @@ class DbStressWritableFileWrapper : public FSWritableFileOwnerWrapper {
 
   IOStatus Append(const Slice& data, const IOOptions& options,
                   IODebugContext* dbg) override {
-#ifndef NDEBUG
-    const ThreadStatus::OperationType thread_op =
-        ThreadStatusUtil::GetThreadOperation();
-    Env::IOActivity io_activity =
-        ThreadStatusUtil::TEST_GetExpectedIOActivity(thread_op);
-    assert(io_activity == Env::IOActivity::kUnknown ||
-           io_activity == options.io_activity);
-#endif
+    CheckIOActivity(options);
     return target()->Append(data, options, dbg);
   }
   IOStatus Append(const Slice& data, const IOOptions& options,
                   const DataVerificationInfo& verification_info,
                   IODebugContext* dbg) override {
-#ifndef NDEBUG
-    const ThreadStatus::OperationType thread_op =
-        ThreadStatusUtil::GetThreadOperation();
-    Env::IOActivity io_activity =
-        ThreadStatusUtil::TEST_GetExpectedIOActivity(thread_op);
-    assert(io_activity == Env::IOActivity::kUnknown ||
-           io_activity == options.io_activity);
-#endif
+    CheckIOActivity(options);
     return target()->Append(data, options, verification_info, dbg);
   }
   IOStatus PositionedAppend(const Slice& data, uint64_t offset,
                             const IOOptions& options,
                             IODebugContext* dbg) override {
-#ifndef NDEBUG
-    const ThreadStatus::OperationType thread_op =
-        ThreadStatusUtil::GetThreadOperation();
-    Env::IOActivity io_activity =
-        ThreadStatusUtil::TEST_GetExpectedIOActivity(thread_op);
-    assert(io_activity == Env::IOActivity::kUnknown ||
-           io_activity == options.io_activity);
-#endif
+    CheckIOActivity(options);
     return target()->PositionedAppend(data, offset, options, dbg);
   }
   IOStatus PositionedAppend(const Slice& data, uint64_t offset,
                             const IOOptions& options,
                             const DataVerificationInfo& verification_info,
                             IODebugContext* dbg) override {
-#ifndef NDEBUG
-    const ThreadStatus::OperationType thread_op =
-        ThreadStatusUtil::GetThreadOperation();
-    Env::IOActivity io_activity =
-        ThreadStatusUtil::TEST_GetExpectedIOActivity(thread_op);
-    assert(io_activity == Env::IOActivity::kUnknown ||
-           io_activity == options.io_activity);
-#endif
+    CheckIOActivity(options);
     return target()->PositionedAppend(data, offset, options, verification_info,
                                       dbg);
   }
 
   IOStatus Truncate(uint64_t size, const IOOptions& options,
                     IODebugContext* dbg) override {
-#ifndef NDEBUG
-    const ThreadStatus::OperationType thread_op =
-        ThreadStatusUtil::GetThreadOperation();
-    Env::IOActivity io_activity =
-        ThreadStatusUtil::TEST_GetExpectedIOActivity(thread_op);
-    assert(io_activity == Env::IOActivity::kUnknown ||
-           io_activity == options.io_activity);
-#endif
+    CheckIOActivity(options);
     return target()->Truncate(size, options, dbg);
   }
 
   IOStatus Close(const IOOptions& options, IODebugContext* dbg) override {
-#ifndef NDEBUG
-    const ThreadStatus::OperationType thread_op =
-        ThreadStatusUtil::GetThreadOperation();
-    Env::IOActivity io_activity =
-        ThreadStatusUtil::TEST_GetExpectedIOActivity(thread_op);
-    assert(io_activity == Env::IOActivity::kUnknown ||
-           io_activity == options.io_activity);
-#endif
+    CheckIOActivity(options);
     return target()->Close(options, dbg);
   }
 
   IOStatus Flush(const IOOptions& options, IODebugContext* dbg) override {
-#ifndef NDEBUG
-    const ThreadStatus::OperationType thread_op =
-        ThreadStatusUtil::GetThreadOperation();
-    Env::IOActivity io_activity =
-        ThreadStatusUtil::TEST_GetExpectedIOActivity(thread_op);
-    assert(io_activity == Env::IOActivity::kUnknown ||
-           io_activity == options.io_activity);
-#endif
+    CheckIOActivity(options);
     return target()->Flush(options, dbg);
   }
 
   IOStatus Sync(const IOOptions& options, IODebugContext* dbg) override {
-#ifndef NDEBUG
-    const ThreadStatus::OperationType thread_op =
-        ThreadStatusUtil::GetThreadOperation();
-    Env::IOActivity io_activity =
-        ThreadStatusUtil::TEST_GetExpectedIOActivity(thread_op);
-    assert(io_activity == Env::IOActivity::kUnknown ||
-           io_activity == options.io_activity);
-#endif
+    CheckIOActivity(options);
     return target()->Sync(options, dbg);
   }
 
   IOStatus Fsync(const IOOptions& options, IODebugContext* dbg) override {
-#ifndef NDEBUG
-    const ThreadStatus::OperationType thread_op =
-        ThreadStatusUtil::GetThreadOperation();
-    Env::IOActivity io_activity =
-        ThreadStatusUtil::TEST_GetExpectedIOActivity(thread_op);
-    assert(io_activity == Env::IOActivity::kUnknown ||
-           io_activity == options.io_activity);
-#endif
+    CheckIOActivity(options);
     return target()->Fsync(options, dbg);
   }
 
 #ifdef ROCKSDB_FALLOCATE_PRESENT
   IOStatus Allocate(uint64_t offset, uint64_t len, const IOOptions& options,
                     IODebugContext* dbg) override {
-#ifndef NDEBUG
-    const ThreadStatus::OperationType thread_op =
-        ThreadStatusUtil::GetThreadOperation();
-    Env::IOActivity io_activity =
-        ThreadStatusUtil::TEST_GetExpectedIOActivity(thread_op);
-    assert(io_activity == Env::IOActivity::kUnknown ||
-           io_activity == options.io_activity);
-#endif
+    CheckIOActivity(options);
     return target()->Allocate(offset, len, options, dbg);
   }
 #endif
 
   IOStatus RangeSync(uint64_t offset, uint64_t nbytes, const IOOptions& options,
                      IODebugContext* dbg) override {
-#ifndef NDEBUG
-    const ThreadStatus::OperationType thread_op =
-        ThreadStatusUtil::GetThreadOperation();
-    Env::IOActivity io_activity =
-        ThreadStatusUtil::TEST_GetExpectedIOActivity(thread_op);
-    assert(io_activity == Env::IOActivity::kUnknown ||
-           io_activity == options.io_activity);
-#endif
+    CheckIOActivity(options);
     return target()->RangeSync(offset, nbytes, options, dbg);
   }
 };
diff --git a/file/writable_file_writer.h b/file/writable_file_writer.h
index 6740b0711..9cf26ead9 100644
--- a/file/writable_file_writer.h
+++ b/file/writable_file_writer.h
@@ -233,13 +233,17 @@ class WritableFileWriter {
   WritableFileWriter& operator=(const WritableFileWriter&) = delete;
 
   ~WritableFileWriter() {
-    ThreadStatus::OperationType cur_op_type =
+    IOOptions io_options;
+#ifndef NDEBUG
+    // This is needed to pass the IOActivity related checks in stress test.
+    // See DbStressWritableFileWrapper.
+    ThreadStatus::OperationType op_type =
         ThreadStatusUtil::GetThreadOperation();
-    ThreadStatusUtil::SetThreadOperation(
-        ThreadStatus::OperationType::OP_UNKNOWN);
-    auto s = Close(IOOptions());
+    io_options.io_activity =
+        ThreadStatusUtil::TEST_GetExpectedIOActivity(op_type);
+#endif
+    auto s = Close(io_options);
     s.PermitUncheckedError();
-    ThreadStatusUtil::SetThreadOperation(cur_op_type);
   }
 
   std::string file_name() const { return file_name_; }