[Backtracing] Add a caching wrapper for MemoryReader.

Currently we read many small chunks from the process we're backtracing.
If it so happens that it's the local process, that isn't really a big
problem, but if it's a remote process, especially on Linux where we have
to use the memory server, it's probably a little slow.

Fix by adding a caching layer.

rdar://117681625

Author: al45tair

stdlib/public/Backtracing/ArrayImageSource.swift

@@ -16,8 +16,6 @@
 
 import Swift
 
-@_implementationOnly import OS.Libc
-
 enum ArrayImageSourceError: Error {
   case outOfBoundsRead(UInt64, UInt64)
 }
@@ -35,16 +33,16 @@ struct ArrayImageSource<T>: ImageSource {
     return Bounds(base: 0, size: Size(array.count * MemoryLayout<T>.stride))
   }
 
-  public func fetch<U>(from addr: Address,
-                       into buffer: UnsafeMutableBufferPointer<U>) throws {
+  public func fetch(from addr: Address,
+                    into buffer: UnsafeMutableRawBufferPointer) throws {
     try array.withUnsafeBytes{
       let size = Size($0.count)
-      let requested = Size(buffer.count * MemoryLayout<U>.stride)
+      let requested = Size(buffer.count)
       if addr > size || requested > size - addr {
         throw ArrayImageSourceError.outOfBoundsRead(addr, requested)
       }
 
-      memcpy(buffer.baseAddress!, $0.baseAddress! + Int(addr), Int(requested))
+      buffer.copyBytes(from: $0[Int(addr)...])
     }
   }
 }

stdlib/public/Backtracing/CMakeLists.txt

@@ -25,6 +25,7 @@ set(BACKTRACING_SOURCES
   Backtrace.swift
   BacktraceFormatter.swift
   ByteSwapping.swift
+  CachingMemoryReader.swift
   Context.swift
   Compression.swift
   CoreSymbolication.swift

stdlib/public/Backtracing/CachingMemoryReader.swift

@@ -0,0 +1,84 @@
+//===--- CachingMemoryReader.swift ----------------------------*- swift -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2023 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+//
+//  Wraps a MemoryReader in a layer that caches memory pages.
+//
+//===----------------------------------------------------------------------===//
+
+import Swift
+
+// The size of the pages in the page cache (must be a power of 2)
+fileprivate let pageSize = 4096
+
+fileprivate let pageMask = pageSize - 1
+
+// The largest chunk we will try to cache data for
+fileprivate let maxCachedSize = pageSize * 8
+
+@_spi(MemoryReaders)
+public class CachingMemoryReader<T: MemoryReader>: MemoryReader {
+  private var reader: T
+  private var cache: [Address:UnsafeRawBufferPointer]
+
+  public init(for reader: T) {
+    self.reader = reader
+    self.cache = [:]
+  }
+
+  deinit {
+    for (_, page) in cache {
+      page.deallocate()
+    }
+  }
+
+  private func getPage(at address: Address) throws -> UnsafeRawBufferPointer {
+    if let page = cache[address] {
+      return page
+    }
+
+    let page = UnsafeMutableRawBufferPointer.allocate(byteCount: pageSize,
+                                                      alignment: pageSize)
+    try reader.fetch(from: address, into: page)
+
+    let result = UnsafeRawBufferPointer(page)
+
+    cache[address] = result
+
+    return result
+  }
+
+  public func fetch(from address: Address,
+                    into buffer: UnsafeMutableRawBufferPointer) throws {
+    guard buffer.count <= maxCachedSize else {
+      try reader.fetch(from: address, into: buffer)
+      return
+    }
+
+    var pageAddress = address & ~Address(pageMask)
+    var done = 0
+    var offset = Int(address - pageAddress)
+    var remaining = buffer.count
+
+    while remaining > 0 {
+      let page = try getPage(at: pageAddress)
+      let maxBytes = pageSize - offset
+      let chunk = min(remaining, maxBytes)
+
+      buffer[done..<done+chunk].copyBytes(from: page[offset...])
+
+      offset = 0
+      done += chunk
+      remaining -= chunk
+      pageAddress += Address(pageSize)
+    }
+  }
+}

stdlib/public/Backtracing/Compression.swift

@@ -427,9 +427,9 @@ internal struct ElfCompressedImageSource<Traits: ElfTraits>: ImageSource {
     }
   }
 
-  public func fetch<T>(from addr: Address,
-                       into buffer: UnsafeMutableBufferPointer<T>) throws {
-    let toFetch = buffer.count * MemoryLayout<T>.stride
+  public func fetch(from addr: Address,
+                    into buffer: UnsafeMutableRawBufferPointer) throws {
+    let toFetch = buffer.count
     if addr < 0 || addr > data.count || data.count - Int(addr) < toFetch {
       throw CompressedImageSourceError.outOfRangeFetch(addr, toFetch)
     }
@@ -474,9 +474,9 @@ internal struct ElfGNUCompressedImageSource: ImageSource {
                           uncompressedSize: uncompressedSize)
   }
 
-  public func fetch<T>(from addr: Address,
-                       into buffer: UnsafeMutableBufferPointer<T>) throws {
-    let toFetch = buffer.count * MemoryLayout<T>.stride
+  public func fetch(from addr: Address,
+                    into buffer: UnsafeMutableRawBufferPointer) throws {
+    let toFetch = buffer.count
     if addr < 0 || addr > data.count || data.count - Int(addr) < toFetch {
       throw CompressedImageSourceError.outOfRangeFetch(addr, toFetch)
     }
@@ -509,9 +509,9 @@ internal struct LZMACompressedImageSource: ImageSource {
                           dataBounds: bounds)
   }
 
-  public func fetch<T>(from addr: Address,
-                       into buffer: UnsafeMutableBufferPointer<T>) throws {
-    let toFetch = buffer.count * MemoryLayout<T>.stride
+  public func fetch(from addr: Address,
+                    into buffer: UnsafeMutableRawBufferPointer) throws {
+    let toFetch = buffer.count
     if addr < 0 || addr > data.count || data.count - Int(addr) < toFetch {
       throw CompressedImageSourceError.outOfRangeFetch(addr, toFetch)
     }

stdlib/public/Backtracing/FileImageSource.swift

@@ -58,8 +58,8 @@ class FileImageSource: ImageSource {
            _mapping.count)
   }
 
-  public func fetch<T>(from addr: Address,
-                       into buffer: UnsafeMutableBufferPointer<T>) throws {
+  public func fetch(from addr: Address,
+                    into buffer: UnsafeMutableRawBufferPointer) throws {
     let start = Int(addr)
     _ = try buffer.withMemoryRebound(to: UInt8.self) { byteBuf in
       guard _mapping.indices.contains(start) else {
@@ -69,7 +69,6 @@ class FileImageSource: ImageSource {
       guard slice.count >= byteBuf.count else {
         throw FileImageSourceError.outOfRangeRead
       }
-      byteBuf.update(from: slice)
     }
   }
 }

stdlib/public/Backtracing/ImageSource.swift

@@ -60,6 +60,11 @@ struct ImageSourceCursor {
     self.pos = offset
   }
 
+  public mutating func read(into buffer: UnsafeMutableRawBufferPointer) throws {
+    try source.fetch(from: pos, into: buffer)
+    pos += UInt64(buffer.count)
+  }
+  
   public mutating func read<T>(into buffer: UnsafeMutableBufferPointer<T>) throws {
     try source.fetch(from: pos, into: buffer)
     pos += UInt64(MemoryLayout<T>.stride * buffer.count)
@@ -138,9 +143,9 @@ struct SubImageSource<S: ImageSource>: ImageSource {
     return parent.isMappedImage
   }
 
-  public func fetch<T>(from addr: Address,
-                       into buffer: UnsafeMutableBufferPointer<T>) throws {
-    let toFetch = buffer.count * MemoryLayout<T>.stride
+  public func fetch(from addr: Address,
+                    into buffer: UnsafeMutableRawBufferPointer) throws {
+    let toFetch = buffer.count
     if addr < 0 || addr > length {
       throw SubImageSourceError.outOfRangeFetch(UInt64(addr), toFetch)
     }

stdlib/public/Backtracing/MemoryImageSource.swift

@@ -28,8 +28,8 @@ class MemoryImageSource<M: MemoryReader>: ImageSource {
     self.reader = reader
   }
 
-  public func fetch<T>(from addr: Address,
-                       into buffer: UnsafeMutableBufferPointer<T>) throws {
+  public func fetch(from addr: Address,
+                    into buffer: UnsafeMutableRawBufferPointer) throws {
     try reader.fetch(from: addr, into: buffer)
   }
 }

stdlib/public/Backtracing/MemoryReader.swift

@@ -30,6 +30,11 @@ import Swift
   typealias Address = UInt64
   typealias Size = UInt64
 
+  /// Fill the specified buffer with data from the specified location in
+  /// the source.
+  func fetch(from address: Address,
+             into buffer: UnsafeMutableRawBufferPointer) throws
+
   /// Fill the specified buffer with data from the specified location in
   /// the source.
   func fetch<T>(from address: Address,
@@ -51,6 +56,11 @@ import Swift
 
 extension MemoryReader {
 
+  public func fetch<T>(from address: Address,
+                       into buffer: UnsafeMutableBufferPointer<T>) throws {
+    try fetch(from: address, into: UnsafeMutableRawBufferPointer(buffer))
+  }
+
   public func fetch<T>(from addr: Address,
                        into pointer: UnsafeMutablePointer<T>) throws {
     try fetch(from: addr,
@@ -96,10 +106,12 @@ extension MemoryReader {
 @_spi(MemoryReaders) public struct UnsafeLocalMemoryReader: MemoryReader {
   public init() {}
 
-  public func fetch<T>(from address: Address,
-                       into buffer: UnsafeMutableBufferPointer<T>) throws {
-    buffer.baseAddress!.update(from: UnsafePointer<T>(bitPattern: UInt(address))!,
-                               count: buffer.count)
+  public func fetch(from address: Address,
+                    into buffer: UnsafeMutableRawBufferPointer) throws {
+    buffer.baseAddress!.copyMemory(
+      from: UnsafeRawPointer(bitPattern: UInt(address))!,
+      byteCount: buffer.count
+    )
   }
 }
 
@@ -116,9 +128,9 @@ extension MemoryReader {
     self.task = task as! task_t
   }
 
-  public func fetch<T>(from address: Address,
-                       into buffer: UnsafeMutableBufferPointer<T>) throws {
-    let size = UInt64(MemoryLayout<T>.stride * buffer.count)
+  public func fetch(from address: Address,
+                    into buffer: UnsafeMutableRawBufferPointer) throws {
+    let size = buffer.count
     var sizeOut = UInt64(0)
     let result = mach_vm_read_overwrite(task,
                                         UInt64(address),
@@ -138,8 +150,8 @@ extension MemoryReader {
   public typealias Address = UInt64
   public typealias Size = UInt64
 
-  public func fetch<T>(from address: Address,
-                       into buffer: UnsafeMutableBufferPointer<T>) throws {
+  public func fetch(from address: Address,
+                    into buffer: UnsafeMutableRawBufferPointer) throws {
     let reader = RemoteMemoryReader(task: mach_task_self())
     return try reader.fetch(from: address, into: buffer)
   }
@@ -221,34 +233,31 @@ extension MemoryReader {
     return response
   }
 
-  public func fetch<T>(from addr: Address,
-                       into buffer: UnsafeMutableBufferPointer<T>) throws {
-    try buffer.withMemoryRebound(to: UInt8.self) {
-      let bytes = UnsafeMutableRawBufferPointer($0)
-      try sendRequest(for: Size(bytes.count), from: addr)
+  public func fetch(from addr: Address,
+                    into buffer: UnsafeMutableRawBufferPointer) throws {
+    try sendRequest(for: Size(buffer.count), from: addr)
 
-      var done = 0
-      while done < bytes.count {
-        let reply = try receiveReply()
-
-        if reply.len < 0 {
-          throw MemserverError(message: "Unreadable at \(hex(addr))")
-        }
+    var done = 0
+    while done < buffer.count {
+      let reply = try receiveReply()
 
-        if done + Int(reply.len) > bytes.count {
-          throw MemserverError(message: "Overrun at \(hex(addr)) trying to read \(bytes.count) bytes")
-        }
+      if reply.len < 0 {
+        throw MemserverError(message: "Unreadable at \(hex(addr))")
+      }
 
-        let ret = try safeRead(fd,
-                               UnsafeMutableRawBufferPointer(
-                                 rebasing: bytes[done..<done+Int(reply.len)]))
+      if done + Int(reply.len) > bytes.count {
+        throw MemserverError(message: "Overrun at \(hex(addr)) trying to read \(bytes.count) bytes")
+      }
 
-        if ret != reply.len {
-          throw MemserverError(message: "Channel closed prematurely")
-        }
+      let ret = try safeRead(fd,
+                             UnsafeMutableRawBufferPointer(
+                               rebasing: buffer[done..<done+Int(reply.len)]))
 
-        done += Int(reply.len)
+      if ret != reply.len {
+        throw MemserverError(message: "Channel closed prematurely")
       }
+
+      done += Int(reply.len)
     }
   }
 }
@@ -260,9 +269,9 @@ extension MemoryReader {
     self.pid = pid as! pid_t
   }
 
-  public func fetch<T>(from address: Address,
-                       into buffer: UnsafeMutableBufferPointer<T>) throws {
-    let size = size_t(MemoryLayout<T>.stride * buffer.count)
+  public func fetch(from address: Address,
+                    into buffer: UnsafeMutableRawBufferPointer) throws {
+    let size = buffer.count
     var fromIOVec = iovec(iov_base: UnsafeMutableRawPointer(
                             bitPattern: UInt(address)),
                           iov_len: size)
@@ -281,8 +290,8 @@ extension MemoryReader {
     reader = RemoteMemoryReader(pid: getpid())
   }
 
-  public func fetch<T>(from address: Address,
-                       into buffer: UnsafeMutableBufferPointer<T>) throws {
+  public func fetch(from address: Address,
+                    into buffer: UnsafeMutableRawBufferPointer) throws {
     return try reader.fetch(from: address, into: buffer)
   }
 }

stdlib/public/libexec/swift-backtrace/TargetLinux.swift

@@ -77,7 +77,7 @@ class Target {
     }
   }
 
-  var reader: MemserverMemoryReader
+  var reader: CachingMemoryReader<MemserverMemoryReader>
 
   // Get the name of a process
   private static func getProcessName(pid: pid_t) -> String {
@@ -110,7 +110,7 @@ class Target {
     let memserverFd: CInt = 4
 
     pid = getppid()
-    reader = MemserverMemoryReader(fd: memserverFd)
+    reader = CachingMemoryReader(for: MemserverMemoryReader(fd: memserverFd))
     name = Self.getProcessName(pid: pid)
 
     let crashInfo: CrashInfo

stdlib/public/libexec/swift-backtrace/TargetMacOS.swift

@@ -98,7 +98,7 @@ class Target {
     }
   }
 
-  var reader: RemoteMemoryReader
+  var reader: CachingMemoryReader<RemoteMemoryReader>
 
   var mcontext: MContext
 
@@ -158,7 +158,7 @@ class Target {
 
     task = parentTask
 
-    reader = RemoteMemoryReader(task: task_t(task))
+    reader = CachingMemoryReader(for: RemoteMemoryReader(task: task_t(task)))
 
     name = Self.getProcessName(pid: pid)