package com.outbrain.gomjabbar.targets; import com.outbrain.gomjabbar.config.ConfigParser; import com.outbrain.ob1k.concurrent.ComposableFuture; import com.outbrain.ob1k.concurrent.ComposableFutures; import com.outbrain.ob1k.consul.ConsulAPI; import com.outbrain.ob1k.consul.ConsulCatalog; import com.outbrain.ob1k.consul.ConsulHealth; import com.outbrain.ob1k.consul.HealthInfoInstance; import org.apache.commons.collections.CollectionUtils; import org.apache.commons.collections.MapUtils; import org.apache.commons.lang3.tuple.Pair; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import java.io.IOException; import java.net.URL; import java.util.Collection; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Objects; import java.util.concurrent.ExecutionException; import java.util.concurrent.ThreadLocalRandom; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.function.Function; import java.util.stream.Collectors; /** * @author Eran Harel */ public class ConsulTargetsCache implements TargetsCollector { private static final Logger log = LoggerFactory.getLogger(ConsulTargetsCache.class); private static final int BATCH_SIZE = 20; private final ConsulHealth health; private final ConsulCatalog catalog; private final TargetFilters targetFilters; private final Reloader reloader = new Reloader(); private volatile Map<String, Map<String, List<HealthInfoInstance>>> cache = null; public ConsulTargetsCache(final ConsulHealth health, final ConsulCatalog catalog, final TargetFilters targetFilters) { this.catalog = catalog; this.health = Objects.requireNonNull(health, "health must not be null"); this.targetFilters = Objects.requireNonNull(targetFilters, "targetFilters must not be null"); reloader.start(); } @Override public ComposableFuture<Target> chooseTarget() { return reloader.reloadFuture.map(__ -> { final String dc = chooseDC(); final String module = chooseModule(dc); return chooseTarget(dc, module); }); } private String chooseDC() { if (cache.isEmpty()) { throw new IllegalStateException("No data centers are present or all are filtered out"); } return randomElement(cache.keySet()); } private String chooseModule(final String dc) { return randomElement(cache.get(dc).keySet()); } private Target chooseTarget(final String dc, final String module) { final List<HealthInfoInstance> instances = cache.get(dc).get(module); final HealthInfoInstance randomInstance = randomElement(instances); return new Target(randomInstance.Node.Node, randomInstance.Service.Service, instances.size(), randomInstance.Service.Tags); } private <T> T randomElement(final Collection<T> collection) { return collection.stream() .skip(ThreadLocalRandom.current().nextInt(collection.size())) .findFirst() .orElseThrow(() -> new RuntimeException("Unexpected missing elemenet o_0")); } private ComposableFuture<?> reloadAsync() { final ComposableFuture<Map<String, Map<String, List<HealthInfoInstance>>>> cacheFuture = catalog.datacenters() .map(this::fetchDcServiceMappingAsync) .flatMap(dc2serviceInstancesFutures -> ComposableFutures.all(dc2serviceInstancesFutures.values()) .map(__ -> transformInstanceFuturesMap(dc2serviceInstancesFutures))); cacheFuture.consume(newCacheTry -> { if (newCacheTry.isFailure()) { log.error("Failed to fetch targets", newCacheTry.getError()); } cache = newCacheTry.getOrElse(HashMap::new); }); return cacheFuture; } private Map<String, ComposableFuture<Map<String, List<HealthInfoInstance>>>> fetchDcServiceMappingAsync(final Collection<String> dcs) { log.debug("consul dcs: {}", dcs); return dcs.stream().filter(targetFilters.dcFilter()) .collect(Collectors.toMap(Function.identity(), this::service2instances)); } private Map<String, Map<String, List<HealthInfoInstance>>> transformInstanceFuturesMap(final Map<String, ComposableFuture<Map<String, List<HealthInfoInstance>>>> dc2serviceInstancesFutures) { final Map<String, Map<String, List<HealthInfoInstance>>> dc2serviceInstances = dc2serviceInstancesFutures.entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, e -> { try { return e.getValue().recover(t -> { log.error("Failed to load targets for DC=" + e.getKey(), t); return new HashMap<>(); }).get(); } catch (InterruptedException | ExecutionException ex) { throw new RuntimeException("shouldn't happen as we're in the future map callback handler", ex); } })); dc2serviceInstances.entrySet().removeIf(e -> MapUtils.isEmpty(e.getValue())); return dc2serviceInstances; } private ComposableFuture<Map<String, List<HealthInfoInstance>>> service2instances(final String dc) { log.debug("Fetching services for {} dc ", dc); return catalog.services(dc) .flatMap(service2tags -> { final List<String> filteredServices = service2tags.keySet() .stream() .filter(targetFilters.moduleFilter()) .collect(Collectors.toList()); return ComposableFutures.batch(filteredServices, BATCH_SIZE, service -> fetchServiceInstancesAsync(dc, service)) .map(pairs -> pairs.stream() .filter(p -> !CollectionUtils.isEmpty(p.getRight())) .collect(Collectors.toMap(Pair::getLeft, Pair::getRight))); }); } private ComposableFuture<Pair<String, List<HealthInfoInstance>>> fetchServiceInstancesAsync(final String dc, final String service) { return health.fetchInstancesHealth(service, dc) .map(instances -> Pair.of(service, filterInstances(instances))); } private List<HealthInfoInstance> filterInstances(final Collection<HealthInfoInstance> instances) { return instances.stream().filter(targetFilters.instanceFilter()).collect(Collectors.toList()); } // debug... private void print() { reloader.reloadFuture.consume(__ -> cache.entrySet().forEach(dc2services -> { System.out.println(dc2services.getKey()); dc2services.getValue().entrySet().forEach(service2instances -> { System.out.println("\t" + service2instances.getKey()); service2instances.getValue().forEach(i -> System.out.println("\t\t" + i.Service.Tags)); }); })); } public static void main(final String[] args) throws IOException, ExecutionException, InterruptedException { final URL configFileUrl = new URL("file:./config-template.yaml"); final TargetFilters targetFilters = ConfigParser.parseConfiguration(configFileUrl).targetFilters; final ConsulTargetsCache consulTargetsCache = new ConsulTargetsCache(ConsulAPI.getHealth(), ConsulAPI.getCatalog(), targetFilters); // consulTargetsCache.print(); for (int i = 0; i < 1000; i++) { System.out.println(consulTargetsCache.chooseTarget().get()); Thread.sleep(2000); } } private class Reloader { private static final int RELOAD_DELAY_MINUTES = 5; private final AtomicBoolean isReloading = new AtomicBoolean(false); private volatile ComposableFuture<?> reloadFuture; private void start() { reload(); } private boolean reload() { if (isReloading.compareAndSet(false, true)) { log.info("Reloading cache..."); final ComposableFuture<?> nextReloadFuture = reloadAsync(); if(null == reloadFuture) { reloadFuture = nextReloadFuture; } nextReloadFuture.consume(result -> { reloadFuture = nextReloadFuture; if(isReloading.compareAndSet(true, false)) { log.info("Scheduling the next cache reload in {} minutes", RELOAD_DELAY_MINUTES); ComposableFutures.schedule(this::reload, RELOAD_DELAY_MINUTES, TimeUnit.MINUTES); } }); return true; } log.debug("Cache reloading is already in progress... skipping..."); return false; } } }