flow_controller: init

src/hstreamdb/src/appender.rs

@@ -3,24 +3,35 @@ use std::sync::Arc;
 use tokio::sync::oneshot;
 
 use crate::common::{self, Record};
+use crate::flow_controller::FlowControllerClient;
 use crate::producer::{self, Request};
 
 #[derive(Clone)]
 pub struct Appender {
     request_sender: tokio::sync::mpsc::UnboundedSender<Request>,
+    flow_controller: Option<FlowControllerClient>,
 }
 
 impl Appender {
-    pub(crate) fn new(request_sender: tokio::sync::mpsc::UnboundedSender<Request>) -> Appender {
-        Appender { request_sender }
+    pub(crate) fn new(
+        request_sender: tokio::sync::mpsc::UnboundedSender<Request>,
+        flow_controller: Option<FlowControllerClient>,
+    ) -> Appender {
+        Appender {
+            request_sender,
+            flow_controller,
+        }
     }
 }
 
 impl Appender {
-    pub fn append(
+    pub async fn append(
         &mut self,
         record: Record,
-    ) -> Result<oneshot::Receiver<Result<String, Arc<common::Error>>>, producer::SendError> {
+    ) -> common::Result<oneshot::Receiver<Result<String, Arc<common::Error>>>> {
+        if let Some(flow_controller) = &self.flow_controller {
+            flow_controller.acquire(record.encoded_len()).await?
+        }
         let (sender, receiver) = oneshot::channel();
         self.request_sender
             .send(Request(record, sender))

src/hstreamdb/src/client.rs

@@ -11,7 +11,7 @@ use url::Url;
 use crate::appender::Appender;
 use crate::channel_provider::{new_channel_provider, ChannelProviderSettings, Channels};
 use crate::producer::{FlushSettings, Producer};
-use crate::{common, format_url, producer};
+use crate::{common, flow_controller, format_url, producer};
 
 pub struct Client {
     pub(crate) channels: Channels,
@@ -169,6 +169,7 @@ impl Client {
         &mut self,
         stream_name: String,
         compression_type: CompressionType,
+        flow_control_bytes_limit: Option<usize>,
         flush_settings: FlushSettings,
         channel_provider_settings: ChannelProviderSettings,
     ) -> common::Result<(Appender, Producer)> {
@@ -177,12 +178,21 @@ impl Client {
 
         let channels = self.new_channel_provider(channel_provider_settings).await?;
 
-        let appender = Appender::new(request_sender.clone());
+        let flow_controller = {
+            if let Some(flow_control_bytes_limit) = flow_control_bytes_limit {
+                Some(flow_controller::start(flow_control_bytes_limit).await)
+            } else {
+                None
+            }
+        };
+
+        let appender = Appender::new(request_sender.clone(), flow_controller.clone());
         let producer = Producer::new(
             channels,
             self.url_scheme.clone(),
             request_receiver,
             stream_name,
+            flow_controller,
             compression_type,
             flush_settings,
         )

src/hstreamdb/src/common.rs

@@ -3,8 +3,11 @@ use std::io;
 use hstreamdb_pb::StreamingFetchRequest;
 pub use hstreamdb_pb::{CompressionType, SpecialOffset, Stream};
 use num_bigint::ParseBigIntError;
+use prost::Message;
 use tonic::transport;
 
+use crate::producer;
+
 #[derive(Debug)]
 pub struct Subscription {
     pub subscription_id: String,
@@ -59,6 +62,8 @@ pub enum Error {
     NoChannelAvailable,
     #[error("{0}")]
     BadArgument(String),
+    #[error(transparent)]
+    AppenderSendError(producer::SendError),
 }
 
 #[derive(Debug, thiserror::Error)]
@@ -87,6 +92,12 @@ impl From<url::ParseError> for Error {
     }
 }
 
+impl From<producer::SendError> for Error {
+    fn from(err: producer::SendError) -> Self {
+        Error::AppenderSendError(err)
+    }
+}
+
 pub type Result<A> = std::result::Result<A, Error>;
 
 #[derive(Debug, Clone)]
@@ -101,6 +112,21 @@ pub enum Payload {
     RawRecord(Vec<u8>),
 }
 
+impl Payload {
+    pub(crate) fn encoded_len(&self) -> usize {
+        match self {
+            Payload::HRecord(x) => x.encoded_len(),
+            Payload::RawRecord(x) => x.len(),
+        }
+    }
+}
+
+impl Record {
+    pub(crate) fn encoded_len(&self) -> usize {
+        self.payload.encoded_len()
+    }
+}
+
 impl From<Vec<u8>> for Payload {
     fn from(payload: Vec<u8>) -> Self {
         Self::RawRecord(payload)

src/hstreamdb/src/consumer.rs

@@ -131,7 +131,7 @@ fn process_streaming_fetch_response(
                     let ack_sender = ack_sender.clone();
                     let subscription_id = subscription_id.clone();
                     let consumer_name = consumer_name.clone();
-                    let ack_fn: AckFn = box (move || {
+                    let ack_fn: AckFn = Box::new(move || {
                         ack_sender.send(StreamingFetchRequest {
                             subscription_id,
                             consumer_name,

src/hstreamdb/src/flow_controller.rs

@@ -0,0 +1,102 @@
+use std::collections::VecDeque;
+
+use tokio::select;
+
+use crate::common;
+
+#[derive(Clone)]
+pub(crate) struct FlowControllerClient {
+    acquire_request_sender: tokio::sync::mpsc::Sender<AcquireRequest>,
+    release_request_sender: tokio::sync::mpsc::Sender<usize>,
+    bytes_limit: usize,
+}
+
+impl FlowControllerClient {
+    pub(crate) async fn acquire(&self, n: usize) -> common::Result<()> {
+        if n > self.bytes_limit {
+            Err(common::Error::BadArgument(format!(
+                "payload size {n} is larger than flow control limit {}",
+                self.bytes_limit
+            )))
+        } else {
+            let (sender, receiver) = tokio::sync::oneshot::channel();
+            self.acquire_request_sender.send((n, sender)).await.unwrap();
+            receiver.await.unwrap();
+            Ok(())
+        }
+    }
+
+    pub(crate) async fn release(&self, n: usize) {
+        self.release_request_sender.send(n).await.unwrap()
+    }
+}
+
+type AcquireRequest = (usize, tokio::sync::oneshot::Sender<()>);
+
+pub(crate) struct FlowControllerServer {
+    acquire_request_receiver: tokio::sync::mpsc::Receiver<AcquireRequest>,
+    release_request_receiver: tokio::sync::mpsc::Receiver<usize>,
+    bytes_available: usize,
+    awaiting_requests: VecDeque<AcquireRequest>,
+}
+
+pub(crate) async fn start(bytes_limit: usize) -> FlowControllerClient {
+    let (acquire_request_sender, acquire_request_receiver) = tokio::sync::mpsc::channel(1);
+    let (release_request_sender, release_request_receiver) = tokio::sync::mpsc::channel(1);
+    _ = tokio::spawn(
+        (FlowControllerServer {
+            acquire_request_receiver,
+            release_request_receiver,
+            bytes_available: bytes_limit,
+            awaiting_requests: VecDeque::new(),
+        })
+        .start(),
+    );
+    FlowControllerClient {
+        acquire_request_sender,
+        release_request_sender,
+        bytes_limit,
+    }
+}
+
+impl FlowControllerServer {
+    async fn start(mut self) {
+        loop {
+            select! {
+                  biased;
+                  request = self.release_request_receiver.recv() => match request {
+                        Some(n) => self.handle_release(n).await,
+                        None => {
+                            break;
+                    }
+                },
+                  request = self.acquire_request_receiver.recv() => match request {
+                        Some(n) => self.handle_acquire(n).await,
+                        None => {
+                            break;
+                    }
+                },
+            }
+        }
+    }
+
+    async fn handle_release(&mut self, n: usize) {
+        self.bytes_available += n;
+        if let Some(head) = self.awaiting_requests.get(0) {
+            if self.bytes_available >= head.0 {
+                let (n, awaiting_request) = self.awaiting_requests.pop_front().unwrap();
+                self.bytes_available -= n;
+                awaiting_request.send(()).unwrap()
+            }
+        }
+    }
+
+    async fn handle_acquire(&mut self, (n, awaiting_request): AcquireRequest) {
+        if self.bytes_available >= n {
+            self.bytes_available -= n;
+            awaiting_request.send(()).unwrap()
+        } else {
+            self.awaiting_requests.push_back((n, awaiting_request))
+        }
+    }
+}

src/hstreamdb/src/lib.rs

@@ -1,11 +1,11 @@
-#![feature(box_syntax)]
 #![feature(default_free_fn)]
 
 pub mod appender;
 mod channel_provider;
 pub mod client;
 pub mod common;
 pub mod consumer;
+mod flow_controller;
 pub mod producer;
 pub mod utils;