diff --git a/core/src/main/java/org/elasticsearch/cluster/ClusterInfo.java b/core/src/main/java/org/elasticsearch/cluster/ClusterInfo.java index 20a95339633..e6e4cf1bb9f 100644 --- a/core/src/main/java/org/elasticsearch/cluster/ClusterInfo.java +++ b/core/src/main/java/org/elasticsearch/cluster/ClusterInfo.java @@ -43,7 +43,7 @@ public class ClusterInfo implements ToXContent, Writeable { private final ImmutableOpenMap mostAvailableSpaceUsage; final ImmutableOpenMap shardSizes; public static final ClusterInfo EMPTY = new ClusterInfo(); - private final ImmutableOpenMap routingToDataPath; + final ImmutableOpenMap routingToDataPath; protected ClusterInfo() { this(ImmutableOpenMap.of(), ImmutableOpenMap.of(), ImmutableOpenMap.of(), ImmutableOpenMap.of()); @@ -68,29 +68,10 @@ public class ClusterInfo implements ToXContent, Writeable { } public ClusterInfo(StreamInput in) throws IOException { - int size = in.readInt(); - Map leastMap = new HashMap<>(size); - for (int i = 0; i < size; i++) { - leastMap.put(in.readString(), new DiskUsage(in)); - } - - size = in.readInt(); - Map mostMap = new HashMap<>(size); - for (int i = 0; i < size; i++) { - mostMap.put(in.readString(), new DiskUsage(in)); - } - - size = in.readInt(); - Map sizeMap = new HashMap<>(size); - for (int i = 0; i < size; i++) { - sizeMap.put(in.readString(), in.readLong()); - } - - size = in.readInt(); - Map routingMap = new HashMap<>(size); - for (int i = 0; i < size; i++) { - routingMap.put(new ShardRouting(in), in.readString()); - } + Map leastMap = in.readMap(StreamInput::readString, DiskUsage::new); + Map mostMap = in.readMap(StreamInput::readString, DiskUsage::new); + Map sizeMap = in.readMap(StreamInput::readString, StreamInput::readLong); + Map routingMap = in.readMap(ShardRouting::new, StreamInput::readString); ImmutableOpenMap.Builder leastBuilder = ImmutableOpenMap.builder(); this.leastAvailableSpaceUsage = leastBuilder.putAll(leastMap).build(); diff --git a/core/src/main/java/org/elasticsearch/cluster/DiskUsage.java b/core/src/main/java/org/elasticsearch/cluster/DiskUsage.java index 32d62ce0c31..868a3f45feb 100644 --- a/core/src/main/java/org/elasticsearch/cluster/DiskUsage.java +++ b/core/src/main/java/org/elasticsearch/cluster/DiskUsage.java @@ -25,10 +25,11 @@ import org.elasticsearch.common.io.stream.StreamOutput; import org.elasticsearch.common.io.stream.Writeable; import org.elasticsearch.common.unit.ByteSizeValue; import org.elasticsearch.common.xcontent.ToXContent; -import org.elasticsearch.common.xcontent.ToXContent.Params; import org.elasticsearch.common.xcontent.XContentBuilder; import java.io.IOException; +import java.util.Objects; + /** * Encapsulation class used to represent the amount of disk used on a node. */ @@ -126,6 +127,24 @@ public class DiskUsage implements ToXContent, Writeable { return getTotalBytes() - getFreeBytes(); } + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + + DiskUsage other = (DiskUsage) o; + return Objects.equals(nodeId, other.nodeId) && + Objects.equals(nodeName, other.nodeName) && + Objects.equals(totalBytes, other.totalBytes) && + Objects.equals(freeBytes, other.freeBytes); + + } + + @Override + public int hashCode() { + return Objects.hash(nodeId, nodeName, path, totalBytes, freeBytes); + } + @Override public String toString() { return "[" + nodeId + "][" + nodeName + "][" + path + "] free: " + new ByteSizeValue(getFreeBytes()) + diff --git a/core/src/main/java/org/elasticsearch/common/io/stream/StreamInput.java b/core/src/main/java/org/elasticsearch/common/io/stream/StreamInput.java index 2e30ff474f2..0f6ab65db1f 100644 --- a/core/src/main/java/org/elasticsearch/common/io/stream/StreamInput.java +++ b/core/src/main/java/org/elasticsearch/common/io/stream/StreamInput.java @@ -420,6 +420,17 @@ public abstract class StreamInput extends InputStream { return null; } + public Map readMap(Writeable.Reader keyReader, Writeable.Reader valueReader) throws IOException { + int size = readVInt(); + Map map = new HashMap<>(size); + for (int i = 0; i < size; i++) { + K key = keyReader.read(this); + V value = valueReader.read(this); + map.put(key, value); + } + return map; + } + @Nullable @SuppressWarnings("unchecked") public Map readMap() throws IOException { diff --git a/core/src/test/java/org/elasticsearch/cluster/ClusterInfoTests.java b/core/src/test/java/org/elasticsearch/cluster/ClusterInfoTests.java new file mode 100644 index 00000000000..99afee8b2c2 --- /dev/null +++ b/core/src/test/java/org/elasticsearch/cluster/ClusterInfoTests.java @@ -0,0 +1,88 @@ +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch licenses this file to you under + * the Apache License, Version 2.0 (the "License"); you may + * not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.elasticsearch.cluster; + +import org.elasticsearch.Version; +import org.elasticsearch.cluster.routing.RestoreSource; +import org.elasticsearch.cluster.routing.ShardRouting; +import org.elasticsearch.cluster.routing.UnassignedInfo; +import org.elasticsearch.common.collect.ImmutableOpenMap; +import org.elasticsearch.common.io.stream.BytesStreamOutput; +import org.elasticsearch.index.shard.ShardId; +import org.elasticsearch.snapshots.Snapshot; +import org.elasticsearch.snapshots.SnapshotId; +import org.elasticsearch.test.ESTestCase; + +public class ClusterInfoTests extends ESTestCase { + + public void testSerialization() throws Exception { + ClusterInfo clusterInfo = new ClusterInfo( + randomDiskUsage(), randomDiskUsage(), randomShardSizes(), randomRoutingToDataPath() + ); + BytesStreamOutput output = new BytesStreamOutput(); + clusterInfo.writeTo(output); + + ClusterInfo result = new ClusterInfo(output.bytes().streamInput()); + assertEquals(clusterInfo.getNodeLeastAvailableDiskUsages(), result.getNodeLeastAvailableDiskUsages()); + assertEquals(clusterInfo.getNodeMostAvailableDiskUsages(), result.getNodeMostAvailableDiskUsages()); + assertEquals(clusterInfo.shardSizes, result.shardSizes); + assertEquals(clusterInfo.routingToDataPath, result.routingToDataPath); + } + + private static ImmutableOpenMap randomDiskUsage() { + int numEntries = randomIntBetween(0, 128); + ImmutableOpenMap.Builder builder = ImmutableOpenMap.builder(numEntries); + for (int i = 0; i < numEntries; i++) { + String key = randomAsciiOfLength(32); + DiskUsage diskUsage = new DiskUsage( + randomAsciiOfLength(4), randomAsciiOfLength(4), randomAsciiOfLength(4), + randomIntBetween(0, Integer.MAX_VALUE), randomIntBetween(0, Integer.MAX_VALUE) + ); + builder.put(key, diskUsage); + } + return builder.build(); + } + + private static ImmutableOpenMap randomShardSizes() { + int numEntries = randomIntBetween(0, 128); + ImmutableOpenMap.Builder builder = ImmutableOpenMap.builder(numEntries); + for (int i = 0; i < numEntries; i++) { + String key = randomAsciiOfLength(32); + long shardSize = randomIntBetween(0, Integer.MAX_VALUE); + builder.put(key, shardSize); + } + return builder.build(); + } + + private static ImmutableOpenMap randomRoutingToDataPath() { + int numEntries = randomIntBetween(0, 128); + ImmutableOpenMap.Builder builder = ImmutableOpenMap.builder(numEntries); + for (int i = 0; i < numEntries; i++) { + RestoreSource restoreSource = new RestoreSource(new Snapshot(randomAsciiOfLength(4), + new SnapshotId(randomAsciiOfLength(4), randomAsciiOfLength(4))), Version.CURRENT, randomAsciiOfLength(4)); + UnassignedInfo.Reason reason = randomFrom(UnassignedInfo.Reason.values()); + UnassignedInfo unassignedInfo = new UnassignedInfo(reason, randomAsciiOfLength(4)); + ShardId shardId = new ShardId(randomAsciiOfLength(32), randomAsciiOfLength(32), randomIntBetween(0, Integer.MAX_VALUE)); + ShardRouting shardRouting = ShardRouting.newUnassigned(shardId, restoreSource, randomBoolean(), unassignedInfo); + builder.put(shardRouting, randomAsciiOfLength(32)); + } + return builder.build(); + } + +}