org.apache.tinkerpop.shaded.kryo.Kryo Java Examples

@Test
public void shouldRegisterMultipleIoRegistryToSerialize() throws Exception {
    final GryoMapper mapper = builder.get().addRegistry(IoXIoRegistry.InstanceBased.instance())
            .addRegistry(IoYIoRegistry.InstanceBased.instance()).create();
    final Kryo kryo = mapper.createMapper();
    try (final OutputStream stream = new ByteArrayOutputStream()) {
        final Output out = new Output(stream);
        final IoX x = new IoX("x");
        final IoY y = new IoY(100, 200);
        kryo.writeClassAndObject(out, x);
        kryo.writeClassAndObject(out, y);

        try (final InputStream inputStream = new ByteArrayInputStream(out.toBytes())) {
            final Input input = new Input(inputStream);
            final IoX readX = (IoX) kryo.readClassAndObject(input);
            final IoY readY = (IoY) kryo.readClassAndObject(input);
            assertEquals(x, readX);
            assertEquals(y, readY);
        }
    }
}
 

@Test
public void shouldSerializeWithCustomClassResolverToHashMap() throws Exception {
    final Supplier<ClassResolver> classResolver = new CustomClassResolverSupplier();
    final GryoMapper mapper = builder.get().classResolver(classResolver).create();
    final Kryo kryo = mapper.createMapper();
    try (final OutputStream stream = new ByteArrayOutputStream()) {
        final Output out = new Output(stream);
        final IoY y = new IoY(100, 200);

        kryo.writeClassAndObject(out, y);

        final GryoMapper mapperWithoutKnowledgeOfIoy = builder.get().create();
        final Kryo kryoWithoutKnowledgeOfIox = mapperWithoutKnowledgeOfIoy.createMapper();
        try (final InputStream inputStream = new ByteArrayInputStream(out.toBytes())) {
            final Input input = new Input(inputStream);
            final Map readY = (HashMap) kryoWithoutKnowledgeOfIox.readClassAndObject(input);
            assertEquals("100-200", readY.get("y"));
        }
    }
}
 

@Override
public void write(final Kryo kryo, final Output output, final Color color) {
    final TinkerGraph graph = TinkerGraph.open();
    final Vertex v = graph.addVertex(T.id, 1, T.label, "color", "name", color.toString());
    final Vertex vRed = graph.addVertex(T.id, 2, T.label, "primary", "name", "red");
    final Vertex vGreen = graph.addVertex(T.id, 3, T.label, "primary", "name", "green");
    final Vertex vBlue = graph.addVertex(T.id, 4, T.label, "primary", "name", "blue");

    v.addEdge("hasComponent", vRed, "amount", color.getRed());
    v.addEdge("hasComponent", vGreen, "amount", color.getGreen());
    v.addEdge("hasComponent", vBlue, "amount", color.getBlue());

    // make some junk so the graph is kinda big
    generate(graph);

    try (final ByteArrayOutputStream stream = new ByteArrayOutputStream()) {
        GryoWriter.build().mapper(() -> kryo).create().writeGraph(stream, graph);
        final byte[] bytes = stream.toByteArray();
        output.writeInt(bytes.length);
        output.write(bytes);
    } catch (Exception ex) {
        ex.printStackTrace();
    }
}
 

@Test
public void shouldSerializeDeserialize() throws Exception {
    final GryoMapper mapper = builder.get().create();
    final Kryo kryo = mapper.createMapper();
    try (final OutputStream stream = new ByteArrayOutputStream()) {
        final Output out = new Output(stream);

        final Map<String,Object> props = new HashMap<>();
        final List<Map<String, Object>> propertyNames = new ArrayList<>(1);
        final Map<String,Object> propertyName = new HashMap<>();
        propertyName.put(GraphSONTokens.ID, "x");
        propertyName.put(GraphSONTokens.KEY, "x");
        propertyName.put(GraphSONTokens.VALUE, "no-way-this-will-ever-work");
        propertyNames.add(propertyName);
        props.put("x", propertyNames);
        final DetachedVertex v = new DetachedVertex(100, Vertex.DEFAULT_LABEL, props);

        kryo.writeClassAndObject(out, v);

        try (final InputStream inputStream = new ByteArrayInputStream(out.toBytes())) {
            final Input input = new Input(inputStream);
            final DetachedVertex readX = (DetachedVertex) kryo.readClassAndObject(input);
            assertEquals("no-way-this-will-ever-work", readX.value("x"));
        }
    }
}
 

@Override
public RequestMessage deserializeRequest(final ByteBuf msg) throws SerializationException {
    try {
        final Kryo kryo = kryoThreadLocal.get();
        final byte[] payload = new byte[msg.readableBytes()];
        msg.readBytes(payload);
        try (final Input input = new Input(payload)) {
            // by the time the message gets here, the mime length/type have been already read, so this part just
            // needs to process the payload.
            final UUID id = kryo.readObject(input, UUID.class);
            final String processor = input.readString();
            final String op = input.readString();

            final RequestMessage.Builder builder = RequestMessage.build(op)
                    .overrideRequestId(id)
                    .processor(processor);

            final Map<String, Object> args = kryo.readObject(input, HashMap.class);
            args.forEach(builder::addArg);
            return builder.create();
        }
    } catch (Exception ex) {
        logger.warn(String.format("Request [%s] could not be deserialized by %s.", msg, AbstractGryoMessageSerializerV1d0.class.getName()), ex);
        throw new SerializationException(ex);
    }
}
 

@Override
public RequestMessage deserializeRequest(final ByteBuf msg) throws SerializationException {
    try {
        final Kryo kryo = kryoThreadLocal.get();
        final byte[] payload = new byte[msg.readableBytes()];
        msg.readBytes(payload);
        try (final Input input = new Input(payload)) {
            // by the time the message gets here, the mime length/type have been already read, so this part just
            // needs to process the payload.
            return kryo.readObject(input, RequestMessage.class);
        }
    } catch (Exception ex) {
        logger.warn(String.format("Request [%s] could not be deserialized by %s.", msg, AbstractGryoMessageSerializerV3d0.class.getName()), ex);
        throw new SerializationException(ex);
    }
}
 

@Override
public void write(final Kryo kryo, final Output output, final IoY ioy) {
    final Map<String, Object> map = new HashMap<>();
    map.put("y", ioy.toString());
    try (final OutputStream stream = new ByteArrayOutputStream()) {
        final Output detachedOutput = new Output(stream);
        kryo.writeObject(detachedOutput, map);

        // have to remove the first byte because it marks a reference we don't want to have as this
        // serializer is trying to "act" like a Map
        final byte[] b = detachedOutput.toBytes();
        output.write(b, 1, b.length - 1);
    } catch (Exception ex) {
        ex.printStackTrace();
    }
}
 

@Override
public Kryo registerWith(final Kryo kryo) {
    if (null != functionOfShadedKryo)
        kryo.register(clazz, functionOfShadedKryo.apply(kryo), id);
    else if (null != shadedSerializer)
        kryo.register(clazz, shadedSerializer, id);
    else if (null != serializerShim)
        kryo.register(clazz, new ShadedSerializerAdapter<>(serializerShim), id);
    else {
        kryo.register(clazz, kryo.getDefaultSerializer(clazz), id);
        // Suprisingly, the preceding call is not equivalent to
        //   kryo.register(clazz, id);
    }

    return kryo;
}
 

@Override
public void read(Kryo kryo, Input input) {
    this.schemaTable = SchemaTable.of(input.readString(), input.readString());
    String s = input.readString();
    if (s.equals("s")) {
        //sequence
        this.id = ID.from(input.readLong());
    } else {
        int size = input.readInt();
        List<Comparable> identifiers = new ArrayList<>();
        for (int i = 0; i < size; i++) {
            String identifier = input.readString();
            identifiers.add(identifier);
        }
        this.id = ID.from(identifiers);
    }
}
 

/**
 * Creates a {@code GryoMapper}.
 */
public GryoMapper create() {
    // consult the registry if provided and inject registry entries as custom classes.
    registries.forEach(registry -> {
        final List<Pair<Class, Object>> serializers = registry.find(GryoIo.class);
        serializers.forEach(p -> {
            if (null == p.getValue1())
                addCustom(p.getValue0());
            else if (p.getValue1() instanceof SerializerShim)
                addCustom(p.getValue0(), new ShadedSerializerAdapter((SerializerShim) p.getValue1()));
            else if (p.getValue1() instanceof Serializer)
                addCustom(p.getValue0(), (Serializer) p.getValue1());
            else if (p.getValue1() instanceof Function)
                addCustom(p.getValue0(), (Function<Kryo, Serializer>) p.getValue1());
            else
                throw new IllegalStateException(String.format(
                        "Unexpected value provided by %s for serializable class %s - expected a parameter in [null, %s implementation or Function<%s, %s>], but received %s",
                        registry.getClass().getSimpleName(), p.getValue0().getClass().getCanonicalName(),
                        Serializer.class.getName(), Kryo.class.getSimpleName(),
                        Serializer.class.getSimpleName(), p.getValue1()));
        });
    });

    return new GryoMapper(this);
}
 

@Override
public void write(Kryo kryo, Output output, Optional<?> optional) {
    if (optional.isPresent()) {
        kryo.writeClassAndObject(output, optional.get());
    } else {
        kryo.writeObject(output, null);
    }
}
 

@Override
public void write(Kryo kryo, Output output, Id id) {
    output.writeByte(id.type().ordinal());
    byte[] idBytes = id.asBytes();
    output.write(idBytes.length);
    output.writeBytes(id.asBytes());
}
 

@Override
public Id read(Kryo kryo, Input input, Class<Id> clazz) {
    int type = input.readByteUnsigned();
    int length = input.read();
    byte[] idBytes = input.readBytes(length);
    return IdGenerator.of(idBytes, IdType.values()[type]);
}
 

private static void writeEntry(Kryo kryo,
                               Output output,
                               Map<HugeKeys, Object> schema) {
    /* Write columns size and data */
    output.writeInt(schema.keySet().size());
    for (Map.Entry<HugeKeys, Object> entry : schema.entrySet()) {
        kryo.writeObject(output, entry.getKey());
        kryo.writeClassAndObject(output, entry.getValue());
    }
}
 

@SuppressWarnings("unused")
private static Map<HugeKeys, Object> readEntry(Kryo kryo, Input input) {
    int columnSize = input.readInt();
    Map<HugeKeys, Object> map = new LinkedHashMap<>();
    for (int i = 0; i < columnSize; i++) {
        HugeKeys key = kryo.readObject(input, HugeKeys.class);
        Object val = kryo.readClassAndObject(input);
        map.put(key, val);
    }
    return map;
}
 

@Override
public void write(final Kryo kryo, final Output output, final TinkerGraph graph) {
    try (final ByteArrayOutputStream stream = new ByteArrayOutputStream()) {
        GryoWriter.build().mapper(() -> kryo).create().writeGraph(stream, graph);
        final byte[] bytes = stream.toByteArray();
        output.writeInt(bytes.length);
        output.write(bytes);
    } catch (Exception io) {
        throw new RuntimeException(io);
    }
}
 

@Override
public TinkerGraph read(final Kryo kryo, final Input input, final Class<TinkerGraph> tinkerGraphClass) {
    final Configuration conf = new BaseConfiguration();
    conf.setProperty("gremlin.tinkergraph.defaultVertexPropertyCardinality", "list");
    final TinkerGraph graph = TinkerGraph.open(conf);
    final int len = input.readInt();
    final byte[] bytes = input.readBytes(len);
    try (final ByteArrayInputStream stream = new ByteArrayInputStream(bytes)) {
        GryoReader.build().mapper(() -> kryo).create().readGraph(stream, graph);
    } catch (Exception io) {
        throw new RuntimeException(io);
    }

    return graph;
}
 

@Override
public void write(final Kryo kryo, final Output output, final TinkerGraph graph) {
    try (final ByteArrayOutputStream stream = new ByteArrayOutputStream()) {
        GryoWriter.build().mapper(() -> kryo).create().writeGraph(stream, graph);
        final byte[] bytes = stream.toByteArray();
        output.writeInt(bytes.length);
        output.write(bytes);
    } catch (Exception io) {
        throw new RuntimeException(io);
    }
}
 

@Override
public void write(final Kryo kryo, final Output output, final TinkerGraph graph) {
    try (final ByteArrayOutputStream stream = new ByteArrayOutputStream()) {
        GryoWriter.build().mapper(() -> kryo).create().writeGraph(stream, graph);
        final byte[] bytes = stream.toByteArray();
        output.writeInt(bytes.length);
        output.write(bytes);
    } catch (Exception io) {
        throw new RuntimeException(io);
    }
}
 

@Test
public void shouldSerializeWithColorClassResolverToTinkerGraph() throws Exception {
    final Map<String,Color> colors = new HashMap<>();
    colors.put("red", Color.RED);
    colors.put("green", Color.GREEN);

    final ArrayList<Color> colorList = new ArrayList<>(Arrays.asList(Color.RED, Color.GREEN));

    final Supplier<ClassResolver> classResolver = new CustomClassResolverSupplier();
    final GryoMapper mapper = GryoMapper.build().version(GryoVersion.V3_0).addRegistry(TinkerIoRegistryV3d0.instance()).classResolver(classResolver).create();
    final Kryo kryo = mapper.createMapper();
    try (final ByteArrayOutputStream stream = new ByteArrayOutputStream()) {
        final Output out = new Output(stream);

        kryo.writeObject(out, colorList);
        out.flush();
        final byte[] b = stream.toByteArray();

        try (final InputStream inputStream = new ByteArrayInputStream(b)) {
            final Input input = new Input(inputStream);
            final List m = kryo.readObject(input, ArrayList.class);
            final TinkerGraph readX = (TinkerGraph) m.get(0);
            assertEquals(104, IteratorUtils.count(readX.vertices()));
            assertEquals(102, IteratorUtils.count(readX.edges()));
        }
    }
}
 

public <T> T serializeDeserialize(final Object o, final Class<T> clazz) throws Exception {
    final Kryo kryo = builder.get().create().createMapper();
    try (final ByteArrayOutputStream stream = new ByteArrayOutputStream()) {
        final Output out = new Output(stream);
        kryo.writeObject(out, o);
        out.flush();

        try (final InputStream inputStream = new ByteArrayInputStream(stream.toByteArray())) {
            final Input input = new Input(inputStream);
            return kryo.readObject(input, clazz);
        }
    }
}
 

@Override
public ByteBuf serializeResponseAsBinary(final ResponseMessage responseMessage, final ByteBufAllocator allocator) throws SerializationException {
    ByteBuf encodedMessage = null;
    try {
        final Kryo kryo = kryoThreadLocal.get();
        try (final ByteArrayOutputStream baos = new ByteArrayOutputStream()) {
            final Output output = new Output(baos, bufferSize);

            // request id - if present
            kryo.writeObjectOrNull(output, responseMessage.getRequestId() != null ? responseMessage.getRequestId() : null, UUID.class);

            // status
            output.writeShort(responseMessage.getStatus().getCode().getValue());
            output.writeString(responseMessage.getStatus().getMessage());
            kryo.writeClassAndObject(output, responseMessage.getStatus().getAttributes());

            // result
            kryo.writeClassAndObject(output, serializeToString ? serializeResultToString(responseMessage) : responseMessage.getResult().getData());
            kryo.writeClassAndObject(output, responseMessage.getResult().getMeta());

            final long size = output.total();
            if (size > Integer.MAX_VALUE)
                throw new SerializationException(String.format("Message size of %s exceeds allocatable space", size));

            output.flush();
            encodedMessage = allocator.buffer((int) size);
            encodedMessage.writeBytes(baos.toByteArray());
        }

        return encodedMessage;
    } catch (Exception ex) {
        if (encodedMessage != null) ReferenceCountUtil.release(encodedMessage);

        logger.warn(String.format("Response [%s] could not be serialized by %s.", responseMessage, AbstractGryoMessageSerializerV1d0.class.getName()), ex);
        throw new SerializationException(ex);
    }
}
 

@Override
public WrappedArray<T> read(final Kryo kryo, final Input input, final Class<WrappedArray<T>> aClass) {
    final int size = input.readVarInt(true);
    final Object[] array = new Object[size];
    for (int i = 0; i < size; i++) {
        array[i] = kryo.readClassAndObject(input);
    }
    return new WrappedArray.ofRef<>((T[]) array);
}
 

@Override
public void write(final Kryo kryo, final Output output, final WrappedArray<T> iterable) {
    output.writeVarInt(iterable.size(), true);
    JavaConversions.asJavaCollection(iterable).forEach(t -> {
        kryo.writeClassAndObject(output, t);
    });
}
 

@Test
public void shouldExpectReadFailureAsIoRegistryOrderIsNotRespected() throws Exception {
    final GryoMapper mapperWrite = builder.get().addRegistry(IoXIoRegistry.InstanceBased.instance())
            .addRegistry(IoYIoRegistry.InstanceBased.instance()).create();

    final GryoMapper mapperRead = GryoMapper.build()
            .addRegistry(IoYIoRegistry.InstanceBased.instance())
            .addRegistry(IoXIoRegistry.InstanceBased.instance()).create();

    final Kryo kryoWriter = mapperWrite.createMapper();
    final Kryo kryoReader = mapperRead.createMapper();
    try (final OutputStream stream = new ByteArrayOutputStream()) {
        final Output out = new Output(stream);
        final IoX x = new IoX("x");
        final IoY y = new IoY(100, 200);
        kryoWriter.writeClassAndObject(out, x);
        kryoWriter.writeClassAndObject(out, y);

        try (final InputStream inputStream = new ByteArrayInputStream(out.toBytes())) {
            final Input input = new Input(inputStream);

            // kryo will read a IoY instance as we've reversed the registries.  it is neither an X or a Y
            // so assert that both are incorrect
            final IoY readY = (IoY) kryoReader.readClassAndObject(input);
            assertNotEquals(y, readY);
            assertNotEquals(x, readY);
        }
    }
}
 

@Override
public ResponseMessage deserializeResponse(final ByteBuf msg) throws SerializationException {
    try {
        final Kryo kryo = kryoThreadLocal.get();
        final byte[] payload = new byte[msg.capacity()];
        msg.readBytes(payload);
        try (final Input input = new Input(payload)) {
            return kryo.readObject(input, ResponseMessage.class);
        }
    } catch (Exception ex) {
        logger.warn(String.format("Response [%s] could not be deserialized by %s.", msg, AbstractGryoMessageSerializerV3d0.class.getName()), ex);
        throw new SerializationException(ex);
    }
}
 

@Override
public ByteBuf serializeResponseAsBinary(final ResponseMessage responseMessage, final ByteBufAllocator allocator) throws SerializationException {
    ByteBuf encodedMessage = null;
    try {
        final Kryo kryo = kryoThreadLocal.get();
        try (final ByteArrayOutputStream baos = new ByteArrayOutputStream()) {
            final Output output = new Output(baos, bufferSize);
            final ResponseMessage msgToWrite = !serializeToString ? responseMessage :
                    ResponseMessage.build(responseMessage.getRequestId())
                        .code(responseMessage.getStatus().getCode())
                        .statusAttributes(responseMessage.getStatus().getAttributes())
                        .responseMetaData(responseMessage.getResult().getMeta())
                        .result(serializeResultToString(responseMessage))
                        .statusMessage(responseMessage.getStatus().getMessage()).create();
            kryo.writeObject(output, msgToWrite);

            final long size = output.total();
            if (size > Integer.MAX_VALUE)
                throw new SerializationException(String.format("Message size of %s exceeds allocatable space", size));

            output.flush();
            encodedMessage = allocator.buffer((int) size);
            encodedMessage.writeBytes(baos.toByteArray());
        }

        return encodedMessage;
    } catch (Exception ex) {
        if (encodedMessage != null) ReferenceCountUtil.release(encodedMessage);

        logger.warn(String.format("Response [%s] could not be serialized by %s.", responseMessage, AbstractGryoMessageSerializerV3d0.class.getName()), ex);
        throw new SerializationException(ex);
    }
}
 

@Override
public void write(final Kryo kryo, final Output output, final TinkerGraph graph) {
    try (final ByteArrayOutputStream stream = new ByteArrayOutputStream()) {
        GryoWriter.build().mapper(() -> kryo).create().writeGraph(stream, graph);
        final byte[] bytes = stream.toByteArray();
        output.writeInt(bytes.length);
        output.write(bytes);
    } catch (Exception io) {
        throw new RuntimeException(io);
    }
}
 

@Override
public void write(Kryo kryo, Output output) {
    output.writeString(this.getSchemaTable().getSchema());
    output.writeString(this.getSchemaTable().getTable());
    if (hasSequenceId()) {
        output.writeString("s");
        output.writeLong(this.getID().sequenceId);
    } else {
        output.writeString("i");
        output.writeInt(getIdentifiers().size());
        for (Comparable identifier : getIdentifiers()) {
            output.writeString((CharSequence) identifier);
        }
    }
}
 

@Override
public T read(final Kryo kryo, final Input input, final Class<T> aClass) {
    // Same caching opportunity as in write(...)
    final ShadedKryoAdapter shadedKryoAdapter = new ShadedKryoAdapter(kryo);
    final ShadedInputAdapter shadedInputAdapter = new ShadedInputAdapter(input);
    return serializer.read(shadedKryoAdapter, shadedInputAdapter, aClass);
}