• Home
• com.graphql-java
• java-dataloader
• 3.3.0
• source code
• DataLoaderHelper.java

×

Project download

Many resources are needed to download a project. Please understand that we have to compensate our server costs. Thank you in advance.
Project price only 1 $

You can buy this project and download/modify it how often you want.

org.dataloader.DataLoaderHelper Maven / Gradle / Ivy

package org.dataloader;

import org.dataloader.annotations.GuardedBy;
import org.dataloader.annotations.Internal;
import org.dataloader.impl.CompletableFutureKit;
import org.dataloader.scheduler.BatchLoaderScheduler;
import org.dataloader.stats.StatisticsCollector;
import org.dataloader.stats.context.IncrementBatchLoadCountByStatisticsContext;
import org.dataloader.stats.context.IncrementBatchLoadExceptionCountStatisticsContext;
import org.dataloader.stats.context.IncrementCacheHitCountStatisticsContext;
import org.dataloader.stats.context.IncrementLoadCountStatisticsContext;
import org.dataloader.stats.context.IncrementLoadErrorCountStatisticsContext;

import java.time.Clock;
import java.time.Instant;
import java.util.ArrayList;
import java.util.Collection;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.atomic.AtomicReference;

import static java.util.Collections.emptyList;
import static java.util.Collections.singletonList;
import static java.util.concurrent.CompletableFuture.allOf;
import static java.util.concurrent.CompletableFuture.completedFuture;
import static java.util.stream.Collectors.toList;
import static org.dataloader.impl.Assertions.assertState;
import static org.dataloader.impl.Assertions.nonNull;

/**
 * This helps break up the large DataLoader class functionality, and it contains the logic to dispatch the
 * promises on behalf of its peer dataloader.
 *
 * @param  the type of keys
 * @param  the type of values
 */
@Internal
class DataLoaderHelper {

    static class LoaderQueueEntry {

        final K key;
        final V value;
        final Object callContext;

        public LoaderQueueEntry(K key, V value, Object callContext) {
            this.key = key;
            this.value = value;
            this.callContext = callContext;
        }

        K getKey() {
            return key;
        }

        V getValue() {
            return value;
        }

        Object getCallContext() {
            return callContext;
        }
    }

    private final DataLoader dataLoader;
    private final Object batchLoadFunction;
    private final DataLoaderOptions loaderOptions;
    private final CacheMap futureCache;
    private final ValueCache valueCache;
    private final List>> loaderQueue;
    private final StatisticsCollector stats;
    private final Clock clock;
    private final AtomicReference lastDispatchTime;

    DataLoaderHelper(DataLoader dataLoader,
                     Object batchLoadFunction,
                     DataLoaderOptions loaderOptions,
                     CacheMap futureCache,
                     ValueCache valueCache,
                     StatisticsCollector stats,
                     Clock clock) {
        this.dataLoader = dataLoader;
        this.batchLoadFunction = batchLoadFunction;
        this.loaderOptions = loaderOptions;
        this.futureCache = futureCache;
        this.valueCache = valueCache;
        this.loaderQueue = new ArrayList<>();
        this.stats = stats;
        this.clock = clock;
        this.lastDispatchTime = new AtomicReference<>();
        this.lastDispatchTime.set(now());
    }

    Instant now() {
        return clock.instant();
    }

    public Instant getLastDispatchTime() {
        return lastDispatchTime.get();
    }

    Optional> getIfPresent(K key) {
        synchronized (dataLoader) {
            boolean cachingEnabled = loaderOptions.cachingEnabled();
            if (cachingEnabled) {
                Object cacheKey = getCacheKey(nonNull(key));
                if (futureCache.containsKey(cacheKey)) {
                    stats.incrementCacheHitCount(new IncrementCacheHitCountStatisticsContext<>(key));
                    return Optional.of(futureCache.get(cacheKey));
                }
            }
        }
        return Optional.empty();
    }

    Optional> getIfCompleted(K key) {
        synchronized (dataLoader) {
            Optional> cachedPromise = getIfPresent(key);
            if (cachedPromise.isPresent()) {
                CompletableFuture promise = cachedPromise.get();
                if (promise.isDone()) {
                    return cachedPromise;
                }
            }
        }
        return Optional.empty();
    }


    CompletableFuture load(K key, Object loadContext) {
        synchronized (dataLoader) {
            boolean batchingEnabled = loaderOptions.batchingEnabled();
            boolean cachingEnabled = loaderOptions.cachingEnabled();

            stats.incrementLoadCount(new IncrementLoadCountStatisticsContext<>(key, loadContext));

            if (cachingEnabled) {
                return loadFromCache(key, loadContext, batchingEnabled);
            } else {
                return queueOrInvokeLoader(key, loadContext, batchingEnabled, false);
            }
        }
    }

    Object getCacheKey(K key) {
        return loaderOptions.cacheKeyFunction().isPresent() ?
                loaderOptions.cacheKeyFunction().get().getKey(key) : key;
    }

    Object getCacheKeyWithContext(K key, Object context) {
        return loaderOptions.cacheKeyFunction().isPresent() ?
                loaderOptions.cacheKeyFunction().get().getKeyWithContext(key, context) : key;
    }

    DispatchResult dispatch() {
        boolean batchingEnabled = loaderOptions.batchingEnabled();
        final List keys;
        final List