harmos_runtime/streams/
publication.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
17use std::any::Any;
18use std::sync::{Arc, Mutex, OnceLock};
19// Third Party Imports
20use tokio::sync::{OwnedSemaphorePermit, Semaphore};
21// Local Imports
22use super::{PublicationRow, StreamPublicationError};
23
24/* === Constants === */
25
26const RETAINED_BYTES: usize = 64 * 1024 * 1024;
27const RETAINED_ROWS: usize = 65_536;
28
29// Includes requested row/receipt headers, queue cell and Arc counters.
30const OWNER_BYTES: usize = 1024;
31
32/* === Macros === */
33
34/* === Type Aliases === */
35
36/* === Structs === */
37
38#[derive(Clone, Debug)]
39pub struct ManagedPublication(Arc<Owned>);
40
41#[derive(Debug)]
42struct Owned {
43    stream: Arc<str>,
44    owner: Arc<str>,
45    schema: u32,
46    generation: Option<u64>,
47    wire_sequence: u64,
48    logical_sequence: OnceLock<u64>,
49    missed: u64,
50    row: PublicationRow,
51    _bytes: OwnedSemaphorePermit,
52    _slot: OwnedSemaphorePermit,
53}
54
55#[derive(Debug)]
56pub(super) struct Pool {
57    bytes: Arc<Semaphore>,
58    rows: Arc<Semaphore>,
59    admission: Mutex<()>,
60}
61
62/* === Enums === */
63
64/* === Traits === */
65
66/* === Implementations === */
67
68impl Pool {
69    pub(super) fn new() -> Self {
70        Self {
71            bytes: Arc::new(Semaphore::new(RETAINED_BYTES)),
72            rows: Arc::new(Semaphore::new(RETAINED_ROWS)),
73            admission: Mutex::new(()),
74        }
75    }
76
77    pub(super) fn retain(
78        &self, definition: &crate::streams::Definition, generation: Option<u64>, row: PublicationRow, wire_sequence: u64, missed: u64,
79        window: Option<(&super::window::Window, u64)>,
80    ) -> Result<ManagedPublication, StreamPublicationError> {
81        let bytes = row
82            .payload
83            .capacity()
84            .checked_add(row.context.as_ref().map_or(0, String::capacity))
85            .and_then(|bytes| bytes.checked_add(OWNER_BYTES))
86            .filter(|bytes| *bytes <= RETAINED_BYTES)
87            .ok_or(StreamPublicationError::Capacity)?;
88        let _admission = self.admission.lock().expect("publication budget poisoned");
89        let missing_bytes = bytes.saturating_sub(self.bytes.available_permits());
90        let missing_rows = usize::from(self.rows.available_permits() == 0);
91        if missing_bytes > 0 || missing_rows > 0 {
92            let (window, taken) = window.ok_or(StreamPublicationError::Capacity)?;
93            window.reclaim(missing_bytes, missing_rows, taken).ok_or(StreamPublicationError::Capacity)?;
94        }
95        let slot = self.rows.clone().try_acquire_owned().map_err(|_| StreamPublicationError::Capacity)?;
96        let bytes = self.bytes.clone().try_acquire_many_owned(bytes as u32).map_err(|_| StreamPublicationError::Capacity)?;
97        Ok(ManagedPublication(Arc::new(Owned {
98            stream: Arc::<str>::from(definition.id()),
99            owner: Arc::from(definition.owner().expect("plugin publication owner")),
100            schema: definition.version(),
101            generation,
102            wire_sequence,
103            logical_sequence: OnceLock::new(),
104            missed,
105            row,
106            _bytes: bytes,
107            _slot: slot,
108        })))
109    }
110}
111
112impl ManagedPublication {
113    pub fn stream(&self) -> &str {
114        &self.0.stream
115    }
116
117    pub fn owner(&self) -> &str {
118        &self.0.owner
119    }
120
121    pub fn schema(&self) -> u32 {
122        self.0.schema
123    }
124
125    pub fn generation(&self) -> Option<u64> {
126        self.0.generation
127    }
128
129    pub fn sequence(&self) -> u64 {
130        *self.0.logical_sequence.get().expect("committed publication")
131    }
132
133    pub(super) fn assign_sequence(&self, sequence: u64) {
134        self.0.logical_sequence.set(sequence).expect("one original publication append");
135    }
136
137    pub(super) fn reclaimable_bytes(&self) -> Option<usize> {
138        (Arc::strong_count(&self.0) == 1).then(|| self.0._bytes.num_permits())
139    }
140
141    pub fn wire_sequence(&self) -> u64 {
142        self.0.wire_sequence
143    }
144
145    pub fn missed(&self) -> u64 {
146        self.0.missed
147    }
148
149    pub fn payload(&self) -> &[u8] {
150        &self.0.row.payload
151    }
152
153    pub fn published_unix_micros(&self) -> Option<i64> {
154        self.0.row.published_unix_micros
155    }
156
157    pub fn context(&self) -> Option<&str> {
158        self.0.row.context.as_deref()
159    }
160
161    pub(super) fn row(&self) -> &PublicationRow {
162        &self.0.row
163    }
164}
165
166/* === Public Functions === */
167
168pub(crate) fn row_value<S: 'static>(row: &dyn Any) -> Option<&S> {
169    row.downcast_ref::<S>().or_else(|| {
170        let managed = row.downcast_ref::<ManagedPublication>()?;
171        (managed.row() as &dyn Any).downcast_ref::<S>()
172    })
173}
174
175/* === Private Functions === */
176
177/* === Tests === */
harmos_runtime/streams/publication.rs

harmos_runtime/streams/
publication.rs
