Block 调用检查

by Ricky Tan

前言

如何确保一个传递给别人的 Block 被调用过，是一个一直困扰我的问题，因为 Block 作为 iOS 的一种回调机制，它可以像函数一样马上被调用，也可以像对象一样被持有、被传递、被释放，并在将来的某个时候被调用。有些时候我们传出去的 Block 必须被调用一次，否则会处于一种不确定的状态而导致程序无法继续，或者出错。例如，之前一篇文章《一种 App 内路由系统的设计》中的路由注册方式，如果使用 Block 方式，那么在路由完成后需要调用一次 complete 以通知路由系统已经完成，否则无法处理新的路由。而在实际开发过程中，确实会遇到条件过多以后，在某些条件下忘记调用 complete 的情况。

WebKit 的实现

如何让系统自动检测出来？这个问题一直没有思路，直到某一天在处理 WebKit 的 delegate 时突发奇想，我不调用它的 handler 会怎么样？

我们知道，相比于 UIWebview 的直接返回布尔值的方式，WKWebview 把决定是否切换导航做成了异步回调的方式。

// UIWebview
- (BOOL)webView:(UIWebView *)webView
shouldStartLoadWithRequest:(NSURLRequest *)request  
 navigationType:(UIWebViewNavigationType)navigationType;

// WKWebview
- (void)webView:(WKWebView *)webView 
decidePolicyForNavigationAction:(WKNavigationAction *)navigationAction  
decisionHandler:(void (^)(WKNavigationActionPolicy))decisionHandler;

这样的改动，使得我们甚至可以在请求一次网络后，根据返回值再决定导航动作。这个 decisionHandler 可以被传递到其他对像上，并把决定权交给它。Block 被传递时像普通对象一样会被引用（拷贝），最终会被释放。但是一旦 decisionHandler 释放前没有被调用过，WebKit 会抛出一个异常：

2018-07-02 18:01:06.625 [13522:489767] *** Terminating app due to uncaught exception 'NSInternalInconsistencyException', reason: 'Completion handler passed to -[RTViewController webView:decidePolicyForNavigationAction:decisionHandler:] was not called'  
*** First throw call stack:
(
    0   CoreFoundation                      0x0000000106c2ef35 __exceptionPreprocess + 165
    1   libobjc.A.dylib                     0x00000001068c9bb7 objc_exception_throw + 45
    2   CoreFoundation                      0x0000000106c2ee6d +[NSException raise:format:] + 205
    3   WebKit                              0x00000001071327da _ZN6WebKit28CompletionHandlerCallCheckerD2Ev + 130
    4   WebKit                              0x000000010720a180 _ZN3WTF20ThreadSafeRefCountedIN6WebKit28CompletionHandlerCallCheckerEE5derefEv + 36
    5   WebKit                              0x0000000107207cb9 _ZN6WebKit15NavigationState12PolicyClient31decidePolicyForNavigationActionEPNS_12WebPageProxyEPNS_13WebFrameProxyERKNS_20NavigationActionDataES5_RKN7WebCore15ResourceRequestESC_N3WTF6RefPtrINS_27WebFramePolicyListenerProxyEEEPN3API6ObjectE + 935
    ...
)
libc++abi.dylib: terminating with uncaught exception of type NSException

很神奇！从堆栈上看，这个异常是一个名为 _ZN6WebKit28CompletionHandlerCallCheckerD2Ev的东西抛出的（名字为什么会这样？见名字修饰）。好在 Apple 早就开源了 WebKit 的源码（https://webkit.org/），下载查阅后确实找到了一个 CompletionHandlerCallChecker：

namespace WebKit {

Ref<CompletionHandlerCallChecker> CompletionHandlerCallChecker::create(id delegate, SEL delegateMethodSelector)  
{
    return adoptRef(*new CompletionHandlerCallChecker(object_getClass(delegate), delegateMethodSelector));
}

CompletionHandlerCallChecker::CompletionHandlerCallChecker(Class delegateClass, SEL delegateMethodSelector)  
    : m_delegateClass(delegateClass)
    , m_delegateMethodSelector(delegateMethodSelector)
    , m_didCallCompletionHandler(false)
{
}

CompletionHandlerCallChecker::~CompletionHandlerCallChecker()  
{
    if (m_didCallCompletionHandler)
        return;
    Class delegateClass = classImplementingDelegateMethod();
    [NSException raise:NSInternalInconsistencyException format:@"Completion handler passed to %c[%@ %@] was not called", class_isMetaClass(delegateClass) ? '+' : '-', NSStringFromClass(delegateClass), NSStringFromSelector(m_delegateMethodSelector)];
}

void CompletionHandlerCallChecker::didCallCompletionHandler()  
{
    ASSERT(!m_didCallCompletionHandler);
    m_didCallCompletionHandler = true;
}

Class CompletionHandlerCallChecker::classImplementingDelegateMethod() const  
{
    Class delegateClass = m_delegateClass;
    Method delegateMethod = class_getInstanceMethod(delegateClass, m_delegateMethodSelector);

    for (Class superclass = class_getSuperclass(delegateClass); superclass; superclass = class_getSuperclass(superclass)) {
        if (class_getInstanceMethod(superclass, m_delegateMethodSelector) != delegateMethod)
            break;
        delegateClass = superclass;
    }
    return delegateClass;
}
} // namespace WebKit
#endif // WK_API_ENABLED

最关键的代码在析构方法上，当 m_didCallCompletionHandler 为 false 时，直接抛异常。而这个 m_didCallCompletionHandler 什么时候会设为 true 呢？查看源码 NavigationState.mm 第 354 行左右：

    // 为了方便理解，去掉了不相关的代码
    ...
    RefPtr<CompletionHandlerCallChecker> checker = CompletionHandlerCallChecker::create(navigationDelegate.get(), @selector(webView:decidePolicyForNavigationAction:decisionHandler:));

    auto decisionHandlerWithPolicies = [localListener = RefPtr<WebFramePolicyListenerProxy>(WTFMove(listener)), localNavigationAction = RefPtr<API::NavigationAction>(&navigationAction), checker = WTFMove(checker), mainFrameURLString](WKNavigationActionPolicy actionPolicy, _WKWebsitePolicies *websitePolicies) mutable {
        if (checker->completionHandlerHasBeenCalled())
            return;
        checker->didCallCompletionHandler();
        ...
    };

    ...
    else {
        auto decisionHandlerWithoutPolicies = [decisionHandlerWithPolicies] (WKNavigationActionPolicy actionPolicy) mutable {
            decisionHandlerWithPolicies(actionPolicy, nil);
        };
        [navigationDelegate webView:m_navigationState.m_webView decidePolicyForNavigationAction:wrapper(navigationAction) decisionHandler:decisionHandlerWithoutPolicies];
    }
}

WebKit 传给开发者的 decisionHandlerWithoutPolicies 实际上是 CPP 闭包：

auto decisionHandlerWithoutPolicies = []() {  
};

不过不用担心，编译器会将它转换为 Objective-C 的 Stack Block。它捕获了 checker 实例，如果它被调用了，则会调用 checker->didCallCompletionHandler();，如果一直没有调用过，当它释放时，checker 实例也被释放从而调用析构方法，并抛出异常！

这样一来其实思路就有了。

Checker

Objective-C 相比于 CPP 有着天然的优势，因为它原生就是引用计数的，不再需要额外的代码支持。那么我们实现一个 RTBlockChecker：

@interface RTBlockChecker : NSObject
@property (nonatomic, readonly, assign) BOOL hasBeenCalled;
@end

@implementation RTBlockChecker

- (void)dealloc
{
    if (_hasBeenCalled)
        return;

    [NSException raise:NSInternalInconsistencyException
                format:@"not called!"];
}

- (void)didCalled {
    _hasBeenCalled = YES;
}

@end

很简单，几行代码。然后，我们使用它的地方代码需要一点点改造：

// 原来的实现
void (^block)(...) = ...;  
[self callAMethodWithBlock:block];

// 新的实现
void (^block)(...) = ...;  
RTBlockChecker *checker = [RTBlockChecker new];  
void (^blockWithChecker)(...) = ^(...) {  
    [checker didCalled];
    block(...);
};
[self callAMethodWithBlock:blockWithChecker];

这样的实现，如果仅仅是在自己的项目中运用是完全够用了，但是如果想当作一个通用的第三方组件，代码入侵性就有点大了，另一方面，这种改动仅适用于有源代码的项目，对于没有源码的库，它返回给开发者的 Block 没法让它拥有调用检查的特性。于是不得不换一种思路。

自定义 Block

Block 对 OC 来说也是一种对象，也有完整的生命周期，可不可能在 Block 自身释放时做一些事情？答案是可以的。

以下是 Apple 对 Block 对象的定义（源码可以在 libclosure 找到）：

// Values for Block_layout->flags to describe block objects
enum {  
    BLOCK_DEALLOCATING =      (0x0001),  // runtime
    BLOCK_REFCOUNT_MASK =     (0xfffe),  // runtime
    BLOCK_NEEDS_FREE =        (1 << 24), // runtime
    BLOCK_HAS_COPY_DISPOSE =  (1 << 25), // compiler
    BLOCK_HAS_CTOR =          (1 << 26), // compiler: helpers have C++ code
    BLOCK_IS_GC =             (1 << 27), // runtime
    BLOCK_IS_GLOBAL =         (1 << 28), // compiler
    BLOCK_USE_STRET =         (1 << 29), // compiler: undefined if !BLOCK_HAS_SIGNATURE
    BLOCK_HAS_SIGNATURE  =    (1 << 30), // compiler
    BLOCK_HAS_EXTENDED_LAYOUT=(1 << 31)  // compiler
};
#define BLOCK_DESCRIPTOR_1 1
struct Block_descriptor_1 {  
    uintptr_t reserved;
    uintptr_t size;
};

#define BLOCK_DESCRIPTOR_2 1
struct Block_descriptor_2 {  
    // requires BLOCK_HAS_COPY_DISPOSE
    void (*copy)(void *dst, const void *src);
    void (*dispose)(const void *);
};

#define BLOCK_DESCRIPTOR_3 1
struct Block_descriptor_3 {  
    // requires BLOCK_HAS_SIGNATURE
    const char *signature;
    const char *layout;     // contents depend on BLOCK_HAS_EXTENDED_LAYOUT
};

struct Block_layout {  
    void *isa;
    volatile int32_t flags; // contains ref count
    int32_t reserved; 
    void (*invoke)(void *, ...);
    struct Block_descriptor_1 *descriptor;
    // imported variables
};

注意这里有三种 descriptor，Block_descriptor_1 是一定存在的，Block_descriptor_2 根据 flags 是否包含 BLOCK_HAS_COPY_DISPOSE 可选，Block_descriptor_3 根据是否包含 BLOCK_HAS_SIGNATURE 可选。那么我们可以构造一个 Block，自定义它的析构函数。 先按 Apple 的定义重写一套自己的类型：

struct RTBlock_Descriptor {  
    uintptr_t reserved;
    uintptr_t size;
    void (*copy)(void *dst, const void *src);
    void (*dispose)(const void *);
    const char *signature;
    const char *layout;
};

enum {  
    BLOCK_NEEDS_FREE       = (1 << 24),
    BLOCK_HAS_COPY_DISPOSE = (1 << 25),
    BLOCK_IS_GC            = (1 << 27),
    BLOCK_IS_GLOBAL        = (1 << 28),
    BLOCK_HAS_STRET        = (1 << 29),
    BLOCK_HAS_SIGNATURE    = (1 << 30),
};

struct RTBlock {  
    Class isa;
    int32_t flags;
    int32_t reserved;
    IMP invoke;
    const struct RTBlock_Descriptor* descriptor;
    void *forwardingBlock;
};
typedef struct RTBlock RTBlock;

定义一个自己的析构函数，实现函数及静态的 descriptor 常量：

static void rt_blockDispose(const void *block) {  
    Block_release(((const RTBlock *)block)->forwardingBlock);
    if (((const RTBlock *)block)->reserved == 0) {
        // exception!
    }
}
static void rt_blockInvoke(void *block, ...) {  
    ((RTBlock *)block)->reserved = 1;
    // pass all parameters to forwardingBlock
}

static const struct RTBlock_Descriptor RTDescriptor = {  
    0,
    sizeof(RTBlock),
    NULL,
    (void (*)(const void *))rt_blockDispose,
};

然后写一个函数构造自己的 Block：

RTBlock *block = (RTBlock *)(malloc(sizeof(RTBlock)));  
block->isa = NSClassFromString(@"__NSMallocBlock__");

const unsigned retainCount = 1;

block->flags = BLOCK_HAS_COPY_DISPOSE | BLOCK_NEEDS_FREE | (retainCount << 1);  
block->reserved = 0;  
block->invoke = rt_blockInvoke;  
block->descriptor = &RTDescriptor;  
block->forwardingBlock = (__bridge void *)[originBlock copy];

这样一来，当 Block 被调用时，实际上进入了 rt_blockInvoke 函数，其中第一个参数 block 相当于 self。我们可以在这里记一次调用（这里暂时利用 reserved），然后再调原始真正干活的 Block（forwardkingBlock），最后在 rt_blockDispose 中检查是否被调用过。

新的问题

听起来这个解决方案十分完美，但马上就遇到新的问题了，原始 Block 可能是任意参数的！ 这里的 rt_blockInvoke 虽然是可变参数的函数，但是它无法依次以正确的类型取出所有参数的，并传给 forwardingBlock！

对于这个问题有想到几个方案：

  1. libffi。libffi 让 C 语言拥有知道函数指针就可以任意调用的能力，且可以任意参数。但是一个简单的 Block 检查功能引入这样一个库肯定是不划算的。
  2. 汇编。函数参数的传递是通过特定的寄存器的，用汇编语言实现 rt_blockInvoke 可以直接绕开那些寄存器，然后用 br 指令直接跳到 forwardingBlock 的实现体。事实上 libffi 底层就是汇编了，另一同事给出了这种方案的实现：BlockCallAssert。

其它解决方法

能不能利用 OC 自身的动态性在不引入汇编的情况下做到？留意下之前的 Block_descriptor_3 中的 signature 是不是很熟悉？对，就是 NSObject selector 的 signature。可以动态地获取 Block 的参数信息：

static NSMethodSignature *rt_blockMethodSignature(id block) {  
    if (!block) {
        return nil;
    }

    RTBlock *layout = (__bridge RTBlock *)block;
    if (!(layout->flags & BLOCK_HAS_SIGNATURE)) {
        return nil;
    }

    char *desc = (char *)layout->descriptor;
    desc += 2 * sizeof(uintptr_t);
    if (layout->flags & BLOCK_HAS_COPY_DISPOSE) {
        desc += 2 * sizeof(void *);
    }
    if (!desc) {
        return nil;
    }
    const char *signature = *(const char **)desc;
    return [NSMethodSignature signatureWithObjCTypes:signature];
}

有了这些信息应该可以处理大部分情况了（为什么是大部分情况，下面会说明），取出原始参数值，传到 NSInvocation，最后调用就好：

static void rt_blockInvoke(id block, ...) {  
    NSMethodSignature *signature = rt_blockMethodSignature(block);
    NSInvocation *invocation = [NSInvocation invocationWithMethodSignature:signature];

    va_list args;
    va_start(args, block);

    NSUInteger numberOfArguments = signature.numberOfArguments;
    for (NSUInteger i = 1; i < numberOfArguments; ++i) {
        const char *argType = [signature getArgumentTypeAtIndex:i];
        switch (argType[0]) {
            case 'c':
            case 's':
            {
                char param = va_arg(args, int);
                [invocation setArgument:&param atIndex:i];
            }
                break;
            case 'f':
            case 'd':
            {
                float param = va_arg(args, double);
                [invocation setArgument:&param atIndex:i];
            }
                break;
            case '@':
            {
                id param = va_arg(args, id);
                [invocation setArgument:&param atIndex:i];
            }
                break;
            case '^':
            {
                void * param = va_arg(args, void *);
                [invocation setArgument:&param atIndex:i];
            }
                break;
            ...
            default:
                if (strcmp(argType, @encode(CGSize))) {
                    CGSize param = va_arg(args, CGSize);
                    [invocation setArgument:&param atIndex:i];
                }
                else ... {

                }
                break;
        }
    }

}

看上去不错，但是等等！NSInvocation 的 target 是 forwardingBlock，但 selector 应该是什么？没有 selector 是无法动态找到 IMP 调用的。后来搜索发现 NSInvocation 有一个私用方法 - (void)invokeUsingIMP:(IMP)imp，天无绝人之路！

[invocation invokeUsingIMP:((RTBlock *)(((__bridge RTBlock *)block)->forwardingBlock))->invoke];

再等等，Block 中有自定义的 struct 时怎么办？union 呢？情况会有点复杂了。

forwardInvocation

自定义的 struct 问题一直没有好的解决方案，事情就卡在这里了，直到有一天搜索发现一个 block forwarding 的办法，瞬间眼前一亮，原来还有这种操作！

简单的说就是把 Block 的 invoke 函数指向 _objc_msgForward，利用 Objective-C 自身的消息转发机制自动填充好参数，最终会走到 - (void)forwardInvocation:(NSInvocation *) 方法。

- (void)forwardInvocation:(NSInvocation *)anInvocation
{
    RTBlock *layout = (__bridge RTBlock *)self;
    layout->reserved = 1;

    [anInvocation setTarget:(__bridge id)layout->forwardingBlock];
    [anInvocation invokeUsingIMP:((RTBlock *)layout->forwardingBlock)->invoke];
}

block forwarding 项目存在一些问题，它只能处理立即被调用的 Block，持有一段时间后调用会有问题，不满足我的要求，不过最后的成品 RTBlockCallChecker 是基于它的实现思路。

成果

最后的成品在使用上非常简单，用一个宏包裹原始 Block，就好，无论是一个变量还是字面量。它是支持任意参数与返回类型的，而且它是类型敏感的，类型错误编译时可以报警。

void (^someBlockMustBeCalled)() = ^{  
   ...
};
// 原来的代码
[self passBlockToAMethod:someBlockMustBeCalled];
// 改为
[self passBlockToAMethod:RT_CHECK_BLOCK_CALLED(someBlockMustBeCalled)];

以上。