package io.vertx.servicediscovery.zookeeper;

import io.vertx.core.CompositeFuture;
import io.vertx.core.Future;
import io.vertx.core.Handler;
import io.vertx.core.Promise;
import io.vertx.core.Vertx;
import io.vertx.core.impl.ConcurrentHashSet;
import io.vertx.core.json.JsonObject;
import io.vertx.core.logging.Logger;
import io.vertx.core.logging.LoggerFactory;
import io.vertx.servicediscovery.Record;
import io.vertx.servicediscovery.spi.ServiceImporter;
import io.vertx.servicediscovery.spi.ServicePublisher;
import io.vertx.servicediscovery.spi.ServiceType;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.recipes.cache.TreeCache;
import org.apache.curator.framework.recipes.cache.TreeCacheEvent;
import org.apache.curator.framework.recipes.cache.TreeCacheListener;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.curator.x.discovery.ServiceDiscovery;
import org.apache.curator.x.discovery.ServiceDiscoveryBuilder;
import org.apache.curator.x.discovery.ServiceInstance;
import org.apache.zookeeper.KeeperException;

import java.util.*;

A service importer retrieving services from Zookeeper.
Author:Clement Escoffier
/** * A service importer retrieving services from Zookeeper. * * @author <a href="http://escoffier.me">Clement Escoffier</a> */
public class ZookeeperServiceImporter implements ServiceImporter, TreeCacheListener { private static final Logger LOGGER = LoggerFactory.getLogger(ZookeeperServiceImporter.class); private ServicePublisher publisher; private CuratorFramework client; private ServiceDiscovery<JsonObject> discovery; private TreeCache cache; private volatile boolean started; private Set<RegistrationHolder<ServiceInstance<JsonObject>>> registrations = new ConcurrentHashSet<>(); @Override public void start(Vertx vertx, ServicePublisher publisher, JsonObject configuration, Promise<Void> future) { this.publisher = publisher; String connection = Objects.requireNonNull(configuration.getString("connection")); int maxRetries = configuration.getInteger("maxRetries", 3); int baseGraceBetweenRetries = configuration.getInteger("baseSleepTimeBetweenRetries", 1000); String basePath = configuration.getString("basePath", "/discovery"); boolean canBeReadOnly = configuration.getBoolean("canBeReadOnly", true); int connectionTimeoutMs = configuration.getInteger("connectionTimeoutMs", 1000); vertx.<Void>executeBlocking( f -> { try { client = CuratorFrameworkFactory.builder() .canBeReadOnly(canBeReadOnly) .connectString(connection) .connectionTimeoutMs(connectionTimeoutMs) .retryPolicy(new ExponentialBackoffRetry(baseGraceBetweenRetries, maxRetries)) .build(); client.start(); discovery = ServiceDiscoveryBuilder.builder(JsonObject.class) .client(client) .basePath(basePath) .serializer(new JsonObjectSerializer()) .watchInstances(true) .build(); discovery.start(); cache = TreeCache.newBuilder(client, basePath).build(); cache.start(); cache.getListenable().addListener(this); f.complete(); } catch (Exception e) { future.fail(e); } }, ar -> { if (ar.failed()) { future.fail(ar.cause()); } else { Promise<Void> p = Promise.promise(); p.future().setHandler(x -> { if (x.failed()) { future.fail(x.cause()); } else { started = true; future.complete(null); } }); compute(p); } } ); } private synchronized void compute(Promise<Void> done) { List<ServiceInstance<JsonObject>> instances = new ArrayList<>(); try { Collection<String> names = discovery.queryForNames(); for (String name : names) { instances.addAll(discovery.queryForInstances(name)); } } catch (KeeperException.NoNodeException e) { // no services // Continue } catch (Exception e) { if (done != null) { done.fail(e); } else { LOGGER.error("Unable to retrieve service instances from Zookeeper", e); return; } } // Current set Set<RegistrationHolder<ServiceInstance<JsonObject>>> registered = new HashSet<>(registrations); Set<ServiceInstance<JsonObject>> remote = new HashSet<>(instances); List<Future> actions = new ArrayList<>(); RegistrationHolder.filter(registered, instances) .stream() .map(reg -> { Promise<Void> promise = Promise.promise(); publisher.unpublish(reg.record().getRegistration(), v -> { registrations.remove(reg); if (v.succeeded()) { promise.complete(null); } else { promise.fail(v.cause()); } }); return promise.future(); }).forEach(actions::add); RegistrationHolder.filter(remote, registrations) .stream() .map(instance -> { Promise<Void> promise = Promise.promise(); publisher.publish(createRecordForInstance(instance), v -> { if (v.succeeded()) { registrations.add(new RegistrationHolder<>(v.result(), instance)); promise.complete(null); } else { promise.fail(v.cause()); } }); return promise.future(); }).forEach(actions::add); if (done != null) { CompositeFuture.all(actions).setHandler(ar -> { if (ar.succeeded()) { done.complete(null); } else { done.fail(ar.cause()); } }); } } static Record createRecordForInstance(ServiceInstance<JsonObject> instance) { Record record = new Record(); record.setName(instance.getName()); JsonObject payload = instance.getPayload(); record.setMetadata(payload); record.getMetadata().put("zookeeper-service-type", instance.getServiceType().toString()); record.getMetadata().put("zookeeper-address", instance.getAddress()); record.getMetadata().put("zookeeper-registration-time", instance.getRegistrationTimeUTC()); record.getMetadata().put("zookeeper-port", instance.getPort()); record.getMetadata().put("zookeeper-ssl-port", instance.getSslPort()); record.getMetadata().put("zookeeper-id", instance.getId()); record.setLocation(new JsonObject()); if (instance.getUriSpec() != null) { String uri = instance.buildUriSpec(); record.getLocation().put("endpoint", uri); } else { String uri = "http"; if (instance.getSslPort() != null) { uri += "s://" + instance.getAddress() + ":" + instance.getSslPort(); } else if (instance.getPort() != null) { uri += "s://" + instance.getAddress() + ":" + instance.getPort(); } else { uri += "://" + instance.getAddress(); } record.getLocation().put("endpoint", uri); } if (instance.getPort() != null) { record.getLocation().put("port", instance.getPort()); } if (instance.getSslPort() != null) { record.getLocation().put("ssl-port", instance.getSslPort()); } if (instance.getAddress() != null) { record.getLocation().put("address", instance.getAddress()); } record.setType(payload.getString("service-type", ServiceType.UNKNOWN)); return record; } @Override public void close(Handler<Void> closeHandler) { Promise<Void> done = Promise.promise(); unregisterAllServices(done); done.future().setHandler(v -> { try { cache.close(); discovery.close(); client.close(); } catch (Exception e) { // Ignore them } closeHandler.handle(null); }); } @Override public void childEvent(CuratorFramework curatorFramework, TreeCacheEvent treeCacheEvent) throws Exception { if (started) { compute(null); } } private synchronized void unregisterAllServices(Promise<Void> done) { List<Future> list = new ArrayList<>(); new HashSet<>(registrations).forEach(reg -> { Promise<Void> unreg = Promise.promise(); publisher.unpublish(reg.record().getRegistration(), unreg); }); registrations.clear(); CompositeFuture.all(list).setHandler(x -> { if (x.failed()) { done.fail(x.cause()); } else { done.complete(); } }); } }