zentol commented on a change in pull request #16064: URL: https://github.com/apache/flink/pull/16064#discussion_r644551809
########## File path: flink-test-utils-parent/flink-test-utils-junit/src/main/java/org/apache/flink/testutils/junit/SharedObjects.java ########## @@ -0,0 +1,156 @@ +/* + * 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.flink.testutils.junit; + +import org.junit.rules.ExternalResource; + +import java.io.ObjectStreamException; +import java.io.Serializable; +import java.util.Map; +import java.util.Objects; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicInteger; + +/** + * This rule allows objects to be used both in the main test case as well as in UDFs by using + * serializable {@link SharedReference}s. Usage: + * + * <pre><code> + * {@literal @Rule} + * public final SharedObjects sharedObjects = SharedObjects.create(); + * + * {@literal @Test} + * public void test() throws Exception { + * StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(); + * SerialRef<Queue<Long>> list = sharedObjects.add(new ConcurrentLinkedQueue<>()); + * int n = 10000; + * env.setParallelism(100); + * env.fromSequence(0, n).map(i -> list.get().add(i)); + * env.execute(); + * assertEquals(n + 1, list.get().size()); + * assertEquals( + * LongStream.rangeClosed(0, n).boxed().collect(Collectors.toList()), + * list.get().stream().sorted().collect(Collectors.toList())); + * } + * </code></pre> + * + * <p>Note that since the shared objects are accessed through multiple threads, they need to be thread-safe or accessed in a thread-safe manner. + */ +public class SharedObjects extends ExternalResource implements Serializable { + /** + * Instance-cache used to make a SharedObjects accessible for multiple threads. + */ + private static final Map<Integer, SharedObjects> INSTANCES = new ConcurrentHashMap<>(); + private static final AtomicInteger LAST_ID = new AtomicInteger(); + /** + * Identifier of the SharedObjects used to retrieve the original instance during deserialization. + */ + private final int id; + /** + * All registered objects for the current test case. The objects are purged upon completion. + */ + private final transient Map<SharedReference<?>, Object> objects = new ConcurrentHashMap<>(); + + private SharedObjects(int id) { + this.id = id; + } + + /** + * Creates a new instance. Usually that should be done inside a JUnit test class as an + * instance-field annotated with {@link org.junit.Rule}. + */ + public static SharedObjects create() { + return new SharedObjects(LAST_ID.getAndIncrement()); + } + + /** + * Adds a new object to this {@code SharedObjects}. Although not necessary, it is recommended + * to only access the object through the returned {@link SharedReference}. + */ + public <T> SharedReference<T> add(T object) { + SharedReference<T> tag = new DefaultTag<>(this, objects.size()); + objects.put(tag, object); + return tag; Review comment: It's _probably_ not necessary, it just stuck out to me because everything else is thread-safe, and I guess wouldn't hurt? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: [email protected]
