MockNM.java
/**
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.hadoop.yarn.server.resourcemanager;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import org.apache.hadoop.yarn.api.records.ApplicationAttemptId;
import org.apache.hadoop.yarn.api.records.ApplicationId;
import org.apache.hadoop.yarn.api.records.Container;
import org.apache.hadoop.yarn.api.records.ContainerId;
import org.apache.hadoop.yarn.api.records.ContainerState;
import org.apache.hadoop.yarn.api.records.ContainerStatus;
import org.apache.hadoop.yarn.api.records.NodeId;
import org.apache.hadoop.yarn.api.records.NodeLabel;
import org.apache.hadoop.yarn.api.records.Resource;
import org.apache.hadoop.yarn.conf.YarnConfiguration;
import org.apache.hadoop.yarn.server.api.protocolrecords.NMContainerStatus;
import org.apache.hadoop.yarn.server.api.protocolrecords.NodeHeartbeatRequest;
import org.apache.hadoop.yarn.server.api.protocolrecords.NodeHeartbeatResponse;
import org.apache.hadoop.yarn.server.api.protocolrecords.RegisterNodeManagerRequest;
import org.apache.hadoop.yarn.server.api.protocolrecords.RegisterNodeManagerResponse;
import org.apache.hadoop.yarn.server.api.records.AppCollectorData;
import org.apache.hadoop.yarn.server.api.protocolrecords.UnRegisterNodeManagerRequest;
import org.apache.hadoop.yarn.server.api.records.MasterKey;
import org.apache.hadoop.yarn.server.api.records.NodeHealthStatus;
import org.apache.hadoop.yarn.server.api.records.NodeStatus;
import org.apache.hadoop.yarn.server.api.records.OpportunisticContainersStatus;
import org.apache.hadoop.yarn.server.utils.BuilderUtils;
import org.apache.hadoop.yarn.util.Records;
import org.apache.hadoop.yarn.util.YarnVersionInfo;
import org.apache.hadoop.yarn.util.resource.Resources;
import org.eclipse.jetty.util.log.Log;
public class MockNM {
private int responseId;
private NodeId nodeId;
private Resource capability;
private ResourceTrackerService resourceTracker;
private int httpPort = 2;
private MasterKey currentContainerTokenMasterKey;
private MasterKey currentNMTokenMasterKey;
private String version;
private Map<ContainerId, ContainerStatus> containerStats =
new HashMap<ContainerId, ContainerStatus>();
private Map<ApplicationId, AppCollectorData> registeringCollectors
= new ConcurrentHashMap<>();
private Set<NodeLabel> nodeLabels;
private long tokenSequenceNo;
public MockNM(String nodeIdStr, int memory, ResourceTrackerService resourceTracker) {
// scale vcores based on the requested memory
this(nodeIdStr, memory,
Math.max(1, (memory * YarnConfiguration.DEFAULT_NM_VCORES) /
YarnConfiguration.DEFAULT_NM_PMEM_MB),
resourceTracker);
}
public MockNM(String nodeIdStr, int memory, int vcores,
ResourceTrackerService resourceTracker) {
this(nodeIdStr, memory, vcores, resourceTracker,
YarnVersionInfo.getVersion());
}
public MockNM(String nodeIdStr, int memory, int vcores,
ResourceTrackerService resourceTracker, String version) {
this(nodeIdStr, Resource.newInstance(memory, vcores), resourceTracker,
version);
}
public MockNM(String nodeIdStr, Resource capability,
ResourceTrackerService resourceTracker) {
this(nodeIdStr, capability, resourceTracker,
YarnVersionInfo.getVersion());
}
public MockNM(String nodeIdStr, Resource capability,
ResourceTrackerService resourceTracker, String version) {
this.capability = capability;
this.resourceTracker = resourceTracker;
this.version = version;
String[] splits = nodeIdStr.split(":");
nodeId = BuilderUtils.newNodeId(splits[0], Integer.parseInt(splits[1]));
}
public MockNM(String nodeIdStr, Resource capability,
ResourceTrackerService resourceTracker, String version, Set<NodeLabel>
nodeLabels) {
this(nodeIdStr, capability, resourceTracker, version);
this.nodeLabels = nodeLabels;
}
public NodeId getNodeId() {
return nodeId;
}
public int getHttpPort() {
return httpPort;
}
public void setHttpPort(int port) {
httpPort = port;
}
public void setResourceTrackerService(ResourceTrackerService resourceTracker) {
this.resourceTracker = resourceTracker;
}
public void containerStatus(ContainerStatus containerStatus) throws Exception {
Map<ApplicationId, List<ContainerStatus>> conts =
new HashMap<ApplicationId, List<ContainerStatus>>();
conts.put(containerStatus.getContainerId().getApplicationAttemptId().getApplicationId(),
Arrays.asList(new ContainerStatus[] { containerStatus }));
nodeHeartbeat(conts, true);
}
public void containerIncreaseStatus(Container container) throws Exception {
ContainerStatus containerStatus = BuilderUtils.newContainerStatus(
container.getId(), ContainerState.RUNNING, "Success", 0,
container.getResource());
List<Container> increasedConts = Collections.singletonList(container);
nodeHeartbeat(Collections.singletonList(containerStatus), increasedConts,
true, responseId);
}
public void addRegisteringCollector(ApplicationId appId,
AppCollectorData data) {
this.registeringCollectors.put(appId, data);
}
public Map<ApplicationId, AppCollectorData> getRegisteringCollectors() {
return this.registeringCollectors;
}
public void unRegisterNode() throws Exception {
UnRegisterNodeManagerRequest request = Records
.newRecord(UnRegisterNodeManagerRequest.class);
request.setNodeId(nodeId);
resourceTracker.unRegisterNodeManager(request);
}
public RegisterNodeManagerResponse registerNode() throws Exception {
return registerNode(null, null);
}
public RegisterNodeManagerResponse registerNode(
List<ApplicationId> runningApplications) throws Exception {
return registerNode(null, runningApplications);
}
public RegisterNodeManagerResponse registerNode(
List<NMContainerStatus> containerReports,
List<ApplicationId> runningApplications) throws Exception {
RegisterNodeManagerRequest req = Records.newRecord(
RegisterNodeManagerRequest.class);
req.setNodeId(nodeId);
req.setHttpPort(httpPort);
req.setResource(capability);
req.setContainerStatuses(containerReports);
req.setNMVersion(version);
req.setRunningApplications(runningApplications);
if ( nodeLabels != null && nodeLabels.size() > 0) {
req.setNodeLabels(nodeLabels);
}
NodeStatus status = Records.newRecord(NodeStatus.class);
status.setResponseId(0);
status.setNodeId(nodeId);
status.setContainersStatuses(new ArrayList<>(containerStats.values()));
NodeHealthStatus healthStatus = Records.newRecord(NodeHealthStatus.class);
healthStatus.setHealthReport("");
healthStatus.setIsNodeHealthy(true);
healthStatus.setLastHealthReportTime(1);
status.setNodeHealthStatus(healthStatus);
req.setNodeStatus(status);
RegisterNodeManagerResponse registrationResponse =
resourceTracker.registerNodeManager(req);
this.currentContainerTokenMasterKey =
registrationResponse.getContainerTokenMasterKey();
this.currentNMTokenMasterKey = registrationResponse.getNMTokenMasterKey();
Resource newResource = registrationResponse.getResource();
if (newResource != null) {
capability = Resources.clone(newResource);
}
containerStats.clear();
if (containerReports != null) {
for (NMContainerStatus report : containerReports) {
if (report.getContainerState() != ContainerState.COMPLETE) {
containerStats.put(report.getContainerId(),
ContainerStatus.newInstance(report.getContainerId(),
report.getContainerState(), report.getDiagnostics(),
report.getContainerExitStatus()));
}
}
}
responseId = 0;
return registrationResponse;
}
public NodeHeartbeatResponse nodeHeartbeat(boolean isHealthy) throws Exception {
return nodeHeartbeat(Collections.<ContainerStatus>emptyList(),
Collections.<Container>emptyList(), isHealthy, responseId);
}
public NodeHeartbeatResponse nodeHeartbeat(ApplicationAttemptId attemptId,
long containerId, ContainerState containerState) throws Exception {
ContainerStatus containerStatus = BuilderUtils.newContainerStatus(
BuilderUtils.newContainerId(attemptId, containerId), containerState,
"Success", 0, capability);
ArrayList<ContainerStatus> containerStatusList =
new ArrayList<ContainerStatus>(1);
containerStatusList.add(containerStatus);
Log.getLog().info("ContainerStatus: " + containerStatus);
return nodeHeartbeat(containerStatusList,
Collections.<Container>emptyList(), true, responseId);
}
public NodeHeartbeatResponse nodeHeartbeat(Map<ApplicationId,
List<ContainerStatus>> conts, boolean isHealthy) throws Exception {
return nodeHeartbeat(conts, isHealthy, responseId);
}
/**
* Sends the heartbeat of the node.
* @param isHealthy whether node is healthy.
* @param resId response id.
* @return response of the heartbeat.
* @throws Exception
*/
public NodeHeartbeatResponse nodeHeartbeat(Map<ApplicationId,
List<ContainerStatus>> conts, boolean isHealthy, int resId) throws Exception {
ArrayList<ContainerStatus> updatedStats = new ArrayList<ContainerStatus>();
for (List<ContainerStatus> stats : conts.values()) {
updatedStats.addAll(stats);
}
return nodeHeartbeat(updatedStats, Collections.<Container>emptyList(),
isHealthy, resId);
}
/**
* Sends the heartbeat of the node.
* @param updatedStats containers with updated status.
* @param isHealthy whether node is healthy.
* @return response of the heartbeat.
* @throws Exception
*/
public NodeHeartbeatResponse nodeHeartbeat(
List<ContainerStatus> updatedStats, boolean isHealthy) throws Exception {
return nodeHeartbeat(updatedStats, Collections.<Container>emptyList(),
isHealthy, responseId);
}
/**
* Sends the heartbeat of the node.
* @param oppContainersStatus opportunistic containers status.
* @param isHealthy whether node is healthy.
* @return response of the heartbeat.
* @throws Exception
*/
public NodeHeartbeatResponse nodeHeartbeat(
OpportunisticContainersStatus oppContainersStatus, boolean isHealthy)
throws Exception {
return nodeHeartbeat(Collections.emptyList(),
Collections.emptyList(), isHealthy, responseId, oppContainersStatus);
}
/**
* Sends the heartbeat of the node.
* @param updatedStats containers with updated status.
* @param increasedConts containers whose resource has been increased.
* @param isHealthy whether node is healthy.
* @param resId response id.
* @return response of the heartbeat.
* @throws Exception
*/
public NodeHeartbeatResponse nodeHeartbeat(
List<ContainerStatus> updatedStats, List<Container> increasedConts,
boolean isHealthy, int resId) throws Exception {
return nodeHeartbeat(updatedStats, increasedConts,
isHealthy, resId, null);
}
/**
* Sends the heartbeat of the node.
* @param updatedStats containers with updated status.
* @param increasedConts containers whose resource has been increased.
* @param isHealthy whether node is healthy.
* @param resId response id.
* @param oppContainersStatus opportunistic containers status.
* @return response of the heartbeat.
* @throws Exception
*/
public NodeHeartbeatResponse nodeHeartbeat(List<ContainerStatus> updatedStats,
List<Container> increasedConts, boolean isHealthy, int resId,
OpportunisticContainersStatus oppContainersStatus) throws Exception {
NodeHeartbeatRequest req = Records.newRecord(NodeHeartbeatRequest.class);
NodeStatus status = Records.newRecord(NodeStatus.class);
status.setResponseId(resId);
status.setNodeId(nodeId);
ArrayList<ContainerId> completedContainers = new ArrayList<ContainerId>();
for (ContainerStatus stat : updatedStats) {
if (stat.getState() == ContainerState.COMPLETE) {
completedContainers.add(stat.getContainerId());
}
containerStats.put(stat.getContainerId(), stat);
}
status.setContainersStatuses(
new ArrayList<ContainerStatus>(containerStats.values()));
for (ContainerId cid : completedContainers) {
containerStats.remove(cid);
}
status.setIncreasedContainers(increasedConts);
status.setOpportunisticContainersStatus(oppContainersStatus);
NodeHealthStatus healthStatus = Records.newRecord(NodeHealthStatus.class);
healthStatus.setHealthReport("");
healthStatus.setIsNodeHealthy(isHealthy);
healthStatus.setLastHealthReportTime(1);
status.setNodeHealthStatus(healthStatus);
req.setNodeStatus(status);
req.setLastKnownContainerTokenMasterKey(this.currentContainerTokenMasterKey);
req.setLastKnownNMTokenMasterKey(this.currentNMTokenMasterKey);
req.setRegisteringCollectors(this.registeringCollectors);
req.setTokenSequenceNo(this.tokenSequenceNo);
NodeHeartbeatResponse heartbeatResponse =
resourceTracker.nodeHeartbeat(req);
responseId = heartbeatResponse.getResponseId();
MasterKey masterKeyFromRM = heartbeatResponse.getContainerTokenMasterKey();
if (masterKeyFromRM != null
&& masterKeyFromRM.getKeyId() != this.currentContainerTokenMasterKey
.getKeyId()) {
this.currentContainerTokenMasterKey = masterKeyFromRM;
}
masterKeyFromRM = heartbeatResponse.getNMTokenMasterKey();
if (masterKeyFromRM != null
&& masterKeyFromRM.getKeyId() != this.currentNMTokenMasterKey
.getKeyId()) {
this.currentNMTokenMasterKey = masterKeyFromRM;
}
Resource newResource = heartbeatResponse.getResource();
if (newResource != null) {
capability = Resources.clone(newResource);
}
this.tokenSequenceNo = heartbeatResponse.getTokenSequenceNo();
return heartbeatResponse;
}
public static NodeStatus createMockNodeStatus() {
NodeStatus mockNodeStatus = mock(NodeStatus.class);
NodeHealthStatus mockNodeHealthStatus = mock(NodeHealthStatus.class);
when(mockNodeStatus.getNodeHealthStatus()).thenReturn(mockNodeHealthStatus);
when(mockNodeHealthStatus.getIsNodeHealthy()).thenReturn(true);
return mockNodeStatus;
}
public long getMemory() {
return capability.getMemorySize();
}
public int getvCores() {
return capability.getVirtualCores();
}
public Resource getCapability() {
return capability;
}
public String getVersion() {
return version;
}
public void setResponseId(int id) {
this.responseId = id;
}
}