From 078f2ab6dad886dd5730ba7f63d0ec75e1c60b45 Mon Sep 17 00:00:00 2001
From: Eugene Brevdo <ebrevdo@google.com>
Date: Thu, 9 May 2019 08:12:53 -0700
Subject: [PATCH] [TF] Check that TF inlined Variants properly call value
 destructor on move.

This new test mimics the behavior of having a class with a custom destructor
and a std::shared_ptr; if the destructor is not called a memory leak occurs.

Here we test the memory leak by doing simple counting with a static class member
variable; we manually override the move-constructor and move-assignment to
mimic the their default behavior (but with liveness counting).

PiperOrigin-RevId: 247427651
---
 tensorflow/core/framework/variant_test.cc | 73 +++++++++++++++++++++++
 1 file changed, 73 insertions(+)

diff --git a/tensorflow/core/framework/variant_test.cc b/tensorflow/core/framework/variant_test.cc
index 096143c6eb6..f12b0ea1e61 100644
--- a/tensorflow/core/framework/variant_test.cc
+++ b/tensorflow/core/framework/variant_test.cc
@@ -45,6 +45,61 @@ using Int = Wrapper<int, BIG>;
 template <bool BIG>
 using Float = Wrapper<float, BIG>;
 
+template <bool BIG>
+class MaybeAlive {
+ public:
+  MaybeAlive() : alive_(false) {}
+
+  explicit MaybeAlive(bool alive) : alive_(alive) {
+    if (alive) ++live_counter_;
+  }
+
+  ~MaybeAlive() {
+    if (alive_) --live_counter_;
+  }
+
+  MaybeAlive(const MaybeAlive& rhs) : alive_(rhs.alive_) {
+    if (alive_) ++live_counter_;
+  }
+
+  MaybeAlive& operator=(const MaybeAlive& rhs) {
+    if (this == &rhs) return *this;
+    if (alive_) --live_counter_;
+    alive_ = rhs.alive_;
+    if (alive_) ++live_counter_;
+    return *this;
+  }
+
+  MaybeAlive(MaybeAlive&& rhs) : alive_(false) {
+    alive_ = std::move(rhs.alive_);
+    if (alive_) ++live_counter_;
+  }
+
+  MaybeAlive& operator=(MaybeAlive&& rhs) {
+    if (this == &rhs) return *this;
+    if (alive_) --live_counter_;
+    alive_ = std::move(rhs.alive_);
+    if (alive_) ++live_counter_;
+    return *this;
+  }
+
+  static int LiveCounter() { return live_counter_; }
+
+  string TypeName() const { return "MaybeAlive"; }
+  void Encode(VariantTensorData* data) const {}
+  bool Decode(VariantTensorData data) { return false; }
+
+ private:
+  bool alive_;
+  char big_[BIG ? 256 : 0];
+  static int live_counter_;
+};
+
+template <>
+int MaybeAlive<false>::live_counter_ = 0;
+template <>
+int MaybeAlive<true>::live_counter_ = 0;
+
 template <bool BIG>
 class DeleteCounter {
  public:
@@ -159,6 +214,24 @@ TEST(VariantTest, MoveAndCopyBetweenBigAndSmallVariants) {
   EXPECT_EQ(deleted_small, 1);
 }
 
+template <bool BIG>
+void TestDestructOnVariantMove() {
+  CHECK_EQ(MaybeAlive<BIG>::LiveCounter(), 0);
+  {
+    Variant a = MaybeAlive<BIG>(true);
+    Variant b = std::move(a);
+  }
+  EXPECT_EQ(MaybeAlive<BIG>::LiveCounter(), 0);
+}
+
+TEST(VariantTest, RHSDestructOnVariantMoveBig) {
+  TestDestructOnVariantMove</*BIG=*/true>();
+}
+
+TEST(VariantTest, RHSDestructOnVariantMoveSmall) {
+  TestDestructOnVariantMove</*BIG=*/false>();
+}
+
 TEST(VariantTest, Int) {
   Variant x;
   EXPECT_EQ(x.get<void>(), nullptr);