faucet_cli/serve/triggers/
object_arrival.rs1use super::context::TriggerEvent;
5use super::enqueue::{self, FireOutcome};
6use super::spec::{ArrivalMode, StartAt, StoreSpec};
7use super::watcher::Watcher;
8use crate::serve::state::ServerState;
9use async_trait::async_trait;
10use futures::StreamExt;
11use object_store::ObjectStore;
12use std::sync::Arc;
13use std::time::Duration;
14
15#[derive(Debug, Clone, PartialEq, Eq)]
18pub struct ListedObject {
19 pub key: String,
20 pub last_modified: chrono::DateTime<chrono::Utc>,
21 pub size: u64,
22 pub etag: Option<String>,
23}
24
25#[derive(Debug, Default)]
29pub struct Cursor {
30 watermark: Option<chrono::DateTime<chrono::Utc>>,
31 seen_at_watermark: std::collections::HashSet<String>,
32}
33
34impl Cursor {
35 pub fn starting_now(now: chrono::DateTime<chrono::Utc>) -> Self {
37 Self {
38 watermark: Some(now),
39 seen_at_watermark: std::collections::HashSet::new(),
40 }
41 }
42
43 pub fn starting_beginning() -> Self {
45 Self::default()
46 }
47
48 pub fn new_objects(&self, listing: &[ListedObject]) -> Vec<ListedObject> {
52 let mut out = Vec::new();
53 for o in listing {
54 match self.watermark {
55 None => out.push(o.clone()),
56 Some(w) if o.last_modified > w => out.push(o.clone()),
57 Some(w) if o.last_modified == w && !self.seen_at_watermark.contains(&o.key) => {
58 out.push(o.clone())
59 }
60 _ => {}
61 }
62 }
63 out
64 }
65
66 pub fn commit(&mut self, o: &ListedObject) {
68 match self.watermark {
69 Some(w) if o.last_modified > w => {
70 self.watermark = Some(o.last_modified);
71 self.seen_at_watermark.clear();
72 self.seen_at_watermark.insert(o.key.clone());
73 }
74 Some(w) if o.last_modified == w => {
75 self.seen_at_watermark.insert(o.key.clone());
76 }
77 None => {
78 self.watermark = Some(o.last_modified);
79 self.seen_at_watermark.insert(o.key.clone());
80 }
81 _ => {}
82 }
83 }
84}
85
86type StoreTriple = (Arc<dyn ObjectStore>, String, Option<String>);
89
90pub struct ObjectArrivalWatcher {
91 name: String,
92 store: Arc<dyn ObjectStore>,
93 bucket: String,
94 prefix: Option<String>,
95 mode: ArrivalMode,
96 poll: Duration,
97 cursor: Cursor,
98 compiled: Arc<super::compiled::CompiledTrigger>,
99}
100
101impl ObjectArrivalWatcher {
102 pub fn build_store(store: &StoreSpec) -> Result<StoreTriple, String> {
104 match store {
105 StoreSpec::S3 {
106 bucket,
107 prefix,
108 region,
109 endpoint,
110 } => {
111 let mut b = object_store::aws::AmazonS3Builder::from_env().with_bucket_name(bucket);
112 if let Some(r) = region {
113 b = b.with_region(r);
114 }
115 if let Some(e) = endpoint {
116 b = b.with_endpoint(e).with_allow_http(true);
117 }
118 let s = b.build().map_err(|e| format!("building S3 client: {e}"))?;
119 Ok((Arc::new(s), bucket.clone(), prefix.clone()))
120 }
121 StoreSpec::Gcs { bucket, prefix } => {
122 let s = object_store::gcp::GoogleCloudStorageBuilder::from_env()
123 .with_bucket_name(bucket)
124 .build()
125 .map_err(|e| format!("building GCS client: {e}"))?;
126 Ok((Arc::new(s), bucket.clone(), prefix.clone()))
127 }
128 }
129 }
130
131 #[allow(clippy::too_many_arguments)]
132 pub fn new(
133 compiled: Arc<super::compiled::CompiledTrigger>,
134 store: Arc<dyn ObjectStore>,
135 bucket: String,
136 prefix: Option<String>,
137 mode: ArrivalMode,
138 poll: Duration,
139 start_at: StartAt,
140 now: chrono::DateTime<chrono::Utc>,
141 ) -> Self {
142 let cursor = match start_at {
143 StartAt::Now => Cursor::starting_now(now),
144 StartAt::Beginning => Cursor::starting_beginning(),
145 };
146 Self {
147 name: compiled.name().to_string(),
148 store,
149 bucket,
150 prefix,
151 mode,
152 poll,
153 cursor,
154 compiled,
155 }
156 }
157
158 async fn list(&self) -> Result<Vec<ListedObject>, String> {
159 let prefix_path = self.prefix.as_deref().map(object_store::path::Path::from);
160 let mut stream = self.store.list(prefix_path.as_ref());
161 let mut out = Vec::new();
162 while let Some(meta) = stream.next().await {
163 let meta = meta.map_err(|e| format!("listing objects: {e}"))?;
164 out.push(ListedObject {
165 key: meta.location.to_string(),
166 last_modified: meta.last_modified,
167 size: meta.size,
168 etag: meta.e_tag,
169 });
170 }
171 Ok(out)
172 }
173}
174
175#[async_trait]
176impl Watcher for ObjectArrivalWatcher {
177 fn name(&self) -> &str {
178 &self.name
179 }
180
181 fn kind(&self) -> &'static str {
182 "object_arrival"
183 }
184
185 fn poll_interval(&self) -> Duration {
186 self.poll
187 }
188
189 async fn poll(&mut self, state: &ServerState) -> Result<bool, String> {
190 let listing = self.list().await?;
191 let mut new = self.cursor.new_objects(&listing);
192 if new.is_empty() {
193 return Ok(false);
194 }
195 new.sort_by(|a, b| {
197 a.last_modified
198 .cmp(&b.last_modified)
199 .then(a.key.cmp(&b.key))
200 });
201 let fired_at = chrono::Utc::now().to_rfc3339();
202 let mut fired = false;
203
204 match self.mode {
205 ArrivalMode::PerObject => {
206 for o in new {
207 let event = TriggerEvent::Object {
208 bucket: self.bucket.clone(),
209 key: o.key.clone(),
210 size: o.size,
211 last_modified: o.last_modified.to_rfc3339(),
212 };
213 match enqueue::fire(state, &self.compiled, event, &fired_at).await {
214 outcome if outcome.committed() => {
215 self.cursor.commit(&o);
216 fired = true;
217 }
218 FireOutcome::Dropped(_) => break, FireOutcome::Error(_) => break,
220 _ => {}
221 }
222 }
223 }
224 ArrivalMode::Batch => {
225 let watermark = new.iter().map(|o| o.last_modified).max().unwrap();
226 let event = TriggerEvent::ObjectBatch {
227 bucket: self.bucket.clone(),
228 count: new.len(),
229 watermark: watermark.to_rfc3339(),
230 };
231 match enqueue::fire(state, &self.compiled, event, &fired_at).await {
232 outcome if outcome.committed() => {
233 for o in &new {
234 self.cursor.commit(o);
235 }
236 fired = true;
237 }
238 _ => {} }
240 }
241 }
242 Ok(fired)
243 }
244}
245
246#[cfg(test)]
247mod cursor_tests {
248 use super::*;
249 use chrono::TimeZone;
250
251 fn obj(key: &str, secs: i64) -> ListedObject {
252 ListedObject {
253 key: key.into(),
254 last_modified: chrono::Utc.timestamp_opt(secs, 0).unwrap(),
255 size: 1,
256 etag: None,
257 }
258 }
259
260 #[test]
261 fn starting_now_ignores_existing() {
262 let now = chrono::Utc.timestamp_opt(1000, 0).unwrap();
263 let c = Cursor::starting_now(now);
264 assert!(c.new_objects(&[obj("a", 900)]).is_empty());
266 assert_eq!(c.new_objects(&[obj("b", 1100)]).len(), 1);
268 }
269
270 #[test]
271 fn starting_beginning_emits_all_then_commits() {
272 let mut c = Cursor::starting_beginning();
273 let listing = vec![obj("a", 100), obj("b", 200)];
274 let new = c.new_objects(&listing);
275 assert_eq!(new.len(), 2);
276 for o in &new {
277 c.commit(o);
278 }
279 assert!(c.new_objects(&listing).is_empty());
281 }
282
283 #[test]
284 fn handles_ties_at_watermark() {
285 let mut c = Cursor::starting_beginning();
286 let a = obj("a", 100);
287 c.commit(&c.new_objects(std::slice::from_ref(&a))[0].clone());
288 let b = obj("b", 100);
290 let new = c.new_objects(&[a.clone(), b.clone()]);
291 assert_eq!(new, vec![b.clone()]);
292 c.commit(&b);
293 assert!(c.new_objects(&[a, b]).is_empty());
294 }
295
296 #[test]
297 fn dropped_fire_is_retried_until_committed() {
298 let mut c = Cursor::starting_beginning();
299 let a = obj("a", 100);
300 assert_eq!(c.new_objects(std::slice::from_ref(&a)).len(), 1);
302 assert_eq!(c.new_objects(std::slice::from_ref(&a)).len(), 1);
303 c.commit(&a);
304 assert!(c.new_objects(&[a]).is_empty());
305 }
306}