Use heartbeat mechanism to detect server ready.

This change detects the moment remote USB server is ready to exchange messages
and attaches the USB device to localhost.

Change-Id: I7c4904a68092c328e55d6f54dff11557c995be1d
(cherry picked from commit d53238a1c4e68dce0578185600f998400f786770)
diff --git a/vadb/BUILD b/vadb/BUILD
index 726f74e..a116670 100644
--- a/vadb/BUILD
+++ b/vadb/BUILD
@@ -28,6 +28,8 @@
         "usb_cmd_data_transfer.h",
         "usb_cmd_device_list.cpp",
         "usb_cmd_device_list.h",
+        "usb_cmd_heartbeat.cpp",
+        "usb_cmd_heartbeat.h",
         "virtual_adb_client.cpp",
         "virtual_adb_client.h",
         "virtual_adb_server.cpp",
diff --git a/vadb/usb_cmd_heartbeat.cpp b/vadb/usb_cmd_heartbeat.cpp
new file mode 100644
index 0000000..864581c
--- /dev/null
+++ b/vadb/usb_cmd_heartbeat.cpp
@@ -0,0 +1,31 @@
+/*
+ * Copyright (C) 2017 The Android Open Source Project
+ *
+ * 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.
+ */
+#include <glog/logging.h>
+
+#include "guest/usbforward/protocol.h"
+#include "host/vadb/usb_cmd_heartbeat.h"
+
+namespace vadb {
+bool USBCmdHeartbeat::OnRequest(const avd::SharedFD& fd) { return true; }
+
+bool USBCmdHeartbeat::OnResponse(bool is_success, const avd::SharedFD& data) {
+  callback_(is_success);
+  return true;
+}
+
+USBCmdHeartbeat::USBCmdHeartbeat(USBCmdHeartbeat::HeartbeatResultCB callback)
+    : callback_(callback) {}
+}  // namespace vadb
diff --git a/vadb/usb_cmd_heartbeat.h b/vadb/usb_cmd_heartbeat.h
new file mode 100644
index 0000000..894c5f2
--- /dev/null
+++ b/vadb/usb_cmd_heartbeat.h
@@ -0,0 +1,48 @@
+/*
+ * Copyright (C) 2017 The Android Open Source Project
+ *
+ * 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.
+ */
+#pragma once
+
+#include "host/vadb/usb_cmd.h"
+
+namespace vadb {
+// Request remote device attach (~open).
+class USBCmdHeartbeat : public USBCommand {
+ public:
+  // Heartbeat result callback receives a boolean argument indicating whether
+  // remote device is ready to be attached.
+  using HeartbeatResultCB = std::function<void(bool)>;
+
+  USBCmdHeartbeat(HeartbeatResultCB callback);
+  ~USBCmdHeartbeat() override = default;
+
+  // Return usbforward command this instance is executing.
+  usb_forward::Command Command() override { return usb_forward::CmdHeartbeat; }
+
+  // Send request body to the server.
+  // Return false, if communication failed.
+  bool OnRequest(const avd::SharedFD& data) override;
+
+  // Receive response data from the server.
+  // Return false, if communication failed.
+  bool OnResponse(bool is_success, const avd::SharedFD& data) override;
+
+ private:
+  HeartbeatResultCB callback_;
+
+  USBCmdHeartbeat(const USBCmdHeartbeat& other) = delete;
+  USBCmdHeartbeat& operator=(const USBCmdHeartbeat& other) = delete;
+};
+}  // namespace vadb
diff --git a/vadb/virtual_adb_client.cpp b/vadb/virtual_adb_client.cpp
index 56dbdac..440db10 100644
--- a/vadb/virtual_adb_client.cpp
+++ b/vadb/virtual_adb_client.cpp
@@ -20,9 +20,19 @@
 #include "host/vadb/usb_cmd_control_transfer.h"
 #include "host/vadb/usb_cmd_data_transfer.h"
 #include "host/vadb/usb_cmd_device_list.h"
+#include "host/vadb/usb_cmd_heartbeat.h"
 #include "host/vadb/virtual_adb_client.h"
 
 namespace vadb {
+namespace {
+constexpr int kHeartbeatTimeoutSeconds = 3;
+}  // namespace
+
+VirtualADBClient::VirtualADBClient(usbip::DevicePool* pool, avd::SharedFD fd)
+    : pool_(pool), fd_(fd) {
+  timer_ = avd::SharedFD::TimerFD(CLOCK_MONOTONIC, 0);
+  SendHeartbeat();
+}
 
 void VirtualADBClient::RegisterDevice(
     const usb_forward::DeviceInfo& dev,
@@ -107,6 +117,53 @@
       std::unique_ptr<USBCommand>(new USBCmdAttach(bus_id, dev_id)));
 }
 
+bool VirtualADBClient::SendHeartbeat() {
+  VLOG(1) << "Sending heartbeat...";
+  struct itimerspec spec{};
+  spec.it_value.tv_sec = kHeartbeatTimeoutSeconds;
+  timer_->TimerSet(0, &spec, nullptr);
+
+  heartbeat_tag_ = tag_;
+
+  return ExecuteCommand(std::unique_ptr<USBCommand>(
+      new USBCmdHeartbeat([this](bool success) { HandleHeartbeat(success); })));
+}
+
+void VirtualADBClient::HandleHeartbeat(bool is_ready) {
+  VLOG(1) << "Remote server status: " << is_ready;
+  if (is_ready && !is_remote_server_ready_) {
+    LOG(INFO) << "Remote server is now ready.";
+    PopulateRemoteDevices();
+  } else if (is_remote_server_ready_ && !is_ready) {
+    LOG(WARNING) << "Remote server connection lost.";
+    // It makes perfect sense to cancel all outstanding USB requests, as device
+    // is not going to answer any of these anyway.
+    for (const auto& pair : commands_) {
+      pair.second->OnResponse(false, fd_);
+    }
+    commands_.clear();
+  }
+  is_remote_server_ready_ = is_ready;
+}
+
+bool VirtualADBClient::HandleHeartbeatTimeout() {
+  uint64_t timer_result;
+  timer_->Read(&timer_result, sizeof(timer_result));
+
+  auto iter = commands_.find(heartbeat_tag_);
+  if (iter != commands_.end()) {
+    // Make sure to erase the value from list of commands prior to running
+    // callback. Particularly important for heartbeat, which cancels all
+    // outstanding USB commands (including self, if found), if device goes
+    // away (eg. reboots).
+    auto command = std::move(iter->second);
+    commands_.erase(iter);
+    command->OnResponse(false, fd_);
+  }
+
+  return SendHeartbeat();
+}
+
 bool VirtualADBClient::ExecuteCommand(std::unique_ptr<USBCommand> cmd) {
   uint32_t this_tag = tag_;
   tag_++;
@@ -126,12 +183,16 @@
 // SharedFDs.
 void VirtualADBClient::BeforeSelect(avd::SharedFDSet* fd_read) const {
   fd_read->Set(fd_);
+  fd_read->Set(timer_);
 }
 
 // AfterSelect is Called right after Select() to detect and respond to changes
 // on affected SharedFDs.
 // Return value indicates whether this client is still valid.
 bool VirtualADBClient::AfterSelect(const avd::SharedFDSet& fd_read) {
+  if (fd_read.IsSet(timer_)) {
+    HandleHeartbeatTimeout();
+  }
   if (fd_read.IsSet(fd_)) {
     usb_forward::ResponseHeader rhdr;
     if (fd_->Read(&rhdr, sizeof(rhdr)) != sizeof(rhdr)) {
@@ -143,15 +204,18 @@
 
     auto iter = commands_.find(rhdr.tag);
     if (iter == commands_.end()) {
-      LOG(ERROR)
-          << "Response does not match any of the previously queued commands!";
-      // TODO(ender): This looks like a protocol error. What should happen now?
-      // Should we cancel all pending commands now?
-      return false;
+      // This is likely a late heartbeat response, but could very well be any of
+      // the remaining commands.
+      LOG(INFO) << "Received response for discarded tag " << rhdr.tag;
+    } else {
+      // Make sure to erase the value from list of commands prior to running
+      // callback. Particularly important for heartbeat, which cancels all
+      // outstanding USB commands (including self, if found), if device goes
+      // away (eg. reboots).
+      auto command = std::move(iter->second);
+      commands_.erase(iter);
+      command->OnResponse(rhdr.status == usb_forward::StatusSuccess, fd_);
     }
-
-    iter->second->OnResponse(rhdr.status == usb_forward::StatusSuccess, fd_);
-    commands_.erase(iter);
   }
 
   return true;
diff --git a/vadb/virtual_adb_client.h b/vadb/virtual_adb_client.h
index 6fbef83..2ade53e 100644
--- a/vadb/virtual_adb_client.h
+++ b/vadb/virtual_adb_client.h
@@ -34,8 +34,7 @@
 // remote USB devices possible with help of USB/IP protocol.
 class VirtualADBClient {
  public:
-  VirtualADBClient(usbip::DevicePool* pool, avd::SharedFD fd)
-      : pool_(pool), fd_(fd) {}
+  VirtualADBClient(usbip::DevicePool* pool, avd::SharedFD fd);
 
   virtual ~VirtualADBClient() = default;
 
@@ -71,6 +70,18 @@
                                std::vector<uint8_t> data,
                                usbip::Device::AsyncTransferReadyCB callback);
 
+  // Send new heartbeat request and arm the heartbeat timer.
+  bool SendHeartbeat();
+
+  // Heartbeat handler receives response to heartbeat request.
+  // Supplied argument indicates, whether remote server is ready to export USB
+  // gadget.
+  void HandleHeartbeat(bool is_ready);
+
+  // Heartbeat timeout detects situation where heartbeat did not receive
+  // matching response. This could be a direct result of device reset.
+  bool HandleHeartbeatTimeout();
+
   // ExecuteCommand creates command header and executes supplied USBCommand.
   // If execution was successful, command will be stored internally until
   // response arrives.
@@ -78,8 +89,13 @@
 
   usbip::DevicePool* pool_;
   avd::SharedFD fd_;
+  avd::SharedFD timer_;
+  bool is_remote_server_ready_ = false;
 
   uint32_t tag_ = 0;
+  // Assign an 'invalid' tag as previously sent heartbeat command. This will
+  // prevent heartbeat timeout handler from finding a command if none was sent.
+  uint32_t heartbeat_tag_ = ~0;
   std::map<uint32_t, std::unique_ptr<USBCommand>> commands_;
 
   VirtualADBClient(const VirtualADBClient& other) = delete;
diff --git a/vadb/virtual_adb_server.cpp b/vadb/virtual_adb_server.cpp
index abd2bac..8d331e7 100644
--- a/vadb/virtual_adb_server.cpp
+++ b/vadb/virtual_adb_server.cpp
@@ -59,7 +59,6 @@
   }
 
   clients_.emplace_back(&pool_, client);
-  clients_.back().PopulateRemoteDevices();
 }
 
 }  // namespace vadb