From 3940d87f1ee4060b4ad8f313888cb48d370cd41b Mon Sep 17 00:00:00 2001
From: Mingwei Zhang <mingwei@mwzhang.com>
Date: Tue, 18 Jun 2024 15:50:02 -0700
Subject: [PATCH 1/2] add heartbeat notifier to cli serve mode

heartbeat URL can be specified with environment variable `BGPKIT_BROKER_HEARTBEAT_URL`
---
 src/cli/main.rs | 44 ++++++++++++++++++++++++++++++++------------
 1 file changed, 32 insertions(+), 12 deletions(-)

diff --git a/src/cli/main.rs b/src/cli/main.rs
index 32198cb..cedcba0 100644
--- a/src/cli/main.rs
+++ b/src/cli/main.rs
@@ -7,7 +7,8 @@ use crate::backup::backup_database;
 use crate::bootstrap::download_file;
 use bgpkit_broker::notifier::NatsNotifier;
 use bgpkit_broker::{
-    crawl_collector, load_collectors, BgpkitBroker, Collector, LocalBrokerDb, DEFAULT_PAGE_SIZE,
+    crawl_collector, load_collectors, BgpkitBroker, BrokerError, Collector, LocalBrokerDb,
+    DEFAULT_PAGE_SIZE,
 };
 use bgpkit_commons::collectors::MrtCollector;
 use chrono::{Duration, NaiveDateTime, Utc};
@@ -225,22 +226,35 @@ fn get_tokio_runtime() -> Runtime {
     rt
 }
 
+async fn try_send_heartbeat(url: Option<String>) -> Result<(), BrokerError> {
+    let url = match url {
+        Some(u) => u,
+        None => match dotenvy::var("BGPKIT_BROKER_HEARTBEAT_URL") {
+            Ok(u) => u,
+            Err(_) => {
+                info!("no heartbeat url specified, skipping");
+                return Ok(());
+            }
+        },
+    };
+    reqwest::get(&url).await?.error_for_status()?;
+    info!("heartbeat sent");
+    Ok(())
+}
+
 /// update the database with data crawled from the given collectors
 async fn update_database(
     db: LocalBrokerDb,
     collectors: Vec<Collector>,
     days: Option<u32>,
-    notify: bool,
+    send_heartbeat: bool,
 ) {
-    let notifier = match notify {
-        true => match NatsNotifier::new(None).await {
-            Ok(n) => Some(n),
-            Err(e) => {
-                error!("want to set up notifier but failed: {}", e);
-                None
-            }
-        },
-        false => None,
+    let notifier = match NatsNotifier::new(None).await {
+        Ok(n) => Some(n),
+        Err(e) => {
+            error!("want to set up notifier but failed: {}", e);
+            None
+        }
     };
 
     let now = Utc::now();
@@ -312,6 +326,12 @@ async fn update_database(
         .await
         .unwrap();
 
+    if send_heartbeat {
+        if let Err(e) = try_send_heartbeat(None).await {
+            error!("{}", e);
+        }
+    }
+
     info!("finished updating broker database");
 }
 
@@ -510,7 +530,7 @@ fn main() {
 
             rt.block_on(async {
                 let db = LocalBrokerDb::new(&db_path).await.unwrap();
-                update_database(db, collectors, days, true).await;
+                update_database(db, collectors, days, false).await;
             });
         }
         Commands::Search { query, json, url } => {

From 3497b2b21b2fe6150c3c89a74475359d16f486dd Mon Sep 17 00:00:00 2001
From: Mingwei Zhang <mingwei@mwzhang.com>
Date: Tue, 18 Jun 2024 15:50:27 -0700
Subject: [PATCH 2/2] add env file param to nomad config template

---
 deployment/nomad_api_raw.hcl | 1 +
 1 file changed, 1 insertion(+)

diff --git a/deployment/nomad_api_raw.hcl b/deployment/nomad_api_raw.hcl
index 1ea2271..df12597 100644
--- a/deployment/nomad_api_raw.hcl
+++ b/deployment/nomad_api_raw.hcl
@@ -9,6 +9,7 @@ job "bgpkit-broker-api" {
         args    = [
           "serve",
           "--port", "40064",
+          "--env", "/usr/local/etc/bgpkit.d/broker.env",
           "/var/db/bgpkit/bgpkit_broker.sqlite3"
         ]
       }