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

com.hazelcast.map.impl.MapKeyLoader Maven / Gradle / Ivy

/*
 * Copyright (c) 2008-2016, Hazelcast, Inc. All Rights Reserved.
 *
 * 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 com.hazelcast.map.impl;

import com.hazelcast.core.ExecutionCallback;
import com.hazelcast.core.IFunction;
import com.hazelcast.core.MapLoader;
import com.hazelcast.core.Member;
import com.hazelcast.map.impl.mapstore.MapStoreContext;
import com.hazelcast.map.impl.operation.LoadStatusOperation;
import com.hazelcast.map.impl.operation.LoadStatusOperationFactory;
import com.hazelcast.map.impl.operation.MapOperation;
import com.hazelcast.map.impl.operation.MapOperationProvider;
import com.hazelcast.map.impl.operation.PartitionCheckIfLoadedOperation;
import com.hazelcast.nio.Address;
import com.hazelcast.nio.serialization.Data;
import com.hazelcast.partition.InternalPartition;
import com.hazelcast.partition.InternalPartitionService;
import com.hazelcast.spi.ExecutionService;
import com.hazelcast.spi.InternalCompletableFuture;
import com.hazelcast.spi.Operation;
import com.hazelcast.spi.OperationService;
import com.hazelcast.spi.impl.AbstractCompletableFuture;
import com.hazelcast.util.FutureUtil;
import com.hazelcast.util.StateMachine;
import com.hazelcast.util.scheduler.CoalescingDelayedTrigger;

import java.io.Closeable;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executor;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

import static com.hazelcast.logging.Logger.getLogger;
import static com.hazelcast.map.impl.MapKeyLoaderUtil.assignRole;
import static com.hazelcast.map.impl.MapKeyLoaderUtil.toBatches;
import static com.hazelcast.map.impl.MapKeyLoaderUtil.toPartition;
import static com.hazelcast.map.impl.MapService.SERVICE_NAME;
import static com.hazelcast.nio.IOUtil.closeResource;
import static com.hazelcast.spi.ExecutionService.MAP_LOAD_ALL_KEYS_EXECUTOR;
import static com.hazelcast.util.IterableUtil.limit;
import static com.hazelcast.util.IterableUtil.map;
import static java.util.concurrent.TimeUnit.SECONDS;

/**
 * Loads keys from a {@link MapLoader} and sends them to all partitions for loading
 */
public class MapKeyLoader {

    private static final long LOADING_TRIGGER_DELAY = SECONDS.toMillis(5);
    private static final long KEY_DISTRIBUTION_TIMEOUT_MINUTES = 15;

    private String mapName;
    private OperationService opService;
    private InternalPartitionService partitionService;
    private IFunction toData;
    private ExecutionService execService;
    private CoalescingDelayedTrigger delayedTrigger;

    private int maxSizePerNode;
    private int maxBatch;
    private int mapNamePartition;
    private int partitionId;
    private boolean hasBackup;

    private LoadFinishedFuture loadFinished = new LoadFinishedFuture(true);
    private MapOperationProvider operationProvider;

    /**
     * Role of this MapKeyLoader
     **/
    enum Role {
        NONE,
        /**
         * Sends out keys to all other partitions
         **/
        SENDER,
        /**
         * Receives keys from sender
         **/
        RECEIVER,
        /**
         * Restarts sending if SENDER fails
         **/
        SENDER_BACKUP
    }

    enum State {
        NOT_LOADED,
        LOADING,
        LOADED
    }

    private final StateMachine role = StateMachine.of(Role.NONE)
            .withTransition(Role.NONE, Role.SENDER, Role.RECEIVER, Role.SENDER_BACKUP)
            .withTransition(Role.SENDER_BACKUP, Role.SENDER);

    private final StateMachine state = StateMachine.of(State.NOT_LOADED)
            .withTransition(State.NOT_LOADED, State.LOADING)
            .withTransition(State.LOADING, State.LOADED, State.NOT_LOADED)
            .withTransition(State.LOADED, State.LOADING);

    public MapKeyLoader(String mapName, OperationService opService, InternalPartitionService ps,
                        ExecutionService execService, IFunction serialize) {
        this.mapName = mapName;
        this.opService = opService;
        this.partitionService = ps;
        this.toData = serialize;
        this.execService = execService;
    }

    public Future startInitialLoad(MapStoreContext mapStoreContext, int partitionId) {

        this.partitionId = partitionId;
        this.mapNamePartition = partitionService.getPartitionId(toData.apply(mapName));
        Role newRole = calculateRole();

        role.nextOrStay(newRole);
        state.next(State.LOADING);

        switch (newRole) {
            case SENDER:
                return sendKeys(mapStoreContext, false);
            case SENDER_BACKUP:
            case RECEIVER:
                return triggerLoading();
            default:
                return loadFinished;
        }
    }

    private Role calculateRole() {
        boolean isPartitionOwner = partitionService.isPartitionOwner(partitionId);
        boolean isMapNamePartition = partitionId == mapNamePartition;
        boolean isMapNamePartitionFirstReplica = false;
        if (hasBackup && isMapNamePartition) {
            InternalPartition partition = partitionService.getPartition(partitionId);
            Address firstReplicaAddress = partition.getReplicaAddress(1);
            Member member = partitionService.getMember(firstReplicaAddress);
            if (member != null) {
                isMapNamePartitionFirstReplica = member.localMember();
            }
        }
        return assignRole(isPartitionOwner, isMapNamePartition, isMapNamePartitionFirstReplica);
    }

    /**
     * Sends keys to all partitions in batches.
     */
    public Future sendKeys(final MapStoreContext mapStoreContext, final boolean replaceExistingValues) {

        if (loadFinished.isDone()) {

            loadFinished = new LoadFinishedFuture();

            Future sent = execService.submit(MAP_LOAD_ALL_KEYS_EXECUTOR, new Callable() {
                @Override
                public Boolean call() throws Exception {
                    sendKeysInBatches(mapStoreContext, replaceExistingValues);
                    return false;
                }
            });

            execService.asCompletableFuture(sent).andThen(loadFinished);
        }

        return loadFinished;
    }

    /**
     * Check if loaded on SENDER partition. Triggers key loading if it hadn't started
     */
    public Future triggerLoading() {

        if (loadFinished.isDone()) {

            loadFinished = new LoadFinishedFuture();

            execService.execute(MAP_LOAD_ALL_KEYS_EXECUTOR, new Runnable() {
                @Override
                public void run() {
                    Operation op = new PartitionCheckIfLoadedOperation(mapName, true);
                    opService.invokeOnPartition(SERVICE_NAME, op, mapNamePartition)
                            .andThen(ifLoadedCallback());
                }
            });
        }

        return loadFinished;
    }

    public Future startLoading(MapStoreContext mapStoreContext, boolean replaceExistingValues) {

        role.nextOrStay(Role.SENDER);

        if (state.is(State.LOADING)) {
            return loadFinished;
        }
        state.next(State.LOADING);

        return sendKeys(mapStoreContext, replaceExistingValues);
    }

    public void trackLoading(boolean lastBatch, Throwable exception) {
        if (lastBatch) {
            state.nextOrStay(State.LOADED);
            if (exception != null) {
                loadFinished.setResult(exception);
            } else {
                loadFinished.setResult(true);
            }
        } else if (state.is(State.LOADED)) {
            state.next(State.LOADING);
        }
    }

    /**
     * Triggers key loading on SENDER if it hadn't started. Delays triggering if invoked multiple times.
     **/
    public void triggerLoadingWithDelay() {
        if (delayedTrigger == null) {
            Runnable runnable = new Runnable() {
                @Override
                public void run() {
                    Operation op = new PartitionCheckIfLoadedOperation(mapName, true);
                    opService.invokeOnPartition(SERVICE_NAME, op, mapNamePartition);
                }
            };
            delayedTrigger = new CoalescingDelayedTrigger(execService, LOADING_TRIGGER_DELAY, LOADING_TRIGGER_DELAY, runnable);
        }

        delayedTrigger.executeWithDelay();
    }

    public boolean shouldDoInitialLoad() {

        if (role.is(Role.SENDER_BACKUP)) {
            // was backup. become primary sender
            role.next(Role.SENDER);

            if (state.is(State.LOADING)) {
                // previous loading was in progress. cancel and start from scratch
                state.next(State.NOT_LOADED);
                loadFinished.setResult(false);
            }
        }

        return state.is(State.NOT_LOADED);
    }

    private void sendKeysInBatches(MapStoreContext mapStoreContext, boolean replaceExistingValues) throws Exception {

        int clusterSize = partitionService.getMemberPartitionsMap().size();
        Iterator keys = null;
        Throwable loadError = null;

        try {
            Iterable allKeys = mapStoreContext.loadAllKeys();
            keys = allKeys.iterator();
            Iterator dataKeys = map(keys, toData);
            int mapMaxSize = clusterSize * maxSizePerNode;

            if (mapMaxSize > 0) {
                dataKeys = limit(dataKeys, mapMaxSize);
            }

            Iterator> partitionsAndKeys = map(dataKeys, toPartition(partitionService));
            Iterator>> batches = toBatches(partitionsAndKeys, maxBatch);

            List futures = new ArrayList();
            while (batches.hasNext()) {
                Map> batch = batches.next();
                futures.addAll(sendBatch(batch, replaceExistingValues));
            }

            // This acts as a barrier to prevent re-ordering of key distribution operations (LoadAllOperation)
            // and LoadStatusOperation(s) which indicates all keys were already loaded.
            // Re-ordering of in-flight operations can happen during a partition migration. We are waiting here
            // for all LoadAllOperation(s) to be ACKed by receivers and only then we send them the LoadStatusOperation
            // See https://github.com/hazelcast/hazelcast/issues/4024 for additional details
            FutureUtil.waitWithDeadline(futures, KEY_DISTRIBUTION_TIMEOUT_MINUTES, TimeUnit.MINUTES);

        } catch (Exception caught) {
            loadError = caught;
        } finally {
            sendLoadCompleted(clusterSize, loadError);

            if (keys instanceof Closeable) {
                closeResource((Closeable) keys);
            }
        }
    }

    private List sendBatch(Map> batch, boolean replaceExistingValues) {
        Set>> entries = batch.entrySet();
        List futures = new ArrayList(entries.size());
        for (Entry> e : entries) {
            int partitionId = e.getKey();
            List keys = e.getValue();

            MapOperation op = operationProvider.createLoadAllOperation(mapName, keys, replaceExistingValues);

            InternalCompletableFuture future = opService.invokeOnPartition(SERVICE_NAME, op, partitionId);
            futures.add(future);
        }
        return futures;
    }

    private void sendLoadCompleted(int clusterSize, Throwable exception) throws Exception {

        // notify all partitions about loading status: finished or exception encountered
        opService.invokeOnAllPartitions(SERVICE_NAME, new LoadStatusOperationFactory(mapName, exception));

        // notify SENDER_BACKUP
        if (hasBackup && clusterSize > 1) {
            Operation op = new LoadStatusOperation(mapName, exception);
            opService.createInvocationBuilder(SERVICE_NAME, op, mapNamePartition).setReplicaIndex(1).invoke();
        }
    }

    private void sendLoadCompleted(Throwable t) {
        Operation op = new LoadStatusOperation(mapName, t);
        // This updates the local record store on the partition thread.
        // If invoked by the SENDER_BACKUP however it's the replica index has to be set to 1, otherwise
        // it will be a remote call to the SENDER who is the owner of the given partitionId.
        if (hasBackup && role.is(Role.SENDER_BACKUP)) {
            opService.createInvocationBuilder(SERVICE_NAME, op, partitionId).setReplicaIndex(1).invoke();
        } else {
            opService.createInvocationBuilder(SERVICE_NAME, op, partitionId).invoke();
        }
    }

    public void setMaxBatch(int maxBatch) {
        this.maxBatch = maxBatch;
    }

    public void setMaxSize(int maxSize) {
        this.maxSizePerNode = maxSize;
    }

    public void setHasBackup(boolean hasBackup) {
        this.hasBackup = hasBackup;
    }

    public void setMapOperationProvider(MapOperationProvider operationProvider) {
        this.operationProvider = operationProvider;
    }

    private ExecutionCallback ifLoadedCallback() {
        return new ExecutionCallback() {
            @Override
            public void onResponse(Boolean response) {
                if (response) {
                    sendLoadCompleted(null);
                }
            }

            @Override
            public void onFailure(Throwable t) {
                sendLoadCompleted(t);
            }
        };
    }

    private static final class LoadFinishedFuture extends AbstractCompletableFuture
            implements ExecutionCallback {

        private LoadFinishedFuture(Boolean result) {
            this();
            setResult(result);
        }

        private LoadFinishedFuture() {
            super((Executor) null, getLogger(LoadFinishedFuture.class));
        }

        @Override
        public Boolean get(long timeout, TimeUnit timeUnit) throws InterruptedException, ExecutionException, TimeoutException {
            if (isDone()) {
                return getResult();
            }
            throw new UnsupportedOperationException("Future is not done yet");
        }

        @Override
        public void onResponse(Boolean loaded) {
            if (loaded) {
                setResult(loaded);
            }
            // if not loaded yet we wait for the last batch to arrive
        }

        @Override
        public void onFailure(Throwable t) {
            setResult(t);
        }

        @Override
        protected boolean shouldCancel(boolean mayInterruptIfRunning) {
            return false;
        }

        @Override
        protected void setResult(Object result) {
            super.setResult(result);
        }

        @Override
        public String toString() {
            return getClass().getSimpleName() + "{done=" + isDone() + "}";
        }
    }

    public void onKeyLoad(ExecutionCallback callback) {
        loadFinished.andThen(callback, execService.getExecutor(MAP_LOAD_ALL_KEYS_EXECUTOR));
    }
}