Skip to content

Commit d514f1c

Browse files
jvh-awsjvh-aws
authored
fix(kinesis): enable full flush for cache (#4169)
* feat(kinesis): fix retry logic for request-level errors * fix(kinesis): enable full flush for cache * Add record storage tests for id exclusion * Update flush api doc * Use retryAfter to enable full flush * Fix nonisolated actor warning in sqliterecordstorage init * Fix test warning * Adapt kinesis client config provider to underlying sdk configuration deprecation * Revert "Adapt kinesis client config provider to underlying sdk configuration deprecation" This reverts commit 4dc335c.
1 parent 6c7fc17 commit d514f1c

8 files changed

Lines changed: 196 additions & 72 deletions

‎AmplifyClients/AmplifyKinesisClient/Sources/AmplifyKinesisClient.swift‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -212,12 +212,14 @@ public class AmplifyKinesisClient {
212212
///
213213
/// Each invocation sends at most one batch per stream, limited by the Kinesis
214214
/// `PutRecords` constraints (up to 500 records or 10 MB per stream). If the cache
215-
/// contains more records than a single batch can hold, the remaining records are
216-
/// sent on subsequent flush invocations — either manually or via the auto-flush
217-
/// scheduler.
215+
/// Flushes all locally stored records to their respective Kinesis streams.
218216
///
219-
/// Records that fail within a batch are marked for retry on the next flush. Records
220-
/// that exceed ``Options/maxRetries`` are removed from the cache.
217+
/// Each flush processes all pending records in batches per stream (limited by
218+
/// record count and byte size). Records that fail or are retryable within a flush
219+
/// cycle are not retried in the same flush — they are skipped and will be picked
220+
/// up in the next flush cycle.
221+
///
222+
/// Records that exceed ``Options/maxRetries`` are removed from the cache.
221223
///
222224
/// If a flush is already in progress, the call returns immediately with
223225
/// `FlushData(recordsFlushed: 0, flushInProgress: true)`.

‎AmplifyClients/AmplifyKinesisClient/Sources/Support/RecordClient.swift‎

Lines changed: 45 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -45,48 +45,59 @@ actor RecordClient {
4545
defer { isFlushing = false }
4646

4747
var totalFlushed = 0
48-
let recordsByStreamList: [[Record]] = try await storage.getRecordsByStream()
49-
logger.debug("Retrieved \(recordsByStreamList.count) stream(s) with records to flush")
48+
var lastIdByStream: [String: Int64] = [:]
5049

51-
for records in recordsByStreamList {
52-
guard !records.isEmpty else { continue }
53-
let streamName = records[0].streamName
54-
let recordCount = records.count
55-
logger.verbose("Flushing \(recordCount) records to stream: \(streamName)")
50+
var recordsByStreamList: [[Record]] = try await storage.getRecordsByStream(afterIdByStream: lastIdByStream)
51+
while !recordsByStreamList.isEmpty {
5652

57-
do {
58-
let response = try await sender.putRecords(streamName: streamName, records: records)
53+
logger.debug("Retrieved \(recordsByStreamList.count) stream(s) with records to flush")
5954

60-
totalFlushed += response.successfulIds.count
55+
for records in recordsByStreamList {
56+
guard !records.isEmpty else { continue }
57+
let streamName = records[0].streamName
58+
let recordCount = records.count
59+
logger.verbose("Flushing \(recordCount) records to stream: \(streamName)")
6160

62-
try await withThrowingTaskGroup(of: Void.self) { group in
63-
group.addTask { try await self.storage.deleteRecords(ids: response.successfulIds) }
64-
group.addTask { try await self.storage.incrementRetryCount(ids: response.retryableIds) }
65-
group.addTask { try await self.storage.deleteRecords(ids: response.failedIds) }
66-
try await group.waitForAll()
67-
}
61+
// Track the last record ID per stream so subsequent batches start after it
62+
let maxId = records.map(\.id).max() ?? 0
63+
lastIdByStream[streamName] = maxId
6864

69-
logger.verbose(
70-
"Stream \(streamName): \(response.successfulIds.count) succeeded, "
71-
+ "\(response.retryableIds.count) retryable, \(response.failedIds.count) failed"
72-
)
73-
} catch {
74-
// Increment retry count for retryable records and delete those at the limit
75-
await handleFailedRequest(records)
76-
77-
// SDK errors are logged but not thrown — one stream shouldn't block others
78-
if error is ModeledError || error is AWSServiceError {
79-
logger.warn(
80-
"Kinesis SDK error flushing stream \(streamName): \(error.localizedDescription)"
81-
)
82-
} else {
83-
// Network errors, storage errors, and unexpected errors — throw to caller
84-
logger.warn(
85-
"Error flushing stream \(streamName): \(error.localizedDescription)"
65+
do {
66+
let response = try await sender.putRecords(streamName: streamName, records: records)
67+
68+
totalFlushed += response.successfulIds.count
69+
70+
try await withThrowingTaskGroup(of: Void.self) { group in
71+
group.addTask { try await self.storage.deleteRecords(ids: response.successfulIds) }
72+
group.addTask { try await self.storage.incrementRetryCount(ids: response.retryableIds) }
73+
group.addTask { try await self.storage.deleteRecords(ids: response.failedIds) }
74+
try await group.waitForAll()
75+
}
76+
77+
logger.verbose(
78+
"Stream \(streamName): \(response.successfulIds.count) succeeded, "
79+
+ "\(response.retryableIds.count) retryable, \(response.failedIds.count) failed"
8680
)
87-
throw error
81+
} catch {
82+
// Increment retry count for retryable records and delete those at the limit
83+
await handleFailedRequest(records)
84+
85+
// SDK errors are logged but not thrown — one stream shouldn't block others
86+
let isSdkError = error is ModeledError || error is AWSServiceError
87+
if isSdkError {
88+
logger.warn(
89+
"Kinesis SDK error flushing stream \(streamName): \(error.localizedDescription)"
90+
)
91+
} else {
92+
// Network errors, storage errors, and unexpected errors — throw to caller
93+
logger.warn(
94+
"Error flushing stream \(streamName): \(error.localizedDescription)"
95+
)
96+
throw error
97+
}
8898
}
8999
}
100+
recordsByStreamList = try await storage.getRecordsByStream(afterIdByStream: lastIdByStream)
90101
}
91102

92103
return FlushData(recordsFlushed: totalFlushed)

‎AmplifyClients/AmplifyKinesisClient/Sources/Support/RecordStorage.swift‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,10 @@ protocol RecordStorage: Actor {
99
/// Adds a new record to storage
1010
func addRecord(_ input: RecordInput) async throws
1111

12-
/// Gets all records grouped by stream name as a list of lists
13-
func getRecordsByStream() async throws -> [[Record]]
12+
/// Gets all records grouped by stream name as a list of lists.
13+
/// - Parameter afterIdByStream: A map of stream name to the last processed record ID.
14+
/// Records with `id <= afterIdByStream[streamName]` are excluded from the results.
15+
func getRecordsByStream(afterIdByStream: [String: Int64]) async throws -> [[Record]]
1416

1517
/// Deletes records by their IDs
1618
func deleteRecords(ids: [Int64]) async throws

‎AmplifyClients/AmplifyKinesisClient/Sources/Support/SQLiteRecordStorage.swift‎

Lines changed: 29 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -44,10 +44,15 @@ actor SQLiteRecordStorage: RecordStorage {
4444
self.maxRecordSizeBytes = maxRecordSizeBytes
4545
self.maxBytesPerStream = maxBytesPerStream
4646
self.maxPartitionKeyLength = maxPartitionKeyLength
47-
self.database = try connection ?? Self.createFileConnection(identifier: identifier)
4847

49-
try Self.setupSchema(on: database)
50-
try resetCacheSizeFromDb()
48+
let db = try connection ?? Self.createFileConnection(identifier: identifier)
49+
self.database = db
50+
51+
try Self.setupSchema(on: db)
52+
let size = try Self.wrapDatabaseError {
53+
try db.scalar(Self.records.select(Self.dataSize.sum)) ?? 0
54+
}
55+
self.cachedSize = Int64(size)
5156
}
5257

5358
private static func createFileConnection(identifier: String) throws -> Connection {
@@ -128,8 +133,19 @@ actor SQLiteRecordStorage: RecordStorage {
128133
cachedSize += Int64(input.dataSize)
129134
}
130135

131-
func getRecordsByStream() throws -> [[Record]] {
136+
func getRecordsByStream(afterIdByStream: [String: Int64] = [:]) throws -> [[Record]] {
132137
try Self.wrapDatabaseError {
138+
// Build per-stream WHERE clauses: id > lastProcessedId for streams we've already seen
139+
let streamFilter: String
140+
if afterIdByStream.isEmpty {
141+
streamFilter = ""
142+
} else {
143+
let conditions = afterIdByStream.map { _ in
144+
"NOT (stream_name = ? AND id <= ?)"
145+
}.joined(separator: " AND ")
146+
streamFilter = "WHERE \(conditions)"
147+
}
148+
133149
// Must use raw SQL for window functions
134150
let query = """
135151
SELECT id, stream_name, partition_key, data, data_size, retry_count, created_at
@@ -138,14 +154,22 @@ actor SQLiteRecordStorage: RecordStorage {
138154
ROW_NUMBER() OVER (PARTITION BY stream_name ORDER BY id) as rn,
139155
SUM(data_size) OVER (PARTITION BY stream_name ORDER BY id) as running_size
140156
FROM records
157+
\(streamFilter)
141158
)
142159
WHERE rn <= ? AND running_size <= ?
143160
ORDER BY stream_name, id
144161
"""
145162

163+
// Bind per-stream after-id filters
164+
var bindings: [Binding?] = afterIdByStream.flatMap { (streamName, afterId) in
165+
[streamName as Binding?, afterId as Binding?]
166+
}
167+
bindings.append(maxRecords as Binding?)
168+
bindings.append(maxBytesPerStream as Binding?)
169+
146170
var recordsByStream: [String: [Record]] = [:]
147171

148-
for row: Statement.Element in try database.prepare(query, maxRecords, maxBytesPerStream) {
172+
for row: Statement.Element in try database.prepare(query, bindings) {
149173
// Ensure row has expected number of columns (in case DB format is invalid)
150174
guard row.count >= 7 else {
151175
throw RecordCacheError.database(

‎AmplifyClients/AmplifyKinesisClient/Tests/UnitTests/AutoFlushSchedulerTests.swift‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -115,7 +115,7 @@ actor MockRecordStorage: RecordStorage {
115115
addRecordCallCount += 1
116116
}
117117

118-
func getRecordsByStream() throws -> [[Record]] {
118+
func getRecordsByStream(afterIdByStream: [String: Int64]) throws -> [[Record]] {
119119
getRecordsByStreamCallCount += 1
120120
return []
121121
}

‎AmplifyClients/AmplifyKinesisClient/Tests/UnitTests/RecordClientConcurrentFlushTests.swift‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ class RecordClientConcurrentFlushTests: XCTestCase {
4545
)
4646
}
4747

48-
let allRecords = try await storage.getRecordsByStream().flatMap { $0 }
48+
let allRecords = try await storage.getRecordsByStream(afterIdByStream: [:]).flatMap { $0 }
4949

5050
// Make the sender slow so the first flush holds the lock
5151
await sender.setDelay(nanoseconds: 500_000_000) // 500ms

‎AmplifyClients/AmplifyKinesisClient/Tests/UnitTests/RecordClientFlushTests.swift‎

Lines changed: 12 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ class RecordClientFlushTests: XCTestCase {
4747
try await storage.addRecord(RecordInput(streamName: streamName, partitionKey: "key3", data: Data([3])))
4848

4949
// Get all records and set retry count for record 3 to max (3)
50-
let allRecordsByStream = try await storage.getRecordsByStream()
50+
let allRecordsByStream = try await storage.getRecordsByStream(afterIdByStream: [:])
5151
let allRecords = allRecordsByStream.flatMap { $0 }
5252
let record3Id = allRecords[2].id
5353
try await storage.incrementRetryCount(ids: [record3Id])
@@ -65,7 +65,7 @@ class RecordClientFlushTests: XCTestCase {
6565
let result = try await recordClient.flush()
6666

6767
XCTAssertEqual(result.recordsFlushed, 1)
68-
let remainingRecords = try await storage.getRecordsByStream().flatMap { $0 }
68+
let remainingRecords = try await storage.getRecordsByStream(afterIdByStream: [:]).flatMap { $0 }
6969
XCTAssertEqual(remainingRecords.count, 1)
7070
XCTAssertEqual(remainingRecords[0].id, allRecords[1].id)
7171
XCTAssertEqual(remainingRecords[0].retryCount, 1)
@@ -86,7 +86,7 @@ class RecordClientFlushTests: XCTestCase {
8686
// Expected — non-SDK errors are critical
8787
}
8888

89-
let remainingRecords = try await storage.getRecordsByStream().flatMap { $0 }
89+
let remainingRecords = try await storage.getRecordsByStream(afterIdByStream: [:]).flatMap { $0 }
9090
XCTAssertEqual(remainingRecords.count, 3)
9191
for record in remainingRecords {
9292
XCTAssertEqual(record.retryCount, 1)
@@ -99,7 +99,7 @@ class RecordClientFlushTests: XCTestCase {
9999
try await storage.addRecord(RecordInput(streamName: streamName, partitionKey: "key2", data: Data([2])))
100100
try await storage.addRecord(RecordInput(streamName: streamName, partitionKey: "key3", data: Data([3])))
101101

102-
let allRecords = try await storage.getRecordsByStream().flatMap { $0 }
102+
let allRecords = try await storage.getRecordsByStream(afterIdByStream: [:]).flatMap { $0 }
103103
let record2Id = allRecords[1].id
104104
let record3Id = allRecords[2].id
105105

@@ -117,7 +117,7 @@ class RecordClientFlushTests: XCTestCase {
117117
// Expected — non-SDK errors are critical
118118
}
119119

120-
let remainingRecords = try await storage.getRecordsByStream().flatMap { $0 }
120+
let remainingRecords = try await storage.getRecordsByStream(afterIdByStream: [:]).flatMap { $0 }
121121
XCTAssertEqual(remainingRecords.count, 1)
122122
XCTAssertEqual(remainingRecords[0].id, allRecords[0].id)
123123
XCTAssertEqual(remainingRecords[0].retryCount, 1)
@@ -142,7 +142,7 @@ class RecordClientFlushTests: XCTestCase {
142142
}
143143

144144
// The first stream processed should have retry incremented, the second should not be processed
145-
let remainingRecords = try await storage.getRecordsByStream().flatMap { $0 }
145+
let remainingRecords = try await storage.getRecordsByStream(afterIdByStream: [:]).flatMap { $0 }
146146
XCTAssertEqual(remainingRecords.count, 2)
147147
let retryCountSum = remainingRecords.map(\.retryCount).reduce(0, +)
148148
XCTAssertEqual(retryCountSum, 1) // Only one stream was processed before throwing
@@ -158,7 +158,7 @@ class RecordClientFlushTests: XCTestCase {
158158
let result = try await recordClient.flush()
159159

160160
XCTAssertEqual(result.recordsFlushed, 0)
161-
let remainingRecords = try await storage.getRecordsByStream().flatMap { $0 }
161+
let remainingRecords = try await storage.getRecordsByStream(afterIdByStream: [:]).flatMap { $0 }
162162
XCTAssertEqual(remainingRecords.count, 2)
163163
for record in remainingRecords {
164164
XCTAssertEqual(record.retryCount, 1)
@@ -171,7 +171,7 @@ class RecordClientFlushTests: XCTestCase {
171171
try await storage.addRecord(RecordInput(streamName: stream1, partitionKey: "key1", data: Data([1])))
172172
try await storage.addRecord(RecordInput(streamName: stream2, partitionKey: "key2", data: Data([2])))
173173

174-
let initialRecords = try await storage.getRecordsByStream().flatMap { $0 }
174+
let initialRecords = try await storage.getRecordsByStream(afterIdByStream: [:]).flatMap { $0 }
175175
let stream2RecordId = try XCTUnwrap(initialRecords.first { $0.streamName == stream2 }?.id)
176176

177177
await sender.setHandler { streamName, _ in
@@ -188,7 +188,7 @@ class RecordClientFlushTests: XCTestCase {
188188
let result = try await recordClient.flush()
189189

190190
XCTAssertEqual(result.recordsFlushed, 1)
191-
let remainingRecords = try await storage.getRecordsByStream().flatMap { $0 }
191+
let remainingRecords = try await storage.getRecordsByStream(afterIdByStream: [:]).flatMap { $0 }
192192
XCTAssertEqual(remainingRecords.count, 1)
193193
XCTAssertTrue(remainingRecords.allSatisfy { $0.streamName != stream2 })
194194
XCTAssertEqual(remainingRecords.first { $0.streamName == stream1 }?.retryCount, 1)
@@ -197,30 +197,23 @@ class RecordClientFlushTests: XCTestCase {
197197

198198
// MARK: - Configurable Mock Sender
199199

200-
final class ConfigurableMockSender: RecordSender, @unchecked Sendable {
201-
private let lock = NSLock()
200+
actor ConfigurableMockSender: RecordSender {
202201
private var handler: (@Sendable (String, [Record]) throws -> PutRecordsResponse)?
203202
private var errorToThrow: Error?
204203

205-
func setHandler(_ handler: @escaping @Sendable (String, [Record]) throws -> PutRecordsResponse) async {
206-
lock.lock()
204+
func setHandler(_ handler: @escaping @Sendable (String, [Record]) throws -> PutRecordsResponse) {
207205
self.handler = handler
208206
errorToThrow = nil
209-
lock.unlock()
210207
}
211208

212-
func setError(_ error: Error) async {
213-
lock.lock()
209+
func setError(_ error: Error) {
214210
errorToThrow = error
215211
handler = nil
216-
lock.unlock()
217212
}
218213

219214
func putRecords(streamName: String, records: [Record]) async throws -> PutRecordsResponse {
220-
lock.lock()
221215
let currentHandler = handler
222216
let currentError = errorToThrow
223-
lock.unlock()
224217

225218
if let error = currentError {
226219
throw error

0 commit comments

Comments
 (0)