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 vote down vote up
@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 vote down vote up
@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 vote down vote up
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 vote down vote up
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 vote down vote up
@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 vote down vote up
@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 vote down vote up
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 vote down vote up
@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 vote down vote up
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 vote down vote up
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 vote down vote up
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 vote down vote up
/**
 * 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 vote down vote up
@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 vote down vote up
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 vote down vote up
@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 vote down vote up
@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 vote down vote up
@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 vote down vote up
/**
 * 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 vote down vote up
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 vote down vote up
public ApplicationMasterProtocolPBServiceImpl(ApplicationMasterProtocol impl) {
  this.real = impl;
}
 
Example #21
Source File: TestLocalContainerAllocator.java    From big-c with Apache License 2.0 4 votes vote down vote up
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 vote down vote up
@Override
protected ApplicationMasterProtocol createSchedulerProxy() {
  return scheduler;
}
 
Example #23
Source File: TestRMContainerAllocator.java    From big-c with Apache License 2.0 4 votes vote down vote up
@Override
protected ApplicationMasterProtocol createSchedulerProxy() {
  return this.rm.getApplicationMasterService();
}
 
Example #24
Source File: TestRMContainerAllocator.java    From big-c with Apache License 2.0 4 votes vote down vote up
@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 vote down vote up
@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 vote down vote up
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 vote down vote up
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 vote down vote up
@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 vote down vote up
@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 vote down vote up
@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());
}