log_store/kafka/worker/
produce.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 common_telemetry::tracing::warn;
16use tokio::sync::oneshot;
17
18use crate::kafka::worker::{
19    BackgroundProducerWorker, PendingRequest, ProduceRequest, ProduceResultReceiver,
20};
21
22impl BackgroundProducerWorker {
23    /// Aggregates records into batches, ensuring that the size of each batch does not exceed a specified maximum (`max_batch_bytes`).
24    ///
25    /// ## Panic
26    /// Panic if any [Record]'s `approximate_size` > `max_batch_bytes`.
27    pub(crate) fn aggregate_records(
28        &self,
29        requests: &mut Vec<ProduceRequest>,
30        max_batch_bytes: usize,
31    ) -> Vec<PendingRequest> {
32        let mut records_buffer = vec![];
33        let mut batch_size = 0;
34        let mut pending_requests = Vec::with_capacity(requests.len());
35
36        for ProduceRequest { batch, sender } in std::mem::take(requests) {
37            let mut receiver = ProduceResultReceiver::default();
38            for record in batch {
39                assert!(record.approximate_size() <= max_batch_bytes);
40                // Yields the `PendingRequest` if buffer is full.
41                if batch_size + record.approximate_size() > max_batch_bytes {
42                    let (tx, rx) = oneshot::channel();
43                    pending_requests.push(PendingRequest {
44                        batch: std::mem::take(&mut records_buffer),
45                        size: batch_size,
46                        sender: tx,
47                    });
48                    batch_size = 0;
49                    receiver.add_receiver(rx);
50                }
51
52                batch_size += record.approximate_size();
53                records_buffer.push(record);
54            }
55            // The remaining records.
56            if batch_size > 0 {
57                // Yields `PendingRequest`
58                let (tx, rx) = oneshot::channel();
59                pending_requests.push(PendingRequest {
60                    batch: std::mem::take(&mut records_buffer),
61                    size: batch_size,
62                    sender: tx,
63                });
64                batch_size = 0;
65                receiver.add_receiver(rx);
66            }
67
68            if sender.send(receiver).is_err() {
69                warn!("The Receiver of ProduceResultReceiver is dropped");
70            }
71        }
72        pending_requests
73    }
74}
log_store/kafka/worker/produce.rs

log_store/kafka/worker/
produce.rs
