Skip to content

Fix continuation memory leak in Ares.query #31

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Feb 21, 2024
Merged

Fix continuation memory leak in Ares.query #31

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
a2990ae
Fix continuation memory leak in Ares.query
dieb Feb 21, 2024
6445386
Use class for QueryReplyHandler
dieb Feb 21, 2024
fe8c851
Deallocate QueryReplyHandler for DNSSD
dieb Feb 21, 2024
5d6e5bf
Move defer block after allocate/initialize, use class
dieb Feb 21, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 7 additions & 8 deletions Sources/AsyncDNSResolver/c-ares/DNSResolver_c-ares.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -160,8 +160,12 @@ class Ares {
preconditionFailure("'arg' is nil. This is a bug.")
}

let handler = QueryReplyHandler(pointer: handlerPointer)
defer { handlerPointer.deallocate() }
let pointer = handlerPointer.assumingMemoryBound(to: QueryReplyHandler.self)
let handler = pointer.pointee
defer {
pointer.deinitialize(count: 1)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cool! I was making similar changes as well but was missing this line of code.

pointer.deallocate()
}

handler.handle(status: status, buffer: buf, length: len)
}
Expand Down Expand Up @@ -258,7 +262,7 @@ extension Ares {

@available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *)
extension Ares {
struct QueryReplyHandler {
class QueryReplyHandler {
private let _handler: (CInt, UnsafeMutablePointer<CUnsignedChar>?, CInt) -> Void

init<Parser: AresQueryReplyParser>(parser: Parser, _ continuation: CheckedContinuation<Parser.Reply, Error>) {
Expand All @@ -276,11 +280,6 @@ extension Ares {
}
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we also change struct QueryReplyHandler to class QueryReplyHandler please? That combined with pointer.deinitialize(count: 1) I see deinit getting called.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Absolutely, just pushed a new commit to reflect those changes.

I'm interested in understanding the rationale behind class vs struct here. Could you please share more about your thought process? Swift noob eager to learn here 😄

My best guess is that class deinit can free up its members, though my tests indicate no leaks with struct.

Copy link
Member

@yim-lee yim-lee Feb 21, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was told that the address of a struct is not stable (we were thinking deallocate might not be working on the right address because of this), plus we are doing all these pointer manipulations already so class probably makes more sense.

}

init(pointer: UnsafeMutableRawPointer) {
let handlerPointer = pointer.assumingMemoryBound(to: Self.self)
self = handlerPointer.pointee
}

func handle(status: CInt, buffer: UnsafeMutablePointer<CUnsignedChar>?, length: CInt) {
self._handler(status, buffer, length)
}
Expand Down
20 changes: 11 additions & 9 deletions Sources/AsyncDNSResolver/dnssd/DNSResolver_dnssd.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -115,18 +115,25 @@ struct DNSSD {
byteCount: MemoryLayout<QueryReplyHandler>.stride,
alignment: MemoryLayout<QueryReplyHandler>.alignment
)
// The handler might be called multiple times so don't deallocate inside `callback`
defer { handlerPointer.deallocate() }

handlerPointer.initializeMemory(as: QueryReplyHandler.self, repeating: handler, count: 1)

// The handler might be called multiple times so don't deallocate inside `callback`
defer {
let pointer = handlerPointer.assumingMemoryBound(to: QueryReplyHandler.self)
pointer.deinitialize(count: 1)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Without this line of code, we have these leaks:

image

The above is modifying test_queryAAAA to run self.resolver.queryAAAA 100K times.

With this line there were no leaks.

pointer.deallocate()
}

// This is called once per record received
let callback: DNSServiceQueryRecordReply = { _, _, _, errorCode, _, _, _, rdlen, rdata, _, context in
guard let handlerPointer = context else {
preconditionFailure("'context' is nil. This is a bug.")
}

let handler = QueryReplyHandler(pointer: handlerPointer)
let pointer = handlerPointer.assumingMemoryBound(to: QueryReplyHandler.self)
let handler = pointer.pointee

// This parses a record then adds it to the stream
handler.handleRecord(errorCode: errorCode, data: rdata, length: rdlen)
}
Expand Down Expand Up @@ -171,7 +178,7 @@ struct DNSSD {

@available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *)
extension DNSSD {
struct QueryReplyHandler {
class QueryReplyHandler {
private let _handleRecord: (DNSServiceErrorType, UnsafeRawPointer?, UInt16) -> Void

init<Handler: DNSSDQueryReplyHandler>(handler: Handler, _ continuation: AsyncThrowingStream<Handler.Record, Error>.Continuation) {
Expand All @@ -189,11 +196,6 @@ extension DNSSD {
}
}

init(pointer: UnsafeMutableRawPointer) {
let handlerPointer = pointer.assumingMemoryBound(to: Self.self)
self = handlerPointer.pointee
}

func handleRecord(errorCode: DNSServiceErrorType, data: UnsafeRawPointer?, length: UInt16) {
self._handleRecord(errorCode, data, length)
}
Expand Down