Skip to main content

ursula_runtime/
request.rs

1use std::sync::Arc;
2
3use bytes::Bytes;
4use serde::Deserialize;
5use serde::Serialize;
6use ursula_shard::BucketStreamId;
7use ursula_shard::ShardPlacement;
8use ursula_stream::ColdChunkRef;
9use ursula_stream::ExternalPayloadRef;
10use ursula_stream::ProducerRequest;
11use ursula_stream::StreamAttrs;
12use ursula_stream::StreamIntegritySnapshot;
13use ursula_stream::StreamReadPlan;
14use ursula_stream::StreamReadSegment;
15use ursula_stream::StreamRecordRange;
16
17use crate::cold_index::ColdIndexPageCache;
18use crate::cold_index::ColdStoreColdIndexPageStore;
19use crate::cold_store::ColdStoreHandle;
20use crate::cold_store::DEFAULT_CONTENT_TYPE;
21use crate::engine::GroupEngineError;
22use crate::engine::in_memory::InMemoryGroupEngine;
23use crate::error::RuntimeError;
24
25#[derive(Debug, Clone, PartialEq, Eq)]
26pub struct CreateStreamRequest {
27    pub stream_id: BucketStreamId,
28    pub content_type: String,
29    pub content_type_explicit: bool,
30    pub initial_payload: Bytes,
31    pub close_after: bool,
32    pub stream_seq: Option<String>,
33    pub producer: Option<ProducerRequest>,
34    pub stream_ttl_seconds: Option<u64>,
35    pub stream_expires_at_ms: Option<u64>,
36    pub attrs: Option<StreamAttrs>,
37    pub now_ms: u64,
38}
39
40#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
41pub struct CreateStreamExternalRequest {
42    pub stream_id: BucketStreamId,
43    pub content_type: String,
44    pub initial_payload: ExternalPayloadRef,
45    #[serde(default)]
46    pub record_ends: Vec<u64>,
47    pub close_after: bool,
48    pub stream_seq: Option<String>,
49    pub producer: Option<ProducerRequest>,
50    pub stream_ttl_seconds: Option<u64>,
51    pub stream_expires_at_ms: Option<u64>,
52    pub attrs: Option<StreamAttrs>,
53    pub now_ms: u64,
54}
55
56impl CreateStreamExternalRequest {
57    pub fn from_create_request(
58        request: CreateStreamRequest,
59        initial_payload: ExternalPayloadRef,
60        record_ends: Vec<u64>,
61    ) -> Self {
62        Self {
63            stream_id: request.stream_id,
64            content_type: request.content_type,
65            initial_payload,
66            record_ends,
67            close_after: request.close_after,
68            stream_seq: request.stream_seq,
69            producer: request.producer,
70            stream_ttl_seconds: request.stream_ttl_seconds,
71            stream_expires_at_ms: request.stream_expires_at_ms,
72            attrs: request.attrs,
73            now_ms: request.now_ms,
74        }
75    }
76}
77
78impl CreateStreamRequest {
79    pub fn canonical_record_ends(&self) -> Vec<u64> {
80        ursula_stream::canonical_json_record_ends(&self.content_type, &self.initial_payload)
81            .unwrap_or_default()
82    }
83
84    pub fn new(stream_id: BucketStreamId, content_type: impl Into<String>) -> Self {
85        Self {
86            stream_id,
87            content_type: content_type.into(),
88            content_type_explicit: true,
89            initial_payload: Bytes::new(),
90            close_after: false,
91            stream_seq: None,
92            producer: None,
93            stream_ttl_seconds: None,
94            stream_expires_at_ms: None,
95            attrs: None,
96            now_ms: 0,
97        }
98    }
99}
100
101#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
102pub struct CreateStreamResponse {
103    pub placement: ShardPlacement,
104    pub next_offset: u64,
105    pub closed: bool,
106    pub already_exists: bool,
107    pub group_commit_index: u64,
108    pub record_range: Option<StreamRecordRange>,
109}
110
111#[derive(Debug, Clone, PartialEq, Eq)]
112pub struct HeadStreamRequest {
113    pub stream_id: BucketStreamId,
114    pub now_ms: u64,
115}
116
117#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
118pub struct HeadStreamResponse {
119    pub placement: ShardPlacement,
120    pub content_type: String,
121    pub tail_offset: u64,
122    pub cold_hot_start_offset: u64,
123    pub closed: bool,
124    pub stream_ttl_seconds: Option<u64>,
125    pub stream_expires_at_ms: Option<u64>,
126    pub snapshot_offset: Option<u64>,
127    pub integrity: StreamIntegritySnapshot,
128    pub record_range: Option<StreamRecordRange>,
129}
130
131#[derive(Debug, Clone, PartialEq, Eq)]
132pub struct GetStreamAttrsRequest {
133    pub stream_id: BucketStreamId,
134    pub now_ms: u64,
135}
136
137#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
138pub struct GetStreamAttrsResponse {
139    pub placement: ShardPlacement,
140    pub attrs: Option<StreamAttrs>,
141}
142
143#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
144pub struct UpdateStreamAttrsRequest {
145    pub stream_id: BucketStreamId,
146    pub attrs: Option<StreamAttrs>,
147    pub now_ms: u64,
148}
149
150#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
151pub struct UpdateStreamAttrsResponse {
152    pub placement: ShardPlacement,
153    pub changed: bool,
154    pub group_commit_index: u64,
155}
156
157#[derive(Debug, Clone, PartialEq, Eq)]
158pub struct ReadStreamRequest {
159    pub stream_id: BucketStreamId,
160    pub offset: u64,
161    pub max_len: usize,
162    pub now_ms: u64,
163    pub record: Option<u64>,
164    pub max_records: Option<u64>,
165}
166
167#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
168pub struct ReadStreamResponse {
169    pub placement: ShardPlacement,
170    pub offset: u64,
171    pub next_offset: u64,
172    pub content_type: String,
173    #[serde(with = "serde_bytes")]
174    pub payload: Vec<u8>,
175    pub up_to_date: bool,
176    pub closed: bool,
177    pub retained_record_range: Option<StreamRecordRange>,
178    pub record_range: Option<StreamRecordRange>,
179}
180
181pub enum GroupReadStreamBody {
182    Materialized(Vec<u8>),
183    Planned {
184        stream_id: BucketStreamId,
185        plan: StreamReadPlan,
186        cold_store: Option<ColdStoreHandle>,
187        cold_index_cache: Option<Arc<ColdIndexPageCache<ColdStoreColdIndexPageStore>>>,
188    },
189    #[cfg(test)]
190    Blocking {
191        entered: Arc<crate::rt::sync::Notify>,
192        materialized: Arc<crate::rt::sync::Notify>,
193        release: Arc<crate::rt::sync::Notify>,
194        payload: Vec<u8>,
195    },
196}
197
198pub struct GroupReadStreamParts {
199    pub placement: ShardPlacement,
200    pub offset: u64,
201    pub next_offset: u64,
202    pub content_type: String,
203    pub up_to_date: bool,
204    pub closed: bool,
205    pub retained_record_range: Option<StreamRecordRange>,
206    pub record_range: Option<StreamRecordRange>,
207    pub body: GroupReadStreamBody,
208}
209
210impl GroupReadStreamParts {
211    pub fn from_response(response: ReadStreamResponse) -> Self {
212        Self {
213            placement: response.placement,
214            offset: response.offset,
215            next_offset: response.next_offset,
216            content_type: response.content_type,
217            up_to_date: response.up_to_date,
218            closed: response.closed,
219            retained_record_range: response.retained_record_range,
220            record_range: response.record_range,
221            body: GroupReadStreamBody::Materialized(response.payload),
222        }
223    }
224
225    pub fn from_plan(
226        placement: ShardPlacement,
227        stream_id: BucketStreamId,
228        plan: StreamReadPlan,
229        cold_store: Option<ColdStoreHandle>,
230        cold_index_cache: Option<Arc<ColdIndexPageCache<ColdStoreColdIndexPageStore>>>,
231    ) -> Self {
232        Self {
233            placement,
234            offset: plan.offset,
235            next_offset: plan.next_offset,
236            content_type: plan.content_type.clone(),
237            up_to_date: plan.up_to_date,
238            closed: plan.closed,
239            retained_record_range: plan.retained_record_range,
240            record_range: plan.record_range,
241            body: GroupReadStreamBody::Planned {
242                stream_id,
243                plan,
244                cold_store,
245                cold_index_cache,
246            },
247        }
248    }
249
250    pub async fn into_response(self) -> Result<ReadStreamResponse, GroupEngineError> {
251        let payload = match &self.body {
252            GroupReadStreamBody::Materialized(payload) => payload.clone(),
253            GroupReadStreamBody::Planned {
254                stream_id,
255                plan,
256                cold_store,
257                cold_index_cache,
258            } => {
259                InMemoryGroupEngine::read_payload_from_plan(
260                    cold_store.as_ref(),
261                    cold_index_cache.as_ref(),
262                    stream_id,
263                    plan,
264                )
265                .await?
266            }
267            #[cfg(test)]
268            GroupReadStreamBody::Blocking {
269                entered,
270                materialized,
271                release,
272                payload,
273            } => {
274                entered.notify_one();
275                materialized.notify_one();
276                release.notified().await;
277                payload.clone()
278            }
279        };
280        Ok(ReadStreamResponse {
281            placement: self.placement,
282            offset: self.offset,
283            next_offset: self.next_offset,
284            content_type: self.content_type,
285            payload,
286            up_to_date: self.up_to_date,
287            closed: self.closed,
288            retained_record_range: self.retained_record_range,
289            record_range: self.record_range,
290        })
291    }
292
293    pub fn payload_is_empty(&self) -> bool {
294        match &self.body {
295            GroupReadStreamBody::Materialized(payload) => payload.is_empty(),
296            GroupReadStreamBody::Planned { plan, .. } => {
297                plan.segments.iter().all(|segment| match segment {
298                    StreamReadSegment::Hot(payload) => payload.is_empty(),
299                    StreamReadSegment::ColdIndex(segment) => segment.len == 0,
300                    StreamReadSegment::Object(segment) => segment.len == 0,
301                })
302            }
303            #[cfg(test)]
304            GroupReadStreamBody::Blocking { payload, .. } => payload.is_empty(),
305        }
306    }
307}
308
309#[derive(Debug, Clone, PartialEq, Eq)]
310pub struct PublishSnapshotRequest {
311    pub stream_id: BucketStreamId,
312    pub snapshot_offset: u64,
313    pub content_type: String,
314    pub payload: Bytes,
315    pub now_ms: u64,
316}
317
318#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
319pub struct PublishSnapshotResponse {
320    pub placement: ShardPlacement,
321    pub snapshot_offset: u64,
322    pub group_commit_index: u64,
323    pub record_range: Option<StreamRecordRange>,
324}
325
326#[derive(Debug, Clone, PartialEq, Eq)]
327pub struct ReadSnapshotRequest {
328    pub stream_id: BucketStreamId,
329    pub snapshot_offset: Option<u64>,
330    pub now_ms: u64,
331}
332
333#[derive(Debug, Clone, PartialEq, Eq)]
334pub struct ReadSnapshotResponse {
335    pub placement: ShardPlacement,
336    pub snapshot_offset: u64,
337    pub next_offset: u64,
338    pub content_type: String,
339    pub payload: Vec<u8>,
340    pub up_to_date: bool,
341    pub record_range: Option<StreamRecordRange>,
342}
343
344#[derive(Debug, Clone, PartialEq, Eq)]
345pub struct DeleteSnapshotRequest {
346    pub stream_id: BucketStreamId,
347    pub snapshot_offset: u64,
348    pub now_ms: u64,
349}
350
351#[derive(Debug, Clone, PartialEq, Eq)]
352pub struct BootstrapStreamRequest {
353    pub stream_id: BucketStreamId,
354    pub now_ms: u64,
355}
356
357#[derive(Debug, Clone, PartialEq, Eq)]
358pub struct BootstrapUpdate {
359    pub start_offset: u64,
360    pub next_offset: u64,
361    pub content_type: String,
362    pub payload: Vec<u8>,
363}
364
365#[derive(Debug, Clone, PartialEq, Eq)]
366pub struct BootstrapStreamResponse {
367    pub placement: ShardPlacement,
368    pub snapshot_offset: Option<u64>,
369    pub snapshot_content_type: String,
370    pub snapshot_payload: Vec<u8>,
371    pub updates: Vec<BootstrapUpdate>,
372    pub next_offset: u64,
373    pub up_to_date: bool,
374    pub closed: bool,
375    pub record_range: Option<StreamRecordRange>,
376}
377
378#[derive(Debug, Clone, PartialEq, Eq)]
379pub struct CloseStreamRequest {
380    pub stream_id: BucketStreamId,
381    pub stream_seq: Option<String>,
382    pub producer: Option<ProducerRequest>,
383    pub now_ms: u64,
384}
385
386#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
387pub struct CloseStreamResponse {
388    pub placement: ShardPlacement,
389    pub next_offset: u64,
390    pub group_commit_index: u64,
391    pub deduplicated: bool,
392    pub record_range: Option<StreamRecordRange>,
393}
394
395#[derive(Debug, Clone, PartialEq, Eq)]
396pub struct DeleteStreamRequest {
397    pub stream_id: BucketStreamId,
398}
399
400#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
401pub struct DeleteStreamResponse {
402    pub placement: ShardPlacement,
403    pub group_commit_index: u64,
404}
405
406#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
407pub struct AckColdGcResponse {
408    pub placement: ShardPlacement,
409    pub removed: u64,
410    pub group_commit_index: u64,
411}
412
413#[derive(Debug, Clone, PartialEq, Eq)]
414pub struct FlushColdRequest {
415    pub stream_id: BucketStreamId,
416    pub chunk: ColdChunkRef,
417}
418
419#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
420pub struct FlushColdResponse {
421    pub placement: ShardPlacement,
422    pub hot_start_offset: u64,
423    pub group_commit_index: u64,
424}
425
426#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
427pub struct TouchStreamAccessResponse {
428    pub placement: ShardPlacement,
429    pub changed: bool,
430    pub expired: bool,
431    pub group_commit_index: u64,
432}
433
434#[derive(Debug, Clone, PartialEq, Eq)]
435pub struct PlanColdFlushRequest {
436    pub stream_id: BucketStreamId,
437    pub min_hot_bytes: usize,
438    pub max_flush_bytes: usize,
439}
440
441#[derive(Debug, Clone, PartialEq, Eq)]
442pub struct PlanGroupColdFlushRequest {
443    pub min_hot_bytes: usize,
444    pub max_flush_bytes: usize,
445}
446
447#[derive(Debug, Clone, PartialEq, Eq)]
448pub struct ColdHotBacklog {
449    pub stream_id: BucketStreamId,
450    pub stream_hot_bytes: u64,
451    pub group_hot_bytes: u64,
452}
453
454#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
455pub struct ColdWriteAdmission {
456    pub max_hot_bytes_per_group: Option<u64>,
457}
458
459impl ColdWriteAdmission {
460    pub(crate) fn is_enabled(self) -> bool {
461        self.max_hot_bytes_per_group.is_some()
462    }
463}
464
465#[derive(Debug, Clone, PartialEq, Eq)]
466pub struct AppendRequest {
467    pub stream_id: BucketStreamId,
468    pub content_type: String,
469    pub payload: Bytes,
470    pub close_after: bool,
471    pub stream_seq: Option<String>,
472    pub producer: Option<ProducerRequest>,
473    pub now_ms: u64,
474    pub record_match: Option<u64>,
475}
476
477#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
478pub struct AppendExternalRequest {
479    pub stream_id: BucketStreamId,
480    pub content_type: String,
481    pub payload: ExternalPayloadRef,
482    #[serde(default)]
483    pub record_ends: Vec<u64>,
484    pub close_after: bool,
485    pub stream_seq: Option<String>,
486    pub producer: Option<ProducerRequest>,
487    pub now_ms: u64,
488    pub record_match: Option<u64>,
489}
490
491impl AppendExternalRequest {
492    pub fn from_append_request(
493        request: AppendRequest,
494        payload: ExternalPayloadRef,
495        record_ends: Vec<u64>,
496    ) -> Self {
497        Self {
498            stream_id: request.stream_id,
499            content_type: request.content_type,
500            payload,
501            record_ends,
502            close_after: request.close_after,
503            stream_seq: request.stream_seq,
504            producer: request.producer,
505            now_ms: request.now_ms,
506            record_match: request.record_match,
507        }
508    }
509}
510
511impl AppendRequest {
512    pub fn canonical_record_ends(&self) -> Vec<u64> {
513        ursula_stream::canonical_json_record_ends(&self.content_type, &self.payload)
514            .unwrap_or_default()
515    }
516
517    pub fn new(stream_id: BucketStreamId, payload_len: u64) -> Self {
518        Self {
519            stream_id,
520            content_type: DEFAULT_CONTENT_TYPE.to_owned(),
521            payload: Bytes::from(vec![
522                0;
523                usize::try_from(payload_len)
524                    .expect("payload_len fits usize")
525            ]),
526            close_after: false,
527            stream_seq: None,
528            producer: None,
529            now_ms: 0,
530            record_match: None,
531        }
532    }
533
534    pub fn from_bytes(stream_id: BucketStreamId, payload: impl Into<Bytes>) -> Self {
535        Self {
536            stream_id,
537            content_type: DEFAULT_CONTENT_TYPE.to_owned(),
538            payload: payload.into(),
539            close_after: false,
540            stream_seq: None,
541            producer: None,
542            now_ms: 0,
543            record_match: None,
544        }
545    }
546
547    pub fn payload_len(&self) -> u64 {
548        u64::try_from(self.payload.len()).expect("payload len fits u64")
549    }
550}
551
552#[derive(Debug, Clone, PartialEq, Eq)]
553pub struct AppendBatchRequest {
554    pub stream_id: BucketStreamId,
555    pub content_type: String,
556    pub payloads: Vec<Bytes>,
557    pub producer: Option<ProducerRequest>,
558    pub now_ms: u64,
559}
560
561impl AppendBatchRequest {
562    pub fn new<P>(stream_id: BucketStreamId, payloads: Vec<P>) -> Self
563    where P: Into<Bytes> {
564        Self {
565            stream_id,
566            content_type: DEFAULT_CONTENT_TYPE.to_owned(),
567            payloads: payloads.into_iter().map(Into::into).collect(),
568            producer: None,
569            now_ms: 0,
570        }
571    }
572}
573
574#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
575pub struct AppendResponse {
576    pub placement: ShardPlacement,
577    pub start_offset: u64,
578    pub next_offset: u64,
579    pub stream_append_count: u64,
580    pub group_commit_index: u64,
581    pub closed: bool,
582    pub deduplicated: bool,
583    pub producer: Option<ProducerRequest>,
584    pub record_range: Option<StreamRecordRange>,
585}
586
587#[derive(Debug, Clone, PartialEq, Eq)]
588pub struct AppendBatchResponse {
589    pub placement: ShardPlacement,
590    pub items: Vec<Result<AppendResponse, RuntimeError>>,
591}
592
593#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
594pub struct StreamAppendCount {
595    pub stream_id: BucketStreamId,
596    pub append_count: u64,
597}