ApplicationCellRouter implements CellRouter { @Override public Optional<Cell> routeKey(JobDescriptor jobDescriptor) { Optional<Cell> pinnedCell = getCellPinnedToJob(jobDescriptor); if (pinnedCell.isPresent()) { return pinnedCell; } String routeKey = routeKeyFor(jobDescriptor); Set<String> antiAffinityNames = StringExt.splitByCommaIntoSet( jobDescriptor.getAttributesMap().get(JobAttributes.JOB_PARAMETER_ATTRIBUTES_CELL_AVOID) ); Optional<Cell> found = selector.select(routeKey, cell -> !antiAffinityNames.contains(cell.getName())); if (found.isPresent()) { return found; } return cellInfoResolver.resolve().stream() .filter(cell -> !antiAffinityNames.contains(cell.getName())) .findAny(); } ApplicationCellRouter(CellInfoResolver cellInfoResolver, TitusFederationConfiguration federationConfiguration); @Override Optional<Cell> routeKey(JobDescriptor jobDescriptor);  }
@Test public void cellRoutingRulesFromConfig() { TitusFederationConfiguration titusFederationConfiguration = mock(TitusFederationConfiguration.class); when(titusFederationConfiguration.getCells()).thenReturn("cell1=hostName1:7001;cell2=hostName2:7002"); when(titusFederationConfiguration.getRoutingRules()).thenReturn("cell1=(app1.*|app2.*);cell2=(app3.*)"); CellInfoResolver cellInfoResolver = new DefaultCellInfoResolver(titusFederationConfiguration); ApplicationCellRouter cellRouter = new ApplicationCellRouter(cellInfoResolver, titusFederationConfiguration); assertThat(cellRouter.routeKey(APP_3).get().getName()).isEqualTo("cell2"); assertThat(cellRouter.routeKey(APP_2).get().getName()).isEqualTo("cell1"); assertThat(cellRouter.routeKey(OTHER).get().getName()).isEqualTo("cell1"); }
@Test public void cellsWithNoRulesCanExist() { TitusFederationConfiguration titusFederationConfiguration = mock(TitusFederationConfiguration.class); when(titusFederationConfiguration.getCells()).thenReturn("cell1=hostName1:7001;cell2=hostName2:7002"); when(titusFederationConfiguration.getRoutingRules()).thenReturn("cell1=(app1.*|app2.*)"); CellInfoResolver cellInfoResolver = new DefaultCellInfoResolver(titusFederationConfiguration); ApplicationCellRouter cellRouter = new ApplicationCellRouter(cellInfoResolver, titusFederationConfiguration); Cell cell = cellRouter.routeKey(APP_2).get(); assertThat(cell.getName()).isEqualTo("cell1"); }
@Test public void rulesWithNonConfiguredCellsAreIgnored() { TitusFederationConfiguration titusFederationConfiguration = mock(TitusFederationConfiguration.class); when(titusFederationConfiguration.getCells()).thenReturn("cell1=hostName1:7001;cell2=hostName2:7002"); when(titusFederationConfiguration.getRoutingRules()).thenReturn("cell1=(app1.*|app2.*);cell2=(app3.*);cell3=(app4.*)"); CellInfoResolver cellInfoResolver = new DefaultCellInfoResolver(titusFederationConfiguration); ApplicationCellRouter cellRouter = new ApplicationCellRouter(cellInfoResolver, titusFederationConfiguration); Cell cell = cellRouter.routeKey(APP_4).get(); assertThat(cell.getName()).isEqualTo("cell1"); }
@Test public void rulesCanBeChangedDynamically() { TitusFederationConfiguration titusFederationConfiguration = mock(TitusFederationConfiguration.class); when(titusFederationConfiguration.getCells()).thenReturn("cell1=hostName1:7001;cell2=hostName2:7002"); when(titusFederationConfiguration.getRoutingRules()).thenReturn("cell1=(app1.*|app2.*);cell2=(app3.*)"); CellInfoResolver cellInfoResolver = new DefaultCellInfoResolver(titusFederationConfiguration); ApplicationCellRouter cellRouter = new ApplicationCellRouter(cellInfoResolver, titusFederationConfiguration); assertThat(cellRouter.routeKey(APP_3).get().getName()).isEqualTo("cell2"); assertThat(cellRouter.routeKey(APP_2).get().getName()).isEqualTo("cell1"); assertThat(cellRouter.routeKey(OTHER).get().getName()).isEqualTo("cell1"); reset(titusFederationConfiguration); when(titusFederationConfiguration.getCells()).thenReturn("cell1=hostName1:7001;cell2=hostName2:7002"); when(titusFederationConfiguration.getRoutingRules()).thenReturn("cell2=(app1.*|app2.*);cell1=(app3.*)"); assertThat(cellRouter.routeKey(APP_3).get().getName()).isEqualTo("cell1"); assertThat(cellRouter.routeKey(APP_2).get().getName()).isEqualTo("cell2"); assertThat(cellRouter.routeKey(OTHER).get().getName()).isEqualTo("cell1"); }
@Test public void jobToCellAffinity() { JobDescriptor withCellAffinity = APP_2.toBuilder() .putAttributes(JobAttributes.JOB_PARAMETER_ATTRIBUTES_CELL_REQUEST, "cell2") .build(); TitusFederationConfiguration titusFederationConfiguration = mock(TitusFederationConfiguration.class); when(titusFederationConfiguration.getCells()).thenReturn("cell1=hostName1:7001;cell2=hostName2:7002"); when(titusFederationConfiguration.getRoutingRules()).thenReturn("cell1=(app1.*|app2.*);cell2=(app3.*)"); CellInfoResolver cellInfoResolver = new DefaultCellInfoResolver(titusFederationConfiguration); ApplicationCellRouter cellRouter = new ApplicationCellRouter(cellInfoResolver, titusFederationConfiguration); assertThat(cellRouter.routeKey(withCellAffinity).get().getName()).isEqualTo("cell2"); }
@Test public void jobCellAntiAffinity() { JobDescriptor withCellAntiAffinity = APP_2.toBuilder() .putAttributes(JobAttributes.JOB_PARAMETER_ATTRIBUTES_CELL_AVOID, "cell4, cell1") .build(); JobDescriptor allRejectedGoesToDefault = APP_2.toBuilder() .putAttributes(JobAttributes.JOB_PARAMETER_ATTRIBUTES_CELL_AVOID, "cell1,cell2") .build(); TitusFederationConfiguration titusFederationConfiguration = mock(TitusFederationConfiguration.class); when(titusFederationConfiguration.getCells()).thenReturn("cell1=hostName1:7001;cell2=hostName2:7002"); when(titusFederationConfiguration.getRoutingRules()).thenReturn("cell1=(app1.*|app2.*);cell2=(app3.*)"); CellInfoResolver cellInfoResolver = new DefaultCellInfoResolver(titusFederationConfiguration); ApplicationCellRouter cellRouter = new ApplicationCellRouter(cellInfoResolver, titusFederationConfiguration); assertThat(cellRouter.routeKey(withCellAntiAffinity).get().getName()).isEqualTo("cell2"); assertThat(cellRouter.routeKey(allRejectedGoesToDefault)).isEmpty(); }