From 879f3663663e538390d0d835919aa73115892e75 Mon Sep 17 00:00:00 2001
From: zhangjinpeng1987 <zhangjinpeng@pingcap.com>
Date: Thu, 3 Nov 2016 11:00:46 -0700
Subject: [PATCH] Add C api for RateLimiter

Summary:
Add C api for RateLimiter.
Closes https://github.com/facebook/rocksdb/pull/1455

Differential Revision: D4116362

Pulled By: yiwu-arbug

fbshipit-source-id: cb05a8d
---
 db/c.cc             | 26 ++++++++++++++++++++++++++
 db/c_test.c         |  4 ++++
 include/rocksdb/c.h |  8 ++++++++
 3 files changed, 38 insertions(+)

diff --git a/db/c.cc b/db/c.cc
index 60d4d570d..cbc77e546 100644
--- a/db/c.cc
+++ b/db/c.cc
@@ -30,6 +30,7 @@
 #include "rocksdb/statistics.h"
 #include "rocksdb/slice_transform.h"
 #include "rocksdb/table.h"
+#include "rocksdb/rate_limiter.h"
 #include "rocksdb/utilities/backupable_db.h"
 #include "utilities/merge_operators.h"
 
@@ -81,6 +82,8 @@ using rocksdb::BackupableDBOptions;
 using rocksdb::BackupInfo;
 using rocksdb::RestoreOptions;
 using rocksdb::CompactRangeOptions;
+using rocksdb::RateLimiter;
+using rocksdb::NewGenericRateLimiter;
 
 using std::shared_ptr;
 
@@ -114,6 +117,7 @@ struct rocksdb_column_family_handle_t  { ColumnFamilyHandle* rep; };
 struct rocksdb_envoptions_t      { EnvOptions        rep; };
 struct rocksdb_ingestexternalfileoptions_t  { IngestExternalFileOptions rep; };
 struct rocksdb_sstfilewriter_t   { SstFileWriter*    rep; };
+struct rocksdb_ratelimiter_t     { RateLimiter*      rep; };
 
 struct rocksdb_compactionfiltercontext_t {
   CompactionFilter::Context rep;
@@ -1880,6 +1884,28 @@ char *rocksdb_options_statistics_get_string(rocksdb_options_t *opt) {
   return nullptr;
 }
 
+void rocksdb_options_set_ratelimiter(rocksdb_options_t *opt, rocksdb_ratelimiter_t *limiter) {
+  opt->rep.rate_limiter.reset(limiter->rep);
+  limiter->rep = nullptr;
+}
+
+rocksdb_ratelimiter_t* rocksdb_ratelimiter_create(
+    int64_t rate_bytes_per_sec,
+    int64_t refill_period_us,
+    int32_t fairness) {
+  rocksdb_ratelimiter_t* rate_limiter = new rocksdb_ratelimiter_t;
+  rate_limiter->rep = NewGenericRateLimiter(rate_bytes_per_sec,
+		  refill_period_us, fairness);
+  return rate_limiter;
+}
+
+void rocksdb_ratelimiter_destroy(rocksdb_ratelimiter_t *limiter) {
+  if (limiter->rep) {
+	delete limiter->rep;
+  }
+  delete limiter;
+}
+
 /*
 TODO:
 DB::OpenForReadOnly
diff --git a/db/c_test.c b/db/c_test.c
index 82339a1fa..5bc8bc8a0 100644
--- a/db/c_test.c
+++ b/db/c_test.c
@@ -290,6 +290,7 @@ int main(int argc, char** argv) {
   rocksdb_block_based_table_options_t* table_options;
   rocksdb_readoptions_t* roptions;
   rocksdb_writeoptions_t* woptions;
+  rocksdb_ratelimiter_t* rate_limiter;
   char* err = NULL;
   int run = -1;
 
@@ -330,6 +331,9 @@ int main(int argc, char** argv) {
   int compression_levels[] = {rocksdb_no_compression, rocksdb_no_compression,
                               rocksdb_no_compression, rocksdb_no_compression};
   rocksdb_options_set_compression_per_level(options, compression_levels, 4);
+  rate_limiter = rocksdb_ratelimiter_create(1000 * 1024 * 1024, 100 * 1000, 10);
+  rocksdb_options_set_ratelimiter(options, rate_limiter);
+  rocksdb_ratelimiter_destroy(rate_limiter);
 
   roptions = rocksdb_readoptions_create();
   rocksdb_readoptions_set_verify_checksums(roptions, 1);
diff --git a/include/rocksdb/c.h b/include/rocksdb/c.h
index 959c8e161..b17917098 100644
--- a/include/rocksdb/c.h
+++ b/include/rocksdb/c.h
@@ -108,6 +108,7 @@ typedef struct rocksdb_column_family_handle_t rocksdb_column_family_handle_t;
 typedef struct rocksdb_envoptions_t      rocksdb_envoptions_t;
 typedef struct rocksdb_ingestexternalfileoptions_t rocksdb_ingestexternalfileoptions_t;
 typedef struct rocksdb_sstfilewriter_t   rocksdb_sstfilewriter_t;
+typedef struct rocksdb_ratelimiter_t     rocksdb_ratelimiter_t;
 
 /* DB operations */
 
@@ -721,6 +722,13 @@ rocksdb_options_set_universal_compaction_options(
     rocksdb_options_t*, rocksdb_universal_compaction_options_t*);
 extern ROCKSDB_LIBRARY_API void rocksdb_options_set_fifo_compaction_options(
     rocksdb_options_t* opt, rocksdb_fifo_compaction_options_t* fifo);
+extern ROCKSDB_LIBRARY_API void rocksdb_options_set_ratelimiter(
+    rocksdb_options_t* opt, rocksdb_ratelimiter_t* limiter);
+
+/* RateLimiter */
+extern ROCKSDB_LIBRARY_API rocksdb_ratelimiter_t* rocksdb_ratelimiter_create(
+    int64_t rate_bytes_per_sec, int64_t refill_period_us, int32_t fairness);
+extern ROCKSDB_LIBRARY_API void rocksdb_ratelimiter_destroy(rocksdb_ratelimiter_t*);
 
 /* Compaction Filter */
 
-- 
GitLab