Java Code Examples for org.apache.hadoop.yarn.api.records.Token

@Private
@Unstable
public static AllocateResponse newInstance(int responseId,
    List<ContainerStatus> completedContainers,
    List<Container> allocatedContainers, List<NodeReport> updatedNodes,
    Resource availResources, AMCommand command, int numClusterNodes,
    PreemptionMessage preempt, List<NMToken> nmTokens, Token amRMToken,
    List<ContainerResourceIncrease> increasedContainers,
    List<ContainerResourceDecrease> decreasedContainers) {
  AllocateResponse response =
      newInstance(responseId, completedContainers, allocatedContainers,
        updatedNodes, availResources, command, numClusterNodes, preempt,
        nmTokens, increasedContainers, decreasedContainers);
  response.setAMRMToken(amRMToken);
  return response;
}
 

/**
* Test CancelDelegationTokenRequestPBImpl.
* Test a transformation to prototype and back
*/
@Test
public void testCancelDelegationTokenRequestPBImpl() {

  Token token = getDelegationToken();

  CancelDelegationTokenRequestPBImpl original = new CancelDelegationTokenRequestPBImpl();
  original.setDelegationToken(token);
  CancelDelegationTokenRequestProto protoType = original.getProto();

  CancelDelegationTokenRequestPBImpl copy = new CancelDelegationTokenRequestPBImpl(protoType);
  assertNotNull(copy.getDelegationToken());
  //compare source and converted
  assertEquals(token, copy.getDelegationToken());

}
 

@Test
public void testResourceIncreaseContext() {
  byte[] identifier = new byte[] { 1, 2, 3, 4 };
  Token token = Token.newInstance(identifier, "", "".getBytes(), "");
  ContainerId containerId = ContainerId
      .newContainerId(ApplicationAttemptId.newInstance(
          ApplicationId.newInstance(1234, 3), 3), 7);
  Resource resource = Resource.newInstance(1023, 3);
  ContainerResourceIncrease ctx = ContainerResourceIncrease.newInstance(
      containerId, resource, token);

  // get proto and recover to ctx
  ContainerResourceIncreaseProto proto = 
      ((ContainerResourceIncreasePBImpl) ctx).getProto();
  ctx = new ContainerResourceIncreasePBImpl(proto);

  // check values
  Assert.assertEquals(ctx.getCapability(), resource);
  Assert.assertEquals(ctx.getContainerId(), containerId);
  Assert.assertTrue(Arrays.equals(ctx.getContainerToken().getIdentifier()
      .array(), identifier));
}
 

public NMToken createAndGetNMToken(String applicationSubmitter,
    ApplicationAttemptId appAttemptId, Container container) {
  try {
    this.readLock.lock();
    HashSet<NodeId> nodeSet = this.appAttemptToNodeKeyMap.get(appAttemptId);
    NMToken nmToken = null;
    if (nodeSet != null) {
      if (!nodeSet.contains(container.getNodeId())) {
        LOG.info("Sending NMToken for nodeId : " + container.getNodeId()
            + " for container : " + container.getId());
        Token token =
            createNMToken(container.getId().getApplicationAttemptId(),
              container.getNodeId(), applicationSubmitter);
        nmToken = NMToken.newInstance(container.getNodeId(), token);
        nodeSet.add(container.getNodeId());
      }
    }
    return nmToken;
  } finally {
    this.readLock.unlock();
  }
}
 

private String testStopContainer(YarnRPC rpc,
    ApplicationAttemptId appAttemptId, NodeId nodeId,
    ContainerId containerId, Token nmToken, boolean isExceptionExpected) {
  try {
    stopContainer(rpc, nmToken,
        Arrays.asList(new ContainerId[] { containerId }), appAttemptId,
        nodeId);
    if (isExceptionExpected) {
      fail("Exception was expected!!");
    }
    return "";
  } catch (Exception e) {
    e.printStackTrace();
    return e.getMessage();
  }
}
 

private void startContainer(final YarnRPC rpc,
    org.apache.hadoop.yarn.api.records.Token nmToken,
    org.apache.hadoop.yarn.api.records.Token containerToken,
    NodeId nodeId, String user) throws Exception {

  ContainerLaunchContext context =
      Records.newRecord(ContainerLaunchContext.class);
  StartContainerRequest scRequest =
      StartContainerRequest.newInstance(context,containerToken);
  List<StartContainerRequest> list = new ArrayList<StartContainerRequest>();
  list.add(scRequest);
  StartContainersRequest allRequests =
      StartContainersRequest.newInstance(list);
  ContainerManagementProtocol proxy = null;
  try {
    proxy = getContainerManagementProtocolProxy(rpc, nmToken, nodeId, user);
    StartContainersResponse response = proxy.startContainers(allRequests);
    for(SerializedException ex : response.getFailedRequests().values()){
      parseAndThrowException(ex.deSerialize());
    }
  } finally {
    if (proxy != null) {
      rpc.stopProxy(proxy, conf);
    }
  }
}
 

private String testStartContainer(YarnRPC rpc,
    ApplicationAttemptId appAttemptId, NodeId nodeId,
    org.apache.hadoop.yarn.api.records.Token containerToken,
    org.apache.hadoop.yarn.api.records.Token nmToken,
    boolean isExceptionExpected) {
  try {
    startContainer(rpc, nmToken, containerToken, nodeId,
        appAttemptId.toString());
    if (isExceptionExpected){
      fail("Exception was expected!!");        
    }
    return "";
  } catch (Exception e) {
    e.printStackTrace();
    return e.getMessage();
  }
}
 

/**
* Test RenewDelegationTokenRequestPBImpl.
* Test a transformation to prototype and back
*/

@Test
public void testRenewDelegationTokenRequestPBImpl() {

  Token token = getDelegationToken();

  RenewDelegationTokenRequestPBImpl original = new RenewDelegationTokenRequestPBImpl();
  original.setDelegationToken(token);
  RenewDelegationTokenRequestProto protoType = original.getProto();

  RenewDelegationTokenRequestPBImpl copy = new RenewDelegationTokenRequestPBImpl(protoType);
  assertNotNull(copy.getDelegationToken());
  //compare source and converted
  assertEquals(token, copy.getDelegationToken());

}
 

public NMToken createAndGetNMToken(String applicationSubmitter,
    ApplicationAttemptId appAttemptId, Container container) {
  try {
    this.readLock.lock();
    HashSet<NodeId> nodeSet = this.appAttemptToNodeKeyMap.get(appAttemptId);
    NMToken nmToken = null;
    if (nodeSet != null) {
      if (!nodeSet.contains(container.getNodeId())) {
        LOG.info("Sending NMToken for nodeId : " + container.getNodeId()
            + " for container : " + container.getId());
        Token token =
            createNMToken(container.getId().getApplicationAttemptId(),
              container.getNodeId(), applicationSubmitter);
        nmToken = NMToken.newInstance(container.getNodeId(), token);
        nodeSet.add(container.getNodeId());
      }
    }
    return nmToken;
  } finally {
    this.readLock.unlock();
  }
}
 

private String testStopContainer(YarnRPC rpc,
    ApplicationAttemptId appAttemptId, NodeId nodeId,
    ContainerId containerId, Token nmToken, boolean isExceptionExpected) {
  try {
    stopContainer(rpc, nmToken,
        Arrays.asList(new ContainerId[] { containerId }), appAttemptId,
        nodeId);
    if (isExceptionExpected) {
      fail("Exception was expected!!");
    }
    return "";
  } catch (Exception e) {
    e.printStackTrace();
    return e.getMessage();
  }
}
 

@Override
public Token getDelegationToken() {
  CancelDelegationTokenRequestProtoOrBuilder p = viaProto ? proto : builder;
  if (this.token != null) {
    return this.token;
  }
  this.token = convertFromProtoFormat(p.getToken());
  return this.token;
}
 

protected void allocateContainersAndValidateNMTokens(MockAM am,
    ArrayList<Container> containersReceived, int totalContainerRequested,
    HashMap<String, Token> nmTokens, MockNM nm) throws Exception,
    InterruptedException {
  ArrayList<ContainerId> releaseContainerList = new ArrayList<ContainerId>();
  AllocateResponse response;
  ArrayList<ResourceRequest> resourceRequest =
      new ArrayList<ResourceRequest>();      
  while (containersReceived.size() < totalContainerRequested) {
    nm.nodeHeartbeat(true);
    LOG.info("requesting containers..");
    response =
        am.allocate(resourceRequest, releaseContainerList);
    containersReceived.addAll(response.getAllocatedContainers());
    if (!response.getNMTokens().isEmpty()) {
      for (NMToken nmToken : response.getNMTokens()) {
        String nodeId = nmToken.getNodeId().toString();
        if (nmTokens.containsKey(nodeId)) {
          Assert.fail("Duplicate NMToken received for : " + nodeId);
        }
        nmTokens.put(nodeId, nmToken.getToken());
      }
    }
    LOG.info("Got " + containersReceived.size()
        + " containers. Waiting to get " + totalContainerRequested);
    Thread.sleep(WAIT_SLEEP_MS);
  }
}
 

@Test
public void testErrorInStartContainerShouldUpdateState() {
  // create mocks
  final int samzaContainerId = 1;
  YarnConfiguration yarnConfiguration = mock(YarnConfiguration.class);
  SamzaAppMasterMetrics metrics = mock(SamzaAppMasterMetrics.class);
  Config config = mock(Config.class);
  AMRMClientAsync asyncClient = mock(AMRMClientAsync.class);
  YarnAppState yarnAppState = new YarnAppState(0, mock(ContainerId.class), "host", 8080, 8081);
  SamzaYarnAppMasterLifecycle lifecycle = mock(SamzaYarnAppMasterLifecycle.class);
  SamzaYarnAppMasterService service = mock(SamzaYarnAppMasterService.class);
  NMClientAsync asyncNMClient = mock(NMClientAsync.class);
  ClusterResourceManager.Callback callback = mock(ClusterResourceManager.Callback.class);

  // start the cluster manager
  YarnClusterResourceManager yarnClusterResourceManager = new YarnClusterResourceManager(asyncClient, asyncNMClient,
      callback, yarnAppState, lifecycle, service, metrics, yarnConfiguration, config);

  yarnAppState.pendingProcessors.put(String.valueOf(samzaContainerId),
      new YarnContainer(Container.newInstance(
          ContainerId.newContainerId(
              ApplicationAttemptId.newInstance(
                  ApplicationId.newInstance(10000L, 1), 1), 1),
          NodeId.newInstance("host1", 8088), "http://host1",
          Resource.newInstance(1024, 1), Priority.newInstance(1),
          Token.newInstance("id".getBytes(), "read", "password".getBytes(), "service"))));

  yarnClusterResourceManager.start();
  assertEquals(1, yarnAppState.pendingProcessors.size());

  yarnClusterResourceManager.onStartContainerError(ContainerId.newContainerId(
      ApplicationAttemptId.newInstance(
          ApplicationId.newInstance(10000L, 1), 1), 1),
      new Exception());

  assertEquals(0, yarnAppState.pendingProcessors.size());
  verify(callback, times(1)).onStreamProcessorLaunchFailure(anyObject(), any(Exception.class));
}
 

@Override
public Token getDelegationToken() {
  RenewDelegationTokenRequestProtoOrBuilder p = viaProto ? proto : builder;
  if (this.token != null) {
    return this.token;
  }
  this.token = convertFromProtoFormat(p.getToken());
  return this.token;
}
 

private void updateAMRMToken(Token token) throws IOException {
  org.apache.hadoop.security.token.Token<AMRMTokenIdentifier> amrmToken =
      new org.apache.hadoop.security.token.Token<AMRMTokenIdentifier>(token
        .getIdentifier().array(), token.getPassword().array(), new Text(
        token.getKind()), new Text(token.getService()));
  UserGroupInformation currentUGI = UserGroupInformation.getCurrentUser();
  currentUGI.addToken(amrmToken);
  amrmToken.setService(ClientRMProxy.getAMRMTokenService(getConfig()));
}
 

@Override
public Token getContainerToken() {
  StartContainerRequestProtoOrBuilder p = viaProto ? proto : builder;
  if (this.containerToken != null) {
    return this.containerToken;
  }
  if (!p.hasContainerToken()) {
    return null;
  }
  this.containerToken = convertFromProtoFormat(p.getContainerToken());
  return this.containerToken;
}
 

private void updateAMRMToken(Token token) throws IOException {
  org.apache.hadoop.security.token.Token<AMRMTokenIdentifier> amrmToken =
      new org.apache.hadoop.security.token.Token<AMRMTokenIdentifier>(token
        .getIdentifier().array(), token.getPassword().array(), new Text(
        token.getKind()), new Text(token.getService()));
  UserGroupInformation currentUGI = UserGroupInformation.getCurrentUser();
  currentUGI.addToken(amrmToken);
  amrmToken.setService(ClientRMProxy.getAMRMTokenService(getConfig()));
}
 

public static Token newContainerToken(NodeId nodeId,
    byte[] password, ContainerTokenIdentifier tokenIdentifier) {
  // RPC layer client expects ip:port as service for tokens
  InetSocketAddress addr =
      NetUtils.createSocketAddrForHost(nodeId.getHost(), nodeId.getPort());
  // NOTE: use SecurityUtil.setTokenService if this becomes a "real" token
  Token containerToken =
      Token.newInstance(tokenIdentifier.getBytes(),
        ContainerTokenIdentifier.KIND.toString(), password, SecurityUtil
          .buildTokenService(addr).toString());
  return containerToken;
}
 

public static Token newInstance(byte[] password,
    NMTokenIdentifier identifier) {
  NodeId nodeId = identifier.getNodeId();
  // RPC layer client expects ip:port as service for tokens
  InetSocketAddress addr =
      NetUtils.createSocketAddrForHost(nodeId.getHost(), nodeId.getPort());
  Token nmToken =
      Token.newInstance(identifier.getBytes(),
        NMTokenIdentifier.KIND.toString(), password, SecurityUtil
          .buildTokenService(addr).toString());
  return nmToken;
}
 

@Override
public Token getDelegationToken() {
  RenewDelegationTokenRequestProtoOrBuilder p = viaProto ? proto : builder;
  if (this.token != null) {
    return this.token;
  }
  this.token = convertFromProtoFormat(p.getToken());
  return this.token;
}
 

@Override
public void setContainerToken(Token containerToken) {
  maybeInitBuilder();
  if (containerToken == null) 
    builder.clearContainerToken();
  this.containerToken = containerToken;
}
 

@Override
public synchronized void setToken(Token token) {
  maybeInitBuilder();
  if (token == null) {
    builder.clearToken();
  }
  this.token = token;
}
 

@Override
public void handle(ContainerAllocatorEvent event) {
  ContainerId cId =
      ContainerId.newContainerId(getContext().getApplicationAttemptId(),
        containerCount++);
  NodeId nodeId = NodeId.newInstance(NM_HOST, NM_PORT);
  Resource resource = Resource.newInstance(1234, 2);
  ContainerTokenIdentifier containerTokenIdentifier =
      new ContainerTokenIdentifier(cId, nodeId.toString(), "user",
      resource, System.currentTimeMillis() + 10000, 42, 42,
      Priority.newInstance(0), 0);
  Token containerToken = newContainerToken(nodeId, "password".getBytes(),
        containerTokenIdentifier);
  Container container = Container.newInstance(cId, nodeId,
      NM_HOST + ":" + NM_HTTP_PORT, resource, null, containerToken);
  JobID id = TypeConverter.fromYarn(applicationId);
  JobId jobId = TypeConverter.toYarn(id);
  getContext().getEventHandler().handle(new JobHistoryEvent(jobId, 
      new NormalizedResourceEvent(
          org.apache.hadoop.mapreduce.TaskType.REDUCE,
      100)));
  getContext().getEventHandler().handle(new JobHistoryEvent(jobId, 
      new NormalizedResourceEvent(
          org.apache.hadoop.mapreduce.TaskType.MAP,
      100)));
  getContext().getEventHandler().handle(
      new TaskAttemptContainerAssignedEvent(event.getAttemptID(),
          container, null));
}
 

@Override
public Token getRMDelegationToken() {
  GetDelegationTokenResponseProtoOrBuilder p = viaProto ? proto : builder;
  if (this.appToken != null) {
    return this.appToken;
  }
  if (!p.hasToken()) {
    return null;
  }
  this.appToken = convertFromProtoFormat(p.getToken());
  return this.appToken;  
}
 

@Override
public Token getContainerToken() {
  StartContainerRequestProtoOrBuilder p = viaProto ? proto : builder;
  if (this.containerToken != null) {
    return this.containerToken;
  }
  if (!p.hasContainerToken()) {
    return null;
  }
  this.containerToken = convertFromProtoFormat(p.getContainerToken());
  return this.containerToken;
}
 

@Override
public Token getAMRMToken() {
  ApplicationReportProtoOrBuilder p = viaProto ? proto : builder;
  if (amRmToken != null) {
    return amRmToken;
  }
  if (!p.hasAmRmToken()) {
    return null;
  }
  amRmToken = convertFromProtoFormat(p.getAmRmToken());
  return amRmToken;
}
 

public ContainerLauncherEvent(ContainerId containerId, NodeId nodeId,
                              Token containerToken, ContainerLauncherEventType type,
                              int launcherId,
                              int schedulerId, int taskCommId) {
  super(type);
  this.containerId = containerId;
  this.nodeId = nodeId;
  this.containerToken = containerToken;
  this.launcherId = launcherId;
  this.schedulerId = schedulerId;
  this.taskCommId = taskCommId;
}
 

protected Token createContainerToken(ContainerId cId, Priority priority,
    long createTime) throws InvalidToken {
  Resource r = BuilderUtils.newResource(1024, 1);
  ContainerTokenIdentifier containerTokenIdentifier =
      new ContainerTokenIdentifier(cId, context.getNodeId().toString(), user,
        r, System.currentTimeMillis() + 10000L, 123, DUMMY_RM_IDENTIFIER,
        priority, createTime);
  Token containerToken =
      BuilderUtils.newContainerToken(
        context.getNodeId(),
        context.getContainerTokenSecretManager().retrievePassword(
          containerTokenIdentifier), containerTokenIdentifier);
  return containerToken;
}
 

@VisibleForTesting
org.apache.hadoop.security.token.Token<TimelineDelegationTokenIdentifier>
    getTimelineDelegationToken() throws IOException, YarnException {
      try {
        return timelineClient.getDelegationToken(timelineDTRenewer);
      } catch (Exception e ) {
        if (timelineServiceBestEffort) {
          LOG.warn("Failed to get delegation token from the timeline server: "
              + e.getMessage());
          return null;
        }
        throw e;
      }
}
 

public static ContainerTokenIdentifier newContainerTokenIdentifier(
    Token containerToken) throws IOException {
  org.apache.hadoop.security.token.Token<ContainerTokenIdentifier> token =
      new org.apache.hadoop.security.token.Token<ContainerTokenIdentifier>(
          containerToken.getIdentifier()
              .array(), containerToken.getPassword().array(), new Text(
              containerToken.getKind()),
          new Text(containerToken.getService()));
  return token.decodeIdentifier();
}