LCOV - code coverage report
Current view: top level - proxy/src/context - parquet.rs (source / functions) Coverage Total Hit
Test: 6df3fc19ec669bcfbbf9aba41d1338898d24eaa0.info Lines: 78.0 % 519 405
Test Date: 2025-03-12 18:28:53 Functions: 55.3 % 103 57

            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 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::{Span, debug, info};
      21              : use utils::backoff;
      22              : 
      23              : use super::{LOG_CHAN, RequestContextInner};
      24              : use crate::config::remote_storage_from_toml;
      25              : use crate::context::LOG_CHAN_DISCONNECT;
      26              : use crate::ext::TaskExt;
      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            3 :     #[clap(long, default_value_t = 8192)]
      41            0 :     parquet_upload_row_group_size: usize,
      42              : 
      43              :     /// How large each column page should be in bytes
      44            3 :     #[clap(long, default_value_t = DEFAULT_PAGE_SIZE)]
      45            0 :     parquet_upload_page_size: usize,
      46              : 
      47              :     /// How large the total parquet file should be in bytes
      48            3 :     #[clap(long, default_value_t = 100_000_000)]
      49            0 :     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            0 :     parquet_upload_maximum_duration: tokio::time::Duration,
      54              : 
      55              :     /// What level of compression to use
      56            3 :     #[clap(long, default_value_t = Compression::UNCOMPRESSED)]
      57            0 :     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      1272000 : #[derive(parquet_derive::ParquetRecordWriter)]
      76              : pub(crate) struct RequestData {
      77              :     region: &'static str,
      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            0 :     {
     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            0 :         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: value.region,
     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            0 :         }
     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 : ) -> anyhow::Result<()> {
     171            0 :     let Some(remote_storage_config) = config.parquet_upload_remote_storage else {
     172            0 :         tracing::warn!("parquet request upload: no s3 bucket configured");
     173            0 :         return Ok(());
     174              :     };
     175              : 
     176            0 :     let (tx, mut rx) = mpsc::unbounded_channel();
     177            0 :     LOG_CHAN
     178            0 :         .set(tx.downgrade())
     179            0 :         .expect("only one worker should set the channel");
     180            0 : 
     181            0 :     // setup row stream that will close on cancellation
     182            0 :     let cancellation_token2 = cancellation_token.clone();
     183            0 :     tokio::spawn(async move {
     184            0 :         cancellation_token2.cancelled().await;
     185              :         // dropping this sender will cause the channel to close only once
     186              :         // all the remaining inflight requests have been completed.
     187            0 :         drop(tx);
     188            0 :     });
     189            0 :     let rx = futures::stream::poll_fn(move |cx| rx.poll_recv(cx));
     190            0 :     let rx = rx.map(RequestData::from);
     191              : 
     192            0 :     let storage = GenericRemoteStorage::from_config(&remote_storage_config)
     193            0 :         .await
     194            0 :         .context("remote storage init")?;
     195              : 
     196            0 :     let properties = WriterProperties::builder()
     197            0 :         .set_data_page_size_limit(config.parquet_upload_page_size)
     198            0 :         .set_compression(config.parquet_upload_compression);
     199            0 : 
     200            0 :     let parquet_config = ParquetConfig {
     201            0 :         propeties: Arc::new(properties.build()),
     202            0 :         rows_per_group: config.parquet_upload_row_group_size,
     203            0 :         file_size: config.parquet_upload_size,
     204            0 :         max_duration: config.parquet_upload_maximum_duration,
     205            0 : 
     206            0 :         #[cfg(any(test, feature = "testing"))]
     207            0 :         test_remote_failures: 0,
     208            0 :     };
     209              : 
     210              :     // TODO(anna): consider moving this to a separate function.
     211            0 :     if let Some(disconnect_events_storage_config) =
     212            0 :         config.parquet_upload_disconnect_events_remote_storage
     213              :     {
     214            0 :         let (tx_disconnect, mut rx_disconnect) = mpsc::unbounded_channel();
     215            0 :         LOG_CHAN_DISCONNECT
     216            0 :             .set(tx_disconnect.downgrade())
     217            0 :             .expect("only one worker should set the channel");
     218            0 : 
     219            0 :         // setup row stream that will close on cancellation
     220            0 :         tokio::spawn(async move {
     221            0 :             cancellation_token.cancelled().await;
     222              :             // dropping this sender will cause the channel to close only once
     223              :             // all the remaining inflight requests have been completed.
     224            0 :             drop(tx_disconnect);
     225            0 :         });
     226            0 :         let rx_disconnect = futures::stream::poll_fn(move |cx| rx_disconnect.poll_recv(cx));
     227            0 :         let rx_disconnect = rx_disconnect.map(RequestData::from);
     228              : 
     229            0 :         let storage_disconnect =
     230            0 :             GenericRemoteStorage::from_config(&disconnect_events_storage_config)
     231            0 :                 .await
     232            0 :                 .context("remote storage for disconnect events init")?;
     233            0 :         let parquet_config_disconnect = parquet_config.clone();
     234            0 :         tokio::try_join!(
     235            0 :             worker_inner(storage, rx, parquet_config),
     236            0 :             worker_inner(storage_disconnect, rx_disconnect, parquet_config_disconnect)
     237            0 :         )
     238            0 :         .map(|_| ())
     239              :     } else {
     240            0 :         worker_inner(storage, rx, parquet_config).await
     241              :     }
     242            0 : }
     243              : 
     244              : #[derive(Clone, Debug)]
     245              : struct ParquetConfig {
     246              :     propeties: WriterPropertiesPtr,
     247              :     rows_per_group: usize,
     248              :     file_size: i64,
     249              : 
     250              :     max_duration: tokio::time::Duration,
     251              : 
     252              :     #[cfg(any(test, feature = "testing"))]
     253              :     test_remote_failures: u64,
     254              : }
     255              : 
     256            4 : async fn worker_inner(
     257            4 :     storage: GenericRemoteStorage,
     258            4 :     rx: impl Stream<Item = RequestData>,
     259            4 :     config: ParquetConfig,
     260            4 : ) -> anyhow::Result<()> {
     261              :     #[cfg(any(test, feature = "testing"))]
     262            4 :     let storage = if config.test_remote_failures > 0 {
     263            2 :         GenericRemoteStorage::unreliable_wrapper(storage, config.test_remote_failures)
     264              :     } else {
     265            2 :         storage
     266              :     };
     267              : 
     268            4 :     let mut rx = std::pin::pin!(rx);
     269            4 : 
     270            4 :     let mut rows = Vec::with_capacity(config.rows_per_group);
     271              : 
     272            4 :     let schema = rows.as_slice().schema()?;
     273            4 :     let buffer = BytesMut::new();
     274            4 :     let w = buffer.writer();
     275            4 :     let mut w = SerializedFileWriter::new(w, schema.clone(), config.propeties.clone())?;
     276              : 
     277            4 :     let mut last_upload = time::Instant::now();
     278            4 : 
     279            4 :     let mut len = 0;
     280       159004 :     while let Some(row) = rx.next().await {
     281       159000 :         rows.push(row);
     282       159000 :         let force = last_upload.elapsed() > config.max_duration;
     283       159000 :         if rows.len() == config.rows_per_group || force {
     284              :             let rg_meta;
     285           80 :             (rows, w, rg_meta) = flush_rows(rows, w).await?;
     286           80 :             len += rg_meta.compressed_size();
     287       158920 :         }
     288       159000 :         if len > config.file_size || force {
     289           23 :             last_upload = time::Instant::now();
     290           23 :             let file = upload_parquet(w, len, &storage).await?;
     291           23 :             w = SerializedFileWriter::new(file, schema.clone(), config.propeties.clone())?;
     292           23 :             len = 0;
     293       158977 :         }
     294              :     }
     295              : 
     296            4 :     if !rows.is_empty() {
     297              :         let rg_meta;
     298            1 :         (_, w, rg_meta) = flush_rows(rows, w).await?;
     299            1 :         len += rg_meta.compressed_size();
     300            3 :     }
     301              : 
     302            4 :     if !w.flushed_row_groups().is_empty() {
     303            3 :         let _rtchk: Writer<BytesMut> = upload_parquet(w, len, &storage).await?;
     304            1 :     }
     305              : 
     306            4 :     Ok(())
     307            4 : }
     308              : 
     309           81 : async fn flush_rows<W>(
     310           81 :     rows: Vec<RequestData>,
     311           81 :     mut w: SerializedFileWriter<W>,
     312           81 : ) -> anyhow::Result<(
     313           81 :     Vec<RequestData>,
     314           81 :     SerializedFileWriter<W>,
     315           81 :     RowGroupMetaDataPtr,
     316           81 : )>
     317           81 : where
     318           81 :     W: std::io::Write + Send + 'static,
     319           81 : {
     320           81 :     let span = Span::current();
     321           81 :     let (mut rows, w, rg_meta) = tokio::task::spawn_blocking(move || {
     322           81 :         let _enter = span.enter();
     323              : 
     324           81 :         let mut rg = w.next_row_group()?;
     325           81 :         rows.as_slice().write_to_row_group(&mut rg)?;
     326           81 :         let rg_meta = rg.close()?;
     327              : 
     328           81 :         let size = rg_meta.compressed_size();
     329           81 :         let compression = rg_meta.compressed_size() as f64 / rg_meta.total_byte_size() as f64;
     330           81 : 
     331           81 :         debug!(size, compression, "flushed row group to parquet file");
     332              : 
     333           81 :         Ok::<_, parquet::errors::ParquetError>((rows, w, rg_meta))
     334           81 :     })
     335           81 :     .await
     336           81 :     .propagate_task_panic()?;
     337              : 
     338           81 :     rows.clear();
     339           81 :     Ok((rows, w, rg_meta))
     340           81 : }
     341              : 
     342           26 : async fn upload_parquet(
     343           26 :     mut w: SerializedFileWriter<Writer<BytesMut>>,
     344           26 :     len: i64,
     345           26 :     storage: &GenericRemoteStorage,
     346           26 : ) -> anyhow::Result<Writer<BytesMut>> {
     347           26 :     let len_uncompressed = w
     348           26 :         .flushed_row_groups()
     349           26 :         .iter()
     350           81 :         .map(|rg| rg.total_byte_size())
     351           26 :         .sum::<i64>();
     352              : 
     353              :     // I don't know how compute intensive this is, although it probably isn't much... better be safe than sorry.
     354              :     // finish method only available on the fork: https://github.com/apache/arrow-rs/issues/5253
     355           26 :     let (mut buffer, metadata) =
     356           26 :         tokio::task::spawn_blocking(move || -> parquet::errors::Result<_> {
     357           26 :             let metadata = w.finish()?;
     358           26 :             let buffer = std::mem::take(w.inner_mut().get_mut());
     359           26 :             Ok((buffer, metadata))
     360           26 :         })
     361           26 :         .await
     362           26 :         .propagate_task_panic()?;
     363              : 
     364           26 :     let data = buffer.split().freeze();
     365           26 : 
     366           26 :     let compression = len as f64 / len_uncompressed as f64;
     367           26 :     let size = data.len();
     368           26 :     let now = chrono::Utc::now();
     369           26 :     let id = uuid::Uuid::new_v7(uuid::Timestamp::from_unix(
     370           26 :         uuid::NoContext,
     371           26 :         // we won't be running this in 1970. this cast is ok
     372           26 :         now.timestamp() as u64,
     373           26 :         now.timestamp_subsec_nanos(),
     374           26 :     ));
     375           26 : 
     376           26 :     info!(
     377              :         %id,
     378              :         rows = metadata.num_rows,
     379            0 :         size, compression, "uploading request parquet file"
     380              :     );
     381              : 
     382           26 :     let year = now.year();
     383           26 :     let month = now.month();
     384           26 :     let day = now.day();
     385           26 :     let hour = now.hour();
     386              :     // segment files by time for S3 performance
     387           26 :     let path = RemotePath::from_string(&format!(
     388           26 :         "{year:04}/{month:02}/{day:02}/{hour:02}/requests_{id}.parquet"
     389           26 :     ))?;
     390           26 :     let cancel = CancellationToken::new();
     391           26 :     let maybe_err = backoff::retry(
     392           38 :         || async {
     393           38 :             let stream = futures::stream::once(futures::future::ready(Ok(data.clone())));
     394           38 :             storage
     395           38 :                 .upload(stream, data.len(), &path, None, &cancel)
     396           38 :                 .await
     397           76 :         },
     398           26 :         TimeoutOrCancel::caused_by_cancel,
     399           26 :         FAILED_UPLOAD_WARN_THRESHOLD,
     400           26 :         FAILED_UPLOAD_MAX_RETRIES,
     401           26 :         "request_data_upload",
     402           26 :         // we don't want cancellation to interrupt here, so we make a dummy cancel token
     403           26 :         &cancel,
     404           26 :     )
     405           26 :     .await
     406           26 :     .ok_or_else(|| anyhow::Error::new(TimeoutOrCancel::Cancel))
     407           26 :     .and_then(|x| x)
     408           26 :     .with_context(|| format!("request_data_upload: path={path}"))
     409           26 :     .err();
     410              : 
     411           26 :     if let Some(err) = maybe_err {
     412            0 :         tracing::error!(%id, %path, error = ?err, "failed to upload request data");
     413           26 :     }
     414              : 
     415           26 :     Ok(buffer.writer())
     416           26 : }
     417              : 
     418              : #[cfg(test)]
     419              : #[expect(clippy::unwrap_used)]
     420              : mod tests {
     421              :     use std::net::Ipv4Addr;
     422              :     use std::num::NonZeroUsize;
     423              :     use std::sync::Arc;
     424              : 
     425              :     use camino::Utf8Path;
     426              :     use clap::Parser;
     427              :     use futures::{Stream, StreamExt};
     428              :     use itertools::Itertools;
     429              :     use parquet::basic::{Compression, ZstdLevel};
     430              :     use parquet::file::properties::{DEFAULT_PAGE_SIZE, WriterProperties};
     431              :     use parquet::file::reader::FileReader;
     432              :     use parquet::file::serialized_reader::SerializedFileReader;
     433              :     use rand::rngs::StdRng;
     434              :     use rand::{Rng, SeedableRng};
     435              :     use remote_storage::{
     436              :         DEFAULT_MAX_KEYS_PER_LIST_RESPONSE, DEFAULT_REMOTE_STORAGE_S3_CONCURRENCY_LIMIT,
     437              :         GenericRemoteStorage, RemoteStorageConfig, RemoteStorageKind, S3Config,
     438              :     };
     439              :     use tokio::sync::mpsc;
     440              :     use tokio::time;
     441              :     use walkdir::WalkDir;
     442              : 
     443              :     use super::{ParquetConfig, ParquetUploadArgs, RequestData, worker_inner};
     444              : 
     445              :     #[derive(Parser)]
     446              :     struct ProxyCliArgs {
     447              :         #[clap(flatten)]
     448              :         parquet_upload: ParquetUploadArgs,
     449              :     }
     450              : 
     451              :     #[test]
     452            1 :     fn default_parser() {
     453            1 :         let ProxyCliArgs { parquet_upload } = ProxyCliArgs::parse_from(["proxy"]);
     454            1 :         assert_eq!(parquet_upload.parquet_upload_remote_storage, None);
     455            1 :         assert_eq!(parquet_upload.parquet_upload_row_group_size, 8192);
     456            1 :         assert_eq!(parquet_upload.parquet_upload_page_size, DEFAULT_PAGE_SIZE);
     457            1 :         assert_eq!(parquet_upload.parquet_upload_size, 100_000_000);
     458            1 :         assert_eq!(
     459            1 :             parquet_upload.parquet_upload_maximum_duration,
     460            1 :             time::Duration::from_secs(20 * 60)
     461            1 :         );
     462            1 :         assert_eq!(
     463            1 :             parquet_upload.parquet_upload_compression,
     464            1 :             Compression::UNCOMPRESSED
     465            1 :         );
     466            1 :     }
     467              : 
     468              :     #[test]
     469            1 :     fn full_parser() {
     470            1 :         let ProxyCliArgs { parquet_upload } = ProxyCliArgs::parse_from([
     471            1 :             "proxy",
     472            1 :             "--parquet-upload-remote-storage",
     473            1 :             "{bucket_name='default',prefix_in_bucket='proxy/',bucket_region='us-east-1',endpoint='http://minio:9000'}",
     474            1 :             "--parquet-upload-row-group-size",
     475            1 :             "100",
     476            1 :             "--parquet-upload-page-size",
     477            1 :             "10000",
     478            1 :             "--parquet-upload-size",
     479            1 :             "10000000",
     480            1 :             "--parquet-upload-maximum-duration",
     481            1 :             "10m",
     482            1 :             "--parquet-upload-compression",
     483            1 :             "zstd(5)",
     484            1 :         ]);
     485            1 :         assert_eq!(
     486            1 :             parquet_upload.parquet_upload_remote_storage,
     487            1 :             Some(RemoteStorageConfig {
     488            1 :                 storage: RemoteStorageKind::AwsS3(S3Config {
     489            1 :                     bucket_name: "default".into(),
     490            1 :                     bucket_region: "us-east-1".into(),
     491            1 :                     prefix_in_bucket: Some("proxy/".into()),
     492            1 :                     endpoint: Some("http://minio:9000".into()),
     493            1 :                     concurrency_limit: NonZeroUsize::new(
     494            1 :                         DEFAULT_REMOTE_STORAGE_S3_CONCURRENCY_LIMIT
     495            1 :                     )
     496            1 :                     .unwrap(),
     497            1 :                     max_keys_per_list_response: DEFAULT_MAX_KEYS_PER_LIST_RESPONSE,
     498            1 :                     upload_storage_class: None,
     499            1 :                 }),
     500            1 :                 timeout: RemoteStorageConfig::DEFAULT_TIMEOUT,
     501            1 :                 small_timeout: RemoteStorageConfig::DEFAULT_SMALL_TIMEOUT,
     502            1 :             })
     503            1 :         );
     504            1 :         assert_eq!(parquet_upload.parquet_upload_row_group_size, 100);
     505            1 :         assert_eq!(parquet_upload.parquet_upload_page_size, 10000);
     506            1 :         assert_eq!(parquet_upload.parquet_upload_size, 10_000_000);
     507            1 :         assert_eq!(
     508            1 :             parquet_upload.parquet_upload_maximum_duration,
     509            1 :             time::Duration::from_secs(10 * 60)
     510            1 :         );
     511            1 :         assert_eq!(
     512            1 :             parquet_upload.parquet_upload_compression,
     513            1 :             Compression::ZSTD(ZstdLevel::try_new(5).unwrap())
     514            1 :         );
     515            1 :     }
     516              : 
     517       159000 :     fn generate_request_data(rng: &mut impl Rng) -> RequestData {
     518       159000 :         RequestData {
     519       159000 :             session_id: uuid::Builder::from_random_bytes(rng.r#gen()).into_uuid(),
     520       159000 :             peer_addr: Ipv4Addr::from(rng.r#gen::<[u8; 4]>()).to_string(),
     521       159000 :             timestamp: chrono::DateTime::from_timestamp_millis(
     522       159000 :                 rng.gen_range(1703862754..1803862754),
     523       159000 :             )
     524       159000 :             .unwrap()
     525       159000 :             .naive_utc(),
     526       159000 :             application_name: Some("test".to_owned()),
     527       159000 :             user_agent: Some("test-user-agent".to_owned()),
     528       159000 :             username: Some(hex::encode(rng.r#gen::<[u8; 4]>())),
     529       159000 :             endpoint_id: Some(hex::encode(rng.r#gen::<[u8; 16]>())),
     530       159000 :             database: Some(hex::encode(rng.r#gen::<[u8; 16]>())),
     531       159000 :             project: Some(hex::encode(rng.r#gen::<[u8; 16]>())),
     532       159000 :             branch: Some(hex::encode(rng.r#gen::<[u8; 16]>())),
     533       159000 :             pg_options: None,
     534       159000 :             auth_method: None,
     535       159000 :             jwt_issuer: None,
     536       159000 :             protocol: ["tcp", "ws", "http"][rng.gen_range(0..3)],
     537       159000 :             region: "us-east-1",
     538       159000 :             error: None,
     539       159000 :             success: rng.r#gen(),
     540       159000 :             cold_start_info: "no",
     541       159000 :             duration_us: rng.gen_range(0..30_000_000),
     542       159000 :             disconnect_timestamp: None,
     543       159000 :         }
     544       159000 :     }
     545              : 
     546            6 :     fn random_stream(len: usize) -> impl Stream<Item = RequestData> + Unpin {
     547            6 :         let mut rng = StdRng::from_seed([0x39; 32]);
     548            6 :         futures::stream::iter(
     549       159000 :             std::iter::repeat_with(move || generate_request_data(&mut rng)).take(len),
     550            6 :         )
     551            6 :     }
     552              : 
     553            4 :     async fn run_test(
     554            4 :         tmpdir: &Utf8Path,
     555            4 :         config: ParquetConfig,
     556            4 :         rx: impl Stream<Item = RequestData>,
     557            4 :     ) -> Vec<(u64, usize, i64)> {
     558            4 :         let remote_storage_config = RemoteStorageConfig {
     559            4 :             storage: RemoteStorageKind::LocalFs {
     560            4 :                 local_path: tmpdir.to_path_buf(),
     561            4 :             },
     562            4 :             timeout: std::time::Duration::from_secs(120),
     563            4 :             small_timeout: std::time::Duration::from_secs(30),
     564            4 :         };
     565            4 :         let storage = GenericRemoteStorage::from_config(&remote_storage_config)
     566            4 :             .await
     567            4 :             .unwrap();
     568            4 : 
     569            4 :         worker_inner(storage, rx, config).await.unwrap();
     570            4 : 
     571            4 :         let mut files = WalkDir::new(tmpdir.as_std_path())
     572            4 :             .into_iter()
     573           46 :             .filter_map(|entry| entry.ok())
     574           46 :             .filter(|entry| entry.file_type().is_file())
     575           26 :             .map(|entry| entry.path().to_path_buf())
     576            4 :             .collect_vec();
     577            4 :         files.sort();
     578            4 : 
     579            4 :         files
     580            4 :             .into_iter()
     581           26 :             .map(|path| std::fs::File::open(tmpdir.as_std_path().join(path)).unwrap())
     582           26 :             .map(|file| {
     583           26 :                 (
     584           26 :                     file.metadata().unwrap(),
     585           26 :                     SerializedFileReader::new(file).unwrap().metadata().clone(),
     586           26 :                 )
     587           26 :             })
     588           26 :             .map(|(file_meta, parquet_meta)| {
     589           26 :                 (
     590           26 :                     file_meta.len(),
     591           26 :                     parquet_meta.num_row_groups(),
     592           26 :                     parquet_meta.file_metadata().num_rows(),
     593           26 :                 )
     594           26 :             })
     595            4 :             .collect()
     596            4 :     }
     597              : 
     598              :     #[tokio::test]
     599            1 :     async fn verify_parquet_no_compression() {
     600            1 :         let tmpdir = camino_tempfile::tempdir().unwrap();
     601            1 : 
     602            1 :         let config = ParquetConfig {
     603            1 :             propeties: Arc::new(WriterProperties::new()),
     604            1 :             rows_per_group: 2_000,
     605            1 :             file_size: 1_000_000,
     606            1 :             max_duration: time::Duration::from_secs(20 * 60),
     607            1 :             test_remote_failures: 0,
     608            1 :         };
     609            1 : 
     610            1 :         let rx = random_stream(50_000);
     611            1 :         let file_stats = run_test(tmpdir.path(), config, rx).await;
     612            1 : 
     613            1 :         assert_eq!(
     614            1 :             file_stats,
     615            1 :             [
     616            1 :                 (1313953, 3, 6000),
     617            1 :                 (1313942, 3, 6000),
     618            1 :                 (1314001, 3, 6000),
     619            1 :                 (1313958, 3, 6000),
     620            1 :                 (1314094, 3, 6000),
     621            1 :                 (1313931, 3, 6000),
     622            1 :                 (1313725, 3, 6000),
     623            1 :                 (1313960, 3, 6000),
     624            1 :                 (438318, 1, 2000)
     625            1 :             ]
     626            1 :         );
     627            1 : 
     628            1 :         tmpdir.close().unwrap();
     629            1 :     }
     630              : 
     631              :     #[tokio::test]
     632            1 :     async fn verify_parquet_strong_compression() {
     633            1 :         let tmpdir = camino_tempfile::tempdir().unwrap();
     634            1 : 
     635            1 :         let config = ParquetConfig {
     636            1 :             propeties: Arc::new(
     637            1 :                 WriterProperties::builder()
     638            1 :                     .set_compression(parquet::basic::Compression::ZSTD(
     639            1 :                         ZstdLevel::try_new(10).unwrap(),
     640            1 :                     ))
     641            1 :                     .build(),
     642            1 :             ),
     643            1 :             rows_per_group: 2_000,
     644            1 :             file_size: 1_000_000,
     645            1 :             max_duration: time::Duration::from_secs(20 * 60),
     646            1 :             test_remote_failures: 0,
     647            1 :         };
     648            1 : 
     649            1 :         let rx = random_stream(50_000);
     650            1 :         let file_stats = run_test(tmpdir.path(), config, rx).await;
     651            1 : 
     652            1 :         // with strong compression, the files are smaller
     653            1 :         assert_eq!(
     654            1 :             file_stats,
     655            1 :             [
     656            1 :                 (1205810, 5, 10000),
     657            1 :                 (1205534, 5, 10000),
     658            1 :                 (1205835, 5, 10000),
     659            1 :                 (1205820, 5, 10000),
     660            1 :                 (1206074, 5, 10000)
     661            1 :             ]
     662            1 :         );
     663            1 : 
     664            1 :         tmpdir.close().unwrap();
     665            1 :     }
     666              : 
     667              :     #[tokio::test]
     668            1 :     async fn verify_parquet_unreliable_upload() {
     669            1 :         let tmpdir = camino_tempfile::tempdir().unwrap();
     670            1 : 
     671            1 :         let config = ParquetConfig {
     672            1 :             propeties: Arc::new(WriterProperties::new()),
     673            1 :             rows_per_group: 2_000,
     674            1 :             file_size: 1_000_000,
     675            1 :             max_duration: time::Duration::from_secs(20 * 60),
     676            1 :             test_remote_failures: 2,
     677            1 :         };
     678            1 : 
     679            1 :         let rx = random_stream(50_000);
     680            1 :         let file_stats = run_test(tmpdir.path(), config, rx).await;
     681            1 : 
     682            1 :         assert_eq!(
     683            1 :             file_stats,
     684            1 :             [
     685            1 :                 (1313953, 3, 6000),
     686            1 :                 (1313942, 3, 6000),
     687            1 :                 (1314001, 3, 6000),
     688            1 :                 (1313958, 3, 6000),
     689            1 :                 (1314094, 3, 6000),
     690            1 :                 (1313931, 3, 6000),
     691            1 :                 (1313725, 3, 6000),
     692            1 :                 (1313960, 3, 6000),
     693            1 :                 (438318, 1, 2000)
     694            1 :             ]
     695            1 :         );
     696            1 : 
     697            1 :         tmpdir.close().unwrap();
     698            1 :     }
     699              : 
     700              :     #[tokio::test(start_paused = true)]
     701            1 :     async fn verify_parquet_regular_upload() {
     702            1 :         let tmpdir = camino_tempfile::tempdir().unwrap();
     703            1 : 
     704            1 :         let config = ParquetConfig {
     705            1 :             propeties: Arc::new(WriterProperties::new()),
     706            1 :             rows_per_group: 2_000,
     707            1 :             file_size: 1_000_000,
     708            1 :             max_duration: time::Duration::from_secs(60),
     709            1 :             test_remote_failures: 2,
     710            1 :         };
     711            1 : 
     712            1 :         let (tx, mut rx) = mpsc::unbounded_channel();
     713            1 : 
     714            1 :         tokio::spawn(async move {
     715            4 :             for _ in 0..3 {
     716            3 :                 let mut s = random_stream(3000);
     717         9003 :                 while let Some(r) = s.next().await {
     718         9000 :                     tx.send(r).unwrap();
     719         9000 :                 }
     720            3 :                 time::sleep(time::Duration::from_secs(70)).await;
     721            1 :             }
     722            1 :         });
     723            1 : 
     724         9071 :         let rx = futures::stream::poll_fn(move |cx| rx.poll_recv(cx));
     725            1 :         let file_stats = run_test(tmpdir.path(), config, rx).await;
     726            1 : 
     727            1 :         // files are smaller than the size threshold, but they took too long to fill so were flushed early
     728            1 :         assert_eq!(
     729            1 :             file_stats,
     730            1 :             [(658584, 2, 3001), (658298, 2, 3000), (658094, 2, 2999)]
     731            1 :         );
     732            1 : 
     733            1 :         tmpdir.close().unwrap();
     734            1 :     }
     735              : }
        

Generated by: LCOV version 2.1-beta