[Concurrency] Fix memory leak around mismatched refcounting in Concurrency

Author: kubamracek

include/swift/ABI/Task.h

@@ -48,7 +48,8 @@ class ContinuationAsyncContext;
 // _swift_concurrency_debug_internal_layout_version and add a comment describing
 // the new version.
 
-extern FullMetadata<DispatchClassMetadata> jobHeapMetadata;
+extern const HeapMetadata *jobHeapMetadataPtr;
+extern const HeapMetadata *taskHeapMetadataPtr;
 
 /// A schedulable job.
 class alignas(2 * alignof(void*)) Job :
@@ -106,14 +107,14 @@ class alignas(2 * alignof(void*)) Job :
   };
 
   Job(JobFlags flags, JobInvokeFunction *invoke,
-      const HeapMetadata *metadata = &jobHeapMetadata)
+      const HeapMetadata *metadata = jobHeapMetadataPtr)
       : HeapObject(metadata), Flags(flags), RunJob(invoke) {
     Voucher = voucher_copy();
     assert(!isAsyncTask() && "wrong constructor for a task");
   }
 
   Job(JobFlags flags, TaskContinuationFunction *invoke,
-      const HeapMetadata *metadata = &jobHeapMetadata,
+      const HeapMetadata *metadata = jobHeapMetadataPtr,
       bool captureCurrentVoucher = true)
       : HeapObject(metadata), Flags(flags), ResumeTask(invoke) {
     if (captureCurrentVoucher)

stdlib/public/Concurrency/CMakeLists.txt

@@ -273,7 +273,7 @@ if(SWIFT_SHOULD_BUILD_EMBEDDED_STDLIB AND SWIFT_SHOULD_BUILD_EMBEDDED_CONCURRENC
         -Xfrontend -emit-empty-object-file
         ${SWIFT_RUNTIME_CONCURRENCY_SWIFT_FLAGS}
       C_COMPILE_FLAGS
-        ${extra_c_compile_flags} ${SWIFT_RUNTIME_CONCURRENCY_C_FLAGS} -DSWIFT_CONCURRENCY_EMBEDDED=1
+        ${extra_c_compile_flags} ${SWIFT_RUNTIME_CONCURRENCY_C_FLAGS} -DSWIFT_CONCURRENCY_EMBEDDED=1 -DSWIFT_RUNTIME_EMBEDDED=1
       MODULE_DIR "${CMAKE_BINARY_DIR}/lib/swift/embedded"
       SDK "embedded"
       ARCHITECTURE "${mod}"

stdlib/public/Concurrency/ExecutorChecks.cpp

@@ -54,7 +54,11 @@ static_assert((SwiftJobPriority)swift::JobPriority::Unspecified
 // Job (has additional fields not exposed via SwiftJob)
 static_assert(sizeof(swift::Job) >= sizeof(SwiftJob));
 static_assert(offsetof(swift::Job, metadata) == offsetof(SwiftJob, metadata));
+#if !SWIFT_CONCURRENCY_EMBEDDED
 static_assert(offsetof(swift::Job, refCounts) == offsetof(SwiftJob, refCounts));
+#else
+static_assert(offsetof(swift::Job, embeddedRefcount) == offsetof(SwiftJob, refCounts));
+#endif
 static_assert(offsetof(swift::Job, SchedulerPrivate) == offsetof(SwiftJob, schedulerPrivate));
 static_assert(offsetof(swift::Job, SchedulerPrivate[0]) == offsetof(SwiftJob, schedulerPrivate[0]));
 static_assert(offsetof(swift::Job, SchedulerPrivate[1]) == offsetof(SwiftJob, schedulerPrivate[1]));

stdlib/public/Concurrency/Task.cpp

@@ -215,7 +215,6 @@ InitialTaskExecutorOwnedPreferenceTaskOptionRecord::getExecutorRefFromUnownedTas
     return executorRef;
 }
 
-
 void NullaryContinuationJob::process(Job *_job) {
   auto *job = cast<NullaryContinuationJob>(_job);
 
@@ -393,7 +392,9 @@ static void jobInvoke(void *obj, void *unused, uint32_t flags) {
 // Magic constant to identify Swift Job vtables to Dispatch.
 static const unsigned long dispatchSwiftObjectType = 1;
 
-FullMetadata<DispatchClassMetadata> swift::jobHeapMetadata = {
+#if !SWIFT_CONCURRENCY_EMBEDDED
+
+FullMetadata<DispatchClassMetadata> jobHeapMetadata = {
   {
     {
       /*type layout*/ nullptr,
@@ -432,10 +433,35 @@ static FullMetadata<DispatchClassMetadata> taskHeapMetadata = {
   }
 };
 
+#else // SWIFT_CONCURRENCY_EMBEDDED
+
+// This matches the embedded class metadata layout in IRGen and in
+// EmbeddedRuntime.swift.
+typedef struct EmbeddedClassMetadata {
+  void *superclass;
+  HeapObjectDestroyer *__ptrauth_swift_heap_object_destructor destroy;
+  void *ivar_destroyer;
+} EmbeddedHeapObject;
+
+static EmbeddedClassMetadata jobHeapMetadata = {
+  0, &destroyJob, 0,
+};
+
+static EmbeddedClassMetadata taskHeapMetadata = {
+  0, &destroyTask, 0,
+};
+
+#endif
+
 const void *const swift::_swift_concurrency_debug_jobMetadata =
-    static_cast<Metadata *>(&jobHeapMetadata);
+    (Metadata *)(&jobHeapMetadata);
 const void *const swift::_swift_concurrency_debug_asyncTaskMetadata =
-    static_cast<Metadata *>(&taskHeapMetadata);
+    (Metadata *)(&taskHeapMetadata);
+
+const HeapMetadata *swift::jobHeapMetadataPtr =
+    (HeapMetadata *)(&jobHeapMetadata);
+const HeapMetadata *swift::taskHeapMetadataPtr =
+    (HeapMetadata *)(&taskHeapMetadata);
 
 static void completeTaskImpl(AsyncTask *task,
                              AsyncContext *context,
@@ -959,14 +985,14 @@ swift_task_create_commonImpl(size_t rawTaskCreateFlags,
   if (asyncLet) {
     // Initialize the refcount bits to "immortal", so that
     // ARC operations don't have any effect on the task.
-    task = new(allocation) AsyncTask(&taskHeapMetadata,
-                             InlineRefCounts::Immortal, jobFlags,
-                             function, initialContext,
-                             captureCurrentVoucher);
+    task = new (allocation)
+        AsyncTask(reinterpret_cast<ClassMetadata *>(&taskHeapMetadata),
+                  InlineRefCounts::Immortal, jobFlags, function, initialContext,
+                  captureCurrentVoucher);
   } else {
-    task = new(allocation) AsyncTask(&taskHeapMetadata, jobFlags,
-                                    function, initialContext,
-                                    captureCurrentVoucher);
+    task = new (allocation)
+        AsyncTask(reinterpret_cast<ClassMetadata *>(&taskHeapMetadata), jobFlags,
+                  function, initialContext, captureCurrentVoucher);
   }
 
   // Initialize the child fragment if applicable.

stdlib/public/SwiftShims/swift/shims/HeapObject.h

@@ -53,6 +53,8 @@ struct HeapObject {
   /// This is always a valid pointer to a metadata object.
   HeapMetadata const *__ptrauth_objc_isa_pointer metadata;
 
+#if !SWIFT_RUNTIME_EMBEDDED
+
   SWIFT_HEAPOBJECT_NON_OBJC_MEMBERS;
 
 #ifndef __swift__
@@ -70,11 +72,40 @@ struct HeapObject {
   : metadata(newMetadata)
   , refCounts(InlineRefCounts::Immortal)
   { }
+#endif // __swift__
+
+#else // SWIFT_RUNTIME_EMBEDDED
+  uintptr_t embeddedRefcount;
+
+  // Note: The immortal refcount value is also hard-coded in IRGen in
+  // `irgen::emitConstantObject`, and in EmbeddedRuntime.swift.
+#if __POINTER_WIDTH__ == 64
+  static const uintptr_t EmbeddedImmortalRefCount = 0x7fffffffffffffffull;
+#elif __POINTER_WIDTH__ == 32
+  static const uintptr_t EmbeddedImmortalRefCount = 0x7fffffff;
+#elif __POINTER_WIDTH__ == 16
+  static const uintptr_t EmbeddedImmortalRefCount = 0x7fff;
+#endif
 
+#ifndef __swift__
+  HeapObject() = default;
+
+  // Initialize a HeapObject header as appropriate for a newly-allocated object.
+  constexpr HeapObject(HeapMetadata const *newMetadata)
+      : metadata(newMetadata), embeddedRefcount(1) {}
+
+  // Initialize a HeapObject header for an immortal object
+  constexpr HeapObject(HeapMetadata const *newMetadata,
+                       InlineRefCounts::Immortal_t immortal)
+      : metadata(newMetadata), embeddedRefcount(EmbeddedImmortalRefCount) {}
+#endif // __swift__
+
+#endif // SWIFT_RUNTIME_EMBEDDED
+
+#ifndef __swift__
 #ifndef NDEBUG
   void dump() const SWIFT_USED;
 #endif
-
 #endif // __swift__
 };
 

stdlib/public/core/EmbeddedRuntime.swift

@@ -99,7 +99,7 @@ public struct HeapObject {
   // to think about supporting (or banning) weak and/or unowned references.
   var refcount: Int
 
-  // Note: The immortalRefCount value is also hard-coded in IRGen in `irgen::emitConstantObject`.
+  // Note: The immortalRefCount value is also hard-coded in IRGen in `irgen::emitConstantObject`, and in HeapObject.h.
 #if _pointerBitWidth(_64)
   static let doNotFreeBit     = Int(bitPattern: 0x8000_0000_0000_0000)
   static let refcountMask     = Int(bitPattern: 0x7fff_ffff_ffff_ffff)

test/embedded/Inputs/debug-malloc.c

@@ -1,22 +1,26 @@
 #include <stdlib.h>
 #include <stdio.h>
 
-#define HEAP_SIZE (8 * 1024)
+#define HEAP_SIZE (128 * 1024)
 
 __attribute__((aligned(16)))
 char heap[HEAP_SIZE] = {};
 
 size_t next_heap_index = 0;
 
 void *calloc(size_t count, size_t size) {
-  printf("malloc(%ld)", count);
+  size_t total_size = count * size;
+  printf("malloc(%ld)", total_size);
 
-  if (next_heap_index + count * size > HEAP_SIZE) {
+  if (total_size % 16 != 0)
+    total_size = total_size + (16 - total_size % 16);
+
+  if (next_heap_index + total_size > HEAP_SIZE) {
     puts("HEAP EXHAUSTED\n");
     __builtin_trap();
   }
   void *p = &heap[next_heap_index];
-  next_heap_index += count * size;
+  next_heap_index += total_size;
   printf("-> %p\n", p);
   return p;
 }

test/embedded/concurrency-leaks.swift

@@ -0,0 +1,46 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend -enable-experimental-feature Embedded -parse-as-library %s -c -o %t/a.o -g -O
+// RUN: %target-clang -x c -std=c11 -c %S/Inputs/debug-malloc.c -o %t/debug-malloc.o -g
+// RUN: %target-clang %t/a.o %t/debug-malloc.o -o %t/a.out -L%swift_obj_root/lib/swift/embedded/%target-cpu-apple-macos -lswift_Concurrency -lswift_ConcurrencyDefaultExecutor -dead_strip -g
+// RUN: %target-run %t/a.out | %FileCheck %s
+
+// REQUIRES: executable_test
+// REQUIRES: swift_in_compiler
+// REQUIRES: optimized_stdlib
+// REQUIRES: OS=macosx
+
+import _Concurrency
+
+@main
+struct Main {
+  static func main() async {
+    print("start")
+    // CHECK: start
+    do {
+        let x = Task {
+            return 42
+        }
+        _ = await x.value
+    }
+    // There should be exactly 5 allocations involved:
+    // - 2x 32 bytes ... closure context for swift_task_create, closure object to pass to Task.init
+    // - 1x 320 bytes ... malloc(amountToAllocate) in swift_task_create_common for the Task heap object itself
+    // - 1x 1016 bytes ... the initial StackAllocator slab in the task-specific allocator
+    // - 1x 40 bytes ... task status record
+    // Check that they are all accounted for and free'd.
+
+    // CHECK: malloc({{[0-9]+}})-> [[M1:0x[0-9a-f]+]]
+    // CHECK: malloc({{[0-9]+}})-> [[M2:0x[0-9a-f]+]]
+    // CHECK: malloc({{[0-9]+}})-> [[M3:0x[0-9a-f]+]]
+    // CHECK: malloc({{[0-9]+}})-> [[M4:0x[0-9a-f]+]]
+    // CHECK: free([[M1]])
+    // CHECK: free([[M2]])
+    // CHECK: malloc({{[0-9]+}})-> [[M5:0x[0-9a-f]+]]
+    // CHECK: free([[M5]])
+    // CHECK: free([[M4]])
+    // CHECK: free([[M3]])
+
+    print("done")
+    // CHECK: done
+  }
+}