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

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();
  }
}
 

@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 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());

}
 

@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));
}
 

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();
  }
}
 

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();
  }
}
 

/**
* 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());

}
 

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

@Override
public StartContainersResponse
    startContainers(StartContainersRequest requests)
        throws YarnException {
  StartContainerRequest request = requests.getStartContainerRequests().get(0);
  LOG.info("Container started by MyContainerManager: " + request);
  launched = true;
  Map<String, String> env =
      request.getContainerLaunchContext().getEnvironment();

  Token containerToken = request.getContainerToken();
  ContainerTokenIdentifier tokenId = null;

  try {
    tokenId = BuilderUtils.newContainerTokenIdentifier(containerToken);
  } catch (IOException e) {
    throw RPCUtil.getRemoteException(e);
  }

  ContainerId containerId = tokenId.getContainerID();
  containerIdAtContainerManager = containerId.toString();
  attemptIdAtContainerManager =
      containerId.getApplicationAttemptId().toString();
  nmHostAtContainerManager = tokenId.getNmHostAddress();
  submitTimeAtContainerManager =
      Long.parseLong(env.get(ApplicationConstants.APP_SUBMIT_TIME_ENV));
  maxAppAttempts =
      Integer.parseInt(env.get(ApplicationConstants.MAX_APP_ATTEMPTS_ENV));
  return StartContainersResponse.newInstance(
    new HashMap<String, ByteBuffer>(), new ArrayList<ContainerId>(),
    new HashMap<ContainerId, SerializedException>());
}
 

public ContainerEvent(ContainerId containerId, NodeId nodeId,
    Token containerToken, ContainerEventType type) {
  super(type);
  this.containerId = containerId;
  this.nodeId = nodeId;
  this.containerToken = containerToken;
}
 

public static Token createContainerToken(ContainerId cId, long rmIdentifier,
    NodeId nodeId, String user,
    NMContainerTokenSecretManager containerTokenSecretManager)
    throws IOException {
  return createContainerToken(cId, rmIdentifier, nodeId, user,
    containerTokenSecretManager, null);
}
 

@Private
@Unstable
public static GetDelegationTokenResponse newInstance(Token rmDTToken) {
  GetDelegationTokenResponse response =
      Records.newRecord(GetDelegationTokenResponse.class);
  response.setRMDelegationToken(rmDTToken);
  return response;
}
 

public ContainerStopRequest(NodeId nodeId,
                            ContainerId containerId,
                            Token containerToken, String schedulerName, String taskCommName) {
  super(nodeId, containerId, containerToken);
  this.schedulerName = schedulerName;
  this.taskCommName = taskCommName;
}
 

private Container mockContainer(int i) {
  ApplicationId appId =
      ApplicationId.newInstance(System.currentTimeMillis(), 1);
  ApplicationAttemptId attemptId =
      ApplicationAttemptId.newInstance(appId, 1);
  ContainerId containerId = ContainerId.newContainerId(attemptId, i);
  nodeId = NodeId.newInstance("localhost", 0);
  // Create an empty record
  containerToken = recordFactory.newRecordInstance(Token.class);
  return Container.newInstance(containerId, nodeId, null, null, null,
    containerToken);
}
 

public ContainerLauncherOperationBase(NodeId nodeId,
                                      ContainerId containerId,
                                      Token containerToken) {
  this.nodeId = nodeId;
  this.containerId = containerId;
  this.containerToken = containerToken;
}
 

private Container mockContainer(int i) {
  ApplicationId appId =
      ApplicationId.newInstance(System.currentTimeMillis(), 1);
  ApplicationAttemptId attemptId =
      ApplicationAttemptId.newInstance(appId, 1);
  ContainerId containerId = ContainerId.newContainerId(attemptId, i);
  nodeId = NodeId.newInstance("localhost", 0);
  // Create an empty record
  containerToken = recordFactory.newRecordInstance(Token.class);
  return Container.newInstance(containerId, nodeId, null, null, null,
    containerToken);
}
 

@Private
@Unstable
public static GetDelegationTokenResponse newInstance(Token rmDTToken) {
  GetDelegationTokenResponse response =
      Records.newRecord(GetDelegationTokenResponse.class);
  response.setRMDelegationToken(rmDTToken);
  return response;
}
 

@Override
public void setClientToAMToken(Token clientToAMToken) {
  maybeInitBuilder();
  if (clientToAMToken == null) 
    builder.clearClientToAmToken();
  this.clientToAMToken = clientToAMToken;
}
 

@Override
public synchronized Token getToken() {
  NMTokenProtoOrBuilder p = viaProto ? proto : builder;
  if (this.token != null) {
    return this.token;
  }
  if (!p.hasToken()) {
    return null;
  }
  this.token = convertFromProtoFormat(p.getToken()); 
  return 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()));
  // Preserve the token service sent by the RM when adding the token
  // to ensure we replace the previous token setup by the RM.
  // Afterwards we can update the service address for the RPC layer.
  UserGroupInformation currentUGI = UserGroupInformation.getCurrentUser();
  currentUGI.addToken(amrmToken);
  amrmToken.setService(ClientRMProxy.getAMRMTokenService(getConfig()));
}
 

public NMCommunicatorEvent(ContainerId containerId, NodeId nodeId,
    Token containerToken, NMCommunicatorEventType type) {
  super(type);
  this.containerId = containerId;
  this.nodeId = nodeId;
  this.containerToken = containerToken;
}
 

@Override
public Token getContainerToken() {
  ContainerProtoOrBuilder 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 getDelegationToken() {
  RenewDelegationTokenRequestProtoOrBuilder p = viaProto ? proto : builder;
  if (this.token != null) {
    return this.token;
  }
  this.token = convertFromProtoFormat(p.getToken());
  return this.token;
}
 

@Public
@Stable
public static StartContainerRequest newInstance(
    ContainerLaunchContext context, Token container) {
  StartContainerRequest request =
      Records.newRecord(StartContainerRequest.class);
  request.setContainerLaunchContext(context);
  request.setContainerToken(container);
  return request;
}
 

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

@Override
public StartContainersResponse
    startContainers(StartContainersRequest requests)
        throws YarnException {
  StartContainerRequest request = requests.getStartContainerRequests().get(0);
  LOG.info("Container started by MyContainerManager: " + request);
  launched = true;
  Map<String, String> env =
      request.getContainerLaunchContext().getEnvironment();

  Token containerToken = request.getContainerToken();
  ContainerTokenIdentifier tokenId = null;

  try {
    tokenId = BuilderUtils.newContainerTokenIdentifier(containerToken);
  } catch (IOException e) {
    throw RPCUtil.getRemoteException(e);
  }

  ContainerId containerId = tokenId.getContainerID();
  containerIdAtContainerManager = containerId.toString();
  attemptIdAtContainerManager =
      containerId.getApplicationAttemptId().toString();
  nmHostAtContainerManager = tokenId.getNmHostAddress();
  submitTimeAtContainerManager =
      Long.parseLong(env.get(ApplicationConstants.APP_SUBMIT_TIME_ENV));
  maxAppAttempts =
      Integer.parseInt(env.get(ApplicationConstants.MAX_APP_ATTEMPTS_ENV));
  return StartContainersResponse.newInstance(
    new HashMap<String, ByteBuffer>(), new ArrayList<ContainerId>(),
    new HashMap<ContainerId, SerializedException>());
}
 

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();
}