hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/util/KryoUtil.java

/*
 * 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.hugegraph.util;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.UUID;

import org.apache.hugegraph.backend.BackendException;
import org.apache.tinkerpop.shaded.kryo.Kryo;
import org.apache.tinkerpop.shaded.kryo.Serializer;
import org.apache.tinkerpop.shaded.kryo.io.Input;
import org.apache.tinkerpop.shaded.kryo.io.Output;

public final class KryoUtil {

    private static final ThreadLocal<Kryo> KRYOS = new ThreadLocal<>();

    public static Kryo kryo() {
        Kryo kryo = KRYOS.get();
        if (kryo != null) {
            return kryo;
        }

        kryo = new Kryo();
        registerSerializers(kryo);
        KRYOS.set(kryo);
        return kryo;
    }

    public static byte[] toKryo(Object value) {
        try (ByteArrayOutputStream bos = new ByteArrayOutputStream();
             Output output = new Output(bos, 256)) {
            kryo().writeObject(output, value);
            output.flush();
            return bos.toByteArray();
        } catch (IOException e) {
            throw new BackendException("Failed to serialize: %s", e, value);
        }
    }

    public static <T> T fromKryo(byte[] value, Class<T> clazz) {
        E.checkState(value != null,
                     "Kryo value can't be null for '%s'",
                     clazz.getSimpleName());
        return kryo().readObject(new Input(value), clazz);
    }

    public static byte[] toKryoWithType(Object value) {
        try (ByteArrayOutputStream bos = new ByteArrayOutputStream();
             Output output = new Output(bos, 256)) {
            kryo().writeClassAndObject(output, value);
            output.flush();
            return bos.toByteArray();
        } catch (IOException e) {
            throw new BackendException("Failed to serialize: %s", e, value);
        }
    }

    @SuppressWarnings("unchecked")
    public static <T> T fromKryoWithType(byte[] value) {
        E.checkState(value != null, "Kryo value can't be null for object");
        return (T) kryo().readClassAndObject(new Input(value));
    }

    private static void registerSerializers(Kryo kryo) {
        kryo.addDefaultSerializer(UUID.class, new Serializer<UUID>() {

            @Override
            public UUID read(Kryo kryo, Input input, Class<UUID> c) {
                return new UUID(input.readLong(), input.readLong());
            }

            @Override
            public void write(Kryo kryo, Output output, UUID uuid) {
                output.writeLong(uuid.getMostSignificantBits());
                output.writeLong(uuid.getLeastSignificantBits());
            }
        });
    }
}