Skip to content

Commit

Permalink
Merge bitcoin#27790: walletdb: Add PrefixCursor
Browse files Browse the repository at this point in the history
ba616b9 wallet: Add GetPrefixCursor to DatabaseBatch (Andrew Chow)
1d858b0 walletdb: Handle when database keys are empty (Ryan Ofsky)
84b2f35 walletdb: Consistently clear key and value streams before writing (Andrew Chow)

Pull request description:

  Split from bitcoin#24914 as suggested in bitcoin#24914 (review)

  This PR adds a wallet database cursor that gives a view over all of the records beginning with the same prefix.

ACKs for top commit:
  ryanofsky:
    Code review ACK ba616b9. Just suggested changes since last review
  furszy:
    ACK ba616b9

Tree-SHA512: 38a61849f108d8003d28c599b1ad0421ac9beb3afe14c02f1253e7b4efc3d4eef483e32647a820fc6636bca3f9efeff9fe062b6b602e0cded69f21f8b26af544
  • Loading branch information
fanquake committed Jun 2, 2023
2 parents e43fdfd + ba616b9 commit 7f20197
Show file tree
Hide file tree
Showing 10 changed files with 253 additions and 17 deletions.
29 changes: 24 additions & 5 deletions src/wallet/bdb.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -668,7 +668,8 @@ void BerkeleyDatabase::ReloadDbEnv()
env->ReloadDbEnv();
}

BerkeleyCursor::BerkeleyCursor(BerkeleyDatabase& database, const BerkeleyBatch& batch)
BerkeleyCursor::BerkeleyCursor(BerkeleyDatabase& database, const BerkeleyBatch& batch, Span<const std::byte> prefix)
: m_key_prefix(prefix.begin(), prefix.end())
{
if (!database.m_db.get()) {
throw std::runtime_error(STR_INTERNAL_BUG("BerkeleyDatabase does not exist"));
Expand All @@ -685,19 +686,30 @@ DatabaseCursor::Status BerkeleyCursor::Next(DataStream& ssKey, DataStream& ssVal
{
if (m_cursor == nullptr) return Status::FAIL;
// Read at cursor
SafeDbt datKey;
SafeDbt datKey(m_key_prefix.data(), m_key_prefix.size());
SafeDbt datValue;
int ret = m_cursor->get(datKey, datValue, DB_NEXT);
int ret = -1;
if (m_first && !m_key_prefix.empty()) {
ret = m_cursor->get(datKey, datValue, DB_SET_RANGE);
} else {
ret = m_cursor->get(datKey, datValue, DB_NEXT);
}
m_first = false;
if (ret == DB_NOTFOUND) {
return Status::DONE;
}
if (ret != 0 || datKey.get_data() == nullptr || datValue.get_data() == nullptr) {
if (ret != 0) {
return Status::FAIL;
}

Span<const std::byte> raw_key = {AsBytePtr(datKey.get_data()), datKey.get_size()};
if (!m_key_prefix.empty() && std::mismatch(raw_key.begin(), raw_key.end(), m_key_prefix.begin(), m_key_prefix.end()).second != m_key_prefix.end()) {
return Status::DONE;
}

// Convert to streams
ssKey.clear();
ssKey.write({AsBytePtr(datKey.get_data()), datKey.get_size()});
ssKey.write(raw_key);
ssValue.clear();
ssValue.write({AsBytePtr(datValue.get_data()), datValue.get_size()});
return Status::MORE;
Expand All @@ -716,6 +728,12 @@ std::unique_ptr<DatabaseCursor> BerkeleyBatch::GetNewCursor()
return std::make_unique<BerkeleyCursor>(m_database, *this);
}

std::unique_ptr<DatabaseCursor> BerkeleyBatch::GetNewPrefixCursor(Span<const std::byte> prefix)
{
if (!pdb) return nullptr;
return std::make_unique<BerkeleyCursor>(m_database, *this, prefix);
}

bool BerkeleyBatch::TxnBegin()
{
if (!pdb || activeTxn)
Expand Down Expand Up @@ -777,6 +795,7 @@ bool BerkeleyBatch::ReadKey(DataStream&& key, DataStream& value)
SafeDbt datValue;
int ret = pdb->get(activeTxn, datKey, datValue, 0);
if (ret == 0 && datValue.get_data() != nullptr) {
value.clear();
value.write({AsBytePtr(datValue.get_data()), datValue.get_size()});
return true;
}
Expand Down
7 changes: 6 additions & 1 deletion src/wallet/bdb.h
Original file line number Diff line number Diff line change
Expand Up @@ -190,9 +190,13 @@ class BerkeleyCursor : public DatabaseCursor
{
private:
Dbc* m_cursor;
std::vector<std::byte> m_key_prefix;
bool m_first{true};

public:
explicit BerkeleyCursor(BerkeleyDatabase& database, const BerkeleyBatch& batch);
// Constructor for cursor for records matching the prefix
// To match all records, an empty prefix may be provided.
explicit BerkeleyCursor(BerkeleyDatabase& database, const BerkeleyBatch& batch, Span<const std::byte> prefix = {});
~BerkeleyCursor() override;

Status Next(DataStream& key, DataStream& value) override;
Expand Down Expand Up @@ -229,6 +233,7 @@ class BerkeleyBatch : public DatabaseBatch
void Close() override;

std::unique_ptr<DatabaseCursor> GetNewCursor() override;
std::unique_ptr<DatabaseCursor> GetNewPrefixCursor(Span<const std::byte> prefix) override;
bool TxnBegin() override;
bool TxnCommit() override;
bool TxnAbort() override;
Expand Down
1 change: 1 addition & 0 deletions src/wallet/db.h
Original file line number Diff line number Diff line change
Expand Up @@ -113,6 +113,7 @@ class DatabaseBatch
virtual bool ErasePrefix(Span<const std::byte> prefix) = 0;

virtual std::unique_ptr<DatabaseCursor> GetNewCursor() = 0;
virtual std::unique_ptr<DatabaseCursor> GetNewPrefixCursor(Span<const std::byte> prefix) = 0;
virtual bool TxnBegin() = 0;
virtual bool TxnCommit() = 0;
virtual bool TxnAbort() = 0;
Expand Down
1 change: 1 addition & 0 deletions src/wallet/salvage.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@ class DummyBatch : public DatabaseBatch
void Close() override {}

std::unique_ptr<DatabaseCursor> GetNewCursor() override { return std::make_unique<DummyCursor>(); }
std::unique_ptr<DatabaseCursor> GetNewPrefixCursor(Span<const std::byte> prefix) override { return GetNewCursor(); }
bool TxnBegin() override { return true; }
bool TxnCommit() override { return true; }
bool TxnAbort() override { return true; }
Expand Down
54 changes: 53 additions & 1 deletion src/wallet/sqlite.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
#include <logging.h>
#include <sync.h>
#include <util/fs_helpers.h>
#include <util/check.h>
#include <util/strencodings.h>
#include <util/translation.h>
#include <wallet/db.h>
Expand Down Expand Up @@ -39,7 +40,11 @@ static bool BindBlobToStatement(sqlite3_stmt* stmt,
Span<const std::byte> blob,
const std::string& description)
{
int res = sqlite3_bind_blob(stmt, index, blob.data(), blob.size(), SQLITE_STATIC);
// Pass a pointer to the empty string "" below instead of passing the
// blob.data() pointer if the blob.data() pointer is null. Passing a null
// data pointer to bind_blob would cause sqlite to bind the SQL NULL value
// instead of the empty blob value X'', which would mess up SQL comparisons.
int res = sqlite3_bind_blob(stmt, index, blob.data() ? static_cast<const void*>(blob.data()) : "", blob.size(), SQLITE_STATIC);
if (res != SQLITE_OK) {
LogPrintf("Unable to bind %s to statement: %s\n", description, sqlite3_errstr(res));
sqlite3_clear_bindings(stmt);
Expand Down Expand Up @@ -409,6 +414,7 @@ bool SQLiteBatch::ReadKey(DataStream&& key, DataStream& value)
// Leftmost column in result is index 0
const std::byte* data{AsBytePtr(sqlite3_column_blob(m_read_stmt, 0))};
size_t data_size(sqlite3_column_bytes(m_read_stmt, 0));
value.clear();
value.write({data, data_size});

sqlite3_clear_bindings(m_read_stmt);
Expand Down Expand Up @@ -495,6 +501,9 @@ DatabaseCursor::Status SQLiteCursor::Next(DataStream& key, DataStream& value)
return Status::FAIL;
}

key.clear();
value.clear();

// Leftmost column in result is index 0
const std::byte* key_data{AsBytePtr(sqlite3_column_blob(m_cursor_stmt, 0))};
size_t key_data_size(sqlite3_column_bytes(m_cursor_stmt, 0));
Expand All @@ -507,6 +516,7 @@ DatabaseCursor::Status SQLiteCursor::Next(DataStream& key, DataStream& value)

SQLiteCursor::~SQLiteCursor()
{
sqlite3_clear_bindings(m_cursor_stmt);
sqlite3_reset(m_cursor_stmt);
int res = sqlite3_finalize(m_cursor_stmt);
if (res != SQLITE_OK) {
Expand All @@ -530,6 +540,48 @@ std::unique_ptr<DatabaseCursor> SQLiteBatch::GetNewCursor()
return cursor;
}

std::unique_ptr<DatabaseCursor> SQLiteBatch::GetNewPrefixCursor(Span<const std::byte> prefix)
{
if (!m_database.m_db) return nullptr;

// To get just the records we want, the SQL statement does a comparison of the binary data
// where the data must be greater than or equal to the prefix, and less than
// the prefix incremented by one (when interpreted as an integer)
std::vector<std::byte> start_range(prefix.begin(), prefix.end());
std::vector<std::byte> end_range(prefix.begin(), prefix.end());
auto it = end_range.rbegin();
for (; it != end_range.rend(); ++it) {
if (*it == std::byte(std::numeric_limits<unsigned char>::max())) {
*it = std::byte(0);
continue;
}
*it = std::byte(std::to_integer<unsigned char>(*it) + 1);
break;
}
if (it == end_range.rend()) {
// If the prefix is all 0xff bytes, clear end_range as we won't need it
end_range.clear();
}

auto cursor = std::make_unique<SQLiteCursor>(start_range, end_range);
if (!cursor) return nullptr;

const char* stmt_text = end_range.empty() ? "SELECT key, value FROM main WHERE key >= ?" :
"SELECT key, value FROM main WHERE key >= ? AND key < ?";
int res = sqlite3_prepare_v2(m_database.m_db, stmt_text, -1, &cursor->m_cursor_stmt, nullptr);
if (res != SQLITE_OK) {
throw std::runtime_error(strprintf(
"SQLiteDatabase: Failed to setup cursor SQL statement: %s\n", sqlite3_errstr(res)));
}

if (!BindBlobToStatement(cursor->m_cursor_stmt, 1, cursor->m_prefix_range_start, "prefix_start")) return nullptr;
if (!end_range.empty()) {
if (!BindBlobToStatement(cursor->m_cursor_stmt, 2, cursor->m_prefix_range_end, "prefix_end")) return nullptr;
}

return cursor;
}

bool SQLiteBatch::TxnBegin()
{
if (!m_database.m_db || sqlite3_get_autocommit(m_database.m_db) == 0) return false;
Expand Down
10 changes: 10 additions & 0 deletions src/wallet/sqlite.h
Original file line number Diff line number Diff line change
Expand Up @@ -15,12 +15,21 @@ struct bilingual_str;
namespace wallet {
class SQLiteDatabase;

/** RAII class that provides a database cursor */
class SQLiteCursor : public DatabaseCursor
{
public:
sqlite3_stmt* m_cursor_stmt{nullptr};
// Copies of the prefix things for the prefix cursor.
// Prevents SQLite from accessing temp variables for the prefix things.
std::vector<std::byte> m_prefix_range_start;
std::vector<std::byte> m_prefix_range_end;

explicit SQLiteCursor() {}
explicit SQLiteCursor(std::vector<std::byte> start_range, std::vector<std::byte> end_range)
: m_prefix_range_start(std::move(start_range)),
m_prefix_range_end(std::move(end_range))
{}
~SQLiteCursor() override;

Status Next(DataStream& key, DataStream& value) override;
Expand Down Expand Up @@ -57,6 +66,7 @@ class SQLiteBatch : public DatabaseBatch
void Close() override;

std::unique_ptr<DatabaseCursor> GetNewCursor() override;
std::unique_ptr<DatabaseCursor> GetNewPrefixCursor(Span<const std::byte> prefix) override;
bool TxnBegin() override;
bool TxnCommit() override;
bool TxnAbort() override;
Expand Down
128 changes: 128 additions & 0 deletions src/wallet/test/db_tests.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -6,13 +6,56 @@

#include <test/util/setup_common.h>
#include <util/fs.h>
#include <util/translation.h>
#ifdef USE_BDB
#include <wallet/bdb.h>
#endif
#ifdef USE_SQLITE
#include <wallet/sqlite.h>
#endif
#include <wallet/test/util.h>
#include <wallet/walletutil.h> // for WALLET_FLAG_DESCRIPTORS

#include <fstream>
#include <memory>
#include <string>

inline std::ostream& operator<<(std::ostream& os, const std::pair<const SerializeData, SerializeData>& kv)
{
Span key{kv.first}, value{kv.second};
os << "(\"" << std::string_view{reinterpret_cast<const char*>(key.data()), key.size()} << "\", \""
<< std::string_view{reinterpret_cast<const char*>(key.data()), key.size()} << "\")";
return os;
}

namespace wallet {

static Span<const std::byte> StringBytes(std::string_view str)
{
return AsBytes<const char>({str.data(), str.size()});
}

static SerializeData StringData(std::string_view str)
{
auto bytes = StringBytes(str);
return SerializeData{bytes.begin(), bytes.end()};
}

static void CheckPrefix(DatabaseBatch& batch, Span<const std::byte> prefix, MockableData expected)
{
std::unique_ptr<DatabaseCursor> cursor = batch.GetNewPrefixCursor(prefix);
MockableData actual;
while (true) {
DataStream key, value;
DatabaseCursor::Status status = cursor->Next(key, value);
if (status == DatabaseCursor::Status::DONE) break;
BOOST_CHECK(status == DatabaseCursor::Status::MORE);
BOOST_CHECK(
actual.emplace(SerializeData(key.begin(), key.end()), SerializeData(value.begin(), value.end())).second);
}
BOOST_CHECK_EQUAL_COLLECTIONS(actual.begin(), actual.end(), expected.begin(), expected.end());
}

BOOST_FIXTURE_TEST_SUITE(db_tests, BasicTestingSetup)

static std::shared_ptr<BerkeleyEnvironment> GetWalletEnv(const fs::path& path, fs::path& database_filename)
Expand Down Expand Up @@ -78,5 +121,90 @@ BOOST_AUTO_TEST_CASE(getwalletenv_g_dbenvs_free_instance)
BOOST_CHECK(env_2_a == env_2_b);
}

static std::vector<std::unique_ptr<WalletDatabase>> TestDatabases(const fs::path& path_root)
{
std::vector<std::unique_ptr<WalletDatabase>> dbs;
DatabaseOptions options;
DatabaseStatus status;
bilingual_str error;
#ifdef USE_BDB
dbs.emplace_back(MakeBerkeleyDatabase(path_root / "bdb", options, status, error));
#endif
#ifdef USE_SQLITE
dbs.emplace_back(MakeSQLiteDatabase(path_root / "sqlite", options, status, error));
#endif
dbs.emplace_back(CreateMockableWalletDatabase());
return dbs;
}

BOOST_AUTO_TEST_CASE(db_cursor_prefix_range_test)
{
// Test each supported db
for (const auto& database : TestDatabases(m_path_root)) {
BOOST_ASSERT(database);

std::vector<std::string> prefixes = {"", "FIRST", "SECOND", "P\xfe\xff", "P\xff\x01", "\xff\xff"};

// Write elements to it
std::unique_ptr<DatabaseBatch> handler = database->MakeBatch();
for (unsigned int i = 0; i < 10; i++) {
for (const auto& prefix : prefixes) {
BOOST_CHECK(handler->Write(std::make_pair(prefix, i), i));
}
}

// Now read all the items by prefix and verify that each element gets parsed correctly
for (const auto& prefix : prefixes) {
DataStream s_prefix;
s_prefix << prefix;
std::unique_ptr<DatabaseCursor> cursor = handler->GetNewPrefixCursor(s_prefix);
DataStream key;
DataStream value;
for (int i = 0; i < 10; i++) {
DatabaseCursor::Status status = cursor->Next(key, value);
BOOST_ASSERT(status == DatabaseCursor::Status::MORE);

std::string key_back;
unsigned int i_back;
key >> key_back >> i_back;
BOOST_CHECK_EQUAL(key_back, prefix);

unsigned int value_back;
value >> value_back;
BOOST_CHECK_EQUAL(value_back, i_back);
}

// Let's now read it once more, it should return DONE
BOOST_CHECK(cursor->Next(key, value) == DatabaseCursor::Status::DONE);
}
}
}

// Lower level DatabaseBase::GetNewPrefixCursor test, to cover cases that aren't
// covered in the higher level test above. The higher level test uses
// serialized strings which are prefixed with string length, so it doesn't test
// truly empty prefixes or prefixes that begin with \xff
BOOST_AUTO_TEST_CASE(db_cursor_prefix_byte_test)
{
const MockableData::value_type
e{StringData(""), StringData("e")},
p{StringData("prefix"), StringData("p")},
ps{StringData("prefixsuffix"), StringData("ps")},
f{StringData("\xff"), StringData("f")},
fs{StringData("\xffsuffix"), StringData("fs")},
ff{StringData("\xff\xff"), StringData("ff")},
ffs{StringData("\xff\xffsuffix"), StringData("ffs")};
for (const auto& database : TestDatabases(m_path_root)) {
std::unique_ptr<DatabaseBatch> batch = database->MakeBatch();
for (const auto& [k, v] : {e, p, ps, f, fs, ff, ffs}) {
batch->Write(MakeUCharSpan(k), MakeUCharSpan(v));
}
CheckPrefix(*batch, StringBytes(""), {e, p, ps, f, fs, ff, ffs});
CheckPrefix(*batch, StringBytes("prefix"), {p, ps});
CheckPrefix(*batch, StringBytes("\xff"), {f, fs, ff, ffs});
CheckPrefix(*batch, StringBytes("\xff\xff"), {ff, ffs});
}
}

BOOST_AUTO_TEST_SUITE_END()
} // namespace wallet
Loading

0 comments on commit 7f20197

Please sign in to comment.