Internal change

PiperOrigin-RevId: 506982474

mediapipe/framework/api2/packet.h

@@ -242,6 +242,8 @@ class Packet : public Packet<internal::Generic> {
   friend Packet<U> PacketAdopting(const U* ptr);
   template <typename U>
   friend Packet<U> PacketAdopting(std::unique_ptr<U> ptr);
+  template <typename U>
+  friend Packet<U> PacketSharingOwnership(std::shared_ptr<const U> ptr);
 };
 
 namespace internal {
@@ -464,6 +466,17 @@ Packet<T> PacketAdopting(std::unique_ptr<T> ptr) {
   return Packet<T>(std::make_shared<packet_internal::Holder<T>>(ptr.release()));
 }
 
+template <typename T>
+Packet<T> PacketSharingOwnership(std::shared_ptr<const T> ptr) {
+  return Packet<T>(
+      std::make_shared<packet_internal::ForeignHolder<T>>(std::move(ptr)));
+}
+
+template <typename T>
+std::shared_ptr<const T> SharedPtrWithPacket(Packet<T> packet) {
+  return mediapipe::SharedPtrWithPacket<T>(std::move(packet));
+}
+
 }  // namespace api2
 }  // namespace mediapipe
 

mediapipe/framework/api2/packet_test.cc

@@ -162,6 +162,21 @@ TEST(PacketTest, PacketAdopting) {
   EXPECT_FALSE(p.IsEmpty());
 }
 
+TEST(PacketTest, PacketSharingOwnership) {
+  bool deleted = false;
+  std::shared_ptr<const int> object(new int(42), [&deleted](const int* p) {
+    delete p;
+    deleted = true;
+  });
+  Packet<int> p = PacketSharingOwnership(object);
+  EXPECT_FALSE(p.IsEmpty());
+  EXPECT_EQ(p.Get(), 42);
+  object = nullptr;
+  EXPECT_FALSE(deleted);  // Packet keeps it alive.
+  p = {};
+  ASSERT_TRUE(deleted);  // last owner expired.
+}
+
 TEST(PacketTest, PacketGeneric) {
   // With C++17, Packet<> could be written simply as Packet.
   Packet<> p = PacketAdopting(new float(1.0));
@@ -281,6 +296,24 @@ TEST(PacketTest, PolymorphismAbstract) {
   EXPECT_EQ(base->name(), "ConcreteDerived");
 }
 
+TEST(PacketTest, ShareSubobjectOwnership) {
+  // Create a packet that contains a vector and tracks deletion.
+  bool deleted = false;
+  std::shared_ptr<const std::vector<int>> ints(new std::vector<int>{0, 1, 2, 3},
+                                               [&deleted](std::vector<int>* p) {
+                                                 delete p;
+                                                 deleted = true;
+                                               });
+  auto vector_packet = PacketSharingOwnership(std::move(ints));
+  // Create a packet that references one of the items in the vector.
+  Packet<int> item_packet = PacketSharingOwnership(std::shared_ptr<const int>(
+      SharedPtrWithPacket(vector_packet), &vector_packet.Get()[1]));
+  vector_packet = {};
+  ASSERT_FALSE(deleted);  // item_packet keeps it alive
+  item_packet = {};
+  ASSERT_TRUE(deleted);
+}
+
 }  // namespace
 }  // namespace api2
 }  // namespace mediapipe

mediapipe/framework/packet.h

@@ -242,6 +242,12 @@ Packet Adopt(const T* ptr);
 // returned Packet but also all of its copies. The timestamp of the returned
 // Packet is Timestamp::Unset(). To set the timestamp, the caller should do
 // PointToForeign(...).At(...).
+// TODO: deprecate PointToForeign in favor of
+// MakePacketSharingOwnership. Currently, we have to provide two separate
+// implementations for handling static array vs. non static array types as
+// the shared_ptr does not work with static array for backward compatibility.
+// Eventually we should encourage the clients to deprecate the usage of these
+// functions.
 template <typename T>
 Packet PointToForeign(const T* ptr);
 
@@ -324,6 +330,17 @@ Packet MakePacket(Args&&... args) {  // NOLINT(build/c++11)
       new T{std::forward<typename std::remove_extent<T>::type>(args)...}));
 }
 
+// Returns a Packet that shares ownership of its data. The packet will hold a
+// reference to the provided shared_ptr throughout its lifetime. Since the
+// payload of packets is expected to be immutable, the caller MUST ensure that
+// the data does not change as long as the Packet is alive.
+// Unlike PointToForeign, which takes a raw pointer, this allows the caller to
+// know when MediaPipe (as well as any other owners) is done using the data.
+// The timestamp of the returned Packet is Timestamp::Unset(). To set the
+// timestamp, the caller should do MakePacketSharingOwnership(...).At(...).
+template <typename T>
+Packet MakePacketSharingOwnership(std::shared_ptr<const T> ptr);
+
 // Returns a mutable pointer to the data in a unique_ptr in a packet. This
 // is useful in combination with AdoptAsUniquePtr.  The caller must
 // exercise caution when mutating the retrieved data, since the data
@@ -579,11 +596,14 @@ class Holder : public HolderBase {
   }
 };
 
-// Like Holder, but does not own its data.
+// Like Holder, but does not exclusively own its data.
 template <typename T>
 class ForeignHolder : public Holder<T> {
  public:
-  using Holder<T>::Holder;
+  explicit ForeignHolder(std::shared_ptr<const T> ptr)
+      : Holder<T>(reinterpret_cast<const T*>(ptr.get())),
+        owner_(std::move(ptr)) {}
+
   ~ForeignHolder() override {
     // Null out ptr_ so it doesn't get deleted by ~Holder.
     // Note that ~Holder cannot call HasForeignOwner because the subclass's
@@ -591,6 +611,9 @@ class ForeignHolder : public Holder<T> {
     this->ptr_ = nullptr;
   }
   bool HasForeignOwner() const final { return true; }
+
+ protected:
+  const std::shared_ptr<const T> owner_;
 };
 
 template <typename T>
@@ -768,7 +791,22 @@ Packet Adopt(const T* ptr) {
 template <typename T>
 Packet PointToForeign(const T* ptr) {
   CHECK(ptr != nullptr);
-  return packet_internal::Create(new packet_internal::ForeignHolder<T>(ptr));
+  using U = typename std::shared_ptr<T>::element_type;
+  // The reinterpret_cast is required here and in the ForeignHolder constructor
+  // in order to handle the type decay introduced by the shared_ptr for the
+  // statically allocated array.
+  return packet_internal::Create(new packet_internal::ForeignHolder<T>(
+      std::shared_ptr<const T>(reinterpret_cast<const U*>(ptr), [](const U*) {
+        // Note: PointToForeign does not own its data in any way, so
+        // the deleter does nothing.
+      })));
+}
+
+template <typename T>
+Packet MakePacketSharingOwnership(std::shared_ptr<const T> ptr) {
+  CHECK(ptr != nullptr);
+  return packet_internal::Create(
+      new packet_internal::ForeignHolder<T>(std::move(ptr)));
 }
 
 // Equal Packets refer to the same memory contents, like equal pointers.

mediapipe/framework/packet_test.cc

@@ -14,6 +14,7 @@
 
 #include "mediapipe/framework/packet.h"
 
+#include <algorithm>
 #include <map>
 #include <memory>
 #include <string>
@@ -262,6 +263,68 @@ TEST(PacketTest, MakePacketOfIntVector) {
             vector_packet2.Get<std::vector<int>>());
 }
 
+TEST(PacketTest, MakePacketSharingOwnership) {
+  bool deleted = false;
+  std::shared_ptr<const int> object(new int(42), [&deleted](const int* p) {
+    delete p;
+    deleted = true;
+  });
+  Packet packet = MakePacketSharingOwnership(object);
+  MP_ASSERT_OK(packet.ValidateAsType<int>());
+  EXPECT_EQ(packet.Get<int>(), 42);
+  EXPECT_FALSE(deleted);
+  object = nullptr;
+  EXPECT_FALSE(deleted);  // Packet keeps it alive.
+  packet = {};
+  EXPECT_TRUE(deleted);  // last owner expired.
+}
+
+TEST(PacketTest, ShareSubobjectOwnership) {
+  // Create a packet that contains a vector and tracks deletion.
+  bool deleted = false;
+  std::shared_ptr<const std::vector<int>> ints(new std::vector<int>{0, 1, 2, 3},
+                                               [&deleted](std::vector<int>* p) {
+                                                 delete p;
+                                                 deleted = true;
+                                               });
+  Packet vector_packet = MakePacketSharingOwnership(std::move(ints));
+  // Create a packet that references one of the items in the vector.
+  Packet item_packet = MakePacketSharingOwnership(std::shared_ptr<const int>(
+      SharedPtrWithPacket<std::vector<int>>(vector_packet),
+      &vector_packet.Get<std::vector<int>>()[1]));
+  vector_packet = {};
+  ASSERT_FALSE(deleted);  // item_packet keeps it alive
+  item_packet = {};
+  ASSERT_TRUE(deleted);
+}
+
+TEST(PacketTest, PointToForeignDynamicArray) {
+  int* input_translation = new int[2];
+  input_translation[0] = 0;
+  input_translation[1] = 1;
+  Packet packet = PointToForeign(&input_translation);
+  const auto& content = packet.Get<int*>();
+  // The packet content should point to the array.
+  EXPECT_EQ(content[0], 0);
+  EXPECT_EQ(content[1], 1);
+  packet = {};
+  // The vector values should be unaffected.
+  EXPECT_EQ(input_translation[0], 0);
+  EXPECT_EQ(input_translation[1], 1);
+  delete[] input_translation;
+}
+
+TEST(PacketTest, PointToForeignStaticArray) {
+  const int input_translation[] = {0, 1};
+  auto packet = PointToForeign(&input_translation);
+  const auto& content = packet.Get<int[2]>();
+  // The packet content should point to the array.
+  EXPECT_THAT(content, testing::ElementsAre(0, 1));
+  packet = {};
+  // The vector values should be unaffected.
+  EXPECT_THAT(input_translation, testing::ElementsAre(0, 1));
+}
+
 TEST(PacketTest, TestPacketMoveConstructor) {
   std::vector<Packet>* packet_vector_ptr = new std::vector<Packet>();
   packet_vector_ptr->push_back(MakePacket<float>(42));