org.springframework.integration.etcd.leader.LeaderInitiator.java Source code

Java tutorial

Introduction

Here is the source code for org.springframework.integration.etcd.leader.LeaderInitiator.java

Source

/*
 * Copyright 2018 the original author or authors.
 *
 * 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
 *
 * https://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 org.springframework.integration.etcd.leader;

import static com.ibm.etcd.client.KeyUtils.bs;

import java.util.concurrent.Callable;
import java.util.concurrent.CancellationException;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

import org.springframework.beans.factory.DisposableBean;
import org.springframework.context.ApplicationEventPublisher;
import org.springframework.context.ApplicationEventPublisherAware;
import org.springframework.context.Lifecycle;
import org.springframework.integration.leader.Candidate;
import org.springframework.integration.leader.Context;
import org.springframework.integration.leader.DefaultCandidate;
import org.springframework.integration.leader.event.DefaultLeaderEventPublisher;
import org.springframework.integration.leader.event.LeaderEventPublisher;
import org.springframework.util.Assert;

import com.ibm.etcd.api.DeleteRangeRequest;
import com.ibm.etcd.api.PutRequest;
import com.ibm.etcd.api.TxnResponse;
import com.ibm.etcd.client.EtcdClient;
import io.grpc.StatusRuntimeException;

/**
 * Bootstrap leadership {@link Candidate candidates} with etcd. Upon construction, {@link #start}
 * must be invoked to register the candidate for leadership election.
 *
 * @author Venil Noronha
 * @author Patrick Peralta
 * @author Lewis Headden
 */
public class LeaderInitiator implements Lifecycle, DisposableBean, ApplicationEventPublisherAware {

    private final Log logger = LogFactory.getLog(getClass());

    /**
     * TTL for etcd entry in seconds.
     */
    private static final int TTL = 10;

    /**
     * Number of seconds to sleep between issuing heartbeats.
     */
    private static final int HEART_BEAT_SLEEP = TTL / 2;

    /**
     * Default namespace for etcd entry.
     */
    private static final String DEFAULT_NAMESPACE = "spring-integration";

    /**
     * {@link EtcdClient} instance.
     */
    private final EtcdClient client;

    /**
     * Candidate for leader election.
     */
    private final Candidate candidate;

    /**
     * Executor service for running leadership daemon.
     */
    private final ExecutorService leaderExecutorService = Executors.newSingleThreadExecutor(r -> {
        Thread thread = new Thread(r, "Etcd-Leadership");
        thread.setDaemon(true);
        return thread;
    });

    /**
     * Executor service for running leadership worker daemon.
     */
    private final ExecutorService workerExecutorService = Executors.newSingleThreadExecutor(r -> {
        Thread thread = new Thread(r, "Etcd-Leadership-Worker");
        thread.setDaemon(true);
        return thread;
    });

    /**
     * Flag that indicates whether the current candidate is the leader.
     */
    private volatile boolean isLeader = false;

    /**
     * Flag that indicates whether the current candidate's leadership should be relinquished.
     */
    private volatile boolean relinquishLeadership = false;

    /**
     * Future returned by submitting a {@link Initiator} to {@link #leaderExecutorService}.
     * This is used to cancel leadership.
     */
    private volatile Future<Void> initiatorFuture;

    /**
     * Future returned by submitting a {@link Worker} to {@link #workerExecutorService}.
     * This is used to notify leadership revocation.
     */
    private volatile Future<Void> workerFuture;

    /**
     * Flag that indicates whether the leadership election for this {@link #candidate} is running.
     */
    private volatile boolean running;

    /**
     * Leader event publisher.
     */
    private volatile LeaderEventPublisher leaderEventPublisher = new DefaultLeaderEventPublisher();

    private boolean customPublisher = false;

    /**
     * The {@link EtcdContext} instance.
     */
    private final EtcdContext context;

    /**
     * The base etcd path where candidate id is to be stored.
     */
    private final String baseEtcdPath;

    /**
     * Construct a {@link LeaderInitiator}.
     * @param client {@link EtcdClient} instance
     * @param namespace Etcd namespace
     */
    public LeaderInitiator(EtcdClient client, String namespace) {
        this(client, new DefaultCandidate(), namespace);
    }

    /**
     * Construct a {@link LeaderInitiator}.
     * @param client {@link EtcdClient} instance
     * @param candidate leadership election candidate
     * @param namespace Etcd namespace
     */
    public LeaderInitiator(EtcdClient client, Candidate candidate, String namespace) {
        Assert.notNull(client, "'etcdClient' must not be null");
        Assert.notNull(candidate, "'candidate' must not be null");
        this.client = client;
        this.candidate = candidate;
        this.context = new EtcdContext();
        this.baseEtcdPath = (namespace == null ? DEFAULT_NAMESPACE : namespace) + "/" + candidate.getRole();
    }

    @Override
    public void setApplicationEventPublisher(ApplicationEventPublisher applicationEventPublisher) {
        if (!this.customPublisher) {
            this.leaderEventPublisher = new DefaultLeaderEventPublisher(applicationEventPublisher);
        }
    }

    /**
     * Start the registration of the {@link #candidate} for leader election.
     */
    @Override
    public synchronized void start() {
        if (!this.running) {
            this.initiatorFuture = this.leaderExecutorService.submit(new Initiator());
            this.running = true;
        }
    }

    /**
     * Stop the registration of the {@link #candidate} for leader election.
     * If the candidate is currently leader, its leadership will be revoked.
     */
    @Override
    public synchronized void stop() {
        if (this.running) {
            this.running = false;
            this.initiatorFuture.cancel(true);
        }
    }

    /**
     * @return true if leadership election for this {@link #candidate} is running
     */
    @Override
    public boolean isRunning() {
        return this.running;
    }

    @Override
    public void destroy() throws Exception {
        stop();
        this.workerExecutorService.shutdown();
        this.leaderExecutorService.shutdown();
        this.workerExecutorService.awaitTermination(30, TimeUnit.SECONDS);
        this.leaderExecutorService.awaitTermination(30, TimeUnit.SECONDS);
    }

    /**
     * Sets the {@link LeaderEventPublisher}.
     * @param leaderEventPublisher the event publisher
     */
    public void setLeaderEventPublisher(LeaderEventPublisher leaderEventPublisher) {
        Assert.notNull(leaderEventPublisher, "leaderEventPublisher cannot be null");
        this.leaderEventPublisher = leaderEventPublisher;
    }

    /**
     * Notifies that the candidate has acquired leadership.
     */
    private void notifyGranted() {
        this.isLeader = true;
        this.leaderEventPublisher.publishOnGranted(LeaderInitiator.this, this.context, this.candidate.getRole());
        this.workerFuture = this.workerExecutorService.submit(new Worker());
    }

    /**
     * Notifies that the candidate's leadership was revoked.
     * @throws InterruptedException if the current thread was interrupted while waiting for the worker
     *     thread to finish.
     */
    private void notifyRevoked() throws InterruptedException {
        this.isLeader = false;
        this.leaderEventPublisher.publishOnRevoked(LeaderInitiator.this, this.context, this.candidate.getRole());
        this.workerFuture.cancel(true);
        try {
            this.workerFuture.get();
        } catch (InterruptedException e) {
            throw e;
        } catch (CancellationException e) {
            // Consume
        } catch (ExecutionException e) {
            logger.error("Exception thrown by candidate", e.getCause());
        }
    }

    /**
     * Tries to delete the candidate's entry from etcd.
     */
    private void tryDeleteCandidateEntry() {
        try {
            final TxnResponse response = this.client.getKvClient().txnIf().cmpEqual(bs(this.baseEtcdPath))
                    .value(bs(this.candidate.getId())).then()
                    .delete(DeleteRangeRequest.newBuilder().setKey(bs(this.baseEtcdPath)).build()).sync();

            if (!response.getSucceeded()) {
                logger.warn("Couldn't delete candidate's entry from etcd because candidate was not leader");
            }
        } catch (StatusRuntimeException e) {
            logger.warn("Failed deleting candidate entry from etcd", e);
        }
    }

    /**
     * Callable that invokes {@link Candidate#onGranted(Context)} when the candidate is granted
     * leadership.
     */
    class Worker implements Callable<Void> {

        @Override
        public Void call() throws InterruptedException {
            try {
                LeaderInitiator.this.candidate.onGranted(LeaderInitiator.this.context);
                Thread.sleep(Long.MAX_VALUE);
            } finally {
                LeaderInitiator.this.relinquishLeadership = true;
                LeaderInitiator.this.candidate.onRevoked(LeaderInitiator.this.context);
            }
            return null;
        }

    }

    /**
     * Callable that manages the etcd heart beats for leadership election.
     */
    class Initiator implements Callable<Void> {

        @Override
        public Void call() throws InterruptedException {
            try {
                while (LeaderInitiator.this.running) {
                    if (LeaderInitiator.this.relinquishLeadership) {
                        relinquishLeadership();
                        LeaderInitiator.this.relinquishLeadership = false;
                    } else if (LeaderInitiator.this.isLeader) {
                        sendHeartBeat();
                    } else {
                        tryAcquire();
                    }
                    TimeUnit.SECONDS.sleep(HEART_BEAT_SLEEP);
                }
            } finally {
                if (LeaderInitiator.this.isLeader) {
                    relinquishLeadership();
                }
            }
            return null;
        }

        /**
         * Relinquishes leadership of current candidate by deleting candidate's entry from etcd and then
         * notifies that the current candidate is no longer leader.
         * @throws InterruptedException if the current thread was interrupted while notifying
         *     revocation.
         */
        private void relinquishLeadership() throws InterruptedException {
            tryDeleteCandidateEntry();
            notifyRevoked();
        }

        /**
         * Sends a heart beat to maintain leadership by refreshing the ttl of the etcd key. If the key
         * has a different value during the call, it is assumed that the current candidate's leadership
         * is revoked. If access to etcd fails, then the the current candidate's leadership is
         * relinquished.
         * @throws InterruptedException if the current thread was interrupted while notifying
         *     revocation.
         */
        private void sendHeartBeat() throws InterruptedException {
            try {
                final TxnResponse response = LeaderInitiator.this.client.getKvClient().txnIf()
                        .cmpEqual(bs(LeaderInitiator.this.baseEtcdPath))
                        .value(bs(LeaderInitiator.this.candidate.getId())).then()
                        .put(PutRequest.newBuilder().setKey(bs(LeaderInitiator.this.baseEtcdPath))
                                .setValue(bs(LeaderInitiator.this.candidate.getId())).build())
                        .sync();

                if (!response.getSucceeded()) {
                    notifyRevoked();
                }
            } catch (StatusRuntimeException e) {
                logger.warn("Failed to send leadership heartbeat", e);
            }
        }

        /**
         * Tries to acquire leadership by posting the candidate's id to etcd. If the etcd call is
         * successful, it is assumed that the current candidate is now leader.
         */
        private void tryAcquire() {
            try {
                final long leaseId = LeaderInitiator.this.client.getLeaseClient().create(TTL).get().getID();

                final TxnResponse response = LeaderInitiator.this.client.getKvClient().txnIf()
                        .notExists(bs(LeaderInitiator.this.baseEtcdPath)).then()
                        .put(PutRequest.newBuilder().setKey(bs(LeaderInitiator.this.baseEtcdPath))
                                .setValue(bs(LeaderInitiator.this.candidate.getId())).setLease(leaseId).build())
                        .sync();
                if (response.getSucceeded()) {
                    notifyGranted();
                } else {
                    logger.info("Tried to acquire leadership but another candidate is leader");
                }
            } catch (InterruptedException | ExecutionException e) {
                logger.warn("Failed trying to acquire leadership", e.getCause());
            }
        }

    }

    /**
     * Implementation of leadership context backed by Etcd.
     */
    class EtcdContext implements Context {

        @Override
        public boolean isLeader() {
            return LeaderInitiator.this.isLeader;
        }

        @Override
        public void yield() {
            if (LeaderInitiator.this.isLeader) {
                LeaderInitiator.this.relinquishLeadership = true;
            }
        }

        @Override
        public String toString() {
            return String.format("EtcdContext{role=%s, id=%s, isLeader=%s}",
                    LeaderInitiator.this.candidate.getRole(), LeaderInitiator.this.candidate.getId(), isLeader());
        }

    }

}