MemoryGrowthLimit.swift

// Keep polling the memory usage (through number of wired pages) every millisecond to discover if usage
// stays below some limit. 
//
//     let limit = MemoryGrowthLimit(megabytes: 5)
//     hardWorkingMemoryConsumingMethodCall()
//     XCTAssertFalse(limit.exceeded)
//
class MemoryGrowthLimit {
    private let pollingInterval = 0.001
    
    private var timer: Timer!
    private let maximumBytes: UInt64
    private let pageSize: UInt64
    private var _exceeded = false

    public var exceeded: Bool {
        return _exceeded
    }
    
    // megabytes is how much memory usage is allowed to grow from the point where init is called.
    init(megabytes: Double) {
        pageSize = MemoryGrowthLimit.pageSizeBytes()
        let startUsage = pageSize * MemoryGrowthLimit.wiredMemoryPages()
        maximumBytes = startUsage + UInt64(1024.0 * 1024.0 * megabytes)
        timer = Timer.scheduledTimer(timeInterval: pollingInterval, target: self, selector: #selector(poll),
                                     userInfo: nil, repeats: true)
    }
    
    @objc private func poll() {
        let usage = MemoryGrowthLimit.wiredMemoryPages() * pageSize
        if usage > maximumBytes {
            _exceeded = true
            timer.invalidate()
        }
    }
    
    private static func wiredMemoryPages() -> UInt64 {
        let count = MemoryLayout<vm_statistics_data_t>.stride / MemoryLayout<integer_t>.stride
        var infoCount = mach_msg_type_number_t(count)
        var info = [integer_t](repeating: 0, count: count)
        
        let kernReturn = host_statistics(mach_host_self(), HOST_VM_INFO, &info, &infoCount)
        XCTAssertEqual(KERN_SUCCESS, kernReturn)
        
        let wiredPages = info.withUnsafeBytes({ (pointer: UnsafeRawBufferPointer) -> UInt64 in
            let vmStats = pointer.bindMemory(to: vm_statistics_data_t.self).baseAddress!.pointee
            return UInt64(vmStats.wire_count)
        })
        return wiredPages
    }
    
    private static func pageSizeBytes() -> UInt64 {
        var size = vm_size_t()
        let kernReturn = host_page_size(mach_host_self(), &size)
        XCTAssertEqual(KERN_SUCCESS, kernReturn)
        
        return UInt64(size)
    }
}