TerrapinControllerServiceImpl implements TerrapinController.ServiceIface { @Override public Future<Void> loadFileSet(final TerrapinLoadRequest request) { return futurePool.apply(new ExceptionalFunction0<Void>() { @Override public Void applyE() throws Throwable { try { String hdfsDir = request.getHdfsDirectory(); LOG.info("Loading resource " + hdfsDir); long startTimeMillis = System.currentTimeMillis(); String resourceName = TerrapinUtil.hdfsDirToHelixResource(hdfsDir); FileSetInfo oldFileSetInfo = zkManager.getFileSetInfo(request.getFileSet()); boolean resourceExists = helixAdmin.getResourcesInCluster(clusterName).contains( resourceName); if (oldFileSetInfo != null && oldFileSetInfo.servingInfo.hdfsPath.equals(hdfsDir) && resourceExists) { LOG.info("Resource and fileset already exist. Not doing anything."); return null; } IdealState idealState = ControllerUtil.buildIdealStateForHdfsDir( hdfsClient, hdfsDir, resourceName, request.getOptions().getPartitioner(), configuration.getInt(Constants.NUM_SERVING_REPLICAS, 3), configuration.getBoolean(Constants.ENABLE_ZK_COMPRESSION, Constants.ENABLE_ZK_COMPRESSION_DEFAULT)); if (idealState.getNumPartitions() != request.getExpectedNumPartitions()) { throw new ControllerException( "Partition count mismatch for hdfs dir " + hdfsDir + " FOUND : " + idealState.getNumPartitions() + " EXPECTED : " + request.getExpectedNumPartitions(), ControllerErrorCode.INVALID_DATA_SET); } try { if (idealState.getBucketSize() > 1) { if (!resourceExists) { helixAdmin.addResource(clusterName, resourceName, idealState.getNumPartitions(), "OnlineOffline", "CUSTOMIZED", idealState.getBucketSize()); } helixAdmin.setResourceIdealState(clusterName, resourceName, idealState); } else { if (!resourceExists) { helixAdmin.addResource(clusterName, resourceName, idealState); } else { helixAdmin.setResourceIdealState(clusterName, resourceName, idealState); } } } catch (Exception e) { LOG.warn("Resource creation failed for " + resourceName + ", rolling back.", e); throw new ControllerException("Resource creation failure for " + resourceName, ControllerErrorCode.HELIX_ERROR); } LOG.info("Time taken for building resource for " + hdfsDir + ": " + (System.currentTimeMillis() - startTimeMillis) + "ms."); LOG.info("New resource " + resourceName); long startLoadTimeMillis = System.currentTimeMillis(); boolean isLoaded = false; while (!isLoaded) { long timeElapsedMillis = System.currentTimeMillis() - startLoadTimeMillis; if (timeElapsedMillis > Constants.LOAD_TIMEOUT_SECONDS * 1000) { throw new ControllerException("Timeout while trying to online resource " + resourceName, ControllerErrorCode.SHARD_LOAD_TIMEOUT); } Thread.sleep(5000); timeElapsedMillis = System.currentTimeMillis() - startLoadTimeMillis; ViewInfo viewInfo = zkManager.getViewInfo(resourceName); if (viewInfo == null) { continue; } int numLoaded = viewInfo.getNumOnlinePartitions(); if (numLoaded == idealState.getNumPartitions()) { LOG.info("Onlined " + resourceName + " in " + timeElapsedMillis + "ms."); isLoaded = true; } else { LOG.info("Onlined " + numLoaded + "/" + idealState.getNumPartitions() + " partitions for resource " + resourceName + " in " + timeElapsedMillis + "ms."); } } Thread.sleep(2000); List<FileSetInfo.ServingInfo> oldServingInfoList = null; if (oldFileSetInfo == null) { oldServingInfoList = Lists.newArrayList(); } else { oldServingInfoList = oldFileSetInfo.oldServingInfoList; oldServingInfoList.add(0, oldFileSetInfo.servingInfo); while (oldServingInfoList.size() >= request.getOptions().getNumVersionsToKeep() && !oldServingInfoList.isEmpty()) { oldServingInfoList.remove(oldServingInfoList.size() - 1); } } FileSetInfo fileSetInfo = new FileSetInfo(request.getFileSet(), hdfsDir, idealState.getNumPartitions(), oldServingInfoList, request.getOptions()); zkManager.setFileSetInfo(request.getFileSet(), fileSetInfo); } catch (Exception e) { LOG.warn("Exception while loading.", e); throw getException(e); } return null; } }); } TerrapinControllerServiceImpl(PropertiesConfiguration configuration,
                                       ZooKeeperManager zkManager,
                                       DFSClient hdfsClient,
                                       HelixAdmin helixAdmin,
                                       String clusterName); @Override Future<Void> loadFileSet(final TerrapinLoadRequest request);  }
@Test @PrepareForTest({ControllerUtil.class}) public void testLoadFileSet() throws Exception { PowerMockito.mockStatic(ControllerUtil.class); TerrapinLoadRequest request = new TerrapinLoadRequest("fileset", "/terrapin/data/fileset/12345564534", 200); Options requestOptions = new Options(); requestOptions.setNumVersionsToKeep(2); request.setOptions(requestOptions); FileSetInfo fileSetInfo = new FileSetInfo( request.getFileSet(), "/terrapin/data/fileset/12345563434", request.getExpectedNumPartitions(), Lists.newArrayList(mock(FileSetInfo.ServingInfo.class)), new Options() ); String resourceName = TerrapinUtil.hdfsDirToHelixResource(request.getHdfsDirectory()); CustomModeISBuilder idealStateBuilder = new CustomModeISBuilder(resourceName); idealStateBuilder.assignInstanceAndState(resourceName + "$0", "host0", "ONLINE"); idealStateBuilder.assignInstanceAndState(resourceName + "$1", "host1", "ONLINE"); idealStateBuilder.setStateModel("OnlineOffline"); idealStateBuilder.setNumReplica(2); idealStateBuilder.setNumPartitions(request.getExpectedNumPartitions()); IdealState is = idealStateBuilder.build(); is.setBucketSize(2); is.setRebalanceMode(IdealState.RebalanceMode.CUSTOMIZED); ViewInfo viewInfo1 = mock(ViewInfo.class); ViewInfo viewInfo2 = mock(ViewInfo.class); when(configuration.getInt(eq(Constants.NUM_SERVING_REPLICAS), eq(3))).thenReturn(3); when(zkManager.getFileSetInfo(eq(request.getFileSet()))).thenReturn(fileSetInfo); when(zkManager.getViewInfo(eq(resourceName))).thenReturn(viewInfo1).thenReturn(viewInfo2); doNothing().when(zkManager).setFileSetInfo(eq(request.getFileSet()), any(FileSetInfo.class)); when(viewInfo1.getNumOnlinePartitions()).thenReturn(request.getExpectedNumPartitions() / 2); when(viewInfo2.getNumOnlinePartitions()).thenReturn(request.getExpectedNumPartitions()); when(helixAdmin.getResourcesInCluster(CLUSTER)).thenReturn(new ArrayList<String>()); when(ControllerUtil.buildIdealStateForHdfsDir(any(DFSClient.class), anyString(), anyString(), any(PartitionerType.class), anyInt(), anyBoolean())). thenReturn(is); doNothing().when(helixAdmin).addResource(eq(CLUSTER), eq(resourceName), eq(is)); doNothing().when(helixAdmin).addResource(eq(CLUSTER), eq(resourceName), eq(is.getNumPartitions()), eq("OnlineOffline"), eq("CUSTOMIZED"), eq(is.getBucketSize())); doNothing().when(helixAdmin).setResourceIdealState(eq(CLUSTER), eq(resourceName), eq(is)); serviceImpl.loadFileSet(request).apply(); ArgumentCaptor<FileSetInfo> fileSetInfoCaptor = ArgumentCaptor.forClass(FileSetInfo.class); verify(zkManager).setFileSetInfo(eq(request.getFileSet()), fileSetInfoCaptor.capture()); FileSetInfo capturedInfo = fileSetInfoCaptor.getValue(); assertEquals(request.getFileSet(), capturedInfo.fileSetName); assertEquals(request.getExpectedNumPartitions(), capturedInfo.servingInfo.numPartitions); assertEquals(Lists.newArrayList(fileSetInfo.servingInfo), (ArrayList)capturedInfo.oldServingInfoList); }