feat: exercise limits for read session

src/main.rs

@@ -204,43 +204,62 @@ enum StreamActions {
 
     /// Append records to a stream. Currently, only newline delimited records are supported.
     Append {
-        /// Newline delimited records to append from a file or stdin (all records are treated as plain text).
+        /// Input newline delimited records to append from a file or stdin.
+        /// All records are treated as plain text.
         /// Use "-" to read from stdin.
-        #[arg(value_parser = parse_records_input_source)]
-        records: RecordsIO,
+        #[arg(value_parser = parse_records_input_source, default_value = "-")]
+        records: RecordsIn,
     },
 
+    /// Read records from a stream.
+    /// If a limit if specified, reading will stop when the limit is reached or there are no more records on the stream.
+    /// If a limit is not specified, the reader will keep tailing and wait for new records.
     Read {
         /// Starting sequence number (inclusive). If not specified, the latest record.
-        start_seq_num: Option<u64>,
+        #[arg(short, long)]
+        start: u64,
 
         /// Output records to a file or stdout.
         /// Use "-" to write to stdout.
-        #[arg(value_parser = parse_records_output_source)]
-        output: Option<RecordsIO>,
+        #[arg(value_parser = parse_records_output_source, default_value = "-")]
+        output: Option<RecordsOut>,
+
+        /// Limit the number of records returned.
+        #[arg(long)]
+        limit_count: Option<u64>,
+
+        /// Limit the number of bytes returned.
+        #[arg(long)]
+        limit_bytes: Option<ByteSize>,
     },
 }
 
-/// Source of records for an append session.
 #[derive(Debug, Clone)]
-pub enum RecordsIO {
+pub enum RecordsIn {
     File(PathBuf),
     Stdin,
+}
+
+/// Sink for records in a read session.
+#[derive(Debug, Clone)]
+pub enum RecordsOut {
+    File(PathBuf),
     Stdout,
 }
 
-impl RecordsIO {
+impl RecordsIn {
     pub async fn into_reader(&self) -> std::io::Result<Box<dyn AsyncBufRead + Send + Unpin>> {
         match self {
-            RecordsIO::File(path) => Ok(Box::new(BufReader::new(File::open(path).await?))),
-            RecordsIO::Stdin => Ok(Box::new(BufReader::new(tokio::io::stdin()))),
-            _ => panic!("unsupported record source"),
+            RecordsIn::File(path) => Ok(Box::new(BufReader::new(File::open(path).await?))),
+            RecordsIn::Stdin => Ok(Box::new(BufReader::new(tokio::io::stdin()))),
         }
     }
+}
 
+impl RecordsOut {
     pub async fn into_writer(&self) -> io::Result<Box<dyn AsyncWrite + Send + Unpin>> {
         match self {
-            RecordsIO::File(path) => {
+            RecordsOut::File(path) => {
                 trace!(?path, "opening file writer");
                 let file = OpenOptions::new()
                     .write(true)
@@ -251,26 +270,25 @@ impl RecordsIO {
 
                 Ok(Box::new(BufWriter::new(file)))
             }
-            RecordsIO::Stdout => {
+            RecordsOut::Stdout => {
                 trace!("stdout writer");
                 Ok(Box::new(BufWriter::new(tokio::io::stdout())))
             }
-            RecordsIO::Stdin => panic!("unsupported record source"),
         }
     }
 }
 
-fn parse_records_input_source(s: &str) -> Result<RecordsIO, std::io::Error> {
+fn parse_records_input_source(s: &str) -> Result<RecordsIn, std::io::Error> {
     match s {
-        "-" => Ok(RecordsIO::Stdin),
-        _ => Ok(RecordsIO::File(PathBuf::from(s))),
+        "" | "-" => Ok(RecordsIn::Stdin),
+        _ => Ok(RecordsIn::File(PathBuf::from(s))),
     }
 }
 
-fn parse_records_output_source(s: &str) -> Result<RecordsIO, std::io::Error> {
+fn parse_records_output_source(s: &str) -> Result<RecordsOut, std::io::Error> {
     match s {
-        "-" => Ok(RecordsIO::Stdout),
-        _ => Ok(RecordsIO::File(PathBuf::from(s))),
+        "" | "-" => Ok(RecordsOut::Stdout),
+        _ => Ok(RecordsOut::File(PathBuf::from(s))),
     }
 }
 
@@ -504,12 +522,14 @@ async fn run() -> Result<(), S2CliError> {
                     }
                 }
                 StreamActions::Read {
-                    start_seq_num,
+                    start,
                     output,
+                    limit_count,
+                    limit_bytes,
                 } => {
                     let stream_client = StreamClient::new(client_config, basin, stream);
                     let mut read_output_stream = StreamService::new(stream_client)
-                        .read_session(start_seq_num)
+                        .read_session(start, limit_count, limit_bytes)
                         .await?;
                     let mut writer = match output {
                         Some(output) => Some(output.into_writer().await.unwrap()),
@@ -540,16 +560,6 @@ async fn run() -> Result<(), S2CliError> {
                                     _ => panic!("empty batch"),
                                 };
                                 for sequenced_record in sequenced_record_batch.records {
-                                    eprintln!(
-                                        "{}",
-                                        format!(
-                                            "✓ [READ] got record batch: seq_num: {}",
-                                            sequenced_record.seq_num,
-                                        )
-                                        .green()
-                                        .bold()
-                                    );
-
                                     let data = &sequenced_record.body;
                                     batch_len += sequenced_record.metered_size();
 
@@ -574,7 +584,7 @@ async fn run() -> Result<(), S2CliError> {
                                 eprintln!(
                                     "{}",
                                     format!(
-                                        "{throughput_mibps:.2} MiB/s \
+                                        "⦿ {throughput_mibps:.2} MiB/s \
                                             ({num_records} records in range {seq_range:?})",
                                     )
                                     .blue()

src/stream.rs

@@ -1,7 +1,7 @@
 use streamstore::{
     batching::AppendRecordsBatchingStream,
     client::StreamClient,
-    types::{AppendOutput, ReadOutput, ReadSessionRequest},
+    types::{AppendOutput, ReadLimit, ReadOutput, ReadSessionRequest},
     Streaming,
 };
 use tokio::io::AsyncBufRead;
@@ -14,6 +14,7 @@ use tokio::io::Lines;
 use tokio_stream::Stream;
 
 use crate::error::s2_status;
+use crate::ByteSize;
 
 pin_project! {
     #[derive(Debug)]
@@ -89,12 +90,20 @@ impl StreamService {
 
     pub async fn read_session(
         &self,
-        start_seq_num: Option<u64>,
+        start: u64,
+        limit_count: Option<u64>,
+        limit_bytes: Option<ByteSize>,
     ) -> Result<Streaming<ReadOutput>, StreamServiceError> {
-        let mut read_session_req = ReadSessionRequest::new();
-        if let Some(start_seq_num) = start_seq_num {
-            read_session_req = read_session_req.with_start_seq_num(start_seq_num);
-        }
+        let read_session_req = ReadSessionRequest {
+            start_seq_num: Some(start),
+            limit: match (limit_count, limit_bytes.map(|b| b.as_u64())) {
+                (Some(count), Some(bytes)) => Some(ReadLimit { count, bytes }),
+                (Some(count), None) => Some(ReadLimit { count, bytes: 0 }),
+                (None, Some(bytes)) => Some(ReadLimit { count: 0, bytes }),
+                (None, None) => None,
+            },
+        };
+
         self.client
             .read_session(read_session_req)
             .await