1 // Copyright (c) 2016-present, Facebook, Inc.  All rights reserved.
2 //  This source code is licensed under both the GPLv2 (found in the
3 //  COPYING file in the root directory) and Apache 2.0 License
4 //  (found in the LICENSE.Apache file in the root directory).
5 
6 #ifndef ROCKSDB_LITE
7 
8 #include "rocksdb/utilities/object_registry.h"
9 
10 #include "rocksdb/customizable.h"
11 #include "test_util/testharness.h"
12 
13 namespace ROCKSDB_NAMESPACE {
14 
15 class EnvRegistryTest : public testing::Test {
16  public:
17   static int num_a, num_b;
18 };
19 
20 int EnvRegistryTest::num_a = 0;
21 int EnvRegistryTest::num_b = 0;
22 static FactoryFunc<Env> test_reg_a = ObjectLibrary::Default()->Register<Env>(
23     "a://.*",
24     [](const std::string& /*uri*/, std::unique_ptr<Env>* /*env_guard*/,
__anon8cda92740102(const std::string& , std::unique_ptr<Env>* , std::string* ) 25        std::string* /* errmsg */) {
26       ++EnvRegistryTest::num_a;
27       return Env::Default();
28     });
29 
30 static FactoryFunc<Env> test_reg_b = ObjectLibrary::Default()->Register<Env>(
31     "b://.*", [](const std::string& /*uri*/, std::unique_ptr<Env>* env_guard,
__anon8cda92740202(const std::string& , std::unique_ptr<Env>* env_guard, std::string* ) 32                  std::string* /* errmsg */) {
33       ++EnvRegistryTest::num_b;
34       // Env::Default() is a singleton so we can't grant ownership directly to
35       // the caller - we must wrap it first.
36       env_guard->reset(new EnvWrapper(Env::Default()));
37       return env_guard->get();
38     });
39 
TEST_F(EnvRegistryTest,Basics)40 TEST_F(EnvRegistryTest, Basics) {
41   std::string msg;
42   std::unique_ptr<Env> env_guard;
43   auto registry = ObjectRegistry::NewInstance();
44   auto res = registry->NewObject<Env>("a://test", &env_guard, &msg);
45   ASSERT_NE(res, nullptr);
46   ASSERT_EQ(env_guard, nullptr);
47   ASSERT_EQ(1, num_a);
48   ASSERT_EQ(0, num_b);
49 
50   res = registry->NewObject<Env>("b://test", &env_guard, &msg);
51   ASSERT_NE(res, nullptr);
52   ASSERT_NE(env_guard, nullptr);
53   ASSERT_EQ(1, num_a);
54   ASSERT_EQ(1, num_b);
55 
56   res = registry->NewObject<Env>("c://test", &env_guard, &msg);
57   ASSERT_EQ(res, nullptr);
58   ASSERT_EQ(env_guard, nullptr);
59   ASSERT_EQ(1, num_a);
60   ASSERT_EQ(1, num_b);
61 }
62 
TEST_F(EnvRegistryTest,LocalRegistry)63 TEST_F(EnvRegistryTest, LocalRegistry) {
64   std::string msg;
65   std::unique_ptr<Env> guard;
66   auto registry = ObjectRegistry::NewInstance();
67   std::shared_ptr<ObjectLibrary> library =
68       std::make_shared<ObjectLibrary>("local");
69   registry->AddLibrary(library);
70   library->Register<Env>(
71       "test-local",
72       [](const std::string& /*uri*/, std::unique_ptr<Env>* /*guard */,
73          std::string* /* errmsg */) { return Env::Default(); });
74 
75   ObjectLibrary::Default()->Register<Env>(
76       "test-global",
77       [](const std::string& /*uri*/, std::unique_ptr<Env>* /*guard */,
78          std::string* /* errmsg */) { return Env::Default(); });
79 
80   ASSERT_EQ(
81       ObjectRegistry::NewInstance()->NewObject<Env>("test-local", &guard, &msg),
82       nullptr);
83   ASSERT_NE(
84       ObjectRegistry::NewInstance()->NewObject("test-global", &guard, &msg),
85       nullptr);
86   ASSERT_NE(registry->NewObject<Env>("test-local", &guard, &msg), nullptr);
87   ASSERT_NE(registry->NewObject<Env>("test-global", &guard, &msg), nullptr);
88 }
89 
TEST_F(EnvRegistryTest,CheckShared)90 TEST_F(EnvRegistryTest, CheckShared) {
91   std::shared_ptr<Env> shared;
92   std::shared_ptr<ObjectRegistry> registry = ObjectRegistry::NewInstance();
93   std::shared_ptr<ObjectLibrary> library =
94       std::make_shared<ObjectLibrary>("shared");
95   registry->AddLibrary(library);
96   library->Register<Env>(
97       "unguarded",
98       [](const std::string& /*uri*/, std::unique_ptr<Env>* /*guard */,
99          std::string* /* errmsg */) { return Env::Default(); });
100 
101   library->Register<Env>(
102       "guarded", [](const std::string& /*uri*/, std::unique_ptr<Env>* guard,
103                     std::string* /* errmsg */) {
104         guard->reset(new EnvWrapper(Env::Default()));
105         return guard->get();
106       });
107 
108   ASSERT_OK(registry->NewSharedObject<Env>("guarded", &shared));
109   ASSERT_NE(shared, nullptr);
110   shared.reset();
111   ASSERT_NOK(registry->NewSharedObject<Env>("unguarded", &shared));
112   ASSERT_EQ(shared, nullptr);
113 }
114 
TEST_F(EnvRegistryTest,CheckStatic)115 TEST_F(EnvRegistryTest, CheckStatic) {
116   Env* env = nullptr;
117   std::shared_ptr<ObjectRegistry> registry = ObjectRegistry::NewInstance();
118   std::shared_ptr<ObjectLibrary> library =
119       std::make_shared<ObjectLibrary>("static");
120   registry->AddLibrary(library);
121   library->Register<Env>(
122       "unguarded",
123       [](const std::string& /*uri*/, std::unique_ptr<Env>* /*guard */,
124          std::string* /* errmsg */) { return Env::Default(); });
125 
126   library->Register<Env>(
127       "guarded", [](const std::string& /*uri*/, std::unique_ptr<Env>* guard,
128                     std::string* /* errmsg */) {
129         guard->reset(new EnvWrapper(Env::Default()));
130         return guard->get();
131       });
132 
133   ASSERT_NOK(registry->NewStaticObject<Env>("guarded", &env));
134   ASSERT_EQ(env, nullptr);
135   env = nullptr;
136   ASSERT_OK(registry->NewStaticObject<Env>("unguarded", &env));
137   ASSERT_NE(env, nullptr);
138 }
139 
TEST_F(EnvRegistryTest,CheckUnique)140 TEST_F(EnvRegistryTest, CheckUnique) {
141   std::unique_ptr<Env> unique;
142   std::shared_ptr<ObjectRegistry> registry = ObjectRegistry::NewInstance();
143   std::shared_ptr<ObjectLibrary> library =
144       std::make_shared<ObjectLibrary>("unique");
145   registry->AddLibrary(library);
146   library->Register<Env>(
147       "unguarded",
148       [](const std::string& /*uri*/, std::unique_ptr<Env>* /*guard */,
149          std::string* /* errmsg */) { return Env::Default(); });
150 
151   library->Register<Env>(
152       "guarded", [](const std::string& /*uri*/, std::unique_ptr<Env>* guard,
153                     std::string* /* errmsg */) {
154         guard->reset(new EnvWrapper(Env::Default()));
155         return guard->get();
156       });
157 
158   ASSERT_OK(registry->NewUniqueObject<Env>("guarded", &unique));
159   ASSERT_NE(unique, nullptr);
160   unique.reset();
161   ASSERT_NOK(registry->NewUniqueObject<Env>("unguarded", &unique));
162   ASSERT_EQ(unique, nullptr);
163 }
164 
TEST_F(EnvRegistryTest,TestRegistryParents)165 TEST_F(EnvRegistryTest, TestRegistryParents) {
166   auto grand = ObjectRegistry::Default();
167   auto parent = ObjectRegistry::NewInstance();  // parent with a grandparent
168   auto uncle = ObjectRegistry::NewInstance(grand);
169   auto child = ObjectRegistry::NewInstance(parent);
170   auto cousin = ObjectRegistry::NewInstance(uncle);
171 
172   auto library = parent->AddLibrary("parent");
173   library->Register<Env>(
174       "parent", [](const std::string& /*uri*/, std::unique_ptr<Env>* guard,
175                    std::string* /* errmsg */) {
176         guard->reset(new EnvWrapper(Env::Default()));
177         return guard->get();
178       });
179   library = cousin->AddLibrary("cousin");
180   library->Register<Env>(
181       "cousin", [](const std::string& /*uri*/, std::unique_ptr<Env>* guard,
182                    std::string* /* errmsg */) {
183         guard->reset(new EnvWrapper(Env::Default()));
184         return guard->get();
185       });
186 
187   std::unique_ptr<Env> guard;
188   std::string msg;
189 
190   // a:://* is registered in Default, so they should all workd
191   ASSERT_NE(parent->NewObject<Env>("a://test", &guard, &msg), nullptr);
192   ASSERT_NE(child->NewObject<Env>("a://test", &guard, &msg), nullptr);
193   ASSERT_NE(uncle->NewObject<Env>("a://test", &guard, &msg), nullptr);
194   ASSERT_NE(cousin->NewObject<Env>("a://test", &guard, &msg), nullptr);
195 
196   // The parent env is only registered for parent, not uncle,
197   // So parent and child should return success and uncle and cousin should fail
198   ASSERT_OK(parent->NewUniqueObject<Env>("parent", &guard));
199   ASSERT_OK(child->NewUniqueObject<Env>("parent", &guard));
200   ASSERT_NOK(uncle->NewUniqueObject<Env>("parent", &guard));
201   ASSERT_NOK(cousin->NewUniqueObject<Env>("parent", &guard));
202 
203   // The cousin is only registered in the cousin, so all of the others should
204   // fail
205   ASSERT_OK(cousin->NewUniqueObject<Env>("cousin", &guard));
206   ASSERT_NOK(parent->NewUniqueObject<Env>("cousin", &guard));
207   ASSERT_NOK(child->NewUniqueObject<Env>("cousin", &guard));
208   ASSERT_NOK(uncle->NewUniqueObject<Env>("cousin", &guard));
209 }
210 class MyCustomizable : public Customizable {
211  public:
Type()212   static const char* Type() { return "MyCustomizable"; }
MyCustomizable(const char * prefix,const std::string & id)213   MyCustomizable(const char* prefix, const std::string& id) : id_(id) {
214     name_ = id_.substr(0, strlen(prefix) - 1);
215   }
Name() const216   const char* Name() const override { return name_.c_str(); }
GetId() const217   std::string GetId() const override { return id_; }
218 
219  private:
220   std::string id_;
221   std::string name_;
222 };
223 
TEST_F(EnvRegistryTest,TestManagedObjects)224 TEST_F(EnvRegistryTest, TestManagedObjects) {
225   auto registry = ObjectRegistry::NewInstance();
226   auto m_a1 = std::make_shared<MyCustomizable>("", "A");
227   auto m_a2 = std::make_shared<MyCustomizable>("", "A");
228 
229   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), nullptr);
230   ASSERT_OK(registry->SetManagedObject<MyCustomizable>(m_a1));
231   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), m_a1);
232 
233   ASSERT_NOK(registry->SetManagedObject<MyCustomizable>(m_a2));
234   ASSERT_OK(registry->SetManagedObject<MyCustomizable>(m_a1));
235   m_a1.reset();
236   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), nullptr);
237   ASSERT_OK(registry->SetManagedObject<MyCustomizable>(m_a2));
238   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), m_a2);
239 }
240 
TEST_F(EnvRegistryTest,TestTwoManagedObjects)241 TEST_F(EnvRegistryTest, TestTwoManagedObjects) {
242   auto registry = ObjectRegistry::NewInstance();
243   auto m_a = std::make_shared<MyCustomizable>("", "A");
244   auto m_b = std::make_shared<MyCustomizable>("", "B");
245   std::vector<std::shared_ptr<MyCustomizable>> objects;
246 
247   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), nullptr);
248   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("B"), nullptr);
249   ASSERT_OK(registry->ListManagedObjects(&objects));
250   ASSERT_EQ(objects.size(), 0U);
251   ASSERT_OK(registry->SetManagedObject(m_a));
252   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("B"), nullptr);
253   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), m_a);
254   ASSERT_OK(registry->ListManagedObjects(&objects));
255   ASSERT_EQ(objects.size(), 1U);
256   ASSERT_EQ(objects.front(), m_a);
257 
258   ASSERT_OK(registry->SetManagedObject(m_b));
259   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), m_a);
260   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("B"), m_b);
261   ASSERT_OK(registry->ListManagedObjects(&objects));
262   ASSERT_EQ(objects.size(), 2U);
263   ASSERT_OK(registry->ListManagedObjects("A", &objects));
264   ASSERT_EQ(objects.size(), 1U);
265   ASSERT_EQ(objects.front(), m_a);
266   ASSERT_OK(registry->ListManagedObjects("B", &objects));
267   ASSERT_EQ(objects.size(), 1U);
268   ASSERT_EQ(objects.front(), m_b);
269   ASSERT_OK(registry->ListManagedObjects("C", &objects));
270   ASSERT_EQ(objects.size(), 0U);
271 
272   m_a.reset();
273   objects.clear();
274 
275   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("B"), m_b);
276   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), nullptr);
277   ASSERT_OK(registry->ListManagedObjects(&objects));
278   ASSERT_EQ(objects.size(), 1U);
279   ASSERT_EQ(objects.front(), m_b);
280 
281   m_b.reset();
282   objects.clear();
283   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), nullptr);
284   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("B"), nullptr);
285 }
286 
TEST_F(EnvRegistryTest,TestAlternateNames)287 TEST_F(EnvRegistryTest, TestAlternateNames) {
288   auto registry = ObjectRegistry::NewInstance();
289   auto m_a = std::make_shared<MyCustomizable>("", "A");
290   auto m_b = std::make_shared<MyCustomizable>("", "B");
291   std::vector<std::shared_ptr<MyCustomizable>> objects;
292   // Test no objects exist
293   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), nullptr);
294   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("B"), nullptr);
295   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("TheOne"), nullptr);
296   ASSERT_OK(registry->ListManagedObjects(&objects));
297   ASSERT_EQ(objects.size(), 0U);
298 
299   // Mark "TheOne" to be A
300   ASSERT_OK(registry->SetManagedObject("TheOne", m_a));
301   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("B"), nullptr);
302   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), nullptr);
303   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("TheOne"), m_a);
304   ASSERT_OK(registry->ListManagedObjects(&objects));
305   ASSERT_EQ(objects.size(), 1U);
306   ASSERT_EQ(objects.front(), m_a);
307 
308   // Try to mark "TheOne" again.
309   ASSERT_NOK(registry->SetManagedObject("TheOne", m_b));
310   ASSERT_OK(registry->SetManagedObject("TheOne", m_a));
311 
312   // Add "A" as a managed object.  Registered 2x
313   ASSERT_OK(registry->SetManagedObject(m_a));
314   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("B"), nullptr);
315   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), m_a);
316   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("TheOne"), m_a);
317   ASSERT_OK(registry->ListManagedObjects(&objects));
318   ASSERT_EQ(objects.size(), 2U);
319 
320   // Delete "A".
321   m_a.reset();
322   objects.clear();
323 
324   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("TheOne"), nullptr);
325   ASSERT_OK(registry->SetManagedObject("TheOne", m_b));
326   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("TheOne"), m_b);
327   ASSERT_OK(registry->ListManagedObjects(&objects));
328   ASSERT_EQ(objects.size(), 1U);
329   ASSERT_EQ(objects.front(), m_b);
330 
331   m_b.reset();
332   objects.clear();
333   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), nullptr);
334   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), nullptr);
335   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("TheOne"), nullptr);
336   ASSERT_OK(registry->ListManagedObjects(&objects));
337   ASSERT_EQ(objects.size(), 0U);
338 }
339 
TEST_F(EnvRegistryTest,TestTwoManagedClasses)340 TEST_F(EnvRegistryTest, TestTwoManagedClasses) {
341   class MyCustomizable2 : public MyCustomizable {
342    public:
343     static const char* Type() { return "MyCustomizable2"; }
344     MyCustomizable2(const char* prefix, const std::string& id)
345         : MyCustomizable(prefix, id) {}
346   };
347 
348   auto registry = ObjectRegistry::NewInstance();
349   auto m_a1 = std::make_shared<MyCustomizable>("", "A");
350   auto m_a2 = std::make_shared<MyCustomizable2>("", "A");
351   std::vector<std::shared_ptr<MyCustomizable>> obj1s;
352   std::vector<std::shared_ptr<MyCustomizable2>> obj2s;
353 
354   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), nullptr);
355   ASSERT_EQ(registry->GetManagedObject<MyCustomizable2>("A"), nullptr);
356 
357   ASSERT_OK(registry->SetManagedObject(m_a1));
358   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), m_a1);
359   ASSERT_EQ(registry->GetManagedObject<MyCustomizable2>("A"), nullptr);
360 
361   ASSERT_OK(registry->SetManagedObject(m_a2));
362   ASSERT_EQ(registry->GetManagedObject<MyCustomizable2>("A"), m_a2);
363   ASSERT_OK(registry->ListManagedObjects(&obj1s));
364   ASSERT_OK(registry->ListManagedObjects(&obj2s));
365   ASSERT_EQ(obj1s.size(), 1U);
366   ASSERT_EQ(obj2s.size(), 1U);
367   ASSERT_EQ(obj1s.front(), m_a1);
368   ASSERT_EQ(obj2s.front(), m_a2);
369   m_a1.reset();
370   obj1s.clear();
371   obj2s.clear();
372   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), nullptr);
373   ASSERT_EQ(registry->GetManagedObject<MyCustomizable2>("A"), m_a2);
374 
375   m_a2.reset();
376   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), nullptr);
377   ASSERT_EQ(registry->GetManagedObject<MyCustomizable2>("A"), nullptr);
378 }
379 
TEST_F(EnvRegistryTest,TestManagedObjectsWithParent)380 TEST_F(EnvRegistryTest, TestManagedObjectsWithParent) {
381   auto base = ObjectRegistry::NewInstance();
382   auto registry = ObjectRegistry::NewInstance(base);
383 
384   auto m_a = std::make_shared<MyCustomizable>("", "A");
385   auto m_b = std::make_shared<MyCustomizable>("", "A");
386 
387   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), nullptr);
388   ASSERT_OK(base->SetManagedObject(m_a));
389   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), m_a);
390 
391   ASSERT_NOK(registry->SetManagedObject(m_b));
392   ASSERT_OK(registry->SetManagedObject(m_a));
393 
394   m_a.reset();
395   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), nullptr);
396   ASSERT_OK(registry->SetManagedObject(m_b));
397   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), m_b);
398 }
399 
TEST_F(EnvRegistryTest,TestGetOrCreateManagedObject)400 TEST_F(EnvRegistryTest, TestGetOrCreateManagedObject) {
401   auto registry = ObjectRegistry::NewInstance();
402   registry->AddLibrary("test")->Register<MyCustomizable>(
403       "MC(@.*)?",
404       [](const std::string& uri, std::unique_ptr<MyCustomizable>* guard,
405          std::string* /* errmsg */) {
406         guard->reset(new MyCustomizable("MC", uri));
407         return guard->get();
408       });
409   std::shared_ptr<MyCustomizable> m_a, m_b, obj;
410   std::vector<std::shared_ptr<MyCustomizable>> objs;
411 
412   std::unordered_map<std::string, std::string> opt_map;
413 
414   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("MC@A"), nullptr);
415   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("MC@B"), nullptr);
416   ASSERT_OK(registry->GetOrCreateManagedObject("MC@A", &m_a));
417   ASSERT_OK(registry->GetOrCreateManagedObject("MC@B", &m_b));
418   ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("MC@A"), m_a);
419   ASSERT_OK(registry->GetOrCreateManagedObject("MC@A", &obj));
420   ASSERT_EQ(obj, m_a);
421   ASSERT_OK(registry->GetOrCreateManagedObject("MC@B", &obj));
422   ASSERT_EQ(obj, m_b);
423   ASSERT_OK(registry->ListManagedObjects(&objs));
424   ASSERT_EQ(objs.size(), 2U);
425 
426   objs.clear();
427   m_a.reset();
428   obj.reset();
429   ASSERT_OK(registry->GetOrCreateManagedObject("MC@A", &m_a));
430   ASSERT_EQ(1, m_a.use_count());
431   ASSERT_OK(registry->GetOrCreateManagedObject("MC@B", &obj));
432   ASSERT_EQ(2, obj.use_count());
433 }
434 }  // namespace ROCKSDB_NAMESPACE
435 
main(int argc,char ** argv)436 int main(int argc, char** argv) {
437   ::testing::InitGoogleTest(&argc, argv);
438   return RUN_ALL_TESTS();
439 }
440 
441 #else  // ROCKSDB_LITE
442 #include <stdio.h>
443 
main(int,char **)444 int main(int /*argc*/, char** /*argv*/) {
445   fprintf(stderr, "SKIPPED as EnvRegistry is not supported in ROCKSDB_LITE\n");
446   return 0;
447 }
448 
449 #endif  // ROCKSDB_LITE
450