package com.ge.snowizard.discovery.client;

import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.base.Preconditions.checkArgument;
import java.util.Collection;
import javax.annotation.concurrent.ThreadSafe;
import org.apache.curator.x.discovery.DownInstancePolicy;
import org.apache.curator.x.discovery.ProviderStrategy;
import org.apache.curator.x.discovery.ServiceCache;
import org.apache.curator.x.discovery.ServiceDiscovery;
import org.apache.curator.x.discovery.ServiceInstance;
import org.apache.curator.x.discovery.ServiceProvider;
import com.ge.snowizard.discovery.core.InstanceMetadata;

@ThreadSafe
public class DiscoveryClient {

    private final ServiceDiscovery<InstanceMetadata> discovery;
    private final ServiceProvider<InstanceMetadata> provider;

    /**
     * Constructor
     *
     * @param serviceName
     *            name of the service to monitor
     * @param discovery
     *            {@link ServiceDiscovery}
     * @param downInstancePolicy
     *            {@link DownInstancePolicy} to use when marking instances as
     *            down
     * @param providerStrategy
     *            {@link ProviderStrategy} to use when selecting an instance
     */
    public DiscoveryClient(final String serviceName,
            final ServiceDiscovery<InstanceMetadata> discovery,
            final DownInstancePolicy downInstancePolicy,
            final ProviderStrategy<InstanceMetadata> providerStrategy) {
        checkNotNull(serviceName);
        checkArgument(!serviceName.isEmpty(), "serviceName cannot be empty");
        checkNotNull(providerStrategy);

        this.discovery = checkNotNull(discovery);

        this.provider = discovery.serviceProviderBuilder()
                .serviceName(serviceName)
                .downInstancePolicy(downInstancePolicy)
                .providerStrategy(providerStrategy).build();
    }

    /**
     * Return a list of discoverable services
     *
     * @return Collection of service names
     */
    public Collection<String> getServices() throws Exception {
        return discovery.queryForNames();
    }

    /**
     * Return the running instances for the service.
     *
     * @return Collection of service instances
     */
    public Collection<ServiceInstance<InstanceMetadata>> getInstances(
            final String serviceName) throws Exception {
        return discovery.queryForInstances(serviceName);
    }

    /**
     * Return an instance of this service.
     *
     * @return ServiceInstance
     * @throws Exception
     */
    public ServiceInstance<InstanceMetadata> getInstance() throws Exception {
        return provider.getInstance();
    }

    /**
     * Note an error when connecting to a service instance.
     *
     * @param instance
     *            {@link ServiceInstance} that is causing the error.
     */
    public void noteError(final ServiceInstance<InstanceMetadata> instance) {
        provider.noteError(instance);
    }

    /**
     * Start the internal {@link ServiceProvider} and {@link ServiceCache}
     *
     * @throws Exception
     */
    public void start() throws Exception {
        provider.start();
    }

    /**
     * Stop the internal {@link ServiceProvider} and {@link ServiceCache}
     *
     * @throws Exception
     */
    public void close() throws Exception {
        provider.close();
    }
}

@@ -168,16 +168,70 @@
 gument;%0A
+import java.io.Closeable;%0Aimport java.io.IOException;%0A
 import j
@@ -630,16 +630,73 @@
 ovider;%0A
+import org.slf4j.Logger;%0Aimport org.slf4j.LoggerFactory;%0A
 import c
@@ -786,18 +786,141 @@
 yClient 
-%7B%0A
+implements Closeable %7B%0A%0A    private static final Logger LOGGER = LoggerFactory%0A            .getLogger(DiscoveryClient.class);
 %0A    pri
@@ -1037,16 +1037,72 @@
 rovider;
+%0A    private final ServiceCache%3CInstanceMetadata%3E cache;
 %0A%0A    /*
@@ -1124,24 +1124,25 @@
 uctor%0A     *
+ 
 %0A     * @par
@@ -2191,16 +2191,97 @@
 build();
+%0A%0A        this.cache = discovery.serviceCacheBuilder().name(serviceName).build();
 %0A    %7D%0A%0A
@@ -2336,24 +2336,25 @@
 vices%0A     *
+ 
 %0A     * @ret
@@ -2564,32 +2564,33 @@
  service.%0A     *
+ 
 %0A     * @return 
@@ -2813,32 +2813,291 @@
 ceName);%0A    %7D%0A%0A
+    /**%0A     * Return a cached list of the running instances for the service.%0A     * %0A     * @return Collection of service instances%0A     */%0A    public Collection%3CServiceInstance%3CInstanceMetadata%3E%3E getInstances() %7B%0A        return cache.getInstances();%0A    %7D%0A%0A
     /**%0A     * R
@@ -3129,32 +3129,33 @@
  service.%0A     *
+ 
 %0A     * @return 
@@ -3397,24 +3397,25 @@
 ance.%0A     *
+ 
 %0A     * @par
@@ -3707,32 +3707,33 @@
 iceCache%7D%0A     *
+ 
 %0A     * @throws 
@@ -3815,24 +3815,47 @@
 er.start();%0A
+        cache.start();%0A
     %7D%0A%0A    /
@@ -3941,73 +3941,146 @@
    *
+/
 %0A    
- * @throws Exception%0A     */%0A    public void close() throws 
+@Override%0A    public void close() %7B%0A        try %7B%0A            cache.close();%0A            provider.close();%0A        %7D catch (final IO
 Exce
@@ -4076,32 +4076,35 @@
 inal IOException
+ e)
  %7B%0A        provi
@@ -4098,33 +4098,74 @@
         
-provider.close();
+    LOGGER.error(%22Unable to close provider%22, e);%0A        %7D
 %0A    %7D%0A%7D

Implement Closeable and use a ServiceCache for frequent lookups