log_store/kafka/util/
record.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::collections::HashMap;
16use std::sync::Arc;
17
18use rskafka::record::Record as KafkaRecord;
19use serde::{Deserialize, Serialize};
20use snafu::{OptionExt, ResultExt, ensure};
21use store_api::logstore::entry::{Entry, MultiplePartEntry, MultiplePartHeader, NaiveEntry};
22use store_api::logstore::provider::{KafkaProvider, Provider};
23use store_api::storage::RegionId;
24
25use crate::error::{
26    DecodeJsonSnafu, EncodeJsonSnafu, IllegalSequenceSnafu, MetaLengthExceededLimitSnafu,
27    MissingKeySnafu, MissingValueSnafu, Result,
28};
29use crate::kafka::{EntryId, NamespaceImpl};
30
31/// The current version of Record.
32pub(crate) const VERSION: u32 = 0;
33
34/// The estimated size in bytes of a serialized RecordMeta.
35/// A record is guaranteed to have sizeof(meta) + sizeof(data) <= max_batch_byte - ESTIMATED_META_SIZE.
36pub(crate) const ESTIMATED_META_SIZE: usize = 256;
37
38/// The type of a record.
39///
40/// - If the entry is able to fit into a Kafka record, it's converted into a Full record.
41///
42/// - If the entry is too large to fit into a Kafka record, it's converted into a collection of records.
43///
44/// Those records must contain exactly one First record and one Last record, and potentially several
45///   Middle records. There may be no Middle record.
46#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
47pub enum RecordType {
48    /// The record is self-contained, i.e. an entry's data is fully stored into this record.
49    Full,
50    /// The record contains the first part of an entry's data.
51    First,
52    /// The record contains one of the middle parts of an entry's data.
53    /// The sequence of the record is identified by the inner field.
54    Middle(usize),
55    /// The record contains the last part of an entry's data.
56    Last,
57}
58
59/// The metadata of a record.
60#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
61pub struct RecordMeta {
62    /// The version of the record. Used for backward compatibility.
63    version: u32,
64    /// The type of the record.
65    pub tp: RecordType,
66    /// The id of the entry the record associated with.
67    pub entry_id: EntryId,
68    /// The namespace of the entry the record associated with.
69    pub ns: NamespaceImpl,
70}
71
72/// The minimal storage unit in the Kafka log store.
73///
74/// An entry will be first converted into several Records before producing.
75/// If an entry is able to fit into a KafkaRecord, it converts to a single Record.
76/// If otherwise an entry cannot fit into a KafkaRecord, it will be split into a collection of Records.
77///
78/// A KafkaRecord is the minimal storage unit used by Kafka client and Kafka server.
79/// The Kafka client produces KafkaRecords and consumes KafkaRecords, and Kafka server stores
80/// a collection of KafkaRecords.
81#[derive(Debug, Clone, PartialEq)]
82pub(crate) struct Record {
83    /// The metadata of the record.
84    pub(crate) meta: RecordMeta,
85    /// The payload of the record.
86    data: Vec<u8>,
87}
88
89impl TryFrom<Record> for KafkaRecord {
90    type Error = crate::error::Error;
91
92    fn try_from(record: Record) -> Result<Self> {
93        let key = serde_json::to_vec(&record.meta).context(EncodeJsonSnafu)?;
94        ensure!(
95            key.len() < ESTIMATED_META_SIZE,
96            MetaLengthExceededLimitSnafu {
97                limit: ESTIMATED_META_SIZE,
98                actual: key.len()
99            }
100        );
101        Ok(KafkaRecord {
102            key: Some(key),
103            value: Some(record.data),
104            timestamp: chrono::Utc::now(),
105            headers: Default::default(),
106        })
107    }
108}
109
110// TODO(niebayes): improve the performance of decoding kafka record.
111impl Record {
112    /// Decodes a record from the key and value of a [`KafkaRecord`].
113    ///
114    /// The key is borrowed so that callers can inspect it if decoding fails, e.g. to
115    /// tell a foreign key apart from a [`RecordMeta`] written by an incompatible version
116    /// of the log store.
117    pub(crate) fn try_from_parts(key: Option<&[u8]>, value: Option<Vec<u8>>) -> Result<Self> {
118        let key = key.context(MissingKeySnafu)?;
119        let meta = serde_json::from_slice(key).context(DecodeJsonSnafu)?;
120        let data = value.context(MissingValueSnafu)?;
121        Ok(Self { meta, data })
122    }
123}
124
125/// Returns whether a key is a JSON object with a `version` field, i.e. whether it may have
126/// been written as a [`RecordMeta`] key.
127///
128/// A `RecordMeta` key that cannot be decoded may come from an incompatible version of the
129/// log store, so the reader must not treat such a key as a foreign record.
130pub(crate) fn is_record_meta_key(key: &[u8]) -> bool {
131    matches!(
132        serde_json::from_slice::<serde_json::Value>(key),
133        Ok(serde_json::Value::Object(fields)) if fields.contains_key("version")
134    )
135}
136
137pub(crate) fn convert_to_kafka_records(entry: Entry) -> Result<Vec<KafkaRecord>> {
138    match entry {
139        Entry::Naive(entry) => Ok(vec![KafkaRecord::try_from(Record {
140            meta: RecordMeta {
141                version: VERSION,
142                tp: RecordType::Full,
143                // TODO(weny): refactor the record meta.
144                entry_id: 0,
145                ns: NamespaceImpl {
146                    region_id: entry.region_id.as_u64(),
147                    // TODO(weny): refactor the record meta.
148                    topic: String::new(),
149                },
150            },
151            data: entry.data,
152        })?]),
153        Entry::MultiplePart(entry) => {
154            let mut entries = Vec::with_capacity(entry.parts.len());
155
156            for (idx, part) in entry.parts.into_iter().enumerate() {
157                let tp = match entry.headers[idx] {
158                    MultiplePartHeader::First => RecordType::First,
159                    MultiplePartHeader::Middle(i) => RecordType::Middle(i),
160                    MultiplePartHeader::Last => RecordType::Last,
161                };
162                entries.push(KafkaRecord::try_from(Record {
163                    meta: RecordMeta {
164                        version: VERSION,
165                        tp,
166                        // TODO(weny): refactor the record meta.
167                        entry_id: 0,
168                        ns: NamespaceImpl {
169                            region_id: entry.region_id.as_u64(),
170                            topic: String::new(),
171                        },
172                    },
173                    data: part,
174                })?)
175            }
176            Ok(entries)
177        }
178    }
179}
180
181fn convert_to_naive_entry(provider: Arc<KafkaProvider>, record: Record) -> Entry {
182    let region_id = RegionId::from_u64(record.meta.ns.region_id);
183
184    Entry::Naive(NaiveEntry {
185        provider: Provider::Kafka(provider),
186        region_id,
187        entry_id: record.meta.entry_id,
188        data: record.data,
189    })
190}
191
192fn convert_to_multiple_entry(
193    provider: Arc<KafkaProvider>,
194    region_id: RegionId,
195    records: Vec<Record>,
196) -> Entry {
197    let mut headers = Vec::with_capacity(records.len());
198    let mut parts = Vec::with_capacity(records.len());
199    let entry_id = records.last().map(|r| r.meta.entry_id).unwrap_or_default();
200
201    for record in records {
202        let header = match record.meta.tp {
203            RecordType::Full => unreachable!(),
204            RecordType::First => MultiplePartHeader::First,
205            RecordType::Middle(i) => MultiplePartHeader::Middle(i),
206            RecordType::Last => MultiplePartHeader::Last,
207        };
208        headers.push(header);
209        parts.push(record.data);
210    }
211
212    Entry::MultiplePart(MultiplePartEntry {
213        provider: Provider::Kafka(provider),
214        region_id,
215        entry_id,
216        headers,
217        parts,
218    })
219}
220
221/// Constructs entries from `buffered_records`
222pub fn remaining_entries(
223    provider: &Arc<KafkaProvider>,
224    buffered_records: &mut HashMap<RegionId, Vec<Record>>,
225) -> Option<Vec<Entry>> {
226    if buffered_records.is_empty() {
227        None
228    } else {
229        let mut entries = Vec::with_capacity(buffered_records.len());
230        for (region_id, records) in buffered_records.drain() {
231            entries.push(convert_to_multiple_entry(
232                provider.clone(),
233                region_id,
234                records,
235            ));
236        }
237        Some(entries)
238    }
239}
240
241/// For type of [Entry::Naive] Entry:
242/// - Emits a [RecordType::Full] type record immediately.
243///
244/// For type of [Entry::MultiplePart] Entry:
245/// - Emits a complete [Entry] immediately when a [RecordType::Last] record arrives with
246///   buffered parts.
247/// - Emits an incomplete [Entry] when a new [RecordType::First] record replaces
248///   buffered incomplete records for the same [RegionId].
249///
250/// **Discarded records:**
251/// - A standalone [RecordType::First] record is discarded when another [RecordType::First]
252///   record with the same payload for the same [RegionId] arrives.
253/// - A [RecordType::Last] record without buffered parts is discarded.
254///
255/// A trailing incomplete entry is emitted by [`remaining_entries`].
256pub(crate) fn maybe_emit_entry(
257    provider: &Arc<KafkaProvider>,
258    record: Record,
259    buffered_records: &mut HashMap<RegionId, Vec<Record>>,
260) -> Result<Option<Entry>> {
261    let mut entry = None;
262    match record.meta.tp {
263        RecordType::Full => entry = Some(convert_to_naive_entry(provider.clone(), record)),
264        RecordType::First => {
265            let region_id = record.meta.ns.region_id.into();
266            let duplicate = buffered_records.get(&region_id).is_some_and(|records| {
267                records.len() == 1
268                    && records[0].meta.tp == RecordType::First
269                    && records[0].data == record.data
270            });
271            if let Some(records) = buffered_records.insert(region_id, vec![record]) {
272                if duplicate {
273                    // A duplicate standalone First cannot form a complete entry.
274                    return Ok(None);
275                }
276                // Incomplete entry
277                entry = Some(convert_to_multiple_entry(
278                    provider.clone(),
279                    region_id,
280                    records,
281                ))
282            }
283        }
284        RecordType::Middle(seq) => {
285            let region_id = record.meta.ns.region_id.into();
286            let records = buffered_records.entry(region_id).or_default();
287
288            // Only validate complete entries.
289            if !records.is_empty() {
290                // Safety: the records are guaranteed not empty if the key exists.
291                let last_record = records.last().unwrap();
292                if matches!(last_record.meta.tp, RecordType::Middle(last_seq) if last_seq == seq)
293                    && last_record.data == record.data
294                {
295                    return Ok(None);
296                }
297                let legal = match last_record.meta.tp {
298                    // Legal if this record follows a First record.
299                    RecordType::First => seq == 1,
300                    // Legal if this record follows a Middle record just prior to this record.
301                    RecordType::Middle(last_seq) => last_seq + 1 == seq,
302                    // Illegal sequence.
303                    _ => false,
304                };
305                ensure!(
306                    legal,
307                    IllegalSequenceSnafu {
308                        error: format!(
309                            "Illegal sequence of a middle record, last record: {:?}, incoming record: {:?}",
310                            last_record.meta.tp, record.meta.tp
311                        )
312                    }
313                );
314            }
315
316            records.push(record);
317        }
318        RecordType::Last => {
319            let region_id = record.meta.ns.region_id.into();
320            if let Some(mut records) = buffered_records.remove(&region_id) {
321                records.push(record);
322                entry = Some(convert_to_multiple_entry(
323                    provider.clone(),
324                    region_id,
325                    records,
326                ));
327            } else {
328                // Intentionally discard a Last record without buffered parts. It is either a
329                // duplicate Last record or the tail of an incomplete multipart entry.
330            }
331        }
332    }
333    Ok(entry)
334}
335
336#[cfg(test)]
337mod tests {
338    use std::assert_matches;
339    use std::sync::Arc;
340
341    use super::*;
342    use crate::error;
343
344    fn new_test_record(tp: RecordType, entry_id: EntryId, region_id: u64, data: Vec<u8>) -> Record {
345        Record {
346            meta: RecordMeta {
347                version: VERSION,
348                tp,
349                ns: NamespaceImpl {
350                    region_id,
351                    topic: "greptimedb_wal_topic".to_string(),
352                },
353                entry_id,
354            },
355            data,
356        }
357    }
358
359    #[test]
360    fn test_maybe_emit_entry_emit_naive_entry() {
361        let provider = Arc::new(KafkaProvider::new("my_topic".to_string()));
362        let region_id = RegionId::new(1, 1);
363        let mut buffer = HashMap::new();
364        let record = new_test_record(RecordType::Full, 1, region_id.as_u64(), vec![1; 100]);
365        let entry = maybe_emit_entry(&provider, record, &mut buffer)
366            .unwrap()
367            .unwrap();
368        assert_eq!(
369            entry,
370            Entry::Naive(NaiveEntry {
371                provider: Provider::Kafka(provider),
372                region_id,
373                entry_id: 1,
374                data: vec![1; 100]
375            })
376        );
377    }
378
379    #[test]
380    fn test_maybe_emit_entry_emit_incomplete_entry() {
381        let provider = Arc::new(KafkaProvider::new("my_topic".to_string()));
382        let region_id = RegionId::new(1, 1);
383        // `First` overwrite `First`
384        let mut buffer = HashMap::new();
385        let record = new_test_record(RecordType::First, 1, region_id.as_u64(), vec![1; 100]);
386        assert!(
387            maybe_emit_entry(&provider, record, &mut buffer)
388                .unwrap()
389                .is_none()
390        );
391        let record = new_test_record(RecordType::First, 2, region_id.as_u64(), vec![1; 100]);
392        // An identical First is a duplicate and does not emit an incomplete entry.
393        assert!(
394            maybe_emit_entry(&provider, record, &mut buffer)
395                .unwrap()
396                .is_none()
397        );
398        assert_eq!(buffer[&region_id][0].meta.entry_id, 2);
399
400        let record = new_test_record(RecordType::First, 3, region_id.as_u64(), vec![2; 100]);
401        let incomplete_entry = maybe_emit_entry(&provider, record, &mut buffer)
402            .unwrap()
403            .unwrap();
404
405        assert_eq!(
406            incomplete_entry,
407            Entry::MultiplePart(MultiplePartEntry {
408                provider: Provider::Kafka(provider.clone()),
409                region_id,
410                entry_id: 2,
411                headers: vec![MultiplePartHeader::First],
412                parts: vec![vec![1; 100]],
413            })
414        );
415
416        // `Last` after `None`
417        let mut buffer = HashMap::new();
418        let record = new_test_record(RecordType::Last, 1, region_id.as_u64(), vec![1; 100]);
419        assert!(
420            maybe_emit_entry(&provider, record, &mut buffer)
421                .unwrap()
422                .is_none()
423        );
424        assert!(buffer.is_empty());
425
426        // `First` overwrite `Middle(0)`
427        let mut buffer = HashMap::new();
428        let record = new_test_record(RecordType::Middle(0), 1, region_id.as_u64(), vec![1; 100]);
429        assert!(
430            maybe_emit_entry(&provider, record, &mut buffer)
431                .unwrap()
432                .is_none()
433        );
434        let record = new_test_record(RecordType::First, 2, region_id.as_u64(), vec![2; 100]);
435        let incomplete_entry = maybe_emit_entry(&provider, record, &mut buffer)
436            .unwrap()
437            .unwrap();
438
439        assert_eq!(
440            incomplete_entry,
441            Entry::MultiplePart(MultiplePartEntry {
442                provider: Provider::Kafka(provider),
443                region_id,
444                entry_id: 1,
445                headers: vec![MultiplePartHeader::Middle(0)],
446                parts: vec![vec![1; 100]],
447            })
448        );
449    }
450
451    #[test]
452    fn test_maybe_emit_entry_illegal_seq() {
453        let provider = Arc::new(KafkaProvider::new("my_topic".to_string()));
454        let region_id = RegionId::new(1, 1);
455        let mut buffer = HashMap::new();
456        let record = new_test_record(RecordType::First, 1, region_id.as_u64(), vec![1; 100]);
457        assert!(
458            maybe_emit_entry(&provider, record, &mut buffer)
459                .unwrap()
460                .is_none()
461        );
462        let record = new_test_record(RecordType::Middle(2), 1, region_id.as_u64(), vec![2; 100]);
463        let err = maybe_emit_entry(&provider, record, &mut buffer).unwrap_err();
464        assert_matches!(err, error::Error::IllegalSequence { .. });
465
466        let mut buffer = HashMap::new();
467        let record = new_test_record(RecordType::First, 1, region_id.as_u64(), vec![1; 100]);
468        assert!(
469            maybe_emit_entry(&provider, record, &mut buffer)
470                .unwrap()
471                .is_none()
472        );
473        let record = new_test_record(RecordType::Middle(1), 1, region_id.as_u64(), vec![2; 100]);
474        assert!(
475            maybe_emit_entry(&provider, record, &mut buffer)
476                .unwrap()
477                .is_none()
478        );
479        let record = new_test_record(RecordType::Middle(3), 1, region_id.as_u64(), vec![2; 100]);
480        let err = maybe_emit_entry(&provider, record, &mut buffer).unwrap_err();
481        assert_matches!(err, error::Error::IllegalSequence { .. });
482    }
483
484    #[test]
485    fn test_maybe_emit_entry_deduplicates_multipart_records() {
486        let provider = Arc::new(KafkaProvider::new("my_topic".to_string()));
487        let region_id = RegionId::new(1, 1);
488        let mut buffer = HashMap::new();
489
490        for record in [
491            new_test_record(RecordType::First, 1, region_id.as_u64(), vec![1; 100]),
492            new_test_record(RecordType::Middle(1), 1, region_id.as_u64(), vec![2; 100]),
493            new_test_record(RecordType::Middle(1), 1, region_id.as_u64(), vec![2; 100]),
494        ] {
495            assert!(
496                maybe_emit_entry(&provider, record, &mut buffer)
497                    .unwrap()
498                    .is_none()
499            );
500        }
501
502        let last = new_test_record(RecordType::Last, 1, region_id.as_u64(), vec![3; 100]);
503        let entry = maybe_emit_entry(&provider, last, &mut buffer)
504            .unwrap()
505            .unwrap();
506
507        assert_eq!(
508            entry,
509            Entry::MultiplePart(MultiplePartEntry {
510                provider: Provider::Kafka(provider.clone()),
511                region_id,
512                entry_id: 1,
513                headers: vec![
514                    MultiplePartHeader::First,
515                    MultiplePartHeader::Middle(1),
516                    MultiplePartHeader::Last,
517                ],
518                parts: vec![vec![1; 100], vec![2; 100], vec![3; 100]],
519            })
520        );
521        let duplicate = new_test_record(RecordType::Last, 1, region_id.as_u64(), vec![3; 100]);
522        assert!(
523            maybe_emit_entry(&provider, duplicate, &mut buffer)
524                .unwrap()
525                .is_none()
526        );
527    }
528
529    #[test]
530    fn test_maybe_emit_entry_rejects_conflicting_duplicate_middle_record() {
531        let provider = Arc::new(KafkaProvider::new("my_topic".to_string()));
532        let region_id = RegionId::new(1, 1);
533        let mut buffer = HashMap::new();
534
535        for record in [
536            new_test_record(RecordType::First, 1, region_id.as_u64(), vec![1; 100]),
537            new_test_record(RecordType::Middle(1), 1, region_id.as_u64(), vec![2; 100]),
538        ] {
539            assert!(
540                maybe_emit_entry(&provider, record, &mut buffer)
541                    .unwrap()
542                    .is_none()
543            );
544        }
545
546        let duplicate = new_test_record(RecordType::Middle(1), 1, region_id.as_u64(), vec![3; 100]);
547        let err = maybe_emit_entry(&provider, duplicate, &mut buffer).unwrap_err();
548
549        assert_matches!(err, error::Error::IllegalSequence { .. });
550    }
551
552    #[test]
553    fn test_maybe_emit_entry_preserves_multipart_entry_order_before_full_record() {
554        let provider = Arc::new(KafkaProvider::new("my_topic".to_string()));
555        let region_id = RegionId::new(1, 1);
556        let mut buffer = HashMap::new();
557
558        let first = new_test_record(RecordType::First, 1, region_id.as_u64(), vec![1; 100]);
559        assert!(
560            maybe_emit_entry(&provider, first, &mut buffer)
561                .unwrap()
562                .is_none()
563        );
564        let last = new_test_record(RecordType::Last, 2, region_id.as_u64(), vec![2; 100]);
565        let multipart_entry = maybe_emit_entry(&provider, last, &mut buffer)
566            .unwrap()
567            .unwrap();
568        let full = new_test_record(RecordType::Full, 3, region_id.as_u64(), vec![3; 100]);
569        let full_entry = maybe_emit_entry(&provider, full, &mut buffer)
570            .unwrap()
571            .unwrap();
572
573        assert_eq!(multipart_entry.entry_id(), 2);
574        assert_eq!(full_entry.entry_id(), 3);
575    }
576
577    #[test]
578    fn test_is_record_meta_key() {
579        let meta = RecordMeta {
580            version: VERSION,
581            tp: RecordType::Full,
582            entry_id: 1,
583            ns: NamespaceImpl {
584                region_id: RegionId::new(1, 1).as_u64(),
585                topic: "greptimedb_wal_topic".to_string(),
586            },
587        };
588        assert!(is_record_meta_key(&serde_json::to_vec(&meta).unwrap()));
589        // A `RecordMeta` key this version cannot decode, e.g. one with a new `RecordType`
590        // variant.
591        assert!(is_record_meta_key(
592            br#"{"version": 0, "tp": "Compressed", "entry_id": 0, "ns": {"region_id": 1, "topic": ""}}"#
593        ));
594        // Keys the log store never writes.
595        assert!(!is_record_meta_key(b"not-a-json-key"));
596        assert!(!is_record_meta_key(b"[1, 2, 3]"));
597        assert!(!is_record_meta_key(br#"{"hello": "world"}"#));
598    }
599
600    #[test]
601    fn test_meta_size() {
602        let meta = RecordMeta {
603            version: VERSION,
604            tp: RecordType::Middle(usize::MAX),
605            entry_id: u64::MAX,
606            ns: NamespaceImpl {
607                region_id: RegionId::new(u32::MAX, u32::MAX).as_u64(),
608                topic: format!("greptime_kafka_cluster/1024/2048/{}", uuid::Uuid::new_v4()),
609            },
610        };
611        let serialized = serde_json::to_vec(&meta).unwrap();
612        // The len of serialized data is 202.
613        assert!(serialized.len() < ESTIMATED_META_SIZE);
614    }
615}
log_store/kafka/util/record.rs

log_store/kafka/util/
record.rs
