ceph/src/rocksdb/utilities/lua/rocks_lua_test.cc

   1 //  Copyright (c) 2016, 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 #include <stdio.h>
   7
   8 #if !defined(ROCKSDB_LITE)
   9
  10 #if defined(LUA)
  11
  12 #include <string>
  13
  14 #include "db/db_test_util.h"
  15 #include "port/stack_trace.h"
  16 #include "rocksdb/compaction_filter.h"
  17 #include "rocksdb/db.h"
  18 #include "rocksdb/utilities/lua/rocks_lua_compaction_filter.h"
  19 #include "util/testharness.h"
  20
  21 namespace rocksdb {
  22
  23 class StopOnErrorLogger : public Logger {
  24  public:
  25   using Logger::Logv;
  26   virtual void Logv(const char* format, va_list ap) override {
  27     vfprintf(stderr, format, ap);
  28     fprintf(stderr, "\n");
  29     FAIL();
  30   }
  31 };
  32
  33
  34 class RocksLuaTest : public testing::Test {
  35  public:
  36   RocksLuaTest() : rnd_(301) {
  37     temp_dir_ = test::TmpDir(Env::Default());
  38     db_ = nullptr;
  39   }
  40
  41   std::string RandomString(int len) {
  42     std::string res;
  43     for (int i = 0; i < len; ++i) {
  44       res += rnd_.Uniform(26) + 'a';
  45     }
  46     return res;
  47   }
  48
  49   void CreateDBWithLuaCompactionFilter(
  50       const lua::RocksLuaCompactionFilterOptions& lua_opt,
  51       const std::string& db_path,
  52       std::unordered_map<std::string, std::string>* kvs,
  53       const int kNumFlushes = 5,
  54       std::shared_ptr<rocksdb::lua::RocksLuaCompactionFilterFactory>*
  55           output_factory = nullptr) {
  56     const int kKeySize = 10;
  57     const int kValueSize = 50;
  58     const int kKeysPerFlush = 2;
  59     auto factory =
  60         std::make_shared<rocksdb::lua::RocksLuaCompactionFilterFactory>(
  61             lua_opt);
  62     if (output_factory != nullptr) {
  63       *output_factory = factory;
  64     }
  65
  66     options_ = Options();
  67     options_.create_if_missing = true;
  68     options_.compaction_filter_factory = factory;
  69     options_.disable_auto_compactions = true;
  70     options_.max_bytes_for_level_base =
  71         (kKeySize + kValueSize) * kKeysPerFlush * 2;
  72     options_.max_bytes_for_level_multiplier = 2;
  73     options_.target_file_size_base = (kKeySize + kValueSize) * kKeysPerFlush;
  74     options_.level0_file_num_compaction_trigger = 2;
  75     DestroyDB(db_path, options_);
  76     ASSERT_OK(DB::Open(options_, db_path, &db_));
  77
  78     for (int f = 0; f < kNumFlushes; ++f) {
  79       for (int i = 0; i < kKeysPerFlush; ++i) {
  80         std::string key = RandomString(kKeySize);
  81         std::string value = RandomString(kValueSize);
  82         kvs->insert({key, value});
  83         ASSERT_OK(db_->Put(WriteOptions(), key, value));
  84       }
  85       db_->Flush(FlushOptions());
  86     }
  87   }
  88
  89   ~RocksLuaTest() {
  90     if (db_) {
  91       delete db_;
  92     }
  93   }
  94   std::string temp_dir_;
  95   DB* db_;
  96   Random rnd_;
  97   Options options_;
  98 };
  99
 100 TEST_F(RocksLuaTest, Default) {
 101   // If nothing is set in the LuaCompactionFilterOptions, then
 102   // RocksDB will keep all the key / value pairs, but it will also
 103   // print our error log indicating failure.
 104   std::string db_path = test::PerThreadDBPath(temp_dir_, "rocks_lua_test");
 105
 106   lua::RocksLuaCompactionFilterOptions lua_opt;
 107
 108   std::unordered_map<std::string, std::string> kvs;
 109   CreateDBWithLuaCompactionFilter(lua_opt, db_path, &kvs);
 110
 111   for (auto const& entry : kvs) {
 112     std::string value;
 113     ASSERT_OK(db_->Get(ReadOptions(), entry.first, &value));
 114     ASSERT_EQ(value, entry.second);
 115   }
 116 }
 117
 118 TEST_F(RocksLuaTest, KeepsAll) {
 119   std::string db_path = test::PerThreadDBPath(temp_dir_, "rocks_lua_test");
 120
 121   lua::RocksLuaCompactionFilterOptions lua_opt;
 122   lua_opt.error_log = std::make_shared<StopOnErrorLogger>();
 123   // keeps all the key value pairs
 124   lua_opt.lua_script =
 125       "function Filter(level, key, existing_value)\n"
 126       "  return false, false, \"\"\n"
 127       "end\n"
 128       "\n"
 129       "function FilterMergeOperand(level, key, operand)\n"
 130       "  return false\n"
 131       "end\n"
 132       "function Name()\n"
 133       "  return \"KeepsAll\"\n"
 134       "end\n"
 135       "\n";
 136
 137   std::unordered_map<std::string, std::string> kvs;
 138   CreateDBWithLuaCompactionFilter(lua_opt, db_path, &kvs);
 139
 140   for (auto const& entry : kvs) {
 141     std::string value;
 142     ASSERT_OK(db_->Get(ReadOptions(), entry.first, &value));
 143     ASSERT_EQ(value, entry.second);
 144   }
 145 }
 146
 147 TEST_F(RocksLuaTest, GetName) {
 148   std::string db_path = test::PerThreadDBPath(temp_dir_, "rocks_lua_test");
 149
 150   lua::RocksLuaCompactionFilterOptions lua_opt;
 151   lua_opt.error_log = std::make_shared<StopOnErrorLogger>();
 152   const std::string kScriptName = "SimpleLuaCompactionFilter";
 153   lua_opt.lua_script =
 154       std::string(
 155           "function Filter(level, key, existing_value)\n"
 156           "  return false, false, \"\"\n"
 157           "end\n"
 158           "\n"
 159           "function FilterMergeOperand(level, key, operand)\n"
 160           "  return false\n"
 161           "end\n"
 162           "function Name()\n"
 163           "  return \"") + kScriptName + "\"\n"
 164       "end\n"
 165       "\n";
 166
 167   std::shared_ptr<CompactionFilterFactory> factory =
 168       std::make_shared<lua::RocksLuaCompactionFilterFactory>(lua_opt);
 169   std::string factory_name(factory->Name());
 170   ASSERT_NE(factory_name.find(kScriptName), std::string::npos);
 171 }
 172
 173 TEST_F(RocksLuaTest, RemovesAll) {
 174   std::string db_path = test::PerThreadDBPath(temp_dir_, "rocks_lua_test");
 175
 176   lua::RocksLuaCompactionFilterOptions lua_opt;
 177   lua_opt.error_log = std::make_shared<StopOnErrorLogger>();
 178   // removes all the key value pairs
 179   lua_opt.lua_script =
 180       "function Filter(level, key, existing_value)\n"
 181       "  return true, false, \"\"\n"
 182       "end\n"
 183       "\n"
 184       "function FilterMergeOperand(level, key, operand)\n"
 185       "  return false\n"
 186       "end\n"
 187       "function Name()\n"
 188       "  return \"RemovesAll\"\n"
 189       "end\n"
 190       "\n";
 191
 192   std::unordered_map<std::string, std::string> kvs;
 193   CreateDBWithLuaCompactionFilter(lua_opt, db_path, &kvs);
 194   // Issue full compaction and expect nothing is in the DB.
 195   ASSERT_OK(db_->CompactRange(CompactRangeOptions(), nullptr, nullptr));
 196
 197   for (auto const& entry : kvs) {
 198     std::string value;
 199     auto s = db_->Get(ReadOptions(), entry.first, &value);
 200     ASSERT_TRUE(s.IsNotFound());
 201   }
 202 }
 203
 204 TEST_F(RocksLuaTest, FilterByKey) {
 205   std::string db_path = test::PerThreadDBPath(temp_dir_, "rocks_lua_test");
 206
 207   lua::RocksLuaCompactionFilterOptions lua_opt;
 208   lua_opt.error_log = std::make_shared<StopOnErrorLogger>();
 209   // removes all keys whose initial is less than 'r'
 210   lua_opt.lua_script =
 211       "function Filter(level, key, existing_value)\n"
 212       "  if key:sub(1,1) < 'r' then\n"
 213       "    return true, false, \"\"\n"
 214       "  end\n"
 215       "  return false, false, \"\"\n"
 216       "end\n"
 217       "\n"
 218       "function FilterMergeOperand(level, key, operand)\n"
 219       "  return false\n"
 220       "end\n"
 221       "function Name()\n"
 222       "  return \"KeepsAll\"\n"
 223       "end\n";
 224
 225   std::unordered_map<std::string, std::string> kvs;
 226   CreateDBWithLuaCompactionFilter(lua_opt, db_path, &kvs);
 227   // Issue full compaction and expect nothing is in the DB.
 228   ASSERT_OK(db_->CompactRange(CompactRangeOptions(), nullptr, nullptr));
 229
 230   for (auto const& entry : kvs) {
 231     std::string value;
 232     auto s = db_->Get(ReadOptions(), entry.first, &value);
 233     if (entry.first[0] < 'r') {
 234       ASSERT_TRUE(s.IsNotFound());
 235     } else {
 236       ASSERT_TRUE(s.ok());
 237       ASSERT_TRUE(value == entry.second);
 238     }
 239   }
 240 }
 241
 242 TEST_F(RocksLuaTest, FilterByValue) {
 243   std::string db_path = test::PerThreadDBPath(temp_dir_, "rocks_lua_test");
 244
 245   lua::RocksLuaCompactionFilterOptions lua_opt;
 246   lua_opt.error_log = std::make_shared<StopOnErrorLogger>();
 247   // removes all values whose initial is less than 'r'
 248   lua_opt.lua_script =
 249       "function Filter(level, key, existing_value)\n"
 250       "  if existing_value:sub(1,1) < 'r' then\n"
 251       "    return true, false, \"\"\n"
 252       "  end\n"
 253       "  return false, false, \"\"\n"
 254       "end\n"
 255       "\n"
 256       "function FilterMergeOperand(level, key, operand)\n"
 257       "  return false\n"
 258       "end\n"
 259       "function Name()\n"
 260       "  return \"FilterByValue\"\n"
 261       "end\n"
 262       "\n";
 263
 264   std::unordered_map<std::string, std::string> kvs;
 265   CreateDBWithLuaCompactionFilter(lua_opt, db_path, &kvs);
 266   // Issue full compaction and expect nothing is in the DB.
 267   ASSERT_OK(db_->CompactRange(CompactRangeOptions(), nullptr, nullptr));
 268
 269   for (auto const& entry : kvs) {
 270     std::string value;
 271     auto s = db_->Get(ReadOptions(), entry.first, &value);
 272     if (entry.second[0] < 'r') {
 273       ASSERT_TRUE(s.IsNotFound());
 274     } else {
 275       ASSERT_TRUE(s.ok());
 276       ASSERT_EQ(value, entry.second);
 277     }
 278   }
 279 }
 280
 281 TEST_F(RocksLuaTest, ChangeValue) {
 282   std::string db_path = test::PerThreadDBPath(temp_dir_, "rocks_lua_test");
 283
 284   lua::RocksLuaCompactionFilterOptions lua_opt;
 285   lua_opt.error_log = std::make_shared<StopOnErrorLogger>();
 286   // Replace all values by their reversed key
 287   lua_opt.lua_script =
 288       "function Filter(level, key, existing_value)\n"
 289       "  return false, true, key:reverse()\n"
 290       "end\n"
 291       "\n"
 292       "function FilterMergeOperand(level, key, operand)\n"
 293       "  return false\n"
 294       "end\n"
 295       "function Name()\n"
 296       "  return \"ChangeValue\"\n"
 297       "end\n"
 298       "\n";
 299
 300   std::unordered_map<std::string, std::string> kvs;
 301   CreateDBWithLuaCompactionFilter(lua_opt, db_path, &kvs);
 302   // Issue full compaction and expect nothing is in the DB.
 303   ASSERT_OK(db_->CompactRange(CompactRangeOptions(), nullptr, nullptr));
 304
 305   for (auto const& entry : kvs) {
 306     std::string value;
 307     ASSERT_OK(db_->Get(ReadOptions(), entry.first, &value));
 308     std::string new_value = entry.first;
 309     std::reverse(new_value.begin(), new_value.end());
 310     ASSERT_EQ(value, new_value);
 311   }
 312 }
 313
 314 TEST_F(RocksLuaTest, ConditionallyChangeAndFilterValue) {
 315   std::string db_path = test::PerThreadDBPath(temp_dir_, "rocks_lua_test");
 316
 317   lua::RocksLuaCompactionFilterOptions lua_opt;
 318   lua_opt.error_log = std::make_shared<StopOnErrorLogger>();
 319   // Performs the following logic:
 320   // If key[0] < 'h' --> replace value by reverse key
 321   // If key[0] >= 'r' --> keep the original key value
 322   // Otherwise, filter the key value
 323   lua_opt.lua_script =
 324       "function Filter(level, key, existing_value)\n"
 325       "  if key:sub(1,1) < 'h' then\n"
 326       "    return false, true, key:reverse()\n"
 327       "  elseif key:sub(1,1) < 'r' then\n"
 328       "    return true, false, \"\"\n"
 329       "  end\n"
 330       "  return false, false, \"\"\n"
 331       "end\n"
 332       "function Name()\n"
 333       "  return \"ConditionallyChangeAndFilterValue\"\n"
 334       "end\n"
 335       "\n";
 336
 337   std::unordered_map<std::string, std::string> kvs;
 338   CreateDBWithLuaCompactionFilter(lua_opt, db_path, &kvs);
 339   // Issue full compaction and expect nothing is in the DB.
 340   ASSERT_OK(db_->CompactRange(CompactRangeOptions(), nullptr, nullptr));
 341
 342   for (auto const& entry : kvs) {
 343     std::string value;
 344     auto s = db_->Get(ReadOptions(), entry.first, &value);
 345     if (entry.first[0] < 'h') {
 346       ASSERT_TRUE(s.ok());
 347       std::string new_value = entry.first;
 348       std::reverse(new_value.begin(), new_value.end());
 349       ASSERT_EQ(value, new_value);
 350     } else if (entry.first[0] < 'r') {
 351       ASSERT_TRUE(s.IsNotFound());
 352     } else {
 353       ASSERT_TRUE(s.ok());
 354       ASSERT_EQ(value, entry.second);
 355     }
 356   }
 357 }
 358
 359 TEST_F(RocksLuaTest, DynamicChangeScript) {
 360   std::string db_path = test::PerThreadDBPath(temp_dir_, "rocks_lua_test");
 361
 362   lua::RocksLuaCompactionFilterOptions lua_opt;
 363   lua_opt.error_log = std::make_shared<StopOnErrorLogger>();
 364   // keeps all the key value pairs
 365   lua_opt.lua_script =
 366       "function Filter(level, key, existing_value)\n"
 367       "  return false, false, \"\"\n"
 368       "end\n"
 369       "\n"
 370       "function FilterMergeOperand(level, key, operand)\n"
 371       "  return false\n"
 372       "end\n"
 373       "function Name()\n"
 374       "  return \"KeepsAll\"\n"
 375       "end\n"
 376       "\n";
 377
 378   std::unordered_map<std::string, std::string> kvs;
 379   std::shared_ptr<rocksdb::lua::RocksLuaCompactionFilterFactory> factory;
 380   CreateDBWithLuaCompactionFilter(lua_opt, db_path, &kvs, 30, &factory);
 381   uint64_t count = 0;
 382   ASSERT_TRUE(db_->GetIntProperty(
 383       rocksdb::DB::Properties::kNumEntriesActiveMemTable, &count));
 384   ASSERT_EQ(count, 0);
 385   ASSERT_TRUE(db_->GetIntProperty(
 386       rocksdb::DB::Properties::kNumEntriesImmMemTables, &count));
 387   ASSERT_EQ(count, 0);
 388
 389   CompactRangeOptions cr_opt;
 390   cr_opt.bottommost_level_compaction =
 391       rocksdb::BottommostLevelCompaction::kForce;
 392
 393   // Issue full compaction and expect everything is in the DB.
 394   ASSERT_OK(db_->CompactRange(cr_opt, nullptr, nullptr));
 395
 396   for (auto const& entry : kvs) {
 397     std::string value;
 398     ASSERT_OK(db_->Get(ReadOptions(), entry.first, &value));
 399     ASSERT_EQ(value, entry.second);
 400   }
 401
 402   // change the lua script to removes all the key value pairs
 403   factory->SetScript(
 404       "function Filter(level, key, existing_value)\n"
 405       "  return true, false, \"\"\n"
 406       "end\n"
 407       "\n"
 408       "function FilterMergeOperand(level, key, operand)\n"
 409       "  return false\n"
 410       "end\n"
 411       "function Name()\n"
 412       "  return \"RemovesAll\"\n"
 413       "end\n"
 414       "\n");
 415   {
 416     std::string key = "another-key";
 417     std::string value = "another-value";
 418     kvs.insert({key, value});
 419     ASSERT_OK(db_->Put(WriteOptions(), key, value));
 420     db_->Flush(FlushOptions());
 421   }
 422
 423   cr_opt.change_level = true;
 424   cr_opt.target_level = 5;
 425   // Issue full compaction and expect nothing is in the DB.
 426   ASSERT_OK(db_->CompactRange(cr_opt, nullptr, nullptr));
 427
 428   for (auto const& entry : kvs) {
 429     std::string value;
 430     auto s = db_->Get(ReadOptions(), entry.first, &value);
 431     ASSERT_TRUE(s.IsNotFound());
 432   }
 433 }
 434
 435 TEST_F(RocksLuaTest, LuaConditionalTypeError) {
 436   std::string db_path = test::PerThreadDBPath(temp_dir_, "rocks_lua_test");
 437
 438   lua::RocksLuaCompactionFilterOptions lua_opt;
 439   // Filter() error when input key's initial >= 'r'
 440   lua_opt.lua_script =
 441       "function Filter(level, key, existing_value)\n"
 442       "  if existing_value:sub(1,1) >= 'r' then\n"
 443       "    return true, 2, \"\" -- incorrect type of 2nd return value\n"
 444       "  end\n"
 445       "  return true, false, \"\"\n"
 446       "end\n"
 447       "\n"
 448       "function FilterMergeOperand(level, key, operand)\n"
 449       "  return false\n"
 450       "end\n"
 451       "function Name()\n"
 452       "  return \"BuggyCode\"\n"
 453       "end\n"
 454       "\n";
 455
 456   std::unordered_map<std::string, std::string> kvs;
 457   // Create DB with 10 files
 458   CreateDBWithLuaCompactionFilter(lua_opt, db_path, &kvs, 10);
 459
 460   // Issue full compaction and expect all keys which initial is < 'r'
 461   // will be deleted as we keep the key value when we hit an error.
 462   ASSERT_OK(db_->CompactRange(CompactRangeOptions(), nullptr, nullptr));
 463
 464   for (auto const& entry : kvs) {
 465     std::string value;
 466     auto s = db_->Get(ReadOptions(), entry.first, &value);
 467     if (entry.second[0] < 'r') {
 468       ASSERT_TRUE(s.IsNotFound());
 469     } else {
 470       ASSERT_TRUE(s.ok());
 471       ASSERT_EQ(value, entry.second);
 472     }
 473   }
 474 }
 475
 476 }  // namespace rocksdb
 477
 478 int main(int argc, char** argv) {
 479   rocksdb::port::InstallStackTraceHandler();
 480   ::testing::InitGoogleTest(&argc, argv);
 481   return RUN_ALL_TESTS();
 482 }
 483
 484 #else
 485
 486 int main(int /*argc*/, char** /*argv*/) {
 487   printf("LUA_PATH is not set.  Ignoring the test.\n");
 488 }
 489
 490 #endif  // defined(LUA)
 491
 492 #else
 493
 494 int main(int /*argc*/, char** /*argv*/) {
 495   printf("Lua is not supported in RocksDBLite.  Ignoring the test.\n");
 496 }
 497
 498 #endif  // !defined(ROCKSDB_LITE)