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