Extending Hazelcast MultiMap

package com.example.hazelcast.mmex;

import com.google.common.collect.*;
import com.hazelcast.core.*;
import java.util.*;
import java.util.concurrent.*;


/**
 * Adds putAll function for MultiMap that batches all puts
 * so that they occur on same node.
 */
public class MultiMapEx<K, V> {

    /**
     * Distributed task to perform batch put.
     */
    static class Task<K, V> implements Callable<Void>,
                                       java.io.Serializable,
                                       HazelcastInstanceAware {
        final K key;
        final List<V> values;
        final String mapName;
        transient volatile HazelcastInstance hz;


        Task(K key, List<V> values, String mapName) {
            this.key = key;
            this.values = values;
            this.mapName = mapName;
        }


        @Override
        public Void call() {
            MultiMap<K, V> map = hz.getMultiMap(mapName);
            for (V value : values) {
                map.put(key, value);
            }
            return null;
        }

        @Override
        public void setHazelcastInstance(HazelcastInstance hz) {
            this.hz = hz;
        }
    }

    /** Name of common executor service for all MultiMapEx instances */
    private static final String EXEC_NAME = "MultiMapEx.executorService";


    final HazelcastInstance hz;
    final String mapName;
    final IExecutorService exec;


    public MultiMapEx(HazelcastInstance hz, String mapName) {
        this.hz = hz;
        this.mapName = mapName;
        this.exec = hz.getExecutorService(EXEC_NAME);
    }


    /**
     * Access to underlying MultiMap.
     */
    public MultiMap<K, V> getMultiMap() {
        return hz.getMultiMap(mapName);
    }

    public void putAll(final K key, List<V> values)
            throws InterruptedException, ExecutionException {

        Task<K, V> task = new Task<K, V>(key, values, mapName);
        exec.submitToKeyOwner(task, key).get();
    }


    /**
     * Test program, creates two HazelcastInstances,
     * and runs a putAll followed by a get on each of them
     * to make sure that we run it on an instance that is
     * not the key owner.
     */
    public static void main(String... args) throws Exception {
        HazelcastInstance hz1 = Hazelcast.newHazelcastInstance();
        HazelcastInstance hz2 = Hazelcast.newHazelcastInstance();
        try {
            for (HazelcastInstance hz : Hazelcast.getAllHazelcastInstances()) {
                MultiMapEx<String, Integer> mme = new MultiMapEx<String, Integer>(hz, "testmap");
                mme.putAll("foo", ImmutableList.<Integer>of(1, 2, 3));
                ImmutableList<Integer> values = ImmutableList.copyOf(mme.getMultiMap().get("foo"));
                System.out.println("foo values: " + values);
            }
        } finally {
            Hazelcast.shutdownAll();
        }
    }
}