diff --git a/rust/cubestore/cubestore-sql-tests/src/files.rs b/rust/cubestore/cubestore-sql-tests/src/files.rs index e42cc6a7ca2dd..cb36708b4369c 100644 --- a/rust/cubestore/cubestore-sql-tests/src/files.rs +++ b/rust/cubestore/cubestore-sql-tests/src/files.rs @@ -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 { let mut file = NamedTempFile::new()?; @@ -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 { + 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); + } + }); + } + }); + 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(()) +} diff --git a/rust/cubestore/cubestore-sql-tests/src/tests.rs b/rust/cubestore/cubestore-sql-tests/src/tests.rs index 205b1420fc3ca..edfa97d66dc09 100644 --- a/rust/cubestore/cubestore-sql-tests/src/tests.rs +++ b/rust/cubestore/cubestore-sql-tests/src/tests.rs @@ -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; @@ -2469,8 +2469,15 @@ async fn create_table_with_csv_no_header_and_quotes( } async fn create_table_with_url(service: Box) -> 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?; + let url = server.url("bikepghpublic.csv"); service .exec_query("CREATE SCHEMA IF NOT EXISTS foo") diff --git a/rust/cubestore/cubestore/src/http/mod.rs b/rust/cubestore/cubestore/src/http/mod.rs index 2f2893d0f8a1d..5c6ec31df3a72 100644 --- a/rust/cubestore/cubestore/src/http/mod.rs +++ b/rust/cubestore/cubestore/src/http/mod.rs @@ -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}; @@ -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; @@ -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::()) + { + 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, @@ -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" ), @@ -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"); @@ -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), + "Websocket error: {:?}", e + ), } break; } @@ -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; } @@ -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; }, @@ -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; } @@ -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; diff --git a/rust/cubestore/cubestore/src/sql/mod.rs b/rust/cubestore/cubestore/src/sql/mod.rs index 19a0dbe0533d9..21fe0f6252adc 100644 --- a/rust/cubestore/cubestore/src/sql/mod.rs +++ b/rust/cubestore/cubestore/src/sql/mod.rs @@ -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); + } + 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?;