Improve Query Subscription Management and Resource Cleanup#591
Conversation
WalkthroughThe changes enhance subscription lifecycle management across several modules. In the Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant MQB as ModelQueryBuilder
participant GraphQL as GraphQL API
participant Backend as Backend Service
Client->>MQB: Call subscribe/countSubscribe/paginateSubscribe
Note right of MQB: Call dispose() first to clear existing subscription
MQB->>MQB: Stop keepalive signals\nUnsubscribe from updates\nNotify backend for cleanup
MQB->>GraphQL: Unsubscribe from updates
GraphQL-->>MQB: Acknowledge unsubscription
MQB->>Backend: Notify resource cleanup
Backend-->>MQB: Confirm cleanup
MQB->>MQB: Create new subscription
MQB-->>Client: Return new subscription
sequenceDiagram
participant PC as PerspectiveClient
participant API as GraphQL API
participant PR as PerspectiveResolver
participant PI as PerspectiveInstance
PC->>API: disposeQuerySubscription(uuid, subscriptionId)
API->>PR: Invoke perspectiveDisposeQuerySubscription(uuid, subscriptionId)
PR->>PI: Call dispose_query_subscription(subscriptionId)
PI-->>PR: Return success (bool)
PR-->>API: Return result (bool)
API-->>PC: Return result (bool)
Possibly related PRs
Suggested reviewers
Poem
✨ Finishing Touches
🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
tests/js/tests/prolog-and-literals.test.ts (2)
2324-2325: Remove debug console.log statements.These console.log statements should be removed before merging to production.
- console.log("countCallback", countCallback.lastCall.args[0]) - let count = countCallback.callCount + let count = countCallback.callCount
2371-2372: Remove debug console.log statements.These console.log statements should be removed before merging to production.
- console.log("countCallback", pageCallback.lastCall.args[0]) - let count = pageCallback.callCount + let count = pageCallback.callCount
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (10)
core/src/model/Ad4mModel.ts(5 hunks)core/src/perspectives/PerspectiveClient.ts(1 hunks)core/src/perspectives/PerspectiveProxy.ts(6 hunks)core/src/perspectives/PerspectiveResolver.ts(1 hunks)docs/pages/developer-guides/model-classes.mdx(1 hunks)docs/pages/perspectives.mdx(1 hunks)rust-executor/src/graphql/mutation_resolvers.rs(1 hunks)rust-executor/src/perspectives/perspective_instance.rs(5 hunks)tests/js/tests/perspective.ts(1 hunks)tests/js/tests/prolog-and-literals.test.ts(2 hunks)
🧰 Additional context used
🧬 Code Definitions (3)
core/src/model/Ad4mModel.ts (1)
core/src/perspectives/PerspectiveProxy.ts (3)
PerspectiveProxy(273-1170)result(156-158)result(185-193)
core/src/perspectives/PerspectiveClient.ts (1)
core/src/perspectives/PerspectiveResolver.ts (1)
perspectiveDisposeQuerySubscription(358-363)
tests/js/tests/prolog-and-literals.test.ts (3)
core/src/model/Ad4mModel.ts (2)
perspective(336-338)Ad4mModel(273-854)core/src/perspectives/PerspectiveProxy.ts (1)
PerspectiveProxy(273-1170)core/src/model/decorators.ts (2)
ModelOptions(544-719)Property(776-783)
🔇 Additional comments (44)
core/src/perspectives/PerspectiveClient.ts (1)
217-226: Implementation for query subscription disposal looks good.This method completes the subscription lifecycle management that was previously missing - we already had methods to create subscriptions, keep them alive, and update them, but lacked proper cleanup. This implementation follows the same pattern as other GraphQL mutation methods in the class.
core/src/perspectives/PerspectiveResolver.ts (1)
357-363:Details
❓ Verification inconclusive
Verify implementation completeness in the resolver.
The implementation simply returns
truewithout performing any actual subscription disposal logic. While this might be acceptable if the actual implementation is delegated to the Rust backend entirely, consider adding a comment to clarify this design decision to help future maintainers understand the intent.
🏁 Script executed:
#!/bin/bash # Check if there are any references to a separate implementation of the subscription disposal logic rg -A 5 -B 5 "dispose_query_subscription" --type ts --type rsLength of output: 88
Based on our previous investigation, it appears that the original grep search did not properly target Rust files (due to the use of "--type rs"). Please run the above script to verify if there are any Rust-source references for the subscription disposal logic. Once confirmed:
Final Review Comment:
Clarify Delegated Subscription Disposal Logic
Incore/src/perspectives/PerspectiveResolver.ts(lines 357–363), the mutationperspectiveDisposeQuerySubscriptioncurrently returnstruewithout executing any disposal logic. If the intention is to delegate the actual disposal to the Rust backend, please add a comment clarifying this decision to aid future maintainers. If not, consider implementing the required logic within this resolver.
🏁 Script executed:
#!/bin/bash # Re-run the search command for Rust files using a glob pattern to target .rs files correctly. rg -g "*.rs" -A 5 -B 5 "dispose_query_subscription"Length of output: 2895
Delegated Subscription Disposal Logic Clarification Needed
Incore/src/perspectives/PerspectiveResolver.ts(lines 357–363), the mutation methodperspectiveDisposeQuerySubscriptioncurrently returnstruewithout handling any disposal logic. Our verification indicates that the actual disposal is implemented in the Rust backend (e.g., seedispose_query_subscriptioninrust-executor/src/perspectives/perspective_instance.rs). To avoid confusion for future maintainers, please add an inline comment in the resolver clarifying that the disposal logic is intentionally delegated to the Rust service.rust-executor/src/graphql/mutation_resolvers.rs (1)
1035-1048:Details
✅ Verification successful
Proper implementation of query subscription disposal.
This implementation correctly:
- Verifies user capabilities before performing the operation
- Retrieves the perspective instance using the UUID
- Delegates subscription disposal to the perspective instance
This matches the pattern used in related methods like
perspective_keep_alive_queryand maintains the same security checks.
🏁 Script executed:
#!/bin/bash # Verify that the dispose_query_subscription method exists in the PerspectiveInstance struct rg -A 5 "fn dispose_query_subscription" --type rustLength of output: 660
Verification Passed: Query Subscription Disposal Implementation is Correct
The implementation in
rust-executor/src/graphql/mutation_resolvers.rscorrectly:
- Verifies user capabilities using
check_capability.- Retrieves the perspective via
get_perspective_with_uuid_field_error.- Delegates the disposal operation to the perspective by calling its
dispose_query_subscriptionmethod (verified inrust-executor/src/perspectives/perspective_instance.rs).No changes are required as it follows the established pattern used by similar methods like
perspective_keep_alive_query.core/src/perspectives/PerspectiveProxy.ts (5)
32-36: Well-documented resource cleanup!The updated documentation clearly explains that resources need to be cleaned up and the backend needs to be notified when a subscription is no longer needed.
53-54: Clear guidance for subscription cleanupThe updated example comment properly emphasizes the need to clean up subscriptions when they're no longer needed, which is crucial for preventing resource leaks.
123-133: Good addition of ID accessor for subscription managementAdding this ID accessor method is important for subscription management, especially for:
- Enabling subscription reference tracking
- Supporting deduplication of identical subscriptions
- Facilitating debugging
The method is well-documented with clear explanations of its purpose and usage.
217-222: Proper backend notification on subscription disposalThis is a critical addition that ensures subscription resources are properly cleaned up on the backend when they're no longer needed on the frontend. Without this, abandoned subscriptions could lead to memory leaks and degraded performance over time.
1127-1152: Enhanced documentation on subscription lifecycle managementThe updated documentation for
subscribeInferclearly communicates that:
- Subscriptions need to be cleaned up on both frontend and backend
- Failing to call
dispose()will cause resource leaks- The proper pattern includes disposing subscriptions when done
This is essential information for developers and the example demonstrates the correct usage pattern.
tests/js/tests/perspective.ts (1)
312-351: Good test for subscription deduplicationThis test effectively validates a key improvement in the subscription system: reusing subscriptions for identical queries. The test:
- Creates two subscriptions with the same query
- Verifies they have the same ID (deduplication)
- Confirms both callbacks receive updates when relevant data changes
- Verifies the consistency of results between both subscriptions
This is important for ensuring efficient resource utilization while maintaining correct behavior.
docs/pages/perspectives.mdx (4)
117-117: Good documentation of new capabilityAdding this bullet point clearly communicates the new real-time query subscription capability with automatic cleanup to users.
120-126: Improved example labelingRenaming this from "Query using Prolog" to "Basic Prolog query" helps differentiate it from the new subscription-based query examples.
127-140: Clear example of subscription usageThis example effectively demonstrates:
- How to create a real-time subscription
- Setting up a callback for updates
- Proper cleanup with dispose()
The code is concise and follows best practices for subscription management.
142-154: Excellent documentation on subscription lifecycleThis section provides critical information on:
- The subscription lifecycle (initialization, keepalive, updates, cleanup)
- The importance of calling
dispose()to prevent resource leaks- What the cleanup process includes (stopping signals, unsubscribing, backend notification)
The explicit reminder to always call
dispose()at the end emphasizes the importance of this step.core/src/model/Ad4mModel.ts (18)
881-881: Added property to track active subscription.This property is essential for maintaining a reference to the current subscription, allowing it to be properly disposed later.
889-906: Added critical resource cleanup method.The new
dispose()method properly handles subscription cleanup by:
- Calling the subscription's dispose method
- Clearing the reference to allow garbage collection
This important addition prevents memory leaks and ensures proper backend resource cleanup.
1057-1059: Updated documentation to highlight resource cleanup.The documentation now properly emphasizes the importance of calling
dispose()when done with subscriptions, improving resource management awareness.
1073-1074: Updated example code to demonstrate proper disposal.The example in the documentation has been updated to show the correct pattern for disposing subscriptions.
1077-1079: Added subscription cleanup before creating a new one.Calling
dispose()at the beginning ofsubscribe()ensures any existing subscription is properly cleaned up before creating a new one, preventing resource leaks.
1081-1082: Updated to store subscription in class property.The subscription is now stored in the
currentSubscriptionproperty rather than a local variable, allowing it to be disposed later.
1088-1088: Updated to use class subscription property.The callback registration now uses the stored
currentSubscriptionproperty.
1092-1093: Updated to use class subscription property for result access.The result is now accessed from the stored
currentSubscriptionproperty.
1123-1124: Updated documentation to highlight resource cleanup.Similar to the previous documentation update, this properly emphasizes the importance of calling
dispose().
1139-1140: Updated example code to demonstrate proper disposal.The example has been updated to show the correct disposal pattern.
1143-1145: Added subscription cleanup before creating a new one.Calling
dispose()at the beginning ofcountSubscribe()ensures any existing subscription is properly cleaned up.
1147-1148: Updated to store subscription in class property.The count subscription is now stored in the
currentSubscriptionproperty.
1154-1155: Updated to use class subscription property.The callback registration and result access now use the stored
currentSubscriptionproperty.
1189-1190: Updated documentation to highlight resource cleanup.This documentation update properly emphasizes the importance of calling
dispose()for paginated subscriptions.
1207-1208: Updated example code to demonstrate proper disposal.The example has been updated to show the correct disposal pattern for paginated subscriptions.
1215-1217: Added subscription cleanup before creating a new one.Calling
dispose()at the beginning ofpaginateSubscribe()ensures any existing subscription is properly cleaned up.
1220-1221: Updated to store subscription in class property.The paginated subscription is now stored in the
currentSubscriptionproperty.
1227-1228: Updated to use class subscription property.The callback registration and result access now use the stored
currentSubscriptionproperty.docs/pages/developer-guides/model-classes.mdx (3)
241-252: Improved subscription example with explicit lifecycle management.The example now properly creates a query builder, subscribes to updates, and explicitly disposes of the subscription when done. This aligns with the implementation changes in the core code and demonstrates the correct pattern for managing subscription resources.
254-265: Added clear documentation about the subscription system.This new section clearly explains how the subscription system works and emphasizes the importance of properly disposing subscriptions. It highlights the key benefits of resource cleanup including preventing memory leaks and stopping background processes.
266-278: Added example for count-only subscriptions.This new example showcases the
countSubscribemethod, demonstrating how to subscribe to just the count of matching items instead of the full result set. This is valuable for UI scenarios where only the count is needed, reducing data transfer and processing overhead.tests/js/tests/prolog-and-literals.test.ts (5)
20-20: Added import for sinon test library.The sinon library is properly imported to create fake callbacks for testing subscription behaviors.
2214-2236: Good test setup for ModelQueryBuilder tests.The test suite is well-structured with:
- A clear describe block for ModelQueryBuilder tests
- A properly defined TestModel class with the necessary decorators
- A beforeEach hook that sets up a clean perspective and registers the model for each test
This ensures isolated and repeatable tests for the subscription functionality.
2238-2302: Comprehensive test for subscription handling and disposal.This test thoroughly validates that:
- Initial subscriptions return empty results
- Callbacks are triggered when matching models are added
- Creating a new subscription automatically disposes the previous one
- Explicitly disposed subscriptions stop receiving updates
This covers the core subscription lifecycle management that's central to this PR.
2304-2341: Well-implemented test for count subscriptions.The test properly verifies count subscription behavior:
- Initial count is zero
- Count updates when models are added
- Count updates stop after subscription disposal
This ensures the count-specific subscription features work correctly.
2343-2388: Comprehensive test for paginated subscriptions.The test properly validates paginated subscription behavior:
- Initial page has empty results and zero total count
- Page data updates when models are added
- Updates stop after subscription disposal
This ensures pagination features work correctly with the subscription lifecycle.
rust-executor/src/perspectives/perspective_instance.rs (5)
2029-2048: Well-implemented subscription update abstraction.This new helper method centralizes the subscription update logic, which improves code maintainability and makes the codebase more DRY. The approach of spawning an async task ensures non-blocking behavior, and the optional delay parameter adds flexibility.
2051-2071: Good implementation of query subscription deduplication.The changes now check for existing subscriptions with the same query, effectively implementing deduplication as intended. This will prevent resource leaks by reusing existing subscriptions instead of creating duplicates.
2108-2111: Simple but effective subscription disposal implementation.This new method provides a clean way to dispose of query subscriptions, which completes the subscription lifecycle management. The implementation is appropriately concise for the task.
2149-2154: Code refactoring to use centralized update method.The use of the new
send_subscription_updatemethod here simplifies the code and ensures consistent behavior across the codebase.
2120-2130: Enhanced logging for improved debugging.The additional logging statements provide better visibility into the subscription checking process, which will be helpful for debugging subscription-related issues.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
rust-executor/src/perspectives/perspective_instance.rs (2)
2057-2079: Good implementation of subscription deduplication.The changes effectively implement subscription reuse by checking for existing subscriptions with the same query. This prevents redundant subscriptions and avoids duplicating processing resources.
One suggestion: consider adding a reference counter to track how many clients are using each subscription, which would allow for better cleanup when all clients disconnect.
struct SubscribedQuery { query: String, last_result: String, last_keepalive: Instant, + reference_count: usize, } // Then in subscribe_and_query: if let Some(existing_id) = existing_subscription { let mut queries = self.subscribed_queries.lock().await; if let Some(query) = queries.get_mut(&existing_id) { + query.reference_count += 1; for delay in [100, 500, 1000] { // ...
2117-2123: Good addition of subscription disposal mechanism.This method is key to properly cleaning up resources when subscriptions are no longer needed. It's simple but effective, removing the subscription from the map when requested.
Building on the reference counter suggestion, consider enhancing this to only fully remove when the count reaches zero:
pub async fn dispose_query_subscription( &self, subscription_id: String, ) -> Result<bool, AnyError> { let mut queries = self.subscribed_queries.lock().await; - Ok(queries.remove(&subscription_id).is_some()) + if let Some(query) = queries.get_mut(&subscription_id) { + query.reference_count -= 1; + if query.reference_count == 0 { + queries.remove(&subscription_id); + return Ok(true); + } + return Ok(false); + } + Ok(false) }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
core/src/perspectives/PerspectiveProxy.ts(8 hunks)rust-executor/src/perspectives/perspective_instance.rs(5 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- core/src/perspectives/PerspectiveProxy.ts
🔇 Additional comments (3)
rust-executor/src/perspectives/perspective_instance.rs (3)
2029-2053: Excellent encapsulation of subscription update logic.This new method centralizes the subscription update logic that was previously duplicated across the codebase. The encapsulation is clean, with proper async handling and smart use of tokio's spawn for non-blocking operations.
2097-2102: Well-refactored subscription update.The initial result is now sent using the centralized
send_subscription_updatemethod, which is more consistent and maintainable than the previous inline implementation.
2132-2133: Improved subscription check and update mechanism.The subscription check has been streamlined with better logging that includes the subscription ID. Using the new
send_subscription_updatemethod for publishing changes is more consistent and better encapsulates the functionality.Also applies to: 2146-2147, 2161-2163
…hey don’t interfere with expected test results
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (3)
tests/js/tests/perspective.ts (3)
353-353: Remove console.log statement.Console logging should be removed from test code as it clutters test output.
- console.log("result1", result1)
312-356: Consider adding a test for subscription disposal.The PR objectives mention implementing subscription disposal, but this test suite doesn't explicitly verify that subscriptions are properly disposed. Consider adding a test that creates a subscription, disposes it, and verifies the cleanup occurred correctly both on the client and server.
334-337: Replace timing-dependent tests with deterministic approach.Using
sleep(1000)makes tests potentially flaky. Consider implementing a more deterministic approach like waiting for a specific event or using a polling mechanism with timeout.- // Wait for the subscriptions to be established - // it's sending the initial result a couple of times - // to allow clients to wait and ensure for the subscription to be established - await sleep(1000) + // Wait for the subscriptions to be established + // Poll until both callbacks have been called or timeout + const timeout = Date.now() + 5000; // 5 second timeout + while (!callback1.called || !callback2.called) { + if (Date.now() > timeout) { + throw new Error('Timed out waiting for subscription callbacks'); + } + await sleep(100); + }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
tests/js/tests/perspective.ts(2 hunks)
🧰 Additional context used
🧬 Code Definitions (1)
tests/js/tests/perspective.ts (2)
core/src/Ad4mClient.ts (1)
Ad4mClient(19-72)core/src/model/Ad4mModel.ts (1)
query(847-853)
🔇 Additional comments (2)
tests/js/tests/perspective.ts (2)
312-356: Well-implemented test for subscription deduplication!This new test correctly validates that identical Prolog queries share the same subscription, which aligns with the PR objective of implementing query subscription deduplication. The test confirms both shared subscription IDs and identical callback results.
386-386:Details
❓ Verification inconclusive
Verify if skipping Batch Operations tests is intentional.
You've changed the Batch Operations test suite to be skipped. Is this intentional or temporary? These tests may be important for verifying that batch operations still work correctly with the new subscription management.
🏁 Script executed:
#!/bin/bash # Check if there are any issues or TODOs related to batch operations rg -i "batch.*TODO|TODO.*batch|FIXME.*batch" --type tsLength of output: 56
Please run this script to see if there is any inline comment or indication of why the test suite was skipped. This will help determine if the skipping is intentional or temporary.
🏁 Script executed:
#!/bin/bash # Display surrounding context for "describe.skip('Batch Operations'" in tests/js/tests/perspective.ts sed -n '380,390p' tests/js/tests/perspective.tsLength of output: 449
Action Needed: Confirm Intentional Skipping of Batch Operations Tests
The
describe.skip('Batch Operations', ...)intests/js/tests/perspective.tscurrently disables these tests without any inline explanation or associated TODO/FIXME. Please confirm whether this is an intentional, long-term decision in light of the new subscription management changes or just a temporary workaround. If temporary, consider adding a comment or creating an issue to track re-enabling these tests.
Improve Query Subscription Management and Resource Cleanup
Overview
This PR implements proper subscription disposal and deduplication for Prolog query subscriptions, preventing resource leaks and unnecessary duplicate subscriptions. The changes ensure that subscriptions are properly cleaned up when no longer needed and that identical queries reuse existing subscriptions.
Key Changes
1. Subscription Disposal Implementation
perspectiveDisposeQuerySubscriptionmutation to GraphQL schemaPerspectiveInstanceto clean up subscription resourcesQuerySubscriptionProxythat:2. Query Subscription Deduplication
PerspectiveInstance3. Resource Management
Summary by CodeRabbit
New Features
disposemethod for improved subscription management, ensuring efficient resource cleanup.Bug Fixes
Documentation
Tests