diff --git a/crates/modelardb_manager/src/main.rs b/crates/modelardb_manager/src/main.rs index 637f1a6c1..7b3b8a678 100644 --- a/crates/modelardb_manager/src/main.rs +++ b/crates/modelardb_manager/src/main.rs @@ -25,7 +25,6 @@ use std::{env, process}; use modelardb_storage::delta_lake::DeltaLake; use modelardb_types::flight::protocol; -use tokio::runtime::Runtime; use tokio::sync::RwLock; use tonic::metadata::errors::InvalidMetadataValue; use tonic::metadata::{Ascii, MetadataValue}; @@ -100,14 +99,12 @@ pub struct Context { /// Flight server. Returns [`ModelarDbManagerError`] if the command line arguments cannot be parsed, /// if the metadata cannot be read from the Delta Lake, or if the Apache Arrow Flight server cannot /// be started. -fn main() -> Result<()> { +#[tokio::main] +async fn main() -> Result<()> { // Initialize a tracing layer that logs events to stdout. let stdout_log = tracing_subscriber::fmt::layer(); tracing_subscriber::registry().with(stdout_log).init(); - // Create a Tokio runtime for executing asynchronous tasks. - let runtime = Arc::new(Runtime::new()?); - let user_arguments = collect_command_line_arguments(3); let user_arguments: Vec<&str> = user_arguments.iter().map(|arg| arg.as_str()).collect(); let remote_data_folder_str = match user_arguments.as_slice() { @@ -115,36 +112,36 @@ fn main() -> Result<()> { _ => print_usage_and_exit_with_error("remote_data_folder"), }; - let context = runtime.block_on(async { - let remote_data_folder = RemoteDataFolder::try_new(remote_data_folder_str).await?; + let remote_data_folder = RemoteDataFolder::try_new(remote_data_folder_str).await?; - let nodes = remote_data_folder.metadata_manager.nodes().await?; + let nodes = remote_data_folder.metadata_manager.nodes().await?; - let mut cluster = Cluster::new(); - for node in nodes { - cluster.register_node(node)?; - } + let mut cluster = Cluster::new(); + for node in nodes { + cluster.register_node(node)?; + } - // Retrieve and parse the key to a tonic metadata value since it is used in tonic requests. - let key = remote_data_folder - .metadata_manager - .manager_key() - .await? - .to_string() - .parse() - .map_err(|error: InvalidMetadataValue| { - ModelarDbManagerError::InvalidArgument(error.to_string()) - })?; - - // Create the Context. - Ok::, ModelarDbManagerError>(Arc::new(Context { - remote_data_folder, - cluster: RwLock::new(cluster), - key, - })) - })?; - - start_apache_arrow_flight_server(context, &runtime, *PORT) + // Retrieve and parse the key to a tonic metadata value since it is used in tonic requests. + let key = remote_data_folder + .metadata_manager + .manager_key() + .await? + .to_string() + .parse() + .map_err(|error: InvalidMetadataValue| { + ModelarDbManagerError::InvalidArgument(error.to_string()) + })?; + + // Create the Context. + let context = Arc::new(Context { + remote_data_folder, + cluster: RwLock::new(cluster), + key, + }); + + start_apache_arrow_flight_server(context, *PORT).await?; + + Ok(()) } /// Collect the command line arguments that this program was started with. diff --git a/crates/modelardb_manager/src/remote.rs b/crates/modelardb_manager/src/remote.rs index e487b1375..6829d3532 100644 --- a/crates/modelardb_manager/src/remote.rs +++ b/crates/modelardb_manager/src/remote.rs @@ -38,7 +38,6 @@ use modelardb_storage::parser::ModelarDbStatement; use modelardb_types::flight::protocol; use modelardb_types::types::{Table, TimeSeriesTableMetadata}; use prost::Message; -use tokio::runtime::Runtime; use tonic::transport::Server; use tonic::{Request, Response, Status, Streaming}; use tracing::info; @@ -47,11 +46,7 @@ use crate::Context; use crate::error::{ModelarDbManagerError, Result}; /// Start an Apache Arrow Flight server on 0.0.0.0:`port`. -pub fn start_apache_arrow_flight_server( - context: Arc, - runtime: &Arc, - port: u16, -) -> Result<()> { +pub async fn start_apache_arrow_flight_server(context: Arc, port: u16) -> Result<()> { let localhost_with_port = "0.0.0.0:".to_owned() + &port.to_string(); let localhost_with_port: SocketAddr = localhost_with_port.parse().map_err(|error| { ModelarDbManagerError::InvalidArgument(format!( @@ -63,13 +58,10 @@ pub fn start_apache_arrow_flight_server( info!("Starting Apache Arrow Flight on {}.", localhost_with_port); - runtime - .block_on(async { - Server::builder() - .add_service(flight_service_server) - .serve(localhost_with_port) - .await - }) + Server::builder() + .add_service(flight_service_server) + .serve(localhost_with_port) + .await .map_err(|error| error.into()) } diff --git a/crates/modelardb_server/src/configuration.rs b/crates/modelardb_server/src/configuration.rs index b85643f9d..61fcc099a 100644 --- a/crates/modelardb_server/src/configuration.rs +++ b/crates/modelardb_server/src/configuration.rs @@ -244,7 +244,6 @@ mod tests { use std::sync::Arc; use tempfile::TempDir; - use tokio::runtime::Runtime; use tokio::sync::RwLock; use uuid::Uuid; @@ -429,13 +428,9 @@ mod tests { ))); let storage_engine = Arc::new(RwLock::new( - StorageEngine::try_new( - Arc::new(Runtime::new().unwrap()), - data_folders, - &configuration_manager, - ) - .await - .unwrap(), + StorageEngine::try_new(data_folders, &configuration_manager) + .await + .unwrap(), )); (storage_engine, configuration_manager) diff --git a/crates/modelardb_server/src/context.rs b/crates/modelardb_server/src/context.rs index a116298aa..f292944da 100644 --- a/crates/modelardb_server/src/context.rs +++ b/crates/modelardb_server/src/context.rs @@ -22,7 +22,6 @@ use datafusion::arrow::datatypes::Schema; use datafusion::catalog::{SchemaProvider, TableProvider}; use datafusion::prelude::SessionContext; use modelardb_types::types::TimeSeriesTableMetadata; -use tokio::runtime::Runtime; use tokio::sync::RwLock; use tracing::info; @@ -48,17 +47,13 @@ impl Context { /// Create the components needed in the [`Context`] and use them to create the [`Context`]. If a /// metadata manager or storage engine could not be created, [`ModelarDbServerError`] is /// returned. - pub async fn try_new( - runtime: Arc, - data_folders: DataFolders, - cluster_mode: ClusterMode, - ) -> Result { + pub async fn try_new(data_folders: DataFolders, cluster_mode: ClusterMode) -> Result { let configuration_manager = Arc::new(RwLock::new(ConfigurationManager::new(cluster_mode))); let session_context = modelardb_storage::create_session_context(); let storage_engine = Arc::new(RwLock::new( - StorageEngine::try_new(runtime, data_folders.clone(), &configuration_manager).await?, + StorageEngine::try_new(data_folders.clone(), &configuration_manager).await?, )); Ok(Context { @@ -1004,7 +999,6 @@ mod tests { Arc::new( Context::try_new( - Arc::new(Runtime::new().unwrap()), DataFolders::new(local_data_folder.clone(), None, local_data_folder), ClusterMode::SingleNode, ) diff --git a/crates/modelardb_server/src/main.rs b/crates/modelardb_server/src/main.rs index 2e1f34fa2..23171b7be 100644 --- a/crates/modelardb_server/src/main.rs +++ b/crates/modelardb_server/src/main.rs @@ -28,7 +28,6 @@ mod storage; use std::sync::{Arc, LazyLock}; use std::{env, process}; -use tokio::runtime::Runtime; use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; use crate::context::Context; @@ -58,54 +57,45 @@ pub enum ClusterMode { /// [`ModelarDbServerError`](error::ModelarDbServerError) if the command line arguments /// cannot be parsed, if the metadata cannot be read from the database, or if the Apache Arrow /// Flight interface cannot be started. -fn main() -> Result<()> { +#[tokio::main] +async fn main() -> Result<()> { // Initialize a tracing layer that logs events to stdout. let stdout_log = tracing_subscriber::fmt::layer(); tracing_subscriber::registry().with(stdout_log).init(); - // Create a Tokio runtime for executing asynchronous tasks. The runtime is not in the context, so - // it can be passed to the components in the context. - let runtime = Arc::new(Runtime::new()?); - let arguments = collect_command_line_arguments(3); let arguments: Vec<&str> = arguments.iter().map(|arg| arg.as_str()).collect(); let (cluster_mode, data_folders) = if let Ok(cluster_mode_and_data_folders) = - runtime.block_on(DataFolders::try_from_command_line_arguments(&arguments)) + DataFolders::try_from_command_line_arguments(&arguments).await { cluster_mode_and_data_folders } else { print_usage_and_exit_with_error("[server_mode] local_data_folder_url [manager_url]"); }; - let context = Arc::new(runtime.block_on(Context::try_new( - runtime.clone(), - data_folders, - cluster_mode.clone(), - ))?); + let context = Arc::new(Context::try_new(data_folders, cluster_mode.clone()).await?); // Register normal tables and time series tables. - runtime.block_on(context.register_normal_tables())?; - runtime.block_on(context.register_time_series_tables())?; + context.register_normal_tables().await?; + context.register_time_series_tables().await?; if let ClusterMode::MultiNode(manager) = &cluster_mode { - runtime.block_on(manager.retrieve_and_create_tables(&context))?; + manager.retrieve_and_create_tables(&context).await?; } // Setup CTRL+C handler. - setup_ctrl_c_handler(&context, &runtime); + setup_ctrl_c_handler(&context); // Initialize storage engine with spilled buffers. - runtime.block_on(async { - context - .storage_engine - .read() - .await - .initialize(&context) - .await - })?; + context + .storage_engine + .read() + .await + .initialize(&context) + .await?; // Start the Apache Arrow Flight interface. - remote::start_apache_arrow_flight_server(context, &runtime, *PORT)?; + remote::start_apache_arrow_flight_server(context, *PORT).await?; Ok(()) } @@ -136,9 +126,9 @@ pub fn print_usage_and_exit_with_error(parameters: &str) -> ! { /// Register a handler to execute when CTRL+C is pressed. The handler takes an exclusive lock for /// the storage engine, flushes the data the storage engine currently buffers, and terminates the /// system without releasing the lock. -fn setup_ctrl_c_handler(context: &Arc, runtime: &Arc) { +fn setup_ctrl_c_handler(context: &Arc) { let ctrl_c_context = context.clone(); - runtime.spawn(async move { + tokio::spawn(async move { // Errors are consciously ignored as the program should terminate if the handler cannot be // registered as buffers otherwise cannot be flushed. tokio::signal::ctrl_c().await.unwrap(); diff --git a/crates/modelardb_server/src/remote.rs b/crates/modelardb_server/src/remote.rs index 66d509e4b..419fb3ec5 100644 --- a/crates/modelardb_server/src/remote.rs +++ b/crates/modelardb_server/src/remote.rs @@ -46,7 +46,6 @@ use modelardb_types::flight::protocol; use modelardb_types::functions; use modelardb_types::types::{Table, TimeSeriesTableMetadata}; use prost::Message; -use tokio::runtime::Runtime; use tokio::sync::mpsc::{self, Sender}; use tokio::task; use tokio_stream::wrappers::ReceiverStream; @@ -61,11 +60,7 @@ use crate::error::{ModelarDbServerError, Result}; /// Start an Apache Arrow Flight server on 0.0.0.0:`port` that passes `context` to the methods that /// process the requests through [`FlightServiceHandler`]. -pub fn start_apache_arrow_flight_server( - context: Arc, - runtime: &Arc, - port: u16, -) -> Result<()> { +pub async fn start_apache_arrow_flight_server(context: Arc, port: u16) -> Result<()> { let localhost_with_port = "0.0.0.0:".to_owned() + &port.to_string(); let localhost_with_port: SocketAddr = localhost_with_port.parse().map_err(|error| { ModelarDbServerError::InvalidArgument(format!( @@ -79,13 +74,11 @@ pub fn start_apache_arrow_flight_server( FlightServiceServer::new(handler).max_decoding_message_size(16777216); info!("Starting Apache Arrow Flight on {}.", localhost_with_port); - runtime - .block_on(async { - Server::builder() - .add_service(flight_service_server) - .serve(localhost_with_port) - .await - }) + + Server::builder() + .add_service(flight_service_server) + .serve(localhost_with_port) + .await .map_err(|error| error.into()) } diff --git a/crates/modelardb_server/src/storage/compressed_data_manager.rs b/crates/modelardb_server/src/storage/compressed_data_manager.rs index b10e944ab..f99f3ae86 100644 --- a/crates/modelardb_server/src/storage/compressed_data_manager.rs +++ b/crates/modelardb_server/src/storage/compressed_data_manager.rs @@ -21,7 +21,7 @@ use std::sync::Arc; use crossbeam_queue::SegQueue; use dashmap::DashMap; use datafusion::arrow::record_batch::RecordBatch; -use tokio::runtime::Runtime; +use tokio::runtime::Handle; use tokio::sync::RwLock; use tracing::{debug, error, info}; @@ -106,20 +106,21 @@ impl CompressedDataManager { /// Read and process messages received from the /// [`UncompressedDataManager`](super::UncompressedDataManager) to either insert compressed /// data, flush buffers, or stop. - pub(super) fn process_compressed_messages(&self, runtime: Arc) -> Result<()> { + pub(super) fn process_compressed_messages(&self, runtime_handle: Handle) -> Result<()> { loop { let message = self.channels.compressed_data_receiver.recv()?; match message { Message::Data(compressed_segment_batch) => { - runtime.block_on(self.insert_compressed_segments(compressed_segment_batch))?; + runtime_handle + .block_on(self.insert_compressed_segments(compressed_segment_batch))?; } Message::Flush => { - self.flush_and_log_errors(&runtime); + self.flush_and_log_errors(&runtime_handle); self.channels.result_sender.send(Ok(()))?; } Message::Stop => { - self.flush_and_log_errors(&runtime); + self.flush_and_log_errors(&runtime_handle); self.channels.result_sender.send(Ok(()))?; break; } @@ -201,8 +202,8 @@ impl CompressedDataManager { /// Flush the data that the [`CompressedDataManager`] is currently managing. Writes a log /// message if some of the data cannot be flushed. - fn flush_and_log_errors(&self, runtime: &Runtime) { - runtime.block_on(async { + fn flush_and_log_errors(&self, runtime_handle: &Handle) { + runtime_handle.block_on(async { if let Err(error) = self.flush().await { error!( "Failed to flush data in compressed data manager due to: {}", diff --git a/crates/modelardb_server/src/storage/mod.rs b/crates/modelardb_server/src/storage/mod.rs index 29952f561..adcd9e1f5 100644 --- a/crates/modelardb_server/src/storage/mod.rs +++ b/crates/modelardb_server/src/storage/mod.rs @@ -35,7 +35,7 @@ use std::thread::{self, JoinHandle}; use datafusion::arrow::record_batch::RecordBatch; use modelardb_types::types::TimeSeriesTableMetadata; -use tokio::runtime::Runtime; +use tokio::runtime::Handle; use tokio::sync::RwLock; use tracing::error; @@ -90,7 +90,6 @@ impl StorageEngine { /// [`ModelarDbServerError`] if `remote_data_folder` is given but [`DataTransfer`] cannot not be /// created. pub(super) async fn try_new( - runtime: Arc, data_folders: DataFolders, configuration_manager: &Arc>, ) -> Result { @@ -102,6 +101,9 @@ impl StorageEngine { configuration_manager.compressed_reserved_memory_in_bytes(), )); + // Create a handle to the Tokio runtime that can be used in the threads. + let runtime_handle = Handle::current(); + // Create threads and shared channels. let mut join_handles = vec![]; let channels = Arc::new(Channels::new()); @@ -114,7 +116,7 @@ impl StorageEngine { )); { - let runtime = runtime.clone(); + let runtime_handle = runtime_handle.clone(); let uncompressed_data_manager = uncompressed_data_manager.clone(); Self::start_threads( @@ -122,7 +124,7 @@ impl StorageEngine { "Ingestion", move || { if let Err(error) = - uncompressed_data_manager.process_uncompressed_messages(runtime) + uncompressed_data_manager.process_uncompressed_messages(runtime_handle) { error!("Failed to receive uncompressed message due to: {}", error); }; @@ -132,7 +134,7 @@ impl StorageEngine { } { - let runtime = runtime.clone(); + let runtime_handle = runtime_handle.clone(); let uncompressed_data_manager = uncompressed_data_manager.clone(); Self::start_threads( @@ -140,7 +142,7 @@ impl StorageEngine { "Compression", move || { if let Err(error) = - uncompressed_data_manager.process_compressor_messages(runtime) + uncompressed_data_manager.process_compressor_messages(runtime_handle) { error!("Failed to receive compressor message due to: {}", error); }; @@ -173,14 +175,15 @@ impl StorageEngine { )); { - let runtime = runtime.clone(); + let runtime_handle = runtime_handle.clone(); let compressed_data_manager = compressed_data_manager.clone(); Self::start_threads( configuration_manager.writer_threads, "Writer", move || { - if let Err(error) = compressed_data_manager.process_compressed_messages(runtime) + if let Err(error) = + compressed_data_manager.process_compressed_messages(runtime_handle) { error!("Failed to receive compressed message due to: {}", error); }; diff --git a/crates/modelardb_server/src/storage/uncompressed_data_manager.rs b/crates/modelardb_server/src/storage/uncompressed_data_manager.rs index 3fe2595a2..6607637b5 100644 --- a/crates/modelardb_server/src/storage/uncompressed_data_manager.rs +++ b/crates/modelardb_server/src/storage/uncompressed_data_manager.rs @@ -26,7 +26,7 @@ use dashmap::DashMap; use futures::StreamExt; use modelardb_types::types::{TimeSeriesTableMetadata, Timestamp, Value}; use object_store::path::{Path, PathPart}; -use tokio::runtime::Runtime; +use tokio::runtime::Handle; use tracing::{debug, error, warn}; use crate::context::Context; @@ -127,13 +127,13 @@ impl UncompressedDataManager { /// Read and process messages received from the [`StorageEngine`](super::StorageEngine) to /// either ingest uncompressed data, flush buffers, or stop. - pub(super) fn process_uncompressed_messages(&self, runtime: Arc) -> Result<()> { + pub(super) fn process_uncompressed_messages(&self, runtime_handle: Handle) -> Result<()> { loop { let message = self.channels.ingested_data_receiver.recv()?; match message { Message::Data(ingested_data_buffer) => { - runtime.block_on(self.insert_data_points(ingested_data_buffer))?; + runtime_handle.block_on(self.insert_data_points(ingested_data_buffer))?; } Message::Flush => { self.flush_and_log_errors(); @@ -524,13 +524,13 @@ impl UncompressedDataManager { /// Read and process messages received from the [`UncompressedDataManager`] to either compress /// uncompressed data, forward a flush message, or stop. - pub(super) fn process_compressor_messages(&self, runtime: Arc) -> Result<()> { + pub(super) fn process_compressor_messages(&self, runtime_handle: Handle) -> Result<()> { loop { let message = self.channels.uncompressed_data_receiver.recv()?; match message { Message::Data(data_buffer) => { - runtime.block_on(self.compress_finished_buffer(data_buffer))?; + runtime_handle.block_on(self.compress_finished_buffer(data_buffer))?; } Message::Flush => { self.channels.compressed_data_sender.send(Message::Flush)?; @@ -663,6 +663,7 @@ mod tests { use modelardb_types::types::{TimestampBuilder, ValueBuilder}; use object_store::local::LocalFileSystem; use tempfile::TempDir; + use tokio::runtime::Runtime; use tokio::time::{Duration, sleep}; use crate::storage::UNCOMPRESSED_DATA_BUFFER_CAPACITY; @@ -681,7 +682,6 @@ mod tests { // Create a context with a storage engine. let context = Arc::new( Context::try_new( - Arc::new(Runtime::new().unwrap()), DataFolders::new(local_data_folder.clone(), None, local_data_folder), ClusterMode::SingleNode, ) @@ -1116,7 +1116,10 @@ mod tests { .uncompressed_data_sender .send(Message::Stop) .unwrap(); - data_manager.process_compressor_messages(runtime).unwrap(); + + data_manager + .process_compressor_messages(runtime.handle().clone()) + .unwrap(); assert!( remaining_memory @@ -1167,7 +1170,9 @@ mod tests { // Since the UncompressedOnDiskDataBuffer is not in memory, the remaining amount of memory // should not increase when it is processed. - data_manager.process_compressor_messages(runtime).unwrap(); + data_manager + .process_compressor_messages(runtime.handle().clone()) + .unwrap(); assert_eq!( remaining_memory, diff --git a/crates/modelardb_storage/src/query/generated_as_exec.rs b/crates/modelardb_storage/src/query/generated_as_exec.rs index 518a2a556..4deab7474 100644 --- a/crates/modelardb_storage/src/query/generated_as_exec.rs +++ b/crates/modelardb_storage/src/query/generated_as_exec.rs @@ -123,7 +123,7 @@ impl ExecutionPlan for GeneratedAsExec { } /// Return the single execution plan batches of rows are read from. - fn children(&self) -> Vec<&Arc<(dyn ExecutionPlan)>> { + fn children(&self) -> Vec<&Arc> { vec![&self.input] } @@ -131,8 +131,8 @@ impl ExecutionPlan for GeneratedAsExec { /// [`DataFusionError::Plan`] is returned if `children` does not contain exactly one element. fn with_new_children( self: Arc, - mut children: Vec>, - ) -> DataFusionResult> { + mut children: Vec>, + ) -> DataFusionResult> { if children.len() == 1 { Ok(GeneratedAsExec::new( self.schema.clone(), diff --git a/crates/modelardb_storage/src/query/sorted_join_exec.rs b/crates/modelardb_storage/src/query/sorted_join_exec.rs index 11bc8e3fe..adae7ddf4 100644 --- a/crates/modelardb_storage/src/query/sorted_join_exec.rs +++ b/crates/modelardb_storage/src/query/sorted_join_exec.rs @@ -118,7 +118,7 @@ impl ExecutionPlan for SortedJoinExec { } /// Return the single execution plan batches of rows are read from. - fn children(&self) -> Vec<&Arc<(dyn ExecutionPlan)>> { + fn children(&self) -> Vec<&Arc> { // iter() returns an iterator that produces elements of type &T. self.inputs.iter().collect() } @@ -128,8 +128,8 @@ impl ExecutionPlan for SortedJoinExec { /// contain at least one element. fn with_new_children( self: Arc, - children: Vec>, - ) -> DataFusionResult> { + children: Vec>, + ) -> DataFusionResult> { if !children.is_empty() { Ok(SortedJoinExec::new( self.schema.clone(),