InMemoryQueue.java

/*
 * Copyright © 2014 Cask Data, Inc.
 *
 * 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 co.cask.cdap.data2.transaction.queue.inmemory;

import co.cask.cdap.common.utils.ImmutablePair;
import co.cask.cdap.data2.queue.ConsumerConfig;
import co.cask.cdap.data2.queue.DequeueStrategy;
import co.cask.cdap.data2.queue.QueueEntry;
import co.cask.cdap.data2.transaction.queue.ConsumerEntryState;
import co.cask.tephra.Transaction;
import com.google.common.base.Objects;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.List;
import java.util.NavigableSet;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ConcurrentNavigableMap;
import java.util.concurrent.ConcurrentSkipListMap;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * Implementation of an in-memory queue.
 */
public class InMemoryQueue {

  private static final Logger LOG = LoggerFactory.getLogger(InMemoryQueue.class);

  private final ConcurrentNavigableMap<Key, Item> entries = new ConcurrentSkipListMap<>();

  public void clear() {
    entries.clear();
  }

  public int getSize() {
    return entries.size();
  }

  public void enqueue(long txId, int seqId, QueueEntry entry) {
    entries.put(new Key(txId, seqId), new Item(entry));
  }

  public void undoEnqueue(long txId, int seqId) {
    entries.remove(new Key(txId, seqId));
  }

  public ImmutablePair<List<Key>, List<byte[]>> dequeue(Transaction tx, ConsumerConfig config,
                                                        ConsumerState consumerState, int maxBatchSize) {

    List<Key> keys = Lists.newArrayListWithCapacity(maxBatchSize);
    List<byte[]> datas = Lists.newArrayListWithCapacity(maxBatchSize);
    NavigableSet<Key> keysToScan = consumerState.startKey == null ? entries.navigableKeySet() :
      entries.tailMap(consumerState.startKey).navigableKeySet();
    boolean updateStartKey = true;

    // navigableKeySet is immune to concurrent modification
    for (Key key : keysToScan) {
      if (keys.size() >= maxBatchSize) {
        break;
      }
      if (updateStartKey && key.txId < tx.getFirstShortInProgress()) {
        // See QueueEntryRow#canCommit for reason.
        consumerState.startKey = key;
      }
      if (tx.getReadPointer() < key.txId) {
        // the entry is newer than the current transaction. so are all subsequent entries. bail out.
        break;
      } else if (tx.isInProgress(key.txId)) {
        // the entry is in the exclude list of current transaction. There is a chance that visible entries follow.
        updateStartKey = false; // next time we have to revisit this entry
        continue;
      }
      Item item = entries.get(key);
      if (item == null) {
        // entry was deleted (evicted or undone) after we started iterating
        continue;
      }
      // check whether this is processed already
      ConsumerEntryState state = item.getConsumerState(config.getGroupId());
      if (ConsumerEntryState.PROCESSED.equals(state)) {
        // already processed but not yet evicted. move on
        continue;
      }
      if (config.getDequeueStrategy().equals(DequeueStrategy.FIFO)) {
        // for FIFO, attempt to claim the entry and return it
        if (item.claim(config)) {
          keys.add(key);
          datas.add(item.entry.getData());
        }
        // else: someone else claimed it, or it was already processed, move on, but we may have to revisit this.
        updateStartKey = false;
        continue;
      }
      // for hash/round robin, if group size is 1, just take it
      if (config.getGroupSize() == 1) {
        keys.add(key);
        datas.add(item.entry.getData());
        updateStartKey = false;
        continue;
      }
      // hash by entry hash key or entry id
      int hash;
      if (config.getDequeueStrategy().equals(DequeueStrategy.ROUND_ROBIN)) {
        hash = key.hashCode();
      } else {
        Integer hashFoundInEntry = item.entry.getHashKey(config.getHashKey());
        hash = hashFoundInEntry == null ? 0 : hashFoundInEntry;
      }
      // modulo of a negative is negative, make sure we're positive or 0.
      if (Math.abs(hash) % config.getGroupSize() == config.getInstanceId()) {
        keys.add(key);
        datas.add(item.entry.getData());
        updateStartKey = false;
      }
    }
    return keys.isEmpty() ? null : ImmutablePair.of(keys, datas);
  }

  public void ack(List<Key> dequeuedKeys, ConsumerConfig config) {
    if (dequeuedKeys == null) {
      return;
    }
    for (Key key : dequeuedKeys) {
      Item item = entries.get(key);
      if (item == null) {
        LOG.warn("Attempting to ack non-existing entry " + key);
        continue;
      }
      item.setConsumerState(config, ConsumerEntryState.PROCESSED);
    }
  }

  public void undoDequeue(List<Key> dequeuedKeys, ConsumerConfig config) {
    if (dequeuedKeys == null) {
      return;
    }
    for (Key key : dequeuedKeys) {
      Item item = entries.get(key);
      if (item == null) {
        LOG.warn("Attempting to undo dequeue for non-existing entry " + key);
        continue;
      }
      item.revokeConsumerState(config, config.getDequeueStrategy() == DequeueStrategy.FIFO);
    }
  }

  public void evict(List<Key> dequeuedKeys, int numGroups) {
    if (numGroups < 1) {
      return; // this means no eviction because number of groups is not known
    }
    if (dequeuedKeys == null) {
      return;
    }
    for (Key key : dequeuedKeys) {
      Item item = entries.get(key);
      if (item == null) {
        LOG.warn("Attempting to evict non-existing entry " + key);
        continue;
      }
      if (item.incrementProcessed() >= numGroups) {
        // all consumer groups have processed _and_ reached the post-commit hook: safe to evict
        entries.remove(key);
      }
    }
  }

  /**
   * Used as the key of each queue item, composed of a transaction id and a sequence number within the transaction.
   */
  public static final class Key implements Comparable<Key> {
    final long txId;
    final int seqNo;

    Key(long tx, int seq) {
      txId = tx;
      seqNo = seq;
    }

    public boolean equals(Object obj) {
      if (obj == this) {
        return true;
      }
      if (obj == null || obj.getClass() != Key.class) {
        return false;
      }
      Key other = (Key) obj;
      return txId == other.txId && seqNo == other.seqNo;
    }

    @Override
    public int hashCode() {
      // return ((int) (txId >> 32)) ^ ((int) txId) ^ seqNo;
      return Objects.hashCode(txId, seqNo);
    }

    @Override
    public int compareTo(Key o) {
      if (txId == o.txId) {
        return seqNo == o.seqNo ? 0 : seqNo < o.seqNo ? -1 : 1;
      } else {
        return txId < o.txId ? -1 : 1;
      }
    }

    @Override
    public String toString() {
      return txId + ":" + seqNo;
    }
  }

  // represents an entry of the queue plus meta data
  private static final class Item {
    final QueueEntry entry;
//    ConcurrentMap<Long, ConsumerEntryState> consumerStates = Maps.newConcurrentMap();
    ConcurrentMap<Long, ItemEntryState> consumerStates = Maps.newConcurrentMap();
    AtomicInteger processedCount = new AtomicInteger();

    Item(QueueEntry entry) {
      this.entry = entry;
    }

    ConsumerEntryState getConsumerState(long consumerGroupId) {
      ItemEntryState entryState = consumerStates.get(consumerGroupId);
      return entryState == null ? null : entryState.getState();
    }

    void setConsumerState(ConsumerConfig config, ConsumerEntryState newState) {
      consumerStates.put(config.getGroupId(), new ItemEntryState(config.getInstanceId(), newState));
    }

    void revokeConsumerState(ConsumerConfig config, boolean revokeToClaim) {
      if (revokeToClaim) {
        consumerStates.put(config.getGroupId(), new ItemEntryState(config.getInstanceId(), ConsumerEntryState.CLAIMED));
      } else {
        consumerStates.remove(config.getGroupId());
      }
    }

    boolean claim(ConsumerConfig config) {
      ItemEntryState state = consumerStates.get(config.getGroupId());
      if (state == null) {
        state = consumerStates.putIfAbsent(config.getGroupId(),
                                           new ItemEntryState(config.getInstanceId(), ConsumerEntryState.CLAIMED));
        if (state == null) {
          return true;
        }
      }
      // If the old claimed consumer is gone or if it has been claimed by the same consumer before,
      // then it can be claimed.
      return state.getInstanceId() >= config.getGroupSize()
        || (state.getInstanceId() == config.getInstanceId() && state.getState() == ConsumerEntryState.CLAIMED);
    }

    int incrementProcessed() {
      return processedCount.incrementAndGet();
    }
  }

  /**
   * Represents the state of an item entry.
   */
  private static final class ItemEntryState {
    final int instanceId;
    ConsumerEntryState state;

    ItemEntryState(int instanceId, ConsumerEntryState state) {
      this.instanceId = instanceId;
      this.state = state;
    }

    int getInstanceId() {
      return instanceId;
    }

    ConsumerEntryState getState() {
      return state;
    }

    void setState(ConsumerEntryState state) {
      this.state = state;
    }
  }

  /**
   * The state of a single consumer, gets modified.
   */
  public static class ConsumerState {
    Key startKey = null;
  }

}