|
| 1 | +use std::sync::Arc; |
1 | 2 | use std::time::Duration;
|
2 | 3 |
|
3 | 4 | use futures::StreamExt;
|
4 |
| -use insta::assert_json_snapshot; |
| 5 | +use insta::{assert_json_snapshot, assert_snapshot}; |
5 | 6 | use libsql::{params, Database};
|
| 7 | +use libsql_server::config::UserApiConfig; |
6 | 8 | use libsql_server::hrana_proto::{Batch, BatchStep, Stmt};
|
| 9 | +use tokio::sync::Notify; |
7 | 10 |
|
8 | 11 | use crate::common::http::Client;
|
9 |
| -use crate::common::net::TurmoilConnector; |
| 12 | +use crate::common::net::{init_tracing, SimServer as _, TestServer, TurmoilConnector}; |
10 | 13 |
|
11 | 14 | #[test]
|
12 | 15 | fn sample_request() {
|
@@ -402,3 +405,179 @@ fn test_simulate_vector_index_load_from_dump() {
|
402 | 405 |
|
403 | 406 | sim.run().unwrap();
|
404 | 407 | }
|
| 408 | + |
| 409 | +#[test] |
| 410 | +fn server_restart_query_execute_invalid_baton() { |
| 411 | + server_restart(|notify, db| async move { |
| 412 | + let conn = db.connect().unwrap(); |
| 413 | + |
| 414 | + conn.query("select 1;", ()).await.unwrap(); |
| 415 | + |
| 416 | + notify.notify_waiters(); |
| 417 | + tokio::time::sleep(std::time::Duration::from_secs(1)).await; |
| 418 | + |
| 419 | + let err = conn.execute("select 1;", ()).await.unwrap_err(); |
| 420 | + assert_snapshot!(err); |
| 421 | + |
| 422 | + Ok(()) |
| 423 | + }); |
| 424 | +} |
| 425 | + |
| 426 | +#[test] |
| 427 | +fn server_restart_txn_execute_execute_invalid_baton() { |
| 428 | + server_restart(|notify, db| async move { |
| 429 | + let conn = db.connect().unwrap(); |
| 430 | + |
| 431 | + let txn = conn.transaction().await.unwrap(); |
| 432 | + |
| 433 | + txn.execute("select 1;", ()).await.unwrap(); |
| 434 | + |
| 435 | + notify.notify_waiters(); |
| 436 | + |
| 437 | + tokio::time::sleep(std::time::Duration::from_secs(1)).await; |
| 438 | + |
| 439 | + let err = txn.execute("select 1;", ()).await.unwrap_err(); |
| 440 | + |
| 441 | + assert_snapshot!(err); |
| 442 | + |
| 443 | + Ok(()) |
| 444 | + }); |
| 445 | +} |
| 446 | + |
| 447 | +#[test] |
| 448 | +fn server_restart_txn_query_execute_invalid_baton() { |
| 449 | + server_restart(|notify, db| async move { |
| 450 | + let conn = db.connect().unwrap(); |
| 451 | + |
| 452 | + let txn = conn.transaction().await.unwrap(); |
| 453 | + |
| 454 | + txn.query("select 1;", ()).await.unwrap(); |
| 455 | + |
| 456 | + notify.notify_waiters(); |
| 457 | + |
| 458 | + tokio::time::sleep(std::time::Duration::from_secs(1)).await; |
| 459 | + |
| 460 | + let err = txn.execute("select 1;", ()).await.unwrap_err(); |
| 461 | + |
| 462 | + assert_snapshot!(err); |
| 463 | + |
| 464 | + Ok(()) |
| 465 | + }); |
| 466 | +} |
| 467 | + |
| 468 | +#[test] |
| 469 | +fn server_restart_txn_execute_query_invalid_baton() { |
| 470 | + server_restart(|notify, db| async move { |
| 471 | + let conn = db.connect().unwrap(); |
| 472 | + |
| 473 | + let txn = conn.transaction().await.unwrap(); |
| 474 | + |
| 475 | + txn.execute("select 1;", ()).await.unwrap(); |
| 476 | + |
| 477 | + notify.notify_waiters(); |
| 478 | + |
| 479 | + tokio::time::sleep(std::time::Duration::from_secs(1)).await; |
| 480 | + |
| 481 | + let err = txn.query("select 1;", ()).await.unwrap_err(); |
| 482 | + |
| 483 | + assert_snapshot!(err); |
| 484 | + |
| 485 | + Ok(()) |
| 486 | + }); |
| 487 | +} |
| 488 | + |
| 489 | +#[test] |
| 490 | +fn server_restart_execute_query() { |
| 491 | + server_restart(|notify, db| async move { |
| 492 | + let conn = db.connect().unwrap(); |
| 493 | + conn.execute("select 1;", ()).await.unwrap(); |
| 494 | + |
| 495 | + notify.notify_waiters(); |
| 496 | + |
| 497 | + tokio::time::sleep(std::time::Duration::from_secs(1)).await; |
| 498 | + |
| 499 | + conn.query("select 1;", ()).await.unwrap(); |
| 500 | + |
| 501 | + Ok(()) |
| 502 | + }); |
| 503 | +} |
| 504 | + |
| 505 | +#[test] |
| 506 | +fn server_timeout() { |
| 507 | + server_restart(|_notify, db| async move { |
| 508 | + let conn = db.connect().unwrap(); |
| 509 | + conn.query("select 1;", ()).await.unwrap(); |
| 510 | + |
| 511 | + tokio::time::sleep(std::time::Duration::from_secs(20)).await; |
| 512 | + |
| 513 | + let err = conn.execute("select 1;", ()).await.unwrap_err(); |
| 514 | + |
| 515 | + assert_snapshot!(err); |
| 516 | + |
| 517 | + let conn = db.connect().unwrap(); |
| 518 | + conn.execute("select 1;", ()).await.unwrap(); |
| 519 | + |
| 520 | + tokio::time::sleep(std::time::Duration::from_secs(20)).await; |
| 521 | + |
| 522 | + conn.execute("select 1;", ()).await.unwrap(); |
| 523 | + |
| 524 | + Ok(()) |
| 525 | + }); |
| 526 | +} |
| 527 | + |
| 528 | +#[track_caller] |
| 529 | +fn server_restart<F, Fut>(f: F) |
| 530 | +where |
| 531 | + F: Fn(Arc<Notify>, Database) -> Fut + 'static, |
| 532 | + Fut: std::future::Future<Output = Result<(), libsql::Error>> + 'static, |
| 533 | +{ |
| 534 | + let mut sim = turmoil::Builder::new() |
| 535 | + .simulation_duration(Duration::from_secs(1000)) |
| 536 | + .build(); |
| 537 | + |
| 538 | + init_tracing(); |
| 539 | + |
| 540 | + let notify = Arc::new(Notify::new()); |
| 541 | + |
| 542 | + let notify_clone = notify.clone(); |
| 543 | + |
| 544 | + sim.host("primary", move || { |
| 545 | + let notify = notify.clone(); |
| 546 | + async move { |
| 547 | + let tmp = tempfile::tempdir()?; |
| 548 | + |
| 549 | + let make_server = || TestServer { |
| 550 | + path: tmp.path().to_owned().into(), |
| 551 | + user_api_config: UserApiConfig { |
| 552 | + hrana_ws_acceptor: None, |
| 553 | + ..Default::default() |
| 554 | + }, |
| 555 | + ..Default::default() |
| 556 | + }; |
| 557 | + |
| 558 | + let server = make_server(); |
| 559 | + |
| 560 | + tokio::select! { |
| 561 | + res = server.start_sim(8080) => { |
| 562 | + res.unwrap() |
| 563 | + } |
| 564 | + _ = notify.notified() => (), |
| 565 | + } |
| 566 | + |
| 567 | + let server = make_server(); |
| 568 | + server.start_sim(8080).await.unwrap(); |
| 569 | + |
| 570 | + Ok(()) |
| 571 | + } |
| 572 | + }); |
| 573 | + |
| 574 | + sim.client("client", async move { |
| 575 | + let db = Database::open_remote_with_connector("http://primary:8080", "", TurmoilConnector)?; |
| 576 | + |
| 577 | + f(notify_clone, db).await.unwrap(); |
| 578 | + |
| 579 | + Ok(()) |
| 580 | + }); |
| 581 | + |
| 582 | + sim.run().unwrap(); |
| 583 | +} |
0 commit comments