Add agent state to DA node id response (#2140)

ydb-platform · Sep 30, 2024 · f4f684f · f4f684f
1 parent e9ec61a
commit f4f684f
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 3 deletions.
diff --git a/cloud/blockstore/libs/storage/disk_registry/disk_registry_actor_get_agent_node_id.cpp b/cloud/blockstore/libs/storage/disk_registry/disk_registry_actor_get_agent_node_id.cpp
@@ -36,7 +36,12 @@ void TDiskRegistryActor::HandleGetAgentNodeId(
         return;
     }
 
+    const auto* info = AgentRegInfo.FindPtr(agent->GetAgentId());
+    const bool connected = info && info->Connected;
+
     response->Record.SetNodeId(agent->GetNodeId());
+    response->Record.SetAgentState(agent->GetState());
+    response->Record.SetConnected(connected);
     NCloud::Reply(ctx, *ev, std::move(response));
 }
 

diff --git a/cloud/blockstore/libs/storage/protos/disk.proto b/cloud/blockstore/libs/storage/protos/disk.proto
@@ -1602,6 +1602,12 @@ message TGetAgentNodeIdResponse
 
     // Node that agent is running on.
     uint32 NodeId = 2;
+
+    // Agent state.
+    EAgentState AgentState = 3;
+
+    // Whether the agent is connected to the DR or not.
+    bool Connected = 4;
 }
 
 ////////////////////////////////////////////////////////////////////////////////

diff --git a/cloud/blockstore/libs/storage/service/service_actor_actions_get_disk_agent_node_id.cpp b/cloud/blockstore/libs/storage/service/service_actor_actions_get_disk_agent_node_id.cpp
@@ -95,11 +95,15 @@ void TGetDiskAgentNodeIdActor::HandleGetAgentNodeIdResponse(
         SetErrorProtoFlag(error, NCloud::NProto::EF_SILENT);
     }
 
-    auto response =
-        std::make_unique<TEvService::TEvExecuteActionResponse>(std::move(error));
+    auto response = std::make_unique<TEvService::TEvExecuteActionResponse>(
+        std::move(error));
+
+    google::protobuf::util::JsonPrintOptions options;
+    options.always_print_primitive_fields = true;
     google::protobuf::util::MessageToJsonString(
         msg->Record,
-        response->Record.MutableOutput());
+        response->Record.MutableOutput(),
+        options);
 
     ReplyAndDie(ctx, std::move(response));
 }

diff --git a/cloud/blockstore/tests/client/test_with_multiple_agents.py b/cloud/blockstore/tests/client/test_with_multiple_agents.py
@@ -227,6 +227,8 @@ def test_wait_dependent_disks_to_switch_node_timeout():
         agent_id = make_agent_id(0)
         node_id_response = json.loads(client.get_disk_agent_node_id(agent_id))
         assert node_id_response["NodeId"] > 50000
+        assert node_id_response["AgentState"] == "AGENT_STATE_ONLINE"
+        assert node_id_response["Connected"]
 
         # This should return immediately.
         wait_response = client.wait_dependent_disks_to_switch_node(