From 9ab482d0afd773703f78d51a2ba8a160896f03c6 Mon Sep 17 00:00:00 2001
From: Fabian Mastenbroek <mail.fabianm@gmail.com>
Date: Tue, 16 Mar 2021 20:13:51 +0100
Subject: utils: Prevent memory leakage in TimerScheduler

This change fixes a possible memory leakage issue in TimerScheduler when
a large number of timers was scheduled for the same timestamp.
---
 .../main/kotlin/org/opendc/utils/TimerScheduler.kt | 27 ++++++++++++++--------
 1 file changed, 18 insertions(+), 9 deletions(-)

(limited to 'simulator')

diff --git a/simulator/opendc-utils/src/main/kotlin/org/opendc/utils/TimerScheduler.kt b/simulator/opendc-utils/src/main/kotlin/org/opendc/utils/TimerScheduler.kt
index bb6f3299..9f40f26a 100644
--- a/simulator/opendc-utils/src/main/kotlin/org/opendc/utils/TimerScheduler.kt
+++ b/simulator/opendc-utils/src/main/kotlin/org/opendc/utils/TimerScheduler.kt
@@ -59,7 +59,9 @@ public class TimerScheduler<T>(private val coroutineScope: CoroutineScope, priva
      * The scheduling job.
      */
     private val job = coroutineScope.launch {
+        val timers = timers
         val queue = queue
+        val clock = clock
         var next: Long? = channel.receive()
 
         while (true) {
@@ -176,17 +178,24 @@ public class TimerScheduler<T>(private val coroutineScope: CoroutineScope, priva
         require(timestamp >= now) { "Timestamp must be in the future" }
         check(job.isActive) { "Timer is stopped" }
 
-        val timer = Timer(key, timestamp, block)
-
         timers.compute(key) { _, old ->
-            old?.isCancelled = true
-            timer
-        }
-        queue.add(timer)
+            if (old?.timestamp == timestamp) {
+                // Fast-path: timer for the same timestamp already exists
+                old
+            } else {
+                // Slow-path: cancel old timer and replace it with new timer
+                val timer = Timer(key, timestamp, block)
 
-        // Check if we need to push the interruption forward
-        if (queue.peek() == timer) {
-            channel.sendBlocking(timer.timestamp)
+                old?.isCancelled = true
+                queue.add(timer)
+
+                // Check if we need to push the interruption forward
+                if (queue.peek() == timer) {
+                    channel.sendBlocking(timer.timestamp)
+                }
+
+                timer
+            }
         }
     }
 
-- 
cgit v1.2.3