<commit_before>package models

import (
	"crypto/rand"
	"encoding/base64"
	"errors"
	"fmt"
	"sync"
	"time"
)

type ChangeDirection struct {
	Index     int
	Direction Direction
	Wait      *sync.WaitGroup
}

const (
	Initializing = iota
	InProgress
	Ended
)

type Field struct {
	mu *sync.Mutex //initially used while adding players

	id            string
	Height, Width int
	needed        int //number of players needed to start the game

	Gophers []*Gopher
	Board   [][]bool

	Change chan ChangeDirection
	Remove chan int

	State     int
	cycles    int
	nextCycle interface{}
}

var (
	ErrInProgress = errors.New("game is already in progress")
	ErrMaxPlayers = errors.New("cannot add more than 4 players")
	mapMu         = new(sync.RWMutex)
	activeFields  = make(map[string]*Field)
)

func GetGames() map[string]*Field {
	games := make(map[string]*Field)
	mapMu.RLock()
	defer mapMu.RUnlock()

	for id, field := range activeFields {
		games[id] = field
	}

	return games
}

func GetGame(id string) (*Field, bool) {
	mapMu.RLock()
	defer mapMu.RUnlock()
	field, ok := activeFields[id]

	return field, ok
}

func NewField(height, width int, needed int) (*Field, error) {
	if needed > 4 {
		return nil, ErrMaxPlayers
	}

	bytes := make([]byte, 10)
	rand.Read(bytes)
	id := base64.URLEncoding.EncodeToString(bytes)

	field := &Field{
		id:        id,
		Height:    height,
		needed:    needed,
		Width:     width,
		Board:     make([][]bool, height),
		Change:    make(chan ChangeDirection),
		Remove:    make(chan int),
		mu:        new(sync.Mutex),
		State:     Initializing,
		nextCycle: time.NewTicker(100 * time.Millisecond),
	}

	for i := range field.Board {
		field.Board[i] = make([]bool, width)
	}

	mapMu.Lock()
	defer mapMu.Unlock()
	activeFields[id] = field

	return field, nil
}

func (f *Field) Add(g *Gopher) (int, error) {
	f.mu.Lock()
	defer f.mu.Unlock()

	if f.State == InProgress {
		return 0, ErrInProgress
	}

	f.Gophers = append(f.Gophers, g)
	//set gopher's initial position and direction

	var initPositions = []struct {
		X int
		Y int
		D Direction
	}{
		{f.Width / 2, 0, North},
		{0, f.Height / 2, East},
		{f.Width / 2, f.Height - 1, South},
		{f.Width - 1, f.Height / 2, West},
	}
	for _, pos := range initPositions {
		if !f.Board[pos.X][pos.Y] {
			f.setPos(g, pos.X, pos.Y)
			g.Direction = pos.D
			break
		}
	}

	if len(f.Gophers) == f.needed {
		f.State = Ended
		go f.start()
	}

	return len(f.Gophers) - 1, nil
}

func (f *Field) setPos(g *Gopher, X, Y int) {
	f.Board[X][Y] = true
	g.X = X
	g.Y = Y
}

func (f *Field) remove(gopher *Gopher) {
	for i, currGopher := range f.Gophers {
		if currGopher == gopher {
			f.Gophers, f.Gophers[len(f.Gophers)-1] = append(f.Gophers[:i], f.Gophers[i+1:]...), nil
		}
	}
}

func (f *Field) PrintBoard() {
	for y := f.Height - 1; y >= 0; y-- {
		for x := 0; x < f.Width; x++ {
			if f.Board[x][y] {
				fmt.Print(1)
			} else {
				fmt.Print(0)
			}
			//fmt.Printf("(%d,%d)", x, y)
		}
		fmt.Println()
	}
	fmt.Println("-----")
}

func (f *Field) increment(g *Gopher) bool {
	g.Path = append(g.Path, Coordinate{g.X, g.Y})
	switch g.Direction {
	case North:
		g.Y++
	case South:
		g.Y--
	case East:
		g.X++
	case West:
		g.X--
	}

	if g.X == f.Height || g.Y == f.Width {
		return true
	}
	if f.Board[g.X][g.Y] {
		return true
	}

	f.Board[g.X][g.Y] = true
	g.Score++

	if f.cycles > 10 {
		f.Board[g.Path[0].X][g.Path[0].Y] = false
	} else {
		f.cycles++
	}

	return false
}

func (f *Field) clearPath(g *Gopher) {
	for _, c := range g.Path {
		f.Board[c.X][c.Y] = false
	}
}

func (f *Field) start() {
	channel := make(chan struct{})

	switch f.nextCycle.(type) {
	case *time.Ticker:
		go func() {
			for {
				<-f.nextCycle.(time.Ticker).C
				channel <- struct{}{}
			}
		}()
	default:
		//testing
		go func() {
			for {
				channel <- <-f.nextCycle.(chan struct{})
			}
		}()
	}

	for {
		select {
		case dir := <-f.Change:
			if f.State != InProgress {
				continue
			}
			f.Gophers[dir.Index].Direction = dir.Direction
			dir.Wait.Done()
		case index := <-f.Remove:
			f.remove(f.Gophers[index])
			if len(f.Gophers) == 1 {
				f.end()
				return
			}
		case <-channel:
			for _, gopher := range f.Gophers {
				if f.increment(gopher) {
					//fmt.Println(i, " collided")
					//gopher collided, clear it's path and remove it
					//from the field
					f.clearPath(gopher)
					f.remove(gopher)

					if len(f.Gophers) == 1 {
						f.end()
						return
					}
				}
			}
			//f.PrintBoard()
		}
	}
}

func (f *Field) end() {
	f.mu.Lock()
	f.State = Ended
	f.mu.Unlock()
	if _, ok := f.nextCycle.(*time.Ticker); ok {
		f.nextCycle.(*time.Ticker).Stop()
	}
	mapMu.Lock()
	delete(activeFields, f.id)
	mapMu.RUnlock()
}

func (f *Field) broadcast() {
	paths := make(map[int][]Coordinate)
	for i, gopher := range f.Gophers {
		copy(paths[i], gopher.Path)
		gopher.Paths <- paths
	}
}
<commit_msg>broadacst(): Use map[string][]Coordinate<commit_after>package models

import (
	"crypto/rand"
	"encoding/base64"
	"errors"
	"fmt"
	"log"
	"strconv"
	"sync"
	"time"
)

type ChangeDirection struct {
	Index     int
	Direction Direction
	Wait      *sync.WaitGroup
}

const (
	Initializing = iota
	InProgress
	Ended
)

type Field struct {
	mu *sync.Mutex //initially used while adding players

	ID            string
	Height, Width int
	needed        int //number of players needed to start the game

	Gophers []*Gopher
	Board   [][]bool

	Change chan ChangeDirection
	Remove chan int

	State     int
	cycles    int
	nextCycle interface{}
}

var (
	ErrInProgress = errors.New("game is already in progress")
	ErrMaxPlayers = errors.New("cannot add more than 4 players")
	mapMu         = new(sync.RWMutex)
	activeFields  = make(map[string]*Field)
)

func GetGames() map[string]*Field {
	games := make(map[string]*Field)
	mapMu.RLock()
	defer mapMu.RUnlock()

	for id, field := range activeFields {
		games[id] = field
	}

	return games
}

func GetGame(id string) (*Field, bool) {
	mapMu.RLock()
	defer mapMu.RUnlock()
	field, ok := activeFields[id]

	return field, ok
}

func NewField(height, width int, needed int) (*Field, error) {
	if needed > 4 {
		return nil, ErrMaxPlayers
	}

	bytes := make([]byte, 10)
	rand.Read(bytes)
	id := base64.URLEncoding.EncodeToString(bytes)

	field := &Field{
		ID:        id,
		Height:    height,
		needed:    needed,
		Width:     width,
		Board:     make([][]bool, height),
		Change:    make(chan ChangeDirection),
		Remove:    make(chan int),
		mu:        new(sync.Mutex),
		State:     Initializing,
		nextCycle: time.NewTicker(100 * time.Millisecond),
	}

	for i := range field.Board {
		field.Board[i] = make([]bool, width)
	}

	mapMu.Lock()
	defer mapMu.Unlock()
	activeFields[id] = field

	return field, nil
}

func (f *Field) Add(g *Gopher) (int, error) {
	f.mu.Lock()
	defer f.mu.Unlock()

	if f.State == InProgress {
		return 0, ErrInProgress
	}

	f.Gophers = append(f.Gophers, g)
	//set gopher's initial position and direction

	var initPositions = []struct {
		X int
		Y int
		D Direction
	}{
		{f.Width / 2, 0, North},
		{0, f.Height / 2, East},
		{f.Width / 2, f.Height - 1, South},
		{f.Width - 1, f.Height / 2, West},
	}
	for _, pos := range initPositions {
		if !f.Board[pos.X][pos.Y] {
			f.setPos(g, pos.X, pos.Y)
			g.Direction = pos.D
			break
		}
	}

	if len(f.Gophers) == f.needed {
		f.State = Ended
		go f.start()
	}

	return len(f.Gophers) - 1, nil
}

func (f *Field) setPos(g *Gopher, X, Y int) {
	f.Board[X][Y] = true
	g.X = X
	g.Y = Y
}

func (f *Field) remove(gopher *Gopher) {
	for i, currGopher := range f.Gophers {
		if currGopher == gopher {
			f.Gophers, f.Gophers[len(f.Gophers)-1] = append(f.Gophers[:i], f.Gophers[i+1:]...), nil
		}
	}
}

func (f *Field) PrintBoard() {
	for y := f.Height - 1; y >= 0; y-- {
		for x := 0; x < f.Width; x++ {
			if f.Board[x][y] {
				fmt.Print(1)
			} else {
				fmt.Print(0)
			}
			//fmt.Printf("(%d,%d)", x, y)
		}
		fmt.Println()
	}
	fmt.Println("-----")
}

func (f *Field) increment(g *Gopher) bool {
	g.Path = append(g.Path, Coordinate{g.X, g.Y})
	switch g.Direction {
	case North:
		g.Y++
	case South:
		g.Y--
	case East:
		g.X++
	case West:
		g.X--
	}

	if g.X == f.Height || g.Y == f.Width {
		return true
	}
	if f.Board[g.X][g.Y] {
		return true
	}

	f.Board[g.X][g.Y] = true
	g.Score++

	if f.cycles > 10 {
		f.Board[g.Path[0].X][g.Path[0].Y] = false
	} else {
		f.cycles++
	}

	return false
}

func (f *Field) clearPath(g *Gopher) {
	for _, c := range g.Path {
		f.Board[c.X][c.Y] = false
	}
}

func (f *Field) start() {
	log.Printf("%s: Starting main game loop.", f.ID)
	tick := time.NewTicker(10 * time.Millisecond)
	defer tick.Stop()
	for {
		select {
		case dir := <-f.Change:
			if f.State != InProgress {
				continue
			}
			f.Gophers[dir.Index].Direction = dir.Direction
			dir.Wait.Done()
		case index := <-f.Remove:
			f.remove(f.Gophers[index])
			if len(f.Gophers) == 1 || len(f.Gophers) == 0 {
				f.end()
				return
			}
		case <-tick.C:
			for _, gopher := range f.Gophers {
				if f.increment(gopher) {
					//fmt.Println(i, " collided")
					//gopher collided, clear it's path and remove it
					//from the field
					f.clearPath(gopher)
					f.remove(gopher)

					if len(f.Gophers) == 1 {
						f.end()
						return
					}
				}
			}
			f.broadcast()
			//f.PrintBoard()
		}
	}
}

func (f *Field) end() {
	log.Printf("%s: ended", f.ID)
	f.mu.Lock()
	f.State = Ended
	f.mu.Unlock()
	if _, ok := f.nextCycle.(*time.Ticker); ok {
		f.nextCycle.(*time.Ticker).Stop()
	}
	mapMu.Lock()
	delete(activeFields, f.ID)
	mapMu.Unlock()
}

func (f *Field) broadcast() {
	paths := make(map[string][]Coordinate)
	for i, gopher := range f.Gophers {
		index := strconv.Itoa(i)
		for _, c := range gopher.Path {
			paths[index] = append(paths[index], c)
		}

		gopher.Paths <- paths
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"io/ioutil"
	"log"
	"os"
	"strconv"
	"syscall"
)

var (
	path    string
	pidfile string
)

func init() {
	flag.StringVar(&path, "path", "/Database/branch/master/watch/com.docker.driver.amd64-linux.node/etc.node/docker.node/daemon.json.node/tree.live", "path of the file to watch")
	flag.StringVar(&pidfile, "pidfile", "/run/docker.pid", "pidfile for process to signal")
}

func main() {
	log.SetFlags(0)
	flag.Parse()

	watch, err := os.Open(path)
	if err != nil {
		log.Fatalln("Failed to open file", path, err)
	}
	// 43 bytes is the record size of the watch
	buf := make([]byte, 43)
	for {
		_, err := watch.Read(buf)
		if err != nil {
			log.Fatalln("Error reading watch file", err)
		}
		bytes, err := ioutil.ReadFile(pidfile)
		pidstring := string(bytes[:])
		if err != nil {
			pid, err := strconv.Atoi(pidstring)
			if err != nil {
				syscall.Kill(pid, syscall.SIGHUP)
			}
		}
	}
}
<commit_msg>reorder test<commit_after>package main

import (
	"flag"
	"io/ioutil"
	"log"
	"os"
	"strconv"
	"syscall"
)

var (
	path    string
	pidfile string
)

func init() {
	flag.StringVar(&path, "path", "/Database/branch/master/watch/com.docker.driver.amd64-linux.node/etc.node/docker.node/daemon.json.node/tree.live", "path of the file to watch")
	flag.StringVar(&pidfile, "pidfile", "/run/docker.pid", "pidfile for process to signal")
}

func main() {
	log.SetFlags(0)
	flag.Parse()

	watch, err := os.Open(path)
	if err != nil {
		log.Fatalln("Failed to open file", path, err)
	}
	// 43 bytes is the record size of the watch
	buf := make([]byte, 43)
	for {
		_, err := watch.Read(buf)
		if err != nil {
			log.Fatalln("Error reading watch file", err)
		}
		bytes, err := ioutil.ReadFile(pidfile)
		if err != nil {
			pidstring := string(bytes[:])
			pid, err := strconv.Atoi(pidstring)
			if err != nil {
				syscall.Kill(pid, syscall.SIGHUP)
			}
		}
	}
}
<|endoftext|>


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

package models

import (
	"path"
	"time"

	git "github.com/gogits/git"
)

type RepoFile struct {
	Id      *git.Oid
	Type    int
	Name    string
	Path    string
	Message string
	Created time.Time
}

func (f *RepoFile) IsFile() bool {
	return f.Type == git.FileModeBlob || f.Type == git.FileModeBlobExec
}

func (f *RepoFile) IsDir() bool {
	return f.Type == git.FileModeTree
}

func GetReposFiles(userName, reposName, branchName, rpath string) ([]*RepoFile, error) {
	f := RepoPath(userName, reposName)

	repo, err := git.OpenRepository(f)
	if err != nil {
		return nil, err
	}

	ref, err := repo.LookupReference("refs/heads/" + branchName)
	if err != nil {
		return nil, err
	}

	lastCommit, err := repo.LookupCommit(ref.Oid)
	if err != nil {
		return nil, err
	}

	var repodirs []*RepoFile
	var repofiles []*RepoFile
	lastCommit.Tree.Walk(func(dirname string, entry *git.TreeEntry) int {
		if dirname == rpath {
			switch entry.Filemode {
			case git.FileModeBlob, git.FileModeBlobExec:
				repofiles = append(repofiles, &RepoFile{
					entry.Id,
					entry.Filemode,
					entry.Name,
					path.Join(dirname, entry.Name),
					lastCommit.Message(),
					lastCommit.Committer.When,
				})
			case git.FileModeTree:
				repodirs = append(repodirs, &RepoFile{
					entry.Id,
					entry.Filemode,
					entry.Name,
					path.Join(dirname, entry.Name),
					lastCommit.Message(),
					lastCommit.Committer.When,
				})
			}
		}
		return 0
	})

	return append(repodirs, repofiles...), nil
}
<commit_msg>refactoring for Repofile<commit_after>// Copyright 2014 The Gogs Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package models

import (
	"path"
	"time"

	git "github.com/gogits/git"
)

type RepoFile struct {
	*git.TreeEntry
	Path    string
	Message string
	Created time.Time
}

func GetReposFiles(userName, reposName, branchName, rpath string) ([]*RepoFile, error) {
	f := RepoPath(userName, reposName)

	repo, err := git.OpenRepository(f)
	if err != nil {
		return nil, err
	}

	ref, err := repo.LookupReference("refs/heads/" + branchName)
	if err != nil {
		return nil, err
	}

	lastCommit, err := repo.LookupCommit(ref.Oid)
	if err != nil {
		return nil, err
	}

	var repodirs []*RepoFile
	var repofiles []*RepoFile
	lastCommit.Tree.Walk(func(dirname string, entry *git.TreeEntry) int {
		if dirname == rpath {
			switch entry.Filemode {
			case git.FileModeBlob, git.FileModeBlobExec:
				repofiles = append(repofiles, &RepoFile{
					entry,
					path.Join(dirname, entry.Name),
					lastCommit.Message(),
					lastCommit.Committer.When,
				})
			case git.FileModeTree:
				repodirs = append(repodirs, &RepoFile{
					entry,
					path.Join(dirname, entry.Name),
					lastCommit.Message(),
					lastCommit.Committer.When,
				})
			}
		}
		return 0
	})

	return append(repodirs, repofiles...), nil
}
<|endoftext|>


<commit_before>package home

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

	"github.com/belogik/goes"
	"github.com/go-martini/martini"

	"github.com/firstrow/logvoyage/common"
	"github.com/firstrow/logvoyage/web/context"
	"github.com/firstrow/logvoyage/web/widgets"
)

const (
	timeLayout = "2006/01/02 15:04" // Users input time format
	perPage    = 100
)

type DateTimeRange struct {
	Start string
	Stop  string
}

func (this *DateTimeRange) IsValid() bool {
	return this.Start != "" || this.Stop != ""
}

// Represents search request to perform in ES
type SearchRequest struct {
	Text      string   // test to search
	Indexes   []string // ES indexeses to perform search
	Types     []string // search types
	Size      int      // home much objects ES must return
	From      int      // how much objects should ES skip from first
	TimeRange DateTimeRange
}

func buildSearchRequest(text string, indexes []string, types []string, size int, from int, datetime DateTimeRange) SearchRequest {
	return SearchRequest{
		Text:      text,
		Indexes:   indexes,
		From:      from,
		Types:     types,
		Size:      perPage,
		TimeRange: datetime,
	}
}

// Detects time range from request and returns
// elastic compatible format string
func buildTimeRange(req *http.Request) DateTimeRange {
	var timeRange DateTimeRange

	switch req.URL.Query().Get("time") {
	case "15m":
		timeRange.Start = "now-15m"
	case "30m":
		timeRange.Start = "now-30m"
	case "60m":
		timeRange.Start = "now-60m"
	case "12h":
		timeRange.Start = "now-12h"
	case "24h":
		timeRange.Start = "now-24h"
	case "week":
		timeRange.Start = "now-7d"
	case "custom":
		timeStart, err := time.Parse(timeLayout, req.URL.Query().Get("time_start"))
		if err == nil {
			timeRange.Start = timeStart.Format(time.RFC3339)
		}
		timeStop, err := time.Parse(timeLayout, req.URL.Query().Get("time_stop"))
		if err == nil {
			timeRange.Stop = timeStop.Format(time.RFC3339)
		}
	}

	return timeRange
}

// Search logs in elastic.
func search(searchRequest SearchRequest) (goes.Response, error) {
	conn := common.GetConnection()

	var query = map[string]interface{}{
		"from": searchRequest.From,
		"size": searchRequest.Size,
		"sort": map[string]string{
			"datetime": "desc",
		},
	}

	if len(searchRequest.Text) > 0 {
		strconv.Quote(searchRequest.Text)
		query["query"] = map[string]interface{}{
			"query_string": map[string]string{
				"default_field": "message",
				"query":         searchRequest.Text,
			},
		}
	}

	// Build time range query
	if searchRequest.TimeRange.IsValid() {
		datetime := make(map[string]string)
		if searchRequest.TimeRange.Start != "" {
			datetime["gte"] = searchRequest.TimeRange.Start
		}
		if searchRequest.TimeRange.Stop != "" {
			datetime["lte"] = searchRequest.TimeRange.Stop
		}
		query["filter"] = map[string]interface{}{
			"range": map[string]interface{}{
				"datetime": datetime,
			},
		}
	}

	extraArgs := make(url.Values, 0)
	searchResults, err := conn.Search(query, searchRequest.Indexes, searchRequest.Types, extraArgs)

	if err != nil {
		return &goes.Response{}, errors.New("No records found.")
	} else {
		return &searchResults, nil
	}
}

// This function handles two routes "/" and "/project/:id"
func ProjectSearch(ctx *context.Context, params martini.Params) {
	var types []string
	var project *common.Project

	query_text := ctx.Request.URL.Query().Get("q")
	selected_types := ctx.Request.URL.Query()["types"]

	// Project scope
	if _, err := params["id"]; err {
		project, err := ctx.User.GetProject(params["id"])
		if err != nil {
			ctx.HTML("shared/error", context.ViewData{
				"message": "Project not found",
			})
			return
		}
		if len(project.Types) == 0 {
			ctx.HTML("home/empty_project", context.ViewData{
				"project": project,
			})
			return
		}
		if len(selected_types) > 0 {
			types = selected_types
		} else {
			types = project.Types
		}
	}

	// Pagination
	pagination := widgets.NewPagination(ctx.Request)
	pagination.SetPerPage(perPage)

	// Load records
	searchRequest := buildSearchRequest(
		query_text,
		[]string{ctx.User.GetIndexName()},
		types,
		pagination.GetPerPage(),
		pagination.DetectFrom(),
		buildTimeRange(ctx.Request),
	)
	// Search data in elastic
	data, _ := search(searchRequest)

	pagination.SetTotalRecords(data.Hits.Total)

	var viewName string
	viewData := context.ViewData{
		"project":    project,
		"logs":       data.Hits.Hits,
		"total":      data.Hits.Total,
		"took":       data.Took,
		"types":      types,
		"time":       ctx.Request.URL.Query().Get("time"),
		"time_start": ctx.Request.URL.Query().Get("time_start"),
		"time_stop":  ctx.Request.URL.Query().Get("time_stop"),
		"query_text": query_text,
		"pagination": pagination,
	}

	if ctx.Request.Header.Get("X-Requested-With") == "XMLHttpRequest" {
		viewName = "home/table"
	} else {
		viewName = "home/index"
	}

	ctx.HTML(viewName, viewData)
}
<commit_msg>return pointer<commit_after>package home

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

	"github.com/belogik/goes"
	"github.com/go-martini/martini"

	"github.com/firstrow/logvoyage/common"
	"github.com/firstrow/logvoyage/web/context"
	"github.com/firstrow/logvoyage/web/widgets"
)

const (
	timeLayout = "2006/01/02 15:04" // Users input time format
	perPage    = 100
)

type DateTimeRange struct {
	Start string
	Stop  string
}

func (this *DateTimeRange) IsValid() bool {
	return this.Start != "" || this.Stop != ""
}

// Represents search request to perform in ES
type SearchRequest struct {
	Text      string   // test to search
	Indexes   []string // ES indexeses to perform search
	Types     []string // search types
	Size      int      // home much objects ES must return
	From      int      // how much objects should ES skip from first
	TimeRange DateTimeRange
}

func buildSearchRequest(text string, indexes []string, types []string, size int, from int, datetime DateTimeRange) SearchRequest {
	return SearchRequest{
		Text:      text,
		Indexes:   indexes,
		From:      from,
		Types:     types,
		Size:      perPage,
		TimeRange: datetime,
	}
}

// Detects time range from request and returns
// elastic compatible format string
func buildTimeRange(req *http.Request) DateTimeRange {
	var timeRange DateTimeRange

	switch req.URL.Query().Get("time") {
	case "15m":
		timeRange.Start = "now-15m"
	case "30m":
		timeRange.Start = "now-30m"
	case "60m":
		timeRange.Start = "now-60m"
	case "12h":
		timeRange.Start = "now-12h"
	case "24h":
		timeRange.Start = "now-24h"
	case "week":
		timeRange.Start = "now-7d"
	case "custom":
		timeStart, err := time.Parse(timeLayout, req.URL.Query().Get("time_start"))
		if err == nil {
			timeRange.Start = timeStart.Format(time.RFC3339)
		}
		timeStop, err := time.Parse(timeLayout, req.URL.Query().Get("time_stop"))
		if err == nil {
			timeRange.Stop = timeStop.Format(time.RFC3339)
		}
	}

	return timeRange
}

// Search logs in elastic.
func search(searchRequest SearchRequest) (goes.Response, error) {
	conn := common.GetConnection()

	var query = map[string]interface{}{
		"from": searchRequest.From,
		"size": searchRequest.Size,
		"sort": map[string]string{
			"datetime": "desc",
		},
	}

	if len(searchRequest.Text) > 0 {
		strconv.Quote(searchRequest.Text)
		query["query"] = map[string]interface{}{
			"query_string": map[string]string{
				"default_field": "message",
				"query":         searchRequest.Text,
			},
		}
	}

	// Build time range query
	if searchRequest.TimeRange.IsValid() {
		datetime := make(map[string]string)
		if searchRequest.TimeRange.Start != "" {
			datetime["gte"] = searchRequest.TimeRange.Start
		}
		if searchRequest.TimeRange.Stop != "" {
			datetime["lte"] = searchRequest.TimeRange.Stop
		}
		query["filter"] = map[string]interface{}{
			"range": map[string]interface{}{
				"datetime": datetime,
			},
		}
	}

	extraArgs := make(url.Values, 0)
	searchResults, err := conn.Search(query, searchRequest.Indexes, searchRequest.Types, extraArgs)

	if err != nil {
		return &goes.Response{}, errors.New("No records found.")
	} else {
		return searchResults, nil
	}
}

// This function handles two routes "/" and "/project/:id"
func ProjectSearch(ctx *context.Context, params martini.Params) {
	var types []string
	var project *common.Project

	query_text := ctx.Request.URL.Query().Get("q")
	selected_types := ctx.Request.URL.Query()["types"]

	// Project scope
	if _, err := params["id"]; err {
		project, err := ctx.User.GetProject(params["id"])
		if err != nil {
			ctx.HTML("shared/error", context.ViewData{
				"message": "Project not found",
			})
			return
		}
		if len(project.Types) == 0 {
			ctx.HTML("home/empty_project", context.ViewData{
				"project": project,
			})
			return
		}
		if len(selected_types) > 0 {
			types = selected_types
		} else {
			types = project.Types
		}
	}

	// Pagination
	pagination := widgets.NewPagination(ctx.Request)
	pagination.SetPerPage(perPage)

	// Load records
	searchRequest := buildSearchRequest(
		query_text,
		[]string{ctx.User.GetIndexName()},
		types,
		pagination.GetPerPage(),
		pagination.DetectFrom(),
		buildTimeRange(ctx.Request),
	)
	// Search data in elastic
	data, _ := search(searchRequest)

	pagination.SetTotalRecords(data.Hits.Total)

	var viewName string
	viewData := context.ViewData{
		"project":    project,
		"logs":       data.Hits.Hits,
		"total":      data.Hits.Total,
		"took":       data.Took,
		"types":      types,
		"time":       ctx.Request.URL.Query().Get("time"),
		"time_start": ctx.Request.URL.Query().Get("time_start"),
		"time_stop":  ctx.Request.URL.Query().Get("time_stop"),
		"query_text": query_text,
		"pagination": pagination,
	}

	if ctx.Request.Header.Get("X-Requested-With") == "XMLHttpRequest" {
		viewName = "home/table"
	} else {
		viewName = "home/index"
	}

	ctx.HTML(viewName, viewData)
}
<|endoftext|>


<commit_before>package restful

// http://jsr311.java.net/nonav/releases/1.1/spec/spec.html

import (
	"encoding/xml"
	"github.com/emicklei/go-restful/wadl"
	"log"
	"net/http"
)

type Dispatcher interface {
	Routes() []Route
	RootPath() string
	//	rootRegEx
}

// Collection of registered Dispatchers that can handle Http requests
var webServices = []Dispatcher{}

// Register a new Dispatcher
func Add(service Dispatcher) {
	webServices = append(webServices, service)
}

// Dispatch the incoming Http Request to a matching Dispatcher.
// Matching algoritm is conform http://jsr311.java.net/nonav/releases/1.1/spec/spec.html

func Dispatch(httpWriter http.ResponseWriter, httpRequest *http.Request) {
	dispatcher, finalMatch, err := detectDispatcher(httpRequest.URL.Path, webServices)
	if err != nil {
		httpWriter.WriteHeader(http.StatusNotFound)
	}
	routes := selectRoutes(dispatcher, finalMatch)
	route, detected := detectRoute(routes, httpWriter, httpRequest)
	if detected {
		route.dispatch(httpWriter, httpRequest)
	}
}

// Hook my Dispatch function as the standard Http handler
func init() {
	log.Printf("restful: register the Dispatch function to the Default Http handlers.\n")
	http.HandleFunc("/", Dispatch)
}

// Return the api in XML
func Wadl(base string) string {
	resources := wadl.Resources{Base: base}
	for _, eachWebService := range webServices {
		for _, eachRoute := range eachWebService.Routes() {
			response := wadl.Response{}
			for _, mimeType := range eachRoute.Produces {
				response.AddRepresentation(wadl.Representation{MediaType: mimeType})
			}
			request := wadl.Request{}
			for _, mimeType := range eachRoute.Consumes {
				request.AddRepresentation(wadl.Representation{MediaType: mimeType})
			}
			method := wadl.Method{Name: eachRoute.Method, Response: response, Request: request, Doc: eachRoute.Doc}
			resource := wadl.Resource{Path: eachRoute.Path, Method: method}
			resources.AddResource(resource)
		}
	}
	app := wadl.Application{Resources: resources}
	bytes, _ := xml.MarshalIndent(app, "", " ")
	return string(bytes)
}
<commit_msg>fixed 404 bug in Dispatch<commit_after>package restful

// http://jsr311.java.net/nonav/releases/1.1/spec/spec.html

import (
	"encoding/xml"
	"github.com/emicklei/go-restful/wadl"
	"log"
	"net/http"
)

type Dispatcher interface {
	Routes() []Route
	RootPath() string
	//	rootRegEx
}

// Collection of registered Dispatchers that can handle Http requests
var webServices = []Dispatcher{}

// Register a new Dispatcher
func Add(service Dispatcher) {
	webServices = append(webServices, service)
}

// Dispatch the incoming Http Request to a matching Dispatcher.
// Matching algoritm is conform http://jsr311.java.net/nonav/releases/1.1/spec/spec.html

func Dispatch(httpWriter http.ResponseWriter, httpRequest *http.Request) {
	dispatcher, finalMatch, err := detectDispatcher(httpRequest.URL.Path, webServices)
	if err != nil {
		httpWriter.WriteHeader(http.StatusNotFound)
		return
	}
	routes := selectRoutes(dispatcher, finalMatch)
	route, detected := detectRoute(routes, httpWriter, httpRequest)
	if detected {
		route.dispatch(httpWriter, httpRequest)
	}
}

// Hook my Dispatch function as the standard Http handler
func init() {
	log.Printf("restful: register the Dispatch function to the Default Http handlers.\n")
	http.HandleFunc("/", Dispatch)
}

// Return the api in XML
func Wadl(base string) string {
	resources := wadl.Resources{Base: base}
	for _, eachWebService := range webServices {
		for _, eachRoute := range eachWebService.Routes() {
			response := wadl.Response{}
			for _, mimeType := range eachRoute.Produces {
				response.AddRepresentation(wadl.Representation{MediaType: mimeType})
			}
			request := wadl.Request{}
			for _, mimeType := range eachRoute.Consumes {
				request.AddRepresentation(wadl.Representation{MediaType: mimeType})
			}
			method := wadl.Method{Name: eachRoute.Method, Response: response, Request: request, Doc: eachRoute.Doc}
			resource := wadl.Resource{Path: eachRoute.Path, Method: method}
			resources.AddResource(resource)
		}
	}
	app := wadl.Application{Resources: resources}
	bytes, _ := xml.MarshalIndent(app, "", " ")
	return string(bytes)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"github.com/monsti/form"
	"github.com/monsti/l10n"
	"github.com/monsti/rpc/client"
	"github.com/monsti/template"
	"github.com/monsti/util"
	"log"
	"os"
)

type settings struct {
	// Absolute paths to used directories.
	Directories struct {
		// HTML Templates
		Templates string
		// Locales, i.e. the gettext machine objects (.mo)
		Locales string
	}
}

var renderer template.Renderer
var logger *log.Logger

func handle(req client.Request, res *client.Response, c client.Connection) {
	switch req.Action {
	case "edit":
		edit(req, res, c)
	default:
		view(req, res, c)
	}
}

type editFormData struct {
	Title string
	Image string
}

func edit(req client.Request, res *client.Response, c client.Connection) {
	G := l10n.UseCatalog(req.Session.Locale)
	data := editFormData{}
	form := form.NewForm(&data, form.Fields{
		"Title": form.Field{G("Title"), "", form.Required(G("Required.")), nil},
		"Image": form.Field{G("Image"), "", nil, new(form.FileWidget)}})
	switch req.Method {
	case "GET":
		data.Title = req.Node.Title
	case "POST":
		imageData, imgerr := c.GetFileData("Image")
		if form.Fill(c.GetFormData()) {
			if imgerr == nil {
				node := req.Node
				node.Title = data.Title
				c.UpdateNode(node)
				c.WriteNodeData(req.Node.Path, "image.data", string(imageData))
				res.Redirect = req.Node.Path
				return
			}
			logger.Println("Image upload failed: ", imgerr)
			form.AddError("Image", G("There was a problem with your image upload."))
		}
	default:
		panic("Request method not supported: " + req.Method)
	}
	fmt.Fprint(res, renderer.Render("edit/image",
		template.Context{"Form": form.RenderData()},
		req.Session.Locale, ""))
}

func view(req client.Request, res *client.Response, c client.Connection) {
	body := c.GetNodeData(req.Node.Path, "image.data")
	res.Raw = true
	res.Write(body)
}

func main() {
	logger = log.New(os.Stderr, "monsti-image", log.LstdFlags)
	flag.Parse()
	cfgPath := util.GetConfigPath("image", flag.Arg(0))
	var settings settings
	if err := util.ParseYAML(cfgPath, &settings); err != nil {
		logger.Fatal("Could not load monsti-image configuration file: ", err)
	}
	l10n.Setup("monsti-image", settings.Directories.Locales)
	renderer.Root = settings.Directories.Templates
	client.NewConnection("monsti-image", logger).Serve(handle)
}
<commit_msg>Fix import path for l10n and template packages.<commit_after>package main

import (
	"flag"
	"fmt"
	"github.com/monsti/form"
	"github.com/monsti/util/l10n"
	"github.com/monsti/rpc/client"
	"github.com/monsti/util/template"
	"github.com/monsti/util"
	"log"
	"os"
)

type settings struct {
	// Absolute paths to used directories.
	Directories struct {
		// HTML Templates
		Templates string
		// Locales, i.e. the gettext machine objects (.mo)
		Locales string
	}
}

var renderer template.Renderer
var logger *log.Logger

func handle(req client.Request, res *client.Response, c client.Connection) {
	switch req.Action {
	case "edit":
		edit(req, res, c)
	default:
		view(req, res, c)
	}
}

type editFormData struct {
	Title string
	Image string
}

func edit(req client.Request, res *client.Response, c client.Connection) {
	G := l10n.UseCatalog(req.Session.Locale)
	data := editFormData{}
	form := form.NewForm(&data, form.Fields{
		"Title": form.Field{G("Title"), "", form.Required(G("Required.")), nil},
		"Image": form.Field{G("Image"), "", nil, new(form.FileWidget)}})
	switch req.Method {
	case "GET":
		data.Title = req.Node.Title
	case "POST":
		imageData, imgerr := c.GetFileData("Image")
		if form.Fill(c.GetFormData()) {
			if imgerr == nil {
				node := req.Node
				node.Title = data.Title
				c.UpdateNode(node)
				c.WriteNodeData(req.Node.Path, "image.data", string(imageData))
				res.Redirect = req.Node.Path
				return
			}
			logger.Println("Image upload failed: ", imgerr)
			form.AddError("Image", G("There was a problem with your image upload."))
		}
	default:
		panic("Request method not supported: " + req.Method)
	}
	fmt.Fprint(res, renderer.Render("edit/image",
		template.Context{"Form": form.RenderData()},
		req.Session.Locale, ""))
}

func view(req client.Request, res *client.Response, c client.Connection) {
	body := c.GetNodeData(req.Node.Path, "image.data")
	res.Raw = true
	res.Write(body)
}

func main() {
	logger = log.New(os.Stderr, "monsti-image", log.LstdFlags)
	flag.Parse()
	cfgPath := util.GetConfigPath("image", flag.Arg(0))
	var settings settings
	if err := util.ParseYAML(cfgPath, &settings); err != nil {
		logger.Fatal("Could not load monsti-image configuration file: ", err)
	}
	l10n.Setup("monsti-image", settings.Directories.Locales)
	renderer.Root = settings.Directories.Templates
	client.NewConnection("monsti-image", logger).Serve(handle)
}
<|endoftext|>


<commit_before>package chaos

import (
	"context"
	"fmt"
	"log"
	"os"
	"os/exec"
	"reflect"
	"testing"
	"time"

	"github.com/davecgh/go-spew/spew"
)

// TODO: cleanup when ctrl-C go test (teardown all containers)

var tracker *Tracker

func TestMain(m *testing.M) {
	ctx, cancelFunc := context.WithCancel(context.Background())
	cmd := exec.CommandContext(ctx, path("docker/launch.sh"), "docker-chaos")
	cmd.Env = append(cmd.Env, "MT_CLUSTER_MIN_AVAILABLE_SHARDS=12")

	var err error
	tracker, err = NewTracker(cmd, false, false, "launch-stdout", "launch-stderr")
	if err != nil {
		log.Fatal(err)
	}

	err = cmd.Start()
	if err != nil {
		log.Fatal(err)
	}

	retcode := m.Run()

	fmt.Println("stopping the docker-compose stack...")
	cancelFunc()
	if err := cmd.Wait(); err != nil {
		log.Printf("ERROR: could not cleanly shutdown running docker-compose command: %s", err)
		retcode = 1
	}

	os.Exit(retcode)
}

func TestClusterStartup(t *testing.T) {
	matchers := []Matcher{
		{Str: "metrictank0_1.*metricIndex initialized.*starting data consumption$"},
		{Str: "metrictank1_1.*metricIndex initialized.*starting data consumption$"},
		{Str: "metrictank2_1.*metricIndex initialized.*starting data consumption$"},
		{Str: "metrictank3_1.*metricIndex initialized.*starting data consumption$"},
		{Str: "metrictank4_1.*metricIndex initialized.*starting data consumption$"},
		{Str: "metrictank5_1.*metricIndex initialized.*starting data consumption$"},
		{Str: "grafana.*Initializing HTTP Server.*:3000"},
	}
	ch := tracker.Match(matchers)
	select {
	case <-ch:
	case <-time.After(time.Second * 40):
		postAnnotation("TestClusterStartup:FAIL")
		t.Fatal("timed out while waiting for all metrictank instances to come up")
	}
}

func TestClusterBaseIngestWorkload(t *testing.T) {
	postAnnotation("TestClusterBaseIngestWorkload:begin")

	go fakeMetrics(t)

	suc6, resp := retryGraphite("perSecond(metrictank.stats.docker-cluster.*.input.kafka-mdm.metrics_received.counter32)", "-8s", 18, func(resp response) bool {
		exp := []string{
			"perSecond(metrictank.stats.docker-cluster.metrictank0.input.kafka-mdm.metrics_received.counter32)",
			"perSecond(metrictank.stats.docker-cluster.metrictank1.input.kafka-mdm.metrics_received.counter32)",
			"perSecond(metrictank.stats.docker-cluster.metrictank2.input.kafka-mdm.metrics_received.counter32)",
			"perSecond(metrictank.stats.docker-cluster.metrictank3.input.kafka-mdm.metrics_received.counter32)",
			"perSecond(metrictank.stats.docker-cluster.metrictank4.input.kafka-mdm.metrics_received.counter32)",
			"perSecond(metrictank.stats.docker-cluster.metrictank5.input.kafka-mdm.metrics_received.counter32)",
		}
		// avg rate must be 4 (metrics ingested per second by each instance)
		return validateTargets(exp)(resp) && validatorAvgWindowed(8, 4)(resp)
	})
	if !suc6 {
		postAnnotation("TestClusterBaseIngestWorkload:FAIL")
		t.Fatalf("cluster did not reach a state where each MT instance receives 4 points per second. last response was: %s", spew.Sdump(resp))
	}

	suc6, resp = retryMT("sum(some.id.of.a.metric.*)", "-16s", 20, validateCorrect(12))
	if !suc6 {
		postAnnotation("TestClusterBaseIngestWorkload:FAIL")
		t.Fatalf("could not query correct result set. sum of 12 series, each valued 1, should result in 12.  last response was: %s", spew.Sdump(resp))
	}
}

func TestQueryWorkload(t *testing.T) {
	postAnnotation("TestQueryWorkload:begin")

	results := checkMT([]int{6060, 6061, 6062, 6063, 6064, 6065}, "sum(some.id.of.a.metric.*)", "-14s", time.Minute, 6000, validateCorrect(12))

	exp := checkResults{
		valid:   []int{6000},
		empty:   0,
		timeout: 0,
		other:   0,
	}
	if !reflect.DeepEqual(exp, results) {
		postAnnotation("TestQueryWorkload:FAIL")
		t.Fatalf("expected only correct results. got %s", spew.Sdump(results))
	}
}

// TestIsolateOneInstance tests what happens during the isolation of one instance, when min-available-shards is 12
// this should happen:
// at all times, all queries to all of the remaining nodes should be successful
// since they have at least 1 instance running for each shard.
// the isolated shard should either return correct replies, or errors (in two cases: when it marks any shards as down,
// but also before it does, but fails to get data via clustered requests from peers)
//. TODO: in production do we stop querying isolated peers?
func TestIsolateOneInstance(t *testing.T) {
	postAnnotation("TestIsolateOneInstance:begin")
	numReqMt4 := 1200

	mt4ResultsChan := make(chan checkResults, 1)
	otherResultsChan := make(chan checkResults, 1)

	go func() {
		mt4ResultsChan <- checkMT([]int{6064}, "sum(some.id.of.a.metric.*)", "-15s", time.Minute, numReqMt4, validateCorrect(12), validateCode(503))
	}()
	go func() {
		otherResultsChan <- checkMT([]int{6060, 6061, 6062, 6063, 6065}, "sum(some.id.of.a.metric.*)", "-15s", time.Minute, 6000, validateCorrect(12))
	}()

	// now go ahead and isolate for 30s
	isolate("metrictank4", "30s", "metrictank0", "metrictank1", "metrictank2", "metrictank3", "metrictank5")

	// collect results of the minute long experiment
	mt4Results := <-mt4ResultsChan
	otherResults := <-otherResultsChan

	// validate results of isolated node
	if mt4Results.valid[0]+mt4Results.valid[1] != numReqMt4 {
		t.Fatalf("expected mt4 to return either correct or erroring responses. got %s", spew.Sdump(mt4Results))
	}
	if mt4Results.valid[1] < numReqMt4*30/100 {
		// the instance is completely down for 30s of the 60s experiment run, but we allow some slack
		t.Fatalf("expected at least 30%% of all mt4 results to succeed. got %s", spew.Sdump(mt4Results))
	}

	// validate results of other cluster nodes
	exp := checkResults{
		valid:   []int{6000},
		empty:   0,
		timeout: 0,
		other:   0,
	}
	if !reflect.DeepEqual(exp, otherResults) {
		postAnnotation("TestIsolateOneInstance:FAIL")
		t.Fatalf("expected only correct results for all cluster nodes. got %s", spew.Sdump(otherResults))
	}
}

func TestHang(t *testing.T) {
	postAnnotation("TestHang:begin")
	t.Log("whatever happens, keep hanging for now, so that we can query grafana dashboards still")
	var ch chan struct{}
	<-ch
}

// maybe useful in the future, test also clean exit and rejoin like so:
//stop("metrictank4")
//time.AfterFunc(30*time.Second, func() {
//	start("metrictank4")
//})
<commit_msg>isolate bi-directionally<commit_after>package chaos

import (
	"context"
	"fmt"
	"log"
	"os"
	"os/exec"
	"reflect"
	"testing"
	"time"

	"github.com/davecgh/go-spew/spew"
)

// TODO: cleanup when ctrl-C go test (teardown all containers)

var tracker *Tracker

func TestMain(m *testing.M) {
	ctx, cancelFunc := context.WithCancel(context.Background())
	cmd := exec.CommandContext(ctx, path("docker/launch.sh"), "docker-chaos")
	cmd.Env = append(cmd.Env, "MT_CLUSTER_MIN_AVAILABLE_SHARDS=12")

	var err error
	tracker, err = NewTracker(cmd, false, false, "launch-stdout", "launch-stderr")
	if err != nil {
		log.Fatal(err)
	}

	err = cmd.Start()
	if err != nil {
		log.Fatal(err)
	}

	retcode := m.Run()

	fmt.Println("stopping the docker-compose stack...")
	cancelFunc()
	if err := cmd.Wait(); err != nil {
		log.Printf("ERROR: could not cleanly shutdown running docker-compose command: %s", err)
		retcode = 1
	}

	os.Exit(retcode)
}

func TestClusterStartup(t *testing.T) {
	matchers := []Matcher{
		{Str: "metrictank0_1.*metricIndex initialized.*starting data consumption$"},
		{Str: "metrictank1_1.*metricIndex initialized.*starting data consumption$"},
		{Str: "metrictank2_1.*metricIndex initialized.*starting data consumption$"},
		{Str: "metrictank3_1.*metricIndex initialized.*starting data consumption$"},
		{Str: "metrictank4_1.*metricIndex initialized.*starting data consumption$"},
		{Str: "metrictank5_1.*metricIndex initialized.*starting data consumption$"},
		{Str: "grafana.*Initializing HTTP Server.*:3000"},
	}
	ch := tracker.Match(matchers)
	select {
	case <-ch:
	case <-time.After(time.Second * 40):
		postAnnotation("TestClusterStartup:FAIL")
		t.Fatal("timed out while waiting for all metrictank instances to come up")
	}
}

func TestClusterBaseIngestWorkload(t *testing.T) {
	postAnnotation("TestClusterBaseIngestWorkload:begin")

	go fakeMetrics(t)

	suc6, resp := retryGraphite("perSecond(metrictank.stats.docker-cluster.*.input.kafka-mdm.metrics_received.counter32)", "-8s", 18, func(resp response) bool {
		exp := []string{
			"perSecond(metrictank.stats.docker-cluster.metrictank0.input.kafka-mdm.metrics_received.counter32)",
			"perSecond(metrictank.stats.docker-cluster.metrictank1.input.kafka-mdm.metrics_received.counter32)",
			"perSecond(metrictank.stats.docker-cluster.metrictank2.input.kafka-mdm.metrics_received.counter32)",
			"perSecond(metrictank.stats.docker-cluster.metrictank3.input.kafka-mdm.metrics_received.counter32)",
			"perSecond(metrictank.stats.docker-cluster.metrictank4.input.kafka-mdm.metrics_received.counter32)",
			"perSecond(metrictank.stats.docker-cluster.metrictank5.input.kafka-mdm.metrics_received.counter32)",
		}
		// avg rate must be 4 (metrics ingested per second by each instance)
		return validateTargets(exp)(resp) && validatorAvgWindowed(8, 4)(resp)
	})
	if !suc6 {
		postAnnotation("TestClusterBaseIngestWorkload:FAIL")
		t.Fatalf("cluster did not reach a state where each MT instance receives 4 points per second. last response was: %s", spew.Sdump(resp))
	}

	suc6, resp = retryMT("sum(some.id.of.a.metric.*)", "-16s", 20, validateCorrect(12))
	if !suc6 {
		postAnnotation("TestClusterBaseIngestWorkload:FAIL")
		t.Fatalf("could not query correct result set. sum of 12 series, each valued 1, should result in 12.  last response was: %s", spew.Sdump(resp))
	}
}

func TestQueryWorkload(t *testing.T) {
	postAnnotation("TestQueryWorkload:begin")

	results := checkMT([]int{6060, 6061, 6062, 6063, 6064, 6065}, "sum(some.id.of.a.metric.*)", "-14s", time.Minute, 6000, validateCorrect(12))

	exp := checkResults{
		valid:   []int{6000},
		empty:   0,
		timeout: 0,
		other:   0,
	}
	if !reflect.DeepEqual(exp, results) {
		postAnnotation("TestQueryWorkload:FAIL")
		t.Fatalf("expected only correct results. got %s", spew.Sdump(results))
	}
}

// TestIsolateOneInstance tests what happens during the isolation of one instance, when min-available-shards is 12
// this should happen:
// at all times, all queries to all of the remaining nodes should be successful
// since they have at least 1 instance running for each shard.
// the isolated shard should either return correct replies, or errors (in two cases: when it marks any shards as down,
// but also before it does, but fails to get data via clustered requests from peers)
//. TODO: in production do we stop querying isolated peers?
func TestIsolateOneInstance(t *testing.T) {
	postAnnotation("TestIsolateOneInstance:begin")
	numReqMt4 := 1200

	mt4ResultsChan := make(chan checkResults, 1)
	otherResultsChan := make(chan checkResults, 1)

	go func() {
		mt4ResultsChan <- checkMT([]int{6064}, "sum(some.id.of.a.metric.*)", "-15s", time.Minute, numReqMt4, validateCorrect(12), validateCode(503))
	}()
	go func() {
		otherResultsChan <- checkMT([]int{6060, 6061, 6062, 6063, 6065}, "sum(some.id.of.a.metric.*)", "-15s", time.Minute, 6000, validateCorrect(12))
	}()

	// now go ahead and isolate for 30s
	isolate("metrictank4", "30s", "metrictank0", "metrictank1", "metrictank2", "metrictank3", "metrictank5")
	isolate("metrictank0", "30s", "metrictank4")
	isolate("metrictank1", "30s", "metrictank4")
	isolate("metrictank2", "30s", "metrictank4")
	isolate("metrictank3", "30s", "metrictank4")
	isolate("metrictank5", "30s", "metrictank4")

	// collect results of the minute long experiment
	mt4Results := <-mt4ResultsChan
	otherResults := <-otherResultsChan

	// validate results of isolated node
	if mt4Results.valid[0]+mt4Results.valid[1] != numReqMt4 {
		t.Fatalf("expected mt4 to return either correct or erroring responses. got %s", spew.Sdump(mt4Results))
	}
	if mt4Results.valid[1] < numReqMt4*30/100 {
		// the instance is completely down for 30s of the 60s experiment run, but we allow some slack
		t.Fatalf("expected at least 30%% of all mt4 results to succeed. got %s", spew.Sdump(mt4Results))
	}

	// validate results of other cluster nodes
	exp := checkResults{
		valid:   []int{6000},
		empty:   0,
		timeout: 0,
		other:   0,
	}
	if !reflect.DeepEqual(exp, otherResults) {
		postAnnotation("TestIsolateOneInstance:FAIL")
		t.Fatalf("expected only correct results for all cluster nodes. got %s", spew.Sdump(otherResults))
	}
}

func TestHang(t *testing.T) {
	postAnnotation("TestHang:begin")
	t.Log("whatever happens, keep hanging for now, so that we can query grafana dashboards still")
	var ch chan struct{}
	<-ch
}

// maybe useful in the future, test also clean exit and rejoin like so:
//stop("metrictank4")
//time.AfterFunc(30*time.Second, func() {
//	start("metrictank4")
//})
<|endoftext|>


<commit_before>// This package implements executable for starting and preconfiguring
// single webrocket server node.
//
// Copyright (C) 2011 by Krzysztof Kowalik <chris@nu7hat.ch>
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
package main

import (
	stepper "../gostepper"
	"../webrocket"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"path/filepath"
	"syscall"
	"time"
)

// Configuration variables.
var (
	// The backend endpoint bind address.
	BackendAddr string
	// The websocket endpoint bind address.
	WebsocketAddr string
	// The admin endpoint bind address.
	AdminAddr string
	// Custom node name.
	NodeName string
	// A path to the websocket endpoint certificate file.
	CertFile string
	// A path to the websocket endpoint key file.
	KeyFile string
	// A path to the The storage directory.
	StorageDir string
	// Shall server be daemonized?
	Daemon bool
)

var (
	// The WebRocket main context.
	ctx *webrocket.Context
	// A stepper instance.
	s stepper.Stepper
)

func init() {
	flag.StringVar(&WebsocketAddr, "websocket-addr", ":8080", "websocket endpoint address")
	flag.StringVar(&BackendAddr, "backend-addr", ":8081", "backend endpoint address")
	flag.StringVar(&AdminAddr, "admin-addr", ":8082", "admin endpoint address")
	flag.StringVar(&NodeName, "node-name", "", "name of the node")
	flag.StringVar(&CertFile, "cert", "", "path to server certificate")
	flag.StringVar(&KeyFile, "key", "", "private key")
	flag.StringVar(&StorageDir, "storage-dir", "/var/lib/webrocket", "path to webrocket's internal data-store")
	flag.BoolVar(&Daemon, "daemon", false, "run in the background")
	flag.Parse()

	StorageDir, _ = filepath.Abs(StorageDir)
}

// SetupContext initializes global WebRocket context, loads configuration
// and all the vhosts data, and generates thean access cookie if its necessary.
func SetupContext() {
	s.Start("Initializing context")
	ctx = webrocket.NewContext()
	if err := ctx.SetStorageDir(StorageDir); err != nil {
		s.Fail(err.Error(), true)
	}
	if NodeName != "" {
		if err := ctx.SetNodeName(NodeName); err != nil {
			s.Fail(err.Error(), true)
		}
	}
	s.Ok()
	s.Start("Locking node")
	if err := ctx.Lock(); err != nil {
		s.Fail(err.Error(), true)
	}
	s.Ok()
	s.Start("Loading configuration")
	if err := ctx.Load(); err != nil {
		s.Fail(err.Error(), true)
	}
	s.Ok()
	s.Start("Generating cookie")
	if err := ctx.GenerateCookie(false); err != nil {
		s.Fail(err.Error(), true)
	}
	s.Ok()
}

// SetupEndpoint is a helper to configure and run a WebRocket endpoint.
//
// kind - The name of the endpoint.
// e    - The enpoint to be started.
//
func SetupEndpoint(kind string, e webrocket.Endpoint) {
	go func() {
		var err error
		s.Start("Starting %s", kind)
		if CertFile != "" && KeyFile != "" {
			err = e.ListenAndServeTLS(CertFile, KeyFile)
		} else {
			err = e.ListenAndServe()
		}
		if err != nil {
			s.Fail(err.Error(), true)
		}
	}()
	for !e.IsAlive() {
		<-time.After(500 * time.Nanosecond)
	}
	s.Ok()
}

// SignalTrap configures a handlers for various system signals, i.a.
// it stops the context and cleans everything up when the app is interrupted.
func SignalTrap() {
	for sig := range signal.Incoming {
		if usig, ok := sig.(os.UnixSignal); ok {
			switch usig {
			case os.SIGQUIT, os.SIGINT:
				fmt.Printf("\n\033[33mInterrupted\033[0m\n")
				if ctx != nil {
					fmt.Printf("\n")
					s.Start("Cleaning up")
					if err := ctx.Kill(); err != nil {
						s.Fail(err.Error(), true)
					}
					s.Ok()
				}
				os.Exit(0)
			case os.SIGTSTP:
				syscall.Kill(syscall.Getpid(), syscall.SIGSTOP)
			case os.SIGHUP:
				// TODO: reload configuration
			}
		}
	}
}

// Daeminize forks the application into a background process and exits
// from the main one.
func Daemonize() int {
	// TODO: ...
	return 0
}

// DisplayAsciiArt as you can see it displays this amazing ASCII art
// spaceship drawing.
func DisplayAsciiArt() {
	fmt.Printf("\n")
	fmt.Printf(AsciiRocket)
	fmt.Printf("WebRocket v%s\n", webrocket.Version())
	fmt.Printf("Copyright (C) 2011-2012 by Krzysztof Kowalik and folks at Cubox.\n")
	fmt.Printf("Released under the AGPL. See http://webrocket.io/ for details.\n\n")
}

// DisplaySystemSettings shows all the information about the running
// webrocket node.
func DisplaySystemSettings() {
	fmt.Printf("\n")
	fmt.Printf("Node               : %s\n", ctx.NodeName())
	fmt.Printf("Cookie             : %s\n", ctx.Cookie())
	fmt.Printf("Data store dir     : %s\n", ctx.StorageDir())
	fmt.Printf("Websocket endpoint : ws://%s\n", WebsocketAddr)
	fmt.Printf("Backend endpoint   : wr://%s\n", BackendAddr)
	fmt.Printf("Admin endpoint     : http://%s\n", AdminAddr)

	fmt.Printf("\n\033[32mWebRocket has been launched!\033[0m\n")
}

func main() {
	if Daemon {
		Daemonize()
	}
	DisplayAsciiArt()
	SetupContext()
	SetupEndpoint("backend endpoint", ctx.NewBackendEndpoint(BackendAddr))
	SetupEndpoint("websocket endpoint", ctx.NewWebsocketEndpoint(WebsocketAddr))
	SetupEndpoint("admin endpoint", ctx.NewAdminEndpoint(AdminAddr))
	DisplaySystemSettings()
	SignalTrap()
}

const AsciiRocket = `` +
	`            /\                                                                    \n` +
	`      ,    /  \      o               .        ___---___                    .      \n` +
	`          /    \            .              .--\        --.     .     .         .  \n` +
	`         /______\                        ./.;_.\     __/~ \.                      \n` +
	`   .    |        |                      /;  / '-'  __\    . \                     \n` +
	`        |        |    .        .       / ,--'     / .   .;   \        |           \n` +
	`        |________|                    | .|       /       __   |      -O-       .  \n` +
	`        |________|                   |__/    __ |  . ;   \ | . |      |           \n` +
	`       /|   ||   |\                  |      /  \\_    . ;| \___|                  \n` +
	`      / |   ||   | \    .    o       |      \  .~\\___,--'     |           .      \n` +
	`     /  |   ||   |  \                 |     | . ; ~~~~\_    __|                   \n` +
	`    /___|:::||:::|___\   |             \    \   .  .  ; \  /_/   .                \n` +
	`        |::::::::|      -O-        .    \   /         . |  ~/                  .  \n` +
	`         \::::::/         |    .          ~\ \   .      /  /~          o          \n` +
	`   o      ||__||       .                   ~--___ ; ___--~                        \n` +
	`            ||                        .          ---         .              .     \n` +
	`            ''                                                                    \n`
<commit_msg>fixed asciiart<commit_after>// This package implements executable for starting and preconfiguring
// single webrocket server node.
//
// Copyright (C) 2011 by Krzysztof Kowalik <chris@nu7hat.ch>
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
package main

import (
	stepper "../gostepper"
	"../webrocket"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"path/filepath"
	"syscall"
	"time"
)

// Configuration variables.
var (
	// The backend endpoint bind address.
	BackendAddr string
	// The websocket endpoint bind address.
	WebsocketAddr string
	// The admin endpoint bind address.
	AdminAddr string
	// Custom node name.
	NodeName string
	// A path to the websocket endpoint certificate file.
	CertFile string
	// A path to the websocket endpoint key file.
	KeyFile string
	// A path to the The storage directory.
	StorageDir string
	// Shall server be daemonized?
	Daemon bool
)

var (
	// The WebRocket main context.
	ctx *webrocket.Context
	// A stepper instance.
	s stepper.Stepper
)

func init() {
	flag.StringVar(&WebsocketAddr, "websocket-addr", ":8080", "websocket endpoint address")
	flag.StringVar(&BackendAddr, "backend-addr", ":8081", "backend endpoint address")
	flag.StringVar(&AdminAddr, "admin-addr", ":8082", "admin endpoint address")
	flag.StringVar(&NodeName, "node-name", "", "name of the node")
	flag.StringVar(&CertFile, "cert", "", "path to server certificate")
	flag.StringVar(&KeyFile, "key", "", "private key")
	flag.StringVar(&StorageDir, "storage-dir", "/var/lib/webrocket", "path to webrocket's internal data-store")
	flag.BoolVar(&Daemon, "daemon", false, "run in the background")
	flag.Parse()

	StorageDir, _ = filepath.Abs(StorageDir)
}

// SetupContext initializes global WebRocket context, loads configuration
// and all the vhosts data, and generates thean access cookie if its necessary.
func SetupContext() {
	s.Start("Initializing context")
	ctx = webrocket.NewContext()
	if err := ctx.SetStorageDir(StorageDir); err != nil {
		s.Fail(err.Error(), true)
	}
	if NodeName != "" {
		if err := ctx.SetNodeName(NodeName); err != nil {
			s.Fail(err.Error(), true)
		}
	}
	s.Ok()
	s.Start("Locking node")
	if err := ctx.Lock(); err != nil {
		s.Fail(err.Error(), true)
	}
	s.Ok()
	s.Start("Loading configuration")
	if err := ctx.Load(); err != nil {
		s.Fail(err.Error(), true)
	}
	s.Ok()
	s.Start("Generating cookie")
	if err := ctx.GenerateCookie(false); err != nil {
		s.Fail(err.Error(), true)
	}
	s.Ok()
}

// SetupEndpoint is a helper to configure and run a WebRocket endpoint.
//
// kind - The name of the endpoint.
// e    - The enpoint to be started.
//
func SetupEndpoint(kind string, e webrocket.Endpoint) {
	go func() {
		var err error
		s.Start("Starting %s", kind)
		if CertFile != "" && KeyFile != "" {
			err = e.ListenAndServeTLS(CertFile, KeyFile)
		} else {
			err = e.ListenAndServe()
		}
		if err != nil {
			s.Fail(err.Error(), true)
		}
	}()
	for !e.IsAlive() {
		<-time.After(500 * time.Nanosecond)
	}
	s.Ok()
}

// SignalTrap configures a handlers for various system signals, i.a.
// it stops the context and cleans everything up when the app is interrupted.
func SignalTrap() {
	for sig := range signal.Incoming {
		if usig, ok := sig.(os.UnixSignal); ok {
			switch usig {
			case os.SIGQUIT, os.SIGINT:
				fmt.Printf("\n\033[33mInterrupted\033[0m\n")
				if ctx != nil {
					fmt.Printf("\n")
					s.Start("Cleaning up")
					if err := ctx.Kill(); err != nil {
						s.Fail(err.Error(), true)
					}
					s.Ok()
				}
				os.Exit(0)
			case os.SIGTSTP:
				syscall.Kill(syscall.Getpid(), syscall.SIGSTOP)
			case os.SIGHUP:
				// TODO: reload configuration
			}
		}
	}
}

// Daeminize forks the application into a background process and exits
// from the main one.
func Daemonize() int {
	// TODO: ...
	return 0
}

// DisplayAsciiArt as you can see it displays this amazing ASCII art
// spaceship drawing.
func DisplayAsciiArt() {
	fmt.Printf("\n")
	fmt.Printf(AsciiRocket)
	fmt.Printf("WebRocket v%s\n", webrocket.Version())
	fmt.Printf("Copyright (C) 2011-2012 by Krzysztof Kowalik and folks at Cubox.\n")
	fmt.Printf("Released under the AGPL. See http://webrocket.io/ for details.\n\n")
}

// DisplaySystemSettings shows all the information about the running
// webrocket node.
func DisplaySystemSettings() {
	fmt.Printf("\n")
	fmt.Printf("Node               : %s\n", ctx.NodeName())
	fmt.Printf("Cookie             : %s\n", ctx.Cookie())
	fmt.Printf("Data store dir     : %s\n", ctx.StorageDir())
	fmt.Printf("Websocket endpoint : ws://%s\n", WebsocketAddr)
	fmt.Printf("Backend endpoint   : wr://%s\n", BackendAddr)
	fmt.Printf("Admin endpoint     : http://%s\n", AdminAddr)

	fmt.Printf("\n\033[32mWebRocket has been launched!\033[0m\n")
}

func main() {
	if Daemon {
		Daemonize()
	}
	DisplayAsciiArt()
	SetupContext()
	SetupEndpoint("backend endpoint", ctx.NewBackendEndpoint(BackendAddr))
	SetupEndpoint("websocket endpoint", ctx.NewWebsocketEndpoint(WebsocketAddr))
	SetupEndpoint("admin endpoint", ctx.NewAdminEndpoint(AdminAddr))
	DisplaySystemSettings()
	SignalTrap()
}

const AsciiRocket = `` +
	`            /\                                                                    ` + "\n" +
	`      ,    /  \      o               .        ___---___                    .      ` + "\n" +
	`          /    \            .              .--\        --.     .     .         .  ` + "\n" +
	`         /______\                        ./.;_.\     __/~ \.                      ` + "\n" +
	`   .    |        |                      /;  / '-'  __\    . \                     ` + "\n" +
	`        |        |    .        .       / ,--'     / .   .;   \        |           ` + "\n" +
	`        |________|                    | .|       /       __   |      -O-       .  ` + "\n" +
	`        |________|                   |__/    __ |  . ;   \ | . |      |           ` + "\n" +
	`       /|   ||   |\                  |      /  \\_    . ;| \___|                  ` + "\n" +
	`      / |   ||   | \    .    o       |      \  .~\\___,--'     |           .      ` + "\n" +
	`     /  |   ||   |  \                 |     | . ; ~~~~\_    __|                   ` + "\n" +
	`    /___|:::||:::|___\   |             \    \   .  .  ; \  /_/   .                ` + "\n" +
	`        |::::::::|      -O-        .    \   /         . |  ~/                  .  ` + "\n" +
	`         \::::::/         |    .          ~\ \   .      /  /~          o          ` + "\n" +
	`   o      ||__||       .                   ~--___ ; ___--~                        ` + "\n" +
	`            ||                        .          ---         .              .     ` + "\n" +
	`            ''                                                                    ` + "\n"
<|endoftext|>


<commit_before>package msgpack

import (
	"errors"
	"io"
	"math"
)

var (
	// ErrLongStringOrBinary is the long string or binary error.
	ErrLongStringOrBinary = errors.New("msgpack: long string or binary")

	// ErrIllegalSize is the illegal array or map size error.
	ErrIllegalSize = errors.New("msgpack: illegal array or map size")
)

// Encoder writes values in MessagePack format.
type Encoder struct {
	buf         [32]byte
	w           io.Writer
	writeString func(string) (int, error)
	err         error // permanent error
}

// NewEncoder allocates and initializes a new Unpacker.
func NewEncoder(w io.Writer) *Encoder {
	e := &Encoder{
		w: w,
	}
	if ws, ok := w.(interface {
		WriteString(string) (int, error)
	}); ok {
		e.writeString = ws.WriteString
	} else {
		e.writeString = e.writeStringUnopt
	}
	return e
}

func (e *Encoder) writeStringUnopt(s string) (int, error) {
	if len(s) <= len(e.buf) {
		copy(e.buf[:], s)
		return e.w.Write(e.buf[:len(s)])
	}
	return e.w.Write([]byte(s))
}

type numCodes struct {
	c8  byte
	c16 byte
	c32 byte
	c64 byte
}

var (
	stringLenEncodings = &numCodes{
		c8:  string8Code,
		c16: string16Code,
		c32: string32Code,
		c64: 0,
	}
	binaryLenEncodings = &numCodes{
		c8:  binary8Code,
		c16: binary16Code,
		c32: binary32Code,
		c64: 0,
	}
	arrayLenEncodings = &numCodes{
		c8:  0,
		c16: array16Code,
		c32: array32Code,
		c64: 0,
	}
	mapLenEncodings = &numCodes{
		c8:  0,
		c16: map16Code,
		c32: map32Code,
		c64: 0,
	}
	extLenEncodings = &numCodes{
		c8:  ext8Code,
		c16: ext16Code,
		c32: ext32Code,
		c64: 0,
	}
	uintEncodings = &numCodes{
		c8:  uint8Code,
		c16: uint16Code,
		c32: uint32Code,
		c64: uint64Code,
	}
)

func (e *Encoder) encodeNum(fc *numCodes, v uint64) []byte {
	switch {
	case fc.c8 != 0 && v <= math.MaxUint8:
		e.buf[0] = fc.c8
		e.buf[1] = byte(v)
		return e.buf[:2]
	case v <= math.MaxUint16:
		e.buf[0] = fc.c16
		e.buf[1] = byte(v >> 8)
		e.buf[2] = byte(v)
		return e.buf[:3]
	case v <= math.MaxUint32:
		e.buf[0] = fc.c32
		e.buf[1] = byte(v >> 24)
		e.buf[2] = byte(v >> 16)
		e.buf[3] = byte(v >> 8)
		e.buf[4] = byte(v)
		return e.buf[:5]
	default:
		e.buf[0] = fc.c64
		e.buf[1] = byte(v >> 56)
		e.buf[2] = byte(v >> 48)
		e.buf[3] = byte(v >> 40)
		e.buf[4] = byte(v >> 32)
		e.buf[5] = byte(v >> 24)
		e.buf[6] = byte(v >> 16)
		e.buf[7] = byte(v >> 8)
		e.buf[8] = byte(v)
		return e.buf[:9]
	}
}

// PackBool writes a Bool value to the MessagePack stream.
func (e *Encoder) PackBool(b bool) error {
	if b {
		e.buf[0] = trueCode
	} else {
		e.buf[0] = falseCode
	}
	_, err := e.w.Write(e.buf[:1])
	return err
}

// PackInt packs an Int value to the MessagePack stream.
func (e *Encoder) PackInt(v int64) error {
	var b []byte
	switch {
	case 0 <= v && v <= math.MaxInt8:
		e.buf[0] = byte(v)
		b = e.buf[:1]
	case v > 0:
		// Pack as unsigned for compatibility with other encoders.
		b = e.encodeNum(uintEncodings, uint64(v))
	case v >= -32:
		e.buf[0] = byte(v)
		b = e.buf[:1]
	case v >= math.MinInt8:
		e.buf[0] = int8Code
		e.buf[1] = byte(v)
		b = e.buf[:2]
	case v >= math.MinInt16:
		e.buf[0] = int16Code
		e.buf[1] = byte(v >> 8)
		e.buf[2] = byte(v)
		b = e.buf[:3]
	case v >= math.MinInt32:
		e.buf[0] = int32Code
		e.buf[1] = byte(v >> 24)
		e.buf[2] = byte(v >> 16)
		e.buf[3] = byte(v >> 8)
		e.buf[4] = byte(v)
		b = e.buf[:5]
	default:
		e.buf[0] = int64Code
		e.buf[1] = byte(v >> 56)
		e.buf[2] = byte(v >> 48)
		e.buf[3] = byte(v >> 40)
		e.buf[4] = byte(v >> 32)
		e.buf[5] = byte(v >> 24)
		e.buf[6] = byte(v >> 16)
		e.buf[7] = byte(v >> 8)
		e.buf[8] = byte(v)
		b = e.buf[:9]
	}
	_, err := e.w.Write(b)
	return err
}

// PackUint packs a Uint value to the message pack stream.
func (e *Encoder) PackUint(v uint64) error {
	var b []byte
	if v <= math.MaxInt8 {
		// Pack as signed for compatibility with other encoders.
		e.buf[0] = byte(v)
		b = e.buf[:1]
	} else {
		b = e.encodeNum(uintEncodings, v)
	}
	_, err := e.w.Write(b)
	return err
}

// PackFloat writes a Float value to the MessagePack stream.
func (e *Encoder) PackFloat(f float64) error {
	n := math.Float64bits(f)
	e.buf[0] = float64Code
	e.buf[1] = byte(n >> 56)
	e.buf[2] = byte(n >> 48)
	e.buf[3] = byte(n >> 40)
	e.buf[4] = byte(n >> 32)
	e.buf[5] = byte(n >> 24)
	e.buf[6] = byte(n >> 16)
	e.buf[7] = byte(n >> 8)
	e.buf[8] = byte(n)
	_, err := e.w.Write(e.buf[:9])
	return err
}

func (e *Encoder) packStringLen(n int64) error {
	var b []byte
	if n < 32 {
		e.buf[0] = byte(fixStringCodeMin + n)
		b = e.buf[:1]
	} else if n <= math.MaxUint32 {
		b = e.encodeNum(stringLenEncodings, uint64(n))
	} else {
		return ErrLongStringOrBinary
	}
	_, err := e.w.Write(b)
	return err
}

// PackString writes a String value to the MessagePack stream.
func (e *Encoder) PackString(v string) error {
	if err := e.packStringLen(int64(len(v))); err != nil {
		return err
	}
	_, err := e.writeString(v)
	return err
}

// PackStringBytes writes a String value to the MessagePack stream.
func (e *Encoder) PackStringBytes(v []byte) error {
	if err := e.packStringLen(int64(len(v))); err != nil {
		return err
	}
	_, err := e.w.Write(v)
	return err
}

// PackBinary writes a Binary value to the MessagePack stream.
func (e *Encoder) PackBinary(v []byte) error {
	n := uint64(len(v))
	if n > math.MaxUint32 {
		return ErrLongStringOrBinary
	}
	if _, err := e.w.Write(e.encodeNum(binaryLenEncodings, n)); err != nil {
		return err
	}
	_, err := e.w.Write(v)
	return err
}

func (e *Encoder) packArrayMapLen(fixMin int64, fc *numCodes, v int64) error {
	if v < 0 || v > math.MaxUint32 {
		return ErrIllegalSize
	}
	var b []byte
	if v < 16 {
		e.buf[0] = byte(fixMin + v)
		b = e.buf[:1]
	} else {
		b = e.encodeNum(fc, uint64(v))
	}
	_, err := e.w.Write(b)
	return err
}

// PackArrayLen write an Array length to the MessagePack stream. The
// application must write n objects to the stream following this call.
func (e *Encoder) PackArrayLen(n int64) error {
	return e.packArrayMapLen(fixArrayCodeMin, arrayLenEncodings, n)
}

// PackMapLen write an Map length to the MessagePack stream. The application
// must write n key-value pairs to the stream following this call.
func (e *Encoder) PackMapLen(n int64) error {
	return e.packArrayMapLen(fixMapCodeMin, mapLenEncodings, n)
}

// PackExtension writes an extension to the MessagePack stream.
func (e *Encoder) PackExtension(kind int, data []byte) error {
	var b []byte
	switch len(data) {
	case 1:
		e.buf[0] = fixext1Code
		e.buf[1] = byte(kind)
		b = e.buf[:2]
	case 2:
		e.buf[0] = fixext2Code
		e.buf[1] = byte(kind)
		b = e.buf[:2]
	case 4:
		e.buf[0] = fixext4Code
		e.buf[1] = byte(kind)
		b = e.buf[:2]
	case 8:
		e.buf[0] = fixext8Code
		e.buf[1] = byte(kind)
		b = e.buf[:2]
	case 16:
		e.buf[0] = fixext16Code
		e.buf[1] = byte(kind)
		b = e.buf[:2]
	default:
		b = e.encodeNum(extLenEncodings, uint64(len(data)))
		b = append(b, byte(kind))
	}
	if _, err := e.w.Write(b); err != nil {
		return err
	}
	_, err := e.w.Write(data)
	return err
}

// PackNil writes a Nil value to the MessagePack stream.
func (e *Encoder) PackNil() error {
	e.buf[0] = nilCode
	_, err := e.w.Write(e.buf[:1])
	return err
}

// PackRaw writes bytes directly to the MessagePack stream. It is the
// application's responsibility to ensure that the bytes are valid.
func (e *Encoder) PackRaw(p []byte) error {
	_, err := e.w.Write(p)
	return err
}
<commit_msg>msgpack: format pack.go<commit_after>package msgpack

import (
	"errors"
	"io"
	"math"
)

var (
	// ErrLongStringOrBinary is the long string or binary error.
	ErrLongStringOrBinary = errors.New("msgpack: long string or binary")

	// ErrIllegalSize is the illegal array or map size error.
	ErrIllegalSize = errors.New("msgpack: illegal array or map size")
)

// Encoder writes values in MessagePack format.
type Encoder struct {
	buf         [32]byte
	w           io.Writer
	writeString func(string) (int, error)
	err         error // permanent error
}

// NewEncoder allocates and initializes a new Unpacker.
func NewEncoder(w io.Writer) *Encoder {
	e := &Encoder{
		w: w,
	}

	if ws, ok := w.(interface {
		WriteString(string) (int, error)
	}); ok {
		e.writeString = ws.WriteString
	} else {
		e.writeString = e.writeStringUnopt
	}

	return e
}

func (e *Encoder) writeStringUnopt(s string) (int, error) {
	if len(s) <= len(e.buf) {
		copy(e.buf[:], s)
		return e.w.Write(e.buf[:len(s)])
	}

	return e.w.Write([]byte(s))
}

type numCodes struct {
	c8  byte
	c16 byte
	c32 byte
	c64 byte
}

var (
	stringLenEncodings = &numCodes{
		c8:  string8Code,
		c16: string16Code,
		c32: string32Code,
		c64: 0,
	}
	binaryLenEncodings = &numCodes{
		c8:  binary8Code,
		c16: binary16Code,
		c32: binary32Code,
		c64: 0,
	}
	arrayLenEncodings = &numCodes{
		c8:  0,
		c16: array16Code,
		c32: array32Code,
		c64: 0,
	}
	mapLenEncodings = &numCodes{
		c8:  0,
		c16: map16Code,
		c32: map32Code,
		c64: 0,
	}
	extLenEncodings = &numCodes{
		c8:  ext8Code,
		c16: ext16Code,
		c32: ext32Code,
		c64: 0,
	}
	uintEncodings = &numCodes{
		c8:  uint8Code,
		c16: uint16Code,
		c32: uint32Code,
		c64: uint64Code,
	}
)

func (e *Encoder) encodeNum(fc *numCodes, v uint64) []byte {
	switch {
	case fc.c8 != 0 && v <= math.MaxUint8:
		e.buf[0] = fc.c8
		e.buf[1] = byte(v)
		return e.buf[:2]

	case v <= math.MaxUint16:
		e.buf[0] = fc.c16
		e.buf[1] = byte(v >> 8)
		e.buf[2] = byte(v)
		return e.buf[:3]

	case v <= math.MaxUint32:
		e.buf[0] = fc.c32
		e.buf[1] = byte(v >> 24)
		e.buf[2] = byte(v >> 16)
		e.buf[3] = byte(v >> 8)
		e.buf[4] = byte(v)
		return e.buf[:5]

	default:
		e.buf[0] = fc.c64
		e.buf[1] = byte(v >> 56)
		e.buf[2] = byte(v >> 48)
		e.buf[3] = byte(v >> 40)
		e.buf[4] = byte(v >> 32)
		e.buf[5] = byte(v >> 24)
		e.buf[6] = byte(v >> 16)
		e.buf[7] = byte(v >> 8)
		e.buf[8] = byte(v)
		return e.buf[:9]
	}
}

// PackBool writes a Bool value to the MessagePack stream.
func (e *Encoder) PackBool(b bool) error {
	if b {
		e.buf[0] = trueCode
	} else {
		e.buf[0] = falseCode
	}

	_, err := e.w.Write(e.buf[:1])
	return err
}

// PackInt packs an Int value to the MessagePack stream.
func (e *Encoder) PackInt(v int64) error {
	var b []byte

	switch {
	case 0 <= v && v <= math.MaxInt8:
		e.buf[0] = byte(v)
		b = e.buf[:1]

	case v > 0:
		// Pack as unsigned for compatibility with other encoders.
		b = e.encodeNum(uintEncodings, uint64(v))

	case v >= -32:
		e.buf[0] = byte(v)
		b = e.buf[:1]

	case v >= math.MinInt8:
		e.buf[0] = int8Code
		e.buf[1] = byte(v)
		b = e.buf[:2]

	case v >= math.MinInt16:
		e.buf[0] = int16Code
		e.buf[1] = byte(v >> 8)
		e.buf[2] = byte(v)
		b = e.buf[:3]

	case v >= math.MinInt32:
		e.buf[0] = int32Code
		e.buf[1] = byte(v >> 24)
		e.buf[2] = byte(v >> 16)
		e.buf[3] = byte(v >> 8)
		e.buf[4] = byte(v)
		b = e.buf[:5]

	default:
		e.buf[0] = int64Code
		e.buf[1] = byte(v >> 56)
		e.buf[2] = byte(v >> 48)
		e.buf[3] = byte(v >> 40)
		e.buf[4] = byte(v >> 32)
		e.buf[5] = byte(v >> 24)
		e.buf[6] = byte(v >> 16)
		e.buf[7] = byte(v >> 8)
		e.buf[8] = byte(v)
		b = e.buf[:9]
	}

	_, err := e.w.Write(b)
	return err
}

// PackUint packs a Uint value to the message pack stream.
func (e *Encoder) PackUint(v uint64) error {
	var b []byte

	if v <= math.MaxInt8 {
		// Pack as signed for compatibility with other encoders.
		e.buf[0] = byte(v)
		b = e.buf[:1]
	} else {
		b = e.encodeNum(uintEncodings, v)
	}

	_, err := e.w.Write(b)
	return err
}

// PackFloat writes a Float value to the MessagePack stream.
func (e *Encoder) PackFloat(f float64) error {
	n := math.Float64bits(f)
	e.buf[0] = float64Code
	e.buf[1] = byte(n >> 56)
	e.buf[2] = byte(n >> 48)
	e.buf[3] = byte(n >> 40)
	e.buf[4] = byte(n >> 32)
	e.buf[5] = byte(n >> 24)
	e.buf[6] = byte(n >> 16)
	e.buf[7] = byte(n >> 8)
	e.buf[8] = byte(n)

	_, err := e.w.Write(e.buf[:9])
	return err
}

func (e *Encoder) packStringLen(n int64) error {
	var b []byte

	if n < 32 {
		e.buf[0] = byte(fixStringCodeMin + n)
		b = e.buf[:1]
	} else if n <= math.MaxUint32 {
		b = e.encodeNum(stringLenEncodings, uint64(n))
	} else {
		return ErrLongStringOrBinary
	}

	_, err := e.w.Write(b)
	return err
}

// PackString writes a String value to the MessagePack stream.
func (e *Encoder) PackString(v string) error {
	if err := e.packStringLen(int64(len(v))); err != nil {
		return err
	}

	_, err := e.writeString(v)
	return err
}

// PackStringBytes writes a String value to the MessagePack stream.
func (e *Encoder) PackStringBytes(v []byte) error {
	if err := e.packStringLen(int64(len(v))); err != nil {
		return err
	}

	_, err := e.w.Write(v)
	return err
}

// PackBinary writes a Binary value to the MessagePack stream.
func (e *Encoder) PackBinary(v []byte) error {
	n := uint64(len(v))

	if n > math.MaxUint32 {
		return ErrLongStringOrBinary
	}

	if _, err := e.w.Write(e.encodeNum(binaryLenEncodings, n)); err != nil {
		return err
	}

	_, err := e.w.Write(v)
	return err
}

func (e *Encoder) packArrayMapLen(fixMin int64, fc *numCodes, v int64) error {
	if v < 0 || v > math.MaxUint32 {
		return ErrIllegalSize
	}

	var b []byte
	if v < 16 {
		e.buf[0] = byte(fixMin + v)
		b = e.buf[:1]
	} else {
		b = e.encodeNum(fc, uint64(v))
	}

	_, err := e.w.Write(b)
	return err
}

// PackArrayLen write an Array length to the MessagePack stream. The
// application must write n objects to the stream following this call.
func (e *Encoder) PackArrayLen(n int64) error {
	return e.packArrayMapLen(fixArrayCodeMin, arrayLenEncodings, n)
}

// PackMapLen write an Map length to the MessagePack stream. The application
// must write n key-value pairs to the stream following this call.
func (e *Encoder) PackMapLen(n int64) error {
	return e.packArrayMapLen(fixMapCodeMin, mapLenEncodings, n)
}

// PackExtension writes an extension to the MessagePack stream.
func (e *Encoder) PackExtension(kind int, data []byte) error {
	var b []byte

	switch len(data) {
	case 1:
		e.buf[0] = fixext1Code
		e.buf[1] = byte(kind)
		b = e.buf[:2]

	case 2:
		e.buf[0] = fixext2Code
		e.buf[1] = byte(kind)
		b = e.buf[:2]

	case 4:
		e.buf[0] = fixext4Code
		e.buf[1] = byte(kind)
		b = e.buf[:2]

	case 8:
		e.buf[0] = fixext8Code
		e.buf[1] = byte(kind)
		b = e.buf[:2]

	case 16:
		e.buf[0] = fixext16Code
		e.buf[1] = byte(kind)
		b = e.buf[:2]

	default:
		b = e.encodeNum(extLenEncodings, uint64(len(data)))
		b = append(b, byte(kind))
	}

	if _, err := e.w.Write(b); err != nil {
		return err
	}

	_, err := e.w.Write(data)
	return err
}

// PackNil writes a Nil value to the MessagePack stream.
func (e *Encoder) PackNil() error {
	e.buf[0] = nilCode
	_, err := e.w.Write(e.buf[:1])
	return err
}

// PackRaw writes bytes directly to the MessagePack stream. It is the
// application's responsibility to ensure that the bytes are valid.
func (e *Encoder) PackRaw(p []byte) error {
	_, err := e.w.Write(p)
	return err
}
<|endoftext|>


<commit_before>// The state package enables reading, observing, and changing
// the state stored in MongoDB of a whole environment
// managed by juju.
package mstate

import (
	"fmt"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
	"labix.org/v2/mgo/txn"
	"launchpad.net/juju-core/charm"
	"launchpad.net/juju-core/trivial"
	"launchpad.net/juju-core/version"
	"net/url"
)

// Tools describes a particular set of juju tools and where to find them.
type Tools struct {
	version.Binary
	URL string
}

// State represents the state of an environment
// managed by juju.
type State struct {
	db        *mgo.Database
	charms    *mgo.Collection
	machines  *mgo.Collection
	relations *mgo.Collection
	services  *mgo.Collection
	settings  *mgo.Collection
	units     *mgo.Collection
	runner    *txn.Runner
}

// EnvironConfig returns the current configuration of the environment.
func (s *State) EnvironConfig() (*ConfigNode, error) {
	return readConfigNode(s, "e")
}

// AddMachine creates a new machine state.
func (s *State) AddMachine() (m *Machine, err error) {
	defer trivial.ErrorContextf(&err, "cannot add a new machine")
	id, err := s.sequence("machine")
	if err != nil {
		return nil, err
	}
	mdoc := machineDoc{
		Id:   id,
		Life: Alive,
	}
	ops := []txn.Op{{
		C:      s.machines.Name,
		Id:     id,
		Assert: txn.DocMissing,
		Insert: mdoc,
	}}
	err = s.runner.Run(ops, "", nil)
	if err != nil {
		return nil, err
	}
	return newMachine(s, &mdoc), nil
}

// RemoveMachine removes the machine with the the given id.
func (s *State) RemoveMachine(id int) (err error) {
	defer trivial.ErrorContextf(&err, "cannot remove machine %d", id)
	m, err := s.Machine(id)
	if err != nil {
		return err
	}
	if m.doc.Life != Dead {
		panic(fmt.Errorf("machine %d is not dead", id))
	}
	sel := bson.D{
		{"_id", id},
		{"life", Dead},
	}
	ops := []txn.Op{{
		C:      s.machines.Name,
		Id:     id,
		Assert: sel,
		Remove: true,
	}}
	err = s.runner.Run(ops, "", nil)
	if err != nil {
		return err
	}
	return nil
}

// AllMachines returns all machines in the environment.
func (s *State) AllMachines() (machines []*Machine, err error) {
	mdocs := []machineDoc{}
	sel := bson.D{}
	err = s.machines.Find(sel).Select(bson.D{{"_id", 1}}).All(&mdocs)
	if err != nil {
		return nil, fmt.Errorf("cannot get all machines: %v", err)
	}
	for _, doc := range mdocs {
		machines = append(machines, newMachine(s, &doc))
	}
	return
}

// Machine returns the machine with the given id.
func (s *State) Machine(id int) (*Machine, error) {
	mdoc := &machineDoc{}
	sel := bson.D{{"_id", id}}
	err := s.machines.Find(sel).One(mdoc)
	if err != nil {
		return nil, fmt.Errorf("cannot get machine %d: %v", id, err)
	}
	return newMachine(s, mdoc), nil
}

// AddCharm adds the ch charm with curl to the state.  bundleUrl must be
// set to a URL where the bundle for ch may be downloaded from.
// On success the newly added charm state is returned.
func (s *State) AddCharm(ch charm.Charm, curl *charm.URL, bundleURL *url.URL, bundleSha256 string) (stch *Charm, err error) {
	cdoc := &charmDoc{
		URL:          curl,
		Meta:         ch.Meta(),
		Config:       ch.Config(),
		BundleURL:    bundleURL,
		BundleSha256: bundleSha256,
	}
	err = s.charms.Insert(cdoc)
	if err != nil {
		return nil, fmt.Errorf("cannot add charm %q: %v", curl, err)
	}
	return newCharm(s, cdoc)
}

// Charm returns the charm with the given URL.
func (s *State) Charm(curl *charm.URL) (*Charm, error) {
	cdoc := &charmDoc{}
	err := s.charms.Find(bson.D{{"_id", curl}}).One(cdoc)
	if err != nil {
		return nil, fmt.Errorf("cannot get charm %q: %v", curl, err)
	}

	return newCharm(s, cdoc)
}

// AddService creates a new service state with the given unique name
// and the charm state.
func (s *State) AddService(name string, ch *Charm) (service *Service, err error) {
	sdoc := serviceDoc{
		Name:     name,
		CharmURL: ch.URL(),
		Life:     Alive,
	}
	ops := []txn.Op{{
		C:      s.services.Name,
		Id:     name,
		Assert: txn.DocMissing,
		Insert: sdoc,
	}}
	err = s.runner.Run(ops, "", nil)
	if err != nil {
		return nil, fmt.Errorf("cannot add service %q:", name, err)
	}
	return &Service{st: s, doc: sdoc}, nil
}

// RemoveService removes a service from the state. It will also remove all
// its units and break any of its existing relations.
func (s *State) RemoveService(svc *Service) (err error) {
	// TODO Integrate with txn and do lifecycle properly.
	defer trivial.ErrorContextf(&err, "cannot remove service %q", svc)

	if svc.doc.Life != Dead {
		panic(fmt.Errorf("service %q is not dead", svc))
	}
	// Remove relations first, to minimize unwanted hook executions.
	rels, err := svc.Relations()
	if err != nil {
		return err
	}
	for _, rel := range rels {
		err = rel.Die()
		if err != nil {
			return err
		}
		err = s.RemoveRelation(rel)
		if err != nil {
			return err
		}
	}
	// TODO Will be deleted with proper lifecycle integration.
	units, err := svc.AllUnits()
	if err != nil {
		return err
	}
	for _, unit := range units {
		err = unit.Die()
		if err != nil {
			return err
		}
		if err = svc.RemoveUnit(unit); err != nil {
			return err
		}
	}
	// Remove the service.
	sel := bson.D{
		{"_id", svc.doc.Name},
		{"life", Dead},
	}
	err = s.services.Remove(sel)
	if err != nil {
		return err
	}
	return nil
}

// Service returns a service state by name.
func (s *State) Service(name string) (service *Service, err error) {
	sdoc := serviceDoc{}
	sel := bson.D{{"_id", name}}
	err = s.services.Find(sel).One(&sdoc)
	if err != nil {
		return nil, fmt.Errorf("cannot get service %q: %v", name, err)
	}
	return &Service{st: s, doc: sdoc}, nil
}

// AllServices returns all deployed services in the environment.
func (s *State) AllServices() (services []*Service, err error) {
	sdocs := []serviceDoc{}
	err = s.services.Find(bson.D{}).All(&sdocs)
	if err != nil {
		return nil, fmt.Errorf("cannot get all services")
	}
	for _, v := range sdocs {
		services = append(services, &Service{st: s, doc: v})
	}
	return services, nil
}

// AddRelation creates a new relation with the given endpoints.
func (s *State) AddRelation(endpoints ...RelationEndpoint) (r *Relation, err error) {
	defer trivial.ErrorContextf(&err, "cannot add relation %q", relationKey(endpoints))
	switch len(endpoints) {
	case 1:
		if endpoints[0].RelationRole != RolePeer {
			return nil, fmt.Errorf("single endpoint must be a peer relation")
		}
	case 2:
		if !endpoints[0].CanRelateTo(&endpoints[1]) {
			return nil, fmt.Errorf("endpoints do not relate")
		}
	default:
		return nil, fmt.Errorf("cannot relate %d endpoints", len(endpoints))
	}

	var scope charm.RelationScope
	for _, v := range endpoints {
		if v.RelationScope == charm.ScopeContainer {
			scope = charm.ScopeContainer
		}
		// BUG(aram): potential race in the time between getting the service
		// to validate the endpoint and actually writting the relation
		// into MongoDB; the service might have disappeared.
		_, err = s.Service(v.ServiceName)
		if err != nil {
			return nil, err
		}
	}
	if scope == charm.ScopeContainer {
		for i := range endpoints {
			endpoints[i].RelationScope = scope
		}
	}
	id, err := s.sequence("relation")
	if err != nil {
		return nil, err
	}
	doc := relationDoc{
		Key:       relationKey(endpoints),
		Id:        id,
		Endpoints: endpoints,
		Life:      Alive,
	}
	ops := []txn.Op{{
		C:      s.relations.Name,
		Id:     id,
		Assert: txn.DocMissing,
		Insert: doc,
	}}
	err = s.runner.Run(ops, "", nil)
	if err != nil {
		return nil, err
	}
	return newRelation(s, &doc), nil
}

// Relation returns the existing relation with the given endpoints.
func (s *State) Relation(endpoints ...RelationEndpoint) (r *Relation, err error) {
	defer trivial.ErrorContextf(&err, "cannot get relation %q", relationKey(endpoints))

	doc := relationDoc{}
	err = s.relations.Find(bson.D{{"_id", relationKey(endpoints)}}).One(&doc)
	if err != nil {
		return nil, err
	}
	return newRelation(s, &doc), nil
}

// RemoveRelation removes the supplied relation.
func (s *State) RemoveRelation(r *Relation) (err error) {
	defer trivial.ErrorContextf(&err, "cannot remove relation %q", r.doc.Key)

	if r.doc.Life != Dead {
		panic(fmt.Errorf("relation %q is not dead", r))
	}
	sel := bson.D{
		{"_id", r.doc.Key},
		{"life", Dead},
	}
	err = s.relations.Remove(sel)
	if err != nil {
		return err
	}
	return nil
}

// Unit returns a unit by name.
func (s *State) Unit(name string) (*Unit, error) {
	doc := unitDoc{}
	err := s.units.FindId(name).One(&doc)
	if err != nil {
		return nil, fmt.Errorf("cannot get unit %q: %v", name, err)
	}
	return newUnit(s, &doc), nil
}
<commit_msg>mstate: use key, not id, in AddRelation txn Op id<commit_after>// The state package enables reading, observing, and changing
// the state stored in MongoDB of a whole environment
// managed by juju.
package mstate

import (
	"fmt"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
	"labix.org/v2/mgo/txn"
	"launchpad.net/juju-core/charm"
	"launchpad.net/juju-core/trivial"
	"launchpad.net/juju-core/version"
	"net/url"
)

// Tools describes a particular set of juju tools and where to find them.
type Tools struct {
	version.Binary
	URL string
}

// State represents the state of an environment
// managed by juju.
type State struct {
	db        *mgo.Database
	charms    *mgo.Collection
	machines  *mgo.Collection
	relations *mgo.Collection
	services  *mgo.Collection
	settings  *mgo.Collection
	units     *mgo.Collection
	runner    *txn.Runner
}

// EnvironConfig returns the current configuration of the environment.
func (s *State) EnvironConfig() (*ConfigNode, error) {
	return readConfigNode(s, "e")
}

// AddMachine creates a new machine state.
func (s *State) AddMachine() (m *Machine, err error) {
	defer trivial.ErrorContextf(&err, "cannot add a new machine")
	id, err := s.sequence("machine")
	if err != nil {
		return nil, err
	}
	mdoc := machineDoc{
		Id:   id,
		Life: Alive,
	}
	ops := []txn.Op{{
		C:      s.machines.Name,
		Id:     id,
		Assert: txn.DocMissing,
		Insert: mdoc,
	}}
	err = s.runner.Run(ops, "", nil)
	if err != nil {
		return nil, err
	}
	return newMachine(s, &mdoc), nil
}

// RemoveMachine removes the machine with the the given id.
func (s *State) RemoveMachine(id int) (err error) {
	defer trivial.ErrorContextf(&err, "cannot remove machine %d", id)
	m, err := s.Machine(id)
	if err != nil {
		return err
	}
	if m.doc.Life != Dead {
		panic(fmt.Errorf("machine %d is not dead", id))
	}
	sel := bson.D{
		{"_id", id},
		{"life", Dead},
	}
	ops := []txn.Op{{
		C:      s.machines.Name,
		Id:     id,
		Assert: sel,
		Remove: true,
	}}
	err = s.runner.Run(ops, "", nil)
	if err != nil {
		return err
	}
	return nil
}

// AllMachines returns all machines in the environment.
func (s *State) AllMachines() (machines []*Machine, err error) {
	mdocs := []machineDoc{}
	sel := bson.D{}
	err = s.machines.Find(sel).Select(bson.D{{"_id", 1}}).All(&mdocs)
	if err != nil {
		return nil, fmt.Errorf("cannot get all machines: %v", err)
	}
	for _, doc := range mdocs {
		machines = append(machines, newMachine(s, &doc))
	}
	return
}

// Machine returns the machine with the given id.
func (s *State) Machine(id int) (*Machine, error) {
	mdoc := &machineDoc{}
	sel := bson.D{{"_id", id}}
	err := s.machines.Find(sel).One(mdoc)
	if err != nil {
		return nil, fmt.Errorf("cannot get machine %d: %v", id, err)
	}
	return newMachine(s, mdoc), nil
}

// AddCharm adds the ch charm with curl to the state.  bundleUrl must be
// set to a URL where the bundle for ch may be downloaded from.
// On success the newly added charm state is returned.
func (s *State) AddCharm(ch charm.Charm, curl *charm.URL, bundleURL *url.URL, bundleSha256 string) (stch *Charm, err error) {
	cdoc := &charmDoc{
		URL:          curl,
		Meta:         ch.Meta(),
		Config:       ch.Config(),
		BundleURL:    bundleURL,
		BundleSha256: bundleSha256,
	}
	err = s.charms.Insert(cdoc)
	if err != nil {
		return nil, fmt.Errorf("cannot add charm %q: %v", curl, err)
	}
	return newCharm(s, cdoc)
}

// Charm returns the charm with the given URL.
func (s *State) Charm(curl *charm.URL) (*Charm, error) {
	cdoc := &charmDoc{}
	err := s.charms.Find(bson.D{{"_id", curl}}).One(cdoc)
	if err != nil {
		return nil, fmt.Errorf("cannot get charm %q: %v", curl, err)
	}

	return newCharm(s, cdoc)
}

// AddService creates a new service state with the given unique name
// and the charm state.
func (s *State) AddService(name string, ch *Charm) (service *Service, err error) {
	sdoc := serviceDoc{
		Name:     name,
		CharmURL: ch.URL(),
		Life:     Alive,
	}
	ops := []txn.Op{{
		C:      s.services.Name,
		Id:     name,
		Assert: txn.DocMissing,
		Insert: sdoc,
	}}
	err = s.runner.Run(ops, "", nil)
	if err != nil {
		return nil, fmt.Errorf("cannot add service %q:", name, err)
	}
	return &Service{st: s, doc: sdoc}, nil
}

// RemoveService removes a service from the state. It will also remove all
// its units and break any of its existing relations.
func (s *State) RemoveService(svc *Service) (err error) {
	// TODO Integrate with txn and do lifecycle properly.
	defer trivial.ErrorContextf(&err, "cannot remove service %q", svc)

	if svc.doc.Life != Dead {
		panic(fmt.Errorf("service %q is not dead", svc))
	}
	// Remove relations first, to minimize unwanted hook executions.
	rels, err := svc.Relations()
	if err != nil {
		return err
	}
	for _, rel := range rels {
		err = rel.Die()
		if err != nil {
			return err
		}
		err = s.RemoveRelation(rel)
		if err != nil {
			return err
		}
	}
	// TODO Will be deleted with proper lifecycle integration.
	units, err := svc.AllUnits()
	if err != nil {
		return err
	}
	for _, unit := range units {
		err = unit.Die()
		if err != nil {
			return err
		}
		if err = svc.RemoveUnit(unit); err != nil {
			return err
		}
	}
	// Remove the service.
	sel := bson.D{
		{"_id", svc.doc.Name},
		{"life", Dead},
	}
	err = s.services.Remove(sel)
	if err != nil {
		return err
	}
	return nil
}

// Service returns a service state by name.
func (s *State) Service(name string) (service *Service, err error) {
	sdoc := serviceDoc{}
	sel := bson.D{{"_id", name}}
	err = s.services.Find(sel).One(&sdoc)
	if err != nil {
		return nil, fmt.Errorf("cannot get service %q: %v", name, err)
	}
	return &Service{st: s, doc: sdoc}, nil
}

// AllServices returns all deployed services in the environment.
func (s *State) AllServices() (services []*Service, err error) {
	sdocs := []serviceDoc{}
	err = s.services.Find(bson.D{}).All(&sdocs)
	if err != nil {
		return nil, fmt.Errorf("cannot get all services")
	}
	for _, v := range sdocs {
		services = append(services, &Service{st: s, doc: v})
	}
	return services, nil
}

// AddRelation creates a new relation with the given endpoints.
func (s *State) AddRelation(endpoints ...RelationEndpoint) (r *Relation, err error) {
	defer trivial.ErrorContextf(&err, "cannot add relation %q", relationKey(endpoints))
	switch len(endpoints) {
	case 1:
		if endpoints[0].RelationRole != RolePeer {
			return nil, fmt.Errorf("single endpoint must be a peer relation")
		}
	case 2:
		if !endpoints[0].CanRelateTo(&endpoints[1]) {
			return nil, fmt.Errorf("endpoints do not relate")
		}
	default:
		return nil, fmt.Errorf("cannot relate %d endpoints", len(endpoints))
	}

	var scope charm.RelationScope
	for _, v := range endpoints {
		if v.RelationScope == charm.ScopeContainer {
			scope = charm.ScopeContainer
		}
		// BUG(aram): potential race in the time between getting the service
		// to validate the endpoint and actually writting the relation
		// into MongoDB; the service might have disappeared.
		_, err = s.Service(v.ServiceName)
		if err != nil {
			return nil, err
		}
	}
	if scope == charm.ScopeContainer {
		for i := range endpoints {
			endpoints[i].RelationScope = scope
		}
	}
	id, err := s.sequence("relation")
	if err != nil {
		return nil, err
	}
	doc := relationDoc{
		Key:       relationKey(endpoints),
		Id:        id,
		Endpoints: endpoints,
		Life:      Alive,
	}
	ops := []txn.Op{{
		C:      s.relations.Name,
		Id:     doc.Key,
		Assert: txn.DocMissing,
		Insert: doc,
	}}
	err = s.runner.Run(ops, "", nil)
	if err != nil {
		return nil, err
	}
	return newRelation(s, &doc), nil
}

// Relation returns the existing relation with the given endpoints.
func (s *State) Relation(endpoints ...RelationEndpoint) (r *Relation, err error) {
	defer trivial.ErrorContextf(&err, "cannot get relation %q", relationKey(endpoints))

	doc := relationDoc{}
	err = s.relations.Find(bson.D{{"_id", relationKey(endpoints)}}).One(&doc)
	if err != nil {
		return nil, err
	}
	return newRelation(s, &doc), nil
}

// RemoveRelation removes the supplied relation.
func (s *State) RemoveRelation(r *Relation) (err error) {
	defer trivial.ErrorContextf(&err, "cannot remove relation %q", r.doc.Key)

	if r.doc.Life != Dead {
		panic(fmt.Errorf("relation %q is not dead", r))
	}
	sel := bson.D{
		{"_id", r.doc.Key},
		{"life", Dead},
	}
	err = s.relations.Remove(sel)
	if err != nil {
		return err
	}
	return nil
}

// Unit returns a unit by name.
func (s *State) Unit(name string) (*Unit, error) {
	doc := unitDoc{}
	err := s.units.FindId(name).One(&doc)
	if err != nil {
		return nil, fmt.Errorf("cannot get unit %q: %v", name, err)
	}
	return newUnit(s, &doc), nil
}
<|endoftext|>