Line data Source code
1 : use std::sync::Arc;
2 : use std::time::SystemTime;
3 :
4 : use anyhow::Context;
5 : use bytes::buf::Writer;
6 : use bytes::{BufMut, BytesMut};
7 : use chrono::{Datelike, Timelike};
8 : use futures::{Stream, StreamExt};
9 : use parquet::basic::Compression;
10 : use parquet::file::metadata::RowGroupMetaDataPtr;
11 : use parquet::file::properties::{DEFAULT_PAGE_SIZE, WriterProperties, WriterPropertiesPtr};
12 : use parquet::file::writer::SerializedFileWriter;
13 : use parquet::record::RecordWriter;
14 : use remote_storage::{GenericRemoteStorage, RemotePath, RemoteStorageConfig, TimeoutOrCancel};
15 : use serde::ser::SerializeMap;
16 : use tokio::sync::mpsc;
17 : use tokio::time;
18 : use tokio_util::sync::CancellationToken;
19 : use tracing::{Span, debug, info};
20 : use utils::backoff;
21 :
22 : use super::{LOG_CHAN, RequestContextInner};
23 : use crate::config::remote_storage_from_toml;
24 : use crate::context::LOG_CHAN_DISCONNECT;
25 : use crate::ext::TaskExt;
26 : use crate::pqproto::StartupMessageParams;
27 :
28 : #[derive(clap::Args, Clone, Debug)]
29 : pub struct ParquetUploadArgs {
30 : /// Storage location to upload the parquet files to.
31 : /// Encoded as toml (same format as pageservers), eg
32 : /// `{bucket_name='the-bucket',bucket_region='us-east-1',prefix_in_bucket='proxy',endpoint='http://minio:9000'}`
33 : #[clap(long, value_parser = remote_storage_from_toml)]
34 : parquet_upload_remote_storage: Option<RemoteStorageConfig>,
35 :
36 : #[clap(long, value_parser = remote_storage_from_toml)]
37 : parquet_upload_disconnect_events_remote_storage: Option<RemoteStorageConfig>,
38 :
39 : /// How many rows to include in a row group
40 : #[clap(long, default_value_t = 8192)]
41 : parquet_upload_row_group_size: usize,
42 :
43 : /// How large each column page should be in bytes
44 : #[clap(long, default_value_t = DEFAULT_PAGE_SIZE)]
45 : parquet_upload_page_size: usize,
46 :
47 : /// How large the total parquet file should be in bytes
48 : #[clap(long, default_value_t = 100_000_000)]
49 : parquet_upload_size: i64,
50 :
51 : /// How long to wait before forcing a file upload
52 : #[clap(long, default_value = "20m", value_parser = humantime::parse_duration)]
53 : parquet_upload_maximum_duration: tokio::time::Duration,
54 :
55 : /// What level of compression to use
56 : #[clap(long, default_value_t = Compression::UNCOMPRESSED)]
57 : parquet_upload_compression: Compression,
58 : }
59 :
60 : // Occasional network issues and such can cause remote operations to fail, and
61 : // that's expected. If a upload fails, we log it at info-level, and retry.
62 : // But after FAILED_UPLOAD_WARN_THRESHOLD retries, we start to log it at WARN
63 : // level instead, as repeated failures can mean a more serious problem. If it
64 : // fails more than FAILED_UPLOAD_RETRIES times, we give up
65 : pub(crate) const FAILED_UPLOAD_WARN_THRESHOLD: u32 = 3;
66 : pub(crate) const FAILED_UPLOAD_MAX_RETRIES: u32 = 10;
67 :
68 : // the parquet crate leaves a lot to be desired...
69 : // what follows is an attempt to write parquet files with minimal allocs.
70 : // complication: parquet is a columnar format, while we want to write in as rows.
71 : // design:
72 : // * we batch up to 1024 rows, then flush them into a 'row group'
73 : // * after each rowgroup write, we check the length of the file and upload to s3 if large enough
74 :
75 : #[derive(parquet_derive::ParquetRecordWriter)]
76 : pub(crate) struct RequestData {
77 : region: String,
78 : protocol: &'static str,
79 : /// Must be UTC. The derive macro doesn't like the timezones
80 : timestamp: chrono::NaiveDateTime,
81 : session_id: uuid::Uuid,
82 : peer_addr: String,
83 : username: Option<String>,
84 : application_name: Option<String>,
85 : user_agent: Option<String>,
86 : endpoint_id: Option<String>,
87 : database: Option<String>,
88 : project: Option<String>,
89 : branch: Option<String>,
90 : pg_options: Option<String>,
91 : auth_method: Option<&'static str>,
92 : jwt_issuer: Option<String>,
93 :
94 : error: Option<&'static str>,
95 : /// Success is counted if we form a HTTP response with sql rows inside
96 : /// Or if we make it to proxy_pass
97 : success: bool,
98 : /// Indicates if the cplane started the new compute node for this request.
99 : cold_start_info: &'static str,
100 : /// Tracks time from session start (HTTP request/libpq TCP handshake)
101 : /// Through to success/failure
102 : duration_us: u64,
103 : /// If the session was successful after the disconnect, will be created one more event with filled `disconnect_timestamp`.
104 : disconnect_timestamp: Option<chrono::NaiveDateTime>,
105 : }
106 :
107 : struct Options<'a> {
108 : options: &'a StartupMessageParams,
109 : }
110 :
111 : impl serde::Serialize for Options<'_> {
112 0 : fn serialize<S>(&self, s: S) -> Result<S::Ok, S::Error>
113 0 : where
114 0 : S: serde::Serializer,
115 : {
116 0 : let mut state = s.serialize_map(None)?;
117 0 : for (k, v) in self.options.iter() {
118 0 : state.serialize_entry(k, v)?;
119 : }
120 0 : state.end()
121 0 : }
122 : }
123 :
124 : impl From<&RequestContextInner> for RequestData {
125 0 : fn from(value: &RequestContextInner) -> Self {
126 : Self {
127 0 : session_id: value.session_id,
128 0 : peer_addr: value.conn_info.addr.ip().to_string(),
129 0 : timestamp: value.first_packet.naive_utc(),
130 0 : username: value.user.as_deref().map(String::from),
131 0 : application_name: value.application.as_deref().map(String::from),
132 0 : user_agent: value.user_agent.as_deref().map(String::from),
133 0 : endpoint_id: value.endpoint_id.as_deref().map(String::from),
134 0 : database: value.dbname.as_deref().map(String::from),
135 0 : project: value.project.as_deref().map(String::from),
136 0 : branch: value.branch.as_deref().map(String::from),
137 0 : pg_options: value
138 0 : .pg_options
139 0 : .as_ref()
140 0 : .and_then(|options| serde_json::to_string(&Options { options }).ok()),
141 0 : auth_method: value.auth_method.as_ref().map(|x| match x {
142 0 : super::AuthMethod::ConsoleRedirect => "console_redirect",
143 0 : super::AuthMethod::ScramSha256 => "scram_sha_256",
144 0 : super::AuthMethod::ScramSha256Plus => "scram_sha_256_plus",
145 0 : super::AuthMethod::Cleartext => "cleartext",
146 0 : super::AuthMethod::Jwt => "jwt",
147 0 : }),
148 0 : jwt_issuer: value.jwt_issuer.clone(),
149 0 : protocol: value.protocol.as_str(),
150 0 : region: String::new(),
151 0 : error: value.error_kind.as_ref().map(|e| e.to_metric_label()),
152 0 : success: value.success,
153 0 : cold_start_info: value.cold_start_info.as_str(),
154 0 : duration_us: SystemTime::from(value.first_packet)
155 0 : .elapsed()
156 0 : .unwrap_or_default()
157 0 : .as_micros() as u64, // 584 millenia... good enough
158 0 : disconnect_timestamp: value.disconnect_timestamp.map(|x| x.naive_utc()),
159 : }
160 0 : }
161 : }
162 :
163 : /// Parquet request context worker
164 : ///
165 : /// It listened on a channel for all completed requests, extracts the data and writes it into a parquet file,
166 : /// then uploads a completed batch to S3
167 0 : pub async fn worker(
168 0 : cancellation_token: CancellationToken,
169 0 : config: ParquetUploadArgs,
170 0 : region: String,
171 0 : ) -> anyhow::Result<()> {
172 0 : let Some(remote_storage_config) = config.parquet_upload_remote_storage else {
173 0 : tracing::warn!("parquet request upload: no s3 bucket configured");
174 0 : return Ok(());
175 : };
176 :
177 0 : let (tx, mut rx) = mpsc::unbounded_channel();
178 0 : LOG_CHAN
179 0 : .set(tx.downgrade())
180 0 : .expect("only one worker should set the channel");
181 :
182 : // setup row stream that will close on cancellation
183 0 : let cancellation_token2 = cancellation_token.clone();
184 0 : tokio::spawn(async move {
185 0 : cancellation_token2.cancelled().await;
186 : // dropping this sender will cause the channel to close only once
187 : // all the remaining inflight requests have been completed.
188 0 : drop(tx);
189 0 : });
190 0 : let rx = futures::stream::poll_fn(move |cx| rx.poll_recv(cx));
191 0 : let rx = rx.map(RequestData::from);
192 :
193 0 : let storage = GenericRemoteStorage::from_config(&remote_storage_config)
194 0 : .await
195 0 : .context("remote storage init")?;
196 :
197 0 : let properties = WriterProperties::builder()
198 0 : .set_data_page_size_limit(config.parquet_upload_page_size)
199 0 : .set_compression(config.parquet_upload_compression);
200 :
201 0 : let parquet_config = ParquetConfig {
202 0 : propeties: Arc::new(properties.build()),
203 0 : rows_per_group: config.parquet_upload_row_group_size,
204 0 : file_size: config.parquet_upload_size,
205 0 : max_duration: config.parquet_upload_maximum_duration,
206 0 :
207 0 : #[cfg(any(test, feature = "testing"))]
208 0 : test_remote_failures: 0,
209 0 : };
210 :
211 : // TODO(anna): consider moving this to a separate function.
212 0 : if let Some(disconnect_events_storage_config) =
213 0 : config.parquet_upload_disconnect_events_remote_storage
214 : {
215 0 : let (tx_disconnect, mut rx_disconnect) = mpsc::unbounded_channel();
216 0 : LOG_CHAN_DISCONNECT
217 0 : .set(tx_disconnect.downgrade())
218 0 : .expect("only one worker should set the channel");
219 :
220 : // setup row stream that will close on cancellation
221 0 : tokio::spawn(async move {
222 0 : cancellation_token.cancelled().await;
223 : // dropping this sender will cause the channel to close only once
224 : // all the remaining inflight requests have been completed.
225 0 : drop(tx_disconnect);
226 0 : });
227 0 : let rx_disconnect = futures::stream::poll_fn(move |cx| rx_disconnect.poll_recv(cx));
228 0 : let rx_disconnect = rx_disconnect.map(RequestData::from);
229 :
230 0 : let storage_disconnect =
231 0 : GenericRemoteStorage::from_config(&disconnect_events_storage_config)
232 0 : .await
233 0 : .context("remote storage for disconnect events init")?;
234 0 : let parquet_config_disconnect = parquet_config.clone();
235 0 : tokio::try_join!(
236 0 : worker_inner(storage, rx, parquet_config, ®ion),
237 0 : worker_inner(
238 0 : storage_disconnect,
239 0 : rx_disconnect,
240 0 : parquet_config_disconnect,
241 0 : ®ion
242 : )
243 : )
244 0 : .map(|_| ())
245 : } else {
246 0 : worker_inner(storage, rx, parquet_config, ®ion).await
247 : }
248 0 : }
249 :
250 : #[derive(Clone, Debug)]
251 : struct ParquetConfig {
252 : propeties: WriterPropertiesPtr,
253 : rows_per_group: usize,
254 : file_size: i64,
255 :
256 : max_duration: tokio::time::Duration,
257 :
258 : #[cfg(any(test, feature = "testing"))]
259 : test_remote_failures: u64,
260 : }
261 :
262 4 : async fn worker_inner(
263 4 : storage: GenericRemoteStorage,
264 4 : rx: impl Stream<Item = RequestData>,
265 4 : config: ParquetConfig,
266 4 : region: &str,
267 4 : ) -> anyhow::Result<()> {
268 : #[cfg(any(test, feature = "testing"))]
269 4 : let storage = if config.test_remote_failures > 0 {
270 2 : GenericRemoteStorage::unreliable_wrapper(storage, config.test_remote_failures, 100)
271 : } else {
272 2 : storage
273 : };
274 :
275 4 : let mut rx = std::pin::pin!(rx);
276 :
277 4 : let mut rows = Vec::with_capacity(config.rows_per_group);
278 :
279 4 : let schema = rows.as_slice().schema()?;
280 4 : let buffer = BytesMut::new();
281 4 : let w = buffer.writer();
282 4 : let mut w = SerializedFileWriter::new(w, schema.clone(), config.propeties.clone())?;
283 :
284 4 : let mut last_upload = time::Instant::now();
285 :
286 4 : let mut len = 0;
287 159004 : while let Some(mut row) = rx.next().await {
288 159000 : region.clone_into(&mut row.region);
289 159000 : rows.push(row);
290 159000 : let force = last_upload.elapsed() > config.max_duration;
291 159000 : if rows.len() == config.rows_per_group || force {
292 : let rg_meta;
293 80 : (rows, w, rg_meta) = flush_rows(rows, w).await?;
294 80 : len += rg_meta.compressed_size();
295 158920 : }
296 159000 : if len > config.file_size || force {
297 23 : last_upload = time::Instant::now();
298 23 : let file = upload_parquet(w, len, &storage).await?;
299 23 : w = SerializedFileWriter::new(file, schema.clone(), config.propeties.clone())?;
300 23 : len = 0;
301 158977 : }
302 : }
303 :
304 4 : if !rows.is_empty() {
305 : let rg_meta;
306 1 : (_, w, rg_meta) = flush_rows(rows, w).await?;
307 1 : len += rg_meta.compressed_size();
308 3 : }
309 :
310 4 : if !w.flushed_row_groups().is_empty() {
311 3 : let _rtchk: Writer<BytesMut> = upload_parquet(w, len, &storage).await?;
312 1 : }
313 :
314 4 : Ok(())
315 4 : }
316 :
317 81 : async fn flush_rows<W>(
318 81 : rows: Vec<RequestData>,
319 81 : mut w: SerializedFileWriter<W>,
320 81 : ) -> anyhow::Result<(
321 81 : Vec<RequestData>,
322 81 : SerializedFileWriter<W>,
323 81 : RowGroupMetaDataPtr,
324 81 : )>
325 81 : where
326 81 : W: std::io::Write + Send + 'static,
327 81 : {
328 81 : let span = Span::current();
329 81 : let (mut rows, w, rg_meta) = tokio::task::spawn_blocking(move || {
330 81 : let _enter = span.enter();
331 :
332 81 : let mut rg = w.next_row_group()?;
333 81 : rows.as_slice().write_to_row_group(&mut rg)?;
334 81 : let rg_meta = rg.close()?;
335 :
336 81 : let size = rg_meta.compressed_size();
337 81 : let compression = rg_meta.compressed_size() as f64 / rg_meta.total_byte_size() as f64;
338 :
339 81 : debug!(size, compression, "flushed row group to parquet file");
340 :
341 81 : Ok::<_, parquet::errors::ParquetError>((rows, w, rg_meta))
342 81 : })
343 81 : .await
344 81 : .propagate_task_panic()?;
345 :
346 81 : rows.clear();
347 81 : Ok((rows, w, rg_meta))
348 81 : }
349 :
350 26 : async fn upload_parquet(
351 26 : mut w: SerializedFileWriter<Writer<BytesMut>>,
352 26 : len: i64,
353 26 : storage: &GenericRemoteStorage,
354 26 : ) -> anyhow::Result<Writer<BytesMut>> {
355 26 : let len_uncompressed = w
356 26 : .flushed_row_groups()
357 26 : .iter()
358 81 : .map(|rg| rg.total_byte_size())
359 26 : .sum::<i64>();
360 :
361 : // I don't know how compute intensive this is, although it probably isn't much... better be safe than sorry.
362 : // finish method only available on the fork: https://github.com/apache/arrow-rs/issues/5253
363 26 : let (mut buffer, metadata) =
364 26 : tokio::task::spawn_blocking(move || -> parquet::errors::Result<_> {
365 26 : let metadata = w.finish()?;
366 26 : let buffer = std::mem::take(w.inner_mut().get_mut());
367 26 : Ok((buffer, metadata))
368 26 : })
369 26 : .await
370 26 : .propagate_task_panic()?;
371 :
372 26 : let data = buffer.split().freeze();
373 :
374 26 : let compression = len as f64 / len_uncompressed as f64;
375 26 : let size = data.len();
376 26 : let now = chrono::Utc::now();
377 26 : let id = uuid::Uuid::new_v7(uuid::Timestamp::from_unix(
378 26 : uuid::NoContext,
379 : // we won't be running this in 1970. this cast is ok
380 26 : now.timestamp() as u64,
381 26 : now.timestamp_subsec_nanos(),
382 : ));
383 :
384 26 : info!(
385 : %id,
386 : rows = metadata.num_rows,
387 0 : size, compression, "uploading request parquet file"
388 : );
389 :
390 26 : let year = now.year();
391 26 : let month = now.month();
392 26 : let day = now.day();
393 26 : let hour = now.hour();
394 : // segment files by time for S3 performance
395 26 : let path = RemotePath::from_string(&format!(
396 26 : "{year:04}/{month:02}/{day:02}/{hour:02}/requests_{id}.parquet"
397 26 : ))?;
398 26 : let cancel = CancellationToken::new();
399 26 : let maybe_err = backoff::retry(
400 38 : || async {
401 38 : let stream = futures::stream::once(futures::future::ready(Ok(data.clone())));
402 38 : storage
403 38 : .upload(stream, data.len(), &path, None, &cancel)
404 38 : .await
405 76 : },
406 : TimeoutOrCancel::caused_by_cancel,
407 : FAILED_UPLOAD_WARN_THRESHOLD,
408 : FAILED_UPLOAD_MAX_RETRIES,
409 26 : "request_data_upload",
410 : // we don't want cancellation to interrupt here, so we make a dummy cancel token
411 26 : &cancel,
412 : )
413 26 : .await
414 26 : .ok_or_else(|| anyhow::Error::new(TimeoutOrCancel::Cancel))
415 26 : .and_then(|x| x)
416 26 : .with_context(|| format!("request_data_upload: path={path}"))
417 26 : .err();
418 :
419 26 : if let Some(err) = maybe_err {
420 0 : tracing::error!(%id, %path, error = ?err, "failed to upload request data");
421 26 : }
422 :
423 26 : Ok(buffer.writer())
424 26 : }
425 :
426 : #[cfg(test)]
427 : mod tests {
428 : use std::net::Ipv4Addr;
429 : use std::num::NonZeroUsize;
430 : use std::sync::Arc;
431 :
432 : use camino::Utf8Path;
433 : use clap::Parser;
434 : use futures::{Stream, StreamExt};
435 : use itertools::Itertools;
436 : use parquet::basic::{Compression, ZstdLevel};
437 : use parquet::file::properties::{DEFAULT_PAGE_SIZE, WriterProperties};
438 : use parquet::file::reader::FileReader;
439 : use parquet::file::serialized_reader::SerializedFileReader;
440 : use rand::rngs::StdRng;
441 : use rand::{Rng, SeedableRng};
442 : use remote_storage::{
443 : DEFAULT_MAX_KEYS_PER_LIST_RESPONSE, DEFAULT_REMOTE_STORAGE_S3_CONCURRENCY_LIMIT,
444 : GenericRemoteStorage, RemoteStorageConfig, RemoteStorageKind, S3Config,
445 : };
446 : use tokio::sync::mpsc;
447 : use tokio::time;
448 : use walkdir::WalkDir;
449 :
450 : use super::{ParquetConfig, ParquetUploadArgs, RequestData, worker_inner};
451 :
452 : #[derive(Parser)]
453 : struct ProxyCliArgs {
454 : #[clap(flatten)]
455 : parquet_upload: ParquetUploadArgs,
456 : }
457 :
458 : #[test]
459 1 : fn default_parser() {
460 1 : let ProxyCliArgs { parquet_upload } = ProxyCliArgs::parse_from(["proxy"]);
461 1 : assert_eq!(parquet_upload.parquet_upload_remote_storage, None);
462 1 : assert_eq!(parquet_upload.parquet_upload_row_group_size, 8192);
463 1 : assert_eq!(parquet_upload.parquet_upload_page_size, DEFAULT_PAGE_SIZE);
464 1 : assert_eq!(parquet_upload.parquet_upload_size, 100_000_000);
465 1 : assert_eq!(
466 : parquet_upload.parquet_upload_maximum_duration,
467 1 : time::Duration::from_secs(20 * 60)
468 : );
469 1 : assert_eq!(
470 : parquet_upload.parquet_upload_compression,
471 : Compression::UNCOMPRESSED
472 : );
473 1 : }
474 :
475 : #[test]
476 1 : fn full_parser() {
477 1 : let ProxyCliArgs { parquet_upload } = ProxyCliArgs::parse_from([
478 1 : "proxy",
479 1 : "--parquet-upload-remote-storage",
480 1 : "{bucket_name='default',prefix_in_bucket='proxy/',bucket_region='us-east-1',endpoint='http://minio:9000'}",
481 1 : "--parquet-upload-row-group-size",
482 1 : "100",
483 1 : "--parquet-upload-page-size",
484 1 : "10000",
485 1 : "--parquet-upload-size",
486 1 : "10000000",
487 1 : "--parquet-upload-maximum-duration",
488 1 : "10m",
489 1 : "--parquet-upload-compression",
490 1 : "zstd(5)",
491 1 : ]);
492 1 : assert_eq!(
493 : parquet_upload.parquet_upload_remote_storage,
494 1 : Some(RemoteStorageConfig {
495 1 : storage: RemoteStorageKind::AwsS3(S3Config {
496 1 : bucket_name: "default".into(),
497 1 : bucket_region: "us-east-1".into(),
498 1 : prefix_in_bucket: Some("proxy/".into()),
499 1 : endpoint: Some("http://minio:9000".into()),
500 1 : concurrency_limit: NonZeroUsize::new(
501 1 : DEFAULT_REMOTE_STORAGE_S3_CONCURRENCY_LIMIT
502 1 : )
503 1 : .unwrap(),
504 1 : max_keys_per_list_response: DEFAULT_MAX_KEYS_PER_LIST_RESPONSE,
505 1 : upload_storage_class: None,
506 1 : }),
507 1 : timeout: RemoteStorageConfig::DEFAULT_TIMEOUT,
508 1 : small_timeout: RemoteStorageConfig::DEFAULT_SMALL_TIMEOUT,
509 1 : })
510 : );
511 1 : assert_eq!(parquet_upload.parquet_upload_row_group_size, 100);
512 1 : assert_eq!(parquet_upload.parquet_upload_page_size, 10000);
513 1 : assert_eq!(parquet_upload.parquet_upload_size, 10_000_000);
514 1 : assert_eq!(
515 : parquet_upload.parquet_upload_maximum_duration,
516 1 : time::Duration::from_secs(10 * 60)
517 : );
518 1 : assert_eq!(
519 : parquet_upload.parquet_upload_compression,
520 1 : Compression::ZSTD(ZstdLevel::try_new(5).unwrap())
521 : );
522 1 : }
523 :
524 159000 : fn generate_request_data(rng: &mut impl Rng) -> RequestData {
525 159000 : RequestData {
526 159000 : session_id: uuid::Builder::from_random_bytes(rng.r#gen()).into_uuid(),
527 159000 : peer_addr: Ipv4Addr::from(rng.r#gen::<[u8; 4]>()).to_string(),
528 159000 : timestamp: chrono::DateTime::from_timestamp_millis(
529 159000 : rng.gen_range(1703862754..1803862754),
530 159000 : )
531 159000 : .unwrap()
532 159000 : .naive_utc(),
533 159000 : application_name: Some("test".to_owned()),
534 159000 : user_agent: Some("test-user-agent".to_owned()),
535 159000 : username: Some(hex::encode(rng.r#gen::<[u8; 4]>())),
536 159000 : endpoint_id: Some(hex::encode(rng.r#gen::<[u8; 16]>())),
537 159000 : database: Some(hex::encode(rng.r#gen::<[u8; 16]>())),
538 159000 : project: Some(hex::encode(rng.r#gen::<[u8; 16]>())),
539 159000 : branch: Some(hex::encode(rng.r#gen::<[u8; 16]>())),
540 159000 : pg_options: None,
541 159000 : auth_method: None,
542 159000 : jwt_issuer: None,
543 159000 : protocol: ["tcp", "ws", "http"][rng.gen_range(0..3)],
544 159000 : region: String::new(),
545 159000 : error: None,
546 159000 : success: rng.r#gen(),
547 159000 : cold_start_info: "no",
548 159000 : duration_us: rng.gen_range(0..30_000_000),
549 159000 : disconnect_timestamp: None,
550 159000 : }
551 159000 : }
552 :
553 6 : fn random_stream(len: usize) -> impl Stream<Item = RequestData> + Unpin {
554 6 : let mut rng = StdRng::from_seed([0x39; 32]);
555 6 : futures::stream::iter(
556 159000 : std::iter::repeat_with(move || generate_request_data(&mut rng)).take(len),
557 : )
558 6 : }
559 :
560 4 : async fn run_test(
561 4 : tmpdir: &Utf8Path,
562 4 : config: ParquetConfig,
563 4 : rx: impl Stream<Item = RequestData>,
564 4 : ) -> Vec<(u64, usize, i64)> {
565 4 : let remote_storage_config = RemoteStorageConfig {
566 4 : storage: RemoteStorageKind::LocalFs {
567 4 : local_path: tmpdir.to_path_buf(),
568 4 : },
569 4 : timeout: std::time::Duration::from_secs(120),
570 4 : small_timeout: std::time::Duration::from_secs(30),
571 4 : };
572 4 : let storage = GenericRemoteStorage::from_config(&remote_storage_config)
573 4 : .await
574 4 : .unwrap();
575 :
576 4 : worker_inner(storage, rx, config, "us-east-1")
577 4 : .await
578 4 : .unwrap();
579 :
580 4 : let mut files = WalkDir::new(tmpdir.as_std_path())
581 4 : .into_iter()
582 46 : .filter_map(|entry| entry.ok())
583 46 : .filter(|entry| entry.file_type().is_file())
584 26 : .map(|entry| entry.path().to_path_buf())
585 4 : .collect_vec();
586 4 : files.sort();
587 :
588 4 : files
589 4 : .into_iter()
590 26 : .map(|path| std::fs::File::open(tmpdir.as_std_path().join(path)).unwrap())
591 26 : .map(|file| {
592 26 : (
593 26 : file.metadata().unwrap(),
594 26 : SerializedFileReader::new(file).unwrap().metadata().clone(),
595 26 : )
596 26 : })
597 26 : .map(|(file_meta, parquet_meta)| {
598 26 : (
599 26 : file_meta.len(),
600 26 : parquet_meta.num_row_groups(),
601 26 : parquet_meta.file_metadata().num_rows(),
602 26 : )
603 26 : })
604 4 : .collect()
605 4 : }
606 :
607 : #[tokio::test]
608 1 : async fn verify_parquet_no_compression() {
609 1 : let tmpdir = camino_tempfile::tempdir().unwrap();
610 :
611 1 : let config = ParquetConfig {
612 1 : propeties: Arc::new(WriterProperties::new()),
613 1 : rows_per_group: 2_000,
614 1 : file_size: 1_000_000,
615 1 : max_duration: time::Duration::from_secs(20 * 60),
616 1 : test_remote_failures: 0,
617 1 : };
618 :
619 1 : let rx = random_stream(50_000);
620 1 : let file_stats = run_test(tmpdir.path(), config, rx).await;
621 :
622 1 : assert_eq!(
623 : file_stats,
624 : [
625 : (1313953, 3, 6000),
626 : (1313942, 3, 6000),
627 : (1314001, 3, 6000),
628 : (1313958, 3, 6000),
629 : (1314094, 3, 6000),
630 : (1313931, 3, 6000),
631 : (1313725, 3, 6000),
632 : (1313960, 3, 6000),
633 : (438318, 1, 2000)
634 : ]
635 : );
636 :
637 1 : tmpdir.close().unwrap();
638 1 : }
639 :
640 : #[tokio::test]
641 1 : async fn verify_parquet_strong_compression() {
642 1 : let tmpdir = camino_tempfile::tempdir().unwrap();
643 :
644 1 : let config = ParquetConfig {
645 1 : propeties: Arc::new(
646 1 : WriterProperties::builder()
647 1 : .set_compression(parquet::basic::Compression::ZSTD(
648 1 : ZstdLevel::try_new(10).unwrap(),
649 1 : ))
650 1 : .build(),
651 1 : ),
652 1 : rows_per_group: 2_000,
653 1 : file_size: 1_000_000,
654 1 : max_duration: time::Duration::from_secs(20 * 60),
655 1 : test_remote_failures: 0,
656 1 : };
657 :
658 1 : let rx = random_stream(50_000);
659 1 : let file_stats = run_test(tmpdir.path(), config, rx).await;
660 :
661 : // with strong compression, the files are smaller
662 1 : assert_eq!(
663 : file_stats,
664 : [
665 : (1205810, 5, 10000),
666 : (1205534, 5, 10000),
667 : (1205835, 5, 10000),
668 : (1205820, 5, 10000),
669 : (1206074, 5, 10000)
670 : ]
671 : );
672 :
673 1 : tmpdir.close().unwrap();
674 1 : }
675 :
676 : #[tokio::test]
677 1 : async fn verify_parquet_unreliable_upload() {
678 1 : let tmpdir = camino_tempfile::tempdir().unwrap();
679 :
680 1 : let config = ParquetConfig {
681 1 : propeties: Arc::new(WriterProperties::new()),
682 1 : rows_per_group: 2_000,
683 1 : file_size: 1_000_000,
684 1 : max_duration: time::Duration::from_secs(20 * 60),
685 1 : test_remote_failures: 2,
686 1 : };
687 :
688 1 : let rx = random_stream(50_000);
689 1 : let file_stats = run_test(tmpdir.path(), config, rx).await;
690 :
691 1 : assert_eq!(
692 : file_stats,
693 : [
694 : (1313953, 3, 6000),
695 : (1313942, 3, 6000),
696 : (1314001, 3, 6000),
697 : (1313958, 3, 6000),
698 : (1314094, 3, 6000),
699 : (1313931, 3, 6000),
700 : (1313725, 3, 6000),
701 : (1313960, 3, 6000),
702 : (438318, 1, 2000)
703 : ]
704 : );
705 :
706 1 : tmpdir.close().unwrap();
707 1 : }
708 :
709 : #[tokio::test(start_paused = true)]
710 1 : async fn verify_parquet_regular_upload() {
711 1 : let tmpdir = camino_tempfile::tempdir().unwrap();
712 :
713 1 : let config = ParquetConfig {
714 1 : propeties: Arc::new(WriterProperties::new()),
715 1 : rows_per_group: 2_000,
716 1 : file_size: 1_000_000,
717 1 : max_duration: time::Duration::from_secs(60),
718 1 : test_remote_failures: 2,
719 1 : };
720 :
721 1 : let (tx, mut rx) = mpsc::unbounded_channel();
722 :
723 1 : tokio::spawn(async move {
724 4 : for _ in 0..3 {
725 3 : let mut s = random_stream(3000);
726 9003 : while let Some(r) = s.next().await {
727 9000 : tx.send(r).unwrap();
728 9000 : }
729 3 : time::sleep(time::Duration::from_secs(70)).await;
730 : }
731 1 : });
732 :
733 9071 : let rx = futures::stream::poll_fn(move |cx| rx.poll_recv(cx));
734 1 : let file_stats = run_test(tmpdir.path(), config, rx).await;
735 :
736 : // files are smaller than the size threshold, but they took too long to fill so were flushed early
737 1 : assert_eq!(
738 : file_stats,
739 : [(658584, 2, 3001), (658298, 2, 3000), (658094, 2, 2999)]
740 : );
741 :
742 1 : tmpdir.close().unwrap();
743 1 : }
744 : }
|
