com.spotify.folsom.ketama.KetamaMemcacheClient.java Source code

Introduction

Here is the source code for com.spotify.folsom.ketama.KetamaMemcacheClient.java

Source

/*
 * Copyright (c) 2014-2015 Spotify AB
 *
 * 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 com.spotify.folsom.ketama;

import com.google.common.base.Function;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.spotify.folsom.GetResult;
import com.spotify.folsom.RawMemcacheClient;
import com.spotify.folsom.client.AbstractMultiMemcacheClient;
import com.spotify.folsom.client.Request;
import com.spotify.folsom.client.MultiRequest;
import com.spotify.folsom.client.Utils;

import java.util.Collection;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

public class KetamaMemcacheClient extends AbstractMultiMemcacheClient {

    private static Collection<RawMemcacheClient> clientsOnly(final Collection<AddressAndClient> addressAndClients) {

        int size = addressAndClients.size();
        final List<RawMemcacheClient> clients = Lists.newArrayListWithCapacity(size);
        for (final AddressAndClient client : addressAndClients) {
            clients.add(client.getClient());
        }
        return clients;
    }

    private final Continuum continuum;

    public KetamaMemcacheClient(final Collection<AddressAndClient> clients) {
        super(clientsOnly(clients));

        this.continuum = new Continuum(clients);
    }

    private RawMemcacheClient getClient(final byte[] key) {
        return continuum.findClient(key);
    }

    @Override
    public <T> ListenableFuture<T> send(final Request<T> request) {
        if (request instanceof MultiRequest) {
            // T for Request should always be List<GetResult<byte[]>> here
            // which means that MultiRequest should have T = GetResult<byte[]>
            final MultiRequest<GetResult<byte[]>> multiRequest = (MultiRequest<GetResult<byte[]>>) request;
            if (multiRequest.getKeys().size() > 1) {
                return (ListenableFuture<T>) sendSplitRequest(multiRequest);
            }
        }

        return getClient(request.getKey()).send(request);
    }

    private <T> ListenableFuture<List<T>> sendSplitRequest(final MultiRequest<T> multiRequest) {
        final List<byte[]> keys = multiRequest.getKeys();

        final Map<RawMemcacheClient, List<byte[]>> routing = Maps.newIdentityHashMap();
        final List<RawMemcacheClient> routing2 = Lists.newArrayListWithCapacity(keys.size());
        for (final byte[] key : keys) {
            final RawMemcacheClient client = getClient(key);
            List<byte[]> subKeys = routing.get(client);
            if (subKeys == null) {
                subKeys = Lists.newArrayList();
                routing.put(client, subKeys);
            }
            subKeys.add(key);
            routing2.add(client);
        }

        final Map<RawMemcacheClient, ListenableFuture<List<T>>> futures = Maps.newIdentityHashMap();

        for (final Map.Entry<RawMemcacheClient, List<byte[]>> entry : routing.entrySet()) {
            final List<byte[]> subKeys = entry.getValue();
            final Request<List<T>> subRequest = multiRequest.create(subKeys);
            final RawMemcacheClient client = entry.getKey();
            ListenableFuture<List<T>> send = client.send(subRequest);
            futures.put(client, send);
        }
        final ListenableFuture<List<List<T>>> allFutures = Futures.allAsList(futures.values());
        return Utils.transform(allFutures, new Assembler<>(futures, routing2));
    }

    private static class Assembler<T, R> implements Function<List<List<T>>, List<T>> {
        private final Map<R, ListenableFuture<List<T>>> futures;
        private final List<R> routing2;

        public Assembler(Map<R, ListenableFuture<List<T>>> futures, List<R> routing2) {
            this.futures = futures;
            this.routing2 = routing2;
        }

        @Override
        public List<T> apply(final List<List<T>> ignored) {
            final Map<R, Iterator<T>> map = Maps.newIdentityHashMap();
            for (final Map.Entry<R, ListenableFuture<List<T>>> entry : futures.entrySet()) {
                final R client = entry.getKey();
                map.put(client, Futures.getUnchecked(entry.getValue()).iterator());
            }
            final List<T> result = Lists.newArrayList();
            for (final R memcacheClient : routing2) {
                final Iterator<T> iterator = map.get(memcacheClient);
                result.add(iterator.next());
            }
            return result;
        }
    }
}