OSGKeyboard/OSGKeyboardShared/Utilities/KeyboardExtensionMemoryTelemetry.swift

// KeyboardExtensionMemoryTelemetry.swift
// OSGKeyboard · Shared
//
// Observes keyboard-extension memory AND sheds load before the system does.
// The host process never starts this monitor, so shared typing code can emit
// extension-only milestones without duplicating host telemetry.
//
// Observation alone was not enough: `didReceiveMemoryWarning` is the only
// system signal an extension gets, and jetsam frequently kills a keyboard at
// the ~60 MiB boundary without ever delivering one. `reliefHandler` lets the
// extension drop its heavy caches on OUR thresholds (40 / 48 MiB), which are
// deliberately below that boundary.

import Darwin
import Foundation

public enum KeyboardExtensionMemoryBudget {
    /// Start preserving evidence before the extension reaches its safe ceiling.
    public static let warningMB: Double = 36
    /// Internal release target. Apple's keyboard-extension limit is not public.
    public static let safePeakMB: Double = 40
    /// Leave headroom below the observed ~60 MiB device jetsam boundary.
    public static let criticalMB: Double = 48

    public enum Level: String, Sendable, Equatable {
        case normal
        case warning
        case high
        case critical
        case unavailable
    }

    public static func level(forPhysFootprintMB footprintMB: Double) -> Level {
        guard footprintMB >= 0 else { return .unavailable }
        if footprintMB >= criticalMB { return .critical }
        if footprintMB >= safePeakMB { return .high }
        if footprintMB >= warningMB { return .warning }
        return .normal
    }
}

@MainActor
public enum KeyboardExtensionMemoryTelemetry {
    private static let peakLogStepMB: Double = 4
    private static let bootSampleInterval = Duration.milliseconds(50)
    private static let bootSampleDuration: TimeInterval = 4

    private static var isActive = false
    private static var processID: Int32 = 0
    private static var context = "surface=- language=-"
    private static var baselineFootprintMB: Double = -1
    private static var peakFootprintMB: Double = -1
    private static var lastLoggedPeakMB: Double = -1
    private static var highestLevel = KeyboardExtensionMemoryBudget.Level.normal
    private static var startedAt: TimeInterval = 0
    private static var samplingTask: Task<Void, Never>?
    /// Highest level already relieved in the current pressure episode. Reset only
    /// when the footprint falls back below `warningMB` — see `requestReliefIfNeeded`.
    private static var lastReliefLevel = KeyboardExtensionMemoryBudget.Level.normal
    /// True while `reliefHandler` is running, so relief can never be requested
    /// from inside relief — see `requestReliefIfNeeded`.
    private static var isRelieving = false

    /// Poll interval once the startup burst is over. Sustained sampling is what
    /// catches growth *between* milestones (a long clipboard session, a big
    /// candidate list) — milestone-only sampling misses it entirely.
    private static let sustainedSampleInterval = Duration.seconds(1)
    /// Invoked on the main actor when the footprint crosses `.high` or
    /// `.critical`. The keyboard extension installs this to release the Rime
    /// engine, the English lexicon and any in-flight pipeline work. Nothing
    /// else in the process is allowed to set it.
    ///
    /// Returns whether relief was actually performed. `false` means the host
    /// declined *this* attempt (it is mid-composition and will not yank the
    /// typing surface out from under the user), and the level stays armed so the
    /// next poll asks again — a declined attempt must not be mistaken for a
    /// completed one, or the keyboard would sit at 48 MiB having shed nothing.
    public static var reliefHandler: (@MainActor @Sendable (KeyboardExtensionMemoryBudget.Level) -> Bool)?

    public static func begin(context initialContext: String) {
        samplingTask?.cancel()
        samplingTask = nil
        lastReliefLevel = .normal
        isActive = true
        processID = getpid()
        context = initialContext
        startedAt = ProcessInfo.processInfo.systemUptime

        let snapshot = OSGDiag.memorySnapshot()
        baselineFootprintMB = snapshot.physFootprintMB
        peakFootprintMB = snapshot.physFootprintMB
        lastLoggedPeakMB = snapshot.physFootprintMB
        highestLevel = .normal
        emit(stage: "process.begin", snapshot: snapshot, alwaysLog: true)
    }

    public static func updateContext(_ newContext: String) {
        guard isActive else { return }
        context = newContext
    }

    public static func record(_ stage: String, details: String? = nil) {
        guard isActive else { return }
        let eventContext = details.map { "\(context) \($0)" } ?? context
        emit(
            stage: stage,
            snapshot: OSGDiag.memorySnapshot(),
            eventContext: eventContext,
            alwaysLog: true
        )
    }

    /// Samples short-lived startup spikes that milestone-only logging can miss,
    /// then keeps polling at a low rate for the rest of the presentation so
    /// pressure that builds up mid-session still reaches `reliefHandler`.
    /// Poll samples log only on a new budget band or each additional 4 MiB peak.
    public static func startBootSampling() {
        guard isActive else { return }
        samplingTask?.cancel()
        let deadline = ProcessInfo.processInfo.systemUptime + bootSampleDuration
        samplingTask = Task { @MainActor in
            while !Task.isCancelled, ProcessInfo.processInfo.systemUptime < deadline {
                try? await Task.sleep(for: bootSampleInterval)
                guard !Task.isCancelled else { return }
                emit(
                    stage: "boot.sample",
                    snapshot: OSGDiag.memorySnapshot(),
                    eventContext: context,
                    alwaysLog: false
                )
            }
            guard !Task.isCancelled else { return }
            record("boot.sample.complete")
            await pollPressure()
        }
    }

    /// Resumes low-rate sampling for a re-presented keyboard without replaying
    /// the startup burst. No-ops while the boot burst is still running.
    public static func startSustainedSamplingIfIdle() {
        guard isActive, samplingTask == nil else { return }
        samplingTask = Task { @MainActor in
            await pollPressure()
        }
    }

    /// Stops sampling (keyboard dismissed). The extension process survives
    /// between presentations, so an uncancelled poll would keep waking a hidden
    /// keyboard forever.
    public static func stopSampling() {
        samplingTask?.cancel()
        samplingTask = nil
    }

    private static func pollPressure() async {
        while !Task.isCancelled {
            try? await Task.sleep(for: sustainedSampleInterval)
            guard !Task.isCancelled else { return }
            emit(
                stage: "pressure.sample",
                snapshot: OSGDiag.memorySnapshot(),
                eventContext: context,
                alwaysLog: false
            )
        }
    }

    private static func emit(
        stage: String,
        snapshot: OSGDiag.MemorySnapshot,
        eventContext: String? = nil,
        alwaysLog: Bool
    ) {
        let footprint = snapshot.physFootprintMB
        if footprint >= 0 {
            peakFootprintMB = max(peakFootprintMB, footprint)
        }
        let level = KeyboardExtensionMemoryBudget.level(forPhysFootprintMB: footprint)
        let crossedLevel = levelRank(level) > levelRank(highestLevel)
        if crossedLevel {
            highestLevel = level
        }
        #if DEBUG
        MemoryDeviceProbe.log("emit.\(stage).\(level.rawValue)")
        #endif
        // Shed BEFORE logging: at `.critical` we are ~12 MiB from the observed
        // jetsam boundary and the log line is the less important half.
        requestReliefIfNeeded(level: level, stage: stage)
        let peakAdvanced = peakFootprintMB >= 0
            && (lastLoggedPeakMB < 0 || peakFootprintMB - lastLoggedPeakMB >= peakLogStepMB)
        guard alwaysLog || crossedLevel || peakAdvanced else { return }
        if peakFootprintMB >= 0 {
            lastLoggedPeakMB = peakFootprintMB
        }

        let elapsedMS = max(
            0,
            Int((ProcessInfo.processInfo.systemUptime - startedAt) * 1_000)
        )
        let delta = baselineFootprintMB >= 0 && peakFootprintMB >= 0
            ? peakFootprintMB - baselineFootprintMB
            : -1
        OSGDiag.log(
            String(
                format: "extMemory pid=%d stage=%@ elapsed=%dms level=%@ crossed=%d "
                    + "rss=%.1fMB foot=%.1fMB peak=%.1fMB delta=%.1fMB "
                    + "safe=%dMB critical=%dMB context={%@}",
                processID,
                stage,
                elapsedMS,
                level.rawValue,
                crossedLevel ? 1 : 0,
                snapshot.rssMB,
                footprint,
                peakFootprintMB,
                delta,
                Int(KeyboardExtensionMemoryBudget.safePeakMB),
                Int(KeyboardExtensionMemoryBudget.criticalMB),
                eventContext ?? context
            ),
            category: "memory"
        )
    }

    /// Asks the host to release heavy resources, at most once per level per
    /// pressure episode. Escalation (`.high` → `.critical`) still fires: the
    /// harder shed must not be swallowed by the soft one that preceded it.
    private static func requestReliefIfNeeded(
        level: KeyboardExtensionMemoryBudget.Level,
        stage: String
    ) {
        // Hysteresis, not a cooldown. A timed gap re-sheds every few seconds
        // while the footprint sits above the threshold, which reads to the user
        // as the keyboard repeatedly resetting itself. Re-arm only once pressure
        // has genuinely receded — back under `warningMB`, a full band below the
        // level that triggered the shed. `.warning` itself neither sheds nor
        // re-arms: it is the band the keyboard lands in right after shedding.
        if level == .normal {
            lastReliefLevel = .normal
        }
        guard level == .high || level == .critical, let handler = reliefHandler else { return }
        guard levelRank(level) > levelRank(lastReliefLevel) else { return }
        // Shedding records a milestone at every teardown step, and recording
        // samples the footprint again — which asks the same handler to shed
        // from inside itself. `lastReliefLevel` cannot stop that, because it is
        // only assigned once the handler returns. The nesting is unbounded and
        // overflowed the main thread's stack instead of relieving anything.
        guard !isRelieving else { return }
        #if DEBUG
        MemoryDeviceProbe.log("relief.fire.\(level.rawValue).\(stage)")
        #endif
        OSGDiag.log(
            "extMemory relief level=\(level.rawValue) stage=\(stage) \(OSGDiag.memoryTag())",
            category: "memory"
        )
        isRelieving = true
        let relieved = handler(level)
        isRelieving = false
        #if DEBUG
        MemoryDeviceProbe.log("relief.done.\(level.rawValue).relieved=\(relieved ? 1 : 0)")
        #endif
        guard relieved else {
            OSGDiag.log(
                "extMemory relief declined level=\(level.rawValue) stage=\(stage)",
                category: "memory"
            )
            return
        }
        lastReliefLevel = level
    }

    #if DEBUG
    /// Test seam: drives the threshold + hysteresis logic without needing the
    /// test process's real footprint to cross 48 MiB.
    static func simulateFootprintForTesting(_ footprintMB: Double, stage: String = "test") {
        requestReliefIfNeeded(
            level: KeyboardExtensionMemoryBudget.level(forPhysFootprintMB: footprintMB),
            stage: stage
        )
    }
    #endif

    private static func levelRank(_ level: KeyboardExtensionMemoryBudget.Level) -> Int {
        switch level {
        case .unavailable:
            return -1
        case .normal:
            return 0
        case .warning:
            return 1
        case .high:
            return 2
        case .critical:
            return 3
        }
    }
}

#if DEBUG
/// DEBUG-only device probe. `log stream --device` is unavailable on this Mac and
/// tool-hosted tests can't run on a physical device, so the extension's real
/// jetsam footprint can't be read from the console. Instead we append every
/// memory sample to an App Group file and pull it back to the Mac:
///
/// ```
/// xcrun devicectl device copy from --device <udid> \
///   --domain-type appGroupDataContainer \
///   --domain-identifier group.com.osgkeyboard.shared \
///   --source memory-probe.log --destination /tmp/
/// ```
///
/// Each line: `<uptime>s <tag> foot=<MB> rss=<MB>`. Temporary — remove once the
/// footprint budget is re-baselined.
public enum MemoryDeviceProbe {
    public static let fileName = "memory-probe.log"

    /// devicectl can only pull from the container's Library/Documents/tmp, so
    /// the probe lives under Library rather than the container root.
    private static var fileURL: URL? {
        guard let base = FileManager.default
            .containerURL(forSecurityApplicationGroupIdentifier: AppGroup.identifier)?
            .appendingPathComponent("Library", isDirectory: true)
        else { return nil }
        try? FileManager.default.createDirectory(at: base, withIntermediateDirectories: true)
        return base.appendingPathComponent(fileName, isDirectory: false)
    }

    /// Appends one sample. Cheap enough for a short DEBUG run; not for shipping.
    public static func log(_ tag: String) {
        guard let url = fileURL else { return }
        let snap = OSGDiag.memorySnapshot()
        let line = String(
            format: "%.2f %@ foot=%.1f rss=%.1f\n",
            ProcessInfo.processInfo.systemUptime,
            tag,
            snap.physFootprintMB,
            snap.rssMB
        )
        guard let data = line.data(using: .utf8) else { return }
        if let handle = try? FileHandle(forWritingTo: url) {
            handle.seekToEndOfFile()
            handle.write(data)
            try? handle.close()
        } else {
            try? data.write(to: url, options: .atomic)
        }
    }
}
#endif