Merge pull request #936 from sgotti/rkt_rmimage_handle_also_stage1

rkt/rmimages: consider also stage1 as referenced image and add functional tests.

rkt/rmimage.go

@@ -127,15 +127,22 @@ func getReferencedImgs(s *store.Store) (map[string]struct{}, error) {
 	walkErrors := []error{}
 	// Consider pods in preparing, prepared, run, exitedgarbage state
 	if err := walkPods(includeMostDirs, func(p *pod) {
-		appImgs, err := p.getAppsHashes()
+		appHashes, err := p.getAppsHashes()
 		if err != nil {
 			// Ignore errors reading/parsing pod file
 			return
 		}
-		for _, appImg := range appImgs {
-			key, err := s.ResolveKey(appImg.String())
+		stage1Hash, err := p.getStage1Hash()
+		if err != nil {
+			// Ignore errors reading/parsing the stage1 hash file
+			return
+		}
+		allHashes := append(appHashes, *stage1Hash)
+
+		for _, imgHash := range allHashes {
+			key, err := s.ResolveKey(imgHash.String())
 			if err != nil && err != store.ErrKeyNotFound {
-				walkErrors = append(walkErrors, fmt.Errorf("bad imageID %q in pod definition: %v", appImg.String(), err))
+				walkErrors = append(walkErrors, fmt.Errorf("bad imageID %q in pod definition: %v", imgHash.String(), err))
 				return
 			}
 			imgs[key] = struct{}{}

tests/rkt_image_rm_test.go

@@ -0,0 +1,134 @@
+// Copyright 2015 The rkt Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package main
+
+import (
+	"bytes"
+	"fmt"
+	"os"
+	"strings"
+	"testing"
+
+	"github.com/coreos/rkt/Godeps/_workspace/src/github.com/ThomasRooney/gexpect"
+)
+
+const (
+	rmImageReferenced = "rkt: imageID is referenced by some containers, cannot remove the tree store"
+	rmImageOk         = "rkt: successfully removed aci for imageID:"
+
+	unreferencedACI = "rkt-unreferencedACI.aci"
+	unreferencedApp = "coreos.com/rkt-unreferenced.aci"
+	referencedACI   = "rkt-inspect.aci"
+	referencedApp   = "coreos.com/rkt-inspect"
+
+	stage1App = "coreos.com/rkt/stage1"
+)
+
+func TestImageRm(t *testing.T) {
+	patchTestACI(unreferencedACI, fmt.Sprintf("--name=%s", unreferencedApp))
+	defer os.Remove(unreferencedACI)
+	ctx := newRktRunCtx()
+	defer ctx.cleanup()
+
+	cmd := fmt.Sprintf("%s --insecure-skip-verify fetch %s", ctx.cmd(), unreferencedACI)
+	t.Logf("Fetching %s: %v", unreferencedACI, cmd)
+	child, err := gexpect.Spawn(cmd)
+	if err != nil {
+		t.Fatalf("Cannot exec: %v", err)
+	}
+	if err := child.Wait(); err != nil {
+		t.Fatalf("rkt didn't terminate correctly: %v", err)
+	}
+
+	cmd = fmt.Sprintf("%s --insecure-skip-verify run %s", ctx.cmd(), referencedACI)
+	t.Logf("Running %s: %v", referencedACI, cmd)
+	child, err = gexpect.Spawn(cmd)
+	if err != nil {
+		t.Fatalf("Cannot exec: %v", err)
+	}
+	if err := child.Wait(); err != nil {
+		t.Fatalf("rkt didn't terminate correctly: %v", err)
+	}
+
+	t.Logf("Retrieving stage1 imageID")
+	stage1ImageID, err := getImageId(ctx, stage1App)
+	if err != nil {
+		t.Fatalf("rkt didn't terminate correctly: %v", err)
+	}
+
+	t.Logf("Retrieving %s imageID", referencedApp)
+	referencedImageID, err := getImageId(ctx, referencedApp)
+	if err != nil {
+		t.Fatalf("rkt didn't terminate correctly: %v", err)
+	}
+
+	t.Logf("Retrieving %s imageID", unreferencedApp)
+	unreferencedImageID, err := getImageId(ctx, unreferencedApp)
+	if err != nil {
+		t.Fatalf("rkt didn't terminate correctly: %v", err)
+	}
+
+	t.Logf("Removing stage1 image (should fail as referenced)")
+	if err := removeImageId(ctx, stage1ImageID, false); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	t.Logf("Removing image for app %s (should fail as referenced)", referencedApp)
+	if err := removeImageId(ctx, referencedImageID, false); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	t.Logf("Removing image for app %s (should work)", unreferencedApp)
+	if err := removeImageId(ctx, unreferencedImageID, true); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+}
+
+func getImageId(ctx *rktRunCtx, name string) (string, error) {
+	cmd := fmt.Sprintf(`/bin/sh -c "%s images -fields=key,appname -no-legend | grep %s | awk '{print $1}'"`, ctx.cmd(), name)
+	child, err := gexpect.Spawn(cmd)
+	if err != nil {
+		return "", fmt.Errorf("Cannot exec rkt: %v", err)
+	}
+	imageID, err := child.ReadLine()
+	imageID = strings.TrimSpace(imageID)
+	imageID = string(bytes.Trim([]byte(imageID), "\x00"))
+	if err != nil {
+		return "", fmt.Errorf("Cannot exec: %v", err)
+	}
+	if err := child.Wait(); err != nil {
+		return "", fmt.Errorf("rkt didn't terminate correctly: %v", err)
+	}
+	return imageID, nil
+}
+
+func removeImageId(ctx *rktRunCtx, imageID string, shouldWork bool) error {
+	expect := rmImageReferenced
+	if shouldWork {
+		expect = rmImageOk
+	}
+	cmd := fmt.Sprintf("%s rmimage %s", ctx.cmd(), imageID)
+	child, err := gexpect.Spawn(cmd)
+	if err != nil {
+		return fmt.Errorf("Cannot exec: %v", err)
+	}
+	if err := child.Expect(expect); err != nil {
+		return fmt.Errorf("Expected %q but not found: %v", expect, err)
+	}
+	if err := child.Wait(); err != nil {
+		return fmt.Errorf("rkt didn't terminate correctly: %v", err)
+	}
+	return nil
+}