diff --git a/HISTORY.md b/HISTORY.md index 201cef2b1..8bcb47505 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -1,5 +1,7 @@ # Rocksdb Change Log ## Unreleased +### New Features +* Support loading custom objects in unit tests. In the affected unit tests, RocksDB will create custom Env objects based on environment variable TEST_ENV_URI. Users need to make sure custom object types are properly registered. For example, a static library should expose a `RegisterCustomObjects` function. By linking the unit test binary with the static library, the unit test can execute this function. ## 6.4.0 (7/30/2019) ### Default Option Change diff --git a/db/column_family_test.cc b/db/column_family_test.cc index 63d987f3c..b55d50a49 100644 --- a/db/column_family_test.cc +++ b/db/column_family_test.cc @@ -17,9 +17,11 @@ #include "memtable/hash_skiplist_rep.h" #include "options/options_parser.h" #include "port/port.h" +#include "port/stack_trace.h" #include "rocksdb/db.h" #include "rocksdb/env.h" #include "rocksdb/iterator.h" +#include "rocksdb/utilities/object_registry.h" #include "test_util/fault_injection_test_env.h" #include "test_util/sync_point.h" #include "test_util/testharness.h" @@ -60,8 +62,18 @@ class EnvCounter : public EnvWrapper { class ColumnFamilyTestBase : public testing::Test { public: - ColumnFamilyTestBase(uint32_t format) : rnd_(139), format_(format) { - env_ = new EnvCounter(Env::Default()); + explicit ColumnFamilyTestBase(uint32_t format) : rnd_(139), format_(format) { + const char* test_env_uri = getenv("TEST_ENV_URI"); + Env* base_env = Env::Default(); + if (test_env_uri) { + Status s = ObjectRegistry::NewInstance()->NewSharedObject(test_env_uri, + &env_guard_); + base_env = env_guard_.get(); + EXPECT_OK(s); + EXPECT_NE(Env::Default(), base_env); + } + EXPECT_NE(nullptr, base_env); + env_ = new EnvCounter(base_env); dbname_ = test::PerThreadDBPath("column_family_test"); db_options_.create_if_missing = true; db_options_.fail_if_options_file_error = true; @@ -532,6 +544,7 @@ class ColumnFamilyTestBase : public testing::Test { std::string dbname_; DB* db_ = nullptr; EnvCounter* env_; + std::shared_ptr env_guard_; Random rnd_; uint32_t format_; }; @@ -3312,7 +3325,17 @@ TEST_P(ColumnFamilyTest, MultipleCFPathsTest) { } // namespace rocksdb +#ifdef ROCKSDB_UNITTESTS_WITH_CUSTOM_OBJECTS_FROM_STATIC_LIBS +extern "C" { +void RegisterCustomObjects(int argc, char** argv); +} +#else +void RegisterCustomObjects(int /*argc*/, char** /*argv*/) {} +#endif // !ROCKSDB_UNITTESTS_WITH_CUSTOM_OBJECTS_FROM_STATIC_LIBS + int main(int argc, char** argv) { + rocksdb::port::InstallStackTraceHandler(); ::testing::InitGoogleTest(&argc, argv); + RegisterCustomObjects(argc, argv); return RUN_ALL_TESTS(); } diff --git a/db/db_basic_test.cc b/db/db_basic_test.cc index dc77fb91a..6104b142a 100644 --- a/db/db_basic_test.cc +++ b/db/db_basic_test.cc @@ -1776,8 +1776,17 @@ INSTANTIATE_TEST_CASE_P(Timestamp, DBBasicTestWithTimestampWithParam, } // namespace rocksdb +#ifdef ROCKSDB_UNITTESTS_WITH_CUSTOM_OBJECTS_FROM_STATIC_LIBS +extern "C" { +void RegisterCustomObjects(int argc, char** argv); +} +#else +void RegisterCustomObjects(int /*argc*/, char** /*argv*/) {} +#endif // !ROCKSDB_UNITTESTS_WITH_CUSTOM_OBJECTS_FROM_STATIC_LIBS + int main(int argc, char** argv) { rocksdb::port::InstallStackTraceHandler(); ::testing::InitGoogleTest(&argc, argv); + RegisterCustomObjects(argc, argv); return RUN_ALL_TESTS(); } diff --git a/db/db_test.cc b/db/db_test.cc index 5c96bec36..ae0481da8 100644 --- a/db/db_test.cc +++ b/db/db_test.cc @@ -6263,8 +6263,17 @@ TEST_F(DBTest, LargeBlockSizeTest) { } // namespace rocksdb +#ifdef ROCKSDB_UNITTESTS_WITH_CUSTOM_OBJECTS_FROM_STATIC_LIBS +extern "C" { +void RegisterCustomObjects(int argc, char** argv); +} +#else +void RegisterCustomObjects(int /*argc*/, char** /*argv*/) {} +#endif // !ROCKSDB_UNITTESTS_WITH_CUSTOM_OBJECTS_FROM_STATIC_LIBS + int main(int argc, char** argv) { rocksdb::port::InstallStackTraceHandler(); ::testing::InitGoogleTest(&argc, argv); + RegisterCustomObjects(argc, argv); return RUN_ALL_TESTS(); } diff --git a/db/db_test2.cc b/db/db_test2.cc index 26604c53a..2c993580a 100644 --- a/db/db_test2.cc +++ b/db/db_test2.cc @@ -3823,8 +3823,17 @@ TEST_F(DBTest2, RowCacheSnapshot) { #endif // ROCKSDB_LITE } // namespace rocksdb +#ifdef ROCKSDB_UNITTESTS_WITH_CUSTOM_OBJECTS_FROM_STATIC_LIBS +extern "C" { +void RegisterCustomObjects(int argc, char** argv); +} +#else +void RegisterCustomObjects(int /*argc*/, char** /*argv*/) {} +#endif // !ROCKSDB_UNITTESTS_WITH_CUSTOM_OBJECTS_FROM_STATIC_LIBS + int main(int argc, char** argv) { rocksdb::port::InstallStackTraceHandler(); ::testing::InitGoogleTest(&argc, argv); + RegisterCustomObjects(argc, argv); return RUN_ALL_TESTS(); } diff --git a/db/db_test_util.cc b/db/db_test_util.cc index ebfc7a9ca..7abee3504 100644 --- a/db/db_test_util.cc +++ b/db/db_test_util.cc @@ -10,6 +10,7 @@ #include "db/db_test_util.h" #include "db/forward_iterator.h" #include "rocksdb/env_encryption.h" +#include "rocksdb/utilities/object_registry.h" namespace rocksdb { @@ -47,20 +48,30 @@ ROT13BlockCipher rot13Cipher_(16); #endif // ROCKSDB_LITE DBTestBase::DBTestBase(const std::string path) - : mem_env_(!getenv("MEM_ENV") ? nullptr : new MockEnv(Env::Default())), -#ifndef ROCKSDB_LITE - encrypted_env_( - !getenv("ENCRYPTED_ENV") - ? nullptr - : NewEncryptedEnv(mem_env_ ? mem_env_ : Env::Default(), - new CTREncryptionProvider(rot13Cipher_))), -#else + : mem_env_(nullptr), encrypted_env_(nullptr), -#endif // ROCKSDB_LITE - env_(new SpecialEnv(encrypted_env_ - ? encrypted_env_ - : (mem_env_ ? mem_env_ : Env::Default()))), option_config_(kDefault) { + const char* test_env_uri = getenv("TEST_ENV_URI"); + Env* base_env = Env::Default(); + if (test_env_uri) { + Status s = ObjectRegistry::NewInstance()->NewSharedObject(test_env_uri, + &env_guard_); + base_env = env_guard_.get(); + EXPECT_OK(s); + EXPECT_NE(Env::Default(), base_env); + } + EXPECT_NE(nullptr, base_env); + if (getenv("MEM_ENV")) { + mem_env_ = new MockEnv(base_env); + } +#ifndef ROCKSDB_LITE + if (getenv("ENCRYPTED_ENV")) { + encrypted_env_ = NewEncryptedEnv(mem_env_ ? mem_env_ : base_env, + new CTREncryptionProvider(rot13Cipher_)); + } +#endif // !ROCKSDB_LITE + env_ = new SpecialEnv(encrypted_env_ ? encrypted_env_ + : (mem_env_ ? mem_env_ : base_env)); env_->SetBackgroundThreads(1, Env::LOW); env_->SetBackgroundThreads(1, Env::HIGH); dbname_ = test::PerThreadDBPath(env_, path); diff --git a/db/db_test_util.h b/db/db_test_util.h index 6e1d0ed7a..667371487 100644 --- a/db/db_test_util.h +++ b/db/db_test_util.h @@ -702,6 +702,7 @@ class DBTestBase : public testing::Test { MockEnv* mem_env_; Env* encrypted_env_; SpecialEnv* env_; + std::shared_ptr env_guard_; DB* db_; std::vector handles_;