From ee1825219b8ccca13df7198d4e9ffb966e44c883 Mon Sep 17 00:00:00 2001
From: FICTURE7 <FICTURE7@gmail.com>
Date: Fri, 9 Sep 2022 03:14:08 +0400
Subject: [PATCH] Clean up rejit queue (#2751)

---
 ARMeilleure/Common/AddressTable.cs            |   4 +-
 .../EventSources/AddressTableEventSource.cs   |  51 --------
 .../Diagnostics/TranslatorEventSource.cs      |  67 ++++++++++
 ARMeilleure/Translation/Translator.cs         | 101 +++++----------
 ARMeilleure/Translation/TranslatorQueue.cs    | 121 ++++++++++++++++++
 5 files changed, 225 insertions(+), 119 deletions(-)
 delete mode 100644 ARMeilleure/Diagnostics/EventSources/AddressTableEventSource.cs
 create mode 100644 ARMeilleure/Diagnostics/TranslatorEventSource.cs
 create mode 100644 ARMeilleure/Translation/TranslatorQueue.cs

diff --git a/ARMeilleure/Common/AddressTable.cs b/ARMeilleure/Common/AddressTable.cs
index e03a38e67..7cb83fdd2 100644
--- a/ARMeilleure/Common/AddressTable.cs
+++ b/ARMeilleure/Common/AddressTable.cs
@@ -1,4 +1,4 @@
-using ARMeilleure.Diagnostics.EventSources;
+using ARMeilleure.Diagnostics;
 using System;
 using System.Collections.Generic;
 using System.Runtime.InteropServices;
@@ -218,7 +218,7 @@ namespace ARMeilleure.Common
 
             _pages.Add(page);
 
-            AddressTableEventSource.Log.Allocated(size, leaf);
+            TranslatorEventSource.Log.AddressTableAllocated(size, leaf);
 
             return page;
         }
diff --git a/ARMeilleure/Diagnostics/EventSources/AddressTableEventSource.cs b/ARMeilleure/Diagnostics/EventSources/AddressTableEventSource.cs
deleted file mode 100644
index 201a25623..000000000
--- a/ARMeilleure/Diagnostics/EventSources/AddressTableEventSource.cs
+++ /dev/null
@@ -1,51 +0,0 @@
-using System.Diagnostics.Tracing;
-
-namespace ARMeilleure.Diagnostics.EventSources
-{
-    [EventSource(Name = "ARMeilleure")]
-    class AddressTableEventSource : EventSource
-    {
-        public static readonly AddressTableEventSource Log = new();
-
-        private ulong _size;
-        private ulong _leafSize;
-        private PollingCounter _sizeCounter;
-        private PollingCounter _leafSizeCounter;
-
-        public AddressTableEventSource()
-        {
-            _sizeCounter = new PollingCounter("addr-tab-alloc", this, () => _size / 1024d / 1024d)
-            {
-                DisplayName = "AddressTable Total Bytes Allocated",
-                DisplayUnits = "MB"
-            };
-
-            _leafSizeCounter = new PollingCounter("addr-tab-leaf-alloc", this, () => _leafSize / 1024d / 1024d)
-            {
-                DisplayName = "AddressTable Total Leaf Bytes Allocated",
-                DisplayUnits = "MB"
-            };
-        }
-
-        public void Allocated(int bytes, bool leaf)
-        {
-            _size += (uint)bytes;
-
-            if (leaf)
-            {
-                _leafSize += (uint)bytes;
-            }
-        }
-
-        protected override void Dispose(bool disposing)
-        {
-            _leafSizeCounter.Dispose();
-            _leafSizeCounter = null;
-
-            _sizeCounter.Dispose();
-            _sizeCounter = null;
-
-            base.Dispose(disposing);
-        }
-    }
-}
diff --git a/ARMeilleure/Diagnostics/TranslatorEventSource.cs b/ARMeilleure/Diagnostics/TranslatorEventSource.cs
new file mode 100644
index 000000000..a0456e984
--- /dev/null
+++ b/ARMeilleure/Diagnostics/TranslatorEventSource.cs
@@ -0,0 +1,67 @@
+using System.Diagnostics.Tracing;
+using System.Threading;
+
+namespace ARMeilleure.Diagnostics
+{
+    [EventSource(Name = "ARMeilleure")]
+    class TranslatorEventSource : EventSource
+    {
+        public static readonly TranslatorEventSource Log = new();
+
+        private int _rejitQueue;
+        private ulong _funcTabSize;
+        private ulong _funcTabLeafSize;
+        private PollingCounter _rejitQueueCounter;
+        private PollingCounter _funcTabSizeCounter;
+        private PollingCounter _funcTabLeafSizeCounter;
+
+        public TranslatorEventSource()
+        {
+            _rejitQueueCounter = new PollingCounter("rejit-queue-length", this, () => _rejitQueue)
+            {
+                DisplayName = "Rejit Queue Length"
+            };
+
+            _funcTabSizeCounter = new PollingCounter("addr-tab-alloc", this, () => _funcTabSize / 1024d / 1024d)
+            {
+                DisplayName = "AddressTable Total Bytes Allocated",
+                DisplayUnits = "MB"
+            };
+
+            _funcTabLeafSizeCounter = new PollingCounter("addr-tab-leaf-alloc", this, () => _funcTabLeafSize / 1024d / 1024d)
+            {
+                DisplayName = "AddressTable Total Leaf Bytes Allocated",
+                DisplayUnits = "MB"
+            };
+        }
+
+        public void RejitQueueAdd(int count)
+        {
+            Interlocked.Add(ref _rejitQueue, count);
+        }
+
+        public void AddressTableAllocated(int bytes, bool leaf)
+        {
+            _funcTabSize += (uint)bytes;
+
+            if (leaf)
+            {
+                _funcTabLeafSize += (uint)bytes;
+            }
+        }
+
+        protected override void Dispose(bool disposing)
+        {
+            _rejitQueueCounter.Dispose();
+            _rejitQueueCounter = null;
+
+            _funcTabLeafSizeCounter.Dispose();
+            _funcTabLeafSizeCounter = null;
+
+            _funcTabSizeCounter.Dispose();
+            _funcTabSizeCounter = null;
+
+            base.Dispose(disposing);
+        }
+    }
+}
diff --git a/ARMeilleure/Translation/Translator.cs b/ARMeilleure/Translation/Translator.cs
index 389adf292..ee8e3e8b5 100644
--- a/ARMeilleure/Translation/Translator.cs
+++ b/ARMeilleure/Translation/Translator.cs
@@ -44,15 +44,11 @@ namespace ARMeilleure.Translation
         private readonly IJitMemoryAllocator _allocator;
         private readonly ConcurrentQueue<KeyValuePair<ulong, TranslatedFunction>> _oldFuncs;
 
-        private readonly ConcurrentDictionary<ulong, object> _backgroundSet;
-        private readonly ConcurrentStack<RejitRequest> _backgroundStack;
-        private readonly AutoResetEvent _backgroundTranslatorEvent;
-        private readonly ReaderWriterLock _backgroundTranslatorLock;
-
         internal TranslatorCache<TranslatedFunction> Functions { get; }
         internal AddressTable<ulong> FunctionTable { get; }
         internal EntryTable<uint> CountTable { get; }
         internal TranslatorStubs Stubs { get; }
+        internal TranslatorQueue Queue { get; }
         internal IMemoryManager Memory { get; }
 
         private volatile int _threadCount;
@@ -67,10 +63,7 @@ namespace ARMeilleure.Translation
 
             _oldFuncs = new ConcurrentQueue<KeyValuePair<ulong, TranslatedFunction>>();
 
-            _backgroundSet = new ConcurrentDictionary<ulong, object>();
-            _backgroundStack = new ConcurrentStack<RejitRequest>();
-            _backgroundTranslatorEvent = new AutoResetEvent(false);
-            _backgroundTranslatorLock = new ReaderWriterLock();
+            Queue = new TranslatorQueue();
 
             JitCache.Initialize(allocator);
 
@@ -87,43 +80,6 @@ namespace ARMeilleure.Translation
             }
         }
 
-        private void TranslateStackedSubs()
-        {
-            while (_threadCount != 0)
-            {
-                _backgroundTranslatorLock.AcquireReaderLock(Timeout.Infinite);
-
-                if (_backgroundStack.TryPop(out RejitRequest request) &&
-                    _backgroundSet.TryRemove(request.Address, out _))
-                {
-                    TranslatedFunction func = Translate(request.Address, request.Mode, highCq: true);
-
-                    Functions.AddOrUpdate(request.Address, func.GuestSize, func, (key, oldFunc) =>
-                    {
-                        EnqueueForDeletion(key, oldFunc);
-                        return func;
-                    });
-
-                    if (PtcProfiler.Enabled)
-                    {
-                        PtcProfiler.UpdateEntry(request.Address, request.Mode, highCq: true);
-                    }
-
-                    RegisterFunction(request.Address, func);
-
-                    _backgroundTranslatorLock.ReleaseReaderLock();
-                }
-                else
-                {
-                    _backgroundTranslatorLock.ReleaseReaderLock();
-                    _backgroundTranslatorEvent.WaitOne();
-                }
-            }
-
-             // Wake up any other background translator threads, to encourage them to exit.
-            _backgroundTranslatorEvent.Set();
-        }
-
         public void Execute(State.ExecutionContext context, ulong address)
         {
             if (Interlocked.Increment(ref _threadCount) == 1)
@@ -155,7 +111,7 @@ namespace ARMeilleure.Translation
                 {
                     bool last = i != 0 && i == unboundedThreadCount - 1;
 
-                    Thread backgroundTranslatorThread = new Thread(TranslateStackedSubs)
+                    Thread backgroundTranslatorThread = new Thread(BackgroundTranslate)
                     {
                         Name = "CPU.BackgroundTranslatorThread." + i,
                         Priority = last ? ThreadPriority.Lowest : ThreadPriority.Normal
@@ -186,10 +142,9 @@ namespace ARMeilleure.Translation
 
             if (Interlocked.Decrement(ref _threadCount) == 0)
             {
-                _backgroundTranslatorEvent.Set();
-
                 ClearJitCache();
 
+                Queue.Dispose();
                 Stubs.Dispose();
                 FunctionTable.Dispose();
                 CountTable.Dispose();
@@ -317,6 +272,27 @@ namespace ARMeilleure.Translation
             return new TranslatedFunction(func, counter, funcSize, highCq);
         }
 
+        private void BackgroundTranslate()
+        {
+            while (_threadCount != 0 && Queue.TryDequeue(out RejitRequest request))
+            {
+                TranslatedFunction func = Translate(request.Address, request.Mode, highCq: true);
+
+                Functions.AddOrUpdate(request.Address, func.GuestSize, func, (key, oldFunc) =>
+                {
+                    EnqueueForDeletion(key, oldFunc);
+                    return func;
+                });
+
+                if (PtcProfiler.Enabled)
+                {
+                    PtcProfiler.UpdateEntry(request.Address, request.Mode, highCq: true);
+                }
+
+                RegisterFunction(request.Address, func);
+            }
+        }
+
         private struct Range
         {
             public ulong Start { get; }
@@ -504,11 +480,7 @@ namespace ARMeilleure.Translation
 
         internal void EnqueueForRejit(ulong guestAddress, ExecutionMode mode)
         {
-            if (_backgroundSet.TryAdd(guestAddress, null))
-            {
-                _backgroundStack.Push(new RejitRequest(guestAddress, mode));
-                _backgroundTranslatorEvent.Set();
-            }
+            Queue.Enqueue(guestAddress, mode);
         }
 
         private void EnqueueForDeletion(ulong guestAddress, TranslatedFunction func)
@@ -542,26 +514,23 @@ namespace ARMeilleure.Translation
 
         private void ClearRejitQueue(bool allowRequeue)
         {
-            _backgroundTranslatorLock.AcquireWriterLock(Timeout.Infinite);
-
-            if (allowRequeue)
+            if (!allowRequeue)
             {
-                while (_backgroundStack.TryPop(out var request))
+                Queue.Clear();
+
+                return;
+            }
+
+            lock (Queue.Sync)
+            {
+                while (Queue.Count > 0 && Queue.TryDequeue(out RejitRequest request))
                 {
                     if (Functions.TryGetValue(request.Address, out var func) && func.CallCounter != null)
                     {
                         Volatile.Write(ref func.CallCounter.Value, 0);
                     }
-
-                    _backgroundSet.TryRemove(request.Address, out _);
                 }
             }
-            else
-            {
-                _backgroundStack.Clear();
-            }
-
-            _backgroundTranslatorLock.ReleaseWriterLock();
         }
     }
 }
diff --git a/ARMeilleure/Translation/TranslatorQueue.cs b/ARMeilleure/Translation/TranslatorQueue.cs
new file mode 100644
index 000000000..fc0aa64f2
--- /dev/null
+++ b/ARMeilleure/Translation/TranslatorQueue.cs
@@ -0,0 +1,121 @@
+using ARMeilleure.Diagnostics;
+using ARMeilleure.State;
+using System;
+using System.Collections.Generic;
+using System.Threading;
+
+namespace ARMeilleure.Translation
+{
+    /// <summary>
+    /// Represents a queue of <see cref="RejitRequest"/>.
+    /// </summary>
+    /// <remarks>
+    /// This does not necessarily behave like a queue, i.e: a FIFO collection.
+    /// </remarks>
+    sealed class TranslatorQueue : IDisposable
+    {
+        private bool _disposed;
+        private readonly Stack<RejitRequest> _requests;
+        private readonly HashSet<ulong> _requestAddresses;
+
+        /// <summary>
+        /// Gets the object used to synchronize access to the <see cref="TranslatorQueue"/>.
+        /// </summary>
+        public object Sync { get; }
+
+        /// <summary>
+        /// Gets the number of requests in the <see cref="TranslatorQueue"/>.
+        /// </summary>
+        public int Count => _requests.Count;
+
+        /// <summary>
+        /// Initializes a new instance of the <see cref="TranslatorQueue"/> class.
+        /// </summary>
+        public TranslatorQueue()
+        {
+            Sync = new object();
+
+            _requests = new Stack<RejitRequest>();
+            _requestAddresses = new HashSet<ulong>();
+        }
+
+        /// <summary>
+        /// Enqueues a request with the specified <paramref name="address"/> and <paramref name="mode"/>.
+        /// </summary>
+        /// <param name="address">Address of request</param>
+        /// <param name="mode"><see cref="ExecutionMode"/> of request</param>
+        public void Enqueue(ulong address, ExecutionMode mode)
+        {
+            lock (Sync)
+            {
+                if (_requestAddresses.Add(address))
+                {
+                    _requests.Push(new RejitRequest(address, mode));
+
+                    TranslatorEventSource.Log.RejitQueueAdd(1);
+
+                    Monitor.Pulse(Sync);
+                }
+            }
+        }
+
+        /// <summary>
+        /// Tries to dequeue a <see cref="RejitRequest"/>. This will block the thread until a <see cref="RejitRequest"/>
+        /// is enqueued or the <see cref="TranslatorQueue"/> is disposed.
+        /// </summary>
+        /// <param name="result"><see cref="RejitRequest"/> dequeued</param>
+        /// <returns><see langword="true"/> on success; otherwise <see langword="false"/></returns>
+        public bool TryDequeue(out RejitRequest result)
+        {
+            while (!_disposed)
+            {
+                lock (Sync)
+                {
+                    if (_requests.TryPop(out result))
+                    {
+                        _requestAddresses.Remove(result.Address);
+
+                        TranslatorEventSource.Log.RejitQueueAdd(-1);
+
+                        return true;
+                    }
+
+                    Monitor.Wait(Sync);
+                }
+            }
+
+            result = default;
+
+            return false;
+        }
+
+        /// <summary>
+        /// Clears the <see cref="TranslatorQueue"/>.
+        /// </summary>
+        public void Clear()
+        {
+            lock (Sync)
+            {
+                TranslatorEventSource.Log.RejitQueueAdd(-_requests.Count);
+
+                _requests.Clear();
+                _requestAddresses.Clear();
+
+                Monitor.PulseAll(Sync);
+            }
+        }
+
+        /// <summary>
+        /// Releases all resources used by the <see cref="TranslatorQueue"/> instance.
+        /// </summary>
+        public void Dispose()
+        {
+            if (!_disposed)
+            {
+                _disposed = true;
+
+                Clear();
+            }
+        }
+    }
+}