4 files changed, 46 insertions, 9 deletions

diff --git a/system/gd/hal/snoop_logger.h b/system/gd/hal/snoop_logger.h
index b009add89c..b0783442b3 100644
--- a/system/gd/hal/snoop_logger.h
+++ b/system/gd/hal/snoop_logger.h
@@ -17,6 +17,7 @@
 #pragma once
 
 #include <bluetooth/log.h>
+#include <com_android_bluetooth_flags.h>
 
 #include <fstream>
 #include <string>
@@ -195,6 +196,13 @@ public:
   };
 
   SnoopLogger(os::Handler* handler);
+  ~SnoopLogger() {
+    if (!com::android::bluetooth::flags::same_handler_for_all_modules()) {
+      GetHandler()->Clear();
+      GetHandler()->WaitUntilStopped(std::chrono::milliseconds(2000));
+      delete GetHandler();
+    }
+  }
 
   // Returns the maximum number of packets per file
   // Changes to this value is only effective after restarting Bluetooth
diff --git a/system/gd/metrics/counter_metrics.h b/system/gd/metrics/counter_metrics.h
index 88c7711f92..2f25f03ed0 100644
--- a/system/gd/metrics/counter_metrics.h
+++ b/system/gd/metrics/counter_metrics.h
@@ -15,6 +15,8 @@
  */
 #pragma once
 
+#include <com_android_bluetooth_flags.h>
+
 #include <unordered_map>
 
 #include "module.h"
@@ -26,6 +28,13 @@ namespace metrics {
 class CounterMetrics : public bluetooth::Module {
 public:
   CounterMetrics(os::Handler* handler) : Module(handler) {}
+  ~CounterMetrics() {
+    if (!com::android::bluetooth::flags::same_handler_for_all_modules()) {
+      GetHandler()->Clear();
+      GetHandler()->WaitUntilStopped(std::chrono::milliseconds(2000));
+      delete GetHandler();
+    }
+  }
 
   bool CacheCount(int32_t key, int64_t value);
   virtual bool Count(int32_t key, int64_t count);
diff --git a/system/gd/storage/storage_module.cc b/system/gd/storage/storage_module.cc
index 5e3459e109..5a81a5cf94 100644
--- a/system/gd/storage/storage_module.cc
+++ b/system/gd/storage/storage_module.cc
@@ -90,6 +90,13 @@ StorageModule::StorageModule(os::Handler* handler, std::string config_file_path,
 
 StorageModule::~StorageModule() {
   std::lock_guard<std::recursive_mutex> lock(mutex_);
+
+  if (!com::android::bluetooth::flags::same_handler_for_all_modules()) {
+    GetHandler()->Clear();
+    GetHandler()->WaitUntilStopped(std::chrono::milliseconds(2000));
+    delete GetHandler();
+  }
+
   pimpl_.reset();
 }
 
diff --git a/system/main/shim/stack.cc b/system/main/shim/stack.cc
index 4560a6180f..b7b93d3690 100644
--- a/system/main/shim/stack.cc
+++ b/system/main/shim/stack.cc
@@ -67,9 +67,9 @@ namespace shim {
 
 struct Stack::impl {
   Acl* acl_ = nullptr;
-  metrics::CounterMetrics* counter_metrics_ = nullptr;
-  storage::StorageModule* storage_ = nullptr;
-  hal::SnoopLogger* snoop_logger_ = nullptr;
+  std::shared_ptr<metrics::CounterMetrics> counter_metrics_ = nullptr;
+  std::shared_ptr<storage::StorageModule> storage_ = nullptr;
+  std::shared_ptr<hal::SnoopLogger> snoop_logger_ = nullptr;
 };
 
 Stack::Stack() { pimpl_ = std::make_shared<Stack::impl>(); }
@@ -89,9 +89,16 @@ void Stack::StartEverything() {
     stack_thread_ = new os::Thread("gd_stack_thread", os::Thread::Priority::REAL_TIME);
     stack_handler_ = new os::Handler(stack_thread_);
 
-    pimpl_->counter_metrics_ = new metrics::CounterMetrics(new Handler(stack_thread_));
-    pimpl_->storage_ = new storage::StorageModule(new Handler(stack_thread_));
-    pimpl_->snoop_logger_ = new hal::SnoopLogger(new Handler(stack_thread_));
+    if (com::android::bluetooth::flags::same_handler_for_all_modules()) {
+      pimpl_->counter_metrics_ = std::make_shared<metrics::CounterMetrics>(stack_handler_);
+      pimpl_->storage_ = std::make_shared<storage::StorageModule>(stack_handler_);
+      pimpl_->snoop_logger_ = std::make_shared<hal::SnoopLogger>(stack_handler_);
+    } else {
+      pimpl_->counter_metrics_ =
+              std::make_shared<metrics::CounterMetrics>(new Handler(stack_thread_));
+      pimpl_->storage_ = std::make_shared<storage::StorageModule>(new Handler(stack_thread_));
+      pimpl_->snoop_logger_ = std::make_shared<hal::SnoopLogger>(new Handler(stack_thread_));
+    }
 
 #if TARGET_FLOSS
     modules.add<sysprops::SyspropsModule>();
@@ -220,19 +227,19 @@ Acl* Stack::GetAcl() const {
 metrics::CounterMetrics* Stack::GetCounterMetrics() const {
   std::lock_guard<std::recursive_mutex> lock(mutex_);
   log::assert_that(is_running_, "assert failed: is_running_");
-  return pimpl_->counter_metrics_;
+  return pimpl_->counter_metrics_.get();
 }
 
 storage::StorageModule* Stack::GetStorage() const {
   std::lock_guard<std::recursive_mutex> lock(mutex_);
   log::assert_that(is_running_, "assert failed: is_running_");
-  return pimpl_->storage_;
+  return pimpl_->storage_.get();
 }
 
 hal::SnoopLogger* Stack::GetSnoopLogger() const {
   std::lock_guard<std::recursive_mutex> lock(mutex_);
   log::assert_that(is_running_, "assert failed: is_running_");
-  return pimpl_->snoop_logger_;
+  return pimpl_->snoop_logger_.get();
 }
 
 os::Handler* Stack::GetHandler() {
@@ -268,9 +275,15 @@ void Stack::handle_start_up(ModuleList* modules, std::promise<void> promise) {
 
 void Stack::handle_shut_down(std::promise<void> promise) {
   registry_.StopAll();
+
   pimpl_->snoop_logger_->Stop();
   pimpl_->storage_->Stop();
   pimpl_->counter_metrics_->Stop();
+
+  pimpl_->snoop_logger_.reset();
+  pimpl_->storage_.reset();
+  pimpl_->counter_metrics_.reset();
+
   promise.set_value();
 }