-
Notifications
You must be signed in to change notification settings - Fork 25.9k
Add utility for tracking licensed persistent tasks #76672
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
Merged
Merged
Changes from all commits
Commits
Show all changes
10 commits
Select commit
Hold shift + click to select a range
00109e5
Add utility for tracking licensed persistent tasks
rjernst 54ec74e
make mark methods final
rjernst 1abe6be
tweak feature name
rjernst 06a534b
ensure init is not overridden
rjernst 8f81bc4
checkstyle
rjernst 78f1b9c
checkstyle
rjernst 557c0c6
test fix
rjernst 7162a77
Merge branch 'master' into license/ml_tracking
rjernst 3398a19
add family
rjernst 5c2674e
add family to tests
rjernst File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
85 changes: 85 additions & 0 deletions
85
.../plugin/core/src/main/java/org/elasticsearch/license/LicensedAllocatedPersistentTask.java
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,85 @@ | ||
| /* | ||
| * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one | ||
| * or more contributor license agreements. Licensed under the Elastic License | ||
| * 2.0; you may not use this file except in compliance with the Elastic License | ||
| * 2.0. | ||
| */ | ||
|
|
||
| package org.elasticsearch.license; | ||
|
|
||
| import org.elasticsearch.persistent.AllocatedPersistentTask; | ||
| import org.elasticsearch.persistent.PersistentTasksService; | ||
| import org.elasticsearch.tasks.TaskId; | ||
| import org.elasticsearch.tasks.TaskManager; | ||
|
|
||
| import java.util.Map; | ||
|
|
||
| /** | ||
| * An AllocatedPersistentTask which automatically tracks as a licensed feature usage. | ||
| */ | ||
| public class LicensedAllocatedPersistentTask extends AllocatedPersistentTask { | ||
| private final LicensedFeature.Persistent licensedFeature; | ||
| private final String featureContext; | ||
| private final XPackLicenseState licenseState; | ||
|
|
||
| public LicensedAllocatedPersistentTask(long id, String type, String action, String description, TaskId parentTask, | ||
| Map<String, String> headers, LicensedFeature.Persistent feature, String featureContext, | ||
| XPackLicenseState licenseState) { | ||
| super(id, type, action, description, parentTask, headers); | ||
| this.licensedFeature = feature; | ||
| this.featureContext = featureContext; | ||
| this.licenseState = licenseState; | ||
| licensedFeature.startTracking(licenseState, featureContext); | ||
| } | ||
|
|
||
| private void stopTracking() { | ||
| licensedFeature.stopTracking(licenseState, featureContext); | ||
| } | ||
|
|
||
| @Override | ||
| protected final boolean markAsCancelled() { | ||
| stopTracking(); | ||
| return doMarkAsCancelled(); | ||
| } | ||
|
|
||
| protected boolean doMarkAsCancelled() { | ||
| return super.markAsCancelled(); | ||
| } | ||
|
|
||
| @Override | ||
| public final void markAsCompleted() { | ||
| stopTracking(); | ||
| doMarkAsCompleted(); | ||
| } | ||
|
|
||
| protected void doMarkAsCompleted() { | ||
| super.markAsCompleted(); | ||
| } | ||
|
|
||
| @Override | ||
| public final void markAsFailed(Exception e) { | ||
| stopTracking(); | ||
| doMarkAsFailed(e); | ||
| } | ||
|
|
||
| protected void doMarkAsFailed(Exception e) { | ||
| super.markAsFailed(e); | ||
| } | ||
|
|
||
| @Override | ||
| public final void markAsLocallyAborted(String localAbortReason) { | ||
| stopTracking(); | ||
| doMarkAsLocallyAborted(localAbortReason); | ||
| } | ||
|
|
||
| protected void doMarkAsLocallyAborted(String localAbortReason) { | ||
| super.markAsLocallyAborted(localAbortReason); | ||
| } | ||
|
|
||
| // this is made public for tests, and final to ensure it is not overridden with something that may throw | ||
| @Override | ||
| public final void init(PersistentTasksService persistentTasksService, TaskManager taskManager, | ||
| String persistentTaskId, long allocationId) { | ||
| super.init(persistentTasksService, taskManager, persistentTaskId, allocationId); | ||
| } | ||
| } | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
93 changes: 93 additions & 0 deletions
93
...in/core/src/test/java/org/elasticsearch/license/LicensedAllocatedPersistentTaskTests.java
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,93 @@ | ||
| /* | ||
| * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one | ||
| * or more contributor license agreements. Licensed under the Elastic License | ||
| * 2.0; you may not use this file except in compliance with the Elastic License | ||
| * 2.0. | ||
| */ | ||
|
|
||
| package org.elasticsearch.license; | ||
|
|
||
| import org.elasticsearch.persistent.PersistentTasksService; | ||
| import org.elasticsearch.tasks.TaskId; | ||
| import org.elasticsearch.tasks.TaskManager; | ||
| import org.elasticsearch.test.ESTestCase; | ||
|
|
||
| import java.util.Map; | ||
| import java.util.concurrent.atomic.AtomicBoolean; | ||
| import java.util.function.Consumer; | ||
|
|
||
| import static org.hamcrest.Matchers.is; | ||
| import static org.mockito.Mockito.mock; | ||
| import static org.mockito.Mockito.times; | ||
| import static org.mockito.Mockito.verify; | ||
|
|
||
| public class LicensedAllocatedPersistentTaskTests extends ESTestCase { | ||
|
|
||
| void assertTrackingComplete(Consumer<LicensedAllocatedPersistentTask> method) { | ||
| XPackLicenseState licenseState = mock(XPackLicenseState.class); | ||
| LicensedFeature.Persistent feature = LicensedFeature.persistent("family", "somefeature", License.OperationMode.PLATINUM); | ||
| var task = new LicensedAllocatedPersistentTask(0, "type", "action", "description", TaskId.EMPTY_TASK_ID, Map.of(), | ||
| feature, "context", licenseState); | ||
| PersistentTasksService service = mock(PersistentTasksService.class); | ||
| TaskManager taskManager = mock(TaskManager.class); | ||
| task.init(service, taskManager, "id", 0); | ||
| verify(licenseState, times(1)).enableUsageTracking(feature, "context"); | ||
| method.accept(task); | ||
| verify(licenseState, times(1)).disableUsageTracking(feature, "context"); | ||
| } | ||
|
|
||
| public void testCompleted() { | ||
| assertTrackingComplete(LicensedAllocatedPersistentTask::markAsCompleted); | ||
| } | ||
|
|
||
| public void testCancelled() { | ||
| assertTrackingComplete(LicensedAllocatedPersistentTask::markAsCancelled); | ||
| } | ||
|
|
||
| public void testFailed() { | ||
| assertTrackingComplete(t -> t.markAsFailed(null)); | ||
| } | ||
|
|
||
| public void testLocallyAborted() { | ||
| assertTrackingComplete(t -> t.markAsLocallyAborted("reason")); | ||
| } | ||
|
|
||
| public void testDoOverrides() { | ||
| XPackLicenseState licenseState = mock(XPackLicenseState.class); | ||
| LicensedFeature.Persistent feature = LicensedFeature.persistent("family", "somefeature", License.OperationMode.PLATINUM); | ||
|
|
||
| AtomicBoolean completedCalled = new AtomicBoolean(); | ||
| AtomicBoolean cancelledCalled = new AtomicBoolean(); | ||
| AtomicBoolean failedCalled = new AtomicBoolean(); | ||
| AtomicBoolean abortedCalled = new AtomicBoolean(); | ||
| var task = new LicensedAllocatedPersistentTask(0, "type", "action", "description", TaskId.EMPTY_TASK_ID, Map.of(), | ||
| feature, "context", licenseState) { | ||
| @Override | ||
| protected boolean doMarkAsCancelled() { | ||
| cancelledCalled.set(true); | ||
| return true; | ||
| } | ||
| @Override | ||
| protected void doMarkAsCompleted() { | ||
| completedCalled.set(true); | ||
| } | ||
| @Override | ||
| protected void doMarkAsFailed(Exception e) { | ||
| failedCalled.set(true); | ||
| } | ||
| @Override | ||
| protected void doMarkAsLocallyAborted(String reason) { | ||
| abortedCalled.set(true); | ||
| } | ||
| }; | ||
|
|
||
| task.markAsCancelled(); | ||
| assertThat(cancelledCalled.get(), is(true)); | ||
| task.markAsCompleted(); | ||
| assertThat(completedCalled.get(), is(true)); | ||
| task.markAsFailed(null); | ||
| assertThat(failedCalled.get(), is(true)); | ||
| task.markAsLocallyAborted("reason"); | ||
| assertThat(abortedCalled.get(), is(true)); | ||
| } | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
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.
Looking at the taskManager code, I think there is a small chance that a task object is destroyed but none of the cancellations methods are called.
elasticsearch/server/src/main/java/org/elasticsearch/persistent/PersistentTasksNodeService.java
Lines 178 to 202 in 7d980ad
Note if
initis called and throws, onlytaskManager.unregister(task);is called. I don't think this calls any of these internal life time tracking methods.I am thinking that
licensedFeature.startTracking(licenseState, featureContext);should probably be called ininitto avoid this weird condition.What do you think?
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 seems like init is being abused currently. It's purpose, IMO is to just set some references, not to do any work? But at least rollup seems to utilize this.
What if instead I were to make init public and final? Public is so that tests can still call it, and then final so that there is no chance some extra work is done that could actually throw an exception.
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.
@rjernst if we can make that guarantee of it not throwing, it makes sense to me. All it's doing is setting internal values...so I don't know why it would ever throw.
So, I think fix init, or somehow make sure that
stopTrackingis called ifinitthrows in this execution path.