Update schedule retry API to return retry token (#125)

awslabs · Dec 15, 2022 · be5bda4 · be5bda4
1 parent 62a3102
commit be5bda4
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git 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<aws_retry_token>?,
                           errorCode: Int32,
                           userData: UnsafeMutableRawPointer!) {
-    let crtRetryStrategyCore = Unmanaged<ContinuationCore<()>>.fromOpaque(userData).takeRetainedValue()
+    let crtRetryStrategyCore = Unmanaged<ContinuationCore<RetryToken>>.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<aws_retry_strategy>?,

diff --git 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)
     }
 }