CloudKitExtras.swift

//
//  John Scalo
//  CloudKitExtras.swift
//

import CloudKit

var ckLogging = true
let ckBatchLimit = 400
let usleepInterval: UInt32 = 500000

// A note on rate limiting: Unfortunately the minimium CKErrorRetryAfterKey interval that CK ever provides seems to be 30s, and more often than not it's 60s. Obviously no user is going to wait that long and in practice waiting only just a few seconds is enough to get past the rate limit. So instead we take the minimum of CK's provided interval and 5s.

extension CKQueryOperation {
    // Use this when you know you want all the records, cursor be damned.
    class func performExhaustiveQuery(query: CKQuery,
                                      database: CKDatabase = CloudKitManager.publicDatabase,
                                      desiredKeys: [String]? = nil,
                                      qualityOfService: QualityOfService = .userInitiated,
                                      funcName: String = #function,
                                      recordFoundHandler: @escaping(_ record: CKRecord) -> Void,
                                      completion: @escaping(_ error: Error?) -> Void) {
        DispatchQueue.global().async {
            var foundError: Error?
            var currentCursor: CKQueryCursor?
            var shouldRetry = false
            repeat {
                let sem = DispatchSemaphore(value: 0)
                
                let queryOp: CKQueryOperation
                
                if let cursor = currentCursor {
                    queryOp = CKQueryOperation(cursor: cursor)
                } else {
                    queryOp = CKQueryOperation(query: query)
                }
                queryOp.qualityOfService = qualityOfService
                if let desiredKeys = desiredKeys {
                    queryOp.desiredKeys = desiredKeys
                }
                
                var cnt = 0
                queryOp.recordFetchedBlock = { record in
                    recordFoundHandler(record)
                    cnt += 1
                }
                
                shouldRetry = false
                queryOp.queryCompletionBlock = { (cursor, error) in
                    if ckLogging { Logger.fileLog("(CKLogging) \(funcName): Finished") }
                    if let error = error {
                        Logger.fileLog("*** Error: \(error)")
                        if let retry = (error as NSError).userInfo[CKErrorRetryAfterKey] as? TimeInterval {
                            let newRetry = min(retry, 5)
                            Logger.fileLog("Will retry after \(newRetry) seconds (actual was: \(retry)")
                            shouldRetry = true
                            usleep(UInt32(ceil(newRetry)) * 1000000)
                            sem.signal()
                            return
                        } else {
                            foundError = error
                        }
                    }
                    if ckLogging { Logger.fileLog("(CKLogging) \(funcName): \(query.recordType): Fetched \(cnt) records") }
                    currentCursor = cursor
                    sem.signal()
                }
                
                if ckLogging { Logger.fileLog("(CKLogging) \(funcName): Start") }
                database.add(queryOp)
                sem.wait()
                if currentCursor != nil && !shouldRetry /* usleep() for retry already happened */ {
                    if ckLogging { Logger.fileLog("(CKLogging) \(funcName): Sleeping for next cursor") }
                    usleep(usleepInterval)
                }
            } while currentCursor != nil || shouldRetry
            
            completion(foundError)
        }
    }
    
    class func performQuery(query: CKQuery,
                          database: CKDatabase = CloudKitManager.publicDatabase,
                          resultsLimit: Int = CKQueryOperationMaximumResults,
                          cursor: CKQueryCursor? = nil,
                          desiredKeys: [String]? = nil,
                          qualityOfService: QualityOfService = .userInitiated,
                          funcName: String = #function,
                          recordFoundHandler: @escaping(_ record: CKRecord) -> Void,
                          completion: @escaping(_ cursor: CKQueryCursor?, _ error: Error?) -> Void) {
        
        func retryAfter(_ t: TimeInterval) {
            let newRetry = min(t, 5)
            Logger.fileLog("Will retry after \(newRetry) seconds (actual was: \(t)")
            DispatchQueue.global().asyncAfter(deadline: .now() + newRetry) {
                CKQueryOperation.performQuery(query: query,
                                              database: database,
                                              resultsLimit: resultsLimit,
                                              cursor: cursor,
                                              desiredKeys: desiredKeys,
                                              qualityOfService: qualityOfService,
                                              recordFoundHandler: recordFoundHandler,
                                              completion: completion)
            }
        }
                
        let queryOp = cursor == nil ? CKQueryOperation(query: query) : CKQueryOperation(cursor: cursor!)
        queryOp.qualityOfService = qualityOfService
        queryOp.resultsLimit = resultsLimit
        if let desiredKeys = desiredKeys {
            queryOp.desiredKeys = desiredKeys
        }
        queryOp.recordFetchedBlock = { record in
            recordFoundHandler(record)
        }
        queryOp.queryCompletionBlock = { (cursor, error) in
            if ckLogging { Logger.fileLog("(CKLogging) \(funcName): Finished") }
            if let error = error {
                Logger.fileLog("*** Error: \(error)")
                if let retry = (error as NSError).userInfo[CKErrorRetryAfterKey] as? TimeInterval {
                    retryAfter(retry)
                    return
                }
            }
            completion(cursor, error)
        }
        
        if ckLogging { Logger.fileLog("(CKLogging) \(funcName): Start") }
        database.add(queryOp)
    }
}

extension CKModifyRecordsOperation {
    /// Handles batching and retries.
    static func performModify(recordsToSave: [CKRecord]?,
                            recordIDsToDelete: [CKRecordID]?,
                            savePolicy: CKRecordSavePolicy = .changedKeys,
                            desiredKeys: [String]? = nil,
                            qualityOfService: QualityOfService = .userInitiated,
                            funcName: String = #function,
                            completion: @escaping(_ records: [CKRecord], _ error: NSError?) -> Void) {
        
        func retryAfter(_ t: TimeInterval) {
            let newRetry = min(t, 5)
            Logger.fileLog("Will retry after \(newRetry) seconds (actual was: \(t)")
            DispatchQueue.global().asyncAfter(deadline: .now() + newRetry) {
                CKModifyRecordsOperation.performModify(recordsToSave: recordsToSave,
                                                     recordIDsToDelete: recordIDsToDelete,
                                                     savePolicy: savePolicy,
                                                     desiredKeys: desiredKeys,
                                                     completion: completion)
            }
        }
        
        DispatchQueue.global().async {
            var returnRecords = [CKRecord]()
            var foundError: NSError?
            if let recordsToSave = recordsToSave {
                var idx = 0
                var upperBound = 0
                repeat {
                    upperBound = min(idx + ckBatchLimit, recordsToSave.count)
                    let batch = Array(recordsToSave[idx..<upperBound])
                    let (savedRecords,error) = CKModifyRecordsOperation.modifyRecordsSync(recordsToSave: batch, recordsIDsToDelete: nil, funcName: funcName, savePolicy: savePolicy, desiredKeys: desiredKeys, qualityOfService: qualityOfService)
                    if let error = error {
                        Logger.fileLog("*** Error: \(error)")
                        if let retry = (error as NSError).userInfo[CKErrorRetryAfterKey] as? TimeInterval {
                            retryAfter(retry)
                            return
                        }
                        foundError = error
                        break
                    } else if let savedRecords = savedRecords {
                        returnRecords += savedRecords
                    }
                    idx = upperBound
                    if idx < recordsToSave.count {
                        usleep(usleepInterval)
                    }
                } while idx < recordsToSave.count
            }
            if let recordIDsToDelete = recordIDsToDelete {
                var idx = 0
                var upperBound = 0
                repeat {
                    upperBound = min(idx + ckBatchLimit, recordIDsToDelete.count)
                    let batch = Array(recordIDsToDelete[idx..<upperBound])
                    let (_,error) = CKModifyRecordsOperation.modifyRecordsSync(recordsToSave: nil, recordsIDsToDelete: batch, funcName: funcName, savePolicy: nil, desiredKeys: desiredKeys, qualityOfService: qualityOfService)
                    if let error = error {
                        Logger.fileLog("*** Error: \(error)")
                        if let retry = (error as NSError).userInfo[CKErrorRetryAfterKey] as? TimeInterval {
                            retryAfter(retry)
                            return
                        }
                        foundError = error
                        break
                    }
                    idx = upperBound
                    if idx < recordIDsToDelete.count {
                        usleep(usleepInterval)
                    }
                } while idx < recordIDsToDelete.count
            }
            
            completion(returnRecords, foundError)
        }
    }
    
    private static func modifyRecordsSync(recordsToSave: [CKRecord]?,
                                  recordsIDsToDelete: [CKRecordID]?,
                                  funcName: String,
                                  savePolicy: CKRecordSavePolicy?,
                                  desiredKeys: [String]?,
                                  qualityOfService: QualityOfService) -> ([CKRecord]?, NSError?) {
        let sem = DispatchSemaphore(value: 0)
        var foundError: NSError?
        var foundRecords: [CKRecord]?
        let op = CKModifyRecordsOperation(recordsToSave: recordsToSave, recordIDsToDelete: recordsIDsToDelete)
        if let savePolicy = savePolicy {
            op.savePolicy = savePolicy
        }
        op.qualityOfService = qualityOfService
        if let desiredKeys = desiredKeys, let recordsToSave = recordsToSave {
            for nextRecord in recordsToSave {
                for nextKey in nextRecord.allKeys() {
                    if !desiredKeys.contains(nextKey) {
                        nextRecord.setObject(nil, forKey: nextKey)
                    }
                }
            }
        }
        op.qualityOfService = .userInitiated
        op.modifyRecordsCompletionBlock = { (records, _, error) in
            if ckLogging { Logger.fileLog("(CKLogging) \(funcName): Finished") }
            foundRecords = records
            foundError = error as NSError?
            sem.signal()
        }
        if ckLogging { Logger.fileLog("(CKLogging) \(funcName): Start") }
        CloudKitManager.publicDatabase.add(op)
        sem.wait()
        return (foundRecords, foundError)
    }
}

extension CKFetchRecordsOperation {
    static func performFetch(recordIDs: [CKRecordID],
                             desiredKeys: [String]? = nil,
                             funcName: String = #function,
                             completion: ((_ records: [CKRecordID : CKRecord], _ error: NSError?) -> Void)? = nil) {
        
        
        DispatchQueue.global().async {
            var idx = 0
            var upperBound = 0
            var foundRecords = [CKRecordID:CKRecord]()
            var foundError: NSError?
            
            repeat {
                upperBound = min(idx + ckBatchLimit, recordIDs.count)
                var shouldRetry = false
                let batchRecordIDs = Array(recordIDs[idx..<upperBound])
                let op = CKFetchRecordsOperation(recordIDs: batchRecordIDs)
                op.desiredKeys = desiredKeys
                op.qualityOfService = .userInitiated
                
                let sem = DispatchSemaphore(value: 0)
                op.fetchRecordsCompletionBlock = { (records, error) in
                    if ckLogging { Logger.fileLog("(CKLogging) \(funcName): Finished") }
                    if let error = error {
                        Logger.fileLog("*** Error: \(error)")
                        if let retry = (error as NSError).userInfo[CKErrorRetryAfterKey] as? TimeInterval {
                            let newRetry = min(retry, 5)
                            Logger.fileLog("Will retry after \(newRetry) seconds (actual was: \(retry)")
                            shouldRetry = true
                            sleep(UInt32(ceil(newRetry)))
                        } else {
                            foundError = error as NSError
                        }
                    } else if let records = records {
                        foundRecords.merge(records) { (_, new) in new }
                    }
                    sem.signal()
                }
                
                if upperBound < recordIDs.count {
                    if ckLogging { Logger.fileLog("(CKLogging) \(funcName): Start batch \(idx)..<\(upperBound)") }
                } else {
                    if ckLogging { Logger.fileLog("(CKLogging) \(funcName): Start") }
                }
                CloudKitManager.publicDatabase.add(op)
                sem.wait()
                
                if foundError != nil {
                    break
                }
                if !shouldRetry {
                    idx = upperBound
                }
                if idx < recordIDs.count {
                    usleep(usleepInterval)
                }
            } while idx < recordIDs.count
            
            completion?(foundRecords, foundError)
        }
    }
}