1use std::sync::Arc;
2
3use chrono::{DateTime, Utc};
4use sqlx::{PgPool, Row};
5
6use eventuary_core::io::OwnerId;
7use eventuary_core::io::reader::{
8 CheckpointScope, Generation, PartitionCoordinator, PartitionLease,
9};
10use eventuary_core::{Error, Partition, Result};
11
12use crate::reader::PgCursor;
13use crate::relation::PgRelationName;
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 TIMESTAMPTZ 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 BIGINT NOT NULL,
32 partition_count BIGINT NULL,
33 owner_id TEXT NULL,
34 lease_until TIMESTAMPTZ NULL,
35 checkpoint_sequence BIGINT NOT NULL DEFAULT 0,
36 generation BIGINT 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 PgPartitionCoordinatorConfig {
54 pub consumers_relation: PgRelationName,
55 pub partitions_relation: PgRelationName,
56}
57
58impl Default for PgPartitionCoordinatorConfig {
59 fn default() -> Self {
60 Self {
61 consumers_relation: PgRelationName::new("event_stream_consumers")
62 .expect("default consumers relation"),
63 partitions_relation: PgRelationName::new("event_stream_partitions")
64 .expect("default partitions relation"),
65 }
66 }
67}
68
69pub struct PgPartitionCoordinator {
70 pool: PgPool,
71 consumers_relation: Arc<String>,
72 partitions_relation: Arc<String>,
73}
74
75impl Clone for PgPartitionCoordinator {
76 fn clone(&self) -> Self {
77 Self {
78 pool: self.pool.clone(),
79 consumers_relation: Arc::clone(&self.consumers_relation),
80 partitions_relation: Arc::clone(&self.partitions_relation),
81 }
82 }
83}
84
85impl PgPartitionCoordinator {
86 pub fn new(pool: PgPool, config: PgPartitionCoordinatorConfig) -> Self {
87 Self {
88 pool,
89 consumers_relation: Arc::new(config.consumers_relation.render()),
90 partitions_relation: Arc::new(config.partitions_relation.render()),
91 }
92 }
93
94 pub async fn connect(pool: PgPool, config: PgPartitionCoordinatorConfig) -> Result<Self> {
95 Self::prepare_schema(&pool, &config).await?;
96 Ok(Self::new(pool, config))
97 }
98
99 pub async fn prepare_schema(
100 pool: &PgPool,
101 config: &PgPartitionCoordinatorConfig,
102 ) -> Result<()> {
103 crate::schema::apply_schema(
104 pool,
105 PARTITION_COORDINATOR_MIGRATIONS,
106 &[
107 RelationReplacement {
108 token: "{consumers}",
109 relation: &config.consumers_relation,
110 },
111 RelationReplacement {
112 token: "{partitions}",
113 relation: &config.partitions_relation,
114 },
115 ],
116 )
117 .await
118 }
119
120 pub fn schema_sql(config: &PgPartitionCoordinatorConfig) -> 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<DateTime<Utc>> {
138 Ok(Utc::now()
139 + chrono::Duration::from_std(lease_duration)
140 .map_err(|_| Error::Config("lease duration out of range".to_owned()))?)
141}
142
143fn lease_until_to_sql(dt: DateTime<Utc>) -> String {
144 dt.format("%Y-%m-%dT%H:%M:%S%.6fZ").to_string()
145}
146
147fn parse_lease_until(s: &str) -> Result<DateTime<Utc>> {
148 DateTime::parse_from_rfc3339(s)
149 .map(|dt| dt.with_timezone(&Utc))
150 .or_else(|_| {
151 DateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S%.f%#z").map(|dt| dt.with_timezone(&Utc))
152 })
153 .map_err(|e| Error::Serialization(format!("lease_until decode: {e}")))
154}
155
156impl PartitionCoordinator<PgCursor> for PgPartitionCoordinator {
157 async fn heartbeat<'a>(
158 &'a self,
159 scope: &'a CheckpointScope,
160 owner_id: &'a OwnerId,
161 lease_duration: std::time::Duration,
162 ) -> Result<()> {
163 let lease_until = lease_until_to_sql(compute_lease_until(lease_duration)?);
164 let sql = format!(
165 "INSERT INTO {consumers} (consumer_group_id, stream_id, owner_id, lease_until) \
166 VALUES ($1, $2, $3, $4::timestamptz) \
167 ON CONFLICT (consumer_group_id, stream_id, owner_id) DO UPDATE \
168 SET lease_until = EXCLUDED.lease_until",
169 consumers = self.consumers_relation
170 );
171 sqlx::query(&sql)
172 .bind(scope.consumer_group_id.as_str())
173 .bind(scope.stream_id.as_str())
174 .bind(owner_id.as_str())
175 .bind(lease_until)
176 .execute(&self.pool)
177 .await
178 .map_err(|e| Error::Store(e.to_string()))?;
179 Ok(())
180 }
181
182 async fn live_consumers<'a>(&'a self, scope: &'a CheckpointScope) -> Result<usize> {
183 let sql = format!(
184 "SELECT COUNT(*) FROM {consumers} \
185 WHERE consumer_group_id = $1 \
186 AND stream_id = $2 \
187 AND lease_until > NOW()",
188 consumers = self.consumers_relation
189 );
190 let row = sqlx::query(&sql)
191 .bind(scope.consumer_group_id.as_str())
192 .bind(scope.stream_id.as_str())
193 .fetch_one(&self.pool)
194 .await
195 .map_err(|e| Error::Store(e.to_string()))?;
196 let count: i64 = row.get(0);
197 Ok(count as usize)
198 }
199
200 async fn release_consumer<'a>(
201 &'a self,
202 scope: &'a CheckpointScope,
203 owner_id: &'a OwnerId,
204 ) -> Result<()> {
205 let sql = format!(
206 "DELETE FROM {consumers} \
207 WHERE consumer_group_id = $1 AND stream_id = $2 AND owner_id = $3",
208 consumers = self.consumers_relation
209 );
210 sqlx::query(&sql)
211 .bind(scope.consumer_group_id.as_str())
212 .bind(scope.stream_id.as_str())
213 .bind(owner_id.as_str())
214 .execute(&self.pool)
215 .await
216 .map_err(|e| Error::Store(e.to_string()))?;
217 Ok(())
218 }
219
220 async fn claim<'a>(
221 &'a self,
222 scope: &'a CheckpointScope,
223 owner_id: &'a OwnerId,
224 partition: Partition,
225 lease_duration: std::time::Duration,
226 ) -> Result<Option<PartitionLease<PgCursor>>> {
227 let lease_until = lease_until_to_sql(compute_lease_until(lease_duration)?);
228 let partition_id_i64 = partition.id() as i64;
229 let partition_count_i64 = partition.count() as i64;
230 let sql = format!(
231 "INSERT INTO {partitions} \
232 (consumer_group_id, stream_id, partition_id, partition_count, owner_id, lease_until, generation, checkpoint_sequence) \
233 VALUES ($1, $2, $3, $4, $5, $6::timestamptz, 1, 0) \
234 ON CONFLICT (consumer_group_id, stream_id, partition_id) DO UPDATE \
235 SET owner_id = EXCLUDED.owner_id, \
236 lease_until = EXCLUDED.lease_until, \
237 partition_count = COALESCE({partitions}.partition_count, EXCLUDED.partition_count), \
238 generation = {partitions}.generation + 1 \
239 WHERE ({partitions}.partition_count IS NULL OR {partitions}.partition_count = EXCLUDED.partition_count) \
240 AND ({partitions}.owner_id IS NULL \
241 OR {partitions}.lease_until IS NULL \
242 OR {partitions}.lease_until < NOW() \
243 OR {partitions}.owner_id = EXCLUDED.owner_id) \
244 RETURNING owner_id, \
245 to_char(lease_until AT TIME ZONE 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS.US\"Z\"') AS lease_until_text, \
246 generation, \
247 checkpoint_sequence",
248 partitions = self.partitions_relation
249 );
250 let row = sqlx::query(&sql)
251 .bind(scope.consumer_group_id.as_str())
252 .bind(scope.stream_id.as_str())
253 .bind(partition_id_i64)
254 .bind(partition_count_i64)
255 .bind(owner_id.as_str())
256 .bind(lease_until)
257 .fetch_optional(&self.pool)
258 .await
259 .map_err(|e| Error::Store(e.to_string()))?;
260 match row {
261 None => {
262 let check_sql = format!(
263 "SELECT partition_count FROM {partitions} \
264 WHERE consumer_group_id = $1 AND stream_id = $2 AND partition_id = $3",
265 partitions = self.partitions_relation
266 );
267 let check_row = sqlx::query(&check_sql)
268 .bind(scope.consumer_group_id.as_str())
269 .bind(scope.stream_id.as_str())
270 .bind(partition_id_i64)
271 .fetch_optional(&self.pool)
272 .await
273 .map_err(|e| Error::Store(e.to_string()))?;
274 if let Some(r) = check_row {
275 let stored: Option<i64> = r.get("partition_count");
276 if let Some(stored) = stored
277 && stored != partition_count_i64
278 {
279 return Err(Error::Config(format!(
280 "partition count mismatch for scope {} stream {} partition {}: stored {}, requested {}",
281 scope.consumer_group_id.as_str(),
282 scope.stream_id.as_str(),
283 partition.id(),
284 stored,
285 partition_count_i64,
286 )));
287 }
288 }
289 Ok(None)
290 }
291 Some(r) => {
292 let lease_until_text: String = r.get("lease_until_text");
293 let returned_lease_until = parse_lease_until(&lease_until_text)?;
294 let generation: i64 = r.get("generation");
295 let checkpoint_sequence: i64 = r.get("checkpoint_sequence");
296 Ok(Some(PartitionLease {
297 scope: scope.clone(),
298 owner_id: owner_id.clone(),
299 partition,
300 generation: Generation::from_i64(generation),
301 checkpoint_cursor: (checkpoint_sequence > 0)
302 .then_some(PgCursor::new(checkpoint_sequence, partition)),
303 lease_until: returned_lease_until,
304 }))
305 }
306 }
307 }
308
309 async fn renew<'a>(
310 &'a self,
311 lease: &'a PartitionLease<PgCursor>,
312 lease_duration: std::time::Duration,
313 ) -> Result<()> {
314 let lease_until = lease_until_to_sql(compute_lease_until(lease_duration)?);
315 let partition_id_i64 = lease.partition.id() as i64;
316 let partition_count_i64 = lease.partition.count() as i64;
317 let generation = lease.generation.get();
318 let sql = format!(
319 "UPDATE {partitions} \
320 SET lease_until = $5::timestamptz \
321 WHERE consumer_group_id = $1 \
322 AND stream_id = $2 \
323 AND partition_id = $3 \
324 AND owner_id = $4 \
325 AND generation = $6 \
326 AND partition_count = $7",
327 partitions = self.partitions_relation
328 );
329 let result = sqlx::query(&sql)
330 .bind(lease.scope.consumer_group_id.as_str())
331 .bind(lease.scope.stream_id.as_str())
332 .bind(partition_id_i64)
333 .bind(lease.owner_id.as_str())
334 .bind(lease_until)
335 .bind(generation)
336 .bind(partition_count_i64)
337 .execute(&self.pool)
338 .await
339 .map_err(|e| Error::Store(e.to_string()))?;
340 if result.rows_affected() == 0 {
341 self.check_partition_count_mismatch(lease).await?;
342 return Err(Error::OwnershipLost(format!(
343 "partition {} generation {}",
344 lease.partition.id(),
345 lease.generation,
346 )));
347 }
348 Ok(())
349 }
350
351 async fn release<'a>(&'a self, lease: &'a PartitionLease<PgCursor>) -> Result<()> {
352 let partition_id_i64 = lease.partition.id() as i64;
353 let partition_count_i64 = lease.partition.count() as i64;
354 let generation = lease.generation.get();
355 let sql = format!(
356 "UPDATE {partitions} \
357 SET owner_id = NULL, \
358 lease_until = NULL, \
359 generation = generation + 1 \
360 WHERE consumer_group_id = $1 \
361 AND stream_id = $2 \
362 AND partition_id = $3 \
363 AND owner_id = $4 \
364 AND generation = $5 \
365 AND partition_count = $6",
366 partitions = self.partitions_relation
367 );
368 let result = sqlx::query(&sql)
369 .bind(lease.scope.consumer_group_id.as_str())
370 .bind(lease.scope.stream_id.as_str())
371 .bind(partition_id_i64)
372 .bind(lease.owner_id.as_str())
373 .bind(generation)
374 .bind(partition_count_i64)
375 .execute(&self.pool)
376 .await
377 .map_err(|e| Error::Store(e.to_string()))?;
378 if result.rows_affected() == 0 {
379 self.check_partition_count_mismatch(lease).await?;
380 return Err(Error::OwnershipLost(format!(
381 "partition {} generation {}",
382 lease.partition.id(),
383 lease.generation,
384 )));
385 }
386 Ok(())
387 }
388
389 async fn checkpoint<'a>(
390 &'a self,
391 lease: &'a PartitionLease<PgCursor>,
392 cursor: PgCursor,
393 ) -> Result<()> {
394 let partition_id_i64 = lease.partition.id() as i64;
395 let partition_count_i64 = lease.partition.count() as i64;
396 let generation = lease.generation.get();
397 let sequence = cursor.sequence;
398 let sql = format!(
399 "UPDATE {partitions} \
400 SET checkpoint_sequence = $6 \
401 WHERE consumer_group_id = $1 \
402 AND stream_id = $2 \
403 AND partition_id = $3 \
404 AND owner_id = $4 \
405 AND generation = $5 \
406 AND partition_count = $7 \
407 AND $6 > {partitions}.checkpoint_sequence",
408 partitions = self.partitions_relation
409 );
410 let result = sqlx::query(&sql)
411 .bind(lease.scope.consumer_group_id.as_str())
412 .bind(lease.scope.stream_id.as_str())
413 .bind(partition_id_i64)
414 .bind(lease.owner_id.as_str())
415 .bind(generation)
416 .bind(sequence)
417 .bind(partition_count_i64)
418 .execute(&self.pool)
419 .await
420 .map_err(|e| Error::Store(e.to_string()))?;
421 if result.rows_affected() == 0 {
422 let check_sql = format!(
423 "SELECT generation, owner_id, partition_count FROM {partitions} \
424 WHERE consumer_group_id = $1 AND stream_id = $2 AND partition_id = $3",
425 partitions = self.partitions_relation
426 );
427 let check_row = sqlx::query(&check_sql)
428 .bind(lease.scope.consumer_group_id.as_str())
429 .bind(lease.scope.stream_id.as_str())
430 .bind(partition_id_i64)
431 .fetch_optional(&self.pool)
432 .await
433 .map_err(|e| Error::Store(e.to_string()))?;
434 match check_row {
435 Some(r) => {
436 let current_generation: i64 = r.get("generation");
437 let current_owner: Option<String> = r.get("owner_id");
438 let current_count: Option<i64> = r.get("partition_count");
439 if let Some(stored) = current_count
440 && stored != partition_count_i64
441 {
442 return Err(Error::Config(format!(
443 "partition count mismatch for scope {} stream {} partition {}: stored {}, requested {}",
444 lease.scope.consumer_group_id.as_str(),
445 lease.scope.stream_id.as_str(),
446 lease.partition.id(),
447 stored,
448 partition_count_i64,
449 )));
450 }
451 if current_generation == generation
452 && current_owner.as_deref() == Some(lease.owner_id.as_str())
453 {
454 return Ok(());
455 }
456 Err(Error::OwnershipLost(format!(
457 "checkpoint rejected for partition {}: stale owner/generation",
458 lease.partition.id(),
459 )))
460 }
461 None => Err(Error::OwnershipLost(format!(
462 "partition {} generation {}",
463 lease.partition.id(),
464 lease.generation,
465 ))),
466 }
467 } else {
468 Ok(())
469 }
470 }
471}
472
473impl PgPartitionCoordinator {
474 async fn check_partition_count_mismatch(&self, lease: &PartitionLease<PgCursor>) -> Result<()> {
475 let partition_id_i64 = lease.partition.id() as i64;
476 let partition_count_i64 = lease.partition.count() as i64;
477 let sql = format!(
478 "SELECT partition_count FROM {partitions} \
479 WHERE consumer_group_id = $1 AND stream_id = $2 AND partition_id = $3",
480 partitions = self.partitions_relation
481 );
482 let row = sqlx::query(&sql)
483 .bind(lease.scope.consumer_group_id.as_str())
484 .bind(lease.scope.stream_id.as_str())
485 .bind(partition_id_i64)
486 .fetch_optional(&self.pool)
487 .await
488 .map_err(|e| Error::Store(e.to_string()))?;
489 if let Some(r) = row {
490 let stored: Option<i64> = r.get("partition_count");
491 if let Some(stored) = stored
492 && stored != partition_count_i64
493 {
494 return Err(Error::Config(format!(
495 "partition count mismatch for scope {} stream {} partition {}: stored {}, requested {}",
496 lease.scope.consumer_group_id.as_str(),
497 lease.scope.stream_id.as_str(),
498 lease.partition.id(),
499 stored,
500 partition_count_i64,
501 )));
502 }
503 }
504 Ok(())
505 }
506}
507
508#[cfg(test)]
509mod schema_tests {
510 use super::*;
511
512 #[test]
513 fn schema_sql_contains_expected_tables() {
514 let sql = PgPartitionCoordinator::schema_sql(&PgPartitionCoordinatorConfig::default());
515 assert!(sql.contains("CREATE TABLE IF NOT EXISTS \"event_stream_consumers\""));
516 assert!(sql.contains("CREATE TABLE IF NOT EXISTS \"event_stream_partitions\""));
517 }
518}