package org.archive.wayback.accesscontrol.robotstxt.redis;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.PrintWriter;
import java.io.UnsupportedEncodingException;
import java.util.LinkedList;
import java.util.List;
import java.util.logging.Level;
import java.util.logging.Logger;
import java.util.zip.Deflater;
import java.util.zip.GZIPOutputStream;
import org.apache.commons.io.IOUtils;
import org.archive.util.zip.OpenJDK7GZIPInputStream;
import org.archive.wayback.exception.LiveWebCacheUnavailableException;
import redis.clients.jedis.Jedis;
import redis.clients.jedis.exceptions.JedisConnectionException;
public class RedisRobotsLogic {
private final static Logger LOGGER =
Logger.getLogger(RedisRobotsLogic.class.getName());
private final static int MIN_GZIP_SIZE = 20;
final static String UTF8 = "UTF-8";
private RedisConnectionManager redisConn;
RedisRobotsLogic(RedisConnectionManager redisConn)
{
this.redisConn = redisConn;
}
public <T> T runJedisCmd(JedisRunner<T> runner) throws LiveWebCacheUnavailableException
{
Jedis jedis = null;
try {
jedis = redisConn.getJedisInstance();
return runner.run(jedis);
} catch (JedisConnectionException jce) {
redisConn.returnBrokenJedis(jedis);
LOGGER.log(Level.SEVERE, "Jedis Exception", jce);
jedis = null;
throw new LiveWebCacheUnavailableException("No Jedis");
} finally {
redisConn.returnJedisInstance(jedis);
}
}
public void runJedisCmd(JedisRunnerVoid runner)
{
Jedis jedis = null;
try {
jedis = redisConn.getJedisInstance();
runner.run(jedis);
} catch (JedisConnectionException jce) {
redisConn.returnBrokenJedis(jedis);
LOGGER.log(Level.SEVERE, "Jedis Exception", jce);
jedis = null;
} finally {
redisConn.returnJedisInstance(jedis);
}
}
static class RedisValue
{
String value;
long ttl;
RedisValue(String value, long ttl)
{
this.value = value;
this.ttl = ttl;
}
}
static class KeyRedisValue extends RedisValue
{
String key;
KeyRedisValue(String key, String value, long ttl)
{
super(value, ttl);
this.key = key;
}
}
public RedisValue getValue(final String key) throws LiveWebCacheUnavailableException
{
long startTime = System.currentTimeMillis();
RedisValue value = null;
try {
value = this.runJedisCmd(new JedisRunner<RedisValue>()
{
public RedisValue run(Jedis jedis)
{
//String value = jedis.get(key);
byte[] binValue = null;
try {
binValue = jedis.get(key.getBytes(UTF8));
} catch (UnsupportedEncodingException e1) {
// TODO Auto-generated catch block
e1.printStackTrace();
}
if (binValue == null) {
return null;
}
String stringValue = null;
try {
if (isGzipStream(binValue)) {
InputStream stream = new OpenJDK7GZIPInputStream(new ByteArrayInputStream(binValue));
stringValue = IOUtils.toString(stream, UTF8);
}
} catch (IOException e) {
}
if (stringValue == null) {
try {
stringValue = new String(binValue, UTF8);
} catch (UnsupportedEncodingException e) {
// TODO Auto-generated catch block
e.printStackTrace();
}
}
long ttl = jedis.ttl(key);
return new RedisValue(stringValue, ttl);
}
});
return value;
} finally {
//PerformanceLogger.noteElapsed("RedisGetTTL", System.currentTimeMillis() - startTime, ((value == null) ? "REDIS MISS: " : "REDIS HIT: ") + key);
}
}
public List<RedisValue> getValue(final String[] keys) throws LiveWebCacheUnavailableException
{
long startTime = System.currentTimeMillis();
List<RedisValue> values = null;
try {
values = this.runJedisCmd(new JedisRunner<List<RedisValue>>()
{
public List<RedisValue> run(Jedis jedis)
{
List<String> values = jedis.mget(keys);
List<RedisValue> redisValues = new LinkedList<RedisValue>();
int index = 0;
for (String value : values) {
if (value == null) {
redisValues.add(null);
} else {
long ttl = jedis.ttl(keys[index]);
redisValues.add(new RedisValue(value, ttl));
}
index++;
}
return redisValues;
}
});
} finally {
//PerformanceLogger.noteElapsed("RedisMultiGetTTL", System.currentTimeMillis() - startTime, Arrays.toString(keys));
}
return values;
}
public void updateValue(final String url, final RedisValue value)
{
updateValue(url, value, false);
}
public void updateValue(final String url, final RedisValue value, final boolean gzip)
{
this.runJedisCmd(new JedisRunnerVoid()
{
public void run(Jedis jedis)
{
if (value.value == null) {
jedis.expire(url, (int)value.ttl);
} else if (!gzip || (value.value.length() < MIN_GZIP_SIZE)) {
jedis.setex(url, (int)value.ttl, value.value);
} else {
try {
byte[] array = value.value.getBytes(UTF8);
ByteArrayOutputStream buff = new ByteArrayOutputStream(array.length + 8);
GZIPOutputStream stream = new GZIPOutputStream(buff) {
{
def.setLevel(Deflater.BEST_COMPRESSION);
}
};
stream.write(array);
stream.finish();
jedis.setex(url.getBytes(UTF8), (int)value.ttl, buff.toByteArray());
} catch (IOException io) {
io.printStackTrace();
}
}
}
});
}
public void pushKey(final String list, final String key)
{
this.runJedisCmd(new JedisRunnerVoid()
{
public void run(Jedis jedis)
{
jedis.rpush(list, key);
}
});
}
public void pushKey(final String list, final String key, final int maxSize)
{
this.runJedisCmd(new JedisRunnerVoid()
{
public void run(Jedis jedis)
{
if (jedis.llen(list) < maxSize) {
jedis.rpush(list, key);
}
}
});
}
public KeyRedisValue popKeyAndGet(final String list) throws LiveWebCacheUnavailableException
{
return this.runJedisCmd(new JedisRunner<KeyRedisValue>()
{
public KeyRedisValue run(Jedis jedis)
{
List<String> values = jedis.blpop(0, list);
if (values == null) {
return null;
}
String key = values.get(1);
String value = jedis.get(key);
if (value == null) {
return null;
}
long ttl = jedis.ttl(key);
return new KeyRedisValue(key, value, ttl);
}
});
}
public void close()
{
redisConn.close();
}
interface JedisRunner<T>
{
public T run(Jedis jedis);
}
interface JedisRunnerVoid
{
public void run(Jedis jedis);
}
public void appendLogInfo(PrintWriter info) {
redisConn.appendLogInfo(info);
}
public static boolean isGzipStream(byte[] bytes) {
if (bytes.length < 2) {
return false;
}
int head = ((int) bytes[0] & 0xff) | ((bytes[1] << 8) & 0xff00);
return (OpenJDK7GZIPInputStream.GZIP_MAGIC == head);
}
}