Resolve race conditions in HealthMonitor and SocketSession

[aleksandar-apostolov]

Goal

Fix two race conditions identified in the core socket infrastructure that could cause incorrect behavior under concurrent access.

Implementation

1. HealthMonitor lastAck atomicity

var lastAck: Long → AtomicLong. The field is written from OkHttp's callback thread (acknowledgeHeartbeat) and read from the monitor coroutine. Without volatile or atomic ops, the reader could see stale values and falsely trigger the liveness threshold, causing premature socket disconnection.

2. SocketSession dual-listener race

During WebSocket handshake, both eventListener and connectListener were subscribed simultaneously — the auth response (connection.ok) was processed by both listeners. Now only the handshake listener is subscribed initially. It handles auth responses and buffers any other messages that arrive during handshake. On success, eventListener subscribes and buffered messages are replayed through the batcher. Zero message loss, no duplicate processing.

Testing

Existing unit tests updated to reflect the new single-listener handshake flow. Full suite passes (all SocketSession + HealthMonitor tests).

Summary by CodeRabbit

• Bug Fixes
  • Improved socket connection reliability through message buffering during the handshake phase
  • Enhanced thread-safe heartbeat monitoring to prevent timeout-related issues
  • Optimized socket subscription lifecycle management to prevent message loss during authentication transitions
  • Strengthened error handling and recovery for socket connection failures

[github-actions]

PR checklist ✅

All required conditions are satisfied:

• Title length is OK (or ignored by label).
• At least one pr: label exists.
• Sections ### Goal, ### Implementation, and ### Testing are filled.

🎉 Great job! This PR is ready for review.

[coderabbitai]

Walkthrough

The changes refactor the socket subscription lifecycle to buffer non-authentication messages during handshake, defer eventListener subscription until after successful authentication, and replay buffered messages to the batcher. Additionally, health monitoring adopts thread-safe atomic operations for timestamp tracking.

Changes

Cohort / File(s)	Summary
Socket Subscription Lifecycle stream-android-core/src/main/java/io/getstream/android/core/internal/socket/StreamSocketSession.kt	Introduced pendingMessages buffering during handshake; removed initial eventListener subscription before handshake and replaced with deferred subscription only after authentication; added logic to replay buffered messages to batcher; updated handshake message handling to buffer non-auth messages; added health monitor acknowledgment in handshake onMessage; refactored deserialization flow from map-based to onSuccess block style.
Health Monitor Thread Safety stream-android-core/src/main/java/io/getstream/android/core/internal/socket/monitor/StreamHealthMonitorImpl.kt	Replaced mutable lastAck: Long with thread-safe AtomicLong; updated acknowledgeHeartbeat() to use set() for atomic writes; updated liveness-check condition to read via get() before computing elapsed time.
Subscription Lifecycle Tests stream-android-core/src/test/java/io/getstream/android/core/internal/socket/StreamSocketSessionTest.kt	Updated connection-failure tests to reflect new lifecycle subscription behavior (expecting 0 non-handshake subscriptions when handshake fails); renamed test cases to align with new expected state emissions; adjusted cancellation verification expectations; reordered stubbing for consistency across lifecycle and handshake subscriptions.

Sequence Diagram

sequenceDiagram
    participant Client
    participant SocketSession as StreamSocketSession
    participant Socket as WebSocket
    participant Batcher
    participant HealthMonitor

    Client->>SocketSession: connect()
    activate SocketSession
    SocketSession->>SocketSession: pendingMessages = []
    SocketSession->>Socket: subscribe(handshakeListener)
    activate Socket
    Socket-->>SocketSession: onOpen()
    deactivate Socket
    
    Socket-->>SocketSession: onMessage(authResponse)
    activate SocketSession
    SocketSession->>HealthMonitor: acknowledgeHeartbeat()
    SocketSession->>SocketSession: deserialize & validate auth
    alt Authentication Success
        SocketSession->>Socket: subscribe(eventListener)
        SocketSession->>Batcher: offer(bufferedMessage) [replay]
        SocketSession->>SocketSession: pendingMessages.clear()
    else Non-Auth Message During Handshake
        SocketSession->>SocketSession: pendingMessages.append(message)
    end
    deactivate SocketSession
    deactivate SocketSession

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Poem

🐰 A whisker-twitch for buffering wise,
Messages queued 'til handshake flies,
Atomic hearts that beat so true,
Subscriptions deferred for steady brew! ✨

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 11.76% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title 'fix(socket): resolve race conditions in HealthMonitor and SocketSession' accurately and specifically describes the main changes across all modified files—fixing race conditions in socket infrastructure.
Description check	✅ Passed	The description covers all required template sections: Goal (fixes two race conditions), Implementation (detailed explanation of both fixes), Testing (tests updated and suite passes), and addresses the Checklist items.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fix/race-conditions

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 2

🧹 Nitpick comments (1)

stream-android-core/src/test/java/io/getstream/android/core/internal/socket/StreamSocketSessionTest.kt (1)

564-594: Please add a direct regression test for the buffered-handshake path.

These additions cover the failure-side subscription counts, but the new behavior in StreamSocketSession.connect is buffering non-auth frames until connection.ok and replaying them afterward. Add one test that delivers a normal event before CONNECTED_JSON and assert it is forwarded exactly once after the steady-state listener is installed.

Based on learnings: "Maintain idempotency and guard callback ordering in connection state, batch processing, and retry flows that feed StreamSubscriptionManager."

Also applies to: 1173-1223

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In
`@stream-android-core/src/test/java/io/getstream/android/core/internal/socket/StreamSocketSessionTest.kt`
around lines 564 - 594, Add a regression test in StreamSocketSessionTest that
verifies the buffered-handshake path: call session.connect(...) as in existing
tests but simulate the handshake listener receiving a non-auth event (a normal
event frame) before emitting CONNECTED_JSON, then emit CONNECTED_JSON and assert
that the normal event is forwarded exactly once to the steady-state subscription
handling (i.e., after the handshake listener is replaced by the steady
listener). Use the same mocking style as the file: stub socket.subscribe(...)
and socket.open(...) to succeed, push a normal event into the initial
StreamWebSocketListener callback before sending the CONNECTED_JSON frame, then
send CONNECTED_JSON and verify the subscription forward path invoked once (use
the same subs.forEach(...) / subscription manager verifications or the steady
listener verification used elsewhere) to ensure the buffered event is replayed
exactly one time.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In
`@stream-android-core/src/main/java/io/getstream/android/core/internal/socket/monitor/StreamHealthMonitorImpl.kt`:
- Around line 52-53: The monitorJob field is racy: ensure start/stop are atomic
by replacing the mutable monitorJob with an AtomicReference<Job?> or by guarding
the check/assign/cancel sequence with a mutex/@Synchronized; update
StreamHealthMonitorImpl.start() to use compare-and-set (CAS) on the
AtomicReference when storing the newly launched Job and change stop() to
atomically getAndSet(null) and cancel the returned Job, referencing monitorJob,
start(), and stop() in your changes (lastAck can remain as-is).

In
`@stream-android-core/src/main/java/io/getstream/android/core/internal/socket/StreamSocketSession.kt`:
- Around line 286-303: The success lambda handling post-handshake must treat
subscription or replay failures as fatal: if
internalSocket.subscribe(eventListener) fails (socketSubRes.isFailure) or any
batcher.offer(message) returns false, cancel handshakeSubscription, avoid
starting healthMonitor or calling connect() success, tear down the session
(stop/close the socket, clear socketSubscription, stop healthMonitor) and
propagate/return a failed connect result instead; also sanitize the replay
failure log so it does not print raw message payloads (use a redacted/summary
message or message id). Ensure you update logic around success,
socketSubscription, pendingMessages, batcher.offer, handshakeSubscription and
healthMonitor to implement this behavior.

---

Nitpick comments:
In
`@stream-android-core/src/test/java/io/getstream/android/core/internal/socket/StreamSocketSessionTest.kt`:
- Around line 564-594: Add a regression test in StreamSocketSessionTest that
verifies the buffered-handshake path: call session.connect(...) as in existing
tests but simulate the handshake listener receiving a non-auth event (a normal
event frame) before emitting CONNECTED_JSON, then emit CONNECTED_JSON and assert
that the normal event is forwarded exactly once to the steady-state subscription
handling (i.e., after the handshake listener is replaced by the steady
listener). Use the same mocking style as the file: stub socket.subscribe(...)
and socket.open(...) to succeed, push a normal event into the initial
StreamWebSocketListener callback before sending the CONNECTED_JSON frame, then
send CONNECTED_JSON and verify the subscription forward path invoked once (use
the same subs.forEach(...) / subscription manager verifications or the steady
listener verification used elsewhere) to ensure the buffered event is replayed
exactly one time.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: 682ba919-9bf8-4ed8-80f9-accfb74c0774

📥 Commits

Reviewing files that changed from the base of the PR and between 1a0a3cd and 4828e21.

📒 Files selected for processing (3)

• stream-android-core/src/main/java/io/getstream/android/core/internal/socket/StreamSocketSession.kt
• stream-android-core/src/main/java/io/getstream/android/core/internal/socket/monitor/StreamHealthMonitorImpl.kt
• stream-android-core/src/test/java/io/getstream/android/core/internal/socket/StreamSocketSessionTest.kt

[sonarqubecloud]

Quality Gate passed

Issues
0 New issues
0 Accepted issues

Measures
0 Security Hotspots
85.2% Coverage on New Code
0.0% Duplication on New Code

See analysis details on SonarQube Cloud

[stream-public-bot]

🚀 Available in v3.0.1