From 34aad1a6c8c379de9fbe4eb6df5cc9ea6a9cd5b7 Mon Sep 17 00:00:00 2001
From: Chris Sullivan <csullivan@octoml.ai>
Date: Tue, 11 Aug 2020 13:04:26 -0700
Subject: [PATCH] [JVM] Support overriding RPCWatchdog termination behavior on
 Android and other platforms (#6216)

* Instead of performing a system exit and leaving unhandled items on
the activity stack, finish the RPCActivity and return cleanly to the
MainActivity where the RPCActivity can be restarted automatically.

* Update doc. string for checkstyle.
---
 .../org/apache/tvm/tvmrpc/RPCActivity.java    |  3 +-
 .../apache/tvm/tvmrpc/RPCAndroidWatchdog.java | 39 +++++++++++++++++++
 .../org/apache/tvm/tvmrpc/RPCProcessor.java   | 10 ++++-
 .../java/org/apache/tvm/rpc/RPCWatchdog.java  |  9 ++++-
 4 files changed, 57 insertions(+), 4 deletions(-)
 create mode 100644 apps/android_rpc/app/src/main/java/org/apache/tvm/tvmrpc/RPCAndroidWatchdog.java

diff --git a/apps/android_rpc/app/src/main/java/org/apache/tvm/tvmrpc/RPCActivity.java b/apps/android_rpc/app/src/main/java/org/apache/tvm/tvmrpc/RPCActivity.java
index 54182d77e170..0d7f04782d3a 100644
--- a/apps/android_rpc/app/src/main/java/org/apache/tvm/tvmrpc/RPCActivity.java
+++ b/apps/android_rpc/app/src/main/java/org/apache/tvm/tvmrpc/RPCActivity.java
@@ -51,7 +51,7 @@ public void onClick(View v) {
     int port = intent.getIntExtra("port", 9090);
     String key = intent.getStringExtra("key");
 
-    tvmServerWorker = new RPCProcessor();
+    tvmServerWorker = new RPCProcessor(this);
     tvmServerWorker.setDaemon(true);
     tvmServerWorker.start();
     tvmServerWorker.connect(host, port, key);
@@ -62,5 +62,6 @@ protected void onDestroy() {
     System.err.println("rpc activity onDestroy");
     tvmServerWorker.disconnect();
     super.onDestroy();
+    android.os.Process.killProcess(android.os.Process.myPid());
   }
 }
diff --git a/apps/android_rpc/app/src/main/java/org/apache/tvm/tvmrpc/RPCAndroidWatchdog.java b/apps/android_rpc/app/src/main/java/org/apache/tvm/tvmrpc/RPCAndroidWatchdog.java
new file mode 100644
index 000000000000..69b7d27d0cf8
--- /dev/null
+++ b/apps/android_rpc/app/src/main/java/org/apache/tvm/tvmrpc/RPCAndroidWatchdog.java
@@ -0,0 +1,39 @@
+/*
+ * 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.tvm.tvmrpc;
+
+import android.app.Activity;
+import org.apache.tvm.rpc.RPCWatchdog;
+
+/**
+ * Watchdog for Android RPC.
+ */
+public class RPCAndroidWatchdog extends RPCWatchdog {
+  public Activity rpc_activity = null;
+  public RPCAndroidWatchdog(Activity activity) {
+    super();
+    rpc_activity = activity;
+  }
+
+  /**
+   * Method to non-destructively terminate the running thread on Android
+   */
+  @Override protected void terminate() {
+    rpc_activity.finish();
+  }
+}
diff --git a/apps/android_rpc/app/src/main/java/org/apache/tvm/tvmrpc/RPCProcessor.java b/apps/android_rpc/app/src/main/java/org/apache/tvm/tvmrpc/RPCProcessor.java
index ca0c9c8f7584..8bb44f10f1a4 100644
--- a/apps/android_rpc/app/src/main/java/org/apache/tvm/tvmrpc/RPCProcessor.java
+++ b/apps/android_rpc/app/src/main/java/org/apache/tvm/tvmrpc/RPCProcessor.java
@@ -17,10 +17,10 @@
 
 package org.apache.tvm.tvmrpc;
 
+import android.app.Activity;
 import android.os.ParcelFileDescriptor;
 import java.net.Socket;
 import org.apache.tvm.rpc.ConnectTrackerServerProcessor;
-import org.apache.tvm.rpc.RPCWatchdog;
 
 /**
  * Connect to RPC proxy and deal with requests.
@@ -33,9 +33,15 @@ class RPCProcessor extends Thread {
   private long startTime;
   private ConnectTrackerServerProcessor currProcessor;
   private boolean first = true;
+  private Activity rpc_activity = null;
+
+  public RPCProcessor(Activity activity) {
+    super();
+    rpc_activity = activity;
+  }
 
   @Override public void run() {
-    RPCWatchdog watchdog = new RPCWatchdog();
+    RPCAndroidWatchdog watchdog = new RPCAndroidWatchdog(rpc_activity);
     watchdog.start();
     while (true) {
       synchronized (this) {
diff --git a/jvm/core/src/main/java/org/apache/tvm/rpc/RPCWatchdog.java b/jvm/core/src/main/java/org/apache/tvm/rpc/RPCWatchdog.java
index 33dbdde4d4c9..ef5bb6058ae1 100644
--- a/jvm/core/src/main/java/org/apache/tvm/rpc/RPCWatchdog.java
+++ b/jvm/core/src/main/java/org/apache/tvm/rpc/RPCWatchdog.java
@@ -71,7 +71,7 @@ public synchronized void finishTimeout() {
             } else {
               System.err.println("watchdog woke up!");
               System.err.println("terminating...");
-              System.exit(0);
+              terminate();
             }
           } catch (InterruptedException e) {
             System.err.println("watchdog interrupted...");
@@ -80,4 +80,11 @@ public synchronized void finishTimeout() {
       }
     }
   }
+
+  /**
+   * Default method to terminate the running RPCActivity process.
+   */
+  protected void terminate() {
+    System.exit(0);
+  }
 }