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

package docker

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/globocom/config"
	"github.com/globocom/tsuru/fs"
	"github.com/globocom/tsuru/log"
	"github.com/globocom/tsuru/provision"
	"github.com/globocom/tsuru/repository"
	"io"
	"io/ioutil"
	"labix.org/v2/mgo/bson"
	"os"
	"os/user"
	"path"
	"strings"
	"time"
)

var fsystem fs.Fs

func filesystem() fs.Fs {
	if fsystem == nil {
		fsystem = fs.OsFs{}
	}
	return fsystem
}

// runCmd executes commands and log the given stdout and stderror.
func runCmd(cmd string, args ...string) (string, error) {
	out := bytes.Buffer{}
	err := executor().Execute(cmd, args, nil, &out, &out)
	log.Printf("running the cmd: %s with the args: %s", cmd, args)
	return out.String(), err
}

func getSSHCommands() ([]string, error) {
	addKeyCommand, err := config.GetString("docker:ssh:add-key-cmd")
	if err != nil {
		return nil, err
	}
	keyFile, err := config.GetString("docker:ssh:public-key")
	if err != nil {
		if u, err := user.Current(); err == nil {
			keyFile = path.Join(u.HomeDir, ".ssh", "id_rsa.pub")
		} else {
			keyFile = os.ExpandEnv("${HOME}/.ssh/id_rsa.pub")
		}
	}
	f, err := filesystem().Open(keyFile)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	keyContent, err := ioutil.ReadAll(f)
	if err != nil {
		return nil, err
	}
	sshdPath, err := config.GetString("docker:ssh:sshd-path")
	if err != nil {
		sshdPath = "/usr/sbin/sshd"
	}
	return []string{
		fmt.Sprintf("%s %s", addKeyCommand, bytes.TrimSpace(keyContent)),
		sshdPath + " -D",
	}, nil
}

func commandToRun(app provision.App) ([]string, error) {
	docker, err := config.GetString("docker:binary")
	if err != nil {
		return nil, err
	}
	repoNamespace, err := config.GetString("docker:repository-namespace")
	if err != nil {
		return nil, err
	}
	port, err := config.GetString("docker:run-cmd:port")
	if err != nil {
		return nil, err
	}
	commands, err := getSSHCommands()
	if err != nil {
		return nil, err
	}
	imageName := fmt.Sprintf("%s/%s", repoNamespace, app.GetPlatform()) // TODO (flaviamissi): should use same algorithm as image.repositoryName
	containerCmd := strings.Join(commands, " && ")
	wholeCmd := []string{docker, "run", "-d", "-t", "-p", port, imageName, "/bin/bash", "-c", containerCmd}
	return wholeCmd, nil
}

func getPort() (string, error) {
	return config.GetString("docker:run-cmd:port")
}

type container struct {
	Id      string `bson:"_id"`
	AppName string
	Type    string
	Ip      string
	Port    string
}

// newContainer creates a new container in Docker and stores it in the database.
//
// TODO (flaviamissi): make it atomic
func newContainer(app provision.App) (*container, error) {
	appName := app.GetName()
	c := container{
		AppName: appName,
		Type:    app.GetPlatform(),
	}
	err := c.create(app)
	if err != nil {
		log.Printf("Error creating container %s", appName)
		log.Printf("Error was: %s", err.Error())
		return nil, err
	}
	return &c, nil
}

func (c *container) inspect() (map[string]interface{}, error) {
	docker, err := config.GetString("docker:binary")
	if err != nil {
		return nil, err
	}
	out, err := runCmd(docker, "inspect", c.Id)
	if err != nil {
		return nil, err
	}
	var r map[string]interface{}
	err = json.Unmarshal([]byte(out), &r)
	if err != nil {
		return nil, err
	}
	return r, nil
}

// hostPort returns the host port mapped for the container.
func (c *container) hostPort() (string, error) {
	if c.Port == "" {
		return "", errors.New("Container does not contain any mapped port")
	}
	data, err := c.inspect()
	if err != nil {
		return "", err
	}
	mappedPorts := data["NetworkSettings"].(map[string]interface{})["PortMapping"].(map[string]interface{})
	if port, ok := mappedPorts[c.Port]; ok {
		return port.(string), nil
	}
	return "", fmt.Errorf("Container port %s is not mapped to any host port", c.Port)
}

// ip returns the ip for the container.
func (c *container) ip() (string, error) {
	result, err := c.inspect()
	if err != nil {
		msg := fmt.Sprintf("error(%s) parsing json from docker when trying to get ipaddress", err)
		log.Print(msg)
		return "", errors.New(msg)
	}
	if ns, ok := result["NetworkSettings"]; !ok || ns == nil {
		msg := "Error when getting container information. NetworkSettings is missing."
		log.Print(msg)
		return "", errors.New(msg)
	}
	networkSettings := result["NetworkSettings"].(map[string]interface{})
	instanceIp := networkSettings["IpAddress"].(string)
	if instanceIp == "" {
		msg := "error: Can't get ipaddress..."
		log.Print(msg)
		return "", errors.New(msg)
	}
	log.Printf("Instance IpAddress: %s", instanceIp)
	return instanceIp, nil
}

// create creates a docker container, stores it on the database and adds a route to it.
//
// It receives the related application in order to choose the correct
// docker image and the repository to pass to the script that will take
// care of the deploy, and a function to generate the correct command ran by
// docker, which might be to deploy a container or to run and expose a
// container for an application.
func (c *container) create(app provision.App) error {
	hostAddr, err := config.Get("docker:host-address")
	if err != nil {
		return err
	}
	cmd, err := commandToRun(app)
	if err != nil {
		return err
	}
	port, err := getPort()
	if err != nil {
		return err
	}
	id, err := runCmd(cmd[0], cmd[1:]...)
	if err != nil {
		return err
	}
	id = strings.TrimSpace(id)
	log.Printf("docker id=%s", id)
	c.Id = strings.TrimSpace(id)
	c.Port = port
	ip, err := c.ip()
	if err != nil {
		return err
	}
	c.Ip = ip
	coll := collection()
	defer coll.Database.Session.Close()
	if err := coll.Insert(c); err != nil {
		log.Print(err)
		return err
	}
	r, err := getRouter()
	if err != nil {
		return err
	}
	hostPort, err := c.hostPort()
	if err != nil {
		hostPort = c.Port
	}
	return r.AddRoute(app.GetName(), fmt.Sprintf("http://%s:%s", hostAddr, hostPort))
}

func (c *container) deploy(w io.Writer) error {
	deployCmd, err := config.GetString("docker:deploy-cmd")
	if err != nil {
		return err
	}
	runBin, err := config.GetString("docker:run-cmd:bin")
	if err != nil {
		return err
	}
	runArgs, _ := config.GetString("docker:run-cmd:args")
	appRepo := repository.GetReadOnlyUrl(c.AppName)
	for i := 0; i < 5; i++ {
		err = c.ssh(w, w, deployCmd, appRepo)
		if err == nil {
			break
		}
		log.Printf("SSH to the container %q failed. Will retry.", c.Id)
		time.Sleep(100e6)
	}
	if err != nil {
		return err
	}
	return c.ssh(w, w, runBin, strings.Fields(runArgs)...)
}

// remove removes a docker container.
func (c *container) remove() error {
	docker, err := config.GetString("docker:binary")
	if err != nil {
		return err
	}
	log.Printf("Removing container %s from docker", c.Id)
	out, err := runCmd(docker, "rm", c.Id)
	if err != nil {
		log.Printf("Failed to remove container from docker: %s", err.Error())
		log.Printf("Command output: %s", out)
		return err
	}
	log.Printf("Removing container %s from database", c.Id)
	coll := collection()
	defer coll.Database.Session.Close()
	if err := coll.RemoveId(c.Id); err != nil {
		log.Printf("Failed to remove container from database: %s", err.Error())
		return err
	}
	r, err := getRouter()
	if err != nil {
		log.Printf("Failed to obtain router: %s", err.Error())
		return err
	}
	if err := r.RemoveRoute(c.AppName); err != nil {
		log.Printf("Failed to remove route: %s", err.Error())
		return err
	}
	return nil
}

func (c *container) ssh(stdout, stderr io.Writer, cmd string, args ...string) error {
	user, err := config.GetString("docker:ssh:user")
	if err != nil {
		return err
	}
	sshArgs := []string{c.Ip, "-l", user, "-o", "StrictHostKeyChecking no"}
	if keyFile, err := config.GetString("docker:ssh:private-key"); err == nil {
		sshArgs = append(sshArgs, "-i", keyFile)
	}
	sshArgs = append(sshArgs, "--", cmd)
	sshArgs = append(sshArgs, args...)
	return executor().Execute("ssh", sshArgs, nil, stdout, stderr)
}

// image represents a docker image.
type image struct {
	Name string
	Id   string
}

// repositoryName returns the image repository name for a given image.
//
// Repository is a docker concept, the image actually does not have a name,
// it has a repository, that is a composed name, e.g.: tsuru/base.
// Tsuru will always use a namespace, defined in tsuru.conf.
// Additionally, tsuru will use the application's name to do that composition.
func (img *image) repositoryName() string {
	repoNamespace, err := config.GetString("docker:repository-namespace")
	if err != nil {
		log.Printf("Tsuru is misconfigured. docker:repository-namespace config is missing.")
		return ""
	}
	return fmt.Sprintf("%s/%s", repoNamespace, img.Name)
}

// commit commits an image in docker
//
// This is another docker concept, in order to generate an image from a container
// one must commit it.
func (img *image) commit(cId string) (string, error) {
	docker, err := config.GetString("docker:binary")
	if err != nil {
		log.Printf("Tsuru is misconfigured. docker:binary config is missing.")
		return "", err
	}
	log.Printf("attempting to commit image from container %s", cId)
	rName := img.repositoryName()
	id, err := runCmd(docker, "commit", cId, rName)
	if err != nil {
		log.Printf("Could not commit docker image: %s", err.Error())
		return "", err
	}
	img.Id = strings.Replace(id, "\n", "", -1)
	if err := imagesCollection().Insert(&img); err != nil {
		log.Printf("Could not store image information %s", err.Error())
		return "", err
	}
	return img.Id, nil
}

// remove removes an image from docker registry
func (img *image) remove() error {
	docker, err := config.GetString("docker:binary")
	if err != nil {
		log.Printf("Tsuru is misconfigured. docker:binary config is missing.")
		return err
	}
	log.Printf("attempting to remove image %s from docker", img.repositoryName())
	_, err = runCmd(docker, "rmi", img.Id)
	if err != nil {
		log.Printf("Could not remove image %s from docker: %s", img.Id, err.Error())
		return err
	}
	err = imagesCollection().Remove(bson.M{"name": img.Name})
	if err != nil {
		log.Printf("Could not remove image %s from mongo: %s", img.Id, err.Error())
		return err
	}
	return nil
}

func getContainer(id string) (*container, error) {
	var c container
	coll := collection()
	defer coll.Database.Session.Close()
	err := coll.Find(bson.M{"_id": id}).One(&c)
	if err != nil {
		return nil, err
	}
	return &c, nil
}

func listAppContainers(appName string) ([]container, error) {
	var containers []container
	err := collection().Find(bson.M{"appname": appName}).All(&containers)
	return containers, err
}
<commit_msg>provision/docker: passing container ip on RemoveRoute call<commit_after>// Copyright 2013 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package docker

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/globocom/config"
	"github.com/globocom/tsuru/fs"
	"github.com/globocom/tsuru/log"
	"github.com/globocom/tsuru/provision"
	"github.com/globocom/tsuru/repository"
	"io"
	"io/ioutil"
	"labix.org/v2/mgo/bson"
	"os"
	"os/user"
	"path"
	"strings"
	"time"
)

var fsystem fs.Fs

func filesystem() fs.Fs {
	if fsystem == nil {
		fsystem = fs.OsFs{}
	}
	return fsystem
}

// runCmd executes commands and log the given stdout and stderror.
func runCmd(cmd string, args ...string) (string, error) {
	out := bytes.Buffer{}
	err := executor().Execute(cmd, args, nil, &out, &out)
	log.Printf("running the cmd: %s with the args: %s", cmd, args)
	return out.String(), err
}

func getSSHCommands() ([]string, error) {
	addKeyCommand, err := config.GetString("docker:ssh:add-key-cmd")
	if err != nil {
		return nil, err
	}
	keyFile, err := config.GetString("docker:ssh:public-key")
	if err != nil {
		if u, err := user.Current(); err == nil {
			keyFile = path.Join(u.HomeDir, ".ssh", "id_rsa.pub")
		} else {
			keyFile = os.ExpandEnv("${HOME}/.ssh/id_rsa.pub")
		}
	}
	f, err := filesystem().Open(keyFile)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	keyContent, err := ioutil.ReadAll(f)
	if err != nil {
		return nil, err
	}
	sshdPath, err := config.GetString("docker:ssh:sshd-path")
	if err != nil {
		sshdPath = "/usr/sbin/sshd"
	}
	return []string{
		fmt.Sprintf("%s %s", addKeyCommand, bytes.TrimSpace(keyContent)),
		sshdPath + " -D",
	}, nil
}

func commandToRun(app provision.App) ([]string, error) {
	docker, err := config.GetString("docker:binary")
	if err != nil {
		return nil, err
	}
	repoNamespace, err := config.GetString("docker:repository-namespace")
	if err != nil {
		return nil, err
	}
	port, err := config.GetString("docker:run-cmd:port")
	if err != nil {
		return nil, err
	}
	commands, err := getSSHCommands()
	if err != nil {
		return nil, err
	}
	imageName := fmt.Sprintf("%s/%s", repoNamespace, app.GetPlatform()) // TODO (flaviamissi): should use same algorithm as image.repositoryName
	containerCmd := strings.Join(commands, " && ")
	wholeCmd := []string{docker, "run", "-d", "-t", "-p", port, imageName, "/bin/bash", "-c", containerCmd}
	return wholeCmd, nil
}

func getPort() (string, error) {
	return config.GetString("docker:run-cmd:port")
}

type container struct {
	Id      string `bson:"_id"`
	AppName string
	Type    string
	Ip      string
	Port    string
}

// newContainer creates a new container in Docker and stores it in the database.
//
// TODO (flaviamissi): make it atomic
func newContainer(app provision.App) (*container, error) {
	appName := app.GetName()
	c := container{
		AppName: appName,
		Type:    app.GetPlatform(),
	}
	err := c.create(app)
	if err != nil {
		log.Printf("Error creating container %s", appName)
		log.Printf("Error was: %s", err.Error())
		return nil, err
	}
	return &c, nil
}

func (c *container) inspect() (map[string]interface{}, error) {
	docker, err := config.GetString("docker:binary")
	if err != nil {
		return nil, err
	}
	out, err := runCmd(docker, "inspect", c.Id)
	if err != nil {
		return nil, err
	}
	var r map[string]interface{}
	err = json.Unmarshal([]byte(out), &r)
	if err != nil {
		return nil, err
	}
	return r, nil
}

// hostPort returns the host port mapped for the container.
func (c *container) hostPort() (string, error) {
	if c.Port == "" {
		return "", errors.New("Container does not contain any mapped port")
	}
	data, err := c.inspect()
	if err != nil {
		return "", err
	}
	mappedPorts := data["NetworkSettings"].(map[string]interface{})["PortMapping"].(map[string]interface{})
	if port, ok := mappedPorts[c.Port]; ok {
		return port.(string), nil
	}
	return "", fmt.Errorf("Container port %s is not mapped to any host port", c.Port)
}

// ip returns the ip for the container.
func (c *container) ip() (string, error) {
	result, err := c.inspect()
	if err != nil {
		msg := fmt.Sprintf("error(%s) parsing json from docker when trying to get ipaddress", err)
		log.Print(msg)
		return "", errors.New(msg)
	}
	if ns, ok := result["NetworkSettings"]; !ok || ns == nil {
		msg := "Error when getting container information. NetworkSettings is missing."
		log.Print(msg)
		return "", errors.New(msg)
	}
	networkSettings := result["NetworkSettings"].(map[string]interface{})
	instanceIp := networkSettings["IpAddress"].(string)
	if instanceIp == "" {
		msg := "error: Can't get ipaddress..."
		log.Print(msg)
		return "", errors.New(msg)
	}
	log.Printf("Instance IpAddress: %s", instanceIp)
	return instanceIp, nil
}

// create creates a docker container, stores it on the database and adds a route to it.
//
// It receives the related application in order to choose the correct
// docker image and the repository to pass to the script that will take
// care of the deploy, and a function to generate the correct command ran by
// docker, which might be to deploy a container or to run and expose a
// container for an application.
func (c *container) create(app provision.App) error {
	hostAddr, err := config.Get("docker:host-address")
	if err != nil {
		return err
	}
	cmd, err := commandToRun(app)
	if err != nil {
		return err
	}
	port, err := getPort()
	if err != nil {
		return err
	}
	id, err := runCmd(cmd[0], cmd[1:]...)
	if err != nil {
		return err
	}
	id = strings.TrimSpace(id)
	log.Printf("docker id=%s", id)
	c.Id = strings.TrimSpace(id)
	c.Port = port
	ip, err := c.ip()
	if err != nil {
		return err
	}
	c.Ip = ip
	coll := collection()
	defer coll.Database.Session.Close()
	if err := coll.Insert(c); err != nil {
		log.Print(err)
		return err
	}
	r, err := getRouter()
	if err != nil {
		return err
	}
	hostPort, err := c.hostPort()
	if err != nil {
		hostPort = c.Port
	}
	return r.AddRoute(app.GetName(), fmt.Sprintf("http://%s:%s", hostAddr, hostPort))
}

func (c *container) deploy(w io.Writer) error {
	deployCmd, err := config.GetString("docker:deploy-cmd")
	if err != nil {
		return err
	}
	runBin, err := config.GetString("docker:run-cmd:bin")
	if err != nil {
		return err
	}
	runArgs, _ := config.GetString("docker:run-cmd:args")
	appRepo := repository.GetReadOnlyUrl(c.AppName)
	for i := 0; i < 5; i++ {
		err = c.ssh(w, w, deployCmd, appRepo)
		if err == nil {
			break
		}
		log.Printf("SSH to the container %q failed. Will retry.", c.Id)
		time.Sleep(100e6)
	}
	if err != nil {
		return err
	}
	return c.ssh(w, w, runBin, strings.Fields(runArgs)...)
}

// remove removes a docker container.
func (c *container) remove() error {
	docker, err := config.GetString("docker:binary")
	if err != nil {
		return err
	}
	log.Printf("Removing container %s from docker", c.Id)
	out, err := runCmd(docker, "rm", c.Id)
	if err != nil {
		log.Printf("Failed to remove container from docker: %s", err.Error())
		log.Printf("Command output: %s", out)
		return err
	}
	log.Printf("Removing container %s from database", c.Id)
	coll := collection()
	defer coll.Database.Session.Close()
	if err := coll.RemoveId(c.Id); err != nil {
		log.Printf("Failed to remove container from database: %s", err.Error())
		return err
	}
	r, err := getRouter()
	if err != nil {
		log.Printf("Failed to obtain router: %s", err.Error())
		return err
	}
	if err := r.RemoveRoute(c.AppName, c.Ip); err != nil {
		log.Printf("Failed to remove route: %s", err.Error())
		return err
	}
	return nil
}

func (c *container) ssh(stdout, stderr io.Writer, cmd string, args ...string) error {
	user, err := config.GetString("docker:ssh:user")
	if err != nil {
		return err
	}
	sshArgs := []string{c.Ip, "-l", user, "-o", "StrictHostKeyChecking no"}
	if keyFile, err := config.GetString("docker:ssh:private-key"); err == nil {
		sshArgs = append(sshArgs, "-i", keyFile)
	}
	sshArgs = append(sshArgs, "--", cmd)
	sshArgs = append(sshArgs, args...)
	return executor().Execute("ssh", sshArgs, nil, stdout, stderr)
}

// image represents a docker image.
type image struct {
	Name string
	Id   string
}

// repositoryName returns the image repository name for a given image.
//
// Repository is a docker concept, the image actually does not have a name,
// it has a repository, that is a composed name, e.g.: tsuru/base.
// Tsuru will always use a namespace, defined in tsuru.conf.
// Additionally, tsuru will use the application's name to do that composition.
func (img *image) repositoryName() string {
	repoNamespace, err := config.GetString("docker:repository-namespace")
	if err != nil {
		log.Printf("Tsuru is misconfigured. docker:repository-namespace config is missing.")
		return ""
	}
	return fmt.Sprintf("%s/%s", repoNamespace, img.Name)
}

// commit commits an image in docker
//
// This is another docker concept, in order to generate an image from a container
// one must commit it.
func (img *image) commit(cId string) (string, error) {
	docker, err := config.GetString("docker:binary")
	if err != nil {
		log.Printf("Tsuru is misconfigured. docker:binary config is missing.")
		return "", err
	}
	log.Printf("attempting to commit image from container %s", cId)
	rName := img.repositoryName()
	id, err := runCmd(docker, "commit", cId, rName)
	if err != nil {
		log.Printf("Could not commit docker image: %s", err.Error())
		return "", err
	}
	img.Id = strings.Replace(id, "\n", "", -1)
	if err := imagesCollection().Insert(&img); err != nil {
		log.Printf("Could not store image information %s", err.Error())
		return "", err
	}
	return img.Id, nil
}

// remove removes an image from docker registry
func (img *image) remove() error {
	docker, err := config.GetString("docker:binary")
	if err != nil {
		log.Printf("Tsuru is misconfigured. docker:binary config is missing.")
		return err
	}
	log.Printf("attempting to remove image %s from docker", img.repositoryName())
	_, err = runCmd(docker, "rmi", img.Id)
	if err != nil {
		log.Printf("Could not remove image %s from docker: %s", img.Id, err.Error())
		return err
	}
	err = imagesCollection().Remove(bson.M{"name": img.Name})
	if err != nil {
		log.Printf("Could not remove image %s from mongo: %s", img.Id, err.Error())
		return err
	}
	return nil
}

func getContainer(id string) (*container, error) {
	var c container
	coll := collection()
	defer coll.Database.Session.Close()
	err := coll.Find(bson.M{"_id": id}).One(&c)
	if err != nil {
		return nil, err
	}
	return &c, nil
}

func listAppContainers(appName string) ([]container, error) {
	var containers []container
	err := collection().Find(bson.M{"appname": appName}).All(&containers)
	return containers, err
}
<|endoftext|>


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

// Package pprof serves via its HTTP server runtime profiling data
// in the format expected by the pprof visualization tool.
// For more information about pprof, see
// http://code.google.com/p/google-perftools/.
//
// The package is typically only imported for the side effect of
// registering its HTTP handlers.
// The handled paths all begin with /debug/pprof/.
//
// To use pprof, link this package into your program:
//	import _ "http/pprof"
//
// Then use the pprof tool to look at the heap profile:
//
//	pprof http://localhost:6060/debug/pprof/heap
//
package pprof

import (
	"bufio"
	"fmt"
	"http"
	"os"
	"runtime"
	"runtime/pprof"
	"strconv"
	"strings"
)

func init() {
	http.Handle("/debug/pprof/cmdline", http.HandlerFunc(Cmdline))
	http.Handle("/debug/pprof/heap", http.HandlerFunc(Heap))
	http.Handle("/debug/pprof/symbol", http.HandlerFunc(Symbol))
}

// Cmdline responds with the running program's
// command line, with arguments separated by NUL bytes.
// The package initialization registers it as /debug/pprof/cmdline.
func Cmdline(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("content-type", "text/plain; charset=utf-8")
	fmt.Fprintf(w, strings.Join(os.Args, "\x00"))
}

// Heap responds with the pprof-formatted heap profile.
// The package initialization registers it as /debug/pprof/heap.
func Heap(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("content-type", "text/plain; charset=utf-8")
	pprof.WriteHeapProfile(w)
}

// Symbol looks up the program counters listed in the request,
// responding with a table mapping program counters to function names.
// The package initialization registers it as /debug/pprof/symbol.
func Symbol(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("content-type", "text/plain; charset=utf-8")

	// We don't know how many symbols we have, but we
	// do have symbol information.  Pprof only cares whether
	// this number is 0 (no symbols available) or > 0.
	fmt.Fprintf(w, "num_symbols: 1\n")

	var b *bufio.Reader
	if r.Method == "POST" {
		b = bufio.NewReader(r.Body)
	} else {
		b = bufio.NewReader(strings.NewReader(r.URL.RawQuery))
	}

	for {
		word, err := b.ReadSlice('+')
		if err == nil {
			word = word[0 : len(word)-1] // trim +
		}
		pc, _ := strconv.Btoui64(string(word), 0)
		if pc != 0 {
			f := runtime.FuncForPC(uintptr(pc))
			if f != nil {
				fmt.Fprintf(w, "%#x %s\n", pc, f.Name())
			}
		}

		// Wait until here to check for err; the last
		// symbol will have an err because it doesn't end in +.
		if err != nil {
			break
		}
	}
}
<commit_msg>http/pprof: cpu profiling support<commit_after>// Copyright 2010 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package pprof serves via its HTTP server runtime profiling data
// in the format expected by the pprof visualization tool.
// For more information about pprof, see
// http://code.google.com/p/google-perftools/.
//
// The package is typically only imported for the side effect of
// registering its HTTP handlers.
// The handled paths all begin with /debug/pprof/.
//
// To use pprof, link this package into your program:
//	import _ "http/pprof"
//
// Then use the pprof tool to look at the heap profile:
//
//	pprof http://localhost:6060/debug/pprof/heap
//
// Or to look at a 30-second CPU profile:
//
//	pprof http://localhost:6060/debug/pprof/profile
//
package pprof

import (
	"bufio"
	"fmt"
	"http"
	"os"
	"runtime"
	"runtime/pprof"
	"strconv"
	"strings"
	"time"
)

func init() {
	http.Handle("/debug/pprof/cmdline", http.HandlerFunc(Cmdline))
	http.Handle("/debug/pprof/profile", http.HandlerFunc(Profile))
	http.Handle("/debug/pprof/heap", http.HandlerFunc(Heap))
	http.Handle("/debug/pprof/symbol", http.HandlerFunc(Symbol))
}

// Cmdline responds with the running program's
// command line, with arguments separated by NUL bytes.
// The package initialization registers it as /debug/pprof/cmdline.
func Cmdline(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
	fmt.Fprintf(w, strings.Join(os.Args, "\x00"))
}

// Heap responds with the pprof-formatted heap profile.
// The package initialization registers it as /debug/pprof/heap.
func Heap(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
	pprof.WriteHeapProfile(w)
}

// Profile responds with the pprof-formatted cpu profile.
// The package initialization registers it as /debug/pprof/profile.
func Profile(w http.ResponseWriter, r *http.Request) {
	sec, _ := strconv.Atoi64(r.FormValue("seconds"))
	if sec == 0 {
		sec = 30
	}

	// Set Content Type assuming StartCPUProfile will work,
	// because if it does it starts writing.
	w.Header().Set("Content-Type", "application/octet-stream")
	if err := pprof.StartCPUProfile(w); err != nil {
		// StartCPUProfile failed, so no writes yet.
		// Can change header back to text content
		// and send error code.
		w.Header().Set("Content-Type", "text/plain; charset=utf-8")
		w.WriteHeader(http.StatusInternalServerError)
		fmt.Fprintf(w, "Could not enable CPU profiling: %s\n", err)
		return
	}
	time.Sleep(sec * 1e9)
	pprof.StopCPUProfile()
}

// Symbol looks up the program counters listed in the request,
// responding with a table mapping program counters to function names.
// The package initialization registers it as /debug/pprof/symbol.
func Symbol(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/plain; charset=utf-8")

	// We don't know how many symbols we have, but we
	// do have symbol information.  Pprof only cares whether
	// this number is 0 (no symbols available) or > 0.
	fmt.Fprintf(w, "num_symbols: 1\n")

	var b *bufio.Reader
	if r.Method == "POST" {
		b = bufio.NewReader(r.Body)
	} else {
		b = bufio.NewReader(strings.NewReader(r.URL.RawQuery))
	}

	for {
		word, err := b.ReadSlice('+')
		if err == nil {
			word = word[0 : len(word)-1] // trim +
		}
		pc, _ := strconv.Btoui64(string(word), 0)
		if pc != 0 {
			f := runtime.FuncForPC(uintptr(pc))
			if f != nil {
				fmt.Fprintf(w, "%#x %s\n", pc, f.Name())
			}
		}

		// Wait until here to check for err; the last
		// symbol will have an err because it doesn't end in +.
		if err != nil {
			break
		}
	}
}
<|endoftext|>


<commit_before>package src
/** qr(m *Matrix, _) returns the QR Decomposition of a matrix m.
    * setting _ to 'q' returns Q, while setting _ to 'r' returns R
*/

func qr(A *Matrix, tag string) (*Matrix, string) {
    if (tag != "q" && tag != "r" && tag != "\"q\"" && tag != "'q'" && tag != "\"r\"" && tag != "'r'") {
        return nil, "ERROR: Invalid tag. Please use either q or r";
    }

    Q, _ := gs(A);
    if (tag == "q") {
        return Q, "";
    } else {
        // Create the transposed R matrix
        Rvals := make([][]float64, Q.N);
        for i := 0; i < len(Rvals); i++ {
            column := make([]float64, Q.N);
            for j := 0; j <= i; j++ {
                column[j] = dotCols(A.cols[i], Q.cols[j]);
            }
            Rvals[i] = column;
        }
        Rtrans, _ := NewMatrix(Rvals);
        return transpose(Rtrans);
    }
}

func dotCols(a []float64, q []float64) float64 {
    result := 0.0;
    for i := 0; i < len(a); i++ {
        result += a[i] * q[i];
    }
    return result;
}
<commit_msg>Round R matrix in qr.go<commit_after>package src
/** qr(m *Matrix, _) returns the QR Decomposition of a matrix m.
    * setting _ to 'q' returns Q, while setting _ to 'r' returns R
*/

func qr(A *Matrix, tag string) (*Matrix, string) {
    if (tag != "q" && tag != "r" && tag != "\"q\"" && tag != "'q'" && tag != "\"r\"" && tag != "'r'") {
        return nil, "ERROR: Invalid tag. Please use either q or r";
    }

    Q, _ := gs(A);
    if (tag == "q") {
        return Q, "";
    } else {
        // Create the transposed R matrix
        Rvals := make([][]float64, Q.N);
        for i := 0; i < len(Rvals); i++ {
            column := make([]float64, Q.N);
            for j := 0; j <= i; j++ {
                column[j] = dotCols(A.cols[i], Q.cols[j]);
            }
            Rvals[i] = column;
        }
        Rtrans, _ := NewMatrix(Rvals);
        round(Rtrans);
        return transpose(Rtrans);
    }
}

func dotCols(a []float64, q []float64) float64 {
    result := 0.0;
    for i := 0; i < len(a); i++ {
        result += a[i] * q[i];
    }
    return result;
}
<|endoftext|>


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

package commands

import (
	"fmt"
	"os"
	"strconv"
	"strings"
	"text/tabwriter"

	log "github.com/sirupsen/logrus"

	"github.com/spf13/cobra"
	"k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/client-go/kubernetes"
)

var (
	showDetails bool
)

type NodeInfo struct {
	node v1.Node
	pods []v1.Pod
}

func NewTopNodeCommand() *cobra.Command {

	var command = &cobra.Command{
		Use:   "node",
		Short: "Display Resource (GPU) usage of nodes.",
		Run: func(cmd *cobra.Command, args []string) {
			setupKubeconfig()
			client, err := initKubeClient()
			if err != nil {
				fmt.Println(err)
				os.Exit(1)
			}
			allPods, err = acquireAllActivePods(client)
			if err != nil {
				fmt.Println(err)
				os.Exit(1)
			}
			nd := newNodeDescriber(client, allPods)
			nodeInfos, err := nd.getAllNodeInfos()
			if err != nil {
				fmt.Println(err)
				os.Exit(1)
			}
			displayTopNode(nodeInfos)
		},
	}

	command.Flags().BoolVarP(&showDetails, "details", "d", false, "Display details")
	return command
}

type NodeDescriber struct {
	client  *kubernetes.Clientset
	allPods []v1.Pod
}

func newNodeDescriber(client *kubernetes.Clientset, pods []v1.Pod) *NodeDescriber {
	return &NodeDescriber{
		client:  client,
		allPods: pods,
	}
}

func (d *NodeDescriber) getAllNodeInfos() ([]NodeInfo, error) {
	nodeInfoList := []NodeInfo{}
	nodeList, err := d.client.CoreV1().Nodes().List(metav1.ListOptions{})

	if err != nil {
		return nodeInfoList, err
	}

	for _, node := range nodeList.Items {
		pods := d.getPodsFromNode(node)
		nodeInfoList = append(nodeInfoList, NodeInfo{
			node: node,
			pods: pods,
		})
	}

	return nodeInfoList, nil
}

func (d *NodeDescriber) getPodsFromNode(node v1.Node) []v1.Pod {
	pods := []v1.Pod{}
	for _, pod := range d.allPods {
		if pod.Spec.NodeName == node.Name {
			pods = append(pods, pod)
		}
	}

	return pods
}

func displayTopNode(nodes []NodeInfo) {
	if showDetails {
		displayTopNodeDetails(nodes)
	} else {
		displayTopNodeSummary(nodes)
	}
}

func displayTopNodeSummary(nodeInfos []NodeInfo) {
	w := tabwriter.NewWriter(os.Stdout, 0, 0, 2, ' ', 0)
	var (
		totalGPUsInCluster     int64
		allocatedGPUsInCluster int64
	)

	fmt.Fprintf(w, "NAME\tIPADDRESS\tROLE\tGPU(Total)\tGPU(Allocated)\n")
	for _, nodeInfo := range nodeInfos {
		totalGPU, allocatedGPU := calculateNodeGPU(nodeInfo)
		totalGPUsInCluster += totalGPU
		allocatedGPUsInCluster += allocatedGPU

		address := "unknown"
		if len(nodeInfo.node.Status.Addresses) > 0 {
			// address = nodeInfo.node.Status.Addresses[0].Address
			for _, addr := range nodeInfo.node.Status.Addresses {
				if addr.Type == v1.NodeInternalIP {
					address = addr.Address
					break
				}
			}
		}

		role := strings.Join(findNodeRoles(&nodeInfo.node), ",")
		if len(role) == 0 {
			role = "<none>"
		}

		fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\n", nodeInfo.node.Name,
			address,
			role,
			strconv.FormatInt(totalGPU, 10),
			strconv.FormatInt(allocatedGPU, 10))
	}
	fmt.Fprintf(w, "-----------------------------------------------------------------------------------------\n")
	fmt.Fprintf(w, "Allocated/Total GPUs In Cluster:\n")
	log.Debugf("gpu: %s, allocated GPUs %s", strconv.FormatInt(totalGPUsInCluster, 10),
		strconv.FormatInt(allocatedGPUsInCluster, 10))
	var gpuUsage float64 = 0
	if totalGPUsInCluster > 0 {
		gpuUsage = float64(allocatedGPUsInCluster) / float64(totalGPUsInCluster) * 100
	}
	fmt.Fprintf(w, "%s/%s (%d%%)\t\n",
		strconv.FormatInt(allocatedGPUsInCluster, 10),
		strconv.FormatInt(totalGPUsInCluster, 10),
		int64(gpuUsage))
	// fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\n", ...)

	_ = w.Flush()
}

func displayTopNodeDetails(nodeInfos []NodeInfo) {
	w := tabwriter.NewWriter(os.Stdout, 0, 0, 2, ' ', 0)
	var (
		totalGPUsInCluster     int64
		allocatedGPUsInCluster int64
	)

	fmt.Fprintf(w, "\n")
	for _, nodeInfo := range nodeInfos {
		totalGPU, allocatedGPU := calculateNodeGPU(nodeInfo)
		totalGPUsInCluster += totalGPU
		allocatedGPUsInCluster += allocatedGPU

		address := "unknown"
		if len(nodeInfo.node.Status.Addresses) > 0 {
			address = nodeInfo.node.Status.Addresses[0].Address
		}

		role := strings.Join(findNodeRoles(&nodeInfo.node), ",")
		if len(role) == 0 {
			role = "<none>"
		}

		fmt.Fprintf(w, "\n")
		fmt.Fprintf(w, "NAME:\t%s\n", nodeInfo.node.Name)
		fmt.Fprintf(w, "IPADDRESS:\t%s\n", address)
		fmt.Fprintf(w, "ROLE:\t%s\n", role)

		pods := gpuPods(nodeInfo.pods)
		if len(pods) > 0 {
			fmt.Fprintf(w, "\n")
			fmt.Fprintf(w, "NAMESPACE\tNAME\tGPU REQUESTS\tGPU LIMITS\n")
			for _, pod := range pods {
				fmt.Fprintf(w, "%s\t%s\t%s\t%s\n", pod.Namespace,
					pod.Name,
					strconv.FormatInt(gpuInPod(pod), 10),
					strconv.FormatInt(gpuInPod(pod), 10))
			}
			fmt.Fprintf(w, "\n")
		}
		var gpuUsageInNode float64 = 0
		if totalGPU > 0 {
			gpuUsageInNode = float64(allocatedGPU) / float64(totalGPU) * 100
		} else {
			fmt.Fprintf(w, "\n")
		}

		fmt.Fprintf(w, "Total GPUs In Node %s:\t%s \t\n", nodeInfo.node.Name, strconv.FormatInt(totalGPU, 10))
		fmt.Fprintf(w, "Allocated GPUs In Node %s:\t%s (%d%%)\t\n", nodeInfo.node.Name, strconv.FormatInt(allocatedGPU, 10), int64(gpuUsageInNode))
		log.Debugf("gpu: %s, allocated GPUs %s", strconv.FormatInt(totalGPU, 10),
			strconv.FormatInt(allocatedGPU, 10))

		fmt.Fprintf(w, "-----------------------------------------------------------------------------------------\n")
	}
	fmt.Fprintf(w, "\n")
	fmt.Fprintf(w, "\n")
	fmt.Fprintf(w, "Allocated/Total GPUs In Cluster:\t")
	log.Debugf("gpu: %s, allocated GPUs %s", strconv.FormatInt(totalGPUsInCluster, 10),
		strconv.FormatInt(allocatedGPUsInCluster, 10))

	var gpuUsage float64 = 0
	if totalGPUsInCluster > 0 {
		gpuUsage = float64(allocatedGPUsInCluster) / float64(totalGPUsInCluster) * 100
	}
	fmt.Fprintf(w, "%s/%s (%d%%)\t\n",
		strconv.FormatInt(allocatedGPUsInCluster, 10),
		strconv.FormatInt(totalGPUsInCluster, 10),
		int64(gpuUsage))
	// fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\n", ...)

	_ = w.Flush()
}

// calculate the GPU count of each node
func calculateNodeGPU(nodeInfo NodeInfo) (totalGPU int64, allocatedGPU int64) {
	node := nodeInfo.node
	totalGPU = gpuInNode(node)
	// allocatedGPU = gpuInPod()

	for _, pod := range nodeInfo.pods {
		allocatedGPU += gpuInPod(pod)
	}

	return totalGPU, allocatedGPU
}

func isMasterNode(node v1.Node) bool {
	if _, ok := node.Labels[masterLabelRole]; ok {
		return true
	}

	return false
}

// findNodeRoles returns the roles of a given node.
// The roles are determined by looking for:
// * a node-role.kubernetes.io/<role>="" label
// * a kubernetes.io/role="<role>" label
func findNodeRoles(node *v1.Node) []string {
	roles := sets.NewString()
	for k, v := range node.Labels {
		switch {
		case strings.HasPrefix(k, labelNodeRolePrefix):
			if role := strings.TrimPrefix(k, labelNodeRolePrefix); len(role) > 0 {
				roles.Insert(role)
			}

		case k == nodeLabelRole && v != "":
			roles.Insert(v)
		}
	}
	return roles.List()
}
<commit_msg>fix details display of node ip address (#130)<commit_after>// Copyright 2018 The Kubeflow Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//       http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package commands

import (
	"fmt"
	"os"
	"strconv"
	"strings"
	"text/tabwriter"

	log "github.com/sirupsen/logrus"

	"github.com/spf13/cobra"
	"k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/client-go/kubernetes"
)

var (
	showDetails bool
)

type NodeInfo struct {
	node v1.Node
	pods []v1.Pod
}

func NewTopNodeCommand() *cobra.Command {

	var command = &cobra.Command{
		Use:   "node",
		Short: "Display Resource (GPU) usage of nodes.",
		Run: func(cmd *cobra.Command, args []string) {
			setupKubeconfig()
			client, err := initKubeClient()
			if err != nil {
				fmt.Println(err)
				os.Exit(1)
			}
			allPods, err = acquireAllActivePods(client)
			if err != nil {
				fmt.Println(err)
				os.Exit(1)
			}
			nd := newNodeDescriber(client, allPods)
			nodeInfos, err := nd.getAllNodeInfos()
			if err != nil {
				fmt.Println(err)
				os.Exit(1)
			}
			displayTopNode(nodeInfos)
		},
	}

	command.Flags().BoolVarP(&showDetails, "details", "d", false, "Display details")
	return command
}

type NodeDescriber struct {
	client  *kubernetes.Clientset
	allPods []v1.Pod
}

func newNodeDescriber(client *kubernetes.Clientset, pods []v1.Pod) *NodeDescriber {
	return &NodeDescriber{
		client:  client,
		allPods: pods,
	}
}

func (d *NodeDescriber) getAllNodeInfos() ([]NodeInfo, error) {
	nodeInfoList := []NodeInfo{}
	nodeList, err := d.client.CoreV1().Nodes().List(metav1.ListOptions{})

	if err != nil {
		return nodeInfoList, err
	}

	for _, node := range nodeList.Items {
		pods := d.getPodsFromNode(node)
		nodeInfoList = append(nodeInfoList, NodeInfo{
			node: node,
			pods: pods,
		})
	}

	return nodeInfoList, nil
}

func (d *NodeDescriber) getPodsFromNode(node v1.Node) []v1.Pod {
	pods := []v1.Pod{}
	for _, pod := range d.allPods {
		if pod.Spec.NodeName == node.Name {
			pods = append(pods, pod)
		}
	}

	return pods
}

func displayTopNode(nodes []NodeInfo) {
	if showDetails {
		displayTopNodeDetails(nodes)
	} else {
		displayTopNodeSummary(nodes)
	}
}

func displayTopNodeSummary(nodeInfos []NodeInfo) {
	w := tabwriter.NewWriter(os.Stdout, 0, 0, 2, ' ', 0)
	var (
		totalGPUsInCluster     int64
		allocatedGPUsInCluster int64
	)

	fmt.Fprintf(w, "NAME\tIPADDRESS\tROLE\tGPU(Total)\tGPU(Allocated)\n")
	for _, nodeInfo := range nodeInfos {
		totalGPU, allocatedGPU := calculateNodeGPU(nodeInfo)
		totalGPUsInCluster += totalGPU
		allocatedGPUsInCluster += allocatedGPU

		address := "unknown"
		if len(nodeInfo.node.Status.Addresses) > 0 {
			// address = nodeInfo.node.Status.Addresses[0].Address
			for _, addr := range nodeInfo.node.Status.Addresses {
				if addr.Type == v1.NodeInternalIP {
					address = addr.Address
					break
				}
			}
		}

		role := strings.Join(findNodeRoles(&nodeInfo.node), ",")
		if len(role) == 0 {
			role = "<none>"
		}

		fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\n", nodeInfo.node.Name,
			address,
			role,
			strconv.FormatInt(totalGPU, 10),
			strconv.FormatInt(allocatedGPU, 10))
	}
	fmt.Fprintf(w, "-----------------------------------------------------------------------------------------\n")
	fmt.Fprintf(w, "Allocated/Total GPUs In Cluster:\n")
	log.Debugf("gpu: %s, allocated GPUs %s", strconv.FormatInt(totalGPUsInCluster, 10),
		strconv.FormatInt(allocatedGPUsInCluster, 10))
	var gpuUsage float64 = 0
	if totalGPUsInCluster > 0 {
		gpuUsage = float64(allocatedGPUsInCluster) / float64(totalGPUsInCluster) * 100
	}
	fmt.Fprintf(w, "%s/%s (%d%%)\t\n",
		strconv.FormatInt(allocatedGPUsInCluster, 10),
		strconv.FormatInt(totalGPUsInCluster, 10),
		int64(gpuUsage))
	// fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\n", ...)

	_ = w.Flush()
}

func displayTopNodeDetails(nodeInfos []NodeInfo) {
	w := tabwriter.NewWriter(os.Stdout, 0, 0, 2, ' ', 0)
	var (
		totalGPUsInCluster     int64
		allocatedGPUsInCluster int64
	)

	fmt.Fprintf(w, "\n")
	for _, nodeInfo := range nodeInfos {
		totalGPU, allocatedGPU := calculateNodeGPU(nodeInfo)
		totalGPUsInCluster += totalGPU
		allocatedGPUsInCluster += allocatedGPU

		address := "unknown"
		if len(nodeInfo.node.Status.Addresses) > 0 {
			//address = nodeInfo.node.Status.Addresses[0].Address
			for _, addr := range nodeInfo.node.Status.Addresses {
				if addr.Type == v1.NodeInternalIP {
					address = addr.Address
				}
			}
		}

		role := strings.Join(findNodeRoles(&nodeInfo.node), ",")
		if len(role) == 0 {
			role = "<none>"
		}

		fmt.Fprintf(w, "\n")
		fmt.Fprintf(w, "NAME:\t%s\n", nodeInfo.node.Name)
		fmt.Fprintf(w, "IPADDRESS:\t%s\n", address)
		fmt.Fprintf(w, "ROLE:\t%s\n", role)

		pods := gpuPods(nodeInfo.pods)
		if len(pods) > 0 {
			fmt.Fprintf(w, "\n")
			fmt.Fprintf(w, "NAMESPACE\tNAME\tGPU REQUESTS\tGPU LIMITS\n")
			for _, pod := range pods {
				fmt.Fprintf(w, "%s\t%s\t%s\t%s\n", pod.Namespace,
					pod.Name,
					strconv.FormatInt(gpuInPod(pod), 10),
					strconv.FormatInt(gpuInPod(pod), 10))
			}
			fmt.Fprintf(w, "\n")
		}
		var gpuUsageInNode float64 = 0
		if totalGPU > 0 {
			gpuUsageInNode = float64(allocatedGPU) / float64(totalGPU) * 100
		} else {
			fmt.Fprintf(w, "\n")
		}

		fmt.Fprintf(w, "Total GPUs In Node %s:\t%s \t\n", nodeInfo.node.Name, strconv.FormatInt(totalGPU, 10))
		fmt.Fprintf(w, "Allocated GPUs In Node %s:\t%s (%d%%)\t\n", nodeInfo.node.Name, strconv.FormatInt(allocatedGPU, 10), int64(gpuUsageInNode))
		log.Debugf("gpu: %s, allocated GPUs %s", strconv.FormatInt(totalGPU, 10),
			strconv.FormatInt(allocatedGPU, 10))

		fmt.Fprintf(w, "-----------------------------------------------------------------------------------------\n")
	}
	fmt.Fprintf(w, "\n")
	fmt.Fprintf(w, "\n")
	fmt.Fprintf(w, "Allocated/Total GPUs In Cluster:\t")
	log.Debugf("gpu: %s, allocated GPUs %s", strconv.FormatInt(totalGPUsInCluster, 10),
		strconv.FormatInt(allocatedGPUsInCluster, 10))

	var gpuUsage float64 = 0
	if totalGPUsInCluster > 0 {
		gpuUsage = float64(allocatedGPUsInCluster) / float64(totalGPUsInCluster) * 100
	}
	fmt.Fprintf(w, "%s/%s (%d%%)\t\n",
		strconv.FormatInt(allocatedGPUsInCluster, 10),
		strconv.FormatInt(totalGPUsInCluster, 10),
		int64(gpuUsage))
	// fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\n", ...)

	_ = w.Flush()
}

// calculate the GPU count of each node
func calculateNodeGPU(nodeInfo NodeInfo) (totalGPU int64, allocatedGPU int64) {
	node := nodeInfo.node
	totalGPU = gpuInNode(node)
	// allocatedGPU = gpuInPod()

	for _, pod := range nodeInfo.pods {
		allocatedGPU += gpuInPod(pod)
	}

	return totalGPU, allocatedGPU
}

func isMasterNode(node v1.Node) bool {
	if _, ok := node.Labels[masterLabelRole]; ok {
		return true
	}

	return false
}

// findNodeRoles returns the roles of a given node.
// The roles are determined by looking for:
// * a node-role.kubernetes.io/<role>="" label
// * a kubernetes.io/role="<role>" label
func findNodeRoles(node *v1.Node) []string {
	roles := sets.NewString()
	for k, v := range node.Labels {
		switch {
		case strings.HasPrefix(k, labelNodeRolePrefix):
			if role := strings.TrimPrefix(k, labelNodeRolePrefix); len(role) > 0 {
				roles.Insert(role)
			}

		case k == nodeLabelRole && v != "":
			roles.Insert(v)
		}
	}
	return roles.List()
}
<|endoftext|>


<commit_before>/*
   Copyright The containerd Authors.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package tasks

import (
	"errors"

	"github.com/containerd/console"
	"github.com/containerd/containerd/cio"
	"github.com/containerd/containerd/cmd/ctr/commands"
	"github.com/sirupsen/logrus"
	"github.com/urfave/cli"
)

//TODO:(jessvalarezo) exec-id is optional here, update to required arg
var execCommand = cli.Command{
	Name:           "exec",
	Usage:          "execute additional processes in an existing container",
	ArgsUsage:      "[flags] CONTAINER CMD [ARG...]",
	SkipArgReorder: true,
	Flags: []cli.Flag{
		cli.StringFlag{
			Name:  "cwd",
			Usage: "working directory of the new process",
		},
		cli.BoolFlag{
			Name:  "tty,t",
			Usage: "allocate a TTY for the container",
		},
		cli.BoolFlag{
			Name:  "detach,d",
			Usage: "detach from the task after it has started execution",
		},
		cli.StringFlag{
			Name:  "exec-id",
			Usage: "exec specific id for the process",
		},
		cli.StringFlag{
			Name:  "fifo-dir",
			Usage: "directory used for storing IO FIFOs",
		},
	},
	Action: func(context *cli.Context) error {
		var (
			id     = context.Args().First()
			args   = context.Args().Tail()
			tty    = context.Bool("tty")
			detach = context.Bool("detach")
		)
		if id == "" {
			return errors.New("container id must be provided")
		}
		client, ctx, cancel, err := commands.NewClient(context)
		if err != nil {
			return err
		}
		defer cancel()
		container, err := client.LoadContainer(ctx, id)
		if err != nil {
			return err
		}
		spec, err := container.Spec(ctx)
		if err != nil {
			return err
		}
		task, err := container.Task(ctx, nil)
		if err != nil {
			return err
		}

		pspec := spec.Process
		pspec.Terminal = tty
		pspec.Args = args

		cioOpts := []cio.Opt{cio.WithStdio, cio.WithFIFODir(context.String("fifo-dir"))}
		if tty {
			cioOpts = append(cioOpts, cio.WithTerminal)
		}
		ioCreator := cio.NewCreator(cioOpts...)
		process, err := task.Exec(ctx, context.String("exec-id"), pspec, ioCreator)
		if err != nil {
			return err
		}
		// if detach, we should not call this defer
		if !detach {
			defer process.Delete(ctx)
		}

		statusC, err := process.Wait(ctx)
		if err != nil {
			return err
		}

		var con console.Console
		if tty {
			con = console.Current()
			defer con.Reset()
			if err := con.SetRaw(); err != nil {
				return err
			}
		}
		if !detach {
			if tty {
				if err := HandleConsoleResize(ctx, process, con); err != nil {
					logrus.WithError(err).Error("console resize")
				}
			} else {
				sigc := commands.ForwardAllSignals(ctx, process)
				defer commands.StopCatch(sigc)
			}
		}

		if err := process.Start(ctx); err != nil {
			return err
		}
		if detach {
			return nil
		}
		status := <-statusC
		code, _, err := status.Result()
		if err != nil {
			return err
		}
		if code != 0 {
			return cli.NewExitError("", int(code))
		}
		return nil
	},
}
<commit_msg>Support --log-uri for exec subcommand<commit_after>/*
   Copyright The containerd Authors.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package tasks

import (
	"errors"
	"net/url"

	"github.com/containerd/console"
	"github.com/containerd/containerd/cio"
	"github.com/containerd/containerd/cmd/ctr/commands"
	"github.com/sirupsen/logrus"
	"github.com/urfave/cli"
)

//TODO:(jessvalarezo) exec-id is optional here, update to required arg
var execCommand = cli.Command{
	Name:           "exec",
	Usage:          "execute additional processes in an existing container",
	ArgsUsage:      "[flags] CONTAINER CMD [ARG...]",
	SkipArgReorder: true,
	Flags: []cli.Flag{
		cli.StringFlag{
			Name:  "cwd",
			Usage: "working directory of the new process",
		},
		cli.BoolFlag{
			Name:  "tty,t",
			Usage: "allocate a TTY for the container",
		},
		cli.BoolFlag{
			Name:  "detach,d",
			Usage: "detach from the task after it has started execution",
		},
		cli.StringFlag{
			Name:  "exec-id",
			Usage: "exec specific id for the process",
		},
		cli.StringFlag{
			Name:  "fifo-dir",
			Usage: "directory used for storing IO FIFOs",
		},
		cli.StringFlag{
			Name:  "log-uri",
			Usage: "log uri for custom shim logging",
		},
	},
	Action: func(context *cli.Context) error {
		var (
			id     = context.Args().First()
			args   = context.Args().Tail()
			tty    = context.Bool("tty")
			detach = context.Bool("detach")
		)
		if id == "" {
			return errors.New("container id must be provided")
		}
		client, ctx, cancel, err := commands.NewClient(context)
		if err != nil {
			return err
		}
		defer cancel()
		container, err := client.LoadContainer(ctx, id)
		if err != nil {
			return err
		}
		spec, err := container.Spec(ctx)
		if err != nil {
			return err
		}
		task, err := container.Task(ctx, nil)
		if err != nil {
			return err
		}

		pspec := spec.Process
		pspec.Terminal = tty
		pspec.Args = args

		var ioCreator cio.Creator

		if logURI := context.String("log-uri"); logURI != "" {
			uri, err := url.Parse(logURI)
			if err != nil {
				return err
			}
			ioCreator = cio.LogURI(uri)
		} else {
			cioOpts := []cio.Opt{cio.WithStdio, cio.WithFIFODir(context.String("fifo-dir"))}
			if tty {
				cioOpts = append(cioOpts, cio.WithTerminal)
			}
			ioCreator = cio.NewCreator(cioOpts...)
		}

		process, err := task.Exec(ctx, context.String("exec-id"), pspec, ioCreator)
		if err != nil {
			return err
		}
		// if detach, we should not call this defer
		if !detach {
			defer process.Delete(ctx)
		}

		statusC, err := process.Wait(ctx)
		if err != nil {
			return err
		}

		var con console.Console
		if tty {
			con = console.Current()
			defer con.Reset()
			if err := con.SetRaw(); err != nil {
				return err
			}
		}
		if !detach {
			if tty {
				if err := HandleConsoleResize(ctx, process, con); err != nil {
					logrus.WithError(err).Error("console resize")
				}
			} else {
				sigc := commands.ForwardAllSignals(ctx, process)
				defer commands.StopCatch(sigc)
			}
		}

		if err := process.Start(ctx); err != nil {
			return err
		}
		if detach {
			return nil
		}
		status := <-statusC
		code, _, err := status.Result()
		if err != nil {
			return err
		}
		if code != 0 {
			return cli.NewExitError("", int(code))
		}
		return nil
	},
}
<|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.
*/
// A basic integration test for the service.
// Assumes that there is a pre-existing etcd server running on localhost.
package main

import (
	"encoding/json"
	"io/ioutil"
	"log"
	"net/http/httptest"
	"time"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/apiserver"
	kube_client "github.com/GoogleCloudPlatform/kubernetes/pkg/client"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/registry"
	"github.com/coreos/go-etcd/etcd"
)

func main() {

	// Setup
	servers := []string{"http://localhost:4001"}
	log.Printf("Creating etcd client pointing to %v", servers)
	etcdClient := etcd.NewClient(servers)
	machineList := []string{"machine"}

	reg := registry.MakeEtcdRegistry(etcdClient, machineList)

	apiserver := apiserver.New(map[string]apiserver.RESTStorage{
		"pods":                  registry.MakePodRegistryStorage(reg, &kube_client.FakeContainerInfo{}, registry.MakeRoundRobinScheduler(machineList)),
		"replicationControllers": registry.MakeControllerRegistryStorage(reg),
	}, "/api/v1beta1")
	server := httptest.NewServer(apiserver)

	controllerManager := registry.MakeReplicationManager(etcd.NewClient(servers),
		kube_client.Client{
			Host: server.URL,
		})

	go controllerManager.Synchronize()
	go controllerManager.WatchControllers()

	// Ok. we're good to go.
	log.Printf("API Server started on %s", server.URL)
	// Wait for the synchronization threads to come up.
	time.Sleep(time.Second * 10)

	kubeClient := kube_client.Client{
		Host: server.URL,
	}
	data, err := ioutil.ReadFile("api/examples/controller.json")
	if err != nil {
		log.Fatalf("Unexpected error: %#v", err)
	}
	var controllerRequest api.ReplicationController
	if err = json.Unmarshal(data, &controllerRequest); err != nil {
		log.Fatalf("Unexpected error: %#v", err)
	}

	if _, err = kubeClient.CreateReplicationController(controllerRequest); err != nil {
		log.Fatalf("Unexpected error: %#v", err)
	}
	// Give the controllers some time to actually create the pods
	time.Sleep(time.Second * 10)

	// Validate that they're truly up.
	pods, err := kubeClient.ListPods(nil)
	if err != nil || len(pods.Items) != 2 {
		log.Fatal("FAILED")
	}
	log.Printf("OK")
}
<commit_msg>remove package name, not needed here<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.
*/

// A basic integration test for the service.
// Assumes that there is a pre-existing etcd server running on localhost.
package main

import (
	"encoding/json"
	"io/ioutil"
	"log"
	"net/http/httptest"
	"time"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/apiserver"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/client"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/registry"
	"github.com/coreos/go-etcd/etcd"
)

func main() {

	// Setup
	servers := []string{"http://localhost:4001"}
	log.Printf("Creating etcd client pointing to %v", servers)
	etcdClient := etcd.NewClient(servers)
	machineList := []string{"machine"}

	reg := registry.MakeEtcdRegistry(etcdClient, machineList)

	apiserver := apiserver.New(map[string]apiserver.RESTStorage{
		"pods": registry.MakePodRegistryStorage(reg, &client.FakeContainerInfo{}, registry.MakeRoundRobinScheduler(machineList)),
		"replicationControllers": registry.MakeControllerRegistryStorage(reg),
	}, "/api/v1beta1")
	server := httptest.NewServer(apiserver)

	controllerManager := registry.MakeReplicationManager(etcd.NewClient(servers),
		client.Client{
			Host: server.URL,
		})

	go controllerManager.Synchronize()
	go controllerManager.WatchControllers()

	// Ok. we're good to go.
	log.Printf("API Server started on %s", server.URL)
	// Wait for the synchronization threads to come up.
	time.Sleep(time.Second * 10)

	kubeClient := client.Client{
		Host: server.URL,
	}
	data, err := ioutil.ReadFile("api/examples/controller.json")
	if err != nil {
		log.Fatalf("Unexpected error: %#v", err)
	}
	var controllerRequest api.ReplicationController
	if err = json.Unmarshal(data, &controllerRequest); err != nil {
		log.Fatalf("Unexpected error: %#v", err)
	}

	if _, err = kubeClient.CreateReplicationController(controllerRequest); err != nil {
		log.Fatalf("Unexpected error: %#v", err)
	}
	// Give the controllers some time to actually create the pods
	time.Sleep(time.Second * 10)

	// Validate that they're truly up.
	pods, err := kubeClient.ListPods(nil)
	if err != nil || len(pods.Items) != 2 {
		log.Fatal("FAILED")
	}
	log.Printf("OK")
}
<|endoftext|>


<commit_before>package main

import (
	"time"

	. "launchpad.net/gocheck"
	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/environs/dummy"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/testing"
	"strings"
)

type ProvisioningSuite struct {
	logging testing.LoggingSuite
	zkSuite
	st *state.State
	op <-chan dummy.Operation
}

var _ = Suite(&ProvisioningSuite{})

var veryShortAttempt = environs.AttemptStrategy{
	Total: 1 * time.Second,
	Delay: 80 * time.Millisecond,
}

func (s *ProvisioningSuite) SetUpTest(c *C) {
	s.logging.SetUpTest(c)
	dummy.Reset()

	// Create the operations channel with more than enough space
	// for those tests that don't listen on it.
	op := make(chan dummy.Operation, 20)
	dummy.Listen(op)
	s.op = op

	env, err := environs.NewEnviron(map[string]interface{}{
		"type":      "dummy",
		"zookeeper": true,
		"name":      "testing",
	})
	c.Assert(err, IsNil)
	err = env.Bootstrap(false)
	c.Assert(err, IsNil)

	s.zkSuite.SetUpTest(c)
	s.st, err = state.Open(s.zkInfo)
	c.Assert(err, IsNil)

	// Make sure that zkInfo holds exactly the info we're passing to
	// NewProvisioner,
	s.zkInfo, err = env.StateInfo()
	c.Assert(err, IsNil)
}

func (s *ProvisioningSuite) TearDownTest(c *C) {
	s.zkSuite.TearDownTest()
	s.logging.TearDownTest(c)
}

// invalidateEnvironment alters the environment configuration
// so the ConfigNode returned from the watcher will not pass
// validation.
func (s *ProvisioningSuite) invalidateEnvironment() error {
	env, err := s.st.EnvironConfig()
	if err != nil {
		return err
	}
	env.Set("name", 1)
	_, err = env.Write()
	return err
}

// fixEnvironment undoes the work of invalidateEnvironment.
func (s *ProvisioningSuite) fixEnvironment() error {
	env, err := s.st.EnvironConfig()
	if err != nil {
		return err
	}
	env.Set("name", "testing")
	_, err = env.Write()
	return err
}

func (s *ProvisioningSuite) stopProvisioner(c *C, p *Provisioner) {
	c.Assert(p.Stop(), IsNil)
}

// checkStartInstance checks that an instance has been started
// with a machine id the same as m's, and that the machine's
// instance id has been set appropriately.
func (s *ProvisioningSuite) checkStartInstance(c *C, m *state.Machine) {
	for {
		select {
		case o := <-s.op:
			switch o := o.(type) {
			case dummy.OpStartInstance:
				c.Assert(o.Info, DeepEquals, s.zkInfo)
				c.Assert(o.MachineId, Equals, m.Id())
				c.Assert(o.Instance, NotNil)
				s.checkMachineId(c, m, o.Instance)
				return
			default:
				c.Logf("ignoring unexpected operation %#v", o)
				// ignore
			}
		case <-time.After(2 * time.Second):
			c.Errorf("provisioner did not start an instance")
			return
		}
	}
}

// checkNotStartInstance checks that an instance was not started
func (s *ProvisioningSuite) checkNotStartInstance(c *C) {
	for {
		select {
		case o := <-s.op:
			switch o.(type) {
			case dummy.OpStartInstance:
				c.Errorf("instance started: %v", o)
				return
			default:
				// ignore
			}
		case <-time.After(200 * time.Millisecond):
			return
		}
	}
}

// checkStopInstance checks that an instance has been stopped.
func (s *ProvisioningSuite) checkStopInstance(c *C) {
	// use the non fatal variants to avoid leaking provisioners.
	for {
		select {
		case o := <-s.op:
			switch o.(type) {
			case dummy.OpStopInstances:
				return
			default:
				//ignore
			}
		case <-time.After(2 * time.Second):
			c.Errorf("provisioner did not stop an instance")
			return
		}
	}
}

// checkMachineIdSet checks that the machine has an instance id
// that matches that of the given instance. If the instance is nil,
// It checks that the instance id is unset.
func (s *ProvisioningSuite) checkMachineId(c *C, m *state.Machine, inst environs.Instance) {
	// TODO(dfc) add machine.WatchConfig() to avoid having to poll.
	instId := ""
	if inst != nil {
		instId = inst.Id()
	}
	for a := veryShortAttempt.Start(); a.Next(); {
		id, err := m.InstanceId()
		c.Assert(err, IsNil)
		if id == instId {
			return
		}
	}
	id, err := m.InstanceId()
	c.Assert(err, IsNil)
	c.Assert(id, Equals, instId)
}

func (s *ProvisioningSuite) TestParseSuccess(c *C) {
	create := func() (cmd.Command, *AgentConf) {
		a := &ProvisioningAgent{}
		return a, &a.Conf
	}
	CheckAgentCommand(c, create, []string{})
}

func (s *ProvisioningSuite) TestParseUnknown(c *C) {
	a := &ProvisioningAgent{}
	err := ParseAgentCommand(a, []string{"nincompoops"})
	c.Assert(err, ErrorMatches, `unrecognized args: \["nincompoops"\]`)
}

func initProvisioningAgent() (*ProvisioningAgent, error) {
	args := []string{"--zookeeper-servers", zkAddr}
	c := &ProvisioningAgent{}
	return c, initCmd(c, args)
}

func (s *ProvisioningSuite) TestProvisionerStartStop(c *C) {
	p, err := NewProvisioner(s.zkInfo)
	c.Assert(err, IsNil)
	c.Assert(p.Stop(), IsNil)
}

func (s *ProvisioningSuite) TestProvisionerDifferentStateInfo(c *C) {
	info := *s.zkInfo

	// Use an equivalent but textually different address and check
	// that the info when new instances are started is that returned
	// from the environment, not the one passed in originally.
	info.Addrs = append([]string{}, info.Addrs...)
	if !strings.HasPrefix(info.Addrs[0], "127.0.0.1") {
		c.Fatalf("local address %q not as expected", info.Addrs)
	}
	info.Addrs[0] = "localhost" + info.Addrs[0][len("127.0.0.1"):]

	p, err := NewProvisioner(&info)
	c.Assert(err, IsNil)
	defer s.stopProvisioner(c, p)

	// place a new machine into the state
	m, err := s.st.AddMachine()
	c.Assert(err, IsNil)

	s.checkStartInstance(c, m)
}

func (s *ProvisioningSuite) TestProvisionerEnvironmentChange(c *C) {
	p, err := NewProvisioner(s.zkInfo)
	c.Assert(err, IsNil)
	defer s.stopProvisioner(c, p)
	// Twiddle with the environment configuration.
	env, err := s.st.EnvironConfig()
	c.Assert(err, IsNil)
	env.Set("name", "testing2")
	_, err = env.Write()
	c.Assert(err, IsNil)
	env.Set("name", "testing3")
	_, err = env.Write()
}

func (s *ProvisioningSuite) TestProvisionerStopOnStateClose(c *C) {
	p, err := NewProvisioner(s.zkInfo)
	c.Assert(err, IsNil)

	p.st.Close()

	// must use Check to avoid leaking PA
	c.Check(p.Wait(), ErrorMatches, ".* zookeeper is closing")
	c.Assert(p.Stop(), ErrorMatches, ".* zookeeper is closing")
}

// Start and stop one machine, watch the PA.
func (s *ProvisioningSuite) TestSimple(c *C) {
	p, err := NewProvisioner(s.zkInfo)
	c.Assert(err, IsNil)
	defer s.stopProvisioner(c, p)

	// place a new machine into the state
	m, err := s.st.AddMachine()
	c.Assert(err, IsNil)

	s.checkStartInstance(c, m)

	// now remove it
	c.Assert(s.st.RemoveMachine(m.Id()), IsNil)

	// watch the PA remove it
	s.checkStopInstance(c)
	s.checkMachineId(c, m, nil)
}

func (s *ProvisioningSuite) TestProvisioningDoesNotOccurWithAnInvalidEnvironment(c *C) {
	err := s.invalidateEnvironment()
	c.Assert(err, IsNil)

	p, err := NewProvisioner(s.zkInfo)
	c.Assert(err, IsNil)
	defer s.stopProvisioner(c, p)

	// try to create a machine
	_, err = s.st.AddMachine()
	c.Assert(err, IsNil)

	// the PA should not create it
	s.checkNotStartInstance(c)
}

func (s *ProvisioningSuite) TestProvisioningOccursWithFixedEnvironment(c *C) {
	err := s.invalidateEnvironment()
	c.Assert(err, IsNil)

	p, err := NewProvisioner(s.zkInfo)
	c.Assert(err, IsNil)
	defer s.stopProvisioner(c, p)

	// try to create a machine
	m, err := s.st.AddMachine()
	c.Assert(err, IsNil)

	// the PA should not create it
	s.checkNotStartInstance(c)

	err = s.fixEnvironment()
	c.Assert(err, IsNil)

	s.checkStartInstance(c, m)
}

func (s *ProvisioningSuite) TestProvisioningDoesOccurAfterInvalidEnvironmentPublished(c *C) {
	p, err := NewProvisioner(s.zkInfo)
	c.Assert(err, IsNil)
	defer s.stopProvisioner(c, p)

	// place a new machine into the state
	m, err := s.st.AddMachine()
	c.Assert(err, IsNil)

	s.checkStartInstance(c, m)

	err = s.invalidateEnvironment()
	c.Assert(err, IsNil)

	// create a second machine
	m, err = s.st.AddMachine()
	c.Assert(err, IsNil)

	// the PA should create it using the old environment
	s.checkStartInstance(c, m)
}

func (s *ProvisioningSuite) TestProvisioningDoesNotProvisionTheSameMachineAfterRestart(c *C) {
	p, err := NewProvisioner(s.zkInfo)
	c.Check(err, IsNil)
	// we are not using defer s.stopProvisioner(c, p) because we need to control when
	// the PA is restarted in this test. tf. Methods like Fatalf and Assert should not be used.

	// create a machine
	m, err := s.st.AddMachine()
	c.Check(err, IsNil)

	s.checkStartInstance(c, m)

	// restart the PA
	c.Check(p.Stop(), IsNil)

	p, err = NewProvisioner(s.zkInfo)
	c.Check(err, IsNil)

	// check that there is only one machine known
	machines, err := p.st.AllMachines()
	c.Check(err, IsNil)
	c.Check(len(machines), Equals, 1)
	c.Check(machines[0].Id(), Equals, 0)

	// the PA should not create it a second time
	s.checkNotStartInstance(c)

	c.Assert(p.Stop(), IsNil)
}

func (s *ProvisioningSuite) TestProvisioningStopsUnknownInstances(c *C) {
	p, err := NewProvisioner(s.zkInfo)
	c.Check(err, IsNil)
	// we are not using defer s.stopProvisioner(c, p) because we need to control when
	// the PA is restarted in this test. Methods like Fatalf and Assert should not be used.

	// create a machine
	m, err := s.st.AddMachine()
	c.Check(err, IsNil)

	s.checkStartInstance(c, m)

	// create a second machine
	m, err = s.st.AddMachine()
	c.Check(err, IsNil)

	s.checkStartInstance(c, m)

	// stop the PA
	c.Check(p.Stop(), IsNil)

	// remove the machine
	err = s.st.RemoveMachine(m.Id())
	c.Check(err, IsNil)

	// start a new provisioner
	p, err = NewProvisioner(s.zkInfo)
	c.Check(err, IsNil)

	s.checkStopInstance(c)

	c.Assert(p.Stop(), IsNil)
}

// This check is different from the one above as it catches the edge case
// where the final machine has been removed from the state while the PA was
// not running.
func (s *ProvisioningSuite) TestProvisioningStopsOnlyUnknownInstances(c *C) {
	p, err := NewProvisioner(s.zkInfo)
	c.Check(err, IsNil)
	// we are not using defer s.stopProvisioner(c, p) because we need to control when
	// the PA is restarted in this test. Methods like Fatalf and Assert should not be used.

	// create a machine
	m, err := s.st.AddMachine()
	c.Check(err, IsNil)

	s.checkStartInstance(c, m)

	// stop the PA
	c.Check(p.Stop(), IsNil)

	// remove the machine
	err = s.st.RemoveMachine(m.Id())
	c.Check(err, IsNil)

	machines, err := s.st.AllMachines()
	c.Check(err, IsNil)
	c.Check(len(machines), Equals, 0) // it's really gone

	// start a new provisioner
	p, err = NewProvisioner(s.zkInfo)
	c.Check(err, IsNil)

	s.checkStopInstance(c)

	c.Assert(p.Stop(), IsNil)
}

func (s *ProvisioningSuite) TestProvisioningRecoversAfterInvalidEnvironmentPublished(c *C) {
	p, err := NewProvisioner(s.zkInfo)
	c.Assert(err, IsNil)
	defer s.stopProvisioner(c, p)

	// place a new machine into the state
	m, err := s.st.AddMachine()
	c.Assert(err, IsNil)

	s.checkStartInstance(c, m)

	err = s.invalidateEnvironment()
	c.Assert(err, IsNil)

	// create a second machine
	m, err = s.st.AddMachine()
	c.Assert(err, IsNil)

	// the PA should create it using the old environment
	s.checkStartInstance(c, m)

	err = s.fixEnvironment()
	c.Assert(err, IsNil)

	// create a third machine
	m, err = s.st.AddMachine()
	c.Assert(err, IsNil)

	// the PA should create it using the new environment
	s.checkStartInstance(c, m)
}
<commit_msg>cmd/jujud: minor changes for review<commit_after>package main

import (
	"time"

	. "launchpad.net/gocheck"
	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/environs/dummy"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/testing"
	"strings"
)

type ProvisioningSuite struct {
	logging testing.LoggingSuite
	zkSuite
	st *state.State
	op <-chan dummy.Operation
}

var _ = Suite(&ProvisioningSuite{})

var veryShortAttempt = environs.AttemptStrategy{
	Total: 1 * time.Second,
	Delay: 80 * time.Millisecond,
}

func (s *ProvisioningSuite) SetUpTest(c *C) {
	s.logging.SetUpTest(c)

	// Create the operations channel with more than enough space
	// for those tests that don't listen on it.
	op := make(chan dummy.Operation, 500)
	dummy.Listen(op)
	s.op = op

	env, err := environs.NewEnviron(map[string]interface{}{
		"type":      "dummy",
		"zookeeper": true,
		"name":      "testing",
	})
	c.Assert(err, IsNil)
	err = env.Bootstrap(false)
	c.Assert(err, IsNil)

	s.zkSuite.SetUpTest(c)
	s.st, err = state.Open(s.zkInfo)
	c.Assert(err, IsNil)

	// Make sure that zkInfo holds exactly the info we're passing to
	// NewProvisioner,
	s.zkInfo, err = env.StateInfo()
	c.Assert(err, IsNil)
}

func (s *ProvisioningSuite) TearDownTest(c *C) {
	dummy.Reset()
	s.zkSuite.TearDownTest()
	s.logging.TearDownTest(c)
}

// invalidateEnvironment alters the environment configuration
// so the ConfigNode returned from the watcher will not pass
// validation.
func (s *ProvisioningSuite) invalidateEnvironment() error {
	env, err := s.st.EnvironConfig()
	if err != nil {
		return err
	}
	env.Set("name", 1)
	_, err = env.Write()
	return err
}

// fixEnvironment undoes the work of invalidateEnvironment.
func (s *ProvisioningSuite) fixEnvironment() error {
	env, err := s.st.EnvironConfig()
	if err != nil {
		return err
	}
	env.Set("name", "testing")
	_, err = env.Write()
	return err
}

func (s *ProvisioningSuite) stopProvisioner(c *C, p *Provisioner) {
	c.Assert(p.Stop(), IsNil)
}

// checkStartInstance checks that an instance has been started
// with a machine id the same as m's, and that the machine's
// instance id has been set appropriately.
func (s *ProvisioningSuite) checkStartInstance(c *C, m *state.Machine) {
	for {
		select {
		case o := <-s.op:
			switch o := o.(type) {
			case dummy.OpStartInstance:
				c.Assert(o.Info, DeepEquals, s.zkInfo)
				c.Assert(o.MachineId, Equals, m.Id())
				c.Assert(o.Instance, NotNil)
				s.checkMachineId(c, m, o.Instance)
				return
			default:
				c.Logf("ignoring unexpected operation %#v", o)
			}
		case <-time.After(2 * time.Second):
			c.Errorf("provisioner did not start an instance")
			return
		}
	}
}

// checkNotStartInstance checks that an instance was not started
func (s *ProvisioningSuite) checkNotStartInstance(c *C) {
	for {
		select {
		case o := <-s.op:
			switch o.(type) {
			case dummy.OpStartInstance:
				c.Errorf("instance started: %v", o)
				return
			default:
				// ignore
			}
		case <-time.After(200 * time.Millisecond):
			return
		}
	}
}

// checkStopInstance checks that an instance has been stopped.
func (s *ProvisioningSuite) checkStopInstance(c *C) {
	// use the non fatal variants to avoid leaking provisioners.
	for {
		select {
		case o := <-s.op:
			switch o.(type) {
			case dummy.OpStopInstances:
				return
			default:
				//ignore
			}
		case <-time.After(2 * time.Second):
			c.Errorf("provisioner did not stop an instance")
			return
		}
	}
}

// checkMachineIdSet checks that the machine has an instance id
// that matches that of the given instance. If the instance is nil,
// It checks that the instance id is unset.
func (s *ProvisioningSuite) checkMachineId(c *C, m *state.Machine, inst environs.Instance) {
	// TODO(dfc) add machine.WatchConfig() to avoid having to poll.
	instId := ""
	if inst != nil {
		instId = inst.Id()
	}
	for a := veryShortAttempt.Start(); a.Next(); {
		id, err := m.InstanceId()
		c.Assert(err, IsNil)
		if id == instId {
			return
		}
	}
	id, err := m.InstanceId()
	c.Assert(err, IsNil)
	c.Assert(id, Equals, instId)
}

func (s *ProvisioningSuite) TestParseSuccess(c *C) {
	create := func() (cmd.Command, *AgentConf) {
		a := &ProvisioningAgent{}
		return a, &a.Conf
	}
	CheckAgentCommand(c, create, []string{})
}

func (s *ProvisioningSuite) TestParseUnknown(c *C) {
	a := &ProvisioningAgent{}
	err := ParseAgentCommand(a, []string{"nincompoops"})
	c.Assert(err, ErrorMatches, `unrecognized args: \["nincompoops"\]`)
}

func initProvisioningAgent() (*ProvisioningAgent, error) {
	args := []string{"--zookeeper-servers", zkAddr}
	c := &ProvisioningAgent{}
	return c, initCmd(c, args)
}

func (s *ProvisioningSuite) TestProvisionerStartStop(c *C) {
	p, err := NewProvisioner(s.zkInfo)
	c.Assert(err, IsNil)
	c.Assert(p.Stop(), IsNil)
}

func (s *ProvisioningSuite) TestProvisionerDifferentStateInfo(c *C) {
	info := *s.zkInfo

	// Use an equivalent but textually different address and check
	// that the info when new instances are started is that returned
	// from the environment, not the one passed in originally.
	info.Addrs = append([]string{}, info.Addrs...)
	if !strings.HasPrefix(info.Addrs[0], "127.0.0.1") {
		c.Fatalf("local address %q not as expected", info.Addrs)
	}
	info.Addrs[0] = "localhost" + info.Addrs[0][len("127.0.0.1"):]

	p, err := NewProvisioner(&info)
	c.Assert(err, IsNil)
	defer s.stopProvisioner(c, p)

	// place a new machine into the state
	m, err := s.st.AddMachine()
	c.Assert(err, IsNil)

	s.checkStartInstance(c, m)
}

func (s *ProvisioningSuite) TestProvisionerEnvironmentChange(c *C) {
	p, err := NewProvisioner(s.zkInfo)
	c.Assert(err, IsNil)
	defer s.stopProvisioner(c, p)
	// Twiddle with the environment configuration.
	env, err := s.st.EnvironConfig()
	c.Assert(err, IsNil)
	env.Set("name", "testing2")
	_, err = env.Write()
	c.Assert(err, IsNil)
	env.Set("name", "testing3")
	_, err = env.Write()
}

func (s *ProvisioningSuite) TestProvisionerStopOnStateClose(c *C) {
	p, err := NewProvisioner(s.zkInfo)
	c.Assert(err, IsNil)

	p.st.Close()

	// must use Check to avoid leaking PA
	c.Check(p.Wait(), ErrorMatches, ".* zookeeper is closing")
	c.Assert(p.Stop(), ErrorMatches, ".* zookeeper is closing")
}

// Start and stop one machine, watch the PA.
func (s *ProvisioningSuite) TestSimple(c *C) {
	p, err := NewProvisioner(s.zkInfo)
	c.Assert(err, IsNil)
	defer s.stopProvisioner(c, p)

	// place a new machine into the state
	m, err := s.st.AddMachine()
	c.Assert(err, IsNil)

	s.checkStartInstance(c, m)

	// now remove it
	c.Assert(s.st.RemoveMachine(m.Id()), IsNil)

	// watch the PA remove it
	s.checkStopInstance(c)
	s.checkMachineId(c, m, nil)
}

func (s *ProvisioningSuite) TestProvisioningDoesNotOccurWithAnInvalidEnvironment(c *C) {
	err := s.invalidateEnvironment()
	c.Assert(err, IsNil)

	p, err := NewProvisioner(s.zkInfo)
	c.Assert(err, IsNil)
	defer s.stopProvisioner(c, p)

	// try to create a machine
	_, err = s.st.AddMachine()
	c.Assert(err, IsNil)

	// the PA should not create it
	s.checkNotStartInstance(c)
}

func (s *ProvisioningSuite) TestProvisioningOccursWithFixedEnvironment(c *C) {
	err := s.invalidateEnvironment()
	c.Assert(err, IsNil)

	p, err := NewProvisioner(s.zkInfo)
	c.Assert(err, IsNil)
	defer s.stopProvisioner(c, p)

	// try to create a machine
	m, err := s.st.AddMachine()
	c.Assert(err, IsNil)

	// the PA should not create it
	s.checkNotStartInstance(c)

	err = s.fixEnvironment()
	c.Assert(err, IsNil)

	s.checkStartInstance(c, m)
}

func (s *ProvisioningSuite) TestProvisioningDoesOccurAfterInvalidEnvironmentPublished(c *C) {
	p, err := NewProvisioner(s.zkInfo)
	c.Assert(err, IsNil)
	defer s.stopProvisioner(c, p)

	// place a new machine into the state
	m, err := s.st.AddMachine()
	c.Assert(err, IsNil)

	s.checkStartInstance(c, m)

	err = s.invalidateEnvironment()
	c.Assert(err, IsNil)

	// create a second machine
	m, err = s.st.AddMachine()
	c.Assert(err, IsNil)

	// the PA should create it using the old environment
	s.checkStartInstance(c, m)
}

func (s *ProvisioningSuite) TestProvisioningDoesNotProvisionTheSameMachineAfterRestart(c *C) {
	p, err := NewProvisioner(s.zkInfo)
	c.Check(err, IsNil)
	// we are not using defer s.stopProvisioner(c, p) because we need to control when
	// the PA is restarted in this test. tf. Methods like Fatalf and Assert should not be used.

	// create a machine
	m, err := s.st.AddMachine()
	c.Check(err, IsNil)

	s.checkStartInstance(c, m)

	// restart the PA
	c.Check(p.Stop(), IsNil)

	p, err = NewProvisioner(s.zkInfo)
	c.Check(err, IsNil)

	// check that there is only one machine known
	machines, err := p.st.AllMachines()
	c.Check(err, IsNil)
	c.Check(len(machines), Equals, 1)
	c.Check(machines[0].Id(), Equals, 0)

	// the PA should not create it a second time
	s.checkNotStartInstance(c)

	c.Assert(p.Stop(), IsNil)
}

func (s *ProvisioningSuite) TestProvisioningStopsUnknownInstances(c *C) {
	p, err := NewProvisioner(s.zkInfo)
	c.Check(err, IsNil)
	// we are not using defer s.stopProvisioner(c, p) because we need to control when
	// the PA is restarted in this test. Methods like Fatalf and Assert should not be used.

	// create a machine
	m, err := s.st.AddMachine()
	c.Check(err, IsNil)

	s.checkStartInstance(c, m)

	// create a second machine
	m, err = s.st.AddMachine()
	c.Check(err, IsNil)

	s.checkStartInstance(c, m)

	// stop the PA
	c.Check(p.Stop(), IsNil)

	// remove the machine
	err = s.st.RemoveMachine(m.Id())
	c.Check(err, IsNil)

	// start a new provisioner
	p, err = NewProvisioner(s.zkInfo)
	c.Check(err, IsNil)

	s.checkStopInstance(c)

	c.Assert(p.Stop(), IsNil)
}

// This check is different from the one above as it catches the edge case
// where the final machine has been removed from the state while the PA was
// not running.
func (s *ProvisioningSuite) TestProvisioningStopsOnlyUnknownInstances(c *C) {
	p, err := NewProvisioner(s.zkInfo)
	c.Check(err, IsNil)
	// we are not using defer s.stopProvisioner(c, p) because we need to control when
	// the PA is restarted in this test. Methods like Fatalf and Assert should not be used.

	// create a machine
	m, err := s.st.AddMachine()
	c.Check(err, IsNil)

	s.checkStartInstance(c, m)

	// stop the PA
	c.Check(p.Stop(), IsNil)

	// remove the machine
	err = s.st.RemoveMachine(m.Id())
	c.Check(err, IsNil)

	machines, err := s.st.AllMachines()
	c.Check(err, IsNil)
	c.Check(len(machines), Equals, 0) // it's really gone

	// start a new provisioner
	p, err = NewProvisioner(s.zkInfo)
	c.Check(err, IsNil)

	s.checkStopInstance(c)

	c.Assert(p.Stop(), IsNil)
}

func (s *ProvisioningSuite) TestProvisioningRecoversAfterInvalidEnvironmentPublished(c *C) {
	p, err := NewProvisioner(s.zkInfo)
	c.Assert(err, IsNil)
	defer s.stopProvisioner(c, p)

	// place a new machine into the state
	m, err := s.st.AddMachine()
	c.Assert(err, IsNil)

	s.checkStartInstance(c, m)

	err = s.invalidateEnvironment()
	c.Assert(err, IsNil)

	// create a second machine
	m, err = s.st.AddMachine()
	c.Assert(err, IsNil)

	// the PA should create it using the old environment
	s.checkStartInstance(c, m)

	err = s.fixEnvironment()
	c.Assert(err, IsNil)

	// create a third machine
	m, err = s.st.AddMachine()
	c.Assert(err, IsNil)

	// the PA should create it using the new environment
	s.checkStartInstance(c, m)
}
<|endoftext|>


<commit_before>package jsonstream

import (
	"bufio"
	"bytes"
	"io"
	"log"
	"strings"
	"testing"
)

type TestObj struct {
	Name string `json:"name,omitempty"`
}
type TestSummary struct {
	Value int `json:"val"`
}

var testHeader1 = "01234567"
var testHeader2 = "0123456x"

func TestStream(t *testing.T) {
	var buf bytes.Buffer

	data := []TestObj{
		{Name: "foo"},
		{Name: "bar"},
		{Name: "baz"},
	}

	w := NewWriter(&buf, testHeader1)
	for _, d := range data {
		if err := w.Write(&d); err != nil {
			t.Errorf("write error: %v", err)
		}
	}
	w.Finalize()
	log.Printf("wrote: %v", buf.String())
	r, err := NewReader(bufio.NewReader(&buf), testHeader1, nil)
	if err != nil {
		t.Errorf("err: %v", err)
		return
	}
	for _, d := range data {
		v := &TestObj{}
		if readerr := r.Read(v); readerr != nil {
			t.Errorf("read error: %v", readerr)
		}
		if v.Name != d.Name {
			t.Errorf("invalid value: '%v', expected '%v'", v.Name, d.Name)
		}
	}
	v := &TestObj{}
	err = r.Read(v)
	if err != io.EOF {
		t.Errorf("expected EOF, got %v", err)
	}
}

func TestStreamWithSummary(t *testing.T) {
	var buf bytes.Buffer

	data := []TestObj{
		{Name: "foo"},
		{Name: "bar"},
		{Name: "baz"},
	}

	w := NewWriter(&buf, testHeader1)
	for _, d := range data {
		if err := w.Write(&d); err != nil {
			t.Errorf("write error: %v", err)
		}
	}
	w.FinalizeWithSummary(TestSummary{Value: 123})
	log.Printf("wrote: %v", buf.String())

	var summary TestSummary
	r, err := NewReader(bufio.NewReader(&buf), testHeader1, &summary)
	if err != nil {
		t.Errorf("err: %v", err)
		return
	}
	for _, d := range data {
		v := &TestObj{}
		if readerr := r.Read(v); readerr != nil {
			t.Errorf("read error: %v", readerr)
		}
		if v.Name != d.Name {
			t.Errorf("invalid value: '%v', expected '%v'", v.Name, d.Name)
		}
	}
	v := &TestObj{}
	err = r.Read(v)
	if err != io.EOF {
		t.Errorf("expected EOF, got %v", err)
	}
	if got, want := summary.Value, 123; got != want {
		t.Errorf("unexpected summary value: %v, wanted %v", got, want)
	}
}

func TestInvalidHeader(t *testing.T) {
	var buf bytes.Buffer

	w := NewWriter(&buf, testHeader1)
	if err := w.Write(&TestObj{Name: "foo"}); err != nil {
		t.Errorf("write error: %v", err)
	}

	_, err := NewReader(bufio.NewReader(&buf), testHeader2, nil)
	if err == nil {
		t.Errorf("expected error, got none")
	} else if !strings.Contains(err.Error(), "invalid stream format") {
		t.Errorf("got incorrect error: %v", err)
	}
}
<commit_msg>jsonstream: added test cases<commit_after>package jsonstream

import (
	"bufio"
	"bytes"
	"io"
	"log"
	"strings"
	"testing"
)

type TestObj struct {
	Name string `json:"name,omitempty"`
}
type TestSummary struct {
	Value int `json:"val"`
}

var testHeader1 = "01234567"
var testHeader2 = "0123456x"

func TestStream(t *testing.T) {
	var buf bytes.Buffer

	data := []TestObj{
		{Name: "foo"},
		{Name: "bar"},
		{Name: "baz"},
	}

	w := NewWriter(&buf, testHeader1)
	for _, d := range data {
		if err := w.Write(&d); err != nil {
			t.Errorf("write error: %v", err)
		}
	}
	w.Finalize()
	log.Printf("wrote: %v", buf.String())
	r, err := NewReader(bufio.NewReader(&buf), testHeader1, nil)
	if err != nil {
		t.Errorf("err: %v", err)
		return
	}
	for _, d := range data {
		v := &TestObj{}
		if readerr := r.Read(v); readerr != nil {
			t.Errorf("read error: %v", readerr)
		}
		if v.Name != d.Name {
			t.Errorf("invalid value: '%v', expected '%v'", v.Name, d.Name)
		}
	}
	v := &TestObj{}
	err = r.Read(v)
	if err != io.EOF {
		t.Errorf("expected EOF, got %v", err)
	}
}

func TestStreamWithSummary(t *testing.T) {
	var buf bytes.Buffer

	data := []TestObj{
		{Name: "foo"},
		{Name: "bar"},
		{Name: "baz"},
	}

	w := NewWriter(&buf, testHeader1)
	for _, d := range data {
		if err := w.Write(&d); err != nil {
			t.Errorf("write error: %v", err)
		}
	}
	w.FinalizeWithSummary(TestSummary{Value: 123})
	log.Printf("wrote: %v", buf.String())

	var summary TestSummary
	r, err := NewReader(bufio.NewReader(&buf), testHeader1, &summary)
	if err != nil {
		t.Errorf("err: %v", err)
		return
	}
	for _, d := range data {
		v := &TestObj{}
		if readerr := r.Read(v); readerr != nil {
			t.Errorf("read error: %v", readerr)
		}
		if v.Name != d.Name {
			t.Errorf("invalid value: '%v', expected '%v'", v.Name, d.Name)
		}
	}
	v := &TestObj{}
	err = r.Read(v)
	if err != io.EOF {
		t.Errorf("expected EOF, got %v", err)
	}
	if got, want := summary.Value, 123; got != want {
		t.Errorf("unexpected summary value: %v, wanted %v", got, want)
	}
}

func TestInvalidHeader(t *testing.T) {
	var buf bytes.Buffer

	w := NewWriter(&buf, testHeader1)
	if err := w.Write(&TestObj{Name: "foo"}); err != nil {
		t.Errorf("write error: %v", err)
	}

	_, err := NewReader(bufio.NewReader(&buf), testHeader2, nil)
	if err == nil {
		t.Errorf("expected error, got none")
	} else if !strings.Contains(err.Error(), "invalid stream format") {
		t.Errorf("got incorrect error: %v", err)
	}
}

func TestInvalidStream(t *testing.T) {
	cases := []string{
		`x`,
		`{}`,
		`{"not-stream":"hdr"}`,
		`{{}}`,
		`{"stream":"non-hdr"}`,
		`{"stream":"hdr","nonEntries":[]}`,
		`{"stream":"hdr","entries":{}}`,
		`{"stream":"hdr","entries":[]}`,
		`{"stream":"hdr","entries":[`,
		`{"stream":"hdr","entries":[}`,
		`{"stream":"hdr","entries":[]`,
		`{"stream":"hdr","entries":[],"summary"`,
		`{"stream":"hdr","entries":[],1.222.33`,
		`{"stream":"hdr","entries":[],"sxummary":{"x":"1",`,
	}

	for _, tc := range cases {
		r, err := NewReader(strings.NewReader(tc), "hdr", nil)
		if err != nil {
			if !isInvalidStream(err) {
				t.Errorf("got invalid error when creating reader: %v", err)
			}
			continue
		}

		for {
			v := map[string]interface{}{}
			if err := r.Read(v); err != nil {
				if err == io.EOF {
					break
				}
				if !isInvalidStream(err) {
					t.Errorf("got invalid error when creating reader: %v", err)
				}
				break
			}
		}
	}
}

func isInvalidStream(e error) bool {
	return e != nil && strings.Contains(e.Error(), "invalid stream format")
}
<|endoftext|>


<commit_before>package memory

import (
	"testing"
)

func TestMemory_Mkdir(t *testing.T) {
	fs := NewMemory()

	fs.Mkdir("test", 0666)
	f, _ := fs.Open("test")
	stat, _ := f.Stat()
	if !stat.IsDir() {
		t.Fail()
	}
	info, _ := f.(*file).Info()
	if info.Mode().IsRegular() {
		t.Fatal("unexpected regular file")
	}
}

func TestMemory_WriteFile(t *testing.T) {
	fs := NewMemory()

	fs.WriteFile("test", []byte{}, 0666)
	file, _ := fs.Open("test")
	stat, _ := file.Stat()
	if !stat.Mode().IsRegular() {
		t.Fail()
	}
}

func TestMemory_Create(t *testing.T) {
	fs := NewMemory()

	fs.Create("test")
	file, _ := fs.Open("test")
	stat, _ := file.Stat()
	if !stat.Mode().IsRegular() {
		t.Fail()
	}
}
<commit_msg>feat: add full test<commit_after>package memory

import (
	"io/ioutil"
	"path/filepath"
	"testing"
)

func TestMemory_Mkdir(t *testing.T) {
	fs := NewMemory()

	fs.Mkdir("test", 0666)
	f, _ := fs.Open("test")
	stat, _ := f.Stat()
	if !stat.IsDir() {
		t.Fail()
	}
	info, _ := f.(*file).Info()
	if info.Mode().IsRegular() {
		t.Fatal("unexpected regular file")
	}
}

func TestMemory_WriteFile(t *testing.T) {
	fs := NewMemory()

	fs.WriteFile("test", []byte{}, 0666)
	file, _ := fs.Open("test")
	stat, _ := file.Stat()
	if !stat.Mode().IsRegular() {
		t.Fail()
	}
}

func TestMemory_Create(t *testing.T) {
	fs := NewMemory()

	fs.Create("test")
	file, _ := fs.Open("test")
	stat, _ := file.Stat()
	if !stat.Mode().IsRegular() {
		t.Fail()
	}
}

func TestMemory(t *testing.T) {
	fs := NewMemory()
	err := fs.Mkdir("directory", 0666)
	if err != nil {
		t.Fatalf("mkdir error %v", err)
	}
	err = fs.WriteFile(filepath.Join("directory", "test"), []byte(`content`), 0666)
	if err != nil {
		t.Fatalf("writefile error: %v", err)
	}
	f, err := fs.Open(filepath.Join("directory", "test"))
	if err != nil {
		t.Fatalf("open error: %v", err)
	}
	content, err := ioutil.ReadAll(f)
	if err != nil {
		t.Fatalf("readall error: %v", err)
	}
	if string(content) != "content" {
		t.Fatal("unexpected content")
	}
	err = fs.RemoveAll("directory")
	if err != nil {
		t.Fatalf("removeall err: %v", err)
	}
	_, err = fs.Open(filepath.Join("directory", "test"))
	if err == nil {
		t.Fatal("file should be gone")
	}
}
<|endoftext|>


<commit_before>type RingBuffer struct {
    buffer []interface{}
    beginPos, endPos int
}

func (rb *RingBuffer) GetDataSlice() []interface{} {
    if rb.beginPos <= rb.endPos {
        return rb.buffer[rb.beginPos:rb.endPos+1]
    } else {
        return append(buffer[rb.beginPos:], rb.buffer[:rb.endPos+1]
    }
}

func (rb *RingBuffer) Append(item interface{}) {
    if rb.beginPos == rb.endPos + 1 {
        rb.beginPos++
    }

    if rb.endPos == len(rb.buffer) - 1 {
        rb.endPos = 0
    } else {
        rb.endPos++
    }

    if rb.beginPos == len(rb.buffer) - 1 {
        rb.beginPos = 0
    }

    rb.buffer[rb.endPos] = item
}
<commit_msg>added constructor for RingBuffer<commit_after>type RingBuffer struct {
    buffer []interface{}
    beginPos, endPos int
}

func (rb *RingBuffer) GetDataSlice() []interface{} {
    if rb.beginPos <= rb.endPos {
        return rb.buffer[rb.beginPos:rb.endPos+1]
    } else {
        return append(buffer[rb.beginPos:], rb.buffer[:rb.endPos+1]
    }
}

func (rb *RingBuffer) Append(item interface{}) {
    if rb.beginPos == rb.endPos + 1 {
        rb.beginPos++
    }

    if rb.endPos == len(rb.buffer) - 1 {
        rb.endPos = 0
    } else {
        rb.endPos++
    }

    if rb.beginPos == len(rb.buffer) - 1 {
        rb.beginPos = 0
    }

    rb.buffer[rb.endPos] = item
}

func NewRingBuffer(size int) *RingBuffer {
    return &RingBuffer{buffer: make([]interface{}, size)}
}
<|endoftext|>


<commit_before>package gosql

import "strings"

func (q *QueryRoot) Parse(m map[string]interface{}) error {
	if m == nil || len(m) == 0 {
		return nil
	}
	q.Value = make([]IQuery, 0, len(m))
	for k, v := range m {
		if v == nil {
		} else if IsQueryKey1(k) || IsQueryKey2(k) {
			elem := &QueryElem{anonymous: IsQueryKey2(k), Key: k}
			if err := elem.Parse(v); err == nil {
				q.Value = append(q.Value, elem)
			}
		}
	}
	return nil
}

func (q *QueryElem) Parse(obj interface{}) error {
	if m, ok := obj.(map[string]interface{}); !ok {
		return ErrTypeMap
	} else if m == nil || len(m) == 0 {
		return ErrTypeMap
	} else {
		q.Value = make([]IQuery, 0, len(m))
		for k, v := range m {
			if v == nil {
			} else if IsQueryKey1(k) || IsQueryKey2(k) {
				elem := &QueryElem{anonymous: IsQueryKey2(k), Key: k}
				if err := elem.Parse(v); err == nil {
					q.Value = append(q.Value, elem)
				}
			} else {
				value := &QueryValue{Key: q.Key, Field: k}
				if err := value.Parse(v); err == nil {
					q.Value = append(q.Value, value)
				}
			}
		}
	}
	if q.Value == nil || len(q.Value) == 0 {
		return ErrTypeMap
	}
	return nil
}

func (q *QueryValue) Parse(obj interface{}) error {
	if obj == nil {
		return ErrTypeString
	} else if strings.HasPrefix(q.Field, "%") {
		return ErrTypeString
	}
	switch v := obj.(type) {
	case int:
		q.Value = v
	case int64:
		q.Value = v
	case float32:
		q.Value = v
	case float64:
		q.Value = v
	case string:
		q.Value = v
	default:
		return ErrTypeString
	}
	return nil
}
<commit_msg>fix: query parser<commit_after>package gosql

import "strings"

func (q *QueryRoot) Parse(m map[string]interface{}) error {
	if m == nil || len(m) == 0 {
		return nil
	}
	q.Value = make([]IQuery, 0, len(m))
	for k, v := range m {
		if v == nil {
		} else if IsQueryKey1(k) || IsQueryKey2(k) {
			elem := &QueryElem{anonymous: IsQueryKey2(k), Key: k}
			if err := elem.Parse(v); err == nil {
				q.Value = append(q.Value, elem)
			}
		}
	}
	return nil
}

func (q *QueryElem) Parse(obj interface{}) error {
	if m, ok := obj.(map[string]interface{}); !ok {
		return ErrTypeMap
	} else if m == nil || len(m) == 0 {
		return ErrTypeMap
	} else {
		q.Value = make([]IQuery, 0, len(m))
		for k, v := range m {
			if v == nil {
			} else if IsQueryKey1(k) || IsQueryKey2(k) {
				if !q.IsAnonymous() {
					elem := &QueryElem{anonymous: IsQueryKey2(k), Key: k}
					if err := elem.Parse(v); err == nil {
						q.Value = append(q.Value, elem)
					}
				}
			} else {
				value := &QueryValue{Key: q.Key, Field: k}
				if err := value.Parse(v); err == nil {
					q.Value = append(q.Value, value)
				}
			}
		}
	}
	if q.Value == nil || len(q.Value) == 0 {
		return ErrTypeMap
	}
	return nil
}

func (q *QueryValue) Parse(obj interface{}) error {
	if obj == nil {
		return ErrTypeString
	} else if strings.HasPrefix(q.Field, "%") {
		return ErrTypeString
	}
	switch v := obj.(type) {
	case int:
		q.Value = v
	case int64:
		q.Value = v
	case float32:
		q.Value = v
	case float64:
		q.Value = v
	case string:
		q.Value = v
	default:
		return ErrTypeString
	}
	return nil
}
<|endoftext|>


<commit_before>package screenshot

import (
	// #cgo LDFLAGS: -framework CoreGraphics
	// #cgo LDFLAGS: -framework CoreFoundation
	// #include <CoreGraphics/CoreGraphics.h>
	// #include <CoreFoundation/CoreFoundation.h>
	"C"
	"image"
	"reflect"
	"unsafe"
)

func ScreenRect() (image.Rectangle, error) {
	displayID := C.CGMainDisplayID()
	width := int(C.CGDisplayPixelsWide(displayID))
	height := int(C.CGDisplayPixelsHigh(displayID))
	return image.Rect(0, 0, width, height), nil
}

func CaptureScreen() (*image.RGBA, error) {
	rect, err := ScreenRect()
	if err != nil {
		return nil, err
	}
	return CaptureRect(rect)
}

func CaptureRect(rect image.Rectangle) (*image.RGBA, error) {
	displayID := C.CGMainDisplayID()
	width := int(C.CGDisplayPixelsWide(displayID))
	rawData := C.CGDataProviderCopyData(C.CGImageGetDataProvider(C.CGDisplayCreateImage(displayID)))

	length := int(C.CFDataGetLength(rawData))
	ptr := unsafe.Pointer(C.CFDataGetBytePtr(rawData))

	var slice []byte
	hdrp := (*reflect.SliceHeader)(unsafe.Pointer(&slice))
	hdrp.Data = uintptr(ptr)
	hdrp.Len = length
	hdrp.Cap = length

	imageBytes := make([]byte, length)

	for i := 0; i < length; i += 4 {
		imageBytes[i], imageBytes[i+2], imageBytes[i+1], imageBytes[i+3] = slice[i+2], slice[i], slice[i+1], slice[i+3]
	}

	C.CFRelease(C.CFTypeRef(rawData))

	img := &image.RGBA{Pix: imageBytes, Stride: 4 * width, Rect: rect}
	return img, nil
}
<commit_msg>garbled screenshots on macOS when screen width is not divisible by 16 (#14)<commit_after>package screenshot

import (
	// #cgo LDFLAGS: -framework CoreGraphics
	// #cgo LDFLAGS: -framework CoreFoundation
	// #include <CoreGraphics/CoreGraphics.h>
	// #include <CoreFoundation/CoreFoundation.h>
	"C"
	"image"
	"reflect"
	"unsafe"
	"math"
)

func ScreenRect() (image.Rectangle, error) {
	displayID := C.CGMainDisplayID()
	width := int(C.CGDisplayPixelsWide(displayID))
	height := int(C.CGDisplayPixelsHigh(displayID))
	return image.Rect(0, 0, width, height), nil
}

func CaptureScreen() (*image.RGBA, error) {
	rect, err := ScreenRect()
	if err != nil {
		return nil, err
	}
	return CaptureRect(rect)
}

func CaptureRect(rect image.Rectangle) (*image.RGBA, error) {
	displayID := C.CGMainDisplayID()
	width := int(math.Ceil(float64(C.CGDisplayPixelsWide(displayID))/16)*16)
	rawData := C.CGDataProviderCopyData(C.CGImageGetDataProvider(C.CGDisplayCreateImage(displayID)))

	length := int(C.CFDataGetLength(rawData))
	ptr := unsafe.Pointer(C.CFDataGetBytePtr(rawData))

	var slice []byte
	hdrp := (*reflect.SliceHeader)(unsafe.Pointer(&slice))
	hdrp.Data = uintptr(ptr)
	hdrp.Len = length
	hdrp.Cap = length

	imageBytes := make([]byte, length)

	for i := 0; i < length; i += 4 {
		imageBytes[i], imageBytes[i+2], imageBytes[i+1], imageBytes[i+3] = slice[i+2], slice[i], slice[i+1], slice[i+3]
	}

	C.CFRelease(C.CFTypeRef(rawData))

	img := &image.RGBA{Pix: imageBytes, Stride: 4 * width, Rect: rect}
	return img, nil
}
<|endoftext|>