uniter, state: handle upgrades better

worker/uniter/uniter.go

 package uniter
 
 import (
         "errors"
         "fmt"
         corecharm "launchpad.net/juju-core/charm"
         "launchpad.net/juju-core/charm/hooks"
         "launchpad.net/juju-core/cmd"
         "launchpad.net/juju-core/environs/agent"
         "launchpad.net/juju-core/log"
@@ skipping 163 matching lines @@
         var hi *hook.Info
         if u.s != nil && (u.s.Op == RunHook || u.s.Op == Upgrade) {
                 // If this upgrade interrupts a RunHook, we need to preserve the hook
                 // info so that we can return to the appropriate error state. However,
                 // if we're resuming (or have force-interrupted) an Upgrade, we also
                 // need to preserve whatever hook info was preserved when we initially
                 // started upgrading, to ensure we still return to the correct state.
                 hi = u.s.Hook
         }
         if u.s == nil || u.s.OpStep != Done {
+                // Set the new charm URL - this returns once it's done and stored.
+                if err := u.f.SetCharm(curl); err != nil {
+                        return err
+                }
+                // Refresh to get the new URL.
+                if err := u.unit.Refresh(); err != nil {
+                        return err
+                }
+
                 log.Printf("worker/uniter: fetching charm %q", curl)
                 sch, err := u.st.Charm(curl)
                 if err != nil {
                         return err
                 }
                 bun, err := u.bundles.Read(sch, u.tomb.Dying())
                 if err != nil {
                         return err
                 }
                 if err = u.deployer.Stage(bun, curl); err != nil {
                         return err
                 }
                 log.Printf("worker/uniter: deploying charm %q", curl)
                 if err = u.writeState(reason, Pending, hi, curl); err != nil {
                         return err
                 }
                 if err = u.deployer.Deploy(u.charm); err != nil {
                         return err
                 }
                 if err = u.writeState(reason, Done, hi, curl); err != nil {
                         return err
                 }
         }
         log.Printf("worker/uniter: charm %q is deployed", curl)
-        // Set the new charm URL - this returns once it's done and stored.
-        if err := u.f.SetCharm(curl); err != nil {
-                return err
-        }
-        // Refresh to get the new URL.
-        u.unit.Refresh()
         status := Queued
         if hi != nil {
                 // If a hook operation was interrupted, restore it.
                 status = Pending
         } else {
                 // Otherwise, queue the relevant post-deploy hook.
                 hi = &hook.Info{}
                 switch reason {
                 case Install:
                         hi.Kind = hooks.Install
@@ skipping 227 matching lines @@
                         } else if err != nil {
                                 return err
                         }
                         log.Printf("worker/uniter: joined relation %q", rel)
                         u.relationers[rel.Id()] = r
                         return nil
                 }
         }
         panic("unreachable")
 }