MacMonitor/Sources/Telemetry/Storage/MMDiskIOProvider.m

#import "MMDiskIOProvider.h"
#import <IOKit/IOKitLib.h>
#import <IOKit/storage/IOBlockStorageDriver.h>
#import <mach/mach_time.h>
#import <os/lock.h>

@interface MMDiskIOProvider () {
    os_unfair_lock _lock;
    NSDictionary<NSString *, NSDictionary<NSString *, NSNumber *> *> *_previousSnapshots;
    uint64_t _previousTimestamp;
    mach_timebase_info_data_t _timebase;
}
@end

@implementation MMDiskIOProvider

- (instancetype)init {
    self = [super init];
    if (self) {
        _lock = OS_UNFAIR_LOCK_INIT;
        _previousSnapshots = nil;
        _previousTimestamp = 0;
        mach_timebase_info(&_timebase);
    }
    return self;
}

- (MMTelemetryDomain)domain {
    return MMTelemetryDomainStorage;
}

- (NSString *)providerIdentifier {
    return @"com.i3omb.macmonitor.telemetry.storage.io";
}

- (BOOL)isAvailable {
    return YES;
}

+ (NSDictionary<NSString *, id> *)calculateDiskIOMetricsWithCurrentSnapshots:(NSDictionary<NSString *, NSDictionary<NSString *, NSNumber *> *> *)current
                                                           previousSnapshots:(nullable NSDictionary<NSString *, NSDictionary<NSString *, NSNumber *> *> *)previous
                                                                   timeDelta:(NSTimeInterval)timeDelta {
    double totalReadBps = 0.0;
    double totalWriteBps = 0.0;
    double totalReadIOPS = 0.0;
    double totalWriteIOPS = 0.0;
    
    NSMutableArray<NSDictionary<NSString *, id> *> *disks = [NSMutableArray arrayWithCapacity:current.count];
    NSArray<NSString *> *sortedBsdNames = [current.allKeys sortedArrayUsingSelector:@selector(compare:)];
    
    for (NSString *bsdName in sortedBsdNames) {
        NSDictionary<NSString *, NSNumber *> *curr = current[bsdName];
        NSDictionary<NSString *, NSNumber *> *prev = previous ? previous[bsdName] : nil;
        
        uint64_t currReadBytes = [curr[@"readBytes"] unsignedLongLongValue];
        uint64_t currWriteBytes = [curr[@"writeBytes"] unsignedLongLongValue];
        uint64_t currReadOps = [curr[@"readOps"] unsignedLongLongValue];
        uint64_t currWriteOps = [curr[@"writeOps"] unsignedLongLongValue];
        
        uint64_t prevReadBytes = prev ? [prev[@"readBytes"] unsignedLongLongValue] : currReadBytes;
        uint64_t prevWriteBytes = prev ? [prev[@"writeBytes"] unsignedLongLongValue] : currWriteBytes;
        uint64_t prevReadOps = prev ? [prev[@"readOps"] unsignedLongLongValue] : currReadOps;
        uint64_t prevWriteOps = prev ? [prev[@"writeOps"] unsignedLongLongValue] : currWriteOps;
        
        double deltaReadBytes = (currReadBytes >= prevReadBytes) ? (double)(currReadBytes - prevReadBytes) : 0.0;
        double deltaWriteBytes = (currWriteBytes >= prevWriteBytes) ? (double)(currWriteBytes - prevWriteBytes) : 0.0;
        double deltaReadOps = (currReadOps >= prevReadOps) ? (double)(currReadOps - prevReadOps) : 0.0;
        double deltaWriteOps = (currWriteOps >= prevWriteOps) ? (double)(currWriteOps - prevWriteOps) : 0.0;
        
        double readBps = (timeDelta > 0.0001) ? (deltaReadBytes / timeDelta) : 0.0;
        double writeBps = (timeDelta > 0.0001) ? (deltaWriteBytes / timeDelta) : 0.0;
        double readIOPS = (timeDelta > 0.0001) ? (deltaReadOps / timeDelta) : 0.0;
        double writeIOPS = (timeDelta > 0.0001) ? (deltaWriteOps / timeDelta) : 0.0;
        
        totalReadBps += readBps;
        totalWriteBps += writeBps;
        totalReadIOPS += readIOPS;
        totalWriteIOPS += writeIOPS;
        
        [disks addObject:@{
            @"bsdName": bsdName,
            @"readBytesPerSec": @(readBps),
            @"writeBytesPerSec": @(writeBps),
            @"readIOPS": @(readIOPS),
            @"writeIOPS": @(writeIOPS),
            @"cumulativeReadBytes": @(currReadBytes),
            @"cumulativeWriteBytes": @(currWriteBytes),
            @"cumulativeReadOps": @(currReadOps),
            @"cumulativeWriteOps": @(currWriteOps)
        }];
    }
    
    return @{
        @"disks": disks,
        @"totalReadBytesPerSec": @(totalReadBps),
        @"totalWriteBytesPerSec": @(totalWriteBps),
        @"totalReadIOPS": @(totalReadIOPS),
        @"totalWriteIOPS": @(totalWriteIOPS),
        @"timeDelta": @(timeDelta)
    };
}

- (nullable NSDictionary<NSString *, id> *)sampleTelemetryWithError:(NSError **)error {
    uint64_t now = mach_absolute_time();
    
    NSMutableDictionary<NSString *, NSDictionary<NSString *, NSNumber *> *> *currentSnapshots = [NSMutableDictionary dictionary];
    
    CFMutableDictionaryRef matching = IOServiceMatching(kIOBlockStorageDriverClass);
    io_iterator_t iterator = IO_OBJECT_NULL;
    if (IOServiceGetMatchingServices(kIOMainPortDefault, matching, &iterator) == KERN_SUCCESS && iterator != IO_OBJECT_NULL) {
        io_registry_entry_t entry;
        while ((entry = IOIteratorNext(iterator)) != IO_OBJECT_NULL) {
            NSString *bsdName = nil;
            io_iterator_t childIterator;
            if (IORegistryEntryGetChildIterator(entry, kIOServicePlane, &childIterator) == KERN_SUCCESS) {
                io_registry_entry_t child;
                while ((child = IOIteratorNext(childIterator)) != IO_OBJECT_NULL) {
                    CFTypeRef nameRef = IORegistryEntryCreateCFProperty(child, CFSTR("BSD Name"), kCFAllocatorDefault, 0);
                    if (nameRef) {
                        bsdName = CFBridgingRelease(nameRef);
                        IOObjectRelease(child);
                        break;
                    }
                    IOObjectRelease(child);
                }
                IOObjectRelease(childIterator);
            }
            
            if (bsdName) {
                CFTypeRef statsRef = IORegistryEntryCreateCFProperty(entry, CFSTR("Statistics"), kCFAllocatorDefault, 0);
                if (statsRef) {
                    NSDictionary *stats = CFBridgingRelease(statsRef);
                    NSNumber *readBytes = stats[@"Bytes (Read)"] ?: @0;
                    NSNumber *writeBytes = stats[@"Bytes (Write)"] ?: @0;
                    NSNumber *readOps = stats[@"Operations (Read)"] ?: @0;
                    NSNumber *writeOps = stats[@"Operations (Write)"] ?: @0;
                    
                    currentSnapshots[bsdName] = @{
                        @"readBytes": readBytes,
                        @"writeBytes": writeBytes,
                        @"readOps": readOps,
                        @"writeOps": writeOps
                    };
                }
            }
            IOObjectRelease(entry);
        }
        IOObjectRelease(iterator);
    }
    
    os_unfair_lock_lock(&_lock);
    NSDictionary<NSString *, NSDictionary<NSString *, NSNumber *> *> *prev = _previousSnapshots;
    uint64_t prevTime = _previousTimestamp;
    
    _previousSnapshots = currentSnapshots;
    _previousTimestamp = now;
    os_unfair_lock_unlock(&_lock);
    
    NSTimeInterval timeDelta = 1.0;
    if (prevTime > 0) {
        uint64_t elapsed = now - prevTime;
        timeDelta = (double)elapsed * _timebase.numer / _timebase.denom / 1e9;
    }
    
    return [MMDiskIOProvider calculateDiskIOMetricsWithCurrentSnapshots:currentSnapshots
                                                      previousSnapshots:prev
                                                              timeDelta:timeDelta];
}

@end