Skip to main content

momento_functions_bytes/
data.rs

1/// A buffer of bytes, which may be inline or on the host.
2///
3/// Some bulk data processing Functions may choose to pass data straight from a
4/// request or response body to another request or response body. To improve
5/// performance on these kinds of Functions, `Data` avoids copying the buffer
6/// into your function's memory.
7#[derive(Debug)]
8pub struct Data {
9    data: Location,
10}
11
12impl Data {
13    /// Turn the Data into a plain `Vec<u8>`.
14    ///
15    /// If the data is buffered on the host, this will read the buffer completely
16    /// into your function's memory. You should use this with caution!
17    ///
18    /// For small buffers, this is inconsequential. For larger buffers, this may
19    /// cause your function to run out of memory or to run slowly.
20    pub fn into_bytes(self) -> Vec<u8> {
21        self.data.into_bytes()
22    }
23}
24
25impl From<crate::wit::momento::bytes::bytes::Data> for Data {
26    fn from(value: crate::wit::momento::bytes::bytes::Data) -> Self {
27        match value {
28            crate::wit::momento::bytes::bytes::Data::Value(buffer) => Self {
29                data: Location::Inline { buffer },
30            },
31            crate::wit::momento::bytes::bytes::Data::Buffer(resource) => Self {
32                data: Location::OnHost { resource },
33            },
34        }
35    }
36}
37
38impl From<Vec<u8>> for Data {
39    fn from(value: Vec<u8>) -> Self {
40        Self {
41            data: Location::Inline { buffer: value },
42        }
43    }
44}
45
46impl From<&[u8]> for Data {
47    fn from(value: &[u8]) -> Self {
48        Self {
49            data: Location::Inline {
50                buffer: value.to_vec(),
51            },
52        }
53    }
54}
55
56impl From<String> for Data {
57    fn from(value: String) -> Self {
58        Self {
59            data: Location::Inline {
60                buffer: value.into_bytes(),
61            },
62        }
63    }
64}
65
66impl From<&str> for Data {
67    fn from(value: &str) -> Self {
68        Self {
69            data: Location::Inline {
70                buffer: value.as_bytes().to_vec(),
71            },
72        }
73    }
74}
75
76enum Location {
77    Inline {
78        buffer: Vec<u8>,
79    },
80    OnHost {
81        resource: crate::wit::momento::bytes::bytes::Buffer,
82    },
83}
84impl Location {
85    fn into_bytes(self) -> Vec<u8> {
86        match self {
87            Location::Inline { buffer } => buffer,
88            Location::OnHost { resource } => {
89                let mut buffer = Vec::with_capacity(resource.remaining() as usize);
90                while let Some(chunk) = resource.read(resource.remaining().max(16384)) {
91                    buffer.extend(chunk);
92                }
93                buffer
94            }
95        }
96    }
97}
98
99impl std::fmt::Debug for Location {
100    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
101        match self {
102            Location::Inline { buffer } => f
103                .debug_struct("Inline")
104                .field("length", &buffer.len())
105                .finish(),
106            Location::OnHost { resource } => f
107                .debug_struct("OnHost")
108                .field("remaining", &resource.remaining())
109                .finish(),
110        }
111    }
112}
113
114impl From<Data> for crate::wit::momento::bytes::bytes::Data {
115    fn from(data: Data) -> Self {
116        match data.data {
117            Location::Inline { buffer } => Self::Value(buffer),
118            Location::OnHost { resource } => Self::Buffer(resource),
119        }
120    }
121}