co.cask.cdap.common.zookeeper.coordination.DefaultResourceAssigner.java Source code

Introduction

Here is the source code for co.cask.cdap.common.zookeeper.coordination.DefaultResourceAssigner.java

Source

/*
 * Copyright  2014 Cask Data, Inc.
 *
 * 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 co.cask.cdap.common.zookeeper.coordination;

import com.google.common.collect.Maps;
import com.google.common.collect.Multimap;
import com.google.common.collect.Multimaps;

import java.util.Map;

/**
 * Implementation of {@link ResourceAssigner} that stores assignment with multimap.
 *
 * @param <T> Type of resource handler.
 */
final class DefaultResourceAssigner<T> implements ResourceAssigner<T> {

    private final Multimap<T, PartitionReplica> assignments;
    private final Map<PartitionReplica, T> replicaToHandler;

    /**
     * Constructor that takes the giving assignments as the current set of assignments.
     *
     * @param assignments Currently set of assignments. This will get modified if any of the set methods are called.
     */
    static <T> ResourceAssigner<T> create(Multimap<T, PartitionReplica> assignments) {
        return new DefaultResourceAssigner<>(assignments);
    }

    private DefaultResourceAssigner(Multimap<T, PartitionReplica> assignments) {
        this.assignments = assignments;
        this.replicaToHandler = Maps.newHashMap();
        for (Map.Entry<T, PartitionReplica> entry : assignments.entries()) {
            replicaToHandler.put(entry.getValue(), entry.getKey());
        }
    }

    @Override
    public Multimap<T, PartitionReplica> get() {
        return Multimaps.unmodifiableMultimap(assignments);
    }

    @Override
    public T getHandler(String partition, int replica) {
        return replicaToHandler.get(new PartitionReplica(partition, replica));
    }

    @Override
    public void set(T handler, String partition, int replica) {
        // Remove existing assignment for the given partitionReplica
        set(handler, new PartitionReplica(partition, replica));
    }

    @Override
    public void set(T handler, PartitionReplica partitionReplica) {
        T oldHandler = replicaToHandler.remove(partitionReplica);
        if (oldHandler != null) {
            assignments.remove(oldHandler, partitionReplica);
        }

        assignments.put(handler, partitionReplica);
        replicaToHandler.put(partitionReplica, handler);
    }
}