Skip to content

Commit 36ee79e

Browse files
authored
fix(amazonq): fix connection manager deadlock when q ListAvailableProfiles is long (#5547)
`CodeWhispererStatusBarWidget` attempts to query active connections, but lock is being held while trying to list profiles ``` "AWT-EventQueue-0" prio=0 tid=0x0 nid=0x0 blocked java.lang.Thread.State: BLOCKED on software.aws.toolkits.jetbrains.core.credentials.DefaultToolkitConnectionManager@37995400 owned by "ApplicationImpl pooled thread 5 @coroutine#17104" Id=128 at software.aws.toolkits.jetbrains.core.credentials.DefaultToolkitConnectionManager.activeConnectionForFeature(DefaultToolkitConnectionManager.kt) at software.aws.toolkits.jetbrains.services.amazonq.QUtilsKt.calculateIfIamIdentityCenterConnection(QUtils.kt:18) at software.aws.toolkits.jetbrains.services.codewhisperer.customization.DefaultCodeWhispererModelConfigurator.activeCustomization(CodeWhispererModelConfigurator.kt:171) at software.aws.toolkits.jetbrains.services.codewhisperer.status.CodeWhispererStatusBarWidget.getSelectedValue(CodeWhispererStatusBarWidget.kt:114) at com.intellij.openapi.wm.impl.status.MultipleTextValues.beforeUpdate(IdeStatusBarImpl.kt:983) ``` ``` - "coroutine#17104":BlockingCoroutine{Active}@7dced952, state: RUNNING [CoroutineId(17104), BlockingEventLoop@2da7461c] at java.base/jdk.internal.misc.Unsafe.park(Native Method) at java.base/java.util.concurrent.locks.LockSupport.parkNanos(LockSupport.java:269) at java.base/java.util.concurrent.CompletableFuture$Signaller.block(CompletableFuture.java:1866) at java.base/java.util.concurrent.ForkJoinPool.unmanagedBlock(ForkJoinPool.java:4013) at java.base/java.util.concurrent.ForkJoinPool.managedBlock(ForkJoinPool.java:3961) at java.base/java.util.concurrent.CompletableFuture.timedGet(CompletableFuture.java:1939) at java.base/java.util.concurrent.CompletableFuture.get(CompletableFuture.java:2095) at migration.software.aws.toolkits.jetbrains.core.AwsResourceCache$Companion.wait(AwsResourceCache.kt:150) at migration.software.aws.toolkits.jetbrains.core.AwsResourceCache$Companion.access$wait(AwsResourceCache.kt:145) at migration.software.aws.toolkits.jetbrains.core.AwsResourceCache.getResourceNow(AwsResourceCache.kt:92) at migration.software.aws.toolkits.jetbrains.core.AwsResourceCache.getResourceNow(AwsResourceCache.kt:105) at software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager.listRegionProfiles(QRegionProfileManager.kt:70) at software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager.validateProfile$lambda$0(QRegionProfileManager.kt:50) at software.aws.toolkits.core.utils.ExceptionUtils.tryOrNull(ExceptionUtils.kt:11) at software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager.validateProfile(QRegionProfileManager.kt:49) at software.aws.toolkits.jetbrains.services.amazonq.toolwindow.AmazonQToolWindowFactory$createToolWindowContent$2.activeConnectionChanged(AmazonQToolWindowFactory.kt:68) [...] at software.aws.toolkits.jetbrains.core.credentials.DefaultToolkitConnectionManager.switchConnection(DefaultToolkitConnectionManager.kt:159) ```
1 parent 289a722 commit 36ee79e

File tree

4 files changed

+47
-9
lines changed

4 files changed

+47
-9
lines changed
Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
{
2+
"type" : "bugfix",
3+
"description" : "Fix issue where IDE freezes when logging into Amazon Q"
4+
}

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/toolwindow/AmazonQToolWindowFactory.kt

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,6 @@ class AmazonQToolWindowFactory : ToolWindowFactory, DumbAware {
6565
override fun activeConnectionChanged(newConnection: ToolkitConnection?) {
6666
ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())?.let { qConn ->
6767
openMeetQPage(project)
68-
QRegionProfileManager.getInstance().validateProfile(project)
6968
}
7069
prepareChatContent(project, qPanel)
7170
}

plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/QRegionProfileManagerTest.kt

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.profile.QProfileSwitchIn
4040
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile
4141
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager
4242
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileSelectedListener
43+
import software.aws.toolkits.jetbrains.utils.satisfiesKt
4344
import software.aws.toolkits.jetbrains.utils.xmlElement
4445
import java.net.URI
4546
import java.util.function.Consumer
@@ -105,6 +106,23 @@ class QRegionProfileManagerTest {
105106
assertThat(sut.activeProfile(project)).isNull()
106107
}
107108

109+
@Test
110+
fun `data is cleared when user logs out`() {
111+
sut.switchProfile(project, QRegionProfile(arn = "arn", profileName = "foo_profile"), QProfileSwitchIntent.User)
112+
assertThat(sut.activeProfile(project)).isEqualTo(QRegionProfile(arn = "arn", profileName = "foo_profile"))
113+
114+
ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())?.let {
115+
if (it is AwsBearerTokenConnection) {
116+
logoutFromSsoConnection(project, it)
117+
}
118+
}
119+
120+
assertThat(sut.state).satisfiesKt {
121+
assertThat(it.connectionIdToActiveProfile).isEmpty()
122+
assertThat(it.connectionIdToProfileList).isEmpty()
123+
}
124+
}
125+
108126
@Test
109127
fun `switch should send message onProfileChanged for active switch`() {
110128
var cnt = 0

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QRegionProfileManager.kt

Lines changed: 25 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -4,13 +4,15 @@
44
package software.aws.toolkits.jetbrains.services.amazonq.profile
55

66
import com.intellij.openapi.Disposable
7+
import com.intellij.openapi.application.ApplicationManager
78
import com.intellij.openapi.components.BaseState
89
import com.intellij.openapi.components.PersistentStateComponent
910
import com.intellij.openapi.components.Service
1011
import com.intellij.openapi.components.State
1112
import com.intellij.openapi.components.Storage
1213
import com.intellij.openapi.components.service
1314
import com.intellij.openapi.project.Project
15+
import com.intellij.util.concurrency.annotations.RequiresBackgroundThread
1416
import com.intellij.util.xmlb.annotations.MapAnnotation
1517
import com.intellij.util.xmlb.annotations.Property
1618
import software.amazon.awssdk.core.SdkClient
@@ -25,6 +27,7 @@ import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
2527
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
2628
import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
2729
import software.aws.toolkits.jetbrains.core.credentials.sono.isSono
30+
import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProviderListener
2831
import software.aws.toolkits.jetbrains.core.region.AwsRegionProvider
2932
import software.aws.toolkits.jetbrains.utils.notifyInfo
3033
import software.aws.toolkits.resources.AmazonQBundle.message
@@ -40,9 +43,23 @@ class QRegionProfileManager : PersistentStateComponent<QProfileState>, Disposabl
4043

4144
// Map to store connectionId to its active profile
4245
private val connectionIdToActiveProfile = Collections.synchronizedMap<String, QRegionProfile>(mutableMapOf())
43-
private val connectionIdToProfileList = mutableMapOf<String, Int>()
46+
private val connectionIdToProfileCount = mutableMapOf<String, Int>()
47+
48+
init {
49+
ApplicationManager.getApplication().messageBus.connect(this)
50+
.subscribe(
51+
BearerTokenProviderListener.TOPIC,
52+
object : BearerTokenProviderListener {
53+
override fun invalidate(providerId: String) {
54+
connectionIdToActiveProfile.remove(providerId)
55+
connectionIdToProfileCount.remove(providerId)
56+
}
57+
}
58+
)
59+
}
4460

4561
// should be call on project startup to validate if profile is still active
62+
@RequiresBackgroundThread
4663
fun validateProfile(project: Project) {
4764
val conn = getIdcConnectionOrNull(project)
4865
val selected = activeProfile(project) ?: return
@@ -78,7 +95,7 @@ class QRegionProfileManager : PersistentStateComponent<QProfileState>, Disposabl
7895
switchProfile(project, mappedProfiles.first(), intent = QProfileSwitchIntent.Update)
7996
}
8097
mappedProfiles.takeIf { it.isNotEmpty() }?.also {
81-
connectionIdToProfileList[connection.id] = it.size
98+
connectionIdToProfileCount[connection.id] = it.size
8299
} ?: error("You don't have access to the resource")
83100
} catch (e: Exception) {
84101
LOG.warn(e) { "Failed to list region profiles: ${e.message}" }
@@ -110,7 +127,7 @@ class QRegionProfileManager : PersistentStateComponent<QProfileState>, Disposabl
110127
Telemetry.amazonq.didSelectProfile.use { span ->
111128
span.source(intent.value)
112129
.amazonQProfileRegion(newProfile.region)
113-
.profileCount(connectionIdToProfileList[conn.id])
130+
.profileCount(connectionIdToProfileCount[conn.id])
114131
.ssoRegion(conn.region)
115132
.credentialStartUrl(conn.startUrl)
116133
.result(MetricResult.Succeeded)
@@ -139,13 +156,13 @@ class QRegionProfileManager : PersistentStateComponent<QProfileState>, Disposabl
139156

140157
// for each idc connection, user should have a profile, otherwise should show the profile selection error page
141158
fun isPendingProfileSelection(project: Project): Boolean = getIdcConnectionOrNull(project)?.let { conn ->
142-
val profileCounts = connectionIdToProfileList[conn.id] ?: 0
159+
val profileCounts = connectionIdToProfileCount[conn.id] ?: 0
143160
val activeProfile = connectionIdToActiveProfile[conn.id]
144161
profileCounts == 0 || (profileCounts > 1 && activeProfile?.arn.isNullOrEmpty())
145162
} ?: false
146163

147164
fun shouldDisplayProfileInfo(project: Project): Boolean = getIdcConnectionOrNull(project)?.let { conn ->
148-
(connectionIdToProfileList[conn.id] ?: 0) > 1
165+
(connectionIdToProfileCount[conn.id] ?: 0) > 1
149166
} ?: false
150167

151168
fun getQClientSettings(project: Project): TokenConnectionSettings {
@@ -191,16 +208,16 @@ class QRegionProfileManager : PersistentStateComponent<QProfileState>, Disposabl
191208
override fun getState(): QProfileState {
192209
val state = QProfileState()
193210
state.connectionIdToActiveProfile.putAll(this.connectionIdToActiveProfile)
194-
state.connectionIdToProfileList.putAll(this.connectionIdToProfileList)
211+
state.connectionIdToProfileList.putAll(this.connectionIdToProfileCount)
195212
return state
196213
}
197214

198215
override fun loadState(state: QProfileState) {
199216
connectionIdToActiveProfile.clear()
200217
connectionIdToActiveProfile.putAll(state.connectionIdToActiveProfile)
201218

202-
connectionIdToProfileList.clear()
203-
connectionIdToProfileList.putAll(state.connectionIdToProfileList)
219+
connectionIdToProfileCount.clear()
220+
connectionIdToProfileCount.putAll(state.connectionIdToProfileList)
204221
}
205222
}
206223

0 commit comments

Comments
 (0)