Make DNSSD queries cancellable

Sources/AsyncDNSResolver/dnssd/DNSResolver_dnssd.swift

@@ -110,20 +110,7 @@ struct DNSSD {
         let recordStream = AsyncThrowingStream<ReplyHandler.Record, Error> { continuation in
             let handler = QueryReplyHandler(handler: replyHandler, continuation)
 
-            // Wrap `handler` into a pointer so we can pass it to DNSServiceQueryRecord
-            let handlerPointer = UnsafeMutableRawPointer.allocate(
-                byteCount: MemoryLayout<QueryReplyHandler>.stride,
-                alignment: MemoryLayout<QueryReplyHandler>.alignment
-            )
-
-            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)
-                pointer.deallocate()
-            }
+            let query = Query(handler: handler)
 
             // This is called once per record received
             let callback: DNSServiceQueryRecordReply = { _, _, _, errorCode, _, _, _, rdlen, rdata, _, context in
@@ -138,32 +125,42 @@ struct DNSSD {
                 handler.handleRecord(errorCode: errorCode, data: rdata, length: rdlen)
             }
 
-            let serviceRefPtr = UnsafeMutablePointer<DNSServiceRef?>.allocate(capacity: 1)
-            defer { serviceRefPtr.deallocate() }
-
             // Run the query
             let _code = DNSServiceQueryRecord(
-                serviceRefPtr,
+                query.serviceRefPtr,
                 kDNSServiceFlagsTimeout,
                 0,
                 name,
                 UInt16(type.kDNSServiceType),
                 UInt16(kDNSServiceClass_IN),
                 callback,
-                handlerPointer
+                query.replyHandlerPointer
             )
 
             // Check if query completed successfully
             guard _code == kDNSServiceErr_NoError else {
                 return continuation.finish(throwing: AsyncDNSResolver.Error(dnssdCode: _code))
             }
 
-            // Read reply from the socket (blocking) then call reply handler
-            DNSServiceProcessResult(serviceRefPtr.pointee)
-            DNSServiceRefDeallocate(serviceRefPtr.pointee)
+            let serviceSockFD = DNSServiceRefSockFD(query.serviceRefPtr.pointee)
+            guard serviceSockFD != -1 else {
+                return continuation.finish(throwing: AsyncDNSResolver.Error(code: .internalError, message: "Failed to access the DNSSD service socket"))
+            }
+
+            let readSource = DispatchSource.makeReadSource(fileDescriptor: serviceSockFD)
+            readSource.setEventHandler {
+                // Read reply from the socket (blocking) then call reply handler
+                DNSServiceProcessResult(query.serviceRefPtr.pointee)
 
-            // Streaming done
-            continuation.finish()
+                // Streaming done
+                continuation.finish()
+            }
+            readSource.resume()
+
+            continuation.onTermination = { _ in
+                readSource.cancel()
+                DNSServiceRefDeallocate(query.serviceRefPtr.pointee)
+            }
         }
 
         // Build reply using records received
@@ -179,6 +176,27 @@ struct DNSSD {
 
 @available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *)
 extension DNSSD {
+    // Class used to manage both the handler and the serviceRef pointers, so that they are deallocated when we are done using them
+    class Query {
+        let serviceRefPtr = UnsafeMutablePointer<DNSServiceRef?>.allocate(capacity: 1)
+
+        let replyHandlerPointer = UnsafeMutableRawPointer.allocate(
+            byteCount: MemoryLayout<QueryReplyHandler>.stride,
+            alignment: MemoryLayout<QueryReplyHandler>.alignment
+        )
+        init(handler: QueryReplyHandler) {
+            // Wrap 'handler' into a pointer so we can pass it to DNSServiceQueryRecord
+            self.replyHandlerPointer.initializeMemory(as: QueryReplyHandler.self, repeating: handler, count: 1)
+        }
+
+        deinit {
+            serviceRefPtr.deallocate()
+            let pointer = replyHandlerPointer.assumingMemoryBound(to: QueryReplyHandler.self)
+            pointer.deinitialize(count: 1)
+            pointer.deallocate()
+        }
+    }
+
     class QueryReplyHandler {
         private let _handleRecord: (DNSServiceErrorType, UnsafeRawPointer?, UInt16) -> Void