Fix race condition when loading the kryon subgraph parallelly just after VajramKryonGraph creation #330
Conversation
|
Warning Rate limit exceeded@RamAnvesh has exceeded the limit for the number of commits or files that can be reviewed per hour. Please wait 17 minutes and 12 seconds before requesting another review. ⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. 📒 Files selected for processing (1)
WalkthroughThis pull request primarily focuses on enhancing thread safety and streamlining executor configurations. The changes replace standard map implementations with concurrent variants in several registry classes, refactor method signatures and internal logic to support these concurrency improvements, and update test cases accordingly. Additionally, the PR updates the Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant VajramKryonGraph
participant DependencyManager
Client->>VajramKryonGraph: getKryonId(vajramId)
VajramKryonGraph->>VajramKryonGraph: loadKryonSubgraph(vajramId, loadingInProgress)
VajramKryonGraph->>DependencyManager: Load dependencies for vajramId
DependencyManager-->>VajramKryonGraph: Return kryon definitions
VajramKryonGraph-->>Client: Return kryonId
Poem
Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media? 🪧 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 (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (9)
vajram/vajram-samples/src/test/java/com/flipkart/krystal/vajram/samples/calculator/adder/SplitAdderTest.java (1)
135-202: Consider adding warmup iterations to benchmark.The benchmark configuration is well documented, but consider adding warmup iterations to ensure JVM optimizations are in place before measuring performance.
@Test void vajram_benchmark() throws Exception { + // Warmup iterations + int warmupCount = 1000; + for (int i = 0; i < warmupCount; i++) { + splitAdd(i); + splitAddAsync(i).get(); + } + int loopCount = 50_000; long javaNativeTimeNs = javaMethodBenchmark(this::splitAdd, loopCount);vajram/vajram-krystex/src/main/java/com/flipkart/krystal/vajramexecutor/krystex/VajramKryonGraph.java (2)
231-234: Inline call toloadKryonSubgraph.This direct call forces a new subgraph load when
kryonIdis not found. The concurrency risk is mitigated byloadKryonSubgraph’s synchronization block, which is good. Confirm that re-entrant calls or repeated attempts do not cause performance bottlenecks.
244-279: Initialize and register subgraph logic.This large block creates or retrieves several definitions (including input resolvers and dependencies). Given the complexity, consider splitting the logic into smaller, dedicated methods or at least adding more in-line documentation to clarify the flow. This helps with readability and testing.
vajram/vajram-samples/src/test/java/com/flipkart/krystal/vajram/samples/calculator/FormulaTest.java (6)
65-70: Pre-leasing executors inbeforeAll().Pre-leasing executor instances for tests can reduce overhead. Confirm that your test environment always has enough resources for the maximum lease requests. Or handle
LeaseUnavailableExceptionmore insightfully if test concurrency scales up.
72-75: Release leases inafterAll().Closing leases ensures tests do not hold resources indefinitely. Great practice. Consider adding a final null-check or a try/catch around the close for extra safety in the face of partial failures.
166-233: Parallel test to highlight race conditions in graph loading.This parameterized test is a great approach for stress-testing concurrency. Check if multiple increments of
Adder.CALL_COUNTERmight occur out of order. Ensure that 1 second (succeedsWithin(ofSeconds(1))) is enough margin under heavy load or lower-end machines.
235-299: Benchmark test with up to a million executors.The approach thoroughly tests concurrency, but it may be time-consuming. Confirm that it runs within acceptable time in your CI pipeline (or mark as a heavy test like you did with
@Disabled). Also note that repeated executor creation can artificially inflate overhead times, so be sure it aligns with real-world usage patterns.
312-312:@Disabled("Long running benchmark (~9s)").Clearly labeling problematic or time-consuming tests is a good practice. Consider a separate performance test profile rather than disabling if you want occasional performance checks.
568-568: Acquire executors from leased array.Directly retrieving the executor from the lease array is straightforward. Just confirm that each index of
EXECUTOR_LEASESis not null and properly closed if an error occurs. Another check for out-of-bounds usage might be prudent.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (7)
.gitignore(1 hunks)krystex/src/main/java/com/flipkart/krystal/krystex/LogicDefinitionRegistry.java(1 hunks)krystex/src/main/java/com/flipkart/krystal/krystex/kryon/KryonDefinitionRegistry.java(1 hunks)vajram/vajram-krystex/src/main/java/com/flipkart/krystal/vajramexecutor/krystex/VajramKryonGraph.java(5 hunks)vajram/vajram-krystex/src/test/java/com/flipkart/krystal/vajramexecutor/krystex/KrystexVajramExecutorTest.java(25 hunks)vajram/vajram-samples/src/test/java/com/flipkart/krystal/vajram/samples/calculator/FormulaTest.java(6 hunks)vajram/vajram-samples/src/test/java/com/flipkart/krystal/vajram/samples/calculator/adder/SplitAdderTest.java(2 hunks)
✅ Files skipped from review due to trivial changes (1)
- .gitignore
🔇 Additional comments (16)
vajram/vajram-samples/src/test/java/com/flipkart/krystal/vajram/samples/calculator/adder/SplitAdderTest.java (2)
155-155: LGTM! The single-threaded executor configuration aligns with race condition fix.The addition of
.singleThreadExecutor(executorLease.get())in both benchmark methods ensures consistent thread handling and helps prevent race conditions when loading the kryon subgraph.Also applies to: 227-227
54-59: Verify executor pool configuration and resource cleanup.The class properly manages thread executors through a shared pool with:
- Fixed pool size of 4 threads
- Proper leasing and releasing of executors
- Cleanup in tearDown
However, let's verify if 4 threads are sufficient for the benchmark load:
Also applies to: 61-80
✅ Verification successful
🏁 Script executed:
#!/bin/bash # Check for other benchmark configurations in the codebase rg -A 5 "SingleThreadExecutorsPool.*Test"Length of output: 13360
Executor Pool Configuration Verified
- The grep results show that many tests use a fixed 4-thread configuration via a call like
new SingleThreadExecutorsPool("Test", 4), consistent with the implementation inSplitAdderTest.java.- In a few tests where heavier loads are anticipated (e.g.,
TooManyQualifiersTest,DoubleMinusOneTest,AddZeroTest), a dynamic thread count (usingRuntime.getRuntime().availableProcessors()) is used. This indicates that the 4-thread setup in the benchmark tests is an intentional and context-specific choice.- Resource cleanup is managed via tearDown routines in these tests (as noted in similar test classes), ensuring that executors are properly released.
Based on the findings, the 4-thread configuration in
SplitAdderTestis consistent with other tests for similar benchmark loads and there is no evidence of resource leaks or mismanagement.vajram/vajram-krystex/src/main/java/com/flipkart/krystal/vajramexecutor/krystex/VajramKryonGraph.java (5)
84-85: Switch toConcurrentHashMapforvajramDefinitionsandvajramDataByClass.Replacing these
LinkedHashMapinstances withConcurrentHashMapimproves thread safety. However, if your code relies on predictable iteration order (asLinkedHashMapguaranteed), be aware thatConcurrentHashMapdoes not preserve insertion order. Please verify that this change does not break any ordering assumptions in the code.
88-92: AdoptingConcurrentHashMapforvajramExecutables.Similar to the above, this change enhances concurrency. Double-check if iteration order was previously significant for
vajramExecutables. If so, consider adding a note or test to confirm correctness under the new map.
177-177: Confirm that concurrency concerns aroundcomputeDependantChainare handled.Even though only one line was changed here, this method potentially interacts with concurrently modified data structures. Ensure that all relevant map and chain accesses remain thread-safe.
235-242: Introduce synchronization for loading subgraph.The synchronized block on
vajramExecutableshelps avoid the race condition. Ensure the lock coverage is sufficient and that no other references tovajramExecutablescan be mutated outside this synchronized context.
457-457: Switch to concurrent method for dependency definitions.Using
loadKryonSubgraphinsidecreateKryonDefinitionsForDependenciesis consistent with the new concurrency model. Re-check if subsequent calls (possibly from multiple threads) do not accidentally introduce cyclical dependencies or deadlocks. If needed, add tests verifying concurrency for multiple dependencies referencing each other.Also applies to: 473-474
vajram/vajram-samples/src/test/java/com/flipkart/krystal/vajram/samples/calculator/FormulaTest.java (4)
12-12: Added imports for parallel execution and concurrency utilities.The new imports (
Arrays.stream,CompletableFuture.runAsync,IntStream.range,atomic.LongAdder,org.checkerframework.checker...) suggest an expanded focus on concurrency testing. No immediate issues spotted, but keep them consistent with any chosen concurrency model (synchronous vs async).Also applies to: 15-15, 17-17, 48-48, 50-50
60-61: Added dynamic thread count and executor lease array.Using
Runtime.getRuntime().availableProcessors()to size the arrays is a good pragmatic approach. Make sure that code does not exceed these leased executors, especially if future concurrency expansions are planned.
84-85: Reusing a leased executor insetUp().Instead of creating a new executor each time, you're using an existing lease. This is more resource-friendly. Just ensure that any test concurrency does not cause cross-test interference.
565-565: RefactoredgetExecutors(int count)method signature.This method taps the pre-leased executors. If count >
MAX_THREADS, be sure to handle the scenario gracefully or document the expectation that count won’t exceed the number of leases available.krystex/src/main/java/com/flipkart/krystal/krystex/kryon/KryonDefinitionRegistry.java (1)
15-15: Switched fromLinkedHashMaptoConcurrentHashMapinKryonDefinitionRegistry.This change improves thread safety but removes guaranteed iteration order. Verify that any dependent code does not rely on insertion order, and consider a thread-safe test if you frequently modify
kryonDefinitions.Also applies to: 20-20
krystex/src/main/java/com/flipkart/krystal/krystex/LogicDefinitionRegistry.java (1)
11-16: LGTM! Thread-safe implementation using ConcurrentHashMap.The change from standard HashMap to ConcurrentHashMap for all three maps (
outputLogicDefinitions,resolverLogicDefinitions,multiResolverDefinitions) improves thread safety and addresses potential race conditions during parallel loading of the kryon subgraph.vajram/vajram-krystex/src/test/java/com/flipkart/krystal/vajramexecutor/krystex/KrystexVajramExecutorTest.java (3)
138-138: LGTM! Test method signatures simplified.The removal of
KryonExecStrategyparameter from test methods and standardization to useBATCHstrategy improves test maintainability while preserving test coverage.Also applies to: 157-157, 177-177, 201-201, 230-230, 260-260, 292-292, 318-318, 355-355, 398-398
544-565: LGTM! Executor configuration enhanced with request level cache.The changes improve thread safety by:
- Standardizing execution strategy to
BATCH- Adding request level cache with proper decorator configuration
873-875: LGTM! Test data provider simplified.The
executorConfigsToTestmethod has been appropriately simplified while maintaining test coverage for both DEPTH and BREADTH traversal strategies.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
vajram/vajram-krystex/src/main/java/com/flipkart/krystal/vajramexecutor/krystex/VajramKryonGraph.java (1)
235-281: Consider adding debug logging for better observability.Adding debug logs for the loading states would help in troubleshooting any potential issues.
private KryonId loadKryonSubgraph(VajramID vajramId, Map<VajramID, KryonId> loadingInProgress) { synchronized (vajramExecutables) { + log.debug("Loading kryon subgraph for vajramId: {}, loadingInProgress: {}", vajramId, loadingInProgress); KryonId kryonId; if ((kryonId = vajramExecutables.get(vajramId)) != null) { + log.debug("Found existing kryon: {} for vajramId: {}", kryonId, vajramId); return kryonId; } else if ((kryonId = loadingInProgress.get(vajramId)) != null) { + log.debug("Found kryon: {} in progress for vajramId: {}", kryonId, vajramId); return kryonId; }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
vajram/vajram-krystex/src/main/java/com/flipkart/krystal/vajramexecutor/krystex/VajramKryonGraph.java(5 hunks)
🔇 Additional comments (4)
vajram/vajram-krystex/src/main/java/com/flipkart/krystal/vajramexecutor/krystex/VajramKryonGraph.java (4)
84-92: LGTM! Thread-safe map implementations.The change from
LinkedHashMaptoConcurrentHashMapforvajramDefinitions,vajramDataByClass, andvajramExecutablesimproves thread safety when accessing these maps concurrently.
226-233: LGTM! Improved thread safety in getKryonId.The method now correctly handles concurrent access by using a new
loadingInProgressmap to track the loading state of kryons.
235-281: LGTM! Well-implemented synchronization for loading kryon subgraphs.The new
loadKryonSubgraphmethod effectively prevents race conditions by:
- Using synchronized block on
vajramExecutables- Checking both
vajramExecutablesandloadingInProgressmaps- Properly tracking loading state to prevent cycles
456-477: LGTM! Thread-safe dependency resolution.The method now accepts a
loadingInProgressmap to track the loading state during recursive calls, ensuring thread safety when resolving dependencies.
RamAnvesh
changed the title
Closes #328
Summary by CodeRabbit
Chores
Refactor
Tests