Skip to content

Commit

Permalink
fix: Rely on on task completion vs http completion to end tasks (#134)
Browse files Browse the repository at this point in the history 
* fix: Rely on on task completion vs http completion to end tasks

* Add tests
  • Loading branch information
@crleona
crleona authored Mar 15, 2024
1 parent 5d214c4 commit b17e704
Show file tree
Hide file tree
Showing 3 changed files with 136 additions and 69 deletions.
59 changes: 21 additions & 38 deletions Sources/Amplitude/Utilities/EventPipeline.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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() {
Expand All @@ -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
}
}
}
82 changes: 52 additions & 30 deletions Tests/AmplitudeTests/Supports/TestUtilities.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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?
Expand All @@ -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 {
Expand Down
64 changes: 63 additions & 1 deletion Tests/AmplitudeTests/Utilities/EventPipelineTests.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand All @@ -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)
}
Expand Down Expand Up @@ -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")
}
}

0 comments on commit b17e704

Please sign in to comment.