dataconnect(change): add core/generated SDK info to request headers of realtime query subscriptions - #8356
Conversation
TODO: add unit tests
…er is sent with subscribe request`
…er is sent with resume request`
… when sending resume or subscribe messages
…eElement context element
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. |
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request updates realtime query subscriptions to include SDK type metadata in request headers, matching standard query executions. It introduces CallerSdkTypeElement as a coroutine context element and refactors ConflatedSignal to be generic (ConflatedSignal<T>) to carry payload values. Feedback on these changes suggests falling back to CallerSdkType.Base instead of throwing an exception if the context element is missing to prevent potential crashes, and optimizing the CAS loop in ConflatedSignal.await() to avoid unnecessary suspension on concurrent updates.
… a deadlock! #8356 (comment) ```kotlin suspend fun await(): T { while (true) { val current = signalState.get() if (current != null) { if (signalState.compareAndSet(current, null)) { return current } } channel.receive() } } ``` ```kotlin suspend fun await(): T { while (true) { val current = signalState.get() if (current == null) { channel.receive() } else if (signalState.compareAndSet(current, null)) { return current } } } ``` --- In the original implementation, if `compareAndSet` failed, the calling thread would unconditionally fall through the `if` block and call the suspending `channel.receive()` method. * **The Old Path:** A failed CAS meant that either another consumer consumed the value, or a producer overwrote it with a new conflated value. Even if a new value was immediately available in `signalState`, the thread was forced to execute `channel.receive()`. While `channel.receive()` might return immediately if a token was buffered, calling a suspending function still incurs the overhead of coroutine state-machine suspension/dispatch machinery. * **The New Path:** By checking `if (current == null)` first and only calling `channel.receive()` when no value is present, the new implementation allows the thread to immediately retry the CAS on any newly set value without ever reaching a suspension point. If the CAS failed because a concurrent producer published a new value, the loop immediately retries and consumes the new value. The most critical impact of this change is the resolution of a highly subtle race condition that could lead to a permanently lost signal and deadlocked waiters when coroutine cancellation is introduced. Consider the following interleaving in the **old** implementation: 1. **Initial State:** `signalState` is `null`. A consumer thread, **Thread B**, calls `await()`, sees `null`, and suspends inside `channel.receive()`. 2. **Signal 1:** A producer calls `signal(value1)`. * `signalState` is set to `value1`. * `channel.trySend(Unit)` is called. Since **Thread B** is suspended, the channel resumes **Thread B**. **Thread B** is placed in the dispatcher's queue to resume. 3. **Concurrent Execution:** Before **Thread B** actually executes its next loop iteration, another consumer, **Thread A**, calls `await()`. * **Thread A** reads `current = value1`. 4. **Signal 2 (Conflation):** Before **Thread A** can execute its CAS, the producer calls `signal(value2)`. * `signalState` is updated to `value2`. * `channel.trySend(Unit)` is called. Since **Thread B** is already active/resuming and **Thread A** is not suspended, the `Unit` token is buffered in the conflated channel. 5. **CAS Failure on Thread A:** **Thread A** now executes `compareAndSet(value1, null)`. This fails because the state is now `value2`. 6. **Suspension on Thread A:** In the old code, **Thread A** falls through and calls `channel.receive()`. * Since the token from *Signal 2* was buffered, **Thread A** consumes it, does not suspend, loops, reads `value2`, and successfully CASes it to `null`, returning `value2`. 7. **Alternative Interleaving (The Bug):** What if **Thread B** was resumed, but during its dispatch or right after waking up, its parent job was cancelled? * In Kotlin Coroutines, a cancelled coroutine throws a `CancellationException` upon resuming from a suspension point (like `channel.receive()`). * Therefore, **Thread B** terminates abruptly with `CancellationException` and never executes the next lines of the loop. * If this happens, the token that resumed **Thread B** is gone. * If **Thread A** then fails its CAS because a new value was set, and calls `channel.receive()`, the channel is empty. **Thread A suspends.** * The new value `value2` remains in `signalState` indefinitely, but the only active waiter (**Thread A**) is now suspended waiting for a channel event that will never come. The signal is **lost**, and **Thread A** is **deadlocked**. --- In the new implementation, the `channel.receive()` call is completely bypassed unless `current == null`: 1. When **Thread A**'s CAS fails (step 5 above), it does **not** fall through to call `channel.receive()`. 2. Instead, it immediately loops back to the beginning of the `while (true)` block. 3. In the next iteration, **Thread A** reads `current = signalState.get()`, which is `value2`. 4. Since `current` is not null, it enters the `else if` branch and performs `compareAndSet(value2, null)`. 5. This CAS succeeds, and **Thread A** returns `value2` successfully. Even if **Thread B** is cancelled during dispatch and the resumption token is lost, **Thread A** makes progress and consumes the value safely. The liveness of the signal is completely decoupled from the cancellation lifecycle of competing threads. --- The loop is guaranteed to be lock-free: * A thread only retries the loop without suspending if `compareAndSet` fails. * A failed CAS mathematically implies that another concurrent operation made progress (either another thread successfully consumed the value, or a producer successfully wrote a new value). * Thus, the system as a whole is guaranteed to make progress, preventing any risk of infinite hot loops or livelocks.
📝 PRs merging into main branchOur main branch should always be in a releasable state. If you are working on a larger change, or if you don't want this change to see the light of the day just yet, consider using a feature branch first, and only merge into the main branch when the code complete and ready to be released. |
…f realtime query subscriptions (#8356)
This PR adds core and generated SDK information to the data connect request headers of realtime query subscriptions. This matches the behavior of standard query executions by sending the
x-goog-api-clientheader withcallerSdkTypemetadata when initiating subscribe or resume requests.Highlights- Realtime Query Metadata: Wires up the propagation of SDK type info so that both subscribe and resume requests on the bidirectional connect stream send the
- Generic Conflated Signal: Refactors
- Coroutine Context SDK Propagation: Introduces
- CHANGELOG.md
- Added a changelog entry documenting that realtime query subscriptions now include SDK type metadata.
- CallerSdkTypeElement.kt
- Created
- DataConnectBidiConnectStream.kt
- Updated the
- Included
- Required the collecting coroutine context to contain
- DataConnectGrpcMetadata.kt
- Exposed
- Added public constants
- DataConnectGrpcRPCs.kt
- Exposed
- Passed
- QuerySubscriptionImpl.kt
- Collected the realtime subscription flow within a coroutine context containing
- RealtimeQueryManager.kt
- Passed the
- GrpcBidiFlow.kt
- Specified the generic type
- ConflatedSignal.kt
- Converted
- Added
- Added a
- DataConnectGrpcRPCsConnectIntegrationTest.kt
- Updated integration tests to pass and assert
- DataConnectGrpcRPCsUnitTest.kt
- Updated unit tests to specify and assert
- QuerySubscriptionImplUnitTest.kt
- Added unit tests verifying the
- ConflatedSignalUnitTest.kt
- Added tests covering the new generic payload support, latest value conflation, and
x-goog-api-clientheader.ConflatedSignalinto a generic classConflatedSignal<T : Any>that buffers and delivers non-nullable payloads (such as the caller SDK type) to competing waiters.CallerSdkTypeElementto carry caller SDK type metadata in the coroutine context, allowing individual flow collectors to specify their caller SDK type.Changelog
CallerSdkTypeElementto carryCallerSdkTypein the coroutine context.subscribemethod and internal connection states to propagateCallerSdkType.x-goog-api-clientheader in subscribe and resume stream requests usinggrpcMetadata.CallerSdkTypeElement.googApiClientHeaderValue()publicly.FIREBASE_AUTH_TOKEN_HEADERandGOOG_API_CLIENT_HEADER.grpcMetadatafor testing.grpcMetadatatoDataConnectBidiConnectStream.CallerSdkTypeElement.callerSdkTypeparameter to the subscription stream connection.Unitwhen instantiatingConflatedSignal.ConflatedSignalto a generic classConflatedSignal<T : Any>that buffers and delivers the latest signal value to a single waiter.signal(value: T)andawait(): Tto handle generic payloads, with values conflating to the most recent signal value.signal()extension method forConflatedSignal<Unit>to support parameterless signaling.CallerSdkTypeusingCallerSdkTypeElement.CallerSdkTypeon stream subscriptions.x-goog-api-clientheader is sent with both subscribe and resume requests.pendingSignalproperty inConflatedSignal.