bharath-techie · eirsep · Jun 22, 2022
diff --git a/server/src/main/java/org/opensearch/action/ActionModule.java b/server/src/main/java/org/opensearch/action/ActionModule.java
@@ -163,7 +163,9 @@
 import org.opensearch.action.admin.indices.rollover.RolloverAction;
 import org.opensearch.action.admin.indices.rollover.TransportRolloverAction;
 import org.opensearch.action.admin.indices.segments.IndicesSegmentsAction;
+import org.opensearch.action.admin.indices.segments.PitSegmentsAction;
 import org.opensearch.action.admin.indices.segments.TransportIndicesSegmentsAction;
+import org.opensearch.action.admin.indices.segments.TransportPitSegmentsAction;
 import org.opensearch.action.admin.indices.settings.get.GetSettingsAction;
 import org.opensearch.action.admin.indices.settings.get.TransportGetSettingsAction;
 import org.opensearch.action.admin.indices.settings.put.TransportUpdateSettingsAction;
@@ -364,6 +366,7 @@
 import org.opensearch.rest.action.cat.RestMasterAction;
 import org.opensearch.rest.action.cat.RestNodeAttrsAction;
 import org.opensearch.rest.action.cat.RestNodesAction;
+import org.opensearch.rest.action.cat.RestPitSegmentsAction;
 import org.opensearch.rest.action.cat.RestPluginsAction;
 import org.opensearch.rest.action.cat.RestRepositoriesAction;
 import org.opensearch.rest.action.cat.RestSegmentsAction;
@@ -644,6 +647,7 @@ public <Request extends ActionRequest, Response extends ActionResponse> void reg
         actions.register(ImportDanglingIndexAction.INSTANCE, TransportImportDanglingIndexAction.class);
         actions.register(DeleteDanglingIndexAction.INSTANCE, TransportDeleteDanglingIndexAction.class);
         actions.register(FindDanglingIndexAction.INSTANCE, TransportFindDanglingIndexAction.class);
+        actions.register(PitSegmentsAction.INSTANCE, TransportPitSegmentsAction.class);
         actions.register(CreatePitAction.INSTANCE, TransportCreatePitAction.class);
         actions.register(DeletePitAction.INSTANCE, TransportDeletePitAction.class);
         actions.register(GetAllPitsAction.INSTANCE, TransportGetAllPitsAction.class);
@@ -819,6 +823,7 @@ public void initRestHandlers(Supplier<DiscoveryNodes> nodesInCluster) {
         registerHandler.accept(new RestRepositoriesAction());
         registerHandler.accept(new RestSnapshotAction());
         registerHandler.accept(new RestTemplatesAction());
+        registerHandler.accept(new RestPitSegmentsAction());
 
         // Point in time API
         registerHandler.accept(new RestCreatePitAction());

diff --git a/server/src/main/java/org/opensearch/action/admin/indices/segments/PitSegmentsAction.java b/server/src/main/java/org/opensearch/action/admin/indices/segments/PitSegmentsAction.java
@@ -0,0 +1,21 @@
+/*
+ * 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.action.admin.indices.segments;
+
+import org.opensearch.action.ActionType;
+
+public class PitSegmentsAction extends ActionType<IndicesSegmentResponse> {
+
+    public static final PitSegmentsAction INSTANCE = new PitSegmentsAction();
+    public static final String NAME = "indices:monitor/point_in_time/segments";
+
+    private PitSegmentsAction() {
+        super(NAME, IndicesSegmentResponse::new);
+    }
+}
diff --git a/server/src/main/java/org/opensearch/action/admin/indices/segments/PitSegmentsRequest.java b/server/src/main/java/org/opensearch/action/admin/indices/segments/PitSegmentsRequest.java
@@ -0,0 +1,76 @@
+/*
+ * 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.action.admin.indices.segments;
+
+import org.opensearch.action.support.broadcast.BroadcastRequest;
+import org.opensearch.common.Strings;
+import org.opensearch.common.io.stream.StreamInput;
+import org.opensearch.common.io.stream.StreamOutput;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+
+public class PitSegmentsRequest extends BroadcastRequest<PitSegmentsRequest> {
+
+    protected boolean verbose = false;
+    private Collection<String> pitIds;
+
+    public PitSegmentsRequest() {
+        this(Strings.EMPTY_ARRAY);
+    }
+
+    public PitSegmentsRequest(StreamInput in) throws IOException {
+        super(in);
+        pitIds = Arrays.asList(in.readStringArray());
+        verbose = in.readBoolean();
+    }
+
+    public PitSegmentsRequest(String... indices) {
+        super(indices);
+        pitIds = Collections.emptyList();
+    }
+
+    /**
+     * <code>true</code> if detailed information about each segment should be returned,
+     * <code>false</code> otherwise.
+     */
+    public boolean verbose() {
+        return verbose;
+    }
+
+    /**
+     * Sets the <code>verbose</code> option.
+     * @see #verbose()
+     */
+    public void verbose(boolean v) {
+        verbose = v;
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        super.writeTo(out);
+        if (pitIds == null) {
+            out.writeVInt(0);
+        } else {
+            out.writeStringArray(pitIds.toArray(new String[pitIds.size()]));
+        }
+        out.writeBoolean(verbose);
+
+    }
+
+    public Collection<String> getPitIds() {
+        return pitIds;
+    }
+
+    public void setPitIds(Collection<String> pitIds) {
+        this.pitIds = pitIds;
+    }
+}
diff --git a/...rc/main/java/org/opensearch/action/admin/indices/segments/TransportPitSegmentsAction.java b/...rc/main/java/org/opensearch/action/admin/indices/segments/TransportPitSegmentsAction.java
@@ -0,0 +1,233 @@
+/*
+ * 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.action.admin.indices.segments;
+
+import com.carrotsearch.hppc.cursors.ObjectCursor;
+import org.opensearch.action.ActionListener;
+import org.opensearch.action.search.GetAllPitNodesRequest;
+import org.opensearch.action.search.GetAllPitNodesResponse;
+import org.opensearch.action.search.GetAllPitsAction;
+import org.opensearch.action.search.ListPitInfo;
+import org.opensearch.action.search.SearchContextId;
+import org.opensearch.action.search.SearchContextIdForNode;
+import org.opensearch.action.support.ActionFilters;
+import org.opensearch.action.support.DefaultShardOperationFailedException;
+import org.opensearch.action.support.broadcast.node.TransportBroadcastByNodeAction;
+import org.opensearch.cluster.ClusterState;
+import org.opensearch.cluster.block.ClusterBlockException;
+import org.opensearch.cluster.block.ClusterBlockLevel;
+import org.opensearch.cluster.metadata.IndexNameExpressionResolver;
+import org.opensearch.cluster.node.DiscoveryNode;
+import org.opensearch.cluster.routing.AllocationId;
+import org.opensearch.cluster.routing.PlainShardsIterator;
+import org.opensearch.cluster.routing.RecoverySource;
+import org.opensearch.cluster.routing.ShardRouting;
+import org.opensearch.cluster.routing.ShardRoutingState;
+import org.opensearch.cluster.routing.ShardsIterator;
+import org.opensearch.cluster.routing.UnassignedInfo;
+import org.opensearch.cluster.service.ClusterService;
+import org.opensearch.common.Strings;
+import org.opensearch.common.inject.Inject;
+import org.opensearch.common.io.stream.NamedWriteableRegistry;
+import org.opensearch.common.io.stream.StreamInput;
+import org.opensearch.common.io.stream.StreamOutput;
+import org.opensearch.common.io.stream.Writeable;
+import org.opensearch.common.xcontent.XContentBuilder;
+import org.opensearch.index.shard.ShardId;
+import org.opensearch.indices.IndicesService;
+import org.opensearch.search.SearchService;
+import org.opensearch.search.internal.PitReaderContext;
+import org.opensearch.tasks.Task;
+import org.opensearch.threadpool.ThreadPool;
+import org.opensearch.transport.TransportException;
+import org.opensearch.transport.TransportResponseHandler;
+import org.opensearch.transport.TransportService;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+import static org.opensearch.action.search.SearchContextId.decode;
+
+public class TransportPitSegmentsAction extends TransportBroadcastByNodeAction<PitSegmentsRequest, IndicesSegmentResponse, ShardSegments> {
+
+    private final ClusterService clusterService;
+    private final IndicesService indicesService;
+    private final SearchService searchService;
+    private final NamedWriteableRegistry namedWriteableRegistry;
+    private final TransportService transportService;
+
+
+    @Inject
+    public TransportPitSegmentsAction(ClusterService clusterService, TransportService transportService,
+                                      IndicesService indicesService, ActionFilters actionFilters,
+                                      IndexNameExpressionResolver indexNameExpressionResolver,
+                                      SearchService searchService,
+                                      NamedWriteableRegistry namedWriteableRegistry) {
+        super(PitSegmentsAction.NAME, clusterService, transportService, actionFilters, indexNameExpressionResolver,
+            PitSegmentsRequest::new, ThreadPool.Names.MANAGEMENT);
+        this.clusterService = clusterService;
+        this.indicesService = indicesService;
+        this.searchService = searchService;
+        this.namedWriteableRegistry = namedWriteableRegistry;
+        this.transportService = transportService;
+
+    }
+
+    @Override
+    protected void doExecute(Task task, PitSegmentsRequest request, ActionListener<IndicesSegmentResponse> listener) {
+        //TODO : Refactor to use a util method to "get all PITs" once such a method is implemented.
+        if (request.getPitIds().isEmpty()) {
+            final List<DiscoveryNode> nodes = new LinkedList<>();
+            for (ObjectCursor<DiscoveryNode> cursor : clusterService.state().nodes().getDataNodes().values()) {
+                DiscoveryNode node = cursor.value;
+                nodes.add(node);
+            }
+            DiscoveryNode[] disNodesArr = new DiscoveryNode[nodes.size()];
+            nodes.toArray(disNodesArr);
+            transportService.sendRequest(transportService.getLocalNode(), GetAllPitsAction.NAME,
+                new GetAllPitNodesRequest(disNodesArr), new TransportResponseHandler<GetAllPitNodesResponse>() {
+                    @Override
+                    public void handleResponse(GetAllPitNodesResponse response) {
+                        request.setPitIds(response.getPITIDs().stream().map(ListPitInfo::getPitId).collect(Collectors.toList()));
+                        getDoExecute(task, request, listener);
+                    }
+
+                    @Override
+                    public void handleException(TransportException exp) {
+                        listener.onFailure(exp);
+                    }
+
+                    @Override
+                    public String executor() {
+                        return ThreadPool.Names.SAME;
+                    }
+
+                    @Override
+                    public GetAllPitNodesResponse read(StreamInput in) throws IOException {
+                        return new GetAllPitNodesResponse(in);
+                    }
+                });
+        } else {
+            getDoExecute(task, request, listener);
+        }
+    }
+
+    private void getDoExecute(Task task, PitSegmentsRequest request, ActionListener<IndicesSegmentResponse> listener) {
+        super.doExecute(task, request, listener);
+    }
+
+    @Override
+    protected ShardsIterator shards(ClusterState clusterState, PitSegmentsRequest request, String[] concreteIndices) {
+        final ArrayList<ShardRouting> iterators = new ArrayList<>();
+        for (String pitId : request.getPitIds()) {
+            SearchContextId searchContext = decode(namedWriteableRegistry, pitId);
+            for (Map.Entry<ShardId, SearchContextIdForNode> entry : searchContext.shards().entrySet()) {
+                final SearchContextIdForNode perNode = entry.getValue();
+                if (Strings.isEmpty(perNode.getClusterAlias())) {
+                    final ShardId shardId = entry.getKey();
+                    iterators.add(new PitAwareShardRouting(pitId, shardId, perNode.getNode(), null, true, ShardRoutingState.STARTED,
+                        null, null, null, -1L));
+                }
+            }
+        }
+        return new PlainShardsIterator(iterators);
+    }
+
+    @Override
+    protected ClusterBlockException checkGlobalBlock(ClusterState state, PitSegmentsRequest request) {
+        return state.blocks().globalBlockedException(ClusterBlockLevel.METADATA_READ);
+    }
+
+    @Override
+    protected ClusterBlockException checkRequestBlock(ClusterState state, PitSegmentsRequest countRequest, String[] concreteIndices) {
+        return state.blocks().indicesBlockedException(ClusterBlockLevel.METADATA_READ, concreteIndices);
+    }
+
+    @Override
+    protected Writeable.Reader<TransportBroadcastByNodeAction<PitSegmentsRequest, IndicesSegmentResponse, ShardSegments>.NodeRequest> getNodeRequestWriteable() {
+        return NodeRequest::new;
+    }
+
+    @Override
+    protected ShardSegments readShardResult(StreamInput in) throws IOException {
+        return new ShardSegments(in);
+    }
+
+    @Override
+    protected IndicesSegmentResponse newResponse(PitSegmentsRequest request, int totalShards, int successfulShards, int failedShards,
+                                                 List<ShardSegments> results, List<DefaultShardOperationFailedException> shardFailures,
+                                                 ClusterState clusterState) {
+        return new IndicesSegmentResponse(results.toArray(new ShardSegments[results.size()]), totalShards, successfulShards, failedShards,
+            shardFailures);
+    }
+
+    @Override
+    protected PitSegmentsRequest readRequestFrom(StreamInput in) throws IOException {
+        return new PitSegmentsRequest(in);
+    }
+
+    @Override
+    protected ShardSegments shardOperation(PitSegmentsRequest request, ShardRouting shardRouting) {
+        PitAwareShardRouting pitAwareShardRouting = (PitAwareShardRouting) shardRouting;
+        SearchContextIdForNode searchContextIdForNode = decode(namedWriteableRegistry,
+            pitAwareShardRouting.getPitId()).shards().get(shardRouting.shardId());
+        PitReaderContext pitReaderContext = searchService.getPitReaderContext(searchContextIdForNode.getSearchContextId());
+        return new ShardSegments(pitReaderContext.getShardRouting(), pitReaderContext.getSegments());
+    }
+
+    public class PitAwareShardRouting extends ShardRouting {
+
+        private final String pitId;
+
+        public PitAwareShardRouting(StreamInput in) throws IOException {
+            super(in);
+            this.pitId = in.readString();
+        }
+
+        public PitAwareShardRouting(
+            String pitId,
+            ShardId shardId,
+            String currentNodeId,
+            String relocatingNodeId,
+            boolean primary,
+            ShardRoutingState state,
+            RecoverySource recoverySource,
+            UnassignedInfo unassignedInfo,
+            AllocationId allocationId,
+            long expectedShardSize
+        ) {
+            super(shardId, currentNodeId, relocatingNodeId, primary, state, recoverySource, unassignedInfo,
+                allocationId, expectedShardSize);
+            this.pitId = pitId;
+        }
+
+        public String getPitId() {
+            return pitId;
+        }
+
+        @Override
+        public void writeTo(StreamOutput out) throws IOException {
+            super.writeTo(out);
+            out.writeString(pitId);
+        }
+
+        @Override
+        public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+            super.toXContent(builder, params);
+            builder.field("pitId", pitId);
+            return builder.endObject();
+        }
+    }
+
+
+}
diff --git a/server/src/main/java/org/opensearch/action/search/SearchContextId.java b/server/src/main/java/org/opensearch/action/search/SearchContextId.java
@@ -41,6 +41,7 @@
 import org.opensearch.common.io.stream.NamedWriteableRegistry;
 import org.opensearch.common.io.stream.StreamInput;
 import org.opensearch.common.io.stream.StreamOutput;
+import org.opensearch.common.io.stream.Writeable;
 import org.opensearch.index.shard.ShardId;
 import org.opensearch.search.SearchPhaseResult;
 import org.opensearch.search.SearchShardTarget;
@@ -62,7 +63,7 @@
  *
  * @opensearch.internal
  */
-public class SearchContextId {
+public class SearchContextId implements Writeable {
     private final Map<ShardId, SearchContextIdForNode> shards;
     private final Map<String, AliasFilter> aliasFilter;