Treat allocations within dispatch_once blocks as global allocations

[fjricci]

Summary

Objective-C programs often treat dispatch_once blocks similarly to global
variables, since they are only executed once over the course of the process.
They are often used to initialize singletons and other process state.

For an example of a simple test case which reveals this behavior in darwin system
libraries:

// clang -fsanitize=leak -framework Foundation test.m
#import <Foundation/Foundation.h>

int main() {
    @autoreleasepool { id defaults = [NSUserDefaults standardUserDefaults]; };
    return 0;
}

This patch treats dispatch_once allocations as global variables, ignoring
them if use_globals is enabled.

[fjricci]

An alternative solution here would be to add a suppression for dispatch_once instead, but that would be less configurable (unless we required the user to provide it) and slightly less performant.

[fjricci]

This doesn't provide as much of a performance benefit as I expected. In addition, the fact that a leak comes from a system library doesn't make it less of a leak. It's not a great user experience to see system library leaks in your program output, but I don't think there is much we can/should do about it.

lib/asan/asan_mac.cc

 #include "sanitizer_common/sanitizer_platform.h"
 #if SANITIZER_MAC
 
 #include "asan_interceptors.h"
 #include "asan_internal.h"
 #include "asan_mapping.h"
 #include "asan_stack.h"
 #include "asan_thread.h"
+#include "lsan/lsan_common.h"
 #include "sanitizer_common/sanitizer_atomic.h"
 #include "sanitizer_common/sanitizer_libc.h"
 #include "sanitizer_common/sanitizer_mac.h"
 
 #include <dlfcn.h>
 #include <fcntl.h>
 #include <libkern/OSAtomic.h>
 #include <mach-o/dyld.h>

 //   http://developer.apple.com/library/mac/#documentation/Performance/Reference/GCD_libdispatch_Ref/Reference/reference.html
 // The implementation details are at
 //   http://libdispatch.macosforge.org/trac/browser/trunk/src/queue.c
 
 typedef void* dispatch_group_t;
 typedef void* dispatch_queue_t;
 typedef void* dispatch_source_t;
 typedef u64 dispatch_time_t;
+typedef long dispatch_once_t;  // NOLINT
 typedef void (*dispatch_function_t)(void *block);
 typedef void* (*worker_t)(void *block);
 
 // A wrapper for the ObjC blocks used to support libdispatch.
 typedef struct {
   void *block;
   dispatch_function_t func;
   u32 parent_tid;
+#if CAN_SANITIZE_LEAKS
+  bool disable_leak_checks;
+#endif
 } asan_block_context_t;
 
 ALWAYS_INLINE
 void asan_register_worker_thread(int parent_tid, StackTrace *stack) {
   AsanThread *t = GetCurrentThread();
   if (!t) {
     t = AsanThread::Create(/* start_routine */ nullptr, /* arg */ nullptr,
                            parent_tid, stack, /* detached */ true);

 void asan_dispatch_call_block_and_release(void *block) {
   GET_STACK_TRACE_THREAD;
   asan_block_context_t *context = (asan_block_context_t*)block;
   VReport(2,
           "asan_dispatch_call_block_and_release(): "
           "context: %p, pthread_self: %p\n",
           block, pthread_self());
   asan_register_worker_thread(context->parent_tid, &stack);
+#if CAN_SANITIZE_LEAKS
+  if (context->disable_leak_checks) __lsan::DisableInThisThread();
+#endif
   // Call the original dispatcher for the block.
   context->func(context->block);
+#if CAN_SANITIZE_LEAKS
+  if (context->disable_leak_checks) __lsan::EnableInThisThread();
+#endif
   asan_free(context, &stack, FROM_MALLOC);
 }
 
 }  // namespace __asan
 
 using namespace __asan;  // NOLINT
 
 // Wrap |ctxt| and |func| into an asan_block_context_t.
 // The caller retains control of the allocated context.
 extern "C"
 asan_block_context_t *alloc_asan_context(void *ctxt, dispatch_function_t func,
                                          BufferedStackTrace *stack) {
   asan_block_context_t *asan_ctxt =
       (asan_block_context_t*) asan_malloc(sizeof(asan_block_context_t), stack);
   asan_ctxt->block = ctxt;
   asan_ctxt->func = func;
   asan_ctxt->parent_tid = GetCurrentTidOrInvalid();
+#if CAN_SANITIZE_LEAKS
+  asan_ctxt->disable_leak_checks = false;
+#endif
   return asan_ctxt;
 }
 
 // Define interceptor for dispatch_*_f function with the three most common
 // parameters: dispatch_queue_t, context, dispatch_function_t.
 #define INTERCEPT_DISPATCH_X_F_3(dispatch_x_f)                                \
   INTERCEPTOR(void, dispatch_x_f, dispatch_queue_t dq, void *ctxt,            \
                                   dispatch_function_t func) {                 \
if (Verbosity() >= 2) {
     Report("dispatch_group_async_f(): context: %p, pthread_self: %p\n",
            asan_ctxt, pthread_self());
     PRINT_CURRENT_STACK();
   }
   REAL(dispatch_group_async_f)(group, dq, (void*)asan_ctxt,
                                asan_dispatch_call_block_and_release);
 }
 
+INTERCEPTOR(void, dispatch_once_f, dispatch_once_t *pred, void *ctxt,
+            dispatch_function_t func) {
+  GET_STACK_TRACE_THREAD;
+  asan_block_context_t *asan_ctxt = alloc_asan_context(ctxt, func, &stack);
+  if (Verbosity() >= 2) {
+    Report("dispatch_once_f(): context: %p, pthread_self: %p\n", asan_ctxt,
+           pthread_self());
+    PRINT_CURRENT_STACK();
+  }
+#if CAN_SANITIZE_LEAKS
+  asan_ctxt->disable_leak_checks = __lsan::flags()->use_globals;
+#endif
+  REAL(dispatch_once_f)
+  (pred, (void *)asan_ctxt, asan_dispatch_call_block_and_release);
+}
+
 #if !defined(MISSING_BLOCKS_SUPPORT)
 extern "C" {
 void dispatch_async(dispatch_queue_t dq, void(^work)(void));
 void dispatch_group_async(dispatch_group_t dg, dispatch_queue_t dq,
                           void(^work)(void));
 void dispatch_after(dispatch_time_t when, dispatch_queue_t queue,
                     void(^work)(void));
 void dispatch_source_set_cancel_handler(dispatch_source_t ds,
                                         void(^work)(void));
 void dispatch_source_set_event_handler(dispatch_source_t ds, void(^work)(void));
+void dispatch_once(dispatch_once_t *pred, void (^work)(void));
 }
 
 #define GET_ASAN_BLOCK(work) \
   void (^asan_block)(void);  \
   int parent_tid = GetCurrentTidOrInvalid(); \
   asan_block = ^(void) { \
     GET_STACK_TRACE_THREAD; \
     asan_register_worker_thread(parent_tid, &stack); \

 }
 
 INTERCEPTOR(void, dispatch_source_set_event_handler,
             dispatch_source_t ds, void(^work)(void)) {
   ENABLE_FRAME_POINTER;
   GET_ASAN_BLOCK(work);
   REAL(dispatch_source_set_event_handler)(ds, asan_block);
 }
+
+// allocations inside dispatch_once blocks are often used to initialize program
+// singletons, and can be treated as global allocations, as they are only
+// executed once per program.
+INTERCEPTOR(void, dispatch_once, dispatch_once_t *pred, void (^work)(void)) {
+  ENABLE_FRAME_POINTER;
+  void (^global_block)(void) = work;
+#if CAN_SANITIZE_LEAKS
+  if (__lsan::flags()->use_globals) {
+    global_block = ^(void) {
+      __lsan::ScopedInterceptorDisabler disabler;
+      work();
+    };
+  }
+#endif
+  GET_ASAN_BLOCK(global_block);
+  REAL(dispatch_once)(pred, asan_block);
+}
 #endif
 
 #endif  // SANITIZER_MAC

lib/lsan/lsan_mac.cc

 //===----------------------------------------------------------------------===//
 
 #include "sanitizer_common/sanitizer_platform.h"
 #if SANITIZER_MAC
 
 #include "interception/interception.h"
 #include "lsan.h"
 #include "lsan_allocator.h"
+#include "lsan_common.h"
 #include "lsan_thread.h"
 
 #include <pthread.h>
 
 namespace __lsan {
 // Support for the following functions from libdispatch on Mac OS:
 //   dispatch_async_f()
 //   dispatch_async()

 //   http://developer.apple.com/library/mac/#documentation/Performance/Reference/GCD_libdispatch_Ref/Reference/reference.html
 // The implementation details are at
 //   http://libdispatch.macosforge.org/trac/browser/trunk/src/queue.c
 
 typedef void *dispatch_group_t;
 typedef void *dispatch_queue_t;
 typedef void *dispatch_source_t;
 typedef u64 dispatch_time_t;
+typedef long dispatch_once_t;  // NOLINT
 typedef void (*dispatch_function_t)(void *block);
 typedef void *(*worker_t)(void *block);
 
 // A wrapper for the ObjC blocks used to support libdispatch.
 typedef struct {
   void *block;
   dispatch_function_t func;
   u32 parent_tid;
+  bool disable_leak_checks;
 } lsan_block_context_t;
 
 ALWAYS_INLINE
 void lsan_register_worker_thread(int parent_tid) {
   if (GetCurrentThread() == kInvalidTid) {
     u32 tid = ThreadCreate(parent_tid, 0, true);
     ThreadStart(tid, GetTid());
     SetCurrentThread(tid);
   }
 }
 
 // For use by only those functions that allocated the context via
 // alloc_lsan_context().
 extern "C" void lsan_dispatch_call_block_and_release(void *block) {
   lsan_block_context_t *context = (lsan_block_context_t *)block;
   VReport(2,
           "lsan_dispatch_call_block_and_release(): "
           "context: %p, pthread_self: %p\n",
           block, pthread_self());
   lsan_register_worker_thread(context->parent_tid);
+  if (context->disable_leak_checks) DisableInThisThread();
   // Call the original dispatcher for the block.
   context->func(context->block);
+  if (context->disable_leak_checks) EnableInThisThread();
   lsan_free(context);
 }
 
 }  // namespace __lsan
 
 using namespace __lsan;  // NOLINT
 
 // Wrap |ctxt| and |func| into an lsan_block_context_t.
 // The caller retains control of the allocated context.
 extern "C" lsan_block_context_t *alloc_lsan_context(void *ctxt,
                                                     dispatch_function_t func) {
   GET_STACK_TRACE_THREAD;
   lsan_block_context_t *lsan_ctxt =
       (lsan_block_context_t *)lsan_malloc(sizeof(lsan_block_context_t), stack);
   lsan_ctxt->block = ctxt;
   lsan_ctxt->func = func;
   lsan_ctxt->parent_tid = GetCurrentThread();
+  lsan_ctxt->disable_leak_checks = false;
   return lsan_ctxt;
 }
 
 // Define interceptor for dispatch_*_f function with the three most common
 // parameters: dispatch_queue_t, context, dispatch_function_t.
 #define INTERCEPT_DISPATCH_X_F_3(dispatch_x_f)                        \
   INTERCEPTOR(void, dispatch_x_f, dispatch_queue_t dq, void *ctxt,    \
               dispatch_function_t func) {                             \

 
 INTERCEPTOR(void, dispatch_group_async_f, dispatch_group_t group,
             dispatch_queue_t dq, void *ctxt, dispatch_function_t func) {
   lsan_block_context_t *lsan_ctxt = alloc_lsan_context(ctxt, func);
   REAL(dispatch_group_async_f)
   (group, dq, (void *)lsan_ctxt, lsan_dispatch_call_block_and_release);
 }
 
+INTERCEPTOR(void, dispatch_once_f, dispatch_once_t *pred, void *ctxt,
+            dispatch_function_t func) {
+  lsan_block_context_t *lsan_ctxt = alloc_lsan_context(ctxt, func);
+  lsan_ctxt->disable_leak_checks = flags()->use_globals;
+  return REAL(dispatch_once_f)(pred, (void *)lsan_ctxt,
+                               lsan_dispatch_call_block_and_release);
+}
+
 #if !defined(MISSING_BLOCKS_SUPPORT)
 extern "C" {
 void dispatch_async(dispatch_queue_t dq, void (^work)(void));
 void dispatch_group_async(dispatch_group_t dg, dispatch_queue_t dq,
                           void (^work)(void));
 void dispatch_after(dispatch_time_t when, dispatch_queue_t queue,
                     void (^work)(void));
 void dispatch_source_set_cancel_handler(dispatch_source_t ds,
                                         void (^work)(void));
 void dispatch_source_set_event_handler(dispatch_source_t ds,
                                        void (^work)(void));
+void dispatch_once(dispatch_once_t *pred, void (^work)(void));
 }
 
 #define GET_LSAN_BLOCK(work)                 \
   void (^lsan_block)(void);                  \
   int parent_tid = GetCurrentThread();       \
   lsan_block = ^(void) {                     \
     lsan_register_worker_thread(parent_tid); \
     work();                                  \
INTERCEPTOR(void, dispatch_source_set_cancel_handler, dispatch_source_t ds,
   REAL(dispatch_source_set_cancel_handler)(ds, lsan_block);
 }
 
 INTERCEPTOR(void, dispatch_source_set_event_handler, dispatch_source_t ds,
             void (^work)(void)) {
   GET_LSAN_BLOCK(work);
   REAL(dispatch_source_set_event_handler)(ds, lsan_block);
 }
+
+// allocations inside dispatch_once blocks are often used to initialize program
+// singletons, and can be treated as global allocations, as they are only
+// executed once per program.
+INTERCEPTOR(void, dispatch_once, dispatch_once_t *pred, void (^work)(void)) {
+  void (^global_block)(void);
+  if (flags()->use_globals) {
+    global_block = ^(void) {
+      ScopedInterceptorDisabler disabler;
+      work();
+    };
+  } else {
+    global_block = work;
+  }
+  GET_LSAN_BLOCK(global_block);
+  REAL(dispatch_once)(pred, lsan_block);
+}
 #endif
 
 #endif  // SANITIZER_MAC

test/lsan/TestCases/Darwin/blocks.mm

+// Test for dispatch_once leaks
+// RUN: LSAN_BASE="report_objects=1:use_stacks=0:use_registers=0"
+// RUN: %clangxx_lsan %s -o %t -framework Foundation
+// RUN: %env_lsan_opts=$LSAN_BASE:use_globals=0 not %run %t 2>&1 | FileCheck %s
+// RUN: %env_lsan_opts=$LSAN_BASE:use_globals=1 %run %t 2>&1
+
+#include <dispatch/dispatch.h>
+#include <stdlib.h>
+
+#include "sanitizer_common/print_address.h"
+
+#define NUM_LOOPS 5
+
+int main() {
+  __block void *singleton;
+  static dispatch_once_t once_token = 0;
+  dispatch_block_t block = ^{
+    singleton = malloc(1337);
+    print_address("Test alloc: ", 1, singleton);
+  };
+
+  for (size_t i = 0; i < NUM_LOOPS; ++i) {
+    dispatch_once(&once_token, block);
+  }
+  return 0;
+}
+
+// CHECK: Test alloc: [[addr:0x[0-9,a-f]+]]
+// CHECK: LeakSanitizer: detected memory leaks
+// CHECK: [[addr]] (1337 bytes)
+// CHECK: SUMMARY: {{(Leak|Address)}}Sanitizer: 1337 byte(s) leaked in 1 allocation(s)