Mailing-List: contact commits-help@beam.incubator.apache.org; run by ezmlm
Precedence: bulk
Reply-To: dev@beam.incubator.apache.org
Content-Type: text/plain; charset="us-ascii"
MIME-Version: 1.0
Content-Transfer-Encoding: 7bit
From: kenn@apache.org
To: commits@beam.incubator.apache.org
Date: Tue, 25 Oct 2016 17:54:54 -0000
Message-Id: <93e1f85cec0849a484a6f88f46d77066@git.apache.org>
In-Reply-To: <e26232ca6de841b1b002d1018e33f381@git.apache.org>
References: <e26232ca6de841b1b002d1018e33f381@git.apache.org>
Subject: [2/2] incubator-beam git commit: Improve teardown behavior in
 DoFnLifecycleManager
archived-at: Tue, 25 Oct 2016 18:05:50 -0000

Improve teardown behavior in DoFnLifecycleManager

Use Cache invalidation hooks to teardown DoFns that are no longer in the
cache. Ensure that remove() and removeAll() report thrown exceptions
even though the exceptions are not thrown by the LoadingCache.


Project: http://git-wip-us.apache.org/repos/asf/incubator-beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/incubator-beam/commit/7239ebb0
Tree: http://git-wip-us.apache.org/repos/asf/incubator-beam/tree/7239ebb0
Diff: http://git-wip-us.apache.org/repos/asf/incubator-beam/diff/7239ebb0

Branch: refs/heads/master
Commit: 7239ebb0c76f539f476cea0b44b1070e765cca41
Parents: 79bb2c2
Author: Thomas Groh <tgroh@google.com>
Authored: Mon Oct 24 13:43:43 2016 -0700
Committer: Kenneth Knowles <klk@google.com>
Committed: Tue Oct 25 10:46:43 2016 -0700

----------------------------------------------------------------------
 .../runners/direct/DoFnLifecycleManager.java    | 56 +++++++++------
 .../direct/DoFnLifecycleManagerTest.java        | 74 ++++++++++++++++++--
 2 files changed, 104 insertions(+), 26 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/7239ebb0/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DoFnLifecycleManager.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DoFnLifecycleManager.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DoFnLifecycleManager.java
index 23460b6..472b28b 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DoFnLifecycleManager.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DoFnLifecycleManager.java
@@ -21,17 +21,17 @@ package org.apache.beam.runners.direct;
 import com.google.common.cache.CacheBuilder;
 import com.google.common.cache.CacheLoader;
 import com.google.common.cache.LoadingCache;
-import java.util.ArrayList;
+import com.google.common.cache.RemovalListener;
+import com.google.common.cache.RemovalNotification;
 import java.util.Collection;
-import java.util.Iterator;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
 import org.apache.beam.sdk.runners.PipelineRunner;
 import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.DoFn.Setup;
 import org.apache.beam.sdk.transforms.DoFn.Teardown;
 import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
 import org.apache.beam.sdk.util.SerializableUtils;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
 /**
  * Manages {@link DoFn} setup, teardown, and serialization.
@@ -42,16 +42,18 @@ import org.slf4j.LoggerFactory;
  * clearing all cached {@link DoFn DoFns}.
  */
 class DoFnLifecycleManager {
-  private static final Logger LOG = LoggerFactory.getLogger(DoFnLifecycleManager.class);
-
   public static DoFnLifecycleManager of(DoFn<?, ?> original) {
     return new DoFnLifecycleManager(original);
   }
 
   private final LoadingCache<Thread, DoFn<?, ?>> outstanding;
+  private final ConcurrentMap<Thread, Exception> thrownOnTeardown;
 
   private DoFnLifecycleManager(DoFn<?, ?> original) {
-    this.outstanding = CacheBuilder.newBuilder().build(new DeserializingCacheLoader(original));
+    this.outstanding = CacheBuilder.newBuilder()
+        .removalListener(new TeardownRemovedFnListener())
+        .build(new DeserializingCacheLoader(original));
+    thrownOnTeardown = new ConcurrentHashMap<>();
   }
 
   public DoFn<?, ?> get() throws Exception {
@@ -61,8 +63,15 @@ class DoFnLifecycleManager {
 
   public void remove() throws Exception {
     Thread currentThread = Thread.currentThread();
-    DoFn<?, ?> fn = outstanding.asMap().remove(currentThread);
-    DoFnInvokers.INSTANCE.invokerFor(fn).invokeTeardown();
+    outstanding.invalidate(currentThread);
+    // Block until the invalidate is fully completed
+    outstanding.cleanUp();
+    // Remove to try too avoid reporting the same teardown exception twice. May still double-report,
+    // but the second will be suppressed.
+    Exception thrown = thrownOnTeardown.remove(currentThread);
+    if (thrown != null) {
+      throw thrown;
+    }
   }
 
   /**
@@ -73,21 +82,13 @@ class DoFnLifecycleManager {
    * DoFn.Teardown @Teardown} method, and the {@link PipelineRunner} should throw an exception.
    */
   public Collection<Exception> removeAll() throws Exception {
-    Iterator<DoFn<?, ?>> fns = outstanding.asMap().values().iterator();
-    Collection<Exception> thrown = new ArrayList<>();
-    while (fns.hasNext()) {
-      DoFn<?, ?> fn = fns.next();
-      fns.remove();
-      try {
-        DoFnInvokers.INSTANCE.invokerFor(fn).invokeTeardown();
-      } catch (Exception e) {
-        thrown.add(e);
-      }
-    }
-    return thrown;
+    outstanding.invalidateAll();
+    // Make sure all of the teardowns are run
+    outstanding.cleanUp();
+    return thrownOnTeardown.values();
   }
 
-  private class DeserializingCacheLoader extends CacheLoader<Thread, DoFn<?, ?>> {
+  private static class DeserializingCacheLoader extends CacheLoader<Thread, DoFn<?, ?>> {
     private final byte[] original;
 
     public DeserializingCacheLoader(DoFn<?, ?> original) {
@@ -102,4 +103,15 @@ class DoFnLifecycleManager {
       return fn;
     }
   }
+
+  private class TeardownRemovedFnListener implements RemovalListener<Thread, DoFn<?, ?>> {
+    @Override
+    public void onRemoval(RemovalNotification<Thread, DoFn<?, ?>> notification) {
+      try {
+        DoFnInvokers.INSTANCE.newByteBuddyInvoker(notification.getValue()).invokeTeardown();
+      } catch (Exception e) {
+        thrownOnTeardown.put(notification.getKey(), e);
+      }
+    }
+  }
 }

http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/7239ebb0/runners/direct-java/src/test/java/org/apache/beam/runners/direct/DoFnLifecycleManagerTest.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/DoFnLifecycleManagerTest.java b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/DoFnLifecycleManagerTest.java
index aef9d29..59e1e16 100644
--- a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/DoFnLifecycleManagerTest.java
+++ b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/DoFnLifecycleManagerTest.java
@@ -21,6 +21,7 @@ package org.apache.beam.runners.direct;
 import static com.google.common.base.Preconditions.checkState;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.isA;
 import static org.hamcrest.Matchers.not;
 import static org.hamcrest.Matchers.theInstance;
 import static org.junit.Assert.assertThat;
@@ -34,8 +35,11 @@ import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.util.UserCodeException;
 import org.hamcrest.Matchers;
+import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.ExpectedException;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
@@ -44,6 +48,8 @@ import org.junit.runners.JUnit4;
  */
 @RunWith(JUnit4.class)
 public class DoFnLifecycleManagerTest {
+  @Rule public ExpectedException thrown = ExpectedException.none();
+
   private TestFn fn = new TestFn();
   private DoFnLifecycleManager mgr = DoFnLifecycleManager.of(fn);
 
@@ -105,6 +111,17 @@ public class DoFnLifecycleManagerTest {
   }
 
   @Test
+  public void teardownThrowsRemoveThrows() throws Exception {
+    TestFn obtained = (TestFn) mgr.get();
+    obtained.teardown();
+
+    thrown.expect(UserCodeException.class);
+    thrown.expectCause(isA(IllegalStateException.class));
+    thrown.expectMessage("Cannot call teardown: already torn down");
+    mgr.remove();
+  }
+
+  @Test
   public void teardownAllOnRemoveAll() throws Exception {
     CountDownLatch startSignal = new CountDownLatch(1);
     ExecutorService executor = Executors.newCachedThreadPool();
@@ -125,6 +142,38 @@ public class DoFnLifecycleManagerTest {
     }
   }
 
+  @Test
+  public void removeAndRemoveAllConcurrent() throws Exception {
+    CountDownLatch startSignal = new CountDownLatch(1);
+    ExecutorService executor = Executors.newCachedThreadPool();
+    List<Future<TestFn>> futures = new ArrayList<>();
+    for (int i = 0; i < 10; i++) {
+      futures.add(executor.submit(new GetFnCallable(mgr, startSignal)));
+    }
+    startSignal.countDown();
+    List<TestFn> fns = new ArrayList<>();
+    for (Future<TestFn> future : futures) {
+      fns.add(future.get(1L, TimeUnit.SECONDS));
+    }
+    CountDownLatch removeSignal = new CountDownLatch(1);
+    List<Future<Void>> removeFutures = new ArrayList<>();
+    for (int i = 0; i < 5; i++) {
+      // These will reuse the threads used in the GetFns
+      removeFutures.add(executor.submit(new TeardownFnCallable(mgr, removeSignal)));
+    }
+    removeSignal.countDown();
+    assertThat(mgr.removeAll(), Matchers.<Exception>emptyIterable());
+    for (Future<Void> removed : removeFutures) {
+      // Should not have thrown an exception.
+      removed.get();
+    }
+
+    for (TestFn fn : fns) {
+      assertThat(fn.setupCalled, is(true));
+      assertThat(fn.teardownCalled, is(true));
+    }
+  }
+
   private static class GetFnCallable implements Callable<TestFn> {
     private final DoFnLifecycleManager mgr;
     private final CountDownLatch startSignal;
@@ -141,6 +190,23 @@ public class DoFnLifecycleManagerTest {
     }
   }
 
+  private static class TeardownFnCallable implements Callable<Void> {
+    private final DoFnLifecycleManager mgr;
+    private final CountDownLatch startSignal;
+
+    private TeardownFnCallable(DoFnLifecycleManager mgr, CountDownLatch startSignal) {
+      this.mgr = mgr;
+      this.startSignal = startSignal;
+    }
+
+    @Override
+    public Void call() throws Exception {
+      startSignal.await();
+      // Will throw an exception if the TestFn has already been removed from this thread
+      mgr.remove();
+      return null;
+    }
+  }
 
   private static class TestFn extends DoFn<Object, Object> {
     boolean setupCalled = false;
@@ -148,8 +214,8 @@ public class DoFnLifecycleManagerTest {
 
     @Setup
     public void setup() {
-      checkState(!setupCalled);
-      checkState(!teardownCalled);
+      checkState(!setupCalled, "Cannot call setup: already set up");
+      checkState(!teardownCalled, "Cannot call setup: already torn down");
 
       setupCalled = true;
     }
@@ -160,8 +226,8 @@ public class DoFnLifecycleManagerTest {
 
     @Teardown
     public void teardown() {
-      checkState(setupCalled);
-      checkState(!teardownCalled);
+      checkState(setupCalled, "Cannot call teardown: not set up");
+      checkState(!teardownCalled, "Cannot call teardown: already torn down");
 
       teardownCalled = true;
     }