Add prioritty for canary rollout

rcaelers · Sep 12, 2024 · 4521096 · 4521096
1 parent e59c5a9
commit 4521096
Show file tree

Hide file tree

Showing 7 changed files with 143 additions and 12 deletions.
diff --git a/include/unfold/Unfold.hh b/include/unfold/Unfold.hh
@@ -91,6 +91,8 @@ namespace unfold
     virtual outcome::std_result<void> set_current_version(const std::string &version) = 0;
     virtual outcome::std_result<void> set_allowed_channels(const std::vector<std::string> &channels) = 0;
     virtual outcome::std_result<void> set_signature_verification_key(const std::string &key) = 0;
+    virtual outcome::std_result<void> set_priority(int prio) = 0;
+    virtual void unset_priority() = 0;
     virtual void set_certificate(const std::string &cert) = 0;
     virtual void set_periodic_update_check_enabled(bool enabled) = 0;
     virtual void set_periodic_update_check_interval(std::chrono::seconds interval) = 0;

diff --git a/src/Settings.cc b/src/Settings.cc
@@ -28,10 +28,11 @@
 
 namespace
 {
-  constexpr const char *last_update_check_time = "LastUpdateCheckTime";
-  constexpr const char *skip_version = "SkipVersion";
-  constexpr const char *periodic_update_check_interval = "PeriodicUpdateCheckInterval";
-  constexpr const char *periodic_update_check_enabled = "PeriodicUpdateCheckEnabled";
+  constexpr const char *setting_last_update_check_time = "LastUpdateCheckTime";
+  constexpr const char *setting_skip_version = "SkipVersion";
+  constexpr const char *setting_periodic_update_check_interval = "PeriodicUpdateCheckInterval";
+  constexpr const char *setting_periodic_update_check_enabled = "PeriodicUpdateCheckEnabled";
+  constexpr const char *setting_priority = "Priority";
 } // namespace
 
 Settings::Settings(std::shared_ptr<SettingsStorage> storage)
@@ -42,7 +43,7 @@ Settings::Settings(std::shared_ptr<SettingsStorage> storage)
 std::optional<std::chrono::system_clock::time_point>
 Settings::get_last_update_check_time() const
 {
-  auto l = storage->get_value(last_update_check_time, SettingType::Int64);
+  auto l = storage->get_value(setting_last_update_check_time, SettingType::Int64);
   if (l)
     {
       return std::chrono::system_clock::time_point(std::chrono::microseconds(std::get<int64_t>(l.value())));
@@ -53,7 +54,7 @@ Settings::get_last_update_check_time() const
 void
 Settings::set_last_update_check_time(std::chrono::system_clock::time_point t)
 {
-  auto rc = storage->set_value(last_update_check_time,
+  auto rc = storage->set_value(setting_last_update_check_time,
                                static_cast<int64_t>(
                                  std::chrono::duration_cast<std::chrono::microseconds>(t.time_since_epoch()).count()));
   if (!rc)
@@ -65,7 +66,7 @@ Settings::set_last_update_check_time(std::chrono::system_clock::time_point t)
 std::string
 Settings::get_skip_version() const
 {
-  auto version = storage->get_value(skip_version, SettingType::String);
+  auto version = storage->get_value(setting_skip_version, SettingType::String);
   if (version)
     {
       return std::get<std::string>(version.value());
@@ -76,7 +77,7 @@ Settings::get_skip_version() const
 void
 Settings::set_skip_version(std::string version)
 {
-  auto rc = storage->set_value(skip_version, version);
+  auto rc = storage->set_value(setting_skip_version, version);
   if (!rc)
     {
       spdlog::error("Failed to set skip version");
@@ -86,7 +87,7 @@ Settings::set_skip_version(std::string version)
 std::chrono::seconds
 Settings::get_periodic_update_check_interval() const
 {
-  auto interval = storage->get_value(periodic_update_check_interval, SettingType::Int64);
+  auto interval = storage->get_value(setting_periodic_update_check_interval, SettingType::Int64);
   if (interval)
     {
       return std::chrono::seconds(std::get<int64_t>(interval.value()));
@@ -97,7 +98,7 @@ Settings::get_periodic_update_check_interval() const
 void
 Settings::set_periodic_update_check_interval(std::chrono::seconds interval)
 {
-  auto rc = storage->set_value(periodic_update_check_interval, static_cast<int64_t>(interval.count()));
+  auto rc = storage->set_value(setting_periodic_update_check_interval, static_cast<int64_t>(interval.count()));
   if (!rc)
     {
       spdlog::error("Failed to set periodic update interval");
@@ -107,7 +108,7 @@ Settings::set_periodic_update_check_interval(std::chrono::seconds interval)
 bool
 Settings::get_periodic_update_check_enabled() const
 {
-  auto enabled = storage->get_value(periodic_update_check_enabled, SettingType::Boolean);
+  auto enabled = storage->get_value(setting_periodic_update_check_enabled, SettingType::Boolean);
   if (enabled)
     {
       return std::get<bool>(enabled.value());
@@ -118,9 +119,30 @@ Settings::get_periodic_update_check_enabled() const
 void
 Settings::set_periodic_update_check_enabled(bool enabled)
 {
-  auto rc = storage->set_value(periodic_update_check_enabled, enabled);
+  auto rc = storage->set_value(setting_periodic_update_check_enabled, enabled);
   if (!rc)
     {
       spdlog::error("Failed to set periodic update check enabled");
     }
 }
+
+int
+Settings::get_priority() const
+{
+  auto priority = storage->get_value(setting_priority, SettingType::Int32);
+  if (priority)
+    {
+      return std::get<int32_t>(priority.value());
+    }
+  return 0;
+}
+
+void
+Settings::set_priority(int priority)
+{
+  auto rc = storage->set_value(setting_priority, static_cast<int32_t>(priority));
+  if (!rc)
+    {
+      spdlog::error("Failed to set priority");
+    }
+}
diff --git a/src/Settings.hh b/src/Settings.hh
@@ -44,6 +44,9 @@ public:
   bool get_periodic_update_check_enabled() const;
   void set_periodic_update_check_enabled(bool enabled);
 
+  int get_priority() const;
+  void set_priority(int priority);
+
 private:
   std::shared_ptr<SettingsStorage> storage;
   std::shared_ptr<spdlog::logger> logger{unfold::utils::Logging::create("unfold:settings")};

diff --git a/src/UpgradeControl.cc b/src/UpgradeControl.cc
@@ -23,6 +23,7 @@
 #include <chrono>
 #include <exception>
 #include <memory>
+#include <random>
 #include <utility>
 #include <fstream>
 #include <ranges>
@@ -71,6 +72,7 @@ UpgradeControl::UpgradeControl(std::shared_ptr<Platform> platform, unfold::coro:
   http->options().set_max_redirects(5);
   http->options().set_follow_redirects(true);
   init_periodic_update_check();
+  init_priority();
 }
 
 UpgradeControl::UpgradeControl(std::shared_ptr<Platform> platform,
@@ -91,6 +93,7 @@ UpgradeControl::UpgradeControl(std::shared_ptr<Platform> platform,
   , check_timer(io_context.get_io_context())
 {
   init_periodic_update_check();
+  init_priority();
 }
 
 outcome::std_result<void>
@@ -129,6 +132,23 @@ UpgradeControl::set_certificate(const std::string &cert)
   http->options().add_ca_cert(cert);
 }
 
+outcome::std_result<void>
+UpgradeControl::set_priority(int priority)
+{
+  if (priority < 1 || priority > 100)
+    {
+      return outcome::failure(unfold::UnfoldErrc::InvalidArgument);
+    }
+  custom_priority = priority;
+  return outcome::success();
+}
+
+void
+UpgradeControl::unset_priority()
+{
+  custom_priority.reset();
+}
+
 void
 UpgradeControl::set_periodic_update_check_enabled(bool enabled)
 {
@@ -342,6 +362,31 @@ UpgradeControl::init_periodic_update_check()
   });
 }
 
+void
+UpgradeControl::init_priority()
+{
+  int priority = state->get_priority();
+  logger->info("current priority: {}", priority);
+  if (priority == 0)
+    {
+      std::random_device rd;
+      std::mt19937 gen(rd());
+      std::uniform_int_distribution<> dis(1, 100);
+      priority = dis(gen);
+      state->set_priority(priority);
+    }
+}
+
+int
+UpgradeControl::get_priority() const
+{
+  if (custom_priority)
+    {
+      return *custom_priority;
+    }
+  return state->get_priority();
+}
+
 void
 UpgradeControl::set_proxy(unfold::ProxyType proxy)
 {

diff --git a/src/UpgradeControl.hh b/src/UpgradeControl.hh
@@ -22,6 +22,7 @@
 #define UPGRADE_CONTROL_HH
 
 #include <memory>
+#include <optional>
 #include <string>
 #include <filesystem>
 #include <chrono>
@@ -60,6 +61,8 @@ public:
   outcome::std_result<void> set_current_version(const std::string &version) override;
   outcome::std_result<void> set_allowed_channels(const std::vector<std::string> &channels) override;
   outcome::std_result<void> set_signature_verification_key(const std::string &key) override;
+  outcome::std_result<void> set_priority(int prio) override;
+  void unset_priority() override;
   void set_certificate(const std::string &cert) override;
   void set_periodic_update_check_enabled(bool enabled) override;
   void set_periodic_update_check_interval(std::chrono::seconds interval) override;
@@ -81,12 +84,14 @@ public:
   void set_proxy(unfold::ProxyType proxy) override;
   void set_custom_proxy(const std::string &proxy) override;
 
+  int get_priority() const;
   boost::asio::awaitable<outcome::std_result<void>> check_for_update_and_notify(bool ignore_skip_version);
 
 private:
   void init_periodic_update_check();
   void update_last_update_check_time();
   void update_check_timer();
+  void init_priority();
 
 private:
   std::shared_ptr<Platform> platform;
@@ -101,6 +106,7 @@ private:
   unfold::utils::OneShotTimer check_timer;
   std::chrono::seconds periodic_update_check_interval{60 * 60 * 24};
   bool periodic_update_check_enabled{false};
+  std::optional<int> custom_priority;
 
   update_available_callback_t update_available_callback;
   update_status_callback_t update_status_callback;

diff --git a/test/SettingsTests.cc b/test/SettingsTests.cc
@@ -192,6 +192,37 @@ BOOST_AUTO_TEST_CASE(settings_set_skip_verion)
   settings.set_skip_version("2.30.75");
 }
 
+BOOST_AUTO_TEST_CASE(settings_get_priority)
+{
+  auto storage = std::make_shared<SettingsStorageMock>();
+
+  EXPECT_CALL(*storage, set_prefix(_)).Times(0);
+  EXPECT_CALL(*storage, remove_key(_)).Times(0);
+  EXPECT_CALL(*storage, set_value(_, _)).Times(0);
+  EXPECT_CALL(*storage, get_value("Priority", SettingType::Int32)).Times(1).WillOnce(Return(outcome::success(10)));
+
+  Settings settings(storage);
+  auto rc = settings.get_priority();
+  BOOST_CHECK_EQUAL(rc, 10);
+}
+
+BOOST_AUTO_TEST_CASE(settings_set_priority)
+{
+  auto storage = std::make_shared<SettingsStorageMock>();
+
+  EXPECT_CALL(*storage, set_prefix(_)).Times(0);
+  EXPECT_CALL(*storage, remove_key(_)).Times(0);
+  EXPECT_CALL(*storage, get_value(_, _)).Times(0);
+  EXPECT_CALL(*storage, set_value("Priority", SettingValue{10})).Times(1).WillOnce(Return(outcome::success()));
+  EXPECT_CALL(*storage, set_value("Priority", SettingValue{101}))
+    .Times(1)
+    .WillOnce(Return(outcome::failure(UnfoldInternalErrc::InvalidSetting)));
+
+  Settings settings(storage);
+  settings.set_priority(10);
+  settings.set_priority(101);
+}
+
 namespace
 {
   template<class T>

diff --git a/test/UpgradeControlTests.cc b/test/UpgradeControlTests.cc
@@ -87,6 +87,10 @@ struct MockedTestFixture
     checker = std::make_shared<CheckerMock>();
     installer = std::make_shared<InstallerMock>();
 
+    EXPECT_CALL(*storage, get_value("Priority", SettingType::Int32)).Times(1).WillOnce(Return(outcome::success(0)));
+    EXPECT_CALL(*storage, set_value("Priority", _)).Times(1).WillRepeatedly(Return(outcome::success()));
+    ;
+
     control = std::make_shared<UpgradeControl>(platform, http, verifier, storage, installer, checker, io_context);
   }
 
@@ -796,4 +800,22 @@ BOOST_AUTO_TEST_CASE(upgrade_control_proxy)
   BOOST_CHECK_EQUAL(http->options().get_proxy(), unfold::http::Options::ProxyType::None);
 }
 
+BOOST_AUTO_TEST_CASE(upgrade_control_priority)
+{
+  EXPECT_CALL(*storage, get_value("Priority", SettingType::Int32)).Times(AtLeast(1)).WillRepeatedly(Return(outcome::success(10)));
+
+  BOOST_CHECK_EQUAL(control->get_priority(), 10);
+  auto ret = control->set_priority(5);
+  BOOST_CHECK_EQUAL(ret.has_error(), false);
+  BOOST_CHECK_EQUAL(control->get_priority(), 5);
+  control->unset_priority();
+  BOOST_CHECK_EQUAL(control->get_priority(), 10);
+  ret = control->set_priority(101);
+  BOOST_CHECK_EQUAL(ret.has_error(), true);
+  ret = control->set_priority(0);
+  BOOST_CHECK_EQUAL(ret.has_error(), true);
+  ret = control->set_priority(-1);
+  BOOST_CHECK_EQUAL(ret.has_error(), true);
+}
+
 BOOST_AUTO_TEST_SUITE_END()