Use #[tokio::main] attribute instead of Runtime::new()

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,51 +99,49 @@ 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() {
         &[remote_data_folder_str] => remote_data_folder_str,
         _ => 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::<Arc<Context>, 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.

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<Context>,
-    runtime: &Arc<Runtime>,
-    port: u16,
-) -> Result<()> {
+pub async fn start_apache_arrow_flight_server(context: Arc<Context>, 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())
 }
 

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)

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<Runtime>,
-        data_folders: DataFolders,
-        cluster_mode: ClusterMode,
-    ) -> Result<Self> {
+    pub async fn try_new(data_folders: DataFolders, cluster_mode: ClusterMode) -> Result<Self> {
         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,
             )

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<Context>, runtime: &Arc<Runtime>) {
+fn setup_ctrl_c_handler(context: &Arc<Context>) {
     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();

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<Context>,
-    runtime: &Arc<Runtime>,
-    port: u16,
-) -> Result<()> {
+pub async fn start_apache_arrow_flight_server(context: Arc<Context>, 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())
 }
 

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<Runtime>) -> 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: {}",