<commit_before>package mongo

import (
	"bytes"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"path"
	"path/filepath"

	"github.com/juju/loggo"
	"labix.org/v2/mgo"

	"launchpad.net/juju-core/instance"
	"launchpad.net/juju-core/replicaset"
	"launchpad.net/juju-core/state/api/params"
	"launchpad.net/juju-core/upstart"
	"launchpad.net/juju-core/utils"
	"launchpad.net/juju-core/version"
)

const (
	maxFiles = 65000
	maxProcs = 20000

	serviceName = "juju-db"

	// SharedSecretFile is the name of the Mongo shared secret file
	// located within the Juju data directory.
	SharedSecretFile = "shared-secret"

	ReplicaSetName = "juju"
)

var (
	logger          = loggo.GetLogger("juju.agent.mongo")
	mongoConfigPath = "/etc/default/mongodb"

	// JujuMongodPath holds the default path to the juju-specific mongod.
	JujuMongodPath = "/usr/lib/juju/bin/mongod"

	upstartConfInstall          = (*upstart.Conf).Install
	upstartServiceStopAndRemove = (*upstart.Service).StopAndRemove
)

// WithAddresses represents an entity that has a set of
// addresses. e.g. a state Machine object
type WithAddresses interface {
	Addresses() []instance.Address
}

// IsMaster returns a boolean that represents whether the given
// machine's peer address is the primary mongo host for the replicaset
func IsMaster(session *mgo.Session, obj WithAddresses) (bool, error) {
	addrs := obj.Addresses()

	masterHostPort, err := replicaset.MasterHostPort(session)

	// If the replica set has not been configured, then we
	// can have only one master and the caller must
	// be that master.
	if err == replicaset.ErrMasterNotConfigured {
		return true, nil
	}

	if err != nil {
		return false, err
	}

	masterAddr, _, err := net.SplitHostPort(masterHostPort)
	if err != nil {
		return false, err
	}

	machinePeerAddr := SelectPeerAddress(addrs)
	return machinePeerAddr == masterAddr, nil
}

// SelectPeerAddress returns the address to use as the
// mongo replica set peer address by selecting it from the given addresses.
func SelectPeerAddress(addrs []instance.Address) string {
	return instance.SelectInternalAddress(addrs, false)
}

// SelectPeerHostPort returns the HostPort to use as the
// mongo replica set peer by selecting it from the given hostPorts.
func SelectPeerHostPort(hostPorts []instance.HostPort) string {
	return instance.SelectInternalHostPort(hostPorts, false)
}

// GenerateSharedSecret generates a pseudo-random shared secret (keyfile)
// for use with Mongo replica sets.
func GenerateSharedSecret() (string, error) {
	// "A key’s length must be between 6 and 1024 characters and may
	// only contain characters in the base64 set."
	//   -- http://docs.mongodb.org/manual/tutorial/generate-key-file/
	buf := make([]byte, base64.StdEncoding.DecodedLen(1024))
	if _, err := rand.Read(buf); err != nil {
		return "", fmt.Errorf("cannot read random secret: %v", err)
	}
	return base64.StdEncoding.EncodeToString(buf), nil
}

// MongoPath returns the executable path to be used to run mongod on this
// machine. If the juju-bundled version of mongo exists, it will return that
// path, otherwise it will return the command to run mongod from the path.
func MongodPath() (string, error) {
	if _, err := os.Stat(JujuMongodPath); err == nil {
		return JujuMongodPath, nil
	}

	path, err := exec.LookPath("mongod")
	if err != nil {
		return "", err
	}
	return path, nil
}

// RemoveService removes the mongoDB upstart service from this machine.
func RemoveService(namespace string) error {
	svc := upstart.NewService(ServiceName(namespace))
	return upstartServiceStopAndRemove(svc)
}

// EnsureMongoServer ensures that the correct mongo upstart script is installed
// and running.
//
// This method will remove old versions of the mongo upstart script as necessary
// before installing the new version.
//
// The namespace is a unique identifier to prevent multiple instances of mongo
// on this machine from colliding. This should be empty unless using
// the local provider.
func EnsureMongoServer(dataDir string, namespace string, info params.StateServingInfo) error {
	logger.Infof("Ensuring mongo server is running; dataDir %s; port %d", dataDir, info.StatePort)
	dbDir := filepath.Join(dataDir, "db")

	if err := os.MkdirAll(dbDir, 0700); err != nil {
		return fmt.Errorf("cannot create mongo dbdir: %v", err)
	}

	certKey := info.Cert + "\n" + info.PrivateKey
	err := utils.AtomicWriteFile(sslKeyPath(dataDir), []byte(certKey), 0600)
	if err != nil {
		return fmt.Errorf("cannot write SSL key: %v", err)
	}

	err = utils.AtomicWriteFile(sharedSecretPath(dataDir), []byte(info.SharedSecret), 0600)
	if err != nil {
		return fmt.Errorf("cannot write mongod shared secret: %v", err)
	}

	// Disable the default mongodb installed by the mongodb-server package.
	// Only do this if the file doesn't exist already, so users can run
	// their own mongodb server if they wish to.
	if _, err := os.Stat(mongoConfigPath); os.IsNotExist(err) {
		err = ioutil.WriteFile(
			mongoConfigPath,
			[]byte("ENABLE_MONGODB=no"),
			0644,
		)
		if err != nil {
			return err
		}
	}

	if err := aptGetInstallMongod(); err != nil {
		return fmt.Errorf("cannot install mongod: %v", err)
	}

	upstartConf, err := mongoUpstartService(namespace, dataDir, dbDir, info.StatePort)
	if err != nil {
		return err
	}
	if err := makeJournalDirs(dbDir); err != nil {
		return fmt.Errorf("Error creating journal directories: %v", err)
	}
	return upstartConfInstall(upstartConf)
}

// ServiceName returns the name of the upstart service config for mongo using
// the given namespace.
func ServiceName(namespace string) string {
	if namespace != "" {
		return fmt.Sprintf("%s-%s", serviceName, namespace)
	}
	return serviceName
}

func makeJournalDirs(dataDir string) error {
	journalDir := path.Join(dataDir, "journal")

	if err := os.MkdirAll(journalDir, 0700); err != nil {
		logger.Errorf("failed to make mongo journal dir %s: %v", journalDir, err)
		return err
	}

	// manually create the prealloc files, since otherwise they get created as 100M files.
	zeroes := make([]byte, 64*1024) // should be enough for anyone
	for x := 0; x < 3; x++ {
		name := fmt.Sprintf("prealloc.%d", x)
		filename := filepath.Join(journalDir, name)
		f, err := os.OpenFile(filename, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0700)
		if err != nil {
			return fmt.Errorf("failed to open mongo prealloc file %q: %v", filename, err)
		}
		defer f.Close()
		for total := 0; total < 1024*1024; {
			n, err := f.Write(zeroes)
			if err != nil {
				return fmt.Errorf("failed to write to mongo prealloc file %q: %v", filename, err)
			}
			total += n
		}
	}
	return nil
}

func sslKeyPath(dataDir string) string {
	return filepath.Join(dataDir, "server.pem")
}

func sharedSecretPath(dataDir string) string {
	return filepath.Join(dataDir, SharedSecretFile)
}

// mongoUpstartService returns the upstart config for the mongo state service.
//
func mongoUpstartService(namespace, dataDir, dbDir string, port int) (*upstart.Conf, error) {
	// NOTE: ensure that the right package is installed?
	svc := upstart.NewService(ServiceName(namespace))

	mongoPath, err := MongodPath()
	if err != nil {
		return nil, err
	}

	conf := &upstart.Conf{
		Service: *svc,
		Desc:    "juju state database",
		Limit: map[string]string{
			"nofile": fmt.Sprintf("%d %d", maxFiles, maxFiles),
			"nproc":  fmt.Sprintf("%d %d", maxProcs, maxProcs),
		},
		Cmd: mongoPath + " --auth" +
			" --dbpath=" + dbDir +
			" --sslOnNormalPorts" +
			" --sslPEMKeyFile " + utils.ShQuote(sslKeyPath(dataDir)) +
			" --sslPEMKeyPassword ignored" +
			" --bind_ip 0.0.0.0" +
			" --port " + fmt.Sprint(port) +
			" --noprealloc" +
			" --syslog" +
			" --smallfiles" +
			" --replSet " + ReplicaSetName +
			" --keyFile " + utils.ShQuote(sharedSecretPath(dataDir)),
	}
	return conf, nil
}

func aptGetInstallMongod() error {
	// Only Quantal requires the PPA (for mongo).
	if version.Current.Series == "quantal" {
		if err := addAptRepository("ppa:juju/stable"); err != nil {
			return err
		}
	}
	cmds := utils.AptGetPreparePackages([]string{"mongodb-server"}, version.Current.Series)
	for _, cmd := range cmds {
		if err := utils.AptGetInstall(cmd...); err != nil {
			return err
		}
	}
	return nil
}

func addAptRepository(name string) error {
	cmd := exec.Command("add-apt-repository", "-y", name)
	out, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("cannot add apt repository: %v (output %s)", err, bytes.TrimSpace(out))
	}
	return nil
}
<commit_msg>code review changes<commit_after>package mongo

import (
	"bytes"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"path"
	"path/filepath"

	"github.com/juju/loggo"
	"labix.org/v2/mgo"

	"launchpad.net/juju-core/instance"
	"launchpad.net/juju-core/replicaset"
	"launchpad.net/juju-core/state/api/params"
	"launchpad.net/juju-core/upstart"
	"launchpad.net/juju-core/utils"
	"launchpad.net/juju-core/version"
)

const (
	maxFiles = 65000
	maxProcs = 20000

	serviceName = "juju-db"

	// SharedSecretFile is the name of the Mongo shared secret file
	// located within the Juju data directory.
	SharedSecretFile = "shared-secret"

	ReplicaSetName = "juju"
)

var (
	logger          = loggo.GetLogger("juju.agent.mongo")
	mongoConfigPath = "/etc/default/mongodb"

	// JujuMongodPath holds the default path to the juju-specific mongod.
	JujuMongodPath = "/usr/lib/juju/bin/mongod"

	upstartConfInstall          = (*upstart.Conf).Install
	upstartServiceStopAndRemove = (*upstart.Service).StopAndRemove
)

// WithAddresses represents an entity that has a set of
// addresses. e.g. a state Machine object
type WithAddresses interface {
	Addresses() []instance.Address
}

// IsMaster returns a boolean that represents whether the given
// machine's peer address is the primary mongo host for the replicaset
func IsMaster(session *mgo.Session, obj WithAddresses) (bool, error) {
	addrs := obj.Addresses()

	masterHostPort, err := replicaset.MasterHostPort(session)

	// If the replica set has not been configured, then we
	// can have only one master and the caller must
	// be that master.
	if err == replicaset.ErrMasterNotConfigured {
		return true, nil
	}

	if err != nil {
		return false, err
	}

	masterAddr, _, err := net.SplitHostPort(masterHostPort)
	if err != nil {
		return false, err
	}

	machinePeerAddr := SelectPeerAddress(addrs)
	return machinePeerAddr == masterAddr, nil
}

// SelectPeerAddress returns the address to use as the
// mongo replica set peer address by selecting it from the given addresses.
func SelectPeerAddress(addrs []instance.Address) string {
	return instance.SelectInternalAddress(addrs, false)
}

// SelectPeerHostPort returns the HostPort to use as the
// mongo replica set peer by selecting it from the given hostPorts.
func SelectPeerHostPort(hostPorts []instance.HostPort) string {
	return instance.SelectInternalHostPort(hostPorts, false)
}

// GenerateSharedSecret generates a pseudo-random shared secret (keyfile)
// for use with Mongo replica sets.
func GenerateSharedSecret() (string, error) {
	// "A key’s length must be between 6 and 1024 characters and may
	// only contain characters in the base64 set."
	//   -- http://docs.mongodb.org/manual/tutorial/generate-key-file/
	buf := make([]byte, base64.StdEncoding.DecodedLen(1024))
	if _, err := rand.Read(buf); err != nil {
		return "", fmt.Errorf("cannot read random secret: %v", err)
	}
	return base64.StdEncoding.EncodeToString(buf), nil
}

// MongoPath returns the executable path to be used to run mongod on this
// machine. If the juju-bundled version of mongo exists, it will return that
// path, otherwise it will return the command to run mongod from the path.
func MongodPath() (string, error) {
	if _, err := os.Stat(JujuMongodPath); err == nil {
		return JujuMongodPath, nil
	}

	path, err := exec.LookPath("mongod")
	if err != nil {
		return "", err
	}
	return path, nil
}

// RemoveService removes the mongoDB upstart service from this machine.
func RemoveService(namespace string) error {
	svc := upstart.NewService(ServiceName(namespace))
	return upstartServiceStopAndRemove(svc)
}

// EnsureMongoServer ensures that the correct mongo upstart script is installed
// and running.
//
// This method will remove old versions of the mongo upstart script as necessary
// before installing the new version.
//
// The namespace is a unique identifier to prevent multiple instances of mongo
// on this machine from colliding. This should be empty unless using
// the local provider.
func EnsureMongoServer(dataDir string, namespace string, info params.StateServingInfo) error {
	logger.Infof("Ensuring mongo server is running; dataDir %s; port %d", dataDir, info.StatePort)
	dbDir := filepath.Join(dataDir, "db")

	if err := os.MkdirAll(dbDir, 0700); err != nil {
		return fmt.Errorf("cannot create mongo dbdir: %v", err)
	}

	certKey := info.Cert + "\n" + info.PrivateKey
	err := utils.AtomicWriteFile(sslKeyPath(dataDir), []byte(certKey), 0600)
	if err != nil {
		return fmt.Errorf("cannot write SSL key: %v", err)
	}

	err = utils.AtomicWriteFile(sharedSecretPath(dataDir), []byte(info.SharedSecret), 0600)
	if err != nil {
		return fmt.Errorf("cannot write mongod shared secret: %v", err)
	}

	// Disable the default mongodb installed by the mongodb-server package.
	// Only do this if the file doesn't exist already, so users can run
	// their own mongodb server if they wish to.
	if _, err := os.Stat(mongoConfigPath); os.IsNotExist(err) {
		err = ioutil.WriteFile(
			mongoConfigPath,
			[]byte("ENABLE_MONGODB=no"),
			0644,
		)
		if err != nil {
			return err
		}
	}

	if err := aptGetInstallMongod(); err != nil {
		return fmt.Errorf("cannot install mongod: %v", err)
	}

	upstartConf, err := mongoUpstartService(namespace, dataDir, dbDir, info.StatePort)
	if err != nil {
		return err
	}
	if err := makeJournalDirs(dbDir); err != nil {
		return fmt.Errorf("Error creating journal directories: %v", err)
	}
	return upstartConfInstall(upstartConf)
}

// ServiceName returns the name of the upstart service config for mongo using
// the given namespace.
func ServiceName(namespace string) string {
	if namespace != "" {
		return fmt.Sprintf("%s-%s", serviceName, namespace)
	}
	return serviceName
}

func makeJournalDirs(dataDir string) error {
	journalDir := path.Join(dataDir, "journal")

	if err := os.MkdirAll(journalDir, 0700); err != nil {
		logger.Errorf("failed to make mongo journal dir %s: %v", journalDir, err)
		return err
	}

	// manually create the prealloc files, since otherwise they get created as 100M files.
	zeroes := make([]byte, 64*1024) // should be enough for anyone
	for x := 0; x < 3; x++ {
		name := fmt.Sprintf("prealloc.%d", x)
		filename := filepath.Join(journalDir, name)
		f, err := os.OpenFile(filename, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0700)
		if err != nil {
			return fmt.Errorf("failed to open mongo prealloc file %q: %v", filename, err)
		}
		defer f.Close()
		for total := 0; total < 1024*1024; {
			n, err := f.Write(zeroes)
			if err != nil {
				return fmt.Errorf("failed to write to mongo prealloc file %q: %v", filename, err)
			}
			total += n
		}
	}
	return nil
}

func sslKeyPath(dataDir string) string {
	return filepath.Join(dataDir, "server.pem")
}

func sharedSecretPath(dataDir string) string {
	return filepath.Join(dataDir, SharedSecretFile)
}

// mongoUpstartService returns the upstart config for the mongo state service.
//
func mongoUpstartService(namespace, dataDir, dbDir string, port int) (*upstart.Conf, error) {
	svc := upstart.NewService(ServiceName(namespace))

	mongoPath, err := MongodPath()
	if err != nil {
		return nil, err
	}

	conf := &upstart.Conf{
		Service: *svc,
		Desc:    "juju state database",
		Limit: map[string]string{
			"nofile": fmt.Sprintf("%d %d", maxFiles, maxFiles),
			"nproc":  fmt.Sprintf("%d %d", maxProcs, maxProcs),
		},
		Cmd: mongoPath + " --auth" +
			" --dbpath=" + dbDir +
			" --sslOnNormalPorts" +
			" --sslPEMKeyFile " + utils.ShQuote(sslKeyPath(dataDir)) +
			" --sslPEMKeyPassword ignored" +
			" --bind_ip 0.0.0.0" +
			" --port " + fmt.Sprint(port) +
			" --noprealloc" +
			" --syslog" +
			" --smallfiles" +
			" --replSet " + ReplicaSetName +
			" --keyFile " + utils.ShQuote(sharedSecretPath(dataDir)),
	}
	return conf, nil
}

func aptGetInstallMongod() error {
	// Only Quantal requires the PPA.
	if version.Current.Series == "quantal" {
		if err := addAptRepository("ppa:juju/stable"); err != nil {
			return err
		}
	}
	cmds := utils.AptGetPreparePackages([]string{"mongodb-server"}, version.Current.Series)
	logger.Infof("installing mongodb-server")
	for _, cmd := range cmds {
		if err := utils.AptGetInstall(cmd...); err != nil {
			return err
		}
	}
	return nil
}

func addAptRepository(name string) error {
	// add-apt-repository requires python-software-properties
	cmds := utils.AptGetPreparePackages(
		[]string{"python-software-properties"},
		version.Current.Series,
	)
	logger.Infof("installing python-software-properties")
	for _, cmd := range cmds {
		if err := utils.AptGetInstall(cmd...); err != nil {
			return err
		}
	}

	logger.Infof("adding apt repository %q", name)
	cmd := exec.Command("add-apt-repository", "-y", name)
	out, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("cannot add apt repository: %v (output %s)", err, bytes.TrimSpace(out))
	}
	return nil
}
<|endoftext|>


<commit_before>// unify-labels is a CLI which will add, rename, or change the color of labels so they match the Jekyll org's requirements.
package main

import (
	"log"
	"strings"

	"github.com/google/go-github/github"
	"github.com/parkr/auto-reply/ctx"
)

var desiredLabels = []*github.Label{
	{Name: github.String("bug"), Color: github.String("d41313")},
	{Name: github.String("discussion"), Color: github.String("006b75")},
	{Name: github.String("documentation"), Color: github.String("006b75")},
	{Name: github.String("enhancement"), Color: github.String("009800")},
	{Name: github.String("feature"), Color: github.String("009800")},
	{Name: github.String("fix"), Color: github.String("eb6420")},
	{Name: github.String("github"), Color: github.String("222222")},
	{Name: github.String("has-pull-request"), Color: github.String("fbca04")},
	{Name: github.String("help-wanted"), Color: github.String("fbca04")},
	{Name: github.String("internal"), Color: github.String("ededed")},
	{Name: github.String("pending-feedback"), Color: github.String("fbca04")},
	{Name: github.String("pending-rebase"), Color: github.String("eb6420")},
	{Name: github.String("pinned"), Color: github.String("f3f4d3")},
	{Name: github.String("release"), Color: github.String("d4c5f9")},
	{Name: github.String("security"), Color: github.String("e11d21")},
	{Name: github.String("stale"), Color: github.String("bfd4f2")},
	{Name: github.String("suggestion"), Color: github.String("0052cc")},
	{Name: github.String("support"), Color: github.String("5319e7")},
	{Name: github.String("tests"), Color: github.String("d4c5f9")},
	{Name: github.String("undetermined"), Color: github.String("fe3868")},
	{Name: github.String("ux"), Color: github.String("006b75")},
	{Name: github.String("windows"), Color: github.String("fbca04")},
	{Name: github.String("wont-fix"), Color: github.String("e11d21")},
}
var listOpts = github.ListOptions{PerPage: 100}

func allPossibleNames(name string) []string {
	return []string{
		name,
		strings.Replace(name, "-", "", -1),
		strings.Replace(name, "-", " ", -1),
		strings.ToLower(name),
		strings.Title(name),
		strings.Title(strings.Replace(name, "-", "", -1)),
		strings.Title(strings.Replace(name, "-", " ", -1)),
	}
}

func findLabel(labels []*github.Label, desiredLabel *github.Label) *github.Label {
	possibleNames := allPossibleNames(*desiredLabel.Name)
	for _, possibleName := range possibleNames {
		for _, label := range labels {
			if *label.Name == possibleName {
				return label
			}
		}
	}

	return nil
}

func main() {
	context := ctx.NewDefaultContext()
	repos, _, err := context.GitHub.Repositories.List("jekyll", &github.RepositoryListOptions{
		Type: "owner", Sort: "full_name", Direction: "asc", ListOptions: listOpts,
	})
	if err != nil {
		log.Fatalln("error fetching repos:", err)
	}

	for _, repo := range repos {
		owner, repoName := *repo.Owner.Login, *repo.Name
		log.Printf("Processing %s", *repo.FullName)
		labels, _, err := context.GitHub.Issues.ListLabels(owner, repoName, &listOpts)
		if err != nil {
			log.Fatalf("error fetching labels for %s: %v", *repo.FullName, err)
		}

		for _, desiredLabel := range desiredLabels {
			matchedLabel := findLabel(labels, desiredLabel)

			if matchedLabel == nil {
				log.Printf("%s: creating %s with color %s",
					*repo.FullName, *desiredLabel.Name, *desiredLabel.Color)
				// It doesn't exist. Create and continue.
				_, _, err := context.GitHub.Issues.CreateLabel(owner, repoName, desiredLabel)
				if err != nil {
					log.Fatalf("error creating '%s' for %s: %v", *desiredLabel.Name, *repo.FullName, err)
				}
				continue
			}

			// It does exist, but possibly with incorrect info. Update it.
			if *matchedLabel.Name != *desiredLabel.Name || *matchedLabel.Color != *desiredLabel.Color {
				log.Printf("%s: updating %s with data: %v",
					*repo.FullName, *matchedLabel.Name, github.Stringify(desiredLabel))
				_, _, err := context.GitHub.Issues.EditLabel(owner, repoName, *matchedLabel.Name, desiredLabel)
				if err != nil {
					log.Fatalf("error updating '%s' for %s: %v", *matchedLabel.Name, *repo.FullName, err)
				}
				continue
			}
		}
	}
}
<commit_msg>cmd/unify-labels: add dry-run option and refactor a teeny bit<commit_after>// unify-labels is a CLI which will add, rename, or change the color of labels so they match the Jekyll org's requirements.
package main

import (
	"flag"
	"log"
	"strings"

	"github.com/google/go-github/github"
	"github.com/parkr/auto-reply/ctx"
)

var desiredLabels = []*github.Label{
	{Name: github.String("bug"), Color: github.String("d41313")},
	{Name: github.String("discussion"), Color: github.String("006b75")},
	{Name: github.String("documentation"), Color: github.String("006b75")},
	{Name: github.String("enhancement"), Color: github.String("009800")},
	{Name: github.String("feature"), Color: github.String("009800")},
	{Name: github.String("fix"), Color: github.String("eb6420")},
	{Name: github.String("github"), Color: github.String("222222")},
	{Name: github.String("has-pull-request"), Color: github.String("fbca04")},
	{Name: github.String("help-wanted"), Color: github.String("fbca04")},
	{Name: github.String("internal"), Color: github.String("ededed")},
	{Name: github.String("pending-feedback"), Color: github.String("fbca04")},
	{Name: github.String("pending-rebase"), Color: github.String("eb6420")},
	{Name: github.String("pinned"), Color: github.String("f3f4d3")},
	{Name: github.String("release"), Color: github.String("d4c5f9")},
	{Name: github.String("security"), Color: github.String("e11d21")},
	{Name: github.String("stale"), Color: github.String("bfd4f2")},
	{Name: github.String("suggestion"), Color: github.String("0052cc")},
	{Name: github.String("support"), Color: github.String("5319e7")},
	{Name: github.String("tests"), Color: github.String("d4c5f9")},
	{Name: github.String("undetermined"), Color: github.String("fe3868")},
	{Name: github.String("ux"), Color: github.String("006b75")},
	{Name: github.String("windows"), Color: github.String("fbca04")},
	{Name: github.String("wont-fix"), Color: github.String("e11d21")},
}
var listOpts = github.ListOptions{PerPage: 100}

func allPossibleNames(name string) []string {
	return []string{
		name,
		strings.Replace(name, "-", "", -1),
		strings.Replace(name, "-", " ", -1),
		strings.ToLower(name),
		strings.Title(name),
		strings.Title(strings.Replace(name, "-", "", -1)),
		strings.Title(strings.Replace(name, "-", " ", -1)),
	}
}

func findLabel(labels []*github.Label, desiredLabel *github.Label) *github.Label {
	possibleNames := allPossibleNames(*desiredLabel.Name)
	for _, possibleName := range possibleNames {
		for _, label := range labels {
			if *label.Name == possibleName {
				return label
			}
		}
	}

	return nil
}

func processRepo(context *ctx.Context, repo *github.Repository, perform bool) error {
	owner, repoName := *repo.Owner.Login, *repo.Name
	context.Log("Processing %s", *repo.FullName)

	// 1. Find labels on GitHub.
	labels, _, err := context.GitHub.Issues.ListLabels(owner, repoName, &listOpts)
	if err != nil {
		return context.NewError("error fetching labels for %s: %v", *repo.FullName, err)
	}

	for _, desiredLabel := range desiredLabels {
		matchedLabel := findLabel(labels, desiredLabel)

		// It doesn't exist. Create and continue.
		if matchedLabel == nil {
			if perform {
				context.Log("%s: creating %s with color %s", *repo.FullName, *desiredLabel.Name, *desiredLabel.Color)
				_, _, err := context.GitHub.Issues.CreateLabel(owner, repoName, desiredLabel)
				if err != nil {
					return context.NewError("error creating '%s' for %s: %v", *desiredLabel.Name, *repo.FullName, err)
				}
			} else {
				context.Log("%s: would create %s with color %s", *repo.FullName, *desiredLabel.Name, *desiredLabel.Color)
			}
			continue
		}

		// It does exist, but possibly with incorrect info. Update it.
		if *matchedLabel.Name != *desiredLabel.Name || *matchedLabel.Color != *desiredLabel.Color {
			if perform {
				context.Log("%s: updating %s with data: %v",
					*repo.FullName, *matchedLabel.Name, github.Stringify(desiredLabel))
				_, _, err := context.GitHub.Issues.EditLabel(owner, repoName, *matchedLabel.Name, desiredLabel)
				if err != nil {
					return context.NewError("%s: error updating '%s': %v", *repo.FullName, *matchedLabel.Name, err)
				}
			} else {
				context.Log("%s: would update %s with data: %v", *repo.FullName, *matchedLabel.Name, github.Stringify(desiredLabel))
			}
			continue
		}
	}

	return nil
}

func main() {
	var perform bool
	flag.BoolVar(&perform, "f", false, "Whether to modify the labels (if true) or show dry-run output (if false).")
	flag.Parse()

	context := ctx.NewDefaultContext()
	repos, _, err := context.GitHub.Repositories.List("jekyll", &github.RepositoryListOptions{
		Type: "owner", Sort: "full_name", Direction: "asc", ListOptions: listOpts,
	})
	if err != nil {
		log.Fatalln("error fetching repos:", err)
	}

	for _, repo := range repos {
		if err := processRepo(context, repo, perform); err != nil {
			context.Log("%s: failed!", *repo.FullName)
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 The Upspin Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build !windows

/*
Package ose is a version of the file ops from the os package using encrypted files.
This version uses on disk files with a simple block encryption scheme to allow
random access of the file. Each 32 byte block of the file is encrypted by xoring
the contents with the AES encryption of the block number. Keys are per file and
kept in memory.

This encryption provides secrecy for files on lost machines but no integrity since
any contents can be changed with impunity.

The arguments to exported functions are the same as the equivalent os pkg functions.
*/

package ose

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/binary"

	"fmt"
	"os"
	"sync"
)

const (
	O_RDONLY = os.O_RDONLY // open the file read-only.
	O_WRONLY = os.O_WRONLY // open the file write-only.
	O_RDWR   = os.O_RDWR   // open the file read-write.
	O_APPEND = os.O_APPEND // append data to the file when writing.
	O_CREATE = os.O_CREATE // create a new file if none exists.
	O_EXCL   = os.O_EXCL   // used with O_CREATE, file must not exist
	O_SYNC   = os.O_SYNC   // open for synchronous I/O.
	O_TRUNC  = os.O_TRUNC  // if possible, truncate file when opened.
)

var state = struct {
	sync.Mutex
	mapping map[string]*File
}{mapping: make(map[string]*File)}

// File represents an encrypted file.
type File struct {
	name string
	f    *os.File
	benc cipher.Block
	refs int
}

// OpenFile opens an encrypted file.
func OpenFile(name string, flag int, mode os.FileMode) (*File, error) {
	f, err := os.OpenFile(name, flag, mode)
	if err != nil {
		return nil, err
	}
	state.Lock()
	defer state.Unlock()
	file, ok := state.mapping[name]
	if ok {
		file.f.Close()
	} else {
		benc, err := newBenc()
		if err != nil {
			return nil, err
		}
		file = &File{name: name, benc: benc}
		state.mapping[name] = file
	}
	file.f = f
	file.refs++
	return file, nil
}

// Create creates an encrypted file.
func Create(name string) (*File, error) {
	f, err := os.Create(name)
	if err != nil {
		return nil, err
	}
	state.Lock()
	defer state.Unlock()
	file, ok := state.mapping[name]
	if ok {
		file.f.Close()
	} else {
		benc, err := newBenc()
		if err != nil {
			return nil, err
		}
		file = &File{name: name, benc: benc}
		state.mapping[name] = file
	}
	file.f = f
	file.refs++
	return file, nil
}

// Rename renames file 'from' to 'to'.
func Rename(from, to string) error {
	state.Lock()
	defer state.Unlock()
	file, ok := state.mapping[from]
	if !ok {
		return fmt.Errorf("old file doesn't exist: %s", from)
	}
	if err := os.Rename(from, to); err != nil {
		return err
	}
	delete(state.mapping, from)
	state.mapping[to] = file
	return nil
}

// Mkdir creates the named directory.
func Mkdir(name string, mode os.FileMode) error {
	return os.Mkdir(name, mode)
}

// MkdirAll creates the named directory and all path elements leading to it.
func MkdirAll(name string, mode os.FileMode) error {
	return os.MkdirAll(name, mode)
}

// Remove removes the named file.
func Remove(name string) error {
	delete(state.mapping, name)
	return os.Remove(name)
}

// RemoveAll is a recursive remove.
func RemoveAll(subtree string) error {
	return os.RemoveAll(subtree)
}

// Truncate shortens a file.
func Truncate(name string, size int64) error {
	return os.Truncate(name, size)
}

// Close closes a file. If the ref count goes to zero, the file is removed.
func (file *File) Close() error {
	state.Lock()
	defer state.Unlock()
	file.refs--
	if file.refs != 0 {
		return nil
	}
	os.Remove(file.name)
	return file.f.Close()
}

// Stat returns the status of a file.
func (file *File) Stat() (os.FileInfo, error) {
	return file.f.Stat()
}

// ReadAt performs a random access read of an encrypted file and returns the
// decrypted content.
func (file *File) ReadAt(b []byte, off int64) (int, error) {
	n, err := file.f.ReadAt(b, off)
	if n > 0 {
		file.xor(b[:n], off)
	}
	return n, err
}

// WriteAt encrypts the content and writes it to the file.
// Unlile os.WriteAt, this changes the contents of b.
func (file *File) WriteAt(b []byte, off int64) (int, error) {
	file.xor(b, off)
	return file.f.WriteAt(b, off)
}

const aesKeyLen = 32

func newBenc() (cipher.Block, error) {
	k := make([]byte, aesKeyLen)
	if _, err := rand.Read(k); err != nil {
		return nil, err
	}

	// Create an Xcrypter.
	benc, err := aes.NewCipher(k)
	if err != nil {
		return nil, err
	}
	return benc, nil
}

func (file *File) xor(b []byte, off int64) {
	bsize := int64(file.benc.BlockSize())
	mask := make([]byte, bsize)
	maskInput := make([]byte, bsize)
	if off%bsize != 0 {
		binary.PutVarint(maskInput, off/bsize)
		file.benc.Encrypt(mask, maskInput)
	}
	for i := range b {
		x := (off + int64(i)) % bsize
		if x == 0 {
			binary.PutVarint(maskInput, (off+int64(i))/bsize)
			file.benc.Encrypt(mask, maskInput)
		}
		b[i] ^= mask[x]
	}
}
<commit_msg>cmd/upspinfs: discard cached files we're done with.<commit_after>// Copyright 2016 The Upspin Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build !windows

/*
Package ose is a version of the file ops from the os package using encrypted files.
This version uses on disk files with a simple block encryption scheme to allow
random access of the file. Each 32 byte block of the file is encrypted by xoring
the contents with the AES encryption of the block number. Keys are per file and
kept in memory.

This encryption provides secrecy for files on lost machines but no integrity since
any contents can be changed with impunity.

The arguments to exported functions are the same as the equivalent os pkg functions.
*/

package ose

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/binary"

	"fmt"
	"os"
	"sync"
)

const (
	O_RDONLY = os.O_RDONLY // open the file read-only.
	O_WRONLY = os.O_WRONLY // open the file write-only.
	O_RDWR   = os.O_RDWR   // open the file read-write.
	O_APPEND = os.O_APPEND // append data to the file when writing.
	O_CREATE = os.O_CREATE // create a new file if none exists.
	O_EXCL   = os.O_EXCL   // used with O_CREATE, file must not exist
	O_SYNC   = os.O_SYNC   // open for synchronous I/O.
	O_TRUNC  = os.O_TRUNC  // if possible, truncate file when opened.
)

var state = struct {
	sync.Mutex
	mapping map[string]*File
}{mapping: make(map[string]*File)}

// File represents an encrypted file.
type File struct {
	name string
	f    *os.File
	benc cipher.Block
	refs int
}

// OpenFile opens an encrypted file.
func OpenFile(name string, flag int, mode os.FileMode) (*File, error) {
	f, err := os.OpenFile(name, flag, mode)
	if err != nil {
		return nil, err
	}
	state.Lock()
	defer state.Unlock()
	file, ok := state.mapping[name]
	if ok {
		file.f.Close()
	} else {
		benc, err := newBenc()
		if err != nil {
			return nil, err
		}
		file = &File{name: name, benc: benc}
		state.mapping[name] = file
	}
	file.f = f
	file.refs++
	return file, nil
}

// Create creates an encrypted file.
func Create(name string) (*File, error) {
	f, err := os.Create(name)
	if err != nil {
		return nil, err
	}
	state.Lock()
	defer state.Unlock()
	file, ok := state.mapping[name]
	if ok {
		file.f.Close()
	} else {
		benc, err := newBenc()
		if err != nil {
			return nil, err
		}
		file = &File{name: name, benc: benc}
		state.mapping[name] = file
	}
	file.f = f
	file.refs++
	return file, nil
}

// Rename renames file 'from' to 'to'.
func Rename(from, to string) error {
	state.Lock()
	defer state.Unlock()
	file, ok := state.mapping[from]
	if !ok {
		return fmt.Errorf("old file doesn't exist: %s", from)
	}
	if err := os.Rename(from, to); err != nil {
		return err
	}
	delete(state.mapping, from)
	state.mapping[to] = file
	file.name = to
	return nil
}

// Mkdir creates the named directory.
func Mkdir(name string, mode os.FileMode) error {
	return os.Mkdir(name, mode)
}

// MkdirAll creates the named directory and all path elements leading to it.
func MkdirAll(name string, mode os.FileMode) error {
	return os.MkdirAll(name, mode)
}

// Remove removes the named file.
func Remove(name string) error {
	delete(state.mapping, name)
	return os.Remove(name)
}

// RemoveAll is a recursive remove.
func RemoveAll(subtree string) error {
	return os.RemoveAll(subtree)
}

// Truncate shortens a file.
func Truncate(name string, size int64) error {
	return os.Truncate(name, size)
}

// Close closes a file. If the ref count goes to zero, the file is removed.
func (file *File) Close() error {
	state.Lock()
	defer state.Unlock()
	file.refs--
	if file.refs != 0 {
		return nil
	}
	os.Remove(file.name)
	return file.f.Close()
}

// Stat returns the status of a file.
func (file *File) Stat() (os.FileInfo, error) {
	return file.f.Stat()
}

// ReadAt performs a random access read of an encrypted file and returns the
// decrypted content.
func (file *File) ReadAt(b []byte, off int64) (int, error) {
	n, err := file.f.ReadAt(b, off)
	if n > 0 {
		file.xor(b[:n], off)
	}
	return n, err
}

// WriteAt encrypts the content and writes it to the file.
// Unlile os.WriteAt, this changes the contents of b.
func (file *File) WriteAt(b []byte, off int64) (int, error) {
	file.xor(b, off)
	return file.f.WriteAt(b, off)
}

const aesKeyLen = 32

func newBenc() (cipher.Block, error) {
	k := make([]byte, aesKeyLen)
	if _, err := rand.Read(k); err != nil {
		return nil, err
	}

	// Create an Xcrypter.
	benc, err := aes.NewCipher(k)
	if err != nil {
		return nil, err
	}
	return benc, nil
}

func (file *File) xor(b []byte, off int64) {
	bsize := int64(file.benc.BlockSize())
	mask := make([]byte, bsize)
	maskInput := make([]byte, bsize)
	if off%bsize != 0 {
		binary.PutVarint(maskInput, off/bsize)
		file.benc.Encrypt(mask, maskInput)
	}
	for i := range b {
		x := (off + int64(i)) % bsize
		if x == 0 {
			binary.PutVarint(maskInput, (off+int64(i))/bsize)
			file.benc.Encrypt(mask, maskInput)
		}
		b[i] ^= mask[x]
	}
}
<|endoftext|>


<commit_before>package xlsx

import (
	"testing"
)

// Test that we can successfully read an XLSX file generated by
// Numbers for Mac.
func TestMacNumbers(t *testing.T) {
	xlsxFile, error := OpenFile("macNumbersTest.xlsx")
	if error != nil {
		t.Error(error.Error())
		return
	}
	if xlsxFile == nil {
		t.Error("OpenFile returned nil FileInterface without generating an os.Error")
		return
	}
	s := xlsxFile.Sheets[0].Cell(0, 0).String()
	if s != "编号" {
		t.Errorf("[TestMacExcel] xlsxFile.Sheets[0].Cell(0,0).String():'%s'", s)
		return
	}
}
<commit_msg>Convert macNumbers_test.go to use gocheck.<commit_after>package xlsx

import (
	. "gopkg.in/check.v1"
)

type MacNumbersSuite struct {}
var _ = Suite(&MacNumbersSuite{})

// Test that we can successfully read an XLSX file generated by
// Numbers for Mac.
func (m *MacNumbersSuite) TestMacNumbers(c *C) {
	xlsxFile, err := OpenFile("macNumbersTest.xlsx")
	c.Assert(err, IsNil)
	c.Assert(xlsxFile, NotNil)
	s := xlsxFile.Sheets[0].Cell(0, 0).String()
	c.Assert(s, Equals, "编号")
}
<|endoftext|>


<commit_before>package docker

import (
	"fmt"
	"strings"
	"time"
)

type NetworkSettings struct {
	Ip          string `json:"IPAddress"`
	IPPrefixLen int
	Gateway     string
	Bridge      string
	PortMapping map[string]map[string]string
	Ports       map[Port][]PortBinding
}

func (settings *NetworkSettings) PortMappingString() string {
	out := []string{}
	tcp := []string{}
	udp := []string{}
	if mappings := settings.TcpPortMapping(); len(mappings) > 0 {
		for v, k := range mappings {
			tcp = append(tcp, k+"=>"+v)
		}
		out = append(out, strings.Join(tcp, ", "))
	}
	if mappings := settings.UdpPortMapping(); len(mappings) > 0 {
		for v, k := range mappings {
			udp = append(udp, k+"=>"+v)
		}
		out = append(out, "UDP: "+strings.Join(udp, ", "))
	}
	return strings.Join(out, " | ")
}

func (settings *NetworkSettings) TcpPortMapping() map[string]string {
	return settings.PortMappingFor("Tcp")
}

func (settings *NetworkSettings) UdpPortMapping() map[string]string {
	return settings.PortMappingFor("Udp")
}

func (settings *NetworkSettings) PortMappingFor(protocol string) map[string]string {
	if mapping := settings.PortMapping[protocol]; mapping != nil {
		return mapping
	}
	return map[string]string{}
}

type PortConfig struct {
	Private  int    `json:"PrivatePort"`
	Public   int    `json:"PublicPort"`
	Protocol string `json:"Type"`
	Ip       string `json:"IP"`
}

type Container struct {
	Id         string
	Image      string
	Command    string
	Created    int64 `json:"Created"`
	Status     string
	Ports      []*PortConfig
	SizeRw     int
	SizeRootFs int
	Names      []string
}

func (self *Container) CreatedAt() time.Time {
	return time.Unix(self.Created, 0)
}

func (container *Container) String() string {
	return fmt.Sprintf("%s: %s", container.Id, container.Status)
}

type ContainerInfo struct {
	Id              string `json:"ID"`
	Image           string
	CreatedAt       time.Time `json:"Created"`
	SysInitPath     string
	ResolvConfPath  string
	Volumes         map[string]string
	VolumesRW       map[string]bool
	Path            string
	Args            []string
	ContainerConfig ContainerConfig `json:"Config"`
	NetworkConfig   NetworkSettings `json:"NetworkSettings"`
	HostConfig      HostConfig      `json:"HostConfig"`
	State           struct {
		Status   string
		ExitCode int
	}
}

// https://github.com/dotcloud/docker/blob/master/container.go#L60-81
type ContainerConfig struct {
	Hostname        string
	Domainname      string
	User            string
	Memory          int64 // Memory limit (in bytes)
	MemorySwap      int64 // Total memory usage (memory + swap); set `-1' to disable swap
	CpuShares       int64 // CPU shares (relative weight vs. other containers)
	AttachStdin     bool
	AttachStdout    bool
	AttachStderr    bool
	PortSpecs       []string
	ExposedPorts    map[Port]struct{}
	Tty             bool // Attach standard streams to a tty, including stdin if it is not closed.
	OpenStdin       bool // Open stdin
	StdinOnce       bool // If true, close stdin after the 1 attached client disconnects.
	Env             []string
	Cmd             []string
	Dns             []string
	Image           string // Name of the image as it was passed by the operator (eg. could be symbolic)
	Volumes         map[string]struct{}
	VolumesFrom     interface{} `json:"VolumesFrom,omitempty"`
	WorkingDir      string
	Entrypoint      []string
	NetworkDisabled bool
}
<commit_msg>dockerlicnet: add label mapping to containers<commit_after>package docker

import (
	"fmt"
	"strings"
	"time"
)

type NetworkSettings struct {
	Ip          string `json:"IPAddress"`
	IPPrefixLen int
	Gateway     string
	Bridge      string
	PortMapping map[string]map[string]string
	Ports       map[Port][]PortBinding
}

func (settings *NetworkSettings) PortMappingString() string {
	out := []string{}
	tcp := []string{}
	udp := []string{}
	if mappings := settings.TcpPortMapping(); len(mappings) > 0 {
		for v, k := range mappings {
			tcp = append(tcp, k+"=>"+v)
		}
		out = append(out, strings.Join(tcp, ", "))
	}
	if mappings := settings.UdpPortMapping(); len(mappings) > 0 {
		for v, k := range mappings {
			udp = append(udp, k+"=>"+v)
		}
		out = append(out, "UDP: "+strings.Join(udp, ", "))
	}
	return strings.Join(out, " | ")
}

func (settings *NetworkSettings) TcpPortMapping() map[string]string {
	return settings.PortMappingFor("Tcp")
}

func (settings *NetworkSettings) UdpPortMapping() map[string]string {
	return settings.PortMappingFor("Udp")
}

func (settings *NetworkSettings) PortMappingFor(protocol string) map[string]string {
	if mapping := settings.PortMapping[protocol]; mapping != nil {
		return mapping
	}
	return map[string]string{}
}

type PortConfig struct {
	Private  int    `json:"PrivatePort"`
	Public   int    `json:"PublicPort"`
	Protocol string `json:"Type"`
	Ip       string `json:"IP"`
}

type Container struct {
	Id         string
	Image      string
	Command    string
	Created    int64 `json:"Created"`
	Status     string
	Ports      []*PortConfig
	SizeRw     int
	SizeRootFs int
	Names      []string
	Labels     map[string]string
}

func (self *Container) CreatedAt() time.Time {
	return time.Unix(self.Created, 0)
}

func (container *Container) String() string {
	return fmt.Sprintf("%s: %s", container.Id, container.Status)
}

type ContainerInfo struct {
	Id              string `json:"ID"`
	Image           string
	CreatedAt       time.Time `json:"Created"`
	SysInitPath     string
	ResolvConfPath  string
	Volumes         map[string]string
	VolumesRW       map[string]bool
	Path            string
	Args            []string
	ContainerConfig ContainerConfig `json:"Config"`
	NetworkConfig   NetworkSettings `json:"NetworkSettings"`
	HostConfig      HostConfig      `json:"HostConfig"`
	State           struct {
		Status   string
		ExitCode int
	}
}

// https://github.com/dotcloud/docker/blob/master/container.go#L60-81
type ContainerConfig struct {
	Hostname        string
	Domainname      string
	User            string
	Memory          int64 // Memory limit (in bytes)
	MemorySwap      int64 // Total memory usage (memory + swap); set `-1' to disable swap
	CpuShares       int64 // CPU shares (relative weight vs. other containers)
	AttachStdin     bool
	AttachStdout    bool
	AttachStderr    bool
	PortSpecs       []string
	ExposedPorts    map[Port]struct{}
	Tty             bool // Attach standard streams to a tty, including stdin if it is not closed.
	OpenStdin       bool // Open stdin
	StdinOnce       bool // If true, close stdin after the 1 attached client disconnects.
	Env             []string
	Cmd             []string
	Dns             []string
	Image           string // Name of the image as it was passed by the operator (eg. could be symbolic)
	Volumes         map[string]struct{}
	VolumesFrom     interface{} `json:"VolumesFrom,omitempty"`
	WorkingDir      string
	Entrypoint      []string
	NetworkDisabled bool
}
<|endoftext|>


<commit_before>// Copyright 2014 Google Inc. All Rights Reserved.
//
// 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 manager

import (
	"fmt"
	"path"
	"strings"
	"sync"
	"time"

	"github.com/golang/glog"
	"github.com/google/cadvisor/container"
	"github.com/google/cadvisor/info"
	"github.com/google/cadvisor/storage"
)

type Manager interface {
	// Start the manager, blocks forever.
	Start() error

	// Get information about a container.
	GetContainerInfo(containerName string, query *info.ContainerInfoRequest) (*info.ContainerInfo, error)

	// Get information about all subcontainers of the specified container (includes self).
	SubcontainersInfo(containerName string, query *info.ContainerInfoRequest) ([]*info.ContainerInfo, error)

	// Get information about the machine.
	GetMachineInfo() (*info.MachineInfo, error)

	// Get version information about different components we depend on.
	GetVersionInfo() (*info.VersionInfo, error)
}

func New(driver storage.StorageDriver) (Manager, error) {
	if driver == nil {
		return nil, fmt.Errorf("nil storage driver!")
	}
	newManager := &manager{}
	newManager.containers = make(map[string]*containerData)

	machineInfo, err := getMachineInfo()
	if err != nil {
		return nil, err
	}
	newManager.machineInfo = *machineInfo
	glog.Infof("Machine: %+v", newManager.machineInfo)

	versionInfo, err := getVersionInfo()
	if err != nil {
		return nil, err
	}
	newManager.versionInfo = *versionInfo
	glog.Infof("Version: %+v", newManager.versionInfo)
	newManager.storageDriver = driver

	return newManager, nil
}

type manager struct {
	containers     map[string]*containerData
	containersLock sync.RWMutex
	storageDriver  storage.StorageDriver
	machineInfo    info.MachineInfo
	versionInfo    info.VersionInfo
}

// Start the container manager.
func (m *manager) Start() error {
	// Create root and then recover all containers.
	_, err := m.createContainer("/")
	if err != nil {
		return err
	}
	glog.Infof("Starting recovery of all containers")
	err = m.detectContainers()
	if err != nil {
		return err
	}
	glog.Infof("Recovery completed")

	// Look for new containers in the main housekeeping thread.
	for t := range time.Tick(time.Second) {
		start := time.Now()

		// Check for new containers.
		err = m.detectContainers()
		if err != nil {
			glog.Errorf("Failed to detect containers: %s", err)
		}

		// Log if housekeeping took more than 100ms.
		duration := time.Since(start)
		if duration >= 100*time.Millisecond {
			glog.V(1).Infof("Global Housekeeping(%d) took %s", t.Unix(), duration)
		}
	}
	return nil
}

// Get a container by name.
func (self *manager) GetContainerInfo(containerName string, query *info.ContainerInfoRequest) (*info.ContainerInfo, error) {
	var cont *containerData
	var ok bool
	func() {
		self.containersLock.RLock()
		defer self.containersLock.RUnlock()

		// Ensure we have the container.
		cont, ok = self.containers[containerName]
	}()
	if !ok {
		return nil, fmt.Errorf("unknown container %q", containerName)
	}

	return self.containerDataToContainerInfo(cont, query)
}

func (self *manager) containerDataToContainerInfo(cont *containerData, query *info.ContainerInfoRequest) (*info.ContainerInfo, error) {
	// Get the info from the container.
	cinfo, err := cont.GetInfo()
	if err != nil {
		return nil, err
	}

	var percentiles *info.ContainerStatsPercentiles
	var samples []*info.ContainerStatsSample
	var stats []*info.ContainerStats
	percentiles, err = self.storageDriver.Percentiles(
		cinfo.Name,
		query.CpuUsagePercentiles,
		query.MemoryUsagePercentiles,
	)
	if err != nil {
		return nil, err
	}
	samples, err = self.storageDriver.Samples(cinfo.Name, query.NumSamples)
	if err != nil {
		return nil, err
	}

	stats, err = self.storageDriver.RecentStats(cinfo.Name, query.NumStats)
	if err != nil {
		return nil, err
	}

	// Make a copy of the info for the user.
	ret := &info.ContainerInfo{
		ContainerReference: info.ContainerReference{
			Name:    cinfo.Name,
			Aliases: cinfo.Aliases,
		},
		Subcontainers:    cinfo.Subcontainers,
		Spec:             cinfo.Spec,
		StatsPercentiles: percentiles,
		Samples:          samples,
		Stats:            stats,
	}

	// Set default value to an actual value
	if ret.Spec.Memory != nil {
		// Memory.Limit is 0 means there's no limit
		if ret.Spec.Memory.Limit == 0 {
			ret.Spec.Memory.Limit = uint64(self.machineInfo.MemoryCapacity)
		}
	}
	return ret, nil
}

func (self *manager) SubcontainersInfo(containerName string, query *info.ContainerInfoRequest) ([]*info.ContainerInfo, error) {
	var containers []*containerData
	func() {
		self.containersLock.RLock()
		defer self.containersLock.RUnlock()
		containers = make([]*containerData, 0, len(self.containers))

		// Get all the subcontainers of the specified container
		matchedName := path.Join(containerName, "/")
		for i := range self.containers {
			name := self.containers[i].info.Name
			if name == containerName || strings.HasPrefix(name, matchedName) {
				containers = append(containers, self.containers[i])
			}
		}
	}()
	if len(containers) == 0 {
		return nil, fmt.Errorf("unknown container %q", containerName)
	}

	// Get the info for each container.
	output := make([]*info.ContainerInfo, 0, len(containers))
	for i := range containers {
		cinfo, err := self.containerDataToContainerInfo(containers[i], query)
		if err != nil {
			// Skip containers with errors, we try to degrade gracefully.
			continue
		}
		output = append(output, cinfo)
	}

	return output, nil
}

func (m *manager) GetMachineInfo() (*info.MachineInfo, error) {
	// Copy and return the MachineInfo.
	ret := m.machineInfo
	return &ret, nil
}

func (m *manager) GetVersionInfo() (*info.VersionInfo, error) {
	ret := m.versionInfo
	return &ret, nil
}

// Create a container. This expects to only be called from the global manager thread.
func (m *manager) createContainer(containerName string) (*containerData, error) {
	cont, err := NewContainerData(containerName, m.storageDriver)
	if err != nil {
		return nil, err
	}

	// Add to the containers map.
	func() {
		m.containersLock.Lock()
		defer m.containersLock.Unlock()

		// Add the container name and all its aliases.
		m.containers[containerName] = cont
		for _, alias := range cont.info.Aliases {
			m.containers[alias] = cont
		}
	}()
	glog.Infof("Added container: %q (aliases: %s)", containerName, cont.info.Aliases)

	// Start the container's housekeeping.
	cont.Start()
	return cont, nil
}

func (m *manager) destroyContainer(containerName string) error {
	m.containersLock.Lock()
	defer m.containersLock.Unlock()

	cont, ok := m.containers[containerName]
	if !ok {
		return fmt.Errorf("Expected container \"%s\" to exist during destroy", containerName)
	}

	// Tell the container to stop.
	err := cont.Stop()
	if err != nil {
		return err
	}

	// Remove the container from our records (and all its aliases).
	delete(m.containers, containerName)
	for _, alias := range cont.info.Aliases {
		delete(m.containers, alias)
	}
	glog.Infof("Destroyed container: %s (aliases: %s)", containerName, cont.info.Aliases)
	return nil
}

// Detect all containers that have been added or deleted.
func (m *manager) getContainersDiff() (added []info.ContainerReference, removed []info.ContainerReference, err error) {
	// TODO(vmarmol): We probably don't need to lock around / since it will always be there.
	m.containersLock.RLock()
	defer m.containersLock.RUnlock()

	// Get all containers on the system.
	cont, ok := m.containers["/"]
	if !ok {
		return nil, nil, fmt.Errorf("Failed to find container \"/\" while checking for new containers")
	}
	allContainers, err := cont.handler.ListContainers(container.LIST_RECURSIVE)
	if err != nil {
		return nil, nil, err
	}
	allContainers = append(allContainers, info.ContainerReference{Name: "/"})

	// Determine which were added and which were removed.
	allContainersSet := make(map[string]*containerData)
	for name, d := range m.containers {
		// Only add the canonical name.
		if d.info.Name == name {
			allContainersSet[name] = d
		}
	}
	for _, c := range allContainers {
		delete(allContainersSet, c.Name)
		_, ok := m.containers[c.Name]
		if !ok {
			added = append(added, c)
		}
	}

	// Removed ones are no longer in the container listing.
	for _, d := range allContainersSet {
		removed = append(removed, d.info.ContainerReference)
	}

	return
}

// Detect the existing containers and reflect the setup here.
func (m *manager) detectContainers() error {
	added, removed, err := m.getContainersDiff()
	if err != nil {
		return err
	}

	// Add the new containers.
	for _, cont := range added {
		_, err = m.createContainer(cont.Name)
		if err != nil {
			glog.Errorf("failed to create existing container: %s: %s", cont.Name, err)
		}
	}

	// Remove the old containers.
	for _, cont := range removed {
		err = m.destroyContainer(cont.Name)
		if err != nil {
			glog.Errorf("failed to destroy existing container: %s: %s", cont.Name, err)
		}
	}

	return nil
}
<commit_msg>Fix capitalization in log files.<commit_after>// Copyright 2014 Google Inc. All Rights Reserved.
//
// 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 manager

import (
	"fmt"
	"path"
	"strings"
	"sync"
	"time"

	"github.com/golang/glog"
	"github.com/google/cadvisor/container"
	"github.com/google/cadvisor/info"
	"github.com/google/cadvisor/storage"
)

type Manager interface {
	// Start the manager, blocks forever.
	Start() error

	// Get information about a container.
	GetContainerInfo(containerName string, query *info.ContainerInfoRequest) (*info.ContainerInfo, error)

	// Get information about all subcontainers of the specified container (includes self).
	SubcontainersInfo(containerName string, query *info.ContainerInfoRequest) ([]*info.ContainerInfo, error)

	// Get information about the machine.
	GetMachineInfo() (*info.MachineInfo, error)

	// Get version information about different components we depend on.
	GetVersionInfo() (*info.VersionInfo, error)
}

func New(driver storage.StorageDriver) (Manager, error) {
	if driver == nil {
		return nil, fmt.Errorf("nil storage driver!")
	}
	newManager := &manager{}
	newManager.containers = make(map[string]*containerData)

	machineInfo, err := getMachineInfo()
	if err != nil {
		return nil, err
	}
	newManager.machineInfo = *machineInfo
	glog.Infof("Machine: %+v", newManager.machineInfo)

	versionInfo, err := getVersionInfo()
	if err != nil {
		return nil, err
	}
	newManager.versionInfo = *versionInfo
	glog.Infof("Version: %+v", newManager.versionInfo)
	newManager.storageDriver = driver

	return newManager, nil
}

type manager struct {
	containers     map[string]*containerData
	containersLock sync.RWMutex
	storageDriver  storage.StorageDriver
	machineInfo    info.MachineInfo
	versionInfo    info.VersionInfo
}

// Start the container manager.
func (m *manager) Start() error {
	// Create root and then recover all containers.
	_, err := m.createContainer("/")
	if err != nil {
		return err
	}
	glog.Infof("Starting recovery of all containers")
	err = m.detectContainers()
	if err != nil {
		return err
	}
	glog.Infof("Recovery completed")

	// Look for new containers in the main housekeeping thread.
	for t := range time.Tick(time.Second) {
		start := time.Now()

		// Check for new containers.
		err = m.detectContainers()
		if err != nil {
			glog.Errorf("Failed to detect containers: %s", err)
		}

		// Log if housekeeping took more than 100ms.
		duration := time.Since(start)
		if duration >= 100*time.Millisecond {
			glog.V(1).Infof("Global Housekeeping(%d) took %s", t.Unix(), duration)
		}
	}
	return nil
}

// Get a container by name.
func (self *manager) GetContainerInfo(containerName string, query *info.ContainerInfoRequest) (*info.ContainerInfo, error) {
	var cont *containerData
	var ok bool
	func() {
		self.containersLock.RLock()
		defer self.containersLock.RUnlock()

		// Ensure we have the container.
		cont, ok = self.containers[containerName]
	}()
	if !ok {
		return nil, fmt.Errorf("unknown container %q", containerName)
	}

	return self.containerDataToContainerInfo(cont, query)
}

func (self *manager) containerDataToContainerInfo(cont *containerData, query *info.ContainerInfoRequest) (*info.ContainerInfo, error) {
	// Get the info from the container.
	cinfo, err := cont.GetInfo()
	if err != nil {
		return nil, err
	}

	var percentiles *info.ContainerStatsPercentiles
	var samples []*info.ContainerStatsSample
	var stats []*info.ContainerStats
	percentiles, err = self.storageDriver.Percentiles(
		cinfo.Name,
		query.CpuUsagePercentiles,
		query.MemoryUsagePercentiles,
	)
	if err != nil {
		return nil, err
	}
	samples, err = self.storageDriver.Samples(cinfo.Name, query.NumSamples)
	if err != nil {
		return nil, err
	}

	stats, err = self.storageDriver.RecentStats(cinfo.Name, query.NumStats)
	if err != nil {
		return nil, err
	}

	// Make a copy of the info for the user.
	ret := &info.ContainerInfo{
		ContainerReference: info.ContainerReference{
			Name:    cinfo.Name,
			Aliases: cinfo.Aliases,
		},
		Subcontainers:    cinfo.Subcontainers,
		Spec:             cinfo.Spec,
		StatsPercentiles: percentiles,
		Samples:          samples,
		Stats:            stats,
	}

	// Set default value to an actual value
	if ret.Spec.Memory != nil {
		// Memory.Limit is 0 means there's no limit
		if ret.Spec.Memory.Limit == 0 {
			ret.Spec.Memory.Limit = uint64(self.machineInfo.MemoryCapacity)
		}
	}
	return ret, nil
}

func (self *manager) SubcontainersInfo(containerName string, query *info.ContainerInfoRequest) ([]*info.ContainerInfo, error) {
	var containers []*containerData
	func() {
		self.containersLock.RLock()
		defer self.containersLock.RUnlock()
		containers = make([]*containerData, 0, len(self.containers))

		// Get all the subcontainers of the specified container
		matchedName := path.Join(containerName, "/")
		for i := range self.containers {
			name := self.containers[i].info.Name
			if name == containerName || strings.HasPrefix(name, matchedName) {
				containers = append(containers, self.containers[i])
			}
		}
	}()
	if len(containers) == 0 {
		return nil, fmt.Errorf("unknown container %q", containerName)
	}

	// Get the info for each container.
	output := make([]*info.ContainerInfo, 0, len(containers))
	for i := range containers {
		cinfo, err := self.containerDataToContainerInfo(containers[i], query)
		if err != nil {
			// Skip containers with errors, we try to degrade gracefully.
			continue
		}
		output = append(output, cinfo)
	}

	return output, nil
}

func (m *manager) GetMachineInfo() (*info.MachineInfo, error) {
	// Copy and return the MachineInfo.
	ret := m.machineInfo
	return &ret, nil
}

func (m *manager) GetVersionInfo() (*info.VersionInfo, error) {
	ret := m.versionInfo
	return &ret, nil
}

// Create a container. This expects to only be called from the global manager thread.
func (m *manager) createContainer(containerName string) (*containerData, error) {
	cont, err := NewContainerData(containerName, m.storageDriver)
	if err != nil {
		return nil, err
	}

	// Add to the containers map.
	func() {
		m.containersLock.Lock()
		defer m.containersLock.Unlock()

		// Add the container name and all its aliases.
		m.containers[containerName] = cont
		for _, alias := range cont.info.Aliases {
			m.containers[alias] = cont
		}
	}()
	glog.Infof("Added container: %q (aliases: %s)", containerName, cont.info.Aliases)

	// Start the container's housekeeping.
	cont.Start()
	return cont, nil
}

func (m *manager) destroyContainer(containerName string) error {
	m.containersLock.Lock()
	defer m.containersLock.Unlock()

	cont, ok := m.containers[containerName]
	if !ok {
		return fmt.Errorf("Expected container \"%s\" to exist during destroy", containerName)
	}

	// Tell the container to stop.
	err := cont.Stop()
	if err != nil {
		return err
	}

	// Remove the container from our records (and all its aliases).
	delete(m.containers, containerName)
	for _, alias := range cont.info.Aliases {
		delete(m.containers, alias)
	}
	glog.Infof("Destroyed container: %s (aliases: %s)", containerName, cont.info.Aliases)
	return nil
}

// Detect all containers that have been added or deleted.
func (m *manager) getContainersDiff() (added []info.ContainerReference, removed []info.ContainerReference, err error) {
	// TODO(vmarmol): We probably don't need to lock around / since it will always be there.
	m.containersLock.RLock()
	defer m.containersLock.RUnlock()

	// Get all containers on the system.
	cont, ok := m.containers["/"]
	if !ok {
		return nil, nil, fmt.Errorf("Failed to find container \"/\" while checking for new containers")
	}
	allContainers, err := cont.handler.ListContainers(container.LIST_RECURSIVE)
	if err != nil {
		return nil, nil, err
	}
	allContainers = append(allContainers, info.ContainerReference{Name: "/"})

	// Determine which were added and which were removed.
	allContainersSet := make(map[string]*containerData)
	for name, d := range m.containers {
		// Only add the canonical name.
		if d.info.Name == name {
			allContainersSet[name] = d
		}
	}
	for _, c := range allContainers {
		delete(allContainersSet, c.Name)
		_, ok := m.containers[c.Name]
		if !ok {
			added = append(added, c)
		}
	}

	// Removed ones are no longer in the container listing.
	for _, d := range allContainersSet {
		removed = append(removed, d.info.ContainerReference)
	}

	return
}

// Detect the existing containers and reflect the setup here.
func (m *manager) detectContainers() error {
	added, removed, err := m.getContainersDiff()
	if err != nil {
		return err
	}

	// Add the new containers.
	for _, cont := range added {
		_, err = m.createContainer(cont.Name)
		if err != nil {
			glog.Errorf("Failed to create existing container: %s: %s", cont.Name, err)
		}
	}

	// Remove the old containers.
	for _, cont := range removed {
		err = m.destroyContainer(cont.Name)
		if err != nil {
			glog.Errorf("Failed to destroy existing container: %s: %s", cont.Name, err)
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package manager

import (
	"bytes"
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"net/http"
	"net/url"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/fsouza/go-dockerclient"

	"github.com/dockpit/dirtar"
)

type Manager struct {
	client *docker.Client
	host   string
}

// the Docker image that is used for mocks
var ImageName = "dockpit/mock:latest"
var MockPrivatePort int64 = 8000

// Manages state for microservice testing by creating
// docker images and starting containers when necessary
func NewManager(host, cert string) (*Manager, error) {
	m := &Manager{
		host: host,
	}

	//parse docker host addr as url
	hurl, err := url.Parse(host)
	if err != nil {
		return nil, err
	}

	//change to http connection
	hurl.Scheme = "https"

	//create docker client
	m.client, err = docker.NewTLSClient(hurl.String(), filepath.Join(cert, "cert.pem"), filepath.Join(cert, "key.pem"), filepath.Join(cert, "ca.pem"))
	if err != nil {
		return nil, err
	}

	return m, nil
}

// hash a path using md5
func containerName(path string) (string, error) {

	//create md5 of full path
	hash := md5.New()
	_, err := hash.Write([]byte(path))
	if err != nil {
		return "", err
	}

	return fmt.Sprintf("pitmock_%s", hex.EncodeToString(hash.Sum(nil))), nil
}

// start a mock container by using examples from the given directory
func (m *Manager) Start(dir string, portb map[docker.Port][]docker.PortBinding) (*MockContainer, error) {

	//create name for container
	cname, err := containerName(dir)
	if err != nil {
		return nil, err
	}

	//@todo, grab the first and set private port to expected mock service
	//to the mock porcess binding
	for _, pb := range portb {
		portb[docker.Port(strconv.FormatInt(MockPrivatePort, 10)+"/tcp")] = pb
		break
	}

	//create the container
	c, err := m.client.CreateContainer(docker.CreateContainerOptions{
		Name: cname,
		Config: &docker.Config{
			Image: ImageName,
		},
	})

	if err != nil {
		return nil, err
	}

	//start the container we created
	err = m.client.StartContainer(c.ID, &docker.HostConfig{PortBindings: portb})
	if err != nil {
		return nil, err
	}

	//get container port mapping
	ci, err := m.client.InspectContainer(c.ID)
	if err != nil {
		return nil, err
	}

	//use docker host location to form url
	hurl, err := url.Parse(m.host)
	if err != nil {
		return nil, err
	}

	//tar examples into memory
	tar := bytes.NewBuffer(nil)
	err = dirtar.Tar(dir, tar)
	if err != nil {
		return nil, err
	}

	//wait for container to settle?
	//@todo very dirty business here
	<-time.After(time.Millisecond * 200)

	//get the external port for 8000 and turn into an url we can send http requests to
	//@todo, here we assume that the first configured port is the http interface, indicate explicetly
	//@todo, use the logic on line 77
	hurl.Scheme = "http"
	for _, pconfig := range ci.NetworkSettings.PortMappingAPI() {
		if pconfig.PrivatePort != 0 {
			hurl.Host = strings.Replace(hurl.Host, ":2376", fmt.Sprintf(":%d", pconfig.PublicPort), 1)
		}
	}

	//send the upload request
	resp, err := http.Post(fmt.Sprintf("%s/_examples", hurl.String()), "application/x-tar", tar)
	if err != nil {
		return nil, err
	}

	//check if we uploaded correctly
	if resp.StatusCode != 201 {
		return nil, fmt.Errorf("Failed to upload examples to container '%s' %s", ci.ID, ci.Name)
	}

	//send HUP signal to 'root' process (dockpit/mock) to reload examples
	err = m.client.KillContainer(docker.KillContainerOptions{
		ID:     c.ID,
		Signal: docker.Signal(syscall.SIGHUP),
	})

	if err != nil {
		return nil, err
	}

	return &MockContainer{c.ID, hurl.String(), dir}, nil
}

// stop a mock container that was started from the given directory
func (m *Manager) Stop(dir string) error {

	//create name for container
	cname, err := containerName(dir)
	if err != nil {
		return err
	}

	//get all containers
	cs, err := m.client.ListContainers(docker.ListContainersOptions{})
	if err != nil {
		return err
	}

	//get container that matches the name
	// var container *docker.APIContainers
	var container docker.APIContainers
	for _, c := range cs {
		for _, n := range c.Names {
			if n[1:] == cname {
				container = c
			}
		}
	}

	//remove hard since mocks are ephemeral
	return m.client.RemoveContainer(docker.RemoveContainerOptions{
		ID:            container.ID,
		RemoveVolumes: true,
		Force:         true,
	})
}
<commit_msg>another consideration<commit_after>package manager

import (
	"bytes"
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"net/http"
	"net/url"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/fsouza/go-dockerclient"

	"github.com/dockpit/dirtar"
)

type Manager struct {
	client *docker.Client
	host   string
}

// the Docker image that is used for mocks
var ImageName = "dockpit/mock:latest"
var MockPrivatePort int64 = 8000

// Manages state for microservice testing by creating
// docker images and starting containers when necessary
func NewManager(host, cert string) (*Manager, error) {
	m := &Manager{
		host: host,
	}

	//parse docker host addr as url
	hurl, err := url.Parse(host)
	if err != nil {
		return nil, err
	}

	//change to http connection
	hurl.Scheme = "https"

	//create docker client
	m.client, err = docker.NewTLSClient(hurl.String(), filepath.Join(cert, "cert.pem"), filepath.Join(cert, "key.pem"), filepath.Join(cert, "ca.pem"))
	if err != nil {
		return nil, err
	}

	return m, nil
}

// hash a path using md5
func containerName(path string) (string, error) {

	//create md5 of full path
	hash := md5.New()
	_, err := hash.Write([]byte(path))
	if err != nil {
		return "", err
	}

	return fmt.Sprintf("pitmock_%s", hex.EncodeToString(hash.Sum(nil))), nil
}

// start a mock container by using examples from the given directory
func (m *Manager) Start(dir string, portb map[docker.Port][]docker.PortBinding) (*MockContainer, error) {

	//create name for container
	cname, err := containerName(dir)
	if err != nil {
		return nil, err
	}

	//@todo, grab the first and set private port to expected mock service
	//to the mock porcess binding
	for _, pb := range portb {
		portb[docker.Port(strconv.FormatInt(MockPrivatePort, 10)+"/tcp")] = pb
		break
	}

	//create the container
	c, err := m.client.CreateContainer(docker.CreateContainerOptions{
		Name: cname,
		Config: &docker.Config{
			Image: ImageName,
		},
	})

	if err != nil {
		return nil, err
	}

	//start the container we created
	err = m.client.StartContainer(c.ID, &docker.HostConfig{PortBindings: portb})
	if err != nil {
		return nil, err
	}

	//get container port mapping
	ci, err := m.client.InspectContainer(c.ID)
	if err != nil {
		return nil, err
	}

	//use docker host location to form url
	hurl, err := url.Parse(m.host)
	if err != nil {
		return nil, err
	}

	//tar examples into memory
	tar := bytes.NewBuffer(nil)
	err = dirtar.Tar(dir, tar)
	if err != nil {
		return nil, err
	}

	//wait for container to settle?
	//@todo very dirty business here
	<-time.After(time.Millisecond * 200)

	//get the external port for 8000 and turn into an url we can send http requests to
	//@todo, here we assume that the first configured port is the http interface, indicate explicetly
	//@todo, use the logic on line 77
	hurl.Scheme = "http"
	for _, pconfig := range ci.NetworkSettings.PortMappingAPI() {
		if pconfig.PrivatePort == MockPrivatePort {
			hurl.Host = strings.Replace(hurl.Host, ":2376", fmt.Sprintf(":%d", pconfig.PublicPort), 1)
		}
	}

	//send the upload request
	resp, err := http.Post(fmt.Sprintf("%s/_examples", hurl.String()), "application/x-tar", tar)
	if err != nil {
		return nil, err
	}

	//check if we uploaded correctly
	if resp.StatusCode != 201 {
		return nil, fmt.Errorf("Failed to upload examples to container '%s' %s", ci.ID, ci.Name)
	}

	//send HUP signal to 'root' process (dockpit/mock) to reload examples
	err = m.client.KillContainer(docker.KillContainerOptions{
		ID:     c.ID,
		Signal: docker.Signal(syscall.SIGHUP),
	})

	if err != nil {
		return nil, err
	}

	return &MockContainer{c.ID, hurl.String(), dir}, nil
}

// stop a mock container that was started from the given directory
func (m *Manager) Stop(dir string) error {

	//create name for container
	cname, err := containerName(dir)
	if err != nil {
		return err
	}

	//get all containers
	cs, err := m.client.ListContainers(docker.ListContainersOptions{})
	if err != nil {
		return err
	}

	//get container that matches the name
	// var container *docker.APIContainers
	var container docker.APIContainers
	for _, c := range cs {
		for _, n := range c.Names {
			if n[1:] == cname {
				container = c
			}
		}
	}

	//remove hard since mocks are ephemeral
	return m.client.RemoveContainer(docker.RemoveContainerOptions{
		ID:            container.ID,
		RemoveVolumes: true,
		Force:         true,
	})
}
<|endoftext|>


<commit_before>//  Copyright (c) 2014 Couchbase, Inc.
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the
//  License. You may obtain a copy of the License at
//    http://www.apache.org/licenses/LICENSE-2.0
//  Unless required by applicable law or agreed to in writing,
//  software distributed under the License is distributed on an "AS
//  IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
//  express or implied. See the License for the specific language
//  governing permissions and limitations under the License.

package main

import (
	"fmt"
	"hash/crc32"
	"io"
	"strings"

	log "github.com/couchbaselabs/clog"
)

// A planner assigns partitions to cbft's and to PIndexes on each cbft.
// NOTE: You *must* update PLANNER_VERSION if the planning algorithm
// or schema changes, following semver rules.

func (mgr *Manager) PlannerNOOP(msg string) {
	SyncWorkReq(mgr.plannerCh, WORK_NOOP, msg, nil)
}

func (mgr *Manager) PlannerKick(msg string) {
	SyncWorkReq(mgr.plannerCh, WORK_KICK, msg, nil)
}

func (mgr *Manager) PlannerLoop() {
	if mgr.cfg != nil { // Might be nil for testing.
		go func() {
			ec := make(chan CfgEvent)
			mgr.cfg.Subscribe(INDEX_DEFS_KEY, ec)
			mgr.cfg.Subscribe(CfgNodeDefsKey(NODE_DEFS_WANTED), ec)
			for e := range ec {
				mgr.PlannerKick("cfg changed, key: " + e.Key)
			}
		}()
	}

	for m := range mgr.plannerCh {
		var err error
		if m.op == WORK_KICK {
			changed, err := mgr.PlannerOnce(m.msg)
			if err != nil {
				log.Printf("error: PlannerOnce, err: %v", err)
				// Keep looping as perhaps it's a transient issue.
			} else if changed {
				mgr.JanitorKick("the plans have changed")
			}
		} else if m.op == WORK_NOOP {
			// NOOP.
		} else {
			err = fmt.Errorf("error: unknown planner op: %s, m: %#v", m.op, m)
		}
		if m.resCh != nil {
			if err != nil {
				m.resCh <- err
			}
			close(m.resCh)
		}
	}
}

func (mgr *Manager) PlannerOnce(reason string) (bool, error) {
	log.Printf("planner awakes, reason: %s", reason)

	if mgr.cfg == nil { // Can occur during testing.
		return false, fmt.Errorf("planner skipped due to nil cfg")
	}
	err := PlannerCheckVersion(mgr.cfg, mgr.version)
	if err != nil {
		return false, err
	}
	indexDefs, err := PlannerGetIndexDefs(mgr.cfg, mgr.version)
	if err != nil {
		return false, err
	}
	nodeDefs, err := PlannerGetNodeDefs(mgr.cfg, mgr.version, mgr.uuid, mgr.bindAddr)
	if err != nil {
		return false, err
	}
	planPIndexesPrev, cas, err := PlannerGetPlanPIndexes(mgr.cfg, mgr.version)
	if err != nil {
		return false, err
	}

	planPIndexes, err :=
		CalcPlan(indexDefs, nodeDefs, planPIndexesPrev, mgr.version, mgr.server)
	if err != nil {
		return false, fmt.Errorf("planner ended on CalcPlan, err: %v", err)
	}
	if planPIndexes == nil {
		return false, fmt.Errorf("planner found no plans from CalcPlan")
	}
	if SamePlanPIndexes(planPIndexes, planPIndexesPrev) {
		return false, nil
	}
	_, err = CfgSetPlanPIndexes(mgr.cfg, planPIndexes, cas)
	if err != nil {
		return false, fmt.Errorf("planner could not save new plan,"+
			" perhaps a concurrent planner won, cas: %d, err: %v",
			cas, err)
	}
	return true, nil
}

func PlannerCheckVersion(cfg Cfg, version string) error {
	ok, err := CheckVersion(cfg, version)
	if err != nil {
		return fmt.Errorf("planner skipped on CheckVersion err: %v", err)
	}
	if !ok {
		return fmt.Errorf("planner skipped with version too low: %v", version)
	}
	return nil
}

func PlannerGetIndexDefs(cfg Cfg, version string) (*IndexDefs, error) {
	indexDefs, _, err := CfgGetIndexDefs(cfg)
	if err != nil {
		return nil, fmt.Errorf("planner skipped on CfgGetIndexDefs err: %v", err)
	}
	if indexDefs == nil {
		return nil, fmt.Errorf("planner ended since no IndexDefs")
	}
	if VersionGTE(version, indexDefs.ImplVersion) == false {
		return nil, fmt.Errorf("planner ended since indexDefs.ImplVersion: %s"+
			" > version: %s", indexDefs.ImplVersion, version)
	}
	return indexDefs, nil
}

func PlannerGetNodeDefs(cfg Cfg, version, uuid, bindAddr string) (*NodeDefs, error) {
	nodeDefs, _, err := CfgGetNodeDefs(cfg, NODE_DEFS_WANTED)
	if err != nil {
		return nil, fmt.Errorf("planner skipped on CfgGetNodeDefs err: %v", err)
	}
	if nodeDefs == nil {
		return nil, fmt.Errorf("planner ended since no NodeDefs")
	}
	if VersionGTE(version, nodeDefs.ImplVersion) == false {
		return nil, fmt.Errorf("planner ended since nodeDefs.ImplVersion: %s"+
			" > version: %s", nodeDefs.ImplVersion, version)
	}
	nodeDef, exists := nodeDefs.NodeDefs[bindAddr]
	if !exists || nodeDef == nil {
		return nil, fmt.Errorf("planner ended since no NodeDef, bindAddr: %s", bindAddr)
	}
	if nodeDef.ImplVersion != version {
		return nil, fmt.Errorf("planner ended since NodeDef, bindAddr: %s,"+
			" NodeDef.ImplVersion: %s != version: %s",
			bindAddr, nodeDef.ImplVersion, version)
	}
	if nodeDef.UUID != uuid {
		return nil, fmt.Errorf("planner ended since NodeDef, bindAddr: %s,"+
			" NodeDef.UUID: %s != uuid: %s",
			bindAddr, nodeDef.UUID, uuid)
	}
	isPlanner := true
	if nodeDef.Tags != nil && len(nodeDef.Tags) > 0 {
		isPlanner = false
		for _, tag := range nodeDef.Tags {
			if tag == "planner" {
				isPlanner = true
			}
		}
	}
	if !isPlanner {
		return nil, fmt.Errorf("planner ended since node, bindAddr: %s,"+
			" is not a planner, tags: %#v", bindAddr, nodeDef.Tags)
	}
	return nodeDefs, nil
}

func PlannerGetPlanPIndexes(cfg Cfg, version string) (*PlanPIndexes, uint64, error) {
	planPIndexesPrev, cas, err := CfgGetPlanPIndexes(cfg)
	if err != nil {
		return nil, 0, fmt.Errorf("planner skipped on CfgGetPlanPIndexes err: %v", err)
	}
	if planPIndexesPrev == nil {
		planPIndexesPrev = NewPlanPIndexes(version)
	}
	if VersionGTE(version, planPIndexesPrev.ImplVersion) == false {
		return nil, 0, fmt.Errorf("planner ended on planPIndexesPrev.ImplVersion: %s"+
			" > version: %s", planPIndexesPrev.ImplVersion, version)
	}
	return planPIndexesPrev, cas, nil
}

// Split logical indexes into PIndexes and assign PIndexes to nodes.
func CalcPlan(indexDefs *IndexDefs, nodeDefs *NodeDefs,
	planPIndexesPrev *PlanPIndexes, version, server string) (
	*PlanPIndexes, error) {
	// This simple planner assigns at most MaxPartitionsPerPIndex
	// number of partitions onto a PIndex.  And then this planner
	// assigns every PIndex onto every node (so every signle node
	// has a replica of that PIndex).
	//
	// TODO: Assign PIndexes to nodes in a fancier way.
	// TODO: This simple planner doesn't handle cbft node membership changes
	// right, and should instead reassign pindexes on leaving nodes,
	// and rebalance pindexes on remaining (and newly added) nodes.
	if indexDefs == nil || nodeDefs == nil {
		return nil, nil
	}

	planPIndexes := NewPlanPIndexes(version)

	for _, indexDef := range indexDefs.IndexDefs {
		maxPartitionsPerPIndex := indexDef.PlanParams.MaxPartitionsPerPIndex

		sourcePartitionsArr, err := DataSourcePartitions(indexDef.SourceType,
			indexDef.SourceName, indexDef.SourceUUID, indexDef.SourceParams, server)
		if err != nil {
			log.Printf("error: planner could not get partitions,"+
				" indexDef: %#v, err: %v", indexDef, err)
			continue
		}

		addPlanPIndex := func(sourcePartitionsCurr []string) {
			sourcePartitions := strings.Join(sourcePartitionsCurr, ",")

			planPIndex := &PlanPIndex{
				Name:             PlanPIndexName(indexDef, sourcePartitions),
				UUID:             NewUUID(),
				IndexType:        indexDef.Type,
				IndexName:        indexDef.Name,
				IndexUUID:        indexDef.UUID,
				IndexSchema:      indexDef.Schema,
				SourceType:       indexDef.SourceType,
				SourceName:       indexDef.SourceName,
				SourceUUID:       indexDef.SourceUUID,
				SourceParams:     indexDef.SourceParams,
				SourcePartitions: sourcePartitions,
				NodeUUIDs:        make(map[string]string),
			}

			// TODO: Assign PIndexes to nodes in a fancier way.
			for _, nodeDef := range nodeDefs.NodeDefs {
				planPIndex.NodeUUIDs[nodeDef.UUID] =
					PLAN_PINDEX_NODE_READ + PLAN_PINDEX_NODE_WRITE
			}

			planPIndexes.PlanPIndexes[planPIndex.Name] = planPIndex
		}

		sourcePartitionsCurr := []string{}
		for _, sourcePartition := range sourcePartitionsArr {
			sourcePartitionsCurr = append(sourcePartitionsCurr, sourcePartition)
			if maxPartitionsPerPIndex > 0 &&
				len(sourcePartitionsCurr) >= maxPartitionsPerPIndex {
				addPlanPIndex(sourcePartitionsCurr)
				sourcePartitionsCurr = []string{}
			}
		}

		if len(sourcePartitionsCurr) > 0 || len(sourcePartitionsArr) <= 0 {
			addPlanPIndex(sourcePartitionsCurr)
		}
	}

	return planPIndexes, nil
}

// NOTE: PlanPIndex.Name must be unique across the cluster and ideally
// functionally based off of the indexDef so that the SamePlanPIndex()
// comparison works even if concurrent planners are racing to
// calculate plans.
//
// NOTE: We can't use sourcePartitions directly as part of a
// PlanPIndex.Name suffix because in vbucket/hash partitioning the
// string would be too long -- since PIndexes might use
// PlanPIndex.Name for filesystem paths.
func PlanPIndexName(indexDef *IndexDef, sourcePartitions string) string {
	h := crc32.NewIEEE()
	io.WriteString(h, sourcePartitions)
	return indexDef.Name + "_" + indexDef.UUID + "_" + fmt.Sprintf("%x", h.Sum32())
}
<commit_msg>skip planPIndexes for indexDef's with no pindex impl type<commit_after>//  Copyright (c) 2014 Couchbase, Inc.
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the
//  License. You may obtain a copy of the License at
//    http://www.apache.org/licenses/LICENSE-2.0
//  Unless required by applicable law or agreed to in writing,
//  software distributed under the License is distributed on an "AS
//  IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
//  express or implied. See the License for the specific language
//  governing permissions and limitations under the License.

package main

import (
	"fmt"
	"hash/crc32"
	"io"
	"strings"

	log "github.com/couchbaselabs/clog"
)

// A planner assigns partitions to cbft's and to PIndexes on each cbft.
// NOTE: You *must* update PLANNER_VERSION if the planning algorithm
// or schema changes, following semver rules.

func (mgr *Manager) PlannerNOOP(msg string) {
	SyncWorkReq(mgr.plannerCh, WORK_NOOP, msg, nil)
}

func (mgr *Manager) PlannerKick(msg string) {
	SyncWorkReq(mgr.plannerCh, WORK_KICK, msg, nil)
}

func (mgr *Manager) PlannerLoop() {
	if mgr.cfg != nil { // Might be nil for testing.
		go func() {
			ec := make(chan CfgEvent)
			mgr.cfg.Subscribe(INDEX_DEFS_KEY, ec)
			mgr.cfg.Subscribe(CfgNodeDefsKey(NODE_DEFS_WANTED), ec)
			for e := range ec {
				mgr.PlannerKick("cfg changed, key: " + e.Key)
			}
		}()
	}

	for m := range mgr.plannerCh {
		var err error
		if m.op == WORK_KICK {
			changed, err := mgr.PlannerOnce(m.msg)
			if err != nil {
				log.Printf("error: PlannerOnce, err: %v", err)
				// Keep looping as perhaps it's a transient issue.
			} else if changed {
				mgr.JanitorKick("the plans have changed")
			}
		} else if m.op == WORK_NOOP {
			// NOOP.
		} else {
			err = fmt.Errorf("error: unknown planner op: %s, m: %#v", m.op, m)
		}
		if m.resCh != nil {
			if err != nil {
				m.resCh <- err
			}
			close(m.resCh)
		}
	}
}

func (mgr *Manager) PlannerOnce(reason string) (bool, error) {
	log.Printf("planner awakes, reason: %s", reason)

	if mgr.cfg == nil { // Can occur during testing.
		return false, fmt.Errorf("planner skipped due to nil cfg")
	}
	err := PlannerCheckVersion(mgr.cfg, mgr.version)
	if err != nil {
		return false, err
	}
	indexDefs, err := PlannerGetIndexDefs(mgr.cfg, mgr.version)
	if err != nil {
		return false, err
	}
	nodeDefs, err := PlannerGetNodeDefs(mgr.cfg, mgr.version, mgr.uuid, mgr.bindAddr)
	if err != nil {
		return false, err
	}
	planPIndexesPrev, cas, err := PlannerGetPlanPIndexes(mgr.cfg, mgr.version)
	if err != nil {
		return false, err
	}

	planPIndexes, err :=
		CalcPlan(indexDefs, nodeDefs, planPIndexesPrev, mgr.version, mgr.server)
	if err != nil {
		return false, fmt.Errorf("planner ended on CalcPlan, err: %v", err)
	}
	if planPIndexes == nil {
		return false, fmt.Errorf("planner found no plans from CalcPlan")
	}
	if SamePlanPIndexes(planPIndexes, planPIndexesPrev) {
		return false, nil
	}
	_, err = CfgSetPlanPIndexes(mgr.cfg, planPIndexes, cas)
	if err != nil {
		return false, fmt.Errorf("planner could not save new plan,"+
			" perhaps a concurrent planner won, cas: %d, err: %v",
			cas, err)
	}
	return true, nil
}

func PlannerCheckVersion(cfg Cfg, version string) error {
	ok, err := CheckVersion(cfg, version)
	if err != nil {
		return fmt.Errorf("planner skipped on CheckVersion err: %v", err)
	}
	if !ok {
		return fmt.Errorf("planner skipped with version too low: %v", version)
	}
	return nil
}

func PlannerGetIndexDefs(cfg Cfg, version string) (*IndexDefs, error) {
	indexDefs, _, err := CfgGetIndexDefs(cfg)
	if err != nil {
		return nil, fmt.Errorf("planner skipped on CfgGetIndexDefs err: %v", err)
	}
	if indexDefs == nil {
		return nil, fmt.Errorf("planner ended since no IndexDefs")
	}
	if VersionGTE(version, indexDefs.ImplVersion) == false {
		return nil, fmt.Errorf("planner ended since indexDefs.ImplVersion: %s"+
			" > version: %s", indexDefs.ImplVersion, version)
	}
	return indexDefs, nil
}

func PlannerGetNodeDefs(cfg Cfg, version, uuid, bindAddr string) (*NodeDefs, error) {
	nodeDefs, _, err := CfgGetNodeDefs(cfg, NODE_DEFS_WANTED)
	if err != nil {
		return nil, fmt.Errorf("planner skipped on CfgGetNodeDefs err: %v", err)
	}
	if nodeDefs == nil {
		return nil, fmt.Errorf("planner ended since no NodeDefs")
	}
	if VersionGTE(version, nodeDefs.ImplVersion) == false {
		return nil, fmt.Errorf("planner ended since nodeDefs.ImplVersion: %s"+
			" > version: %s", nodeDefs.ImplVersion, version)
	}
	nodeDef, exists := nodeDefs.NodeDefs[bindAddr]
	if !exists || nodeDef == nil {
		return nil, fmt.Errorf("planner ended since no NodeDef, bindAddr: %s", bindAddr)
	}
	if nodeDef.ImplVersion != version {
		return nil, fmt.Errorf("planner ended since NodeDef, bindAddr: %s,"+
			" NodeDef.ImplVersion: %s != version: %s",
			bindAddr, nodeDef.ImplVersion, version)
	}
	if nodeDef.UUID != uuid {
		return nil, fmt.Errorf("planner ended since NodeDef, bindAddr: %s,"+
			" NodeDef.UUID: %s != uuid: %s",
			bindAddr, nodeDef.UUID, uuid)
	}
	isPlanner := true
	if nodeDef.Tags != nil && len(nodeDef.Tags) > 0 {
		isPlanner = false
		for _, tag := range nodeDef.Tags {
			if tag == "planner" {
				isPlanner = true
			}
		}
	}
	if !isPlanner {
		return nil, fmt.Errorf("planner ended since node, bindAddr: %s,"+
			" is not a planner, tags: %#v", bindAddr, nodeDef.Tags)
	}
	return nodeDefs, nil
}

func PlannerGetPlanPIndexes(cfg Cfg, version string) (*PlanPIndexes, uint64, error) {
	planPIndexesPrev, cas, err := CfgGetPlanPIndexes(cfg)
	if err != nil {
		return nil, 0, fmt.Errorf("planner skipped on CfgGetPlanPIndexes err: %v", err)
	}
	if planPIndexesPrev == nil {
		planPIndexesPrev = NewPlanPIndexes(version)
	}
	if VersionGTE(version, planPIndexesPrev.ImplVersion) == false {
		return nil, 0, fmt.Errorf("planner ended on planPIndexesPrev.ImplVersion: %s"+
			" > version: %s", planPIndexesPrev.ImplVersion, version)
	}
	return planPIndexesPrev, cas, nil
}

// Split logical indexes into PIndexes and assign PIndexes to nodes.
func CalcPlan(indexDefs *IndexDefs, nodeDefs *NodeDefs,
	planPIndexesPrev *PlanPIndexes, version, server string) (
	*PlanPIndexes, error) {
	// This simple planner assigns at most MaxPartitionsPerPIndex
	// number of partitions onto a PIndex.  And then this planner
	// assigns every PIndex onto every node (so every signle node
	// has a replica of that PIndex).
	//
	// TODO: Assign PIndexes to nodes in a fancier way.
	// TODO: This simple planner doesn't handle cbft node membership changes
	// right, and should instead reassign pindexes on leaving nodes,
	// and rebalance pindexes on remaining (and newly added) nodes.
	if indexDefs == nil || nodeDefs == nil {
		return nil, nil
	}

	planPIndexes := NewPlanPIndexes(version)

	for _, indexDef := range indexDefs.IndexDefs {
		_, exists := pindexImplTypes[indexDef.Type]
		if !exists {
			continue // Skip indexDef's that have no pindexImplType.
		}

		maxPartitionsPerPIndex := indexDef.PlanParams.MaxPartitionsPerPIndex

		sourcePartitionsArr, err := DataSourcePartitions(indexDef.SourceType,
			indexDef.SourceName, indexDef.SourceUUID, indexDef.SourceParams, server)
		if err != nil {
			log.Printf("error: planner could not get partitions,"+
				" indexDef: %#v, err: %v", indexDef, err)
			continue
		}

		addPlanPIndex := func(sourcePartitionsCurr []string) {
			sourcePartitions := strings.Join(sourcePartitionsCurr, ",")

			planPIndex := &PlanPIndex{
				Name:             PlanPIndexName(indexDef, sourcePartitions),
				UUID:             NewUUID(),
				IndexType:        indexDef.Type,
				IndexName:        indexDef.Name,
				IndexUUID:        indexDef.UUID,
				IndexSchema:      indexDef.Schema,
				SourceType:       indexDef.SourceType,
				SourceName:       indexDef.SourceName,
				SourceUUID:       indexDef.SourceUUID,
				SourceParams:     indexDef.SourceParams,
				SourcePartitions: sourcePartitions,
				NodeUUIDs:        make(map[string]string),
			}

			// TODO: Assign PIndexes to nodes in a fancier way.
			for _, nodeDef := range nodeDefs.NodeDefs {
				planPIndex.NodeUUIDs[nodeDef.UUID] =
					PLAN_PINDEX_NODE_READ + PLAN_PINDEX_NODE_WRITE
			}

			planPIndexes.PlanPIndexes[planPIndex.Name] = planPIndex
		}

		sourcePartitionsCurr := []string{}
		for _, sourcePartition := range sourcePartitionsArr {
			sourcePartitionsCurr = append(sourcePartitionsCurr, sourcePartition)
			if maxPartitionsPerPIndex > 0 &&
				len(sourcePartitionsCurr) >= maxPartitionsPerPIndex {
				addPlanPIndex(sourcePartitionsCurr)
				sourcePartitionsCurr = []string{}
			}
		}

		if len(sourcePartitionsCurr) > 0 || len(sourcePartitionsArr) <= 0 {
			addPlanPIndex(sourcePartitionsCurr)
		}
	}

	return planPIndexes, nil
}

// NOTE: PlanPIndex.Name must be unique across the cluster and ideally
// functionally based off of the indexDef so that the SamePlanPIndex()
// comparison works even if concurrent planners are racing to
// calculate plans.
//
// NOTE: We can't use sourcePartitions directly as part of a
// PlanPIndex.Name suffix because in vbucket/hash partitioning the
// string would be too long -- since PIndexes might use
// PlanPIndex.Name for filesystem paths.
func PlanPIndexName(indexDef *IndexDef, sourcePartitions string) string {
	h := crc32.NewIEEE()
	io.WriteString(h, sourcePartitions)
	return indexDef.Name + "_" + indexDef.UUID + "_" + fmt.Sprintf("%x", h.Sum32())
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"runtime/pprof"
	"time"

	"github.com/mrfuxi/neural"
	"github.com/petar/GoMNIST"
)

var (
	cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")
	nnSaveFile = flag.String("save-file", "", "Save neural network to file")
	nnLoadFile = flag.String("load-file", "", "Load neural network to file")
	inputSize  = GoMNIST.Width * GoMNIST.Height
)

func prepareMnistData(rawData *GoMNIST.Set) []neural.TrainExample {
	trainData := make([]neural.TrainExample, rawData.Count())
	for i := range trainData {
		image, label := rawData.Get(i)
		trainData[i].Input = make([]float64, inputSize, inputSize)
		trainData[i].Output = make([]float64, 10, 10)
		for j, pix := range image {
			trainData[i].Input[j] = (float64(pix)/255)*0.9 + 0.1
		}

		for j := range trainData[i].Output {
			trainData[i].Output[j] = 0.1
		}
		trainData[i].Output[label] = 0.9
	}
	return trainData
}

func loadTestData() ([]neural.TrainExample, []neural.TrainExample) {
	train, test, err := GoMNIST.Load("./data")
	if err != nil {
		panic(err)
	}

	trainData := prepareMnistData(train)
	testData := prepareMnistData(test)
	return trainData, testData
}

func epocheCallback(nn neural.Evaluator, cost neural.Cost, trainData, testData []neural.TrainExample) neural.EpocheCallback {
	return func(epoche int, dt time.Duration) {
		avgCost, errors := neural.CalculateCorrectness(nn, cost, testData)
		fmt.Printf("%v,%v,%v\n", epoche, avgCost, errors)
	}
}

func main() {
	trainData, testData := loadTestData()

	activator := neural.NewSigmoidActivator()
	nn := neural.NewNeuralNetwork(
		[]int{inputSize, 100, 10},
		neural.NewFullyConnectedLayer(activator),
		neural.NewFullyConnectedLayer(activator),
	)

	flag.Parse()

	if *nnLoadFile != "" {
		fn, err := os.Open(*nnLoadFile)
		if err != nil {
			log.Fatalln(err)
		}
		if err := neural.Load(nn, fn); err != nil {
			log.Fatalln(err)
		}
	}

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	cost := neural.NewCrossEntropyCost()
	options := neural.TrainOptions{
		Epochs:         10,
		MiniBatchSize:  10,
		LearningRate:   4,
		TrainerFactory: neural.NewBackpropagationTrainer,
		EpocheCallback: epocheCallback(nn, cost, trainData, testData),
		Cost:           cost,
	}

	t0 := time.Now()
	neural.Train(nn, trainData, options)
	dt := time.Since(t0)

	fmt.Println("Training complete in", dt)

	if *nnSaveFile != "" {
		fn, err := os.OpenFile(*nnSaveFile, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0666)
		if err != nil {
			log.Fatalln(err)
		}
		if err := neural.Save(nn, fn); err != nil {
			log.Fatalln(err)
		}
	}
}
<commit_msg>MNIST example to use softmax and L2 regularilisation<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"runtime/pprof"
	"time"

	"github.com/mrfuxi/neural"
	"github.com/petar/GoMNIST"
)

var (
	cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")
	nnSaveFile = flag.String("save-file", "", "Save neural network to file")
	nnLoadFile = flag.String("load-file", "", "Load neural network to file")
	inputSize  = GoMNIST.Width * GoMNIST.Height
)

func prepareMnistData(rawData *GoMNIST.Set) []neural.TrainExample {
	trainData := make([]neural.TrainExample, rawData.Count())
	for i := range trainData {
		image, label := rawData.Get(i)
		trainData[i].Input = make([]float64, inputSize, inputSize)
		trainData[i].Output = make([]float64, 10, 10)
		for j, pix := range image {
			trainData[i].Input[j] = (float64(pix)/255)*0.9 + 0.1
		}

		for j := range trainData[i].Output {
			trainData[i].Output[j] = 0
			// trainData[i].Output[j] = 0.1
		}
		trainData[i].Output[label] = 1
		// trainData[i].Output[label] = 0.9
	}
	return trainData
}

func loadTestData() ([]neural.TrainExample, []neural.TrainExample) {
	train, test, err := GoMNIST.Load("./data")
	if err != nil {
		panic(err)
	}

	trainData := prepareMnistData(train)
	testData := prepareMnistData(test)
	return trainData, testData
}

func epocheCallback(nn neural.Evaluator, cost neural.Cost, trainData, testData []neural.TrainExample) neural.EpocheCallback {
	return func(epoche int, dt time.Duration) {
		avgCost, errors := neural.CalculateCorrectness(nn, cost, testData)
		fmt.Printf("%v,%v,%v\n", epoche, avgCost, errors)
	}
}

func main() {
	trainData, testData := loadTestData()

	activator := neural.NewSigmoidActivator()
	outActivator := neural.NewSoftmaxFunction()
	nn := neural.NewNeuralNetwork(
		[]int{inputSize, 100, 10},
		neural.NewFullyConnectedLayer(activator),
		neural.NewFullyConnectedLayer(outActivator),
	)

	flag.Parse()

	if *nnLoadFile != "" {
		fn, err := os.Open(*nnLoadFile)
		if err != nil {
			log.Fatalln(err)
		}
		if err := neural.Load(nn, fn); err != nil {
			log.Fatalln(err)
		}
	}

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	cost := neural.NewLogLikelihoodCost()
	options := neural.TrainOptions{
		Epochs:         40,
		MiniBatchSize:  10,
		LearningRate:   0.1,
		Regularization: 5,
		TrainerFactory: neural.NewBackpropagationTrainer,
		EpocheCallback: epocheCallback(nn, cost, trainData, testData),
		Cost:           cost,
	}

	t0 := time.Now()
	neural.Train(nn, trainData, options)
	dt := time.Since(t0)

	fmt.Println("Training complete in", dt)

	if *nnSaveFile != "" {
		fn, err := os.OpenFile(*nnSaveFile, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0666)
		if err != nil {
			log.Fatalln(err)
		}
		if err := neural.Save(nn, fn); err != nil {
			log.Fatalln(err)
		}
	}
}
<|endoftext|>