graph.go

package docker

import (
	"encoding/json"
	"fmt"
	"github.com/dotcloud/docker/registry"
	"github.com/dotcloud/docker/utils"
	"io"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

// A Graph is a store for versioned filesystem images and the relationship between them.
type Graph struct {
	Root         string
	idIndex      *utils.TruncIndex
	checksumLock map[string]*sync.Mutex
	lockSumFile  *sync.Mutex
	lockSumMap   *sync.Mutex
}

// NewGraph instantiates a new graph at the given root path in the filesystem.
// `root` will be created if it doesn't exist.
func NewGraph(root string) (*Graph, error) {
	abspath, err := filepath.Abs(root)
	if err != nil {
		return nil, err
	}
	// Create the root directory if it doesn't exists
	if err := os.MkdirAll(root, 0700); err != nil && !os.IsExist(err) {
		return nil, err
	}
	graph := &Graph{
		Root:         abspath,
		idIndex:      utils.NewTruncIndex(),
		checksumLock: make(map[string]*sync.Mutex),
		lockSumFile:  &sync.Mutex{},
		lockSumMap:   &sync.Mutex{},
	}
	if err := graph.restore(); err != nil {
		return nil, err
	}
	return graph, nil
}

func (graph *Graph) restore() error {
	dir, err := ioutil.ReadDir(graph.Root)
	if err != nil {
		return err
	}
	for _, v := range dir {
		id := v.Name()
		graph.idIndex.Add(id)
	}
	return nil
}

// FIXME: Implement error subclass instead of looking at the error text
// Note: This is the way golang implements os.IsNotExists on Plan9
func (graph *Graph) IsNotExist(err error) bool {
	return err != nil && (strings.Contains(err.Error(), "does not exist") || strings.Contains(err.Error(), "No such"))
}

// Exists returns true if an image is registered at the given id.
// If the image doesn't exist or if an error is encountered, false is returned.
func (graph *Graph) Exists(id string) bool {
	if _, err := graph.Get(id); err != nil {
		return false
	}
	return true
}

// Get returns the image with the given id, or an error if the image doesn't exist.
func (graph *Graph) Get(name string) (*Image, error) {
	id, err := graph.idIndex.Get(name)
	if err != nil {
		return nil, err
	}
	// FIXME: return nil when the image doesn't exist, instead of an error
	img, err := LoadImage(graph.imageRoot(id))
	if err != nil {
		return nil, err
	}
	if img.ID != id {
		return nil, fmt.Errorf("Image stored at '%s' has wrong id '%s'", id, img.ID)
	}
	img.graph = graph
	if img.Size == 0 {
		root, err := img.root()
		if err != nil {
			return nil, err
		}
		if err := StoreSize(img, root); err != nil {
			return nil, err
		}
	}
	graph.lockSumMap.Lock()
	defer graph.lockSumMap.Unlock()
	if _, exists := graph.checksumLock[img.ID]; !exists {
		graph.checksumLock[img.ID] = &sync.Mutex{}
	}
	return img, nil
}

// Create creates a new image and registers it in the graph.
func (graph *Graph) Create(layerData Archive, container *Container, comment, author string, config *Config) (*Image, error) {
	img := &Image{
		ID:            GenerateID(),
		Comment:       comment,
		Created:       time.Now(),
		DockerVersion: VERSION,
		Author:        author,
		Config:        config,
		Architecture:  "x86_64",
	}
	if container != nil {
		img.Parent = container.Image
		img.Container = container.ID
		img.ContainerConfig = *container.Config
	}
	if err := graph.Register(layerData, layerData != nil, img); err != nil {
		return nil, err
	}
	go img.Checksum()
	return img, nil
}

// Register imports a pre-existing image into the graph.
// FIXME: pass img as first argument
func (graph *Graph) Register(layerData Archive, store bool, img *Image) error {
	if err := ValidateID(img.ID); err != nil {
		return err
	}
	// (This is a convenience to save time. Race conditions are taken care of by os.Rename)
	if graph.Exists(img.ID) {
		return fmt.Errorf("Image %s already exists", img.ID)
	}
	tmp, err := graph.Mktemp("")
	defer os.RemoveAll(tmp)
	if err != nil {
		return fmt.Errorf("Mktemp failed: %s", err)
	}
	if err := StoreImage(img, layerData, tmp, store); err != nil {
		return err
	}
	// Commit
	if err := os.Rename(tmp, graph.imageRoot(img.ID)); err != nil {
		return err
	}
	img.graph = graph
	graph.idIndex.Add(img.ID)
	graph.checksumLock[img.ID] = &sync.Mutex{}
	return nil
}

// TempLayerArchive creates a temporary archive of the given image's filesystem layer.
//   The archive is stored on disk and will be automatically deleted as soon as has been read.
//   If output is not nil, a human-readable progress bar will be written to it.
//   FIXME: does this belong in Graph? How about MktempFile, let the caller use it for archives?
func (graph *Graph) TempLayerArchive(id string, compression Compression, output io.Writer) (*TempArchive, error) {
	image, err := graph.Get(id)
	if err != nil {
		return nil, err
	}
	tmp, err := graph.tmp()
	if err != nil {
		return nil, err
	}
	archive, err := image.TarLayer(compression)
	if err != nil {
		return nil, err
	}
	sf := utils.NewStreamFormatter(false)
	return NewTempArchive(utils.ProgressReader(ioutil.NopCloser(archive), 0, output, sf.FormatProgress("Buffering to disk", "%v/%v (%v)"), sf), tmp.Root)
}

// Mktemp creates a temporary sub-directory inside the graph's filesystem.
func (graph *Graph) Mktemp(id string) (string, error) {
	if id == "" {
		id = GenerateID()
	}
	tmp, err := graph.tmp()
	if err != nil {
		return "", fmt.Errorf("Couldn't create temp: %s", err)
	}
	if tmp.Exists(id) {
		return "", fmt.Errorf("Image %d already exists", id)
	}
	return tmp.imageRoot(id), nil
}

func (graph *Graph) tmp() (*Graph, error) {
	return NewGraph(path.Join(graph.Root, ":tmp:"))
}

// Check if given error is "not empty".
// Note: this is the way golang does it internally with os.IsNotExists.
func isNotEmpty(err error) bool {
	switch pe := err.(type) {
	case nil:
		return false
	case *os.PathError:
		err = pe.Err
	case *os.LinkError:
		err = pe.Err
	}
	return strings.Contains(err.Error(), " not empty")
}

// Delete atomically removes an image from the graph.
func (graph *Graph) Delete(name string) error {
	id, err := graph.idIndex.Get(name)
	if err != nil {
		return err
	}
	tmp, err := graph.Mktemp("")
	if err != nil {
		return err
	}
	graph.idIndex.Delete(id)
	err = os.Rename(graph.imageRoot(id), tmp)
	if err != nil {
		return err
	}
	return os.RemoveAll(tmp)
}

// Map returns a list of all images in the graph, addressable by ID.
func (graph *Graph) Map() (map[string]*Image, error) {
	// FIXME: this should replace All()
	all, err := graph.All()
	if err != nil {
		return nil, err
	}
	images := make(map[string]*Image, len(all))
	for _, image := range all {
		images[image.ID] = image
	}
	return images, nil
}

// All returns a list of all images in the graph.
func (graph *Graph) All() ([]*Image, error) {
	var images []*Image
	err := graph.WalkAll(func(image *Image) {
		images = append(images, image)
	})
	return images, err
}

// WalkAll iterates over each image in the graph, and passes it to a handler.
// The walking order is undetermined.
func (graph *Graph) WalkAll(handler func(*Image)) error {
	files, err := ioutil.ReadDir(graph.Root)
	if err != nil {
		return err
	}
	for _, st := range files {
		if img, err := graph.Get(st.Name()); err != nil {
			// Skip image
			continue
		} else if handler != nil {
			handler(img)
		}
	}
	return nil
}

// ByParent returns a lookup table of images by their parent.
// If an image of id ID has 3 children images, then the value for key ID
// will be a list of 3 images.
// If an image has no children, it will not have an entry in the table.
func (graph *Graph) ByParent() (map[string][]*Image, error) {
	byParent := make(map[string][]*Image)
	err := graph.WalkAll(func(image *Image) {
		parent, err := graph.Get(image.Parent)
		if err != nil {
			return
		}
		if children, exists := byParent[parent.ID]; exists {
			byParent[parent.ID] = []*Image{image}
		} else {
			byParent[parent.ID] = append(children, image)
		}
	})
	return byParent, err
}

// Heads returns all heads in the graph, keyed by id.
// A head is an image which is not the parent of another image in the graph.
func (graph *Graph) Heads() (map[string]*Image, error) {
	heads := make(map[string]*Image)
	byParent, err := graph.ByParent()
	if err != nil {
		return nil, err
	}
	err = graph.WalkAll(func(image *Image) {
		// If it's not in the byParent lookup table, then
		// it's not a parent -> so it's a head!
		if _, exists := byParent[image.ID]; !exists {
			heads[image.ID] = image
		}
	})
	return heads, err
}

func (graph *Graph) imageRoot(id string) string {
	return path.Join(graph.Root, id)
}

func (graph *Graph) getStoredChecksums() (map[string]string, error) {
	checksums := make(map[string]string)
	// FIXME: Store the checksum in memory

	if checksumDict, err := ioutil.ReadFile(path.Join(graph.Root, "checksums")); err == nil {
		if err := json.Unmarshal(checksumDict, &checksums); err != nil {
			return nil, err
		}
	}
	return checksums, nil
}

func (graph *Graph) storeChecksums(checksums map[string]string) error {
	checksumJSON, err := json.Marshal(checksums)
	if err != nil {
		return err
	}
	if err := ioutil.WriteFile(path.Join(graph.Root, "checksums"), checksumJSON, 0600); err != nil {
		return err
	}
	return nil
}

func (graph *Graph) UpdateChecksums(newChecksums map[string]*registry.ImgData) error {
	graph.lockSumFile.Lock()
	defer graph.lockSumFile.Unlock()

	localChecksums, err := graph.getStoredChecksums()
	if err != nil {
		return err
	}
	for id, elem := range newChecksums {
		localChecksums[id] = elem.Checksum
	}
	return graph.storeChecksums(localChecksums)
}