<commit_before>package boomer

import (
	"time"
)

type requestStats struct {
	entries              map[string]*statsEntry
	errors               map[string]*statsError
	numRequests          int64
	numFailures          int64
	maxRequests          int64
	lastRequestTimestamp int64
	startTime            int64
}

func (s *requestStats) get(name string, method string) (entry *statsEntry) {
	entry, ok := s.entries[name+method]
	if !ok {
		newEntry := &statsEntry{
			stats:         s,
			name:          name,
			method:        method,
			numReqsPerSec: make(map[int64]int64),
			responseTimes: make(map[float64]int64),
		}
		newEntry.reset()
		s.entries[name+method] = newEntry
		return newEntry
	} else {
		return entry
	}

}

func (s *requestStats) clearAll() {
	s.numRequests = 0
	s.numFailures = 0
	s.entries = make(map[string]*statsEntry)
	s.errors = make(map[string]*statsError)
	s.maxRequests = 0
	s.lastRequestTimestamp = 0
	s.startTime = 0
}

type statsEntry struct {
	stats                *requestStats
	name                 string
	method               string
	numRequests          int64
	numFailures          int64
	totalResponseTime    float64
	minResponseTime      float64
	maxResponseTime      float64
	numReqsPerSec        map[int64]int64
	responseTimes        map[float64]int64
	totalContentLength   int64
	startTime            int64
	lastRequestTimestamp int64
}

func (s *statsEntry) reset() {
	s.startTime = int64(time.Now().Unix())
	s.numRequests = 0
	s.numFailures = 0
	s.totalResponseTime = 0
	s.responseTimes = make(map[float64]int64)
	s.minResponseTime = 0
	s.maxResponseTime = 0
	s.lastRequestTimestamp = int64(time.Now().Unix())
	s.numReqsPerSec = make(map[int64]int64)
	s.totalContentLength = 0
}

func (s *statsEntry) log(responseTime float64, contentLength int64) {

	s.numRequests += 1

	s.logTimeOfRequest()
	s.logResponseTime(responseTime)

	s.totalContentLength += contentLength

}

func (s *statsEntry) logTimeOfRequest() {

	now := int64(time.Now().Unix())

	_, ok := s.numReqsPerSec[now]
	if !ok {
		s.numReqsPerSec[now] = 0
	} else {
		s.numReqsPerSec[now] += 1
	}

	s.lastRequestTimestamp = now

}

func (s *statsEntry) logResponseTime(responseTime float64) {
	s.totalResponseTime += responseTime

	if s.minResponseTime == 0 {
		s.minResponseTime = responseTime
	}

	if responseTime < s.minResponseTime {
		s.minResponseTime = responseTime
	}

	if responseTime > s.maxResponseTime {
		s.maxResponseTime = responseTime
	}

	roundedResponseTime := float64(0)

	if responseTime < 100 {
		roundedResponseTime = responseTime
	} else if responseTime < 1000 {
		roundedResponseTime = float64(Round(responseTime, .5, -1))
	} else if responseTime < 10000 {
		roundedResponseTime = float64(Round(responseTime, .5, -2))
	} else {
		roundedResponseTime = float64(Round(responseTime, .5, -3))
	}

	_, ok := s.responseTimes[roundedResponseTime]
	if !ok {
		s.responseTimes[roundedResponseTime] = 0
	} else {
		s.responseTimes[roundedResponseTime] += 1
	}

}

func (s *statsEntry) logError(err string) {
	s.numFailures += 1
	key := MD5(s.method, s.name, err)
	entry, ok := s.stats.errors[key]
	if !ok {
		entry = &statsError{
			name:   s.name,
			method: s.method,
			error:  err,
		}
		s.stats.errors[key] = entry
	}
	entry.occured()
}

func (s *statsEntry) serialize() map[string]interface{} {
	result := make(map[string]interface{})
	result["name"] = s.name
	result["method"] = s.method
	result["last_request_timestamp"] = s.lastRequestTimestamp
	result["start_time"] = s.startTime
	result["num_requests"] = s.numRequests
	result["num_failures"] = s.numFailures
	result["total_response_time"] = s.totalResponseTime
	result["max_response_time"] = s.maxResponseTime
	result["min_response_time"] = s.minResponseTime
	result["total_content_length"] = s.totalContentLength
	result["response_times"] = s.responseTimes
	result["num_reqs_per_sec"] = s.numReqsPerSec
	return result
}

func (s *statsEntry) getStrippedReport() map[string]interface{} {
	report := s.serialize()
	s.reset()
	return report
}

type statsError struct {
	name       string
	method     string
	error      string
	occurences int64
}

func (err *statsError) occured() {
	err.occurences += 1
}

func (err *statsError) toMap() map[string]interface{} {
	m := make(map[string]interface{})
	m["method"] = err.method
	m["name"] = err.name
	m["error"] = err.error
	m["occurences"] = err.occurences
	return m
}

func collectReportData() map[string]interface{} {
	data := make(map[string]interface{})
	entries := make([]interface{}, 0, len(stats.entries))
	for _, v := range stats.entries {
		if !(v.numRequests == 0 && v.numFailures == 0) {
			entries = append(entries, v.getStrippedReport())
		}
	}

	errors := make(map[string]map[string]interface{})
	for k, v := range stats.errors {
		errors[k] = v.toMap()
	}

	data["stats"] = entries
	data["errors"] = errors
	stats.entries = make(map[string]*statsEntry)
	stats.errors = make(map[string]*statsError)

	return data
}

type requestSuccess struct {
	requestType    string
	name           string
	responseTime   float64
	responseLength int64
}

type requestFailure struct {
	requestType  string
	name         string
	responseTime float64
	error        string
}

var stats = new(requestStats)
var requestSuccessChannel = make(chan *requestSuccess, 100)
var requestFailureChannel = make(chan *requestFailure, 100)
var clearStatsChannel = make(chan bool)
var messageToServerChannel = make(chan map[string]interface{}, 10)

func init() {
	stats.entries = make(map[string]*statsEntry)
	stats.errors = make(map[string]*statsError)
	go func() {
		var ticker = time.NewTicker(SLAVE_REPORT_INTERVAL)
		for {
			select {
			case m := <-requestSuccessChannel:
				entry := stats.get(m.name, m.requestType)
				entry.log(m.responseTime, m.responseLength)
			case n := <-requestFailureChannel:
				stats.get(n.name, n.requestType).logError(n.error)
			case <-clearStatsChannel:
				stats.clearAll()
			case <-ticker.C:
				data := collectReportData()
				// send data to channel, no network IO in this goroutine
				messageToServerChannel <- data
			}
		}
	}()
}
<commit_msg>first response time and num_reqs should be 1<commit_after>package boomer

import (
	"time"
)

type requestStats struct {
	entries              map[string]*statsEntry
	errors               map[string]*statsError
	numRequests          int64
	numFailures          int64
	maxRequests          int64
	lastRequestTimestamp int64
	startTime            int64
}

func (s *requestStats) get(name string, method string) (entry *statsEntry) {
	entry, ok := s.entries[name+method]
	if !ok {
		newEntry := &statsEntry{
			stats:         s,
			name:          name,
			method:        method,
			numReqsPerSec: make(map[int64]int64),
			responseTimes: make(map[float64]int64),
		}
		newEntry.reset()
		s.entries[name+method] = newEntry
		return newEntry
	} else {
		return entry
	}

}

func (s *requestStats) clearAll() {
	s.numRequests = 0
	s.numFailures = 0
	s.entries = make(map[string]*statsEntry)
	s.errors = make(map[string]*statsError)
	s.maxRequests = 0
	s.lastRequestTimestamp = 0
	s.startTime = 0
}

type statsEntry struct {
	stats                *requestStats
	name                 string
	method               string
	numRequests          int64
	numFailures          int64
	totalResponseTime    float64
	minResponseTime      float64
	maxResponseTime      float64
	numReqsPerSec        map[int64]int64
	responseTimes        map[float64]int64
	totalContentLength   int64
	startTime            int64
	lastRequestTimestamp int64
}

func (s *statsEntry) reset() {
	s.startTime = int64(time.Now().Unix())
	s.numRequests = 0
	s.numFailures = 0
	s.totalResponseTime = 0
	s.responseTimes = make(map[float64]int64)
	s.minResponseTime = 0
	s.maxResponseTime = 0
	s.lastRequestTimestamp = int64(time.Now().Unix())
	s.numReqsPerSec = make(map[int64]int64)
	s.totalContentLength = 0
}

func (s *statsEntry) log(responseTime float64, contentLength int64) {

	s.numRequests += 1

	s.logTimeOfRequest()
	s.logResponseTime(responseTime)

	s.totalContentLength += contentLength

}

func (s *statsEntry) logTimeOfRequest() {

	now := int64(time.Now().Unix())

	_, ok := s.numReqsPerSec[now]
	if !ok {
		s.numReqsPerSec[now] = 1
	} else {
		s.numReqsPerSec[now] += 1
	}

	s.lastRequestTimestamp = now

}

func (s *statsEntry) logResponseTime(responseTime float64) {
	s.totalResponseTime += responseTime

	if s.minResponseTime == 0 {
		s.minResponseTime = responseTime
	}

	if responseTime < s.minResponseTime {
		s.minResponseTime = responseTime
	}

	if responseTime > s.maxResponseTime {
		s.maxResponseTime = responseTime
	}

	roundedResponseTime := float64(0)

	if responseTime < 100 {
		roundedResponseTime = responseTime
	} else if responseTime < 1000 {
		roundedResponseTime = float64(Round(responseTime, .5, -1))
	} else if responseTime < 10000 {
		roundedResponseTime = float64(Round(responseTime, .5, -2))
	} else {
		roundedResponseTime = float64(Round(responseTime, .5, -3))
	}

	_, ok := s.responseTimes[roundedResponseTime]
	if !ok {
		s.responseTimes[roundedResponseTime] = 1
	} else {
		s.responseTimes[roundedResponseTime] += 1
	}

}

func (s *statsEntry) logError(err string) {
	s.numFailures += 1
	key := MD5(s.method, s.name, err)
	entry, ok := s.stats.errors[key]
	if !ok {
		entry = &statsError{
			name:   s.name,
			method: s.method,
			error:  err,
		}
		s.stats.errors[key] = entry
	}
	entry.occured()
}

func (s *statsEntry) serialize() map[string]interface{} {
	result := make(map[string]interface{})
	result["name"] = s.name
	result["method"] = s.method
	result["last_request_timestamp"] = s.lastRequestTimestamp
	result["start_time"] = s.startTime
	result["num_requests"] = s.numRequests
	result["num_failures"] = s.numFailures
	result["total_response_time"] = s.totalResponseTime
	result["max_response_time"] = s.maxResponseTime
	result["min_response_time"] = s.minResponseTime
	result["total_content_length"] = s.totalContentLength
	result["response_times"] = s.responseTimes
	result["num_reqs_per_sec"] = s.numReqsPerSec
	return result
}

func (s *statsEntry) getStrippedReport() map[string]interface{} {
	report := s.serialize()
	s.reset()
	return report
}

type statsError struct {
	name       string
	method     string
	error      string
	occurences int64
}

func (err *statsError) occured() {
	err.occurences += 1
}

func (err *statsError) toMap() map[string]interface{} {
	m := make(map[string]interface{})
	m["method"] = err.method
	m["name"] = err.name
	m["error"] = err.error
	m["occurences"] = err.occurences
	return m
}

func collectReportData() map[string]interface{} {
	data := make(map[string]interface{})
	entries := make([]interface{}, 0, len(stats.entries))
	for _, v := range stats.entries {
		if !(v.numRequests == 0 && v.numFailures == 0) {
			entries = append(entries, v.getStrippedReport())
		}
	}

	errors := make(map[string]map[string]interface{})
	for k, v := range stats.errors {
		errors[k] = v.toMap()
	}

	data["stats"] = entries
	data["errors"] = errors
	stats.entries = make(map[string]*statsEntry)
	stats.errors = make(map[string]*statsError)

	return data
}

type requestSuccess struct {
	requestType    string
	name           string
	responseTime   float64
	responseLength int64
}

type requestFailure struct {
	requestType  string
	name         string
	responseTime float64
	error        string
}

var stats = new(requestStats)
var requestSuccessChannel = make(chan *requestSuccess, 100)
var requestFailureChannel = make(chan *requestFailure, 100)
var clearStatsChannel = make(chan bool)
var messageToServerChannel = make(chan map[string]interface{}, 10)

func init() {
	stats.entries = make(map[string]*statsEntry)
	stats.errors = make(map[string]*statsError)
	go func() {
		var ticker = time.NewTicker(SLAVE_REPORT_INTERVAL)
		for {
			select {
			case m := <-requestSuccessChannel:
				entry := stats.get(m.name, m.requestType)
				entry.log(m.responseTime, m.responseLength)
			case n := <-requestFailureChannel:
				stats.get(n.name, n.requestType).logError(n.error)
			case <-clearStatsChannel:
				stats.clearAll()
			case <-ticker.C:
				data := collectReportData()
				// send data to channel, no network IO in this goroutine
				messageToServerChannel <- data
			}
		}
	}()
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"strings"
)

// Command represents a CLI command.
// For example, in the command `heroku apps:create` the command would be `create`.
// They must have a Topic name that links to a real topic's name.
type Command struct {
	Topic        string             `json:"topic"`
	Command      string             `json:"command,omitempty"`
	Plugin       string             `json:"plugin"`
	Usage        string             `json:"usage"`
	Description  string             `json:"description"`
	Default      bool               `json:"default"`
	Help         string             `json:"help"`
	FullHelp     string             `json:"fullHelp"`
	Hidden       bool               `json:"hidden"`
	NeedsApp     bool               `json:"needsApp"`
	WantsApp     bool               `json:"wantsApp"`
	NeedsAuth    bool               `json:"needsAuth"`
	VariableArgs bool               `json:"variableArgs"`
	Args         []Arg              `json:"args"`
	Flags        []Flag             `json:"flags"`
	Run          func(ctx *Context) `json:"-"`
}

func (c *Command) String() string {
	if c.Command == "" {
		return c.Topic
	}
	return c.Topic + ":" + c.Command
}

func commandUsage(c *Command) string {
	return c.String() + argsString(c.Args)
}

func (c *Command) buildFullHelp() string {
	flags := c.Flags
	if c.NeedsApp || c.WantsApp {
		flags = append(flags, *appFlag, *remoteFlag)
	}
	lines := make([]string, 0, len(flags))
	if c.Description != "" {
		lines = append(lines, c.Description, "")
	}
	if len(flags) > 0 {
		for _, flag := range flags {
			if flag.Hidden {
				continue
			}
			if flag.Description == "" {
				lines = append(lines, flag.String())
			} else {
				lines = append(lines, fmt.Sprintf("%-20s # %s", flag.String(), flag.Description))
			}
		}
		lines = append(lines, "")
	}
	if c.Help != "" {
		lines = append(lines, c.Help)
	}
	return strings.TrimSuffix(strings.Join(lines, "\n"), "\n")
}

// CommandSet is a slice of Command structs with some helper methods.
type CommandSet []*Command

// ByTopicAndCommand returns a command that matches the passed topic and command.
func (commands CommandSet) ByTopicAndCommand(topic, command string) *Command {
	for _, c := range commands {
		if c.Topic == topic {
			if c.Command == command || c.Default && command == "" {
				return c
			}
		}
	}
	return nil
}

func (commands CommandSet) loadUsages() {
	for _, c := range commands {
		if c.Usage == "" {
			c.Usage = commandUsage(c)
		}
	}
}

func (commands CommandSet) loadFullHelp() {
	for _, c := range commands {
		if c.FullHelp == "" {
			c.FullHelp = c.buildFullHelp()
		}
	}
}

func (commands CommandSet) Len() int {
	return len(commands)
}

func (commands CommandSet) Less(i, j int) bool {
	return commands[i].Command < commands[j].Command
}

func (commands CommandSet) Swap(i, j int) {
	commands[i], commands[j] = commands[j], commands[i]
}

// Arg defines an argument for a command.
// These will be parsed in Go and passed to the Run method in the Context struct.
type Arg struct {
	Name     string `json:"name"`
	Optional bool   `json:"optional"`
	Hidden   bool   `json:"hidden"`
}

func (a *Arg) String() string {
	if a.Optional {
		return "[" + strings.ToUpper(a.Name) + "]"
	}
	return strings.ToUpper(a.Name)
}

func argsString(args []Arg) string {
	var buffer bytes.Buffer
	for _, arg := range args {
		if arg.Hidden {
			continue
		}
		if arg.Optional {
			buffer.WriteString(" [" + strings.ToUpper(arg.Name) + "]")
		} else {
			buffer.WriteString(" " + strings.ToUpper(arg.Name))
		}
	}
	return buffer.String()
}

var commandsTopic = &Topic{
	Name:        "commands",
	Description: "list all commands",
	Hidden:      true,
}

var commandsListCmd = &Command{
	Topic:       "commands",
	Description: "list all commands",
	Flags:       []Flag{{Name: "json"}},
	Run: func(ctx *Context) {
		cli.LoadPlugins(GetPlugins())
		if ctx.Flags["json"] == true {
			cli.Commands.loadUsages()
			cli.Commands.loadFullHelp()
			doc := map[string]interface{}{"topics": cli.Topics, "commands": cli.Commands}
			s, _ := json.Marshal(doc)
			Println(string(s))
			return
		}
		for _, command := range cli.Commands {
			if command.Hidden {
				continue
			}
			if command.Command == "" {
				Printf("%s\n", command.Topic)
			} else {
				Printf("%s:%s\n", command.Topic, command.Command)
			}
		}
	},
}
<commit_msg>setup builtin plugins before running commands command<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"strings"
)

// Command represents a CLI command.
// For example, in the command `heroku apps:create` the command would be `create`.
// They must have a Topic name that links to a real topic's name.
type Command struct {
	Topic        string             `json:"topic"`
	Command      string             `json:"command,omitempty"`
	Plugin       string             `json:"plugin"`
	Usage        string             `json:"usage"`
	Description  string             `json:"description"`
	Default      bool               `json:"default"`
	Help         string             `json:"help"`
	FullHelp     string             `json:"fullHelp"`
	Hidden       bool               `json:"hidden"`
	NeedsApp     bool               `json:"needsApp"`
	WantsApp     bool               `json:"wantsApp"`
	NeedsAuth    bool               `json:"needsAuth"`
	VariableArgs bool               `json:"variableArgs"`
	Args         []Arg              `json:"args"`
	Flags        []Flag             `json:"flags"`
	Run          func(ctx *Context) `json:"-"`
}

func (c *Command) String() string {
	if c.Command == "" {
		return c.Topic
	}
	return c.Topic + ":" + c.Command
}

func commandUsage(c *Command) string {
	return c.String() + argsString(c.Args)
}

func (c *Command) buildFullHelp() string {
	flags := c.Flags
	if c.NeedsApp || c.WantsApp {
		flags = append(flags, *appFlag, *remoteFlag)
	}
	lines := make([]string, 0, len(flags))
	if c.Description != "" {
		lines = append(lines, c.Description, "")
	}
	if len(flags) > 0 {
		for _, flag := range flags {
			if flag.Hidden {
				continue
			}
			if flag.Description == "" {
				lines = append(lines, flag.String())
			} else {
				lines = append(lines, fmt.Sprintf("%-20s # %s", flag.String(), flag.Description))
			}
		}
		lines = append(lines, "")
	}
	if c.Help != "" {
		lines = append(lines, c.Help)
	}
	return strings.TrimSuffix(strings.Join(lines, "\n"), "\n")
}

// CommandSet is a slice of Command structs with some helper methods.
type CommandSet []*Command

// ByTopicAndCommand returns a command that matches the passed topic and command.
func (commands CommandSet) ByTopicAndCommand(topic, command string) *Command {
	for _, c := range commands {
		if c.Topic == topic {
			if c.Command == command || c.Default && command == "" {
				return c
			}
		}
	}
	return nil
}

func (commands CommandSet) loadUsages() {
	for _, c := range commands {
		if c.Usage == "" {
			c.Usage = commandUsage(c)
		}
	}
}

func (commands CommandSet) loadFullHelp() {
	for _, c := range commands {
		if c.FullHelp == "" {
			c.FullHelp = c.buildFullHelp()
		}
	}
}

func (commands CommandSet) Len() int {
	return len(commands)
}

func (commands CommandSet) Less(i, j int) bool {
	return commands[i].Command < commands[j].Command
}

func (commands CommandSet) Swap(i, j int) {
	commands[i], commands[j] = commands[j], commands[i]
}

// Arg defines an argument for a command.
// These will be parsed in Go and passed to the Run method in the Context struct.
type Arg struct {
	Name     string `json:"name"`
	Optional bool   `json:"optional"`
	Hidden   bool   `json:"hidden"`
}

func (a *Arg) String() string {
	if a.Optional {
		return "[" + strings.ToUpper(a.Name) + "]"
	}
	return strings.ToUpper(a.Name)
}

func argsString(args []Arg) string {
	var buffer bytes.Buffer
	for _, arg := range args {
		if arg.Hidden {
			continue
		}
		if arg.Optional {
			buffer.WriteString(" [" + strings.ToUpper(arg.Name) + "]")
		} else {
			buffer.WriteString(" " + strings.ToUpper(arg.Name))
		}
	}
	return buffer.String()
}

var commandsTopic = &Topic{
	Name:        "commands",
	Description: "list all commands",
	Hidden:      true,
}

var commandsListCmd = &Command{
	Topic:       "commands",
	Description: "list all commands",
	Flags:       []Flag{{Name: "json"}},
	Run: func(ctx *Context) {
		SetupBuiltinPlugins()
		cli.LoadPlugins(GetPlugins())
		if ctx.Flags["json"] == true {
			cli.Commands.loadUsages()
			cli.Commands.loadFullHelp()
			doc := map[string]interface{}{"topics": cli.Topics, "commands": cli.Commands}
			s, _ := json.Marshal(doc)
			Println(string(s))
			return
		}
		for _, command := range cli.Commands {
			if command.Hidden {
				continue
			}
			if command.Command == "" {
				Printf("%s\n", command.Topic)
			} else {
				Printf("%s:%s\n", command.Topic, command.Command)
			}
		}
	},
}
<|endoftext|>


<commit_before>package hut

import (
	"github.com/cactus/go-statsd-client/statsd"
)

type Statter statsd.Statter

func (s *Service) NewStatsd() (client Statter) {
	address := s.Env.GetUDPServiceAddress("STATSD", 8125)
	if address == "" {
		client, _ = statsd.NewNoop()
		return client
	}
	// TODO: put the right prefix in here
	client, err := statsd.New(address, s.name)
	if err != nil {
		s.Log.Error().Printf("Could not connect to statsd: %s\n", err)
		client, _ = statsd.NewNoop()
	}
	return
}
<commit_msg>We have a prefix, remove the TODO<commit_after>package hut

import (
	"github.com/cactus/go-statsd-client/statsd"
)

type Statter statsd.Statter

func (s *Service) NewStatsd() (client Statter) {
	address := s.Env.GetUDPServiceAddress("STATSD", 8125)
	if address == "" {
		client, _ = statsd.NewNoop()
		return client
	}
	client, err := statsd.New(address, s.name)
	if err != nil {
		s.Log.Error().Printf("Could not connect to statsd: %s\n", err)
		client, _ = statsd.NewNoop()
	}
	return
}
<|endoftext|>


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

package hostagent

import (
	"encoding/json"
	"fmt"
	"net"
	"reflect"

	etcdclient "github.com/coreos/etcd/client"
	"golang.org/x/net/context"

	etcd "github.com/noironetworks/aci-containers/pkg/cf_etcd"
	md "github.com/noironetworks/aci-containers/pkg/metadata"
)

func (env *CfEnvironment) updateContainerMetadata(metadataKey *string) {
	ctId := extractContainerIdFromMetadataKey(metadataKey)
	if ctId == "" {
		return
	}

	env.agent.indexMutex.Lock()
	md, ok := env.agent.epMetadata[*metadataKey]
	env.agent.indexMutex.Unlock()

	kapi := env.etcdKeysApi
	key := etcd.CONTROLLER_KEY_BASE + "/containers/" + ctId
	var err error
	if !ok {
		_, err = kapi.Delete(context.Background(), key, &etcdclient.DeleteOptions{Recursive: true})
		if err != nil {
			keyerr, ok := err.(etcdclient.Error)
			if ok && keyerr.Code == etcdclient.ErrorCodeKeyNotFound {
				env.log.Info(fmt.Sprintf("Etcd subtree %s doesn't exist yet", key))
				err = nil
			}
		}
	} else {
		if md[ctId] != nil && md[ctId].Ifaces != nil {
			var md_json []byte
			md_json, err = json.Marshal(md[ctId].Ifaces)
			if err == nil {
				_, err = kapi.Set(context.Background(), key+"/metadata", string(md_json), nil)
			}
		}
	}
	if err != nil {
		env.log.Error("Failed to update container metadata in etcd: ", err)
	}
}

func (env *CfEnvironment) cfAppContainerChanged(ctId *string, ep *etcd.EpInfo) {
	if ep == nil {
		return
	}
	metaKey := "_cf_/" + *ctId

	epGroup := &md.OpflexGroup{PolicySpace: ep.EpgTenant, Name: ep.Epg}
	secGroup := make([]md.OpflexGroup, len(ep.SecurityGroups))
	for i, s := range ep.SecurityGroups {
		secGroup[i].PolicySpace = s.Tenant
		secGroup[i].Name = s.Group
	}

	epAttributes := make(map[string]string)
	epAttributes["vm-name"] = ep.EpName(*ctId)
	epAttributes["app-id"] = ep.AppId
	epAttributes["space-id"] = ep.SpaceId
	epAttributes["org-id"] = ep.OrgId
	epAttributes["container-id"] = *ctId

	// Update iptables rules and container ports-set
	cportset := make(map[uint32]struct{})
	env.indexLock.Lock()
	for p := range env.cfNetContainerPorts {
		cportset[p] = struct{}{}
	}
	// pre-routing DNAT rules
	env.updatePreNatRule(ctId, ep, ep.PortMapping)
	// post-routing SNAT rules
	for _, pmap := range ep.PortMapping {
		cport := fmt.Sprintf("%d", pmap.ContainerPort)
		err := env.iptbl.AppendUnique("nat", NAT_POST_CHAIN, "-o", env.cfconfig.CfNetOvsPort, "-p", "tcp",
			"-m", "tcp", "--dport", cport, "-j", "SNAT", "--to-source",
			env.cfconfig.CfNetIntfAddress)
		if err != nil {
			env.log.Warning("Failed to add post-routing iptables rule: ", err)
		}
		cportset[pmap.ContainerPort] = struct{}{}
	}
	cfnet_update := !reflect.DeepEqual(env.cfNetContainerPorts, cportset)
	if cfnet_update {
		env.cfNetContainerPorts = cportset
	}
	env.indexLock.Unlock()

	env.agent.indexMutex.Lock()
	env.agent.epChanged(ctId, &metaKey, epGroup, secGroup, epAttributes, nil)
	if cfnet_update {
		env.updateLegacyCfNetService(cportset)
	}
	env.agent.indexMutex.Unlock()
}

// must be called with env.indexLock
func (env *CfEnvironment) updatePreNatRule(ctId *string, ep *etcd.EpInfo, portmap []etcd.PortMap) {
	ctIp := net.ParseIP(ep.IpAddress)
	if ctIp == nil || (env.cfNetv4 && ctIp.To4() == nil) {
		return
	}
	old_pm := env.ctPortMap[*ctId]
	new_pm := make(map[uint32]uint32)
	for _, ch := range portmap {
		err := env.iptbl.AppendUnique("nat", NAT_PRE_CHAIN, "-d", env.cfconfig.CellAddress, "-p", "tcp",
			"--dport", fmt.Sprintf("%d", ch.HostPort),
			"-j", "DNAT", "--to-destination",
			ep.IpAddress+":"+fmt.Sprintf("%d", ch.ContainerPort))
		if err != nil {
			env.log.Warning(fmt.Sprintf("Failed to add pre-routing iptables rule for %d: %v", *ctId, err))
		}
		new_pm[ch.HostPort] = ch.ContainerPort
		delete(old_pm, ch.HostPort)
	}
	for hp, cp := range old_pm {
		args := []string{"-d", env.cfconfig.CellAddress, "-p", "tcp", "--dport",
			fmt.Sprintf("%d", hp), "-j", "DNAT", "--to-destination",
			ep.IpAddress + ":" + fmt.Sprintf("%d", cp)}
		exist, _ := env.iptbl.Exists("nat", NAT_PRE_CHAIN, args...)
		if !exist {
			continue
		}
		err := env.iptbl.Delete("nat", NAT_PRE_CHAIN, args...)
		if err != nil {
			env.log.Warning(fmt.Sprintf("Failed to delete pre-routing iptables rule for %d: %v", *ctId, err))
		}
	}
	env.ctPortMap[*ctId] = new_pm
}

func (env *CfEnvironment) cfAppContainerDeleted(ctId *string, ep *etcd.EpInfo) {
	env.agent.indexMutex.Lock()
	env.agent.epDeleted(ctId)
	env.agent.indexMutex.Unlock()

	if ep == nil {
		return
	}
	env.indexLock.Lock()
	defer env.indexLock.Unlock()
	env.updatePreNatRule(ctId, ep, nil)
	delete(env.ctPortMap, *ctId)
}

func (env *CfEnvironment) updateLegacyCfNetService(portmap map[uint32]struct{}) error {
	// should be called with agent.indexMutex held
	uuid := "cf-net-" + env.cfconfig.CellID
	new_svc := opflexService{Uuid: uuid,
		DomainPolicySpace: env.agent.config.AciVrfTenant,
		DomainName:        env.agent.config.AciVrf,
		ServiceMac:        env.cfNetLink.Attrs().HardwareAddr.String(),
		InterfaceName:     env.cfconfig.CfNetOvsPort}
	for p := range portmap {
		svc_map := opflexServiceMapping{ServiceIp: env.cfconfig.CfNetIntfAddress,
			ServicePort: uint16(p),
			NextHopIps:  make([]string, 0)}
		new_svc.ServiceMappings = append(new_svc.ServiceMappings, svc_map)
	}
	exist, ok := env.agent.opflexServices[uuid]
	if !ok || !reflect.DeepEqual(*exist, new_svc) {
		env.log.Debug("Updating CF legacy-networking service ", uuid)
		env.agent.opflexServices[uuid] = &new_svc
		env.agent.scheduleSyncServices()
	}
	return nil
}

func (env *CfEnvironment) cfAppDeleted(appId *string, app *etcd.AppInfo) {
	env.agent.indexMutex.Lock()
	defer env.agent.indexMutex.Unlock()
	uuid := *appId
	_, vip_ok := env.agent.opflexServices[uuid]
	if vip_ok {
		env.log.Debug("Removing service CF app vip/ext-ip ", uuid)
		delete(env.agent.opflexServices, uuid)
	}
	uuid += "-external"
	_, ext_ip_ok := env.agent.opflexServices[uuid]
	if ext_ip_ok {
		env.log.Debug("Removing service CF app vip/ext-ip ", uuid)
		delete(env.agent.opflexServices, uuid)
	}
	if vip_ok || ext_ip_ok {
		env.agent.scheduleSyncServices()
	}
}

// 0 -> ipv4, 1 -> ipv6, anything else -> invalid IP
func getIpType(ip_str string) int {
	ip := net.ParseIP(ip_str)
	if ip == nil {
		return -1
	}
	if ip.To4() != nil {
		return 0
	}
	if ip.To16() != nil {
		return 1
	}
	return -2
}

func (env *CfEnvironment) cfAppIdChanged(appId *string) {
	env.indexLock.Lock()
	appInfo := env.appIdx[*appId]
	env.indexLock.Unlock()
	if appInfo != nil {
		env.cfAppChanged(appId, appInfo)
	}
}

func (env *CfEnvironment) cfAppChanged(appId *string, app *etcd.AppInfo) {
	env.updateCfAppServiceEp(appId, app, false)
	env.updateCfAppServiceEp(appId, app, true)
}

func (env *CfEnvironment) updateCfAppServiceEp(appId *string, app *etcd.AppInfo, external bool) {
	agent := env.agent
	uuid := *appId
	if external {
		uuid += "-external"
	}
	appas := opflexService{
		Uuid:              uuid,
		DomainPolicySpace: agent.config.AciVrfTenant,
		DomainName:        agent.config.AciVrf,
		ServiceMode:       "loadbalancer",
		ServiceMappings:   make([]opflexServiceMapping, 0),
	}
	if external && agent.config.UplinkIface != "" && agent.serviceEp.Mac != "" &&
		(agent.serviceEp.Ipv4 != nil || agent.serviceEp.Ipv6 != nil) {

		appas.InterfaceName = agent.config.UplinkIface
		appas.InterfaceVlan = uint16(agent.config.ServiceVlan)
		appas.ServiceMac = agent.serviceEp.Mac
		if agent.serviceEp.Ipv4 != nil {
			appas.InterfaceIp = agent.serviceEp.Ipv4.String()
		} else {
			appas.InterfaceIp = agent.serviceEp.Ipv6.String() // TODO dual stack?
		}
	}
	ips := app.VirtualIp
	if external {
		ips = app.ExternalIp
	}
	localContainerIps := make([]string, 0)
	agent.indexMutex.Lock()
	defer agent.indexMutex.Unlock()
	if external {
		for _, c := range env.epIdx {
			if c != nil && c.AppId == *appId {
				localContainerIps = append(localContainerIps, c.IpAddress)
			}
		}
	}
	for _, vip := range ips {
		ipt := getIpType(vip)
		if ipt != 0 && ipt != 1 {
			continue
		}
		sm := opflexServiceMapping{
			ServiceIp:  vip,
			NextHopIps: make([]string, 0),
			Conntrack:  true,
		}
		containers := app.ContainerIps
		if external {
			containers = localContainerIps
		}
		for _, cip := range containers {
			if ipt != getIpType(cip) {
				continue
			}
			sm.NextHopIps = append(sm.NextHopIps, cip)
		}
		if len(sm.NextHopIps) > 0 {
			appas.ServiceMappings = append(appas.ServiceMappings, sm)
		}
	}
	valid := len(appas.ServiceMappings) > 0

	exist, ok := env.agent.opflexServices[uuid]
	if valid && (!ok || !reflect.DeepEqual(*exist, appas)) {
		env.log.Debug("Updating CF app vip/ext-ip service ", uuid)
		env.agent.opflexServices[uuid] = &appas
		env.agent.scheduleSyncServices()
	} else if !valid && ok {
		env.log.Debug("Removing CF app vip/ext-ip service ", uuid)
		delete(env.agent.opflexServices, uuid)
		env.agent.scheduleSyncServices()
	}
}
<commit_msg>Fix format strings "%d" -> "%s"<commit_after>// Copyright 2017 Cisco Systems, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package hostagent

import (
	"encoding/json"
	"fmt"
	"net"
	"reflect"

	etcdclient "github.com/coreos/etcd/client"
	"golang.org/x/net/context"

	etcd "github.com/noironetworks/aci-containers/pkg/cf_etcd"
	md "github.com/noironetworks/aci-containers/pkg/metadata"
)

func (env *CfEnvironment) updateContainerMetadata(metadataKey *string) {
	ctId := extractContainerIdFromMetadataKey(metadataKey)
	if ctId == "" {
		return
	}

	env.agent.indexMutex.Lock()
	md, ok := env.agent.epMetadata[*metadataKey]
	env.agent.indexMutex.Unlock()

	kapi := env.etcdKeysApi
	key := etcd.CONTROLLER_KEY_BASE + "/containers/" + ctId
	var err error
	if !ok {
		_, err = kapi.Delete(context.Background(), key, &etcdclient.DeleteOptions{Recursive: true})
		if err != nil {
			keyerr, ok := err.(etcdclient.Error)
			if ok && keyerr.Code == etcdclient.ErrorCodeKeyNotFound {
				env.log.Info(fmt.Sprintf("Etcd subtree %s doesn't exist yet", key))
				err = nil
			}
		}
	} else {
		if md[ctId] != nil && md[ctId].Ifaces != nil {
			var md_json []byte
			md_json, err = json.Marshal(md[ctId].Ifaces)
			if err == nil {
				_, err = kapi.Set(context.Background(), key+"/metadata", string(md_json), nil)
			}
		}
	}
	if err != nil {
		env.log.Error("Failed to update container metadata in etcd: ", err)
	}
}

func (env *CfEnvironment) cfAppContainerChanged(ctId *string, ep *etcd.EpInfo) {
	if ep == nil {
		return
	}
	metaKey := "_cf_/" + *ctId

	epGroup := &md.OpflexGroup{PolicySpace: ep.EpgTenant, Name: ep.Epg}
	secGroup := make([]md.OpflexGroup, len(ep.SecurityGroups))
	for i, s := range ep.SecurityGroups {
		secGroup[i].PolicySpace = s.Tenant
		secGroup[i].Name = s.Group
	}

	epAttributes := make(map[string]string)
	epAttributes["vm-name"] = ep.EpName(*ctId)
	epAttributes["app-id"] = ep.AppId
	epAttributes["space-id"] = ep.SpaceId
	epAttributes["org-id"] = ep.OrgId
	epAttributes["container-id"] = *ctId

	// Update iptables rules and container ports-set
	cportset := make(map[uint32]struct{})
	env.indexLock.Lock()
	for p := range env.cfNetContainerPorts {
		cportset[p] = struct{}{}
	}
	// pre-routing DNAT rules
	env.updatePreNatRule(ctId, ep, ep.PortMapping)
	// post-routing SNAT rules
	for _, pmap := range ep.PortMapping {
		cport := fmt.Sprintf("%d", pmap.ContainerPort)
		err := env.iptbl.AppendUnique("nat", NAT_POST_CHAIN, "-o", env.cfconfig.CfNetOvsPort, "-p", "tcp",
			"-m", "tcp", "--dport", cport, "-j", "SNAT", "--to-source",
			env.cfconfig.CfNetIntfAddress)
		if err != nil {
			env.log.Warning("Failed to add post-routing iptables rule: ", err)
		}
		cportset[pmap.ContainerPort] = struct{}{}
	}
	cfnet_update := !reflect.DeepEqual(env.cfNetContainerPorts, cportset)
	if cfnet_update {
		env.cfNetContainerPorts = cportset
	}
	env.indexLock.Unlock()

	env.agent.indexMutex.Lock()
	env.agent.epChanged(ctId, &metaKey, epGroup, secGroup, epAttributes, nil)
	if cfnet_update {
		env.updateLegacyCfNetService(cportset)
	}
	env.agent.indexMutex.Unlock()
}

// must be called with env.indexLock
func (env *CfEnvironment) updatePreNatRule(ctId *string, ep *etcd.EpInfo, portmap []etcd.PortMap) {
	ctIp := net.ParseIP(ep.IpAddress)
	if ctIp == nil || (env.cfNetv4 && ctIp.To4() == nil) {
		return
	}
	old_pm := env.ctPortMap[*ctId]
	new_pm := make(map[uint32]uint32)
	for _, ch := range portmap {
		err := env.iptbl.AppendUnique("nat", NAT_PRE_CHAIN, "-d",
			env.cfconfig.CellAddress, "-p", "tcp",
			"--dport", fmt.Sprintf("%d", ch.HostPort),
			"-j", "DNAT", "--to-destination",
			ep.IpAddress+":"+fmt.Sprintf("%d", ch.ContainerPort))
		if err != nil {
			env.log.Warning(fmt.Sprintf("Failed to add pre-routing "+
				"iptables rule for %s: %v", *ctId, err))
		}
		new_pm[ch.HostPort] = ch.ContainerPort
		delete(old_pm, ch.HostPort)
	}
	for hp, cp := range old_pm {
		args := []string{"-d", env.cfconfig.CellAddress, "-p", "tcp", "--dport",
			fmt.Sprintf("%d", hp), "-j", "DNAT", "--to-destination",
			ep.IpAddress + ":" + fmt.Sprintf("%d", cp)}
		exist, _ := env.iptbl.Exists("nat", NAT_PRE_CHAIN, args...)
		if !exist {
			continue
		}
		err := env.iptbl.Delete("nat", NAT_PRE_CHAIN, args...)
		if err != nil {
			env.log.Warning(fmt.Sprintf("Failed to delete pre-routing "+
				"iptables rule for %s: %v", *ctId, err))
		}
	}
	env.ctPortMap[*ctId] = new_pm
}

func (env *CfEnvironment) cfAppContainerDeleted(ctId *string, ep *etcd.EpInfo) {
	env.agent.indexMutex.Lock()
	env.agent.epDeleted(ctId)
	env.agent.indexMutex.Unlock()

	if ep == nil {
		return
	}
	env.indexLock.Lock()
	defer env.indexLock.Unlock()
	env.updatePreNatRule(ctId, ep, nil)
	delete(env.ctPortMap, *ctId)
}

func (env *CfEnvironment) updateLegacyCfNetService(portmap map[uint32]struct{}) error {
	// should be called with agent.indexMutex held
	uuid := "cf-net-" + env.cfconfig.CellID
	new_svc := opflexService{Uuid: uuid,
		DomainPolicySpace: env.agent.config.AciVrfTenant,
		DomainName:        env.agent.config.AciVrf,
		ServiceMac:        env.cfNetLink.Attrs().HardwareAddr.String(),
		InterfaceName:     env.cfconfig.CfNetOvsPort}
	for p := range portmap {
		svc_map := opflexServiceMapping{ServiceIp: env.cfconfig.CfNetIntfAddress,
			ServicePort: uint16(p),
			NextHopIps:  make([]string, 0)}
		new_svc.ServiceMappings = append(new_svc.ServiceMappings, svc_map)
	}
	exist, ok := env.agent.opflexServices[uuid]
	if !ok || !reflect.DeepEqual(*exist, new_svc) {
		env.log.Debug("Updating CF legacy-networking service ", uuid)
		env.agent.opflexServices[uuid] = &new_svc
		env.agent.scheduleSyncServices()
	}
	return nil
}

func (env *CfEnvironment) cfAppDeleted(appId *string, app *etcd.AppInfo) {
	env.agent.indexMutex.Lock()
	defer env.agent.indexMutex.Unlock()
	uuid := *appId
	_, vip_ok := env.agent.opflexServices[uuid]
	if vip_ok {
		env.log.Debug("Removing service CF app vip/ext-ip ", uuid)
		delete(env.agent.opflexServices, uuid)
	}
	uuid += "-external"
	_, ext_ip_ok := env.agent.opflexServices[uuid]
	if ext_ip_ok {
		env.log.Debug("Removing service CF app vip/ext-ip ", uuid)
		delete(env.agent.opflexServices, uuid)
	}
	if vip_ok || ext_ip_ok {
		env.agent.scheduleSyncServices()
	}
}

// 0 -> ipv4, 1 -> ipv6, anything else -> invalid IP
func getIpType(ip_str string) int {
	ip := net.ParseIP(ip_str)
	if ip == nil {
		return -1
	}
	if ip.To4() != nil {
		return 0
	}
	if ip.To16() != nil {
		return 1
	}
	return -2
}

func (env *CfEnvironment) cfAppIdChanged(appId *string) {
	env.indexLock.Lock()
	appInfo := env.appIdx[*appId]
	env.indexLock.Unlock()
	if appInfo != nil {
		env.cfAppChanged(appId, appInfo)
	}
}

func (env *CfEnvironment) cfAppChanged(appId *string, app *etcd.AppInfo) {
	env.updateCfAppServiceEp(appId, app, false)
	env.updateCfAppServiceEp(appId, app, true)
}

func (env *CfEnvironment) updateCfAppServiceEp(appId *string, app *etcd.AppInfo, external bool) {
	agent := env.agent
	uuid := *appId
	if external {
		uuid += "-external"
	}
	appas := opflexService{
		Uuid:              uuid,
		DomainPolicySpace: agent.config.AciVrfTenant,
		DomainName:        agent.config.AciVrf,
		ServiceMode:       "loadbalancer",
		ServiceMappings:   make([]opflexServiceMapping, 0),
	}
	if external && agent.config.UplinkIface != "" && agent.serviceEp.Mac != "" &&
		(agent.serviceEp.Ipv4 != nil || agent.serviceEp.Ipv6 != nil) {

		appas.InterfaceName = agent.config.UplinkIface
		appas.InterfaceVlan = uint16(agent.config.ServiceVlan)
		appas.ServiceMac = agent.serviceEp.Mac
		if agent.serviceEp.Ipv4 != nil {
			appas.InterfaceIp = agent.serviceEp.Ipv4.String()
		} else {
			appas.InterfaceIp = agent.serviceEp.Ipv6.String() // TODO dual stack?
		}
	}
	ips := app.VirtualIp
	if external {
		ips = app.ExternalIp
	}
	localContainerIps := make([]string, 0)
	agent.indexMutex.Lock()
	defer agent.indexMutex.Unlock()
	if external {
		for _, c := range env.epIdx {
			if c != nil && c.AppId == *appId {
				localContainerIps = append(localContainerIps, c.IpAddress)
			}
		}
	}
	for _, vip := range ips {
		ipt := getIpType(vip)
		if ipt != 0 && ipt != 1 {
			continue
		}
		sm := opflexServiceMapping{
			ServiceIp:  vip,
			NextHopIps: make([]string, 0),
			Conntrack:  true,
		}
		containers := app.ContainerIps
		if external {
			containers = localContainerIps
		}
		for _, cip := range containers {
			if ipt != getIpType(cip) {
				continue
			}
			sm.NextHopIps = append(sm.NextHopIps, cip)
		}
		if len(sm.NextHopIps) > 0 {
			appas.ServiceMappings = append(appas.ServiceMappings, sm)
		}
	}
	valid := len(appas.ServiceMappings) > 0

	exist, ok := env.agent.opflexServices[uuid]
	if valid && (!ok || !reflect.DeepEqual(*exist, appas)) {
		env.log.Debug("Updating CF app vip/ext-ip service ", uuid)
		env.agent.opflexServices[uuid] = &appas
		env.agent.scheduleSyncServices()
	} else if !valid && ok {
		env.log.Debug("Removing CF app vip/ext-ip service ", uuid)
		delete(env.agent.opflexServices, uuid)
		env.agent.scheduleSyncServices()
	}
}
<|endoftext|>


<commit_before>// Package command implements helper functions for multi-command programs.
//
// This package includes types and functions for easily defining multiple
// subcommands with different options. A "help" subcommand is also automatically
// generated, which might be used to list all the available subcommands or
// to view all the help about a specific subcommand.
//
// Clients should just define a list of the command they implement and the call
// Run or RunArgs directly from main().
//
//  package main
//
//  import (
//	"gopkgs.com/command.v1"
//  )
//
//  var (
//	cmds = []*command.Cmd{
//	    {
//		Name: "awesome",
//		Help: "do something awesome"
//		Options: &awesomeOptions{Value:42},
//	    },
//	}
//  )
//
//  type awesomeOptions struct {
//	Value int `name:"v" help:"Some arbitrary value"`
//  }
//
//  func awesomeCommand(args []string, opts *AwesomeOptions) error {
//  ...
//  }
//
//  func main() {
//	commands.Run(cmds)
//  }
package command

import (
	"errors"
	"flag"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"reflect"
	"text/tabwriter"
)

var (
	argsType = reflect.TypeOf([]string(nil))
	// ErrNoCommand is returned from Run when no command
	// has been specified (i.e. there are no arguments).
	ErrNoCommand = errors.New("no command provided")
	// ErrHelp is returned from Run when the help is shown,
	// either the brief help or the detailed help for a
	// command.
	ErrHelp = errors.New("help has been shown")
	// ErrUnusedArguments is returned form Run when the command
	// does not accept any arguments, but the user has provided
	// some.
	ErrUnusedArguments = errors.New("arguments provided but not used")
)

// UnknownCommandError is returned from Run when the specified
// command does not exist.
type UnknownCommandError string

func (e UnknownCommandError) Error() string {
	return fmt.Sprintf("unknown command %s", string(e))
}

// Cmd represents an available command.
type Cmd struct {
	// Name is the name of the command, case sensitive.
	Name string
	// Help is a short, one line help string displayed by
	// the help command when listing all the commands.
	Help string
	// LongHelp is a long, potentially multi-line, help message
	// displayed when using the help command for a specific
	// command (e.g. myprogram help somecommand)
	LongHelp string
	// Usage is displayed when showing the help for a specific
	// command. The program name (os.Args[0]) and the command
	// name are prepended to it when displaying it to the user.
	// (e.g. Usage = <some-argument> shows "usage: myprog subcmd <some-argument>")
	Usage string
	// Func is the handler function for the command. The function must take either
	// one or two arguments. The first one must be a []string, which will contain
	// any non-flag arguments. If the function accepts a second argument it must
	// be of the exact same type than the value provided in the Options field.
	// Handler functions might optionally return an error value.
	Func interface{}
	// Options might be either nil or a pointer to a struct type. Command flags
	// will be generated from this struct, in the same order as the fields are
	// defined. The current value of the field will be used as the default value
	// for the flag. Each field might also include two struct tags:
	//
	//  - name: The name of the flag. If not present, it will default to the field name.
	//  - help: The short help shown the flag package for the given field.
	Options interface{}
}

// Run is a shorthand for RunArgs(os.Args[1:], commands).
func Run(commands []*Cmd) error {
	return RunArgs(os.Args[1:], commands)
}

// RunArgs tries to run a command from the specified list using the
// given arguments, interpreting the first argument as the command name.
// If the returned error is non-nil, it will be one of:
//
//  - ErrNoCommand when no arguments are provided
//  - ErrHelp when the user has requested any help to be shown
//  - ErrUnusedArguments when the command doesn't accept any arguments, but the user has provided some
//  - An UnknownCommandError when the command (the first argument) does not exist
//  - Any error returned by the command handler
//
// Any user error will be printed to Stderr by RunArgs, so callers don't need to print any
// error messages themselves.
//
// Note that RunArgs will panic in case of a programming error. This usually happens
// when Func or Options don't match the required constraints. See the documentation on
// those fields in the Cmd type for more information.
func RunArgs(args []string, commands []*Cmd) error {
	if len(args) == 0 || args[0] == "help" {
		return printHelp(os.Stderr, args, commands)
	}
	name := args[0]
	rem := args[1:]
	cmd := commandByName(commands, name)
	if cmd == nil {
		return printHelp(os.Stderr, args, commands)
	}
	fn := reflect.ValueOf(cmd.Func)
	if fn.Kind() != reflect.Func {
		panic(fmt.Errorf("command handler %s is not a function, it's %T", name, cmd.Func))
	}
	var optsVal reflect.Value
	var fnArgs []reflect.Value
	typ := fn.Type()
	numIn := typ.NumIn()
	if cmd.Options != nil {
		optsVal = reflect.ValueOf(cmd.Options)
		optsType := optsVal.Type()
		if numIn == 0 || (typ.In(0) != optsType && (numIn < 2 || typ.In(1) != optsType)) {
			panic(fmt.Errorf("command %s (%s) declares options of type %T but does not accept them", name, typ, cmd.Options))
		}
		flags, err := setupOptionsFlags(name, optsVal)
		if err != nil {
			panic(err)
		}
		if err := flags.Parse(rem); err != nil {
			return err
		}
		rem = flags.Args()
	}
	if numIn > 0 && typ.In(0) == argsType {
		fnArgs = append(fnArgs, reflect.ValueOf(rem))
	} else if len(rem) > 0 {
		fmt.Fprintf(os.Stderr, "command %s does not accept any arguments\n", name)
		return ErrUnusedArguments
	}
	if cmd.Options != nil {
		fnArgs = append(fnArgs, optsVal)
	}
	res := fn.Call(fnArgs)
	if len(res) > 0 {
		if err, ok := res[0].Interface().(error); ok {
			fmt.Fprintf(os.Stderr, "error running command %s: %s\n", name, err)
			return err
		}
	}
	return nil
}

func setupOptionsFlags(name string, val reflect.Value) (*flag.FlagSet, error) {
	if val.Kind() != reflect.Ptr {
		return nil, fmt.Errorf("invalid options %s, must be a pointer", val.Type())
	}
	val = reflect.Indirect(val)
	typ := val.Type()
	if typ.Kind() != reflect.Struct {
		return nil, fmt.Errorf("invalid Options type %s, must be an struct", typ)
	}
	flagsName := fmt.Sprintf("%s %s subcommand", filepath.Base(os.Args[0]), name)
	flags := flag.NewFlagSet(flagsName, flag.ContinueOnError)
	for ii := 0; ii < typ.NumField(); ii++ {
		field := typ.Field(ii)
		fieldVal := val.Field(ii)
		ptr := fieldVal.Addr().Interface()
		name := field.Name
		var help string
		if n := field.Tag.Get("name"); n != "" {
			name = n
		}
		if h := field.Tag.Get("help"); h != "" {
			help = h
		}
		if name == "" {
			return nil, fmt.Errorf("no name provided for field #%d in type %s", ii, typ)
		}
		if value, ok := ptr.(flag.Value); ok {
			flags.Var(value, name, help)
			continue
		}
		switch field.Type.Kind() {
		case reflect.Bool:
			flags.BoolVar(ptr.(*bool), name, fieldVal.Bool(), help)
		case reflect.Float64:
			flags.Float64Var(ptr.(*float64), name, fieldVal.Float(), help)
		case reflect.Int:
			flags.IntVar(ptr.(*int), name, int(fieldVal.Int()), help)
		case reflect.Uint:
			flags.UintVar(ptr.(*uint), name, uint(fieldVal.Uint()), help)
		case reflect.Int64:
			flags.Int64Var(ptr.(*int64), name, fieldVal.Int(), help)
		case reflect.Uint64:
			flags.Uint64Var(ptr.(*uint64), name, fieldVal.Uint(), help)
		case reflect.String:
			flags.StringVar(ptr.(*string), name, fieldVal.String(), help)
		default:
			return nil, fmt.Errorf("field %s has invalid option type %s", field.Name, field.Type)
		}
	}
	return flags, nil
}

func commandByName(commands []*Cmd, name string) *Cmd {
	for _, v := range commands {
		if v.Name == name {
			return v
		}
	}
	return nil
}

func printCommandHelp(w io.Writer, cmd *Cmd) {
	fmt.Fprintf(w, "%s: %s\n", cmd.Name, cmd.Help)
	if cmd.Usage != "" {
		fmt.Fprintf(w, "\nusage: %s %s %s\n", filepath.Base(os.Args[0]), cmd.Name, cmd.Usage)
	}
	if cmd.LongHelp != "" {
		fmt.Fprintf(w, "\n%s\n\n", cmd.LongHelp)
	}
	if cmd.Options != nil {
		opts := reflect.ValueOf(cmd.Options)
		if fs, err := setupOptionsFlags(cmd.Name, opts); err == nil {
			fs.PrintDefaults()
		}
	}
}

func printHelp(w io.Writer, args []string, commands []*Cmd) error {
	var err error
	if len(args) == 0 {
		fmt.Fprintln(w, "missing command, available ones are:\n")
		err = ErrNoCommand
	} else {
		var unknown string
		if args[0] != "help" && commandByName(commands, args[0]) == nil {
			unknown = args[0]
		}
		if len(args) > 1 && args[0] == "help" {
			if cmd := commandByName(commands, args[1]); cmd != nil {
				printCommandHelp(w, cmd)
				return ErrHelp
			}
			unknown = args[1]
		}
		if unknown != "" {
			fmt.Fprintf(w, "unknown command %s, available ones are:\n\n", unknown)
			err = UnknownCommandError(unknown)
		}
	}
	tw := tabwriter.NewWriter(w, 0, 8, 0, '\t', 0)
	for _, v := range commands {
		fmt.Fprintf(tw, "%s\t%s\n", v.Name, v.Help)
	}
	fmt.Fprint(tw, "help\tPrint this help\n")
	tw.Flush()
	fmt.Fprint(w, "\nTo view additional help for each command use help <command_name>\n")
	if err == nil {
		err = ErrHelp
	}
	return err
}
<commit_msg>Default to the field name in lowercase when there's no name tag<commit_after>// Package command implements helper functions for multi-command programs.
//
// This package includes types and functions for easily defining multiple
// subcommands with different options. A "help" subcommand is also automatically
// generated, which might be used to list all the available subcommands or
// to view all the help about a specific subcommand.
//
// Clients should just define a list of the command they implement and the call
// Run or RunArgs directly from main().
//
//  package main
//
//  import (
//	"gopkgs.com/command.v1"
//  )
//
//  var (
//	cmds = []*command.Cmd{
//	    {
//		Name: "awesome",
//		Help: "do something awesome"
//		Options: &awesomeOptions{Value:42},
//	    },
//	}
//  )
//
//  type awesomeOptions struct {
//	Value int `name:"v" help:"Some arbitrary value"`
//  }
//
//  func awesomeCommand(args []string, opts *AwesomeOptions) error {
//  ...
//  }
//
//  func main() {
//	commands.Run(cmds)
//  }
package command

import (
	"errors"
	"flag"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"reflect"
	"strings"
	"text/tabwriter"
)

var (
	argsType = reflect.TypeOf([]string(nil))
	// ErrNoCommand is returned from Run when no command
	// has been specified (i.e. there are no arguments).
	ErrNoCommand = errors.New("no command provided")
	// ErrHelp is returned from Run when the help is shown,
	// either the brief help or the detailed help for a
	// command.
	ErrHelp = errors.New("help has been shown")
	// ErrUnusedArguments is returned form Run when the command
	// does not accept any arguments, but the user has provided
	// some.
	ErrUnusedArguments = errors.New("arguments provided but not used")
)

// UnknownCommandError is returned from Run when the specified
// command does not exist.
type UnknownCommandError string

func (e UnknownCommandError) Error() string {
	return fmt.Sprintf("unknown command %s", string(e))
}

// Cmd represents an available command.
type Cmd struct {
	// Name is the name of the command, case sensitive.
	Name string
	// Help is a short, one line help string displayed by
	// the help command when listing all the commands.
	Help string
	// LongHelp is a long, potentially multi-line, help message
	// displayed when using the help command for a specific
	// command (e.g. myprogram help somecommand)
	LongHelp string
	// Usage is displayed when showing the help for a specific
	// command. The program name (os.Args[0]) and the command
	// name are prepended to it when displaying it to the user.
	// (e.g. Usage = <some-argument> shows "usage: myprog subcmd <some-argument>")
	Usage string
	// Func is the handler function for the command. The function must take either
	// one or two arguments. The first one must be a []string, which will contain
	// any non-flag arguments. If the function accepts a second argument it must
	// be of the exact same type than the value provided in the Options field.
	// Handler functions might optionally return an error value.
	Func interface{}
	// Options might be either nil or a pointer to a struct type. Command flags
	// will be generated from this struct, in the same order as the fields are
	// defined. The current value of the field will be used as the default value
	// for the flag. Each field might also include two struct tags:
	//
	//  - name: The name of the flag. If not present, it will default to the field name in lowercase.
	//  - help: The short help shown the flag package for the given field.
	Options interface{}
}

// Run is a shorthand for RunArgs(os.Args[1:], commands).
func Run(commands []*Cmd) error {
	return RunArgs(os.Args[1:], commands)
}

// RunArgs tries to run a command from the specified list using the
// given arguments, interpreting the first argument as the command name.
// If the returned error is non-nil, it will be one of:
//
//  - ErrNoCommand when no arguments are provided
//  - ErrHelp when the user has requested any help to be shown
//  - ErrUnusedArguments when the command doesn't accept any arguments, but the user has provided some
//  - An UnknownCommandError when the command (the first argument) does not exist
//  - Any error returned by the command handler
//
// Any user error will be printed to Stderr by RunArgs, so callers don't need to print any
// error messages themselves.
//
// Note that RunArgs will panic in case of a programming error. This usually happens
// when Func or Options don't match the required constraints. See the documentation on
// those fields in the Cmd type for more information.
func RunArgs(args []string, commands []*Cmd) error {
	if len(args) == 0 || args[0] == "help" {
		return printHelp(os.Stderr, args, commands)
	}
	name := args[0]
	rem := args[1:]
	cmd := commandByName(commands, name)
	if cmd == nil {
		return printHelp(os.Stderr, args, commands)
	}
	fn := reflect.ValueOf(cmd.Func)
	if fn.Kind() != reflect.Func {
		panic(fmt.Errorf("command handler %s is not a function, it's %T", name, cmd.Func))
	}
	var optsVal reflect.Value
	var fnArgs []reflect.Value
	typ := fn.Type()
	numIn := typ.NumIn()
	if cmd.Options != nil {
		optsVal = reflect.ValueOf(cmd.Options)
		optsType := optsVal.Type()
		if numIn == 0 || (typ.In(0) != optsType && (numIn < 2 || typ.In(1) != optsType)) {
			panic(fmt.Errorf("command %s (%s) declares options of type %T but does not accept them", name, typ, cmd.Options))
		}
		flags, err := setupOptionsFlags(name, optsVal)
		if err != nil {
			panic(err)
		}
		if err := flags.Parse(rem); err != nil {
			return err
		}
		rem = flags.Args()
	}
	if numIn > 0 && typ.In(0) == argsType {
		fnArgs = append(fnArgs, reflect.ValueOf(rem))
	} else if len(rem) > 0 {
		fmt.Fprintf(os.Stderr, "command %s does not accept any arguments\n", name)
		return ErrUnusedArguments
	}
	if cmd.Options != nil {
		fnArgs = append(fnArgs, optsVal)
	}
	res := fn.Call(fnArgs)
	if len(res) > 0 {
		if err, ok := res[0].Interface().(error); ok {
			fmt.Fprintf(os.Stderr, "error running command %s: %s\n", name, err)
			return err
		}
	}
	return nil
}

func setupOptionsFlags(name string, val reflect.Value) (*flag.FlagSet, error) {
	if val.Kind() != reflect.Ptr {
		return nil, fmt.Errorf("invalid options %s, must be a pointer", val.Type())
	}
	val = reflect.Indirect(val)
	typ := val.Type()
	if typ.Kind() != reflect.Struct {
		return nil, fmt.Errorf("invalid Options type %s, must be an struct", typ)
	}
	flagsName := fmt.Sprintf("%s %s subcommand", filepath.Base(os.Args[0]), name)
	flags := flag.NewFlagSet(flagsName, flag.ContinueOnError)
	for ii := 0; ii < typ.NumField(); ii++ {
		field := typ.Field(ii)
		fieldVal := val.Field(ii)
		ptr := fieldVal.Addr().Interface()
		name := strings.ToLower(field.Name)
		var help string
		if n := field.Tag.Get("name"); n != "" {
			name = n
		}
		if h := field.Tag.Get("help"); h != "" {
			help = h
		}
		if name == "" {
			return nil, fmt.Errorf("no name provided for field #%d in type %s", ii, typ)
		}
		if value, ok := ptr.(flag.Value); ok {
			flags.Var(value, name, help)
			continue
		}
		switch field.Type.Kind() {
		case reflect.Bool:
			flags.BoolVar(ptr.(*bool), name, fieldVal.Bool(), help)
		case reflect.Float64:
			flags.Float64Var(ptr.(*float64), name, fieldVal.Float(), help)
		case reflect.Int:
			flags.IntVar(ptr.(*int), name, int(fieldVal.Int()), help)
		case reflect.Uint:
			flags.UintVar(ptr.(*uint), name, uint(fieldVal.Uint()), help)
		case reflect.Int64:
			flags.Int64Var(ptr.(*int64), name, fieldVal.Int(), help)
		case reflect.Uint64:
			flags.Uint64Var(ptr.(*uint64), name, fieldVal.Uint(), help)
		case reflect.String:
			flags.StringVar(ptr.(*string), name, fieldVal.String(), help)
		default:
			return nil, fmt.Errorf("field %s has invalid option type %s", field.Name, field.Type)
		}
	}
	return flags, nil
}

func commandByName(commands []*Cmd, name string) *Cmd {
	for _, v := range commands {
		if v.Name == name {
			return v
		}
	}
	return nil
}

func printCommandHelp(w io.Writer, cmd *Cmd) {
	fmt.Fprintf(w, "%s: %s\n", cmd.Name, cmd.Help)
	if cmd.Usage != "" {
		fmt.Fprintf(w, "\nusage: %s %s %s\n", filepath.Base(os.Args[0]), cmd.Name, cmd.Usage)
	}
	if cmd.LongHelp != "" {
		fmt.Fprintf(w, "\n%s\n\n", cmd.LongHelp)
	}
	if cmd.Options != nil {
		opts := reflect.ValueOf(cmd.Options)
		if fs, err := setupOptionsFlags(cmd.Name, opts); err == nil {
			fs.PrintDefaults()
		}
	}
}

func printHelp(w io.Writer, args []string, commands []*Cmd) error {
	var err error
	if len(args) == 0 {
		fmt.Fprintln(w, "missing command, available ones are:\n")
		err = ErrNoCommand
	} else {
		var unknown string
		if args[0] != "help" && commandByName(commands, args[0]) == nil {
			unknown = args[0]
		}
		if len(args) > 1 && args[0] == "help" {
			if cmd := commandByName(commands, args[1]); cmd != nil {
				printCommandHelp(w, cmd)
				return ErrHelp
			}
			unknown = args[1]
		}
		if unknown != "" {
			fmt.Fprintf(w, "unknown command %s, available ones are:\n\n", unknown)
			err = UnknownCommandError(unknown)
		}
	}
	tw := tabwriter.NewWriter(w, 0, 8, 0, '\t', 0)
	for _, v := range commands {
		fmt.Fprintf(tw, "%s\t%s\n", v.Name, v.Help)
	}
	fmt.Fprint(tw, "help\tPrint this help\n")
	tw.Flush()
	fmt.Fprint(w, "\nTo view additional help for each command use help <command_name>\n")
	if err == nil {
		err = ErrHelp
	}
	return err
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/md5"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/brotherlogic/goserver"
	"golang.org/x/net/context"
	"google.golang.org/grpc"

	pbd "github.com/brotherlogic/discovery/proto"
	pb "github.com/brotherlogic/gobuildslave/proto"
	pbs "github.com/brotherlogic/goserver/proto"
)

// Server the main server type
type Server struct {
	*goserver.GoServer
	runner *Runner
	disk   diskChecker
	jobs   map[string]*pb.JobDetails
}

func (s *Server) monitor(job *pb.JobDetails) {
	for true {
		log.Printf("MONITOR RUN %v", job)
		switch job.State {
		case pb.JobDetails_ACKNOWLEDGED:
			job.State = pb.JobDetails_BUILDING
			s.runner.Checkout(job.GetSpec().GetName())
			job.State = pb.JobDetails_BUILT
		case pb.JobDetails_BUILT:
			s.runner.Run(job.GetSpec())
			job.State = pb.JobDetails_RUNNING
		case pb.JobDetails_KILLING:
			s.runner.kill(job.GetSpec())
			if !isAlive(job.GetSpec()) {
				job.State = pb.JobDetails_DEAD
			}
		case pb.JobDetails_UPDATE_STARTING:
			s.runner.Update(job.GetSpec())
			job.State = pb.JobDetails_RUNNING
		case pb.JobDetails_RUNNING:
			time.Sleep(waitTime)
			if !isAlive(job.GetSpec()) {
				job.State = pb.JobDetails_DEAD
			}
		case pb.JobDetails_DEAD:
			time.Sleep(waitTime)
		}
	}
}

func getHash(file string) (string, error) {
	env := os.Environ()
	home := ""
	for _, s := range env {
		if strings.HasPrefix(s, "HOME=") {
			home = s[5:]
		}
	}

	if len(home) == 0 {
		log.Printf("Error in home: %v", home)
	}
	gpath := home + "/gobuild"

	f, err := os.Open(strings.Replace(file, "$GOPATH", gpath, 1))
	if err != nil {
		return "", err
	}
	defer f.Close()

	h := md5.New()
	if _, err := io.Copy(h, f); err != nil {
		return "", err
	}

	return string(h.Sum(nil)), nil
}

func getIP(name string, server string) (string, int) {
	conn, _ := grpc.Dial("192.168.86.64:50055", grpc.WithInsecure())
	defer conn.Close()

	registry := pbd.NewDiscoveryServiceClient(conn)
	entry := pbd.RegistryEntry{Name: name, Identifier: server}
	log.Printf("Searching for %v", entry)
	r, err := registry.Discover(context.Background(), &entry)

	if err != nil {
		log.Printf("Lookup failed for %v,%v -> %v", name, server, err)
		return "", -1
	}

	return r.Ip, int(r.Port)
}

// updateState of the runner command
func isAlive(spec *pb.JobSpec) bool {
	elems := strings.Split(spec.Name, "/")
	dServer, dPort := getIP(elems[len(elems)-1], spec.Server)

	log.Printf("Unable to find: %v -> %v", elems, dPort)
	if dPort > 0 {
		dConn, err := grpc.Dial(dServer+":"+strconv.Itoa(dPort), grpc.WithInsecure())
		if err != nil {
			return false
		}
		defer dConn.Close()

		c := pbs.NewGoserverServiceClient(dConn)
		_, err = c.IsAlive(context.Background(), &pbs.Alive{})
		log.Printf("UPDATED: %v,%v -> %v", dServer, dPort, err)
		return err == nil
	}

	//Mark as false if we can't locate the job
	return false
}

// DoRegister Registers this server
func (s Server) DoRegister(server *grpc.Server) {
	pb.RegisterGoBuildSlaveServer(server, &s)
}

// ReportHealth determines if the server is healthy
func (s Server) ReportHealth() bool {
	return true
}

// Mote promotes/demotes this server
func (s Server) Mote(master bool) error {
	return nil
}

//Init builds the default runner framework
func Init() *Runner {
	r := &Runner{gopath: "goautobuild"}
	r.runner = runCommand
	go r.run()
	return r
}

func runCommand(c *runnerCommand) {
	log.Printf("RUNNING COMMAND: %v", c)

	if c.command == nil {
		return
	}

	env := os.Environ()
	home := ""
	for _, s := range env {
		if strings.HasPrefix(s, "HOME=") {
			home = s[5:]
		}
	}

	if len(home) == 0 {
		log.Printf("Error in home: %v", home)
	}
	gpath := home + "/gobuild"
	c.command.Path = strings.Replace(c.command.Path, "$GOPATH", gpath, -1)
	for i := range c.command.Args {
		c.command.Args[i] = strings.Replace(c.command.Args[i], "$GOPATH", gpath, -1)
	}

	path := fmt.Sprintf("GOPATH=" + home + "/gobuild")
	found := false
	log.Printf("HERE = %v", c.command.Env)
	envl := os.Environ()
	for i, blah := range envl {
		if strings.HasPrefix(blah, "GOPATH") {
			envl[i] = path
			found = true
		}
	}
	if !found {
		envl = append(envl, path)
	}
	log.Printf("ENV = %v", envl)
	c.command.Env = envl

	out, err := c.command.StdoutPipe()
	out2, err2 := c.command.StderrPipe()
	if err != nil {
		log.Printf("Blah: %v", err)
	}

	if err2 != nil {
		log.Printf("Blah2: %v", err)
	}

	log.Printf("RUNNING %v, %v and %v", c.command.Path, c.command.Args, c.command.Env)
	c.command.Start()
	log.Printf("RUN STARTED:  %v", c.background)

	if !c.background {
		str := ""

		if out != nil {
			buf := new(bytes.Buffer)
			log.Printf("RUN READING 1:%v", out)
			buf.ReadFrom(out)
			str = buf.String()
			log.Printf("RUN IS HERE: %v", str)
		}

		if out2 != nil {
			buf2 := new(bytes.Buffer)
			log.Printf("RUN READING 2")
			buf2.ReadFrom(out2)
			str2 := buf2.String()
			log.Printf("RUN NOW %v and %v", str, str2)

		}
		log.Print("RUN HAS STARTING TO WAIT")
		c.command.Wait()
		log.Printf("RUN DONE WAITING")
		c.output = str
		c.complete = true
	}
	log.Printf("RUN IS DONE")
}

func (diskChecker prodDiskChecker) diskUsage(path string) int64 {
	return diskUsage(path)
}

func (s *Server) rebuildLoop() {
	for true {
		time.Sleep(time.Minute)

		var rebuildList []*pb.JobSpec
		var hashList []string
		for _, job := range s.runner.backgroundTasks {
			log.Printf("Job (started %v, now %v) %v", job.started, time.Now(), job)
			if time.Since(job.started) > time.Hour {
				log.Printf("Added to rebuild list (%v)", job)
				rebuildList = append(rebuildList, job.details.Spec)
				hashList = append(hashList, job.hash)
			}
		}

		log.Printf("Rebuilding %v", rebuildList)
		for i := range rebuildList {
			s.runner.Rebuild(rebuildList[i], hashList[i])
		}
	}
}

func main() {
	var quiet = flag.Bool("quiet", false, "Show all output")
	flag.Parse()

	if *quiet {
		log.SetFlags(0)
		log.SetOutput(ioutil.Discard)
	}

	s := Server{&goserver.GoServer{}, Init(), prodDiskChecker{}, make(map[string]*pb.JobDetails)}
	s.Register = s
	s.PrepServer()
	s.GoServer.Killme = false
	s.RegisterServingTask(s.rebuildLoop)
	s.RegisterServer("gobuildslave", false)
	s.Serve()
}
<commit_msg>Added some logging<commit_after>package main

import (
	"bytes"
	"crypto/md5"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/brotherlogic/goserver"
	"golang.org/x/net/context"
	"google.golang.org/grpc"

	pbd "github.com/brotherlogic/discovery/proto"
	pb "github.com/brotherlogic/gobuildslave/proto"
	pbs "github.com/brotherlogic/goserver/proto"
)

// Server the main server type
type Server struct {
	*goserver.GoServer
	runner *Runner
	disk   diskChecker
	jobs   map[string]*pb.JobDetails
}

func (s *Server) monitor(job *pb.JobDetails) {
	for true {
		log.Printf("MONITOR RUN %v", job)
		switch job.State {
		case pb.JobDetails_ACKNOWLEDGED:
			job.State = pb.JobDetails_BUILDING
			s.runner.Checkout(job.GetSpec().GetName())
			job.State = pb.JobDetails_BUILT
		case pb.JobDetails_BUILT:
			s.runner.Run(job.GetSpec())
			job.State = pb.JobDetails_RUNNING
		case pb.JobDetails_KILLING:
			s.runner.kill(job.GetSpec())
			if !isAlive(job.GetSpec()) {
				job.State = pb.JobDetails_DEAD
			}
		case pb.JobDetails_UPDATE_STARTING:
			s.runner.Update(job.GetSpec())
			job.State = pb.JobDetails_RUNNING
		case pb.JobDetails_RUNNING:
			time.Sleep(waitTime)
			if !isAlive(job.GetSpec()) {
				job.State = pb.JobDetails_DEAD
			}
		case pb.JobDetails_DEAD:
			job.State = pb.JobDetails_ACKNOWLEDGED
		}
	}
}

func getHash(file string) (string, error) {
	env := os.Environ()
	home := ""
	for _, s := range env {
		if strings.HasPrefix(s, "HOME=") {
			home = s[5:]
		}
	}

	if len(home) == 0 {
		log.Printf("Error in home: %v", home)
	}
	gpath := home + "/gobuild"

	f, err := os.Open(strings.Replace(file, "$GOPATH", gpath, 1))
	if err != nil {
		return "", err
	}
	defer f.Close()

	h := md5.New()
	if _, err := io.Copy(h, f); err != nil {
		return "", err
	}

	return string(h.Sum(nil)), nil
}

func getIP(name string, server string) (string, int) {
	conn, _ := grpc.Dial("192.168.86.64:50055", grpc.WithInsecure())
	defer conn.Close()

	registry := pbd.NewDiscoveryServiceClient(conn)
	entry := pbd.RegistryEntry{Name: name, Identifier: server}
	log.Printf("Searching for %v", entry)
	r, err := registry.Discover(context.Background(), &entry)

	if err != nil {
		log.Printf("Lookup failed for %v,%v -> %v", name, server, err)
		return "", -1
	}

	return r.Ip, int(r.Port)
}

// updateState of the runner command
func isAlive(spec *pb.JobSpec) bool {
	elems := strings.Split(spec.Name, "/")
	dServer, dPort := getIP(elems[len(elems)-1], spec.Server)

	log.Printf("Unable to find: %v -> %v", elems, dPort)
	if dPort > 0 {
		dConn, err := grpc.Dial(dServer+":"+strconv.Itoa(dPort), grpc.WithInsecure())
		if err != nil {
			return false
		}
		defer dConn.Close()

		c := pbs.NewGoserverServiceClient(dConn)
		_, err = c.IsAlive(context.Background(), &pbs.Alive{})
		log.Printf("UPDATED: %v,%v -> %v", dServer, dPort, err)
		return err == nil
	}

	//Mark as false if we can't locate the job
	return false
}

// DoRegister Registers this server
func (s Server) DoRegister(server *grpc.Server) {
	pb.RegisterGoBuildSlaveServer(server, &s)
}

// ReportHealth determines if the server is healthy
func (s Server) ReportHealth() bool {
	return true
}

// Mote promotes/demotes this server
func (s Server) Mote(master bool) error {
	return nil
}

//Init builds the default runner framework
func Init() *Runner {
	r := &Runner{gopath: "goautobuild"}
	r.runner = runCommand
	go r.run()
	return r
}

func runCommand(c *runnerCommand) {
	log.Printf("RUNNING COMMAND: %v", c)

	if c.command == nil {
		return
	}

	env := os.Environ()
	home := ""
	for _, s := range env {
		if strings.HasPrefix(s, "HOME=") {
			home = s[5:]
		}
	}

	if len(home) == 0 {
		log.Printf("Error in home: %v", home)
	}
	gpath := home + "/gobuild"
	c.command.Path = strings.Replace(c.command.Path, "$GOPATH", gpath, -1)
	for i := range c.command.Args {
		c.command.Args[i] = strings.Replace(c.command.Args[i], "$GOPATH", gpath, -1)
	}

	path := fmt.Sprintf("GOPATH=" + home + "/gobuild")
	found := false
	log.Printf("HERE = %v", c.command.Env)
	envl := os.Environ()
	for i, blah := range envl {
		if strings.HasPrefix(blah, "GOPATH") {
			envl[i] = path
			found = true
		}
	}
	if !found {
		envl = append(envl, path)
	}
	log.Printf("ENV = %v", envl)
	c.command.Env = envl

	out, err := c.command.StdoutPipe()
	out2, err2 := c.command.StderrPipe()
	if err != nil {
		log.Printf("Blah: %v", err)
	}

	if err2 != nil {
		log.Printf("Blah2: %v", err)
	}

	log.Printf("RUNNING %v, %v and %v", c.command.Path, c.command.Args, c.command.Env)
	c.command.Start()
	log.Printf("RUN STARTED:  %v", c.background)

	if !c.background {
		str := ""

		if out != nil {
			buf := new(bytes.Buffer)
			log.Printf("RUN READING 1:%v", out)
			buf.ReadFrom(out)
			str = buf.String()
			log.Printf("RUN IS HERE: %v", str)
		}

		if out2 != nil {
			buf2 := new(bytes.Buffer)
			log.Printf("RUN READING 2")
			buf2.ReadFrom(out2)
			str2 := buf2.String()
			log.Printf("RUN NOW %v and %v", str, str2)

		}
		log.Print("RUN HAS STARTING TO WAIT")
		c.command.Wait()
		log.Printf("RUN DONE WAITING")
		c.output = str
		c.complete = true
	}
	log.Printf("RUN IS DONE")
}

func (diskChecker prodDiskChecker) diskUsage(path string) int64 {
	return diskUsage(path)
}

func (s *Server) rebuildLoop() {
	for true {
		time.Sleep(time.Minute)

		var rebuildList []*pb.JobSpec
		var hashList []string
		for _, job := range s.runner.backgroundTasks {
			log.Printf("Job (started %v, now %v) %v", job.started, time.Now(), job)
			if time.Since(job.started) > time.Hour {
				log.Printf("Added to rebuild list (%v)", job)
				rebuildList = append(rebuildList, job.details.Spec)
				hashList = append(hashList, job.hash)
			}
		}

		log.Printf("Rebuilding %v", rebuildList)
		for i := range rebuildList {
			s.runner.Rebuild(rebuildList[i], hashList[i])
		}
	}
}

func main() {
	var quiet = flag.Bool("quiet", false, "Show all output")
	flag.Parse()

	if *quiet {
		log.SetFlags(0)
		log.SetOutput(ioutil.Discard)
	}

	s := Server{&goserver.GoServer{}, Init(), prodDiskChecker{}, make(map[string]*pb.JobDetails)}
	s.Register = s
	s.PrepServer()
	s.GoServer.Killme = false
	s.RegisterServingTask(s.rebuildLoop)
	s.RegisterServer("gobuildslave", false)
	s.Serve()
}
<|endoftext|>


<commit_before>package initials

import (
	"bytes"
	"context"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"
	"unicode"
	"unicode/utf8"

	"github.com/cozy/cozy-stack/pkg/config/config"
	"github.com/cozy/cozy-stack/pkg/logger"
)

// Image returns an image with the initials for the given name (and the
// content-type to use for the HTTP response).
// TODO add cache
func Image(publicName string) ([]byte, string, error) {
	name := strings.TrimSpace(publicName)
	info := extractInfo(name)
	bytes, err := draw(info)
	if err != nil {
		return nil, "", err
	}
	return bytes, "image/png", nil
}

// See https://github.com/cozy/cozy-ui/blob/master/react/Avatar/index.jsx#L9-L26
// and https://docs.cozy.io/cozy-ui/styleguide/section-settings.html#kssref-settings-colors
var colors = []string{
	"#1FA8F1",
	"#FD7461",
	"#FC6D00",
	"#F52D2D",
	"#FF962F",
	"#FF7F1B",
	"#6984CE",
	"#7F6BEE",
	"#B449E7",
	"#40DE8E",
	"#0DCBCF",
	"#35CE68",
	"#3DA67E",
	"#C2ADF4",
	"#FFC644",
	"#FC4C83",
}

type info struct {
	initials string
	color    string
}

func extractInfo(name string) info {
	initials := getInitials(name)
	color := getColor(name)
	return info{initials: initials, color: color}
}

func getInitials(name string) string {
	parts := strings.Split(name, " ")
	initials := make([]rune, 0, len(parts))
	for _, part := range parts {
		r, size := utf8.DecodeRuneInString(part)
		if size > 0 && unicode.IsLetter(r) {
			initials = append(initials, r)
		}
	}
	switch len(initials) {
	case 0:
		return "?"
	case 1:
		return string(initials)
	default:
		return string(initials[0]) + string(initials[len(initials)-1])
	}
}

func getColor(name string) string {
	sum := 0
	for i := 0; i < len(name); i++ {
		sum += int(name[i])
	}
	return colors[sum%len(colors)]
}

func draw(info info) ([]byte, error) {
	var env []string
	{
		tempDir, err := ioutil.TempDir("", "magick")
		if err == nil {
			defer os.RemoveAll(tempDir)
			envTempDir := fmt.Sprintf("MAGICK_TEMPORARY_PATH=%s", tempDir)
			env = []string{envTempDir}
		}
	}

	convertCmd := config.GetConfig().Jobs.ImageMagickConvertCmd
	if convertCmd == "" {
		convertCmd = "convert"
	}

	// convert -size 128x128 null: -fill blue -draw 'circle 64,64 0,64' -fill white -font Lato-Regular
	// -pointsize 64 -gravity center -annotate "+0,+0" "AM" foo.png
	args := []string{
		"-limit", "Memory", "1GB",
		"-limit", "Map", "1GB",
		// Use a transparent background
		"-size", "128x128",
		"null:",
		// Add a cicle of color
		"-fill", info.color,
		"-draw", "circle 64,64 0,64",
		// Add the initials
		"-fill", "white",
		"-font", "Lato-Regular",
		"-pointsize", "64",
		"-gravity", "center",
		"-annotate", "+0,+0",
		info.initials,
		// Use the colorspace recommended for web, sRGB
		"-colorspace", "sRGB",
		// Send the output on stdout, in PNG format
		"png:-",
	}

	var stdout, stderr bytes.Buffer
	cmd := exec.CommandContext(context.Background(), convertCmd, args...)
	cmd.Env = env
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	if err := cmd.Run(); err != nil {
		logger.WithNamespace("initials").
			WithField("stderr", stderr.String()).
			WithField("initials", info.initials).
			WithField("color", info.color).
			Errorf("imagemagick failed: %s", err)
		return nil, err
	}
	return stdout.Bytes(), nil
}
<commit_msg>Add cache for the initials images<commit_after>package initials

import (
	"bytes"
	"context"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"
	"time"
	"unicode"
	"unicode/utf8"

	"github.com/cozy/cozy-stack/pkg/config/config"
	"github.com/cozy/cozy-stack/pkg/logger"
)

const (
	cacheTTL    = 30 * 24 * time.Hour // 1 month
	contentType = "image/png"
)

// Image returns an image with the initials for the given name (and the
// content-type to use for the HTTP response).
func Image(publicName string) ([]byte, string, error) {
	name := strings.TrimSpace(publicName)
	info := extractInfo(name)
	cache := config.GetConfig().CacheStorage
	key := "initials:" + info.initials + info.color

	if bytes, ok := cache.Get(key); ok {
		return bytes, contentType, nil
	}

	bytes, err := draw(info)
	if err != nil {
		return nil, "", err
	}
	cache.Set(key, bytes, cacheTTL)
	return bytes, contentType, nil
}

// See https://github.com/cozy/cozy-ui/blob/master/react/Avatar/index.jsx#L9-L26
// and https://docs.cozy.io/cozy-ui/styleguide/section-settings.html#kssref-settings-colors
var colors = []string{
	"#1FA8F1",
	"#FD7461",
	"#FC6D00",
	"#F52D2D",
	"#FF962F",
	"#FF7F1B",
	"#6984CE",
	"#7F6BEE",
	"#B449E7",
	"#40DE8E",
	"#0DCBCF",
	"#35CE68",
	"#3DA67E",
	"#C2ADF4",
	"#FFC644",
	"#FC4C83",
}

type info struct {
	initials string
	color    string
}

func extractInfo(name string) info {
	initials := getInitials(name)
	color := getColor(name)
	return info{initials: initials, color: color}
}

func getInitials(name string) string {
	parts := strings.Split(name, " ")
	initials := make([]rune, 0, len(parts))
	for _, part := range parts {
		r, size := utf8.DecodeRuneInString(part)
		if size > 0 && unicode.IsLetter(r) {
			initials = append(initials, r)
		}
	}
	switch len(initials) {
	case 0:
		return "?"
	case 1:
		return string(initials)
	default:
		return string(initials[0]) + string(initials[len(initials)-1])
	}
}

func getColor(name string) string {
	sum := 0
	for i := 0; i < len(name); i++ {
		sum += int(name[i])
	}
	return colors[sum%len(colors)]
}

func draw(info info) ([]byte, error) {
	var env []string
	{
		tempDir, err := ioutil.TempDir("", "magick")
		if err == nil {
			defer os.RemoveAll(tempDir)
			envTempDir := fmt.Sprintf("MAGICK_TEMPORARY_PATH=%s", tempDir)
			env = []string{envTempDir}
		}
	}

	convertCmd := config.GetConfig().Jobs.ImageMagickConvertCmd
	if convertCmd == "" {
		convertCmd = "convert"
	}

	// convert -size 128x128 null: -fill blue -draw 'circle 64,64 0,64' -fill white -font Lato-Regular
	// -pointsize 64 -gravity center -annotate "+0,+0" "AM" foo.png
	args := []string{
		"-limit", "Memory", "1GB",
		"-limit", "Map", "1GB",
		// Use a transparent background
		"-size", "128x128",
		"null:",
		// Add a cicle of color
		"-fill", info.color,
		"-draw", "circle 64,64 0,64",
		// Add the initials
		"-fill", "white",
		"-font", "Lato-Regular",
		"-pointsize", "64",
		"-gravity", "center",
		"-annotate", "+0,+0",
		info.initials,
		// Use the colorspace recommended for web, sRGB
		"-colorspace", "sRGB",
		// Send the output on stdout, in PNG format
		"png:-",
	}

	var stdout, stderr bytes.Buffer
	cmd := exec.CommandContext(context.Background(), convertCmd, args...)
	cmd.Env = env
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	if err := cmd.Run(); err != nil {
		logger.WithNamespace("initials").
			WithField("stderr", stderr.String()).
			WithField("initials", info.initials).
			WithField("color", info.color).
			Errorf("imagemagick failed: %s", err)
		return nil, err
	}
	return stdout.Bytes(), nil
}
<|endoftext|>


<commit_before>package ioutils

import (
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
)

// AtomicFileWriterOptions specifies options for creating the atomic file writer.
type AtomicFileWriterOptions struct {
	// NoSync specifies whether the sync call must be skipped for the file.
	// If NoSync is not specified, the file is synced to the
	// storage after it has been written and before it is moved to
	// the specified path.
	NoSync bool
}

var defaultWriterOptions AtomicFileWriterOptions = AtomicFileWriterOptions{}

// SetDefaultOptions overrides the default options used when creating an
// atomic file writer.
func SetDefaultOptions(opts AtomicFileWriterOptions) {
	defaultWriterOptions = opts
}

// NewAtomicFileWriterWithOpts returns WriteCloser so that writing to it writes to a
// temporary file and closing it atomically changes the temporary file to
// destination path. Writing and closing concurrently is not allowed.
func NewAtomicFileWriterWithOpts(filename string, perm os.FileMode, opts *AtomicFileWriterOptions) (io.WriteCloser, error) {
	f, err := ioutil.TempFile(filepath.Dir(filename), ".tmp-"+filepath.Base(filename))
	if err != nil {
		return nil, err
	}
	if opts == nil {
		opts = &defaultWriterOptions
	}
	abspath, err := filepath.Abs(filename)
	if err != nil {
		return nil, err
	}
	return &atomicFileWriter{
		f:      f,
		fn:     abspath,
		perm:   perm,
		noSync: opts.NoSync,
	}, nil
}

// NewAtomicFileWriter returns WriteCloser so that writing to it writes to a
// temporary file and closing it atomically changes the temporary file to
// destination path. Writing and closing concurrently is not allowed.
func NewAtomicFileWriter(filename string, perm os.FileMode) (io.WriteCloser, error) {
	return NewAtomicFileWriterWithOpts(filename, perm, nil)
}

// AtomicWriteFile atomically writes data to a file named by filename.
func AtomicWriteFile(filename string, data []byte, perm os.FileMode) error {
	f, err := NewAtomicFileWriter(filename, perm)
	if err != nil {
		return err
	}
	n, err := f.Write(data)
	if err == nil && n < len(data) {
		err = io.ErrShortWrite
		f.(*atomicFileWriter).writeErr = err
	}
	if err1 := f.Close(); err == nil {
		err = err1
	}
	return err
}

type atomicFileWriter struct {
	f        *os.File
	fn       string
	writeErr error
	perm     os.FileMode
	noSync   bool
}

func (w *atomicFileWriter) Write(dt []byte) (int, error) {
	n, err := w.f.Write(dt)
	if err != nil {
		w.writeErr = err
	}
	return n, err
}

func (w *atomicFileWriter) Close() (retErr error) {
	defer func() {
		if retErr != nil || w.writeErr != nil {
			os.Remove(w.f.Name())
		}
	}()
	if !w.noSync {
		if err := fdatasync(w.f); err != nil {
			w.f.Close()
			return err
		}
	}
	if err := w.f.Close(); err != nil {
		return err
	}
	if err := os.Chmod(w.f.Name(), w.perm); err != nil {
		return err
	}
	if w.writeErr == nil {
		return os.Rename(w.f.Name(), w.fn)
	}
	return nil
}

// AtomicWriteSet is used to atomically write a set
// of files and ensure they are visible at the same time.
// Must be committed to a new directory.
type AtomicWriteSet struct {
	root string
}

// NewAtomicWriteSet creates a new atomic write set to
// atomically create a set of files. The given directory
// is used as the base directory for storing files before
// commit. If no temporary directory is given the system
// default is used.
func NewAtomicWriteSet(tmpDir string) (*AtomicWriteSet, error) {
	td, err := ioutil.TempDir(tmpDir, "write-set-")
	if err != nil {
		return nil, err
	}

	return &AtomicWriteSet{
		root: td,
	}, nil
}

// WriteFile writes a file to the set, guaranteeing the file
// has been synced.
func (ws *AtomicWriteSet) WriteFile(filename string, data []byte, perm os.FileMode) error {
	f, err := ws.FileWriter(filename, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, perm)
	if err != nil {
		return err
	}
	n, err := f.Write(data)
	if err == nil && n < len(data) {
		err = io.ErrShortWrite
	}
	if err1 := f.Close(); err == nil {
		err = err1
	}
	return err
}

type syncFileCloser struct {
	*os.File
}

func (w syncFileCloser) Close() error {
	err := fdatasync(w.File)
	if err1 := w.File.Close(); err == nil {
		err = err1
	}
	return err
}

// FileWriter opens a file writer inside the set. The file
// should be synced and closed before calling commit.
func (ws *AtomicWriteSet) FileWriter(name string, flag int, perm os.FileMode) (io.WriteCloser, error) {
	f, err := os.OpenFile(filepath.Join(ws.root, name), flag, perm)
	if err != nil {
		return nil, err
	}
	return syncFileCloser{f}, nil
}

// Cancel cancels the set and removes all temporary data
// created in the set.
func (ws *AtomicWriteSet) Cancel() error {
	return os.RemoveAll(ws.root)
}

// Commit moves all created files to the target directory. The
// target directory must not exist and the parent of the target
// directory must exist.
func (ws *AtomicWriteSet) Commit(target string) error {
	return os.Rename(ws.root, target)
}

// String returns the location the set is writing to.
func (ws *AtomicWriteSet) String() string {
	return ws.root
}
<commit_msg>fswriters: honor nosync option<commit_after>package ioutils

import (
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
)

// AtomicFileWriterOptions specifies options for creating the atomic file writer.
type AtomicFileWriterOptions struct {
	// NoSync specifies whether the sync call must be skipped for the file.
	// If NoSync is not specified, the file is synced to the
	// storage after it has been written and before it is moved to
	// the specified path.
	NoSync bool
}

var defaultWriterOptions AtomicFileWriterOptions = AtomicFileWriterOptions{}

// SetDefaultOptions overrides the default options used when creating an
// atomic file writer.
func SetDefaultOptions(opts AtomicFileWriterOptions) {
	defaultWriterOptions = opts
}

// NewAtomicFileWriterWithOpts returns WriteCloser so that writing to it writes to a
// temporary file and closing it atomically changes the temporary file to
// destination path. Writing and closing concurrently is not allowed.
func NewAtomicFileWriterWithOpts(filename string, perm os.FileMode, opts *AtomicFileWriterOptions) (io.WriteCloser, error) {
	f, err := ioutil.TempFile(filepath.Dir(filename), ".tmp-"+filepath.Base(filename))
	if err != nil {
		return nil, err
	}
	if opts == nil {
		opts = &defaultWriterOptions
	}
	abspath, err := filepath.Abs(filename)
	if err != nil {
		return nil, err
	}
	return &atomicFileWriter{
		f:      f,
		fn:     abspath,
		perm:   perm,
		noSync: opts.NoSync,
	}, nil
}

// NewAtomicFileWriter returns WriteCloser so that writing to it writes to a
// temporary file and closing it atomically changes the temporary file to
// destination path. Writing and closing concurrently is not allowed.
func NewAtomicFileWriter(filename string, perm os.FileMode) (io.WriteCloser, error) {
	return NewAtomicFileWriterWithOpts(filename, perm, nil)
}

// AtomicWriteFile atomically writes data to a file named by filename.
func AtomicWriteFile(filename string, data []byte, perm os.FileMode) error {
	f, err := NewAtomicFileWriter(filename, perm)
	if err != nil {
		return err
	}
	n, err := f.Write(data)
	if err == nil && n < len(data) {
		err = io.ErrShortWrite
		f.(*atomicFileWriter).writeErr = err
	}
	if err1 := f.Close(); err == nil {
		err = err1
	}
	return err
}

type atomicFileWriter struct {
	f        *os.File
	fn       string
	writeErr error
	perm     os.FileMode
	noSync   bool
}

func (w *atomicFileWriter) Write(dt []byte) (int, error) {
	n, err := w.f.Write(dt)
	if err != nil {
		w.writeErr = err
	}
	return n, err
}

func (w *atomicFileWriter) Close() (retErr error) {
	defer func() {
		if retErr != nil || w.writeErr != nil {
			os.Remove(w.f.Name())
		}
	}()
	if !w.noSync {
		if err := fdatasync(w.f); err != nil {
			w.f.Close()
			return err
		}
	}
	if err := w.f.Close(); err != nil {
		return err
	}
	if err := os.Chmod(w.f.Name(), w.perm); err != nil {
		return err
	}
	if w.writeErr == nil {
		return os.Rename(w.f.Name(), w.fn)
	}
	return nil
}

// AtomicWriteSet is used to atomically write a set
// of files and ensure they are visible at the same time.
// Must be committed to a new directory.
type AtomicWriteSet struct {
	root string
}

// NewAtomicWriteSet creates a new atomic write set to
// atomically create a set of files. The given directory
// is used as the base directory for storing files before
// commit. If no temporary directory is given the system
// default is used.
func NewAtomicWriteSet(tmpDir string) (*AtomicWriteSet, error) {
	td, err := ioutil.TempDir(tmpDir, "write-set-")
	if err != nil {
		return nil, err
	}

	return &AtomicWriteSet{
		root: td,
	}, nil
}

// WriteFile writes a file to the set, guaranteeing the file
// has been synced.
func (ws *AtomicWriteSet) WriteFile(filename string, data []byte, perm os.FileMode) error {
	f, err := ws.FileWriter(filename, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, perm)
	if err != nil {
		return err
	}
	n, err := f.Write(data)
	if err == nil && n < len(data) {
		err = io.ErrShortWrite
	}
	if err1 := f.Close(); err == nil {
		err = err1
	}
	return err
}

type syncFileCloser struct {
	*os.File
}

func (w syncFileCloser) Close() error {
	if !defaultWriterOptions.NoSync {
		return w.File.Close()
	}
	err := fdatasync(w.File)
	if err1 := w.File.Close(); err == nil {
		err = err1
	}
	return err
}

// FileWriter opens a file writer inside the set. The file
// should be synced and closed before calling commit.
func (ws *AtomicWriteSet) FileWriter(name string, flag int, perm os.FileMode) (io.WriteCloser, error) {
	f, err := os.OpenFile(filepath.Join(ws.root, name), flag, perm)
	if err != nil {
		return nil, err
	}
	return syncFileCloser{f}, nil
}

// Cancel cancels the set and removes all temporary data
// created in the set.
func (ws *AtomicWriteSet) Cancel() error {
	return os.RemoveAll(ws.root)
}

// Commit moves all created files to the target directory. The
// target directory must not exist and the parent of the target
// directory must exist.
func (ws *AtomicWriteSet) Commit(target string) error {
	return os.Rename(ws.root, target)
}

// String returns the location the set is writing to.
func (ws *AtomicWriteSet) String() string {
	return ws.root
}
<|endoftext|>