villebro commented on code in PR #2577: URL: https://github.com/apache/datafusion-ballista/pull/2577#discussion_r4249316775
########## ballista/result-service/src/main.rs: ########## @@ -0,0 +1,251 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Ballista Result Service: serves query results to clients by forwarding each +//! partition fetch to the executor that holds the partition. Point the +//! scheduler's `--advertise-flight-endpoint` at it. + +mod http; +mod metered; +mod metrics; + +use std::future::Future; +use std::net::SocketAddr; +use std::sync::Arc; +use std::time::Duration; + +use arrow_flight::flight_service_server::FlightServiceServer; +use ballista_core::serving::{ForwardingBackend, ServingFlightService}; +use ballista_core::utils::{GrpcServerConfig, create_grpc_server}; +use clap::Parser; +use log::{error, info, warn}; +use tonic::service::Routes; +use tonic::transport::server::Router; + +use crate::metered::MeteredBackend; +use crate::metrics::{ResultMetricsCollector, default_metrics_collector}; + +/// Command-line configuration for the result service. +#[derive(Debug, Parser)] +#[command(name = "ballista-result-service", version, about)] +struct Config { + /// Host or IP address the Flight service binds to. + #[arg(long, default_value = "0.0.0.0")] + bind_host: String, + + /// Port the Flight service, health probes (/healthz and /readyz) and + /// metrics (/api/metrics) bind to. + #[arg(long, default_value_t = 50055)] + bind_port: u16, + + /// Use TLS when connecting to executors. + #[arg(long, default_value_t = false)] + use_tls: bool, Review Comment: Good catch, thanks. The stock scheduler binary has no executor TLS option for the embedded proxy either. `SchedulerConfig::use_tls` is library-only, and real TLS needs a build with tonic's TLS features, plus an endpoint override for private CAs or mTLS, as in the mTLS example. So rather than ship a flag that looks like it works, d39f49ca drops `--use-tls`, along with the `ca-certificates` install that only existed for it. The binary connects to executors in plaintext, like the stock scheduler, and the docs say so. TLS for the binary is listed as a follow-up in the description: system roots and private CAs, a client certificate and key, a server-name override, and tests against plaintext, TLS and mTLS executors. ########## ballista/result-service/src/main.rs: ########## @@ -0,0 +1,251 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Ballista Result Service: serves query results to clients by forwarding each +//! partition fetch to the executor that holds the partition. Point the +//! scheduler's `--advertise-flight-endpoint` at it. + +mod http; +mod metered; +mod metrics; + +use std::future::Future; +use std::net::SocketAddr; +use std::sync::Arc; +use std::time::Duration; + +use arrow_flight::flight_service_server::FlightServiceServer; +use ballista_core::serving::{ForwardingBackend, ServingFlightService}; +use ballista_core::utils::{GrpcServerConfig, create_grpc_server}; +use clap::Parser; +use log::{error, info, warn}; +use tonic::service::Routes; +use tonic::transport::server::Router; + +use crate::metered::MeteredBackend; +use crate::metrics::{ResultMetricsCollector, default_metrics_collector}; + +/// Command-line configuration for the result service. +#[derive(Debug, Parser)] +#[command(name = "ballista-result-service", version, about)] +struct Config { + /// Host or IP address the Flight service binds to. + #[arg(long, default_value = "0.0.0.0")] + bind_host: String, + + /// Port the Flight service, health probes (/healthz and /readyz) and + /// metrics (/api/metrics) bind to. + #[arg(long, default_value_t = 50055)] + bind_port: u16, + + /// Use TLS when connecting to executors. + #[arg(long, default_value_t = false)] + use_tls: bool, + + /// Maximum gRPC message size, in bytes, this service will decode. + #[arg(long, default_value_t = 16777216)] + grpc_max_decoding_message_size: usize, + + /// Maximum gRPC message size, in bytes, this service will encode. + #[arg(long, default_value_t = 16777216)] + grpc_max_encoding_message_size: usize, + + /// Time in seconds to let in-flight result streams finish after a + /// shutdown signal (SIGTERM or Ctrl-C) before exiting anyway. Keep it below + /// your orchestrator's termination grace period (Kubernetes' + /// terminationGracePeriodSeconds, 30 by default), so this bound is reached + /// before a SIGKILL is. + #[arg(long, default_value_t = 10)] + graceful_shutdown_timeout_seconds: u64, +} + +#[tokio::main] +async fn main() -> Result<(), Box<dyn std::error::Error>> { + env_logger::init(); + let config = Config::parse(); + + let addr: SocketAddr = format!("{}:{}", config.bind_host, config.bind_port) + .parse() + .map_err(|e| { + format!( + "invalid bind address {}:{}: {e}", + config.bind_host, config.bind_port + ) + })?; + + let server = server(&config, default_metrics_collector()?); + + info!("Ballista Result Service listening on {addr}"); + + let (stop_tx, stop_rx) = tokio::sync::oneshot::channel::<()>(); + let server = server.serve_with_shutdown(addr, async { + let _ = stop_rx.await; + }); + tokio::pin!(server); + + // Serve until a shutdown signal arrives, unless the server fails first. + tokio::select! { + result = &mut server => { + result?; + info!("Ballista Result Service stopped"); + return Ok(()); + } + () = shutdown_signal() => {} + } + + // Stop accepting connections, then give in-flight streams a bounded time to + // finish. + let _ = stop_tx.send(()); + drain_or_time_out( + server, + Duration::from_secs(config.graceful_shutdown_timeout_seconds), + ) + .await?; + + info!("Ballista Result Service stopped"); + Ok(()) +} + +/// The Flight service, health probes and metrics, served on one port. +fn server(config: &Config, metrics: Arc<dyn ResultMetricsCollector>) -> Router { + let backend = MeteredBackend::new( + ForwardingBackend::new( + config.grpc_max_decoding_message_size, + config.grpc_max_encoding_message_size, + config.use_tls, + None, Review Comment: Agreed, that was a gap for embedders. 8e149961 adds a small library target: `ballista_result_service::serve(ResultServiceConfig, shutdown)`. Its config takes the same `use_tls` and `override_create_grpc_client_endpoint` settings as `SchedulerConfig`, so embedders that customize the proxy's executor connections, for example for mTLS, can move to the result service with the same settings. As with the proxy, `use_tls` only encrypts in a build with tonic's TLS features. The 56.0.0 upgrade guide now documents this. ########## chaos-testing/src/bin/chaos-result-service.rs: ########## @@ -0,0 +1,59 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! result service for the chaos harness. It is configured from the environment +//! because `TestCluster` spawns it as a child process, like `chaos-scheduler` +//! and `chaos-executor`. + +use arrow_flight::flight_service_server::FlightServiceServer; +use ballista_core::serving::{ForwardingBackend, ServingFlightService}; +use ballista_core::utils::{GrpcServerConfig, create_grpc_server}; +use std::net::SocketAddr; + +const MAX_MESSAGE_SIZE: usize = 16 * 1024 * 1024; + +#[tokio::main] +async fn main() -> ballista_core::error::Result<()> { + env_logger::init(); + + let bind_host = + std::env::var("CHAOS_BIND_HOST").unwrap_or_else(|_| "127.0.0.1".into()); + let bind_port: u16 = std::env::var("CHAOS_RESULT_SERVICE_PORT") + .expect("CHAOS_RESULT_SERVICE_PORT must be set") + .parse() + .expect("CHAOS_RESULT_SERVICE_PORT must be a u16"); + + let addr: SocketAddr = format!("{bind_host}:{bind_port}") + .parse() + .expect("result service address must parse"); + + // Plaintext: the harness runs everything on loopback. + let backend = ForwardingBackend::new(MAX_MESSAGE_SIZE, MAX_MESSAGE_SIZE, false, None); Review Comment: Agreed. With the library from 8e149961, `chaos-result-service` is a thin wrapper around the shipped `serve`, configured from the environment, the same way `chaos-scheduler` calls the scheduler's `start_server`. It now runs the binary's code for the shared HTTP/gRPC listener, metering, health endpoints, forwarding and shutdown, and the k8s readiness probe hits `/readyz`. The drain itself is covered by a new unit test; see the drain thread. ########## ballista/result-service/src/metrics.rs: ########## @@ -0,0 +1,228 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Metrics for result fetches. + +use std::sync::Arc; +use std::time::Duration; + +use ballista_core::error::Result; + +/// How a result fetch ended. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum FetchOutcome { + /// The whole stream was delivered. + Completed, + /// The executor or the stream returned an error. + Failed, + /// The client went away before the stream ended. + Cancelled, +} + +/// Records result fetches. +pub trait ResultMetricsCollector: Send + Sync { + /// A fetch started. + fn record_started(&self); + + /// A fetch relayed `bytes` of Flight data to its client. + fn record_bytes(&self, bytes: u64); + + /// A fetch ended after `elapsed`. + fn record_finished(&self, outcome: FetchOutcome, elapsed: Duration); + + /// The metrics and their content type, or `None` if this collector exports + /// none. + fn gather_metrics(&self) -> Result<Option<(Vec<u8>, String)>>; +} + +/// A collector that records nothing. +#[cfg(any(test, not(feature = "prometheus-metrics")))] +#[derive(Default)] +pub struct NoopMetricsCollector; + +#[cfg(any(test, not(feature = "prometheus-metrics")))] +impl ResultMetricsCollector for NoopMetricsCollector { + fn record_started(&self) {} + + fn record_bytes(&self, _bytes: u64) {} + + fn record_finished(&self, _outcome: FetchOutcome, _elapsed: Duration) {} + + fn gather_metrics(&self) -> Result<Option<(Vec<u8>, String)>> { + Ok(None) + } +} + +/// The Prometheus collector with the `prometheus-metrics` feature, otherwise +/// one that records nothing. +pub fn default_metrics_collector() -> Result<Arc<dyn ResultMetricsCollector>> { + #[cfg(feature = "prometheus-metrics")] + { + Ok(Arc::new(PrometheusMetricsCollector::new( + ::prometheus::default_registry(), + )?)) + } + #[cfg(not(feature = "prometheus-metrics"))] + { + Ok(Arc::new(NoopMetricsCollector)) + } +} + +#[cfg(feature = "prometheus-metrics")] +pub use prometheus_metrics::PrometheusMetricsCollector; + +#[cfg(feature = "prometheus-metrics")] +mod prometheus_metrics { + use super::{FetchOutcome, ResultMetricsCollector}; + use ballista_core::error::{BallistaError, Result}; + use prometheus::{ + Counter, Encoder, Gauge, Histogram, Registry, TextEncoder, + register_counter_with_registry, register_gauge_with_registry, + register_histogram_with_registry, + }; + use std::time::Duration; + + /// A [`ResultMetricsCollector`] that exports Prometheus metrics: + /// + /// - `fetch_started_total`, `fetch_completed_total`, `fetch_failed_total` + /// and `fetch_cancelled_total` count fetches by how they ended + /// - `fetch_in_flight` is the number of fetches currently streaming + /// - `fetch_time_seconds` is a histogram of fetch durations + /// - `fetch_bytes_total` counts the Flight data relayed to clients + pub struct PrometheusMetricsCollector { + registry: Registry, + started: Counter, + completed: Counter, + failed: Counter, + cancelled: Counter, + in_flight: Gauge, + time: Histogram, + bytes: Counter, + } + + fn registration_error(e: prometheus::Error) -> BallistaError { + BallistaError::Internal(format!("Error registering metric: {e:?}")) + } + + impl PrometheusMetricsCollector { + /// Creates the collector, registering its metrics in `registry`. + pub fn new(registry: &Registry) -> Result<Self> { + let counter = |name: &str, help: &str| { + register_counter_with_registry!(name, help, registry) + .map_err(registration_error) + }; + Ok(Self { + registry: registry.clone(), + started: counter( + "fetch_started_total", Review Comment: Good call, now is the cheapest time. ebc2839f renames them to: - `ballista_result_service_fetches_total{outcome="completed|failed|cancelled"}` - `ballista_result_service_fetches_in_flight` - `ballista_result_service_fetch_duration_seconds` - `ballista_result_service_relayed_bytes_total` It also drops the separate started counter, since starts can be derived from the outcome counts plus the in-flight gauge. Each outcome series is exported from startup. The same commit adds `ballista-result-service/prometheus-metrics` to the Linux test job, so the metrics tests run in CI. ########## ballista/result-service/src/main.rs: ########## @@ -0,0 +1,251 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Ballista Result Service: serves query results to clients by forwarding each +//! partition fetch to the executor that holds the partition. Point the +//! scheduler's `--advertise-flight-endpoint` at it. + +mod http; +mod metered; +mod metrics; + +use std::future::Future; +use std::net::SocketAddr; +use std::sync::Arc; +use std::time::Duration; + +use arrow_flight::flight_service_server::FlightServiceServer; +use ballista_core::serving::{ForwardingBackend, ServingFlightService}; +use ballista_core::utils::{GrpcServerConfig, create_grpc_server}; +use clap::Parser; +use log::{error, info, warn}; +use tonic::service::Routes; +use tonic::transport::server::Router; + +use crate::metered::MeteredBackend; +use crate::metrics::{ResultMetricsCollector, default_metrics_collector}; + +/// Command-line configuration for the result service. +#[derive(Debug, Parser)] +#[command(name = "ballista-result-service", version, about)] +struct Config { + /// Host or IP address the Flight service binds to. + #[arg(long, default_value = "0.0.0.0")] + bind_host: String, + + /// Port the Flight service, health probes (/healthz and /readyz) and + /// metrics (/api/metrics) bind to. + #[arg(long, default_value_t = 50055)] + bind_port: u16, + + /// Use TLS when connecting to executors. + #[arg(long, default_value_t = false)] + use_tls: bool, + + /// Maximum gRPC message size, in bytes, this service will decode. + #[arg(long, default_value_t = 16777216)] + grpc_max_decoding_message_size: usize, + + /// Maximum gRPC message size, in bytes, this service will encode. + #[arg(long, default_value_t = 16777216)] + grpc_max_encoding_message_size: usize, + + /// Time in seconds to let in-flight result streams finish after a + /// shutdown signal (SIGTERM or Ctrl-C) before exiting anyway. Keep it below + /// your orchestrator's termination grace period (Kubernetes' + /// terminationGracePeriodSeconds, 30 by default), so this bound is reached + /// before a SIGKILL is. + #[arg(long, default_value_t = 10)] + graceful_shutdown_timeout_seconds: u64, +} + +#[tokio::main] +async fn main() -> Result<(), Box<dyn std::error::Error>> { + env_logger::init(); + let config = Config::parse(); + + let addr: SocketAddr = format!("{}:{}", config.bind_host, config.bind_port) + .parse() + .map_err(|e| { + format!( + "invalid bind address {}:{}: {e}", + config.bind_host, config.bind_port + ) + })?; + + let server = server(&config, default_metrics_collector()?); + + info!("Ballista Result Service listening on {addr}"); + + let (stop_tx, stop_rx) = tokio::sync::oneshot::channel::<()>(); + let server = server.serve_with_shutdown(addr, async { + let _ = stop_rx.await; + }); + tokio::pin!(server); + + // Serve until a shutdown signal arrives, unless the server fails first. + tokio::select! { + result = &mut server => { + result?; + info!("Ballista Result Service stopped"); + return Ok(()); + } + () = shutdown_signal() => {} + } + + // Stop accepting connections, then give in-flight streams a bounded time to + // finish. + let _ = stop_tx.send(()); Review Comment: Thanks for testing this. I've fixed it rather than soften the docs. 3d1a8c04 serves from an incoming stream that ends when shutdown begins. That drops the listener, so new connections are refused. tonic's own shutdown signal never fires, so the connections it has already accepted still drain, within `graceful_shutdown_timeout`. A new test checks both sides: a connection attempted during the drain is refused, while an in-flight result stream completes. ########## ballista/scheduler/src/config.rs: ########## @@ -364,6 +367,11 @@ pub struct SchedulerConfig { /// *told*. When a non-empty endpoint is also set it takes precedence for /// the advertisement, which is how you put a load balancer in front of an /// embedded proxy. + /// + /// Deprecated, and to be removed in 57.0.0: the proxy puts result traffic + /// on the scheduler. Run a result service and advertise it with + /// [`Self::advertise_flight_endpoint`] instead. The scheduler logs a warning + /// at startup while this is set. pub enable_embedded_flight_proxy: bool, Review Comment: Agreed. 6d10de46 deprecates the `SchedulerConfig::enable_embedded_flight_proxy` field too, alongside the builder method. The scheduler's own uses and `chaos-scheduler.rs` get narrowly scoped `#[allow(deprecated)]`. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
