add date utilities

Graylog2 · Jul 8, 2024 · 404aec3 · 404aec3
1 parent a5d4c07
commit 404aec3
Showing 1 changed file with 14 additions and 1 deletion.
diff --git a/graylog2-server/src/main/java/org/graylog2/system/traffic/TrafficUpdater.java b/graylog2-server/src/main/java/org/graylog2/system/traffic/TrafficUpdater.java
@@ -19,6 +19,8 @@
 import org.graylog2.plugin.system.NodeId;
 import org.joda.time.DateTime;
 
+import java.time.ZonedDateTime;
+import java.time.temporal.ChronoUnit;
 import java.util.Map;
 import java.util.TreeMap;
 import java.util.stream.Collectors;
@@ -34,10 +36,21 @@ static DateTime getHourBucketStart(DateTime observationTime) {
         return observationTime.withMinuteOfHour(0).withSecondOfMinute(0).withMillisOfSecond(0);
     }
 
-    static TreeMap<DateTime, Long> aggregateToDaily(Map<DateTime, Long> histogram) {
+    static Map<DateTime, Long> aggregateToDaily(Map<DateTime, Long> histogram) {
         return histogram.entrySet().stream()
                 .collect(Collectors.groupingBy(entry -> entry.getKey().withTimeAtStartOfDay(),
                         TreeMap::new,
                         Collectors.mapping(Map.Entry::getValue, Collectors.summingLong(Long::valueOf))));
     }
+
+    static Map<ZonedDateTime, Long> aggregateToDailyZDT(Map<ZonedDateTime, Long> histogram) {
+        return histogram.entrySet().stream()
+                .collect(Collectors.groupingBy(entry -> timeAtStartOfDay(entry.getKey()),
+                        TreeMap::new,
+                        Collectors.mapping(Map.Entry::getValue, Collectors.summingLong(Long::valueOf))));
+    }
+
+    static private ZonedDateTime timeAtStartOfDay(ZonedDateTime dateTime) {
+        return dateTime.truncatedTo(ChronoUnit.DAYS);
+    }
 }