Skip to content

Instantly share code, notes, and snippets.

@KaQuMiQ
Last active January 28, 2022 13:22
Show Gist options
  • Save KaQuMiQ/b567dc8aa0e6fcf769d4ab7dfd9dedf9 to your computer and use it in GitHub Desktop.
Save KaQuMiQ/b567dc8aa0e6fcf769d4ab7dfd9dedf9 to your computer and use it in GitHub Desktop.
import struct Foundation.UUID
import os
import libkern
public enum MemoryLeaks {
public static func listActiveMarkers() -> Array<Marker> {
let markersRegistryCopy: Dictionary<AnyHashable, WeakMarker> = markersRegistryLock
.withLock { Self.markersRegistry }
return markersRegistryCopy.values.compactMap(\.marker)
}
}
extension MemoryLeaks {
private static let logger: os.Logger = .init(
subsystem: "Memory",
category: "LeakCheck"
)
fileprivate static let markersRegistryLock: SpinLock = .init()
fileprivate static var markersRegistry: Dictionary<AnyHashable, WeakMarker> = .init()
fileprivate static func register(
marker: Marker
) {
let id: AnyHashable = marker.id
let weakMarker: WeakMarker = .init(marker: marker)
markersRegistryLock
.withLock {
Self.markersRegistry[id] = weakMarker
}
Self.logger.debug("NewMarker:\(id, privacy: .public)")
}
fileprivate static func unregister(
markerWithID markerID: AnyHashable
) {
markersRegistryLock
.withLock {
Self.markersRegistry[markerID] = .none
}
Self.logger.debug("RemovedMarker:\(markerID, privacy: .public)")
}
}
extension MemoryLeaks {
fileprivate final class WeakMarker {
fileprivate weak var marker: Marker?
fileprivate init(
marker: Marker
) {
self.marker = marker
}
}
fileprivate final class SpinLock {
private let ptr: UnsafeMutablePointer<atomic_flag>
fileprivate init() {
self.ptr = .allocate(capacity: 1)
self.ptr.pointee = atomic_flag()
}
deinit {
self.ptr.deinitialize(count: 1)
self.ptr.deallocate()
}
fileprivate func lock() {
while !atomic_flag_test_and_set(self.ptr) {
// keep spinning
}
}
fileprivate func unlock() {
atomic_flag_clear(self.ptr)
}
fileprivate func withLock<R>(
_ execute: () throws -> R
) rethrows -> R {
self.lock()
defer { self.unlock() }
return try execute()
}
}
}
extension MemoryLeaks {
public final class Marker {
fileprivate let id: String
private let customID: String
private let uuid: String
private let timestamp: Int
private let file: StaticString
private let line: UInt
fileprivate init(
customID: String,
file: StaticString,
line: UInt
) {
let uuidString: String = UUID().uuidString
let timestamp: Int = time(nil)
self.id = "\(file):\(line)[\(uuidString)]<\(timestamp)>-\(customID)"
self.customID = customID
self.uuid = uuidString
self.timestamp = timestamp
self.file = file
self.line = line
MemoryLeaks
.register(marker: self)
}
deinit {
MemoryLeaks
.unregister(markerWithID: self.id)
}
fileprivate func wrappedReturn<R>(
_ value: @autoclosure () throws -> R
) rethrows -> R {
try value()
}
}
}
extension MemoryLeaks.Marker: CustomStringConvertible {
public var description: String {
"MemoryLeaksMarker:\(self.id)"
}
}
extension MemoryLeaks.Marker: CustomDebugStringConvertible {
public var debugDescription: String {
"MemoryLeaksMarker:\(self.id)"
}
}
extension MemoryLeaks.Marker: CustomReflectable {
public var customMirror: Mirror {
.init(
self,
children: ["id": self.id],
displayStyle: .tuple,
ancestorRepresentation: .suppressed
)
}
}
extension MemoryLeaks {
public static func mark<R>(
_ function: @escaping () -> R,
customID: String = "",
file: StaticString = #fileID,
line: UInt = #line
) -> () -> R {
#if DEBUG
let marker: MemoryLeaks.Marker = .init(
customID: customID,
file: file,
line: line
)
return {
marker.wrappedReturn(function())
}
#else
return function
#endif
}
public static func markThrowing<R>(
_ function: @escaping () throws -> R,
customID: String = "",
file: StaticString = #fileID,
line: UInt = #line
) -> () throws -> R {
#if DEBUG
let marker: MemoryLeaks.Marker = .init(
customID: customID,
file: file,
line: line
)
return {
try marker.wrappedReturn(try function())
}
#else
return function
#endif
}
public static func mark<R, A1>(
_ function: @escaping (A1) -> R,
customID: String = "",
file: StaticString = #fileID,
line: UInt = #line
) -> (A1) -> R {
#if DEBUG
let marker: MemoryLeaks.Marker = .init(
customID: customID,
file: file,
line: line
)
return { a1 in
marker.wrappedReturn(function(a1))
}
#else
return function
#endif
}
public static func mark<R, A1, A2>(
_ function: @escaping (A1, A2) -> R,
customID: String = "",
file: StaticString = #fileID,
line: UInt = #line
) -> (A1, A2) -> R {
#if DEBUG
let marker: MemoryLeaks.Marker = .init(
customID: customID,
file: file,
line: line
)
return { a1, a2 in
marker.wrappedReturn(function(a1, a2))
}
#else
return function
#endif
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment