1use std::{
8 any::TypeId, future::Future, marker::PhantomData, mem, panic::AssertUnwindSafe, pin::pin,
9 task::Poll,
10};
11
12use saddle_admission::DbRequestPermit;
13use saddle_runtime::db_finalizer::{
14 DbQueryPoll, DbQueryTransition, DbTransitionRequest, drive_db_finalizer_with_transition,
15};
16use sqlx::MySql;
17
18use crate::{
19 Database,
20 c6_query_optional::{
21 QueryOptionalContractError, QueryOptionalParameterShape, sealed::ParameterShape as _,
22 },
23};
24
25pub trait StaticWriteOperation: Send + 'static {
27 type Parameters: QueryOptionalParameterShape;
28
29 const OPERATION: &'static str;
30 const SQL: &'static str;
31}
32
33#[derive(Clone, Copy, Debug, Eq, PartialEq)]
34pub struct WriteCreditDemand {
35 connections: u32,
36 operations: u32,
37}
38
39impl WriteCreditDemand {
40 const ONE: Self = Self {
41 connections: 1,
42 operations: 1,
43 };
44
45 pub const fn connections(self) -> u32 {
46 self.connections
47 }
48
49 pub const fn operations(self) -> u32 {
50 self.operations
51 }
52
53 pub const fn merge_route_max(self, other: Self) -> Self {
54 Self {
55 connections: if self.connections > other.connections {
56 self.connections
57 } else {
58 other.connections
59 },
60 operations: if self.operations > other.operations {
61 self.operations
62 } else {
63 other.operations
64 },
65 }
66 }
67}
68
69#[derive(Clone, Copy, Debug, Eq, PartialEq)]
70pub struct WriteLayout {
71 operation: TypeId,
72 parameter_bytes: usize,
73 credits: WriteCreditDemand,
74}
75
76impl WriteLayout {
77 pub const fn parameter_bytes(self) -> usize {
78 self.parameter_bytes
79 }
80
81 pub const fn credits(self) -> WriteCreditDemand {
82 self.credits
83 }
84
85 pub fn belongs_to<O: StaticWriteOperation>(self) -> bool {
86 self.operation == TypeId::of::<O>()
87 }
88}
89
90#[derive(Clone, Copy, Debug)]
91pub struct WriteOperationProof<O: StaticWriteOperation> {
92 layout: WriteLayout,
93 _operation: PhantomData<fn() -> O>,
94}
95
96impl<O: StaticWriteOperation> WriteOperationProof<O> {
97 pub fn bind() -> Result<Self, QueryOptionalContractError> {
98 validate_operation(O::OPERATION)?;
99 validate_sql(O::SQL)?;
100 let parameter_bytes = mem::size_of::<O::Parameters>();
101 if parameter_bytes == 0 {
102 return Err(QueryOptionalContractError::EmptyShape);
103 }
104 Ok(Self {
105 layout: WriteLayout {
106 operation: TypeId::of::<O>(),
107 parameter_bytes,
108 credits: WriteCreditDemand::ONE,
109 },
110 _operation: PhantomData,
111 })
112 }
113
114 pub const fn layout(&self) -> WriteLayout {
115 self.layout
116 }
117
118 pub fn invocation(self, parameters: O::Parameters) -> WriteInvocation<O> {
119 WriteInvocation {
120 layout: self.layout,
121 parameters,
122 _operation: PhantomData,
123 }
124 }
125}
126
127pub struct WriteInvocation<O: StaticWriteOperation> {
128 layout: WriteLayout,
129 parameters: O::Parameters,
130 _operation: PhantomData<fn() -> O>,
131}
132
133impl<O: StaticWriteOperation> WriteInvocation<O> {
134 pub const fn layout(&self) -> WriteLayout {
135 self.layout
136 }
137}
138
139pub struct WriteExecution<O: StaticWriteOperation> {
141 permit: DbRequestPermit,
142 invocation: WriteInvocation<O>,
143}
144
145impl<O: StaticWriteOperation> WriteExecution<O> {
146 #[doc(hidden)]
147 pub fn from_compiled_handoff(permit: DbRequestPermit, invocation: WriteInvocation<O>) -> Self {
148 Self { permit, invocation }
149 }
150
151 fn into_parts(self) -> (DbRequestPermit, WriteInvocation<O>) {
152 (self.permit, self.invocation)
153 }
154}
155
156#[derive(Clone, Copy, Debug, Eq, PartialEq)]
158pub struct ManagedWriteResult {
159 rows_affected: u64,
160 last_insert_id: u64,
161}
162
163impl ManagedWriteResult {
164 pub const fn rows_affected(self) -> u64 {
165 self.rows_affected
166 }
167
168 pub const fn last_insert_id(self) -> u64 {
169 self.last_insert_id
170 }
171}
172
173#[derive(Clone, Copy, Debug, Eq, PartialEq)]
174pub enum WriteExecutionError {
175 ConnectionUnavailable,
176 WriteFailed,
177 Cancelled,
178 Shutdown,
179 FinalizerFailed,
180}
181
182impl WriteExecutionError {
183 pub const fn code(self) -> &'static str {
184 match self {
185 Self::ConnectionUnavailable => "db.connection_unavailable",
186 Self::WriteFailed => "db.write_failed",
187 Self::Cancelled => "db.write_cancelled",
188 Self::Shutdown => "db.write_shutdown",
189 Self::FinalizerFailed => "db.finalizer_failed",
190 }
191 }
192}
193
194#[derive(Clone, Copy, Debug, Eq, PartialEq)]
196pub enum TransactionDecision {
197 Commit,
198 Rollback,
199}
200
201pub struct TransactionInvocation<O: StaticWriteOperation> {
206 write: WriteInvocation<O>,
207 decision: TransactionDecision,
208}
209
210impl<O: StaticWriteOperation> TransactionInvocation<O> {
211 pub const fn layout(&self) -> WriteLayout {
212 self.write.layout
213 }
214
215 pub const fn decision(&self) -> TransactionDecision {
216 self.decision
217 }
218}
219
220#[derive(Debug)]
221pub struct TransactionOperationProof<O: StaticWriteOperation> {
222 write: WriteOperationProof<O>,
223}
224
225impl<O: StaticWriteOperation> TransactionOperationProof<O> {
226 pub fn bind() -> Result<Self, QueryOptionalContractError> {
227 WriteOperationProof::bind().map(|write| Self { write })
228 }
229
230 pub const fn layout(&self) -> WriteLayout {
231 self.write.layout()
232 }
233
234 pub fn invocation(
235 self,
236 parameters: O::Parameters,
237 decision: TransactionDecision,
238 ) -> TransactionInvocation<O> {
239 TransactionInvocation {
240 write: self.write.invocation(parameters),
241 decision,
242 }
243 }
244}
245
246pub struct TransactionExecution<O: StaticWriteOperation> {
247 permit: DbRequestPermit,
248 invocation: TransactionInvocation<O>,
249}
250
251impl<O: StaticWriteOperation> TransactionExecution<O> {
252 #[doc(hidden)]
253 pub fn from_compiled_handoff(
254 permit: DbRequestPermit,
255 invocation: TransactionInvocation<O>,
256 ) -> Self {
257 Self { permit, invocation }
258 }
259
260 fn into_parts(self) -> (DbRequestPermit, TransactionInvocation<O>) {
261 (self.permit, self.invocation)
262 }
263}
264
265#[derive(Clone, Copy, Debug, Eq, PartialEq)]
266pub enum TransactionExecutionError {
267 ConnectionUnavailable,
268 BeginFailed,
269 WriteFailed,
270 BusinessRollback,
271 CommitFailed,
272 RollbackFailed,
273 Cancelled,
274 Shutdown,
275 FinalizerFailed,
276}
277
278impl TransactionExecutionError {
279 pub const fn code(self) -> &'static str {
280 match self {
281 Self::ConnectionUnavailable => "db.connection_unavailable",
282 Self::BeginFailed => "db.transaction_begin_failed",
283 Self::WriteFailed => "db.write_failed",
284 Self::BusinessRollback => "db.transaction_business_rollback",
285 Self::CommitFailed => "db.transaction_commit_failed",
286 Self::RollbackFailed => "db.transaction_rollback_failed",
287 Self::Cancelled => "db.transaction_cancelled",
288 Self::Shutdown => "db.transaction_shutdown",
289 Self::FinalizerFailed => "db.finalizer_failed",
290 }
291 }
292}
293
294#[derive(Clone, Copy, Debug, Eq, PartialEq)]
295pub struct WriteTransactionFinalizationProof {
296 cancel_external_io_awaits: u8,
297 shutdown_external_io_awaits: u8,
298 panic_external_io_awaits: u8,
299 single_connection: bool,
300 nested_transactions: bool,
301 releases_pool_size_before_permit: bool,
302 normal_return_requires_termination_bound: bool,
303}
304
305impl WriteTransactionFinalizationProof {
306 const PRODUCTION: Self = Self {
307 cancel_external_io_awaits: 0,
308 shutdown_external_io_awaits: 0,
309 panic_external_io_awaits: 0,
310 single_connection: true,
311 nested_transactions: false,
312 releases_pool_size_before_permit: true,
313 normal_return_requires_termination_bound: true,
314 };
315
316 pub const fn cancel_external_io_awaits(self) -> u8 {
317 self.cancel_external_io_awaits
318 }
319
320 pub const fn shutdown_external_io_awaits(self) -> u8 {
321 self.shutdown_external_io_awaits
322 }
323
324 pub const fn panic_external_io_awaits(self) -> u8 {
325 self.panic_external_io_awaits
326 }
327
328 pub const fn single_connection(self) -> bool {
329 self.single_connection
330 }
331
332 pub const fn nested_transactions(self) -> bool {
333 self.nested_transactions
334 }
335
336 pub const fn releases_pool_size_before_permit(self) -> bool {
337 self.releases_pool_size_before_permit
338 }
339
340 pub const fn normal_return_requires_termination_bound(self) -> bool {
341 self.normal_return_requires_termination_bound
342 }
343}
344
345#[doc(hidden)]
346pub const fn write_transaction_finalization_proof() -> WriteTransactionFinalizationProof {
347 WriteTransactionFinalizationProof::PRODUCTION
348}
349
350impl Database {
351 #[doc(hidden)]
352 pub async fn execute_write<O, C, S>(
353 &self,
354 execution: WriteExecution<O>,
355 cancel: C,
356 shutdown: S,
357 ) -> Result<ManagedWriteResult, WriteExecutionError>
358 where
359 O: StaticWriteOperation,
360 C: Future + Unpin + Send + 'static,
361 S: Future + Unpin + Send + 'static,
362 {
363 let pool = self.pool.clone();
364 drive_db_finalizer_with_transition(cancel, shutdown, move |transition| {
365 execute_write_with_transition(pool, execution, transition)
366 })
367 .await
368 .map_err(|_| WriteExecutionError::FinalizerFailed)
369 .and_then(|result| result)
370 }
371
372 #[doc(hidden)]
373 pub async fn execute_transaction<O, C, S>(
374 &self,
375 execution: TransactionExecution<O>,
376 cancel: C,
377 shutdown: S,
378 ) -> Result<ManagedWriteResult, TransactionExecutionError>
379 where
380 O: StaticWriteOperation,
381 C: Future + Unpin + Send + 'static,
382 S: Future + Unpin + Send + 'static,
383 {
384 let pool = self.pool.clone();
385 drive_db_finalizer_with_transition(cancel, shutdown, move |transition| {
386 execute_transaction_with_transition(pool, execution, transition)
387 })
388 .await
389 .map_err(|_| TransactionExecutionError::FinalizerFailed)
390 .and_then(|result| result)
391 }
392}
393
394async fn execute_write_with_transition<O, C, S>(
395 pool: sqlx::MySqlPool,
396 execution: WriteExecution<O>,
397 mut transition: DbQueryTransition<C, S>,
398) -> Result<
399 saddle_runtime::db_finalizer::DbFinalizingOutput<
400 Result<ManagedWriteResult, WriteExecutionError>,
401 impl Future<Output = ()> + Send + 'static,
402 >,
403 saddle_admission::AdmissionError,
404>
405where
406 O: StaticWriteOperation,
407 C: Future + Unpin + Send + 'static,
408 S: Future + Unpin + Send + 'static,
409{
410 let (permit, invocation) = execution.into_parts();
411 let mut connection = pool.try_acquire();
412 let (value, physical) = if let Some(connection) = connection.as_mut() {
413 let query = async {
414 invocation
415 .parameters
416 .bind(sqlx::query::<MySql>(O::SQL))
417 .execute(&mut **connection)
418 .await
419 };
420 let mut query = pin!(query);
421 let outcome = poll_operation(&mut transition, &permit, query.as_mut()).await;
422 match outcome {
423 Ok(DbQueryPoll::Ready(Ok(result))) => (
424 Ok(ManagedWriteResult {
425 rows_affected: result.rows_affected(),
426 last_insert_id: result.last_insert_id(),
427 }),
428 PhysicalFinalization::ReturnToPool,
429 ),
430 Ok(DbQueryPoll::Ready(Err(_))) | Err(()) => (
431 Err(WriteExecutionError::WriteFailed),
432 PhysicalFinalization::PoisonDiscard,
433 ),
434 Ok(DbQueryPoll::Transition(DbTransitionRequest::Cancel)) => (
435 Err(WriteExecutionError::Cancelled),
436 PhysicalFinalization::PoisonDiscard,
437 ),
438 Ok(DbQueryPoll::Transition(DbTransitionRequest::Shutdown)) => (
439 Err(WriteExecutionError::Shutdown),
440 PhysicalFinalization::PoisonDiscard,
441 ),
442 }
443 } else {
444 (
445 Err(WriteExecutionError::ConnectionUnavailable),
446 PhysicalFinalization::ReturnToPool,
447 )
448 };
449 let finalizer = physical_finalizer(connection, physical);
450 transition.begin_finalizing(value, permit, finalizer)
451}
452
453async fn execute_transaction_with_transition<O, C, S>(
454 pool: sqlx::MySqlPool,
455 execution: TransactionExecution<O>,
456 mut transition: DbQueryTransition<C, S>,
457) -> Result<
458 saddle_runtime::db_finalizer::DbFinalizingOutput<
459 Result<ManagedWriteResult, TransactionExecutionError>,
460 impl Future<Output = ()> + Send + 'static,
461 >,
462 saddle_admission::AdmissionError,
463>
464where
465 O: StaticWriteOperation,
466 C: Future + Unpin + Send + 'static,
467 S: Future + Unpin + Send + 'static,
468{
469 let (permit, invocation) = execution.into_parts();
470 let mut connection = pool.try_acquire();
471 let (value, physical) = if let Some(connection) = connection.as_mut() {
472 let transaction = async {
473 sqlx::query("BEGIN")
474 .execute(&mut **connection)
475 .await
476 .map_err(|_| TransactionExecutionError::BeginFailed)?;
477 let write = invocation
478 .write
479 .parameters
480 .bind(sqlx::query::<MySql>(O::SQL))
481 .execute(&mut **connection)
482 .await
483 .map_err(|_| TransactionExecutionError::WriteFailed)?;
484 let result = ManagedWriteResult {
485 rows_affected: write.rows_affected(),
486 last_insert_id: write.last_insert_id(),
487 };
488 match invocation.decision {
489 TransactionDecision::Commit => {
490 sqlx::query("COMMIT")
491 .execute(&mut **connection)
492 .await
493 .map_err(|_| TransactionExecutionError::CommitFailed)?;
494 Ok(result)
495 }
496 TransactionDecision::Rollback => {
497 sqlx::query("ROLLBACK")
498 .execute(&mut **connection)
499 .await
500 .map_err(|_| TransactionExecutionError::RollbackFailed)?;
501 Err(TransactionExecutionError::BusinessRollback)
502 }
503 }
504 };
505 let mut transaction = pin!(transaction);
506 let outcome = poll_operation(&mut transition, &permit, transaction.as_mut()).await;
507 match outcome {
508 Ok(DbQueryPoll::Ready(Ok(result))) => (Ok(result), PhysicalFinalization::ReturnToPool),
509 Ok(DbQueryPoll::Ready(Err(TransactionExecutionError::BusinessRollback))) => (
510 Err(TransactionExecutionError::BusinessRollback),
511 PhysicalFinalization::ReturnToPool,
512 ),
513 Ok(DbQueryPoll::Ready(Err(error))) => (Err(error), PhysicalFinalization::PoisonDiscard),
514 Ok(DbQueryPoll::Transition(DbTransitionRequest::Cancel)) => (
515 Err(TransactionExecutionError::Cancelled),
516 PhysicalFinalization::PoisonDiscard,
517 ),
518 Ok(DbQueryPoll::Transition(DbTransitionRequest::Shutdown)) => (
519 Err(TransactionExecutionError::Shutdown),
520 PhysicalFinalization::PoisonDiscard,
521 ),
522 Err(()) => (
523 Err(TransactionExecutionError::WriteFailed),
524 PhysicalFinalization::PoisonDiscard,
525 ),
526 }
527 } else {
528 (
529 Err(TransactionExecutionError::ConnectionUnavailable),
530 PhysicalFinalization::ReturnToPool,
531 )
532 };
533 let finalizer = physical_finalizer(connection, physical);
534 transition.begin_finalizing(value, permit, finalizer)
535}
536
537async fn poll_operation<C, S, Q>(
538 transition: &mut DbQueryTransition<C, S>,
539 permit: &DbRequestPermit,
540 mut query: std::pin::Pin<&mut Q>,
541) -> Result<DbQueryPoll<Q::Output>, ()>
542where
543 C: Future + Unpin,
544 S: Future + Unpin,
545 Q: Future,
546{
547 std::future::poll_fn(|context| {
548 match std::panic::catch_unwind(AssertUnwindSafe(|| {
549 transition.poll_query(permit, query.as_mut(), context)
550 })) {
551 Ok(Poll::Ready(output)) => Poll::Ready(Ok(output)),
552 Ok(Poll::Pending) => Poll::Pending,
553 Err(_) => Poll::Ready(Err(())),
554 }
555 })
556 .await
557}
558
559async fn physical_finalizer(
560 mut connection: Option<sqlx::pool::PoolConnection<MySql>>,
561 physical: PhysicalFinalization,
562) {
563 if let Some(mut connection) = connection.take() {
564 match physical {
565 PhysicalFinalization::ReturnToPool => connection.return_to_pool().await,
566 PhysicalFinalization::PoisonDiscard => drop(connection.detach()),
567 }
568 }
569}
570
571#[derive(Clone, Copy)]
572enum PhysicalFinalization {
573 ReturnToPool,
574 PoisonDiscard,
575}
576
577fn validate_operation(operation: &str) -> Result<(), QueryOptionalContractError> {
578 if operation.is_empty()
579 || operation.len() > 128
580 || !operation
581 .bytes()
582 .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'.' | b'_' | b'-'))
583 {
584 return Err(QueryOptionalContractError::InvalidOperation);
585 }
586 Ok(())
587}
588
589fn validate_sql(sql: &str) -> Result<(), QueryOptionalContractError> {
590 if sql.trim().is_empty() || sql.len() > 65_536 {
591 return Err(QueryOptionalContractError::InvalidSql);
592 }
593 Ok(())
594}
595
596#[cfg(test)]
597mod tests {
598 use std::{
599 env,
600 future::Future,
601 io,
602 io::{Read, Write},
603 net::{Shutdown, TcpListener, TcpStream},
604 pin::Pin,
605 sync::{Arc, Mutex},
606 task::{Context, Poll},
607 thread,
608 };
609
610 use saddle_admission::{
611 AdmissionError, DbCreditProfile, DbRouteCreditDemand, DbRouteResources, EntryIoAuditPlan,
612 EntryReadPoll, ManagedBytes, ManagedResponse, OfficialTokioEntryIoAttemptOutcome,
613 OfficialTokioRegistrationProfile, ProcessLedger, RequestMemory, ResourceConfig,
614 ResponseWritePoll,
615 };
616 use saddle_core::ComponentLifecycle;
617 use saddle_observability::{Observer, ObserverConfig};
618 use sqlx::Connection;
619
620 use super::*;
621 use crate::{
622 DatabaseConfig,
623 c6_query_optional::{DbPair, DbU64, ManagedDbField, ManagedQueryParameters, sealed},
624 };
625
626 struct Upsert;
627
628 impl StaticWriteOperation for Upsert {
629 type Parameters = ManagedQueryParameters<DbPair<DbU64, DbU64>>;
630
631 const OPERATION: &'static str = "orders.upsert";
632 const SQL: &'static str = "INSERT INTO saddle_c1_write (id, value_number) VALUES (?, ?) ON DUPLICATE KEY UPDATE value_number = VALUES(value_number)";
633 }
634
635 #[derive(Clone, Copy)]
636 struct PanicField(u8);
637
638 impl sealed::Field for PanicField {
639 fn bind<'q>(
640 &'q self,
641 _: sqlx::query::Query<'q, MySql, sqlx::mysql::MySqlArguments>,
642 ) -> sqlx::query::Query<'q, MySql, sqlx::mysql::MySqlArguments> {
643 let _ = self.0;
644 panic!("generated bind panic")
645 }
646
647 fn decode(_: &sqlx::mysql::MySqlRow, _: &mut usize) -> std::result::Result<Self, ()> {
648 unreachable!()
649 }
650 }
651
652 impl ManagedDbField for PanicField {}
653
654 struct PanicWrite;
655
656 impl StaticWriteOperation for PanicWrite {
657 type Parameters = ManagedQueryParameters<PanicField>;
658
659 const OPERATION: &'static str = "orders.panic";
660 const SQL: &'static str = "SELECT ?";
661 }
662
663 struct LockWrite;
664
665 impl StaticWriteOperation for LockWrite {
666 type Parameters = ManagedQueryParameters<DbPair<DbU64, DbU64>>;
667
668 const OPERATION: &'static str = "orders.lock";
669 const SQL: &'static str = "UPDATE saddle_c1_write SET value_number = ? WHERE id = ?";
670 }
671
672 struct MissingWrite;
673
674 impl StaticWriteOperation for MissingWrite {
675 type Parameters = ManagedQueryParameters<DbU64>;
676
677 const OPERATION: &'static str = "orders.missing";
678 const SQL: &'static str = "UPDATE saddle_c1_missing SET value_number = ? WHERE id = 1";
679 }
680
681 struct EntryConnection;
682 struct ReadyRead;
683 struct ReadyWrite;
684
685 impl EntryReadPoll<EntryConnection> for ReadyRead {
686 fn poll_read(
687 &mut self,
688 _: &mut EntryConnection,
689 memory: &RequestMemory,
690 _: &mut Context<'_>,
691 ) -> Poll<Result<ManagedBytes, AdmissionError>> {
692 Poll::Ready(memory.try_bytes(&[]))
693 }
694 }
695
696 impl ResponseWritePoll<EntryConnection> for ReadyWrite {
697 fn poll_write(
698 &mut self,
699 _: &mut EntryConnection,
700 _: &ManagedResponse,
701 _: &mut Context<'_>,
702 ) -> Poll<Result<(), AdmissionError>> {
703 Poll::Ready(Ok(()))
704 }
705 }
706
707 struct FixedSignal {
708 pending_polls: Option<u8>,
709 }
710
711 impl FixedSignal {
712 const fn pending() -> Self {
713 Self {
714 pending_polls: None,
715 }
716 }
717
718 const fn after_one_poll() -> Self {
719 Self {
720 pending_polls: Some(1),
721 }
722 }
723 }
724
725 impl Future for FixedSignal {
726 type Output = ();
727
728 fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<()> {
729 match self.pending_polls {
730 None => Poll::Pending,
731 Some(0) => Poll::Ready(()),
732 Some(remaining) => {
733 self.pending_polls = Some(remaining - 1);
734 context.waker().wake_by_ref();
735 Poll::Pending
736 }
737 }
738 }
739 }
740
741 fn resource_config(
742 registration: OfficialTokioRegistrationProfile,
743 task_reserve: usize,
744 ) -> ResourceConfig {
745 let process_state_reserve = ProcessLedger::minimum_process_state_reserve_with_waiters(1, 1)
746 .unwrap()
747 + ProcessLedger::official_tokio_state_reserve(registration).unwrap();
748 ResourceConfig {
749 managed_capacity: 4096,
750 entry_reserve: 64,
751 framework_reserve: 1024 * 1024,
752 task_reserve,
753 process_state_reserve,
754 system_estimate: 1024 * 1024,
755 safety_margin: 1024 * 1024,
756 process_limit: 4096
757 + 64
758 + 1024 * 1024
759 + task_reserve
760 + process_state_reserve
761 + 2 * 1024 * 1024,
762 max_active_requests: 1,
763 }
764 }
765
766 fn entry_plan() -> EntryIoAuditPlan {
767 EntryIoAuditPlan::locked_linux_x86_64_tokio_1_53_1(64, 64, &[]).unwrap()
768 }
769
770 async fn admitted_permit<T, F, B>(
771 task_reserve: usize,
772 run: F,
773 ) -> (T, saddle_admission::DbCreditSnapshot)
774 where
775 T: Send + 'static,
776 F: FnOnce(DbRequestPermit) -> B + Send + Unpin + 'static,
777 B: Future<Output = T> + Send + 'static,
778 {
779 let registration = OfficialTokioRegistrationProfile {
780 listener: 1,
781 transport_connections: 1,
782 runtime_fixed: 1,
783 };
784 let ledger =
785 ProcessLedger::new_with_waiters(resource_config(registration, task_reserve), 1)
786 .unwrap();
787 let runtime_domain = ledger.prepare_official_tokio_domain(registration).unwrap();
788 let allocation = ledger
789 .prepare_process_allocation_profile(usize::MAX)
790 .unwrap();
791 let db_domain = ledger
792 .prepare_db_domain(DbCreditProfile {
793 connections: 1,
794 operations: 1,
795 })
796 .unwrap();
797 let demand = DbRouteCreditDemand::new(1, 1).unwrap();
798 let result = Arc::new(Mutex::new(None));
799 let result_for_task = result.clone();
800 let outcome = ledger.attempt_official_tokio_entry_io(
801 &runtime_domain,
802 DbRouteResources::required(&db_domain, demand),
803 4096,
804 task_reserve,
805 entry_plan(),
806 entry_plan(),
807 |_| (EntryConnection, ReadyRead, ReadyWrite),
808 move |_, permit, memory| {
809 let operation = run(permit.unwrap());
810 let response = memory.try_response(&[]).unwrap();
811 async move {
812 *result_for_task.lock().unwrap() = Some(operation.await);
813 response
814 }
815 },
816 );
817 let ready = match outcome {
818 OfficialTokioEntryIoAttemptOutcome::Ready(ready) => ready,
819 _ => panic!("fixed DB resources must admit"),
820 };
821 let (envelope, task_slot) = ready.into_runtime_parts();
822 let envelope_size = std::mem::size_of_val(&envelope);
823 let envelope_align = std::mem::align_of_val(&envelope);
824 assert!(envelope_size <= task_reserve);
825 assert!(envelope_align.is_power_of_two());
826 eprintln!(
827 "C1 envelope_size={envelope_size} envelope_align={envelope_align} tested_reserve={task_reserve}"
828 );
829 tokio::spawn(envelope).await.unwrap().unwrap();
830 drop(task_slot);
831 let snapshot = db_domain.snapshot().unwrap();
832 drop(db_domain);
833 drop(runtime_domain);
834 assert!(!allocation.finish().unwrap().breached);
835 assert_eq!(ledger.try_shutdown().unwrap().active_accounts, 0);
836 let result = Arc::try_unwrap(result)
837 .ok()
838 .unwrap()
839 .into_inner()
840 .unwrap()
841 .unwrap();
842 (result, snapshot)
843 }
844
845 async fn run_write<O, C, S>(
846 database: Database,
847 invocation: WriteInvocation<O>,
848 cancel: C,
849 shutdown: S,
850 ) -> Result<ManagedWriteResult, WriteExecutionError>
851 where
852 O: StaticWriteOperation,
853 O::Parameters: Unpin,
854 C: Future<Output = ()> + Unpin + Send + 'static,
855 S: Future<Output = ()> + Unpin + Send + 'static,
856 {
857 let (result, snapshot) = admitted_permit(1024 * 1024, move |permit| {
858 let execution = WriteExecution::from_compiled_handoff(permit, invocation);
859 async move { database.execute_write(execution, cancel, shutdown).await }
860 })
861 .await;
862 assert_eq!(snapshot.connections_in_use, 0);
863 assert_eq!(snapshot.operations_in_use, 0);
864 result
865 }
866
867 async fn run_transaction<O, C, S>(
868 database: Database,
869 invocation: TransactionInvocation<O>,
870 cancel: C,
871 shutdown: S,
872 ) -> Result<ManagedWriteResult, TransactionExecutionError>
873 where
874 O: StaticWriteOperation,
875 O::Parameters: Unpin,
876 C: Future<Output = ()> + Unpin + Send + 'static,
877 S: Future<Output = ()> + Unpin + Send + 'static,
878 {
879 let (result, snapshot) = admitted_permit(1024 * 1024, move |permit| {
880 let execution = TransactionExecution::from_compiled_handoff(permit, invocation);
881 async move {
882 database
883 .execute_transaction(execution, cancel, shutdown)
884 .await
885 }
886 })
887 .await;
888 assert_eq!(snapshot.connections_in_use, 0);
889 assert_eq!(snapshot.operations_in_use, 0);
890 result
891 }
892
893 async fn database(url: &str) -> Database {
894 Database::connect(
895 DatabaseConfig::new(url).max_connections(1),
896 Observer::with_writer(ObserverConfig::default(), io::sink()).unwrap(),
897 )
898 .await
899 .unwrap()
900 }
901
902 fn parameters(id: u64, value: u64) -> ManagedQueryParameters<DbPair<DbU64, DbU64>> {
903 ManagedQueryParameters(DbPair(DbU64(id), DbU64(value)))
904 }
905
906 fn commit_response_cut_proxy(database_url: &str) -> (String, thread::JoinHandle<()>) {
907 let backend = database_url
908 .split_once('@')
909 .and_then(|(_, suffix)| suffix.split_once('/'))
910 .map(|(authority, _)| authority)
911 .expect("test database URL has an authority");
912 let listener = TcpListener::bind("127.0.0.1:0").unwrap();
913 let proxy = listener.local_addr().unwrap();
914 let proxy_url = database_url.replacen(backend, &proxy.to_string(), 1);
915 let backend = backend.to_owned();
916 let worker = thread::spawn(move || {
917 loop {
918 let (client, _) = listener.accept().unwrap();
919 let server = TcpStream::connect(&backend).unwrap();
920 let mut request_client = client.try_clone().unwrap();
921 let mut request_server = server.try_clone().unwrap();
922 let response_client = client.try_clone().unwrap();
923 let response_server = server.try_clone().unwrap();
924 let response = thread::spawn(move || {
925 let _ = io::copy(&mut &response_server, &mut &response_client);
926 });
927 let mut cut_commit_response = false;
928 let mut buffer = [0_u8; 4096];
929 loop {
930 let count = request_client.read(&mut buffer).unwrap();
931 if count == 0 {
932 break;
933 }
934 request_server.write_all(&buffer[..count]).unwrap();
935 if buffer[..count]
936 .windows(b"COMMIT".len())
937 .any(|window| window.eq_ignore_ascii_case(b"COMMIT"))
938 {
939 cut_commit_response = true;
942 let _ = request_client.shutdown(Shutdown::Both);
943 let _ = request_server.shutdown(Shutdown::Both);
944 break;
945 }
946 }
947 let _ = response.join();
948 if cut_commit_response {
949 break;
950 }
951 }
952 });
953 (proxy_url, worker)
954 }
955
956 #[test]
957 fn fixed_contract_has_one_credit_and_no_nested_transaction() {
958 let write = WriteOperationProof::<Upsert>::bind().unwrap();
959 assert_eq!(write.layout().credits(), WriteCreditDemand::ONE);
960 let transaction = TransactionOperationProof::<Upsert>::bind().unwrap();
961 assert_eq!(transaction.layout().credits(), WriteCreditDemand::ONE);
962 let finalization = write_transaction_finalization_proof();
963 assert!(finalization.single_connection());
964 assert!(!finalization.nested_transactions());
965 assert_eq!(finalization.cancel_external_io_awaits(), 0);
966 assert!(finalization.releases_pool_size_before_permit());
967 }
968
969 #[tokio::test]
970 async fn real_mariadb_write_commit_rollback_cancel_shutdown_and_panic_close() {
971 let Ok(url) = env::var("SADDLE_TEST_DATABASE_URL") else {
972 eprintln!("skipping C1 write/transaction: SADDLE_TEST_DATABASE_URL is not set");
973 return;
974 };
975 let setup = database(&url).await;
976 sqlx::query(
977 "CREATE TABLE IF NOT EXISTS saddle_c1_write (id BIGINT UNSIGNED PRIMARY KEY, value_number BIGINT UNSIGNED NOT NULL)",
978 )
979 .execute(&setup.pool)
980 .await
981 .unwrap();
982 sqlx::query("DELETE FROM saddle_c1_write")
983 .execute(&setup.pool)
984 .await
985 .unwrap();
986 setup.shutdown().await.unwrap();
987
988 let write_db = database(&url).await;
989 let result = run_write(
990 write_db.clone(),
991 WriteOperationProof::<Upsert>::bind()
992 .unwrap()
993 .invocation(parameters(1, 10)),
994 FixedSignal::pending(),
995 FixedSignal::pending(),
996 )
997 .await
998 .unwrap();
999 assert_eq!(result.rows_affected(), 1);
1000 assert_eq!((write_db.pool.size(), write_db.pool.num_idle()), (1, 1));
1001 write_db.shutdown().await.unwrap();
1002
1003 let error_db = database(&url).await;
1004 let error = run_write(
1005 error_db.clone(),
1006 WriteOperationProof::<MissingWrite>::bind()
1007 .unwrap()
1008 .invocation(ManagedQueryParameters(DbU64(10))),
1009 FixedSignal::pending(),
1010 FixedSignal::pending(),
1011 )
1012 .await
1013 .unwrap_err();
1014 assert_eq!(error, WriteExecutionError::WriteFailed);
1015 assert_eq!((error_db.pool.size(), error_db.pool.num_idle()), (0, 0));
1016 error_db.shutdown().await.unwrap();
1017
1018 let commit_db = database(&url).await;
1019 run_transaction(
1020 commit_db.clone(),
1021 TransactionOperationProof::<Upsert>::bind()
1022 .unwrap()
1023 .invocation(parameters(2, 20), TransactionDecision::Commit),
1024 FixedSignal::pending(),
1025 FixedSignal::pending(),
1026 )
1027 .await
1028 .unwrap();
1029 assert_eq!((commit_db.pool.size(), commit_db.pool.num_idle()), (1, 1));
1030 commit_db.shutdown().await.unwrap();
1031
1032 let (ambiguous_url, proxy) = commit_response_cut_proxy(&url);
1033 let ambiguous_db = database(&ambiguous_url).await;
1034 let ambiguous = run_transaction(
1035 ambiguous_db.clone(),
1036 TransactionOperationProof::<Upsert>::bind()
1037 .unwrap()
1038 .invocation(parameters(4, 40), TransactionDecision::Commit),
1039 FixedSignal::pending(),
1040 FixedSignal::pending(),
1041 )
1042 .await
1043 .unwrap_err();
1044 assert_eq!(ambiguous, TransactionExecutionError::CommitFailed);
1045 assert_eq!(
1046 (ambiguous_db.pool.size(), ambiguous_db.pool.num_idle()),
1047 (0, 0)
1048 );
1049 ambiguous_db.shutdown().await.unwrap();
1050 proxy.join().unwrap();
1051
1052 let rollback_db = database(&url).await;
1053 let rollback = run_transaction(
1054 rollback_db.clone(),
1055 TransactionOperationProof::<Upsert>::bind()
1056 .unwrap()
1057 .invocation(parameters(3, 30), TransactionDecision::Rollback),
1058 FixedSignal::pending(),
1059 FixedSignal::pending(),
1060 )
1061 .await
1062 .unwrap_err();
1063 assert_eq!(rollback, TransactionExecutionError::BusinessRollback);
1064 assert_eq!(
1065 (rollback_db.pool.size(), rollback_db.pool.num_idle()),
1066 (1, 1)
1067 );
1068 rollback_db.shutdown().await.unwrap();
1069
1070 let panic_db = database(&url).await;
1071 let panic = run_write(
1072 panic_db.clone(),
1073 WriteOperationProof::<PanicWrite>::bind()
1074 .unwrap()
1075 .invocation(ManagedQueryParameters(PanicField(0))),
1076 FixedSignal::pending(),
1077 FixedSignal::pending(),
1078 )
1079 .await
1080 .unwrap_err();
1081 assert_eq!(panic, WriteExecutionError::WriteFailed);
1082 assert_eq!((panic_db.pool.size(), panic_db.pool.num_idle()), (0, 0));
1083 panic_db.shutdown().await.unwrap();
1084
1085 for (cancel, shutdown, expected) in [
1086 (
1087 FixedSignal::after_one_poll(),
1088 FixedSignal::pending(),
1089 TransactionExecutionError::Cancelled,
1090 ),
1091 (
1092 FixedSignal::pending(),
1093 FixedSignal::after_one_poll(),
1094 TransactionExecutionError::Shutdown,
1095 ),
1096 ] {
1097 let blocked = database(&url).await;
1098 let mut blocker = sqlx::mysql::MySqlConnection::connect(&url).await.unwrap();
1099 sqlx::query("BEGIN").execute(&mut blocker).await.unwrap();
1100 sqlx::query("UPDATE saddle_c1_write SET value_number = 99 WHERE id = 1")
1101 .execute(&mut blocker)
1102 .await
1103 .unwrap();
1104 let error = run_transaction(
1105 blocked.clone(),
1106 TransactionOperationProof::<LockWrite>::bind()
1107 .unwrap()
1108 .invocation(parameters(100, 1), TransactionDecision::Commit),
1109 cancel,
1110 shutdown,
1111 )
1112 .await
1113 .unwrap_err();
1114 assert_eq!(error, expected);
1115 assert_eq!((blocked.pool.size(), blocked.pool.num_idle()), (0, 0));
1116 sqlx::query("ROLLBACK").execute(&mut blocker).await.unwrap();
1117 blocked.shutdown().await.unwrap();
1118 }
1119
1120 let verify = database(&url).await;
1121 let rows: Vec<(u64, u64)> =
1122 sqlx::query_as("SELECT id, value_number FROM saddle_c1_write ORDER BY id")
1123 .fetch_all(&verify.pool)
1124 .await
1125 .unwrap();
1126 assert_eq!(rows, vec![(1, 10), (2, 20)]);
1127 verify.shutdown().await.unwrap();
1128 }
1129}