log_store/kafka/
producer.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::sync::Arc;
16
17use common_telemetry::warn;
18use dashmap::DashMap;
19use rskafka::client::partition::{Compression, OffsetAt, PartitionClient};
20use rskafka::record::Record;
21use store_api::logstore::provider::KafkaProvider;
22use tokio::sync::mpsc::{self, Receiver, Sender};
23
24use crate::error::{self, Result};
25use crate::kafka::log_store::TopicStat;
26use crate::kafka::worker::{BackgroundProducerWorker, ProduceResultHandle, WorkerRequest};
27use crate::metrics::{
28    METRIC_KAFKA_CLIENT_BYTES_TOTAL, METRIC_KAFKA_CLIENT_PRODUCE_ELAPSED,
29    METRIC_KAFKA_CLIENT_TRAFFIC_TOTAL,
30};
31
32pub type OrderedBatchProducerRef = Arc<OrderedBatchProducer>;
33
34// Max batch size for a `OrderedBatchProducer` to handle requests.
35const REQUEST_BATCH_SIZE: usize = 64;
36
37// Producer channel size
38const PRODUCER_CHANNEL_SIZE: usize = REQUEST_BATCH_SIZE * 2;
39
40/// [`OrderedBatchProducer`] attempts to aggregate multiple produce requests together
41#[derive(Debug)]
42pub(crate) struct OrderedBatchProducer {
43    pub(crate) sender: Sender<WorkerRequest>,
44}
45
46impl OrderedBatchProducer {
47    pub(crate) fn channel() -> (Sender<WorkerRequest>, Receiver<WorkerRequest>) {
48        mpsc::channel(PRODUCER_CHANNEL_SIZE)
49    }
50
51    /// Constructs a new [`OrderedBatchProducer`].
52    pub(crate) fn new(
53        (tx, rx): (Sender<WorkerRequest>, Receiver<WorkerRequest>),
54        provider: Arc<KafkaProvider>,
55        client: Arc<dyn ProducerClient>,
56        compression: Compression,
57        max_batch_bytes: usize,
58        topic_stats: Arc<DashMap<Arc<KafkaProvider>, TopicStat>>,
59    ) -> Self {
60        let mut worker = BackgroundProducerWorker {
61            provider,
62            client,
63            compression,
64            receiver: rx,
65            request_batch_size: REQUEST_BATCH_SIZE,
66            max_batch_bytes,
67            topic_stats,
68        };
69        tokio::spawn(async move { worker.run().await });
70        Self { sender: tx }
71    }
72
73    /// Writes `data` to the [`OrderedBatchProducer`].
74    ///
75    /// Returns the [ProduceResultHandle], which will receive a result when data has been committed to Kafka
76    /// or an unrecoverable error has been encountered.
77    ///
78    /// ## Panic
79    /// Panic if any [Record]'s `approximate_size` > `max_batch_bytes`.
80    pub(crate) async fn produce(&self, batch: Vec<Record>) -> Result<ProduceResultHandle> {
81        let (req, handle) = WorkerRequest::new_produce_request(batch);
82        if self.sender.send(req).await.is_err() {
83            warn!("OrderedBatchProducer is already exited");
84            return error::OrderedBatchProducerStoppedSnafu {}.fail();
85        }
86
87        Ok(handle)
88    }
89
90    /// Sends an [WorkerRequest::FetchLatestOffset] request to the producer.
91    /// This is used to fetch the latest offset for the topic.
92    pub(crate) async fn fetch_latest_offset(&self) -> Result<()> {
93        if self
94            .sender
95            .send(WorkerRequest::FetchLatestOffset)
96            .await
97            .is_err()
98        {
99            warn!("OrderedBatchProducer is already exited");
100            return error::OrderedBatchProducerStoppedSnafu {}.fail();
101        }
102        Ok(())
103    }
104}
105
106pub struct ProduceResult {
107    pub offsets: Vec<i64>,
108    pub encoded_request_size: usize,
109}
110
111#[async_trait::async_trait]
112pub trait ProducerClient: std::fmt::Debug + Send + Sync {
113    async fn produce(
114        &self,
115        records: Vec<Record>,
116        compression: Compression,
117    ) -> rskafka::client::error::Result<ProduceResult>;
118
119    async fn get_offset(&self, at: OffsetAt) -> rskafka::client::error::Result<i64>;
120}
121
122#[async_trait::async_trait]
123impl ProducerClient for PartitionClient {
124    async fn produce(
125        &self,
126        records: Vec<Record>,
127        compression: Compression,
128    ) -> rskafka::client::error::Result<ProduceResult> {
129        let partition = self.partition().to_string();
130        let _timer = METRIC_KAFKA_CLIENT_PRODUCE_ELAPSED
131            .with_label_values(&[self.topic(), &partition])
132            .start_timer();
133
134        let result = self.produce(records, compression).await?;
135
136        METRIC_KAFKA_CLIENT_BYTES_TOTAL
137            .with_label_values(&[self.topic(), &partition])
138            .inc_by(result.encoded_request_size as u64);
139        METRIC_KAFKA_CLIENT_TRAFFIC_TOTAL
140            .with_label_values(&[self.topic(), &partition])
141            .inc();
142
143        Ok(ProduceResult {
144            offsets: result.offsets,
145            encoded_request_size: result.encoded_request_size,
146        })
147    }
148
149    async fn get_offset(&self, at: OffsetAt) -> rskafka::client::error::Result<i64> {
150        self.get_offset(at).await
151    }
152}
153
154#[cfg(test)]
155mod tests {
156    use std::sync::{Arc, Mutex};
157    use std::time::Duration;
158
159    use common_base::readable_size::ReadableSize;
160    use common_telemetry::debug;
161    use futures::stream::FuturesUnordered;
162    use futures::{FutureExt, StreamExt};
163    use rskafka::client::error::{Error as ClientError, RequestContext};
164    use rskafka::client::partition::Compression;
165    use rskafka::protocol::error::Error as ProtocolError;
166    use rskafka::record::Record;
167
168    use super::*;
169    use crate::kafka::producer::OrderedBatchProducer;
170    use crate::kafka::test_util::record;
171
172    #[derive(Debug)]
173    struct MockClient {
174        error: Option<ProtocolError>,
175        panic: Option<String>,
176        delay: Duration,
177        batch_sizes: Mutex<Vec<usize>>,
178    }
179
180    #[async_trait::async_trait]
181    impl ProducerClient for MockClient {
182        async fn produce(
183            &self,
184            records: Vec<Record>,
185            _compression: Compression,
186        ) -> rskafka::client::error::Result<ProduceResult> {
187            tokio::time::sleep(self.delay).await;
188
189            if let Some(e) = self.error {
190                return Err(ClientError::ServerError {
191                    protocol_error: e,
192                    error_message: None,
193                    request: RequestContext::Partition("foo".into(), 1),
194                    response: None,
195                    is_virtual: false,
196                });
197            }
198
199            if let Some(p) = self.panic.as_ref() {
200                panic!("{}", p);
201            }
202
203            let mut batch_sizes = self.batch_sizes.lock().unwrap();
204            let offset_base = batch_sizes.iter().sum::<usize>();
205            let offsets = (0..records.len())
206                .map(|x| (x + offset_base) as i64)
207                .collect();
208            batch_sizes.push(records.len());
209            debug!("Return offsets: {offsets:?}");
210            Ok(ProduceResult {
211                offsets,
212                encoded_request_size: 0,
213            })
214        }
215
216        async fn get_offset(&self, _at: OffsetAt) -> rskafka::client::error::Result<i64> {
217            todo!()
218        }
219    }
220
221    #[tokio::test]
222    async fn test_producer() {
223        common_telemetry::init_default_ut_logging();
224        let record = record();
225        let delay = Duration::from_secs(0);
226        let client = Arc::new(MockClient {
227            error: None,
228            panic: None,
229            delay,
230            batch_sizes: Default::default(),
231        });
232        let provider = Arc::new(KafkaProvider::new(String::new()));
233        let producer = OrderedBatchProducer::new(
234            OrderedBatchProducer::channel(),
235            provider,
236            client.clone(),
237            Compression::NoCompression,
238            ReadableSize((record.approximate_size() * 2) as u64).as_bytes() as usize,
239            Arc::new(DashMap::new()),
240        );
241
242        // Produces 3 records
243        let handle = producer
244            .produce(vec![record.clone(), record.clone(), record.clone()])
245            .await
246            .unwrap();
247        assert_eq!(handle.wait().await.unwrap(), 2);
248        assert_eq!(client.batch_sizes.lock().unwrap().as_slice(), &[2, 1]);
249
250        // Produces 2 records
251        let handle = producer
252            .produce(vec![record.clone(), record.clone()])
253            .await
254            .unwrap();
255        assert_eq!(handle.wait().await.unwrap(), 4);
256        assert_eq!(client.batch_sizes.lock().unwrap().as_slice(), &[2, 1, 2]);
257
258        // Produces 1 records
259        let handle = producer.produce(vec![record.clone()]).await.unwrap();
260        assert_eq!(handle.wait().await.unwrap(), 5);
261        assert_eq!(client.batch_sizes.lock().unwrap().as_slice(), &[2, 1, 2, 1]);
262    }
263
264    #[tokio::test]
265    async fn test_producer_client_error() {
266        let record = record();
267        let client = Arc::new(MockClient {
268            error: Some(ProtocolError::NetworkException),
269            panic: None,
270            delay: Duration::from_millis(1),
271            batch_sizes: Default::default(),
272        });
273        let provider = Arc::new(KafkaProvider::new(String::new()));
274        let producer = OrderedBatchProducer::new(
275            OrderedBatchProducer::channel(),
276            provider,
277            client.clone(),
278            Compression::NoCompression,
279            ReadableSize((record.approximate_size() * 2) as u64).as_bytes() as usize,
280            Arc::new(DashMap::new()),
281        );
282
283        let mut futures = FuturesUnordered::new();
284        futures.push(
285            producer
286                .produce(vec![record.clone(), record.clone(), record.clone()])
287                .await
288                .unwrap()
289                .wait(),
290        );
291        futures.push(
292            producer
293                .produce(vec![record.clone(), record.clone()])
294                .await
295                .unwrap()
296                .wait(),
297        );
298        futures.push(producer.produce(vec![record.clone()]).await.unwrap().wait());
299
300        futures.next().await.unwrap().unwrap_err();
301        futures.next().await.unwrap().unwrap_err();
302        futures.next().await.unwrap().unwrap_err();
303    }
304
305    #[tokio::test]
306    async fn test_producer_cancel() {
307        let record = record();
308        let client = Arc::new(MockClient {
309            error: None,
310            panic: None,
311            delay: Duration::from_millis(1),
312            batch_sizes: Default::default(),
313        });
314
315        let provider = Arc::new(KafkaProvider::new(String::new()));
316        let producer = OrderedBatchProducer::new(
317            OrderedBatchProducer::channel(),
318            provider,
319            client.clone(),
320            Compression::NoCompression,
321            ReadableSize((record.approximate_size() * 2) as u64).as_bytes() as usize,
322            Arc::new(DashMap::new()),
323        );
324
325        let a = producer
326            .produce(vec![record.clone(), record.clone(), record.clone()])
327            .await
328            .unwrap()
329            .wait()
330            .fuse();
331
332        let b = producer.produce(vec![record]).await.unwrap().wait().fuse();
333
334        let mut b = Box::pin(b);
335
336        {
337            // Cancel a when it exits this block
338            let mut a = Box::pin(a);
339
340            // Select biased to encourage `a` to be the one with the linger that
341            // expires first and performs the produce operation
342            futures::select_biased! {
343                _ = &mut a => panic!("a should not have flushed"),
344                _ = &mut b => panic!("b should not have flushed"),
345                _ = tokio::time::sleep(Duration::from_millis(1)).fuse() => {},
346            }
347        }
348
349        // But `b` should still complete successfully
350        tokio::time::timeout(Duration::from_secs(1), b)
351            .await
352            .unwrap()
353            .unwrap();
354
355        assert_eq!(
356            client
357                .batch_sizes
358                .lock()
359                .unwrap()
360                .as_slice()
361                .iter()
362                .sum::<usize>(),
363            4
364        );
365    }
366}
log_store/kafka/producer.rs

log_store/kafka/
producer.rs
