Skip to main content

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 { bytes: Arc::new(Semaphore::new(RETAINED_BYTES)), rows: Arc::new(Semaphore::new(RETAINED_ROWS)), admission: Mutex::new(()) }
71    }
72
73    pub(super) fn retain(
74        &self, definition: &crate::streams::Definition, generation: Option<u64>, row: PublicationRow, wire_sequence: u64, missed: u64,
75        window: Option<(&super::window::Window, u64)>,
76    ) -> Result<ManagedPublication, StreamPublicationError> {
77        let bytes = row
78            .payload
79            .capacity()
80            .checked_add(row.context.as_ref().map_or(0, String::capacity))
81            .and_then(|bytes| bytes.checked_add(OWNER_BYTES))
82            .filter(|bytes| *bytes <= RETAINED_BYTES)
83            .ok_or(StreamPublicationError::Capacity)?;
84        let _admission = self.admission.lock().expect("publication budget poisoned");
85        let missing_bytes = bytes.saturating_sub(self.bytes.available_permits());
86        let missing_rows = usize::from(self.rows.available_permits() == 0);
87        if missing_bytes > 0 || missing_rows > 0 {
88            let (window, taken) = window.ok_or(StreamPublicationError::Capacity)?;
89            window.reclaim(missing_bytes, missing_rows, taken).ok_or(StreamPublicationError::Capacity)?;
90        }
91        let slot = self.rows.clone().try_acquire_owned().map_err(|_| StreamPublicationError::Capacity)?;
92        let bytes = self.bytes.clone().try_acquire_many_owned(bytes as u32).map_err(|_| StreamPublicationError::Capacity)?;
93        Ok(ManagedPublication(Arc::new(Owned {
94            stream: Arc::<str>::from(definition.id()),
95            owner: Arc::from(definition.owner().expect("plugin publication owner")),
96            schema: definition.version(),
97            generation,
98            wire_sequence,
99            logical_sequence: OnceLock::new(),
100            missed,
101            row,
102            _bytes: bytes,
103            _slot: slot,
104        })))
105    }
106}
107
108impl ManagedPublication {
109    pub fn stream(&self) -> &str {
110        &self.0.stream
111    }
112
113    pub fn owner(&self) -> &str {
114        &self.0.owner
115    }
116
117    pub fn schema(&self) -> u32 {
118        self.0.schema
119    }
120
121    pub fn generation(&self) -> Option<u64> {
122        self.0.generation
123    }
124
125    pub fn sequence(&self) -> u64 {
126        *self.0.logical_sequence.get().expect("committed publication")
127    }
128
129    pub(super) fn assign_sequence(&self, sequence: u64) {
130        self.0.logical_sequence.set(sequence).expect("one original publication append");
131    }
132
133    pub(super) fn reclaimable_bytes(&self) -> Option<usize> {
134        (Arc::strong_count(&self.0) == 1).then(|| self.0._bytes.num_permits())
135    }
136
137    pub fn wire_sequence(&self) -> u64 {
138        self.0.wire_sequence
139    }
140
141    pub fn missed(&self) -> u64 {
142        self.0.missed
143    }
144
145    pub fn payload(&self) -> &[u8] {
146        &self.0.row.payload
147    }
148
149    pub fn published_unix_micros(&self) -> Option<i64> {
150        self.0.row.published_unix_micros
151    }
152
153    pub fn context(&self) -> Option<&str> {
154        self.0.row.context.as_deref()
155    }
156
157    pub(super) fn row(&self) -> &PublicationRow {
158        &self.0.row
159    }
160}
161
162/* === Public Functions === */
163
164pub(crate) fn row_value<S: 'static>(row: &dyn Any) -> Option<&S> {
165    row.downcast_ref::<S>().or_else(|| {
166        let managed = row.downcast_ref::<ManagedPublication>()?;
167        (managed.row() as &dyn Any).downcast_ref::<S>()
168    })
169}
170
171/* === Private Functions === */
172
173/* === Tests === */