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