ContainerLocalizer { @SuppressWarnings("deprecation") public void runLocalization(final InetSocketAddress nmAddr) throws IOException, InterruptedException { initDirs(conf, user, appId, lfs, localDirs); final Credentials creds = new Credentials(); DataInputStream credFile = null; try { Path tokenPath = new Path(String.format(TOKEN_FILE_NAME_FMT, localizerId)); credFile = lfs.open(tokenPath); creds.readTokenStorageStream(credFile); lfs.delete(tokenPath, false); } finally { if (credFile != null) { credFile.close(); } } UserGroupInformation remoteUser = UserGroupInformation.createRemoteUser(user); remoteUser.addToken(creds.getToken(LocalizerTokenIdentifier.KIND)); final LocalizationProtocol nodeManager = remoteUser.doAs(new PrivilegedAction<LocalizationProtocol>() { @Override public LocalizationProtocol run() { return getProxy(nmAddr); } }); UserGroupInformation ugi = UserGroupInformation.createRemoteUser(user); for (Token<? extends TokenIdentifier> token : creds.getAllTokens()) { ugi.addToken(token); } ExecutorService exec = null; try { exec = createDownloadThreadPool(); CompletionService<Path> ecs = createCompletionService(exec); localizeFiles(nodeManager, ecs, ugi); } catch (Throwable e) { throw new IOException(e); } finally { try { if (exec != null) { exec.shutdown(); destroyShellProcesses(getAllShells()); exec.awaitTermination(10, TimeUnit.SECONDS); } LocalDirAllocator.removeContext(appCacheDirContextName); } finally { closeFileSystems(ugi); } } } ContainerLocalizer(FileContext lfs, String user, String appId,
      String localizerId, List<Path> localDirs,
      RecordFactory recordFactory); @Private @VisibleForTesting LocalizationProtocol getProxy(final InetSocketAddress nmAddr); @SuppressWarnings("deprecation") void runLocalization(final InetSocketAddress nmAddr); static List<String> getJavaOpts(Configuration conf); static void buildMainArgs(List<String> command,
      String user, String appId, String locId,
      InetSocketAddress nmAddr, List<String> localDirs, Configuration conf); static void main(String[] argv); static final String FILECACHE; static final String APPCACHE; static final String USERCACHE; static final String TOKEN_FILE_NAME_FMT; }
@Test(timeout = 15000) public void testMainFailure() throws Exception { ContainerLocalizerWrapper wrapper = new ContainerLocalizerWrapper(); ContainerLocalizer localizer = wrapper.setupContainerLocalizerForTest(); LocalizationProtocol nmProxy = wrapper.nmProxy; mockOutDownloads(localizer); when(nmProxy.heartbeat(isA(LocalizerStatus.class))).thenThrow( new YarnException("Sigh, no token!")); try { localizer.runLocalization(nmAddr); Assert.fail("Localization succeeded unexpectedly!"); } catch (IOException e) { Assert.assertTrue(e.getMessage().contains("Sigh, no token!")); } }
@Test public void testMultipleLocalizers() throws Exception { FakeContainerLocalizerWrapper testA = new FakeContainerLocalizerWrapper(); FakeContainerLocalizerWrapper testB = new FakeContainerLocalizerWrapper(); FakeContainerLocalizer localizerA = testA.init(); FakeContainerLocalizer localizerB = testB.init(); Thread threadA = new Thread() { @Override public void run() { try { localizerA.runLocalization(nmAddr); } catch (Exception e) { LOG.warn(e.toString()); } } }; Thread threadB = new Thread() { @Override public void run() { try { localizerB.runLocalization(nmAddr); } catch (Exception e) { LOG.warn(e.toString()); } } }; ShellCommandExecutor shexcA = null; ShellCommandExecutor shexcB = null; try { threadA.start(); threadB.start(); GenericTestUtils.waitFor(new Supplier<Boolean>() { @Override public Boolean get() { FakeContainerLocalizer.FakeLongDownload downloader = localizerA.getDownloader(); return downloader != null && downloader.getShexc() != null && downloader.getShexc().getProcess() != null; } }, 10, 30000); GenericTestUtils.waitFor(new Supplier<Boolean>() { @Override public Boolean get() { FakeContainerLocalizer.FakeLongDownload downloader = localizerB.getDownloader(); return downloader != null && downloader.getShexc() != null && downloader.getShexc().getProcess() != null; } }, 10, 30000); shexcA = localizerA.getDownloader().getShexc(); shexcB = localizerB.getDownloader().getShexc(); assertTrue("Localizer A process not running, but should be", shexcA.getProcess().isAlive()); assertTrue("Localizer B process not running, but should be", shexcB.getProcess().isAlive()); testA.heartbeatResponse++; testB.heartbeatResponse++; testA.heartbeatResponse++; threadA.join(); shexcA.getProcess().waitFor(10000, TimeUnit.MILLISECONDS); assertFalse("Localizer A process is still running, but shouldn't be", shexcA.getProcess().isAlive()); assertTrue("Localizer B process not running, but should be", shexcB.getProcess().isAlive()); } finally { shexcA.getProcess().destroy(); shexcB.getProcess().destroy(); shexcA.getProcess().waitFor(10000, TimeUnit.MILLISECONDS); shexcB.getProcess().waitFor(10000, TimeUnit.MILLISECONDS); threadA.join(); testB.heartbeatResponse++; threadB.join(); } }