Java Code Examples for org.apache.flink.api.common.ExecutionConfig#registerKryoType()

@Test
public void testListStateInstantiation() throws Exception {

	final ExecutionConfig config = new ExecutionConfig();
	config.registerKryoType(Path.class);

	final AtomicReference<Object> descriptorCapture = new AtomicReference<>();

	StreamingRuntimeContext context = createRuntimeContext(descriptorCapture, config);

	ListStateDescriptor<TaskInfo> descr = new ListStateDescriptor<>("name", TaskInfo.class);
	context.getListState(descr);

	ListStateDescriptor<?> descrIntercepted = (ListStateDescriptor<?>) descriptorCapture.get();
	TypeSerializer<?> serializer = descrIntercepted.getSerializer();

	// check that the Path class is really registered, i.e., the execution config was applied
	assertTrue(serializer instanceof ListSerializer);

	TypeSerializer<?> elementSerializer = descrIntercepted.getElementSerializer();
	assertTrue(elementSerializer instanceof KryoSerializer);
	assertTrue(((KryoSerializer<?>) elementSerializer).getKryo().getRegistration(Path.class).getId() > 0);
}
 

@Test
public void testValueStateInstantiation() throws Exception {

	final ExecutionConfig config = new ExecutionConfig();
	config.registerKryoType(Path.class);

	final AtomicReference<Object> descriptorCapture = new AtomicReference<>();

	StreamingRuntimeContext context = new StreamingRuntimeContext(
			createDescriptorCapturingMockOp(descriptorCapture, config),
			createMockEnvironment(),
			Collections.<String, Accumulator<?, ?>>emptyMap());

	ValueStateDescriptor<TaskInfo> descr = new ValueStateDescriptor<>("name", TaskInfo.class);
	context.getState(descr);

	StateDescriptor<?, ?> descrIntercepted = (StateDescriptor<?, ?>) descriptorCapture.get();
	TypeSerializer<?> serializer = descrIntercepted.getSerializer();

	// check that the Path class is really registered, i.e., the execution config was applied
	assertTrue(serializer instanceof KryoSerializer);
	assertTrue(((KryoSerializer<?>) serializer).getKryo().getRegistration(Path.class).getId() > 0);
}
 

@Test
public void testAggregatingStateInstantiation() throws Exception {
	final ExecutionConfig config = new ExecutionConfig();
	config.registerKryoType(Path.class);

	final AtomicReference<Object> descriptorCapture = new AtomicReference<>();

	StreamingRuntimeContext context = createRuntimeContext(descriptorCapture, config);

	@SuppressWarnings("unchecked")
	AggregateFunction<String, TaskInfo, String> aggregate = (AggregateFunction<String, TaskInfo, String>) mock(AggregateFunction.class);

	AggregatingStateDescriptor<String, TaskInfo, String> descr =
			new AggregatingStateDescriptor<>("name", aggregate, TaskInfo.class);

	context.getAggregatingState(descr);

	AggregatingStateDescriptor<?, ?, ?> descrIntercepted = (AggregatingStateDescriptor<?, ?, ?>) descriptorCapture.get();
	TypeSerializer<?> serializer = descrIntercepted.getSerializer();

	// check that the Path class is really registered, i.e., the execution config was applied
	assertTrue(serializer instanceof KryoSerializer);
	assertTrue(((KryoSerializer<?>) serializer).getKryo().getRegistration(Path.class).getId() > 0);
}
 

@Test
public void testInitializeSerializerAfterSerializationWithCustomConfig() throws Exception {
	// guard our test assumptions.
	assertEquals("broken test assumption", -1,
			new KryoSerializer<>(String.class, new ExecutionConfig()).getKryo()
					.getRegistration(File.class).getId());

	final ExecutionConfig config = new ExecutionConfig();
	config.registerKryoType(File.class);

	final TestStateDescriptor<Path> original = new TestStateDescriptor<>("test", Path.class);
	TestStateDescriptor<Path> clone = CommonTestUtils.createCopySerializable(original);

	clone.initializeSerializerUnlessSet(config);

	// serialized one (later initialized) carries the registration
	assertTrue(((KryoSerializer<?>) clone.getSerializer()).getKryo()
			.getRegistration(File.class).getId() > 0);
}
 

@Test
public void testListStateInstantiation() throws Exception {

	final ExecutionConfig config = new ExecutionConfig();
	config.registerKryoType(Path.class);

	final AtomicReference<Object> descriptorCapture = new AtomicReference<>();

	StreamingRuntimeContext context = new StreamingRuntimeContext(
			createDescriptorCapturingMockOp(descriptorCapture, config),
			createMockEnvironment(),
			Collections.<String, Accumulator<?, ?>>emptyMap());

	ListStateDescriptor<TaskInfo> descr = new ListStateDescriptor<>("name", TaskInfo.class);
	context.getListState(descr);

	ListStateDescriptor<?> descrIntercepted = (ListStateDescriptor<?>) descriptorCapture.get();
	TypeSerializer<?> serializer = descrIntercepted.getSerializer();

	// check that the Path class is really registered, i.e., the execution config was applied
	assertTrue(serializer instanceof ListSerializer);

	TypeSerializer<?> elementSerializer = descrIntercepted.getElementSerializer();
	assertTrue(elementSerializer instanceof KryoSerializer);
	assertTrue(((KryoSerializer<?>) elementSerializer).getKryo().getRegistration(Path.class).getId() > 0);
}
 

@Test
public void testMapStateInstantiation() throws Exception {

	final ExecutionConfig config = new ExecutionConfig();
	config.registerKryoType(Path.class);

	final AtomicReference<Object> descriptorCapture = new AtomicReference<>();

	StreamingRuntimeContext context = new StreamingRuntimeContext(
			createDescriptorCapturingMockOp(descriptorCapture, config),
			createMockEnvironment(),
			Collections.<String, Accumulator<?, ?>>emptyMap());

	MapStateDescriptor<String, TaskInfo> descr =
			new MapStateDescriptor<>("name", String.class, TaskInfo.class);

	context.getMapState(descr);

	MapStateDescriptor<?, ?> descrIntercepted = (MapStateDescriptor<?, ?>) descriptorCapture.get();
	TypeSerializer<?> valueSerializer = descrIntercepted.getValueSerializer();

	// check that the Path class is really registered, i.e., the execution config was applied
	assertTrue(valueSerializer instanceof KryoSerializer);
	assertTrue(((KryoSerializer<?>) valueSerializer).getKryo().getRegistration(Path.class).getId() > 0);
}
 

@Test
public void testValueStateInstantiation() throws Exception {

	final ExecutionConfig config = new ExecutionConfig();
	config.registerKryoType(Path.class);

	final AtomicReference<Object> descriptorCapture = new AtomicReference<>();

	StreamingRuntimeContext context = new StreamingRuntimeContext(
			createDescriptorCapturingMockOp(descriptorCapture, config),
			createMockEnvironment(),
			Collections.<String, Accumulator<?, ?>>emptyMap());

	ValueStateDescriptor<TaskInfo> descr = new ValueStateDescriptor<>("name", TaskInfo.class);
	context.getState(descr);

	StateDescriptor<?, ?> descrIntercepted = (StateDescriptor<?, ?>) descriptorCapture.get();
	TypeSerializer<?> serializer = descrIntercepted.getSerializer();

	// check that the Path class is really registered, i.e., the execution config was applied
	assertTrue(serializer instanceof KryoSerializer);
	assertTrue(((KryoSerializer<?>) serializer).getKryo().getRegistration(Path.class).getId() > 0);
}
 

@Test
public void testInitializeSerializerAfterSerializationWithCustomConfig() throws Exception {
	// guard our test assumptions.
	assertEquals("broken test assumption", -1,
			new KryoSerializer<>(String.class, new ExecutionConfig()).getKryo()
					.getRegistration(File.class).getId());

	final ExecutionConfig config = new ExecutionConfig();
	config.registerKryoType(File.class);

	final TestStateDescriptor<Path> original = new TestStateDescriptor<>("test", Path.class);
	TestStateDescriptor<Path> clone = CommonTestUtils.createCopySerializable(original);

	clone.initializeSerializerUnlessSet(config);

	// serialized one (later initialized) carries the registration
	assertTrue(((KryoSerializer<?>) clone.getSerializer()).getKryo()
			.getRegistration(File.class).getId() > 0);
}
 

@Benchmark
@OperationsPerInvocation(value = SerializationFrameworkMiniBenchmarks.RECORDS_PER_INVOCATION)
public void serializerKryo(FlinkEnvironmentContext context) throws Exception {
	StreamExecutionEnvironment env = context.env;
	env.setParallelism(4);
	ExecutionConfig executionConfig = env.getConfig();
	executionConfig.enableForceKryo();
	executionConfig.registerKryoType(MyPojo.class);
	executionConfig.registerKryoType(MyOperation.class);

	env.addSource(new PojoSource(RECORDS_PER_INVOCATION, 10))
			.rebalance()
			.addSink(new DiscardingSink<>());

	env.execute();
}
 

@Override
public TypeSerializer<Animal> createPriorSerializer() {
	ExecutionConfig executionConfig = new ExecutionConfig();
	executionConfig.registerKryoType(Dog.class);
	executionConfig.registerKryoType(Cat.class);
	executionConfig.registerKryoType(Parrot.class);

	return new KryoSerializer<>(Animal.class, executionConfig);
}
 

@Override
public TypeSerializer<Animal> createPriorSerializer() {
	ExecutionConfig executionConfig = new ExecutionConfig();
	executionConfig.registerTypeWithKryoSerializer(
			Dog.class,
			KryoPojosForMigrationTests.DogKryoSerializer.class);
	executionConfig.registerKryoType(Cat.class);
	executionConfig.registerTypeWithKryoSerializer(
			Parrot.class,
			KryoPojosForMigrationTests.ParrotKryoSerializer.class);

	return new KryoSerializer<>(Animal.class, executionConfig);
}
 

@Test
public void testAggregatingStateInstantiation() throws Exception {

	final ExecutionConfig config = new ExecutionConfig();
	config.registerKryoType(Path.class);

	final AtomicReference<Object> descriptorCapture = new AtomicReference<>();

	StreamingRuntimeContext context = new StreamingRuntimeContext(
			createDescriptorCapturingMockOp(descriptorCapture, config),
			createMockEnvironment(),
			Collections.<String, Accumulator<?, ?>>emptyMap());

	@SuppressWarnings("unchecked")
	AggregateFunction<String, TaskInfo, String> aggregate = (AggregateFunction<String, TaskInfo, String>) mock(AggregateFunction.class);

	AggregatingStateDescriptor<String, TaskInfo, String> descr =
			new AggregatingStateDescriptor<>("name", aggregate, TaskInfo.class);

	context.getAggregatingState(descr);

	AggregatingStateDescriptor<?, ?, ?> descrIntercepted = (AggregatingStateDescriptor<?, ?, ?>) descriptorCapture.get();
	TypeSerializer<?> serializer = descrIntercepted.getSerializer();

	// check that the Path class is really registered, i.e., the execution config was applied
	assertTrue(serializer instanceof KryoSerializer);
	assertTrue(((KryoSerializer<?>) serializer).getKryo().getRegistration(Path.class).getId() > 0);
}
 

@Test
public void testReducingStateInstantiation() throws Exception {

	final ExecutionConfig config = new ExecutionConfig();
	config.registerKryoType(Path.class);

	final AtomicReference<Object> descriptorCapture = new AtomicReference<>();

	StreamingRuntimeContext context = new StreamingRuntimeContext(
			createDescriptorCapturingMockOp(descriptorCapture, config),
			createMockEnvironment(),
			Collections.<String, Accumulator<?, ?>>emptyMap());

	@SuppressWarnings("unchecked")
	ReduceFunction<TaskInfo> reducer = (ReduceFunction<TaskInfo>) mock(ReduceFunction.class);

	ReducingStateDescriptor<TaskInfo> descr =
			new ReducingStateDescriptor<>("name", reducer, TaskInfo.class);

	context.getReducingState(descr);

	StateDescriptor<?, ?> descrIntercepted = (StateDescriptor<?, ?>) descriptorCapture.get();
	TypeSerializer<?> serializer = descrIntercepted.getSerializer();

	// check that the Path class is really registered, i.e., the execution config was applied
	assertTrue(serializer instanceof KryoSerializer);
	assertTrue(((KryoSerializer<?>) serializer).getKryo().getRegistration(Path.class).getId() > 0);
}
 

@Override
public TypeSerializer<Animal> createUpgradedSerializer() {
	ExecutionConfig executionConfig = new ExecutionConfig();
	executionConfig.registerKryoType(DummyClassOne.class);
	executionConfig.registerTypeWithKryoSerializer(
			Dog.class,
			KryoPojosForMigrationTests.DogV2KryoSerializer.class);
	executionConfig.registerKryoType(DummyClassTwo.class);
	executionConfig.registerKryoType(Cat.class);
	executionConfig.registerTypeWithKryoSerializer(
			Parrot.class,
			KryoPojosForMigrationTests.ParrotKryoSerializer.class);

	return new KryoSerializer<>(Animal.class, executionConfig);
}
 

@Test
public void testReducingStateInstantiation() throws Exception {

	final ExecutionConfig config = new ExecutionConfig();
	config.registerKryoType(Path.class);

	final AtomicReference<Object> descriptorCapture = new AtomicReference<>();

	StreamingRuntimeContext context = new StreamingRuntimeContext(
			createDescriptorCapturingMockOp(descriptorCapture, config),
			createMockEnvironment(),
			Collections.<String, Accumulator<?, ?>>emptyMap());

	@SuppressWarnings("unchecked")
	ReduceFunction<TaskInfo> reducer = (ReduceFunction<TaskInfo>) mock(ReduceFunction.class);

	ReducingStateDescriptor<TaskInfo> descr =
			new ReducingStateDescriptor<>("name", reducer, TaskInfo.class);

	context.getReducingState(descr);

	StateDescriptor<?, ?> descrIntercepted = (StateDescriptor<?, ?>) descriptorCapture.get();
	TypeSerializer<?> serializer = descrIntercepted.getSerializer();

	// check that the Path class is really registered, i.e., the execution config was applied
	assertTrue(serializer instanceof KryoSerializer);
	assertTrue(((KryoSerializer<?>) serializer).getKryo().getRegistration(Path.class).getId() > 0);
}
 

private static ExecutionConfig registerClassThatIsNotInClassPath(ClassLoader tempClassLoader) {
	Object objectForClassNotInClassPath =
		CommonTestUtils.createObjectForClassNotInClassPath(tempClassLoader);

	ExecutionConfig conf = new ExecutionConfig();
	conf.registerKryoType(objectForClassNotInClassPath.getClass());
	return conf;
}
 

/**
 * Tests that the kryo output buffer is cleared in case of an exception. Flink uses the
 * EOFException to signal that a buffer is full. In such a case, the record which was tried
 * to be written will be rewritten. Therefore, eventually buffered data of this record has
 * to be cleared.
 */
@Test
public void testOutputBufferedBeingClearedInCaseOfException() throws Exception {
	ExecutionConfig executionConfig = new ExecutionConfig();
	executionConfig.registerTypeWithKryoSerializer(TestRecord.class, new TestRecordSerializer());
	executionConfig.registerKryoType(TestRecord.class);

	KryoSerializer<TestRecord> kryoSerializer = new KryoSerializer<TestRecord>(
		TestRecord.class,
		executionConfig);

	int size = 94;
	int bufferSize = 150;

	TestRecord testRecord = new TestRecord(size);

	TestDataOutputView target = new TestDataOutputView(bufferSize);

	kryoSerializer.serialize(testRecord, target);

	try {
		kryoSerializer.serialize(testRecord, target);
		Assert.fail("Expected an EOFException.");
	} catch(EOFException eofException) {
		// expected exception
		// now the Kryo Output should have been cleared
	}

	TestRecord actualRecord = kryoSerializer.deserialize(
			new DataInputViewStreamWrapper(new ByteArrayInputStream(target.getBuffer())));

	Assert.assertEquals(testRecord, actualRecord);

	target.clear();

	// if the kryo output has been cleared then we can serialize our test record into the target
	// because the target buffer 150 bytes can host one TestRecord (total serialization size 100)
	kryoSerializer.serialize(testRecord, target);

	byte[] buffer = target.getBuffer();
	int counter = 0;

	for (int i = 0; i < buffer.length; i++) {
		if(buffer[i] == 42) {
			counter++;
		}
	}

	Assert.assertEquals(size, counter);
}
 

@Test
public void testReducingStateInstantiation() throws Exception {

	final ExecutionConfig config = new ExecutionConfig();
	config.registerKryoType(Path.class);

	final AtomicReference<Object> descriptorCapture = new AtomicReference<>();

	StreamingRuntimeContext context = createRuntimeContext(descriptorCapture, config);

	@SuppressWarnings("unchecked")
	ReduceFunction<TaskInfo> reducer = (ReduceFunction<TaskInfo>) mock(ReduceFunction.class);

	ReducingStateDescriptor<TaskInfo> descr =
			new ReducingStateDescriptor<>("name", reducer, TaskInfo.class);

	context.getReducingState(descr);

	StateDescriptor<?, ?> descrIntercepted = (StateDescriptor<?, ?>) descriptorCapture.get();
	TypeSerializer<?> serializer = descrIntercepted.getSerializer();

	// check that the Path class is really registered, i.e., the execution config was applied
	assertTrue(serializer instanceof KryoSerializer);
	assertTrue(((KryoSerializer<?>) serializer).getKryo().getRegistration(Path.class).getId() > 0);
}
 

@Test
public void testMapStateInstantiation() throws Exception {

	final ExecutionConfig config = new ExecutionConfig();
	config.registerKryoType(Path.class);

	final AtomicReference<Object> descriptorCapture = new AtomicReference<>();

	StreamingRuntimeContext context = createRuntimeContext(descriptorCapture, config);

	MapStateDescriptor<String, TaskInfo> descr =
			new MapStateDescriptor<>("name", String.class, TaskInfo.class);

	context.getMapState(descr);

	MapStateDescriptor<?, ?> descrIntercepted = (MapStateDescriptor<?, ?>) descriptorCapture.get();
	TypeSerializer<?> valueSerializer = descrIntercepted.getValueSerializer();

	// check that the Path class is really registered, i.e., the execution config was applied
	assertTrue(valueSerializer instanceof KryoSerializer);
	assertTrue(((KryoSerializer<?>) valueSerializer).getKryo().getRegistration(Path.class).getId() > 0);
}
 

/**
 * Tests that the kryo output buffer is cleared in case of an exception. Flink uses the
 * EOFException to signal that a buffer is full. In such a case, the record which was tried
 * to be written will be rewritten. Therefore, eventually buffered data of this record has
 * to be cleared.
 */
@Test
public void testOutputBufferedBeingClearedInCaseOfException() throws Exception {
	ExecutionConfig executionConfig = new ExecutionConfig();
	executionConfig.registerTypeWithKryoSerializer(TestRecord.class, new TestRecordSerializer());
	executionConfig.registerKryoType(TestRecord.class);

	KryoSerializer<TestRecord> kryoSerializer = new KryoSerializer<TestRecord>(
		TestRecord.class,
		executionConfig);

	int size = 94;
	int bufferSize = 150;

	TestRecord testRecord = new TestRecord(size);

	TestDataOutputView target = new TestDataOutputView(bufferSize);

	kryoSerializer.serialize(testRecord, target);

	try {
		kryoSerializer.serialize(testRecord, target);
		Assert.fail("Expected an EOFException.");
	} catch(EOFException eofException) {
		// expected exception
		// now the Kryo Output should have been cleared
	}

	TestRecord actualRecord = kryoSerializer.deserialize(
			new DataInputViewStreamWrapper(new ByteArrayInputStream(target.getBuffer())));

	Assert.assertEquals(testRecord, actualRecord);

	target.clear();

	// if the kryo output has been cleared then we can serialize our test record into the target
	// because the target buffer 150 bytes can host one TestRecord (total serialization size 100)
	kryoSerializer.serialize(testRecord, target);

	byte[] buffer = target.getBuffer();
	int counter = 0;

	for (int i = 0; i < buffer.length; i++) {
		if(buffer[i] == 42) {
			counter++;
		}
	}

	Assert.assertEquals(size, counter);
}