如何判断block回调未被调用

2018-03-20

在处理异步过程中，我们经常会碰到这种情况，需要异步处理并异步回调completionHandler，但是有些场景下，如果你在处理完异步逻辑，而不回调completion的时候，会产生逻辑上的bug或者内存泄露问题，那么我们就需要知道调用方是否调用了completion。

这里举几个比较典型的例子，比如WKUIDelegate中的回调：

-                     (void)webView:(WKWebView *)webView
 runJavaScriptAlertPanelWithMessage:(NSString *)message
                   initiatedByFrame:(WKFrameInfo *)frame
                  completionHandler:(void (^)(void))completionHandler;

如果不回调其completionHandler，会导致其逻辑上的错误，那么这里我们来看看如何动态监测completionHandler是否被调用过。

这里说一下，WK是通过WTF的C++模板来实现的，我这里采用C语言来实现，其思路是大致相同的。

Block

首先我们来看看Block是什么。虽然我们平时可以像OC对象那样去使用它，但它严格意义上来说并不是一个OC对象，或者说它是一中极为特殊的OC对象。

struct Block_layout {
    void *isa;
    volatile int32_t flags; // contains ref count
    int32_t reserved;
    void (*invoke)(void *, ...);
    Descriptor *descriptor;
    // imported variables
};
struct Descriptor {
    uintptr_t reserved;
    uintptr_t size;
    void (*copy)(void *dst, const void *src);
    void (*dispose)(const void *);
};

上面就是Block的内存布局，其中Block_layout是一个不定长的结构体，我们平时看到的捕获变量都会存在结构尾部。这里我们看到和OC对象一样，也有isa指针，但是这里的指针永远只会指向几个地方，这个之后会说。

其实我们在调用Block的时候，实际上调用的是block->invoke()，第一个参数是Block本身，然后是入参按顺序排下去，这一部分编译器都会给我们做好，所以一个block调用实际是这样的：

1	block->invoke(block, arg1, arg2, arg3);

可以看到和OC的objc_msgSend方法相同的是第一个参数是对象本身，但是不同的是第二个参数不再是SEL。

既然知道了Block的结构，那么我们就可以自定义block了。

Block类型

Block定义的类型有：

BLOCK_EXPORT void * _NSConcreteGlobalBlock[32]
    __OSX_AVAILABLE_STARTING(__MAC_10_6, __IPHONE_3_2);
BLOCK_EXPORT void * _NSConcreteStackBlock[32]
    __OSX_AVAILABLE_STARTING(__MAC_10_6, __IPHONE_3_2);

BLOCK_EXPORT void * _NSConcreteMallocBlock[32]
    __OSX_AVAILABLE_STARTING(__MAC_10_6, __IPHONE_3_2);
BLOCK_EXPORT void * _NSConcreteAutoBlock[32]
    __OSX_AVAILABLE_STARTING(__MAC_10_6, __IPHONE_3_2);
BLOCK_EXPORT void * _NSConcreteFinalizingBlock[32]
    __OSX_AVAILABLE_STARTING(__MAC_10_6, __IPHONE_3_2);
BLOCK_EXPORT void * _NSConcreteWeakBlockVariable[32]
    __OSX_AVAILABLE_STARTING(__MAC_10_6, __IPHONE_3_2);

其中只有前2种是公开的，而我们平时会碰到的基本都是前3种类型，其中Global是永远不会被释放的，Stack是在栈上，所以只要栈销毁了就会被释放，Malloc和普通OC对象一样，采用引用计数来决定生命周期的。

那么我们回到最初的目的，如何判断是否被调用了呢？因为这个调用有可能是异步的，所以不可能通过__block bool called这样的临时对象来判断，也不能通过其是否由Stack拷贝成Malloc来判断，因为copy了并不一定会被调用。

Block Wrap

这里要判断Block是否被调用，肯定是需要在原始Block基础上包裹一层可以计数调用次数的Block。C++会方便的多，可以直接通过模板来构造一个签名一样的Block。

这里我们利用了MallocBlock在未被任何人引用的时候会销毁的特性，在其被释放之前，来监测计数是否为0。如果是0则说明从来没有被调用过，不是0则说明被调用了。

那么接下来我们来看看如何动态构建这样一个Block，以及如果去包裹其实现体。

动态构建Block

struct Block_layout {
    void *isa;
    volatile int32_t flags; // contains ref count
    int32_t reserved;
    void (*invoke)(void);
    void *descriptor;
    
    // imported variables
    void *block;
    int64_t called;
    char *message;
};

首先我们将我们所需要的几个参数定义在Block末尾，分别是原始的Block，调用计数，以及错误信息（这个在报错的时候使用，和该方案关系不大）。

然后，我们需要定义自己的descriptor。这里重写了dispose方法，我们需要在这里判断是否计数为0，同时也要在这里将对象释放掉（由于在C环境中，所以block也需要手动将其释放）。

void block_call_assert_wrap_dispose(const void * ptr) {
    struct Block_layout *self = (struct Block_layout *)ptr;
    if (!((struct Block_layout *)ptr)->called) {
        if (exception_handler) {
            if (self->message) {
                char *buf = (char *)malloc((strlen(self->message) + 64) * sizeof(char));
                sprintf(buf, "ERROR: Block must be called at %s!\n", self->message);
                exception_handler(buf);
                free(buf);
            }
            else {
                exception_handler("ERROR: Block must be called at %s!\n");
            }
        }
    }
    Block_release(self->block);
    if (self->message) free(self->message);
}
static const struct Descriptor descriptor = {
    0,
    sizeof(struct Block_layout),
    NULL,
    block_call_assert_wrap_dispose
};

接下来就是将我们的所有数据内容填入Block_layout，来合成一个Block对象。

void *block_call_assert_wrap_block(void *orig_blk, char *message) {
    struct Block_layout *block = (struct Block_layout *)malloc(sizeof(struct Block_layout));
    block->isa = _NSConcreteMallocBlock;
    
    enum {
        BLOCK_NEEDS_FREE = (1 << 24),
        BLOCK_HAS_COPY_DISPOSE = (1 << 25),
    };
    const unsigned retainCount = 1;
    
    block->flags = BLOCK_HAS_COPY_DISPOSE | BLOCK_NEEDS_FREE | (retainCount << 1);
    block->reserved = 0;
    block->invoke = (void (*)(void))block_call_assert_wrap_invoke;
    block->descriptor = (void *)&descriptor;
    
    block->block = (void *)Block_copy(orig_blk);
    block->called = 0;
    
    size_t len = strlen(message)*sizeof(char);
    char *buf = (char *)malloc(len);
    memcpy(buf, message, len);
    block->message = buf;
    
    return block;
}

其中invoke方法被我们的新方法block_call_assert_wrap_invoke所替换，在这个方法里面，会更新计数，并且调用原始block的invoke方法。

block_call_assert_wrap_invoke的实现

block的方法是非常灵活的，参数个数以及返回值不一样的时候，经过前几篇内容，我们知道不能简单的通过方法调用来实现参数的传递，而且在这里我们也无法知道参数的个数以及类型。那么我们要怎么做才能简单而又实用呢？

这时候，我们想到objc_msgSend方法，它就实现了非常技巧的实现了arguments forward的功能（其功能特性可以参考C++模板的多参传递template <typename Args...>）。

由于这里找不到i386的系统和arm32的系统了，所以只给出x86_64和arm64的实现方案。

#if __x86_64__

.align 4

.global _block_call_assert_wrap_invoke
_block_call_assert_wrap_invoke:

mov  %rdi, %r10

movq $1, 0x28(%r10)         // called

movq 0x20(%r10), %r11       // block
movq %r11, %rdi
movq 0x10(%r11), %r11        // block->block->invoke

jmp *%r11

#endif

#ifdef __arm64__
.align 4

.global _block_call_assert_wrap_invoke
_block_call_assert_wrap_invoke:

mov x9, x0
add x10, x9, #0x20   // &block
add x11, x9, #0x28   // called

mov x12, #1
str x12, [x11]

ldr x12, [x10]        // block
add x12, x12, #0x10 // block->invoke
ldr x12, [x12]
mov x0, x11

br x12
ret
#endif

这里简单的说明一下段汇编的逻辑。

取出block->called，并置为1（可能改为真正的计数会比较好）。
取出原始block block->block，并放到第一个参数位置。
调用原始block的invoke call block->block->invoke。

这样我们就非常简单的包裹了原始invoke方法，并且插入了自己的逻辑。

使用

首先我们需要设置上述的exception_handler。

void exception_log(const char *str) {
    NSLog(@"%s", str);
}
block_call_assert_set_exception_handler(exception_log);

这里我只是让他打印出错误，更好的应该是直接抛出异常[NSException raise:]。

在此基础上，定义一个宏以方便使用，以及可以加入#if DEBUG，来禁用线上环境的该功能，并且把当前的位置传递给exception_message：

#define BLOCK_CALL_ASSERT(x) ({                 \
    typeof ((x)) blk = x;                       \
    char *message = (char *)malloc(512);        \
    memset(message, 0, 512);                    \
    sprintf(message, "(%s:%d %s)", __FILE__, __LINE__, __FUNCTION__); \
    typeof (blk) ret = (__bridge_transfer typeof(blk))block_call_assert_wrap_block((__bridge void *)blk, message); \
    free(message);                              \
    ret;                                        \
})

bridge，恩我们是支持的ARC，所以在此为了防止类型转换的warning和error，在此使用宏来定义。（好像Objc++会有警告）

那么在使用的时候就是这样：

- (void)doAsyncWithCompletion:(block_t)completionBlock {
  dispatch_async(..., ^{
      completionBlock(...)
  });
}

[self doAsyncWithCompletion:BLOCK_CALL_ASSERT(^{
    do_after_completion();
    do_clear();
})];

那么在此时，如果被调用者没有调用过completionBlock()时，就会触发exception_handler。这样我们就可以检测到是否出现可能的逻辑错误和内存泄露了。

1	ERROR: Block must be called at (BlockCallAssert/BlockCallAssert/BlockCallAssert/ViewController.mm:41 -[ViewController test2])!

最后

一般来说，我们一旦设计了包含completionBlock这样的接口，基本是需要回调方100%的回调的，如果可以不用回调，那么我们为什么不改变设计方案呢。

当我们的调用方是自己的时候，我们可以确保，而如果是SDK，我们就很难确保，文档这个东西是不靠谱的，那么我们就让调用方在忽略了回调的时候给他一个重拳吧（exception）。

这个方案的实现我放在github，和cocoaPods BlockCallAssert。