CassandraLoadBalancerStore implements LoadBalancerStore { @Override public Mono<Void> addOrUpdateTargets(Collection<LoadBalancerTargetState> targets) { List<Mono<Void>> insertOperations = targets.stream() .map(this::addOrUpdateTarget) .collect(Collectors.toList()); int limit = configuration.getLoadBalancerWriteConcurrencyLimit(); return Flux.mergeSequentialDelayError(insertOperations, limit, limit) .ignoreElements() .doOnSubscribe(ignored -> { Map<String, Long> countPerLoadBalancer = targets.stream().collect(Collectors.groupingBy( target -> target.getLoadBalancerTarget().getLoadBalancerId(), Collectors.counting() )); logger.info("Inserting/updating targets: {}", countPerLoadBalancer); logger.debug("Inserting/updating {} targets. Details: {}", targets.size(), targets); }); } @Inject  CassandraLoadBalancerStore(CassandraStoreConfiguration configuration,
                                      @Named(LOAD_BALANCER_SANITIZER) EntitySanitizer entitySanitizer,
                                      @Named(LOAD_BALANCER_CASSANDRA_SESSION) Session session); @Activator void init(); @Override Observable<JobLoadBalancer> getAssociatedLoadBalancersForJob(String jobId); @Override Set<JobLoadBalancer> getAssociatedLoadBalancersSetForJob(String jobId); @Override List<JobLoadBalancerState> getAssociations(); @Override List<JobLoadBalancer> getAssociationsPage(int offset, int limit); @Override Completable addOrUpdateLoadBalancer(JobLoadBalancer jobLoadBalancer, JobLoadBalancer.State state); @Override Completable removeLoadBalancer(JobLoadBalancer jobLoadBalancer); @Override int getNumLoadBalancersForJob(String jobId); @Override Mono<Void> addOrUpdateTargets(Collection<LoadBalancerTargetState> targets); @Override Mono<Void> removeDeregisteredTargets(Collection<LoadBalancerTarget> toRemove); @Override Flux<LoadBalancerTargetState> getLoadBalancerTargets(String loadBalancerId); static final String LOAD_BALANCER_CASSANDRA_SESSION; }
@Test(timeout = TIMEOUT_MS) public void testUpdateTarget() throws Exception { Session session = cassandraCQLUnit.getSession(); BoundStatement countStmt = session.prepare("SELECT COUNT(*) FROM load_balancer_targets;").bind(); PreparedStatement stateStmt = session.prepare("SELECT state FROM load_balancer_targets WHERE load_balancer_id = ? AND ip_address = ?;"); assertThat(session.execute(countStmt).one().getLong(0)).isEqualTo(0); LoadBalancerTarget target = new LoadBalancerTarget("lb-1", "task-1", "1.1.1.1"); CassandraLoadBalancerStore store = getInitdStore(); store.addOrUpdateTargets(target.withState(LoadBalancerTarget.State.REGISTERED)).block(); assertThat(session.execute(countStmt).one().getLong(0)).isEqualTo(1); Row registered = session.execute(stateStmt.bind("lb-1", "1.1.1.1")).one(); assertThat(registered.getString("state")).isEqualTo("REGISTERED"); store.addOrUpdateTargets(target.withState(LoadBalancerTarget.State.DEREGISTERED)).block(); assertThat(session.execute(countStmt).one().getLong(0)).isEqualTo(1); Row deregistered = session.execute(stateStmt.bind("lb-1", "1.1.1.1")).one(); assertThat(deregistered.getString("state")).isEqualTo("DEREGISTERED"); }