Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: handle common signals for streams #32

Merged
merged 4 commits into from
Nov 23, 2024
Merged

fix: handle common signals for streams #32

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
8255f87
feat: handle common signals for streams
infiniteregrets Nov 22, 2024
84cc813
..
infiniteregrets Nov 22, 2024
72b80b3
.
infiniteregrets Nov 22, 2024
f4fdeda
.
infiniteregrets Nov 23, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 24 additions & 0 deletions Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 2 additions & 0 deletions Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@ pin-project-lite = "0.2"
pin-utils = "0.1.0"
serde = { version = "1.0.214", features = ["derive"] }
serde_json = "1.0.132"
signal-hook = "0.3.17"
signal-hook-tokio = { version = "0.3.1", features = ["futures-v0_3"] }
streamstore = { git = "https://github.com/s2-streamstore/s2-sdk-rust.git", rev = "63b4964b66503f705e7c73ae07ba47f81019b79a" }
thiserror = "1.0.67"
tokio = { version = "*", features = ["full"] }
Expand Down
198 changes: 123 additions & 75 deletions src/main.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,8 @@ use clap::{builder::styling, Parser, Subcommand};
use colored::*;
use config::{config_path, create_config};
use error::{S2CliError, ServiceError, ServiceErrorContext};
use signal_hook::consts::{SIGINT, SIGTERM, SIGTSTP};
use signal_hook_tokio::Signals;
use stream::{RecordStream, StreamService};
use streamstore::{
bytesize::ByteSize,
Expand All @@ -16,6 +18,7 @@ use streamstore::{
use tokio::{
fs::{File, OpenOptions},
io::{self, AsyncBufRead, AsyncBufReadExt, AsyncWrite, AsyncWriteExt, BufReader, BufWriter},
select,
time::Instant,
};
use tokio_stream::StreamExt;
Expand Down Expand Up @@ -501,27 +504,51 @@ async fn run() -> Result<(), S2CliError> {
.lines(),
);

let mut signals =
Signals::new([SIGTSTP, SIGINT, SIGTERM]).expect("valid signals");

let mut append_output_stream = StreamService::new(stream_client)
.append_session(append_input_stream)
.await?;
while let Some(append_result) = append_output_stream.next().await {
append_result
.map(|append_result| {
eprintln!(
"{}",
format!(
"✓ [APPENDED] start: {}, end: {}, next: {}",
append_result.start_seq_num,
append_result.end_seq_num,
append_result.next_seq_num
)
.green()
.bold()
);
})
.map_err(|e| {
ServiceError::new(ServiceErrorContext::AppendSession, e)
})?;
loop {
select! {
maybe_append_result = append_output_stream.next() => {
match maybe_append_result {
Some(append_result) => {
match append_result {
Ok(append_result) => {
eprintln!(
"{}",
format!(
"✓ [APPENDED] start: {}, end: {}, next: {}",
append_result.start_seq_num,
append_result.end_seq_num,
append_result.next_seq_num
)
.green()
.bold()
);
},
Err(e) => {
return Err(ServiceError::new(ServiceErrorContext::AppendSession, e).into());
}
}
}
None => break,
}
}

Some(signal) = signals.next() => {
match signal {
SIGTSTP | SIGINT | SIGTERM => {
drop(append_output_stream);
eprintln!("{}", "■ [ABORTED]".red().bold());
break;
}
_ => {}
}
}
}
}
}
StreamActions::Read {
Expand All @@ -531,6 +558,8 @@ async fn run() -> Result<(), S2CliError> {
limit_bytes,
} => {
let stream_client = StreamClient::new(client_config, basin, stream);
let mut signals =
Signals::new([SIGTSTP, SIGINT, SIGTERM]).expect("valid signals");
let mut read_output_stream = StreamService::new(stream_client)
.read_session(start_seq_num, limit_count, limit_bytes)
.await?;
Expand All @@ -539,65 +568,84 @@ async fn run() -> Result<(), S2CliError> {
let mut start = None;
let mut total_data_len = ByteSize::b(0);

while let Some(read_result) = read_output_stream.next().await {
if start.is_none() {
start = Some(Instant::now());
}

let read_result = read_result
.map_err(|e| ServiceError::new(ServiceErrorContext::ReadSession, e))?;

match read_result {
ReadOutput::Batch(sequenced_record_batch) => {
let num_records = sequenced_record_batch.records.len();
let mut batch_len = ByteSize::b(0);

let seq_range = match (
sequenced_record_batch.records.first(),
sequenced_record_batch.records.last(),
) {
(Some(first), Some(last)) => first.seq_num..=last.seq_num,
_ => panic!("empty batch"),
};
for sequenced_record in sequenced_record_batch.records {
let data = &sequenced_record.body;
batch_len += sequenced_record.metered_size();

writer
.write_all(data)
.await
.map_err(|e| S2CliError::RecordWrite(e.to_string()))?;
writer
.write_all(b"\n")
.await
.map_err(|e| S2CliError::RecordWrite(e.to_string()))?;
loop {
select! {
maybe_read_result = read_output_stream.next() => {
match maybe_read_result {
Some(read_result) => {
if start.is_none() {
start = Some(Instant::now());
}
match read_result {
Ok(ReadOutput::Batch(sequenced_record_batch)) => {
let num_records = sequenced_record_batch.records.len();
let mut batch_len = ByteSize::b(0);

let seq_range = match (
sequenced_record_batch.records.first(),
sequenced_record_batch.records.last(),
) {
(Some(first), Some(last)) => first.seq_num..=last.seq_num,
_ => panic!("empty batch"),
};
for sequenced_record in sequenced_record_batch.records {
let data = &sequenced_record.body;
batch_len += sequenced_record.metered_size();

writer
.write_all(data)
.await
.map_err(|e| S2CliError::RecordWrite(e.to_string()))?;
writer
.write_all(b"\n")
.await
.map_err(|e| S2CliError::RecordWrite(e.to_string()))?;
}
total_data_len += batch_len;

let throughput_mibps = (total_data_len.0 as f64
/ start.unwrap().elapsed().as_secs_f64())
/ 1024.0
/ 1024.0;

eprintln!(
"{}",
format!(
"⦿ {throughput_mibps:.2} MiB/s \
({num_records} records in range {seq_range:?})",
)
.blue()
.bold()
);
}

Ok(ReadOutput::FirstSeqNum(seq_num)) => {
eprintln!("{}", format!("first_seq_num: {seq_num}").blue().bold());
}

Ok(ReadOutput::NextSeqNum(seq_num)) => {
eprintln!("{}", format!("next_seq_num: {seq_num}").blue().bold());
}

Err(e) => {
return Err(ServiceError::new(ServiceErrorContext::ReadSession, e).into());
}
}
}
None => break,
}
},
Some(signal) = signals.next() => {
match signal {
SIGTSTP | SIGINT | SIGTERM => {
drop(read_output_stream);
eprintln!("{}", "■ [ABORTED]".red().bold());
break;
}
_ => {}
}
total_data_len += batch_len;

let throughput_mibps = (total_data_len.0 as f64
/ start.unwrap().elapsed().as_secs_f64())
/ 1024.0
/ 1024.0;

eprintln!(
"{}",
format!(
"⦿ {throughput_mibps:.2} MiB/s \
({num_records} records in range {seq_range:?})",
)
.blue()
.bold()
);
}
// TODO: better message for these cases
ReadOutput::FirstSeqNum(seq_num) => {
eprintln!("{}", format!("first_seq_num: {seq_num}").blue().bold());
}
ReadOutput::NextSeqNum(seq_num) => {
eprintln!("{}", format!("next_seq_num: {seq_num}").blue().bold());
}
}

let total_elapsed_time = start.unwrap().elapsed().as_secs_f64();

let total_throughput_mibps =
Expand All @@ -620,5 +668,5 @@ async fn run() -> Result<(), S2CliError> {
}
}
}
Ok(())
std::process::exit(0);
}