All Downloads are FREE. Search and download functionalities are using the official Maven repository.

io.soabase.zookeeper.discovery.ZooKeeperDiscovery Maven / Gradle / Ivy

There is a newer version: 0.11.2
Show newest version
/**
 * Copyright 2014 Jordan Zimmerman
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package io.soabase.zookeeper.discovery;

import com.google.common.base.Function;
import com.google.common.base.Objects;
import com.google.common.base.Predicate;
import com.google.common.cache.CacheBuilder;
import com.google.common.cache.CacheLoader;
import com.google.common.cache.LoadingCache;
import com.google.common.cache.RemovalListener;
import com.google.common.cache.RemovalNotification;
import com.google.common.collect.Collections2;
import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.net.HostAndPort;
import io.dropwizard.lifecycle.Managed;
import io.dropwizard.setup.Environment;
import io.soabase.core.SoaBundle;
import io.soabase.core.SoaFeatures;
import io.soabase.core.SoaInfo;
import io.soabase.core.features.discovery.DiscoveryInstance;
import io.soabase.core.features.discovery.ExtendedDiscovery;
import io.soabase.core.features.discovery.ForcedState;
import io.soabase.core.features.discovery.HealthyState;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.utils.CloseableUtils;
import org.apache.curator.x.discovery.InstanceFilter;
import org.apache.curator.x.discovery.ServiceDiscovery;
import org.apache.curator.x.discovery.ServiceDiscoveryBuilder;
import org.apache.curator.x.discovery.ServiceInstance;
import org.apache.curator.x.discovery.ServiceInstanceBuilder;
import org.apache.curator.x.discovery.ServiceProvider;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import javax.annotation.Nullable;
import java.util.Collection;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;

// TODO
public class ZooKeeperDiscovery extends CacheLoader> implements ExtendedDiscovery, Managed, RemovalListener>
{
    private final Logger log = LoggerFactory.getLogger(getClass());
    private final ServiceDiscovery discovery;
    private final LoadingCache> providers;
    private final AtomicReference> us = new AtomicReference<>();
    private final String bindAddress;
    private final SoaInfo soaInfo;
    private final Environment environment;

    private static class FoundInstance
    {
        final ServiceInstance instance;
        final ServiceProvider provider;

        FoundInstance(ServiceInstance instance, ServiceProvider provider)
        {
            this.instance = instance;
            this.provider = provider;
        }
    }

    public ZooKeeperDiscovery(CuratorFramework curator, ZooKeeperDiscoveryFactory factory, SoaInfo soaInfo, Environment environment, Collection instanceDeploymentGroups)
    {
        this.soaInfo = soaInfo;
        this.environment = environment;
        bindAddress = factory.getBindAddress();
        providers = CacheBuilder.newBuilder()
            .expireAfterWrite(5, TimeUnit.MINUTES)  // TODO config
            .removalListener(this)
            .build(this);

        try
        {
            HashMap metaData = Maps.newHashMap();
            Payload.addDeploymentGroups(metaData, instanceDeploymentGroups);
            Payload payload = new Payload(soaInfo.getAdminPort().getHostText(), soaInfo.getAdminPort().getPort(), metaData, ForcedState.CLEARED, HealthyState.UNHEALTHY);  // initially unhealthy

            us.set(buildInstance(payload, null));

            discovery = ServiceDiscoveryBuilder
                .builder(Payload.class)
                .basePath(factory.getZookeeperPath())
                .client(curator)
                .watchInstances(true)
                .build();
        }
        catch ( Exception e )
        {
            log.error("Could not build discovery instance", e);
            throw new RuntimeException(e);
        }
    }

    @Override
    public Collection getServiceNames()
    {
        return providers.asMap().keySet();
    }

    @Override
    public Collection queryForServiceNames()
    {
        try
        {
            // TODO - possibly cache this
            return discovery.queryForNames();
        }
        catch ( Exception e )
        {
            log.error("Could not query for names", e);
            throw new RuntimeException(e);
        }
    }

    @Override
    public void setHealthyState(HealthyState newHealthyState)
    {
        Payload payload = us.get().getPayload();
        updateRegistration(new Payload(null, payload.getAdminPort(), payload.getMetaData(), payload.getForcedState(), newHealthyState));
    }

    @Override
    public void setMetaData(Map newMetaData)
    {
        Payload payload = us.get().getPayload();
        updateRegistration(new Payload(null, payload.getAdminPort(), newMetaData, payload.getForcedState(), payload.getHealthyState()));
    }

    @Override
    public void setForcedState(String serviceName, String instanceId, ForcedState forcedState)
    {
        try
        {
            ServiceInstance foundInstance = discovery.queryForInstance(serviceName, instanceId);
            if ( foundInstance != null )
            {
                DiscoveryInstance soaInstance = toSoaInstance(foundInstance);
                Payload oldPayload = foundInstance.getPayload();
                Payload newPayload = new Payload(null, oldPayload.getAdminPort(), oldPayload.getMetaData(), forcedState, oldPayload.getHealthyState());
                ServiceInstance updatedInstance = buildInstance(serviceName, HostAndPort.fromParts(soaInstance.getHost(), soaInstance.getPort()), newPayload, instanceId, soaInstance.getHost());
                discovery.updateService(updatedInstance);
            } // TODO else?
        }
        catch ( Exception e )
        {
            log.error("Could not update service: " + (serviceName + ":" + instanceId), e);
            throw new RuntimeException(e);
        }
    }

    @Override
    public ServiceProvider load(String serviceName) throws Exception
    {
        InstanceFilter filter = new InstanceFilter()
        {
            @Override
            public boolean apply(ServiceInstance instance)
            {
                Payload payload = instance.getPayload();
                if ( payload.getForcedState() == ForcedState.CLEARED )
                {
                    return (payload.getHealthyState() == HealthyState.HEALTHY);
                }
                return (payload.getForcedState() == ForcedState.REGISTER);
            }
        };
        ServiceProvider provider = discovery
            .serviceProviderBuilder()
            .serviceName(serviceName)
            .additionalFilter(filter)
            .build();
        provider.start();
        return provider;
    }

    @Override
    public void onRemoval(RemovalNotification> notification)
    {
        CloseableUtils.closeQuietly(notification.getValue());
    }

    @Override
    public Collection queryForAllInstances(String serviceName)
    {
        try
        {
            Collection> serviceInstances = discovery.queryForInstances(serviceName);
            Iterable transformed = Iterables.transform(serviceInstances, new Function, DiscoveryInstance>()
            {
                @Nullable
                @Override
                public DiscoveryInstance apply(ServiceInstance instance)
                {
                    return toSoaInstance(instance);
                }
            });
            return Lists.newArrayList(transformed);
        }
        catch ( Exception e )
        {
            log.error("Could query all instances for service: " + serviceName, e);
            throw new RuntimeException(e);
        }
    }

    @Override
    public Collection getAllInstances(String serviceName)
    {
        try
        {
            // TODO - validate service name
            ServiceProvider provider = providers.get(serviceName);
            Collection> allInstances = provider.getAllInstances();
            return Collections2.transform(allInstances, new Function, DiscoveryInstance>()
            {
                @Nullable
                @Override
                public DiscoveryInstance apply(@Nullable ServiceInstance instance)
                {
                    return toSoaInstance(instance);
                }
            });
        }
        catch ( Exception e )
        {
            log.error("Could not get service: " + serviceName, e);
            throw new RuntimeException(e);
        }
    }

    @Override
    public DiscoveryInstance getInstance(String serviceName)
    {
        ServiceInstance instance;
        try
        {
            // TODO - validate service name
            ServiceProvider provider = providers.get(serviceName);
            instance = provider.getInstance();
            return toSoaInstance(instance);
        }
        catch ( Exception e )
        {
            log.error("Could not service instance: " + serviceName, e);
            throw new RuntimeException(e);
        }
    }

    @Override
    public void noteError(String serviceName, final DiscoveryInstance errorInstance, int statusCode, Throwable exception)
    {
        FoundInstance foundInstance = findInstanceFromProvider(serviceName, errorInstance);
        if ( foundInstance != null )
        {
            foundInstance.provider.noteError(foundInstance.instance);
        }
    }

    @Override
    public void noteSuccess(String serviceName, DiscoveryInstance instance, int statusCode)
    {
        // NOP
    }

    @Override
    public void start() throws Exception
    {
        discovery.start();
        if ( soaInfo.isRegisterInDiscovery() )
        {
            discovery.registerService(us.get());
        }
    }

    @Override
    public void stop() throws Exception
    {
        providers.invalidateAll();
        CloseableUtils.closeQuietly(discovery);
    }

    private FoundInstance findInstanceFromProvider(final String serviceName, final DiscoveryInstance instanceToFind)
    {
        ServiceInstance foundInstance = null;
        ServiceProvider provider = providers.getUnchecked(serviceName);
        if ( provider != null )
        {
            try
            {
                foundInstance = Iterables.find
                    (
                        provider.getAllInstances(),
                        new Predicate>()
                        {
                            @Override
                            public boolean apply(ServiceInstance instance)
                            {
                                SoaFeatures soaFeatures = SoaBundle.getFeatures(environment);
                                return soaFeatures.getDeploymentGroupManager().isAnyGroupEnabled(serviceName, instance.getPayload().getDeploymentGroups()) && instanceToFind.getId().equals(instance.getId());
                            }
                        },
                        null
                    );
            }
            catch ( Exception e )
            {
                log.error("Could not find service: " + (serviceName + ":" + instanceToFind.getId()), e);
                throw new RuntimeException(e);
            }
        }
        return (foundInstance != null) ? new FoundInstance(foundInstance, provider) : null;
    }

    private DiscoveryInstance toSoaInstance(ServiceInstance instance)
    {
        if ( instance == null )
        {
            return null;
        }

        Payload payload = instance.getPayload();
        int port = Objects.firstNonNull(instance.getPort(), Objects.firstNonNull(instance.getSslPort(), 0));
        return new DiscoveryInstanceImpl(instance.getId(), instance.getAddress(), port, instance.getSslPort() != null, payload);
    }

    private void updateRegistration(Payload newPayload)
    {
        if ( !soaInfo.isRegisterInDiscovery() )
        {
            return;
        }

        ServiceInstance localUs = us.get();
        Payload currentPayload = localUs.getPayload();
        if ( !newPayload.equals(currentPayload) )
        {
            try
            {
                ServiceInstance updatedInstance = buildInstance(newPayload, localUs.getId());
                us.set(updatedInstance);
                discovery.updateService(updatedInstance);
            }
            catch ( Exception e )
            {
                log.error("Could not update registration for local instance: " + localUs, e);
                throw new RuntimeException(e);
            }
        }
    }

    private ServiceInstance buildInstance(Payload payload, String id) throws Exception
    {
        return buildInstance(soaInfo.getServiceName(), soaInfo.getMainPort(), payload, id, null);
    }

    private ServiceInstance buildInstance(String serviceName, HostAndPort mainPort, Payload payload, String id, String address) throws Exception
    {
        ServiceInstanceBuilder builder = ServiceInstance.builder()
            .name(serviceName)
            .payload(payload)
            .address(mainPort.getHostText())
            .port(mainPort.getPort())
            ;
        if ( id != null )
        {
            builder = builder.id(id);
        }
        if ( address != null )
        {
            builder = builder.address(address);
        }
        else if ( bindAddress != null )
        {
            builder = builder.address(bindAddress);
        }
        return builder.build();
    }
}




© 2015 - 2024 Weber Informatics LLC | Privacy Policy