harmos_storage/manager/
writes.rs1use super::*;
22
23impl<C> StorageManager<C> {
38 pub async fn write_schema(&self, key: &SchemaKey, schema: &crate::schema::graph::Schema) -> Result<StorageCursor> {
39 self.admit_schema(key, schema, None).await
40 }
41
42 pub(crate) async fn admit_schema(
43 &self, key: &SchemaKey, schema: &crate::schema::graph::Schema, mutation: Option<&crate::storage::MutationId>,
44 ) -> Result<StorageCursor> {
45 let scratch = schema
46 .validation_memory_bytes(&harmos_domain::config::SchemaSettings::default())
47 .map_err(|_| Error::InvalidInput { field: "storage schema", reason: "invalid message field budget" })?;
48 let mutation = mutation.map(|id| crate::admissions::schema_mutation(id, key, schema)).transpose()?;
49 if let Some(mutation) = mutation {
52 let state = self.inner.state.lock().unwrap();
53 if state.status.closing {
54 return Err(Error::Closing { domain: "storage" });
55 }
56 if !state.ready || !state.status.running {
57 return Err(Error::NotReady { domain: "storage" });
58 }
59 if let ReceiptState::Complete(value) = state.receipts.lookup(mutation.id, &mutation.fingerprint, None)? {
60 return value.cursor();
61 }
62 }
63 let scratch = self.reserve_scratch(scratch).await?;
64 schema
65 .prepare(&harmos_domain::config::SchemaSettings::default())
66 .map_err(|_| Error::InvalidInput { field: "storage schema", reason: "invalid message fields" })?;
67 drop(scratch);
68 let bytes = StorageSchema::resident_bytes(key, schema, self.inner.primary.history())?;
69 let sequence = self
70 .admit_mutation(Domain::Data, bytes, 1, mutation, false, |sequence| {
71 Ok(Operation::Schema(StorageSchema::new(self.position(sequence), key.clone(), schema.clone())))
72 })
73 .await?;
74 Ok(self.position(sequence))
75 }
76
77 pub fn record_batch_len(&self, records: &[RecordWrite<'_>]) -> Result<usize> {
78 self.batch_prefix(records.iter().map(|record| {
79 let value = StoredRecord {
80 position: self.position(u64::MAX),
81 schema: record.schema.clone(),
82 scope: record.scope.cloned(),
83 item: record.item.clone(),
84 source: record.source.cloned(),
85 source_time: record.source_time,
86 payload: record.payload,
87 };
88 Ok((encoded_size(&value)?, record.resident_bytes(self.inner.primary.history())?))
89 }))
90 }
91
92 pub fn stream_batch_len(&self, rows: &[StreamWrite<'_>]) -> Result<usize> {
93 self.batch_prefix(rows.iter().map(|row| {
94 let value = StoredRow {
95 record: StoredRecord {
96 position: self.position(u64::MAX),
97 schema: row.record.schema.clone(),
98 scope: row.record.scope.cloned(),
99 item: row.record.item.clone(),
100 source: row.record.source.cloned(),
101 source_time: row.record.source_time,
102 payload: row.record.payload,
103 },
104 stream: row.stream.clone(),
105 sequence: row.sequence,
106 };
107 let resident = row
108 .record
109 .resident_bytes(self.inner.primary.history())?
110 .checked_add(std::mem::size_of::<StoredRow>() - std::mem::size_of::<StoredRecord>())
111 .and_then(|bytes| bytes.checked_add(row.stream.as_str().len()))
112 .ok_or(Error::ResourceExhausted { resource: "stream batch memory" })?;
113 Ok((encoded_size(&value)?, resident))
114 }))
115 }
116
117 pub(super) fn batch_prefix(&self, sizes: impl Iterator<Item = Result<(usize, usize)>>) -> Result<usize> {
118 let mut bytes = 0usize;
119 let mut rows = 0usize;
120 let mut resident = 0usize;
121 let memory = self.inner.settings.request.max_pending_write_bytes.get();
122 let maximum = self.inner.settings.max_batch_bytes().get();
123 for size in sizes {
124 let (size, held) = size?;
125 if size > maximum || held > memory {
126 return if rows == 0 { Err(Error::ResourceExhausted { resource: "oversized atomic write row" }) } else { Ok(rows) };
127 }
128 if size > maximum - bytes || held > memory - resident {
129 break;
130 }
131 bytes += size;
132 resident += held;
133 rows += 1;
134 }
135 Ok(rows)
136 }
137
138 pub async fn write_records(&self, records: &[RecordWrite<'_>]) -> Result<StorageCursor> {
139 self.admit_records(records.iter().copied(), None).await
140 }
141
142 pub(crate) async fn admit_records<'a>(
143 &self, records: impl ExactSizeIterator<Item = RecordWrite<'a>> + Clone, mutation: Option<Mutation<'_>>,
144 ) -> Result<StorageCursor> {
145 if records.len() == 0 {
146 return Ok(self.status().accepted);
147 }
148 let bytes = records.clone().try_fold(0usize, |bytes, record| {
149 bytes
150 .checked_add(record.resident_bytes(self.inner.primary.history())?)
151 .ok_or(Error::InvalidInput { field: "storage admission", reason: "size overflow" })
152 })?;
153 let sequence = self
154 .admit_mutation(Domain::Data, bytes, records.len(), mutation, false, |first| {
155 Ok(Operation::Records(
156 records.enumerate().map(|(index, record)| record.stored(self.position(first + index as u64))).collect(),
157 ))
158 })
159 .await?;
160 Ok(self.position(sequence))
161 }
162
163 pub async fn write_stream(&self, rows: &[StreamWrite<'_>]) -> Result<StorageCursor> {
164 self.admit_stream(rows.iter().copied(), None).await
165 }
166
167 pub(crate) async fn admit_stream<'a>(
168 &self, rows: impl ExactSizeIterator<Item = StreamWrite<'a>> + Clone, mutation: Option<Mutation<'_>>,
169 ) -> Result<StorageCursor> {
170 if rows.len() == 0 {
171 return Ok(self.status().accepted);
172 }
173 let bytes = rows.clone().try_fold(0usize, |bytes, row| {
174 bytes
175 .checked_add(row.record.resident_bytes(self.inner.primary.history())?)
176 .and_then(|bytes| bytes.checked_add(std::mem::size_of::<StoredRow>() - std::mem::size_of::<StoredRecord>()))
177 .and_then(|bytes| bytes.checked_add(row.stream.as_str().len()))
178 .ok_or(Error::InvalidInput { field: "storage admission", reason: "size overflow" })
179 })?;
180 let sequence = self
181 .admit_mutation(Domain::Data, bytes, rows.len(), mutation, true, |first| {
182 Ok(Operation::Stream(
183 rows.enumerate()
184 .map(|(index, row)| StoredRow {
185 record: row.record.stored(self.position(first + index as u64)),
186 stream: row.stream.clone(),
187 sequence: row.sequence,
188 })
189 .collect(),
190 ))
191 })
192 .await?;
193 Ok(self.position(sequence))
194 }
195}
196
197