com.redis.spring.batch.reader.ScanSizeEstimator Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of spring-batch-redis Show documentation
Spring Batch Redis
There is a newer version: 4.0.7
package com.redis.spring.batch.reader;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.function.LongSupplier;
import java.util.function.Predicate;

import org.springframework.util.StringUtils;

import com.redis.spring.batch.util.ConnectionUtils;
import com.redis.spring.batch.util.Predicates;

import io.lettuce.core.AbstractRedisClient;
import io.lettuce.core.RedisFuture;
import io.lettuce.core.ScriptOutputType;
import io.lettuce.core.api.StatefulConnection;
import io.lettuce.core.api.async.RedisScriptingAsyncCommands;
import io.lettuce.core.api.sync.BaseRedisCommands;
import io.lettuce.core.api.sync.RedisServerCommands;

public class ScanSizeEstimator implements LongSupplier {

    public static final long UNKNOWN_SIZE = -1;

    public static final int DEFAULT_SAMPLES = 100;

    private static final String LUA_FILE = "randomkeytype.lua";

    private final AbstractRedisClient client;

    private String scanMatch;

    private String scanType;

    private int samples = DEFAULT_SAMPLES;

    public ScanSizeEstimator(AbstractRedisClient client) {
        this.client = client;
    }

    public String getScanMatch() {
        return scanMatch;
    }

    public void setScanMatch(String match) {
        this.scanMatch = match;
    }

    public int getSamples() {
        return samples;
    }

    public void setSamples(int samples) {
        this.samples = samples;
    }

    public void setScanType(String type) {
        this.scanType = type;
    }

    public String getScanType() {
        return scanType;
    }

    /**
     * Estimates the number of keys that match the given pattern and type.
     * 
     * @return Estimated number of keys matching the given pattern and type. Returns null if database is empty or any error
     *         occurs
     * @throws TimeoutException
     * @throws ExecutionException
     */
    @SuppressWarnings("unchecked")
    @Override
    public long getAsLong() {
        StatefulConnection connection = ConnectionUtils.supplier(client).get();
        BaseRedisCommands sync = ConnectionUtils.sync(connection);
        Long dbsize = ((RedisServerCommands) sync).dbsize();
        if (dbsize == null) {
            return UNKNOWN_SIZE;
        }
        if (!StringUtils.hasLength(scanMatch) && !StringUtils.hasLength(scanType)) {
            return dbsize;
        }
        String digest = ConnectionUtils.loadScript(client, LUA_FILE);
        RedisScriptingAsyncCommands commands = ConnectionUtils.async(connection);
        try {
            connection.setAutoFlushCommands(false);
            List>> futures = new ArrayList<>();
            for (int index = 0; index < samples; index++) {
                futures.add(commands.evalsha(digest, ScriptOutputType.MULTI));
            }
            connection.flushCommands();
            Predicate matchPredicate = Predicates.glob(scanMatch);
            Predicate typePredicate = typePredicate();
            int total = 0;
            int matchCount = 0;
            for (RedisFuture> future : futures) {
                List