org.apache.tinkerpop.gremlin.structure.io.IoRegistry Java Examples

/**
 * Reads a list of fully qualified class names from the value of the {@link #TOKEN_IO_REGISTRIES} configuration
 * key. These classes should equate to {@link IoRegistry} implementations that will be assigned to the
 * {@link Mapper.Builder}.  The assumption is that the {@link IoRegistry} either has a static {@code instance()}
 * method or has a zero-arg constructor from which it can be instantiated.
 */
protected void addIoRegistries(final Map<String, Object> config, final Mapper.Builder builder) {
    final List<String> classNameList = getListStringFromConfig(TOKEN_IO_REGISTRIES, config);

    classNameList.stream().forEach(className -> {
        try {
            final Class<?> clazz = Class.forName(className);
            try {
                // try instance() first and then getInstance() which was deprecated in 3.2.4
                final Method instanceMethod = tryInstanceMethod(clazz);
                if (IoRegistry.class.isAssignableFrom(instanceMethod.getReturnType()))
                    builder.addRegistry((IoRegistry) instanceMethod.invoke(null));
                else
                    throw new Exception();
            } catch (Exception methodex) {
                // tried instance() and that failed so try newInstance() no-arg constructor
                builder.addRegistry((IoRegistry) clazz.newInstance());
            }
        } catch (Exception ex) {
            throw new IllegalStateException(ex);
        }
    });
}
 

@Test
public void shouldConfigPoolOnConstructionWithPoolSizeOneAndNoIoRegistry() throws Exception {
    final Configuration conf = new BaseConfiguration();
    final GryoPool pool = GryoPool.build().poolSize(1).ioRegistries(conf.getList(IoRegistry.IO_REGISTRY, Collections.emptyList())).create();
    final GryoReader reader = pool.takeReader();
    final GryoWriter writer = pool.takeWriter();

    pool.offerReader(reader);
    pool.offerWriter(writer);

    for (int ix = 0; ix < 100; ix++) {
        final GryoReader r = pool.takeReader();
        final GryoWriter w = pool.takeWriter();
        assertReaderWriter(w, r, 1, Integer.class);

        // should always return the same original instance
        assertEquals(reader, r);
        assertEquals(writer, w);

        pool.offerReader(r);
        pool.offerWriter(w);
    }
}
 

public synchronized static void initialize(final Configuration configuration) {
    if (!INITIALIZED) {
        INITIALIZED = true;
        GRYO_POOL = GryoPool.build().
                poolSize(configuration.getInt(GryoPool.CONFIG_IO_GRYO_POOL_SIZE, 256)).
                version(GryoVersion.valueOf(configuration.getString(GryoPool.CONFIG_IO_GRYO_VERSION, GryoPool.CONFIG_IO_GRYO_POOL_VERSION_DEFAULT.name()))).
                ioRegistries(configuration.getList(IoRegistry.IO_REGISTRY, Collections.emptyList())).
                initializeMapper(m -> m.registrationRequired(false)).
                create();
    }
}
 

public void checkGraphSONIoRegistryCompliance(final HadoopGraph graph, final Class<? extends GraphComputer> graphComputerClass) throws Exception {
    final File input = TestHelper.generateTempFile(this.getClass(), "graphson-io-registry", ".json");
    graph.configuration().setProperty(Constants.GREMLIN_HADOOP_GRAPH_READER, GraphSONInputFormat.class.getCanonicalName());
    graph.configuration().setProperty(Constants.GREMLIN_HADOOP_GRAPH_WRITER, GraphSONOutputFormat.class.getCanonicalName());
    graph.configuration().setProperty(Constants.GREMLIN_HADOOP_INPUT_LOCATION, Storage.toPath(input));
    graph.configuration().setProperty(IoRegistry.IO_REGISTRY, ToyIoRegistry.class.getCanonicalName());
    final GraphSONRecordWriter writer = new GraphSONRecordWriter(new DataOutputStream(new FileOutputStream(input)), ConfUtil.makeHadoopConfiguration(graph.configuration()));
    validateIoRegistryGraph(graph, graphComputerClass, writer);
    assertTrue(input.delete());
}
 

public void checkGryoV3d0IoRegistryCompliance(final HadoopGraph graph, final Class<? extends GraphComputer> graphComputerClass) throws Exception {
    final File input = TestHelper.generateTempFile(this.getClass(), "gryo-io-registry", ".kryo");
    graph.configuration().setProperty(Constants.GREMLIN_HADOOP_GRAPH_READER, GryoInputFormat.class.getCanonicalName());
    graph.configuration().setProperty(Constants.GREMLIN_HADOOP_GRAPH_WRITER, GryoOutputFormat.class.getCanonicalName());
    graph.configuration().setProperty(Constants.GREMLIN_HADOOP_INPUT_LOCATION, Storage.toPath(input));
    graph.configuration().setProperty(GryoPool.CONFIG_IO_GRYO_VERSION, GryoVersion.V3_0.name());
    graph.configuration().setProperty(IoRegistry.IO_REGISTRY, ToyIoRegistry.class.getCanonicalName());
    final GryoRecordWriter writer = new GryoRecordWriter(new DataOutputStream(new FileOutputStream(input)), ConfUtil.makeHadoopConfiguration(graph.configuration()));
    validateIoRegistryGraph(graph, graphComputerClass, writer);
    assertTrue(input.delete());
}
 

public void checkGryoV1d0IoRegistryCompliance(final HadoopGraph graph, final Class<? extends GraphComputer> graphComputerClass) throws Exception {
    final File input = TestHelper.generateTempFile(this.getClass(), "gryo-io-registry", ".kryo");
    graph.configuration().setProperty(Constants.GREMLIN_HADOOP_GRAPH_READER, GryoInputFormat.class.getCanonicalName());
    graph.configuration().setProperty(Constants.GREMLIN_HADOOP_GRAPH_WRITER, GryoOutputFormat.class.getCanonicalName());
    graph.configuration().setProperty(GryoPool.CONFIG_IO_GRYO_VERSION, GryoVersion.V1_0.name());
    graph.configuration().setProperty(Constants.GREMLIN_HADOOP_INPUT_LOCATION, Storage.toPath(input));
    graph.configuration().setProperty(IoRegistry.IO_REGISTRY, ToyIoRegistry.class.getCanonicalName());
    final GryoRecordWriter writer = new GryoRecordWriter(new DataOutputStream(new FileOutputStream(input)), ConfUtil.makeHadoopConfiguration(graph.configuration()));
    validateIoRegistryGraph(graph, graphComputerClass, writer);
    assertTrue(input.delete());
}
 

public GryoSerializer(final SparkConf sparkConfiguration) {
    final long bufferSizeKb = sparkConfiguration.getSizeAsKb("spark.kryoserializer.buffer", "64k");
    final long maxBufferSizeMb = sparkConfiguration.getSizeAsMb("spark.kryoserializer.buffer.max", "64m");
    this.referenceTracking = sparkConfiguration.getBoolean("spark.kryo.referenceTracking", true);
    this.registrationRequired = sparkConfiguration.getBoolean(Constants.SPARK_KRYO_REGISTRATION_REQUIRED, false);
    if (bufferSizeKb >= ByteUnit.GiB.toKiB(2L)) {
        throw new IllegalArgumentException("spark.kryoserializer.buffer must be less than 2048 mb, got: " + bufferSizeKb + " mb.");
    } else {
        this.bufferSize = (int) ByteUnit.KiB.toBytes(bufferSizeKb);
        if (maxBufferSizeMb >= ByteUnit.GiB.toMiB(2L)) {
            throw new IllegalArgumentException("spark.kryoserializer.buffer.max must be less than 2048 mb, got: " + maxBufferSizeMb + " mb.");
        } else {
            this.maxBufferSize = (int) ByteUnit.MiB.toBytes(maxBufferSizeMb);
            //this.userRegistrator = sparkConfiguration.getOption("spark.kryo.registrator");
        }
    }
    // create a GryoPool and store it in static HadoopPools
    final List<Object> ioRegistries = new ArrayList<>();
    ioRegistries.addAll(makeApacheConfiguration(sparkConfiguration).getList(IoRegistry.IO_REGISTRY, Collections.emptyList()));
    ioRegistries.add(SparkIoRegistry.class.getCanonicalName().replace("." + SparkIoRegistry.class.getSimpleName(), "$" + SparkIoRegistry.class.getSimpleName()));
    HadoopPools.initialize(GryoPool.build().
            version(GryoVersion.valueOf(sparkConfiguration.get(GryoPool.CONFIG_IO_GRYO_VERSION, GryoPool.CONFIG_IO_GRYO_POOL_VERSION_DEFAULT.name()))).
            poolSize(sparkConfiguration.getInt(GryoPool.CONFIG_IO_GRYO_POOL_SIZE, GryoPool.CONFIG_IO_GRYO_POOL_SIZE_DEFAULT)).
            ioRegistries(ioRegistries).
            initializeMapper(builder ->
                    builder.referenceTracking(this.referenceTracking).
                            registrationRequired(this.registrationRequired)).
            create());
}
 

@Test
public void shouldConfigPoolOnConstructionWithMultipleCustomIoRegistries() throws Exception {
    final Configuration conf = new BaseConfiguration();
    ((BaseConfiguration) conf).setListDelimiterHandler(new LegacyListDelimiterHandler(','));
    conf.setProperty(IoRegistry.IO_REGISTRY,
            IoXIoRegistry.InstanceBased.class.getName() + "," + IoYIoRegistry.InstanceBased.class.getName());
    final GryoPool pool = GryoPool.build().ioRegistries(conf.getList(IoRegistry.IO_REGISTRY, Collections.emptyList())).create();
    assertReaderWriter(pool.takeWriter(), pool.takeReader(), new IoX("test"), IoX.class);
    assertReaderWriter(pool.takeWriter(), pool.takeReader(), new IoY(100, 200), IoY.class);
}
 

@Test
public void shouldConfigPoolOnConstructionWithCustomIoRegistryInstance() throws Exception {
    final Configuration conf = new BaseConfiguration();
    conf.setProperty(IoRegistry.IO_REGISTRY, IoXIoRegistry.InstanceBased.class.getName());
    final GryoPool pool = GryoPool.build().ioRegistries(conf.getList(IoRegistry.IO_REGISTRY, Collections.emptyList())).create();
    assertReaderWriter(pool.takeWriter(), pool.takeReader(), new IoX("test"), IoX.class);
}
 

@Test
public void shouldConfigPoolOnConstructionWithCustomIoRegistryConstructor() throws Exception {
    final Configuration conf = new BaseConfiguration();
    conf.setProperty(IoRegistry.IO_REGISTRY, IoXIoRegistry.ConstructorBased.class.getName());
    final GryoPool pool = GryoPool.build().ioRegistries(conf.getList(IoRegistry.IO_REGISTRY, Collections.emptyList())).create();
    assertReaderWriter(pool.takeWriter(), pool.takeReader(), new IoX("test"), IoX.class);
}
 

@Test
public void shouldDoWithReaderWriterMethods() throws Exception {
    final Configuration conf = new BaseConfiguration();
    final GryoPool pool = GryoPool.build().ioRegistries(conf.getList(IoRegistry.IO_REGISTRY, Collections.emptyList())).create();
    try (final ByteArrayOutputStream os = new ByteArrayOutputStream()) {
        pool.doWithWriter(writer -> writer.writeObject(os, 1));
        os.flush();
        try (final ByteArrayInputStream is = new ByteArrayInputStream(os.toByteArray())) {
            assertEquals(1, pool.<Integer>doWithReader(FunctionUtils.wrapFunction(reader -> reader.readObject(is, Integer.class))).intValue());
        }
    }
    assertReaderWriter(pool.takeWriter(), pool.takeReader(), 1, Integer.class);
}
 

protected List<IoRegistry> detectRegistries() {
    final List<Object> k = parameters.get(IO.registry, null);
    return k.stream().map(cn -> {
        try {
            if (cn instanceof IoRegistry)
                return (IoRegistry) cn;
            else {
                final Class<?> clazz = Class.forName(cn.toString());
                return (IoRegistry) clazz.getMethod("instance").invoke(null);
            }
        } catch (Exception ex) {
            throw new IllegalStateException(ex);
        }
    }).collect(Collectors.toList());
}
 

public static List<IoRegistry> createRegistries(final Configuration configuration) {
    if (configuration.containsKey(IoRegistry.IO_REGISTRY)) {
        final Object property = configuration.getProperty(IoRegistry.IO_REGISTRY);
        if (property instanceof IoRegistry)
            return Collections.singletonList((IoRegistry) property);
        else if (property instanceof List)
            return createRegistries((List) property);
        else if (property instanceof String)
            return createRegistries(Arrays.asList(((String) property).split(",")));
        else
            throw new IllegalArgumentException("The provided registry object can not be resolved to an instance: " + property);
    } else
        return Collections.emptyList();
}
 

/**
 * Add {@link CustomTypeSerializer} by way of an {@link IoRegistry}. The registry entries should be bound to
 * {@link GraphBinaryIo}.
 */
public Builder addRegistry(final IoRegistry registry) {
    if (null == registry) throw new IllegalArgumentException("The registry cannot be null");

    final List<Pair<Class, CustomTypeSerializer>> classSerializers = registry.find(GraphBinaryIo.class, CustomTypeSerializer.class);
    for (Pair<Class,CustomTypeSerializer> cs : classSerializers) {
        addCustomType(cs.getValue0(), cs.getValue1());
    }

    return this;
}
 

/**
 * {@inheritDoc}
 */
@Override
public Builder addRegistry(final IoRegistry registry) {
    if (null == registry) throw new IllegalArgumentException("The registry cannot be null");
    this.registries.add(registry);
    return this;
}
 

public static Cluster cluster(String host, int port, IoRegistry registry) {
    // GryoMapper.Builder builder = GryoMapper.build().addRegistry(JanusGraphIoRegistry.getInstance());
    GryoMapper.Builder builder = GryoMapper.build().addRegistry(registry);
    MessageSerializer serializer = new GryoMessageSerializerV3d0(builder);
    return Cluster.build().
            addContactPoint(host).
            port(port).
            serializer(serializer).
            create();
}
 

@Test
public void shouldConfigPoolOnConstructionWithDefaults() throws Exception {
    final Configuration conf = new BaseConfiguration();
    final GryoPool pool = GryoPool.build().ioRegistries(conf.getList(IoRegistry.IO_REGISTRY, Collections.emptyList())).create();
    assertReaderWriter(pool.takeWriter(), pool.takeReader(), 1, Integer.class);
}
 

@Override
public Builder addRegistry(final IoRegistry registry) {
    throw new UnsupportedOperationException("GraphML does not accept custom serializers - it is a format for full Graph serialization");
}
 

@Test(expected = IllegalArgumentException.class)
public void shouldConfigPoolOnConstructionWithoutCustomIoRegistryAndFail() throws Exception {
    final Configuration conf = new BaseConfiguration();
    final GryoPool pool = GryoPool.build().ioRegistries(conf.getList(IoRegistry.IO_REGISTRY, Collections.emptyList())).create();
    assertReaderWriter(pool.takeWriter(), pool.takeReader(), new IoX("test"), IoX.class);
}
 

@Test(expected = IllegalStateException.class)
public void shouldConfigPoolOnConstructionWithoutBadIoRegistryAndFail() throws Exception {
    final Configuration conf = new BaseConfiguration();
    conf.setProperty(IoRegistry.IO_REGISTRY, "some.class.that.does.not.exist");
    GryoPool.build().ioRegistries(conf.getList(IoRegistry.IO_REGISTRY, Collections.emptyList())).create();
}
 

/**
 * {@inheritDoc}
 */
@Override
public Builder addRegistry(final IoRegistry registry) {
    registries.add(registry);
    return this;
}