diff --git a/CMakeLists.txt b/CMakeLists.txt index be9b2ddd6..54a608648 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -254,6 +254,7 @@ set(SOURCES utilities/document/json_document.cc utilities/document/json_document_builder.cc utilities/env_mirror.cc + utilities/env_registry.cc utilities/flashcache/flashcache.cc utilities/geodb/geodb_impl.cc utilities/leveldb_options/leveldb_options.cc @@ -421,6 +422,7 @@ set(TESTS utilities/checkpoint/checkpoint_test.cc utilities/document/document_db_test.cc utilities/document/json_document_test.cc + utilities/env_registry_test.cc utilities/geodb/geodb_test.cc utilities/memory/memory_test.cc utilities/merge_operators/string_append/stringappend_test.cc diff --git a/Makefile b/Makefile index c1c003303..d4f41247c 100644 --- a/Makefile +++ b/Makefile @@ -970,6 +970,9 @@ spatial_db_test: utilities/spatialdb/spatial_db_test.o $(LIBOBJECTS) $(TESTHARNE env_mirror_test: utilities/env_mirror_test.o $(LIBOBJECTS) $(TESTHARNESS) $(AM_LINK) +env_registry_test: utilities/env_registry_test.o $(LIBOBJECTS) $(TESTHARNESS) + $(AM_LINK) + ttl_test: utilities/ttl/ttl_test.o $(LIBOBJECTS) $(TESTHARNESS) $(AM_LINK) diff --git a/include/rocksdb/utilities/env_registry.h b/include/rocksdb/utilities/env_registry.h new file mode 100644 index 000000000..4074c87f0 --- /dev/null +++ b/include/rocksdb/utilities/env_registry.h @@ -0,0 +1,45 @@ +// Copyright (c) 2016-present, Facebook, Inc. All rights reserved. +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. An additional grant +// of patent rights can be found in the PATENTS file in the same directory. + +#pragma once + +#ifndef ROCKSDB_LITE + +#include +#include +#include + +#include "rocksdb/env.h" + +namespace rocksdb { + +// Returns a new Env when called with a URI string. Populates the unique_ptr +// argument if granting ownership to caller. +typedef std::function*)> + EnvFactoryFunc; + +// Creates a new Env using the registered factory function corresponding to a +// prefix of uri. +// +// If no prefixes match, returns nullptr. If multiple prefixes match, the +// factory function used is unspecified. +// +// Populates env_guard with result pointer if caller is granted ownership. +Env* NewEnvFromUri(const std::string& uri, std::unique_ptr* env_guard); + +// To register an Env factory function, initialize an EnvRegistrar object with +// static storage duration. For example: +// +// static EnvRegistrar hdfs_reg("hdfs://", &CreateHdfsEnv); +// +// Then, calling NewEnvFromUri("hdfs://some_path", ...) will use CreateHdfsEnv +// to make a new Env. +class EnvRegistrar { + public: + explicit EnvRegistrar(std::string uri_prefix, EnvFactoryFunc env_factory); +}; + +} // namespace rocksdb +#endif // ROCKSDB_LITE diff --git a/src.mk b/src.mk index 16b1001bd..63639c419 100644 --- a/src.mk +++ b/src.mk @@ -120,6 +120,7 @@ LIB_SOURCES = \ utilities/document/json_document_builder.cc \ utilities/document/json_document.cc \ utilities/env_mirror.cc \ + utilities/env_registry.cc \ utilities/flashcache/flashcache.cc \ utilities/geodb/geodb_impl.cc \ utilities/leveldb_options/leveldb_options.cc \ @@ -269,6 +270,7 @@ TEST_BENCH_SOURCES = \ utilities/checkpoint/checkpoint_test.cc \ utilities/document/document_db_test.cc \ utilities/document/json_document_test.cc \ + utilities/env_registry_test.cc \ utilities/geodb/geodb_test.cc \ utilities/memory/memory_test.cc \ utilities/merge_operators/string_append/stringappend_test.cc \ diff --git a/utilities/env_registry.cc b/utilities/env_registry.cc new file mode 100644 index 000000000..1b8e556a4 --- /dev/null +++ b/utilities/env_registry.cc @@ -0,0 +1,47 @@ +// Copyright (c) 2016-present, Facebook, Inc. All rights reserved. +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. An additional grant +// of patent rights can be found in the PATENTS file in the same directory. + +#ifndef ROCKSDB_LITE + +#include "rocksdb/utilities/env_registry.h" + +#include +#include + +namespace rocksdb { + +struct EnvRegistryEntry { + std::string prefix; + EnvFactoryFunc env_factory; +}; + +struct EnvRegistry { + static EnvRegistry* Get() { + static EnvRegistry instance; + return &instance; + } + std::vector entries; + + private: + EnvRegistry() = default; +}; + +Env* NewEnvFromUri(const std::string& uri, std::unique_ptr* env_guard) { + env_guard->reset(); + for (const auto& entry : EnvRegistry::Get()->entries) { + if (uri.compare(0, entry.prefix.size(), entry.prefix) == 0) { + return entry.env_factory(uri, env_guard); + } + } + return nullptr; +} + +EnvRegistrar::EnvRegistrar(std::string uri_prefix, EnvFactoryFunc env_factory) { + EnvRegistry::Get()->entries.emplace_back( + EnvRegistryEntry{std::move(uri_prefix), std::move(env_factory)}); +} + +} // namespace rocksdb +#endif // ROCKSDB_LITE diff --git a/utilities/env_registry_test.cc b/utilities/env_registry_test.cc new file mode 100644 index 000000000..e4f16c096 --- /dev/null +++ b/utilities/env_registry_test.cc @@ -0,0 +1,72 @@ +// Copyright (c) 2016-present, Facebook, Inc. All rights reserved. +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. An additional grant +// of patent rights can be found in the PATENTS file in the same directory. + +#ifndef ROCKSDB_LITE + +#include "rocksdb/utilities/env_registry.h" +#include "util/testharness.h" + +namespace rocksdb { + +class EnvRegistryTest : public testing::Test { + public: + static int num_a, num_b; +}; + +int EnvRegistryTest::num_a = 0; +int EnvRegistryTest::num_b = 0; + +static EnvRegistrar test_reg_a("a://", [](const std::string& uri, + std::unique_ptr* env_guard) { + ++EnvRegistryTest::num_a; + return Env::Default(); +}); + +static EnvRegistrar test_reg_b("b://", [](const std::string& uri, + std::unique_ptr* env_guard) { + ++EnvRegistryTest::num_b; + // Env::Default() is a singleton so we can't grant ownership directly to the + // caller - we must wrap it first. + env_guard->reset(new EnvWrapper(Env::Default())); + return env_guard->get(); +}); + +TEST_F(EnvRegistryTest, Basics) { + std::unique_ptr env_guard; + auto res = NewEnvFromUri("a://test", &env_guard); + ASSERT_NE(res, nullptr); + ASSERT_EQ(env_guard, nullptr); + ASSERT_EQ(1, num_a); + ASSERT_EQ(0, num_b); + + res = NewEnvFromUri("b://test", &env_guard); + ASSERT_NE(res, nullptr); + ASSERT_NE(env_guard, nullptr); + ASSERT_EQ(1, num_a); + ASSERT_EQ(1, num_b); + + res = NewEnvFromUri("c://test", &env_guard); + ASSERT_EQ(res, nullptr); + ASSERT_EQ(env_guard, nullptr); + ASSERT_EQ(1, num_a); + ASSERT_EQ(1, num_b); +} + +} // namespace rocksdb + +int main(int argc, char** argv) { + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} + +#else // ROCKSDB_LITE +#include + +int main(int argc, char** argv) { + fprintf(stderr, "SKIPPED as EnvRegistry is not supported in ROCKSDB_LITE\n"); + return 0; +} + +#endif // ROCKSDB_LITE