Create DiskSize

7 files changed, 117 insertions, 38 deletions

diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/NodeSpec.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/NodeSpec.java
index 85c35369a22..76616a3a8f5 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/NodeSpec.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/NodeSpec.java
@@ -7,6 +7,7 @@ import com.yahoo.config.provision.ApplicationId;
 import com.yahoo.config.provision.DockerImage;
 import com.yahoo.config.provision.NodeResources;
 import com.yahoo.config.provision.NodeType;
+import com.yahoo.vespa.hosted.node.admin.task.util.file.DiskSize;
 
 import java.time.Instant;
 import java.util.EnumSet;
@@ -220,6 +221,10 @@ public class NodeSpec {
         return resources.memoryGb();
     }
 
+    public DiskSize diskSize() {
+        return DiskSize.of(resources.diskGb(), DiskSize.Unit.GB);
+    }
+
     public double diskGb() {
         return resources.diskGb();
     }
diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/maintenance/StorageMaintainer.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/maintenance/StorageMaintainer.java
index 7a7a0a8e2bc..9c718955c0e 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/maintenance/StorageMaintainer.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/maintenance/StorageMaintainer.java
@@ -16,6 +16,7 @@ import com.yahoo.vespa.hosted.node.admin.maintenance.disk.LinearCleanupRule;
 import com.yahoo.vespa.hosted.node.admin.nodeadmin.ConvergenceException;
 import com.yahoo.vespa.hosted.node.admin.nodeagent.NodeAgentContext;
 import com.yahoo.vespa.hosted.node.admin.task.util.file.FileFinder;
+import com.yahoo.vespa.hosted.node.admin.task.util.file.DiskSize;
 import com.yahoo.vespa.hosted.node.admin.task.util.file.UnixPath;
 import com.yahoo.vespa.hosted.node.admin.task.util.process.Terminal;
 
@@ -55,7 +56,7 @@ public class StorageMaintainer {
     private final Clock clock;
 
     // We cache disk usage to avoid doing expensive disk operations so often
-    private final Cache<ContainerName, Long> diskUsage = CacheBuilder.newBuilder()
+    private final Cache<ContainerName, DiskSize> diskUsage = CacheBuilder.newBuilder()
             .maximumSize(100)
             .expireAfterWrite(5, TimeUnit.MINUTES)
             .build();
@@ -72,12 +73,17 @@ public class StorageMaintainer {
         this.clock = clock;
     }
 
+    // TODO: Remove, use diskUsageFor() instead
     public Optional<Long> getDiskUsageFor(NodeAgentContext context) {
+        return diskUsageFor(context).map(DiskSize::bytes);
+    }
+
+    public Optional<DiskSize> diskUsageFor(NodeAgentContext context) {
         try {
-            Long cachedDiskUsage = diskUsage.getIfPresent(context.containerName());
+            DiskSize cachedDiskUsage = diskUsage.getIfPresent(context.containerName());
             if (cachedDiskUsage != null) return Optional.of(cachedDiskUsage);
 
-            long diskUsageBytes = getDiskUsedInBytes(context, context.pathOnHostFromPathInNode("/"));
+            DiskSize diskUsageBytes = getDiskUsed(context, context.pathOnHostFromPathInNode("/"));
             diskUsage.put(context.containerName(), diskUsageBytes);
             return Optional.of(diskUsageBytes);
         } catch (Exception e) {
@@ -86,9 +92,8 @@ public class StorageMaintainer {
         }
     }
 
-    // Public for testing
-    long getDiskUsedInBytes(TaskContext context, Path path) {
-        if (!Files.exists(path)) return 0;
+    DiskSize getDiskUsed(TaskContext context, Path path) {
+        if (!Files.exists(path)) return DiskSize.ZERO;
 
         String output = terminal.newCommandLine(context)
                 .add("du", "-xsk", path.toString())
@@ -101,14 +106,14 @@ public class StorageMaintainer {
             throw new ConvergenceException("Result from disk usage command not as expected: " + output);
         }
 
-        return 1024 * Long.parseLong(results[0]);
+        return DiskSize.of(Long.parseLong(results[0]), DiskSize.Unit.kiB);
     }
 
     public boolean cleanDiskIfFull(NodeAgentContext context) {
-        double totalBytes = context.node().diskGb() * 1_000_000_000L;
+        double totalBytes = context.node().diskSize().bytes();
         // Delete enough bytes to get below 80% disk usage, but only if we are already using more than 90% disk
-        long bytesToRemove = getDiskUsageFor(context)
-                .map(diskUsage -> (long) (diskUsage - 0.8 * totalBytes))
+        long bytesToRemove = diskUsageFor(context)
+                .map(diskUsage -> (long) (diskUsage.bytes() - 0.8 * totalBytes))
                 .filter(bytes -> bytes > totalBytes * 0.1)
                 .orElse(0L);
 
diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/maintenance/disk/DiskCleanup.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/maintenance/disk/DiskCleanup.java
index a9c1f8a9600..c9a35d09bff 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/maintenance/disk/DiskCleanup.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/maintenance/disk/DiskCleanup.java
@@ -2,6 +2,7 @@
 package com.yahoo.vespa.hosted.node.admin.maintenance.disk;
 
 import com.yahoo.vespa.hosted.node.admin.component.TaskContext;
+import com.yahoo.vespa.hosted.node.admin.task.util.file.DiskSize;
 
 import java.nio.file.Files;
 import java.nio.file.Path;
@@ -43,8 +44,8 @@ public class DiskCleanup {
                     });
 
         } finally {
-            String wantedDeleteSize = bytesToDisplayCount(bytesToRemove);
-            String deletedSize = bytesToDisplayCount(bytesToRemove - btr[0]);
+            String wantedDeleteSize = DiskSize.of(bytesToRemove).asString();
+            String deletedSize = DiskSize.of(bytesToRemove - btr[0]).asString();
             if (deletedPaths.size() > 20) {
                 context.log(logger, "Deleted %d files (%s) because disk was getting full", deletedPaths.size(), deletedSize);
             } else if (deletedPaths.size() > 0) {
@@ -56,12 +57,4 @@ public class DiskCleanup {
 
         return !deletedPaths.isEmpty();
     }
-
-    static String bytesToDisplayCount(long bytes) {
-        if (bytes < 1000) return bytes + " bytes";
-
-        int unit = -1;
-        for (; bytes >= 1000; unit++) bytes /= 1000;
-        return bytes + " " + UNITS[unit] + "B";
-    }
 }
diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/task/util/file/DiskSize.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/task/util/file/DiskSize.java
new file mode 100644
index 00000000000..e6d8835ef9c
--- /dev/null
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/task/util/file/DiskSize.java
@@ -0,0 +1,63 @@
+package com.yahoo.vespa.hosted.node.admin.task.util.file;
+
+import java.util.Objects;
+
+/**
+ * @author freva
+ */
+public class DiskSize {
+
+    public static final DiskSize ZERO = DiskSize.of(0);
+    private static final char[] UNITS = "kMGTPE".toCharArray();
+
+    public enum Unit {
+        kB(1000), kiB(1 << 10),
+        MB(1_000_000), MiB(1 << 20),
+        GB(1_000_000_000), GiB(1 << 30),
+        PB(1_000_000_000_000L), PiB(1L << 40);
+
+        private final long size;
+        Unit(long size) { this.size = size; }
+    }
+
+    private final long bytes;
+    private DiskSize(long bytes) { this.bytes = bytes; }
+
+    public long bytes() { return bytes; }
+    public long as(Unit unit) { return bytes / unit.size; }
+    public double asDouble(Unit unit) { return (double) bytes / unit.size; }
+
+    public DiskSize add(DiskSize other) { return new DiskSize(bytes + other.bytes); }
+
+    public static DiskSize of(long bytes) { return new DiskSize(bytes); }
+    public static DiskSize of(double bytes, Unit unit) { return new DiskSize((long) (bytes * unit.size)); }
+    public static DiskSize of(long bytes, Unit unit) { return new DiskSize(bytes * unit.size); }
+
+    public String asString() { return asString(0); }
+    public String asString(int decimals) {
+        if (bytes < 1000) return bytes + " bytes";
+
+        int unit = -1;
+        double remaining = bytes;
+        for (; remaining >= 1000; unit++) remaining /= 1000;
+        return String.format("%." + decimals + "f %sB", remaining, UNITS[unit]);
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        DiskSize size = (DiskSize) o;
+        return bytes == size.bytes;
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(bytes);
+    }
+
+    @Override
+    public String toString() {
+        return asString();
+    }
+}
diff --git a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/maintenance/StorageMaintainerTest.java b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/maintenance/StorageMaintainerTest.java
index 47f24ba67c3..c17d0017269 100644
--- a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/maintenance/StorageMaintainerTest.java
+++ b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/maintenance/StorageMaintainerTest.java
@@ -8,6 +8,7 @@ import com.yahoo.vespa.hosted.node.admin.maintenance.disk.DiskCleanup;
 import com.yahoo.vespa.hosted.node.admin.nodeagent.NodeAgentContext;
 import com.yahoo.vespa.hosted.node.admin.nodeagent.NodeAgentContextImpl;
 import com.yahoo.vespa.hosted.node.admin.task.util.file.FileFinder;
+import com.yahoo.vespa.hosted.node.admin.task.util.file.DiskSize;
 import com.yahoo.vespa.hosted.node.admin.task.util.process.TestTerminal;
 import com.yahoo.vespa.test.file.TestFileSystem;
 import org.junit.After;
@@ -52,17 +53,17 @@ public class StorageMaintainerTest {
             Files.createDirectories(context.pathOnHostFromPathInNode("/"));
 
             terminal.expectCommand("du -xsk /home/docker/host-1 2>&1", 0, "321\t/home/docker/host-1/");
-            assertEquals(Optional.of(328_704L), storageMaintainer.getDiskUsageFor(context));
+            assertEquals(Optional.of(DiskSize.of(328_704)), storageMaintainer.diskUsageFor(context));
 
             // Value should still be cached, no new execution against the terminal
-            assertEquals(Optional.of(328_704L), storageMaintainer.getDiskUsageFor(context));
+            assertEquals(Optional.of(DiskSize.of(328_704)), storageMaintainer.diskUsageFor(context));
         }
 
         @Test
         public void testNonExistingDiskUsed() {
             StorageMaintainer storageMaintainer = new StorageMaintainer(terminal, null, null);
-            long usedBytes = storageMaintainer.getDiskUsedInBytes(null, Paths.get("/fake/path"));
-            assertEquals(0L, usedBytes);
+            DiskSize size = storageMaintainer.getDiskUsed(null, Paths.get("/fake/path"));
+            assertEquals(DiskSize.ZERO, size);
         }
 
         @After
diff --git a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/maintenance/disk/DiskCleanupTest.java b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/maintenance/disk/DiskCleanupTest.java
index 3f383a68b84..7065b261b1b 100644
--- a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/maintenance/disk/DiskCleanupTest.java
+++ b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/maintenance/disk/DiskCleanupTest.java
@@ -88,19 +88,6 @@ public class DiskCleanupTest {
                 "/opt/vespa/var/crash/core4", "/opt/vespa/var/crash/vespa-proton-bin.core-232", "/opt/vespa/logs/vespa-2.log", "/opt/vespa/var/crash/core3");
     }
 
-    @Test
-    public void bytes_to_display_count_test() {
-        assertEquals("-1 bytes", DiskCleanup.bytesToDisplayCount(-1));
-        assertEquals("123 bytes", DiskCleanup.bytesToDisplayCount(123));
-        assertEquals("1 kB", DiskCleanup.bytesToDisplayCount(1_000));
-        assertEquals("15 MB", DiskCleanup.bytesToDisplayCount(15_000_000));
-        assertEquals("123 GB", DiskCleanup.bytesToDisplayCount(123_456_789_012L));
-        assertEquals("987 TB", DiskCleanup.bytesToDisplayCount(987_654_321_098_765L));
-        assertEquals("2 PB", DiskCleanup.bytesToDisplayCount(2_000_000_000_000_000L));
-        assertEquals("9 EB", DiskCleanup.bytesToDisplayCount(Long.MAX_VALUE));
-
-    }
-
     private static class DiskCleanupRuleMock implements DiskCleanupRule {
         private final ArrayList<PrioritizedFileAttributes> pfa = new ArrayList<>();
 
@@ -139,4 +126,4 @@ public class DiskCleanupTest {
             assertEquals(expected, actual);
         }
     }
-}
\ No newline at end of file
+}
diff --git a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/task/util/file/DiskSizeTest.java b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/task/util/file/DiskSizeTest.java
new file mode 100644
index 00000000000..bcf3651fdfe
--- /dev/null
+++ b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/task/util/file/DiskSizeTest.java
@@ -0,0 +1,25 @@
+package com.yahoo.vespa.hosted.node.admin.task.util.file;
+
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+
+/**
+ * @author freva
+ */
+public class DiskSizeTest {
+
+    @Test
+    public void bytes_to_display_count_test() {
+        assertEquals("-1 bytes", DiskSize.of(-1).asString());
+        assertEquals("123 bytes", DiskSize.of(123).asString());
+        assertEquals("1 kB", DiskSize.of(1_000).asString());
+        assertEquals("15 MB", DiskSize.of(15_000_000).asString());
+        assertEquals("123 GB", DiskSize.of(123_456_789_012L).asString());
+        assertEquals("988 TB", DiskSize.of(987_654_321_098_765L).asString());
+        assertEquals("987.7 TB", DiskSize.of(987_654_321_098_765L).asString(1));
+        assertEquals("987.65 TB", DiskSize.of(987_654_321_098_765L).asString(2));
+        assertEquals("2 PB", DiskSize.of(2_000_000_000_000_000L).asString());
+        assertEquals("9 EB", DiskSize.of(Long.MAX_VALUE).asString());
+    }
+}