diff --git a/Source/AwsCommonRuntimeKit/io/retryer/RetryStrategy.swift b/Source/AwsCommonRuntimeKit/io/retryer/RetryStrategy.swift index c8eb6df04..46b12a48f 100644 --- a/Source/AwsCommonRuntimeKit/io/retryer/RetryStrategy.swift +++ b/Source/AwsCommonRuntimeKit/io/retryer/RetryStrategy.swift @@ -71,8 +71,8 @@ public class RetryStrategy { } } - public func scheduleRetry(token: RetryToken, errorType: RetryError) async throws { - try await withCheckedThrowingContinuation({ (continuation: CheckedContinuation<(), Error>) in + public func scheduleRetry(token: RetryToken, errorType: RetryError) async throws -> RetryToken { + try await withCheckedThrowingContinuation { continuation in let continuationCore = ContinuationCore(continuation: continuation) if aws_retry_strategy_schedule_retry(token.rawValue, errorType.rawValue, @@ -81,7 +81,7 @@ public class RetryStrategy { continuationCore.release() continuation.resume(throwing: CommonRunTimeError.crtError(.makeFromLastError())) } - }) + } } /// Records a successful retry.You should always call it after a successful operation @@ -98,14 +98,15 @@ public class RetryStrategy { private func onRetryReady(token: UnsafeMutablePointer?, errorCode: Int32, userData: UnsafeMutableRawPointer!) { - let crtRetryStrategyCore = Unmanaged>.fromOpaque(userData).takeRetainedValue() + let crtRetryStrategyCore = Unmanaged>.fromOpaque(userData).takeRetainedValue() if errorCode != AWS_OP_SUCCESS { crtRetryStrategyCore.continuation.resume(throwing: CommonRunTimeError.crtError(CRTError(code: errorCode))) return } // Success - crtRetryStrategyCore.continuation.resume() + aws_retry_token_acquire(token!) + crtRetryStrategyCore.continuation.resume(returning: RetryToken(rawValue: token!)) } private func onRetryTokenAcquired(retry_strategy: UnsafeMutablePointer?, diff --git a/Test/AwsCommonRuntimeKitTests/io/RetryerTests.swift b/Test/AwsCommonRuntimeKitTests/io/RetryerTests.swift index da55b4dc9..989ed83dd 100644 --- a/Test/AwsCommonRuntimeKitTests/io/RetryerTests.swift +++ b/Test/AwsCommonRuntimeKitTests/io/RetryerTests.swift @@ -22,6 +22,6 @@ class RetryerTests: XCBaseTestCase { let retryer = try RetryStrategy(eventLoopGroup: elg, allocator: allocator) let token = try await retryer.acquireToken(partitionId: "partition1") XCTAssertNotNil(token) - try await retryer.scheduleRetry(token: token, errorType: RetryError.serverError) + _ = try await retryer.scheduleRetry(token: token, errorType: RetryError.serverError) } }