Adding support for context aware segments

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

Author: RS146BIJAY

CHANGELOG.md

@@ -23,6 +23,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Implement GRPC Search params `Highlight`and `Sort` ([#19868](https://github.com/opensearch-project/OpenSearch/pull/19868))
 - Implement GRPC ConstantScoreQuery, FuzzyQuery, MatchBoolPrefixQuery, MatchPhrasePrefix, PrefixQuery, MatchQuery ([#19854](https://github.com/opensearch-project/OpenSearch/pull/19854))
 - Add async periodic flush task support for pull-based ingestion ([#19878](https://github.com/opensearch-project/OpenSearch/pull/19878))
+- Add support for context aware segments ([#19098](https://github.com/opensearch-project/OpenSearch/pull/19098))
 
 ### Changed
 - Faster `terms` query creation for `keyword` field with index and docValues enabled ([#19350](https://github.com/opensearch-project/OpenSearch/pull/19350))

modules/store-subdirectory/src/main/java/org/opensearch/plugin/store/subdirectory/SubdirectoryAwareStore.java

@@ -26,6 +26,7 @@
 import org.opensearch.index.shard.ShardPath;
 import org.opensearch.index.store.Store;
 import org.opensearch.index.store.StoreFileMetadata;
+import org.opensearch.plugins.IndexStorePlugin;
 
 import java.io.IOException;
 import java.nio.file.FileVisitResult;
@@ -79,6 +80,37 @@ public SubdirectoryAwareStore(
         super(shardId, indexSettings, new SubdirectoryAwareDirectory(directory, shardPath), shardLock, onClose, shardPath);
     }
 
+    /**
+     * Constructor for SubdirectoryAwareStore.
+     *
+     * @param shardId the shard ID
+     * @param indexSettings the index settings
+     * @param directory the directory to use for the store
+     * @param shardLock the shard lock
+     * @param onClose the on close callback
+     * @param shardPath the shard path
+     * @param directoryFactory the directory factory
+     */
+    public SubdirectoryAwareStore(
+        ShardId shardId,
+        IndexSettings indexSettings,
+        Directory directory,
+        ShardLock shardLock,
+        OnClose onClose,
+        ShardPath shardPath,
+        IndexStorePlugin.DirectoryFactory directoryFactory
+    ) {
+        super(
+            shardId,
+            indexSettings,
+            new SubdirectoryAwareDirectory(directory, shardPath),
+            shardLock,
+            onClose,
+            shardPath,
+            directoryFactory
+        );
+    }
+
     @Override
     public MetadataSnapshot getMetadata(IndexCommit commit) throws IOException {
         long totalNumDocs = 0;

modules/store-subdirectory/src/main/java/org/opensearch/plugin/store/subdirectory/SubdirectoryStorePlugin.java

@@ -81,5 +81,31 @@ public Store newStore(
         ) {
             return new SubdirectoryAwareStore(shardId, indexSettings, directory, shardLock, onClose, shardPath);
         }
+
+        /**
+         * Creates a new {@link SubdirectoryAwareStore} instance.
+         *
+         * @param shardId the shard identifier
+         * @param indexSettings the index settings
+         * @param directory the underlying Lucene directory
+         * @param shardLock the shard lock
+         * @param onClose callback to execute when the store is closed
+         * @param shardPath the path information for the shard
+         * @param directoryFactory the directory factory to create child level directory.
+         *                         Used for Context Aware Segments enabled indices.
+         * @return a new SubdirectoryAwareStore instance
+         */
+        @Override
+        public Store newStore(
+            ShardId shardId,
+            IndexSettings indexSettings,
+            Directory directory,
+            ShardLock shardLock,
+            Store.OnClose onClose,
+            ShardPath shardPath,
+            DirectoryFactory directoryFactory
+        ) {
+            return new SubdirectoryAwareStore(shardId, indexSettings, directory, shardLock, onClose, shardPath, directoryFactory);
+        }
     }
 }

modules/store-subdirectory/src/test/java/org/opensearch/plugin/store/subdirectory/SubdirectoryStorePluginTests.java

@@ -20,6 +20,7 @@
 import org.opensearch.common.util.io.IOUtils;
 import org.opensearch.core.index.shard.ShardId;
 import org.opensearch.index.shard.ShardPath;
+import org.opensearch.index.store.FsDirectoryFactory;
 import org.opensearch.index.store.Store;
 import org.opensearch.index.store.StoreStats;
 import org.opensearch.plugins.IndexStorePlugin;
@@ -67,7 +68,8 @@ public void testStats() throws IOException {
             SubdirectoryStorePluginTests.newFSDirectory(path.resolve("index")),
             new DummyShardLock(shardId),
             Store.OnClose.EMPTY,
-            new ShardPath(false, path, path, shardId)
+            new ShardPath(false, path, path, shardId),
+            new FsDirectoryFactory()
         );
 
         long initialStoreSize = 0;

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/OpenSearchServerException.java

@@ -24,6 +24,7 @@
 import static org.opensearch.Version.V_2_7_0;
 import static org.opensearch.Version.V_3_0_0;
 import static org.opensearch.Version.V_3_2_0;
+import static org.opensearch.Version.V_3_3_0;
 
 /**
  * Utility class to register server exceptions
@@ -1241,5 +1242,13 @@ public static void registerExceptions() {
                 V_3_2_0
             )
         );
+        registerExceptionHandle(
+            new OpenSearchExceptionHandle(
+                org.opensearch.index.engine.LookupMapLockAcquisitionException.class,
+                org.opensearch.index.engine.LookupMapLockAcquisitionException::new,
+                CUSTOM_ELASTICSEARCH_EXCEPTIONS_BASE_ID + 2,
+                V_3_3_0
+            )
+        );
     }
 }

server/src/main/java/org/opensearch/action/bulk/TransportShardBulkAction.java

@@ -87,6 +87,7 @@
 import org.opensearch.index.IndexingPressureService;
 import org.opensearch.index.SegmentReplicationPressureService;
 import org.opensearch.index.engine.Engine;
+import org.opensearch.index.engine.LookupMapLockAcquisitionException;
 import org.opensearch.index.engine.VersionConflictEngineException;
 import org.opensearch.index.get.GetResult;
 import org.opensearch.index.mapper.MapperException;
@@ -727,7 +728,15 @@ && isConflictException(executionResult.getFailure().getCause())
             && context.getRetryCounter() < ((UpdateRequest) docWriteRequest).retryOnConflict()) {
             context.resetForExecutionForRetry();
             return;
-        }
+        } else if (isFailed
+            && context.getPrimary() != null
+            && context.getPrimary().indexSettings() != null
+            && context.getPrimary().indexSettings().isContextAwareEnabled()
+            && isLookupMapLockAcquisitionException(executionResult.getFailure().getCause())
+            && context.getRetryCounter() < context.getPrimary().indexSettings().getMaxRetryOnLookupMapAcquisitionException()) {
+                context.resetForExecutionForRetry();
+                return;
+            }
         final BulkItemResponse response;
         if (isUpdate) {
             response = processUpdateResponse((UpdateRequest) docWriteRequest, context.getConcreteIndex(), executionResult, updateResult);
@@ -756,6 +765,10 @@ private static boolean isConflictException(final Exception e) {
         return ExceptionsHelper.unwrapCause(e) instanceof VersionConflictEngineException;
     }
 
+    private static boolean isLookupMapLockAcquisitionException(final Exception e) {
+        return ExceptionsHelper.unwrapCause(e) instanceof LookupMapLockAcquisitionException;
+    }
+
     /**
      * Creates a new bulk item result from the given requests and result of performing the update operation on the shard.
      */

server/src/main/java/org/opensearch/common/lucene/Lucene.java

@@ -91,6 +91,7 @@
 import org.opensearch.core.common.io.stream.StreamOutput;
 import org.opensearch.index.analysis.AnalyzerScope;
 import org.opensearch.index.analysis.NamedAnalyzer;
+import org.opensearch.index.codec.CriteriaBasedCodec;
 import org.opensearch.index.fielddata.IndexFieldData;
 import org.opensearch.index.fielddata.plain.NonPruningSortedSetOrdinalsIndexFieldData.NonPruningSortField;
 import org.opensearch.search.sort.SortedWiderNumericSortField;
@@ -939,10 +940,23 @@ public LeafReader wrap(LeafReader leaf) {
                     // Two scenarios that we have hard-deletes: (1) from old segments where soft-deletes was disabled,
                     // (2) when IndexWriter hits non-aborted exceptions. These two cases, IW flushes SegmentInfos
                     // before exposing the hard-deletes, thus we can use the hard-delete count of SegmentInfos.
-                    final int numDocs = segmentReader.maxDoc() - segmentReader.getSegmentInfo().getDelCount();
+
+                    // With CAS enabled segments, hard deletes can also be present, so correcting numDocs.
+                    // We are using attribute value here to identify whether segment has CAS enabled or not.
+                    int numDocs;
+                    if (isContextAwareEnabled(segmentReader)) {
+                        numDocs = popCount(hardLiveDocs);
+                    } else {
+                        numDocs = segmentReader.maxDoc() - segmentReader.getSegmentInfo().getDelCount();
+                    }
+
                     assert numDocs == popCount(hardLiveDocs) : numDocs + " != " + popCount(hardLiveDocs);
                     return new LeafReaderWithLiveDocs(segmentReader, hardLiveDocs, numDocs);
                 }
+
+                private boolean isContextAwareEnabled(SegmentReader reader) {
+                    return reader.getSegmentInfo().info.getAttribute(CriteriaBasedCodec.BUCKET_NAME) != null;
+                }
             });
         }
 

server/src/main/java/org/opensearch/common/settings/IndexScopedSettings.java

@@ -176,6 +176,8 @@ public final class IndexScopedSettings extends AbstractScopedSettings {
                 ShardsLimitAllocationDecider.INDEX_TOTAL_REMOTE_CAPABLE_PRIMARY_SHARDS_PER_NODE_SETTING,
                 IndexSettings.INDEX_GC_DELETES_SETTING,
                 IndexSettings.INDEX_SOFT_DELETES_SETTING,
+                IndexSettings.INDEX_CONTEXT_AWARE_ENABLED_SETTING,
+                IndexSettings.INDEX_MAX_RETRY_ON_LOOKUP_MAP_LOCK_ACQUISITION_EXCEPTION,
                 IndexSettings.INDEX_SOFT_DELETES_RETENTION_OPERATIONS_SETTING,
                 IndexSettings.INDEX_SOFT_DELETES_RETENTION_LEASE_PERIOD_SETTING,
                 IndicesRequestCache.INDEX_CACHE_REQUEST_ENABLED_SETTING,