fix: preview system fails to recover from the loss of the primary proxy object

CHANGELOG.md

@@ -11,6 +11,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- [#460] Preview system fails to recover when the primary proxy is destroyed
+
 ### Changed
 
 ### Removed

Editor/PreviewSystem/Rendering/ProxyObjectCache.cs

@@ -31,6 +31,7 @@ private class ProxyHandleImpl : IProxyHandle
             private readonly Renderer _key;
             private readonly Func<Renderer> _createFunc;
             private readonly RendererState _state;
+            private bool _disposed;
 
             public ProxyHandleImpl(ProxyObjectCache cache, Renderer key, Func<Renderer> createFunc, RendererState state)
             {
@@ -70,6 +71,14 @@ public void ReturnSetupProxy(Renderer proxy)
 
             public void Dispose()
             {
+                if (_disposed)
+                {
+                    Debug.LogWarning("Proxy handle was disposed twice!");
+                    throw new ObjectDisposedException(nameof(ProxyHandleImpl));
+                }
+
+                _disposed = true;
+
                 if (--_state.ActivePrimaryCount == 0)
                 {
                     _cache.MaybeDisposeProxy(_key);
@@ -118,30 +127,43 @@ private void OnPlayModeStateChanged(PlayModeStateChange obj)
         public IProxyHandle GetHandle(Renderer original, Func<Renderer> create)
         {
             IsRegistered = true;
+
+            Func<Renderer> createShimmed = () =>
+            {
+                var newProxy = create();
+                newProxy.gameObject.AddComponent<ProxyTagComponent>();
+                _proxyObjectInstanceIds.Add(newProxy.gameObject.GetInstanceID());
+
+                return newProxy;
+            };
 
             if (!_renderers.TryGetValue(original, out var state))
             {
                 state = new RendererState();
-                state.PrimaryProxy = create();
+                state.PrimaryProxy = createShimmed();
                 _renderers.Add(original, state);
             }
 
-            state.ActivePrimaryCount++;
-
-            return new ProxyHandleImpl(this, original, () =>
+            if (state.PrimaryProxy == null)
             {
-                var newProxy = create();
-                _proxyObjectInstanceIds.Add(newProxy.gameObject.GetInstanceID());
+                // Recover from loss of the primary proxy
+                state.PrimaryProxy = createShimmed();
+            }
 
-                return newProxy;
-            }, state);
+            state.ActivePrimaryCount++;
+
+            return new ProxyHandleImpl(this, original, createShimmed, state);
         }
 
         private static void DestroyProxy(Renderer proxy)
         {
             if (proxy == null) return;
 
             var gameObject = proxy.gameObject;
+            if (gameObject.TryGetComponent<ProxyTagComponent>(out var tag))
+            {
+                tag.Armed = false;
+            }
             _proxyObjectInstanceIds.Remove(gameObject.GetInstanceID());
             Object.DestroyImmediate(gameObject);
         }

Editor/PreviewSystem/Rendering/ProxyObjectController.cs

@@ -1,11 +1,13 @@
 #region
 
 using System;
+using System.Diagnostics;
 using System.Linq;
 using UnityEditor;
 using UnityEngine;
 using UnityEngine.Profiling;
 using UnityEngine.SceneManagement;
+using Debug = UnityEngine.Debug;
 using Object = UnityEngine.Object;
 
 #endregion
@@ -33,6 +35,8 @@ internal class ProxyObjectController : IDisposable
         private bool _visibilityOffOriginal;
         private bool _pickingOffOriginal, _pickingOffReplacement;
         private long _lastVisibilityCheck = long.MinValue;
+
+        private readonly Stopwatch _lastWarning = new();
 
         private static CustomSampler _onPreFrameSampler = CustomSampler.Create("ProxyObjectController.OnPreFrame");
 
@@ -137,7 +141,11 @@ internal bool OnPreFrame()
             {
                 if (Renderer == null)
                 {
-                    Debug.LogWarning("Proxy object was destroyed improperly! Resetting pipeline...");
+                    if (!_lastWarning.IsRunning || _lastWarning.ElapsedMilliseconds > 1000)
+                    {
+                        Debug.LogWarning("Proxy object was destroyed improperly! Resetting pipeline...");
+                        _lastWarning.Restart();
+                    }
                 }
                 return false;
             }

Editor/PreviewSystem/Rendering/ProxyPipeline.cs

@@ -303,6 +303,12 @@ await Task.WhenAll(_stages.SelectMany(s => s.NodeTasks))
             {
                 proxy.FinishSetup();
 
+                if (proxy.Renderer == null)
+                {
+                    Invalidate();
+                    continue;
+                }
+
                 OriginalToProxyRenderer = OriginalToProxyRenderer.Add(r, proxy.Renderer);
                 OriginalToProxyObject = OriginalToProxyObject.Add(r.gameObject, proxy.Renderer.gameObject);
                 ProxyToOriginalObject = ProxyToOriginalObject.Add(proxy.Renderer.gameObject, r.gameObject);

Runtime/ProxyTagComponent.cs

@@ -0,0 +1,20 @@
+using UnityEngine;
+
+namespace nadena.dev.ndmf.preview
+{
+    [AddComponentMenu("/")]
+    internal class ProxyTagComponent : MonoBehaviour
+    {
+        internal bool Armed = true;
+
+        private void OnDestroy()
+        {
+#if NDMF_DEBUG
+            if (Armed)
+            {
+                Debug.LogWarning("Proxy object was destroyed improperly here!");
+            }
+#endif
+        }
+    }
+}