Skip to content

Commit 5470418

Browse files
ashishrp-awslaileni-aws
ashishrp-aws
and
laileni-aws
authored
fix(amazonq): for /test adding backoff and retry for payload upload APIs. (#5310)
* Adding backoff and retry for payload upload APIs. Co-authored-by: Laxman Reddy <[email protected]>
1 parent 02ec702 commit 5470418

File tree

2 files changed

+149
-80
lines changed

2 files changed

+149
-80
lines changed

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeWhispererUTGChatManager.kt‎

Lines changed: 32 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -16,14 +16,14 @@ import kotlinx.coroutines.CoroutineScope
1616
import kotlinx.coroutines.Job
1717
import kotlinx.coroutines.delay
1818
import kotlinx.coroutines.launch
19-
import software.amazon.awssdk.core.exception.SdkServiceException
2019
import software.amazon.awssdk.services.codewhispererruntime.model.GetTestGenerationResponse
2120
import software.amazon.awssdk.services.codewhispererruntime.model.Range
2221
import software.amazon.awssdk.services.codewhispererruntime.model.StartTestGenerationResponse
2322
import software.amazon.awssdk.services.codewhispererruntime.model.TargetCode
2423
import software.amazon.awssdk.services.codewhispererruntime.model.TestGenerationJobStatus
2524
import software.amazon.awssdk.services.codewhispererstreaming.model.ExportContext
2625
import software.amazon.awssdk.services.codewhispererstreaming.model.ExportIntent
26+
import software.aws.toolkits.core.utils.Waiters.waitUntil
2727
import software.aws.toolkits.core.utils.debug
2828
import software.aws.toolkits.core.utils.error
2929
import software.aws.toolkits.core.utils.getLogger
@@ -58,6 +58,7 @@ import java.io.ByteArrayOutputStream
5858
import java.io.File
5959
import java.io.IOException
6060
import java.nio.file.Paths
61+
import java.time.Duration
6162
import java.time.Instant
6263
import java.util.concurrent.atomic.AtomicBoolean
6364
import java.util.zip.ZipInputStream
@@ -109,29 +110,38 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
109110

110111
// 2nd API call: StartTestGeneration
111112
val startTestGenerationResponse = try {
112-
startTestGeneration(
113-
uploadId = createUploadUrlResponse.uploadId(),
114-
targetCode = listOf(
115-
TargetCode.builder()
116-
.relativeTargetPath(codeTestResponseContext.currentFileRelativePath.toString())
117-
.targetLineRangeList(
118-
if (selectionRange != null) {
119-
listOf(
120-
selectionRange
113+
var response: StartTestGenerationResponse? = null
114+
115+
waitUntil(
116+
succeedOn = { response?.sdkHttpResponse()?.statusCode() == 200 },
117+
maxDuration = Duration.ofSeconds(1), // 1 second timeout
118+
) {
119+
try {
120+
response = startTestGeneration(
121+
uploadId = createUploadUrlResponse.uploadId(),
122+
targetCode = listOf(
123+
TargetCode.builder()
124+
.relativeTargetPath(codeTestResponseContext.currentFileRelativePath.toString())
125+
.targetLineRangeList(
126+
if (selectionRange != null) {
127+
listOf(selectionRange)
128+
} else {
129+
emptyList()
130+
}
121131
)
122-
} else {
123-
emptyList()
124-
}
125-
)
126-
.build()
127-
),
128-
userInput = prompt
129-
)
130-
} catch (e: Exception) {
131-
val statusCode = when {
132-
e is SdkServiceException -> e.statusCode()
133-
else -> 400
132+
.build()
133+
),
134+
userInput = prompt
135+
)
136+
delay(200)
137+
response?.testGenerationJob() != null
138+
} catch (e: Exception) {
139+
throw e
140+
}
134141
}
142+
143+
response ?: throw RuntimeException("Failed to start test generation")
144+
} catch (e: Exception) {
135145
LOG.error(e) { "Unexpected error while creating test generation job" }
136146
val errorMessage = getTelemetryErrorMessage(e, CodeWhispererConstants.FeatureName.TEST_GENERATION)
137147
throw CodeTestException(

‎plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererZipUploadManager.kt‎

Lines changed: 117 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@ import software.amazon.awssdk.services.codewhispererruntime.model.CodeAnalysisUp
1212
import software.amazon.awssdk.services.codewhispererruntime.model.CodeFixUploadContext
1313
import software.amazon.awssdk.services.codewhispererruntime.model.CreateUploadUrlRequest
1414
import software.amazon.awssdk.services.codewhispererruntime.model.CreateUploadUrlResponse
15+
import software.amazon.awssdk.services.codewhispererruntime.model.InternalServerException
16+
import software.amazon.awssdk.services.codewhispererruntime.model.ThrottlingException
1517
import software.amazon.awssdk.services.codewhispererruntime.model.UploadContext
1618
import software.amazon.awssdk.services.codewhispererruntime.model.UploadIntent
1719
import software.amazon.awssdk.utils.IoUtils
@@ -82,40 +84,50 @@ class CodeWhispererZipUploadManager(private val project: Project) {
8284
requestHeaders: Map<String, String>?,
8385
featureUseCase: CodeWhispererConstants.FeatureName,
8486
) {
85-
try {
86-
val uploadIdJson = """{"uploadId":"$uploadId"}"""
87-
HttpRequests.put(url, "application/zip").userAgent(AwsClientManager.getUserAgent()).tuner {
88-
if (requestHeaders.isNullOrEmpty()) {
89-
it.setRequestProperty(CONTENT_MD5, md5)
90-
it.setRequestProperty(CONTENT_TYPE, APPLICATION_ZIP)
91-
it.setRequestProperty(SERVER_SIDE_ENCRYPTION, AWS_KMS)
92-
if (kmsArn?.isNotEmpty() == true) {
93-
it.setRequestProperty(SERVER_SIDE_ENCRYPTION_AWS_KMS_KEY_ID, kmsArn)
94-
}
95-
it.setRequestProperty(SERVER_SIDE_ENCRYPTION_CONTEXT, Base64.getEncoder().encodeToString(uploadIdJson.toByteArray()))
96-
} else {
97-
requestHeaders.forEach { entry ->
98-
it.setRequestProperty(entry.key, entry.value)
87+
RetryableOperation<Unit>().execute(
88+
operation = {
89+
val uploadIdJson = """{"uploadId":"$uploadId"}"""
90+
HttpRequests.put(url, "application/zip").userAgent(AwsClientManager.getUserAgent()).tuner {
91+
if (requestHeaders.isNullOrEmpty()) {
92+
it.setRequestProperty(CONTENT_MD5, md5)
93+
it.setRequestProperty(CONTENT_TYPE, APPLICATION_ZIP)
94+
it.setRequestProperty(SERVER_SIDE_ENCRYPTION, AWS_KMS)
95+
if (kmsArn?.isNotEmpty() == true) {
96+
it.setRequestProperty(SERVER_SIDE_ENCRYPTION_AWS_KMS_KEY_ID, kmsArn)
97+
}
98+
it.setRequestProperty(SERVER_SIDE_ENCRYPTION_CONTEXT, Base64.getEncoder().encodeToString(uploadIdJson.toByteArray()))
99+
} else {
100+
requestHeaders.forEach { entry ->
101+
it.setRequestProperty(entry.key, entry.value)
102+
}
99103
}
104+
}.connect {
105+
val connection = it.connection as HttpURLConnection
106+
connection.setFixedLengthStreamingMode(fileToUpload.length())
107+
IoUtils.copy(fileToUpload.inputStream(), connection.outputStream)
108+
}
109+
},
110+
isRetryable = { e ->
111+
when (e) {
112+
is IOException -> true
113+
else -> false
114+
}
115+
},
116+
errorHandler = { e, attempts ->
117+
val errorMessage = getTelemetryErrorMessage(e, featureUseCase)
118+
when (featureUseCase) {
119+
CodeWhispererConstants.FeatureName.CODE_REVIEW ->
120+
codeScanServerException("CreateUploadUrlException: $errorMessage")
121+
CodeWhispererConstants.FeatureName.TEST_GENERATION ->
122+
throw CodeTestException(
123+
"UploadTestArtifactToS3Error: $errorMessage",
124+
"UploadTestArtifactToS3Error",
125+
message("testgen.error.generic_technical_error_message")
126+
)
127+
else -> throw RuntimeException("$errorMessage (after $attempts attempts)")
100128
}
101-
}.connect {
102-
val connection = it.connection as HttpURLConnection
103-
connection.setFixedLengthStreamingMode(fileToUpload.length())
104-
IoUtils.copy(fileToUpload.inputStream(), connection.outputStream)
105-
}
106-
} catch (e: Exception) {
107-
LOG.debug { "$featureUseCase: Artifact failed to upload in the S3 bucket: ${e.message}" }
108-
val errorMessage = getTelemetryErrorMessage(e, featureUseCase)
109-
when (featureUseCase) {
110-
CodeWhispererConstants.FeatureName.CODE_REVIEW -> codeScanServerException("CreateUploadUrlException: $errorMessage")
111-
CodeWhispererConstants.FeatureName.TEST_GENERATION -> throw CodeTestException(
112-
"UploadTestArtifactToS3Error: $errorMessage",
113-
"UploadTestArtifactToS3Error",
114-
message("testgen.error.generic_technical_error_message")
115-
)
116-
else -> throw RuntimeException(errorMessage) // Adding else for safety check
117129
}
118-
}
130+
)
119131
}
120132

121133
fun createUploadUrl(
@@ -124,35 +136,44 @@ class CodeWhispererZipUploadManager(private val project: Project) {
124136
uploadTaskType: CodeWhispererConstants.UploadTaskType,
125137
taskName: String,
126138
featureUseCase: CodeWhispererConstants.FeatureName,
127-
): CreateUploadUrlResponse = try {
128-
CodeWhispererClientAdaptor.getInstance(project).createUploadUrl(
129-
CreateUploadUrlRequest.builder()
130-
.contentMd5(md5Content)
131-
.artifactType(artifactType)
132-
.uploadIntent(getUploadIntent(uploadTaskType))
133-
.uploadContext(
134-
// For UTG we don't need uploadContext but sending else case as UploadContext
135-
if (uploadTaskType == CodeWhispererConstants.UploadTaskType.CODE_FIX) {
136-
UploadContext.fromCodeFixUploadContext(CodeFixUploadContext.builder().codeFixName(taskName).build())
137-
} else {
138-
UploadContext.fromCodeAnalysisUploadContext(CodeAnalysisUploadContext.builder().codeScanName(taskName).build())
139-
}
140-
)
141-
.build()
142-
)
143-
} catch (e: Exception) {
144-
LOG.debug { "$featureUseCase: Create Upload URL failed: ${e.message}" }
145-
val errorMessage = getTelemetryErrorMessage(e, featureUseCase)
146-
when (featureUseCase) {
147-
CodeWhispererConstants.FeatureName.CODE_REVIEW -> codeScanServerException("CreateUploadUrlException: $errorMessage")
148-
CodeWhispererConstants.FeatureName.TEST_GENERATION -> throw CodeTestException(
149-
"CreateUploadUrlError: $errorMessage",
150-
"CreateUploadUrlError",
151-
message("testgen.error.generic_technical_error_message")
139+
): CreateUploadUrlResponse = RetryableOperation<CreateUploadUrlResponse>().execute(
140+
operation = {
141+
CodeWhispererClientAdaptor.getInstance(project).createUploadUrl(
142+
CreateUploadUrlRequest.builder()
143+
.contentMd5(md5Content)
144+
.artifactType(artifactType)
145+
.uploadIntent(getUploadIntent(uploadTaskType))
146+
.uploadContext(
147+
// For UTG we don't need uploadContext but sending else case as UploadContext
148+
if (uploadTaskType == CodeWhispererConstants.UploadTaskType.CODE_FIX) {
149+
UploadContext.fromCodeFixUploadContext(CodeFixUploadContext.builder().codeFixName(taskName).build())
150+
} else {
151+
UploadContext.fromCodeAnalysisUploadContext(CodeAnalysisUploadContext.builder().codeScanName(taskName).build())
152+
}
153+
)
154+
.build()
152155
)
153-
else -> throw RuntimeException(errorMessage) // Adding else for safety check
156+
},
157+
isRetryable = { e ->
158+
e is ThrottlingException || e is InternalServerException
159+
},
160+
errorHandler = { e, attempts ->
161+
val errorMessage = getTelemetryErrorMessage(e, featureUseCase)
162+
when (featureUseCase) {
163+
CodeWhispererConstants.FeatureName.CODE_REVIEW ->
164+
codeScanServerException("CreateUploadUrlException after $attempts attempts: $errorMessage")
165+
166+
CodeWhispererConstants.FeatureName.TEST_GENERATION ->
167+
throw CodeTestException(
168+
"CreateUploadUrlError after $attempts attempts: $errorMessage",
169+
"CreateUploadUrlError",
170+
message("testgen.error.generic_technical_error_message")
171+
)
172+
173+
else -> throw RuntimeException("$errorMessage (after $attempts attempts)")
174+
}
154175
}
155-
}
176+
)
156177

157178
private fun getUploadIntent(uploadTaskType: CodeWhispererConstants.UploadTaskType): UploadIntent = when (uploadTaskType) {
158179
CodeWhispererConstants.UploadTaskType.SCAN_FILE -> UploadIntent.AUTOMATIC_FILE_SECURITY_SCAN
@@ -187,3 +208,41 @@ fun getTelemetryErrorMessage(e: Exception, featureUseCase: CodeWhispererConstant
187208
else -> message("testgen.message.failed")
188209
}
189210
}
211+
212+
class RetryableOperation<T> {
213+
private var attempts = 0
214+
private var currentDelay = INITIAL_DELAY
215+
private var lastException: Exception? = null
216+
217+
fun execute(
218+
operation: () -> T,
219+
isRetryable: (Exception) -> Boolean,
220+
errorHandler: (Exception, Int) -> Nothing,
221+
): T {
222+
while (attempts < MAX_RETRY_ATTEMPTS) {
223+
try {
224+
return operation()
225+
} catch (e: Exception) {
226+
lastException = e
227+
228+
attempts++
229+
if (attempts < MAX_RETRY_ATTEMPTS && isRetryable(e)) {
230+
Thread.sleep(currentDelay)
231+
currentDelay = (currentDelay * 2).coerceAtMost(MAX_BACKOFF)
232+
continue
233+
}
234+
235+
errorHandler(e, attempts)
236+
}
237+
}
238+
239+
// This line should never be reached due to errorHandler throwing exception
240+
throw RuntimeException("Unexpected state after $attempts attempts")
241+
}
242+
243+
companion object {
244+
private const val INITIAL_DELAY = 100L // milliseconds
245+
private const val MAX_BACKOFF = 10000L // milliseconds
246+
private const val MAX_RETRY_ATTEMPTS = 3
247+
}
248+
}

0 commit comments

Comments
 (0)