* etc/grub.d/09_snappy:

  - Set snappy_stamp rather than relying on recordfail
   (since that sets timeout=-1).
  - Always explicitly set the default value if the stamp variable is not
    set for safety.
  - Added comments.
* partition/bootloader_grub.go:
  - MarkCurrentBootSuccessful(): Set "snappy_stamp=1".
* partition/bootloader_grub_test.go:
  - TestGrubMarkCurrentBootSuccessful(): New test for grub's version of MarkCurentBootSuccessful().
* partition/bootloader_uboot_test.go:
  - TestUbootMarkCurrentBootSuccessful(): new test for uboots's version of MarkCurentBootSuccessful().
* partition/partition_test.go:
  - Code to mock bootloaderUbootStampFile.

etc/grub.d/09_snappy

@@ -57,9 +57,6 @@ linux_entry_ext()
 
   echo "menuentry '$name' ${CLASS} \$menuentry_id_option 'gnulinux-simple-$boot_device_id' {" | sed "s/^/$submenu_indentation/"
 
-  if [ "$quick_boot" = 1 ]; then
-      echo "	recordfail" | sed "s/^/$submenu_indentation/"
-  fi
   if [ x$type != xrecovery ] ; then
       save_default_entry | grub_add_tab
   fi
@@ -376,10 +373,11 @@ handle_menu_entry "$other_root" "$other_label" "$other_mountpoint"
 
 # Toggle rootfs if previous boot failed.
 #
-# Since grub sets recordfail, if it is _already_ set when grub starts
-# and we're in try mode, the previous boot must have failed to unset
-# recordfail, hence toggle the rootfs.
+# Since grub sets snappy_stamp, if it is _already_ set when grub starts
+# and we're in try mode, the previous boot must have failed to unset it,
+# so toggle the rootfs.
 sed "s/^/$submenu_indentation/" << EOF
+  # set defaults
   if [ -z "\$snappy_mode" ]; then
       set snappy_mode=default
       save_env snappy_mode
@@ -390,16 +388,28 @@ sed "s/^/$submenu_indentation/" << EOF
   fi
 
   if [ "\$snappy_mode" = "try" ]; then
-    if [ "\$recordfail" = "1" ]; then
-        if [ "\$snappy_ab" = "a" ]; then
-            set default="$(make_name system-b)"
-            set snappy_ab=b
-        else
-            set snappy_ab=a
-            set default="$(make_name system-a)"
-        fi
-        save_env snappy_ab
-    fi
+    if [ "\$snappy_stamp" = "1" ]; then
+          # previous boot failed to clear snappy_stamp, so toggle rootfs
+          if [ "\$snappy_ab" = "a" ]; then
+              set default="$(make_name system-b)"
+              set snappy_ab=b
+          else
+              set snappy_ab=a
+              set default="$(make_name system-a)"
+          fi
+          save_env snappy_ab
+      else
+          # Note: don't use the standard recordfail variable since that forces
+          # the menu to be displayed and sets an infinite timeout if set!
+          set snappy_stamp=1
+          save_env snappy_stamp
+
+          if [ "\$snappy_ab" = "a" ]; then
+              set default="$(make_name system-a)"
+          else
+              set default="$(make_name system-b)"
+          fi
+      fi
   else
       if [ "\$snappy_ab" = "a" ]; then
           set default="$(make_name system-a)"

partition/bootloader_grub.go

@@ -32,12 +32,14 @@ const (
 
 	bootloaderGrubEnvCmdReal    = "/usr/bin/grub-editenv"
 	bootloaderGrubUpdateCmdReal = "/usr/sbin/update-grub"
+	bootloaderGrubStampVarReal  = "snappy_stamp"
 )
 
 // var to make it testable
 var (
 	bootloaderGrubDir        = bootloaderGrubDirReal
 	bootloaderGrubConfigFile = bootloaderGrubConfigFileReal
+	bootloaderGrubStampVar   = bootloaderGrubStampVarReal
 	bootloaderGrubEnvFile    = bootloaderGrubEnvFileReal
 
 	bootloaderGrubEnvCmd    = bootloaderGrubEnvCmdReal
@@ -128,6 +130,10 @@ func (g *grub) GetOtherRootFSName() string {
 }
 
 func (g *grub) MarkCurrentBootSuccessful() (err error) {
+	// disable the variable set by grub on boot.
+	if err := g.setBootVar(bootloaderGrubStampVar, "0"); err != nil {
+		return err
+	}
 	return g.setBootVar(bootloaderBootmodeVar, bootloaderBootmodeSuccess)
 }
 

partition/bootloader_grub_test.go

@@ -131,3 +131,26 @@ func (s *PartitionTestSuite) TestGetBootloaderWithGrub(c *C) {
 	c.Assert(err, IsNil)
 	c.Assert(bootloader.Name(), Equals, bootloaderNameGrub)
 }
+
+func (s *PartitionTestSuite) TestGrubMarkCurrentBootSuccessful(c *C) {
+	s.makeFakeGrubEnv(c)
+	allCommands = []singleCommand{}
+
+	partition := New()
+	g := newGrub(partition)
+	c.Assert(g, NotNil)
+	err := g.MarkCurrentBootSuccessful()
+	c.Assert(err, IsNil)
+
+	// this is always called
+	mp := singleCommand{"/bin/mountpoint", "/writable/cache/system"}
+	c.Assert(allCommands[0], DeepEquals, mp)
+
+	expectedGrubSet := singleCommand{bootloaderGrubEnvCmd, bootloaderGrubEnvFile, "set", "snappy_stamp=0"}
+
+	c.Assert(allCommands[1], DeepEquals, expectedGrubSet)
+
+	expectedGrubSet2 := singleCommand{bootloaderGrubEnvCmd, bootloaderGrubEnvFile, "set", "snappy_mode=default"}
+
+	c.Assert(allCommands[2], DeepEquals, expectedGrubSet2)
+}

partition/bootloader_uboot_test.go

@@ -247,3 +247,37 @@ bootloader u-boot
 
 	c.Assert(bootloader.HandleAssets(), NotNil)
 }
+
+func (s *PartitionTestSuite) TestUbootMarkCurrentBootSuccessful(c *C) {
+	s.makeFakeUbootEnv(c)
+
+	// create stamp file
+	err := ioutil.WriteFile(bootloaderUbootStampFile, []byte(""), 0640)
+	c.Assert(err, IsNil)
+	c.Assert(helpers.FileExists(bootloaderUbootStampFile), Equals, true)
+
+	partition := New()
+	u := newUboot(partition)
+	c.Assert(u, NotNil)
+
+	// enter try mode
+	err = u.ToggleRootFS()
+	c.Assert(err, IsNil)
+
+	c.Assert(helpers.FileExists(bootloaderUbootEnvFile), Equals, true)
+	bytes, err := ioutil.ReadFile(bootloaderUbootEnvFile)
+	c.Assert(err, IsNil)
+	c.Assert(strings.Contains(string(bytes), "snappy_mode=try"), Equals, true)
+	c.Assert(strings.Contains(string(bytes), "snappy_mode=default"), Equals, false)
+
+	err = u.MarkCurrentBootSuccessful()
+	c.Assert(err, IsNil)
+
+	c.Assert(helpers.FileExists(bootloaderUbootStampFile), Equals, false)
+	c.Assert(helpers.FileExists(bootloaderUbootEnvFile), Equals, true)
+
+	bytes, err = ioutil.ReadFile(bootloaderUbootEnvFile)
+	c.Assert(err, IsNil)
+	c.Assert(strings.Contains(string(bytes), "snappy_mode=try"), Equals, false)
+	c.Assert(strings.Contains(string(bytes), "snappy_mode=default"), Equals, true)
+}

partition/partition_test.go

@@ -60,6 +60,7 @@ func (s *PartitionTestSuite) SetUpTest(c *C) {
 	bootloaderUbootDir = filepath.Join(s.tempdir, "boot", "uboot")
 	bootloaderUbootConfigFile = filepath.Join(bootloaderUbootDir, "uEnv.txt")
 	bootloaderUbootEnvFile = filepath.Join(bootloaderUbootDir, "uEnv.txt")
+	bootloaderUbootStampFile = filepath.Join(bootloaderUbootDir, "snappy-stamp.txt")
 
 	c.Assert(mounts, DeepEquals, mountEntryArray(nil))
 }
@@ -81,6 +82,7 @@ func (s *PartitionTestSuite) TearDownTest(c *C) {
 	bootloaderUbootDir = bootloaderUbootDirReal
 	bootloaderUbootConfigFile = bootloaderUbootConfigFileReal
 	bootloaderUbootEnvFile = bootloaderUbootEnvFileReal
+	bootloaderUbootStampFile = bootloaderUbootStampFileReal
 
 	c.Assert(mounts, DeepEquals, mountEntryArray(nil))
 }