8 files changed, 88 insertions, 16 deletions

diff --git a/absl/strings/internal/cordz_handle.cc b/absl/strings/internal/cordz_handle.cc
index 5297ec8e..a73fefed 100644
--- a/absl/strings/internal/cordz_handle.cc
+++ b/absl/strings/internal/cordz_handle.cc
@@ -68,11 +68,16 @@ CordzHandle::~CordzHandle() {
   }
 }
 
+bool CordzHandle::SafeToDelete() const {
+  return is_snapshot_ || queue_->IsEmpty();
+}
+
 void CordzHandle::Delete(CordzHandle* handle) {
+  assert(handle);
   if (handle) {
     handle->ODRCheck();
     Queue* const queue = handle->queue_;
-    if (!handle->is_snapshot_ && !queue->IsEmpty()) {
+    if (!handle->SafeToDelete()) {
       SpinLockHolder lock(&queue->mutex);
       CordzHandle* dq_tail = queue->dq_tail.load(std::memory_order_acquire);
       if (dq_tail != nullptr) {
diff --git a/absl/strings/internal/cordz_handle.h b/absl/strings/internal/cordz_handle.h
index 93076cfd..5df53c78 100644
--- a/absl/strings/internal/cordz_handle.h
+++ b/absl/strings/internal/cordz_handle.h
@@ -40,9 +40,20 @@ class CordzHandle {
 
   bool is_snapshot() const { return is_snapshot_; }
 
+  // Returns true if this instance is safe to be deleted because it is either a
+  // snapshot, which is always safe to delete, or not included in the global
+  // delete queue and thus not included in any snapshot.
+  // Callers are responsible for making sure this instance can not be newly
+  // discovered by other threads. For example, CordzInfo instances first de-list
+  // themselves from the global CordzInfo list before determining if they are
+  // safe to be deleted directly.
+  // If SafeToDelete returns false, callers MUST use the Delete() method to
+  // safely queue CordzHandle instances for deletion.
+  bool SafeToDelete() const;
+
   // Deletes the provided instance, or puts it on the delete queue to be deleted
   // once there are no more sample tokens (snapshot) instances potentially
-  // referencing the instance. `handle` may be null.
+  // referencing the instance. `handle` should not be null.
   static void Delete(CordzHandle* handle);
 
   // Returns the current entries in the delete queue in LIFO order.
diff --git a/absl/strings/internal/cordz_handle_test.cc b/absl/strings/internal/cordz_handle_test.cc
index c04240e8..4eefd72c 100644
--- a/absl/strings/internal/cordz_handle_test.cc
+++ b/absl/strings/internal/cordz_handle_test.cc
@@ -52,6 +52,7 @@ TEST(CordzHandleTest, CordzHandleCreateDelete) {
   bool deleted = false;
   auto* handle = new CordzHandleDeleteTracker(&deleted);
   EXPECT_FALSE(handle->is_snapshot());
+  EXPECT_TRUE(handle->SafeToDelete());
   EXPECT_THAT(DeleteQueue(), SizeIs(0));
 
   CordzHandle::Delete(handle);
@@ -62,6 +63,7 @@ TEST(CordzHandleTest, CordzHandleCreateDelete) {
 TEST(CordzHandleTest, CordzSnapshotCreateDelete) {
   auto* snapshot = new CordzSnapshot();
   EXPECT_TRUE(snapshot->is_snapshot());
+  EXPECT_TRUE(snapshot->SafeToDelete());
   EXPECT_THAT(DeleteQueue(), ElementsAre(snapshot));
   delete snapshot;
   EXPECT_THAT(DeleteQueue(), SizeIs(0));
@@ -71,10 +73,12 @@ TEST(CordzHandleTest, CordzHandleCreateDeleteWithSnapshot) {
   bool deleted = false;
   auto* snapshot = new CordzSnapshot();
   auto* handle = new CordzHandleDeleteTracker(&deleted);
+  EXPECT_FALSE(handle->SafeToDelete());
 
   CordzHandle::Delete(handle);
   EXPECT_THAT(DeleteQueue(), ElementsAre(handle, snapshot));
   EXPECT_FALSE(deleted);
+  EXPECT_FALSE(handle->SafeToDelete());
 
   delete snapshot;
   EXPECT_THAT(DeleteQueue(), SizeIs(0));
@@ -219,8 +223,8 @@ TEST(CordzHandleTest, MultiThreaded) {
           if (safe_to_inspect.size() > max_safe_to_inspect) {
             max_safe_to_inspect = safe_to_inspect.size();
           }
+          CordzHandle::Delete(old_handle);
         }
-        CordzHandle::Delete(old_handle);
       }
 
       // Confirm that the test did *something*. This check will be satisfied as
@@ -234,9 +238,10 @@ TEST(CordzHandleTest, MultiThreaded) {
       // Have each thread attempt to clean up everything. Some thread will be
       // the last to reach this cleanup code, and it will be guaranteed to clean
       // up everything because nothing remains to create new handles.
-      for (size_t i = 0; i < handles.size(); i++) {
-        CordzHandle* handle = handles[i].exchange(nullptr);
-        CordzHandle::Delete(handle);
+      for (auto& h : handles) {
+        if (CordzHandle* handle = h.exchange(nullptr)) {
+          CordzHandle::Delete(handle);
+        }
       }
   });
   }
diff --git a/absl/strings/internal/cordz_info.cc b/absl/strings/internal/cordz_info.cc
index 0461ec46..0f657aaf 100644
--- a/absl/strings/internal/cordz_info.cc
+++ b/absl/strings/internal/cordz_info.cc
@@ -126,13 +126,6 @@ void CordzInfo::Track() {
 }
 
 void CordzInfo::Untrack() {
-  {
-    // TODO(b/117940323): change this to assuming ownership instead once all
-    // Cord logic is properly keeping `rep_` in sync with the Cord root rep.
-    absl::MutexLock lock(&mutex_);
-    rep_ = nullptr;
-  }
-
   ODRCheck();
   {
     SpinLockHolder l(&list_->mutex);
@@ -154,6 +147,20 @@ void CordzInfo::Untrack() {
       list_->head.store(next, std::memory_order_release);
     }
   }
+
+  // We can no longer be discovered: perform a fast path check if we are not
+  // listed on any delete queue, so we can directly delete this instance.
+  if (SafeToDelete()) {
+    UnsafeSetCordRep(nullptr);
+    delete this;
+    return;
+  }
+
+  // We are likely part of a snapshot, extend the life of the CordRep
+  {
+    absl::MutexLock lock(&mutex_);
+    if (rep_) CordRep::Ref(rep_);
+  }
   CordzHandle::Delete(this);
 }
 
diff --git a/absl/strings/internal/cordz_info.h b/absl/strings/internal/cordz_info.h
index f7682cb3..1fe046ec 100644
--- a/absl/strings/internal/cordz_info.h
+++ b/absl/strings/internal/cordz_info.h
@@ -110,7 +110,7 @@ class ABSL_LOCKABLE CordzInfo : public CordzHandle {
   // Asserts that this CordzInfo instance is locked.
   void AssertHeld() ABSL_ASSERT_EXCLUSIVE_LOCK(mutex_);
 
-  // Updates the `rep' property of this instance. This methods is invoked by
+  // Updates the `rep` property of this instance. This methods is invoked by
   // Cord logic each time the root node of a sampled Cord changes, and before
   // the old root reference count is deleted. This guarantees that collection
   // code can always safely take a reference on the tracked cord.
@@ -119,6 +119,11 @@ class ABSL_LOCKABLE CordzInfo : public CordzHandle {
   // Cord code is in a state where this can be proven true by the compiler.
   void SetCordRep(CordRep* rep);
 
+  // Returns the current `rep` property of this instance with a reference
+  // added, or null if this instance represents a cord that has since been
+  // deleted or untracked.
+  CordRep* RefCordRep() const ABSL_LOCKS_EXCLUDED(mutex_);
+
   // Returns the current value of `rep_` for testing purposes only.
   CordRep* GetCordRepForTesting() const ABSL_NO_THREAD_SAFETY_ANALYSIS {
     return rep_;
@@ -148,6 +153,9 @@ class ABSL_LOCKABLE CordzInfo : public CordzHandle {
   }
 
  private:
+  using SpinLock = absl::base_internal::SpinLock;
+  using SpinLockHolder = ::absl::base_internal::SpinLockHolder;
+
   // Global cordz info list. CordzInfo stores a pointer to the global list
   // instance to harden against ODR violations.
   struct List {
@@ -155,7 +163,7 @@ class ABSL_LOCKABLE CordzInfo : public CordzHandle {
         : mutex(absl::kConstInit,
                 absl::base_internal::SCHEDULE_COOPERATIVE_AND_KERNEL) {}
 
-    absl::base_internal::SpinLock mutex;
+    SpinLock mutex;
     std::atomic<CordzInfo*> head ABSL_GUARDED_BY(mutex){nullptr};
   };
 
@@ -165,6 +173,9 @@ class ABSL_LOCKABLE CordzInfo : public CordzHandle {
                      MethodIdentifier method);
   ~CordzInfo() override;
 
+  // Sets `rep_` without holding a lock.
+  void UnsafeSetCordRep(CordRep* rep) ABSL_NO_THREAD_SAFETY_ANALYSIS;
+
   void Track();
 
   // Returns the parent method from `src`, which is either `parent_method_` or
@@ -244,6 +255,13 @@ inline void CordzInfo::SetCordRep(CordRep* rep) {
   }
 }
 
+inline void CordzInfo::UnsafeSetCordRep(CordRep* rep) { rep_ = rep; }
+
+inline CordRep* CordzInfo::RefCordRep() const ABSL_LOCKS_EXCLUDED(mutex_) {
+  MutexLock lock(&mutex_);
+  return rep_ ? CordRep::Ref(rep_) : nullptr;
+}
+
 }  // namespace cord_internal
 ABSL_NAMESPACE_END
 }  // namespace absl
diff --git a/absl/strings/internal/cordz_info_test.cc b/absl/strings/internal/cordz_info_test.cc
index 5eaf4b9c..5bb23e47 100644
--- a/absl/strings/internal/cordz_info_test.cc
+++ b/absl/strings/internal/cordz_info_test.cc
@@ -38,6 +38,7 @@ using ::testing::ElementsAre;
 using ::testing::Eq;
 using ::testing::HasSubstr;
 using ::testing::Ne;
+using ::testing::SizeIs;
 
 // Used test values
 auto constexpr kUnknownMethod = CordzUpdateTracker::kUnknown;
@@ -78,10 +79,19 @@ TEST(CordzInfoTest, UntrackCord) {
   CordzInfo::TrackCord(data.data, kTrackCordMethod);
   CordzInfo* info = data.data.cordz_info();
 
+  info->Untrack();
+  EXPECT_THAT(DeleteQueue(), SizeIs(0));
+}
+
+TEST(CordzInfoTest, UntrackCordWithSnapshot) {
+  TestCordData data;
+  CordzInfo::TrackCord(data.data, kTrackCordMethod);
+  CordzInfo* info = data.data.cordz_info();
+
   CordzSnapshot snapshot;
   info->Untrack();
   EXPECT_THAT(CordzInfo::Head(CordzSnapshot()), Eq(nullptr));
-  EXPECT_THAT(info->GetCordRepForTesting(), Eq(nullptr));
+  EXPECT_THAT(info->GetCordRepForTesting(), Eq(data.rep.rep));
   EXPECT_THAT(DeleteQueue(), ElementsAre(info, &snapshot));
 }
 
@@ -113,6 +123,18 @@ TEST(CordzInfoTest, SetCordRepNullUntracksCordOnUnlock) {
   EXPECT_THAT(CordzInfo::Head(CordzSnapshot()), Eq(nullptr));
 }
 
+TEST(CordzInfoTest, RefCordRep) {
+  TestCordData data;
+  CordzInfo::TrackCord(data.data, kTrackCordMethod);
+  CordzInfo* info = data.data.cordz_info();
+
+  size_t refcount = data.rep.rep->refcount.Get();
+  EXPECT_THAT(info->RefCordRep(), Eq(data.rep.rep));
+  EXPECT_THAT(data.rep.rep->refcount.Get(), Eq(refcount + 1));
+  CordRep::Unref(data.rep.rep);
+  info->Untrack();
+}
+
 #if GTEST_HAS_DEATH_TEST
 
 TEST(CordzInfoTest, SetCordRepRequiresMutex) {
diff --git a/absl/strings/internal/cordz_update_tracker.h b/absl/strings/internal/cordz_update_tracker.h
index 741950b2..d5b14067 100644
--- a/absl/strings/internal/cordz_update_tracker.h
+++ b/absl/strings/internal/cordz_update_tracker.h
@@ -47,8 +47,10 @@ class CordzUpdateTracker {
     kClear,
     kConstructorCord,
     kConstructorString,
+    kCordReader,
     kFlatten,
     kGetAppendRegion,
+    kMakeCordFromExternal,
     kMoveAppendCord,
     kMoveAssignCord,
     kMovePrependCord,
diff --git a/absl/strings/internal/cordz_update_tracker_test.cc b/absl/strings/internal/cordz_update_tracker_test.cc
index eda662f2..fcd17df7 100644
--- a/absl/strings/internal/cordz_update_tracker_test.cc
+++ b/absl/strings/internal/cordz_update_tracker_test.cc
@@ -45,8 +45,10 @@ Methods AllMethods() {
                  Method::kClear,
                  Method::kConstructorCord,
                  Method::kConstructorString,
+                 Method::kCordReader,
                  Method::kFlatten,
                  Method::kGetAppendRegion,
+                 Method::kMakeCordFromExternal,
                  Method::kMoveAppendCord,
                  Method::kMoveAssignCord,
                  Method::kMovePrependCord,