Skip to content

Commit

Permalink
Add to and from XContent to ClusterBlock and ClusterBlocks
Browse files Browse the repository at this point in the history 
Signed-off-by: Shivansh Arora <[email protected]>
  • Loading branch information
@shiv0408
shiv0408 committed May 15, 2024
1 parent febe8c7 commit 75ce997
Show file tree
Hide file tree
Showing 5 changed files with 385 additions and 8 deletions.
83 changes: 78 additions & 5 deletions server/src/main/java/org/opensearch/cluster/block/ClusterBlock.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,17 +34,22 @@

import org.opensearch.common.Nullable;
import org.opensearch.common.annotation.PublicApi;
import org.opensearch.common.util.set.Sets;
import org.opensearch.core.common.io.stream.StreamInput;
import org.opensearch.core.common.io.stream.StreamOutput;
import org.opensearch.core.common.io.stream.Writeable;
import org.opensearch.core.rest.RestStatus;
import org.opensearch.core.xcontent.ToXContentFragment;
import org.opensearch.core.xcontent.XContentBuilder;
import org.opensearch.core.xcontent.XContentParseException;
import org.opensearch.core.xcontent.XContentParser;
import org.opensearch.core.xcontent.XContentParserUtils;

import java.io.IOException;
import java.util.EnumSet;
import java.util.Locale;
import java.util.Objects;
import java.util.Set;

/**
* Blocks the cluster for concurrency
Expand All @@ -54,6 +59,12 @@
@PublicApi(since = "1.0.0")
public class ClusterBlock implements Writeable, ToXContentFragment {

static final String KEY_UUID = "uuid";
static final String KEY_DESCRIPTION = "description";
static final String KEY_RETRYABLE = "retryable";
static final String KEY_DISABLE_STATE_PERSISTENCE = "disable_state_persistence";
static final String KEY_LEVELS = "levels";
private static final Set<String> VALID_FIELDS = Sets.newHashSet(KEY_UUID, KEY_DESCRIPTION, KEY_RETRYABLE, KEY_DISABLE_STATE_PERSISTENCE, KEY_LEVELS);
private final int id;
@Nullable
private final String uuid;
Expand Down Expand Up @@ -156,14 +167,14 @@ public boolean disableStatePersistence() {
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
builder.startObject(Integer.toString(id));
if (uuid != null) {
builder.field("uuid", uuid);
builder.field(KEY_UUID, uuid);
}
builder.field("description", description);
builder.field("retryable", retryable);
builder.field(KEY_DESCRIPTION, description);
builder.field(KEY_RETRYABLE, retryable);
if (disableStatePersistence) {
builder.field("disable_state_persistence", disableStatePersistence);
builder.field(KEY_DISABLE_STATE_PERSISTENCE, disableStatePersistence);
}
builder.startArray("levels");
builder.startArray(KEY_LEVELS);
for (ClusterBlockLevel level : levels) {
builder.value(level.name().toLowerCase(Locale.ROOT));
}
Expand All @@ -172,6 +183,68 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
return builder;
}

public static ClusterBlock fromXContent(XContentParser parser, int id) throws IOException {
String uuid = null;
String description = null;
boolean retryable = false;
boolean disableStatePersistence = false;
EnumSet<ClusterBlockLevel> levels = EnumSet.noneOf(ClusterBlockLevel.class);
String currentFieldName = skipBlockID(parser);
XContentParser.Token token;
while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
if (token == XContentParser.Token.FIELD_NAME) {
currentFieldName = parser.currentName();
} else if (token.isValue()) {
switch (Objects.requireNonNull(currentFieldName)) {
case KEY_UUID:
uuid = parser.text();
break;
case KEY_DESCRIPTION:
description = parser.text();
break;
case KEY_RETRYABLE:
retryable = parser.booleanValue();
break;
case KEY_DISABLE_STATE_PERSISTENCE:
disableStatePersistence = parser.booleanValue();
break;
default:
throw new IllegalArgumentException("unknown field [" + currentFieldName + "]");
}
} else if (token == XContentParser.Token.START_ARRAY) {
if (currentFieldName.equals(KEY_LEVELS)) {
while ((token = parser.nextToken()) != XContentParser.Token.END_ARRAY) {
levels.add(ClusterBlockLevel.fromString(parser.text(), Locale.ROOT));
}
} else {
throw new IllegalArgumentException("unknown field [" + currentFieldName + "]");
}
} else {
throw new IllegalArgumentException("unexpected token [" + token + "]");
}
}
return new ClusterBlock(id, uuid, description, retryable, disableStatePersistence, false, null, levels);
}

private static String skipBlockID(XContentParser parser) throws IOException {
if (parser.currentToken() == null) {
parser.nextToken();
}
if (parser.currentToken() == XContentParser.Token.START_OBJECT) {
parser.nextToken();
if (parser.currentToken() == XContentParser.Token.FIELD_NAME) {
String currentFieldName = parser.currentName();
if (VALID_FIELDS.contains(currentFieldName)) {
return currentFieldName;
} else {
// we have hit block id, just move on
parser.nextToken();
}
}
}
return null;
}

@Override
public void writeTo(StreamOutput out) throws IOException {
out.writeVInt(id);
Expand Down
8 changes: 8 additions & 0 deletions server/src/main/java/org/opensearch/cluster/block/ClusterBlockLevel.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@
import org.opensearch.common.annotation.PublicApi;

import java.util.EnumSet;
import java.util.Locale;

/**
* What level to block the cluster
Expand All @@ -51,4 +52,11 @@ public enum ClusterBlockLevel {

public static final EnumSet<ClusterBlockLevel> ALL = EnumSet.allOf(ClusterBlockLevel.class);
public static final EnumSet<ClusterBlockLevel> READ_WRITE = EnumSet.of(READ, WRITE);

/*
* This method is used to convert a string to a ClusterBlockLevel.
* */
public static ClusterBlockLevel fromString(String level, Locale locale) {
return ClusterBlockLevel.valueOf(level.toUpperCase(locale));
}
}
102 changes: 100 additions & 2 deletions server/src/main/java/org/opensearch/cluster/block/ClusterBlocks.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,11 @@
import org.opensearch.core.common.io.stream.StreamInput;
import org.opensearch.core.common.io.stream.StreamOutput;
import org.opensearch.core.rest.RestStatus;
import org.opensearch.core.xcontent.ToXContent;
import org.opensearch.core.xcontent.ToXContentFragment;
import org.opensearch.core.xcontent.XContentBuilder;
import org.opensearch.core.xcontent.XContentParser;
import org.opensearch.core.xcontent.XContentParserUtils;
import org.opensearch.index.IndexModule;

import java.io.IOException;
Expand All @@ -63,7 +68,7 @@
* @opensearch.api
*/
@PublicApi(since = "1.0.0")
public class ClusterBlocks extends AbstractDiffable<ClusterBlocks> {
public class ClusterBlocks extends AbstractDiffable<ClusterBlocks> implements ToXContentFragment {
public static final ClusterBlocks EMPTY_CLUSTER_BLOCK = new ClusterBlocks(emptySet(), Map.of());

private final Set<ClusterBlock> global;
Expand Down Expand Up @@ -326,6 +331,16 @@ public static Diff<ClusterBlocks> readDiffFrom(StreamInput in) throws IOExceptio
return AbstractDiffable.readDiffFrom(ClusterBlocks::readFrom, in);
}

@Override
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
Builder.toXContext(this, builder, params);
return builder;
}

public static ClusterBlocks fromXContent(XContentParser parser) throws IOException {
return Builder.fromXContent(parser);
}

/**
* An immutable level holder.
*
Expand Down Expand Up @@ -427,10 +442,16 @@ public Builder removeGlobalBlock(int blockId) {
}

public Builder addIndexBlock(String index, ClusterBlock block) {
prepareIndexForBlocks(index);
indices.get(index).add(block);
return this;
}

// initialize an index adding further blocks
private Builder prepareIndexForBlocks(String index) {
if (!indices.containsKey(index)) {
indices.put(index, new HashSet<>());
}
indices.get(index).add(block);
return this;
}

Expand Down Expand Up @@ -480,5 +501,82 @@ public ClusterBlocks build() {
}
return new ClusterBlocks(unmodifiableSet(new HashSet<>(global)), indicesBuilder);
}

public static void toXContext(ClusterBlocks blocks, XContentBuilder builder, ToXContent.Params params) throws IOException {
builder.startObject("blocks");
if (blocks.global().isEmpty() == false) {
builder.startObject("global");
for (ClusterBlock block : blocks.global()) {
block.toXContent(builder, params);
}
builder.endObject();
}

if (blocks.indices().isEmpty() == false) {
builder.startObject("indices");
for (Map.Entry<String, Set<ClusterBlock>> entry : blocks.indices().entrySet()) {
builder.startObject(entry.getKey());
for (ClusterBlock block : entry.getValue()) {
block.toXContent(builder, params);
}
builder.endObject();
}
builder.endObject();
}
builder.endObject();
}

public static ClusterBlocks fromXContent(XContentParser parser) throws IOException {
Builder builder = new Builder();
String currentFieldName = skipBlocksField(parser);
XContentParser.Token token;
while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
XContentParserUtils.ensureExpectedToken(XContentParser.Token.FIELD_NAME, token, parser);
currentFieldName = parser.currentName();
parser.nextToken();
switch (currentFieldName) {
case "global":
while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
currentFieldName = parser.currentName();
parser.nextToken();
builder.addGlobalBlock(ClusterBlock.fromXContent(parser, Integer.parseInt(currentFieldName)));
}
break;
case "indices":
while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
String indexName = parser.currentName();
parser.nextToken();
// prepare for this index as we want to add this to ClusterBlocks even if there is no Block associated with it
builder.prepareIndexForBlocks(indexName);
while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
currentFieldName = parser.currentName();
parser.nextToken();
builder.addIndexBlock(indexName, ClusterBlock.fromXContent(parser, Integer.parseInt(currentFieldName)));
}
}
break;
default:
throw new IllegalArgumentException("unknown field [" + currentFieldName + "]");
}
}
return builder.build();
}

private static String skipBlocksField(XContentParser parser) throws IOException {
if (parser.currentToken() == null) {
parser.nextToken();
}
if (parser.currentToken() == XContentParser.Token.START_OBJECT) {
parser.nextToken();
if (parser.currentToken() == XContentParser.Token.FIELD_NAME) {
if ("blocks".equals(parser.currentName())) {
parser.nextToken();
} else {
return parser.currentName();
}
}
}
return null;
}
}
}
76 changes: 75 additions & 1 deletion server/src/test/java/org/opensearch/cluster/block/ClusterBlockTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,13 +35,22 @@
import org.opensearch.Version;
import org.opensearch.common.UUIDs;
import org.opensearch.common.io.stream.BytesStreamOutput;
import org.opensearch.common.xcontent.json.JsonXContent;
import org.opensearch.core.common.bytes.BytesReference;
import org.opensearch.core.common.io.stream.StreamInput;
import org.opensearch.core.rest.RestStatus;
import org.opensearch.core.xcontent.MediaType;
import org.opensearch.core.xcontent.MediaTypeRegistry;
import org.opensearch.core.xcontent.ToXContent;
import org.opensearch.core.xcontent.XContentBuilder;
import org.opensearch.test.OpenSearchTestCase;
import org.opensearch.test.XContentTestUtils;

import java.io.IOException;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Locale;
import java.util.Map;

import static java.util.EnumSet.copyOf;
Expand Down Expand Up @@ -136,7 +145,72 @@ public void testGetIndexBlockWithId() {
assertThat(builder.build().getIndexBlockWithId("index", randomValueOtherThan(blockId, OpenSearchTestCase::randomInt)), nullValue());
}

private ClusterBlock randomClusterBlock() {
public void testToXContent() throws IOException {
ClusterBlock clusterBlock = randomClusterBlock();
XContentBuilder builder = JsonXContent.contentBuilder().prettyPrint();
builder.startObject();
clusterBlock.toXContent(builder, ToXContent.EMPTY_PARAMS);
builder.endObject();

String expectedString = "{\n" + getExpectedXContentFragment(clusterBlock, " ") + "\n}";

assertEquals(expectedString, builder.toString());
}

public void testFromXContent() throws IOException {
doFromXContentTestWithRandomFields(false);
}

public void testFromXContentWithRandomFields() throws IOException {
doFromXContentTestWithRandomFields(true);
}

private void doFromXContentTestWithRandomFields(boolean addRandomFields) throws IOException {
ClusterBlock clusterBlock = randomClusterBlock();
boolean humanReadable = randomBoolean();
final MediaType mediaType = MediaTypeRegistry.JSON;
BytesReference originalBytes = toShuffledXContent(clusterBlock, mediaType, ToXContent.EMPTY_PARAMS, humanReadable);

if (addRandomFields) {
String unsupportedField = "unsupported_field";
BytesReference mutated = BytesReference.bytes(
XContentTestUtils.insertIntoXContent(
mediaType.xContent(),
originalBytes,
Collections.singletonList(Integer.toString(clusterBlock.id())),
() -> unsupportedField,
() -> randomAlphaOfLengthBetween(3, 10)
)
);
IllegalArgumentException e = expectThrows(
IllegalArgumentException.class,
() -> ClusterBlock.fromXContent(createParser(mediaType.xContent(), mutated), clusterBlock.id())
);
assertEquals(e.getMessage(), "unknown field [" + unsupportedField + "]");
} else {
ClusterBlock parsedBlock = ClusterBlock.fromXContent(createParser(mediaType.xContent(), originalBytes), clusterBlock.id());
assertEquals(clusterBlock, parsedBlock);
assertEquals(clusterBlock.description(), parsedBlock.description());
assertEquals(clusterBlock.retryable(), parsedBlock.retryable());
assertEquals(clusterBlock.disableStatePersistence(), parsedBlock.disableStatePersistence());
assertArrayEquals(clusterBlock.levels().toArray(), parsedBlock.levels().toArray());
}
}

static String getExpectedXContentFragment(ClusterBlock clusterBlock, String indent) {
return indent + "\"" + clusterBlock.id() + "\" : {\n"
+ (clusterBlock.uuid() != null ?
indent + " \"uuid\" : \""+ clusterBlock.uuid() + "\",\n" : "")
+ indent + " \"description\" : \"" + clusterBlock.description() + "\",\n"
+ indent + " \"retryable\" : " + clusterBlock.retryable() + ",\n"
+ (clusterBlock.disableStatePersistence() ?
indent + " \"disable_state_persistence\" : " + clusterBlock.disableStatePersistence() + ",\n" : "")
+ String.format(indent + " \"levels\" : [%s]\n", clusterBlock.levels().isEmpty() ? " " :
"\n" + String.join(",\n", clusterBlock.levels().stream().map(level -> indent + " \"" + level.name().toLowerCase(Locale.ROOT) + "\"").toArray(String[]::new)) + "\n " + indent)
+ indent + "}";
}

static ClusterBlock randomClusterBlock() {
final String uuid = randomBoolean() ? UUIDs.randomBase64UUID() : null;
final List<ClusterBlockLevel> levels = Arrays.asList(ClusterBlockLevel.values());
return new ClusterBlock(
Expand Down
Loading

0 comments on commit 75ce997

Please sign in to comment.