Transition to async I/O for most things

Mixing async and sync I/O leads to a multitude of complications, and
generally speaking it's much more convenient to stick to one paradigm
or the other. Since axum (and many other HTTP servers) use async,
and since async is a convenient model for performing operations like
"handle an MPSC message or file read, whichever happens first", let's
commit to an async interface.
This commit is contained in:
Will Greenberg
2024-02-13 20:23:02 -08:00
parent abd3b98cff
commit 775cbcda1e
14 changed files with 550 additions and 629 deletions

View File

@@ -1,26 +1,24 @@
use crate::ServerState;
use rayhunter::diag::DataType;
use rayhunter::gsmtap_parser::GsmtapParser;
use rayhunter::pcap::GsmtapPcapWriter;
use rayhunter::qmdl::{QmdlReader, QmdlReaderError};
use rayhunter::diag_reader::DiagReader;
use rayhunter::qmdl::QmdlReader;
use axum::body::Body;
use axum::http::header::CONTENT_TYPE;
use axum::extract::{State, Path};
use axum::http::StatusCode;
use axum::response::{Response, IntoResponse};
use std::io::Write;
use std::pin::Pin;
use tokio::io::duplex;
use tokio_util::io::ReaderStream;
use std::{future, pin::pin};
use std::sync::Arc;
use std::task::{Poll, Context};
use futures_core::Stream;
use log::error;
use tokio::sync::mpsc;
use futures::TryStreamExt;
// Streams a pcap file chunk-by-chunk to the client by reading the QMDL data
// written so far. This is done by spawning a blocking thread (a tokio thread
// capable of handling blocking operations) which streams chunks of pcap data to
// a channel that's piped to the client.
// written so far. This is done by spawning a thread which streams chunks of
// pcap data to a channel that's piped to the client.
pub async fn get_pcap(State(state): State<Arc<ServerState>>, Path(qmdl_name): Path<String>) -> Result<Response, (StatusCode, String)> {
let qmdl_store = state.qmdl_store_lock.read().await;
let entry = qmdl_store.entry_for_name(&qmdl_name)
@@ -31,82 +29,39 @@ pub async fn get_pcap(State(state): State<Arc<ServerState>>, Path(qmdl_name): Pa
"QMDL file is empty, try again in a bit!".to_string()
));
}
let qmdl_file = qmdl_store.open_entry(&entry).await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("{:?}", e)))?
.into_std().await;
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("{:?}", e)))?;
// the QMDL reader should stop at the last successfully written data chunk
// (entry.size_bytes)
let mut gsmtap_parser = GsmtapParser::new();
let (reader, writer) = duplex(1024);
let mut pcap_writer = GsmtapPcapWriter::new(writer).await.unwrap();
pcap_writer.write_iface_header().await.unwrap();
let (tx, rx) = mpsc::channel(1);
let channel_reader = ChannelReader { rx };
let channel_writer = ChannelWriter { tx };
tokio::task::spawn_blocking(move || {
// the QMDL reader should stop at the last successfully written data
// chunk (qmdl_bytes_written)
let mut qmdl_reader = QmdlReader::new(qmdl_file, Some(entry.size_bytes));
let mut gsmtap_parser = GsmtapParser::new();
let mut pcap_writer = GsmtapPcapWriter::new(channel_writer).unwrap();
pcap_writer.write_iface_header().unwrap();
loop {
match qmdl_reader.read_response() {
Ok(messages) => {
for maybe_msg in messages {
match maybe_msg {
Ok(msg) => {
let maybe_gsmtap_msg = gsmtap_parser.recv_message(msg)
.expect("error parsing gsmtap message");
if let Some((timestamp, gsmtap_msg)) = maybe_gsmtap_msg {
pcap_writer.write_gsmtap_message(gsmtap_msg, timestamp)
.expect("error writing pcap packet");
}
},
Err(e) => error!("error parsing message: {:?}", e),
tokio::spawn(async move {
let mut reader = QmdlReader::new(qmdl_file, Some(entry.size_bytes));
let mut messages_stream = pin!(reader.as_stream()
.try_filter(|container| future::ready(container.data_type == DataType::UserSpace)));
while let Some(container) = messages_stream.try_next().await.expect("failed getting QMDL container") {
for maybe_msg in container.into_messages() {
match maybe_msg {
Ok(msg) => {
let maybe_gsmtap_msg = gsmtap_parser.recv_message(msg)
.expect("error parsing gsmtap message");
if let Some((timestamp, gsmtap_msg)) = maybe_gsmtap_msg {
pcap_writer.write_gsmtap_message(gsmtap_msg, timestamp).await
.expect("error writing pcap packet");
}
}
},
// this is expected, and just means we've reached the end of the
// safely written QMDL data
Err(QmdlReaderError::MaxBytesReached(_)) => break,
Err(e) => {
error!("error reading qmdl file: {:?}", e);
break;
},
},
Err(e) => error!("error parsing message: {:?}", e),
}
}
}
});
let headers = [(CONTENT_TYPE, "application/vnd.tcpdump.pcap")];
let body = Body::from_stream(channel_reader);
let body = Body::from_stream(ReaderStream::new(reader));
Ok((headers, body).into_response())
}
struct ChannelWriter {
tx: mpsc::Sender<Vec<u8>>,
}
impl Write for ChannelWriter {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
self.tx.blocking_send(buf.to_vec())
.map_err(|_| std::io::Error::new(std::io::ErrorKind::Other, "channel closed"))?;
Ok(buf.len())
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
struct ChannelReader {
rx: mpsc::Receiver<Vec<u8>>,
}
impl Stream for ChannelReader {
type Item = Result<Vec<u8>, String>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
match self.rx.poll_recv(cx) {
Poll::Ready(Some(msg)) => Poll::Ready(Some(Ok(msg))),
Poll::Ready(None) => Poll::Ready(None),
Poll::Pending => Poll::Pending,
}
}
}