diff --git a/crates/modelardb_server/src/cluster.rs b/crates/modelardb_server/src/cluster.rs index 10959a02..af97f882 100644 --- a/crates/modelardb_server/src/cluster.rs +++ b/crates/modelardb_server/src/cluster.rs @@ -48,7 +48,7 @@ pub(crate) struct Cluster { /// The remote data folder that each node in the cluster should be synchronized with. /// When a table is created, dropped, vacuumed, or truncated, it is done in the /// remote data folder first. - remote_data_folder: DataFolder, + remote_data_folder: Arc, } impl Cluster { @@ -56,7 +56,7 @@ impl Cluster { /// It is assumed that `node` corresponds to the local system running `modelardbd`. If the /// cluster metadata tables do not exist and could not be created or the node could not be /// saved, return [`ModelarDbServerError`]. - pub(crate) async fn try_new(node: Node, remote_data_folder: DataFolder) -> Result { + pub(crate) async fn try_new(node: Node, remote_data_folder: Arc) -> Result { remote_data_folder .create_and_register_cluster_metadata_tables() .await?; @@ -499,24 +499,19 @@ mod test { // Create a normal table in the remote data folder that should be retrieved and created // in the local data folder and one that already exists. - create_normal_table( - "normal_table_1", - "column", - data_folders.local_data_folder.clone(), - ) - .await; + create_normal_table("normal_table_1", "column", &data_folders.local_data_folder).await; create_normal_table( "normal_table_1", "column", - data_folders.maybe_remote_data_folder.clone().unwrap(), + data_folders.maybe_remote_data_folder.as_ref().unwrap(), ) .await; create_normal_table( "normal_table_2", "column", - data_folders.maybe_remote_data_folder.clone().unwrap(), + data_folders.maybe_remote_data_folder.as_ref().unwrap(), ) .await; @@ -538,17 +533,12 @@ mod test { // Create a normal table in the local data folder with the same name as a normal table in // the remote data folder, but with a different schema. - create_normal_table( - NORMAL_TABLE_NAME, - "local", - data_folders.local_data_folder.clone(), - ) - .await; + create_normal_table(NORMAL_TABLE_NAME, "local", &data_folders.local_data_folder).await; create_normal_table( NORMAL_TABLE_NAME, "remote", - data_folders.maybe_remote_data_folder.clone().unwrap(), + data_folders.maybe_remote_data_folder.as_ref().unwrap(), ) .await; @@ -576,21 +566,21 @@ mod test { create_time_series_table( "time_series_table_1", "field", - data_folders.local_data_folder.clone(), + &data_folders.local_data_folder, ) .await; create_time_series_table( "time_series_table_1", "field", - data_folders.maybe_remote_data_folder.clone().unwrap(), + data_folders.maybe_remote_data_folder.as_ref().unwrap(), ) .await; create_time_series_table( "time_series_table_2", "field", - data_folders.maybe_remote_data_folder.clone().unwrap(), + data_folders.maybe_remote_data_folder.as_ref().unwrap(), ) .await; @@ -615,14 +605,14 @@ mod test { create_time_series_table( TIME_SERIES_TABLE_NAME, "local", - data_folders.local_data_folder.clone(), + &data_folders.local_data_folder, ) .await; create_time_series_table( TIME_SERIES_TABLE_NAME, "remote", - data_folders.maybe_remote_data_folder.clone().unwrap(), + data_folders.maybe_remote_data_folder.as_ref().unwrap(), ) .await; @@ -646,17 +636,12 @@ mod test { let data_folders = context.data_folders.clone(); // Create tables in the local data folder that are not in the remote data folder. - create_normal_table( - NORMAL_TABLE_NAME, - "local", - data_folders.local_data_folder.clone(), - ) - .await; + create_normal_table(NORMAL_TABLE_NAME, "local", &data_folders.local_data_folder).await; create_time_series_table( TIME_SERIES_TABLE_NAME, "local", - data_folders.local_data_folder.clone(), + &data_folders.local_data_folder, ) .await; @@ -673,7 +658,7 @@ mod test { /// Create a normal table named `table_name` with a single column named `column_name` in /// `data_folder`. - async fn create_normal_table(table_name: &str, column_name: &str, data_folder: DataFolder) { + async fn create_normal_table(table_name: &str, column_name: &str, data_folder: &DataFolder) { let schema = Schema::new(vec![Field::new(column_name, ArrowValue::DATA_TYPE, false)]); data_folder @@ -687,7 +672,7 @@ mod test { async fn create_time_series_table( table_name: &str, column_name: &str, - data_folder: DataFolder, + data_folder: &DataFolder, ) { let query_schema = Arc::new(Schema::new(vec![ Field::new("timestamp", ArrowTimestamp::DATA_TYPE, false), @@ -724,7 +709,7 @@ mod test { /// data folder in the context uses a local [`DataFolder`]. async fn create_context(local_temp_dir: &TempDir, remote_temp_dir: &TempDir) -> Arc { let temp_dir_url = local_temp_dir.path().to_str().unwrap(); - let local_data_folder = DataFolder::open_local_url(temp_dir_url).await.unwrap(); + let local_data_folder = Arc::new(DataFolder::open_local_url(temp_dir_url).await.unwrap()); let edge_node = Node::new("edge".to_owned(), ServerMode::Edge); let cluster = create_cluster_with_node(remote_temp_dir, edge_node).await; @@ -813,7 +798,7 @@ mod test { /// Create a [`Cluster`] that uses a local [`DataFolder`] for the remote data folder. async fn create_cluster_with_node(temp_dir: &TempDir, node: Node) -> Cluster { let temp_dir_url = temp_dir.path().to_str().unwrap(); - let local_data_folder = DataFolder::open_local_url(temp_dir_url).await.unwrap(); + let local_data_folder = Arc::new(DataFolder::open_local_url(temp_dir_url).await.unwrap()); Cluster::try_new(node, local_data_folder.clone()) .await diff --git a/crates/modelardb_server/src/configuration.rs b/crates/modelardb_server/src/configuration.rs index 39b1ae59..6b1386d7 100644 --- a/crates/modelardb_server/src/configuration.rs +++ b/crates/modelardb_server/src/configuration.rs @@ -195,7 +195,7 @@ pub struct ConfigurationManager { /// The mode of the write-ahead log used to determine whether data is logged before ingestion. wal_mode: WalMode, /// The local data folder that stores the configuration file at the root. - local_data_folder: DataFolder, + local_data_folder: Arc, /// The configuration of the system. This is stored in a separate type to allow for easier /// serialization and deserialization. configuration: Configuration, @@ -208,7 +208,7 @@ impl ConfigurationManager { /// if the corresponding CLI flags or environment variables are set. If the configuration file /// could not be read or created, [`ModelarDbServerError`] is returned. pub async fn try_new( - local_data_folder: DataFolder, + local_data_folder: Arc, cluster_mode: ClusterMode, args: &ServerArgs, ) -> Result { @@ -589,7 +589,7 @@ mod tests { async fn test_invalid_configuration_in_configuration_file() { let temp_dir = tempfile::tempdir().unwrap(); let local_url = temp_dir.path().to_str().unwrap(); - let local_data_folder = DataFolder::open_local_url(local_url).await.unwrap(); + let local_data_folder = Arc::new(DataFolder::open_local_url(local_url).await.unwrap()); // Multiple threads per component are not supported. let invalid_configuration = Configuration { @@ -619,7 +619,7 @@ mod tests { async fn test_invalid_toml_in_configuration_file() { let temp_dir = tempfile::tempdir().unwrap(); let local_url = temp_dir.path().to_str().unwrap(); - let local_data_folder = DataFolder::open_local_url(local_url).await.unwrap(); + let local_data_folder = Arc::new(DataFolder::open_local_url(local_url).await.unwrap()); // Write invalid TOML to the configuration file. let path = temp_dir.path().join(CONFIGURATION_FILE_NAME); @@ -1000,11 +1000,11 @@ mod tests { Arc>, ) { let local_url = temp_dir.path().to_str().unwrap(); - let local_data_folder = DataFolder::open_local_url(local_url).await.unwrap(); + let local_data_folder = Arc::new(DataFolder::open_local_url(local_url).await.unwrap()); let target_dir = tempfile::tempdir().unwrap(); let target_url = target_dir.path().to_str().unwrap(); - let remote_data_folder = DataFolder::open_local_url(target_url).await.unwrap(); + let remote_data_folder = Arc::new(DataFolder::open_local_url(target_url).await.unwrap()); let data_folders = DataFolders::new( local_data_folder.clone(), diff --git a/crates/modelardb_server/src/context.rs b/crates/modelardb_server/src/context.rs index 9d600959..3244459a 100644 --- a/crates/modelardb_server/src/context.rs +++ b/crates/modelardb_server/src/context.rs @@ -1161,7 +1161,7 @@ mod tests { /// Create a simple [`Context`] that uses `temp_dir` as the local data folder and query data folder. async fn create_context(temp_dir: &TempDir) -> Arc { let temp_dir_url = temp_dir.path().to_str().unwrap(); - let local_data_folder = DataFolder::open_local_url(temp_dir_url).await.unwrap(); + let local_data_folder = Arc::new(DataFolder::open_local_url(temp_dir_url).await.unwrap()); Arc::new( Context::try_new( diff --git a/crates/modelardb_server/src/data_folders.rs b/crates/modelardb_server/src/data_folders.rs index 8c01d2f9..43f5ea8b 100644 --- a/crates/modelardb_server/src/data_folders.rs +++ b/crates/modelardb_server/src/data_folders.rs @@ -15,6 +15,8 @@ //! Implementation of a struct that provides access to the local and remote data storage components. +use std::sync::Arc; + use modelardb_storage::data_folder::DataFolder; use modelardb_types::types::{Node, ServerMode}; @@ -25,20 +27,20 @@ use crate::{ClusterMode, Result, ServerMode as ServerModeArg}; #[derive(Clone)] pub struct DataFolders { /// Folder for storing metadata and data in Apache Parquet files on the local file system. - pub local_data_folder: DataFolder, + pub local_data_folder: Arc, /// Folder for storing metadata and data in Apache Parquet files in a remote object store. - pub maybe_remote_data_folder: Option, + pub maybe_remote_data_folder: Option>, /// Folder from which metadata and data in Apache Parquet files will be read during query /// execution. It is equivalent to `local_data_folder` when deployed on the edge and /// `remote_data_folder` when deployed in the cloud. - pub query_data_folder: DataFolder, + pub query_data_folder: Arc, } impl DataFolders { pub fn new( - local_data_folder: DataFolder, - maybe_remote_data_folder: Option, - query_data_folder: DataFolder, + local_data_folder: Arc, + maybe_remote_data_folder: Option>, + query_data_folder: Arc, ) -> Self { Self { local_data_folder, @@ -65,7 +67,8 @@ impl DataFolders { remote_data_folder: None, .. } => { - let local_data_folder = DataFolder::open_local_url(local_data_folder).await?; + let local_data_folder = + Arc::new(DataFolder::open_local_url(local_data_folder).await?); Ok(( ClusterMode::SingleNode, @@ -79,9 +82,10 @@ impl DataFolders { credentials, } => { let remote_data_folder = - DataFolder::open_remote_url(remote_data_folder, credentials).await?; + Arc::new(DataFolder::open_remote_url(remote_data_folder, credentials).await?); - let local_data_folder = DataFolder::open_local_url(local_data_folder).await?; + let local_data_folder = + Arc::new(DataFolder::open_local_url(local_data_folder).await?); let node = Node::new(url_with_port, ServerMode::Edge); let cluster = Cluster::try_new(node, remote_data_folder.clone()).await?; @@ -102,9 +106,10 @@ impl DataFolders { credentials, } => { let remote_data_folder = - DataFolder::open_remote_url(remote_data_folder, credentials).await?; + Arc::new(DataFolder::open_remote_url(remote_data_folder, credentials).await?); - let local_data_folder = DataFolder::open_local_url(local_data_folder).await?; + let local_data_folder = + Arc::new(DataFolder::open_local_url(local_data_folder).await?); let node = Node::new(url_with_port, ServerMode::Cloud); let cluster = Cluster::try_new(node, remote_data_folder.clone()).await?; diff --git a/crates/modelardb_server/src/storage/compressed_data_manager.rs b/crates/modelardb_server/src/storage/compressed_data_manager.rs index 94cb11e4..a726149b 100644 --- a/crates/modelardb_server/src/storage/compressed_data_manager.rs +++ b/crates/modelardb_server/src/storage/compressed_data_manager.rs @@ -43,7 +43,7 @@ pub(super) struct CompressedDataManager { /// Component that transfers saved compressed data to the remote data folder when it is necessary. pub(super) data_transfer: Arc>>, /// Folder containing all compressed data managed by the [`StorageEngine`](crate::storage::StorageEngine). - pub(crate) local_data_folder: DataFolder, + pub(crate) local_data_folder: Arc, /// The compressed segments before they are saved to persistent storage. The key is the name of /// the time series table the compressed segments represents data points for so the Apache Parquet /// files can be partitioned by table. @@ -63,7 +63,7 @@ impl CompressedDataManager { pub(super) fn new( data_storage_compactor: Arc>, data_transfer: Arc>>, - local_data_folder: DataFolder, + local_data_folder: Arc, channels: Arc, memory_pool: Arc, wal_mode: WalMode, @@ -580,7 +580,7 @@ mod tests { // Create a local data folder and save a single time series table to the Delta Lake. let temp_dir_url = temp_dir.path().to_str().unwrap(); - let local_data_folder = DataFolder::open_local_url(temp_dir_url).await.unwrap(); + let local_data_folder = Arc::new(DataFolder::open_local_url(temp_dir_url).await.unwrap()); let time_series_table_metadata = table::time_series_table_metadata(); local_data_folder diff --git a/crates/modelardb_server/src/storage/data_storage_compactor.rs b/crates/modelardb_server/src/storage/data_storage_compactor.rs index 4b502351..43cd59fa 100644 --- a/crates/modelardb_server/src/storage/data_storage_compactor.rs +++ b/crates/modelardb_server/src/storage/data_storage_compactor.rs @@ -18,6 +18,8 @@ //! table by merging those small files into fewer larger ones and vacuuming the files left behind, //! reducing storage use and query time. +use std::sync::Arc; + use dashmap::DashMap; use modelardb_storage::data_folder::DataFolder; use tracing::debug; @@ -31,7 +33,7 @@ use crate::error::Result; pub(super) struct DataStorageCompactor { /// The data folder containing all compressed data managed by the /// [`StorageEngine`](crate::storage::StorageEngine). - local_data_folder: DataFolder, + local_data_folder: Arc, /// The target size, in bytes, of the files produced when a table is optimized. A table is /// compacted once its `estimated_compactable_size_in_bytes` reaches this size, so the same /// value decides both when to compact and how large the optimized files are. @@ -56,7 +58,7 @@ impl DataStorageCompactor { /// so small files written before a restart are not forgotten. If the files in `local_data_folder` /// could not be read, return [`ModelarDbServerError`](crate::error::ModelarDbServerError). pub(super) async fn try_new( - local_data_folder: DataFolder, + local_data_folder: Arc, optimize_target_file_size_in_bytes: u64, vacuum_retention_period_in_seconds: u64, ) -> Result { @@ -292,10 +294,10 @@ mod tests { } /// Create a [`DataFolder`] in a local [`TempDir`] containing a single time series table. - async fn create_local_data_folder_with_table() -> (TempDir, DataFolder) { + async fn create_local_data_folder_with_table() -> (TempDir, Arc) { let temp_dir = tempfile::tempdir().unwrap(); let temp_dir_url = temp_dir.path().to_str().unwrap(); - let local_data_folder = DataFolder::open_local_url(temp_dir_url).await.unwrap(); + let local_data_folder = Arc::new(DataFolder::open_local_url(temp_dir_url).await.unwrap()); let time_series_table_metadata = table::time_series_table_metadata(); local_data_folder @@ -333,7 +335,9 @@ mod tests { } /// Create a [`DataStorageCompactor`] that compacts the tables in `local_data_folder`. - async fn create_data_storage_compactor(local_data_folder: DataFolder) -> DataStorageCompactor { + async fn create_data_storage_compactor( + local_data_folder: Arc, + ) -> DataStorageCompactor { DataStorageCompactor::try_new( local_data_folder, OPTIMIZE_TARGET_FILE_SIZE_IN_BYTES, diff --git a/crates/modelardb_server/src/storage/data_transfer.rs b/crates/modelardb_server/src/storage/data_transfer.rs index 885d530c..d29b510e 100644 --- a/crates/modelardb_server/src/storage/data_transfer.rs +++ b/crates/modelardb_server/src/storage/data_transfer.rs @@ -17,6 +17,7 @@ //! is managed here until it is of a sufficient size to be transferred efficiently. use std::collections::HashSet; +use std::sync::Arc; use dashmap::DashMap; use deltalake::arrow::array::RecordBatch; @@ -33,9 +34,9 @@ use crate::error::Result; pub struct DataTransfer { /// The data folder containing all compressed data managed by the /// [`StorageEngine`](crate::storage::StorageEngine). - local_data_folder: DataFolder, + local_data_folder: Arc, /// The data folder that the data should be transferred to. - remote_data_folder: DataFolder, + remote_data_folder: Arc, /// Map from table names to the current size of the table in bytes. table_size_in_bytes: DashMap, /// The number of bytes that are required before transferring a batch of data to the remote @@ -51,8 +52,8 @@ impl DataTransfer { /// `local_data_folder_path` could not be read, return /// [`ModelarDbServerError`](crate::error::ModelarDbServerError). pub async fn try_new( - local_data_folder: DataFolder, - remote_data_folder: DataFolder, + local_data_folder: Arc, + remote_data_folder: Arc, transfer_batch_size_in_bytes: Option, ) -> Result { let table_names = local_data_folder.table_names().await?; @@ -405,10 +406,10 @@ mod tests { /// Create a [`DataFolder`] in a local [`TempDir`] and create a single normal table and a /// single time series table in it. - async fn create_local_data_folder_with_tables() -> (TempDir, DataFolder) { + async fn create_local_data_folder_with_tables() -> (TempDir, Arc) { let temp_dir = tempfile::tempdir().unwrap(); let temp_dir_url = temp_dir.path().to_str().unwrap(); - let local_data_folder = DataFolder::open_local_url(temp_dir_url).await.unwrap(); + let local_data_folder = Arc::new(DataFolder::open_local_url(temp_dir_url).await.unwrap()); // Create a normal table. local_data_folder @@ -467,11 +468,12 @@ mod tests { /// Create a data transfer component with a target object store that is deleted once the test is finished. async fn create_data_transfer_component( - local_data_folder: DataFolder, + local_data_folder: Arc, ) -> (TempDir, DataTransfer) { let target_dir = tempfile::tempdir().unwrap(); let target_dir_url = target_dir.path().to_str().unwrap(); - let remote_data_folder = DataFolder::open_local_url(target_dir_url).await.unwrap(); + let remote_data_folder = + Arc::new(DataFolder::open_local_url(target_dir_url).await.unwrap()); // Set the transfer batch size so that data is transferred if three batches are written. let data_transfer = DataTransfer::try_new( diff --git a/crates/modelardb_server/src/storage/uncompressed_data_buffer.rs b/crates/modelardb_server/src/storage/uncompressed_data_buffer.rs index c2dd9a7b..f573fe11 100644 --- a/crates/modelardb_server/src/storage/uncompressed_data_buffer.rs +++ b/crates/modelardb_server/src/storage/uncompressed_data_buffer.rs @@ -227,7 +227,7 @@ impl UncompressedInMemoryDataBuffer { /// an [`UncompressedOnDiskDataBuffer`] when finished. pub(super) async fn spill_to_apache_parquet( &mut self, - local_data_folder: Arc, + local_object_store: Arc, ) -> Result { let data_points = self.record_batch().await?; @@ -235,7 +235,7 @@ impl UncompressedInMemoryDataBuffer { self.tag_hash, self.time_series_table_metadata.clone(), self.updated_by_batch_index, - local_data_folder, + local_object_store, data_points, self.batch_ids.clone(), ) @@ -272,8 +272,8 @@ pub(super) struct UncompressedOnDiskDataBuffer { time_series_table_metadata: Arc, /// Index of the last batch that added data points to this buffer. updated_by_batch_index: u64, - /// Location for storing spilled buffers at `file_path`. - local_data_folder: Arc, + /// Object store the spilled buffer is written to at `file_path`. + local_object_store: Arc, /// Path to the Apache Parquet file containing the uncompressed data in the /// [`UncompressedOnDiskDataBuffer`]. file_path: Path, @@ -283,14 +283,14 @@ pub(super) struct UncompressedOnDiskDataBuffer { impl UncompressedOnDiskDataBuffer { /// Spill the in-memory `data_points` from the time series with `tag_hash` to an Apache Parquet - /// file in `local_data_folder`. If the Apache Parquet file is written successfully, return an + /// file in `local_object_store`. If the Apache Parquet file is written successfully, return an /// [`UncompressedOnDiskDataBuffer`], otherwise return /// [`ModelarDbServerError`](crate::error::ModelarDbServerError). pub(super) async fn try_spill( tag_hash: u64, time_series_table_metadata: Arc, updated_by_batch_index: u64, - local_data_folder: Arc, + local_object_store: Arc, data_points: RecordBatch, batch_ids: HashSet, ) -> Result { @@ -307,7 +307,7 @@ impl UncompressedOnDiskDataBuffer { &file_path, &data_points, None, - &(local_data_folder.clone() as Arc), + &local_object_store, ) .await?; @@ -315,7 +315,7 @@ impl UncompressedOnDiskDataBuffer { tag_hash, time_series_table_metadata, updated_by_batch_index, - local_data_folder, + local_object_store, file_path, batch_ids, }) @@ -328,11 +328,11 @@ impl UncompressedOnDiskDataBuffer { pub(super) async fn record_batch(&self) -> Result { let data_points = modelardb_storage::read_record_batch_from_apache_parquet_file( &self.file_path, - self.local_data_folder.clone(), + self.local_object_store.clone(), ) .await?; - self.local_data_folder.delete(&self.file_path).await?; + self.local_object_store.delete(&self.file_path).await?; Ok(data_points) } @@ -766,7 +766,7 @@ mod tests { /// Create an on-disk data buffer in `temp_dir` from a full `UncompressedInMemoryDataBuffer`. async fn create_on_disk_data_buffer(temp_dir: &TempDir) -> UncompressedOnDiskDataBuffer { - let local_data_folder = + let local_object_store = Arc::new(LocalFileSystem::new_with_prefix(temp_dir.path()).unwrap()); let mut uncompressed_in_memory_buffer_to_be_spilled = UncompressedInMemoryDataBuffer::new( @@ -783,7 +783,7 @@ mod tests { ); uncompressed_in_memory_buffer_to_be_spilled - .spill_to_apache_parquet(local_data_folder) + .spill_to_apache_parquet(local_object_store) .await .unwrap() } diff --git a/crates/modelardb_server/src/storage/uncompressed_data_manager.rs b/crates/modelardb_server/src/storage/uncompressed_data_manager.rs index 369d2742..e1bd73d4 100644 --- a/crates/modelardb_server/src/storage/uncompressed_data_manager.rs +++ b/crates/modelardb_server/src/storage/uncompressed_data_manager.rs @@ -45,7 +45,7 @@ use crate::storage::uncompressed_data_buffer::{ /// files. When an uncompressed data buffer is finished the data is made available for compression. pub(super) struct UncompressedDataManager { /// Folder for storing metadata and data in Apache Parquet files on the local file system. - pub local_data_folder: DataFolder, + pub local_data_folder: Arc, /// Counter incremented for each [`RecordBatch`](datafusion::arrow::array::RecordBatch) of data /// points ingested. The value is assigned to buffers that are created or updated and is used to /// flush buffers that are no longer used. @@ -71,7 +71,7 @@ impl UncompressedDataManager { /// there are any. If the existing buffers could not be deleted, return /// [`ModelarDbServerError`](crate::error::ModelarDbServerError). pub(super) async fn try_new( - local_data_folder: DataFolder, + local_data_folder: Arc, memory_pool: Arc, channels: Arc, ) -> Result { @@ -1196,7 +1196,7 @@ mod tests { temp_dir: &TempDir, ) -> (UncompressedDataManager, Arc) { let temp_dir_url = temp_dir.path().to_str().unwrap(); - let local_data_folder = DataFolder::open_local_url(temp_dir_url).await.unwrap(); + let local_data_folder = Arc::new(DataFolder::open_local_url(temp_dir_url).await.unwrap()); // Ensure the expected metadata is available through the metadata manager. let time_series_table_metadata = table::time_series_table_metadata(); diff --git a/crates/modelardb_server/tests/integration_test.rs b/crates/modelardb_server/tests/integration_test.rs index ea5dad15..93855e5e 100644 --- a/crates/modelardb_server/tests/integration_test.rs +++ b/crates/modelardb_server/tests/integration_test.rs @@ -111,13 +111,13 @@ impl TestContext { self.client = Self::create_client(self.port).await; } - /// Create a server that stores data in `local_data_folder` and listens on `port` and ensure it - /// is ready to receive requests. - async fn create_server(local_data_folder: &TempDir, port: u16) -> Child { + /// Create a server that stores data in the local data folder at `temp_dir` and listens on + /// `port` and ensure it is ready to receive requests. + async fn create_server(temp_dir: &TempDir, port: u16) -> Child { // The server's stdout and stderr are piped so the log messages (stdout) and expected errors // (stderr) are not printed when all the tests are run using the "cargo test" command. // modelardbd is run using dev-release so the tests can use larger more realistic data sets. - let local_data_folder = local_data_folder.path().to_str().unwrap(); + let local_data_folder_path = temp_dir.path().to_str().unwrap(); let mut server = Command::new("cargo") .env("MODELARDBD_PORT", port.to_string()) .args([ @@ -127,7 +127,7 @@ impl TestContext { "--bin", "modelardbd", "edge", - local_data_folder, + local_data_folder_path, ]) .kill_on_drop(true) .stdout(Stdio::piped()) diff --git a/crates/modelardb_storage/src/data_folder/mod.rs b/crates/modelardb_storage/src/data_folder/mod.rs index 4d9b4f2b..02248147 100644 --- a/crates/modelardb_storage/src/data_folder/mod.rs +++ b/crates/modelardb_storage/src/data_folder/mod.rs @@ -66,8 +66,9 @@ enum TableType { TimeSeriesTable, } -/// Functionality for managing Delta Lake tables in a local folder or an object store. -#[derive(Clone)] +/// Functionality for managing Delta Lake tables in a local folder or an object store. A single +/// instance should be shared through an [`Arc`], so the cache and session context are only +/// created once. pub struct DataFolder { /// URL to access the root of the Delta Lake. location: String, @@ -78,7 +79,7 @@ pub struct DataFolder { /// Cache of Delta tables to avoid opening the same table multiple times. delta_table_cache: DashMap, /// Session context used to query the tables using Apache DataFusion. - session_context: Arc, + session_context: SessionContext, } impl DataFolder { @@ -261,7 +262,7 @@ impl DataFolder { storage_options, object_store, delta_table_cache: DashMap::new(), - session_context: Arc::new(crate::create_session_context()), + session_context: crate::create_session_context(), }; data_folder.create_and_register_metadata_tables().await?; @@ -824,10 +825,17 @@ impl DataFolder { /// [`ModelarDbStorageError`] if a connection to the Delta Lake cannot be established or the /// table does not exist. async fn delta_table_from_path(&self, table_path: &str) -> Result { - // Use the cache if possible and load to get the latest table data. - if let Some(mut delta_table) = self.delta_table_cache.get_mut(table_path) { + // Clone the cached table if possible and drop the DashMap guard before loading. load() is + // an async IO call, and get_mut() holds a write guard on the shard, so loading while + // holding it would serialize access to every table in the same shard. + let maybe_delta_table = self + .delta_table_cache + .get(table_path) + .map(|delta_table| delta_table.clone()); + + if let Some(mut delta_table) = maybe_delta_table { delta_table.load().await?; - Ok(delta_table.clone()) + Ok(delta_table) } else { // Return a clear error message if the table does not exist instead of the internal // error message from deltalake.