<commit_before>package diamond

import (
	"fmt"
	"math/rand"
	"reflect"
	"strings"
	"testing"
	"testing/quick"
	"time"
)

var config = &quick.Config{Rand: rand.New(rand.NewSource(time.Now().UnixNano()))}

type correctChar byte

func (c correctChar) Generate(r *rand.Rand, _ int) reflect.Value {
	return reflect.ValueOf(correctChar('A' + r.Intn('Z'-'A'+1)))
}

func checkCorrect(requirement func(byte, []string) bool, keepSeparator bool, t *testing.T) {
	assertion := func(char correctChar) bool {
		d, err := Gen(byte(char))
		if err != nil {
			return false
		}
		separator := strings.Split
		if keepSeparator {
			separator = strings.SplitAfter
		}
		rows := separator(d, "\n")
		if len(rows) < 2 {
			return false
		}
		return requirement(byte(char), rows[:len(rows)-1])
	}
	if err := quick.Check(assertion, config); err != nil {
		t.Error(err)
	}
}

func TestFirstRowContainsOneA(t *testing.T) {
	requirement := func(char byte, rows []string) bool {
		return len(rows) > 0 && strings.Count(rows[0], "A") == 1
	}
	checkCorrect(requirement, false, t)
}

func TestLastRowContainsOneA(t *testing.T) {
	requirement := func(char byte, rows []string) bool {
		fmt.Printf("row %v count %v\n", rows[len(rows)-1], strings.Count(rows[len(rows)-1], "A"))
		return len(rows) > 0 && strings.Count(rows[len(rows)-1], "A") == 1
	}
	checkCorrect(requirement, false, t)
}

func TestAllRowsIdenticalLettersExceptFirstAndLast(t *testing.T) {
	requirement := func(char byte, rows []string) bool {
		for i, row := range rows {
			r := strings.TrimSpace(row)
			if r[0] != r[len(r)-1] {
				return false
			}
			if len(r) < 2 && i != 0 && i != len(rows)-1 {
				return false
			}
		}
		return true
	}
	checkCorrect(requirement, false, t)
}

func TestAllRowsHaveSameTrailingSpaces(t *testing.T) {
	requirement := func(char byte, rows []string) bool {
		for _, row := range rows {
			if row == "" {
				return false
			}
			for i, j := 0, len(row)-1; i < j && row[i] == ' '; i, j = i+1, j-1 {
				if row[j] != ' ' {
					return false
				}
			}
		}
		return true
	}
	checkCorrect(requirement, false, t)
}

func TestDiamondIsHorizontallySymmetric(t *testing.T) {
	requirement := func(char byte, rows []string) bool {
		for _, row := range rows {
			l := len(row)
			for i := l/2 - 1; i >= 0; i-- {
				if row[i] != row[l-1-i] {
					return false
				}
			}
		}
		return true
	}
	checkCorrect(requirement, false, t)
}

func TestDiamondIsVerticallySymmetric(t *testing.T) {
	requirement := func(char byte, rows []string) bool {
		for i, j := 0, len(rows)-1; i < j; i, j = i+1, j-1 {
			if rows[i] != rows[j] {
				return false
			}
		}
		return true
	}
	checkCorrect(requirement, true, t)
}

func TestDiamondIsSquare(t *testing.T) {
	requirement := func(char byte, rows []string) bool {
		if int(char-'A')*2+1 != len(rows) {
			return false
		}
		for _, row := range rows {
			if len(row) != len(rows) {
				return false
			}
		}
		return true
	}
	checkCorrect(requirement, false, t)
}

func TestDiamondHasItsShape(t *testing.T) {
	requirement := func(char byte, rows []string) bool {
		var n int
		for i, row := range rows {
			s := len(strings.TrimSpace(row))
			if i > len(rows)/2 && n <= s {
				return false
			} else if i <= len(rows)/2 && n >= s {
				return false
			}
			n = s
		}
		return true
	}
	checkCorrect(requirement, false, t)
}

func TestTopHalfHasAscendingLetters(t *testing.T) {
	requirement := func(char byte, rows []string) bool {
		var start byte = 'A' - 1
		for i := 0; i <= len(rows)/2; i++ {
			s := strings.TrimLeft(rows[i], " ")
			if s == "" || s[0] <= start {
				return false
			}
			start = s[0]
		}
		return true
	}
	checkCorrect(requirement, false, t)
}

func TestBottomHalfHasDescendingLetters(t *testing.T) {
	requirement := func(char byte, rows []string) bool {
		var start byte = 'A' - 1
		for i := len(rows) - 1; i > len(rows)/2; i-- {
			s := strings.TrimLeft(rows[i], " ")
			if s == "" || s[0] <= start {
				return false
			}
			start = s[0]
		}
		return true
	}
	checkCorrect(requirement, false, t)
}

func TestDiamondFourCornersAreTriangle(t *testing.T) {
	requirement := func(char byte, rows []string) bool {
		notSpace := func(r rune) bool { return r <= 'Z' && r >= 'A' }
		var n int
		for i, row := range rows {
			s := strings.IndexFunc(row, notSpace)
			e := len(row) - strings.LastIndexFunc(row, notSpace) - 1
			switch {
			case s != e:
				return false
			case i == 0:
				n = s
			default:
				if i > len(rows)/2 && n >= s {
					return false
				} else if i <= len(rows)/2 && n <= s {
					return false
				}
				n = s
			}
		}
		return true
	}
	checkCorrect(requirement, false, t)
}

type wrongChar byte

func (c wrongChar) Generate(r *rand.Rand, _ int) reflect.Value {
	b := rand.Intn(256)
	for ; b >= 'A' && b <= 'Z'; b = r.Intn(256) {
	}
	return reflect.ValueOf(wrongChar(b))
}

func TestCharOutOfRangeShouldGiveError(t *testing.T) {
	assertion := func(char wrongChar) bool {
		_, err := Gen(byte(char))
		return err != nil
	}
	if err := quick.Check(assertion, config); err != nil {
		t.Error(err)
	}
}
<commit_msg>Improvement idea for Go tests<commit_after>package diamond

import (
	"fmt"
	"math/rand"
	"reflect"
	"strings"
	"testing"
	"testing/quick"
	"time"
)

type test struct {
	input rune
	want  string
}

func TestDiamond(t *testing.T) {
	tests := []test{
		{'A', "A\n"},
		{'B', " A \nB B\n A \n"},
		{'C', "  A  \n B B \nC   C\n B B \n  A  \n"},
		{'D', "   A   \n  B B  \n C   C \nD     D\n C   C \n  B B  \n   A   \n"},
	}
	for _, test := range tests {
		got, err := Gen(byte(test.input))

		if err != nil {
			t.Errorf("Gen(byte(%v)) err %v", test.input, err)
		}
		if got != test.want {
			t.Errorf("Gen(byte(%v) got\n%v want\n%v", test.input, got, test.want)
		}
	}
}

var config = &quick.Config{Rand: rand.New(rand.NewSource(time.Now().UnixNano()))}

type correctChar byte

func (c correctChar) Generate(r *rand.Rand, _ int) reflect.Value {
	return reflect.ValueOf(correctChar('A' + r.Intn('Z'-'A'+1)))
}

func checkCorrect(requirement func(byte, []string) bool, keepSeparator bool, t *testing.T) {
	assertion := func(char correctChar) bool {
		d, err := Gen(byte(char))
		if err != nil {
			return false
		}
		separator := strings.Split
		if keepSeparator {
			separator = strings.SplitAfter
		}
		rows := separator(d, "\n")
		if len(rows) < 2 {
			return false
		}
		return requirement(byte(char), rows[:len(rows)-1])
	}
	if err := quick.Check(assertion, config); err != nil {
		t.Error(err)
	}
}

func TestFirstRowContainsOneA(t *testing.T) {
	requirement := func(char byte, rows []string) bool {
		return len(rows) > 0 && strings.Count(rows[0], "A") == 1
	}
	checkCorrect(requirement, false, t)
}

func TestLastRowContainsOneA(t *testing.T) {
	requirement := func(char byte, rows []string) bool {
		fmt.Printf("row %v count %v\n", rows[len(rows)-1], strings.Count(rows[len(rows)-1], "A"))
		return len(rows) > 0 && strings.Count(rows[len(rows)-1], "A") == 1
	}
	checkCorrect(requirement, false, t)
}

func TestAllRowsIdenticalLettersExceptFirstAndLast(t *testing.T) {
	requirement := func(char byte, rows []string) bool {
		for i, row := range rows {
			r := strings.TrimSpace(row)
			if r[0] != r[len(r)-1] {
				return false
			}
			if len(r) < 2 && i != 0 && i != len(rows)-1 {
				return false
			}
		}
		return true
	}
	checkCorrect(requirement, false, t)
}

func TestAllRowsHaveSameTrailingSpaces(t *testing.T) {
	requirement := func(char byte, rows []string) bool {
		for _, row := range rows {
			if row == "" {
				return false
			}
			for i, j := 0, len(row)-1; i < j && row[i] == ' '; i, j = i+1, j-1 {
				if row[j] != ' ' {
					return false
				}
			}
		}
		return true
	}
	checkCorrect(requirement, false, t)
}

func TestDiamondIsHorizontallySymmetric(t *testing.T) {
	requirement := func(char byte, rows []string) bool {
		for _, row := range rows {
			l := len(row)
			for i := l/2 - 1; i >= 0; i-- {
				if row[i] != row[l-1-i] {
					return false
				}
			}
		}
		return true
	}
	checkCorrect(requirement, false, t)
}

func TestDiamondIsVerticallySymmetric(t *testing.T) {
	requirement := func(char byte, rows []string) bool {
		for i, j := 0, len(rows)-1; i < j; i, j = i+1, j-1 {
			if rows[i] != rows[j] {
				return false
			}
		}
		return true
	}
	checkCorrect(requirement, true, t)
}

func TestDiamondIsSquare(t *testing.T) {
	requirement := func(char byte, rows []string) bool {
		if int(char-'A')*2+1 != len(rows) {
			return false
		}
		for _, row := range rows {
			if len(row) != len(rows) {
				return false
			}
		}
		return true
	}
	checkCorrect(requirement, false, t)
}

func TestDiamondHasItsShape(t *testing.T) {
	requirement := func(char byte, rows []string) bool {
		var n int
		for i, row := range rows {
			s := len(strings.TrimSpace(row))
			if i > len(rows)/2 && n <= s {
				return false
			} else if i <= len(rows)/2 && n >= s {
				return false
			}
			n = s
		}
		return true
	}
	checkCorrect(requirement, false, t)
}

func TestTopHalfHasAscendingLetters(t *testing.T) {
	requirement := func(char byte, rows []string) bool {
		var start byte = 'A' - 1
		for i := 0; i <= len(rows)/2; i++ {
			s := strings.TrimLeft(rows[i], " ")
			if s == "" || s[0] <= start {
				return false
			}
			start = s[0]
		}
		return true
	}
	checkCorrect(requirement, false, t)
}

func TestBottomHalfHasDescendingLetters(t *testing.T) {
	requirement := func(char byte, rows []string) bool {
		var start byte = 'A' - 1
		for i := len(rows) - 1; i > len(rows)/2; i-- {
			s := strings.TrimLeft(rows[i], " ")
			if s == "" || s[0] <= start {
				return false
			}
			start = s[0]
		}
		return true
	}
	checkCorrect(requirement, false, t)
}

func TestDiamondFourCornersAreTriangle(t *testing.T) {
	requirement := func(char byte, rows []string) bool {
		notSpace := func(r rune) bool { return r <= 'Z' && r >= 'A' }
		var n int
		for i, row := range rows {
			s := strings.IndexFunc(row, notSpace)
			e := len(row) - strings.LastIndexFunc(row, notSpace) - 1
			switch {
			case s != e:
				return false
			case i == 0:
				n = s
			default:
				if i > len(rows)/2 && n >= s {
					return false
				} else if i <= len(rows)/2 && n <= s {
					return false
				}
				n = s
			}
		}
		return true
	}
	checkCorrect(requirement, false, t)
}

type wrongChar byte

func (c wrongChar) Generate(r *rand.Rand, _ int) reflect.Value {
	b := rand.Intn(256)
	for ; b >= 'A' && b <= 'Z'; b = r.Intn(256) {
	}
	return reflect.ValueOf(wrongChar(b))
}

func TestCharOutOfRangeShouldGiveError(t *testing.T) {
	assertion := func(char wrongChar) bool {
		_, err := Gen(byte(char))
		return err != nil
	}
	if err := quick.Check(assertion, config); err != nil {
		t.Error(err)
	}
}
<|endoftext|>


<commit_before>package storage

import (
	"bytes"
	"errors"
	"time"

	"github.com/keybase/client/go/gregor"
	"github.com/keybase/client/go/protocol/gregor1"
	"golang.org/x/net/context"

	rpc "github.com/keybase/go-framed-msgpack-rpc"
)

type LocalStorageEngine interface {
	Store(gregor.UID, []byte) error
	Load(gregor.UID) ([]byte, error)
}

type Client struct {
	user    gregor.UID
	device  gregor.DeviceID
	sm      gregor.StateMachine
	storage LocalStorageEngine
	log     rpc.LogOutput

	saveTimer <-chan time.Time
}

func NewClient(user gregor.UID, device gregor.DeviceID, sm gregor.StateMachine,
	storage LocalStorageEngine, saveInterval time.Duration, log rpc.LogOutput) *Client {
	c := &Client{
		user:      user,
		device:    device,
		sm:        sm,
		storage:   storage,
		log:       log,
		saveTimer: time.Tick(saveInterval), // How often we save to local storage
	}
	return c
}

func (c *Client) Save() error {
	if !c.sm.IsEphemeral() {
		return errors.New("state machine is non-ephemeral")
	}

	state, err := c.sm.State(c.user, c.device, nil)
	if err != nil {
		return err
	}

	b, err := state.Marshal()
	if err != nil {
		return err
	}

	return c.storage.Store(c.user, b)
}

func (c *Client) Restore() error {
	if !c.sm.IsEphemeral() {
		return errors.New("state machine is non-ephemeral")
	}

	value, err := c.storage.Load(c.user)
	if err != nil {
		return err
	}

	state, err := c.sm.ObjFactory().UnmarshalState(value)
	if err != nil {
		return err
	}

	if err := c.sm.InitState(state); err != nil {
		return err
	}

	return nil
}

type errHashMismatch struct{}

func (e errHashMismatch) Error() string {
	return "local state hash != server state hash"
}

func (c *Client) syncFromTime(cli gregor1.IncomingInterface, t *time.Time) (msgs []gregor.InBandMessage, err error) {

	ctx, _ := context.WithTimeout(context.Background(), time.Second)
	arg := gregor1.SyncArg{
		Uid:      gregor1.UID(c.user.Bytes()),
		Deviceid: gregor1.DeviceID(c.device.Bytes()),
	}
	if t != nil {
		arg.Ctime = gregor1.ToTime(*t)
	}

	// Grab the events from gregord
	c.log.Debug("syncFromTime from: %s", gregor1.FromTime(arg.Ctime))
	res, err := cli.Sync(ctx, arg)
	if err != nil {
		return nil, err
	}

	c.log.Debug("syncFromTime consuming %d messages", len(res.Msgs))
	for _, ibm := range res.Msgs {
		m := gregor1.Message{Ibm_: &ibm}
		msgs = append(msgs, ibm)
		c.sm.ConsumeMessage(m)
	}

	// Check to make sure the server state is legit
	state, err := c.sm.State(c.user, c.device, nil)
	if err != nil {
		return nil, err
	}
	hash, err := state.Hash()
	if err != nil {
		return nil, err
	}
	if !bytes.Equal(res.Hash, hash) {
		return nil, errHashMismatch{}
	}

	return msgs, nil
}

func (c *Client) freshSync(cli gregor1.IncomingInterface) ([]gregor.InBandMessage, error) {

	var msgs []gregor.InBandMessage
	var err error

	c.sm.Clear()
	state, err := c.State(cli)
	if err != nil {
		return msgs, err
	}
	if msgs, err = c.InBandMessagesFromState(state); err != nil {
		return msgs, err
	}
	if err = c.sm.InitState(state); err != nil {
		return msgs, err
	}

	return msgs, nil
}

func (c *Client) Sync(cli gregor1.IncomingInterface) ([]gregor.InBandMessage, error) {
	latestCtime := c.sm.LatestCTime(c.user, c.device)
	if latestCtime == nil || latestCtime.IsZero() {
		c.log.Debug("Sync(): fresh server sync: using State()")
		return c.freshSync(cli)
	} else {
		c.log.Debug("Sync(): incremental server sync: using Sync()")
		if msgs, err := c.syncFromTime(cli, latestCtime); err != nil {
			if _, ok := err.(errHashMismatch); ok {
				c.log.Info("Sync failure: %v\nResetting StateMachine and retrying", err)
				return c.freshSync(cli)
			}
			return msgs, err
		} else {
			return msgs, nil
		}
	}
}

func (c *Client) InBandMessagesFromState(s gregor.State) ([]gregor.InBandMessage, error) {
	items, err := s.Items()
	if err != nil {
		return nil, err
	}

	var res []gregor.InBandMessage
	for _, i := range items {
		if ibm, err := c.sm.ObjFactory().MakeInBandMessageFromItem(i); err == nil {
			res = append(res, ibm)
		}
	}
	return res, nil
}

func (c *Client) State(cli gregor1.IncomingInterface) (gregor.State, error) {
	ctx, _ := context.WithTimeout(context.Background(), 5*time.Second)
	arg := gregor1.StateArg{
		Uid:          gregor1.UID(c.user.Bytes()),
		Deviceid:     gregor1.DeviceID(c.device.Bytes()),
		TimeOrOffset: gregor1.TimeOrOffset{},
	}
	res, err := cli.State(ctx, arg)
	if err != nil {
		return nil, err
	}
	return res, nil
}

func (c *Client) StateMachineConsumeMessage(m gregor1.Message) error {
	if _, err := c.sm.ConsumeMessage(m); err != nil {
		return err
	}

	// Check to see if we should save
	select {
	case <-c.saveTimer:
		c.log.Debug("StateMachineConsumeMessage(): saving local state")
		return c.Save()
	default:
		c.log.Debug("StateMachineConsumeMessage(): not saving local state")
		// Plow through if the timer isn't up
	}

	return nil
}

func (c *Client) StateMachineLatestCTime() *time.Time {
	return c.sm.LatestCTime(c.user, c.device)
}

func (c *Client) StateMachineInBandMessagesSince(t time.Time) ([]gregor.InBandMessage, error) {
	return c.sm.InBandMessagesSince(c.user, c.device, t)
}

func (c *Client) StateMachineState(t gregor.TimeOrOffset) (gregor.State, error) {
	return c.sm.State(c.user, c.device, t)
}
<commit_msg>export<commit_after>package storage

import (
	"bytes"
	"errors"
	"time"

	"github.com/keybase/client/go/gregor"
	"github.com/keybase/client/go/protocol/gregor1"
	"golang.org/x/net/context"

	rpc "github.com/keybase/go-framed-msgpack-rpc"
)

type LocalStorageEngine interface {
	Store(gregor.UID, []byte) error
	Load(gregor.UID) ([]byte, error)
}

type Client struct {
	User    gregor.UID
	Device  gregor.DeviceID
	Sm      gregor.StateMachine
	Storage LocalStorageEngine
	Log     rpc.LogOutput

	SaveTimer <-chan time.Time
}

func NewClient(user gregor.UID, device gregor.DeviceID, sm gregor.StateMachine,
	storage LocalStorageEngine, saveInterval time.Duration, log rpc.LogOutput) *Client {
	c := &Client{
		User:      user,
		Device:    device,
		Sm:        sm,
		Storage:   storage,
		Log:       log,
		SaveTimer: time.Tick(saveInterval), // How often we save to local storage
	}
	return c
}

func (c *Client) Save() error {
	if !c.Sm.IsEphemeral() {
		return errors.New("state machine is non-ephemeral")
	}

	state, err := c.Sm.State(c.User, c.Device, nil)
	if err != nil {
		return err
	}

	b, err := state.Marshal()
	if err != nil {
		return err
	}

	return c.Storage.Store(c.User, b)
}

func (c *Client) Restore() error {
	if !c.Sm.IsEphemeral() {
		return errors.New("state machine is non-ephemeral")
	}

	value, err := c.Storage.Load(c.User)
	if err != nil {
		return err
	}

	state, err := c.Sm.ObjFactory().UnmarshalState(value)
	if err != nil {
		return err
	}

	if err := c.Sm.InitState(state); err != nil {
		return err
	}

	return nil
}

type ErrHashMismatch struct{}

func (e ErrHashMismatch) Error() string {
	return "local state hash != server state hash"
}

func (c *Client) SyncFromTime(cli gregor1.IncomingInterface, t *time.Time) (msgs []gregor.InBandMessage, err error) {

	ctx, _ := context.WithTimeout(context.Background(), time.Second)
	arg := gregor1.SyncArg{
		Uid:      gregor1.UID(c.User.Bytes()),
		Deviceid: gregor1.DeviceID(c.Device.Bytes()),
	}
	if t != nil {
		arg.Ctime = gregor1.ToTime(*t)
	}

	// Grab the events from gregord
	c.Log.Debug("syncFromTime from: %s", gregor1.FromTime(arg.Ctime))
	res, err := cli.Sync(ctx, arg)
	if err != nil {
		return nil, err
	}

	c.Log.Debug("syncFromTime consuming %d messages", len(res.Msgs))
	for _, ibm := range res.Msgs {
		m := gregor1.Message{Ibm_: &ibm}
		msgs = append(msgs, ibm)
		c.Sm.ConsumeMessage(m)
	}

	// Check to make sure the server state is legit
	state, err := c.Sm.State(c.User, c.Device, nil)
	if err != nil {
		return nil, err
	}
	hash, err := state.Hash()
	if err != nil {
		return nil, err
	}
	if !bytes.Equal(res.Hash, hash) {
		return nil, ErrHashMismatch{}
	}

	return msgs, nil
}

func (c *Client) freshSync(cli gregor1.IncomingInterface) ([]gregor.InBandMessage, error) {

	var msgs []gregor.InBandMessage
	var err error

	c.Sm.Clear()
	state, err := c.State(cli)
	if err != nil {
		return msgs, err
	}
	if msgs, err = c.InBandMessagesFromState(state); err != nil {
		return msgs, err
	}
	if err = c.Sm.InitState(state); err != nil {
		return msgs, err
	}

	return msgs, nil
}

func (c *Client) Sync(cli gregor1.IncomingInterface) ([]gregor.InBandMessage, error) {
	latestCtime := c.Sm.LatestCTime(c.User, c.Device)
	if latestCtime == nil || latestCtime.IsZero() {
		c.Log.Debug("Sync(): fresh server sync: using State()")
		return c.freshSync(cli)
	} else {
		c.Log.Debug("Sync(): incremental server sync: using Sync()")
		if msgs, err := c.SyncFromTime(cli, latestCtime); err != nil {
			if _, ok := err.(ErrHashMismatch); ok {
				c.Log.Info("Sync failure: %v\nResetting StateMachine and retrying", err)
				return c.freshSync(cli)
			}
			return msgs, err
		} else {
			return msgs, nil
		}
	}
}

func (c *Client) InBandMessagesFromState(s gregor.State) ([]gregor.InBandMessage, error) {
	items, err := s.Items()
	if err != nil {
		return nil, err
	}

	var res []gregor.InBandMessage
	for _, i := range items {
		if ibm, err := c.Sm.ObjFactory().MakeInBandMessageFromItem(i); err == nil {
			res = append(res, ibm)
		}
	}
	return res, nil
}

func (c *Client) State(cli gregor1.IncomingInterface) (gregor.State, error) {
	ctx, _ := context.WithTimeout(context.Background(), 5*time.Second)
	arg := gregor1.StateArg{
		Uid:          gregor1.UID(c.User.Bytes()),
		Deviceid:     gregor1.DeviceID(c.Device.Bytes()),
		TimeOrOffset: gregor1.TimeOrOffset{},
	}
	res, err := cli.State(ctx, arg)
	if err != nil {
		return nil, err
	}
	return res, nil
}

func (c *Client) StateMachineConsumeMessage(m gregor1.Message) error {
	if _, err := c.Sm.ConsumeMessage(m); err != nil {
		return err
	}

	// Check to see if we should save
	select {
	case <-c.SaveTimer:
		c.Log.Debug("StateMachineConsumeMessage(): saving local state")
		return c.Save()
	default:
		c.Log.Debug("StateMachineConsumeMessage(): not saving local state")
		// Plow through if the timer isn't up
	}

	return nil
}

func (c *Client) StateMachineLatestCTime() *time.Time {
	return c.Sm.LatestCTime(c.User, c.Device)
}

func (c *Client) StateMachineInBandMessagesSince(t time.Time) ([]gregor.InBandMessage, error) {
	return c.Sm.InBandMessagesSince(c.User, c.Device, t)
}

func (c *Client) StateMachineState(t gregor.TimeOrOffset) (gregor.State, error) {
	return c.Sm.State(c.User, c.Device, t)
}
<|endoftext|>


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

package doorman

import (
	"flag"
	"time"

	log "github.com/golang/glog"
)

var cleaningInterval = flag.Duration("leasestore-cleanup-interval", 30*time.Second, "Cleaning interval of the lease store")

// Lease represents a lease on capacity for some resource for some client.
type Lease struct {
	// Expiry is the time at which this lease expires.
	Expiry time.Time

	// RefreshInterval is the interval at which the client should
	// poll for updates for this resource.
	RefreshInterval time.Duration

	// Has is how much capacity was given to this client.
	Has float64

	// Wants is how much capacity was requested.
	Wants float64

	// Subclients is the number of subclients of this client.
	Subclients int64
}

// ClientLeaseStatus contains information about a lease that
// a client has.
type ClientLeaseStatus struct {
	ClientID string
	Lease    Lease
}

// ResourceLeaseStatus is a read-only copy of information about all leases
// for a particular resource. It can be used for display purposes.
type ResourceLeaseStatus struct {
	// Resource id.
	ID string

	// The sum of all outstanding obligations.
	SumHas float64

	// The sum of all client wants.
	SumWants float64

	// A slice containing information about outstanding leases
	Leases []ClientLeaseStatus
}

// IsZero returns true if the lease is a zero instance.
func (lease Lease) IsZero() bool {
	return lease.Expiry.IsZero()
}

// LeaseStore is the set of outstanding leases for some particular
// resource.
type LeaseStore interface {
	// SumHas returns the total capacity given out to clients
	// at the moment.
	SumHas() float64

	// SumWants returns the total capacity requested by clients.
	SumWants() float64

	// Get returns the lease currently granted to this client.
	Get(client string) Lease

	// HasClient returns whether this client currently has a lease.
	HasClient(client string) bool

	// Assign updates the lease store to represent capacity given
	// to a client.
	Assign(client string, leaseLength, refreshInterval time.Duration, has, wants float64, subclients int64) Lease

	// Release releases any leases granted to a client.
	Release(client string)

	// Clean removes any out of date leases. Returns the number of leases cleaned
	Clean() int

	// Count returns the number of subclients leases in the store.
	Count() int64

	// Returns a read-only copy of summary information about outstanding leases.
	ResourceLeaseStatus() ResourceLeaseStatus

	// Map executes a function for every lease in the store.
	Map(func(id string, lease Lease))

	// Subclients returns the number of subclients of the client with the given id.
	Subclients(id string) int64
}

type leaseStoreImpl struct {
	id       string
	leases   map[string]Lease
	sumWants float64
	sumHas   float64
}

// New returns a fresh LeaseStore implementation.
func NewLeaseStore(id string) LeaseStore {
	store := &leaseStoreImpl{
		id:     id,
		leases: make(map[string]Lease),
	}

	// Starts a goroutine that regularly cleans the store of
	// expired leases.
	go func() {
		tick := time.Tick(*cleaningInterval)

		for {
			select {
			case <-tick:
				if n := store.Clean(); n > 0 {
					log.Infof("cleaning lease store for resource %s removed %d lease(s)", store.id, n)
				}
			}
		}
	}()

	return store
}

func (store *leaseStoreImpl) Count() int64 {
	var subclients int64
	for _, lease := range store.leases {
		subclients += lease.Subclients
	}
	return subclients
}

func (store *leaseStoreImpl) SumWants() float64 {
	return store.sumWants
}

func (store *leaseStoreImpl) modifySumWants(diff float64) {
	store.sumWants += diff
}

func (store *leaseStoreImpl) SumHas() float64 {
	return store.sumHas
}

func (store *leaseStoreImpl) modifySumHas(diff float64) {
	store.sumHas += diff
}

func (store *leaseStoreImpl) HasClient(client string) bool {
	_, ok := store.leases[client]
	return ok
}

func (store *leaseStoreImpl) Get(client string) Lease {
	return store.leases[client]
}

func (store *leaseStoreImpl) Release(client string) {
	lease, ok := store.leases[client]
	if !ok {
		return
	}
	store.modifySumWants(-lease.Wants)
	store.modifySumHas(-lease.Has)
	delete(store.leases, client)
}

func (store *leaseStoreImpl) Assign(client string, leaseLength, refreshInterval time.Duration, has, wants float64, subclients int64) Lease {
	lease := store.leases[client]

	store.modifySumHas(has - lease.Has)
	store.modifySumWants(wants - lease.Wants)

	lease.Has, lease.Wants = has, wants
	lease.Expiry = time.Now().Add(leaseLength)
	lease.RefreshInterval = refreshInterval
	lease.Subclients = subclients
	store.leases[client] = lease

	return lease
}

func (store *leaseStoreImpl) Clean() int {
	when := time.Now()
	result := 0

	for client, lease := range store.leases {
		if when.After(lease.Expiry) {
			store.Release(client)
			result++
		}
	}

	return result
}

// ResourceLeaseStatus returns a read-only copy with summary information about outstanding leases.
func (store *leaseStoreImpl) ResourceLeaseStatus() ResourceLeaseStatus {
	status := ResourceLeaseStatus{
		ID:       store.id,
		SumHas:   store.sumHas,
		SumWants: store.sumWants,
		Leases:   make([]ClientLeaseStatus, 0, len(store.leases)),
	}

	for client, lease := range store.leases {
		status.Leases = append(status.Leases, ClientLeaseStatus{
			ClientID: client,
			Lease:    lease,
		})
	}

	return status
}

// Map executes a function for every lease in the store.
func (store *leaseStoreImpl) Map(fun func(string, Lease)) {
	for id, lease := range store.leases {
		fun(id, lease)
	}
}

// Subclients returns the number of subclients of the client with the given id.
// Every client has at least one subclient.
func (store *leaseStoreImpl) Subclients(id string) int64 {
	return store.leases[id].Subclients
}
<commit_msg>Remove racy code.<commit_after>// Copyright 2016 Google, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package doorman

import (
	"flag"
	"time"
)

var cleaningInterval = flag.Duration("leasestore-cleanup-interval", 30*time.Second, "Cleaning interval of the lease store")

// Lease represents a lease on capacity for some resource for some client.
type Lease struct {
	// Expiry is the time at which this lease expires.
	Expiry time.Time

	// RefreshInterval is the interval at which the client should
	// poll for updates for this resource.
	RefreshInterval time.Duration

	// Has is how much capacity was given to this client.
	Has float64

	// Wants is how much capacity was requested.
	Wants float64

	// Subclients is the number of subclients of this client.
	Subclients int64
}

// ClientLeaseStatus contains information about a lease that
// a client has.
type ClientLeaseStatus struct {
	ClientID string
	Lease    Lease
}

// ResourceLeaseStatus is a read-only copy of information about all leases
// for a particular resource. It can be used for display purposes.
type ResourceLeaseStatus struct {
	// Resource id.
	ID string

	// The sum of all outstanding obligations.
	SumHas float64

	// The sum of all client wants.
	SumWants float64

	// A slice containing information about outstanding leases
	Leases []ClientLeaseStatus
}

// IsZero returns true if the lease is a zero instance.
func (lease Lease) IsZero() bool {
	return lease.Expiry.IsZero()
}

// LeaseStore is the set of outstanding leases for some particular
// resource.
type LeaseStore interface {
	// SumHas returns the total capacity given out to clients
	// at the moment.
	SumHas() float64

	// SumWants returns the total capacity requested by clients.
	SumWants() float64

	// Get returns the lease currently granted to this client.
	Get(client string) Lease

	// HasClient returns whether this client currently has a lease.
	HasClient(client string) bool

	// Assign updates the lease store to represent capacity given
	// to a client.
	Assign(client string, leaseLength, refreshInterval time.Duration, has, wants float64, subclients int64) Lease

	// Release releases any leases granted to a client.
	Release(client string)

	// Clean removes any out of date leases. Returns the number of leases cleaned
	Clean() int

	// Count returns the number of subclients leases in the store.
	Count() int64

	// Returns a read-only copy of summary information about outstanding leases.
	ResourceLeaseStatus() ResourceLeaseStatus

	// Map executes a function for every lease in the store.
	Map(func(id string, lease Lease))

	// Subclients returns the number of subclients of the client with the given id.
	Subclients(id string) int64
}

type leaseStoreImpl struct {
	id       string
	leases   map[string]Lease
	sumWants float64
	sumHas   float64
}

// New returns a fresh LeaseStore implementation.
func NewLeaseStore(id string) LeaseStore {
	return &leaseStoreImpl{
		id:     id,
		leases: make(map[string]Lease),
	}
}

func (store *leaseStoreImpl) Count() int64 {
	var subclients int64
	for _, lease := range store.leases {
		subclients += lease.Subclients
	}
	return subclients
}

func (store *leaseStoreImpl) SumWants() float64 {
	return store.sumWants
}

func (store *leaseStoreImpl) modifySumWants(diff float64) {
	store.sumWants += diff
}

func (store *leaseStoreImpl) SumHas() float64 {
	return store.sumHas
}

func (store *leaseStoreImpl) modifySumHas(diff float64) {
	store.sumHas += diff
}

func (store *leaseStoreImpl) HasClient(client string) bool {
	_, ok := store.leases[client]
	return ok
}

func (store *leaseStoreImpl) Get(client string) Lease {
	return store.leases[client]
}

func (store *leaseStoreImpl) Release(client string) {
	lease, ok := store.leases[client]
	if !ok {
		return
	}
	store.modifySumWants(-lease.Wants)
	store.modifySumHas(-lease.Has)
	delete(store.leases, client)
}

func (store *leaseStoreImpl) Assign(client string, leaseLength, refreshInterval time.Duration, has, wants float64, subclients int64) Lease {
	lease := store.leases[client]

	store.modifySumHas(has - lease.Has)
	store.modifySumWants(wants - lease.Wants)

	lease.Has, lease.Wants = has, wants
	lease.Expiry = time.Now().Add(leaseLength)
	lease.RefreshInterval = refreshInterval
	lease.Subclients = subclients
	store.leases[client] = lease

	return lease
}

func (store *leaseStoreImpl) Clean() int {
	when := time.Now()
	result := 0

	for client, lease := range store.leases {
		if when.After(lease.Expiry) {
			store.Release(client)
			result++
		}
	}

	return result
}

// ResourceLeaseStatus returns a read-only copy with summary information about outstanding leases.
func (store *leaseStoreImpl) ResourceLeaseStatus() ResourceLeaseStatus {
	status := ResourceLeaseStatus{
		ID:       store.id,
		SumHas:   store.sumHas,
		SumWants: store.sumWants,
		Leases:   make([]ClientLeaseStatus, 0, len(store.leases)),
	}

	for client, lease := range store.leases {
		status.Leases = append(status.Leases, ClientLeaseStatus{
			ClientID: client,
			Lease:    lease,
		})
	}

	return status
}

// Map executes a function for every lease in the store.
func (store *leaseStoreImpl) Map(fun func(string, Lease)) {
	for id, lease := range store.leases {
		fun(id, lease)
	}
}

// Subclients returns the number of subclients of the client with the given id.
// Every client has at least one subclient.
func (store *leaseStoreImpl) Subclients(id string) int64 {
	return store.leases[id].Subclients
}
<|endoftext|>


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

package wrangler

import (
	"time"

	"code.google.com/p/vitess/go/relog"
	tm "code.google.com/p/vitess/go/vt/tabletmanager"
	"code.google.com/p/vitess/go/zk"
)

const (
	DefaultActionTimeout = 30 * time.Second
	DefaultLockTimeout   = 30 * time.Second
)

type Wrangler struct {
	zconn       zk.Conn
	ai          *tm.ActionInitiator
	deadline    time.Time
	lockTimeout time.Duration
}

// actionTimeout: how long should we wait for an action to complete?
// lockTimeout: how long should we wait for the initial lock to start a complex action?
//   This is distinct from actionTimeout because most of the time, we want to immediately
//   know that out action will fail. However, automated action will need some time to
//   arbitrate the locks.
func NewWrangler(zconn zk.Conn, actionTimeout, lockTimeout time.Duration) *Wrangler {
	return &Wrangler{zconn, tm.NewActionInitiator(zconn), time.Now().Add(actionTimeout), lockTimeout}
}

func (wr *Wrangler) actionTimeout() time.Duration {
	return wr.deadline.Sub(time.Now())
}

func (wr *Wrangler) readTablet(zkTabletPath string) (*tm.TabletInfo, error) {
	return tm.ReadTablet(wr.zconn, zkTabletPath)
}

func (wr *Wrangler) ZkConn() zk.Conn {
	return wr.zconn
}

func (wr *Wrangler) ActionInitiator() *tm.ActionInitiator {
	return wr.ai
}

// Change the type of tablet and recompute all necessary derived paths in the
// serving graph.
// force: Bypass the vtaction system and make the data change directly, and
// do not run the idle_server_check nor live_server_check hooks
func (wr *Wrangler) ChangeType(zkTabletPath string, dbType tm.TabletType, force bool) error {
	// Load tablet to find keyspace and shard assignment.
	// Don't load after the ChangeType which might have unassigned
	// the tablet.
	ti, err := tm.ReadTablet(wr.zconn, zkTabletPath)
	if err != nil {
		return err
	}
	rebuildRequired := ti.Tablet.IsServingType()

	if force {
		// with --force, we do not run any hook
		err = tm.ChangeType(wr.zconn, zkTabletPath, dbType)
	} else {
		// the remote action will run the hooks
		actionPath, err := wr.ai.ChangeType(zkTabletPath, dbType)
		// You don't have a choice - you must wait for
		// completion before rebuilding.
		if err == nil {
			err = wr.ai.WaitForCompletion(actionPath, DefaultActionTimeout)
		}
	}

	if err != nil {
		return err
	}

	// we rebuild if the tablet was serving, or if it is now
	var shardToRebuild string
	if rebuildRequired {
		shardToRebuild = ti.ShardPath()
	} else {
		// re-read the tablet, see if we become serving
		ti, err := tm.ReadTablet(wr.zconn, zkTabletPath)
		if err != nil {
			return err
		}
		if ti.Tablet.IsServingType() {
			rebuildRequired = true
			shardToRebuild = ti.ShardPath()
		}
	}

	if rebuildRequired {
		if _, err := wr.RebuildShardGraph(shardToRebuild); err != nil {
			return err
		}
	}
	return nil
}

// same as ChangeType, but assume we already have the shard lock,
// and do not have the option to force anything
// FIXME(alainjobart): doesn't rebuild the Keyspace, as that part has locks,
// so the local serving graphs will be wrong. To do that, I need to refactor
// some code, might be a bigger change.
// Mike says: Updating the shard should be good enough. I'm debating dropping the entire
// keyspace rollup, since I think that is adding complexity and feels like it might
// be a premature optimization.
func (wr *Wrangler) changeTypeInternal(zkTabletPath string, dbType tm.TabletType) error {
	ti, err := tm.ReadTablet(wr.zconn, zkTabletPath)
	if err != nil {
		return err
	}
	rebuildRequired := ti.Tablet.IsServingType()

	// change the type
	actionPath, err := wr.ai.ChangeType(ti.Path(), dbType)
	if err != nil {
		return err
	}
	err = wr.ai.WaitForCompletion(actionPath, wr.actionTimeout())
	if err != nil {
		return err
	}

	// rebuild if necessary
	if rebuildRequired {
		err = wr.rebuildShard(ti.ShardPath(), false)
		if err != nil {
			return err
		}
		// FIXME(alainjobart) We already have the lock on one shard, so this is not
		// possible. But maybe it's not necessary anyway.
		// We could pass in a shard path we already have the lock on, and skip it?
		//		err = wr.rebuildKeyspace(ti.KeyspacePath())
		//		if err != nil {
		//			return err
		//		}
	}
	return nil
}

// Cleanup an action node and write back status/error to zk.
// Only returns an error if something went wrong with zk.
func (wr *Wrangler) handleActionError(actionPath string, actionErr error) error {
	// re-read the action node
	data, _, err := wr.zconn.Get(actionPath)
	if err != nil {
		return err
	}
	var actionNode *tm.ActionNode
	actionNode, err = tm.ActionNodeFromJson(data, actionPath)
	if err != nil {
		return err
	}

	// write what happened to the action log
	err = tm.StoreActionResponse(wr.zconn, actionNode, actionPath, actionErr)
	if err != nil {
		return err
	}

	// no error, we can unblock the action queue
	if actionErr == nil {
		return zk.DeleteRecursive(wr.zconn, actionPath, -1)
	}
	return nil
}

// Scrap a tablet. If force is used, we write to ZK directly and don't
// remote-execute the command.
func (wr *Wrangler) Scrap(zkTabletPath string, force, skipRebuild bool) (actionPath string, err error) {
	// load the tablet, see if we'll need to rebuild
	ti, err := tm.ReadTablet(wr.zconn, zkTabletPath)
	if err != nil {
		return "", err
	}
	rebuildRequired := ti.Tablet.IsServingType()

	if force {
		err = tm.Scrap(wr.zconn, zkTabletPath, force)
	} else {
		actionPath, err = wr.ai.Scrap(zkTabletPath)
	}

	if !rebuildRequired {
		relog.Info("Rebuild not required")
		return
	}
	if skipRebuild {
		relog.Warning("Rebuild required, but skipping it")
		return
	}

	// wait for the remote Scrap if necessary
	if actionPath != "" {
		err = wr.ai.WaitForCompletion(actionPath, wr.actionTimeout())
		if err != nil {
			return "", nil
		}
	}

	// and rebuild the original shard / keyspace
	return wr.RebuildShardGraph(ti.ShardPath())
}
<commit_msg>Fixing an error reporting bug.<commit_after>// Copyright 2012, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package wrangler

import (
	"time"

	"code.google.com/p/vitess/go/relog"
	tm "code.google.com/p/vitess/go/vt/tabletmanager"
	"code.google.com/p/vitess/go/zk"
)

const (
	DefaultActionTimeout = 30 * time.Second
	DefaultLockTimeout   = 30 * time.Second
)

type Wrangler struct {
	zconn       zk.Conn
	ai          *tm.ActionInitiator
	deadline    time.Time
	lockTimeout time.Duration
}

// actionTimeout: how long should we wait for an action to complete?
// lockTimeout: how long should we wait for the initial lock to start a complex action?
//   This is distinct from actionTimeout because most of the time, we want to immediately
//   know that out action will fail. However, automated action will need some time to
//   arbitrate the locks.
func NewWrangler(zconn zk.Conn, actionTimeout, lockTimeout time.Duration) *Wrangler {
	return &Wrangler{zconn, tm.NewActionInitiator(zconn), time.Now().Add(actionTimeout), lockTimeout}
}

func (wr *Wrangler) actionTimeout() time.Duration {
	return wr.deadline.Sub(time.Now())
}

func (wr *Wrangler) readTablet(zkTabletPath string) (*tm.TabletInfo, error) {
	return tm.ReadTablet(wr.zconn, zkTabletPath)
}

func (wr *Wrangler) ZkConn() zk.Conn {
	return wr.zconn
}

func (wr *Wrangler) ActionInitiator() *tm.ActionInitiator {
	return wr.ai
}

// Change the type of tablet and recompute all necessary derived paths in the
// serving graph.
// force: Bypass the vtaction system and make the data change directly, and
// do not run the idle_server_check nor live_server_check hooks
func (wr *Wrangler) ChangeType(zkTabletPath string, dbType tm.TabletType, force bool) error {
	// Load tablet to find keyspace and shard assignment.
	// Don't load after the ChangeType which might have unassigned
	// the tablet.
	ti, err := tm.ReadTablet(wr.zconn, zkTabletPath)
	if err != nil {
		return err
	}
	rebuildRequired := ti.Tablet.IsServingType()

	if force {
		// with --force, we do not run any hook
		err = tm.ChangeType(wr.zconn, zkTabletPath, dbType)
	} else {
		// the remote action will run the hooks
		var actionPath string
		actionPath, err = wr.ai.ChangeType(zkTabletPath, dbType)
		// You don't have a choice - you must wait for
		// completion before rebuilding.
		if err == nil {
			err = wr.ai.WaitForCompletion(actionPath, DefaultActionTimeout)
		}
	}

	if err != nil {
		return err
	}

	// we rebuild if the tablet was serving, or if it is now
	var shardToRebuild string
	if rebuildRequired {
		shardToRebuild = ti.ShardPath()
	} else {
		// re-read the tablet, see if we become serving
		ti, err := tm.ReadTablet(wr.zconn, zkTabletPath)
		if err != nil {
			return err
		}
		if ti.Tablet.IsServingType() {
			rebuildRequired = true
			shardToRebuild = ti.ShardPath()
		}
	}

	if rebuildRequired {
		if _, err := wr.RebuildShardGraph(shardToRebuild); err != nil {
			return err
		}
	}
	return nil
}

// same as ChangeType, but assume we already have the shard lock,
// and do not have the option to force anything
// FIXME(alainjobart): doesn't rebuild the Keyspace, as that part has locks,
// so the local serving graphs will be wrong. To do that, I need to refactor
// some code, might be a bigger change.
// Mike says: Updating the shard should be good enough. I'm debating dropping the entire
// keyspace rollup, since I think that is adding complexity and feels like it might
// be a premature optimization.
func (wr *Wrangler) changeTypeInternal(zkTabletPath string, dbType tm.TabletType) error {
	ti, err := tm.ReadTablet(wr.zconn, zkTabletPath)
	if err != nil {
		return err
	}
	rebuildRequired := ti.Tablet.IsServingType()

	// change the type
	actionPath, err := wr.ai.ChangeType(ti.Path(), dbType)
	if err != nil {
		return err
	}
	err = wr.ai.WaitForCompletion(actionPath, wr.actionTimeout())
	if err != nil {
		return err
	}

	// rebuild if necessary
	if rebuildRequired {
		err = wr.rebuildShard(ti.ShardPath(), false)
		if err != nil {
			return err
		}
		// FIXME(alainjobart) We already have the lock on one shard, so this is not
		// possible. But maybe it's not necessary anyway.
		// We could pass in a shard path we already have the lock on, and skip it?
		//		err = wr.rebuildKeyspace(ti.KeyspacePath())
		//		if err != nil {
		//			return err
		//		}
	}
	return nil
}

// Cleanup an action node and write back status/error to zk.
// Only returns an error if something went wrong with zk.
func (wr *Wrangler) handleActionError(actionPath string, actionErr error) error {
	// re-read the action node
	data, _, err := wr.zconn.Get(actionPath)
	if err != nil {
		return err
	}
	var actionNode *tm.ActionNode
	actionNode, err = tm.ActionNodeFromJson(data, actionPath)
	if err != nil {
		return err
	}

	// write what happened to the action log
	err = tm.StoreActionResponse(wr.zconn, actionNode, actionPath, actionErr)
	if err != nil {
		return err
	}

	// no error, we can unblock the action queue
	if actionErr == nil {
		return zk.DeleteRecursive(wr.zconn, actionPath, -1)
	}
	return nil
}

// Scrap a tablet. If force is used, we write to ZK directly and don't
// remote-execute the command.
func (wr *Wrangler) Scrap(zkTabletPath string, force, skipRebuild bool) (actionPath string, err error) {
	// load the tablet, see if we'll need to rebuild
	ti, err := tm.ReadTablet(wr.zconn, zkTabletPath)
	if err != nil {
		return "", err
	}
	rebuildRequired := ti.Tablet.IsServingType()

	if force {
		err = tm.Scrap(wr.zconn, zkTabletPath, force)
	} else {
		actionPath, err = wr.ai.Scrap(zkTabletPath)
	}

	if !rebuildRequired {
		relog.Info("Rebuild not required")
		return
	}
	if skipRebuild {
		relog.Warning("Rebuild required, but skipping it")
		return
	}

	// wait for the remote Scrap if necessary
	if actionPath != "" {
		err = wr.ai.WaitForCompletion(actionPath, wr.actionTimeout())
		if err != nil {
			return "", nil
		}
	}

	// and rebuild the original shard / keyspace
	return wr.RebuildShardGraph(ti.ShardPath())
}
<|endoftext|>


<commit_before>/**
Message-oriented socket.

Provides a TCP connection for transferring messages.
Each message has its 32bit length.

Sending of messages is synchronous (see Send() method).
Reading of incoming messages is implemented as background
thread. 'handler' callback is called for each incoming
message.
*/

package tcpcall

import (
	"bufio"
	"encoding/binary"
	"errors"
	"io"
	"net"
	"sync"
	"tcpcall/pools"
	"time"
)

// Packet header Len
const HEADER_LEN = 4

// Errors
var (
	MsgConnNotConnectedError = errors.New("msg conn: not connected")
	MsgTooLongError          = errors.New("incoming message is too long")
)

// Message oriented connection
type MsgConn struct {
	socket    net.Conn
	buffer    *bufio.Writer
	lastFlush time.Time
	// socket write mutex
	socketMu sync.Mutex
	// maximum allowed length for incoming packets
	MaxPacketLen int
	// Minimum time between write buffer flushes
	minFlushPeriod time.Duration
	// incoming package handler
	handler func([]byte)
	// callback for disconnect event
	onDisconnect func()
	closed       bool
}

// Create new message oriented connection.
func NewMsgConn(socket net.Conn, minFlushPeriod time.Duration,
	writeBufferSize int,
	handler func([]byte), onClose func()) (*MsgConn, error) {
	if err := socket.SetReadDeadline(time.Time{}); err != nil {
		return nil, err
	}
	conn := &MsgConn{
		socket:         socket,
		buffer:         bufio.NewWriterSize(socket, writeBufferSize),
		handler:        handler,
		onDisconnect:   onClose,
		minFlushPeriod: minFlushPeriod,
	}
	go conn.readLoop()
	return conn, nil
}

// SetWriteDeadline sets the deadline for future Write calls
// and any currently-blocked Write call.
// A zero value for t means Write will not time out.
func (c *MsgConn) SetWriteDeadline(t time.Time) error {
	return c.socket.SetWriteDeadline(t)
}

// Send message to the other side.
func (c *MsgConn) Send(msg [][]byte) error {
	if c == nil {
		return MsgConnNotConnectedError
	}
	msgLen := 0
	for _, e := range msg {
		msgLen += len(e)
	}
	c.socketMu.Lock()
	if err := binary.Write(c.buffer, binary.BigEndian, uint32(msgLen)); err != nil {
		c.Close()
		c.socketMu.Unlock()
		return err
	}
	// write chunks one by one
	for _, e := range msg {
		if _, err := c.buffer.Write(e); err != nil {
			c.Close()
			c.socketMu.Unlock()
			return err
		}
	}
	// flush the buffer
	if c.minFlushPeriod <= 0 ||
		time.Now().After(c.lastFlush.Add(c.minFlushPeriod)) {
		if err := c.buffer.Flush(); err != nil {
			c.Close()
			c.socketMu.Unlock()
			return err
		}
		c.lastFlush = time.Now()
	}
	c.socketMu.Unlock()
	return nil
}

// Close connection.
func (c *MsgConn) Close() {
	c.socket.Close()
	c.closed = true
	if c.onDisconnect != nil {
		c.onDisconnect()
	}
}

// Return truth if connection is already closed.
func (c *MsgConn) Closed() bool {
	return c.closed
}

// Goroutine.
// Receive incoming messages from the other side
// and call callback function for each.
func (c *MsgConn) readLoop() {
	defer c.Close()
	for {
		packet, err := c.readPacket()
		if err != nil {
			return
		}
		c.handler(packet)
	}
}

// Receive next message from the other side.
func (c *MsgConn) readPacket() ([]byte, error) {
	header := pools.GetFreeBuffer(HEADER_LEN)
	if _, err := io.ReadAtLeast(c.socket, header, len(header)); err != nil {
		return nil, err
	}
	len := int(binary.BigEndian.Uint32(header))
	pools.AppendToBuffer(header)
	if 0 < c.MaxPacketLen && c.MaxPacketLen < len {
		return nil, MsgTooLongError
	}
	buffer := make([]byte, len)
	if _, err := io.ReadAtLeast(c.socket, buffer, len); err != nil {
		return nil, err
	}
	return buffer[:len], nil
}
<commit_msg>Style fix for function signature<commit_after>/**
Message-oriented socket.

Provides a TCP connection for transferring messages.
Each message has its 32bit length.

Sending of messages is synchronous (see Send() method).
Reading of incoming messages is implemented as background
thread. 'handler' callback is called for each incoming
message.
*/

package tcpcall

import (
	"bufio"
	"encoding/binary"
	"errors"
	"io"
	"net"
	"sync"
	"tcpcall/pools"
	"time"
)

// Packet header Len
const HEADER_LEN = 4

// Errors
var (
	MsgConnNotConnectedError = errors.New("msg conn: not connected")
	MsgTooLongError          = errors.New("incoming message is too long")
)

// Message oriented connection
type MsgConn struct {
	socket    net.Conn
	buffer    *bufio.Writer
	lastFlush time.Time
	// socket write mutex
	socketMu sync.Mutex
	// maximum allowed length for incoming packets
	MaxPacketLen int
	// Minimum time between write buffer flushes
	minFlushPeriod time.Duration
	// incoming package handler
	handler func([]byte)
	// callback for disconnect event
	onDisconnect func()
	closed       bool
}

// Create new message oriented connection.
func NewMsgConn(
	socket net.Conn,
	minFlushPeriod time.Duration,
	writeBufferSize int,
	handler func([]byte),
	onClose func(),
) (*MsgConn, error) {
	if err := socket.SetReadDeadline(time.Time{}); err != nil {
		return nil, err
	}
	conn := &MsgConn{
		socket:         socket,
		buffer:         bufio.NewWriterSize(socket, writeBufferSize),
		handler:        handler,
		onDisconnect:   onClose,
		minFlushPeriod: minFlushPeriod,
	}
	go conn.readLoop()
	return conn, nil
}

// SetWriteDeadline sets the deadline for future Write calls
// and any currently-blocked Write call.
// A zero value for t means Write will not time out.
func (c *MsgConn) SetWriteDeadline(t time.Time) error {
	return c.socket.SetWriteDeadline(t)
}

// Send message to the other side.
func (c *MsgConn) Send(msg [][]byte) error {
	if c == nil {
		return MsgConnNotConnectedError
	}
	msgLen := 0
	for _, e := range msg {
		msgLen += len(e)
	}
	c.socketMu.Lock()
	if err := binary.Write(c.buffer, binary.BigEndian, uint32(msgLen)); err != nil {
		c.Close()
		c.socketMu.Unlock()
		return err
	}
	// write chunks one by one
	for _, e := range msg {
		if _, err := c.buffer.Write(e); err != nil {
			c.Close()
			c.socketMu.Unlock()
			return err
		}
	}
	// flush the buffer
	if c.minFlushPeriod <= 0 ||
		time.Now().After(c.lastFlush.Add(c.minFlushPeriod)) {
		if err := c.buffer.Flush(); err != nil {
			c.Close()
			c.socketMu.Unlock()
			return err
		}
		c.lastFlush = time.Now()
	}
	c.socketMu.Unlock()
	return nil
}

// Close connection.
func (c *MsgConn) Close() {
	c.socket.Close()
	c.closed = true
	if c.onDisconnect != nil {
		c.onDisconnect()
	}
}

// Return truth if connection is already closed.
func (c *MsgConn) Closed() bool {
	return c.closed
}

// Goroutine.
// Receive incoming messages from the other side
// and call callback function for each.
func (c *MsgConn) readLoop() {
	defer c.Close()
	for {
		packet, err := c.readPacket()
		if err != nil {
			return
		}
		c.handler(packet)
	}
}

// Receive next message from the other side.
func (c *MsgConn) readPacket() ([]byte, error) {
	header := pools.GetFreeBuffer(HEADER_LEN)
	if _, err := io.ReadAtLeast(c.socket, header, len(header)); err != nil {
		return nil, err
	}
	len := int(binary.BigEndian.Uint32(header))
	pools.AppendToBuffer(header)
	if 0 < c.MaxPacketLen && c.MaxPacketLen < len {
		return nil, MsgTooLongError
	}
	buffer := make([]byte, len)
	if _, err := io.ReadAtLeast(c.socket, buffer, len); err != nil {
		return nil, err
	}
	return buffer[:len], nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"
)

var shaTypes = [...]string{"sha256", "sha384", "sha512"}

func main() {

	shaType := flag.String("sha", "sha256", "type of sha")
	flag.Parse()
	// flag.Visit(func(f *flag.Flag) {
	// 	val := f.Value.String()
	// 	for _, ss := range shaTypes {
	// 		if val == ss {
	// 			return
	// 		}
	// 	}
	// })
	switch *shaType {
	case "sha256":
		fmt.Println("sha256")
	// sha256.Sum256()
	case "sha384":
		fmt.Println("sha384")
	// sha512.Sum384()
	case "sha512":
		fmt.Println("sha512")
		// sha512.Sum512()
	default:
		fmt.Fprintf(os.Stderr, "error: invalid sha type '%s'\n", *shaType)
		os.Exit(1)
	}
}
<commit_msg>Draft solution for task 4.2<commit_after>package main

import (
	"bytes"
	"crypto/sha256"
	"crypto/sha512"
	"flag"
	"fmt"
	"io"
	"os"
)

var shaTypes = [...]string{"sha256", "sha384", "sha512"}

func main() {

	shaType := flag.String("sha", "sha256", "type of sha")
	flag.Parse()
	flag.Visit(func(f *flag.Flag) {
		val := f.Value.String()
		for _, ss := range shaTypes {
			if val == ss {
				return
			}
		}
		fmt.Fprintf(os.Stderr, "error: invalid sha '%s'\n", *shaType)
		os.Exit(1)
	})

	var buff bytes.Buffer
	io.Copy(&buff, os.Stdin)

	switch *shaType {
	case "sha256":
		fmt.Printf("%x\n", sha256.Sum256(buff.Bytes()))
	case "sha384":
		fmt.Printf("%x\n", sha512.Sum384(buff.Bytes()))
	case "sha512":
		fmt.Printf("%x\n", sha512.Sum512(buff.Bytes()))
	}
}
<|endoftext|>


<commit_before>package convertutil

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/grokify/gotilla/image/imageutil"
	"github.com/grokify/gotilla/io/ioutilmore"
	"github.com/pkg/errors"
)

const (
	PDFSpecs    = "1950pxw300dpi"
	KindleSpecs = "600pxw300dpi"
)

type CopyType int

const (
	PDFFormat    CopyType = iota // 0 convert cli value
	KindleFormat                 // 1 convert cli value
)

// ReformatImages converts images in one dir to another using default
// formats for Kindle and PDF.
func ReformatImages(baseSrcDir, baseOutDir string, copyType CopyType, rewrite bool) error {
	var err error
	baseSrcDir, err = filepath.Abs(strings.TrimSpace(baseSrcDir))
	if err != nil {
		return err
	}
	baseOutDir, err = filepath.Abs(strings.TrimSpace(baseOutDir))
	if err != nil {
		return err
	}
	return reformatImagesSubdir(baseSrcDir, baseOutDir, "", copyType, rewrite)
}

func reformatImagesSubdir(baseSrcDir, baseOutDir, dirPart string, copyType CopyType, rewrite bool) error {
	thisSrcDir := baseSrcDir
	thisOutDir := baseOutDir
	dirPart = strings.TrimSpace(dirPart)
	if len(dirPart) > 0 {
		thisSrcDir = filepath.Join(thisSrcDir, dirPart)
		thisOutDir = filepath.Join(thisOutDir, dirPart)

		isDir, err := ioutilmore.IsDir(thisSrcDir)
		if err != nil {
			return err
		}
		if !isDir {
			return fmt.Errorf("Need Dir [%s]", thisSrcDir)
		}
	}

	if err := os.MkdirAll(thisOutDir, 0755); err != nil {
		return err
	}

	sdirs, files, err := ioutilmore.ReadDirSplit(thisSrcDir, false)
	if err != nil {
		return err
	}
	for _, file := range files {
		thisSrcFile := filepath.Join(thisSrcDir, file.Name())
		thisOutFile := filepath.Join(thisOutDir, file.Name())
		if !imageutil.IsImageExt(thisSrcFile) {
			continue
		}
		if !rewrite {
			isFile, err := ioutilmore.IsFile(thisOutFile)
			if err == nil && isFile {
				continue
			}
		}

		switch copyType {
		case PDFFormat:
			_, _, err := ConvertToPDF(thisSrcFile, thisOutFile)
			if err != nil {
				return errors.Wrap(err, fmt.Sprintf("ConvertToPDF failed for [%s]", thisSrcFile))
			}
		case KindleFormat:
			_, stderr, err := ConvertToKindle(thisSrcFile, thisOutFile)
			err = CheckError(err, stderr)
			if err != nil {
				return err
			}
		}
	}
	for _, sdir := range sdirs {
		subDir := sdir.Name()
		if len(dirPart) > 0 {
			subDir = filepath.Join(dirPart, subDir)
		}
		err := reformatImagesSubdir(baseSrcDir, baseOutDir, subDir, copyType, rewrite)
		if err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>lint: image/convertutil: reduce cyclomatic complexity<commit_after>package convertutil

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/grokify/gotilla/image/imageutil"
	"github.com/grokify/gotilla/io/ioutilmore"
	"github.com/pkg/errors"
)

const (
	PDFSpecs    = "1950pxw300dpi"
	KindleSpecs = "600pxw300dpi"
)

type CopyType int

const (
	PDFFormat    CopyType = iota // 0 convert cli value
	KindleFormat                 // 1 convert cli value
)

// ReformatImages converts images in one dir to another using default
// formats for Kindle and PDF.
func ReformatImages(baseSrcDir, baseOutDir string, copyType CopyType, rewrite bool) error {
	var err error
	baseSrcDir, err = filepath.Abs(strings.TrimSpace(baseSrcDir))
	if err != nil {
		return err
	}
	baseOutDir, err = filepath.Abs(strings.TrimSpace(baseOutDir))
	if err != nil {
		return err
	}
	return reformatImagesSubdir(baseSrcDir, baseOutDir, "", copyType, rewrite)
}

func reformatImagesSubdir(baseSrcDir, baseOutDir, dirPart string, copyType CopyType, rewrite bool) error {
	thisSrcDir := baseSrcDir
	thisOutDir := baseOutDir
	dirPart = strings.TrimSpace(dirPart)
	if len(dirPart) > 0 {
		thisSrcDir = filepath.Join(thisSrcDir, dirPart)
		thisOutDir = filepath.Join(thisOutDir, dirPart)

		isDir, err := ioutilmore.IsDir(thisSrcDir)
		if err != nil {
			return err
		}
		if !isDir {
			return fmt.Errorf("Need Dir [%s]", thisSrcDir)
		}
	}

	if err := os.MkdirAll(thisOutDir, 0755); err != nil {
		return err
	}

	sdirs, files, err := ioutilmore.ReadDirSplit(thisSrcDir, false)
	if err != nil {
		return err
	}
	for _, file := range files {
		err := reformatImagesSubdirFile(
			filepath.Join(thisSrcDir, file.Name()),
			filepath.Join(thisOutDir, file.Name()),
			copyType, rewrite)
		if err != nil {
			return err
		}
		/*
			thisSrcFile := filepath.Join(thisSrcDir, file.Name())
			thisOutFile := filepath.Join(thisOutDir, file.Name())

			if !imageutil.IsImageExt(thisSrcFile) {
				continue
			}
			if !rewrite {
				isFile, err := ioutilmore.IsFile(thisOutFile)
				if err == nil && isFile {
					continue
				}
			}

			switch copyType {
			case PDFFormat:
				_, _, err := ConvertToPDF(thisSrcFile, thisOutFile)
				if err != nil {
					return errors.Wrap(err, fmt.Sprintf("ConvertToPDF failed for [%s]", thisSrcFile))
				}
			case KindleFormat:
				_, stderr, err := ConvertToKindle(thisSrcFile, thisOutFile)
				err = CheckError(err, stderr)
				if err != nil {
					return err
				}
			}
		*/
	}
	for _, sdir := range sdirs {
		subDir := sdir.Name()
		if len(dirPart) > 0 {
			subDir = filepath.Join(dirPart, subDir)
		}
		err := reformatImagesSubdir(baseSrcDir, baseOutDir, subDir, copyType, rewrite)
		if err != nil {
			return err
		}
	}
	return nil
}

func reformatImagesSubdirFile(thisSrcFile, thisOutFile string, copyType CopyType, rewrite bool) error {
	if !imageutil.IsImageExt(thisSrcFile) {
		return nil
	}
	if !rewrite {
		isFile, err := ioutilmore.IsFile(thisOutFile)
		if err == nil && isFile {
			return nil
		}
	}

	switch copyType {
	case PDFFormat:
		_, _, err := ConvertToPDF(thisSrcFile, thisOutFile)
		if err != nil {
			return errors.Wrap(err, fmt.Sprintf("ConvertToPDF failed for [%s]", thisSrcFile))
		}
	case KindleFormat:
		_, stderr, err := ConvertToKindle(thisSrcFile, thisOutFile)
		err = CheckError(err, stderr)
		if err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2018 TiKV Project 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package cases

import (
	"bytes"
	"fmt"
	"math/rand"
	"os"

	"github.com/go-echarts/go-echarts/charts"
	"github.com/pingcap/kvproto/pkg/metapb"
	"github.com/pingcap/log"
	"github.com/tikv/pd/pkg/codec"
	"github.com/tikv/pd/server/core"
	"github.com/tikv/pd/tools/pd-simulator/simulator/info"
	"github.com/tikv/pd/tools/pd-simulator/simulator/simutil"
	"go.uber.org/zap"
)

func newImportData() *Case {
	var simCase Case
	// Initialize the cluster
	for i := 1; i <= 10; i++ {
		simCase.Stores = append(simCase.Stores, &Store{
			ID:        IDAllocator.nextID(),
			Status:    metapb.StoreState_Up,
			Capacity:  1 * TB,
			Available: 900 * GB,
			Version:   "2.1.0",
		})
	}

	for i := 0; i < getRegionNum(); i++ {
		storeIDs := rand.Perm(10)
		peers := []*metapb.Peer{
			{Id: IDAllocator.nextID(), StoreId: uint64(storeIDs[0] + 1)},
			{Id: IDAllocator.nextID(), StoreId: uint64(storeIDs[1] + 1)},
			{Id: IDAllocator.nextID(), StoreId: uint64(storeIDs[2] + 1)},
		}
		simCase.Regions = append(simCase.Regions, Region{
			ID:     IDAllocator.nextID(),
			Peers:  peers,
			Leader: peers[0],
			Size:   32 * MB,
			Keys:   320000,
		})
	}

	simCase.RegionSplitSize = 64 * MB
	simCase.RegionSplitKeys = 640000
	simCase.TableNumber = 10
	// Events description
	e := &WriteFlowOnSpotDescriptor{}
	table12 := string(codec.EncodeBytes(codec.GenerateTableKey(12)))
	table13 := string(codec.EncodeBytes(codec.GenerateTableKey(13)))
	e.Step = func(tick int64) map[string]int64 {
		if tick > int64(getRegionNum())/10 {
			return nil
		}
		return map[string]int64{
			table12: 32 * MB,
		}
	}
	simCase.Events = []EventDescriptor{e}

	// Checker description
	checkCount := uint64(0)
	var newRegionCount [][3]int
	var allRegionCount [][3]int
	simCase.Checker = func(regions *core.RegionsInfo, stats []info.StoreStats) bool {
		leaderDist := make(map[uint64]int)
		peerDist := make(map[uint64]int)
		leaderTotal := 0
		peerTotal := 0
		res := make([]*core.RegionInfo, 0, 100)
		regions.ScanRangeWithIterator([]byte(table12), func(region *core.RegionInfo) bool {
			if bytes.Compare(region.GetEndKey(), []byte(table13)) < 0 {
				res = append(res, regions.GetRegion(region.GetID()))
				return true
			}
			return false
		})

		for _, r := range res {
			leaderTotal++
			leaderDist[r.GetLeader().GetStoreId()]++
			for _, p := range r.GetPeers() {
				peerDist[p.GetStoreId()]++
				peerTotal++
			}
		}
		if leaderTotal == 0 || peerTotal == 0 {
			return false
		}
		tableLeaderLog := fmt.Sprintf("%d leader:", leaderTotal)
		tablePeerLog := fmt.Sprintf("%d peer: ", peerTotal)
		for storeID := 1; storeID <= 10; storeID++ {
			if leaderCount, ok := leaderDist[uint64(storeID)]; ok {
				tableLeaderLog = fmt.Sprintf("%s [store %d]:%.2f%%", tableLeaderLog, storeID, float64(leaderCount)/float64(leaderTotal)*100)
			}
		}
		for storeID := 1; storeID <= 10; storeID++ {
			if peerCount, ok := peerDist[uint64(storeID)]; ok {
				newRegionCount = append(newRegionCount, [3]int{storeID, int(checkCount), peerCount})
				tablePeerLog = fmt.Sprintf("%s [store %d]:%.2f%%", tablePeerLog, storeID, float64(peerCount)/float64(peerTotal)*100)
			}
		}
		regionTotal := regions.GetRegionCount()
		totalLeaderLog := fmt.Sprintf("%d leader:", regionTotal)
		totalPeerLog := fmt.Sprintf("%d peer:", regionTotal*3)
		isEnd := false
		var regionProps []float64
		for storeID := uint64(1); storeID <= 10; storeID++ {
			totalLeaderLog = fmt.Sprintf("%s [store %d]:%.2f%%", totalLeaderLog, storeID, float64(regions.GetStoreLeaderCount(storeID))/float64(regionTotal)*100)
			regionProp := float64(regions.GetStoreRegionCount(storeID)) / float64(regionTotal*3) * 100
			regionProps = append(regionProps, regionProp)
			totalPeerLog = fmt.Sprintf("%s [store %d]:%.2f%%", totalPeerLog, storeID, regionProp)
			allRegionCount = append(allRegionCount, [3]int{int(storeID), int(checkCount), regions.GetStoreRegionCount(storeID)})
		}
		simutil.Logger.Info("import data information",
			zap.String("table-leader", tableLeaderLog),
			zap.String("table-peer", tablePeerLog),
			zap.String("total-leader", totalLeaderLog),
			zap.String("total-peer", totalPeerLog))
		checkCount++
		dev := 0.0
		for _, p := range regionProps {
			dev += (p - 10) * (p - 10) / 100
		}
		if dev > 0.005 {
			simutil.Logger.Warn("Not balanced, change scheduler or store limit", zap.Float64("dev score", dev))
		}
		if checkCount > uint64(getRegionNum())/10 {
			isEnd = dev < 0.002
		}
		if isEnd {
			renderPlot("new_region.html", newRegionCount, int(checkCount), 0, getRegionNum()/10)
			renderPlot("all_region.html", allRegionCount, int(checkCount), 28*getRegionNum()/100, getRegionNum()/3)
		}
		return isEnd
	}
	return &simCase
}

func renderPlot(name string, data [][3]int, len, minCount, maxCount int) {
	var rangeColor = []string{
		"#313695", "#4575b4", "#74add1", "#abd9e9", "#e0f3f8",
		"#fee090", "#fdae61", "#f46d43", "#d73027", "#a50026",
	}
	bar3d := charts.NewBar3D()
	bar3d.SetGlobalOptions(
		charts.TitleOpts{Title: "Region count"},
		charts.VisualMapOpts{
			Range:      []float32{float32(minCount), float32(maxCount)},
			Calculable: true,
			InRange:    charts.VMInRange{Color: rangeColor},
			Min:        float32(minCount),
			Max:        float32(maxCount),
		},
		charts.Grid3DOpts{BoxDepth: 100, BoxWidth: 300},
	)
	xAxis := make([]int, 10)
	for i := 1; i <= 10; i++ {
		xAxis[i-1] = i
	}
	yAxis := make([]int, len)
	for i := 1; i <= len; i++ {
		yAxis[i-1] = i
	}
	bar3d.AddXYAxis(xAxis, yAxis).AddZAxis("bar3d", data)
	f, _ := os.Create(name)
	err := bar3d.Render(f)
	if err != nil {
		log.Error("Render error", zap.Error(err))
	}
}
<commit_msg>Force stop sim (#3142)<commit_after>// Copyright 2018 TiKV Project 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package cases

import (
	"bytes"
	"fmt"
	"math/rand"
	"os"

	"github.com/go-echarts/go-echarts/charts"
	"github.com/pingcap/kvproto/pkg/metapb"
	"github.com/pingcap/log"
	"github.com/tikv/pd/pkg/codec"
	"github.com/tikv/pd/server/core"
	"github.com/tikv/pd/tools/pd-simulator/simulator/info"
	"github.com/tikv/pd/tools/pd-simulator/simulator/simutil"
	"go.uber.org/zap"
)

func newImportData() *Case {
	var simCase Case
	// Initialize the cluster
	for i := 1; i <= 10; i++ {
		simCase.Stores = append(simCase.Stores, &Store{
			ID:        IDAllocator.nextID(),
			Status:    metapb.StoreState_Up,
			Capacity:  1 * TB,
			Available: 900 * GB,
			Version:   "2.1.0",
		})
	}

	for i := 0; i < getRegionNum(); i++ {
		storeIDs := rand.Perm(10)
		peers := []*metapb.Peer{
			{Id: IDAllocator.nextID(), StoreId: uint64(storeIDs[0] + 1)},
			{Id: IDAllocator.nextID(), StoreId: uint64(storeIDs[1] + 1)},
			{Id: IDAllocator.nextID(), StoreId: uint64(storeIDs[2] + 1)},
		}
		simCase.Regions = append(simCase.Regions, Region{
			ID:     IDAllocator.nextID(),
			Peers:  peers,
			Leader: peers[0],
			Size:   32 * MB,
			Keys:   320000,
		})
	}

	simCase.RegionSplitSize = 64 * MB
	simCase.RegionSplitKeys = 640000
	simCase.TableNumber = 10
	// Events description
	e := &WriteFlowOnSpotDescriptor{}
	table12 := string(codec.EncodeBytes(codec.GenerateTableKey(12)))
	table13 := string(codec.EncodeBytes(codec.GenerateTableKey(13)))
	e.Step = func(tick int64) map[string]int64 {
		if tick > int64(getRegionNum())/10 {
			return nil
		}
		return map[string]int64{
			table12: 32 * MB,
		}
	}
	simCase.Events = []EventDescriptor{e}

	// Checker description
	checkCount := uint64(0)
	var newRegionCount [][3]int
	var allRegionCount [][3]int
	simCase.Checker = func(regions *core.RegionsInfo, stats []info.StoreStats) bool {
		leaderDist := make(map[uint64]int)
		peerDist := make(map[uint64]int)
		leaderTotal := 0
		peerTotal := 0
		res := make([]*core.RegionInfo, 0, 100)
		regions.ScanRangeWithIterator([]byte(table12), func(region *core.RegionInfo) bool {
			if bytes.Compare(region.GetEndKey(), []byte(table13)) < 0 {
				res = append(res, regions.GetRegion(region.GetID()))
				return true
			}
			return false
		})

		for _, r := range res {
			leaderTotal++
			leaderDist[r.GetLeader().GetStoreId()]++
			for _, p := range r.GetPeers() {
				peerDist[p.GetStoreId()]++
				peerTotal++
			}
		}
		if leaderTotal == 0 || peerTotal == 0 {
			return false
		}
		tableLeaderLog := fmt.Sprintf("%d leader:", leaderTotal)
		tablePeerLog := fmt.Sprintf("%d peer: ", peerTotal)
		for storeID := 1; storeID <= 10; storeID++ {
			if leaderCount, ok := leaderDist[uint64(storeID)]; ok {
				tableLeaderLog = fmt.Sprintf("%s [store %d]:%.2f%%", tableLeaderLog, storeID, float64(leaderCount)/float64(leaderTotal)*100)
			}
		}
		for storeID := 1; storeID <= 10; storeID++ {
			if peerCount, ok := peerDist[uint64(storeID)]; ok {
				newRegionCount = append(newRegionCount, [3]int{storeID, int(checkCount), peerCount})
				tablePeerLog = fmt.Sprintf("%s [store %d]:%.2f%%", tablePeerLog, storeID, float64(peerCount)/float64(peerTotal)*100)
			}
		}
		regionTotal := regions.GetRegionCount()
		totalLeaderLog := fmt.Sprintf("%d leader:", regionTotal)
		totalPeerLog := fmt.Sprintf("%d peer:", regionTotal*3)
		isEnd := false
		var regionProps []float64
		for storeID := uint64(1); storeID <= 10; storeID++ {
			totalLeaderLog = fmt.Sprintf("%s [store %d]:%.2f%%", totalLeaderLog, storeID, float64(regions.GetStoreLeaderCount(storeID))/float64(regionTotal)*100)
			regionProp := float64(regions.GetStoreRegionCount(storeID)) / float64(regionTotal*3) * 100
			regionProps = append(regionProps, regionProp)
			totalPeerLog = fmt.Sprintf("%s [store %d]:%.2f%%", totalPeerLog, storeID, regionProp)
			allRegionCount = append(allRegionCount, [3]int{int(storeID), int(checkCount), regions.GetStoreRegionCount(storeID)})
		}
		simutil.Logger.Info("import data information",
			zap.String("table-leader", tableLeaderLog),
			zap.String("table-peer", tablePeerLog),
			zap.String("total-leader", totalLeaderLog),
			zap.String("total-peer", totalPeerLog))
		checkCount++
		dev := 0.0
		for _, p := range regionProps {
			dev += (p - 10) * (p - 10) / 100
		}
		if dev > 0.02 {
			simutil.Logger.Warn("Not balanced, change scheduler or store limit", zap.Float64("dev score", dev))
		}
		if checkCount > uint64(getRegionNum())/5 {
			isEnd = true
		} else if checkCount > uint64(getRegionNum())/10 {
			isEnd = dev < 0.01
		}
		if isEnd {
			renderPlot("new_region.html", newRegionCount, int(checkCount), 0, getRegionNum()/10)
			renderPlot("all_region.html", allRegionCount, int(checkCount), 28*getRegionNum()/100, getRegionNum()/3)
		}
		return isEnd
	}
	return &simCase
}

func renderPlot(name string, data [][3]int, len, minCount, maxCount int) {
	var rangeColor = []string{
		"#313695", "#4575b4", "#74add1", "#abd9e9", "#e0f3f8",
		"#fee090", "#fdae61", "#f46d43", "#d73027", "#a50026",
	}
	bar3d := charts.NewBar3D()
	bar3d.SetGlobalOptions(
		charts.TitleOpts{Title: "Region count"},
		charts.VisualMapOpts{
			Range:      []float32{float32(minCount), float32(maxCount)},
			Calculable: true,
			InRange:    charts.VMInRange{Color: rangeColor},
			Min:        float32(minCount),
			Max:        float32(maxCount),
		},
		charts.Grid3DOpts{BoxDepth: 100, BoxWidth: 300},
	)
	xAxis := make([]int, 10)
	for i := 1; i <= 10; i++ {
		xAxis[i-1] = i
	}
	yAxis := make([]int, len)
	for i := 1; i <= len; i++ {
		yAxis[i-1] = i
	}
	bar3d.AddXYAxis(xAxis, yAxis).AddZAxis("bar3d", data)
	f, _ := os.Create(name)
	err := bar3d.Render(f)
	if err != nil {
		log.Error("Render error", zap.Error(err))
	}
}
<|endoftext|>


<commit_before>package infodownloader

import (
	"bytes"
	"fmt"

	"github.com/cenkalti/rain/torrent/internal/peerconn"
	"github.com/cenkalti/rain/torrent/internal/peerprotocol"
)

const blockSize = 16 * 1024

const maxQueuedBlocks = 10

// InfoDownloader downloads all blocks of a piece from a peer.
type InfoDownloader struct {
	extID          uint8
	totalSize      uint32
	nextBlockIndex uint32
	requested      map[uint32]struct{}
	Peer           *peerconn.Conn
	blocks         []block
	DataC          chan Data
	// RejectC chan *piece.Block
	resultC chan Result
	closeC  chan struct{}
}

type Data struct {
	Index uint32
	Data  []byte
}

type block struct {
	index uint32
	size  uint32
	data  []byte
}

type Result struct {
	Peer  *peerconn.Conn
	Bytes []byte
	Error error
}

func New(pe *peerconn.Conn, extID uint8, totalSize uint32, resultC chan Result) *InfoDownloader {
	numBlocks := totalSize / blockSize
	mod := totalSize % blockSize
	if mod != 0 {
		numBlocks++
	}
	blocks := make([]block, numBlocks)
	for i := range blocks {
		blocks[i] = block{
			index: uint32(i),
			size:  blockSize,
		}
	}
	if mod != 0 && len(blocks) > 0 {
		blocks[len(blocks)-1].size = mod
	}
	return &InfoDownloader{
		extID:     extID,
		totalSize: totalSize,
		Peer:      pe,
		blocks:    blocks,
		requested: make(map[uint32]struct{}),
		DataC:     make(chan Data),
		// RejectC: make(chan *piece.Block),
		resultC: resultC,
		closeC:  make(chan struct{}),
	}
}

func (d *InfoDownloader) Close() {
	close(d.closeC)
}

func (d *InfoDownloader) requestBlocks() {
	for ; d.nextBlockIndex < uint32(len(d.blocks)) && len(d.requested) < maxQueuedBlocks; d.nextBlockIndex++ {
		d.requested[d.nextBlockIndex] = struct{}{}
		msg := peerprotocol.ExtensionMessage{
			ExtendedMessageID: d.extID,
			Payload: peerprotocol.ExtensionMetadataMessage{
				Type:  peerprotocol.ExtensionMetadataMessageTypeRequest,
				Piece: d.nextBlockIndex,
			},
		}
		d.Peer.SendMessage(msg)
	}
}

func (d *InfoDownloader) Run() {
	result := Result{
		Peer: d.Peer,
	}
	defer func() {
		select {
		case d.resultC <- result:
		case <-d.closeC:
		}
	}()
	for {
		d.requestBlocks()
		select {
		case msg := <-d.DataC:
			if _, ok := d.requested[msg.Index]; !ok {
				result.Error = fmt.Errorf("peer sent unrequested index for metadata message: %q", msg.Index)
				return
			}
			b := &d.blocks[msg.Index]
			if uint32(len(msg.Data)) != b.size {
				result.Error = fmt.Errorf("peer sent invalid size for metadata message: %q", len(msg.Data))
				return
			}
			delete(d.requested, msg.Index)
			b.data = msg.Data
			if d.allDone() {
				result.Bytes = d.assembleBlocks().Bytes()
				return
			}
		// TODO handle rejects in info downloader
		// case blk := <-d.RejectC:
		// 	b := d.blocks[blk.Index]
		// 	if !b.requested {
		// 		d.Peer.Close()
		// 		d.ErrC <- errors.New("received invalid reject message")
		// 		return
		// 	}
		// 	d.blocks[blk.Index].requested = false
		case <-d.closeC:
			return
		}
	}
}

func (d *InfoDownloader) allDone() bool {
	return d.nextBlockIndex == uint32(len(d.blocks)) && len(d.requested) == 0
}

func (d *InfoDownloader) assembleBlocks() *bytes.Buffer {
	buf := bytes.NewBuffer(make([]byte, 0, d.totalSize))
	for i := range d.blocks {
		buf.Write(d.blocks[i].data)
	}
	return buf
}
<commit_msg>Revert "remove semaphore usage"<commit_after>package infodownloader

import (
	"bytes"
	"fmt"

	"github.com/cenkalti/rain/torrent/internal/peerconn"
	"github.com/cenkalti/rain/torrent/internal/peerprotocol"
)

const blockSize = 16 * 1024

const maxQueuedBlocks = 10

// InfoDownloader downloads all blocks of a piece from a peer.
type InfoDownloader struct {
	extID     uint8
	totalSize uint32
	limiter   chan struct{}
	Peer      *peerconn.Conn
	blocks    []block
	DataC     chan Data
	// RejectC chan *piece.Block
	resultC chan Result
	closeC  chan struct{}
}

type Data struct {
	Index uint32
	Data  []byte
}

type block struct {
	index     uint32
	size      uint32
	requested bool
	data      []byte
}

type Result struct {
	Peer  *peerconn.Conn
	Bytes []byte
	Error error
}

func New(pe *peerconn.Conn, extID uint8, totalSize uint32, resultC chan Result) *InfoDownloader {
	numBlocks := totalSize / blockSize
	mod := totalSize % blockSize
	if mod != 0 {
		numBlocks++
	}
	blocks := make([]block, numBlocks)
	for i := range blocks {
		blocks[i] = block{
			index: uint32(i),
			size:  blockSize,
		}
	}
	if mod != 0 && len(blocks) > 0 {
		blocks[len(blocks)-1].size = mod
	}
	return &InfoDownloader{
		extID:     extID,
		totalSize: totalSize,
		Peer:      pe,
		blocks:    blocks,
		limiter:   make(chan struct{}, maxQueuedBlocks),
		DataC:     make(chan Data),
		// RejectC: make(chan *piece.Block),
		resultC: resultC,
		closeC:  make(chan struct{}),
	}
}

func (d *InfoDownloader) Close() {
	close(d.closeC)
}

func (d *InfoDownloader) Run() {
	result := Result{
		Peer: d.Peer,
	}
	defer func() {
		select {
		case d.resultC <- result:
		case <-d.closeC:
		}
	}()
	for {
		select {
		case d.limiter <- struct{}{}:
			b := d.nextBlock()
			if b == nil {
				d.limiter = nil
				break
			}
			msg := peerprotocol.ExtensionMessage{
				ExtendedMessageID: d.extID,
				Payload: peerprotocol.ExtensionMetadataMessage{
					Type:  peerprotocol.ExtensionMetadataMessageTypeRequest,
					Piece: b.index,
				},
			}
			d.Peer.SendMessage(msg)
		case msg := <-d.DataC:
			if msg.Index >= uint32(len(d.blocks)) {
				result.Error = fmt.Errorf("peer sent invalid index for metadata message: %q", msg.Index)
				return
			}
			b := &d.blocks[msg.Index]
			if uint32(len(msg.Data)) != b.size {
				result.Error = fmt.Errorf("peer sent invalid size for metadata message: %q", len(msg.Data))
				return
			}
			if b.requested && b.data == nil {
				if d.limiter == nil {
					d.limiter = make(chan struct{}, maxQueuedBlocks)
				}
			}
			b.data = msg.Data
			if d.allDone() {
				result.Bytes = d.assembleBlocks().Bytes()
				return
			}
		// TODO handle rejects in info downloader
		// case blk := <-d.RejectC:
		// 	b := d.blocks[blk.Index]
		// 	if !b.requested {
		// 		d.Peer.Close()
		// 		d.ErrC <- errors.New("received invalid reject message")
		// 		return
		// 	}
		// 	d.blocks[blk.Index].requested = false
		case <-d.closeC:
			return
		}
	}
}

func (d *InfoDownloader) nextBlock() *block {
	for i := range d.blocks {
		if !d.blocks[i].requested {
			d.blocks[i].requested = true
			return &d.blocks[i]
		}
	}
	return nil
}

func (d *InfoDownloader) allDone() bool {
	for i := range d.blocks {
		if d.blocks[i].data == nil {
			return false
		}
	}
	return true
}

func (d *InfoDownloader) assembleBlocks() *bytes.Buffer {
	buf := bytes.NewBuffer(make([]byte, 0, d.totalSize))
	for i := range d.blocks {
		buf.Write(d.blocks[i].data)
	}
	return buf
}
<|endoftext|>


<commit_before>/*

   Copyright 2016 Wenhui Shen <www.webx.top>

   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 captcha

import (
	"strconv"
	"strings"

	"github.com/webx-top/captcha"
	"github.com/webx-top/echo"
)

type Options struct {
	EnableImage bool
	EnableAudio bool
	AudioLangs  []string
	Prefix      string
}

func (o Options) Wrapper(e echo.RouteRegister) {
	if o.AudioLangs == nil || len(o.AudioLangs) == 0 {
		o.AudioLangs = []string{`zh`, `ru`, `en`}
	}
	if len(o.Prefix) == 0 {
		o.Prefix = `captcha`
	} else {
		o.Prefix = strings.TrimRight(o.Prefix, "/")
	}
	e.Get(o.Prefix+"/*", func(ctx echo.Context) error {
		var id, ext string
		param := ctx.P(0)
		if p := strings.LastIndex(param, `.`); p > 0 {
			id = param[0:p]
			ext = param[p:]
		}
		if len(ext) == 0 || len(id) == 0 {
			ctx.Response().NotFound()
			return nil
		}
		if len(ctx.Query("reload")) > 0 {
			captcha.Reload(id)
		}
		w := ctx.Response()
		header := w.Header()
		download := len(ctx.Query("download")) > 0

		switch ext {
		case ".png":
			if !o.EnableImage {
				ctx.Response().NotFound()
				return nil
			}
			if download {
				header.Set(echo.HeaderContentType, "application/octet-stream")
			} else {
				header.Set(echo.HeaderContentType, "image/png")
			}
			return captcha.WriteImage(w.Writer(), id, captcha.StdWidth, captcha.StdHeight)
		case ".wav":
			if !o.EnableAudio {
				ctx.Response().NotFound()
				return nil
			}
			lang := strings.ToLower(ctx.Query("lang"))
			supported := false
			for _, supportedLang := range o.AudioLangs {
				if supportedLang == lang {
					supported = true
					break
				}
			}
			if !supported && len(o.AudioLangs) > 0 {
				lang = o.AudioLangs[0]
			}
			au, err := captcha.GetAudio(id, lang)
			if err != nil {
				return err
			}
			if download {
				header.Set(echo.HeaderContentType, "application/octet-stream")
			} else {
				header.Set(echo.HeaderContentType, "audio/x-wav")
			}
			header.Set("Content-Length", strconv.Itoa(au.EncodedLen()))
			_, err = au.WriteTo(w.Writer())
			return err
		}
		return nil
	})
}
<commit_msg>update<commit_after>/*

   Copyright 2016 Wenhui Shen <www.webx.top>

   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 captcha

import (
	"strconv"
	"strings"

	"github.com/webx-top/captcha"
	"github.com/webx-top/echo"
)

type Options struct {
	EnableImage bool
	EnableAudio bool
	AudioLangs  []string
	Prefix      string
}

func (o Options) Wrapper(e echo.RouteRegister) {
	if o.AudioLangs == nil || len(o.AudioLangs) == 0 {
		o.AudioLangs = []string{`zh`, `ru`, `en`}
	}
	if len(o.Prefix) == 0 {
		o.Prefix = `/captcha`
	} else {
		o.Prefix = strings.TrimRight(o.Prefix, "/")
	}
	e.Get(o.Prefix+"/*", func(ctx echo.Context) error {
		var id, ext string
		param := ctx.P(0)
		if p := strings.LastIndex(param, `.`); p > 0 {
			id = param[0:p]
			ext = param[p:]
		}
		if len(ext) == 0 || len(id) == 0 {
			ctx.Response().NotFound()
			return nil
		}
		if len(ctx.Query("reload")) > 0 {
			captcha.Reload(id)
		}
		w := ctx.Response()
		header := w.Header()
		download := len(ctx.Query("download")) > 0

		switch ext {
		case ".png":
			if !o.EnableImage {
				ctx.Response().NotFound()
				return nil
			}
			if download {
				header.Set(echo.HeaderContentType, "application/octet-stream")
			} else {
				header.Set(echo.HeaderContentType, "image/png")
			}
			return captcha.WriteImage(w.Writer(), id, captcha.StdWidth, captcha.StdHeight)
		case ".wav":
			if !o.EnableAudio {
				ctx.Response().NotFound()
				return nil
			}
			lang := strings.ToLower(ctx.Query("lang"))
			supported := false
			for _, supportedLang := range o.AudioLangs {
				if supportedLang == lang {
					supported = true
					break
				}
			}
			if !supported && len(o.AudioLangs) > 0 {
				lang = o.AudioLangs[0]
			}
			au, err := captcha.GetAudio(id, lang)
			if err != nil {
				return err
			}
			if download {
				header.Set(echo.HeaderContentType, "application/octet-stream")
			} else {
				header.Set(echo.HeaderContentType, "audio/x-wav")
			}
			header.Set("Content-Length", strconv.Itoa(au.EncodedLen()))
			_, err = au.WriteTo(w.Writer())
			return err
		}
		return nil
	})
}
<|endoftext|>


<commit_before>package harvester

import (
	"github.com/cloudfoundry-incubator/executor"
	"github.com/cloudfoundry-incubator/rep"
	"github.com/cloudfoundry-incubator/runtime-schema/bbs"
	"github.com/cloudfoundry-incubator/runtime-schema/bbs/bbserrors"
	"github.com/pivotal-golang/lager"
)

type lrpProcessor struct {
	cellID         string
	logger         lager.Logger
	bbs            bbs.RepBBS
	executorClient executor.Client
}

func NewLRPProcessor(
	cellID string,
	logger lager.Logger,
	bbs bbs.RepBBS,
	executorClient executor.Client,
) Processor {
	return &lrpProcessor{
		cellID:         cellID,
		logger:         logger,
		bbs:            bbs,
		executorClient: executorClient,
	}
}

func (p *lrpProcessor) Process(container executor.Container) {
	logger := p.logger.Session("process-lrp", lager.Data{
		"container-guid":  container.Guid,
		"container-state": container.State,
	})

	if container.State == executor.StateReserved {
		logger.Debug("ignoring-reserved-container")
		return
	}

	lrpKey, err := rep.ActualLRPKeyFromContainer(container)
	if err != nil {
		logger.Error("failed-to-generate-lrp-key-from-container", err)
		return
	}

	lrpContainerKey, err := rep.ActualLRPContainerKeyFromContainer(container, p.cellID)
	if err != nil {
		logger.Error("failed-to-generate-container-key-from-container", err)
		return
	}

	lrpNetInfo, err := rep.ActualLRPNetInfoFromContainer(container)
	if err != nil {
		logger.Error("failed-to-generate-net-info-from-container", err)
		return
	}

	switch container.State {
	case executor.StateInitializing, executor.StateCreated:
		err = p.bbs.ClaimActualLRP(lrpKey, lrpContainerKey, logger)
		if err == bbserrors.ErrActualLRPCannotBeClaimed {
			logger.Debug("failed-to-claim-actual-lrp")
			p.deleteContainer(lrpContainerKey.InstanceGuid, logger)
			return
		}
		if err != nil {
			logger.Error("failed-to-claim-actual-lrp", err)
			return
		}
		logger.Debug("claimed-actual-lrp")

	case executor.StateRunning:
		err := p.bbs.StartActualLRP(lrpKey, lrpContainerKey, lrpNetInfo, logger)
		if err == bbserrors.ErrActualLRPCannotBeStarted {
			logger.Debug("failed-to-start-actual-lrp")
			p.deleteContainer(lrpContainerKey.InstanceGuid, logger)
			return
		}
		if err != nil {
			logger.Error("failed-to-start-actual-lrp", err)
			return
		}
		logger.Debug("started-actual-lrp")

	case executor.StateCompleted:
		err := p.bbs.RemoveActualLRP(lrpKey, lrpContainerKey, logger)
		if err != nil {
			logger.Error("failed-to-remove-actual-lrp", err)
		} else {
			logger.Info("removed-actual-lrp")
		}

		p.deleteContainer(lrpContainerKey.InstanceGuid, logger)
	default:
		logger.Debug("unknown-container-state")
	}
}

func (p *lrpProcessor) deleteContainer(guid string, logger lager.Logger) {
	err := p.executorClient.DeleteContainer(guid)
	if err != nil {
		logger.Error("failed-to-delete-container", err)
	} else {
		logger.Info("deleted-container")
	}
}
<commit_msg>only get net info if actuallrp is running<commit_after>package harvester

import (
	"github.com/cloudfoundry-incubator/executor"
	"github.com/cloudfoundry-incubator/rep"
	"github.com/cloudfoundry-incubator/runtime-schema/bbs"
	"github.com/cloudfoundry-incubator/runtime-schema/bbs/bbserrors"
	"github.com/pivotal-golang/lager"
)

type lrpProcessor struct {
	cellID         string
	logger         lager.Logger
	bbs            bbs.RepBBS
	executorClient executor.Client
}

func NewLRPProcessor(
	cellID string,
	logger lager.Logger,
	bbs bbs.RepBBS,
	executorClient executor.Client,
) Processor {
	return &lrpProcessor{
		cellID:         cellID,
		logger:         logger,
		bbs:            bbs,
		executorClient: executorClient,
	}
}

func (p *lrpProcessor) Process(container executor.Container) {
	logger := p.logger.Session("process-lrp", lager.Data{
		"container-guid":  container.Guid,
		"container-state": container.State,
	})

	if container.State == executor.StateReserved {
		logger.Debug("ignoring-reserved-container")
		return
	}

	lrpKey, err := rep.ActualLRPKeyFromContainer(container)
	if err != nil {
		logger.Error("failed-to-generate-lrp-key-from-container", err)
		return
	}

	lrpContainerKey, err := rep.ActualLRPContainerKeyFromContainer(container, p.cellID)
	if err != nil {
		logger.Error("failed-to-generate-container-key-from-container", err)
		return
	}

	switch container.State {
	case executor.StateInitializing, executor.StateCreated:
		err = p.bbs.ClaimActualLRP(lrpKey, lrpContainerKey, logger)
		if err == bbserrors.ErrActualLRPCannotBeClaimed {
			logger.Debug("failed-to-claim-actual-lrp")
			p.deleteContainer(lrpContainerKey.InstanceGuid, logger)
			return
		}
		if err != nil {
			logger.Error("failed-to-claim-actual-lrp", err)
			return
		}
		logger.Debug("claimed-actual-lrp")

	case executor.StateRunning:
		lrpNetInfo, err := rep.ActualLRPNetInfoFromContainer(container)
		if err != nil {
			logger.Error("failed-to-generate-net-info-from-container", err)
			return
		}

		err = p.bbs.StartActualLRP(lrpKey, lrpContainerKey, lrpNetInfo, logger)
		if err == bbserrors.ErrActualLRPCannotBeStarted {
			logger.Debug("failed-to-start-actual-lrp")
			p.deleteContainer(lrpContainerKey.InstanceGuid, logger)
			return
		}
		if err != nil {
			logger.Error("failed-to-start-actual-lrp", err)
			return
		}
		logger.Debug("started-actual-lrp")

	case executor.StateCompleted:
		err := p.bbs.RemoveActualLRP(lrpKey, lrpContainerKey, logger)
		if err != nil {
			logger.Error("failed-to-remove-actual-lrp", err)
		} else {
			logger.Info("removed-actual-lrp")
		}

		p.deleteContainer(lrpContainerKey.InstanceGuid, logger)
	default:
		logger.Debug("unknown-container-state")
	}
}

func (p *lrpProcessor) deleteContainer(guid string, logger lager.Logger) {
	err := p.executorClient.DeleteContainer(guid)
	if err != nil {
		logger.Error("failed-to-delete-container", err)
	} else {
		logger.Info("deleted-container")
	}
}
<|endoftext|>


<commit_before>// Copyright 2019 Google LLC
//
// 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
//
//     https://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 snippets

// [START healthcare_export_dicom_instance]
import (
	"context"
	"fmt"
	"io"

	healthcare "google.golang.org/api/healthcare/v1beta1"
)

// exportDICOMInstance exports DICOM objects to GCS.
func exportDICOMInstance(w io.Writer, projectID, location, datasetID, dicomStoreID, destination string) error {
	ctx := context.Background()

	healthcareService, err := healthcare.NewService(ctx)
	if err != nil {
		return fmt.Errorf("healthcare.NewService: %v", err)
	}

	storesService := healthcareService.Projects.Locations.Datasets.DicomStores

	req := &healthcare.ExportDicomDataRequest{
		GcsDestination: &healthcare.GoogleCloudHealthcareV1beta1DicomGcsDestination{
			UriPrefix: destination, // "gs://my-bucket/path/to/prefix/"
		},
	}
	name := fmt.Sprintf("projects/%s/locations/%s/datasets/%s/dicomStores/%s", projectID, location, datasetID, dicomStoreID)

	lro, err := storesService.Export(name, req).Do()
	if err != nil {
		return fmt.Errorf("Export: %v", err)
	}

	fmt.Fprintf(w, "Export to DICOM store started. Operation: %q\n", lro.Name)
	return nil
}

// [END healthcare_export_dicom_instance]
<commit_msg>healthcare: update region tag (#892)<commit_after>// Copyright 2019 Google LLC
//
// 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
//
//     https://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 snippets

// [START healthcare_export_dicom_instance_gcs]
import (
	"context"
	"fmt"
	"io"

	healthcare "google.golang.org/api/healthcare/v1beta1"
)

// exportDICOMInstance exports DICOM objects to GCS.
func exportDICOMInstance(w io.Writer, projectID, location, datasetID, dicomStoreID, destination string) error {
	ctx := context.Background()

	healthcareService, err := healthcare.NewService(ctx)
	if err != nil {
		return fmt.Errorf("healthcare.NewService: %v", err)
	}

	storesService := healthcareService.Projects.Locations.Datasets.DicomStores

	req := &healthcare.ExportDicomDataRequest{
		GcsDestination: &healthcare.GoogleCloudHealthcareV1beta1DicomGcsDestination{
			UriPrefix: destination, // "gs://my-bucket/path/to/prefix/"
		},
	}
	name := fmt.Sprintf("projects/%s/locations/%s/datasets/%s/dicomStores/%s", projectID, location, datasetID, dicomStoreID)

	lro, err := storesService.Export(name, req).Do()
	if err != nil {
		return fmt.Errorf("Export: %v", err)
	}

	fmt.Fprintf(w, "Export to DICOM store started. Operation: %q\n", lro.Name)
	return nil
}

// [END healthcare_export_dicom_instance_gcs]
<|endoftext|>