rust(feat): Add test server to CLI

Author: marc-sift

rust/crates/sift_cli/Cargo.toml

@@ -23,13 +23,24 @@ flate2 = "1.1.2"
 indicatif = "0.18.0"
 parquet = "56.2.0"
 pbjson-types = { workspace = true }
+prost = "0.13.5"
 reqwest = "0.12.23"
 sift_pbfs = { workspace = true }
 sift_rs = { workspace = true }
-tokio = { version = "1.47.1", features = ["full", "net", "time"] }
+sift_stream = { workspace = true }
+tokio = { version = "1.47.1", features = ["full", "net", "time", "macros", "rt-multi-thread", "signal"] }
 tokio-stream = "0.1.17"
+tonic = { workspace = true }
+tonic-reflection = "0.12"
 toml = "0.8.23"
 zip = "6.0.0"
 
 [dev-dependencies]
 indoc = "2.0.6"
+
+[dependencies.uuid]
+version = "1.19.0"
+features = ["v4"]
+
+[build-dependencies]
+tonic-build = "0.12"

rust/crates/sift_cli/build.rs

@@ -0,0 +1,21 @@
+/// Build descriptor's so that the Black Hole gRPC server can
+/// stand up the reflection service.
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    tonic_build::configure()
+        .file_descriptor_set_path("descriptor.bin")
+        .build_client(false)
+        .build_server(false)
+        .compile_protos(
+            &[
+                "/tmp/exported-protos/sift/assets/v1/assets.proto",
+                "/tmp/exported-protos/sift/ping/v1/ping.proto",
+                "/tmp/exported-protos/sift/ingest/v1/ingest.proto",
+                "/tmp/exported-protos/sift/ingestion_configs/v2/ingestion_configs.proto",
+            ],
+            // Run the following command to generate exported-protos:
+            // buf export ../../../protos --output /tmp/exported-protos
+            &["/tmp/exported-protos"],
+        )?;
+
+    Ok(())
+}

rust/crates/sift_cli/src/cli/mod.rs

@@ -48,6 +48,10 @@ pub enum Cmd {
     /// Import time series files into Sift
     #[command(subcommand)]
     Import(ImportCmd),
+
+    /// Start a test gRPC server for streaming.
+    #[command(subcommand)]
+    TestServer(TestServerCmd),
 }
 
 #[derive(Subcommand)]
@@ -164,6 +168,27 @@ pub enum ConfigCmd {
     Update(ConfigUpdateArgs),
 }
 
+#[derive(Subcommand)]
+pub enum TestServerCmd {
+    /// Start a test ingestion server.
+    Run(TestServerArgs),
+}
+
+#[derive(clap::Args)]
+pub struct TestServerArgs {
+    /// The address to serve gRPC server. Default is 127.0.0.1:50051.
+    #[arg(short, long)]
+    pub local_address: Option<String>,
+
+    /// Whether to stream metrics to Sift.
+    #[arg(short, long)]
+    pub stream_metrics: Option<bool>,
+
+    /// The asset name to use when streaming server ingestion metrics.
+    #[arg(short, long)]
+    pub metrics_asset_name: Option<String>,
+}
+
 #[derive(clap::Args)]
 pub struct ConfigUpdateArgs {
     /// Edit or create a profile interactively (ignores other flags)

rust/crates/sift_cli/src/cmd/mod.rs

@@ -8,6 +8,7 @@ pub mod completions;
 pub mod config;
 pub mod export;
 pub mod import;
+pub mod test_server;
 
 pub struct Context {
     pub grpc_uri: String,

rust/crates/sift_cli/src/cmd/test_server/metrics_streaming_client.rs

@@ -0,0 +1,127 @@
+use super::Context;
+use anyhow::{Ok, anyhow};
+use crossterm::style::Stylize;
+use sift_stream::{
+    ChannelConfig, ChannelDataType, ChannelValue, Credentials, Flow, FlowBuilder, FlowConfig,
+    IngestionConfigEncoder, IngestionConfigForm, RecoveryStrategy, RetryPolicy, RunForm,
+    SiftStream, SiftStreamBuilder, TimeValue,
+};
+
+/// Streams metrics to Sift.
+pub struct MetricsStreamingClient {
+    ctx: Context,
+    asset_name: String,
+    sift_stream: Option<SiftStream<IngestionConfigEncoder>>,
+}
+
+impl MetricsStreamingClient {
+    pub fn build(
+        ctx: Context,
+        stream_metrics: &Option<bool>,
+        asset_name: &Option<String>,
+    ) -> Result<Option<MetricsStreamingClient>, anyhow::Error> {
+        if !stream_metrics.unwrap_or(false) {
+            return Ok(None);
+        }
+
+        let Some(asset_name) = asset_name else {
+            return Err(anyhow!(
+                "must specify {} with streaming enabled",
+                "--metrics_asset_name".cyan()
+            ));
+        };
+
+        Ok(Some(MetricsStreamingClient {
+            ctx: ctx,
+            asset_name: asset_name.clone(),
+            sift_stream: None,
+        }))
+    }
+
+    /// Initialize SiftStream and create ingestion config.
+    pub async fn initialize(&mut self) -> Result<(), anyhow::Error> {
+        let credentials = Credentials::Config {
+            apikey: self.ctx.api_key.clone(),
+            uri: self.ctx.grpc_uri.clone(),
+        };
+
+        let ingestion_config = IngestionConfigForm {
+            asset_name: self.asset_name.to_string(),
+            client_key: "stress-test-ingestion-config-test".into(),
+            flows: vec![FlowConfig {
+                name: "metrics".into(),
+                channels: vec![
+                    ChannelConfig {
+                        name: "total_num_streams".into(),
+                        description: "Total number of streams created".into(),
+                        data_type: ChannelDataType::Uint32.into(),
+                        ..Default::default()
+                    },
+                    ChannelConfig {
+                        name: "total_num_bytes_read".into(),
+                        description: "Total number of bytes read".into(),
+                        unit: "B".into(),
+                        data_type: ChannelDataType::Uint64.into(),
+                        ..Default::default()
+                    },
+                    ChannelConfig {
+                        name: "total_num_messages".into(),
+                        description: "Total number of messages received".into(),
+                        unit: "message".into(),
+                        data_type: ChannelDataType::Uint64.into(),
+                        ..Default::default()
+                    },
+                    ChannelConfig {
+                        name: "bytes_per_s".into(),
+                        description: "Number of bytes received per second".into(),
+                        data_type: ChannelDataType::Double.into(),
+                        unit: "B/s".into(),
+                        ..Default::default()
+                    },
+                    ChannelConfig {
+                        name: "messages_per_s".into(),
+                        description: "Number of messages received per second".into(),
+                        unit: "message/s".into(),
+                        data_type: ChannelDataType::Double.into(),
+                        ..Default::default()
+                    },
+                ],
+            }],
+        };
+
+        let sift_stream = SiftStreamBuilder::new(credentials)
+            .ingestion_config(ingestion_config)
+            .recovery_strategy(RecoveryStrategy::RetryOnly(RetryPolicy::default()))
+            .build()
+            .await?;
+
+        self.sift_stream = Some(sift_stream);
+
+        Ok(())
+    }
+
+    /// Send metrics to Sift.
+    pub async fn ingest(&mut self, metrics: Metrics) {
+        let flow = Flow::new(
+            "metrics",
+            TimeValue::now(),
+            &[
+                ChannelValue::new("total_num_streams", metrics.total_num_streams),
+                ChannelValue::new("total_num_bytes_read", metrics.total_num_bytes_read),
+                ChannelValue::new("total_num_messages", metrics.total_num_messages),
+                ChannelValue::new("bytes_per_s", metrics.bytes_per_s),
+                ChannelValue::new("messages_per_s", metrics.messages_per_s),
+            ],
+        );
+
+        self.sift_stream.as_mut().unwrap().send(flow).await.unwrap();
+    }
+}
+
+pub struct Metrics {
+    pub total_num_streams: u32,
+    pub total_num_bytes_read: u64,
+    pub total_num_messages: u64,
+    pub bytes_per_s: f64,
+    pub messages_per_s: f64,
+}

rust/crates/sift_cli/src/cmd/test_server/mod.rs

@@ -0,0 +1,101 @@
+use super::Context;
+use crate::cmd::test_server::metrics_streaming_client::MetricsStreamingClient;
+use crate::{cli::TestServerArgs, util::tty::Output};
+use anyhow::Result;
+use server::TestServer;
+use sift_rs::assets::v1::asset_service_server::AssetServiceServer;
+use sift_rs::ingest::v1::ingest_service_server::IngestServiceServer;
+use sift_rs::ingestion_configs::v2::ingestion_config_service_server::IngestionConfigServiceServer;
+use sift_rs::ping::v1::ping_service_server::PingServiceServer;
+use std::process::ExitCode;
+use std::sync::Arc;
+use tokio::sync::mpsc;
+use tokio::sync::watch;
+use tonic::transport::Server;
+use tonic_reflection::server::Builder;
+
+pub const FILE_DESCRIPTOR_SET: &[u8] = include_bytes!("../../../descriptor.bin");
+pub mod metrics_streaming_client;
+pub mod server;
+use crate::cmd::test_server::metrics_streaming_client::Metrics;
+
+pub async fn run(ctx: Context, args: TestServerArgs) -> Result<ExitCode> {
+    let local_address = args.local_address.unwrap_or("127.0.0.1:50051".into());
+    let addr = local_address.parse()?;
+
+    // Initialize streaming client.
+    let mut streaming_client =
+        MetricsStreamingClient::build(ctx, &args.stream_metrics, &args.metrics_asset_name)?;
+    if streaming_client.is_some() {
+        streaming_client.as_mut().unwrap().initialize().await?;
+    }
+
+    // Channel to signal program exit.
+    let (shutdown_tx, mut shutdown_rx) = watch::channel(false);
+    let mut shutdown_rx2 = shutdown_rx.clone();
+
+    // Channel to send metrics.
+    let (metrics_tx, mut metrics_rx) = mpsc::channel::<Metrics>(1024);
+
+    // Initialize gRPC server.
+    let server = Arc::new(TestServer::default());
+
+    // Start task to calculate ingestion metrics.
+    let server_arc = Arc::clone(&server);
+    let calc_stats_task = tokio::spawn(async move {
+        server_arc
+            .calculate_metrics(
+                &mut shutdown_rx,
+                metrics_tx,
+                args.stream_metrics.unwrap_or(false),
+            )
+            .await;
+    });
+
+    // Start task to ingest metrics to Sift.
+    let ingest_metrics_task = tokio::spawn(async move {
+        if streaming_client.is_none() {
+            return;
+        }
+
+        let mut client = streaming_client.unwrap();
+        loop {
+            tokio::select! {
+                _ = shutdown_rx2.changed() => {
+                    Output::new().line("Ingest task shutting down").print();
+                    break;
+                }
+                Some(metrics) = metrics_rx.recv() => {
+                    client.ingest(metrics).await;
+                }
+            };
+        }
+    });
+
+    let reflection_service = Builder::configure()
+        .register_encoded_file_descriptor_set(FILE_DESCRIPTOR_SET)
+        .build_v1()?;
+
+    Output::new()
+        .line(format!("Server listening on {addr}"))
+        .print();
+
+    Server::builder()
+        .add_service(reflection_service)
+        .add_service(PingServiceServer::from_arc(server.clone()))
+        .add_service(IngestServiceServer::from_arc(server.clone()))
+        .add_service(IngestionConfigServiceServer::from_arc(server.clone()))
+        .add_service(AssetServiceServer::from_arc(server.clone()))
+        .serve_with_shutdown(addr, async move {
+            tokio::signal::ctrl_c().await.unwrap();
+            let _ = shutdown_tx.send(true);
+        })
+        .await?;
+
+    calc_stats_task.await?;
+    ingest_metrics_task.await?;
+
+    Output::new().line("Exiting.").print();
+
+    Ok(ExitCode::SUCCESS)
+}