minor improvments for simian army

quickwit-oss · Dec 8, 2023 · 9d39e1b · 9d39e1b
1 parent baf224c
commit 9d39e1b
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 1 deletion.
diff --git a/quickwit/quickwit-common/src/rate_limiter.rs b/quickwit/quickwit-common/src/rate_limiter.rs
@@ -101,6 +101,19 @@ impl RateLimiter {
         }
     }
 
+    pub fn acquire_with_duration(&mut self, num_permits: u64) -> Option<Duration> {
+        if self.acquire_inner(num_permits) {
+            return None;
+        }
+        self.refill(Instant::now());
+        if self.acquire_inner(num_permits) {
+            return None;
+        }
+        let missing = num_permits - self.available_permits;
+        let wait = Duration::from_micros(missing * self.refill_period_micros / self.refill_amount);
+        Some(wait)
+    }
+
     pub fn acquire_bytes(&mut self, bytes: ByteSize) -> bool {
         self.acquire(bytes.as_u64())
     }

diff --git a/quickwit/quickwit-rest-client/src/lib.rs b/quickwit/quickwit-rest-client/src/lib.rs
@@ -29,6 +29,10 @@ pub mod error;
 pub mod models;
 pub mod rest_client;
 
+// re-exports
+pub use quickwit_config::ConfigFormat;
+pub use reqwest::Url;
+
 pub(crate) struct BatchLineReader {
     buf_reader: BufReader<Box<dyn AsyncRead + Send + Sync + Unpin>>,
     buffer: Vec<u8>,

diff --git a/quickwit/quickwit-rest-client/src/rest_client.rs b/quickwit/quickwit-rest-client/src/rest_client.rs
@@ -329,11 +329,12 @@ impl<'a> IndexClient<'a> {
 
     pub async fn create(
         &self,
-        body: Bytes,
+        index_config: &str,
         config_format: ConfigFormat,
         overwrite: bool,
     ) -> Result<IndexMetadata, Error> {
         let header_map = header_from_config_format(config_format);
+        let body = bytes::Bytes::from(index_config.to_string());
         let response = self
             .transport
             .send(