<commit_before>//
// Use and distribution licensed under the Apache license version 2.
//
// See the COPYING file in the root project directory for full text.
//

package ghw

import (
	"github.com/StackExchange/wmi"
)

const wqlOperatingSystem = "SELECT FreePhysicalMemory, FreeSpaceInPagingFiles, FreeVirtualMemory, TotalSwapSpaceSize, TotalVirtualMemorySize, TotalVisibleMemorySize FROM Win32_OperatingSystem"

type win32OperatingSystem struct {
	FreePhysicalMemory     uint64
	FreeSpaceInPagingFiles uint64
	FreeVirtualMemory      uint64
	TotalSwapSpaceSize     uint64
	TotalVirtualMemorySize uint64
	TotalVisibleMemorySize uint64
}

const wqlPhysicalMemory = "SELECT BankLabel, Capacity, DataWidth, Description, DeviceLocator, Manufacturer, Model, Name, PartNumber, PositionInRow, SerialNumber, Speed, Tag, TotalWidth FROM Win32_PhysicalMemory"

type win32PhysicalMemory struct {
	BankLabel     string
	Capacity      uint64
	DataWidth     uint16
	Description   string
	DeviceLocator string
	Manufacturer  string
	Model         string
	Name          string
	PartNumber    string
	PositionInRow uint32
	SerialNumber  string
	Speed         uint32
	Tag           string
	TotalWidth    uint16
}

func (ctx *context) memFillInfo(info *MemoryInfo) error {
	// Getting info from WMI
	var win32OSDescriptions []win32OperatingSystem
	if err := wmi.Query(wqlOperatingSystem, &win32OSDescriptions); err != nil {
		return err
	}
	var win32MemDescriptions []win32PhysicalMemory
	if err := wmi.Query(wqlPhysicalMemory, &win32MemDescriptions); err != nil {
		return err
	}
	// Converting into standard structures
	// Handling physical memory modules
	info.Modules = make([]*MemoryModule, 0, len(win32MemDescriptions))
	for _, description := range win32MemDescriptions {
		info.Modules = append(info.Modules, &MemoryModule{
			Label:        description.BankLabel,
			Location:     description.DeviceLocator,
			SerialNumber: description.SerialNumber,
			SizeBytes:    int64(description.Capacity),
			Vendor:       description.Manufacturer,
		})
	}
	// Handling physical memory total/free size (as seen by OS)
	var totalUsableBytes uint64
	var totalPhysicalBytes uint64
	for _, description := range win32OSDescriptions {
		totalUsableBytes += description.FreePhysicalMemory
		totalPhysicalBytes += description.TotalVisibleMemorySize
	}
	info.TotalUsableBytes = int64(totalUsableBytes)
	info.TotalPhysicalBytes = int64(totalPhysicalBytes)
	return nil
}
<commit_msg>fix Windows memory calculations<commit_after>//
// Use and distribution licensed under the Apache license version 2.
//
// See the COPYING file in the root project directory for full text.
//

package ghw

import (
	"github.com/StackExchange/wmi"
)

const wqlOperatingSystem = "SELECT TotalVisibleMemorySize FROM Win32_OperatingSystem"

type win32OperatingSystem struct {
	TotalVisibleMemorySize uint64
}

const wqlPhysicalMemory = "SELECT BankLabel, Capacity, DataWidth, Description, DeviceLocator, Manufacturer, Model, Name, PartNumber, PositionInRow, SerialNumber, Speed, Tag, TotalWidth FROM Win32_PhysicalMemory"

type win32PhysicalMemory struct {
	BankLabel     string
	Capacity      uint64
	DataWidth     uint16
	Description   string
	DeviceLocator string
	Manufacturer  string
	Model         string
	Name          string
	PartNumber    string
	PositionInRow uint32
	SerialNumber  string
	Speed         uint32
	Tag           string
	TotalWidth    uint16
}

func (ctx *context) memFillInfo(info *MemoryInfo) error {
	// Getting info from WMI
	var win32OSDescriptions []win32OperatingSystem
	if err := wmi.Query(wqlOperatingSystem, &win32OSDescriptions); err != nil {
		return err
	}
	var win32MemDescriptions []win32PhysicalMemory
	if err := wmi.Query(wqlPhysicalMemory, &win32MemDescriptions); err != nil {
		return err
	}
	// We calculate total physical memory size by summing the DIMM sizes
	var totalPhysicalBytes uint64
	info.Modules = make([]*MemoryModule, 0, len(win32MemDescriptions))
	for _, description := range win32MemDescriptions {
		totalPhysicalBytes += description.Capacity
		info.Modules = append(info.Modules, &MemoryModule{
			Label:        description.BankLabel,
			Location:     description.DeviceLocator,
			SerialNumber: description.SerialNumber,
			SizeBytes:    int64(description.Capacity),
			Vendor:       description.Manufacturer,
		})
	}
	var totalUsableBytes uint64
	for _, description := range win32OSDescriptions {
		// TotalVisibleMemorySize is the amount of memory available for us by
		// the operating system **in Kilobytes**
		totalUsableBytes += description.TotalVisibleMemorySize * uint64(KB)
	}
	info.TotalUsableBytes = int64(totalUsableBytes)
	info.TotalPhysicalBytes = int64(totalPhysicalBytes)
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/streadway/amqp"
	"sync"
	"time"
)

type MessageBroker interface {
	DeclareQueue(string) error
	Subscribe(string, int, chan bool, func() MessageProcessor) error
	SubscribeFanout(string, func() MessageProcessor) error
	Publish(string, string, string, []byte) error
	Close() error
}

type MessageProcessor interface {
	Process(message []byte)
}

type RabbitMessageBroker struct {
	conn *amqp.Connection
}

type TestMessageBroker struct {
	queues map[string]chan []byte
}

func (mb *RabbitMessageBroker) DeclareQueue(queueName string) error {
	ch, err := mb.conn.Channel()
	if err != nil {
		return err
	}

	_, err = ch.QueueDeclare(queueName, true, false, false, false, nil)
	return err
}

func (mb *RabbitMessageBroker) Publish(exchange, routingKey, msgType string, message []byte) error {
	ch, err := mb.conn.Channel()
	if err != nil {
		return err
	}
	defer ch.Close()

	msg := amqp.Publishing{
		Type:      msgType,
		Timestamp: time.Now(),
		Body:      message,
	}

	return ch.Publish(exchange, routingKey, false, false, msg)
}

// Subscribe will start pulling messages off the given queue and process up to
// subCount messages concurrently by passing them to the given function.
//
// When the passed gracefulQuitChan is closed, the subscribers shut down after
// finishing the message it is currently processing.
func (mb *RabbitMessageBroker) Subscribe(queueName string, subCount int, gracefulQuitChan chan bool, f func() MessageProcessor) error {
	ch, err := mb.conn.Channel()
	if err != nil {
		return err
	}
	defer ch.Close()

	err = ch.Qos(subCount, 0, false)
	if err != nil {
		return err
	}

	messages, err := ch.Consume(queueName, "processor", false, false, false, false, nil)
	if err != nil {
		return err
	}

	var wg sync.WaitGroup
	wg.Add(subCount)
	for i := 0; i < subCount; i++ {
		go func() {
			defer wg.Done()

			for {
				select {
				case <-gracefulQuitChan:
					return
				case message, ok := <-messages:
					if !ok {
						return
					}

					f().Process(message.Body)
					message.Ack(false)
				}
			}
		}()
	}
	wg.Wait()

	return nil
}

func (mb *RabbitMessageBroker) SubscribeFanout(exchange string, f func() MessageProcessor) error {
	ch, err := mb.conn.Channel()
	if err != nil {
		return err
	}
	defer ch.Close()

	err = ch.Qos(1, 0, false)
	if err != nil {
		return err
	}

	err = ch.ExchangeDeclare(exchange, "fanout", false, false, false, false, nil)
	if err != nil {
		return err
	}

	queue, err := ch.QueueDeclare("", false, false, true, false, nil)
	if err != nil {
		return err
	}

	err = ch.QueueBind(queue.Name, "", exchange, false, nil)
	if err != nil {
		return err
	}

	messages, err := ch.Consume(queue.Name, "", false, false, false, false, nil)
	if err != nil {
		return err
	}

	for message := range messages {
		f().Process(message.Body)
		message.Ack(false)
	}

	return nil
}

func (mb *RabbitMessageBroker) Close() error {
	return mb.conn.Close()
}

func NewMessageBroker(url string) (MessageBroker, error) {
	if url == "" {
		return nil, fmt.Errorf("URL is blank")
	}

	conn, err := amqp.Dial(url)
	if err != nil {
		return nil, err
	}

	return &RabbitMessageBroker{conn}, nil
}

func NewTestMessageBroker() MessageBroker {
	return &TestMessageBroker{
		queues: make(map[string]chan []byte),
	}
}

func (mb *TestMessageBroker) DeclareQueue(queueName string) error {
	mb.queues[queueName] = make(chan []byte, 1)
	return nil
}

func (mb *TestMessageBroker) Subscribe(queueName string, subCount int, gracefulQuitChan chan bool, f func() MessageProcessor) error {
	var wg sync.WaitGroup
	wg.Add(subCount)
	for i := 0; i < subCount; i++ {
		go func() {
			defer wg.Done()

			processor := f()

			for body := range mb.queues[queueName] {
				processor.Process(body)
			}
		}()
	}
	wg.Wait()

	return nil
}

func (mb *TestMessageBroker) SubscribeFanout(exchange string, f func() MessageProcessor) error {
	return fmt.Errorf("TestMessageBroker.SubscribeFanout needs to be implemented")
}

func (mb *TestMessageBroker) Publish(exchange, routingKey, msgType string, msg []byte) error {
	mb.queues[routingKey] <- msg
	return nil
}

func (mb *TestMessageBroker) Close() error {
	for _, ch := range mb.queues {
		close(ch)
	}
	return nil
}
<commit_msg>refactor(message_broker): keep a single channel around for publishing<commit_after>package main

import (
	"fmt"
	"github.com/streadway/amqp"
	"sync"
	"time"
)

type MessageBroker interface {
	DeclareQueue(string) error
	Subscribe(string, int, chan bool, func() MessageProcessor) error
	SubscribeFanout(string, func() MessageProcessor) error
	Publish(string, string, string, []byte) error
	Close() error
}

type MessageProcessor interface {
	Process(message []byte)
}

type RabbitMessageBroker struct {
	conn        *amqp.Connection
	publishChan *amqp.Channel
}

type TestMessageBroker struct {
	queues map[string]chan []byte
}

func (mb *RabbitMessageBroker) DeclareQueue(queueName string) error {
	ch, err := mb.conn.Channel()
	if err != nil {
		return err
	}

	_, err = ch.QueueDeclare(queueName, true, false, false, false, nil)
	return err
}

func (mb *RabbitMessageBroker) Publish(exchange, routingKey, msgType string, message []byte) error {
	msg := amqp.Publishing{
		Type:      msgType,
		Timestamp: time.Now(),
		Body:      message,
	}

	return mb.publishChan.Publish(exchange, routingKey, false, false, msg)
}

// Subscribe will start pulling messages off the given queue and process up to
// subCount messages concurrently by passing them to the given function.
//
// When the passed gracefulQuitChan is closed, the subscribers shut down after
// finishing the message it is currently processing.
func (mb *RabbitMessageBroker) Subscribe(queueName string, subCount int, gracefulQuitChan chan bool, f func() MessageProcessor) error {
	ch, err := mb.conn.Channel()
	if err != nil {
		return err
	}
	defer ch.Close()

	err = ch.Qos(subCount, 0, false)
	if err != nil {
		return err
	}

	messages, err := ch.Consume(queueName, "processor", false, false, false, false, nil)
	if err != nil {
		return err
	}

	var wg sync.WaitGroup
	wg.Add(subCount)
	for i := 0; i < subCount; i++ {
		go func() {
			defer wg.Done()

			for {
				select {
				case <-gracefulQuitChan:
					return
				case message, ok := <-messages:
					if !ok {
						return
					}

					f().Process(message.Body)
					message.Ack(false)
				}
			}
		}()
	}
	wg.Wait()

	return nil
}

func (mb *RabbitMessageBroker) SubscribeFanout(exchange string, f func() MessageProcessor) error {
	ch, err := mb.conn.Channel()
	if err != nil {
		return err
	}
	defer ch.Close()

	err = ch.Qos(1, 0, false)
	if err != nil {
		return err
	}

	err = ch.ExchangeDeclare(exchange, "fanout", false, false, false, false, nil)
	if err != nil {
		return err
	}

	queue, err := ch.QueueDeclare("", false, false, true, false, nil)
	if err != nil {
		return err
	}

	err = ch.QueueBind(queue.Name, "", exchange, false, nil)
	if err != nil {
		return err
	}

	messages, err := ch.Consume(queue.Name, "", false, false, false, false, nil)
	if err != nil {
		return err
	}

	for message := range messages {
		f().Process(message.Body)
		message.Ack(false)
	}

	return nil
}

func (mb *RabbitMessageBroker) Close() error {
	err := mb.publishChan.Close()
	if err != nil {
		return err
	}

	return mb.conn.Close()
}

func NewMessageBroker(url string) (MessageBroker, error) {
	if url == "" {
		return nil, fmt.Errorf("URL is blank")
	}

	conn, err := amqp.Dial(url)
	if err != nil {
		return nil, err
	}

	publishChan, err := conn.Channel()
	if err != nil {
		return nil, err
	}

	return &RabbitMessageBroker{conn, publishChan}, nil
}

func NewTestMessageBroker() MessageBroker {
	return &TestMessageBroker{
		queues: make(map[string]chan []byte),
	}
}

func (mb *TestMessageBroker) DeclareQueue(queueName string) error {
	mb.queues[queueName] = make(chan []byte, 1)
	return nil
}

func (mb *TestMessageBroker) Subscribe(queueName string, subCount int, gracefulQuitChan chan bool, f func() MessageProcessor) error {
	var wg sync.WaitGroup
	wg.Add(subCount)
	for i := 0; i < subCount; i++ {
		go func() {
			defer wg.Done()

			processor := f()

			for body := range mb.queues[queueName] {
				processor.Process(body)
			}
		}()
	}
	wg.Wait()

	return nil
}

func (mb *TestMessageBroker) SubscribeFanout(exchange string, f func() MessageProcessor) error {
	return fmt.Errorf("TestMessageBroker.SubscribeFanout needs to be implemented")
}

func (mb *TestMessageBroker) Publish(exchange, routingKey, msgType string, msg []byte) error {
	mb.queues[routingKey] <- msg
	return nil
}

func (mb *TestMessageBroker) Close() error {
	for _, ch := range mb.queues {
		close(ch)
	}
	return nil
}
<|endoftext|>


<commit_before>package oauth

import (
	"errors"
	"fmt"
	"github.com/telemetryapp/gotelemetry"
	"github.com/telemetryapp/gotelemetry_agent/agent/aggregations"
	"github.com/telemetryapp/gotelemetry_agent/agent/config"
	"golang.org/x/oauth2"
)

var entries map[string]config.OAuthConfigEntry

func Init(e map[string]config.OAuthConfigEntry) {
	entries = e

	aggregations.InitOAuthStorage()
}

func configForEntryWithName(name string) (*oauth2.Config, error) {
	entry, ok := entries[name]

	if !ok {
		return nil, fmt.Errorf("oAuth entry %s not found", name)
	}

	res := &oauth2.Config{
		ClientID:     entry.ClientID,
		ClientSecret: entry.ClientSecret,
		Scopes:       entry.Scopes,
		Endpoint: oauth2.Endpoint{
			AuthURL:  entry.AuthorizationURL,
			TokenURL: entry.TokenURL,
		},
		RedirectURL: "https://qa-www.telemetryapp.com/oauth_response",
	}

	return res, nil
}

func tokenForEntryWithName(name string) (*oauth2.Token, error) {
	if _, ok := entries[name]; !ok {
		return nil, fmt.Errorf("oAuth entry %s not found", name)
	}

	return aggregations.ReadOAuthToken(name)
}

func writeTokenForEntryWithName(name string, token *oauth2.Token) error {
	if _, ok := entries[name]; !ok {
		return fmt.Errorf("oAuth entry %s not found", name)
	}

	return aggregations.WriteOAuthToken(name, token)
}

func RunCommand(cfg config.CLIConfigType, errorChannel chan error, completionChannel chan bool) {
	switch cfg.OAuthCommand {
	case config.OAuthCommands.None:
		// Do nothing
		break

	case config.OAuthCommands.Request:
		entry, err := configForEntryWithName(cfg.OAuthName)

		if err != nil {
			errorChannel <- err
			break
		}

		errorChannel <- gotelemetry.NewLogError("Please visit this URL and authorize the Agent:\n\n---\n%s\n---\n\n", entry.AuthCodeURL("s", oauth2.AccessTypeOffline, oauth2.ApprovalForce))
		errorChannel <- gotelemetry.NewLogError("When you are done, please run the agent with the oauth-exchange command to set the new token.\n\n")

	case config.OAuthCommands.Exchange:

		entry, err := configForEntryWithName(cfg.OAuthName)

		if err != nil {
			errorChannel <- err
			break
		}

		code := cfg.OAuthCode

		if code == "" {
			errorChannel <- errors.New("No authorization code found. Please provide one with -c.")
			break
		}

		token, err := entry.Exchange(oauth2.NoContext, code)

		if err != nil {
			errorChannel <- err
			break
		}

		if err := writeTokenForEntryWithName(cfg.OAuthName, token); err != nil {
			errorChannel <- err
			break
		}

		errorChannel <- gotelemetry.NewLogError("Token exchanged successfully. The entry %s can now be used to make authenticated calls.", cfg.OAuthName)

	default:
		errorChannel <- fmt.Errorf("Unknown oauth command %s", cfg.OAuthCommand)
	}

	completionChannel <- true
}
<commit_msg>- Use proper redirect<commit_after>package oauth

import (
	"errors"
	"fmt"
	"github.com/telemetryapp/gotelemetry"
	"github.com/telemetryapp/gotelemetry_agent/agent/aggregations"
	"github.com/telemetryapp/gotelemetry_agent/agent/config"
	"golang.org/x/oauth2"
)

var entries map[string]config.OAuthConfigEntry

func Init(e map[string]config.OAuthConfigEntry) {
	entries = e

	aggregations.InitOAuthStorage()
}

func configForEntryWithName(name string) (*oauth2.Config, error) {
	entry, ok := entries[name]

	if !ok {
		return nil, fmt.Errorf("oAuth entry %s not found", name)
	}

	res := &oauth2.Config{
		ClientID:     entry.ClientID,
		ClientSecret: entry.ClientSecret,
		Scopes:       entry.Scopes,
		Endpoint: oauth2.Endpoint{
			AuthURL:  entry.AuthorizationURL,
			TokenURL: entry.TokenURL,
		},
		RedirectURL: "https://telemetryapp.com/oauth_response",
	}

	return res, nil
}

func tokenForEntryWithName(name string) (*oauth2.Token, error) {
	if _, ok := entries[name]; !ok {
		return nil, fmt.Errorf("oAuth entry %s not found", name)
	}

	return aggregations.ReadOAuthToken(name)
}

func writeTokenForEntryWithName(name string, token *oauth2.Token) error {
	if _, ok := entries[name]; !ok {
		return fmt.Errorf("oAuth entry %s not found", name)
	}

	return aggregations.WriteOAuthToken(name, token)
}

func RunCommand(cfg config.CLIConfigType, errorChannel chan error, completionChannel chan bool) {
	switch cfg.OAuthCommand {
	case config.OAuthCommands.None:
		// Do nothing
		break

	case config.OAuthCommands.Request:
		entry, err := configForEntryWithName(cfg.OAuthName)

		if err != nil {
			errorChannel <- err
			break
		}

		errorChannel <- gotelemetry.NewLogError("Please visit this URL and authorize the Agent:\n\n---\n%s\n---\n\n", entry.AuthCodeURL("s", oauth2.AccessTypeOffline, oauth2.ApprovalForce))
		errorChannel <- gotelemetry.NewLogError("When you are done, please run the agent with the oauth-exchange command to set the new token.\n\n")

	case config.OAuthCommands.Exchange:

		entry, err := configForEntryWithName(cfg.OAuthName)

		if err != nil {
			errorChannel <- err
			break
		}

		code := cfg.OAuthCode

		if code == "" {
			errorChannel <- errors.New("No authorization code found. Please provide one with -c.")
			break
		}

		token, err := entry.Exchange(oauth2.NoContext, code)

		if err != nil {
			errorChannel <- err
			break
		}

		if err := writeTokenForEntryWithName(cfg.OAuthName, token); err != nil {
			errorChannel <- err
			break
		}

		errorChannel <- gotelemetry.NewLogError("Token exchanged successfully. The entry %s can now be used to make authenticated calls.", cfg.OAuthName)

	default:
		errorChannel <- fmt.Errorf("Unknown oauth command %s", cfg.OAuthCommand)
	}

	completionChannel <- true
}
<|endoftext|>


<commit_before>package alarmy

import (
	"net/http"

	"github.com/pressly/chi/render"
)

func (a *Api) ProjectAll(w http.ResponseWriter, r *http.Request) {
	projects, err := a.store.ProjectAll()
	if err != nil {
		http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
		return
	}
	render.JSON(w, r, projects)
}

func (a *Api) ProjectCreate(w http.ResponseWriter, r *http.Request) {
	data := &ProjectRequest{}
	if err := render.Bind(r, data); err != nil {
		http.Error(w, http.StatusText(http.StatusUnprocessableEntity), http.StatusUnprocessableEntity)
		return
	}

	project, err := a.store.ProjectCreate(data.Project)
	if err != nil {
		http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
		return
	}

	render.Status(r, http.StatusCreated)
	render.JSON(w, r, project)
}
<commit_msg>Return empty array instead of nil<commit_after>package alarmy

import (
	"net/http"

	"github.com/pressly/chi/render"
)

func (a *Api) ProjectAll(w http.ResponseWriter, r *http.Request) {
	projects, err := a.store.ProjectAll()
	if err != nil {
		http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
		return
	}

	if projects == nil {
		projects = []Project{}
	}

	render.JSON(w, r, projects)
}

func (a *Api) ProjectCreate(w http.ResponseWriter, r *http.Request) {
	data := &ProjectRequest{}
	if err := render.Bind(r, data); err != nil {
		http.Error(w, http.StatusText(http.StatusUnprocessableEntity), http.StatusUnprocessableEntity)
		return
	}

	project, err := a.store.ProjectCreate(data.Project)
	if err != nil {
		http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
		return
	}

	render.Status(r, http.StatusCreated)
	render.JSON(w, r, project)
}
<|endoftext|>


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

package models

import (
	"bytes"
	"errors"
	"fmt"
	"path"
	"strings"
	"time"

	"github.com/Unknwon/com"

	"github.com/Unknwon/gowalker/modules/base"
	"github.com/Unknwon/gowalker/modules/setting"
)

var (
	ErrEmptyPackagePath     = errors.New("Package import path is empty")
	ErrPackageNotFound      = errors.New("Package does not found")
	ErrPackageVersionTooOld = errors.New("Package version is too old")
)

// PkgInfo represents the package information.
type PkgInfo struct {
	ID         int64  `xorm:"pk autoincr"`
	Name       string `xorm:"-"`
	ImportPath string `xorm:"UNIQUE"`
	Etag       string

	ProjectPath string
	ViewDirPath string
	Synopsis    string

	IsCmd       bool
	IsCgo       bool
	IsGoRepo    bool
	IsGoSubrepo bool

	PkgVer int

	Views int64
	// Indicate how many JS should be downloaded(JsNum=total num - 1)
	JsNum int

	ImportNum int64
	ImportIDs string `xorm:"import_ids TEXT"`
	// Import num usually is small so save it to reduce a database query.
	ImportPaths string `xorm:"TEXT"`

	RefNum int64
	RefIDs string `xorm:"ref_ids TEXT"`

	Subdirs string `xorm:"TEXT"`

	LastView int64 `xorm:"-"`
	Created  int64
}

func (p *PkgInfo) JSPath() string {
	return path.Join(setting.DocsJsPath, p.ImportPath) + ".js"
}

// CanRefresh returns true if package is available to refresh.
func (p *PkgInfo) CanRefresh() bool {
	return time.Now().UTC().Add(-1*setting.RefreshInterval).Unix() > p.Created
}

// GetRefs returns a list of packages that import this one.
func (p *PkgInfo) GetRefs() []*PkgInfo {
	pinfos := make([]*PkgInfo, 0, p.RefNum)
	refIDs := strings.Split(p.RefIDs, "|")
	for i := range refIDs {
		if len(refIDs[i]) == 0 {
			continue
		}

		id := com.StrTo(refIDs[i][1:]).MustInt64()
		if pinfo, _ := GetPkgInfoById(id); pinfo != nil {
			pinfo.Name = path.Base(pinfo.ImportPath)
			pinfos = append(pinfos, pinfo)
		}
	}
	return pinfos
}

// PACKAGE_VER is modified when previously stored packages are invalid.
const PACKAGE_VER = 1

// PkgRef represents temporary reference information of a package.
type PkgRef struct {
	ID         int64  `xorm:"pk autoincr"`
	ImportPath string `xorm:"UNIQUE"`
	RefNum     int64
	RefIDs     string `xorm:"ref_ids TEXT"`
}

func updatePkgRef(pid int64, refPath string) error {
	if base.IsGoRepoPath(refPath) {
		return nil
	}

	ref := new(PkgRef)
	has, err := x.Where("import_path=?", refPath).Get(ref)
	if err != nil {
		return fmt.Errorf("get PkgRef: %v", err)
	}

	queryStr := "$" + com.ToStr(pid) + "|"
	if !has {
		if _, err = x.Insert(&PkgRef{
			ImportPath: refPath,
			RefNum:     1,
			RefIDs:     queryStr,
		}); err != nil {
			return fmt.Errorf("insert PkgRef: %v", err)
		}
		return nil
	}

	i := strings.Index(ref.RefIDs, queryStr)
	if i > -1 {
		return nil
	}

	ref.RefIDs += queryStr
	ref.RefNum++
	_, err = x.Id(ref.ID).AllCols().Update(ref)
	return err
}

// checkRefs checks if given packages are still referencing this one.
func checkRefs(pinfo *PkgInfo) {
	var buf bytes.Buffer
	pinfo.RefNum = 0
	refIDs := strings.Split(pinfo.RefIDs, "|")
	for i := range refIDs {
		if len(refIDs[i]) == 0 {
			continue
		}

		fmt.Println(com.StrTo(refIDs[i][1:]).MustInt64())
		pkg, _ := GetPkgInfoById(com.StrTo(refIDs[i][1:]).MustInt64())
		if pkg == nil {
			continue
		}

		if strings.Index(pkg.ImportIDs, "$"+com.ToStr(pinfo.ID)+"|") == -1 {
			continue
		}

		buf.WriteString("$")
		buf.WriteString(com.ToStr(pkg.ID))
		buf.WriteString("|")
		pinfo.RefNum++
	}
	pinfo.RefIDs = buf.String()
}

// updateRef updates or crates corresponding reference import information.
func updateRef(pid int64, refPath string) (int64, error) {
	pinfo, err := GetPkgInfo(refPath)
	if err != nil && pinfo == nil {
		if err == ErrPackageNotFound ||
			err == ErrPackageVersionTooOld {
			// Package hasn't existed yet, save to temporary place.
			return 0, updatePkgRef(pid, refPath)
		}
		return 0, fmt.Errorf("GetPkgInfo(%s): %v", refPath, err)
	}

	// Check if reference information has beed recorded.
	queryStr := "$" + com.ToStr(pid) + "|"
	i := strings.Index(pinfo.RefIDs, queryStr)
	if i > -1 {
		return pinfo.ID, nil
	}

	// Add new as needed.
	pinfo.RefIDs += queryStr
	pinfo.RefNum++
	_, err = x.Id(pinfo.ID).AllCols().Update(pinfo)
	return pinfo.ID, err
}

// SavePkgInfo saves package information.
func SavePkgInfo(pinfo *PkgInfo, updateRefs bool) (err error) {
	pinfo.PkgVer = PACKAGE_VER

	// Create or update package info itself.
	// Note(Unknwon): do this because we need ID field later.
	if pinfo.ID == 0 {
		pinfo.Views = 1

		// First time created, check PkgRef.
		ref := new(PkgRef)
		has, err := x.Where("import_path=?", pinfo.ImportPath).Get(ref)
		if err != nil {
			return fmt.Errorf("get PkgRef: %v", err)
		} else if has {
			pinfo.RefNum = ref.RefNum
			pinfo.RefIDs = ref.RefIDs
			if _, err = x.Id(ref.ID).Delete(ref); err != nil {
				return fmt.Errorf("delete PkgRef: %v", err)
			}
		}

		_, err = x.Insert(pinfo)
	} else {
		_, err = x.Id(pinfo.ID).AllCols().Update(pinfo)
	}
	if err != nil {
		return fmt.Errorf("update package info: %v", err)
	}

	// Update package import references.
	// Note(Unknwon): I just don't see the value of who imports STD
	//	when you don't even import and uses what objects.
	if updateRefs && !pinfo.IsGoRepo {
		var buf bytes.Buffer
		paths := strings.Split(pinfo.ImportPaths, "|")
		for i := range paths {
			refID, err := updateRef(pinfo.ID, paths[i])
			if err != nil {
				return fmt.Errorf("updateRef: %v", err)
			} else if refID == 0 {
				continue
			}
			buf.WriteString("$")
			buf.WriteString(com.ToStr(refID))
			buf.WriteString("|")
		}
		pinfo.ImportIDs = buf.String()

		// Check packages who import this is still importing.
		checkRefs(pinfo)
		_, err = x.Id(pinfo.ID).AllCols().Update(pinfo)
		return err
	}
	return nil
}

// GetPkgInfo returns package information by given import path.
func GetPkgInfo(importPath string) (*PkgInfo, error) {
	if len(importPath) == 0 {
		return nil, ErrEmptyPackagePath
	}

	pinfo := new(PkgInfo)
	has, err := x.Where("import_path=?", importPath).Get(pinfo)
	if err != nil {
		return nil, err
	} else if !has {
		return nil, ErrPackageNotFound
	} else if pinfo.PkgVer < PACKAGE_VER {
		pinfo.Etag = ""
		return pinfo, ErrPackageVersionTooOld
	}

	if !com.IsFile(pinfo.JSPath()) {
		pinfo.Etag = ""
		return pinfo, ErrPackageVersionTooOld
	}

	return pinfo, nil
}

// GetSubPkgs returns sub-projects by given sub-directories.
func GetSubPkgs(importPath string, dirs []string) []*PkgInfo {
	pinfos := make([]*PkgInfo, 0, len(dirs))
	for _, dir := range dirs {
		if len(dir) == 0 {
			continue
		}

		fullPath := importPath + "/" + dir
		if pinfo, err := GetPkgInfo(fullPath); err == nil {
			pinfo.Name = dir
			pinfos = append(pinfos, pinfo)
		} else {
			pinfos = append(pinfos, &PkgInfo{
				Name:       dir,
				ImportPath: fullPath,
			})
		}
	}
	return pinfos
}

// GetPkgInfosByPaths returns a list of packages by given import paths.
func GetPkgInfosByPaths(paths []string) []*PkgInfo {
	pinfos := make([]*PkgInfo, 0, len(paths))
	for _, p := range paths {
		if len(p) == 0 {
			continue
		}

		if pinfo, err := GetPkgInfo(p); err == nil {
			pinfo.Name = path.Base(p)
			pinfos = append(pinfos, pinfo)
		} else {
			pinfos = append(pinfos, &PkgInfo{
				Name:       path.Base(p),
				ImportPath: p,
			})
		}
	}
	return pinfos
}

// GetPkgInfoById returns package information by given ID.
func GetPkgInfoById(id int64) (*PkgInfo, error) {
	pinfo := new(PkgInfo)
	has, err := x.Id(id).Get(pinfo)
	if err != nil {
		return nil, err
	} else if !has {
		return nil, ErrPackageNotFound
	} else if pinfo.PkgVer < PACKAGE_VER {
		return pinfo, ErrPackageVersionTooOld
	}

	if !com.IsFile(pinfo.JSPath()) {
		return pinfo, ErrPackageVersionTooOld
	}

	return pinfo, nil
}

// SearchPkgInfo searches package information by given keyword.
func SearchPkgInfo(limit int, keyword string) ([]*PkgInfo, error) {
	if len(keyword) == 0 {
		return nil, nil
	}
	pkgs := make([]*PkgInfo, 0, limit)
	return pkgs, x.Limit(limit).Desc("views").Where("import_path like ?", "%"+keyword+"%").Find(&pkgs)
}
<commit_msg>fix missing check<commit_after>// Copyright 2015 Unknwon
//
// Licensed under the Apache License, Version 2.0 (the "License"): you may
// not use this file except in compliance with the License. You may obtain
// a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
// License for the specific language governing permissions and limitations
// under the License.

package models

import (
	"bytes"
	"errors"
	"fmt"
	"path"
	"strings"
	"time"

	"github.com/Unknwon/com"

	"github.com/Unknwon/gowalker/modules/base"
	"github.com/Unknwon/gowalker/modules/setting"
)

var (
	ErrEmptyPackagePath     = errors.New("Package import path is empty")
	ErrPackageNotFound      = errors.New("Package does not found")
	ErrPackageVersionTooOld = errors.New("Package version is too old")
)

// PkgInfo represents the package information.
type PkgInfo struct {
	ID         int64  `xorm:"pk autoincr"`
	Name       string `xorm:"-"`
	ImportPath string `xorm:"UNIQUE"`
	Etag       string

	ProjectPath string
	ViewDirPath string
	Synopsis    string

	IsCmd       bool
	IsCgo       bool
	IsGoRepo    bool
	IsGoSubrepo bool

	PkgVer int

	Views int64
	// Indicate how many JS should be downloaded(JsNum=total num - 1)
	JsNum int

	ImportNum int64
	ImportIDs string `xorm:"import_ids TEXT"`
	// Import num usually is small so save it to reduce a database query.
	ImportPaths string `xorm:"TEXT"`

	RefNum int64
	RefIDs string `xorm:"ref_ids TEXT"`

	Subdirs string `xorm:"TEXT"`

	LastView int64 `xorm:"-"`
	Created  int64
}

func (p *PkgInfo) JSPath() string {
	return path.Join(setting.DocsJsPath, p.ImportPath) + ".js"
}

// CanRefresh returns true if package is available to refresh.
func (p *PkgInfo) CanRefresh() bool {
	return time.Now().UTC().Add(-1*setting.RefreshInterval).Unix() > p.Created
}

// GetRefs returns a list of packages that import this one.
func (p *PkgInfo) GetRefs() []*PkgInfo {
	pinfos := make([]*PkgInfo, 0, p.RefNum)
	refIDs := strings.Split(p.RefIDs, "|")
	for i := range refIDs {
		if len(refIDs[i]) == 0 {
			continue
		}

		id := com.StrTo(refIDs[i][1:]).MustInt64()
		if pinfo, _ := GetPkgInfoById(id); pinfo != nil {
			pinfo.Name = path.Base(pinfo.ImportPath)
			pinfos = append(pinfos, pinfo)
		}
	}
	return pinfos
}

// PACKAGE_VER is modified when previously stored packages are invalid.
const PACKAGE_VER = 1

// PkgRef represents temporary reference information of a package.
type PkgRef struct {
	ID         int64  `xorm:"pk autoincr"`
	ImportPath string `xorm:"UNIQUE"`
	RefNum     int64
	RefIDs     string `xorm:"ref_ids TEXT"`
}

func updatePkgRef(pid int64, refPath string) error {
	if base.IsGoRepoPath(refPath) || refPath == "C" {
		return nil
	}

	ref := new(PkgRef)
	has, err := x.Where("import_path=?", refPath).Get(ref)
	if err != nil {
		return fmt.Errorf("get PkgRef: %v", err)
	}

	queryStr := "$" + com.ToStr(pid) + "|"
	if !has {
		if _, err = x.Insert(&PkgRef{
			ImportPath: refPath,
			RefNum:     1,
			RefIDs:     queryStr,
		}); err != nil {
			return fmt.Errorf("insert PkgRef: %v", err)
		}
		return nil
	}

	i := strings.Index(ref.RefIDs, queryStr)
	if i > -1 {
		return nil
	}

	ref.RefIDs += queryStr
	ref.RefNum++
	_, err = x.Id(ref.ID).AllCols().Update(ref)
	return err
}

// checkRefs checks if given packages are still referencing this one.
func checkRefs(pinfo *PkgInfo) {
	var buf bytes.Buffer
	pinfo.RefNum = 0
	refIDs := strings.Split(pinfo.RefIDs, "|")
	for i := range refIDs {
		if len(refIDs[i]) == 0 {
			continue
		}

		fmt.Println(com.StrTo(refIDs[i][1:]).MustInt64())
		pkg, _ := GetPkgInfoById(com.StrTo(refIDs[i][1:]).MustInt64())
		if pkg == nil {
			continue
		}

		if strings.Index(pkg.ImportIDs, "$"+com.ToStr(pinfo.ID)+"|") == -1 {
			continue
		}

		buf.WriteString("$")
		buf.WriteString(com.ToStr(pkg.ID))
		buf.WriteString("|")
		pinfo.RefNum++
	}
	pinfo.RefIDs = buf.String()
}

// updateRef updates or crates corresponding reference import information.
func updateRef(pid int64, refPath string) (int64, error) {
	pinfo, err := GetPkgInfo(refPath)
	if err != nil && pinfo == nil {
		if err == ErrPackageNotFound ||
			err == ErrPackageVersionTooOld {
			// Package hasn't existed yet, save to temporary place.
			return 0, updatePkgRef(pid, refPath)
		}
		return 0, fmt.Errorf("GetPkgInfo(%s): %v", refPath, err)
	}

	// Check if reference information has beed recorded.
	queryStr := "$" + com.ToStr(pid) + "|"
	i := strings.Index(pinfo.RefIDs, queryStr)
	if i > -1 {
		return pinfo.ID, nil
	}

	// Add new as needed.
	pinfo.RefIDs += queryStr
	pinfo.RefNum++
	_, err = x.Id(pinfo.ID).AllCols().Update(pinfo)
	return pinfo.ID, err
}

// SavePkgInfo saves package information.
func SavePkgInfo(pinfo *PkgInfo, updateRefs bool) (err error) {
	pinfo.PkgVer = PACKAGE_VER

	// Create or update package info itself.
	// Note(Unknwon): do this because we need ID field later.
	if pinfo.ID == 0 {
		pinfo.Views = 1

		// First time created, check PkgRef.
		ref := new(PkgRef)
		has, err := x.Where("import_path=?", pinfo.ImportPath).Get(ref)
		if err != nil {
			return fmt.Errorf("get PkgRef: %v", err)
		} else if has {
			pinfo.RefNum = ref.RefNum
			pinfo.RefIDs = ref.RefIDs
			if _, err = x.Id(ref.ID).Delete(ref); err != nil {
				return fmt.Errorf("delete PkgRef: %v", err)
			}
		}

		_, err = x.Insert(pinfo)
	} else {
		_, err = x.Id(pinfo.ID).AllCols().Update(pinfo)
	}
	if err != nil {
		return fmt.Errorf("update package info: %v", err)
	}

	// Update package import references.
	// Note(Unknwon): I just don't see the value of who imports STD
	//	when you don't even import and uses what objects.
	if updateRefs && !pinfo.IsGoRepo {
		var buf bytes.Buffer
		paths := strings.Split(pinfo.ImportPaths, "|")
		for i := range paths {
			refID, err := updateRef(pinfo.ID, paths[i])
			if err != nil {
				return fmt.Errorf("updateRef: %v", err)
			} else if refID == 0 {
				continue
			}
			buf.WriteString("$")
			buf.WriteString(com.ToStr(refID))
			buf.WriteString("|")
		}
		pinfo.ImportIDs = buf.String()

		// Check packages who import this is still importing.
		checkRefs(pinfo)
		_, err = x.Id(pinfo.ID).AllCols().Update(pinfo)
		return err
	}
	return nil
}

// GetPkgInfo returns package information by given import path.
func GetPkgInfo(importPath string) (*PkgInfo, error) {
	if len(importPath) == 0 {
		return nil, ErrEmptyPackagePath
	}

	pinfo := new(PkgInfo)
	has, err := x.Where("import_path=?", importPath).Get(pinfo)
	if err != nil {
		return nil, err
	} else if !has {
		return nil, ErrPackageNotFound
	} else if pinfo.PkgVer < PACKAGE_VER {
		pinfo.Etag = ""
		return pinfo, ErrPackageVersionTooOld
	}

	if !com.IsFile(pinfo.JSPath()) {
		pinfo.Etag = ""
		return pinfo, ErrPackageVersionTooOld
	}

	return pinfo, nil
}

// GetSubPkgs returns sub-projects by given sub-directories.
func GetSubPkgs(importPath string, dirs []string) []*PkgInfo {
	pinfos := make([]*PkgInfo, 0, len(dirs))
	for _, dir := range dirs {
		if len(dir) == 0 {
			continue
		}

		fullPath := importPath + "/" + dir
		if pinfo, err := GetPkgInfo(fullPath); err == nil {
			pinfo.Name = dir
			pinfos = append(pinfos, pinfo)
		} else {
			pinfos = append(pinfos, &PkgInfo{
				Name:       dir,
				ImportPath: fullPath,
			})
		}
	}
	return pinfos
}

// GetPkgInfosByPaths returns a list of packages by given import paths.
func GetPkgInfosByPaths(paths []string) []*PkgInfo {
	pinfos := make([]*PkgInfo, 0, len(paths))
	for _, p := range paths {
		if len(p) == 0 {
			continue
		}

		if pinfo, err := GetPkgInfo(p); err == nil {
			pinfo.Name = path.Base(p)
			pinfos = append(pinfos, pinfo)
		} else {
			pinfos = append(pinfos, &PkgInfo{
				Name:       path.Base(p),
				ImportPath: p,
			})
		}
	}
	return pinfos
}

// GetPkgInfoById returns package information by given ID.
func GetPkgInfoById(id int64) (*PkgInfo, error) {
	pinfo := new(PkgInfo)
	has, err := x.Id(id).Get(pinfo)
	if err != nil {
		return nil, err
	} else if !has {
		return nil, ErrPackageNotFound
	} else if pinfo.PkgVer < PACKAGE_VER {
		return pinfo, ErrPackageVersionTooOld
	}

	if !com.IsFile(pinfo.JSPath()) {
		return pinfo, ErrPackageVersionTooOld
	}

	return pinfo, nil
}

// SearchPkgInfo searches package information by given keyword.
func SearchPkgInfo(limit int, keyword string) ([]*PkgInfo, error) {
	if len(keyword) == 0 {
		return nil, nil
	}
	pkgs := make([]*PkgInfo, 0, limit)
	return pkgs, x.Limit(limit).Desc("views").Where("import_path like ?", "%"+keyword+"%").Find(&pkgs)
}
<|endoftext|>


<commit_before>package mstate

import (
	"labix.org/v2/mgo"
	"launchpad.net/juju-core/mstate/watcher"
	"launchpad.net/tomb"
)

// commonWatcher is part of all client watchers.
type commonWatcher struct {
	st   *State
	tomb tomb.Tomb
}

// MachineWatcher observes changes to the settings of a machine.
type MachineWatcher struct {
	commonWatcher
	changeChan chan *Machine
}

// MachinesWatcher notifies about machines being added or removed
// from the environment.
type MachinesWatcher struct {
	commonWatcher
	changeChan    chan *MachinesChange
	knownMachines map[int]*Machine
}

// MachinesChange contains information about
// machines that have been added or deleted.
type MachinesChange struct {
	Added   []*Machine
	Removed []*Machine
}

// newMachineWatcher creates and starts a watcher to watch information
// about the machine.
func newMachineWatcher(m *Machine) *MachineWatcher {
	w := &MachineWatcher{
		changeChan:    make(chan *Machine),
		commonWatcher: commonWatcher{st: m.st},
	}
	go func() {
		defer w.tomb.Done()
		defer close(w.changeChan)
		w.tomb.Kill(w.loop(m))
	}()
	return w
}

// Changes returns a channel that will receive the new
// *Machine when a change is detected. Note that multiple
// changes may be observed as a single event in the channel.
// The first event on the channel holds the initial state
// as returned by Machine.Info.
func (w *MachineWatcher) Changes() <-chan *Machine {
	return w.changeChan
}

func (w *MachineWatcher) Stop() error {
	w.tomb.Kill(nil)
	return w.tomb.Wait()
}

func (w *MachineWatcher) loop(m *Machine) (err error) {
	ch := make(chan watcher.Change)
	id := m.Id()
	st := m.st
	st.watcher.Watch(st.machines.Name, id, m.doc.TxnRevno, ch)
	defer st.watcher.Unwatch(st.machines.Name, id, ch)
	for {
		select {
		case <-st.watcher.Dead():
			return watcher.MustErr(st.watcher)
		case <-w.tomb.Dying():
			return tomb.ErrDying
		case <-ch:
		}
		if m, err = st.Machine(id); err != nil {
			return err
		}
		for {
			select {
			case <-st.watcher.Dead():
				return watcher.MustErr(st.watcher)
			case <-w.tomb.Dying():
				return tomb.ErrDying
			case <-ch:
				if err := m.Refresh(); err != nil {
					return err
				}
				continue
			case w.changeChan <- m:
			}
			break
		}
	}
	return nil
}

// WatchMachines returns a watcher for observing machines being
// added or removed.
func (s *State) WatchMachines() *MachinesWatcher {
	return newMachinesWatcher(s)
}

// newMachinesWatcher creates and starts a watcher to watch information
// about machines being added or deleted.
func newMachinesWatcher(st *State) *MachinesWatcher {
	w := &MachinesWatcher{
		changeChan:    make(chan *MachinesChange),
		knownMachines: make(map[int]*Machine),
		commonWatcher: commonWatcher{st: st},
	}
	go func() {
		defer w.tomb.Done()
		defer close(w.changeChan)
		w.tomb.Kill(w.loop())
	}()
	return w
}

// Changes returns a channel that will receive changes when machines are
// added or deleted. The Added field in the first event on the channel
// holds the initial state as returned by State.AllMachines.
func (w *MachinesWatcher) Changes() <-chan *MachinesChange {
	return w.changeChan
}

// Stop stops the watcher and returns any errors encountered while watching.
func (w *MachinesWatcher) Stop() error {
	w.tomb.Kill(nil)
	return w.tomb.Wait()
}

func (w *MachinesWatcher) appendChange(changes *MachinesChange, ch watcher.Change) (err error) {
	id := ch.Id.(int)
	if m, ok := w.knownMachines[id]; ch.Revno == -1 && ok {
		m.doc.Life = Dead
		changes.Removed = append(changes.Removed, m)
		delete(w.knownMachines, id)
		return nil
	}
	doc := &machineDoc{}
	err = w.st.machines.FindId(id).One(doc)
	if err == mgo.ErrNotFound {
		return nil
	}
	if err != nil {
		return err
	}
	m := newMachine(w.st, doc)
	if _, ok := w.knownMachines[id]; !ok {
		changes.Added = append(changes.Added, m)
	}
	w.knownMachines[id] = m
	return nil
}

func (changes *MachinesChange) isEmpty() bool {
	return len(changes.Added)+len(changes.Removed) == 0
}

func (w *MachinesWatcher) getInitialEvent() (initial *MachinesChange, err error) {
	changes := &MachinesChange{}
	docs := []machineDoc{}
	err = w.st.machines.Find(nil).All(&docs)
	if err != nil {
		return nil, err
	}
	for _, doc := range docs {
		m := newMachine(w.st, &doc)
		w.knownMachines[doc.Id] = m
		changes.Added = append(changes.Added, m)
	}
	return changes, nil
}

func (w *MachinesWatcher) loop() (err error) {
	ch := make(chan watcher.Change)
	w.st.watcher.WatchCollection(w.st.machines.Name, ch)
	defer w.st.watcher.UnwatchCollection(w.st.machines.Name, ch)
	changes, err := w.getInitialEvent()
	if err != nil {
		return err
	}
	for {
		if changes == nil {
			select {
			case <-w.st.watcher.Dead():
				return watcher.MustErr(w.st.watcher)
			case <-w.tomb.Dying():
				return tomb.ErrDying
			case c := <-ch:
				changes = &MachinesChange{}
				err := w.appendChange(changes, c)
				if err != nil {
					return err
				}
				if changes.isEmpty() {
					changes = nil
				}
				continue
			}
		}
		select {
		case <-w.st.watcher.Dead():
			return watcher.MustErr(w.st.watcher)
		case <-w.tomb.Dying():
			return tomb.ErrDying
		case c := <-ch:
			err := w.appendChange(changes, c)
			if err != nil {
				return err
			}
		case w.changeChan <- changes:
			changes = nil
		}
	}
	return nil
}
<commit_msg>mstate: s/appendChange/mergeChange/g<commit_after>package mstate

import (
	"labix.org/v2/mgo"
	"launchpad.net/juju-core/mstate/watcher"
	"launchpad.net/tomb"
)

// commonWatcher is part of all client watchers.
type commonWatcher struct {
	st   *State
	tomb tomb.Tomb
}

// MachineWatcher observes changes to the settings of a machine.
type MachineWatcher struct {
	commonWatcher
	changeChan chan *Machine
}

// MachinesWatcher notifies about machines being added or removed
// from the environment.
type MachinesWatcher struct {
	commonWatcher
	changeChan    chan *MachinesChange
	knownMachines map[int]*Machine
}

// MachinesChange contains information about
// machines that have been added or deleted.
type MachinesChange struct {
	Added   []*Machine
	Removed []*Machine
}

// newMachineWatcher creates and starts a watcher to watch information
// about the machine.
func newMachineWatcher(m *Machine) *MachineWatcher {
	w := &MachineWatcher{
		changeChan:    make(chan *Machine),
		commonWatcher: commonWatcher{st: m.st},
	}
	go func() {
		defer w.tomb.Done()
		defer close(w.changeChan)
		w.tomb.Kill(w.loop(m))
	}()
	return w
}

// Changes returns a channel that will receive the new
// *Machine when a change is detected. Note that multiple
// changes may be observed as a single event in the channel.
// The first event on the channel holds the initial state
// as returned by Machine.Info.
func (w *MachineWatcher) Changes() <-chan *Machine {
	return w.changeChan
}

func (w *MachineWatcher) Stop() error {
	w.tomb.Kill(nil)
	return w.tomb.Wait()
}

func (w *MachineWatcher) loop(m *Machine) (err error) {
	ch := make(chan watcher.Change)
	id := m.Id()
	st := m.st
	st.watcher.Watch(st.machines.Name, id, m.doc.TxnRevno, ch)
	defer st.watcher.Unwatch(st.machines.Name, id, ch)
	for {
		select {
		case <-st.watcher.Dead():
			return watcher.MustErr(st.watcher)
		case <-w.tomb.Dying():
			return tomb.ErrDying
		case <-ch:
		}
		if m, err = st.Machine(id); err != nil {
			return err
		}
		for {
			select {
			case <-st.watcher.Dead():
				return watcher.MustErr(st.watcher)
			case <-w.tomb.Dying():
				return tomb.ErrDying
			case <-ch:
				if err := m.Refresh(); err != nil {
					return err
				}
				continue
			case w.changeChan <- m:
			}
			break
		}
	}
	return nil
}

// WatchMachines returns a watcher for observing machines being
// added or removed.
func (s *State) WatchMachines() *MachinesWatcher {
	return newMachinesWatcher(s)
}

// newMachinesWatcher creates and starts a watcher to watch information
// about machines being added or deleted.
func newMachinesWatcher(st *State) *MachinesWatcher {
	w := &MachinesWatcher{
		changeChan:    make(chan *MachinesChange),
		knownMachines: make(map[int]*Machine),
		commonWatcher: commonWatcher{st: st},
	}
	go func() {
		defer w.tomb.Done()
		defer close(w.changeChan)
		w.tomb.Kill(w.loop())
	}()
	return w
}

// Changes returns a channel that will receive changes when machines are
// added or deleted. The Added field in the first event on the channel
// holds the initial state as returned by State.AllMachines.
func (w *MachinesWatcher) Changes() <-chan *MachinesChange {
	return w.changeChan
}

// Stop stops the watcher and returns any errors encountered while watching.
func (w *MachinesWatcher) Stop() error {
	w.tomb.Kill(nil)
	return w.tomb.Wait()
}

func (w *MachinesWatcher) mergeChange(changes *MachinesChange, ch watcher.Change) (err error) {
	id := ch.Id.(int)
	if m, ok := w.knownMachines[id]; ch.Revno == -1 && ok {
		m.doc.Life = Dead
		changes.Removed = append(changes.Removed, m)
		delete(w.knownMachines, id)
		return nil
	}
	doc := &machineDoc{}
	err = w.st.machines.FindId(id).One(doc)
	if err == mgo.ErrNotFound {
		return nil
	}
	if err != nil {
		return err
	}
	m := newMachine(w.st, doc)
	if _, ok := w.knownMachines[id]; !ok {
		changes.Added = append(changes.Added, m)
	}
	w.knownMachines[id] = m
	return nil
}

func (changes *MachinesChange) isEmpty() bool {
	return len(changes.Added)+len(changes.Removed) == 0
}

func (w *MachinesWatcher) getInitialEvent() (initial *MachinesChange, err error) {
	changes := &MachinesChange{}
	docs := []machineDoc{}
	err = w.st.machines.Find(nil).All(&docs)
	if err != nil {
		return nil, err
	}
	for _, doc := range docs {
		m := newMachine(w.st, &doc)
		w.knownMachines[doc.Id] = m
		changes.Added = append(changes.Added, m)
	}
	return changes, nil
}

func (w *MachinesWatcher) loop() (err error) {
	ch := make(chan watcher.Change)
	w.st.watcher.WatchCollection(w.st.machines.Name, ch)
	defer w.st.watcher.UnwatchCollection(w.st.machines.Name, ch)
	changes, err := w.getInitialEvent()
	if err != nil {
		return err
	}
	for {
		if changes == nil {
			select {
			case <-w.st.watcher.Dead():
				return watcher.MustErr(w.st.watcher)
			case <-w.tomb.Dying():
				return tomb.ErrDying
			case c := <-ch:
				changes = &MachinesChange{}
				err := w.mergeChange(changes, c)
				if err != nil {
					return err
				}
				if changes.isEmpty() {
					changes = nil
				}
				continue
			}
		}
		select {
		case <-w.st.watcher.Dead():
			return watcher.MustErr(w.st.watcher)
		case <-w.tomb.Dying():
			return tomb.ErrDying
		case c := <-ch:
			err := w.mergeChange(changes, c)
			if err != nil {
				return err
			}
		case w.changeChan <- changes:
			changes = nil
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package mstate

import (
	"labix.org/v2/mgo"
	"launchpad.net/juju-core/mstate/watcher"
	"launchpad.net/tomb"
)

// commonWatcher is part of all client watchers.
type commonWatcher struct {
	st   *State
	tomb tomb.Tomb
}

// MachineWatcher observes changes to the settings of a machine.
type MachineWatcher struct {
	commonWatcher
	changeChan chan *Machine
}

// MachinesWatcher notifies about machines being added or removed
// from the environment.
type MachinesWatcher struct {
	commonWatcher
	changeChan    chan *MachinesChange
	knownMachines map[int]*Machine
}

// MachinesChange contains information about
// machines that have been added or deleted.
type MachinesChange struct {
	Added   []*Machine
	Removed []*Machine
}

// newMachineWatcher creates and starts a watcher to watch information
// about the machine.
func newMachineWatcher(m *Machine) *MachineWatcher {
	w := &MachineWatcher{
		changeChan:    make(chan *Machine),
		commonWatcher: commonWatcher{st: m.st},
	}
	go func() {
		defer w.tomb.Done()
		defer close(w.changeChan)
		w.tomb.Kill(w.loop(m))
	}()
	return w
}

// Changes returns a channel that will receive the new
// *Machine when a change is detected. Note that multiple
// changes may be observed as a single event in the channel.
// The first event on the channel holds the initial state
// as returned by Machine.Info.
func (w *MachineWatcher) Changes() <-chan *Machine {
	return w.changeChan
}

func (w *MachineWatcher) Stop() error {
	w.tomb.Kill(nil)
	return w.tomb.Wait()
}

func (w *MachineWatcher) loop(m *Machine) (err error) {
	ch := make(chan watcher.Change)
	id := m.Id()
	st := m.st
	st.watcher.Watch(st.machines.Name, id, m.doc.TxnRevno, ch)
	defer st.watcher.Unwatch(st.machines.Name, id, ch)
	for {
		select {
		case <-st.watcher.Dead():
			return watcher.MustErr(st.watcher)
		case <-w.tomb.Dying():
			return tomb.ErrDying
		case <-ch:
		}
		if m, err = st.Machine(id); err != nil {
			return err
		}
		for {
			select {
			case <-st.watcher.Dead():
				return watcher.MustErr(st.watcher)
			case <-w.tomb.Dying():
				return tomb.ErrDying
			case <-ch:
				if err := m.Refresh(); err != nil {
					return err
				}
				continue
			case w.changeChan <- m:
			}
			break
		}
	}
	return nil
}

// WatchMachines returns a watcher for observing machines being
// added or removed.
func (s *State) WatchMachines() *MachinesWatcher {
	return newMachinesWatcher(s)
}

// newMachinesWatcher creates and starts a watcher to watch information
// about machines being added or deleted.
func newMachinesWatcher(st *State) *MachinesWatcher {
	w := &MachinesWatcher{
		changeChan:    make(chan *MachinesChange),
		knownMachines: make(map[int]*Machine),
		commonWatcher: commonWatcher{st: st},
	}
	go func() {
		defer w.tomb.Done()
		defer close(w.changeChan)
		w.tomb.Kill(w.loop())
	}()
	return w
}

// Changes returns a channel that will receive changes when machines are
// added or deleted. The Added field in the first event on the channel
// holds the initial state as returned by State.AllMachines.
func (w *MachinesWatcher) Changes() <-chan *MachinesChange {
	return w.changeChan
}

// Stop stops the watcher and returns any errors encountered while watching.
func (w *MachinesWatcher) Stop() error {
	w.tomb.Kill(nil)
	return w.tomb.Wait()
}

func (w *MachinesWatcher) mergeChange(changes *MachinesChange, ch watcher.Change) (err error) {
	id := ch.Id.(int)
	if m, ok := w.knownMachines[id]; ch.Revno == -1 && ok {
		m.doc.Life = Dead
		changes.Removed = append(changes.Removed, m)
		delete(w.knownMachines, id)
		return nil
	}
	doc := &machineDoc{}
	err = w.st.machines.FindId(id).One(doc)
	if err == mgo.ErrNotFound {
		return nil
	}
	if err != nil {
		return err
	}
	m := newMachine(w.st, doc)
	if _, ok := w.knownMachines[id]; !ok {
		changes.Added = append(changes.Added, m)
	}
	w.knownMachines[id] = m
	return nil
}

func (changes *MachinesChange) isEmpty() bool {
	return len(changes.Added)+len(changes.Removed) == 0
}

func (w *MachinesWatcher) getInitialEvent() (initial *MachinesChange, err error) {
	changes := &MachinesChange{}
	docs := []machineDoc{}
	err = w.st.machines.Find(nil).All(&docs)
	if err != nil {
		return nil, err
	}
	for _, doc := range docs {
		m := newMachine(w.st, &doc)
		w.knownMachines[doc.Id] = m
		changes.Added = append(changes.Added, m)
	}
	return changes, nil
}

func (w *MachinesWatcher) loop() (err error) {
	ch := make(chan watcher.Change)
	w.st.watcher.WatchCollection(w.st.machines.Name, ch)
	defer w.st.watcher.UnwatchCollection(w.st.machines.Name, ch)
	changes, err := w.getInitialEvent()
	if err != nil {
		return err
	}
	for {
		for changes != nil {
			select {
			case <-w.st.watcher.Dead():
				return watcher.MustErr(w.st.watcher)
			case <-w.tomb.Dying():
				return tomb.ErrDying
			case c := <-ch:
				err := w.mergeChange(changes, c)
				if err != nil {
					return err
				}
			case w.changeChan <- changes:
				changes = nil
			}
		}
		select {
		case <-w.st.watcher.Dead():
			return watcher.MustErr(w.st.watcher)
		case <-w.tomb.Dying():
			return tomb.ErrDying
		case c := <-ch:
			changes = &MachinesChange{}
			err := w.mergeChange(changes, c)
			if err != nil {
				return err
			}
			if changes.isEmpty() {
				changes = nil
			}
		}
	}
	return nil
}
<commit_msg>mstate: add services watcher<commit_after>package mstate

import (
	"labix.org/v2/mgo"
	"launchpad.net/juju-core/mstate/watcher"
	"launchpad.net/tomb"
)

// commonWatcher is part of all client watchers.
type commonWatcher struct {
	st   *State
	tomb tomb.Tomb
}

// MachineWatcher observes changes to the settings of a machine.
type MachineWatcher struct {
	commonWatcher
	changeChan chan *Machine
}

// MachinesWatcher notifies about machines being added or removed
// from the environment.
type MachinesWatcher struct {
	commonWatcher
	changeChan    chan *MachinesChange
	knownMachines map[int]*Machine
}

// MachinesChange contains information about
// machines that have been added or deleted.
type MachinesChange struct {
	Added   []*Machine
	Removed []*Machine
}

// ServicesWatcher observes the addition and removal of services.
type ServicesWatcher struct {
	commonWatcher
	changeChan    chan *ServicesChange
	knownServices map[string]*Service
}

// ServicesChange holds services that were added or removed
// from the environment.
type ServicesChange struct {
	Added   []*Service
	Removed []*Service
}

// newMachineWatcher creates and starts a watcher to watch information
// about the machine.
func newMachineWatcher(m *Machine) *MachineWatcher {
	w := &MachineWatcher{
		changeChan:    make(chan *Machine),
		commonWatcher: commonWatcher{st: m.st},
	}
	go func() {
		defer w.tomb.Done()
		defer close(w.changeChan)
		w.tomb.Kill(w.loop(m))
	}()
	return w
}

// Changes returns a channel that will receive the new
// *Machine when a change is detected. Note that multiple
// changes may be observed as a single event in the channel.
// The first event on the channel holds the initial state
// as returned by Machine.Info.
func (w *MachineWatcher) Changes() <-chan *Machine {
	return w.changeChan
}

func (w *MachineWatcher) Stop() error {
	w.tomb.Kill(nil)
	return w.tomb.Wait()
}

func (w *MachineWatcher) loop(m *Machine) (err error) {
	ch := make(chan watcher.Change)
	id := m.Id()
	st := m.st
	st.watcher.Watch(st.machines.Name, id, m.doc.TxnRevno, ch)
	defer st.watcher.Unwatch(st.machines.Name, id, ch)
	for {
		select {
		case <-st.watcher.Dead():
			return watcher.MustErr(st.watcher)
		case <-w.tomb.Dying():
			return tomb.ErrDying
		case <-ch:
		}
		if m, err = st.Machine(id); err != nil {
			return err
		}
		for {
			select {
			case <-st.watcher.Dead():
				return watcher.MustErr(st.watcher)
			case <-w.tomb.Dying():
				return tomb.ErrDying
			case <-ch:
				if err := m.Refresh(); err != nil {
					return err
				}
				continue
			case w.changeChan <- m:
			}
			break
		}
	}
	return nil
}

// WatchMachines returns a watcher for observing machines being
// added or removed.
func (s *State) WatchMachines() *MachinesWatcher {
	return newMachinesWatcher(s)
}

// newMachinesWatcher creates and starts a watcher to watch information
// about machines being added or deleted.
func newMachinesWatcher(st *State) *MachinesWatcher {
	w := &MachinesWatcher{
		changeChan:    make(chan *MachinesChange),
		knownMachines: make(map[int]*Machine),
		commonWatcher: commonWatcher{st: st},
	}
	go func() {
		defer w.tomb.Done()
		defer close(w.changeChan)
		w.tomb.Kill(w.loop())
	}()
	return w
}

// Changes returns a channel that will receive changes when machines are
// added or deleted. The Added field in the first event on the channel
// holds the initial state as returned by State.AllMachines.
func (w *MachinesWatcher) Changes() <-chan *MachinesChange {
	return w.changeChan
}

// Stop stops the watcher and returns any errors encountered while watching.
func (w *MachinesWatcher) Stop() error {
	w.tomb.Kill(nil)
	return w.tomb.Wait()
}

func (w *MachinesWatcher) mergeChange(changes *MachinesChange, ch watcher.Change) (err error) {
	id := ch.Id.(int)
	if m, ok := w.knownMachines[id]; ch.Revno == -1 && ok {
		m.doc.Life = Dead
		changes.Removed = append(changes.Removed, m)
		delete(w.knownMachines, id)
		return nil
	}
	doc := &machineDoc{}
	err = w.st.machines.FindId(id).One(doc)
	if err == mgo.ErrNotFound {
		return nil
	}
	if err != nil {
		return err
	}
	m := newMachine(w.st, doc)
	if _, ok := w.knownMachines[id]; !ok {
		changes.Added = append(changes.Added, m)
	}
	w.knownMachines[id] = m
	return nil
}

func (changes *MachinesChange) isEmpty() bool {
	return len(changes.Added)+len(changes.Removed) == 0
}

func (w *MachinesWatcher) getInitialEvent() (initial *MachinesChange, err error) {
	changes := &MachinesChange{}
	docs := []machineDoc{}
	err = w.st.machines.Find(nil).All(&docs)
	if err != nil {
		return nil, err
	}
	for _, doc := range docs {
		m := newMachine(w.st, &doc)
		w.knownMachines[doc.Id] = m
		changes.Added = append(changes.Added, m)
	}
	return changes, nil
}

func (w *MachinesWatcher) loop() (err error) {
	ch := make(chan watcher.Change)
	w.st.watcher.WatchCollection(w.st.machines.Name, ch)
	defer w.st.watcher.UnwatchCollection(w.st.machines.Name, ch)
	changes, err := w.getInitialEvent()
	if err != nil {
		return err
	}
	for {
		for changes != nil {
			select {
			case <-w.st.watcher.Dead():
				return watcher.MustErr(w.st.watcher)
			case <-w.tomb.Dying():
				return tomb.ErrDying
			case c := <-ch:
				err := w.mergeChange(changes, c)
				if err != nil {
					return err
				}
			case w.changeChan <- changes:
				changes = nil
			}
		}
		select {
		case <-w.st.watcher.Dead():
			return watcher.MustErr(w.st.watcher)
		case <-w.tomb.Dying():
			return tomb.ErrDying
		case c := <-ch:
			changes = &MachinesChange{}
			err := w.mergeChange(changes, c)
			if err != nil {
				return err
			}
			if changes.isEmpty() {
				changes = nil
			}
		}
	}
	return nil
}

// WatchServices returns a watcher for observing services being
// added or removed.
func (s *State) WatchServices() *ServicesWatcher {
	return newServicesWatcher(s)
}

// newServicesWatcher creates and starts a watcher to watch information
// about services being added or deleted.
func newServicesWatcher(st *State) *ServicesWatcher {
	w := &ServicesWatcher{
		changeChan:    make(chan *ServicesChange),
		knownServices: make(map[string]*Service),
		commonWatcher: commonWatcher{st: st},
	}
	go func() {
		defer w.tomb.Done()
		defer close(w.changeChan)
		w.tomb.Kill(w.loop())
	}()
	return w
}

// Changes returns a channel that will receive changes when services are
// added or deleted. The Added field in the first event on the channel
// holds the initial state as returned by State.AllServices.
func (w *ServicesWatcher) Changes() <-chan *ServicesChange {
	return w.changeChan
}

// Stop stops the watcher and returns any errors encountered while watching.
func (w *ServicesWatcher) Stop() error {
	w.tomb.Kill(nil)
	return w.tomb.Wait()
}

func (w *ServicesWatcher) mergeChange(changes *ServicesChange, ch watcher.Change) (err error) {
	name := ch.Id.(string)
	if svc, ok := w.knownServices[name]; ch.Revno == -1 && ok {
		svc.doc.Life = Dead
		changes.Removed = append(changes.Removed, svc)
		delete(w.knownServices, name)
		return nil
	}
	doc := serviceDoc{}
	err = w.st.services.FindId(name).One(&doc)
	if err == mgo.ErrNotFound {
		return nil
	}
	if err != nil {
		return err
	}
	svc := &Service{st: w.st, doc: doc}
	if _, ok := w.knownServices[name]; !ok {
		changes.Added = append(changes.Added, svc)
	}
	w.knownServices[name] = svc
	return nil
}

func (changes *ServicesChange) isEmpty() bool {
	return len(changes.Added)+len(changes.Removed) == 0
}

func (w *ServicesWatcher) getInitialEvent() (initial *ServicesChange, err error) {
	changes := &ServicesChange{}
	docs := []serviceDoc{}
	err = w.st.services.Find(nil).All(&docs)
	if err != nil {
		return nil, err
	}
	for _, doc := range docs {
		svc := &Service{st: w.st, doc: doc}
		w.knownServices[doc.Name] = svc
		changes.Added = append(changes.Added, svc)
	}
	return changes, nil
}

func (w *ServicesWatcher) loop() (err error) {
	ch := make(chan watcher.Change)
	w.st.watcher.WatchCollection(w.st.services.Name, ch)
	defer w.st.watcher.UnwatchCollection(w.st.services.Name, ch)
	changes, err := w.getInitialEvent()
	if err != nil {
		return err
	}
	for {
		for changes != nil {
			select {
			case <-w.st.watcher.Dead():
				return watcher.MustErr(w.st.watcher)
			case <-w.tomb.Dying():
				return tomb.ErrDying
			case c := <-ch:
				err := w.mergeChange(changes, c)
				if err != nil {
					return err
				}
			case w.changeChan <- changes:
				changes = nil
			}
		}
		select {
		case <-w.st.watcher.Dead():
			return watcher.MustErr(w.st.watcher)
		case <-w.tomb.Dying():
			return tomb.ErrDying
		case c := <-ch:
			changes = &ServicesChange{}
			err := w.mergeChange(changes, c)
			if err != nil {
				return err
			}
			if changes.isEmpty() {
				changes = nil
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/prometheus/client_golang/prometheus"
)

const (
	proto         = "tcp"
	retryInterval = 1
)

var (
	listeningAddress    = flag.String("listeningAddress", ":8080", "Address on which to expose Prometheus metrics.")
	muninAddress        = flag.String("muninAddress", "localhost:4949", "munin-node address.")
	muninScrapeInterval = flag.Int("muninScrapeInterval", 60, "Interval in seconds between scrapes.")
	globalConn          net.Conn
	hostname            string
	graphs              []string
	gaugePerMetric      map[string]*prometheus.GaugeVec
	counterPerMetric     map[string]*prometheus.CounterVec
	muninBanner         *regexp.Regexp
)

func init() {
	flag.Parse()
	var err error
	gaugePerMetric = map[string]*prometheus.GaugeVec{}
	counterPerMetric = map[string]*prometheus.CounterVec{}
	muninBanner = regexp.MustCompile(`# munin node at (.*)`)

	err = connect()
	if err != nil {
		log.Fatalf("Could not connect to %s: %s", *muninAddress, err)
	}
}

func serveStatus() {
	http.Handle("/metrics", prometheus.Handler())
	http.ListenAndServe(*listeningAddress, nil)
}

func connect() (err error) {
	log.Printf("Connecting...")
	globalConn, err = net.Dial(proto, *muninAddress)
	if err != nil {
		return
	}
	log.Printf("connected!")

	reader := bufio.NewReader(globalConn)
	head, err := reader.ReadString('\n')
	if err != nil {
		return
	}

	matches := muninBanner.FindStringSubmatch(head)
	if len(matches) != 2 { // expect: # munin node at <hostname>
		return fmt.Errorf("Unexpected line: %s", head)
	}
	hostname = matches[1]
	log.Printf("Found hostname: %s", hostname)
	return
}

func muninCommand(cmd string) (reader *bufio.Reader, err error) {
	reader = bufio.NewReader(globalConn)

	fmt.Fprintf(globalConn, cmd+"\n")

	_, err = reader.Peek(1)
	switch err {
	case io.EOF:
		log.Printf("not connected anymore, closing connection")
		globalConn.Close()
		for {
			err = connect()
			if err == nil {
				break
			}
			log.Printf("Couldn't reconnect: %s", err)
			time.Sleep(retryInterval * time.Second)
		}

		return muninCommand(cmd)
	case nil: //no error
		break
	default:
		log.Fatalf("Unexpected error: %s", err)
	}

	return
}

func muninList() (items []string, err error) {
	munin, err := muninCommand("list")
	if err != nil {
		log.Printf("couldn't get list")
		return
	}

	response, err := munin.ReadString('\n') // we are only interested in the first line
	if err != nil {
		log.Printf("couldn't read response")
		return
	}

	if response[0] == '#' { // # not expected here
		err = fmt.Errorf("Error getting items: %s", response)
		return
	}
	items = strings.Fields(strings.TrimRight(response, "\n"))
	return
}

func muninConfig(name string) (config map[string]map[string]string, graphConfig map[string]string, err error) {
	graphConfig = make(map[string]string)
	config = make(map[string]map[string]string)

	resp, err := muninCommand("config " + name)
	if err != nil {
		log.Printf("couldn't get config for %s", name)
		return
	}

	for {
		line, err := resp.ReadString('\n')
		if err == io.EOF {
			log.Fatalf("unexpected EOF, retrying")
			return muninConfig(name)
		}
		if err != nil {
			return nil, nil, err
		}
		if line == ".\n" { // munin end marker
			break
		}
		if line[0] == '#' { // here it's just a comment, so ignore it
			continue
		}
		parts := strings.Fields(line)
		if len(parts) < 2 {
			return nil, nil, fmt.Errorf("Line unexpected: %s", line)
		}
		key, value := parts[0], strings.TrimRight(strings.Join(parts[1:], " "), "\n")

		keyParts := strings.Split(key, ".")
		if len(keyParts) > 1 { // it's a metric config (metric.label etc)
			if _, ok := config[keyParts[0]]; !ok { //FIXME: is there no better way?
				config[keyParts[0]] = make(map[string]string)
			}
			config[keyParts[0]][keyParts[1]] = value
		} else {
			graphConfig[keyParts[0]] = value
		}
	}
	return
}

func registerMetrics() (err error) {
	items, err := muninList()
	if err != nil {
		return
	}

	for _, name := range items {
		graphs = append(graphs, name)
		configs, graphConfig, err := muninConfig(name)
		if err != nil {
			return err
		}

		for metric, config := range configs {
			metricName := strings.Replace(name + "_" + metric, "-","_",-1)
			desc := graphConfig["graph_title"] + ": " + config["label"]
			if config["info"] != "" {
				desc = desc + ", " + config["info"]
			}
			muninType := strings.ToLower(config["type"])
			// muninType can be empty and defaults to gauge
			if muninType == "counter" || muninType == "derive" {
	                        gv := prometheus.NewCounterVec(
        	                        prometheus.CounterOpts{
                	                        Name: metricName,
                        	                Help: desc,
                                	},
                                	[]string{"hostname"},
                        	)
				log.Printf("Registered counter %s: %s", metricName, desc)
                        	counterPerMetric[metricName] = gv
                        	prometheus.Register(gv)

			} else {
                        	gv := prometheus.NewGaugeVec(
                                	prometheus.GaugeOpts{
                                        	Name: metricName,
	                                        Help: desc,
                	                },
                        	        []string{"hostname"},
                        	)
				log.Printf("Registered gauge %s: %s", metricName, desc)
        	                gaugePerMetric[metricName] = gv
                	        prometheus.Register(gv)
			}
		}
	}
	return nil
}

func fetchMetrics() (err error) {
	for _, graph := range graphs {
		munin, err := muninCommand("fetch " + graph)
		if err != nil {
			return err
		}

		for {
			line, err := munin.ReadString('\n')
			line = strings.TrimRight(line, "\n")
			if err == io.EOF {
				log.Fatalf("unexpected EOF, retrying")
				return fetchMetrics()
			}
			if err != nil {
				return err
			}
			if len(line) == 1 && line[0] == '.' {
				log.Printf("End of list")
				break
			}

			parts := strings.Fields(line)
			if len(parts) != 2 {
				log.Printf("unexpected line: %s", line)
				continue
			}
			key, valueString := strings.Split(parts[0], ".")[0], parts[1]
			value, err := strconv.ParseFloat(valueString, 64)
			if err != nil {
				log.Printf("Couldn't parse value in line %s, malformed?", line)
				continue
			}
			name := strings.Replace(graph + "_" + key, "-","_",-1)
			log.Printf("%s: %f\n", name, value)
			_, isGauge := gaugePerMetric[name]
			if isGauge {
	                        gaugePerMetric[name].WithLabelValues(hostname).Set(value)
			} else {
				counterPerMetric[name].WithLabelValues(hostname).Set(value)
			}
		}
	}
	return
}

func main() {
	flag.Parse()
	err := registerMetrics()
	if err != nil {
		log.Fatalf("Could not register metrics: %s", err)
	}

	go serveStatus()

	func() {
		for {
			log.Printf("Scraping")
			err := fetchMetrics()
			if err != nil {
				log.Printf("Error occured when trying to fetch metrics: %s", err)
			}
			time.Sleep(time.Duration(*muninScrapeInterval) * time.Second)
		}
	}()
}
<commit_msg>Introduced a ConstLabel "type" to be easily able to differentiate between counter, gauge and drive. Added labels "graphname" and "muninlabel" as well, to simplify aggregation.<commit_after>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/prometheus/client_golang/prometheus"
)

const (
	proto         = "tcp"
	retryInterval = 1
)

var (
	listeningAddress    = flag.String("listeningAddress", ":8080", "Address on which to expose Prometheus metrics.")
	muninAddress        = flag.String("muninAddress", "localhost:4949", "munin-node address.")
	muninScrapeInterval = flag.Int("muninScrapeInterval", 60, "Interval in seconds between scrapes.")
	globalConn          net.Conn
	hostname            string
	graphs              []string
	gaugePerMetric      map[string]*prometheus.GaugeVec
	counterPerMetric     map[string]*prometheus.CounterVec
	muninBanner         *regexp.Regexp
)

func init() {
	flag.Parse()
	var err error
	gaugePerMetric = map[string]*prometheus.GaugeVec{}
	counterPerMetric = map[string]*prometheus.CounterVec{}
	muninBanner = regexp.MustCompile(`# munin node at (.*)`)

	err = connect()
	if err != nil {
		log.Fatalf("Could not connect to %s: %s", *muninAddress, err)
	}
}

func serveStatus() {
	http.Handle("/metrics", prometheus.Handler())
	http.ListenAndServe(*listeningAddress, nil)
}

func connect() (err error) {
	log.Printf("Connecting...")
	globalConn, err = net.Dial(proto, *muninAddress)
	if err != nil {
		return
	}
	log.Printf("connected!")

	reader := bufio.NewReader(globalConn)
	head, err := reader.ReadString('\n')
	if err != nil {
		return
	}

	matches := muninBanner.FindStringSubmatch(head)
	if len(matches) != 2 { // expect: # munin node at <hostname>
		return fmt.Errorf("Unexpected line: %s", head)
	}
	hostname = matches[1]
	log.Printf("Found hostname: %s", hostname)
	return
}

func muninCommand(cmd string) (reader *bufio.Reader, err error) {
	reader = bufio.NewReader(globalConn)

	fmt.Fprintf(globalConn, cmd+"\n")

	_, err = reader.Peek(1)
	switch err {
	case io.EOF:
		log.Printf("not connected anymore, closing connection")
		globalConn.Close()
		for {
			err = connect()
			if err == nil {
				break
			}
			log.Printf("Couldn't reconnect: %s", err)
			time.Sleep(retryInterval * time.Second)
		}

		return muninCommand(cmd)
	case nil: //no error
		break
	default:
		log.Fatalf("Unexpected error: %s", err)
	}

	return
}

func muninList() (items []string, err error) {
	munin, err := muninCommand("list")
	if err != nil {
		log.Printf("couldn't get list")
		return
	}

	response, err := munin.ReadString('\n') // we are only interested in the first line
	if err != nil {
		log.Printf("couldn't read response")
		return
	}

	if response[0] == '#' { // # not expected here
		err = fmt.Errorf("Error getting items: %s", response)
		return
	}
	items = strings.Fields(strings.TrimRight(response, "\n"))
	return
}

func muninConfig(name string) (config map[string]map[string]string, graphConfig map[string]string, err error) {
	graphConfig = make(map[string]string)
	config = make(map[string]map[string]string)

	resp, err := muninCommand("config " + name)
	if err != nil {
		log.Printf("couldn't get config for %s", name)
		return
	}

	for {
		line, err := resp.ReadString('\n')
		if err == io.EOF {
			log.Fatalf("unexpected EOF, retrying")
			return muninConfig(name)
		}
		if err != nil {
			return nil, nil, err
		}
		if line == ".\n" { // munin end marker
			break
		}
		if line[0] == '#' { // here it's just a comment, so ignore it
			continue
		}
		parts := strings.Fields(line)
		if len(parts) < 2 {
			return nil, nil, fmt.Errorf("Line unexpected: %s", line)
		}
		key, value := parts[0], strings.TrimRight(strings.Join(parts[1:], " "), "\n")

		keyParts := strings.Split(key, ".")
		if len(keyParts) > 1 { // it's a metric config (metric.label etc)
			if _, ok := config[keyParts[0]]; !ok { //FIXME: is there no better way?
				config[keyParts[0]] = make(map[string]string)
			}
			config[keyParts[0]][keyParts[1]] = value
		} else {
			graphConfig[keyParts[0]] = value
		}
	}
	return
}

func registerMetrics() (err error) {
	items, err := muninList()
	if err != nil {
		return
	}

	for _, name := range items {
		graphs = append(graphs, name)
		configs, graphConfig, err := muninConfig(name)
		if err != nil {
			return err
		}

		for metric, config := range configs {
			metricName := strings.Replace(name + "_" + metric, "-","_",-1)
			desc := graphConfig["graph_title"] + ": " + config["label"]
			if config["info"] != "" {
				desc = desc + ", " + config["info"]
			}
			muninType := strings.ToLower(config["type"])
			// muninType can be empty and defaults to gauge
			if muninType == "counter" || muninType == "derive" {
	                        gv := prometheus.NewCounterVec(
        	                        prometheus.CounterOpts{
                	                        Name: metricName,
                        	                Help: desc,
						ConstLabels: prometheus.Labels{"type":muninType},
                                	},
                                	[]string{"hostname","graphname","muninlabel"},
                        	)
				log.Printf("Registered counter %s: %s", metricName, desc)
                        	counterPerMetric[metricName] = gv
                        	prometheus.Register(gv)

			} else {
                        	gv := prometheus.NewGaugeVec(
                                	prometheus.GaugeOpts{
                                        	Name: metricName,
	                                        Help: desc,
						ConstLabels: prometheus.Labels{"type":"counter"},
                	                },
                        	        []string{"hostname","graphname","muninlabel"},
                        	)
				log.Printf("Registered gauge %s: %s", metricName, desc)
        	                gaugePerMetric[metricName] = gv
                	        prometheus.Register(gv)
			}
		}
	}
	return nil
}

func fetchMetrics() (err error) {
	for _, graph := range graphs {
		munin, err := muninCommand("fetch " + graph)
		if err != nil {
			return err
		}

		for {
			line, err := munin.ReadString('\n')
			line = strings.TrimRight(line, "\n")
			if err == io.EOF {
				log.Fatalf("unexpected EOF, retrying")
				return fetchMetrics()
			}
			if err != nil {
				return err
			}
			if len(line) == 1 && line[0] == '.' {
				log.Printf("End of list")
				break
			}

			parts := strings.Fields(line)
			if len(parts) != 2 {
				log.Printf("unexpected line: %s", line)
				continue
			}
			key, valueString := strings.Split(parts[0], ".")[0], parts[1]
			value, err := strconv.ParseFloat(valueString, 64)
			if err != nil {
				log.Printf("Couldn't parse value in line %s, malformed?", line)
				continue
			}
			name := strings.Replace(graph + "_" + key, "-","_",-1)
			log.Printf("%s: %f\n", name, value)
			_, isGauge := gaugePerMetric[name]
			if isGauge {
	                        gaugePerMetric[name].WithLabelValues(hostname, graph, key).Set(value)
			} else {
				counterPerMetric[name].WithLabelValues(hostname, graph, key).Set(value)
			}
		}
	}
	return
}

func main() {
	flag.Parse()
	err := registerMetrics()
	if err != nil {
		log.Fatalf("Could not register metrics: %s", err)
	}

	go serveStatus()

	func() {
		for {
			log.Printf("Scraping")
			err := fetchMetrics()
			if err != nil {
				log.Printf("Error occured when trying to fetch metrics: %s", err)
			}
			time.Sleep(time.Duration(*muninScrapeInterval) * time.Second)
		}
	}()
}
<|endoftext|>


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

package api

import (
    "net/http"
    "encoding/json"
    "io"
    "bytes"
    "log"
)

type catalogHandler interface {
    authenticate(w http.ResponseWriter, r *http.Request) bool
}

func (c *CatalogV2) authenticate(w http.ResponseWriter, r *http.Request) bool {
    if r.URL.User != nil && r.URL.User.Username() != "" {
        log.Print("Basic Authentication: ", r.URL.User.Username())
        if password, _ := r.URL.User.Password(); password != "" {
            log.Print(" : ", password)
        }
        log.Println()
    }
    return true
}

var catalog *CatalogV2 = &CatalogV2 {
    Services : []ServiceV2 {
        ServiceV2 {
            Id: "9c372bbc-1e7b-472b-bcb6-eeda5b21eb35",
            Name: "redis-cluster-managed-by-kubernetes",
            Description: `The Redis is a high reliable and scalable cluster deployed upon` +
                ` Kubernetes v1, it failovers in master/slave, and load balancing with multiple` +
                ` sentinel nodes`,
            Bindable: false,
            Tags: []string{"redis", "cluster", "k-v", "database"},
            Plans: []ServicePlanV2 {
                ServicePlanV2 {
                    Id: "8cfbbaf5-efdb-41c1-89ab-f797185f7818",
                    Name: "demo",
                    Description: "this is a redis cluster demo",
                    Free: true,
                },
            },
        },
    },
}

/*
    curl -H "X-Broker-API-Version: 2.6" http://username:password@broker-url/v2/catalog
*/

func HandleCatalog(w http.ResponseWriter, r *http.Request) {
    if r.Method != "GET" {
        http.Error(w, "Api only support GET method.", http.StatusMethodNotAllowed)
        return
    }
    if v := r.Header["X-Broker-API-Version"]; len(v) > 0 && v[0] != "2.6" {
        http.Error(w, "Unmatched API version.", http.StatusPreconditionFailed)
        return
    }

    if !catalog.authenticate(w, r) {
        http.Error(w, "Not authorized", http.StatusUnauthorized)
        return
    }

    //enc := json.NewEncoder(w)
    //enc.Encode(catalog)
    js, err := json.Marshal(catalog)
    if err != nil {
        http.Error(w, err.Error(), http.StatusInternalServerError)
        return
    }
    w.Header().Set("Content-Type", "application/json")
    //w.Write(js)
    io.Copy(w, bytes.NewBuffer(js))
}<commit_msg>change redis catalog<commit_after>/*
Copyright 2015 All rights reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
    http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package api

import (
    "net/http"
    "encoding/json"
    "io"
    "bytes"
    "log"
)

type catalogHandler interface {
    authenticate(w http.ResponseWriter, r *http.Request) bool
}

func (c *CatalogV2) authenticate(w http.ResponseWriter, r *http.Request) bool {
    if r.URL.User != nil && r.URL.User.Username() != "" {
        log.Print("Basic Authentication: ", r.URL.User.Username())
        if password, _ := r.URL.User.Password(); password != "" {
            log.Print(" : ", password)
        }
        log.Println()
    }
    return true
}

var catalog *CatalogV2 = &CatalogV2 {
    Services : []ServiceV2 {
        ServiceV2 {
            Id: "9c372bbc-1e7b-472b-bcb6-eeda5b21eb35",
            Name: "redis-cluster-managed-by-kubernetes",
            Description: `The Redis is a high reliable and scalable cluster deployed upon` +
                ` Kubernetes v1, it failovers in master/slave, and load balancing with multiple` +
                ` sentinel nodes`,
            Bindable: true,
            Tags: []string{"redis", "cluster", "k-v", "database"},
            Plans: []ServicePlanV2 {
                ServicePlanV2 {
                    Id: "8cfbbaf5-efdb-41c1-89ab-f797185f7818",
                    Name: "demo",
                    Description: "this is a redis cluster demo",
                    Free: true,
                },
            },
        },
    },
}

/*
    curl -H "X-Broker-API-Version: 2.6" http://username:password@broker-url/v2/catalog
*/

func HandleCatalog(w http.ResponseWriter, r *http.Request) {
    if r.Method != "GET" {
        http.Error(w, "Api only support GET method.", http.StatusMethodNotAllowed)
        return
    }
    if v := r.Header["X-Broker-API-Version"]; len(v) > 0 && v[0] != "2.6" {
        http.Error(w, "Unmatched API version.", http.StatusPreconditionFailed)
        return
    }

    if !catalog.authenticate(w, r) {
        http.Error(w, "Not authorized", http.StatusUnauthorized)
        return
    }

    //enc := json.NewEncoder(w)
    //enc.Encode(catalog)
    js, err := json.Marshal(catalog)
    if err != nil {
        http.Error(w, err.Error(), http.StatusInternalServerError)
        return
    }
    w.Header().Set("Content-Type", "application/json")
    //w.Write(js)
    io.Copy(w, bytes.NewBuffer(js))
}<|endoftext|>


<commit_before>package echo

import (
	"testing"

	"github.com/labstack/echo"
	"github.com/open-gtd/server/api"
	"github.com/stretchr/testify/mock"
)

type register struct {
	mock.Mock
}

type context struct {
	mock.Mock
}

func (r register) Add(method, path string, handler echo.HandlerFunc, middleware ...echo.MiddlewareFunc) *echo.Route {
	r.Called(method, path, handler)
	return nil
}

func TestRegisterer_GET_ShouldCallRegisterAddWithGETMethodAndParameters(t *testing.T) {
	path := "/xcx"

	groups := map[string]Router{}
	r := prepareRouterMock("GET", path)

	sut := NewEchoRegisterer(r, groups)
	sut.GET("", path, handler)

	r.AssertExpectations(t)
}

func TestRegisterer_POST_ShouldCallRegisterAddWithPOSTMethodAndParameters(t *testing.T) {
	path := "/xzx"

	groups := map[string]Router{}
	r := prepareRouterMock("POST", path)

	sut := NewEchoRegisterer(r, groups)
	sut.POST("", path, handler)

	r.AssertExpectations(t)
}

func TestRegisterer_PATCH_ShouldCallRegisterAddWithPATCHMethodAndParameters(t *testing.T) {
	path := "/xxz"

	groups := map[string]Router{}
	r := prepareRouterMock("PATCH", path)

	sut := NewEchoRegisterer(r, groups)
	sut.PATCH("", path, handler)

	r.AssertExpectations(t)
}

func TestRegisterer_PUT_ShouldCallRegisterAddWithPUTMethodAndParameters(t *testing.T) {
	path := "/zxx"

	groups := map[string]Router{}
	r := prepareRouterMock("PUT", path)

	sut := NewEchoRegisterer(r, groups)
	sut.PUT("", path, handler)

	r.AssertExpectations(t)
}

func TestRegisterer_DELETE_ShouldCallRegisterAddWithDELETEMethodAndParameters(t *testing.T) {
	prefix := ""
	path := "/xyy"

	groups := map[string]Router{}
	r := prepareRouterMock("DELETE", path)

	sut := NewEchoRegisterer(r, groups)
	sut.DELETE(prefix, path, handler)

	r.AssertExpectations(t)
}

func TestRegisterer_GET_ShouldCallRegisterAddWithGETMethodAndParametersAndPrefix(t *testing.T) {
	path := "/xcx"
	prefix := "/prefix"

	groups := map[string]Router{}
	r := prepareRouterMock("GET", prefix+path)

	sut := NewEchoRegisterer(r, groups)
	sut.GET(prefix, path, handler)

	r.AssertExpectations(t)
}

func TestRegisterer_POST_ShouldCallRegisterAddWithPOSTMethodAndParametersAndPrefix(t *testing.T) {
	path := "/xzx"
	prefix := "/prefix"

	groups := map[string]Router{}
	r := prepareRouterMock("POST", prefix+path)

	sut := NewEchoRegisterer(r, groups)
	sut.POST(prefix, path, handler)

	r.AssertExpectations(t)
}

func TestRegisterer_PATCH_ShouldCallRegisterAddWithPATCHMethodAndParametersAndPrefix(t *testing.T) {
	path := "/xxz"
	prefix := "/prefix"

	groups := map[string]Router{}
	r := prepareRouterMock("PATCH", prefix+path)

	sut := NewEchoRegisterer(r, groups)
	sut.PATCH(prefix, path, handler)

	r.AssertExpectations(t)
}

func TestRegisterer_PUT_ShouldCallRegisterAddWithPUTMethodAndParametersAndPrefix(t *testing.T) {
	path := "/zxx"
	prefix := "/prefix"

	groups := map[string]Router{}
	r := prepareRouterMock("PUT", prefix+path)

	sut := NewEchoRegisterer(r, groups)
	sut.PUT(prefix, path, handler)

	r.AssertExpectations(t)
}

func TestRegisterer_DELETE_ShouldCallRegisterAddWithDELETEMethodAndParametersAndPrefix(t *testing.T) {
	prefix := "/prefix"
	path := "/xyy"

	groups := map[string]Router{}
	r := prepareRouterMock("DELETE", prefix+path)

	sut := NewEchoRegisterer(r, groups)
	sut.DELETE(prefix, path, handler)

	r.AssertExpectations(t)
}

func TestRegisterer_GET_ShouldCallRegisterAddOnGroupWithGETMethodAndParameters(t *testing.T) {
	path := "/xxz"
	prefix := "/prefix"

	r := register{}
	groups := map[string]Router{}
	groups["/prefix"] = prepareRouterMock("GET", path)

	sut := NewEchoRegisterer(r, groups)
	sut.GET(prefix, path, handler)

	r.AssertExpectations(t)
}

func TestRegisterer_POST_ShouldCallRegisterAddOnGroupWithPOSTMethodAndParameters(t *testing.T) {
	path := "/xxz"
	prefix := "/prefix"

	r := register{}
	groups := map[string]Router{}
	groups["/prefix"] = prepareRouterMock("POST", path)

	sut := NewEchoRegisterer(r, groups)
	sut.POST(prefix, path, handler)

	r.AssertExpectations(t)
}

func TestRegisterer_PATCH_ShouldCallRegisterAddOnGroupWithPATCHMethodAndParameters(t *testing.T) {
	path := "/xxz"
	prefix := "/prefix"

	r := register{}
	groups := map[string]Router{}
	groups["/prefix"] = prepareRouterMock("PATCH", path)

	sut := NewEchoRegisterer(r, groups)
	sut.PATCH(prefix, path, handler)

	r.AssertExpectations(t)
}

func TestRegisterer_PUT_ShouldCallRegisterAddOnGroupWithPUTMethodAndParameters(t *testing.T) {
	path := "/xxz"
	prefix := "/prefix"

	r := register{}
	groups := map[string]Router{}
	groups["/prefix"] = prepareRouterMock("PUT", path)

	sut := NewEchoRegisterer(r, groups)
	sut.PUT(prefix, path, handler)

	r.AssertExpectations(t)
}

func TestRegisterer_DELETE_ShouldCallRegisterAddOnGroupWithDELETEMethodAndParameters(t *testing.T) {
	path := "/xxz"
	prefix := "/prefix"

	r := register{}
	groups := map[string]Router{}
	groups["/prefix"] = prepareRouterMock("DELETE", path)

	sut := NewEchoRegisterer(r, groups)
	sut.DELETE(prefix, path, handler)

	r.AssertExpectations(t)
}

func prepareRouterMock(method string, path string) register {
	r := register{}

	r.On("Add", method, path, mock.AnythingOfType("echo.HandlerFunc"))
	return r
}

var handler = func(api.Request, api.Response) error {
	return nil
}
<commit_msg>echo register mock fixed<commit_after>package echo

import (
	"testing"

	"github.com/labstack/echo"
	"github.com/open-gtd/server/api"
	"github.com/stretchr/testify/mock"
	"fmt"
)

type register struct {
	mock.Mock
}

type context struct {
	mock.Mock
}

func (r register) Add(method, path string, handler echo.HandlerFunc, middleware ...echo.MiddlewareFunc) *echo.Route {
	args:=r.Called(method, path, handler)
	return route(args.Get(0))
}

func TestRegisterer_GET_ShouldCallRegisterAddWithGETMethodAndParameters(t *testing.T) {
	path := "/xcx"

	groups := map[string]Router{}
	r := prepareRouterMock("GET", path)

	sut := NewEchoRegisterer(r, groups)
	sut.GET("", path, handler)

	r.AssertExpectations(t)
}

func TestRegisterer_POST_ShouldCallRegisterAddWithPOSTMethodAndParameters(t *testing.T) {
	path := "/xzx"

	groups := map[string]Router{}
	r := prepareRouterMock("POST", path)

	sut := NewEchoRegisterer(r, groups)
	sut.POST("", path, handler)

	r.AssertExpectations(t)
}

func TestRegisterer_PATCH_ShouldCallRegisterAddWithPATCHMethodAndParameters(t *testing.T) {
	path := "/xxz"

	groups := map[string]Router{}
	r := prepareRouterMock("PATCH", path)

	sut := NewEchoRegisterer(r, groups)
	sut.PATCH("", path, handler)

	r.AssertExpectations(t)
}

func TestRegisterer_PUT_ShouldCallRegisterAddWithPUTMethodAndParameters(t *testing.T) {
	path := "/zxx"

	groups := map[string]Router{}
	r := prepareRouterMock("PUT", path)

	sut := NewEchoRegisterer(r, groups)
	sut.PUT("", path, handler)

	r.AssertExpectations(t)
}

func TestRegisterer_DELETE_ShouldCallRegisterAddWithDELETEMethodAndParameters(t *testing.T) {
	prefix := ""
	path := "/xyy"

	groups := map[string]Router{}
	r := prepareRouterMock("DELETE", path)

	sut := NewEchoRegisterer(r, groups)
	sut.DELETE(prefix, path, handler)

	r.AssertExpectations(t)
}

func TestRegisterer_GET_ShouldCallRegisterAddWithGETMethodAndParametersAndPrefix(t *testing.T) {
	path := "/xcx"
	prefix := "/prefix"

	groups := map[string]Router{}
	r := prepareRouterMock("GET", prefix+path)

	sut := NewEchoRegisterer(r, groups)
	sut.GET(prefix, path, handler)

	r.AssertExpectations(t)
}

func TestRegisterer_POST_ShouldCallRegisterAddWithPOSTMethodAndParametersAndPrefix(t *testing.T) {
	path := "/xzx"
	prefix := "/prefix"

	groups := map[string]Router{}
	r := prepareRouterMock("POST", prefix+path)

	sut := NewEchoRegisterer(r, groups)
	sut.POST(prefix, path, handler)

	r.AssertExpectations(t)
}

func TestRegisterer_PATCH_ShouldCallRegisterAddWithPATCHMethodAndParametersAndPrefix(t *testing.T) {
	path := "/xxz"
	prefix := "/prefix"

	groups := map[string]Router{}
	r := prepareRouterMock("PATCH", prefix+path)

	sut := NewEchoRegisterer(r, groups)
	sut.PATCH(prefix, path, handler)

	r.AssertExpectations(t)
}

func TestRegisterer_PUT_ShouldCallRegisterAddWithPUTMethodAndParametersAndPrefix(t *testing.T) {
	path := "/zxx"
	prefix := "/prefix"

	groups := map[string]Router{}
	r := prepareRouterMock("PUT", prefix+path)

	sut := NewEchoRegisterer(r, groups)
	sut.PUT(prefix, path, handler)

	r.AssertExpectations(t)
}

func TestRegisterer_DELETE_ShouldCallRegisterAddWithDELETEMethodAndParametersAndPrefix(t *testing.T) {
	prefix := "/prefix"
	path := "/xyy"

	groups := map[string]Router{}
	r := prepareRouterMock("DELETE", prefix+path)

	sut := NewEchoRegisterer(r, groups)
	sut.DELETE(prefix, path, handler)

	r.AssertExpectations(t)
}

func TestRegisterer_GET_ShouldCallRegisterAddOnGroupWithGETMethodAndParameters(t *testing.T) {
	path := "/xxz"
	prefix := "/prefix"

	r := register{}
	groups := map[string]Router{}
	groups["/prefix"] = prepareRouterMock("GET", path)

	sut := NewEchoRegisterer(r, groups)
	sut.GET(prefix, path, handler)

	r.AssertExpectations(t)
}

func TestRegisterer_POST_ShouldCallRegisterAddOnGroupWithPOSTMethodAndParameters(t *testing.T) {
	path := "/xxz"
	prefix := "/prefix"

	r := register{}
	groups := map[string]Router{}
	groups["/prefix"] = prepareRouterMock("POST", path)

	sut := NewEchoRegisterer(r, groups)
	sut.POST(prefix, path, handler)

	r.AssertExpectations(t)
}

func TestRegisterer_PATCH_ShouldCallRegisterAddOnGroupWithPATCHMethodAndParameters(t *testing.T) {
	path := "/xxz"
	prefix := "/prefix"

	r := register{}
	groups := map[string]Router{}
	groups["/prefix"] = prepareRouterMock("PATCH", path)

	sut := NewEchoRegisterer(r, groups)
	sut.PATCH(prefix, path, handler)

	r.AssertExpectations(t)
}

func TestRegisterer_PUT_ShouldCallRegisterAddOnGroupWithPUTMethodAndParameters(t *testing.T) {
	path := "/xxz"
	prefix := "/prefix"

	r := register{}
	groups := map[string]Router{}
	groups["/prefix"] = prepareRouterMock("PUT", path)

	sut := NewEchoRegisterer(r, groups)
	sut.PUT(prefix, path, handler)

	r.AssertExpectations(t)
}

func TestRegisterer_DELETE_ShouldCallRegisterAddOnGroupWithDELETEMethodAndParameters(t *testing.T) {
	path := "/xxz"
	prefix := "/prefix"

	r := register{}
	groups := map[string]Router{}
	groups["/prefix"] = prepareRouterMock("DELETE", path)

	sut := NewEchoRegisterer(r, groups)
	sut.DELETE(prefix, path, handler)

	r.AssertExpectations(t)
}

func prepareRouterMock(method string, path string) register {
	r := register{}

	r.On("Add", method, path, mock.AnythingOfType("echo.HandlerFunc")).Return(nil)
	return r
}

var handler = func(api.Request, api.Response) error {
	return nil
}

func route(obj interface{}) *echo.Route {
	var r *echo.Route
	var ok bool
	if obj == nil {
		return nil
	}
	if r, ok = obj.(*echo.Route); !ok {
		panic(fmt.Sprintf("assert: arguments: Controller failed because object wasn't correct type: %v", obj))
	}
	return r
}
<|endoftext|>


<commit_before>package handler

import (
	"io"

	"net"
	"net/http"
	"net/http/httputil"
	"net/url"

	"github.com/pivotal-golang/lager"
)

func (handler *Handler) HijackBuild(w http.ResponseWriter, r *http.Request) {
	guid := r.FormValue(":guid")

	handler.buildsMutex.RLock()
	build, found := handler.builds[guid]
	handler.buildsMutex.RUnlock()

	if !found {
		w.WriteHeader(http.StatusNotFound)
		return
	}

	log := handler.logger.Session("hijack", lager.Data{
		"build": build,
	})

	log.Info("hijacking")

	hijackURL, err := url.Parse(build.HijackURL)
	if err != nil {
		log.Error("failed-to-parse-url", err)
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	conn, err := net.Dial("tcp", hijackURL.Host)
	if err != nil {
		log.Error("failed-to-dial-turbine", err)
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	client := httputil.NewClientConn(conn, nil)

	req, err := http.NewRequest(r.Method, build.HijackURL, r.Body)
	if err != nil {
		log.Error("failed-to-create-request", err)
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	resp, err := client.Do(req)
	if err != nil {
		log.Error("failed-to-hijack", err)
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	w.WriteHeader(resp.StatusCode)

	sconn, sbr, err := w.(http.Hijacker).Hijack()
	if err != nil {
		log.Error("failed-to-hijack", err)
		return
	}

	cconn, cbr := client.Hijack()

	go io.Copy(cconn, sbr)

	io.Copy(sconn, cbr)
}
<commit_msg>close properly in hijack handler<commit_after>package handler

import (
	"io"

	"net"
	"net/http"
	"net/http/httputil"
	"net/url"

	"github.com/pivotal-golang/lager"
)

func (handler *Handler) HijackBuild(w http.ResponseWriter, r *http.Request) {
	guid := r.FormValue(":guid")

	handler.buildsMutex.RLock()
	build, found := handler.builds[guid]
	handler.buildsMutex.RUnlock()

	if !found {
		w.WriteHeader(http.StatusNotFound)
		return
	}

	log := handler.logger.Session("hijack", lager.Data{
		"build": build,
	})

	log.Info("hijacking")

	hijackURL, err := url.Parse(build.HijackURL)
	if err != nil {
		log.Error("failed-to-parse-url", err)
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	conn, err := net.Dial("tcp", hijackURL.Host)
	if err != nil {
		log.Error("failed-to-dial-turbine", err)
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	req, err := http.NewRequest(r.Method, build.HijackURL, r.Body)
	if err != nil {
		log.Error("failed-to-create-request", err)
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	client := httputil.NewClientConn(conn, nil)

	resp, err := client.Do(req)
	if err != nil {
		log.Error("failed-to-hijack", err)
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	if resp.StatusCode != http.StatusOK {
		resp.Write(w)
		return
	}

	w.WriteHeader(http.StatusOK)

	sconn, sbr, err := w.(http.Hijacker).Hijack()
	if err != nil {
		log.Error("failed-to-hijack", err)
		return
	}

	cconn, cbr := client.Hijack()

	defer cconn.Close()
	defer sconn.Close()

	go io.Copy(cconn, sbr)

	io.Copy(sconn, cbr)
}
<|endoftext|>