Update cvd stop to handle targeted instances

Author: SuperStrongDinosaur

base/cvd/cuttlefish/flag_parser/gflags_compat.cpp

@@ -422,6 +422,11 @@ Flag GflagsCompatFlag(const std::string& name,
   return GflagsCompatFlagImpl<std::string>(name, value, "");
 }
 
+Flag GflagsCompatFlag(const std::string& name,
+                      std::vector<unsigned>& value) {
+  return GflagsCompatFlagImpl<unsigned>(name, value, 0);
+}
+
 Flag GflagsCompatFlag(const std::string& name, std::vector<bool>& value,
                       const bool default_value) {
   return GflagsCompatFlagImpl(name, value, default_value);

base/cvd/cuttlefish/flag_parser/gflags_compat.h

@@ -40,6 +40,7 @@ Flag GflagsCompatFlag(const std::string& name, int32_t& value);
 Flag GflagsCompatFlag(const std::string& name, size_t& value);
 Flag GflagsCompatFlag(const std::string& name, bool& value);
 Flag GflagsCompatFlag(const std::string& name, std::vector<std::string>& value);
+Flag GflagsCompatFlag(const std::string& name, std::vector<unsigned>& value);
 Flag GflagsCompatFlag(const std::string& name, std::vector<bool>& value,
                       bool default_value);
 // Indicates when to assign std::nullopt to the std::optional backing the flag.

base/cvd/cuttlefish/host/commands/cvd/cli/commands/stop.cpp

@@ -34,6 +34,7 @@
 #include "cuttlefish/host/commands/cvd/cli/types.h"
 #include "cuttlefish/host/commands/cvd/cli/utils.h"
 #include "cuttlefish/host/commands/cvd/instances/instance_manager.h"
+#include "cuttlefish/host/commands/cvd/instances/local_instance.h"
 #include "cuttlefish/host/libs/metrics/device_metrics_orchestration.h"
 #include "cuttlefish/result/result.h"
 
@@ -94,10 +95,24 @@ Result<void> CvdStopCommandHandler::Handle(const CommandRequest& request) {
   if (flags.wait_for_launcher_secs > 0) {
     launcher_timeout.emplace(flags.wait_for_launcher_secs);
   }
+
+  std::vector<unsigned> instance_nums;
+  if (request.Selectors().instance_names) {
+    for (const auto& name : *request.Selectors().instance_names) {
+      std::vector<LocalInstance> instances = group.FindByInstanceName(name);
+      CF_EXPECTF(!instances.empty(), "Instance '{}' not found in group '{}'",
+                 name, group.GroupName());
+      for (const auto& inst : instances) {
+        instance_nums.push_back(inst.Id());
+      }
+    }
+  }
+
   Result<void> stop_outcome = instance_manager_.StopInstanceGroup(
       group, launcher_timeout,
       flags.clear_instance_dirs ? InstanceDirActionOnStop::Clear
-                                : InstanceDirActionOnStop::Keep);
+                                : InstanceDirActionOnStop::Keep,
+      instance_nums);
 
   GatherVmStopMetrics(group);
 

base/cvd/cuttlefish/host/commands/cvd/instances/BUILD.bazel

@@ -241,6 +241,7 @@ cf_cc_library(
         "//cuttlefish/result",
         "//libbase",
         "@abseil-cpp//absl/log",
+        "@abseil-cpp//absl/strings",
         "@fmt",
     ],
 )

base/cvd/cuttlefish/host/commands/cvd/instances/instance_manager.cpp

@@ -214,7 +214,21 @@ Result<void> InstanceManager::UpdateInstanceGroup(
 Result<void> InstanceManager::StopInstanceGroup(
     LocalInstanceGroup& group,
     std::optional<std::chrono::seconds> launcher_timeout,
-    InstanceDirActionOnStop instance_dir_action) {
+    InstanceDirActionOnStop instance_dir_action,
+    const std::vector<unsigned>& instance_nums) {
+  // Validate that the requested instances actually belong to this group
+  std::set<unsigned> valid_ids;
+  for (const auto& instance : group.Instances()) {
+    valid_ids.insert(instance.Id());
+  }
+  std::set<unsigned> ids_to_stop;
+  for (unsigned num : instance_nums) {
+    CF_EXPECTF(valid_ids.count(num) != 0,
+               "Instance ID {} does not belong to group '{}'", num,
+               group.GroupName());
+    ids_to_stop.insert(num);
+  }
+
   const auto stop_bin = CF_EXPECT(StopBin(group.HostArtifactsPath()));
   const auto stop_bin_path = group.HostArtifactsPath() + "/bin/" + stop_bin;
   int wait_for_launcher_secs = 0;
@@ -227,13 +241,23 @@ Result<void> InstanceManager::StopInstanceGroup(
       .wait_for_launcher_secs = wait_for_launcher_secs,
       .clear_runtime_dirs =
           instance_dir_action == InstanceDirActionOnStop::Clear,
+      .instance_nums = instance_nums,
   });
   if (!cmd_result.ok()) {
     LOG(WARNING)
         << "Warning: error stopping instances for dir \"" + group.HomeDir() +
                "\".\nThis can happen if instances are already stopped.\n";
+    return cmd_result;
+  }
+  if (instance_nums.empty()) {
+    group.SetAllStates(cvd::INSTANCE_STATE_STOPPED);
+  } else {
+    for (auto& instance : group.Instances()) {
+      if (ids_to_stop.count(instance.Id()) > 0) {
+        instance.SetState(cvd::INSTANCE_STATE_STOPPED);
+      }
+    }
   }
-  group.SetAllStates(cvd::INSTANCE_STATE_STOPPED);
   // TODO: b/471069557 - diagnose unused
   Result<void> unused = instance_db_.UpdateInstanceGroup(group);
   return {};

base/cvd/cuttlefish/host/commands/cvd/instances/instance_manager.h

@@ -85,7 +85,8 @@ class InstanceManager {
   Result<void> StopInstanceGroup(
       LocalInstanceGroup& group,
       std::optional<std::chrono::seconds> launcher_timeout,
-      InstanceDirActionOnStop instance_dir_action);
+      InstanceDirActionOnStop instance_dir_action,
+      const std::vector<unsigned>& instance_nums = {});
 
  private:
   struct InternalInstanceDesc {

base/cvd/cuttlefish/host/commands/cvd/instances/stop.cpp

@@ -32,6 +32,7 @@
 #include <fmt/core.h>
 #include <fmt/ranges.h>  // NOLINT(misc-include-cleaner): version difference
 #include "absl/log/log.h"
+#include "absl/strings/str_join.h"
 
 #include "cuttlefish/common/libs/utils/contains.h"
 #include "cuttlefish/common/libs/utils/files.h"
@@ -266,6 +267,9 @@ Result<void> RunStopCvd(StopCvdParams params) {
   if (params.clear_runtime_dirs) {
     args.push_back("--clear_instance_dirs=true");
   }
+  if (!params.instance_nums.empty()) {
+    args.push_back(fmt::format("--instance_nums={}", absl::StrJoin(params.instance_nums, ",")));
+  }
   Result<void> cmd_res = RunStopCvdCmd(stopper_path, stop_cvd_envs, args);
   if (cmd_res.ok()) {
     return {};

base/cvd/cuttlefish/host/commands/cvd/instances/stop.h

@@ -37,6 +37,7 @@ struct StopCvdParams{
   std::string home_dir;
   int wait_for_launcher_secs;
   bool clear_runtime_dirs;
+  std::vector<unsigned> instance_nums;
 };
 Result<void> RunStopCvd(StopCvdParams params);
 

base/cvd/cuttlefish/host/commands/stop/main.cc

@@ -159,12 +159,14 @@ int StopInstance(const CuttlefishConfig& config,
 struct FlagVaules {
   std::int32_t wait_for_launcher;
   bool clear_instance_dirs;
+  std::vector<unsigned> instance_nums;
   bool helpxml;
 };
 
 FlagVaules GetFlagValues(int argc, char** argv) {
   std::int32_t wait_for_launcher = 5;
   bool clear_instance_dirs = false;
+  std::vector<unsigned> instance_nums;
   std::vector<Flag> flags;
   flags.emplace_back(
       GflagsCompatFlag("wait_for_launcher", wait_for_launcher)
@@ -174,6 +176,9 @@ FlagVaules GetFlagValues(int argc, char** argv) {
       GflagsCompatFlag("clear_instance_dirs", clear_instance_dirs)
           .Help("If provided, deletes the instance dir after attempting to "
                 "stop each instance."));
+  flags.emplace_back(
+      GflagsCompatFlag("instance_nums", instance_nums)
+          .Help("Comma-separated list of instance numbers to stop."));
   flags.emplace_back(HelpFlag(flags));
   bool helpxml = false;
   flags.emplace_back(HelpXmlFlag(flags, std::cout, helpxml));
@@ -182,22 +187,48 @@ FlagVaules GetFlagValues(int argc, char** argv) {
   auto parse_res = ConsumeFlags(flags, args);
   CHECK(parse_res.ok() || helpxml) << "Could not process command line flags.";
 
-  return {wait_for_launcher, clear_instance_dirs, helpxml};
+  return {wait_for_launcher, clear_instance_dirs, instance_nums, helpxml};
 }
 
 int StopCvdMain(const std::int32_t wait_for_launcher,
-                const bool clear_instance_dirs) {
+                const bool clear_instance_dirs,
+                const std::vector<unsigned>& instance_nums) {
   auto config = CuttlefishConfig::Get();
   if (!config) {
     LOG(ERROR) << "Failed to obtain config object";
     return FallBackStop(FallbackDirs());
   }
 
+  std::set<unsigned> instance_ids(instance_nums.begin(), instance_nums.end());
+
   int exit_code = 0;
   auto instances = config->Instances();
+
+  if (!instance_ids.empty() && !instances.empty()) {
+    unsigned first_instance_id;
+    if (absl::SimpleAtoi(instances[0].id(), &first_instance_id)) {
+      if (instance_ids.count(first_instance_id) > 0 &&
+          instance_ids.size() < instances.size()) {
+        LOG(ERROR) << "Stopping the first instance (ID: " << first_instance_id
+                   << ") is not allowed while other instances are remaining. "
+                   << "Stop the entire group instead.";
+        return 1;
+      }
+    }
+  }
+
   std::vector<std::future<int>> exit_state_futures;
   exit_state_futures.reserve(instances.size());
   for (const auto& instance : instances) {
+    unsigned id;
+    if (!absl::SimpleAtoi(instance.id(), &id)) {
+      LOG(ERROR) << "Failed to parse instance ID \"" << instance.id() << "\" as unsigned";
+      exit_code |= 1;
+      continue;
+    }
+    if (!instance_ids.empty() && instance_ids.count(id) == 0) {
+      continue;
+    }
     std::future<int> exit_code_from_thread = std::async(
         std::launch::async,
         [&instance, &config, &wait_for_launcher,
@@ -223,7 +254,7 @@ int StopCvdMain(const std::int32_t wait_for_launcher,
 } // namespace cuttlefish
 
 int main(int argc, char** argv) {
-  const auto [wait_for_launcher, clear_instance_dirs, helpxml] =
+  const auto [wait_for_launcher, clear_instance_dirs, instance_nums, helpxml] =
       cuttlefish::GetFlagValues(argc, argv);
   cuttlefish::LogToStderr();
 
@@ -244,5 +275,5 @@ int main(int argc, char** argv) {
     cuttlefish::MetricsReceiver::LogMetricsVMStop();
   }
 
-  return cuttlefish::StopCvdMain(wait_for_launcher, clear_instance_dirs);
+  return cuttlefish::StopCvdMain(wait_for_launcher, clear_instance_dirs, instance_nums);
 }