kv: add bounded snapshot scan API

Signed-off-by: lance6716 <lance6716@gmail.com>

Author: lance6716

include/pingcap/kv/Snapshot.h

@@ -3,12 +3,31 @@
 #include <pingcap/kv/Cluster.h>
 #include <pingcap/kv/RegionClient.h>
 
+#include <vector>
+
 namespace pingcap
 {
 namespace kv
 {
 struct Scanner;
 
+struct ScanOptions
+{
+    std::string start_key;
+    std::string end_key;
+    uint32_t limit = 0;
+    bool reverse = false;
+    bool key_only = false;
+    uint64_t version = 0;
+};
+
+struct ScanResult
+{
+    std::vector<kvrpcpb::KvPair> pairs;
+    bool has_more = false;
+    std::string next_start_key;
+};
+
 struct Snapshot
 {
     Cluster * cluster;
@@ -31,6 +50,9 @@ struct Snapshot
 
     kvrpcpb::MvccInfo mvccGet(Backoffer & bo, const std::string & key);
 
+    ScanResult ScanOnce(const ScanOptions & options);
+    ScanResult ScanOnce(Backoffer & bo, const ScanOptions & options);
+
     Scanner Scan(const std::string & begin, const std::string & end);
 };
 

src/kv/Snapshot.cc

@@ -4,12 +4,67 @@
 #include <pingcap/kv/Scanner.h>
 #include <pingcap/kv/Snapshot.h>
 
+#include <algorithm>
+#include <limits>
+
 namespace pingcap
 {
 namespace kv
 {
 constexpr int scan_batch_size = 256;
 
+namespace
+{
+
+uint32_t requestLimit(uint32_t remaining)
+{
+    if (remaining == std::numeric_limits<uint32_t>::max())
+        return remaining;
+    return remaining + 1;
+}
+
+std::string keyBefore(const std::string & upper_bound)
+{
+    // Used to locate the region containing keys immediately below a reverse scan upper bound.
+    if (upper_bound.empty())
+        return std::string(1, static_cast<char>(0xff));
+
+    std::string key = upper_bound;
+    auto last = static_cast<unsigned char>(key.back());
+    if (last == 0)
+    {
+        key.pop_back();
+        return key;
+    }
+
+    key.back() = static_cast<char>(last - 1);
+    key.push_back(static_cast<char>(0xff));
+    return key;
+}
+
+bool forwardRangeFinished(const std::string & cursor, const std::string & end_key)
+{
+    return cursor.empty() || (!end_key.empty() && cursor >= end_key);
+}
+
+bool reverseRangeFinished(const std::string & cursor, const std::string & end_key)
+{
+    return cursor.empty() || (!end_key.empty() && cursor <= end_key);
+}
+
+void setScanContext(kvrpcpb::ScanRequest & request, const MinCommitTSPushed & min_commit_ts_pushed)
+{
+    auto * context = request.mutable_context();
+    context->set_priority(::kvrpcpb::Normal);
+    context->set_not_fill_cache(false);
+    for (auto ts : min_commit_ts_pushed.getTimestamps())
+    {
+        context->add_resolved_locks(ts);
+    }
+}
+
+} // namespace
+
 //bool extractLockFromKeyErr()
 
 kvrpcpb::MvccInfo Snapshot::mvccGet(const std::string & key)
@@ -113,6 +168,164 @@ std::string Snapshot::Get(Backoffer & bo, const std::string & key)
     }
 }
 
+ScanResult Snapshot::ScanOnce(const ScanOptions & options)
+{
+    Backoffer bo(scanMaxBackoff);
+    return ScanOnce(bo, options);
+}
+
+ScanResult Snapshot::ScanOnce(Backoffer & bo, const ScanOptions & options)
+{
+    ScanResult result;
+    if (options.limit == 0)
+        return result;
+
+    const uint64_t read_version = options.version == 0 ? version : options.version;
+    std::string cursor = options.start_key;
+
+    if (!options.reverse && !options.end_key.empty() && cursor >= options.end_key)
+        return result;
+    if (options.reverse && !cursor.empty() && !options.end_key.empty() && options.end_key >= cursor)
+        return result;
+
+    while (result.pairs.size() < options.limit)
+    {
+        auto loc = options.reverse ? cluster->region_cache->locateKey(bo, keyBefore(cursor)) : cluster->region_cache->locateKey(bo, cursor);
+
+        kvrpcpb::ScanRequest request;
+        if (options.reverse)
+        {
+            std::string request_end_key = loc.start_key;
+            if (!options.end_key.empty() && options.end_key > request_end_key)
+                request_end_key = options.end_key;
+
+            std::string request_start_key = cursor;
+            if (!loc.end_key.empty() && (request_start_key.empty() || request_start_key > loc.end_key))
+                request_start_key = loc.end_key;
+
+            if (!request_start_key.empty() && request_start_key <= request_end_key)
+            {
+                cursor = loc.start_key;
+                if (reverseRangeFinished(cursor, options.end_key))
+                    return result;
+                continue;
+            }
+
+            request.set_start_key(request_start_key);
+            request.set_end_key(request_end_key);
+            request.set_reverse(true);
+        }
+        else
+        {
+            std::string request_end_key = options.end_key;
+            if (!request_end_key.empty() && !loc.end_key.empty() && loc.end_key < request_end_key)
+                request_end_key = loc.end_key;
+
+            request.set_start_key(cursor);
+            request.set_end_key(request_end_key);
+        }
+
+        const auto remaining = static_cast<uint32_t>(options.limit - result.pairs.size());
+        request.set_limit(requestLimit(remaining));
+        request.set_version(read_version);
+        request.set_key_only(options.key_only);
+        setScanContext(request, min_commit_ts_pushed);
+
+        kvrpcpb::ScanResponse response;
+        RegionClient region_client(cluster, loc.region);
+        try
+        {
+            region_client.sendReqToRegion<RPC_NAME(KvScan)>(bo, request, &response);
+        }
+        catch (Exception & e)
+        {
+            bo.backoff(boRegionMiss, e);
+            continue;
+        }
+
+        if (response.has_error())
+        {
+            auto lock = extractLockFromKeyErr(response.error());
+            std::vector<LockPtr> locks{lock};
+            std::vector<uint64_t> pushed{};
+            auto ms_before_expired = cluster->lock_resolver->resolveLocks(bo, read_version, locks, pushed);
+            if (!pushed.empty())
+            {
+                min_commit_ts_pushed.addTimestamps(pushed);
+            }
+            if (ms_before_expired > 0)
+            {
+                bo.backoffWithMaxSleep(
+                    BackoffType::boTxnLockFast,
+                    ms_before_expired,
+                    Exception("key is locked during scanning", ErrorCodes::LockError));
+            }
+            continue;
+        }
+
+        std::vector<LockPtr> locks;
+        for (int i = 0; i < response.pairs_size(); i++)
+        {
+            const auto & pair = response.pairs(i);
+            if (pair.has_error())
+                locks.push_back(extractLockFromKeyErr(pair.error()));
+        }
+        if (!locks.empty())
+        {
+            std::vector<uint64_t> pushed{};
+            auto ms_before_expired = cluster->lock_resolver->resolveLocks(bo, read_version, locks, pushed);
+            if (!pushed.empty())
+            {
+                min_commit_ts_pushed.addTimestamps(pushed);
+            }
+            if (ms_before_expired > 0)
+            {
+                bo.backoffWithMaxSleep(
+                    BackoffType::boTxnLockFast,
+                    ms_before_expired,
+                    Exception("key is locked during scanning", ErrorCodes::LockError));
+            }
+            continue;
+        }
+
+        const auto pairs_size = static_cast<size_t>(response.pairs_size());
+        const auto append_size = std::min<size_t>(pairs_size, remaining);
+        for (size_t i = 0; i < append_size; i++)
+        {
+            result.pairs.push_back(response.pairs(static_cast<int>(i)));
+        }
+
+        if (pairs_size > remaining)
+        {
+            result.has_more = true;
+            result.next_start_key = options.reverse ? result.pairs.back().key() : alphabeticalNext(result.pairs.back().key());
+            return result;
+        }
+
+        if (options.reverse)
+        {
+            cursor = loc.start_key;
+            if (reverseRangeFinished(cursor, options.end_key))
+                return result;
+        }
+        else
+        {
+            cursor = loc.end_key;
+            if (forwardRangeFinished(cursor, options.end_key))
+                return result;
+        }
+
+        if (result.pairs.size() == options.limit)
+        {
+            result.has_more = true;
+            result.next_start_key = cursor;
+            return result;
+        }
+    }
+
+    return result;
+}
+
 Scanner Snapshot::Scan(const std::string & begin, const std::string & end)
 {
     return Scanner(*this, begin, end, scan_batch_size);

src/test/region_split_test.cc

@@ -13,6 +13,16 @@ namespace pingcap::tests
 using namespace pingcap;
 using namespace pingcap::kv;
 
+std::vector<std::string> scanKeys(const ScanResult & result)
+{
+    std::vector<std::string> keys;
+    for (const auto & pair : result.pairs)
+    {
+        keys.push_back(pair.key());
+    }
+    return keys;
+}
+
 class TestWithMockKVRegionSplit : public testing::Test
 {
 protected:
@@ -116,4 +126,57 @@ TEST_F(TestWithMockKVRegionSplit, testSplitRegionScan)
     ASSERT_EQ(answer, 6);
 }
 
+TEST_F(TestWithMockKVRegionSplit, testScanOnce)
+{
+    Txn txn(test_cluster.get());
+
+    txn.set("abc", "1");
+    txn.set("abd", "2");
+    txn.set("abe", "3");
+    txn.set("abf", "4");
+    txn.set("abg", "5");
+    txn.set("abh", "6");
+    txn.set("zzz", "7");
+    txn.commit();
+
+    control_cluster->splitRegion("abe");
+
+    Snapshot snap(test_cluster.get(), test_cluster->pd_client->getTS());
+
+    ScanOptions forward_options;
+    forward_options.start_key = "ab";
+    forward_options.end_key = "ac";
+    forward_options.limit = 3;
+
+    auto forward_result = snap.ScanOnce(forward_options);
+    ASSERT_EQ(scanKeys(forward_result), std::vector<std::string>({"abc", "abd", "abe"}));
+    ASSERT_TRUE(forward_result.has_more);
+
+    forward_options.start_key = forward_result.next_start_key;
+    forward_result = snap.ScanOnce(forward_options);
+    ASSERT_EQ(scanKeys(forward_result), std::vector<std::string>({"abf", "abg", "abh"}));
+    ASSERT_FALSE(forward_result.has_more);
+
+    ScanOptions reverse_options;
+    reverse_options.start_key = "ac";
+    reverse_options.end_key = "ab";
+    reverse_options.limit = 10;
+    reverse_options.reverse = true;
+
+    auto reverse_result = snap.ScanOnce(reverse_options);
+    ASSERT_EQ(scanKeys(reverse_result), std::vector<std::string>({"abh", "abg", "abf", "abe", "abd", "abc"}));
+    ASSERT_FALSE(reverse_result.has_more);
+
+    reverse_options.limit = 2;
+    reverse_result = snap.ScanOnce(reverse_options);
+    ASSERT_EQ(scanKeys(reverse_result), std::vector<std::string>({"abh", "abg"}));
+    ASSERT_TRUE(reverse_result.has_more);
+
+    reverse_options.start_key = reverse_result.next_start_key;
+    reverse_options.limit = 10;
+    reverse_result = snap.ScanOnce(reverse_options);
+    ASSERT_EQ(scanKeys(reverse_result), std::vector<std::string>({"abf", "abe", "abd", "abc"}));
+    ASSERT_FALSE(reverse_result.has_more);
+}
+
 } // namespace pingcap::tests