Add on_stall_conditions_changed to EventListener (#226)

* Add on_stall_conditions_changed to EventListener;
Signed-off-by: Tong Zhigao <tongzhigao@pingcap.com>

* Test stall condition changes from normal to other;
Signed-off-by: Tong Zhigao <tongzhigao@pingcap.com>

* Optimize test;
Signed-off-by: Tong Zhigao <tongzhigao@pingcap.com>

librocksdb_sys/crocksdb/c.cc

@@ -76,6 +76,8 @@ using rocksdb::InfoLogLevel;
 using rocksdb::FileLock;
 using rocksdb::FilterPolicy;
 using rocksdb::FlushJobInfo;
+using rocksdb::WriteStallInfo;
+using rocksdb::WriteStallCondition;
 using rocksdb::FlushOptions;
 using rocksdb::IngestExternalFileOptions;
 using rocksdb::Iterator;
@@ -193,6 +195,12 @@ struct crocksdb_pinnableslice_t   { PinnableSlice     rep; };
 struct crocksdb_flushjobinfo_t {
   FlushJobInfo rep;
 };
+struct crocksdb_writestallcondition_t {
+  WriteStallCondition rep;
+};
+struct crocksdb_writestallinfo_t {
+  WriteStallInfo rep;
+};
 struct crocksdb_compactionjobinfo_t {
   CompactionJobInfo rep;
 };
@@ -1792,6 +1800,14 @@ const crocksdb_table_properties_t* crocksdb_flushjobinfo_table_properties(
       &info->rep.table_properties);
 }
 
+bool crocksdb_flushjobinfo_triggered_writes_slowdown(const crocksdb_flushjobinfo_t* info) {
+  return info->rep.triggered_writes_slowdown;
+}
+
+bool crocksdb_flushjobinfo_triggered_writes_stop(const crocksdb_flushjobinfo_t* info) {
+  return info->rep.triggered_writes_stop;
+}
+
 /* CompactionJobInfo */
 
 const char* crocksdb_compactionjobinfo_cf_name(
@@ -1887,6 +1903,26 @@ crocksdb_externalfileingestioninfo_table_properties(
       &info->rep.table_properties);
 }
 
+/* External write stall info */
+extern C_ROCKSDB_LIBRARY_API const char*
+crocksdb_writestallinfo_cf_name(
+    const crocksdb_writestallinfo_t* info, size_t* size) {
+  *size = info->rep.cf_name.size();
+  return info->rep.cf_name.data();
+}
+
+const crocksdb_writestallcondition_t* crocksdb_writestallinfo_cur(
+    const crocksdb_writestallinfo_t* info) {
+  return reinterpret_cast<const crocksdb_writestallcondition_t*>(
+      &info->rep.condition.cur);
+}
+
+const crocksdb_writestallcondition_t* crocksdb_writestallinfo_prev(
+    const crocksdb_writestallinfo_t* info) {
+  return reinterpret_cast<const crocksdb_writestallcondition_t*>(
+      &info->rep.condition.prev);
+}
+
 /* event listener */
 
 struct crocksdb_eventlistener_t : public EventListener {
@@ -1898,6 +1934,7 @@ struct crocksdb_eventlistener_t : public EventListener {
                                   const crocksdb_compactionjobinfo_t*);
   void (*on_external_file_ingested)(
       void*, crocksdb_t*, const crocksdb_externalfileingestioninfo_t*);
+  void (*on_stall_conditions_changed)(void*, const crocksdb_writestallinfo_t*);
 
   virtual void OnFlushCompleted(DB* db, const FlushJobInfo& info) {
     crocksdb_t c_db = {db};
@@ -1920,6 +1957,12 @@ struct crocksdb_eventlistener_t : public EventListener {
         reinterpret_cast<const crocksdb_externalfileingestioninfo_t*>(&info));
   }
 
+  virtual void OnStallConditionsChanged(const WriteStallInfo& info) {
+    on_stall_conditions_changed(
+        state_,
+        reinterpret_cast<const crocksdb_writestallinfo_t*>(&info));
+  }
+
   virtual ~crocksdb_eventlistener_t() { destructor_(state_); }
 };
 
@@ -1927,13 +1970,15 @@ crocksdb_eventlistener_t* crocksdb_eventlistener_create(
     void* state_, void (*destructor_)(void*),
     on_flush_completed_cb on_flush_completed,
     on_compaction_completed_cb on_compaction_completed,
-    on_external_file_ingested_cb on_external_file_ingested) {
+    on_external_file_ingested_cb on_external_file_ingested,
+    on_stall_conditions_changed_cb on_stall_conditions_changed) {
   crocksdb_eventlistener_t* et = new crocksdb_eventlistener_t;
   et->state_ = state_;
   et->destructor_ = destructor_;
   et->on_flush_completed = on_flush_completed;
   et->on_compaction_completed = on_compaction_completed;
   et->on_external_file_ingested = on_external_file_ingested;
+  et->on_stall_conditions_changed = on_stall_conditions_changed;
   return et;
 }
 

librocksdb_sys/crocksdb/crocksdb/c.h

@@ -139,6 +139,8 @@ typedef struct crocksdb_level_meta_data_t crocksdb_level_meta_data_t;
 typedef struct crocksdb_sst_file_meta_data_t crocksdb_sst_file_meta_data_t;
 typedef struct crocksdb_compaction_options_t crocksdb_compaction_options_t;
 typedef struct crocksdb_perf_context_t crocksdb_perf_context_t;
+typedef struct crocksdb_writestallinfo_t crocksdb_writestallinfo_t;
+typedef struct crocksdb_writestallcondition_t crocksdb_writestallcondition_t;
 
 typedef enum crocksdb_table_property_t {
   kDataSize = 1,
@@ -656,6 +658,10 @@ extern C_ROCKSDB_LIBRARY_API const char* crocksdb_flushjobinfo_file_path(
     const crocksdb_flushjobinfo_t*, size_t*);
 extern C_ROCKSDB_LIBRARY_API const crocksdb_table_properties_t*
 crocksdb_flushjobinfo_table_properties(const crocksdb_flushjobinfo_t*);
+extern C_ROCKSDB_LIBRARY_API bool
+crocksdb_flushjobinfo_triggered_writes_slowdown(const crocksdb_flushjobinfo_t*);
+extern C_ROCKSDB_LIBRARY_API bool
+crocksdb_flushjobinfo_triggered_writes_stop(const crocksdb_flushjobinfo_t*);
 
 /* Compaction job info */
 
@@ -709,6 +715,17 @@ extern C_ROCKSDB_LIBRARY_API const crocksdb_table_properties_t*
 crocksdb_externalfileingestioninfo_table_properties(
     const crocksdb_externalfileingestioninfo_t*);
 
+/* External write stall info */
+extern C_ROCKSDB_LIBRARY_API const char*
+crocksdb_writestallinfo_cf_name(
+    const crocksdb_writestallinfo_t*, size_t*);
+extern C_ROCKSDB_LIBRARY_API const crocksdb_writestallcondition_t*
+crocksdb_writestallinfo_cur(
+    const crocksdb_writestallinfo_t*);
+extern C_ROCKSDB_LIBRARY_API const crocksdb_writestallcondition_t*
+crocksdb_writestallinfo_prev(
+    const crocksdb_writestallinfo_t*);
+
 /* Event listener */
 
 typedef void (*on_flush_completed_cb)(void*, crocksdb_t*,
@@ -717,13 +734,15 @@ typedef void (*on_compaction_completed_cb)(void*, crocksdb_t*,
                                            const crocksdb_compactionjobinfo_t*);
 typedef void (*on_external_file_ingested_cb)(
     void*, crocksdb_t*, const crocksdb_externalfileingestioninfo_t*);
+typedef void (*on_stall_conditions_changed_cb)(void*, const crocksdb_writestallinfo_t*);
 
 extern C_ROCKSDB_LIBRARY_API crocksdb_eventlistener_t*
 crocksdb_eventlistener_create(
     void* state_, void (*destructor_)(void*),
     on_flush_completed_cb on_flush_completed,
     on_compaction_completed_cb on_compaction_completed,
-    on_external_file_ingested_cb on_external_file_ingested);
+    on_external_file_ingested_cb on_external_file_ingested,
+    on_stall_conditions_changed_cb on_stall_conditions_changed);
 extern C_ROCKSDB_LIBRARY_API void crocksdb_eventlistener_destroy(
     crocksdb_eventlistener_t*);
 extern C_ROCKSDB_LIBRARY_API void crocksdb_options_add_eventlistener(

librocksdb_sys/src/lib.rs

@@ -68,6 +68,15 @@ pub enum DBLevelMetaData {}
 pub enum DBSstFileMetaData {}
 pub enum DBCompactionOptions {}
 pub enum DBPerfContext {}
+pub enum DBWriteStallInfo {}
+
+#[derive(Copy, Clone, Debug, Eq, PartialEq)]
+#[repr(C)]
+pub enum WriteStallCondition {
+    Normal = 0,
+    Delayed = 1,
+    Stopped = 2,
+}
 
 mod generated;
 pub use generated::*;
@@ -1432,6 +1441,8 @@ extern "C" {
     pub fn crocksdb_flushjobinfo_table_properties(
         info: *const DBFlushJobInfo,
     ) -> *const DBTableProperties;
+    pub fn crocksdb_flushjobinfo_triggered_writes_slowdown(info: *const DBFlushJobInfo) -> bool;
+    pub fn crocksdb_flushjobinfo_triggered_writes_stop(info: *const DBFlushJobInfo) -> bool;
 
     pub fn crocksdb_compactionjobinfo_cf_name(
         info: *const DBCompactionJobInfo,
@@ -1481,12 +1492,23 @@ extern "C" {
         info: *const DBIngestionInfo,
     ) -> *const DBTableProperties;
 
+    pub fn crocksdb_writestallinfo_cf_name(
+        info: *const DBWriteStallInfo,
+        size: *mut size_t,
+    ) -> *const c_char;
+    pub fn crocksdb_writestallinfo_prev(
+        info: *const DBWriteStallInfo,
+    ) -> *const WriteStallCondition;
+    pub fn crocksdb_writestallinfo_cur(info: *const DBWriteStallInfo)
+        -> *const WriteStallCondition;
+
     pub fn crocksdb_eventlistener_create(
         state: *mut c_void,
         destructor: extern "C" fn(*mut c_void),
         flush: extern "C" fn(*mut c_void, *mut DBInstance, *const DBFlushJobInfo),
         compact: extern "C" fn(*mut c_void, *mut DBInstance, *const DBCompactionJobInfo),
         ingest: extern "C" fn(*mut c_void, *mut DBInstance, *const DBIngestionInfo),
+        stall_conditions: extern "C" fn(*mut c_void, *const DBWriteStallInfo),
     ) -> *mut DBEventListener;
     pub fn crocksdb_eventlistener_destroy(et: *mut DBEventListener);
     pub fn crocksdb_options_add_eventlistener(opt: *mut Options, et: *mut DBEventListener);

src/event_listener.rs

@@ -13,6 +13,7 @@
 
 use crocksdb_ffi::{
     self, DBCompactionJobInfo, DBEventListener, DBFlushJobInfo, DBIngestionInfo, DBInstance,
+    DBWriteStallInfo, WriteStallCondition,
 };
 use libc::c_void;
 use std::path::Path;
@@ -46,6 +47,14 @@ impl FlushJobInfo {
             TableProperties::from_ptr(prop)
         }
     }
+
+    pub fn triggered_writes_slowdown(&self) -> bool {
+        unsafe { crocksdb_ffi::crocksdb_flushjobinfo_triggered_writes_slowdown(&self.0) }
+    }
+
+    pub fn triggered_writes_stop(&self) -> bool {
+        unsafe { crocksdb_ffi::crocksdb_flushjobinfo_triggered_writes_stop(&self.0) }
+    }
 }
 
 pub struct CompactionJobInfo(DBCompactionJobInfo);
@@ -133,6 +142,20 @@ impl IngestionInfo {
     }
 }
 
+pub struct WriteStallInfo(DBWriteStallInfo);
+
+impl WriteStallInfo {
+    pub fn cf_name(&self) -> &str {
+        unsafe { fetch_str!(crocksdb_writestallinfo_cf_name(&self.0)) }
+    }
+    pub fn cur(&self) -> WriteStallCondition {
+        unsafe { *crocksdb_ffi::crocksdb_writestallinfo_cur(&self.0) }
+    }
+    pub fn prev(&self) -> WriteStallCondition {
+        unsafe { *crocksdb_ffi::crocksdb_writestallinfo_prev(&self.0) }
+    }
+}
+
 /// EventListener trait contains a set of call-back functions that will
 /// be called when specific RocksDB event happens such as flush.  It can
 /// be used as a building block for developing custom features such as
@@ -146,6 +169,7 @@ pub trait EventListener: Send + Sync {
     fn on_flush_completed(&self, _: &FlushJobInfo) {}
     fn on_compaction_completed(&self, _: &CompactionJobInfo) {}
     fn on_external_file_ingested(&self, _: &IngestionInfo) {}
+    fn on_stall_conditions_changed(&self, _: &WriteStallInfo) {}
 }
 
 extern "C" fn destructor(ctx: *mut c_void) {
@@ -183,6 +207,11 @@ extern "C" fn on_external_file_ingested(
     ctx.on_external_file_ingested(info);
 }
 
+extern "C" fn on_stall_conditions_changed(ctx: *mut c_void, info: *const DBWriteStallInfo) {
+    let (ctx, info) = unsafe { (&*(ctx as *mut Box<EventListener>), mem::transmute(&*info)) };
+    ctx.on_stall_conditions_changed(info);
+}
+
 pub fn new_event_listener<L: EventListener>(l: L) -> *mut DBEventListener {
     let p: Box<EventListener> = Box::new(l);
     unsafe {
@@ -192,6 +221,7 @@ pub fn new_event_listener<L: EventListener>(l: L) -> *mut DBEventListener {
             on_flush_completed,
             on_compaction_completed,
             on_external_file_ingested,
+            on_stall_conditions_changed,
         )
     }
 }

src/lib.rs

@@ -21,11 +21,13 @@ pub extern crate librocksdb_sys;
 extern crate tempdir;
 
 pub use compaction_filter::CompactionFilter;
-pub use event_listener::{CompactionJobInfo, EventListener, FlushJobInfo, IngestionInfo};
+pub use event_listener::{
+    CompactionJobInfo, EventListener, FlushJobInfo, IngestionInfo, WriteStallInfo,
+};
 pub use librocksdb_sys::{
     self as crocksdb_ffi, new_bloom_filter, CompactionPriority, DBBottommostLevelCompaction,
     DBCompactionStyle, DBCompressionType, DBEntryType, DBInfoLogLevel, DBRecoveryMode,
-    DBStatisticsHistogramType, DBStatisticsTickerType,
+    DBStatisticsHistogramType, DBStatisticsTickerType, WriteStallCondition,
 };
 pub use merge_operator::MergeOperands;
 pub use metadata::{ColumnFamilyMetaData, LevelMetaData, SstFileMetaData};

tests/test_event_listener.rs

@@ -85,6 +85,79 @@ impl EventListener for EventCounter {
     }
 }
 
+#[derive(Default, Clone)]
+struct StallEventCounter {
+    flush: Arc<AtomicUsize>,
+    stall_conditions_changed_num: Arc<AtomicUsize>,
+    triggered_writes_slowdown: Arc<AtomicUsize>,
+    triggered_writes_stop: Arc<AtomicUsize>,
+    stall_change_from_normal_to_other: Arc<AtomicUsize>,
+}
+
+impl EventListener for StallEventCounter {
+    fn on_flush_completed(&self, info: &FlushJobInfo) {
+        assert!(!info.cf_name().is_empty());
+        self.flush.fetch_add(1, Ordering::SeqCst);
+        self.triggered_writes_slowdown
+            .fetch_add(info.triggered_writes_slowdown() as usize, Ordering::SeqCst);
+        self.triggered_writes_stop
+            .fetch_add(info.triggered_writes_stop() as usize, Ordering::SeqCst);
+    }
+
+    fn on_stall_conditions_changed(&self, info: &WriteStallInfo) {
+        assert!(info.cf_name() == "test_cf");
+        self.stall_conditions_changed_num
+            .fetch_add(1, Ordering::SeqCst);
+        if info.prev() == WriteStallCondition::Normal && info.cur() != WriteStallCondition::Normal {
+            self.stall_change_from_normal_to_other
+                .fetch_add(1, Ordering::SeqCst);
+        }
+    }
+}
+
+#[test]
+fn test_event_listener_stall_conditions_changed() {
+    let path = TempDir::new("_rust_rocksdb_event_listener_stall_conditions").expect("");
+    let path_str = path.path().to_str().unwrap();
+
+    let mut opts = DBOptions::new();
+    let counter = StallEventCounter::default();
+    opts.add_event_listener(counter.clone());
+    opts.create_if_missing(true);
+    let mut cf_opts = ColumnFamilyOptions::new();
+    cf_opts.set_level_zero_slowdown_writes_trigger(1);
+    cf_opts.set_level_zero_stop_writes_trigger(1);
+    cf_opts.set_level_zero_file_num_compaction_trigger(1);
+    let mut db = DB::open_cf(
+        opts,
+        path_str,
+        vec![("default", ColumnFamilyOptions::new())],
+    ).unwrap();
+    db.create_cf(("test_cf", cf_opts)).unwrap();
+
+    let test_cf = db.cf_handle("test_cf").unwrap();
+    for i in 1..5 {
+        db.put_cf(
+            test_cf,
+            format!("{:04}", i).as_bytes(),
+            format!("{:04}", i).as_bytes(),
+        ).unwrap();
+        db.flush_cf(test_cf, true).unwrap();
+    }
+    let flush_cnt = counter.flush.load(Ordering::SeqCst);
+    assert_ne!(flush_cnt, 0);
+    let stall_conditions_changed_num = counter.stall_conditions_changed_num.load(Ordering::SeqCst);
+    let triggered_writes_slowdown = counter.triggered_writes_slowdown.load(Ordering::SeqCst);
+    let triggered_writes_stop = counter.triggered_writes_stop.load(Ordering::SeqCst);
+    let stall_change_from_normal_to_other = counter
+        .stall_change_from_normal_to_other
+        .load(Ordering::SeqCst);
+    assert_ne!(stall_conditions_changed_num, 0);
+    assert_ne!(triggered_writes_slowdown, 0);
+    assert_ne!(triggered_writes_stop, 0);
+    assert_ne!(stall_change_from_normal_to_other, 0);
+}
+
 #[test]
 fn test_event_listener_basic() {
     let path = TempDir::new("_rust_rocksdb_event_listener_flush").expect("");