diff --git a/core/src/main/java/org/elasticsearch/common/unit/ByteSizeUnit.java b/core/src/main/java/org/elasticsearch/common/unit/ByteSizeUnit.java index 4a159957d5e..7a412aac090 100644 --- a/core/src/main/java/org/elasticsearch/common/unit/ByteSizeUnit.java +++ b/core/src/main/java/org/elasticsearch/common/unit/ByteSizeUnit.java @@ -19,16 +19,20 @@ package org.elasticsearch.common.unit; +import org.elasticsearch.common.io.stream.StreamInput; +import org.elasticsearch.common.io.stream.StreamOutput; +import org.elasticsearch.common.io.stream.Writeable; + +import java.io.IOException; + /** * A SizeUnit represents size at a given unit of * granularity and provides utility methods to convert across units. * A SizeUnit does not maintain size information, but only * helps organize and use size representations that may be maintained * separately across various contexts. - * - * */ -public enum ByteSizeUnit { +public enum ByteSizeUnit implements Writeable { BYTES { @Override public long toBytes(long size) { @@ -225,6 +229,13 @@ public enum ByteSizeUnit { static final long MAX = Long.MAX_VALUE; + public static ByteSizeUnit fromId(int id) { + if (id < 0 || id >= values().length) { + throw new IllegalArgumentException("No byte size unit found for id [" + id + "]"); + } + return values()[id]; + } + /** * Scale d by m, checking for overflow. * This has a short name to make above code more readable. @@ -235,7 +246,6 @@ public enum ByteSizeUnit { return d * m; } - public abstract long toBytes(long size); public abstract long toKB(long size); @@ -247,4 +257,16 @@ public enum ByteSizeUnit { public abstract long toTB(long size); public abstract long toPB(long size); + + @Override + public void writeTo(StreamOutput out) throws IOException { + out.writeVInt(this.ordinal()); + } + + /** + * Reads a {@link ByteSizeUnit} from a given {@link StreamInput} + */ + public static ByteSizeUnit readFrom(StreamInput in) throws IOException { + return ByteSizeUnit.fromId(in.readVInt()); + } } \ No newline at end of file diff --git a/core/src/main/java/org/elasticsearch/common/unit/ByteSizeValue.java b/core/src/main/java/org/elasticsearch/common/unit/ByteSizeValue.java index db66885a974..7d2be6fee3e 100644 --- a/core/src/main/java/org/elasticsearch/common/unit/ByteSizeValue.java +++ b/core/src/main/java/org/elasticsearch/common/unit/ByteSizeValue.java @@ -32,11 +32,11 @@ import java.util.Objects; public class ByteSizeValue implements Writeable { private final long size; - private final ByteSizeUnit sizeUnit; + private final ByteSizeUnit unit; public ByteSizeValue(StreamInput in) throws IOException { size = in.readVLong(); - sizeUnit = ByteSizeUnit.BYTES; + unit = ByteSizeUnit.BYTES; } @Override @@ -48,9 +48,9 @@ public class ByteSizeValue implements Writeable { this(bytes, ByteSizeUnit.BYTES); } - public ByteSizeValue(long size, ByteSizeUnit sizeUnit) { + public ByteSizeValue(long size, ByteSizeUnit unit) { this.size = size; - this.sizeUnit = sizeUnit; + this.unit = unit; } public int bytesAsInt() { @@ -62,27 +62,27 @@ public class ByteSizeValue implements Writeable { } public long getBytes() { - return sizeUnit.toBytes(size); + return unit.toBytes(size); } public long getKb() { - return sizeUnit.toKB(size); + return unit.toKB(size); } public long getMb() { - return sizeUnit.toMB(size); + return unit.toMB(size); } public long getGb() { - return sizeUnit.toGB(size); + return unit.toGB(size); } public long getTb() { - return sizeUnit.toTB(size); + return unit.toTB(size); } public long getPb() { - return sizeUnit.toPB(size); + return unit.toPB(size); } public double getKbFrac() { @@ -199,7 +199,7 @@ public class ByteSizeValue implements Writeable { @Override public int hashCode() { int result = Long.hashCode(size); - result = 31 * result + (sizeUnit != null ? sizeUnit.hashCode() : 0); + result = 31 * result + (unit != null ? unit.hashCode() : 0); return result; } } diff --git a/core/src/test/java/org/elasticsearch/common/unit/ByteSizeUnitTests.java b/core/src/test/java/org/elasticsearch/common/unit/ByteSizeUnitTests.java index ab6d2818946..719313d1c86 100644 --- a/core/src/test/java/org/elasticsearch/common/unit/ByteSizeUnitTests.java +++ b/core/src/test/java/org/elasticsearch/common/unit/ByteSizeUnitTests.java @@ -19,20 +19,23 @@ package org.elasticsearch.common.unit; +import org.elasticsearch.common.io.stream.BytesStreamOutput; +import org.elasticsearch.common.io.stream.StreamInput; import org.elasticsearch.test.ESTestCase; +import java.io.IOException; + import static org.elasticsearch.common.unit.ByteSizeUnit.BYTES; import static org.elasticsearch.common.unit.ByteSizeUnit.GB; import static org.elasticsearch.common.unit.ByteSizeUnit.KB; import static org.elasticsearch.common.unit.ByteSizeUnit.MB; import static org.elasticsearch.common.unit.ByteSizeUnit.PB; import static org.elasticsearch.common.unit.ByteSizeUnit.TB; +import static org.hamcrest.Matchers.containsString; import static org.hamcrest.Matchers.equalTo; -/** - * - */ public class ByteSizeUnitTests extends ESTestCase { + public void testBytes() { assertThat(BYTES.toBytes(1), equalTo(1L)); assertThat(BYTES.toKB(1024), equalTo(1L)); @@ -77,4 +80,23 @@ public class ByteSizeUnitTests extends ESTestCase { assertThat(PB.toTB(1), equalTo(1024L)); assertThat(PB.toPB(1), equalTo(1L)); } + + public void testSerialization() throws IOException { + for (ByteSizeUnit unit : ByteSizeUnit.values()) { + try (BytesStreamOutput out = new BytesStreamOutput()) { + unit.writeTo(out); + + try (StreamInput in = out.bytes().streamInput()) { + ByteSizeUnit deserialized = ByteSizeUnit.readFrom(in); + assertEquals(unit, deserialized); + } + } + } + } + + public void testFromUnknownId() throws IOException { + final byte randomId = (byte) randomIntBetween(ByteSizeUnit.values().length + 1, 100); + IllegalArgumentException e = expectThrows(IllegalArgumentException.class, () -> ByteSizeUnit.fromId(randomId)); + assertThat(e.getMessage(), containsString("No byte size unit found for id [" + String.valueOf(randomId) + "]")); + } }