Add a function to check if the client is dead.

AcsAgentClient would try to re-create the reactor_ and register the connection when the rpc is closed down by the server.

However, this self-saving effort could still fail, and then the caller of the AcsAgentClient should stop calling the SendMessage function and instead kill the object.

PiperOrigin-RevId: 693190583

Author: GuestPlatform Team

cpp/acs_agent_client.cc

@@ -136,6 +136,18 @@ absl::Status AcsAgentClient::SendMessage(MessageBody message_body) {
   return AddRequest(request);
 }
 
+bool AcsAgentClient::IsDead() {
+  // If the stream_state_ is kStreamFailedToInitialize or kShutdown, the client
+  // is dead, no point to retry sending messages. If the stream_state_ is
+  // kStreamClosed, the restart_client_thread_ will restart the client soon, the
+  // caller can retry sending messages later with a backoff mechanism. If the
+  // stream_state_ is kStreamNotInitialized, the client is not dead yet, just
+  // waiting for the successful registration.
+  absl::MutexLock lock(&reactor_mtx_);
+  return stream_state_ == ClientState::kStreamFailedToInitialize ||
+         stream_state_ == ClientState::kShutdown;
+}
+
 absl::Status AcsAgentClient::Init() {
   // Register the connection. The registration request should only be sent once.
   std::unique_ptr<Request> registration_request =
@@ -169,7 +181,7 @@ absl::Status AcsAgentClient::AddRequestAndWaitForResponse(
       absl::MutexLock lock(&reactor_mtx_);
       if (request.has_register_connection() &&
           (stream_state_ != ClientState::kStreamNotInitialized &&
-           stream_state_ != ClientState::kStreamClosed)) {
+           stream_state_ != ClientState::kStreamTemporarilyDown)) {
         return absl::InternalError(
             "The stream is not in the correct state to accept new registration "
             "request.");
@@ -272,7 +284,7 @@ void AcsAgentClient::ReactorReadCallback(Response response, bool ok) {
     ABSL_LOG(WARNING) << "ReactorReadCallback not ok";
     // Wakes up RestartReactor() to restart the stream.
     absl::MutexLock lock(&reactor_mtx_);
-    stream_state_ = ClientState::kStreamClosed;
+    stream_state_ = ClientState::kStreamTemporarilyDown;
     return;
   }
   // Wake up ClientReadMessage().
@@ -346,7 +358,7 @@ void AcsAgentClient::RestartClient() {
   while (true) {
     reactor_mtx_.LockWhen(absl::Condition(
         +[](ClientState* stream_state) {
-          return *stream_state == ClientState::kStreamClosed ||
+          return *stream_state == ClientState::kStreamTemporarilyDown ||
                  *stream_state == ClientState::kShutdown;
         },
         &stream_state_));
@@ -369,7 +381,7 @@ void AcsAgentClient::RestartClient() {
     }
     std::unique_ptr<AcsStub> stub = GenerateConnectionIdAndStub();
     if (stub == nullptr) {
-      stream_state_ = ClientState::kStreamFailedToRestart;
+      stream_state_ = ClientState::kStreamFailedToInitialize;
       reactor_mtx_.Unlock();
       return;
     }
@@ -378,7 +390,7 @@ void AcsAgentClient::RestartClient() {
         absl::bind_front(&AcsAgentClient::ReactorReadCallback, this),
         connection_id_);
     if (reactor_ == nullptr) {
-      stream_state_ = ClientState::kStreamFailedToRestart;
+      stream_state_ = ClientState::kStreamFailedToInitialize;
       reactor_mtx_.Unlock();
       ABSL_LOG(WARNING) << "Failed to generate connection id and reactor.";
       return;
@@ -387,7 +399,7 @@ void AcsAgentClient::RestartClient() {
     // Initialize the client.
     absl::Status init_status = Init();
     if (!init_status.ok()) {
-      stream_state_ = ClientState::kStreamFailedToRestart;
+      stream_state_ = ClientState::kStreamFailedToInitialize;
       reactor_mtx_.Unlock();
       return;
     }

cpp/acs_agent_client.h

@@ -67,6 +67,14 @@ class AcsAgentClient {
       ABSL_LOCKS_EXCLUDED(request_delivery_status_mtx_)
           ABSL_LOCKS_EXCLUDED(reactor_mtx_);
 
+  // Checks if the client is dead. If the caller of this class has failure of
+  // sending or receiving messages, it can call this function to check if the
+  // client is dead and needs a restart. If the client is dead, the caller can
+  // call Shutdown() or directly invoke destructor to clean up the client. If
+  // the client is not dead, the caller can retry sending messages later with a
+  // self-determined backoff mechanism.
+  bool IsDead() ABSL_LOCKS_EXCLUDED(reactor_mtx_);
+
   // Shuts down the client by joining the restart client thread and the
   // read_response_thread_, and then cancel the RPC.
   void Shutdown() ABSL_LOCKS_EXCLUDED(reactor_mtx_)
@@ -218,12 +226,12 @@ class AcsAgentClient {
   enum class ClientState {
     // The client is ready to read any Response from the server.
     kReady,
-    // Stream not initialized,
+    // Stream not initialized, waiting for the first registration request.
     kStreamNotInitialized,
-    // The RPC is down, needs a restart.
-    kStreamClosed,
-    // The RPC failed to be restarted.
-    kStreamFailedToRestart,
+    // The RPC is temporarily down, waiting for a restart.
+    kStreamTemporarilyDown,
+    // The RPC failed to be initialized.
+    kStreamFailedToInitialize,
     // The client is being shutdown.
     kShutdown,
   };

cpp/acs_agent_client_test.cc

@@ -580,8 +580,11 @@ TEST_F(AcsAgentClientTest, TestClientRecreatedAfterServerCancellation) {
 }
 
 TEST_F(AcsAgentClientTest, TestFailureToRegisterConnection) {
-  // Shutdown the client.
+  // Shutdown the client and test if IsClientDead() returns true.
+  (*client_)->Shutdown();
+  EXPECT_TRUE((*client_)->IsDead());
   *client_ = nullptr;
+
   // Make sure server does delay response.
   SetServerDelay(true, absl::Seconds(5));
   std::chrono::system_clock::time_point deadline =
@@ -620,5 +623,33 @@ TEST_F(AcsAgentClientTest, TestFailureToRegisterConnection) {
                   "Timeout waiting for promise to be set for message with id"));
 }
 
+TEST_F(AcsAgentClientTest, TestFailureToRestartClientAndClientIsDead) {
+  // Make sure server does delay response, then registration will fail to
+  // complete due to timeout.
+  SetServerDelay(true, absl::Seconds(5));
+  std::chrono::system_clock::time_point deadline =
+      std::chrono::system_clock::now() + std::chrono::seconds(1);
+  server_->Shutdown(deadline);
+  server_->Wait();
+  server_ = nullptr;
+  server_ = std::make_unique<FakeAcsAgentServer>(&service_);
+
+  // Create a channel to the server, and ensure server is connectable.
+  grpc::ChannelArguments channel_args;
+  // Keepalive settings
+  channel_args.SetInt(GRPC_ARG_KEEPALIVE_TIME_MS, 600 * 1000);  // 600 seconds
+  channel_args.SetInt(GRPC_ARG_KEEPALIVE_TIMEOUT_MS,
+                      100 * 1000);  // 100 seconds
+  std::shared_ptr<grpc::Channel> channel = grpc::CreateCustomChannel(
+      server_->GetServerAddress(), grpc::InsecureChannelCredentials(),
+      channel_args);
+  deadline = std::chrono::system_clock::now() + std::chrono::seconds(10);
+  ASSERT_TRUE(channel->WaitForConnected(deadline));
+
+  // Client will transition to kStreamFailedToInitialize eventually.
+  EXPECT_TRUE(WaitUntil([this]() { return (*client_)->IsDead(); },
+                        absl::Seconds(20), absl::Seconds(0.1)));
+}
+
 }  // namespace
 }  // namespace agent_communication