harmos_storage/manager/
arrow.rs

1/* ***
2 *
3 * Harmos
4 *
5 * SPDX-License-Identifier: Apache-2.0
6 *
7 * ***
8 */
9
10/* === Modules === */
11
12/* === Re-Exports === */
13
14/* === Imports === */
15
16// Standard Library Imports
17
18// Third Party Imports
19
20// Local Imports
21use super::*;
22
23/* === Constants === */
24
25/* === Macros === */
26
27/* === Type Aliases === */
28
29/* === Structs === */
30
31/* === Enums === */
32
33/* === Traits === */
34
35/* === Implementations === */
36
37impl<C> StorageManager<C> {
38    #[tracing::instrument(name = "storage.manager.arrow.lookup_schema", level = "debug", skip_all)]
39    pub async fn lookup_schema(&self, key: &SchemaKey) -> Result<(StorageSchema, RetainedMemory)> {
40        let held = Held::default();
41        let snapshot = self.snapshot(&self.cursor_saved(), &held).await?;
42        crate::query::schema_at(&self.inner, &held, &snapshot, key).await
43    }
44
45    #[tracing::instrument(name = "storage.manager.arrow.admit_arrow_records", level = "debug", skip_all)]
46    pub async fn admit_arrow_records(&self, scope: Option<&Identifier>, batch: &crate::writes::DecodedBatch) -> Result<StorageCursor> {
47        let scratch = self.reserve_arrow_write(batch).await?;
48        let (schema, _schema_memory) = self.lookup_schema(&batch.schema).await?;
49        let mut records = crate::writes::convert::records(scope, batch, &schema.schema)?;
50        for record in &mut records {
51            if record.source.as_ref().is_some_and(|source| source != &batch.destination) {
52                return Err(Error::InvalidInput {
53                    field: "record destination",
54                    reason: "record source differs from authorized destination",
55                });
56            }
57            record.source = Some(batch.destination.clone());
58        }
59        let result = self.admit_records(records, None).await;
60        drop(scratch);
61        result
62    }
63
64    #[tracing::instrument(name = "storage.manager.arrow.admit_arrow_streams", level = "debug", skip_all)]
65    pub async fn admit_arrow_streams(&self, scope: Option<&Identifier>, batch: &crate::writes::DecodedBatch) -> Result<StorageCursor> {
66        let scratch = self.reserve_arrow_write(batch).await?;
67        let (schema, _schema_memory) = self.lookup_schema(&batch.schema).await?;
68        let rows = crate::writes::convert::streams(scope, batch, &schema.schema)?;
69        let result = self.admit_stream(rows, None).await;
70        drop(scratch);
71        result
72    }
73
74    #[tracing::instrument(name = "storage.manager.arrow.admit_arrow_objects", level = "debug", skip_all)]
75    pub async fn admit_arrow_objects(&self, _scope: Option<&Identifier>, batch: &crate::writes::DecodedBatch) -> Result<StorageCursor> {
76        use crate::MemorySize;
77        if batch.destination != batch.schema.id {
78            return Err(Error::InvalidInput { field: "object destination", reason: "destination differs from object schema" });
79        }
80        let scratch = self.reserve_arrow_write(batch).await?;
81        let _writer = tracing::Instrument::instrument(self.inner.object_writes.lock(), tracing::trace_span!("storage.lock.wait")).await;
82        let definition = (self.inner.object_definition)(&batch.schema)?;
83        let graph = definition.graph(&harmos_domain::config::SchemaSettings::default())?;
84        let values = crate::writes::convert::objects(batch, &graph, self.inner.settings.max_batch_bytes())?;
85        let count =
86            std::num::NonZeroUsize::new(values.len()).ok_or(Error::InvalidInput { field: "object batch", reason: "empty atomic batch" })?;
87        let metadata = graph
88            .memory_size()
89            .checked_mul(4)
90            .and_then(|bytes| bytes.checked_add(ENTRY_OVERHEAD * 8))
91            .and_then(|bytes| bytes.checked_mul(values.len()))
92            .and_then(|bytes| values.memory_size().checked_mul(4).and_then(|payload| bytes.checked_add(payload)))
93            .ok_or(Error::ResourceExhausted { resource: "object atomic batch" })?;
94        let allowance = std::num::NonZeroUsize::new(metadata).ok_or(Error::ResourceExhausted { resource: "object atomic batch" })?;
95        if scratch
96            .bytes()
97            .checked_add(allowance.get())
98            .is_none_or(|bytes| bytes > self.inner.settings.request.max_pending_write_bytes.get())
99            || values.memory_size() > allowance.get()
100        {
101            return Err(Error::ResourceExhausted { resource: "object atomic batch" });
102        }
103        self.check_object_preconditions(batch, &values).await?;
104        let receipts = self
105            .admit_object_group_locked(allowance, count, None, |group| {
106                for value in values {
107                    group.write_value(value)?;
108                }
109                Ok(())
110            })
111            .await?;
112        drop(scratch);
113        receipts
114            .last()
115            .map(|receipt| receipt.position.clone())
116            .ok_or(Error::InvalidInput { field: "object batch", reason: "accepted batch has no receipt" })
117    }
118
119    #[tracing::instrument(name = "storage.manager.arrow.check_object_preconditions", level = "debug", skip_all)]
120    pub(super) async fn check_object_preconditions(
121        &self, batch: &crate::writes::DecodedBatch, values: &[crate::ObjectValue],
122    ) -> Result<()> {
123        use arrow_array::Array;
124        let Some(metadata) = batch
125            .data
126            .as_record_batch()
127            .column_by_name("_storage")
128            .and_then(|array| array.as_any().downcast_ref::<arrow_array::StructArray>())
129        else {
130            return Err(Error::InvalidInput { field: "object batch", reason: "missing object metadata" });
131        };
132        let Some(expected) = metadata.column_by_name("expected_revision") else {
133            return Ok(());
134        };
135        let expected = expected
136            .as_any()
137            .downcast_ref::<arrow_array::UInt64Array>()
138            .ok_or(Error::InvalidInput { field: "object revision", reason: "expected_revision must be nullable UInt64" })?;
139        let held = Held::default();
140        let snapshot = self.snapshot(&self.cursor_saved(), &held).await?;
141        let page_bytes = self.inner.settings.scan_bytes();
142        let store = metadata.column_by_name("store").and_then(|array| array.as_any().downcast_ref::<arrow_array::StringArray>());
143        let epoch = metadata.column_by_name("epoch").and_then(|array| array.as_any().downcast_ref::<arrow_array::StringArray>());
144        for (row, value) in values.iter().enumerate() {
145            if expected.is_null(row) {
146                continue;
147            }
148            if expected.value(row) != 0
149                && (store.is_none_or(|store| store.is_null(row) || store.value(row) != self.inner.primary.history().store.as_str())
150                    || epoch.is_none_or(|epoch| epoch.is_null(row) || epoch.value(row) != self.inner.primary.history().epoch.as_str()))
151            {
152                return Err(Error::InvalidInput {
153                    field: "object revision",
154                    reason: "expected revision belongs to a different storage history",
155                });
156            }
157            let selection = crate::query::ObjectSelection {
158                schema: Some(value.schema.clone()),
159                key: Some(value.item.clone()),
160                versions: crate::query::ObjectVersions::Latest,
161            };
162            let page = crate::ScanPage { after: 0, rows: std::num::NonZeroUsize::MIN, bytes: page_bytes };
163            let page_memory = Reservation::wait(&self.inner, Pool::Read, &held, page_bytes.get().saturating_add(ENTRY_OVERHEAD), None)
164                .await
165                .map_err(Refusal::into_error)?
166                .retain();
167            let (latest, _) = snapshot.scan::<crate::ObjectEntry>(&selection, page, page_memory.clone()).await?;
168            let revision = latest.first().map_or(0, |entry| entry.version.revision().get());
169            if revision != expected.value(row) {
170                return Err(Error::InvalidInput { field: "object revision", reason: "atomic batch revision precondition failed" });
171            }
172        }
173        Ok(())
174    }
175
176    #[tracing::instrument(name = "storage.manager.arrow.reserve_arrow_write", level = "debug", skip_all)]
177    pub(super) async fn reserve_arrow_write(&self, batch: &crate::writes::DecodedBatch) -> Result<Reservation> {
178        use crate::MemorySize;
179        if batch.data.as_record_batch().num_rows() == 0 {
180            return Err(Error::InvalidInput { field: "storage write", reason: "empty atomic batch" });
181        }
182        let bytes = batch.memory_size().checked_mul(8).ok_or(Error::ResourceExhausted { resource: "storage write conversion" })?;
183        // The original arrays, conversion scratch and admitted native copy coexist.
184        // Refuse an impossible working set before allocating or waiting on ourselves.
185        if bytes.checked_mul(2).is_none_or(|bytes| bytes > self.inner.settings.request.max_pending_write_bytes.get()) {
186            return Err(Error::ResourceExhausted { resource: "storage write conversion" });
187        }
188        self.reserve_conversion(bytes).await
189    }
190
191    // Conversion Scratch Draws On At Most Half The Write Budget So The Admitted
192    // Native Copy Always Has Coexistence Space.
193    #[tracing::instrument(name = "storage.manager.arrow.reserve_conversion", level = "debug", skip_all)]
194    pub(super) async fn reserve_conversion(&self, bytes: usize) -> Result<Reservation> {
195        Reservation::wait(&self.inner, Pool::Write, &Held::default(), bytes, None).await.map_err(Refusal::into_error)
196    }
197}
198
199/* === Public Functions === */
200
201/* === Private Functions === */
202
203/* === Tests === */
harmos_storage/manager/arrow.rs

harmos_storage/manager/
arrow.rs
