harmos_storage/
read_session.rs1use crate::capacity::{Held, Pool, Refusal, Reservation};
18use crate::manager::ENTRY_OVERHEAD;
19use crate::{MemorySize, ReadBatch, Result, StorageManager};
20
21pub struct ReadSession<C = ()> {
36 pub(crate) manager: StorageManager<C>,
37 pub(crate) held: Held,
38 pub(crate) queued: bool,
39}
40
41impl<C> Clone for ReadSession<C> {
48 fn clone(&self) -> Self {
49 Self { manager: self.manager.clone(), held: self.held.clone(), queued: self.queued }
50 }
51}
52
53impl<C> StorageManager<C> {
54 #[tracing::instrument(name = "storage.read_session.new", level = "trace", skip_all)]
55 pub fn read_session(&self) -> ReadSession<C> {
56 ReadSession { manager: self.clone(), held: Held::default(), queued: false }
57 }
58}
59
60impl<C> ReadSession<C> {
61 pub(crate) fn standalone(manager: &StorageManager<C>) -> Self {
62 Self { manager: manager.clone(), held: Held::default(), queued: true }
63 }
64
65 #[tracing::instrument(name = "storage.read_session.try_memory", level = "trace", skip_all, fields(bytes = payload_bytes.get(), outcome = "cancelled"))]
72 pub fn try_memory(&self, payload_bytes: std::num::NonZeroUsize) -> Result<crate::arrow::RetainedMemory> {
73 let result: Result<crate::arrow::RetainedMemory> = (|| {
74 let total = payload_bytes
75 .get()
76 .checked_add(ENTRY_OVERHEAD)
77 .ok_or(crate::Error::ResourceExhausted { resource: "session read memory" })?;
78 Reservation::try_new(&self.manager.inner, Pool::Read, &self.held, total).map(Reservation::retain).map_err(Refusal::into_error)
79 })();
80 tracing::Span::current().record("outcome", if result.is_ok() { "completed" } else { "failed" });
81 result
82 }
83
84 #[tracing::instrument(name = "storage.read_session.read_file", level = "trace", skip_all, fields(outcome = "cancelled"))]
85 pub async fn read_file(&self, request: harmos_domain::operations::storage::ReadFile) -> Result<crate::FileReader> {
86 let result = self.manager.read_file_with_held(request, self.held.clone()).await;
87 tracing::Span::current().record("outcome", if result.is_ok() { "completed" } else { "failed" });
88 result
89 }
90
91 #[tracing::instrument(name = "storage.read_session.collection", level = "trace", skip_all, fields(outcome = "cancelled"))]
97 pub fn collection<U: MemorySize>(&self) -> Result<ReadBatch<Vec<ReadBatch<U>>>> {
98 let result: Result<ReadBatch<Vec<ReadBatch<U>>>> = (|| {
99 let reservation =
100 Reservation::try_new(&self.manager.inner, Pool::Read, &self.held, ENTRY_OVERHEAD).map_err(Refusal::into_error)?;
101 Ok(ReadBatch::new(Vec::new(), reservation))
102 })();
103 tracing::Span::current().record("outcome", if result.is_ok() { "completed" } else { "failed" });
104 result
105 }
106}
107
108