fix OOM crash on long audio files (#93)

The end-of-loop backfill in WaveformAnalyzer.extract padded the input
sample buffer with `(target - leftSamples.count) * samplesPerPixel * 2`
bytes of zeros when the read produced fewer samples than the target -
trivially gigabytes for multi-hour files when the reader ended early
(e.g. after backgrounding). This crashed on memory exhaustion.

Pad the output array with silence-equivalent dB values instead, so the
allocation is O(targetSampleCount) regardless of audio duration. Skip
the pad entirely when the reader didn't reach .completed - the result
is discarded in that case anyway. Wrap the per-chunk read body in an
autoreleasepool so CMSampleBuffer references drain promptly during long
reads instead of pinning gigabytes until the loop exits. Guard the
unconditional startReading() call so a non-.unknown reader fails
gracefully rather than via uncatchable ObjC exception.

Adds a test target with a reproducer that exercises the buggy condition
by inflating samplesPerPixel against a short fixture, plus a smoke test
and a defensive crash-resilience test.

Author: dmrschmidt

Package.swift

@@ -28,5 +28,9 @@ let package = Package(
             name: "DSWaveformImageViews",
             dependencies: ["DSWaveformImage"]
         ),
+        .testTarget(
+            name: "DSWaveformImageTests",
+            dependencies: ["DSWaveformImage"]
+        ),
     ]
 )

Sources/DSWaveformImage/WaveformAnalyzer.swift

@@ -96,7 +96,7 @@ public struct WaveformAnalyzer: Sendable {
 
 // MARK: - Private
 
-fileprivate extension WaveformAnalyzer {
+internal extension WaveformAnalyzer {
     func waveformSamples(
             track audioAssetTrack: AVAssetTrack,
             reader assetReader: AVAssetReader,
@@ -147,58 +147,73 @@ fileprivate extension WaveformAnalyzer {
         let samplesPerPixel = max(1, totalSamples / targetSampleCount)
         let samplesPerFFT = 4096 // ~100ms at 44.1kHz, rounded to closest pow(2) for FFT
 
+        // `startReading()` throws an uncatchable ObjC exception if the reader isn't in `.unknown`
+        // (e.g. already cancelled or failed). Normal callers always pass a fresh reader, but bail
+        // gracefully if that contract is violated so we surface as `readerError` rather than crash.
+        guard assetReader.status == .unknown else {
+            return WaveformAnalysis(amplitudes: [], fft: outputFFT)
+        }
         assetReader.startReading()
         while assetReader.status == .reading {
-            let trackOutput = assetReader.outputs.first!
-
-            guard let nextSampleBuffer = trackOutput.copyNextSampleBuffer(),
-                let blockBuffer = CMSampleBufferGetDataBuffer(nextSampleBuffer) else {
-                    break
-            }
+            // CMSampleBuffer is a Core Foundation type that lives in the autorelease pool.
+            // Without an explicit drain per iteration, long files iterate thousands of times and
+            // can keep gigabytes of buffer memory pinned until the loop exits.
+            let continueReading = autoreleasepool { () -> Bool in
+                let trackOutput = assetReader.outputs.first!
+
+                guard let nextSampleBuffer = trackOutput.copyNextSampleBuffer(),
+                    let blockBuffer = CMSampleBufferGetDataBuffer(nextSampleBuffer) else {
+                        return false
+                }
 
-            var readBufferLength = 0
-            var readBufferPointer: UnsafeMutablePointer<Int8>? = nil
-            CMBlockBufferGetDataPointer(blockBuffer, atOffset: 0, lengthAtOffsetOut: &readBufferLength, totalLengthOut: nil, dataPointerOut: &readBufferPointer)
-            sampleBuffer.append(UnsafeBufferPointer(start: readBufferPointer, count: readBufferLength))
-            if fftBands != nil {
-                // don't append data to this buffer unless we're going to use it.
-                sampleBufferFFT.append(UnsafeBufferPointer(start: readBufferPointer, count: readBufferLength))
-            }
-            CMSampleBufferInvalidate(nextSampleBuffer)
+                var readBufferLength = 0
+                var readBufferPointer: UnsafeMutablePointer<Int8>? = nil
+                CMBlockBufferGetDataPointer(blockBuffer, atOffset: 0, lengthAtOffsetOut: &readBufferLength, totalLengthOut: nil, dataPointerOut: &readBufferPointer)
+                sampleBuffer.append(UnsafeBufferPointer(start: readBufferPointer, count: readBufferLength))
+                if fftBands != nil {
+                    // don't append data to this buffer unless we're going to use it.
+                    sampleBufferFFT.append(UnsafeBufferPointer(start: readBufferPointer, count: readBufferLength))
+                }
+                CMSampleBufferInvalidate(nextSampleBuffer)
 
-            let result = process(sampleBuffer, from: assetReader, downsampleTo: samplesPerPixel, channelSelection: channelSelection)
-            leftSamples += result.left
-            rightSamples += result.right
+                let result = process(sampleBuffer, from: assetReader, downsampleTo: samplesPerPixel, channelSelection: channelSelection)
+                leftSamples += result.left
+                rightSamples += result.right
 
-            if result.bytesConsumed > 0 {
-                sampleBuffer.removeFirst(result.bytesConsumed)
+                if result.bytesConsumed > 0 {
+                    sampleBuffer.removeFirst(result.bytesConsumed)
 
-                // this takes care of a memory leak where Memory continues to increase even though it should clear after calling .removeFirst(…) above.
-                sampleBuffer = Data(sampleBuffer)
-            }
+                    // this takes care of a memory leak where Memory continues to increase even though it should clear after calling .removeFirst(…) above.
+                    sampleBuffer = Data(sampleBuffer)
+                }
 
-            if let fftBands = fftBands, sampleBufferFFT.count / MemoryLayout<Int16>.size >= samplesPerFFT {
-                let processedFFTs = process(sampleBufferFFT, samplesPerFFT: samplesPerFFT, fftBands: fftBands)
-                sampleBufferFFT.removeFirst(processedFFTs.count * samplesPerFFT * MemoryLayout<Int16>.size)
-                outputFFT? += processedFFTs
+                if let fftBands = fftBands, sampleBufferFFT.count / MemoryLayout<Int16>.size >= samplesPerFFT {
+                    let processedFFTs = process(sampleBufferFFT, samplesPerFFT: samplesPerFFT, fftBands: fftBands)
+                    sampleBufferFFT.removeFirst(processedFFTs.count * samplesPerFFT * MemoryLayout<Int16>.size)
+                    outputFFT? += processedFFTs
+                }
+                return true
             }
+            if !continueReading { break }
         }
 
-        // if we don't have enough pixels yet,
-        // process leftover samples with padding (to reach multiple of samplesPerPixel for vDSP_desamp)
-        if leftSamples.count < targetSampleCount {
-            // each output sample for a single rendered "channel" consumes `samplesPerPixel * inputUnitsPerOutputSample`
-            // Int16s from the interleaved buffer.
-            let channelCount = channelInfo(from: assetReader)?.channelCount ?? 1
-            let inputUnitsPerOutputSample = (channelSelection == .merged) ? 1 : channelCount
-            let missingSampleCount = (targetSampleCount - leftSamples.count) * samplesPerPixel * inputUnitsPerOutputSample
-            let backfillPaddingSampleCount = max(0, missingSampleCount - (sampleBuffer.count / MemoryLayout<Int16>.size))
-            let backfillPaddingByteCount = backfillPaddingSampleCount * MemoryLayout<Int16>.size
-            let backfillPaddingSamples = [UInt8](repeating: 0, count: backfillPaddingByteCount)
-            sampleBuffer.append(backfillPaddingSamples, count: backfillPaddingByteCount)
-            let result = process(sampleBuffer, from: assetReader, downsampleTo: samplesPerPixel, channelSelection: channelSelection)
-            leftSamples += result.left
-            rightSamples += result.right
+        // Pad the *output* with silence-equivalent dB values when the read produced fewer samples
+        // than the target — e.g. a short tail or a reader that ended early (failed/cancelled after
+        // backgrounding). These become 1.0 (silence) after `normalize`. Allocation is
+        // O(targetSampleCount), independent of audio duration — the previous implementation padded
+        // the *input* buffer with up to `target × samplesPerPixel × 2` bytes of zeros, which
+        // crashed on multi-hour files (issue #93). We only pad on a clean read; a non-`.completed`
+        // status means `waveformSamples` will throw and the result is discarded anyway, so skip the
+        // wasted work.
+        if assetReader.status == .completed {
+            if leftSamples.count < targetSampleCount {
+                let missing = targetSampleCount - leftSamples.count
+                leftSamples.append(contentsOf: repeatElement(noiseFloorDecibelCutoff, count: missing))
+            }
+            if isStereo, rightSamples.count < targetSampleCount {
+                let missing = targetSampleCount - rightSamples.count
+                rightSamples.append(contentsOf: repeatElement(noiseFloorDecibelCutoff, count: missing))
+            }
         }
 
         let amplitudes: [Float]

Tests/DSWaveformImageTests/TestSupport.swift

@@ -0,0 +1,89 @@
+import AVFoundation
+import Darwin
+import Foundation
+
+/// Writes a silent mono 16-bit PCM WAV file to a temporary location and returns the URL.
+/// Bytes-on-disk = `durationSeconds * sampleRate * 2`. A 10-minute file is ~53 MB.
+func makeSilentAudioFile(durationSeconds: Double, sampleRate: Double = 44_100) throws -> URL {
+    let url = FileManager.default.temporaryDirectory
+        .appendingPathComponent("dswaveform-test-\(UUID().uuidString).wav")
+
+    let settings: [String: Any] = [
+        AVFormatIDKey: kAudioFormatLinearPCM,
+        AVSampleRateKey: sampleRate,
+        AVNumberOfChannelsKey: 1,
+        AVLinearPCMBitDepthKey: 16,
+        AVLinearPCMIsBigEndianKey: false,
+        AVLinearPCMIsFloatKey: false,
+        AVLinearPCMIsNonInterleaved: false,
+    ]
+    let file = try AVAudioFile(forWriting: url, settings: settings)
+
+    // Write in chunks so we don't allocate the whole audio file as a single PCM buffer.
+    let format = file.processingFormat
+    let chunkFrames: AVAudioFrameCount = 44_100 // 1 s of audio at 44.1 kHz
+    let totalFrames = AVAudioFrameCount(durationSeconds * sampleRate)
+    var written: AVAudioFrameCount = 0
+    while written < totalFrames {
+        let frames = min(chunkFrames, totalFrames - written)
+        guard let buffer = AVAudioPCMBuffer(pcmFormat: format, frameCapacity: frames) else {
+            throw NSError(domain: "TestSupport", code: 1, userInfo: [NSLocalizedDescriptionKey: "buffer alloc failed"])
+        }
+        buffer.frameLength = frames
+        // The buffer is zero-initialized — that's our silence.
+        try file.write(from: buffer)
+        written += frames
+    }
+    return url
+}
+
+/// Current process physical memory footprint in bytes, via `task_vm_info`.
+func currentPhysFootprint() -> Int64 {
+    var info = task_vm_info_data_t()
+    var count = mach_msg_type_number_t(MemoryLayout<task_vm_info_data_t>.size / MemoryLayout<integer_t>.size)
+    let kr = withUnsafeMutablePointer(to: &info) {
+        $0.withMemoryRebound(to: integer_t.self, capacity: Int(count)) {
+            task_info(mach_task_self_, task_flavor_t(TASK_VM_INFO), $0, &count)
+        }
+    }
+    return kr == KERN_SUCCESS ? Int64(info.phys_footprint) : 0
+}
+
+/// Runs `operation` while polling the process's physical memory footprint, returning the operation's
+/// result along with the peak delta over the baseline measured immediately before the call.
+func measurePeakMemoryDelta<T>(
+    pollIntervalNanos: UInt64 = 5_000_000,
+    during operation: () async throws -> T
+) async throws -> (result: T?, error: Error?, peakDeltaBytes: Int64) {
+    let baseline = currentPhysFootprint()
+    let peakBox = PeakBox()
+    let monitor = Task { [peakBox] in
+        while !Task.isCancelled {
+            let delta = currentPhysFootprint() - baseline
+            await peakBox.update(delta)
+            try? await Task.sleep(nanoseconds: pollIntervalNanos)
+        }
+    }
+
+    let result: T?
+    let caught: Error?
+    do {
+        result = try await operation()
+        caught = nil
+    } catch {
+        result = nil
+        caught = error
+    }
+
+    monitor.cancel()
+    _ = await monitor.value
+    // Pick up any final spike that may have landed between the last poll and cancel.
+    await peakBox.update(currentPhysFootprint() - baseline)
+    let peak = await peakBox.peak
+    return (result, caught, peak)
+}
+
+private actor PeakBox {
+    private(set) var peak: Int64 = 0
+    func update(_ value: Int64) { peak = max(peak, value) }
+}

Tests/DSWaveformImageTests/WaveformAnalyzerTests.swift

@@ -0,0 +1,98 @@
+import AVFoundation
+import XCTest
+@testable import DSWaveformImage
+
+final class WaveformAnalyzerTests: XCTestCase {
+    /// Defensive: `extract` previously called `startReading()` unconditionally, which throws an
+    /// uncatchable ObjC exception when the reader isn't in `.unknown`. Normal callers always pass
+    /// a fresh reader, but the guard prevents a future caller from triggering a hard crash.
+    func testExtractDoesNotCrashOnCancelledReader() async throws {
+        let url = try makeSilentAudioFile(durationSeconds: 1)
+        defer { try? FileManager.default.removeItem(at: url) }
+
+        let asset = AVURLAsset(url: url)
+        let tracks = try await asset.loadTracks(withMediaType: .audio)
+        let track = try XCTUnwrap(tracks.first)
+        let reader = try AVAssetReader(asset: asset)
+        let trackOutput = AVAssetReaderTrackOutput(track: track, outputSettings: [
+            AVFormatIDKey: kAudioFormatLinearPCM,
+            AVLinearPCMBitDepthKey: 16,
+            AVLinearPCMIsBigEndianKey: false,
+            AVLinearPCMIsFloatKey: false,
+            AVLinearPCMIsNonInterleaved: false,
+        ])
+        reader.add(trackOutput)
+        reader.cancelReading()
+
+        let analysis = WaveformAnalyzer().extract(44_100, downsampledTo: 100, from: reader, channelSelection: .merged, fftBands: nil)
+        XCTAssertEqual(analysis.amplitudes, [], "extract on a non-fresh reader should bail with empty amplitudes")
+    }
+
+    /// Sanity check: a short file analyzes cleanly via the public surface.
+    func testShortFileAnalyzesCleanly() async throws {
+        let url = try makeSilentAudioFile(durationSeconds: 2)
+        defer { try? FileManager.default.removeItem(at: url) }
+
+        let samples = try await WaveformAnalyzer().samples(fromAudioAt: url, count: 100)
+        XCTAssertEqual(samples.count, 100)
+    }
+
+    /// Reproduces issue #93: when the read loop produces far fewer output samples than the target
+    /// (e.g. the reader fails or is cancelled mid-way after backgrounding), the analyzer's
+    /// end-of-loop backfill path padded the *input* buffer with
+    /// `(targetSampleCount - leftSamples.count) * samplesPerPixel * 2` bytes of zeros — easily
+    /// gigabytes for long files when `samplesPerPixel` is large.
+    ///
+    /// We simulate the condition without needing a long file or a cancelled reader: pass a
+    /// deliberately inflated `totalSamples` to `extract` while reading from a short file. That
+    /// makes `samplesPerPixel` huge (`totalSamples / targetSampleCount`), and the actual read
+    /// only produces a handful of output samples — exactly the bug's preconditions.
+    ///
+    /// Post-fix expectation: backfill pads the *output* with silence-equivalent floats and
+    /// doesn't allocate anywhere near the input-byte-equivalent footprint.
+    func testBackfillDoesNotAllocateInputBufferWhenReadProducesFewSamples() async throws {
+        let url = try makeSilentAudioFile(durationSeconds: 2) // 88_200 real samples
+        defer { try? FileManager.default.removeItem(at: url) }
+
+        let asset = AVURLAsset(url: url)
+        let tracks = try await asset.loadTracks(withMediaType: .audio)
+        let track = try XCTUnwrap(tracks.first)
+        let reader = try AVAssetReader(asset: asset)
+        let trackOutput = AVAssetReaderTrackOutput(track: track, outputSettings: [
+            AVFormatIDKey: kAudioFormatLinearPCM,
+            AVLinearPCMBitDepthKey: 16,
+            AVLinearPCMIsBigEndianKey: false,
+            AVLinearPCMIsFloatKey: false,
+            AVLinearPCMIsNonInterleaved: false,
+        ])
+        reader.add(trackOutput)
+
+        let targetSampleCount = 400
+        // Lie about totalSamples to inflate samplesPerPixel. With samplesPerPixel = 25_000 and
+        // only 88_200 real samples to consume, the read produces ~3 leftSamples — well under the
+        // 400 target. The buggy backfill then allocates ~(397 × 25_000 × 2) = 19.85 MB of zeros
+        // and a 39.7 MB Float copy. The fixed backfill allocates target-count Floats (~1.6 KB).
+        let fakeTotalSamples = 10_000_000
+        let buggyAllocBytes: Int64 = 19_850_000 + 39_700_000
+
+        let (analysis, _, peakDelta) = try await measurePeakMemoryDelta {
+            WaveformAnalyzer().extract(
+                fakeTotalSamples,
+                downsampledTo: targetSampleCount,
+                from: reader,
+                channelSelection: .merged,
+                fftBands: nil
+            )
+        }
+
+        XCTAssertNotNil(analysis)
+        XCTAssertEqual(analysis?.amplitudes.count, targetSampleCount, "backfill should bring count to target")
+
+        // Threshold at 25 % of the buggy footprint — well above any noise but well below the bug.
+        let threshold = buggyAllocBytes / 4
+        XCTAssertLessThan(
+            peakDelta, threshold,
+            "Peak alloc delta \(peakDelta / 1_000_000) MB exceeds threshold \(threshold / 1_000_000) MB — backfill is still padding the input buffer"
+        )
+    }
+}