1use bytes::Bytes;
4use notedthat_core::{ConditionalHeaders, KbSlug, ObjectPath, PutOutcome, Storage, StorageError};
5use notedthat_indexer::IndexEvent;
6use tokio::sync::mpsc::Sender;
7use tokio::sync::mpsc::error::TrySendError;
8
9use crate::WriteError;
10use crate::mime::sniff_content_type;
11
12pub const MAX_UPLOAD_BYTES: u64 = 5 * 1024 * 1024 * 1024;
14
15pub fn check_size(size: u64, limit: u64) -> Result<(), WriteError> {
17 if size > limit {
18 Err(WriteError::TooLarge { size, limit })
19 } else {
20 Ok(())
21 }
22}
23
24pub async fn commit(
26 storage: &dyn Storage,
27 indexer_tx: &Sender<IndexEvent>,
28 kb: &KbSlug,
29 path: &ObjectPath,
30 bytes: Bytes,
31 caller_content_type: Option<&str>,
32 conditionals: ConditionalHeaders,
33) -> Result<PutOutcome, WriteError> {
34 check_size(bytes.len() as u64, MAX_UPLOAD_BYTES)?;
35 let mime = sniff_content_type(caller_content_type, path);
36 let outcome = storage
37 .put_object(kb, path, bytes, Some(&mime), conditionals)
38 .await?;
39
40 let event = IndexEvent::Upsert {
41 kb: kb.clone(),
42 object_key: path.clone(),
43 etag: outcome.etag.clone().unwrap_or_default(),
44 mtime: current_unix_seconds(),
45 };
46 match indexer_tx.try_send(event) {
47 Ok(()) => {}
48 Err(TrySendError::Full(ev)) => {
49 tracing::warn!(target: "notedthat::indexing", kb = %kb, path = %path, "INDEX_QUEUE_FULL");
50 let _ = ev;
51 return Err(WriteError::IndexerBackpressureUpsert);
52 }
53 Err(TrySendError::Closed(ev)) => {
54 tracing::error!(target: "notedthat::indexing", kb = %kb, path = %path, "INDEX_QUEUE_CLOSED");
55 let _ = ev;
56 }
59 }
60
61 Ok(outcome)
62}
63
64pub async fn commit_delete(
66 storage: &dyn Storage,
67 indexer_tx: &Sender<IndexEvent>,
68 kb: &KbSlug,
69 path: &ObjectPath,
70 conditionals: ConditionalHeaders,
71) -> Result<(), WriteError> {
72 match storage.delete_object(kb, path, conditionals).await {
73 Ok(()) | Err(StorageError::NotFound { .. }) => {}
74 Err(e) => return Err(WriteError::Storage(e)),
75 }
76
77 let event = IndexEvent::Tombstone {
78 kb: kb.clone(),
79 object_key: path.clone(),
80 };
81 match indexer_tx.try_send(event) {
82 Ok(()) => {}
83 Err(TrySendError::Full(ev)) => {
84 tracing::warn!(target: "notedthat::indexing", kb = %kb, path = %path, "INDEX_QUEUE_FULL");
85 let _ = ev;
86 return Err(WriteError::IndexerBackpressureTombstone);
87 }
88 Err(TrySendError::Closed(ev)) => {
90 tracing::error!(target: "notedthat::indexing", kb = %kb, path = %path, "INDEX_QUEUE_CLOSED");
91 let _ = ev;
92 }
93 }
94
95 Ok(())
96}
97
98fn current_unix_seconds() -> i64 {
99 std::time::SystemTime::now()
100 .duration_since(std::time::UNIX_EPOCH)
101 .map_or(0, |d| i64::try_from(d.as_secs()).unwrap_or(i64::MAX))
102}
103
104#[cfg(test)]
105mod tests {
106 use super::*;
107 use async_trait::async_trait;
108 use notedthat_core::{KbManifest, ListResponse, ObjectMeta, ObjectRead};
109 use std::collections::HashMap;
110 use std::sync::{Arc, Mutex};
111 use tokio::sync::mpsc;
112
113 #[derive(Default)]
114 struct TestStorage {
115 objects: Mutex<HashMap<String, String>>,
116 }
117
118 #[async_trait]
119 impl Storage for TestStorage {
120 async fn ensure_bucket(&self, _kb: &KbSlug) -> Result<(), StorageError> {
121 unimplemented!()
122 }
123
124 async fn read_manifest(&self, _kb: &KbSlug) -> Result<KbManifest, StorageError> {
125 unimplemented!()
126 }
127
128 async fn write_manifest(
129 &self,
130 _kb: &KbSlug,
131 _manifest: &KbManifest,
132 ) -> Result<(), StorageError> {
133 unimplemented!()
134 }
135
136 async fn head_object(
137 &self,
138 _kb: &KbSlug,
139 _path: &ObjectPath,
140 _conditionals: ConditionalHeaders,
141 ) -> Result<ObjectMeta, StorageError> {
142 unimplemented!()
143 }
144
145 async fn get_object(
146 &self,
147 _kb: &KbSlug,
148 _path: &ObjectPath,
149 _range: Option<Vec<notedthat_core::ByteRange>>,
150 _conditionals: ConditionalHeaders,
151 ) -> Result<ObjectRead, StorageError> {
152 unimplemented!()
153 }
154
155 async fn put_object(
156 &self,
157 kb: &KbSlug,
158 path: &ObjectPath,
159 _bytes: Bytes,
160 _content_type: Option<&str>,
161 conditionals: ConditionalHeaders,
162 ) -> Result<PutOutcome, StorageError> {
163 let key = format!("{}/{}", kb.as_str(), path.as_str());
164 let mut objects = self.objects.lock().expect("mutex not poisoned");
165 let existing = objects.get(&key);
166 if let Some(if_match) = conditionals.if_match
167 && existing.is_none_or(|etag| etag != &if_match)
168 {
169 return Err(StorageError::PreconditionFailed);
170 }
171
172 let etag = format!("\"etag-{}\"", objects.len() + 1);
173 objects.insert(key, etag.clone());
174 Ok(PutOutcome { etag: Some(etag) })
175 }
176
177 async fn delete_object(
178 &self,
179 kb: &KbSlug,
180 path: &ObjectPath,
181 _conditionals: ConditionalHeaders,
182 ) -> Result<(), StorageError> {
183 let key = format!("{}/{}", kb.as_str(), path.as_str());
184 self.objects
185 .lock()
186 .expect("mutex not poisoned")
187 .remove(&key);
188 Ok(())
189 }
190
191 async fn list_objects(
192 &self,
193 _kb: &KbSlug,
194 _prefix: Option<&str>,
195 _limit: u32,
196 _cursor: Option<&str>,
197 ) -> Result<ListResponse, StorageError> {
198 unimplemented!()
199 }
200 }
201
202 fn kb() -> KbSlug {
203 KbSlug::try_new("test-kb").expect("valid kb slug")
204 }
205
206 fn path() -> ObjectPath {
207 ObjectPath::try_from_str("test.md").expect("valid path")
208 }
209
210 fn path_named(value: &str) -> ObjectPath {
211 ObjectPath::try_from_str(value).expect("valid path")
212 }
213
214 #[tokio::test]
215 async fn successful_put_enqueues_event() {
216 let storage = TestStorage::default();
217 let kb = kb();
218 let path = path();
219 let (indexer_tx, mut rx) = mpsc::channel(1024);
220
221 let outcome = commit(
222 &storage,
223 &indexer_tx,
224 &kb,
225 &path,
226 Bytes::from_static(b"# Test"),
227 Some("text/markdown"),
228 ConditionalHeaders::default(),
229 )
230 .await;
231
232 assert!(outcome.is_ok());
233 assert!(outcome.unwrap().etag.is_some());
234
235 let event = rx.recv().await.expect("event should be enqueued");
236 assert_eq!(event.kb().as_str(), "test-kb");
237 assert_eq!(event.object_key().as_str(), "test.md");
238 }
239
240 #[tokio::test]
241 async fn full_queue_returns_indexer_backpressure() {
242 let storage = TestStorage::default();
243 let kb = kb();
244 let path = path();
245 let (indexer_tx, _rx) = mpsc::channel(1);
246
247 let dummy_event = IndexEvent::Upsert {
248 kb: kb.clone(),
249 object_key: path.clone(),
250 etag: "dummy".to_string(),
251 mtime: 0,
252 };
253 indexer_tx
254 .try_send(dummy_event)
255 .expect("first send should succeed");
256
257 let outcome = commit(
258 &storage,
259 &indexer_tx,
260 &kb,
261 &path,
262 Bytes::from_static(b"# Test"),
263 Some("text/markdown"),
264 ConditionalHeaders::default(),
265 )
266 .await;
267
268 let err = outcome.unwrap_err();
269 assert!(
270 matches!(err, WriteError::IndexerBackpressureUpsert),
271 "expected IndexerBackpressureUpsert, got {err:?}"
272 );
273 }
274
275 #[tokio::test]
276 async fn burst_write_returns_backpressure_after_capacity() {
277 let storage = Arc::new(TestStorage::default());
278 let kb = kb();
279 let path_a = path_named("a.md");
280 let path_b = path_named("b.md");
281 let path_c = path_named("c.md");
282 let (indexer_tx, _rx) = mpsc::channel::<IndexEvent>(2);
283
284 let first = commit(
285 storage.as_ref(),
286 &indexer_tx,
287 &kb,
288 &path_a,
289 Bytes::from_static(b"# A"),
290 Some("text/markdown"),
291 ConditionalHeaders::default(),
292 )
293 .await;
294 let second = commit(
295 storage.as_ref(),
296 &indexer_tx,
297 &kb,
298 &path_b,
299 Bytes::from_static(b"# B"),
300 Some("text/markdown"),
301 ConditionalHeaders::default(),
302 )
303 .await;
304 let third = commit(
305 storage.as_ref(),
306 &indexer_tx,
307 &kb,
308 &path_c,
309 Bytes::from_static(b"# C"),
310 Some("text/markdown"),
311 ConditionalHeaders::default(),
312 )
313 .await;
314
315 assert!(first.is_ok(), "first write should fill queue slot one");
316 assert!(second.is_ok(), "second write should fill queue slot two");
317 let err = third.unwrap_err();
318 assert!(
319 matches!(err, WriteError::IndexerBackpressureUpsert),
320 "expected IndexerBackpressureUpsert, got {err:?}"
321 );
322 assert!(
323 storage
324 .objects
325 .lock()
326 .expect("mutex not poisoned")
327 .contains_key("test-kb/c.md"),
328 "stored object should remain after enqueue backpressure"
329 );
330 }
331
332 #[tokio::test]
333 async fn commit_delete_full_queue_returns_backpressure() {
334 let storage = TestStorage::default();
335 let kb = kb();
336 let path = path_named("delete.md");
337 storage
338 .put_object(
339 &kb,
340 &path,
341 Bytes::from_static(b"# Delete"),
342 Some("text/markdown"),
343 ConditionalHeaders::default(),
344 )
345 .await
346 .expect("prepopulate object");
347 let (indexer_tx, _rx) = mpsc::channel(1);
348 let dummy_event = IndexEvent::Tombstone {
349 kb: kb.clone(),
350 object_key: path.clone(),
351 };
352 indexer_tx
353 .try_send(dummy_event)
354 .expect("first send should succeed");
355
356 let outcome = commit_delete(
357 &storage,
358 &indexer_tx,
359 &kb,
360 &path,
361 ConditionalHeaders::default(),
362 )
363 .await;
364
365 let err = outcome.unwrap_err();
366 assert!(
367 matches!(err, WriteError::IndexerBackpressureTombstone),
368 "expected IndexerBackpressureTombstone, got {err:?}"
369 );
370 assert!(
371 !storage
372 .objects
373 .lock()
374 .expect("mutex not poisoned")
375 .contains_key("test-kb/delete.md"),
376 "deleted object should remain deleted after enqueue backpressure"
377 );
378 }
379
380 #[tokio::test]
381 async fn closed_queue_returns_write_success() {
382 let storage = TestStorage::default();
383 let kb = kb();
384 let path = path();
385 let (indexer_tx, rx) = mpsc::channel(1024);
386
387 drop(rx);
388
389 let outcome = commit(
390 &storage,
391 &indexer_tx,
392 &kb,
393 &path,
394 Bytes::from_static(b"# Test"),
395 Some("text/markdown"),
396 ConditionalHeaders::default(),
397 )
398 .await;
399
400 assert!(
401 outcome.is_ok(),
402 "write should succeed even if queue is closed"
403 );
404 }
405
406 #[tokio::test]
407 async fn put_failure_returns_error_no_event() {
408 let storage = TestStorage::default();
409 let kb = kb();
410 let path = path();
411 let (indexer_tx, mut rx) = mpsc::channel(1024);
412
413 let conditionals = ConditionalHeaders {
414 if_match: Some("\"wrong-etag\"".to_string()),
415 ..ConditionalHeaders::default()
416 };
417
418 let outcome = commit(
419 &storage,
420 &indexer_tx,
421 &kb,
422 &path,
423 Bytes::from_static(b"# Test"),
424 Some("text/markdown"),
425 conditionals,
426 )
427 .await;
428
429 assert!(outcome.is_err(), "put should fail with precondition");
430 assert!(
431 rx.try_recv().is_err(),
432 "no event should be enqueued on put failure"
433 );
434 }
435
436 #[test]
437 fn check_size_over_limit_returns_too_large() {
438 let err = check_size(MAX_UPLOAD_BYTES + 1, MAX_UPLOAD_BYTES).expect_err("too large");
439 assert!(matches!(
440 err,
441 WriteError::TooLarge {
442 size,
443 limit
444 } if size == MAX_UPLOAD_BYTES + 1 && limit == MAX_UPLOAD_BYTES
445 ));
446 }
447
448 #[test]
449 fn check_size_at_limit_returns_ok() {
450 assert!(check_size(MAX_UPLOAD_BYTES, MAX_UPLOAD_BYTES).is_ok());
451 }
452
453 #[test]
454 fn check_size_below_limit_returns_ok() {
455 assert!(check_size(1024, MAX_UPLOAD_BYTES).is_ok());
456 }
457}