Try to maintain backward compatibility of APIs

Instead of extracting an interface from QueryRewriteContext, we
can keep it as a base class and subclass it for QueryCoordinatorContext.

Signed-off-by: Michael Froh <[email protected]>

Author: msfroh

server/src/main/java/org/opensearch/index/query/BaseQueryRewriteContext.java

@@ -8,16 +8,11 @@
 
 package org.opensearch.index.query;
 
-import org.opensearch.client.Client;
 import org.opensearch.common.annotation.PublicApi;
-import org.opensearch.core.action.ActionListener;
-import org.opensearch.core.common.io.stream.NamedWriteableRegistry;
-import org.opensearch.core.xcontent.NamedXContentRegistry;
 import org.opensearch.search.pipeline.PipelinedRequest;
 
 import java.util.HashMap;
 import java.util.Map;
-import java.util.function.BiConsumer;
 
 /**
  * The QueryCoordinatorContext class implements the QueryRewriteContext interface and provides
@@ -29,53 +24,12 @@
  * @since 2.19.0
  */
 @PublicApi(since = "2.19.0")
-public class QueryCoordinatorContext implements QueryRewriteContext {
-    private final QueryRewriteContext rewriteContext;
+public class QueryCoordinatorContext extends QueryRewriteContext {
     private final PipelinedRequest searchRequest;
 
-    public QueryCoordinatorContext(QueryRewriteContext rewriteContext, PipelinedRequest searchRequest) {
-        this.rewriteContext = rewriteContext;
-        this.searchRequest = searchRequest;
-    }
-
-    @Override
-    public NamedXContentRegistry getXContentRegistry() {
-        return rewriteContext.getXContentRegistry();
-    }
-
-    @Override
-    public long nowInMillis() {
-        return rewriteContext.nowInMillis();
-    }
-
-    @Override
-    public NamedWriteableRegistry getWriteableRegistry() {
-        return rewriteContext.getWriteableRegistry();
-    }
-
-    @Override
-    public QueryShardContext convertToShardContext() {
-        return rewriteContext.convertToShardContext();
-    }
-
-    @Override
-    public void registerAsyncAction(BiConsumer<Client, ActionListener<?>> asyncAction) {
-        rewriteContext.registerAsyncAction(asyncAction);
-    }
-
-    @Override
-    public boolean hasAsyncActions() {
-        return rewriteContext.hasAsyncActions();
-    }
-
-    @Override
-    public void executeAsyncActions(ActionListener listener) {
-        rewriteContext.executeAsyncActions(listener);
-    }
-
-    @Override
-    public boolean validate() {
-        return rewriteContext.validate();
+    public QueryCoordinatorContext(QueryRewriteContext parent, PipelinedRequest pipelinedRequest) {
+        super(parent.getXContentRegistry(), parent.getWriteableRegistry(), parent.client, parent.nowInMillis, parent.validate());
+        this.searchRequest = pipelinedRequest;
     }
 
     @Override

server/src/main/java/org/opensearch/index/query/QueryCoordinatorContext.java

@@ -33,40 +33,85 @@
 
 import org.opensearch.client.Client;
 import org.opensearch.common.annotation.PublicApi;
+import org.opensearch.common.util.concurrent.CountDown;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.common.io.stream.NamedWriteableRegistry;
 import org.opensearch.core.xcontent.NamedXContentRegistry;
 import org.opensearch.core.xcontent.XContentParser;
 
+import java.util.ArrayList;
+import java.util.List;
 import java.util.function.BiConsumer;
+import java.util.function.LongSupplier;
 
 /**
  * Context object used to rewrite {@link QueryBuilder} instances into simplified version.
  *
  * @opensearch.api
  */
 @PublicApi(since = "1.0.0")
-public interface QueryRewriteContext {
+public class QueryRewriteContext {
+    private final NamedXContentRegistry xContentRegistry;
+    private final NamedWriteableRegistry writeableRegistry;
+    protected final Client client;
+    protected final LongSupplier nowInMillis;
+    private final List<BiConsumer<Client, ActionListener<?>>> asyncActions = new ArrayList<>();
+    private final boolean validate;
+
+    public QueryRewriteContext(
+        NamedXContentRegistry xContentRegistry,
+        NamedWriteableRegistry writeableRegistry,
+        Client client,
+        LongSupplier nowInMillis
+    ) {
+        this(xContentRegistry, writeableRegistry, client, nowInMillis, false);
+    }
+
+    public QueryRewriteContext(
+        NamedXContentRegistry xContentRegistry,
+        NamedWriteableRegistry writeableRegistry,
+        Client client,
+        LongSupplier nowInMillis,
+        boolean validate
+    ) {
+
+        this.xContentRegistry = xContentRegistry;
+        this.writeableRegistry = writeableRegistry;
+        this.client = client;
+        this.nowInMillis = nowInMillis;
+        this.validate = validate;
+    }
+
     /**
      * The registry used to build new {@link XContentParser}s. Contains registered named parsers needed to parse the query.
      */
-    NamedXContentRegistry getXContentRegistry();
+    public NamedXContentRegistry getXContentRegistry() {
+        return xContentRegistry;
+    }
 
     /**
      * Returns the time in milliseconds that is shared across all resources involved. Even across shards and nodes.
      */
-    long nowInMillis();
+    public long nowInMillis() {
+        return nowInMillis.getAsLong();
+    }
 
-    NamedWriteableRegistry getWriteableRegistry();
+    public NamedWriteableRegistry getWriteableRegistry() {
+        return writeableRegistry;
+    }
 
     /**
-     * Returns an instance of {@link QueryShardContext} if available of null otherwise
+     * Returns an instance of {@link QueryShardContext} if available or null otherwise
      */
-    default QueryShardContext convertToShardContext() {
+    public QueryShardContext convertToShardContext() {
         return null;
     }
 
-    default QueryCoordinatorContext convertToCoordinatorContext() {
+    /**
+     * Returns an instance of {@link QueryCoordinatorContext} if available or null otherwise
+     * @return
+     */
+    public QueryCoordinatorContext convertToCoordinatorContext() {
         return null;
     }
 
@@ -75,18 +120,52 @@ default QueryCoordinatorContext convertToCoordinatorContext() {
      * This should be used if a rewriteabel needs to fetch some external resources in order to be executed ie. a document
      * from an index.
      */
-    void registerAsyncAction(BiConsumer<Client, ActionListener<?>> asyncAction);
+    public void registerAsyncAction(BiConsumer<Client, ActionListener<?>> asyncAction) {
+        asyncActions.add(asyncAction);
+    }
 
     /**
      * Returns <code>true</code> if there are any registered async actions.
      */
-    boolean hasAsyncActions();
+    public boolean hasAsyncActions() {
+        return asyncActions.isEmpty() == false;
+    }
 
     /**
      * Executes all registered async actions and notifies the listener once it's done. The value that is passed to the listener is always
      * <code>null</code>. The list of registered actions is cleared once this method returns.
      */
-    void executeAsyncActions(ActionListener listener);
+    public void executeAsyncActions(ActionListener listener) {
+        if (asyncActions.isEmpty()) {
+            listener.onResponse(null);
+            return;
+        }
 
-    boolean validate();
+        CountDown countDown = new CountDown(asyncActions.size());
+        ActionListener<?> internalListener = new ActionListener() {
+            @Override
+            public void onResponse(Object o) {
+                if (countDown.countDown()) {
+                    listener.onResponse(null);
+                }
+            }
+
+            @Override
+            public void onFailure(Exception e) {
+                if (countDown.fastForward()) {
+                    listener.onFailure(e);
+                }
+            }
+        };
+        // make a copy to prevent concurrent modification exception
+        List<BiConsumer<Client, ActionListener<?>>> biConsumers = new ArrayList<>(asyncActions);
+        asyncActions.clear();
+        for (BiConsumer<Client, ActionListener<?>> action : biConsumers) {
+            action.accept(client, internalListener);
+        }
+    }
+
+    public boolean validate() {
+        return validate;
+    }
 }

server/src/main/java/org/opensearch/index/query/QueryRewriteContext.java

@@ -101,7 +101,7 @@
  * @opensearch.api
  */
 @PublicApi(since = "1.0.0")
-public class QueryShardContext extends BaseQueryRewriteContext {
+public class QueryShardContext extends QueryRewriteContext {
 
     private final ScriptService scriptService;
     private final IndexSettings indexSettings;