diff --git a/main.cc b/main.cc
index 48ea9f1..769850c 100755
--- a/main.cc
+++ b/main.cc
@@ -41,21 +41,7 @@
 
 gboolean UpdateBootFlags(void* arg) {
   UpdateAttempter* attempter = reinterpret_cast<UpdateAttempter*>(arg);
-  if (attempter->status() == UPDATE_STATUS_UPDATED_NEED_REBOOT) {
-    // Don't update the flags if there's an update that's just been applied and
-    // we're waiting for reboot because we may end up reverting the update.
-    return FALSE;  // Don't call this callback again.
-  }
-  if (attempter->status() != UPDATE_STATUS_IDLE) {
-    // To avoid races (e.g., setting a good kernel right after post-install but
-    // before the status changes), update the boot flag only if the attempter is
-    // idle.
-    return TRUE;  // Call this callback again.
-  }
-  // This is purely best effort. Failures should be logged by Subprocess.
-  int unused = 0;
-  vector<string> cmd(1, "/usr/sbin/chromeos-setgoodkernel");
-  Subprocess::SynchronousExec(cmd, &unused);
+  attempter->UpdateBootFlags();
   return FALSE;  // Don't call this callback again
 }
 
diff --git a/update_attempter.cc b/update_attempter.cc
index f56c0d1..e1ab0b1 100644
--- a/update_attempter.cc
+++ b/update_attempter.cc
@@ -28,6 +28,7 @@
 #include "update_engine/omaha_response_handler_action.h"
 #include "update_engine/postinstall_runner_action.h"
 #include "update_engine/prefs_interface.h"
+#include "update_engine/subprocess.h"
 #include "update_engine/update_check_scheduler.h"
 
 using base::TimeDelta;
@@ -114,7 +115,8 @@
       is_full_update_(false),
       proxy_manual_checks_(0),
       obeying_proxies_(true),
-      chrome_proxy_resolver_(dbus_iface) {
+      chrome_proxy_resolver_(dbus_iface),
+      updated_boot_flags_(false) {
   if (utils::FileExists(kUpdateCompletedMarker))
     status_ = UPDATE_STATUS_UPDATED_NEED_REBOOT;
 }
@@ -127,6 +129,7 @@
                              const std::string& omaha_url,
                              bool obey_proxies) {
   chrome_proxy_resolver_.Init();
+  UpdateBootFlags();  // Just in case we didn't do this yet.
   if (status_ == UPDATE_STATUS_UPDATED_NEED_REBOOT) {
     LOG(INFO) << "Not updating b/c we already updated and we're waiting for "
               << "reboot";
@@ -451,6 +454,19 @@
   return true;
 }
 
+void UpdateAttempter::UpdateBootFlags() {
+  if (updated_boot_flags_) {
+    LOG(INFO) << "Already updated boot flags. Skipping.";
+    return;
+  }
+  // This is purely best effort. Failures should be logged by Subprocess.
+  int unused = 0;
+  vector<string> cmd(1, "/usr/sbin/chromeos-setgoodkernel");
+  Subprocess::SynchronousExec(cmd, &unused);
+  updated_boot_flags_ = true;
+  return;
+}
+
 void UpdateAttempter::SetStatusAndNotify(UpdateStatus status) {
   status_ = status;
   if (update_check_scheduler_) {
diff --git a/update_attempter.h b/update_attempter.h
index e41cab9..8a0bd26 100644
--- a/update_attempter.h
+++ b/update_attempter.h
@@ -84,6 +84,10 @@
                  std::string* new_version,
                  int64_t* new_size);
 
+  // Runs setgoodkernel, whose responsibility it is to mark the currently
+  // booted partition has high priority/permanent/etc.
+  void UpdateBootFlags();
+
   UpdateStatus status() const { return status_; }
 
   int http_response_code() const { return http_response_code_; }
@@ -240,6 +244,9 @@
   DirectProxyResolver direct_proxy_resolver_;
   ChromeBrowserProxyResolver chrome_proxy_resolver_;
 
+  // True if UpdateBootFlags has already been called
+  bool updated_boot_flags_;
+
   DISALLOW_COPY_AND_ASSIGN(UpdateAttempter);
 };