Skip to main content

cloud_sdk/client/
workspace.rs

1use core::fmt;
2use core::sync::atomic::{AtomicUsize, Ordering};
3
4use cloud_sdk_sanitization::{SecretBuffer, sanitize_bytes};
5
6/// Maximum number of atomic workspace leases in one pool.
7pub const MAX_CLIENT_WORKSPACE_LEASES: usize = usize::BITS as usize;
8
9/// Invalid compile-time workspace-pool capacity.
10#[derive(Clone, Copy, Debug, Eq, PartialEq)]
11pub enum WorkspacePoolError {
12    /// A pool must admit at least one workspace.
13    ZeroCapacity,
14    /// The requested capacity exceeds the atomic lease bitmap.
15    CapacityTooLarge,
16}
17
18impl_static_error!(WorkspacePoolError,
19    Self::ZeroCapacity => "client workspace pool capacity is zero",
20    Self::CapacityTooLarge => "client workspace pool capacity exceeds the atomic bound",
21);
22
23/// Failure to admit one caller-owned workspace immediately.
24#[derive(Clone, Copy, Debug, Eq, PartialEq)]
25pub enum WorkspaceAcquireError {
26    /// Every bounded pool slot is currently leased.
27    Exhausted,
28}
29
30impl_static_error!(WorkspaceAcquireError,
31    Self::Exhausted => "client workspace pool is exhausted",
32);
33
34/// Four independent caller-owned buffers for one complete client request.
35///
36/// The constructor borrows every region mutably, so aliasing request and
37/// response storage is rejected by safe Rust:
38///
39/// ```compile_fail
40/// use cloud_sdk::client::ClientWorkspace;
41///
42/// let mut shared = [0_u8; 128];
43/// let mut body = [0_u8; 128];
44/// let mut headers = [0_u8; 128];
45/// let _ = ClientWorkspace::new(&mut shared, &mut body, &mut shared, &mut headers);
46/// ```
47pub struct ClientWorkspace<'storage> {
48    target: SecretBuffer<'storage>,
49    request_body: SecretBuffer<'storage>,
50    response_body: SecretBuffer<'storage>,
51    response_headers: SecretBuffer<'storage>,
52}
53
54impl<'storage> ClientWorkspace<'storage> {
55    /// Admits and immediately clears four independent storage regions.
56    #[must_use]
57    pub fn new(
58        target: &'storage mut [u8],
59        request_body: &'storage mut [u8],
60        response_body: &'storage mut [u8],
61        response_headers: &'storage mut [u8],
62    ) -> Self {
63        sanitize_bytes(target);
64        sanitize_bytes(request_body);
65        sanitize_bytes(response_body);
66        sanitize_bytes(response_headers);
67        Self {
68            target: SecretBuffer::new(target),
69            request_body: SecretBuffer::new(request_body),
70            response_body: SecretBuffer::new(response_body),
71            response_headers: SecretBuffer::new(response_headers),
72        }
73    }
74
75    /// Returns capacities without exposing stored bytes.
76    #[must_use]
77    pub fn capacities(&self) -> (usize, usize, usize, usize) {
78        (
79            self.target.as_slice().len(),
80            self.request_body.as_slice().len(),
81            self.response_body.as_slice().len(),
82            self.response_headers.as_slice().len(),
83        )
84    }
85
86    pub(crate) fn parts_mut(&mut self) -> ClientWorkspaceParts<'_> {
87        ClientWorkspaceParts {
88            target: self.target.as_mut_slice(),
89            request_body: self.request_body.as_mut_slice(),
90            response_body: self.response_body.as_mut_slice(),
91            response_headers: self.response_headers.as_mut_slice(),
92        }
93    }
94}
95
96impl fmt::Debug for ClientWorkspace<'_> {
97    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
98        formatter
99            .debug_struct("ClientWorkspace")
100            .field("capacities", &self.capacities())
101            .finish_non_exhaustive()
102    }
103}
104
105pub(crate) struct ClientWorkspaceParts<'storage> {
106    pub(crate) target: &'storage mut [u8],
107    pub(crate) request_body: &'storage mut [u8],
108    pub(crate) response_body: &'storage mut [u8],
109    pub(crate) response_headers: &'storage mut [u8],
110}
111
112impl ClientWorkspaceParts<'_> {
113    pub(crate) fn clear(&mut self) {
114        sanitize_bytes(self.target);
115        sanitize_bytes(self.request_body);
116        sanitize_bytes(self.response_body);
117        sanitize_bytes(self.response_headers);
118    }
119}
120
121/// Fixed-capacity atomic admission for caller-supplied workspaces.
122///
123/// The pool stores no buffers and has no wait queue. Exhaustion is immediate.
124pub struct ClientWorkspacePool<const N: usize> {
125    leased: AtomicUsize,
126}
127
128impl<const N: usize> ClientWorkspacePool<N> {
129    /// Creates a pool when `N` fits the platform atomic bitmap.
130    pub const fn new() -> Result<Self, WorkspacePoolError> {
131        if N == 0 {
132            return Err(WorkspacePoolError::ZeroCapacity);
133        }
134        if N > MAX_CLIENT_WORKSPACE_LEASES {
135            return Err(WorkspacePoolError::CapacityTooLarge);
136        }
137        Ok(Self {
138            leased: AtomicUsize::new(0),
139        })
140    }
141
142    /// Admits one workspace immediately or returns exhaustion without queuing.
143    pub fn try_acquire<'pool, 'storage>(
144        &'pool self,
145        workspace: ClientWorkspace<'storage>,
146    ) -> Result<ClientWorkspaceLease<'pool, 'storage, N>, WorkspaceAcquireError> {
147        let valid = valid_mask::<N>();
148        let mut observed = self.leased.load(Ordering::Acquire);
149        loop {
150            let available = !observed & valid;
151            if available == 0 {
152                return Err(WorkspaceAcquireError::Exhausted);
153            }
154            let index = available.trailing_zeros() as usize;
155            let bit = 1_usize << index;
156            match self.leased.compare_exchange_weak(
157                observed,
158                observed | bit,
159                Ordering::AcqRel,
160                Ordering::Acquire,
161            ) {
162                Ok(_) => {
163                    return Ok(ClientWorkspaceLease {
164                        workspace,
165                        _slot: LeaseSlot { pool: self, bit },
166                    });
167                }
168                Err(current) => observed = current,
169            }
170        }
171    }
172
173    /// Returns the number of currently admitted workspaces.
174    #[must_use]
175    pub fn active_leases(&self) -> usize {
176        (self.leased.load(Ordering::Acquire) & valid_mask::<N>()).count_ones() as usize
177    }
178}
179
180impl<const N: usize> fmt::Debug for ClientWorkspacePool<N> {
181    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
182        formatter
183            .debug_struct("ClientWorkspacePool")
184            .field("capacity", &N)
185            .field("active_leases", &self.active_leases())
186            .finish()
187    }
188}
189
190/// One admitted workspace whose buffers remain uniquely borrowed until drop.
191///
192/// The lease owns all mutable borrows when moved into an async execution
193/// future, so caller code cannot touch any buffer while that request is live.
194///
195/// ```compile_fail
196/// use cloud_sdk::client::{ClientWorkspace, ClientWorkspaceLease, ClientWorkspacePool};
197///
198/// async fn hold<const N: usize>(_: ClientWorkspaceLease<'_, '_, N>) {}
199///
200/// let pool = ClientWorkspacePool::<1>::new().unwrap();
201/// let mut target = [0_u8; 16];
202/// let mut request = [0_u8; 16];
203/// let mut response = [0_u8; 16];
204/// let mut headers = [0_u8; 16];
205/// let workspace = ClientWorkspace::new(
206///     &mut target,
207///     &mut request,
208///     &mut response,
209///     &mut headers,
210/// );
211/// let lease = pool.try_acquire(workspace).unwrap();
212/// let future = hold(lease);
213/// response.fill(0xa5);
214/// drop(future);
215/// ```
216pub struct ClientWorkspaceLease<'pool, 'storage, const N: usize> {
217    workspace: ClientWorkspace<'storage>,
218    _slot: LeaseSlot<'pool, N>,
219}
220
221impl<const N: usize> ClientWorkspaceLease<'_, '_, N> {
222    /// Returns all workspace capacities without exposing stored bytes.
223    #[must_use]
224    pub fn capacities(&self) -> (usize, usize, usize, usize) {
225        self.workspace.capacities()
226    }
227
228    pub(crate) fn parts_mut(&mut self) -> ClientWorkspaceParts<'_> {
229        self.workspace.parts_mut()
230    }
231}
232
233impl<const N: usize> fmt::Debug for ClientWorkspaceLease<'_, '_, N> {
234    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
235        formatter
236            .debug_struct("ClientWorkspaceLease")
237            .field("capacities", &self.capacities())
238            .field("slot", &"[leased]")
239            .finish()
240    }
241}
242
243struct LeaseSlot<'pool, const N: usize> {
244    pool: &'pool ClientWorkspacePool<N>,
245    bit: usize,
246}
247
248impl<const N: usize> Drop for LeaseSlot<'_, N> {
249    fn drop(&mut self) {
250        self.pool.leased.fetch_and(!self.bit, Ordering::Release);
251    }
252}
253
254fn valid_mask<const N: usize>() -> usize {
255    let shift = match u32::try_from(N) {
256        Ok(value) => value,
257        Err(_) => return usize::MAX,
258    };
259    match 1_usize.checked_shl(shift) {
260        Some(upper_bit) => upper_bit.saturating_sub(1),
261        None => usize::MAX,
262    }
263}