Adding support for context aware segments

Signed-off-by: RS146BIJAY <[email protected]>

Author: RS146BIJAY

CHANGELOG-3.0.md

@@ -25,6 +25,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Add filter function for AbstractQueryBuilder, BoolQueryBuilder, ConstantScoreQueryBuilder([#17409](https://github.com/opensearch-project/OpenSearch/pull/17409))
 - [Star Tree] [Search] Resolving keyword & numeric bucket aggregation with metric aggregation using star-tree ([#17165](https://github.com/opensearch-project/OpenSearch/pull/17165))
 - Added error handling support for the pull-based ingestion ([#17427](https://github.com/opensearch-project/OpenSearch/pull/17427))
+- Add support for context aware segments ([#19098](https://github.com/opensearch-project/OpenSearch/pull/19098)
 
 
 ### Dependencies

plugins/store-smb/src/main/java/org/opensearch/index/store/smbmmapfs/SmbMmapFsDirectoryFactory.java

@@ -47,7 +47,7 @@
 public final class SmbMmapFsDirectoryFactory extends FsDirectoryFactory {
 
     @Override
-    protected Directory newFSDirectory(Path location, LockFactory lockFactory, IndexSettings indexSettings) throws IOException {
+    public Directory newFSDirectory(Path location, LockFactory lockFactory, IndexSettings indexSettings) throws IOException {
         return new SmbDirectoryWrapper(
             setPreload(
                 new MMapDirectory(location, lockFactory),

plugins/store-smb/src/main/java/org/opensearch/index/store/smbniofs/SmbNIOFsDirectoryFactory.java

@@ -24,7 +24,7 @@
 public final class SmbNIOFsDirectoryFactory extends FsDirectoryFactory {
 
     @Override
-    protected Directory newFSDirectory(Path location, LockFactory lockFactory, IndexSettings indexSettings) throws IOException {
+    public Directory newFSDirectory(Path location, LockFactory lockFactory, IndexSettings indexSettings) throws IOException {
         return new SmbDirectoryWrapper(new NIOFSDirectory(location, lockFactory));
     }
 }

server/src/main/java/org/opensearch/index/CriteriaBasedCompositeDirectory.java

@@ -0,0 +1,31 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.index;
+
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.store.FilterDirectory;
+
+import java.io.IOException;
+import java.util.Arrays;
+
+public class CriteriaBasedCompositeDirectory extends FilterDirectory {
+    /**
+     * Sole constructor, typically called from sub-classes.
+     *
+     * @param in
+     */
+    protected CriteriaBasedCompositeDirectory(Directory in) {
+        super(in);
+    }
+
+    @Override
+    public String[] listAll() throws IOException {
+        return Arrays.stream(super.listAll()).filter(fileName -> !fileName.startsWith("temp_")).distinct().toArray(String[]::new);
+    }
+}

server/src/main/java/org/opensearch/index/CriteriaBasedMergePolicy.java

@@ -0,0 +1,60 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.index;
+
+import org.apache.lucene.index.MergeTrigger;
+import org.apache.lucene.index.SegmentCommitInfo;
+import org.apache.lucene.index.SegmentInfos;
+import org.apache.lucene.index.TieredMergePolicy;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+public class CriteriaBasedMergePolicy extends TieredMergePolicy {
+    @Override
+    public MergeSpecification findMerges(
+        MergeTrigger mergeTrigger, SegmentInfos infos, MergeContext mergeContext) throws IOException {
+        final Set<SegmentCommitInfo> merging = mergeContext.getMergingSegments();
+        MergeSpecification spec = null;
+        final Map<String, List<SegmentCommitInfo>> commitInfos = new HashMap<>();
+        for (SegmentCommitInfo si : infos) {
+            if (merging.contains(si)) {
+                continue;
+            }
+
+            final String dwptGroupNumber = si.info.getAttribute("criteria");
+            commitInfos.computeIfAbsent(dwptGroupNumber, k -> new ArrayList<>()).add(si);
+        }
+
+        for (String dwptGroupNumber : commitInfos.keySet()) {
+            if (commitInfos.get(dwptGroupNumber).size() > 1) {
+                final SegmentInfos newSIS = new SegmentInfos(infos.getIndexCreatedVersionMajor());
+                for (SegmentCommitInfo info : commitInfos.get(dwptGroupNumber)) {
+                    newSIS.add(info);
+                }
+
+                final MergeSpecification tieredMergePolicySpec =
+                    super.findMerges(mergeTrigger, newSIS, mergeContext);
+                if (tieredMergePolicySpec != null) {
+                    if (spec == null) {
+                        spec = new MergeSpecification();
+                    }
+
+                    spec.merges.addAll(tieredMergePolicySpec.merges);
+                }
+            }
+        }
+
+        return spec;
+    }
+}

server/src/main/java/org/opensearch/index/IndexService.java

@@ -637,7 +637,7 @@ protected void closeInternal() {
                         // Do nothing for shard lock on remote store
                     }
                 };
-                remoteStore = new Store(shardId, this.indexSettings, remoteDirectory, remoteStoreLock, Store.OnClose.EMPTY, path);
+                remoteStore = new Store(shardId, this.indexSettings, remoteDirectory, remoteStoreLock, Store.OnClose.EMPTY, path, directoryFactory);
             } else {
                 // Disallow shards with remote store based settings to be created on non-remote store enabled nodes
                 // Even though we have `RemoteStoreMigrationAllocationDecider` in place to prevent something like this from happening at the
@@ -658,15 +658,17 @@ protected void closeInternal() {
                 Directory localDirectory = directoryFactory.newDirectory(this.indexSettings, path);
                 directory = new CompositeDirectory(localDirectory, remoteDirectory, fileCache);
             } else {
-                directory = directoryFactory.newDirectory(this.indexSettings, path);
+                Directory localDirectory = directoryFactory.newDirectory(this.indexSettings, path);
+                directory = new CriteriaBasedCompositeDirectory(localDirectory);
             }
             store = new Store(
                 shardId,
                 this.indexSettings,
                 directory,
                 lock,
                 new StoreCloseListener(shardId, () -> eventListener.onStoreClosed(shardId)),
-                path
+                path,
+                directoryFactory
             );
             eventListener.onStoreCreated(shardId);
             indexShard = new IndexShard(

server/src/main/java/org/opensearch/index/codec/CriteriaBasedCodec.java

@@ -0,0 +1,49 @@
+/*Add commentMore actions
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.index.codec;
+
+import org.apache.lucene.codecs.Codec;
+import org.apache.lucene.codecs.FilterCodec;
+import org.apache.lucene.codecs.SegmentInfoFormat;
+import org.apache.lucene.codecs.lucene101.Lucene101Codec;
+import org.apache.lucene.index.SegmentInfo;
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.store.IOContext;
+
+import java.io.IOException;
+
+public class CriteriaBasedCodec extends FilterCodec {
+
+    private final String criteria;
+    public CriteriaBasedCodec() {
+        super("CriteriaBasedCodec", new Lucene101Codec());
+        criteria = null;
+    }
+
+    public CriteriaBasedCodec(Codec delegate, String criteria) {
+        super("CriteriaBasedCodec", delegate);
+        this.criteria = criteria;
+    }
+
+    @Override
+    public SegmentInfoFormat segmentInfoFormat() {
+        return new SegmentInfoFormat() {
+            @Override
+            public SegmentInfo read(Directory directory, String segmentName, byte[] segmentID, IOContext context) throws IOException {
+                return delegate.segmentInfoFormat().read(directory, segmentName, segmentID, context);
+            }
+
+            @Override
+            public void write(Directory directory, SegmentInfo info, IOContext ioContext) throws IOException {
+                info.putAttribute("criteria", criteria);
+                delegate.segmentInfoFormat().write(directory, info, ioContext);
+            }
+        };
+    }
+}