Add `BytesQuantity`

2 files changed, 131 insertions, 0 deletions

diff --git a/vespajlib/src/main/java/ai/vespa/utils/BytesQuantity.java b/vespajlib/src/main/java/ai/vespa/utils/BytesQuantity.java
new file mode 100644
index 00000000000..b9d8426938e
--- /dev/null
+++ b/vespajlib/src/main/java/ai/vespa/utils/BytesQuantity.java
@@ -0,0 +1,89 @@
+package ai.vespa.utils;
+
+import java.util.Locale;
+import java.util.Objects;
+import java.util.regex.Pattern;
+
+/**
+ * Reprents a quantity of bytes with a human-readable string representation.
+ * Currently only supports binary units (e.g. 1 kB = 1024 bytes).
+ *
+ * @author bjorncs
+ */
+public class BytesQuantity {
+    public enum Unit {
+        BYTES, KB, MB, GB, TB;
+
+        public long binarySize() {
+            return switch (this) {
+                case BYTES -> 1;
+                case KB -> 1 << 10;
+                case MB -> 1 << 20;
+                case GB -> 1 << 30;
+                case TB -> 1L << 40;
+            };
+        }
+
+        public String toUnitString() {
+            return switch (this) {
+                case BYTES -> "bytes";
+                case KB -> "kB";
+                case MB -> "MB";
+                case GB -> "GB";
+                case TB -> "TB";
+            };
+        }
+
+        static Unit fromString(String s) {
+            return switch (s) {
+                case "", "b", "B", "bytes" -> BYTES;
+                case "kB", "k", "K", "KB" -> KB;
+                case "MB", "m", "M" -> MB;
+                case "GB", "g", "G" -> GB;
+                case "TB", "t", "T" -> TB;
+                default -> throw new IllegalArgumentException("Invalid unit: " + s);
+            };
+        }
+    }
+
+    private final long bytes;
+    private BytesQuantity(long bytes) { this.bytes = bytes; }
+
+    public long toBytes() { return bytes; }
+
+    public static BytesQuantity ofBytes(long bytes) { return new BytesQuantity(bytes); }
+    public static BytesQuantity ofKB(long kb) { return BytesQuantity.of(kb, Unit.KB); }
+    public static BytesQuantity ofMB(long mb) { return BytesQuantity.of(mb, Unit.MB); }
+    public static BytesQuantity ofGB(long gb) { return BytesQuantity.of(gb, Unit.GB); }
+    public static BytesQuantity ofTB(long tb) { return BytesQuantity.of(tb, Unit.TB); }
+    public static BytesQuantity of(long value, Unit unit) { return new BytesQuantity(value * unit.binarySize()); }
+
+    private static final Pattern PATTERN = Pattern.compile("^(?<digits>\\d+)\\s*(?<unit>[a-zA-Z]*)$");
+    public static BytesQuantity fromString(String value) {
+        var matcher = PATTERN.matcher(value);
+        if (!matcher.matches())
+            throw new IllegalArgumentException(
+                    "Bytes quantity '%s' does not match pattern '%s'".formatted(value, PATTERN.pattern()));
+        var digits = Long.parseLong(matcher.group("digits"));
+        var unit = Unit.fromString(matcher.group("unit"));
+        return BytesQuantity.of(digits, unit);
+    }
+
+    public String asPrettyString() {
+        long remaining = bytes;
+        int unit = 0;
+        for (; remaining % 1024 == 0 && unit < Unit.values().length - 1; unit++) remaining /= 1024;
+        return String.format(Locale.ENGLISH, "%d %s", remaining, Unit.values()[unit].toUnitString());
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        BytesQuantity that = (BytesQuantity) o;
+        return bytes == that.bytes;
+    }
+
+    @Override public int hashCode() { return Objects.hashCode(bytes); }
+    @Override public String toString() { return asPrettyString(); }
+}
diff --git a/vespajlib/src/test/java/ai/vespa/utils/BytesQuantityTest.java b/vespajlib/src/test/java/ai/vespa/utils/BytesQuantityTest.java
new file mode 100644
index 00000000000..ba888528b85
--- /dev/null
+++ b/vespajlib/src/test/java/ai/vespa/utils/BytesQuantityTest.java
@@ -0,0 +1,42 @@
+package ai.vespa.utils;
+
+import org.junit.jupiter.api.Test;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+/**
+ * @author bjorncs
+ */
+class BytesQuantityTest {
+
+    @Test
+    void from_string() {
+        assertEquals(1L, BytesQuantity.fromString("1 bytes").toBytes());
+        assertEquals(1L, BytesQuantity.fromString("1 B").toBytes());
+        assertEquals(1L, BytesQuantity.fromString("1B").toBytes());
+        assertEquals(1L, BytesQuantity.fromString("1").toBytes());
+        assertEquals(1024L, BytesQuantity.fromString("1kB").toBytes());
+        assertEquals(1024L, BytesQuantity.fromString("1k").toBytes());
+        assertEquals(1024L, BytesQuantity.fromString("1K").toBytes());
+        assertEquals(1024L, BytesQuantity.fromString("1KB").toBytes());
+        assertEquals(1024L * 1024, BytesQuantity.fromString("1MB").toBytes());
+        assertEquals(1024L * 1024, BytesQuantity.fromString("1M").toBytes());
+        assertEquals(1024L * 1024 * 1024, BytesQuantity.fromString("1GB").toBytes());
+        assertEquals(1024L * 1024 * 1024, BytesQuantity.fromString("1G").toBytes());
+        assertEquals(1024L * 1024 * 1024 * 1024, BytesQuantity.fromString("1TB").toBytes());
+        assertEquals(1024L * 1024 * 1024 * 1024, BytesQuantity.fromString("1T").toBytes());
+    }
+
+    @Test
+    void as_pretty_string() {
+        assertEquals("1 bytes", BytesQuantity.ofBytes(1).asPrettyString());
+        assertEquals("2 kB", BytesQuantity.ofKB(2).asPrettyString());
+        assertEquals("3 MB", BytesQuantity.ofMB(3).asPrettyString());
+        assertEquals("4 GB", BytesQuantity.ofGB(4).asPrettyString());
+        assertEquals("5 TB", BytesQuantity.ofTB(5).asPrettyString());
+
+        assertEquals("2560 bytes", BytesQuantity.ofBytes(2048 + 512).asPrettyString());
+        assertEquals("3073 kB", BytesQuantity.ofBytes(3 * 1024 * 1024 + 1024).asPrettyString());
+        assertEquals("5120 TB", BytesQuantity.ofBytes(1024L * 1024 * 1024 * 1024 * 1024 * 5).asPrettyString());
+    }
+}