1use std::sync::Arc;
2
3use chrono::{Duration, Utc};
4
5use eventuary_core::io::OwnerId;
6use eventuary_core::io::reader::{
7 CheckpointScope, Generation, PartitionCoordinator, PartitionLease,
8};
9use eventuary_core::{Error, Partition, Result};
10
11use crate::database::SqliteConn;
12use crate::reader::SqliteCursor;
13use crate::relation::SqliteRelationName;
14use crate::schema::{Migration, RelationReplacement};
15
16const PARTITION_COORDINATOR_0001_INIT_SQL: &str = r#"
17CREATE TABLE IF NOT EXISTS {consumers} (
18 consumer_group_id TEXT NOT NULL,
19 stream_id TEXT NOT NULL,
20 owner_id TEXT NOT NULL,
21 lease_until TEXT NOT NULL,
22 PRIMARY KEY (consumer_group_id, stream_id, owner_id)
23);
24
25CREATE INDEX IF NOT EXISTS idx_event_stream_consumers_group_stream_lease
26ON {consumers} (consumer_group_id, stream_id, lease_until);
27
28CREATE TABLE IF NOT EXISTS {partitions} (
29 consumer_group_id TEXT NOT NULL,
30 stream_id TEXT NOT NULL,
31 partition_id INTEGER NOT NULL,
32 partition_count INTEGER NULL,
33 owner_id TEXT NULL,
34 lease_until TEXT NULL,
35 checkpoint_sequence INTEGER NOT NULL DEFAULT 0,
36 generation INTEGER NOT NULL DEFAULT 0,
37 PRIMARY KEY (consumer_group_id, stream_id, partition_id)
38);
39
40CREATE INDEX IF NOT EXISTS idx_event_stream_partitions_group_stream_owner
41ON {partitions} (consumer_group_id, stream_id, owner_id);
42
43CREATE INDEX IF NOT EXISTS idx_event_stream_partitions_group_stream_count
44ON {partitions} (consumer_group_id, stream_id, partition_count, partition_id);
45"#;
46
47const PARTITION_COORDINATOR_MIGRATIONS: &[Migration] = &[Migration {
48 name: "0001_init",
49 sql: PARTITION_COORDINATOR_0001_INIT_SQL,
50}];
51
52#[derive(Debug, Clone)]
53pub struct SqlitePartitionCoordinatorConfig {
54 pub consumers_relation: SqliteRelationName,
55 pub partitions_relation: SqliteRelationName,
56}
57
58impl Default for SqlitePartitionCoordinatorConfig {
59 fn default() -> Self {
60 Self {
61 consumers_relation: SqliteRelationName::new("event_stream_consumers")
62 .expect("default consumers relation"),
63 partitions_relation: SqliteRelationName::new("event_stream_partitions")
64 .expect("default partitions relation"),
65 }
66 }
67}
68
69pub struct SqlitePartitionCoordinator {
70 conn: SqliteConn,
71 consumers_relation: Arc<String>,
72 partitions_relation: Arc<String>,
73}
74
75impl Clone for SqlitePartitionCoordinator {
76 fn clone(&self) -> Self {
77 Self {
78 conn: Arc::clone(&self.conn),
79 consumers_relation: Arc::clone(&self.consumers_relation),
80 partitions_relation: Arc::clone(&self.partitions_relation),
81 }
82 }
83}
84
85impl SqlitePartitionCoordinator {
86 pub fn new(conn: SqliteConn, config: SqlitePartitionCoordinatorConfig) -> Self {
87 Self {
88 conn,
89 consumers_relation: Arc::new(config.consumers_relation.render()),
90 partitions_relation: Arc::new(config.partitions_relation.render()),
91 }
92 }
93
94 pub fn connect(conn: SqliteConn, config: SqlitePartitionCoordinatorConfig) -> Result<Self> {
95 Self::prepare_schema(&conn, &config)?;
96 Ok(Self::new(conn, config))
97 }
98
99 pub fn prepare_schema(
100 conn: &SqliteConn,
101 config: &SqlitePartitionCoordinatorConfig,
102 ) -> Result<()> {
103 let guard = conn.lock().map_err(|e| Error::Store(e.to_string()))?;
104 crate::schema::apply_schema(
105 &guard,
106 PARTITION_COORDINATOR_MIGRATIONS,
107 &[
108 RelationReplacement {
109 token: "{consumers}",
110 relation: &config.consumers_relation,
111 },
112 RelationReplacement {
113 token: "{partitions}",
114 relation: &config.partitions_relation,
115 },
116 ],
117 )
118 }
119
120 pub fn schema_sql(config: &SqlitePartitionCoordinatorConfig) -> String {
121 crate::schema::render_schema_sql(
122 PARTITION_COORDINATOR_MIGRATIONS,
123 &[
124 RelationReplacement {
125 token: "{consumers}",
126 relation: &config.consumers_relation,
127 },
128 RelationReplacement {
129 token: "{partitions}",
130 relation: &config.partitions_relation,
131 },
132 ],
133 )
134 }
135}
136
137fn compute_lease_until(lease_duration: std::time::Duration) -> Result<String> {
138 let dt = Utc::now()
139 + Duration::from_std(lease_duration)
140 .map_err(|_| Error::Config("lease duration out of range".to_owned()))?;
141 Ok(dt.to_rfc3339())
142}
143
144impl PartitionCoordinator<SqliteCursor> for SqlitePartitionCoordinator {
145 async fn heartbeat<'a>(
146 &'a self,
147 scope: &'a CheckpointScope,
148 owner_id: &'a OwnerId,
149 lease_duration: std::time::Duration,
150 ) -> Result<()> {
151 let lease_until = compute_lease_until(lease_duration)?;
152 let conn = Arc::clone(&self.conn);
153 let consumers = Arc::clone(&self.consumers_relation);
154 let group = scope.consumer_group_id.as_str().to_owned();
155 let stream = scope.stream_id.as_str().to_owned();
156 let owner = owner_id.as_str().to_owned();
157 tokio::task::spawn_blocking(move || {
158 let guard = conn.lock().map_err(|e| Error::Store(e.to_string()))?;
159 let sql = format!(
160 "INSERT INTO {consumers} (consumer_group_id, stream_id, owner_id, lease_until) \
161 VALUES (?1, ?2, ?3, ?4) \
162 ON CONFLICT (consumer_group_id, stream_id, owner_id) \
163 DO UPDATE SET lease_until = excluded.lease_until"
164 );
165 guard
166 .execute(&sql, rusqlite::params![group, stream, owner, lease_until])
167 .map_err(|e| Error::Store(e.to_string()))?;
168 Ok(())
169 })
170 .await
171 .map_err(|e| Error::Store(format!("blocking task panicked: {e}")))?
172 }
173
174 async fn live_consumers<'a>(&'a self, scope: &'a CheckpointScope) -> Result<usize> {
175 let conn = Arc::clone(&self.conn);
176 let consumers = Arc::clone(&self.consumers_relation);
177 let group = scope.consumer_group_id.as_str().to_owned();
178 let stream = scope.stream_id.as_str().to_owned();
179 let now = Utc::now().to_rfc3339();
180 tokio::task::spawn_blocking(move || {
181 let guard = conn.lock().map_err(|e| Error::Store(e.to_string()))?;
182 let sql = format!(
183 "SELECT COUNT(*) FROM {consumers} \
184 WHERE consumer_group_id = ?1 \
185 AND stream_id = ?2 \
186 AND lease_until > ?3"
187 );
188 let count: i64 = guard
189 .query_row(&sql, rusqlite::params![group, stream, now], |r| r.get(0))
190 .map_err(|e| Error::Store(e.to_string()))?;
191 Ok(count as usize)
192 })
193 .await
194 .map_err(|e| Error::Store(format!("blocking task panicked: {e}")))?
195 }
196
197 async fn release_consumer<'a>(
198 &'a self,
199 scope: &'a CheckpointScope,
200 owner_id: &'a OwnerId,
201 ) -> Result<()> {
202 let conn = Arc::clone(&self.conn);
203 let consumers = Arc::clone(&self.consumers_relation);
204 let group = scope.consumer_group_id.as_str().to_owned();
205 let stream = scope.stream_id.as_str().to_owned();
206 let owner = owner_id.as_str().to_owned();
207 tokio::task::spawn_blocking(move || {
208 let guard = conn.lock().map_err(|e| Error::Store(e.to_string()))?;
209 let sql = format!(
210 "DELETE FROM {consumers} \
211 WHERE consumer_group_id = ?1 AND stream_id = ?2 AND owner_id = ?3"
212 );
213 guard
214 .execute(&sql, rusqlite::params![group, stream, owner])
215 .map_err(|e| Error::Store(e.to_string()))?;
216 Ok(())
217 })
218 .await
219 .map_err(|e| Error::Store(format!("blocking task panicked: {e}")))?
220 }
221
222 async fn claim<'a>(
223 &'a self,
224 scope: &'a CheckpointScope,
225 owner_id: &'a OwnerId,
226 partition: Partition,
227 lease_duration: std::time::Duration,
228 ) -> Result<Option<PartitionLease<SqliteCursor>>> {
229 let lease_until = compute_lease_until(lease_duration)?;
230 let conn = Arc::clone(&self.conn);
231 let partitions = Arc::clone(&self.partitions_relation);
232 let group = scope.consumer_group_id.as_str().to_owned();
233 let stream = scope.stream_id.as_str().to_owned();
234 let owner = owner_id.as_str().to_owned();
235 let partition_id_i64 = partition.id() as i64;
236 let partition_count_i64 = partition.count() as i64;
237 let now = Utc::now().to_rfc3339();
238 let scope = scope.clone();
239 let owner_id = owner_id.clone();
240 tokio::task::spawn_blocking(move || {
241 let guard = conn.lock().map_err(|e| Error::Store(e.to_string()))?;
242 let sql = format!(
243 "INSERT INTO {partitions} \
244 (consumer_group_id, stream_id, partition_id, partition_count, owner_id, lease_until, generation, checkpoint_sequence) \
245 VALUES (?1, ?2, ?3, ?4, ?5, ?6, 1, 0) \
246 ON CONFLICT (consumer_group_id, stream_id, partition_id) DO UPDATE \
247 SET owner_id = excluded.owner_id, \
248 lease_until = excluded.lease_until, \
249 partition_count = COALESCE({partitions}.partition_count, excluded.partition_count), \
250 generation = {partitions}.generation + 1 \
251 WHERE ({partitions}.partition_count IS NULL OR {partitions}.partition_count = excluded.partition_count) \
252 AND ({partitions}.owner_id IS NULL \
253 OR {partitions}.lease_until IS NULL \
254 OR {partitions}.lease_until < ?7 \
255 OR {partitions}.owner_id = excluded.owner_id) \
256 RETURNING owner_id, lease_until, generation, checkpoint_sequence"
257 );
258 let row = guard
259 .query_row(
260 &sql,
261 rusqlite::params![
262 group,
263 stream,
264 partition_id_i64,
265 partition_count_i64,
266 owner,
267 lease_until,
268 now
269 ],
270 |r| {
271 Ok((
272 r.get::<_, String>(0)?,
273 r.get::<_, String>(1)?,
274 r.get::<_, i64>(2)?,
275 r.get::<_, i64>(3)?,
276 ))
277 },
278 )
279 .map(Some)
280 .or_else(|e| match e {
281 rusqlite::Error::QueryReturnedNoRows => Ok(None),
282 other => Err(other),
283 })
284 .map_err(|e| Error::Store(e.to_string()))?;
285 match row {
286 None => {
287 let check_sql = format!(
288 "SELECT partition_count FROM {partitions} \
289 WHERE consumer_group_id = ?1 AND stream_id = ?2 AND partition_id = ?3"
290 );
291 let stored: Option<i64> = guard
292 .query_row(
293 &check_sql,
294 rusqlite::params![group, stream, partition_id_i64],
295 |r| r.get::<_, Option<i64>>(0),
296 )
297 .map(Some)
298 .or_else(|e| match e {
299 rusqlite::Error::QueryReturnedNoRows => Ok(None),
300 other => Err(other),
301 })
302 .map_err(|e| Error::Store(e.to_string()))?
303 .flatten();
304 if let Some(stored) = stored
305 && stored != partition_count_i64
306 {
307 return Err(Error::Config(format!(
308 "partition count mismatch for scope {} stream {} partition {}: stored {}, requested {}",
309 group, stream, partition_id_i64, stored, partition_count_i64,
310 )));
311 }
312 Ok(None)
313 }
314 Some((_returned_owner, lease_until_text, generation, checkpoint_sequence)) => {
315 let lease_until = chrono::DateTime::parse_from_rfc3339(&lease_until_text)
316 .map(|dt| dt.with_timezone(&Utc))
317 .map_err(|e| {
318 Error::Serialization(format!("lease_until decode: {e}"))
319 })?;
320 Ok(Some(PartitionLease {
321 scope,
322 owner_id,
323 partition,
324 generation: Generation::from_i64(generation),
325 checkpoint_cursor: (checkpoint_sequence > 0)
326 .then_some(SqliteCursor::new(checkpoint_sequence, partition)),
327 lease_until,
328 }))
329 }
330 }
331 })
332 .await
333 .map_err(|e| Error::Store(format!("blocking task panicked: {e}")))?
334 }
335
336 async fn renew<'a>(
337 &'a self,
338 lease: &'a PartitionLease<SqliteCursor>,
339 lease_duration: std::time::Duration,
340 ) -> Result<()> {
341 let lease_until = compute_lease_until(lease_duration)?;
342 let conn = Arc::clone(&self.conn);
343 let partitions = Arc::clone(&self.partitions_relation);
344 let group = lease.scope.consumer_group_id.as_str().to_owned();
345 let stream = lease.scope.stream_id.as_str().to_owned();
346 let owner = lease.owner_id.as_str().to_owned();
347 let partition_id_i64 = lease.partition.id() as i64;
348 let partition_count_i64 = lease.partition.count() as i64;
349 let generation = lease.generation.get();
350 tokio::task::spawn_blocking(move || {
351 let guard = conn.lock().map_err(|e| Error::Store(e.to_string()))?;
352 let sql = format!(
353 "UPDATE {partitions} \
354 SET lease_until = ?5 \
355 WHERE consumer_group_id = ?1 \
356 AND stream_id = ?2 \
357 AND partition_id = ?3 \
358 AND owner_id = ?4 \
359 AND generation = ?6 \
360 AND partition_count = ?7"
361 );
362 let affected = guard
363 .execute(
364 &sql,
365 rusqlite::params![
366 group,
367 stream,
368 partition_id_i64,
369 owner,
370 lease_until,
371 generation,
372 partition_count_i64
373 ],
374 )
375 .map_err(|e| Error::Store(e.to_string()))?;
376 if affected == 0 {
377 check_partition_count_mismatch(
378 &guard,
379 &partitions,
380 &group,
381 &stream,
382 partition_id_i64,
383 partition_count_i64,
384 )?;
385 return Err(Error::OwnershipLost(format!(
386 "partition {partition_id_i64} generation {generation}"
387 )));
388 }
389 Ok(())
390 })
391 .await
392 .map_err(|e| Error::Store(format!("blocking task panicked: {e}")))?
393 }
394
395 async fn release<'a>(&'a self, lease: &'a PartitionLease<SqliteCursor>) -> Result<()> {
396 let conn = Arc::clone(&self.conn);
397 let partitions = Arc::clone(&self.partitions_relation);
398 let group = lease.scope.consumer_group_id.as_str().to_owned();
399 let stream = lease.scope.stream_id.as_str().to_owned();
400 let owner = lease.owner_id.as_str().to_owned();
401 let partition_id_i64 = lease.partition.id() as i64;
402 let partition_count_i64 = lease.partition.count() as i64;
403 let generation = lease.generation.get();
404 tokio::task::spawn_blocking(move || {
405 let guard = conn.lock().map_err(|e| Error::Store(e.to_string()))?;
406 let sql = format!(
407 "UPDATE {partitions} \
408 SET owner_id = NULL, \
409 lease_until = NULL, \
410 generation = generation + 1 \
411 WHERE consumer_group_id = ?1 \
412 AND stream_id = ?2 \
413 AND partition_id = ?3 \
414 AND owner_id = ?4 \
415 AND generation = ?5 \
416 AND partition_count = ?6"
417 );
418 let affected = guard
419 .execute(
420 &sql,
421 rusqlite::params![
422 group,
423 stream,
424 partition_id_i64,
425 owner,
426 generation,
427 partition_count_i64
428 ],
429 )
430 .map_err(|e| Error::Store(e.to_string()))?;
431 if affected == 0 {
432 check_partition_count_mismatch(
433 &guard,
434 &partitions,
435 &group,
436 &stream,
437 partition_id_i64,
438 partition_count_i64,
439 )?;
440 return Err(Error::OwnershipLost(format!(
441 "partition {partition_id_i64} generation {generation}"
442 )));
443 }
444 Ok(())
445 })
446 .await
447 .map_err(|e| Error::Store(format!("blocking task panicked: {e}")))?
448 }
449
450 async fn checkpoint<'a>(
451 &'a self,
452 lease: &'a PartitionLease<SqliteCursor>,
453 cursor: SqliteCursor,
454 ) -> Result<()> {
455 let conn = Arc::clone(&self.conn);
456 let partitions = Arc::clone(&self.partitions_relation);
457 let group = lease.scope.consumer_group_id.as_str().to_owned();
458 let stream = lease.scope.stream_id.as_str().to_owned();
459 let owner = lease.owner_id.as_str().to_owned();
460 let partition_id_i64 = lease.partition.id() as i64;
461 let partition_count_i64 = lease.partition.count() as i64;
462 let generation = lease.generation.get();
463 let sequence = cursor.sequence;
464 tokio::task::spawn_blocking(move || {
465 let guard = conn.lock().map_err(|e| Error::Store(e.to_string()))?;
466 let sql = format!(
467 "UPDATE {partitions} \
468 SET checkpoint_sequence = ?6 \
469 WHERE consumer_group_id = ?1 \
470 AND stream_id = ?2 \
471 AND partition_id = ?3 \
472 AND owner_id = ?4 \
473 AND generation = ?5 \
474 AND partition_count = ?7 \
475 AND ?6 > {partitions}.checkpoint_sequence"
476 );
477 let affected = guard
478 .execute(
479 &sql,
480 rusqlite::params![
481 group,
482 stream,
483 partition_id_i64,
484 owner,
485 generation,
486 sequence,
487 partition_count_i64
488 ],
489 )
490 .map_err(|e| Error::Store(e.to_string()))?;
491 if affected == 0 {
492 let check_sql = format!(
493 "SELECT owner_id, generation, partition_count FROM {partitions} \
494 WHERE consumer_group_id = ?1 AND stream_id = ?2 AND partition_id = ?3"
495 );
496 let check_row = guard
497 .query_row(
498 &check_sql,
499 rusqlite::params![group, stream, partition_id_i64],
500 |r| {
501 Ok((
502 r.get::<_, Option<String>>(0)?,
503 r.get::<_, i64>(1)?,
504 r.get::<_, Option<i64>>(2)?,
505 ))
506 },
507 )
508 .map(Some)
509 .or_else(|e| match e {
510 rusqlite::Error::QueryReturnedNoRows => Ok(None),
511 other => Err(other),
512 })
513 .map_err(|e| Error::Store(e.to_string()))?;
514 match check_row {
515 Some((current_owner, current_generation, current_count)) => {
516 if let Some(stored) = current_count
517 && stored != partition_count_i64
518 {
519 return Err(Error::Config(format!(
520 "partition count mismatch for scope {group} stream {stream} partition {partition_id_i64}: stored {stored}, requested {partition_count_i64}"
521 )));
522 }
523 if current_generation == generation
524 && current_owner.as_deref() == Some(owner.as_str())
525 {
526 return Ok(());
527 }
528 Err(Error::OwnershipLost(format!(
529 "checkpoint rejected for partition {partition_id_i64}: stale owner/generation"
530 )))
531 }
532 None => Err(Error::OwnershipLost(format!(
533 "partition {partition_id_i64} generation {generation}"
534 ))),
535 }
536 } else {
537 Ok(())
538 }
539 })
540 .await
541 .map_err(|e| Error::Store(format!("blocking task panicked: {e}")))?
542 }
543}
544
545fn check_partition_count_mismatch(
546 guard: &rusqlite::Connection,
547 partitions: &str,
548 group: &str,
549 stream: &str,
550 partition_id: i64,
551 requested_count: i64,
552) -> Result<()> {
553 let sql = format!(
554 "SELECT partition_count FROM {partitions} \
555 WHERE consumer_group_id = ?1 AND stream_id = ?2 AND partition_id = ?3"
556 );
557 let stored: Option<i64> = guard
558 .query_row(&sql, rusqlite::params![group, stream, partition_id], |r| {
559 r.get::<_, Option<i64>>(0)
560 })
561 .map(Some)
562 .or_else(|e| match e {
563 rusqlite::Error::QueryReturnedNoRows => Ok(None),
564 other => Err(other),
565 })
566 .map_err(|e| Error::Store(e.to_string()))?
567 .flatten();
568 if let Some(stored) = stored
569 && stored != requested_count
570 {
571 return Err(Error::Config(format!(
572 "partition count mismatch for scope {group} stream {stream} partition {partition_id}: stored {stored}, requested {requested_count}"
573 )));
574 }
575 Ok(())
576}
577
578#[cfg(test)]
579mod tests {
580 use std::num::NonZeroU32;
581
582 use super::*;
583 use eventuary_core::io::reader::PartitionCoordinator;
584 use eventuary_core::io::{ConsumerGroupId, OwnerId, StreamId};
585
586 use crate::database::SqliteDatabase;
587
588 #[test]
589 fn schema_sql_contains_expected_tables() {
590 let sql =
591 SqlitePartitionCoordinator::schema_sql(&SqlitePartitionCoordinatorConfig::default());
592 assert!(sql.contains("CREATE TABLE IF NOT EXISTS \"event_stream_consumers\""));
593 assert!(sql.contains("CREATE TABLE IF NOT EXISTS \"event_stream_partitions\""));
594 }
595
596 fn make_coordinator() -> SqlitePartitionCoordinator {
597 let db = SqliteDatabase::open_in_memory().unwrap();
598 let conn = db.conn();
599 SqlitePartitionCoordinator::prepare_schema(
600 &conn,
601 &SqlitePartitionCoordinatorConfig::default(),
602 )
603 .unwrap();
604 SqlitePartitionCoordinator::new(conn, SqlitePartitionCoordinatorConfig::default())
605 }
606
607 fn scope() -> CheckpointScope {
608 CheckpointScope::new(
609 ConsumerGroupId::new("group-1").unwrap(),
610 StreamId::new("orders").unwrap(),
611 )
612 }
613
614 fn partition(id: u32) -> Partition {
615 Partition::new(id, NonZeroU32::new(64).unwrap()).unwrap()
616 }
617
618 #[tokio::test]
619 async fn heartbeat_and_live_count() {
620 let coord = make_coordinator();
621 let s = scope();
622 let owner_a = OwnerId::new("worker-a").unwrap();
623 let owner_b = OwnerId::new("worker-b").unwrap();
624 let long_lease = std::time::Duration::from_secs(60);
625
626 coord.heartbeat(&s, &owner_a, long_lease).await.unwrap();
627 coord.heartbeat(&s, &owner_b, long_lease).await.unwrap();
628
629 let live = coord.live_consumers(&s).await.unwrap();
630 assert_eq!(live, 2);
631 }
632
633 #[tokio::test]
634 async fn live_consumers_excludes_expired() {
635 let coord = make_coordinator();
636 let s = scope();
637 let owner_a = OwnerId::new("worker-a").unwrap();
638 let owner_b = OwnerId::new("worker-b").unwrap();
639 let long_lease = std::time::Duration::from_secs(60);
640 let short_lease = std::time::Duration::from_millis(50);
641
642 coord.heartbeat(&s, &owner_a, long_lease).await.unwrap();
643 coord.heartbeat(&s, &owner_b, short_lease).await.unwrap();
644
645 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
646
647 let live = coord.live_consumers(&s).await.unwrap();
648 assert_eq!(live, 1);
649 }
650
651 #[tokio::test]
652 async fn claim_free_partition_succeeds() {
653 let coord = make_coordinator();
654 let s = scope();
655 let owner = OwnerId::new("worker-a").unwrap();
656 let lease_dur = std::time::Duration::from_secs(60);
657
658 let lease = coord
659 .claim(&s, &owner, partition(0), lease_dur)
660 .await
661 .unwrap()
662 .expect("should get lease");
663
664 assert_eq!(lease.generation.get(), 1);
665 assert!(lease.checkpoint_cursor.is_none());
666 }
667
668 #[tokio::test]
669 async fn claim_contested_returns_none() {
670 let coord = make_coordinator();
671 let s = scope();
672 let owner_a = OwnerId::new("worker-a").unwrap();
673 let owner_b = OwnerId::new("worker-b").unwrap();
674 let lease_dur = std::time::Duration::from_secs(60);
675
676 let first = coord
677 .claim(&s, &owner_a, partition(0), lease_dur)
678 .await
679 .unwrap();
680 assert!(first.is_some());
681
682 let second = coord
683 .claim(&s, &owner_b, partition(0), lease_dur)
684 .await
685 .unwrap();
686 assert!(second.is_none());
687 }
688
689 #[tokio::test]
690 async fn claim_after_expiry_succeeds() {
691 let coord = make_coordinator();
692 let s = scope();
693 let owner_a = OwnerId::new("worker-a").unwrap();
694 let owner_b = OwnerId::new("worker-b").unwrap();
695 let short_lease = std::time::Duration::from_millis(50);
696 let long_lease = std::time::Duration::from_secs(60);
697
698 let first = coord
699 .claim(&s, &owner_a, partition(0), short_lease)
700 .await
701 .unwrap();
702 assert_eq!(first.unwrap().generation.get(), 1);
703
704 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
705
706 let second = coord
707 .claim(&s, &owner_b, partition(0), long_lease)
708 .await
709 .unwrap()
710 .expect("should succeed after expiry");
711 assert_eq!(second.generation.get(), 2);
712 }
713
714 #[tokio::test]
715 async fn renew_with_matching_generation() {
716 let coord = make_coordinator();
717 let s = scope();
718 let owner = OwnerId::new("worker-a").unwrap();
719 let lease_dur = std::time::Duration::from_secs(60);
720
721 let lease = coord
722 .claim(&s, &owner, partition(0), lease_dur)
723 .await
724 .unwrap()
725 .unwrap();
726
727 coord.renew(&lease, lease_dur).await.unwrap();
728 }
729
730 #[tokio::test]
731 async fn renew_with_stale_generation_returns_ownership_lost() {
732 let coord = make_coordinator();
733 let s = scope();
734 let owner_a = OwnerId::new("worker-a").unwrap();
735 let owner_b = OwnerId::new("worker-b").unwrap();
736 let short_lease = std::time::Duration::from_millis(50);
737 let long_lease = std::time::Duration::from_secs(60);
738
739 let lease_a = coord
740 .claim(&s, &owner_a, partition(0), short_lease)
741 .await
742 .unwrap()
743 .unwrap();
744 assert_eq!(lease_a.generation.get(), 1);
745
746 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
747
748 coord
749 .claim(&s, &owner_b, partition(0), long_lease)
750 .await
751 .unwrap()
752 .unwrap();
753
754 let err = coord.renew(&lease_a, long_lease).await.unwrap_err();
755 assert!(matches!(err, eventuary_core::Error::OwnershipLost(_)));
756 }
757
758 #[tokio::test]
759 async fn release_with_matching_generation() {
760 let coord = make_coordinator();
761 let s = scope();
762 let owner_a = OwnerId::new("worker-a").unwrap();
763 let owner_b = OwnerId::new("worker-b").unwrap();
764 let lease_dur = std::time::Duration::from_secs(60);
765
766 let lease = coord
767 .claim(&s, &owner_a, partition(0), lease_dur)
768 .await
769 .unwrap()
770 .unwrap();
771 assert_eq!(lease.generation.get(), 1);
772
773 coord.release(&lease).await.unwrap();
774
775 let new_lease = coord
776 .claim(&s, &owner_b, partition(0), lease_dur)
777 .await
778 .unwrap()
779 .expect("should claim after release");
780 assert_eq!(new_lease.generation.get(), 3);
781 }
782
783 #[tokio::test]
784 async fn release_with_stale_generation_returns_ownership_lost() {
785 let coord = make_coordinator();
786 let s = scope();
787 let owner_a = OwnerId::new("worker-a").unwrap();
788 let owner_b = OwnerId::new("worker-b").unwrap();
789 let short_lease = std::time::Duration::from_millis(50);
790 let long_lease = std::time::Duration::from_secs(60);
791
792 let lease_a = coord
793 .claim(&s, &owner_a, partition(0), short_lease)
794 .await
795 .unwrap()
796 .unwrap();
797
798 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
799
800 coord
801 .claim(&s, &owner_b, partition(0), long_lease)
802 .await
803 .unwrap()
804 .unwrap();
805
806 let err = coord.release(&lease_a).await.unwrap_err();
807 assert!(matches!(err, eventuary_core::Error::OwnershipLost(_)));
808 }
809
810 #[tokio::test]
811 async fn checkpoint_with_matching_generation_advances() {
812 let coord = make_coordinator();
813 let s = scope();
814 let owner_a = OwnerId::new("worker-a").unwrap();
815 let owner_b = OwnerId::new("worker-b").unwrap();
816 let long_lease = std::time::Duration::from_secs(60);
817
818 let lease = coord
819 .claim(&s, &owner_a, partition(0), long_lease)
820 .await
821 .unwrap()
822 .unwrap();
823
824 coord
825 .checkpoint(&lease, SqliteCursor::new(100, partition(0)))
826 .await
827 .unwrap();
828
829 coord.release(&lease).await.unwrap();
830
831 let new_lease = coord
832 .claim(&s, &owner_b, partition(0), long_lease)
833 .await
834 .unwrap()
835 .expect("should claim after release");
836 assert_eq!(new_lease.checkpoint_cursor.unwrap().sequence, 100);
837 }
838
839 #[tokio::test]
840 async fn checkpoint_is_monotonic() {
841 let coord = make_coordinator();
842 let s = scope();
843 let owner_a = OwnerId::new("worker-a").unwrap();
844 let owner_b = OwnerId::new("worker-b").unwrap();
845 let long_lease = std::time::Duration::from_secs(60);
846
847 let lease = coord
848 .claim(&s, &owner_a, partition(0), long_lease)
849 .await
850 .unwrap()
851 .unwrap();
852
853 coord
854 .checkpoint(&lease, SqliteCursor::new(100, partition(0)))
855 .await
856 .unwrap();
857 coord
858 .checkpoint(&lease, SqliteCursor::new(50, partition(0)))
859 .await
860 .unwrap();
861
862 coord.release(&lease).await.unwrap();
863
864 let new_lease = coord
865 .claim(&s, &owner_b, partition(0), long_lease)
866 .await
867 .unwrap()
868 .expect("should claim after release");
869 assert_eq!(new_lease.checkpoint_cursor.unwrap().sequence, 100);
870 }
871
872 #[tokio::test]
873 async fn checkpoint_with_stale_generation_returns_ownership_lost() {
874 let coord = make_coordinator();
875 let s = scope();
876 let owner_a = OwnerId::new("worker-a").unwrap();
877 let owner_b = OwnerId::new("worker-b").unwrap();
878 let short_lease = std::time::Duration::from_millis(50);
879 let long_lease = std::time::Duration::from_secs(60);
880
881 let lease_a = coord
882 .claim(&s, &owner_a, partition(0), short_lease)
883 .await
884 .unwrap()
885 .unwrap();
886
887 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
888
889 coord
890 .claim(&s, &owner_b, partition(0), long_lease)
891 .await
892 .unwrap()
893 .unwrap();
894
895 let err = coord
896 .checkpoint(&lease_a, SqliteCursor::new(100, partition(0)))
897 .await
898 .unwrap_err();
899 assert!(matches!(err, eventuary_core::Error::OwnershipLost(_)));
900 }
901
902 #[tokio::test]
903 async fn claim_rejects_partition_count_mismatch() {
904 let coord = make_coordinator();
905 let s = scope();
906 let owner_a = OwnerId::new("worker-a").unwrap();
907 let owner_b = OwnerId::new("worker-b").unwrap();
908 let p_four = Partition::new(0, NonZeroU32::new(4).unwrap()).unwrap();
909 let p_eight = Partition::new(0, NonZeroU32::new(8).unwrap()).unwrap();
910
911 coord
912 .claim(&s, &owner_a, p_four, std::time::Duration::from_millis(1))
913 .await
914 .unwrap()
915 .unwrap();
916 tokio::time::sleep(std::time::Duration::from_millis(5)).await;
917
918 let err = coord
919 .claim(&s, &owner_b, p_eight, std::time::Duration::from_secs(10))
920 .await
921 .unwrap_err();
922 assert!(
923 matches!(err, eventuary_core::Error::Config(ref message) if message.contains("partition count mismatch")),
924 "expected partition count mismatch error, got {err:?}"
925 );
926 }
927
928 #[tokio::test]
929 async fn checkpoint_after_release_returns_ownership_lost() {
930 let coord = make_coordinator();
931 let s = scope();
932 let owner_a = OwnerId::new("worker-a").unwrap();
933 let lease_dur = std::time::Duration::from_secs(60);
934
935 let lease = coord
936 .claim(&s, &owner_a, partition(0), lease_dur)
937 .await
938 .unwrap()
939 .unwrap();
940 assert_eq!(lease.generation.get(), 1);
941
942 coord.release(&lease).await.unwrap();
943
944 let err = coord
945 .checkpoint(&lease, SqliteCursor::new(100, partition(0)))
946 .await
947 .unwrap_err();
948 assert!(matches!(err, eventuary_core::Error::OwnershipLost(_)));
949 }
950}