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.objenesis.strategy.StdInstantiatorStrategy;
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);

    private static final int OUTPUT_BUFFER_SIZE = 1<<12;
    private static final int MAX_OUTPUT_BUFFER_SIZE = 1<<24;

    private static final int TIDY_FACTOR = 1<<4;

    private static final int SWITCH_LIMIT = Math.max(MAX_OUTPUT_BUFFER_SIZE, MAX_OUTPUT_BUFFER_SIZE / TIDY_FACTOR);

    private static final ThreadLocal<Output> output = new ThreadLocal<Output>();
    private static final ThreadLocal<Kryo> kryo = new ThreadLocal<Kryo>();

    private int prevPosition;

    private Iterable<KryoFactory.ClassPair> kryoPairs = new ArrayList<KryoFactory.ClassPair>();

    public KryoSerializer() {
    }

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

    public void tidyBuffer() {
        int currentPosition = getOutput().position();

        // If the previous serialized object was large (greater than the switch size) and the current
        // object falls below the switch size, reset the buffer to be small again. If both objects
        // are small (or large), no reallocation occurs.
        if (prevPosition > SWITCH_LIMIT && currentPosition <= SWITCH_LIMIT) {
            getOutput().setBuffer(new byte[OUTPUT_BUFFER_SIZE], MAX_OUTPUT_BUFFER_SIZE);
        }

        prevPosition = currentPosition;
        getOutput().clear();
    }

    private Kryo freshKryo() {
        Kryo k = new Kryo();
        KryoFactory factory = new KryoFactory(new Configuration());

        k.setInstantiatorStrategy(new StdInstantiatorStrategy());
        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.get() == null)
            kryo.set(freshKryo());

        return kryo.get();
    }

    public Output getOutput() {
        if (output.get() == null)
            output.set(new Output(OUTPUT_BUFFER_SIZE, MAX_OUTPUT_BUFFER_SIZE));

        return output.get();
    }

    public byte[] serialize(Object o) {
        LOG.debug("Serializing object: " + o);
        LOG.debug("Registration for " + o + " is: " + getKryo().getRegistration(o.getClass()));
        Output output = getOutput();
        getKryo().writeClassAndObject(output, o);
        byte[] bytes =  output.toBytes();

        tidyBuffer();

        LOG.debug("returning " + bytes.length + " serialized bytes");
        return bytes;

    }

    public Object deserialize(byte[] bytes) {
        LOG.debug("deserializing " + bytes.length + " bytes");
        return getKryo().readClassAndObject(new Input(bytes));
    }

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

@@ -345,16 +345,82 @@
 rayList;
+%0Aimport java.io.ByteArrayOutputStream;%0Aimport java.io.IOException;
 %0A%0A/**%0A *
@@ -606,375 +606,12 @@
 s);%0A
-%0A
     
-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    private static final ThreadLocal%3COutput%3E output = new ThreadLocal%3COutput%3E();
 %0A   
@@ -686,40 +686,8 @@
 ();%0A
-%0A    private int prevPosition;%0A%0A
     
@@ -777,16 +777,62 @@
 Pair%3E();
+%0A    private ByteArrayOutputStream byteStream;
 %0A%0A    pu
@@ -977,606 +977,8 @@
  %7D%0A%0A
-    public void tidyBuffer() %7B%0A        int currentPosition = getOutput().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) %7B%0A            getOutput().setBuffer(new byte%5BOUTPUT_BUFFER_SIZE%5D, MAX_OUTPUT_BUFFER_SIZE);%0A        %7D%0A%0A        prevPosition = currentPosition;%0A        getOutput().clear();%0A    %7D%0A%0A
     
@@ -1029,24 +1029,74 @@
 new Kryo();%0A
+        byteStream = new ByteArrayOutputStream();%0A
         Kryo
@@ -1147,25 +1147,24 @@
 uration());%0A
-%0A
         k.se
@@ -1676,397 +1676,106 @@
 lic 
-Output getOutput() %7B%0A        if (output.get() == null)%0A            output.set(new Output(OUTPUT_BUFFER_SIZE, MAX_OUTPUT_BUFFER_SIZE));%0A%0A        return output.get();%0A    %7D%0A%0A    public byte%5B%5D serialize(Object o) %7B%0A        LOG.debug(%22Serializing object: %22 + o);%0A        LOG.debug(%22Registration for %22 + o + %22 is: %22 + getKryo().getRegistration(o.getClass()));%0A        Output output = getOutput(
+byte%5B%5D serialize(Object o) %7B%0A        byteStream.reset();%0A        Output ko = new Output(byteStream
 );%0A 
@@ -1815,20 +1815,36 @@
 ect(
+k
 o
-utput
 , o);%0A
+        ko.flush();%0A
     
@@ -1866,15 +1866,10 @@
 s = 
- output
+ko
 .toB
@@ -1876,25 +1876,24 @@
 ytes();%0A
-%0A
         
 tidyBuff
@@ -1888,92 +1888,8 @@
     
-tidyBuffer();%0A%0A        LOG.debug(%22returning %22 + bytes.length + %22 serialized bytes%22);
 %0A   
@@ -1907,17 +1907,16 @@
  bytes;%0A
-%0A
     %7D%0A%0A 
@@ -1964,71 +1964,8 @@
 ) %7B%0A
-        LOG.debug(%22deserializing %22 + bytes.length + %22 bytes%22);%0A
     

remove fixed buffer size in kryo serializer