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