From db4def45f94908209be211ab27931ade0e334323 Mon Sep 17 00:00:00 2001 From: Cody Henthorne Date: Wed, 13 May 2026 10:33:08 -0400 Subject: [PATCH] Reduce blocking calls in happy network call path. --- .../api/websocket/SignalWebSocket.kt | 19 +++- .../websocket/LibSignalChatConnection.kt | 98 ++++++++++--------- 2 files changed, 70 insertions(+), 47 deletions(-) diff --git a/lib/libsignal-service/src/main/java/org/whispersystems/signalservice/api/websocket/SignalWebSocket.kt b/lib/libsignal-service/src/main/java/org/whispersystems/signalservice/api/websocket/SignalWebSocket.kt index bf1648dd6a..5bac0f70cb 100644 --- a/lib/libsignal-service/src/main/java/org/whispersystems/signalservice/api/websocket/SignalWebSocket.kt +++ b/lib/libsignal-service/src/main/java/org/whispersystems/signalservice/api/websocket/SignalWebSocket.kt @@ -64,6 +64,7 @@ sealed class SignalWebSocket( const val FOREGROUND_KEEPALIVE = "Foregrounded" } + @Volatile private var connection: WebSocketConnection? = null val connectionName get() = connection?.name ?: "[null]" @@ -75,6 +76,8 @@ sealed class SignalWebSocket( private val keepAliveChangeListeners: MutableSet = CopyOnWriteArraySet() private val scope = CoroutineScope(SupervisorJob() + Dispatchers.IO) + + @Volatile private var delayedDisconnectJob: Job? = null val state: Observable = _state @@ -209,13 +212,24 @@ sealed class SignalWebSocket( } } - @Synchronized @Throws(WebSocketUnavailableException::class) protected fun getWebSocket(): WebSocketConnection { if (!canConnect.canConnect()) { throw WebSocketUnavailableException() } + connection?.takeIf { !it.isDead() }?.let { return it } + + return getOrCreateWebSocketLocked() + } + + @Synchronized + @Throws(WebSocketUnavailableException::class) + private fun getOrCreateWebSocketLocked(): WebSocketConnection { + if (!canConnect.canConnect()) { + throw WebSocketUnavailableException() + } + if (connection == null || connection?.isDead() == true) { connection?.shutdown() disposable.dispose() @@ -252,6 +266,9 @@ sealed class SignalWebSocket( } private fun restartDelayedDisconnectIfNecessary() { + if (delayedDisconnectJob?.isActive != true) { + return + } synchronized(this) { if (delayedDisconnectJob?.isActive == true) { startDelayedDisconnectIfNecessary() diff --git a/lib/libsignal-service/src/main/java/org/whispersystems/signalservice/internal/websocket/LibSignalChatConnection.kt b/lib/libsignal-service/src/main/java/org/whispersystems/signalservice/internal/websocket/LibSignalChatConnection.kt index d69afe7942..5418c17a19 100644 --- a/lib/libsignal-service/src/main/java/org/whispersystems/signalservice/internal/websocket/LibSignalChatConnection.kt +++ b/lib/libsignal-service/src/main/java/org/whispersystems/signalservice/internal/websocket/LibSignalChatConnection.kt @@ -99,6 +99,8 @@ class LibSignalChatConnection( // chatConnectionFuture: Set only when state == CONNECTING private val CHAT_SERVICE_LOCK = ReentrantLock() private val stateChangedOrMessageReceivedCondition = CHAT_SERVICE_LOCK.newCondition() + + @Volatile private var chatConnection: ChatConnection? = null private var chatConnectionFuture: CompletableFuture? = null @@ -320,19 +322,17 @@ class LibSignalChatConnection( } override fun isDead(): Boolean { - CHAT_SERVICE_LOCK.withLock { - return when (state.value) { - WebSocketConnectionState.DISCONNECTED, - WebSocketConnectionState.DISCONNECTING, - WebSocketConnectionState.FAILED, - WebSocketConnectionState.AUTHENTICATION_FAILED, - WebSocketConnectionState.REMOTE_DEPRECATED -> true + return when (state.value) { + WebSocketConnectionState.DISCONNECTED, + WebSocketConnectionState.DISCONNECTING, + WebSocketConnectionState.FAILED, + WebSocketConnectionState.AUTHENTICATION_FAILED, + WebSocketConnectionState.REMOTE_DEPRECATED -> true - WebSocketConnectionState.CONNECTING, - WebSocketConnectionState.CONNECTED -> false + WebSocketConnectionState.CONNECTING, + WebSocketConnectionState.CONNECTED -> false - null -> throw IllegalStateException("LibSignalChatConnection.state can never be null") - } + null -> throw IllegalStateException("LibSignalChatConnection.state can never be null") } } @@ -591,48 +591,54 @@ class LibSignalChatConnection( } @OptIn(InternalCoroutinesApi::class) - override suspend fun runWithChatConnection(callback: (ChatConnection) -> T): T = suspendCancellableCoroutine { continuation -> - CHAT_SERVICE_LOCK.withLock { - when (state.value) { - WebSocketConnectionState.CONNECTED -> { - try { - val result = callback(chatConnection!!) - continuation.resume(result) - } catch (e: Exception) { - continuation.resumeWithException(e) - } - } + override suspend fun runWithChatConnection(callback: (ChatConnection) -> T): T { + if (state.value == WebSocketConnectionState.CONNECTED) { + chatConnection?.let { return callback(it) } + } - WebSocketConnectionState.CONNECTING -> { - val action = PendingAction( - onConnectionSuccess = { connection -> - CHAT_SERVICE_LOCK.withLock { - try { - val result = callback(connection) - // NB: We use the experimental tryResume* methods here to avoid crashing if the continuation is - // canceled before we finish the connection attempt, but the PendingAction cannot be removed from - // pendingActions before we get to executing it. - continuation.tryResume(result)?.let(continuation::completeResume) - } catch (e: Throwable) { - continuation.tryResumeWithException(e)?.let(continuation::completeResume) + return suspendCancellableCoroutine { continuation -> + CHAT_SERVICE_LOCK.withLock { + when (state.value) { + WebSocketConnectionState.CONNECTED -> { + try { + val result = callback(chatConnection!!) + continuation.resume(result) + } catch (e: Exception) { + continuation.resumeWithException(e) + } + } + + WebSocketConnectionState.CONNECTING -> { + val action = PendingAction( + onConnectionSuccess = { connection -> + CHAT_SERVICE_LOCK.withLock { + try { + val result = callback(connection) + // NB: We use the experimental tryResume* methods here to avoid crashing if the continuation is + // canceled before we finish the connection attempt, but the PendingAction cannot be removed from + // pendingActions before we get to executing it. + continuation.tryResume(result)?.let(continuation::completeResume) + } catch (e: Throwable) { + continuation.tryResumeWithException(e)?.let(continuation::completeResume) + } } + }, + onFailure = { error -> + continuation.tryResumeWithException(error)?.let(continuation::completeResume) } - }, - onFailure = { error -> - continuation.tryResumeWithException(error)?.let(continuation::completeResume) - } - ) - pendingCallbacks.add(action) + ) + pendingCallbacks.add(action) - continuation.invokeOnCancellation { - CHAT_SERVICE_LOCK.withLock { - pendingCallbacks.removeIf { it === action } + continuation.invokeOnCancellation { + CHAT_SERVICE_LOCK.withLock { + pendingCallbacks.removeIf { it === action } + } } } - } - else -> { - continuation.resumeWithException(IOException("WebSocket is not connected (state: ${state.value})")) + else -> { + continuation.resumeWithException(IOException("WebSocket is not connected (state: ${state.value})")) + } } } }