-
Notifications
You must be signed in to change notification settings - Fork 4.5k
[BEAM-12524] Ensure that failed BundleProcessor objects are not re-added to the cache. #15062
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -309,12 +309,13 @@ public BeamFnApi.InstructionResponse.Builder processBundle(BeamFnApi.Instruction | |
| throw new RuntimeException(e); | ||
| } | ||
| }); | ||
| PTransformFunctionRegistry startFunctionRegistry = bundleProcessor.getStartFunctionRegistry(); | ||
| PTransformFunctionRegistry finishFunctionRegistry = bundleProcessor.getFinishFunctionRegistry(); | ||
| ExecutionStateTracker stateTracker = bundleProcessor.getStateTracker(); | ||
| QueueingBeamFnDataClient queueingClient = bundleProcessor.getQueueingClient(); | ||
|
|
||
| try { | ||
| PTransformFunctionRegistry startFunctionRegistry = bundleProcessor.getStartFunctionRegistry(); | ||
| PTransformFunctionRegistry finishFunctionRegistry = | ||
| bundleProcessor.getFinishFunctionRegistry(); | ||
| ExecutionStateTracker stateTracker = bundleProcessor.getStateTracker(); | ||
| QueueingBeamFnDataClient queueingClient = bundleProcessor.getQueueingClient(); | ||
|
|
||
| try (HandleStateCallsForBundle beamFnStateClient = bundleProcessor.getBeamFnStateClient()) { | ||
| try (Closeable closeTracker = stateTracker.activate()) { | ||
| // Already in reverse topological order so we don't need to do anything. | ||
|
|
@@ -354,14 +355,16 @@ public BeamFnApi.InstructionResponse.Builder processBundle(BeamFnApi.Instruction | |
| response.setRequiresFinalization(true); | ||
| } | ||
| } | ||
|
|
||
| // Mark the bundle processor as re-usable. | ||
| bundleProcessorCache.release( | ||
| request.getProcessBundle().getProcessBundleDescriptorId(), bundleProcessor); | ||
| return BeamFnApi.InstructionResponse.newBuilder().setProcessBundle(response); | ||
| } catch (Exception e) { | ||
| bundleProcessorCache.release( | ||
| request.getProcessBundle().getProcessBundleDescriptorId(), bundleProcessor); | ||
| // Make sure we clean-up from the active set of bundle processors. | ||
| bundleProcessorCache.discard(bundleProcessor); | ||
| throw e; | ||
| } | ||
| return BeamFnApi.InstructionResponse.newBuilder().setProcessBundle(response); | ||
| } | ||
|
|
||
| public BeamFnApi.InstructionResponse.Builder progress(BeamFnApi.InstructionRequest request) | ||
|
|
@@ -648,13 +651,12 @@ public BundleProcessor find(String instructionId) { | |
| } | ||
|
|
||
| /** | ||
| * Add a {@link BundleProcessor} to cache. The {@link BundleProcessor} will be reset before | ||
| * being added to the cache and will be marked as inactive. | ||
| * Add a {@link BundleProcessor} to cache. The {@link BundleProcessor} will be marked as | ||
| * inactive and reset before being added to the cache. | ||
| */ | ||
| void release(String bundleDescriptorId, BundleProcessor bundleProcessor) { | ||
| activeBundleProcessors.remove(bundleProcessor.getInstructionId()); | ||
| try { | ||
| bundleProcessor.setInstructionId(null); | ||
| bundleProcessor.reset(); | ||
| cachedBundleProcessors.get(bundleDescriptorId).add(bundleProcessor); | ||
| } catch (Exception e) { | ||
|
|
@@ -665,6 +667,11 @@ void release(String bundleDescriptorId, BundleProcessor bundleProcessor) { | |
| } | ||
| } | ||
|
|
||
| /** Discard an active {@link BundleProcessor} instead of being re-used. */ | ||
| void discard(BundleProcessor bundleProcessor) { | ||
| activeBundleProcessors.remove(bundleProcessor.getInstructionId()); | ||
| } | ||
|
|
||
| /** Shutdown all the cached {@link BundleProcessor}s, running the tearDown() functions. */ | ||
| void shutdown() throws Exception { | ||
| cachedBundleProcessors.invalidateAll(); | ||
|
|
@@ -742,6 +749,7 @@ synchronized void setInstructionId(String instructionId) { | |
| } | ||
|
|
||
| void reset() throws Exception { | ||
| this.instructionId = null; | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I think this broke Java PreCommit (it looks like this was submitted while precommits weren't running). Running
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Opened #15085 |
||
| getStartFunctionRegistry().reset(); | ||
| getFinishFunctionRegistry().reset(); | ||
| getSplitListener().clear(); | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Do we want to do
bundleProcessor.setInstructionId(null);before discarding it to avoid the same race condition asreset?There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
It is unnecessary since reset will not be invoked.