Make ByteSizeUnit implements Writeable (#20557)

This commit makes ByteSizeUnit implement Writeable.

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 <tt>SizeUnit</tt> represents size at a given unit of
  * granularity and provides utility methods to convert across units.
  * A <tt>SizeUnit</tt> 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());
+    }
 }

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;
     }
 }

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) + "]"));
+    }
 }