diff --git a/server/src/test/java/com/metamx/druid/master/DruidMasterBalancerTester.java b/server/src/test/java/com/metamx/druid/master/DruidMasterBalancerTester.java index 690fa991fd9..c27230f4a60 100644 --- a/server/src/test/java/com/metamx/druid/master/DruidMasterBalancerTester.java +++ b/server/src/test/java/com/metamx/druid/master/DruidMasterBalancerTester.java @@ -57,7 +57,7 @@ public class DruidMasterBalancerTester extends DruidMasterBalancer { serverMap.put(toServerName,0); } - serverMap.put(fromServerName,serverMap.get(fromServerName)-1); +// serverMap.put(fromServerName,serverMap.get(fromServerName)-1); serverMap.put(toServerName, serverMap.get(toServerName)+1); } catch (Exception e) { diff --git a/server/src/test/java/com/metamx/druid/utils/DruidMasterBalancerProfiler.java b/server/src/test/java/com/metamx/druid/utils/DruidMasterBalancerProfiler.java index b46f19b7b06..01736e4fe7b 100644 --- a/server/src/test/java/com/metamx/druid/utils/DruidMasterBalancerProfiler.java +++ b/server/src/test/java/com/metamx/druid/utils/DruidMasterBalancerProfiler.java @@ -1,26 +1,34 @@ package com.metamx.druid.utils; import com.google.common.base.Stopwatch; +import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.Lists; import com.google.common.collect.Maps; import com.google.common.collect.MinMaxPriorityQueue; import com.metamx.druid.client.DataSegment; import com.metamx.druid.client.DruidServer; +import com.metamx.druid.db.DatabaseRuleManager; import com.metamx.druid.master.DruidCluster; import com.metamx.druid.master.DruidMaster; import com.metamx.druid.master.DruidMasterBalancerTester; +import com.metamx.druid.master.DruidMasterRuleRunner; import com.metamx.druid.master.DruidMasterRuntimeParams; import com.metamx.druid.master.LoadPeonCallback; import com.metamx.druid.master.LoadQueuePeon; import com.metamx.druid.master.LoadQueuePeonTester; +import com.metamx.druid.master.ReplicationThrottler; +import com.metamx.druid.master.SegmentReplicantLookup; import com.metamx.druid.master.ServerHolder; +import com.metamx.druid.master.rules.PeriodLoadRule; +import com.metamx.druid.master.rules.Rule; import com.metamx.druid.shard.NoneShardSpec; import com.metamx.emitter.EmittingLogger; import com.metamx.emitter.service.ServiceEmitter; import org.easymock.EasyMock; import org.joda.time.DateTime; import org.joda.time.Interval; +import org.joda.time.Period; import org.junit.Before; import org.junit.Test; @@ -39,6 +47,9 @@ public class DruidMasterBalancerProfiler DateTime start1 = new DateTime("2012-01-01"); DateTime version = new DateTime("2012-03-01"); ServiceEmitter emitter; + DatabaseRuleManager manager; + PeriodLoadRule loadRule = new PeriodLoadRule(new Period("P5000Y"),3,"normal"); + List rules = ImmutableList.of(loadRule); @Before public void setUp() throws Exception { @@ -47,6 +58,7 @@ public class DruidMasterBalancerProfiler druidServer2 = EasyMock.createMock(DruidServer.class); emitter = EasyMock.createMock(ServiceEmitter.class); EmittingLogger.registerEmitter(emitter); + manager = EasyMock.createMock(DatabaseRuleManager.class); } @Test @@ -54,13 +66,18 @@ public class DruidMasterBalancerProfiler { Stopwatch watch = new Stopwatch(); int numServers = 1000; + + EasyMock.expect(manager.getAllRules()).andReturn(ImmutableMap.>of("test", rules)).anyTimes(); + EasyMock.expect(manager.getRules(EasyMock.anyObject())).andReturn(rules).anyTimes(); + EasyMock.expect(manager.getRulesWithDefault(EasyMock.anyObject())).andReturn(rules).anyTimes(); + EasyMock.replay(manager); + master.moveSegment( EasyMock.anyObject(), EasyMock.anyObject(), EasyMock.anyObject(), EasyMock.anyObject() ); - EasyMock.expectLastCall().anyTimes(); EasyMock.replay(master); @@ -92,7 +109,9 @@ public class DruidMasterBalancerProfiler EasyMock.expect(server.getMetadata()).andReturn(null).anyTimes(); EasyMock.expect(server.getCurrSize()).andReturn(30L).atLeastOnce(); EasyMock.expect(server.getMaxSize()).andReturn(100L).atLeastOnce(); + EasyMock.expect(server.getTier()).andReturn("normal").anyTimes(); EasyMock.expect(server.getName()).andReturn(Integer.toString(i)).atLeastOnce(); + EasyMock.expect(server.getHost()).andReturn(Integer.toString(i)).anyTimes(); if (i==0) { EasyMock.expect(server.getSegments()).andReturn(segmentMap).anyTimes(); @@ -107,7 +126,7 @@ public class DruidMasterBalancerProfiler LoadQueuePeon peon = new LoadQueuePeonTester(); peonMap.put(Integer.toString(i),peon); serverHolderList.add(new ServerHolder(server, peon)); - } + } DruidMasterRuntimeParams params = DruidMasterRuntimeParams.newBuilder() @@ -129,13 +148,29 @@ public class DruidMasterBalancerProfiler .withMaxSegmentsToMove(MAX_SEGMENTS_TO_MOVE) .withBalancerReferenceTimestamp(new DateTime("2013-01-01")) .withEmitter(emitter) + .withDatabaseRuleManager(manager) + .withReplicationManager(new ReplicationThrottler(2, 500)) + .withSegmentReplicantLookup( + SegmentReplicantLookup.make(new DruidCluster( + ImmutableMap.>of( + "normal", + MinMaxPriorityQueue.orderedBy(DruidMasterBalancerTester.percentUsedComparator) + .create( + serverHolderList + ) + ) + ) + ) + ) .build(); DruidMasterBalancerTester tester = new DruidMasterBalancerTester(master); + DruidMasterRuleRunner runner = new DruidMasterRuleRunner(master,500,5); watch.start(); - while (!tester.isBalanced(20,50)) + while (!tester.isBalanced(40,50)) { params = tester.run(params); + params = runner.run(params); } System.out.println(watch.stop()); }