<commit_before>package tarexport

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"

	"github.com/Sirupsen/logrus"
	"github.com/docker/distribution"
	"github.com/docker/docker/image"
	"github.com/docker/docker/image/v1"
	"github.com/docker/docker/layer"
	"github.com/docker/docker/pkg/archive"
	"github.com/docker/docker/pkg/chrootarchive"
	"github.com/docker/docker/pkg/progress"
	"github.com/docker/docker/pkg/streamformatter"
	"github.com/docker/docker/pkg/stringid"
	"github.com/docker/docker/pkg/symlink"
	"github.com/docker/docker/pkg/system"
	"github.com/docker/docker/reference"
	"github.com/opencontainers/go-digest"
)

func (l *tarexporter) Load(inTar io.ReadCloser, outStream io.Writer, quiet bool) error {
	var (
		sf             = streamformatter.NewJSONStreamFormatter()
		progressOutput progress.Output
	)
	if !quiet {
		progressOutput = sf.NewProgressOutput(outStream, false)
	}
	outStream = &streamformatter.StdoutFormatter{Writer: outStream, StreamFormatter: streamformatter.NewJSONStreamFormatter()}

	tmpDir, err := ioutil.TempDir("", "docker-import-")
	if err != nil {
		return err
	}
	defer os.RemoveAll(tmpDir)

	if err := chrootarchive.Untar(inTar, tmpDir, nil); err != nil {
		return err
	}
	// read manifest, if no file then load in legacy mode
	manifestPath, err := safePath(tmpDir, manifestFileName)
	if err != nil {
		return err
	}
	manifestFile, err := os.Open(manifestPath)
	if err != nil {
		if os.IsNotExist(err) {
			return l.legacyLoad(tmpDir, outStream, progressOutput)
		}
		return err
	}
	defer manifestFile.Close()

	var manifest []manifestItem
	if err := json.NewDecoder(manifestFile).Decode(&manifest); err != nil {
		return err
	}

	var parentLinks []parentLink
	var imageIDsStr string
	var imageRefCount int

	for _, m := range manifest {
		configPath, err := safePath(tmpDir, m.Config)
		if err != nil {
			return err
		}
		config, err := ioutil.ReadFile(configPath)
		if err != nil {
			return err
		}
		img, err := image.NewFromJSON(config)
		if err != nil {
			return err
		}
		var rootFS image.RootFS
		rootFS = *img.RootFS
		rootFS.DiffIDs = nil

		if expected, actual := len(m.Layers), len(img.RootFS.DiffIDs); expected != actual {
			return fmt.Errorf("invalid manifest, layers length mismatch: expected %q, got %q", expected, actual)
		}

		for i, diffID := range img.RootFS.DiffIDs {
			layerPath, err := safePath(tmpDir, m.Layers[i])
			if err != nil {
				return err
			}
			r := rootFS
			r.Append(diffID)
			newLayer, err := l.ls.Get(r.ChainID())
			if err != nil {
				newLayer, err = l.loadLayer(layerPath, rootFS, diffID.String(), m.LayerSources[diffID], progressOutput)
				if err != nil {
					return err
				}
			}
			defer layer.ReleaseAndLog(l.ls, newLayer)
			if expected, actual := diffID, newLayer.DiffID(); expected != actual {
				return fmt.Errorf("invalid diffID for layer %d: expected %q, got %q", i, expected, actual)
			}
			rootFS.Append(diffID)
		}

		imgID, err := l.is.Create(config)
		if err != nil {
			return err
		}
		imageIDsStr += fmt.Sprintf("Loaded image ID: %s\n", imgID)

		imageRefCount = 0
		for _, repoTag := range m.RepoTags {
			named, err := reference.ParseNamed(repoTag)
			if err != nil {
				return err
			}
			ref, ok := named.(reference.NamedTagged)
			if !ok {
				return fmt.Errorf("invalid tag %q", repoTag)
			}
			l.setLoadedTag(ref, imgID.Digest(), outStream)
			outStream.Write([]byte(fmt.Sprintf("Loaded image: %s\n", ref)))
			imageRefCount++
		}

		parentLinks = append(parentLinks, parentLink{imgID, m.Parent})
		l.loggerImgEvent.LogImageEvent(imgID.String(), imgID.String(), "load")
	}

	for _, p := range validatedParentLinks(parentLinks) {
		if p.parentID != "" {
			if err := l.setParentID(p.id, p.parentID); err != nil {
				return err
			}
		}
	}

	if imageRefCount == 0 {
		outStream.Write([]byte(imageIDsStr))
	}

	return nil
}

func (l *tarexporter) setParentID(id, parentID image.ID) error {
	img, err := l.is.Get(id)
	if err != nil {
		return err
	}
	parent, err := l.is.Get(parentID)
	if err != nil {
		return err
	}
	if !checkValidParent(img, parent) {
		return fmt.Errorf("image %v is not a valid parent for %v", parent.ID(), img.ID())
	}
	return l.is.SetParent(id, parentID)
}

func (l *tarexporter) loadLayer(filename string, rootFS image.RootFS, id string, foreignSrc distribution.Descriptor, progressOutput progress.Output) (layer.Layer, error) {
	// We use system.OpenSequential to use sequential file access on Windows, avoiding
	// depleting the standby list. On Linux, this equates to a regular os.Open.
	rawTar, err := system.OpenSequential(filename)
	if err != nil {
		logrus.Debugf("Error reading embedded tar: %v", err)
		return nil, err
	}
	defer rawTar.Close()

	var r io.Reader
	if progressOutput != nil {
		fileInfo, err := rawTar.Stat()
		if err != nil {
			logrus.Debugf("Error statting file: %v", err)
			return nil, err
		}

		r = progress.NewProgressReader(rawTar, progressOutput, fileInfo.Size(), stringid.TruncateID(id), "Loading layer")
	} else {
		r = rawTar
	}

	inflatedLayerData, err := archive.DecompressStream(r)
	if err != nil {
		return nil, err
	}
	defer inflatedLayerData.Close()

	if ds, ok := l.ls.(layer.DescribableStore); ok {
		return ds.RegisterWithDescriptor(inflatedLayerData, rootFS.ChainID(), foreignSrc)
	}
	return l.ls.Register(inflatedLayerData, rootFS.ChainID())
}

func (l *tarexporter) setLoadedTag(ref reference.NamedTagged, imgID digest.Digest, outStream io.Writer) error {
	if prevID, err := l.rs.Get(ref); err == nil && prevID != imgID {
		fmt.Fprintf(outStream, "The image %s already exists, renaming the old one with ID %s to empty string\n", ref.String(), string(prevID)) // todo: this message is wrong in case of multiple tags
	}

	if err := l.rs.AddTag(ref, imgID, true); err != nil {
		return err
	}
	return nil
}

func (l *tarexporter) legacyLoad(tmpDir string, outStream io.Writer, progressOutput progress.Output) error {
	legacyLoadedMap := make(map[string]image.ID)

	dirs, err := ioutil.ReadDir(tmpDir)
	if err != nil {
		return err
	}

	// every dir represents an image
	for _, d := range dirs {
		if d.IsDir() {
			if err := l.legacyLoadImage(d.Name(), tmpDir, legacyLoadedMap, progressOutput); err != nil {
				return err
			}
		}
	}

	// load tags from repositories file
	repositoriesPath, err := safePath(tmpDir, legacyRepositoriesFileName)
	if err != nil {
		return err
	}
	repositoriesFile, err := os.Open(repositoriesPath)
	if err != nil {
		return err
	}
	defer repositoriesFile.Close()

	repositories := make(map[string]map[string]string)
	if err := json.NewDecoder(repositoriesFile).Decode(&repositories); err != nil {
		return err
	}

	for name, tagMap := range repositories {
		for tag, oldID := range tagMap {
			imgID, ok := legacyLoadedMap[oldID]
			if !ok {
				return fmt.Errorf("invalid target ID: %v", oldID)
			}
			named, err := reference.WithName(name)
			if err != nil {
				return err
			}
			ref, err := reference.WithTag(named, tag)
			if err != nil {
				return err
			}
			l.setLoadedTag(ref, imgID.Digest(), outStream)
		}
	}

	return nil
}

func (l *tarexporter) legacyLoadImage(oldID, sourceDir string, loadedMap map[string]image.ID, progressOutput progress.Output) error {
	if _, loaded := loadedMap[oldID]; loaded {
		return nil
	}
	configPath, err := safePath(sourceDir, filepath.Join(oldID, legacyConfigFileName))
	if err != nil {
		return err
	}
	imageJSON, err := ioutil.ReadFile(configPath)
	if err != nil {
		logrus.Debugf("Error reading json: %v", err)
		return err
	}

	var img struct{ Parent string }
	if err := json.Unmarshal(imageJSON, &img); err != nil {
		return err
	}

	var parentID image.ID
	if img.Parent != "" {
		for {
			var loaded bool
			if parentID, loaded = loadedMap[img.Parent]; !loaded {
				if err := l.legacyLoadImage(img.Parent, sourceDir, loadedMap, progressOutput); err != nil {
					return err
				}
			} else {
				break
			}
		}
	}

	// todo: try to connect with migrate code
	rootFS := image.NewRootFS()
	var history []image.History

	if parentID != "" {
		parentImg, err := l.is.Get(parentID)
		if err != nil {
			return err
		}

		rootFS = parentImg.RootFS
		history = parentImg.History
	}

	layerPath, err := safePath(sourceDir, filepath.Join(oldID, legacyLayerFileName))
	if err != nil {
		return err
	}
	newLayer, err := l.loadLayer(layerPath, *rootFS, oldID, distribution.Descriptor{}, progressOutput)
	if err != nil {
		return err
	}
	rootFS.Append(newLayer.DiffID())

	h, err := v1.HistoryFromConfig(imageJSON, false)
	if err != nil {
		return err
	}
	history = append(history, h)

	config, err := v1.MakeConfigFromV1Config(imageJSON, rootFS, history)
	if err != nil {
		return err
	}
	imgID, err := l.is.Create(config)
	if err != nil {
		return err
	}

	metadata, err := l.ls.Release(newLayer)
	layer.LogReleaseMetadata(metadata)
	if err != nil {
		return err
	}

	if parentID != "" {
		if err := l.is.SetParent(imgID, parentID); err != nil {
			return err
		}
	}

	loadedMap[oldID] = imgID
	return nil
}

func safePath(base, path string) (string, error) {
	return symlink.FollowSymlinkInScope(filepath.Join(base, path), base)
}

type parentLink struct {
	id, parentID image.ID
}

func validatedParentLinks(pl []parentLink) (ret []parentLink) {
mainloop:
	for i, p := range pl {
		ret = append(ret, p)
		for _, p2 := range pl {
			if p2.id == p.parentID && p2.id != p.id {
				continue mainloop
			}
		}
		ret[i].parentID = ""
	}
	return
}

func checkValidParent(img, parent *image.Image) bool {
	if len(img.History) == 0 && len(parent.History) == 0 {
		return true // having history is not mandatory
	}
	if len(img.History)-len(parent.History) != 1 {
		return false
	}
	for i, h := range parent.History {
		if !reflect.DeepEqual(h, img.History[i]) {
			return false
		}
	}
	return true
}
<commit_msg>image: tarexport: do not quote integers in format string<commit_after>package tarexport

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"

	"github.com/Sirupsen/logrus"
	"github.com/docker/distribution"
	"github.com/docker/docker/image"
	"github.com/docker/docker/image/v1"
	"github.com/docker/docker/layer"
	"github.com/docker/docker/pkg/archive"
	"github.com/docker/docker/pkg/chrootarchive"
	"github.com/docker/docker/pkg/progress"
	"github.com/docker/docker/pkg/streamformatter"
	"github.com/docker/docker/pkg/stringid"
	"github.com/docker/docker/pkg/symlink"
	"github.com/docker/docker/pkg/system"
	"github.com/docker/docker/reference"
	"github.com/opencontainers/go-digest"
)

func (l *tarexporter) Load(inTar io.ReadCloser, outStream io.Writer, quiet bool) error {
	var (
		sf             = streamformatter.NewJSONStreamFormatter()
		progressOutput progress.Output
	)
	if !quiet {
		progressOutput = sf.NewProgressOutput(outStream, false)
	}
	outStream = &streamformatter.StdoutFormatter{Writer: outStream, StreamFormatter: streamformatter.NewJSONStreamFormatter()}

	tmpDir, err := ioutil.TempDir("", "docker-import-")
	if err != nil {
		return err
	}
	defer os.RemoveAll(tmpDir)

	if err := chrootarchive.Untar(inTar, tmpDir, nil); err != nil {
		return err
	}
	// read manifest, if no file then load in legacy mode
	manifestPath, err := safePath(tmpDir, manifestFileName)
	if err != nil {
		return err
	}
	manifestFile, err := os.Open(manifestPath)
	if err != nil {
		if os.IsNotExist(err) {
			return l.legacyLoad(tmpDir, outStream, progressOutput)
		}
		return err
	}
	defer manifestFile.Close()

	var manifest []manifestItem
	if err := json.NewDecoder(manifestFile).Decode(&manifest); err != nil {
		return err
	}

	var parentLinks []parentLink
	var imageIDsStr string
	var imageRefCount int

	for _, m := range manifest {
		configPath, err := safePath(tmpDir, m.Config)
		if err != nil {
			return err
		}
		config, err := ioutil.ReadFile(configPath)
		if err != nil {
			return err
		}
		img, err := image.NewFromJSON(config)
		if err != nil {
			return err
		}
		var rootFS image.RootFS
		rootFS = *img.RootFS
		rootFS.DiffIDs = nil

		if expected, actual := len(m.Layers), len(img.RootFS.DiffIDs); expected != actual {
			return fmt.Errorf("invalid manifest, layers length mismatch: expected %d, got %d", expected, actual)
		}

		for i, diffID := range img.RootFS.DiffIDs {
			layerPath, err := safePath(tmpDir, m.Layers[i])
			if err != nil {
				return err
			}
			r := rootFS
			r.Append(diffID)
			newLayer, err := l.ls.Get(r.ChainID())
			if err != nil {
				newLayer, err = l.loadLayer(layerPath, rootFS, diffID.String(), m.LayerSources[diffID], progressOutput)
				if err != nil {
					return err
				}
			}
			defer layer.ReleaseAndLog(l.ls, newLayer)
			if expected, actual := diffID, newLayer.DiffID(); expected != actual {
				return fmt.Errorf("invalid diffID for layer %d: expected %q, got %q", i, expected, actual)
			}
			rootFS.Append(diffID)
		}

		imgID, err := l.is.Create(config)
		if err != nil {
			return err
		}
		imageIDsStr += fmt.Sprintf("Loaded image ID: %s\n", imgID)

		imageRefCount = 0
		for _, repoTag := range m.RepoTags {
			named, err := reference.ParseNamed(repoTag)
			if err != nil {
				return err
			}
			ref, ok := named.(reference.NamedTagged)
			if !ok {
				return fmt.Errorf("invalid tag %q", repoTag)
			}
			l.setLoadedTag(ref, imgID.Digest(), outStream)
			outStream.Write([]byte(fmt.Sprintf("Loaded image: %s\n", ref)))
			imageRefCount++
		}

		parentLinks = append(parentLinks, parentLink{imgID, m.Parent})
		l.loggerImgEvent.LogImageEvent(imgID.String(), imgID.String(), "load")
	}

	for _, p := range validatedParentLinks(parentLinks) {
		if p.parentID != "" {
			if err := l.setParentID(p.id, p.parentID); err != nil {
				return err
			}
		}
	}

	if imageRefCount == 0 {
		outStream.Write([]byte(imageIDsStr))
	}

	return nil
}

func (l *tarexporter) setParentID(id, parentID image.ID) error {
	img, err := l.is.Get(id)
	if err != nil {
		return err
	}
	parent, err := l.is.Get(parentID)
	if err != nil {
		return err
	}
	if !checkValidParent(img, parent) {
		return fmt.Errorf("image %v is not a valid parent for %v", parent.ID(), img.ID())
	}
	return l.is.SetParent(id, parentID)
}

func (l *tarexporter) loadLayer(filename string, rootFS image.RootFS, id string, foreignSrc distribution.Descriptor, progressOutput progress.Output) (layer.Layer, error) {
	// We use system.OpenSequential to use sequential file access on Windows, avoiding
	// depleting the standby list. On Linux, this equates to a regular os.Open.
	rawTar, err := system.OpenSequential(filename)
	if err != nil {
		logrus.Debugf("Error reading embedded tar: %v", err)
		return nil, err
	}
	defer rawTar.Close()

	var r io.Reader
	if progressOutput != nil {
		fileInfo, err := rawTar.Stat()
		if err != nil {
			logrus.Debugf("Error statting file: %v", err)
			return nil, err
		}

		r = progress.NewProgressReader(rawTar, progressOutput, fileInfo.Size(), stringid.TruncateID(id), "Loading layer")
	} else {
		r = rawTar
	}

	inflatedLayerData, err := archive.DecompressStream(r)
	if err != nil {
		return nil, err
	}
	defer inflatedLayerData.Close()

	if ds, ok := l.ls.(layer.DescribableStore); ok {
		return ds.RegisterWithDescriptor(inflatedLayerData, rootFS.ChainID(), foreignSrc)
	}
	return l.ls.Register(inflatedLayerData, rootFS.ChainID())
}

func (l *tarexporter) setLoadedTag(ref reference.NamedTagged, imgID digest.Digest, outStream io.Writer) error {
	if prevID, err := l.rs.Get(ref); err == nil && prevID != imgID {
		fmt.Fprintf(outStream, "The image %s already exists, renaming the old one with ID %s to empty string\n", ref.String(), string(prevID)) // todo: this message is wrong in case of multiple tags
	}

	if err := l.rs.AddTag(ref, imgID, true); err != nil {
		return err
	}
	return nil
}

func (l *tarexporter) legacyLoad(tmpDir string, outStream io.Writer, progressOutput progress.Output) error {
	legacyLoadedMap := make(map[string]image.ID)

	dirs, err := ioutil.ReadDir(tmpDir)
	if err != nil {
		return err
	}

	// every dir represents an image
	for _, d := range dirs {
		if d.IsDir() {
			if err := l.legacyLoadImage(d.Name(), tmpDir, legacyLoadedMap, progressOutput); err != nil {
				return err
			}
		}
	}

	// load tags from repositories file
	repositoriesPath, err := safePath(tmpDir, legacyRepositoriesFileName)
	if err != nil {
		return err
	}
	repositoriesFile, err := os.Open(repositoriesPath)
	if err != nil {
		return err
	}
	defer repositoriesFile.Close()

	repositories := make(map[string]map[string]string)
	if err := json.NewDecoder(repositoriesFile).Decode(&repositories); err != nil {
		return err
	}

	for name, tagMap := range repositories {
		for tag, oldID := range tagMap {
			imgID, ok := legacyLoadedMap[oldID]
			if !ok {
				return fmt.Errorf("invalid target ID: %v", oldID)
			}
			named, err := reference.WithName(name)
			if err != nil {
				return err
			}
			ref, err := reference.WithTag(named, tag)
			if err != nil {
				return err
			}
			l.setLoadedTag(ref, imgID.Digest(), outStream)
		}
	}

	return nil
}

func (l *tarexporter) legacyLoadImage(oldID, sourceDir string, loadedMap map[string]image.ID, progressOutput progress.Output) error {
	if _, loaded := loadedMap[oldID]; loaded {
		return nil
	}
	configPath, err := safePath(sourceDir, filepath.Join(oldID, legacyConfigFileName))
	if err != nil {
		return err
	}
	imageJSON, err := ioutil.ReadFile(configPath)
	if err != nil {
		logrus.Debugf("Error reading json: %v", err)
		return err
	}

	var img struct{ Parent string }
	if err := json.Unmarshal(imageJSON, &img); err != nil {
		return err
	}

	var parentID image.ID
	if img.Parent != "" {
		for {
			var loaded bool
			if parentID, loaded = loadedMap[img.Parent]; !loaded {
				if err := l.legacyLoadImage(img.Parent, sourceDir, loadedMap, progressOutput); err != nil {
					return err
				}
			} else {
				break
			}
		}
	}

	// todo: try to connect with migrate code
	rootFS := image.NewRootFS()
	var history []image.History

	if parentID != "" {
		parentImg, err := l.is.Get(parentID)
		if err != nil {
			return err
		}

		rootFS = parentImg.RootFS
		history = parentImg.History
	}

	layerPath, err := safePath(sourceDir, filepath.Join(oldID, legacyLayerFileName))
	if err != nil {
		return err
	}
	newLayer, err := l.loadLayer(layerPath, *rootFS, oldID, distribution.Descriptor{}, progressOutput)
	if err != nil {
		return err
	}
	rootFS.Append(newLayer.DiffID())

	h, err := v1.HistoryFromConfig(imageJSON, false)
	if err != nil {
		return err
	}
	history = append(history, h)

	config, err := v1.MakeConfigFromV1Config(imageJSON, rootFS, history)
	if err != nil {
		return err
	}
	imgID, err := l.is.Create(config)
	if err != nil {
		return err
	}

	metadata, err := l.ls.Release(newLayer)
	layer.LogReleaseMetadata(metadata)
	if err != nil {
		return err
	}

	if parentID != "" {
		if err := l.is.SetParent(imgID, parentID); err != nil {
			return err
		}
	}

	loadedMap[oldID] = imgID
	return nil
}

func safePath(base, path string) (string, error) {
	return symlink.FollowSymlinkInScope(filepath.Join(base, path), base)
}

type parentLink struct {
	id, parentID image.ID
}

func validatedParentLinks(pl []parentLink) (ret []parentLink) {
mainloop:
	for i, p := range pl {
		ret = append(ret, p)
		for _, p2 := range pl {
			if p2.id == p.parentID && p2.id != p.id {
				continue mainloop
			}
		}
		ret[i].parentID = ""
	}
	return
}

func checkValidParent(img, parent *image.Image) bool {
	if len(img.History) == 0 && len(parent.History) == 0 {
		return true // having history is not mandatory
	}
	if len(img.History)-len(parent.History) != 1 {
		return false
	}
	for i, h := range parent.History {
		if !reflect.DeepEqual(h, img.History[i]) {
			return false
		}
	}
	return true
}
<|endoftext|>


<commit_before>package tarexport

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"time"

	"github.com/docker/distribution"
	"github.com/docker/distribution/digest"
	"github.com/docker/docker/image"
	"github.com/docker/docker/image/v1"
	"github.com/docker/docker/layer"
	"github.com/docker/docker/pkg/archive"
	"github.com/docker/docker/pkg/system"
	"github.com/docker/docker/reference"
)

type imageDescriptor struct {
	refs   []reference.NamedTagged
	layers []string
}

type saveSession struct {
	*tarexporter
	outDir      string
	images      map[image.ID]*imageDescriptor
	savedLayers map[string]struct{}
	diffIDPaths map[layer.DiffID]string // cache every diffID blob to avoid duplicates
}

func (l *tarexporter) Save(names []string, outStream io.Writer) error {
	images, err := l.parseNames(names)
	if err != nil {
		return err
	}

	return (&saveSession{tarexporter: l, images: images}).save(outStream)
}

func (l *tarexporter) parseNames(names []string) (map[image.ID]*imageDescriptor, error) {
	imgDescr := make(map[image.ID]*imageDescriptor)

	addAssoc := func(id image.ID, ref reference.Named) {
		if _, ok := imgDescr[id]; !ok {
			imgDescr[id] = &imageDescriptor{}
		}

		if ref != nil {
			var tagged reference.NamedTagged
			if _, ok := ref.(reference.Canonical); ok {
				return
			}
			var ok bool
			if tagged, ok = ref.(reference.NamedTagged); !ok {
				var err error
				if tagged, err = reference.WithTag(ref, reference.DefaultTag); err != nil {
					return
				}
			}

			for _, t := range imgDescr[id].refs {
				if tagged.String() == t.String() {
					return
				}
			}
			imgDescr[id].refs = append(imgDescr[id].refs, tagged)
		}
	}

	for _, name := range names {
		id, ref, err := reference.ParseIDOrReference(name)
		if err != nil {
			return nil, err
		}
		if id != "" {
			_, err := l.is.Get(image.IDFromDigest(id))
			if err != nil {
				return nil, err
			}
			addAssoc(image.IDFromDigest(id), nil)
			continue
		}
		if ref.Name() == string(digest.Canonical) {
			imgID, err := l.is.Search(name)
			if err != nil {
				return nil, err
			}
			addAssoc(imgID, nil)
			continue
		}
		if reference.IsNameOnly(ref) {
			assocs := l.rs.ReferencesByName(ref)
			for _, assoc := range assocs {
				addAssoc(image.IDFromDigest(assoc.ID), assoc.Ref)
			}
			if len(assocs) == 0 {
				imgID, err := l.is.Search(name)
				if err != nil {
					return nil, err
				}
				addAssoc(imgID, nil)
			}
			continue
		}
		id, err = l.rs.Get(ref)
		if err != nil {
			return nil, err
		}
		addAssoc(image.IDFromDigest(id), ref)

	}
	return imgDescr, nil
}

func (s *saveSession) save(outStream io.Writer) error {
	s.savedLayers = make(map[string]struct{})
	s.diffIDPaths = make(map[layer.DiffID]string)

	// get image json
	tempDir, err := ioutil.TempDir("", "docker-export-")
	if err != nil {
		return err
	}
	defer os.RemoveAll(tempDir)

	s.outDir = tempDir
	reposLegacy := make(map[string]map[string]string)

	var manifest []manifestItem
	var parentLinks []parentLink

	for id, imageDescr := range s.images {
		foreignSrcs, err := s.saveImage(id)
		if err != nil {
			return err
		}

		var repoTags []string
		var layers []string

		for _, ref := range imageDescr.refs {
			if _, ok := reposLegacy[ref.Name()]; !ok {
				reposLegacy[ref.Name()] = make(map[string]string)
			}
			reposLegacy[ref.Name()][ref.Tag()] = imageDescr.layers[len(imageDescr.layers)-1]
			repoTags = append(repoTags, ref.String())
		}

		for _, l := range imageDescr.layers {
			layers = append(layers, filepath.Join(l, legacyLayerFileName))
		}

		manifest = append(manifest, manifestItem{
			Config:       id.Digest().Hex() + ".json",
			RepoTags:     repoTags,
			Layers:       layers,
			LayerSources: foreignSrcs,
		})

		parentID, _ := s.is.GetParent(id)
		parentLinks = append(parentLinks, parentLink{id, parentID})
		s.tarexporter.loggerImgEvent.LogImageEvent(id.String(), id.String(), "save")
	}

	for i, p := range validatedParentLinks(parentLinks) {
		if p.parentID != "" {
			manifest[i].Parent = p.parentID
		}
	}

	if len(reposLegacy) > 0 {
		reposFile := filepath.Join(tempDir, legacyRepositoriesFileName)
		rf, err := os.OpenFile(reposFile, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0644)
		if err != nil {
			return err
		}

		if err := json.NewEncoder(rf).Encode(reposLegacy); err != nil {
			rf.Close()
			return err
		}

		rf.Close()

		if err := system.Chtimes(reposFile, time.Unix(0, 0), time.Unix(0, 0)); err != nil {
			return err
		}
	}

	manifestFileName := filepath.Join(tempDir, manifestFileName)
	f, err := os.OpenFile(manifestFileName, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0644)
	if err != nil {
		return err
	}

	if err := json.NewEncoder(f).Encode(manifest); err != nil {
		f.Close()
		return err
	}

	f.Close()

	if err := system.Chtimes(manifestFileName, time.Unix(0, 0), time.Unix(0, 0)); err != nil {
		return err
	}

	fs, err := archive.Tar(tempDir, archive.Uncompressed)
	if err != nil {
		return err
	}
	defer fs.Close()

	if _, err := io.Copy(outStream, fs); err != nil {
		return err
	}
	return nil
}

func (s *saveSession) saveImage(id image.ID) (map[layer.DiffID]distribution.Descriptor, error) {
	img, err := s.is.Get(id)
	if err != nil {
		return nil, err
	}

	if len(img.RootFS.DiffIDs) == 0 {
		return nil, fmt.Errorf("empty export - not implemented")
	}

	var parent digest.Digest
	var layers []string
	var foreignSrcs map[layer.DiffID]distribution.Descriptor
	for i := range img.RootFS.DiffIDs {
		v1Img := image.V1Image{}
		if i == len(img.RootFS.DiffIDs)-1 {
			v1Img = img.V1Image
		}
		rootFS := *img.RootFS
		rootFS.DiffIDs = rootFS.DiffIDs[:i+1]
		v1ID, err := v1.CreateID(v1Img, rootFS.ChainID(), parent)
		if err != nil {
			return nil, err
		}

		v1Img.ID = v1ID.Hex()
		if parent != "" {
			v1Img.Parent = parent.Hex()
		}

		src, err := s.saveLayer(rootFS.ChainID(), v1Img, img.Created)
		if err != nil {
			return nil, err
		}
		layers = append(layers, v1Img.ID)
		parent = v1ID
		if src.Digest != "" {
			if foreignSrcs == nil {
				foreignSrcs = make(map[layer.DiffID]distribution.Descriptor)
			}
			foreignSrcs[img.RootFS.DiffIDs[i]] = src
		}
	}

	configFile := filepath.Join(s.outDir, id.Digest().Hex()+".json")
	if err := ioutil.WriteFile(configFile, img.RawJSON(), 0644); err != nil {
		return nil, err
	}
	if err := system.Chtimes(configFile, img.Created, img.Created); err != nil {
		return nil, err
	}

	s.images[id].layers = layers
	return foreignSrcs, nil
}

func (s *saveSession) saveLayer(id layer.ChainID, legacyImg image.V1Image, createdTime time.Time) (distribution.Descriptor, error) {
	if _, exists := s.savedLayers[legacyImg.ID]; exists {
		return distribution.Descriptor{}, nil
	}

	outDir := filepath.Join(s.outDir, legacyImg.ID)
	if err := os.Mkdir(outDir, 0755); err != nil {
		return distribution.Descriptor{}, err
	}

	// todo: why is this version file here?
	if err := ioutil.WriteFile(filepath.Join(outDir, legacyVersionFileName), []byte("1.0"), 0644); err != nil {
		return distribution.Descriptor{}, err
	}

	imageConfig, err := json.Marshal(legacyImg)
	if err != nil {
		return distribution.Descriptor{}, err
	}

	if err := ioutil.WriteFile(filepath.Join(outDir, legacyConfigFileName), imageConfig, 0644); err != nil {
		return distribution.Descriptor{}, err
	}

	// serialize filesystem
	layerPath := filepath.Join(outDir, legacyLayerFileName)
	l, err := s.ls.Get(id)
	if err != nil {
		return distribution.Descriptor{}, err
	}
	defer layer.ReleaseAndLog(s.ls, l)

	if oldPath, exists := s.diffIDPaths[l.DiffID()]; exists {
		relPath, err := filepath.Rel(outDir, oldPath)
		if err != nil {
			return distribution.Descriptor{}, err
		}
		os.Symlink(relPath, layerPath)
	} else {
		// Use system.CreateSequential rather than os.Create. This ensures sequential
		// file access on Windows to avoid eating into MM standby list.
		// On Linux, this equates to a regular os.Create.
		tarFile, err := system.CreateSequential(layerPath)
		if err != nil {
			return distribution.Descriptor{}, err
		}
		defer tarFile.Close()

		arch, err := l.TarStream()
		if err != nil {
			return distribution.Descriptor{}, err
		}
		defer arch.Close()

		if _, err := io.Copy(tarFile, arch); err != nil {
			return distribution.Descriptor{}, err
		}

		for _, fname := range []string{"", legacyVersionFileName, legacyConfigFileName, legacyLayerFileName} {
			// todo: maybe save layer created timestamp?
			if err := system.Chtimes(filepath.Join(outDir, fname), createdTime, createdTime); err != nil {
				return distribution.Descriptor{}, err
			}
		}

		s.diffIDPaths[l.DiffID()] = layerPath
	}
	s.savedLayers[legacyImg.ID] = struct{}{}

	var src distribution.Descriptor
	if fs, ok := l.(distribution.Describable); ok {
		src = fs.Descriptor()
	}
	return src, nil
}
<commit_msg>Fix docker save with empty timestamp of layer created time<commit_after>package tarexport

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"time"

	"github.com/docker/distribution"
	"github.com/docker/distribution/digest"
	"github.com/docker/docker/image"
	"github.com/docker/docker/image/v1"
	"github.com/docker/docker/layer"
	"github.com/docker/docker/pkg/archive"
	"github.com/docker/docker/pkg/system"
	"github.com/docker/docker/reference"
)

type imageDescriptor struct {
	refs   []reference.NamedTagged
	layers []string
}

type saveSession struct {
	*tarexporter
	outDir      string
	images      map[image.ID]*imageDescriptor
	savedLayers map[string]struct{}
	diffIDPaths map[layer.DiffID]string // cache every diffID blob to avoid duplicates
}

func (l *tarexporter) Save(names []string, outStream io.Writer) error {
	images, err := l.parseNames(names)
	if err != nil {
		return err
	}

	return (&saveSession{tarexporter: l, images: images}).save(outStream)
}

func (l *tarexporter) parseNames(names []string) (map[image.ID]*imageDescriptor, error) {
	imgDescr := make(map[image.ID]*imageDescriptor)

	addAssoc := func(id image.ID, ref reference.Named) {
		if _, ok := imgDescr[id]; !ok {
			imgDescr[id] = &imageDescriptor{}
		}

		if ref != nil {
			var tagged reference.NamedTagged
			if _, ok := ref.(reference.Canonical); ok {
				return
			}
			var ok bool
			if tagged, ok = ref.(reference.NamedTagged); !ok {
				var err error
				if tagged, err = reference.WithTag(ref, reference.DefaultTag); err != nil {
					return
				}
			}

			for _, t := range imgDescr[id].refs {
				if tagged.String() == t.String() {
					return
				}
			}
			imgDescr[id].refs = append(imgDescr[id].refs, tagged)
		}
	}

	for _, name := range names {
		id, ref, err := reference.ParseIDOrReference(name)
		if err != nil {
			return nil, err
		}
		if id != "" {
			_, err := l.is.Get(image.IDFromDigest(id))
			if err != nil {
				return nil, err
			}
			addAssoc(image.IDFromDigest(id), nil)
			continue
		}
		if ref.Name() == string(digest.Canonical) {
			imgID, err := l.is.Search(name)
			if err != nil {
				return nil, err
			}
			addAssoc(imgID, nil)
			continue
		}
		if reference.IsNameOnly(ref) {
			assocs := l.rs.ReferencesByName(ref)
			for _, assoc := range assocs {
				addAssoc(image.IDFromDigest(assoc.ID), assoc.Ref)
			}
			if len(assocs) == 0 {
				imgID, err := l.is.Search(name)
				if err != nil {
					return nil, err
				}
				addAssoc(imgID, nil)
			}
			continue
		}
		id, err = l.rs.Get(ref)
		if err != nil {
			return nil, err
		}
		addAssoc(image.IDFromDigest(id), ref)

	}
	return imgDescr, nil
}

func (s *saveSession) save(outStream io.Writer) error {
	s.savedLayers = make(map[string]struct{})
	s.diffIDPaths = make(map[layer.DiffID]string)

	// get image json
	tempDir, err := ioutil.TempDir("", "docker-export-")
	if err != nil {
		return err
	}
	defer os.RemoveAll(tempDir)

	s.outDir = tempDir
	reposLegacy := make(map[string]map[string]string)

	var manifest []manifestItem
	var parentLinks []parentLink

	for id, imageDescr := range s.images {
		foreignSrcs, err := s.saveImage(id)
		if err != nil {
			return err
		}

		var repoTags []string
		var layers []string

		for _, ref := range imageDescr.refs {
			if _, ok := reposLegacy[ref.Name()]; !ok {
				reposLegacy[ref.Name()] = make(map[string]string)
			}
			reposLegacy[ref.Name()][ref.Tag()] = imageDescr.layers[len(imageDescr.layers)-1]
			repoTags = append(repoTags, ref.String())
		}

		for _, l := range imageDescr.layers {
			layers = append(layers, filepath.Join(l, legacyLayerFileName))
		}

		manifest = append(manifest, manifestItem{
			Config:       id.Digest().Hex() + ".json",
			RepoTags:     repoTags,
			Layers:       layers,
			LayerSources: foreignSrcs,
		})

		parentID, _ := s.is.GetParent(id)
		parentLinks = append(parentLinks, parentLink{id, parentID})
		s.tarexporter.loggerImgEvent.LogImageEvent(id.String(), id.String(), "save")
	}

	for i, p := range validatedParentLinks(parentLinks) {
		if p.parentID != "" {
			manifest[i].Parent = p.parentID
		}
	}

	if len(reposLegacy) > 0 {
		reposFile := filepath.Join(tempDir, legacyRepositoriesFileName)
		rf, err := os.OpenFile(reposFile, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0644)
		if err != nil {
			return err
		}

		if err := json.NewEncoder(rf).Encode(reposLegacy); err != nil {
			rf.Close()
			return err
		}

		rf.Close()

		if err := system.Chtimes(reposFile, time.Unix(0, 0), time.Unix(0, 0)); err != nil {
			return err
		}
	}

	manifestFileName := filepath.Join(tempDir, manifestFileName)
	f, err := os.OpenFile(manifestFileName, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0644)
	if err != nil {
		return err
	}

	if err := json.NewEncoder(f).Encode(manifest); err != nil {
		f.Close()
		return err
	}

	f.Close()

	if err := system.Chtimes(manifestFileName, time.Unix(0, 0), time.Unix(0, 0)); err != nil {
		return err
	}

	fs, err := archive.Tar(tempDir, archive.Uncompressed)
	if err != nil {
		return err
	}
	defer fs.Close()

	if _, err := io.Copy(outStream, fs); err != nil {
		return err
	}
	return nil
}

func (s *saveSession) saveImage(id image.ID) (map[layer.DiffID]distribution.Descriptor, error) {
	img, err := s.is.Get(id)
	if err != nil {
		return nil, err
	}

	if len(img.RootFS.DiffIDs) == 0 {
		return nil, fmt.Errorf("empty export - not implemented")
	}

	var parent digest.Digest
	var layers []string
	var foreignSrcs map[layer.DiffID]distribution.Descriptor
	for i := range img.RootFS.DiffIDs {
		v1Img := image.V1Image{
			Created: img.Created,
		}
		if i == len(img.RootFS.DiffIDs)-1 {
			v1Img = img.V1Image
		}
		rootFS := *img.RootFS
		rootFS.DiffIDs = rootFS.DiffIDs[:i+1]
		v1ID, err := v1.CreateID(v1Img, rootFS.ChainID(), parent)
		if err != nil {
			return nil, err
		}

		v1Img.ID = v1ID.Hex()
		if parent != "" {
			v1Img.Parent = parent.Hex()
		}

		src, err := s.saveLayer(rootFS.ChainID(), v1Img, img.Created)
		if err != nil {
			return nil, err
		}
		layers = append(layers, v1Img.ID)
		parent = v1ID
		if src.Digest != "" {
			if foreignSrcs == nil {
				foreignSrcs = make(map[layer.DiffID]distribution.Descriptor)
			}
			foreignSrcs[img.RootFS.DiffIDs[i]] = src
		}
	}

	configFile := filepath.Join(s.outDir, id.Digest().Hex()+".json")
	if err := ioutil.WriteFile(configFile, img.RawJSON(), 0644); err != nil {
		return nil, err
	}
	if err := system.Chtimes(configFile, img.Created, img.Created); err != nil {
		return nil, err
	}

	s.images[id].layers = layers
	return foreignSrcs, nil
}

func (s *saveSession) saveLayer(id layer.ChainID, legacyImg image.V1Image, createdTime time.Time) (distribution.Descriptor, error) {
	if _, exists := s.savedLayers[legacyImg.ID]; exists {
		return distribution.Descriptor{}, nil
	}

	outDir := filepath.Join(s.outDir, legacyImg.ID)
	if err := os.Mkdir(outDir, 0755); err != nil {
		return distribution.Descriptor{}, err
	}

	// todo: why is this version file here?
	if err := ioutil.WriteFile(filepath.Join(outDir, legacyVersionFileName), []byte("1.0"), 0644); err != nil {
		return distribution.Descriptor{}, err
	}

	imageConfig, err := json.Marshal(legacyImg)
	if err != nil {
		return distribution.Descriptor{}, err
	}

	if err := ioutil.WriteFile(filepath.Join(outDir, legacyConfigFileName), imageConfig, 0644); err != nil {
		return distribution.Descriptor{}, err
	}

	// serialize filesystem
	layerPath := filepath.Join(outDir, legacyLayerFileName)
	l, err := s.ls.Get(id)
	if err != nil {
		return distribution.Descriptor{}, err
	}
	defer layer.ReleaseAndLog(s.ls, l)

	if oldPath, exists := s.diffIDPaths[l.DiffID()]; exists {
		relPath, err := filepath.Rel(outDir, oldPath)
		if err != nil {
			return distribution.Descriptor{}, err
		}
		os.Symlink(relPath, layerPath)
	} else {
		// Use system.CreateSequential rather than os.Create. This ensures sequential
		// file access on Windows to avoid eating into MM standby list.
		// On Linux, this equates to a regular os.Create.
		tarFile, err := system.CreateSequential(layerPath)
		if err != nil {
			return distribution.Descriptor{}, err
		}
		defer tarFile.Close()

		arch, err := l.TarStream()
		if err != nil {
			return distribution.Descriptor{}, err
		}
		defer arch.Close()

		if _, err := io.Copy(tarFile, arch); err != nil {
			return distribution.Descriptor{}, err
		}

		for _, fname := range []string{"", legacyVersionFileName, legacyConfigFileName, legacyLayerFileName} {
			// todo: maybe save layer created timestamp?
			if err := system.Chtimes(filepath.Join(outDir, fname), createdTime, createdTime); err != nil {
				return distribution.Descriptor{}, err
			}
		}

		s.diffIDPaths[l.DiffID()] = layerPath
	}
	s.savedLayers[legacyImg.ID] = struct{}{}

	var src distribution.Descriptor
	if fs, ok := l.(distribution.Describable); ok {
		src = fs.Descriptor()
	}
	return src, nil
}
<|endoftext|>


<commit_before>package utils

import (
	"crypto/rand"
	"math/big"

	"github.com/freeusd/solebtc/models"
)

// RandomReward generates a random reward with rates given
func RandomReward(rates []models.RewardRate) int64 {
	sum := sumOfWeights(rates)
	if sum < 1 {
		panic("sum of reward rates weight should be greater than 0")
	}

	i := 0
	for r := randInt64(0, sum); i < len(rates); i++ {
		r -= rates[i].Weight
		if r < 0 {
			break
		}
	}

	rate := rates[i]
	return randInt64(rate.Min, rate.Max)
}

func sumOfWeights(rates []models.RewardRate) (sum int64) {
	for i := range rates {
		sum += rates[i].Weight
	}
	return
}

func randInt64(min, max int64) int64 {
	// panic if rand.Int returns error, fail fast here
	n, _ := rand.Int(rand.Reader, big.NewInt(max-min))
	return min + n.Int64()
}
<commit_msg>Move sumOfWeights inline of RandomReward<commit_after>package utils

import (
	"crypto/rand"
	"math/big"

	"github.com/freeusd/solebtc/models"
)

// RandomReward generates a random reward with rates given
func RandomReward(rates []models.RewardRate) int64 {
	var sum int64
	for i := range rates {
		sum += rates[i].Weight
	}
	if sum < 1 {
		panic("sum of reward rates weight should be greater than 0")
	}

	i := 0
	for r := randInt64(0, sum); i < len(rates); i++ {
		r -= rates[i].Weight
		if r < 0 {
			break
		}
	}

	rate := rates[i]
	return randInt64(rate.Min, rate.Max)
}

func randInt64(min, max int64) int64 {
	// panic if rand.Int returns error, fail fast here
	n, _ := rand.Int(rand.Reader, big.NewInt(max-min))
	return min + n.Int64()
}
<|endoftext|>


<commit_before>// Copyright 2017 Pilosa Corp.
//
// 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 pilosa

import (
	"bufio"
	"bytes"
	"fmt"
	"io/ioutil"
	"path/filepath"
	"sync"
	"time"

	"github.com/gogo/protobuf/proto"
	"github.com/pkg/errors"
)

// NewTestCluster returns a cluster with n nodes and uses a mod-based hasher.
func NewTestCluster(n int) *cluster {
	path, err := ioutil.TempDir("", "pilosa-cluster-")
	if err != nil {
		panic(err)
	}

	c := newCluster()
	c.ReplicaN = 1
	c.Hasher = NewTestModHasher()
	c.Path = path
	c.Topology = newTopology()

	for i := 0; i < n; i++ {
		c.nodes = append(c.nodes, &Node{
			ID:  fmt.Sprintf("node%d", i),
			URI: NewTestURI("http", fmt.Sprintf("host%d", i), uint16(0)),
		})
	}

	c.Node = c.nodes[0]
	c.Coordinator = c.nodes[0].ID
	c.SetState(ClusterStateNormal)

	return c
}

// NewTestURI is a test URI creator that intentionally swallows errors.
func NewTestURI(scheme, host string, port uint16) URI {
	uri := defaultURI()
	_ = uri.setScheme(scheme)
	_ = uri.setHost(host)
	uri.SetPort(port)
	return *uri
}

func NewTestURIFromHostPort(host string, port uint16) URI {
	uri := defaultURI()
	_ = uri.setHost(host)
	uri.SetPort(port)
	return *uri
}

// ModHasher represents a simple, mod-based hashing.
type TestModHasher struct{}

// NewTestModHasher returns a new instance of ModHasher with n buckets.
func NewTestModHasher() *TestModHasher { return &TestModHasher{} }

func (*TestModHasher) Hash(key uint64, n int) int { return int(key) % n }

// ClusterCluster represents a cluster of test nodes, each of which
// has a Cluster.
// ClusterCluster implements Broadcaster interface.
type ClusterCluster struct {
	Clusters []*cluster

	common *commonClusterSettings

	mu         sync.RWMutex
	resizing   bool
	resizeDone chan struct{}
}

type commonClusterSettings struct {
	Nodes []*Node
}

func (t *ClusterCluster) CreateIndex(name string) error {
	for _, c := range t.Clusters {
		if _, err := c.holder.CreateIndexIfNotExists(name, IndexOptions{}); err != nil {
			return err
		}
	}
	return nil
}

func (t *ClusterCluster) CreateField(index, field string, opts FieldOption) error {
	for _, c := range t.Clusters {
		idx, err := c.holder.CreateIndexIfNotExists(index, IndexOptions{})
		if err != nil {
			return err
		}
		if _, err := idx.CreateField(field, opts); err != nil {
			return err
		}
	}
	return nil
}

func (t *ClusterCluster) SetBit(index, field string, rowID, colID uint64, x *time.Time) error {
	// Determine which node should receive the SetBit.
	c0 := t.Clusters[0] // use the first node's cluster to determine shard location.
	shard := colID / ShardWidth
	nodes := c0.shardNodes(index, shard)

	for _, node := range nodes {
		c := t.clusterByID(node.ID)
		if c == nil {
			continue
		}
		f := c.holder.Field(index, field)
		if f == nil {
			return fmt.Errorf("index/field does not exist: %s/%s", index, field)
		}
		_, err := f.SetBit(rowID, colID, x)
		if err != nil {
			return err
		}
	}

	return nil
}

func (t *ClusterCluster) clusterByID(id string) *cluster {
	for _, c := range t.Clusters {
		if c.Node.ID == id {
			return c
		}
	}
	return nil
}

// addNode adds a node to the cluster and (potentially) starts a resize job.
func (t *ClusterCluster) addNode() error {
	id := len(t.Clusters)

	c, err := t.addCluster(id, false)
	if err != nil {
		return err
	}

	// Send NodeJoin event to coordinator.
	if id > 0 {
		coord := t.Clusters[0]
		ev := &NodeEvent{
			Event: NodeJoin,
			Node:  c.Node,
		}

		if err := coord.ReceiveEvent(ev); err != nil {
			return err
		}

		// Wait for the AddNode job to finish.
		if c.State() != ClusterStateNormal {
			t.resizeDone = make(chan struct{})
			t.mu.Lock()
			t.resizing = true
			t.mu.Unlock()
			<-t.resizeDone
		}
	}

	return nil
}

// WriteTopology writes the given topology to disk.
func (t *ClusterCluster) WriteTopology(path string, top *Topology) error {
	if buf, err := proto.Marshal(top.encode()); err != nil {
		return err
	} else if err := ioutil.WriteFile(filepath.Join(path, ".topology"), buf, 0666); err != nil {
		return err
	}
	return nil
}

func (t *ClusterCluster) addCluster(i int, saveTopology bool) (*cluster, error) {

	id := fmt.Sprintf("node%d", i)
	uri := NewTestURI("http", fmt.Sprintf("host%d", i), uint16(0))

	node := &Node{
		ID:  id,
		URI: uri,
	}

	// add URI to common
	//t.common.NodeIDs = append(t.common.NodeIDs, id)
	//sort.Sort(t.common.NodeIDs)

	// add node to common
	t.common.Nodes = append(t.common.Nodes, node)

	// create node-specific temp directory
	path, err := ioutil.TempDir(*TempDir, fmt.Sprintf("pilosa-cluster-node-%d-", i))
	if err != nil {
		return nil, err
	}

	// holder
	h := NewHolder()
	h.Path = path

	// cluster
	c := newCluster()
	c.ReplicaN = 1
	c.Hasher = NewTestModHasher()
	c.Path = path
	c.Topology = newTopology()
	c.holder = h
	c.Node = node
	c.Coordinator = t.common.Nodes[0].ID // the first node is the coordinator
	c.broadcaster = t.broadcaster(c)

	// add nodes
	if saveTopology {
		for _, n := range t.common.Nodes {
			c.addNode(n)
		}
	}

	// Add this node to the ClusterCluster.
	t.Clusters = append(t.Clusters, c)

	return c, nil
}

// NewClusterCluster returns a new instance of test.Cluster.
func NewClusterCluster(n int) *ClusterCluster {

	tc := &ClusterCluster{
		common: &commonClusterSettings{},
	}

	// add clusters
	for i := 0; i < n; i++ {
		_, err := tc.addCluster(i, true)
		if err != nil {
			panic(err)
		}
	}
	return tc
}

// SetState sets the state of the cluster on each node.
func (t *ClusterCluster) SetState(state string) {
	for _, c := range t.Clusters {
		c.SetState(state)
	}
}

// Open opens all clusters in the test cluster.
func (t *ClusterCluster) Open() error {
	for _, c := range t.Clusters {
		if err := c.open(); err != nil {
			return err
		}
		if err := c.holder.Open(); err != nil {
			return err
		}
		if err := c.setNodeState(nodeStateReady); err != nil {
			return err
		}
	}

	// Start the listener on the coordinator.
	if len(t.Clusters) == 0 {
		return nil
	}
	t.Clusters[0].listenForJoins()

	return nil
}

// Close closes all clusters in the test cluster.
func (t *ClusterCluster) Close() error {
	for _, c := range t.Clusters {
		err := c.close()
		if err != nil {
			return err
		}
	}
	return nil
}

type bcast struct {
	t *ClusterCluster
	c *cluster
}

func (b bcast) SendSync(m Message) error {
	switch obj := m.(type) {
	case *ClusterStatus:
		// Apply the send message to all nodes (except the coordinator).
		for _, c := range b.t.Clusters {
			if c != b.c {
				c.mergeClusterStatus(obj)
			}
		}
		b.t.mu.RLock()
		if obj.State == ClusterStateNormal && b.t.resizing {
			close(b.t.resizeDone)
		}
		b.t.mu.RUnlock()
	}
	return nil
}

func (t *ClusterCluster) broadcaster(c *cluster) broadcaster {
	return bcast{
		t: t,
		c: c,
	}
}

// SendAsync is a test implemenetation of Broadcaster SendAsync method.
func (bcast) SendAsync(Message) error {
	return nil
}

// SendTo is a test implementation of Broadcaster SendTo method.
func (b bcast) SendTo(to *Node, m Message) error {
	switch obj := m.(type) {
	case *ResizeInstruction:
		err := b.t.FollowResizeInstruction(obj)
		if err != nil {
			return err
		}
	case *ResizeInstructionComplete:
		coord := b.t.clusterByID(to.ID)
		go coord.markResizeInstructionComplete(obj)
	case *ClusterStatus:
		// Apply the send message to the node.
		for _, c := range b.t.Clusters {
			if c.Node.ID == to.ID {
				c.mergeClusterStatus(obj)
			}
		}
		b.t.mu.RLock()
		if obj.State == ClusterStateNormal && b.t.resizing {
			close(b.t.resizeDone)
		}
		b.t.mu.RUnlock()
	default:
		panic(fmt.Sprintf("message not handled:\n%#v\n", obj))
	}
	return nil
}

// FollowResizeInstruction is a version of cluster.FollowResizeInstruction used for testing.
func (t *ClusterCluster) FollowResizeInstruction(instr *ResizeInstruction) error {

	// Prepare the return message.
	complete := &ResizeInstructionComplete{
		JobID: instr.JobID,
		Node:  instr.Node,
		Error: "",
	}

	// Stop processing on any error.
	if err := func() error {

		// figure out which node it was meant for, then call the operation on that cluster
		// basically need to mimic this: client.RetrieveShardFromURI(context.Background(), src.Index, src.Field, src.View, src.Shard, srcURI)
		instrNode := instr.Node
		destCluster := t.clusterByID(instrNode.ID)

		// Sync the schema received in the resize instruction.
		if err := destCluster.holder.applySchema(instr.NodeStatus.Schema); err != nil {
			return err
		}

		// Sync available shards.
		for _, is := range instr.NodeStatus.Indexes {
			for _, fs := range is.Fields {
				f := destCluster.holder.Field(is.Name, fs.Name)

				// if we don't know about a field locally, log an error because
				// fields should be created and synced prior to shard creation
				if f == nil {
					continue
				}
				if err := f.AddRemoteAvailableShards(fs.AvailableShards); err != nil {
					return errors.Wrap(err, "adding remote available shards")
				}
			}
		}

		for _, src := range instr.Sources {
			srcCluster := t.clusterByID(src.Node.ID)

			srcFragment := srcCluster.holder.fragment(src.Index, src.Field, src.View, src.Shard)
			destFragment := destCluster.holder.fragment(src.Index, src.Field, src.View, src.Shard)
			if destFragment == nil {
				// Create fragment on destination if it doesn't exist.
				f := destCluster.holder.Field(src.Index, src.Field)
				v := f.view(src.View)
				var err error
				destFragment, err = v.CreateFragmentIfNotExists(src.Shard)
				if err != nil {
					return err
				}
			}

			buf := bytes.NewBuffer(nil)

			bw := bufio.NewWriter(buf)
			br := bufio.NewReader(buf)

			// Get the fragment from source.
			if _, err := srcFragment.WriteTo(bw); err != nil {
				return err
			}

			// Flush the bufio.buf to the io.Writer (buf).
			bw.Flush()

			// Write data to destination.
			if _, err := destFragment.ReadFrom(br); err != nil {
				return err
			}
		}

		return nil
	}(); err != nil {
		complete.Error = err.Error()
	}

	node := instr.Coordinator
	return bcast{t: t}.SendTo(node, complete)
}
<commit_msg>lint fixes to cluster behavior in utils test<commit_after>// Copyright 2017 Pilosa Corp.
//
// 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 pilosa

import (
	"bufio"
	"bytes"
	"fmt"
	"io/ioutil"
	"path/filepath"
	"sync"
	"time"

	"github.com/gogo/protobuf/proto"
	"github.com/pkg/errors"
)

// NewTestCluster returns a cluster with n nodes and uses a mod-based hasher.
func NewTestCluster(n int) *cluster {
	path, err := ioutil.TempDir("", "pilosa-cluster-")
	if err != nil {
		panic(err)
	}

	c := newCluster()
	c.ReplicaN = 1
	c.Hasher = NewTestModHasher()
	c.Path = path
	c.Topology = newTopology()

	for i := 0; i < n; i++ {
		c.nodes = append(c.nodes, &Node{
			ID:  fmt.Sprintf("node%d", i),
			URI: NewTestURI("http", fmt.Sprintf("host%d", i), uint16(0)),
		})
	}

	c.Node = c.nodes[0]
	c.Coordinator = c.nodes[0].ID
	c.SetState(ClusterStateNormal)

	return c
}

// NewTestURI is a test URI creator that intentionally swallows errors.
func NewTestURI(scheme, host string, port uint16) URI {
	uri := defaultURI()
	_ = uri.setScheme(scheme)
	_ = uri.setHost(host)
	uri.SetPort(port)
	return *uri
}

func NewTestURIFromHostPort(host string, port uint16) URI {
	uri := defaultURI()
	_ = uri.setHost(host)
	uri.SetPort(port)
	return *uri
}

// ModHasher represents a simple, mod-based hashing.
type TestModHasher struct{}

// NewTestModHasher returns a new instance of ModHasher with n buckets.
func NewTestModHasher() *TestModHasher { return &TestModHasher{} }

func (*TestModHasher) Hash(key uint64, n int) int { return int(key) % n }

// ClusterCluster represents a cluster of test nodes, each of which
// has a Cluster.
// ClusterCluster implements Broadcaster interface.
type ClusterCluster struct {
	Clusters []*cluster

	common *commonClusterSettings

	mu         sync.RWMutex
	resizing   bool
	resizeDone chan struct{}
}

type commonClusterSettings struct {
	Nodes []*Node
}

func (t *ClusterCluster) CreateIndex(name string) error {
	for _, c := range t.Clusters {
		if _, err := c.holder.CreateIndexIfNotExists(name, IndexOptions{}); err != nil {
			return err
		}
	}
	return nil
}

func (t *ClusterCluster) CreateField(index, field string, opts FieldOption) error {
	for _, c := range t.Clusters {
		idx, err := c.holder.CreateIndexIfNotExists(index, IndexOptions{})
		if err != nil {
			return err
		}
		if _, err := idx.CreateField(field, opts); err != nil {
			return err
		}
	}
	return nil
}

func (t *ClusterCluster) SetBit(index, field string, rowID, colID uint64, x *time.Time) error {
	// Determine which node should receive the SetBit.
	c0 := t.Clusters[0] // use the first node's cluster to determine shard location.
	shard := colID / ShardWidth
	nodes := c0.shardNodes(index, shard)

	for _, node := range nodes {
		c := t.clusterByID(node.ID)
		if c == nil {
			continue
		}
		f := c.holder.Field(index, field)
		if f == nil {
			return fmt.Errorf("index/field does not exist: %s/%s", index, field)
		}
		_, err := f.SetBit(rowID, colID, x)
		if err != nil {
			return err
		}
	}

	return nil
}

func (t *ClusterCluster) clusterByID(id string) *cluster {
	for _, c := range t.Clusters {
		if c.Node.ID == id {
			return c
		}
	}
	return nil
}

// addNode adds a node to the cluster and (potentially) starts a resize job.
func (t *ClusterCluster) addNode() error {
	id := len(t.Clusters)

	c, err := t.addCluster(id, false)
	if err != nil {
		return err
	}

	// Send NodeJoin event to coordinator.
	if id > 0 {
		coord := t.Clusters[0]
		ev := &NodeEvent{
			Event: NodeJoin,
			Node:  c.Node,
		}

		if err := coord.ReceiveEvent(ev); err != nil {
			return err
		}

		// Wait for the AddNode job to finish.
		if c.State() != ClusterStateNormal {
			t.resizeDone = make(chan struct{})
			t.mu.Lock()
			t.resizing = true
			t.mu.Unlock()
			<-t.resizeDone
		}
	}

	return nil
}

// WriteTopology writes the given topology to disk.
func (t *ClusterCluster) WriteTopology(path string, top *Topology) error {
	if buf, err := proto.Marshal(top.encode()); err != nil {
		return err
	} else if err := ioutil.WriteFile(filepath.Join(path, ".topology"), buf, 0666); err != nil {
		return err
	}
	return nil
}

func (t *ClusterCluster) addCluster(i int, saveTopology bool) (*cluster, error) {

	id := fmt.Sprintf("node%d", i)
	uri := NewTestURI("http", fmt.Sprintf("host%d", i), uint16(0))

	node := &Node{
		ID:  id,
		URI: uri,
	}

	// add URI to common
	//t.common.NodeIDs = append(t.common.NodeIDs, id)
	//sort.Sort(t.common.NodeIDs)

	// add node to common
	t.common.Nodes = append(t.common.Nodes, node)

	// create node-specific temp directory
	path, err := ioutil.TempDir(*TempDir, fmt.Sprintf("pilosa-cluster-node-%d-", i))
	if err != nil {
		return nil, err
	}

	// holder
	h := NewHolder()
	h.Path = path

	// cluster
	c := newCluster()
	c.ReplicaN = 1
	c.Hasher = NewTestModHasher()
	c.Path = path
	c.Topology = newTopology()
	c.holder = h
	c.Node = node
	c.Coordinator = t.common.Nodes[0].ID // the first node is the coordinator
	c.broadcaster = t.broadcaster(c)

	// add nodes
	if saveTopology {
		for _, n := range t.common.Nodes {
			if err := c.addNode(n); err != nil {
				return nil, err
			}
		}
	}

	// Add this node to the ClusterCluster.
	t.Clusters = append(t.Clusters, c)

	return c, nil
}

// NewClusterCluster returns a new instance of test.Cluster.
func NewClusterCluster(n int) *ClusterCluster {

	tc := &ClusterCluster{
		common: &commonClusterSettings{},
	}

	// add clusters
	for i := 0; i < n; i++ {
		_, err := tc.addCluster(i, true)
		if err != nil {
			panic(err)
		}
	}
	return tc
}

// SetState sets the state of the cluster on each node.
func (t *ClusterCluster) SetState(state string) {
	for _, c := range t.Clusters {
		c.SetState(state)
	}
}

// Open opens all clusters in the test cluster.
func (t *ClusterCluster) Open() error {
	for _, c := range t.Clusters {
		if err := c.open(); err != nil {
			return err
		}
		if err := c.holder.Open(); err != nil {
			return err
		}
		if err := c.setNodeState(nodeStateReady); err != nil {
			return err
		}
	}

	// Start the listener on the coordinator.
	if len(t.Clusters) == 0 {
		return nil
	}
	t.Clusters[0].listenForJoins()

	return nil
}

// Close closes all clusters in the test cluster.
func (t *ClusterCluster) Close() error {
	for _, c := range t.Clusters {
		err := c.close()
		if err != nil {
			return err
		}
	}
	return nil
}

type bcast struct {
	t *ClusterCluster
	c *cluster
}

func (b bcast) SendSync(m Message) error {
	switch obj := m.(type) {
	case *ClusterStatus:
		// Apply the send message to all nodes (except the coordinator).
		for _, c := range b.t.Clusters {
			if c != b.c {
				err := c.mergeClusterStatus(obj)
				if err != nil {
					return err
				}
			}
		}
		b.t.mu.RLock()
		if obj.State == ClusterStateNormal && b.t.resizing {
			close(b.t.resizeDone)
		}
		b.t.mu.RUnlock()
	}
	return nil
}

func (t *ClusterCluster) broadcaster(c *cluster) broadcaster {
	return bcast{
		t: t,
		c: c,
	}
}

// SendAsync is a test implemenetation of Broadcaster SendAsync method.
func (bcast) SendAsync(Message) error {
	return nil
}

// SendTo is a test implementation of Broadcaster SendTo method.
func (b bcast) SendTo(to *Node, m Message) error {
	switch obj := m.(type) {
	case *ResizeInstruction:
		err := b.t.FollowResizeInstruction(obj)
		if err != nil {
			return err
		}
	case *ResizeInstructionComplete:
		coord := b.t.clusterByID(to.ID)
		// this used to be async, but that prevented us from checking
		// its error status...
		return coord.markResizeInstructionComplete(obj)
	case *ClusterStatus:
		// Apply the send message to the node.
		for _, c := range b.t.Clusters {
			if c.Node.ID == to.ID {
				c.mergeClusterStatus(obj)
			}
		}
		b.t.mu.RLock()
		if obj.State == ClusterStateNormal && b.t.resizing {
			close(b.t.resizeDone)
		}
		b.t.mu.RUnlock()
	default:
		panic(fmt.Sprintf("message not handled:\n%#v\n", obj))
	}
	return nil
}

// FollowResizeInstruction is a version of cluster.FollowResizeInstruction used for testing.
func (t *ClusterCluster) FollowResizeInstruction(instr *ResizeInstruction) error {

	// Prepare the return message.
	complete := &ResizeInstructionComplete{
		JobID: instr.JobID,
		Node:  instr.Node,
		Error: "",
	}

	// Stop processing on any error.
	if err := func() error {

		// figure out which node it was meant for, then call the operation on that cluster
		// basically need to mimic this: client.RetrieveShardFromURI(context.Background(), src.Index, src.Field, src.View, src.Shard, srcURI)
		instrNode := instr.Node
		destCluster := t.clusterByID(instrNode.ID)

		// Sync the schema received in the resize instruction.
		if err := destCluster.holder.applySchema(instr.NodeStatus.Schema); err != nil {
			return err
		}

		// Sync available shards.
		for _, is := range instr.NodeStatus.Indexes {
			for _, fs := range is.Fields {
				f := destCluster.holder.Field(is.Name, fs.Name)

				// if we don't know about a field locally, log an error because
				// fields should be created and synced prior to shard creation
				if f == nil {
					continue
				}
				if err := f.AddRemoteAvailableShards(fs.AvailableShards); err != nil {
					return errors.Wrap(err, "adding remote available shards")
				}
			}
		}

		for _, src := range instr.Sources {
			srcCluster := t.clusterByID(src.Node.ID)

			srcFragment := srcCluster.holder.fragment(src.Index, src.Field, src.View, src.Shard)
			destFragment := destCluster.holder.fragment(src.Index, src.Field, src.View, src.Shard)
			if destFragment == nil {
				// Create fragment on destination if it doesn't exist.
				f := destCluster.holder.Field(src.Index, src.Field)
				v := f.view(src.View)
				var err error
				destFragment, err = v.CreateFragmentIfNotExists(src.Shard)
				if err != nil {
					return err
				}
			}

			buf := bytes.NewBuffer(nil)

			bw := bufio.NewWriter(buf)
			br := bufio.NewReader(buf)

			// Get the fragment from source.
			if _, err := srcFragment.WriteTo(bw); err != nil {
				return err
			}

			// Flush the bufio.buf to the io.Writer (buf).
			bw.Flush()

			// Write data to destination.
			if _, err := destFragment.ReadFrom(br); err != nil {
				return err
			}
		}

		return nil
	}(); err != nil {
		complete.Error = err.Error()
	}

	node := instr.Coordinator
	return bcast{t: t}.SendTo(node, complete)
}
<|endoftext|>


<commit_before>package backends

import (
	"reflect"
	"time"

	"github.com/RichardKnop/machinery/v1/config"
	"github.com/RichardKnop/machinery/v1/log"
	"github.com/RichardKnop/machinery/v1/tasks"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

// MongodbBackend represents a MongoDB result backend
type MongodbBackend struct {
	Backend
	session              *mgo.Session
	tasksCollection      *mgo.Collection
	groupMetasCollection *mgo.Collection
}

// NewMongodbBackend creates MongodbBackend instance
func NewMongodbBackend(cnf *config.Config) Interface {
	return &MongodbBackend{Backend: New(cnf)}
}

// InitGroup creates and saves a group meta data object
func (b *MongodbBackend) InitGroup(groupUUID string, taskUUIDs []string) error {
	if err := b.connect(); err != nil {
		return err
	}

	groupMeta := &tasks.GroupMeta{
		GroupUUID: groupUUID,
		TaskUUIDs: taskUUIDs,
	}
	return b.groupMetasCollection.Insert(groupMeta)
}

// GroupCompleted returns true if all tasks in a group finished
func (b *MongodbBackend) GroupCompleted(groupUUID string, groupTaskCount int) (bool, error) {
	groupMeta, err := b.getGroupMeta(groupUUID)
	if err != nil {
		return false, err
	}

	taskStates, err := b.getStates(groupMeta.TaskUUIDs...)
	if err != nil {
		return false, err
	}

	var countSuccessTasks = 0
	for _, taskState := range taskStates {
		if taskState.IsCompleted() {
			countSuccessTasks++
		}
	}

	return countSuccessTasks == groupTaskCount, nil
}

// GroupTaskStates returns states of all tasks in the group
func (b *MongodbBackend) GroupTaskStates(groupUUID string, groupTaskCount int) ([]*tasks.TaskState, error) {
	groupMeta, err := b.getGroupMeta(groupUUID)
	if err != nil {
		return []*tasks.TaskState{}, err
	}

	return b.getStates(groupMeta.TaskUUIDs...)
}

// TriggerChord flags chord as triggered in the backend storage to make sure
// chord is never trigerred multiple times. Returns a boolean flag to indicate
// whether the worker should trigger chord (true) or no if it has been triggered
// already (false)
func (b *MongodbBackend) TriggerChord(groupUUID string) (bool, error) {
	if err := b.connect(); err != nil {
		return false, err
	}
	query := bson.M{
		"_id":             groupUUID,
		"chord_triggered": false,
	}
	change := mgo.Change{
		Update: bson.M{
			"$set": bson.M{
				"chord_triggered": true,
			},
		},
		ReturnNew: false,
	}
	_, err := b.groupMetasCollection.
		Find(query).
		Apply(change, nil)
	if err != nil {
		if err == mgo.ErrNotFound {
			log.WARNING.Printf("Chord already triggered for group %s", groupUUID)
			return false, nil
		}
		return false, err
	}
	return true, nil
}

// SetStatePending updates task state to PENDING
func (b *MongodbBackend) SetStatePending(signature *tasks.Signature) error {
	update := bson.M{"state": tasks.StatePending}
	return b.updateState(signature, update)
}

// SetStateReceived updates task state to RECEIVED
func (b *MongodbBackend) SetStateReceived(signature *tasks.Signature) error {
	update := bson.M{"state": tasks.StateReceived}
	return b.updateState(signature, update)
}

// SetStateStarted updates task state to STARTED
func (b *MongodbBackend) SetStateStarted(signature *tasks.Signature) error {
	update := bson.M{"state": tasks.StateStarted}
	return b.updateState(signature, update)
}

// SetStateRetry updates task state to RETRY
func (b *MongodbBackend) SetStateRetry(signature *tasks.Signature) error {
	update := bson.M{"state": tasks.StateRetry}
	return b.updateState(signature, update)
}

// SetStateSuccess updates task state to SUCCESS
func (b *MongodbBackend) SetStateSuccess(signature *tasks.Signature, results []*tasks.TaskResult) error {
	//edited by surendra tiwari
	var err error
	bsonResults := make([]bson.M, len(results))
	for i, result := range results {
		//to hold the json result
		bsonResult := new(bson.M)
		resultType := reflect.TypeOf(result.Value).Kind()
		if resultType == reflect.String {
			//convert type to json
			err = bson.UnmarshalJSON([]byte(result.Value.(string)), bsonResult)
			if err == nil {
				bsonResults[i] = bson.M{
					"type":  "Json",
					"value": bsonResult,
				}
			} else {
				bsonResults[i] = bson.M{
					"type":  result.Type,
					"value": result.Value,
				}
			}
		} else {
			bsonResults[i] = bson.M{
				"type":  result.Type,
				"value": result.Value,
			}
		}
	}
	update := bson.M{
		"state":   tasks.StateSuccess,
		"results": bsonResults,
	}
	return b.updateState(signature, update)
}

// SetStateFailure updates task state to FAILURE
func (b *MongodbBackend) SetStateFailure(signature *tasks.Signature, err string) error {
	update := bson.M{"state": tasks.StateFailure, "error": err}
	return b.updateState(signature, update)
}

// GetState returns the latest task state
func (b *MongodbBackend) GetState(taskUUID string) (*tasks.TaskState, error) {
	if err := b.connect(); err != nil {
		return nil, err
	}

	state := new(tasks.TaskState)
	if err := b.tasksCollection.FindId(taskUUID).One(state); err != nil {
		return nil, err
	}
	return state, nil
}

// PurgeState deletes stored task state
func (b *MongodbBackend) PurgeState(taskUUID string) error {
	if err := b.connect(); err != nil {
		return err
	}

	return b.tasksCollection.RemoveId(taskUUID)
}

// PurgeGroupMeta deletes stored group meta data
func (b *MongodbBackend) PurgeGroupMeta(groupUUID string) error {
	if err := b.connect(); err != nil {
		return err
	}

	return b.groupMetasCollection.RemoveId(groupUUID)
}

// lockGroupMeta acquires lock on groupUUID document
func (b *MongodbBackend) lockGroupMeta(groupUUID string) error {
	query := bson.M{
		"_id":  groupUUID,
		"lock": false,
	}
	change := mgo.Change{
		Update:    bson.M{
			"$set":
			bson.M{
				"lock": true
				},
			},
		ReturnNew: false,
	}
	_, err := b.groupMetasCollection.
		Find(query).
		Apply(change, nil)
	return err
}

// unlockGroupMeta releases lock on groupUUID document
func (b *MongodbBackend) unlockGroupMeta(groupUUID string) error {
	update := bson.M{"$set": bson.M{"lock": false}}
	_, err := b.groupMetasCollection.UpsertId(groupUUID, update)
	return err
}

// getGroupMeta retrieves group meta data, convenience function to avoid repetition
func (b *MongodbBackend) getGroupMeta(groupUUID string) (*tasks.GroupMeta, error) {
	if err := b.connect(); err != nil {
		return nil, err
	}

	query := bson.M{"_id": groupUUID}

	groupMeta := new(tasks.GroupMeta)
	if err := b.groupMetasCollection.Find(query).One(groupMeta); err != nil {
		return nil, err
	}
	return groupMeta, nil
}

// getStates returns multiple task states
func (b *MongodbBackend) getStates(taskUUIDs ...string) ([]*tasks.TaskState, error) {
	if err := b.connect(); err != nil {
		return nil, err
	}

	states := make([]*tasks.TaskState, 0, len(taskUUIDs))

	iter := b.tasksCollection.Find(bson.M{"_id": bson.M{"$in": taskUUIDs}}).Iter()

	state := new(tasks.TaskState)
	for iter.Next(state) {
		states = append(states, state)

		// otherwise we would end up with the last task being every element of the slice
		state = new(tasks.TaskState)
	}

	return states, nil
}

// updateState saves current task state
func (b *MongodbBackend) updateState(signature *tasks.Signature, update bson.M) error {
	if err := b.connect(); err != nil {
		return err
	}

	update = bson.M{"$set": update}
	_, err := b.tasksCollection.UpsertId(signature.UUID, update)
	if err != nil {
		return err
	}
	return nil
}

// connect returns a session if we are already connected to mongo, otherwise
// (when called for the first time) it will open a new session and ensure
// all required indexes for our collections exist
func (b *MongodbBackend) connect() error {
	if b.session != nil {
		return nil
	}

	session, err := mgo.Dial(b.cnf.ResultBackend)
	if err != nil {
		return err
	}
	b.session = session

	b.tasksCollection = b.session.DB("").C("tasks")
	b.groupMetasCollection = b.session.DB("").C("group_metas")

	return b.createMongoIndexes()
}

// createMongoIndexes ensures all indexes are in place
func (b *MongodbBackend) createMongoIndexes() error {
	indexes := []mgo.Index{
		{
			Key:         []string{"state"},
			Background:  true, // can be used while index is being built
			ExpireAfter: time.Duration(b.cnf.ResultsExpireIn) * time.Second,
		},
		{
			Key:         []string{"lock"},
			Background:  true, // can be used while index is being built
			ExpireAfter: time.Duration(b.cnf.ResultsExpireIn) * time.Second,
		},
	}

	for _, index := range indexes {
		// Check if index already exists, if it does, skip
		if err := b.tasksCollection.EnsureIndex(index); err == nil {
			log.INFO.Printf("%s index already exist, skipping create step", index.Key[0])
			continue
		}

		// Create index (keep in mind EnsureIndex is blocking operation)
		log.INFO.Printf("Creating %s index", index.Key[0])
		if err := b.tasksCollection.DropIndex(index.Key[0]); err != nil {
			return err
		}
		if err := b.tasksCollection.EnsureIndex(index); err != nil {
			return err
		}
	}

	return nil
}
<commit_msg>missed comma<commit_after>package backends

import (
	"reflect"
	"time"

	"github.com/RichardKnop/machinery/v1/config"
	"github.com/RichardKnop/machinery/v1/log"
	"github.com/RichardKnop/machinery/v1/tasks"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

// MongodbBackend represents a MongoDB result backend
type MongodbBackend struct {
	Backend
	session              *mgo.Session
	tasksCollection      *mgo.Collection
	groupMetasCollection *mgo.Collection
}

// NewMongodbBackend creates MongodbBackend instance
func NewMongodbBackend(cnf *config.Config) Interface {
	return &MongodbBackend{Backend: New(cnf)}
}

// InitGroup creates and saves a group meta data object
func (b *MongodbBackend) InitGroup(groupUUID string, taskUUIDs []string) error {
	if err := b.connect(); err != nil {
		return err
	}

	groupMeta := &tasks.GroupMeta{
		GroupUUID: groupUUID,
		TaskUUIDs: taskUUIDs,
	}
	return b.groupMetasCollection.Insert(groupMeta)
}

// GroupCompleted returns true if all tasks in a group finished
func (b *MongodbBackend) GroupCompleted(groupUUID string, groupTaskCount int) (bool, error) {
	groupMeta, err := b.getGroupMeta(groupUUID)
	if err != nil {
		return false, err
	}

	taskStates, err := b.getStates(groupMeta.TaskUUIDs...)
	if err != nil {
		return false, err
	}

	var countSuccessTasks = 0
	for _, taskState := range taskStates {
		if taskState.IsCompleted() {
			countSuccessTasks++
		}
	}

	return countSuccessTasks == groupTaskCount, nil
}

// GroupTaskStates returns states of all tasks in the group
func (b *MongodbBackend) GroupTaskStates(groupUUID string, groupTaskCount int) ([]*tasks.TaskState, error) {
	groupMeta, err := b.getGroupMeta(groupUUID)
	if err != nil {
		return []*tasks.TaskState{}, err
	}

	return b.getStates(groupMeta.TaskUUIDs...)
}

// TriggerChord flags chord as triggered in the backend storage to make sure
// chord is never trigerred multiple times. Returns a boolean flag to indicate
// whether the worker should trigger chord (true) or no if it has been triggered
// already (false)
func (b *MongodbBackend) TriggerChord(groupUUID string) (bool, error) {
	if err := b.connect(); err != nil {
		return false, err
	}
	query := bson.M{
		"_id":             groupUUID,
		"chord_triggered": false,
	}
	change := mgo.Change{
		Update: bson.M{
			"$set": bson.M{
				"chord_triggered": true,
			},
		},
		ReturnNew: false,
	}
	_, err := b.groupMetasCollection.
		Find(query).
		Apply(change, nil)
	if err != nil {
		if err == mgo.ErrNotFound {
			log.WARNING.Printf("Chord already triggered for group %s", groupUUID)
			return false, nil
		}
		return false, err
	}
	return true, nil
}

// SetStatePending updates task state to PENDING
func (b *MongodbBackend) SetStatePending(signature *tasks.Signature) error {
	update := bson.M{"state": tasks.StatePending}
	return b.updateState(signature, update)
}

// SetStateReceived updates task state to RECEIVED
func (b *MongodbBackend) SetStateReceived(signature *tasks.Signature) error {
	update := bson.M{"state": tasks.StateReceived}
	return b.updateState(signature, update)
}

// SetStateStarted updates task state to STARTED
func (b *MongodbBackend) SetStateStarted(signature *tasks.Signature) error {
	update := bson.M{"state": tasks.StateStarted}
	return b.updateState(signature, update)
}

// SetStateRetry updates task state to RETRY
func (b *MongodbBackend) SetStateRetry(signature *tasks.Signature) error {
	update := bson.M{"state": tasks.StateRetry}
	return b.updateState(signature, update)
}

// SetStateSuccess updates task state to SUCCESS
func (b *MongodbBackend) SetStateSuccess(signature *tasks.Signature, results []*tasks.TaskResult) error {
	//edited by surendra tiwari
	var err error
	bsonResults := make([]bson.M, len(results))
	for i, result := range results {
		//to hold the json result
		bsonResult := new(bson.M)
		resultType := reflect.TypeOf(result.Value).Kind()
		if resultType == reflect.String {
			//convert type to json
			err = bson.UnmarshalJSON([]byte(result.Value.(string)), bsonResult)
			if err == nil {
				bsonResults[i] = bson.M{
					"type":  "Json",
					"value": bsonResult,
				}
			} else {
				bsonResults[i] = bson.M{
					"type":  result.Type,
					"value": result.Value,
				}
			}
		} else {
			bsonResults[i] = bson.M{
				"type":  result.Type,
				"value": result.Value,
			}
		}
	}
	update := bson.M{
		"state":   tasks.StateSuccess,
		"results": bsonResults,
	}
	return b.updateState(signature, update)
}

// SetStateFailure updates task state to FAILURE
func (b *MongodbBackend) SetStateFailure(signature *tasks.Signature, err string) error {
	update := bson.M{"state": tasks.StateFailure, "error": err}
	return b.updateState(signature, update)
}

// GetState returns the latest task state
func (b *MongodbBackend) GetState(taskUUID string) (*tasks.TaskState, error) {
	if err := b.connect(); err != nil {
		return nil, err
	}

	state := new(tasks.TaskState)
	if err := b.tasksCollection.FindId(taskUUID).One(state); err != nil {
		return nil, err
	}
	return state, nil
}

// PurgeState deletes stored task state
func (b *MongodbBackend) PurgeState(taskUUID string) error {
	if err := b.connect(); err != nil {
		return err
	}

	return b.tasksCollection.RemoveId(taskUUID)
}

// PurgeGroupMeta deletes stored group meta data
func (b *MongodbBackend) PurgeGroupMeta(groupUUID string) error {
	if err := b.connect(); err != nil {
		return err
	}

	return b.groupMetasCollection.RemoveId(groupUUID)
}

// lockGroupMeta acquires lock on groupUUID document
func (b *MongodbBackend) lockGroupMeta(groupUUID string) error {
	query := bson.M{
		"_id":  groupUUID,
		"lock": false,
	}
	change := mgo.Change{
		Update:    bson.M{
			"$set":
			bson.M{
				"lock": true,
			},
		},
		ReturnNew: false,
	}
	_, err := b.groupMetasCollection.
		Find(query).
		Apply(change, nil)
	return err
}

// unlockGroupMeta releases lock on groupUUID document
func (b *MongodbBackend) unlockGroupMeta(groupUUID string) error {
	update := bson.M{"$set": bson.M{"lock": false}}
	_, err := b.groupMetasCollection.UpsertId(groupUUID, update)
	return err
}

// getGroupMeta retrieves group meta data, convenience function to avoid repetition
func (b *MongodbBackend) getGroupMeta(groupUUID string) (*tasks.GroupMeta, error) {
	if err := b.connect(); err != nil {
		return nil, err
	}

	query := bson.M{"_id": groupUUID}

	groupMeta := new(tasks.GroupMeta)
	if err := b.groupMetasCollection.Find(query).One(groupMeta); err != nil {
		return nil, err
	}
	return groupMeta, nil
}

// getStates returns multiple task states
func (b *MongodbBackend) getStates(taskUUIDs ...string) ([]*tasks.TaskState, error) {
	if err := b.connect(); err != nil {
		return nil, err
	}

	states := make([]*tasks.TaskState, 0, len(taskUUIDs))

	iter := b.tasksCollection.Find(bson.M{"_id": bson.M{"$in": taskUUIDs}}).Iter()

	state := new(tasks.TaskState)
	for iter.Next(state) {
		states = append(states, state)

		// otherwise we would end up with the last task being every element of the slice
		state = new(tasks.TaskState)
	}

	return states, nil
}

// updateState saves current task state
func (b *MongodbBackend) updateState(signature *tasks.Signature, update bson.M) error {
	if err := b.connect(); err != nil {
		return err
	}

	update = bson.M{"$set": update}
	_, err := b.tasksCollection.UpsertId(signature.UUID, update)
	if err != nil {
		return err
	}
	return nil
}

// connect returns a session if we are already connected to mongo, otherwise
// (when called for the first time) it will open a new session and ensure
// all required indexes for our collections exist
func (b *MongodbBackend) connect() error {
	if b.session != nil {
		return nil
	}

	session, err := mgo.Dial(b.cnf.ResultBackend)
	if err != nil {
		return err
	}
	b.session = session

	b.tasksCollection = b.session.DB("").C("tasks")
	b.groupMetasCollection = b.session.DB("").C("group_metas")

	return b.createMongoIndexes()
}

// createMongoIndexes ensures all indexes are in place
func (b *MongodbBackend) createMongoIndexes() error {
	indexes := []mgo.Index{
		{
			Key:         []string{"state"},
			Background:  true, // can be used while index is being built
			ExpireAfter: time.Duration(b.cnf.ResultsExpireIn) * time.Second,
		},
		{
			Key:         []string{"lock"},
			Background:  true, // can be used while index is being built
			ExpireAfter: time.Duration(b.cnf.ResultsExpireIn) * time.Second,
		},
	}

	for _, index := range indexes {
		// Check if index already exists, if it does, skip
		if err := b.tasksCollection.EnsureIndex(index); err == nil {
			log.INFO.Printf("%s index already exist, skipping create step", index.Key[0])
			continue
		}

		// Create index (keep in mind EnsureIndex is blocking operation)
		log.INFO.Printf("Creating %s index", index.Key[0])
		if err := b.tasksCollection.DropIndex(index.Key[0]); err != nil {
			return err
		}
		if err := b.tasksCollection.EnsureIndex(index); err != nil {
			return err
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/cloudformation"
	"github.com/hashicorp/terraform/helper/acctest"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

func TestAccAwsServerlessRepositoryChangeSet_basic(t *testing.T) {
	var stack cloudformation.Stack
	stackName := fmt.Sprintf("tf-acc-test-basic-%s", acctest.RandString(10))

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSCloudFormationDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAwsServerlessRepositoryApplicationConfig(stackName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckerverlessRepositoryApplicationExists("aws_serverlessrepository_stack.postgres-rotator", &stack),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "application_id", "arn:aws:serverlessrepo:us-east-1:297356227824:applications/SecretsManagerRDSPostgreSQLRotationSingleUser"),
					resource.TestCheckResourceAttrSet("aws_serverlessrepository_stack.postgres-rotator", "semantic_version"),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "parameters.%", "2"),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "parameters.functionName", fmt.Sprintf("func-%s", stackName)),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "parameters.endpoint", "secretsmanager.us-west-2.amazonaws.com"),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "outputs.%", "1"),
					resource.TestCheckResourceAttrSet("aws_serverlessrepository_stack.postgres-rotator", "outputs.RotationLambdaARN"),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "capabilities.#", "1"),
					//resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "capabilities.0", "CAPABILITY_NAMED_IAM"),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "tags.%", "0"),
				),
			},
		},
	})
}

func TestAccAwsServerlessRepositoryApplication_versioned(t *testing.T) {
	var stack cloudformation.Stack
	stackName := fmt.Sprintf("tf-acc-test-versioned-%s", acctest.RandString(10))
	const version = "1.0.15"

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSCloudFormationDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSServerlessRepositoryApplicationConfig_versioned(stackName, version),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckerverlessRepositoryApplicationExists("aws_serverlessrepository_stack.postgres-rotator", &stack),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "semantic_version", version),
				),
			},
		},
	})
}

func TestAccAwsServerlessRepositoryApplication_tagged(t *testing.T) {
	var stack cloudformation.Stack
	stackName := fmt.Sprintf("tf-acc-test-tagged-%s", acctest.RandString(10))

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSCloudFormationDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAwsServerlessRepositoryApplicationConfig_tagged(stackName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckerverlessRepositoryApplicationExists("aws_serverlessrepository_stack.postgres-rotator", &stack),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "tags.%", "1"),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "tags.MyTag", "My value"),
				),
			},
		},
	})
}

func TestAccAwsServerlessRepositoryApplication_versionUpdate(t *testing.T) {
	var stack cloudformation.Stack
	stackName := fmt.Sprintf("tf-acc-test-update-%s", acctest.RandString(10))
	const initialVersion = "1.0.15"
	const updateVersion = "1.0.36"

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSCloudFormationDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSServerlessRepositoryApplicationConfig_versioned(stackName, initialVersion),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckerverlessRepositoryApplicationExists("aws_serverlessrepository_stack.postgres-rotator", &stack),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "application_id", "arn:aws:serverlessrepo:us-east-1:297356227824:applications/SecretsManagerRDSPostgreSQLRotationSingleUser"),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "semantic_version", initialVersion),
				),
			},
			{
				Config: testAccAWSServerlessRepositoryApplicationConfig_versioned(stackName, updateVersion),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckerverlessRepositoryApplicationExists("aws_serverlessrepository_stack.postgres-rotator", &stack),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "application_id", "arn:aws:serverlessrepo:us-east-1:297356227824:applications/SecretsManagerRDSPostgreSQLRotationSingleUser"),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "semantic_version", updateVersion),
				),
			},
		},
	})
}

func TestAccAwsServerlessRepositoryApplication_update(t *testing.T) {
	var stack cloudformation.Stack
	stackName := fmt.Sprintf("tf-acc-test-update-name-%s", acctest.RandString(10))
	const initialName = "FuncName1"
	const updatedName = "FuncName2"

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSCloudFormationDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSServerlessRepositoryApplicationConfig_updateInitial(stackName, initialName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckerverlessRepositoryApplicationExists("aws_serverlessrepository_stack.postgres-rotator", &stack),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "application_id", "arn:aws:serverlessrepo:us-east-1:297356227824:applications/SecretsManagerRDSPostgreSQLRotationSingleUser"),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "parameters.functionName", initialName),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "tags.%", "2"),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "tags.ToDelete", "ToBeDeleted"),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "tags.ToUpdate", "InitialValue"),
				),
			},
			{
				Config: testAccAWSServerlessRepositoryApplicationConfig_updateUpdated(stackName, updatedName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckerverlessRepositoryApplicationExists("aws_serverlessrepository_stack.postgres-rotator", &stack),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "parameters.functionName", updatedName),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "tags.%", "2"),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "tags.ToUpdate", "UpdatedValue"),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "tags.ToAdd", "AddedValue"),
				),
			},
		},
	})
}

func testAccAwsServerlessRepositoryApplicationConfig(stackName string) string {
	return fmt.Sprintf(`
resource "aws_serverlessrepository_stack" "postgres-rotator" {
  name           = "%[1]s"
  application_id = "arn:aws:serverlessrepo:us-east-1:297356227824:applications/SecretsManagerRDSPostgreSQLRotationSingleUser"
  parameters = {
    functionName = "func-%[1]s"
    endpoint     = "secretsmanager.us-west-2.amazonaws.com"
  }
}`, stackName)
}

func testAccAWSServerlessRepositoryApplicationConfig_updateInitial(stackName, functionName string) string {
	return fmt.Sprintf(`
resource "aws_serverlessrepository_stack" "postgres-rotator" {
  name           = "%[1]s"
  application_id = "arn:aws:serverlessrepo:us-east-1:297356227824:applications/SecretsManagerRDSPostgreSQLRotationSingleUser"
  parameters = {
    functionName = "%[2]s"
    endpoint     = "secretsmanager.us-west-2.amazonaws.com"
  }
  tags = {
	ToDelete = "ToBeDeleted"
	ToUpdate = "InitialValue"
  }
}`, stackName, functionName)
}

func testAccAWSServerlessRepositoryApplicationConfig_updateUpdated(stackName, functionName string) string {
	return fmt.Sprintf(`
resource "aws_serverlessrepository_stack" "postgres-rotator" {
  name           = "%[1]s"
  application_id = "arn:aws:serverlessrepo:us-east-1:297356227824:applications/SecretsManagerRDSPostgreSQLRotationSingleUser"
  parameters = {
    functionName = "%[2]s"
    endpoint     = "secretsmanager.us-west-2.amazonaws.com"
  }
  tags = {
	ToUpdate = "UpdatedValue"
	ToAdd    = "AddedValue"
  }
}`, stackName, functionName)
}

func testAccAWSServerlessRepositoryApplicationConfig_versioned(stackName, version string) string {
	return fmt.Sprintf(`
resource "aws_serverlessrepository_stack" "postgres-rotator" {
  name             = "%[1]s"
  application_id   = "arn:aws:serverlessrepo:us-east-1:297356227824:applications/SecretsManagerRDSPostgreSQLRotationSingleUser"
  semantic_version = "%[2]s"
  parameters = {
    functionName = "func-%[1]s"
    endpoint     = "secretsmanager.us-west-2.amazonaws.com"
  }
}`, stackName, version)
}

func testAccAwsServerlessRepositoryApplicationConfig_tagged(stackName string) string {
	return fmt.Sprintf(`
resource "aws_serverlessrepository_stack" "postgres-rotator" {
  name           = "%[1]s"
  application_id = "arn:aws:serverlessrepo:us-east-1:297356227824:applications/SecretsManagerRDSPostgreSQLRotationSingleUser"
  parameters = {
    functionName = "func-%[1]s"
    endpoint     = "secretsmanager.us-west-2.amazonaws.com"
  }
  tags = {
    MyTag = "My value"
  }
}`, stackName)
}

func testAccCheckerverlessRepositoryApplicationExists(n string, stack *cloudformation.Stack) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[n]
		if !ok {
			return fmt.Errorf("Not found: %s", n)
		}

		conn := testAccProvider.Meta().(*AWSClient).cfconn
		params := &cloudformation.DescribeStacksInput{
			StackName: aws.String(rs.Primary.ID),
		}
		resp, err := conn.DescribeStacks(params)
		if err != nil {
			return err
		}
		if len(resp.Stacks) == 0 {
			return fmt.Errorf("CloudFormation stack not found")
		}

		return nil
	}
}
<commit_msg>Renames tests for `aws_serverlessrepository_stack`<commit_after>package aws

import (
	"fmt"
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/cloudformation"
	"github.com/hashicorp/terraform/helper/acctest"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

func TestAccAwsServerlessRepositoryStack_basic(t *testing.T) {
	var stack cloudformation.Stack
	stackName := fmt.Sprintf("tf-acc-test-basic-%s", acctest.RandString(10))

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSCloudFormationDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAwsServerlessRepositoryStackConfig(stackName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckerverlessRepositoryStackExists("aws_serverlessrepository_stack.postgres-rotator", &stack),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "application_id", "arn:aws:serverlessrepo:us-east-1:297356227824:applications/SecretsManagerRDSPostgreSQLRotationSingleUser"),
					resource.TestCheckResourceAttrSet("aws_serverlessrepository_stack.postgres-rotator", "semantic_version"),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "parameters.%", "2"),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "parameters.functionName", fmt.Sprintf("func-%s", stackName)),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "parameters.endpoint", "secretsmanager.us-west-2.amazonaws.com"),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "outputs.%", "1"),
					resource.TestCheckResourceAttrSet("aws_serverlessrepository_stack.postgres-rotator", "outputs.RotationLambdaARN"),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "capabilities.#", "1"),
					//resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "capabilities.0", "CAPABILITY_NAMED_IAM"),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "tags.%", "0"),
				),
			},
		},
	})
}

func TestAccAwsServerlessRepositoryStack_versioned(t *testing.T) {
	var stack cloudformation.Stack
	stackName := fmt.Sprintf("tf-acc-test-versioned-%s", acctest.RandString(10))
	const version = "1.0.15"

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSCloudFormationDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSServerlessRepositoryStackConfig_versioned(stackName, version),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckerverlessRepositoryStackExists("aws_serverlessrepository_stack.postgres-rotator", &stack),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "semantic_version", version),
				),
			},
		},
	})
}

func TestAccAwsServerlessRepositoryStack_tagged(t *testing.T) {
	var stack cloudformation.Stack
	stackName := fmt.Sprintf("tf-acc-test-tagged-%s", acctest.RandString(10))

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSCloudFormationDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAwsServerlessRepositoryStackConfig_tagged(stackName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckerverlessRepositoryStackExists("aws_serverlessrepository_stack.postgres-rotator", &stack),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "tags.%", "1"),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "tags.MyTag", "My value"),
				),
			},
		},
	})
}

func TestAccAwsServerlessRepositoryStack_versionUpdate(t *testing.T) {
	var stack cloudformation.Stack
	stackName := fmt.Sprintf("tf-acc-test-update-%s", acctest.RandString(10))
	const initialVersion = "1.0.15"
	const updateVersion = "1.0.36"

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSCloudFormationDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSServerlessRepositoryStackConfig_versioned(stackName, initialVersion),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckerverlessRepositoryStackExists("aws_serverlessrepository_stack.postgres-rotator", &stack),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "application_id", "arn:aws:serverlessrepo:us-east-1:297356227824:applications/SecretsManagerRDSPostgreSQLRotationSingleUser"),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "semantic_version", initialVersion),
				),
			},
			{
				Config: testAccAWSServerlessRepositoryStackConfig_versioned(stackName, updateVersion),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckerverlessRepositoryStackExists("aws_serverlessrepository_stack.postgres-rotator", &stack),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "application_id", "arn:aws:serverlessrepo:us-east-1:297356227824:applications/SecretsManagerRDSPostgreSQLRotationSingleUser"),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "semantic_version", updateVersion),
				),
			},
		},
	})
}

func TestAccAwsServerlessRepositoryStack_update(t *testing.T) {
	var stack cloudformation.Stack
	stackName := fmt.Sprintf("tf-acc-test-update-name-%s", acctest.RandString(10))
	const initialName = "FuncName1"
	const updatedName = "FuncName2"

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSCloudFormationDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSServerlessRepositoryStackConfig_updateInitial(stackName, initialName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckerverlessRepositoryStackExists("aws_serverlessrepository_stack.postgres-rotator", &stack),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "application_id", "arn:aws:serverlessrepo:us-east-1:297356227824:applications/SecretsManagerRDSPostgreSQLRotationSingleUser"),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "parameters.functionName", initialName),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "tags.%", "2"),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "tags.ToDelete", "ToBeDeleted"),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "tags.ToUpdate", "InitialValue"),
				),
			},
			{
				Config: testAccAWSServerlessRepositoryStackConfig_updateUpdated(stackName, updatedName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckerverlessRepositoryStackExists("aws_serverlessrepository_stack.postgres-rotator", &stack),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "parameters.functionName", updatedName),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "tags.%", "2"),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "tags.ToUpdate", "UpdatedValue"),
					resource.TestCheckResourceAttr("aws_serverlessrepository_stack.postgres-rotator", "tags.ToAdd", "AddedValue"),
				),
			},
		},
	})
}

func testAccAwsServerlessRepositoryStackConfig(stackName string) string {
	return fmt.Sprintf(`
resource "aws_serverlessrepository_stack" "postgres-rotator" {
  name           = "%[1]s"
  application_id = "arn:aws:serverlessrepo:us-east-1:297356227824:applications/SecretsManagerRDSPostgreSQLRotationSingleUser"
  parameters = {
    functionName = "func-%[1]s"
    endpoint     = "secretsmanager.us-west-2.amazonaws.com"
  }
}`, stackName)
}

func testAccAWSServerlessRepositoryStackConfig_updateInitial(stackName, functionName string) string {
	return fmt.Sprintf(`
resource "aws_serverlessrepository_stack" "postgres-rotator" {
  name           = "%[1]s"
  application_id = "arn:aws:serverlessrepo:us-east-1:297356227824:applications/SecretsManagerRDSPostgreSQLRotationSingleUser"
  parameters = {
    functionName = "%[2]s"
    endpoint     = "secretsmanager.us-west-2.amazonaws.com"
  }
  tags = {
	ToDelete = "ToBeDeleted"
	ToUpdate = "InitialValue"
  }
}`, stackName, functionName)
}

func testAccAWSServerlessRepositoryStackConfig_updateUpdated(stackName, functionName string) string {
	return fmt.Sprintf(`
resource "aws_serverlessrepository_stack" "postgres-rotator" {
  name           = "%[1]s"
  application_id = "arn:aws:serverlessrepo:us-east-1:297356227824:applications/SecretsManagerRDSPostgreSQLRotationSingleUser"
  parameters = {
    functionName = "%[2]s"
    endpoint     = "secretsmanager.us-west-2.amazonaws.com"
  }
  tags = {
	ToUpdate = "UpdatedValue"
	ToAdd    = "AddedValue"
  }
}`, stackName, functionName)
}

func testAccAWSServerlessRepositoryStackConfig_versioned(stackName, version string) string {
	return fmt.Sprintf(`
resource "aws_serverlessrepository_stack" "postgres-rotator" {
  name             = "%[1]s"
  application_id   = "arn:aws:serverlessrepo:us-east-1:297356227824:applications/SecretsManagerRDSPostgreSQLRotationSingleUser"
  semantic_version = "%[2]s"
  parameters = {
    functionName = "func-%[1]s"
    endpoint     = "secretsmanager.us-west-2.amazonaws.com"
  }
}`, stackName, version)
}

func testAccAwsServerlessRepositoryStackConfig_tagged(stackName string) string {
	return fmt.Sprintf(`
resource "aws_serverlessrepository_stack" "postgres-rotator" {
  name           = "%[1]s"
  application_id = "arn:aws:serverlessrepo:us-east-1:297356227824:applications/SecretsManagerRDSPostgreSQLRotationSingleUser"
  parameters = {
    functionName = "func-%[1]s"
    endpoint     = "secretsmanager.us-west-2.amazonaws.com"
  }
  tags = {
    MyTag = "My value"
  }
}`, stackName)
}

func testAccCheckerverlessRepositoryStackExists(n string, stack *cloudformation.Stack) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[n]
		if !ok {
			return fmt.Errorf("Not found: %s", n)
		}

		conn := testAccProvider.Meta().(*AWSClient).cfconn
		params := &cloudformation.DescribeStacksInput{
			StackName: aws.String(rs.Primary.ID),
		}
		resp, err := conn.DescribeStacks(params)
		if err != nil {
			return err
		}
		if len(resp.Stacks) == 0 {
			return fmt.Errorf("CloudFormation stack not found")
		}

		return nil
	}
}
<|endoftext|>