add limit for content_scan api

include/sharksfin/api.h

@@ -788,6 +788,7 @@ inline std::ostream& operator<<(std::ostream& out, EndPointKind value) {
  * @param end_key the content key of ending position
  * @param end_kind end-point kind of the ending position
  * @param result [OUT] an iterator handle over the key range
+ * @param limit the max number of entries to be fetched. 0 indicates no limit.
  * @param reverse whether or not the iterator scans in reverse order (from end to begin)
  * Current limitation on reverse = true is that end_kind must be EndPointKind::UNBOUND and at most one entry is
  * fetched as the scan result
@@ -801,7 +802,9 @@ extern "C" StatusCode content_scan(
         StorageHandle storage,
         Slice begin_key, EndPointKind begin_kind,
         Slice end_key, EndPointKind end_kind,
-        IteratorHandle* result, bool reverse = false);
+        IteratorHandle* result,
+        std::size_t limit = 0,
+        bool reverse = false);
 
 /**
  * @brief advances the given iterator.

memory/src/Iterator.h

@@ -46,17 +46,19 @@ class Iterator {
      * @param begin_kind end-point kind of the beginning position
      * @param end_key the content key of ending position
      * @param end_kind end-point kind of the ending position
+     * @param limit the max number of entries to be fetched. 0 indicates no limit.
      * @param reverse whether or not the iterator scans in reverse order (from end to begin)
      */
     Iterator(
             Storage* owner,
             Slice begin_key, EndPointKind begin_kind,
-            Slice end_key, EndPointKind end_kind, bool reverse = false)
+            Slice end_key, EndPointKind end_kind, std::size_t limit = 0, bool reverse = false)
         : owner_(owner)
         , next_key_(begin_kind == EndPointKind::UNBOUND ? std::string_view {} : begin_key.to_string_view())
         , end_key_(end_kind == EndPointKind::UNBOUND ? Slice {} : end_key)
         , end_type_(interpret_end_kind(end_kind))
         , state_(interpret_begin_kind(begin_kind))
+        , limit_(limit)
         , reverse_(reverse)
     {}
 
@@ -107,6 +109,7 @@ class Iterator {
     Buffer end_key_;
     End end_type_;
     State state_;
+    std::size_t limit_;
     bool reverse_;
 
     Slice payload_ {};

memory/src/api.cpp

@@ -399,16 +399,18 @@ StatusCode content_scan(
     Slice begin_key, EndPointKind begin_kind,
     Slice end_key, EndPointKind end_kind,
     IteratorHandle* result,
+    std::size_t limit,
     bool reverse) {
     log_entry << fn_name << " transaction:" << transaction << " storage:" << storage <<
         binstring(begin_key) << " begin_kind:" << begin_kind <<
-        binstring(end_key) << " end_kind:" << end_kind;
+        binstring(end_key) << " end_kind:" << end_kind <<
+        " limit:" << limit << " reverse:" << reverse;
     auto rc = impl::content_scan(
         transaction,
         storage,
         begin_key, begin_kind,
         end_key, end_kind,
-        result, reverse);
+        result, limit, reverse);
     log_rc(rc, fn_name);
     log_exit << fn_name << " rc:" << rc << " result:" << *result;
     return rc;

memory/src/api_helper.cpp

@@ -505,6 +505,7 @@ StatusCode content_scan(
         Slice begin_key, EndPointKind begin_kind,
         Slice end_key, EndPointKind end_kind,
         IteratorHandle* result,
+        std::size_t limit,
         bool reverse) {
     auto tx = unwrap(transaction);
     auto st = unwrap(storage);
@@ -514,7 +515,7 @@ StatusCode content_scan(
     auto iterator = std::make_unique<memory::Iterator>(
             st,
             begin_key, begin_kind,
-            end_key, end_kind, reverse);
+            end_key, end_kind, limit, reverse);
     *result = wrap(iterator.release());
     return StatusCode::OK;
 }

memory/src/api_helper.h

@@ -179,7 +179,9 @@ StatusCode content_scan(
         StorageHandle storage,
         Slice begin_key, EndPointKind begin_kind,
         Slice end_key, EndPointKind end_kind,
-        IteratorHandle* result, bool reverse);
+        IteratorHandle* result,
+        std::size_t limit,
+        bool reverse);
 
 StatusCode iterator_next(IteratorHandle handle);
 

shirakami/src/Iterator.cpp

@@ -36,6 +36,7 @@ Iterator::Iterator(
     EndPointKind begin_kind,
     Slice end_key,
     EndPointKind end_kind,
+    std::size_t limit,
     bool reverse) :
     owner_(owner),
     state_(State::INIT),
@@ -44,6 +45,7 @@ Iterator::Iterator(
     begin_kind_(begin_kind),
     end_key_(end_kind == EndPointKind::UNBOUND ? std::string_view{} : end_key.to_string_view()),
     end_kind_(end_kind),
+    limit_(limit),
     reverse_(reverse) {}
 
 Iterator::~Iterator() {
@@ -201,13 +203,10 @@ StatusCode Iterator::open_cursor() {
             break;
     }
 
-    // reverse scan can fetch only one entry for now
-    std::size_t max_size = reverse_ ? 1 : 0;
-
     auto res = api::open_scan(tx_->native_handle(),
         owner_->handle(),
         begin_key_, begin_endpoint,
-        end_key_, end_endpoint, handle_, max_size, reverse_);
+        end_key_, end_endpoint, handle_, limit_, reverse_);
     tx_->last_call_status(res);
     correct_transaction_state(*tx_, res);
     if(res == Status::WARN_NOT_FOUND) {

shirakami/src/Iterator.h

@@ -46,13 +46,15 @@ class Iterator {
      * @param end_key the content key of ending position
      * @param end_kind whether or not ending position is exclusive
      * If end_key is not empty and end kind UNBOUND, the end_kind is reduced to PREFIXED_INCLUSIVE
+     * @param limit the max number of entries to be fetched. 0 indicates no limit.
      * @param reverse whether or not the iterator scans in reverse order (from end to begin)
      */
     Iterator( // NOLINT(performance-unnecessary-value-param)
         Storage* owner,
         Transaction* tx,
         Slice begin_key, EndPointKind begin_kind,
         Slice end_key, EndPointKind end_kind,
+        std::size_t limit = 0,
         bool reverse = false
     );
 
@@ -99,6 +101,7 @@ class Iterator {
     EndPointKind begin_kind_{};
     std::string end_key_{};
     EndPointKind end_kind_{};
+    std::size_t limit_{};
     bool reverse_{};
     bool key_value_readable_{false};
     bool need_scan_close_{false};

shirakami/src/Storage.cpp

@@ -66,12 +66,13 @@ StatusCode Storage::scan(Transaction* tx,
         Slice begin_key, EndPointKind begin_kind,
         Slice end_key, EndPointKind end_kind,
         std::unique_ptr<Iterator>& out,
+        std::size_t limit,
         bool reverse
 ) {
     if(! tx->active()) return StatusCode::ERR_INACTIVE_TRANSACTION;
     out = std::make_unique<Iterator>(this, tx,
             begin_key, begin_kind,
-            end_key, end_kind, reverse);
+            end_key, end_kind, limit, reverse);
     return StatusCode::OK;
 }
 

shirakami/src/Storage.h

@@ -105,13 +105,15 @@ class Storage {
      * @param end_key the content key of ending position
      * @param end_kind end-point kind of the ending position
      * @param out [OUT] the created iterator
+     * @param limit the max number of entries to be fetched. 0 indicates no limit.
      * @param reverse whether or not the scan in reverse order (from end to begin)
      * @return the operation status
      */
     StatusCode scan(Transaction* tx,
             Slice begin_key, EndPointKind begin_kind,
             Slice end_key, EndPointKind end_kind,
             std::unique_ptr<Iterator>& out,
+            std::size_t limit = 0,
             bool reverse = false
     );
 

shirakami/src/api.cpp

@@ -465,6 +465,7 @@ StatusCode content_scan(
         Slice begin_key, EndPointKind begin_kind,
         Slice end_key, EndPointKind end_kind,
         IteratorHandle* result,
+        std::size_t limit,
         bool reverse) {
     auto tx = unwrap(transaction);
     if (! tx->active()) return StatusCode::ERR_INACTIVE_TRANSACTION;
@@ -474,7 +475,7 @@ StatusCode content_scan(
         return StatusCode::ERR_INVALID_STATE;
     }
     std::unique_ptr<shirakami::Iterator> iter{};
-    auto rc = stg->scan(tx, begin_key, begin_kind, end_key, end_kind, iter, reverse);
+    auto rc = stg->scan(tx, begin_key, begin_kind, end_key, end_kind, iter, limit, reverse);
     if(rc != StatusCode::OK) {
         return rc;
     }

shirakami/test/ApiTest.cpp

@@ -937,6 +937,7 @@ TEST_F(ShirakamiApiTest, scan_range) {
     EXPECT_EQ(transaction_exec(db, {}, &S::test, &s), StatusCode::OK);
     EXPECT_EQ(database_close(db), StatusCode::OK);
 }
+
 TEST_F(ShirakamiApiTest, scan) {
     DatabaseOptions options;
     options.attribute(KEY_LOCATION, path());
@@ -1026,6 +1027,83 @@ TEST_F(ShirakamiApiTest, scan) {
     EXPECT_EQ(database_close(db), StatusCode::OK);
 }
 
+TEST_F(ShirakamiApiTest, scan_with_limit) {
+    DatabaseOptions options;
+    options.attribute(KEY_LOCATION, path());
+    DatabaseHandle db;
+    ASSERT_EQ(database_open(options, &db), StatusCode::OK);
+    HandleHolder dbh { db };
+
+    struct S {
+        static TransactionOperation prepare(TransactionHandle tx, void* args) {
+            auto st = extract<S>(args);
+            if (content_put(tx, st, "a", "NG") != StatusCode::OK) {
+                return TransactionOperation::ERROR;
+            }
+            if (content_put(tx, st, "a1", "NG") != StatusCode::OK) {
+                return TransactionOperation::ERROR;
+            }
+            if (content_put(tx, st, "b", "B") != StatusCode::OK) {
+                return TransactionOperation::ERROR;
+            }
+            if (content_put(tx, st, "c", "C") != StatusCode::OK) {
+                return TransactionOperation::ERROR;
+            }
+            if (content_put(tx, st, "d", "NG") != StatusCode::OK) {
+                return TransactionOperation::ERROR;
+            }
+            return TransactionOperation::COMMIT;
+        }
+        static TransactionOperation test(TransactionHandle tx, void* args) {
+            auto st = extract<S>(args);
+            IteratorHandle iter;
+            if (content_scan(
+                    tx, st,
+                    "b", EndPointKind::PREFIXED_INCLUSIVE,
+                    "", EndPointKind::UNBOUND,
+                    &iter,
+                    2) != StatusCode::OK) {
+                return TransactionOperation::ERROR;
+            }
+            HandleHolder closer { iter };
+
+            Slice s;
+            if (iterator_next(iter) != StatusCode::OK) {
+                return TransactionOperation::ERROR;
+            }
+            if (iterator_get_key(iter, &s) != StatusCode::OK || s != "b") {
+                return TransactionOperation::ERROR;
+            }
+            if (iterator_get_value(iter, &s) != StatusCode::OK || s != "B") {
+                return TransactionOperation::ERROR;
+            }
+
+            if (iterator_next(iter) != StatusCode::OK) {
+                return TransactionOperation::ERROR;
+            }
+            if (iterator_get_key(iter, &s) != StatusCode::OK || s != "c") {
+                return TransactionOperation::ERROR;
+            }
+            if (iterator_get_value(iter, &s) != StatusCode::OK || s != "C") {
+                return TransactionOperation::ERROR;
+            }
+
+            if (iterator_next(iter) != StatusCode::NOT_FOUND) {
+                return TransactionOperation::ERROR;
+            }
+            return TransactionOperation::COMMIT;
+        }
+        StorageHandle st;
+    };
+    S s;
+    ASSERT_EQ(storage_create(db, "s", &s.st), StatusCode::OK);
+    HandleHolder sth { s.st };
+
+    EXPECT_EQ(transaction_exec(db, {}, &S::prepare, &s), StatusCode::OK);
+    EXPECT_EQ(transaction_exec(db, {}, &S::test, &s), StatusCode::OK);
+    EXPECT_EQ(database_close(db), StatusCode::OK);
+}
+
 TEST_F(ShirakamiApiTest, reverse_scan) {
     // currently reverse scan can fetch only one entry with unbonded end point
     DatabaseOptions options;
@@ -1056,6 +1134,7 @@ TEST_F(ShirakamiApiTest, reverse_scan) {
                     "a", EndPointKind::PREFIXED_INCLUSIVE,
                     "", EndPointKind::UNBOUND,
                     &iter,
+                    1,  // reverse scan must have limit 1
                     true // reverse scan
                     ) != StatusCode::OK) {
                 return TransactionOperation::ERROR;