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

harmos_storage/manager/
arrow.rs
