1use super::*;
27
28#[async_trait::async_trait]
29impl SessionCommitStore for Store {
30 fn durability_tier(&self) -> DurabilityTier {
31 DurabilityTier::Durable
32 }
33
34 async fn load_session(
35 &self,
36 scope: SessionReadScope,
37 ) -> Result<Option<PersistedSessionRead>, StoreError> {
38 self.conn
39 .call(move |conn| {
40 let outcome: Result<Option<PersistedSessionRead>, StoreError> = (|| {
41 let Some(meta) = try_load_session_head_meta_from_conn(conn)? else {
42 return Ok(None);
43 };
44 let leaf_node_id = match &scope {
45 SessionReadScope::FullGraph => meta.leaf_node_id.clone(),
46 SessionReadScope::ActivePath { leaf_node_id } => {
47 leaf_node_id.clone().or_else(|| meta.leaf_node_id.clone())
48 }
49 };
50 let mut graph = match scope {
51 SessionReadScope::FullGraph => {
52 Self::load_session_graph_from_conn(conn, meta.leaf_node_id.clone())
53 }
54 SessionReadScope::ActivePath { .. } => {
55 Self::load_active_path_session_graph_from_conn(
56 conn,
57 leaf_node_id.clone(),
58 )
59 .map_err(sqlite_error)?
60 }
61 };
62 graph.set_leaf_node_id(leaf_node_id);
63 let checkpoint = meta
64 .checkpoint_ref
65 .as_ref()
66 .map(|blob_ref| Self::get_checkpoint_conn(conn, blob_ref))
67 .transpose()?
68 .flatten();
69 Ok(Some(PersistedSessionRead {
70 session_id: meta.session_id,
71 head_revision: meta.head_revision,
72 config: meta.config,
73 agent_frames: meta.agent_frames,
74 current_agent_frame_id: meta.current_agent_frame_id,
75 graph,
76 checkpoint_ref: meta.checkpoint_ref,
77 checkpoint,
78 token_ledger: merge_token_ledger_entries(Self::load_usage_deltas_conn(
79 conn,
80 )),
81 }))
82 })(
83 );
84 Ok(outcome)
85 })
86 .await
87 .map_err(sqlite_error)?
88 }
89
90 async fn load_node(
91 &self,
92 node_id: &str,
93 ) -> Result<Option<lash_core::SessionNodeRecord>, StoreError> {
94 let node_id = node_id.to_string();
95 let row: Option<String> = self
96 .conn
97 .call(move |conn| {
98 conn.query_row(
99 "SELECT node_json FROM graph_nodes WHERE node_id = ?1 AND tombstoned = 0",
100 params![node_id],
101 |row| row.get(0),
102 )
103 .optional()
104 })
105 .await
106 .map_err(sqlite_error)?;
107 Ok(row.and_then(|json| serde_json::from_str(&json).ok()))
108 }
109
110 async fn commit_runtime_state(
111 &self,
112 commit: RuntimeCommit,
113 ) -> Result<RuntimeCommitResult, StoreError> {
114 let blob_profile = self.options.blob_profile;
115 let now = self.clock.timestamp_ms();
116 let enqueue_nonce_start = self.commit_count.fetch_add(
117 commit.enqueued_queue_batches.len() as u64,
118 AtomicOrdering::Relaxed,
119 );
120 let result = self
121 .conn
122 .write_flow(move |tx| {
123 let outcome: Result<RuntimeCommitResult, StoreError> = (|| {
124 let existing = try_load_session_head_meta_from_conn(tx)?;
125 if let Some(bound_session_id) =
126 existing.as_ref().map(|meta| meta.session_id.as_str())
127 && bound_session_id != commit.session_id
128 {
129 return Err(StoreError::SessionBindingMismatch {
130 bound_session_id: bound_session_id.to_string(),
131 attempted_session_id: commit.session_id.clone(),
132 });
133 }
134 if let Some(completed) = &commit.turn_commit {
135 if completed.session_id != commit.session_id {
136 return Err(StoreError::RuntimeTurnCommitConflict {
137 session_id: completed.session_id.clone(),
138 turn_id: completed.turn_id.clone(),
139 });
140 }
141 let prior: Option<(String, String)> = tx
142 .query_row(
143 "SELECT turn_commit_hash, result_json FROM runtime_turn_commits
144 WHERE session_id = ?1 AND turn_id = ?2",
145 params![completed.session_id, completed.turn_id],
146 |row| Ok((row.get(0)?, row.get(1)?)),
147 )
148 .optional()
149 .map_err(sqlite_error)?;
150 if let Some((turn_commit_hash, result_json)) = prior {
151 if turn_commit_hash == completed.turn_commit_hash {
152 let result: RuntimeCommitResult =
153 serde_json::from_str(&result_json).map_err(|err| {
154 StoreError::Backend(format!(
155 "failed to decode runtime turn commit result: {err}"
156 ))
157 })?;
158 if let Some(completion) =
159 commit.release_session_execution_lease.as_ref()
160 {
161 release_session_execution_lease_conn(tx, completion)?;
162 }
163 return Ok(result);
164 }
165 return Err(StoreError::RuntimeTurnCommitConflict {
166 session_id: completed.session_id.clone(),
167 turn_id: completed.turn_id.clone(),
168 });
169 }
170 }
171 let Some(session_execution_lease) = commit.session_execution_lease.as_ref()
172 else {
173 return Err(StoreError::SessionExecutionLeaseExpired {
174 session_id: commit.session_id.clone(),
175 });
176 };
177 ensure_session_execution_lease_conn(
178 tx,
179 &commit.session_id,
180 session_execution_lease,
181 now,
182 )?;
183 let actual_revision = existing.as_ref().map_or(0, |meta| meta.head_revision);
184 if commit.expected_head_revision.is_some()
185 && commit.expected_head_revision != Some(actual_revision)
186 {
187 return Err(StoreError::HeadRevisionConflict {
188 expected: commit.expected_head_revision,
189 actual: actual_revision,
190 });
191 }
192 for completed in &commit.completed_queue_claims {
193 if completed.session_id != commit.session_id {
194 return Err(StoreError::QueuedWorkClaimSuperseded {
195 session_id: completed.session_id.clone(),
196 claim_id: completed.claim_id.clone(),
197 });
198 }
199 ensure_queued_work_completion_conn(tx, completed)?;
200 }
201 for completed in &commit.completed_turn_input_claims {
202 if completed.session_id != commit.session_id {
203 return Err(StoreError::TurnInputClaimSuperseded {
204 session_id: completed.session_id.clone(),
205 claim_id: completed.claim_id.clone(),
206 });
207 }
208 let owned_rows: usize = tx
209 .query_row(
210 "SELECT COUNT(*)
211 FROM pending_turn_inputs
212 WHERE session_id = ?1
213 AND claim_id = ?2
214 AND claim_token = ?3",
215 params![
216 completed.session_id,
217 completed.claim_id,
218 completed.lease_token
219 ],
220 |row| row.get::<_, i64>(0),
221 )
222 .map_err(sqlite_error)? as usize;
223 ensure_turn_input_completion_owns_all_inputs(completed, owned_rows)?;
224 }
225
226 let stored_checkpoint =
227 Self::put_checkpoint_conn(tx, &commit.checkpoint, blob_profile)
228 .map_err(sqlite_error)?;
229
230 if !commit.usage_deltas.is_empty() {
231 let mut stmt = tx
232 .prepare(
233 "INSERT INTO usage_deltas (
234 source, model, input_tokens, output_tokens, cache_read_input_tokens, cache_write_input_tokens, reasoning_output_tokens
235 ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
236 )
237 .map_err(sqlite_error)?;
238 for entry in &commit.usage_deltas {
239 stmt.execute(params![
240 entry.source,
241 entry.model,
242 entry.usage.input_tokens,
243 entry.usage.output_tokens,
244 entry.usage.cache_read_input_tokens,
245 entry.usage.cache_write_input_tokens,
246 entry.usage.reasoning_output_tokens,
247 ])
248 .map_err(sqlite_error)?;
249 }
250 }
251
252 let leaf_node_id = match &commit.graph {
253 GraphCommitDelta::Unchanged { leaf_node_id } => leaf_node_id.clone(),
254 GraphCommitDelta::Append {
255 nodes,
256 leaf_node_id,
257 } => {
258 for node in nodes {
259 let node_json = encode_json(node);
260 tx.execute(
261 "INSERT INTO graph_nodes (node_id, node_json) VALUES (?1, ?2)",
262 params![node.node_id, node_json],
263 )
264 .map_err(sqlite_error)?;
265 }
266 leaf_node_id.clone()
267 }
268 GraphCommitDelta::ReplaceFull(graph) => {
269 tx.execute("DELETE FROM graph_nodes", [])
270 .map_err(sqlite_error)?;
271 for node in &graph.nodes {
272 let node_json = encode_json(node);
273 tx.execute(
274 "INSERT INTO graph_nodes (node_id, node_json) VALUES (?1, ?2)",
275 params![node.node_id, node_json],
276 )
277 .map_err(sqlite_error)?;
278 }
279 graph.leaf_node_id.clone()
280 }
281 };
282 let graph_node_count: usize = tx
283 .query_row(
284 "SELECT COUNT(*) FROM graph_nodes WHERE tombstoned = 0",
285 [],
286 |row| row.get::<_, i64>(0),
287 )
288 .map_err(sqlite_error)? as usize;
289 let next_revision = actual_revision + 1;
290 let meta = SessionHeadMeta {
291 schema_version: lash_core::store::SESSION_HEAD_META_SCHEMA_VERSION,
292 session_id: commit.session_id.clone(),
293 head_revision: next_revision,
294 config: commit.config.clone(),
295 agent_frames: commit.agent_frames.clone(),
296 current_agent_frame_id: commit.current_agent_frame_id.clone(),
297 checkpoint_ref: Some(stored_checkpoint.checkpoint_ref.clone()),
298 leaf_node_id,
299 graph_node_count,
300 token_ledger: Vec::new(),
301 };
302 tx.execute(
303 "INSERT OR REPLACE INTO session_head (singleton, session_id, head_json, head_revision)
304 VALUES (1, ?1, ?2, ?3)",
305 params![
306 meta.session_id,
307 encode_json(&meta),
308 meta.head_revision as i64
309 ],
310 )
311 .map_err(sqlite_error)?;
312 for completed in &commit.completed_queue_claims {
313 for batch_id in &completed.batch_ids {
314 tx.execute(
315 "DELETE FROM queued_work_batches
316 WHERE session_id = ?1
317 AND batch_id = ?2
318 AND claim_id = ?3
319 AND claim_token = ?4",
320 params![
321 completed.session_id,
322 batch_id,
323 completed.claim_id,
324 completed.lease_token
325 ],
326 )
327 .map_err(sqlite_error)?;
328 }
329 }
330 for completed in &commit.completed_turn_input_claims {
331 for input_id in &completed.input_ids {
332 tx.execute(
333 "UPDATE pending_turn_inputs
334 SET state = ?5,
335 claim_id = NULL,
336 claim_owner_id = NULL,
337 claim_owner_incarnation_id = NULL,
338 claim_owner_liveness_json = NULL,
339 claim_token = NULL,
340 claim_session_lease_generation = 0
341 WHERE session_id = ?1
342 AND input_id = ?2
343 AND claim_id = ?3
344 AND claim_token = ?4",
345 params![
346 completed.session_id,
347 input_id,
348 completed.claim_id,
349 completed.lease_token,
350 lash_core::TurnInputState::Completed.as_str(),
351 ],
352 )
353 .map_err(sqlite_error)?;
354 }
355 }
356 if let Some(turn_id) = commit.interrupted_turn_input_turn_id.as_deref() {
357 let input_ids = {
358 let mut stmt = tx
359 .prepare(
360 "SELECT input_id, ingress_json
361 FROM pending_turn_inputs
362 WHERE session_id = ?1 AND state = ?2",
363 )
364 .map_err(sqlite_error)?;
365 let rows = stmt
366 .query_map(
367 params![
368 commit.session_id,
369 lash_core::TurnInputState::PendingActive.as_str()
370 ],
371 |row| {
372 Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?))
373 },
374 )
375 .map_err(sqlite_error)?;
376 let mut input_ids = Vec::new();
377 for row in rows {
378 let (input_id, ingress_json) = row.map_err(sqlite_error)?;
379 let ingress = decode_turn_input_ingress(ingress_json)?;
380 if ingress
381 .active_turn_id()
382 .is_some_and(|active| active == turn_id)
383 {
384 input_ids.push(input_id);
385 }
386 }
387 input_ids
388 };
389 let next_turn_ingress = encode_json(&lash_core::TurnInputIngress::NextTurn);
390 let mut stmt = tx
391 .prepare(
392 "UPDATE pending_turn_inputs
393 SET state = ?3,
394 ingress_json = ?4,
395 claim_id = NULL,
396 claim_owner_id = NULL,
397 claim_owner_incarnation_id = NULL,
398 claim_owner_liveness_json = NULL,
399 claim_token = NULL,
400 claim_session_lease_generation = 0
401 WHERE session_id = ?1 AND input_id = ?2",
402 )
403 .map_err(sqlite_error)?;
404 for input_id in input_ids {
405 stmt.execute(params![
406 commit.session_id,
407 input_id,
408 lash_core::TurnInputState::DeferredNextTurn.as_str(),
409 next_turn_ingress
410 ])
411 .map_err(sqlite_error)?;
412 }
413 }
414 if !commit.committed_attachment_ids.is_empty() {
415 let now = now as i64;
416 let mut stmt = tx
417 .prepare(
418 "UPDATE attachment_manifest
419 SET committed_at_ms = COALESCE(committed_at_ms, ?1)
420 WHERE attachment_id = ?2 AND session_id = ?3",
421 )
422 .map_err(sqlite_error)?;
423 for id in &commit.committed_attachment_ids {
424 stmt.execute(params![now, id.as_str(), commit.session_id])
425 .map_err(sqlite_error)?;
426 }
427 }
428 let mut enqueued_queue_batches = Vec::new();
429 for (index, batch) in commit.enqueued_queue_batches.iter().enumerate() {
430 if batch.session_id != commit.session_id {
431 return Err(StoreError::SessionBindingMismatch {
432 bound_session_id: commit.session_id.clone(),
433 attempted_session_id: batch.session_id.clone(),
434 });
435 }
436 enqueued_queue_batches.push(enqueue_queued_work_conn(
437 tx,
438 batch,
439 now,
440 enqueue_nonce_start.saturating_add(index as u64),
441 )?);
442 }
443 let result = RuntimeCommitResult {
444 head_revision: next_revision,
445 checkpoint_ref: stored_checkpoint.checkpoint_ref,
446 manifest: stored_checkpoint.manifest,
447 enqueued_queue_batches,
448 };
449 if let Some(completed) = &commit.turn_commit {
450 tx.execute(
451 "INSERT INTO runtime_turn_commits (
452 session_id, turn_id, turn_commit_hash, result_json, committed_at_ms
453 )
454 VALUES (?1, ?2, ?3, ?4, ?5)",
455 params![
456 completed.session_id,
457 completed.turn_id,
458 completed.turn_commit_hash,
459 encode_json(&result),
460 now as i64
461 ],
462 )
463 .map_err(sqlite_error)?;
464 }
465 if let Some(completion) = commit.release_session_execution_lease.as_ref() {
466 release_session_execution_lease_conn(tx, completion)?;
467 }
468 Ok(result)
469 })();
470 match outcome {
475 Ok(value) => Ok(TxOutcome::Commit(Ok(value))),
476 Err(err) => Ok(TxOutcome::Rollback(Err(err))),
477 }
478 })
479 .await
480 .map_err(sqlite_error)??;
481 self.maybe_auto_gc().await;
482 Ok(result)
483 }
484
485 async fn save_session_meta(&self, meta: SessionMeta) -> Result<(), StoreError> {
486 Store::save_session_meta(self, meta).await;
487 Ok(())
488 }
489
490 async fn load_session_meta(&self) -> Result<Option<SessionMeta>, StoreError> {
491 Ok(Store::load_session_meta(self).await)
492 }
493}
494
495#[async_trait::async_trait]
496impl SessionExecutionLeaseStore for Store {
497 async fn try_claim_session_execution_lease(
498 &self,
499 session_id: &str,
500 owner: &LeaseOwnerIdentity,
501 lease_ttl_ms: u64,
502 ) -> Result<SessionExecutionLeaseClaimOutcome, StoreError> {
503 let session_id = session_id.to_string();
504 let owner = owner.clone();
505 let now = self.clock.timestamp_ms();
506 self.conn
507 .write_flow(move |tx| {
508 let outcome: Result<SessionExecutionLeaseClaimOutcome, StoreError> = (|| {
509 let current = load_session_execution_lease_row_conn(tx, &session_id)?;
510 if current.as_ref().is_some_and(|lease| {
511 lease.lease_token.is_some() && lease.expires_at_ms > now
512 }) {
513 let current = current.expect("checked current lease is present");
514 if current
515 .owner
516 .as_ref()
517 .is_some_and(|current_owner| current_owner.same_incarnation(&owner))
518 {
519 let expires_at = now.saturating_add(lease_ttl_ms);
520 tx.execute(
521 "UPDATE session_execution_leases
522 SET lease_expires_at_ms = ?2
523 WHERE session_id = ?1",
524 params![session_id, expires_at as i64],
525 )
526 .map_err(sqlite_error)?;
527 return Ok(SessionExecutionLeaseClaimOutcome::Acquired(
528 SessionExecutionLease {
529 session_id,
530 owner,
531 lease_token: current.lease_token.expect("live lease token set"),
532 fencing_token: current.fencing_token,
533 claimed_at_epoch_ms: current.claimed_at_ms,
534 expires_at_epoch_ms: expires_at,
535 },
536 ));
537 }
538 return Ok(SessionExecutionLeaseClaimOutcome::Busy {
539 holder: row_to_session_execution_lease(&session_id, current)?,
540 });
541 }
542 Ok(SessionExecutionLeaseClaimOutcome::Acquired(
543 acquire_session_execution_lease_conn(
544 tx,
545 &session_id,
546 &owner,
547 current.as_ref().map_or(0, |lease| lease.fencing_token),
548 now,
549 lease_ttl_ms,
550 )?,
551 ))
552 })(
553 );
554 match outcome {
555 Ok(value) => Ok(TxOutcome::Commit(Ok(value))),
556 Err(err) => Ok(TxOutcome::Rollback(Err(err))),
557 }
558 })
559 .await
560 .map_err(sqlite_error)?
561 }
562
563 async fn reclaim_session_execution_lease(
564 &self,
565 session_id: &str,
566 owner: &LeaseOwnerIdentity,
567 observed_holder: &SessionExecutionLeaseFence,
568 lease_ttl_ms: u64,
569 ) -> Result<SessionExecutionLeaseClaimOutcome, StoreError> {
570 let session_id = session_id.to_string();
571 let owner = owner.clone();
572 let observed_holder = observed_holder.clone();
573 let now = self.clock.timestamp_ms();
574 self.conn
575 .write_flow(move |tx| {
576 let outcome: Result<SessionExecutionLeaseClaimOutcome, StoreError> = (|| {
577 let current = load_session_execution_lease_row_conn(tx, &session_id)?;
578 let Some(current) = current else {
579 return Ok(SessionExecutionLeaseClaimOutcome::Acquired(
580 acquire_session_execution_lease_conn(
581 tx,
582 &session_id,
583 &owner,
584 0,
585 now,
586 lease_ttl_ms,
587 )?,
588 ));
589 };
590 if current.lease_token.is_none() || current.expires_at_ms <= now {
591 return Ok(SessionExecutionLeaseClaimOutcome::Acquired(
592 acquire_session_execution_lease_conn(
593 tx,
594 &session_id,
595 &owner,
596 current.fencing_token,
597 now,
598 lease_ttl_ms,
599 )?,
600 ));
601 }
602 let holder = row_to_session_execution_lease(&session_id, current)?;
603 if observed_holder.session_id == session_id
604 && holder.owner.same_incarnation(&observed_holder.owner)
605 && holder.lease_token == observed_holder.lease_token
606 && holder.fencing_token == observed_holder.fencing_token
607 && holder.owner.is_definitely_dead_for_claimant(&owner)
608 {
609 let fencing_token = holder.fencing_token.saturating_add(1);
610 let lease_token = format!(
611 "{}:{}:{}:{now}:{fencing_token}",
612 session_id, owner.owner_id, owner.incarnation_id
613 );
614 let expires_at = now.saturating_add(lease_ttl_ms);
615 let liveness_json = encode_liveness(&owner.liveness)?;
616 let changed = tx
617 .execute(
618 "UPDATE session_execution_leases
619 SET lease_owner_id = ?1,
620 lease_owner_incarnation_id = ?2,
621 lease_owner_liveness_json = ?3,
622 lease_token = ?4,
623 lease_fencing_token = ?5,
624 lease_claimed_at_ms = ?6,
625 lease_expires_at_ms = ?7
626 WHERE session_id = ?8
627 AND lease_owner_id = ?9
628 AND lease_owner_incarnation_id = ?10
629 AND lease_token = ?11
630 AND lease_fencing_token = ?12",
631 params![
632 owner.owner_id,
633 owner.incarnation_id,
634 liveness_json,
635 lease_token,
636 fencing_token as i64,
637 now as i64,
638 expires_at as i64,
639 session_id,
640 observed_holder.owner.owner_id,
641 observed_holder.owner.incarnation_id,
642 observed_holder.lease_token,
643 observed_holder.fencing_token as i64,
644 ],
645 )
646 .map_err(sqlite_error)?;
647 if changed == 1 {
648 return Ok(SessionExecutionLeaseClaimOutcome::Acquired(
649 SessionExecutionLease {
650 session_id,
651 owner,
652 lease_token,
653 fencing_token,
654 claimed_at_epoch_ms: now,
655 expires_at_epoch_ms: expires_at,
656 },
657 ));
658 }
659 let current = load_session_execution_lease_row_conn(tx, &session_id)?;
660 if current.as_ref().is_some_and(|lease| {
661 lease.lease_token.is_some() && lease.expires_at_ms > now
662 }) {
663 let current = current.expect("checked current lease is present");
664 return Ok(SessionExecutionLeaseClaimOutcome::Busy {
665 holder: row_to_session_execution_lease(&session_id, current)?,
666 });
667 }
668 let previous_fencing_token =
669 current.as_ref().map_or(0, |lease| lease.fencing_token);
670 return Ok(SessionExecutionLeaseClaimOutcome::Acquired(
671 acquire_session_execution_lease_conn(
672 tx,
673 &session_id,
674 &owner,
675 previous_fencing_token,
676 now,
677 lease_ttl_ms,
678 )?,
679 ));
680 }
681 Ok(SessionExecutionLeaseClaimOutcome::Busy { holder })
682 })(
683 );
684 match outcome {
685 Ok(value) => Ok(TxOutcome::Commit(Ok(value))),
686 Err(err) => Ok(TxOutcome::Rollback(Err(err))),
687 }
688 })
689 .await
690 .map_err(sqlite_error)?
691 }
692
693 async fn renew_session_execution_lease(
694 &self,
695 fence: &SessionExecutionLeaseFence,
696 lease_ttl_ms: u64,
697 ) -> Result<SessionExecutionLease, StoreError> {
698 let fence = fence.clone();
699 let now = self.clock.timestamp_ms();
700 self.conn
701 .write_flow(move |tx| {
702 let outcome: Result<SessionExecutionLease, StoreError> = (|| {
703 let current = load_session_execution_lease_row_conn(tx, &fence.session_id)?;
704 let Some(current) = current else {
705 return Err(StoreError::SessionExecutionLeaseExpired {
706 session_id: fence.session_id.clone(),
707 });
708 };
709 if !current
710 .owner
711 .as_ref()
712 .is_some_and(|owner| owner.same_incarnation(&fence.owner))
713 || current.lease_token.as_deref() != Some(fence.lease_token.as_str())
714 || current.fencing_token != fence.fencing_token
715 || current.expires_at_ms <= now
716 {
717 return Err(StoreError::SessionExecutionLeaseExpired {
718 session_id: fence.session_id.clone(),
719 });
720 }
721 let expires_at = now.saturating_add(lease_ttl_ms);
722 tx.execute(
723 "UPDATE session_execution_leases
724 SET lease_expires_at_ms = ?5
725 WHERE session_id = ?1
726 AND lease_owner_id = ?2
727 AND lease_owner_incarnation_id = ?3
728 AND lease_token = ?4
729 AND lease_fencing_token = ?6",
730 params![
731 fence.session_id,
732 fence.owner.owner_id,
733 fence.owner.incarnation_id,
734 fence.lease_token,
735 expires_at as i64,
736 fence.fencing_token as i64
737 ],
738 )
739 .map_err(sqlite_error)?;
740 Ok(SessionExecutionLease {
741 session_id: fence.session_id,
742 owner: fence.owner,
743 lease_token: fence.lease_token,
744 fencing_token: fence.fencing_token,
745 claimed_at_epoch_ms: current.claimed_at_ms,
746 expires_at_epoch_ms: expires_at,
747 })
748 })();
749 match outcome {
750 Ok(value) => Ok(TxOutcome::Commit(Ok(value))),
751 Err(err) => Ok(TxOutcome::Rollback(Err(err))),
752 }
753 })
754 .await
755 .map_err(sqlite_error)?
756 }
757
758 async fn release_session_execution_lease(
759 &self,
760 completion: &SessionExecutionLeaseCompletion,
761 ) -> Result<(), StoreError> {
762 let completion = completion.clone();
763 self.conn
764 .write_flow(move |tx| {
765 let outcome = release_session_execution_lease_conn(tx, &completion);
766 match outcome {
767 Ok(()) => Ok(TxOutcome::Commit(Ok(()))),
768 Err(err) => Ok(TxOutcome::Rollback(Err(err))),
769 }
770 })
771 .await
772 .map_err(sqlite_error)?
773 }
774}
775
776#[async_trait::async_trait]
777impl QueuedWorkStore for Store {
778 async fn enqueue_queued_work(
779 &self,
780 batch: QueuedWorkBatchDraft,
781 ) -> Result<QueuedWorkBatch, StoreError> {
782 let nonce = self.commit_count.fetch_add(1, AtomicOrdering::Relaxed);
783 let now = self.clock.timestamp_ms();
784 self.conn
785 .write_flow(move |tx| {
786 let outcome = enqueue_queued_work_conn(tx, &batch, now, nonce);
787 match outcome {
790 Ok(value) => Ok(TxOutcome::Commit(Ok(value))),
791 Err(err) => Ok(TxOutcome::Rollback(Err(err))),
792 }
793 })
794 .await
795 .map_err(sqlite_error)?
796 }
797
798 async fn claim_leading_ready_session_command(
799 &self,
800 session_id: &str,
801 session_execution_lease: &SessionExecutionLeaseFence,
802 owner: &LeaseOwnerIdentity,
803 ) -> Result<Option<QueuedWorkClaim>, StoreError> {
804 let session_id = session_id.to_string();
805 let session_execution_lease = session_execution_lease.clone();
806 let owner = owner.clone();
807 let now = self.clock.timestamp_ms();
808 self.conn
809 .write_flow(move |tx| {
810 let outcome: Result<TxOutcome<Option<QueuedWorkClaim>>, StoreError> = (|| {
811 ensure_session_execution_lease_conn(
812 tx,
813 &session_id,
814 &session_execution_lease,
815 now,
816 )?;
817 let generation = session_execution_lease.fencing_token;
821 let candidate_rows = {
822 let mut stmt = tx
823 .prepare(
824 "SELECT enqueue_seq, batch_id, session_id, source_key, delivery_policy,
825 slot_policy, merge_key_json, available_at_ms, enqueued_at_ms,
826 claim_fencing_token, claim_owner_id, claim_owner_incarnation_id,
827 claim_owner_liveness_json, claim_token, claim_session_lease_generation
828 FROM queued_work_batches
829 WHERE session_id = ?1
830 AND available_at_ms <= ?2
831 AND (
832 claim_token IS NULL
833 OR claim_session_lease_generation <> ?3
834 )
835 ORDER BY enqueue_seq ASC
836 LIMIT ?4",
837 )
838 .map_err(sqlite_error)?;
839 let rows = stmt
840 .query_map(
841 params![
842 session_id,
843 now as i64,
844 generation as i64,
845 claim_scan_limit(1)
846 ],
847 queued_batch_row_from_sql,
848 )
849 .map_err(sqlite_error)?;
850 rows.collect::<Result<Vec<_>, _>>().map_err(sqlite_error)?
851 };
852 let candidate_rows = candidate_rows
853 .into_iter()
854 .filter(|row| {
855 row.claim_token.is_none()
856 || row.claim_session_lease_generation != generation
857 })
858 .collect::<Vec<_>>();
859 let candidate_batches = candidate_rows
860 .iter()
861 .map(|row| queued_work_batch_from_conn(tx, row.clone()))
862 .collect::<Result<Vec<_>, StoreError>>()?;
863 let candidates = candidate_rows
864 .iter()
865 .zip(candidate_batches.iter())
866 .map(|(row, batch)| {
867 Ok(ClaimCandidate {
868 enqueue_seq: row.enqueue_seq,
869 claim_fencing_token: row.claim_fencing_token,
870 work_class: batch.work_class().ok_or_else(|| {
871 StoreError::Backend(format!(
872 "queued-work batch `{}` has mixed or empty payload classes",
873 batch.batch_id
874 ))
875 })?,
876 delivery_policy: decode_delivery_policy(
877 row.delivery_policy.clone(),
878 )?,
879 slot_policy: decode_slot_policy(row.slot_policy.clone())?,
880 merge_key: decode_merge_key(row.merge_key_json.clone())?,
881 })
882 })
883 .collect::<Result<Vec<_>, StoreError>>()?;
884 let selected_len = select_leading_session_command(&candidates);
885 if selected_len == 0 {
886 return Ok(TxOutcome::Commit(None));
887 }
888 let mut selected = candidate_rows;
889 selected.truncate(selected_len);
890 let mut selected_batches = candidate_batches;
891 selected_batches.truncate(selected_len);
892 let lease = QueuedWorkClaimLease::derive(
893 &candidates[0],
894 &session_id,
895 &owner,
896 now,
897 generation,
898 );
899 let liveness_json = encode_liveness(&owner.liveness)?;
900 for row in &selected {
901 let claimed = tx
902 .execute(
903 "UPDATE queued_work_batches
904 SET claim_id = ?3,
905 claim_owner_id = ?4,
906 claim_owner_incarnation_id = ?5,
907 claim_owner_liveness_json = ?6,
908 claim_token = ?7,
909 claim_fencing_token = claim_fencing_token + 1,
910 claim_session_lease_generation = ?8
911 WHERE session_id = ?1
912 AND batch_id = ?2
913 AND (
914 claim_token IS NULL
915 OR claim_session_lease_generation <> ?8
916 )",
917 params![
918 session_id,
919 row.batch_id,
920 lease.claim_id,
921 owner.owner_id.as_str(),
922 owner.incarnation_id.as_str(),
923 liveness_json.as_str(),
924 lease.lease_token,
925 lease.session_lease_generation as i64,
926 ],
927 )
928 .map_err(sqlite_error)?;
929 if claimed == 0 {
930 return Ok(TxOutcome::Rollback(None));
931 }
932 }
933 Ok(TxOutcome::Commit(Some(QueuedWorkClaim {
934 session_id: session_id.clone(),
935 claim_id: lease.claim_id,
936 owner: owner.clone(),
937 lease_token: lease.lease_token,
938 fencing_token: lease.fencing_token,
939 session_lease_generation: lease.session_lease_generation,
940 batches: selected_batches,
941 })))
942 })();
943 match outcome {
944 Ok(TxOutcome::Commit(value)) => Ok(TxOutcome::Commit(Ok(value))),
945 Ok(TxOutcome::Rollback(value)) => Ok(TxOutcome::Rollback(Ok(value))),
946 Err(err) => Ok(TxOutcome::Rollback(Err(err))),
947 }
948 })
949 .await
950 .map_err(sqlite_error)?
951 }
952
953 async fn claim_ready_queued_work(
954 &self,
955 session_id: &str,
956 session_execution_lease: &SessionExecutionLeaseFence,
957 owner: &LeaseOwnerIdentity,
958 boundary: QueuedWorkClaimBoundary,
959 max_batches: usize,
960 ) -> Result<Option<QueuedWorkClaim>, StoreError> {
961 if max_batches == 0 {
962 return Ok(None);
963 }
964 let session_id = session_id.to_string();
965 let session_execution_lease = session_execution_lease.clone();
966 let owner = owner.clone();
967 let now = self.clock.timestamp_ms();
968 self.conn
969 .write_flow(move |tx| {
970 let outcome: Result<TxOutcome<Option<QueuedWorkClaim>>, StoreError> = (|| {
971 ensure_session_execution_lease_conn(
972 tx,
973 &session_id,
974 &session_execution_lease,
975 now,
976 )?;
977 let generation = session_execution_lease.fencing_token;
978 let candidate_rows = {
979 let mut stmt = tx
980 .prepare(
981 "SELECT enqueue_seq, batch_id, session_id, source_key, delivery_policy,
982 slot_policy, merge_key_json, available_at_ms, enqueued_at_ms,
983 claim_fencing_token, claim_owner_id, claim_owner_incarnation_id,
984 claim_owner_liveness_json, claim_token, claim_session_lease_generation
985 FROM queued_work_batches
986 WHERE session_id = ?1
987 AND available_at_ms <= ?2
988 AND (
989 claim_token IS NULL
990 OR claim_session_lease_generation <> ?3
991 )
992 ORDER BY enqueue_seq ASC
993 LIMIT ?4",
994 )
995 .map_err(sqlite_error)?;
996 let rows = stmt
997 .query_map(
998 params![
999 session_id,
1000 now as i64,
1001 generation as i64,
1002 claim_scan_limit(max_batches)
1003 ],
1004 queued_batch_row_from_sql,
1005 )
1006 .map_err(sqlite_error)?;
1007 rows.collect::<Result<Vec<_>, _>>().map_err(sqlite_error)?
1008 };
1009 let candidate_rows = candidate_rows
1010 .into_iter()
1011 .filter(|row| {
1012 row.claim_token.is_none()
1013 || row.claim_session_lease_generation != generation
1014 })
1015 .collect::<Vec<_>>();
1016 let candidate_batches = candidate_rows
1017 .iter()
1018 .map(|row| queued_work_batch_from_conn(tx, row.clone()))
1019 .collect::<Result<Vec<_>, StoreError>>()?;
1020 let candidates = candidate_rows
1021 .iter()
1022 .zip(candidate_batches.iter())
1023 .map(|(row, batch)| {
1024 Ok(ClaimCandidate {
1025 enqueue_seq: row.enqueue_seq,
1026 claim_fencing_token: row.claim_fencing_token,
1027 work_class: batch.work_class().ok_or_else(|| {
1028 StoreError::Backend(format!(
1029 "queued-work batch `{}` has mixed or empty payload classes",
1030 batch.batch_id
1031 ))
1032 })?,
1033 delivery_policy: decode_delivery_policy(
1034 row.delivery_policy.clone(),
1035 )?,
1036 slot_policy: decode_slot_policy(row.slot_policy.clone())?,
1037 merge_key: decode_merge_key(row.merge_key_json.clone())?,
1038 })
1039 })
1040 .collect::<Result<Vec<_>, StoreError>>()?;
1041 let selected_len =
1042 select_turn_work_claim_prefix(&candidates, boundary, max_batches);
1043 if selected_len == 0 {
1044 return Ok(TxOutcome::Commit(None));
1045 }
1046 let mut selected = candidate_rows;
1047 selected.truncate(selected_len);
1048 let mut selected_batches = candidate_batches;
1049 selected_batches.truncate(selected_len);
1050 let lease = QueuedWorkClaimLease::derive(
1051 &candidates[0],
1052 &session_id,
1053 &owner,
1054 now,
1055 generation,
1056 );
1057 let liveness_json = encode_liveness(&owner.liveness)?;
1058 for row in &selected {
1059 let claimed = tx
1068 .execute(
1069 "UPDATE queued_work_batches
1070 SET claim_id = ?3,
1071 claim_owner_id = ?4,
1072 claim_owner_incarnation_id = ?5,
1073 claim_owner_liveness_json = ?6,
1074 claim_token = ?7,
1075 claim_fencing_token = claim_fencing_token + 1,
1076 claim_session_lease_generation = ?8
1077 WHERE session_id = ?1
1078 AND batch_id = ?2
1079 AND (
1080 claim_token IS NULL
1081 OR claim_session_lease_generation <> ?8
1082 )",
1083 params![
1084 session_id,
1085 row.batch_id,
1086 lease.claim_id,
1087 owner.owner_id.as_str(),
1088 owner.incarnation_id.as_str(),
1089 liveness_json.as_str(),
1090 lease.lease_token,
1091 lease.session_lease_generation as i64,
1092 ],
1093 )
1094 .map_err(sqlite_error)?;
1095 if claimed == 0 {
1096 return Ok(TxOutcome::Rollback(None));
1100 }
1101 }
1102 Ok(TxOutcome::Commit(Some(QueuedWorkClaim {
1103 session_id: session_id.clone(),
1104 claim_id: lease.claim_id,
1105 owner: owner.clone(),
1106 lease_token: lease.lease_token,
1107 fencing_token: lease.fencing_token,
1108 session_lease_generation: lease.session_lease_generation,
1109 batches: selected_batches,
1110 })))
1111 })();
1112 match outcome {
1116 Ok(TxOutcome::Commit(value)) => Ok(TxOutcome::Commit(Ok(value))),
1117 Ok(TxOutcome::Rollback(value)) => Ok(TxOutcome::Rollback(Ok(value))),
1118 Err(err) => Ok(TxOutcome::Rollback(Err(err))),
1119 }
1120 })
1121 .await
1122 .map_err(sqlite_error)?
1123 }
1124
1125 async fn claim_ready_queued_work_by_batch_ids(
1126 &self,
1127 session_id: &str,
1128 session_execution_lease: &SessionExecutionLeaseFence,
1129 owner: &LeaseOwnerIdentity,
1130 boundary: QueuedWorkClaimBoundary,
1131 batch_ids: &[String],
1132 ) -> Result<Option<QueuedWorkClaim>, StoreError> {
1133 if batch_ids.is_empty() {
1134 return Ok(None);
1135 }
1136 let session_id = session_id.to_string();
1137 let fence = session_execution_lease.clone();
1138 let owner = owner.clone();
1139 let batch_ids = batch_ids.to_vec();
1140 let now = self.clock.timestamp_ms();
1141 self.conn
1142 .write_flow(move |tx| {
1143 let outcome: Result<Option<QueuedWorkClaim>, StoreError> = (|| {
1144 ensure_session_execution_lease_conn(tx, &session_id, &fence, now)?;
1145 let generation = fence.fencing_token;
1146 let mut rows = Vec::new();
1147 let mut batches = Vec::new();
1148 for batch_id in &batch_ids {
1149 let row = tx
1150 .query_row(
1151 "SELECT enqueue_seq, batch_id, session_id, source_key,
1152 delivery_policy, slot_policy, merge_key_json,
1153 available_at_ms, enqueued_at_ms, claim_fencing_token,
1154 claim_owner_id, claim_owner_incarnation_id,
1155 claim_owner_liveness_json, claim_token,
1156 claim_session_lease_generation
1157 FROM queued_work_batches
1158 WHERE session_id = ?1 AND batch_id = ?2
1159 AND available_at_ms <= ?3
1160 AND (claim_token IS NULL
1161 OR claim_session_lease_generation <> ?4)",
1162 params![session_id, batch_id, now as i64, generation as i64],
1163 queued_batch_row_from_sql,
1164 )
1165 .optional()
1166 .map_err(sqlite_error)?;
1167 let Some(row) = row else {
1168 return Ok(None);
1169 };
1170 let batch = queued_work_batch_from_conn(tx, row.clone())?;
1171 if batch.work_class() != Some(lash_core::runtime::QueuedWorkClass::TurnWork)
1172 {
1173 return Ok(None);
1174 }
1175 rows.push(row);
1176 batches.push(batch);
1177 }
1178 let candidates = rows
1179 .iter()
1180 .map(|row| {
1181 Ok(ClaimCandidate {
1182 enqueue_seq: row.enqueue_seq,
1183 claim_fencing_token: row.claim_fencing_token,
1184 work_class: lash_core::runtime::QueuedWorkClass::TurnWork,
1185 delivery_policy: decode_delivery_policy(
1186 row.delivery_policy.clone(),
1187 )?,
1188 slot_policy: decode_slot_policy(row.slot_policy.clone())?,
1189 merge_key: decode_merge_key(row.merge_key_json.clone())?,
1190 })
1191 })
1192 .collect::<Result<Vec<_>, StoreError>>()?;
1193 if select_turn_work_claim_prefix(&candidates, boundary, candidates.len())
1194 != candidates.len()
1195 {
1196 return Ok(None);
1197 }
1198 let lease = QueuedWorkClaimLease::derive(
1199 &candidates[0],
1200 &session_id,
1201 &owner,
1202 now,
1203 generation,
1204 );
1205 let owner_liveness_json = encode_liveness(&owner.liveness)?;
1206 for row in &rows {
1207 let changed = tx
1208 .execute(
1209 "UPDATE queued_work_batches
1210 SET claim_id = ?3, claim_owner_id = ?4,
1211 claim_owner_incarnation_id = ?5,
1212 claim_owner_liveness_json = ?6, claim_token = ?7,
1213 claim_fencing_token = claim_fencing_token + 1,
1214 claim_session_lease_generation = ?8
1215 WHERE session_id = ?1 AND batch_id = ?2
1216 AND (claim_token IS NULL
1217 OR claim_session_lease_generation <> ?8)",
1218 params![
1219 session_id,
1220 row.batch_id,
1221 lease.claim_id,
1222 owner.owner_id,
1223 owner.incarnation_id,
1224 owner_liveness_json,
1225 lease.lease_token,
1226 generation as i64,
1227 ],
1228 )
1229 .map_err(sqlite_error)?;
1230 if changed != 1 {
1231 return Ok(None);
1232 }
1233 }
1234 Ok(Some(QueuedWorkClaim {
1235 session_id,
1236 claim_id: lease.claim_id,
1237 owner,
1238 lease_token: lease.lease_token,
1239 fencing_token: lease.fencing_token,
1240 session_lease_generation: lease.session_lease_generation,
1241 batches,
1242 }))
1243 })();
1244 match outcome {
1245 Ok(Some(value)) => Ok(TxOutcome::Commit(Ok(Some(value)))),
1246 Ok(None) => Ok(TxOutcome::Rollback(Ok(None))),
1247 Err(err) => Ok(TxOutcome::Rollback(Err(err))),
1248 }
1249 })
1250 .await
1251 .map_err(sqlite_error)?
1252 }
1253
1254 async fn abandon_queued_work_claim(&self, claim: &QueuedWorkClaim) -> Result<(), StoreError> {
1255 let session_id = claim.session_id.clone();
1256 let claim_id = claim.claim_id.clone();
1257 let lease_token = claim.lease_token.clone();
1258 self.conn
1259 .write(move |tx| {
1260 tx.execute(
1261 "UPDATE queued_work_batches
1262 SET claim_id = NULL,
1263 claim_owner_id = NULL,
1264 claim_owner_incarnation_id = NULL,
1265 claim_owner_liveness_json = NULL,
1266 claim_token = NULL,
1267 claim_session_lease_generation = 0
1268 WHERE session_id = ?1 AND claim_id = ?2 AND claim_token = ?3",
1269 params![session_id, claim_id, lease_token],
1270 )
1271 })
1272 .await
1273 .map_err(sqlite_error)?;
1274 Ok(())
1275 }
1276
1277 async fn cancel_queued_work_batch(
1278 &self,
1279 session_id: &str,
1280 batch_id: &str,
1281 ) -> Result<Option<QueuedWorkBatch>, StoreError> {
1282 let session_id = session_id.to_string();
1283 let batch_id = batch_id.to_string();
1284 let now = self.clock.timestamp_ms() as i64;
1285 self.conn
1286 .write_flow(move |tx| {
1287 let outcome: Result<Option<QueuedWorkBatch>, StoreError> = (|| {
1288 let row = tx
1289 .query_row(
1290 "SELECT enqueue_seq, batch_id, session_id, source_key, delivery_policy,
1291 slot_policy, merge_key_json, available_at_ms, enqueued_at_ms,
1292 claim_fencing_token, claim_owner_id, claim_owner_incarnation_id,
1293 claim_owner_liveness_json, claim_token, claim_session_lease_generation
1294 FROM queued_work_batches
1295 WHERE session_id = ?1
1296 AND batch_id = ?2
1297 AND (claim_token IS NULL OR NOT EXISTS (
1298 SELECT 1 FROM session_execution_leases sel
1299 WHERE sel.session_id = ?1
1300 AND sel.lease_token IS NOT NULL
1301 AND sel.lease_expires_at_ms > ?3
1302 AND sel.lease_fencing_token
1303 = queued_work_batches.claim_session_lease_generation
1304 ))",
1305 params![session_id, batch_id, now],
1306 queued_batch_row_from_sql,
1307 )
1308 .optional()
1309 .map_err(sqlite_error)?;
1310 let Some(row) = row else {
1311 return Ok(None);
1312 };
1313 let batch = queued_work_batch_from_conn(tx, row)?;
1314 tx.execute(
1315 "DELETE FROM queued_work_batches
1316 WHERE session_id = ?1
1317 AND batch_id = ?2
1318 AND (claim_token IS NULL OR NOT EXISTS (
1319 SELECT 1 FROM session_execution_leases sel
1320 WHERE sel.session_id = ?1
1321 AND sel.lease_token IS NOT NULL
1322 AND sel.lease_expires_at_ms > ?3
1323 AND sel.lease_fencing_token
1324 = queued_work_batches.claim_session_lease_generation
1325 ))",
1326 params![session_id, batch_id, now],
1327 )
1328 .map_err(sqlite_error)?;
1329 Ok(Some(batch))
1330 })();
1331 match outcome {
1332 Ok(value) => Ok(TxOutcome::Commit(Ok(value))),
1333 Err(err) => Ok(TxOutcome::Rollback(Err(err))),
1334 }
1335 })
1336 .await
1337 .map_err(sqlite_error)?
1338 }
1339
1340 async fn list_queued_work(&self, session_id: &str) -> Result<Vec<QueuedWorkBatch>, StoreError> {
1341 let session_id = session_id.to_string();
1342 self.conn
1343 .call(move |conn| {
1344 let outcome: Result<Vec<QueuedWorkBatch>, StoreError> = (|| {
1345 let rows = {
1346 let mut stmt = conn
1347 .prepare(
1348 "SELECT enqueue_seq, batch_id, session_id, source_key, delivery_policy,
1349 slot_policy, merge_key_json, available_at_ms, enqueued_at_ms,
1350 claim_fencing_token, claim_owner_id, claim_owner_incarnation_id,
1351 claim_owner_liveness_json, claim_token, claim_session_lease_generation
1352 FROM queued_work_batches
1353 WHERE session_id = ?1
1354 ORDER BY enqueue_seq ASC",
1355 )
1356 .map_err(sqlite_error)?;
1357 let rows = stmt
1358 .query_map(params![session_id], queued_batch_row_from_sql)
1359 .map_err(sqlite_error)?;
1360 rows.collect::<Result<Vec<_>, _>>().map_err(sqlite_error)?
1361 };
1362 rows.into_iter()
1363 .map(|row| queued_work_batch_from_conn(conn, row))
1364 .collect()
1365 })();
1366 Ok(outcome)
1367 })
1368 .await
1369 .map_err(sqlite_error)?
1370 }
1371
1372 async fn list_pending_queued_work(
1373 &self,
1374 session_id: &str,
1375 ) -> Result<Vec<QueuedWorkBatch>, StoreError> {
1376 let session_id = session_id.to_string();
1377 let now = self.clock.timestamp_ms();
1378 self.conn
1379 .call(move |conn| {
1380 let outcome: Result<Vec<QueuedWorkBatch>, StoreError> = (|| {
1381 let rows = {
1382 let mut stmt = conn
1383 .prepare(
1384 "SELECT enqueue_seq, batch_id, session_id, source_key, delivery_policy,
1385 slot_policy, merge_key_json, available_at_ms, enqueued_at_ms,
1386 claim_fencing_token, claim_owner_id, claim_owner_incarnation_id,
1387 claim_owner_liveness_json, claim_token, claim_session_lease_generation
1388 FROM queued_work_batches
1389 WHERE session_id = ?1
1390 AND (claim_token IS NULL OR NOT EXISTS (
1391 SELECT 1 FROM session_execution_leases sel
1392 WHERE sel.session_id = ?1
1393 AND sel.lease_token IS NOT NULL
1394 AND sel.lease_expires_at_ms > ?2
1395 AND sel.lease_fencing_token
1396 = queued_work_batches.claim_session_lease_generation
1397 ))
1398 ORDER BY enqueue_seq ASC",
1399 )
1400 .map_err(sqlite_error)?;
1401 let rows = stmt
1402 .query_map(
1403 params![session_id, now as i64],
1404 queued_batch_row_from_sql,
1405 )
1406 .map_err(sqlite_error)?;
1407 rows.collect::<Result<Vec<_>, _>>().map_err(sqlite_error)?
1408 };
1409 rows.into_iter()
1410 .map(|row| queued_work_batch_from_conn(conn, row))
1411 .collect()
1412 })();
1413 Ok(outcome)
1414 })
1415 .await
1416 .map_err(sqlite_error)?
1417 }
1418}
1419
1420#[async_trait::async_trait]
1421impl TurnInputStore for Store {
1422 async fn enqueue_pending_turn_input(
1423 &self,
1424 draft: lash_core::PendingTurnInputDraft,
1425 ) -> Result<lash_core::PendingTurnInput, StoreError> {
1426 let nonce = self.commit_count.fetch_add(1, AtomicOrdering::Relaxed);
1427 let now = self.clock.timestamp_ms();
1428 self.conn
1429 .write_flow(move |tx| {
1430 let outcome: Result<lash_core::PendingTurnInput, StoreError> = (|| {
1431 if let Some(source_key) = draft.source_key.as_deref() {
1432 let existing_id: Option<String> = tx
1433 .query_row(
1434 "SELECT input_id
1435 FROM pending_turn_inputs
1436 WHERE session_id = ?1 AND source_key = ?2",
1437 params![draft.session_id, source_key],
1438 |row| row.get(0),
1439 )
1440 .optional()
1441 .map_err(sqlite_error)?;
1442 if let Some(input_id) = existing_id {
1443 let existing = load_pending_turn_input_by_id_conn(
1444 tx,
1445 &draft.session_id,
1446 &input_id,
1447 )?
1448 .ok_or_else(|| {
1449 StoreError::Backend(
1450 "pending turn input source row disappeared".to_string(),
1451 )
1452 })?;
1453 if !draft.submitted_content_matches(&existing).map_err(|err| {
1454 StoreError::Backend(format!(
1455 "failed to compare pending turn input submission: {err}"
1456 ))
1457 })? {
1458 return Err(StoreError::PendingTurnInputSourceKeyConflict {
1459 session_id: draft.session_id.clone(),
1460 source_key: source_key.to_string(),
1461 existing_input_id: existing.input_id.clone(),
1462 });
1463 }
1464 return Ok(existing);
1465 }
1466 }
1467 let input_id = draft.input_id.clone().unwrap_or_else(|| {
1468 derive_pending_turn_input_id(
1469 &draft.session_id,
1470 draft.source_key.as_deref(),
1471 now,
1472 nonce,
1473 )
1474 });
1475 let state = match draft.ingress {
1476 lash_core::TurnInputIngress::ActiveTurn { .. } => {
1477 lash_core::TurnInputState::PendingActive
1478 }
1479 lash_core::TurnInputIngress::NextTurn => {
1480 lash_core::TurnInputState::DeferredNextTurn
1481 }
1482 };
1483 tx.execute(
1484 "INSERT INTO pending_turn_inputs (
1485 input_id, session_id, source_key, ingress_json, state,
1486 input_json, enqueued_at_ms
1487 )
1488 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
1489 params![
1490 input_id,
1491 draft.session_id,
1492 draft.source_key.as_deref(),
1493 encode_json(&draft.ingress),
1494 state.as_str(),
1495 encode_json(&draft.input),
1496 now as i64,
1497 ],
1498 )
1499 .map_err(sqlite_error)?;
1500 load_pending_turn_input_by_id_conn(tx, &draft.session_id, &input_id)?
1501 .ok_or_else(|| {
1502 StoreError::Backend("pending turn input insert disappeared".to_string())
1503 })
1504 })();
1505 match outcome {
1506 Ok(value) => Ok(TxOutcome::Commit(Ok(value))),
1507 Err(err) => Ok(TxOutcome::Rollback(Err(err))),
1508 }
1509 })
1510 .await
1511 .map_err(sqlite_error)?
1512 }
1513
1514 async fn list_pending_turn_inputs(
1515 &self,
1516 session_id: &str,
1517 ) -> Result<Vec<lash_core::PendingTurnInput>, StoreError> {
1518 let session_id = session_id.to_string();
1519 let now = self.clock.timestamp_ms();
1520 self.conn
1521 .call(move |conn| {
1522 let outcome: Result<Vec<lash_core::PendingTurnInput>, StoreError> = (|| {
1523 let rows = {
1524 let mut stmt = conn
1525 .prepare(
1526 "SELECT enqueue_seq, input_id, session_id, source_key, ingress_json,
1527 state, input_json, enqueued_at_ms, claim_id, claim_fencing_token,
1528 claim_owner_id, claim_owner_incarnation_id,
1529 claim_owner_liveness_json, claim_token, claim_session_lease_generation
1530 FROM pending_turn_inputs
1531 WHERE session_id = ?1
1532 AND state IN (?2, ?3)
1533 AND (claim_token IS NULL OR NOT EXISTS (
1534 SELECT 1 FROM session_execution_leases sel
1535 WHERE sel.session_id = ?1
1536 AND sel.lease_token IS NOT NULL
1537 AND sel.lease_expires_at_ms > ?4
1538 AND sel.lease_fencing_token
1539 = pending_turn_inputs.claim_session_lease_generation
1540 ))
1541 ORDER BY enqueue_seq ASC",
1542 )
1543 .map_err(sqlite_error)?;
1544 let rows = stmt
1545 .query_map(
1546 params![
1547 session_id,
1548 lash_core::TurnInputState::PendingActive.as_str(),
1549 lash_core::TurnInputState::DeferredNextTurn.as_str(),
1550 now as i64
1551 ],
1552 pending_turn_input_row_from_sql,
1553 )
1554 .map_err(sqlite_error)?;
1555 rows.collect::<Result<Vec<_>, _>>().map_err(sqlite_error)?
1556 };
1557 rows.into_iter().map(pending_turn_input_from_row).collect()
1558 })(
1559 );
1560 Ok(outcome)
1561 })
1562 .await
1563 .map_err(sqlite_error)?
1564 }
1565
1566 async fn cancel_pending_turn_inputs(
1567 &self,
1568 session_id: &str,
1569 targets: &[lash_core::PendingTurnInputCancelTarget],
1570 ) -> Result<Vec<lash_core::PendingTurnInputCancelResult>, StoreError> {
1571 let session_id = session_id.to_string();
1572 let targets = targets.to_vec();
1573 let now = self.clock.timestamp_ms();
1574 self.conn
1575 .write_flow(move |tx| {
1576 let outcome: Result<Vec<lash_core::PendingTurnInputCancelResult>, StoreError> =
1577 (|| {
1578 let mut results = Vec::with_capacity(targets.len());
1579 for target in targets {
1580 let outcome = match load_pending_turn_input_row_by_target_conn(
1581 tx,
1582 &session_id,
1583 &target,
1584 )? {
1585 Some(row) => cancel_pending_turn_input_row_conn(tx, row, now)?,
1586 None => lash_core::PendingTurnInputCancelOutcome::NotFound,
1587 };
1588 results
1589 .push(lash_core::PendingTurnInputCancelResult { target, outcome });
1590 }
1591 Ok(results)
1592 })();
1593 match outcome {
1594 Ok(value) => Ok(TxOutcome::Commit(Ok(value))),
1595 Err(err) => Ok(TxOutcome::Rollback(Err(err))),
1596 }
1597 })
1598 .await
1599 .map_err(sqlite_error)?
1600 }
1601
1602 async fn cancel_pending_turn_input_suffix(
1603 &self,
1604 session_id: &str,
1605 anchor: &lash_core::PendingTurnInputCancelTarget,
1606 ) -> Result<lash_core::PendingTurnInputSuffixCancelOutcome, StoreError> {
1607 let session_id = session_id.to_string();
1608 let anchor = anchor.clone();
1609 let now = self.clock.timestamp_ms();
1610 self.conn
1611 .write_flow(move |tx| {
1612 let outcome: Result<lash_core::PendingTurnInputSuffixCancelOutcome, StoreError> =
1613 (|| {
1614 let Some(anchor_row) =
1615 load_pending_turn_input_row_by_target_conn(tx, &session_id, &anchor)?
1616 else {
1617 return Ok(
1618 lash_core::PendingTurnInputSuffixCancelOutcome::AnchorNotFound {
1619 anchor,
1620 },
1621 );
1622 };
1623 let rows = {
1624 let mut stmt = tx
1625 .prepare(
1626 "SELECT enqueue_seq, input_id, session_id, source_key, ingress_json,
1627 state, input_json, enqueued_at_ms, claim_id, claim_fencing_token,
1628 claim_owner_id, claim_owner_incarnation_id,
1629 claim_owner_liveness_json, claim_token, claim_session_lease_generation
1630 FROM pending_turn_inputs
1631 WHERE session_id = ?1 AND enqueue_seq >= ?2
1632 ORDER BY enqueue_seq ASC",
1633 )
1634 .map_err(sqlite_error)?;
1635 let rows = stmt
1636 .query_map(
1637 params![session_id, anchor_row.enqueue_seq as i64],
1638 pending_turn_input_row_from_sql,
1639 )
1640 .map_err(sqlite_error)?;
1641 rows.collect::<Result<Vec<_>, _>>().map_err(sqlite_error)?
1642 };
1643 let mut outcomes = Vec::with_capacity(rows.len());
1644 for row in rows {
1645 outcomes.push(cancel_pending_turn_input_row_conn(tx, row, now)?);
1646 }
1647 Ok(lash_core::PendingTurnInputSuffixCancelOutcome::Outcomes {
1648 anchor,
1649 outcomes,
1650 })
1651 })();
1652 match outcome {
1653 Ok(value) => Ok(TxOutcome::Commit(Ok(value))),
1654 Err(err) => Ok(TxOutcome::Rollback(Err(err))),
1655 }
1656 })
1657 .await
1658 .map_err(sqlite_error)?
1659 }
1660
1661 async fn claim_active_turn_inputs(
1662 &self,
1663 session_id: &str,
1664 session_execution_lease: &SessionExecutionLeaseFence,
1665 owner: &LeaseOwnerIdentity,
1666 turn_id: &str,
1667 checkpoint: lash_core::CheckpointKind,
1668 max_inputs: usize,
1669 ) -> Result<Option<lash_core::TurnInputClaim>, StoreError> {
1670 claim_pending_turn_inputs_sqlite(
1671 &self.conn,
1672 self.clock.timestamp_ms(),
1673 session_id,
1674 session_execution_lease,
1675 owner,
1676 max_inputs,
1677 lash_core::TurnInputClaimMode::ActiveTurn {
1678 turn_id: turn_id.to_string(),
1679 checkpoint,
1680 },
1681 )
1682 .await
1683 }
1684
1685 async fn claim_next_turn_inputs(
1686 &self,
1687 session_id: &str,
1688 session_execution_lease: &SessionExecutionLeaseFence,
1689 owner: &LeaseOwnerIdentity,
1690 max_inputs: usize,
1691 ) -> Result<Option<lash_core::TurnInputClaim>, StoreError> {
1692 claim_pending_turn_inputs_sqlite(
1693 &self.conn,
1694 self.clock.timestamp_ms(),
1695 session_id,
1696 session_execution_lease,
1697 owner,
1698 max_inputs,
1699 lash_core::TurnInputClaimMode::NextTurn,
1700 )
1701 .await
1702 }
1703
1704 async fn abandon_turn_input_claim(
1705 &self,
1706 claim: &lash_core::TurnInputClaim,
1707 ) -> Result<(), StoreError> {
1708 let session_id = claim.session_id.clone();
1709 let claim_id = claim.claim_id.clone();
1710 let lease_token = claim.lease_token.clone();
1711 let restored_state = match claim.mode {
1712 lash_core::TurnInputClaimMode::ActiveTurn { .. } => {
1713 lash_core::TurnInputState::PendingActive
1714 }
1715 lash_core::TurnInputClaimMode::NextTurn => lash_core::TurnInputState::DeferredNextTurn,
1716 };
1717 self.conn
1718 .write(move |tx| {
1719 tx.execute(
1720 "UPDATE pending_turn_inputs
1721 SET state = CASE
1722 WHEN state = ?4 THEN ?5
1723 ELSE state
1724 END,
1725 claim_id = NULL,
1726 claim_owner_id = NULL,
1727 claim_owner_incarnation_id = NULL,
1728 claim_owner_liveness_json = NULL,
1729 claim_token = NULL,
1730 claim_session_lease_generation = 0
1731 WHERE session_id = ?1 AND claim_id = ?2 AND claim_token = ?3",
1732 params![
1733 session_id,
1734 claim_id,
1735 lease_token,
1736 lash_core::TurnInputState::Accepted.as_str(),
1737 restored_state.as_str(),
1738 ],
1739 )
1740 })
1741 .await
1742 .map_err(sqlite_error)?;
1743 Ok(())
1744 }
1745}
1746
1747#[async_trait::async_trait]
1748impl StoreMaintenance for Store {
1749 async fn tombstone_nodes(&self, ids: &[String]) -> Result<(), StoreError> {
1750 if ids.is_empty() {
1751 return Ok(());
1752 }
1753 let ids = ids.to_vec();
1754 self.conn
1755 .write(move |tx| {
1756 let mut stmt =
1757 tx.prepare("UPDATE graph_nodes SET tombstoned = 1 WHERE node_id = ?1")?;
1758 for id in &ids {
1759 stmt.execute(params![id])?;
1760 }
1761 Ok(())
1762 })
1763 .await
1764 .map_err(sqlite_error)
1765 }
1766
1767 async fn vacuum(&self) -> Result<VacuumReport, StoreError> {
1768 let (removed_node_count, removed_pending_turn_input_tombstone_count) = self
1769 .conn
1770 .write(move |tx| {
1771 let removed_node_count =
1772 tx.execute("DELETE FROM graph_nodes WHERE tombstoned = 1", [])?;
1773 let removed_pending_turn_input_tombstone_count = tx.execute(
1774 "DELETE FROM pending_turn_inputs
1775 WHERE state IN (?1, ?2)",
1776 params![
1777 lash_core::TurnInputState::Cancelled.as_str(),
1778 lash_core::TurnInputState::Completed.as_str()
1779 ],
1780 )?;
1781 Ok((
1782 removed_node_count,
1783 removed_pending_turn_input_tombstone_count,
1784 ))
1785 })
1786 .await
1787 .map_err(sqlite_error)?;
1788 Ok(VacuumReport {
1789 removed_node_count,
1790 removed_pending_turn_input_tombstone_count,
1791 })
1792 }
1793
1794 async fn gc_unreachable(&self) -> Result<GcReport, StoreError> {
1795 Ok(Store::gc_unreachable(self).await)
1796 }
1797}
1798
1799fn derive_pending_turn_input_id(
1800 session_id: &str,
1801 source_key: Option<&str>,
1802 now_epoch_ms: u64,
1803 nonce: u64,
1804) -> String {
1805 format!(
1806 "ti:{:x}",
1807 Sha256::digest(format!("{session_id}:{source_key:?}:{now_epoch_ms}:{nonce}").as_bytes())
1808 )
1809}
1810
1811fn cancel_pending_turn_input_row_conn(
1812 conn: &Connection,
1813 row: PendingTurnInputRow,
1814 now_epoch_ms: u64,
1815) -> Result<lash_core::PendingTurnInputCancelOutcome, StoreError> {
1816 let mut input = pending_turn_input_from_row(row.clone())?;
1817 match input.state {
1818 lash_core::TurnInputState::Cancelled => Ok(
1819 lash_core::PendingTurnInputCancelOutcome::AlreadyCancelled(input),
1820 ),
1821 lash_core::TurnInputState::Completed => Ok(
1822 lash_core::PendingTurnInputCancelOutcome::AlreadyCompleted(input),
1823 ),
1824 lash_core::TurnInputState::Accepted => {
1825 Ok(lash_core::PendingTurnInputCancelOutcome::AlreadyClaimed {
1826 claim: pending_turn_input_claim_diagnostics_from_row(&row, input.state),
1827 input,
1828 })
1829 }
1830 lash_core::TurnInputState::PendingActive | lash_core::TurnInputState::DeferredNextTurn => {
1831 let live_claim = row.claim_token.is_some()
1834 && load_session_execution_lease_row_conn(conn, &row.session_id)?.is_some_and(
1835 |lease| {
1836 lease.lease_token.is_some()
1837 && lease.expires_at_ms > now_epoch_ms
1838 && lease.fencing_token == row.claim_session_lease_generation
1839 },
1840 );
1841 if live_claim {
1842 return Ok(lash_core::PendingTurnInputCancelOutcome::AlreadyClaimed {
1843 claim: pending_turn_input_claim_diagnostics_from_row(&row, input.state),
1844 input,
1845 });
1846 }
1847 conn.execute(
1848 "UPDATE pending_turn_inputs
1849 SET state = ?3,
1850 claim_id = NULL,
1851 claim_owner_id = NULL,
1852 claim_owner_incarnation_id = NULL,
1853 claim_owner_liveness_json = NULL,
1854 claim_token = NULL,
1855 claim_session_lease_generation = 0
1856 WHERE session_id = ?1 AND input_id = ?2",
1857 params![
1858 row.session_id,
1859 row.input_id,
1860 lash_core::TurnInputState::Cancelled.as_str(),
1861 ],
1862 )
1863 .map_err(sqlite_error)?;
1864 input.state = lash_core::TurnInputState::Cancelled;
1865 Ok(lash_core::PendingTurnInputCancelOutcome::Cancelled(input))
1866 }
1867 }
1868}
1869
1870async fn claim_pending_turn_inputs_sqlite(
1871 conn: &SqliteConnection,
1872 now: u64,
1873 session_id: &str,
1874 session_execution_lease: &SessionExecutionLeaseFence,
1875 owner: &LeaseOwnerIdentity,
1876 max_inputs: usize,
1877 mode: lash_core::TurnInputClaimMode,
1878) -> Result<Option<lash_core::TurnInputClaim>, StoreError> {
1879 if max_inputs == 0 {
1880 return Ok(None);
1881 }
1882 let session_id = session_id.to_string();
1883 let session_execution_lease = session_execution_lease.clone();
1884 let owner = owner.clone();
1885 conn.write_flow(move |tx| {
1886 let outcome: Result<TxOutcome<Option<lash_core::TurnInputClaim>>, StoreError> = (|| {
1887 ensure_session_execution_lease_conn(
1888 tx,
1889 &session_id,
1890 &session_execution_lease,
1891 now,
1892 )?;
1893 let generation = session_execution_lease.fencing_token;
1894 let wanted_state = match &mode {
1895 lash_core::TurnInputClaimMode::ActiveTurn { .. } => {
1896 lash_core::TurnInputState::PendingActive
1897 }
1898 lash_core::TurnInputClaimMode::NextTurn => {
1899 lash_core::TurnInputState::DeferredNextTurn
1900 }
1901 };
1902 let candidate_rows = {
1903 let mut stmt = tx
1904 .prepare(
1905 "SELECT enqueue_seq, input_id, session_id, source_key, ingress_json,
1906 state, input_json, enqueued_at_ms, claim_id, claim_fencing_token,
1907 claim_owner_id, claim_owner_incarnation_id,
1908 claim_owner_liveness_json, claim_token, claim_session_lease_generation
1909 FROM pending_turn_inputs
1910 WHERE session_id = ?1 AND state = ?2
1911 AND (
1912 claim_token IS NULL
1913 OR claim_session_lease_generation <> ?3
1914 )
1915 ORDER BY enqueue_seq ASC
1916 LIMIT ?4",
1917 )
1918 .map_err(sqlite_error)?;
1919 let rows = stmt
1920 .query_map(
1921 params![
1922 session_id,
1923 wanted_state.as_str(),
1924 generation as i64,
1925 (max_inputs as i64).saturating_add(32)
1926 ],
1927 pending_turn_input_row_from_sql,
1928 )
1929 .map_err(sqlite_error)?;
1930 rows.collect::<Result<Vec<_>, _>>().map_err(sqlite_error)?
1931 };
1932 let mut selected = Vec::new();
1933 for row in candidate_rows {
1934 let claim_available = row.claim_token.is_none()
1935 || row.claim_session_lease_generation != generation;
1936 if !claim_available {
1937 continue;
1938 }
1939 let input = pending_turn_input_from_row(row.clone())?;
1940 let matches_mode = match &mode {
1941 lash_core::TurnInputClaimMode::ActiveTurn {
1942 turn_id,
1943 checkpoint,
1944 } => {
1945 input
1946 .ingress
1947 .active_turn_id()
1948 .is_some_and(|active| active == turn_id)
1949 && input.ingress.admits_checkpoint(*checkpoint)
1950 }
1951 lash_core::TurnInputClaimMode::NextTurn => input.state.is_next_turn_pending(),
1952 };
1953 if matches_mode {
1954 selected.push((row, input));
1955 if selected.len() >= max_inputs {
1956 break;
1957 }
1958 }
1959 }
1960 let Some((head, _)) = selected.first() else {
1961 return Ok(TxOutcome::Commit(None));
1962 };
1963 let lease = TurnInputClaimLease::derive(head, &session_id, &owner, now, generation);
1964 let liveness_json = encode_liveness(&owner.liveness)?;
1965 let state_after_claim = match &mode {
1966 lash_core::TurnInputClaimMode::ActiveTurn { .. } => {
1967 lash_core::TurnInputState::Accepted
1968 }
1969 lash_core::TurnInputClaimMode::NextTurn => {
1970 lash_core::TurnInputState::DeferredNextTurn
1971 }
1972 };
1973 let mut inputs = Vec::new();
1974 for (row, mut input) in selected {
1975 let claimed = tx
1976 .execute(
1977 "UPDATE pending_turn_inputs
1978 SET state = ?3,
1979 claim_id = ?4,
1980 claim_owner_id = ?5,
1981 claim_owner_incarnation_id = ?6,
1982 claim_owner_liveness_json = ?7,
1983 claim_token = ?8,
1984 claim_fencing_token = claim_fencing_token + 1,
1985 claim_session_lease_generation = ?9
1986 WHERE session_id = ?1
1987 AND input_id = ?2
1988 AND (
1989 claim_token IS NULL
1990 OR claim_session_lease_generation <> ?9
1991 )",
1992 params![
1993 session_id,
1994 row.input_id,
1995 state_after_claim.as_str(),
1996 lease.claim_id,
1997 owner.owner_id.as_str(),
1998 owner.incarnation_id.as_str(),
1999 liveness_json.as_str(),
2000 lease.lease_token,
2001 lease.session_lease_generation as i64,
2002 ],
2003 )
2004 .map_err(sqlite_error)?;
2005 if claimed == 0 {
2006 return Ok(TxOutcome::Rollback(None));
2007 }
2008 input.state = state_after_claim;
2009 inputs.push(input);
2010 }
2011 Ok(TxOutcome::Commit(Some(lash_core::TurnInputClaim {
2012 session_id: session_id.clone(),
2013 claim_id: lease.claim_id,
2014 owner: owner.clone(),
2015 lease_token: lease.lease_token,
2016 fencing_token: lease.fencing_token,
2017 session_lease_generation: lease.session_lease_generation,
2018 mode,
2019 inputs,
2020 })))
2021 })(
2022 );
2023 match outcome {
2024 Ok(TxOutcome::Commit(value)) => Ok(TxOutcome::Commit(Ok(value))),
2025 Ok(TxOutcome::Rollback(value)) => Ok(TxOutcome::Rollback(Ok(value))),
2026 Err(err) => Ok(TxOutcome::Rollback(Err(err))),
2027 }
2028 })
2029 .await
2030 .map_err(sqlite_error)?
2031}
2032
2033struct SessionExecutionLeaseRow {
2034 owner: Option<LeaseOwnerIdentity>,
2035 lease_token: Option<String>,
2036 fencing_token: u64,
2037 claimed_at_ms: u64,
2038 expires_at_ms: u64,
2039}
2040
2041fn load_session_execution_lease_row_conn(
2042 conn: &Connection,
2043 session_id: &str,
2044) -> Result<Option<SessionExecutionLeaseRow>, StoreError> {
2045 let row = conn
2046 .query_row(
2047 "SELECT lease_owner_id, lease_token, lease_fencing_token,
2048 lease_claimed_at_ms, lease_expires_at_ms,
2049 lease_owner_incarnation_id, lease_owner_liveness_json
2050 FROM session_execution_leases
2051 WHERE session_id = ?1",
2052 params![session_id],
2053 |row| {
2054 let owner_id: Option<String> = row.get(0)?;
2055 let incarnation_id: Option<String> = row.get(5)?;
2056 let liveness_json: Option<String> = row.get(6)?;
2057 Ok(SessionExecutionLeaseRow {
2058 owner: lease_owner_from_columns(owner_id, incarnation_id, liveness_json),
2059 lease_token: row.get(1)?,
2060 fencing_token: row.get::<_, i64>(2)? as u64,
2061 claimed_at_ms: row.get::<_, i64>(3)? as u64,
2062 expires_at_ms: row.get::<_, i64>(4)? as u64,
2063 })
2064 },
2065 )
2066 .optional()
2067 .map_err(sqlite_error)?;
2068 Ok(row)
2069}
2070
2071fn lease_owner_from_columns(
2072 owner_id: Option<String>,
2073 incarnation_id: Option<String>,
2074 liveness_json: Option<String>,
2075) -> Option<LeaseOwnerIdentity> {
2076 owner_id.map(|owner_id| LeaseOwnerIdentity {
2077 incarnation_id: incarnation_id.unwrap_or_else(|| owner_id.clone()),
2078 owner_id,
2079 liveness: liveness_json
2080 .as_deref()
2081 .and_then(|json| serde_json::from_str(json).ok())
2082 .unwrap_or(LeaseOwnerLiveness::Opaque),
2083 })
2084}
2085
2086fn encode_liveness(liveness: &LeaseOwnerLiveness) -> Result<String, StoreError> {
2087 serde_json::to_string(liveness)
2088 .map_err(|err| StoreError::Backend(format!("failed to encode lease liveness: {err}")))
2089}
2090
2091fn row_to_session_execution_lease(
2092 session_id: &str,
2093 row: SessionExecutionLeaseRow,
2094) -> Result<SessionExecutionLease, StoreError> {
2095 Ok(SessionExecutionLease {
2096 session_id: session_id.to_string(),
2097 owner: row
2098 .owner
2099 .ok_or_else(|| StoreError::Backend("live session lease missing owner".to_string()))?,
2100 lease_token: row.lease_token.ok_or_else(|| {
2101 StoreError::Backend("live session lease missing lease token".to_string())
2102 })?,
2103 fencing_token: row.fencing_token,
2104 claimed_at_epoch_ms: row.claimed_at_ms,
2105 expires_at_epoch_ms: row.expires_at_ms,
2106 })
2107}
2108
2109fn acquire_session_execution_lease_conn(
2110 conn: &Connection,
2111 session_id: &str,
2112 owner: &LeaseOwnerIdentity,
2113 previous_fencing_token: u64,
2114 now: u64,
2115 lease_ttl_ms: u64,
2116) -> Result<SessionExecutionLease, StoreError> {
2117 let fencing_token = previous_fencing_token.saturating_add(1);
2118 let lease_token = format!(
2119 "{}:{}:{}:{now}:{fencing_token}",
2120 session_id, owner.owner_id, owner.incarnation_id
2121 );
2122 let expires_at = now.saturating_add(lease_ttl_ms);
2123 let liveness_json = encode_liveness(&owner.liveness)?;
2124 conn.execute(
2125 "INSERT INTO session_execution_leases (
2126 session_id, lease_owner_id, lease_owner_incarnation_id, lease_owner_liveness_json,
2127 lease_token, lease_fencing_token, lease_claimed_at_ms, lease_expires_at_ms
2128 )
2129 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)
2130 ON CONFLICT(session_id) DO UPDATE SET
2131 lease_owner_id = excluded.lease_owner_id,
2132 lease_owner_incarnation_id = excluded.lease_owner_incarnation_id,
2133 lease_owner_liveness_json = excluded.lease_owner_liveness_json,
2134 lease_token = excluded.lease_token,
2135 lease_fencing_token = excluded.lease_fencing_token,
2136 lease_claimed_at_ms = excluded.lease_claimed_at_ms,
2137 lease_expires_at_ms = excluded.lease_expires_at_ms",
2138 params![
2139 session_id,
2140 owner.owner_id,
2141 owner.incarnation_id,
2142 liveness_json,
2143 lease_token,
2144 fencing_token as i64,
2145 now as i64,
2146 expires_at as i64
2147 ],
2148 )
2149 .map_err(sqlite_error)?;
2150 Ok(SessionExecutionLease {
2151 session_id: session_id.to_string(),
2152 owner: owner.clone(),
2153 lease_token,
2154 fencing_token,
2155 claimed_at_epoch_ms: now,
2156 expires_at_epoch_ms: expires_at,
2157 })
2158}
2159
2160fn ensure_session_execution_lease_conn(
2161 conn: &Connection,
2162 session_id: &str,
2163 fence: &SessionExecutionLeaseFence,
2164 now: u64,
2165) -> Result<(), StoreError> {
2166 if fence.session_id != session_id {
2167 return Err(StoreError::SessionExecutionLeaseExpired {
2168 session_id: session_id.to_string(),
2169 });
2170 }
2171 let current = load_session_execution_lease_row_conn(conn, session_id)?;
2172 let Some(current) = current else {
2173 return Err(StoreError::SessionExecutionLeaseExpired {
2174 session_id: session_id.to_string(),
2175 });
2176 };
2177 if current
2178 .owner
2179 .as_ref()
2180 .is_some_and(|owner| owner.same_incarnation(&fence.owner))
2181 && current.lease_token.as_deref() == Some(fence.lease_token.as_str())
2182 && current.fencing_token == fence.fencing_token
2183 && current.expires_at_ms > now
2184 {
2185 Ok(())
2186 } else {
2187 Err(StoreError::SessionExecutionLeaseExpired {
2188 session_id: session_id.to_string(),
2189 })
2190 }
2191}
2192
2193fn release_session_execution_lease_conn(
2194 conn: &Connection,
2195 completion: &SessionExecutionLeaseCompletion,
2196) -> Result<(), StoreError> {
2197 conn.execute(
2198 "UPDATE session_execution_leases
2199 SET lease_owner_id = NULL,
2200 lease_owner_incarnation_id = NULL,
2201 lease_owner_liveness_json = NULL,
2202 lease_token = NULL,
2203 lease_claimed_at_ms = 0,
2204 lease_expires_at_ms = 0
2205 WHERE session_id = ?1
2206 AND lease_owner_id = ?2
2207 AND lease_owner_incarnation_id = ?3
2208 AND lease_token = ?4
2209 AND lease_fencing_token = ?5",
2210 params![
2211 completion.session_id,
2212 completion.owner.owner_id,
2213 completion.owner.incarnation_id,
2214 completion.lease_token,
2215 completion.fencing_token as i64
2216 ],
2217 )
2218 .map_err(sqlite_error)?;
2219 Ok(())
2220}