org.wso2.andes.thrift.MBThriftClient.java Source code

Java tutorial

Introduction

Here is the source code for org.wso2.andes.thrift.MBThriftClient.java

Source

/*
 * Copyright (c) 2016, WSO2 Inc. (http://www.wso2.org) All Rights Reserved.
 *
 * WSO2 Inc. licenses this file to you 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 org.wso2.andes.thrift;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.commons.pool2.ObjectPool;
import org.apache.commons.pool2.PooledObjectFactory;
import org.apache.commons.pool2.impl.GenericObjectPool;
import org.apache.thrift.TException;
import org.apache.thrift.transport.TTransportException;
import org.wso2.andes.configuration.AndesConfigurationManager;
import org.wso2.andes.configuration.enums.AndesConfiguration;
import org.wso2.andes.kernel.slot.ConnectionException;
import org.wso2.andes.kernel.slot.CoordinatorConnectionListener;
import org.wso2.andes.kernel.slot.Slot;
import org.wso2.andes.thrift.slot.gen.SlotInfo;
import org.wso2.andes.thrift.slot.gen.SlotManagementService;

import java.util.Queue;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * A wrapper client for the native thrift client with a client connection pool and retry logic.
 */

public class MBThriftClient {

    /**
     * A state variable to indicate whether the reconnecting  to the thrift server is started or
     * not
     */
    private boolean reconnecting = false;

    private final Log log = LogFactory.getLog(MBThriftClient.class);

    private final Queue<CoordinatorConnectionListener> connectionListenerQueue = new ConcurrentLinkedQueue<>();

    private AtomicBoolean isConnected = new AtomicBoolean(false);

    /**
     * Thrift client factory to build thrift client pool.
     */
    private PooledObjectFactory<SlotManagementService.Client> thriftClientFactory = new ThriftClientFactory();

    /**
     * Thrift socket timeout value from broker.xml file
     */
    private int socketTimeout = 0;

    /**
     * The thrift client pool.
     */
    private ObjectPool<SlotManagementService.Client> thriftClientPool;

    /**
     * Number of times each thrift operation retries before checking for a coordinator change.
     */
    private static final int RETRY_COUNT = 1;

    /**
     * Keeps track of the timestamp of the last successful reconnect so that reconnect requests received before this
     * can be ignored.
     */
    private long lastReconnectionSuccessTimestamp = 0;

    public MBThriftClient() {
        socketTimeout = AndesConfigurationManager.readValue(AndesConfiguration.COORDINATION_THRIFT_SO_TIMEOUT);
        initializeThriftConnectionPool();
    }

    /**
     * getSlot method. Returns Slot Object, when the
     * queue name is given
     *
     * @param queueName name of the queue
     * @param nodeId    of this node
     * @return slot object
     * @throws ConnectionException Throws when thrift connection fails
     */
    public Slot getSlot(String queueName, String nodeId) throws ConnectionException {
        SlotInfo slotInfo;

        for (int i = 0; i <= RETRY_COUNT; i++) {
            SlotManagementService.Client client = null;

            try {
                client = getServiceClient();
                slotInfo = client.getSlotInfo(queueName, nodeId);
                return convertSlotInforToSlot(slotInfo);
            } catch (TException e) {
                invalidateServiceClient(client);
                log.error("Attempt " + i + " failed requesting a slot from coordinator", e);
            } finally {
                if (client != null) {
                    returnServiceClient(client);
                }
            }
        }

        handleCoordinatorChanges();
        throw new ConnectionException("Coordinator has changed");
    }

    /**
     * Add Thrift connection listener
     *
     * @param connectionListener {@link CoordinatorConnectionListener}
     *
     */
    public void addConnectionListener(CoordinatorConnectionListener connectionListener) {
        connectionListenerQueue.add(connectionListener);
    }

    /**
     * Convert SlotInfo object to Slot object
     *
     * @param slotInfo object generated by thrift
     * @return slot object
     */
    private Slot convertSlotInforToSlot(SlotInfo slotInfo) {
        Slot slot = new Slot();
        slot.setStartMessageId(slotInfo.getStartMessageId());
        slot.setEndMessageId(slotInfo.getEndMessageId());
        slot.setStorageQueueName(slotInfo.getQueueName());
        return slot;
    }

    /**
     * updateMessageId method. This method will pass the locally chosen slot range to the SlotManagerClusterMode. Slot manager
     * maintains a list of slot ranges in a map along with the queue. This messageId will
     * be stored in that map.
     *
     * @param queueName name of the queue
     * @param nodeId unique hazelcast identifier of node.
     * @param startMessageId start message Id of the locally chosen slot.
     * @param endMessageId end message Id of the locally chosen slot.
     * @param localSafeZone Minimum message ID of the node that is deemed safe.
     * @throws ConnectionException in case of an connection error
     */
    public synchronized void updateMessageId(String queueName, String nodeId, long startMessageId,
            long endMessageId, long localSafeZone) throws ConnectionException {

        boolean updateSuccess = false;

        for (int i = 0; i <= RETRY_COUNT; i++) {
            if (updateSuccess) {
                break;
            }
            SlotManagementService.Client client = null;

            try {
                client = getServiceClient();
                client.updateMessageId(queueName, nodeId, startMessageId, endMessageId, localSafeZone);
                updateSuccess = true;
            } catch (TException e) {
                invalidateServiceClient(client);
                log.error("Attempt " + i + " failed updating message Id", e);
            } finally {
                if (client != null) {
                    returnServiceClient(client);
                }
            }
        }

        if (!updateSuccess) {
            handleCoordinatorChanges();
            throw new ConnectionException("Coordinator has changed");
        }
    }

    /**
     * Delete the slot from SlotAssignmentMap when all the messages in the slot has been sent and
     * all the acks are received.
     *
     * @param queueName  name of the queue where slot belongs to
     * @param slot      to be deleted
     * @throws ConnectionException Throws when thrift connection fails
     */
    public boolean deleteSlot(String queueName, Slot slot, String nodeId) throws ConnectionException {

        SlotInfo slotInfo = new SlotInfo(slot.getStartMessageId(), slot.getEndMessageId(),
                slot.getStorageQueueName(), nodeId, slot.isAnOverlappingSlot());

        for (int i = 0; i <= RETRY_COUNT; i++) {
            SlotManagementService.Client client = null;

            try {
                client = getServiceClient();
                return client.deleteSlot(queueName, slotInfo, nodeId);
            } catch (TException e) {
                invalidateServiceClient(client);
                log.error("Attempt " + i + " failed deleting slot");
            } finally {
                if (client != null) {
                    returnServiceClient(client);
                }
            }
        }

        handleCoordinatorChanges();
        throw new ConnectionException("Coordinator has changed");
    }

    /**
     * Re-assign the slot when the last subscriber leaves the node
     *
     * @param nodeId    of this node
     * @param queueName name of the queue
     * @throws ConnectionException Throws when thrift connection fails
     */
    public synchronized void reAssignSlotWhenNoSubscribers(String nodeId, String queueName)
            throws ConnectionException {

        boolean reassignSuccess = false;

        for (int i = 0; i <= RETRY_COUNT; i++) {
            if (reassignSuccess) {
                break;
            }
            SlotManagementService.Client client = null;

            try {
                client = getServiceClient();
                client.reAssignSlotWhenNoSubscribers(nodeId, queueName);
                reassignSuccess = true;
            } catch (TException e) {
                invalidateServiceClient(client);
                log.error("Attempt " + i + " failed reassigning slot", e);
            } finally {
                if (client != null) {
                    returnServiceClient(client);
                }
            }
        }

        if (!reassignSuccess) {
            handleCoordinatorChanges();
            throw new ConnectionException("Coordinator has changed");
        }
    }

    /**
     * Delete all in-memory slot associations with a given queue. This is required to handle a queue purge event.
     *
     * @param queueName name of destination queue
     * @throws ConnectionException Throws when thrift connection fails
     */
    public synchronized void clearAllActiveSlotRelationsToQueue(String queueName) throws ConnectionException {

        boolean success = false;

        for (int i = 0; i <= RETRY_COUNT; i++) {
            if (success) {
                break;
            }
            SlotManagementService.Client client = null;

            try {
                client = getServiceClient();
                client.clearAllActiveSlotRelationsToQueue(queueName);
                success = true;
            } catch (TException e) {
                invalidateServiceClient(client);
                log.error("Attempt " + i + " failed clearing active slot relations to queue", e);
            } finally {
                if (client != null) {
                    returnServiceClient(client);
                }
            }
        }

        if (!success) {
            handleCoordinatorChanges();
            throw new ConnectionException("Coordinator has changed");
        }
    }

    /**
     * Update the safeZone message ID of this node
     * @param safeZoneMessageID safe zone message ID
     * @param nodeID node ID of this node
     * @return global safeZone
     * @throws ConnectionException when MB thrift server is down
     */
    public synchronized long updateSlotDeletionSafeZone(long safeZoneMessageID, String nodeID)
            throws ConnectionException {

        for (int i = 0; i <= RETRY_COUNT; i++) {
            SlotManagementService.Client client = null;

            try {
                client = getServiceClient();
                return client.updateCurrentMessageIdForSafeZone(safeZoneMessageID, nodeID);
            } catch (TException e) {
                invalidateServiceClient(client);
                log.error("Attempt " + i + " failed updating slot delete safe zone", e);
            } finally {
                if (client != null) {
                    returnServiceClient(client);
                }
            }
        }

        handleCoordinatorChanges();
        throw new ConnectionException("Coordinator has changed");
    }

    /**
     * Start the thrift server reconnecting thread when the coordinator of the cluster is changed.
     */
    private void handleCoordinatorChanges() {

        long attemptStartTimeMillis = System.currentTimeMillis();

        synchronized (this) {

            if (!isReconnecting() && lastReconnectionSuccessTimestamp < attemptStartTimeMillis) {
                notifyDisconnection();
                setReconnectingFlag();
                startReconnecting();
            }
        }
    }

    /**
     * Notify Thrift connection disconnect to the {@link CoordinatorConnectionListener}s
     */
    private void notifyDisconnection() {
        if (isConnected.compareAndSet(true, false)) {
            for (CoordinatorConnectionListener listener : connectionListenerQueue) {
                listener.onCoordinatorDisconnect();
            }
        }
    }

    /**
     * This will try to connect to thrift server until successful.
     * reconnecting flag will be set to false when reconnecting is successful.
     *
     * @throws TTransportException when connecting to thrift server is unsuccessful
     */
    private void reConnectToServer() throws TTransportException {
        Long reconnectTimeout = (Long) AndesConfigurationManager
                .readValue(AndesConfiguration.COORDINATOR_THRIFT_RECONNECT_TIMEOUT) * 1000;
        try {
            //Reconnect timeout set because Hazelcast coordinator may still not elected in failover scenario
            Thread.sleep(reconnectTimeout);

            // re-initialize thrift connection pool discarding the previous pool
            thriftClientPool.close();
            initializeThriftConnectionPool();

            // If following call succeeds that means a new thrift connection can be made to the coordinator.
            SlotManagementService.Client client = thriftClientPool.borrowObject();

            reconnecting = false;
            lastReconnectionSuccessTimestamp = System.currentTimeMillis();

            // Return the borrowed client back to the pool
            thriftClientPool.returnObject(client);
            notifyConnection();
        } catch (TTransportException e) {
            log.error("Could not connect to the Thrift Server", e);
            throw new TTransportException("Could not connect to the Thrift Server", e);
        } catch (InterruptedException ignore) {
        } catch (Exception e) {
            log.error("Could not connect to the Thrift Server", e);
        }
    }

    /**
     * Notify Thrift connection establish signal to the {@link CoordinatorConnectionListener}
     */
    private void notifyConnection() {
        if (isConnected.compareAndSet(false, true)) {
            for (CoordinatorConnectionListener listener : connectionListenerQueue) {
                listener.onCoordinatorReconnect();
            }
        }
    }

    /**
     * This thread is responsible of reconnecting to the thrift server of the coordinator until it
     * gets succeeded
     */
    private void startReconnecting() {
        while (reconnecting) {

            try {
                reConnectToServer();
                // If re connect to server is successful, following code segment will be executed
                reconnecting = false;
            } catch (Throwable e) {
                log.error("Error occurred while reconnecting to slot coordinator", e);

                try {
                    TimeUnit.SECONDS.sleep(2);
                } catch (InterruptedException interruptedException) {
                    Thread.currentThread().interrupt();
                }
            }
        }
    }

    /**
     * A flag to specify whether the reconnecting to thrift server is happening or not
     *
     * @return whether the reconnecting to thrift server is happening or not
     */
    public boolean isReconnecting() {
        return reconnecting;
    }

    /**
     * Set reconnecting flag to true.
     *
     */
    public void setReconnectingFlag() {
        reconnecting = true;
    }

    /**
     * Initialize a new thrift client connection pool.
     */
    private void initializeThriftConnectionPool() {

        int thriftClientPoolSize = AndesConfigurationManager
                .readValue(AndesConfiguration.PERFORMANCE_TUNING_THRIFT_CLIENT_POOL_SIZE);

        GenericObjectPool<SlotManagementService.Client> thriftConnectionPool = new GenericObjectPool<>(
                thriftClientFactory);
        thriftConnectionPool.setMaxTotal(thriftClientPoolSize);
        thriftConnectionPool.setTestOnBorrow(true);
        if (socketTimeout > 0) {
            thriftConnectionPool.setTimeBetweenEvictionRunsMillis(socketTimeout / 2);
            thriftConnectionPool.setMinEvictableIdleTimeMillis(socketTimeout);
        }
        thriftClientPool = thriftConnectionPool;
    }

    /**
     * Returns an instance of Slot Management service client from the pool which is used to communicate to the
     * thrift server. If it does not succeed in connecting to the server, it throws a  ConnectionException
     *
     * @return a SlotManagementService client
     */
    private SlotManagementService.Client getServiceClient() throws ConnectionException {
        try {
            return thriftClientPool.borrowObject();
        } catch (Exception e) {
            throw new ConnectionException("Error borrowing a connection from thrift connection pool", e);
        }
    }

    /**
     * Return thrift client connection back to the connection pool.
     *
     * @param client The client to return to the pool
     */
    private void returnServiceClient(SlotManagementService.Client client) {
        try {
            thriftClientPool.returnObject(client);
        } catch (Exception e) {
            log.warn("Error returning thrift client back to the pool. Coordinator might have changed.", e);
        }
    }

    /**
     * Invalidate a thrift connection and ask to remove it from the pool.
     *
     * @param client Invalid thrift client connection
     * @throws ConnectionException Throws when thrift client pool fails to invalidate the given object
     */
    private void invalidateServiceClient(SlotManagementService.Client client) throws ConnectionException {
        if (client != null) {
            try {
                thriftClientPool.invalidateObject(client);
            } catch (Exception e) {
                throw new ConnectionException("Error invalidating a connection from thrift connection pool", e);
            }
        }
    }
}