@Test(timeout = 100000)
@SuppressWarnings("unchecked")
public void testParallelDownloadAttemptsForPrivateResource() throws Exception {
DrainDispatcher dispatcher1 = null;
try {
dispatcher1 = new DrainDispatcher();
String user = "testuser";
ApplicationId appId = BuilderUtils.newApplicationId(1, 1);
// creating one local directory
List<Path> localDirs = new ArrayList<Path>();
String[] sDirs = new String[1];
for (int i = 0; i < 1; ++i) {
localDirs.add(lfs.makeQualified(new Path(basedir, i + "")));
sDirs[i] = localDirs.get(i).toString();
}
conf.setStrings(YarnConfiguration.NM_LOCAL_DIRS, sDirs);
LocalDirsHandlerService localDirHandler = new LocalDirsHandlerService();
localDirHandler.init(conf);
// Registering event handlers
EventHandler<ApplicationEvent> applicationBus = mock(EventHandler.class);
dispatcher1.register(ApplicationEventType.class, applicationBus);
EventHandler<ContainerEvent> containerBus = mock(EventHandler.class);
dispatcher1.register(ContainerEventType.class, containerBus);
ContainerExecutor exec = mock(ContainerExecutor.class);
DeletionService delService = mock(DeletionService.class);
LocalDirsHandlerService dirsHandler = new LocalDirsHandlerService();
// initializing directory handler.
dirsHandler.init(conf);
dispatcher1.init(conf);
dispatcher1.start();
ResourceLocalizationService rls =
new ResourceLocalizationService(dispatcher1, exec, delService,
localDirHandler, new NMNullStateStoreService());
dispatcher1.register(LocalizationEventType.class, rls);
rls.init(conf);
rls.handle(createApplicationLocalizationEvent(user, appId));
LocalResourceRequest req =
new LocalResourceRequest(new Path("file:///tmp"), 123L,
LocalResourceType.FILE, LocalResourceVisibility.PRIVATE, "");
// We need to pre-populate the LocalizerRunner as the
// Resource Localization Service code internally starts them which
// definitely we don't want.
// creating new containers and populating corresponding localizer runners
// Container - 1
ContainerImpl container1 = createMockContainer(user, 1);
String localizerId1 = container1.getContainerId().toString();
rls.getPrivateLocalizers().put(
localizerId1,
rls.new LocalizerRunner(new LocalizerContext(user, container1
.getContainerId(), null), localizerId1));
LocalizerRunner localizerRunner1 = rls.getLocalizerRunner(localizerId1);
dispatcher1.getEventHandler().handle(
createContainerLocalizationEvent(container1,
LocalResourceVisibility.PRIVATE, req));
Assert
.assertTrue(waitForPrivateDownloadToStart(rls, localizerId1, 1, 200));
// Container - 2 now makes the request.
ContainerImpl container2 = createMockContainer(user, 2);
String localizerId2 = container2.getContainerId().toString();
rls.getPrivateLocalizers().put(
localizerId2,
rls.new LocalizerRunner(new LocalizerContext(user, container2
.getContainerId(), null), localizerId2));
LocalizerRunner localizerRunner2 = rls.getLocalizerRunner(localizerId2);
dispatcher1.getEventHandler().handle(
createContainerLocalizationEvent(container2,
LocalResourceVisibility.PRIVATE, req));
Assert
.assertTrue(waitForPrivateDownloadToStart(rls, localizerId2, 1, 200));
// Retrieving localized resource.
LocalResourcesTracker tracker =
rls.getLocalResourcesTracker(LocalResourceVisibility.PRIVATE, user,
appId);
LocalizedResource lr = tracker.getLocalizedResource(req);
// Resource would now have moved into DOWNLOADING state
Assert.assertEquals(ResourceState.DOWNLOADING, lr.getState());
// Resource should have one permit
Assert.assertEquals(1, lr.sem.availablePermits());
// Resource Localization Service receives first heart beat from
// ContainerLocalizer for container1
LocalizerHeartbeatResponse response1 =
rls.heartbeat(createLocalizerStatus(localizerId1));
// Resource must have been added to scheduled map
Assert.assertEquals(1, localizerRunner1.scheduled.size());
// Checking resource in the response and also available permits for it.
Assert.assertEquals(req.getResource(), response1.getResourceSpecs()
.get(0).getResource().getResource());
Assert.assertEquals(0, lr.sem.availablePermits());
// Resource Localization Service now receives first heart beat from
// ContainerLocalizer for container2
LocalizerHeartbeatResponse response2 =
rls.heartbeat(createLocalizerStatus(localizerId2));
// Resource must not have been added to scheduled map
Assert.assertEquals(0, localizerRunner2.scheduled.size());
// No resource is returned in response
Assert.assertEquals(0, response2.getResourceSpecs().size());
// ContainerLocalizer - 1 now sends failed resource heartbeat.
rls.heartbeat(createLocalizerStatusForFailedResource(localizerId1, req));
// Resource Localization should fail and state is modified accordingly.
// Also Local should be release on the LocalizedResource.
Assert
.assertTrue(waitForResourceState(lr, rls, req,
LocalResourceVisibility.PRIVATE, user, appId, ResourceState.FAILED,
200));
Assert.assertTrue(lr.getState().equals(ResourceState.FAILED));
Assert.assertEquals(0, localizerRunner1.scheduled.size());
// Now Container-2 once again sends heart beat to resource localization
// service
// Now container-2 again try to download the resource it should still
// not get the resource as the resource is now not in DOWNLOADING state.
response2 = rls.heartbeat(createLocalizerStatus(localizerId2));
// Resource must not have been added to scheduled map.
// Also as the resource has failed download it will be removed from
// pending list.
Assert.assertEquals(0, localizerRunner2.scheduled.size());
Assert.assertEquals(0, localizerRunner2.pending.size());
Assert.assertEquals(0, response2.getResourceSpecs().size());
} finally {
if (dispatcher1 != null) {
dispatcher1.stop();
}
}
}