1use std::{
2 fmt,
3 num::{NonZeroU64, NonZeroUsize},
4 path::{Path, PathBuf},
5 slice,
6};
7
8use serde::{Deserialize, Deserializer, Serialize, de};
9use snafu::{ResultExt, Snafu};
10use tracing::Span;
11use vector_common::{config::ComponentKey, finalization::Finalizable};
12use vector_config::configurable_component;
13
14use crate::{
15 Bufferable, WhenFull,
16 topology::{
17 builder::{TopologyBuilder, TopologyError},
18 channel::{BufferReceiver, BufferSender},
19 },
20 variants::{DiskV2Buffer, MemoryBuffer},
21};
22
23#[derive(Debug, Snafu)]
24pub enum BufferBuildError {
25 #[snafu(display("the configured buffer type requires `data_dir` be specified"))]
26 RequiresDataDir,
27 #[snafu(display("error occurred when building buffer: {}", source))]
28 FailedToBuildTopology { source: TopologyError },
29 #[snafu(display("`max_events` must be greater than zero"))]
30 InvalidMaxEvents,
31}
32
33#[derive(Deserialize, Serialize)]
34enum BufferTypeKind {
35 #[serde(rename = "memory")]
36 Memory,
37 #[serde(rename = "disk")]
38 DiskV2,
39}
40
41const ALL_FIELDS: [&str; 4] = ["type", "max_events", "max_size", "when_full"];
42
43struct BufferTypeVisitor;
44
45impl BufferTypeVisitor {
46 fn visit_map_impl<'de, A>(mut map: A) -> Result<BufferType, A::Error>
47 where
48 A: de::MapAccess<'de>,
49 {
50 let mut kind: Option<BufferTypeKind> = None;
51 let mut max_events: Option<NonZeroUsize> = None;
52 let mut max_size: Option<NonZeroU64> = None;
53 let mut when_full: Option<WhenFull> = None;
54 while let Some(key) = map.next_key::<String>()? {
55 match key.as_str() {
56 "type" => {
57 if kind.is_some() {
58 return Err(de::Error::duplicate_field("type"));
59 }
60 kind = Some(map.next_value()?);
61 }
62 "max_events" => {
63 if max_events.is_some() {
64 return Err(de::Error::duplicate_field("max_events"));
65 }
66 max_events = Some(map.next_value()?);
67 }
68 "max_size" => {
69 if max_size.is_some() {
70 return Err(de::Error::duplicate_field("max_size"));
71 }
72 max_size = Some(map.next_value()?);
73 }
74 "when_full" => {
75 if when_full.is_some() {
76 return Err(de::Error::duplicate_field("when_full"));
77 }
78 when_full = Some(map.next_value()?);
79 }
80 other => {
81 return Err(de::Error::unknown_field(other, &ALL_FIELDS));
82 }
83 }
84 }
85 let kind = kind.unwrap_or(BufferTypeKind::Memory);
86 let when_full = when_full.unwrap_or_default();
87 match kind {
88 BufferTypeKind::Memory => {
89 let size = match (max_events, max_size) {
90 (Some(_), Some(_)) => {
91 return Err(de::Error::unknown_field(
92 "max_events",
93 &["type", "max_size", "when_full"],
94 ));
95 }
96 (_, Some(max_size)) => {
97 if let Ok(bounded_max_bytes) = usize::try_from(max_size.get()) {
98 MemoryBufferSize::MaxSize(NonZeroUsize::new(bounded_max_bytes).unwrap())
99 } else {
100 return Err(de::Error::invalid_value(
101 de::Unexpected::Unsigned(max_size.into()),
102 &format!(
103 "Value for max_bytes must be a positive integer <= {}",
104 usize::MAX
105 )
106 .as_str(),
107 ));
108 }
109 }
110 _ => MemoryBufferSize::MaxEvents(
111 max_events.unwrap_or_else(memory_buffer_default_max_events),
112 ),
113 };
114 Ok(BufferType::Memory { size, when_full })
115 }
116 BufferTypeKind::DiskV2 => {
117 if max_events.is_some() {
118 return Err(de::Error::unknown_field(
119 "max_events",
120 &["type", "max_size", "when_full"],
121 ));
122 }
123 Ok(BufferType::DiskV2 {
124 max_size: max_size.ok_or_else(|| de::Error::missing_field("max_size"))?,
125 when_full,
126 })
127 }
128 }
129 }
130}
131
132impl<'de> de::Visitor<'de> for BufferTypeVisitor {
133 type Value = BufferType;
134
135 fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
136 formatter.write_str("enum BufferType")
137 }
138
139 fn visit_map<A>(self, map: A) -> Result<Self::Value, A::Error>
140 where
141 A: de::MapAccess<'de>,
142 {
143 BufferTypeVisitor::visit_map_impl(map)
144 }
145}
146
147impl<'de> Deserialize<'de> for BufferType {
148 fn deserialize<D>(deserializer: D) -> Result<BufferType, D::Error>
149 where
150 D: Deserializer<'de>,
151 {
152 deserializer.deserialize_map(BufferTypeVisitor)
153 }
154}
155
156pub const fn memory_buffer_default_max_events() -> NonZeroUsize {
160 NonZeroUsize::new(500).expect("500 is non-zero")
161}
162
163#[derive(Debug)]
165pub struct DiskUsage {
166 id: ComponentKey,
167 data_dir: PathBuf,
168 max_size: NonZeroU64,
169}
170
171impl DiskUsage {
172 pub fn new(id: ComponentKey, data_dir: PathBuf, max_size: NonZeroU64) -> Self {
174 Self {
175 id,
176 data_dir,
177 max_size,
178 }
179 }
180
181 pub fn id(&self) -> &ComponentKey {
183 &self.id
184 }
185
186 pub fn max_size(&self) -> u64 {
188 self.max_size.get()
189 }
190
191 pub fn data_dir(&self) -> &Path {
193 self.data_dir.as_path()
194 }
195}
196
197#[configurable_component(no_deser)]
200#[serde(rename_all = "snake_case")]
201#[derive(Clone, Copy, Debug, PartialEq, Eq)]
202pub enum MemoryBufferSize {
203 MaxEvents(#[serde(default = "memory_buffer_default_max_events")] NonZeroUsize),
205
206 MaxSize(#[configurable(metadata(docs::type_unit = "bytes"))] NonZeroUsize),
212}
213
214#[configurable_component(no_deser)]
216#[derive(Clone, Copy, Debug, PartialEq, Eq)]
217#[serde(rename_all = "snake_case", tag = "type")]
218#[configurable(metadata(docs::enum_tag_description = "The type of buffer to use."))]
219pub enum BufferType {
220 #[configurable(title = "Events are buffered in memory.")]
225 Memory {
226 #[serde(flatten)]
228 size: MemoryBufferSize,
229
230 #[configurable(derived)]
231 #[serde(default)]
232 when_full: WhenFull,
233 },
234
235 #[configurable(title = "Events are buffered on disk.")]
242 #[serde(rename = "disk")]
243 DiskV2 {
244 #[configurable(
248 validation(range(min = 268435488)),
249 metadata(docs::type_unit = "bytes")
250 )]
251 max_size: NonZeroU64,
252
253 #[configurable(derived)]
254 #[serde(default)]
255 when_full: WhenFull,
256 },
257}
258
259impl BufferType {
260 pub fn disk_usage(
267 &self,
268 global_data_dir: Option<PathBuf>,
269 id: &ComponentKey,
270 ) -> Option<DiskUsage> {
271 match global_data_dir {
284 None => None,
285 Some(global_data_dir) => match self {
286 Self::Memory { .. } => None,
287 Self::DiskV2 { max_size, .. } => {
288 let data_dir = crate::variants::disk_v2::get_disk_v2_data_dir_path(
289 &global_data_dir,
290 id.id(),
291 );
292
293 Some(DiskUsage::new(id.clone(), data_dir, *max_size))
294 }
295 },
296 }
297 }
298
299 pub fn add_to_builder<T>(
306 &self,
307 builder: &mut TopologyBuilder<T>,
308 data_dir: Option<PathBuf>,
309 id: String,
310 ) -> Result<(), BufferBuildError>
311 where
312 T: Bufferable + Clone + Finalizable,
313 {
314 match *self {
315 BufferType::Memory { size, when_full } => {
316 builder.stage(MemoryBuffer::new(size), when_full);
317 }
318 BufferType::DiskV2 {
319 when_full,
320 max_size,
321 } => {
322 let data_dir = data_dir.ok_or(BufferBuildError::RequiresDataDir)?;
323 builder.stage(DiskV2Buffer::new(id, data_dir, max_size), when_full);
324 }
325 }
326
327 Ok(())
328 }
329}
330
331#[configurable_component]
350#[derive(Clone, Debug, PartialEq, Eq)]
351#[serde(untagged)]
352#[configurable(
353 title = "Configures the buffering behavior for this sink.",
354 description = r#"More information about the individual buffer types, and buffer behavior, can be found in the
355[Buffering Model][buffering_model] section.
356
357[buffering_model]: /docs/architecture/buffering-model/"#
358)]
359pub enum BufferConfig {
360 Single(BufferType),
362
363 Chained(Vec<BufferType>),
365}
366
367impl Default for BufferConfig {
368 fn default() -> Self {
369 Self::Single(BufferType::Memory {
370 size: MemoryBufferSize::MaxEvents(memory_buffer_default_max_events()),
371 when_full: WhenFull::default(),
372 })
373 }
374}
375
376impl BufferConfig {
377 pub fn has_disk_stage(&self) -> bool {
379 self.stages()
380 .iter()
381 .any(|stage| matches!(stage, BufferType::DiskV2 { .. }))
382 }
383
384 pub fn stages(&self) -> &[BufferType] {
386 match self {
387 Self::Single(stage) => slice::from_ref(stage),
388 Self::Chained(stages) => stages.as_slice(),
389 }
390 }
391
392 #[allow(clippy::needless_pass_by_value)]
405 pub async fn build<T>(
406 &self,
407 data_dir: Option<PathBuf>,
408 buffer_id: String,
409 span: Span,
410 ) -> Result<(BufferSender<T>, BufferReceiver<T>), BufferBuildError>
411 where
412 T: Bufferable + Clone + Finalizable,
413 {
414 let mut builder = TopologyBuilder::default();
415
416 for stage in self.stages() {
417 stage.add_to_builder(&mut builder, data_dir.clone(), buffer_id.clone())?;
418 }
419
420 builder
421 .build(buffer_id, span)
422 .await
423 .context(FailedToBuildTopologySnafu)
424 }
425}
426
427#[cfg(test)]
428mod test {
429 use std::num::{NonZeroU64, NonZeroUsize};
430
431 use crate::{BufferConfig, BufferType, MemoryBufferSize, WhenFull};
432
433 fn check_single_stage(source: &str, expected: BufferType) {
434 let config: BufferConfig = serde_yaml::from_str(source).unwrap();
435 assert_eq!(config.stages().len(), 1);
436 let actual = config.stages().first().unwrap();
437 assert_eq!(actual, &expected);
438 }
439
440 fn check_multiple_stages(source: &str, expected_stages: &[BufferType]) {
441 let config: BufferConfig = serde_yaml::from_str(source).unwrap();
442 assert_eq!(config.stages().len(), expected_stages.len());
443 for (actual, expected) in config.stages().iter().zip(expected_stages) {
444 assert_eq!(actual, expected);
445 }
446 }
447
448 const BUFFER_CONFIG_NO_MATCH_ERR: &str =
449 "data did not match any variant of untagged enum BufferConfig";
450
451 #[test]
452 fn parse_empty() {
453 let source = "";
454 let error = serde_yaml::from_str::<BufferConfig>(source).unwrap_err();
455 assert_eq!(error.to_string(), BUFFER_CONFIG_NO_MATCH_ERR);
456 }
457
458 #[test]
459 fn parse_only_invalid_keys() {
460 let source = "foo: 314";
461 let error = serde_yaml::from_str::<BufferConfig>(source).unwrap_err();
462 assert_eq!(error.to_string(), BUFFER_CONFIG_NO_MATCH_ERR);
463 }
464
465 #[test]
466 fn parse_partial_invalid_keys() {
467 let source = r"max_size: 100
468max_events: 42
469";
470 let error = serde_yaml::from_str::<BufferConfig>(source).unwrap_err();
471 assert_eq!(error.to_string(), BUFFER_CONFIG_NO_MATCH_ERR);
472 }
473
474 #[test]
475 fn parse_without_type_tag() {
476 check_single_stage(
477 r"
478 max_events: 100
479 ",
480 BufferType::Memory {
481 size: MemoryBufferSize::MaxEvents(NonZeroUsize::new(100).unwrap()),
482 when_full: WhenFull::Block,
483 },
484 );
485 }
486
487 #[test]
488 fn parse_memory_with_byte_size_option() {
489 check_single_stage(
490 r"
491 max_size: 4096
492 ",
493 BufferType::Memory {
494 size: MemoryBufferSize::MaxSize(NonZeroUsize::new(4096).unwrap()),
495 when_full: WhenFull::Block,
496 },
497 );
498 }
499
500 #[test]
501 fn parse_multiple_stages() {
502 check_multiple_stages(
503 r"
504 - max_events: 42
505 - max_events: 100
506 when_full: drop_newest
507 ",
508 &[
509 BufferType::Memory {
510 size: MemoryBufferSize::MaxEvents(NonZeroUsize::new(42).unwrap()),
511 when_full: WhenFull::Block,
512 },
513 BufferType::Memory {
514 size: MemoryBufferSize::MaxEvents(NonZeroUsize::new(100).unwrap()),
515 when_full: WhenFull::DropNewest,
516 },
517 ],
518 );
519 }
520
521 #[test]
522 fn ensure_field_defaults_for_all_types() {
523 check_single_stage(
524 r"
525 type: memory
526 ",
527 BufferType::Memory {
528 size: MemoryBufferSize::MaxEvents(NonZeroUsize::new(500).unwrap()),
529 when_full: WhenFull::Block,
530 },
531 );
532
533 check_single_stage(
534 r"
535 type: memory
536 max_events: 100
537 ",
538 BufferType::Memory {
539 size: MemoryBufferSize::MaxEvents(NonZeroUsize::new(100).unwrap()),
540 when_full: WhenFull::Block,
541 },
542 );
543
544 check_single_stage(
545 r"
546 type: memory
547 when_full: drop_newest
548 ",
549 BufferType::Memory {
550 size: MemoryBufferSize::MaxEvents(NonZeroUsize::new(500).unwrap()),
551 when_full: WhenFull::DropNewest,
552 },
553 );
554
555 check_single_stage(
556 r"
557 type: memory
558 when_full: overflow
559 ",
560 BufferType::Memory {
561 size: MemoryBufferSize::MaxEvents(NonZeroUsize::new(500).unwrap()),
562 when_full: WhenFull::Overflow,
563 },
564 );
565
566 check_single_stage(
567 r"
568 type: disk
569 max_size: 1024
570 ",
571 BufferType::DiskV2 {
572 max_size: NonZeroU64::new(1024).unwrap(),
573 when_full: WhenFull::Block,
574 },
575 );
576 }
577}