<commit_before>package uploadedfile

import (
	"errors"
	"fmt"
	"math"
	"os"
	"regexp"
	"strconv"

	"github.com/Imgur/mandible/imageprocessor/processorcommand"
	"github.com/Imgur/mandible/imageprocessor/thumbType"
)

var (
	defaultQuality   = 83
	maxImageSideSize = 10000
)

type ThumbFile struct {
	localPath string

	Name          string
	Width         int
	MaxWidth      int
	Height        int
	MaxHeight     int
	Shape         string
	CropGravity   string
	CropWidth     int
	CropHeight    int
	CropRatio     string
	Quality       int
	Format        string
	StoreURI      string
	DesiredFormat string
}

func NewThumbFile(width, maxWidth, height, maxHeight int, name, shape, path, cropGravity string, cropWidth, cropHeight int, cropRatio string, quality int, desiredFormat string) *ThumbFile {
	if quality == 0 {
		quality = defaultQuality
	}

	return &ThumbFile{
		localPath: path,

		Name:          name,
		Width:         width,
		MaxWidth:      maxWidth,
		Height:        height,
		MaxHeight:     maxHeight,
		Shape:         shape,
		CropGravity:   cropGravity,
		CropWidth:     cropWidth,
		CropHeight:    cropHeight,
		CropRatio:     cropRatio,
		Quality:       quality,
		Format:        "",
		StoreURI:      "",
		DesiredFormat: desiredFormat,
	}
}

func (this *ThumbFile) SetPath(path string) error {
	if _, err := os.Stat(path); os.IsNotExist(err) {
		return errors.New(fmt.Sprintf("Error when creating thumbnail %s", this.Name))
	}

	this.localPath = path

	return nil
}

func (this *ThumbFile) GetPath() string {
	return this.localPath
}

func (this *ThumbFile) GetOutputFormat(original *UploadedFile) thumbType.ThumbType {
	if this.DesiredFormat != "" {
		return thumbType.FromString(this.DesiredFormat)
	}

	return thumbType.FromMime(original.GetMime())
}

func (this *ThumbFile) ComputeWidth(original *UploadedFile) int {
	width := this.Width

	oWidth, _, err := original.Dimensions()
	if err != nil {
		return 0
	}

	if this.MaxWidth > 0 {
		width = int(math.Min(float64(oWidth), float64(this.MaxWidth)))
	}

	return width
}

func (this *ThumbFile) ComputeHeight(original *UploadedFile) int {
	height := this.Height

	_, oHeight, err := original.Dimensions()
	if err != nil {
		return 0
	}

	if this.MaxHeight > 0 {
		height = int(math.Min(float64(oHeight), float64(this.MaxHeight)))
	}

	return height
}

func (this *ThumbFile) ComputeCrop(original *UploadedFile) (int, int, error) {
	re := regexp.MustCompile("(.*):(.*)")
	matches := re.FindStringSubmatch(this.CropRatio)
	if len(matches) != 3 {
		return 0, 0, errors.New("Invalid crop_ratio")
	}

	wRatio, werr := strconv.ParseFloat(matches[1], 64)
	hRatio, herr := strconv.ParseFloat(matches[2], 64)
	if werr != nil || herr != nil {
		return 0, 0, errors.New("Invalid crop_ratio")
	}

	var cropWidth, cropHeight float64

	if wRatio >= hRatio {
		wRatio = wRatio / hRatio
		hRatio = 1
		cropWidth = math.Ceil(float64(this.ComputeHeight(original)) * wRatio)
		cropHeight = math.Ceil(float64(this.ComputeHeight(original)) * hRatio)
	} else {
		hRatio = hRatio / wRatio
		wRatio = 1
		cropWidth = math.Ceil(float64(this.ComputeWidth(original)) * wRatio)
		cropHeight = math.Ceil(float64(this.ComputeWidth(original)) * hRatio)
	}

	return int(cropWidth), int(cropHeight), nil
}

func (this *ThumbFile) Process(original *UploadedFile) error {
	switch this.Shape {
	case "circle":
		return this.processCircle(original)
	case "thumb":
		return this.processThumb(original)
	case "square":
		return this.processSquare(original)
	case "custom":
		return this.processCustom(original)
	default:
		return this.processFull(original)
	}
}

func (this *ThumbFile) String() string {
	return fmt.Sprintf("Thumbnail of <%s>", this.Name)
}

func (this *ThumbFile) processSquare(original *UploadedFile) error {
	if this.Width == 0 {
		return errors.New("Width cannot be 0")
	}
	if this.Width > maxImageSideSize {
		return errors.New("Width too large")
	}

	filename, err := processorcommand.SquareThumb(original.GetPath(), this.Name, this.Width, this.GetOutputFormat(original))
	if err != nil {
		return err
	}

	if err := this.SetPath(filename); err != nil {
		return err
	}

	return nil
}

func (this *ThumbFile) processCircle(original *UploadedFile) error {
	if this.Width == 0 {
		return errors.New("Width cannot be 0")
	}
	if this.Width > maxImageSideSize {
		return errors.New("Width too large")
	}

	filename, err := processorcommand.CircleThumb(original.GetPath(), this.Name, this.Width, this.GetOutputFormat(original))
	if err != nil {
		return err
	}

	if err := this.SetPath(filename); err != nil {
		return err
	}

	return nil
}

func (this *ThumbFile) processThumb(original *UploadedFile) error {
	if this.Width == 0 {
		return errors.New("Width cannot be 0")
	}
	if this.Width > maxImageSideSize {
		return errors.New("Width too large")
	}
	if this.Height == 0 {
		return errors.New("Height cannot be 0")
	}
	if this.Height > maxImageSideSize {
		return errors.New("Height too large")
	}

	filename, err := processorcommand.Thumb(original.GetPath(), this.Name, this.Width, this.Height, this.GetOutputFormat(original))
	if err != nil {
		return err
	}

	if err := this.SetPath(filename); err != nil {
		return err
	}

	return nil
}

func (this *ThumbFile) processCustom(original *UploadedFile) error {
	cropWidth := this.CropWidth
	cropHeight := this.CropHeight
	var err error

	if this.CropRatio != "" {
		cropWidth, cropHeight, err = this.ComputeCrop(original)
		if err != nil {
			return err
		}
	}

	width := this.ComputeWidth(original)
	height := this.ComputeHeight(original)
	if (width == 0 || width > maxImageSideSize) && this.CropRatio == "" {
		return errors.New("Invalid width")
	}
	if (height == 0 || height > maxImageSideSize) && this.CropRatio == "" {
		return errors.New("Invalid height")
	}

	filename, err := processorcommand.CustomThumb(original.GetPath(), this.Name, width, height, this.CropGravity, cropWidth, cropHeight, this.Quality, this.GetOutputFormat(original))
	if err != nil {
		return err
	}

	if err := this.SetPath(filename); err != nil {
		return err
	}

	return nil
}

func (this *ThumbFile) processFull(original *UploadedFile) error {
	filename, err := processorcommand.Full(original.GetPath(), this.Name, this.GetOutputFormat(original))
	if err != nil {
		return err
	}

	if err := this.SetPath(filename); err != nil {
		return err
	}

	return nil
}
<commit_msg>Allow for custom thumbnails to only need a width or height and not both<commit_after>package uploadedfile

import (
	"errors"
	"fmt"
	"math"
	"os"
	"regexp"
	"strconv"

	"github.com/Imgur/mandible/imageprocessor/processorcommand"
	"github.com/Imgur/mandible/imageprocessor/thumbType"
)

var (
	defaultQuality   = 83
	maxImageSideSize = 10000
)

type ThumbFile struct {
	localPath string

	Name          string
	Width         int
	MaxWidth      int
	Height        int
	MaxHeight     int
	Shape         string
	CropGravity   string
	CropWidth     int
	CropHeight    int
	CropRatio     string
	Quality       int
	Format        string
	StoreURI      string
	DesiredFormat string
}

func NewThumbFile(width, maxWidth, height, maxHeight int, name, shape, path, cropGravity string, cropWidth, cropHeight int, cropRatio string, quality int, desiredFormat string) *ThumbFile {
	if quality == 0 {
		quality = defaultQuality
	}

	return &ThumbFile{
		localPath: path,

		Name:          name,
		Width:         width,
		MaxWidth:      maxWidth,
		Height:        height,
		MaxHeight:     maxHeight,
		Shape:         shape,
		CropGravity:   cropGravity,
		CropWidth:     cropWidth,
		CropHeight:    cropHeight,
		CropRatio:     cropRatio,
		Quality:       quality,
		Format:        "",
		StoreURI:      "",
		DesiredFormat: desiredFormat,
	}
}

func (this *ThumbFile) SetPath(path string) error {
	if _, err := os.Stat(path); os.IsNotExist(err) {
		return errors.New(fmt.Sprintf("Error when creating thumbnail %s", this.Name))
	}

	this.localPath = path

	return nil
}

func (this *ThumbFile) GetPath() string {
	return this.localPath
}

func (this *ThumbFile) GetOutputFormat(original *UploadedFile) thumbType.ThumbType {
	if this.DesiredFormat != "" {
		return thumbType.FromString(this.DesiredFormat)
	}

	return thumbType.FromMime(original.GetMime())
}

func (this *ThumbFile) ComputeWidth(original *UploadedFile) int {
	width := this.Width

	oWidth, _, err := original.Dimensions()
	if err != nil {
		return 0
	}

	if this.MaxWidth > 0 {
		width = int(math.Min(float64(oWidth), float64(this.MaxWidth)))
	}

	return width
}

func (this *ThumbFile) ComputeHeight(original *UploadedFile) int {
	height := this.Height

	_, oHeight, err := original.Dimensions()
	if err != nil {
		return 0
	}

	if this.MaxHeight > 0 {
		height = int(math.Min(float64(oHeight), float64(this.MaxHeight)))
	}

	return height
}

func (this *ThumbFile) ComputeCrop(original *UploadedFile) (int, int, error) {
	re := regexp.MustCompile("(.*):(.*)")
	matches := re.FindStringSubmatch(this.CropRatio)
	if len(matches) != 3 {
		return 0, 0, errors.New("Invalid crop_ratio")
	}

	wRatio, werr := strconv.ParseFloat(matches[1], 64)
	hRatio, herr := strconv.ParseFloat(matches[2], 64)
	if werr != nil || herr != nil {
		return 0, 0, errors.New("Invalid crop_ratio")
	}

	var cropWidth, cropHeight float64

	if wRatio >= hRatio {
		wRatio = wRatio / hRatio
		hRatio = 1
		cropWidth = math.Ceil(float64(this.ComputeHeight(original)) * wRatio)
		cropHeight = math.Ceil(float64(this.ComputeHeight(original)) * hRatio)
	} else {
		hRatio = hRatio / wRatio
		wRatio = 1
		cropWidth = math.Ceil(float64(this.ComputeWidth(original)) * wRatio)
		cropHeight = math.Ceil(float64(this.ComputeWidth(original)) * hRatio)
	}

	return int(cropWidth), int(cropHeight), nil
}

func (this *ThumbFile) Process(original *UploadedFile) error {
	switch this.Shape {
	case "circle":
		return this.processCircle(original)
	case "thumb":
		return this.processThumb(original)
	case "square":
		return this.processSquare(original)
	case "custom":
		return this.processCustom(original)
	default:
		return this.processFull(original)
	}
}

func (this *ThumbFile) String() string {
	return fmt.Sprintf("Thumbnail of <%s>", this.Name)
}

func (this *ThumbFile) processSquare(original *UploadedFile) error {
	if this.Width == 0 {
		return errors.New("Width cannot be 0")
	}
	if this.Width > maxImageSideSize {
		return errors.New("Width too large")
	}

	filename, err := processorcommand.SquareThumb(original.GetPath(), this.Name, this.Width, this.GetOutputFormat(original))
	if err != nil {
		return err
	}

	if err := this.SetPath(filename); err != nil {
		return err
	}

	return nil
}

func (this *ThumbFile) processCircle(original *UploadedFile) error {
	if this.Width == 0 {
		return errors.New("Width cannot be 0")
	}
	if this.Width > maxImageSideSize {
		return errors.New("Width too large")
	}

	filename, err := processorcommand.CircleThumb(original.GetPath(), this.Name, this.Width, this.GetOutputFormat(original))
	if err != nil {
		return err
	}

	if err := this.SetPath(filename); err != nil {
		return err
	}

	return nil
}

func (this *ThumbFile) processThumb(original *UploadedFile) error {
	if this.Width == 0 {
		return errors.New("Width cannot be 0")
	}
	if this.Width > maxImageSideSize {
		return errors.New("Width too large")
	}
	if this.Height == 0 {
		return errors.New("Height cannot be 0")
	}
	if this.Height > maxImageSideSize {
		return errors.New("Height too large")
	}

	filename, err := processorcommand.Thumb(original.GetPath(), this.Name, this.Width, this.Height, this.GetOutputFormat(original))
	if err != nil {
		return err
	}

	if err := this.SetPath(filename); err != nil {
		return err
	}

	return nil
}

func (this *ThumbFile) processCustom(original *UploadedFile) error {
	cropWidth := this.CropWidth
	cropHeight := this.CropHeight
	var err error

	if this.CropRatio != "" {
		cropWidth, cropHeight, err = this.ComputeCrop(original)
		if err != nil {
			return err
		}
	}

	width := this.ComputeWidth(original)
	height := this.ComputeHeight(original)
	validWidth := width > 0 && width <= maxImageSideSize
	validHeight := height > 0 && height <= maxImageSideSize

	if !validWidth && !validHeight {
		if !validWidth {
			return errors.New("Invalid width")
		}

		return errors.New("Invalid height")
	}

	filename, err := processorcommand.CustomThumb(original.GetPath(), this.Name, width, height, this.CropGravity, cropWidth, cropHeight, this.Quality, this.GetOutputFormat(original))
	if err != nil {
		return err
	}

	if err := this.SetPath(filename); err != nil {
		return err
	}

	return nil
}

func (this *ThumbFile) processFull(original *UploadedFile) error {
	filename, err := processorcommand.Full(original.GetPath(), this.Name, this.GetOutputFormat(original))
	if err != nil {
		return err
	}

	if err := this.SetPath(filename); err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>/*
 Copyright 2010-2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.

 This file is licensed under the Apache License, Version 2.0 (the "License").
 You may not use this file except in compliance with the License. A copy of the
 License is located at

 http://aws.amazon.com/apache2.0/

 This file 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 main

import (
	"encoding/json"
	"fmt"
	"os"
	"strconv"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/lambda"
)

type getItemsRequest struct {
	SortBy     string
	SortOrder  string
	ItemsToGet int
}

type getItemsResponseError struct {
	Message string `json:"message"`
}

type getItemsResponseData struct {
	Item string `json:"item"`
}

type getItemsResponseBody struct {
	Result string                 `json:"result"`
	Data   []getItemsResponseData `json:"data"`
	Error  getItemsResponseError  `json:"error"`
}

type getItemsResponseHeaders struct {
	ContentType string `json:"Content-Type"`
}

type getItemsResponse struct {
	StatusCode int                     `json:"statusCode"`
	Headers    getItemsResponseHeaders `json:"headers"`
	Body       getItemsResponseBody    `json:"body"`
}

func main() {
	// Create Lambda service client
	sess := session.Must(session.NewSessionWithOptions(session.Options{
		SharedConfigState: session.SharedConfigEnable,
	}))

    client := lambda.New(sess, &aws.Config{Region: aws.String("us-west-2")})

    // Get the 10 most recent items
    request := getItemsRequest{"time", "descending", 10}

    payload, err := json.Marshal(request)

	if err != nil {
		fmt.Println("Error marshalling MyGetItemsFunction request")
		os.Exit(0)
	}

	result, err := client.Invoke(&lambda.InvokeInput{FunctionName: aws.String("MyGetItemsFunction"), Payload: payload})

	if err != nil {
		fmt.Println("Error calling MyGetItemsFunction")
		os.Exit(0)
	}

	var resp getItemsResponse

	err = json.Unmarshal(result.Payload, &resp)

	if err != nil {
		fmt.Println("Error unmarshalling MyGetItemsFunction response")
		os.Exit(0)
	}

	// If the status code is NOT 200, the call failed
	if resp.StatusCode != 200 {
		fmt.Println("Error getting items, StatusCode: " + strconv.Itoa(resp.StatusCode))
		os.Exit(0)
	}

	// If the result is failure, we got an error
	if resp.Body.Result == "failure" {
		fmt.Println("Failed to get items")
		os.Exit(0)
	}

	// Print out items
	if len(resp.Body.Data) > 0 {
		for i := range resp.Body.Data {
			fmt.Println(resp.Body.Data[i].Item)
		}
	} else {
		fmt.Println("There were no items")
	}
}
<commit_msg>Untabified Go Lambda example<commit_after>/*
 Copyright 2010-2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.

 This file is licensed under the Apache License, Version 2.0 (the "License").
 You may not use this file except in compliance with the License. A copy of the
 License is located at

 http://aws.amazon.com/apache2.0/

 This file 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 main

import (
    "encoding/json"
    "fmt"
    "os"
    "strconv"

    "github.com/aws/aws-sdk-go/aws"
    "github.com/aws/aws-sdk-go/aws/session"
    "github.com/aws/aws-sdk-go/service/lambda"
)

type getItemsRequest struct {
    SortBy     string
    SortOrder  string
    ItemsToGet int
}

type getItemsResponseError struct {
    Message string `json:"message"`
}

type getItemsResponseData struct {
    Item string `json:"item"`
}

type getItemsResponseBody struct {
    Result string                 `json:"result"`
    Data   []getItemsResponseData `json:"data"`
    Error  getItemsResponseError  `json:"error"`
}

type getItemsResponseHeaders struct {
    ContentType string `json:"Content-Type"`
}

type getItemsResponse struct {
    StatusCode int                     `json:"statusCode"`
    Headers    getItemsResponseHeaders `json:"headers"`
    Body       getItemsResponseBody    `json:"body"`
}

func main() {
    // Create Lambda service client
    sess := session.Must(session.NewSessionWithOptions(session.Options{
        SharedConfigState: session.SharedConfigEnable,
    }))

    client := lambda.New(sess, &aws.Config{Region: aws.String("us-west-2")})

    // Get the 10 most recent items
    request := getItemsRequest{"time", "descending", 10}

    payload, err := json.Marshal(request)

    if err != nil {
        fmt.Println("Error marshalling MyGetItemsFunction request")
        os.Exit(0)
    }

    result, err := client.Invoke(&lambda.InvokeInput{FunctionName: aws.String("MyGetItemsFunction"), Payload: payload})

    if err != nil {
        fmt.Println("Error calling MyGetItemsFunction")
        os.Exit(0)
    }

    var resp getItemsResponse

    err = json.Unmarshal(result.Payload, &resp)

    if err != nil {
        fmt.Println("Error unmarshalling MyGetItemsFunction response")
        os.Exit(0)
    }

    // If the status code is NOT 200, the call failed
    if resp.StatusCode != 200 {
        fmt.Println("Error getting items, StatusCode: " + strconv.Itoa(resp.StatusCode))
        os.Exit(0)
    }

    // If the result is failure, we got an error
    if resp.Body.Result == "failure" {
        fmt.Println("Failed to get items")
        os.Exit(0)
    }

    // Print out items
    if len(resp.Body.Data) > 0 {
        for i := range resp.Body.Data {
            fmt.Println(resp.Body.Data[i].Item)
        }
    } else {
        fmt.Println("There were no items")
    }
}
<|endoftext|>


<commit_before>package catalog

import (
	"fmt"
	"log"
	"sync"

	"github.com/dnaeon/gru/graph"
	"github.com/dnaeon/gru/resource"
	"github.com/dnaeon/gru/utils"
	"github.com/layeh/gopher-luar"
	"github.com/yuin/gopher-lua"
)

// Catalog type contains a collection of resources
type Catalog struct {
	// Unsorted contains the list of resources created by Lua
	Unsorted []resource.Resource `luar:"-"`

	// Collection contains the unsorted resources as a collection
	collection resource.Collection `luar:"-"`

	// Sorted contains the resources after a topological sort.
	sorted []*graph.Node `luar:"-"`

	// Reversed contains the resource dependency graph in reverse
	// order. It is used for finding the reverse dependencies of
	// resources.
	reversed *graph.Graph `luar:"-"`

	// Status contains status information about resources
	status *status `luar:"-"`

	// Configuration settings
	config *Config `luar:"-"`
}

// Config type represents a set of settings to use when
// creating and processing the catalog
type Config struct {
	// Name of the Lua module to load and execute
	Module string

	// Do not take any actions, just report what would be done
	DryRun bool

	// Writer used to log events
	Logger *log.Logger

	// Path to the site repo containing module and data files
	SiteRepo string

	// The Lua state
	L *lua.LState

	// Number of goroutines to use for concurrent processing
	Concurrency int
}

// status type contains status information about processed resources
type status struct {
	sync.RWMutex

	// Items contain the result of resource processing and any
	// errors that might have occurred during processing.
	// Keys of the map are the resource ids and their
	// values are the errors returned by resources.
	items map[string]error
}

// set sets the status for a resource
func (s *status) set(id string, err error) {
	s.Lock()
	defer s.Unlock()
	s.items[id] = err
}

// get retrieves the status of a resource
func (s *status) get(id string) (error, bool) {
	s.Lock()
	defer s.Unlock()
	err, ok := s.items[id]

	return err, ok
}

// New creates a new empty catalog with the provided configuration
func New(config *Config) *Catalog {
	c := &Catalog{
		config:     config,
		collection: make(resource.Collection),
		sorted:     make([]*graph.Node, 0),
		reversed:   graph.New(),
		status: &status{
			items: make(map[string]error),
		},
		Unsorted: make([]resource.Resource, 0),
	}

	// Inject the configuration for resources
	resource.DefaultConfig = &resource.Config{
		Logger:   config.Logger,
		SiteRepo: config.SiteRepo,
	}

	// Register the catalog type in Lua and also register
	// metamethods for the catalog, so that we can use
	// the catalog in a more Lua-friendly way
	mt := luar.MT(config.L, c)
	mt.RawSetString("__len", luar.New(config.L, (*Catalog).Len))
	config.L.SetGlobal("catalog", luar.New(config.L, c))

	return c
}

// Add adds a resource to the catalog.
// This method is called from Lua when adding new resources
func (c *Catalog) Add(resources ...resource.Resource) {
	for _, r := range resources {
		if r != nil {
			c.Unsorted = append(c.Unsorted, r)
		}
	}
}

// Len returns the number of unsorted resources in catalog
func (c *Catalog) Len() int {
	return len(c.Unsorted)
}

// Load loads resources into the catalog
func (c *Catalog) Load() error {
	// Register the resource providers and catalog in Lua
	resource.LuaRegisterBuiltin(c.config.L)
	if err := c.config.L.DoFile(c.config.Module); err != nil {
		return err
	}

	// Perform a topological sort of the resources
	collection, err := resource.CreateCollection(c.Unsorted)
	if err != nil {
		return err
	}

	collectionGraph, err := collection.DependencyGraph()
	if err != nil {
		return err
	}

	reversed, err := collection.ReversedGraph()
	if err != nil {
		return err
	}

	sorted, err := collectionGraph.Sort()
	if err != nil {
		return err
	}

	// Set catalog fields
	c.collection = collection
	c.sorted = sorted
	c.reversed = reversed

	c.config.Logger.Printf("Loaded %d resources\n", len(c.sorted))

	return nil
}

// Run processes the resources from catalog
func (c *Catalog) Run() error {
	// process executes a single resource
	process := func(r resource.Resource) {
		id := r.ID()
		err := c.execute(r)
		c.status.set(id, err)
		if err != nil {
			c.config.Logger.Printf("%s %s\n", id, err)
		}
	}

	// Start goroutines for concurrent processing
	var wg sync.WaitGroup
	ch := make(chan resource.Resource, 1024)
	c.config.Logger.Printf("Starting %d goroutines for concurrent processing\n", c.config.Concurrency)
	for i := 0; i < c.config.Concurrency; i++ {
		wg.Add(1)
		worker := func() {
			defer wg.Done()
			for r := range ch {
				c.config.Logger.Printf("%s is concurrent", r.ID())
				process(r)
			}
		}
		go worker()
	}

	// Process the resources
	for _, node := range c.sorted {
		r := c.collection[node.Name]
		switch {
		// Resource is concurrent and has no dependencies
		case r.IsConcurrent() && len(r.Dependencies()) == 0:
			ch <- r
			continue
		// Resource is concurrent and have no reverse dependencies
		case r.IsConcurrent() && len(c.reversed.Nodes[r.ID()].Edges) == 0:
			ch <- r
			continue
		// Resource is not concurrent
		default:
			process(r)
		}
	}

	close(ch)
	wg.Wait()

	return nil
}

// execute processes a single resource
func (c *Catalog) execute(r resource.Resource) error {
	// Check if the resource has failed dependencies
	for _, dep := range r.Dependencies() {
		if err, _ := c.status.get(dep); err != nil {
			return fmt.Errorf("failed dependency for %s", dep)
		}
	}

	if err := r.Validate(); err != nil {
		return err
	}

	state, err := r.Evaluate()
	if err != nil {
		return err
	}

	if c.config.DryRun {
		return nil
	}

	// Current and wanted states for the resource
	want := utils.NewString(state.Want)
	current := utils.NewString(state.Current)

	// The list of present and absent states for the resource
	present := utils.NewList(r.GetPresentStates()...)
	absent := utils.NewList(r.GetAbsentStates()...)

	var action func() error
	switch {
	case want.IsInList(present) && current.IsInList(absent):
		action = r.Create
	case want.IsInList(absent) && current.IsInList(present):
		action = r.Delete
	case state.Outdated:
		action = r.Update
	}

	if action != nil {
		return action()
	}

	return nil
}
<commit_msg>catalog: use Graph.Reversed() to get the reversed graph<commit_after>package catalog

import (
	"fmt"
	"log"
	"sync"

	"github.com/dnaeon/gru/graph"
	"github.com/dnaeon/gru/resource"
	"github.com/dnaeon/gru/utils"
	"github.com/layeh/gopher-luar"
	"github.com/yuin/gopher-lua"
)

// Catalog type contains a collection of resources
type Catalog struct {
	// Unsorted contains the list of resources created by Lua
	Unsorted []resource.Resource `luar:"-"`

	// Collection contains the unsorted resources as a collection
	collection resource.Collection `luar:"-"`

	// Sorted contains the resources after a topological sort.
	sorted []*graph.Node `luar:"-"`

	// Reversed contains the resource dependency graph in reverse
	// order. It is used for finding the reverse dependencies of
	// resources.
	reversed *graph.Graph `luar:"-"`

	// Status contains status information about resources
	status *status `luar:"-"`

	// Configuration settings
	config *Config `luar:"-"`
}

// Config type represents a set of settings to use when
// creating and processing the catalog
type Config struct {
	// Name of the Lua module to load and execute
	Module string

	// Do not take any actions, just report what would be done
	DryRun bool

	// Writer used to log events
	Logger *log.Logger

	// Path to the site repo containing module and data files
	SiteRepo string

	// The Lua state
	L *lua.LState

	// Number of goroutines to use for concurrent processing
	Concurrency int
}

// status type contains status information about processed resources
type status struct {
	sync.RWMutex

	// Items contain the result of resource processing and any
	// errors that might have occurred during processing.
	// Keys of the map are the resource ids and their
	// values are the errors returned by resources.
	items map[string]error
}

// set sets the status for a resource
func (s *status) set(id string, err error) {
	s.Lock()
	defer s.Unlock()
	s.items[id] = err
}

// get retrieves the status of a resource
func (s *status) get(id string) (error, bool) {
	s.Lock()
	defer s.Unlock()
	err, ok := s.items[id]

	return err, ok
}

// New creates a new empty catalog with the provided configuration
func New(config *Config) *Catalog {
	c := &Catalog{
		config:     config,
		collection: make(resource.Collection),
		sorted:     make([]*graph.Node, 0),
		reversed:   graph.New(),
		status: &status{
			items: make(map[string]error),
		},
		Unsorted: make([]resource.Resource, 0),
	}

	// Inject the configuration for resources
	resource.DefaultConfig = &resource.Config{
		Logger:   config.Logger,
		SiteRepo: config.SiteRepo,
	}

	// Register the catalog type in Lua and also register
	// metamethods for the catalog, so that we can use
	// the catalog in a more Lua-friendly way
	mt := luar.MT(config.L, c)
	mt.RawSetString("__len", luar.New(config.L, (*Catalog).Len))
	config.L.SetGlobal("catalog", luar.New(config.L, c))

	return c
}

// Add adds a resource to the catalog.
// This method is called from Lua when adding new resources
func (c *Catalog) Add(resources ...resource.Resource) {
	for _, r := range resources {
		if r != nil {
			c.Unsorted = append(c.Unsorted, r)
		}
	}
}

// Len returns the number of unsorted resources in catalog
func (c *Catalog) Len() int {
	return len(c.Unsorted)
}

// Load loads resources into the catalog
func (c *Catalog) Load() error {
	// Register the resource providers and catalog in Lua
	resource.LuaRegisterBuiltin(c.config.L)
	if err := c.config.L.DoFile(c.config.Module); err != nil {
		return err
	}

	// Perform a topological sort of the resources
	collection, err := resource.CreateCollection(c.Unsorted)
	if err != nil {
		return err
	}

	collectionGraph, err := collection.DependencyGraph()
	if err != nil {
		return err
	}

	reversed := collectionGraph.Reversed()

	sorted, err := collectionGraph.Sort()
	if err != nil {
		return err
	}

	// Set catalog fields
	c.collection = collection
	c.sorted = sorted
	c.reversed = reversed

	c.config.Logger.Printf("Loaded %d resources\n", len(c.sorted))

	return nil
}

// Run processes the resources from catalog
func (c *Catalog) Run() error {
	// process executes a single resource
	process := func(r resource.Resource) {
		id := r.ID()
		err := c.execute(r)
		c.status.set(id, err)
		if err != nil {
			c.config.Logger.Printf("%s %s\n", id, err)
		}
	}

	// Start goroutines for concurrent processing
	var wg sync.WaitGroup
	ch := make(chan resource.Resource, 1024)
	c.config.Logger.Printf("Starting %d goroutines for concurrent processing\n", c.config.Concurrency)
	for i := 0; i < c.config.Concurrency; i++ {
		wg.Add(1)
		worker := func() {
			defer wg.Done()
			for r := range ch {
				c.config.Logger.Printf("%s is concurrent", r.ID())
				process(r)
			}
		}
		go worker()
	}

	// Process the resources
	for _, node := range c.sorted {
		r := c.collection[node.Name]
		switch {
		// Resource is concurrent and has no dependencies
		case r.IsConcurrent() && len(r.Dependencies()) == 0:
			ch <- r
			continue
		// Resource is concurrent and have no reverse dependencies
		case r.IsConcurrent() && len(c.reversed.Nodes[r.ID()].Edges) == 0:
			ch <- r
			continue
		// Resource is not concurrent
		default:
			process(r)
		}
	}

	close(ch)
	wg.Wait()

	return nil
}

// execute processes a single resource
func (c *Catalog) execute(r resource.Resource) error {
	// Check if the resource has failed dependencies
	for _, dep := range r.Dependencies() {
		if err, _ := c.status.get(dep); err != nil {
			return fmt.Errorf("failed dependency for %s", dep)
		}
	}

	if err := r.Validate(); err != nil {
		return err
	}

	state, err := r.Evaluate()
	if err != nil {
		return err
	}

	if c.config.DryRun {
		return nil
	}

	// Current and wanted states for the resource
	want := utils.NewString(state.Want)
	current := utils.NewString(state.Current)

	// The list of present and absent states for the resource
	present := utils.NewList(r.GetPresentStates()...)
	absent := utils.NewList(r.GetAbsentStates()...)

	var action func() error
	switch {
	case want.IsInList(present) && current.IsInList(absent):
		action = r.Create
	case want.IsInList(absent) && current.IsInList(present):
		action = r.Delete
	case state.Outdated:
		action = r.Update
	}

	if action != nil {
		return action()
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"text/template"
	"time"

	"github.com/couchbaselabs/cbfs/config"
	"sync"
)

type StorageNode struct {
	Addr      string
	AddrRaw   string    `json:"addr_raw"`
	Started   time.Time `json:"starttime"`
	HBTime    time.Time `json:"hbtime"`
	BindAddr  string
	FrameBind string
	HBAgeStr  string `json:"hbage_str"`
	Used      int64
	Free      int64
	Size      int64
	UptimeStr string `json:"uptime_str"`
}

type Tasks map[string]map[string]struct {
	State string
	TS    time.Time
}

type Backup struct {
	Filename string
	OID      string
	When     time.Time
	Conf     cbfsconfig.CBFSConfig
}

type Nodes map[string]StorageNode

var infoFlags = flag.NewFlagSet("info", flag.ExitOnError)
var infoTemplate = infoFlags.String("t", "", "Display template")
var infoTemplateFile = infoFlags.String("T", "", "Display template filename")

const defaultInfoTemplate = `nodes:
{{ range $name, $nodeinfo := .Nodes }}  {{$name}} up {{$nodeinfo.HBAgeStr}}
{{ end }}
{{if .Tasks}}tasks:{{end}}{{ range $node, $tasks := .Tasks }}
  {{$node}}
  {{ range $task, $info := $tasks }}    {{$task}} - {{$info.State}} - {{$info.TS}}
  {{end}}{{end}}
backups:
  Found {{len .Backups.Previous }} backups.
  Current:  {{.Backups.Latest.Filename}} ({{.Backups.Latest.OID}})
            as of {{.Backups.Latest.When}}

config:
{{ range $k, $v := .Conf.ToMap}}  {{$k}}: {{$v}}
{{end}}
`

func getJsonData(u string, into interface{}) error {
	res, err := http.Get(u)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	if res.StatusCode != 200 {
		return fmt.Errorf("HTTP Error: %v", res.Status)
	}

	d := json.NewDecoder(res.Body)
	return d.Decode(into)
}

func infoCommand(base string, args []string) {
	infoFlags.Parse(args)

	tmplstr := *infoTemplate
	if tmplstr == "" {
		if *infoTemplateFile == "" {
			tmplstr = defaultInfoTemplate
		} else {
			td, err := ioutil.ReadFile(*infoTemplateFile)
			if err != nil {
				log.Fatalf("Error reading template file: %v", err)
			}
			tmplstr = string(td)
		}
	}

	tmpl, err := template.New("").Parse(tmplstr)
	if err != nil {
		log.Fatalf("Error parsing template: %v", err)
	}

	u, err := url.Parse(base)
	if err != nil {
		log.Fatalf("Error parsing URL: %v", err)
	}

	result := struct {
		Nodes   Nodes
		Tasks   Tasks
		Backups struct {
			Previous []Backup `json:"backups"`
			Latest   Backup
		}
		Conf cbfsconfig.CBFSConfig
	}{}

	todo := map[string]interface{}{
		"/.cbfs/nodes/":  &result.Nodes,
		"/.cbfs/tasks/":  &result.Tasks,
		"/.cbfs/backup/": &result.Backups,
		"/.cbfs/config/": &result.Conf,
	}

	wg := sync.WaitGroup{}

	for k, v := range todo {
		u.Path = k
		wg.Add(1)
		go func(s string, to interface{}) {
			defer wg.Done()
			err = getJsonData(s, to)
			if err != nil {
				log.Fatalf("Error getting node info: %v", err)
			}
		}(u.String(), v)
	}

	wg.Wait()

	err = tmpl.Execute(os.Stdout, result)
	if err != nil {
		log.Fatalf("Error executing template: %v", err)
	}
}
<commit_msg>Fix uptime display in the default template.<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"text/template"
	"time"

	"github.com/couchbaselabs/cbfs/config"
	"sync"
)

type StorageNode struct {
	Addr      string
	AddrRaw   string    `json:"addr_raw"`
	Started   time.Time `json:"starttime"`
	HBTime    time.Time `json:"hbtime"`
	BindAddr  string
	FrameBind string
	HBAgeStr  string `json:"hbage_str"`
	Used      int64
	Free      int64
	Size      int64
	UptimeStr string `json:"uptime_str"`
}

type Tasks map[string]map[string]struct {
	State string
	TS    time.Time
}

type Backup struct {
	Filename string
	OID      string
	When     time.Time
	Conf     cbfsconfig.CBFSConfig
}

type Nodes map[string]StorageNode

var infoFlags = flag.NewFlagSet("info", flag.ExitOnError)
var infoTemplate = infoFlags.String("t", "", "Display template")
var infoTemplateFile = infoFlags.String("T", "", "Display template filename")

const defaultInfoTemplate = `nodes:
{{ range $name, $nodeinfo := .Nodes }}  {{$name}} up {{$nodeinfo.UptimeStr}} (age: {{$nodeinfo.HBAgeStr}})
{{ end }}
{{if .Tasks}}tasks:{{end}}{{ range $node, $tasks := .Tasks }}
  {{$node}}
  {{ range $task, $info := $tasks }}    {{$task}} - {{$info.State}} - {{$info.TS}}
  {{end}}{{end}}
backups:
  Found {{len .Backups.Previous }} backups.
  Current:  {{.Backups.Latest.Filename}} ({{.Backups.Latest.OID}})
            as of {{.Backups.Latest.When}}

config:
{{ range $k, $v := .Conf.ToMap}}  {{$k}}: {{$v}}
{{end}}
`

func getJsonData(u string, into interface{}) error {
	res, err := http.Get(u)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	if res.StatusCode != 200 {
		return fmt.Errorf("HTTP Error: %v", res.Status)
	}

	d := json.NewDecoder(res.Body)
	return d.Decode(into)
}

func infoCommand(base string, args []string) {
	infoFlags.Parse(args)

	tmplstr := *infoTemplate
	if tmplstr == "" {
		if *infoTemplateFile == "" {
			tmplstr = defaultInfoTemplate
		} else {
			td, err := ioutil.ReadFile(*infoTemplateFile)
			if err != nil {
				log.Fatalf("Error reading template file: %v", err)
			}
			tmplstr = string(td)
		}
	}

	tmpl, err := template.New("").Parse(tmplstr)
	if err != nil {
		log.Fatalf("Error parsing template: %v", err)
	}

	u, err := url.Parse(base)
	if err != nil {
		log.Fatalf("Error parsing URL: %v", err)
	}

	result := struct {
		Nodes   Nodes
		Tasks   Tasks
		Backups struct {
			Previous []Backup `json:"backups"`
			Latest   Backup
		}
		Conf cbfsconfig.CBFSConfig
	}{}

	todo := map[string]interface{}{
		"/.cbfs/nodes/":  &result.Nodes,
		"/.cbfs/tasks/":  &result.Tasks,
		"/.cbfs/backup/": &result.Backups,
		"/.cbfs/config/": &result.Conf,
	}

	wg := sync.WaitGroup{}

	for k, v := range todo {
		u.Path = k
		wg.Add(1)
		go func(s string, to interface{}) {
			defer wg.Done()
			err = getJsonData(s, to)
			if err != nil {
				log.Fatalf("Error getting node info: %v", err)
			}
		}(u.String(), v)
	}

	wg.Wait()

	err = tmpl.Execute(os.Stdout, result)
	if err != nil {
		log.Fatalf("Error executing template: %v", err)
	}
}
<|endoftext|>


<commit_before>package config

import (
  "bitbucket.org/classroomsystems/ini"
  "github.com/phzfi/RIC/server/logging"
  "strconv"
  "errors"
)

type Conf struct{
  conf ini.Config
}

type DefaultConf struct{
   minHeight int
   minWidth int
   maxHeight int
   maxWidth int
   addMark bool
   imgpath string
   tokens int
   vertical float64
   horizontal float64
}

func ReadConfig(path string) (config Conf, err error) {
  conf, err := ini.LoadFile(path)
  if err != nil {
    logging.Debug("Error reading config " + err.Error())
    return
  }
  config = Conf {
    conf: conf,
  }
  return
}

func (conf Conf) GetString(section, key string) (value string, err error) {
  value, success := conf.conf.Get(section, key)
  if (success) {
    return
  } else {
    return "", errors.New("Value not found for "+ key +" in "+ section)
  }
}

func (conf Conf) GetInt(section, key string) (value int, err error) {
  str, success := conf.conf.Get(section, key)
  if success {
    return strconv.Atoi(str)
  } else {
    return 0, errors.New("Value not found for "+ key +" in "+ section)
  }
}

func (conf Conf) GetUint64(section, key string) (value uint64, err error) {
  str, success := conf.conf.Get(section, key)
  if success {
    return strconv.ParseUint(str, 10, 64)
  } else {
    return 0, errors.New("Value not found for "+ key +" in "+ section)
  }
}

func (conf Conf) GetFloat64(section, key string) (value float64, err error) {
  str, success := conf.conf.Get(section, key)
  if (success) {
    return strconv.ParseFloat(str, 64)
  } else {
    return 0.0, errors.New("Value not found for "+ key +" in "+ section)
  }
}

func (conf Conf) GetBool(section, key string) (value bool, err error) {
  str, success := conf.conf.Get(section, key)
  if (success) {
    return strconv.ParseBool(str)
  } else {
    return false, errors.New("Value not found for "+ key +" in "+ section)
  }
}
<commit_msg>refactoring stuff alot<commit_after>package config

import (
  "bitbucket.org/classroomsystems/ini"
  "github.com/phzfi/RIC/server/logging"
  "strconv"
  "errors"
)

type Conf struct{
  conf ini.Config
}

type ConfValues struct{
   MinHeight int
   MinWidth int
   MaxHeight int
   MaxWidth int
   AddMark bool
   Imgpath string
   Tokens int
   Vertical float64
   Horizontal float64
   Mem uint64
}

func ReadConfig(path string) (config Conf, err error) {
  conf, err := ini.LoadFile(path)
  if err != nil {
    logging.Debug("Error reading config " + err.Error())
    return
  }
  config = Conf {
    conf: conf,
  }
  return
}

func (conf Conf) GetString(section, key string) (value string, err error) {
  value, success := conf.conf.Get(section, key)
  if (success) {
    return
  } else {
    return "", errors.New("Value not found for "+ key +" in "+ section)
  }
}

func (conf Conf) GetInt(section, key string) (value int, err error) {
  str, success := conf.conf.Get(section, key)
  if success {
    return strconv.Atoi(str)
  } else {
    return 0, errors.New("Value not found for "+ key +" in "+ section)
  }
}

func (conf Conf) GetUint64(section, key string) (value uint64, err error) {
  str, success := conf.conf.Get(section, key)
  if success {
    return strconv.ParseUint(str, 10, 64)
  } else {
    return 0, errors.New("Value not found for "+ key +" in "+ section)
  }
}

func (conf Conf) GetFloat64(section, key string) (value float64, err error) {
  str, success := conf.conf.Get(section, key)
  if (success) {
    return strconv.ParseFloat(str, 64)
  } else {
    return 0.0, errors.New("Value not found for "+ key +" in "+ section)
  }
}

func (conf Conf) GetBool(section, key string) (value bool, err error) {
  str, success := conf.conf.Get(section, key)
  if (success) {
    return strconv.ParseBool(str)
  } else {
    return false, errors.New("Value not found for "+ key +" in "+ section)
  }
}
<|endoftext|>


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

package main

import (
	"bytes"
	"fmt"
	"math/rand"
	"strconv"
	"sync/atomic"
	"time"

	"github.com/cockroachdb/cockroach/client"
	"github.com/cockroachdb/cockroach/proto"
	"github.com/cockroachdb/cockroach/security"
	"github.com/cockroachdb/cockroach/security/securitytest"
	"github.com/cockroachdb/cockroach/server"
	"github.com/cockroachdb/cockroach/storage"
	"github.com/cockroachdb/cockroach/testutils"
	"github.com/cockroachdb/cockroach/util/log"
)

// makes an id string from an id int.
func makeAccountID(id int) []byte {
	return []byte(fmt.Sprintf("%09d", id))
}

// Bank stores all the bank related state.
type Bank struct {
	kvClient        *client.KV
	numAccounts     int
	numTransactions int32
}

// moveMoney() moves an amount between two accounts if the amount
// is available in the from account. Returns true on success.
func (bank *Bank) moveMoney(from, to []byte, amount int64) bool {
	// Early exit when from == to
	if bytes.Compare(from, to) == 0 {
		return true
	}
	txnOpts := &client.TransactionOptions{Name: fmt.Sprintf("Transferring %s-%s-%d", from, to, amount)}
	err := bank.kvClient.RunTransaction(txnOpts, func(txn *client.Txn) error {
		fromGet := client.Get(proto.Key(from))
		fromResp := fromGet.Reply.(*proto.GetResponse)
		toGet := client.Get(proto.Key(to))
		toResp := toGet.Reply.(*proto.GetResponse)
		if err := txn.Run(fromGet, toGet); err != nil {
			return err
		}
		// Read from value.
		var fromValue int64
		if fromResp.Value != nil && fromResp.Value.Bytes != nil {
			readValue, err := strconv.ParseInt(string(fromResp.Value.Bytes), 10, 64)
			if err != nil {
				return err
			}
			fromValue = readValue
		}
		// Ensure there is enough cash.
		if fromValue < amount {
			return nil
		}
		// Read to value.
		var toValue int64
		if toResp.Value != nil && toResp.Value.Bytes != nil {
			readValue, err := strconv.ParseInt(string(toResp.Value.Bytes), 10, 64)
			if err != nil {
				return err
			}
			toValue = readValue
		}
		// Update both accounts.
		txn.Prepare(client.Put(proto.Key(from), []byte(fmt.Sprintf("%d", fromValue-amount))))
		txn.Prepare(client.Put(proto.Key(to), []byte(fmt.Sprintf("%d", toValue+amount))))
		return nil
	})
	if err != nil {
		log.Fatal(err)
	}
	atomic.AddInt32(&bank.numTransactions, 1)
	return true
}

// Read the balances in all the accounts and return them.
func (bank *Bank) readAllAccounts() []int64 {
	balances := make([]int64, bank.numAccounts)
	calls := make([]client.Call, bank.numAccounts)
	txnOpts := &client.TransactionOptions{Name: "Reading all balances"}
	err := bank.kvClient.RunTransaction(txnOpts, func(txn *client.Txn) error {
		for i := 0; i < bank.numAccounts; i++ {
			calls[i] = client.Get(proto.Key(makeAccountID(i)))
		}
		if err := txn.Run(calls...); err != nil {
			log.Fatal(err)
		}
		// Copy responses into balances
		for i := 0; i < bank.numAccounts; i++ {
			gr := calls[i].Reply.(*proto.GetResponse)
			if gr.Value != nil && gr.Value.Bytes != nil {
				balance, err := strconv.ParseInt(string(gr.Value.Bytes), 10, 64)
				if err != nil {
					log.Fatal(err)
				}
				balances[i] = balance
			}
		}
		return nil
	})
	if err != nil {
		log.Fatal(err)
	}
	return balances
}

// continuouslyTransferMoney() keeps moving random amounts between
// random accounts.
func (bank *Bank) continuousMoneyTransfer() {
	for {
		from := makeAccountID(rand.Intn(bank.numAccounts))
		to := makeAccountID(rand.Intn(bank.numAccounts))
		exchange := rand.Int63n(100)
		bank.moveMoney(from, to, exchange)
	}
}

// Initialize all the bank accounts with cash
func (bank *Bank) initBankAccounts(cash int64) {
	calls := make([]client.Call, bank.numAccounts)
	for i := 0; i < bank.numAccounts; i++ {
		calls[i] = client.Put(proto.Key(makeAccountID(i)), []byte(fmt.Sprintf("%d", cash)))
	}
	if err := bank.kvClient.Run(calls...); err != nil {
		log.Fatal(err)
	}
}

func main() {
	fmt.Printf("A simple program that keeps moving money between bank accounts\n\n")
	// Run a test cockroach instance to represent the bank
	security.SetReadFileFn(securitytest.Asset)
	serv := server.StartTestServer(nil)
	defer serv.Stop()
	// Initialize the bank
	var bank Bank
	bank.numAccounts = 1000
	// Key Value Client initialization.
	sender, err := client.NewHTTPSender(serv.ServingAddr(), testutils.NewTestBaseContext())
	if err != nil {
		log.Fatal(err)
	}
	bank.kvClient = client.NewKV(nil, sender)
	bank.kvClient.User = storage.UserRoot
	// Initialize all the bank accounts
	const initCash = 1000
	bank.initBankAccounts(initCash)

	// Start all the money transfer routines
	const numTransferRoutines = 10
	for i := 0; i < numTransferRoutines; i++ {
		go bank.continuousMoneyTransfer()
	}
	// Sleep for a bit to allow money transfers to happen in the background and then exit
	time.Sleep(10 * time.Second)

	fmt.Printf("%d transactions were executed\n\n", bank.numTransactions)

	// Check that all the money is accounted for
	balances := bank.readAllAccounts()
	var totalAmount int64
	for i := 0; i < bank.numAccounts; i++ {
		fmt.Printf("Account %d contains %d$\n", i, balances[i])
		totalAmount += balances[i]
	}
	if totalAmount != int64(bank.numAccounts)*initCash {
		err := fmt.Sprintf("\nTotal cash in the bank = %d\n", totalAmount)
		log.Fatal(err)
	}
	fmt.Printf("\nThe bank is in good order\n\n")
}
<commit_msg>Changed a few comments<commit_after>// Copyright 2015 The Cockroach Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License. See the AUTHORS file
// for names of contributors.
//
// Author: Vivek Menezes (vivek.menezes@gmail.com)

package main

import (
	"bytes"
	"fmt"
	"math/rand"
	"strconv"
	"sync/atomic"
	"time"

	"github.com/cockroachdb/cockroach/client"
	"github.com/cockroachdb/cockroach/proto"
	"github.com/cockroachdb/cockroach/security"
	"github.com/cockroachdb/cockroach/security/securitytest"
	"github.com/cockroachdb/cockroach/server"
	"github.com/cockroachdb/cockroach/storage"
	"github.com/cockroachdb/cockroach/testutils"
	"github.com/cockroachdb/cockroach/util/log"
)

// Makes an id string from an id int.
func makeAccountID(id int) []byte {
	return []byte(fmt.Sprintf("%09d", id))
}

// Bank stores all the bank related state.
type Bank struct {
	kvClient        *client.KV
	numAccounts     int
	numTransactions int32
}

// moveMoney() moves an amount between two accounts if the amount
// is available in the from account. Returns true on success.
func (bank *Bank) moveMoney(from, to []byte, amount int64) bool {
	// Early exit when from == to
	if bytes.Compare(from, to) == 0 {
		return true
	}
	txnOpts := &client.TransactionOptions{Name: fmt.Sprintf("Transferring %s-%s-%d", from, to, amount)}
	err := bank.kvClient.RunTransaction(txnOpts, func(txn *client.Txn) error {
		fromGet := client.Get(proto.Key(from))
		fromResp := fromGet.Reply.(*proto.GetResponse)
		toGet := client.Get(proto.Key(to))
		toResp := toGet.Reply.(*proto.GetResponse)
		if err := txn.Run(fromGet, toGet); err != nil {
			return err
		}
		// Read from value.
		var fromValue int64
		if fromResp.Value != nil && fromResp.Value.Bytes != nil {
			readValue, err := strconv.ParseInt(string(fromResp.Value.Bytes), 10, 64)
			if err != nil {
				return err
			}
			fromValue = readValue
		}
		// Ensure there is enough cash.
		if fromValue < amount {
			return nil
		}
		// Read to value.
		var toValue int64
		if toResp.Value != nil && toResp.Value.Bytes != nil {
			readValue, err := strconv.ParseInt(string(toResp.Value.Bytes), 10, 64)
			if err != nil {
				return err
			}
			toValue = readValue
		}
		// Update both accounts.
		txn.Prepare(client.Put(proto.Key(from), []byte(fmt.Sprintf("%d", fromValue-amount))))
		txn.Prepare(client.Put(proto.Key(to), []byte(fmt.Sprintf("%d", toValue+amount))))
		return nil
	})
	if err != nil {
		log.Fatal(err)
	}
	atomic.AddInt32(&bank.numTransactions, 1)
	return true
}

// Read the balances in all the accounts and return them.
func (bank *Bank) readAllAccounts() []int64 {
	balances := make([]int64, bank.numAccounts)
	calls := make([]client.Call, bank.numAccounts)
	txnOpts := &client.TransactionOptions{Name: "Reading all balances"}
	err := bank.kvClient.RunTransaction(txnOpts, func(txn *client.Txn) error {
		for i := 0; i < bank.numAccounts; i++ {
			calls[i] = client.Get(proto.Key(makeAccountID(i)))
		}
		if err := txn.Run(calls...); err != nil {
			log.Fatal(err)
		}
		// Copy responses into balances.
		for i := 0; i < bank.numAccounts; i++ {
			gr := calls[i].Reply.(*proto.GetResponse)
			if gr.Value != nil && gr.Value.Bytes != nil {
				balance, err := strconv.ParseInt(string(gr.Value.Bytes), 10, 64)
				if err != nil {
					log.Fatal(err)
				}
				balances[i] = balance
			}
		}
		return nil
	})
	if err != nil {
		log.Fatal(err)
	}
	return balances
}

// continuouslyTransferMoney() keeps moving random amounts between
// random accounts.
func (bank *Bank) continuousMoneyTransfer() {
	for {
		from := makeAccountID(rand.Intn(bank.numAccounts))
		to := makeAccountID(rand.Intn(bank.numAccounts))
		exchange := rand.Int63n(100)
		bank.moveMoney(from, to, exchange)
	}
}

// Initialize all the bank accounts with cash.
func (bank *Bank) initBankAccounts(cash int64) {
	calls := make([]client.Call, bank.numAccounts)
	for i := 0; i < bank.numAccounts; i++ {
		calls[i] = client.Put(proto.Key(makeAccountID(i)), []byte(fmt.Sprintf("%d", cash)))
	}
	if err := bank.kvClient.Run(calls...); err != nil {
		log.Fatal(err)
	}
}

func main() {
	fmt.Printf("A simple program that keeps moving money between bank accounts\n\n")
	// Run a test cockroach instance to represent the bank.
	security.SetReadFileFn(securitytest.Asset)
	serv := server.StartTestServer(nil)
	defer serv.Stop()
	// Initialize the bank.
	var bank Bank
	bank.numAccounts = 1000
	// Key Value Client initialization.
	sender, err := client.NewHTTPSender(serv.ServingAddr(), testutils.NewTestBaseContext())
	if err != nil {
		log.Fatal(err)
	}
	bank.kvClient = client.NewKV(nil, sender)
	bank.kvClient.User = storage.UserRoot
	// Initialize all the bank accounts.
	const initCash = 1000
	bank.initBankAccounts(initCash)

	// Start all the money transfer routines.
	const numTransferRoutines = 10
	for i := 0; i < numTransferRoutines; i++ {
		go bank.continuousMoneyTransfer()
	}
	// Sleep for a bit to allow money transfers to happen in the background and then exit.
	time.Sleep(10 * time.Second)

	fmt.Printf("%d transactions were executed\n\n", bank.numTransactions)

	// Check that all the money is accounted for.
	balances := bank.readAllAccounts()
	var totalAmount int64
	for i := 0; i < bank.numAccounts; i++ {
		fmt.Printf("Account %d contains %d$\n", i, balances[i])
		totalAmount += balances[i]
	}
	if totalAmount != int64(bank.numAccounts)*initCash {
		err := fmt.Sprintf("\nTotal cash in the bank = %d\n", totalAmount)
		log.Fatal(err)
	}
	fmt.Printf("\nThe bank is in good order\n\n")
}
<|endoftext|>


<commit_before>package sarama

import (
	"sync"
	"time"
)

// Offset Manager

// OffsetManager uses Kafka to store and fetch consumed partition offsets.
type OffsetManager interface {
	// ManagePartition creates a PartitionOffsetManager on the given topic/partition. It will
	// return an error if this OffsetManager is already managing the given topic/partition.
	ManagePartition(topic string, partition int32) (PartitionOffsetManager, error)

	// Close stops the OffsetManager from managing offsets. It is required to call this function
	// before an OffsetManager object passes out of scope, as it will otherwise
	// leak memory. You must call this after all the PartitionOffsetManagers are closed.
	Close() error
}

type offsetManager struct {
	client Client
	conf   *Config
	group  string

	lock sync.Mutex
	poms map[string]map[int32]*partitionOffsetManager
	boms map[*Broker]*brokerOffsetManager
}

// NewOffsetManagerFromClient creates a new OffsetManager from the given client.
// It is still necessary to call Close() on the underlying client when finished with the partition manager.
func NewOffsetManagerFromClient(group string, client Client) (OffsetManager, error) {
	// Check that we are not dealing with a closed Client before processing any other arguments
	if client.Closed() {
		return nil, ErrClosedClient
	}

	om := &offsetManager{
		client: client,
		conf:   client.Config(),
		group:  group,
		poms:   make(map[string]map[int32]*partitionOffsetManager),
		boms:   make(map[*Broker]*brokerOffsetManager),
	}

	return om, nil
}

func (om *offsetManager) ManagePartition(topic string, partition int32) (PartitionOffsetManager, error) {
	pom, err := om.newPartitionOffsetManager(topic, partition)
	if err != nil {
		return nil, err
	}

	om.lock.Lock()
	defer om.lock.Unlock()

	topicManagers := om.poms[topic]
	if topicManagers == nil {
		topicManagers = make(map[int32]*partitionOffsetManager)
		om.poms[topic] = topicManagers
	}

	if topicManagers[partition] != nil {
		return nil, ConfigurationError("That topic/partition is already being managed")
	}

	topicManagers[partition] = pom
	return pom, nil
}

func (om *offsetManager) Close() error {
	return nil
}

func (om *offsetManager) refBrokerOffsetManager(broker *Broker) *brokerOffsetManager {
	om.lock.Lock()
	defer om.lock.Unlock()

	bom := om.boms[broker]
	if bom == nil {
		bom = om.newBrokerOffsetManager(broker)
		om.boms[broker] = bom
	}

	bom.refs++

	return bom
}

func (om *offsetManager) unrefBrokerOffsetManager(bom *brokerOffsetManager) {
	om.lock.Lock()
	defer om.lock.Unlock()

	bom.refs--

	if bom.refs == 0 {
		close(bom.updateSubscriptions)
		if om.boms[bom.broker] == bom {
			delete(om.boms, bom.broker)
		}
	}
}

func (om *offsetManager) abandonBroker(bom *brokerOffsetManager) {
	om.lock.Lock()
	defer om.lock.Unlock()

	delete(om.boms, bom.broker)
}

func (om *offsetManager) abandonPartitionOffsetManager(pom *partitionOffsetManager) {
	om.lock.Lock()
	defer om.lock.Unlock()

	delete(om.poms[pom.topic], pom.partition)
	if len(om.poms[pom.topic]) == 0 {
		delete(om.poms, pom.topic)
	}
}

// Partition Offset Manager

// PartitionOffsetManager uses Kafka to store and fetch consumed partition offsets. You MUST call Close()
// on a partition offset manager to avoid leaks, it will not be garbage-collected automatically when it passes
// out of scope.
type PartitionOffsetManager interface {
	// NextOffset returns the next offset that should be consumed for the managed partition, accompanied
	// by metadata which can be used to reconstruct the state of the partition consumer when it resumes.
	// NextOffset() will return `config.Consumer.Offsets.Initial` and an empty metadata string if no
	// offset was committed for this partition yet.
	NextOffset() (int64, string)

	// MarkOffset marks the provided offset as processed, alongside a metadata string that represents
	// the state of the partition consumer at that point in time. The metadata string can be used by
	// another consumer to restore that state, so it can resume consumption.
	//
	// Note: calling MarkOffset does not necessarily commit the offset to the backend store immediately
	// for efficiency reasons, and it may never be committed if your application crashes. This means that
	// you may end up processing the same message twice, and your processing should ideally be idempotent.
	MarkOffset(offset int64, metadata string)

	// Errors returns a read channel of errors that occur during offset management, if enabled. By default,
	// errors are logged and not returned over this channel. If you want to implement any custom error
	// handling, set your config's Consumer.Return.Errors setting to true, and read from this channel.
	Errors() <-chan *ConsumerError

	// AsyncClose initiates a shutdown of the PartitionOffsetManager. This method will return immediately,
	// after which you should wait until the 'errors' channel has been drained and closed.
	// It is required to call this function, or Close before a consumer object passes out of scope,
	// as it will otherwise leak memory.  You must call this before calling Close on the underlying
	// client.
	AsyncClose()

	// Close stops the PartitionOffsetManager from managing offsets. It is required to call this function
	// (or AsyncClose) before a PartitionOffsetManager object passes out of scope, as it will otherwise
	// leak memory. You must call this before calling Close on the underlying client.
	Close() error
}

type partitionOffsetManager struct {
	parent    *offsetManager
	topic     string
	partition int32

	lock     sync.Mutex
	offset   int64
	metadata string
	dirty    bool
	clean    chan none
	broker   *brokerOffsetManager

	errors    chan *ConsumerError
	rebalance chan none
	dying     chan none
}

func (om *offsetManager) newPartitionOffsetManager(topic string, partition int32) (*partitionOffsetManager, error) {
	pom := &partitionOffsetManager{
		parent:    om,
		topic:     topic,
		partition: partition,
		clean:     make(chan none),
		errors:    make(chan *ConsumerError, om.conf.ChannelBufferSize),
		rebalance: make(chan none, 1),
		dying:     make(chan none),
	}

	if err := pom.selectBroker(); err != nil {
		return nil, err
	}

	if err := pom.fetchInitialOffset(om.conf.Metadata.Retry.Max); err != nil {
		return nil, err
	}

	pom.broker.updateSubscriptions <- pom

	go withRecover(pom.mainLoop)

	return pom, nil
}

func (pom *partitionOffsetManager) mainLoop() {
	for {
		select {
		case <-pom.rebalance:
			if err := pom.selectBroker(); err != nil {
				pom.handleError(err)
				pom.rebalance <- none{}
			} else {
				pom.broker.updateSubscriptions <- pom
			}
		case <-pom.dying:
			if pom.broker != nil {
				select {
				case <-pom.rebalance:
				case pom.broker.updateSubscriptions <- pom:
				}
				pom.parent.unrefBrokerOffsetManager(pom.broker)
			}
			pom.parent.abandonPartitionOffsetManager(pom)
			close(pom.errors)
			return
		}
	}
}

func (pom *partitionOffsetManager) selectBroker() error {
	if pom.broker != nil {
		pom.parent.unrefBrokerOffsetManager(pom.broker)
		pom.broker = nil
	}

	var broker *Broker
	var err error

	if err = pom.parent.client.RefreshCoordinator(pom.parent.group); err != nil {
		return err
	}

	if broker, err = pom.parent.client.Coordinator(pom.parent.group); err != nil {
		return err
	}

	pom.broker = pom.parent.refBrokerOffsetManager(broker)
	return nil
}

func (pom *partitionOffsetManager) fetchInitialOffset(retries int) error {
	request := new(OffsetFetchRequest)
	request.Version = 1
	request.ConsumerGroup = pom.parent.group
	request.AddPartition(pom.topic, pom.partition)

	response, err := pom.broker.broker.FetchOffset(request)
	if err != nil {
		return err
	}

	block := response.GetBlock(pom.topic, pom.partition)
	if block == nil {
		return ErrIncompleteResponse
	}

	switch block.Err {
	case ErrNoError:
		pom.offset = block.Offset
		pom.metadata = block.Metadata
		return nil
	case ErrNotCoordinatorForConsumer:
		if retries <= 0 {
			return block.Err
		}
		if err := pom.selectBroker(); err != nil {
			return err
		}
		return pom.fetchInitialOffset(retries - 1)
	case ErrOffsetsLoadInProgress:
		if retries <= 0 {
			return block.Err
		}
		time.Sleep(pom.parent.conf.Metadata.Retry.Backoff)
		return pom.fetchInitialOffset(retries - 1)
	default:
		return block.Err
	}
}

func (pom *partitionOffsetManager) handleError(err error) {
	cErr := &ConsumerError{
		Topic:     pom.topic,
		Partition: pom.partition,
		Err:       err,
	}

	if pom.parent.conf.Consumer.Return.Errors {
		pom.errors <- cErr
	} else {
		Logger.Println(cErr)
	}
}

func (pom *partitionOffsetManager) Errors() <-chan *ConsumerError {
	return pom.errors
}

func (pom *partitionOffsetManager) MarkOffset(offset int64, metadata string) {
	pom.lock.Lock()
	defer pom.lock.Unlock()

	if offset > pom.offset {
		pom.offset = offset
		pom.metadata = metadata
		pom.dirty = true
	}
}

func (pom *partitionOffsetManager) updateCommitted(offset int64, metadata string) {
	pom.lock.Lock()
	defer pom.lock.Unlock()

	if pom.offset == offset && pom.metadata == metadata {
		pom.dirty = false

		select {
		case pom.clean <- none{}:
		default:
		}
	}
}

func (pom *partitionOffsetManager) NextOffset() (int64, string) {
	pom.lock.Lock()
	defer pom.lock.Unlock()

	if pom.offset >= 0 {
		return pom.offset + 1, pom.metadata
	} else {
		return pom.parent.conf.Consumer.Offsets.Initial, ""
	}
}

func (pom *partitionOffsetManager) AsyncClose() {
	go func() {
		pom.lock.Lock()
		dirty := pom.dirty
		pom.lock.Unlock()

		if dirty {
			<-pom.clean
		}

		close(pom.dying)
	}()
}

func (pom *partitionOffsetManager) Close() error {
	pom.AsyncClose()

	var errors ConsumerErrors
	for err := range pom.errors {
		errors = append(errors, err)
	}

	if len(errors) > 0 {
		return errors
	}
	return nil
}

// Broker Offset Manager

type brokerOffsetManager struct {
	parent              *offsetManager
	broker              *Broker
	timer               *time.Ticker
	updateSubscriptions chan *partitionOffsetManager
	subscriptions       map[*partitionOffsetManager]none
	refs                int
}

func (om *offsetManager) newBrokerOffsetManager(broker *Broker) *brokerOffsetManager {
	bom := &brokerOffsetManager{
		parent:              om,
		broker:              broker,
		timer:               time.NewTicker(om.conf.Consumer.Offsets.CommitInterval),
		updateSubscriptions: make(chan *partitionOffsetManager),
		subscriptions:       make(map[*partitionOffsetManager]none),
	}

	go withRecover(bom.mainLoop)

	return bom
}

func (bom *brokerOffsetManager) mainLoop() {
	for {
		select {
		case <-bom.timer.C:
			if len(bom.subscriptions) > 0 {
				bom.flushToBroker()
			}
		case s, ok := <-bom.updateSubscriptions:
			if !ok {
				bom.timer.Stop()
				return
			}
			if _, ok := bom.subscriptions[s]; ok {
				delete(bom.subscriptions, s)
			} else {
				bom.subscriptions[s] = none{}
			}
		}
	}
}

func (bom *brokerOffsetManager) flushToBroker() {
	request := bom.constructRequest()
	if request == nil {
		return
	}

	response, err := bom.broker.CommitOffset(request)

	if err != nil {
		bom.abort(err)
		return
	}

	for s := range bom.subscriptions {
		if request.blocks[s.topic] == nil || request.blocks[s.topic][s.partition] == nil {
			continue
		}

		var err KError
		var ok bool

		if response.Errors[s.topic] == nil {
			s.handleError(ErrIncompleteResponse)
			delete(bom.subscriptions, s)
			s.rebalance <- none{}
			continue
		}
		if err, ok = response.Errors[s.topic][s.partition]; !ok {
			s.handleError(ErrIncompleteResponse)
			delete(bom.subscriptions, s)
			s.rebalance <- none{}
			continue
		}

		switch err {
		case ErrNoError:
			block := request.blocks[s.topic][s.partition]
			s.updateCommitted(block.offset, block.metadata)
			break
		case ErrUnknownTopicOrPartition, ErrNotLeaderForPartition, ErrLeaderNotAvailable:
			delete(bom.subscriptions, s)
			s.rebalance <- none{}
		default:
			s.handleError(err)
			delete(bom.subscriptions, s)
			s.rebalance <- none{}
		}
	}
}

func (bom *brokerOffsetManager) constructRequest() *OffsetCommitRequest {
	r := &OffsetCommitRequest{
		Version:       1,
		ConsumerGroup: bom.parent.group,
	}

	for s := range bom.subscriptions {
		s.lock.Lock()
		if s.dirty {
			r.AddBlock(s.topic, s.partition, s.offset, 0, s.metadata)
		}
		s.lock.Unlock()
	}

	if len(r.blocks) > 0 {
		return r
	}

	return nil
}

func (bom *brokerOffsetManager) abort(err error) {
	_ = bom.broker.Close() // we don't care about the error this might return, we already have one
	bom.parent.abandonBroker(bom)

	for pom := range bom.subscriptions {
		pom.handleError(err)
		pom.rebalance <- none{}
	}

	for s := range bom.updateSubscriptions {
		if _, ok := bom.subscriptions[s]; !ok {
			s.handleError(err)
			s.rebalance <- none{}
		}
	}

	bom.subscriptions = make(map[*partitionOffsetManager]none)
}
<commit_msg>Fix offset-manager expiry timestamps to 'now'<commit_after>package sarama

import (
	"sync"
	"time"
)

// Offset Manager

// OffsetManager uses Kafka to store and fetch consumed partition offsets.
type OffsetManager interface {
	// ManagePartition creates a PartitionOffsetManager on the given topic/partition. It will
	// return an error if this OffsetManager is already managing the given topic/partition.
	ManagePartition(topic string, partition int32) (PartitionOffsetManager, error)

	// Close stops the OffsetManager from managing offsets. It is required to call this function
	// before an OffsetManager object passes out of scope, as it will otherwise
	// leak memory. You must call this after all the PartitionOffsetManagers are closed.
	Close() error
}

type offsetManager struct {
	client Client
	conf   *Config
	group  string

	lock sync.Mutex
	poms map[string]map[int32]*partitionOffsetManager
	boms map[*Broker]*brokerOffsetManager
}

// NewOffsetManagerFromClient creates a new OffsetManager from the given client.
// It is still necessary to call Close() on the underlying client when finished with the partition manager.
func NewOffsetManagerFromClient(group string, client Client) (OffsetManager, error) {
	// Check that we are not dealing with a closed Client before processing any other arguments
	if client.Closed() {
		return nil, ErrClosedClient
	}

	om := &offsetManager{
		client: client,
		conf:   client.Config(),
		group:  group,
		poms:   make(map[string]map[int32]*partitionOffsetManager),
		boms:   make(map[*Broker]*brokerOffsetManager),
	}

	return om, nil
}

func (om *offsetManager) ManagePartition(topic string, partition int32) (PartitionOffsetManager, error) {
	pom, err := om.newPartitionOffsetManager(topic, partition)
	if err != nil {
		return nil, err
	}

	om.lock.Lock()
	defer om.lock.Unlock()

	topicManagers := om.poms[topic]
	if topicManagers == nil {
		topicManagers = make(map[int32]*partitionOffsetManager)
		om.poms[topic] = topicManagers
	}

	if topicManagers[partition] != nil {
		return nil, ConfigurationError("That topic/partition is already being managed")
	}

	topicManagers[partition] = pom
	return pom, nil
}

func (om *offsetManager) Close() error {
	return nil
}

func (om *offsetManager) refBrokerOffsetManager(broker *Broker) *brokerOffsetManager {
	om.lock.Lock()
	defer om.lock.Unlock()

	bom := om.boms[broker]
	if bom == nil {
		bom = om.newBrokerOffsetManager(broker)
		om.boms[broker] = bom
	}

	bom.refs++

	return bom
}

func (om *offsetManager) unrefBrokerOffsetManager(bom *brokerOffsetManager) {
	om.lock.Lock()
	defer om.lock.Unlock()

	bom.refs--

	if bom.refs == 0 {
		close(bom.updateSubscriptions)
		if om.boms[bom.broker] == bom {
			delete(om.boms, bom.broker)
		}
	}
}

func (om *offsetManager) abandonBroker(bom *brokerOffsetManager) {
	om.lock.Lock()
	defer om.lock.Unlock()

	delete(om.boms, bom.broker)
}

func (om *offsetManager) abandonPartitionOffsetManager(pom *partitionOffsetManager) {
	om.lock.Lock()
	defer om.lock.Unlock()

	delete(om.poms[pom.topic], pom.partition)
	if len(om.poms[pom.topic]) == 0 {
		delete(om.poms, pom.topic)
	}
}

// Partition Offset Manager

// PartitionOffsetManager uses Kafka to store and fetch consumed partition offsets. You MUST call Close()
// on a partition offset manager to avoid leaks, it will not be garbage-collected automatically when it passes
// out of scope.
type PartitionOffsetManager interface {
	// NextOffset returns the next offset that should be consumed for the managed partition, accompanied
	// by metadata which can be used to reconstruct the state of the partition consumer when it resumes.
	// NextOffset() will return `config.Consumer.Offsets.Initial` and an empty metadata string if no
	// offset was committed for this partition yet.
	NextOffset() (int64, string)

	// MarkOffset marks the provided offset as processed, alongside a metadata string that represents
	// the state of the partition consumer at that point in time. The metadata string can be used by
	// another consumer to restore that state, so it can resume consumption.
	//
	// Note: calling MarkOffset does not necessarily commit the offset to the backend store immediately
	// for efficiency reasons, and it may never be committed if your application crashes. This means that
	// you may end up processing the same message twice, and your processing should ideally be idempotent.
	MarkOffset(offset int64, metadata string)

	// Errors returns a read channel of errors that occur during offset management, if enabled. By default,
	// errors are logged and not returned over this channel. If you want to implement any custom error
	// handling, set your config's Consumer.Return.Errors setting to true, and read from this channel.
	Errors() <-chan *ConsumerError

	// AsyncClose initiates a shutdown of the PartitionOffsetManager. This method will return immediately,
	// after which you should wait until the 'errors' channel has been drained and closed.
	// It is required to call this function, or Close before a consumer object passes out of scope,
	// as it will otherwise leak memory.  You must call this before calling Close on the underlying
	// client.
	AsyncClose()

	// Close stops the PartitionOffsetManager from managing offsets. It is required to call this function
	// (or AsyncClose) before a PartitionOffsetManager object passes out of scope, as it will otherwise
	// leak memory. You must call this before calling Close on the underlying client.
	Close() error
}

type partitionOffsetManager struct {
	parent    *offsetManager
	topic     string
	partition int32

	lock     sync.Mutex
	offset   int64
	metadata string
	dirty    bool
	clean    chan none
	broker   *brokerOffsetManager

	errors    chan *ConsumerError
	rebalance chan none
	dying     chan none
}

func (om *offsetManager) newPartitionOffsetManager(topic string, partition int32) (*partitionOffsetManager, error) {
	pom := &partitionOffsetManager{
		parent:    om,
		topic:     topic,
		partition: partition,
		clean:     make(chan none),
		errors:    make(chan *ConsumerError, om.conf.ChannelBufferSize),
		rebalance: make(chan none, 1),
		dying:     make(chan none),
	}

	if err := pom.selectBroker(); err != nil {
		return nil, err
	}

	if err := pom.fetchInitialOffset(om.conf.Metadata.Retry.Max); err != nil {
		return nil, err
	}

	pom.broker.updateSubscriptions <- pom

	go withRecover(pom.mainLoop)

	return pom, nil
}

func (pom *partitionOffsetManager) mainLoop() {
	for {
		select {
		case <-pom.rebalance:
			if err := pom.selectBroker(); err != nil {
				pom.handleError(err)
				pom.rebalance <- none{}
			} else {
				pom.broker.updateSubscriptions <- pom
			}
		case <-pom.dying:
			if pom.broker != nil {
				select {
				case <-pom.rebalance:
				case pom.broker.updateSubscriptions <- pom:
				}
				pom.parent.unrefBrokerOffsetManager(pom.broker)
			}
			pom.parent.abandonPartitionOffsetManager(pom)
			close(pom.errors)
			return
		}
	}
}

func (pom *partitionOffsetManager) selectBroker() error {
	if pom.broker != nil {
		pom.parent.unrefBrokerOffsetManager(pom.broker)
		pom.broker = nil
	}

	var broker *Broker
	var err error

	if err = pom.parent.client.RefreshCoordinator(pom.parent.group); err != nil {
		return err
	}

	if broker, err = pom.parent.client.Coordinator(pom.parent.group); err != nil {
		return err
	}

	pom.broker = pom.parent.refBrokerOffsetManager(broker)
	return nil
}

func (pom *partitionOffsetManager) fetchInitialOffset(retries int) error {
	request := new(OffsetFetchRequest)
	request.Version = 1
	request.ConsumerGroup = pom.parent.group
	request.AddPartition(pom.topic, pom.partition)

	response, err := pom.broker.broker.FetchOffset(request)
	if err != nil {
		return err
	}

	block := response.GetBlock(pom.topic, pom.partition)
	if block == nil {
		return ErrIncompleteResponse
	}

	switch block.Err {
	case ErrNoError:
		pom.offset = block.Offset
		pom.metadata = block.Metadata
		return nil
	case ErrNotCoordinatorForConsumer:
		if retries <= 0 {
			return block.Err
		}
		if err := pom.selectBroker(); err != nil {
			return err
		}
		return pom.fetchInitialOffset(retries - 1)
	case ErrOffsetsLoadInProgress:
		if retries <= 0 {
			return block.Err
		}
		time.Sleep(pom.parent.conf.Metadata.Retry.Backoff)
		return pom.fetchInitialOffset(retries - 1)
	default:
		return block.Err
	}
}

func (pom *partitionOffsetManager) handleError(err error) {
	cErr := &ConsumerError{
		Topic:     pom.topic,
		Partition: pom.partition,
		Err:       err,
	}

	if pom.parent.conf.Consumer.Return.Errors {
		pom.errors <- cErr
	} else {
		Logger.Println(cErr)
	}
}

func (pom *partitionOffsetManager) Errors() <-chan *ConsumerError {
	return pom.errors
}

func (pom *partitionOffsetManager) MarkOffset(offset int64, metadata string) {
	pom.lock.Lock()
	defer pom.lock.Unlock()

	if offset > pom.offset {
		pom.offset = offset
		pom.metadata = metadata
		pom.dirty = true
	}
}

func (pom *partitionOffsetManager) updateCommitted(offset int64, metadata string) {
	pom.lock.Lock()
	defer pom.lock.Unlock()

	if pom.offset == offset && pom.metadata == metadata {
		pom.dirty = false

		select {
		case pom.clean <- none{}:
		default:
		}
	}
}

func (pom *partitionOffsetManager) NextOffset() (int64, string) {
	pom.lock.Lock()
	defer pom.lock.Unlock()

	if pom.offset >= 0 {
		return pom.offset + 1, pom.metadata
	} else {
		return pom.parent.conf.Consumer.Offsets.Initial, ""
	}
}

func (pom *partitionOffsetManager) AsyncClose() {
	go func() {
		pom.lock.Lock()
		dirty := pom.dirty
		pom.lock.Unlock()

		if dirty {
			<-pom.clean
		}

		close(pom.dying)
	}()
}

func (pom *partitionOffsetManager) Close() error {
	pom.AsyncClose()

	var errors ConsumerErrors
	for err := range pom.errors {
		errors = append(errors, err)
	}

	if len(errors) > 0 {
		return errors
	}
	return nil
}

// Broker Offset Manager

type brokerOffsetManager struct {
	parent              *offsetManager
	broker              *Broker
	timer               *time.Ticker
	updateSubscriptions chan *partitionOffsetManager
	subscriptions       map[*partitionOffsetManager]none
	refs                int
}

func (om *offsetManager) newBrokerOffsetManager(broker *Broker) *brokerOffsetManager {
	bom := &brokerOffsetManager{
		parent:              om,
		broker:              broker,
		timer:               time.NewTicker(om.conf.Consumer.Offsets.CommitInterval),
		updateSubscriptions: make(chan *partitionOffsetManager),
		subscriptions:       make(map[*partitionOffsetManager]none),
	}

	go withRecover(bom.mainLoop)

	return bom
}

func (bom *brokerOffsetManager) mainLoop() {
	for {
		select {
		case <-bom.timer.C:
			if len(bom.subscriptions) > 0 {
				bom.flushToBroker()
			}
		case s, ok := <-bom.updateSubscriptions:
			if !ok {
				bom.timer.Stop()
				return
			}
			if _, ok := bom.subscriptions[s]; ok {
				delete(bom.subscriptions, s)
			} else {
				bom.subscriptions[s] = none{}
			}
		}
	}
}

func (bom *brokerOffsetManager) flushToBroker() {
	request := bom.constructRequest()
	if request == nil {
		return
	}

	response, err := bom.broker.CommitOffset(request)

	if err != nil {
		bom.abort(err)
		return
	}

	for s := range bom.subscriptions {
		if request.blocks[s.topic] == nil || request.blocks[s.topic][s.partition] == nil {
			continue
		}

		var err KError
		var ok bool

		if response.Errors[s.topic] == nil {
			s.handleError(ErrIncompleteResponse)
			delete(bom.subscriptions, s)
			s.rebalance <- none{}
			continue
		}
		if err, ok = response.Errors[s.topic][s.partition]; !ok {
			s.handleError(ErrIncompleteResponse)
			delete(bom.subscriptions, s)
			s.rebalance <- none{}
			continue
		}

		switch err {
		case ErrNoError:
			block := request.blocks[s.topic][s.partition]
			s.updateCommitted(block.offset, block.metadata)
			break
		case ErrUnknownTopicOrPartition, ErrNotLeaderForPartition, ErrLeaderNotAvailable:
			delete(bom.subscriptions, s)
			s.rebalance <- none{}
		default:
			s.handleError(err)
			delete(bom.subscriptions, s)
			s.rebalance <- none{}
		}
	}
}

func (bom *brokerOffsetManager) constructRequest() *OffsetCommitRequest {
	r := &OffsetCommitRequest{
		Version:       1,
		ConsumerGroup: bom.parent.group,
	}

	for s := range bom.subscriptions {
		s.lock.Lock()
		if s.dirty {
			r.AddBlock(s.topic, s.partition, s.offset, ReceiveTime, s.metadata)
		}
		s.lock.Unlock()
	}

	if len(r.blocks) > 0 {
		return r
	}

	return nil
}

func (bom *brokerOffsetManager) abort(err error) {
	_ = bom.broker.Close() // we don't care about the error this might return, we already have one
	bom.parent.abandonBroker(bom)

	for pom := range bom.subscriptions {
		pom.handleError(err)
		pom.rebalance <- none{}
	}

	for s := range bom.updateSubscriptions {
		if _, ok := bom.subscriptions[s]; !ok {
			s.handleError(err)
			s.rebalance <- none{}
		}
	}

	bom.subscriptions = make(map[*partitionOffsetManager]none)
}
<|endoftext|>


<commit_before>package udsstorage

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"pfi/sensorbee/sensorbee/bql/udf"
	"regexp"
)

// fsUDSStorage is a UDSStorage which store states as files on a filesystem.
// This simple storage doesn't provide any version controlling capability.
// It doesn't provide checksum, either. If such capability is required, another
// UDSStorage should be implemented.
type fsUDSStorage struct {
	dirPath    string
	tmpDirPath string
}

var (
	_ udf.UDSStorage = &fsUDSStorage{}
)

func NewFS(dir, tmpDir string) udf.UDSStorage {
	return &fsUDSStorage{
		dirPath:    dir,
		tmpDirPath: tmpDir,
	}
}

func (s *fsUDSStorage) Save(topology, state string) (udf.UDSStorageWriter, error) {
	f, err := s.stateTmpFile(topology, state)
	if err != nil {
		return nil, err
	}
	return &onDiskUDSStorageWriter{
		s:          s,
		f:          f,
		w:          bufio.NewWriter(f),
		targetPath: s.stateFilepath(topology, state),
	}, nil
}

func (s *fsUDSStorage) stateTmpFile(topology, state string) (*os.File, error) {
	return ioutil.TempFile(s.tmpDirPath, s.stateFilename(topology, state))
}

func (s *fsUDSStorage) Load(topology, state string) (io.ReadCloser, error) {
	f, err := os.Open(s.stateFilepath(topology, state))
	if err != nil {
		return nil, err
	}
	return f, nil
}

var (
	fsUDSStorageFilePathRegexp = regexp.MustCompile(`^(.+)-(.+).state$`)
)

func (s *fsUDSStorage) List() (map[string][]string, error) {
	fs, err := ioutil.ReadDir(s.dirPath)
	if err != nil {
		return nil, err
	}

	res := map[string][]string{}
	for _, f := range fs {
		m := fsUDSStorageFilePathRegexp.FindStringSubmatch(f.Name())
		if m == nil {
			continue
		}
		res[m[1]] = append(res[m[1]], m[2])
	}
	return res, nil
}

func (s *fsUDSStorage) stateFilename(topology, state string) string {
	return fmt.Sprintf("%v-%v.state", topology, state)
}

func (s *fsUDSStorage) stateFilepath(topology, state string) string {
	return filepath.Join(s.dirPath, s.stateFilename(topology, state))
}

type onDiskUDSStorageWriter struct {
	s *fsUDSStorage
	f *os.File
	w *bufio.Writer

	targetPath string
}

func (w *onDiskUDSStorageWriter) Write(data []byte) (int, error) {
	if w.w == nil {
		return 0, errors.New("writer is already closed")
	}
	return w.w.Write(data)
}

func (w *onDiskUDSStorageWriter) Commit() (err error) {
	if w.w == nil {
		return errors.New("writer is already closed")
	}
	defer func() {
		if w.f != nil {
			if e := w.f.Close(); e != nil && err == nil {
				err = e
			}
			w.f = nil
		}
	}()

	if e := w.w.Flush(); e != nil {
		return e
	}
	w.w = nil

	f := w.f
	// TODO: Name doesn't return an absolute path, so this code might fail with
	// some scenarios.
	fn := f.Name()
	w.f = nil
	if e := f.Close(); e != nil {
		return e
	}

	// This Rename might not always be atomic and application level locking
	// might be required.
	if e := os.Rename(fn, w.targetPath); e != nil {
		return e
	}
	return nil
}

func (w *onDiskUDSStorageWriter) Abort() error {
	if w.w == nil {
		return errors.New("writer is already closed")
	}
	w.w = nil

	f := w.f
	w.f = nil
	fn := f.Name()
	if err := f.Close(); err != nil {
		return err
	}
	return os.Remove(fn)
}
<commit_msg>Rename tmp to temp to be consistent with Go's standard libraries<commit_after>package udsstorage

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"pfi/sensorbee/sensorbee/bql/udf"
	"regexp"
)

// fsUDSStorage is a UDSStorage which store states as files on a filesystem.
// This simple storage doesn't provide any version controlling capability.
// It doesn't provide checksum, either. If such capability is required, another
// UDSStorage should be implemented.
type fsUDSStorage struct {
	dirPath     string
	tempDirPath string
}

var (
	_ udf.UDSStorage = &fsUDSStorage{}
)

func NewFS(dir, tempDir string) udf.UDSStorage {
	return &fsUDSStorage{
		dirPath:     dir,
		tempDirPath: tempDir,
	}
}

func (s *fsUDSStorage) Save(topology, state string) (udf.UDSStorageWriter, error) {
	f, err := s.stateTempFile(topology, state)
	if err != nil {
		return nil, err
	}
	return &onDiskUDSStorageWriter{
		s:          s,
		f:          f,
		w:          bufio.NewWriter(f),
		targetPath: s.stateFilepath(topology, state),
	}, nil
}

func (s *fsUDSStorage) stateTempFile(topology, state string) (*os.File, error) {
	return ioutil.TempFile(s.tempDirPath, s.stateFilename(topology, state))
}

func (s *fsUDSStorage) Load(topology, state string) (io.ReadCloser, error) {
	f, err := os.Open(s.stateFilepath(topology, state))
	if err != nil {
		return nil, err
	}
	return f, nil
}

var (
	fsUDSStorageFilePathRegexp = regexp.MustCompile(`^(.+)-(.+).state$`)
)

func (s *fsUDSStorage) List() (map[string][]string, error) {
	fs, err := ioutil.ReadDir(s.dirPath)
	if err != nil {
		return nil, err
	}

	res := map[string][]string{}
	for _, f := range fs {
		m := fsUDSStorageFilePathRegexp.FindStringSubmatch(f.Name())
		if m == nil {
			continue
		}
		res[m[1]] = append(res[m[1]], m[2])
	}
	return res, nil
}

func (s *fsUDSStorage) stateFilename(topology, state string) string {
	return fmt.Sprintf("%v-%v.state", topology, state)
}

func (s *fsUDSStorage) stateFilepath(topology, state string) string {
	return filepath.Join(s.dirPath, s.stateFilename(topology, state))
}

type onDiskUDSStorageWriter struct {
	s *fsUDSStorage
	f *os.File
	w *bufio.Writer

	targetPath string
}

func (w *onDiskUDSStorageWriter) Write(data []byte) (int, error) {
	if w.w == nil {
		return 0, errors.New("writer is already closed")
	}
	return w.w.Write(data)
}

func (w *onDiskUDSStorageWriter) Commit() (err error) {
	if w.w == nil {
		return errors.New("writer is already closed")
	}
	defer func() {
		if w.f != nil {
			if e := w.f.Close(); e != nil && err == nil {
				err = e
			}
			w.f = nil
		}
	}()

	if e := w.w.Flush(); e != nil {
		return e
	}
	w.w = nil

	f := w.f
	// TODO: Name doesn't return an absolute path, so this code might fail with
	// some scenarios.
	fn := f.Name()
	w.f = nil
	if e := f.Close(); e != nil {
		return e
	}

	// This Rename might not always be atomic and application level locking
	// might be required.
	if e := os.Rename(fn, w.targetPath); e != nil {
		return e
	}
	return nil
}

func (w *onDiskUDSStorageWriter) Abort() error {
	if w.w == nil {
		return errors.New("writer is already closed")
	}
	w.w = nil

	f := w.f
	w.f = nil
	fn := f.Name()
	if err := f.Close(); err != nil {
		return err
	}
	return os.Remove(fn)
}
<|endoftext|>


<commit_before>package main

import (
	"os"
	"os/exec"
)

func cmdRun(args []string) error {
	dir, err := os.Getwd()
	if err != nil {
		return err
	}

	subargs := []string{"-lsst=" + dir, "fcs-run"}
	if len(args) <= 0 {
		subargs = append(subargs, "bash")
	} else {
		subargs = append(subargs, args...)
	}

	// make sure the distrib is up-to-date
	err = cmdDist(nil)
	if err != nil {
		return err
	}

	cmd := exec.Command(
		"fcs-boot",
		subargs...,
	)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	err = cmd.Run()
	return err
}
<commit_msg>fcs-mgr: fixup 'run' w/o args<commit_after>package main

import (
	"os"
	"os/exec"
)

func cmdRun(args []string) error {
	dir, err := os.Getwd()
	if err != nil {
		return err
	}

	subargs := []string{"-lsst=" + dir, "fcs-run"}
	if len(args) <= 0 {
		subargs = append(subargs, "shell")
	} else {
		subargs = append(subargs, args...)
	}

	// make sure the distrib is up-to-date
	err = cmdDist(nil)
	if err != nil {
		return err
	}

	cmd := exec.Command(
		"fcs-boot",
		subargs...,
	)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	err = cmd.Run()
	return err
}
<|endoftext|>


<commit_before>package cmd

import (
	"os"

	"github.com/markbates/pop"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"
)

var migrationPath string

var migrateCmd = &cobra.Command{
	Use:     "migrate",
	Aliases: []string{"m"},
	Short:   "Runs migrations against your database.",
	PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
		RootCmd.PersistentPreRun(cmd, args)
		return os.MkdirAll(migrationPath, 0766)
	},
	RunE: func(cmd *cobra.Command, args []string) error {
		mig, err := pop.NewFileMigrator(migrationPath, getConn())
		if err != nil {
			return errors.WithStack(err)
		}
		err = mig.Up()
		if err != nil {
			return errors.WithStack(err)
		}

		return mig.DumpMigrationSchema()
	},
}

func init() {
	RootCmd.AddCommand(migrateCmd)
	RootCmd.PersistentFlags().StringVarP(&migrationPath, "path", "p", "./migrations", "Path to the migrations folder")
}
<commit_msg>removed double DumpMigrationSchema call<commit_after>package cmd

import (
	"os"

	"github.com/markbates/pop"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"
)

var migrationPath string

var migrateCmd = &cobra.Command{
	Use:     "migrate",
	Aliases: []string{"m"},
	Short:   "Runs migrations against your database.",
	PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
		RootCmd.PersistentPreRun(cmd, args)
		return os.MkdirAll(migrationPath, 0766)
	},
	RunE: func(cmd *cobra.Command, args []string) error {
		mig, err := pop.NewFileMigrator(migrationPath, getConn())
		if err != nil {
			return errors.WithStack(err)
		}
		return mig.Up()
	},
}

func init() {
	RootCmd.AddCommand(migrateCmd)
	RootCmd.PersistentFlags().StringVarP(&migrationPath, "path", "p", "./migrations", "Path to the migrations folder")
}
<|endoftext|>


<commit_before>package cmd

const Version = "3.42.1"
<commit_msg>version bump<commit_after>package cmd

const Version = "3.50.0"
<|endoftext|>


<commit_before>// Copyright 2015 CoreOS, 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 main

import (
	"flag"
	"fmt"
	"os"
	"path"
	"strings"
	"text/tabwriter"

	"github.com/coreos/locksmith/etcd"
	"github.com/coreos/locksmith/lock"
	"github.com/coreos/locksmith/version"
)

const (
	cliName         = "locksmithctl"
	cliDescription  = `Manage the cluster wide reboot lock.`
	defaultEndpoint = "http://127.0.0.1:4001"
)

var (
	out *tabwriter.Writer

	commands      []*Command
	globalFlagSet *flag.FlagSet = flag.NewFlagSet("locksmithctl", flag.ExitOnError)

	globalFlags = struct {
		Debug        bool
		Endpoints    endpoints
		EtcdKeyFile  string
		EtcdCertFile string
		EtcdCAFile   string
		Version      bool
	}{}
)

type endpoints []string

func (e *endpoints) String() string {
	if len(*e) == 0 {
		return defaultEndpoint
	}

	return strings.Join(*e, ",")
}

func (e *endpoints) Set(value string) error {
	for _, url := range strings.Split(value, ",") {
		*e = append(*e, strings.TrimSpace(url))
	}

	return nil
}

func init() {
	out = new(tabwriter.Writer)
	out.Init(os.Stdout, 0, 8, 1, '\t', 0)

	globalFlagSet.BoolVar(&globalFlags.Debug, "debug", false, "Print out debug information to stderr.")
	globalFlagSet.Var(&globalFlags.Endpoints, "endpoint", "etcd endpoint for locksmith. Specify multiple times to use multiple endpoints.")
	globalFlagSet.StringVar(&globalFlags.EtcdKeyFile, "etcd-keyfile", "", "etcd key file authentication")
	globalFlagSet.StringVar(&globalFlags.EtcdCertFile, "etcd-certfile", "", "etcd cert file authentication")
	globalFlagSet.StringVar(&globalFlags.EtcdCAFile, "etcd-cafile", "", "etcd CA file authentication")
	globalFlagSet.BoolVar(&globalFlags.Version, "version", false, "Print the version and exit.")

	commands = []*Command{
		cmdHelp,
		cmdLock,
		cmdReboot,
		cmdSendNeedReboot,
		cmdSetMax,
		cmdStatus,
		cmdUnlock,
	}
}

type Command struct {
	Name        string                  // Name of the Command and the string to use to invoke it
	Summary     string                  // One-sentence summary of what the Command does
	Usage       string                  // Usage options/arguments
	Description string                  // Detailed description of command
	Flags       flag.FlagSet            // Set of flags associated with this command
	Run         func(args []string) int // Run a command with the given arguments, return exit status
}

func getAllFlags() (flags []*flag.Flag) {
	return getFlags(globalFlagSet)
}

func getFlags(flagSet *flag.FlagSet) (flags []*flag.Flag) {
	flags = make([]*flag.Flag, 0)
	flagSet.VisitAll(func(f *flag.Flag) {
		flags = append(flags, f)
	})
	return
}

func main() {
	globalFlagSet.Parse(os.Args[1:])
	var args = globalFlagSet.Args()

	if len(globalFlags.Endpoints) == 0 {
		globalFlags.Endpoints = []string{defaultEndpoint}
	}

	progName := path.Base(os.Args[0])

	if globalFlags.Version {
		fmt.Printf("%s version %s\n", progName, version.Version)
		os.Exit(0)
	}

	if progName == "locksmithd" {
		flagsFromEnv("LOCKSMITHD", globalFlagSet)
		os.Exit(runDaemon())
	}

	// no command specified - trigger help
	if len(args) < 1 {
		args = append(args, "help")
	}

	flagsFromEnv("LOCKSMITHCTL", globalFlagSet)

	var cmd *Command

	// determine which Command should be run
	for _, c := range commands {
		if c.Name == args[0] {
			cmd = c
			if err := c.Flags.Parse(args[1:]); err != nil {
				fmt.Println(err.Error())
				os.Exit(2)
			}
			break
		}
	}

	if cmd == nil {
		fmt.Printf("%v: unknown subcommand: %q\n", cliName, args[0])
		fmt.Printf("Run '%v help' for usage.\n", cliName)
		os.Exit(2)
	}

	os.Exit(cmd.Run(cmd.Flags.Args()))
}

// getLockClient returns an initialized EtcdLockClient, using an etcd
// client configured from the global etcd flags
func getClient() (*lock.EtcdLockClient, error) {
	var ti *etcd.TLSInfo
	if globalFlags.EtcdCAFile != "" || globalFlags.EtcdCertFile != "" || globalFlags.EtcdKeyFile != "" {
		ti = &etcd.TLSInfo{
			CertFile: globalFlags.EtcdCertFile,
			KeyFile:  globalFlags.EtcdKeyFile,
			CAFile:   globalFlags.EtcdCAFile,
		}
	}
	ec, err := etcd.NewClient(globalFlags.Endpoints, ti)
	if err != nil {
		return nil, err
	}
	lc, err := lock.NewEtcdLockClient(ec)
	if err != nil {
		return nil, err
	}
	return lc, err
}

// flagsFromEnv parses all registered flags in the given flagSet,
// and if they are not already set it attempts to set their values from
// environment variables. Environment variables take the name of the flag but
// are UPPERCASE, have the given prefix, and any dashes are replaced by
// underscores - for example: some-flag => PREFIX_SOME_FLAG
func flagsFromEnv(prefix string, fs *flag.FlagSet) {
	alreadySet := make(map[string]bool)
	fs.Visit(func(f *flag.Flag) {
		alreadySet[f.Name] = true
	})
	fs.VisitAll(func(f *flag.Flag) {
		if !alreadySet[f.Name] {
			key := strings.ToUpper(prefix + "_" + strings.Replace(f.Name, "-", "_", -1))
			val := os.Getenv(key)
			if val != "" {
				fs.Set(f.Name, val)
			}
		}
	})
}
<commit_msg>locksmithctl: use IANA port<commit_after>// Copyright 2015 CoreOS, 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 main

import (
	"flag"
	"fmt"
	"os"
	"path"
	"strings"
	"text/tabwriter"

	"github.com/coreos/locksmith/etcd"
	"github.com/coreos/locksmith/lock"
	"github.com/coreos/locksmith/version"
)

const (
	cliName        = "locksmithctl"
	cliDescription = `Manage the cluster wide reboot lock.`
)

var (
	out *tabwriter.Writer

	commands      []*Command
	globalFlagSet *flag.FlagSet = flag.NewFlagSet("locksmithctl", flag.ExitOnError)

	globalFlags = struct {
		Debug        bool
		Endpoints    endpoints
		EtcdKeyFile  string
		EtcdCertFile string
		EtcdCAFile   string
		Version      bool
	}{}

	defaultEndpoints = []string{
		"http://127.0.0.1:2379",
		"http://127.0.0.1:4001",
	}
)

type endpoints []string

func (e *endpoints) String() string {
	if len(*e) == 0 {
		return strings.Join(defaultEndpoints, ",")
	}

	return strings.Join(*e, ",")
}

func (e *endpoints) Set(value string) error {
	for _, url := range strings.Split(value, ",") {
		*e = append(*e, strings.TrimSpace(url))
	}

	return nil
}

func init() {
	out = new(tabwriter.Writer)
	out.Init(os.Stdout, 0, 8, 1, '\t', 0)

	globalFlagSet.BoolVar(&globalFlags.Debug, "debug", false, "Print out debug information to stderr.")
	globalFlagSet.Var(&globalFlags.Endpoints, "endpoint", "etcd endpoint for locksmith. Specify multiple times to use multiple endpoints.")
	globalFlagSet.StringVar(&globalFlags.EtcdKeyFile, "etcd-keyfile", "", "etcd key file authentication")
	globalFlagSet.StringVar(&globalFlags.EtcdCertFile, "etcd-certfile", "", "etcd cert file authentication")
	globalFlagSet.StringVar(&globalFlags.EtcdCAFile, "etcd-cafile", "", "etcd CA file authentication")
	globalFlagSet.BoolVar(&globalFlags.Version, "version", false, "Print the version and exit.")

	commands = []*Command{
		cmdHelp,
		cmdLock,
		cmdReboot,
		cmdSendNeedReboot,
		cmdSetMax,
		cmdStatus,
		cmdUnlock,
	}
}

type Command struct {
	Name        string                  // Name of the Command and the string to use to invoke it
	Summary     string                  // One-sentence summary of what the Command does
	Usage       string                  // Usage options/arguments
	Description string                  // Detailed description of command
	Flags       flag.FlagSet            // Set of flags associated with this command
	Run         func(args []string) int // Run a command with the given arguments, return exit status
}

func getAllFlags() (flags []*flag.Flag) {
	return getFlags(globalFlagSet)
}

func getFlags(flagSet *flag.FlagSet) (flags []*flag.Flag) {
	flags = make([]*flag.Flag, 0)
	flagSet.VisitAll(func(f *flag.Flag) {
		flags = append(flags, f)
	})
	return
}

func main() {
	globalFlagSet.Parse(os.Args[1:])
	var args = globalFlagSet.Args()

	if len(globalFlags.Endpoints) == 0 {
		globalFlags.Endpoints = defaultEndpoints
	}

	progName := path.Base(os.Args[0])

	if globalFlags.Version {
		fmt.Printf("%s version %s\n", progName, version.Version)
		os.Exit(0)
	}

	if progName == "locksmithd" {
		flagsFromEnv("LOCKSMITHD", globalFlagSet)
		os.Exit(runDaemon())
	}

	// no command specified - trigger help
	if len(args) < 1 {
		args = append(args, "help")
	}

	flagsFromEnv("LOCKSMITHCTL", globalFlagSet)

	var cmd *Command

	// determine which Command should be run
	for _, c := range commands {
		if c.Name == args[0] {
			cmd = c
			if err := c.Flags.Parse(args[1:]); err != nil {
				fmt.Println(err.Error())
				os.Exit(2)
			}
			break
		}
	}

	if cmd == nil {
		fmt.Printf("%v: unknown subcommand: %q\n", cliName, args[0])
		fmt.Printf("Run '%v help' for usage.\n", cliName)
		os.Exit(2)
	}

	os.Exit(cmd.Run(cmd.Flags.Args()))
}

// getLockClient returns an initialized EtcdLockClient, using an etcd
// client configured from the global etcd flags
func getClient() (*lock.EtcdLockClient, error) {
	var ti *etcd.TLSInfo
	if globalFlags.EtcdCAFile != "" || globalFlags.EtcdCertFile != "" || globalFlags.EtcdKeyFile != "" {
		ti = &etcd.TLSInfo{
			CertFile: globalFlags.EtcdCertFile,
			KeyFile:  globalFlags.EtcdKeyFile,
			CAFile:   globalFlags.EtcdCAFile,
		}
	}
	ec, err := etcd.NewClient(globalFlags.Endpoints, ti)
	if err != nil {
		return nil, err
	}
	lc, err := lock.NewEtcdLockClient(ec)
	if err != nil {
		return nil, err
	}
	return lc, err
}

// flagsFromEnv parses all registered flags in the given flagSet,
// and if they are not already set it attempts to set their values from
// environment variables. Environment variables take the name of the flag but
// are UPPERCASE, have the given prefix, and any dashes are replaced by
// underscores - for example: some-flag => PREFIX_SOME_FLAG
func flagsFromEnv(prefix string, fs *flag.FlagSet) {
	alreadySet := make(map[string]bool)
	fs.Visit(func(f *flag.Flag) {
		alreadySet[f.Name] = true
	})
	fs.VisitAll(func(f *flag.Flag) {
		if !alreadySet[f.Name] {
			key := strings.ToUpper(prefix + "_" + strings.Replace(f.Name, "-", "_", -1))
			val := os.Getenv(key)
			if val != "" {
				fs.Set(f.Name, val)
			}
		}
	})
}
<|endoftext|>