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

package mongostore

import (
	"errors"
	"github.com/gorilla/securecookie"
	"github.com/gorilla/sessions"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
	"net/http"
	"time"
)

var (
	ErrInvalidId = errors.New("mgostore: invalid session id")
)

// Session object store in MongoDB
type Session struct {
	Id       bson.ObjectId `bson:"_id,omitempty"`
	Data     string
	Modified time.Time
}

// MongoStore stores sessions in MongoDB
type MongoStore struct {
	Codecs  []securecookie.Codec
	Options *sessions.Options
	Token   TokenGetSeter
	coll    *mgo.Collection
}

// NewMongoStore returns a new MongoStore.
// Set ensureTTL to true let the database auto-remove expired object by maxAge.
func NewMongoStore(c *mgo.Collection, maxAge int, ensureTTL bool,
	keyPairs ...[]byte) *MongoStore {
	store := &MongoStore{
		Codecs: securecookie.CodecsFromPairs(keyPairs...),
		Options: &sessions.Options{
			Path: "/",
			MaxAge: maxAge,
		},
		Token: &CookieToken{},
		coll:  c,
	}

	if ensureTTL {
		c.EnsureIndex(mgo.Index{
			Key:         []string{"modified"},
			Background:  true,
			Sparse:      true,
			ExpireAfter: time.Duration(maxAge) * time.Second,
		})
	}

	return store
}

// Get registers and returns a session for the given name and session store.
// It returns a new session if there are no sessions registered for the name.
func (m *MongoStore) Get(r *http.Request, name string) (
	*sessions.Session, error) {
	return sessions.GetRegistry(r).Get(m, name)
}

// New returns a session for the given name without adding it to the registry.
func (m *MongoStore) New(r *http.Request, name string) (
	*sessions.Session, error) {
	session := sessions.NewSession(m, name)
	session.Options = &sessions.Options{
		Path:   m.Options.Path,
		MaxAge: m.Options.MaxAge,
	}
	session.IsNew = true
	var err error
	if cook, errToken := m.Token.GetToken(r, name); errToken == nil {
		err = securecookie.DecodeMulti(name, cook, &session.ID, m.Codecs...)
		if err == nil {
			err = m.load(session)
			if err == nil {
				session.IsNew = false
			} else {
				err = nil
			}
		}
	}
	return session, err
}

// Save saves all sessions registered for the current request.
func (m *MongoStore) Save(r *http.Request, w http.ResponseWriter,
	session *sessions.Session) error {
	if session.Options.MaxAge < 0 {
		if err := m.delete(session); err != nil {
			return err
		}
		m.Token.SetToken(w, session.Name(), "", session.Options)
		return nil
	}

	if session.ID == "" {
		session.ID = bson.NewObjectId().Hex()
	}

	if err := m.upsert(session); err != nil {
		return err
	}

	encoded, err := securecookie.EncodeMulti(session.Name(), session.ID,
		m.Codecs...)
	if err != nil {
		return err
	}

	m.Token.SetToken(w, session.Name(), encoded, session.Options)
	return nil
}

func (m *MongoStore) load(session *sessions.Session) error {
	if !bson.IsObjectIdHex(session.ID) {
		return ErrInvalidId
	}

	s := Session{}
	err := m.coll.FindId(bson.ObjectIdHex(session.ID)).One(&s)
	if err != nil {
		return err
	}

	if err := securecookie.DecodeMulti(session.Name(), s.Data, &session.Values,
		m.Codecs...); err != nil {
		return err
	}

	return nil
}

func (m *MongoStore) upsert(session *sessions.Session) error {
	if !bson.IsObjectIdHex(session.ID) {
		return ErrInvalidId
	}

	var modified time.Time
	if val, ok := session.Values["modified"]; ok {
		modified, ok = val.(time.Time)
		if !ok {
			return errors.New("mongostore: invalid modified value")
		}
	} else {
		modified = time.Now()
	}

	encoded, err := securecookie.EncodeMulti(session.Name(), session.Values,
		m.Codecs...)
	if err != nil {
		return err
	}

	s := Session{
		Id:       bson.ObjectIdHex(session.ID),
		Data:     encoded,
		Modified: modified,
	}

	_, err = m.coll.UpsertId(s.Id, &s)
	if err != nil {
		return err
	}

	return nil
}

func (m *MongoStore) delete(session *sessions.Session) error {
	if !bson.IsObjectIdHex(session.ID) {
		return ErrInvalidId
	}

	return m.coll.RemoveId(bson.ObjectIdHex(session.ID))
}
<commit_msg>Set MaxAge for underlying securecookie on creation.<commit_after>// Copyright 2012 The KidStuff Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package mongostore

import (
	"errors"
	"net/http"
	"time"

	"github.com/gorilla/securecookie"
	"github.com/gorilla/sessions"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

var (
	ErrInvalidId = errors.New("mgostore: invalid session id")
)

// Session object store in MongoDB
type Session struct {
	Id       bson.ObjectId `bson:"_id,omitempty"`
	Data     string
	Modified time.Time
}

// MongoStore stores sessions in MongoDB
type MongoStore struct {
	Codecs  []securecookie.Codec
	Options *sessions.Options
	Token   TokenGetSeter
	coll    *mgo.Collection
}

// NewMongoStore returns a new MongoStore.
// Set ensureTTL to true let the database auto-remove expired object by maxAge.
func NewMongoStore(c *mgo.Collection, maxAge int, ensureTTL bool,
	keyPairs ...[]byte) *MongoStore {
	store := &MongoStore{
		Codecs: securecookie.CodecsFromPairs(keyPairs...),
		Options: &sessions.Options{
			Path:   "/",
			MaxAge: maxAge,
		},
		Token: &CookieToken{},
		coll:  c,
	}

	store.MaxAge(maxAge)

	if ensureTTL {
		c.EnsureIndex(mgo.Index{
			Key:         []string{"modified"},
			Background:  true,
			Sparse:      true,
			ExpireAfter: time.Duration(maxAge) * time.Second,
		})
	}

	return store
}

// Get registers and returns a session for the given name and session store.
// It returns a new session if there are no sessions registered for the name.
func (m *MongoStore) Get(r *http.Request, name string) (
	*sessions.Session, error) {
	return sessions.GetRegistry(r).Get(m, name)
}

// New returns a session for the given name without adding it to the registry.
func (m *MongoStore) New(r *http.Request, name string) (
	*sessions.Session, error) {
	session := sessions.NewSession(m, name)
	session.Options = &sessions.Options{
		Path:   m.Options.Path,
		MaxAge: m.Options.MaxAge,
	}
	session.IsNew = true
	var err error
	if cook, errToken := m.Token.GetToken(r, name); errToken == nil {
		err = securecookie.DecodeMulti(name, cook, &session.ID, m.Codecs...)
		if err == nil {
			err = m.load(session)
			if err == nil {
				session.IsNew = false
			} else {
				err = nil
			}
		}
	}
	return session, err
}

// Save saves all sessions registered for the current request.
func (m *MongoStore) Save(r *http.Request, w http.ResponseWriter,
	session *sessions.Session) error {
	if session.Options.MaxAge < 0 {
		if err := m.delete(session); err != nil {
			return err
		}
		m.Token.SetToken(w, session.Name(), "", session.Options)
		return nil
	}

	if session.ID == "" {
		session.ID = bson.NewObjectId().Hex()
	}

	if err := m.upsert(session); err != nil {
		return err
	}

	encoded, err := securecookie.EncodeMulti(session.Name(), session.ID,
		m.Codecs...)
	if err != nil {
		return err
	}

	m.Token.SetToken(w, session.Name(), encoded, session.Options)
	return nil
}

// MaxAge sets the maximum age for the store and the underlying cookie
// implementation. Individual sessions can be deleted by setting Options.MaxAge
// = -1 for that session.
func (m *MongoStore) MaxAge(age int) {
	m.Options.MaxAge = age

	// Set the maxAge for each securecookie instance.
	for _, codec := range m.Codecs {
		if sc, ok := codec.(*securecookie.SecureCookie); ok {
			sc.MaxAge(age)
		}
	}
}

func (m *MongoStore) load(session *sessions.Session) error {
	if !bson.IsObjectIdHex(session.ID) {
		return ErrInvalidId
	}

	s := Session{}
	err := m.coll.FindId(bson.ObjectIdHex(session.ID)).One(&s)
	if err != nil {
		return err
	}

	if err := securecookie.DecodeMulti(session.Name(), s.Data, &session.Values,
		m.Codecs...); err != nil {
		return err
	}

	return nil
}

func (m *MongoStore) upsert(session *sessions.Session) error {
	if !bson.IsObjectIdHex(session.ID) {
		return ErrInvalidId
	}

	var modified time.Time
	if val, ok := session.Values["modified"]; ok {
		modified, ok = val.(time.Time)
		if !ok {
			return errors.New("mongostore: invalid modified value")
		}
	} else {
		modified = time.Now()
	}

	encoded, err := securecookie.EncodeMulti(session.Name(), session.Values,
		m.Codecs...)
	if err != nil {
		return err
	}

	s := Session{
		Id:       bson.ObjectIdHex(session.ID),
		Data:     encoded,
		Modified: modified,
	}

	_, err = m.coll.UpsertId(s.Id, &s)
	if err != nil {
		return err
	}

	return nil
}

func (m *MongoStore) delete(session *sessions.Session) error {
	if !bson.IsObjectIdHex(session.ID) {
		return ErrInvalidId
	}

	return m.coll.RemoveId(bson.ObjectIdHex(session.ID))
}
<|endoftext|>


<commit_before>package gocube

// Phase1Heuristic stores the data needed to effectively prune the search for a
// solution for phase-1.
type Phase1Heuristic struct {
	// This stores the number of moves needed to orient the corners. All the
	// fields in this array should be filled in.
	CO [2187]int8

	// This stores the number of moves needed to solve a given EO + E Slice
	// case. The fields are combined using slice*2048 + eo. If a value in this
	// array is -1, it should be interpreted as 8. This is because depths
	// greater than 7 will not be searched if the index was made under time
	// pressure.
	EOSlice [1013760]int8
}

// NewPhase1Heuristic generates a heuristic for the phase-1 solver.
// The complete argument specifies if the entire EOSlice table should be
// generated.
func NewPhase1Heuristic(moves *Phase1Moves, complete bool) *Phase1Heuristic {
	res := new(Phase1Heuristic)
	res.computeCO(moves)
	res.computeEOSlice(moves, complete)
	return res
}

// AllLowerBound returns the minimum number of moves needed to solve all the
// phase-1 axes at once.
func (p *Phase1Heuristic) AllLowerBound(c *Phase1Cube) int {
	var result int8

	// Corner orientation heuristic.
	if r := p.CO[c.XCornerOrientation]; r > result {
		result = r
	}
	if r := p.CO[c.YCornerOrientation]; r > result {
		result = r
	}
	if r := p.CO[c.ZCornerOrientation]; r > result {
		result = r
	}

	// EOSlice heuristic.
	sliceValues := []int{
		c.MSlicePermutation*2048 + c.XEdgeOrientation(),
		c.ESlicePermutation*2048 + c.FBEdgeOrientation,
		c.SSlicePermutation*2048 + c.UDEdgeOrientation,
	}
	for _, eoSlice := range sliceValues {
		if r := p.EOSlice[eoSlice]; r > result {
			result = r
		} else if r < 0 && result < 8 {
			result = 8
		}
	}

	return int(result)
}

// LowerBound returns the minimum number of moves needed to solve at least one
// phase-1 axis.
func (p *Phase1Heuristic) LowerBound(c *Phase1Cube) int {
	finalResult := int8(127)

	sliceValues := []int{
		c.MSlicePermutation*2048 + c.XEdgeOrientation(),
		c.ESlicePermutation*2048 + c.FBEdgeOrientation,
		c.SSlicePermutation*2048 + c.UDEdgeOrientation,
	}
	coValues := []int{c.XCornerOrientation, c.YCornerOrientation,
		c.ZCornerOrientation}

	for axis := 0; axis < 3; axis++ {
		result := p.CO[coValues[axis]]
		if r := p.EOSlice[sliceValues[axis]]; r > result {
			result = r
		} else if r < 0 && result < 8 {
			result = 8
		}
		if result < finalResult {
			finalResult = result
		}
	}

	return int(finalResult)
}

func (p *Phase1Heuristic) computeCO(moves *Phase1Moves) {
	for i := 0; i < 2187; i++ {
		p.CO[i] = -1
	}
	nodes := []phase1CONode{phase1CONode{1093, 0}}
	for len(nodes) > 0 {
		node := nodes[0]
		nodes = nodes[1:]
		if p.CO[node.corners] != -1 {
			continue
		}
		p.CO[node.corners] = node.depth
		for move := 0; move < 18; move++ {
			applied := moves.COMoves[node.corners][move]
			nodes = append(nodes, phase1CONode{applied, node.depth + 1})
		}
	}
}

func (p *Phase1Heuristic) computeEOSlice(moves *Phase1Moves, complete bool) {
	for i := 0; i < 1013760; i++ {
		p.EOSlice[i] = -1
	}
	nodes := []phase1EOSliceNode{phase1EOSliceNode{0, 220, 0}}
	for len(nodes) > 0 {
		node := nodes[0]
		nodes = nodes[1:]
		hash := node.slice*2048 + node.eo
		if p.EOSlice[hash] != -1 {
			continue
		}
		p.EOSlice[hash] = node.depth

		// Stop searching after 7 moves, which makes indexing faster and should
		// not greatly affect search.
		if !complete && node.depth == 7 {
			continue
		}

		for move := 0; move < 18; move++ {
			newEO := moves.EOMoves[node.eo][move]
			newSlice := moves.ESliceMoves[node.slice][move]
			node := phase1EOSliceNode{newEO, newSlice, node.depth + 1}
			nodes = append(nodes, node)
		}
	}
}

// A Phase1Solver finds solutions to a specific phase-1 state.
type Phase1Solver struct {
	stopped   chan struct{}
	solutions <-chan []Move

	heuristic *Phase1Heuristic
	moves     *Phase1Moves
}

// NewPhase1Solver creates and starts a Phase1Solver.
func NewPhase1Solver(c Phase1Cube, h *Phase1Heuristic,
	m *Phase1Moves) *Phase1Solver {
	solutions := make(chan []Move)
	res := &Phase1Solver{make(chan struct{}), solutions, h, m}
	go res.search(solutions, c)
	return res
}

func (p *Phase1Solver) Solutions() <-chan []Move {
	return p.solutions
}

func (p *Phase1Solver) Stop() {
	close(p.stopped)
}

func (p *Phase1Solver) depthFirst(solutions chan<- []Move, c Phase1Cube,
	moves []Move, depth int) bool {
	// If the depth is zero, we may have a solution.
	if depth == 0 {
		if c.AnySolved() {
			res := make([]Move, len(moves))
			copy(res, moves)
			select {
			case <-p.stopped:
				return false
			case solutions <- res:
			}
		}
		return true
	}

	// Check the heuristic.
	if p.heuristic.LowerBound(&c) > depth {
		return true
	}

	// Apply every move and recurse.
	for m := 0; m < 18; m++ {
		cube := c
		move := Move(m)
		cube.Move(move, p.moves)
		moves = append(moves, move)
		if !p.depthFirst(solutions, cube, moves, depth-1) {
			return false
		}
		moves = moves[:len(moves)-1]
		if depth >= 7 && p.isStopped() {
			return false
		}
	}

	return true
}

func (p *Phase1Solver) isStopped() bool {
	select {
	case <-p.stopped:
		return true
	default:
		return false
	}
}

func (p *Phase1Solver) search(solutions chan<- []Move, c Phase1Cube) {
	depth := 0
	for {
		moves := make([]Move, 0, depth)
		if !p.depthFirst(solutions, c, moves, depth) {
			return
		}
		depth++
	}
}

type phase1CONode struct {
	corners int
	depth   int8
}

type phase1EOSliceNode struct {
	eo    int
	slice int
	depth int8
}
<commit_msg>phase1 solver is roughly twice as fast (avg 17ms)<commit_after>package gocube

// Phase1Heuristic stores the data needed to effectively prune the search for a
// solution for phase-1.
type Phase1Heuristic struct {
	// This stores the number of moves needed to orient the corners. All the
	// fields in this array should be filled in.
	CO [2187]int8

	// This stores the number of moves needed to solve a given EO + E Slice
	// case. The fields are combined using slice*2048 + eo. If a value in this
	// array is -1, it should be interpreted as 8. This is because depths
	// greater than 7 will not be searched if the index was made under time
	// pressure.
	EOSlice [1013760]int8
}

// NewPhase1Heuristic generates a heuristic for the phase-1 solver.
// The complete argument specifies if the entire EOSlice table should be
// generated.
func NewPhase1Heuristic(moves *Phase1Moves, complete bool) *Phase1Heuristic {
	res := new(Phase1Heuristic)
	res.computeCO(moves)
	res.computeEOSlice(moves, complete)
	return res
}

// AllLowerBound returns the minimum number of moves needed to solve all the
// phase-1 axes at once.
func (p *Phase1Heuristic) AllLowerBound(c *Phase1Cube) int {
	var result int8

	// Corner orientation heuristic.
	if r := p.CO[c.XCornerOrientation]; r > result {
		result = r
	}
	if r := p.CO[c.YCornerOrientation]; r > result {
		result = r
	}
	if r := p.CO[c.ZCornerOrientation]; r > result {
		result = r
	}

	// EOSlice heuristic.
	sliceValues := []int{
		c.MSlicePermutation*2048 + c.XEdgeOrientation(),
		c.ESlicePermutation*2048 + c.FBEdgeOrientation,
		c.SSlicePermutation*2048 + c.UDEdgeOrientation,
	}
	for _, eoSlice := range sliceValues {
		if r := p.EOSlice[eoSlice]; r > result {
			result = r
		} else if r < 0 && result < 8 {
			result = 8
		}
	}

	return int(result)
}

// LowerBound returns the minimum number of moves needed to solve at least one
// phase-1 axis.
func (p *Phase1Heuristic) LowerBound(c *Phase1Cube) int {
	finalResult := int8(127)

	sliceValues := []int{
		c.MSlicePermutation*2048 + c.XEdgeOrientation(),
		c.ESlicePermutation*2048 + c.FBEdgeOrientation,
		c.SSlicePermutation*2048 + c.UDEdgeOrientation,
	}
	coValues := []int{c.XCornerOrientation, c.YCornerOrientation,
		c.ZCornerOrientation}

	for axis := 0; axis < 3; axis++ {
		result := p.CO[coValues[axis]]
		if r := p.EOSlice[sliceValues[axis]]; r > result {
			result = r
		} else if r < 0 && result < 8 {
			result = 8
		}
		if result < finalResult {
			finalResult = result
		}
	}

	return int(finalResult)
}

func (p *Phase1Heuristic) computeCO(moves *Phase1Moves) {
	for i := 0; i < 2187; i++ {
		p.CO[i] = -1
	}
	nodes := []phase1CONode{phase1CONode{1093, 0}}
	for len(nodes) > 0 {
		node := nodes[0]
		nodes = nodes[1:]
		if p.CO[node.corners] != -1 {
			continue
		}
		p.CO[node.corners] = node.depth
		for move := 0; move < 18; move++ {
			applied := moves.COMoves[node.corners][move]
			nodes = append(nodes, phase1CONode{applied, node.depth + 1})
		}
	}
}

func (p *Phase1Heuristic) computeEOSlice(moves *Phase1Moves, complete bool) {
	for i := 0; i < 1013760; i++ {
		p.EOSlice[i] = -1
	}
	nodes := []phase1EOSliceNode{phase1EOSliceNode{0, 220, 0}}
	for len(nodes) > 0 {
		node := nodes[0]
		nodes = nodes[1:]
		hash := node.slice*2048 + node.eo
		if p.EOSlice[hash] != -1 {
			continue
		}
		p.EOSlice[hash] = node.depth

		// Stop searching after 7 moves, which makes indexing faster and should
		// not greatly affect search.
		if !complete && node.depth == 7 {
			continue
		}

		for move := 0; move < 18; move++ {
			newEO := moves.EOMoves[node.eo][move]
			newSlice := moves.ESliceMoves[node.slice][move]
			node := phase1EOSliceNode{newEO, newSlice, node.depth + 1}
			nodes = append(nodes, node)
		}
	}
}

// A Phase1Solver finds solutions to a specific phase-1 state.
type Phase1Solver struct {
	stopped   chan struct{}
	solutions <-chan []Move

	heuristic *Phase1Heuristic
	moves     *Phase1Moves
}

// NewPhase1Solver creates and starts a Phase1Solver.
func NewPhase1Solver(c Phase1Cube, h *Phase1Heuristic,
	m *Phase1Moves) *Phase1Solver {
	solutions := make(chan []Move)
	res := &Phase1Solver{make(chan struct{}), solutions, h, m}
	go res.search(solutions, c)
	return res
}

func (p *Phase1Solver) Solutions() <-chan []Move {
	return p.solutions
}

func (p *Phase1Solver) Stop() {
	close(p.stopped)
}

func (p *Phase1Solver) depthFirst(solutions chan<- []Move, c Phase1Cube,
	moves []Move, depth int, lastFace int) bool {
	// If the depth is zero, we may have a solution.
	if depth == 0 {
		if c.AnySolved() {
			res := make([]Move, len(moves))
			copy(res, moves)
			select {
			case <-p.stopped:
				return false
			case solutions <- res:
			}
		}
		return true
	}

	// Check the heuristic.
	if p.heuristic.LowerBound(&c) > depth {
		return true
	}

	// Apply every move and recurse.
	for m := 0; m < 18; m++ {
		move := Move(m)
		if move.Face() == lastFace {
			continue
		}
		cube := c
		cube.Move(move, p.moves)
		moves = append(moves, move)
		if !p.depthFirst(solutions, cube, moves, depth-1, move.Face()) {
			return false
		}
		moves = moves[:len(moves)-1]
		if depth >= 7 && p.isStopped() {
			return false
		}
	}

	return true
}

func (p *Phase1Solver) isStopped() bool {
	select {
	case <-p.stopped:
		return true
	default:
		return false
	}
}

func (p *Phase1Solver) search(solutions chan<- []Move, c Phase1Cube) {
	depth := 0
	for {
		moves := make([]Move, 0, depth)
		if !p.depthFirst(solutions, c, moves, depth, 0) {
			return
		}
		depth++
	}
}

type phase1CONode struct {
	corners int
	depth   int8
}

type phase1EOSliceNode struct {
	eo    int
	slice int
	depth int8
}
<|endoftext|>


<commit_before>package migrations

import (
	"errors"
	"fmt"
	"os"
	"strings"
	"sync"
	"sync/atomic"
	"testing"

	"github.com/go-sql-driver/mysql"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"xorm.io/xorm"

	. "github.com/grafana/grafana/pkg/services/sqlstore/migrator"
	"github.com/grafana/grafana/pkg/services/sqlstore/sqlutil"
	"github.com/grafana/grafana/pkg/setting"
)

func TestMigrations(t *testing.T) {
	testDB := sqlutil.SQLite3TestDB()
	const query = `select count(*) as count from migration_log`
	result := struct{ Count int }{}

	x, err := xorm.NewEngine(testDB.DriverName, testDB.ConnStr)
	require.NoError(t, err)

	err = NewDialect(x).CleanDB()
	require.NoError(t, err)

	_, err = x.SQL(query).Get(&result)
	require.Error(t, err)

	mg := NewMigrator(x, &setting.Cfg{})
	migrations := &OSSMigrations{}
	migrations.AddMigration(mg)
	expectedMigrations := mg.GetMigrationIDs(true)

	err = mg.Start(false, 0)
	require.NoError(t, err)

	has, err := x.SQL(query).Get(&result)
	require.NoError(t, err)
	require.True(t, has)

	checkStepsAndDatabaseMatch(t, mg, expectedMigrations)

	mg = NewMigrator(x, &setting.Cfg{})
	migrations.AddMigration(mg)

	err = mg.Start(false, 0)
	require.NoError(t, err)

	has, err = x.SQL(query).Get(&result)
	require.NoError(t, err)
	require.True(t, has)
	checkStepsAndDatabaseMatch(t, mg, expectedMigrations)
}

func TestMigrationLock(t *testing.T) {
	dbType := getDBType()
	if dbType == SQLite {
		t.Skip()
	}

	testDB := getTestDB(t, dbType)

	x, err := xorm.NewEngine(testDB.DriverName, testDB.ConnStr)
	require.NoError(t, err)

	dialect := NewDialect(x)

	sess := x.NewSession()
	t.Cleanup(func() {
		sess.Close()
	})

	cfg := LockCfg{Session: sess}

	t.Run("obtaining lock should succeed", func(t *testing.T) {
		err := dialect.Lock(cfg)
		require.NoError(t, err)

		t.Run("releasing previously obtained lock should succeed", func(t *testing.T) {
			err := dialect.Unlock(cfg)
			require.NoError(t, err)

			t.Run("releasing already released lock should fail", func(t *testing.T) {
				err := dialect.Unlock(cfg)
				require.Error(t, err)
				assert.ErrorIs(t, err, ErrReleaseLockDB)
			})
		})
	})

	t.Run("obtaining lock twice should succeed", func(t *testing.T) {
		err = dialect.Lock(cfg)
		require.NoError(t, err)

		err = dialect.Lock(cfg)
		require.NoError(t, err)

		t.Cleanup(func() {
			err := dialect.Unlock(cfg)
			require.NoError(t, err)

			err = dialect.Unlock(cfg)
			require.NoError(t, err)
		})
	})

	t.Run("obtaining same lock from another session should fail", func(t *testing.T) {
		x2, err := xorm.NewEngine(testDB.DriverName, testDB.ConnStr)
		require.NoError(t, err)
		sess2 := x2.NewSession()

		d2 := NewDialect(x2)

		err = dialect.Lock(cfg)
		require.NoError(t, err)

		err = d2.Lock(LockCfg{Session: sess2})
		require.Error(t, err)
		assert.ErrorIs(t, err, ErrLockDB)

		t.Cleanup(func() {
			err := dialect.Unlock(cfg)
			require.NoError(t, err)
		})
	})

	t.Run("obtaining lock for a another database should succeed", func(t *testing.T) {
		err := dialect.Lock(cfg)
		require.NoError(t, err)

		x, err := xorm.NewEngine(testDB.DriverName, replaceDBName(t, testDB.ConnStr, dbType))
		require.NoError(t, err)

		d := NewDialect(x)
		err = d.Lock(cfg)
		require.NoError(t, err)

		t.Cleanup(func() {
			err := dialect.Unlock(cfg)
			require.NoError(t, err)

			err = d.Unlock(cfg)
			require.NoError(t, err)
		})
	})
}

func TestMigratorLocking(t *testing.T) {
	dbType := getDBType()
	testDB := getTestDB(t, dbType)

	x, err := xorm.NewEngine(testDB.DriverName, testDB.ConnStr)
	require.NoError(t, err)

	err = NewDialect(x).CleanDB()
	require.NoError(t, err)

	mg := NewMigrator(x, &setting.Cfg{})
	migrations := &OSSMigrations{}
	migrations.AddMigration(mg)

	var errorNum int64
	t.Run("when concurrent migrations for the same migrator occur, the second one should fail", func(t *testing.T) {
		for i := 0; i < 2; i++ {
			i := i // capture i variable
			t.Run(fmt.Sprintf("run migration %d", i), func(t *testing.T) {
				t.Parallel()
				err := mg.Start(true, 0)
				if err != nil {
					if errors.Is(err, ErrMigratorIsLocked) {
						atomic.AddInt64(&errorNum, 1)
					}
				}
			})
		}
	})
	assert.Equal(t, int64(1), errorNum)
}

func TestDatabaseLocking(t *testing.T) {
	dbType := getDBType()
	// skip for SQLite since there is no database locking (only migrator locking)
	if dbType == SQLite {
		t.Skip()
	}

	testDB := getTestDB(t, dbType)

	x, err := xorm.NewEngine(testDB.DriverName, testDB.ConnStr)
	require.NoError(t, err)

	err = NewDialect(x).CleanDB()
	require.NoError(t, err)

	mg1 := NewMigrator(x, &setting.Cfg{})
	migrations := &OSSMigrations{}
	migrations.AddMigration(mg1)
	reg := registry{
		migrators: make(map[int]*Migrator, 2),
	}
	reg.set(0, mg1)

	mg2 := NewMigrator(x, &setting.Cfg{})
	migrations.AddMigration(mg2)
	reg.set(1, mg2)

	var errorNum int64
	t.Run("when concurrent migrations occur for different migrators occur, the second one should fail", func(t *testing.T) {
		for i := 0; i < 2; i++ {
			i := i // capture i variable
			t.Run(fmt.Sprintf("run migration %d", i), func(t *testing.T) {
				mg, err := reg.get(i)
				require.NoError(t, err)
				t.Parallel()
				err = mg.Start(true, 0)
				if err != nil {
					assert.ErrorIs(t, err, ErrLockDB)
					if errors.Is(err, ErrLockDB) {
						atomic.AddInt64(&errorNum, 1)
					}
				}
			})
		}
	})
	assert.Equal(t, int64(1), errorNum)
}

func checkStepsAndDatabaseMatch(t *testing.T, mg *Migrator, expected []string) {
	t.Helper()
	log, err := mg.GetMigrationLog()
	require.NoError(t, err)
	missing := make([]string, 0)
	for _, id := range expected {
		_, ok := log[id]
		if !ok {
			missing = append(missing, id)
		}
	}
	notIntended := make([]string, 0)
	for logId := range log {
		found := false
		for _, s := range expected {
			found = s == logId
			if found {
				break
			}
		}
		if !found {
			notIntended = append(notIntended, logId)
		}
	}

	if len(missing) == 0 && len(notIntended) == 0 {
		return
	}

	var msg string
	if len(missing) > 0 {
		msg = fmt.Sprintf("was not executed [%v], ", strings.Join(missing, ", "))
	}
	if len(notIntended) > 0 {
		msg += fmt.Sprintf("executed but should not [%v]", strings.Join(notIntended, ", "))
	}
	require.Failf(t, "the number of migrations does not match log in database", msg)
}

func getDBType() string {
	dbType := SQLite

	// environment variable present for test db?
	if db, present := os.LookupEnv("GRAFANA_TEST_DB"); present {
		dbType = db
	}
	return dbType
}

func getTestDB(t *testing.T, dbType string) sqlutil.TestDB {
	switch dbType {
	case "mysql":
		return sqlutil.MySQLTestDB()
	case "postgres":
		return sqlutil.PostgresTestDB()
	default:
		f, err := os.CreateTemp(".", "grafana-test-db-")
		require.NoError(t, err)
		t.Cleanup(func() {
			err := os.Remove(f.Name())
			require.NoError(t, err)
		})

		return sqlutil.TestDB{
			DriverName: "sqlite3",
			ConnStr:    f.Name(),
		}
	}
}

func replaceDBName(t *testing.T, connStr, dbType string) string {
	switch dbType {
	case "mysql":
		cfg, err := mysql.ParseDSN(connStr)
		require.NoError(t, err)
		cfg.DBName = "grafana_ds_tests"
		return cfg.FormatDSN()
	case "postgres":
		return strings.Replace(connStr, "dbname=grafanatest", "dbname=grafanadstest", 1)
	default:
		return connStr
	}
}

type registry struct {
	mu        sync.Mutex
	migrators map[int]*Migrator
}

func (r *registry) get(i int) (*Migrator, error) {
	r.mu.Lock()
	defer r.mu.Unlock()

	m, ok := r.migrators[i]
	if !ok {
		return nil, fmt.Errorf("invalid index: %d", i)
	}
	return m, nil
}

func (r *registry) set(i int, mg *Migrator) {
	r.mu.Lock()
	defer r.mu.Unlock()

	r.migrators[i] = mg
}
<commit_msg>Skip flaky test (#45785)<commit_after>package migrations

import (
	"errors"
	"fmt"
	"os"
	"strings"
	"sync"
	"sync/atomic"
	"testing"

	"github.com/go-sql-driver/mysql"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"xorm.io/xorm"

	. "github.com/grafana/grafana/pkg/services/sqlstore/migrator"
	"github.com/grafana/grafana/pkg/services/sqlstore/sqlutil"
	"github.com/grafana/grafana/pkg/setting"
)

func TestMigrations(t *testing.T) {
	testDB := sqlutil.SQLite3TestDB()
	const query = `select count(*) as count from migration_log`
	result := struct{ Count int }{}

	x, err := xorm.NewEngine(testDB.DriverName, testDB.ConnStr)
	require.NoError(t, err)

	err = NewDialect(x).CleanDB()
	require.NoError(t, err)

	_, err = x.SQL(query).Get(&result)
	require.Error(t, err)

	mg := NewMigrator(x, &setting.Cfg{})
	migrations := &OSSMigrations{}
	migrations.AddMigration(mg)
	expectedMigrations := mg.GetMigrationIDs(true)

	err = mg.Start(false, 0)
	require.NoError(t, err)

	has, err := x.SQL(query).Get(&result)
	require.NoError(t, err)
	require.True(t, has)

	checkStepsAndDatabaseMatch(t, mg, expectedMigrations)

	mg = NewMigrator(x, &setting.Cfg{})
	migrations.AddMigration(mg)

	err = mg.Start(false, 0)
	require.NoError(t, err)

	has, err = x.SQL(query).Get(&result)
	require.NoError(t, err)
	require.True(t, has)
	checkStepsAndDatabaseMatch(t, mg, expectedMigrations)
}

func TestMigrationLock(t *testing.T) {
	dbType := getDBType()
	if dbType == SQLite {
		t.Skip()
	}

	testDB := getTestDB(t, dbType)

	x, err := xorm.NewEngine(testDB.DriverName, testDB.ConnStr)
	require.NoError(t, err)

	dialect := NewDialect(x)

	sess := x.NewSession()
	t.Cleanup(func() {
		sess.Close()
	})

	cfg := LockCfg{Session: sess}

	t.Run("obtaining lock should succeed", func(t *testing.T) {
		err := dialect.Lock(cfg)
		require.NoError(t, err)

		t.Run("releasing previously obtained lock should succeed", func(t *testing.T) {
			err := dialect.Unlock(cfg)
			require.NoError(t, err)

			t.Run("releasing already released lock should fail", func(t *testing.T) {
				err := dialect.Unlock(cfg)
				require.Error(t, err)
				assert.ErrorIs(t, err, ErrReleaseLockDB)
			})
		})
	})

	t.Run("obtaining lock twice should succeed", func(t *testing.T) {
		err = dialect.Lock(cfg)
		require.NoError(t, err)

		err = dialect.Lock(cfg)
		require.NoError(t, err)

		t.Cleanup(func() {
			err := dialect.Unlock(cfg)
			require.NoError(t, err)

			err = dialect.Unlock(cfg)
			require.NoError(t, err)
		})
	})

	t.Run("obtaining same lock from another session should fail", func(t *testing.T) {
		x2, err := xorm.NewEngine(testDB.DriverName, testDB.ConnStr)
		require.NoError(t, err)
		sess2 := x2.NewSession()

		d2 := NewDialect(x2)

		err = dialect.Lock(cfg)
		require.NoError(t, err)

		err = d2.Lock(LockCfg{Session: sess2})
		require.Error(t, err)
		assert.ErrorIs(t, err, ErrLockDB)

		t.Cleanup(func() {
			err := dialect.Unlock(cfg)
			require.NoError(t, err)
		})
	})

	t.Run("obtaining lock for a another database should succeed", func(t *testing.T) {
		err := dialect.Lock(cfg)
		require.NoError(t, err)

		x, err := xorm.NewEngine(testDB.DriverName, replaceDBName(t, testDB.ConnStr, dbType))
		require.NoError(t, err)

		d := NewDialect(x)
		err = d.Lock(cfg)
		require.NoError(t, err)

		t.Cleanup(func() {
			err := dialect.Unlock(cfg)
			require.NoError(t, err)

			err = d.Unlock(cfg)
			require.NoError(t, err)
		})
	})
}

func TestMigratorLocking(t *testing.T) {
	dbType := getDBType()
	testDB := getTestDB(t, dbType)
	// skip for SQLite for now since it occasionally fails for not clear reason
	// anyway starting migrations concurretly for the same migrator is impossible use case
	if dbType == SQLite {
		t.Skip()
	}

	x, err := xorm.NewEngine(testDB.DriverName, testDB.ConnStr)
	require.NoError(t, err)

	err = NewDialect(x).CleanDB()
	require.NoError(t, err)

	mg := NewMigrator(x, &setting.Cfg{})
	migrations := &OSSMigrations{}
	migrations.AddMigration(mg)

	var errorNum int64
	t.Run("when concurrent migrations for the same migrator occur, the second one should fail", func(t *testing.T) {
		for i := 0; i < 2; i++ {
			i := i // capture i variable
			t.Run(fmt.Sprintf("run migration %d", i), func(t *testing.T) {
				t.Parallel()
				err := mg.Start(true, 0)
				if err != nil {
					if errors.Is(err, ErrMigratorIsLocked) {
						atomic.AddInt64(&errorNum, 1)
					}
				}
			})
		}
	})
	assert.Equal(t, int64(1), atomic.LoadInt64(&errorNum))
}

func TestDatabaseLocking(t *testing.T) {
	dbType := getDBType()
	// skip for SQLite since there is no database locking (only migrator locking)
	if dbType == SQLite {
		t.Skip()
	}

	testDB := getTestDB(t, dbType)

	x, err := xorm.NewEngine(testDB.DriverName, testDB.ConnStr)
	require.NoError(t, err)

	err = NewDialect(x).CleanDB()
	require.NoError(t, err)

	mg1 := NewMigrator(x, &setting.Cfg{})
	migrations := &OSSMigrations{}
	migrations.AddMigration(mg1)
	reg := registry{
		migrators: make(map[int]*Migrator, 2),
	}
	reg.set(0, mg1)

	mg2 := NewMigrator(x, &setting.Cfg{})
	migrations.AddMigration(mg2)
	reg.set(1, mg2)

	var errorNum int64
	t.Run("when concurrent migrations occur for different migrators occur, the second one should fail", func(t *testing.T) {
		for i := 0; i < 2; i++ {
			i := i // capture i variable
			t.Run(fmt.Sprintf("run migration %d", i), func(t *testing.T) {
				mg, err := reg.get(i)
				require.NoError(t, err)
				t.Parallel()
				err = mg.Start(true, 0)
				if err != nil {
					assert.ErrorIs(t, err, ErrLockDB)
					if errors.Is(err, ErrLockDB) {
						atomic.AddInt64(&errorNum, 1)
					}
				}
			})
		}
	})
	assert.Equal(t, int64(1), errorNum)
}

func checkStepsAndDatabaseMatch(t *testing.T, mg *Migrator, expected []string) {
	t.Helper()
	log, err := mg.GetMigrationLog()
	require.NoError(t, err)
	missing := make([]string, 0)
	for _, id := range expected {
		_, ok := log[id]
		if !ok {
			missing = append(missing, id)
		}
	}
	notIntended := make([]string, 0)
	for logId := range log {
		found := false
		for _, s := range expected {
			found = s == logId
			if found {
				break
			}
		}
		if !found {
			notIntended = append(notIntended, logId)
		}
	}

	if len(missing) == 0 && len(notIntended) == 0 {
		return
	}

	var msg string
	if len(missing) > 0 {
		msg = fmt.Sprintf("was not executed [%v], ", strings.Join(missing, ", "))
	}
	if len(notIntended) > 0 {
		msg += fmt.Sprintf("executed but should not [%v]", strings.Join(notIntended, ", "))
	}
	require.Failf(t, "the number of migrations does not match log in database", msg)
}

func getDBType() string {
	dbType := SQLite

	// environment variable present for test db?
	if db, present := os.LookupEnv("GRAFANA_TEST_DB"); present {
		dbType = db
	}
	return dbType
}

func getTestDB(t *testing.T, dbType string) sqlutil.TestDB {
	switch dbType {
	case "mysql":
		return sqlutil.MySQLTestDB()
	case "postgres":
		return sqlutil.PostgresTestDB()
	default:
		f, err := os.CreateTemp(".", "grafana-test-db-")
		require.NoError(t, err)
		t.Cleanup(func() {
			err := os.Remove(f.Name())
			require.NoError(t, err)
		})

		return sqlutil.TestDB{
			DriverName: "sqlite3",
			ConnStr:    f.Name(),
		}
	}
}

func replaceDBName(t *testing.T, connStr, dbType string) string {
	switch dbType {
	case "mysql":
		cfg, err := mysql.ParseDSN(connStr)
		require.NoError(t, err)
		cfg.DBName = "grafana_ds_tests"
		return cfg.FormatDSN()
	case "postgres":
		return strings.Replace(connStr, "dbname=grafanatest", "dbname=grafanadstest", 1)
	default:
		return connStr
	}
}

type registry struct {
	mu        sync.Mutex
	migrators map[int]*Migrator
}

func (r *registry) get(i int) (*Migrator, error) {
	r.mu.Lock()
	defer r.mu.Unlock()

	m, ok := r.migrators[i]
	if !ok {
		return nil, fmt.Errorf("invalid index: %d", i)
	}
	return m, nil
}

func (r *registry) set(i int, mg *Migrator) {
	r.mu.Lock()
	defer r.mu.Unlock()

	r.migrators[i] = mg
}
<|endoftext|>


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

package kube

import (
	"context"
	"fmt"
	"sync"
	"time"

	"github.com/hashicorp/go-multierror"
	kubeCore "k8s.io/api/core/v1"

	istioKube "istio.io/istio/pkg/kube"
	"istio.io/istio/pkg/test"
	echoClient "istio.io/istio/pkg/test/echo"
	"istio.io/istio/pkg/test/echo/common"
	"istio.io/istio/pkg/test/echo/proto"
	"istio.io/istio/pkg/test/framework/components/cluster"
	"istio.io/istio/pkg/test/framework/components/echo"
	"istio.io/istio/pkg/test/framework/errors"
	"istio.io/istio/pkg/test/framework/resource"
	"istio.io/istio/pkg/test/scopes"
	"istio.io/istio/pkg/test/util/retry"
)

const (
	appContainerName = "app"
)

var _ echo.Workload = &workload{}

type workloadConfig struct {
	pod        kubeCore.Pod
	hasSidecar bool
	grpcPort   uint16
	cluster    cluster.Cluster
	tls        *common.TLSSettings
}

type workload struct {
	client *echoClient.Client

	workloadConfig
	forwarder  istioKube.PortForwarder
	sidecar    *sidecar
	ctx        resource.Context
	mutex      sync.Mutex
	connectErr error
}

func newWorkload(cfg workloadConfig, ctx resource.Context) (*workload, error) {
	w := &workload{
		workloadConfig: cfg,
		ctx:            ctx,
	}

	// If the pod is ready, connect.
	if err := w.Update(cfg.pod); err != nil {
		return nil, err
	}

	return w, nil
}

func (w *workload) IsReady() bool {
	w.mutex.Lock()
	ready := w.isConnected()
	w.mutex.Unlock()
	return ready
}

func (w *workload) Client() (c *echoClient.Client, err error) {
	w.mutex.Lock()
	c = w.client
	if c == nil {
		err = fmt.Errorf("attempt to use disconnected client for echo pod %s/%s (in cluster %s)",
			w.pod.Namespace, w.pod.Name, w.cluster.Name())
	}
	w.mutex.Unlock()
	return
}

func (w *workload) Update(pod kubeCore.Pod) error {
	w.mutex.Lock()
	defer w.mutex.Unlock()

	if isPodReady(pod) && !w.isConnected() {
		if err := w.connect(pod); err != nil {
			w.connectErr = err
			return err
		}
	} else if !isPodReady(pod) && w.isConnected() {
		scopes.Framework.Infof("echo pod %s/%s (in cluster %s) transitioned to NOT READY. Pod Status=%v",
			pod.Namespace, pod.Name, w.cluster.Name(), pod.Status.Phase)
		w.pod = pod
		return w.disconnect()
	}

	// Update the pod.
	w.pod = pod
	return nil
}

func (w *workload) Close() (err error) {
	w.mutex.Lock()
	defer w.mutex.Unlock()

	if w.isConnected() {
		return w.disconnect()
	}
	return nil
}

func (w *workload) PodName() string {
	w.mutex.Lock()
	n := w.pod.Name
	w.mutex.Unlock()
	return n
}

func (w *workload) Address() string {
	w.mutex.Lock()
	ip := w.pod.Status.PodIP
	w.mutex.Unlock()
	return ip
}

func (w *workload) ForwardEcho(ctx context.Context, request *proto.ForwardEchoRequest) (echoClient.Responses, error) {
	w.mutex.Lock()
	c := w.client
	if c == nil {
		return nil, fmt.Errorf("failed forwarding echo for disconnected pod %s/%s",
			w.pod.Namespace, w.pod.Name)
	}
	w.mutex.Unlock()

	return c.ForwardEcho(ctx, request)
}

func (w *workload) Sidecar() echo.Sidecar {
	w.mutex.Lock()
	s := w.sidecar
	w.mutex.Unlock()
	return s
}

func (w *workload) Cluster() cluster.Cluster {
	return w.cluster
}

func (w *workload) Logs() (string, error) {
	return w.cluster.PodLogs(context.TODO(), w.pod.Name, w.pod.Namespace, appContainerName, false)
}

func (w *workload) LogsOrFail(t test.Failer) string {
	t.Helper()
	logs, err := w.Logs()
	if err != nil {
		t.Fatal(err)
	}
	return logs
}

func isPodReady(pod kubeCore.Pod) bool {
	return istioKube.CheckPodReady(&pod) == nil
}

func (w *workload) isConnected() bool {
	return w.forwarder != nil
}

func (w *workload) connect(pod kubeCore.Pod) (err error) {
	defer func() {
		if err != nil {
			_ = w.disconnect()
		}
	}()

	// Create a forwarder to the command port of the app.
	if err = retry.UntilSuccess(func() error {
		w.forwarder, err = w.cluster.NewPortForwarder(pod.Name, pod.Namespace, "", 0, int(w.grpcPort))
		if err != nil {
			return fmt.Errorf("failed creating new port forwarder for pod %s/%s: %v",
				pod.Namespace, pod.Name, err)
		}
		if err = w.forwarder.Start(); err != nil {
			return fmt.Errorf("failed starting port forwarder for pod %s/%s: %v",
				pod.Namespace, pod.Name, err)
		}
		return nil
	}, retry.BackoffDelay(100*time.Millisecond), retry.Timeout(10*time.Second)); err != nil {
		return err
	}

	// Create a gRPC client to this workload.
	w.client, err = echoClient.New(w.forwarder.Address(), w.tls)
	if err != nil {
		return fmt.Errorf("failed connecting to grpc client to pod %s/%s : %v",
			pod.Namespace, pod.Name, err)
	}

	if w.hasSidecar {
		w.sidecar = newSidecar(pod, w.cluster)
	}

	return nil
}

func (w *workload) disconnect() (err error) {
	if w.client != nil {
		err = multierror.Append(err, w.client.Close()).ErrorOrNil()
		w.client = nil
	}
	if w.forwarder != nil {
		w.forwarder.Close()
		w.forwarder = nil
	}
	if w.ctx.Settings().FailOnDeprecation && w.sidecar != nil {
		err = multierror.Append(err, w.checkDeprecation()).ErrorOrNil()
		w.sidecar = nil
	}
	return err
}

func (w *workload) checkDeprecation() error {
	logs, err := w.sidecar.Logs()
	if err != nil {
		return fmt.Errorf("could not get sidecar logs to inspect for deprecation messages: %v", err)
	}

	info := fmt.Sprintf("pod: %s/%s", w.pod.Namespace, w.pod.Name)
	return errors.FindDeprecatedMessagesInEnvoyLog(logs, info)
}
<commit_msg>Fix missing unlock (#40817)<commit_after>// Copyright Istio Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package kube

import (
	"context"
	"fmt"
	"sync"
	"time"

	"github.com/hashicorp/go-multierror"
	kubeCore "k8s.io/api/core/v1"

	istioKube "istio.io/istio/pkg/kube"
	"istio.io/istio/pkg/test"
	echoClient "istio.io/istio/pkg/test/echo"
	"istio.io/istio/pkg/test/echo/common"
	"istio.io/istio/pkg/test/echo/proto"
	"istio.io/istio/pkg/test/framework/components/cluster"
	"istio.io/istio/pkg/test/framework/components/echo"
	"istio.io/istio/pkg/test/framework/errors"
	"istio.io/istio/pkg/test/framework/resource"
	"istio.io/istio/pkg/test/scopes"
	"istio.io/istio/pkg/test/util/retry"
)

const (
	appContainerName = "app"
)

var _ echo.Workload = &workload{}

type workloadConfig struct {
	pod        kubeCore.Pod
	hasSidecar bool
	grpcPort   uint16
	cluster    cluster.Cluster
	tls        *common.TLSSettings
}

type workload struct {
	client *echoClient.Client

	workloadConfig
	forwarder  istioKube.PortForwarder
	sidecar    *sidecar
	ctx        resource.Context
	mutex      sync.Mutex
	connectErr error
}

func newWorkload(cfg workloadConfig, ctx resource.Context) (*workload, error) {
	w := &workload{
		workloadConfig: cfg,
		ctx:            ctx,
	}

	// If the pod is ready, connect.
	if err := w.Update(cfg.pod); err != nil {
		return nil, err
	}

	return w, nil
}

func (w *workload) IsReady() bool {
	w.mutex.Lock()
	ready := w.isConnected()
	w.mutex.Unlock()
	return ready
}

func (w *workload) Client() (c *echoClient.Client, err error) {
	w.mutex.Lock()
	c = w.client
	if c == nil {
		err = fmt.Errorf("attempt to use disconnected client for echo pod %s/%s (in cluster %s)",
			w.pod.Namespace, w.pod.Name, w.cluster.Name())
	}
	w.mutex.Unlock()
	return
}

func (w *workload) Update(pod kubeCore.Pod) error {
	w.mutex.Lock()
	defer w.mutex.Unlock()

	if isPodReady(pod) && !w.isConnected() {
		if err := w.connect(pod); err != nil {
			w.connectErr = err
			return err
		}
	} else if !isPodReady(pod) && w.isConnected() {
		scopes.Framework.Infof("echo pod %s/%s (in cluster %s) transitioned to NOT READY. Pod Status=%v",
			pod.Namespace, pod.Name, w.cluster.Name(), pod.Status.Phase)
		w.pod = pod
		return w.disconnect()
	}

	// Update the pod.
	w.pod = pod
	return nil
}

func (w *workload) Close() (err error) {
	w.mutex.Lock()
	defer w.mutex.Unlock()

	if w.isConnected() {
		return w.disconnect()
	}
	return nil
}

func (w *workload) PodName() string {
	w.mutex.Lock()
	n := w.pod.Name
	w.mutex.Unlock()
	return n
}

func (w *workload) Address() string {
	w.mutex.Lock()
	ip := w.pod.Status.PodIP
	w.mutex.Unlock()
	return ip
}

func (w *workload) ForwardEcho(ctx context.Context, request *proto.ForwardEchoRequest) (echoClient.Responses, error) {
	w.mutex.Lock()
	c := w.client
	if c == nil {
		w.mutex.Unlock()
		return nil, fmt.Errorf("failed forwarding echo for disconnected pod %s/%s",
			w.pod.Namespace, w.pod.Name)
	}
	w.mutex.Unlock()

	return c.ForwardEcho(ctx, request)
}

func (w *workload) Sidecar() echo.Sidecar {
	w.mutex.Lock()
	s := w.sidecar
	w.mutex.Unlock()
	return s
}

func (w *workload) Cluster() cluster.Cluster {
	return w.cluster
}

func (w *workload) Logs() (string, error) {
	return w.cluster.PodLogs(context.TODO(), w.pod.Name, w.pod.Namespace, appContainerName, false)
}

func (w *workload) LogsOrFail(t test.Failer) string {
	t.Helper()
	logs, err := w.Logs()
	if err != nil {
		t.Fatal(err)
	}
	return logs
}

func isPodReady(pod kubeCore.Pod) bool {
	return istioKube.CheckPodReady(&pod) == nil
}

func (w *workload) isConnected() bool {
	return w.forwarder != nil
}

func (w *workload) connect(pod kubeCore.Pod) (err error) {
	defer func() {
		if err != nil {
			_ = w.disconnect()
		}
	}()

	// Create a forwarder to the command port of the app.
	if err = retry.UntilSuccess(func() error {
		w.forwarder, err = w.cluster.NewPortForwarder(pod.Name, pod.Namespace, "", 0, int(w.grpcPort))
		if err != nil {
			return fmt.Errorf("failed creating new port forwarder for pod %s/%s: %v",
				pod.Namespace, pod.Name, err)
		}
		if err = w.forwarder.Start(); err != nil {
			return fmt.Errorf("failed starting port forwarder for pod %s/%s: %v",
				pod.Namespace, pod.Name, err)
		}
		return nil
	}, retry.BackoffDelay(100*time.Millisecond), retry.Timeout(10*time.Second)); err != nil {
		return err
	}

	// Create a gRPC client to this workload.
	w.client, err = echoClient.New(w.forwarder.Address(), w.tls)
	if err != nil {
		return fmt.Errorf("failed connecting to grpc client to pod %s/%s : %v",
			pod.Namespace, pod.Name, err)
	}

	if w.hasSidecar {
		w.sidecar = newSidecar(pod, w.cluster)
	}

	return nil
}

func (w *workload) disconnect() (err error) {
	if w.client != nil {
		err = multierror.Append(err, w.client.Close()).ErrorOrNil()
		w.client = nil
	}
	if w.forwarder != nil {
		w.forwarder.Close()
		w.forwarder = nil
	}
	if w.ctx.Settings().FailOnDeprecation && w.sidecar != nil {
		err = multierror.Append(err, w.checkDeprecation()).ErrorOrNil()
		w.sidecar = nil
	}
	return err
}

func (w *workload) checkDeprecation() error {
	logs, err := w.sidecar.Logs()
	if err != nil {
		return fmt.Errorf("could not get sidecar logs to inspect for deprecation messages: %v", err)
	}

	info := fmt.Sprintf("pod: %s/%s", w.pod.Namespace, w.pod.Name)
	return errors.FindDeprecatedMessagesInEnvoyLog(logs, info)
}
<|endoftext|>


<commit_before>package coap

import (
	"bufio"
	"encoding/binary"
	"errors"
)

// TCPMessage is a CoAP Message that can encode itself for TCP
// transport.
type TcpMessage struct {
	Message
}

func (m *TcpMessage) MarshalBinary() ([]byte, error) {

	bin, err := m.Message.MarshalBinary()
	if err != nil {
		return nil, err
	}

	/*
		A CoAP TCP message looks like:

		     0                   1                   2                   3
		    0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
		   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
		   |        Message Length         |Ver| T |  TKL  |      Code     |
		   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
		   |   Token (if any, TKL bytes) ...
		   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
		   |   Options (if any) ...
		   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
		   |1 1 1 1 1 1 1 1|    Payload (if any) ...
		   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
	*/

	bin[2] = bin[0]
	bin[3] = bin[1]

	// insert len
	binary.BigEndian.PutUint16(bin, uint16(len(bin)-2))

	return bin, nil
}

func (m *TcpMessage) UnmarshalBinary(data []byte) error {
	if len(data) < 4 {
		return errors.New("short packet")
	}

	data[0] = data[2]
	data[1] = data[3]

	return m.Message.UnmarshalBinary(data)
}

func Decode(reader *bufio.Reader) (*TcpMessage, error) {
	header := []byte{0, 0}

	nr, err := reader.Read(header)

	if err != nil {
		return nil, err
	}

	if nr < 2 {
		return nil, errors.New("can't read 2 bytes")
	}

	ln := binary.BigEndian.Uint16(header)

	packet := make([]byte, ln)

	m := TcpMessage{}

	err = m.UnmarshalBinary(packet)
	return &m, err
}
<commit_msg>Document and fix Decode(io.Reader)<commit_after>package coap

import (
	"encoding/binary"
	"errors"
	"io"
)

// TcpMessage is a CoAP Message that can encode itself for TCP
// transport.
type TcpMessage struct {
	Message
}

func (m *TcpMessage) MarshalBinary() ([]byte, error) {

	bin, err := m.Message.MarshalBinary()
	if err != nil {
		return nil, err
	}

	/*
		A CoAP TCP message looks like:

		     0                   1                   2                   3
		    0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
		   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
		   |        Message Length         |Ver| T |  TKL  |      Code     |
		   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
		   |   Token (if any, TKL bytes) ...
		   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
		   |   Options (if any) ...
		   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
		   |1 1 1 1 1 1 1 1|    Payload (if any) ...
		   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
	*/

	bin[2] = bin[0]
	bin[3] = bin[1]

	// insert len
	binary.BigEndian.PutUint16(bin, uint16(len(bin)-2))

	return bin, nil
}

func (m *TcpMessage) UnmarshalBinary(data []byte) error {
	if len(data) < 4 {
		return errors.New("short packet")
	}

	data[0] = data[2]
	data[1] = data[3]

	return m.Message.UnmarshalBinary(data)
}

// Decode reads a single message from its input.
func Decode(r io.Reader) (*TcpMessage, error) {
	header := []byte{0, 0}

	_, err := io.ReadFull(r, header)

	if err != nil {
		return nil, err
	}

	ln := binary.BigEndian.Uint16(header)

	packet := make([]byte, ln)
	_, err = io.ReadFull(r, packet)
	if err != nil {
		return nil, err
	}

	m := TcpMessage{}

	err = m.UnmarshalBinary(packet)
	return &m, err
}
<|endoftext|>


<commit_before>package fuse_test

import (
	"io/ioutil"
	"os"
	"path"
	"strings"
	"testing"

	"golang.org/x/net/context"

	"github.com/jacobsa/fuse"
	"github.com/jacobsa/fuse/fuseops"
	"github.com/jacobsa/fuse/fuseutil"
)

////////////////////////////////////////////////////////////////////////
// minimalFS
////////////////////////////////////////////////////////////////////////

// A minimal fuseutil.FileSystem that can successfully mount but do nothing
// else.
type minimalFS struct {
	fuseutil.NotImplementedFileSystem
}

func (fs *minimalFS) StatFS(
	ctx context.Context,
	op *fuseops.StatFSOp) (err error) {
	return
}

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func TestSuccessfulMount(t *testing.T) {
	ctx := context.Background()

	// Set up a temporary directory.
	dir, err := ioutil.TempDir("", "mount_test")
	if err != nil {
		t.Fatal("ioutil.TempDir: %v", err)
	}

	defer os.RemoveAll(dir)

	// Mount.
	fs := &minimalFS{}
	mfs, err := fuse.Mount(
		dir,
		fuseutil.NewFileSystemServer(fs),
		&fuse.MountConfig{})

	if err != nil {
		t.Fatalf("fuse.Mount: %v", err)
	}

	defer func() {
		if err := mfs.Join(ctx); err != nil {
			t.Errorf("Joining: %v", err)
		}
	}()

	defer fuse.Unmount(mfs.Dir())
}

func TestNonEmptyMountPoint(t *testing.T) {
	ctx := context.Background()

	// Set up a temporary directory.
	dir, err := ioutil.TempDir("", "mount_test")
	if err != nil {
		t.Fatal("ioutil.TempDir: %v", err)
	}

	defer os.RemoveAll(dir)

	// Add a file within it.
	err = ioutil.WriteFile(path.Join(dir, "foo"), []byte{}, 0600)
	if err != nil {
		t.Fatalf("ioutil.WriteFile: %v", err)
	}

	// Attempt to mount.
	fs := &minimalFS{}
	mfs, err := fuse.Mount(
		dir,
		fuseutil.NewFileSystemServer(fs),
		&fuse.MountConfig{})

	if err == nil {
		fuse.Unmount(mfs.Dir())
		mfs.Join(ctx)
		t.Fatal("fuse.Mount returned nil")
	}

	const want = "not empty"
	if got := err.Error(); !strings.Contains(got, want) {
		t.Errorf("Unexpected error: %v", got)
	}
}
<commit_msg>Switch to a test that can pass on OS X, too.<commit_after>package fuse_test

import (
	"io/ioutil"
	"os"
	"path"
	"strings"
	"testing"

	"golang.org/x/net/context"

	"github.com/jacobsa/fuse"
	"github.com/jacobsa/fuse/fuseops"
	"github.com/jacobsa/fuse/fuseutil"
)

////////////////////////////////////////////////////////////////////////
// minimalFS
////////////////////////////////////////////////////////////////////////

// A minimal fuseutil.FileSystem that can successfully mount but do nothing
// else.
type minimalFS struct {
	fuseutil.NotImplementedFileSystem
}

func (fs *minimalFS) StatFS(
	ctx context.Context,
	op *fuseops.StatFSOp) (err error) {
	return
}

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func TestSuccessfulMount(t *testing.T) {
	ctx := context.Background()

	// Set up a temporary directory.
	dir, err := ioutil.TempDir("", "mount_test")
	if err != nil {
		t.Fatal("ioutil.TempDir: %v", err)
	}

	defer os.RemoveAll(dir)

	// Mount.
	fs := &minimalFS{}
	mfs, err := fuse.Mount(
		dir,
		fuseutil.NewFileSystemServer(fs),
		&fuse.MountConfig{})

	if err != nil {
		t.Fatalf("fuse.Mount: %v", err)
	}

	defer func() {
		if err := mfs.Join(ctx); err != nil {
			t.Errorf("Joining: %v", err)
		}
	}()

	defer fuse.Unmount(mfs.Dir())
}

func TestNonexistentMountPoint(t *testing.T) {
	ctx := context.Background()

	// Set up a temporary directory.
	dir, err := ioutil.TempDir("", "mount_test")
	if err != nil {
		t.Fatal("ioutil.TempDir: %v", err)
	}

	defer os.RemoveAll(dir)

	// Attempt to mount into a sub-directory that doesn't exist.
	fs := &minimalFS{}
	mfs, err := fuse.Mount(
		path.Join(dir, "foo"),
		fuseutil.NewFileSystemServer(fs),
		&fuse.MountConfig{})

	if err == nil {
		fuse.Unmount(mfs.Dir())
		mfs.Join(ctx)
		t.Fatal("fuse.Mount returned nil")
	}

	const want = "no such file"
	if got := err.Error(); !strings.Contains(got, want) {
		t.Errorf("Unexpected error: %v", got)
	}
}
<|endoftext|>


<commit_before>// Copyright 2009 The GoMPD Authors. All rights reserved.
// Use of this source code is governed by the MIT
// license that can be found in the LICENSE file.

// This package provides the client side interface to MPD (Music Player Daemon).
// The protocol reference can be found at http://www.musicpd.org/doc/protocol/index.html
package mpd

import (
	"net/textproto"
	"os"
	"strconv"
	"strings"
)

type Client struct {
	text *textproto.Conn
}

type Attrs map[string]string

// Dial connects to MPD listening on address addr (e.g. "127.0.0.1:6600")
// on network network (e.g. "tcp").
func Dial(network, addr string) (c *Client, err os.Error) {
	text, err := textproto.Dial(network, addr)
	if err != nil {
		return nil, err
	}
	line, err := text.ReadLine()
	if err != nil {
		return nil, err
	}
	if line[0:6] != "OK MPD" {
		return nil, textproto.ProtocolError("no greeting")
	}
	return &Client{text: text}, nil
}

// Close terminates the connection with MPD.
func (c *Client) Close() (err os.Error) {
	if c.text != nil {
		c.text.PrintfLine("close")
		err = c.text.Close()
		c.text = nil
	}
	return
}

// Ping sends a no-op message to MPD. It's useful for keeping the connection alive.
func (c *Client) Ping() os.Error {
	return c.okCmd("ping")
}

func (c *Client) readPlaylist() (pls []Attrs, err os.Error) {
	pls = []Attrs{}

	for {
		line, err := c.text.ReadLine()
		if err != nil {
			return nil, err
		}
		if line == "OK" {
			break
		}
		if strings.HasPrefix(line, "file:") { // new song entry begins
			pls = append(pls, Attrs{})
		}
		if len(pls) == 0 {
			return nil, textproto.ProtocolError("unexpected: " + line)
		}
		z := strings.Index(line, ": ")
		if z < 0 {
			return nil, textproto.ProtocolError("can't parse line: " + line)
		}
		key := line[0:z]
		pls[len(pls)-1][key] = line[z+2:]
	}
	return pls, nil
}

func (c *Client) readAttrs() (attrs Attrs, err os.Error) {
	attrs = make(Attrs)
	for {
		line, err := c.text.ReadLine()
		if err != nil {
			return nil, err
		}
		if line == "OK" {
			break
		}
		z := strings.Index(line, ": ")
		if z < 0 {
			return nil, textproto.ProtocolError("can't parse line: " + line)
		}
		key := line[0:z]
		attrs[key] = line[z+2:]
	}
	return
}

// CurrentSong returns information about the current song in the playlist.
func (c *Client) CurrentSong() (Attrs, os.Error) {
	id, err := c.text.Cmd("currentsong")
	if err != nil {
		return nil, err
	}
	c.text.StartResponse(id)
	defer c.text.EndResponse(id)
	return c.readAttrs()
}

// Status returns information about the current status of MPD.
func (c *Client) Status() (Attrs, os.Error) {
	id, err := c.text.Cmd("status")
	if err != nil {
		return nil, err
	}
	c.text.StartResponse(id)
	defer c.text.EndResponse(id)
	return c.readAttrs()
}

func (c *Client) readOKLine() (err os.Error) {
	line, err := c.text.ReadLine()
	if err != nil {
		return
	}
	if line == "OK" {
		return nil
	}
	return textproto.ProtocolError("unexpected response: " + line)
}

func (c *Client) okCmd(format string, args ...interface{}) os.Error {
	id, err := c.text.Cmd(format, args...)
	if err != nil {
		return err
	}
	c.text.StartResponse(id)
	defer c.text.EndResponse(id)
	return c.readOKLine()
}

//
// Playback control
//

// Next plays next song in the playlist.
func (c *Client) Next() os.Error {
	return c.okCmd("next")
}

// Pause pauses playback if pause is true; resumes playback otherwise.
func (c *Client) Pause(pause bool) os.Error {
	if pause {
		return c.okCmd("pause 1")
	}
	return c.okCmd("pause 0")
}

// Play starts playing the song at playlist position pos. If pos is negative,
// start playing at the current position in the playlist.
func (c *Client) Play(pos int) os.Error {
	if pos < 0 {
		c.okCmd("play")
	}
	return c.okCmd("play %d", pos)
}

// PlayId plays the song identified by id. If id is negative, start playing
// at the currect position in playlist.
func (c *Client) PlayId(id int) os.Error {
	if id < 0 {
		return c.okCmd("playid")
	}
	return c.okCmd("playid %d", id)
}

// Previous plays previous song in the playlist.
func (c *Client) Previous() os.Error {
	return c.okCmd("next")
}

// Seek seeks to the position time (in seconds) of the song at playlist position pos.
func (c *Client) Seek(pos, time int) os.Error {
	return c.okCmd("seek %d %d", pos, time)
}

// SeekId is identical to Seek except the song is identified by it's id
// (not position in playlist).
func (c *Client) SeekId(id, time int) os.Error {
	return c.okCmd("seekid %d %d", id, time)
}

// Stop stops playback.
func (c *Client) Stop() os.Error {
	return c.okCmd("stop")
}

//
// Playlist related functions
//

// PlaylistInfo returns attributes for songs in the current playlist. If
// both start and end are negative, it does this for all songs in
// playlist. If end is negative but start is positive, it does it for the
// song at position start. If both start and end are positive, it does it
// for positions in range [start, end).
func (c *Client) PlaylistInfo(start, end int) (pls []Attrs, err os.Error) {
	if start < 0 && end >= 0 {
		return nil, os.NewError("negative start index")
	}
	if start >= 0 && end < 0 {
		id, err := c.text.Cmd("playlistinfo %d", start)
		if err != nil {
			return nil, err
		}
		c.text.StartResponse(id)
		defer c.text.EndResponse(id)
		return c.readPlaylist()
	}
	id, err := c.text.Cmd("playlistinfo")
	if err != nil {
		return nil, err
	}
	c.text.StartResponse(id)
	defer c.text.EndResponse(id)
	pls, err = c.readPlaylist()
	if err != nil || start < 0 || end < 0 {
		return
	}
	return pls[start:end], nil
}

// Delete deletes songs from playlist. If both start and end are positive,
// it deletes those at positions in range [start, end). If end is negative,
// it deletes the song at position start.
func (c *Client) Delete(start, end int) os.Error {
	if start < 0 {
		return os.NewError("negative start index")
	}
	if end < 0 {
		return c.okCmd("delete %d", start)
	}
	return c.okCmd("delete %d %d", start, end)
}

// DeleteId deletes the song identified by id.
func (c *Client) DeleteId(id int) os.Error {
	return c.okCmd("deleteid %d", id)
}

// Add adds the file/directory uri to playlist. Directories add recursively.
func (c *Client) Add(uri string) os.Error {
	return c.okCmd("add %q", uri)
}

// AddId adds the file/directory uri to playlist and returns the identity
// id of the song added. If pos is positive, the song is added to position
// pos.
func (c *Client) AddId(uri string, pos int) (int, os.Error) {
	var id uint
	var err os.Error
	if pos >= 0 {
		id, err = c.text.Cmd("addid %q %d", uri, pos)
	}
	id, err = c.text.Cmd("addid %q", uri)
	if err != nil {
		return -1, err
	}

	c.text.StartResponse(id)
	defer c.text.EndResponse(id)

	attrs, err := c.readAttrs()
	if err != nil {
		return -1, err
	}
	tok, ok := attrs["Id"]
	if !ok {
		return -1, textproto.ProtocolError("addid did not return Id")
	}
	return strconv.Atoi(tok)
}

// Clear clears the current playlist.
func (c *Client) Clear() os.Error {
	return c.okCmd("clear")
}
<commit_msg>gofix error<commit_after>// Copyright 2009 The GoMPD Authors. All rights reserved.
// Use of this source code is governed by the MIT
// license that can be found in the LICENSE file.

// This package provides the client side interface to MPD (Music Player Daemon).
// The protocol reference can be found at http://www.musicpd.org/doc/protocol/index.html
package mpd

import (
	"errors"
	"net/textproto"
	"strconv"
	"strings"
)

type Client struct {
	text *textproto.Conn
}

type Attrs map[string]string

// Dial connects to MPD listening on address addr (e.g. "127.0.0.1:6600")
// on network network (e.g. "tcp").
func Dial(network, addr string) (c *Client, err error) {
	text, err := textproto.Dial(network, addr)
	if err != nil {
		return nil, err
	}
	line, err := text.ReadLine()
	if err != nil {
		return nil, err
	}
	if line[0:6] != "OK MPD" {
		return nil, textproto.ProtocolError("no greeting")
	}
	return &Client{text: text}, nil
}

// Close terminates the connection with MPD.
func (c *Client) Close() (err error) {
	if c.text != nil {
		c.text.PrintfLine("close")
		err = c.text.Close()
		c.text = nil
	}
	return
}

// Ping sends a no-op message to MPD. It's useful for keeping the connection alive.
func (c *Client) Ping() error {
	return c.okCmd("ping")
}

func (c *Client) readPlaylist() (pls []Attrs, err error) {
	pls = []Attrs{}

	for {
		line, err := c.text.ReadLine()
		if err != nil {
			return nil, err
		}
		if line == "OK" {
			break
		}
		if strings.HasPrefix(line, "file:") { // new song entry begins
			pls = append(pls, Attrs{})
		}
		if len(pls) == 0 {
			return nil, textproto.ProtocolError("unexpected: " + line)
		}
		z := strings.Index(line, ": ")
		if z < 0 {
			return nil, textproto.ProtocolError("can't parse line: " + line)
		}
		key := line[0:z]
		pls[len(pls)-1][key] = line[z+2:]
	}
	return pls, nil
}

func (c *Client) readAttrs() (attrs Attrs, err error) {
	attrs = make(Attrs)
	for {
		line, err := c.text.ReadLine()
		if err != nil {
			return nil, err
		}
		if line == "OK" {
			break
		}
		z := strings.Index(line, ": ")
		if z < 0 {
			return nil, textproto.ProtocolError("can't parse line: " + line)
		}
		key := line[0:z]
		attrs[key] = line[z+2:]
	}
	return
}

// CurrentSong returns information about the current song in the playlist.
func (c *Client) CurrentSong() (Attrs, error) {
	id, err := c.text.Cmd("currentsong")
	if err != nil {
		return nil, err
	}
	c.text.StartResponse(id)
	defer c.text.EndResponse(id)
	return c.readAttrs()
}

// Status returns information about the current status of MPD.
func (c *Client) Status() (Attrs, error) {
	id, err := c.text.Cmd("status")
	if err != nil {
		return nil, err
	}
	c.text.StartResponse(id)
	defer c.text.EndResponse(id)
	return c.readAttrs()
}

func (c *Client) readOKLine() (err error) {
	line, err := c.text.ReadLine()
	if err != nil {
		return
	}
	if line == "OK" {
		return nil
	}
	return textproto.ProtocolError("unexpected response: " + line)
}

func (c *Client) okCmd(format string, args ...interface{}) error {
	id, err := c.text.Cmd(format, args...)
	if err != nil {
		return err
	}
	c.text.StartResponse(id)
	defer c.text.EndResponse(id)
	return c.readOKLine()
}

//
// Playback control
//

// Next plays next song in the playlist.
func (c *Client) Next() error {
	return c.okCmd("next")
}

// Pause pauses playback if pause is true; resumes playback otherwise.
func (c *Client) Pause(pause bool) error {
	if pause {
		return c.okCmd("pause 1")
	}
	return c.okCmd("pause 0")
}

// Play starts playing the song at playlist position pos. If pos is negative,
// start playing at the current position in the playlist.
func (c *Client) Play(pos int) error {
	if pos < 0 {
		c.okCmd("play")
	}
	return c.okCmd("play %d", pos)
}

// PlayId plays the song identified by id. If id is negative, start playing
// at the currect position in playlist.
func (c *Client) PlayId(id int) error {
	if id < 0 {
		return c.okCmd("playid")
	}
	return c.okCmd("playid %d", id)
}

// Previous plays previous song in the playlist.
func (c *Client) Previous() error {
	return c.okCmd("next")
}

// Seek seeks to the position time (in seconds) of the song at playlist position pos.
func (c *Client) Seek(pos, time int) error {
	return c.okCmd("seek %d %d", pos, time)
}

// SeekId is identical to Seek except the song is identified by it's id
// (not position in playlist).
func (c *Client) SeekId(id, time int) error {
	return c.okCmd("seekid %d %d", id, time)
}

// Stop stops playback.
func (c *Client) Stop() error {
	return c.okCmd("stop")
}

//
// Playlist related functions
//

// PlaylistInfo returns attributes for songs in the current playlist. If
// both start and end are negative, it does this for all songs in
// playlist. If end is negative but start is positive, it does it for the
// song at position start. If both start and end are positive, it does it
// for positions in range [start, end).
func (c *Client) PlaylistInfo(start, end int) (pls []Attrs, err error) {
	if start < 0 && end >= 0 {
		return nil, errors.New("negative start index")
	}
	if start >= 0 && end < 0 {
		id, err := c.text.Cmd("playlistinfo %d", start)
		if err != nil {
			return nil, err
		}
		c.text.StartResponse(id)
		defer c.text.EndResponse(id)
		return c.readPlaylist()
	}
	id, err := c.text.Cmd("playlistinfo")
	if err != nil {
		return nil, err
	}
	c.text.StartResponse(id)
	defer c.text.EndResponse(id)
	pls, err = c.readPlaylist()
	if err != nil || start < 0 || end < 0 {
		return
	}
	return pls[start:end], nil
}

// Delete deletes songs from playlist. If both start and end are positive,
// it deletes those at positions in range [start, end). If end is negative,
// it deletes the song at position start.
func (c *Client) Delete(start, end int) error {
	if start < 0 {
		return errors.New("negative start index")
	}
	if end < 0 {
		return c.okCmd("delete %d", start)
	}
	return c.okCmd("delete %d %d", start, end)
}

// DeleteId deletes the song identified by id.
func (c *Client) DeleteId(id int) error {
	return c.okCmd("deleteid %d", id)
}

// Add adds the file/directory uri to playlist. Directories add recursively.
func (c *Client) Add(uri string) error {
	return c.okCmd("add %q", uri)
}

// AddId adds the file/directory uri to playlist and returns the identity
// id of the song added. If pos is positive, the song is added to position
// pos.
func (c *Client) AddId(uri string, pos int) (int, error) {
	var id uint
	var err error
	if pos >= 0 {
		id, err = c.text.Cmd("addid %q %d", uri, pos)
	}
	id, err = c.text.Cmd("addid %q", uri)
	if err != nil {
		return -1, err
	}

	c.text.StartResponse(id)
	defer c.text.EndResponse(id)

	attrs, err := c.readAttrs()
	if err != nil {
		return -1, err
	}
	tok, ok := attrs["Id"]
	if !ok {
		return -1, textproto.ProtocolError("addid did not return Id")
	}
	return strconv.Atoi(tok)
}

// Clear clears the current playlist.
func (c *Client) Clear() error {
	return c.okCmd("clear")
}
<|endoftext|>


<commit_before>package testify

import (
	"bufio"
	"fmt"
	"go/format"
	"go/types"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/littledot/mockhiato/lib"
	log "github.com/sirupsen/logrus"
)

const (
	magic              = "// Code generated by Mockhiato. DO NOT EDIT."
	interfaceNameToken = "{interface}"
)

// NewTestifyFormatter creates a new testifyFormatter
func NewTestifyFormatter(config lib.Config) lib.Formatter {
	return &testifyFormatter{
		config: config,
	}
}

type testifyFormatter struct {
	config lib.Config
}

func (r *testifyFormatter) IsMockFile(file *os.File) bool {
	fileScanner := bufio.NewScanner(file)
	for i := 0; i < 6 && fileScanner.Scan(); i++ { // Check first 5 lines or EOF
		line := fileScanner.Text()
		if strings.HasPrefix(line, "//") { // Check 1st comment, which should be the 3rd line
			if line == magic { // Found magic string? File is a mock
				return true
			}
			break
		}
	}
	return false
}

func (r *testifyFormatter) GenerateMocks(project *lib.Project) {
	for _, pack := range project.Packages {
		r.generateMocksForDefinedInterfaces(project, pack)
	}
	r.generateMocksForDependentInterfaces(project, project.DependentPackage)
}

func (r *testifyFormatter) generateMocksForDefinedInterfaces(project *lib.Project, pack *lib.Package) {
	r.generateMock(project, pack.Context.Path(), pack.Context.Name(), pack.Interfaces)
}

func (r *testifyFormatter) generateMocksForDependentInterfaces(project *lib.Project, pack *lib.GeneratedPackage) {
	r.generateMock(project, pack.ContextPath, pack.ContextName, pack.Interfaces)
}

func (r *testifyFormatter) generateMock(project *lib.Project, contextPath, contextName string, interfaces []*lib.Interface) {
	if len(interfaces) == 0 {
		return
	}
	log.Debugf("Generating mocks for %s", contextPath)

	buf := &lib.Bufferw{}
	pm := lib.NewPackageMapper(contextPath)
	pm.RecordDependency(types.NewPackage("github.com/stretchr/testify/mock", "mock"))
	for _, iface := range interfaces {
		pm.IndexInterface(iface.TInterface)
	}

	// Write package
	packageName := contextName
	log.Debugf("Writing package: %s", packageName)
	buf.WriteString(fmt.Sprintf("package %s\n", packageName))

	// Write magic string
	log.Debugf("Writing magic: %s", magic)
	buf.WriteString(magic + "\n")

	// Write imports
	imports := []string{}
	for depPath, depAlias := range pm.PathToAlias {
		imports = append(imports, fmt.Sprintf(`%s "%s"`, depAlias, strings.TrimPrefix(depPath, project.VendorPath+"/")))
	}
	log.Debugf("Writing imports: %s", strings.Join(imports, ", "))
	buf.WriteString("import(\n")
	buf.WriteString(strings.Join(imports, "\n"))
	buf.WriteString("\n)\n")

	for _, iface := range interfaces {

		interfaceName := iface.TObject.Name()
		mockName := strings.Replace(r.config.StructNameFormat, interfaceNameToken, interfaceName, -1)

		// Write constructor
		buf.WriteString(fmt.Sprintf("// New%s creates a new %s\n", mockName, mockName))
		buf.WriteString(fmt.Sprintf("func New%s() *%s {return &%s{}}\n", mockName, mockName, mockName))

		// Write struct
		log.Debugf("Writing struct: %s", mockName)
		buf.WriteString(fmt.Sprintf("// %s implements %s.%s\n", mockName, packageName, interfaceName))
		buf.WriteString(fmt.Sprintf("type %s struct { mock.Mock }\n", mockName))

		for i := 0; i < iface.TInterface.NumMethods(); i++ {
			// Write method
			method := iface.TInterface.Method(i)
			signature := method.Type().(*types.Signature)

			paramNames := []string{}
			paramExprs := []string{}
			for j := 0; j < signature.Params().Len(); j++ {
				param := signature.Params().At(j)
				varName := "p" + strconv.Itoa(j)
				paramNames = append(paramNames, varName)
				paramExprs = append(paramExprs, varName+" "+pm.ObjectTypeString(param))
			}

			if signature.Variadic() { // Variadic method? Replace last parameter's [] with ... ("p1 []int" -> "p1 ...int")
				pos := signature.Params().Len() - 1
				paramExprs[pos] = strings.Replace(paramExprs[pos], "[]", "...", 1)
			}

			returnNames := []string{}
			returnTypes := []string{}
			verifyReturnLines := []string{}
			for j := 0; j < signature.Results().Len(); j++ {
				result := signature.Results().At(j)
				resultTypeString := pm.ObjectTypeString(result)
				varName := "ret" + strconv.Itoa(j)
				returnNames = append(returnNames, varName)
				returnTypes = append(returnTypes, resultTypeString)
				verifyReturnLine := fmt.Sprintf("var %s %s; if a := ret.Get(%d); a != nil { %s = a.(%s) }\n", varName, resultTypeString, j, varName, resultTypeString)
				verifyReturnLines = append(verifyReturnLines, verifyReturnLine)
			}

			log.Debugf("Writing method: %s()", method.Name())
			commentLine := fmt.Sprintf("// %s implements (%s.%s).%s\n", method.Name(), packageName, interfaceName, method.Name())
			signatureLine := fmt.Sprintf("func (r *%s) %s(%s) (%s) {\n", mockName, method.Name(), strings.Join(paramExprs, ", "), strings.Join(returnTypes, ","))
			verifyInvokedLine := fmt.Sprintf("r.Called(%s)\n", strings.Join(paramNames, ", "))
			if signature.Results().Len() > 0 {
				verifyInvokedLine = "ret := " + verifyInvokedLine
			}
			returnLine := fmt.Sprintf("return %s\n}\n", strings.Join(returnNames, ", "))

			buf.WriteString(commentLine)
			buf.WriteString(signatureLine)
			buf.WriteString(verifyInvokedLine)
			for _, verifyReturnLine := range verifyReturnLines {
				buf.WriteString(verifyReturnLine)
			}
			buf.WriteString(returnLine)
		}
	}

	mockPath := filepath.Join(project.GoSrcAbsPath, contextPath, r.config.MockFileName)
	log.Debugf("Creating file: %s", mockPath)
	if err := os.MkdirAll(filepath.Dir(mockPath), 0755); err != nil {
		panic(err)
	}
	mockFile, err := os.Create(mockPath)
	if err != nil {
		panic(err)
	}

	// Format generated code
	sourceCode, err := format.Source(buf.Bytes())
	if err != nil {
		panic(err)
	}
	if _, err := mockFile.Write(sourceCode); err != nil {
		panic(err)
	}
	if err := mockFile.Close(); err != nil {
		panic(err)
	}

	project.GenAbsPaths = append(project.GenAbsPaths, mockPath)
}
<commit_msg>Fixes #20. Move `format.Source()` above file creation.<commit_after>package testify

import (
	"bufio"
	"fmt"
	"go/format"
	"go/types"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/littledot/mockhiato/lib"
	log "github.com/sirupsen/logrus"
)

const (
	magic              = "// Code generated by Mockhiato. DO NOT EDIT."
	interfaceNameToken = "{interface}"
)

// NewTestifyFormatter creates a new testifyFormatter
func NewTestifyFormatter(config lib.Config) lib.Formatter {
	return &testifyFormatter{
		config: config,
	}
}

type testifyFormatter struct {
	config lib.Config
}

func (r *testifyFormatter) IsMockFile(file *os.File) bool {
	fileScanner := bufio.NewScanner(file)
	for i := 0; i < 6 && fileScanner.Scan(); i++ { // Check first 5 lines or EOF
		line := fileScanner.Text()
		if strings.HasPrefix(line, "//") { // Check 1st comment, which should be the 3rd line
			if line == magic { // Found magic string? File is a mock
				return true
			}
			break
		}
	}
	return false
}

func (r *testifyFormatter) GenerateMocks(project *lib.Project) {
	for _, pack := range project.Packages {
		r.generateMocksForDefinedInterfaces(project, pack)
	}
	r.generateMocksForDependentInterfaces(project, project.DependentPackage)
}

func (r *testifyFormatter) generateMocksForDefinedInterfaces(project *lib.Project, pack *lib.Package) {
	r.generateMock(project, pack.Context.Path(), pack.Context.Name(), pack.Interfaces)
}

func (r *testifyFormatter) generateMocksForDependentInterfaces(project *lib.Project, pack *lib.GeneratedPackage) {
	r.generateMock(project, pack.ContextPath, pack.ContextName, pack.Interfaces)
}

func (r *testifyFormatter) generateMock(project *lib.Project, contextPath, contextName string, interfaces []*lib.Interface) {
	if len(interfaces) == 0 {
		return
	}
	log.Debugf("Generating mocks for %s", contextPath)

	buf := &lib.Bufferw{}
	pm := lib.NewPackageMapper(contextPath)
	pm.RecordDependency(types.NewPackage("github.com/stretchr/testify/mock", "mock"))
	for _, iface := range interfaces {
		pm.IndexInterface(iface.TInterface)
	}

	// Write package
	packageName := contextName
	log.Debugf("Writing package: %s", packageName)
	buf.WriteString(fmt.Sprintf("package %s\n", packageName))

	// Write magic string
	log.Debugf("Writing magic: %s", magic)
	buf.WriteString(magic + "\n")

	// Write imports
	imports := []string{}
	for depPath, depAlias := range pm.PathToAlias {
		imports = append(imports, fmt.Sprintf(`%s "%s"`, depAlias, strings.TrimPrefix(depPath, project.VendorPath+"/")))
	}
	log.Debugf("Writing imports: %s", strings.Join(imports, ", "))
	buf.WriteString("import(\n")
	buf.WriteString(strings.Join(imports, "\n"))
	buf.WriteString("\n)\n")

	for _, iface := range interfaces {

		interfaceName := iface.TObject.Name()
		mockName := strings.Replace(r.config.StructNameFormat, interfaceNameToken, interfaceName, -1)

		// Write constructor
		buf.WriteString(fmt.Sprintf("// New%s creates a new %s\n", mockName, mockName))
		buf.WriteString(fmt.Sprintf("func New%s() *%s {return &%s{}}\n", mockName, mockName, mockName))

		// Write struct
		log.Debugf("Writing struct: %s", mockName)
		buf.WriteString(fmt.Sprintf("// %s implements %s.%s\n", mockName, packageName, interfaceName))
		buf.WriteString(fmt.Sprintf("type %s struct { mock.Mock }\n", mockName))

		for i := 0; i < iface.TInterface.NumMethods(); i++ {
			// Write method
			method := iface.TInterface.Method(i)
			signature := method.Type().(*types.Signature)

			paramNames := []string{}
			paramExprs := []string{}
			for j := 0; j < signature.Params().Len(); j++ {
				param := signature.Params().At(j)
				varName := "p" + strconv.Itoa(j)
				paramNames = append(paramNames, varName)
				paramExprs = append(paramExprs, varName+" "+pm.ObjectTypeString(param))
			}

			if signature.Variadic() { // Variadic method? Replace last parameter's [] with ... ("p1 []int" -> "p1 ...int")
				pos := signature.Params().Len() - 1
				paramExprs[pos] = strings.Replace(paramExprs[pos], "[]", "...", 1)
			}

			returnNames := []string{}
			returnTypes := []string{}
			verifyReturnLines := []string{}
			for j := 0; j < signature.Results().Len(); j++ {
				result := signature.Results().At(j)
				resultTypeString := pm.ObjectTypeString(result)
				varName := "ret" + strconv.Itoa(j)
				returnNames = append(returnNames, varName)
				returnTypes = append(returnTypes, resultTypeString)
				verifyReturnLine := fmt.Sprintf("var %s %s; if a := ret.Get(%d); a != nil { %s = a.(%s) }\n", varName, resultTypeString, j, varName, resultTypeString)
				verifyReturnLines = append(verifyReturnLines, verifyReturnLine)
			}

			log.Debugf("Writing method: %s()", method.Name())
			commentLine := fmt.Sprintf("// %s implements (%s.%s).%s\n", method.Name(), packageName, interfaceName, method.Name())
			signatureLine := fmt.Sprintf("func (r *%s) %s(%s) (%s) {\n", mockName, method.Name(), strings.Join(paramExprs, ", "), strings.Join(returnTypes, ","))
			verifyInvokedLine := fmt.Sprintf("r.Called(%s)\n", strings.Join(paramNames, ", "))
			if signature.Results().Len() > 0 {
				verifyInvokedLine = "ret := " + verifyInvokedLine
			}
			returnLine := fmt.Sprintf("return %s\n}\n", strings.Join(returnNames, ", "))

			buf.WriteString(commentLine)
			buf.WriteString(signatureLine)
			buf.WriteString(verifyInvokedLine)
			for _, verifyReturnLine := range verifyReturnLines {
				buf.WriteString(verifyReturnLine)
			}
			buf.WriteString(returnLine)
		}
	}

	// Format generated code
	sourceCode, err := format.Source(buf.Bytes())
	if err != nil {
		log.Debugf("Format code: %s", buf.Bytes())
		panic(err)
	}

	// Flush code to disk
	mockPath := filepath.Join(project.GoSrcAbsPath, contextPath, r.config.MockFileName)
	log.Debugf("Creating file: %s", mockPath)
	if err := os.MkdirAll(filepath.Dir(mockPath), 0755); err != nil {
		panic(err)
	}
	mockFile, err := os.Create(mockPath)
	if err != nil {
		panic(err)
	}
	if _, err := mockFile.Write(sourceCode); err != nil {
		panic(err)
	}
	if err := mockFile.Close(); err != nil {
		panic(err)
	}

	project.GenAbsPaths = append(project.GenAbsPaths, mockPath)
}
<|endoftext|>


<commit_before>package pila

import (
	"encoding/json"
	"fmt"
	"sort"

	"github.com/fern4lvarez/piladb/pkg/uuid"
)

// Database represents a piladb database.
type Database struct {
	// ID is a unique identifier of the database
	ID fmt.Stringer
	// Name of the database
	Name string
	// Pointer to the current piladb instance
	Pila *Pila
	// Stacks associated to Database mapped by their ID
	Stacks map[fmt.Stringer]*Stack
}

// NewDatabase creates a new Database given a name,
// without any link to the piladb instance.
func NewDatabase(name string) *Database {
	stacks := make(map[fmt.Stringer]*Stack)
	return &Database{
		ID:     uuid.New(name),
		Name:   name,
		Stacks: stacks,
	}
}

// CreateStack creates a new Stack, given a name, which is associated
// to the Database.
func (db *Database) CreateStack(name string) fmt.Stringer {
	stack := NewStack(name)
	stack.SetDatabase(db)
	db.Stacks[stack.ID] = stack
	return stack.ID
}

// AddStack adds a given Stack to the Database, returning
// an error if any was found.
func (db *Database) AddStack(stack *Stack) error {
	if stack.Database != nil {
		return fmt.Errorf("stack %v already added to database %v", stack.Name, stack.Database.Name)
	}

	stack.SetDatabase(db)
	if _, ok := db.Stacks[stack.ID]; ok {
		stack.Database = nil
		return fmt.Errorf("database %v already contains stack %v", db.Name, stack.Name)
	}

	db.Stacks[stack.ID] = stack
	return nil
}

// RemoveStack removes a Stack from the Database given an id,
// returning true if it succeeded. It will return false if the
// Stack wasn't added to the Database.
func (db *Database) RemoveStack(id fmt.Stringer) bool {
	stack, ok := db.Stacks[id]
	if !ok {
		return false
	}
	stack.Database = nil
	stack.base = nil
	delete(db.Stacks, id)
	return true
}

// Status returns the status of the Database.
func (db *Database) Status() DatabaseStatus {
	dbs := DatabaseStatus{}
	dbs.ID = db.ID.String()
	dbs.Name = db.Name
	dbs.NumberStacks = len(db.Stacks)

	var ss sort.StringSlice = make([]string, len(db.Stacks))
	n := 0
	for sID := range db.Stacks {
		ss[n] = sID.String()
		n++
	}
	ss.Sort()
	dbs.Stacks = ss

	return dbs
}

// StacksStatus returns the status of the Stacks of Database.
func (db *Database) StacksStatus() StacksStatus {
	var n int
	ss := make([]StackStatus, len(db.Stacks))
	for _, s := range db.Stacks {
		ss[n] = s.Status()
		n++
	}

	status := StacksStatus{}
	status.Stacks = ss
	sort.Sort(status)
	return status
}

// StacksKV returns the status of the Stacks of Database
// in a key-value format.
func (db *Database) StacksKV() StacksKV {
	kv := make(map[string]interface{})
	for _, s := range db.Stacks {
		kv[s.Name] = s.Peek()
	}

	stacksKV := StacksKV{}
	stacksKV.Stacks = kv
	return stacksKV
}

// DatabaseStatus represents the status of a Database.
type DatabaseStatus struct {
	ID           string   `json:"id"`
	Name         string   `json:"name"`
	NumberStacks int      `json:"number_of_stacks"`
	Stacks       []string `json:"stacks,omitempty"`
}

// ToJSON converts a DatabaseStatus into JSON.
func (databaseStatus DatabaseStatus) ToJSON() []byte {
	// Do not check error as the Status type does
	// not contain types that could cause such case.
	// See http://golang.org/src/encoding/json/encode.go?s=5438:5481#L125
	b, _ := json.Marshal(databaseStatus)
	return b
}
<commit_msg>Fixup for PR<commit_after>package pila

import (
	"encoding/json"
	"fmt"
	"sort"

	"github.com/fern4lvarez/piladb/pkg/uuid"
)

// Database represents a piladb database.
type Database struct {
	// ID is a unique identifier of the database
	ID fmt.Stringer
	// Name of the database
	Name string
	// Pointer to the current piladb instance
	Pila *Pila
	// Stacks associated to Database mapped by their ID
	Stacks map[fmt.Stringer]*Stack
}

// NewDatabase creates a new Database given a name,
// without any link to the piladb instance.
func NewDatabase(name string) *Database {
	stacks := make(map[fmt.Stringer]*Stack)
	return &Database{
		ID:     uuid.New(name),
		Name:   name,
		Stacks: stacks,
	}
}

// CreateStack creates a new Stack, given a name, which is associated
// to the Database.
func (db *Database) CreateStack(name string) fmt.Stringer {
	stack := NewStack(name)
	stack.SetDatabase(db)
	db.Stacks[stack.ID] = stack
	return stack.ID
}

// AddStack adds a given Stack to the Database, returning
// an error if any was found.
func (db *Database) AddStack(stack *Stack) error {
	if stack.Database != nil {
		return fmt.Errorf("stack %v already added to database %v", stack.Name, stack.Database.Name)
	}

	stack.SetDatabase(db)
	if _, ok := db.Stacks[stack.ID]; ok {
		stack.Database = nil
		return fmt.Errorf("database %v already contains stack %v", db.Name, stack.Name)
	}

	db.Stacks[stack.ID] = stack
	return nil
}

// RemoveStack removes a Stack from the Database given an id,
// returning true if it succeeded. It will return false if the
// Stack wasn't added to the Database.
func (db *Database) RemoveStack(id fmt.Stringer) bool {
	stack, ok := db.Stacks[id]
	if !ok {
		return false
	}
	stack.Database = nil
	stack.base = nil
	delete(db.Stacks, id)
	return true
}

// Status returns the status of the Database.
func (db *Database) Status() DatabaseStatus {
	dbs := DatabaseStatus{}
	dbs.ID = db.ID.String()
	dbs.Name = db.Name
	dbs.NumberStacks = len(db.Stacks)

	var ss sort.StringSlice = make([]string, len(db.Stacks))
	n := 0
	for sID := range db.Stacks {
		ss[n] = sID.String()
		n++
	}
	ss.Sort()
	dbs.Stacks = ss

	return dbs
}

// StacksStatus returns the status of the Stacks of Database.
func (db *Database) StacksStatus() StacksStatus {
	var n int
	ss := make([]StackStatus, len(db.Stacks))
	for _, s := range db.Stacks {
		ss[n] = s.Status()
		n++
	}

	status := StacksStatus{Stacks: ss}
	sort.Sort(status)
	return status
}

// StacksKV returns the status of the Stacks of Database
// in a key-value format.
func (db *Database) StacksKV() StacksKV {
	kv := make(map[string]interface{})
	for _, s := range db.Stacks {
		kv[s.Name] = s.Peek()
	}

	stacksKV := StacksKV{Stacks: kv}
	return stacksKV
}

// DatabaseStatus represents the status of a Database.
type DatabaseStatus struct {
	ID           string   `json:"id"`
	Name         string   `json:"name"`
	NumberStacks int      `json:"number_of_stacks"`
	Stacks       []string `json:"stacks,omitempty"`
}

// ToJSON converts a DatabaseStatus into JSON.
func (databaseStatus DatabaseStatus) ToJSON() []byte {
	// Do not check error as the Status type does
	// not contain types that could cause such case.
	// See http://golang.org/src/encoding/json/encode.go?s=5438:5481#L125
	b, _ := json.Marshal(databaseStatus)
	return b
}
<|endoftext|>