MemcachedStableReadExternalCache

package com.atlassian.vcache.internal.memcached;

import com.atlassian.marshalling.api.MarshallingPair;
import com.atlassian.vcache.ExternalCacheException;
import com.atlassian.vcache.ExternalCacheSettings;
import com.atlassian.vcache.PutPolicy;
import com.atlassian.vcache.internal.MetricLabel;
import com.atlassian.vcache.internal.RequestContext;
import com.atlassian.vcache.internal.core.ExternalCacheKeyGenerator;
import com.atlassian.vcache.internal.core.VCacheCoreUtils;
import com.atlassian.vcache.internal.core.metrics.CacheType;
import com.atlassian.vcache.internal.core.metrics.MetricsRecorder;
import com.atlassian.vcache.internal.core.service.AbstractStableReadExternalCache;
import com.atlassian.vcache.internal.core.service.VersionedExternalCacheRequestContext;
import net.spy.memcached.MemcachedClientIF;
import net.spy.memcached.OperationTimeoutException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.time.Duration;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;
import java.util.function.Supplier;
import java.util.stream.Collectors;
import java.util.stream.StreamSupport;

import static com.atlassian.vcache.internal.core.VCacheCoreUtils.isEmpty;
import static com.atlassian.vcache.internal.core.VCacheCoreUtils.unmarshall;
import static com.atlassian.vcache.internal.memcached.MemcachedUtils.expiryTime;
import static com.atlassian.vcache.internal.memcached.MemcachedUtils.putOperationForPolicy;
import static java.util.Objects.requireNonNull;

/**
 * Implementation that backs onto Memcached.
 *
 * @param <V> the value type
 * @since 1.0.0
 */
class MemcachedStableReadExternalCache<V>
        extends AbstractStableReadExternalCache<V> {
    private static final Logger log = LoggerFactory.getLogger(MemcachedStableReadExternalCache.class);

    private final Supplier<MemcachedClientIF> clientSupplier;
    private final Supplier<RequestContext> contextSupplier;
    private final ExternalCacheKeyGenerator keyGenerator;
    private final MarshallingPair<V> valueMarshalling;
    private final int ttlSeconds;

    MemcachedStableReadExternalCache(
            Supplier<MemcachedClientIF> clientSupplier,
            Supplier<RequestContext> contextSupplier,
            ExternalCacheKeyGenerator keyGenerator,
            String name,
            MarshallingPair<V> valueMarshalling,
            ExternalCacheSettings settings,
            MetricsRecorder metricsRecorder,
            Duration lockTimeout) {
        super(name, metricsRecorder, lockTimeout);
        this.clientSupplier = requireNonNull(clientSupplier);
        this.contextSupplier = requireNonNull(contextSupplier);
        this.keyGenerator = requireNonNull(keyGenerator);
        this.valueMarshalling = requireNonNull(valueMarshalling);
        this.ttlSeconds = VCacheCoreUtils.roundUpToSeconds(settings.getDefaultTtl().get());
    }

    @Override
    public boolean internalPut(String internalKey, V value, PutPolicy policy) {
        final String externalKey = ensureCacheContext().externalEntryKeyFor(internalKey);
        final byte[] valueBytes = valueMarshalling.getMarshaller().marshallToBytes(value);

        try {
            final Future<Boolean> putOp = putOperationForPolicy(
                    policy, clientSupplier.get(), externalKey, expiryTime(ttlSeconds), valueBytes);
            return putOp.get();
        } catch (InterruptedException | ExecutionException e) {
            throw new ExternalCacheException(ExternalCacheException.Reason.UNCLASSIFIED_FAILURE, e);
        }
    }

    @Override
    protected void internalRemove(Iterable<String> internalKeys) {
        // There is no bulk delete in the api, so need to remove each one async
        if (isEmpty(internalKeys)) {
            return;
        }

        // Lodge all the requests for delete
        final VersionedExternalCacheRequestContext<V> cacheContext = ensureCacheContext();
        final Map<String, Future<Boolean>> deleteOps =
                StreamSupport.stream(internalKeys.spliterator(), false)
                        .distinct()
                        .collect(Collectors.toMap(
                                k -> k,
                                k -> clientSupplier.get().delete(cacheContext.externalEntryKeyFor(k))
                        ));

        // Need to loop verifying the outcome. If an exception was thrown for an operation,
        // then do not record the removal.
        Exception failureException = null;
        for (Map.Entry<String, Future<Boolean>> delOp : deleteOps.entrySet()) {
            try {
                // Do not care whether the delOp returns true or false. Either way, no entry with the specified key
                // now exists on the Memcache server.
                delOp.getValue().get();
                cacheContext.recordValue(delOp.getKey(), Optional.empty());
            } catch (ExecutionException | InterruptedException ex) {
                log.info("Cache {}: unable to remove key {}", name, delOp.getKey(), ex);
                failureException = ex;
            }
        }

        // Finally, if there was a failure, then re-throw the last one reported (as good as any!)
        if (failureException != null) {
            throw new ExternalCacheException(ExternalCacheException.Reason.NETWORK_FAILURE, failureException);
        }
    }

    @Override
    protected void internalRemoveAll() {
        final VersionedExternalCacheRequestContext<V> cacheContext = ensureCacheContext();
        cacheContext.updateCacheVersion(
                MemcachedUtils.incrementCacheVersion(clientSupplier, cacheContext.externalCacheVersionKey()));
    }

    @Override
    protected Logger getLogger() {
        return log;
    }

    @Override
    protected VersionedExternalCacheRequestContext<V> ensureCacheContext() throws OperationTimeoutException {
        final RequestContext requestContext = contextSupplier.get();

        return requestContext.computeIfAbsent(this, () -> {
            // Need to build a new context, which involves getting the current cache version, or setting it if it does
            // not exist.
            log.trace("Cache {}: Setting up a new context", name);
            final VersionedExternalCacheRequestContext<V> newCacheContext = new VersionedExternalCacheRequestContext<>(
                    keyGenerator, name, requestContext::partitionIdentifier, lockTimeout);
            newCacheContext.updateCacheVersion(
                    MemcachedUtils.obtainCacheVersion(clientSupplier, newCacheContext.externalCacheVersionKey()));
            return newCacheContext;
        });
    }

    @Override
    protected V handleCreation(String internalKey, V candidateValue)
            throws ExecutionException, InterruptedException {
        final VersionedExternalCacheRequestContext<V> cacheContext = ensureCacheContext();
        final byte[] candidateValueBytes = valueMarshalling.getMarshaller().marshallToBytes(candidateValue);
        final String externalKey = cacheContext.externalEntryKeyFor(internalKey);

        // Loop until either able to add the candidate value, or retrieve one that has been added by another thread
        for (; ; ) {
            metricsRecorder.record(name, CacheType.EXTERNAL, MetricLabel.NUMBER_OF_REMOTE_GET, 1);
            final Future<Boolean> addOp = clientSupplier.get().add(externalKey, expiryTime(ttlSeconds), candidateValueBytes);
            if (addOp.get()) {
                // I break here, rather than just return, due to battling with the compiler. Unless written
                // this way, the compiler will not allow the lambda structure.
                break;
            }

            getLogger().info("Cache {}, unable to add candidate for key {}, retrieve what was added", name, internalKey);
            metricsRecorder.record(name, CacheType.EXTERNAL, MetricLabel.NUMBER_OF_REMOTE_GET, 1);
            final Optional<V> otherAddedValue = unmarshall((byte[]) clientSupplier.get().get(externalKey), valueMarshalling);
            if (otherAddedValue.isPresent()) {
                return otherAddedValue.get();
            }

            getLogger().info("Cache {}, unable to retrieve recently added candidate for key {}, looping", name, internalKey);
        }

        return candidateValue;
    }

    @Override
    protected final ExternalCacheException mapException(Exception ex) {
        return MemcachedUtils.mapException(ex);
    }

    @Override
    protected final Optional<V> directGet(String externalKey) {
        return unmarshall((byte[]) clientSupplier.get().get(externalKey), valueMarshalling);
    }

    @Override
    protected final Map<String, Optional<V>> directGetBulk(Set<String> externalKeys) {
        return MemcachedUtils.directGetBulk(externalKeys, clientSupplier, valueMarshalling);
    }
}