feat: opensearch-project#87 implemented rest endpoint to make stats a…

…vailable

Signed-off-by: Johannes Peter <[email protected]>

src/main/java/com/o19s/es/ltr/LtrQueryParserPlugin.java

@@ -19,13 +19,16 @@
 import ciir.umass.edu.learning.RankerFactory;
 import org.opensearch.ltr.breaker.LTRCircuitBreakerService;
 import org.opensearch.ltr.settings.LTRSettings;
+import org.opensearch.ltr.rest.RestStatsLTRAction;
 import org.opensearch.ltr.stats.LTRStat;
 import org.opensearch.ltr.stats.LTRStats;
 import org.opensearch.ltr.stats.StatName;
 import org.opensearch.ltr.stats.suppliers.CacheStatsOnNodeSupplier;
 import org.opensearch.ltr.stats.suppliers.PluginHealthStatusSupplier;
 import org.opensearch.ltr.stats.suppliers.StoreStatsSupplier;
 import org.opensearch.ltr.stats.suppliers.CounterSupplier;
+import org.opensearch.ltr.transport.LTRStatsAction;
+import org.opensearch.ltr.transport.TransportLTRStatsAction;
 import com.o19s.es.explore.ExplorerQueryBuilder;
 import com.o19s.es.ltr.action.AddFeaturesToSetAction;
 import com.o19s.es.ltr.action.CachesStatsAction;
@@ -201,6 +204,7 @@ public List<RestHandler> getRestHandlers(Settings settings, RestController restC
         list.add(new RestFeatureStoreCaches());
         list.add(new RestCreateModelFromSet());
         list.add(new RestAddFeatureToSet());
+        list.add(new RestStatsLTRAction(ltrStats));
         return unmodifiableList(list);
     }
 
@@ -212,7 +216,8 @@ public List<RestHandler> getRestHandlers(Settings settings, RestController restC
                 new ActionHandler<>(ClearCachesAction.INSTANCE, TransportClearCachesAction.class),
                 new ActionHandler<>(AddFeaturesToSetAction.INSTANCE, TransportAddFeatureToSetAction.class),
                 new ActionHandler<>(CreateModelFromSetAction.INSTANCE, TransportCreateModelFromSetAction.class),
-                new ActionHandler<>(ListStoresAction.INSTANCE, TransportListStoresAction.class)));
+                new ActionHandler<>(ListStoresAction.INSTANCE, TransportListStoresAction.class),
+                new ActionHandler<>(LTRStatsAction.INSTANCE, TransportLTRStatsAction.class)));
     }
 
     @Override

src/main/java/org/opensearch/ltr/rest/RestStatsLTRAction.java

@@ -0,0 +1,164 @@
+/*
+ * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package org.opensearch.ltr.rest;
+
+import org.opensearch.client.node.NodeClient;
+import org.opensearch.ltr.stats.LTRStats;
+import org.opensearch.ltr.transport.LTRStatsAction;
+import org.opensearch.ltr.transport.LTRStatsRequest;
+import org.opensearch.rest.BaseRestHandler;
+import org.opensearch.rest.RestRequest;
+import org.opensearch.rest.action.RestActions;
+
+import java.io.IOException;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Locale;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+import static com.o19s.es.ltr.LtrQueryParserPlugin.LTR_BASE_URI;
+import static com.o19s.es.ltr.LtrQueryParserPlugin.LTR_LEGACY_BASE_URI;
+
+/**
+ * Provide an API to get information on the plugin usage and
+ * performance, such as
+ * <ul>
+ *     <li>statistics on plugin's cache performance</li>
+ *     <li>statistics on indices used to store features, feature sets and model definitions.</li>
+ *     <li>information on overall plugin status</li>
+ * </ul>
+ */
+public class RestStatsLTRAction extends BaseRestHandler {
+    private static final String NAME = "learning_to_rank_stats";
+    private final LTRStats ltrStats;
+
+    public RestStatsLTRAction(final LTRStats ltrStats) {
+        this.ltrStats = ltrStats;
+    }
+
+    @Override
+    public String getName() {
+        return NAME;
+    }
+
+    @Override
+    public List<Route> routes() {
+        return List.of();
+    }
+
+    @Override
+    public List<ReplacedRoute> replacedRoutes() {
+        return List.of(
+                new ReplacedRoute(
+                        RestRequest.Method.GET,
+                        String.format(Locale.ROOT, "%s%s", LTR_BASE_URI, "/{nodeId}/stats/"),
+                        RestRequest.Method.GET,
+                        String.format(Locale.ROOT, "%s%s", LTR_LEGACY_BASE_URI, "/{nodeId}/stats/")
+                ),
+                new ReplacedRoute(
+                        RestRequest.Method.GET,
+                        String.format(Locale.ROOT, "%s%s", LTR_BASE_URI, "/{nodeId}/stats/{stat}"),
+                        RestRequest.Method.GET,
+                        String.format(Locale.ROOT, "%s%s", LTR_LEGACY_BASE_URI, "/{nodeId}/stats/{stat}")
+                ),
+                new ReplacedRoute(
+                        RestRequest.Method.GET,
+                        String.format(Locale.ROOT, "%s%s", LTR_BASE_URI, "/stats/"),
+                        RestRequest.Method.GET,
+                        String.format(Locale.ROOT, "%s%s", LTR_LEGACY_BASE_URI, "/stats/")
+                ),
+                new ReplacedRoute(
+                        RestRequest.Method.GET,
+                        String.format(Locale.ROOT, "%s%s", LTR_BASE_URI, "/stats/{stat}"),
+                        RestRequest.Method.GET,
+                        String.format(Locale.ROOT, "%s%s", LTR_LEGACY_BASE_URI, "/stats/{stat}")
+                )
+        );
+    }
+
+    @Override
+    @SuppressWarnings({"rawtypes", "unchecked"})
+    protected RestChannelConsumer prepareRequest(final RestRequest request, final NodeClient client) throws IOException {
+        final LTRStatsRequest ltrStatsRequest = getRequest(request);
+        return (channel) -> client.execute(LTRStatsAction.INSTANCE,
+                ltrStatsRequest,
+                new RestActions.NodesResponseRestListener(channel));
+    }
+
+    /**
+     * Creates a LTRStatsRequest from a RestRequest
+     *
+     * @param request RestRequest
+     * @return LTRStatsRequest
+     */
+    private LTRStatsRequest getRequest(final RestRequest request) {
+        final LTRStatsRequest ltrStatsRequest = new LTRStatsRequest(
+                splitCommaSeparatedParam(request, "nodeId")
+        );
+        ltrStatsRequest.timeout(request.param("timeout"));
+
+        final List<String> requestedStats = List.of(splitCommaSeparatedParam(request, "stat"));
+
+        final Set<String> validStats = ltrStats.getStats().keySet();
+
+        if (isAllStatsRequested(requestedStats)) {
+            ltrStatsRequest.addAll(validStats);
+
+        } else {
+            ltrStatsRequest.addAll(getStatsToBeRetrieved(request, validStats, requestedStats));
+        }
+
+        return ltrStatsRequest;
+    }
+
+    private Set<String> getStatsToBeRetrieved(
+            final RestRequest request,
+            final Set<String> validStats,
+            final List<String> requestedStats) {
+
+        if (requestedStats.contains(LTRStatsRequest.ALL_STATS_KEY)) {
+            throw new IllegalArgumentException(String.format("Request %s contains both %s and individual stats",
+                    request.path(), LTRStatsRequest.ALL_STATS_KEY));
+        }
+
+        final Set<String> invalidStats = requestedStats.stream()
+                .filter(s -> !validStats.contains(s))
+                .collect(Collectors.toSet());
+
+        if (!invalidStats.isEmpty()) {
+            throw new IllegalArgumentException(
+                    unrecognized(request, invalidStats, new HashSet<>(requestedStats), "stat"));
+        }
+
+        return new HashSet<>(requestedStats);
+    }
+
+    private boolean isAllStatsRequested(final List<String> requestedStats) {
+        return requestedStats.isEmpty()
+                || (requestedStats.size() == 1 && requestedStats.contains(LTRStatsRequest.ALL_STATS_KEY));
+    }
+
+    private String[] splitCommaSeparatedParam(final RestRequest request, final String paramName) {
+        final String param = request.param(paramName);
+
+        if (param == null) {
+            return new String[0];
+        } else {
+            return param.split(",");
+        }
+    }
+}

src/main/java/org/opensearch/ltr/transport/LTRStatsAction.java

@@ -0,0 +1,38 @@
+/*
+ * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package org.opensearch.ltr.transport;
+
+import org.opensearch.action.ActionRequestBuilder;
+import org.opensearch.action.ActionType;
+import org.opensearch.client.OpenSearchClient;
+
+public class LTRStatsAction extends ActionType<LTRStatsNodesResponse> {
+    public static final String NAME = "cluster:admin/ltr/stats";
+    public static final LTRStatsAction INSTANCE = new LTRStatsAction();
+
+    protected LTRStatsAction() {
+        super(NAME, LTRStatsNodesResponse::new);
+    }
+
+    public static class LTRStatsRequestBuilder
+            extends ActionRequestBuilder<LTRStatsRequest, LTRStatsNodesResponse> {
+        private static final String[] nodeIds = null;
+
+        protected LTRStatsRequestBuilder(final OpenSearchClient client) {
+            super(client, INSTANCE, new LTRStatsRequest(nodeIds));
+        }
+    }
+}

src/main/java/org/opensearch/ltr/transport/LTRStatsNodeRequest.java

@@ -0,0 +1,48 @@
+/*
+ * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package org.opensearch.ltr.transport;
+
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.transport.TransportRequest;
+
+import java.io.IOException;
+
+/**
+ * LTRStatsNodeRequest to get a node stat
+ */
+public class LTRStatsNodeRequest extends TransportRequest {
+    private final LTRStatsRequest ltrStatsRequest;
+
+    public LTRStatsNodeRequest(final LTRStatsRequest ltrStatsRequest) {
+        this.ltrStatsRequest = ltrStatsRequest;
+    }
+
+    public LTRStatsNodeRequest(final StreamInput in) throws IOException {
+        super(in);
+        ltrStatsRequest = new LTRStatsRequest(in);
+    }
+
+    public LTRStatsRequest getLTRStatsNodesRequest() {
+        return ltrStatsRequest;
+    }
+
+    @Override
+    public void writeTo(final StreamOutput out) throws IOException {
+        super.writeTo(out);
+        ltrStatsRequest.writeTo(out);
+    }
+}

src/main/java/org/opensearch/ltr/transport/LTRStatsNodeResponse.java

@@ -0,0 +1,63 @@
+/*
+ * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package org.opensearch.ltr.transport;
+
+import org.opensearch.action.support.nodes.BaseNodeResponse;
+import org.opensearch.cluster.node.DiscoveryNode;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.core.xcontent.ToXContentFragment;
+import org.opensearch.core.xcontent.XContentBuilder;
+
+import java.io.IOException;
+import java.util.Map;
+
+public class LTRStatsNodeResponse extends BaseNodeResponse implements ToXContentFragment {
+
+    private final Map<String, Object> statsMap;
+
+    LTRStatsNodeResponse(final StreamInput in) throws IOException {
+        super(in);
+        this.statsMap = in.readMap(StreamInput::readString, StreamInput::readGenericValue);
+    }
+
+    LTRStatsNodeResponse(final DiscoveryNode node, final Map<String, Object> statsToValues) {
+        super(node);
+        this.statsMap = statsToValues;
+    }
+
+    public static LTRStatsNodeResponse readStats(final StreamInput in) throws IOException {
+        return new LTRStatsNodeResponse(in);
+    }
+
+    public Map<String, Object> getStatsMap() {
+        return statsMap;
+    }
+
+    @Override
+    public void writeTo(final StreamOutput out) throws IOException {
+        super.writeTo(out);
+        out.writeMap(statsMap, StreamOutput::writeString, StreamOutput::writeGenericValue);
+    }
+
+    public XContentBuilder toXContent(final XContentBuilder builder, final Params params) throws IOException {
+        for (Map.Entry<String, Object> stat : statsMap.entrySet()) {
+            builder.field(stat.getKey(), stat.getValue());
+        }
+
+        return builder;
+    }
+}