fix: start server before DB download to prevent health check timeout

Decouple HTTP server startup from database initialization so the
/health endpoint responds immediately. On first deploy with an empty
volume, the DB download blocked the server from binding, causing
Docker health checks to time out after 60s.

src/db.rs

@@ -7,22 +7,32 @@ use maxminddb::Reader;
 use serde::Deserialize;
 
 use crate::models::{
-    IpInfo, MmdbAsn, MmdbCity, MmdbCountry, MmdbProxy, MmdbRecord, ProxyInfo, from_mmdb_record,
+    from_mmdb_record, IpInfo, MmdbAsn, MmdbCity, MmdbCountry, MmdbProxy, MmdbRecord, ProxyInfo,
 };
 
 pub type DbReader = Reader<Vec<u8>>;
-pub type SharedDb = Arc<ArcSwap<DbReader>>;
+pub type SharedDb = Arc<ArcSwap<Option<DbReader>>>;
 
 pub fn load_db(path: &Path) -> Result<DbReader, maxminddb::MaxMindDbError> {
     Reader::open_readfile(path)
 }
 
 pub fn new_shared(reader: DbReader) -> SharedDb {
-    Arc::new(ArcSwap::from_pointee(reader))
+    Arc::new(ArcSwap::from_pointee(Some(reader)))
 }
 
-pub fn lookup(db: &DbReader, ip: IpAddr) -> Option<IpInfo> {
-    let result = db.lookup(ip).ok()?;
+pub fn new_empty() -> SharedDb {
+    Arc::new(ArcSwap::from_pointee(None))
+}
+
+pub fn is_ready(db: &SharedDb) -> bool {
+    db.load().is_some()
+}
+
+pub fn lookup(db: &SharedDb, ip: IpAddr) -> Option<IpInfo> {
+    let guard = db.load();
+    let reader = Option::as_ref(&guard)?;
+    let result = reader.lookup(ip).ok()?;
     let record: MmdbRecord<'_> = result.decode().ok()??;
     Some(from_mmdb_record(ip, &record))
 }
@@ -54,8 +64,10 @@ struct CityOnly<'a> {
     city: MmdbCity<'a>,
 }
 
-pub fn lookup_proxy(db: &DbReader, ip: IpAddr) -> Option<ProxyInfo> {
-    let result = db.lookup(ip).ok()?;
+pub fn lookup_proxy(db: &SharedDb, ip: IpAddr) -> Option<ProxyInfo> {
+    let guard = db.load();
+    let reader = Option::as_ref(&guard)?;
+    let result = reader.lookup(ip).ok()?;
     let record: ProxyOnly = result.decode().ok()??;
     Some(ProxyInfo {
         is_proxy: record.proxy.is_proxy,
@@ -68,26 +80,34 @@ pub fn lookup_proxy(db: &DbReader, ip: IpAddr) -> Option<ProxyInfo> {
     })
 }
 
-pub fn lookup_asn_number(db: &DbReader, ip: IpAddr) -> Option<u32> {
-    let result = db.lookup(ip).ok()?;
+pub fn lookup_asn_number(db: &SharedDb, ip: IpAddr) -> Option<u32> {
+    let guard = db.load();
+    let reader = Option::as_ref(&guard)?;
+    let result = reader.lookup(ip).ok()?;
     let record: AsnOnly<'_> = result.decode().ok()??;
     record.asn.autonomous_system_number
 }
 
-pub fn lookup_asn_org(db: &DbReader, ip: IpAddr) -> Option<String> {
-    let result = db.lookup(ip).ok()?;
+pub fn lookup_asn_org(db: &SharedDb, ip: IpAddr) -> Option<String> {
+    let guard = db.load();
+    let reader = Option::as_ref(&guard)?;
+    let result = reader.lookup(ip).ok()?;
     let record: AsnOnly<'_> = result.decode().ok()??;
     record.asn.autonomous_system_organization.map(str::to_owned)
 }
 
-pub fn lookup_country_name(db: &DbReader, ip: IpAddr) -> Option<String> {
-    let result = db.lookup(ip).ok()?;
+pub fn lookup_country_name(db: &SharedDb, ip: IpAddr) -> Option<String> {
+    let guard = db.load();
+    let reader = Option::as_ref(&guard)?;
+    let result = reader.lookup(ip).ok()?;
     let record: CountryOnly<'_> = result.decode().ok()??;
     record.country.names.english.map(str::to_owned)
 }
 
-pub fn lookup_city_name(db: &DbReader, ip: IpAddr) -> Option<String> {
-    let result = db.lookup(ip).ok()?;
+pub fn lookup_city_name(db: &SharedDb, ip: IpAddr) -> Option<String> {
+    let guard = db.load();
+    let reader = Option::as_ref(&guard)?;
+    let result = reader.lookup(ip).ok()?;
     let record: CityOnly<'_> = result.decode().ok()??;
     record.city.names.english.map(str::to_owned)
 }
@@ -106,6 +126,10 @@ mod tests {
         load_db(Path::new(&test_db_path())).expect("failed to load test DB")
     }
 
+    fn test_shared_db() -> SharedDb {
+        new_shared(test_reader())
+    }
+
     #[test]
     fn load_db_with_valid_path() {
         let reader = load_db(Path::new(&test_db_path()));
@@ -122,18 +146,25 @@ mod tests {
     fn new_shared_creates_valid_shared_db() {
         let reader = test_reader();
         let shared = new_shared(reader);
-        let loaded = shared.load();
-        // Verify the shared DB is functional by doing a simple lookup
+        assert!(is_ready(&shared));
+        let guard = shared.load();
+        let reader = Option::as_ref(&guard).unwrap();
         let ip: IpAddr = "1.1.1.1".parse().unwrap();
-        let result = loaded.lookup(ip);
+        let result = reader.lookup(ip);
         assert!(result.is_ok());
     }
 
+    #[test]
+    fn new_empty_creates_unready_shared_db() {
+        let shared = new_empty();
+        assert!(!is_ready(&shared));
+    }
+
     #[test]
     fn lookup_returns_some_for_known_ip() {
-        let reader = test_reader();
+        let shared = test_shared_db();
         let ip: IpAddr = "1.1.1.1".parse().unwrap();
-        let result = lookup(&reader, ip);
+        let result = lookup(&shared, ip);
         assert!(result.is_some());
         let info = result.unwrap();
         assert_eq!(info.ip, "1.1.1.1");
@@ -141,17 +172,24 @@ mod tests {
 
     #[test]
     fn lookup_returns_none_for_unroutable_ip() {
-        let reader = test_reader();
+        let shared = test_shared_db();
         let ip: IpAddr = "0.0.0.0".parse().unwrap();
-        let result = lookup(&reader, ip);
+        let result = lookup(&shared, ip);
         assert!(result.is_none());
     }
 
+    #[test]
+    fn lookup_returns_none_when_db_not_loaded() {
+        let shared = new_empty();
+        let ip: IpAddr = "1.1.1.1".parse().unwrap();
+        assert!(lookup(&shared, ip).is_none());
+    }
+
     #[test]
     fn lookup_proxy_returns_data_for_known_ip() {
-        let reader = test_reader();
+        let shared = test_shared_db();
         let ip: IpAddr = "45.77.77.77".parse().unwrap();
-        let proxy = lookup_proxy(&reader, ip);
+        let proxy = lookup_proxy(&shared, ip);
         assert!(proxy.is_some());
         let proxy = proxy.unwrap();
         assert!(proxy.is_proxy);
@@ -160,82 +198,82 @@ mod tests {
 
     #[test]
     fn lookup_proxy_returns_none_for_unroutable_ip() {
-        let reader = test_reader();
+        let shared = test_shared_db();
         let ip: IpAddr = "0.0.0.0".parse().unwrap();
-        let proxy = lookup_proxy(&reader, ip);
+        let proxy = lookup_proxy(&shared, ip);
         assert!(proxy.is_none());
     }
 
     #[test]
     fn lookup_asn_number_for_known_ip() {
-        let reader = test_reader();
+        let shared = test_shared_db();
         let ip: IpAddr = "1.1.1.1".parse().unwrap();
-        let asn = lookup_asn_number(&reader, ip);
+        let asn = lookup_asn_number(&shared, ip);
         assert_eq!(asn, Some(13335));
     }
 
     #[test]
     fn lookup_asn_number_returns_none_for_unroutable_ip() {
-        let reader = test_reader();
+        let shared = test_shared_db();
         let ip: IpAddr = "0.0.0.0".parse().unwrap();
-        let asn = lookup_asn_number(&reader, ip);
+        let asn = lookup_asn_number(&shared, ip);
         assert!(asn.is_none());
     }
 
     #[test]
     fn lookup_asn_org_for_known_ip() {
-        let reader = test_reader();
+        let shared = test_shared_db();
         let ip: IpAddr = "1.1.1.1".parse().unwrap();
-        let org = lookup_asn_org(&reader, ip);
+        let org = lookup_asn_org(&shared, ip);
         assert!(org.is_some());
         assert_eq!(org.unwrap(), "Cloudflare, Inc.");
     }
 
     #[test]
     fn lookup_asn_org_returns_none_for_unroutable_ip() {
-        let reader = test_reader();
+        let shared = test_shared_db();
         let ip: IpAddr = "0.0.0.0".parse().unwrap();
-        let org = lookup_asn_org(&reader, ip);
+        let org = lookup_asn_org(&shared, ip);
         assert!(org.is_none());
     }
 
     #[test]
     fn lookup_country_name_for_known_ip() {
-        let reader = test_reader();
+        let shared = test_shared_db();
         let ip: IpAddr = "45.77.77.77".parse().unwrap();
-        let country = lookup_country_name(&reader, ip);
+        let country = lookup_country_name(&shared, ip);
         assert_eq!(country.as_deref(), Some("United States"));
     }
 
     #[test]
     fn lookup_country_name_returns_none_for_unroutable_ip() {
-        let reader = test_reader();
+        let shared = test_shared_db();
         let ip: IpAddr = "0.0.0.0".parse().unwrap();
-        let country = lookup_country_name(&reader, ip);
+        let country = lookup_country_name(&shared, ip);
         assert!(country.is_none());
     }
 
     #[test]
     fn lookup_city_name_for_known_ip() {
-        let reader = test_reader();
+        let shared = test_shared_db();
         let ip: IpAddr = "45.77.77.77".parse().unwrap();
-        let city = lookup_city_name(&reader, ip);
+        let city = lookup_city_name(&shared, ip);
         assert_eq!(city.as_deref(), Some("Piscataway"));
     }
 
     #[test]
     fn lookup_city_name_returns_none_for_unroutable_ip() {
-        let reader = test_reader();
+        let shared = test_shared_db();
         let ip: IpAddr = "0.0.0.0".parse().unwrap();
-        let city = lookup_city_name(&reader, ip);
+        let city = lookup_city_name(&shared, ip);
         assert!(city.is_none());
     }
 
     #[test]
     fn lookup_works_with_ipv6() {
-        let reader = test_reader();
+        let shared = test_shared_db();
         let ip: IpAddr = "2606:4700:4700::1111".parse().unwrap();
-        let result = lookup(&reader, ip);
+        let result = lookup(&shared, ip);
         assert!(result.is_some());
         let info = result.unwrap();
         assert_eq!(info.ip, "2606:4700:4700::1111");
@@ -248,10 +286,10 @@ mod tests {
 
     #[test]
     fn partial_lookups_work_with_ipv6() {
-        let reader = test_reader();
+        let shared = test_shared_db();
         let ip: IpAddr = "2606:4700:4700::1111".parse().unwrap();
-        assert_eq!(lookup_asn_number(&reader, ip), Some(13335));
-        assert!(lookup_asn_org(&reader, ip).is_some());
-        assert!(lookup_proxy(&reader, ip).is_some());
+        assert_eq!(lookup_asn_number(&shared, ip), Some(13335));
+        assert!(lookup_asn_org(&shared, ip).is_some());
+        assert!(lookup_proxy(&shared, ip).is_some());
     }
 }

src/error.rs

@@ -6,6 +6,7 @@ pub enum AppError {
     InvalidIp,
     NonPublicIp,
     MissingClientIp,
+    DbNotReady,
     DbLookupFailed,
     TemplateRenderFailed,
 }
@@ -16,6 +17,7 @@ impl AppError {
             Self::IpNotFound => StatusCode::NOT_FOUND,
             Self::InvalidIp | Self::NonPublicIp => StatusCode::BAD_REQUEST,
             Self::MissingClientIp => StatusCode::BAD_REQUEST,
+            Self::DbNotReady => StatusCode::SERVICE_UNAVAILABLE,
             Self::DbLookupFailed => StatusCode::INTERNAL_SERVER_ERROR,
             Self::TemplateRenderFailed => StatusCode::INTERNAL_SERVER_ERROR,
         }
@@ -27,6 +29,7 @@ impl AppError {
             Self::InvalidIp => "Invalid IP address",
             Self::NonPublicIp => "Only public IP addresses can be queried",
             Self::MissingClientIp => "Missing client IP address",
+            Self::DbNotReady => "Service is starting up, please try again shortly",
             Self::DbLookupFailed => "Database lookup failed",
             Self::TemplateRenderFailed => "Template render failed",
         }
@@ -96,6 +99,16 @@ mod tests {
         .await;
     }
 
+    #[tokio::test]
+    async fn db_not_ready_returns_503() {
+        check_error_response(
+            AppError::DbNotReady,
+            StatusCode::SERVICE_UNAVAILABLE,
+            "Service is starting up, please try again shortly",
+        )
+        .await;
+    }
+
     #[tokio::test]
     async fn db_lookup_failed_returns_500() {
         check_error_response(

src/handlers.rs

@@ -177,8 +177,10 @@ fn extract_client_ip(headers: &HeaderMap, dev_mode: bool) -> Result<IpAddr, AppE
 }
 
 fn lookup_ip(db: &SharedDb, ip: IpAddr) -> Result<IpInfo, AppError> {
-    let reader = db.load();
-    db::lookup(&reader, ip).ok_or(AppError::DbLookupFailed)
+    if !db::is_ready(db) {
+        return Err(AppError::DbNotReady);
+    }
+    db::lookup(db, ip).ok_or(AppError::DbLookupFailed)
 }
 
 fn wants_json(headers: &HeaderMap) -> bool {
@@ -675,9 +677,11 @@ pub async fn asn_handler(
     headers: HeaderMap,
     Query(query): Query<IpQuery>,
 ) -> Result<String, AppError> {
+    if !db::is_ready(&state.db) {
+        return Err(AppError::DbNotReady);
+    }
     let ip = resolve_ip(&headers, &query, state.dev_mode)?;
-    let reader = state.db.load();
-    let text = db::lookup_asn_number(&reader, ip)
+    let text = db::lookup_asn_number(&state.db, ip)
         .map(|n| format!("AS{n}"))
         .unwrap_or_default();
     Ok(format!("{text}\n"))
@@ -688,9 +692,11 @@ pub async fn org_handler(
     headers: HeaderMap,
     Query(query): Query<IpQuery>,
 ) -> Result<String, AppError> {
+    if !db::is_ready(&state.db) {
+        return Err(AppError::DbNotReady);
+    }
     let ip = resolve_ip(&headers, &query, state.dev_mode)?;
-    let reader = state.db.load();
-    let org = db::lookup_asn_org(&reader, ip).unwrap_or_default();
+    let org = db::lookup_asn_org(&state.db, ip).unwrap_or_default();
     Ok(format!("{org}\n"))
 }
 
@@ -699,9 +705,11 @@ pub async fn country_handler(
     headers: HeaderMap,
     Query(query): Query<IpQuery>,
 ) -> Result<String, AppError> {
+    if !db::is_ready(&state.db) {
+        return Err(AppError::DbNotReady);
+    }
     let ip = resolve_ip(&headers, &query, state.dev_mode)?;
-    let reader = state.db.load();
-    let country = db::lookup_country_name(&reader, ip).unwrap_or_default();
+    let country = db::lookup_country_name(&state.db, ip).unwrap_or_default();
     Ok(format!("{country}\n"))
 }
 
@@ -710,9 +718,11 @@ pub async fn city_handler(
     headers: HeaderMap,
     Query(query): Query<IpQuery>,
 ) -> Result<String, AppError> {
+    if !db::is_ready(&state.db) {
+        return Err(AppError::DbNotReady);
+    }
     let ip = resolve_ip(&headers, &query, state.dev_mode)?;
-    let reader = state.db.load();
-    let city = db::lookup_city_name(&reader, ip).unwrap_or_default();
+    let city = db::lookup_city_name(&state.db, ip).unwrap_or_default();
     Ok(format!("{city}\n"))
 }
 
@@ -721,9 +731,11 @@ pub async fn proxy_handler(
     headers: HeaderMap,
     Query(query): Query<IpQuery>,
 ) -> Result<String, AppError> {
+    if !db::is_ready(&state.db) {
+        return Err(AppError::DbNotReady);
+    }
     let ip = resolve_ip(&headers, &query, state.dev_mode)?;
-    let reader = state.db.load();
-    let proxy = db::lookup_proxy(&reader, ip).ok_or(AppError::DbLookupFailed)?;
+    let proxy = db::lookup_proxy(&state.db, ip).ok_or(AppError::DbLookupFailed)?;
     Ok(format!("{}\n", proxy.is_proxy))
 }
 
@@ -732,9 +744,11 @@ pub async fn vpn_handler(
     headers: HeaderMap,
     Query(query): Query<IpQuery>,
 ) -> Result<String, AppError> {
+    if !db::is_ready(&state.db) {
+        return Err(AppError::DbNotReady);
+    }
     let ip = resolve_ip(&headers, &query, state.dev_mode)?;
-    let reader = state.db.load();
-    let proxy = db::lookup_proxy(&reader, ip).ok_or(AppError::DbLookupFailed)?;
+    let proxy = db::lookup_proxy(&state.db, ip).ok_or(AppError::DbLookupFailed)?;
     Ok(format!("{}\n", proxy.is_vpn))
 }
 
@@ -743,9 +757,11 @@ pub async fn hosting_handler(
     headers: HeaderMap,
     Query(query): Query<IpQuery>,
 ) -> Result<String, AppError> {
+    if !db::is_ready(&state.db) {
+        return Err(AppError::DbNotReady);
+    }
     let ip = resolve_ip(&headers, &query, state.dev_mode)?;
-    let reader = state.db.load();
-    let proxy = db::lookup_proxy(&reader, ip).ok_or(AppError::DbLookupFailed)?;
+    let proxy = db::lookup_proxy(&state.db, ip).ok_or(AppError::DbLookupFailed)?;
     Ok(format!("{}\n", proxy.is_hosting))
 }
 
@@ -754,9 +770,11 @@ pub async fn tor_handler(
     headers: HeaderMap,
     Query(query): Query<IpQuery>,
 ) -> Result<String, AppError> {
+    if !db::is_ready(&state.db) {
+        return Err(AppError::DbNotReady);
+    }
     let ip = resolve_ip(&headers, &query, state.dev_mode)?;
-    let reader = state.db.load();
-    let proxy = db::lookup_proxy(&reader, ip).ok_or(AppError::DbLookupFailed)?;
+    let proxy = db::lookup_proxy(&state.db, ip).ok_or(AppError::DbLookupFailed)?;
     Ok(format!("{}\n", proxy.is_tor))
 }
 

src/main.rs

@@ -19,35 +19,30 @@ async fn main() {
     let config = config::Config::from_env();
     let db_path = PathBuf::from(&config.db_path);
 
-    if let Err(e) = updater::ensure_db_exists(&db_path, &config.db_update_url).await {
-        error!("failed to initialize database: {e}");
-        std::process::exit(1);
-    }
-
-    let reader = match db::load_db(&db_path) {
-        Ok(r) => r,
-        Err(e) => {
-            error!("failed to load database: {e}");
-            std::process::exit(1);
-        }
+    // Create shared DB in an empty (not-yet-loaded) state so the server can
+    // start immediately and respond to health checks while the database is
+    // still being downloaded on first deploy.
+    let shared_db = if db_path.exists() {
+        let reader = match db::load_db(&db_path) {
+            Ok(r) => r,
+            Err(e) => {
+                error!("failed to load database: {e}");
+                std::process::exit(1);
+            }
+        };
+        db::new_shared(reader)
+    } else {
+        info!("database file not found, server will start while download proceeds");
+        db::new_empty()
     };
 
-    let shared_db = db::new_shared(reader);
     let cancel = CancellationToken::new();
 
-    let updater_handle = updater::spawn_updater(
-        shared_db.clone(),
-        db_path,
-        config.db_update_url,
-        config.db_update_interval_hours,
-        cancel.clone(),
-    );
-
     let site_domain: std::sync::Arc<str> = config.site_domain.into();
     let openapi_json = handlers::build_openapi_json(&site_domain);
 
     let state = handlers::AppState {
-        db: shared_db,
+        db: shared_db.clone(),
         site_domain,
         dev_mode: config.dev_mode,
         openapi_json,
@@ -59,6 +54,33 @@ async fn main() {
         .await
         .expect("failed to bind listener");
 
+    // Spawn background database initialization and periodic updater.
+    let init_db = shared_db.clone();
+    let init_path = db_path.clone();
+    let init_url = config.db_update_url.clone();
+    let init_cancel = cancel.clone();
+    let updater_handle = tokio::spawn(async move {
+        // Only download if the DB was not loaded synchronously above.
+        if !db::is_ready(&init_db) {
+            if let Err(e) = updater::init_db(&init_db, &init_path, &init_url).await {
+                error!("failed to initialize database: {e}");
+                std::process::exit(1);
+            }
+        }
+
+        // Start periodic updater after the initial load succeeds.
+        updater::run_updater(
+            init_db,
+            init_path,
+            init_url,
+            config.db_update_interval_hours,
+            init_cancel,
+        )
+        .await;
+    });
+
+    info!("listening on {}", config.listen_addr);
+
     axum::serve(listener, app)
         .with_graceful_shutdown(shutdown_signal())
         .await

src/updater.rs

@@ -10,6 +10,9 @@ use tracing::{error, info};
 
 use crate::db::{DbReader, SharedDb, load_db};
 
+#[cfg(test)]
+use crate::db::new_shared;
+
 pub async fn ensure_db_exists(
     db_path: &Path,
     update_url: &str,
@@ -26,6 +29,21 @@ pub async fn ensure_db_exists(
     Ok(())
 }
 
+/// Ensures the database file exists (downloading if needed), loads it, and
+/// stores the reader into the shared DB. Called during startup to initialize
+/// the database before the application can serve requests.
+pub async fn init_db(
+    shared_db: &SharedDb,
+    db_path: &Path,
+    update_url: &str,
+) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
+    ensure_db_exists(db_path, update_url).await?;
+    let reader = load_db(db_path)?;
+    shared_db.store(Arc::new(Some(reader)));
+    info!("database loaded successfully");
+    Ok(())
+}
+
 /// Downloads the database to a temp file, validates it, and atomically
 /// renames it into place. Returns the validated reader for immediate use.
 /// Cleans up the temp file on any failure.
@@ -86,21 +104,37 @@ pub fn spawn_updater(
     interval_hours: u64,
     cancel: CancellationToken,
 ) -> JoinHandle<()> {
-    tokio::spawn(async move {
-        let interval = Duration::from_secs(interval_hours * 3600);
-        loop {
-            tokio::select! {
-                () = tokio::time::sleep(interval) => {}
-                () = cancel.cancelled() => {
-                    info!("updater task cancelled, stopping");
-                    break;
-                }
-            }
-            if let Err(e) = update_db(&shared_db, &db_path, &update_url).await {
-                error!("database update failed: {e}");
+    tokio::spawn(run_updater(
+        shared_db,
+        db_path,
+        update_url,
+        interval_hours,
+        cancel,
+    ))
+}
+
+/// Runs the periodic database updater loop. This is the inner async function
+/// that can be called directly when already inside a spawned task.
+pub async fn run_updater(
+    shared_db: SharedDb,
+    db_path: PathBuf,
+    update_url: String,
+    interval_hours: u64,
+    cancel: CancellationToken,
+) {
+    let interval = Duration::from_secs(interval_hours * 3600);
+    loop {
+        tokio::select! {
+            () = tokio::time::sleep(interval) => {}
+            () = cancel.cancelled() => {
+                info!("updater task cancelled, stopping");
+                break;
             }
         }
-    })
+        if let Err(e) = update_db(&shared_db, &db_path, &update_url).await {
+            error!("database update failed: {e}");
+        }
+    }
 }
 
 async fn update_db(
@@ -109,7 +143,7 @@ async fn update_db(
     update_url: &str,
 ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
     let reader = download_db(update_url, db_path).await?;
-    shared_db.store(Arc::new(reader));
+    shared_db.store(Arc::new(Some(reader)));
     info!("database updated successfully");
     Ok(())
 }
@@ -210,7 +244,7 @@ mod tests {
     #[tokio::test]
     async fn spawn_updater_stops_on_cancellation() {
         let reader = load_db(&test_db_path()).unwrap();
-        let shared_db = crate::db::new_shared(reader);
+        let shared_db = new_shared(reader);
         let cancel = CancellationToken::new();
 
         let handle = spawn_updater(
@@ -239,16 +273,17 @@ mod tests {
         std::fs::copy(test_db_path(), &dest).unwrap();
 
         let original_reader = load_db(&dest).unwrap();
-        let shared_db = crate::db::new_shared(original_reader);
+        let shared_db = new_shared(original_reader);
 
         // Load a new reader and manually swap it to verify the mechanism
         let new_reader = load_db(&dest).unwrap();
-        shared_db.store(Arc::new(new_reader));
+        shared_db.store(Arc::new(Some(new_reader)));
 
         // Verify the shared DB still works after the swap
         let guard = shared_db.load();
+        let reader = Option::as_ref(&guard).unwrap();
         let ip: IpAddr = "1.1.1.1".parse().unwrap();
-        let result = guard.lookup(ip);
+        let result = reader.lookup(ip);
         assert!(result.is_ok());
     }
 }