Fix leaks

apple · Nov 4, 2024 · c3884fb · c3884fb
1 parent da8a204
commit c3884fb
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 5 deletions.
diff --git a/Sources/NIOResumableUpload/HTTPResumableUpload.swift b/Sources/NIOResumableUpload/HTTPResumableUpload.swift
@@ -142,6 +142,7 @@ extension HTTPResumableUpload {
             if close {
                 uploadHandler.close(mode: .all, promise: nil)
             }
+            uploadHandler.detach()
 
             if self.uploadChannel != nil {
                 self.idleTimer?.cancel()
@@ -469,6 +470,10 @@ extension HTTPResumableUpload {
         precondition(mode != .input)
         self.destroyChannel(error: nil)
         self.uploadHandler?.close(mode: mode, promise: promise)
+        self.uploadHandler?.detach()
+        self.uploadHandler = nil
+        self.idleTimer?.cancel()
+        self.idleTimer = nil
     }
 }
 

diff --git a/Sources/NIOResumableUpload/HTTPResumableUploadHandler.swift b/Sources/NIOResumableUpload/HTTPResumableUploadHandler.swift
@@ -66,13 +66,13 @@ public final class HTTPResumableUploadHandler: ChannelDuplexHandler {
         }
     }
 
-    private func resetUpload() {
+    private func resetUpload(context: ChannelHandlerContext) {
         if let existingUpload = self.upload {
             existingUpload.end(handler: self, error: nil)
         }
         let upload = self.createUpload()
         upload.scheduleOnEventLoop(self.eventLoop)
-        upload.attachUploadHandler(self, channel: self.context.channel)
+        upload.attachUploadHandler(self, channel: context.channel)
         self.upload = upload
         self.shouldReset = false
     }
@@ -81,7 +81,7 @@ public final class HTTPResumableUploadHandler: ChannelDuplexHandler {
         self.context = context
         self.eventLoop = context.eventLoop
 
-        self.resetUpload()
+        self.resetUpload(context: context)
     }
 
     public func channelActive(context: ChannelHandlerContext) {
@@ -94,13 +94,13 @@ public final class HTTPResumableUploadHandler: ChannelDuplexHandler {
 
     public func channelRead(context: ChannelHandlerContext, data: NIOAny) {
         if self.shouldReset {
-            self.resetUpload()
+            self.resetUpload(context: context)
         }
         let part = self.unwrapInboundIn(data)
         if case .end = part {
             self.shouldReset = true
         }
-        self.upload?.receive(handler: self, channel: self.context.channel, part: part)
+        self.upload?.receive(handler: self, channel: context.channel, part: part)
     }
 
     public func channelReadComplete(context: ChannelHandlerContext) {
@@ -163,4 +163,10 @@ extension HTTPResumableUploadHandler {
             self.context.close(mode: mode, promise: promise)
         }
     }
+
+    func detach() {
+        self.runInEventLoop {
+            self.context = nil
+        }
+    }
 }