DefaultLoadBalancerService implements LoadBalancerService { @VisibleForTesting void backfillTargetsToStore() { Observable<Single<LoadBalancer>> fetchLoadBalancerOperations = loadBalancerStore.getAssociations().stream() .map(a -> a.getJobLoadBalancer().getLoadBalancerId()) .map(loadBalancerId -> loadBalancerConnector.getLoadBalancer(loadBalancerId) .retryWhen(RetryHandlerBuilder.retryHandler() .withRetryDelay(10, TimeUnit.MILLISECONDS) .withRetryCount(10) .buildExponentialBackoff() ) ) .collect(Collectors.collectingAndThen(Collectors.toList(), Observable::from)); try { int concurrency = configuration.getStoreBackfillConcurrencyLimit(); ReactorExt.toFlux(Single.mergeDelayError(fetchLoadBalancerOperations, concurrency)) .flatMap(this::backfillTargetsToStore) .ignoreElements() .block(Duration.ofMillis(configuration.getStoreBackfillTimeoutMs())); } catch (Exception e) { Registry registry = runtime.getRegistry(); registry.counter(registry.createId("titus.loadbalancer.service.backfillErrors", "error", e.getClass().getSimpleName() )).increment(); logger.error("Backfill did not complete successfully, missing targets may be orphaned and leak on load balancers (i.e. never deregistered)", e); } } @Inject  DefaultLoadBalancerService(TitusRuntime runtime,
                                      LoadBalancerConfiguration configuration,
                                      LoadBalancerConnector loadBalancerConnector,
                                      LoadBalancerStore loadBalancerStore,
                                      V3JobOperations v3JobOperations,
                                      LoadBalancerJobValidator validator); @VisibleForTesting  DefaultLoadBalancerService(TitusRuntime runtime,
                               LoadBalancerConfiguration configuration,
                               LoadBalancerConnector loadBalancerConnector,
                               LoadBalancerStore loadBalancerStore,
                               LoadBalancerJobOperations loadBalancerJobOperations,
                               LoadBalancerReconciler reconciler,
                               LoadBalancerJobValidator validator,
                               Scheduler scheduler); @Override Observable<String> getJobLoadBalancers(String jobId); @Override Pair<List<JobLoadBalancer>, Pagination> getAllLoadBalancers(Page page); @Override Completable addLoadBalancer(String jobId, String loadBalancerId); @Override Completable removeLoadBalancer(String jobId, String loadBalancerId); @Activator void activate(); @Deactivator void deactivate();  }
@Test(timeout = 30_000) public void backfillsCurrentTargetsToStore() { String jobId = UUID.randomUUID().toString(); String associatedId = "lb-" + UUID.randomUUID().toString(); String dissociatedId = "lb-" + UUID.randomUUID().toString(); String removedId = "lb-" + UUID.randomUUID().toString(); LoadBalancerConfiguration configuration = LoadBalancerTests.mockConfiguration(MIN_TIME_IN_QUEUE_MS); when(configuration.isTargetsToStoreBackfillEnabled()).thenReturn(true); when(configuration.getStoreBackfillConcurrencyLimit()).thenReturn(10); when(configuration.getStoreBackfillTimeoutMs()).thenReturn(5000L); when(client.getLoadBalancer(associatedId)).thenReturn(Single.just( new LoadBalancer(associatedId, LoadBalancer.State.ACTIVE, CollectionsExt.asSet( "1.1.1.1", "2.2.2.2", "3.3.3.3" )) )); when(client.getLoadBalancer(dissociatedId)).thenReturn(Single.just( new LoadBalancer(dissociatedId, LoadBalancer.State.ACTIVE, CollectionsExt.asSet( "4.4.4.4", "5.5.5.5", "6.6.6.6" )) )); when(client.getLoadBalancer(removedId)).thenReturn(Single.just( new LoadBalancer(removedId, LoadBalancer.State.REMOVED, Collections.emptySet()) )); loadBalancerStore.addOrUpdateLoadBalancer( new JobLoadBalancer(jobId, associatedId), JobLoadBalancer.State.ASSOCIATED ).await(); loadBalancerStore.addOrUpdateLoadBalancer( new JobLoadBalancer(jobId, dissociatedId), JobLoadBalancer.State.DISSOCIATED ).await(); loadBalancerStore.addOrUpdateLoadBalancer( new JobLoadBalancer(jobId, removedId), JobLoadBalancer.State.ASSOCIATED ).await(); DefaultLoadBalancerService service = new DefaultLoadBalancerService(runtime, configuration, client, loadBalancerStore, loadBalancerJobOperations, reconciler, validator, testScheduler); service.backfillTargetsToStore(); assertThat(loadBalancerStore.getLoadBalancerTargets(associatedId).collectList().block()) .containsExactlyInAnyOrder( new LoadBalancerTargetState( new LoadBalancerTarget(associatedId, "BACKFILLED", "1.1.1.1"), REGISTERED ), new LoadBalancerTargetState( new LoadBalancerTarget(associatedId, "BACKFILLED", "2.2.2.2"), REGISTERED ), new LoadBalancerTargetState( new LoadBalancerTarget(associatedId, "BACKFILLED", "3.3.3.3"), REGISTERED ) ); assertThat(loadBalancerStore.getLoadBalancerTargets(dissociatedId).collectList().block()) .containsExactlyInAnyOrder( new LoadBalancerTargetState( new LoadBalancerTarget(dissociatedId, "BACKFILLED", "4.4.4.4"), REGISTERED ), new LoadBalancerTargetState( new LoadBalancerTarget(dissociatedId, "BACKFILLED", "5.5.5.5"), REGISTERED ), new LoadBalancerTargetState( new LoadBalancerTarget(dissociatedId, "BACKFILLED", "6.6.6.6"), REGISTERED ) ); assertThat(loadBalancerStore.getLoadBalancerTargets(removedId).collectList().block()).isEmpty(); }