From 9f7bb2310865b5a8d1f55eff3a001a0efdc4e12a Mon Sep 17 00:00:00 2001
From: Gustavo Inacio <gustavo@semiotic.ai>
Date: Wed, 13 Dec 2023 17:23:51 -0300
Subject: [PATCH] fix: drop the manager handler correctly

---
 tap-agent/src/agent.rs |  6 +++---
 tap-agent/src/main.rs  | 27 +++++++++++++++------------
 2 files changed, 18 insertions(+), 15 deletions(-)

diff --git a/tap-agent/src/agent.rs b/tap-agent/src/agent.rs
index 4b54c932..d1ecde26 100644
--- a/tap-agent/src/agent.rs
+++ b/tap-agent/src/agent.rs
@@ -13,7 +13,7 @@ use crate::{
     tap::sender_allocation_relationships_manager::SenderAllocationRelationshipsManager,
 };
 
-pub async fn start_agent(config: &'static config::Cli) {
+pub async fn start_agent(config: &'static config::Cli) -> SenderAllocationRelationshipsManager {
     let pgpool = database::connect(&config.postgres).await;
 
     let http_client = reqwest::Client::new();
@@ -79,7 +79,7 @@ pub async fn start_agent(config: &'static config::Cli) {
         verifying_contract: config.receipts.receipts_verifier_address,
     };
 
-    let _sender_allocation_relationships_manager = SenderAllocationRelationshipsManager::new(
+    SenderAllocationRelationshipsManager::new(
         config,
         pgpool,
         indexer_allocations,
@@ -88,5 +88,5 @@ pub async fn start_agent(config: &'static config::Cli) {
         tap_eip712_domain_separator,
         sender_aggregator_endpoints,
     )
-    .await;
+    .await
 }
diff --git a/tap-agent/src/main.rs b/tap-agent/src/main.rs
index 68f34e88..a99ac21d 100644
--- a/tap-agent/src/main.rs
+++ b/tap-agent/src/main.rs
@@ -24,19 +24,22 @@ async fn main() -> Result<()> {
     lazy_static::initialize(&CONFIG);
     debug!("Config: {:?}", *CONFIG);
 
-    agent::start_agent(&CONFIG).await;
-    info!("TAP Agent started.");
-
-    // Have tokio wait for SIGTERM or SIGINT.
-    let mut signal_sigint = signal(SignalKind::interrupt())?;
-    let mut signal_sigterm = signal(SignalKind::terminate())?;
-    tokio::select! {
-        _ = signal_sigint.recv() => debug!("Received SIGINT."),
-        _ = signal_sigterm.recv() => debug!("Received SIGTERM."),
+    {
+        let _manager = agent::start_agent(&CONFIG).await;
+        info!("TAP Agent started.");
+
+        // Have tokio wait for SIGTERM or SIGINT.
+        let mut signal_sigint = signal(SignalKind::interrupt())?;
+        let mut signal_sigterm = signal(SignalKind::terminate())?;
+        tokio::select! {
+            _ = signal_sigint.recv() => debug!("Received SIGINT."),
+            _ = signal_sigterm.recv() => debug!("Received SIGTERM."),
+        }
+
+        // If we're here, we've received a signal to exit.
+        info!("Shutting down...");
     }
-
-    // If we're here, we've received a signal to exit.
-    info!("Shutting down...");
+    // Manager should be successfully dropped here.
 
     // Stop the server and wait for it to finish gracefully.
     debug!("Goodbye!");