<commit_before>package mpdocker

import (
	"testing"

	"github.com/fsouza/go-dockerclient"
)

func TestNormalizeMetricName(t *testing.T) {
	testSets := [][]string{
		{"foo/bar", "foo_bar"},
		{"foo:bar", "foo_bar"},
	}

	for _, testSet := range testSets {
		if normalizeMetricName(testSet[0]) != testSet[1] {
			t.Errorf("normalizeMetricName: '%s' should be normalized to '%s', but '%s'", testSet[0], testSet[1], normalizeMetricName(testSet[0]))
		}
	}
}

func TestGraphDefinition(t *testing.T) {
	var docker DockerPlugin

	graphdef := docker.GraphDefinition()
	if len(graphdef) != 6 {
		t.Errorf("GraphDefinition: %d should be 6", len(graphdef))
	}
}

func TestGenerateName(t *testing.T) {
	stub := docker.APIContainers{
		ID:      "bab2b03c736de41ecba6470eba736c5109436f706eedca4f3e0d93d6530eccd4",
		Image:   "tutum/mongodb",
		Command: "/run.sh",
		Created: 1456995574,
		Status:  "Up 4 days",
		Ports: []docker.APIPort{
			{PrivatePort: 28017, Type: "tcp"},
			{PrivatePort: 27017, Type: "tcp"},
		},
		Names:  []string{"/my-mongodb"},
		Labels: map[string]string{"foo": "bar"},
	}
	/* {"Id":"5b963f266d609d2b02aee8f57d664e04d35aa8c23afcbc6bb73bc4a5b2e7c44d",
	   "Image":"memcached",
	   "Command":"/entrypoint.sh memcached",
	   "Created":1456994862,
	   "Status":"Up 4 days",
	   "Ports":[{"PrivatePort":11211,
	   "Type":"tcp"}],
	   "Names":["/my-memcache"]}]`
	*/
	var docker DockerPlugin
	docker.NameFormat = "name_id"
	if docker.generateName(stub) != "my-mongodb_bab2b0" {
		t.Errorf("generateName(name): %s should be 'my-mongodb_bab2b0'", docker.generateName(stub))
	}
	docker.NameFormat = "name"
	if docker.generateName(stub) != "my-mongodb" {
		t.Errorf("generateName(name): %s should be 'my-mongodb'", docker.generateName(stub))
	}
	docker.NameFormat = "id"
	if docker.generateName(stub) != "bab2b03c736de41ecba6470eba736c5109436f706eedca4f3e0d93d6530eccd4" {
		t.Errorf("generateName(name): %s should be 'bab2b03c736de41ecba6470eba736c5109436f706eedca4f3e0d93d6530eccd4'", docker.generateName(stub))
	}
	docker.NameFormat = "image"
	if docker.generateName(stub) != "tutum/mongodb" {
		t.Errorf("generateName(name): %s should be 'tutum/mongodb'", docker.generateName(stub))
	}
	docker.NameFormat = "image_id"
	if docker.generateName(stub) != "tutum/mongodb_bab2b0" {
		t.Errorf("generateName(name): %s should be 'tutum/mongodb_bab2b0'", docker.generateName(stub))
	}
	docker.NameFormat = "image_name"
	if docker.generateName(stub) != "tutum/mongodb_my-mongodb" {
		t.Errorf("generateName(name): %s should be 'tutum/mongodb_my-mongodb'", docker.generateName(stub))
	}
	docker.NameFormat = "label"
	docker.Label = "foo"
	if docker.generateName(stub) != "bar" {
		t.Errorf("generateName(name): %s should be 'bar'", docker.generateName(stub))
	}

}

func TestAddCPUPercentageStats(t *testing.T) {
	stats := map[string]interface{}{
		"docker.cpuacct.containerA._host":  uint64(100000),
		"docker.cpuacct.containerA.user":   uint64(3000),
		"docker.cpuacct.containerA.system": uint64(2000),
		"docker.cpuacct.containerB._host":  uint64(100000),
		"docker.cpuacct.containerB.user":   uint64(3500),
		"docker.cpuacct.containerC.user":   uint64(3300),
		"docker.cpuacct.containerC.system": uint64(2300),
		"docker.cpuacct.containerD._host":  uint64(100000),
		"docker.cpuacct.containerD.user":   uint64(3000),
		"docker.cpuacct.containerD.system": uint64(2000),
	}
	oldStats := map[string]interface{}{
		"docker.cpuacct.containerA._host":  float64(90000),
		"docker.cpuacct.containerA.user":   float64(1000),
		"docker.cpuacct.containerA.system": float64(1500),
		"docker.cpuacct.containerB._host":  float64(90000),
		"docker.cpuacct.containerB.user":   float64(3000),
		"docker.cpuacct.containerC.user":   float64(3000),
		"docker.cpuacct.containerC.system": float64(2000),
		"docker.cpuacct.containerE._host":  float64(100000),
		"docker.cpuacct.containerE.user":   float64(3000),
		"docker.cpuacct.containerE.system": float64(2000),
	}
	addCPUPercentageStats(&stats, oldStats)

	if stat, ok := stats["docker.cpuacct_percentage.containerA.user"]; !ok {
		t.Errorf("docker.cpuacct_percentage.containerA.user should be calculated")
	} else if stat != float64(20.0) {
		t.Errorf("docker.cpuacct_percentage.containerA.user should be %s, but %s", stat, float64(20.0))
	}

	if _, ok := stats["docker.cpuacct_percentage.containerC.user"]; ok {
		t.Errorf("docker.cpuacct_percentage.containerC.user should not be calculated")
	}

	if _, ok := stats["docker.cpuacct_percentage.containerD.user"]; ok {
		t.Errorf("docker.cpuacct_percentage.containerD.user should not be calculated")
	}

	if _, ok := stats["docker.cpuacct_percentage.containerE.user"]; ok {
		t.Errorf("docker.cpuacct_percentage.containerE.user should not be calculated")
	}
}
<commit_msg>fix tests<commit_after>package mpdocker

import (
	"testing"

	"github.com/fsouza/go-dockerclient"
)

func TestNormalizeMetricName(t *testing.T) {
	testSets := [][]string{
		{"foo/bar", "foo_bar"},
		{"foo:bar", "foo_bar"},
	}

	for _, testSet := range testSets {
		if normalizeMetricName(testSet[0]) != testSet[1] {
			t.Errorf("normalizeMetricName: '%s' should be normalized to '%s', but '%s'", testSet[0], testSet[1], normalizeMetricName(testSet[0]))
		}
	}
}

func TestGraphDefinition(t *testing.T) {
	var docker DockerPlugin

	graphdef := docker.GraphDefinition()
	if len(graphdef) != 6 {
		t.Errorf("GraphDefinition: %d should be 6", len(graphdef))
	}
}

func TestGenerateName(t *testing.T) {
	stub := docker.APIContainers{
		ID:      "bab2b03c736de41ecba6470eba736c5109436f706eedca4f3e0d93d6530eccd4",
		Image:   "tutum/mongodb",
		Command: "/run.sh",
		Created: 1456995574,
		Status:  "Up 4 days",
		Ports: []docker.APIPort{
			{PrivatePort: 28017, Type: "tcp"},
			{PrivatePort: 27017, Type: "tcp"},
		},
		Names:  []string{"/my-mongodb"},
		Labels: map[string]string{"foo": "bar"},
	}
	/* {"Id":"5b963f266d609d2b02aee8f57d664e04d35aa8c23afcbc6bb73bc4a5b2e7c44d",
	   "Image":"memcached",
	   "Command":"/entrypoint.sh memcached",
	   "Created":1456994862,
	   "Status":"Up 4 days",
	   "Ports":[{"PrivatePort":11211,
	   "Type":"tcp"}],
	   "Names":["/my-memcache"]}]`
	*/
	var docker DockerPlugin
	docker.NameFormat = "name_id"
	if docker.generateName(stub) != "my-mongodb_bab2b0" {
		t.Errorf("generateName(name): %s should be 'my-mongodb_bab2b0'", docker.generateName(stub))
	}
	docker.NameFormat = "name"
	if docker.generateName(stub) != "my-mongodb" {
		t.Errorf("generateName(name): %s should be 'my-mongodb'", docker.generateName(stub))
	}
	docker.NameFormat = "id"
	if docker.generateName(stub) != "bab2b03c736de41ecba6470eba736c5109436f706eedca4f3e0d93d6530eccd4" {
		t.Errorf("generateName(name): %s should be 'bab2b03c736de41ecba6470eba736c5109436f706eedca4f3e0d93d6530eccd4'", docker.generateName(stub))
	}
	docker.NameFormat = "image"
	if docker.generateName(stub) != "tutum/mongodb" {
		t.Errorf("generateName(name): %s should be 'tutum/mongodb'", docker.generateName(stub))
	}
	docker.NameFormat = "image_id"
	if docker.generateName(stub) != "tutum/mongodb_bab2b0" {
		t.Errorf("generateName(name): %s should be 'tutum/mongodb_bab2b0'", docker.generateName(stub))
	}
	docker.NameFormat = "image_name"
	if docker.generateName(stub) != "tutum/mongodb_my-mongodb" {
		t.Errorf("generateName(name): %s should be 'tutum/mongodb_my-mongodb'", docker.generateName(stub))
	}
	docker.NameFormat = "label"
	docker.Label = "foo"
	if docker.generateName(stub) != "bar" {
		t.Errorf("generateName(name): %s should be 'bar'", docker.generateName(stub))
	}

}

func TestAddCPUPercentageStats(t *testing.T) {
	stats := map[string]interface{}{
		"docker.cpuacct.containerA._host":       uint64(100000),
		"docker.cpuacct.containerA.user":        uint64(3000),
		"docker.cpuacct.containerA.system":      uint64(2000),
		"docker.cpuacct.containerA._onlineCPUs": int(2),
		"docker.cpuacct.containerB._host":       uint64(100000),
		"docker.cpuacct.containerB.user":        uint64(3500),
		"docker.cpuacct.containerC.user":        uint64(3300),
		"docker.cpuacct.containerC.system":      uint64(2300),
		"docker.cpuacct.containerD._host":       uint64(100000),
		"docker.cpuacct.containerD.user":        uint64(3000),
		"docker.cpuacct.containerD.system":      uint64(2000),
	}
	oldStats := map[string]interface{}{
		"docker.cpuacct.containerA._host":  float64(90000),
		"docker.cpuacct.containerA.user":   float64(1000),
		"docker.cpuacct.containerA.system": float64(1500),
		"docker.cpuacct.containerB._host":  float64(90000),
		"docker.cpuacct.containerB.user":   float64(3000),
		"docker.cpuacct.containerC.user":   float64(3000),
		"docker.cpuacct.containerC.system": float64(2000),
		"docker.cpuacct.containerE._host":  float64(100000),
		"docker.cpuacct.containerE.user":   float64(3000),
		"docker.cpuacct.containerE.system": float64(2000),
	}
	addCPUPercentageStats(&stats, oldStats)

	if stat, ok := stats["docker.cpuacct_percentage.containerA.user"]; !ok {
		t.Errorf("docker.cpuacct_percentage.containerA.user should be calculated")
	} else if stat != float64(40.0) {
		t.Errorf("docker.cpuacct_percentage.containerA.user should be %s, but %s", stat, float64(40.0))
	}

	if _, ok := stats["docker.cpuacct_percentage.containerC.user"]; ok {
		t.Errorf("docker.cpuacct_percentage.containerC.user should not be calculated")
	}

	if _, ok := stats["docker.cpuacct_percentage.containerB.user"]; ok {
		t.Errorf("docker.cpuacct_percentage.containerB.user should not be calculated")
	}

	if _, ok := stats["docker.cpuacct_percentage.containerD.user"]; ok {
		t.Errorf("docker.cpuacct_percentage.containerD.user should not be calculated")
	}

	if _, ok := stats["docker.cpuacct_percentage.containerE.user"]; ok {
		t.Errorf("docker.cpuacct_percentage.containerE.user should not be calculated")
	}
}
<|endoftext|>


<commit_before>package main

import (
	"testing"

	"github.com/erikstmartin/go-testdb"
	"github.com/jmoiron/sqlx"
)

func TestFetchStatDatabase(t *testing.T) {
	db, _ := sqlx.Connect("testdb", "")

	columns := []string{"xact_commit", "xact_rollback", "blks_read", "blks_hit", "blk_read_time", "blk_write_time",
		"tup_returned", "tup_fetched", "tup_inserted", "tup_updated", "tup_deleted", "deadlocks", "temp_bytes"}

	testdb.StubQuery(`
		select xact_commit, xact_rollback, blks_read, blks_hit, blk_read_time, blk_write_time,
		tup_returned, tup_fetched, tup_inserted, tup_updated, tup_deleted, deadlocks, temp_bytes
		from pg_stat_database
	`, testdb.RowsFromCSVString(columns, `
	1,2,3,4,5,6,7,8,9,10,11,12,13
	10,20,30,40,50,60,70,80,90,100,110,120,130
	`))

	stat, err := fetchStatDatabase(db)

	expected := map[string]interface{}{
		"xact_commit":  11.0,
		"blks_hit":     44.0,
		"tup_returned": 77.0,
	}

	if err != nil {
		t.Errorf("Expected no error, but got %s instead", err)
	}
	if err = db.Close(); err != nil {
		t.Errorf("Error '%s' was not expected while closing the database", err)
	}
	if stat["xact_commit"] != expected["xact_commit"] {
		t.Error("should be 11")
	}
	if stat["blks_hit"] != expected["blks_hit"] {
		t.Error("should be 44")
	}
	if stat["tup_returned"] != expected["tup_returned"] {
		t.Error("should be 77")
	}
}
<commit_msg>fix test<commit_after>package main

import (
	"testing"

	"github.com/erikstmartin/go-testdb"
	"github.com/jmoiron/sqlx"
)

func TestFetchStatDatabase(t *testing.T) {
	db, _ := sqlx.Connect("testdb", "")

	columns := []string{"xact_commit", "xact_rollback", "blks_read", "blks_hit", "blk_read_time", "blk_write_time",
		"tup_returned", "tup_fetched", "tup_inserted", "tup_updated", "tup_deleted", "deadlocks", "temp_bytes"}

	testdb.StubQuery(`SELECT * FROM pg_stat_database`, testdb.RowsFromCSVString(columns, `
	1,2,3,4,5,6,7,8,9,10,11,12,13
	10,20,30,40,50,60,70,80,90,100,110,120,130
	`))

	stat, err := fetchStatDatabase(db)

	expected := map[string]interface{}{
		"xact_commit":  uint64(11),
		"blks_hit":     uint64(44),
		"tup_returned": uint64(77),
	}

	if err != nil {
		t.Errorf("Expected no error, but got %s instead", err)
	}
	if err = db.Close(); err != nil {
		t.Errorf("Error '%s' was not expected while closing the database", err)
	}
	if stat["xact_commit"] != expected["xact_commit"] {
		t.Error("should be 11")
	}
	if stat["blks_hit"] != expected["blks_hit"] {
		t.Error("should be 44")
	}
	if stat["tup_returned"] != expected["tup_returned"] {
		t.Error("should be 77")
	}
}
<|endoftext|>


<commit_before>package api_test

import (
	"fmt"
	. "launchpad.net/gocheck"
	"launchpad.net/juju-core/juju/testing"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/state/api"
	coretesting "launchpad.net/juju-core/testing"
	"net"
	"runtime"
	"sync"
	stdtesting "testing"
	"time"
)

func TestAll(t *stdtesting.T) {
	coretesting.MgoTestPackage(t)
}

type suite struct {
	testing.JujuConnSuite
	APIState *api.State
	listener net.Listener
	srv      *api.Server
}

var _ = Suite(&suite{})

func (s *suite) SetUpSuite(c *C) {
	s.JujuConnSuite.SetUpSuite(c)
}

func (s *suite) TearDownSuite(c *C) {
	s.JujuConnSuite.TearDownSuite(c)
}

func (s *suite) SetUpTest(c *C) {
	s.JujuConnSuite.SetUpTest(c)
	var err error
	s.srv, err = api.NewServer(s.State, "localhost:0", []byte(coretesting.ServerCert), []byte(coretesting.ServerKey))
	c.Assert(err, IsNil)
	s.APIState, err = api.Open(&api.Info{
		Addr:   s.srv.Addr(),
		CACert: []byte(coretesting.CACert),
	})
	c.Assert(err, IsNil)
}

func (s *suite) TearDownTest(c *C) {
	err := s.srv.Stop()
	c.Assert(err, IsNil)
	s.JujuConnSuite.TearDownTest(c)
}

func (s *suite) TestRequest(c *C) {
	m, err := s.State.AddMachine(state.MachinerWorker)
	c.Assert(err, IsNil)
	instId, err := s.APIState.Request(m.Id())
	c.Check(instId, Equals, "")
	c.Assert(err, ErrorMatches, "instance id for machine 0 not found")

	err = m.SetInstanceId("foo")
	c.Assert(err, IsNil)

	instId, err = s.APIState.Request(m.Id())
	c.Assert(err, IsNil)
	c.Assert(instId, Equals, "foo")
}
<commit_msg>state/api: add TestStop<commit_after>package api_test

import (
	. "launchpad.net/gocheck"
	"launchpad.net/juju-core/juju/testing"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/state/api"
	coretesting "launchpad.net/juju-core/testing"
	"net"
	stdtesting "testing"
)

func TestAll(t *stdtesting.T) {
	coretesting.MgoTestPackage(t)
}

type suite struct {
	testing.JujuConnSuite
	APIState *api.State
	listener net.Listener
	srv      *api.Server
}

var _ = Suite(&suite{})

func (s *suite) SetUpSuite(c *C) {
	s.JujuConnSuite.SetUpSuite(c)
}

func (s *suite) TearDownSuite(c *C) {
	s.JujuConnSuite.TearDownSuite(c)
}

func (s *suite) SetUpTest(c *C) {
	s.JujuConnSuite.SetUpTest(c)
	var err error
	s.srv, err = api.NewServer(s.State, "localhost:0", []byte(coretesting.ServerCert), []byte(coretesting.ServerKey))
	c.Assert(err, IsNil)
	s.APIState, err = api.Open(&api.Info{
		Addr:   s.srv.Addr(),
		CACert: []byte(coretesting.CACert),
	})
	c.Assert(err, IsNil)
}

func (s *suite) TearDownTest(c *C) {
	err := s.srv.Stop()
	c.Assert(err, IsNil)
	s.JujuConnSuite.TearDownTest(c)
}

func (s *suite) TestRequest(c *C) {
	m, err := s.State.AddMachine(state.MachinerWorker)
	c.Assert(err, IsNil)
	instId, err := s.APIState.Request(m.Id())
	c.Check(instId, Equals, "")
	c.Assert(err, ErrorMatches, "instance id for machine 0 not found")

	err = m.SetInstanceId("foo")
	c.Assert(err, IsNil)

	instId, err = s.APIState.Request(m.Id())
	c.Assert(err, IsNil)
	c.Assert(instId, Equals, "foo")
}

func (s *suite) TestStop(c *C) {
	m, err := s.State.AddMachine(state.MachinerWorker)
	c.Assert(err, IsNil)
	err = m.SetInstanceId("foo")
	c.Assert(err, IsNil)

	err = s.srv.Stop()
	c.Assert(err, IsNil)
	_, err = s.APIState.Request(m.Id())
	c.Assert(err, ErrorMatches, "cannot receive response: EOF")

	// Check it can be stopped twice.
	err = s.srv.Stop()
	c.Assert(err, IsNil)
}
<|endoftext|>


<commit_before>package state_test

import (
	. "launchpad.net/gocheck"
	"launchpad.net/juju/go/state"
	"launchpad.net/juju/go/state/watcher"
	"time"
)

var serviceWatchConfigData = []map[string]interface{}{
	{},
	{"foo": "bar", "baz": "yadda"},
	{"baz": "yadda"},
}

func (s *StateSuite) TestServiceWatchConfig(c *C) {
	dummy := s.addDummyCharm(c)
	wordpress, err := s.st.AddService("wordpress", dummy)
	c.Assert(err, IsNil)
	c.Assert(wordpress.Name(), Equals, "wordpress")

	config, err := wordpress.Config()
	c.Assert(err, IsNil)
	c.Assert(config.Keys(), HasLen, 0)
	configWatcher := wordpress.WatchConfig()

	// Two change events.
	config.Set("foo", "bar")
	config.Set("baz", "yadda")
	_, err = config.Write()
	c.Assert(err, IsNil)
	time.Sleep(100 * time.Millisecond)
	config.Delete("foo")
	_, err = config.Write()
	c.Assert(err, IsNil)

	for _, want := range serviceWatchConfigData {
		select {
		case got, ok := <-configWatcher.Changes():
			c.Assert(ok, Equals, true)
			c.Assert(got.Map(), DeepEquals, want)
		case <-time.After(200 * time.Millisecond):
			c.Fatalf("didn't get change: %#v", want)
		}
	}

	select {
	case got, _ := <-configWatcher.Changes():
		c.Fatalf("got unexpected change: %#v", got)
	case <-time.After(100 * time.Millisecond):
	}

	err = configWatcher.Stop()
	c.Assert(err, IsNil)
}

func (s *StateSuite) TestServiceWatchConfigIllegalData(c *C) {
	dummy := s.addDummyCharm(c)
	wordpress, err := s.st.AddService("wordpress", dummy)
	c.Assert(err, IsNil)
	c.Assert(wordpress.Name(), Equals, "wordpress")
	configWatcher := wordpress.WatchConfig()

	// Receive empty change after service adding.
	select {
	case got, ok := <-configWatcher.Changes():
		c.Assert(ok, Equals, true)
		c.Assert(got.Map(), DeepEquals, map[string]interface{}{})
	case <-time.After(100 * time.Millisecond):
		c.Fatalf("unexpected timeout")
	}

	// Set config to illegal data.
	_, err = s.zkConn.Set("/services/service-0000000000/config", "---", -1)
	c.Assert(err, IsNil)

	select {
	case _, ok := <-configWatcher.Changes():
		c.Assert(ok, Equals, false)
	case <-time.After(100 * time.Millisecond):
	}

	err = configWatcher.Stop()
	c.Assert(err, ErrorMatches, "YAML error: .*")
}

type unitWatchNeedsUpgradeTest struct {
	test func(*state.Unit) error
	want state.NeedsUpgrade
}

var unitWatchNeedsUpgradeTests = []unitWatchNeedsUpgradeTest{
	{func(u *state.Unit) error { return u.SetNeedsUpgrade(false) }, state.NeedsUpgrade{true, false}},
	{func(u *state.Unit) error { return u.ClearNeedsUpgrade() }, state.NeedsUpgrade{false, false}},
	{func(u *state.Unit) error { return u.SetNeedsUpgrade(true) }, state.NeedsUpgrade{true, true}},
}

func (s *StateSuite) TestUnitWatchNeedsUpgrade(c *C) {
	dummy := s.addDummyCharm(c)
	wordpress, err := s.st.AddService("wordpress", dummy)
	c.Assert(err, IsNil)
	c.Assert(wordpress.Name(), Equals, "wordpress")
	unit, err := wordpress.AddUnit()
	c.Assert(err, IsNil)
	needsUpgradeWatcher := unit.WatchNeedsUpgrade()

	for _, test := range unitWatchNeedsUpgradeTests {
		err := test.test(unit)
		c.Assert(err, IsNil)
		select {
		case got, ok := <-needsUpgradeWatcher.Changes():
			c.Assert(ok, Equals, true)
			c.Assert(got, DeepEquals, test.want)
		case <-time.After(200 * time.Millisecond):
			c.Fatalf("didn't get change: %#v", test.want)
		}
	}

	select {
	case got, _ := <-needsUpgradeWatcher.Changes():
		c.Fatalf("got unexpected change: %#v", got)
	case <-time.After(100 * time.Millisecond):
	}

	err = needsUpgradeWatcher.Stop()
	c.Assert(err, IsNil)
}

type unitWatchResolvedTest struct {
	test func(*state.Unit) error
	want state.ResolvedMode
}

var unitWatchResolvedTests = []unitWatchResolvedTest{
	{func(u *state.Unit) error { return u.SetResolved(state.ResolvedRetryHooks) }, state.ResolvedRetryHooks},
	{func(u *state.Unit) error { return u.ClearResolved() }, state.ResolvedNone},
	{func(u *state.Unit) error { return u.SetResolved(state.ResolvedNoHooks) }, state.ResolvedNoHooks},
}

func (s *StateSuite) TestUnitWatchResolved(c *C) {
	dummy := s.addDummyCharm(c)
	wordpress, err := s.st.AddService("wordpress", dummy)
	c.Assert(err, IsNil)
	c.Assert(wordpress.Name(), Equals, "wordpress")
	unit, err := wordpress.AddUnit()
	c.Assert(err, IsNil)
	resolvedWatcher := unit.WatchResolved()

	for _, test := range unitWatchResolvedTests {
		err := test.test(unit)
		c.Assert(err, IsNil)
		select {
		case got, ok := <-resolvedWatcher.Changes():
			c.Assert(ok, Equals, true)
			c.Assert(got, Equals, test.want)
		case <-time.After(200 * time.Millisecond):
			c.Fatalf("didn't get change: %#v", test.want)
		}
	}

	select {
	case got, _ := <-resolvedWatcher.Changes():
		c.Fatalf("got unexpected change: %#v", got)
	case <-time.After(100 * time.Millisecond):
	}

	err = resolvedWatcher.Stop()
	c.Assert(err, IsNil)
}

type unitWatchPortsTest struct {
	test func(*state.Unit) error
	want []state.Port
}

var unitWatchPortsTests = []unitWatchPortsTest{
	{func(u *state.Unit) error { return u.OpenPort("tcp", 80) }, []state.Port{{"tcp", 80}}},
	{func(u *state.Unit) error { return u.OpenPort("udp", 53) }, []state.Port{{"tcp", 80}, {"udp", 53}}},
	{func(u *state.Unit) error { return u.ClosePort("tcp", 80) }, []state.Port{{"udp", 53}}},
}

func (s *StateSuite) TestUnitWatchPorts(c *C) {
	dummy := s.addDummyCharm(c)
	wordpress, err := s.st.AddService("wordpress", dummy)
	c.Assert(err, IsNil)
	c.Assert(wordpress.Name(), Equals, "wordpress")
	unit, err := wordpress.AddUnit()
	c.Assert(err, IsNil)
	portsWatcher := unit.WatchPorts()

	for _, test := range unitWatchPortsTests {
		err := test.test(unit)
		c.Assert(err, IsNil)
		select {
		case got, ok := <-portsWatcher.Changes():
			c.Assert(ok, Equals, true)
			c.Assert(got, DeepEquals, test.want)
		case <-time.After(200 * time.Millisecond):
			c.Fatalf("didn't get change: %#v", test.want)
		}
	}

	select {
	case got, _ := <-portsWatcher.Changes():
		c.Fatalf("got unexpected change: %#v", got)
	case <-time.After(100 * time.Millisecond):
	}

	err = portsWatcher.Stop()
	c.Assert(err, IsNil)
}

type machinesWatchTest struct {
	test func(*state.State) error
	want watcher.ChildrenChange
}

var machinesWatchTests = []machinesWatchTest{
	{func(s *state.State) error { _, err := s.AddMachine(); return err }, watcher.ChildrenChange{Added: []string{"machine-0000000000"}}},
	{func(s *state.State) error { _, err := s.AddMachine(); return err }, watcher.ChildrenChange{Added: []string{"machine-0000000001"}}},
	{func(s *state.State) error { return s.RemoveMachine(1) }, watcher.ChildrenChange{Deleted: []string{"machine-0000000001"}}},
}

func (s *StateSuite) TestWatchMachines(c *C) {
	w := s.st.WatchMachines()

	for _, test := range machinesWatchTests {
		err := test.test(s.st)
		c.Assert(err, IsNil)
		select {
		case got, ok := <-w.Changes:
			c.Assert(ok, Equals, true)
			c.Assert(got, DeepEquals, test.want)
		case <-time.After(200 * time.Millisecond):
			c.Fatalf("didn't get change: %#v", test.want)
		}
	}

	select {
	case got, _ := <-w.Changes:
		c.Fatalf("got unexpected change: %#v", got)
	case <-time.After(100 * time.Millisecond):
	}

	c.Assert(w.Stop(), IsNil)
}
<commit_msg>reformat tests<commit_after>package state_test

import (
	. "launchpad.net/gocheck"
	"launchpad.net/juju/go/state"
	"launchpad.net/juju/go/state/watcher"
	"time"
)

var serviceWatchConfigData = []map[string]interface{}{
	{},
	{"foo": "bar", "baz": "yadda"},
	{"baz": "yadda"},
}

func (s *StateSuite) TestServiceWatchConfig(c *C) {
	dummy := s.addDummyCharm(c)
	wordpress, err := s.st.AddService("wordpress", dummy)
	c.Assert(err, IsNil)
	c.Assert(wordpress.Name(), Equals, "wordpress")

	config, err := wordpress.Config()
	c.Assert(err, IsNil)
	c.Assert(config.Keys(), HasLen, 0)
	configWatcher := wordpress.WatchConfig()

	// Two change events.
	config.Set("foo", "bar")
	config.Set("baz", "yadda")
	_, err = config.Write()
	c.Assert(err, IsNil)
	time.Sleep(100 * time.Millisecond)
	config.Delete("foo")
	_, err = config.Write()
	c.Assert(err, IsNil)

	for _, want := range serviceWatchConfigData {
		select {
		case got, ok := <-configWatcher.Changes():
			c.Assert(ok, Equals, true)
			c.Assert(got.Map(), DeepEquals, want)
		case <-time.After(200 * time.Millisecond):
			c.Fatalf("didn't get change: %#v", want)
		}
	}

	select {
	case got, _ := <-configWatcher.Changes():
		c.Fatalf("got unexpected change: %#v", got)
	case <-time.After(100 * time.Millisecond):
	}

	err = configWatcher.Stop()
	c.Assert(err, IsNil)
}

func (s *StateSuite) TestServiceWatchConfigIllegalData(c *C) {
	dummy := s.addDummyCharm(c)
	wordpress, err := s.st.AddService("wordpress", dummy)
	c.Assert(err, IsNil)
	c.Assert(wordpress.Name(), Equals, "wordpress")
	configWatcher := wordpress.WatchConfig()

	// Receive empty change after service adding.
	select {
	case got, ok := <-configWatcher.Changes():
		c.Assert(ok, Equals, true)
		c.Assert(got.Map(), DeepEquals, map[string]interface{}{})
	case <-time.After(100 * time.Millisecond):
		c.Fatalf("unexpected timeout")
	}

	// Set config to illegal data.
	_, err = s.zkConn.Set("/services/service-0000000000/config", "---", -1)
	c.Assert(err, IsNil)

	select {
	case _, ok := <-configWatcher.Changes():
		c.Assert(ok, Equals, false)
	case <-time.After(100 * time.Millisecond):
	}

	err = configWatcher.Stop()
	c.Assert(err, ErrorMatches, "YAML error: .*")
}

type unitWatchNeedsUpgradeTest struct {
	test func(*state.Unit) error
	want state.NeedsUpgrade
}

var unitWatchNeedsUpgradeTests = []unitWatchNeedsUpgradeTest{
	{func(u *state.Unit) error { return u.SetNeedsUpgrade(false) }, state.NeedsUpgrade{true, false}},
	{func(u *state.Unit) error { return u.ClearNeedsUpgrade() }, state.NeedsUpgrade{false, false}},
	{func(u *state.Unit) error { return u.SetNeedsUpgrade(true) }, state.NeedsUpgrade{true, true}},
}

func (s *StateSuite) TestUnitWatchNeedsUpgrade(c *C) {
	dummy := s.addDummyCharm(c)
	wordpress, err := s.st.AddService("wordpress", dummy)
	c.Assert(err, IsNil)
	c.Assert(wordpress.Name(), Equals, "wordpress")
	unit, err := wordpress.AddUnit()
	c.Assert(err, IsNil)
	needsUpgradeWatcher := unit.WatchNeedsUpgrade()

	for _, test := range unitWatchNeedsUpgradeTests {
		err := test.test(unit)
		c.Assert(err, IsNil)
		select {
		case got, ok := <-needsUpgradeWatcher.Changes():
			c.Assert(ok, Equals, true)
			c.Assert(got, DeepEquals, test.want)
		case <-time.After(200 * time.Millisecond):
			c.Fatalf("didn't get change: %#v", test.want)
		}
	}

	select {
	case got, _ := <-needsUpgradeWatcher.Changes():
		c.Fatalf("got unexpected change: %#v", got)
	case <-time.After(100 * time.Millisecond):
	}

	err = needsUpgradeWatcher.Stop()
	c.Assert(err, IsNil)
}

type unitWatchResolvedTest struct {
	test func(*state.Unit) error
	want state.ResolvedMode
}

var unitWatchResolvedTests = []unitWatchResolvedTest{
	{func(u *state.Unit) error { return u.SetResolved(state.ResolvedRetryHooks) }, state.ResolvedRetryHooks},
	{func(u *state.Unit) error { return u.ClearResolved() }, state.ResolvedNone},
	{func(u *state.Unit) error { return u.SetResolved(state.ResolvedNoHooks) }, state.ResolvedNoHooks},
}

func (s *StateSuite) TestUnitWatchResolved(c *C) {
	dummy := s.addDummyCharm(c)
	wordpress, err := s.st.AddService("wordpress", dummy)
	c.Assert(err, IsNil)
	c.Assert(wordpress.Name(), Equals, "wordpress")
	unit, err := wordpress.AddUnit()
	c.Assert(err, IsNil)
	resolvedWatcher := unit.WatchResolved()

	for _, test := range unitWatchResolvedTests {
		err := test.test(unit)
		c.Assert(err, IsNil)
		select {
		case got, ok := <-resolvedWatcher.Changes():
			c.Assert(ok, Equals, true)
			c.Assert(got, Equals, test.want)
		case <-time.After(200 * time.Millisecond):
			c.Fatalf("didn't get change: %#v", test.want)
		}
	}

	select {
	case got, _ := <-resolvedWatcher.Changes():
		c.Fatalf("got unexpected change: %#v", got)
	case <-time.After(100 * time.Millisecond):
	}

	err = resolvedWatcher.Stop()
	c.Assert(err, IsNil)
}

type unitWatchPortsTest struct {
	test func(*state.Unit) error
	want []state.Port
}

var unitWatchPortsTests = []unitWatchPortsTest{
	{func(u *state.Unit) error { return u.OpenPort("tcp", 80) }, []state.Port{{"tcp", 80}}},
	{func(u *state.Unit) error { return u.OpenPort("udp", 53) }, []state.Port{{"tcp", 80}, {"udp", 53}}},
	{func(u *state.Unit) error { return u.ClosePort("tcp", 80) }, []state.Port{{"udp", 53}}},
}

func (s *StateSuite) TestUnitWatchPorts(c *C) {
	dummy := s.addDummyCharm(c)
	wordpress, err := s.st.AddService("wordpress", dummy)
	c.Assert(err, IsNil)
	c.Assert(wordpress.Name(), Equals, "wordpress")
	unit, err := wordpress.AddUnit()
	c.Assert(err, IsNil)
	portsWatcher := unit.WatchPorts()

	for _, test := range unitWatchPortsTests {
		err := test.test(unit)
		c.Assert(err, IsNil)
		select {
		case got, ok := <-portsWatcher.Changes():
			c.Assert(ok, Equals, true)
			c.Assert(got, DeepEquals, test.want)
		case <-time.After(200 * time.Millisecond):
			c.Fatalf("didn't get change: %#v", test.want)
		}
	}

	select {
	case got, _ := <-portsWatcher.Changes():
		c.Fatalf("got unexpected change: %#v", got)
	case <-time.After(100 * time.Millisecond):
	}

	err = portsWatcher.Stop()
	c.Assert(err, IsNil)
}

type machinesWatchTest struct {
	test func(*state.State) error
	want watcher.ChildrenChange
}

var machinesWatchTests = []machinesWatchTest{
	{
		func(s *state.State) error { _, err := s.AddMachine(); return err },
		watcher.ChildrenChange{
			Added: []string{"machine-0000000000"},
		},
	},
	{
		func(s *state.State) error { _, err := s.AddMachine(); return err },
		watcher.ChildrenChange{
			Added: []string{"machine-0000000001"},
		},
	},
	{
		func(s *state.State) error { return s.RemoveMachine(1) },
		watcher.ChildrenChange{
			Deleted: []string{"machine-0000000001"},
		},
	},
}

func (s *StateSuite) TestWatchMachines(c *C) {
	w := s.st.WatchMachines()

	for _, test := range machinesWatchTests {
		err := test.test(s.st)
		c.Assert(err, IsNil)
		select {
		case got, ok := <-w.Changes:
			c.Assert(ok, Equals, true)
			c.Assert(got, DeepEquals, test.want)
		case <-time.After(200 * time.Millisecond):
			c.Fatalf("didn't get change: %#v", test.want)
		}
	}

	select {
	case got, _ := <-w.Changes:
		c.Fatalf("got unexpected change: %#v", got)
	case <-time.After(100 * time.Millisecond):
	}

	c.Assert(w.Stop(), IsNil)
}
<|endoftext|>


<commit_before>package point

import (
	"sync"
	"time"

	"github.com/v2ray/v2ray-core/app"
	"github.com/v2ray/v2ray-core/common/dice"
	"github.com/v2ray/v2ray-core/common/log"
	v2net "github.com/v2ray/v2ray-core/common/net"
	"github.com/v2ray/v2ray-core/common/retry"
	"github.com/v2ray/v2ray-core/proxy"
	proxyrepo "github.com/v2ray/v2ray-core/proxy/repo"
)

type InboundDetourHandlerDynamic struct {
	sync.RWMutex
	space       app.Space
	config      *InboundDetourConfig
	portsInUse  map[v2net.Port]bool
	ichInUse    []*InboundConnectionHandlerWithPort
	ich2Recycle []*InboundConnectionHandlerWithPort
	lastRefresh time.Time
	started     bool
}

func NewInboundDetourHandlerDynamic(space app.Space, config *InboundDetourConfig) (*InboundDetourHandlerDynamic, error) {
	handler := &InboundDetourHandlerDynamic{
		space:      space,
		config:     config,
		portsInUse: make(map[v2net.Port]bool),
	}
	if err := handler.refresh(); err != nil {
		return nil, err
	}
	return handler, nil
}

func (this *InboundDetourHandlerDynamic) refresh() error {
	this.Lock()
	defer this.Unlock()

	this.ich2Recycle = this.ichInUse
	if this.ich2Recycle != nil {
		time.AfterFunc(time.Minute, func() {
			for _, ich := range this.ich2Recycle {
				if ich != nil {
					ich.handler.Close()
					delete(this.portsInUse, ich.port)
				}
			}
		})
	}

	ichCount := this.config.Allocation.Concurrency
	// TODO: check ichCount
	this.ichInUse = make([]*InboundConnectionHandlerWithPort, ichCount)
	for idx, _ := range this.ichInUse {
		port := this.pickUnusedPort()
		ich, err := proxyrepo.CreateInboundConnectionHandler(this.config.Protocol, this.space, this.config.Settings)
		if err != nil {
			log.Error("Point: Failed to create inbound connection handler: ", err)
			return err
		}
		this.ichInUse[idx] = &InboundConnectionHandlerWithPort{
			port:    port,
			handler: ich,
		}
	}
	if this.started {
		this.Start()
	}

	this.lastRefresh = time.Now()
	time.AfterFunc(time.Duration(this.config.Allocation.Refresh)*time.Minute, func() {
		this.refresh()
	})

	return nil
}

func (this *InboundDetourHandlerDynamic) pickUnusedPort() v2net.Port {
	delta := int(this.config.PortRange.To) - int(this.config.PortRange.From) + 1
	for {
		r := dice.Roll(delta)
		port := this.config.PortRange.From + v2net.Port(r)
		_, used := this.portsInUse[port]
		if !used {
			this.portsInUse[port] = true
			return port
		}
	}
}

func (this *InboundDetourHandlerDynamic) GetConnectionHandler() (proxy.InboundConnectionHandler, int) {
	this.RLock()
	defer this.RUnlock()
	ich := this.ichInUse[dice.Roll(len(this.ichInUse))]
	until := this.config.Allocation.Refresh - int((time.Now().Unix()-this.lastRefresh.Unix())/60/1000)
	if until < 0 {
		until = 0
	}
	return ich.handler, int(until)
}

func (this *InboundDetourHandlerDynamic) Close() {
	this.Lock()
	defer this.Unlock()
	for _, ich := range this.ichInUse {
		ich.handler.Close()
	}
	if this.ich2Recycle != nil {
		for _, ich := range this.ich2Recycle {
			if ich != nil && ich.handler != nil {
				ich.handler.Close()
			}
		}
	}
}

func (this *InboundDetourHandlerDynamic) Start() error {
	for _, ich := range this.ichInUse {
		err := retry.Timed(100 /* times */, 100 /* ms */).On(func() error {
			err := ich.handler.Listen(ich.port)
			if err != nil {
				log.Error("Point: Failed to start inbound detour on port ", ich.port, ": ", err)
				return err
			}
			return nil
		})
		if err != nil {
			return err
		}
	}
	this.started = true
	return nil
}
<commit_msg>reuse buffer<commit_after>package point

import (
	"sync"
	"time"

	"github.com/v2ray/v2ray-core/app"
	"github.com/v2ray/v2ray-core/common/dice"
	"github.com/v2ray/v2ray-core/common/log"
	v2net "github.com/v2ray/v2ray-core/common/net"
	"github.com/v2ray/v2ray-core/common/retry"
	"github.com/v2ray/v2ray-core/proxy"
	proxyrepo "github.com/v2ray/v2ray-core/proxy/repo"
)

type InboundDetourHandlerDynamic struct {
	sync.RWMutex
	space       app.Space
	config      *InboundDetourConfig
	portsInUse  map[v2net.Port]bool
	ichInUse    []*InboundConnectionHandlerWithPort
	ich2Recycle []*InboundConnectionHandlerWithPort
	lastRefresh time.Time
	started     bool
}

func NewInboundDetourHandlerDynamic(space app.Space, config *InboundDetourConfig) (*InboundDetourHandlerDynamic, error) {
	handler := &InboundDetourHandlerDynamic{
		space:      space,
		config:     config,
		portsInUse: make(map[v2net.Port]bool),
	}
	if err := handler.refresh(); err != nil {
		return nil, err
	}
	return handler, nil
}

func (this *InboundDetourHandlerDynamic) refresh() error {
	this.Lock()
	defer this.Unlock()

	this.ich2Recycle, this.ichInUse = this.ichInUse, this.ich2Recycle
	if this.ich2Recycle != nil {
		time.AfterFunc(time.Minute, func() {
			for i := 0; i < len(this.ich2Recycle); i++ {
				ich := this.ich2Recycle[i]
				if ich != nil {
					ich.handler.Close()
					delete(this.portsInUse, ich.port)
				}
				this.ich2Recycle[i] = nil
			}
		})
	}

	ichCount := this.config.Allocation.Concurrency
	// TODO: check ichCount
	if this.ichInUse == nil {
		this.ichInUse = make([]*InboundConnectionHandlerWithPort, ichCount)
	}

	for idx, _ := range this.ichInUse {
		port := this.pickUnusedPort()
		ich, err := proxyrepo.CreateInboundConnectionHandler(this.config.Protocol, this.space, this.config.Settings)
		if err != nil {
			log.Error("Point: Failed to create inbound connection handler: ", err)
			return err
		}
		this.ichInUse[idx] = &InboundConnectionHandlerWithPort{
			port:    port,
			handler: ich,
		}
	}
	if this.started {
		this.Start()
	}

	this.lastRefresh = time.Now()
	time.AfterFunc(time.Duration(this.config.Allocation.Refresh)*time.Minute, func() {
		this.refresh()
	})

	return nil
}

func (this *InboundDetourHandlerDynamic) pickUnusedPort() v2net.Port {
	delta := int(this.config.PortRange.To) - int(this.config.PortRange.From) + 1
	for {
		r := dice.Roll(delta)
		port := this.config.PortRange.From + v2net.Port(r)
		_, used := this.portsInUse[port]
		if !used {
			this.portsInUse[port] = true
			return port
		}
	}
}

func (this *InboundDetourHandlerDynamic) GetConnectionHandler() (proxy.InboundConnectionHandler, int) {
	this.RLock()
	defer this.RUnlock()
	ich := this.ichInUse[dice.Roll(len(this.ichInUse))]
	until := this.config.Allocation.Refresh - int((time.Now().Unix()-this.lastRefresh.Unix())/60/1000)
	if until < 0 {
		until = 0
	}
	return ich.handler, int(until)
}

func (this *InboundDetourHandlerDynamic) Close() {
	this.Lock()
	defer this.Unlock()
	for _, ich := range this.ichInUse {
		ich.handler.Close()
	}
	if this.ich2Recycle != nil {
		for _, ich := range this.ich2Recycle {
			if ich != nil && ich.handler != nil {
				ich.handler.Close()
			}
		}
	}
}

func (this *InboundDetourHandlerDynamic) Start() error {
	for _, ich := range this.ichInUse {
		err := retry.Timed(100 /* times */, 100 /* ms */).On(func() error {
			err := ich.handler.Listen(ich.port)
			if err != nil {
				log.Error("Point: Failed to start inbound detour on port ", ich.port, ": ", err)
				return err
			}
			return nil
		})
		if err != nil {
			return err
		}
	}
	this.started = true
	return nil
}
<|endoftext|>


<commit_before>package upgrade

import (
	"fmt"

	"github.com/Sirupsen/logrus"
	"github.com/docker/libcompose/project"
	rancherClient "github.com/rancherio/go-rancher/client"
	"github.com/rancherio/rancher-compose/rancher"
)

type UpgradeOpts struct {
	BatchSize      int
	IntervalMillis int
	FinalScale     int
	UpdateLinks    bool
	Wait           bool
}

func Upgrade(p *project.Project, from, to string, opts UpgradeOpts) error {
	fromService, err := p.CreateService(from)
	if err != nil {
		return err
	}

	toService, err := p.CreateService(to)
	if err != nil {
		return err
	}

	rFromService, ok := fromService.(*rancher.RancherService)
	if !ok {
		return fmt.Errorf("%s is not a Rancher service", from)
	}

	rToService, ok := toService.(*rancher.RancherService)
	if !ok {
		return fmt.Errorf("%s is not a Rancher service", to)
	}

	if err := rToService.Create(); err != nil {
		return err
	}

	if err := rToService.Scale(0); err != nil {
		return err
	}

	if err := rToService.Up(); err != nil {
		return err
	}

	source, err := rFromService.RancherService()
	if err != nil {
		return err
	}

	dest, err := rToService.RancherService()
	if err != nil {
		return err
	}

	if source == nil {
		return fmt.Errorf("Failed to find service %s", from)
	}

	if dest == nil {
		return fmt.Errorf("Failed to find service %s", to)
	}

	upgradeOpts := &rancherClient.ServiceUpgrade{
		UpdateLinks:    opts.UpdateLinks,
		FinalScale:     int64(opts.FinalScale),
		BatchSize:      int64(opts.BatchSize),
		IntervalMillis: int64(opts.IntervalMillis),
		ToServiceId:    dest.Id,
	}
	if upgradeOpts.FinalScale == -1 {
		upgradeOpts.FinalScale = source.Scale
	}

	client := rFromService.Client()

	logrus.Infof("Upgrading %s to %s, scale=%d", from, to, upgradeOpts.FinalScale)
	service, err := client.Service.ActionUpgrade(source, upgradeOpts)
	if err != nil {
		return err
	}

	if opts.Wait {
		return rFromService.Wait(service)
	}

	return nil
}

func upgradeInfo(up bool, p *project.Project, from, to string, opts UpgradeOpts) (*rancherClient.Service, *rancherClient.Service, *rancherClient.RancherClient, error) {
	fromService, err := p.CreateService(from)
	if err != nil {
		return nil, nil, nil, err
	}

	toService, err := p.CreateService(to)
	if err != nil {
		return nil, nil, nil, err
	}

	rFromService, ok := fromService.(*rancher.RancherService)
	if !ok {
		return nil, nil, nil, fmt.Errorf("%s is not a Rancher service", from)
	}

	rToService, ok := toService.(*rancher.RancherService)
	if !ok {
		return nil, nil, nil, fmt.Errorf("%s is not a Rancher service", to)
	}

	if up {
		if err := rToService.Up(); err != nil {
			return nil, nil, nil, err
		}
	}

	source, err := rFromService.RancherService()
	if err != nil {
		return nil, nil, nil, err
	}

	dest, err := rToService.RancherService()
	if err != nil {
		return nil, nil, nil, err
	}

	return source, dest, rFromService.Client(), nil
}
<commit_msg>Fix setting target scale to zero on start<commit_after>package upgrade

import (
	"fmt"

	"github.com/Sirupsen/logrus"
	"github.com/docker/libcompose/project"
	rancherClient "github.com/rancherio/go-rancher/client"
	"github.com/rancherio/rancher-compose/rancher"
)

type UpgradeOpts struct {
	BatchSize      int
	IntervalMillis int
	FinalScale     int
	UpdateLinks    bool
	Wait           bool
}

func Upgrade(p *project.Project, from, to string, opts UpgradeOpts) error {
	fromService, err := p.CreateService(from)
	if err != nil {
		return err
	}

	toService, err := p.CreateService(to)
	if err != nil {
		return err
	}

	rFromService, ok := fromService.(*rancher.RancherService)
	if !ok {
		return fmt.Errorf("%s is not a Rancher service", from)
	}

	rToService, ok := toService.(*rancher.RancherService)
	if !ok {
		return fmt.Errorf("%s is not a Rancher service", to)
	}

	if service, err := rToService.RancherService(); err != nil {
		return err
	} else if service == nil {
		if err := rToService.Create(); err != nil {
			return err
		}

		if err := rToService.Scale(0); err != nil {
			return err
		}
	}

	if err := rToService.Up(); err != nil {
		return err
	}

	source, err := rFromService.RancherService()
	if err != nil {
		return err
	}

	dest, err := rToService.RancherService()
	if err != nil {
		return err
	}

	if source == nil {
		return fmt.Errorf("Failed to find service %s", from)
	}

	if dest == nil {
		return fmt.Errorf("Failed to find service %s", to)
	}

	upgradeOpts := &rancherClient.ServiceUpgrade{
		UpdateLinks:    opts.UpdateLinks,
		FinalScale:     int64(opts.FinalScale),
		BatchSize:      int64(opts.BatchSize),
		IntervalMillis: int64(opts.IntervalMillis),
		ToServiceId:    dest.Id,
	}
	if upgradeOpts.FinalScale == -1 {
		upgradeOpts.FinalScale = source.Scale
	}

	client := rFromService.Client()

	logrus.Infof("Upgrading %s to %s, scale=%d", from, to, upgradeOpts.FinalScale)
	service, err := client.Service.ActionUpgrade(source, upgradeOpts)
	if err != nil {
		return err
	}

	if opts.Wait {
		return rFromService.Wait(service)
	}

	return nil
}

func upgradeInfo(up bool, p *project.Project, from, to string, opts UpgradeOpts) (*rancherClient.Service, *rancherClient.Service, *rancherClient.RancherClient, error) {
	fromService, err := p.CreateService(from)
	if err != nil {
		return nil, nil, nil, err
	}

	toService, err := p.CreateService(to)
	if err != nil {
		return nil, nil, nil, err
	}

	rFromService, ok := fromService.(*rancher.RancherService)
	if !ok {
		return nil, nil, nil, fmt.Errorf("%s is not a Rancher service", from)
	}

	rToService, ok := toService.(*rancher.RancherService)
	if !ok {
		return nil, nil, nil, fmt.Errorf("%s is not a Rancher service", to)
	}

	if up {
		if err := rToService.Up(); err != nil {
			return nil, nil, nil, err
		}
	}

	source, err := rFromService.RancherService()
	if err != nil {
		return nil, nil, nil, err
	}

	dest, err := rToService.RancherService()
	if err != nil {
		return nil, nil, nil, err
	}

	return source, dest, rFromService.Client(), nil
}
<|endoftext|>


<commit_before>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package upstart

import (
	"fmt"
	"path/filepath"

	"launchpad.net/juju-core/utils"
)

const (
	maxMongoFiles = 65000
	maxAgentFiles = 20000
)

// MongoUpstartService returns the upstart config for the mongo state service.
func MongoUpstartService(name, dataDir, dbDir string, port int) *Conf {
	keyFile := filepath.Join(dataDir, "server.pem")
	svc := NewService(name)
	return &Conf{
		Service: *svc,
		Desc:    "juju state database",
		Limit: map[string]string{
			"nofile": fmt.Sprintf("%d %d", maxMongoFiles, maxMongoFiles),
			"nproc":  fmt.Sprintf("%d %d", maxAgentFiles, maxAgentFiles),
		},
		Cmd: "/usr/bin/mongod" +
			" --auth" +
			" --dbpath=" + dbDir +
			" --sslOnNormalPorts" +
			" --sslPEMKeyFile " + utils.ShQuote(keyFile) +
			" --sslPEMKeyPassword ignored" +
			" --bind_ip 0.0.0.0" +
			" --port " + fmt.Sprint(port) +
			" --noprealloc" +
			" --syslog" +
			" --smallfiles",
	}
}

// MachineAgentUpstartService returns the upstart config for a machine agent
// based on the tag and machineId passed in.
func MachineAgentUpstartService(name, toolsDir, dataDir, logDir, tag, machineId, logConfig string) *Conf {
	svc := NewService(name)
	logFile := filepath.Join(logDir, tag+".log")
	return &Conf{
		Service: *svc,
		Desc:    fmt.Sprintf("juju %s agent", tag),
		Limit: map[string]string{
			"nofile": fmt.Sprintf("%d %d", maxAgentFiles, maxAgentFiles),
		},
		Cmd: cmd,
		Out: filepath.Join(toolsDir, "jujud") +
			" machine" +
			" --log-file " + utils.ShQuote(logFile) +
			" --data-dir " + utils.ShQuote(dataDir) +
			" --machine-id " + machineId +
			" " + logConfig,
	}
}
<commit_msg>Restructure just a tad more.<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package upstart

import (
	"fmt"
	"path/filepath"

	"launchpad.net/juju-core/utils"
)

const (
	maxMongoFiles = 65000
	maxAgentFiles = 20000
)

// MongoUpstartService returns the upstart config for the mongo state service.
func MongoUpstartService(name, dataDir, dbDir string, port int) *Conf {
	keyFile := filepath.Join(dataDir, "server.pem")
	svc := NewService(name)
	return &Conf{
		Service: *svc,
		Desc:    "juju state database",
		Limit: map[string]string{
			"nofile": fmt.Sprintf("%d %d", maxMongoFiles, maxMongoFiles),
			"nproc":  fmt.Sprintf("%d %d", maxAgentFiles, maxAgentFiles),
		},
		Cmd: "/usr/bin/mongod" +
			" --auth" +
			" --dbpath=" + dbDir +
			" --sslOnNormalPorts" +
			" --sslPEMKeyFile " + utils.ShQuote(keyFile) +
			" --sslPEMKeyPassword ignored" +
			" --bind_ip 0.0.0.0" +
			" --port " + fmt.Sprint(port) +
			" --noprealloc" +
			" --syslog" +
			" --smallfiles",
	}
}

// MachineAgentUpstartService returns the upstart config for a machine agent
// based on the tag and machineId passed in.
func MachineAgentUpstartService(name, toolsDir, dataDir, logDir, tag, machineId, logConfig string) *Conf {
	svc := NewService(name)
	return &Conf{
		Service: *svc,
		Desc:    fmt.Sprintf("juju %s agent", tag),
		Limit: map[string]string{
			"nofile": fmt.Sprintf("%d %d", maxAgentFiles, maxAgentFiles),
		},
		Cmd: filepath.Join(toolsDir, "jujud") +
			" machine" +
			" --log-file " + utils.ShQuote(logFile) +
			" --data-dir " + utils.ShQuote(dataDir) +
			" --machine-id " + machineId +
			" " + logConfig,
		Out: filepath.Join(logDir, tag+".log"),
	}
}
<|endoftext|>


<commit_before>package tempFluc

import (
	"math"
)
import (
	"../bzone"
	"../tempAll"
	"../tempCrit"
	vec "../vector"
)

// Calculate U_{1}/N = 1/N \sum_k \epsilon_h(k) f_h(\xi_h(k))
func HolonEnergy(env *tempAll.Environment) (float64, error) {
	inner := func(k vec.Vector) float64 {
		return env.Epsilon_h(k) * env.Fermi(env.Xi_h(k))
	}
	dim := 2
	avg := bzone.Avg(env.PointsPerSide, dim, inner)
	return avg, nil
}

// Calculate U_{2}/N = 1/N \sum_k (\omega_+(k) + \mu_b) n_b(\omega_+(k))
func PairEnergy(env *tempAll.Environment) (float64, error) {
	/*
		// kz^2 version
		oc, err := tempCrit.OmegaFit(env, tempCrit.OmegaPlus)
		if err != nil {
			return 0.0, err
		}
		integrand := func(y float64) float64 {
			num := math.Pow(y, 1.5)
			denom := math.Exp(y-env.Beta*env.Mu_b) - 1.0
			return num / denom
		}
		integral, err := tempCrit.OmegaIntegralY(env, oc, integrand)
		if err != nil {
			return 0.0, err
		}
		return integral / math.Pow(env.Beta, 2.5), nil
	*/
	// cos(kz) version
	oc, err := tempCrit.OmegaFit(env, tempCrit.OmegaPlus)
	if err != nil {
		return 0.0, err
	}
	integrand := func(y, kz float64) float64 {
		bterm := 2.0 * oc[2] * (1.0 - math.Cos(kz))
		num := y/env.Beta + bterm
		denom := math.Exp(y+env.Beta*(bterm-env.Mu_b)) - 1.0
		return num / denom
	}
	integral, err := tempCrit.OmegaIntegralCos(env, oc, integrand)
	if err != nil {
		return 0.0, err
	}
	return integral, nil
}
<commit_msg>implement pair energy for finite magnetic field<commit_after>package tempFluc

import (
	"math"
	"fmt"
)
import (
	"../bzone"
	"../tempAll"
	"../tempCrit"
	vec "../vector"
	"../bessel"
	"../seriesaccel"
)

// Calculate U_{1}/N = 1/N \sum_k \epsilon_h(k) f_h(\xi_h(k))
func HolonEnergy(env *tempAll.Environment) (float64, error) {
	inner := func(k vec.Vector) float64 {
		return env.Epsilon_h(k) * env.Fermi(env.Xi_h(k))
	}
	dim := 2
	avg := bzone.Avg(env.PointsPerSide, dim, inner)
	return avg, nil
}

// Calculate U_{2}/N = 1/N \sum_k (\omega_+(k) + \mu_b) n_b(\omega_+(k))
func PairEnergy(env *tempAll.Environment) (float64, error) {
	/*
		// kz^2 version
		oc, err := tempCrit.OmegaFit(env, tempCrit.OmegaPlus)
		if err != nil {
			return 0.0, err
		}
		integrand := func(y float64) float64 {
			num := math.Pow(y, 1.5)
			denom := math.Exp(y-env.Beta*env.Mu_b) - 1.0
			return num / denom
		}
		integral, err := tempCrit.OmegaIntegralY(env, oc, integrand)
		if err != nil {
			return 0.0, err
		}
		return integral / math.Pow(env.Beta, 2.5), nil
	*/
	// cos(kz) version
	oc, err := tempCrit.OmegaFit(env, tempCrit.OmegaPlus)
	if err != nil {
		return 0.0, err
	}
	if math.Abs(env.Be_field) > 1e-9 {
		integrand := func(y, kz float64) float64 {
			bterm := 2.0 * oc[2] * (1.0 - math.Cos(kz))
			num := y/env.Beta + bterm
			denom := math.Exp(y+env.Beta*(bterm-env.Mu_b)) - 1.0
			return num / denom
		}
		integral, err := tempCrit.OmegaIntegralCos(env, oc, integrand)
		if err != nil {
			return 0.0, err
		}
		return integral, nil
	}
	// if we get here, math.Abs(env.Be_field) >= 1e-9
	E2BSumTerm := func(ri int) float64 {
		r := float64(ri)
		a, b := oc[0], oc[2]
		I0 := bessel.ModifiedBesselFirstKindZeroth(2.0 * b * env.Beta * r)
		I1 := bessel.ModifiedBesselFirstKindFirst(2.0 * b * env.Beta * r)
		omega_c := 4.0 * env.Be_field * a
		mu_tilde := env.Mu_b - omega_c/2.0
		expL := math.Exp(r * env.Beta * (mu_tilde - 2.0*b))
		expR := math.Exp(-env.Beta * omega_c * r)
		expm1 := -math.Expm1(-env.Beta * omega_c * r)
		return expL * ((I0 * (0.5 + 2.0*b) - 2.0*b*I1)*expm1 + (I0 * omega_c * expR * expm1 * expm1))
	}
	sum, absErr := seriesaccel.Levin_u(E2BSumTerm, 1, 20)
	fmt.Printf("E2 B sum %e, absErr %e\n", sum, absErr)
	return 2.0 * env.Be_field * sum / math.Pi, nil

}
<|endoftext|>


<commit_before>package digitalocean

import (
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"time"

	"github.com/digitalocean/godo"
	"github.com/docker/machine/libmachine/drivers"
	"github.com/docker/machine/libmachine/log"
	"github.com/docker/machine/libmachine/mcnflag"
	"github.com/docker/machine/libmachine/ssh"
	"github.com/docker/machine/libmachine/state"
	"golang.org/x/oauth2"
)

type Driver struct {
	*drivers.BaseDriver
	AccessToken       string
	DropletID         int
	DropletName       string
	Image             string
	Region            string
	SSHKeyID          int
	Size              string
	IPv6              bool
	Backups           bool
	PrivateNetworking bool
	UserDataFile      string
}

const (
	defaultSSHPort = 22
	defaultSSHUSer = "root"
	defaultImage   = "ubuntu-15-10-x64"
	defaultRegion  = "nyc3"
	defaultSize    = "512mb"
)

// GetCreateFlags registers the flags this driver adds to
// "docker hosts create"
func (d *Driver) GetCreateFlags() []mcnflag.Flag {
	return []mcnflag.Flag{
		mcnflag.StringFlag{
			EnvVar: "DIGITALOCEAN_ACCESS_TOKEN",
			Name:   "digitalocean-access-token",
			Usage:  "Digital Ocean access token",
		},
		mcnflag.StringFlag{
			EnvVar: "DIGITALOCEAN_SSH_USER",
			Name:   "digitalocean-ssh-user",
			Usage:  "SSH username",
			Value:  defaultSSHUSer,
		},
		mcnflag.IntFlag{
			EnvVar: "DIGITALOCEAN_SSH_PORT",
			Name:   "digitalocean-ssh-port",
			Usage:  "SSH port",
			Value:  defaultSSHPort,
		},
		mcnflag.StringFlag{
			EnvVar: "DIGITALOCEAN_IMAGE",
			Name:   "digitalocean-image",
			Usage:  "Digital Ocean Image",
			Value:  defaultImage,
		},
		mcnflag.StringFlag{
			EnvVar: "DIGITALOCEAN_REGION",
			Name:   "digitalocean-region",
			Usage:  "Digital Ocean region",
			Value:  defaultRegion,
		},
		mcnflag.StringFlag{
			EnvVar: "DIGITALOCEAN_SIZE",
			Name:   "digitalocean-size",
			Usage:  "Digital Ocean size",
			Value:  defaultSize,
		},
		mcnflag.BoolFlag{
			EnvVar: "DIGITALOCEAN_IPV6",
			Name:   "digitalocean-ipv6",
			Usage:  "enable ipv6 for droplet",
		},
		mcnflag.BoolFlag{
			EnvVar: "DIGITALOCEAN_PRIVATE_NETWORKING",
			Name:   "digitalocean-private-networking",
			Usage:  "enable private networking for droplet",
		},
		mcnflag.BoolFlag{
			EnvVar: "DIGITALOCEAN_BACKUPS",
			Name:   "digitalocean-backups",
			Usage:  "enable backups for droplet",
		},
		mcnflag.StringFlag{
			EnvVar: "DIGITALOCEAN_USERDATA",
			Name:   "digitalocean-userdata",
			Usage:  "path to file with cloud-init user-data",
		},
	}
}

func NewDriver(hostName, storePath string) *Driver {
	return &Driver{
		Image:  defaultImage,
		Size:   defaultSize,
		Region: defaultRegion,
		BaseDriver: &drivers.BaseDriver{
			MachineName: hostName,
			StorePath:   storePath,
		},
	}
}

func (d *Driver) GetSSHHostname() (string, error) {
	return d.GetIP()
}

// DriverName returns the name of the driver
func (d *Driver) DriverName() string {
	return "digitalocean"
}

func (d *Driver) SetConfigFromFlags(flags drivers.DriverOptions) error {
	d.AccessToken = flags.String("digitalocean-access-token")
	d.Image = flags.String("digitalocean-image")
	d.Region = flags.String("digitalocean-region")
	d.Size = flags.String("digitalocean-size")
	d.IPv6 = flags.Bool("digitalocean-ipv6")
	d.PrivateNetworking = flags.Bool("digitalocean-private-networking")
	d.Backups = flags.Bool("digitalocean-backups")
	d.UserDataFile = flags.String("digitalocean-userdata")
	d.SSHUser = flags.String("digitalocean-ssh-user")
	d.SSHPort = flags.Int("digitalocean-ssh-port")
	d.SetSwarmConfigFromFlags(flags)

	if d.AccessToken == "" {
		return fmt.Errorf("digitalocean driver requires the --digitalocean-access-token option")
	}

	return nil
}

func (d *Driver) PreCreateCheck() error {
	if d.UserDataFile != "" {
		if _, err := os.Stat(d.UserDataFile); os.IsNotExist(err) {
			return fmt.Errorf("user-data file %s could not be found", d.UserDataFile)
		}
	}

	client := d.getClient()
	regions, _, err := client.Regions.List(nil)
	if err != nil {
		return err
	}
	for _, region := range regions {
		if region.Slug == d.Region {
			return nil
		}
	}

	return fmt.Errorf("digitalocean requires a valid region")
}

func (d *Driver) Create() error {
	var userdata string
	if d.UserDataFile != "" {
		buf, err := ioutil.ReadFile(d.UserDataFile)
		if err != nil {
			return err
		}
		userdata = string(buf)
	}

	log.Infof("Creating SSH key...")

	key, err := d.createSSHKey()
	if err != nil {
		return err
	}

	d.SSHKeyID = key.ID

	log.Infof("Creating Digital Ocean droplet...")

	client := d.getClient()

	createRequest := &godo.DropletCreateRequest{
		Image:             godo.DropletCreateImage{Slug: d.Image},
		Name:              d.MachineName,
		Region:            d.Region,
		Size:              d.Size,
		IPv6:              d.IPv6,
		PrivateNetworking: d.PrivateNetworking,
		Backups:           d.Backups,
		UserData:          userdata,
		SSHKeys:           []godo.DropletCreateSSHKey{{ID: d.SSHKeyID}},
	}

	newDroplet, _, err := client.Droplets.Create(createRequest)
	if err != nil {
		return err
	}

	d.DropletID = newDroplet.ID

	log.Info("Waiting for IP address to be assigned to the Droplet...")
	for {
		newDroplet, _, err = client.Droplets.Get(d.DropletID)
		if err != nil {
			return err
		}
		for _, network := range newDroplet.Networks.V4 {
			if network.Type == "public" {
				d.IPAddress = network.IPAddress
			}
		}

		if d.IPAddress != "" {
			break
		}

		time.Sleep(1 * time.Second)
	}

	log.Debugf("Created droplet ID %d, IP address %s",
		newDroplet.ID,
		d.IPAddress)

	return nil
}

func (d *Driver) createSSHKey() (*godo.Key, error) {
	if err := ssh.GenerateSSHKey(d.GetSSHKeyPath()); err != nil {
		return nil, err
	}

	publicKey, err := ioutil.ReadFile(d.publicSSHKeyPath())
	if err != nil {
		return nil, err
	}

	createRequest := &godo.KeyCreateRequest{
		Name:      d.MachineName,
		PublicKey: string(publicKey),
	}

	key, _, err := d.getClient().Keys.Create(createRequest)
	if err != nil {
		return key, err
	}

	return key, nil
}

func (d *Driver) GetURL() (string, error) {
	if err := drivers.MustBeRunning(d); err != nil {
		return "", err
	}

	ip, err := d.GetIP()
	if err != nil {
		return "", err
	}

	return fmt.Sprintf("tcp://%s", net.JoinHostPort(ip, "2376")), nil
}

func (d *Driver) GetState() (state.State, error) {
	droplet, _, err := d.getClient().Droplets.Get(d.DropletID)
	if err != nil {
		return state.Error, err
	}
	switch droplet.Status {
	case "new":
		return state.Starting, nil
	case "active":
		return state.Running, nil
	case "off":
		return state.Stopped, nil
	}
	return state.None, nil
}

func (d *Driver) Start() error {
	_, _, err := d.getClient().DropletActions.PowerOn(d.DropletID)
	return err
}

func (d *Driver) Stop() error {
	_, _, err := d.getClient().DropletActions.Shutdown(d.DropletID)
	return err
}

func (d *Driver) Restart() error {
	_, _, err := d.getClient().DropletActions.Reboot(d.DropletID)
	return err
}

func (d *Driver) Kill() error {
	_, _, err := d.getClient().DropletActions.PowerOff(d.DropletID)
	return err
}

func (d *Driver) Remove() error {
	client := d.getClient()
	if resp, err := client.Keys.DeleteByID(d.SSHKeyID); err != nil {
		if resp.StatusCode == 404 {
			log.Infof("Digital Ocean SSH key doesn't exist, assuming it is already deleted")
		} else {
			return err
		}
	}
	if resp, err := client.Droplets.Delete(d.DropletID); err != nil {
		if resp.StatusCode == 404 {
			log.Infof("Digital Ocean droplet doesn't exist, assuming it is already deleted")
		} else {
			return err
		}
	}
	return nil
}

func (d *Driver) getClient() *godo.Client {
	token := &oauth2.Token{AccessToken: d.AccessToken}
	tokenSource := oauth2.StaticTokenSource(token)
	client := oauth2.NewClient(oauth2.NoContext, tokenSource)

	return godo.NewClient(client)
}

func (d *Driver) publicSSHKeyPath() string {
	return d.GetSSHKeyPath() + ".pub"
}
<commit_msg>Fix typo in DigitalOcean driver<commit_after>package digitalocean

import (
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"time"

	"github.com/digitalocean/godo"
	"github.com/docker/machine/libmachine/drivers"
	"github.com/docker/machine/libmachine/log"
	"github.com/docker/machine/libmachine/mcnflag"
	"github.com/docker/machine/libmachine/ssh"
	"github.com/docker/machine/libmachine/state"
	"golang.org/x/oauth2"
)

type Driver struct {
	*drivers.BaseDriver
	AccessToken       string
	DropletID         int
	DropletName       string
	Image             string
	Region            string
	SSHKeyID          int
	Size              string
	IPv6              bool
	Backups           bool
	PrivateNetworking bool
	UserDataFile      string
}

const (
	defaultSSHPort = 22
	defaultSSHUser = "root"
	defaultImage   = "ubuntu-15-10-x64"
	defaultRegion  = "nyc3"
	defaultSize    = "512mb"
)

// GetCreateFlags registers the flags this driver adds to
// "docker hosts create"
func (d *Driver) GetCreateFlags() []mcnflag.Flag {
	return []mcnflag.Flag{
		mcnflag.StringFlag{
			EnvVar: "DIGITALOCEAN_ACCESS_TOKEN",
			Name:   "digitalocean-access-token",
			Usage:  "Digital Ocean access token",
		},
		mcnflag.StringFlag{
			EnvVar: "DIGITALOCEAN_SSH_USER",
			Name:   "digitalocean-ssh-user",
			Usage:  "SSH username",
			Value:  defaultSSHUser,
		},
		mcnflag.IntFlag{
			EnvVar: "DIGITALOCEAN_SSH_PORT",
			Name:   "digitalocean-ssh-port",
			Usage:  "SSH port",
			Value:  defaultSSHPort,
		},
		mcnflag.StringFlag{
			EnvVar: "DIGITALOCEAN_IMAGE",
			Name:   "digitalocean-image",
			Usage:  "Digital Ocean Image",
			Value:  defaultImage,
		},
		mcnflag.StringFlag{
			EnvVar: "DIGITALOCEAN_REGION",
			Name:   "digitalocean-region",
			Usage:  "Digital Ocean region",
			Value:  defaultRegion,
		},
		mcnflag.StringFlag{
			EnvVar: "DIGITALOCEAN_SIZE",
			Name:   "digitalocean-size",
			Usage:  "Digital Ocean size",
			Value:  defaultSize,
		},
		mcnflag.BoolFlag{
			EnvVar: "DIGITALOCEAN_IPV6",
			Name:   "digitalocean-ipv6",
			Usage:  "enable ipv6 for droplet",
		},
		mcnflag.BoolFlag{
			EnvVar: "DIGITALOCEAN_PRIVATE_NETWORKING",
			Name:   "digitalocean-private-networking",
			Usage:  "enable private networking for droplet",
		},
		mcnflag.BoolFlag{
			EnvVar: "DIGITALOCEAN_BACKUPS",
			Name:   "digitalocean-backups",
			Usage:  "enable backups for droplet",
		},
		mcnflag.StringFlag{
			EnvVar: "DIGITALOCEAN_USERDATA",
			Name:   "digitalocean-userdata",
			Usage:  "path to file with cloud-init user-data",
		},
	}
}

func NewDriver(hostName, storePath string) *Driver {
	return &Driver{
		Image:  defaultImage,
		Size:   defaultSize,
		Region: defaultRegion,
		BaseDriver: &drivers.BaseDriver{
			MachineName: hostName,
			StorePath:   storePath,
		},
	}
}

func (d *Driver) GetSSHHostname() (string, error) {
	return d.GetIP()
}

// DriverName returns the name of the driver
func (d *Driver) DriverName() string {
	return "digitalocean"
}

func (d *Driver) SetConfigFromFlags(flags drivers.DriverOptions) error {
	d.AccessToken = flags.String("digitalocean-access-token")
	d.Image = flags.String("digitalocean-image")
	d.Region = flags.String("digitalocean-region")
	d.Size = flags.String("digitalocean-size")
	d.IPv6 = flags.Bool("digitalocean-ipv6")
	d.PrivateNetworking = flags.Bool("digitalocean-private-networking")
	d.Backups = flags.Bool("digitalocean-backups")
	d.UserDataFile = flags.String("digitalocean-userdata")
	d.SSHUser = flags.String("digitalocean-ssh-user")
	d.SSHPort = flags.Int("digitalocean-ssh-port")
	d.SetSwarmConfigFromFlags(flags)

	if d.AccessToken == "" {
		return fmt.Errorf("digitalocean driver requires the --digitalocean-access-token option")
	}

	return nil
}

func (d *Driver) PreCreateCheck() error {
	if d.UserDataFile != "" {
		if _, err := os.Stat(d.UserDataFile); os.IsNotExist(err) {
			return fmt.Errorf("user-data file %s could not be found", d.UserDataFile)
		}
	}

	client := d.getClient()
	regions, _, err := client.Regions.List(nil)
	if err != nil {
		return err
	}
	for _, region := range regions {
		if region.Slug == d.Region {
			return nil
		}
	}

	return fmt.Errorf("digitalocean requires a valid region")
}

func (d *Driver) Create() error {
	var userdata string
	if d.UserDataFile != "" {
		buf, err := ioutil.ReadFile(d.UserDataFile)
		if err != nil {
			return err
		}
		userdata = string(buf)
	}

	log.Infof("Creating SSH key...")

	key, err := d.createSSHKey()
	if err != nil {
		return err
	}

	d.SSHKeyID = key.ID

	log.Infof("Creating Digital Ocean droplet...")

	client := d.getClient()

	createRequest := &godo.DropletCreateRequest{
		Image:             godo.DropletCreateImage{Slug: d.Image},
		Name:              d.MachineName,
		Region:            d.Region,
		Size:              d.Size,
		IPv6:              d.IPv6,
		PrivateNetworking: d.PrivateNetworking,
		Backups:           d.Backups,
		UserData:          userdata,
		SSHKeys:           []godo.DropletCreateSSHKey{{ID: d.SSHKeyID}},
	}

	newDroplet, _, err := client.Droplets.Create(createRequest)
	if err != nil {
		return err
	}

	d.DropletID = newDroplet.ID

	log.Info("Waiting for IP address to be assigned to the Droplet...")
	for {
		newDroplet, _, err = client.Droplets.Get(d.DropletID)
		if err != nil {
			return err
		}
		for _, network := range newDroplet.Networks.V4 {
			if network.Type == "public" {
				d.IPAddress = network.IPAddress
			}
		}

		if d.IPAddress != "" {
			break
		}

		time.Sleep(1 * time.Second)
	}

	log.Debugf("Created droplet ID %d, IP address %s",
		newDroplet.ID,
		d.IPAddress)

	return nil
}

func (d *Driver) createSSHKey() (*godo.Key, error) {
	if err := ssh.GenerateSSHKey(d.GetSSHKeyPath()); err != nil {
		return nil, err
	}

	publicKey, err := ioutil.ReadFile(d.publicSSHKeyPath())
	if err != nil {
		return nil, err
	}

	createRequest := &godo.KeyCreateRequest{
		Name:      d.MachineName,
		PublicKey: string(publicKey),
	}

	key, _, err := d.getClient().Keys.Create(createRequest)
	if err != nil {
		return key, err
	}

	return key, nil
}

func (d *Driver) GetURL() (string, error) {
	if err := drivers.MustBeRunning(d); err != nil {
		return "", err
	}

	ip, err := d.GetIP()
	if err != nil {
		return "", err
	}

	return fmt.Sprintf("tcp://%s", net.JoinHostPort(ip, "2376")), nil
}

func (d *Driver) GetState() (state.State, error) {
	droplet, _, err := d.getClient().Droplets.Get(d.DropletID)
	if err != nil {
		return state.Error, err
	}
	switch droplet.Status {
	case "new":
		return state.Starting, nil
	case "active":
		return state.Running, nil
	case "off":
		return state.Stopped, nil
	}
	return state.None, nil
}

func (d *Driver) Start() error {
	_, _, err := d.getClient().DropletActions.PowerOn(d.DropletID)
	return err
}

func (d *Driver) Stop() error {
	_, _, err := d.getClient().DropletActions.Shutdown(d.DropletID)
	return err
}

func (d *Driver) Restart() error {
	_, _, err := d.getClient().DropletActions.Reboot(d.DropletID)
	return err
}

func (d *Driver) Kill() error {
	_, _, err := d.getClient().DropletActions.PowerOff(d.DropletID)
	return err
}

func (d *Driver) Remove() error {
	client := d.getClient()
	if resp, err := client.Keys.DeleteByID(d.SSHKeyID); err != nil {
		if resp.StatusCode == 404 {
			log.Infof("Digital Ocean SSH key doesn't exist, assuming it is already deleted")
		} else {
			return err
		}
	}
	if resp, err := client.Droplets.Delete(d.DropletID); err != nil {
		if resp.StatusCode == 404 {
			log.Infof("Digital Ocean droplet doesn't exist, assuming it is already deleted")
		} else {
			return err
		}
	}
	return nil
}

func (d *Driver) getClient() *godo.Client {
	token := &oauth2.Token{AccessToken: d.AccessToken}
	tokenSource := oauth2.StaticTokenSource(token)
	client := oauth2.NewClient(oauth2.NoContext, tokenSource)

	return godo.NewClient(client)
}

func (d *Driver) publicSSHKeyPath() string {
	return d.GetSSHKeyPath() + ".pub"
}
<|endoftext|>


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

// Package agent implements a service that runs on a serviced node. It is
// responsible for ensuring that a particular node is running the correct services
// and reporting the state and health of those services back to the master
// serviced.

package utils

import (
	"fmt"
	"os"
	"os/user"
	"path"
	"path/filepath"
	"strings"
	"github.com/control-center/serviced/config"
	"github.com/zenoss/glog"
	"runtime"
)

//ServiceDHome gets the home location of serviced by looking at the environment
func ServiceDHome() string {
	homeDir := config.GetOptions().HomePath
	if len(homeDir) == 0 {
		// This fallback is used in unit-tests, but in actual practice,
		// we should not hit this case, because there is a default
		// value defined via code in cli/api/options.go.  But just in
		// case that somehow get's undone in the future, the log message
		// will let us know that we used a fallback.
		_, filename, _, _ := runtime.Caller(1)
		homeDir = strings.Replace(path.Dir(filename), "utils", "", 1)
		plog.Warnf("SERVICED_HOME not set; defaulting to %s", homeDir)
	}
	return homeDir
}

//LocalDir gets the absolute path to a particular directory under ServiceDHome
// if SERVICED_HOME is not defined then we use the location of the caller
func LocalDir(p string) string {
	homeDir := ServiceDHome()
	return path.Join(homeDir, p)
}

// ResourcesDir points to internal services resources directory
func ResourcesDir() string {
	return LocalDir("isvcs/resources")
}

// BackupDir gets the directory where backup files are stored
func BackupDir(basepath string) string {
	if backupDir := strings.TrimSpace(basepath); backupDir == "" {
		return TempDir("backups")
	} else {
		return filepath.Join(filepath.Clean(backupDir), "backups")
	}
}

// TempDir gets the temp serviced directory
func TempDir(p string) string {
	var tmp string

	if user, err := user.Current(); err == nil {
		tmp = path.Join(os.TempDir(), fmt.Sprintf("serviced-%s", user.Username), p)
	} else {
		tmp = path.Join(os.TempDir(), fmt.Sprintf("serviced"), p)
		glog.Warningf("Defaulting home to %s", tmp)
	}

	return tmp
}

// ServicedLogDir gets the serviced log directory
func ServicedLogDir() string {
        if config.GetOptions().LogPath != "" {
		return config.GetOptions().LogPath
	} else{
		return os.Getenv("SERVICED_LOG_PATH")
	}
}
<commit_msg>Remove unused method utils.BackupDir()<commit_after>// Copyright 2014 The Serviced Authors.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package agent implements a service that runs on a serviced node. It is
// responsible for ensuring that a particular node is running the correct services
// and reporting the state and health of those services back to the master
// serviced.

package utils

import (
	"fmt"
	"os"
	"os/user"
	"path"
	"runtime"
	"strings"

	"github.com/control-center/serviced/config"
	"github.com/zenoss/glog"
)

//ServiceDHome gets the home location of serviced by looking at the environment
func ServiceDHome() string {
	homeDir := config.GetOptions().HomePath
	if len(homeDir) == 0 {
		// This fallback is used in unit-tests, but in actual practice,
		// we should not hit this case, because there is a default
		// value defined via code in cli/api/options.go.  But just in
		// case that somehow get's undone in the future, the log message
		// will let us know that we used a fallback.
		_, filename, _, _ := runtime.Caller(1)
		homeDir = strings.Replace(path.Dir(filename), "utils", "", 1)
		plog.Warnf("SERVICED_HOME not set; defaulting to %s", homeDir)
	}
	return homeDir
}

//LocalDir gets the absolute path to a particular directory under ServiceDHome
// if SERVICED_HOME is not defined then we use the location of the caller
func LocalDir(p string) string {
	homeDir := ServiceDHome()
	return path.Join(homeDir, p)
}

// ResourcesDir points to internal services resources directory
func ResourcesDir() string {
	return LocalDir("isvcs/resources")
}

// TempDir gets the temp serviced directory
func TempDir(p string) string {
	var tmp string

	if user, err := user.Current(); err == nil {
		tmp = path.Join(os.TempDir(), fmt.Sprintf("serviced-%s", user.Username), p)
	} else {
		tmp = path.Join(os.TempDir(), fmt.Sprintf("serviced"), p)
		glog.Warningf("Defaulting home to %s", tmp)
	}

	return tmp
}

// ServicedLogDir gets the serviced log directory
func ServicedLogDir() string {
        if config.GetOptions().LogPath != "" {
		return config.GetOptions().LogPath
	} else{
		return os.Getenv("SERVICED_LOG_PATH")
	}
}
<|endoftext|>


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

package reverseproxy

import (
	"encoding/json"
	"flag"
	"fmt"
	"net/http"
	"strconv"

	"github.com/caddyserver/caddy/v2"
	"github.com/caddyserver/caddy/v2/caddyconfig"
	"github.com/caddyserver/caddy/v2/caddyconfig/httpcaddyfile"
	caddycmd "github.com/caddyserver/caddy/v2/cmd"
	"github.com/caddyserver/caddy/v2/modules/caddyhttp"
	"github.com/caddyserver/caddy/v2/modules/caddyhttp/headers"
)

func init() {
	caddycmd.RegisterCommand(caddycmd.Command{
		Name:  "reverse-proxy",
		Func:  cmdReverseProxy,
		Usage: "[--from <addr>] [--to <addr>] [--change-host-header]",
		Short: "A quick and production-ready reverse proxy",
		Long: `
A simple but production-ready reverse proxy. Useful for quick deployments,
demos, and development.

Simply shuttles HTTP(S) traffic from the --from address to the --to address.

Unless otherwise specified in the addresses, the --from address will be
assumed to be HTTPS if a hostname is given, and the --to address will be
assumed to be HTTP.

If the --from address has a host or IP, Caddy will attempt to serve the
proxy over HTTPS with a certificate (unless overridden by the HTTP scheme
or port).

If --change-host-header is set, the Host header on the request will be modified
from its original incoming value to the address of the upstream. (Otherwise, by
default, all incoming headers are passed through unmodified.)
`,
		Flags: func() *flag.FlagSet {
			fs := flag.NewFlagSet("reverse-proxy", flag.ExitOnError)
			fs.String("from", "localhost", "Address on which to receive traffic")
			fs.String("to", "", "Upstream address to which traffic should be sent")
			fs.Bool("change-host-header", false, "Set upstream Host header to address of upstream")
			fs.Bool("insecure", false, "Disable TLS verification (WARNING: DISABLES SECURITY BY NOT VERIFYING SSL CERTIFICATES!)")
			return fs
		}(),
	})
}

func cmdReverseProxy(fs caddycmd.Flags) (int, error) {
	caddy.TrapSignals()

	from := fs.String("from")
	to := fs.String("to")
	changeHost := fs.Bool("change-host-header")
	insecure := fs.Bool("insecure")

	httpPort := strconv.Itoa(caddyhttp.DefaultHTTPPort)
	httpsPort := strconv.Itoa(caddyhttp.DefaultHTTPSPort)

	if to == "" {
		return caddy.ExitCodeFailedStartup, fmt.Errorf("--to is required")
	}

	// set up the downstream address; assume missing information from given parts
	fromAddr, err := httpcaddyfile.ParseAddress(from)
	if err != nil {
		return caddy.ExitCodeFailedStartup, fmt.Errorf("invalid downstream address %s: %v", from, err)
	}
	if fromAddr.Path != "" {
		return caddy.ExitCodeFailedStartup, fmt.Errorf("paths are not allowed: %s", from)
	}
	if fromAddr.Scheme == "" {
		if fromAddr.Port == httpPort || fromAddr.Host == "" {
			fromAddr.Scheme = "http"
		} else {
			fromAddr.Scheme = "https"
		}
	}
	if fromAddr.Port == "" {
		if fromAddr.Scheme == "http" {
			fromAddr.Port = httpPort
		} else if fromAddr.Scheme == "https" {
			fromAddr.Port = httpsPort
		}
	}

	// set up the upstream address; assume missing information from given parts
	toAddr, toScheme, err := parseUpstreamDialAddress(to)
	if err != nil {
		return caddy.ExitCodeFailedStartup, fmt.Errorf("invalid upstream address %s: %v", to, err)
	}

	// proceed to build the handler and server
	ht := HTTPTransport{}
	if toScheme == "https" {
		ht.TLS = new(TLSConfig)
		if insecure {
			ht.TLS.InsecureSkipVerify = true
		}
	}

	handler := Handler{
		TransportRaw: caddyconfig.JSONModuleObject(ht, "protocol", "http", nil),
		Upstreams:    UpstreamPool{{Dial: toAddr}},
	}

	if changeHost {
		handler.Headers = &headers.Handler{
			Request: &headers.HeaderOps{
				Set: http.Header{
					"Host": []string{"{http.reverse_proxy.upstream.hostport}"},
				},
			},
		}
	}

	route := caddyhttp.Route{
		HandlersRaw: []json.RawMessage{
			caddyconfig.JSONModuleObject(handler, "handler", "reverse_proxy", nil),
		},
	}
	if fromAddr.Host != "" {
		route.MatcherSetsRaw = []caddy.ModuleMap{
			{
				"host": caddyconfig.JSON(caddyhttp.MatchHost{fromAddr.Host}, nil),
			},
		}
	}

	server := &caddyhttp.Server{
		Routes: caddyhttp.RouteList{route},
		Listen: []string{":" + fromAddr.Port},
	}

	httpApp := caddyhttp.App{
		Servers: map[string]*caddyhttp.Server{"proxy": server},
	}

	cfg := &caddy.Config{
		Admin: &caddy.AdminConfig{Disabled: true},
		AppsRaw: caddy.ModuleMap{
			"http": caddyconfig.JSON(httpApp, nil),
		},
	}

	err = caddy.Run(cfg)
	if err != nil {
		return caddy.ExitCodeFailedStartup, err
	}

	fmt.Printf("Caddy proxying %s -> %s\n", fromAddr.String(), toAddr)

	select {}
}
<commit_msg>reverseproxy: Add --internal-certs CLI flag #3589 (#4817)<commit_after>// Copyright 2015 Matthew Holt and The Caddy Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package reverseproxy

import (
	"encoding/json"
	"flag"
	"fmt"
	"net/http"
	"strconv"

	"github.com/caddyserver/caddy/v2"
	"github.com/caddyserver/caddy/v2/caddyconfig"
	"github.com/caddyserver/caddy/v2/caddyconfig/httpcaddyfile"
	caddycmd "github.com/caddyserver/caddy/v2/cmd"
	"github.com/caddyserver/caddy/v2/modules/caddyhttp"
	"github.com/caddyserver/caddy/v2/modules/caddyhttp/headers"
	"github.com/caddyserver/caddy/v2/modules/caddytls"
)

func init() {
	caddycmd.RegisterCommand(caddycmd.Command{
		Name:  "reverse-proxy",
		Func:  cmdReverseProxy,
		Usage: "[--from <addr>] [--to <addr>] [--change-host-header]",
		Short: "A quick and production-ready reverse proxy",
		Long: `
A simple but production-ready reverse proxy. Useful for quick deployments,
demos, and development.

Simply shuttles HTTP(S) traffic from the --from address to the --to address.

Unless otherwise specified in the addresses, the --from address will be
assumed to be HTTPS if a hostname is given, and the --to address will be
assumed to be HTTP.

If the --from address has a host or IP, Caddy will attempt to serve the
proxy over HTTPS with a certificate (unless overridden by the HTTP scheme
or port).

If --change-host-header is set, the Host header on the request will be modified
from its original incoming value to the address of the upstream. (Otherwise, by
default, all incoming headers are passed through unmodified.)
`,
		Flags: func() *flag.FlagSet {
			fs := flag.NewFlagSet("reverse-proxy", flag.ExitOnError)
			fs.String("from", "localhost", "Address on which to receive traffic")
			fs.String("to", "", "Upstream address to which traffic should be sent")
			fs.Bool("change-host-header", false, "Set upstream Host header to address of upstream")
			fs.Bool("insecure", false, "Disable TLS verification (WARNING: DISABLES SECURITY BY NOT VERIFYING SSL CERTIFICATES!)")
			fs.Bool("internal-certs", false, "Use internal CA for issuing certs")
			return fs
		}(),
	})
}

func cmdReverseProxy(fs caddycmd.Flags) (int, error) {
	caddy.TrapSignals()

	from := fs.String("from")
	to := fs.String("to")
	changeHost := fs.Bool("change-host-header")
	insecure := fs.Bool("insecure")
	internalCerts := fs.Bool("internal-certs")

	httpPort := strconv.Itoa(caddyhttp.DefaultHTTPPort)
	httpsPort := strconv.Itoa(caddyhttp.DefaultHTTPSPort)

	if to == "" {
		return caddy.ExitCodeFailedStartup, fmt.Errorf("--to is required")
	}

	// set up the downstream address; assume missing information from given parts
	fromAddr, err := httpcaddyfile.ParseAddress(from)
	if err != nil {
		return caddy.ExitCodeFailedStartup, fmt.Errorf("invalid downstream address %s: %v", from, err)
	}
	if fromAddr.Path != "" {
		return caddy.ExitCodeFailedStartup, fmt.Errorf("paths are not allowed: %s", from)
	}
	if fromAddr.Scheme == "" {
		if fromAddr.Port == httpPort || fromAddr.Host == "" {
			fromAddr.Scheme = "http"
		} else {
			fromAddr.Scheme = "https"
		}
	}
	if fromAddr.Port == "" {
		if fromAddr.Scheme == "http" {
			fromAddr.Port = httpPort
		} else if fromAddr.Scheme == "https" {
			fromAddr.Port = httpsPort
		}
	}

	// set up the upstream address; assume missing information from given parts
	toAddr, toScheme, err := parseUpstreamDialAddress(to)
	if err != nil {
		return caddy.ExitCodeFailedStartup, fmt.Errorf("invalid upstream address %s: %v", to, err)
	}

	// proceed to build the handler and server
	ht := HTTPTransport{}
	if toScheme == "https" {
		ht.TLS = new(TLSConfig)
		if insecure {
			ht.TLS.InsecureSkipVerify = true
		}
	}

	handler := Handler{
		TransportRaw: caddyconfig.JSONModuleObject(ht, "protocol", "http", nil),
		Upstreams:    UpstreamPool{{Dial: toAddr}},
	}

	if changeHost {
		handler.Headers = &headers.Handler{
			Request: &headers.HeaderOps{
				Set: http.Header{
					"Host": []string{"{http.reverse_proxy.upstream.hostport}"},
				},
			},
		}
	}

	route := caddyhttp.Route{
		HandlersRaw: []json.RawMessage{
			caddyconfig.JSONModuleObject(handler, "handler", "reverse_proxy", nil),
		},
	}
	if fromAddr.Host != "" {
		route.MatcherSetsRaw = []caddy.ModuleMap{
			{
				"host": caddyconfig.JSON(caddyhttp.MatchHost{fromAddr.Host}, nil),
			},
		}
	}

	server := &caddyhttp.Server{
		Routes: caddyhttp.RouteList{route},
		Listen: []string{":" + fromAddr.Port},
	}

	httpApp := caddyhttp.App{
		Servers: map[string]*caddyhttp.Server{"proxy": server},
	}

	appsRaw := caddy.ModuleMap{
		"http": caddyconfig.JSON(httpApp, nil),
	}
	if internalCerts && fromAddr.Host != "" {
		tlsApp := caddytls.TLS{
			Automation: &caddytls.AutomationConfig{
				Policies: []*caddytls.AutomationPolicy{{
					Subjects:   []string{fromAddr.Host},
					IssuersRaw: []json.RawMessage{json.RawMessage(`{"module":"internal"}`)},
				}},
			},
		}
		appsRaw["tls"] = caddyconfig.JSON(tlsApp, nil)
	}

	cfg := &caddy.Config{
		Admin:   &caddy.AdminConfig{Disabled: true},
		AppsRaw: appsRaw,
	}

	err = caddy.Run(cfg)
	if err != nil {
		return caddy.ExitCodeFailedStartup, err
	}

	fmt.Printf("Caddy proxying %s -> %s\n", fromAddr.String(), toAddr)

	select {}
}
<|endoftext|>


<commit_before>package git

import (
	"io"

	"github.com/github/git-lfs/tools"
)

type pktlineReader struct {
	pl *pktline

	buf []byte
}

var _ io.Reader = new(pktlineReader)

func (r *pktlineReader) Read(p []byte) (int, error) {
	var n int

	if len(r.buf) > 0 {
		// If there is data in the buffer, shift as much out of it and
		// into the given "p" as we can.
		n = tools.MinInt(len(p), len(r.buf))

		copy(p, r.buf[:n])
		r.buf = r.buf[n:]
	}

	// Loop and grab as many packets as we can in a given "run", until we
	// have either, a) overfilled the given buffer "p", or we have started
	// to internally buffer in "r.buf".
	for len(r.buf) == 0 {
		chunk, err := r.pl.readPacket()
		if err != nil {
			return n, err
		}

		if len(chunk) == 0 {
			// If we got an empty chunk, then we know that we have
			// reached the end of processing for this particular
			// packet, so let's terminate.

			return n, io.EOF
		}

		// Figure out how much of the packet we can read into "p".
		nn := tools.MinInt(len(chunk), len(p))

		// Move that amount into "p", from where we left off.
		copy(p[n:], chunk[:nn])
		// And move the rest into the buffer.
		r.buf = append(r.buf, chunk[nn:]...)

		// Mark that we have read "nn" bytes into "p"
		n += nn

		if n >= len(p) {
			break
		}
	}

	return n, nil
}
<commit_msg>git/pr: data is determined from how much remaining space "p" has, not total<commit_after>package git

import (
	"io"

	"github.com/github/git-lfs/tools"
)

type pktlineReader struct {
	pl *pktline

	buf []byte
}

var _ io.Reader = new(pktlineReader)

func (r *pktlineReader) Read(p []byte) (int, error) {
	var n int

	if len(r.buf) > 0 {
		// If there is data in the buffer, shift as much out of it and
		// into the given "p" as we can.
		n = tools.MinInt(len(p), len(r.buf))

		copy(p, r.buf[:n])
		r.buf = r.buf[n:]
	}

	// Loop and grab as many packets as we can in a given "run", until we
	// have either, a) overfilled the given buffer "p", or we have started
	// to internally buffer in "r.buf".
	for len(r.buf) == 0 {
		chunk, err := r.pl.readPacket()
		if err != nil {
			return n, err
		}

		if len(chunk) == 0 {
			// If we got an empty chunk, then we know that we have
			// reached the end of processing for this particular
			// packet, so let's terminate.

			return n, io.EOF
		}

		// Figure out how much of the packet we can read into "p".
		nn := tools.MinInt(len(chunk), len(p[n:]))

		// Move that amount into "p", from where we left off.
		copy(p[n:], chunk[:nn])
		// And move the rest into the buffer.
		r.buf = append(r.buf, chunk[nn:]...)

		// Mark that we have read "nn" bytes into "p"
		n += nn

		if n >= len(p) {
			break
		}
	}

	return n, nil
}
<|endoftext|>


<commit_before>package git

import (
	"io"

	"github.com/github/git-lfs/tools"
)

// PktlineWriter is an implementation of `io.Writer` which writes data buffers
// "p" to an underlying pkt-line stream for use with the Git pkt-line format.
type PktlineWriter struct {
	// buf is an internal buffer used to store data until enough has been
	// collected to write a full packet, or the buffer was instructed to
	// flush.
	buf []byte
	// pl is the place where packets get written.
	pl *pktline
}

var _ io.Writer = new(PktlineWriter)

// NewPktlineWriter returns a new *PktlineWriter, which will write to the
// underlying data stream "w". The internal buffer is initialized with the given
// capacity, "c".
//
// If "w" is already a `*PktlineWriter`, it will be returned as-is.
func NewPktlineWriter(w io.Writer, c int) *PktlineWriter {
	if pw, ok := w.(*PktlineWriter); ok {
		return pw
	}

	return &PktlineWriter{
		buf: make([]byte, 0, c),
		pl:  newPktline(nil, w),
	}
}

// Write implements the io.Writer interface's `Write` method by providing a
// packet-based backend to the given buffer "p".
//
// As many bytes are removed from "p" as possible and stored in an internal
// buffer until the amount of data in the internal buffer is enough to write a
// single packet. Once the internal buffer is full, a packet is written to the
// underlying stream of data, and the process repeats.
//
// When the caller has no more data to write in the given chunk of packets, a
// subsequent call to `Flush()` MUST be made in order to flush the remaining
// data in the buffer, and write the terminating bytes to the underlying packet
// stream. Alternatively, the caller can use the public `Flush() error` func to
// perform the same operation.
//
// Write returns the number of bytes in "p" accepted into the writer, which
// _MAY_ be written to the underlying protocol stream, or may be written into
// the internal buffer.
//
// If any error was encountered while either buffering or writing, that
// error is returned, along with the number of bytes written to the underlying
// protocol stream, as described above.
func (w *PktlineWriter) Write(p []byte) (int, error) {
	var n int

	for len(p[n:]) > 0 {
		// While there is still data left to process in "p", grab as
		// much of it as we can while not allowing the internal buffer
		// to exceed the MaxPacketLength const.
		m := tools.MinInt(len(p[n:]), MaxPacketLength-len(w.buf))

		// Append on all of the data that we could into the internal
		// buffer.
		w.buf = append(w.buf, p[n:n+m]...)

		n += m

		if len(w.buf) == MaxPacketLength {
			// If we were able to grab an entire packet's worth of
			// data, flush the buffer.

			if _, err := w.flush(); err != nil {
				return n, err
			}

		}
	}

	return n, nil
}

// Flush empties the internal buffer used to store data temporarily and then
// writes the pkt-line's FLUSH packet, to signal that it is done writing this
// chunk of data.
func (w *PktlineWriter) Flush() error {
	if _, err := w.flush(); err != nil {
		return err
	}

	if err := w.pl.writeFlush(); err != nil {
		return err
	}

	return nil
}

// flush writes any data in the internal buffer out to the underlying protocol
// stream. If the amount of data in the internal buffer exceeds the
// MaxPacketLength, the data will be written in multiple packets to accommodate.
//
// flush returns the number of bytes written to the underlying packet stream,
// and any error that it encountered along the way.
func (w *PktlineWriter) flush() (int, error) {
	var n int

	for len(w.buf) > 0 {
		if err := w.pl.writePacket(w.buf); err != nil {
			return 0, err
		}

		m := tools.MinInt(len(w.buf), MaxPacketLength)

		w.buf = w.buf[m:]

		n = n + m
	}

	return n, nil
}
<commit_msg>git/pw: :nail_care: documentation<commit_after>package git

import (
	"io"

	"github.com/github/git-lfs/tools"
)

// PktlineWriter is an implementation of `io.Writer` which writes data buffers
// "p" to an underlying pkt-line stream for use with the Git pkt-line format.
type PktlineWriter struct {
	// buf is an internal buffer used to store data until enough has been
	// collected to write a full packet, or the buffer was instructed to
	// flush.
	buf []byte
	// pl is the place where packets get written.
	pl *pktline
}

var _ io.Writer = new(PktlineWriter)

// NewPktlineWriter returns a new *PktlineWriter, which will write to the
// underlying data stream "w". The internal buffer is initialized with the given
// capacity, "c".
//
// If "w" is already a `*PktlineWriter`, it will be returned as-is.
func NewPktlineWriter(w io.Writer, c int) *PktlineWriter {
	if pw, ok := w.(*PktlineWriter); ok {
		return pw
	}

	return &PktlineWriter{
		buf: make([]byte, 0, c),
		pl:  newPktline(nil, w),
	}
}

// Write implements the io.Writer interface's `Write` method by providing a
// packet-based backend to the given buffer "p".
//
// As many bytes are removed from "p" as possible and stored in an internal
// buffer until the amount of data in the internal buffer is enough to write a
// single packet. Once the internal buffer is full, a packet is written to the
// underlying stream of data, and the process repeats.
//
// When the caller has no more data to write in the given chunk of packets, a
// subsequent call to `Flush()` SHOULD be made in order to signify that the
// current pkt sequence has terminated, and a new one can begin.
//
// Write returns the number of bytes in "p" accepted into the writer, which
// _MAY_ be written to the underlying protocol stream, or may be written into
// the internal buffer.
//
// If any error was encountered while either buffering or writing, that
// error is returned, along with the number of bytes written to the underlying
// protocol stream, as described above.
func (w *PktlineWriter) Write(p []byte) (int, error) {
	var n int

	for len(p[n:]) > 0 {
		// While there is still data left to process in "p", grab as
		// much of it as we can while not allowing the internal buffer
		// to exceed the MaxPacketLength const.
		m := tools.MinInt(len(p[n:]), MaxPacketLength-len(w.buf))

		// Append on all of the data that we could into the internal
		// buffer.
		w.buf = append(w.buf, p[n:n+m]...)

		n += m

		if len(w.buf) == MaxPacketLength {
			// If we were able to grab an entire packet's worth of
			// data, flush the buffer.

			if _, err := w.flush(); err != nil {
				return n, err
			}

		}
	}

	return n, nil
}

// Flush empties the internal buffer used to store data temporarily and then
// writes the pkt-line's FLUSH packet, to signal that it is done writing this
// chunk of data.
func (w *PktlineWriter) Flush() error {
	if _, err := w.flush(); err != nil {
		return err
	}

	if err := w.pl.writeFlush(); err != nil {
		return err
	}

	return nil
}

// flush writes any data in the internal buffer out to the underlying protocol
// stream. If the amount of data in the internal buffer exceeds the
// MaxPacketLength, the data will be written in multiple packets to accommodate.
//
// flush returns the number of bytes written to the underlying packet stream,
// and any error that it encountered along the way.
func (w *PktlineWriter) flush() (int, error) {
	var n int

	for len(w.buf) > 0 {
		if err := w.pl.writePacket(w.buf); err != nil {
			return 0, err
		}

		m := tools.MinInt(len(w.buf), MaxPacketLength)

		w.buf = w.buf[m:]

		n = n + m
	}

	return n, nil
}
<|endoftext|>


<commit_before>package flip

import (
	"context"
	"io"
	"time"
)

type GameMessageReplayed struct {
	GameMessageWrappedEncoded
	Time time.Time
}

type GameHistory []GameMessageReplayed

type GameSummary struct {
	Err     error
	Players []UserDevice
	Result  Result
}

func (g GameHistory) start(rh ReplayHelper) (game *Game, rest GameHistory, err error) {
	if len(g) == 0 {
		return nil, nil, NewReplayError("cannot reply 0-length game")
	}
	first := &g[0]
	rest = g[1:]
	gmw, err := first.Decode()
	if err != nil {
		return nil, nil, err
	}
	t, err := gmw.Msg.Body.T()
	if err != nil {
		return nil, nil, err
	}
	if t != MessageType_START {
		return nil, nil, NewReplayError("expected first message to be of type START")
	}
	start := gmw.Msg.Body.Start()

	md := gmw.Msg.Md
	if !md.Initiator.Eq(gmw.Sender) {
		return nil, nil, NewReplayError("bad initiator; didn't match sender")
	}

	_, err = computeClockSkew(md, first.Time, start.StartTime.Time(), first.Time)
	if err != nil {
		return nil, nil, err
	}

	game = &Game{
		md:           md,
		isLeader:     false,
		start:        first.Time,
		key:          md.ToKey(),
		params:       start,
		gameUpdateCh: make(chan GameStateUpdateMessage),
		players:      make(map[UserDeviceKey]*GamePlayerState),
		stage:        Stage_ROUND1,
		clogf:        rh.CLogf,
	}

	return game, rest, nil
}

func runReplayLoop(ctx context.Context, game *Game, gh GameHistory) (err error) {
	for _, m := range gh {
		gmw, err := m.GameMessageWrappedEncoded.Decode()
		if err != nil {
			return err
		}
		err = game.handleMessage(ctx, gmw, m.Time)
		if err == io.EOF {
			return nil
		}
		if err != nil {
			return err
		}
	}
	return nil
}

func extractUserDevices(v []UserDeviceCommitment) []UserDevice {
	ret := make([]UserDevice, len(v))
	for i, e := range v {
		ret[i] = e.Ud
	}
	return ret
}

func Replay(ctx context.Context, rh ReplayHelper, gh GameHistory) (*GameSummary, error) {

	var game *Game
	var err error
	game, gh, err = gh.start(rh)
	if err != nil {
		return nil, err
	}

	errCh := make(chan error)
	go func() {
		err := runReplayLoop(ctx, game, gh)
		close(game.gameUpdateCh)
		errCh <- err
	}()

	summaryCh := make(chan GameSummary)
	go func() {
		var ret GameSummary
		found := false
		players := make(map[UserDeviceKey]UserDevice)
		for msg := range game.gameUpdateCh {
			switch {
			case msg.Err != nil:
				ret.Err = msg.Err
			case msg.CommitmentComplete != nil:
				ret.Players = extractUserDevices(msg.CommitmentComplete.Players)
				for _, p := range ret.Players {
					players[p.ToKey()] = p
				}
			case msg.Reveal != nil:
				delete(players, msg.Reveal.User.ToKey())
			case msg.Result != nil:
				ret.Result = *msg.Result
				found = true
			}
		}
		if !found && ret.Err == nil {
			var ea AbsenteesError
			for _, v := range players {
				ea.Absentees = append(ea.Absentees, v)
			}
			ret.Err = ea
		}
		summaryCh <- ret
	}()

	err = <-errCh
	if err != nil {
		return nil, err
	}
	ret := <-summaryCh
	if ret.Err != nil {
		return nil, ret.Err
	}

	return &ret, nil
}
<commit_msg>debug the replay (#16174)<commit_after>package flip

import (
	"context"
	"io"
	"time"
)

type GameMessageReplayed struct {
	GameMessageWrappedEncoded
	Time time.Time
}

type GameHistory []GameMessageReplayed

type GameSummary struct {
	Err     error
	Players []UserDevice
	Result  Result
}

func (g GameHistory) start(rh ReplayHelper) (game *Game, rest GameHistory, err error) {
	if len(g) == 0 {
		return nil, nil, NewReplayError("cannot reply 0-length game")
	}
	first := &g[0]
	rest = g[1:]
	gmw, err := first.Decode()
	if err != nil {
		return nil, nil, err
	}
	t, err := gmw.Msg.Body.T()
	if err != nil {
		return nil, nil, err
	}
	if t != MessageType_START {
		return nil, nil, NewReplayError("expected first message to be of type START")
	}
	start := gmw.Msg.Body.Start()

	md := gmw.Msg.Md
	if !md.Initiator.Eq(gmw.Sender) {
		return nil, nil, NewReplayError("bad initiator; didn't match sender")
	}

	_, err = computeClockSkew(md, first.Time, start.StartTime.Time(), first.Time)
	if err != nil {
		return nil, nil, err
	}

	game = &Game{
		md:           md,
		isLeader:     false,
		start:        first.Time,
		key:          md.ToKey(),
		params:       start,
		gameUpdateCh: make(chan GameStateUpdateMessage),
		players:      make(map[UserDeviceKey]*GamePlayerState),
		stage:        Stage_ROUND1,
		clogf:        rh.CLogf,
	}

	return game, rest, nil
}

func runReplayLoop(ctx context.Context, game *Game, gh GameHistory) (err error) {
	for _, m := range gh {
		gmw, err := m.GameMessageWrappedEncoded.Decode()
		if err != nil {
			return err
		}
		err = game.handleMessage(ctx, gmw, m.Time)
		if err == io.EOF {
			return nil
		}
		if err != nil {
			return err
		}
	}
	return nil
}

func extractUserDevices(v []UserDeviceCommitment) []UserDevice {
	ret := make([]UserDevice, len(v))
	for i, e := range v {
		ret[i] = e.Ud
	}
	return ret
}

func Replay(ctx context.Context, rh ReplayHelper, gh GameHistory) (*GameSummary, error) {
	ret, err := replay(ctx, rh, gh)
	if err != nil {
		rh.CLogf(ctx, "Replay failure (%s); game dump: %+v", err, gh)
	}
	return ret, err
}

func replay(ctx context.Context, rh ReplayHelper, gh GameHistory) (*GameSummary, error) {

	var game *Game
	var err error
	game, gh, err = gh.start(rh)
	if err != nil {
		return nil, err
	}

	errCh := make(chan error)
	go func() {
		err := runReplayLoop(ctx, game, gh)
		close(game.gameUpdateCh)
		errCh <- err
	}()

	summaryCh := make(chan GameSummary)
	go func() {
		var ret GameSummary
		found := false
		players := make(map[UserDeviceKey]UserDevice)
		for msg := range game.gameUpdateCh {
			switch {
			case msg.Err != nil:
				ret.Err = msg.Err
			case msg.CommitmentComplete != nil:
				ret.Players = extractUserDevices(msg.CommitmentComplete.Players)
				for _, p := range ret.Players {
					players[p.ToKey()] = p
				}
			case msg.Reveal != nil:
				delete(players, msg.Reveal.User.ToKey())
			case msg.Result != nil:
				ret.Result = *msg.Result
				found = true
			}
		}
		if !found && ret.Err == nil {
			var ea AbsenteesError
			for _, v := range players {
				ea.Absentees = append(ea.Absentees, v)
			}
			ret.Err = ea
		}
		summaryCh <- ret
	}()

	err = <-errCh
	if err != nil {
		return nil, err
	}
	ret := <-summaryCh
	if ret.Err != nil {
		return nil, ret.Err
	}

	return &ret, nil
}
<|endoftext|>