Merge pull request #45 from ainestal/burst

Burst

Author: hugues31

src/bitstamp/api.rs

@@ -12,6 +12,8 @@ use serde_json::value::Map;
 
 use std::collections::HashMap;
 use std::io::Read;
+use std::thread;
+use std::time::Duration;
 
 use coinnect::Credentials;
 use exchange::Exchange;
@@ -44,6 +46,7 @@ pub struct BitstampApi {
     api_secret: String,
     customer_id: String,
     http_client: Client,
+    burst: bool,
 }
 
 
@@ -69,9 +72,30 @@ impl BitstampApi {
                api_secret: creds.get("api_secret").unwrap_or_default(),
                customer_id: creds.get("customer_id").unwrap_or_default(),
                http_client: Client::with_connector(connector),
+               burst: false, // No burst by default
            })
     }
 
+    /// The number of calls in a given period is limited. In order to avoid a ban we limit
+    /// by default the number of api requests.
+    /// This function sets or removes the limitation.
+    /// Burst false implies no block.
+    /// Burst true implies there is a control over the number of calls allowed to the exchange
+    pub fn set_burst(&mut self, burst: bool) {
+        self.burst = burst
+    }
+
+    fn block_or_continue(&self) {
+        if ! self.burst {
+            let threshold: u64 = 1000; // 600 requests per 10 mins = 1 request per second
+            let offset: u64 = helpers::get_unix_timestamp_ms() as u64 - self.last_request as u64;
+            if offset < threshold {
+                let wait_ms = Duration::from_millis(threshold - offset);
+                thread::sleep(wait_ms);
+            }
+        }
+    }
+
     fn public_query(&mut self, params: &HashMap<&str, &str>) -> Result<Map<String, Value>> {
 
         let method: &str = params
@@ -80,7 +104,7 @@ impl BitstampApi {
         let pair: &str = params.get("pair").ok_or_else(|| "Missing \"pair\" field.")?;
         let url: String = utils::build_url(method, pair);
 
-        utils::block_or_continue(self.last_request);
+        self.block_or_continue();
         let mut response = self.http_client.get(&url).send()?;
         self.last_request = helpers::get_unix_timestamp_ms();
         let mut buffer = String::new();
@@ -327,3 +351,45 @@ impl BitstampApi {
         self.private_query(&params)
     }
 }
+
+
+#[cfg(test)]
+mod bitstamp_api_tests {
+    use super::*;
+
+    #[test]
+    fn should_block_or_not_block_when_enabled_or_disabled() {
+        let mut api = BitstampApi {
+            last_request: helpers::get_unix_timestamp_ms(),
+            api_key: "".to_string(),
+            api_secret: "".to_string(),
+            customer_id: "".to_string(),
+            http_client: Client::new(),
+            burst: false,
+        };
+
+        let mut counter = 0;
+        loop {
+            api.set_burst(false);
+            let start = helpers::get_unix_timestamp_ms();
+            api.block_or_continue();
+            api.last_request = helpers::get_unix_timestamp_ms();
+
+            let difference = api.last_request - start;
+            assert!(difference >= 999);
+            assert!(difference < 10000);
+
+
+            api.set_burst(true);
+            let start = helpers::get_unix_timestamp_ms();
+            api.block_or_continue();
+            api.last_request = helpers::get_unix_timestamp_ms();
+
+            let difference = api.last_request - start;
+            assert!(difference < 10);
+
+            counter = counter + 1;
+            if counter >= 3 { break; }
+        }
+    }
+}

src/bitstamp/utils.rs

@@ -7,9 +7,6 @@ use serde_json;
 use serde_json::Value;
 use serde_json::value::Map;
 
-use std::thread;
-use std::time::Duration;
-
 use error::*;
 use helpers;
 use types::Currency;
@@ -41,15 +38,6 @@ pub fn get_pair_enum(pair: &str) -> Option<&Pair> {
     PAIRS_STRING.get_by_second(&pair)
 }
 
-pub fn block_or_continue(last_request: i64) {
-    let threshold: u64 = 1000; // 600 requests per 10 mins = 1 request per second
-    let offset: u64 = helpers::get_unix_timestamp_ms() as u64 - last_request as u64;
-    if offset < threshold {
-        let wait_ms = Duration::from_millis(threshold - offset);
-        thread::sleep(wait_ms);
-    }
-}
-
 pub fn build_signature(nonce: &str,
                        customer_id: &str,
                        api_key: &str,
@@ -160,16 +148,3 @@ pub fn get_currency_string(currency: Currency) -> Option<String> {
         _ => None,
     }
 }
-
-#[cfg(test)]
-mod utils_tests {
-    use super::*;
-
-    #[test]
-    fn should_block_when_enabled() {
-        let start = helpers::get_unix_timestamp_ms();
-        block_or_continue(start);
-        let end = helpers::get_unix_timestamp_ms();
-        assert!(end - start >= 1000)
-    }
-}

src/kraken/api.rs

@@ -46,6 +46,7 @@ pub struct KrakenApi {
     api_secret: String,
     otp: Option<String>, // two-factor password (if two-factor enabled, otherwise not required)
     http_client: Client,
+    burst: bool,
 }
 
 
@@ -69,6 +70,7 @@ impl KrakenApi {
                api_secret: creds.get("api_secret").unwrap_or_default(),
                otp: None,
                http_client: Client::with_connector(connector),
+               burst: false,
            })
     }
 
@@ -78,12 +80,23 @@ impl KrakenApi {
         self.otp = Some(otp);
     }
 
+    /// The number of calls in a given period is limited. In order to avoid a ban we limit
+    /// by default the number of api requests.
+    /// This function sets or removes the limitation.
+    /// Burst false implies no block.
+    /// Burst true implies there is a control over the number of calls allowed to the exchange
+    pub fn set_burst(&mut self, burst: bool) {
+        self.burst = burst
+    }
+
     pub fn block_or_continue(&self) {
-        let threshold: u64 = 2000; // 1 request/2sec
-        let offset: u64 = helpers::get_unix_timestamp_ms() as u64 - self.last_request as u64;
-        if offset < threshold {
-            let wait_ms = Duration::from_millis(threshold - offset);
-            thread::sleep(wait_ms);
+        if ! self.burst {
+            let threshold: u64 = 2000; // 1 request/2sec
+            let offset: u64 = helpers::get_unix_timestamp_ms() as u64 - self.last_request as u64;
+            if offset < threshold {
+                let wait_ms = Duration::from_millis(threshold - offset);
+                thread::sleep(wait_ms);
+            }
         }
     }
 
@@ -1111,3 +1124,44 @@ impl KrakenApi {
         self.private_query("WithdrawCancel", &mut params)
     }
 }
+
+#[cfg(test)]
+mod kraken_api_tests {
+    use super::*;
+
+    #[test]
+    fn should_block_or_not_block_when_enabled_or_disabled() {
+        let mut api = KrakenApi {
+            last_request: helpers::get_unix_timestamp_ms(),
+            api_key: "".to_string(),
+            api_secret: "".to_string(),
+            otp: None,
+            http_client: Client::new(),
+            burst: false,
+        };
+
+        let mut counter = 0;
+        loop {
+            api.set_burst(false);
+            let start = helpers::get_unix_timestamp_ms();
+            api.block_or_continue();
+            api.last_request = helpers::get_unix_timestamp_ms();
+
+            let difference = api.last_request - start;
+            assert!(difference >= 1999);
+            assert!(difference < 10000);
+
+
+            api.set_burst(true);
+            let start = helpers::get_unix_timestamp_ms();
+            api.block_or_continue();
+            api.last_request = helpers::get_unix_timestamp_ms();
+
+            let difference = api.last_request - start;
+            assert!(difference < 10);
+
+            counter = counter + 1;
+            if counter >= 3 { break; }
+        }
+    }
+}

src/poloniex/api.rs

@@ -47,6 +47,7 @@ pub struct PoloniexApi {
     api_key: String,
     api_secret: String,
     http_client: Client,
+    burst: bool,
 }
 
 
@@ -69,15 +70,27 @@ impl PoloniexApi {
                api_key: creds.get("api_key").unwrap_or_default(),
                api_secret: creds.get("api_secret").unwrap_or_default(),
                http_client: Client::with_connector(connector),
+               burst: false,
            })
     }
 
+    /// The number of calls in a given period is limited. In order to avoid a ban we limit
+    /// by default the number of api requests.
+    /// This function sets or removes the limitation.
+    /// Burst false implies no block.
+    /// Burst true implies there is a control over the number of calls allowed to the exchange
+    pub fn set_burst(&mut self, burst: bool) {
+        self.burst = burst
+    }
+
     fn block_or_continue(&self) {
-        let threshold: u64 = 167; // 6 requests/sec = 1/6*1000
-        let offset: u64 = helpers::get_unix_timestamp_ms() as u64 - self.last_request as u64;
-        if offset < threshold {
-            let wait_ms = Duration::from_millis(threshold - offset);
-            thread::sleep(wait_ms);
+        if ! self.burst {
+            let threshold: u64 = 167; // 6 requests/sec = 1/6*1000
+            let offset: u64 = helpers::get_unix_timestamp_ms() as u64 - self.last_request as u64;
+            if offset < threshold {
+                let wait_ms = Duration::from_millis(threshold - offset);
+                thread::sleep(wait_ms);
+            }
         }
     }
 
@@ -787,3 +800,44 @@ impl PoloniexApi {
         self.private_query("toggleAutoRenew", &params)
     }
 }
+
+
+#[cfg(test)]
+mod poloniex_api_tests {
+    use super::*;
+
+    #[test]
+    fn should_block_or_not_block_when_enabled_or_disabled() {
+        let mut api = PoloniexApi {
+            last_request: helpers::get_unix_timestamp_ms(),
+            api_key: "".to_string(),
+            api_secret: "".to_string(),
+            http_client: Client::new(),
+            burst: false,
+        };
+
+        let mut counter = 0;
+        loop {
+            api.set_burst(false);
+            let start = helpers::get_unix_timestamp_ms();
+            api.block_or_continue();
+            api.last_request = helpers::get_unix_timestamp_ms();
+
+            let difference = api.last_request - start;
+            assert!(difference >= 166);
+            assert!(difference < 1000);
+
+
+            api.set_burst(true);
+            let start = helpers::get_unix_timestamp_ms();
+            api.block_or_continue();
+            api.last_request = helpers::get_unix_timestamp_ms();
+
+            let difference = api.last_request - start;
+            assert!(difference < 10);
+
+            counter = counter + 1;
+            if counter >= 3 { break; }
+        }
+    }
+}

tests/coinnect.rs

@@ -24,7 +24,7 @@ mod coinnect_tests {
                    "BitstampApi { last_request: 0, api_key: \"bs_api_key\", api_secret: \
         \"bs_api_secret\", customer_id: \"bs_cust_id\", http_client: Client { \
                     redirect_policy: FollowAll, read_timeout: None, write_timeout: None, proxy: \
-                    None } }");
+                    None }, burst: false }");
     }
     #[test]
     fn can_create_new_api_connection_to_kraken() {