krabka_client_admin/
telemetry.rs

1//! Client metrics push (KIP-714) of the admin client: Kafka's
2//! `enable.metrics.push` on `KafkaAdminClient`.
3//!
4//! The reporter of `krabka-client-core` sends on the admin client's current
5//! bootstrap or controller connection, and every connection of the client
6//! counts in the `admin-client-metrics` network metrics.
7
8use std::sync::Arc;
9
10use krabka_client_core::{
11    Connection,
12    telemetry::{TelemetryConnections, client_instance_id},
13};
14use krabka_units::Time;
15
16use crate::{AdminClient, AdminError, RecoveringConnection};
17
18/// The connection that the admin client's reporter sends on. It has no
19/// broker id in the metadata, so the reporter sees it as `-1`, as a
20/// bootstrap connection.
21pub struct AdminTelemetryConnection(pub(crate) Arc<RecoveringConnection>);
22
23impl TelemetryConnections for AdminTelemetryConnection {
24    fn open_connection(&self, _preferred: Option<i32>) -> Option<(i32, Arc<Connection>)> {
25        // A connection that is being replaced is skipped; the reporter tries
26        // again after the reconnect backoff.
27        let connection = self.0.inner.try_read().ok()?;
28        (!connection.is_closed()).then(|| (-1, Arc::clone(&connection)))
29    }
30}
31
32impl AdminClient {
33    /// The client instance id that the broker assigned for metrics push
34    /// (KIP-714), as Kafka's `KafkaAdminClient.clientInstanceId` returns it.
35    ///
36    /// The call waits up to `timeout` for the first
37    /// `GetTelemetrySubscriptions` response. It returns `Ok(None)` when none
38    /// has come by then, as Kafka returns `null`. A zero `timeout` does not
39    /// wait.
40    ///
41    /// # Errors
42    /// Returns [`AdminError::Transport`] with
43    /// [`ClientError::InvalidArgument`] for a negative `timeout`, and with
44    /// [`ClientError::TelemetryDisabled`] when
45    /// [`AdminClientConfig::enable_metrics_push`] is `false`, as Kafka throws
46    /// an `IllegalArgumentException` and an `IllegalStateException`.
47    ///
48    /// [`ClientError::InvalidArgument`]: krabka_client_core::ClientError::InvalidArgument
49    /// [`ClientError::TelemetryDisabled`]: krabka_client_core::ClientError::TelemetryDisabled
50    /// [`AdminClientConfig::enable_metrics_push`]: crate::AdminClientConfig::enable_metrics_push
51    pub async fn client_instance_id(
52        &self,
53        timeout: Time,
54    ) -> Result<Option<uuid::Uuid>, AdminError> {
55        Ok(client_instance_id(self.telemetry.as_ref(), timeout).await?)
56    }
57
58    /// Close the client: send the terminating telemetry push (KIP-714), if
59    /// the client pushes metrics and has a subscription, and wait for it, as
60    /// Kafka's `KafkaAdminClient.close` does. The wait is at most about two
61    /// request timeouts. Dropping the client sends the push in the
62    /// background instead.
63    pub async fn close(self) {
64        if let Some(telemetry) = &self.telemetry {
65            telemetry.close().await;
66        }
67    }
68}
69
70#[cfg(test)]
71mod tests {
72    use std::time::Duration;
73
74    use assert2::assert;
75    use bytes::BytesMut;
76    use krabka_client_core::{
77        MockBroker,
78        telemetry::otlp::{MetricData, MetricsData, NumberValue},
79    };
80    use krabka_protocol::{
81        Decode, Encode,
82        owned::{
83            api_versions_request,
84            api_versions_response::{ApiVersion, ApiVersionsResponse},
85            get_telemetry_subscriptions_request::{self, GetTelemetrySubscriptionsRequest},
86            get_telemetry_subscriptions_response::GetTelemetrySubscriptionsResponse,
87            push_telemetry_request::{self, PushTelemetryRequest},
88            push_telemetry_response::PushTelemetryResponse,
89        },
90        primitives::uuid::Uuid as WireUuid,
91    };
92    use tokio::{sync::mpsc, time};
93
94    use crate::AdminClientConfig;
95
96    const INSTANCE: WireUuid = WireUuid([7; 16]);
97    const CONNECTION_COUNT: &str = "org.apache.kafka.admin.client.connection.count";
98
99    /// A telemetry request that the broker received. A push holds each
100    /// metric's name and value.
101    #[derive(Clone, Debug, PartialEq)]
102    enum Seen {
103        Subscriptions(WireUuid),
104        Push {
105            client_instance_id: WireUuid,
106            terminating: bool,
107            metrics: Vec<(String, Option<NumberValue>)>,
108        },
109    }
110
111    /// The body of a flexible request after the request header v2.
112    fn request_body(frame: &[u8]) -> &[u8] {
113        let client_id_len = usize::try_from(i16::from_be_bytes([frame[0], frame[1]])).unwrap_or(0);
114        &frame[2 + client_id_len + 1..]
115    }
116
117    fn flexible_response(response: &impl Encode) -> Vec<u8> {
118        let mut body = BytesMut::from(&[0u8][..]);
119        response.encode(&mut body, 0).unwrap();
120        body.to_vec()
121    }
122
123    fn seen_push(request: &PushTelemetryRequest) -> Seen {
124        let data = MetricsData::decode(&request.metrics).unwrap();
125        let metrics = data
126            .resource_metrics
127            .into_iter()
128            .flat_map(|resource| resource.scope_metrics)
129            .flat_map(|scope| scope.metrics)
130            .map(|metric| {
131                let value = match metric.data {
132                    Some(
133                        MetricData::Gauge { data_points } | MetricData::Sum { data_points, .. },
134                    ) => data_points.first().and_then(|point| point.value),
135                    None => None,
136                };
137                (metric.name, value)
138            })
139            .collect();
140        Seen::Push {
141            client_instance_id: request.client_instance_id,
142            terminating: request.terminating,
143            metrics,
144        }
145    }
146
147    /// A broker that assigns [`INSTANCE`] and subscribes the admin client's
148    /// connection count, pushed every 200 ms. With `advertise_telemetry`
149    /// false it does not list the KIP-714 APIs.
150    async fn start_broker(
151        advertise_telemetry: bool,
152    ) -> (MockBroker, mpsc::UnboundedReceiver<Seen>) {
153        let (seen_tx, seen_rx) = mpsc::unbounded_channel();
154        let broker = MockBroker::start(move |api_key, version, _corr, frame| match api_key {
155            api_versions_request::API_KEY => {
156                let mut keys = vec![(api_versions_request::API_KEY, 3)];
157                if advertise_telemetry {
158                    keys.push((get_telemetry_subscriptions_request::API_KEY, 0));
159                    keys.push((push_telemetry_request::API_KEY, 0));
160                }
161                let response = ApiVersionsResponse {
162                    api_keys: keys
163                        .into_iter()
164                        .map(|(api_key, max_version)| ApiVersion {
165                            api_key,
166                            min_version: 0,
167                            max_version,
168                            ..Default::default()
169                        })
170                        .collect(),
171                    ..Default::default()
172                };
173                let mut body = BytesMut::new();
174                response.encode(&mut body, version).unwrap();
175                Some(body.to_vec())
176            }
177            get_telemetry_subscriptions_request::API_KEY => {
178                let mut body = request_body(frame);
179                let request = GetTelemetrySubscriptionsRequest::decode(&mut body, version).unwrap();
180                let _ = seen_tx.send(Seen::Subscriptions(request.client_instance_id));
181                Some(flexible_response(&GetTelemetrySubscriptionsResponse {
182                    client_instance_id: INSTANCE,
183                    subscription_id: 5,
184                    push_interval_ms: 200,
185                    telemetry_max_bytes: 1 << 20,
186                    requested_metrics: vec![CONNECTION_COUNT.to_owned()],
187                    ..Default::default()
188                }))
189            }
190            push_telemetry_request::API_KEY => {
191                let mut body = request_body(frame);
192                let request = PushTelemetryRequest::decode(&mut body, version).unwrap();
193                let _ = seen_tx.send(seen_push(&request));
194                Some(flexible_response(&PushTelemetryResponse::default()))
195            }
196            _ => None,
197        })
198        .await
199        .unwrap();
200        (broker, seen_rx)
201    }
202
203    async fn next(seen: &mut mpsc::UnboundedReceiver<Seen>) -> Seen {
204        time::timeout(Duration::from_secs(10), seen.recv())
205            .await
206            .expect("a telemetry request within 10 s")
207            .expect("the broker is running")
208    }
209
210    /// Kafka's admin client with `enable.metrics.push`: it asks for the
211    /// subscription, pushes the subscribed metrics (its one open
212    /// connection), returns the assigned id from `clientInstanceId`, and
213    /// sends a terminating push on close. With the setting off, or with a
214    /// broker without KIP-714, it sends no telemetry request.
215    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
216    async fn admin_client_pushes_metrics_and_reports_its_instance_id() {
217        let disabled =
218            "client-core: Telemetry is not enabled. Set config `enable.metrics.push` to \
219                        `true`."
220                .to_owned();
221        let push = |terminating| Seen::Push {
222            client_instance_id: INSTANCE,
223            terminating,
224            metrics: vec![(CONNECTION_COUNT.to_owned(), Some(NumberValue::Double(1.0)))],
225        };
226        let cases = [
227            (
228                "enable.metrics.push=true",
229                true,
230                true,
231                2,
232                Ok(Some(uuid::Uuid::from_bytes(INSTANCE.0))),
233                vec![Seen::Subscriptions(WireUuid::ZERO), push(false), push(true)],
234            ),
235            (
236                "enable.metrics.push=false",
237                false,
238                true,
239                0,
240                Err(disabled),
241                vec![],
242            ),
243            ("broker without KIP-714", true, false, 0, Ok(None), vec![]),
244        ];
245        for (name, enable_metrics_push, advertise_telemetry, count, expected_id, expected) in cases
246        {
247            let (broker, mut seen) = start_broker(advertise_telemetry).await;
248            let admin = crate::AdminClient::connect_with_config(
249                &[broker.addr.to_string()],
250                AdminClientConfig {
251                    enable_metrics_push,
252                    ..AdminClientConfig::default()
253                },
254            )
255            .await
256            .expect("admin connects");
257            let timeout = if count == 0 {
258                krabka_units::millis(300)
259            } else {
260                krabka_units::secs(10)
261            };
262            let id = admin
263                .client_instance_id(timeout)
264                .await
265                .map_err(|error| error.to_string());
266            let mut requests = Vec::new();
267            for _ in 0..count {
268                requests.push(next(&mut seen).await);
269            }
270            time::timeout(Duration::from_secs(20), admin.close())
271                .await
272                .expect("the close ends");
273            broker.stop();
274            while let Ok(request) = seen.try_recv() {
275                requests.push(request);
276            }
277            assert!((id, requests) == (expected_id, expected), "{name}");
278        }
279    }
280}
krabka_client_admin/telemetry.rs

krabka_client_admin/
telemetry.rs
