Refactor license feature tracking

x-pack/plugin/core/src/main/java/org/elasticsearch/license/GetFeatureUsageResponse.java

@@ -7,43 +7,56 @@
 
 package org.elasticsearch.license;
 
+import org.elasticsearch.Version;
 import org.elasticsearch.action.ActionResponse;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.common.xcontent.ToXContentObject;
 import org.elasticsearch.common.xcontent.XContentBuilder;
+import org.elasticsearch.core.Nullable;
 
 import java.io.IOException;
 import java.time.Instant;
 import java.time.ZoneOffset;
 import java.time.ZonedDateTime;
 import java.util.Collections;
 import java.util.List;
+import java.util.Objects;
 
 public class GetFeatureUsageResponse extends ActionResponse implements ToXContentObject {
 
     public static class FeatureUsageInfo implements Writeable {
-        public final String name;
-        public final ZonedDateTime lastUsedTime;
+        private final String name;
+        private final ZonedDateTime lastUsedTime;
+        private final String context;
         public final String licenseLevel;
 
-        public FeatureUsageInfo(String name, ZonedDateTime lastUsedTime, String licenseLevel) {
-            this.name = name;
-            this.lastUsedTime = lastUsedTime;
-            this.licenseLevel = licenseLevel;
+        public FeatureUsageInfo(String name, ZonedDateTime lastUsedTime, @Nullable String context, String licenseLevel) {
+            this.name = Objects.requireNonNull(name, "Feature name may not be null");
+            this.lastUsedTime = Objects.requireNonNull(lastUsedTime, "Last used time may not be null");
+            this.context = context;
+            this.licenseLevel = Objects.requireNonNull(licenseLevel, "License level may not be null");
         }
 
         public FeatureUsageInfo(StreamInput in) throws IOException {
             this.name = in.readString();
             this.lastUsedTime = ZonedDateTime.ofInstant(Instant.ofEpochSecond(in.readLong()), ZoneOffset.UTC);
+            if (in.getVersion().onOrAfter(Version.V_8_0_0)) {
+                this.context = in.readOptionalString();
+            } else {
+                this.context = null;
+            }
             this.licenseLevel = in.readString();
         }
 
         @Override
         public void writeTo(StreamOutput out) throws IOException {
             out.writeString(name);
             out.writeLong(lastUsedTime.toEpochSecond());
+            if (out.getVersion().onOrAfter(Version.V_8_0_0)) {
+                out.writeOptionalString(this.context);
+            }
             out.writeString(licenseLevel);
         }
     }
@@ -74,6 +87,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         for (FeatureUsageInfo feature : features) {
             builder.startObject();
             builder.field("name", feature.name);
+            builder.field("context", feature.context);
             builder.field("last_used", feature.lastUsedTime.toString());
             builder.field("license_level", feature.licenseLevel);
             builder.endObject();

x-pack/plugin/core/src/main/java/org/elasticsearch/license/LicensedFeature.java

@@ -0,0 +1,124 @@
+/*
+ * 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 java.util.Objects;
+
+/**
+ * A base class for checking licensed features against the license.
+ */
+public abstract class LicensedFeature {
+
+    public static class Momentary extends LicensedFeature {
+
+        private Momentary(String name, License.OperationMode minimumOperationMode, boolean needsActive) {
+            super(name, minimumOperationMode, needsActive);
+        }
+
+        /**
+         * Checks whether the feature is allowed by the given license state, and
+         * updates the last time the feature was used.
+         */
+        public boolean check(XPackLicenseState state) {
+            if (state.isAllowed(this)) {
+                state.featureUsed(this);
+                return true;
+            } else {
+                return false;
+            }
+        }
+    }
+
+    public static class Persistent extends LicensedFeature {
+        private Persistent(String name, License.OperationMode minimumOperationMode, boolean needsActive) {
+            super(name, minimumOperationMode, needsActive);
+        }
+
+        public boolean checkAndStartTracking(XPackLicenseState state, String contextName) {
+            if (state.isAllowed(this)) {
+                state.enableUsageTracking(this, contextName);
+                return true;
+            } else {
+                return false;
+            }
+        }
+
+        public void stopTracking(XPackLicenseState state, String contextName) {
+            state.disableUsageTracking(this, contextName);
+        }
+    }
+
+    public static class Untracked extends LicensedFeature {
+        private Untracked(String name, License.OperationMode minimumOperationMode, boolean needsActive) {
+            super(name, minimumOperationMode, needsActive);
+        }
+
+        public boolean check(XPackLicenseState state) {
+            return state.isAllowed(this);
+        }
+    }
+
+    final String name;
+    final License.OperationMode minimumOperationMode;
+    final boolean needsActive;
+
+    public LicensedFeature(String name, License.OperationMode minimumOperationMode, boolean needsActive) {
+        this.name = name;
+        this.minimumOperationMode = minimumOperationMode;
+        this.needsActive = needsActive;
+    }
+
+    /**
+     * Creates a feature that is tracked at the moment it is checked.
+     * @param name A unique name for the feature that will be returned in
+     *             the tracking API. This should not change.
+     * @param licenseLevel The lowest level of license in which this feature should be allowed.
+     */
+    public static Momentary momentary(String name, License.OperationMode licenseLevel) {
+        return new Momentary(name, licenseLevel, true);
+    }
+
+    public static Persistent persistent(String name, License.OperationMode licenseLevel) {
+        return new Persistent(name, licenseLevel, true);
+    }
+
+    /**
+     * Creates a feature that is tracked at the moment it is checked, but that is lenient as
+     * to whether the license needs to be active to allow the feature.
+     */
+    @Deprecated
+    public static Momentary momentaryLenient(String name, License.OperationMode licenseLevel) {
+        return new Momentary(name, licenseLevel, false);
+    }
+
+    @Deprecated
+    public static Persistent persistentLenient(String name, License.OperationMode licenseLevel) {
+        return new Persistent(name, licenseLevel, false);
+    }
+
+    /**
+     * Returns whether the feature is allowed by the current license
+     * without affecting feature tracking.
+     */
+    public final boolean checkWithoutTracking(XPackLicenseState state) {
+        return state.isAllowed(this);
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        LicensedFeature that = (LicensedFeature) o;
+        return Objects.equals(name, that.name);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(name);
+    }
+}

x-pack/plugin/core/src/main/java/org/elasticsearch/license/TransportGetFeatureUsageAction.java

@@ -20,7 +20,6 @@
 import java.time.ZonedDateTime;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.Locale;
 import java.util.Map;
 
 public class TransportGetFeatureUsageAction extends HandledTransportAction<GetFeatureUsageRequest, GetFeatureUsageResponse> {
@@ -40,15 +39,19 @@ public TransportGetFeatureUsageAction(TransportService transportService, ActionF
 
     @Override
     protected void doExecute(Task task, GetFeatureUsageRequest request, ActionListener<GetFeatureUsageResponse> listener) {
-        Map<XPackLicenseState.Feature, Long> featureUsage = licenseState.getLastUsed();
-        List<GetFeatureUsageResponse.FeatureUsageInfo> usageInfos = new ArrayList<>();
-        for (var entry : featureUsage.entrySet()) {
-            XPackLicenseState.Feature feature = entry.getKey();
-            String name = feature.name().toLowerCase(Locale.ROOT);
-            ZonedDateTime lastUsedTime = Instant.ofEpochMilli(entry.getValue()).atZone(ZoneOffset.UTC);
-            String licenseLevel = feature.minimumOperationMode.name().toLowerCase(Locale.ROOT);
-            usageInfos.add(new GetFeatureUsageResponse.FeatureUsageInfo(name, lastUsedTime, licenseLevel));
-        }
+        Map<XPackLicenseState.FeatureUsage, Long> featureUsage = licenseState.getLastUsed();
+        List<GetFeatureUsageResponse.FeatureUsageInfo> usageInfos = new ArrayList<>(featureUsage.size());
+        featureUsage.forEach((usage, lastUsed) -> {
+            ZonedDateTime lastUsedTime = Instant.ofEpochMilli(lastUsed).atZone(ZoneOffset.UTC);
+            usageInfos.add(
+                new GetFeatureUsageResponse.FeatureUsageInfo(
+                    usage.featureName(),
+                    lastUsedTime,
+                    usage.contextName(),
+                    usage.minimumOperationMode().description()
+                )
+            );
+        });
         listener.onResponse(new GetFeatureUsageResponse(usageInfos));
     }
 }