vector/sinks/splunk_hec/common/
acknowledgements.rs

1use http_body::{Body as _, Collected};
2use hyper::Body;
3use serde::{Deserialize, Serialize};
4use std::{
5    collections::HashMap,
6    io::Write,
7    num::{NonZeroU8, NonZeroU64},
8    sync::Arc,
9    time::Duration,
10};
11use tokio::sync::{mpsc::Receiver, oneshot::Sender};
12use vector_lib::{configurable::configurable_component, event::EventStatus};
13
14use super::service::{HttpRequestBuilder, MetadataFields};
15use crate::{
16    config::AcknowledgementsConfig,
17    http::HttpClient,
18    internal_events::{
19        SplunkIndexerAcknowledgementAPIError, SplunkIndexerAcknowledgementAckAdded,
20        SplunkIndexerAcknowledgementAcksRemoved,
21    },
22    sinks::util::Compressor,
23};
24
25/// Splunk HEC acknowledgement configuration.
26#[configurable_component]
27#[derive(Clone, Debug)]
28#[serde(default)]
29pub struct HecClientAcknowledgementsConfig {
30    /// Controls if the sink integrates with [Splunk HEC indexer acknowledgements][splunk_indexer_ack_docs] for end-to-end acknowledgements.
31    ///
32    /// [splunk_indexer_ack_docs]: https://docs.splunk.com/Documentation/Splunk/8.2.3/Data/AboutHECIDXAck
33    pub indexer_acknowledgements_enabled: bool,
34
35    /// The amount of time to wait between queries to the Splunk HEC indexer acknowledgement endpoint.
36    #[configurable(metadata(docs::type_unit = "seconds"))]
37    pub query_interval: NonZeroU8,
38
39    /// The maximum number of times an acknowledgement ID is queried for its status.
40    pub retry_limit: NonZeroU8,
41
42    /// The maximum number of pending acknowledgements from events sent to the Splunk HEC collector.
43    ///
44    /// Once reached, the sink begins applying backpressure.
45    pub max_pending_acks: NonZeroU64,
46
47    #[serde(
48        default,
49        deserialize_with = "crate::serde::bool_or_struct",
50        flatten,
51        skip_serializing_if = "crate::serde::is_default"
52    )]
53    pub inner: AcknowledgementsConfig,
54}
55
56impl Default for HecClientAcknowledgementsConfig {
57    fn default() -> Self {
58        Self {
59            indexer_acknowledgements_enabled: true,
60            query_interval: NonZeroU8::new(10).unwrap(),
61            retry_limit: NonZeroU8::new(30).unwrap(),
62            max_pending_acks: NonZeroU64::new(1_000_000).unwrap(),
63            inner: Default::default(),
64        }
65    }
66}
67
68#[derive(Deserialize, Serialize, Eq, PartialEq, Debug)]
69pub struct HecAckStatusRequest {
70    pub acks: Vec<u64>,
71}
72
73#[derive(Deserialize, Serialize, Debug)]
74pub struct HecAckStatusResponse {
75    pub acks: HashMap<u64, bool>,
76}
77
78#[derive(Debug)]
79pub enum HecAckApiError {
80    ClientBuildRequest,
81    ClientParseResponse,
82    ClientSendQuery,
83    ServerSendQuery,
84}
85
86struct HecAckClient {
87    acks: HashMap<u64, (u8, Sender<EventStatus>)>,
88    retry_limit: u8,
89    client: HttpClient,
90    http_request_builder: Arc<HttpRequestBuilder>,
91}
92
93impl HecAckClient {
94    fn new(
95        retry_limit: u8,
96        client: HttpClient,
97        http_request_builder: Arc<HttpRequestBuilder>,
98    ) -> Self {
99        Self {
100            acks: HashMap::new(),
101            retry_limit,
102            client,
103            http_request_builder,
104        }
105    }
106
107    /// Adds an ack id to be queried
108    fn add(&mut self, ack_id: u64, ack_event_status_sender: Sender<EventStatus>) {
109        self.acks
110            .insert(ack_id, (self.retry_limit, ack_event_status_sender));
111        emit!(SplunkIndexerAcknowledgementAckAdded);
112    }
113
114    /// Queries Splunk HEC with stored ack ids and finalizes events that are successfully acked
115    async fn run(&mut self) {
116        let ack_query_body = self.get_ack_query_body();
117        if !ack_query_body.acks.is_empty() {
118            let ack_query_response = self.send_ack_query_request(&ack_query_body).await;
119
120            match ack_query_response {
121                Ok(ack_query_response) => {
122                    debug!(message = "Received ack statuses.", ?ack_query_response);
123                    let acked_ack_ids = ack_query_response
124                        .acks
125                        .iter()
126                        .filter(|&(_ack_id, ack_status)| *ack_status)
127                        .map(|(ack_id, _ack_status)| *ack_id)
128                        .collect::<Vec<u64>>();
129                    self.finalize_delivered_ack_ids(acked_ack_ids.as_slice());
130                    self.expire_ack_ids_with_status(EventStatus::Rejected);
131                }
132                Err(error) => {
133                    match error {
134                        HecAckApiError::ClientParseResponse | HecAckApiError::ClientSendQuery => {
135                            // If we are permanently unable to interact with
136                            // Splunk HEC indexer acknowledgements (e.g. due to
137                            // request/response format changes in future
138                            // versions), log an error and fall back to default
139                            // behavior.
140                            emit!(SplunkIndexerAcknowledgementAPIError {
141                                message: "Unable to use indexer acknowledgements. Acknowledging based on initial 200 OK.",
142                                error,
143                            });
144                            self.finalize_delivered_ack_ids(
145                                self.acks.keys().copied().collect::<Vec<_>>().as_slice(),
146                            );
147                        }
148                        _ => {
149                            emit!(SplunkIndexerAcknowledgementAPIError {
150                                message: "Unable to send acknowledgement query request. Will retry.",
151                                error,
152                            });
153                            self.expire_ack_ids_with_status(EventStatus::Errored);
154                        }
155                    }
156                }
157            };
158        }
159    }
160
161    /// Removes successfully acked ack ids and finalizes associated events
162    fn finalize_delivered_ack_ids(&mut self, ack_ids: &[u64]) {
163        let mut removed_count = 0.0;
164        for ack_id in ack_ids {
165            if let Some((_, ack_event_status_sender)) = self.acks.remove(ack_id) {
166                _ = ack_event_status_sender.send(EventStatus::Delivered);
167                removed_count += 1.0;
168                debug!(message = "Finalized ack id.", ?ack_id);
169            }
170        }
171        emit!(SplunkIndexerAcknowledgementAcksRemoved {
172            count: removed_count
173        });
174    }
175
176    /// Builds an ack query body with stored ack ids
177    fn get_ack_query_body(&mut self) -> HecAckStatusRequest {
178        HecAckStatusRequest {
179            acks: self.acks.keys().copied().collect::<Vec<u64>>(),
180        }
181    }
182
183    /// Decrements retry count on all stored ack ids by 1
184    fn decrement_retries(&mut self) {
185        for (retries, _) in self.acks.values_mut() {
186            *retries = retries.checked_sub(1).unwrap_or(0);
187        }
188    }
189
190    /// Removes all expired ack ids (those with a retry count of 0) and
191    /// finalizes associated events with the given status
192    fn expire_ack_ids_with_status(&mut self, status: EventStatus) {
193        let expired_ack_ids = self
194            .acks
195            .iter()
196            .filter_map(|(ack_id, (retries, _))| (*retries == 0).then_some(*ack_id))
197            .collect::<Vec<_>>();
198        let mut removed_count = 0.0;
199        for ack_id in expired_ack_ids {
200            if let Some((_, ack_event_status_sender)) = self.acks.remove(&ack_id) {
201                _ = ack_event_status_sender.send(status);
202                removed_count += 1.0;
203            }
204        }
205        emit!(SplunkIndexerAcknowledgementAcksRemoved {
206            count: removed_count
207        });
208    }
209
210    // Sends an ack status query request to Splunk HEC
211    async fn send_ack_query_request(
212        &mut self,
213        request_body: &HecAckStatusRequest,
214    ) -> Result<HecAckStatusResponse, HecAckApiError> {
215        self.decrement_retries();
216        let request_body_bytes = crate::serde::json::to_bytes(request_body)
217            .map_err(|_| HecAckApiError::ClientBuildRequest)?
218            .freeze();
219        let mut compressor = Compressor::from(self.http_request_builder.compression);
220        compressor
221            .write_all(request_body_bytes.as_ref())
222            .map_err(|_| HecAckApiError::ClientBuildRequest)?;
223        let payload = compressor
224            .finish()
225            .map_err(|_| HecAckApiError::ClientBuildRequest)?
226            .freeze();
227        let request = self
228            .http_request_builder
229            .build_request(
230                payload,
231                "/services/collector/ack",
232                None,
233                MetadataFields::default(),
234                false,
235            )
236            .map_err(|_| HecAckApiError::ClientBuildRequest)?;
237
238        let response = self
239            .client
240            .send(request.map(Body::from))
241            .await
242            .map_err(|_| HecAckApiError::ServerSendQuery)?;
243
244        let status = response.status();
245        if status.is_success() {
246            let response_body = response
247                .into_body()
248                .collect()
249                .await
250                .map(Collected::to_bytes)
251                .map_err(|_| HecAckApiError::ClientParseResponse)?;
252            serde_json::from_slice::<HecAckStatusResponse>(&response_body)
253                .map_err(|_| HecAckApiError::ClientParseResponse)
254        } else if status.is_client_error() {
255            Err(HecAckApiError::ClientSendQuery)
256        } else {
257            Err(HecAckApiError::ServerSendQuery)
258        }
259    }
260}
261
262pub async fn run_acknowledgements(
263    mut receiver: Receiver<(u64, Sender<EventStatus>)>,
264    client: HttpClient,
265    http_request_builder: Arc<HttpRequestBuilder>,
266    indexer_acknowledgements: HecClientAcknowledgementsConfig,
267) {
268    let mut interval = tokio::time::interval(Duration::from_secs(
269        indexer_acknowledgements.query_interval.get() as u64,
270    ));
271    let mut ack_client = HecAckClient::new(
272        indexer_acknowledgements.retry_limit.get(),
273        client,
274        http_request_builder,
275    );
276
277    loop {
278        tokio::select! {
279            _ = interval.tick() => {
280                ack_client.run().await;
281            },
282            ack_info = receiver.recv() => {
283                match ack_info {
284                    Some((ack_id, tx)) => {
285                        ack_client.add(ack_id, tx);
286                        debug!(message = "Stored ack id.", ?ack_id);
287                    },
288                    None => break,
289                }
290            }
291        }
292    }
293}
294
295#[cfg(test)]
296mod tests {
297    use std::sync::Arc;
298
299    use futures_util::{StreamExt, stream::FuturesUnordered};
300    use tokio::sync::oneshot::{self, Receiver};
301    use vector_lib::{config::proxy::ProxyConfig, event::EventStatus};
302
303    use super::HecAckClient;
304    use crate::{
305        http::HttpClient,
306        sinks::{
307            splunk_hec::common::{
308                EndpointTarget, acknowledgements::HecAckStatusRequest, service::HttpRequestBuilder,
309            },
310            util::Compression,
311        },
312    };
313
314    fn get_ack_client(retry_limit: u8) -> HecAckClient {
315        let client = HttpClient::new(None, &ProxyConfig::default()).unwrap();
316        let http_request_builder = HttpRequestBuilder::new(
317            String::from(""),
318            EndpointTarget::default(),
319            String::from(""),
320            Compression::default(),
321        );
322        HecAckClient::new(retry_limit, client, Arc::new(http_request_builder))
323    }
324
325    fn populate_ack_client(
326        ack_client: &mut HecAckClient,
327        ack_ids: &[u64],
328    ) -> Vec<Receiver<EventStatus>> {
329        let mut ack_status_rxs = Vec::new();
330        for ack_id in ack_ids {
331            let (tx, rx) = oneshot::channel();
332            ack_client.add(*ack_id, tx);
333            ack_status_rxs.push(rx);
334        }
335        ack_status_rxs
336    }
337
338    #[test]
339    fn test_get_ack_query_body() {
340        let mut ack_client = get_ack_client(1);
341        let ack_ids = (0..100).collect::<Vec<u64>>();
342        _ = populate_ack_client(&mut ack_client, &ack_ids);
343        let expected_ack_body = HecAckStatusRequest { acks: ack_ids };
344
345        let mut ack_request_body = ack_client.get_ack_query_body();
346        ack_request_body.acks.sort_unstable();
347        assert_eq!(expected_ack_body, ack_request_body);
348    }
349
350    #[test]
351    fn test_decrement_retries() {
352        let mut ack_client = get_ack_client(1);
353        let ack_ids = (0..100).collect::<Vec<u64>>();
354        _ = populate_ack_client(&mut ack_client, &ack_ids);
355
356        let mut ack_request_body = ack_client.get_ack_query_body();
357        ack_request_body.acks.sort_unstable();
358        assert_eq!(ack_ids, ack_request_body.acks);
359        ack_client.decrement_retries();
360        ack_client.expire_ack_ids_with_status(EventStatus::Rejected);
361
362        let ack_request_body = ack_client.get_ack_query_body();
363        assert!(ack_request_body.acks.is_empty())
364    }
365
366    #[tokio::test]
367    async fn test_finalize_delivered_ack_ids() {
368        let mut ack_client = get_ack_client(1);
369        let ack_ids = (0..100).collect::<Vec<u64>>();
370        let ack_status_rxs = populate_ack_client(&mut ack_client, &ack_ids);
371
372        ack_client.finalize_delivered_ack_ids(ack_ids.as_slice());
373        let mut statuses = ack_status_rxs.into_iter().collect::<FuturesUnordered<_>>();
374        while let Some(status) = statuses.next().await {
375            assert_eq!(EventStatus::Delivered, status.unwrap());
376        }
377    }
378
379    #[tokio::test]
380    async fn test_expire_ack_ids_with_status() {
381        let mut ack_client = get_ack_client(1);
382        let ack_ids = (0..100).collect::<Vec<u64>>();
383        let ack_status_rxs = populate_ack_client(&mut ack_client, &ack_ids);
384
385        ack_client.decrement_retries();
386        ack_client.expire_ack_ids_with_status(EventStatus::Rejected);
387        let mut statuses = ack_status_rxs.into_iter().collect::<FuturesUnordered<_>>();
388        while let Some(status) = statuses.next().await {
389            assert_eq!(EventStatus::Rejected, status.unwrap());
390        }
391    }
392}
vector/sinks/splunk_hec/common/acknowledgements.rs

vector/sinks/splunk_hec/common/
acknowledgements.rs
