diff --git a/Cargo.lock b/Cargo.lock index bae02ae5..832c23c9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3558,6 +3558,7 @@ dependencies = [ "modelardb_compression", "modelardb_storage", "modelardb_types", + "prost", "tempfile", "tokio", "tonic", @@ -3599,6 +3600,7 @@ dependencies = [ "serde", "snmalloc-rs", "sqlparser", + "sysinfo", "tempfile", "tokio", "tokio-stream", diff --git a/crates/modelardb_embedded/Cargo.toml b/crates/modelardb_embedded/Cargo.toml index cd33b310..f376fb18 100644 --- a/crates/modelardb_embedded/Cargo.toml +++ b/crates/modelardb_embedded/Cargo.toml @@ -35,6 +35,7 @@ modelardb_auth.workspace = true modelardb_compression.workspace = true modelardb_storage.workspace = true modelardb_types.workspace = true +prost.workspace = true tokio.workspace = true tonic.workspace = true diff --git a/crates/modelardb_embedded/src/error.rs b/crates/modelardb_embedded/src/error.rs index e3d1c78e..9fd0db40 100644 --- a/crates/modelardb_embedded/src/error.rs +++ b/crates/modelardb_embedded/src/error.rs @@ -29,6 +29,7 @@ use deltalake::{DeltaTableError, ObjectStoreError}; use modelardb_compression::error::ModelarDbCompressionError; use modelardb_storage::error::ModelarDbStorageError; use modelardb_types::error::ModelarDbTypesError; +use prost::DecodeError; use tonic::Status as TonicStatusError; use tonic::transport::Error as TonicTransportError; @@ -58,6 +59,8 @@ pub enum ModelarDbEmbeddedError { ObjectStore(ObjectStoreError), /// Error returned by Apache Parquet. Parquet(ParquetError), + /// Error returned by Prost when decoding a message that is not valid. + ProstDecode(DecodeError), /// Status returned by Tonic. TonicStatus(Box), /// Error returned by Tonic. @@ -83,6 +86,7 @@ impl Display for ModelarDbEmbeddedError { Self::ModelarDbTypes(reason) => write!(f, "ModelarDB Types Error: {reason}"), Self::ObjectStore(reason) => write!(f, "Object Store Error: {reason}"), Self::Parquet(reason) => write!(f, "Parquet Error: {reason}"), + Self::ProstDecode(reason) => write!(f, "Prost Decode Error: {reason}"), Self::TonicStatus(reason) => write!(f, "Tonic Status Error: {reason}"), Self::TonicTransport(reason) => write!(f, "Tonic Transport Error: {reason}"), Self::Unimplemented(reason) => write!(f, "Unimplemented Error: {reason}"), @@ -104,6 +108,7 @@ impl Error for ModelarDbEmbeddedError { Self::ModelarDbTypes(reason) => Some(reason), Self::ObjectStore(reason) => Some(reason), Self::Parquet(reason) => Some(reason), + Self::ProstDecode(reason) => Some(reason), Self::TonicStatus(reason) => Some(reason), Self::TonicTransport(reason) => Some(reason), Self::Unimplemented(_reason) => None, @@ -166,6 +171,12 @@ impl From for ModelarDbEmbeddedError { } } +impl From for ModelarDbEmbeddedError { + fn from(error: DecodeError) -> Self { + Self::ProstDecode(error) + } +} + impl From for ModelarDbEmbeddedError { fn from(error: TonicStatusError) -> Self { Self::TonicStatus(Box::new(error)) diff --git a/crates/modelardb_embedded/src/operations/client.rs b/crates/modelardb_embedded/src/operations/client.rs index 4456d067..642bafd6 100644 --- a/crates/modelardb_embedded/src/operations/client.rs +++ b/crates/modelardb_embedded/src/operations/client.rs @@ -28,16 +28,20 @@ use arrow::record_batch::RecordBatch; use arrow_flight::decode::FlightRecordBatchStream; use arrow_flight::encode::FlightDataEncoderBuilder; use arrow_flight::flight_service_client::FlightServiceClient; -use arrow_flight::{Action, Criteria, FlightDescriptor, Ticket}; +use arrow_flight::{Action, Criteria, FlightDescriptor, Result as FlightResult, Ticket}; use async_trait::async_trait; use datafusion::error::DataFusionError; use datafusion::execution::RecordBatchStream; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use futures::{StreamExt, TryStreamExt, stream}; use modelardb_auth::BearerInterceptor; +use modelardb_types::flight::protocol; +use modelardb_types::types::Node; +use prost::Message; +use prost::bytes::Bytes; use tonic::codegen::InterceptedService; use tonic::transport::{Channel, Endpoint}; -use tonic::{Request, Status}; +use tonic::{Request, Status, Streaming}; use crate::error::{ModelarDbEmbeddedError, Result}; use crate::operations::{ @@ -65,6 +69,134 @@ impl Client { Ok(Client { flight_client }) } + + /// Returns the current configuration of the node. If the configuration could not be retrieved, + /// [`ModelarDbEmbeddedError`] is returned. + pub async fn configuration(&mut self) -> Result { + let bytes = self.retrieve_action_bytes("GetConfiguration").await?; + + Ok(protocol::Configuration::decode(bytes)?) + } + + /// Updates `setting` in the node configuration to `new_value`. If the setting could not be + /// updated, [`ModelarDbEmbeddedError`] is returned. + pub async fn update_configuration( + &mut self, + setting: protocol::update_configuration::Setting, + new_value: Option, + ) -> Result<()> { + let update_configuration = protocol::UpdateConfiguration { + setting: setting as i32, + new_value, + }; + + self.send_action("UpdateConfiguration", update_configuration.encode_to_vec()) + .await?; + + Ok(()) + } + + /// Flushes all data in memory to disk. If the data could not be flushed, + /// [`ModelarDbEmbeddedError`] is returned. + pub async fn flush_memory(&mut self) -> Result<()> { + self.send_action("FlushMemory", vec![]).await?; + + Ok(()) + } + + /// Flushes all data in memory to disk and then transfers all compressed data to the remote + /// object store. If the data could not be flushed, [`ModelarDbEmbeddedError`] is returned. + pub async fn flush_node(&mut self) -> Result<()> { + self.send_action("FlushNode", vec![]).await?; + + Ok(()) + } + + /// Flushes all data to disk, transfers it to the remote object store, removes the node from the + /// cluster if necessary, and kills the node process. Since the process is killed, a + /// conventional response cannot be returned, so a dropped connection is not treated as an + /// error. If the data could not be flushed before the node was killed, + /// [`ModelarDbEmbeddedError`] is returned. + pub async fn kill_node(&mut self) -> Result<()> { + // The node exits while handling this action, so the response stream is dropped. + let _ = self.send_action("KillNode", vec![]).await; + + Ok(()) + } + + /// Returns the nodes that are currently part of the cluster. A single node returns only itself. + /// If the nodes could not be retrieved, [`ModelarDbEmbeddedError`] is returned. + pub async fn list_nodes(&mut self) -> Result> { + let bytes = self.retrieve_action_bytes("ListNodes").await?; + + Ok(modelardb_types::flight::deserialize_and_extract_cluster_nodes(&bytes)?) + } + + /// Returns the current resource usage metrics of the node. If the metrics could not be + /// retrieved, [`ModelarDbEmbeddedError`] is returned. + pub async fn node_metrics(&mut self) -> Result { + let bytes = self.retrieve_action_bytes("NodeMetrics").await?; + + Ok(protocol::NodeMetrics::decode(bytes)?) + } + + /// Sends the action with the type `action_type` and an empty body to the node and returns the + /// body of the response. If the action could not be performed, [`ModelarDbEmbeddedError`] is + /// returned. + async fn retrieve_action_bytes(&mut self, action_type: &str) -> Result { + let mut response = self.send_action(action_type, vec![]).await?; + + let message = response.message().await?.ok_or_else(|| { + ModelarDbEmbeddedError::from(Status::internal(format!( + "Action '{action_type}' did not return a response message." + ))) + })?; + + Ok(message.body) + } + + /// Sends the action with the type `action_type` and `body` to the node and returns the response + /// stream. If the action could not be performed, [`ModelarDbEmbeddedError`] is returned. + async fn send_action( + &mut self, + action_type: &str, + body: Vec, + ) -> Result> { + let action = Action { + r#type: action_type.to_owned(), + body: body.into(), + }; + + let response = self.flight_client.do_action(Request::new(action)).await?; + + Ok(response.into_inner()) + } + + /// Returns the URL of the cloud node that the node assigns to execute the SQL in `sql`. If the + /// node is not running in a cluster, or a cloud node could not be assigned, + /// [`ModelarDbEmbeddedError`] is returned. + pub async fn cloud_query_node(&mut self, sql: &str) -> Result { + let flight_descriptor = FlightDescriptor::new_cmd(sql.to_owned()); + let flight_info = self + .flight_client + .get_flight_info(Request::new(flight_descriptor)) + .await? + .into_inner(); + + let endpoint = flight_info.endpoint.into_iter().next().ok_or_else(|| { + ModelarDbEmbeddedError::InvalidArgument( + "The node did not return an endpoint for the query.".to_owned(), + ) + })?; + + let location = endpoint.location.into_iter().next().ok_or_else(|| { + ModelarDbEmbeddedError::InvalidArgument( + "The endpoint did not return a cloud node location for the query.".to_owned(), + ) + })?; + + Ok(location.uri) + } } #[async_trait] @@ -76,21 +208,9 @@ impl Operations for Client { /// Returns the type of the ModelarDB node that the client is connected to. async fn modelardb_type(&mut self) -> Result { - // Retrieve the node type from the ModelarDB node. - let action = Action { - r#type: "NodeType".to_owned(), - body: vec![].into(), - }; + let bytes = self.retrieve_action_bytes("NodeType").await?; - let response = self.flight_client.do_action(Request::new(action)).await?; - - let message = response - .into_inner() - .message() - .await? - .expect("Flight message should exist."); - - ModelarDBType::from_str(str::from_utf8(&message.body)?) + ModelarDBType::from_str(str::from_utf8(&bytes)?) } /// Creates a table with the name in `table_name` and the information in `table_type`. If the @@ -118,12 +238,7 @@ impl Operations for Client { } }; - let action = Action { - r#type: "CreateTable".to_owned(), - body: protobuf_bytes.into(), - }; - - self.flight_client.do_action(action).await?; + self.send_action("CreateTable", protobuf_bytes).await?; Ok(()) } diff --git a/crates/modelardb_server/Cargo.toml b/crates/modelardb_server/Cargo.toml index 8d837053..77e82bae 100644 --- a/crates/modelardb_server/Cargo.toml +++ b/crates/modelardb_server/Cargo.toml @@ -46,6 +46,7 @@ rand.workspace = true serde.workspace = true snmalloc-rs = { workspace = true, features = ["build_cc"] } sqlparser.workspace = true +sysinfo.workspace = true tokio = { workspace = true, features = ["rt-multi-thread", "signal"] } tokio-stream.workspace = true toml.workspace = true diff --git a/crates/modelardb_server/src/cluster.rs b/crates/modelardb_server/src/cluster.rs index af97f882..d0bbfb65 100644 --- a/crates/modelardb_server/src/cluster.rs +++ b/crates/modelardb_server/src/cluster.rs @@ -36,6 +36,26 @@ use tonic::transport::Endpoint; use crate::context::Context; use crate::error::{ModelarDbServerError, Result}; +/// The different possible modes that a ModelarDB server can be deployed in, assigned when the +/// server is started. +#[derive(Clone)] +pub(crate) enum ClusterMode { + SingleNode(Node), + MultiNode(Cluster), +} + +impl ClusterMode { + /// Return all nodes in the deployment. A single node returns only itself, while a node in a + /// cluster returns every node in the cluster. If the nodes could not be retrieved, return + /// [`ModelarDbServerError`]. + pub(crate) async fn nodes(&self) -> Result> { + match self { + ClusterMode::SingleNode(node) => Ok(vec![node.clone()]), + ClusterMode::MultiNode(cluster) => cluster.nodes().await, + } + } +} + /// Stores the node that represents the local system and allows for performing operations that need /// to be applied to every peer node in the cluster. #[derive(Clone)] @@ -357,10 +377,19 @@ impl Cluster { Ok(url.to_owned()) } + /// Return all nodes currently in the cluster. If the nodes could not be retrieved, return + /// [`ModelarDbServerError`]. + pub(crate) async fn nodes(&self) -> Result> { + self.remote_data_folder + .nodes() + .await + .map_err(|error| error.into()) + } + /// Return all nodes in the cluster except the node that was saved when the [`Cluster`] was /// created. If the nodes could not be retrieved, return [`ModelarDbServerError`]. async fn peer_nodes(&self) -> Result> { - let nodes = self.remote_data_folder.nodes().await?; + let nodes = self.nodes().await?; Ok(nodes .into_iter() @@ -485,8 +514,8 @@ mod test { use modelardb_types::types::{ArrowTimestamp, ArrowValue, ErrorBound, ServerMode}; use tempfile::TempDir; + use crate::ServerArgs; use crate::data_folders::DataFolders; - use crate::{ClusterMode, ServerArgs}; // Tests for Cluster. #[tokio::test] @@ -721,7 +750,7 @@ mod test { Some(cluster.remote_data_folder.clone()), local_data_folder, ), - ClusterMode::MultiNode(Box::new(cluster)), + ClusterMode::MultiNode(cluster), &ServerArgs::parse_from(["modelardbd", "edge", "data", "s3://bucket"]), ) .await diff --git a/crates/modelardb_server/src/configuration.rs b/crates/modelardb_server/src/configuration.rs index 6b1386d7..77d37f8d 100644 --- a/crates/modelardb_server/src/configuration.rs +++ b/crates/modelardb_server/src/configuration.rs @@ -34,9 +34,10 @@ use prost::Message; use serde::{Deserialize, Serialize}; use tokio::sync::RwLock; +use crate::ServerArgs; +use crate::cluster::ClusterMode; use crate::error::{ModelarDbServerError, Result}; use crate::storage::StorageEngine; -use crate::{ClusterMode, ServerArgs}; const CONFIGURATION_FILE_NAME: &str = "modelardbd.toml"; @@ -602,9 +603,10 @@ mod tests { .await .unwrap(); + let node = Node::new("edge".to_owned(), ServerMode::Edge); let result = ConfigurationManager::try_new( local_data_folder, - ClusterMode::SingleNode, + ClusterMode::SingleNode(node), &default_args(), ) .await; @@ -625,9 +627,10 @@ mod tests { let path = temp_dir.path().join(CONFIGURATION_FILE_NAME); std::fs::write(path, "invalid_toml").unwrap(); + let node = Node::new("edge".to_owned(), ServerMode::Edge); let result = ConfigurationManager::try_new( local_data_folder, - ClusterMode::SingleNode, + ClusterMode::SingleNode(node), &default_args(), ) .await; @@ -1020,7 +1023,7 @@ mod tests { let configuration_manager = Arc::new(RwLock::new( ConfigurationManager::try_new( local_data_folder.clone(), - ClusterMode::MultiNode(Box::new(cluster)), + ClusterMode::MultiNode(cluster), &default_args(), ) .await diff --git a/crates/modelardb_server/src/context.rs b/crates/modelardb_server/src/context.rs index 3244459a..c88cd777 100644 --- a/crates/modelardb_server/src/context.rs +++ b/crates/modelardb_server/src/context.rs @@ -16,19 +16,23 @@ //! Implementation of a [`Context`] that provides access to the system's configuration and //! components. +use std::path::Path as StdPath; use std::sync::Arc; use datafusion::arrow::datatypes::Schema; use datafusion::catalog::SchemaProvider; +use modelardb_types::flight::protocol; use modelardb_types::types::TimeSeriesTableMetadata; +use sysinfo::{Disks, System}; use tokio::sync::RwLock; use tracing::{info, warn}; +use crate::cluster::ClusterMode; use crate::configuration::{ConfigurationManager, WalMode}; use crate::error::{ModelarDbServerError, Result}; use crate::storage::StorageEngine; use crate::storage::data_sinks::{NormalTableDataSink, TimeSeriesTableDataSink}; -use crate::{ClusterMode, DataFolders, ServerArgs}; +use crate::{DataFolders, ServerArgs}; /// Provides access to the system's configuration and components. pub struct Context { @@ -490,6 +494,91 @@ impl Context { Ok(schema) } + + /// Collect the current resource usage metrics of the node, including CPU, memory, disk, and + /// storage engine memory usage. + pub(crate) async fn node_metrics(&self) -> protocol::NodeMetrics { + let mut system = System::new(); + + // Refresh the CPU usage twice, separated by the minimum update interval, since a single + // refresh results in global_cpu_usage() always returning 0. + system.refresh_cpu_usage(); + tokio::time::sleep(sysinfo::MINIMUM_CPU_UPDATE_INTERVAL).await; + system.refresh_cpu_usage(); + + let cpu_usage_percentage = system.global_cpu_usage() as f64; + let cpu_count = system.cpus().len() as u32; + + system.refresh_memory(); + let used_memory_in_bytes = system.used_memory(); + let total_memory_in_bytes = system.total_memory(); + + let (used_disk_space_in_bytes, total_disk_space_in_bytes) = + self.local_data_folder_disk_space(); + + let configuration_manager = self.configuration_manager.read().await; + let storage_engine = self.storage_engine.read().await; + + let ingested_reserved_memory_in_bytes = + configuration_manager.ingested_reserved_memory_in_bytes(); + let uncompressed_reserved_memory_in_bytes = + configuration_manager.uncompressed_reserved_memory_in_bytes(); + let compressed_reserved_memory_in_bytes = + configuration_manager.compressed_reserved_memory_in_bytes(); + + // The used memory is the reserved memory minus the available remaining memory. The + // remaining memory can be negative when the reserved memory is decreased below what is + // currently in use, so treat a negative value as zero. + let ingested_used_memory_in_bytes = ingested_reserved_memory_in_bytes + - storage_engine.remaining_ingested_memory_in_bytes().max(0) as u64; + let uncompressed_used_memory_in_bytes = uncompressed_reserved_memory_in_bytes + - storage_engine + .remaining_uncompressed_memory_in_bytes() + .max(0) as u64; + let compressed_used_memory_in_bytes = compressed_reserved_memory_in_bytes + - storage_engine.remaining_compressed_memory_in_bytes().max(0) as u64; + + protocol::NodeMetrics { + cpu_usage_percentage, + cpu_count, + used_memory_in_bytes, + total_memory_in_bytes, + used_disk_space_in_bytes, + total_disk_space_in_bytes, + ingested_used_memory_in_bytes, + ingested_reserved_memory_in_bytes, + uncompressed_used_memory_in_bytes, + uncompressed_reserved_memory_in_bytes, + compressed_used_memory_in_bytes, + compressed_reserved_memory_in_bytes, + } + } + + /// Return the used and total disk space in bytes of the disk holding the local data folder. If + /// no disk holds the data folder, the disk with the most capacity is used. If no disks are + /// found, `(0, 0)` is returned. + fn local_data_folder_disk_space(&self) -> (u64, u64) { + let disks = Disks::new_with_refreshed_list(); + let location = StdPath::new(self.data_folders.local_data_folder.location()); + + // A path can sit under multiple mount points when one volume is mounted inside another, so + // pick the disk whose mount point is the longest prefix of the location, as that is the + // most specific match. If the data folder is in memory, no mount point matches, so fall + // back to the disk with the most capacity. + let maybe_disk = disks + .iter() + .filter(|disk| location.starts_with(disk.mount_point())) + .max_by_key(|disk| disk.mount_point().as_os_str().len()) + .or_else(|| disks.iter().max_by_key(|disk| disk.total_space())); + + if let Some(disk) = maybe_disk { + let total = disk.total_space(); + let used = total - disk.available_space(); + (used, total) + } else { + (0, 0) + } + } } /// Return a [`ModelarDbServerError`] indicating that a table with `table_name` does not exist. @@ -504,7 +593,7 @@ mod tests { use clap::Parser; use modelardb_storage::data_folder::DataFolder; use modelardb_test::table::{self, NORMAL_TABLE_NAME, TIME_SERIES_TABLE_NAME}; - use modelardb_types::types::MAX_RETENTION_PERIOD_IN_SECONDS; + use modelardb_types::types::{MAX_RETENTION_PERIOD_IN_SECONDS, Node, ServerMode}; use tempfile::TempDir; // Tests for Context. @@ -1162,11 +1251,12 @@ mod tests { async fn create_context(temp_dir: &TempDir) -> Arc { let temp_dir_url = temp_dir.path().to_str().unwrap(); let local_data_folder = Arc::new(DataFolder::open_local_url(temp_dir_url).await.unwrap()); + let node = Node::new("edge".to_owned(), ServerMode::Edge); Arc::new( Context::try_new( DataFolders::new(local_data_folder.clone(), None, local_data_folder), - ClusterMode::SingleNode, + ClusterMode::SingleNode(node), &ServerArgs::parse_from(["modelardbd", "edge", "data"]), ) .await diff --git a/crates/modelardb_server/src/data_folders.rs b/crates/modelardb_server/src/data_folders.rs index 43f5ea8b..f0067dcb 100644 --- a/crates/modelardb_server/src/data_folders.rs +++ b/crates/modelardb_server/src/data_folders.rs @@ -20,8 +20,8 @@ use std::sync::Arc; use modelardb_storage::data_folder::DataFolder; use modelardb_types::types::{Node, ServerMode}; -use crate::cluster::Cluster; -use crate::{ClusterMode, Result, ServerMode as ServerModeArg}; +use crate::cluster::{Cluster, ClusterMode}; +use crate::{Result, ServerMode as ServerModeArg}; /// Folders for storing metadata and data in Apache Parquet files locally and remotely. #[derive(Clone)] @@ -69,9 +69,10 @@ impl DataFolders { } => { let local_data_folder = Arc::new(DataFolder::open_local_url(local_data_folder).await?); + let node = Node::new(url_with_port, ServerMode::Edge); Ok(( - ClusterMode::SingleNode, + ClusterMode::SingleNode(node), Self::new(local_data_folder.clone(), None, local_data_folder), )) } @@ -91,7 +92,7 @@ impl DataFolders { let cluster = Cluster::try_new(node, remote_data_folder.clone()).await?; Ok(( - ClusterMode::MultiNode(Box::new(cluster)), + ClusterMode::MultiNode(cluster), Self::new( local_data_folder.clone(), Some(remote_data_folder), @@ -115,7 +116,7 @@ impl DataFolders { let cluster = Cluster::try_new(node, remote_data_folder.clone()).await?; Ok(( - ClusterMode::MultiNode(Box::new(cluster)), + ClusterMode::MultiNode(cluster), Self::new( local_data_folder, Some(remote_data_folder.clone()), @@ -149,7 +150,8 @@ mod tests { .await .unwrap(); - assert!(matches!(cluster_mode, ClusterMode::SingleNode)); + let expected_node = Node::new("grpc://127.0.0.1:9999".to_owned(), ServerMode::Edge); + assert!(matches!(cluster_mode, ClusterMode::SingleNode(node) if node == expected_node)); assert!(data_folders.maybe_remote_data_folder.is_none()); } diff --git a/crates/modelardb_server/src/main.rs b/crates/modelardb_server/src/main.rs index 673e2847..d4695afd 100644 --- a/crates/modelardb_server/src/main.rs +++ b/crates/modelardb_server/src/main.rs @@ -29,7 +29,7 @@ use clap::{Parser, Subcommand}; use modelardb_types::types::CloudCredentials; use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; -use crate::cluster::Cluster; +use crate::cluster::ClusterMode; use crate::context::Context; use crate::data_folders::DataFolders; use crate::error::Result; @@ -37,14 +37,6 @@ use crate::error::Result; #[global_allocator] static ALLOC: snmalloc_rs::SnMalloc = snmalloc_rs::SnMalloc; -/// The different possible modes that a ModelarDB server can be deployed in, assigned when the -/// server is started. -#[derive(Clone)] -pub(crate) enum ClusterMode { - SingleNode, - MultiNode(Box), -} - /// Command line arguments for the ModelarDB server. #[derive(Parser)] #[command( diff --git a/crates/modelardb_server/src/remote/mod.rs b/crates/modelardb_server/src/remote/mod.rs index 8f245963..65c776d4 100644 --- a/crates/modelardb_server/src/remote/mod.rs +++ b/crates/modelardb_server/src/remote/mod.rs @@ -59,7 +59,7 @@ use tonic::transport::{Endpoint, Server}; use tonic::{Request, Response, Status, Streaming}; use tracing::{debug, error, info}; -use crate::ClusterMode; +use crate::cluster::ClusterMode; use crate::context::Context; use crate::error::{ModelarDbServerError, Result}; use crate::remote::auth_layer::AuthLayer; @@ -83,7 +83,7 @@ pub async fn start_apache_arrow_flight_server( let maybe_cluster_key = match context.configuration_manager.read().await.cluster_mode() { ClusterMode::MultiNode(cluster) => Some(cluster.key().clone()), - ClusterMode::SingleNode => None, + ClusterMode::SingleNode(_) => None, }; let auth_layer = AuthLayer::new(maybe_authenticator, maybe_cluster_key); @@ -898,6 +898,12 @@ impl FlightService for FlightServiceHandler { /// and the change is persisted in the configuration file. /// * `NodeType`: Get the type of the node. The type is `SingleEdge`, `ClusterEdge`, or /// `ClusterCloud`. The type of the node is returned as a string. + /// * `ListNodes`: Get the nodes that are currently part of the cluster. The nodes are returned + /// in a [`ClusterNodes`](protocol::ClusterNodes) protobuf message. A single node returns only + /// itself. + /// * `NodeMetrics`: Get the current resource usage metrics of the node, including CPU, memory, + /// disk, and storage engine memory usage. The metrics are returned in a + /// [`NodeMetrics`](protocol::NodeMetrics) protobuf message. async fn do_action( &self, request: Request, @@ -1060,7 +1066,7 @@ impl FlightService for FlightServiceHandler { let configuration_manager = self.context.configuration_manager.read().await; let node_type = match configuration_manager.cluster_mode() { - ClusterMode::SingleNode => "SingleEdge", + ClusterMode::SingleNode(_) => "SingleEdge", ClusterMode::MultiNode(cluster) => match cluster.node().mode { ServerMode::Edge => "ClusterEdge", ServerMode::Cloud => "ClusterCloud", @@ -1074,6 +1080,29 @@ impl FlightService for FlightServiceHandler { Ok(Response::new(Box::pin(stream::once(async { Ok(flight_result) })))) + } else if action.r#type == "ListNodes" { + let configuration_manager = self.context.configuration_manager.read().await; + let nodes = configuration_manager + .cluster_mode() + .nodes() + .await + .map_err(error_to_status_internal)?; + + let protobuf_bytes = modelardb_types::flight::encode_and_serialize_cluster_nodes(nodes); + + Ok(Response::new(Box::pin(stream::once(async { + Ok(FlightResult { + body: protobuf_bytes.into(), + }) + })))) + } else if action.r#type == "NodeMetrics" { + let protobuf_bytes = self.context.node_metrics().await.encode_to_vec(); + + Ok(Response::new(Box::pin(stream::once(async { + Ok(FlightResult { + body: protobuf_bytes.into(), + }) + })))) } else { Err(Status::unimplemented("Action not implemented.")) } @@ -1126,6 +1155,16 @@ impl FlightService for FlightServiceHandler { description: "Get the type of the node.".to_owned(), }; + let list_nodes_action = ActionType { + r#type: "ListNodes".to_owned(), + description: "Get the nodes that are currently part of the cluster.".to_owned(), + }; + + let node_metrics_action = ActionType { + r#type: "NodeMetrics".to_owned(), + description: "Get the current resource usage metrics of the node.".to_owned(), + }; + let output = stream::iter(vec![ Ok(create_tables_action), Ok(flush_memory_action), @@ -1134,6 +1173,8 @@ impl FlightService for FlightServiceHandler { Ok(get_configuration_action), Ok(update_configuration_action), Ok(node_type_action), + Ok(list_nodes_action), + Ok(node_metrics_action), ]); Ok(Response::new(Box::pin(output))) diff --git a/crates/modelardb_server/src/storage/mod.rs b/crates/modelardb_server/src/storage/mod.rs index 19eb2718..bb914efe 100644 --- a/crates/modelardb_server/src/storage/mod.rs +++ b/crates/modelardb_server/src/storage/mod.rs @@ -357,6 +357,21 @@ impl StorageEngine { .await } + /// Return the amount of memory available for ingested data in bytes. + pub(super) fn remaining_ingested_memory_in_bytes(&self) -> i64 { + self.memory_pool.remaining_ingested_memory_in_bytes() + } + + /// Return the amount of memory available for uncompressed data in bytes. + pub(super) fn remaining_uncompressed_memory_in_bytes(&self) -> i64 { + self.memory_pool.remaining_uncompressed_memory_in_bytes() + } + + /// Return the amount of memory available for compressed data in bytes. + pub(super) fn remaining_compressed_memory_in_bytes(&self) -> i64 { + self.memory_pool.remaining_compressed_memory_in_bytes() + } + /// Mark the table with `table_name` as dropped in the data transfer component. This will prevent /// data related to the table from being transferred to the remote data folder. pub(super) async fn mark_table_as_dropped(&self, table_name: &str) { diff --git a/crates/modelardb_server/src/storage/types.rs b/crates/modelardb_server/src/storage/types.rs index 73d196cc..05c3568c 100644 --- a/crates/modelardb_server/src/storage/types.rs +++ b/crates/modelardb_server/src/storage/types.rs @@ -98,7 +98,6 @@ impl MemoryPool { } /// Return the amount of memory available for ingested data in bytes. - #[cfg(test)] #[must_use] pub(super) fn remaining_ingested_memory_in_bytes(&self) -> i64 { *self diff --git a/crates/modelardb_server/tests/integration_test.rs b/crates/modelardb_server/tests/integration_test.rs index 93855e5e..bc65c239 100644 --- a/crates/modelardb_server/tests/integration_test.rs +++ b/crates/modelardb_server/tests/integration_test.rs @@ -40,7 +40,7 @@ use futures::{StreamExt, stream}; use modelardb_test::data_generation; use modelardb_test::table::{self, NORMAL_TABLE_NAME, TIME_SERIES_TABLE_NAME}; use modelardb_types::flight::protocol; -use modelardb_types::types::ErrorBound; +use modelardb_types::types::{ErrorBound, Node, ServerMode}; use prost::Message; use tempfile::TempDir; use tokio::io::{AsyncBufReadExt, BufReader}; @@ -988,6 +988,8 @@ async fn test_can_list_actions() { "FlushNode", "GetConfiguration", "KillNode", + "ListNodes", + "NodeMetrics", "NodeType", "UpdateConfiguration", ] @@ -1612,3 +1614,53 @@ async fn test_can_create_time_series_table_from_metadata() { let retrieved_table_names = test_context.retrieve_all_table_names().await.unwrap(); assert_eq!(retrieved_table_names[0], TIME_SERIES_TABLE_NAME); } + +#[tokio::test] +async fn test_can_list_nodes() { + let mut test_context = TestContext::new().await; + let nodes_bytes = test_context.retrieve_action_bytes("ListNodes").await; + let nodes = + modelardb_types::flight::deserialize_and_extract_cluster_nodes(&nodes_bytes).unwrap(); + + assert_eq!(nodes.len(), 1); + assert_eq!( + nodes[0], + Node::new( + format!("grpc://{HOST}:{}", test_context.port), + ServerMode::Edge + ) + ); +} + +#[tokio::test] +async fn test_can_get_node_metrics() { + let mut test_context = TestContext::new().await; + let metrics_bytes = test_context.retrieve_action_bytes("NodeMetrics").await; + let metrics = protocol::NodeMetrics::decode(metrics_bytes).unwrap(); + + // Only stable fields are asserted exactly. CPU usage, used memory, and disk usage vary per run + // and per machine. + assert!(metrics.cpu_usage_percentage > 0.0); + assert!(metrics.cpu_count > 0); + + assert!(metrics.used_memory_in_bytes > 0); + assert!(metrics.total_memory_in_bytes > 0); + + assert!(metrics.used_disk_space_in_bytes > 0); + assert!(metrics.total_disk_space_in_bytes > 0); + + assert_eq!(metrics.ingested_used_memory_in_bytes, 0); + assert_eq!(metrics.ingested_reserved_memory_in_bytes, 512 * 1024 * 1024); + + assert_eq!(metrics.uncompressed_used_memory_in_bytes, 0); + assert_eq!( + metrics.uncompressed_reserved_memory_in_bytes, + 512 * 1024 * 1024 + ); + + assert_eq!(metrics.compressed_used_memory_in_bytes, 0); + assert_eq!( + metrics.compressed_reserved_memory_in_bytes, + 512 * 1024 * 1024 + ); +} diff --git a/crates/modelardb_types/src/flight/mod.rs b/crates/modelardb_types/src/flight/mod.rs index aa077450..96bf4808 100644 --- a/crates/modelardb_types/src/flight/mod.rs +++ b/crates/modelardb_types/src/flight/mod.rs @@ -17,6 +17,7 @@ //! defined in `flight/protocol.proto`. The module also provides functions to serialize and //! deserialize encoded messages to and from bytes. +use std::str::FromStr; use std::sync::Arc; use arrow::datatypes::Schema; @@ -28,7 +29,7 @@ use prost::bytes::Bytes; use crate::error::{ModelarDbTypesError, Result}; use crate::functions::{try_convert_bytes_to_schema, try_convert_schema_to_bytes}; -use crate::types::{ErrorBound, GeneratedColumn, Table, TimeSeriesTableMetadata}; +use crate::types::{ErrorBound, GeneratedColumn, Node, ServerMode, Table, TimeSeriesTableMetadata}; pub mod protocol { include!(concat!(env!("OUT_DIR"), "/modelardb.flight.protocol.rs")); @@ -215,6 +216,34 @@ fn decode_generated_column_expressions( Ok(expressions) } +/// Encode `nodes` into a [`ClusterNodes`](protocol::ClusterNodes) protobuf message and serialize it. +pub fn encode_and_serialize_cluster_nodes(nodes: Vec) -> Vec { + let cluster_nodes = protocol::ClusterNodes { + nodes: nodes + .into_iter() + .map(|node| protocol::NodeMetadata { + url: node.url, + mode: node.mode.to_string(), + }) + .collect(), + }; + + cluster_nodes.encode_to_vec() +} + +/// Deserialize `bytes` into a [`ClusterNodes`](protocol::ClusterNodes) protobuf message and extract +/// a vector of [`Node`]. If `bytes` cannot be deserialized or a node has an invalid mode, return +/// [`ModelarDbTypesError`]. +pub fn deserialize_and_extract_cluster_nodes(bytes: &[u8]) -> Result> { + let cluster_nodes = protocol::ClusterNodes::decode(bytes)?; + + cluster_nodes + .nodes + .into_iter() + .map(|node| Ok(Node::new(node.url, ServerMode::from_str(&node.mode)?))) + .collect() +} + #[cfg(test)] mod test { use super::*; @@ -259,4 +288,18 @@ mod test { _ => panic!("Expected time series table."), } } + + // Test for serializing and deserializing cluster nodes. + #[test] + fn test_serialize_and_deserialize_cluster_nodes() { + let expected_nodes = vec![ + Node::new("grpc://127.0.0.1:9999".to_owned(), ServerMode::Edge), + Node::new("grpc://127.0.0.1:9998".to_owned(), ServerMode::Cloud), + ]; + + let bytes = encode_and_serialize_cluster_nodes(expected_nodes.clone()); + let nodes = deserialize_and_extract_cluster_nodes(&bytes).unwrap(); + + assert_eq!(nodes, expected_nodes); + } } diff --git a/crates/modelardb_types/src/flight/protocol.proto b/crates/modelardb_types/src/flight/protocol.proto index 141816e3..8f7dbd64 100644 --- a/crates/modelardb_types/src/flight/protocol.proto +++ b/crates/modelardb_types/src/flight/protocol.proto @@ -108,3 +108,56 @@ message UpdateConfiguration { // New value for the setting. optional uint64 new_value = 2; } + +// Metadata identifying a single node in a ModelarDB cluster. +message NodeMetadata { + // Apache Arrow Flight URL that uniquely identifies the node. + string url = 1; + + // Mode the node was started in, either "edge" or "cloud". + string mode = 2; +} + +// The nodes that are currently part of a ModelarDB cluster. +message ClusterNodes { + repeated NodeMetadata nodes = 1; +} + +// Resource usage metrics for a single ModelarDB node. +message NodeMetrics { + // Percentage of the CPU that is currently in use across all cores. + double cpu_usage_percentage = 1; + + // Number of logical CPUs available to the node. + uint32 cpu_count = 2; + + // Amount of memory currently in use on the node. + uint64 used_memory_in_bytes = 3; + + // Total amount of memory on the node. + uint64 total_memory_in_bytes = 4; + + // Amount of disk space currently in use on the disk holding the local data folder. + uint64 used_disk_space_in_bytes = 5; + + // Total amount of disk space on the disk holding the local data folder. + uint64 total_disk_space_in_bytes = 6; + + // Amount of memory currently in use for storing ingested time series. + uint64 ingested_used_memory_in_bytes = 7; + + // Total amount of memory reserved for storing ingested time series. + uint64 ingested_reserved_memory_in_bytes = 8; + + // Amount of memory currently in use for storing uncompressed data buffers. + uint64 uncompressed_used_memory_in_bytes = 9; + + // Total amount of memory reserved for storing uncompressed data buffers. + uint64 uncompressed_reserved_memory_in_bytes = 10; + + // Amount of memory currently in use for storing compressed data buffers. + uint64 compressed_used_memory_in_bytes = 11; + + // Total amount of memory reserved for storing compressed data buffers. + uint64 compressed_reserved_memory_in_bytes = 12; +}