package elephantdb.serialize;

import cascading.kryo.KryoFactory;
import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import org.apache.hadoop.conf.Configuration;
import org.apache.log4j.Logger;

import java.util.ArrayList;

/**
 * User: sritchie
 * Date: 1/2/12
 * Time: 12:22 PM
 */
public class KryoSerializer implements Serializer {
    public static final Logger LOG = Logger.getLogger(KryoSerializer.class);

    /**
     * Initial capacity of the Kryo object buffer.
     */
    private static final int INIT_CAPACITY = 2000;

    /**
     * Maximum capacity of the Kryo object buffer.
     */
    private static final int FINAL_CAPACITY = 2000000000;

    private transient Output output;
    private transient Kryo kryo;
    private Iterable<KryoFactory.ClassPair> kryoPairs = new ArrayList<KryoFactory.ClassPair>();

    public KryoSerializer() {
    }

    public KryoSerializer(Iterable<KryoFactory.ClassPair> kryoPairs) {
        setKryoPairs(kryoPairs);
    }

    private Kryo freshKryo() {
        Kryo k = new Kryo();
        KryoFactory factory = new KryoFactory(new Configuration());
        k.setRegistrationRequired(false);
        factory.registerBasic(k, getKryoPairs());
        return k;
    }

    public void setKryoPairs(Iterable<KryoFactory.ClassPair> pairs) {
        kryoPairs = pairs;
    }

    public Iterable<KryoFactory.ClassPair> getKryoPairs() {
        return kryoPairs;
    }

    public Kryo getKryo() {
        if (kryo == null)
            kryo = freshKryo();

        return kryo;
    }

    public Output getOutput() {
        if (output == null)
            output = new Output(INIT_CAPACITY, FINAL_CAPACITY);

        return output;
    }

    public byte[] serialize(Object o) {
        LOG.debug("Serializing object: " + o);
        Output output = getOutput();
        output.clear();
        freshKryo().writeClassAndObject(output, o);
        return output.toBytes();
    }

    public Object deserialize(byte[] bytes) {
        return freshKryo().readClassAndObject(new Input(bytes));
    }

    public <T> T deserialize(byte[] bytes, Class<T> klass) {
        return freshKryo().readObject(new Input(bytes), klass);
    }
}

@@ -490,246 +490,287 @@
     
-/**%0A     * Initial capacity of the Kryo object buffer.%0A     */%0A    private static final int INIT_CAPACITY = 2000;%0A%0A    /**%0A     * Maximum capacity of the Kryo object buffer.%0A     */%0A    private static final int FINAL_CAPACITY = 2000000000
+private static final int OUTPUT_BUFFER_SIZE = 1%3C%3C12;%0A    private static final int MAX_OUTPUT_BUFFER_SIZE = 1%3C%3C24;%0A%0A    private static final int TIDY_FACTOR = 1%3C%3C4;%0A%0A    private static final int SWITCH_LIMIT = Math.max(MAX_OUTPUT_BUFFER_SIZE, MAX_OUTPUT_BUFFER_SIZE / TIDY_FACTOR)
 ;%0A%0A 
@@ -834,24 +834,56 @@
  Kryo kryo;%0A
+%0A    private int prevPosition;%0A%0A
     private 
@@ -1111,24 +1111,595 @@
 rs);%0A    %7D%0A%0A
+    public void tidyBuffer() %7B%0A        int currentPosition = output.position();%0A%0A        // If the previous serialized object was large (greater than the switch size) and the current%0A        // object falls below the switch size, reset the buffer to be small again. If both objects%0A        // are small (or large), no reallocation occurs.%0A        if (prevPosition %3E SWITCH_LIMIT && currentPosition %3C= SWITCH_LIMIT)%0A            output.setBuffer(new byte%5BOUTPUT_BUFFER_SIZE%5D, MAX_OUTPUT_BUFFER_SIZE);%0A%0A        prevPosition = currentPosition;%0A        output.clear();%0A    %7D%0A%0A
     private 
@@ -2339,37 +2339,73 @@
 put(
-INIT_CAPACITY, FINAL_CAPACITY
+OUTPUT_BUFFER_SIZE, MAX_OUTPUT_BUFFER_SIZE);%0A%0A        tidyBuffer(
 );%0A%0A
@@ -2562,32 +2562,8 @@
 ();%0A
-        output.clear();%0A
     
@@ -2639,24 +2639,25 @@
 .toBytes();%0A
+%0A
     %7D%0A%0A    p

change kryo serde to behave like cascading.kryo