Merge "Add a timeout option in shell subscriber."
diff --git a/cmds/statsd/Android.bp b/cmds/statsd/Android.bp
index d610f66..59b2aa6 100644
--- a/cmds/statsd/Android.bp
+++ b/cmds/statsd/Android.bp
@@ -332,6 +332,8 @@
         "src/stats_log.proto",
         "src/statsd_config.proto",
         "src/atoms.proto",
+        "src/shell/shell_config.proto",
+        "src/shell/shell_data.proto",
     ],
 
     static_libs: [
diff --git a/cmds/statsd/src/StatsService.cpp b/cmds/statsd/src/StatsService.cpp
index f2a4663..3107b4d 100644
--- a/cmds/statsd/src/StatsService.cpp
+++ b/cmds/statsd/src/StatsService.cpp
@@ -360,7 +360,11 @@
             if (mShellSubscriber == nullptr) {
                 mShellSubscriber = new ShellSubscriber(mUidMap, mPullerManager);
             }
-            mShellSubscriber->startNewSubscription(in, out, resultReceiver);
+            int timeoutSec = -1;
+            if (argCount >= 2) {
+                timeoutSec = atoi(args[1].c_str());
+            }
+            mShellSubscriber->startNewSubscription(in, out, resultReceiver, timeoutSec);
             return NO_ERROR;
         }
     }
diff --git a/cmds/statsd/src/shell/ShellSubscriber.cpp b/cmds/statsd/src/shell/ShellSubscriber.cpp
index 22883f3..52d5ffc 100644
--- a/cmds/statsd/src/shell/ShellSubscriber.cpp
+++ b/cmds/statsd/src/shell/ShellSubscriber.cpp
@@ -30,7 +30,8 @@
 
 const static int FIELD_ID_ATOM = 1;
 
-void ShellSubscriber::startNewSubscription(int in, int out, sp<IResultReceiver> resultReceiver) {
+void ShellSubscriber::startNewSubscription(int in, int out, sp<IResultReceiver> resultReceiver,
+                                           int timeoutSec) {
     VLOG("start new shell subscription");
     {
         std::lock_guard<std::mutex> lock(mMutex);
@@ -50,11 +51,18 @@
     // Read config forever until EOF is reached. Clients may send multiple configs -- each new
     // config replace the previous one.
     readConfig(in);
+    VLOG("timeout : %d", timeoutSec);
 
     // Now we have read an EOF we now wait for the semaphore until the client exits.
     VLOG("Now wait for client to exit");
     std::unique_lock<std::mutex> lk(mMutex);
-    mShellDied.wait(lk, [this, resultReceiver] { return mResultReceiver != resultReceiver; });
+
+    if (timeoutSec > 0) {
+        mShellDied.wait_for(lk, timeoutSec * 1s,
+                            [this, resultReceiver] { return mResultReceiver != resultReceiver; });
+    } else {
+        mShellDied.wait(lk, [this, resultReceiver] { return mResultReceiver != resultReceiver; });
+    }
 }
 
 void ShellSubscriber::updateConfig(const ShellSubscription& config) {
diff --git a/cmds/statsd/src/shell/ShellSubscriber.h b/cmds/statsd/src/shell/ShellSubscriber.h
index 5401f31..8e54a8b 100644
--- a/cmds/statsd/src/shell/ShellSubscriber.h
+++ b/cmds/statsd/src/shell/ShellSubscriber.h
@@ -65,7 +65,8 @@
     /**
      * Start a new subscription.
      */
-    void startNewSubscription(int inFd, int outFd, sp<IResultReceiver> resultReceiver);
+    void startNewSubscription(int inFd, int outFd, sp<IResultReceiver> resultReceiver,
+                              int timeoutSec);
 
     void binderDied(const wp<IBinder>& who);
 
diff --git a/cmds/statsd/tests/shell/ShellSubscriber_test.cpp b/cmds/statsd/tests/shell/ShellSubscriber_test.cpp
index a184f56..73d1fd7 100644
--- a/cmds/statsd/tests/shell/ShellSubscriber_test.cpp
+++ b/cmds/statsd/tests/shell/ShellSubscriber_test.cpp
@@ -83,7 +83,7 @@
 
     // mimic a binder thread that a shell subscriber runs on. it would block.
     std::thread reader([&resultReceiver, &fds_config, &fds_data, &shellClient] {
-        shellClient->startNewSubscription(fds_config[0], fds_data[1], resultReceiver);
+        shellClient->startNewSubscription(fds_config[0], fds_data[1], resultReceiver, -1);
     });
     reader.detach();