log_store/kafka/
worker.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
15pub(crate) mod fetch_latest_offset;
16pub(crate) mod flush;
17pub(crate) mod produce;
18
19use std::sync::Arc;
20
21use common_telemetry::debug;
22use dashmap::DashMap;
23use futures::future::try_join_all;
24use rskafka::client::partition::Compression;
25use rskafka::record::Record;
26use snafu::{OptionExt, ResultExt};
27use store_api::logstore::provider::KafkaProvider;
28use tokio::sync::mpsc::Receiver;
29use tokio::sync::oneshot::{self};
30
31use crate::error::{self, NoMaxValueSnafu, Result};
32use crate::kafka::log_store::TopicStat;
33use crate::kafka::producer::ProducerClient;
34
35pub(crate) enum WorkerRequest {
36    Produce(ProduceRequest),
37    FetchLatestOffset,
38}
39
40impl WorkerRequest {
41    pub(crate) fn new_produce_request(batch: Vec<Record>) -> (WorkerRequest, ProduceResultHandle) {
42        let (tx, rx) = oneshot::channel();
43
44        (
45            WorkerRequest::Produce(ProduceRequest { batch, sender: tx }),
46            ProduceResultHandle { receiver: rx },
47        )
48    }
49}
50
51pub(crate) struct ProduceRequest {
52    batch: Vec<Record>,
53    sender: oneshot::Sender<ProduceResultReceiver>,
54}
55
56/// Receives the committed offsets when data has been committed to Kafka
57/// or an unrecoverable error has been encountered.
58pub(crate) struct ProduceResultHandle {
59    receiver: oneshot::Receiver<ProduceResultReceiver>,
60}
61
62impl ProduceResultHandle {
63    /// Waits for the data has been committed to Kafka.
64    /// Returns the **max** committed offsets.
65    pub(crate) async fn wait(self) -> Result<u64> {
66        self.receiver
67            .await
68            .context(error::WaitProduceResultReceiverSnafu)?
69            .wait()
70            .await
71    }
72}
73
74#[derive(Default)]
75pub(crate) struct ProduceResultReceiver {
76    receivers: Vec<oneshot::Receiver<Result<Vec<i64>>>>,
77}
78
79impl ProduceResultReceiver {
80    fn add_receiver(&mut self, receiver: oneshot::Receiver<Result<Vec<i64>>>) {
81        self.receivers.push(receiver)
82    }
83
84    async fn wait(self) -> Result<u64> {
85        Ok(try_join_all(self.receivers)
86            .await
87            .into_iter()
88            .flatten()
89            .collect::<Result<Vec<_>>>()?
90            .into_iter()
91            .flatten()
92            .max()
93            .context(NoMaxValueSnafu)? as u64)
94    }
95}
96
97pub(crate) struct PendingRequest {
98    batch: Vec<Record>,
99    size: usize,
100    sender: oneshot::Sender<Result<Vec<i64>>>,
101}
102
103pub(crate) struct BackgroundProducerWorker {
104    pub(crate) provider: Arc<KafkaProvider>,
105    /// The [`ProducerClient`].
106    pub(crate) client: Arc<dyn ProducerClient>,
107    // The compression configuration.
108    pub(crate) compression: Compression,
109    /// Receiver of [ProduceRequest].
110    pub(crate) receiver: Receiver<WorkerRequest>,
111    /// Max batch size for a worker to handle requests.
112    pub(crate) request_batch_size: usize,
113    /// Max bytes size for a single flush.
114    pub(crate) max_batch_bytes: usize,
115    /// The stats of each topic.
116    pub(crate) topic_stats: Arc<DashMap<Arc<KafkaProvider>, TopicStat>>,
117}
118
119impl BackgroundProducerWorker {
120    pub(crate) async fn run(&mut self) {
121        let mut buffer = Vec::with_capacity(self.request_batch_size);
122        loop {
123            match self.receiver.recv().await {
124                Some(req) => {
125                    buffer.clear();
126                    buffer.push(req);
127                    for _ in 1..self.request_batch_size {
128                        match self.receiver.try_recv() {
129                            Ok(req) => buffer.push(req),
130                            Err(_) => break,
131                        }
132                    }
133                    self.handle_requests(&mut buffer).await;
134                }
135                None => {
136                    debug!("The sender is dropped, BackgroundProducerWorker exited");
137                    // Exits the loop if the `sender` is dropped.
138                    break;
139                }
140            }
141        }
142    }
143
144    async fn handle_requests(&mut self, buffer: &mut Vec<WorkerRequest>) {
145        let mut produce_requests = Vec::with_capacity(buffer.len());
146        for req in buffer.drain(..) {
147            match req {
148                WorkerRequest::Produce(req) => produce_requests.push(req),
149                WorkerRequest::FetchLatestOffset => {
150                    self.fetch_latest_offset().await;
151                }
152            }
153        }
154
155        let pending_requests = self.aggregate_records(&mut produce_requests, self.max_batch_bytes);
156        self.try_flush_pending_requests(pending_requests).await;
157    }
158}
log_store/kafka/worker.rs

log_store/kafka/
worker.rs
