Associate node stat agg with request

Signed-off-by: John Mazanec <[email protected]>
opensearch-project · Mar 4, 2025 · c8c8edc · c8c8edc
1 parent 4cf10b8
commit c8c8edc
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 4 deletions.
diff --git a/src/main/java/org/opensearch/knn/plugin/transport/KNNStatsRequest.java b/src/main/java/org/opensearch/knn/plugin/transport/KNNStatsRequest.java
@@ -5,9 +5,12 @@
 
 package org.opensearch.knn.plugin.transport;
 
+import lombok.Getter;
+import lombok.Setter;
 import org.opensearch.action.support.nodes.BaseNodesRequest;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.knn.plugin.stats.KNNNodeStatAggregation;
 import org.opensearch.knn.plugin.stats.StatNames;
 
 import java.io.IOException;
@@ -26,6 +29,12 @@ public class KNNStatsRequest extends BaseNodesRequest<KNNStatsRequest> {
     public static final String ALL_STATS_KEY = "_all";
     private final Set<String> validStats;
     private final Set<String> statsToBeRetrieved;
+    /**
+     * Node stat aggregation associated with the request. Not serialized between nodes. Can be null.
+     */
+    @Getter
+    @Setter
+    private KNNNodeStatAggregation aggregation;
 
     /**
      * Empty constructor needed for KNNStatsTransportAction
@@ -34,6 +43,7 @@ public KNNStatsRequest() {
         super((String[]) null);
         validStats = StatNames.getNames();
         statsToBeRetrieved = new HashSet<>();
+        aggregation = null;
     }
 
     /**
@@ -46,6 +56,7 @@ public KNNStatsRequest(StreamInput in) throws IOException {
         super(in);
         validStats = in.readSet(StreamInput::readString);
         statsToBeRetrieved = in.readSet(StreamInput::readString);
+        aggregation = null;
     }
 
     /**
@@ -57,6 +68,7 @@ public KNNStatsRequest(String... nodeIds) {
         super(nodeIds);
         validStats = StatNames.getNames();
         statsToBeRetrieved = new HashSet<>();
+        aggregation = null;
     }
 
     /**

diff --git a/src/main/java/org/opensearch/knn/plugin/transport/KNNStatsTransportAction.java b/src/main/java/org/opensearch/knn/plugin/transport/KNNStatsTransportAction.java
@@ -38,7 +38,6 @@ public class KNNStatsTransportAction extends TransportNodesAction<
 
     private final KNNStats knnStats;
     private final Client client;
-    private KNNNodeStatAggregation knnNodeStatAggregation;
 
     /**
      * Constructor
@@ -71,7 +70,6 @@ public KNNStatsTransportAction(
         );
         this.knnStats = knnStats;
         this.client = client;
-        this.knnNodeStatAggregation = null;
     }
 
     @Override
@@ -98,7 +96,7 @@ protected void doExecute(Task task, KNNStatsRequest request, ActionListener<KNNS
             // Add the stats makes sure that we dont recurse infinitely.
             dependentStats.forEach(knnStatsRequest::addStat);
             client.execute(KNNStatsAction.INSTANCE, knnStatsRequest, ActionListener.wrap(knnStatsResponse -> {
-                knnNodeStatAggregation = new KNNNodeStatAggregation(knnStatsResponse.getNodes());
+                request.setAggregation(new KNNNodeStatAggregation(knnStatsResponse.getNodes()));
                 contextListener.onResponse(null);
             }, contextListener::onFailure));
         } else {
@@ -118,7 +116,7 @@ protected KNNStatsResponse newResponse(
 
         for (String statName : knnStats.getClusterStats().keySet()) {
             if (statsToBeRetrieved.contains(statName)) {
-                clusterStats.put(statName, knnStats.getStats().get(statName).getValue(knnNodeStatAggregation));
+                clusterStats.put(statName, knnStats.getStats().get(statName).getValue(request.getAggregation()));
             }
         }