<commit_before>package app_test

import (
	"database/sql"
	_ "github.com/mattn/go-sqlite3"
	"github.com/timeredbull/tsuru/api/app"
	. "launchpad.net/gocheck"
	"os"
	"testing"
)

func Test(t *testing.T) { TestingT(t) }

type S struct {
	db *sql.DB
}

var _ = Suite(&S{})

func (s *S) SetUpSuite(c *C) {
	s.db, _ = sql.Open("sqlite3", "./tsuru.db")
	_, err := s.db.Exec("CREATE TABLE 'apps' ('id' INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL, 'name' varchar(255), 'framework' varchar(255), 'state' varchar(255), ip varchar(100))")
	c.Check(err, IsNil)
}

func (s *S) TearDownSuite(c *C) {
	os.Remove("./tsuru.db")
	s.db.Close()
}

func (s *S) TearDownTest(c *C) {
	s.db.Exec("DELETE FROM apps")
}

func (s *S) TestAll(c *C) {
	expected := make([]app.App, 0)
	app1 := app.App{Name: "app1"}
	app1.Create()
	expected = append(expected, app1)
	app2 := app.App{Name: "app2"}
	app2.Create()
	expected = append(expected, app2)
	app3 := app.App{Name: "app3"}
	app3.Create()
	expected = append(expected, app3)

	appList, err := app.AllApps()
	c.Assert(err, IsNil)
	c.Assert(expected, DeepEquals, appList)

	app1.Destroy()
	app2.Destroy()
	app3.Destroy()
}

func (s *S) TestGet(c *C) {
	newApp := app.App{Name: "myApp", Framework: "django"}
	err := newApp.Create()
	c.Assert(err, IsNil)

	myApp := app.App{Name: "myApp"}
	err = myApp.Get()
	c.Assert(err, IsNil)
	c.Assert(myApp, Equals, newApp)

	err = myApp.Destroy()
	c.Assert(err, IsNil)
}

func (s *S) TestDestroy(c *C) {
	app := app.App{}
	app.Name = "appName"
	app.Framework = "django"

	err := app.Create()
	c.Assert(err, IsNil)

	err = app.Destroy()
	c.Assert(err, IsNil)

	rows, err := s.db.Query("SELECT count(*) FROM apps WHERE name = 'appName'")

	if err != nil {
		panic(err)
	}

	var qtd int

	for rows.Next() {
		rows.Scan(&qtd)
	}

	c.Assert(qtd, Equals, 0)
}

func (s *S) TestCreate(c *C) {
	app := app.App{}
	app.Name = "appName"
	app.Framework = "django"

	err := app.Create()
	c.Assert(err, IsNil)

	c.Assert(app.State, Equals, "Pending")
	c.Assert(app.Id, Not(Equals), int64(0))

	rows, err := s.db.Query("SELECT id, name, framework, state FROM apps WHERE name = 'appName'")

	if err != nil {
		panic(err)
	}

	var state string
	var name string
	var framework string
	var id int

	for rows.Next() {
		rows.Scan(&id, &name, &framework, &state)
	}

	c.Assert(id, Equals, int(app.Id))
	c.Assert(name, Equals, app.Name)
	c.Assert(framework, Equals, app.Framework)
	c.Assert(state, Equals, app.State)

	app.Destroy()
}
<commit_msg>assert err is nil instead panic them<commit_after>package app_test

import (
	"database/sql"
	_ "github.com/mattn/go-sqlite3"
	"github.com/timeredbull/tsuru/api/app"
	. "launchpad.net/gocheck"
	"os"
	"testing"
)

func Test(t *testing.T) { TestingT(t) }

type S struct {
	db *sql.DB
}

var _ = Suite(&S{})

func (s *S) SetUpSuite(c *C) {
	s.db, _ = sql.Open("sqlite3", "./tsuru.db")
	_, err := s.db.Exec("CREATE TABLE 'apps' ('id' INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL, 'name' varchar(255), 'framework' varchar(255), 'state' varchar(255), ip varchar(100))")
	c.Check(err, IsNil)
}

func (s *S) TearDownSuite(c *C) {
	os.Remove("./tsuru.db")
	s.db.Close()
}

func (s *S) TearDownTest(c *C) {
	s.db.Exec("DELETE FROM apps")
}

func (s *S) TestAll(c *C) {
	expected := make([]app.App, 0)
	app1 := app.App{Name: "app1"}
	app1.Create()
	expected = append(expected, app1)
	app2 := app.App{Name: "app2"}
	app2.Create()
	expected = append(expected, app2)
	app3 := app.App{Name: "app3"}
	app3.Create()
	expected = append(expected, app3)

	appList, err := app.AllApps()
	c.Assert(err, IsNil)
	c.Assert(expected, DeepEquals, appList)

	app1.Destroy()
	app2.Destroy()
	app3.Destroy()
}

func (s *S) TestGet(c *C) {
	newApp := app.App{Name: "myApp", Framework: "django"}
	err := newApp.Create()
	c.Assert(err, IsNil)

	myApp := app.App{Name: "myApp"}
	err = myApp.Get()
	c.Assert(err, IsNil)
	c.Assert(myApp, Equals, newApp)

	err = myApp.Destroy()
	c.Assert(err, IsNil)
}

func (s *S) TestDestroy(c *C) {
	app := app.App{}
	app.Name = "appName"
	app.Framework = "django"

	err := app.Create()
	c.Assert(err, IsNil)

	err = app.Destroy()
	c.Assert(err, IsNil)

	rows, err := s.db.Query("SELECT count(*) FROM apps WHERE name = 'appName'")
	c.Assert(err, IsNil)

	var qtd int
	for rows.Next() {
		rows.Scan(&qtd)
	}

	c.Assert(qtd, Equals, 0)
}

func (s *S) TestCreate(c *C) {
	app := app.App{}
	app.Name = "appName"
	app.Framework = "django"

	err := app.Create()
	c.Assert(err, IsNil)

	c.Assert(app.State, Equals, "Pending")
	c.Assert(app.Id, Not(Equals), int64(0))

	rows, err := s.db.Query("SELECT id, name, framework, state FROM apps WHERE name = 'appName'")
	c.Assert(err, IsNil)

	var state string
	var name string
	var framework string
	var id int

	for rows.Next() {
		rows.Scan(&id, &name, &framework, &state)
	}

	c.Assert(id, Equals, int(app.Id))
	c.Assert(name, Equals, app.Name)
	c.Assert(framework, Equals, app.Framework)
	c.Assert(state, Equals, app.State)

	app.Destroy()
}
<|endoftext|>


<commit_before>/*
Copyright 2013 CoreOS Inc.

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

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

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

package dbus

import (
	"fmt"
	"github.com/guelfey/go.dbus"
	"math/rand"
	"os"
	"path/filepath"
	"reflect"
	"testing"
)

func setupConn(t *testing.T) *Conn {
	conn, err := New()
	if err != nil {
		t.Fatal(err)
	}

	return conn
}

func setupUnit(target string, conn *Conn, t *testing.T) {
	// Blindly stop the unit in case it is running
	conn.StopUnit(target, "replace")

	// Blindly remove the symlink in case it exists
	targetRun := filepath.Join("/run/systemd/system/", target)
	err := os.Remove(targetRun)

	// 1. Enable the unit
	abs, err := filepath.Abs("../fixtures/" + target)
	if err != nil {
		t.Fatal(err)
	}

	fixture := []string{abs}

	install, changes, err := conn.EnableUnitFiles(fixture, true, true)
	if err != nil {
		t.Fatal(err)
	}

	if install != false {
		t.Fatal("Install was true")
	}

	if len(changes) < 1 {
		t.Fatalf("Expected one change, got %v", changes)
	}

	if changes[0].Filename != targetRun {
		t.Fatal("Unexpected target filename")
	}
}

// Ensure that basic unit starting and stopping works.
func TestStartStopUnit(t *testing.T) {
	target := "start-stop.service"
	conn := setupConn(t)

	setupUnit(target, conn, t)

	// 2. Start the unit
	job, err := conn.StartUnit(target, "replace")
	if err != nil {
		t.Fatal(err)
	}

	if job != "done" {
		t.Fatal("Job is not done, %v", job)
	}

	units, err := conn.ListUnits()

	var unit *UnitStatus
	for _, u := range units {
		if u.Name == target {
			unit = &u
		}
	}

	if unit == nil {
		t.Fatalf("Test unit not found in list")
	}

	if unit.ActiveState != "active" {
		t.Fatalf("Test unit not active")
	}

	// 3. Stop the unit
	job, err = conn.StopUnit(target, "replace")
	if err != nil {
		t.Fatal(err)
	}

	units, err = conn.ListUnits()

	unit = nil
	for _, u := range units {
		if u.Name == target {
			unit = &u
		}
	}

	if unit != nil {
		t.Fatalf("Test unit found in list, should be stopped")
	}
}

// Enables a unit and then immediately tears it down
func TestEnableDisableUnit(t *testing.T) {
	target := "enable-disable.service"
	conn := setupConn(t)

	setupUnit(target, conn, t)

	abs, err := filepath.Abs("../fixtures/" + target)
	if err != nil {
		t.Fatal(err)
	}

	path := filepath.Join("/run/systemd/system/", target)

	// 2. Disable the unit
	changes, err := conn.DisableUnitFiles([]string{abs}, true)
	if err != nil {
		t.Fatal(err)
	}

	if len(changes) != 1 {
		t.Fatalf("Changes should include the path, %v", changes)
	}
	if changes[0].Filename != path {
		t.Fatalf("Change should include correct filename, %+v", changes[0])
	}
	if changes[0].Destination != "" {
		t.Fatalf("Change destination should be empty, %+v", changes[0])
	}
}

// TestGetUnitProperties reads the `-.mount` which should exist on all systemd
// systems and ensures that one of its properties is valid.
func TestGetUnitProperties(t *testing.T) {
	conn := setupConn(t)

	unit := "-.mount"

	info, err := conn.GetUnitProperties(unit)
	if err != nil {
		t.Fatal(err)
	}

	names := info["Wants"].([]string)

	if len(names) < 1 {
		t.Fatal("/ is unwanted")
	}

	if names[0] != "system.slice" {
		t.Fatal("unexpected wants for /")
	}

	prop, err := conn.GetUnitProperty(unit, "Wants")
	if err != nil {
		t.Fatal(err)
	}

	if prop.Name != "Wants" {
		t.Fatal("unexpected property name")
	}

	val := prop.Value.Value().([]string)
	if !reflect.DeepEqual(val, names) {
		t.Fatal("unexpected property value")
	}
}

// TestGetUnitPropertiesRejectsInvalidName attempts to get the properties for a
// unit with an invalid name. This test should be run with --test.timeout set,
// as a fail will manifest as GetUnitProperties hanging indefinitely.
func TestGetUnitPropertiesRejectsInvalidName(t *testing.T) {
	conn := setupConn(t)

	unit := "//invalid#$^/"

	_, err := conn.GetUnitProperties(unit)
	if err == nil {
		t.Fatal("Expected an error, got nil")
	}

	_, err = conn.GetUnitProperty(unit, "Wants")
	if err == nil {
		t.Fatal("Expected an error, got nil")
	}
}

// TestSetUnitProperties changes a cgroup setting on the `tmp.mount`
// which should exist on all systemd systems and ensures that the
// property was set.
func TestSetUnitProperties(t *testing.T) {
	conn := setupConn(t)

	unit := "tmp.mount"

	if err := conn.SetUnitProperties(unit, true, Property{"CPUShares", dbus.MakeVariant(uint64(1023))}); err != nil {
		t.Fatal(err)
	}

	info, err := conn.GetUnitTypeProperties(unit, "Mount")
	if err != nil {
		t.Fatal(err)
	}

	value := info["CPUShares"].(uint64)
	if value != 1023 {
		t.Fatal("CPUShares of unit is not 1023, %s", value)
	}
}

// Ensure that basic transient unit starting and stopping works.
func TestStartStopTransientUnit(t *testing.T) {
	conn := setupConn(t)

	props := []Property{
		PropExecStart([]string{"/bin/sleep", "400"}, false),
	}
	target := fmt.Sprintf("testing-transient-%d.service", rand.Int())

	// Start the unit
	job, err := conn.StartTransientUnit(target, "replace", props...)
	if err != nil {
		t.Fatal(err)
	}

	if job != "done" {
		t.Fatal("Job is not done, %v", job)
	}

	units, err := conn.ListUnits()

	var unit *UnitStatus
	for _, u := range units {
		if u.Name == target {
			unit = &u
		}
	}

	if unit == nil {
		t.Fatalf("Test unit not found in list")
	}

	if unit.ActiveState != "active" {
		t.Fatalf("Test unit not active")
	}

	// 3. Stop the unit
	job, err = conn.StopUnit(target, "replace")
	if err != nil {
		t.Fatal(err)
	}

	units, err = conn.ListUnits()

	unit = nil
	for _, u := range units {
		if u.Name == target {
			unit = &u
		}
	}

	if unit != nil {
		t.Fatalf("Test unit found in list, should be stopped")
	}
}

func TestConnJobListener(t *testing.T) {
	target := "start-stop.service"
	conn := setupConn(t)

	setupUnit(target, conn, t)

	jobSize := len(conn.jobListener.jobs)

	_, err := conn.StartUnit(target, "replace")
	if err != nil {
		t.Fatal(err)
	}

	_, err = conn.StopUnit(target, "replace")
	if err != nil {
		t.Fatal(err)
	}

	currentJobSize := len(conn.jobListener.jobs)
	if jobSize != currentJobSize {
		t.Fatal("JobListener jobs leaked")
	}
}
<commit_msg>fix(dbus): use github.com/godbus/dbus in methods_test<commit_after>/*
Copyright 2013 CoreOS Inc.

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

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

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

package dbus

import (
	"fmt"
	"math/rand"
	"os"
	"path/filepath"
	"reflect"
	"testing"

	"github.com/godbus/dbus"
)

func setupConn(t *testing.T) *Conn {
	conn, err := New()
	if err != nil {
		t.Fatal(err)
	}

	return conn
}

func setupUnit(target string, conn *Conn, t *testing.T) {
	// Blindly stop the unit in case it is running
	conn.StopUnit(target, "replace")

	// Blindly remove the symlink in case it exists
	targetRun := filepath.Join("/run/systemd/system/", target)
	err := os.Remove(targetRun)

	// 1. Enable the unit
	abs, err := filepath.Abs("../fixtures/" + target)
	if err != nil {
		t.Fatal(err)
	}

	fixture := []string{abs}

	install, changes, err := conn.EnableUnitFiles(fixture, true, true)
	if err != nil {
		t.Fatal(err)
	}

	if install != false {
		t.Fatal("Install was true")
	}

	if len(changes) < 1 {
		t.Fatalf("Expected one change, got %v", changes)
	}

	if changes[0].Filename != targetRun {
		t.Fatal("Unexpected target filename")
	}
}

// Ensure that basic unit starting and stopping works.
func TestStartStopUnit(t *testing.T) {
	target := "start-stop.service"
	conn := setupConn(t)

	setupUnit(target, conn, t)

	// 2. Start the unit
	job, err := conn.StartUnit(target, "replace")
	if err != nil {
		t.Fatal(err)
	}

	if job != "done" {
		t.Fatal("Job is not done, %v", job)
	}

	units, err := conn.ListUnits()

	var unit *UnitStatus
	for _, u := range units {
		if u.Name == target {
			unit = &u
		}
	}

	if unit == nil {
		t.Fatalf("Test unit not found in list")
	}

	if unit.ActiveState != "active" {
		t.Fatalf("Test unit not active")
	}

	// 3. Stop the unit
	job, err = conn.StopUnit(target, "replace")
	if err != nil {
		t.Fatal(err)
	}

	units, err = conn.ListUnits()

	unit = nil
	for _, u := range units {
		if u.Name == target {
			unit = &u
		}
	}

	if unit != nil {
		t.Fatalf("Test unit found in list, should be stopped")
	}
}

// Enables a unit and then immediately tears it down
func TestEnableDisableUnit(t *testing.T) {
	target := "enable-disable.service"
	conn := setupConn(t)

	setupUnit(target, conn, t)

	abs, err := filepath.Abs("../fixtures/" + target)
	if err != nil {
		t.Fatal(err)
	}

	path := filepath.Join("/run/systemd/system/", target)

	// 2. Disable the unit
	changes, err := conn.DisableUnitFiles([]string{abs}, true)
	if err != nil {
		t.Fatal(err)
	}

	if len(changes) != 1 {
		t.Fatalf("Changes should include the path, %v", changes)
	}
	if changes[0].Filename != path {
		t.Fatalf("Change should include correct filename, %+v", changes[0])
	}
	if changes[0].Destination != "" {
		t.Fatalf("Change destination should be empty, %+v", changes[0])
	}
}

// TestGetUnitProperties reads the `-.mount` which should exist on all systemd
// systems and ensures that one of its properties is valid.
func TestGetUnitProperties(t *testing.T) {
	conn := setupConn(t)

	unit := "-.mount"

	info, err := conn.GetUnitProperties(unit)
	if err != nil {
		t.Fatal(err)
	}

	names := info["Wants"].([]string)

	if len(names) < 1 {
		t.Fatal("/ is unwanted")
	}

	if names[0] != "system.slice" {
		t.Fatal("unexpected wants for /")
	}

	prop, err := conn.GetUnitProperty(unit, "Wants")
	if err != nil {
		t.Fatal(err)
	}

	if prop.Name != "Wants" {
		t.Fatal("unexpected property name")
	}

	val := prop.Value.Value().([]string)
	if !reflect.DeepEqual(val, names) {
		t.Fatal("unexpected property value")
	}
}

// TestGetUnitPropertiesRejectsInvalidName attempts to get the properties for a
// unit with an invalid name. This test should be run with --test.timeout set,
// as a fail will manifest as GetUnitProperties hanging indefinitely.
func TestGetUnitPropertiesRejectsInvalidName(t *testing.T) {
	conn := setupConn(t)

	unit := "//invalid#$^/"

	_, err := conn.GetUnitProperties(unit)
	if err == nil {
		t.Fatal("Expected an error, got nil")
	}

	_, err = conn.GetUnitProperty(unit, "Wants")
	if err == nil {
		t.Fatal("Expected an error, got nil")
	}
}

// TestSetUnitProperties changes a cgroup setting on the `tmp.mount`
// which should exist on all systemd systems and ensures that the
// property was set.
func TestSetUnitProperties(t *testing.T) {
	conn := setupConn(t)

	unit := "tmp.mount"

	if err := conn.SetUnitProperties(unit, true, Property{"CPUShares", dbus.MakeVariant(uint64(1023))}); err != nil {
		t.Fatal(err)
	}

	info, err := conn.GetUnitTypeProperties(unit, "Mount")
	if err != nil {
		t.Fatal(err)
	}

	value := info["CPUShares"].(uint64)
	if value != 1023 {
		t.Fatal("CPUShares of unit is not 1023, %s", value)
	}
}

// Ensure that basic transient unit starting and stopping works.
func TestStartStopTransientUnit(t *testing.T) {
	conn := setupConn(t)

	props := []Property{
		PropExecStart([]string{"/bin/sleep", "400"}, false),
	}
	target := fmt.Sprintf("testing-transient-%d.service", rand.Int())

	// Start the unit
	job, err := conn.StartTransientUnit(target, "replace", props...)
	if err != nil {
		t.Fatal(err)
	}

	if job != "done" {
		t.Fatal("Job is not done, %v", job)
	}

	units, err := conn.ListUnits()

	var unit *UnitStatus
	for _, u := range units {
		if u.Name == target {
			unit = &u
		}
	}

	if unit == nil {
		t.Fatalf("Test unit not found in list")
	}

	if unit.ActiveState != "active" {
		t.Fatalf("Test unit not active")
	}

	// 3. Stop the unit
	job, err = conn.StopUnit(target, "replace")
	if err != nil {
		t.Fatal(err)
	}

	units, err = conn.ListUnits()

	unit = nil
	for _, u := range units {
		if u.Name == target {
			unit = &u
		}
	}

	if unit != nil {
		t.Fatalf("Test unit found in list, should be stopped")
	}
}

func TestConnJobListener(t *testing.T) {
	target := "start-stop.service"
	conn := setupConn(t)

	setupUnit(target, conn, t)

	jobSize := len(conn.jobListener.jobs)

	_, err := conn.StartUnit(target, "replace")
	if err != nil {
		t.Fatal(err)
	}

	_, err = conn.StopUnit(target, "replace")
	if err != nil {
		t.Fatal(err)
	}

	currentJobSize := len(conn.jobListener.jobs)
	if jobSize != currentJobSize {
		t.Fatal("JobListener jobs leaked")
	}
}
<|endoftext|>


<commit_before>package builds

import "fmt"

type Status string

const (
	StatusStarted   Status = "started"
	StatusSucceeded Status = "succeeded"
	StatusFailed    Status = "failed"
	StatusErrored   Status = "errored"
)

type Build struct {
	Guid string `json:"guid"`

	Privileged bool `json:"privileged"`

	Config Config `json:"config"`

	LogsURL  string `json:"logs_url"`
	Callback string `json:"callback"`

	Inputs  []Input  `json:"inputs"`
	Outputs []Output `json:"outputs"`

	Status Status `json:"status"`
}

type Config struct {
	Image string `json:"image"`

	Env    [][2]string `json:"env"`
	Script string      `json:"script"`
}

type Input struct {
	Name string `json:"name"`

	Type   string `json:"type"`
	Source Source `json:"source,omitempty"`

	ConfigPath      string `json:"configPath"`
	DestinationPath string `json:"destinationPath"`
}

type Source []byte

func (source Source) MarshalJSON() ([]byte, error) {
	return []byte(source), nil
}

func (source *Source) UnmarshalJSON(data []byte) error {
	*source = append((*source)[0:0], data...)
	return nil
}

func (source Source) String() string {
	return string(source)
}

func (source Source) GoString() string {
	return fmt.Sprintf("builds.Source(%q)", source)
}

type Output struct {
	Name string `json:"name"`

	Type   string `json:"type"`
	Params Params `json:"params,omitempty"`

	Source Source `json:"source,omitempty"`

	SourcePath string `json:"sourcePath"`
}

type Params []byte

func (params Params) MarshalJSON() ([]byte, error) {
	return []byte(params), nil
}

func (params *Params) UnmarshalJSON(data []byte) error {
	*params = append((*params)[0:0], data...)
	return nil
}

func (params Params) String() string {
	return string(params)
}

func (params Params) GoString() string {
	return fmt.Sprintf("builds.Params(%q)", params)
}
<commit_msg>add metadata to input, make json keys consistent<commit_after>package builds

import "fmt"

type Status string

const (
	StatusStarted   Status = "started"
	StatusSucceeded Status = "succeeded"
	StatusFailed    Status = "failed"
	StatusErrored   Status = "errored"
)

type Build struct {
	Guid string `json:"guid"`

	Privileged bool `json:"privileged"`

	Config Config `json:"config"`

	LogsURL  string `json:"logs_url"`
	Callback string `json:"callback"`

	Inputs  []Input  `json:"inputs"`
	Outputs []Output `json:"outputs"`

	Status Status `json:"status"`
}

type Config struct {
	Image string `json:"image"`

	Env    [][2]string `json:"env"`
	Script string      `json:"script"`
}

type Input struct {
	Name string `json:"name"`

	Type   string `json:"type"`
	Source Source `json:"source,omitempty"`

	ConfigPath      string `json:"config_path"`
	DestinationPath string `json:"destination_path"`

	Metadata []MetadataField `json:"metadata"`
}

type MetadataField struct {
	Name  string `json:"name"`
	Value string `json:"value"`
}

type Source []byte

func (source Source) MarshalJSON() ([]byte, error) {
	return []byte(source), nil
}

func (source *Source) UnmarshalJSON(data []byte) error {
	*source = append((*source)[0:0], data...)
	return nil
}

func (source Source) String() string {
	return string(source)
}

func (source Source) GoString() string {
	return fmt.Sprintf("builds.Source(%q)", source)
}

type Output struct {
	Name string `json:"name"`

	Type   string `json:"type"`
	Params Params `json:"params,omitempty"`

	Source Source `json:"source,omitempty"`

	SourcePath string `json:"sourcePath"`
}

type Params []byte

func (params Params) MarshalJSON() ([]byte, error) {
	return []byte(params), nil
}

func (params *Params) UnmarshalJSON(data []byte) error {
	*params = append((*params)[0:0], data...)
	return nil
}

func (params Params) String() string {
	return string(params)
}

func (params Params) GoString() string {
	return fmt.Sprintf("builds.Params(%q)", params)
}
<|endoftext|>


<commit_before>// deepcopy deep copies maps, slices, etc. A standard copy will copy the
// pointers: deep copy copies the values pointed to.
//
// Only what is needed has been implemented. Could make more dynamic, at the
// cost of reflection. Either adjust as needed or create a new function.
//
// Copyright (c)2014, Joel Scoble (github.com/mohae), all rights reserved.
// License: MIT, for more details check the included LICENSE.txt.
package deepcopy

import (
	"reflect"
)

// InterfaceToSliceOfStrings takes an interface that is either a slice of
// strings or a string and returns a deep copy of it as a slice of strings.
func InterfaceToSliceOfStrings(v interface{}) []string {
	if v == nil {
		return nil
	}
	var sl []string
	switch reflect.TypeOf(v).Kind() {
	case reflect.Slice:
		s := reflect.ValueOf(v)
		sl = make([]string, s.Len(), s.Len())
		for i := 0; i < s.Len(); i++ {
			sl[i] = s.Index(i).Interface().(string)
		}
	case reflect.String:
		sl = append(sl, reflect.ValueOf(v).Interface().(string))
	default:
		return nil
	}
	return sl
}

// InterfaceToSliceOfInts takes an interface that is a slice of ints and returns
// a deep copy of it as a slice of strings. An error is returned if the
// interface is not a slice of strings.
func InterfaceToSliceOfInts(v interface{}) []int {
	if v == nil {
		return nil
	}
	var sl []int
	switch reflect.TypeOf(v).Kind() {
	case reflect.Slice:
		s := reflect.ValueOf(v)
		sl = make([]int, s.Len(), s.Len())
		for i := 0; i < s.Len(); i++ {
			sl[i] = s.Index(i).Interface().(int)
		}
	case reflect.Int:
		sl = append(sl, reflect.ValueOf(v).Interface().(int))
	default:
		return nil
	}
	return sl
}

// Iface recursively deep copies an interface{}
func Iface(iface interface{}) interface{} {
	if iface == nil {
		return nil
	}
	// Make the interface a reflect.Value
	original := reflect.ValueOf(iface)
	// Make a copy of the same type as the original.
	cpy := reflect.New(original.Type()).Elem()
	// Recursively copy the original.
	copyRecursive(original, cpy)
	// Return theb copy as an interface.
	return copy.Interface()
}

// copyRecursive does the actual copying of the interface. It currently has
// limited support for what it can handle. Add as needed.
func copyRecursive(original, cpy reflect.Value) {
	// handle according to original's Kind
	switch original.Kind() {
	case reflect.Ptr:
		// Get the actual value being pointed to.
		originalValue := original.Elem()
		// if  it isn't valid, return.
		if !originalValue.IsValid() {
			return
		}
		cpy.Set(reflect.New(originalValue.Type()))
		copyRecursive(originalValue, cpy.Elem())
	case reflect.Interface:
		// Get the value for the interface, not the pointer.
		originalValue := original.Elem()
		if !originalValue.IsValid() {
			return
		}
		// Get the value by calling Elem().
		copyValue := reflect.New(originalValue.Type()).Elem()
		copyRecursive(originalValue, copyValue)
		cpy.Set(copyValue)
	case reflect.Struct:
		// Go through each field of the struct and copy it.
		for i := 0; i < original.NumField(); i++ {
			if cpy.Field(i).CanSet() {
				copyRecursive(original.Field(i), cpy.Field(i))
			}
		}
	case reflect.Slice:
		// Make a new slice and copy each element.
		cpy.Set(reflect.MakeSlice(original.Type(), original.Len(), original.Cap()))
		for i := 0; i < original.Len(); i++ {
			copyRecursive(original.Index(i), cpy.Index(i))
		}
	case reflect.Map:
		cpy.Set(reflect.MakeMap(original.Type()))
		for _, key := range original.MapKeys() {
			originalValue := original.MapIndex(key)
			copyValue := reflect.New(originalValue.Type()).Elem()
			copyRecursive(originalValue, copyValue)
			cpy.SetMapIndex(key, copyValue)
		}
	// Set the actual values from here on.
	case reflect.String:
		cpy.SetString(original.Interface().(string))
	case reflect.Int:
		cpy.SetInt(int64(original.Interface().(int)))
	case reflect.Bool:
		cpy.SetBool(original.Interface().(bool))
	case reflect.Float64:
		cpy.SetFloat(original.Interface().(float64))

	default:
		cpy.Set(original)
	}
}
<commit_msg>fix var name error in deepcopy<commit_after>// deepcopy deep copies maps, slices, etc. A standard copy will copy the
// pointers: deep copy copies the values pointed to.
//
// Only what is needed has been implemented. Could make more dynamic, at the
// cost of reflection. Either adjust as needed or create a new function.
//
// Copyright (c)2014, Joel Scoble (github.com/mohae), all rights reserved.
// License: MIT, for more details check the included LICENSE.txt.
package deepcopy

import (
	"reflect"
)

// InterfaceToSliceOfStrings takes an interface that is either a slice of
// strings or a string and returns a deep copy of it as a slice of strings.
func InterfaceToSliceOfStrings(v interface{}) []string {
	if v == nil {
		return nil
	}
	var sl []string
	switch reflect.TypeOf(v).Kind() {
	case reflect.Slice:
		s := reflect.ValueOf(v)
		sl = make([]string, s.Len(), s.Len())
		for i := 0; i < s.Len(); i++ {
			sl[i] = s.Index(i).Interface().(string)
		}
	case reflect.String:
		sl = append(sl, reflect.ValueOf(v).Interface().(string))
	default:
		return nil
	}
	return sl
}

// InterfaceToSliceOfInts takes an interface that is a slice of ints and returns
// a deep copy of it as a slice of strings. An error is returned if the
// interface is not a slice of strings.
func InterfaceToSliceOfInts(v interface{}) []int {
	if v == nil {
		return nil
	}
	var sl []int
	switch reflect.TypeOf(v).Kind() {
	case reflect.Slice:
		s := reflect.ValueOf(v)
		sl = make([]int, s.Len(), s.Len())
		for i := 0; i < s.Len(); i++ {
			sl[i] = s.Index(i).Interface().(int)
		}
	case reflect.Int:
		sl = append(sl, reflect.ValueOf(v).Interface().(int))
	default:
		return nil
	}
	return sl
}

// Iface recursively deep copies an interface{}
func Iface(iface interface{}) interface{} {
	if iface == nil {
		return nil
	}
	// Make the interface a reflect.Value
	original := reflect.ValueOf(iface)
	// Make a copy of the same type as the original.
	cpy := reflect.New(original.Type()).Elem()
	// Recursively copy the original.
	copyRecursive(original, cpy)
	// Return theb copy as an interface.
	return cpy.Interface()
}

// copyRecursive does the actual copying of the interface. It currently has
// limited support for what it can handle. Add as needed.
func copyRecursive(original, cpy reflect.Value) {
	// handle according to original's Kind
	switch original.Kind() {
	case reflect.Ptr:
		// Get the actual value being pointed to.
		originalValue := original.Elem()
		// if  it isn't valid, return.
		if !originalValue.IsValid() {
			return
		}
		cpy.Set(reflect.New(originalValue.Type()))
		copyRecursive(originalValue, cpy.Elem())
	case reflect.Interface:
		// Get the value for the interface, not the pointer.
		originalValue := original.Elem()
		if !originalValue.IsValid() {
			return
		}
		// Get the value by calling Elem().
		copyValue := reflect.New(originalValue.Type()).Elem()
		copyRecursive(originalValue, copyValue)
		cpy.Set(copyValue)
	case reflect.Struct:
		// Go through each field of the struct and copy it.
		for i := 0; i < original.NumField(); i++ {
			if cpy.Field(i).CanSet() {
				copyRecursive(original.Field(i), cpy.Field(i))
			}
		}
	case reflect.Slice:
		// Make a new slice and copy each element.
		cpy.Set(reflect.MakeSlice(original.Type(), original.Len(), original.Cap()))
		for i := 0; i < original.Len(); i++ {
			copyRecursive(original.Index(i), cpy.Index(i))
		}
	case reflect.Map:
		cpy.Set(reflect.MakeMap(original.Type()))
		for _, key := range original.MapKeys() {
			originalValue := original.MapIndex(key)
			copyValue := reflect.New(originalValue.Type()).Elem()
			copyRecursive(originalValue, copyValue)
			cpy.SetMapIndex(key, copyValue)
		}
	// Set the actual values from here on.
	case reflect.String:
		cpy.SetString(original.Interface().(string))
	case reflect.Int:
		cpy.SetInt(int64(original.Interface().(int)))
	case reflect.Bool:
		cpy.SetBool(original.Interface().(bool))
	case reflect.Float64:
		cpy.SetFloat(original.Interface().(float64))

	default:
		cpy.Set(original)
	}
}
<|endoftext|>


<commit_before>package disruptor

import "runtime"

type DefaultSequencer struct {
	written  *Sequence // the ring buffer has been written up to this sequence
	upstream Barrier   // all of the readers have advanced up to this sequence
	capacity int64
	previous int64
	gate     int64
}

func NewSequencer(written *Sequence, upstream Barrier, capacity int64) *DefaultSequencer {
	return &DefaultSequencer{
		upstream: upstream,
		written:  written,
		capacity: capacity,
		previous: defaultSequenceValue,
		gate:     defaultSequenceValue,
	}
}

func (this *DefaultSequencer) Reserve(count int64) int64 {
	this.previous += count

	for spin := int64(0); this.previous-this.capacity > this.gate; spin++ {
		if spin&SpinMask == 0 {
			runtime.Gosched() // LockSupport.parkNanos(1L); http://bit.ly/1xiDINZ
		}

		this.gate = this.upstream.Load()

	}

	return this.previous
}

func (this *DefaultSequencer) Commit(_, upper int64) { this.written.Store(upper) }

const SpinMask = 1024*16 - 1 // arbitrary; we'll want to experiment with different values
<commit_msg>Min reservation size.<commit_after>package disruptor

import "runtime"

type DefaultSequencer struct {
	written  *Sequence // the ring buffer has been written up to this sequence
	upstream Barrier   // all of the readers have advanced up to this sequence
	capacity int64
	previous int64
	gate     int64
}

func NewSequencer(written *Sequence, upstream Barrier, capacity int64) *DefaultSequencer {
	return &DefaultSequencer{
		upstream: upstream,
		written:  written,
		capacity: capacity,
		previous: defaultSequenceValue,
		gate:     defaultSequenceValue,
	}
}

func (this *DefaultSequencer) Reserve(count int64) int64 {
	if count <= 0 {
		panic(ErrMinimumReservationSize)
	}

	this.previous += count
	for spin := int64(0); this.previous-this.capacity > this.gate; spin++ {
		if spin&SpinMask == 0 {
			runtime.Gosched() // LockSupport.parkNanos(1L); http://bit.ly/1xiDINZ
		}

		this.gate = this.upstream.Load()
	}
	return this.previous
}

func (this *DefaultSequencer) Commit(_, upper int64) { this.written.Store(upper) }

const SpinMask = 1024*16 - 1 // arbitrary; we'll want to experiment with different values
<|endoftext|>


<commit_before>package delegate

import (
	"errors"
	"time"

	"github.com/gonfire/oauth2"
)

// ErrNotFound can be returned by the delegate to indicate that the requested
// client, resource owner, authorization code or refresh token has not been found,
var ErrNotFound = errors.New("not found")

// ErrMalformed can be returned by the delegate to indicate that the provided
// authorization code or refresh token is malformed.
var ErrMalformed = errors.New("malformed")

// ErrRejected can be returned by the delegate to indicate that the provided
// scope has been rejected and the current request should be denied.
var ErrRejected = errors.New("rejected")

type Delegate interface {
	LookupClient(string) (Client, error)
	LookupResourceOwner(string) (ResourceOwner, error)

	GrantScope(Client, ResourceOwner, oauth2.Scope) (oauth2.Scope, error)
	IssueAccessToken(Client, ResourceOwner, oauth2.Scope) (string, int, error)
}

type AuthorizationDelegate interface {
	Delegate

	ParseConsent(r *oauth2.AuthorizationRequest) (string, string, oauth2.Scope, error)
}

type AuthorizationCodeDelegate interface {
	AuthorizationDelegate

	LookupAuthorizationCode(string) (AuthorizationCode, error)
	IssueAuthorizationCode(Client, ResourceOwner, oauth2.Scope, string) (string, error)
	RemoveAuthorizationCode(string) error
}

type RefreshTokenDelegate interface {
	Delegate

	LookupRefreshToken(string) (RefreshToken, error)
	IssueRefreshToken(Client, ResourceOwner, oauth2.Scope) (string, error)
	RemoveRefreshToken(string) error
}

type Client interface {
	ID() string
	Confidential() bool
	ValidSecret(string) bool
	ValidRedirectURI(string) bool
}

type ResourceOwner interface {
	ID() string
	ValidSecret(string) bool
}

type AuthorizationCode interface {
	ClientID() string
	ResourceOwnerID() string
	ExpiresAt() time.Time
	Scope() oauth2.Scope
	RedirectURI() string
}

type RefreshToken interface {
	ClientID() string
	ResourceOwnerID() string
	ExpiresAt() time.Time
	Scope() oauth2.Scope
}
<commit_msg>use common credential interface<commit_after>package delegate

import (
	"errors"
	"time"

	"github.com/gonfire/oauth2"
)

// ErrNotFound can be returned by the delegate to indicate that the requested
// client, resource owner, authorization code or refresh token has not been found,
var ErrNotFound = errors.New("not found")

// ErrMalformed can be returned by the delegate to indicate that the provided
// authorization code or refresh token is malformed.
var ErrMalformed = errors.New("malformed")

// ErrRejected can be returned by the delegate to indicate that the provided
// scope has been rejected and the current request should be denied.
var ErrRejected = errors.New("rejected")

type Delegate interface {
	LookupClient(string) (Client, error)
	LookupResourceOwner(string) (ResourceOwner, error)

	GrantScope(Client, ResourceOwner, oauth2.Scope) (oauth2.Scope, error)
	IssueAccessToken(Client, ResourceOwner, oauth2.Scope) (string, int, error)
}

type AuthorizationDelegate interface {
	Delegate

	ParseConsent(r *oauth2.AuthorizationRequest) (string, string, oauth2.Scope, error)
}

type AuthorizationCodeDelegate interface {
	AuthorizationDelegate

	LookupAuthorizationCode(string) (AuthorizationCode, error)
	IssueAuthorizationCode(Client, ResourceOwner, oauth2.Scope, string) (string, error)
	RemoveAuthorizationCode(string) error
}

type RefreshTokenDelegate interface {
	Delegate

	LookupRefreshToken(string) (RefreshToken, error)
	IssueRefreshToken(Client, ResourceOwner, oauth2.Scope) (string, error)
	RemoveRefreshToken(string) error
}

type Client interface {
	ID() string
	Confidential() bool
	ValidSecret(string) bool
	ValidRedirectURI(string) bool
}

type ResourceOwner interface {
	ID() string
	ValidSecret(string) bool
}

type Credential interface {
	ClientID() string
	ResourceOwnerID() string
	ExpiresAt() time.Time
	Scope() oauth2.Scope
}

type AuthorizationCode interface {
	Credential

	RedirectURI() string
}

type RefreshToken interface {
	Credential
}
<|endoftext|>


<commit_before>package http

import (
	"fmt"
	"io"
	"log"
	gohttp "net/http"
	"os"
	"path/filepath"
	"time"
)

type Fetcher struct{}

func (f *Fetcher) Fetch(url string, destination string) error {
	if _, err := os.Stat(destination); os.IsNotExist(err) {
		start := time.Now()

		resp, err := gohttp.Get(url)

		if err != nil || resp.StatusCode != 200 {
			log.Printf("Unable to download image: %s, status code: %d", url, resp.StatusCode)
			log.Println(err)
			return fmt.Errorf("Unable to download image: %s, status code: %d", url, resp.StatusCode)
		}
		log.Printf("Downloaded from %s with code %d", url, resp.StatusCode)
		defer resp.Body.Close()

		dir := filepath.Dir(destination)
		os.MkdirAll(dir, 0700)

		out, err := os.Create(destination)
		defer out.Close()
		if err != nil {
			log.Printf("Unable to create file: %s", destination)
			log.Println(err)
			return fmt.Errorf("Unable to create file: %s", destination)
		}

		io.Copy(out, resp.Body)
		log.Printf("Took %s to download image: %s", time.Since(start), destination)
	}
	return nil
}
<commit_msg>Handle error downloading file and bad status code separately<commit_after>package http

import (
	"fmt"
	"io"
	"log"
	gohttp "net/http"
	"os"
	"path/filepath"
	"time"
)

type Fetcher struct{}

func (f *Fetcher) Fetch(url string, destination string) error {
	if _, err := os.Stat(destination); os.IsNotExist(err) {
		start := time.Now()

		resp, err := gohttp.Get(url)

		if err != nil {
			return err
		}

		if resp.StatusCode != 200 {
			return fmt.Errorf("Unable to download image: %s, status code: %d", url, resp.StatusCode)
		}

		log.Printf("Downloaded from %s with code %d", url, resp.StatusCode)

		defer resp.Body.Close()

		dir := filepath.Dir(destination)
		os.MkdirAll(dir, 0700)

		out, err := os.Create(destination)
		defer out.Close()
		if err != nil {
			log.Printf("Unable to create file: %s", destination)
			log.Println(err)
			return fmt.Errorf("Unable to create file: %s", destination)
		}

		io.Copy(out, resp.Body)
		log.Printf("Took %s to download image: %s", time.Since(start), destination)
	}
	return nil
}
<|endoftext|>


<commit_before>package yarn_client

import (
	"github.com/hortonworks/gohadoop/hadoop_yarn"
	yarn_conf "github.com/hortonworks/gohadoop/hadoop_yarn/conf"
	"log"
	"sync"
	"time"
)

type AMRMClient struct {
	applicationAttemptId *hadoop_yarn.ApplicationAttemptIdProto
	client               hadoop_yarn.ApplicationMasterProtocolService
	responseId           int32
	conf                 yarn_conf.YarnConfiguration
}

type resource_to_request struct {
	capability    *hadoop_yarn.ResourceProto
	numContainers int32
}

var allocationRequests = struct {
	sync.RWMutex
	resourceRequests map[int32]map[string]*resource_to_request
	releaseRequests  map[*hadoop_yarn.ContainerIdProto]bool
}{resourceRequests: make(map[int32]map[string]*resource_to_request),
	releaseRequests: make(map[*hadoop_yarn.ContainerIdProto]bool)}

func CreateAMRMClient(conf yarn_conf.YarnConfiguration, applicationAttemptId *hadoop_yarn.ApplicationAttemptIdProto) (*AMRMClient, error) {
	c, err := hadoop_yarn.DialApplicationMasterProtocolService(conf)
	return &AMRMClient{applicationAttemptId: applicationAttemptId, client: c, conf: conf}, err
}

func (c *AMRMClient) RegisterApplicationMaster(host string, port int32, url string) error {
	request := hadoop_yarn.RegisterApplicationMasterRequestProto{Host: &host, RpcPort: &port, TrackingUrl: &url, ApplicationAttemptId: c.applicationAttemptId}
	response := hadoop_yarn.RegisterApplicationMasterResponseProto{}
	err := c.client.RegisterApplicationMaster(&request, &response)

	if err != nil {
		return err
	}

	go c.periodicPingWithEmptyAllocate()
	return nil
}

func (c *AMRMClient) FinishApplicationMaster(finalStatus *hadoop_yarn.FinalApplicationStatusProto, message string, url string) error {
	request := hadoop_yarn.FinishApplicationMasterRequestProto{FinalApplicationStatus: finalStatus, Diagnostics: &message, TrackingUrl: &url, ApplicationAttemptId: c.applicationAttemptId}
	response := hadoop_yarn.FinishApplicationMasterResponseProto{}
	return c.client.FinishApplicationMaster(&request, &response)
}

func (c *AMRMClient) ReleaseAssignedContainer(containerId *hadoop_yarn.ContainerIdProto) {
	if containerId != nil {
		allocationRequests.Lock()
		allocationRequests.releaseRequests[containerId] = true
		allocationRequests.Unlock()
	}
}

func (c *AMRMClient) AddRequest(priority int32, resourceName string, capability *hadoop_yarn.ResourceProto, numContainers int32) error {
	allocationRequests.Lock()
	existingResourceRequests, exists := allocationRequests.resourceRequests[priority]
	if !exists {
		existingResourceRequests = make(map[string]*resource_to_request)
		allocationRequests.resourceRequests[priority] = existingResourceRequests
	}
	request, exists := existingResourceRequests[resourceName]
	if !exists {
		request = &resource_to_request{capability: capability, numContainers: numContainers}
		existingResourceRequests[resourceName] = request
	} else {
		request.numContainers += numContainers
	}
	allocationRequests.Unlock()

	return nil
}

func (c *AMRMClient) Allocate() (*hadoop_yarn.AllocateResponseProto, error) {
	// Increment responseId
	c.responseId++
	log.Println("ResponseId: ", c.responseId)

	asks := []*hadoop_yarn.ResourceRequestProto{}

	// Set up resource-requests
	allocationRequests.Lock()
	for priority, requests := range allocationRequests.resourceRequests {
		for host, request := range requests {
			log.Println("priority: ", priority)
			log.Println("host: ", host)
			log.Println("request: ", request)

			resourceRequest := hadoop_yarn.ResourceRequestProto{Priority: &hadoop_yarn.PriorityProto{Priority: &priority}, ResourceName: &host, Capability: request.capability, NumContainers: &request.numContainers}
			asks = append(asks, &resourceRequest)
		}
	}

	var releases []*hadoop_yarn.ContainerIdProto
	for containerId, _ := range allocationRequests.releaseRequests {
		releases = append(releases, containerId)
	}

	log.Printf("AMRMClient.Allocate #asks: %d #releases: %d", len(asks), len(releases))

	// Clear
	allocationRequests.resourceRequests = make(map[int32]map[string]*resource_to_request)
	allocationRequests.releaseRequests = make(map[*hadoop_yarn.ContainerIdProto]bool)
	allocationRequests.Unlock()

	request := hadoop_yarn.AllocateRequestProto{ApplicationAttemptId: c.applicationAttemptId, Ask: asks, Release: releases, ResponseId: &c.responseId}
	response := hadoop_yarn.AllocateResponseProto{}
	err := c.client.Allocate(&request, &response)
	return &response, err
}

//We need to periodically "ping" the Resource Manager in order to ensure the AM isn't timed out.
func (c *AMRMClient) periodicPingWithEmptyAllocate() {
	sleepIntervalMs, err := c.conf.GetInt(yarn_conf.RM_AM_EXPIRY_INTERVAL_MS, yarn_conf.DEFAULT_RM_AM_EXPIRY_INTERVAL_MS)

	if err != nil {
		log.Println("failed to read expiry configuration. ping routine will NOT run")
		return
	}

	//keep the sleep interval shorter than the expiry timeout
	sleepIntervalMs = sleepIntervalMs / 2

	for {
		log.Println("ping with empty allocate.")

		//lock allocation requests
		request := hadoop_yarn.AllocateRequestProto{ApplicationAttemptId: c.applicationAttemptId}
		response := hadoop_yarn.AllocateResponseProto{}
		err := c.client.Allocate(&request, &response)

		if err == nil {
			log.Println("allocate response(ping): ", response)
		} else {
			log.Println("ping allocate failed! Error: ", err)
		}

		time.Sleep(time.Duration(sleepIntervalMs) * time.Millisecond)
	}
}
<commit_msg>cleaned up synchronization to avoid operation overlap<commit_after>package yarn_client

import (
	"github.com/hortonworks/gohadoop/hadoop_yarn"
	yarn_conf "github.com/hortonworks/gohadoop/hadoop_yarn/conf"
	"log"
	"sync"
	"time"
)

type AMRMClient struct {
	applicationAttemptId *hadoop_yarn.ApplicationAttemptIdProto
	client               hadoop_yarn.ApplicationMasterProtocolService
	responseId           int32
	conf                 yarn_conf.YarnConfiguration
}

type resource_to_request struct {
	capability    *hadoop_yarn.ResourceProto
	numContainers int32
}

var allocationRequests = struct {
	sync.RWMutex
	resourceRequests map[int32]map[string]*resource_to_request
	releaseRequests  map[*hadoop_yarn.ContainerIdProto]bool
}{resourceRequests: make(map[int32]map[string]*resource_to_request),
	releaseRequests: make(map[*hadoop_yarn.ContainerIdProto]bool)}

func CreateAMRMClient(conf yarn_conf.YarnConfiguration, applicationAttemptId *hadoop_yarn.ApplicationAttemptIdProto) (*AMRMClient, error) {
	c, err := hadoop_yarn.DialApplicationMasterProtocolService(conf)
	return &AMRMClient{applicationAttemptId: applicationAttemptId, client: c, conf: conf}, err
}

func (c *AMRMClient) RegisterApplicationMaster(host string, port int32, url string) error {
	request := hadoop_yarn.RegisterApplicationMasterRequestProto{Host: &host, RpcPort: &port, TrackingUrl: &url, ApplicationAttemptId: c.applicationAttemptId}
	response := hadoop_yarn.RegisterApplicationMasterResponseProto{}
	err := c.client.RegisterApplicationMaster(&request, &response)

	if err != nil {
		return err
	}

	go c.periodicPingWithEmptyAllocate()
	return nil
}

func (c *AMRMClient) FinishApplicationMaster(finalStatus *hadoop_yarn.FinalApplicationStatusProto, message string, url string) error {
	request := hadoop_yarn.FinishApplicationMasterRequestProto{FinalApplicationStatus: finalStatus, Diagnostics: &message, TrackingUrl: &url, ApplicationAttemptId: c.applicationAttemptId}
	response := hadoop_yarn.FinishApplicationMasterResponseProto{}
	return c.client.FinishApplicationMaster(&request, &response)
}

func (c *AMRMClient) ReleaseAssignedContainer(containerId *hadoop_yarn.ContainerIdProto) {
	if containerId != nil {
		allocationRequests.Lock()
    defer allocationRequests.Unlock()
		allocationRequests.releaseRequests[containerId] = true
	}
}

func (c *AMRMClient) AddRequest(priority int32, resourceName string, capability *hadoop_yarn.ResourceProto, numContainers int32) error {
	allocationRequests.Lock()
	existingResourceRequests, exists := allocationRequests.resourceRequests[priority]
	if !exists {
		existingResourceRequests = make(map[string]*resource_to_request)
		allocationRequests.resourceRequests[priority] = existingResourceRequests
	}
	request, exists := existingResourceRequests[resourceName]
	if !exists {
		request = &resource_to_request{capability: capability, numContainers: numContainers}
		existingResourceRequests[resourceName] = request
	} else {
		request.numContainers += numContainers
	}
	allocationRequests.Unlock()

	return nil
}

func (c *AMRMClient) Allocate() (*hadoop_yarn.AllocateResponseProto, error) {
  allocationRequests.Lock()
  defer allocationRequests.Unlock()

	// Increment responseId
	c.responseId++
	log.Println("ResponseId: ", c.responseId)

	asks := []*hadoop_yarn.ResourceRequestProto{}

	// Set up resource-requests
	for priority, requests := range allocationRequests.resourceRequests {
		for host, request := range requests {
			log.Println("priority: ", priority)
			log.Println("host: ", host)
			log.Println("request: ", request)

			resourceRequest := hadoop_yarn.ResourceRequestProto{Priority: &hadoop_yarn.PriorityProto{Priority: &priority}, ResourceName: &host, Capability: request.capability, NumContainers: &request.numContainers}
			asks = append(asks, &resourceRequest)
		}
	}

	var releases []*hadoop_yarn.ContainerIdProto
	for containerId, _ := range allocationRequests.releaseRequests {
		releases = append(releases, containerId)
	}

	log.Printf("AMRMClient.Allocate #asks: %d #releases: %d", len(asks), len(releases))

	// Clear
	allocationRequests.resourceRequests = make(map[int32]map[string]*resource_to_request)
	allocationRequests.releaseRequests = make(map[*hadoop_yarn.ContainerIdProto]bool)

	request := hadoop_yarn.AllocateRequestProto{ApplicationAttemptId: c.applicationAttemptId, Ask: asks, Release: releases, ResponseId: &c.responseId}
	response := hadoop_yarn.AllocateResponseProto{}
	err := c.client.Allocate(&request, &response)

	return &response, err
}

//We need to periodically "ping" the Resource Manager in order to ensure the AM isn't timed out.
func (c *AMRMClient) periodicPingWithEmptyAllocate() {
	sleepIntervalMs, err := c.conf.GetInt(yarn_conf.RM_AM_EXPIRY_INTERVAL_MS, yarn_conf.DEFAULT_RM_AM_EXPIRY_INTERVAL_MS)

	if err != nil {
		log.Println("failed to read expiry configuration. ping routine will NOT run")
		return
	}

	//keep the sleep interval shorter than the expiry timeout
	sleepIntervalMs = sleepIntervalMs / 2

	for {
		log.Println("ping with empty allocate.")
	  c.sendPingRequest()
		time.Sleep(time.Duration(sleepIntervalMs) * time.Millisecond)
	}
}

func(c * AMRMClient) sendPingRequest() {
  //ensure no other operations are in progress
  allocationRequests.Lock()
  defer allocationRequests.Unlock()

  request := hadoop_yarn.AllocateRequestProto{ApplicationAttemptId: c.applicationAttemptId}
  response := hadoop_yarn.AllocateResponseProto{}
  err := c.client.Allocate(&request, &response)

  if err == nil {
    log.Println("allocate response(ping): ", response)
  } else {
    log.Println("ping allocate failed! Error: ", err)
  }
}

<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"fmt"
	"log"
	// Register hdb driver.
	_ "github.com/SAP/go-hdb/driver"
	// cli helper
	"github.com/mkideal/cli"
	// ini config
	"github.com/go-ini/ini"
	// internal
	"github.com/morxs/go-hana/utils"
)

type argT struct {
	cli.Helper
	ArgFile   string `cli:"*f" usage:"Bulk data (CSV, semicolon separated) file"`
	ArgConfig string `cli:"c" usage:"Custom config file" dft:"config.ini"`
}

const (
	driverName = "hdb"
)

func main() {
	cli.Run(new(argT), func(ctx *cli.Context) error {
		argv := ctx.Argv().(*argT)

		// read config file
		utils.WriteMsg("READ CONFIG")
		iniCfg, err := ini.Load(argv.ArgConfig)
		if err != nil {
			utils.WriteMsg("CONFIG")
			log.Fatal(err)
		}
		iniSection := iniCfg.Section("server")
		/*
			fmt.Println(iniSection)
			fmt.Println(iniSection.KeyStrings())
			fmt.Println(iniSection.Key("uid").String())
			fmt.Println(iniSection.GetKey("uid"))
		*/
		iniKeyUsername := iniSection.Key("uid").String()
		iniKeyPassword := iniSection.Key("pwd").String()
		iniKeyHost := iniSection.Key("host").String()
		// iniKeyHost = "10.11.1.53"
		iniKeyPort := iniSection.Key("port").String()
		hdbDsn := "hdb://" + iniKeyUsername + ":" + iniKeyPassword + "@" + iniKeyHost + ":" + iniKeyPort

		utils.WriteMsg("OPEN HDB")
		//fmt.Print("OPENDB...")
		db, err := sql.Open(driverName, hdbDsn)
		if err != nil {
			//fmt.Print("OPENDB")
			log.Fatal(err)
		}
		defer db.Close()

		if err := db.Ping(); err != nil {
			log.Fatal(err)
		}

		stmt, err := db.Prepare("bulk insert into Z_WILMAR_CONSODB.GL_CONSOL_PACK_MAP values (?,?,?,?,?,?,?,?,?,?,?,?,?,?)")

		// baca file
		rec, _ := utils.ReadCsv(argv.ArgFile, ';')

		for i := 0; i < len(rec); i++ {
			/*
				if rec[i][10] == "" {
					rec[i][10] = sql.NullString{}
				}
			*/
			if _, err := stmt.Exec(
				rec[i][0],
				rec[i][1],
				rec[i][2],
				rec[i][3],
				rec[i][4],
				rec[i][5],
				rec[i][6],
				NewNullString(rec[i][7]),
				NewNullString(rec[i][8]),
				rec[i][9],
				NewNullString(rec[i][10]),
				rec[i][11],
				rec[i][12],
				rec[i][13]); err != nil {
				log.Fatal(err)
			}
		}

		if _, err := stmt.Exec(); err != nil {
			log.Fatal(err)
		}
		fmt.Println("DONE")
		return nil
	})
}

func NewNullString(s string) sql.NullString {
	if len(s) == 0 {
		return sql.NullString{}
	}
	return sql.NullString{
		String: s,
		Valid:  true,
	}
}
<commit_msg>update upload consodb code<commit_after>package main

import (
	"database/sql"
	"fmt"
	"log"
	// Register hdb driver.
	_ "github.com/SAP/go-hdb/driver"
	// cli helper
	"github.com/mkideal/cli"
	// ini config
	"github.com/go-ini/ini"
	// internal
	"github.com/morxs/go-hana/utils"
)

type argT struct {
	cli.Helper
	ArgFile   string `cli:"*f" usage:"Bulk data (CSV, semicolon separated) file"`
	ArgConfig string `cli:"c" usage:"Custom config file" dft:"config.ini"`
}

const (
	driverName = "hdb"
)

func main() {
	cli.Run(new(argT), func(ctx *cli.Context) error {
		argv := ctx.Argv().(*argT)

		// read config file
		utils.WriteMsg("READ CONFIG")
		iniCfg, err := ini.Load(argv.ArgConfig)
		if err != nil {
			utils.WriteMsg("CONFIG")
			log.Fatal(err)
		}
		iniSection := iniCfg.Section("server")
		/*
			fmt.Println(iniSection)
			fmt.Println(iniSection.KeyStrings())
			fmt.Println(iniSection.Key("uid").String())
			fmt.Println(iniSection.GetKey("uid"))
		*/
		iniKeyUsername := iniSection.Key("uid").String()
		iniKeyPassword := iniSection.Key("pwd").String()
		iniKeyHost := iniSection.Key("host").String()
		iniKeyHost = "10.11.1.53"
		iniKeyPort := iniSection.Key("port").String()
		hdbDsn := "hdb://" + iniKeyUsername + ":" + iniKeyPassword + "@" + iniKeyHost + ":" + iniKeyPort

		utils.WriteMsg("OPEN HDB")
		//fmt.Print("OPENDB...")
		db, err := sql.Open(driverName, hdbDsn)
		if err != nil {
			//fmt.Print("OPENDB")
			log.Fatal(err)
		}
		defer db.Close()

		if err := db.Ping(); err != nil {
			log.Fatal(err)
		}

		stmt, err := db.Prepare("bulk insert into Z_WILMAR_CONSODB.GL_CONSOL_PACK_MAP values (?,?,?,?,?,?,?,?,?,?,?,?,?,?)")

		// baca file
		rec, _ := utils.ReadCsv(argv.ArgFile, ';')

		for i := 0; i < len(rec); i++ {
			/*
				if rec[i][10] == "" {
					rec[i][10] = sql.NullString{}
				}
			*/
			if _, err := stmt.Exec(
				rec[i][0],
				rec[i][1],
				rec[i][2],
				rec[i][3],
				rec[i][4],
				rec[i][5],
				rec[i][6],
				NewNullString(rec[i][7]),
				NewNullString(rec[i][8]),
				rec[i][9],
				NewNullString(rec[i][10]),
				rec[i][11],
				rec[i][12],
				rec[i][13]); err != nil {
				log.Fatal(err)
			}
		}

		if _, err := stmt.Exec(); err != nil {
			log.Fatal(err)
		}
		fmt.Println("DONE")
		return nil
	})
}

func NewNullString(s string) sql.NullString {
	if len(s) == 0 {
		return sql.NullString{}
	}
	return sql.NullString{
		String: s,
		Valid:  true,
	}
}
<|endoftext|>


<commit_before>// Copyright 2012 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 flate

import (
	"io/ioutil"
	"runtime"
	"testing"
)

func benchmarkEncoder(b *testing.B, testfile, level, n int) {
	b.StopTimer()
	b.SetBytes(int64(n))
	buf0, err := ioutil.ReadFile(testfiles[testfile])
	if err != nil {
		b.Fatal(err)
	}
	if len(buf0) == 0 {
		b.Fatalf("test file %q has no data", testfiles[testfile])
	}
	buf1 := make([]byte, n)
	for i := 0; i < n; i += len(buf0) {
		if len(buf0) > n-i {
			buf0 = buf0[:n-i]
		}
		copy(buf1[i:], buf0)
	}
	buf0 = nil
	runtime.GC()
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		w, err := NewWriter(ioutil.Discard, level)
		if err != nil {
			b.Fatal(err)
		}
		w.Write(buf1)
		w.Close()
	}
}

func BenchmarkEncodeDigitsSpeed1e4(b *testing.B)    { benchmarkEncoder(b, digits, speed, 1e4) }
func BenchmarkEncodeDigitsSpeed1e5(b *testing.B)    { benchmarkEncoder(b, digits, speed, 1e5) }
func BenchmarkEncodeDigitsSpeed1e6(b *testing.B)    { benchmarkEncoder(b, digits, speed, 1e6) }
func BenchmarkEncodeDigitsDefault1e4(b *testing.B)  { benchmarkEncoder(b, digits, default_, 1e4) }
func BenchmarkEncodeDigitsDefault1e5(b *testing.B)  { benchmarkEncoder(b, digits, default_, 1e5) }
func BenchmarkEncodeDigitsDefault1e6(b *testing.B)  { benchmarkEncoder(b, digits, default_, 1e6) }
func BenchmarkEncodeDigitsCompress1e4(b *testing.B) { benchmarkEncoder(b, digits, compress, 1e4) }
func BenchmarkEncodeDigitsCompress1e5(b *testing.B) { benchmarkEncoder(b, digits, compress, 1e5) }
func BenchmarkEncodeDigitsCompress1e6(b *testing.B) { benchmarkEncoder(b, digits, compress, 1e6) }
func BenchmarkEncodeTwainSpeed1e4(b *testing.B)     { benchmarkEncoder(b, twain, speed, 1e4) }
func BenchmarkEncodeTwainSpeed1e5(b *testing.B)     { benchmarkEncoder(b, twain, speed, 1e5) }
func BenchmarkEncodeTwainSpeed1e6(b *testing.B)     { benchmarkEncoder(b, twain, speed, 1e6) }
func BenchmarkEncodeTwainDefault1e4(b *testing.B)   { benchmarkEncoder(b, twain, default_, 1e4) }
func BenchmarkEncodeTwainDefault1e5(b *testing.B)   { benchmarkEncoder(b, twain, default_, 1e5) }
func BenchmarkEncodeTwainDefault1e6(b *testing.B)   { benchmarkEncoder(b, twain, default_, 1e6) }
func BenchmarkEncodeTwainCompress1e4(b *testing.B)  { benchmarkEncoder(b, twain, compress, 1e4) }
func BenchmarkEncodeTwainCompress1e5(b *testing.B)  { benchmarkEncoder(b, twain, compress, 1e5) }
func BenchmarkEncodeTwainCompress1e6(b *testing.B)  { benchmarkEncoder(b, twain, compress, 1e6) }
<commit_msg>Re-use writer in benchmark.<commit_after>// Copyright 2012 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 flate

import (
	"io/ioutil"
	"runtime"
	"testing"
)

func benchmarkEncoder(b *testing.B, testfile, level, n int) {
	b.StopTimer()
	b.SetBytes(int64(n))
	buf0, err := ioutil.ReadFile(testfiles[testfile])
	if err != nil {
		b.Fatal(err)
	}
	if len(buf0) == 0 {
		b.Fatalf("test file %q has no data", testfiles[testfile])
	}
	buf1 := make([]byte, n)
	for i := 0; i < n; i += len(buf0) {
		if len(buf0) > n-i {
			buf0 = buf0[:n-i]
		}
		copy(buf1[i:], buf0)
	}
	buf0 = nil
	runtime.GC()
	w, err := NewWriter(ioutil.Discard, level)
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		w.Reset(ioutil.Discard)
		_, err = w.Write(buf1)
		if err != nil {
			b.Fatal(err)
		}
		err = w.Close()
		if err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkEncodeDigitsSpeed1e4(b *testing.B)    { benchmarkEncoder(b, digits, speed, 1e4) }
func BenchmarkEncodeDigitsSpeed1e5(b *testing.B)    { benchmarkEncoder(b, digits, speed, 1e5) }
func BenchmarkEncodeDigitsSpeed1e6(b *testing.B)    { benchmarkEncoder(b, digits, speed, 1e6) }
func BenchmarkEncodeDigitsDefault1e4(b *testing.B)  { benchmarkEncoder(b, digits, default_, 1e4) }
func BenchmarkEncodeDigitsDefault1e5(b *testing.B)  { benchmarkEncoder(b, digits, default_, 1e5) }
func BenchmarkEncodeDigitsDefault1e6(b *testing.B)  { benchmarkEncoder(b, digits, default_, 1e6) }
func BenchmarkEncodeDigitsCompress1e4(b *testing.B) { benchmarkEncoder(b, digits, compress, 1e4) }
func BenchmarkEncodeDigitsCompress1e5(b *testing.B) { benchmarkEncoder(b, digits, compress, 1e5) }
func BenchmarkEncodeDigitsCompress1e6(b *testing.B) { benchmarkEncoder(b, digits, compress, 1e6) }
func BenchmarkEncodeTwainSpeed1e4(b *testing.B)     { benchmarkEncoder(b, twain, speed, 1e4) }
func BenchmarkEncodeTwainSpeed1e5(b *testing.B)     { benchmarkEncoder(b, twain, speed, 1e5) }
func BenchmarkEncodeTwainSpeed1e6(b *testing.B)     { benchmarkEncoder(b, twain, speed, 1e6) }
func BenchmarkEncodeTwainDefault1e4(b *testing.B)   { benchmarkEncoder(b, twain, default_, 1e4) }
func BenchmarkEncodeTwainDefault1e5(b *testing.B)   { benchmarkEncoder(b, twain, default_, 1e5) }
func BenchmarkEncodeTwainDefault1e6(b *testing.B)   { benchmarkEncoder(b, twain, default_, 1e6) }
func BenchmarkEncodeTwainCompress1e4(b *testing.B)  { benchmarkEncoder(b, twain, compress, 1e4) }
func BenchmarkEncodeTwainCompress1e5(b *testing.B)  { benchmarkEncoder(b, twain, compress, 1e5) }
func BenchmarkEncodeTwainCompress1e6(b *testing.B)  { benchmarkEncoder(b, twain, compress, 1e6) }
<|endoftext|>


<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 local

import (
	"bytes"
	"fmt"
	"github.com/globocom/commandmocker"
	"github.com/globocom/config"
	fstesting "github.com/globocom/tsuru/fs/testing"
	"github.com/globocom/tsuru/provision"
	"github.com/globocom/tsuru/testing"
	"io/ioutil"
	"labix.org/v2/mgo/bson"
	"launchpad.net/gocheck"
	"os"
	"time"
)

func (s *S) TestShouldBeRegistered(c *gocheck.C) {
	p, err := provision.Get("local")
	c.Assert(err, gocheck.IsNil)
	c.Assert(p, gocheck.FitsTypeOf, &LocalProvisioner{})
}

func (s *S) TestProvisionerProvision(c *gocheck.C) {
	config.Set("local:authorized-key-path", "somepath")
	rfs := &fstesting.RecordingFs{}
	fsystem = rfs
	defer func() {
		fsystem = nil
	}()
	f, _ := os.Open("testdata/dnsmasq.leases")
	data, err := ioutil.ReadAll(f)
	c.Assert(err, gocheck.IsNil)
	file, err := rfs.Open("/var/lib/misc/dnsmasq.leases")
	c.Assert(err, gocheck.IsNil)
	_, err = file.Write(data)
	c.Assert(err, gocheck.IsNil)
	tmpdir, err := commandmocker.Add("sudo", "$*")
	c.Assert(err, gocheck.IsNil)
	defer commandmocker.Remove(tmpdir)
	sshTempDir, err := commandmocker.Add("ssh", "$*")
	c.Assert(err, gocheck.IsNil)
	defer commandmocker.Remove(sshTempDir)
	scpTempDir, err := commandmocker.Add("scp", "$*")
	c.Assert(err, gocheck.IsNil)
	defer commandmocker.Remove(scpTempDir)
	var p LocalProvisioner
	app := testing.NewFakeApp("myapp", "python", 0)
	c.Assert(p.Provision(app), gocheck.IsNil)
	time.Sleep(5 * time.Second)
	c.Assert(commandmocker.Ran(tmpdir), gocheck.Equals, true)
	expected := "lxc-create -t ubuntu -n myapp -- -S somepath"
	expected += "lxc-start --daemon -n myapp"
	expected += "service nginx restart"
	c.Assert(commandmocker.Output(tmpdir), gocheck.Equals, expected)
	var unit provision.Unit
	err = p.collection().Find(bson.M{"name": "myapp"}).One(&unit)
	c.Assert(err, gocheck.IsNil)
	c.Assert(unit.Ip, gocheck.Equals, "10.10.10.15")
	defer p.collection().Remove(bson.M{"name": "myapp"})
}

func (s *S) TestProvisionerRestart(c *gocheck.C) {
	var p LocalProvisioner
	tmpdir, err := commandmocker.Add("ssh", "ok")
	c.Assert(err, gocheck.IsNil)
	defer commandmocker.Remove(tmpdir)
	app := testing.NewFakeApp("almah", "static", 1)
	err = p.Restart(app)
	c.Assert(err, gocheck.IsNil)
	ip := app.ProvisionUnits()[0].GetIp()
	expected := []string{
		"-l", "ubuntu", "-q", "-o", "StrictHostKeyChecking no", ip, "/var/lib/tsuru/hooks/restart",
	}
	c.Assert(commandmocker.Ran(tmpdir), gocheck.Equals, true)
	c.Assert(commandmocker.Parameters(tmpdir), gocheck.DeepEquals, expected)
}

func (s *S) TestProvisionerRestartFailure(c *gocheck.C) {
	tmpdir, err := commandmocker.Error("ssh", "fatal unexpected failure", 25)
	c.Assert(err, gocheck.IsNil)
	defer commandmocker.Remove(tmpdir)
	app := testing.NewFakeApp("cribcaged", "python", 1)
	p := LocalProvisioner{}
	err = p.Restart(app)
	c.Assert(err, gocheck.NotNil)
	pErr, ok := err.(*provision.Error)
	c.Assert(ok, gocheck.Equals, true)
	c.Assert(pErr.Reason, gocheck.Equals, "fatal unexpected failure")
	c.Assert(pErr.Err.Error(), gocheck.Equals, "exit status 25")
}

func (s *S) TestProvisionerDestroy(c *gocheck.C) {
	config.Set("local:authorized-key-path", "somepath")
	rfs := &fstesting.RecordingFs{}
	fsystem = rfs
	defer func() {
		fsystem = nil
	}()
	f, _ := os.Open("testdata/dnsmasq.leases")
	data, err := ioutil.ReadAll(f)
	c.Assert(err, gocheck.IsNil)
	file, err := rfs.Open("/var/lib/misc/dnsmasq.leases")
	c.Assert(err, gocheck.IsNil)
	_, err = file.Write(data)
	c.Assert(err, gocheck.IsNil)
	tmpdir, err := commandmocker.Add("sudo", "$*")
	c.Assert(err, gocheck.IsNil)
	defer commandmocker.Remove(tmpdir)
	sshTempDir, err := commandmocker.Add("ssh", "$*")
	c.Assert(err, gocheck.IsNil)
	defer commandmocker.Remove(sshTempDir)
	scpTempDir, err := commandmocker.Add("scp", "$*")
	c.Assert(err, gocheck.IsNil)
	defer commandmocker.Remove(scpTempDir)
	var p LocalProvisioner
	app := testing.NewFakeApp("myapp", "python", 0)
	err = p.Provision(app)
	time.Sleep(5 * time.Second)
	c.Assert(err, gocheck.IsNil)
	c.Assert(p.Destroy(app), gocheck.IsNil)
	time.Sleep(5 * time.Second)
	c.Assert(commandmocker.Ran(tmpdir), gocheck.Equals, true)
	expected := "lxc-create -t ubuntu -n myapp -- -S somepath"
	expected += "lxc-start --daemon -n myapp"
	expected += "service nginx restart"
	expected += "lxc-stop -n myapp"
	expected += "lxc-destroy -n myapp"
	c.Assert(commandmocker.Output(tmpdir), gocheck.Equals, expected)
	length, err := p.collection().Find(bson.M{"name": "myapp"}).Count()
	c.Assert(err, gocheck.IsNil)
	c.Assert(length, gocheck.Equals, 0)
}

func (s *S) TestProvisionerAddr(c *gocheck.C) {
	var p LocalProvisioner
	app := testing.NewFakeApp("myapp", "python", 1)
	addr, err := p.Addr(app)
	c.Assert(err, gocheck.IsNil)
	c.Assert(addr, gocheck.Equals, app.ProvisionUnits()[0].GetIp())
}

func (s *S) TestProvisionerAddUnits(c *gocheck.C) {
	var p LocalProvisioner
	app := testing.NewFakeApp("myapp", "python", 0)
	units, err := p.AddUnits(app, 2)
	c.Assert(err, gocheck.IsNil)
	c.Assert(units, gocheck.DeepEquals, []provision.Unit{})
}

func (s *S) TestProvisionerRemoveUnit(c *gocheck.C) {
	var p LocalProvisioner
	app := testing.NewFakeApp("myapp", "python", 0)
	err := p.RemoveUnit(app, "")
	c.Assert(err, gocheck.IsNil)
}

func (s *S) TestProvisionerExecuteCommand(c *gocheck.C) {
	var p LocalProvisioner
	var buf bytes.Buffer
	tmpdir, err := commandmocker.Add("ssh", "$*")
	c.Assert(err, gocheck.IsNil)
	defer commandmocker.Remove(tmpdir)
	app := testing.NewFakeApp("almah", "static", 2)
	err = p.ExecuteCommand(&buf, &buf, app, "ls", "-lh")
	c.Assert(err, gocheck.IsNil)
	cmdOutput := fmt.Sprintf("-l ubuntu -q -o StrictHostKeyChecking no %s ls -lh", app.ProvisionUnits()[0].GetIp())
	c.Assert(commandmocker.Ran(tmpdir), gocheck.Equals, true)
	c.Assert(commandmocker.Output(tmpdir), gocheck.Equals, cmdOutput)
}

func (s *S) TestCollectStatus(c *gocheck.C) {
	var p LocalProvisioner
	expected := []provision.Unit{
		{
			Name:       "vm1",
			AppName:    "vm1",
			Type:       "django",
			Machine:    0,
			InstanceId: "vm1",
			Ip:         "10.10.10.9",
			Status:     provision.StatusStarted,
		},
		{
			Name:       "vm2",
			AppName:    "vm2",
			Type:       "gunicorn",
			Machine:    0,
			InstanceId: "vm2",
			Ip:         "10.10.10.10",
			Status:     provision.StatusInstalling,
		},
	}
	for _, u := range expected {
		err := p.collection().Insert(u)
		c.Assert(err, gocheck.IsNil)
	}
	units, err := p.CollectStatus()
	c.Assert(err, gocheck.IsNil)
	c.Assert(units, gocheck.DeepEquals, expected)
}

func (s *S) TestProvisionCollection(c *gocheck.C) {
	var p LocalProvisioner
	collection := p.collection()
	c.Assert(collection.Name, gocheck.Equals, s.collName)
}

func (s *S) TestProvisionInstall(c *gocheck.C) {
	tmpdir, err := commandmocker.Add("ssh", "$*")
	c.Assert(err, gocheck.IsNil)
	defer commandmocker.Remove(tmpdir)
	p := LocalProvisioner{}
	err = p.install("10.10.10.10")
	c.Assert(err, gocheck.IsNil)
	c.Assert(commandmocker.Ran(tmpdir), gocheck.Equals, true)
	cmds := []string{
		"-q",
		"-o",
		"StrictHostKeyChecking no",
		"-l",
		"ubuntu",
		"10.10.10.10",
		"sudo /var/lib/tsuru/hooks/install",
	}
	c.Assert(commandmocker.Parameters(tmpdir), gocheck.DeepEquals, cmds)
}

func (s *S) TestProvisionStart(c *gocheck.C) {
	tmpdir, err := commandmocker.Add("ssh", "$*")
	c.Assert(err, gocheck.IsNil)
	defer commandmocker.Remove(tmpdir)
	p := LocalProvisioner{}
	err = p.start("10.10.10.10")
	c.Assert(err, gocheck.IsNil)
	c.Assert(commandmocker.Ran(tmpdir), gocheck.Equals, true)
	cmds := []string{
		"-q",
		"-o",
		"StrictHostKeyChecking no",
		"-l",
		"ubuntu",
		"10.10.10.10",
		"sudo /var/lib/tsuru/hooks/start",
	}
	c.Assert(commandmocker.Parameters(tmpdir), gocheck.DeepEquals, cmds)
}

func (s *S) TestProvisionSetup(c *gocheck.C) {
	tmpdir, err := commandmocker.Add("scp", "$*")
	c.Assert(err, gocheck.IsNil)
	defer commandmocker.Remove(tmpdir)
	sshTempDir, err := commandmocker.Add("ssh", "$*")
	c.Assert(err, gocheck.IsNil)
	defer commandmocker.Remove(sshTempDir)
	p := LocalProvisioner{}
	formulasPath := "/home/ubuntu/formulas"
	config.Set("local:formulas-path", formulasPath)
	err = p.setup("10.10.10.10", "static")
	c.Assert(err, gocheck.IsNil)
	c.Assert(commandmocker.Ran(tmpdir), gocheck.Equals, true)
	cmds := []string{
		"-q",
		"-o",
		"StrictHostKeyChecking no",
		"-r",
		formulasPath + "/static/hooks",
		"ubuntu@10.10.10.10:/var/lib/tsuru",
	}
	c.Assert(commandmocker.Parameters(tmpdir), gocheck.DeepEquals, cmds)
	c.Assert(commandmocker.Ran(sshTempDir), gocheck.Equals, true)
	cmds = []string{
		"-q",
		"-o",
		"StrictHostKeyChecking no",
		"-l",
		"ubuntu",
		"10.10.10.10",
		"sudo mkdir -p /var/lib/tsuru/hooks",
		"-q",
		"-o",
		"StrictHostKeyChecking no",
		"-l",
		"ubuntu",
		"10.10.10.10",
		"sudo chown -R ubuntu /var/lib/tsuru/hooks",
	}
	c.Assert(commandmocker.Parameters(sshTempDir), gocheck.DeepEquals, cmds)
}
<commit_msg>provision/local: make teste more reliable<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 local

import (
	"bytes"
	"fmt"
	"github.com/globocom/commandmocker"
	"github.com/globocom/config"
	fstesting "github.com/globocom/tsuru/fs/testing"
	"github.com/globocom/tsuru/provision"
	"github.com/globocom/tsuru/testing"
	"io/ioutil"
	"labix.org/v2/mgo/bson"
	"launchpad.net/gocheck"
	"os"
	"time"
)

func (s *S) TestShouldBeRegistered(c *gocheck.C) {
	p, err := provision.Get("local")
	c.Assert(err, gocheck.IsNil)
	c.Assert(p, gocheck.FitsTypeOf, &LocalProvisioner{})
}

func (s *S) TestProvisionerProvision(c *gocheck.C) {
	config.Set("local:authorized-key-path", "somepath")
	rfs := &fstesting.RecordingFs{}
	fsystem = rfs
	defer func() {
		fsystem = nil
	}()
	f, _ := os.Open("testdata/dnsmasq.leases")
	data, err := ioutil.ReadAll(f)
	c.Assert(err, gocheck.IsNil)
	file, err := rfs.Open("/var/lib/misc/dnsmasq.leases")
	c.Assert(err, gocheck.IsNil)
	_, err = file.Write(data)
	c.Assert(err, gocheck.IsNil)
	tmpdir, err := commandmocker.Add("sudo", "$*")
	c.Assert(err, gocheck.IsNil)
	defer commandmocker.Remove(tmpdir)
	sshTempDir, err := commandmocker.Add("ssh", "$*")
	c.Assert(err, gocheck.IsNil)
	defer commandmocker.Remove(sshTempDir)
	scpTempDir, err := commandmocker.Add("scp", "$*")
	c.Assert(err, gocheck.IsNil)
	defer commandmocker.Remove(scpTempDir)
	var p LocalProvisioner
	app := testing.NewFakeApp("myapp", "python", 0)
	defer p.collection().Remove(bson.M{"name": "myapp"})
	c.Assert(p.Provision(app), gocheck.IsNil)
	ok := make(chan bool, 1)
	go func() {
		for {
			coll := s.conn.Collection(s.collName)
			ct, err := coll.Find(bson.M{"name": "myapp", "status": provision.StatusStarted}).Count()
			if err != nil {
				c.Fatal(err)
			}
			if ct > 0 {
				ok <- true
				return
			}
			time.Sleep(1e3)
		}
	}()
	select {
	case <-ok:
	case <-time.After(10e9):
		c.Fatal("Timed out waiting for the container to be provisioned (10 seconds)")
	}
	c.Assert(commandmocker.Ran(tmpdir), gocheck.Equals, true)
	expected := "lxc-create -t ubuntu -n myapp -- -S somepath"
	expected += "lxc-start --daemon -n myapp"
	expected += "service nginx restart"
	c.Assert(commandmocker.Output(tmpdir), gocheck.Equals, expected)
	var unit provision.Unit
	err = s.conn.Collection(s.collName).Find(bson.M{"name": "myapp"}).One(&unit)
	c.Assert(err, gocheck.IsNil)
	c.Assert(unit.Ip, gocheck.Equals, "10.10.10.15")
}

func (s *S) TestProvisionerRestart(c *gocheck.C) {
	var p LocalProvisioner
	tmpdir, err := commandmocker.Add("ssh", "ok")
	c.Assert(err, gocheck.IsNil)
	defer commandmocker.Remove(tmpdir)
	app := testing.NewFakeApp("almah", "static", 1)
	err = p.Restart(app)
	c.Assert(err, gocheck.IsNil)
	ip := app.ProvisionUnits()[0].GetIp()
	expected := []string{
		"-l", "ubuntu", "-q", "-o", "StrictHostKeyChecking no", ip, "/var/lib/tsuru/hooks/restart",
	}
	c.Assert(commandmocker.Ran(tmpdir), gocheck.Equals, true)
	c.Assert(commandmocker.Parameters(tmpdir), gocheck.DeepEquals, expected)
}

func (s *S) TestProvisionerRestartFailure(c *gocheck.C) {
	tmpdir, err := commandmocker.Error("ssh", "fatal unexpected failure", 25)
	c.Assert(err, gocheck.IsNil)
	defer commandmocker.Remove(tmpdir)
	app := testing.NewFakeApp("cribcaged", "python", 1)
	p := LocalProvisioner{}
	err = p.Restart(app)
	c.Assert(err, gocheck.NotNil)
	pErr, ok := err.(*provision.Error)
	c.Assert(ok, gocheck.Equals, true)
	c.Assert(pErr.Reason, gocheck.Equals, "fatal unexpected failure")
	c.Assert(pErr.Err.Error(), gocheck.Equals, "exit status 25")
}

func (s *S) TestProvisionerDestroy(c *gocheck.C) {
	config.Set("local:authorized-key-path", "somepath")
	rfs := &fstesting.RecordingFs{}
	fsystem = rfs
	defer func() {
		fsystem = nil
	}()
	f, _ := os.Open("testdata/dnsmasq.leases")
	data, err := ioutil.ReadAll(f)
	c.Assert(err, gocheck.IsNil)
	file, err := rfs.Open("/var/lib/misc/dnsmasq.leases")
	c.Assert(err, gocheck.IsNil)
	_, err = file.Write(data)
	c.Assert(err, gocheck.IsNil)
	tmpdir, err := commandmocker.Add("sudo", "$*")
	c.Assert(err, gocheck.IsNil)
	defer commandmocker.Remove(tmpdir)
	sshTempDir, err := commandmocker.Add("ssh", "$*")
	c.Assert(err, gocheck.IsNil)
	defer commandmocker.Remove(sshTempDir)
	scpTempDir, err := commandmocker.Add("scp", "$*")
	c.Assert(err, gocheck.IsNil)
	defer commandmocker.Remove(scpTempDir)
	var p LocalProvisioner
	app := testing.NewFakeApp("myapp", "python", 0)
	err = p.Provision(app)
	time.Sleep(5 * time.Second)
	c.Assert(err, gocheck.IsNil)
	c.Assert(p.Destroy(app), gocheck.IsNil)
	time.Sleep(5 * time.Second)
	c.Assert(commandmocker.Ran(tmpdir), gocheck.Equals, true)
	expected := "lxc-create -t ubuntu -n myapp -- -S somepath"
	expected += "lxc-start --daemon -n myapp"
	expected += "service nginx restart"
	expected += "lxc-stop -n myapp"
	expected += "lxc-destroy -n myapp"
	c.Assert(commandmocker.Output(tmpdir), gocheck.Equals, expected)
	length, err := p.collection().Find(bson.M{"name": "myapp"}).Count()
	c.Assert(err, gocheck.IsNil)
	c.Assert(length, gocheck.Equals, 0)
}

func (s *S) TestProvisionerAddr(c *gocheck.C) {
	var p LocalProvisioner
	app := testing.NewFakeApp("myapp", "python", 1)
	addr, err := p.Addr(app)
	c.Assert(err, gocheck.IsNil)
	c.Assert(addr, gocheck.Equals, app.ProvisionUnits()[0].GetIp())
}

func (s *S) TestProvisionerAddUnits(c *gocheck.C) {
	var p LocalProvisioner
	app := testing.NewFakeApp("myapp", "python", 0)
	units, err := p.AddUnits(app, 2)
	c.Assert(err, gocheck.IsNil)
	c.Assert(units, gocheck.DeepEquals, []provision.Unit{})
}

func (s *S) TestProvisionerRemoveUnit(c *gocheck.C) {
	var p LocalProvisioner
	app := testing.NewFakeApp("myapp", "python", 0)
	err := p.RemoveUnit(app, "")
	c.Assert(err, gocheck.IsNil)
}

func (s *S) TestProvisionerExecuteCommand(c *gocheck.C) {
	var p LocalProvisioner
	var buf bytes.Buffer
	tmpdir, err := commandmocker.Add("ssh", "$*")
	c.Assert(err, gocheck.IsNil)
	defer commandmocker.Remove(tmpdir)
	app := testing.NewFakeApp("almah", "static", 2)
	err = p.ExecuteCommand(&buf, &buf, app, "ls", "-lh")
	c.Assert(err, gocheck.IsNil)
	cmdOutput := fmt.Sprintf("-l ubuntu -q -o StrictHostKeyChecking no %s ls -lh", app.ProvisionUnits()[0].GetIp())
	c.Assert(commandmocker.Ran(tmpdir), gocheck.Equals, true)
	c.Assert(commandmocker.Output(tmpdir), gocheck.Equals, cmdOutput)
}

func (s *S) TestCollectStatus(c *gocheck.C) {
	var p LocalProvisioner
	expected := []provision.Unit{
		{
			Name:       "vm1",
			AppName:    "vm1",
			Type:       "django",
			Machine:    0,
			InstanceId: "vm1",
			Ip:         "10.10.10.9",
			Status:     provision.StatusStarted,
		},
		{
			Name:       "vm2",
			AppName:    "vm2",
			Type:       "gunicorn",
			Machine:    0,
			InstanceId: "vm2",
			Ip:         "10.10.10.10",
			Status:     provision.StatusInstalling,
		},
	}
	for _, u := range expected {
		err := p.collection().Insert(u)
		c.Assert(err, gocheck.IsNil)
	}
	units, err := p.CollectStatus()
	c.Assert(err, gocheck.IsNil)
	c.Assert(units, gocheck.DeepEquals, expected)
}

func (s *S) TestProvisionCollection(c *gocheck.C) {
	var p LocalProvisioner
	collection := p.collection()
	c.Assert(collection.Name, gocheck.Equals, s.collName)
}

func (s *S) TestProvisionInstall(c *gocheck.C) {
	tmpdir, err := commandmocker.Add("ssh", "$*")
	c.Assert(err, gocheck.IsNil)
	defer commandmocker.Remove(tmpdir)
	p := LocalProvisioner{}
	err = p.install("10.10.10.10")
	c.Assert(err, gocheck.IsNil)
	c.Assert(commandmocker.Ran(tmpdir), gocheck.Equals, true)
	cmds := []string{
		"-q",
		"-o",
		"StrictHostKeyChecking no",
		"-l",
		"ubuntu",
		"10.10.10.10",
		"sudo /var/lib/tsuru/hooks/install",
	}
	c.Assert(commandmocker.Parameters(tmpdir), gocheck.DeepEquals, cmds)
}

func (s *S) TestProvisionStart(c *gocheck.C) {
	tmpdir, err := commandmocker.Add("ssh", "$*")
	c.Assert(err, gocheck.IsNil)
	defer commandmocker.Remove(tmpdir)
	p := LocalProvisioner{}
	err = p.start("10.10.10.10")
	c.Assert(err, gocheck.IsNil)
	c.Assert(commandmocker.Ran(tmpdir), gocheck.Equals, true)
	cmds := []string{
		"-q",
		"-o",
		"StrictHostKeyChecking no",
		"-l",
		"ubuntu",
		"10.10.10.10",
		"sudo /var/lib/tsuru/hooks/start",
	}
	c.Assert(commandmocker.Parameters(tmpdir), gocheck.DeepEquals, cmds)
}

func (s *S) TestProvisionSetup(c *gocheck.C) {
	tmpdir, err := commandmocker.Add("scp", "$*")
	c.Assert(err, gocheck.IsNil)
	defer commandmocker.Remove(tmpdir)
	sshTempDir, err := commandmocker.Add("ssh", "$*")
	c.Assert(err, gocheck.IsNil)
	defer commandmocker.Remove(sshTempDir)
	p := LocalProvisioner{}
	formulasPath := "/home/ubuntu/formulas"
	config.Set("local:formulas-path", formulasPath)
	err = p.setup("10.10.10.10", "static")
	c.Assert(err, gocheck.IsNil)
	c.Assert(commandmocker.Ran(tmpdir), gocheck.Equals, true)
	cmds := []string{
		"-q",
		"-o",
		"StrictHostKeyChecking no",
		"-r",
		formulasPath + "/static/hooks",
		"ubuntu@10.10.10.10:/var/lib/tsuru",
	}
	c.Assert(commandmocker.Parameters(tmpdir), gocheck.DeepEquals, cmds)
	c.Assert(commandmocker.Ran(sshTempDir), gocheck.Equals, true)
	cmds = []string{
		"-q",
		"-o",
		"StrictHostKeyChecking no",
		"-l",
		"ubuntu",
		"10.10.10.10",
		"sudo mkdir -p /var/lib/tsuru/hooks",
		"-q",
		"-o",
		"StrictHostKeyChecking no",
		"-l",
		"ubuntu",
		"10.10.10.10",
		"sudo chown -R ubuntu /var/lib/tsuru/hooks",
	}
	c.Assert(commandmocker.Parameters(sshTempDir), gocheck.DeepEquals, cmds)
}
<|endoftext|>


<commit_before>//
// Written by Maxim Khitrov (November 2012)
//

// Package flowrate provides the tools for monitoring and limiting the flow rate
// of an arbitrary data stream.
package flowrate

import (
	"math"
	"sync"
	"time"
)

// Monitor monitors and limits the transfer rate of a data stream.
type Monitor struct {
	mu      sync.Mutex    // Mutex guarding access to all internal fields
	active  bool          // Flag indicating an active transfer
	start   time.Duration // Transfer start time (clock() value)
	bytes   int64         // Total number of bytes transferred
	samples int64         // Total number of samples taken

	rSample float64 // Most recent transfer rate sample (bytes per second)
	rEMA    float64 // Exponential moving average of rSample
	rPeak   float64 // Peak transfer rate (max of all rSamples)
	rWindow float64 // rEMA window (seconds)

	sBytes int64         // Number of bytes transferred since sLast
	sLast  time.Duration // Most recent sample time (stop time when inactive)
	sRate  time.Duration // Sampling rate

	tBytes int64         // Number of bytes expected in the current transfer
	tLast  time.Duration // Time of the most recent transfer of at least 1 byte
}

// New creates a new flow control monitor. Instantaneous transfer rate is
// measured and updated for each sampleRate interval. windowSize determines the
// weight of each sample in the exponential moving average (EMA) calculation.
// The exact formulas are:
//
// 	sampleTime = currentTime - prevSampleTime
// 	sampleRate = byteCount / sampleTime
// 	weight     = 1 - exp(-sampleTime/windowSize)
// 	newRate    = weight*sampleRate + (1-weight)*oldRate
//
// The default values for sampleRate and windowSize (if <= 0) are 100ms and 1s,
// respectively.
func New(sampleRate, windowSize time.Duration) *Monitor {
	if sampleRate = clockRound(sampleRate); sampleRate <= 0 {
		sampleRate = 5 * clockRate
	}
	if windowSize <= 0 {
		windowSize = 1 * time.Second
	}
	now := clock()
	return &Monitor{
		active:  true,
		start:   now,
		rWindow: windowSize.Seconds(),
		sLast:   now,
		sRate:   sampleRate,
		tLast:   now,
	}
}

// Update records the transfer of n bytes and returns n. It should be called
// after each Read/Write operation, even if n is 0.
func (m *Monitor) Update(n int) int {
	m.mu.Lock()
	m.update(n)
	m.mu.Unlock()
	return n
}

// IO is a convenience method intended to wrap io.Reader and io.Writer method
// execution. It calls m.Update(n) and then returns (n, err) unmodified.
func (m *Monitor) IO(n int, err error) (int, error) {
	return m.Update(n), err
}

// Done marks the transfer as finished and prevents any further updates or
// limiting. Instantaneous and current transfer rates drop to 0. Update, IO, and
// Limit methods become NOOPs. It returns the total number of bytes transferred.
func (m *Monitor) Done() int64 {
	m.mu.Lock()
	if now := m.update(0); m.sBytes > 0 {
		m.reset(now)
	}
	m.active = false
	m.tLast = 0
	n := m.bytes
	m.mu.Unlock()
	return n
}

// timeRemLimit is the maximum Status.TimeRem value.
const timeRemLimit = 999*time.Hour + 59*time.Minute + 59*time.Second

// Status represents the current Monitor status. All transfer rates are in bytes
// per second rounded to the nearest byte.
type Status struct {
	Active   bool          // Flag indicating an active transfer
	Start    time.Time     // Transfer start time
	Duration time.Duration // Time period covered by the statistics
	Idle     time.Duration // Time since the last transfer of at least 1 byte
	Bytes    int64         // Total number of bytes transferred
	Samples  int64         // Total number of samples taken
	InstRate int64         // Instantaneous transfer rate
	CurRate  int64         // Current transfer rate (EMA of InstRate)
	AvgRate  int64         // Average transfer rate (Bytes / Duration)
	PeakRate int64         // Maximum instantaneous transfer rate
	BytesRem int64         // Number of bytes remaining in the transfer
	TimeRem  time.Duration // Estimated time to completion
	Progress Percent       // Overall transfer progress
}

// Status returns current transfer status information. The returned value
// becomes static after a call to Done.
func (m *Monitor) Status() Status {
	m.mu.Lock()
	now := m.update(0)
	s := Status{
		Active:   m.active,
		Start:    clockToTime(m.start),
		Duration: m.sLast - m.start,
		Idle:     now - m.tLast,
		Bytes:    m.bytes,
		Samples:  m.samples,
		PeakRate: round(m.rPeak),
		BytesRem: m.tBytes - m.bytes,
		Progress: percentOf(float64(m.bytes), float64(m.tBytes)),
	}
	if s.BytesRem < 0 {
		s.BytesRem = 0
	}
	if s.Duration > 0 {
		rAvg := float64(s.Bytes) / s.Duration.Seconds()
		s.AvgRate = round(rAvg)
		if s.Active {
			s.InstRate = round(m.rSample)
			s.CurRate = round(m.rEMA)
			if s.BytesRem > 0 {
				if tRate := 0.8*m.rEMA + 0.2*rAvg; tRate > 0 {
					ns := float64(s.BytesRem) / tRate * 1e9
					if ns > float64(timeRemLimit) {
						ns = float64(timeRemLimit)
					}
					s.TimeRem = clockRound(time.Duration(ns))
				}
			}
		}
	}
	m.mu.Unlock()
	return s
}

// Limit restricts the instantaneous (per-sample) data flow to rate bytes per
// second. It returns the maximum number of bytes (0 <= n <= want) that may be
// transferred immediately without exceeding the limit. If block == true, the
// call blocks until n > 0. want is returned unmodified if want < 1, rate < 1,
// or the transfer is inactive (after a call to Done).
//
// At least one byte is always allowed to be transferred in any given sampling
// period. Thus, if the sampling rate is 100ms, the lowest achievable flow rate
// is 10 bytes per second.
//
// For usage examples, see the implementation of Reader and Writer in io.go.
func (m *Monitor) Limit(want int, rate int64, block bool) (n int) {
	if want < 1 || rate < 1 {
		return want
	}
	m.mu.Lock()

	// Determine the maximum number of bytes that can be sent in one sample
	limit := round(float64(rate) * m.sRate.Seconds())
	if limit <= 0 {
		limit = 1
	}

	// If block == true, wait until m.sBytes < limit
	if now := m.update(0); block {
		for m.sBytes >= limit && m.active {
			now = m.waitNextSample(now)
		}
	}

	// Make limit <= want (unlimited if the transfer is no longer active)
	if limit -= m.sBytes; limit > int64(want) || !m.active {
		limit = int64(want)
	}
	m.mu.Unlock()

	if limit < 0 {
		limit = 0
	}
	return int(limit)
}

// SetTransferSize specifies the total size of the data transfer, which allows
// the Monitor to calculate the overall progress and time to completion.
func (m *Monitor) SetTransferSize(bytes int64) {
	if bytes < 0 {
		bytes = 0
	}
	m.mu.Lock()
	m.tBytes = bytes
	m.mu.Unlock()
}

// update accumulates the transferred byte count for the current sample until
// clock() - m.sLast >= m.sRate. The monitor status is updated once the current
// sample is done.
func (m *Monitor) update(n int) (now time.Duration) {
	if !m.active {
		return
	}
	if now = clock(); n > 0 {
		m.tLast = now
	}
	m.sBytes += int64(n)
	if sTime := now - m.sLast; sTime >= m.sRate {
		t := sTime.Seconds()
		if m.rSample = float64(m.sBytes) / t; m.rSample > m.rPeak {
			m.rPeak = m.rSample
		}

		// Exponential moving average using a method similar to *nix load
		// average calculation. Longer sampling periods carry greater weight.
		if m.samples > 0 {
			w := math.Exp(-t / m.rWindow)
			m.rEMA = m.rSample + w*(m.rEMA-m.rSample)
		} else {
			m.rEMA = m.rSample
		}
		m.reset(now)
	}
	return
}

// reset clears the current sample state in preparation for the next sample.
func (m *Monitor) reset(sampleTime time.Duration) {
	m.bytes += m.sBytes
	m.samples++
	m.sBytes = 0
	m.sLast = sampleTime
}

// waitNextSample sleeps for the remainder of the current sample. The lock is
// released and reacquired during the actual sleep period, so it's possible for
// the transfer to be inactive when this method returns.
func (m *Monitor) waitNextSample(now time.Duration) time.Duration {
	const minWait = 5 * time.Millisecond
	current := m.sLast

	// sleep until the last sample time changes (ideally, just one iteration)
	for m.sLast == current && m.active {
		d := current + m.sRate - now
		m.mu.Unlock()
		if d < minWait {
			d = minWait
		}
		time.Sleep(d)
		m.mu.Lock()
		now = m.update(0)
	}
	return now
}
<commit_msg>Add Monitor.SetREMA()<commit_after>//
// Written by Maxim Khitrov (November 2012)
//

// Package flowrate provides the tools for monitoring and limiting the flow rate
// of an arbitrary data stream.
package flowrate

import (
	"math"
	"sync"
	"time"
)

// Monitor monitors and limits the transfer rate of a data stream.
type Monitor struct {
	mu      sync.Mutex    // Mutex guarding access to all internal fields
	active  bool          // Flag indicating an active transfer
	start   time.Duration // Transfer start time (clock() value)
	bytes   int64         // Total number of bytes transferred
	samples int64         // Total number of samples taken

	rSample float64 // Most recent transfer rate sample (bytes per second)
	rEMA    float64 // Exponential moving average of rSample
	rPeak   float64 // Peak transfer rate (max of all rSamples)
	rWindow float64 // rEMA window (seconds)

	sBytes int64         // Number of bytes transferred since sLast
	sLast  time.Duration // Most recent sample time (stop time when inactive)
	sRate  time.Duration // Sampling rate

	tBytes int64         // Number of bytes expected in the current transfer
	tLast  time.Duration // Time of the most recent transfer of at least 1 byte
}

// New creates a new flow control monitor. Instantaneous transfer rate is
// measured and updated for each sampleRate interval. windowSize determines the
// weight of each sample in the exponential moving average (EMA) calculation.
// The exact formulas are:
//
// 	sampleTime = currentTime - prevSampleTime
// 	sampleRate = byteCount / sampleTime
// 	weight     = 1 - exp(-sampleTime/windowSize)
// 	newRate    = weight*sampleRate + (1-weight)*oldRate
//
// The default values for sampleRate and windowSize (if <= 0) are 100ms and 1s,
// respectively.
func New(sampleRate, windowSize time.Duration) *Monitor {
	if sampleRate = clockRound(sampleRate); sampleRate <= 0 {
		sampleRate = 5 * clockRate
	}
	if windowSize <= 0 {
		windowSize = 1 * time.Second
	}
	now := clock()
	return &Monitor{
		active:  true,
		start:   now,
		rWindow: windowSize.Seconds(),
		sLast:   now,
		sRate:   sampleRate,
		tLast:   now,
	}
}

// Update records the transfer of n bytes and returns n. It should be called
// after each Read/Write operation, even if n is 0.
func (m *Monitor) Update(n int) int {
	m.mu.Lock()
	m.update(n)
	m.mu.Unlock()
	return n
}

// Hack to set the current rEMA.
func (m *Monitor) SetREMA(rEMA float64) {
	m.mu.Lock()
	m.rEMA = rEMA
	m.samples++
	m.mu.Unlock()
}

// IO is a convenience method intended to wrap io.Reader and io.Writer method
// execution. It calls m.Update(n) and then returns (n, err) unmodified.
func (m *Monitor) IO(n int, err error) (int, error) {
	return m.Update(n), err
}

// Done marks the transfer as finished and prevents any further updates or
// limiting. Instantaneous and current transfer rates drop to 0. Update, IO, and
// Limit methods become NOOPs. It returns the total number of bytes transferred.
func (m *Monitor) Done() int64 {
	m.mu.Lock()
	if now := m.update(0); m.sBytes > 0 {
		m.reset(now)
	}
	m.active = false
	m.tLast = 0
	n := m.bytes
	m.mu.Unlock()
	return n
}

// timeRemLimit is the maximum Status.TimeRem value.
const timeRemLimit = 999*time.Hour + 59*time.Minute + 59*time.Second

// Status represents the current Monitor status. All transfer rates are in bytes
// per second rounded to the nearest byte.
type Status struct {
	Active   bool          // Flag indicating an active transfer
	Start    time.Time     // Transfer start time
	Duration time.Duration // Time period covered by the statistics
	Idle     time.Duration // Time since the last transfer of at least 1 byte
	Bytes    int64         // Total number of bytes transferred
	Samples  int64         // Total number of samples taken
	InstRate int64         // Instantaneous transfer rate
	CurRate  int64         // Current transfer rate (EMA of InstRate)
	AvgRate  int64         // Average transfer rate (Bytes / Duration)
	PeakRate int64         // Maximum instantaneous transfer rate
	BytesRem int64         // Number of bytes remaining in the transfer
	TimeRem  time.Duration // Estimated time to completion
	Progress Percent       // Overall transfer progress
}

// Status returns current transfer status information. The returned value
// becomes static after a call to Done.
func (m *Monitor) Status() Status {
	m.mu.Lock()
	now := m.update(0)
	s := Status{
		Active:   m.active,
		Start:    clockToTime(m.start),
		Duration: m.sLast - m.start,
		Idle:     now - m.tLast,
		Bytes:    m.bytes,
		Samples:  m.samples,
		PeakRate: round(m.rPeak),
		BytesRem: m.tBytes - m.bytes,
		Progress: percentOf(float64(m.bytes), float64(m.tBytes)),
	}
	if s.BytesRem < 0 {
		s.BytesRem = 0
	}
	if s.Duration > 0 {
		rAvg := float64(s.Bytes) / s.Duration.Seconds()
		s.AvgRate = round(rAvg)
		if s.Active {
			s.InstRate = round(m.rSample)
			s.CurRate = round(m.rEMA)
			if s.BytesRem > 0 {
				if tRate := 0.8*m.rEMA + 0.2*rAvg; tRate > 0 {
					ns := float64(s.BytesRem) / tRate * 1e9
					if ns > float64(timeRemLimit) {
						ns = float64(timeRemLimit)
					}
					s.TimeRem = clockRound(time.Duration(ns))
				}
			}
		}
	}
	m.mu.Unlock()
	return s
}

// Limit restricts the instantaneous (per-sample) data flow to rate bytes per
// second. It returns the maximum number of bytes (0 <= n <= want) that may be
// transferred immediately without exceeding the limit. If block == true, the
// call blocks until n > 0. want is returned unmodified if want < 1, rate < 1,
// or the transfer is inactive (after a call to Done).
//
// At least one byte is always allowed to be transferred in any given sampling
// period. Thus, if the sampling rate is 100ms, the lowest achievable flow rate
// is 10 bytes per second.
//
// For usage examples, see the implementation of Reader and Writer in io.go.
func (m *Monitor) Limit(want int, rate int64, block bool) (n int) {
	if want < 1 || rate < 1 {
		return want
	}
	m.mu.Lock()

	// Determine the maximum number of bytes that can be sent in one sample
	limit := round(float64(rate) * m.sRate.Seconds())
	if limit <= 0 {
		limit = 1
	}

	// If block == true, wait until m.sBytes < limit
	if now := m.update(0); block {
		for m.sBytes >= limit && m.active {
			now = m.waitNextSample(now)
		}
	}

	// Make limit <= want (unlimited if the transfer is no longer active)
	if limit -= m.sBytes; limit > int64(want) || !m.active {
		limit = int64(want)
	}
	m.mu.Unlock()

	if limit < 0 {
		limit = 0
	}
	return int(limit)
}

// SetTransferSize specifies the total size of the data transfer, which allows
// the Monitor to calculate the overall progress and time to completion.
func (m *Monitor) SetTransferSize(bytes int64) {
	if bytes < 0 {
		bytes = 0
	}
	m.mu.Lock()
	m.tBytes = bytes
	m.mu.Unlock()
}

// update accumulates the transferred byte count for the current sample until
// clock() - m.sLast >= m.sRate. The monitor status is updated once the current
// sample is done.
func (m *Monitor) update(n int) (now time.Duration) {
	if !m.active {
		return
	}
	if now = clock(); n > 0 {
		m.tLast = now
	}
	m.sBytes += int64(n)
	if sTime := now - m.sLast; sTime >= m.sRate {
		t := sTime.Seconds()
		if m.rSample = float64(m.sBytes) / t; m.rSample > m.rPeak {
			m.rPeak = m.rSample
		}

		// Exponential moving average using a method similar to *nix load
		// average calculation. Longer sampling periods carry greater weight.
		if m.samples > 0 {
			w := math.Exp(-t / m.rWindow)
			m.rEMA = m.rSample + w*(m.rEMA-m.rSample)
		} else {
			m.rEMA = m.rSample
		}
		m.reset(now)
	}
	return
}

// reset clears the current sample state in preparation for the next sample.
func (m *Monitor) reset(sampleTime time.Duration) {
	m.bytes += m.sBytes
	m.samples++
	m.sBytes = 0
	m.sLast = sampleTime
}

// waitNextSample sleeps for the remainder of the current sample. The lock is
// released and reacquired during the actual sleep period, so it's possible for
// the transfer to be inactive when this method returns.
func (m *Monitor) waitNextSample(now time.Duration) time.Duration {
	const minWait = 5 * time.Millisecond
	current := m.sLast

	// sleep until the last sample time changes (ideally, just one iteration)
	for m.sLast == current && m.active {
		d := current + m.sRate - now
		m.mu.Unlock()
		if d < minWait {
			d = minWait
		}
		time.Sleep(d)
		m.mu.Lock()
		now = m.update(0)
	}
	return now
}
<|endoftext|>