From f1d5e900e6d5da39452c387a0a00ef73af894aff Mon Sep 17 00:00:00 2001
From: Simon Farre <simon.farre.cx@gmail.com>
Date: Mon, 15 Jan 2024 10:57:08 +0100
Subject: [PATCH] Fixes #3678

Refactor so that marks_with_checkpoints is just changed in one place, not arbitrarily access it. Ref counts had the same changes in a previous commit.

Fixes a bug for loaded persistent checkpoints where the re-created checkpoints did not get their reference counting correct.

This closes #3678
---
 src/GdbServer.cc      |  4 ++++
 src/ReplayTimeline.cc | 29 ++++++++++++++---------------
 src/ReplayTimeline.h  |  3 +++
 3 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/src/GdbServer.cc b/src/GdbServer.cc
index ff8d329132c..ce097e4cf20 100644
--- a/src/GdbServer.cc
+++ b/src/GdbServer.cc
@@ -2250,6 +2250,10 @@ Checkpoint::Checkpoint(ReplayTimeline& timeline, TaskUid last_continue_tuid,
     mark = timeline.add_explicit_checkpoint();
   } else {
     mark = timeline.mark();
+    const auto prior = timeline.find_closest_mark_with_clone(mark);
+    if(prior) {
+      prior->get_internal()->inc_refcount();
+    }
   }
 }
 
diff --git a/src/ReplayTimeline.cc b/src/ReplayTimeline.cc
index bdbce3d44e3..3e3a2c1822c 100644
--- a/src/ReplayTimeline.cc
+++ b/src/ReplayTimeline.cc
@@ -225,7 +225,7 @@ ReplayTimeline::Mark ReplayTimeline::mark() {
 ReplayTimeline::Mark ReplayTimeline::recreate_mark_from_data(const MarkData& mark_data, ReplaySession::shr_ptr session) {
   Mark result;
   auto m = make_shared<InternalMark>(this, mark_data, session);
-
+  m->inc_refcount();
   auto& mark_vector = marks[m->proto.key];
   mark_vector.push_back(m);
   result.ptr = mark_vector.back();
@@ -235,11 +235,7 @@ ReplayTimeline::Mark ReplayTimeline::recreate_mark_from_data(const MarkData& mar
 void ReplayTimeline::register_mark_as_checkpoint(Mark& m) {
   DEBUG_ASSERT(m.ptr && m.ptr->checkpoint && "Can't register mark as checkpoint if no checkpoint exists");
   auto key = m.ptr->proto.key;
-  if (marks_with_checkpoints.find(key) == marks_with_checkpoints.end()) {
-    marks_with_checkpoints[key] = 1;
-  } else {
-    marks_with_checkpoints[key]++;
-  }
+  increase_mark_with_checkpoints(key);
   m.ptr->inc_refcount();
 }
 
@@ -313,6 +309,15 @@ ReplayTimeline::Mark ReplayTimeline::lazy_reverse_singlestep(const Mark& from,
   return Mark();
 }
 
+void ReplayTimeline::increase_mark_with_checkpoints(
+    const MarkKey& key) noexcept {
+  if (marks_with_checkpoints.find(key) == marks_with_checkpoints.end()) {
+    marks_with_checkpoints[key] = 1;
+  } else {
+    marks_with_checkpoints[key]++;
+  }
+}
+
 ReplayTimeline::Mark ReplayTimeline::add_explicit_checkpoint() {
   DEBUG_ASSERT(current->can_clone());
 
@@ -321,11 +326,7 @@ ReplayTimeline::Mark ReplayTimeline::add_explicit_checkpoint() {
     unapply_breakpoints_and_watchpoints();
     m.ptr->checkpoint = current->clone();
     auto key = m.ptr->proto.key;
-    if (marks_with_checkpoints.find(key) == marks_with_checkpoints.end()) {
-      marks_with_checkpoints[key] = 1;
-    } else {
-      marks_with_checkpoints[key]++;
-    }
+    increase_mark_with_checkpoints(key);
   }
   m.ptr->inc_refcount();
   return m;
@@ -798,11 +799,9 @@ ReplayTimeline::Mark ReplayTimeline::recreate_marks_for_non_explicit(const Check
   // first add the mark with an actual clone, this is not a GDB checkpoint, but an RR checkpoint
   auto mark = recreate_mark_from_data(cp.clone_data, clone);
   reverse_exec_checkpoints[mark] = estimate_progress_of(*clone);
-  mark.get_internal()->inc_refcount();
   register_mark_as_checkpoint(mark);
   // then add the mark with no clone, the one that will be visible to GDB, i.e. non explicit checkpoint
   Mark result = recreate_mark_from_data(*cp.non_explicit_mark_data, nullptr);
-  // auto internal = make_shared<InternalMark>(this, *cp.non_explicit_data, nullptr);
   return result;
 }
 
@@ -1736,8 +1735,8 @@ std::shared_ptr<ReplayTimeline::Mark> ReplayTimeline::find_closest_mark_with_clo
   for(auto it = std::rbegin(marks_found->second); it != std::rend(marks_found->second); it++) {
     DEBUG_ASSERT(it->get() != nullptr);
     if((*it)->checkpoint) {
-      std::shared_ptr<Mark> result = std::make_shared<Mark>();
-      result->ptr = *it;
+      auto result = std::make_shared<Mark>();
+      result->ptr = *(it);
       return result;
     }
   }
diff --git a/src/ReplayTimeline.h b/src/ReplayTimeline.h
index a4b586a6125..778a90dde9a 100644
--- a/src/ReplayTimeline.h
+++ b/src/ReplayTimeline.h
@@ -147,6 +147,9 @@ class ReplayTimeline {
    */
   void remove_explicit_checkpoint(const Mark& mark);
 
+  /** Find mark that has `key` and increase the checkpoint count for that mark. */
+  void increase_mark_with_checkpoints(const MarkKey& key) noexcept;
+
   /**
    * Return true if we're currently at the given mark.
    */