Walle API adapter replace specific devices (#18257)

SammyVimes · web-flow · commit d11ac2887a77 · 2025-05-15T13:40:48.000+04:00
(cherry-picked from 53455f3)
diff --git a/ydb/core/cms/walle_api_handler.cpp b/ydb/core/cms/walle_api_handler.cpp
@@ -80,6 +80,10 @@ class TWalleCrateTaskHandler : public TActorBootstrapped<TWalleCrateTaskHandler>
         auto &hosts = map["hosts"].GetArray();
         for (auto &host : hosts)
             *request->Record.AddHosts() = host.GetString();
+        
+        auto &devices = map["devices"].GetArray();
+        for (auto &device : devices)
+            *request->Record.AddDevices() = device.GetString();
 
         const auto &params = http.GetParams();
         if (params.contains("dry_run"))
@@ -109,9 +113,15 @@ class TWalleCrateTaskHandler : public TActorBootstrapped<TWalleCrateTaskHandler>
         NJson::TJsonValue hosts(NJson::JSON_ARRAY);
         for (auto &host : resp.GetHosts())
             hosts.AppendValue(host);
+
+        NJson::TJsonValue devices(NJson::JSON_ARRAY);
+        for (auto &device : resp.GetDevices())
+            devices.AppendValue(device);
+
         NJson::TJsonValue map;
         map.InsertValue("id", resp.GetTaskId());
         map.InsertValue("hosts", hosts);
+        map.InsertValue("devices", devices);
         map.InsertValue("status", status);
         map.InsertValue("message", resp.GetStatus().GetReason());
 
@@ -147,12 +157,19 @@ class TWalleCrateTaskHandler : public TActorBootstrapped<TWalleCrateTaskHandler>
     void Reply(const TWalleListTasksResponse &resp, const TActorContext &ctx) {
         NJson::TJsonValue tasks(NJson::JSON_ARRAY);
         for (auto &task : resp.GetTasks()) {
+
             NJson::TJsonValue hosts(NJson::JSON_ARRAY);
             for (auto &host : task.GetHosts())
                 hosts.AppendValue(host);
+
+            NJson::TJsonValue devices(NJson::JSON_ARRAY);
+            for (auto &device : task.GetDevices())
+                devices.AppendValue(device);
+
             NJson::TJsonValue map;
             map.InsertValue("id", task.GetTaskId());
             map.InsertValue("hosts", hosts);
+            map.InsertValue("devices", devices);
             map.InsertValue("status", task.GetStatus());
 
             tasks.AppendValue(map);
@@ -199,9 +216,15 @@ class TWalleCrateTaskHandler : public TActorBootstrapped<TWalleCrateTaskHandler>
         NJson::TJsonValue hosts(NJson::JSON_ARRAY);
         for (auto &host : resp.GetTask().GetHosts())
             hosts.AppendValue(host);
+
+        NJson::TJsonValue devices(NJson::JSON_ARRAY);
+        for (auto &device : resp.GetTask().GetDevices())
+            devices.AppendValue(device);
+
         NJson::TJsonValue map;
         map.InsertValue("id", resp.GetTask().GetTaskId());
         map.InsertValue("hosts", hosts);
+        map.InsertValue("devices", devices);
         map.InsertValue("status", status);
         map.InsertValue("message", resp.GetStatus().GetReason());
 
diff --git a/ydb/core/cms/walle_check_task_adapter.cpp b/ydb/core/cms/walle_check_task_adapter.cpp
@@ -41,8 +41,11 @@ class TWalleCheckTaskAdapter : public TActorBootstrapped<TWalleCheckTaskAdapter>
         if (State->ScheduledRequests.contains(task.RequestId)) {
             auto &req = State->ScheduledRequests.find(task.RequestId)->second;
 
-            for (auto &action : req.Request.GetActions())
+            for (auto &action : req.Request.GetActions()) {
                 *info.AddHosts() = action.GetHost();
+                for (auto &device : action.GetDevices())
+                    *info.AddDevices() = device;
+            }
 
             TAutoPtr<TEvCms::TEvCheckRequest> event = new TEvCms::TEvCheckRequest;
             event->Record.SetUser(WALLE_CMS_USER);
@@ -53,8 +56,13 @@ class TWalleCheckTaskAdapter : public TActorBootstrapped<TWalleCheckTaskAdapter>
             Become(&TThis::StateWork, ctx, TDuration::Seconds(10), new TEvents::TEvWakeup());
         } else {
             for (auto &id : task.Permissions) {
-                if (State->Permissions.contains(id))
-                    *info.AddHosts() = State->Permissions.find(id)->second.Action.GetHost();
+                if (State->Permissions.contains(id)) {
+                    const auto &action = State->Permissions.find(id)->second.Action;
+                    *info.AddHosts() = action.GetHost();
+
+                    for (auto &device : action.GetDevices())
+                        *info.AddDevices() = device;
+                }
             }
 
             if (!info.HostsSize()) {
diff --git a/ydb/core/cms/walle_create_task_adapter.cpp b/ydb/core/cms/walle_create_task_adapter.cpp
@@ -65,6 +65,7 @@ class TWalleCreateTaskAdapter : public TActorBootstrapped<TWalleCreateTaskAdapte
         resp->Record.MutableStatus()->SetReason(err);
         resp->Record.SetTaskId(rec.GetTaskId());
         resp->Record.MutableHosts()->CopyFrom(rec.GetHosts());
+        resp->Record.MutableDevices()->CopyFrom(rec.GetDevices());
         ReplyAndDie(resp.Release(), ctx);
     }
 
@@ -81,6 +82,7 @@ class TWalleCreateTaskAdapter : public TActorBootstrapped<TWalleCreateTaskAdapte
         Response->Record.MutableStatus()->CopyFrom(rec.GetStatus());
         Response->Record.SetTaskId(RequestEvent->Get()->Record.GetTaskId());
         Response->Record.MutableHosts()->CopyFrom(RequestEvent->Get()->Record.GetHosts());
+        Response->Record.MutableDevices()->CopyFrom(RequestEvent->Get()->Record.GetDevices());
 
         // In case of success or scheduled request we have to store
         // task information.
@@ -101,6 +103,52 @@ class TWalleCreateTaskAdapter : public TActorBootstrapped<TWalleCreateTaskAdapte
         ReplyAndDie(Response, ctx);
     }
 
+    void HandleReplaceDevicesAction(const TClusterInfoPtr &cluster, 
+                                    TAutoPtr<TEvCms::TEvPermissionRequest> &request, 
+                                    ui64 duration, 
+                                    const TActorContext &ctx) {
+        auto &task = RequestEvent->Get()->Record;
+
+        if (task.DevicesSize() > 0) {
+            // Replace specified devices on the specified host.
+            if (task.HostsSize() != 1) {
+                ReplyWithErrorAndDie(TStatus::WRONG_REQUEST, "Exactly one host must be specified if \"devices\" are set", ctx);
+                return;
+            }
+            TString host = task.get_idx_hosts(0);
+            auto &action = *request->Record.AddActions();
+            action.SetHost(host);
+            action.SetType(TAction::REPLACE_DEVICES);
+            action.SetDuration(duration);
+
+            for (const auto& device : task.GetDevices())
+                *action.AddDevices() = device;
+        } else {
+            // Replace all devices on all provided hosts.
+            for (auto &host : task.GetHosts()) {
+                auto &action = *request->Record.AddActions();
+                action.SetHost(host);
+                action.SetType(TAction::REPLACE_DEVICES);
+                action.SetDuration(duration);
+                for (const auto node : cluster->HostNodes(host)) {
+                    for (auto &pdiskId : node->PDisks)
+                        *action.AddDevices() = cluster->PDisk(pdiskId).GetDeviceName();
+                }
+            }
+        }
+    }
+
+    void HandleGenericAction(TAutoPtr<TEvCms::TEvPermissionRequest> &cmsRequest, TAction::EType actionType, ui64 duration) {
+        auto &task = RequestEvent->Get()->Record;
+
+        for (auto &host : task.GetHosts()) {
+            auto &action = *cmsRequest->Record.AddActions();
+            action.SetHost(host);
+            action.SetType(actionType);
+            action.SetDuration(duration);
+        }
+    }
+
     void Handle(TEvCms::TEvGetClusterInfoResponse::TPtr &ev, const TActorContext &ctx) {
         if (ev->Get()->Info->IsOutdated()) {
             ReplyWithErrorAndDie(TStatus::ERROR_TEMP, "Cannot collect cluster info", ctx);
@@ -135,23 +183,20 @@ class TWalleCreateTaskAdapter : public TActorBootstrapped<TWalleCreateTaskAdapte
             TAutoPtr<TEvCms::TEvWalleCreateTaskResponse> resp = new TEvCms::TEvWalleCreateTaskResponse;
             resp->Record.SetTaskId(task.GetTaskId());
             resp->Record.MutableHosts()->CopyFrom(task.GetHosts());
+            resp->Record.MutableDevices()->CopyFrom(task.GetDevices());
             resp->Record.MutableStatus()->SetCode(TStatus::OK);
             ReplyAndDie(resp.Release(), ctx);
             return;
         } else {
-            for (auto &host : task.GetHosts()) {
-                auto &action = *request->Record.AddActions();
-                action.SetHost(host);
-                action.SetType(*it->second);
-                // We always use infinite duration.
-                // Wall-E MUST delete processed tasks.
-                action.SetDuration(TDuration::Max().GetValue());
-                if (action.GetType() == TAction::REPLACE_DEVICES) {
-                    for (const auto node : cluster->HostNodes(host)) {
-                        for (auto &pdiskId : node->PDisks)
-                            *action.AddDevices() = cluster->PDisk(pdiskId).GetDeviceName();
-                    }
-                }
+            // We always use infinite duration.
+            // Wall-E MUST delete processed tasks.
+            ui64 duration = TDuration::Max().GetValue();
+
+            TAction::EType actionType = *it->second;
+            if (actionType == TAction::REPLACE_DEVICES) {
+                HandleReplaceDevicesAction(cluster, request, duration, ctx);
+            } else {
+                HandleGenericAction(request, actionType, duration);
             }
         }
 
diff --git a/ydb/core/cms/walle_list_tasks_adapter.cpp b/ydb/core/cms/walle_list_tasks_adapter.cpp
@@ -35,13 +35,21 @@ class TWalleListTasksAdapter : public TActorBootstrapped<TWalleListTasksAdapter>
             info.SetTaskId(task.TaskId);
             if (State->ScheduledRequests.contains(task.RequestId)) {
                 auto &req = State->ScheduledRequests.find(task.RequestId)->second;
-                for (auto &action : req.Request.GetActions())
+                for (auto &action : req.Request.GetActions()) {
                     *info.AddHosts() = action.GetHost();
+                    for (auto &device : action.GetDevices())
+                        *info.AddDevices() = device;
+                }
                 info.SetStatus("in-process");
             } else {
                 for (auto &id : task.Permissions) {
-                    if (State->Permissions.contains(id))
-                        *info.AddHosts() = State->Permissions.find(id)->second.Action.GetHost();
+                    if (State->Permissions.contains(id)) {
+                        const auto &action = State->Permissions.find(id)->second.Action;
+                        *info.AddHosts() = action.GetHost();
+
+                        for (auto &device : action.GetDevices())
+                            *info.AddDevices() = device;
+                    }
                 }
                 info.SetStatus("ok");
             }
diff --git a/ydb/core/protos/cms.proto b/ydb/core/protos/cms.proto
@@ -355,12 +355,14 @@ message TWalleCreateTaskRequest {
   optional string Action = 4;
   repeated string Hosts = 5;
   optional bool   DryRun = 6;
+  repeated string Devices = 7;
 }
 
 message TWalleCreateTaskResponse {
   optional TStatus Status = 1;
   optional string  TaskId = 2;
   repeated string  Hosts = 3;
+  repeated string  Devices = 4;
 }
 
 message TWalleListTasksRequest {
@@ -370,6 +372,7 @@ message TWalleTaskInfo {
   optional string TaskId = 1;
   repeated string Hosts = 2;
   optional string Status = 3;
+  repeated string Devices = 4;
 }
 
 message TWalleListTasksResponse {

Original file line number	Diff line number	Diff line change
`@@ -355,12 +355,14 @@ message TWalleCreateTaskRequest {`
`355`	`355`	`optional string Action = 4;`
`356`	`356`	`repeated string Hosts = 5;`
`357`	`357`	`optional bool DryRun = 6;`
	`358`	`+ repeated string Devices = 7;`
`358`	`359`	`}`
`359`	`360`
`360`	`361`	`message TWalleCreateTaskResponse {`
`361`	`362`	`optional TStatus Status = 1;`
`362`	`363`	`optional string TaskId = 2;`
`363`	`364`	`repeated string Hosts = 3;`
	`365`	`+ repeated string Devices = 4;`
`364`	`366`	`}`
`365`	`367`
`366`	`368`	`message TWalleListTasksRequest {`
`@@ -370,6 +372,7 @@ message TWalleTaskInfo {`
`370`	`372`	`optional string TaskId = 1;`
`371`	`373`	`repeated string Hosts = 2;`
`372`	`374`	`optional string Status = 3;`
	`375`	`+ repeated string Devices = 4;`
`373`	`376`	`}`
`374`	`377`
`375`	`378`	`message TWalleListTasksResponse {`