com.redis.spring.batch.common.BatchUtils Maven / Gradle / Ivy

Go to download
package com.redis.spring.batch.common;

import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.nio.ByteBuffer;
import java.time.Duration;
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.Function;
import java.util.function.Supplier;

import org.springframework.util.FileCopyUtils;

import com.redis.lettucemod.RedisModulesClient;
import com.redis.lettucemod.api.StatefulRedisModulesConnection;
import com.redis.lettucemod.cluster.RedisModulesClusterClient;
import com.redis.lettucemod.cluster.api.StatefulRedisModulesClusterConnection;

import io.lettuce.core.AbstractRedisClient;
import io.lettuce.core.ReadFrom;
import io.lettuce.core.RedisFuture;
import io.lettuce.core.codec.ByteArrayCodec;
import io.lettuce.core.codec.RedisCodec;
import io.lettuce.core.codec.StringCodec;

public abstract class BatchUtils {

	private BatchUtils() {
	}

	public static String readFile(String filename) throws IOException {
		try (InputStream inputStream = BatchUtils.class.getClassLoader().getResourceAsStream(filename)) {
			return FileCopyUtils.copyToString(new InputStreamReader(inputStream));
		}
	}

	public static  Function stringKeyFunction(RedisCodec codec) {
		Function encode = StringCodec.UTF8::encodeKey;
		return encode.andThen(codec::decodeKey);
	}

	public static  Function toStringKeyFunction(RedisCodec codec) {
		Function encode = codec::encodeKey;
		return encode.andThen(StringCodec.UTF8::decodeKey);
	}

	public static  Function stringValueFunction(RedisCodec codec) {
		Function encode = StringCodec.UTF8::encodeValue;
		return encode.andThen(codec::decodeValue);
	}

	public static  Function toStringValueFunction(RedisCodec codec) {
		Function encode = codec::encodeValue;
		return encode.andThen(StringCodec.UTF8::decodeValue);
	}

	public static  Function toByteArrayKeyFunction(RedisCodec codec) {
		Function encode = codec::encodeKey;
		return encode.andThen(ByteArrayCodec.INSTANCE::decodeKey);
	}

	@SuppressWarnings("resource")
	public static  Supplier> supplier(AbstractRedisClient client,
			RedisCodec codec, ReadFrom readFrom) {
		if (client instanceof RedisModulesClusterClient) {
			RedisModulesClusterClient clusterClient = (RedisModulesClusterClient) client;
			return () -> connection(clusterClient, codec, readFrom);
		}
		RedisModulesClient redisClient = (RedisModulesClient) client;
		return () -> redisClient.connect(codec);
	}

	public static  StatefulRedisModulesConnection connection(AbstractRedisClient client,
			RedisCodec codec, ReadFrom readFrom) {
		if (client instanceof RedisModulesClusterClient) {
			return connection((RedisModulesClusterClient) client, codec, readFrom);
		}
		return ((RedisModulesClient) client).connect(codec);
	}

	public static  StatefulRedisModulesClusterConnection connection(RedisModulesClusterClient client,
			RedisCodec codec, ReadFrom readFrom) {
		StatefulRedisModulesClusterConnection connection = client.connect(codec);
		if (readFrom != null) {
			connection.setReadFrom(readFrom);
		}
		return connection;
	}

	public static  List getAll(Duration timeout, Iterable> futures)
			throws TimeoutException, InterruptedException, ExecutionException {
		List items = new ArrayList<>();
		long nanos = timeout.toNanos();
		long time = System.nanoTime();
		for (RedisFuture f : futures) {
			if (timeout.isNegative()) {
				items.add(f.get());
			} else {
				if (nanos < 0) {
					throw new TimeoutException(String.format("Timed out after %s", timeout));
				}
				T item = f.get(nanos, TimeUnit.NANOSECONDS);
				items.add(item);
				long now = System.nanoTime();
				nanos -= now - time;
				time = now;
			}
		}
		return items;
	}

}