krabka_client_core/
telemetry.rs

1//! Client metrics push (KIP-714).
2//!
3//! When Kafka's `enable.metrics.push` is on (the default), a producer,
4//! consumer or admin client asks a broker for its client metrics
5//! subscription with `GetTelemetrySubscriptions`, and pushes the metrics that
6//! the subscription names with `PushTelemetry` at the push interval. On close
7//! it sends one last push with `terminating` set.
8//!
9//! [`ClientTelemetry`] is that reporter, as Kafka's `ClientTelemetryReporter`
10//! and the telemetry sender of its `NetworkClient` are. A [`Client`] built
11//! with [`ClientTelemetryConfig`] starts one. The metrics come from a
12//! [`ClientMetrics`] registry; client-core registers the network metrics of
13//! Kafka's `Selector` in it.
14//!
15//! Differences from Kafka:
16//!
17//! - The reporter sends on a connection that the client already has open. It
18//!   does not open a connection of its own, so a client that has no open
19//!   connection pushes nothing until it opens one.
20//! - A subscription with the zero client instance id stops telemetry. Kafka's
21//!   reporter throws from its response handler, after which it sends no
22//!   further request.
23//!
24//! [`Client`]: crate::Client
25
26mod metrics;
27pub mod otlp;
28mod sender;
29
30use std::{sync::Arc, time::Duration};
31
32use krabka_protocol::ProtocolRequest;
33use krabka_units::{Time, convert::TimeExt as _};
34use tokio::{
35    sync::watch,
36    time::{self, Instant, sleep},
37};
38use tokio_util::sync::CancellationToken;
39
40pub(crate) use self::metrics::ConnectionMetrics;
41pub use self::metrics::{
42    ClientMetrics, Counter, Gauge, MetricKey, NetworkMetrics, TELEMETRY_METRIC_DOMAIN,
43};
44use self::{
45    metrics::{Collector, unix_nanos_now},
46    otlp::{KeyValue, Resource},
47    sender::{Failure, TelemetryRequest, TelemetrySender, TelemetryState},
48};
49use crate::{backoff, connection::Connection, error::ClientError};
50
51/// The kind of client, which names its metric group.
52#[derive(Clone, Copy, Debug, PartialEq, Eq)]
53pub enum TelemetryClientType {
54    Producer,
55    Consumer,
56    Admin,
57}
58
59impl TelemetryClientType {
60    /// The group of the client's own metrics, such as the network metrics:
61    /// `producer-metrics`, `consumer-metrics` or `admin-client-metrics`.
62    #[must_use]
63    pub const fn metric_group(self) -> &'static str {
64        match self {
65            Self::Producer => "producer-metrics",
66            Self::Consumer => "consumer-metrics",
67            Self::Admin => "admin-client-metrics",
68        }
69    }
70}
71
72/// What a client pushes: its type and the resource attributes of every
73/// metric (`ClientTelemetryProvider.contextChange`).
74#[derive(Clone, Debug, PartialEq, Eq)]
75pub struct ClientTelemetryConfig {
76    pub client_type: TelemetryClientType,
77    /// The attributes of the OTLP resource, in order.
78    pub resource_attributes: Vec<(String, String)>,
79}
80
81impl ClientTelemetryConfig {
82    /// A producer. Kafka adds `transactional_id` when it is set.
83    #[must_use]
84    pub fn producer(transactional_id: Option<&str>) -> Self {
85        Self::with_attributes(
86            TelemetryClientType::Producer,
87            [("transactional_id", transactional_id)],
88        )
89    }
90
91    /// A consumer. Kafka adds `group_id`, `group_instance_id` and
92    /// `client_rack` when they are set.
93    #[must_use]
94    pub fn consumer(
95        group_id: Option<&str>,
96        group_instance_id: Option<&str>,
97        client_rack: Option<&str>,
98    ) -> Self {
99        Self::with_attributes(
100            TelemetryClientType::Consumer,
101            [
102                ("group_id", group_id),
103                ("group_instance_id", group_instance_id),
104                ("client_rack", client_rack),
105            ],
106        )
107    }
108
109    /// An admin client.
110    #[must_use]
111    pub fn admin() -> Self {
112        Self::with_attributes(TelemetryClientType::Admin, [])
113    }
114
115    fn with_attributes<'a>(
116        client_type: TelemetryClientType,
117        attributes: impl IntoIterator<Item = (&'a str, Option<&'a str>)>,
118    ) -> Self {
119        Self {
120            client_type,
121            resource_attributes: attributes
122                .into_iter()
123                .filter_map(|(key, value)| {
124                    value
125                        .filter(|value| !value.is_empty())
126                        .map(|value| (key.to_owned(), value.to_owned()))
127                })
128                .collect(),
129        }
130    }
131
132    fn resource(&self) -> Resource {
133        Resource {
134            attributes: self
135                .resource_attributes
136                .iter()
137                .map(|(key, value)| KeyValue::new(key, value))
138                .collect(),
139        }
140    }
141}
142
143/// The connections that a reporter sends on.
144pub trait TelemetryConnections: Send + Sync + 'static {
145    /// An open connection with its broker id, preferring the broker of
146    /// `preferred` while it stays open. Per KIP-714 a client keeps pushing to
147    /// the same broker. `None` when no connection is open.
148    fn open_connection(&self, preferred: Option<i32>) -> Option<(i32, Arc<Connection>)>;
149}
150
151/// The timing of a reporter.
152#[derive(Clone, Copy, Debug, PartialEq, Eq)]
153pub struct TelemetryTiming {
154    /// Kafka's `request.timeout.ms`: the wait for a response, and the bound
155    /// on the terminating push after a close.
156    pub request_timeout: Duration,
157    /// Kafka's `reconnect.backoff.ms`: the wait for an open connection.
158    pub reconnect_backoff: Duration,
159}
160
161/// A running KIP-714 reporter. Dropping it starts the close, as
162/// [`close`](Self::close) does, without the wait.
163#[derive(Debug)]
164pub struct ClientTelemetry {
165    close: CancellationToken,
166    done: CancellationToken,
167    /// The client instance id of the last subscription, `None` until the
168    /// first one loads. The reporter task publishes it.
169    instance_id: Arc<watch::Sender<Option<uuid::Uuid>>>,
170}
171
172impl ClientTelemetry {
173    /// Start a reporter task that sends on `connections` and pushes the
174    /// metrics of `metrics`.
175    ///
176    /// # Panics
177    /// Panics outside a Tokio runtime.
178    #[must_use]
179    pub fn start(
180        connections: impl TelemetryConnections,
181        config: &ClientTelemetryConfig,
182        metrics: ClientMetrics,
183        timing: TelemetryTiming,
184    ) -> Self {
185        let close = CancellationToken::new();
186        let done = CancellationToken::new();
187        let instance_id = Arc::new(watch::Sender::new(None));
188        let sender = TelemetrySender::new(Collector::new(metrics, config.resource()));
189        tokio::spawn(run(
190            connections,
191            sender,
192            timing,
193            Arc::clone(&instance_id),
194            (close.clone(), done.clone()),
195        ));
196        Self {
197            close,
198            done,
199            instance_id,
200        }
201    }
202
203    /// The client instance id that the broker assigned, waiting up to
204    /// `timeout` for the first subscription to load. `None` when none has
205    /// loaded by then, as Kafka's `ClientTelemetrySender.clientInstanceId`
206    /// returns an empty `Optional`. A zero `timeout` does not wait.
207    pub async fn client_instance_id(&self, timeout: Duration) -> Option<uuid::Uuid> {
208        let mut loaded = self.instance_id.subscribe();
209        time::timeout(timeout, loaded.wait_for(Option::is_some))
210            .await
211            .ok()
212            .and_then(Result::ok)
213            .and_then(|id| *id)
214    }
215
216    /// Send the terminating push, if the client has a subscription, and wait
217    /// for the reporter to end. The wait is at most about two request
218    /// timeouts: one for a request in flight and one for the terminating
219    /// push.
220    pub async fn close(&self) {
221        self.close.cancel();
222        self.done.cancelled().await;
223    }
224}
225
226impl Drop for ClientTelemetry {
227    fn drop(&mut self) {
228        self.close.cancel();
229    }
230}
231
232/// Kafka's `clientInstanceId` of a producer, consumer or admin client
233/// (`ClientTelemetryUtils.fetchClientInstanceId`).
234///
235/// The broker-assigned
236/// client instance id, waiting up to `timeout` for the first subscription.
237/// `Ok(None)` when none has loaded by then, as Kafka returns `null`.
238///
239/// # Errors
240/// Returns [`ClientError::InvalidArgument`] for a negative or non-finite
241/// `timeout`, and [`ClientError::TelemetryDisabled`] when the client does not
242/// push metrics (`telemetry` is `None`), as Kafka throws an
243/// `IllegalArgumentException` and an `IllegalStateException`.
244pub async fn client_instance_id(
245    telemetry: Option<&ClientTelemetry>,
246    timeout: Time,
247) -> Result<Option<uuid::Uuid>, ClientError> {
248    if !timeout.secs_f64().is_finite() || timeout < Time::ZERO {
249        return Err(ClientError::InvalidArgument(
250            "The timeout cannot be negative.".to_owned(),
251        ));
252    }
253    let telemetry = telemetry.ok_or(ClientError::TelemetryDisabled)?;
254    Ok(telemetry.client_instance_id(timeout.to_std()).await)
255}
256
257/// The reporter task: Kafka's `NetworkClient.TelemetrySender.maybeUpdate`
258/// in a loop.
259async fn run(
260    connections: impl TelemetryConnections,
261    mut sender: TelemetrySender,
262    timing: TelemetryTiming,
263    instance_id: Arc<watch::Sender<Option<uuid::Uuid>>>,
264    (close, done): (CancellationToken, CancellationToken),
265) {
266    let mut closing: Option<Instant> = None;
267    let mut sticky: Option<i32> = None;
268    loop {
269        if closing.is_none() && close.is_cancelled() {
270            closing = Some(Instant::now() + timing.request_timeout);
271            if !sender.initiate_close() {
272                break;
273            }
274        }
275        let now = Instant::now();
276        if closing.is_some_and(|deadline| now >= deadline) {
277            break;
278        }
279        match sender.time_to_next_update(now, timing.request_timeout) {
280            Some(wait) if wait.is_zero() => {}
281            // During the close the only request is the terminating push,
282            // which is due at once.
283            _ if closing.is_some() => break,
284            Some(wait) => {
285                tokio::select! {
286                    () = sleep(wait) => {}
287                    () = close.cancelled() => {}
288                }
289                continue;
290            }
291            None => {
292                close.cancelled().await;
293                continue;
294            }
295        }
296        let open = connections.open_connection(sticky);
297        let Some((broker_id, connection)) = open else {
298            sticky = None;
299            if closing.is_some() {
300                break;
301            }
302            tokio::select! {
303                () = sleep(timing.reconnect_backoff) => {}
304                () = close.cancelled() => {}
305            }
306            continue;
307        };
308        sticky = Some(broker_id);
309        let Some(request) = sender.create_request(unix_nanos_now()) else {
310            continue;
311        };
312        let failed = match request {
313            TelemetryRequest::Subscriptions(request) => match send(&connection, request).await {
314                Ok(response) => {
315                    sender.handle_subscriptions_response(
316                        &response,
317                        Instant::now(),
318                        backoff::random_unit(),
319                    );
320                    // Kafka's `subscriptionLoaded.signalAll()`.
321                    let loaded = sender
322                        .client_instance_id()
323                        .map(|id| uuid::Uuid::from_bytes(id.0));
324                    instance_id.send_if_modified(|current| {
325                        let changed = loaded.is_some() && *current != loaded;
326                        if changed {
327                            *current = loaded;
328                        }
329                        changed
330                    });
331                    None
332                }
333                Err(failure) => Some(failure),
334            },
335            TelemetryRequest::Push(request) => match send(&connection, request).await {
336                Ok(response) => {
337                    sender.handle_push_response(&response, Instant::now());
338                    None
339                }
340                Err(failure) => Some(failure),
341            },
342        };
343        if let Some(failure) = failed {
344            // Per KIP-714, a failed broker is left for another one.
345            sticky = None;
346            sender.handle_failed_request(failure, Instant::now());
347        }
348        if sender.state() == TelemetryState::TerminatingPushInProgress {
349            break;
350        }
351    }
352    sender.close();
353    done.cancel();
354}
355
356/// Send a telemetry request on `connection`. A broker that does not list
357/// the API fails the request as Kafka's `UnsupportedVersionException` does.
358async fn send<R: ProtocolRequest>(
359    connection: &Connection,
360    request: R,
361) -> Result<R::Response, Failure> {
362    if connection.advertised_api_range(R::API_KEY).is_none() {
363        tracing::debug!(
364            api_key = R::API_KEY,
365            "broker does not support client telemetry"
366        );
367        return Err(Failure::Fatal);
368    }
369    connection.send(request).await.map_err(|error| {
370        tracing::debug!(%error, api_key = R::API_KEY, "telemetry request failed");
371        failure(&error)
372    })
373}
374
375/// A disconnect or timeout is Kafka's `RetriableException`.
376fn failure(error: &ClientError) -> Failure {
377    match error {
378        ClientError::Disconnected
379        | ClientError::Timeout(_)
380        | ClientError::Io(_)
381        | ClientError::Connect { .. }
382        | ClientError::Tls { .. }
383        | ClientError::Sasl { .. } => Failure::Retriable,
384        _ => Failure::Fatal,
385    }
386}
387
388#[cfg(test)]
389mod tests;
krabka_client_core/telemetry.rs

krabka_client_core/
telemetry.rs
