More refactoring

Author: alchemist51

sandbox/plugins/composite-engine/src/main/java/org/opensearch/composite/CompositeDataFormatWriterPool.java

@@ -23,14 +23,15 @@
 import java.util.Set;
 import java.util.function.Supplier;
 
-public class CompositeWriterPool implements Iterable<CompositeWriter>, Closeable {
+public class CompositeDataFormatWriterPool implements Iterable<CompositeWriter>, Closeable {
 
     private final Set<CompositeWriter> writers;
     private final LockableConcurrentQueue<CompositeWriter> availableWriters;
     private final Supplier<CompositeWriter> writerSupplier;
     private volatile boolean closed;
 
-    public CompositeWriterPool(
+
+    public CompositeDataFormatWriterPool(
         Supplier<CompositeWriter> writerSupplier,
         Supplier<Queue<CompositeWriter>> queueSupplier,
         int concurrency
@@ -79,6 +80,7 @@ public void releaseAndUnlock(CompositeWriter state) {
         availableWriters.addAndUnlock(state);
     }
 
+
     /**
      * Lock and checkout all CompositeWriters from the pool for flush.
      *

sandbox/plugins/composite-engine/src/main/java/org/opensearch/composite/CompositeDocumentInput.java

@@ -16,9 +16,11 @@
 import org.opensearch.index.engine.dataformat.FieldTypeCapabilities;
 import org.opensearch.index.mapper.MappedFieldType;
 
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
@@ -36,13 +38,13 @@
  * @opensearch.experimental
  */
 @ExperimentalApi
-public class CompositeDocumentInput implements DocumentInput<Map<DataFormat, Object>> {
+public class CompositeDocumentInput implements DocumentInput<List<? extends DocumentInput<?>>> {
 
     private static final Logger logger = LogManager.getLogger(CompositeDocumentInput.class);
 
-    private final DocumentInput<?> primaryInput;
+    private final DocumentInput<?> primaryDocumentInput;
     private final DataFormat primaryFormat;
-    private final Map<DataFormat, DocumentInput<?>> secondaryInputs;
+    private final Map<DataFormat, DocumentInput<?>> secondaryDocumentInputs;
     private final Map<String, Set<DataFormat>> fieldTypeToFormats;
 
     /**
@@ -53,21 +55,23 @@ public class CompositeDocumentInput implements DocumentInput<Map<DataFormat, Obj
      * {@link DataFormat#supportedFields()}.
      *
      * @param primaryFormat the primary data format
-     * @param primaryInput the document input for the primary format
-     * @param secondaryInputs a map of secondary data formats to their corresponding document inputs
+     * @param primaryDocumentInput the document input for the primary format
+     * @param secondaryDocumentInputs a map of secondary data formats to their corresponding document inputs
      */
     public CompositeDocumentInput(
         DataFormat primaryFormat,
-        DocumentInput<?> primaryInput,
-        Map<DataFormat, DocumentInput<?>> secondaryInputs
+        DocumentInput<?> primaryDocumentInput,
+        Map<DataFormat, DocumentInput<?>> secondaryDocumentInputs
     ) {
         this.primaryFormat = Objects.requireNonNull(primaryFormat, "primaryFormat must not be null");
-        this.primaryInput = Objects.requireNonNull(primaryInput, "primaryInput must not be null");
-        this.secondaryInputs = Map.copyOf(Objects.requireNonNull(secondaryInputs, "secondaryInputs must not be null"));
+        this.primaryDocumentInput = Objects.requireNonNull(primaryDocumentInput, "primaryDocumentInput must not be null");
+        this.secondaryDocumentInputs = Map.copyOf(
+            Objects.requireNonNull(secondaryDocumentInputs, "secondaryDocumentInputs must not be null")
+        );
 
         Map<String, Set<DataFormat>> routing = new HashMap<>();
         addFormatToRouting(routing, primaryFormat);
-        for (DataFormat format : this.secondaryInputs.keySet()) {
+        for (DataFormat format : this.secondaryDocumentInputs.keySet()) {
             addFormatToRouting(routing, format);
         }
         this.fieldTypeToFormats = Collections.unmodifiableMap(routing);
@@ -88,9 +92,9 @@ public void addField(MappedFieldType fieldType, Object value) {
             return;
         }
         if (formats.contains(primaryFormat)) {
-            primaryInput.addField(fieldType, value);
+            primaryDocumentInput.addField(fieldType, value);
         }
-        for (Map.Entry<DataFormat, DocumentInput<?>> entry : secondaryInputs.entrySet()) {
+        for (Map.Entry<DataFormat, DocumentInput<?>> entry : secondaryDocumentInputs.entrySet()) {
             if (formats.contains(entry.getKey())) {
                 entry.getValue().addField(fieldType, value);
             }
@@ -99,30 +103,25 @@ public void addField(MappedFieldType fieldType, Object value) {
 
     @Override
     public void setRowId(String rowIdFieldName, long rowId) {
-        primaryInput.setRowId(rowIdFieldName, rowId);
-        for (DocumentInput<?> input : secondaryInputs.values()) {
+        primaryDocumentInput.setRowId(rowIdFieldName, rowId);
+        for (DocumentInput<?> input : secondaryDocumentInputs.values()) {
             input.setRowId(rowIdFieldName, rowId);
         }
     }
 
     @Override
-    public Map<DataFormat, Object> getFinalInput() {
-        Map<DataFormat, Object> result = new HashMap<>();
-        result.put(primaryFormat, primaryInput.getFinalInput());
-        for (Map.Entry<DataFormat, DocumentInput<?>> entry : secondaryInputs.entrySet()) {
-            result.put(entry.getKey(), entry.getValue().getFinalInput());
-        }
-        return Collections.unmodifiableMap(result);
+    public List<? extends DocumentInput<?>> getFinalInput() {
+        return null;
     }
 
     @Override
     public void close() {
         try {
-            primaryInput.close();
+            primaryDocumentInput.close();
         } catch (Exception e) {
             logger.warn("Failed to close primary DocumentInput", e);
         }
-        for (DocumentInput<?> input : secondaryInputs.values()) {
+        for (DocumentInput<?> input : secondaryDocumentInputs.values()) {
             try {
                 input.close();
             } catch (Exception e) {
@@ -137,7 +136,7 @@ public void close() {
      * @return the primary document input
      */
     public DocumentInput<?> getPrimaryInput() {
-        return primaryInput;
+        return primaryDocumentInput;
     }
 
     /**
@@ -155,6 +154,6 @@ public DataFormat getPrimaryFormat() {
      * @return the secondary inputs
      */
     public Map<DataFormat, DocumentInput<?>> getSecondaryInputs() {
-        return secondaryInputs;
+        return secondaryDocumentInputs;
     }
 }

sandbox/plugins/composite-engine/src/main/java/org/opensearch/composite/CompositeIndexingExecutionEngine.java

@@ -28,11 +28,12 @@
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collection;
-import java.util.HashMap;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.atomic.AtomicLong;
 
 /**
  * A composite {@link IndexingExecutionEngine} that orchestrates indexing across
@@ -52,10 +53,11 @@ public class CompositeIndexingExecutionEngine implements IndexingExecutionEngine
 
     private static final Logger logger = LogManager.getLogger(CompositeIndexingExecutionEngine.class);
 
-    private final IndexingExecutionEngine<?, ?> primaryEngine;
-    private final List<IndexingExecutionEngine<?, ?>> secondaryEngines;
-    private final List<IndexingExecutionEngine<?, ?>> allEngines;
+    private final IndexingExecutionEngine<? extends DataFormat, ? extends DocumentInput<?>> primaryEngine;
+    private final List<IndexingExecutionEngine<? extends DataFormat, ? extends DocumentInput<?>>> secondaryEngines;
     private final CompositeDataFormat compositeDataFormat;
+    private final CompositeDataFormatWriterPool dataFormatWriterPool;
+    private final AtomicLong writerGenerationCounter;
 
     /**
      * Constructs a CompositeIndexingExecutionEngine by reading index settings to
@@ -110,16 +112,18 @@ public CompositeIndexingExecutionEngine(
         }
         this.secondaryEngines = List.copyOf(secondaries);
 
-        List<IndexingExecutionEngine<?, ?>> all = new ArrayList<>();
-        all.add(this.primaryEngine);
-        all.addAll(this.secondaryEngines);
-        this.allEngines = List.copyOf(all);
-
         List<DataFormat> allFormats = new ArrayList<>();
-        for (IndexingExecutionEngine<?, ?> engine : this.allEngines) {
+        allFormats.add(primaryEngine.getDataFormat());
+        for (IndexingExecutionEngine<?, ?> engine : secondaryEngines) {
             allFormats.add(engine.getDataFormat());
         }
         this.compositeDataFormat = new CompositeDataFormat(allFormats);
+        this.writerGenerationCounter = new AtomicLong(0);
+        this.dataFormatWriterPool = new CompositeDataFormatWriterPool(
+            () -> new CompositeWriter(this, primaryEngine.getDataFormat(), writerGenerationCounter.getAndIncrement()),
+            ConcurrentLinkedQueue::new,
+            Runtime.getRuntime().availableProcessors()
+        );
     }
 
     /**
@@ -160,12 +164,7 @@ static void validateFormatsRegistered(
 
     @Override
     public Writer<CompositeDocumentInput> createWriter(long writerGeneration) {
-        Map<DataFormat, Writer<?>> writerMap = new LinkedHashMap<>();
-        for (IndexingExecutionEngine<?, ?> engine : allEngines) {
-            Writer<?> writer = engine.createWriter(writerGeneration);
-            writerMap.put(engine.getDataFormat(), writer);
-        }
-        return new CompositeWriter(writerMap);
+        return new CompositeWriter(this, primaryEngine.getDataFormat(), writerGeneration);
     }
 
     @Override
@@ -176,7 +175,9 @@ public Merger getMerger() {
     @Override
     public RefreshResult refresh(RefreshInput refreshInput) throws IOException {
         List<Segment> allSegments = new ArrayList<>();
-        for (IndexingExecutionEngine<?, ?> engine : allEngines) {
+        RefreshResult primaryResult = primaryEngine.refresh(refreshInput);
+        allSegments.addAll(primaryResult.refreshedSegments());
+        for (IndexingExecutionEngine<?, ?> engine : secondaryEngines) {
             RefreshResult result = engine.refresh(refreshInput);
             allSegments.addAll(result.refreshedSegments());
         }
@@ -191,27 +192,55 @@ public CompositeDataFormat getDataFormat() {
 
     @Override
     public long getNativeBytesUsed() {
-        long total = 0;
-        for (IndexingExecutionEngine<?, ?> engine : allEngines) {
+        long total = primaryEngine.getNativeBytesUsed();
+        for (IndexingExecutionEngine<?, ?> engine : secondaryEngines) {
             total += engine.getNativeBytesUsed();
         }
         return total;
     }
 
     @Override
     public void deleteFiles(Map<String, Collection<String>> filesToDelete) throws IOException {
-        for (IndexingExecutionEngine<?, ?> engine : allEngines) {
+        primaryEngine.deleteFiles(filesToDelete);
+        for (IndexingExecutionEngine<?, ?> engine : secondaryEngines) {
             engine.deleteFiles(filesToDelete);
         }
     }
 
     @Override
     public CompositeDocumentInput newDocumentInput() {
         DocumentInput<?> primaryInput = primaryEngine.newDocumentInput();
-        Map<DataFormat, DocumentInput<?>> secondaryInputMap = new HashMap<>();
+        Map<DataFormat, DocumentInput<?>> secondaryInputMap = new LinkedHashMap<>();
         for (IndexingExecutionEngine<?, ?> engine : secondaryEngines) {
             secondaryInputMap.put(engine.getDataFormat(), engine.newDocumentInput());
         }
         return new CompositeDocumentInput(primaryEngine.getDataFormat(), primaryInput, secondaryInputMap);
     }
+
+    /**
+     * Returns the primary delegate engine.
+     *
+     * @return the primary engine
+     */
+    public IndexingExecutionEngine<?, ?> getPrimaryDelegate() {
+        return primaryEngine;
+    }
+
+    /**
+     * Returns the secondary delegate engines.
+     *
+     * @return the secondary engines
+     */
+    public List<IndexingExecutionEngine<? extends DataFormat, ? extends DocumentInput<?>>> getSecondaryDelegates() {
+        return secondaryEngines;
+    }
+
+    /**
+     * Returns the composite writer pool for this engine.
+     *
+     * @return the writer pool
+     */
+    public CompositeDataFormatWriterPool getDataFormatWriterPool() {
+        return dataFormatWriterPool;
+    }
 }