Fix TimePartitionUtils Overflow (#13881) (#13885)

apache · Oct 23, 2024 · 0f63b46 · 0f63b46
1 parent cb0d2a5
commit 0f63b46
Show file tree

Hide file tree

Showing 2 changed files with 110 additions and 10 deletions.
diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/utils/TimePartitionUtils.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/utils/TimePartitionUtils.java
@@ -23,6 +23,8 @@
 
 import org.apache.tsfile.read.filter.basic.Filter;
 
+import java.math.BigInteger;
+
 public class TimePartitionUtils {
 
   /**
@@ -36,6 +38,39 @@ public class TimePartitionUtils {
   private static long timePartitionInterval =
       CommonDescriptor.getInstance().getConfig().getTimePartitionInterval();
 
+  private static final BigInteger bigTimePartitionOrigin = BigInteger.valueOf(timePartitionOrigin);
+  private static final BigInteger bigTimePartitionInterval =
+      BigInteger.valueOf(timePartitionInterval);
+  private static final boolean originMayCauseOverflow = (timePartitionOrigin != 0);
+  private static final long timePartitionLowerBoundWithoutOverflow;
+  private static final long timePartitionUpperBoundWithoutOverflow;
+
+  static {
+    long minPartition = getTimePartitionIdWithoutOverflow(Long.MIN_VALUE);
+    long maxPartition = getTimePartitionIdWithoutOverflow(Long.MAX_VALUE);
+    BigInteger minPartitionStartTime =
+        BigInteger.valueOf(minPartition)
+            .multiply(bigTimePartitionInterval)
+            .add(bigTimePartitionOrigin);
+    BigInteger maxPartitionEndTime =
+        BigInteger.valueOf(maxPartition)
+            .multiply(bigTimePartitionInterval)
+            .add(bigTimePartitionInterval)
+            .add(bigTimePartitionOrigin);
+    if (minPartitionStartTime.compareTo(BigInteger.valueOf(Long.MIN_VALUE)) < 0) {
+      timePartitionLowerBoundWithoutOverflow =
+          minPartitionStartTime.add(bigTimePartitionInterval).longValue();
+    } else {
+      timePartitionLowerBoundWithoutOverflow = minPartitionStartTime.longValue();
+    }
+    if (maxPartitionEndTime.compareTo(BigInteger.valueOf(Long.MAX_VALUE)) > 0) {
+      timePartitionUpperBoundWithoutOverflow =
+          maxPartitionEndTime.subtract(bigTimePartitionInterval).longValue();
+    } else {
+      timePartitionUpperBoundWithoutOverflow = maxPartitionEndTime.longValue();
+    }
+  }
+
   public static TTimePartitionSlot getTimePartitionSlot(long time) {
     TTimePartitionSlot timePartitionSlot = new TTimePartitionSlot();
     timePartitionSlot.setStartTime(getTimePartitionLowerBound(time));
@@ -47,14 +82,27 @@ public static long getTimePartitionInterval() {
   }
 
   public static long getTimePartitionLowerBound(long time) {
-    long lowerBoundOfTimePartition;
-    lowerBoundOfTimePartition =
-        getTimePartitionId(time) * timePartitionInterval + timePartitionOrigin;
-    return lowerBoundOfTimePartition;
+    if (time < timePartitionLowerBoundWithoutOverflow) {
+      return Long.MIN_VALUE;
+    }
+    if (originMayCauseOverflow) {
+      return BigInteger.valueOf(getTimePartitionIdWithoutOverflow(time))
+          .multiply(bigTimePartitionInterval)
+          .add(bigTimePartitionOrigin)
+          .longValue();
+    } else {
+      return getTimePartitionId(time) * timePartitionInterval + timePartitionOrigin;
+    }
   }
 
   public static long getTimePartitionUpperBound(long time) {
-    return getTimePartitionLowerBound(time) + timePartitionInterval;
+    if (time >= timePartitionUpperBoundWithoutOverflow) {
+      return Long.MAX_VALUE;
+    }
+    long lowerBound = getTimePartitionLowerBound(time);
+    return lowerBound == Long.MIN_VALUE
+        ? timePartitionLowerBoundWithoutOverflow
+        : lowerBound + timePartitionInterval;
   }
 
   public static long getTimePartitionId(long time) {
@@ -64,19 +112,48 @@ public static long getTimePartitionId(long time) {
         : time / timePartitionInterval - 1;
   }
 
+  public static long getTimePartitionIdWithoutOverflow(long time) {
+    BigInteger bigTime = BigInteger.valueOf(time).subtract(bigTimePartitionOrigin);
+    BigInteger partitionId =
+        bigTime.compareTo(BigInteger.ZERO) > 0
+                || bigTime.remainder(bigTimePartitionInterval).equals(BigInteger.ZERO)
+            ? bigTime.divide(bigTimePartitionInterval)
+            : bigTime.divide(bigTimePartitionInterval).subtract(BigInteger.ONE);
+    return partitionId.longValue();
+  }
+
   public static boolean satisfyPartitionId(long startTime, long endTime, long partitionId) {
-    return getTimePartitionId(startTime) <= partitionId
-        && getTimePartitionId(endTime) >= partitionId;
+    long startPartition =
+        originMayCauseOverflow
+            ? getTimePartitionIdWithoutOverflow(startTime)
+            : getTimePartitionId(startTime);
+    long endPartition =
+        originMayCauseOverflow
+            ? getTimePartitionIdWithoutOverflow(endTime)
+            : getTimePartitionId(endTime);
+    return startPartition <= partitionId && endPartition >= partitionId;
   }
 
   public static boolean satisfyPartitionStartTime(Filter timeFilter, long partitionStartTime) {
+    long partitionEndTime =
+        partitionStartTime >= timePartitionLowerBoundWithoutOverflow
+            ? Long.MAX_VALUE
+            : (partitionStartTime + timePartitionInterval - 1);
     return timeFilter == null
-        || timeFilter.satisfyStartEndTime(
-            partitionStartTime, partitionStartTime + timePartitionInterval - 1);
+        || timeFilter.satisfyStartEndTime(partitionStartTime, partitionEndTime);
   }
 
   public static boolean satisfyTimePartition(Filter timeFilter, long partitionId) {
-    long partitionStartTime = partitionId * timePartitionInterval + timePartitionOrigin;
+    long partitionStartTime;
+    if (originMayCauseOverflow) {
+      partitionStartTime =
+          BigInteger.valueOf(partitionId)
+              .multiply(bigTimePartitionInterval)
+              .add(bigTimePartitionOrigin)
+              .longValue();
+    } else {
+      partitionStartTime = partitionId * timePartitionInterval + timePartitionOrigin;
+    }
     return satisfyPartitionStartTime(timeFilter, partitionStartTime);
   }
 
@@ -85,6 +162,13 @@ public static void setTimePartitionInterval(long timePartitionInterval) {
   }
 
   public static long getEstimateTimePartitionSize(long startTime, long endTime) {
+    if (endTime > 0 && startTime < 0) {
+      return BigInteger.valueOf(endTime)
+              .subtract(BigInteger.valueOf(startTime))
+              .divide(bigTimePartitionInterval)
+              .longValue()
+          + 1;
+    }
     return (endTime - startTime) / timePartitionInterval + 1;
   }
 }
diff --git a/...ore/node-commons/src/test/java/org/apache/iotdb/commons/utils/TimePartitionUtilsTest.java b/...ore/node-commons/src/test/java/org/apache/iotdb/commons/utils/TimePartitionUtilsTest.java
@@ -22,6 +22,7 @@
 import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot;
 import org.apache.iotdb.commons.conf.CommonDescriptor;
 
+import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
 
@@ -89,4 +90,19 @@ public void testGetTimePartitionSlot_NegativeBoundaryTime() {
     TTimePartitionSlot actualSlot = TimePartitionUtils.getTimePartitionSlot(testTime);
     assertEquals(expectedSlot.getStartTime(), actualSlot.getStartTime());
   }
+
+  @Test
+  public void testOverflow() {
+    long testTime = Long.MIN_VALUE;
+    TTimePartitionSlot actualSlot = TimePartitionUtils.getTimePartitionSlot(testTime);
+    Assert.assertTrue(actualSlot.getStartTime() < 0);
+    testTime += 1;
+    long lowerBound = TimePartitionUtils.getTimePartitionLowerBound(testTime);
+    assertEquals(Long.MIN_VALUE, lowerBound);
+    testTime = Long.MAX_VALUE;
+    actualSlot = TimePartitionUtils.getTimePartitionSlot(testTime);
+    Assert.assertTrue(actualSlot.getStartTime() > 0);
+    long upperBound = TimePartitionUtils.getTimePartitionUpperBound(testTime);
+    assertEquals(Long.MAX_VALUE, upperBound);
+  }
 }