Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
77 changes: 77 additions & 0 deletions rust/cubestore/cubestore-sql-tests/src/files.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,14 @@ use cubestore::CubeError;
use flate2::read::GzDecoder;
use std::io::Cursor;
use std::io::Write;
use std::net::SocketAddr;
use std::path::Path;
use std::time::Duration;
use tar::Archive;
use tempfile::NamedTempFile;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio::task::JoinHandle;

pub fn write_tmp_file(text: &str) -> Result<NamedTempFile, CubeError> {
let mut file = NamedTempFile::new()?;
Expand Down Expand Up @@ -55,3 +60,75 @@ pub fn recursive_copy_directory(from: &Path, to: &Path) -> Result<(), CubeError>

Ok(())
}

/// Serves one in-memory file over HTTP, for as long as this value is alive.
pub struct TestFileServer {
addr: SocketAddr,
task: JoinHandle<()>,
}

impl TestFileServer {
pub fn url(&self, name: &str) -> String {
format!("http://{}/{}", self.addr, name)
}
}

impl Drop for TestFileServer {
fn drop(&mut self) {
self.task.abort();
}
}

/// `download_delay` holds back the body, so a caller can observe the state a
/// location is in while its download is still running. Head requests are
/// answered immediately.
pub async fn serve_file(
body: String,
download_delay: Duration,
) -> Result<TestFileServer, CubeError> {
let listener = TcpListener::bind("127.0.0.1:0").await?;
let addr = listener.local_addr()?;
let task = tokio::spawn(async move {
while let Ok((socket, _)) = listener.accept().await {
let body = body.clone();
tokio::spawn(async move {
if let Err(e) = serve_one_request(socket, &body, download_delay).await {
log::error!("Test file server: {}", e);

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Two small things about the accept loop, neither blocking:

  • log::error! here fires on a benign disconnect too — Drop aborts only the accept task, so an in-flight serve_one_request whose peer went away (the test finished, the importer dropped the connection after a retry) writes to a closed socket and logs an error. Given what this PR is about, debug! seems more consistent; nothing an operator or a CI reader can act on.
  • while let Ok(..) exits the loop on the first accept error, so a transient failure silently turns the server into a black hole and the test then fails as an import timeout rather than as "the server stopped". Errcontinue (or a break with a log) would make that legible.

Fix this →

}
});
}
});
Ok(TestFileServer { addr, task })
}

async fn serve_one_request(
mut socket: TcpStream,
body: &str,
download_delay: Duration,
) -> Result<(), CubeError> {
let mut request = Vec::new();
let mut buf = [0u8; 1024];
while !request.windows(4).any(|w| w == b"\r\n\r\n") {
let read = socket.read(&mut buf).await?;
if read == 0 {
return Ok(());
}
request.extend_from_slice(&buf[..read]);
}

let mut response = format!(
"HTTP/1.1 200 OK\r\n\
Content-Type: text/csv\r\n\
Content-Length: {}\r\n\
Connection: close\r\n\r\n",
body.len()
);
if !request.starts_with(b"HEAD ") {
tokio::time::sleep(download_delay).await;
response += body;
}

socket.write_all(response.as_bytes()).await?;
socket.shutdown().await?;
Ok(())
}
13 changes: 10 additions & 3 deletions rust/cubestore/cubestore-sql-tests/src/tests.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use crate::files::write_tmp_file;
use crate::files::{serve_file, write_tmp_file};
use crate::rows::{rows, NULL};
use crate::SqlClient;
use async_compression::tokio::write::GzipEncoder;
Expand Down Expand Up @@ -2469,8 +2469,15 @@ async fn create_table_with_csv_no_header_and_quotes(
}

async fn create_table_with_url(service: Box<dyn SqlClient>) -> Result<(), CubeError> {
// TODO serve this data ourselves
let url = "https://data.wprdc.org/dataset/0b584c84-7e35-4f4d-a5a2-b01697470c0f/resource/e95dd941-8e47-4460-9bd8-1e51c194370b/download/bikepghpublic.csv";
let mut csv = "Response ID,Start Date,End Date\n".to_string();
for id in 0..813 {
csv += &format!("{},2020-01-01T00:00:00.000Z,2020-01-02T00:00:00.000Z\n", id);
}
// The body is held back because the query below has to run against a table
// whose import has not finished: only ready tables are visible to the
// planner.
let server = serve_file(csv, Duration::from_millis(500)).await?;

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Serving the file ourselves is the right fix, and the assertion is more robust than it looks — if the SELECT is planned before CREATE TABLE registers the table it fails with "not found" rather than "not ready", so either ordering satisfies is_err(). The one way it can flip is the SELECT future not being polled to completion within the 500 ms window on a loaded runner, at which point the import has finished and the table is queryable.

500 ms is probably enough in practice, but the window is a magic number that only the prose above explains. Releasing the body on a signal instead would remove the timing dependency entirely — have serve_file take a oneshot::Receiver, and fire the sender in the SELECT branch of the join! once the query has returned. Failing that, a named constant would at least say what the number is protecting.

Not blocking — the failure mode is a re-run, not a wrong result.

Fix this →

let url = server.url("bikepghpublic.csv");

service
.exec_query("CREATE SCHEMA IF NOT EXISTS foo")
Expand Down
65 changes: 57 additions & 8 deletions rust/cubestore/cubestore/src/http/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -30,10 +30,12 @@ use http_auth_basic::Credentials;
use log::error;
use log::info;
use log::trace;
use log::Level;
use serde::Deserialize;
use std::collections::{BTreeMap, HashMap};
use std::convert::TryFrom;
use std::error::Error as StdError;
use std::io;
use std::net::SocketAddr;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, SystemTime};
Expand All @@ -43,6 +45,7 @@ use tokio::io::{AsyncWriteExt, BufReader};
use tokio::sync::mpsc::Sender;
use tokio::sync::{mpsc, Mutex};
use tokio_tungstenite::tungstenite;
use tokio_tungstenite::tungstenite::error::ProtocolError;
use tokio_util::sync::CancellationToken;
use warp::filters::ws::{Message, Ws};
use warp::http::StatusCode;
Expand Down Expand Up @@ -149,6 +152,49 @@ fn message_too_large_reason(
}
}

/// Level at which a websocket transport error is reported, in either
/// direction. A peer that vanished is not actionable, and a fleet of them going
/// away at once would otherwise bury the transport failures that are. `warp`
/// boxes the `tungstenite` error, so it comes back through `source()`; anything
/// else is not known to be benign and stays at `Error`.
fn websocket_error_level(e: &warp::Error) -> Level {
match e
.source()
.and_then(|s| s.downcast_ref::<tungstenite::Error>())
{
Some(e) => tungstenite_error_level(e),
None => Level::Error,
}
}

fn tungstenite_error_level(e: &tungstenite::Error) -> Level {
match e {
// How a client normally goes away: `terminate()`, a killed process, a
// closed browser tab. Nothing else about it is ever known.
tungstenite::Error::Protocol(ProtocolError::ResetWithoutClosingHandshake) => Level::Debug,
// Reading from, or writing to, a peer that has already closed.
tungstenite::Error::ConnectionClosed | tungstenite::Error::AlreadyClosed => Level::Debug,
tungstenite::Error::Io(io) if is_peer_gone(io.kind()) => Level::Debug,
// A frame that raced the peer's own close frame: harmless once, a
// client that closes incorrectly if it keeps happening.
tungstenite::Error::Protocol(ProtocolError::ReceivedAfterClosing) => Level::Warn,
_ => Level::Error,
}
}

/// Whether an IO error means the peer is simply gone, as opposed to the
/// connection failing while the peer is still there.
fn is_peer_gone(kind: io::ErrorKind) -> bool {
matches!(
kind,
io::ErrorKind::ConnectionReset
| io::ErrorKind::ConnectionAborted
| io::ErrorKind::BrokenPipe
| io::ErrorKind::NotConnected
| io::ErrorKind::UnexpectedEof
)
}

pub struct HttpServer {
bind_address: String,
sql_service: Arc<dyn SqlService>,
Expand Down Expand Up @@ -459,7 +505,7 @@ impl HttpServer {
));
match tokio::time::timeout(EVICTED_CLOSE_TIMEOUT, close).await {
Ok(Ok(())) => {}
Ok(Err(e)) => error!("Websocket close send error: {:?}", e),
Ok(Err(e)) => log::log!(websocket_error_level(&e), "Websocket close send error: {:?}", e),
Err(_) => log::warn!(
"Timed out sending the close frame of an evicted websocket connection"
),
Expand All @@ -470,7 +516,7 @@ impl HttpServer {
trace!("Sending web socket response (process_id: {})", process_id);
let send_res = web_socket.send(Message::binary(res.bytes())).await;
if let Err(e) = send_res {
error!("Websocket message send error: {:?}", e)
log::log!(websocket_error_level(&e), "Websocket message send error: {:?}", e)
}
if res.should_close_connection() {
log::warn!("Websocket connection closed");
Expand All @@ -496,10 +542,13 @@ impl HttpServer {
Message::close_with(MESSAGE_TOO_BIG_CLOSE_CODE, reason)
).await;
if let Err(e) = send_res {
error!("Websocket close send error: {:?}", e)
log::log!(websocket_error_level(&e), "Websocket close send error: {:?}", e)
}
}
None => error!("Websocket error: {:?}", e),
None => log::log!(
websocket_error_level(&e),
Comment thread
claude[bot] marked this conversation as resolved.
"Websocket error: {:?}", e
),
}
break;
}
Expand Down Expand Up @@ -531,7 +580,7 @@ impl HttpServer {
Message::binary(HttpMessage { message_id, connection_id, command: HttpCommand::Error { error } }.bytes())
).await;
if let Err(e) = send_res {
error!("Websocket message send error: {:?}", e)
log::log!(websocket_error_level(&e), "Websocket message send error: {:?}", e)
}
continue;
}
Expand All @@ -544,7 +593,7 @@ impl HttpServer {
Message::binary(HttpMessage { message_id, connection_id, command: HttpCommand::Error { error: e.to_string() } }.bytes())
).await;
if let Err(e) = send_res {
error!("Websocket message send error: {:?}", e)
log::log!(websocket_error_level(&e), "Websocket message send error: {:?}", e)
}
break;
},
Expand All @@ -557,7 +606,7 @@ impl HttpServer {
Message::binary(HttpMessage { message_id, connection_id, command: HttpCommand::Error { error: e.to_string() } }.bytes())
).await;
if let Err(e) = send_res {
error!("Websocket message send error: {:?}", e)
log::log!(websocket_error_level(&e), "Websocket message send error: {:?}", e)
}
break;
}
Expand All @@ -566,7 +615,7 @@ impl HttpServer {
} else if msg.is_ping() {
let send_res = web_socket.send(Message::pong(Vec::new())).await;
if let Err(e) = send_res {
error!("Websocket ping send error: {:?}", e)
log::log!(websocket_error_level(&e), "Websocket ping send error: {:?}", e)
}
} else if msg.is_close() {
break;
Expand Down
15 changes: 13 additions & 2 deletions rust/cubestore/cubestore/src/sql/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4878,13 +4878,24 @@ mod tests {
config.max_partition_split_threshold = 200;
config
}).start_test_worker(async move |_| {
let url = "https://data.wprdc.org/dataset/0b584c84-7e35-4f4d-a5a2-b01697470c0f/resource/e95dd941-8e47-4460-9bd8-1e51c194370b/download/bikepghpublic.csv";
// The threshold is derived from the size of the location,
// not from its contents: see ImportServiceImpl::estimate_rows.
// Rows are padded so the estimate lands above
// max_partition_split_threshold per select worker, which is
// what makes the cap observable.
let path = env::temp_dir().join(format!("{}.csv", test_name));
let padding = "x".repeat(256);
let mut csv = "Response ID,Start Date,End Date\n".to_string();
for id in 0..813 {
csv += &format!("{},2020-01-01T00:00:00.000Z,{}\n", id, padding);
}
Comment on lines +4881 to +4891

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The generated file is the right call, and the arithmetic checks out (232,440 bytes → 907 estimated rows → 453 per select worker, comfortably above the 200 cap). Two small things about how that arithmetic is expressed:

The comment is 5 lines for 6 lines of code, and its second half restates what the constants below already do. The load-bearing fact is the first sentence — estimate_rows never opens the file, only stats it — because that is what makes a synthetic file valid here at all. The rest would read better as names on the numbers: 256 is not an arbitrary pad width, it is exactly estimate_rows's average_row_length, so one padded row ≈ one estimated row, and 813 is the row count the count(*) assert below pins. Something like:

// estimate_rows only stats the location, it never reads it.
let rows = 813;
// One padded row is about one estimated row: estimate_rows divides by 256.
let padding = "x".repeat(256);

makes the "813 rows / 2 workers > 200" chain visible without prose. Not blocking either way.

Fix this →

tokio::fs::write(&path, csv).await?;

service
.exec_query("CREATE SCHEMA IF NOT EXISTS foo")
.await?.collect().await?;

let create_table_sql = format!("CREATE TABLE foo.bikes (`Response ID` int, `Start Date` text, `End Date` text) LOCATION '{}'", url);
let create_table_sql = format!("CREATE TABLE foo.bikes (`Response ID` int, `Start Date` text, `End Date` text) LOCATION '{}'", path.to_string_lossy());

service.exec_query(&create_table_sql).await?.collect().await?;

Expand Down
Loading