Invoke RPCs on the main thread if they are tagged with @RunOnUiThread. (#47)

Fixes concurrency handling in SnippetManager.

Fixes #51 

7 files changed, 188 insertions, 76 deletions

diff --git a/third_party/sl4a/src/main/java/com/google/android/mobly/snippet/event/EventSnippet.java b/third_party/sl4a/src/main/java/com/google/android/mobly/snippet/event/EventSnippet.java
index 1a29b39..fb11720 100644
--- a/third_party/sl4a/src/main/java/com/google/android/mobly/snippet/event/EventSnippet.java
+++ b/third_party/sl4a/src/main/java/com/google/android/mobly/snippet/event/EventSnippet.java
@@ -28,8 +28,8 @@ import org.json.JSONException;
 import org.json.JSONObject;
 
 public class EventSnippet implements Snippet {
-
     private static class EventSnippetException extends Exception {
+        private static final long serialVersionUID = 1L;
         public EventSnippetException(String msg) {
             super(msg);
         }
diff --git a/third_party/sl4a/src/main/java/com/google/android/mobly/snippet/manager/SnippetManager.java b/third_party/sl4a/src/main/java/com/google/android/mobly/snippet/manager/SnippetManager.java
index 8944a84..4befdd5 100644
--- a/third_party/sl4a/src/main/java/com/google/android/mobly/snippet/manager/SnippetManager.java
+++ b/third_party/sl4a/src/main/java/com/google/android/mobly/snippet/manager/SnippetManager.java
@@ -20,38 +20,49 @@ import android.os.Build;
 import com.google.android.mobly.snippet.Snippet;
 import com.google.android.mobly.snippet.rpc.MethodDescriptor;
 import com.google.android.mobly.snippet.rpc.RpcMinSdk;
+import com.google.android.mobly.snippet.rpc.RunOnUiThread;
 import com.google.android.mobly.snippet.util.Log;
+import com.google.android.mobly.snippet.util.MainThread;
 import com.google.android.mobly.snippet.util.SnippetLibException;
 import java.lang.reflect.Constructor;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Map.Entry;
 import java.util.SortedSet;
 import java.util.TreeSet;
+import java.util.concurrent.Callable;
 
 public class SnippetManager {
-    private final Map<Class<? extends Snippet>, Snippet> mReceivers;
+    private final Map<Class<? extends Snippet>, Snippet> mSnippets;
     /** A map of strings to known RPCs. */
-    private final Map<String, MethodDescriptor> mKnownRpcs =
-            new HashMap<String, MethodDescriptor>();
+    private final Map<String, MethodDescriptor> mKnownRpcs;
 
     public SnippetManager(Collection<Class<? extends Snippet>> classList) {
-        mReceivers = new HashMap<>();
+        // Synchronized for multiple connections on the same session. Can't use ConcurrentHashMap
+        // because we have to put in a value of 'null' before the class is constructed, but
+        // ConcurrentHashMap does not allow null values.
+        mSnippets = Collections.synchronizedMap(new HashMap<Class<? extends Snippet>, Snippet>());
+        Map<String, MethodDescriptor> knownRpcs = new HashMap<>();
         for (Class<? extends Snippet> receiverClass : classList) {
-            mReceivers.put(receiverClass, null);
+            mSnippets.put(receiverClass, null);
             Collection<MethodDescriptor> methodList = MethodDescriptor.collectFrom(receiverClass);
             for (MethodDescriptor m : methodList) {
-                if (mKnownRpcs.containsKey(m.getName())) {
+                if (knownRpcs.containsKey(m.getName())) {
                     // We already know an RPC of the same name. We don't catch this anywhere because
                     // this is a programming error.
                     throw new RuntimeException(
                             "An RPC with the name " + m.getName() + " is already known.");
                 }
-                mKnownRpcs.put(m.getName(), m);
+                knownRpcs.put(m.getName(), m);
             }
         }
+        // Does not need to be concurrent because this map is read only, so it is safe to access
+        // from multiple threads. Wrap in an unmodifiableMap to enforce this.
+        mKnownRpcs = Collections.unmodifiableMap(knownRpcs);
     }
 
     public MethodDescriptor getMethodDescriptor(String methodName) {
@@ -73,35 +84,78 @@ public class SnippetManager {
                                 method.getName(), requiredSdkLevel, Build.VERSION.SDK_INT));
             }
         }
+        Snippet object;
         try {
-            Snippet object = get(clazz);
-            return method.invoke(object, args);
+            object = get(clazz);
+            return invoke(object, method, args);
         } catch (InvocationTargetException e) {
             throw e.getCause();
         }
     }
 
-    public void shutdown() {
-        for (Snippet receiver : mReceivers.values()) {
-            try {
-                if (receiver != null) {
-                    receiver.shutdown();
-                }
-            } catch (Exception e) {
-                Log.e("Failed to shut down an Snippet", e);
+    public void shutdown() throws Exception {
+        for (final Entry<Class<? extends Snippet>, Snippet> entry : mSnippets.entrySet()) {
+            if (entry.getValue() == null) {
+                continue;
+            }
+            Method method = entry.getKey().getMethod("shutdown");
+            if (method.isAnnotationPresent(RunOnUiThread.class)) {
+                Log.d("Shutting down " + entry.getKey().getName() + " on the main thread");
+                MainThread.run(new Callable<Void>() {
+                    @Override
+                    public Void call() throws Exception {
+                        entry.getValue().shutdown();
+                        return null;
+                    }
+                });
+            } else {
+                Log.d("Shutting down " + entry.getKey().getName());
+                entry.getValue().shutdown();
             }
         }
     }
 
     private Snippet get(Class<? extends Snippet> clazz) throws Exception {
-        Snippet object = mReceivers.get(clazz);
-        if (object != null) {
-            return object;
+        Snippet snippetImpl = mSnippets.get(clazz);
+        if (snippetImpl == null) {
+            // First time calling an RPC for this snippet; construct an instance under lock.
+            synchronized (clazz) {
+                snippetImpl = mSnippets.get(clazz);
+                if (snippetImpl == null) {
+                    final Constructor<? extends Snippet> constructor = clazz.getConstructor();
+                    if (constructor.isAnnotationPresent(RunOnUiThread.class)) {
+                        Log.d("Constructing " + clazz + " on the main thread");
+                        snippetImpl = MainThread.run(new Callable<Snippet>() {
+                            @Override
+                            public Snippet call() throws Exception {
+                                return constructor.newInstance();
+                            }
+                        });
+                    } else {
+                        Log.d("Constructing " + clazz);
+                        snippetImpl = constructor.newInstance();
+                    }
+                    mSnippets.put(clazz, snippetImpl);
+                }
+            }
+        }
+        return snippetImpl;
+    }
+
+    private Object invoke(final Snippet snippetImpl, final Method method, final Object[] args)
+            throws Exception {
+        if (method.isAnnotationPresent(RunOnUiThread.class)) {
+            Log.d("Invoking RPC method " + method.getDeclaringClass() + "#" + method.getName()
+                      + " on the main thread");
+            return MainThread.run(new Callable<Object>() {
+                @Override
+                public Object call() throws Exception {
+                    return method.invoke(snippetImpl, args);
+                }
+            });
+        } else {
+            Log.d("Invoking RPC method " + method.getDeclaringClass() + "#" + method.getName());
+            return method.invoke(snippetImpl, args);
         }
-        Constructor<? extends Snippet> constructor;
-        constructor = clazz.getConstructor();
-        object = constructor.newInstance();
-        mReceivers.put(clazz, object);
-        return object;
     }
 }
diff --git a/third_party/sl4a/src/main/java/com/google/android/mobly/snippet/rpc/JsonRpcServer.java b/third_party/sl4a/src/main/java/com/google/android/mobly/snippet/rpc/JsonRpcServer.java
index e031dcb..0486482 100644
--- a/third_party/sl4a/src/main/java/com/google/android/mobly/snippet/rpc/JsonRpcServer.java
+++ b/third_party/sl4a/src/main/java/com/google/android/mobly/snippet/rpc/JsonRpcServer.java
@@ -44,15 +44,6 @@ public class JsonRpcServer extends SimpleServer {
     }
 
     @Override
-    public void shutdown() {
-        super.shutdown();
-        // Notify all RPC receiving objects. They may have to clean up some of their state.
-        for (SnippetManager manager : mSnippetManagerFactory.getSnippetManagers().values()) {
-            manager.shutdown();
-        }
-    }
-
-    @Override
     protected void handleRPCConnection(
             Socket sock, Integer UID, BufferedReader reader, PrintWriter writer) throws Exception {
         SnippetManager receiverManager = null;
@@ -82,14 +73,24 @@ public class JsonRpcServer extends SimpleServer {
                 continue;
             } else if (method.equals(CMD_CLOSE_SESSION)) {
                 Log.d("Got shutdown signal");
-                send(writer, JsonRpcResult.empty(id), UID);
                 synchronized (writer) {
-                    receiverManager.shutdown();
+                    // Shut down all RPC receivers.
+                    for (SnippetManager manager :
+                             mSnippetManagerFactory.getSnippetManagers().values()) {
+                        manager.shutdown();
+                    }
+
+                    // Shut down this client connection. As soon as this happens, the client will
+                    // kill us by triggering the 'stop' action from another instrumentation, so no
+                    // other cleanup steps are guaranteed to execute.
+                    send(writer, JsonRpcResult.empty(id), UID);
                     reader.close();
                     writer.close();
                     sock.close();
+
+                    // Shut down this server.
                     shutdown();
-                    mgrs.remove(UID);
+                    mgrs.clear();
                 }
                 return;
             }
diff --git a/third_party/sl4a/src/main/java/com/google/android/mobly/snippet/rpc/MethodDescriptor.java b/third_party/sl4a/src/main/java/com/google/android/mobly/snippet/rpc/MethodDescriptor.java
index 4c04bb6..2d822f3 100644
--- a/third_party/sl4a/src/main/java/com/google/android/mobly/snippet/rpc/MethodDescriptor.java
+++ b/third_party/sl4a/src/main/java/com/google/android/mobly/snippet/rpc/MethodDescriptor.java
@@ -63,7 +63,7 @@ public final class MethodDescriptor {
      *
      * @param parameters {@code JSONArray} containing the parameters
      * @return result
-     * @throws Throwable
+     * @throws Throwable the exception raised from executing the RPC method.
      */
     public Object invoke(SnippetManager manager, final JSONArray parameters) throws Throwable {
         final Type[] parameterTypes = getGenericParameterTypes();
diff --git a/third_party/sl4a/src/main/java/com/google/android/mobly/snippet/rpc/RunOnUiThread.java b/third_party/sl4a/src/main/java/com/google/android/mobly/snippet/rpc/RunOnUiThread.java
new file mode 100644
index 0000000..3e2d31e
--- /dev/null
+++ b/third_party/sl4a/src/main/java/com/google/android/mobly/snippet/rpc/RunOnUiThread.java
@@ -0,0 +1,34 @@
+/*
+ * Copyright (C) 2017 Google Inc.
+ *
+ * Licensed 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 com.google.android.mobly.snippet.rpc;
+
+import com.google.android.mobly.snippet.Snippet;
+import java.lang.annotation.Documented;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+
+/** This annotation will cause the RPC to execute on the main app thread.
+ *
+ * This annotation can be applied to:
+ * <ul>
+ *   <li>The constructor of a class implementing the {@link Snippet} interface.
+ *   <li>A method annotated with the {@link Rpc} or {@link AsyncRpc} annotation.
+ *   <li>The {@link Snippet#shutdown()} method.
+ * </ul>
+ */
+@Retention(RetentionPolicy.RUNTIME)
+@Documented
+public @interface RunOnUiThread {}
diff --git a/third_party/sl4a/src/main/java/com/google/android/mobly/snippet/rpc/SimpleServer.java b/third_party/sl4a/src/main/java/com/google/android/mobly/snippet/rpc/SimpleServer.java
index 61880be..4a87386 100644
--- a/third_party/sl4a/src/main/java/com/google/android/mobly/snippet/rpc/SimpleServer.java
+++ b/third_party/sl4a/src/main/java/com/google/android/mobly/snippet/rpc/SimpleServer.java
@@ -284,7 +284,7 @@ public abstract class SimpleServer {
         }
     }
 
-    public void shutdown() {
+    public void shutdown() throws Exception {
         // Stop listening on the server socket to ensure that
         // beyond this point there are no incoming requests.
         mStopServer = true;
diff --git a/third_party/sl4a/src/main/java/com/google/android/mobly/snippet/util/MainThread.java b/third_party/sl4a/src/main/java/com/google/android/mobly/snippet/util/MainThread.java
index 480498a..44b4f2b 100644
--- a/third_party/sl4a/src/main/java/com/google/android/mobly/snippet/util/MainThread.java
+++ b/third_party/sl4a/src/main/java/com/google/android/mobly/snippet/util/MainThread.java
@@ -16,52 +16,75 @@
 
 package com.google.android.mobly.snippet.util;
 
-import android.content.Context;
 import android.os.Handler;
-import com.google.android.mobly.snippet.future.FutureResult;
+import android.os.Looper;
 import java.util.concurrent.Callable;
+import java.util.concurrent.CountDownLatch;
 
 public class MainThread {
+    /**
+    * Wraps a {@link Callable} in a {@link Runnable} that has a way to get the return value and
+    * exception after the fact.
+    */
+    private static class CallableWrapper<T> implements Runnable {
+        private final Callable<T> mCallable;
+        private final CountDownLatch mLatch = new CountDownLatch(1);
+        private T mReturnValue;
+        private Throwable mException;
+
+        public CallableWrapper(Callable<T> callable) {
+            mCallable = callable;
+        }
+
+        @Override
+        public final void run() {
+            try {
+                mReturnValue = mCallable.call();
+            } catch (Throwable t) {
+                mException = t;
+            } finally {
+                mLatch.countDown();
+            }
+        }
+
+        public void awaitTermination() throws InterruptedException {
+            mLatch.await();
+        }
+
+        public T getReturnValue() {
+            return mReturnValue;
+        }
+
+        public Throwable getException() {
+            return mException;
+        }
+    }
+
+    private static final Handler sMainThreadHandler = new Handler(Looper.getMainLooper());
 
     private MainThread() {
         // Utility class.
     }
 
-    /**
-     * Executed in the main thread, returns the result of an execution. Anything that runs here
-     * should finish quickly to avoid hanging the UI thread.
-     */
-    public static <T> T run(Context context, final Callable<T> task) {
-        final FutureResult<T> result = new FutureResult<T>();
-        Handler handler = new Handler(context.getMainLooper());
-        handler.post(
-                new Runnable() {
-                    @Override
-                    public void run() {
-                        try {
-                            result.set(task.call());
-                        } catch (Exception e) {
-                            Log.e(e);
-                            result.set(null);
-                        }
-                    }
-                });
-        try {
-            return result.get();
-        } catch (InterruptedException e) {
-            Log.e(e);
-        }
-        return null;
+    /** Executed in the main thread. Returns the result of an execution or any exception thrown. */
+    public static <T> T run(final Callable<T> task) throws Exception {
+        CallableWrapper<T> wrapper = new CallableWrapper<>(task);
+        return runCallableWrapper(wrapper);
     }
 
-    public static void run(Context context, final Runnable task) {
-        Handler handler = new Handler(context.getMainLooper());
-        handler.post(
-                new Runnable() {
-                    @Override
-                    public void run() {
-                        task.run();
-                    }
-                });
+    private static <T> T runCallableWrapper(CallableWrapper<T> wrapper) throws Exception {
+        sMainThreadHandler.post(wrapper);
+        wrapper.awaitTermination();
+        Throwable exception = wrapper.getException();
+        if (exception != null) {
+            if (exception instanceof RuntimeException) {
+                throw (RuntimeException) exception;
+            }
+            if (exception instanceof Error) {
+                throw (Error) exception;
+            }
+            throw (Exception) exception;
+        }
+        return wrapper.getReturnValue();
     }
 }