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.of(), + Arrays.asList("d"), + Arrays.asList("m"), + null, + null, + size + ); + } +}