Shmem: track number of message objects, throw if non-zero at reset

fairmq/shmem/Message.cxx

@@ -9,6 +9,7 @@
 #include "Region.h"
 #include "Message.h"
 #include "UnmanagedRegion.h"
+#include "TransportFactory.h"
 
 #include <FairMQLogger.h>
 
@@ -39,6 +40,7 @@ Message::Message(Manager& manager, FairMQTransportFactory* factory)
     , fRegionPtr(nullptr)
     , fLocalPtr(nullptr)
 {
+    static_cast<TransportFactory*>(GetTransport())->IncrementMsgCounter();
 }
 
 Message::Message(Manager& manager, const size_t size, FairMQTransportFactory* factory)
@@ -50,6 +52,7 @@ Message::Message(Manager& manager, const size_t size, FairMQTransportFactory* fa
     , fLocalPtr(nullptr)
 {
     InitializeChunk(size);
+    static_cast<TransportFactory*>(GetTransport())->IncrementMsgCounter();
 }
 
 Message::Message(Manager& manager, MetaHeader& hdr, FairMQTransportFactory* factory)
@@ -60,6 +63,7 @@ Message::Message(Manager& manager, MetaHeader& hdr, FairMQTransportFactory* fact
     , fRegionPtr(nullptr)
     , fLocalPtr(nullptr)
 {
+    static_cast<TransportFactory*>(GetTransport())->IncrementMsgCounter();
 }
 
 Message::Message(Manager& manager, void* data, const size_t size, fairmq_free_fn* ffn, void* hint, FairMQTransportFactory* factory)
@@ -78,6 +82,7 @@ Message::Message(Manager& manager, void* data, const size_t size, fairmq_free_fn
             free(data);
         }
     }
+    static_cast<TransportFactory*>(GetTransport())->IncrementMsgCounter();
 }
 
 Message::Message(Manager& manager, UnmanagedRegionPtr& region, void* data, const size_t size, void* hint, FairMQTransportFactory* factory)
@@ -95,6 +100,7 @@ Message::Message(Manager& manager, UnmanagedRegionPtr& region, void* data, const
         LOG(error) << "trying to create region message with data from outside the region";
         throw runtime_error("trying to create region message with data from outside the region");
     }
+    static_cast<TransportFactory*>(GetTransport())->IncrementMsgCounter();
 }
 
 bool Message::InitializeChunk(const size_t size)
@@ -225,6 +231,8 @@ void Message::CloseMessage()
             }
         }
     }
+
+    static_cast<TransportFactory*>(GetTransport())->DecrementMsgCounter();
 }
 
 }

fairmq/shmem/Socket.cxx

@@ -10,6 +10,7 @@
 #include "Socket.h"
 #include "Message.h"
 #include "UnmanagedRegion.h"
+#include "TransportFactory.h"
 
 #include <FairMQLogger.h>
 #include <fairmq/Tools.h>

fairmq/shmem/TransportFactory.cxx

@@ -47,6 +47,7 @@ TransportFactory::TransportFactory(const string& id, const ProgOptions* config)
     , fManager(nullptr)
     , fHeartbeatThread()
     , fSendHeartbeats(true)
+    , fMsgCounter(0)
 {
     int major, minor, patch;
     zmq_version(&major, &minor, &patch);
@@ -168,6 +169,15 @@ Transport TransportFactory::GetType() const
     return fTransportType;
 }
 
+void TransportFactory::Reset()
+{
+    if (fMsgCounter.load() != 0) {
+        LOG(error) << "Message counter during Reset expected to be 0, found: " << fMsgCounter.load();
+        throw MessageError(tools::ToString("Message counter during Reset expected to be 0, found: ", fMsgCounter.load()));
+    }
+}
+
+
 TransportFactory::~TransportFactory()
 {
     LOG(debug) << "Destroying Shared Memory transport...";

fairmq/shmem/TransportFactory.h

@@ -55,7 +55,10 @@ class TransportFactory final : public fair::mq::TransportFactory
 
     void Interrupt() override { Socket::Interrupt(); }
     void Resume() override { Socket::Resume(); }
-    void Reset() override {}
+    void Reset() override;
+
+    void IncrementMsgCounter() { ++fMsgCounter; }
+    void DecrementMsgCounter() { --fMsgCounter; }
 
     ~TransportFactory() override;
 
@@ -69,6 +72,7 @@ class TransportFactory final : public fair::mq::TransportFactory
     std::unique_ptr<Manager> fManager;
     std::thread fHeartbeatThread;
     std::atomic<bool> fSendHeartbeats;
+    std::atomic<int32_t> fMsgCounter;
 };
 
 } // namespace shmem

test/device/_error_state.cxx

@@ -11,6 +11,7 @@
 #include <gtest/gtest.h>
 #include <boost/process.hpp>
 #include <fairmq/Tools.h>
+#include <FairMQDevice.h>
 
 #include <string>
 #include <thread>
@@ -23,6 +24,39 @@ using namespace std;
 using namespace fair::mq::test;
 using namespace fair::mq::tools;
 
+class BadDevice : public FairMQDevice
+{
+  public:
+    BadDevice()
+    {
+        fDeviceThread = thread([&](){
+            EXPECT_THROW(RunStateMachine(), fair::mq::MessageError);
+        });
+
+        SetTransport("shmem");
+
+        ChangeState(fair::mq::Transition::InitDevice);
+        WaitForState(fair::mq::State::InitializingDevice);
+        ChangeState(fair::mq::Transition::CompleteInit);
+        WaitForState(fair::mq::State::Initialized);
+
+        parts.AddPart(NewMessage());
+    }
+
+    ~BadDevice()
+    {
+        ChangeState(fair::mq::Transition::ResetDevice);
+
+        if (fDeviceThread.joinable()) {
+            fDeviceThread.join();
+        }
+    }
+
+  private:
+    thread fDeviceThread;
+    FairMQParts parts;
+};
+
 void RunErrorStateIn(const string& state, const string& control, const string& input = "")
 {
     size_t session{fair::mq::tools::UuidHash()};
@@ -118,4 +152,9 @@ TEST(ErrorState, interactive_InReset)
     EXPECT_EXIT(RunErrorStateIn("Reset", "interactive", "q"), ::testing::ExitedWithCode(1), "");
 }
 
+TEST(ErrorState, OrphanMessages)
+{
+    BadDevice badDevice;
+}
+
 } // namespace