bogdan Add large import benchmark harness

01414eb a day ago History

import HealthKit
import SQLite3
import XCTest
@testable import HealthProbe

final class SQLiteLargeImportPerformanceTests: XCTestCase {
    private final class AsyncResultBox<T>: @unchecked Sendable {
        var result: Result<T, Error>?
    }

    private var temporaryDirectory: URL!

    override func setUpWithError() throws {
        temporaryDirectory = FileManager.default.temporaryDirectory
            .appending(path: "HealthProbeLargeImportTests-\(UUID().uuidString)", directoryHint: .isDirectory)
        try FileManager.default.createDirectory(at: temporaryDirectory, withIntermediateDirectories: true)
    }

    override func tearDownWithError() throws {
        if let temporaryDirectory {
            try? FileManager.default.removeItem(at: temporaryDirectory)
        }
        temporaryDirectory = nil
    }

    func testLargeSyntheticFullImportSmoke() throws {
        let definitions = syntheticImportDefinitions()
        let samples = makeSyntheticImportSamples(totalCount: 5_000, definitions: definitions)
        let url = databaseURL(named: "LargeImportSmoke-\(UUID().uuidString).sqlite")

        try runSyntheticImport(
            samples: samples,
            definitions: definitions,
            databaseURL: url
        )
    }

    func testLargeSyntheticFullImportBenchmark() throws {
        let environment = ProcessInfo.processInfo.environment
        let defaults = UserDefaults.standard
        let isEnabled = environment["HP_ENABLE_LARGE_IMPORT_BENCHMARK"] == "1"
            || defaults.bool(forKey: "HP_ENABLE_LARGE_IMPORT_BENCHMARK")

        guard isEnabled else {
            throw XCTSkip("""
            Large import benchmark is opt-in. Run with environment variable \
            HP_ENABLE_LARGE_IMPORT_BENCHMARK=1 or launch argument \
            -HP_ENABLE_LARGE_IMPORT_BENCHMARK YES, plus optional \
            HP_LARGE_IMPORT_SAMPLE_COUNT / HP_LARGE_IMPORT_ITERATIONS overrides.
            """)
        }

        let sampleCount = max(
            1_000,
            Int(environment["HP_LARGE_IMPORT_SAMPLE_COUNT"] ?? "")
                ?? defaults.integer(forKey: "HP_LARGE_IMPORT_SAMPLE_COUNT").nonZero
                ?? 250_000
        )
        let iterationCount = max(
            1,
            Int(environment["HP_LARGE_IMPORT_ITERATIONS"] ?? "")
                ?? defaults.integer(forKey: "HP_LARGE_IMPORT_ITERATIONS").nonZero
                ?? 1
        )
        let definitions = syntheticImportDefinitions()
        let samples = makeSyntheticImportSamples(totalCount: sampleCount, definitions: definitions)
        let options = XCTMeasureOptions()
        options.iterationCount = iterationCount

        measure(metrics: [XCTClockMetric(), XCTMemoryMetric()], options: options) {
            let url = databaseURL(named: "LargeImport-\(UUID().uuidString).sqlite")
            do {
                try runSyntheticImport(
                    samples: samples,
                    definitions: definitions,
                    databaseURL: url
                )
            } catch {
                XCTFail("Large synthetic import benchmark failed: \(error)")
            }
        }
    }

    private func databaseURL(named name: String) -> URL {
        temporaryDirectory.appending(path: name)
    }

    private func syntheticImportDefinitions() -> [(sampleType: HKQuantityType, unit: HKUnit, baseline: Double)] {
        [
            (HKQuantityType.quantityType(forIdentifier: .heartRate)!, HKUnit.count().unitDivided(by: .minute()), 55),
            (HKQuantityType.quantityType(forIdentifier: .activeEnergyBurned)!, .kilocalorie(), 1.5),
            (HKQuantityType.quantityType(forIdentifier: .basalEnergyBurned)!, .kilocalorie(), 1.1),
            (HKQuantityType.quantityType(forIdentifier: .stepCount)!, .count(), 12),
            (HKQuantityType.quantityType(forIdentifier: .distanceWalkingRunning)!, .meter(), 8)
        ]
    }

    private func makeSyntheticImportSamples(
        totalCount: Int,
        definitions: [(sampleType: HKQuantityType, unit: HKUnit, baseline: Double)]
    ) -> [HKQuantitySample] {
        return (0..<totalCount).map { index in
            let definition = definitions[index % definitions.count]
            let value = definition.baseline + Double(index % 17)
            let start = Date(timeIntervalSince1970: 10_000 + Double(index * 60))
            let end = start.addingTimeInterval(definition.sampleType.identifier == HKQuantityTypeIdentifier.heartRate.rawValue ? 5 : 60)
            let quantity = HKQuantity(unit: definition.unit, doubleValue: value)
            return HKQuantitySample(type: definition.sampleType, quantity: quantity, start: start, end: end)
        }
    }

    private func runSyntheticImport(
        samples: [HKQuantitySample],
        definitions: [(sampleType: HKQuantityType, unit: HKUnit, baseline: Double)],
        databaseURL url: URL
    ) throws {
        let store = SQLiteHealthArchiveStore(databaseURL: url)
        let observedAt = Date(timeIntervalSince1970: 5_000_000)

        defer {
            try? FileManager.default.removeItem(at: url)
            try? FileManager.default.removeItem(at: URL(fileURLWithPath: url.path + "-shm"))
            try? FileManager.default.removeItem(at: URL(fileURLWithPath: url.path + "-wal"))
        }

        try waitForArchiveOperation {
            let observationID = try await store.beginObservation(
                observedAt: observedAt,
                triggerReason: "benchmarkLargeSyntheticImport",
                selectedTypeSetHash: "synthetic-large-import"
            )
            _ = try await store.upsertSamples(
                samples,
                observedAt: observedAt,
                observationID: observationID
            )
            for definition in definitions {
                _ = try await store.markVerification(
                    sampleType: definition.sampleType,
                    verifiedAt: observedAt,
                    observationID: observationID
                )
            }
            try await store.finishObservation(
                observationID: observationID,
                status: "completed",
                endedAt: observedAt.addingTimeInterval(1)
            )
            return ()
        }

        XCTAssertEqual(try countRows(in: "samples", at: url), samples.count)
        XCTAssertEqual(try countRows(in: "sample_versions", at: url), samples.count)
        XCTAssertEqual(try countRows(in: "observation_type_runs", at: url), definitions.count)
    }

    private func waitForArchiveOperation<T>(_ operation: @escaping () async throws -> T) throws -> T {
        let expectation = expectation(description: "archive operation")
        let box = AsyncResultBox<T>()

        Task {
            do {
                box.result = .success(try await operation())
            } catch {
                box.result = .failure(error)
            }
            expectation.fulfill()
        }

        wait(for: [expectation], timeout: 120)
        return try XCTUnwrap(box.result).get()
    }

    private func countRows(in tableName: String, at url: URL) throws -> Int {
        var db: OpaquePointer?
        guard sqlite3_open_v2(url.path, &db, SQLITE_OPEN_READONLY | SQLITE_OPEN_FULLMUTEX, nil) == SQLITE_OK else {
            sqlite3_close(db)
            XCTFail("Could not open benchmark database")
            return 0
        }
        defer { sqlite3_close(db) }

        var statement: OpaquePointer?
        guard sqlite3_prepare_v2(db, "SELECT COUNT(*) FROM \(tableName)", -1, &statement, nil) == SQLITE_OK else {
            sqlite3_finalize(statement)
            XCTFail("Could not prepare benchmark count query")
            return 0
        }
        defer { sqlite3_finalize(statement) }

        guard sqlite3_step(statement) == SQLITE_ROW else {
            return 0
        }
        return Int(sqlite3_column_int(statement, 0))
    }
}

private extension Int {
    var nonZero: Int? {
        self == 0 ? nil : self
    }
}