@colyseus/schema: improve instance sharing and ref count tracking.

Author: endel

colyseus/serializer/schema/decoder.lua

@@ -203,7 +203,7 @@ function Decoder:new(state, context)
 
   -- set root state refid
   state.__refid = 1
-  instance.refs:set(state.__refid, state)
+  instance.refs:add(state.__refid, state, true)
 
 	return instance
 end
@@ -279,15 +279,15 @@ function Decoder:decode_value(decoder, operation, ref, field_index, field_type,
   -- DELETE operations
   --
   if bit.band(operation, OPERATION.DELETE) == OPERATION.DELETE then
-    if operation ~= OPERATION.DELETE_AND_ADD then
-      ref:delete_by_index(field_index)
-    end
-
     -- Flag `ref_id` for garbage collection.
     if type(previous_value) == "table" and previous_value.__refid ~= nil then
       decoder.refs:remove(previous_value.__refid)
     end
 
+    if operation ~= OPERATION.DELETE_AND_ADD then
+      ref:delete_by_index(field_index)
+    end
+
     value = nil
   end
 
@@ -302,26 +302,14 @@ function Decoder:decode_value(decoder, operation, ref, field_index, field_type,
     local __refid = decode.number(bytes, it) + 1
     value = decoder.refs:get(__refid)
 
-    if previous_value ~= nil then
-      local previous_refid = previous_value.__refid
-      if (
-        previous_refid > 0 and
-        previous_refid ~= __refid and
-        -- FIXME: we may need to check for REPLACE operation as well
-        bit.band(operation, OPERATION.DELETE) == OPERATION.DELETE
-      ) then
-        decoder.refs:remove(previous_refid)
-      end
-    end
-
     if bit.band(operation, OPERATION.ADD) == OPERATION.ADD then
       local concrete_child_type = decoder:get_schema_type(bytes, it, field_type);
       if value == nil then
         value = concrete_child_type:new()
         value.__refid = __refid
       end
 
-      decoder.refs:set(__refid, value, (
+      decoder.refs:add(__refid, value, (
         value ~= previous_value or
         (operation == OPERATION.DELETE_AND_ADD and value == previous_value)
       ));
@@ -348,12 +336,15 @@ function Decoder:decode_value(decoder, operation, ref, field_index, field_type,
 
     if previous_value ~= nil then
       if previous_value.__refid ~= nil and previous_value.__refid ~= __refid then
-        decoder.refs:remove(previous_value.__refid)
 
         --
         -- Trigger on_remove() if structure has been replaced.
         --
         previous_value:each(function(val, key)
+          if type(val) == "table" and val.__refid ~= nil then
+            decoder.refs:remove(val.__refid)
+          end
+
           table.insert(all_changes, {
             __refid = previous_value.__refid,
             op = OPERATION.DELETE,
@@ -365,7 +356,10 @@ function Decoder:decode_value(decoder, operation, ref, field_index, field_type,
       end
     end
 
-    decoder.refs:set(__refid, value, (value_ref ~= previous_value))
+    decoder.refs:add(__refid, value, (
+      value_ref ~= previous_value or
+      (operation == OPERATION.DELETE_AND_ADD and value_ref == previous_value)
+    ))
   end
 
   return value, previous_value

colyseus/serializer/schema/reference_tracker.lua

@@ -26,10 +26,10 @@ function reference_tracker:get(ref_id)
   return self.refs[ref_id]
 end
 
-function reference_tracker:set(ref_id, ref, increment_count)
+function reference_tracker:add(ref_id, ref, increment_count)
   self.refs[ref_id] = ref
 
-  if increment_count == nil or increment_count == true then
+  if increment_count == true then
     self.ref_counts[ref_id] = (self.ref_counts[ref_id] or 0) + 1
   end
 
@@ -39,12 +39,27 @@ function reference_tracker:set(ref_id, ref, increment_count)
 end
 
 function reference_tracker:remove(ref_id)
-  self.ref_counts[ref_id] = self.ref_counts[ref_id] - 1;
+  local ref_count = self.ref_counts[ref_id]
 
-  local added_to_deleted_refs = (self.deleted_refs[ref_id] == nil)
-  self.deleted_refs[ref_id] = true
+  -- skip if ref_id is not being tracked
+  if ref_count == nil then
+    print("trying to remove ref_id that doesn't exist: " .. tostring(ref_id))
+    return false
+  end
+
+  if ref_count == 0 then
+    print("trying to remove ref_id with 0 ref_count: " .. tostring(ref_id))
+    return false
+  end
+
+  self.ref_counts[ref_id] = ref_count - 1;
+
+  if self.ref_counts[ref_id] <= 0 then
+    self.deleted_refs[ref_id] = true
+    return true
+  end
 
-  return added_to_deleted_refs
+  return false
 end
 
 function reference_tracker:count()
@@ -72,20 +87,20 @@ function reference_tracker:garbage_collection()
       --
       if ref._schema ~= nil then
         for field, field_type in pairs(ref._schema) do
-          if (
-            type(field_type) ~= "string" and
-            ref[field] ~= nil and
-            ref[field].__refid ~= nil and
-            self:remove(ref[field].__refid)
-          ) then
-            table.insert(deleted_refs, ref[field].__refid)
+          local child_ref_id = type(field_type) ~= "string" and ref[field] ~= nil and ref[field].__refid
+          if (child_ref_id ~= nil and self.deleted_refs[child_ref_id] == nil and self:remove(child_ref_id)) then
+            table.insert(deleted_refs, child_ref_id)
           end
         end
 
       elseif ref._child_type['new'] ~= nil then
         ref:each(function(value)
-          if self:remove(value.__refid) then
-            table.insert(deleted_refs, value.__refid)
+          local child_ref_id = value.__refid
+          if (
+            self.deleted_refs[child_ref_id] == nil and
+            self:remove(child_ref_id)
+          ) then
+            table.insert(deleted_refs, child_ref_id)
           end
         end)
       end

test/schema_serializer.lua

@@ -355,18 +355,18 @@ return function()
       local state = InstanceSharingTypes:new()
       local decoder = Decoder:new(state)
 
-      decoder:decode({ 130, 1, 131, 2, 128, 3, 129, 3, 255, 1, 255, 2, 255, 3, 128, 4, 255, 3, 128, 4, 255, 4, 128, 10, 129, 10, 255, 4, 128, 10, 129, 10 });
+      decoder:decode({ 130, 1, 131, 2, 128, 3, 129, 3, 255, 3, 128, 4, 255, 4, 128, 10, 129, 10 });
       assert_equal(state.player1, state.player2);
       assert_equal(state.player1.position, state.player2.position);
       assert_equal(decoder.refs.ref_counts[state.player1.__refid], 2);
       assert_equal(5, decoder.refs:count());
 
-      decoder:decode({ 130, 1, 131, 2, 64, 65 });
+      decoder:decode({ 64, 65, 255, 0, 64, 65 });
       assert_equal(nil, state.player1);
       assert_equal(nil, state.player2);
       assert_equal(3, decoder.refs:count());
 
-      decoder:decode({ 255, 1, 128, 0, 5, 128, 1, 5, 128, 2, 5, 128, 3, 6, 255, 5, 128, 7, 255, 6, 128, 8, 255, 7, 128, 10, 129, 10, 255, 8, 128, 10, 129, 10 });
+      decoder:decode({ 255, 1, 128, 0, 5, 128, 1, 5, 128, 2, 5, 128, 3, 7, 255, 5, 128, 6, 255, 6, 128, 10, 129, 10, 255, 7, 128, 8, 255, 8, 128, 10, 129, 10 });
       assert_equal(state.arrayOfPlayers[1], state.arrayOfPlayers[2]);
       assert_equal(state.arrayOfPlayers[2], state.arrayOfPlayers[3]);
       assert_not_equal(state.arrayOfPlayers[3], state.arrayOfPlayers[4]);
@@ -378,7 +378,7 @@ return function()
       local previous_array_schema__refid = state.arrayOfPlayers.__refid;
 
       -- Replacing ArraySchema
-      decoder:decode({ 130, 9, 255, 9, 128, 0, 10, 255, 10, 128, 11, 255, 11, 128, 10, 129, 20 });
+      decoder:decode({ 194, 9, 255, 9, 128, 0, 10, 255, 10, 128, 11, 255, 11, 128, 10, 129, 20 });
       assert_equal(false, decoder.refs:has(previous_array_schema__refid));
       assert_equal(1, state.arrayOfPlayers:length());
       assert_equal(5, decoder.refs:count());