Skip to content

Commit 7402d31

Browse files
varunbharadwajvarunbharadwaj
committed
Refactor and support shard level acks
Signed-off-by: Varun Bharadwaj <[email protected]>
1 parent dcccf22 commit 7402d31

26 files changed

+926
-437
lines changed

‎plugins/ingestion-kafka/src/internalClusterTest/java/org/opensearch/plugin/kafka/RemoteStoreKafkaIT.java

+5-3
Original file line numberDiff line numberDiff line change
@@ -172,9 +172,10 @@ public void testPauseAndResumeIngestion() throws Exception {
172172
// pause ingestion
173173
PauseIngestionResponse pauseResponse = pauseIngestion(indexName);
174174
assertTrue(pauseResponse.isAcknowledged());
175+
assertTrue(pauseResponse.isShardsAcknowledged());
175176
waitForState(() -> {
176177
GetIngestionStateResponse ingestionState = getIngestionState(indexName);
177-
return Arrays.stream(ingestionState.getShardStates()).allMatch(state -> state.getPollerState().equalsIgnoreCase("paused"));
178+
return Arrays.stream(ingestionState.getShardStates()).allMatch(state -> state.isPollerPaused() && state.getPollerState().equalsIgnoreCase("paused"));
178179
});
179180

180181
// verify ingestion state is persisted
@@ -191,17 +192,18 @@ public void testPauseAndResumeIngestion() throws Exception {
191192
assertEquals(2, getSearchableDocCount(nodeB));
192193
waitForState(() -> {
193194
GetIngestionStateResponse ingestionState = getIngestionState(indexName);
194-
return Arrays.stream(ingestionState.getShardStates()).allMatch(state -> state.getPollerState().equalsIgnoreCase("paused"));
195+
return Arrays.stream(ingestionState.getShardStates()).allMatch(state -> state.isPollerPaused() && state.getPollerState().equalsIgnoreCase("paused"));
195196
});
196197

197198
// resume ingestion
198199
ResumeIngestionResponse resumeResponse = resumeIngestion(indexName);
199200
assertTrue(resumeResponse.isAcknowledged());
201+
assertTrue(resumeResponse.isShardsAcknowledged());
200202
waitForState(() -> {
201203
GetIngestionStateResponse ingestionState = getIngestionState(indexName);
202204
return Arrays.stream(ingestionState.getShardStates())
203205
.allMatch(
204-
state -> state.getPollerState().equalsIgnoreCase("polling") || state.getPollerState().equalsIgnoreCase("processing")
206+
state -> state.isPollerPaused() == false && (state.getPollerState().equalsIgnoreCase("polling") || state.getPollerState().equalsIgnoreCase("processing"))
205207
);
206208
});
207209
waitForSearchableDocs(4, Arrays.asList(nodeB, nodeC));

‎server/src/main/java/org/opensearch/action/ActionModule.java

+3
Original file line numberDiff line numberDiff line change
@@ -204,6 +204,8 @@
204204
import org.opensearch.action.admin.indices.streamingingestion.resume.TransportResumeIngestionAction;
205205
import org.opensearch.action.admin.indices.streamingingestion.state.GetIngestionStateAction;
206206
import org.opensearch.action.admin.indices.streamingingestion.state.TransportGetIngestionStateAction;
207+
import org.opensearch.action.admin.indices.streamingingestion.state.TransportUpdateIngestionStateAction;
208+
import org.opensearch.action.admin.indices.streamingingestion.state.UpdateIngestionStateAction;
207209
import org.opensearch.action.admin.indices.template.delete.DeleteComponentTemplateAction;
208210
import org.opensearch.action.admin.indices.template.delete.DeleteComposableIndexTemplateAction;
209211
import org.opensearch.action.admin.indices.template.delete.DeleteIndexTemplateAction;
@@ -819,6 +821,7 @@ public <Request extends ActionRequest, Response extends ActionResponse> void reg
819821
actions.register(PauseIngestionAction.INSTANCE, TransportPauseIngestionAction.class);
820822
actions.register(ResumeIngestionAction.INSTANCE, TransportResumeIngestionAction.class);
821823
actions.register(GetIngestionStateAction.INSTANCE, TransportGetIngestionStateAction.class);
824+
actions.register(UpdateIngestionStateAction.INSTANCE, TransportUpdateIngestionStateAction.class);
822825

823826
return unmodifiableMap(actions.getRegistry());
824827
}

‎server/src/main/java/org/opensearch/action/admin/indices/streamingingestion/IngestionStateShardFailure.java

+110
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
9+
package org.opensearch.action.admin.indices.streamingingestion;
10+
11+
import org.opensearch.common.annotation.ExperimentalApi;
12+
import org.opensearch.core.common.Strings;
13+
import org.opensearch.core.common.io.stream.StreamInput;
14+
import org.opensearch.core.common.io.stream.StreamOutput;
15+
import org.opensearch.core.common.io.stream.Writeable;
16+
import org.opensearch.core.xcontent.MediaTypeRegistry;
17+
import org.opensearch.core.xcontent.ToXContentFragment;
18+
import org.opensearch.core.xcontent.XContentBuilder;
19+
20+
import java.io.IOException;
21+
import java.util.ArrayList;
22+
import java.util.HashMap;
23+
import java.util.List;
24+
import java.util.Map;
25+
import java.util.Objects;
26+
27+
/**
28+
* Indicates ingestion failures at index and shard level.
29+
*
30+
* @opensearch.experimental
31+
*/
32+
@ExperimentalApi
33+
public class IngestionStateShardFailure implements Writeable, ToXContentFragment {
34+
private static final String SHARD = "shard";
35+
private static final String ERROR = "error";
36+
37+
private final String index;
38+
private final int shard;
39+
private String errorMessage;
40+
41+
public IngestionStateShardFailure(String index, int shard, String errorMessage) {
42+
this.index = index;
43+
this.shard = shard;
44+
this.errorMessage = errorMessage;
45+
}
46+
47+
public IngestionStateShardFailure(StreamInput in) throws IOException {
48+
this.index = in.readString();
49+
this.shard = in.readInt();
50+
this.errorMessage = in.readString();
51+
}
52+
53+
public String getIndex() {
54+
return index;
55+
}
56+
57+
public int getShard() {
58+
return shard;
59+
}
60+
61+
@Override
62+
public void writeTo(StreamOutput out) throws IOException {
63+
out.writeString(index);
64+
out.writeInt(shard);
65+
out.writeString(errorMessage);
66+
}
67+
68+
@Override
69+
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
70+
builder.startObject();
71+
builder.field(SHARD, shard);
72+
builder.field(ERROR, errorMessage);
73+
return builder.endObject();
74+
}
75+
76+
@Override
77+
public boolean equals(Object o) {
78+
if (this == o) return true;
79+
if (o == null || getClass() != o.getClass()) return false;
80+
IngestionStateShardFailure that = (IngestionStateShardFailure) o;
81+
return Objects.equals(index, that.index) && shard == that.shard && Objects.equals(errorMessage, that.errorMessage);
82+
}
83+
84+
@Override
85+
public int hashCode() {
86+
int result = Objects.hashCode(index);
87+
result = 31 * result + Objects.hashCode(shard);
88+
result = 31 * result + Objects.hashCode(errorMessage);
89+
return result;
90+
}
91+
92+
@Override
93+
public String toString() {
94+
return Strings.toString(MediaTypeRegistry.JSON, this);
95+
}
96+
97+
/**
98+
* Groups provided shard ingestion state failures by index name.
99+
*/
100+
public static Map<String, List<IngestionStateShardFailure>> groupShardFailuresByIndex(IngestionStateShardFailure[] shardFailures) {
101+
Map<String, List<IngestionStateShardFailure>> shardFailuresByIndex = new HashMap<>();
102+
103+
for (IngestionStateShardFailure shardFailure : shardFailures) {
104+
shardFailuresByIndex.computeIfAbsent(shardFailure.getIndex(), (index) -> new ArrayList<>());
105+
shardFailuresByIndex.get(shardFailure.getIndex()).add(shardFailure);
106+
}
107+
108+
return shardFailuresByIndex;
109+
}
110+
}

‎server/src/main/java/org/opensearch/action/admin/indices/streamingingestion/IngestionUpdateStateResponse.java

+96
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
9+
package org.opensearch.action.admin.indices.streamingingestion;
10+
11+
import org.opensearch.action.support.clustermanager.AcknowledgedResponse;
12+
import org.opensearch.common.annotation.ExperimentalApi;
13+
import org.opensearch.core.common.Strings;
14+
import org.opensearch.core.common.io.stream.StreamInput;
15+
import org.opensearch.core.common.io.stream.StreamOutput;
16+
import org.opensearch.core.xcontent.MediaTypeRegistry;
17+
import org.opensearch.core.xcontent.XContentBuilder;
18+
19+
import java.io.IOException;
20+
import java.util.List;
21+
import java.util.Map;
22+
23+
/**
24+
* Transport response for pausing ingestion.
25+
*
26+
* @opensearch.experimental
27+
*/
28+
@ExperimentalApi
29+
public class IngestionUpdateStateResponse extends AcknowledgedResponse {
30+
protected static final String SHARD_ACK = "shards_acknowledged";
31+
protected static final String ERROR = "error";
32+
protected static final String FAILURES = "failures";
33+
34+
protected boolean shardsAcknowledged;
35+
protected IngestionStateShardFailure[] shardFailuresList;
36+
protected String errorMessage;
37+
38+
public IngestionUpdateStateResponse(StreamInput in) throws IOException {
39+
super(in);
40+
shardFailuresList = in.readArray(IngestionStateShardFailure::new, IngestionStateShardFailure[]::new);
41+
errorMessage = in.readString();
42+
shardsAcknowledged = in.readBoolean();
43+
}
44+
45+
public IngestionUpdateStateResponse(final boolean acknowledged, final boolean shardsAcknowledged, final IngestionStateShardFailure[] shardFailuresList, String errorMessage) {
46+
super(acknowledged);
47+
this.shardFailuresList = shardFailuresList;
48+
this.shardsAcknowledged = shardsAcknowledged;
49+
this.errorMessage = errorMessage;
50+
}
51+
52+
@Override
53+
public void writeTo(StreamOutput out) throws IOException {
54+
super.writeTo(out);
55+
out.writeArray(shardFailuresList);
56+
out.writeString(errorMessage);
57+
out.writeBoolean(shardsAcknowledged);
58+
}
59+
60+
@Override
61+
protected void addCustomFields(final XContentBuilder builder, final Params params) throws IOException {
62+
super.addCustomFields(builder, params);
63+
builder.field(SHARD_ACK, shardsAcknowledged);
64+
builder.field(ERROR, errorMessage);
65+
66+
if (shardFailuresList.length > 0) {
67+
Map<String, List<IngestionStateShardFailure>> shardFailuresByIndex = IngestionStateShardFailure.groupShardFailuresByIndex(shardFailuresList);
68+
builder.startObject(FAILURES);
69+
for (Map.Entry<String, List<IngestionStateShardFailure>> indexShardFailures : shardFailuresByIndex.entrySet()) {
70+
builder.startArray(indexShardFailures.getKey());
71+
for (IngestionStateShardFailure shardFailure : indexShardFailures.getValue()) {
72+
shardFailure.toXContent(builder, params);
73+
}
74+
builder.endArray();
75+
}
76+
builder.endObject();
77+
}
78+
}
79+
80+
@Override
81+
public String toString() {
82+
return Strings.toString(MediaTypeRegistry.JSON, this);
83+
}
84+
85+
public boolean isShardsAcknowledged() {
86+
return shardsAcknowledged;
87+
}
88+
89+
public IngestionStateShardFailure[] getShardFailures() {
90+
return shardFailuresList;
91+
}
92+
93+
public String getErrorMessage() {
94+
return errorMessage;
95+
}
96+
}

‎server/src/main/java/org/opensearch/action/admin/indices/streamingingestion/pause/PauseIngestionClusterStateUpdateRequest.java

-34
This file was deleted.

0 commit comments

Comments
 (0)