| /* |
| * 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.service; |
| |
| import com.google.common.base.Supplier; |
| import com.google.common.collect.Lists; |
| import org.apache.hadoop.fs.Path; |
| import org.apache.hadoop.io.DataOutputBuffer; |
| import org.apache.hadoop.registry.client.api.RegistryOperations; |
| import org.apache.hadoop.registry.client.binding.RegistryPathUtils; |
| import org.apache.hadoop.registry.client.types.ServiceRecord; |
| import org.apache.hadoop.registry.client.types.yarn.PersistencePolicies; |
| import org.apache.hadoop.registry.client.types.yarn.YarnRegistryAttributes; |
| import org.apache.hadoop.security.Credentials; |
| import org.apache.hadoop.test.GenericTestUtils; |
| import org.apache.hadoop.yarn.api.protocolrecords.AllocateResponse; |
| import org.apache.hadoop.yarn.api.protocolrecords.RegisterApplicationMasterResponse; |
| 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.FinalApplicationStatus; |
| import org.apache.hadoop.yarn.api.records.NodeId; |
| import org.apache.hadoop.yarn.api.records.Priority; |
| import org.apache.hadoop.yarn.api.records.Resource; |
| import org.apache.hadoop.yarn.client.api.AMRMClient; |
| import org.apache.hadoop.yarn.client.api.NMClient; |
| import org.apache.hadoop.yarn.client.api.async.AMRMClientAsync; |
| import org.apache.hadoop.yarn.client.api.async.NMClientAsync; |
| import org.apache.hadoop.yarn.client.api.impl.AMRMClientImpl; |
| import org.apache.hadoop.yarn.exceptions.YarnException; |
| import org.apache.hadoop.yarn.proto.ClientAMProtocol; |
| import org.apache.hadoop.yarn.service.api.records.Service; |
| import org.apache.hadoop.yarn.service.component.Component; |
| import org.apache.hadoop.yarn.service.component.ComponentState; |
| import org.apache.hadoop.yarn.service.component.instance.ComponentInstance; |
| import org.apache.hadoop.yarn.service.component.instance.ComponentInstanceState; |
| import org.apache.hadoop.yarn.service.exceptions.BadClusterStateException; |
| import org.apache.hadoop.yarn.service.registry.YarnRegistryViewForProviders; |
| import org.apache.hadoop.yarn.service.utils.SliderFileSystem; |
| import org.apache.hadoop.yarn.util.Records; |
| import org.apache.hadoop.yarn.util.resource.ResourceUtils; |
| import org.slf4j.Logger; |
| import org.slf4j.LoggerFactory; |
| |
| import java.io.IOException; |
| import java.nio.ByteBuffer; |
| import java.util.Collections; |
| import java.util.Iterator; |
| import java.util.LinkedList; |
| import java.util.List; |
| import java.util.Map; |
| import java.util.concurrent.ConcurrentHashMap; |
| import java.util.concurrent.TimeoutException; |
| |
| import static org.mockito.Matchers.anyObject; |
| import static org.mockito.Matchers.anyString; |
| import static org.mockito.Mockito.mock; |
| import static org.mockito.Mockito.when; |
| |
| public class MockServiceAM extends ServiceMaster { |
| |
| private static final Logger LOG = |
| LoggerFactory.getLogger(MockServiceAM.class); |
| |
| Service service; |
| // The list of containers fed by tests to be returned on |
| // AMRMClientCallBackHandler#onContainersAllocated |
| final List<Container> feedContainers = |
| Collections.synchronizedList(new LinkedList<>()); |
| |
| final List<ContainerStatus> failedContainers = |
| Collections.synchronizedList(new LinkedList<>()); |
| |
| private final List<Container> recoveredContainers = |
| Collections.synchronizedList(new LinkedList<>()); |
| |
| private final Map<String, ServiceRecord> registryComponents = |
| new ConcurrentHashMap<>(); |
| |
| private Map<ContainerId, ContainerStatus> containerStatuses = |
| new ConcurrentHashMap<>(); |
| |
| private Credentials amCreds; |
| |
| public MockServiceAM(Service service) { |
| super(service.getName()); |
| this.service = service; |
| } |
| |
| public MockServiceAM(Service service, Credentials amCreds) { |
| super(service.getName()); |
| this.service = service; |
| this.amCreds = amCreds; |
| } |
| |
| @Override |
| protected ContainerId getAMContainerId() |
| throws BadClusterStateException { |
| return ContainerId.newContainerId(ApplicationAttemptId |
| .newInstance(ApplicationId.fromString(service.getId()), 1), 1); |
| } |
| |
| @Override |
| protected Path getAppDir() { |
| Path path = new Path(new Path("target", "apps"), service.getName()); |
| LOG.info("Service path: {}", path); |
| return path; |
| } |
| |
| @Override |
| protected ServiceScheduler createServiceScheduler(ServiceContext context) |
| throws IOException, YarnException { |
| return new ServiceScheduler(context) { |
| |
| @SuppressWarnings("SuspiciousMethodCalls") |
| @Override |
| protected YarnRegistryViewForProviders createYarnRegistryOperations( |
| ServiceContext context, RegistryOperations registryClient) { |
| YarnRegistryViewForProviders yarnRegistryView = mock( |
| YarnRegistryViewForProviders.class); |
| if (!registryComponents.isEmpty()) { |
| try { |
| when(yarnRegistryView.listComponents()) |
| .thenReturn(new LinkedList<>(registryComponents.keySet())); |
| when(yarnRegistryView.getComponent(anyString())).thenAnswer( |
| invocation -> |
| registryComponents.get(invocation.getArguments()[0])); |
| } catch (IOException e) { |
| throw new RuntimeException(e); |
| } |
| } |
| return yarnRegistryView; |
| } |
| |
| @Override |
| protected AMRMClientAsync<AMRMClient.ContainerRequest> createAMRMClient() { |
| AMRMClientImpl client1 = new AMRMClientImpl() { |
| @Override public AllocateResponse allocate(float progressIndicator) |
| throws YarnException, IOException { |
| |
| AllocateResponse.AllocateResponseBuilder builder = |
| AllocateResponse.newBuilder(); |
| // add new containers if any |
| synchronized (feedContainers) { |
| if (feedContainers.isEmpty()) { |
| LOG.info("Allocating........ no containers"); |
| } else { |
| // The AMRMClient will return containers for compoenent that are |
| // at FLEXING state |
| List<Container> allocatedContainers = new LinkedList<>(); |
| Iterator<Container> itor = feedContainers.iterator(); |
| while (itor.hasNext()) { |
| Container c = itor.next(); |
| org.apache.hadoop.yarn.service.component.Component component = |
| componentsById.get(c.getAllocationRequestId()); |
| if (component.getState() == ComponentState.FLEXING) { |
| LOG.info("Allocated container {} ", c.getId()); |
| allocatedContainers.add(c); |
| itor.remove(); |
| } |
| } |
| builder.allocatedContainers(allocatedContainers); |
| } |
| } |
| |
| // add recovered containers if any |
| synchronized (recoveredContainers) { |
| if (!recoveredContainers.isEmpty()) { |
| List<Container> containersFromPrevAttempt = new LinkedList<>(); |
| containersFromPrevAttempt.addAll(recoveredContainers); |
| recoveredContainers.clear(); |
| builder.containersFromPreviousAttempt( |
| containersFromPrevAttempt); |
| } |
| } |
| |
| // add failed containers if any |
| synchronized (failedContainers) { |
| if (!failedContainers.isEmpty()) { |
| List<ContainerStatus> failed = |
| new LinkedList<>(failedContainers); |
| failedContainers.clear(); |
| builder.completedContainersStatuses(failed); |
| } |
| } |
| return builder.build(); |
| } |
| |
| @Override |
| public RegisterApplicationMasterResponse registerApplicationMaster( |
| String appHostName, int appHostPort, String appTrackingUrl, |
| Map placementConstraintsMap) throws YarnException, IOException { |
| return this.registerApplicationMaster(appHostName, appHostPort, |
| appTrackingUrl); |
| } |
| |
| @Override |
| public RegisterApplicationMasterResponse registerApplicationMaster( |
| String appHostName, int appHostPort, String appTrackingUrl) { |
| RegisterApplicationMasterResponse response = mock( |
| RegisterApplicationMasterResponse.class); |
| when(response.getResourceTypes()).thenReturn( |
| ResourceUtils.getResourcesTypeInfo()); |
| return response; |
| } |
| |
| @Override public void unregisterApplicationMaster( |
| FinalApplicationStatus appStatus, String appMessage, |
| String appTrackingUrl) { |
| // DO nothing |
| } |
| }; |
| |
| AMRMClientAsync<AMRMClient.ContainerRequest> amrmClientAsync = |
| AMRMClientAsync.createAMRMClientAsync(client1, 1000, |
| this.new AMRMClientCallback()); |
| |
| return amrmClientAsync; |
| } |
| |
| @SuppressWarnings("SuspiciousMethodCalls") |
| @Override |
| public NMClientAsync createNMClient() { |
| NMClientAsync nmClientAsync = super.createNMClient(); |
| NMClient nmClient = mock(NMClient.class); |
| try { |
| when(nmClient.getContainerStatus(anyObject(), anyObject())) |
| .thenAnswer(invocation -> |
| containerStatuses.get(invocation.getArguments()[0])); |
| } catch (YarnException | IOException e) { |
| throw new RuntimeException(e); |
| } |
| nmClientAsync.setClient(nmClient); |
| return nmClientAsync; |
| } |
| }; |
| } |
| |
| @Override protected void loadApplicationJson(ServiceContext context, |
| SliderFileSystem fs) throws IOException { |
| context.service = service; |
| } |
| |
| public void feedRegistryComponent(ContainerId containerId, String compName, |
| String compInstName) { |
| ServiceRecord record = new ServiceRecord(); |
| record.set(YarnRegistryAttributes.YARN_ID, containerId.toString()); |
| record.description = compInstName; |
| record.set(YarnRegistryAttributes.YARN_PERSISTENCE, |
| PersistencePolicies.CONTAINER); |
| record.set(YarnRegistryAttributes.YARN_IP, "localhost"); |
| record.set(YarnRegistryAttributes.YARN_HOSTNAME, "localhost"); |
| record.set(YarnRegistryAttributes.YARN_COMPONENT, compName); |
| registryComponents.put(RegistryPathUtils.encodeYarnID( |
| containerId.toString()), record); |
| } |
| |
| /** |
| * Simulates a recovered container that is sent to the AM in the heartbeat |
| * response. |
| * |
| * @param containerId The ID for the container |
| * @param compName The component to which the recovered container is fed. |
| */ |
| public void feedRecoveredContainer(ContainerId containerId, String compName) { |
| Container container = createContainer(containerId, compName); |
| recoveredContainers.add(container); |
| addContainerStatus(container, ContainerState.RUNNING); |
| } |
| |
| /** |
| * |
| * @param service The service for the component |
| * @param id The id for the container |
| * @param compName The component to which the container is fed |
| * @return |
| */ |
| public Container feedContainerToComp(Service service, int id, |
| String compName) { |
| ContainerId containerId = createContainerId(id); |
| Container container = createContainer(containerId, compName); |
| synchronized (feedContainers) { |
| feedContainers.add(container); |
| } |
| addContainerStatus(container, ContainerState.RUNNING); |
| return container; |
| } |
| |
| public void feedFailedContainerToComp(Service service, int id, String |
| compName) { |
| ApplicationId applicationId = ApplicationId.fromString(service.getId()); |
| ContainerId containerId = ContainerId |
| .newContainerId(ApplicationAttemptId.newInstance(applicationId, 1), id); |
| ContainerStatus status = Records.newRecord(ContainerStatus.class); |
| status.setContainerId(containerId); |
| synchronized (failedContainers) { |
| failedContainers.add(status); |
| } |
| } |
| |
| public ContainerId createContainerId(int id) { |
| ApplicationId applicationId = ApplicationId.fromString(service.getId()); |
| return ContainerId.newContainerId( |
| ApplicationAttemptId.newInstance(applicationId, 1), id); |
| } |
| |
| private Container createContainer(ContainerId containerId, String compName) { |
| NodeId nodeId = NodeId.newInstance("localhost", 1234); |
| Container container = Container.newInstance( |
| containerId, nodeId, "localhost", |
| Resource.newInstance(100, 1), |
| Priority.newInstance(0), null); |
| long allocateId = |
| context.scheduler.getAllComponents().get(compName).getAllocateId(); |
| container.setAllocationRequestId(allocateId); |
| return container; |
| } |
| |
| public void flexComponent(String compName, long numberOfContainers) |
| throws IOException { |
| ClientAMProtocol.ComponentCountProto componentCountProto = |
| ClientAMProtocol.ComponentCountProto.newBuilder().setName(compName) |
| .setNumberOfContainers(numberOfContainers).build(); |
| ClientAMProtocol.FlexComponentsRequestProto requestProto = |
| ClientAMProtocol.FlexComponentsRequestProto.newBuilder() |
| .addComponents(componentCountProto).build(); |
| context.clientAMService.flexComponents(requestProto); |
| } |
| |
| public Component getComponent(String compName) { |
| return context.scheduler.getAllComponents().get(compName); |
| } |
| |
| public void waitForDependenciesSatisfied(String compName) |
| throws TimeoutException, InterruptedException { |
| GenericTestUtils.waitFor(new Supplier<Boolean>() { |
| @Override public Boolean get() { |
| return context.scheduler.getAllComponents().get(compName) |
| .areDependenciesReady(); |
| } |
| }, 1000, 20000); |
| } |
| |
| public void waitForNumDesiredContainers(String compName, |
| int numDesiredContainers) throws TimeoutException, InterruptedException { |
| GenericTestUtils.waitFor(new Supplier<Boolean>() { |
| @Override public Boolean get() { |
| return context.scheduler.getAllComponents().get(compName) |
| .getNumDesiredInstances() == numDesiredContainers; |
| } |
| }, 1000, 20000); |
| } |
| |
| |
| public ComponentInstance getCompInstance(String compName, String |
| instanceName) { |
| return context.scheduler.getAllComponents().get(compName) |
| .getComponentInstance(instanceName); |
| } |
| |
| public void waitForCompInstanceState(ComponentInstance instance, |
| ComponentInstanceState state) |
| throws TimeoutException, InterruptedException { |
| GenericTestUtils.waitFor(new Supplier<Boolean>() { |
| @Override |
| public Boolean get() { |
| return instance.getState().equals(state); |
| } |
| }, 1000, 20000); |
| } |
| |
| private void addContainerStatus(Container container, ContainerState state) { |
| ContainerStatus status = ContainerStatus.newInstance(container.getId(), |
| state, "", 0); |
| status.setHost(container.getNodeId().getHost()); |
| status.setIPs(Lists.newArrayList(container.getNodeId().getHost())); |
| containerStatuses.put(container.getId(), status); |
| } |
| |
| @Override |
| protected ByteBuffer recordTokensForContainers() |
| throws IOException { |
| DataOutputBuffer dob = new DataOutputBuffer(); |
| if (amCreds == null) { |
| return ByteBuffer.wrap(dob.getData(), 0, dob.getLength()); |
| } |
| try { |
| amCreds.writeTokenStorageToStream(dob); |
| } finally { |
| dob.close(); |
| } |
| return ByteBuffer.wrap(dob.getData(), 0, dob.getLength()); |
| } |
| } |
