/*

  * Licensed to the Apache Software Foundation (ASF) under one

  * or more contributor license agreements.  See the NOTICE file

  * distributed with this work for additional information

  * regarding copyright ownership.  The ASF licenses this file

  * to you under the Apache License, Version 2.0 (the

  * "License"); you may not use this file except in compliance

  * with the License.  You may obtain a copy of the License at

  *

  *     http://www.apache.org/licenses/LICENSE-2.0

  *

  * Unless required by applicable law or agreed to in writing, software

  * distributed under the License is distributed on an "AS IS" BASIS,

  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

  * See the License for the specific language governing permissions and

  * limitations under the License.

  */

 package org.apache.giraph.writable.kryo;

 import java.io.DataInput;

 import java.io.DataOutput;

 import java.util.Arrays;

 import java.util.Collections;

 import java.util.LinkedHashMap;

 import java.util.Map;

 import java.util.Map.Entry;

 import java.util.Random;

 import com.esotericsoftware.kryo.util.DefaultClassResolver;

 import de.javakaffee.kryoserializers.guava.ImmutableMapSerializer;

 import org.apache.giraph.conf.GiraphConfigurationSettable;

 import com.esotericsoftware.kryo.ClassResolver;

 import com.esotericsoftware.kryo.ReferenceResolver;

 import com.esotericsoftware.kryo.util.MapReferenceResolver;

 import org.apache.giraph.types.ops.collections.Basic2ObjectMap;

 import org.apache.giraph.types.ops.collections.BasicSet;

 import org.apache.giraph.writable.kryo.markers.NonKryoWritable;

 import org.apache.giraph.writable.kryo.markers.KryoIgnoreWritable;

 import org.apache.giraph.writable.kryo.serializers.ArraysAsListSerializer;

 import org.apache.giraph.writable.kryo.serializers.CollectionsNCopiesSerializer;

 import org.apache.giraph.writable.kryo.serializers.DirectWritableSerializer;

 import org.apache.giraph.writable.kryo.serializers.FastUtilSerializer;

 import org.apache.giraph.writable.kryo.serializers.ImmutableBiMapSerializerUtils;

 import org.apache.giraph.writable.kryo.serializers.ReusableFieldSerializer;

 import org.apache.hadoop.conf.Configurable;

 import org.apache.hadoop.conf.Configuration;

 import org.apache.hadoop.io.Writable;

 import org.apache.log4j.Logger;

 import org.objenesis.strategy.StdInstantiatorStrategy;

 import com.esotericsoftware.kryo.Kryo;

 import com.esotericsoftware.kryo.Serializer;

 import com.esotericsoftware.kryo.factories.SerializerFactory;

 import com.esotericsoftware.kryo.io.Input;

 import com.esotericsoftware.kryo.io.InputChunked;

 import com.esotericsoftware.kryo.io.Output;

 import com.esotericsoftware.kryo.io.OutputChunked;

 import com.esotericsoftware.kryo.pool.KryoCallback;

 import com.esotericsoftware.kryo.pool.KryoFactory;

 import com.esotericsoftware.kryo.pool.KryoPool;

 import com.esotericsoftware.kryo.serializers.ClosureSerializer;

 import com.esotericsoftware.kryo.serializers.FieldSerializer;

 import com.esotericsoftware.kryo.util.ObjectMap;

 import com.google.common.base.Preconditions;

 import de.javakaffee.kryoserializers.guava.ImmutableListSerializer;

 /**

  * Kryo instance that provides serialization through DataInput/DataOutput

  * that org.apache.hadoop.io.Writable uses.

  *

  * All public APIs are static.

  *

  * It extends Kryo to reuse KryoPool functionality, but have additional needed

  * objects cached as well. If we move to ThreadLocal or other caching

  * technique, we can use composition, instead of inheritance here.

  *

  * TODO: Refactor this class into two separate classes depending on

  * whether the reference tracking is enabled or disabled.

  */

 public class HadoopKryo extends Kryo {

   /** Pool of reusable Kryo objects, since they are expensive to create */

   private static final KryoPool KRYO_POOL = new KryoPool.Builder(

       new KryoFactory() {

         @Override

         public Kryo create() {

           return createKryo(true, true);

         }

       }).build();

   /** Thread local HadoopKryo object */

   private static final ThreadLocal<HadoopKryo> KRYO =

     new ThreadLocal<HadoopKryo>() {

       @Override protected HadoopKryo initialValue() {

         return  createKryo(false, false);

       }

     };

   /**

    * List of interfaces/parent classes that will not be allowed to be

    * serialized, together with explanation of why, that will be shown

    * when throwing such exception

    */

   private static final Map<Class<?>, String> NON_SERIALIZABLE;

   static {

     NON_SERIALIZABLE = new LinkedHashMap<>();

     NON_SERIALIZABLE.put(

         NonKryoWritable.class,

         "it is marked to not allow serialization, " +

         "look at the class for more details");

     NON_SERIALIZABLE.put(

         KryoWritableWrapper.class, "recursion is disallowed");

     NON_SERIALIZABLE.put(

         Configuration.class,

         "it cannot be supported since it contains ClassLoader");

     NON_SERIALIZABLE.put(

         GiraphConfigurationSettable.class, "configuration cannot be set");

     NON_SERIALIZABLE.put(

         Configurable.class, "configuration cannot be set");

     NON_SERIALIZABLE.put(

         Random.class,

         "it should be rarely serialized, since it would create same stream " +

         "of numbers everywhere, use TransientRandom instead");

     NON_SERIALIZABLE.put(

         Logger.class,

         "Logger must be a static field");

   }

   /** Reusable Input object */

   private InputChunked input;

   /** Reusable Output object */

   private OutputChunked output;

   /** Reusable DataInput wrapper stream */

   private DataInputWrapperStream dataInputWrapperStream;

   /** Reusable DataOutput wrapper stream */

   private DataOutputWrapperStream dataOutputWrapperStream;

   /**

    * Map of already initialized serializers used

    * for readIntoObject/writeOutOfObject pair of methods

    */

   private final ObjectMap<Class<?>, ReusableFieldSerializer<Object>>

   classToIntoSerializer = new ObjectMap<>();

   /** Hide constructor, so all access go through pool of cached objects */

   private HadoopKryo() {

   }

   /**

    * Constructor that takes custom class resolver and reference resolver.

    * @param classResolver Class resolver

    * @param referenceResolver Reference resolver

    */

   private HadoopKryo(ClassResolver classResolver,

                      ReferenceResolver referenceResolver) {

     super(classResolver, referenceResolver);

   }

   // Public API:

   /**

    * Write type of given object and the object itself to the output stream.

    * Inverse of readClassAndObj.

    *

    * @param out Output stream

    * @param object Object to write

    */

   public static void writeClassAndObj(

           final DataOutput out, final Object object) {

     writeInternal(out, object, false);

   }

   /**

    * Read object from the input stream, by reading first type of the object,

    * and then all of its fields.

    * Inverse of writeClassAndObject.

    *

    * @param in Input stream

    * @return Deserialized object

    * @param <T> Type of the object being read

    */

   public static <T> T readClassAndObj(DataInput in) {

     return readInternal(in, null, false);

   }

   /**

    * Write an object to output, in a way that can be read by readIntoObject.

    *

    * @param out Output stream

    * @param object Object to be written

    */

   public static void writeOutOfObject(

       final DataOutput out, final Object object) {

     writeInternal(out, object, true);

   }

   /**

    * Reads an object, from input, into a given object,

    * allowing object reuse.

    * Inverse of writeOutOfObject.

    *

    * @param in Input stream

    * @param object Object to fill from input

    */

   public static void readIntoObject(DataInput in, Object object) {

     readInternal(in, object, true);

   }

   /**

    * Writes class and object to specified output stream with specified

    * Kryo object. It does not use interim buffers.

    * @param kryo Kryo object

    * @param out Output stream

    * @param object Object

    */

   public static void writeWithKryo(

           final HadoopKryo kryo, final Output out,

           final Object object) {

     kryo.writeClassAndObject(out, object);

     out.close();

   }

   /**

    * Write out of object with given kryo

    * @param kryo Kryo object

    * @param out Output

    * @param object Object to write

    */

   public static void writeWithKryoOutOfObject(

           final HadoopKryo kryo, final Output out,

           final Object object) {

     kryo.writeOutOfObject(out, object);

     out.close();

   }

   /**

    * Reads class and object from specified input stream with

    * specified kryo object.

    * it does not use interim buffers.

    * @param kryo Kryo object

    * @param in Input buffer

    * @param <T> Object type parameter

    * @return Object

    */

   public static <T> T readWithKryo(

           final HadoopKryo kryo, final Input in) {

     T object;

     object = (T) kryo.readClassAndObject(in);

     in.close();

     return object;

   }

   /**

    * Read into object with given kryo.

    * @param kryo Kryo object

    * @param in Input

    * @param object Object to read into

    */

   public static void readWithKryoIntoObject(

           final HadoopKryo kryo, final Input in, Object object) {

     kryo.readIntoObject(in, object);

     in.close();

   }

   /**

    * Create copy of the object, by magically recursively copying

    * all of its fields, keeping reference structures (like cycles)

    *

    * @param object Object to be copied

    * @return Copy of the object.

    * @param <T> Type of the object

    */

   public static <T> T createCopy(final T object) {

     return KRYO_POOL.run(new KryoCallback<T>() {

       @Override

       public T execute(Kryo kryo) {

         return kryo.copy(object);

       }

     });

   }

   /**

    * Returns a kryo which doesn't track objects, hence

    * serialization of recursive/nested objects is not

    * supported.

    *

    * Reference tracking significantly degrades the performance

    * since kryo has to store all serialized objects and search

    * the history to check if an object has been already serialized.

    *

    * @return Hadoop kryo which doesn't track objects.

    */

   public static HadoopKryo getNontrackingKryo() {

     return KRYO.get();

   }

   // Private implementation:

   /**

    * Create new instance of HadoopKryo, properly initialized.

    * @param trackReferences if true, object references are tracked.

    * @param hasBuffer if true, an interim buffer is used.

    * @return new HadoopKryo instance

    */

   private static HadoopKryo createKryo(boolean trackReferences,

                                        boolean hasBuffer) {

     HadoopKryo kryo;

     if (trackReferences) {

       kryo = new HadoopKryo();

     } else {

       // Only use GiraphClassResolver if it is properly initialized.

       // This is to enable test cases which use KryoSimpleWrapper

       // but don't start ZK.

       kryo = new HadoopKryo(

               GiraphClassResolver.isInitialized() ? new GiraphClassResolver() :

                                                     new DefaultClassResolver(),

               new MapReferenceResolver());

     }

     try {

       kryo.register(Class.forName("java.lang.invoke.SerializedLambda"));

       kryo.register(Class.forName(

         "com.esotericsoftware.kryo.serializers.ClosureSerializer$Closure"),

         new ClosureSerializer());

     } catch (ClassNotFoundException e) {

       throw new IllegalStateException(

         "Trying to use Kryo on Java version " +

           System.getProperty("java.version") +

           ", but unable to find needed classes", e);

     }

     kryo.register(Arrays.asList().getClass(), new ArraysAsListSerializer());

     kryo.register(Collections.nCopies(1, new Object()).getClass(),

         new CollectionsNCopiesSerializer());

     ImmutableListSerializer.registerSerializers(kryo);

     ImmutableMapSerializer.registerSerializers(kryo);

     ImmutableBiMapSerializerUtils.registerSerializers(kryo);

     // There are many fastutil classes, register them at the end,

     // so they don't use up small registration numbers

     FastUtilSerializer.registerAll(kryo);

     kryo.setInstantiatorStrategy(new DefaultInstantiatorStrategy(

         new StdInstantiatorStrategy()));

     SerializerFactory customSerializerFactory = new SerializerFactory() {

       @SuppressWarnings("rawtypes")

       @Override

       public Serializer makeSerializer(Kryo kryo, final Class<?> type) {

         for (final Entry<Class<?>, String> entry :

             NON_SERIALIZABLE.entrySet()) {

           if (entry.getKey().isAssignableFrom(type)) {

             // Allow Class object to be serialized, but not a live instance.

             return new Serializer() {

               @Override

               public Object read(Kryo kryo, Input input, Class type) {

                 throw new RuntimeException("Cannot serialize " + type +

                     ". Objects being serialized cannot capture " +

                     entry.getKey() + " because " + entry.getValue() +

                     ". Either remove field in question" +

                     ", or make it transient (so that it isn't serialized)");

               }

               @Override

               public void write(Kryo kryo, Output output, Object object) {

                 throw new RuntimeException("Cannot serialize " + type +

                     ". Objects being serialized cannot capture " +

                     entry.getKey() + " because " + entry.getValue() +

                     ". Either remove field in question" +

                     ", or make it transient (so that it isn't serialized)");

               }

             };

           }

         }

         if (Writable.class.isAssignableFrom(type) &&

             !KryoIgnoreWritable.class.isAssignableFrom(type) &&

             // remove BasicSet, BasicArrayList and Basic2ObjectMap temporarily,

             // for lack of constructors

             !BasicSet.class.isAssignableFrom(type) &&

             !Basic2ObjectMap.class.isAssignableFrom(type)) {

           // use the Writable method defined by the type

           DirectWritableSerializer serializer = new DirectWritableSerializer();

           return serializer;

         } else {

           FieldSerializer serializer = new FieldSerializer<>(kryo, type);

           serializer.setIgnoreSyntheticFields(false);

           return serializer;

         }

       }

     };

     kryo.addDefaultSerializer(Writable.class, customSerializerFactory);

     kryo.setDefaultSerializer(customSerializerFactory);

     if (hasBuffer) {

       kryo.input = new InputChunked(4096);

       kryo.output = new OutputChunked(4096);

       kryo.dataInputWrapperStream = new DataInputWrapperStream();

       kryo.dataOutputWrapperStream = new DataOutputWrapperStream();

     }

     if (!trackReferences) {

       kryo.setReferences(false);

       // Auto reset can only be disabled if the GiraphClassResolver is

       // properly initialized.

       if (GiraphClassResolver.isInitialized()) {

         kryo.setAutoReset(false);

       }

     }

     return kryo;

   }

   /**

    * Register serializer for class with class name

    *

    * @param kryo HadoopKryo

    * @param className Name of the class for which to register serializer

    * @param serializer Serializer to use

    */

   private static void registerSerializer(HadoopKryo kryo, String className,

       Serializer serializer) {

     try {

       kryo.register(Class.forName(className), serializer);

     } catch (ClassNotFoundException e) {

       throw new IllegalStateException("Class " + className + " is missing", e);

     }

   }

   /**

    * Initialize reusable objects for reading from given DataInput.

    *

    * @param in Input stream

    */

   private void setDataInput(DataInput in) {

     dataInputWrapperStream.setDataInput(in);

     input.setInputStream(dataInputWrapperStream);

   }

   /**

    * Initialize reusable objects for writing into given DataOutput.

    *

    *  @param out Output stream

    */

   private void setDataOutput(DataOutput out) {

     dataOutputWrapperStream.setDataOutput(out);

     output.setOutputStream(dataOutputWrapperStream);

   }

   /**

    * Get or create reusable serializer for given class.

    *

    * @param type Type of the object

    * @return Serializer

    */

   private ReusableFieldSerializer<Object> getOrCreateReusableSerializer(

       Class<?> type) {

     ReusableFieldSerializer<Object> serializer =

         classToIntoSerializer.get(type);

     if (serializer == null) {

       serializer = new ReusableFieldSerializer<>(this, type);

       classToIntoSerializer.put(type, serializer);

     }

     return serializer;

   }

   /**

    * Internal write implementation, that reuses HadoopKryo objects

    * from the pool.

    *

    * @param out Output stream

    * @param object Object to be written

    * @param outOf whether we are writing reusable objects,

    *              or full objects with class name

    */

   private static void writeInternal(

       final DataOutput out, final Object object, final boolean outOf) {

     KRYO_POOL.run(new KryoCallback<Void>() {

       @Override

       public Void execute(Kryo kryo) {

         HadoopKryo hkryo = (HadoopKryo) kryo;

         hkryo.setDataOutput(out);

         if (outOf) {

           hkryo.writeOutOfObject(hkryo.output, object);

         } else {

           hkryo.writeClassAndObject(hkryo.output, object);

         }

         hkryo.output.endChunks();

         hkryo.output.close();

         return null;

       }

     });

   }

   /**

    * Internal read implementation, that reuses HadoopKryo objects

    * from the pool.

    *

    * @param in Input stream

    * @param outObject Object to fill from input (if not null)

    * @param into whether we are reading reusable objects,

    *             or full objects with class name

    * @return Read object (new one, or same passed in if we use reusable)

    * @param <T> Type of the object to read

    */

   @SuppressWarnings("unchecked")

   private static <T> T readInternal(

       final DataInput in, final T outObject, final boolean into) {

     return KRYO_POOL.run(new KryoCallback<T>() {

       @Override

       public T execute(Kryo kryo) {

         HadoopKryo hkryo = (HadoopKryo) kryo;

         hkryo.setDataInput(in);

         T object;

         if (into) {

           hkryo.readIntoObject(hkryo.input, outObject);

           object = outObject;

         } else {

           object = (T) hkryo.readClassAndObject(hkryo.input);

         }

         hkryo.input.nextChunks();

         hkryo.input.close();

         return object;

       }

     });

   }

   /**

    * Reads an object, from input, into a given object,

    * allowing object reuse.

    *

    * @param input Input stream

    * @param object Object to fill from input

    */

   private void readIntoObject(Input input, Object object) {

     Preconditions.checkNotNull(object);

     Class<?> type = object.getClass();

     ReusableFieldSerializer<Object> serializer =

         getOrCreateReusableSerializer(type);

     serializer.setReadIntoObject(object);

     Object result = readObject(input, type, serializer);

     Preconditions.checkState(result == object);

   }

   /**

    * Write an object to output, in a way that can be read

    * using readIntoObject.

    * @param output Output stream

    * @param object Object to be written

    */

   private void writeOutOfObject(Output output, Object object) {

     ReusableFieldSerializer<Object> serializer =

         getOrCreateReusableSerializer(object.getClass());

     writeObject(output, object, serializer);

   }

 }