tools/stress/src/org/apache/cassandra/stress/util/SmartThriftClient.java

package org.apache.cassandra.stress.util;
/*
 * 
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF 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.
 * 
 */


import java.net.InetAddress;
import java.nio.ByteBuffer;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.atomic.AtomicInteger;

import com.google.common.collect.Iterators;

import com.datastax.driver.core.Host;
import com.datastax.driver.core.Metadata;
import org.apache.cassandra.stress.settings.StressSettings;
import org.apache.cassandra.thrift.*;
import org.apache.cassandra.utils.ByteBufferUtil;
import org.apache.thrift.TException;

public class SmartThriftClient implements ThriftClient
{

    final String keyspace;
    final Metadata metadata;
    final StressSettings settings;
    final ConcurrentHashMap<InetAddress, ConcurrentLinkedQueue<Client>> cache = new ConcurrentHashMap<>();

    final AtomicInteger queryIdCounter = new AtomicInteger();
    final ConcurrentHashMap<Integer, String> queryStrings = new ConcurrentHashMap<>();
    final ConcurrentHashMap<String, Integer> queryIds = new ConcurrentHashMap<>();
    final Set<InetAddress> whiteset;
    final List<InetAddress> whitelist;

    public SmartThriftClient(StressSettings settings, String keyspace, Metadata metadata)
    {
        this.metadata = metadata;
        this.keyspace = keyspace;
        this.settings = settings;
        if (!settings.node.isWhiteList)
        {
            whiteset = null;
            whitelist = null;
        }
        else
        {
            whiteset = settings.node.resolveAllSpecified();
            whitelist = Arrays.asList(whiteset.toArray(new InetAddress[0]));
        }
    }

    private final AtomicInteger roundrobin = new AtomicInteger();

    private Integer getId(String query)
    {
        Integer r;
        if ((r = queryIds.get(query)) != null)
            return r;
        r = queryIdCounter.incrementAndGet();
        if (queryIds.putIfAbsent(query, r) == null)
        {
            queryStrings.put(r, query);
            return r;
        }
        return queryIds.get(query);
    }

    final class Client
    {
        final Cassandra.Client client;
        final InetAddress server;
        final Map<Integer, Integer> queryMap = new HashMap<>();

        Client(Cassandra.Client client, InetAddress server)
        {
            this.client = client;
            this.server = server;
        }

        Integer get(Integer id, boolean cql3) throws TException
        {
            Integer serverId = queryMap.get(id);
            if (serverId != null)
                return serverId;
            prepare(id, cql3);
            return queryMap.get(id);
        }

       void prepare(Integer id, boolean cql3) throws TException
       {
           String query;
           while ( null == (query = queryStrings.get(id)) ) ;
           if (cql3)
           {
               Integer serverId = client.prepare_cql3_query(ByteBufferUtil.bytes(query), Compression.NONE).itemId;
               queryMap.put(id, serverId);
           }
           else
           {
               Integer serverId = client.prepare_cql_query(ByteBufferUtil.bytes(query), Compression.NONE).itemId;
               queryMap.put(id, serverId);
           }
       }
    }

    private Client get(ByteBuffer pk)
    {
        Set<Host> hosts = metadata.getReplicas(metadata.quote(keyspace), pk);
        InetAddress address = null;
        if (hosts.size() > 0)
        {
            int pos = roundrobin.incrementAndGet() % hosts.size();
            for (int i = 0 ; address == null && i < hosts.size() ; i++)
            {
                if (pos < 0)
                    pos = -pos;
                Host host = Iterators.get(hosts.iterator(), (pos + i) % hosts.size());
                if (whiteset == null || whiteset.contains(host.getAddress()))
                    address = host.getAddress();
            }
        }
        if (address == null)
            address = whitelist.get(ThreadLocalRandom.current().nextInt(whitelist.size()));
        ConcurrentLinkedQueue<Client> q = cache.get(address);
        if (q == null)
        {
            ConcurrentLinkedQueue<Client> newQ = new ConcurrentLinkedQueue<Client>();
            q = cache.putIfAbsent(address, newQ);
            if (q == null)
                q = newQ;
        }
        Client tclient = q.poll();
        if (tclient != null)
            return tclient;
        return new Client(settings.getRawThriftClient(address.getHostAddress()), address);
    }

    @Override
    public void batch_mutate(Map<ByteBuffer, Map<String, List<Mutation>>> record, ConsistencyLevel consistencyLevel) throws TException
    {
        for (Map.Entry<ByteBuffer, Map<String, List<Mutation>>> e : record.entrySet())
        {
            Client client = get(e.getKey());
            try
            {
                client.client.batch_mutate(Collections.singletonMap(e.getKey(), e.getValue()), consistencyLevel);
            } finally
            {
                cache.get(client.server).add(client);
            }
        }
    }

    @Override
    public List<ColumnOrSuperColumn> get_slice(ByteBuffer key, ColumnParent parent, SlicePredicate predicate, ConsistencyLevel consistencyLevel) throws TException
    {
        Client client = get(key);
        try
        {
            return client.client.get_slice(key, parent, predicate, consistencyLevel);
        } finally
        {
            cache.get(client.server).add(client);
        }
    }

    @Override
    public void insert(ByteBuffer key, ColumnParent column_parent, Column column, ConsistencyLevel consistency_level) throws TException
    {
        Client client = get(key);
        try
        {
            client.client.insert(key, column_parent, column, consistency_level);
        } finally
        {
            cache.get(client.server).add(client);
        }
    }

    @Override
    public CqlResult execute_cql_query(String query, ByteBuffer key, Compression compression) throws TException
    {
        Client client = get(key);
        try
        {
            return client.client.execute_cql_query(ByteBufferUtil.bytes(query), compression);
        } finally
        {
            cache.get(client.server).add(client);
        }
    }

    @Override
    public CqlResult execute_cql3_query(String query, ByteBuffer key, Compression compression, ConsistencyLevel consistency) throws TException
    {
        Client client = get(key);
        try
        {
            return client.client.execute_cql3_query(ByteBufferUtil.bytes(query), compression, consistency);
        } finally
        {
            cache.get(client.server).add(client);
        }
    }

    @Override
    public Integer prepare_cql3_query(String query, Compression compression) throws TException
    {
        return getId(query);
    }

    @Override
    public CqlResult execute_prepared_cql3_query(int queryId, ByteBuffer key, List<ByteBuffer> values, ConsistencyLevel consistency) throws TException
    {
        Client client = get(key);
        try
        {
            return client.client.execute_prepared_cql3_query(client.get(queryId, true), values, consistency);
        } finally
        {
            cache.get(client.server).add(client);
        }
    }

    @Override
    public Integer prepare_cql_query(String query, Compression compression) throws TException
    {
        return getId(query);
    }

    @Override
    public CqlResult execute_prepared_cql_query(int queryId, ByteBuffer key, List<ByteBuffer> values) throws TException
    {
        Client client = get(key);
        try
        {
            return client.client.execute_prepared_cql_query(client.get(queryId, true), values);
        } finally
        {
            cache.get(client.server).add(client);
        }
    }

    @Override
    public Map<ByteBuffer, List<ColumnOrSuperColumn>> multiget_slice(List<ByteBuffer> keys, ColumnParent column_parent, SlicePredicate predicate, ConsistencyLevel consistency_level) throws TException
    {
        throw new UnsupportedOperationException();
    }

    @Override
    public List<KeySlice> get_range_slices(ColumnParent column_parent, SlicePredicate predicate, KeyRange range, ConsistencyLevel consistency_level) throws TException
    {
        throw new UnsupportedOperationException();
    }

    @Override
    public List<KeySlice> get_indexed_slices(ColumnParent column_parent, IndexClause index_clause, SlicePredicate column_predicate, ConsistencyLevel consistency_level) throws TException
    {
        throw new UnsupportedOperationException();
    }

}