From 8666aef39f403da08daabcb5c83e78316f17f3bc Mon Sep 17 00:00:00 2001 From: Kuba Mracek Date: Fri, 18 Oct 2024 23:19:40 -0700 Subject: [PATCH] [embedded] Fix a memory leak caused by incorrect refcounting logic around doNotFreeBit --- lib/IRGen/GenConstant.cpp | 3 +- stdlib/public/core/EmbeddedRuntime.swift | 29 ++++++++++-------- test/embedded/Inputs/debug-malloc.c | 38 ++++++++++++++++++++++++ test/embedded/refcounting-leak.swift | 35 ++++++++++++++++++++++ 4 files changed, 91 insertions(+), 14 deletions(-) create mode 100644 test/embedded/Inputs/debug-malloc.c create mode 100644 test/embedded/refcounting-leak.swift diff --git a/lib/IRGen/GenConstant.cpp b/lib/IRGen/GenConstant.cpp index 84eb7553e66..871b64fbad2 100644 --- a/lib/IRGen/GenConstant.cpp +++ b/lib/IRGen/GenConstant.cpp @@ -458,7 +458,8 @@ llvm::Constant *irgen::emitConstantObject(IRGenModule &IGM, ObjectInst *OI, if (IGM.canMakeStaticObjectReadOnly(OI->getType())) { if (!IGM.swiftImmortalRefCount) { if (IGM.Context.LangOpts.hasFeature(Feature::Embedded)) { - // = HeapObject.immortalRefCount + // = HeapObject.immortalRefCount | HeapObject.doNotFreeBit + // 0xffff_ffff on 32-bit, 0xffff_ffff_ffff_ffff on 64-bit IGM.swiftImmortalRefCount = llvm::ConstantInt::get(IGM.IntPtrTy, -1); } else { IGM.swiftImmortalRefCount = llvm::ConstantExpr::getPtrToInt( diff --git a/stdlib/public/core/EmbeddedRuntime.swift b/stdlib/public/core/EmbeddedRuntime.swift index 49fbe55cdd6..a71880bb535 100644 --- a/stdlib/public/core/EmbeddedRuntime.swift +++ b/stdlib/public/core/EmbeddedRuntime.swift @@ -38,20 +38,21 @@ 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`. #if _pointerBitWidth(_64) - static let doNotFreeBit = Int(bitPattern: 0x8000_0000_0000_0000) - static let refcountMask = Int(bitPattern: 0x7fff_ffff_ffff_ffff) + static let doNotFreeBit = Int(bitPattern: 0x8000_0000_0000_0000) + static let refcountMask = Int(bitPattern: 0x7fff_ffff_ffff_ffff) + static let immortalRefCount = Int(bitPattern: 0x7fff_ffff_ffff_ffff) // Make sure we don't have doNotFreeBit set #elseif _pointerBitWidth(_32) - static let doNotFreeBit = Int(bitPattern: 0x8000_0000) - static let refcountMask = Int(bitPattern: 0x7fff_ffff) + static let doNotFreeBit = Int(bitPattern: 0x8000_0000) + static let refcountMask = Int(bitPattern: 0x7fff_ffff) + static let immortalRefCount = Int(bitPattern: 0x7fff_ffff) // Make sure we don't have doNotFreeBit set #elseif _pointerBitWidth(_16) - static let doNotFreeBit = Int(bitPattern: 0x8000) - static let refcountMask = Int(bitPattern: 0x7fff) + static let doNotFreeBit = Int(bitPattern: 0x8000) + static let refcountMask = Int(bitPattern: 0x7fff) + static let immortalRefCount = Int(bitPattern: 0x7fff) // Make sure we don't have doNotFreeBit set #endif - // Note: The immortalRefCount value of -1 is also hard-coded in IRGen in `irgen::emitConstantObject`. - static let immortalRefCount = -1 - #if _pointerBitWidth(_64) static let immortalObjectPointerBit = UInt(0x8000_0000_0000_0000) #endif @@ -158,7 +159,7 @@ public func swift_initStaticObject(metadata: Builtin.RawPointer, object: Builtin func swift_initStaticObject(metadata: UnsafeMutablePointer, object: UnsafeMutablePointer) -> UnsafeMutablePointer { _swift_embedded_set_heap_object_metadata_pointer(object, metadata) - object.pointee.refcount = HeapObject.immortalRefCount + object.pointee.refcount = HeapObject.immortalRefCount | HeapObject.doNotFreeBit return object } @@ -251,7 +252,7 @@ public func swift_retain_n(object: Builtin.RawPointer, n: UInt32) -> Builtin.Raw func swift_retain_n_(object: UnsafeMutablePointer, n: UInt32) -> UnsafeMutablePointer { let refcount = refcountPointer(for: object) - if loadRelaxed(refcount) == HeapObject.immortalRefCount { + if loadRelaxed(refcount) & HeapObject.refcountMask == HeapObject.immortalRefCount { return object } @@ -294,7 +295,8 @@ func swift_release_n_(object: UnsafeMutablePointer?, n: UInt32) { } let refcount = refcountPointer(for: object) - if loadRelaxed(refcount) == HeapObject.immortalRefCount { + let loadedRefcount = loadRelaxed(refcount) + if loadedRefcount & HeapObject.refcountMask == HeapObject.immortalRefCount { return } @@ -309,7 +311,8 @@ func swift_release_n_(object: UnsafeMutablePointer?, n: UInt32) { // There can only be one thread with a reference at this point (because // we're releasing the last existing reference), so a relaxed store is // enough. - storeRelaxed(refcount, newValue: HeapObject.immortalRefCount) + let doNotFree = (loadedRefcount & HeapObject.doNotFreeBit) != 0 + storeRelaxed(refcount, newValue: HeapObject.immortalRefCount | (doNotFree ? HeapObject.doNotFreeBit : 0)) _swift_embedded_invoke_heap_object_destroy(object) } else if resultingRefcount < 0 { diff --git a/test/embedded/Inputs/debug-malloc.c b/test/embedded/Inputs/debug-malloc.c new file mode 100644 index 00000000000..b780d9e5ee6 --- /dev/null +++ b/test/embedded/Inputs/debug-malloc.c @@ -0,0 +1,38 @@ +#include +#include + +#define HEAP_SIZE (8 * 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); + + if (next_heap_index + count * size > HEAP_SIZE) { + puts("HEAP EXHAUSTED\n"); + __builtin_trap(); + } + void *p = &heap[next_heap_index]; + next_heap_index += count * size; + printf("-> %p\n", p); + return p; +} + +void *malloc(size_t count) { + return calloc(count, 1); +} + +int posix_memalign(void **memptr, size_t alignment, size_t size) { + *memptr = calloc(size + alignment, 1); + if (((uintptr_t)*memptr) % alignment == 0) return 0; + *(uintptr_t *)memptr += alignment - ((uintptr_t)*memptr % alignment); + return 0; +} + +void free(void *ptr) { + // don't actually free + printf("free(%p)\n", ptr); +} diff --git a/test/embedded/refcounting-leak.swift b/test/embedded/refcounting-leak.swift new file mode 100644 index 00000000000..8b3b513c360 --- /dev/null +++ b/test/embedded/refcounting-leak.swift @@ -0,0 +1,35 @@ +// RUN: %empty-directory(%t) +// RUN: %target-swift-frontend -enable-experimental-feature Embedded -parse-as-library %s -c -o %t/a.o +// RUN: %target-clang -x c -std=c11 -c %S/Inputs/debug-malloc.c -o %t/debug-malloc.o +// RUN: %target-clang %t/a.o %t/debug-malloc.o -o %t/a.out +// RUN: %target-run %t/a.out | %FileCheck %s + +// REQUIRES: executable_test +// REQUIRES: swift_in_compiler +// REQUIRES: optimized_stdlib +// REQUIRES: OS=macosx + +var x: UInt = 1 +func randomNumber() -> UInt { + x = ((x >> 16) ^ x) &* 0x45d9f3b + x = ((x >> 16) ^ x) &* 0x45d9f3b + x = (x >> 16) ^ x + return x +} + +@main +struct Main { + static func main() { + for _ in 0 ..< 3 { + _ = randomNumber().description + } + print("OK!") + // CHECK: malloc + // CHECK: free + // CHECK: malloc + // CHECK: free + // CHECK: malloc + // CHECK: free + // CHECK: OK! + } +}