io.lettuce.core.codec.CompressionCodec Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of lettuce-core Show documentation
Advanced and thread-safe Java Redis client for synchronous, asynchronous, and reactive usage. Supports Cluster, Sentinel, Pipelining, Auto-Reconnect, Codecs and much more.
The newest version!
package io.lettuce.core.codec;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.ByteBuffer;
import java.util.zip.DeflaterOutputStream;
import java.util.zip.GZIPInputStream;
import java.util.zip.GZIPOutputStream;
import java.util.zip.InflaterInputStream;

import io.lettuce.core.internal.LettuceAssert;

/**
 * A compressing/decompressing {@link RedisCodec} that wraps a typed {@link RedisCodec codec} and compresses values using GZIP
 * or Deflate. See {@link io.lettuce.core.codec.CompressionCodec.CompressionType} for supported compression types.
 *
 * @author Mark Paluch
 */
public abstract class CompressionCodec {

    private CompressionCodec() {
    }

    /**
     * A {@link RedisCodec} that compresses values from a delegating {@link RedisCodec}.
     *
     * @param delegate codec used for key-value encoding/decoding, must not be {@code null}.
     * @param compressionType the compression type, must not be {@code null}.
     * @param  Key type.
     * @param  Value type.
     * @return Value-compressing codec.
     */
    @SuppressWarnings({ "rawtypes", "unchecked" })
    public static  RedisCodec valueCompressor(RedisCodec delegate, CompressionType compressionType) {
        LettuceAssert.notNull(delegate, "RedisCodec must not be null");
        LettuceAssert.notNull(compressionType, "CompressionType must not be null");
        return (RedisCodec) new CompressingValueCodecWrapper((RedisCodec) delegate, compressionType);
    }

    private static class CompressingValueCodecWrapper implements RedisCodec {

        private RedisCodec delegate;

        private CompressionType compressionType;

        public CompressingValueCodecWrapper(RedisCodec delegate, CompressionType compressionType) {
            this.delegate = delegate;
            this.compressionType = compressionType;
        }

        @Override
        public Object decodeKey(ByteBuffer bytes) {
            return delegate.decodeKey(bytes);
        }

        @Override
        public Object decodeValue(ByteBuffer bytes) {
            try {
                return delegate.decodeValue(decompress(bytes));
            } catch (IOException e) {
                throw new IllegalStateException(e);
            }
        }

        @Override
        public ByteBuffer encodeKey(Object key) {
            return delegate.encodeKey(key);
        }

        @Override
        public ByteBuffer encodeValue(Object value) {
            try {
                return compress(delegate.encodeValue(value));
            } catch (IOException e) {
                throw new IllegalStateException(e);
            }
        }

        private ByteBuffer compress(ByteBuffer source) throws IOException {
            if (source.remaining() == 0) {
                return source;
            }

            ByteArrayOutputStream outputStream = new ByteArrayOutputStream(source.remaining() / 2);
            OutputStream compressor = null;

            try {
                try (ByteBufferInputStream sourceStream = new ByteBufferInputStream(source)) {
                    if (compressionType == CompressionType.GZIP) {
                        compressor = new GZIPOutputStream(outputStream);
                    }

                    if (compressionType == CompressionType.DEFLATE) {
                        compressor = new DeflaterOutputStream(outputStream);
                    }
                    copy(sourceStream, compressor);
                } finally {

                    if (compressor != null) {
                        compressor.close();
                    }
                }

                return ByteBuffer.wrap(outputStream.toByteArray());
            } finally {
                outputStream.close();
            }
        }

        private ByteBuffer decompress(ByteBuffer source) throws IOException {
            if (source.remaining() == 0) {
                return source;
            }

            InputStream decompressor = null;
            ByteArrayOutputStream outputStream = new ByteArrayOutputStream(source.remaining() * 2);

            try {
                try (ByteBufferInputStream sourceStream = new ByteBufferInputStream(source);) {

                    if (compressionType == CompressionType.GZIP) {
                        decompressor = new GZIPInputStream(sourceStream);
                    }

                    if (compressionType == CompressionType.DEFLATE) {
                        decompressor = new InflaterInputStream(sourceStream);
                    }

                    copy(decompressor, outputStream);
                } finally {
                    if (decompressor != null) {
                        decompressor.close();
                    }
                }

                return ByteBuffer.wrap(outputStream.toByteArray());
            } finally {
                outputStream.close();
            }
        }

    }

    /**
     * Copies all bytes from the input stream to the output stream. Does not close or flush either stream.
     *
     * @param from the input stream to read from
     * @param to the output stream to write to
     * @return the number of bytes copied
     * @throws IOException if an I/O error occurs
     */
    private static long copy(InputStream from, OutputStream to) throws IOException {
        LettuceAssert.notNull(from, "From must not be null");
        LettuceAssert.notNull(to, "From must not be null");
        byte[] buf = new byte[4096];
        long total = 0;
        while (true) {
            int r = from.read(buf);
            if (r == -1) {
                break;
            }
            to.write(buf, 0, r);
            total += r;
        }
        return total;
    }

    public enum CompressionType {
        GZIP, DEFLATE;
    }

}