rocksdb/tools/ldb_cmd_test.cc
Yanqin Jin 167cdc9f17 Support custom env in sst_dump (#5845)
Summary:
This PR allows for the creation of custom env when using sst_dump. If
the user does not set options.env or set options.env to nullptr, then sst_dump
will automatically try to create a custom env depending on the path to the sst
file or db directory. In order to use this feature, the user must call
ObjectRegistry::Register() beforehand.

Test Plan (on devserver):
```
$make all && make check
```
All tests must pass to ensure this change does not break anything.
Pull Request resolved: https://github.com/facebook/rocksdb/pull/5845

Differential Revision: D17678038

Pulled By: riversand963

fbshipit-source-id: 58ecb4b3f75246d52b07c4c924a63ee61c1ee626
2019-10-08 19:19:12 -07:00

256 lines
7.7 KiB
C++

// Copyright (c) 2011-present, Facebook, Inc. All rights reserved.
// This source code is licensed under both the GPLv2 (found in the
// COPYING file in the root directory) and Apache 2.0 License
// (found in the LICENSE.Apache file in the root directory).
//
#ifndef ROCKSDB_LITE
#include "rocksdb/utilities/ldb_cmd.h"
#include "port/stack_trace.h"
#include "test_util/sync_point.h"
#include "test_util/testharness.h"
using std::string;
using std::vector;
using std::map;
namespace rocksdb {
class LdbCmdTest : public testing::Test {
public:
LdbCmdTest() : testing::Test() {}
Env* TryLoadCustomOrDefaultEnv() {
const char* test_env_uri = getenv("TEST_ENV_URI");
if (!test_env_uri) {
return Env::Default();
}
Env* env = Env::Default();
Env::LoadEnv(test_env_uri, &env, &env_guard_);
return env;
}
private:
std::shared_ptr<Env> env_guard_;
};
TEST_F(LdbCmdTest, HexToString) {
// map input to expected outputs.
// odd number of "hex" half bytes doesn't make sense
map<string, vector<int>> inputMap = {
{"0x07", {7}}, {"0x5050", {80, 80}}, {"0xFF", {-1}},
{"0x1234", {18, 52}}, {"0xaaAbAC", {-86, -85, -84}}, {"0x1203", {18, 3}},
};
for (const auto& inPair : inputMap) {
auto actual = rocksdb::LDBCommand::HexToString(inPair.first);
auto expected = inPair.second;
for (unsigned int i = 0; i < actual.length(); i++) {
EXPECT_EQ(expected[i], static_cast<int>((signed char) actual[i]));
}
auto reverse = rocksdb::LDBCommand::StringToHex(actual);
EXPECT_STRCASEEQ(inPair.first.c_str(), reverse.c_str());
}
}
TEST_F(LdbCmdTest, HexToStringBadInputs) {
const vector<string> badInputs = {
"0xZZ", "123", "0xx5", "0x111G", "0x123", "Ox12", "0xT", "0x1Q1",
};
for (const auto badInput : badInputs) {
try {
rocksdb::LDBCommand::HexToString(badInput);
std::cerr << "Should fail on bad hex value: " << badInput << "\n";
FAIL();
} catch (...) {
}
}
}
TEST_F(LdbCmdTest, MemEnv) {
Env* base_env = TryLoadCustomOrDefaultEnv();
std::unique_ptr<Env> env(NewMemEnv(base_env));
Options opts;
opts.env = env.get();
opts.create_if_missing = true;
DB* db = nullptr;
std::string dbname = test::TmpDir();
ASSERT_OK(DB::Open(opts, dbname, &db));
WriteOptions wopts;
for (int i = 0; i < 100; i++) {
char buf[16];
snprintf(buf, sizeof(buf), "%08d", i);
ASSERT_OK(db->Put(wopts, buf, buf));
}
FlushOptions fopts;
fopts.wait = true;
ASSERT_OK(db->Flush(fopts));
delete db;
char arg1[] = "./ldb";
char arg2[1024];
snprintf(arg2, sizeof(arg2), "--db=%s", dbname.c_str());
char arg3[] = "dump_live_files";
char* argv[] = {arg1, arg2, arg3};
ASSERT_EQ(0,
LDBCommandRunner::RunCommand(3, argv, opts, LDBOptions(), nullptr));
}
TEST_F(LdbCmdTest, OptionParsing) {
// test parsing flags
Options opts;
opts.env = TryLoadCustomOrDefaultEnv();
{
std::vector<std::string> args;
args.push_back("scan");
args.push_back("--ttl");
args.push_back("--timestamp");
LDBCommand* command = rocksdb::LDBCommand::InitFromCmdLineArgs(
args, opts, LDBOptions(), nullptr);
const std::vector<std::string> flags = command->TEST_GetFlags();
EXPECT_EQ(flags.size(), 2);
EXPECT_EQ(flags[0], "ttl");
EXPECT_EQ(flags[1], "timestamp");
delete command;
}
// test parsing options which contains equal sign in the option value
{
std::vector<std::string> args;
args.push_back("scan");
args.push_back("--db=/dev/shm/ldbtest/");
args.push_back(
"--from='abcd/efg/hijk/lmn/"
"opq:__rst.uvw.xyz?a=3+4+bcd+efghi&jk=lm_no&pq=rst-0&uv=wx-8&yz=a&bcd_"
"ef=gh.ijk'");
LDBCommand* command = rocksdb::LDBCommand::InitFromCmdLineArgs(
args, opts, LDBOptions(), nullptr);
const std::map<std::string, std::string> option_map =
command->TEST_GetOptionMap();
EXPECT_EQ(option_map.at("db"), "/dev/shm/ldbtest/");
EXPECT_EQ(option_map.at("from"),
"'abcd/efg/hijk/lmn/"
"opq:__rst.uvw.xyz?a=3+4+bcd+efghi&jk=lm_no&pq=rst-0&uv=wx-8&yz="
"a&bcd_ef=gh.ijk'");
delete command;
}
}
TEST_F(LdbCmdTest, ListFileTombstone) {
Env* base_env = TryLoadCustomOrDefaultEnv();
std::unique_ptr<Env> env(NewMemEnv(base_env));
Options opts;
opts.env = env.get();
opts.create_if_missing = true;
DB* db = nullptr;
std::string dbname = test::TmpDir();
ASSERT_OK(DB::Open(opts, dbname, &db));
WriteOptions wopts;
ASSERT_OK(db->Put(wopts, "foo", "1"));
ASSERT_OK(db->Put(wopts, "bar", "2"));
FlushOptions fopts;
fopts.wait = true;
ASSERT_OK(db->Flush(fopts));
ASSERT_OK(db->DeleteRange(wopts, db->DefaultColumnFamily(), "foo", "foo2"));
ASSERT_OK(db->DeleteRange(wopts, db->DefaultColumnFamily(), "bar", "foo2"));
ASSERT_OK(db->Flush(fopts));
delete db;
{
char arg1[] = "./ldb";
char arg2[1024];
snprintf(arg2, sizeof(arg2), "--db=%s", dbname.c_str());
char arg3[] = "list_file_range_deletes";
char* argv[] = {arg1, arg2, arg3};
rocksdb::SyncPoint::GetInstance()->SetCallBack(
"ListFileRangeDeletesCommand::DoCommand:BeforePrint", [&](void* arg) {
std::string* out_str = reinterpret_cast<std::string*>(arg);
// Count number of tombstones printed
int num_tb = 0;
const std::string kFingerprintStr = "start: ";
auto offset = out_str->find(kFingerprintStr);
while (offset != std::string::npos) {
num_tb++;
offset =
out_str->find(kFingerprintStr, offset + kFingerprintStr.size());
}
EXPECT_EQ(2, num_tb);
});
rocksdb::SyncPoint::GetInstance()->EnableProcessing();
ASSERT_EQ(
0, LDBCommandRunner::RunCommand(3, argv, opts, LDBOptions(), nullptr));
rocksdb::SyncPoint::GetInstance()->ClearAllCallBacks();
rocksdb::SyncPoint::GetInstance()->DisableProcessing();
}
// Test the case of limiting tombstones
{
char arg1[] = "./ldb";
char arg2[1024];
snprintf(arg2, sizeof(arg2), "--db=%s", dbname.c_str());
char arg3[] = "list_file_range_deletes";
char arg4[] = "--max_keys=1";
char* argv[] = {arg1, arg2, arg3, arg4};
rocksdb::SyncPoint::GetInstance()->SetCallBack(
"ListFileRangeDeletesCommand::DoCommand:BeforePrint", [&](void* arg) {
std::string* out_str = reinterpret_cast<std::string*>(arg);
// Count number of tombstones printed
int num_tb = 0;
const std::string kFingerprintStr = "start: ";
auto offset = out_str->find(kFingerprintStr);
while (offset != std::string::npos) {
num_tb++;
offset =
out_str->find(kFingerprintStr, offset + kFingerprintStr.size());
}
EXPECT_EQ(1, num_tb);
});
rocksdb::SyncPoint::GetInstance()->EnableProcessing();
ASSERT_EQ(
0, LDBCommandRunner::RunCommand(4, argv, opts, LDBOptions(), nullptr));
rocksdb::SyncPoint::GetInstance()->ClearAllCallBacks();
rocksdb::SyncPoint::GetInstance()->DisableProcessing();
}
}
} // 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();
}
#else
#include <stdio.h>
int main(int /*argc*/, char** /*argv*/) {
fprintf(stderr, "SKIPPED as LDBCommand is not supported in ROCKSDB_LITE\n");
return 0;
}
#endif // ROCKSDB_LITE