From 424e22b41a414a24a02058c3ce9ffb14037e813f Mon Sep 17 00:00:00 2001
From: Alexey Rybalchenko <alexryba@gmail.com>
Date: Wed, 15 Nov 2023 10:10:42 +0100
Subject: [PATCH] shm: Throw RefCountBadAlloc if insufficient space in the ref
 count segment

---
 fairmq/Message.h       |  5 +++++
 fairmq/shmem/Message.h | 10 +++++++++-
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/fairmq/Message.h b/fairmq/Message.h
index f065085e..9ce28c96 100644
--- a/fairmq/Message.h
+++ b/fairmq/Message.h
@@ -76,6 +76,11 @@ struct MessageBadAlloc : std::runtime_error
     using std::runtime_error::runtime_error;
 };
 
+struct RefCountBadAlloc : std::runtime_error
+{
+    using std::runtime_error::runtime_error;
+};
+
 }   // namespace fair::mq
 
 using fairmq_free_fn [[deprecated("Use fair::mq::FreeFn")]] = fair::mq::FreeFn;
diff --git a/fairmq/shmem/Message.h b/fairmq/shmem/Message.h
index 04d0d75c..acaf4200 100644
--- a/fairmq/shmem/Message.h
+++ b/fairmq/shmem/Message.h
@@ -279,7 +279,15 @@ class Message final : public fair::mq::Message
             }
             if (otherMsg.fShared < 0) {
                 // UR msg not yet shared, create the reference counting object with count 2
-                otherMsg.fShared = fRegionPtr->HandleFromAddress(&(fRegionPtr->MakeRefCount(2)));
+                try {
+                    otherMsg.fShared = fRegionPtr->HandleFromAddress(&(fRegionPtr->MakeRefCount(2)));
+                } catch (boost::interprocess::bad_alloc& ba) {
+                    throw RefCountBadAlloc(tools::ToString(
+                        "Insufficient space in the reference count segment ",
+                        otherMsg.fRegionId,
+                        ", original exception: bad_alloc: ", 
+                        ba.what()));
+                }
             } else {
                 fRegionPtr->GetRefCountAddressFromHandle(otherMsg.fShared)->Increment();
             }