From a7ef5b8b70e0fa642f437ad0b75b38639352686d Mon Sep 17 00:00:00 2001
From: cheddar <cheddar@metamarkets.com>
Date: Wed, 14 Aug 2013 10:03:47 -0700
Subject: [PATCH] 1) Fix bug with SingleSegmentLoader.StorageLocation keeping
 track of its storage size incorrectly.  Add unit test ftw.

---
 .../druid/loading/SingleSegmentLoader.java    | 20 +++---
 .../loading/SingleSegmentLoaderTest.java      | 71 +++++++++++++++++++
 2 files changed, 81 insertions(+), 10 deletions(-)
 create mode 100644 server/src/test/java/com/metamx/druid/loading/SingleSegmentLoaderTest.java

diff --git a/server/src/main/java/com/metamx/druid/loading/SingleSegmentLoader.java b/server/src/main/java/com/metamx/druid/loading/SingleSegmentLoader.java
index f0e9a7f20e8..987939bb069 100644
--- a/server/src/main/java/com/metamx/druid/loading/SingleSegmentLoader.java
+++ b/server/src/main/java/com/metamx/druid/loading/SingleSegmentLoader.java
@@ -176,7 +176,7 @@ public class SingleSegmentLoader implements SegmentLoader
     }
   }
 
-  private static class StorageLocation
+  static class StorageLocation
   {
     private final File path;
     private final long maxSize;
@@ -195,41 +195,41 @@ public class SingleSegmentLoader implements SegmentLoader
       this.segments = Sets.newHashSet();
     }
 
-    private File getPath()
+    File getPath()
     {
       return path;
     }
 
-    private Long getMaxSize()
+    Long getMaxSize()
     {
       return maxSize;
     }
 
-    private synchronized void addSegment(DataSegment segment)
+    synchronized void addSegment(DataSegment segment)
     {
-      if (! segments.add(segment)) {
+      if (segments.add(segment)) {
         currSize += segment.getSize();
       }
     }
 
-    private synchronized void removeSegment(DataSegment segment)
+    synchronized void removeSegment(DataSegment segment)
     {
       if (segments.remove(segment)) {
         currSize -= segment.getSize();
       }
     }
 
-    private boolean canHandle(long size)
+    boolean canHandle(long size)
     {
-      return available() > size;
+      return available() >= size;
     }
 
-    private synchronized long available()
+    synchronized long available()
     {
       return maxSize - currSize;
     }
 
-    private StorageLocation mostEmpty(StorageLocation other)
+    StorageLocation mostEmpty(StorageLocation other)
     {
       return available() > other.available() ? this : other;
     }
diff --git a/server/src/test/java/com/metamx/druid/loading/SingleSegmentLoaderTest.java b/server/src/test/java/com/metamx/druid/loading/SingleSegmentLoaderTest.java
new file mode 100644
index 00000000000..f82311609ac
--- /dev/null
+++ b/server/src/test/java/com/metamx/druid/loading/SingleSegmentLoaderTest.java
@@ -0,0 +1,71 @@
+package com.metamx.druid.loading;
+
+import com.google.common.collect.ImmutableMap;
+import com.metamx.druid.client.DataSegment;
+import org.joda.time.Interval;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.io.File;
+import java.util.Arrays;
+
+/**
+ */
+public class SingleSegmentLoaderTest
+{
+  @Test
+  public void testStorageLocation() throws Exception
+  {
+    long expectedAvail = 1000l;
+    SingleSegmentLoader.StorageLocation loc = new SingleSegmentLoader.StorageLocation(new File("/tmp"), expectedAvail);
+
+    verifyLoc(expectedAvail, loc);
+
+    final DataSegment secondSegment = makeSegment("2012-01-02/2012-01-03", 23);
+
+    loc.addSegment(makeSegment("2012-01-01/2012-01-02", 10));
+    expectedAvail -= 10;
+    verifyLoc(expectedAvail, loc);
+
+    loc.addSegment(makeSegment("2012-01-01/2012-01-02", 10));
+    verifyLoc(expectedAvail, loc);
+
+    loc.addSegment(secondSegment);
+    expectedAvail -= 23;
+    verifyLoc(expectedAvail, loc);
+
+    loc.removeSegment(makeSegment("2012-01-01/2012-01-02", 10));
+    expectedAvail += 10;
+    verifyLoc(expectedAvail, loc);
+
+    loc.removeSegment(makeSegment("2012-01-01/2012-01-02", 10));
+    verifyLoc(expectedAvail, loc);
+
+    loc.removeSegment(secondSegment);
+    expectedAvail += 23;
+    verifyLoc(expectedAvail, loc);
+  }
+
+  private void verifyLoc(long maxSize, SingleSegmentLoader.StorageLocation loc)
+  {
+    Assert.assertEquals(maxSize, loc.available());
+    for (int i = 0; i <= maxSize; ++i) {
+      Assert.assertTrue(String.valueOf(i), loc.canHandle(i));
+    }
+  }
+
+  private DataSegment makeSegment(String intervalString, long size)
+  {
+    return new DataSegment(
+        "test",
+        new Interval(intervalString),
+        "1",
+        ImmutableMap.<String, Object>of(),
+        Arrays.asList("d"),
+        Arrays.asList("m"),
+        null,
+        null,
+        size
+    );
+  }
+}