org.apache.hadoop.yarn.api.ApplicationMasterProtocol Java Examples
The following examples show how to use
org.apache.hadoop.yarn.api.ApplicationMasterProtocol.
You can vote up the ones you like or vote down the ones you don't like,
and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.
Example #1
Source File: YarnTajoResourceManager.java From incubator-tajo with Apache License 2.0 | 6 votes |
@Override public void init(Configuration conf) { this.conf = conf; connectYarnClient(); final YarnConfiguration yarnConf = new YarnConfiguration(conf); final YarnRPC rpc = YarnRPC.create(conf); final InetSocketAddress rmAddress = conf.getSocketAddr( YarnConfiguration.RM_SCHEDULER_ADDRESS, YarnConfiguration.DEFAULT_RM_SCHEDULER_ADDRESS, YarnConfiguration.DEFAULT_RM_SCHEDULER_PORT); UserGroupInformation currentUser; try { currentUser = UserGroupInformation.getCurrentUser(); } catch (IOException e) { throw new YarnRuntimeException(e); } rmClient = currentUser.doAs(new PrivilegedAction<ApplicationMasterProtocol>() { @Override public ApplicationMasterProtocol run() { return (ApplicationMasterProtocol) rpc.getProxy(ApplicationMasterProtocol.class, rmAddress, yarnConf); } }); }
Example #2
Source File: ClientRMProxy.java From hadoop with Apache License 2.0 | 6 votes |
@Private @Override protected InetSocketAddress getRMAddress(YarnConfiguration conf, Class<?> protocol) throws IOException { if (protocol == ApplicationClientProtocol.class) { return conf.getSocketAddr(YarnConfiguration.RM_ADDRESS, YarnConfiguration.DEFAULT_RM_ADDRESS, YarnConfiguration.DEFAULT_RM_PORT); } else if (protocol == ResourceManagerAdministrationProtocol.class) { return conf.getSocketAddr( YarnConfiguration.RM_ADMIN_ADDRESS, YarnConfiguration.DEFAULT_RM_ADMIN_ADDRESS, YarnConfiguration.DEFAULT_RM_ADMIN_PORT); } else if (protocol == ApplicationMasterProtocol.class) { setAMRMTokenService(conf); return conf.getSocketAddr(YarnConfiguration.RM_SCHEDULER_ADDRESS, YarnConfiguration.DEFAULT_RM_SCHEDULER_ADDRESS, YarnConfiguration.DEFAULT_RM_SCHEDULER_PORT); } else { String message = "Unsupported protocol found when creating the proxy " + "connection to ResourceManager: " + ((protocol != null) ? protocol.getClass().getName() : "null"); LOG.error(message); throw new IllegalStateException(message); } }
Example #3
Source File: TestRPCFactories.java From hadoop with Apache License 2.0 | 6 votes |
private void testPbServerFactory() { InetSocketAddress addr = new InetSocketAddress(0); Configuration conf = new Configuration(); ApplicationMasterProtocol instance = new AMRMProtocolTestImpl(); Server server = null; try { server = RpcServerFactoryPBImpl.get().getServer( ApplicationMasterProtocol.class, instance, addr, conf, null, 1); server.start(); } catch (YarnRuntimeException e) { e.printStackTrace(); Assert.fail("Failed to create server"); } finally { if (server != null) { server.stop(); } } }
Example #4
Source File: TestRPCFactories.java From big-c with Apache License 2.0 | 6 votes |
private void testPbServerFactory() { InetSocketAddress addr = new InetSocketAddress(0); Configuration conf = new Configuration(); ApplicationMasterProtocol instance = new AMRMProtocolTestImpl(); Server server = null; try { server = RpcServerFactoryPBImpl.get().getServer( ApplicationMasterProtocol.class, instance, addr, conf, null, 1); server.start(); } catch (YarnRuntimeException e) { e.printStackTrace(); Assert.fail("Failed to create server"); } finally { if (server != null) { server.stop(); } } }
Example #5
Source File: ClientRMProxy.java From big-c with Apache License 2.0 | 6 votes |
@Private @Override protected InetSocketAddress getRMAddress(YarnConfiguration conf, Class<?> protocol) throws IOException { if (protocol == ApplicationClientProtocol.class) { return conf.getSocketAddr(YarnConfiguration.RM_ADDRESS, YarnConfiguration.DEFAULT_RM_ADDRESS, YarnConfiguration.DEFAULT_RM_PORT); } else if (protocol == ResourceManagerAdministrationProtocol.class) { return conf.getSocketAddr( YarnConfiguration.RM_ADMIN_ADDRESS, YarnConfiguration.DEFAULT_RM_ADMIN_ADDRESS, YarnConfiguration.DEFAULT_RM_ADMIN_PORT); } else if (protocol == ApplicationMasterProtocol.class) { setAMRMTokenService(conf); return conf.getSocketAddr(YarnConfiguration.RM_SCHEDULER_ADDRESS, YarnConfiguration.DEFAULT_RM_SCHEDULER_ADDRESS, YarnConfiguration.DEFAULT_RM_SCHEDULER_PORT); } else { String message = "Unsupported protocol found when creating the proxy " + "connection to ResourceManager: " + ((protocol != null) ? protocol.getClass().getName() : "null"); LOG.error(message); throw new IllegalStateException(message); } }
Example #6
Source File: TestApplicationMasterServiceProtocolOnHA.java From big-c with Apache License 2.0 | 6 votes |
@Before public void initialize() throws Exception { startHACluster(0, false, false, true); attemptId = this.cluster.createFakeApplicationAttemptId(); amClient = ClientRMProxy .createRMProxy(this.conf, ApplicationMasterProtocol.class); Token<AMRMTokenIdentifier> appToken = this.cluster.getResourceManager().getRMContext() .getAMRMTokenSecretManager().createAndGetAMRMToken(attemptId); appToken.setService(ClientRMProxy.getAMRMTokenService(conf)); UserGroupInformation.setLoginUser(UserGroupInformation .createRemoteUser(UserGroupInformation.getCurrentUser() .getUserName())); UserGroupInformation.getCurrentUser().addToken(appToken); syncToken(appToken); }
Example #7
Source File: TestUnmanagedAMLauncher.java From big-c with Apache License 2.0 | 6 votes |
public static void main(String[] args) throws Exception { if (args[0].equals("success")) { ApplicationMasterProtocol client = ClientRMProxy.createRMProxy(conf, ApplicationMasterProtocol.class); client.registerApplicationMaster(RegisterApplicationMasterRequest .newInstance(NetUtils.getHostname(), -1, "")); Thread.sleep(1000); FinishApplicationMasterResponse resp = client.finishApplicationMaster(FinishApplicationMasterRequest .newInstance(FinalApplicationStatus.SUCCEEDED, "success", null)); assertTrue(resp.getIsUnregistered()); System.exit(0); } else { System.exit(1); } }
Example #8
Source File: TestApplicationMasterServiceProtocolOnHA.java From hadoop with Apache License 2.0 | 6 votes |
@Before public void initialize() throws Exception { startHACluster(0, false, false, true); attemptId = this.cluster.createFakeApplicationAttemptId(); amClient = ClientRMProxy .createRMProxy(this.conf, ApplicationMasterProtocol.class); Token<AMRMTokenIdentifier> appToken = this.cluster.getResourceManager().getRMContext() .getAMRMTokenSecretManager().createAndGetAMRMToken(attemptId); appToken.setService(ClientRMProxy.getAMRMTokenService(conf)); UserGroupInformation.setLoginUser(UserGroupInformation .createRemoteUser(UserGroupInformation.getCurrentUser() .getUserName())); UserGroupInformation.getCurrentUser().addToken(appToken); syncToken(appToken); }
Example #9
Source File: TestUnmanagedAMLauncher.java From hadoop with Apache License 2.0 | 6 votes |
public static void main(String[] args) throws Exception { if (args[0].equals("success")) { ApplicationMasterProtocol client = ClientRMProxy.createRMProxy(conf, ApplicationMasterProtocol.class); client.registerApplicationMaster(RegisterApplicationMasterRequest .newInstance(NetUtils.getHostname(), -1, "")); Thread.sleep(1000); FinishApplicationMasterResponse resp = client.finishApplicationMaster(FinishApplicationMasterRequest .newInstance(FinalApplicationStatus.SUCCEEDED, "success", null)); assertTrue(resp.getIsUnregistered()); System.exit(0); } else { System.exit(1); } }
Example #10
Source File: RMCommunicator.java From hadoop with Apache License 2.0 | 5 votes |
protected ApplicationMasterProtocol createSchedulerProxy() { final Configuration conf = getConfig(); try { return ClientRMProxy.createRMProxy(conf, ApplicationMasterProtocol.class); } catch (IOException e) { throw new YarnRuntimeException(e); } }
Example #11
Source File: TestAMRMTokens.java From big-c with Apache License 2.0 | 5 votes |
private ApplicationMasterProtocol createRMClient(final MockRM rm, final Configuration conf, final YarnRPC rpc, UserGroupInformation currentUser) { return currentUser.doAs(new PrivilegedAction<ApplicationMasterProtocol>() { @Override public ApplicationMasterProtocol run() { return (ApplicationMasterProtocol) rpc.getProxy(ApplicationMasterProtocol.class, rm .getApplicationMasterService().getBindAddress(), conf); } }); }
Example #12
Source File: StramClientUtils.java From Bats with Apache License 2.0 | 5 votes |
/** * Connect to the Resource Manager<p> * * @return Handle to communicate with the RM */ public ApplicationMasterProtocol connectToRM() { InetSocketAddress rmAddress = conf.getSocketAddr( YarnConfiguration.RM_SCHEDULER_ADDRESS, YarnConfiguration.DEFAULT_RM_SCHEDULER_ADDRESS, YarnConfiguration.DEFAULT_RM_SCHEDULER_PORT); LOG.debug("Connecting to ResourceManager at " + rmAddress); return ((ApplicationMasterProtocol)rpc.getProxy(ApplicationMasterProtocol.class, rmAddress, conf)); }
Example #13
Source File: AMRMClientImpl.java From big-c with Apache License 2.0 | 5 votes |
@Override protected void serviceStart() throws Exception { final YarnConfiguration conf = new YarnConfiguration(getConfig()); try { rmClient = ClientRMProxy.createRMProxy(conf, ApplicationMasterProtocol.class); } catch (IOException e) { throw new YarnRuntimeException(e); } super.serviceStart(); }
Example #14
Source File: RMCommunicator.java From big-c with Apache License 2.0 | 5 votes |
protected ApplicationMasterProtocol createSchedulerProxy() { final Configuration conf = getConfig(); try { return ClientRMProxy.createRMProxy(conf, ApplicationMasterProtocol.class); } catch (IOException e) { throw new YarnRuntimeException(e); } }
Example #15
Source File: TestLocalContainerAllocator.java From big-c with Apache License 2.0 | 5 votes |
@Test public void testAllocResponseId() throws Exception { ApplicationMasterProtocol scheduler = new MockScheduler(); Configuration conf = new Configuration(); LocalContainerAllocator lca = new StubbedLocalContainerAllocator(scheduler); lca.init(conf); lca.start(); // do two heartbeats to verify the response ID is being tracked lca.heartbeat(); lca.heartbeat(); lca.close(); }
Example #16
Source File: TestLocalContainerAllocator.java From hadoop with Apache License 2.0 | 5 votes |
@Test public void testAllocResponseId() throws Exception { ApplicationMasterProtocol scheduler = new MockScheduler(); Configuration conf = new Configuration(); LocalContainerAllocator lca = new StubbedLocalContainerAllocator(scheduler); lca.init(conf); lca.start(); // do two heartbeats to verify the response ID is being tracked lca.heartbeat(); lca.heartbeat(); lca.close(); }
Example #17
Source File: AMRMClientImpl.java From hadoop with Apache License 2.0 | 5 votes |
@Override protected void serviceStart() throws Exception { final YarnConfiguration conf = new YarnConfiguration(getConfig()); try { rmClient = ClientRMProxy.createRMProxy(conf, ApplicationMasterProtocol.class); } catch (IOException e) { throw new YarnRuntimeException(e); } super.serviceStart(); }
Example #18
Source File: StramClientUtils.java From attic-apex-core with Apache License 2.0 | 5 votes |
/** * Connect to the Resource Manager<p> * * @return Handle to communicate with the RM */ public ApplicationMasterProtocol connectToRM() { InetSocketAddress rmAddress = conf.getSocketAddr( YarnConfiguration.RM_SCHEDULER_ADDRESS, YarnConfiguration.DEFAULT_RM_SCHEDULER_ADDRESS, YarnConfiguration.DEFAULT_RM_SCHEDULER_PORT); LOG.debug("Connecting to ResourceManager at " + rmAddress); return ((ApplicationMasterProtocol)rpc.getProxy(ApplicationMasterProtocol.class, rmAddress, conf)); }
Example #19
Source File: TestAMRMTokens.java From hadoop with Apache License 2.0 | 5 votes |
private ApplicationMasterProtocol createRMClient(final MockRM rm, final Configuration conf, final YarnRPC rpc, UserGroupInformation currentUser) { return currentUser.doAs(new PrivilegedAction<ApplicationMasterProtocol>() { @Override public ApplicationMasterProtocol run() { return (ApplicationMasterProtocol) rpc.getProxy(ApplicationMasterProtocol.class, rm .getApplicationMasterService().getBindAddress(), conf); } }); }
Example #20
Source File: ApplicationMasterProtocolPBServiceImpl.java From big-c with Apache License 2.0 | 4 votes |
public ApplicationMasterProtocolPBServiceImpl(ApplicationMasterProtocol impl) { this.real = impl; }
Example #21
Source File: TestLocalContainerAllocator.java From big-c with Apache License 2.0 | 4 votes |
public StubbedLocalContainerAllocator(ApplicationMasterProtocol scheduler) { super(mock(ClientService.class), createAppContext(), "nmhost", 1, 2, null); this.scheduler = scheduler; }
Example #22
Source File: TestLocalContainerAllocator.java From big-c with Apache License 2.0 | 4 votes |
@Override protected ApplicationMasterProtocol createSchedulerProxy() { return scheduler; }
Example #23
Source File: TestRMContainerAllocator.java From big-c with Apache License 2.0 | 4 votes |
@Override protected ApplicationMasterProtocol createSchedulerProxy() { return this.rm.getApplicationMasterService(); }
Example #24
Source File: TestRMContainerAllocator.java From big-c with Apache License 2.0 | 4 votes |
@Test public void testHeartbeatHandler() throws Exception { LOG.info("Running testHeartbeatHandler"); Configuration conf = new Configuration(); conf.setInt(MRJobConfig.MR_AM_TO_RM_HEARTBEAT_INTERVAL_MS, 1); ControlledClock clock = new ControlledClock(new SystemClock()); AppContext appContext = mock(AppContext.class); when(appContext.getClock()).thenReturn(clock); when(appContext.getApplicationID()).thenReturn( ApplicationId.newInstance(1, 1)); RMContainerAllocator allocator = new RMContainerAllocator( mock(ClientService.class), appContext) { @Override protected void register() { } @Override protected ApplicationMasterProtocol createSchedulerProxy() { return mock(ApplicationMasterProtocol.class); } @Override protected synchronized void heartbeat() throws Exception { } }; allocator.init(conf); allocator.start(); clock.setTime(5); int timeToWaitMs = 5000; while (allocator.getLastHeartbeatTime() != 5 && timeToWaitMs > 0) { Thread.sleep(10); timeToWaitMs -= 10; } Assert.assertEquals(5, allocator.getLastHeartbeatTime()); clock.setTime(7); timeToWaitMs = 5000; while (allocator.getLastHeartbeatTime() != 7 && timeToWaitMs > 0) { Thread.sleep(10); timeToWaitMs -= 10; } Assert.assertEquals(7, allocator.getLastHeartbeatTime()); final AtomicBoolean callbackCalled = new AtomicBoolean(false); allocator.runOnNextHeartbeat(new Runnable() { @Override public void run() { callbackCalled.set(true); } }); clock.setTime(8); timeToWaitMs = 5000; while (allocator.getLastHeartbeatTime() != 8 && timeToWaitMs > 0) { Thread.sleep(10); timeToWaitMs -= 10; } Assert.assertEquals(8, allocator.getLastHeartbeatTime()); Assert.assertTrue(callbackCalled.get()); }
Example #25
Source File: TestSchedulerUtils.java From big-c with Apache License 2.0 | 4 votes |
@Test public void testValidateResourceBlacklistRequest() throws Exception { MyContainerManager containerManager = new MyContainerManager(); final MockRMWithAMS rm = new MockRMWithAMS(new YarnConfiguration(), containerManager); rm.start(); MockNM nm1 = rm.registerNode("localhost:1234", 5120); Map<ApplicationAccessType, String> acls = new HashMap<ApplicationAccessType, String>(2); acls.put(ApplicationAccessType.VIEW_APP, "*"); RMApp app = rm.submitApp(1024, "appname", "appuser", acls); nm1.nodeHeartbeat(true); RMAppAttempt attempt = app.getCurrentAppAttempt(); ApplicationAttemptId applicationAttemptId = attempt.getAppAttemptId(); waitForLaunchedState(attempt); // Create a client to the RM. final Configuration conf = rm.getConfig(); final YarnRPC rpc = YarnRPC.create(conf); UserGroupInformation currentUser = UserGroupInformation.createRemoteUser(applicationAttemptId.toString()); Credentials credentials = containerManager.getContainerCredentials(); final InetSocketAddress rmBindAddress = rm.getApplicationMasterService().getBindAddress(); Token<? extends TokenIdentifier> amRMToken = MockRMWithAMS.setupAndReturnAMRMToken(rmBindAddress, credentials.getAllTokens()); currentUser.addToken(amRMToken); ApplicationMasterProtocol client = currentUser.doAs(new PrivilegedAction<ApplicationMasterProtocol>() { @Override public ApplicationMasterProtocol run() { return (ApplicationMasterProtocol) rpc.getProxy( ApplicationMasterProtocol.class, rmBindAddress, conf); } }); RegisterApplicationMasterRequest request = Records .newRecord(RegisterApplicationMasterRequest.class); client.registerApplicationMaster(request); ResourceBlacklistRequest blacklistRequest = ResourceBlacklistRequest.newInstance( Collections.singletonList(ResourceRequest.ANY), null); AllocateRequest allocateRequest = AllocateRequest.newInstance(0, 0.0f, null, null, blacklistRequest); boolean error = false; try { client.allocate(allocateRequest); } catch (InvalidResourceBlacklistRequestException e) { error = true; } rm.stop(); Assert.assertTrue( "Didn't not catch InvalidResourceBlacklistRequestException", error); }
Example #26
Source File: MockAM.java From big-c with Apache License 2.0 | 4 votes |
public void setAMRMProtocol(ApplicationMasterProtocol amRMProtocol, RMContext context) { this.context = context; this.amRMProtocol = amRMProtocol; }
Example #27
Source File: MockAM.java From big-c with Apache License 2.0 | 4 votes |
public MockAM(RMContext context, ApplicationMasterProtocol amRMProtocol, ApplicationAttemptId attemptId) { this.context = context; this.amRMProtocol = amRMProtocol; this.attemptId = attemptId; }
Example #28
Source File: TestAMAuthorization.java From big-c with Apache License 2.0 | 4 votes |
@Test public void testUnauthorizedAccess() throws Exception { MyContainerManager containerManager = new MyContainerManager(); rm = new MockRMWithAMS(conf, containerManager); rm.start(); MockNM nm1 = rm.registerNode("localhost:1234", 5120); RMApp app = rm.submitApp(1024); nm1.nodeHeartbeat(true); int waitCount = 0; while (containerManager.containerTokens == null && waitCount++ < 40) { LOG.info("Waiting for AM Launch to happen.."); Thread.sleep(1000); } Assert.assertNotNull(containerManager.containerTokens); RMAppAttempt attempt = app.getCurrentAppAttempt(); ApplicationAttemptId applicationAttemptId = attempt.getAppAttemptId(); waitForLaunchedState(attempt); final Configuration conf = rm.getConfig(); final YarnRPC rpc = YarnRPC.create(conf); final InetSocketAddress serviceAddr = conf.getSocketAddr( YarnConfiguration.RM_SCHEDULER_ADDRESS, YarnConfiguration.DEFAULT_RM_SCHEDULER_ADDRESS, YarnConfiguration.DEFAULT_RM_SCHEDULER_PORT); UserGroupInformation currentUser = UserGroupInformation .createRemoteUser(applicationAttemptId.toString()); // First try contacting NM without tokens ApplicationMasterProtocol client = currentUser .doAs(new PrivilegedAction<ApplicationMasterProtocol>() { @Override public ApplicationMasterProtocol run() { return (ApplicationMasterProtocol) rpc.getProxy(ApplicationMasterProtocol.class, serviceAddr, conf); } }); RegisterApplicationMasterRequest request = Records .newRecord(RegisterApplicationMasterRequest.class); try { client.registerApplicationMaster(request); Assert.fail("Should fail with authorization error"); } catch (Exception e) { if (isCause(AccessControlException.class, e)) { // Because there are no tokens, the request should be rejected as the // server side will assume we are trying simple auth. String expectedMessage = ""; if (UserGroupInformation.isSecurityEnabled()) { expectedMessage = "Client cannot authenticate via:[TOKEN]"; } else { expectedMessage = "SIMPLE authentication is not enabled. Available:[TOKEN]"; } Assert.assertTrue(e.getCause().getMessage().contains(expectedMessage)); } else { throw e; } } // TODO: Add validation of invalid authorization when there's more data in // the AMRMToken }
Example #29
Source File: TestAMAuthorization.java From big-c with Apache License 2.0 | 4 votes |
@Test public void testAuthorizedAccess() throws Exception { MyContainerManager containerManager = new MyContainerManager(); rm = new MockRMWithAMS(conf, containerManager); rm.start(); MockNM nm1 = rm.registerNode("localhost:1234", 5120); Map<ApplicationAccessType, String> acls = new HashMap<ApplicationAccessType, String>(2); acls.put(ApplicationAccessType.VIEW_APP, "*"); RMApp app = rm.submitApp(1024, "appname", "appuser", acls); nm1.nodeHeartbeat(true); int waitCount = 0; while (containerManager.containerTokens == null && waitCount++ < 20) { LOG.info("Waiting for AM Launch to happen.."); Thread.sleep(1000); } Assert.assertNotNull(containerManager.containerTokens); RMAppAttempt attempt = app.getCurrentAppAttempt(); ApplicationAttemptId applicationAttemptId = attempt.getAppAttemptId(); waitForLaunchedState(attempt); // Create a client to the RM. final Configuration conf = rm.getConfig(); final YarnRPC rpc = YarnRPC.create(conf); UserGroupInformation currentUser = UserGroupInformation .createRemoteUser(applicationAttemptId.toString()); Credentials credentials = containerManager.getContainerCredentials(); final InetSocketAddress rmBindAddress = rm.getApplicationMasterService().getBindAddress(); Token<? extends TokenIdentifier> amRMToken = MockRMWithAMS.setupAndReturnAMRMToken(rmBindAddress, credentials.getAllTokens()); currentUser.addToken(amRMToken); ApplicationMasterProtocol client = currentUser .doAs(new PrivilegedAction<ApplicationMasterProtocol>() { @Override public ApplicationMasterProtocol run() { return (ApplicationMasterProtocol) rpc.getProxy(ApplicationMasterProtocol.class, rm .getApplicationMasterService().getBindAddress(), conf); } }); RegisterApplicationMasterRequest request = Records .newRecord(RegisterApplicationMasterRequest.class); RegisterApplicationMasterResponse response = client.registerApplicationMaster(request); Assert.assertNotNull(response.getClientToAMTokenMasterKey()); if (UserGroupInformation.isSecurityEnabled()) { Assert .assertTrue(response.getClientToAMTokenMasterKey().array().length > 0); } Assert.assertEquals("Register response has bad ACLs", "*", response.getApplicationACLs().get(ApplicationAccessType.VIEW_APP)); }
Example #30
Source File: TestRMContainerAllocator.java From hadoop with Apache License 2.0 | 4 votes |
@Test public void testHeartbeatHandler() throws Exception { LOG.info("Running testHeartbeatHandler"); Configuration conf = new Configuration(); conf.setInt(MRJobConfig.MR_AM_TO_RM_HEARTBEAT_INTERVAL_MS, 1); ControlledClock clock = new ControlledClock(new SystemClock()); AppContext appContext = mock(AppContext.class); when(appContext.getClock()).thenReturn(clock); when(appContext.getApplicationID()).thenReturn( ApplicationId.newInstance(1, 1)); RMContainerAllocator allocator = new RMContainerAllocator( mock(ClientService.class), appContext) { @Override protected void register() { } @Override protected ApplicationMasterProtocol createSchedulerProxy() { return mock(ApplicationMasterProtocol.class); } @Override protected synchronized void heartbeat() throws Exception { } }; allocator.init(conf); allocator.start(); clock.setTime(5); int timeToWaitMs = 5000; while (allocator.getLastHeartbeatTime() != 5 && timeToWaitMs > 0) { Thread.sleep(10); timeToWaitMs -= 10; } Assert.assertEquals(5, allocator.getLastHeartbeatTime()); clock.setTime(7); timeToWaitMs = 5000; while (allocator.getLastHeartbeatTime() != 7 && timeToWaitMs > 0) { Thread.sleep(10); timeToWaitMs -= 10; } Assert.assertEquals(7, allocator.getLastHeartbeatTime()); final AtomicBoolean callbackCalled = new AtomicBoolean(false); allocator.runOnNextHeartbeat(new Runnable() { @Override public void run() { callbackCalled.set(true); } }); clock.setTime(8); timeToWaitMs = 5000; while (allocator.getLastHeartbeatTime() != 8 && timeToWaitMs > 0) { Thread.sleep(10); timeToWaitMs -= 10; } Assert.assertEquals(8, allocator.getLastHeartbeatTime()); Assert.assertTrue(callbackCalled.get()); }