producer: flush callback

src/hstreamdb/src/client.rs

@@ -10,7 +10,7 @@ use url::Url;
 
 use crate::appender::Appender;
 use crate::channel_provider::{new_channel_provider, ChannelProviderSettings, Channels};
-use crate::producer::{FlushSettings, Producer};
+use crate::producer::{BoxedFlushCallback, FlushSettings, Producer};
 use crate::{common, flow_controller, format_url, producer};
 
 pub struct Client {
@@ -172,6 +172,7 @@ impl Client {
         flow_control_bytes_limit: Option<usize>,
         flush_settings: FlushSettings,
         channel_provider_settings: ChannelProviderSettings,
+        flush_callback: Option<BoxedFlushCallback>,
     ) -> common::Result<(Appender, Producer)> {
         let (request_sender, request_receiver) =
             tokio::sync::mpsc::unbounded_channel::<producer::Request>();
@@ -195,6 +196,7 @@ impl Client {
             flow_controller,
             compression_type,
             flush_settings,
+            flush_callback,
         )
         .await?;
         Ok((appender, producer))

src/hstreamdb/src/producer.rs

@@ -50,8 +50,11 @@ pub struct Producer {
     compression_type: CompressionType,
     flush_settings: FlushSettings,
     shards: Vec<Shard>,
+    flush_callback: Option<BoxedFlushCallback>,
 }
 
+pub type BoxedFlushCallback = Arc<dyn Fn(bool, usize, usize) + Send + Sync>;
+
 #[derive(Default)]
 struct BufferState {
     len: usize,
@@ -132,6 +135,7 @@ impl BufferState {
 }
 
 impl Producer {
+    #[allow(clippy::too_many_arguments)]
     pub(crate) async fn new(
         channels: Channels,
         url_scheme: String,
@@ -140,6 +144,7 @@ impl Producer {
         flow_controller: Option<FlowControllerClient>,
         compression_type: CompressionType,
         flush_settings: FlushSettings,
+        flush_callback: Option<BoxedFlushCallback>,
     ) -> common::Result<Self> {
         let shards = channels
             .channel()
@@ -169,6 +174,7 @@ impl Producer {
             compression_type,
             flush_settings,
             shards,
+            flush_callback,
         };
         Ok(producer)
     }
@@ -296,7 +302,9 @@ impl Producer {
             shard_url,
             self.compression_type,
             buffer,
+            buffer_size,
             results,
+            self.flush_callback.clone(),
         );
         let task = tokio::spawn(async move {
             task.await;
@@ -326,14 +334,17 @@ impl Producer {
     }
 }
 
+#[allow(clippy::too_many_arguments)]
 async fn flush(
     channels: Channels,
     stream_name: String,
     shard_id: ShardId,
     shard_url: String,
     compression_type: CompressionType,
     buffer: Vec<Record>,
+    buffer_size: usize,
     results: ResultVec,
+    flush_callback: Option<BoxedFlushCallback>,
 ) -> Result<(), String> {
     if buffer.is_empty() {
         Ok(())
@@ -342,15 +353,20 @@ async fn flush(
             .channel_at(shard_url.clone())
             .await
             .map_err(|err| format!("producer connect error: url = {shard_url}, {err}"))?;
-        match append(
+        let append_result = append(
             channel,
             stream_name,
             shard_id,
             compression_type,
             buffer.to_vec(),
         )
-        .await
-        {
+        .await;
+
+        if let Some(flush_callback) = flush_callback.as_ref() {
+            flush_callback(append_result.is_ok(), buffer.len(), buffer_size)
+        }
+
+        match append_result {
             Err(err) => {
                 let err = Arc::new(err);
                 for sender in results.into_iter() {
@@ -377,14 +393,17 @@ async fn flush(
     }
 }
 
+#[allow(clippy::too_many_arguments)]
 async fn flush_(
     channels: Channels,
     stream_name: String,
     shard_id: ShardId,
     shard_url: String,
     compression_type: CompressionType,
     buffer: Vec<Record>,
+    buffer_size: usize,
     results: ResultVec,
+    flush_callback: Option<BoxedFlushCallback>,
 ) {
     flush(
         channels,
@@ -393,7 +412,9 @@ async fn flush_(
         shard_url,
         compression_type,
         buffer,
+        buffer_size,
         results,
+        flush_callback,
     )
     .await
     .unwrap_or_else(|err| log::error!("{err}"))

src/hstreamdb/tests/consumer_test.rs

@@ -54,6 +54,7 @@ async fn test_consumer() {
             ChannelProviderSettings {
                 concurrency_limit: Some(8),
             },
+            None,
         )
         .await
         .unwrap();

src/hstreamdb/tests/producer_test.rs

@@ -43,6 +43,7 @@ async fn test_producer() {
             ChannelProviderSettings {
                 concurrency_limit: Some(8),
             },
+            None,
         )
         .await
         .unwrap();

src/x/hstreamdb-erl-nifs/src/lib.rs

@@ -16,7 +16,8 @@ mod runtime;
 rustler::atoms! {
     compression_type, none, gzip, zstd,
     concurrency_limit,
-    max_batch_len, max_batch_size, batch_deadline
+    max_batch_len, max_batch_size, batch_deadline,
+    record_id
 }
 
 rustler::init!(
@@ -33,7 +34,7 @@ rustler::init!(
 
 #[derive(Debug)]
 enum AppendResult {
-    RecordId(String),
+    RecordId(RecordId),
     Error(String),
 }
 
@@ -129,6 +130,7 @@ pub fn try_start_producer(
                     None,
                     flush_settings,
                     ChannelProviderSettings { concurrency_limit },
+                    None,
                 )
                 .await?;
 
@@ -207,23 +209,28 @@ fn await_append_result(env: Env, x: ResourceArc<AppendResultFuture>) -> Term {
         let receiver: &Mutex<_> = &x.0;
         let mut receiver: MutexGuard<Option<_>> = receiver.lock();
         let receiver = mem::take(&mut (*receiver));
-        let append_result: Result<String, Arc<_>> = receiver
-            .unwrap()
-            .blocking_recv()
-            .unwrap()
-            .map(|x| x.to_string());
+        let append_result: Result<hstreamdb::RecordId, Arc<_>> =
+            receiver.unwrap().blocking_recv().unwrap();
         let append_result = match append_result {
             Ok(record_id) => RecordId(record_id),
             Err(err) => Error(err.to_string()),
         };
         result.set(append_result).unwrap()
     }
 
-    let result = match result.get().unwrap() {
-        RecordId(record_id) => (ok(), record_id.to_string()),
-        Error(err) => (error(), err.to_string()),
-    };
-    result.encode(env)
+    match result.get().unwrap() {
+        RecordId(record_id_v) => (
+            ok(),
+            (
+                record_id(),
+                record_id_v.shard_id,
+                record_id_v.batch_id,
+                record_id_v.batch_index,
+            ),
+        )
+            .encode(env),
+        Error(err) => (error(), err.to_string()).encode(env),
+    }
 }
 
 fn atom_to_compression_type(compression_type: Atom) -> Option<CompressionType> {