Skip to content

fix(amazonq): retry S3 upload #5208

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 11 commits into from
Dec 19, 2024
Merged

fix(amazonq): retry S3 upload #5208

Show file tree
Hide file tree
Changes from 9 commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
f18e5a9
fix(amazonq): retry S3 upload
Dec 15, 2024
85b0901
fix test
Dec 16, 2024
5182c60
fix detekt
Dec 16, 2024
71f09d8
use waitUntil
Dec 17, 2024
4e70278
update changelog
Dec 17, 2024
51aeedf
fix detekt
Dec 17, 2024
4de333b
fix detekt again
Dec 17, 2024
c2a824a
remove unused import
Dec 17, 2024
3a53721
fix tests
Dec 17, 2024
8767d37
Update .changes/next-release/bugfix-552a4676-3214-4321-a332-344bdf9f1…
dhasani23 Dec 19, 2024
f566e6e
Merge branch 'main' into retries
dhasani23 Dec 19, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions .changes/next-release/bugfix-552a4676-3214-4321-a332-344bdf9f1830.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
{
"type" : "bugfix",
"description" : "Amazon Q Code Transformation: retry initial project upload"
}
28 changes: 21 additions & 7 deletions ...nity/src/software/aws/toolkits/jetbrains/services/codemodernizer/CodeModernizerSession.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import com.intellij.openapi.application.runInEdt
import com.intellij.serviceContainer.AlreadyDisposedException
import com.intellij.util.io.HttpRequests
import kotlinx.coroutines.delay
import kotlinx.coroutines.withContext
import org.apache.commons.codec.digest.DigestUtils
import software.amazon.awssdk.core.exception.SdkClientException
import software.amazon.awssdk.services.codewhispererruntime.model.ResumeTransformationResponse
Expand All @@ -20,11 +21,13 @@ import software.amazon.awssdk.services.codewhispererruntime.model.Transformation
import software.amazon.awssdk.services.codewhispererruntime.model.TransformationUserActionStatus
import software.amazon.awssdk.services.codewhispererstreaming.model.TransformationDownloadArtifactType
import software.amazon.awssdk.services.ssooidc.model.SsoOidcException
import software.aws.toolkits.core.utils.Waiters.waitUntil
import software.aws.toolkits.core.utils.error
import software.aws.toolkits.core.utils.exists
import software.aws.toolkits.core.utils.getLogger
import software.aws.toolkits.core.utils.info
import software.aws.toolkits.core.utils.warn
import software.aws.toolkits.jetbrains.core.coroutines.getCoroutineBgContext
import software.aws.toolkits.jetbrains.services.codemodernizer.client.GumbyClient
import software.aws.toolkits.jetbrains.services.codemodernizer.commands.CodeTransformMessageListener
import software.aws.toolkits.jetbrains.services.codemodernizer.model.CodeModernizerException
Expand Down Expand Up @@ -55,6 +58,8 @@ import java.io.File
import java.io.FileInputStream
import java.io.IOException
import java.net.ConnectException
import java.net.SocketTimeoutException
import java.net.UnknownHostException
import java.nio.file.Path
import java.time.Instant
import java.util.Base64
Expand Down Expand Up @@ -142,7 +147,7 @@ class CodeModernizerSession(
*
* Based on [CodeWhispererCodeScanSession]
*/
fun createModernizationJob(copyResult: MavenCopyCommandsResult?): CodeModernizerStartJobResult {
suspend fun createModernizationJob(copyResult: MavenCopyCommandsResult?): CodeModernizerStartJobResult {
LOG.info { "Compressing local project" }
val payload: File?
var payloadSize = 0
Expand Down Expand Up @@ -377,8 +382,12 @@ class CodeModernizerSession(
/**
* Adapted from [CodeWhispererCodeScanSession]
*/
fun uploadPayload(payload: File): String {
val sha256checksum: String = Base64.getEncoder().encodeToString(DigestUtils.sha256(FileInputStream(payload)))
suspend fun uploadPayload(payload: File): String {
val sha256checksum: String = Base64.getEncoder().encodeToString(
withContext(getCoroutineBgContext()) {
DigestUtils.sha256(FileInputStream(payload))
}
)
if (isDisposed.get()) {
throw AlreadyDisposedException("Disposed when about to create upload URL")
}
Expand All @@ -394,17 +403,22 @@ class CodeModernizerSession(
throw AlreadyDisposedException("Disposed when about to upload project artifact to s3")
}
val uploadStartTime = Instant.now()
try {
waitUntil(
exceptionsToIgnore = setOf(
UnknownHostException::class,
SocketTimeoutException::class,
HttpRequests.HttpStatusException::class,
ConnectException::class
Copy link
Contributor Author

@dhasani23 dhasani23 Dec 17, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

UnknownHostException and SocketTimeoutException can be reproduced by turning off your WiFi, HttpStatusException could be a 500 or something, and I've occasionally seen ConnectException too, so I think all of these are worth retrying on.

)
) {
clientAdaptor.uploadArtifactToS3(
createUploadUrlResponse.uploadUrl(),
payload,
sha256checksum,
createUploadUrlResponse.kmsKeyArn().orEmpty(),
) { shouldStop.get() }
} catch (e: Exception) {
LOG.error { "Unexpected error when uploading project artifact to S3: $e" }
throw e // pass along error to callee
}
LOG.info { "Upload to S3 succeeded" }
Copy link
Contributor Author

@dhasani23 dhasani23 Dec 15, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not a big fan of implementing our own retry logic rather than doing it at the SDK layer, but the existing S3 client's upload functionality requires a bucket name as a part of the request, and 1) I don't think we should have the internal S3 bucket names we use floating around this repo, 2) we use multiple different buckets, and 3) we don't have access to the bucket name in the response of our CreateUploadUrl — this API would have to be changed and then we can use the existing S3 client.

Copy link
Contributor Author

@dhasani23 dhasani23 Dec 17, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Update: use waitUntil instead of a for loop; makes this retry logic even simpler

if (!shouldStop.get()) {
LOG.info { "Uploaded artifact. Latency: ${calculateTotalLatency(uploadStartTime, Instant.now())}ms" }
}
Expand Down
1 change: 1 addition & 0 deletions ...mmunity/src/software/aws/toolkits/jetbrains/services/codemodernizer/client/GumbyClient.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -148,6 +148,7 @@ class GumbyClient(private val project: Project) {
var result: CodeWhispererRuntimeResponse? = null
try {
result = apiCall()
LOG.info { "$apiName request ID: ${result.responseMetadata()?.requestId()}" }
return result
} catch (e: Exception) {
LOG.error(e) { "$apiName failed: ${e.message}" }
Expand Down
26 changes: 12 additions & 14 deletions .../aws/toolkits/jetbrains/services/codemodernizer/CodeWhispererCodeModernizerSessionTest.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -410,7 +410,7 @@ class CodeWhispererCodeModernizerSessionTest : CodeWhispererCodeModernizerTestBa
}

@Test
fun `CodeModernizer can create modernization job`() {
fun `CodeModernizer can create modernization job`() = runTest {
doReturn(ZipCreationResult.Succeeded(File("./tst-resources/codemodernizer/test.txt")))
.whenever(testSessionContextSpy).createZipWithModuleFiles(any())
doReturn(exampleCreateUploadUrlResponse).whenever(clientAdaptorSpy).createGumbyUploadUrl(any())
Expand All @@ -425,7 +425,7 @@ class CodeWhispererCodeModernizerSessionTest : CodeWhispererCodeModernizerTestBa
}

@Test
fun `CodeModernizer cannot upload payload due to already disposed`() {
fun `CodeModernizer cannot upload payload due to already disposed`() = runTest {
doReturn(ZipCreationResult.Succeeded(File("./tst-resources/codemodernizer/test.txt")))
.whenever(testSessionContextSpy).createZipWithModuleFiles(any())
doReturn(exampleCreateUploadUrlResponse).whenever(clientAdaptorSpy).createGumbyUploadUrl(any())
Expand All @@ -435,7 +435,7 @@ class CodeWhispererCodeModernizerSessionTest : CodeWhispererCodeModernizerTestBa
}

@Test
fun `CodeModernizer returns credentials expired when SsoOidcException during upload`() {
fun `CodeModernizer returns credentials expired when SsoOidcException during upload`() = runTest {
setupConnection(BearerTokenAuthState.AUTHORIZED)
doReturn(ZipCreationResult.Succeeded(File("./tst-resources/codemodernizer/test.txt")))
.whenever(testSessionContextSpy).createZipWithModuleFiles(any())
Expand All @@ -445,7 +445,7 @@ class CodeWhispererCodeModernizerSessionTest : CodeWhispererCodeModernizerTestBa
}

@Test
fun `CodeModernizer returns credentials expired when expired before upload`() {
fun `CodeModernizer returns credentials expired when expired before upload`() = runTest {
listOf(BearerTokenAuthState.NEEDS_REFRESH, BearerTokenAuthState.NOT_AUTHENTICATED).forEach {
setupConnection(it)
val result = testSessionSpy.createModernizationJob(MavenCopyCommandsResult.Success(File("./mock/path/")))
Expand All @@ -454,33 +454,31 @@ class CodeWhispererCodeModernizerSessionTest : CodeWhispererCodeModernizerTestBa
}

@Test
fun `CodeModernizer cannot upload payload due to presigned url issue`() {
fun `CodeModernizer cannot upload payload due to presigned url issue`() = runTest {
doReturn(ZipCreationResult.Succeeded(File("./tst-resources/codemodernizer/test.txt")))
.whenever(testSessionContextSpy).createZipWithModuleFiles(any())
doReturn(exampleCreateUploadUrlResponse).whenever(clientAdaptorSpy).createGumbyUploadUrl(any())
doAnswer { throw HttpRequests.HttpStatusException("mock error", 403, "mock url") }
.whenever(clientAdaptorSpy).uploadArtifactToS3(any(), any(), any(), any(), any())
doAnswer { throw HttpRequests.HttpStatusException("mock error", 403, "mock url") }.whenever(testSessionSpy).uploadPayload(any())
val result = testSessionSpy.createModernizationJob(MavenCopyCommandsResult.Success(File("./mock/path/")))
assertEquals(CodeModernizerStartJobResult.ZipUploadFailed(UploadFailureReason.PRESIGNED_URL_EXPIRED), result)
verify(testSessionStateSpy, times(1)).putJobHistory(any(), eq(TransformationStatus.FAILED), any(), any())
assertEquals(testSessionStateSpy.currentJobStatus, TransformationStatus.FAILED)
}

@Test
fun `CodeModernizer cannot upload payload due to other status code`() {
fun `CodeModernizer cannot upload payload due to other status code`() = runTest {
doReturn(ZipCreationResult.Succeeded(File("./tst-resources/codemodernizer/test.txt")))
.whenever(testSessionContextSpy).createZipWithModuleFiles(any())
doReturn(exampleCreateUploadUrlResponse).whenever(clientAdaptorSpy).createGumbyUploadUrl(any())
doAnswer { throw HttpRequests.HttpStatusException("mock error", 407, "mock url") }
.whenever(clientAdaptorSpy).uploadArtifactToS3(any(), any(), any(), any(), any())
doAnswer { throw HttpRequests.HttpStatusException("mock error", 407, "mock url") }.whenever(testSessionSpy).uploadPayload(any())
val result = testSessionSpy.createModernizationJob(MavenCopyCommandsResult.Success(File("./mock/path/")))
assertEquals(CodeModernizerStartJobResult.ZipUploadFailed(UploadFailureReason.HTTP_ERROR(407)), result)
verify(testSessionStateSpy, times(1)).putJobHistory(any(), eq(TransformationStatus.FAILED), any(), any())
assertEquals(testSessionStateSpy.currentJobStatus, TransformationStatus.FAILED)
}

@Test
fun `CodeModernizer cannot upload payload due to unknown issue`() {
fun `CodeModernizer cannot upload payload due to unknown issue`() = runTest {
doReturn(ZipCreationResult.Succeeded(File("./tst-resources/codemodernizer/test.txt")))
.whenever(testSessionContextSpy).createZipWithModuleFiles(any())
doReturn(exampleCreateUploadUrlResponse).whenever(clientAdaptorSpy).createGumbyUploadUrl(any())
Expand All @@ -492,11 +490,11 @@ class CodeWhispererCodeModernizerSessionTest : CodeWhispererCodeModernizerTestBa
}

@Test
fun `CodeModernizer cannot upload payload due to connection refused`() {
fun `CodeModernizer cannot upload payload due to connection refused`() = runTest {
doReturn(ZipCreationResult.Succeeded(File("./tst-resources/codemodernizer/test.txt")))
.whenever(testSessionContextSpy).createZipWithModuleFiles(any())
doReturn(exampleCreateUploadUrlResponse).whenever(clientAdaptorSpy).createGumbyUploadUrl(any())
doAnswer { throw ConnectException("mock exception") }.whenever(clientAdaptorSpy).uploadArtifactToS3(any(), any(), any(), any(), any())
doAnswer { throw ConnectException("mock exception") }.whenever(testSessionSpy).uploadPayload(any())
val result = testSessionSpy.createModernizationJob(MavenCopyCommandsResult.Success(File("./mock/path/")))
assertEquals(CodeModernizerStartJobResult.ZipUploadFailed(UploadFailureReason.CONNECTION_REFUSED), result)
verify(testSessionStateSpy, times(1)).putJobHistory(any(), eq(TransformationStatus.FAILED), any(), any())
Expand Down Expand Up @@ -549,7 +547,7 @@ class CodeWhispererCodeModernizerSessionTest : CodeWhispererCodeModernizerTestBa
}

@Test
fun `test uploadPayload()`() {
fun `test uploadPayload()`() = runTest {
val s3endpoint = "http://127.0.0.1:${wireMock.port()}"
val gumbyUploadUrlResponse = CreateUploadUrlResponse.builder()
.uploadUrl(s3endpoint)
Expand Down
Loading