org.apache.hadoop.yarn.server.nodemanager.containermanager.localizer.TestContainerLocalizer.java Source code

Java tutorial

Introduction

Here is the source code for org.apache.hadoop.yarn.server.nodemanager.containermanager.localizer.TestContainerLocalizer.java

Source

/**
* 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.nodemanager.containermanager.localizer;

import static org.junit.Assert.fail;
import static org.mockito.Matchers.any;
import static org.mockito.Matchers.anyBoolean;
import static org.mockito.Matchers.anyInt;
import static org.mockito.Matchers.argThat;
import static org.mockito.Matchers.eq;
import static org.mockito.Matchers.isA;
import static org.mockito.Matchers.same;
import static org.mockito.Mockito.doNothing;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

import java.io.File;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Random;
import java.util.concurrent.Callable;
import java.util.concurrent.CompletionService;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Future;

import org.apache.commons.io.FileUtils;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.AbstractFileSystem;
import org.apache.hadoop.fs.CommonConfigurationKeys;
import org.apache.hadoop.fs.FSDataInputStream;
import org.apache.hadoop.fs.FileContext;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.fs.permission.FsPermission;
import org.apache.hadoop.io.DataInputBuffer;
import org.apache.hadoop.io.DataOutputBuffer;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.security.Credentials;
import org.apache.hadoop.security.UserGroupInformation;
import org.apache.hadoop.security.token.Token;
import org.apache.hadoop.yarn.api.records.LocalResource;
import org.apache.hadoop.yarn.api.records.LocalResourceType;
import org.apache.hadoop.yarn.api.records.LocalResourceVisibility;
import org.apache.hadoop.yarn.api.records.URL;
import org.apache.hadoop.yarn.exceptions.YarnException;
import org.apache.hadoop.yarn.exceptions.YarnRuntimeException;
import org.apache.hadoop.yarn.factories.RecordFactory;
import org.apache.hadoop.yarn.server.nodemanager.api.LocalizationProtocol;
import org.apache.hadoop.yarn.server.nodemanager.api.ResourceLocalizationSpec;
import org.apache.hadoop.yarn.server.nodemanager.api.protocolrecords.LocalResourceStatus;
import org.apache.hadoop.yarn.server.nodemanager.api.protocolrecords.LocalizerAction;
import org.apache.hadoop.yarn.server.nodemanager.api.protocolrecords.LocalizerStatus;
import org.junit.After;
import org.junit.Assert;
import org.junit.Test;
import org.mockito.ArgumentMatcher;
import org.mockito.invocation.InvocationOnMock;
import org.mockito.stubbing.Answer;

public class TestContainerLocalizer {

    static final Log LOG = LogFactory.getLog(TestContainerLocalizer.class);
    static final Path basedir = new Path("target", TestContainerLocalizer.class.getName());
    static final FsPermission CACHE_DIR_PERM = new FsPermission((short) 0710);
    static final FsPermission USERCACHE_DIR_PERM = new FsPermission((short) 0755);

    static final String appUser = "yak";
    static final String appUserFolder = "yakFolder";
    static final String appId = "app_RM_0";
    static final String containerId = "container_0";
    static final InetSocketAddress nmAddr = new InetSocketAddress("foobar", 8040);

    private AbstractFileSystem spylfs;
    private Random random;
    private List<Path> localDirs;
    private Path tokenPath;
    private LocalizationProtocol nmProxy;

    @After
    public void cleanUp() throws IOException {
        FileUtils.deleteDirectory(new File(basedir.toUri().getRawPath()));
    }

    @Test
    public void testMain() throws Exception {
        FileContext fs = FileContext.getLocalFSFileContext();
        spylfs = spy(fs.getDefaultFileSystem());
        ContainerLocalizer localizer = setupContainerLocalizerForTest();

        // verify created cache
        List<Path> privCacheList = new ArrayList<Path>();
        List<Path> appCacheList = new ArrayList<Path>();
        for (Path p : localDirs) {
            Path base = new Path(new Path(p, ContainerLocalizer.USERCACHE), appUser);
            Path privcache = new Path(base, ContainerLocalizer.FILECACHE);
            privCacheList.add(privcache);
            Path appDir = new Path(base, new Path(ContainerLocalizer.APPCACHE, appId));
            Path appcache = new Path(appDir, ContainerLocalizer.FILECACHE);
            appCacheList.add(appcache);
        }

        // mock heartbeat responses from NM
        ResourceLocalizationSpec rsrcA = getMockRsrc(random, LocalResourceVisibility.PRIVATE, privCacheList.get(0));
        ResourceLocalizationSpec rsrcB = getMockRsrc(random, LocalResourceVisibility.PRIVATE, privCacheList.get(0));
        ResourceLocalizationSpec rsrcC = getMockRsrc(random, LocalResourceVisibility.APPLICATION,
                appCacheList.get(0));
        ResourceLocalizationSpec rsrcD = getMockRsrc(random, LocalResourceVisibility.PRIVATE, privCacheList.get(0));

        when(nmProxy.heartbeat(isA(LocalizerStatus.class)))
                .thenReturn(
                        new MockLocalizerHeartbeatResponse(LocalizerAction.LIVE, Collections.singletonList(rsrcA)))
                .thenReturn(
                        new MockLocalizerHeartbeatResponse(LocalizerAction.LIVE, Collections.singletonList(rsrcB)))
                .thenReturn(
                        new MockLocalizerHeartbeatResponse(LocalizerAction.LIVE, Collections.singletonList(rsrcC)))
                .thenReturn(
                        new MockLocalizerHeartbeatResponse(LocalizerAction.LIVE, Collections.singletonList(rsrcD)))
                .thenReturn(new MockLocalizerHeartbeatResponse(LocalizerAction.LIVE,
                        Collections.<ResourceLocalizationSpec>emptyList()))
                .thenReturn(new MockLocalizerHeartbeatResponse(LocalizerAction.DIE, null));

        LocalResource tRsrcA = rsrcA.getResource();
        LocalResource tRsrcB = rsrcB.getResource();
        LocalResource tRsrcC = rsrcC.getResource();
        LocalResource tRsrcD = rsrcD.getResource();
        doReturn(new FakeDownload(rsrcA.getResource().getResource().getFile(), true)).when(localizer)
                .download(isA(Path.class), eq(tRsrcA), isA(UserGroupInformation.class));
        doReturn(new FakeDownload(rsrcB.getResource().getResource().getFile(), true)).when(localizer)
                .download(isA(Path.class), eq(tRsrcB), isA(UserGroupInformation.class));
        doReturn(new FakeDownload(rsrcC.getResource().getResource().getFile(), true)).when(localizer)
                .download(isA(Path.class), eq(tRsrcC), isA(UserGroupInformation.class));
        doReturn(new FakeDownload(rsrcD.getResource().getResource().getFile(), true)).when(localizer)
                .download(isA(Path.class), eq(tRsrcD), isA(UserGroupInformation.class));

        // run localization
        localizer.runLocalization(nmAddr);
        for (Path p : localDirs) {
            Path base = new Path(new Path(p, ContainerLocalizer.USERCACHE), appUserFolder);
            Path privcache = new Path(base, ContainerLocalizer.FILECACHE);
            // $x/usercache/$user/filecache
            verify(spylfs).mkdir(eq(privcache), eq(CACHE_DIR_PERM), eq(false));
            Path appDir = new Path(base, new Path(ContainerLocalizer.APPCACHE, appId));
            // $x/usercache/$user/appcache/$appId/filecache
            Path appcache = new Path(appDir, ContainerLocalizer.FILECACHE);
            verify(spylfs).mkdir(eq(appcache), eq(CACHE_DIR_PERM), eq(false));
        }
        // verify tokens read at expected location
        verify(spylfs).open(tokenPath);

        // verify downloaded resources reported to NM
        verify(nmProxy).heartbeat(argThat(new HBMatches(rsrcA.getResource())));
        verify(nmProxy).heartbeat(argThat(new HBMatches(rsrcB.getResource())));
        verify(nmProxy).heartbeat(argThat(new HBMatches(rsrcC.getResource())));
        verify(nmProxy).heartbeat(argThat(new HBMatches(rsrcD.getResource())));

        // verify all HB use localizerID provided
        verify(nmProxy, never()).heartbeat(argThat(new ArgumentMatcher<LocalizerStatus>() {
            @Override
            public boolean matches(Object o) {
                LocalizerStatus status = (LocalizerStatus) o;
                return !containerId.equals(status.getLocalizerId());
            }
        }));
    }

    @Test(timeout = 15000)
    public void testMainFailure() throws Exception {

        FileContext fs = FileContext.getLocalFSFileContext();
        spylfs = spy(fs.getDefaultFileSystem());
        ContainerLocalizer localizer = setupContainerLocalizerForTest();

        // Assume the NM heartbeat fails say because of absent tokens.
        when(nmProxy.heartbeat(isA(LocalizerStatus.class))).thenThrow(new YarnException("Sigh, no token!"));

        // run localization, it should fail
        try {
            localizer.runLocalization(nmAddr);
            Assert.fail("Localization succeeded unexpectedly!");
        } catch (IOException e) {
            Assert.assertTrue(e.getMessage().contains("Sigh, no token!"));
        }
    }

    @Test
    @SuppressWarnings("unchecked")
    public void testLocalizerTokenIsGettingRemoved() throws Exception {
        FileContext fs = FileContext.getLocalFSFileContext();
        spylfs = spy(fs.getDefaultFileSystem());
        ContainerLocalizer localizer = setupContainerLocalizerForTest();
        doNothing().when(localizer).localizeFiles(any(LocalizationProtocol.class), any(CompletionService.class),
                any(UserGroupInformation.class));
        localizer.runLocalization(nmAddr);
        verify(spylfs, times(1)).delete(tokenPath, false);
    }

    @Test
    @SuppressWarnings("unchecked") // mocked generics
    public void testContainerLocalizerClosesFilesystems() throws Exception {

        // verify filesystems are closed when localizer doesn't fail
        FileContext fs = FileContext.getLocalFSFileContext();
        spylfs = spy(fs.getDefaultFileSystem());

        ContainerLocalizer localizer = setupContainerLocalizerForTest();
        doNothing().when(localizer).localizeFiles(any(LocalizationProtocol.class), any(CompletionService.class),
                any(UserGroupInformation.class));
        verify(localizer, never()).closeFileSystems(any(UserGroupInformation.class));

        localizer.runLocalization(nmAddr);
        verify(localizer).closeFileSystems(any(UserGroupInformation.class));

        spylfs = spy(fs.getDefaultFileSystem());

        // verify filesystems are closed when localizer fails
        localizer = setupContainerLocalizerForTest();
        doThrow(new YarnRuntimeException("Forced Failure")).when(localizer).localizeFiles(
                any(LocalizationProtocol.class), any(CompletionService.class), any(UserGroupInformation.class));
        verify(localizer, never()).closeFileSystems(any(UserGroupInformation.class));
        try {
            localizer.runLocalization(nmAddr);
            Assert.fail("Localization succeeded unexpectedly!");
        } catch (IOException e) {
            verify(localizer).closeFileSystems(any(UserGroupInformation.class));
        }
    }

    @SuppressWarnings("unchecked") // mocked generics
    private ContainerLocalizer setupContainerLocalizerForTest() throws Exception {
        // don't actually create dirs
        doNothing().when(spylfs).mkdir(isA(Path.class), isA(FsPermission.class), anyBoolean());

        Configuration conf = new Configuration();
        FileContext lfs = FileContext.getFileContext(spylfs, conf);
        localDirs = new ArrayList<Path>();
        for (int i = 0; i < 4; ++i) {
            localDirs.add(lfs.makeQualified(new Path(basedir, i + "")));
        }
        RecordFactory mockRF = getMockLocalizerRecordFactory();
        ContainerLocalizer concreteLoc = new ContainerLocalizer(lfs, appUser, appId, containerId, localDirs, mockRF,
                appUserFolder);
        ContainerLocalizer localizer = spy(concreteLoc);

        // return credential stream instead of opening local file
        random = new Random();
        long seed = random.nextLong();
        System.out.println("SEED: " + seed);
        random.setSeed(seed);
        DataInputBuffer appTokens = createFakeCredentials(random, 10);
        tokenPath = lfs.makeQualified(new Path(String.format(ContainerLocalizer.TOKEN_FILE_NAME_FMT, containerId)));
        doReturn(new FSDataInputStream(new FakeFSDataInputStream(appTokens))).when(spylfs).open(tokenPath);
        nmProxy = mock(LocalizationProtocol.class);
        doReturn(nmProxy).when(localizer).getProxy(nmAddr);
        doNothing().when(localizer).sleep(anyInt());

        // return result instantly for deterministic test
        ExecutorService syncExec = mock(ExecutorService.class);
        CompletionService<Path> cs = mock(CompletionService.class);
        when(cs.submit(isA(Callable.class))).thenAnswer(new Answer<Future<Path>>() {
            @Override
            public Future<Path> answer(InvocationOnMock invoc) throws Throwable {
                Future<Path> done = mock(Future.class);
                when(done.isDone()).thenReturn(true);
                FakeDownload d = (FakeDownload) invoc.getArguments()[0];
                when(done.get()).thenReturn(d.call());
                return done;
            }
        });
        doReturn(syncExec).when(localizer).createDownloadThreadPool();
        doReturn(cs).when(localizer).createCompletionService(syncExec);

        return localizer;
    }

    static class HBMatches extends ArgumentMatcher<LocalizerStatus> {
        final LocalResource rsrc;

        HBMatches(LocalResource rsrc) {
            this.rsrc = rsrc;
        }

        @Override
        public boolean matches(Object o) {
            LocalizerStatus status = (LocalizerStatus) o;
            for (LocalResourceStatus localized : status.getResources()) {
                switch (localized.getStatus()) {
                case FETCH_SUCCESS:
                    if (localized.getLocalPath().getFile().contains(rsrc.getResource().getFile())) {
                        return true;
                    }
                    break;
                default:
                    fail("Unexpected: " + localized.getStatus());
                    break;
                }
            }
            return false;
        }
    }

    static class FakeDownload implements Callable<Path> {
        private final Path localPath;
        private final boolean succeed;

        FakeDownload(String absPath, boolean succeed) {
            this.localPath = new Path("file:///localcache" + absPath);
            this.succeed = succeed;
        }

        @Override
        public Path call() throws IOException {
            if (!succeed) {
                throw new IOException("FAIL " + localPath);
            }
            return localPath;
        }
    }

    static RecordFactory getMockLocalizerRecordFactory() {
        RecordFactory mockRF = mock(RecordFactory.class);
        when(mockRF.newRecordInstance(same(LocalResourceStatus.class)))
                .thenAnswer(new Answer<LocalResourceStatus>() {
                    @Override
                    public LocalResourceStatus answer(InvocationOnMock invoc) throws Throwable {
                        return new MockLocalResourceStatus();
                    }
                });
        when(mockRF.newRecordInstance(same(LocalizerStatus.class))).thenAnswer(new Answer<LocalizerStatus>() {
            @Override
            public LocalizerStatus answer(InvocationOnMock invoc) throws Throwable {
                return new MockLocalizerStatus();
            }
        });
        return mockRF;
    }

    static ResourceLocalizationSpec getMockRsrc(Random r, LocalResourceVisibility vis, Path p) {
        ResourceLocalizationSpec resourceLocalizationSpec = mock(ResourceLocalizationSpec.class);

        LocalResource rsrc = mock(LocalResource.class);
        String name = Long.toHexString(r.nextLong());
        URL uri = mock(org.apache.hadoop.yarn.api.records.URL.class);
        when(uri.getScheme()).thenReturn("file");
        when(uri.getHost()).thenReturn(null);
        when(uri.getFile()).thenReturn("/local/" + vis + "/" + name);

        when(rsrc.getResource()).thenReturn(uri);
        when(rsrc.getSize()).thenReturn(r.nextInt(1024) + 1024L);
        when(rsrc.getTimestamp()).thenReturn(r.nextInt(1024) + 2048L);
        when(rsrc.getType()).thenReturn(LocalResourceType.FILE);
        when(rsrc.getVisibility()).thenReturn(vis);

        when(resourceLocalizationSpec.getResource()).thenReturn(rsrc);
        when(resourceLocalizationSpec.getDestinationDirectory()).thenReturn(URL.fromPath(p));
        return resourceLocalizationSpec;
    }

    @SuppressWarnings({ "rawtypes", "unchecked" })
    static DataInputBuffer createFakeCredentials(Random r, int nTok) throws IOException {
        Credentials creds = new Credentials();
        byte[] password = new byte[20];
        Text kind = new Text();
        Text service = new Text();
        Text alias = new Text();
        for (int i = 0; i < nTok; ++i) {
            byte[] identifier = ("idef" + i).getBytes();
            r.nextBytes(password);
            kind.set("kind" + i);
            service.set("service" + i);
            alias.set("token" + i);
            Token token = new Token(identifier, password, kind, service);
            creds.addToken(alias, token);
        }
        DataOutputBuffer buf = new DataOutputBuffer();
        creds.writeTokenStorageToStream(buf);
        DataInputBuffer ret = new DataInputBuffer();
        ret.reset(buf.getData(), 0, buf.getLength());
        return ret;
    }

    @Test(timeout = 10000)
    public void testUserCacheDirPermission() throws Exception {
        Configuration conf = new Configuration();
        conf.set(CommonConfigurationKeys.FS_PERMISSIONS_UMASK_KEY, "077");
        FileContext lfs = FileContext.getLocalFSFileContext(conf);
        Path fileCacheDir = lfs.makeQualified(new Path(basedir, "filecache"));
        lfs.mkdir(fileCacheDir, FsPermission.getDefault(), true);
        RecordFactory recordFactory = mock(RecordFactory.class);
        ContainerLocalizer localizer = new ContainerLocalizer(lfs,
                UserGroupInformation.getCurrentUser().getUserName(), "application_01", "container_01",
                new ArrayList<Path>(), recordFactory, UserGroupInformation.getCurrentUser().getUserName());
        LocalResource rsrc = mock(LocalResource.class);
        when(rsrc.getVisibility()).thenReturn(LocalResourceVisibility.PRIVATE);
        Path destDirPath = new Path(fileCacheDir, "0/0/85");
        //create one of the parent directories with the wrong permissions first
        FsPermission wrongPerm = new FsPermission((short) 0700);
        lfs.mkdir(destDirPath.getParent().getParent(), wrongPerm, false);
        lfs.mkdir(destDirPath.getParent(), wrongPerm, false);
        //Localize and check the directory permission are correct.
        localizer.download(destDirPath, rsrc, UserGroupInformation.getCurrentUser());
        Assert.assertEquals("Cache directory permissions filecache/0/0 is incorrect", USERCACHE_DIR_PERM,
                lfs.getFileStatus(destDirPath.getParent()).getPermission());
        Assert.assertEquals("Cache directory permissions filecache/0 is incorrect", USERCACHE_DIR_PERM,
                lfs.getFileStatus(destDirPath.getParent().getParent()).getPermission());
    }

}