Fix race condition in WebSocketCoreImpl by kmod-midori · Pull Request #6948 · home-assistant/android

kmod-midori · 2026-06-01T11:35:36Z

Summary

Fix #6947. Upon close inspection there is a race condition in WebSocketCoreImpl.

android/common/src/main/kotlin/io/homeassistant/companion/android/common/data/websocket/impl/WebSocketCoreImpl.kt

Lines 764 to 769 in ba2371c

    
           if (activeMessages.isEmpty()) { 
        
               Timber.i("No more subscriptions, closing connection.") 
        
               connectionHolder.get()?.webSocket?.close(1001, "Done listening to subscriptions.") 
        
           } else { 
        
               Timber.i("Still ${activeMessages.size} messages in the queue, not closing connection.") 
        
           }

Here the WS connection is closed (no more message can be sent). However, it triggers onClosing and onClosed asynchronously, which means that when it returns, WebSocketCoreImpl still considers the socket connected.

android/common/src/main/kotlin/io/homeassistant/companion/android/common/data/websocket/impl/WebSocketCoreImpl.kt

Lines 305 to 328 in ba2371c

    
           suspend fun connect(): Boolean { 
        
               val connectDeferred: CompletableDeferred<Boolean> 
        
               // Track pending connection state locally within the urlObserverJob scope. 
        
               // This is read across collectLatest invocations to cancel partial connections on URL change. 
        
               var pendingWebSocket: WebSocket? = null 
        
               var urlObserverJob: Job? = null 
        
               connectedMutex.withLock { 
        
                   // Already connected? 
        
                   if (connectionHolder.get() != null && authCompleted.isCompleted) { 
        
                       if (authCompleted.isCancelled) Timber.w("Trying to connect but was cancelled") 
        
                       return !authCompleted.isCancelled 
        
                   } 
        
                   // Connection already in progress? Reuse its deferred 
        
                   pendingConnectDeferred?.takeIf { !it.isCompleted }?.let { existing -> 
        
                       Timber.d("Connection already in progress, reusing existing deferred and release lock") 
        
                       connectDeferred = existing 
        
                       return@withLock 
        
                   } 
        
                   // Start new connection attempt 
        
                   connectDeferred = CompletableDeferred()

If a new subscription is made after close and before onClosing, connect() sees that connectionHolder still contains something and return without making a new connection. Therefore, the new messages will be sent to the closing/closed connection, causing exceptions.

Xiaomi's SystemUI implementation creates and cancels subscriptions in HaControlsProviderService in quick succession, therefore triggering this.

Checklist

New or updated tests have been added to cover the changes following the testing guidelines.
I have no idea how to create a unit test to cover such a race condition, but at least I no longer see exceptions (consistently) after these changes.
The code follows the project's code style and best_practices.
The changes have been thoroughly tested, and edge cases have been considered.
Changes are backward compatible whenever feasible. Any breaking changes are documented in the changelog for users and/or in the code for developers depending on the relevance.

Any other notes

N/A

Copilot

Pull request overview

Note

Copilot was unable to run its full agentic suite in this review.

Adds a single, centralized close(...) helper to gracefully tear down the current WebSocket session and reset internal connection state so future connection attempts can proceed.

Changes:

Introduced suspend fun close(code, reason) that cancels any pending connection attempt and clears the current connection reference.
Updated the “no more subscriptions” path to use the new close(...) helper instead of calling webSocket.close(...) directly.

kmod-midori · 2026-06-17T11:20:36Z

Hi, is there anything I need to do?

TimoPtr · 2026-06-17T14:06:17Z

Hi, is there anything I need to do?

We need a test that replicate the issue (failing if we remove your logic and pass when adding it). I didn't look at the test in the comment but if you think it does what we ask submit it and put the PR back in ready for review.

kmod-midori · 2026-06-18T04:18:30Z

@TimoPtr The tests are ready. It fails on ba2371c (before this PR) and passes now.

home-assistant · 2026-06-18T10:04:19Z

Please take a look at the requested changes, and use the Ready for review button when you are done, thanks 👍

Learn more about our pull request process.

…sure all resources are properly cleaned up

TimoPtr

The most important question I have is that did you test this change on your scenario where it was an issue?

kmod-midori · 2026-06-23T13:55:35Z

Yes, and it fixes my problem.

Copilot

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.

kmod-midori · 2026-06-23T14:11:12Z

+    private fun close(code: Int, reason: String?) = wsScope.launch {
+        connectedMutex.withLock {
+            val holder = connectionHolder.getAndSet(null)
+            holder?.webSocket?.close(code, reason)
+            holder?.urlObserverJob?.cancel()


@TimoPtr Should I do this? I've never used CoroutineStart.UNDISPATCHED.

the window is very very small, IO is not limited in terms of number of concurrent execution like Default is, so it should be exectuted right away. And I also never used CoroutineStart.UNDISPATCHED

TimoPtr

From the code I don't see any issue. Let's wait @jpelgrom for a sanity check.

jpelgrom

I don't see any issues with it either. Thanks for taking the time to work on this, we wouldn't have found this with Xiaomi's specific behavior being the reason.

kmod-midori added 2 commits June 1, 2026 19:03

Fix race conditions in WebSocketCoreImpl

9a75c21

Undo move

58f57a6

Copilot AI review requested due to automatic review settings June 1, 2026 11:35

home-assistant Bot added the cla-signed label Jun 1, 2026

Copilot AI reviewed Jun 1, 2026

View reviewed changes

Comment thread ...in/kotlin/io/homeassistant/companion/android/common/data/websocket/impl/WebSocketCoreImpl.kt Outdated

Comment thread ...in/kotlin/io/homeassistant/companion/android/common/data/websocket/impl/WebSocketCoreImpl.kt Outdated

TimoPtr reviewed Jun 1, 2026

View reviewed changes

Comment thread ...in/kotlin/io/homeassistant/companion/android/common/data/websocket/impl/WebSocketCoreImpl.kt Outdated

Merge branch 'home-assistant:main' into bugfix/ws-race

b7bb8f2

TimoPtr marked this pull request as draft June 1, 2026 14:52

Make close private

9d38ad7

kmod-midori marked this pull request as ready for review June 17, 2026 11:20

TimoPtr marked this pull request as draft June 17, 2026 14:06

Add reconnectsAfterShutdown test

f667bf0

kmod-midori marked this pull request as ready for review June 17, 2026 14:09

Fix close and test

4a8793e

Merge branch 'main' into bugfix/ws-race

5a4a2b1

TimoPtr requested changes Jun 18, 2026

View reviewed changes

Comment thread ...otlin/io/homeassistant/companion/android/common/data/websocket/impl/WebSocketCoreImplTest.kt Outdated

Comment thread ...in/kotlin/io/homeassistant/companion/android/common/data/websocket/impl/WebSocketCoreImpl.kt Outdated

home-assistant Bot marked this pull request as draft June 18, 2026 10:04

Make it explicit that we are intentionally closing the socket, and en…

fe3055b

…sure all resources are properly cleaned up

kmod-midori marked this pull request as ready for review June 18, 2026 15:19

home-assistant Bot requested a review from TimoPtr June 18, 2026 15:19

kmod-midori added 2 commits June 18, 2026 23:20

Merge branch 'main' into bugfix/ws-race

7fe234c

Merge branch 'main' into bugfix/ws-race

98e4038

TimoPtr requested a review from Copilot June 23, 2026 13:52

TimoPtr reviewed Jun 23, 2026

View reviewed changes

Copilot started reviewing on behalf of TimoPtr June 23, 2026 13:53 View session

Copilot AI reviewed Jun 23, 2026

View reviewed changes

kmod-midori added 2 commits June 23, 2026 22:09

Do not throw exception in pendingConnectDeferred

7937193

Use runCurrent so we are not advancing time

af1b9b4

TimoPtr reviewed Jun 23, 2026

View reviewed changes

Comment thread ...in/kotlin/io/homeassistant/companion/android/common/data/websocket/impl/WebSocketCoreImpl.kt

TimoPtr approved these changes Jun 23, 2026

View reviewed changes

jpelgrom approved these changes Jun 23, 2026

View reviewed changes

jpelgrom merged commit 35f447a into home-assistant:main Jun 23, 2026
20 checks passed

	if (activeMessages.isEmpty()) {
	Timber.i("No more subscriptions, closing connection.")
	connectionHolder.get()?.webSocket?.close(1001, "Done listening to subscriptions.")
	} else {
	Timber.i("Still ${activeMessages.size} messages in the queue, not closing connection.")
	}

	suspend fun connect(): Boolean {
	val connectDeferred: CompletableDeferred<Boolean>

	// Track pending connection state locally within the urlObserverJob scope.
	// This is read across collectLatest invocations to cancel partial connections on URL change.
	var pendingWebSocket: WebSocket? = null
	var urlObserverJob: Job? = null

	connectedMutex.withLock {
	// Already connected?
	if (connectionHolder.get() != null && authCompleted.isCompleted) {
	if (authCompleted.isCancelled) Timber.w("Trying to connect but was cancelled")
	return !authCompleted.isCancelled
	}

	// Connection already in progress? Reuse its deferred
	pendingConnectDeferred?.takeIf { !it.isCompleted }?.let { existing ->
	Timber.d("Connection already in progress, reusing existing deferred and release lock")
	connectDeferred = existing
	return@withLock
	}

	// Start new connection attempt
	connectDeferred = CompletableDeferred()

Uh oh!

Conversation

kmod-midori commented Jun 1, 2026

Summary

Checklist

Any other notes

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

kmod-midori commented Jun 17, 2026

Uh oh!

TimoPtr commented Jun 17, 2026

Uh oh!

kmod-midori commented Jun 18, 2026

Uh oh!

Uh oh!

Uh oh!

home-assistant Bot commented Jun 18, 2026

Uh oh!

TimoPtr left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kmod-midori commented Jun 23, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

kmod-midori Jun 23, 2026

Choose a reason for hiding this comment

Uh oh!

TimoPtr Jun 23, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

TimoPtr left a comment

Choose a reason for hiding this comment

Uh oh!

jpelgrom left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

TimoPtr left a comment •

edited

Loading