o/snapstate, o/hookstate: run component remove hooks

overlord/hookstate/hooks.go

@@ -37,6 +37,7 @@ func init() {
 	snapstate.SetupInstallComponentHook = SetupInstallComponentHook
 	snapstate.SetupPreRefreshComponentHook = SetupPreRefreshComponentHook
 	snapstate.SetupPostRefreshComponentHook = SetupPostRefreshComponentHook
+	snapstate.SetupRemoveComponentHook = SetupRemoveComponentHook
 	snapstate.SetupPreRefreshHook = SetupPreRefreshHook
 	snapstate.SetupPostRefreshHook = SetupPostRefreshHook
 	snapstate.SetupRemoveHook = SetupRemoveHook
@@ -98,6 +99,20 @@ func SetupPostRefreshComponentHook(st *state.State, snap, component string) *sta
 	return task
 }
 
+func SetupRemoveComponentHook(st *state.State, snap, component string) *state.Task {
+	hooksup := &HookSetup{
+		Snap:      snap,
+		Component: component,
+		Hook:      "remove",
+		Optional:  true,
+	}
+
+	summary := fmt.Sprintf(i18n.G(`Run remove hook of "%s+%s" component if present`), hooksup.Snap, hooksup.Component)
+	task := HookTask(st, summary, hooksup, nil)
+
+	return task
+}
+
 func SetupPostRefreshHook(st *state.State, snapName string) *state.Task {
 	hooksup := &HookSetup{
 		Snap:     snapName,

overlord/snapstate/component.go

@@ -554,28 +554,26 @@ func removeComponentTasks(st *state.State, snapst *SnapState, compst *sequence.C
 		CompType:     compst.CompType,
 	}
 
-	// TODO:COMPS: Run component remove hook. This will change the first task run,
-	// changing uses of unlink below to the new task.
-
-	// Unlink component
-	unlink := st.NewTask("unlink-current-component", fmt.Sprintf(i18n.G(
-		"Make current revision for component %q unavailable"),
-		compst.SideInfo.Component))
-	unlink.Set("component-setup", compSetup)
-	unlink.Set("snap-setup", snapSup)
-
-	var prev *state.Task
-	tasks := []*state.Task{unlink}
-	prev = unlink
+	removeHook := SetupRemoveComponentHook(st, instName, compst.SideInfo.Component.ComponentName)
+	removeHook.Set("component-setup", compSetup)
+	removeHook.Set("snap-setup", snapSup)
 
+	prev := removeHook
+	tasks := []*state.Task{removeHook}
 	addTask := func(t *state.Task) {
-		t.Set("component-setup-task", unlink.ID())
-		t.Set("snap-setup-task", unlink.ID())
+		t.Set("component-setup-task", removeHook.ID())
+		t.Set("snap-setup-task", removeHook.ID())
 		t.WaitFor(prev)
 		tasks = append(tasks, t)
 		prev = t
 	}
 
+	// Unlink component
+	unlink := st.NewTask("unlink-current-component", fmt.Sprintf(i18n.G(
+		"Make current revision for component %q unavailable"),
+		compst.SideInfo.Component))
+	addTask(unlink)
+
 	// For kernel-modules, regenerate drivers tree
 	revisionStr := fmt.Sprintf(" (%s)", compst.SideInfo.Revision)
 	if compst.CompType == snap.KernelModulesComponent {

overlord/snapstate/component_remove_test.go

@@ -32,8 +32,7 @@ import (
 )
 
 func expectedComponentRemoveTasks(opts int) []string {
-	var removeTasks []string
-	removeTasks = append(removeTasks, "unlink-current-component")
+	removeTasks := []string{"run-hook[remove]", "unlink-current-component"}
 	if opts&compTypeIsKernMods != 0 {
 		removeTasks = append(removeTasks, "prepare-kernel-modules-components")
 	}
@@ -292,9 +291,10 @@ func (s *snapmgrTestSuite) TestRemoveComponentsPathRunWithError(c *C) {
 	// component tasks are undone/hold
 	for i := 0; i < 2; i++ {
 		ts := tss[i].Tasks()
-		c.Check(ts[0].Status(), Equals, state.UndoneStatus)
 		c.Check(ts[1].Status(), Equals, state.UndoneStatus)
-		c.Check(ts[2].Status(), Equals, state.HoldStatus)
+		c.Check(ts[1].Status(), Equals, state.UndoneStatus)
+		c.Check(ts[2].Status(), Equals, state.UndoneStatus)
+		c.Check(ts[3].Status(), Equals, state.HoldStatus)
 	}
 	// update profile is hold
 	c.Check(updateProfTask.Status(), Equals, state.HoldStatus)

overlord/snapstate/snapstate.go

@@ -945,6 +945,10 @@ var SetupPostRefreshComponentHook = func(st *state.State, snap, component string
 	panic("internal error: snapstate.SetupPostRefreshComponentHook is unset")
 }
 
+var SetupRemoveComponentHook = func(st *state.State, snap, component string) *state.Task {
+	panic("internal error: snapstate.SetupRemoveComponentHook is unset")
+}
+
 var SetupPreRefreshHook = func(st *state.State, snapName string) *state.Task {
 	panic("internal error: snapstate.SetupPreRefreshHook is unset")
 }
@@ -3617,7 +3621,11 @@ func removeTasks(st *state.State, name string, revision snap.Revision, flags *Re
 
 	// only run remove hook if uninstalling the snap completely
 	if removeAll {
-		// TODO:COMPS: Run component remove hooks
+		for _, comp := range snapst.Sequence.ComponentsForRevision(snapst.Current) {
+			removeCompHook := SetupRemoveComponentHook(st, snapsup.InstanceName(), comp.SideInfo.Component.ComponentName)
+			addNext(state.NewTaskSet(removeCompHook))
+			prev = removeCompHook
+		}
 
 		removeHook := SetupRemoveHook(st, snapsup.InstanceName())
 		addNext(state.NewTaskSet(removeHook))

overlord/snapstate/snapstate_remove_test.go

@@ -2148,7 +2148,9 @@ func (s *snapmgrTestSuite) TestRemoveWithCompsTasks(c *C) {
 	c.Assert(s.state.TaskCount(), Equals, len(ts.Tasks()))
 	c.Assert(taskKinds(ts.Tasks()), DeepEquals, []string{
 		"stop-snap-services",
-		"run-hook[remove]",
+		"run-hook[remove]", // component remove hook
+		"run-hook[remove]", // component remove hook
+		"run-hook[remove]", // snap remove hook
 		"auto-disconnect",
 		"save-snapshot",
 		"remove-aliases",

overlord/snapstate/snapstate_test.go

@@ -187,6 +187,7 @@ func (s *snapmgrBaseTest) SetUpTest(c *C) {
 	oldSetupInstallComponentHook := snapstate.SetupInstallComponentHook
 	oldSetupPostRefreshComponentHook := snapstate.SetupPostRefreshComponentHook
 	oldSetupPreRefreshComponentHook := snapstate.SetupPreRefreshComponentHook
+	oldSetupRemoveComponentHook := snapstate.SetupRemoveComponentHook
 	oldSetupPreRefreshHook := snapstate.SetupPreRefreshHook
 	oldSetupPostRefreshHook := snapstate.SetupPostRefreshHook
 	oldSetupRemoveHook := snapstate.SetupRemoveHook
@@ -236,6 +237,7 @@ func (s *snapmgrBaseTest) SetUpTest(c *C) {
 		snapstate.SetupInstallComponentHook = oldSetupInstallComponentHook
 		snapstate.SetupPostRefreshComponentHook = oldSetupPostRefreshComponentHook
 		snapstate.SetupPreRefreshComponentHook = oldSetupPreRefreshComponentHook
+		snapstate.SetupRemoveComponentHook = oldSetupRemoveComponentHook
 		snapstate.SetupPreRefreshHook = oldSetupPreRefreshHook
 		snapstate.SetupPostRefreshHook = oldSetupPostRefreshHook
 		snapstate.SetupRemoveHook = oldSetupRemoveHook