From 6e8aa2dcbcd08d42b535885715e294e069aafbd4 Mon Sep 17 00:00:00 2001
From: Matthew Chen <matthew@signal.org>
Date: Sat, 23 Oct 2021 10:21:42 -0300
Subject: [PATCH] Elaborate memory usage logging.

---
 .../NotificationService.swift                 | 46 +++++++++++++------
 .../src/Messages/MessageProcessor.swift       |  4 +-
 SignalServiceKit/src/Util/PendingTasks.swift  |  2 +-
 3 files changed, 34 insertions(+), 18 deletions(-)

diff --git a/NotificationServiceExtension/NotificationService.swift b/NotificationServiceExtension/NotificationService.swift
index 100bea00a3..6b7a86a90d 100644
--- a/NotificationServiceExtension/NotificationService.swift
+++ b/NotificationServiceExtension/NotificationService.swift
@@ -36,21 +36,43 @@ class NotificationService: UNNotificationServiceExtension {
     private typealias ContentHandler = (UNNotificationContent) -> Void
     private var contentHandler = AtomicOptional<ContentHandler>(nil)
 
-    private var logTimer: OffMainThreadTimer?
+    // MARK: -
 
-    private static let nseCounter = AtomicUInt(0)
+    private static let unfairLock = UnfairLock()
+    private static var _logTimer: OffMainThreadTimer?
+    private static var _nseCounter: Int = 0
 
-    deinit {
-        logTimer?.invalidate()
-        logTimer = nil
+    private static func nseDidStart() -> Int {
+        unfairLock.withLock {
+            if _logTimer == nil {
+                _logTimer = OffMainThreadTimer(timeInterval: 1.0, repeats: true) { _ in
+                    Logger.info("... memoryUsage: \(LocalDevice.memoryUsage)")
+                }
+            }
+
+            _nseCounter = _nseCounter + 1
+            return _nseCounter
+        }
     }
 
+    private static func nseDidComplete() -> Int {
+        unfairLock.withLock {
+            _nseCounter = _nseCounter > 0 ? _nseCounter - 1 : 0
+
+            if _nseCounter == 0, _logTimer != nil {
+                _logTimer?.invalidate()
+                _logTimer = nil
+            }
+            return _nseCounter
+        }
+    }
+
+    // MARK: -
+
     // This method is thread-safe.
     func completeSilenty(timeHasExpired: Bool = false) {
-        let nseCount = Self.nseCounter.decrementOrZero()
 
-        logTimer?.invalidate()
-        logTimer = nil
+        let nseCount = Self.nseDidComplete()
 
         guard let contentHandler = contentHandler.swap(nil) else {
             if DebugFlags.internalLogging {
@@ -114,16 +136,10 @@ class NotificationService: UNNotificationServiceExtension {
 
         owsAssertDebug(FeatureFlags.notificationServiceExtension)
 
-        let nseCount = Self.nseCounter.increment()
+        let nseCount = Self.nseDidStart()
 
         Logger.info("Received notification in class: \(self), thread: \(Thread.current), pid: \(ProcessInfo.processInfo.processIdentifier), memoryUsage: \(LocalDevice.memoryUsage), nseCount: \(nseCount)")
 
-        owsAssertDebug(logTimer == nil)
-        logTimer?.invalidate()
-        logTimer = OffMainThreadTimer(timeInterval: 2.0, repeats: true) { _ in
-            Logger.info("... memoryUsage: \(LocalDevice.memoryUsage)")
-        }
-
         AppReadiness.runNowOrWhenAppDidBecomeReadySync {
             environment.askMainAppToHandleReceipt { [weak self] mainAppHandledReceipt in
                 guard !mainAppHandledReceipt else {
diff --git a/SignalServiceKit/src/Messages/MessageProcessor.swift b/SignalServiceKit/src/Messages/MessageProcessor.swift
index f845ce2242..552598fb57 100644
--- a/SignalServiceKit/src/Messages/MessageProcessor.swift
+++ b/SignalServiceKit/src/Messages/MessageProcessor.swift
@@ -266,12 +266,12 @@ public class MessageProcessor: NSObject {
 
             guard !batchEnvelopes.isEmpty, messagePipelineSupervisor.isMessageProcessingPermitted else {
                 if DebugFlags.internalLogging {
-                    Logger.info("Processing complete: \(self.queuedContentCount).")
+                    Logger.info("Processing complete: \(self.queuedContentCount) (memoryUsage: \(LocalDevice.memoryUsage).")
                 }
                 return false
             }
 
-            Logger.info("Processing batch of \(batchEnvelopes.count)/\(pendingEnvelopesCount) received envelope(s).")
+            Logger.info("Processing batch of \(batchEnvelopes.count)/\(pendingEnvelopesCount) received envelope(s). (memoryUsage: \(LocalDevice.memoryUsage)")
 
             var processedEnvelopes: [PendingEnvelope] = []
             SDSDatabaseStorage.shared.write { transaction in
diff --git a/SignalServiceKit/src/Util/PendingTasks.swift b/SignalServiceKit/src/Util/PendingTasks.swift
index e4bf5fc2bc..0ff171ed69 100644
--- a/SignalServiceKit/src/Util/PendingTasks.swift
+++ b/SignalServiceKit/src/Util/PendingTasks.swift
@@ -30,7 +30,7 @@ public class PendingTasks: NSObject {
             Promise.when(resolved: promises).asVoid()
         }.map(on: .global()) {
             if DebugFlags.internalLogging {
-                Logger.info("Complete \(label).")
+                Logger.info("Complete \(label) (memoryUsage: \(LocalDevice.memoryUsage).")
             }
         }
     }