fix: Rely on on task completion vs http completion to end tasks (#134)

crleona · web-flow · commit b17e7044e20e · 2024-03-15T11:43:00.000-07:00
* fix: Rely on on task completion vs http completion to end tasks

* Add tests
diff --git a/Sources/Amplitude/Utilities/EventPipeline.swift b/Sources/Amplitude/Utilities/EventPipeline.swift
@@ -19,7 +19,7 @@ public class EventPipeline {
         let events: String
         let task: URLSessionDataTask
     }
-    private var uploads = [UploadTaskInfo]()
+    private var uploads = [URL: UploadTaskInfo]()
 
     init(amplitude: Amplitude) {
         self.amplitude = amplitude
@@ -54,32 +54,36 @@ public class EventPipeline {
         eventCount = 0
         guard let storage = self.storage else { return }
         storage.rollover()
-        guard let eventFiles: [URL]? = storage.read(key: StorageKey.EVENTS) else { return }
-        cleanupUploads()
-        if pendingUploads == 0 {
-            for eventFile in eventFiles! {
-                guard let eventsString = storage.getEventsString(eventBlock: eventFile) else {
-                    continue
+        guard let eventFiles: [URL] = storage.read(key: StorageKey.EVENTS) else { return }
+        for eventFile in eventFiles {
+            uploadsQueue.sync {
+                guard uploads[eventFile] == nil else {
+                    amplitude.logger?.log(message: "Existing upload in progress, skipping...")
+                    return
                 }
-                if eventsString.isEmpty {
-                    continue
+                guard let eventsString = storage.getEventsString(eventBlock: eventFile),
+                      !eventsString.isEmpty else {
+                    return
                 }
                 let uploadTask = httpClient.upload(events: eventsString) { [weak self] result in
+                    guard let self else {
+                        return
+                    }
                     let responseHandler = storage.getResponseHandler(
-                        configuration: self!.amplitude.configuration,
-                        eventPipeline: self!,
+                        configuration: self.amplitude.configuration,
+                        eventPipeline: self,
                         eventBlock: eventFile,
                         eventsString: eventsString
                     )
                     responseHandler.handle(result: result)
-                    self?.cleanupUploads()
+                    self.completeUpload(for: eventFile)
                 }
-                if let upload = uploadTask {
-                    add(uploadTask: UploadTaskInfo(events: eventsString, task: upload))
+                if let uploadTask {
+                    uploads[eventFile] = UploadTaskInfo(events: eventsString, task: uploadTask)
                 }
             }
-            completion?()
         }
+        completion?()
     }
 
     func start() {
@@ -101,31 +105,10 @@ public class EventPipeline {
 }
 
 extension EventPipeline {
-    internal func cleanupUploads() {
-        uploadsQueue.sync {
-            let before = uploads.count
-            var newPending = uploads
-            newPending.removeAll { uploadInfo in
-                let shouldRemove = uploadInfo.task.state != .running
-                return shouldRemove
-            }
-            uploads = newPending
-            let after = uploads.count
-            amplitude.logger?.log(message: "Cleaned up \(before - after) non-running uploads.")
-        }
-    }
-
-    internal var pendingUploads: Int {
-        var uploadsCount = 0
-        uploadsQueue.sync {
-            uploadsCount = uploads.count
-        }
-        return uploadsCount
-    }
 
-    internal func add(uploadTask: UploadTaskInfo) {
+    func completeUpload(for eventBlock: URL) {
         uploadsQueue.sync {
-            uploads.append(uploadTask)
+            uploads[eventBlock] = nil
         }
     }
 }
diff --git a/Tests/AmplitudeTests/Supports/TestUtilities.swift b/Tests/AmplitudeTests/Supports/TestUtilities.swift
@@ -47,62 +47,75 @@ class FakeInMemoryStorage: Storage {
     var keyValueStore = [String: Any?]()
     var eventsStore = [URL: [BaseEvent]]()
     var index = URL(string: "0")!
+    let storageQueue = DispatchQueue(label: "Amplitude.FakeInMemoryStorage")
 
     func write(key: StorageKey, value: Any?) throws {
-        switch key {
-        case .EVENTS:
-            if let event = value as? BaseEvent {
-                var chunk = eventsStore[index, default: [BaseEvent]()]
-                chunk.append(event)
-                eventsStore[index] = chunk
+        storageQueue.sync {
+            switch key {
+            case .EVENTS:
+                if let event = value as? BaseEvent {
+                    var chunk = eventsStore[index, default: [BaseEvent]()]
+                    chunk.append(event)
+                    eventsStore[index] = chunk
+                }
+            default:
+                keyValueStore[key.rawValue] = value
             }
-        default:
-            keyValueStore[key.rawValue] = value
         }
     }
 
     func read<T>(key: StorageKey) -> T? {
-        var result: T?
-        switch key {
-        case .EVENTS:
-            result = Array(eventsStore.keys) as? T
-        default:
-            result = keyValueStore[key.rawValue] as? T
+        storageQueue.sync {
+            var result: T?
+            switch key {
+            case .EVENTS:
+                result = Array(eventsStore.keys) as? T
+            default:
+                result = keyValueStore[key.rawValue] as? T
+            }
+            return result
         }
-        return result
     }
 
     func getEventsString(eventBlock: EventBlock) -> String? {
-        var content: String?
-        content = "["
-        content = content! + (eventsStore[eventBlock] ?? []).map { $0.toString() }.joined(separator: ", ")
-        content = content! + "]"
-        return content
+        storageQueue.sync {
+            var content: String?
+            content = "["
+            content = content! + (eventsStore[eventBlock] ?? []).map { $0.toString() }.joined(separator: ", ")
+            content = content! + "]"
+            return content
+        }
     }
 
     func rollover() {
     }
 
     func reset() {
-        keyValueStore.removeAll()
-        eventsStore.removeAll()
+        storageQueue.sync {
+            keyValueStore.removeAll()
+            eventsStore.removeAll()
+        }
     }
 
     func remove(eventBlock: EventBlock) {
-        eventsStore.removeValue(forKey: eventBlock)
+        storageQueue.sync {
+            _ = eventsStore.removeValue(forKey: eventBlock)
+        }
     }
 
     func splitBlock(eventBlock: EventBlock, events: [BaseEvent]) {
     }
 
     func events() -> [BaseEvent] {
-        var result: [BaseEvent] = []
-        for (_, value) in eventsStore {
-            for event in value {
-                result.append(event)
+        storageQueue.sync {
+            var result: [BaseEvent] = []
+            for (_, value) in eventsStore {
+                for event in value {
+                    result.append(event)
+                }
             }
+            return result
         }
-        return result
     }
 
     nonisolated func getResponseHandler(
@@ -121,6 +134,7 @@ class FakeHttpClient: HttpClient {
     var uploadCount: Int = 0
     var uploadedEvents: [String] = []
     var uploadExpectations: [XCTestExpectation] = []
+    var uploadResults: [Result<Int, Error>] = []
 
     override func upload(events: String, completion: @escaping (_ result: Result<Int, Error>) -> Void)
         -> URLSessionDataTask?
@@ -130,8 +144,16 @@ class FakeHttpClient: HttpClient {
         if !uploadExpectations.isEmpty {
             uploadExpectations.removeFirst().fulfill()
         }
-        completion(Result.success(200))
-        return nil
+        let result: Result<Int, Error>
+        if uploadResults.isEmpty {
+            result = .success(200)
+        } else {
+            result = uploadResults.removeFirst()
+        }
+        DispatchQueue.global().async {
+            completion(result)
+        }
+        return URLSession.shared.dataTask(with: URLRequest(url: URL(string: "https://www.amplitude.com")!))
     }
 
     override func getDate() -> Date {
diff --git a/Tests/AmplitudeTests/Utilities/EventPipelineTests.swift b/Tests/AmplitudeTests/Utilities/EventPipelineTests.swift
@@ -15,10 +15,14 @@ final class EventPipelineTests: XCTestCase {
     private var configuration: Configuration!
     private var pipeline: EventPipeline!
     private var httpClient: FakeHttpClient!
+    private var storage: PersistentStorage!
 
     override func setUp() {
         super.setUp()
-        let storage = FakeInMemoryStorage()
+        storage = PersistentStorage(
+            storagePrefix: "event-pipeline-tests",
+            logger: nil,
+            diagonostics: Diagnostics())
         configuration = Configuration(
             apiKey: "testApiKey",
             flushIntervalMillis: Int(Self.FLUSH_INTERVAL_SECONDS * 1000),
@@ -30,6 +34,11 @@ final class EventPipelineTests: XCTestCase {
         pipeline.httpClient = httpClient
     }
 
+    override func tearDown() {
+        super.tearDown()
+        storage.reset()
+    }
+
     func testInit() {
         XCTAssertEqual(pipeline.amplitude.configuration.apiKey, configuration.apiKey)
     }
@@ -77,4 +86,57 @@ final class EventPipelineTests: XCTestCase {
         XCTAssertEqual(pipeline.amplitude.configuration.offline, true)
         XCTAssertEqual(httpClient.uploadCount, 0, "There should be no uploads when offline")
     }
+
+    func testSimultaneousFlush() {
+        let testEvent = BaseEvent(userId: "unit-test", deviceId: "unit-test-machine", eventType: "testEvent")
+        try? pipeline.storage?.write(key: StorageKey.EVENTS, value: testEvent)
+
+        let flushExpectations = (0..<2).map { _ in
+            let expectation = expectation(description: "flush")
+            pipeline.flush {
+                expectation.fulfill()
+            }
+            return expectation
+        }
+
+        let waitResult = XCTWaiter.wait(for: flushExpectations, timeout: 1)
+        XCTAssertNotEqual(waitResult, .timedOut)
+        XCTAssertEqual(httpClient.uploadCount, 1)
+        let uploadedEvents = BaseEvent.fromArrayString(jsonString: httpClient.uploadedEvents[0])
+        XCTAssertEqual(uploadedEvents?.count, 1)
+        XCTAssertEqual(uploadedEvents![0].eventType, "testEvent")
+    }
+
+    func testInvalidEventUpload() {
+        (0..<2).forEach { i in
+            let testEvent = BaseEvent(userId: "test", deviceId: "test-machine", eventType: "testEvent-\(i)")
+            try? pipeline.storage?.write(key: StorageKey.EVENTS, value: testEvent)
+        }
+
+        let invalidResponseData = "{\"events_with_invalid_fields\": {\"user_id\": [0]}}".data(using: .utf8)!
+
+        httpClient.uploadResults = [
+            .failure(HttpClient.Exception.httpError(code: 400, data: invalidResponseData))
+        ]
+
+        let uploadExpectations = (0..<2).map { _ in expectation(description: "httpresponse") }
+        httpClient.uploadExpectations = uploadExpectations
+
+        pipeline.flush()
+        wait(for: [uploadExpectations[0]], timeout: 1)
+
+        pipeline.flush()
+        wait(for: [uploadExpectations[1]], timeout: 1)
+
+        XCTAssertEqual(httpClient.uploadCount, 2)
+
+        let uploadedEvents0 = BaseEvent.fromArrayString(jsonString: httpClient.uploadedEvents[0])
+        XCTAssertEqual(uploadedEvents0?.count, 2)
+        XCTAssertEqual(uploadedEvents0?[0].eventType, "testEvent-0")
+        XCTAssertEqual(uploadedEvents0?[1].eventType, "testEvent-1")
+
+        let uploadedEvents1 = BaseEvent.fromArrayString(jsonString: httpClient.uploadedEvents[1])
+        XCTAssertEqual(uploadedEvents1?.count, 1)
+        XCTAssertEqual(uploadedEvents1?[0].eventType, "testEvent-1")
+    }
 }
