Poller.swift

private class Poller<ResultType, ErrorType: JiveError, ErrorCodeType where ErrorType.CodeType == ErrorCodeType> {
    private let retryInterval: NSTimeInterval
    private let timeoutErrorCode: ErrorCodeType
    private let expectation: XCTestExpectation
    private var state: PollerState<ResultType, ErrorType>
    
    init(
        createJiveFuture: () -> JiveFuture<ResultType, ErrorType>,
        acceptJiveResult: (JiveResult<ResultType, ErrorType>) -> JiveResult<ResultType, ErrorType>?,
        retryInterval: NSTimeInterval,
        timeoutErrorCode: ErrorCodeType,
        expectation: XCTestExpectation) {
            self.retryInterval = retryInterval
            self.timeoutErrorCode = timeoutErrorCode
            self.expectation = expectation
            state = .AwaitingFirstPoll(
                createJiveFuture,
                acceptJiveResult)
    }
    
    func startPolling() {
        if !NSThread.isMainThread() {
            fatalError("Can only startPolling on the main thread")
        }
        
        _performPollRecursively()
    }
    
    func finishPolling() -> JiveResult<ResultType, ErrorType> {
        if !NSThread.isMainThread() {
            fatalError("Can only finishPolling on the main thread")
        }
        
        switch state {
        case .AwaitingFirstPoll:
            fatalError("Must startPolling before you can finishPolling.")
        case .Polling(
            let jiveFuture,
            _,
            _):
            jiveFuture.cancel()
            return .Failure(JiveErrorBox(code: timeoutErrorCode))
        case .Done(let jiveResult):
            return jiveResult
        }
    }
    
    func _performPollRecursively() {
        switch state {
        case .AwaitingFirstPoll(
            let createJiveFuture,
            let acceptJiveResult):
            _performPoll(
                createJiveFuture: createJiveFuture,
                acceptJiveResult: acceptJiveResult)
        case .Polling(
            let jiveFuture,
            let createJiveFuture,
            let acceptJiveResult):
            if !jiveFuture.done {
                fatalError("JiveFuture should have already completed, so cancellation should fail")
            }
            _performPoll(
                createJiveFuture: createJiveFuture,
                acceptJiveResult: acceptJiveResult)
        case .Done:
            break
        }
    }
    
    func _performPoll(
        #createJiveFuture: () -> JiveFuture<ResultType, ErrorType>,
        acceptJiveResult: (JiveResult<ResultType, ErrorType>) -> JiveResult<ResultType, ErrorType>?) {
            let jiveFuture = createJiveFuture()
            state = .Polling(
                jiveFuture,
                createJiveFuture,
                acceptJiveResult)
            jiveFuture.onComplete() { jiveResult in
                switch self.state {
                case .AwaitingFirstPoll:
                    fatalError("Can't be awaiting first poll during a JiveFuture onComplete!")
                case .Polling(
                    let expectedJiveFuture,
                    _,
                    _):
                    if jiveFuture === expectedJiveFuture {
                        if let acceptedJiveResult = acceptJiveResult(jiveResult) {
                            self.state = .Done(jiveResult)
                            self.expectation.fulfill()
                        } else {
                            // recursion!!
                            dispatch_after(
                                (dispatch_time(DISPATCH_TIME_NOW,
                                    Int64(self.retryInterval * Double(NSEC_PER_SEC)))),
                                dispatch_get_main_queue(),
                                self._performPollRecursively)
                        }
                    } else {
                        fatalError("Can only poll one JiveFuture at a time. Got an onComplete for: \(jiveFuture) but expected: \(expectedJiveFuture)")
                    }
                case .Done:
                    fatalError("A JiveFuture should never complete when Done because we should cancel old JiveFutures before polling for new ones and after waiting")
                }
            }
    }
}

private enum PollerState<ResultType, ErrorType: JiveError>: Printable {
    // keep the closures in PollerState so they are immediately
    // freed when not needed. If kept in Poller, they'll live as long
    // as the Poller. When holding onto a closure, it seems safest
    // to free the closure as early as possible.
    case AwaitingFirstPoll(
        () -> JiveFuture<ResultType, ErrorType>,
        (JiveResult<ResultType, ErrorType>) -> JiveResult<ResultType, ErrorType>?)
    case Polling(
        JiveFuture<ResultType, ErrorType>,
        () -> JiveFuture<ResultType, ErrorType>,
        (JiveResult<ResultType, ErrorType>) -> JiveResult<ResultType, ErrorType>?)
    case Done(JiveResult<ResultType, ErrorType>)
    
    // MARK: Printable
    
    var description: String {
        switch self {
        case .AwaitingFirstPoll:
            return "AwaitingFirstPoll"
        case .Polling:
            return "Polling"
        case .Done:
            return "Done"
        }
    }
}