From 4225e45a33855710437aaf154cb7b16eed0feac4 Mon Sep 17 00:00:00 2001
From: Zhennan Zhou <zhouz46@uw.edu>
Date: Thu, 5 Sep 2024 23:35:01 -0700
Subject: [PATCH] reduce memory footprint for icmp ping handler

---
 Package.swift                                 |  4 +--
 Sources/LCLPing/ICMP/ICMPHandler.swift        | 31 +++++++++----------
 .../HTTPTracingHandlerTests.swift             |  1 +
 3 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/Package.swift b/Package.swift
index 4608db7..1a26e82 100644
--- a/Package.swift
+++ b/Package.swift
@@ -16,10 +16,10 @@ let package = Package(
     ],
     dependencies: [
         // Dependencies declare other packages that this package depends on.
-        .package(url: "https://github.com/apple/swift-nio.git", from: "2.62.0"),
+        .package(url: "https://github.com/apple/swift-nio.git", from: "2.72.0"),
         .package(url: "https://github.com/apple/swift-log.git", from: "1.5.3"),
         .package(url: "https://github.com/apple/swift-nio-ssl.git", from: "2.25.0"),
-        .package(url: "https://github.com/apple/swift-collections.git", from: "1.0.4")
+        .package(url: "https://github.com/apple/swift-collections.git", from: "1.1.1")
     ],
     targets: [
         // Targets are the basic building blocks of a package. A target can define a module or a test suite.
diff --git a/Sources/LCLPing/ICMP/ICMPHandler.swift b/Sources/LCLPing/ICMP/ICMPHandler.swift
index 9dc329c..6e9060b 100644
--- a/Sources/LCLPing/ICMP/ICMPHandler.swift
+++ b/Sources/LCLPing/ICMP/ICMPHandler.swift
@@ -12,6 +12,7 @@
 
 import Foundation
 import NIOCore
+import Collections
 
 extension ICMPPingClient {
 
@@ -139,11 +140,11 @@ final class ICMPHandler: PingHandler {
     // sequence number to ICMP request
     private var seqToRequest: [Int: ICMPPingClient.ICMPHeader]
 
-    // sequence number to an optional ICMP response
-    private var seqToResponse: [Int: ICMPPingClient.ICMPHeader?]
+    // a bit set that contains the response sequence number seen so far
+    private var seen: BitSet
 
-    // a set that contains the response sequence number received by the handler
-    private var responseSeqNumSet: Set<Int>
+    // a bit set that contains the response sequence number received by the handler
+    private var hasResponses: BitSet
 
     // a list of `PingResponse`
     private var result: [PingResponse]
@@ -154,10 +155,10 @@ final class ICMPHandler: PingHandler {
     init(totalCount: Int, promise: EventLoopPromise<[PingResponse]>) {
         self.totalCount = totalCount
         self.seqToRequest = [:]
-        self.seqToResponse = [:]
-        self.responseSeqNumSet = Set()
         self.result = []
         self.icmpPingPromise = promise
+        self.seen = BitSet(reservingCapacity: self.totalCount)
+        self.hasResponses = BitSet(reservingCapacity: self.totalCount)
     }
 
     func handleRead(response: ICMPPingClient.ICMPHeader) {
@@ -176,17 +177,16 @@ final class ICMPHandler: PingHandler {
             return
         }
 
-        if self.responseSeqNumSet.contains(sequenceNum) {
-            let pingResponse: PingResponse = self.seqToResponse[sequenceNum] ==
-                                                nil ? .timeout(sequenceNum) : .duplicated(sequenceNum)
-            logger.debug("[\(#fileID)][\(#line)][\(#function)]:: response for #\(sequenceNum) is \(self.seqToResponse[sequenceNum] == nil ? "timeout" : "duplicate")")
+        if self.seen.contains(sequenceNum) {
+            let pingResponse: PingResponse = self.hasResponses.contains(sequenceNum) ? .duplicated(sequenceNum) : .timeout(sequenceNum)
+            logger.debug("[\(#fileID)][\(#line)][\(#function)]:: response for #\(sequenceNum) is \(self.hasResponses.contains(sequenceNum) ? "timeout" : "duplicate")")
             result.append(pingResponse)
             shouldCloseHandler()
             return
         }
 
-        self.seqToResponse[sequenceNum] = response
-        self.responseSeqNumSet.insert(sequenceNum)
+        self.hasResponses.insert(sequenceNum)
+        self.seen.insert(sequenceNum)
 
         switch (type, code) {
         case (ICMPPingClient.ICMPType.echoReply.rawValue, 0):
@@ -334,9 +334,9 @@ final class ICMPHandler: PingHandler {
     }
 
     func handleTimeout(sequenceNumber: Int) {
-        if !self.responseSeqNumSet.contains(sequenceNumber) {
+        if !self.seen.contains(sequenceNumber) {
             logger.debug("[\(#fileID)][\(#line)][\(#function)]: #\(sequenceNumber) timed out")
-            self.responseSeqNumSet.insert(sequenceNumber)
+            self.seen.insert(sequenceNumber)
             self.result.append(.timeout(sequenceNumber))
             shouldCloseHandler()
         }
@@ -353,12 +353,11 @@ final class ICMPHandler: PingHandler {
 
     func reset() {
         self.seqToRequest.removeAll()
-        self.seqToResponse.removeAll()
         self.result.removeAll()
     }
 
     func shouldCloseHandler(shouldForceClose: Bool = false) {
-        if self.responseSeqNumSet.count == self.totalCount || shouldForceClose {
+        if self.seen.count == self.totalCount || shouldForceClose {
             logger.debug("[\(#fileID)][\(#line)][\(#function)]: should close icmp handler")
             self.icmpPingPromise.succeed(self.result)
         }
diff --git a/Tests/HTTPChannelTests/HTTPTracingHandlerTests.swift b/Tests/HTTPChannelTests/HTTPTracingHandlerTests.swift
index 3abd50f..b58c0b2 100644
--- a/Tests/HTTPChannelTests/HTTPTracingHandlerTests.swift
+++ b/Tests/HTTPChannelTests/HTTPTracingHandlerTests.swift
@@ -68,6 +68,7 @@ final class HTTPTracingHandlerTests: XCTestCase {
             XCTAssertEqual(request.method, HTTPMethod.GET)
             XCTAssertEqual(request.uri, config.url.uri)
             XCTAssertEqual(request.headers, config.httpHeaders)
+            promise.succeed(.ok(2, 0, 0))
         default:
             XCTFail("Should receive a head and end. But received head = \(String(describing: head)), end = \(String(describing: end))")
         }