diff --git a/client/src/test/java/com/metamx/druid/client/BatchServerInventoryViewTest.java b/client/src/test/java/com/metamx/druid/client/BatchServerInventoryViewTest.java new file mode 100644 index 00000000000..7690b0ab06b --- /dev/null +++ b/client/src/test/java/com/metamx/druid/client/BatchServerInventoryViewTest.java @@ -0,0 +1,231 @@ +/* + * Druid - a distributed column store. + * Copyright (C) 2012 Metamarkets Group Inc. + * + * This program is free software; you can redistribute it and/or + * modify it under the terms of the GNU General Public License + * as published by the Free Software Foundation; either version 2 + * of the License, or (at your option) any later version. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU General Public License for more details. + * + * You should have received a copy of the GNU General Public License + * along with this program; if not, write to the Free Software + * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. + */ + +package com.metamx.druid.client; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.base.Joiner; +import com.google.common.base.Stopwatch; +import com.google.common.collect.Iterables; +import com.google.common.collect.Sets; +import com.google.common.util.concurrent.MoreExecutors; +import com.metamx.common.ISE; +import com.metamx.druid.coordination.BatchDataSegmentAnnouncer; +import com.metamx.druid.coordination.DruidServerMetadata; +import com.metamx.druid.curator.PotentiallyGzippedCompressionProvider; +import com.metamx.druid.curator.announcement.Announcer; +import com.metamx.druid.initialization.ZkDataSegmentAnnouncerConfig; +import com.metamx.druid.initialization.ZkPathsConfig; +import com.metamx.druid.jackson.DefaultObjectMapper; +import junit.framework.Assert; +import org.apache.curator.framework.CuratorFramework; +import org.apache.curator.framework.CuratorFrameworkFactory; +import org.apache.curator.retry.ExponentialBackoffRetry; +import org.apache.curator.test.TestingCluster; +import org.joda.time.DateTime; +import org.joda.time.Interval; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +import java.util.Set; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; + +/** + */ +public class BatchServerInventoryViewTest +{ + private static final String testBasePath = "/test"; + private static final Joiner joiner = Joiner.on("/"); + + private TestingCluster testingCluster; + private CuratorFramework cf; + private ObjectMapper jsonMapper; + private Announcer announcer; + private BatchDataSegmentAnnouncer segmentAnnouncer; + private Set testSegments; + private BatchServerInventoryView batchServerInventoryView; + + @Before + public void setUp() throws Exception + { + testingCluster = new TestingCluster(1); + testingCluster.start(); + + cf = CuratorFrameworkFactory.builder() + .connectString(testingCluster.getConnectString()) + .retryPolicy(new ExponentialBackoffRetry(1, 10)) + .compressionProvider(new PotentiallyGzippedCompressionProvider(true)) + .build(); + cf.start(); + cf.create().creatingParentsIfNeeded().forPath(testBasePath); + + jsonMapper = new DefaultObjectMapper(); + + announcer = new Announcer( + cf, + MoreExecutors.sameThreadExecutor() + ); + announcer.start(); + + segmentAnnouncer = new BatchDataSegmentAnnouncer( + new DruidServerMetadata( + "id", + "host", + Long.MAX_VALUE, + "type", + "tier" + ), + new ZkDataSegmentAnnouncerConfig() + { + @Override + public String getZkBasePath() + { + return testBasePath; + } + + @Override + public int getSegmentsPerNode() + { + return 50; + } + + @Override + public long getMaxNumBytes() + { + return 100000; + } + + @Override + public String getAnnouncerType() + { + return "batch"; + } + }, + announcer, + jsonMapper + ); + segmentAnnouncer.start(); + + testSegments = Sets.newHashSet(); + for (int i = 0; i < 100; i++) { + testSegments.add(makeSegment(i)); + } + + batchServerInventoryView = new BatchServerInventoryView( + new ServerInventoryViewConfig() + { + @Override + public int getRemovedSegmentLifetime() + { + return 0; + } + + @Override + public String getAnnouncerType() + { + return "batch"; + } + }, + new ZkPathsConfig() + { + @Override + public String getZkBasePath() + { + return testBasePath; + } + }, + cf, + Executors.newSingleThreadExecutor(), + jsonMapper + ); + + batchServerInventoryView.start(); + } + + @After + public void tearDown() throws Exception + { + batchServerInventoryView.stop(); + segmentAnnouncer.stop(); + announcer.stop(); + cf.close(); + testingCluster.stop(); + } + + @Test + public void testRun() throws Exception + { + segmentAnnouncer.announceSegments(testSegments); + + waitForSync(); + + DruidServer server = Iterables.get(batchServerInventoryView.getInventory(), 0); + Set segments = Sets.newHashSet(server.getSegments().values()); + + Assert.assertEquals(testSegments, segments); + + DataSegment segment1 = makeSegment(101); + DataSegment segment2 = makeSegment(102); + + segmentAnnouncer.announceSegment(segment1); + segmentAnnouncer.announceSegment(segment2); + testSegments.add(segment1); + testSegments.add(segment2); + + waitForSync(); + + Assert.assertEquals(testSegments, Sets.newHashSet(server.getSegments().values())); + + segmentAnnouncer.unannounceSegment(segment1); + segmentAnnouncer.unannounceSegment(segment2); + testSegments.remove(segment1); + testSegments.remove(segment2); + + waitForSync(); + + Assert.assertEquals(testSegments, Sets.newHashSet(server.getSegments().values())); + } + + private DataSegment makeSegment(int offset) + { + return DataSegment.builder() + .dataSource("foo") + .interval( + new Interval( + new DateTime("2013-01-01").plusDays(offset), + new DateTime("2013-01-02").plusDays(offset) + ) + ) + .version(new DateTime().toString()) + .build(); + } + + private void waitForSync() throws Exception + { + Stopwatch stopwatch = new Stopwatch().start(); + while (Iterables.get(batchServerInventoryView.getInventory(), 0).getSegments().size() != testSegments.size()) { + Thread.sleep(500); + if (stopwatch.elapsed(TimeUnit.MILLISECONDS) > 5000) { + throw new ISE("BatchServerInventoryView is not updating"); + } + } + } +}