Clear up ownership for shm segments and ringbuffers

Summary:
I found it odd that we were using shared_ptrs to store some parts of our ringbuffer/shm-segment code. I can see two reasons:
- It allows us to have both the consumer and producer hold a shared_ptr to the ringbuffer, so that one could have multiple of them sharing the same ringbuffer, which would be kept alive as long as there were users. In practice I don't think this mattered: each ringbuffer had only one user, both "in the code" (one consumer or one producer, never both) and "logically" (it was owned by one side of the connection, or by reactor, ... and destroyed when they were gone). shared_ptrs weren't buying us anything and in fact were weakening the ownership model.
- Its aliasing constructor allowed us to bind the lifetime of the shared memory segment to the one of the ringbuffer object that was stored in it, thus allowing us to keep a bunch of things alive by holding just one pointer.

So in the end this resulted in a lot of implicit links, with the end user (the connection, the reactor, ...) holding just a shared_ptr to a consumer/producer and this in turn owning (in a pretend-shared way) a lot of resources, including mmapped memory and file descriptors.

Here I try to make the ownership model more explicit, by separating the resource (the Segment class, which now becomes a movable RAII wrapper which owns the fd and the mmap) from the shallow and stateless helpers used to access it (the ringbuffer, which is a pair of pointers, and the consumer/producer, which are little more than a collection of methods). Consumers/producers are in fact now so simple that they don't have a common base class anymore and are created on-demand when needed.

This removes the usage of shared_ptrs entirely, and I think clarifies the ownership and gives us stronger guarantees that resources will be cleaned up and when that'll occur.

Also, it allows us to store the ringbuffer header on the stack or as a field of another object without having to do some shared_ptr contortions (an aliased shared_ptr with a custom empty destructor). This will come in handy later, for InfiniBand.

Reviewed By: beauby

Differential Revision: D23567125

fbshipit-source-id: c229c7f96655324787eda02c8545ab7bafadb885

Author: lw

tensorpipe/test/util/ringbuffer/ringbuffer_test.cc

@@ -22,20 +22,22 @@
 #include <gtest/gtest.h>
 
 using namespace tensorpipe::util::ringbuffer;
+using namespace tensorpipe::util::shm;
 using namespace tensorpipe::transport::shm;
 
 // Same process produces and consumes share memory through different mappings.
 TEST(ShmRingBuffer, SameProducerConsumer) {
-  // This must stay alive for the file descriptors to remain open.
-  std::shared_ptr<RingBuffer> producer_rb;
   int header_fd = -1;
   int data_fd = -1;
   {
     // Producer part.
     // Buffer large enough to fit all data and persistent
     // (needs to be unlinked up manually).
-    std::tie(header_fd, data_fd, producer_rb) = shm::create(256 * 1024);
-    Producer prod{producer_rb};
+    Segment header_segment;
+    Segment data_segment;
+    RingBuffer rb;
+    std::tie(header_segment, data_segment, rb) = shm::create(256 * 1024);
+    Producer prod{rb};
 
     // Producer loop. It all fits in buffer.
     int i = 0;
@@ -44,12 +46,20 @@ TEST(ShmRingBuffer, SameProducerConsumer) {
       EXPECT_EQ(ret, sizeof(i));
       ++i;
     }
+
+    // Duplicate the file descriptors so that the shared memory remains alive
+    // when the original fds are closed by the segments' destructors.
+    header_fd = ::dup(header_segment.getFd());
+    data_fd = ::dup(data_segment.getFd());
   }
 
   {
     // Consumer part.
     // Map file again (to a different address) and consume it.
-    auto rb = shm::load(header_fd, data_fd);
+    Segment header_segment;
+    Segment data_segment;
+    RingBuffer rb;
+    std::tie(header_segment, data_segment, rb) = shm::load(header_fd, data_fd);
     Consumer cons{rb};
 
     int i = 0;
@@ -84,16 +94,17 @@ TEST(ShmRingBuffer, SingleProducer_SingleConsumer) {
 
   if (pid == 0) {
     // child, the producer
-    // Make a scope so shared_ptr's are released even on exit(0).
+    // Make a scope so segments are destroyed even on exit(0).
     {
-      int header_fd;
-      int data_fd;
-      std::shared_ptr<RingBuffer> rb;
-      std::tie(header_fd, data_fd, rb) = shm::create(1024);
+      Segment header_segment;
+      Segment data_segment;
+      RingBuffer rb;
+      std::tie(header_segment, data_segment, rb) = shm::create(1024);
       Producer prod{rb};
 
       {
-        auto err = sendFdsToSocket(sock_fds[0], header_fd, data_fd);
+        auto err = sendFdsToSocket(
+            sock_fds[0], header_segment.getFd(), data_segment.getFd());
         if (err) {
           TP_THROW_ASSERT() << err.what();
         }
@@ -134,7 +145,10 @@ TEST(ShmRingBuffer, SingleProducer_SingleConsumer) {
       TP_THROW_ASSERT() << err.what();
     }
   }
-  auto rb = shm::load(header_fd, data_fd);
+  Segment header_segment;
+  Segment data_segment;
+  RingBuffer rb;
+  std::tie(header_segment, data_segment, rb) = shm::load(header_fd, data_fd);
   Consumer cons{rb};
 
   int i = 0;

tensorpipe/test/util/ringbuffer/shm_ringbuffer_test.cc

@@ -29,24 +29,29 @@ TEST(Segment, SameProducerConsumer_Scalar) {
   sched_setaffinity(0, sizeof(cpu_set_t), &cpuset);
 
   // This must stay alive for the file descriptor to remain open.
-  std::shared_ptr<Segment> producer_segment;
+  int fd = -1;
   {
     // Producer part.
-    std::shared_ptr<int> my_int_ptr;
-    std::tie(my_int_ptr, producer_segment) =
+    Segment segment;
+    int* my_int_ptr;
+    std::tie(segment, my_int_ptr) =
         Segment::create<int>(true, PageType::Default);
     int& my_int = *my_int_ptr;
     my_int = 1000;
+
+    // Duplicate the file descriptor so that the shared memory remains alive
+    // when the original fd is closed by the segment's destructor.
+    fd = ::dup(segment.getFd());
   }
 
   {
     // Consumer part.
     // Map file again (to a different address) and consume it.
-    std::shared_ptr<int> my_int_ptr;
-    std::shared_ptr<Segment> segment;
-    std::tie(my_int_ptr, segment) =
-        Segment::load<int>(producer_segment->getFd(), true, PageType::Default);
-    EXPECT_EQ(segment->getSize(), sizeof(int));
+    Segment segment;
+    int* my_int_ptr;
+    std::tie(segment, my_int_ptr) =
+        Segment::load<int>(fd, true, PageType::Default);
+    EXPECT_EQ(segment.getSize(), sizeof(int));
     EXPECT_EQ(*my_int_ptr, 1000);
   }
 };
@@ -78,17 +83,17 @@ TEST(SegmentManager, SingleProducer_SingleConsumer_Array) {
     {
       // use huge pages in creation and not in loading. This should only affects
       // TLB overhead.
-      std::shared_ptr<float> my_floats;
-      std::shared_ptr<Segment> segment;
-      std::tie(my_floats, segment) =
+      Segment segment;
+      float* my_floats;
+      std::tie(segment, my_floats) =
           Segment::create<float[]>(num_floats, true, PageType::HugeTLB_2MB);
 
       for (int i = 0; i < num_floats; ++i) {
-        my_floats.get()[i] = i;
+        my_floats[i] = i;
       }
 
       {
-        auto err = sendFdsToSocket(sock_fds[0], segment->getFd());
+        auto err = sendFdsToSocket(sock_fds[0], segment.getFd());
         if (err) {
           TP_THROW_ASSERT() << err.what();
         }
@@ -112,13 +117,13 @@ TEST(SegmentManager, SingleProducer_SingleConsumer_Array) {
       TP_THROW_ASSERT() << err.what();
     }
   }
-  std::shared_ptr<float> my_floats;
-  std::shared_ptr<Segment> segment;
-  std::tie(my_floats, segment) =
+  Segment segment;
+  float* my_floats;
+  std::tie(segment, my_floats) =
       Segment::load<float[]>(segment_fd, false, PageType::Default);
-  EXPECT_EQ(num_floats * sizeof(float), segment->getSize());
+  EXPECT_EQ(num_floats * sizeof(float), segment.getSize());
   for (int i = 0; i < num_floats; ++i) {
-    EXPECT_EQ(my_floats.get()[i], i);
+    EXPECT_EQ(my_floats[i], i);
   }
   {
     uint64_t c = 1;

tensorpipe/test/util/shm/segment_test.cc

@@ -361,13 +361,15 @@ class Connection::Impl : public std::enable_shared_from_this<Connection::Impl>,
   ClosingReceiver closingReceiver_;
 
   // Inbox.
-  int inboxHeaderFd_;
-  int inboxDataFd_;
-  optional<util::ringbuffer::Consumer> inbox_;
+  util::shm::Segment inboxHeaderSegment_;
+  util::shm::Segment inboxDataSegment_;
+  util::ringbuffer::RingBuffer inboxRb_;
   optional<Reactor::TToken> inboxReactorToken_;
 
   // Outbox.
-  optional<util::ringbuffer::Producer> outbox_;
+  util::shm::Segment outboxHeaderSegment_;
+  util::shm::Segment outboxDataSegment_;
+  util::ringbuffer::RingBuffer outboxRb_;
   optional<Reactor::TToken> outboxReactorToken_;
 
   // Peer trigger/tokens.
@@ -514,10 +516,8 @@ void Connection::Impl::initFromLoop() {
   }
 
   // Create ringbuffer for inbox.
-  std::shared_ptr<util::ringbuffer::RingBuffer> inboxRingBuffer;
-  std::tie(inboxHeaderFd_, inboxDataFd_, inboxRingBuffer) =
+  std::tie(inboxHeaderSegment_, inboxDataSegment_, inboxRb_) =
       util::ringbuffer::shm::create(kBufferSize);
-  inbox_.emplace(std::move(inboxRingBuffer));
 
   // Register method to be called when our peer writes to our inbox.
   inboxReactorToken_ = context_->addReaction(runIfAlive(*this, [](Impl& impl) {
@@ -919,8 +919,9 @@ void Connection::Impl::handleEventInFromLoop() {
     }
 
     // Load ringbuffer for outbox.
-    outbox_.emplace(util::ringbuffer::shm::load(
-        outboxHeaderFd.release(), outboxDataFd.release()));
+    std::tie(outboxHeaderSegment_, outboxDataSegment_, outboxRb_) =
+        util::ringbuffer::shm::load(
+            outboxHeaderFd.release(), outboxDataFd.release());
 
     // Initialize remote reactor trigger.
     peerReactorTrigger_.emplace(
@@ -963,8 +964,8 @@ void Connection::Impl::handleEventOutFromLoop() {
         outboxReactorToken_.value(),
         reactorHeaderFd,
         reactorDataFd,
-        inboxHeaderFd_,
-        inboxDataFd_);
+        inboxHeaderSegment_.getFd(),
+        inboxDataSegment_.getFd());
     if (err) {
       setError_(std::move(err));
       return;
@@ -988,10 +989,11 @@ void Connection::Impl::processReadOperationsFromLoop() {
     return;
   }
   // Serve read operations
+  util::ringbuffer::Consumer inboxConsumer(inboxRb_);
   while (!readOperations_.empty()) {
     auto readOperation = std::move(readOperations_.front());
     readOperations_.pop_front();
-    if (readOperation.handleRead(*inbox_)) {
+    if (readOperation.handleRead(inboxConsumer)) {
       peerReactorTrigger_->run(peerOutboxReactorToken_.value());
     }
     if (!readOperation.completed()) {
@@ -1008,10 +1010,11 @@ void Connection::Impl::processWriteOperationsFromLoop() {
     return;
   }
 
+  util::ringbuffer::Producer outboxProducer(outboxRb_);
   while (!writeOperations_.empty()) {
     auto writeOperation = std::move(writeOperations_.front());
     writeOperations_.pop_front();
-    if (writeOperation.handleWrite(*outbox_)) {
+    if (writeOperation.handleWrite(outboxProducer)) {
       peerReactorTrigger_->run(peerInboxReactorToken_.value());
     }
     if (!writeOperation.completed()) {

tensorpipe/transport/shm/connection.cc

@@ -43,14 +43,8 @@ void writeToken(util::ringbuffer::Producer& producer, Reactor::TToken token) {
 } // namespace
 
 Reactor::Reactor() {
-  int headerFd;
-  int dataFd;
-  std::shared_ptr<util::ringbuffer::RingBuffer> rb;
-  std::tie(headerFd, dataFd, rb) = util::ringbuffer::shm::create(kSize);
-  headerFd_ = Fd(headerFd);
-  dataFd_ = Fd(dataFd);
-  consumer_.emplace(rb);
-  producer_.emplace(rb);
+  std::tie(headerSegment_, dataSegment_, rb_) =
+      util::ringbuffer::shm::create(kSize);
   thread_ = std::thread(&Reactor::run, this);
 }
 
@@ -106,23 +100,19 @@ void Reactor::remove(TToken token) {
   functionCount_--;
 }
 
-void Reactor::trigger(TToken token) {
-  std::unique_lock<std::mutex> lock(mutex_);
-  writeToken(producer_.value(), token);
-}
-
 std::tuple<int, int> Reactor::fds() const {
-  return std::make_tuple(headerFd_.fd(), dataFd_.fd());
+  return std::make_tuple(headerSegment_.getFd(), dataSegment_.getFd());
 }
 
 void Reactor::run() {
   setThreadName("TP_SHM_reactor");
 
+  util::ringbuffer::Consumer reactorConsumer(rb_);
   // Stop when another thread has asked the reactor the close and when
   // all functions have been removed.
   while (!closed_ || functionCount_ > 0) {
     uint32_t token;
-    auto ret = consumer_->read(&token, sizeof(token));
+    auto ret = reactorConsumer.read(&token, sizeof(token));
     if (ret == -ENODATA) {
       if (deferredFunctionCount_ > 0) {
         decltype(deferredFunctionList_) fns;
@@ -179,15 +169,16 @@ void Reactor::run() {
   }
 }
 
-Reactor::Trigger::Trigger(Fd&& headerFd, Fd&& dataFd)
-    : producer_(util::ringbuffer::shm::load(
-          // The header and data segment objects take over ownership
-          // of file descriptors. Release them to avoid double close.
-          headerFd.release(),
-          dataFd.release())) {}
+Reactor::Trigger::Trigger(Fd&& headerFd, Fd&& dataFd) {
+  // The header and data segment objects take over ownership
+  // of file descriptors. Release them to avoid double close.
+  std::tie(headerSegment_, dataSegment_, rb_) =
+      util::ringbuffer::shm::load(headerFd.release(), dataFd.release());
+}
 
 void Reactor::Trigger::run(TToken token) {
-  writeToken(producer_, token);
+  util::ringbuffer::Producer producer(rb_);
+  writeToken(producer, token);
 }
 
 void Reactor::deferToLoop(TDeferredFunction fn) {

tensorpipe/transport/shm/reactor.cc

@@ -22,6 +22,7 @@
 #include <tensorpipe/transport/shm/fd.h>
 #include <tensorpipe/util/ringbuffer/consumer.h>
 #include <tensorpipe/util/ringbuffer/producer.h>
+#include <tensorpipe/util/shm/segment.h>
 
 namespace tensorpipe {
 namespace transport {
@@ -87,9 +88,6 @@ class Reactor final {
   // Removes function associated with token from reactor.
   void remove(TToken token);
 
-  // Trigger reactor with specified token.
-  void trigger(TToken token);
-
   // Returns the file descriptors for the underlying ring buffer.
   std::tuple<int, int> fds() const;
 
@@ -110,10 +108,9 @@ class Reactor final {
   ~Reactor();
 
  private:
-  Fd headerFd_;
-  Fd dataFd_;
-  optional<util::ringbuffer::Consumer> consumer_;
-  optional<util::ringbuffer::Producer> producer_;
+  util::shm::Segment headerSegment_;
+  util::shm::Segment dataSegment_;
+  util::ringbuffer::RingBuffer rb_;
 
   std::mutex mutex_;
   std::thread thread_;
@@ -162,7 +159,9 @@ class Reactor final {
     void run(TToken token);
 
    private:
-    util::ringbuffer::Producer producer_;
+    util::shm::Segment headerSegment_;
+    util::shm::Segment dataSegment_;
+    util::ringbuffer::RingBuffer rb_;
   };
 };
 

tensorpipe/transport/shm/reactor.h

@@ -19,15 +19,21 @@ namespace ringbuffer {
 ///
 /// Provides methods to read data from a ringbuffer.
 ///
-class Consumer : public RingBufferWrapper {
+class Consumer {
  public:
-  // Use base class constructor.
-  using RingBufferWrapper::RingBufferWrapper;
+  Consumer() = delete;
+
+  Consumer(RingBuffer& rb) : header_{rb.getHeader()}, data_{rb.getData()} {
+    TP_THROW_IF_NULLPTR(data_);
+  }
 
   Consumer(const Consumer&) = delete;
   Consumer(Consumer&&) = delete;
 
-  virtual ~Consumer() noexcept {
+  Consumer& operator=(const Consumer&) = delete;
+  Consumer& operator=(Consumer&&) = delete;
+
+  ~Consumer() noexcept {
     TP_THROW_ASSERT_IF(inTx());
   }
 
@@ -192,7 +198,10 @@ class Consumer : public RingBufferWrapper {
     return size;
   }
 
- protected:
+ private:
+  RingBufferHeader& header_;
+  const uint8_t* const data_;
+  unsigned tx_size_ = 0;
   bool inTx_{false};
 };
 

tensorpipe/util/ringbuffer/consumer.h

@@ -19,15 +19,21 @@ namespace ringbuffer {
 ///
 /// Provides methods to write data into a ringbuffer.
 ///
-class Producer : public RingBufferWrapper {
+class Producer {
  public:
-  // Use base class constructor.
-  using RingBufferWrapper::RingBufferWrapper;
+  Producer() = delete;
+
+  Producer(RingBuffer& rb) : header_{rb.getHeader()}, data_{rb.getData()} {
+    TP_THROW_IF_NULLPTR(data_);
+  }
 
   Producer(const Producer&) = delete;
   Producer(Producer&&) = delete;
 
-  virtual ~Producer() noexcept {
+  Producer& operator=(const Producer&) = delete;
+  Producer& operator=(Producer&&) = delete;
+
+  ~Producer() noexcept {
     TP_THROW_ASSERT_IF(inTx());
   }
 
@@ -196,7 +202,10 @@ class Producer : public RingBufferWrapper {
     return size;
   }
 
- protected:
+ private:
+  RingBufferHeader& header_;
+  uint8_t* const data_;
+  unsigned tx_size_ = 0;
   bool inTx_{false};
 };
 

tensorpipe/util/ringbuffer/producer.h

@@ -133,13 +133,10 @@ class RingBufferHeader {
 ///
 class RingBuffer final {
  public:
-  RingBuffer(const RingBuffer&) = delete;
-  RingBuffer(RingBuffer&&) = delete;
+  RingBuffer() = default;
 
-  RingBuffer(
-      std::shared_ptr<RingBufferHeader> header,
-      std::shared_ptr<uint8_t> data)
-      : header_{std::move(header)}, data_{std::move(data)} {
+  RingBuffer(RingBufferHeader* header, uint8_t* data)
+      : header_(header), data_(data) {
     TP_THROW_IF_NULLPTR(header_) << "Header cannot be nullptr";
     TP_THROW_IF_NULLPTR(data_) << "Data cannot be nullptr";
   }
@@ -153,51 +150,16 @@ class RingBuffer final {
   }
 
   const uint8_t* getData() const {
-    return data_.get();
+    return data_;
   }
 
   uint8_t* getData() {
-    return data_.get();
+    return data_;
   }
 
  protected:
-  std::shared_ptr<RingBufferHeader> header_;
-
-  // Note: this is a std::shared_ptr<uint8_t[]> semantically.
-  // A shared_ptr with array type is supported in C++17 and higher.
-  std::shared_ptr<uint8_t> data_;
-};
-
-///
-/// Ringbuffer wrapper
-///
-class RingBufferWrapper {
- public:
-  RingBufferWrapper(const RingBufferWrapper&) = delete;
-
-  RingBufferWrapper(std::shared_ptr<RingBuffer> rb)
-      : rb_{rb}, header_{rb->getHeader()}, data_{rb->getData()} {
-    TP_THROW_IF_NULLPTR(rb);
-    TP_THROW_IF_NULLPTR(data_);
-  }
-
-  auto& getHeader() {
-    return header_;
-  }
-
-  const auto& getHeader() const {
-    return header_;
-  }
-
-  auto getRingBuffer() {
-    return rb_;
-  }
-
- protected:
-  std::shared_ptr<RingBuffer> rb_;
-  RingBufferHeader& header_;
-  uint8_t* const data_;
-  unsigned tx_size_ = 0;
+  RingBufferHeader* header_ = nullptr;
+  uint8_t* data_ = nullptr;
 };
 
 } // namespace ringbuffer

tensorpipe/util/ringbuffer/ringbuffer.h

@@ -13,57 +13,55 @@ namespace util {
 namespace ringbuffer {
 namespace shm {
 
-std::tuple<int, int, std::shared_ptr<RingBuffer>> create(
+std::tuple<util::shm::Segment, util::shm::Segment, RingBuffer> create(
     size_t min_rb_byte_size,
-    optional<tensorpipe::util::shm::PageType> data_page_type,
+    optional<util::shm::PageType> data_page_type,
     bool perm_write) {
-  std::shared_ptr<RingBufferHeader> header;
-  std::shared_ptr<tensorpipe::util::shm::Segment> header_segment;
-  std::tie(header, header_segment) =
-      tensorpipe::util::shm::Segment::create<RingBufferHeader>(
-          perm_write,
-          tensorpipe::util::shm::PageType::Default,
-          min_rb_byte_size);
+  util::shm::Segment header_segment;
+  RingBufferHeader* header;
+  std::tie(header_segment, header) =
+      util::shm::Segment::create<RingBufferHeader>(
+          perm_write, util::shm::PageType::Default, min_rb_byte_size);
 
-  std::shared_ptr<uint8_t> data;
-  std::shared_ptr<tensorpipe::util::shm::Segment> data_segment;
-  std::tie(data, data_segment) =
-      tensorpipe::util::shm::Segment::create<uint8_t[]>(
-          header->kDataPoolByteSize, perm_write, data_page_type);
+  util::shm::Segment data_segment;
+  uint8_t* data;
+  std::tie(data_segment, data) = util::shm::Segment::create<uint8_t[]>(
+      header->kDataPoolByteSize, perm_write, data_page_type);
 
   // Note: cannot use implicit construction from initializer list on GCC 5.5:
   // "converting to XYZ from initializer list would use explicit constructor".
   return std::make_tuple(
-      header_segment->getFd(),
-      data_segment->getFd(),
-      std::make_shared<RingBuffer>(std::move(header), std::move(data)));
+      std::move(header_segment),
+      std::move(data_segment),
+      RingBuffer(header, data));
 }
 
-std::shared_ptr<RingBuffer> load(
+std::tuple<util::shm::Segment, util::shm::Segment, RingBuffer> load(
     int header_fd,
     int data_fd,
-    optional<tensorpipe::util::shm::PageType> data_page_type,
+    optional<util::shm::PageType> data_page_type,
     bool perm_write) {
-  std::shared_ptr<RingBufferHeader> header;
-  std::shared_ptr<tensorpipe::util::shm::Segment> header_segment;
-  std::tie(header, header_segment) =
-      tensorpipe::util::shm::Segment::load<RingBufferHeader>(
-          header_fd, perm_write, tensorpipe::util::shm::PageType::Default);
+  util::shm::Segment header_segment;
+  RingBufferHeader* header;
+  std::tie(header_segment, header) = util::shm::Segment::load<RingBufferHeader>(
+      header_fd, perm_write, util::shm::PageType::Default);
   constexpr auto kHeaderSize = sizeof(RingBufferHeader);
-  if (unlikely(kHeaderSize != header_segment->getSize())) {
+  if (unlikely(kHeaderSize != header_segment.getSize())) {
     TP_THROW_SYSTEM(EPERM) << "Header segment of unexpected size";
   }
 
-  std::shared_ptr<uint8_t> data;
-  std::shared_ptr<tensorpipe::util::shm::Segment> data_segment;
-  std::tie(data, data_segment) =
-      tensorpipe::util::shm::Segment::load<uint8_t[]>(
-          data_fd, perm_write, data_page_type);
-  if (unlikely(header->kDataPoolByteSize != data_segment->getSize())) {
+  util::shm::Segment data_segment;
+  uint8_t* data;
+  std::tie(data_segment, data) =
+      util::shm::Segment::load<uint8_t[]>(data_fd, perm_write, data_page_type);
+  if (unlikely(header->kDataPoolByteSize != data_segment.getSize())) {
     TP_THROW_SYSTEM(EPERM) << "Data segment of unexpected size";
   }
 
-  return std::make_shared<RingBuffer>(std::move(header), std::move(data));
+  return std::make_tuple(
+      std::move(header_segment),
+      std::move(data_segment),
+      RingBuffer(header, data));
 }
 
 } // namespace shm

tensorpipe/util/ringbuffer/shm.cc

@@ -18,7 +18,7 @@ namespace shm {
 
 /// Creates ringbuffer on shared memory.
 ///
-/// RingBuffer's data can have any <tensorpipe::util::shm::PageType>
+/// RingBuffer's data can have any <util::shm::PageType>
 /// (e.g. 4KB or a HugeTLB Page of 2MB or 1GB). If  <data_page_type> is not
 /// provided, then choose the largest page that would result in
 /// close to full occupancy.
@@ -29,15 +29,15 @@ namespace shm {
 /// <min_rb_byte_size> is the minimum size of the data section
 /// of a RingBuffer (or each CPU's RingBuffer).
 ///
-std::tuple<int, int, std::shared_ptr<RingBuffer>> create(
+std::tuple<util::shm::Segment, util::shm::Segment, RingBuffer> create(
     size_t min_rb_byte_size,
-    optional<tensorpipe::util::shm::PageType> data_page_type = nullopt,
+    optional<util::shm::PageType> data_page_type = nullopt,
     bool perm_write = true);
 
-std::shared_ptr<RingBuffer> load(
+std::tuple<util::shm::Segment, util::shm::Segment, RingBuffer> load(
     int header_fd,
     int data_fd,
-    optional<tensorpipe::util::shm::PageType> data_page_type = nullopt,
+    optional<util::shm::PageType> data_page_type = nullopt,
     bool perm_write = true);
 
 } // namespace shm

tensorpipe/util/ringbuffer/shm.h

@@ -82,6 +82,21 @@ update to obtain the latest correctness checks."
 // Mention static constexpr char to export the symbol.
 constexpr char Segment::kBasePath[];
 
+Segment::Segment(Segment&& other) noexcept {
+  std::swap(page_type_, other.page_type_);
+  std::swap(fd_, other.fd_);
+  std::swap(byte_size_, other.byte_size_);
+  std::swap(base_ptr_, other.base_ptr_);
+}
+
+Segment& Segment::operator=(Segment&& other) noexcept {
+  std::swap(page_type_, other.page_type_);
+  std::swap(fd_, other.fd_);
+  std::swap(byte_size_, other.byte_size_);
+  std::swap(base_ptr_, other.base_ptr_);
+  return *this;
+}
+
 Segment::Segment(
     size_t byte_size,
     bool perm_write,
@@ -118,17 +133,17 @@ void Segment::mmap(bool perm_write, optional<PageType> page_type) {
 }
 
 Segment::~Segment() {
-  int ret = munmap(base_ptr_, byte_size_);
-  if (ret == -1) {
-    TP_LOG_ERROR() << "Error while munmapping shared memory segment. Error: "
-                   << toErrorCode(errno).message();
+  if (base_ptr_ != nullptr) {
+    int ret = munmap(base_ptr_, byte_size_);
+    if (ret == -1) {
+      TP_LOG_ERROR() << "Error while munmapping shared memory segment. Error: "
+                     << toErrorCode(errno).message();
+    }
   }
-  if (0 > fd_) {
-    TP_LOG_ERROR() << "Attempt to destroy segment with negative file "
-                   << "descriptor";
-    return;
+  if (fd_ >= 0) {
+    int ret = ::close(fd_);
+    TP_THROW_SYSTEM_IF(ret != 0, errno);
   }
-  ::close(fd_);
 }
 
 } // namespace shm

tensorpipe/util/shm/segment.cc

@@ -61,20 +61,25 @@ class Segment {
   // Default base path for all segments created.
   static constexpr char kBasePath[] = "/dev/shm";
 
+  Segment() = default;
+
   Segment(const Segment&) = delete;
-  Segment(Segment&&) = delete;
+  Segment(Segment&&) noexcept;
+
+  Segment& operator=(const Segment& other) = delete;
+  Segment& operator=(Segment&& other) noexcept;
 
   Segment(size_t byte_size, bool perm_write, optional<PageType> page_type);
 
   Segment(int fd, bool perm_write, optional<PageType> page_type);
 
-  /// Create read and size shared memory to contain an object of class T.
+  /// Allocate shared memory to contain an object of type T and construct it.
   ///
-  /// The created object's shared_ptr will own the lifetime of the
-  /// Segment and will call Segment destructor.
-  /// Caller can use the shared_ptr to the underlying Segment.
+  /// The Segment object owns the memory and frees it when destructed.
+  /// The raw pointer to the object provides a view into the Segment but doesn't
+  /// own it and may thus become invalid if the Segment isn't kept alive.
   template <class T, class... Args>
-  static std::pair<std::shared_ptr<T>, std::shared_ptr<Segment>> create(
+  static std::pair<Segment, T*> create(
       bool perm_write,
       optional<PageType> page_type,
       Args&&... args) {
@@ -84,25 +89,24 @@ class Segment {
     static_assert(std::is_trivially_copyable<T>::value, "!");
 
     const auto byte_size = sizeof(T);
-    auto segment = std::make_shared<Segment>(byte_size, perm_write, page_type);
-    TP_DCHECK_EQ(segment->getSize(), byte_size);
+    Segment segment(byte_size, perm_write, page_type);
+    TP_DCHECK_EQ(segment.getSize(), byte_size);
 
     // Initialize in place. Forward T's constructor arguments.
-    T* ptr = new (segment->getPtr()) T(std::forward<Args>(args)...);
-    if (ptr != segment->getPtr()) {
+    T* ptr = new (segment.getPtr()) T(std::forward<Args>(args)...);
+    if (ptr != segment.getPtr()) {
       TP_THROW_SYSTEM(EPERM)
-          << "new's address cannot be different from segment->getPtr() "
+          << "new's address cannot be different from segment.getPtr() "
           << " address. Some aligment assumption was incorrect";
     }
 
-    return {std::shared_ptr<T>(segment, ptr), segment};
+    return {std::move(segment), ptr};
   }
 
   /// One-dimensional array version of create<T, ...Args>.
-  /// Caller can use the shared_ptr to the underlying Segment.
   // XXX: Fuse all versions of create.
   template <class T, typename TScalar = typename std::remove_extent<T>::type>
-  static std::pair<std::shared_ptr<TScalar>, std::shared_ptr<Segment>> create(
+  static std::pair<Segment, TScalar*> create(
       size_t num_elements,
       bool perm_write,
       optional<PageType> page_type) {
@@ -125,55 +129,48 @@ class Segment {
     static_assert(std::is_same<TScalar[], T>::value, "Type mismatch");
 
     size_t byte_size = sizeof(TScalar) * num_elements;
-    auto segment = std::make_shared<Segment>(byte_size, perm_write, page_type);
-    TP_DCHECK_EQ(segment->getSize(), byte_size);
+    Segment segment(byte_size, perm_write, page_type);
+    TP_DCHECK_EQ(segment.getSize(), byte_size);
 
     // Initialize in place.
-    TScalar* ptr = new (segment->getPtr()) TScalar[num_elements]();
-    if (ptr != segment->getPtr()) {
+    TScalar* ptr = new (segment.getPtr()) TScalar[num_elements]();
+    if (ptr != segment.getPtr()) {
       TP_THROW_SYSTEM(EPERM)
-          << "new's address cannot be different from segment->getPtr() "
+          << "new's address cannot be different from segment.getPtr() "
           << " address. Some aligment assumption was incorrect";
     }
 
-    return {std::shared_ptr<TScalar>(segment, ptr), segment};
+    return {std::move(segment), ptr};
   }
 
-  /// Load an already created shared memory Segment that holds an
-  /// object of type T, where T is an array type.
-  ///
-  /// Lifecycle of shared_ptr and Segment's reference_wrapper is
-  /// identical to create<>().
+  /// Load an existing shared memory region that already holds an object of type
+  /// T, where T is an array type.
   template <
       class T,
       typename TScalar = typename std::remove_extent<T>::type,
       std::enable_if_t<std::is_array<T>::value, int> = 0>
-  static std::pair<std::shared_ptr<TScalar>, std::shared_ptr<Segment>> load(
+  static std::pair<Segment, TScalar*> load(
       int fd,
       bool perm_write,
       optional<PageType> page_type) {
-    auto segment = std::make_shared<Segment>(fd, perm_write, page_type);
-    const size_t size = segment->getSize();
+    Segment segment(fd, perm_write, page_type);
     static_assert(
         std::rank<T>::value == 1,
         "Currently only rank one arrays are supported");
     static_assert(std::is_trivially_copyable<TScalar>::value, "!");
-    auto ptr = static_cast<TScalar*>(segment->getPtr());
-    return {std::shared_ptr<TScalar>(segment, ptr), segment};
+    auto ptr = static_cast<TScalar*>(segment.getPtr());
+    return {std::move(segment), ptr};
   }
 
-  /// Load an already created shared memory Segment that holds an
-  /// object of type T, where T is NOT an array type.
-  ///
-  /// Lifecycle of shared_ptr and Segment's reference_wrapper is
-  /// identical to create<>().
+  /// Load an existing shared memory region that already holds an object of type
+  /// T, where T is NOT an array type.
   template <class T, std::enable_if_t<!std::is_array<T>::value, int> = 0>
-  static std::pair<std::shared_ptr<T>, std::shared_ptr<Segment>> load(
+  static std::pair<Segment, T*> load(
       int fd,
       bool perm_write,
       optional<PageType> page_type) {
-    auto segment = std::make_shared<Segment>(fd, perm_write, page_type);
-    const size_t size = segment->getSize();
+    Segment segment(fd, perm_write, page_type);
+    const size_t size = segment.getSize();
     // XXX: Do some checking other than the size that we are loading
     // the right type.
     if (size != sizeof(T)) {
@@ -184,8 +181,8 @@ class Segment {
           << "consider linking segment after it has been fully initialized.";
     }
     static_assert(std::is_trivially_copyable<T>::value, "!");
-    auto ptr = static_cast<T*>(segment->getPtr());
-    return {std::shared_ptr<T>(segment, ptr), segment};
+    auto ptr = static_cast<T*>(segment.getPtr());
+    return {std::move(segment), ptr};
   }
 
   const int getFd() const {
@@ -212,16 +209,16 @@ class Segment {
 
  protected:
   // The page used to mmap the segment.
-  PageType page_type_;
+  PageType page_type_ = PageType::Default;
 
   // The file descriptor of the shared memory file.
   int fd_ = -1;
 
   // Byte size of shared memory segment.
-  size_t byte_size_;
+  size_t byte_size_ = 0;
 
   // Base pointer of mmmap'ed shared memory segment.
-  void* base_ptr_;
+  void* base_ptr_ = nullptr;
 
   void mmap(bool perm_write, optional<PageType> page_type);
 };