Refactoring away the use of AZ::TransformNotificationBus. Added unittests for deactivating parents of inner roots.

Signed-off-by: AMZN-Olex <5432499+AMZN-Olex@users.noreply.github.com>
2021-09-15 11:29:10 -04:00
parent fcd5b3f184
commit 714dc0a0e8
6 changed files with 296 additions and 111 deletions
@@ -19,6 +19,10 @@ namespace Multiplayer
        : public AZ::ComponentBus
    {
    public:
+        //! @returns true if the entity a hierarchical component attached should be considered for inclusion in a hierarchy
+        //!          this should return false when an entity is deactivating
+        virtual bool IsHierarchyEnabled() const = 0;
+
        //! @returns hierarchical entities, the first element is the top level root
        virtual AZStd::vector<AZ::Entity*> GetHierarchicalEntities() const = 0;

@@ -9,6 +9,7 @@
 #pragma once

 #include <AzCore/Component/Component.h>
+#include <AzCore/Component/TransformBus.h>
 #include <Multiplayer/Components/NetworkHierarchyBus.h>
 #include <Source/AutoGen/NetworkHierarchyChildComponent.AutoComponent.h>

@@ -42,12 +43,13 @@ namespace Multiplayer
        //! NetworkHierarchyChildComponentBase overrides.
        //! @{
        void OnInit() override;
-        void OnActivate(Multiplayer::EntityIsMigrating entityIsMigrating) override;
-        void OnDeactivate(Multiplayer::EntityIsMigrating entityIsMigrating) override;
+        void OnActivate(EntityIsMigrating entityIsMigrating) override;
+        void OnDeactivate(EntityIsMigrating entityIsMigrating) override;
        //! @}

        //! NetworkHierarchyRequestBus overrides.
        //! @{
+        bool IsHierarchyEnabled() const override;
        bool IsHierarchicalChild() const override;
        bool IsHierarchicalRoot() const override { return false; }
        AZ::Entity* GetHierarchicalRoot() const override;
@@ -58,15 +60,26 @@ namespace Multiplayer

    protected:
        //! Used by @NetworkHierarchyRootComponent
-        void SetTopLevelHierarchyRootComponent(NetworkHierarchyRootComponent* hierarchyRoot);
+        void SetTopLevelHierarchyRootEntity(AZ::Entity* hierarchyRoot);

    private:
-        const NetworkHierarchyRootComponent* m_hierarchyRootComponent = nullptr;
+        AZ::ChildChangedEvent::Handler m_childChangedHandler;
+        AZ::ParentChangedEvent::Handler m_parentChangedHandler;
+
+        void OnChildChanged(AZ::ChildChangeType type, AZ::EntityId child);
+        void OnParentChanged(AZ::EntityId oldParent, AZ::EntityId parent);
+
+        //! Points to the top level root.
+        AZ::Entity* m_rootEntity = nullptr;

        AZ::Event<NetEntityId>::Handler m_hierarchyRootNetIdChanged;
        void OnHierarchyRootNetIdChanged(NetEntityId rootNetId);

        NetworkHierarchyChangedEvent m_networkHierarchyChangedEvent;
        NetworkHierarchyLeaveEvent m_networkHierarchyLeaveEvent;
+
+        bool m_isDeactivating = false;
+
+        void NotifyChildrenHierarchyDisbanded();
    };
 }
@@ -10,8 +10,8 @@

 #include <AzCore/Component/Component.h>
 #include <AzCore/Component/TransformBus.h>
-#include <Source/AutoGen/NetworkHierarchyRootComponent.AutoComponent.h>
 #include <Multiplayer/Components/NetworkHierarchyBus.h>
+#include <Source/AutoGen/NetworkHierarchyRootComponent.AutoComponent.h>

 namespace Multiplayer
 {
@@ -27,8 +27,8 @@ namespace Multiplayer
    class NetworkHierarchyRootComponent final
        : public NetworkHierarchyRootComponentBase
        , public NetworkHierarchyRequestBus::Handler
-        , protected AZ::TransformNotificationBus::MultiHandler
    {
+        friend class NetworkHierarchyChildComponent;
    public:
        AZ_MULTIPLAYER_COMPONENT(Multiplayer::NetworkHierarchyRootComponent, s_networkHierarchyRootComponentConcreteUuid, Multiplayer::NetworkHierarchyRootComponentBase);

@@ -37,6 +37,8 @@ namespace Multiplayer
        static void GetProvidedServices(AZ::ComponentDescriptor::DependencyArrayType& provided);
        static void GetIncompatibleServices(AZ::ComponentDescriptor::DependencyArrayType& incompatible);

+        NetworkHierarchyRootComponent();
+
        //! NetworkHierarchyRootComponentBase overrides.
        //! @{
        void OnInit() override;
@@ -46,6 +48,7 @@ namespace Multiplayer

        //! NetworkHierarchyRequestBus overrides.
        //! @{
+        bool IsHierarchyEnabled() const override;
        bool IsHierarchicalRoot() const override;
        bool IsHierarchicalChild() const override;
        AZStd::vector<AZ::Entity*> GetHierarchicalEntities() const override;
@@ -55,22 +58,24 @@ namespace Multiplayer
        //! @}

    protected:
-        //! AZ::TransformNotificationBus::Handler overrides.
-        //! @{
-        void OnParentChanged(AZ::EntityId oldParent, AZ::EntityId newParent) override;
-        void OnChildAdded(AZ::EntityId child) override;
-        void OnChildRemoved(AZ::EntityId childRemovedId) override;
-        //! @}
-
        void SetTopLevelHierarchyRootEntity(AZ::Entity* hierarchyRoot);

    private:
+        AZ::ChildChangedEvent::Handler m_childChangedHandler;
+        AZ::ParentChangedEvent::Handler m_parentChangedHandler;
+
+        void OnChildChanged(AZ::ChildChangeType type, AZ::EntityId child);
+        void OnParentChanged(AZ::EntityId oldParent, AZ::EntityId parent);
+
        NetworkHierarchyChangedEvent m_networkHierarchyChangedEvent;
        NetworkHierarchyLeaveEvent m_networkHierarchyLeaveEvent;

-        AZ::Entity* m_higherRootEntity = nullptr;
+        //! Points to the top level root, if this root is an inner root in this hierarchy.
+        AZ::Entity* m_rootEntity = nullptr;
+
        AZStd::vector<AZ::Entity*> m_hierarchicalEntities;

+        //! Rebuilds hierarchy starting from this root component's entity.
        void RebuildHierarchy();
        
        //! @param underEntity Walk the child entities that belong to @underEntity and consider adding them to the hierarchy
@@ -86,5 +91,9 @@ namespace Multiplayer
        //!            @currentEntityCount will be modified to reflect the total entity count upon completion of this method.
        //! @returns false if an attempt was made to go beyond the maximum supported hierarchy size, true otherwise
        bool RecursiveAttachHierarchicalChild(AZ::EntityId entity, uint32_t& currentEntityCount);
+
+        void SetRootForEntity(AZ::Entity* root, const AZ::Entity* childEntity);
+
+        bool m_isDeactivating = false;
    };
 }
@@ -7,6 +7,7 @@
 */

 #include <AzCore/Serialization/EditContext.h>
+#include <AzFramework/Components/TransformComponent.h>
 #include <Multiplayer/IMultiplayer.h>
 #include <Multiplayer/Components/NetBindComponent.h>
 #include <Multiplayer/Components/NetworkHierarchyBus.h>
@@ -53,7 +54,9 @@ namespace Multiplayer
    }

    NetworkHierarchyChildComponent::NetworkHierarchyChildComponent()
-        : m_hierarchyRootNetIdChanged([this](NetEntityId rootNetId) {OnHierarchyRootNetIdChanged(rootNetId); })
+        : m_childChangedHandler([this](AZ::ChildChangeType type, AZ::EntityId child) { OnChildChanged(type, child); })
+        , m_parentChangedHandler([this](AZ::EntityId oldParent, AZ::EntityId parent) { OnParentChanged(oldParent, parent); })
+        , m_hierarchyRootNetIdChanged([this](NetEntityId rootNetId) {OnHierarchyRootNetIdChanged(rootNetId); })
    {

    }
@@ -66,13 +69,36 @@ namespace Multiplayer
    {
        HierarchyRootAddEvent(m_hierarchyRootNetIdChanged);
        NetworkHierarchyRequestBus::Handler::BusConnect(GetEntityId());
+
+        if (AzFramework::TransformComponent* transformComponent = GetEntity()->FindComponent<AzFramework::TransformComponent>())
+        {
+            transformComponent->BindChildChangedEventHandler(m_childChangedHandler);
+            transformComponent->BindParentChangedEventHandler(m_parentChangedHandler);
+        }
    }

    void NetworkHierarchyChildComponent::OnDeactivate([[maybe_unused]] EntityIsMigrating entityIsMigrating)
    {
+        m_isDeactivating = true;
+
+        if (m_rootEntity)
+        {
+            if (NetworkHierarchyRootComponent* root = m_rootEntity->FindComponent<NetworkHierarchyRootComponent>())
+            {
+                root->RebuildHierarchy();
+            }
+        }
+
+        NotifyChildrenHierarchyDisbanded();
+
        NetworkHierarchyRequestBus::Handler::BusDisconnect();
    }

+    bool NetworkHierarchyChildComponent::IsHierarchyEnabled() const
+    {
+        return !m_isDeactivating;
+    }
+
    bool NetworkHierarchyChildComponent::IsHierarchicalChild() const
    {
        if (GetHierarchyRoot() != InvalidNetEntityId)
@@ -85,14 +111,14 @@ namespace Multiplayer

    AZ::Entity* NetworkHierarchyChildComponent::GetHierarchicalRoot() const
    {
-        return m_hierarchyRootComponent ? m_hierarchyRootComponent->GetEntity() : nullptr;
+        return m_rootEntity;
    }

    AZStd::vector<AZ::Entity*> NetworkHierarchyChildComponent::GetHierarchicalEntities() const
    {
-        if (m_hierarchyRootComponent)
+        if (m_rootEntity)
        {
-            return m_hierarchyRootComponent->GetHierarchicalEntities();
+            return m_rootEntity->FindComponent<NetworkHierarchyRootComponent>()->GetHierarchicalEntities();
        }

        return {};
@@ -108,18 +134,18 @@ namespace Multiplayer
        handler.Connect(m_networkHierarchyLeaveEvent);
    }

-    void NetworkHierarchyChildComponent::SetTopLevelHierarchyRootComponent(NetworkHierarchyRootComponent* hierarchyRoot)
+    void NetworkHierarchyChildComponent::SetTopLevelHierarchyRootEntity(AZ::Entity* hierarchyRoot)
    {
-        m_hierarchyRootComponent = hierarchyRoot;
+        m_rootEntity = hierarchyRoot;
        if (HasController() && GetNetBindComponent()->GetNetEntityRole() == NetEntityRole::Authority)
        {
            NetworkHierarchyChildComponentController* controller = static_cast<NetworkHierarchyChildComponentController*>(GetController());
-            if (hierarchyRoot)
+            if (m_rootEntity)
            {
-                const NetEntityId netRootId = GetNetworkEntityManager()->GetNetEntityIdById(hierarchyRoot->GetEntityId());
+                const NetEntityId netRootId = GetNetworkEntityManager()->GetNetEntityIdById(m_rootEntity->GetId());
                controller->SetHierarchyRoot(netRootId);

-                m_networkHierarchyChangedEvent.Signal(hierarchyRoot->GetEntityId());
+                m_networkHierarchyChangedEvent.Signal(m_rootEntity->GetId());
            }
            else
            {
@@ -128,20 +154,74 @@ namespace Multiplayer
                m_networkHierarchyLeaveEvent.Signal();
            }
        }
+
+        if (m_rootEntity == nullptr)
+        {
+            NotifyChildrenHierarchyDisbanded();
+        }
+    }
+
+    void NetworkHierarchyChildComponent::OnChildChanged([[maybe_unused]] AZ::ChildChangeType type, [[maybe_unused]] AZ::EntityId child)
+    {
+        if (m_rootEntity)
+        {
+            if (NetworkHierarchyRootComponent* root = m_rootEntity->FindComponent<NetworkHierarchyRootComponent>())
+            {
+                root->RebuildHierarchy();
+            }
+        }
+    }
+
+    void NetworkHierarchyChildComponent::OnParentChanged([[maybe_unused]] AZ::EntityId oldParent, [[maybe_unused]] AZ::EntityId parent)
+    {
+        if (m_rootEntity)
+        {
+            if (NetworkHierarchyRootComponent* root = m_rootEntity->FindComponent<NetworkHierarchyRootComponent>())
+            {
+                root->RebuildHierarchy();
+            }
+        }
    }

    void NetworkHierarchyChildComponent::OnHierarchyRootNetIdChanged(NetEntityId rootNetId)
    {
-        const ConstNetworkEntityHandle rootHandle = GetNetworkEntityManager()->GetEntity(rootNetId);
+        ConstNetworkEntityHandle rootHandle = GetNetworkEntityManager()->GetEntity(rootNetId);
        if (rootHandle.Exists())
        {
-            m_hierarchyRootComponent = rootHandle.FindComponent<NetworkHierarchyRootComponent>();
-            m_networkHierarchyChangedEvent.Signal(m_hierarchyRootComponent->GetEntityId());
+            AZ::Entity* newRoot = rootHandle.GetEntity();
+            if (m_rootEntity != newRoot)
+            {
+                m_rootEntity = newRoot;
+                m_networkHierarchyChangedEvent.Signal(m_rootEntity->GetId());
+            }
        }
        else
        {
-            m_hierarchyRootComponent = nullptr;
-            m_networkHierarchyLeaveEvent.Signal();
+            if (!m_rootEntity)
+            {
+                m_rootEntity = nullptr;
+                m_networkHierarchyLeaveEvent.Signal();
+            }
+        }
+    }
+
+    void NetworkHierarchyChildComponent::NotifyChildrenHierarchyDisbanded()
+    {
+        AZStd::vector<AZ::EntityId> allChildren;
+        AZ::TransformBus::EventResult(allChildren, GetEntityId(), &AZ::TransformBus::Events::GetChildren);
+        for (const AZ::EntityId& childEntityId : allChildren)
+        {
+            if (const AZ::Entity* childEntity = AZ::Interface<AZ::ComponentApplicationRequests>::Get()->FindEntity(childEntityId))
+            {
+                if (auto* hierarchyChildComponent = childEntity->FindComponent<NetworkHierarchyChildComponent>())
+                {
+                    hierarchyChildComponent->SetTopLevelHierarchyRootEntity(nullptr);
+                }
+                else if (auto* hierarchyRootComponent = childEntity->FindComponent<NetworkHierarchyRootComponent>())
+                {
+                    hierarchyRootComponent->SetTopLevelHierarchyRootEntity(nullptr);
+                }
+            }
        }
    }
 }
@@ -11,6 +11,7 @@
 #include <AzCore/Console/ILogger.h>
 #include <AzCore/Interface/Interface.h>
 #include <AzCore/Serialization/EditContext.h>
+#include <AzFramework/Components/TransformComponent.h>
 #include <Multiplayer/IMultiplayer.h>
 #include <Multiplayer/Components/NetBindComponent.h>
 #include <Multiplayer/Components/NetworkHierarchyChildComponent.h>
@@ -58,6 +59,12 @@ namespace Multiplayer
        incompatible.push_back(AZ_CRC_CE("NetworkHierarchyRootComponent"));
    }

+    NetworkHierarchyRootComponent::NetworkHierarchyRootComponent()
+        : m_childChangedHandler([this](AZ::ChildChangeType type, AZ::EntityId child) { OnChildChanged(type, child); })
+        , m_parentChangedHandler([this](AZ::EntityId oldParent, AZ::EntityId parent) { OnParentChanged(oldParent, parent); })
+    {
+    }
+
    void NetworkHierarchyRootComponent::OnInit()
    {
    }
@@ -67,31 +74,54 @@ namespace Multiplayer
        m_hierarchicalEntities.push_back(GetEntity());

        NetworkHierarchyRequestBus::Handler::BusConnect(GetEntityId());
-        AZ::TransformNotificationBus::MultiHandler::BusConnect(GetEntityId());
+
+        if (AzFramework::TransformComponent* transformComponent = GetEntity()->FindComponent<AzFramework::TransformComponent>())
+        {
+            transformComponent->BindChildChangedEventHandler(m_childChangedHandler);
+            transformComponent->BindParentChangedEventHandler(m_parentChangedHandler);
+        }
    }

    void NetworkHierarchyRootComponent::OnDeactivate([[maybe_unused]] Multiplayer::EntityIsMigrating entityIsMigrating)
    {
-        AZ::TransformNotificationBus::MultiHandler::BusDisconnect();
-        NetworkHierarchyRequestBus::Handler::BusDisconnect();
+        m_isDeactivating = true;

-        for (const AZ::Entity* childEntity : m_hierarchicalEntities)
+        if (m_rootEntity)
        {
-            auto* hierarchyChildComponent = childEntity->FindComponent<NetworkHierarchyChildComponent>();
-            auto* hierarchyRootComponent = childEntity->FindComponent<NetworkHierarchyRootComponent>();
-
-            if (hierarchyChildComponent)
+            // tell parent to re-build the hierarchy
+            if (NetworkHierarchyRootComponent* root = m_rootEntity->FindComponent<NetworkHierarchyRootComponent>())
            {
-                hierarchyChildComponent->SetTopLevelHierarchyRootComponent(nullptr);
+                root->RebuildHierarchy();
            }
-            if (hierarchyRootComponent)
+        }
+        else
+        {
+            // notify children to that a hierarchy is disbanded
+
+            AZStd::vector<AZ::EntityId> allChildren;
+            AZ::TransformBus::EventResult(allChildren, GetEntityId(), &AZ::TransformBus::Events::GetChildren);
+
+            for (const AZ::EntityId& childEntityId : allChildren)
            {
-                hierarchyRootComponent->SetTopLevelHierarchyRootEntity(nullptr);
+                if (const AZ::Entity* childEntity = AZ::Interface<AZ::ComponentApplicationRequests>::Get()->FindEntity(childEntityId))
+                {
+                    SetRootForEntity(nullptr, childEntity);
+                }
            }
        }

+        m_childChangedHandler.Disconnect();
+        m_parentChangedHandler.Disconnect();
+
+        NetworkHierarchyRequestBus::Handler::BusDisconnect();
+
        m_hierarchicalEntities.clear();
-        m_higherRootEntity = nullptr;
+        m_rootEntity = nullptr;
+    }
+
+    bool NetworkHierarchyRootComponent::IsHierarchyEnabled() const
+    {
+        return !m_isDeactivating;
    }

    bool NetworkHierarchyRootComponent::IsHierarchicalRoot() const
@@ -116,9 +146,9 @@ namespace Multiplayer

    AZ::Entity* NetworkHierarchyRootComponent::GetHierarchicalRoot() const
    {
-        if (m_higherRootEntity)
+        if (m_rootEntity)
        {
-            return m_higherRootEntity;
+            return m_rootEntity;
        }

        return GetEntity();
@@ -134,24 +164,36 @@ namespace Multiplayer
        handler.Connect(m_networkHierarchyLeaveEvent);
    }

+    void NetworkHierarchyRootComponent::OnChildChanged([[maybe_unused]] AZ::ChildChangeType type, [[maybe_unused]] AZ::EntityId child)
+    {
+        if (IsHierarchicalRoot())
+        {
+            // Parent-child notifications are not reliable enough to avoid duplicate notifications,
+            // so we will rebuild from scratch to avoid duplicate entries in @m_hierarchicalEntities.
+            RebuildHierarchy();
+        }
+        else if (NetworkHierarchyRootComponent* root = GetHierarchicalRoot()->FindComponent<NetworkHierarchyRootComponent>())
+        {
+            root->RebuildHierarchy();
+        }
+    }
+
    void NetworkHierarchyRootComponent::OnParentChanged([[maybe_unused]] AZ::EntityId oldParent, AZ::EntityId newParent)
    {
-        const AZ::EntityId entityBusId = *AZ::TransformNotificationBus::GetCurrentBusId();
-        if (GetEntityId() != entityBusId)
-        {
-            return; // ignore parent changes of child entities
-        }
+        // If the parent is part of a hierarchy, it will detect this entity as a new child and rebuild hierarchy.
+        // Thus, we only need to take care of a case when the parent is not part of a hierarchy,
+        // in which case, this entity will be a new root of a new hierarchy.

        if (AZ::Entity* parentEntity = AZ::Interface<AZ::ComponentApplicationRequests>::Get()->FindEntity(newParent))
        {
-            if (parentEntity->FindComponent<NetworkHierarchyRootComponent>())
+            if (parentEntity->FindComponent<NetworkHierarchyRootComponent>() == nullptr &&
+                parentEntity->FindComponent<NetworkHierarchyChildComponent>() == nullptr)
+            {
+                RebuildHierarchy();
+            }
+            else
            {
-                m_higherRootEntity = parentEntity;
                m_hierarchicalEntities.clear();
-                AZ::TransformNotificationBus::MultiHandler::BusDisconnect();
-
-                // Should still listen for its events, such as when this root detaches
-                AZ::TransformNotificationBus::MultiHandler::BusConnect(GetEntityId());
            }
        }
        else
@@ -161,25 +203,63 @@ namespace Multiplayer
        }
    }

-    void NetworkHierarchyRootComponent::OnChildAdded([[maybe_unused]] AZ::EntityId child)
-    {
-        // Parent-child notifications from TransformNotificationBus are not reliable enough to avoid duplicate notifications,
-        // so we will rebuild from scratch to avoid duplicate entries in @m_hierarchicalEntities.
-        RebuildHierarchy();
-    }
-
    void NetworkHierarchyRootComponent::RebuildHierarchy()
    {
-        m_hierarchicalEntities.clear();
-        m_hierarchicalEntities.push_back(GetEntity()); // add the root itself
+        AZStd::vector<AZ::Entity*> previousEntities;
+        m_hierarchicalEntities.swap(previousEntities);

-        AZ::TransformNotificationBus::MultiHandler::BusDisconnect();
-        AZ::TransformNotificationBus::MultiHandler::BusConnect(GetEntityId());
+        m_hierarchicalEntities.push_back(GetEntity()); // Add the root.

        uint32_t currentEntityCount = aznumeric_cast<uint32_t>(m_hierarchicalEntities.size());
        RecursiveAttachHierarchicalEntities(GetEntityId(), currentEntityCount);

-        m_networkHierarchyChangedEvent.Signal(GetEntityId());
+        bool hierarchyChanged = false;
+
+        // Send out join and leave events.
+
+        for (AZ::Entity* currentEntity : m_hierarchicalEntities)
+        {
+            const auto prevEntityIterator = AZStd::find(previousEntities.begin(), previousEntities.end(), currentEntity);
+            if (prevEntityIterator != previousEntities.end())
+            {
+                // This entity was here before the build of the hierarchy.
+                previousEntities.erase(prevEntityIterator);
+            }
+            else
+            {
+                // This is a newly added entity to the network hierarchy.
+                hierarchyChanged = true;
+                SetRootForEntity(GetEntity(), currentEntity);
+            }
+        }
+
+        // These entities were removed since last rebuild.
+        for (const AZ::Entity* previousEntity : previousEntities)
+        {
+            SetRootForEntity(nullptr, previousEntity);
+        }
+
+        if (!previousEntities.empty())
+        {
+            hierarchyChanged = true;
+        }
+
+        if (hierarchyChanged)
+        {
+            m_networkHierarchyChangedEvent.Signal(GetEntityId());
+        }
+    }
+
+    void NetworkHierarchyRootComponent::SetRootForEntity(AZ::Entity* root, const AZ::Entity* childEntity)
+    {
+        if (auto* hierarchyChildComponent = childEntity->FindComponent<NetworkHierarchyChildComponent>())
+        {
+            hierarchyChildComponent->SetTopLevelHierarchyRootEntity(root);
+        }
+        else if (auto* hierarchyRootComponent = childEntity->FindComponent<NetworkHierarchyRootComponent>())
+        {
+            hierarchyRootComponent->SetTopLevelHierarchyRootEntity(root);
+        }
    }

    bool NetworkHierarchyRootComponent::RecursiveAttachHierarchicalEntities(AZ::EntityId underEntity, uint32_t& currentEntityCount)
@@ -212,21 +292,12 @@ namespace Multiplayer
            auto* hierarchyChildComponent = childEntity->FindComponent<NetworkHierarchyChildComponent>();
            auto* hierarchyRootComponent = childEntity->FindComponent<NetworkHierarchyRootComponent>();

-            if (hierarchyChildComponent || hierarchyRootComponent)
+            if ((hierarchyChildComponent && hierarchyChildComponent->IsHierarchyEnabled()) ||
+                (hierarchyRootComponent && hierarchyRootComponent->IsHierarchyEnabled()))
            {
-                AZ::TransformNotificationBus::MultiHandler::BusConnect(entity);
                m_hierarchicalEntities.push_back(childEntity);
                ++currentEntityCount;

-                if (hierarchyChildComponent)
-                {
-                    hierarchyChildComponent->SetTopLevelHierarchyRootComponent(this);
-                }
-                else if (hierarchyRootComponent)
-                {
-                    hierarchyRootComponent->SetTopLevelHierarchyRootEntity(GetEntity());
-                }
-
                if (!RecursiveAttachHierarchicalEntities(entity, currentEntityCount))
                {
                    return false;
@@ -237,43 +308,10 @@ namespace Multiplayer
        return true;
    }

-    void NetworkHierarchyRootComponent::OnChildRemoved(AZ::EntityId childRemovedId)
-    {
-        AZStd::vector<AZ::EntityId> allChildren;
-        AZ::TransformBus::EventResult(allChildren, childRemovedId, &AZ::TransformBus::Events::GetEntityAndAllDescendants);
-
-        for (AZ::EntityId childId : allChildren)
-        {
-            AZ::TransformNotificationBus::MultiHandler::BusDisconnect(childId);
-
-            const AZ::Entity* childEntity = nullptr;
-
-            AZStd::erase_if(m_hierarchicalEntities, [childId, &childEntity](const AZ::Entity* entity)
-                {
-                    if (entity->GetId() == childId)
-                    {
-                        childEntity = entity;
-                        return true;
-                    }
-
-                    return false;
-                });
-
-            if (childEntity)
-            {
-                if (NetworkHierarchyChildComponent* childComponent = childEntity->FindComponent<NetworkHierarchyChildComponent>())
-                {
-                    childComponent->SetTopLevelHierarchyRootComponent(nullptr);
-                }
-            }
-        }
-
-        m_networkHierarchyChangedEvent.Signal(GetEntityId());
-    }
-
    void NetworkHierarchyRootComponent::SetTopLevelHierarchyRootEntity(AZ::Entity* hierarchyRoot)
    {
-        m_higherRootEntity = hierarchyRoot;
+        m_rootEntity = hierarchyRoot;
+
        if (HasController() && GetNetBindComponent()->GetNetEntityRole() == NetEntityRole::Authority)
        {
            NetworkHierarchyChildComponentController* controller = static_cast<NetworkHierarchyChildComponentController*>(GetController());
@@ -287,5 +325,11 @@ namespace Multiplayer
                controller->SetHierarchyRoot(InvalidNetEntityId);
            }
        }
+
+        if (m_rootEntity == nullptr)
+        {
+            // We lost the parent hierarchical entity, so as a root we need to re-build our own hierarchy.
+            RebuildHierarchy();
+        }
    }
 }
@@ -847,6 +847,41 @@ namespace Multiplayer
        );
    }

+    TEST_F(ServerHierarchyOfHierarchyTests, Inner_Root_Has_Child_References_After_Top_Root_Deactivates)
+    {
+        m_rootEntity2->FindComponent<AzFramework::TransformComponent>()->SetParent(m_childOfChildEntity->GetId());
+
+        StopAndDeleteEntity(m_rootEntity);
+
+        EXPECT_EQ(
+            m_rootEntity2->FindComponent<NetworkHierarchyRootComponent>()->GetHierarchicalEntities().size(),
+            3
+        );
+    }
+
+    TEST_F(ServerHierarchyOfHierarchyTests, Inner_Root_Has_Child_References_After_Child_Of_Top_Root_Deactivates)
+    {
+        m_rootEntity2->FindComponent<AzFramework::TransformComponent>()->SetParent(m_childOfChildEntity->GetId());
+
+        StopAndDeleteEntity(m_childEntity);
+
+        EXPECT_EQ(
+            m_rootEntity2->FindComponent<NetworkHierarchyRootComponent>()->GetHierarchicalEntities().size(),
+            3
+        );
+    }
+
+    TEST_F(ServerHierarchyOfHierarchyTests, Inner_Root_Has_Child_References_After_Child_Of_Child_Deactivates)
+    {
+        m_rootEntity2->FindComponent<AzFramework::TransformComponent>()->SetParent(m_childOfChildEntity->GetId());
+        StopAndDeleteEntity(m_childOfChildEntity);
+
+        EXPECT_EQ(
+            m_rootEntity2->FindComponent<NetworkHierarchyRootComponent>()->GetHierarchicalEntities().size(),
+            3
+        );
+    }
+
    TEST_F(ServerHierarchyOfHierarchyTests, Stress_Test_Inner_Root_Has_Child_References_After_Detachment_From_Child_Of_Child)
    {
        for (int i = 0; i < 100; ++i)