From 7cf53f3c3fb4999d12807b75749826e76b17c307 Mon Sep 17 00:00:00 2001
From: Alex Owens <alexowens90@gmail.com>
Date: Thu, 31 Oct 2024 15:30:23 +0000
Subject: [PATCH] Fix C++ tests compilation

---
 cpp/arcticdb/processing/component_manager.cpp | 21 +++++++++++++------
 cpp/arcticdb/processing/component_manager.hpp | 15 ++++++-------
 2 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/cpp/arcticdb/processing/component_manager.cpp b/cpp/arcticdb/processing/component_manager.cpp
index 445fd0c9f6..36544db71b 100644
--- a/cpp/arcticdb/processing/component_manager.cpp
+++ b/cpp/arcticdb/processing/component_manager.cpp
@@ -19,12 +19,21 @@ std::vector<EntityId> ComponentManager::get_new_entity_ids(size_t count) {
     return ids;
 }
 
-void ComponentManager::erase_entity(EntityId id) {
-    // Ideally would call registry_.destroy(id), or at least registry_.erase<std::shared_ptr<SegmentInMemory>>(id)
-    // at this point. However, they are both slower than this, so just decrement the ref count of the only
-    // sizeable component, so that when the shared pointer goes out of scope in the calling function, the
-    // memory is freed
-    registry_.get<std::shared_ptr<SegmentInMemory>>(id).reset();
+void ComponentManager::decrement_entity_fetch_count(EntityId id) {
+    if (registry_.get<std::atomic<EntityFetchCount>>(id).fetch_sub(1) == 0) {
+        // This entity will never be accessed again
+        // Ideally would call registry_.destroy(id), or at least registry_.erase<std::shared_ptr<SegmentInMemory>>(id)
+        // at this point. However, they are both slower than this, and would require taking a unique_lock on the
+        // shared_mutex, so just decrement the ref count of the only sizeable component, so that when the shared pointer
+        // goes out of scope in the calling function, the memory is freed
+        registry_.get<std::shared_ptr<SegmentInMemory>>(id).reset();
+        debug::check<ErrorCode::E_ASSERTION_FAILURE>(!registry_.get<std::shared_ptr<SegmentInMemory>>(id),
+                                                     "SegmentInMemory memory retained in ComponentManager");
+    }
+}
+
+void ComponentManager::update_entity_fetch_count(EntityId id, EntityFetchCount count) {
+    registry_.get<std::atomic<EntityFetchCount>>(id).store(count);
 }
 
 
diff --git a/cpp/arcticdb/processing/component_manager.hpp b/cpp/arcticdb/processing/component_manager.hpp
index 44273e53a1..500906ced2 100644
--- a/cpp/arcticdb/processing/component_manager.hpp
+++ b/cpp/arcticdb/processing/component_manager.hpp
@@ -8,6 +8,7 @@
 #pragma once
 
 #include <atomic>
+#include <shared_mutex>
 
 #include <entt/entity/registry.hpp>
 
@@ -84,7 +85,7 @@ class ComponentManager {
         for (auto id: ids) {
             registry_.replace<T>(id, value);
             if constexpr (std::is_same_v<T, EntityFetchCount>) {
-                registry_.get<std::atomic<EntityFetchCount>>(id).store(value);
+                update_entity_fetch_count(id, value);
             }
         }
     }
@@ -96,7 +97,7 @@ class ComponentManager {
         for (auto [idx, id]: folly::enumerate(ids)) {
             registry_.replace<T>(id, values[idx]);
             if constexpr (std::is_same_v<T, EntityFetchCount>) {
-                registry_.get<std::atomic<EntityFetchCount>>(id).store(values[idx]);
+                update_entity_fetch_count(id, values[idx]);
             }
         }
     }
@@ -115,12 +116,7 @@ class ComponentManager {
             }
             if (decrement_fetch_count) {
                 for (auto id: ids) {
-                    // This entity will never be accessed again
-                    if (registry_.get<std::atomic<EntityFetchCount>>(id).fetch_sub(1) == 0) {
-                        erase_entity(id);
-                        debug::check<ErrorCode::E_ASSERTION_FAILURE>(!registry_.get<std::shared_ptr<SegmentInMemory>>(id),
-                                                                     "SegmentInMemory memory retained in ComponentManager");
-                    }
+                    decrement_entity_fetch_count(id);
                 }
             }
         }
@@ -138,7 +134,8 @@ class ComponentManager {
     }
 
 private:
-    void erase_entity(EntityId id);
+    void decrement_entity_fetch_count(EntityId id);
+    void update_entity_fetch_count(EntityId id, EntityFetchCount count);
 
     entt::registry registry_;
     std::shared_mutex mtx_;