From 549301d0f01976eb321055dc2a168368089050ef Mon Sep 17 00:00:00 2001
From: rdb <git@rdb.name>
Date: Sat, 5 May 2018 23:10:41 +0200
Subject: [PATCH] putil: keep reference to objects queued for writing

This prevents a crash when a reference counted object is destroyed right after a write_pointer call.

Fixes #310
---
 panda/src/putil/bamWriter.cxx | 30 ++++++++++++++++++++++++++++--
 panda/src/putil/bamWriter.h   |  3 ++-
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/panda/src/putil/bamWriter.cxx b/panda/src/putil/bamWriter.cxx
index 679600b664..0fcfd2fe1f 100644
--- a/panda/src/putil/bamWriter.cxx
+++ b/panda/src/putil/bamWriter.cxx
@@ -94,6 +94,10 @@ BamWriter::
   for (si = _state_map.begin(); si != _state_map.end(); ++si) {
     TypedWritable *object = (TypedWritable *)(*si).first;
     object->remove_bam_writer(this);
+
+    if ((*si).second._refcount != nullptr) {
+      unref_delete((*si).second._refcount);
+    }
   }
 }
 
@@ -529,6 +533,9 @@ object_destructs(TypedWritable *object) {
     // we're in trouble when we do write it out.
     nassertv(!(*si).second._written_seq.is_initial());
 
+    // This cannot be called if we are still holding a reference to it.
+    nassertv((*si).second._refcount == nullptr);
+
     int object_id = (*si).second._object_id;
     _freed_object_ids.push_back(object_id);
 
@@ -606,8 +613,10 @@ enqueue_object(const TypedWritable *object) {
     // No, it hasn't, so assign it the next number in sequence arbitrarily.
     object_id = _next_object_id;
 
-    bool inserted =
-      _state_map.insert(StateMap::value_type(object, StoreState(_next_object_id))).second;
+    StateMap::iterator si;
+    bool inserted;
+    tie(si, inserted) =
+      _state_map.insert(StateMap::value_type(object, StoreState(_next_object_id)));
     nassertr(inserted, false);
 
     // Store ourselves on the TypedWritable so that we get notified when it
@@ -615,6 +624,14 @@ enqueue_object(const TypedWritable *object) {
     (const_cast<TypedWritable*>(object))->add_bam_writer(this);
     _next_object_id++;
 
+    // Increase the reference count if this inherits from ReferenceCount,
+    // until we get a chance to write this object for the first time.
+    const ReferenceCount *rc = ((TypedWritable *)object)->as_reference_count();
+    if (rc != nullptr) {
+      rc->ref();
+      (*si).second._refcount = rc;
+    }
+
   } else {
     // Yes, it has; get the object ID.
     object_id = (*si).second._object_id;
@@ -703,6 +720,15 @@ flush_queue() {
       (*si).second._written_seq = _writing_seq;
       (*si).second._modified = object->get_bam_modified();
 
+      // Now release any reference we hold to it, so that it may destruct.
+      const ReferenceCount *rc = (*si).second._refcount;
+      if (rc != nullptr) {
+        // We need to assign this pointer to null before deleting the object,
+        // since that may end up calling object_destructs.
+        (*si).second._refcount = nullptr;
+        unref_delete(rc);
+      }
+
     } else {
       // On subsequent times when we write a particular object, we write
       // simply TypeHandle::none(), followed by the object ID. The occurrence
diff --git a/panda/src/putil/bamWriter.h b/panda/src/putil/bamWriter.h
index c2b04b3757..bcd4a76ce6 100644
--- a/panda/src/putil/bamWriter.h
+++ b/panda/src/putil/bamWriter.h
@@ -140,8 +140,9 @@ private:
     int _object_id;
     UpdateSeq _written_seq;
     UpdateSeq _modified;
+    const ReferenceCount *_refcount;
 
-    StoreState(int object_id) : _object_id(object_id) {}
+    StoreState(int object_id) : _object_id(object_id), _refcount(nullptr) {}
   };
   typedef phash_map<const TypedWritable *, StoreState, pointer_hash> StateMap;
   StateMap _state_map;