path: root/src/main/java/com/thinkaurelius/titan/diskstorage/cassandra/thrift/CassandraThriftKeyColumnValueStore.java

package com.thinkaurelius.titan.diskstorage.cassandra.thrift;

import com.google.common.base.Preconditions;
import com.google.common.base.Predicate;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Iterables;
import com.google.common.collect.Iterators;
import com.thinkaurelius.titan.diskstorage.*;
import com.thinkaurelius.titan.diskstorage.cassandra.thrift.thriftpool.CTConnection;
import com.thinkaurelius.titan.diskstorage.cassandra.thrift.thriftpool.CTConnectionPool;
import com.thinkaurelius.titan.diskstorage.cassandra.utils.CassandraHelper;
import com.thinkaurelius.titan.diskstorage.keycolumnvalue.*;
import com.thinkaurelius.titan.diskstorage.util.*;
import org.apache.cassandra.dht.*;
import org.apache.cassandra.thrift.*;
import org.apache.cassandra.thrift.ConsistencyLevel;
import org.apache.commons.lang.ArrayUtils;
import org.apache.thrift.TException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nullable;
import java.nio.ByteBuffer;
import java.util.*;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import static com.thinkaurelius.titan.diskstorage.cassandra.CassandraTransaction.getTx;

/**
 * A Titan {@code KeyColumnValueStore} backed by Cassandra.
 * This uses the Cassandra Thrift API.
 *
 * @author Dan LaRocque <dalaro@hopcount.org>
 * @see CassandraThriftStoreManager
 */
public class CassandraThriftKeyColumnValueStore implements KeyColumnValueStore {

    private static final Logger logger =
            LoggerFactory.getLogger(CassandraThriftKeyColumnValueStore.class);

    private static final Pattern BROKEN_BYTE_TOKEN_PATTERN = Pattern.compile("^Token\\(bytes\\[(.+)\\]\\)$");

    // Cassandra access
    private final CassandraThriftStoreManager storeManager;
    private final String keyspace;
    private final String columnFamily;
    private final CTConnectionPool pool;
    private final ThriftGetter entryGetter;

    public CassandraThriftKeyColumnValueStore(String keyspace, String columnFamily, CassandraThriftStoreManager storeManager,
                                              CTConnectionPool pool) {
        this.storeManager = storeManager;
        this.keyspace = keyspace;
        this.columnFamily = columnFamily;
        this.pool = pool;
        this.entryGetter = new ThriftGetter(storeManager.getMetaDataSchema(columnFamily));
    }

    /**
     * Call Cassandra's Thrift get_slice() method.
     * <p/>
     * When columnEnd equals columnStart and either startInclusive
     * or endInclusive is false (or both are false), then this
     * method returns an empty list without making any Thrift calls.
     * <p/>
     * If columnEnd = columnStart + 1, and both startInclusive and
     * startExclusive are false, then the arguments effectively form
     * an empty interval.  In this case, as in the one previous,
     * an empty list is returned.  However, it may not necessarily
     * be handled efficiently; a Thrift call might still be made
     * before returning the empty list.
     *
     * @throws com.thinkaurelius.titan.diskstorage.BackendException
     *          when columnEnd < columnStart
     */
    @Override
    public EntryList getSlice(KeySliceQuery query, StoreTransaction txh) throws BackendException {
        Map<StaticBuffer, EntryList> result = getNamesSlice(query.getKey(), query, txh);
        return Iterables.getOnlyElement(result.values(), EntryList.EMPTY_LIST);
    }

    @Override
    public Map<StaticBuffer, EntryList> getSlice(List<StaticBuffer> keys, SliceQuery query, StoreTransaction txh) throws BackendException {
        return getNamesSlice(keys, query, txh);
    }

    public Map<StaticBuffer, EntryList> getNamesSlice(StaticBuffer key,
                                                      SliceQuery query, StoreTransaction txh) throws BackendException {
        return getNamesSlice(ImmutableList.of(key),query,txh);
    }

    public Map<StaticBuffer, EntryList> getNamesSlice(List<StaticBuffer> keys,
                                                      SliceQuery query,
                                                      StoreTransaction txh) throws BackendException {
        ColumnParent parent = new ColumnParent(columnFamily);
        /*
         * Cassandra cannot handle columnStart = columnEnd.
		 * Cassandra's Thrift getSlice() throws InvalidRequestException
		 * if columnStart = columnEnd.
		 */
        if (query.getSliceStart().compareTo(query.getSliceEnd()) >= 0) {
            // Check for invalid arguments where columnEnd < columnStart
            if (query.getSliceEnd().compareTo(query.getSliceStart())<0) {
                throw new PermanentBackendException("columnStart=" + query.getSliceStart() +
                        " is greater than columnEnd=" + query.getSliceEnd() + ". " +
                        "columnStart must be less than or equal to columnEnd");
            }
            if (0 != query.getSliceStart().length() && 0 != query.getSliceEnd().length()) {
                logger.debug("Return empty list due to columnEnd==columnStart and neither empty");
                return KCVSUtil.emptyResults(keys);
            }
        }

        assert query.getSliceStart().compareTo(query.getSliceEnd()) < 0;
        ConsistencyLevel consistency = getTx(txh).getReadConsistencyLevel().getThrift();
        SlicePredicate predicate = new SlicePredicate();
        SliceRange range = new SliceRange();
        range.setCount(query.getLimit() + (query.hasLimit()?1:0)); //Add one for potentially removed last column
        range.setStart(query.getSliceStart().asByteBuffer());
        range.setFinish(query.getSliceEnd().asByteBuffer());
        predicate.setSlice_range(range);

        CTConnection conn = null;
        try {
            conn = pool.borrowObject(keyspace);
            Cassandra.Client client = conn.getClient();
            Map<ByteBuffer, List<ColumnOrSuperColumn>> rows = client.multiget_slice(CassandraHelper.convert(keys),
                    parent,
                    predicate,
                    consistency);

			/*
			 * The final size of the "result" List may be at most rows.size().
			 * However, "result" could also be up to two elements smaller than
			 * rows.size(), depending on startInclusive and endInclusive
			 */
            Map<StaticBuffer, EntryList> results = new HashMap<StaticBuffer, EntryList>();

            for (ByteBuffer key : rows.keySet()) {
                results.put(StaticArrayBuffer.of(key),
                        CassandraHelper.makeEntryList(rows.get(key), entryGetter, query.getSliceEnd(), query.getLimit()));
            }

            return results;
        } catch (Exception e) {
            throw convertException(e);
        } finally {
            pool.returnObjectUnsafe(keyspace, conn);
        }
    }

    private static class ThriftGetter implements StaticArrayEntry.GetColVal<ColumnOrSuperColumn,ByteBuffer> {

        private final EntryMetaData[] schema;

        private ThriftGetter(EntryMetaData[] schema) {
            this.schema = schema;
        }

        @Override
        public ByteBuffer getColumn(ColumnOrSuperColumn element) {
            return element.getColumn().bufferForName();
        }

        @Override
        public ByteBuffer getValue(ColumnOrSuperColumn element) {
            return element.getColumn().bufferForValue();
        }

        @Override
        public EntryMetaData[] getMetaSchema(ColumnOrSuperColumn element) {
            return schema;
        }

        @Override
        public Object getMetaData(ColumnOrSuperColumn element, EntryMetaData meta) {
            switch(meta) {
                case TIMESTAMP:
                    return element.getColumn().getTimestamp();
                case TTL:
                    return element.getColumn().getTtl();
                default:
                    throw new UnsupportedOperationException("Unsupported meta data: " + meta);
            }
        }
    }

    @Override
    public void close() {
        // Do nothing
    }

    @Override
    public void acquireLock(StaticBuffer key, StaticBuffer column, StaticBuffer expectedValue,
                            StoreTransaction txh) throws BackendException {
        throw new UnsupportedOperationException();
    }

    @Override
    public KeyIterator getKeys(@Nullable SliceQuery sliceQuery, StoreTransaction txh) throws BackendException {
        final IPartitioner partitioner = storeManager.getCassandraPartitioner();

        if (!(partitioner instanceof RandomPartitioner) && !(partitioner instanceof Murmur3Partitioner))
            throw new PermanentBackendException("This operation is only allowed when random partitioner (md5 or murmur3) is used.");

        try {
            return new AllTokensIterator(partitioner, sliceQuery, storeManager.getPageSize());
        } catch (Exception e) {
            throw convertException(e);
        }
    }

    @Override
    public KeyIterator getKeys(KeyRangeQuery keyRangeQuery, StoreTransaction txh) throws BackendException {
        final IPartitioner partitioner = storeManager.getCassandraPartitioner();

        // see rant about the reason of this limitation in Astyanax implementation of this method.
        if (!(partitioner instanceof AbstractByteOrderedPartitioner))
            throw new PermanentBackendException("This operation is only allowed when byte-ordered partitioner is used.");

        try {
            return new KeyRangeIterator(partitioner, keyRangeQuery, storeManager.getPageSize(),
                    keyRangeQuery.getKeyStart().asByteBuffer(),
                    keyRangeQuery.getKeyEnd().asByteBuffer());
        } catch (Exception e) {
            throw convertException(e);
        }
    }

    @Override
    public String getName() {
        return columnFamily;
    }

    @Override
    public void mutate(StaticBuffer key, List<Entry> additions, List<StaticBuffer> deletions, StoreTransaction txh) throws BackendException {
        Map<StaticBuffer, KCVMutation> mutations = ImmutableMap.of(key, new KCVMutation(additions, deletions));
        mutateMany(mutations, txh);
    }

    public void mutateMany(Map<StaticBuffer, KCVMutation> mutations, StoreTransaction txh) throws BackendException {
        storeManager.mutateMany(ImmutableMap.of(columnFamily, mutations), txh);
    }

    static BackendException convertException(Throwable e) {
        if (e instanceof TException) {
            return new PermanentBackendException(e);
        } else if (e instanceof TimedOutException) {
            return new TemporaryBackendException(e);
        } else if (e instanceof UnavailableException) {
            return new TemporaryBackendException(e);
        } else if (e instanceof InvalidRequestException) {
            return new PermanentBackendException(e);
        } else {
            return new PermanentBackendException(e);
        }
    }

    @Override
    public String toString() {
        return "CassandraThriftKeyColumnValueStore[ks="
                + keyspace + ", cf=" + columnFamily + "]";
    }


    private List<KeySlice> getKeySlice(ByteBuffer startKey,
                                       ByteBuffer endKey,
                                       SliceQuery columnSlice,
                                       int count) throws BackendException {
        return getRangeSlices(new org.apache.cassandra.thrift.KeyRange().setStart_key(startKey).setEnd_key(endKey).setCount(count), columnSlice);
    }

    private <T extends Token> List<KeySlice> getTokenSlice(T startToken, T endToken,
            SliceQuery sliceQuery, int count) throws BackendException {

        String st = sanitizeBrokenByteToken(startToken);
        String et = sanitizeBrokenByteToken(endToken);

        org.apache.cassandra.thrift.KeyRange kr = new org.apache.cassandra.thrift.KeyRange().setStart_token(st).setEnd_token(et).setCount(count);

        return getRangeSlices(kr, sliceQuery);
    }

    private String sanitizeBrokenByteToken(Token tok) {
        /*
         * Background: https://issues.apache.org/jira/browse/CASSANDRA-5566
         *
         * This check is useful for compatibility with Cassandra server versions
         * 1.2.4 and earlier.
         */
        String st = tok.toString();
        if (!(tok instanceof BytesToken))
            return st;

        // Do a cheap 1-character startsWith before unleashing the regex
        if (st.startsWith("T")) {
            Matcher m = BROKEN_BYTE_TOKEN_PATTERN.matcher(st);
            if (!m.matches()) {
                logger.warn("Unknown token string format: \"{}\"", st);
            } else {
                String old = st;
                st = m.group(1);
                logger.debug("Rewrote token string: \"{}\" -> \"{}\"", old, st);
            }
        }
        return st;
    }

    private List<KeySlice> getRangeSlices(org.apache.cassandra.thrift.KeyRange keyRange, @Nullable SliceQuery sliceQuery) throws BackendException {
        SliceRange sliceRange = new SliceRange();

        if (sliceQuery == null) {
            sliceRange.setStart(ArrayUtils.EMPTY_BYTE_ARRAY)
                    .setFinish(ArrayUtils.EMPTY_BYTE_ARRAY)
                    .setCount(5);
        } else {
            sliceRange.setStart(sliceQuery.getSliceStart().asByteBuffer())
                    .setFinish(sliceQuery.getSliceEnd().asByteBuffer())
                    .setCount((sliceQuery.hasLimit()) ? sliceQuery.getLimit() : Integer.MAX_VALUE);
        }


        CTConnection connection = null;
        try {
            connection = pool.borrowObject(keyspace);

            List<KeySlice> slices =
                    connection.getClient().get_range_slices(new ColumnParent(columnFamily),
                            new SlicePredicate()
                                    .setSlice_range(sliceRange),
                            keyRange,
                            ConsistencyLevel.QUORUM);

            for (KeySlice s : slices) {
                logger.debug("Key {}", ByteBufferUtil.toString(s.key, "-"));
            }

            /* Note: we need to fetch columns for each row as well to remove "range ghosts" */
            List<KeySlice> result = new ArrayList<>(slices.size());
            KeyIterationPredicate pred = new KeyIterationPredicate();
            for (KeySlice ks : slices)
                if (pred.apply(ks))
                    result.add(ks);
            return result;
        } catch (Exception e) {
            throw convertException(e);
        } finally {
            if (connection != null)
                pool.returnObjectUnsafe(keyspace, connection);
        }
    }

    private static class KeyIterationPredicate implements Predicate<KeySlice> {

        @Override
        public boolean apply(@Nullable KeySlice row) {
            return (row != null) && row.getColumns().size() > 0;
        }
    }

    /**
     * Slices rows and columns using tokens. Recall that the partitioner turns
     * keys into tokens. For instance, under RandomPartitioner, tokens are the
     * MD5 hashes of keys.
     */
    public class AbstractBufferedRowIter implements KeyIterator {

        private final int pageSize;
        private final SliceQuery columnSlice;

        private boolean isClosed;
        private boolean seenEnd;
        protected Iterator<KeySlice> ksIter;
        private KeySlice mostRecentRow;

        private final IPartitioner partitioner;
        private Token nextStartToken;
        private final Token endToken;
        private ByteBuffer nextStartKey;

        private boolean omitEndToken;

        public AbstractBufferedRowIter(IPartitioner partitioner,
                SliceQuery columnSlice, int pageSize, Token startToken, Token endToken, boolean omitEndToken) {
            this.pageSize = pageSize;
            this.partitioner = partitioner;
            this.nextStartToken = startToken;
            this.endToken = endToken;
            this.columnSlice = columnSlice;

            this.seenEnd = false;
            this.isClosed = false;
            this.ksIter = Iterators.emptyIterator();
            this.mostRecentRow = null;
            this.omitEndToken = omitEndToken;
        }

        @Override
        public boolean hasNext() {
            ensureOpen();

            if (!ksIter.hasNext() && !seenEnd) {
                try {
                    ksIter = rebuffer().iterator();
                } catch (BackendException e) {
                    throw new RuntimeException(e);
                }
            }

            return ksIter.hasNext();
        }

        @Override
        public StaticBuffer next() {
            ensureOpen();

            if (!hasNext())
                throw new NoSuchElementException();

            mostRecentRow = ksIter.next();

            Preconditions.checkNotNull(mostRecentRow);
            return StaticArrayBuffer.of(mostRecentRow.bufferForKey());
        }

        @Override
        public void close() {
            closeIterator();
        }

        @Override
        public void remove() {
            throw new UnsupportedOperationException();
        }

        @Override
        public RecordIterator<Entry> getEntries() {
            ensureOpen();

            return new RecordIterator<Entry>() {
                final Iterator<Entry> columns =
                        CassandraHelper.makeEntryIterator(mostRecentRow.getColumns(),
                                entryGetter, columnSlice.getSliceEnd(),
                                columnSlice.getLimit());

                @Override
                public boolean hasNext() {
                    ensureOpen();
                    return columns.hasNext();
                }

                @Override
                public Entry next() {
                    ensureOpen();
                    return columns.next();
                }

                @Override
                public void close() {
                    closeIterator();
                }

                @Override
                public void remove() {
                    throw new UnsupportedOperationException();
                }
            };
        }

        private void ensureOpen() {
            if (isClosed)
                throw new IllegalStateException("Iterator has been closed.");
        }

        private void closeIterator() {
            if (!isClosed) {
                isClosed = true;
            }
        }

        private List<KeySlice> rebuffer() throws BackendException {

            Preconditions.checkArgument(!seenEnd);

            return checkFreshSlices(getNextKeySlices());
        }

        protected List<KeySlice> checkFreshSlices(List<KeySlice> ks) {

            if (0 == ks.size()) {
                seenEnd = true;
                return Collections.emptyList();
            }

            nextStartKey = ks.get(ks.size() - 1).bufferForKey();
            nextStartToken = partitioner.getToken(nextStartKey);

            if (nextStartToken.equals(endToken)) {
                seenEnd = true;
                if (omitEndToken)
                    ks.remove(ks.size() - 1);
            }

            return ks;
        }

        protected final List<KeySlice> getNextKeySlices() throws BackendException {
            return getTokenSlice(nextStartToken, endToken, columnSlice, pageSize);
        }
    }

    private final class AllTokensIterator extends AbstractBufferedRowIter {
        public AllTokensIterator(IPartitioner partitioner, SliceQuery columnSlice, int pageSize) {
            super(partitioner, columnSlice, pageSize, partitioner.getMinimumToken(), partitioner.getMinimumToken(), false);
        }
    }

    private final class KeyRangeIterator extends AbstractBufferedRowIter {
        public KeyRangeIterator(IPartitioner partitioner, SliceQuery columnSlice,
                int pageSize, ByteBuffer startKey, ByteBuffer endKey) throws BackendException {
            super(partitioner, columnSlice, pageSize, partitioner.getToken(startKey), partitioner.getToken(endKey), true);

            Preconditions.checkArgument(partitioner instanceof AbstractByteOrderedPartitioner);

            // Get first slice with key range instead of token range. Token
            // ranges are start-exclusive, key ranges are start-inclusive. Both
            // are end-inclusive. If we don't make the call below, then we will
            // erroneously miss startKey.
            List<KeySlice> ks = getKeySlice(startKey, endKey, columnSlice, pageSize);

            this.ksIter = checkFreshSlices(ks).iterator();
        }
    }
}