From 14c21aedd2eb9b98ca2c139c70f4c322f24ea750 Mon Sep 17 00:00:00 2001 From: Yang Wang Date: Thu, 27 Feb 2020 14:14:16 +1100 Subject: [PATCH] Simplify ml license checking with XpackLicenseState internals (#52684) (#52863) This change removes TrainedModelConfig#isAvailableWithLicense method with calls to XPackLicenseState#isAllowedByLicense. Please note there are subtle changes to the code logic. But they are the right changes: * Instead of Platinum license, Enterprise license nows guarantees availability. * No explicit check when the license requirement is basic. Since basic license is always available, this check is unnecessary. * Trial license is always allowed. --- .../core/ml/inference/TrainedModelConfig.java | 16 ----- .../ml/inference/TrainedModelConfigTests.java | 67 ------------------- .../TransportInternalInferModelAction.java | 4 +- 3 files changed, 2 insertions(+), 85 deletions(-) diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/TrainedModelConfig.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/TrainedModelConfig.java index e2511c39db1..33e19cb2882 100644 --- a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/TrainedModelConfig.java +++ b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/TrainedModelConfig.java @@ -21,7 +21,6 @@ import org.elasticsearch.common.xcontent.ToXContentObject; import org.elasticsearch.common.xcontent.XContentBuilder; import org.elasticsearch.common.xcontent.XContentParser; import org.elasticsearch.license.License; -import org.elasticsearch.license.XPackLicenseState; import org.elasticsearch.xpack.core.common.time.TimeUtils; import org.elasticsearch.xpack.core.ml.inference.persistence.InferenceIndexConstants; import org.elasticsearch.xpack.core.ml.job.messages.Messages; @@ -236,21 +235,6 @@ public class TrainedModelConfig implements ToXContentObject, Writeable { return licenseLevel; } - public boolean isAvailableWithLicense(XPackLicenseState licenseState) { - // Basic is always true - if (licenseLevel.equals(License.OperationMode.BASIC)) { - return true; - } - - // The model license does not matter, Platinum license gets the same functions as the highest license - if (licenseState.isAllowedByLicense(License.OperationMode.PLATINUM)) { - return true; - } - - // catch the rest, if the license is active and is at least the required model license - return licenseState.isAllowedByLicense(licenseLevel, true, false); - } - @Override public void writeTo(StreamOutput out) throws IOException { out.writeString(modelId); diff --git a/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/inference/TrainedModelConfigTests.java b/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/inference/TrainedModelConfigTests.java index e0af20ebc8b..daaa4a04aa3 100644 --- a/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/inference/TrainedModelConfigTests.java +++ b/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/inference/TrainedModelConfigTests.java @@ -20,7 +20,6 @@ import org.elasticsearch.common.xcontent.XContentHelper; import org.elasticsearch.common.xcontent.XContentParser; import org.elasticsearch.common.xcontent.XContentType; import org.elasticsearch.license.License; -import org.elasticsearch.license.XPackLicenseState; import org.elasticsearch.search.SearchModule; import org.elasticsearch.test.AbstractSerializingTestCase; import org.elasticsearch.xpack.core.ml.job.messages.Messages; @@ -44,10 +43,6 @@ import static org.elasticsearch.xpack.core.ml.utils.ToXContentParams.FOR_INTERNA import static org.hamcrest.Matchers.containsString; import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.not; -import static org.mockito.Matchers.any; -import static org.mockito.Matchers.anyBoolean; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; public class TrainedModelConfigTests extends AbstractSerializingTestCase { @@ -306,66 +301,4 @@ public class TrainedModelConfigTests extends AbstractSerializingTestCase { - final Object[] arguments = invocationOnMock.getArguments(); - assertTrue((boolean) arguments[1]); // ensure the call is made to require active license - return false; - } - ); - assertFalse(builder.setLicenseLevel(License.OperationMode.ENTERPRISE.description()).build().isAvailableWithLicense(licenseState)); - assertFalse(builder.setLicenseLevel(License.OperationMode.PLATINUM.description()).build().isAvailableWithLicense(licenseState)); - assertFalse(builder.setLicenseLevel(License.OperationMode.GOLD.description()).build().isAvailableWithLicense(licenseState)); - // Basic license always works not matter what - assertTrue(builder.setLicenseLevel(License.OperationMode.BASIC.description()).build().isAvailableWithLicense(licenseState)); - } - - public void testActivePlatinumLicenseAlwaysWorks() { - TrainedModelConfig.Builder builder = createTestInstance(randomAlphaOfLength(10)); - XPackLicenseState licenseState = mock(XPackLicenseState.class); - - when(licenseState.isAllowedByLicense(License.OperationMode.PLATINUM)).thenReturn(true); - - // Active Platinum license functions the same as Enterprise license (highest) and should always work - when(licenseState.isAllowedByLicense(any(License.OperationMode.class), anyBoolean(), anyBoolean())).thenAnswer( - invocationOnMock -> { - final Object[] arguments = invocationOnMock.getArguments(); - assertEquals(License.OperationMode.PLATINUM, arguments[0]); - assertTrue((boolean) arguments[1]); // ensure the call is made to require active license - assertTrue((boolean) arguments[2]); - return true; - } - ); - assertTrue(builder.setLicenseLevel(License.OperationMode.ENTERPRISE.description()).build().isAvailableWithLicense(licenseState)); - assertTrue(builder.setLicenseLevel(License.OperationMode.PLATINUM.description()).build().isAvailableWithLicense(licenseState)); - assertTrue(builder.setLicenseLevel(License.OperationMode.BASIC.description()).build().isAvailableWithLicense(licenseState)); - assertTrue(builder.setLicenseLevel(License.OperationMode.GOLD.description()).build().isAvailableWithLicense(licenseState)); - } - - public void testActiveGoldLicenseWillWorkWhenRequiredLevelIsGold() { - TrainedModelConfig.Builder builder = createTestInstance(randomAlphaOfLength(10)); - XPackLicenseState licenseState = mock(XPackLicenseState.class); - - // Active Gold license should work when required level is gold - when(licenseState.isAllowedByLicense(any(License.OperationMode.class), anyBoolean(), anyBoolean())).thenAnswer( - invocationOnMock -> { - final Object[] arguments = invocationOnMock.getArguments(); - assertTrue((boolean) arguments[1]); // ensure the call is made to require active license - if (License.OperationMode.PLATINUM == arguments[0] && Boolean.TRUE.equals(arguments[2])) { - return false; - } else - return License.OperationMode.GOLD == arguments[0] && Boolean.FALSE.equals(arguments[2]); - } - ); - assertFalse(builder.setLicenseLevel(License.OperationMode.ENTERPRISE.description()).build().isAvailableWithLicense(licenseState)); - assertFalse(builder.setLicenseLevel(License.OperationMode.PLATINUM.description()).build().isAvailableWithLicense(licenseState)); - assertTrue(builder.setLicenseLevel(License.OperationMode.BASIC.description()).build().isAvailableWithLicense(licenseState)); - assertTrue(builder.setLicenseLevel(License.OperationMode.GOLD.description()).build().isAvailableWithLicense(licenseState)); - } } diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportInternalInferModelAction.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportInternalInferModelAction.java index 248169b37f6..a32367a6766 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportInternalInferModelAction.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportInternalInferModelAction.java @@ -79,8 +79,8 @@ public class TransportInternalInferModelAction extends HandledTransportAction { - responseBuilder.setLicensed(trainedModelConfig.isAvailableWithLicense(licenseState)); - if (trainedModelConfig.isAvailableWithLicense(licenseState) || request.isPreviouslyLicensed()) { + responseBuilder.setLicensed(licenseState.isAllowedByLicense(trainedModelConfig.getLicenseLevel())); + if (licenseState.isAllowedByLicense(trainedModelConfig.getLicenseLevel()) || request.isPreviouslyLicensed()) { this.modelLoadingService.getModel(request.getModelId(), getModelListener); } else { listener.onFailure(LicenseUtils.newComplianceException(XPackField.MACHINE_LEARNING));