Fix fullscan hang when losing connection (#231)

We have detected an annoying issue that stops the fullscan without a way
to resume when the connection is lost, losing hours of index building
time. This PR fixes the issue by adding a retry loop to the fullscan.

Fixes: VECTOR-193

Author: ewienik

crates/validator/src/common.rs

@@ -10,6 +10,7 @@ use crate::vector_store_cluster::VectorStoreClusterExt;
 use httpclient::HttpClient;
 use scylla::client::session::Session;
 use scylla::client::session_builder::SessionBuilder;
+use std::net::Ipv4Addr;
 use std::sync::Arc;
 use std::time::Duration;
 use tokio::time;
@@ -23,21 +24,29 @@ pub(crate) const DB_PORT: u16 = 9042;
 pub(crate) const VS_OCTET: u8 = 1;
 pub(crate) const DB_OCTET: u8 = 2;
 
+pub(crate) async fn get_default_vs_url(actors: &TestActors) -> String {
+    format!(
+        "http://{}.{}:{}",
+        VS_NAME,
+        actors.dns.domain().await,
+        VS_PORT
+    )
+}
+
+pub(crate) fn get_default_db_ip(actors: &TestActors) -> Ipv4Addr {
+    actors.services_subnet.ip(DB_OCTET)
+}
+
 pub(crate) async fn init(actors: TestActors) {
     info!("started");
 
     let vs_ip = actors.services_subnet.ip(VS_OCTET);
 
     actors.dns.upsert(VS_NAME.to_string(), vs_ip).await;
 
-    let vs_url = format!(
-        "http://{}.{}:{}",
-        VS_NAME,
-        actors.dns.domain().await,
-        VS_PORT
-    );
+    let vs_url = get_default_vs_url(&actors).await;
 
-    let db_ip = actors.services_subnet.ip(DB_OCTET);
+    let db_ip = get_default_db_ip(&actors);
 
     actors.db.start(vs_url, db_ip, None).await;
     assert!(actors.db.wait_for_ready().await);
@@ -59,7 +68,7 @@ pub(crate) async fn cleanup(actors: TestActors) {
     info!("finished");
 }
 
-pub(crate) async fn prepare_connection(actors: TestActors) -> (Arc<Session>, HttpClient) {
+pub(crate) async fn prepare_connection(actors: &TestActors) -> (Arc<Session>, HttpClient) {
     let session = Arc::new(
         SessionBuilder::new()
             .known_node(actors.services_subnet.ip(DB_OCTET).to_string())

crates/validator/src/scylla_cluster.rs

@@ -4,6 +4,7 @@
  */
 
 use std::net::Ipv4Addr;
+use std::path::Path;
 use std::path::PathBuf;
 use std::process::Stdio;
 use std::time::Duration;
@@ -31,6 +32,11 @@ pub(crate) enum ScyllaCluster {
         tx: oneshot::Sender<bool>,
     },
     Stop,
+    Up {
+        vs_uri: String,
+        conf: Option<Vec<u8>>,
+    },
+    Down,
 }
 
 pub(crate) trait ScyllaClusterExt {
@@ -45,6 +51,12 @@ pub(crate) trait ScyllaClusterExt {
 
     /// Waits for the ScyllaDB cluster to be ready.
     async fn wait_for_ready(&self) -> bool;
+
+    /// Starts a paused instance back again.
+    async fn up(&self, vs_uri: String, conf: Option<Vec<u8>>);
+
+    /// Pauses an instance
+    async fn down(&self);
 }
 
 impl ScyllaClusterExt for mpsc::Sender<ScyllaCluster> {
@@ -81,6 +93,18 @@ impl ScyllaClusterExt for mpsc::Sender<ScyllaCluster> {
         rx.await
             .expect("ScyllaClusterExt::wait_for_ready: internal actor should send response")
     }
+
+    async fn up(&self, vs_uri: String, conf: Option<Vec<u8>>) {
+        self.send(ScyllaCluster::Up { vs_uri, conf })
+            .await
+            .expect("ScyllaClusterExt::up: internal actor should receive request")
+    }
+
+    async fn down(&self) {
+        self.send(ScyllaCluster::Down)
+            .await
+            .expect("ScyllaClusterExt::down: internal actor should receive request")
+    }
 }
 
 pub(crate) async fn new(
@@ -175,13 +199,26 @@ async fn process(msg: ScyllaCluster, state: &mut State) {
             tx.send(wait_for_ready(state).await)
                 .expect("process ScyllaCluster::WaitForReady: failed to send a response");
         }
+
+        ScyllaCluster::Up { vs_uri, conf } => {
+            up(vs_uri, conf, state).await;
+        }
+
+        ScyllaCluster::Down => {
+            down(state).await;
+        }
     }
 }
 
-async fn start(vs_uri: String, db_ip: Ipv4Addr, conf: Option<Vec<u8>>, state: &mut State) {
-    let workdir = TempDir::new().expect("start: failed to create temporary directory for scylladb");
+async fn run_cluster(
+    vs_uri: &String,
+    db_ip: &Ipv4Addr,
+    conf: &Option<Vec<u8>>,
+    path: &Path,
+    state: &mut State,
+) {
     let conf = if let Some(conf) = conf {
-        let conf_path = workdir.path().join("scylla.conf");
+        let conf_path = path.join("scylla.conf");
         fs::write(&conf_path, conf)
             .await
             .expect("start: failed to write scylla config");
@@ -198,7 +235,7 @@ async fn start(vs_uri: String, db_ip: Ipv4Addr, conf: Option<Vec<u8>>, state: &m
             .arg("--options-file")
             .arg(&conf)
             .arg("--workdir")
-            .arg(workdir.path())
+            .arg(path)
             .arg("--listen-address")
             .arg(db_ip.to_string())
             .arg("--rpc-address")
@@ -216,6 +253,11 @@ async fn start(vs_uri: String, db_ip: Ipv4Addr, conf: Option<Vec<u8>>, state: &m
             .spawn()
             .expect("start: failed to spawn scylladb"),
     );
+}
+
+async fn start(vs_uri: String, db_ip: Ipv4Addr, conf: Option<Vec<u8>>, state: &mut State) {
+    let workdir = TempDir::new().expect("start: failed to create temporary directory for scylladb");
+    run_cluster(&vs_uri, &db_ip, &conf, workdir.path(), state).await;
     state.workdir = Some(workdir);
     state.db_ip = Some(db_ip);
 }
@@ -236,6 +278,20 @@ async fn stop(state: &mut State) {
     state.db_ip = None;
 }
 
+async fn down(state: &mut State) {
+    let Some(mut child) = state.child.take() else {
+        return;
+    };
+    child
+        .start_kill()
+        .expect("stop: failed to send SIGTERM to scylladb process");
+    child
+        .wait()
+        .await
+        .expect("stop: failed to wait for scylladb process to exit");
+    state.child = None;
+}
+
 /// Waits for ScyllaDB to be ready by checking the nodetool status.
 async fn wait_for_ready(state: &State) -> bool {
     let Some(db_ip) = state.db_ip else {
@@ -262,3 +318,15 @@ async fn wait_for_ready(state: &State) -> bool {
         time::sleep(Duration::from_millis(100)).await;
     }
 }
+
+async fn up(vs_uri: String, conf: Option<Vec<u8>>, state: &mut State) {
+    let db_ip = state.db_ip.expect("State should have DB IP");
+    let path = state
+        .workdir
+        .as_ref()
+        .expect("State should have workdir")
+        .path()
+        .to_path_buf();
+
+    run_cluster(&vs_uri, &db_ip, &conf, &path, state).await;
+}

crates/validator/src/tests/crud.rs

@@ -28,7 +28,7 @@ pub(crate) async fn new() -> TestCase {
 async fn simple_create_drop_index(actors: TestActors) {
     info!("started");
 
-    let (session, client) = prepare_connection(actors).await;
+    let (session, client) = prepare_connection(&actors).await;
 
     session.query_unpaged(
         "CREATE KEYSPACE ks WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 1}",
@@ -83,7 +83,7 @@ async fn simple_create_drop_index(actors: TestActors) {
 async fn simple_create_drop_multiple_indexes(actors: TestActors) {
     info!("started");
 
-    let (session, client) = prepare_connection(actors).await;
+    let (session, client) = prepare_connection(&actors).await;
 
     // Create keyspace
     // Different keyspace name have to be used until the issue VECTOR-213 is fixed.

crates/validator/src/tests/full_scan.rs

@@ -23,7 +23,7 @@ pub(crate) async fn new() -> TestCase {
 async fn full_scan_is_completed_when_responding_to_messages_concurrently(actors: TestActors) {
     info!("started");
 
-    let (session, client) = prepare_connection(actors).await;
+    let (session, client) = prepare_connection(&actors).await;
 
     session.query_unpaged(
         "CREATE KEYSPACE ks WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 1}",

crates/validator/src/tests/mod.rs

@@ -5,6 +5,7 @@
 
 mod crud;
 mod full_scan;
+mod reconnect;
 mod serde;
 
 use crate::ServicesSubnet;
@@ -197,6 +198,7 @@ pub(crate) async fn register() -> Vec<(String, TestCase)> {
         ("crud", crud::new().await),
         ("full_scan", full_scan::new().await),
         ("serde", serde::new().await),
+        ("reconnect", reconnect::new().await),
     ]
     .into_iter()
     .map(|(name, test_case)| (name.to_string(), test_case))

crates/validator/src/tests/reconnect.rs

@@ -0,0 +1,120 @@
+/*
+ * Copyright 2025-present ScyllaDB
+ * SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.0
+ */
+
+use crate::common::*;
+use crate::scylla_cluster::ScyllaClusterExt;
+use crate::tests::*;
+use std::time::Duration;
+use tokio::time::sleep;
+use tracing::info;
+
+pub(crate) async fn new() -> TestCase {
+    let timeout = Duration::from_secs(30);
+    TestCase::empty()
+        .with_init(timeout, init)
+        .with_cleanup(timeout, cleanup)
+        .with_test(
+            "reconnect_doesnt_break_fullscan",
+            timeout,
+            reconnect_doesnt_break_fullscan,
+        )
+}
+
+async fn reconnect_doesnt_break_fullscan(actors: TestActors) {
+    info!("started");
+
+    let (session, client) = prepare_connection(&actors).await;
+
+    session.query_unpaged(
+        "CREATE KEYSPACE ks WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 1}",
+        (),
+    ).await.expect("failed to create a keyspace");
+
+    session
+        .use_keyspace("ks", false)
+        .await
+        .expect("failed to use a keyspace");
+
+    session
+        .query_unpaged(
+            "
+            CREATE TABLE tbl (id INT PRIMARY KEY, embedding VECTOR<FLOAT, 3>)
+            WITH cdc = {'enabled': true }
+            ",
+            (),
+        )
+        .await
+        .expect("failed to create a table");
+
+    let stmt = session
+        .prepare("INSERT INTO tbl (id, embedding) VALUES (?, [1.0, 2.0, 3.0])")
+        .await
+        .expect("failed to prepare a statement");
+
+    for id in 0..1000 {
+        session
+            .execute_unpaged(&stmt, (id,))
+            .await
+            .expect("failed to insert a row");
+    }
+
+    session
+        .query_unpaged(
+            "CREATE INDEX idx ON tbl(embedding) USING 'vector_index'",
+            (),
+        )
+        .await
+        .expect("failed to create an index");
+
+    while client.indexes().await.is_empty() {}
+    let indexes = client.indexes().await;
+    assert_eq!(indexes.len(), 1);
+    let index = &indexes[0];
+    assert_eq!(index.keyspace.as_ref(), "ks");
+    assert_eq!(index.index.as_ref(), "idx");
+
+    let result = session
+        .query_unpaged(
+            "SELECT * FROM tbl ORDER BY embedding ANN OF [1.0, 2.0, 3.0] LIMIT 1",
+            (),
+        )
+        .await;
+
+    match &result {
+        Err(e) if format!("{e:?}").contains("503 Service Unavailable") => {}
+        _ => panic!("Expected SERVICE_UNAVAILABLE error, got: {result:?}"),
+    }
+
+    actors.db.down().await;
+
+    sleep(Duration::from_secs(1)).await;
+    let count = client.count(&index.keyspace, &index.index).await;
+    assert!(count.is_some() && count.unwrap() < 1000);
+    actors.db.up(get_default_vs_url(&actors).await, None).await;
+
+    assert!(actors.db.wait_for_ready().await);
+
+    wait_for(
+        || async {
+            session
+                .query_unpaged(
+                    "SELECT * FROM tbl ORDER BY embedding ANN OF [1.0, 2.0, 3.0] LIMIT 1",
+                    (),
+                )
+                .await
+                .is_ok()
+        },
+        "Waiting for index build",
+        Duration::from_secs(10),
+    )
+    .await;
+
+    session
+        .query_unpaged("DROP KEYSPACE ks", ())
+        .await
+        .expect("failed to drop a keyspace");
+
+    info!("finished");
+}

crates/validator/src/tests/serde.rs

@@ -19,7 +19,7 @@ pub(crate) async fn new() -> TestCase {
 }
 
 async fn test_serialization_deserialization_all_types(actors: TestActors) {
-    let (session, _client) = crate::common::prepare_connection(actors).await;
+    let (session, _client) = crate::common::prepare_connection(&actors).await;
 
     let cases = vec![
         ("ascii", "'random_text'"),