From d80226e7bd801fcbccf71f19ba373b7baf71a49e Mon Sep 17 00:00:00 2001
From: Takashi Kokubun <takashikkbn@gmail.com>
Date: Fri, 27 Nov 2020 22:07:02 -0800
Subject: [PATCH] Avoid unloading units which have enough total_calls

instead of just unloading worst 10% methods.
---
 mjit_worker.c         | 44 ++++++++++++++++++++++++-------------------
 test/ruby/test_jit.rb | 10 +++++++---
 2 files changed, 32 insertions(+), 22 deletions(-)

diff --git a/mjit_worker.c b/mjit_worker.c
index c708f01b63..426f1ace9f 100644
--- a/mjit_worker.c
+++ b/mjit_worker.c
@@ -1259,9 +1259,9 @@ static struct mjit_cont *first_cont;
 static void
 unload_units(void)
 {
-    struct rb_mjit_unit *unit = 0, *next, *worst;
+    struct rb_mjit_unit *unit = 0, *next;
     struct mjit_cont *cont;
-    int delete_num, units_num = active_units.length;
+    int units_num = active_units.length;
 
     // For now, we don't unload units when ISeq is GCed. We should
     // unload such ISeqs first here.
@@ -1284,29 +1284,35 @@ unload_units(void)
     }
     // TODO: check stale_units and unload unused ones! (note that the unit is not associated to ISeq anymore)
 
-    // Remove 1/10 units more to decrease unloading calls.
-    // TODO: Calculate max total_calls in unit_queue and don't unload units
-    // whose total_calls are larger than the max.
-    delete_num = active_units.length / 10;
-    for (; active_units.length > mjit_opts.max_cache_size - delete_num;) {
-        // Find one unit that has the minimum total_calls.
-        worst = NULL;
+    // Unload units whose total_calls is smaller than any total_calls in unit_queue.
+    // TODO: make the algorithm more efficient
+    long unsigned prev_queue_calls = -1;
+    while (true) {
+        // Calculate the next max total_calls in unit_queue
+        long unsigned max_queue_calls = 0;
+        list_for_each(&unit_queue.head, unit, unode) {
+            if (unit->iseq != NULL && max_queue_calls < unit->iseq->body->total_calls
+                    && unit->iseq->body->total_calls < prev_queue_calls) {
+                max_queue_calls = unit->iseq->body->total_calls;
+            }
+        }
+        prev_queue_calls = max_queue_calls;
+
+        bool unloaded_p = false;
         list_for_each(&active_units.head, unit, unode) {
             if (unit->used_code_p) // We can't unload code on stack.
                 continue;
 
-            if (worst == NULL || worst->iseq->body->total_calls > unit->iseq->body->total_calls) {
-                worst = unit;
+            if (max_queue_calls > unit->iseq->body->total_calls) {
+                verbose(2, "Unloading unit %d (calls=%lu, threshold=%lu)",
+                        unit->id, unit->iseq->body->total_calls, max_queue_calls);
+                assert(unit->handle != NULL);
+                remove_from_list(unit, &active_units);
+                free_unit(unit);
+                unloaded_p = true;
             }
         }
-        if (worst == NULL)
-            break;
-
-        // Unload the worst node.
-        verbose(2, "Unloading unit %d (calls=%lu)", worst->id, worst->iseq->body->total_calls);
-        assert(worst->handle != NULL);
-        remove_from_list(worst, &active_units);
-        free_unit(worst);
+        if (!unloaded_p) break;
     }
 
     if (units_num > active_units.length) {
diff --git a/test/ruby/test_jit.rb b/test/ruby/test_jit.rb
index f1019b92c2..5585264de6 100644
--- a/test/ruby/test_jit.rb
+++ b/test/ruby/test_jit.rb
@@ -690,11 +690,15 @@ class TestJIT < Test::Unit::TestCase
         assert_match(/\A#{JIT_SUCCESS_PREFIX}: mjit#{i}@\(eval\):/, errs[i], debug_info)
       end
 
-      assert_equal("Too many JIT code -- 1 units unloaded\n", errs[10], debug_info)
-      assert_match(/\A#{JIT_SUCCESS_PREFIX}: mjit10@\(eval\):/, errs[11], debug_info)
       # On --jit-wait, when the number of JIT-ed code reaches --jit-max-cache,
       # it should trigger compaction.
-      unless RUBY_PLATFORM.match?(/mswin|mingw/) # compaction is not supported on Windows yet
+      if RUBY_PLATFORM.match?(/mswin|mingw/) # compaction is not supported on Windows yet
+        assert_equal("Too many JIT code -- 1 units unloaded\n", errs[10], debug_info)
+        assert_match(/\A#{JIT_SUCCESS_PREFIX}: mjit10@\(eval\):/, errs[11], debug_info)
+      else
+        assert_equal("No units can be unloaded -- incremented max-cache-size to 11 for --jit-wait\n", errs[10], debug_info)
+        assert_match(/\A#{JIT_SUCCESS_PREFIX}: mjit10@\(eval\):/, errs[11], debug_info)
+
         assert_equal(3, compactions.size, debug_info)
       end