From 62449d7c14ad91558f968b1d280a3b17107188ad Mon Sep 17 00:00:00 2001
From: bd_ <bd_@nadena.dev>
Date: Sun, 1 Sep 2024 14:58:42 -0700
Subject: [PATCH] feat: add AsyncProfiler (#360)

* feat: add AsyncProfiler

This allows for profiling async tasks in the unity profiler view

* perf: profile NodeController.Create/Refresh

* chore: CHANGELOG update
---
 CHANGELOG.md                                  |  2 +
 Editor/AsyncProfiler.cs                       | 89 +++++++++++++++++++
 Editor/AsyncProfiler.cs.meta                  |  3 +
 .../PreviewSystem/Rendering/NodeController.cs |  2 +
 4 files changed, 96 insertions(+)
 create mode 100644 Editor/AsyncProfiler.cs
 create mode 100644 Editor/AsyncProfiler.cs.meta

diff --git a/CHANGELOG.md b/CHANGELOG.md
index a982e08..6f5eba0 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 
 ### Added
+- [#360] Added `AsyncProfiler` to help profile code running in Tasks
+  - Also added profiler scopes for `IRenderFilter.Create`/`IRenderFilterNode.Refresh`. 
 - [#361] Added `IRenderFilterNode.OnFrameGroup`
 
 ### Fixed
diff --git a/Editor/AsyncProfiler.cs b/Editor/AsyncProfiler.cs
new file mode 100644
index 0000000..871476d
--- /dev/null
+++ b/Editor/AsyncProfiler.cs
@@ -0,0 +1,89 @@
+using System;
+using System.Threading;
+using UnityEngine.Profiling;
+using Object = UnityEngine.Object;
+
+namespace nadena.dev.ndmf
+{
+    public static class AsyncProfiler
+    {
+        private class ProfilerFrame
+        {
+            public ProfilerFrame Parent;
+            public ProfilerFrame Root;
+            public int Depth;
+
+            public string Context;
+            public Object Object;
+        }
+
+        private static readonly AsyncLocal<ProfilerFrame> _currentFrame = new(OnFrameChange);
+
+        private static void OnFrameChange(AsyncLocalValueChangedArgs<ProfilerFrame> obj)
+        {
+            var currentFrame = obj.PreviousValue;
+            var newFrame = obj.CurrentValue;
+
+            while (currentFrame != null && (newFrame == null || newFrame.Root != currentFrame.Root ||
+                                            newFrame.Depth <= currentFrame.Depth))
+            {
+                Profiler.EndSample();
+                currentFrame = currentFrame.Parent;
+            }
+
+            EnterFrameRecursive(newFrame, currentFrame?.Depth ?? -1);
+        }
+
+        private static void EnterFrameRecursive(ProfilerFrame newFrame, int currentFrameDepth)
+        {
+            if (newFrame == null || newFrame.Depth <= currentFrameDepth) return;
+
+            if (newFrame.Parent != null) EnterFrameRecursive(newFrame.Parent, currentFrameDepth);
+
+            Profiler.BeginSample(newFrame.Context, newFrame.Object);
+        }
+
+        public static IDisposable PushProfilerContext(string context, Object obj = null)
+        {
+            var currentFrame = _currentFrame.Value;
+
+            var newFrame = new ProfilerFrame
+            {
+                Parent = currentFrame,
+                Root = currentFrame?.Root,
+                Depth = currentFrame?.Depth + 1 ?? 0,
+                Context = context,
+                Object = obj
+            };
+
+            if (newFrame.Root == null) newFrame.Root = newFrame;
+
+            _currentFrame.Value = newFrame;
+
+            return new PopFrame(newFrame);
+        }
+
+        public static void PopProfilerContext()
+        {
+            _currentFrame.Value = _currentFrame.Value?.Parent;
+        }
+
+        private class PopFrame : IDisposable
+        {
+            private readonly ProfilerFrame _targetFrame;
+
+            public PopFrame(ProfilerFrame targetFrame)
+            {
+                _targetFrame = targetFrame;
+            }
+
+            public void Dispose()
+            {
+                var currentFrame = _currentFrame.Value;
+                if (currentFrame.Root != _targetFrame.Root || currentFrame.Depth < _targetFrame.Depth) return;
+
+                _currentFrame.Value = _targetFrame;
+            }
+        }
+    }
+}
\ No newline at end of file
diff --git a/Editor/AsyncProfiler.cs.meta b/Editor/AsyncProfiler.cs.meta
new file mode 100644
index 0000000..e1caa54
--- /dev/null
+++ b/Editor/AsyncProfiler.cs.meta
@@ -0,0 +1,3 @@
+fileFormatVersion: 2
+guid: 57aa949093c3444a9bca50f3c4c99930
+timeCreated: 1725221587
\ No newline at end of file
diff --git a/Editor/PreviewSystem/Rendering/NodeController.cs b/Editor/PreviewSystem/Rendering/NodeController.cs
index e6c3144..1ecb2f0 100644
--- a/Editor/PreviewSystem/Rendering/NodeController.cs
+++ b/Editor/PreviewSystem/Rendering/NodeController.cs
@@ -104,6 +104,7 @@ public static async Task<NodeController> Create(
             string trace
         )
         {
+            AsyncProfiler.PushProfilerContext("NodeController.Create[" + filter + "]", group.Renderers[0].gameObject);
             var context =
                 new ComputeContext("NodeController " + trace + " for " + filter + " on " +
                                    group.Renderers[0].gameObject.name);
@@ -149,6 +150,7 @@ public async Task<NodeController> Refresh(
             string trace
         )
         {
+            AsyncProfiler.PushProfilerContext("NodeController.Refresh[" + _filter + "]", _group.Renderers[0].gameObject);
             var registry = ObjectRegistry.Merge(null, proxies.Select(p => p.Item3)
                 .Append(ObjectRegistry));
             var context = new ComputeContext("NodeController (refresh) for " + _filter + " on " +