subprojects/s4-comm/src/main/java/org/apache/s4/comm/topology/AssignmentFromHelix.java

package org.apache.s4.comm.topology;

import java.net.InetAddress;
import java.net.UnknownHostException;
import java.util.List;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

import org.apache.helix.BaseDataAccessor;
import org.apache.helix.HelixAdmin;
import org.apache.helix.HelixDataAccessor;
import org.apache.helix.HelixManager;
import org.apache.helix.HelixManagerFactory;
import org.apache.helix.InstanceType;
import org.apache.helix.PropertyKey.Builder;
import org.apache.helix.ZNRecord;
import org.apache.helix.manager.zk.ZKHelixAdmin;
import org.apache.helix.manager.zk.ZKHelixDataAccessor;
import org.apache.helix.manager.zk.ZNRecordSerializer;
import org.apache.helix.manager.zk.ZkBaseDataAccessor;
import org.apache.helix.manager.zk.ZkClient;
import org.apache.helix.model.InstanceConfig;
import org.apache.helix.participant.statemachine.StateModel;
import org.apache.helix.participant.statemachine.StateModelFactory;
import org.apache.s4.comm.helix.TaskStateModelFactory;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.inject.Inject;
import com.google.inject.Singleton;
import com.google.inject.name.Named;

@Singleton
public class AssignmentFromHelix implements Assignment {
    private static final Logger logger = LoggerFactory.getLogger(AssignmentFromHelix.class);

    private String clusterName;
    private final String zookeeperAddress;
    private String machineId;
    private HelixManager zkHelixManager;

    private HelixDataAccessor helixDataAccessor;
    AtomicReference<ClusterNode> clusterNodeRef;
    private final Lock lock;
    private final AtomicBoolean currentlyOwningTask;
    private final Condition taskAcquired;

    private final StateModelFactory<? extends StateModel> taskStateModelFactory;

    // private final StateModelFactory<? extends StateModel> appStateModelFactory;

    @Inject
    public AssignmentFromHelix(@Named("s4.cluster.name") String clusterName,
            @Named("s4.instance.name") String instanceName, @Named("s4.cluster.zk_address") String zookeeperAddress)
            throws Exception {
        this.taskStateModelFactory = new TaskStateModelFactory();
        // this.appStateModelFactory = appStateModelFactory;
        this.clusterName = clusterName;
        this.zookeeperAddress = zookeeperAddress;
        machineId = "localhost";
        lock = new ReentrantLock();
        ZkClient zkClient = new ZkClient(zookeeperAddress);
        zkClient.setZkSerializer(new ZNRecordSerializer());
        zkClient.waitUntilConnected(60, TimeUnit.SECONDS);
        BaseDataAccessor<ZNRecord> baseDataAccessor = new ZkBaseDataAccessor<ZNRecord>(zkClient);
        helixDataAccessor = new ZKHelixDataAccessor(clusterName, baseDataAccessor);
        clusterNodeRef = new AtomicReference<ClusterNode>();
        taskAcquired = lock.newCondition();
        currentlyOwningTask = new AtomicBoolean(true);
        try {
            machineId = InetAddress.getLocalHost().getCanonicalHostName();
        } catch (UnknownHostException e) {
            logger.warn("Unable to get hostname", e);
            machineId = "UNKNOWN";
        }
        ClusterNode node = new ClusterNode(-1, Integer.parseInt(instanceName.split("_")[1]), machineId, instanceName);
        clusterNodeRef.set(node);
        currentlyOwningTask.set(true);
    }

    @Inject
    public void init() {
        // joinCluster();
    }

    @Override
    public ClusterNode assignClusterNode() {
        lock.lock();
        try {
            while (!currentlyOwningTask.get()) {
                taskAcquired.awaitUninterruptibly();
            }
        } catch (Exception e) {
            logger.error("Exception while waiting to join the cluster");
            return null;
        } finally {
            lock.unlock();
        }
        return clusterNodeRef.get();
    }

    public void joinClusterOld() {
        lock.lock();
        try {
            Builder keyBuilder = helixDataAccessor.keyBuilder();
            do {
                List<InstanceConfig> instances = helixDataAccessor.getChildValues(keyBuilder.instanceConfigs());
                List<String> liveInstances = helixDataAccessor.getChildNames(keyBuilder.liveInstances());
                for (InstanceConfig instanceConfig : instances) {
                    String instanceName = instanceConfig.getInstanceName();
                    if (!liveInstances.contains(instanceName)) {
                        zkHelixManager = HelixManagerFactory.getZKHelixManager(clusterName, instanceName,
                                InstanceType.PARTICIPANT, zookeeperAddress);
                        zkHelixManager.getStateMachineEngine().registerStateModelFactory("LeaderStandby",
                                taskStateModelFactory);

                        zkHelixManager.connect();
                        ClusterNode node = new ClusterNode(-1, Integer.parseInt(instanceConfig.getPort()), machineId,
                                instanceName);
                        clusterNodeRef.set(node);
                        currentlyOwningTask.set(true);
                        taskAcquired.signalAll();
                        break;
                    }
                }
                if (instances.size() == liveInstances.size()) {
                    System.out.println("No more nodes can join the cluster. Will wait for some node to die.");
                    Thread.sleep(100000);
                }
            } while (!currentlyOwningTask.get());
            System.out.println("Joined the cluster:" + clusterName + " as " + clusterNodeRef.get().getTaskId());
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            lock.unlock();
        }
    }

}