org.elasticsearch.action.admin.cluster.snapshots.restore.RestoreClusterStateListener Maven / Gradle / Ivy

/*
 * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
 * or more contributor license agreements. Licensed under the Elastic License
 * 2.0 and the Server Side Public License, v 1; you may not use this file except
 * in compliance with, at your election, the Elastic License 2.0 or the Server
 * Side Public License, v 1.
 */

package org.elasticsearch.action.admin.cluster.snapshots.restore;

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.elasticsearch.action.ActionListener;
import org.elasticsearch.cluster.ClusterChangedEvent;
import org.elasticsearch.cluster.ClusterStateListener;
import org.elasticsearch.cluster.RestoreInProgress;
import org.elasticsearch.cluster.service.ClusterService;
import org.elasticsearch.common.collect.ImmutableOpenMap;
import org.elasticsearch.common.util.concurrent.ThreadContext;
import org.elasticsearch.index.shard.ShardId;
import org.elasticsearch.snapshots.RestoreInfo;
import org.elasticsearch.snapshots.RestoreService;

import java.util.function.Supplier;

import static org.elasticsearch.snapshots.RestoreService.restoreInProgress;

public class RestoreClusterStateListener implements ClusterStateListener {

    private static final Logger logger = LogManager.getLogger(RestoreClusterStateListener.class);

    private final ClusterService clusterService;
    private final String uuid;
    private final ActionListener listener;
    private final Supplier contextSupplier;

    private RestoreClusterStateListener(
        ClusterService clusterService,
        RestoreService.RestoreCompletionResponse response,
        ActionListener listener,
        Supplier contextSupplier
    ) {
        this.clusterService = clusterService;
        this.uuid = response.getUuid();
        this.listener = listener;
        this.contextSupplier = contextSupplier;
    }

    @Override
    public void clusterChanged(ClusterChangedEvent changedEvent) {
        try (ThreadContext.StoredContext stored = contextSupplier.get()) {
            final RestoreInProgress.Entry prevEntry = restoreInProgress(changedEvent.previousState(), uuid);
            final RestoreInProgress.Entry newEntry = restoreInProgress(changedEvent.state(), uuid);
            if (prevEntry == null) {
                // When there is a master failure after a restore has been started, this listener might not be registered
                // on the current master and as such it might miss some intermediary cluster states due to batching.
                // Clean up listener in that case and acknowledge completion of restore operation to client.
                clusterService.removeListener(this);
                listener.onResponse(new RestoreSnapshotResponse((RestoreInfo) null));
            } else if (newEntry == null) {
                clusterService.removeListener(this);
                ImmutableOpenMap shards = prevEntry.shards();
                assert prevEntry.state().completed() : "expected completed snapshot state but was " + prevEntry.state();
                assert RestoreService.completed(shards) : "expected all restore entries to be completed";
                RestoreInfo ri = new RestoreInfo(
                    prevEntry.snapshot().getSnapshotId().getName(),
                    prevEntry.indices(),
                    shards.size(),
                    shards.size() - RestoreService.failedShards(shards)
                );
                RestoreSnapshotResponse response = new RestoreSnapshotResponse(ri);
                logger.debug("restore of [{}] completed", prevEntry.snapshot().getSnapshotId());
                listener.onResponse(response);
            } else {
                // restore not completed yet, wait for next cluster state update
            }
        }
    }

    /**
     * Creates a cluster state listener and registers it with the cluster service. The listener passed as a
     * parameter will be called when the restore is complete.
     */
    public static void createAndRegisterListener(
        ClusterService clusterService,
        RestoreService.RestoreCompletionResponse response,
        ActionListener listener,
        ThreadContext threadContext
    ) {
        clusterService.addListener(
            new RestoreClusterStateListener(clusterService, response, listener, threadContext.newRestorableContext(true))
        );
    }
}