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

com.netflix.astyanax.recipes.queue.MessageConsumerImpl Maven / Gradle / Ivy

The newest version!
package com.netflix.astyanax.recipes.queue;

import com.google.common.collect.Lists;
import com.netflix.astyanax.ColumnListMutation;
import com.netflix.astyanax.MutationBatch;
import com.netflix.astyanax.connectionpool.exceptions.ConnectionException;
import com.netflix.astyanax.connectionpool.exceptions.NotFoundException;
import com.netflix.astyanax.model.Column;
import com.netflix.astyanax.model.ColumnList;
import com.netflix.astyanax.model.Equality;
import com.netflix.astyanax.model.RangeEndpoint;
import com.netflix.astyanax.recipes.locks.BusyLockException;
import com.netflix.astyanax.recipes.queue.triggers.Trigger;
import com.netflix.astyanax.util.RangeBuilder;
import com.netflix.astyanax.util.TimeUUIDUtils;
import java.util.Collection;
import java.util.List;
import java.util.concurrent.TimeUnit;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Message consumer implementation based on the sharded queue.
 *
 * @author pbhattacharyya
 */
class MessageConsumerImpl implements MessageConsumer {
    private static final Logger LOG = LoggerFactory.getLogger(MessageConsumerImpl.class);

    private final ShardedDistributedMessageQueue queue;

    public MessageConsumerImpl(ShardedDistributedMessageQueue q) {
        this.queue = q;
    }

    @Override
    public List readMessages(int itemsToPop) throws MessageQueueException, BusyLockException, InterruptedException {
        return readMessages(itemsToPop, 0, null);
    }

    @Override
    public List readMessages(int itemsToPop, long timeout, TimeUnit units) throws MessageQueueException, BusyLockException, InterruptedException {
        long timeoutTime = (timeout == 0) ? 0 : System.currentTimeMillis() + TimeUnit.MILLISECONDS.convert(timeout, units);
        // Loop while trying to get messages.
        // TODO: Make it possible to cancel this loop
        // TODO: Read full itemsToPop instead of just stopping when we get the first successful set
        List messages = null;
        while (true) {
            boolean success = false;
            MessageQueueShard partition = queue.shardReaderPolicy.nextShard();
            if (partition != null) {
                try {
                    messages = readAndReturnShard(partition, itemsToPop);
                    success = true;
                    if (messages != null && !messages.isEmpty()) {
                        return messages;
                    }
                } finally {
                    // releaseShard needs to know how many messages were proceed OR if there was an error fetching messages (-1)
                    queue.shardReaderPolicy.releaseShard(partition, success ? (messages == null ? 0 : messages.size()) : -1);
                }
            }
            if (timeoutTime != 0 && System.currentTimeMillis() > timeoutTime) {
                return Lists.newLinkedList();
            }
            Thread.sleep(queue.shardReaderPolicy.getPollInterval());
        }
    }

    @Override
    public List peekMessages(int itemsToPeek) throws MessageQueueException {
        return queue.peekMessages(itemsToPeek);
    }

    private List readAndReturnShard(MessageQueueShard shard, int itemsToPop) throws MessageQueueException, BusyLockException, InterruptedException {
        List messages = null;
        try {
            messages = readMessagesFromShard(shard.getName(), itemsToPop);
        } finally {
            if (messages == null || messages.isEmpty()) {
                queue.stats.incEmptyPartitionCount();
            }
        }
        return messages;
    }

    @Override
    public List readMessagesFromShard(String shardName, int itemsToPop) throws MessageQueueException, BusyLockException {
        if(queue.lockManager != null) {
            return readMessagesFromShardUsingLockManager(shardName, itemsToPop);
        }
        return readMessagesFromShardUsingDefaultLock(shardName, itemsToPop);
    }

    List readMessagesFromShardUsingLockManager(String shardName, int itemToPop) throws MessageQueueException, BusyLockException {
        ShardLock lock = null;
        try {
            lock = queue.lockManager.acquireLock(shardName);
            MutationBatch m = queue.keyspace.prepareMutationBatch().setConsistencyLevel(queue.consistencyLevel);
            ColumnListMutation rowMutation = m.withRow(queue.queueColumnFamily, shardName);
            long curTimeMicros = TimeUUIDUtils.getMicrosTimeFromUUID(TimeUUIDUtils.getUniqueTimeUUIDinMicros());
            return readMessagesInternal(shardName, itemToPop, 0, null, rowMutation, m, curTimeMicros);
        } catch (BusyLockException e) {
            queue.stats.incLockContentionCount();
            throw e;
        } catch (Exception e) {
            LOG.error("Error reading shard " + shardName, e);
            throw new MessageQueueException("Error", e);
        } finally {
            queue.lockManager.releaseLock(lock);
        }
    }

    List readMessagesFromShardUsingDefaultLock(String shardName, int itemsToPop) throws MessageQueueException, BusyLockException {
        MutationBatch m = null;
        MessageQueueEntry lockColumn = null;
        ColumnListMutation rowMutation = null;
        int lockColumnCount = 0;
        // Try locking first
        try {
            // 1. Write the lock column
            lockColumn = MessageQueueEntry.newLockEntry(MessageQueueEntryState.None);
            long curTimeMicros = TimeUUIDUtils.getTimeFromUUID(lockColumn.getTimestamp());
            m = queue.keyspace.prepareMutationBatch().setConsistencyLevel(queue.consistencyLevel);
            m.withRow(queue.queueColumnFamily, shardName).putColumn(lockColumn, curTimeMicros + queue.lockTimeout, queue.lockTtl);
            m.execute();
            // 2. Read back lock columns and entries
            ColumnList result = queue.keyspace.prepareQuery(queue.queueColumnFamily).setConsistencyLevel(queue.consistencyLevel).getKey(shardName)
                    .withColumnRange(ShardedDistributedMessageQueue.entrySerializer
                                                                                   .buildRange()
                                                                                   .greaterThanEquals((byte) MessageQueueEntryType.Lock.ordinal())
                                                                                   .lessThanEquals((byte) MessageQueueEntryType.Lock.ordinal())
                                                                                   .build()
                                                                   )
                    .execute()
                    .getResult();
            m = queue.keyspace.prepareMutationBatch().setConsistencyLevel(queue.consistencyLevel);
            rowMutation = m.withRow(queue.queueColumnFamily, shardName);
            rowMutation.deleteColumn(lockColumn);
            int lockCount = 0;
            boolean lockAcquired = false;
            lockColumnCount = result.size();
            for (Column column : result) {
                MessageQueueEntry lock = column.getName();
                if (lock.getType() == MessageQueueEntryType.Lock) {
                    lockColumnCount++;
                    // Stale lock so we can discard it
                    if (column.getLongValue() < curTimeMicros) {
                        queue.stats.incExpiredLockCount();
                        rowMutation.deleteColumn(lock);
                    } else if (lock.getState() == MessageQueueEntryState.Acquired) {
                        throw new BusyLockException("Not first lock");
                    } else {
                        lockCount++;
                        if (lockCount == 1 && lock.getTimestamp().equals(lockColumn.getTimestamp())) {
                            lockAcquired = true;
                        }
                    }
                    if (!lockAcquired) {
                        throw new BusyLockException("Not first lock");
                    }
                    // Write the acquired lock column
                    lockColumn = MessageQueueEntry.newLockEntry(lockColumn.getTimestamp(), MessageQueueEntryState.Acquired);
                    rowMutation.putColumn(lockColumn, curTimeMicros + queue.lockTimeout, queue.lockTtl);
                }
            }
        } catch (BusyLockException e) {
            queue.stats.incLockContentionCount();
            throw e;
        } catch (ConnectionException e) {
            LOG.error("Error reading shard " + shardName, e);
            throw new MessageQueueException("Error", e);
        } finally {
            try {
                m.execute();
            } catch (Exception e) {
                throw new MessageQueueException("Error committing lock", e);
            }
        }
        long curTimeMicros = TimeUUIDUtils.getMicrosTimeFromUUID(lockColumn.getTimestamp());
        m = queue.keyspace.prepareMutationBatch().setConsistencyLevel(queue.consistencyLevel);
        // First, release the lock column
        rowMutation = m.withRow(queue.queueColumnFamily, shardName);
        rowMutation.deleteColumn(lockColumn);
        return readMessagesInternal(shardName, itemsToPop, lockColumnCount, lockColumn, rowMutation, m, curTimeMicros);
    }

    @Override
    public void ackMessage(MessageContext context) throws MessageQueueException {
        MutationBatch mb = queue.keyspace.prepareMutationBatch().setConsistencyLevel(queue.consistencyLevel);
        fillAckMutation(context, mb);
        try {
            mb.execute();
        } catch (ConnectionException e) {
            throw new MessageQueueException("Failed to ack message", e);
        }
    }

    @Override
    public void ackMessages(Collection messages) throws MessageQueueException {
        MutationBatch mb = queue.keyspace.prepareMutationBatch().setConsistencyLevel(queue.consistencyLevel);
        for (MessageContext context : messages) {
            fillAckMutation(context, mb);
        }
        try {
            mb.execute();
        } catch (ConnectionException e) {
            throw new MessageQueueException("Failed to ack messages", e);
        }
    }

    private void fillAckMutation(MessageContext context, MutationBatch mb) {
        queue.stats.incAckMessageCount();
        Message message = context.getMessage();
        // Token refers to the timeout event.  If 0 (i.e. no) timeout was specified
        // then the token will not exist
        if (message.getToken() != null) {
            MessageQueueEntry entry = MessageQueueEntry.newBusyEntry(message);
            // Remove timeout entry from the queue
            mb.withRow(queue.queueColumnFamily, queue.getShardKey(message)).deleteColumn(entry);
            // Remove entry lookup from the key, if one exists
            if (message.hasKey()) {
                mb.withRow(queue.keyIndexColumnFamily, queue.getCompositeKey(queue.getName(), message.getKey()))
                        .putEmptyColumn(MessageMetadataEntry.newMessageId(queue.getCompositeKey(queue.getShardKey(message), entry.getMessageId())), queue.metadataDeleteTTL);
                if (message.isKeepHistory()) {
                    MessageHistory history = context.getHistory();
                    if (history.getStatus() == MessageStatus.RUNNING) {
                        history.setStatus(MessageStatus.DONE);
                    }
                    history.setEndTime(TimeUnit.MICROSECONDS.convert(System.currentTimeMillis(), TimeUnit.MILLISECONDS));
                    try {
                        mb.withRow(queue.historyColumnFamily, message.getKey())
                                .putColumn(history.getToken(), queue.serializeToString(context.getHistory()), queue.metadata.getHistoryTtl()); // TTL
                    } catch (Exception e) {
                        LOG.warn("Error serializing message history for " + message.getKey(), e);
                    }
                }
            }
            // Run hooks
            for (MessageQueueHooks hook : queue.hooks) {
                hook.beforeAckMessage(message, mb);
            }
        }
        if (context.getNextMessage() != null) {
            try {
                queue.fillMessageMutation(mb, context.getNextMessage());
            } catch (MessageQueueException e) {
                LOG.warn("Error filling nextMessage for " + message.getKey(), e);
            }
        }
    }

    @Override
    public void ackPoisonMessage(MessageContext context) throws MessageQueueException {
        // TODO: Remove bad message and add to poison queue
        MutationBatch mb = queue.keyspace.prepareMutationBatch().setConsistencyLevel(queue.consistencyLevel);
        fillAckMutation(context, mb);
        try {
            mb.execute();
        } catch (ConnectionException e) {
            queue.stats.incPersistError();
            throw new MessageQueueException("Failed to ack messages", e);
        }
    }

    private List readMessagesInternal(String shardName,
                                                      int itemsToPop,
                                                      int lockColumnCount,
                                                      MessageQueueEntry lockColumn,
                                                      ColumnListMutation rowMutation,
                                                      MutationBatch m,
                                                      long curTimeMicros) throws BusyLockException, MessageQueueException {

        try {
            List entries = Lists.newArrayList();
            RangeEndpoint re = ShardedDistributedMessageQueue.entrySerializer
                                              .makeEndpoint((byte) MessageQueueEntryType.Message.ordinal(), Equality.EQUAL)
                                              .append((byte) 0, Equality.EQUAL);
            if(lockColumn!=null) {
                re.append(lockColumn.getTimestamp(), Equality.LESS_THAN_EQUALS);
            } else {
                re.append(TimeUUIDUtils.getMicrosTimeUUID(curTimeMicros), Equality.LESS_THAN_EQUALS);
            }

            ColumnList result = queue.keyspace.prepareQuery(queue.queueColumnFamily)
                    .setConsistencyLevel(queue.consistencyLevel).getKey(shardName).
                    withColumnRange(new RangeBuilder()
                                        .setLimit(itemsToPop + (lockColumn == null? 0:(lockColumnCount + 1)))
                                        .setEnd(re.toBytes())
                                        .build()).execute().getResult();
            for (Column column : result) {
                if (itemsToPop == 0) {
                    break;
                }
                MessageQueueEntry entry = column.getName();
                switch (entry.getType()) {
                    case Lock:
                        // TODO: Track number of locks read and make sure we don't exceed itemsToPop
                        // We have the lock
                        if (lockColumn != null && entry.getState() == MessageQueueEntryState.Acquired) {
                            if (!entry.getTimestamp().equals(lockColumn.getTimestamp())) {
                                throw new BusyLockException("Someone else snuck in");
                            }
                        }
                        break;
                    case Message:
                        {
                            try {
                                itemsToPop--;
                                // First, we always want to remove the old item
                                String messageId = queue.getCompositeKey(shardName, entry.getMessageId());
                                rowMutation.deleteColumn(entry);
                                // Next, parse the message metadata and add a timeout entry
                                final Message message = queue.extractMessageFromColumn(column);
                                // Update the message state
                                if (message != null) {
                                    MessageContext context = new MessageContext();
                                    context.setMessage(message);
                                    // Message has a trigger so we need to figure out if it is an
                                    // unfinished repeating trigger and re-add it.
                                    if (message.hasTrigger()) {
                                        // Read back all messageIds associated with this key and check to see if we have duplicates.
                                        String groupRowKey = queue.getCompositeKey(queue.getName(), message.getKey());
                                        try {
                                            // Use consistency level
                                            ColumnList columns = queue.keyspace.prepareQuery(queue.keyIndexColumnFamily).getRow(groupRowKey).withColumnRange(ShardedDistributedMessageQueue.metadataSerializer.buildRange().greaterThanEquals((byte) MessageMetadataEntryType.MessageId.ordinal()).lessThanEquals((byte) MessageMetadataEntryType.MessageId.ordinal()).build()).execute().getResult();
                                            MessageMetadataEntry mostRecentMessageMetadata = null;
                                            long mostRecentTriggerTime = 0;
                                            for (Column currMessageEntry : columns) {
                                                MessageQueueEntry pendingMessageEntry = MessageQueueEntry.fromMetadata(currMessageEntry.getName());
                                                if (currMessageEntry.getTtl() == 0) {
                                                    long currMessageTriggerTime = pendingMessageEntry.getTimestamp(TimeUnit.MICROSECONDS);
                                                    // First message we found, so treat as the most recent
                                                    if (mostRecentMessageMetadata == null) {
                                                        mostRecentMessageMetadata = currMessageEntry.getName();
                                                        mostRecentTriggerTime = currMessageTriggerTime;
                                                    } else {
                                                        // This message's trigger time is after what we thought was the most recent.
                                                        // Discard the previous 'most' recent and accept this one instead
                                                        if (currMessageTriggerTime > mostRecentTriggerTime) {
                                                            LOG.warn("Need to discard : " + entry.getMessageId() + " => " + mostRecentMessageMetadata.getName());
                                                            m.withRow(queue.keyIndexColumnFamily,
                                                                    queue.getCompositeKey(queue.getName(), message.getKey())).putEmptyColumn(mostRecentMessageMetadata, queue.metadataDeleteTTL);
                                                            mostRecentTriggerTime = currMessageTriggerTime;
                                                            mostRecentMessageMetadata = currMessageEntry.getName();
                                                        } else {
                                                            LOG.warn("Need to discard : " + entry.getMessageId() + " => " + currMessageEntry.getName());
                                                            m.withRow(queue.keyIndexColumnFamily,
                                                                    queue.getCompositeKey(queue.getName(), message.getKey())).putEmptyColumn(currMessageEntry.getName(), queue.metadataDeleteTTL);
                                                        }
                                                    }
                                                }
                                            }
                                            if (mostRecentMessageMetadata != null) {
                                                if (!mostRecentMessageMetadata.getName().endsWith(entry.getMessageId())) {
                                                    throw new DuplicateMessageException("Duplicate trigger for " + messageId);
                                                }
                                            }
                                        } catch (NotFoundException e) {
                                        } catch (ConnectionException e) {
                                            throw new MessageQueueException("Error fetching row " + groupRowKey, e);
                                        }
                                        // Update the trigger
                                        final Message nextMessage;
                                        Trigger trigger = message.getTrigger().nextTrigger();
                                        if (trigger != null) {
                                            nextMessage = message.clone();
                                            nextMessage.setTrigger(trigger);
                                            context.setNextMessage(nextMessage);
                                            if (message.isAutoCommitTrigger()) {
                                                queue.fillMessageMutation(m, nextMessage);
                                            }
                                        }
                                    }
                                    // Message has a key so we remove this item from the messages by key index.
                                    // A timeout item will be added later
                                    if (message.hasKey()) {
                                        m.withRow(queue.keyIndexColumnFamily,
                                                queue.getCompositeKey(queue.getName(), message.getKey()))
                                                .putEmptyColumn(MessageMetadataEntry.newMessageId(messageId), queue.metadataDeleteTTL);
                                        LOG.debug("Removing from key  :  " + queue.getCompositeKey(queue.getName(), message.getKey()) + " : " + messageId);
                                        if (message.isKeepHistory()) {
                                            MessageHistory history = context.getHistory();
                                            history.setToken(entry.getTimestamp());
                                            history.setStartTime(curTimeMicros);
                                            history.setTriggerTime(message.getTrigger().getTriggerTime());
                                            history.setStatus(MessageStatus.RUNNING);
                                            try {
                                                m.withRow(queue.historyColumnFamily, message.getKey()).putColumn(entry.getTimestamp(), queue.serializeToString(history)
                                                        , queue.metadata.getHistoryTtl());
                                            } catch (Exception e) {
                                                LOG.warn("Error serializing history for key '" + message.getKey() + "'", e);
                                            }
                                        }
                                    }
                                    // Message has a timeout so we add a timeout event.
                                    if (message.getTimeout() > 0) {
                                        MessageQueueEntry timeoutEntry = MessageQueueEntry.newMessageEntry((byte) 0,
                                                TimeUUIDUtils.getMicrosTimeUUID(curTimeMicros
                                                + TimeUnit.MICROSECONDS.convert(message.getTimeout(), TimeUnit.SECONDS)
                                                + (queue.counter.incrementAndGet() % 1000)), MessageQueueEntryState.Busy);
                                        message.setToken(timeoutEntry.getTimestamp());
                                        message.setRandom(timeoutEntry.getRandom());
                                        m.withRow(queue.queueColumnFamily, queue.getShardKey(message))
                                                .putColumn(timeoutEntry, column.getStringValue(), queue.metadata.getRetentionTimeout());
                                        MessageMetadataEntry messageIdEntry = MessageMetadataEntry.newMessageId(queue.getCompositeKey(queue.getShardKey(message), timeoutEntry.getMessageId()));
                                        // Add the timeout column to the key
                                        if (message.hasKey()) {
                                            m.withRow(queue.keyIndexColumnFamily, queue.getCompositeKey(queue.getName(), message.getKey()))
                                                    .putEmptyColumn(messageIdEntry, queue.metadata.getRetentionTimeout());
                                        }
                                        context.setAckMessageId(messageIdEntry.getName());
                                    } else {
                                        message.setToken(null);
                                    }
                                    // Update some stats
                                    switch (entry.getState()) {
                                        case Waiting:
                                            queue.stats.incProcessCount();
                                            break;
                                        case Busy:
                                            queue.stats.incReprocessCount();
                                            break;
                                        default:
                                            LOG.warn("Unknown message state: " + entry.getState());
                                            // TODO:
                                            break;
                                    }
                                    entries.add(context);
                                } else {
                                    queue.stats.incInvalidMessageCount();
                                    // TODO: Add to poison queue
                                }
                            } catch (DuplicateMessageException e) {
                                // OK to ignore this error.  All the proper columns will have been deleted in the batch.
                            }
                            break;
                        }
                    default:
                        {
                            // TODO: Error: Unknown type
                            break;
                        }
                }
            }
            return entries;
        } catch (BusyLockException e) {
            queue.stats.incLockContentionCount();
            throw e;
        } catch (Exception e) {
            throw new MessageQueueException("Error processing queue shard : " + shardName, e);
        } finally {
            try {
                m.execute();
            } catch (Exception e) {
                throw new MessageQueueException("Error processing queue shard : " + shardName, e);
            }
        }
    }

}




© 2015 - 2024 Weber Informatics LLC | Privacy Policy