Segfault when restarting RMQ broker with MPI

CMakeLists.txt

@@ -50,6 +50,11 @@ option(WITH_AMS_LIB        "Install C++ library to support scientific applicatio
 option(WITH_ADIAK          "Use Adiak for recording metadata" OFF)
 option(BUILD_SHARED_LIBS   "Build using shared libraries" ON)
 
+#if (CMAKE_CXX_COMPILER_ID MATCHES "Clang|GNU")
+#  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fsanitize=thread")
+#  set(CMAKE_EXE_LINKER_FLAGS "${CMAKE_EXE_LINKER_FLAGS} -fsanitize=thread")
+#endif()
+
 if (WITH_MPI)
   #  SET(CMAKE_CXX_COMPILER "${MPI_CXX_COMPILER}" CACHE FILEPATH "CXX compiler overridden with MPI C++ wrapper")
   #SET(CMAKE_C_COMPILER "${MPI_C_COMPILER}" CACHE FILEPATH "CXX compiler overridden with MPI C++ wrapper")

src/AMSlib/AMS.cpp

@@ -512,6 +512,9 @@ class AMSWrap
 public:
   AMSWrap() : memManager(ams::ResourceManager::getInstance())
   {
+  }
+
+  void init() {
     auto log_stats = setup_loggers();
     DBG(AMS,
         "Enable Log %d stored under %s",
@@ -596,6 +599,7 @@ class AMSWrap
 };
 
 static AMSWrap _amsWrap;
+static std::once_flag _AMSInit;
 
 void _AMSExecute(AMSExecutor executor,
                  void *probDescr,
@@ -605,6 +609,10 @@ void _AMSExecute(AMSExecutor executor,
                  int inputDim,
                  int outputDim)
 {
+  std::call_once(_AMSInit, [&]() {
+    _amsWrap.init();
+  });
+
   int64_t index = static_cast<int64_t>(executor);
   if (index >= _amsWrap.executors.size())
     throw std::runtime_error("AMS Executor identifier does not exist\n");
@@ -648,6 +656,10 @@ ams::AMSWorkflow<FPTypeValue> *_AMSCreateExecutor(AMSCAbstrModel model,
     rm.init();
   });
 
+  std::call_once(_AMSInit, [&]() {
+    _amsWrap.init();
+  });
+
   auto &model_descr = _amsWrap.get_model(model);
 
   ams::AMSWorkflow<FPTypeValue> *WF =
@@ -670,6 +682,10 @@ template <typename FPTypeValue>
 AMSExecutor _AMSRegisterExecutor(AMSDType data_type,
                                  ams::AMSWorkflow<FPTypeValue> *workflow)
 {
+  std::call_once(_AMSInit, [&]() {
+    _amsWrap.init();
+  });
+
   _amsWrap.executors.push_back(
       std::make_pair(data_type, static_cast<void *>(workflow)));
   return static_cast<AMSExecutor>(_amsWrap.executors.size()) - 1L;
@@ -748,6 +764,10 @@ void AMSExecute(AMSExecutor executor,
 
 void AMSDestroyExecutor(AMSExecutor executor)
 {
+  std::call_once(_AMSInit, [&]() {
+    _amsWrap.init();
+  });
+
   int64_t index = static_cast<int64_t>(executor);
   if (index >= _amsWrap.executors.size())
     throw std::runtime_error("AMS Executor identifier does not exist\n");
@@ -785,6 +805,10 @@ AMSCAbstrModel AMSRegisterAbstractModel(const char *domain_name,
                                         const char *db_label,
                                         int num_clusters)
 {
+  std::call_once(_AMSInit, [&]() {
+    _amsWrap.init();
+  });
+
   auto id = _amsWrap.get_model_index(domain_name);
   if (id == -1) {
     id = _amsWrap.register_model(domain_name,
@@ -802,6 +826,10 @@ AMSCAbstrModel AMSRegisterAbstractModel(const char *domain_name,
 
 AMSCAbstrModel AMSQueryModel(const char *domain_model)
 {
+  std::call_once(_AMSInit, [&]() {
+    _amsWrap.init();
+  });
+
   return _amsWrap.get_model_index(domain_model);
 }
 

src/AMSlib/wf/basedb.hpp

@@ -959,10 +959,10 @@ class AMSMessage
  */
 class AMSMessageRecords
 {
+private:
   using record_t = std::pair<std::shared_ptr<uint8_t>, size_t>;
   using iterator_t = std::unordered_map<int, record_t>::iterator;
 
-private:
   /** @brief Internal data structure that keeps messages nack */
   std::unordered_map<int, record_t> _msgs;
   /** @brief Shared mutex to ensure thread-safe access */
@@ -977,6 +977,22 @@ class AMSMessageRecords
   AMSMessageRecords(AMSMessageRecords&&) = delete;
   AMSMessageRecords& operator=(AMSMessageRecords&&) = delete;
 
+  /**
+  * @brief Custom destructor for our shared_ptr
+  */
+  struct AMSMessageDeleter
+  {
+      void operator()(void* x) {
+        DBG(AMSMessageDeleter, "Freeing %p", x)
+        free(x);
+        x = nullptr;
+      }
+  };
+
+  static AMSMessageDeleter getDeleter() {
+    return AMSMessageDeleter();
+  }
+
   /**
   * @brief Return an iterator at the beggining of the records
   */
@@ -1096,6 +1112,9 @@ class RMQHandler : public AMQP::LibEventHandler
 
   std::atomic<bool> error_connection{0};  // Atomic variable
 
+  // FIXME: TEMPORARY for debug: removed after bug is fixed
+  bool _make_broker_crash = false;
+
 public:
   /**
    *  @brief Constructor
@@ -1107,7 +1126,9 @@ class RMQHandler : public AMQP::LibEventHandler
              std::shared_ptr<struct event_base> loop,
              std::string cacert = "");
 
-  ~RMQHandler() = default;
+  ~RMQHandler() {
+    DBG(RMQHandler, "In ~RMQHandler()");
+  };
 
   /**
    *  @brief  Wait (blocking call) until connection has been established or that ms * repeat is over.
@@ -1402,7 +1423,9 @@ class RMQPublisherHandler final : public RMQHandler
                       std::string cacert,
                       std::string queue);
 
-  ~RMQPublisherHandler() = default;
+  ~RMQPublisherHandler() {
+    DBG(RMQPublisherHandler, "In ~RMQPublisherHandler()");
+  }
 
   /**
    *  @brief  Publish data on RMQ queue.
@@ -1477,6 +1500,12 @@ class RMQPublisher
                std::string cacert,
                std::string queue);
 
+  ~RMQPublisher() {
+    DBG(RMQPublisher, "In ~RMQPublisher(%p)", _connection)
+    // delete _connection; // this leads to double free segfault
+    // _connection = nullptr;
+  };
+
   /**
    * @brief Check if the underlying RabbitMQ connection is ready and usable
    * @return True if the publisher is ready to publish
@@ -1536,9 +1565,6 @@ class RMQPublisher
    *  @return   Number of messages
    */
   bool close(unsigned ms, int repeat = 1);
-
-  ~RMQPublisher() = default;
-
 };  // class RMQPublisher
 
 /**
@@ -1715,7 +1741,7 @@ class RMQInterface
 
     if (!_publisher->connectionValid()) restartPublisher();
 
-    std::shared_ptr<uint8_t> ptr(msg.data());
+    std::shared_ptr<uint8_t> ptr(msg.data(), AMSMessageRecords::getDeleter());
     auto record = std::make_pair(std::move(ptr), msg.size());
 
     // if we have some messages to send first (from a potential restart)

src/AMSlib/wf/rmqdb.cpp

@@ -190,10 +190,9 @@ void AMSMessageRecords::print()
         e.second.second);
 }
 
-
 void AMSMessageRecords::publishUnacknowledged(RMQPublisher& publisher)
 {
-  std::shared_lock<std::shared_mutex> lock(_mutex);
+  std::unique_lock<std::shared_mutex> lock(_mutex);
   if (_msgs.size() == 0) return;
 
   for (auto& item : _msgs) {
@@ -306,7 +305,9 @@ bool RMQHandler::waitToClose(unsigned ms, int repeat)
 
 bool RMQHandler::connectionValid()
 {
-  return !error_connection.load(std::memory_order_acquire);
+  auto status = !error_connection.load(std::memory_order_acquire);
+  // FIXME: TEMPORARY for debug: removed after bug is fixed
+  return status && !_make_broker_crash;
 }
 
 bool RMQHandler::onSecuring(AMQP::TcpConnection* connection, SSL* ssl)
@@ -708,19 +709,19 @@ void RMQPublisherHandler::publish(
                   _queue,
                   reinterpret_cast<char*>(message_content.first.get()),
                   message_content.second)
-        .onAck([this, &_nb_msg_ack = _nb_msg_ack, id = message_id]() {
+        .onAck([_rId = _rId, &_nb_msg_ack = _nb_msg_ack, id = message_id, ptr = message_content.first]() {
           DBG(RMQPublisherHandler,
-              "[r%ld] message #%d got acknowledged "
+              "[r%ld] message #%d (%p) got acknowledged "
               "successfully "
               "by "
               "RMQ "
               "server",
               _rId,
-              id)
+              id,
+              ptr.get())
           _nb_msg_ack++;
         })
-        .onNack([this,
-                 id = message_id,
+        .onNack([_rId = _rId, id = message_id,
                  ptr = message_content.first,
                  size = message_content.second]() {
           DBG(RMQPublisherHandler,
@@ -737,7 +738,7 @@ void RMQPublisherHandler::publish(
                                                   std::make_pair(std::move(ptr),
                                                                  size));
         })
-        .onError([this,
+        .onError([_rId = _rId,
                   id = message_id,
                   ptr = message_content.first,
                   size =
@@ -764,6 +765,13 @@ void RMQPublisherHandler::publish(
                        message_content.second));
   }
   _nb_msg++;
+
+  // FIXME: TEMPORARY for debug: removed after bug is fixed
+  if (_nb_msg >= 3) {
+    printf("Crashing broker\n");
+    _make_broker_crash = true;
+  }
+
   CALIPER(CALI_MARK_END("RMQ_PUBLISH");)
 }
 
@@ -859,7 +867,9 @@ RMQPublisher::RMQPublisher(uint64_t rId,
 
   _loop = std::shared_ptr<struct event_base>(event_base_new(),
                                              [](struct event_base* event) {
+                                               DBG(event_base, "In ~event_base(%p)", event);
                                                event_base_free(event);
+                                               DBG(event_base, "freed (%p)", event);
                                              });
 
   _handler =
@@ -950,7 +960,7 @@ std::pair<bool, bool> RMQInterface::connect(std::string rmq_password,
 
   _publisher_thread = std::thread([&]() { _publisher->start(); });
 
-  if (!_publisher->waitToEstablish(100, 10)) {
+  if (!_publisher->waitToEstablish(100, 30)) {
     _publisher->stop();
     _publisher_thread.join();
     FATAL(RMQInterface, "Could not establish publisher connection");
@@ -962,7 +972,7 @@ std::pair<bool, bool> RMQInterface::connect(std::string rmq_password,
         _rId, *_address, _cacert, _exchange, _routing_key);
     _consumer_thread = std::thread([&]() { _consumer->start(); });
 
-    if (!_consumer->waitToEstablish(100, 10)) {
+    if (!_consumer->waitToEstablish(100, 30)) {
       _consumer->stop();
       _consumer_thread.join();
       FATAL(RabbitMQDB, "Could not establish consumer connection");
@@ -981,7 +991,7 @@ void RMQInterface::restartPublisher()
   _publisher_connected = false;
 
   // Stop the faulty publisher
-  _publisher->close(100, 10);
+  _publisher->close(100, 50);
   _publisher->stop();
   if (_publisher_thread.joinable()) _publisher_thread.join();
   _publisher.reset();
@@ -996,7 +1006,7 @@ void RMQInterface::restartPublisher()
       std::make_shared<RMQPublisher>(_rId, *_address, _cacert, _queue_sender);
   _publisher_thread = std::thread([&]() { _publisher->start(); });
 
-  if (!_publisher->waitToEstablish(100, 10)) {
+  if (!_publisher->waitToEstablish(100, 50)) {
     _publisher->stop();
     if (_publisher_thread.joinable()) _publisher_thread.join();
     FATAL(RMQInterface,
@@ -1009,7 +1019,7 @@ void RMQInterface::restartPublisher()
 void RMQInterface::close()
 {
   if (isPublisherConnected()) {
-    bool status = _publisher->close(100, 10);
+    bool status = _publisher->close(100, 30);
     CWARNING(RMQInterface,
              !status,
              "Could not gracefully close publisher TCP connection")
@@ -1027,7 +1037,7 @@ void RMQInterface::close()
   }
 
   if (isConsumerConnected()) {
-    bool status = _consumer->close(100, 10);
+    bool status = _consumer->close(100, 30);
     CWARNING(RabbitMQDB,
              !status,
              "Could not gracefully close consumer TCP connection")

tests/AMSlib/CMakeLists.txt

@@ -200,6 +200,7 @@ ADDTEST(ams_packing_test AMSPack)
 
 # AMS Database benchmark (RMQ and/or HDF5 + MPI / No ML models used)
 BUILD_TEST(ams_benchmark_db ams_bench_db.cpp)
+BUILD_TEST(ams_new_rmq_interface rmq_redesign.cpp)
 # The AMS DB Benchmark requires mfem
 # TODO: Remove mfem requirement from the benchmark
 target_link_libraries(ams_benchmark_db PRIVATE AMS ${AMS_EXAMPLE_LIBRARIES})