[GitHub] [geode] kirklund commented on a change in pull request #5664: GEODE-8650: Support multiple instances of DistributedReference

[GitBox <gi...@apache.org>]

kirklund commented on a change in pull request #5664:
URL: https://github.com/apache/geode/pull/5664#discussion_r511235182



##########
File path: geode-dunit/src/main/java/org/apache/geode/test/dunit/rules/DistributedReference.java
##########
@@ -187,47 +197,76 @@ public V get() {
    * @param newValue the new value
    */
   public DistributedReference<V> set(V newValue) {
-    REFERENCE.set(newValue);
+    REFERENCE.get().put(identity, newValue);
     return this;
   }
 
+  @Override
+  protected void before() {
+    invoker().invokeInEveryVMAndController(() -> invokeBefore());
+  }
+
   @Override
   protected void after() {
-    invoker().invokeInEveryVMAndController(this::invokeAfter);
+    invoker().invokeInEveryVMAndController(() -> invokeAfter());
+  }
+
+  protected void afterCreateVM(VM vm) {
+    vm.invoke(() -> invokeBefore());
+  }
+
+  protected void beforeBounceVM(VM vm) {
+    // override if needed
+  }
+
+  protected void afterBounceVM(VM vm) {
+    vm.invoke(() -> invokeBefore());
+  }
+
+  private void invokeBefore() {
+    REFERENCE.compareAndSet(null, new HashMap<>());
+    REFERENCE.get().putIfAbsent(identity, null);
   }
 
   private void invokeAfter() {
-    V value = get();
-    if (value == null) {
+    Map<Integer, Object> references = REFERENCE.getAndSet(null);
+    if (references == null) {
+      return;
+    }
+
+    for (Object object : references.values()) {
+      invokeAfter(object);
+    }
+  }
+
+  private void invokeAfter(Object object) {
+    if (object == null) {
       return;
     }
-    REFERENCE.set(null);
 
     if (autoClose.get()) {
-      autoClose(value);
+      autoClose(object);
     }
   }
 
-  private void autoClose(V value) {
-    if (value instanceof AutoCloseable) {
-      close((AutoCloseable) value);
+  private void autoClose(Object object) {
+    if (object instanceof AutoCloseable) {
+      close((AutoCloseable) object);
 
-    } else if (hasMethod(value.getClass(), "close")) {
-      invokeMethod(value, "close");
+    } else if (object instanceof AtomicBoolean) {
+      toFalse((AtomicBoolean) object);
 
-    } else if (hasMethod(value.getClass(), "disconnect")) {
-      invokeMethod(value, "disconnect");
+    } else if (object instanceof CountDownLatch) {
+      openLatch((CountDownLatch) object);
 
-    } else if (hasMethod(value.getClass(), "stop")) {
-      invokeMethod(value, "stop");
-    }
-  }
+    } else if (hasMethod(object.getClass(), "close")) {
+      invokeMethod(object, "close");

Review comment:
       Do we really want a method called `close(AtomicBoolean)` that sets the specified value to false? I started out with multiple `close` methods and then I changed them from that to:
   * toFalse(AtomicBoolean)
   * openLatch(CountDownLatch)
   
   And then what about opening `CountDownLatch`? Do we change that to `close` as well even tho close is the opposite of open?
   
   I was thinking that my creation of `CompletionUtils` might be slightly more controversial because it's of dubious value outside `DistributedReference`.




----------------------------------------------------------------
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:
users@infra.apache.org