<commit_before>package SimpleCUPSPrintService

import (
	"fmt"
	"log"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"text/template"
	"time"
)

const (
	maxNameLength = 16
)

//LabelPrinter ...
type LabelPrinter struct {
	isToBeSentToPrinter bool
	printerName,
	templatePath,
	outputPath string
	errorLog,
	infoLog *log.Logger
	labelTemplate *template.Template
}

//NewLabelPrinter  ...
func NewLabelPrinter(isToBeSentToPrinter bool,
	printerName,
	templatePath,
	outputPath string,
	errorLog,
	infoLog *log.Logger) *LabelPrinter {
	templatePathParts := strings.Split(templatePath, string(os.PathSeparator))
	templateName := templatePathParts[len(templatePathParts)-1]
	labelTemplate, err := template.New(templateName).ParseFiles(templatePath)
	if err != nil {
		return nil
	}
	var result = LabelPrinter{isToBeSentToPrinter,
		printerName,
		templatePath,
		outputPath,
		errorLog,
		infoLog,
		labelTemplate}

	return &result
}

func formatLabelPart(maxLength int, labelPart string) string {
	partLength := len(labelPart)
	if partLength >= maxLength {
		return labelPart[0:maxLength-3] + "..."
	}
	return strings.Title(fmt.Sprintf("%s", labelPart))
}

func sendLabelToPrinter(filePath string, printerName string) (err error) {
	args := []string{filePath, printerName, "media=w36h89"}
	err = exec.Command("lpr", args[0], "-P", args[1], "-o", args[2]).Run()
	return err
}

func prepareLabelInfoForPrinting(lableInfo *LabelInfo) {
	lableInfo.Code = strings.ToUpper(formatLabelPart(maxNameLength, lableInfo.Id))
	lableInfo.Date = time.Now().Format("2006-01-02 15:04:05")
	lableInfo.FirstName = formatLabelPart(maxNameLength, lableInfo.FirstName)
	lableInfo.LastName = formatLabelPart(maxNameLength, lableInfo.LastName)
	lableInfo.Notes = lableInfo.AdditionalInfo
	if lableInfo.CanPhotograph {
		lableInfo.Photography = ""
	} else {
		lableInfo.Photography = "Don't photograph"
	}
	if lableInfo.HasAllergies {
		lableInfo.Allergies = "Allergies"
	} else {
		lableInfo.Allergies = ""
	}
}

func writeLabel(lableInfo *LabelInfo,
	labelTemplate *template.Template,
	outputPath string) (err error) {
	prepareLabelInfoForPrinting(lableInfo)
	w, err := os.OpenFile(outputPath, os.O_CREATE|os.O_RDWR, os.ModePerm)
	if err != nil {
		return err
	}
	defer w.Close()
	labelTemplate.Execute(w, lableInfo)
	return nil
}

//Print prints labels
func (labelPrinter *LabelPrinter) Print(printJob LabelInfo,
	numberOfCopies int) (err error) {

	for i := 1; i <= numberOfCopies; i++ {
		labelPrinter.infoLog.Printf("Printing label: %s ... \n", printJob)
		timeStamp := time.Now()
		mockGUID := strconv.Itoa(timeStamp.Nanosecond())
		labelFileName := labelPrinter.outputPath + printJob.Id + "_" + timeStamp.Format("2006_01_02_15_04_05") + "_" + mockGUID + ".ps"
		if i > 1 {
			printJob.Title = "Parent / Guardian Copy"
		}

		err = writeLabel(&printJob, labelPrinter.labelTemplate, labelFileName)
		if err != nil {
			labelPrinter.errorLog.Printf("Error printing label  %v \n", err)
			return err
		}
		if !labelPrinter.isToBeSentToPrinter {
			return nil
		}

		err = sendLabelToPrinter(labelFileName, labelPrinter.printerName)
		if err != nil {
			labelPrinter.errorLog.Printf("Error printing label  %v \n", err)
			return err
		}

		labelPrinter.infoLog.Println("Printed label successfully. Please collect at Label Station.")
	}
	return nil
}
<commit_msg>Stopped returning when print-out not desired, instead continues.<commit_after>package SimpleCUPSPrintService

import (
	"fmt"
	"log"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"text/template"
	"time"
)

const (
	maxNameLength = 16
)

//LabelPrinter ...
type LabelPrinter struct {
	isToBeSentToPrinter bool
	printerName,
	templatePath,
	outputPath string
	errorLog,
	infoLog *log.Logger
	labelTemplate *template.Template
}

//NewLabelPrinter  ...
func NewLabelPrinter(isToBeSentToPrinter bool,
	printerName,
	templatePath,
	outputPath string,
	errorLog,
	infoLog *log.Logger) *LabelPrinter {
	templatePathParts := strings.Split(templatePath, string(os.PathSeparator))
	templateName := templatePathParts[len(templatePathParts)-1]
	labelTemplate, err := template.New(templateName).ParseFiles(templatePath)
	if err != nil {
		return nil
	}
	var result = LabelPrinter{isToBeSentToPrinter,
		printerName,
		templatePath,
		outputPath,
		errorLog,
		infoLog,
		labelTemplate}

	return &result
}

func formatLabelPart(maxLength int, labelPart string) string {
	partLength := len(labelPart)
	if partLength >= maxLength {
		return labelPart[0:maxLength-3] + "..."
	}
	return strings.Title(fmt.Sprintf("%s", labelPart))
}

func sendLabelToPrinter(filePath string, printerName string) (err error) {
	args := []string{filePath, printerName, "media=w36h89"}
	err = exec.Command("lpr", args[0], "-P", args[1], "-o", args[2]).Run()
	return err
}

func prepareLabelInfoForPrinting(lableInfo *LabelInfo) {
	lableInfo.Code = strings.ToUpper(formatLabelPart(maxNameLength, lableInfo.Id))
	lableInfo.Date = time.Now().Format("2006-01-02 15:04:05")
	lableInfo.FirstName = formatLabelPart(maxNameLength, lableInfo.FirstName)
	lableInfo.LastName = formatLabelPart(maxNameLength, lableInfo.LastName)
	lableInfo.Notes = lableInfo.AdditionalInfo
	if lableInfo.CanPhotograph {
		lableInfo.Photography = ""
	} else {
		lableInfo.Photography = "Don't photograph"
	}
	if lableInfo.HasAllergies {
		lableInfo.Allergies = "Allergies"
	} else {
		lableInfo.Allergies = ""
	}
}

func writeLabel(lableInfo *LabelInfo,
	labelTemplate *template.Template,
	outputPath string) (err error) {
	prepareLabelInfoForPrinting(lableInfo)
	w, err := os.OpenFile(outputPath, os.O_CREATE|os.O_RDWR, os.ModePerm)
	if err != nil {
		return err
	}
	defer w.Close()
	labelTemplate.Execute(w, lableInfo)
	return nil
}

//Print prints labels
func (labelPrinter *LabelPrinter) Print(printJob LabelInfo,
	numberOfCopies int) (err error) {

	for i := 1; i <= numberOfCopies; i++ {
		labelPrinter.infoLog.Printf("Label Printer: Printing label: %s ... \n", printJob)
		timeStamp := time.Now()
		mockGUID := strconv.Itoa(timeStamp.Nanosecond())

		labelFileName := labelPrinter.outputPath + printJob.Id + "_" + timeStamp.Format("2006_01_02_15_04_05") + "_" + mockGUID + ".ps"
		if i > 1 {
			labelFileName = labelPrinter.outputPath + printJob.Id + fmt.Sprintf("_ %d", i) + "_" + timeStamp.Format("2006_01_02_15_04_05") + "_" + mockGUID + ".ps"
			printJob.Title = "Parent / Guardian Copy"
		}

		err = writeLabel(&printJob, labelPrinter.labelTemplate, labelFileName)
		if err != nil {
			labelPrinter.errorLog.Printf("Error printing label  %v \n", err)
			return err
		}
		if !labelPrinter.isToBeSentToPrinter {
			labelPrinter.infoLog.Println("Printed label successfully. Please collect at Label Station.")
			continue
		}

		err = sendLabelToPrinter(labelFileName, labelPrinter.printerName)
		if err != nil {
			labelPrinter.errorLog.Printf("Error printing label  %v \n", err)
			return err
		}

		labelPrinter.infoLog.Println("Printed label successfully. Please collect at Label Station.")
	}
	return nil
}
<|endoftext|>


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

package app

import (
	"fmt"
	"github.com/globocom/tsuru/db"
	"github.com/globocom/tsuru/log"
	"github.com/globocom/tsuru/provision"
	"github.com/globocom/tsuru/queue"
	"github.com/globocom/tsuru/service"
	"io/ioutil"
	"labix.org/v2/mgo/bson"
)

const (
	// queue actions
	regenerateApprc         = "regenerate-apprc"
	startApp                = "start-app"
	RegenerateApprcAndStart = "regenerate-apprc-start-app"
	bindService             = "bind-service"

	// name of the queue for internal messages.
	queueName = "tsuru-app"

	// Name of the queue for external messages.
	QueueName = "tsuru-app-public"
)

// ensureAppIsStarted make sure that the app and all units present in the given
// message are started.
func ensureAppIsStarted(msg *queue.Message) (App, error) {
	a := App{Name: msg.Args[0]}
	err := a.Get()
	if err != nil {
		return a, fmt.Errorf("Error handling %q: app %q does not exist.", msg.Action, a.Name)
	}
	units := getUnits(&a, msg.Args[1:])
	if !a.Available() || !units.Started() {
		format := "Error handling %q for the app %q:"
		uState := units.State()
		if uState == "error" || uState == "down" {
			format += fmt.Sprintf(" units are in %q state.", uState)
			msg.Delete()
		} else {
			format += " all units must be started."
		}
		return a, fmt.Errorf(format, msg.Action, a.Name)
	}
	return a, nil
}

// bindUnit handles the bind-service message, binding a unit to all service
// instances binded to the app.
func bindUnit(msg *queue.Message) error {
	a := App{Name: msg.Args[0]}
	err := a.Get()
	if err != nil {
		msg.Delete()
		return fmt.Errorf("Error handling %q: app %q does not exist.", msg.Action, a.Name)
	}
	conn, err := db.Conn()
	if err != nil {
		return fmt.Errorf("Error handling %q: %s", msg.Action, err)
	}
	unit := getUnits(&a, msg.Args[1:])[0]
	var instances []service.ServiceInstance
	q := bson.M{"apps": bson.M{"$in": []string{msg.Args[0]}}}
	err = conn.ServiceInstances().Find(q).All(&instances)
	if err != nil {
		return err
	}
	for _, instance := range instances {
		_, err = instance.BindUnit(&unit)
		if err != nil {
			log.Printf("Error binding the unit %s with the service instance %s.", unit.Name, instance.Name)
		}
	}
	return nil
}

// handle is the function called by the queue handler on each message.
func handle(msg *queue.Message) {
	switch msg.Action {
	case RegenerateApprcAndStart:
		fallthrough
	case regenerateApprc:
		if len(msg.Args) < 1 {
			log.Printf("Error handling %q: this action requires at least 1 argument.", msg.Action)
			return
		}
		app, err := ensureAppIsStarted(msg)
		if err != nil {
			log.Print(err)
			return
		}
		msg.Delete()
		app.serializeEnvVars()
		fallthrough
	case startApp:
		if msg.Action == regenerateApprc {
			break
		}
		if len(msg.Args) < 1 {
			log.Printf("Error handling %q: this action requires at least 1 argument.", msg.Action)
		}
		app, err := ensureAppIsStarted(msg)
		if err != nil {
			log.Print(err)
			return
		}
		err = app.Restart(ioutil.Discard)
		if err != nil {
			log.Printf("Error handling %q. App failed to start:\n%s.", msg.Action, err)
			return
		}
		msg.Delete()
	case bindService:
		err := bindUnit(msg)
		if err != nil {
			log.Print(err)
			return
		}
		msg.Delete()
	default:
		log.Printf("Error handling %q: invalid action.", msg.Action)
		msg.Delete()
	}
}

// unitList is a simple slice of units, with special methods to handle state.
type unitList []Unit

// Started returns true if all units in the list is started.
func (l unitList) Started() bool {
	for _, unit := range l {
		if unit.State != string(provision.StatusStarted) {
			return false
		}
	}
	return true
}

// State returns a string if all units have the same state. Otherwise it
// returns an empty string.
func (l unitList) State() string {
	if len(l) == 0 {
		return ""
	}
	state := l[0].State
	for i := 1; i < len(l); i++ {
		if l[i].State != state {
			return ""
		}
	}
	return state
}

// getUnits builds a unitList from the given app and the names in the string
// slice.
func getUnits(a *App, names []string) unitList {
	var units []Unit
	if len(names) > 0 {
		units = make([]Unit, len(names))
		i := 0
		for _, unitName := range names {
			for _, appUnit := range a.Units {
				if appUnit.Name == unitName {
					units[i] = appUnit
					i++
					break
				}
			}
		}
	}
	return unitList(units)
}

var (
	qfactory queue.QFactory
	_queue   queue.Q
	_handler queue.Handler
)

func handler() queue.Handler {
	if _handler != nil {
		return _handler
	}
	var err error
	if qfactory == nil {
		qfactory, err = queue.Factory()
		if err != nil {
			log.Fatalf("Failed to get the queue instance: %s", err)
		}
	}
	_handler, err = qfactory.Handler(handle, queueName, QueueName)
	if err != nil {
		log.Fatalf("Failed to create the queue handler: %s", err)
	}
	return _handler
}

func aqueue() queue.Q {
	if _queue != nil {
		return _queue
	}
	var err error
	if qfactory == nil {
		qfactory, err = queue.Factory()
		if err != nil {
			log.Fatalf("Failed to get the queue instance: %s", err)
		}
	}
	_queue, err = qfactory.Get(queueName)
	if err != nil {
		log.Fatalf("Failed to get the queue: %s", err)
	}
	return _queue
}

// enqueues the given messages and start the handler.
func enqueue(msgs ...queue.Message) {
	for _, msg := range msgs {
		copy := msg
		aqueue().Put(&copy, 0)
	}
	handler().Start()
}
<commit_msg>app: fix data races<commit_after>// Copyright 2013 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package app

import (
	"fmt"
	"github.com/globocom/tsuru/db"
	"github.com/globocom/tsuru/log"
	"github.com/globocom/tsuru/provision"
	"github.com/globocom/tsuru/queue"
	"github.com/globocom/tsuru/service"
	"io/ioutil"
	"labix.org/v2/mgo/bson"
	"sync"
)

const (
	// queue actions
	regenerateApprc         = "regenerate-apprc"
	startApp                = "start-app"
	RegenerateApprcAndStart = "regenerate-apprc-start-app"
	bindService             = "bind-service"

	// name of the queue for internal messages.
	queueName = "tsuru-app"

	// Name of the queue for external messages.
	QueueName = "tsuru-app-public"
)

// ensureAppIsStarted make sure that the app and all units present in the given
// message are started.
func ensureAppIsStarted(msg *queue.Message) (App, error) {
	a := App{Name: msg.Args[0]}
	err := a.Get()
	if err != nil {
		return a, fmt.Errorf("Error handling %q: app %q does not exist.", msg.Action, a.Name)
	}
	units := getUnits(&a, msg.Args[1:])
	if !a.Available() || !units.Started() {
		format := "Error handling %q for the app %q:"
		uState := units.State()
		if uState == "error" || uState == "down" {
			format += fmt.Sprintf(" units are in %q state.", uState)
			msg.Delete()
		} else {
			format += " all units must be started."
		}
		return a, fmt.Errorf(format, msg.Action, a.Name)
	}
	return a, nil
}

// bindUnit handles the bind-service message, binding a unit to all service
// instances binded to the app.
func bindUnit(msg *queue.Message) error {
	a := App{Name: msg.Args[0]}
	err := a.Get()
	if err != nil {
		msg.Delete()
		return fmt.Errorf("Error handling %q: app %q does not exist.", msg.Action, a.Name)
	}
	conn, err := db.Conn()
	if err != nil {
		return fmt.Errorf("Error handling %q: %s", msg.Action, err)
	}
	unit := getUnits(&a, msg.Args[1:])[0]
	var instances []service.ServiceInstance
	q := bson.M{"apps": bson.M{"$in": []string{msg.Args[0]}}}
	err = conn.ServiceInstances().Find(q).All(&instances)
	if err != nil {
		return err
	}
	for _, instance := range instances {
		_, err = instance.BindUnit(&unit)
		if err != nil {
			log.Printf("Error binding the unit %s with the service instance %s.", unit.Name, instance.Name)
		}
	}
	return nil
}

// handle is the function called by the queue handler on each message.
func handle(msg *queue.Message) {
	switch msg.Action {
	case RegenerateApprcAndStart:
		fallthrough
	case regenerateApprc:
		if len(msg.Args) < 1 {
			log.Printf("Error handling %q: this action requires at least 1 argument.", msg.Action)
			return
		}
		app, err := ensureAppIsStarted(msg)
		if err != nil {
			log.Print(err)
			return
		}
		msg.Delete()
		app.serializeEnvVars()
		fallthrough
	case startApp:
		if msg.Action == regenerateApprc {
			break
		}
		if len(msg.Args) < 1 {
			log.Printf("Error handling %q: this action requires at least 1 argument.", msg.Action)
		}
		app, err := ensureAppIsStarted(msg)
		if err != nil {
			log.Print(err)
			return
		}
		err = app.Restart(ioutil.Discard)
		if err != nil {
			log.Printf("Error handling %q. App failed to start:\n%s.", msg.Action, err)
			return
		}
		msg.Delete()
	case bindService:
		err := bindUnit(msg)
		if err != nil {
			log.Print(err)
			return
		}
		msg.Delete()
	default:
		log.Printf("Error handling %q: invalid action.", msg.Action)
		msg.Delete()
	}
}

// unitList is a simple slice of units, with special methods to handle state.
type unitList []Unit

// Started returns true if all units in the list is started.
func (l unitList) Started() bool {
	for _, unit := range l {
		if unit.State != string(provision.StatusStarted) {
			return false
		}
	}
	return true
}

// State returns a string if all units have the same state. Otherwise it
// returns an empty string.
func (l unitList) State() string {
	if len(l) == 0 {
		return ""
	}
	state := l[0].State
	for i := 1; i < len(l); i++ {
		if l[i].State != state {
			return ""
		}
	}
	return state
}

// getUnits builds a unitList from the given app and the names in the string
// slice.
func getUnits(a *App, names []string) unitList {
	var units []Unit
	if len(names) > 0 {
		units = make([]Unit, len(names))
		i := 0
		for _, unitName := range names {
			for _, appUnit := range a.Units {
				if appUnit.Name == unitName {
					units[i] = appUnit
					i++
					break
				}
			}
		}
	}
	return unitList(units)
}

var (
	qfactory queue.QFactory
	_queue   queue.Q
	_handler queue.Handler
	fmut     sync.Mutex
	hmut     sync.Mutex
	qmut     sync.Mutex
)

func handler() queue.Handler {
	// TODO(fss): avoid these mutexes contention.
	hmut.Lock()
	defer hmut.Unlock()
	if _handler != nil {
		return _handler
	}
	var err error
	fmut.Lock()
	if qfactory == nil {
		qfactory, err = queue.Factory()
		if err != nil {
			fmut.Unlock()
			log.Fatalf("Failed to get the queue instance: %s", err)
		}
	}
	_handler, err = qfactory.Handler(handle, queueName, QueueName)
	fmut.Unlock()
	if err != nil {
		log.Fatalf("Failed to create the queue handler: %s", err)
	}
	return _handler
}

func aqueue() queue.Q {
	// TODO(fss): avoid these mutexes contention.
	qmut.Lock()
	defer qmut.Unlock()
	if _queue != nil {
		return _queue
	}
	var err error
	fmut.Lock()
	if qfactory == nil {
		qfactory, err = queue.Factory()
		if err != nil {
			fmut.Unlock()
			log.Fatalf("Failed to get the queue instance: %s", err)
		}
	}
	_queue, err = qfactory.Get(queueName)
	fmut.Unlock()
	if err != nil {
		log.Fatalf("Failed to get the queue: %s", err)
	}
	return _queue
}

// enqueues the given messages and start the handler.
func enqueue(msgs ...queue.Message) {
	for _, msg := range msgs {
		copy := msg
		aqueue().Put(&copy, 0)
	}
	handler().Start()
}
<|endoftext|>


<commit_before>// Copyright 2011 Google Inc. All rights reserved.
// Use of this source code is governed by the Apache 2.0
// license that can be found in the LICENSE file.

// Package appengine provides basic functionality for Google App Engine.
//
// For more information on how to write Go apps for Google App Engine, see:
// https://cloud.google.com/appengine/docs/go/
package appengine // import "google.golang.org/appengine"

import (
	"net/http"

	"github.com/golang/protobuf/proto"
	"golang.org/x/net/context"

	"google.golang.org/appengine/internal"
)

// The gophers party all night; the rabbits provide the beats.

// Main is the principal entry point for an app running in App Engine.
//
// On App Engine Flexible it installs a trivial health checker if one isn't
// already registered, and starts listening on port 8080 (overridden by the
// $PORT environment variable).
//
// See https://cloud.google.com/appengine/docs/flexible/custom-runtimes#health_check_requests
// for details on how to do your own health checking.
//
// On App Engine Standard it ensures the server has started and is prepared to
// receive requests.
//
// Main never returns.
//
// Main is designed so that the app's main package looks like this:
//
//      package main
//
//      import (
//              "google.golang.org/appengine"
//
//              _ "myapp/package0"
//              _ "myapp/package1"
//      )
//
//      func main() {
//              appengine.Main()
//      }
//
// The "myapp/packageX" packages are expected to register HTTP handlers
// in their init functions.
func Main() {
	internal.Main()
}

// IsDevAppServer reports whether the App Engine app is running in the
// development App Server.
func IsDevAppServer() bool {
	return internal.IsDevAppServer()
}

// IsStandard reports whether the App Engine app is running in the standard
// environment. This includes both the first generation runtimes (<= Go 1.9)
// and the second generation runtimes (>= Go 1.11).
func IsStandard() bool {
	return internal.IsStandard()
}

// IsFlex reports whether the App Engine app is running in the flexible environment.
func IsFlex() bool {
	return internal.IsFlex()
}

// IsAppEngine reports whether the App Engine app is running on App Engine, in either
// the standard or flexible environment.
func IsAppEngine() bool {
	return internal.IsAppEngine()
}

// IsSecondGen reports whether the App Engine app is running on the second generation
// runtimes (>= Go 1.11).
func IsSecondGen() bool {
	return internal.IsSecondGen()
}

// NewContext returns a context for an in-flight HTTP request.
// This function is cheap.
func NewContext(req *http.Request) context.Context {
	return internal.ReqContext(req)
}

// WithContext returns a copy of the parent context
// and associates it with an in-flight HTTP request.
// This function is cheap.
func WithContext(parent context.Context, req *http.Request) context.Context {
	return internal.WithContext(parent, req)
}

// TODO(dsymonds): Add a Call function here? Otherwise other packages can't access internal.Call.

// BlobKey is a key for a blobstore blob.
//
// Conceptually, this type belongs in the blobstore package, but it lives in
// the appengine package to avoid a circular dependency: blobstore depends on
// datastore, and datastore needs to refer to the BlobKey type.
type BlobKey string

// GeoPoint represents a location as latitude/longitude in degrees.
type GeoPoint struct {
	Lat, Lng float64
}

// Valid returns whether a GeoPoint is within [-90, 90] latitude and [-180, 180] longitude.
func (g GeoPoint) Valid() bool {
	return -90 <= g.Lat && g.Lat <= 90 && -180 <= g.Lng && g.Lng <= 180
}

// APICallFunc defines a function type for handling an API call.
// See WithCallOverride.
type APICallFunc func(ctx context.Context, service, method string, in, out proto.Message) error

// WithAPICallFunc returns a copy of the parent context
// that will cause API calls to invoke f instead of their normal operation.
//
// This is intended for advanced users only.
func WithAPICallFunc(ctx context.Context, f APICallFunc) context.Context {
	return internal.WithCallOverride(ctx, internal.CallOverrideFunc(f))
}

// APICall performs an API call.
//
// This is not intended for general use; it is exported for use in conjunction
// with WithAPICallFunc.
func APICall(ctx context.Context, service, method string, in, out proto.Message) error {
	return internal.Call(ctx, service, method, in, out)
}
<commit_msg>Remove TODO that has been implemented. (#193)<commit_after>// Copyright 2011 Google Inc. All rights reserved.
// Use of this source code is governed by the Apache 2.0
// license that can be found in the LICENSE file.

// Package appengine provides basic functionality for Google App Engine.
//
// For more information on how to write Go apps for Google App Engine, see:
// https://cloud.google.com/appengine/docs/go/
package appengine // import "google.golang.org/appengine"

import (
	"net/http"

	"github.com/golang/protobuf/proto"
	"golang.org/x/net/context"

	"google.golang.org/appengine/internal"
)

// The gophers party all night; the rabbits provide the beats.

// Main is the principal entry point for an app running in App Engine.
//
// On App Engine Flexible it installs a trivial health checker if one isn't
// already registered, and starts listening on port 8080 (overridden by the
// $PORT environment variable).
//
// See https://cloud.google.com/appengine/docs/flexible/custom-runtimes#health_check_requests
// for details on how to do your own health checking.
//
// On App Engine Standard it ensures the server has started and is prepared to
// receive requests.
//
// Main never returns.
//
// Main is designed so that the app's main package looks like this:
//
//      package main
//
//      import (
//              "google.golang.org/appengine"
//
//              _ "myapp/package0"
//              _ "myapp/package1"
//      )
//
//      func main() {
//              appengine.Main()
//      }
//
// The "myapp/packageX" packages are expected to register HTTP handlers
// in their init functions.
func Main() {
	internal.Main()
}

// IsDevAppServer reports whether the App Engine app is running in the
// development App Server.
func IsDevAppServer() bool {
	return internal.IsDevAppServer()
}

// IsStandard reports whether the App Engine app is running in the standard
// environment. This includes both the first generation runtimes (<= Go 1.9)
// and the second generation runtimes (>= Go 1.11).
func IsStandard() bool {
	return internal.IsStandard()
}

// IsFlex reports whether the App Engine app is running in the flexible environment.
func IsFlex() bool {
	return internal.IsFlex()
}

// IsAppEngine reports whether the App Engine app is running on App Engine, in either
// the standard or flexible environment.
func IsAppEngine() bool {
	return internal.IsAppEngine()
}

// IsSecondGen reports whether the App Engine app is running on the second generation
// runtimes (>= Go 1.11).
func IsSecondGen() bool {
	return internal.IsSecondGen()
}

// NewContext returns a context for an in-flight HTTP request.
// This function is cheap.
func NewContext(req *http.Request) context.Context {
	return internal.ReqContext(req)
}

// WithContext returns a copy of the parent context
// and associates it with an in-flight HTTP request.
// This function is cheap.
func WithContext(parent context.Context, req *http.Request) context.Context {
	return internal.WithContext(parent, req)
}

// BlobKey is a key for a blobstore blob.
//
// Conceptually, this type belongs in the blobstore package, but it lives in
// the appengine package to avoid a circular dependency: blobstore depends on
// datastore, and datastore needs to refer to the BlobKey type.
type BlobKey string

// GeoPoint represents a location as latitude/longitude in degrees.
type GeoPoint struct {
	Lat, Lng float64
}

// Valid returns whether a GeoPoint is within [-90, 90] latitude and [-180, 180] longitude.
func (g GeoPoint) Valid() bool {
	return -90 <= g.Lat && g.Lat <= 90 && -180 <= g.Lng && g.Lng <= 180
}

// APICallFunc defines a function type for handling an API call.
// See WithCallOverride.
type APICallFunc func(ctx context.Context, service, method string, in, out proto.Message) error

// WithAPICallFunc returns a copy of the parent context
// that will cause API calls to invoke f instead of their normal operation.
//
// This is intended for advanced users only.
func WithAPICallFunc(ctx context.Context, f APICallFunc) context.Context {
	return internal.WithCallOverride(ctx, internal.CallOverrideFunc(f))
}

// APICall performs an API call.
//
// This is not intended for general use; it is exported for use in conjunction
// with WithAPICallFunc.
func APICall(ctx context.Context, service, method string, in, out proto.Message) error {
	return internal.Call(ctx, service, method, in, out)
}
<|endoftext|>


<commit_before>package fire

import (
	"time"

	"github.com/gin-gonic/gin"
	"github.com/kr/pretty"
	"github.com/ory-am/fosite"
	"github.com/ory-am/fosite/handler/core"
	"github.com/ory-am/fosite/handler/core/client"
	"github.com/ory-am/fosite/handler/core/implicit"
	"github.com/ory-am/fosite/handler/core/owner"
	"github.com/ory-am/fosite/handler/core/strategy"
	"github.com/ory-am/fosite/token/hmac"
	"gopkg.in/mgo.v2"
)

var callbackTemplate = []byte(`<!DOCTYPE html>

    <title>Authorize</title>

`)

type AccessToken struct {
	Base          `bson:",inline" fire:"access-token:access-tokens:access_tokens"`
	Signature     string    `json:"-" valid:"required"`
	RequestedAt   time.Time `json:"requested-at" valid:"required" bson:"requested_at"`
	GrantedScopes []string  `json:"granted-scopes" valid:"required" bson:"granted_scopes"`
}

type Authenticator struct {
	storage *authenticatorStorage

	strategy     *strategy.HMACSHAStrategy
	handleHelper *core.HandleHelper
	fosite       *fosite.Fosite
}

var accessTokenModel *AccessToken

func init() {
	accessTokenModel = &AccessToken{}
	Init(accessTokenModel)
}

func NewAuthenticator(db *mgo.Database, ownerModel, clientModel Model, secret string) *Authenticator {
	// initialize models
	Init(ownerModel)
	Init(clientModel)

	// extract attributes from owner
	ownerIdentifiable := ownerModel.getBase().attributesByTag("identifiable")
	if len(ownerIdentifiable) != 1 {
		panic("expected to find exactly one 'identifiable' tag on model")
	}
	ownerVerifiable := ownerModel.getBase().attributesByTag("verifiable")
	if len(ownerVerifiable) != 1 {
		panic("expected to find exactly one 'verifiable' tag on model")
	}

	// extract attributes from client
	clientIdentifiable := clientModel.getBase().attributesByTag("identifiable")
	if len(clientIdentifiable) != 1 {
		panic("expected to find exactly one 'identifiable' tag on model")
	}
	clientVerifiable := clientModel.getBase().attributesByTag("verifiable")
	if len(clientVerifiable) != 1 {
		panic("expected to find exactly one 'verifiable' tag on model")
	}
	clientCallable := clientModel.getBase().attributesByTag("callable")
	if len(clientCallable) != 1 {
		panic("expected to find exactly one 'callable' tag on model")
	}

	// create storage
	s := &authenticatorStorage{
		db:                 db,
		ownerModel:         ownerModel,
		ownerIDAttr:        ownerIdentifiable[0],
		ownerSecretAttr:    ownerVerifiable[0],
		clientModel:        clientModel,
		clientIDAttr:       clientIdentifiable[0],
		clientSecretAttr:   clientVerifiable[0],
		clientCallableAttr: clientCallable[0],
	}

	// set the default token lifespan to one hour
	tokenLifespan := time.Hour

	// create a new token generation strategy
	strategy := &strategy.HMACSHAStrategy{
		Enigma: &hmac.HMACStrategy{
			GlobalSecret: []byte(secret),
		},
		AccessTokenLifespan:   tokenLifespan,
		AuthorizeCodeLifespan: tokenLifespan,
	}

	// instantiate a new fosite instance
	f := fosite.NewFosite(s)

	// set mandatory scope
	f.MandatoryScope = "fire"

	// this little helper is used by some of the handlers later
	handleHelper := &core.HandleHelper{
		AccessTokenStrategy: strategy,
		AccessTokenStorage:  s,
		AccessTokenLifespan: tokenLifespan,
	}

	// TODO: Enable refresh token.
	// this handler is responsible for the refresh token grant
	//refreshHandler := &refresh.RefreshTokenGrantHandler{
	//	AccessTokenStrategy:      strategy,
	//	RefreshTokenStrategy:     strategy,
	//	RefreshTokenGrantStorage: s,
	//	AccessTokenLifespan:      accessTokenLifespan,
	//}

	// add handler to fosite
	//f.TokenEndpointHandlers.Append(refreshHandler)

	// add a request validator for access tokens to fosite
	f.AuthorizedRequestValidators.Append(&core.CoreValidator{
		AccessTokenStrategy: strategy,
		AccessTokenStorage:  s,
	})

	return &Authenticator{
		storage:      s,
		fosite:       f,
		handleHelper: handleHelper,
		strategy:     strategy,
	}
}

// EnablePasswordGrant enables the usage of the OAuth 2.0 Resource Owner Password
// Credentials Grant.
func (a *Authenticator) EnablePasswordGrant() {
	// create handler
	passwordHandler := &owner.ResourceOwnerPasswordCredentialsGrantHandler{
		HandleHelper:                                 a.handleHelper,
		ResourceOwnerPasswordCredentialsGrantStorage: a.storage,
	}

	// add handler to fosite
	a.fosite.TokenEndpointHandlers.Append(passwordHandler)
}

// EnableCredentialsGrant enables the usage of the OAuth 2.0 Client Credentials Grant.
func (a *Authenticator) EnableCredentialsGrant() {
	// create handler
	credentialsHandler := &client.ClientCredentialsGrantHandler{
		HandleHelper: a.handleHelper,
	}

	// add handler to fosite
	a.fosite.TokenEndpointHandlers.Append(credentialsHandler)
}

// EnableImplicitGrant enables the usage of the OAuth 2.0 Implicit Grant.
func (a *Authenticator) EnableImplicitGrant() {
	// create handler
	implicitHandler := &implicit.AuthorizeImplicitGrantTypeHandler{
		AccessTokenStrategy: a.handleHelper.AccessTokenStrategy,
		AccessTokenStorage:  a.handleHelper.AccessTokenStorage,
		AccessTokenLifespan: a.handleHelper.AccessTokenLifespan,
	}

	// add handler to fosite
	a.fosite.AuthorizeEndpointHandlers.Append(implicitHandler)
}

// HashPassword returns an Authenticator compatible hash of the password.
func (a *Authenticator) HashPassword(password string) ([]byte, error) {
	return a.fosite.Hasher.Hash([]byte(password))
}

// MustHashPassword is the same as HashPassword except that it raises and error
// when the hashing failed.
func (a *Authenticator) MustHashPassword(password string) []byte {
	bytes, err := a.HashPassword(password)
	if err != nil {
		panic(err)
	}

	return bytes
}

func (a *Authenticator) Register(prefix string, router gin.IRouter) {
	router.POST(prefix+"/token", a.tokenEndpoint)
	router.GET(prefix+"/authorize", a.authorizeEndpoint)
	router.POST(prefix+"/authorize", a.authorizeEndpoint)
	router.GET(prefix+"/callback", a.callbackEndpoint)
}

func (a *Authenticator) Authorizer() Callback {
	return func(ctx *Context) (error, error) {
		// prepare fosite
		f := fosite.NewContext()
		session := &strategy.HMACSession{}

		// validate request
		_, err := a.fosite.ValidateRequestAuthorization(f, ctx.GinContext.Request, session, "fire")
		if err != nil {
			return err, nil
		}

		return nil, nil
	}
}

func (a *Authenticator) tokenEndpoint(ctx *gin.Context) {
	// create new context
	f := fosite.NewContext()

	// create new session
	s := &strategy.HMACSession{}

	// obtain access request
	req, err := a.fosite.NewAccessRequest(f, ctx.Request, s)
	if err != nil {
		pretty.Println(err)
		a.fosite.WriteAccessError(ctx.Writer, req, err)
		return
	}

	// grant the mandatory scope
	if req.GetScopes().Has("fire") {
		req.GrantScope("fire")
	}

	// obtain access response
	res, err := a.fosite.NewAccessResponse(f, ctx.Request, req)
	if err != nil {
		pretty.Println(err)
		a.fosite.WriteAccessError(ctx.Writer, req, err)
		return
	}

	// write response
	a.fosite.WriteAccessResponse(ctx.Writer, req, res)
}

func (a *Authenticator) authorizeEndpoint(ctx *gin.Context) {
	// create new context
	f := fosite.NewContext()

	// obtain authorize request
	req, err := a.fosite.NewAuthorizeRequest(f, ctx.Request)
	if err != nil {
		pretty.Println(err)
		a.fosite.WriteAuthorizeError(ctx.Writer, req, err)
		return
	}
	// You have now access to authorizeRequest, Code ResponseTypes, Scopes ...

	// Normally, this would be the place where you would check if the user is logged in and gives his consent.
	// We're simplifying things and just checking if the request includes a valid username and password
	if ctx.Request.Form.Get("username") != "peter" {
		ctx.Writer.Header().Set("Content-Type", "text/html; charset=utf-8")
		ctx.Writer.Write([]byte(`<h1>Login page</h1>`))
		ctx.Writer.Write([]byte(`

Howdy! This is the log in page. For this example, it is enough to supply the username.

			<form method="post">
				<input type="text" name="username" /> try peter

				<input type="submit">
			</form>
		`))
		return
	}

	// we allow issuing of refresh tokens per default
	if req.GetScopes().Has("offline") {
		req.GrantScope("offline")
	}

	// create new session
	s := &strategy.HMACSession{}

	// obtain authorize response
	res, err := a.fosite.NewAuthorizeResponse(ctx, ctx.Request, req, s)
	if err != nil {
		pretty.Println(err)
		a.fosite.WriteAuthorizeError(ctx.Writer, req, err)
		return
	}

	// write response
	a.fosite.WriteAuthorizeResponse(ctx.Writer, req, res)
}

func (a *Authenticator) callbackEndpoint(ctx *gin.Context) {
	ctx.Writer.Write(callbackTemplate)
}
<commit_msg>redirect back instead of presenting an ui on auth error<commit_after>package fire

import (
	"net/http"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/kr/pretty"
	"github.com/ory-am/fosite"
	"github.com/ory-am/fosite/handler/core"
	"github.com/ory-am/fosite/handler/core/client"
	"github.com/ory-am/fosite/handler/core/implicit"
	"github.com/ory-am/fosite/handler/core/owner"
	"github.com/ory-am/fosite/handler/core/strategy"
	"github.com/ory-am/fosite/token/hmac"
	"gopkg.in/mgo.v2"
)

var callbackTemplate = []byte(`<!DOCTYPE html>

    <title>Authorize</title>

`)

type AccessToken struct {
	Base          `bson:",inline" fire:"access-token:access-tokens:access_tokens"`
	Signature     string    `json:"-" valid:"required"`
	RequestedAt   time.Time `json:"requested-at" valid:"required" bson:"requested_at"`
	GrantedScopes []string  `json:"granted-scopes" valid:"required" bson:"granted_scopes"`
}

type Authenticator struct {
	storage *authenticatorStorage

	strategy     *strategy.HMACSHAStrategy
	handleHelper *core.HandleHelper
	fosite       *fosite.Fosite
}

var accessTokenModel *AccessToken

func init() {
	accessTokenModel = &AccessToken{}
	Init(accessTokenModel)
}

func NewAuthenticator(db *mgo.Database, ownerModel, clientModel Model, secret string) *Authenticator {
	// initialize models
	Init(ownerModel)
	Init(clientModel)

	// extract attributes from owner
	ownerIdentifiable := ownerModel.getBase().attributesByTag("identifiable")
	if len(ownerIdentifiable) != 1 {
		panic("expected to find exactly one 'identifiable' tag on model")
	}
	ownerVerifiable := ownerModel.getBase().attributesByTag("verifiable")
	if len(ownerVerifiable) != 1 {
		panic("expected to find exactly one 'verifiable' tag on model")
	}

	// extract attributes from client
	clientIdentifiable := clientModel.getBase().attributesByTag("identifiable")
	if len(clientIdentifiable) != 1 {
		panic("expected to find exactly one 'identifiable' tag on model")
	}
	clientVerifiable := clientModel.getBase().attributesByTag("verifiable")
	if len(clientVerifiable) != 1 {
		panic("expected to find exactly one 'verifiable' tag on model")
	}
	clientCallable := clientModel.getBase().attributesByTag("callable")
	if len(clientCallable) != 1 {
		panic("expected to find exactly one 'callable' tag on model")
	}

	// create storage
	s := &authenticatorStorage{
		db:                 db,
		ownerModel:         ownerModel,
		ownerIDAttr:        ownerIdentifiable[0],
		ownerSecretAttr:    ownerVerifiable[0],
		clientModel:        clientModel,
		clientIDAttr:       clientIdentifiable[0],
		clientSecretAttr:   clientVerifiable[0],
		clientCallableAttr: clientCallable[0],
	}

	// set the default token lifespan to one hour
	tokenLifespan := time.Hour

	// create a new token generation strategy
	strategy := &strategy.HMACSHAStrategy{
		Enigma: &hmac.HMACStrategy{
			GlobalSecret: []byte(secret),
		},
		AccessTokenLifespan:   tokenLifespan,
		AuthorizeCodeLifespan: tokenLifespan,
	}

	// instantiate a new fosite instance
	f := fosite.NewFosite(s)

	// set mandatory scope
	f.MandatoryScope = "fire"

	// this little helper is used by some of the handlers later
	handleHelper := &core.HandleHelper{
		AccessTokenStrategy: strategy,
		AccessTokenStorage:  s,
		AccessTokenLifespan: tokenLifespan,
	}

	// TODO: Enable refresh token.
	// this handler is responsible for the refresh token grant
	//refreshHandler := &refresh.RefreshTokenGrantHandler{
	//	AccessTokenStrategy:      strategy,
	//	RefreshTokenStrategy:     strategy,
	//	RefreshTokenGrantStorage: s,
	//	AccessTokenLifespan:      accessTokenLifespan,
	//}

	// add handler to fosite
	//f.TokenEndpointHandlers.Append(refreshHandler)

	// add a request validator for access tokens to fosite
	f.AuthorizedRequestValidators.Append(&core.CoreValidator{
		AccessTokenStrategy: strategy,
		AccessTokenStorage:  s,
	})

	return &Authenticator{
		storage:      s,
		fosite:       f,
		handleHelper: handleHelper,
		strategy:     strategy,
	}
}

// EnablePasswordGrant enables the usage of the OAuth 2.0 Resource Owner Password
// Credentials Grant.
func (a *Authenticator) EnablePasswordGrant() {
	// create handler
	passwordHandler := &owner.ResourceOwnerPasswordCredentialsGrantHandler{
		HandleHelper:                                 a.handleHelper,
		ResourceOwnerPasswordCredentialsGrantStorage: a.storage,
	}

	// add handler to fosite
	a.fosite.TokenEndpointHandlers.Append(passwordHandler)
}

// EnableCredentialsGrant enables the usage of the OAuth 2.0 Client Credentials Grant.
func (a *Authenticator) EnableCredentialsGrant() {
	// create handler
	credentialsHandler := &client.ClientCredentialsGrantHandler{
		HandleHelper: a.handleHelper,
	}

	// add handler to fosite
	a.fosite.TokenEndpointHandlers.Append(credentialsHandler)
}

// EnableImplicitGrant enables the usage of the OAuth 2.0 Implicit Grant.
func (a *Authenticator) EnableImplicitGrant() {
	// create handler
	implicitHandler := &implicit.AuthorizeImplicitGrantTypeHandler{
		AccessTokenStrategy: a.handleHelper.AccessTokenStrategy,
		AccessTokenStorage:  a.handleHelper.AccessTokenStorage,
		AccessTokenLifespan: a.handleHelper.AccessTokenLifespan,
	}

	// add handler to fosite
	a.fosite.AuthorizeEndpointHandlers.Append(implicitHandler)
}

// HashPassword returns an Authenticator compatible hash of the password.
func (a *Authenticator) HashPassword(password string) ([]byte, error) {
	return a.fosite.Hasher.Hash([]byte(password))
}

// MustHashPassword is the same as HashPassword except that it raises and error
// when the hashing failed.
func (a *Authenticator) MustHashPassword(password string) []byte {
	bytes, err := a.HashPassword(password)
	if err != nil {
		panic(err)
	}

	return bytes
}

func (a *Authenticator) Register(prefix string, router gin.IRouter) {
	router.POST(prefix+"/token", a.tokenEndpoint)
	router.POST(prefix+"/authorize", a.authorizeEndpoint)
	router.GET(prefix+"/callback", a.callbackEndpoint)
}

func (a *Authenticator) Authorizer() Callback {
	return func(ctx *Context) (error, error) {
		// prepare fosite
		f := fosite.NewContext()
		session := &strategy.HMACSession{}

		// validate request
		_, err := a.fosite.ValidateRequestAuthorization(f, ctx.GinContext.Request, session, "fire")
		if err != nil {
			return err, nil
		}

		return nil, nil
	}
}

func (a *Authenticator) tokenEndpoint(ctx *gin.Context) {
	// create new context
	f := fosite.NewContext()

	// create new session
	s := &strategy.HMACSession{}

	// obtain access request
	req, err := a.fosite.NewAccessRequest(f, ctx.Request, s)
	if err != nil {
		pretty.Println(err)
		a.fosite.WriteAccessError(ctx.Writer, req, err)
		return
	}

	// grant the mandatory scope
	if req.GetScopes().Has("fire") {
		req.GrantScope("fire")
	}

	// obtain access response
	res, err := a.fosite.NewAccessResponse(f, ctx.Request, req)
	if err != nil {
		pretty.Println(err)
		a.fosite.WriteAccessError(ctx.Writer, req, err)
		return
	}

	// write response
	a.fosite.WriteAccessResponse(ctx.Writer, req, res)
}

func (a *Authenticator) authorizeEndpoint(ctx *gin.Context) {
	// create new context
	f := fosite.NewContext()

	// obtain authorize request
	req, err := a.fosite.NewAuthorizeRequest(f, ctx.Request)
	if err != nil {
		pretty.Println(err)
		a.fosite.WriteAuthorizeError(ctx.Writer, req, err)
		return
	}

	// validate user
	if ctx.Request.Form.Get("username") != "peter" {
		uri := ctx.Request.Referer() + "&error=invalid_credentials"
		ctx.Redirect(http.StatusTemporaryRedirect, uri)
		return
	}

	// we allow issuing of refresh tokens per default
	if req.GetScopes().Has("offline") {
		req.GrantScope("offline")
	}

	// create new session
	s := &strategy.HMACSession{}

	// obtain authorize response
	res, err := a.fosite.NewAuthorizeResponse(ctx, ctx.Request, req, s)
	if err != nil {
		pretty.Println(err)
		a.fosite.WriteAuthorizeError(ctx.Writer, req, err)
		return
	}

	// write response
	a.fosite.WriteAuthorizeResponse(ctx.Writer, req, res)
}

func (a *Authenticator) callbackEndpoint(ctx *gin.Context) {
	ctx.Writer.Write(callbackTemplate)
}
<|endoftext|>


<commit_before>/*
NotifyDispatcher attempts to make working with a single Listener easy for a
dynamic set of independent listeners.

Usage

Viz:

    import (
        "github.com/lib/pq"
        "github.com/johto/notifyutils/notifydispatcher"
        "fmt"
        "time"
    )

    func listener(dispatcher *notifydispatcher.NotifyDispatcher) {
        ch := make(chan *pq.Notification, 8)
        err := dispatcher.Listen("listenerchannel", ch)
        if err != nil {
            panic(err)
        }
        for n := range ch {
            if n == nil {
                fmt.Println("lost connection, but we're fine now!")
                continue
            }

            fmt.Println("received notification!")
            // do something with notification
        }
        panic("could not keep up!")
    }

    func main() {
        dispatcher := notifydispatcher.NewNotifyDispatcher(pq.NewListener("", time.Second, time.Minute, nil))
        for i := 0; i < 8; i++ {
            go listener(dispatcher)
        }
        select{}
    }

*/
package notifydispatcher

import (
	"errors"
	"fmt"
	"github.com/lib/pq"
	"sync"
)

var (
	ErrChannelAlreadyActive = errors.New("channel is already active")
	ErrChannelNotActive = errors.New("channel is not active")
)

var errClosed = errors.New("NotifyDispatcher has been closed")

// SlowReaderEliminationStrategy controls the behaviour of the dispatcher in
// case the buffer of a listener's channel is full and attempting to send to it
// would block the dispatcher, preventing it from delivering notifications for
// unrelated listeners.  The default is CloseSlowReaders, but it can be changed
// at any point during a dispatcher's lifespan using
// SetSlowReaderEliminationStrategy.
type SlowReaderEliminationStrategy int
const (
	// When a send would block, the listener's channel is removed from the set
	// of listeners for that notification channel, and the channel is closed.
	// This is the default strategy.
	CloseSlowReaders SlowReaderEliminationStrategy = iota

	// When a send would block, the notification is not delivered.  Delivery is
	// not reattempted.
	NeglectSlowReaders
)

type listenRequest struct {
	channel string
	unlisten bool
}

type NotifyDispatcher struct {
	listener *pq.Listener
	slowReaderEliminationStrategy SlowReaderEliminationStrategy

	listenRequestch chan listenRequest

	lock sync.Mutex
	channels map[string] *listenSet
	closed bool
	// provide an escape hatch for goroutines sending on listenRequestch
	closeChannel chan bool
}

// NewNotifyDispatcher creates a new NotifyDispatcher, using the supplied
// pq.Listener underneath.  The ownership of the Listener is transferred to
// NotifyDispatcher.  You should not use it after calling NewNotifyDispatcher.
func NewNotifyDispatcher(l *pq.Listener) *NotifyDispatcher {
	d := &NotifyDispatcher{
		listener: l,
		slowReaderEliminationStrategy: CloseSlowReaders,
		listenRequestch: make(chan listenRequest, 64),
		channels: make(map[string] *listenSet),
		closeChannel: make(chan bool),
	}
	go d.dispatcherLoop()
	go d.listenRequestHandlerLoop()
	return d
}

// Sets the strategy for mitigating the adverse effects slow readers might have
// on the dispatcher.  See SlowReaderEliminationStrategy.
func (d *NotifyDispatcher) SetSlowReaderEliminationStrategy(strategy SlowReaderEliminationStrategy) {
	d.lock.Lock()
	defer d.lock.Unlock()
	d.slowReaderEliminationStrategy = strategy
}

func (d *NotifyDispatcher) broadcast() {
	reapchans := []string{}

	d.lock.Lock()
	for channel, set := range d.channels {
		if !set.broadcast(d.slowReaderEliminationStrategy) {
			reapchans = append(reapchans, channel)
		}
	}
	d.lock.Unlock()

	for _, ch := range reapchans {
		d.listenRequestch <- listenRequest{ch, true}
	}
}

func (d *NotifyDispatcher) dispatch(n *pq.Notification) {
	reap := false

	d.lock.Lock()
	set, ok := d.channels[n.Channel]
	if ok {
		reap = !set.broadcast(d.slowReaderEliminationStrategy)
	}
	d.lock.Unlock()

	if reap {
		d.listenRequestch <- listenRequest{n.Channel, true}
	}
}

func (d *NotifyDispatcher) dispatcherLoop() {
	for {
		n := <-d.listener.Notify
		if n == nil {
			d.broadcast()
		} else {
			d.dispatch(n)
		}
	}
}

func (d *NotifyDispatcher) execListen(channel string) {
	for {
		err := d.listener.Listen(channel)
		// ErrChannelAlreadyOpen is a valid return value here; we could have
		// abandoned a channel in Unlisten() if the server returned an error
		// for no apparent reason.
		if err == nil ||
		   err == pq.ErrChannelAlreadyOpen {
			break
		}
	}

	d.lock.Lock()
	defer d.lock.Unlock()
	set, ok := d.channels[channel]
	if !ok {
		panic("oops")
	}
	set.setState(listenSetStateActive)
}

func (d *NotifyDispatcher) execUnlisten(channel string) {
	// we don't really care about the error
	_ = d.listener.Unlisten(channel)

	d.lock.Lock()
	set, ok := d.channels[channel]
	if !ok {
		panic("oops")
	}
	if set.state != listenSetStateZombie {
		panic("oops")
	}
	if set.reap() {
		delete(d.channels, channel)
		d.lock.Unlock()
	} else {
		// Couldn't reap the set because it got new listeners while we were
		// waiting for the UNLISTEN to go through.  Re-LISTEN it, but remember
		// to release the lock first.
		d.lock.Unlock()

		d.execListen(channel)
	}
}

func (d *NotifyDispatcher) listenRequestHandlerLoop() {
	for {
		// check closeChannel, just in case we've been closed and there's a
		// backlog of requests
		select {
			case <-d.closeChannel:
				return
			default:
		}

		select {
			case <-d.closeChannel:
				return

			case req := <-d.listenRequestch:
				if req.unlisten {
					d.execUnlisten(req.channel)
				} else {
					d.execListen(req.channel)
				}
		}
	}
}

func (d *NotifyDispatcher) requestListen(channel string, unlisten bool) error {
	select {
		// make sure we don't get stuck here if someone Close()s us
		case <-d.closeChannel:
			return errClosed
		case d.listenRequestch <- listenRequest{channel, unlisten}:
			return nil
	}

	panic("not reached")
}

// Listen adds ch to the set of listeners for notification channel channel.  ch
// should be a buffered channel.  If SlowReaderEliminationStrategy is
// CloseSlowReaders, ch should not already be a listener of a different
// channel.  If ch is already in the set of listeners for channel,
// ErrChannelAlreadyActive is returned.  After Listen has returned, the
// notification channel is open and the dispatcher will attempt to deliver all
// notifications received for that channel to ch.
func (d *NotifyDispatcher) Listen(channel string, ch chan<- *pq.Notification) error {
	d.lock.Lock()
	defer d.lock.Unlock()

	if d.closed {
		return errClosed
	}

	set, ok := d.channels[channel]
	if ok {
		err := set.add(ch)
		if err != nil {
			return err
		}
	} else {
		set = d.newListenSet(ch)
		d.channels[channel] = set

		// must not be holding the lock while requesting a listen
		d.lock.Unlock()
		err := d.requestListen(channel, false)
		if err != nil {
			return err
		}
		d.lock.Lock()
	}

	return set.waitForActive(&d.closed)
}

// Removes ch from the set of listeners for notification channel channel.  If
// ch is not in the set of listeners for channel, ErrChannelNotActive is
// returned.
func (d *NotifyDispatcher) Unlisten(channel string, ch chan<- *pq.Notification) error {
	d.lock.Lock()

	if d.closed {
		d.lock.Unlock()
		return errClosed
	}

	set, ok := d.channels[channel]
	if !ok {
		d.lock.Unlock()
		return ErrChannelNotActive
	}
	last, err := set.remove(ch)
	if err != nil {
		d.lock.Unlock()
		return err
	}
	if !last {
		// the set isn't empty; nothing further for us to do
		d.lock.Unlock()
		return nil
	}
	d.lock.Unlock()

	// we were the last listener, cue the reaper
	return d.requestListen(channel, true)
}

func (d *NotifyDispatcher) Close() error {
	d.lock.Lock()
	defer d.lock.Unlock()

	if d.closed {
		return errClosed
	}

	d.closed = true
	close(d.closeChannel)

	for _, set := range d.channels {
		set.activeOrClosedCond.Broadcast()
	}

	// TODO: kill the dispatcher as well

	return nil
}

type listenSetState int
const (
	// The set was recently spawned or respawned, and it's waiting for a call
	// to Listen() to succeed.
	listenSetStateNewborn listenSetState = iota
	// The set is ready and any notifications from the database will be
	// dispatched to the set.
	listenSetStateActive
	// The set has recently been emptied, and it's waiting for a call to
	// Unlisten() to finish.
	listenSetStateZombie
)

type listenSet struct {
	channels map[chan<- *pq.Notification] struct{}
	state listenSetState
	activeOrClosedCond *sync.Cond
}

func (d *NotifyDispatcher) newListenSet(firstInhabitant chan<- *pq.Notification) *listenSet {
	s := &listenSet{
		channels: make(map[chan<- *pq.Notification] struct{}),
		state: listenSetStateNewborn,
	}
	s.activeOrClosedCond = sync.NewCond(&d.lock)
	s.channels[firstInhabitant] = struct{}{}
	return s
}

func (s *listenSet) setState(newState listenSetState) {
	var expectedState listenSetState
	switch newState {
		case listenSetStateNewborn:
			expectedState = listenSetStateZombie
		case listenSetStateActive:
			expectedState = listenSetStateNewborn
		case listenSetStateZombie:
			expectedState = listenSetStateActive
	}
	if s.state != expectedState {
		panic(fmt.Sprintf("illegal state transition from %v to %v", s.state, newState))
	}
	s.state = newState
	if s.state == listenSetStateActive {
		s.activeOrClosedCond.Broadcast()
	}
}

func (s *listenSet) add(ch chan<- *pq.Notification) error {
	_, ok := s.channels[ch]
	if ok {
		return ErrChannelAlreadyActive
	}
	s.channels[ch] = struct{}{}
	return nil
}

func (s *listenSet) remove(ch chan<- *pq.Notification) (last bool, err error) {
	_, ok := s.channels[ch]
	if !ok {
		return false, ErrChannelNotActive
	}
	delete(s.channels, ch)

	if len(s.channels) == 0 {
		s.setState(listenSetStateZombie)
		return true, nil
	}
	return false, nil
}

func (s *listenSet) broadcast(strategy SlowReaderEliminationStrategy) bool {
	// must be active
	if s.state != listenSetStateActive {
		return true
	}

	for ch := range s.channels {
		select {
			case ch <- nil:

			default:
				if strategy == CloseSlowReaders {
					delete(s.channels, ch)
					close(ch)
				}
		}
	}

	if len(s.channels) == 0 {
		s.setState(listenSetStateZombie)
		return false
	}

	return true
}

// Marks the set active after a successful call to Listen().
func (s *listenSet) markActive() {
	s.setState(listenSetStateActive)
}

// Wait for the listen set to become "active".  Returns nil if successfull, or
// errClosed if the dispatcher was closed while waiting.  The caller should be
// holding d.lock.
func (s *listenSet) waitForActive(closed *bool) error {
	for {
		if *closed {
			return errClosed
		}
		if s.state == listenSetStateActive {
			return nil
		}
		s.activeOrClosedCond.Wait()
	}
}

// Try to reap a zombie set after Unlisten().  Returns true if the set should
// be removed, false otherwise.
func (s *listenSet) reap() bool {
	if s.state != listenSetStateZombie {
		panic("unexpected state in reap")
	}

	if len(s.channels) > 0 {
		// we need to be respawned
		s.setState(listenSetStateNewborn)
		return false
	}

	return true
}
<commit_msg>Add broadcasting alternatives to NotifyDispatcher<commit_after>/*
NotifyDispatcher attempts to make working with a single Listener easy for a
dynamic set of independent listeners.

Usage

Viz:

    import (
        "github.com/lib/pq"
        "github.com/johto/notifyutils/notifydispatcher"
        "fmt"
        "time"
    )

    func listener(dispatcher *notifydispatcher.NotifyDispatcher) {
        ch := make(chan *pq.Notification, 8)
        err := dispatcher.Listen("listenerchannel", ch)
        if err != nil {
            panic(err)
        }
        for n := range ch {
            if n == nil {
                fmt.Println("lost connection, but we're fine now!")
                continue
            }

            fmt.Println("received notification!")
            // do something with notification
        }
        panic("could not keep up!")
    }

    func main() {
        dispatcher := notifydispatcher.NewNotifyDispatcher(pq.NewListener("", time.Second, time.Minute, nil))
        for i := 0; i < 8; i++ {
            go listener(dispatcher)
        }
        select{}
    }

*/
package notifydispatcher

import (
	"container/list"
	"errors"
	"fmt"
	"github.com/lib/pq"
	"sync"
)

var (
	ErrChannelAlreadyActive = errors.New("channel is already active")
	ErrChannelNotActive = errors.New("channel is not active")
)

var errClosed = errors.New("NotifyDispatcher has been closed")

// SlowReaderEliminationStrategy controls the behaviour of the dispatcher in
// case the buffer of a listener's channel is full and attempting to send to it
// would block the dispatcher, preventing it from delivering notifications for
// unrelated listeners.  The default is CloseSlowReaders, but it can be changed
// at any point during a dispatcher's lifespan using
// SetSlowReaderEliminationStrategy.
type SlowReaderEliminationStrategy int
const (
	// When a send would block, the listener's channel is removed from the set
	// of listeners for that notification channel, and the channel is closed.
	// This is the default strategy.
	CloseSlowReaders SlowReaderEliminationStrategy = iota

	// When a send would block, the notification is not delivered.  Delivery is
	// not reattempted.
	NeglectSlowReaders
)

type listenRequest struct {
	channel string
	unlisten bool
}

type BroadcastChannel struct {
	Channel chan struct{}
	elem *list.Element
}

type NotifyDispatcher struct {
	listener *pq.Listener

	// Some details about the behaviour.  Only touch or look at while holding
	// "lock".
	slowReaderEliminationStrategy SlowReaderEliminationStrategy
	broadcastOnConnectionLoss bool
	broadcastChannels *list.List

	listenRequestch chan listenRequest

	lock sync.Mutex
	channels map[string] *listenSet
	closed bool
	// provide an escape hatch for goroutines sending on listenRequestch
	closeChannel chan bool
}

// NewNotifyDispatcher creates a new NotifyDispatcher, using the supplied
// pq.Listener underneath.  The ownership of the Listener is transferred to
// NotifyDispatcher.  You should not use it after calling NewNotifyDispatcher.
func NewNotifyDispatcher(l *pq.Listener) *NotifyDispatcher {
	d := &NotifyDispatcher{
		listener: l,
		slowReaderEliminationStrategy: CloseSlowReaders,
		broadcastOnConnectionLoss: true,
		broadcastChannels: list.New(),
		listenRequestch: make(chan listenRequest, 64),
		channels: make(map[string] *listenSet),
		closeChannel: make(chan bool),
	}
	go d.dispatcherLoop()
	go d.listenRequestHandlerLoop()
	return d
}

// Sets the strategy for mitigating the adverse effects slow readers might have
// on the dispatcher.  See SlowReaderEliminationStrategy.
func (d *NotifyDispatcher) SetSlowReaderEliminationStrategy(strategy SlowReaderEliminationStrategy) {
	d.lock.Lock()
	defer d.lock.Unlock()
	d.slowReaderEliminationStrategy = strategy
}

// Controls whether a nil notification from the underlying Listener is
// broadcast to all channels in the set.
func (d *NotifyDispatcher) SetBroadcastOnConnectionLoss(value bool) {
	d.lock.Lock()
	defer d.lock.Unlock()
	d.broadcastOnConnectionLoss = value
}

// Opens a new "broadcast channel".  A broadcast channel is sent to by the
// NotifyDispatcher every time the underlying Listener has re-established its
// server connection.
func (d *NotifyDispatcher) OpenBroadcastChannel() BroadcastChannel {
	d.lock.Lock()
	defer d.lock.Unlock()

	ch := BroadcastChannel{
		Channel: make(chan struct{}, 1),
	}
	ch.elem = d.broadcastChannels.PushFront(ch)
	return ch
}

// Closes the broadcast channel ch.
func (d *NotifyDispatcher) CloseBroadcastChannel(ch BroadcastChannel) {
	d.lock.Lock()
	defer d.lock.Unlock()

	if d.broadcastChannels.Remove(ch.elem) != ch {
		panic("oops")
	}
	close(ch.Channel)
}

// Broadcast on all broadcastChannels and on all channels unless
// broadcastOnConnectionLoss is disabled.
func (d *NotifyDispatcher) maybeBroadcast() {
	reapchans := []string{}

	d.lock.Lock()
	for e := d.broadcastChannels.Front(); e != nil; e = e.Next() {
		select {
			case e.Value.(BroadcastChannel).Channel <- struct{}{}:
			default:
				// nothing to do
		}
	}

	if !d.broadcastOnConnectionLoss {
		d.lock.Unlock()
		return
	}

	for channel, set := range d.channels {
		if !set.broadcast(d.slowReaderEliminationStrategy) {
			reapchans = append(reapchans, channel)
		}
	}
	d.lock.Unlock()

	for _, ch := range reapchans {
		d.listenRequestch <- listenRequest{ch, true}
	}
}

func (d *NotifyDispatcher) dispatch(n *pq.Notification) {
	reap := false

	d.lock.Lock()
	set, ok := d.channels[n.Channel]
	if ok {
		reap = !set.broadcast(d.slowReaderEliminationStrategy)
	}
	d.lock.Unlock()

	if reap {
		d.listenRequestch <- listenRequest{n.Channel, true}
	}
}

func (d *NotifyDispatcher) dispatcherLoop() {
	for {
		n := <-d.listener.Notify
		if n == nil {
			d.maybeBroadcast()
		} else {
			d.dispatch(n)
		}
	}
}

// Attempt to start listening on channel.  The caller should not be holding
// lock.
func (d *NotifyDispatcher) execListen(channel string) {
	for {
		err := d.listener.Listen(channel)
		// ErrChannelAlreadyOpen is a valid return value here; we could have
		// abandoned a channel in Unlisten() if the server returned an error
		// for no apparent reason.
		if err == nil ||
		   err == pq.ErrChannelAlreadyOpen {
			break
		}
	}

	d.lock.Lock()
	defer d.lock.Unlock()
	set, ok := d.channels[channel]
	if !ok {
		panic("oops")
	}
	set.setState(listenSetStateActive)
}

func (d *NotifyDispatcher) execUnlisten(channel string) {
	// we don't really care about the error
	_ = d.listener.Unlisten(channel)

	d.lock.Lock()
	set, ok := d.channels[channel]
	if !ok {
		panic("oops")
	}
	if set.state != listenSetStateZombie {
		panic("oops")
	}
	if set.reap() {
		delete(d.channels, channel)
		d.lock.Unlock()
	} else {
		// Couldn't reap the set because it got new listeners while we were
		// waiting for the UNLISTEN to go through.  Re-LISTEN it, but remember
		// to release the lock first.
		d.lock.Unlock()

		d.execListen(channel)
	}
}

func (d *NotifyDispatcher) listenRequestHandlerLoop() {
	for {
		// check closeChannel, just in case we've been closed and there's a
		// backlog of requests
		select {
			case <-d.closeChannel:
				return
			default:
		}

		select {
			case <-d.closeChannel:
				return

			case req := <-d.listenRequestch:
				if req.unlisten {
					d.execUnlisten(req.channel)
				} else {
					d.execListen(req.channel)
				}
		}
	}
}

func (d *NotifyDispatcher) requestListen(channel string, unlisten bool) error {
	select {
		// make sure we don't get stuck here if someone Close()s us
		case <-d.closeChannel:
			return errClosed
		case d.listenRequestch <- listenRequest{channel, unlisten}:
			return nil
	}

	panic("not reached")
}

// Listen adds ch to the set of listeners for notification channel channel.  ch
// should be a buffered channel.  If ch is already in the set of listeners for
// channel, ErrChannelAlreadyActive is returned.  After Listen has returned,
// the notification channel is open and the dispatcher will attempt to deliver
// all notifications received for that channel to ch.
//
// If SlowReaderEliminationStrategy is CloseSlowReaders, reusing the same ch
// for multiple notification channels is not allowed.
func (d *NotifyDispatcher) Listen(channel string, ch chan<- *pq.Notification) error {
	d.lock.Lock()
	defer d.lock.Unlock()

	if d.closed {
		return errClosed
	}

	set, ok := d.channels[channel]
	if ok {
		err := set.add(ch)
		if err != nil {
			return err
		}
	} else {
		set = d.newListenSet(ch)
		d.channels[channel] = set

		// must not be holding the lock while requesting a listen
		d.lock.Unlock()
		err := d.requestListen(channel, false)
		if err != nil {
			return err
		}
		d.lock.Lock()
	}

	return set.waitForActive(&d.closed)
}

// Removes ch from the set of listeners for notification channel channel.  If
// ch is not in the set of listeners for channel, ErrChannelNotActive is
// returned.
func (d *NotifyDispatcher) Unlisten(channel string, ch chan<- *pq.Notification) error {
	d.lock.Lock()

	if d.closed {
		d.lock.Unlock()
		return errClosed
	}

	set, ok := d.channels[channel]
	if !ok {
		d.lock.Unlock()
		return ErrChannelNotActive
	}
	last, err := set.remove(ch)
	if err != nil {
		d.lock.Unlock()
		return err
	}
	if !last {
		// the set isn't empty; nothing further for us to do
		d.lock.Unlock()
		return nil
	}
	d.lock.Unlock()

	// we were the last listener, cue the reaper
	return d.requestListen(channel, true)
}

func (d *NotifyDispatcher) Close() error {
	d.lock.Lock()
	defer d.lock.Unlock()

	if d.closed {
		return errClosed
	}

	d.closed = true
	close(d.closeChannel)

	for _, set := range d.channels {
		set.activeOrClosedCond.Broadcast()
	}

	// TODO: kill the dispatcher as well

	return nil
}

type listenSetState int
const (
	// The set was recently spawned or respawned, and it's waiting for a call
	// to Listen() to succeed.
	listenSetStateNewborn listenSetState = iota
	// The set is ready and any notifications from the database will be
	// dispatched to the set.
	listenSetStateActive
	// The set has recently been emptied, and it's waiting for a call to
	// Unlisten() to finish.
	listenSetStateZombie
)

type listenSet struct {
	channels map[chan<- *pq.Notification] struct{}
	state listenSetState
	activeOrClosedCond *sync.Cond
}

func (d *NotifyDispatcher) newListenSet(firstInhabitant chan<- *pq.Notification) *listenSet {
	s := &listenSet{
		channels: make(map[chan<- *pq.Notification] struct{}),
		state: listenSetStateNewborn,
	}
	s.activeOrClosedCond = sync.NewCond(&d.lock)
	s.channels[firstInhabitant] = struct{}{}
	return s
}

func (s *listenSet) setState(newState listenSetState) {
	var expectedState listenSetState
	switch newState {
		case listenSetStateNewborn:
			expectedState = listenSetStateZombie
		case listenSetStateActive:
			expectedState = listenSetStateNewborn
		case listenSetStateZombie:
			expectedState = listenSetStateActive
	}
	if s.state != expectedState {
		panic(fmt.Sprintf("illegal state transition from %v to %v", s.state, newState))
	}
	s.state = newState
	if s.state == listenSetStateActive {
		s.activeOrClosedCond.Broadcast()
	}
}

func (s *listenSet) add(ch chan<- *pq.Notification) error {
	_, ok := s.channels[ch]
	if ok {
		return ErrChannelAlreadyActive
	}
	s.channels[ch] = struct{}{}
	return nil
}

func (s *listenSet) remove(ch chan<- *pq.Notification) (last bool, err error) {
	_, ok := s.channels[ch]
	if !ok {
		return false, ErrChannelNotActive
	}
	delete(s.channels, ch)

	if len(s.channels) == 0 {
		s.setState(listenSetStateZombie)
		return true, nil
	}
	return false, nil
}

func (s *listenSet) broadcast(strategy SlowReaderEliminationStrategy) bool {
	// must be active
	if s.state != listenSetStateActive {
		return true
	}

	for ch := range s.channels {
		select {
			case ch <- nil:

			default:
				if strategy == CloseSlowReaders {
					delete(s.channels, ch)
					close(ch)
				}
		}
	}

	if len(s.channels) == 0 {
		s.setState(listenSetStateZombie)
		return false
	}

	return true
}

// Marks the set active after a successful call to Listen().
func (s *listenSet) markActive() {
	s.setState(listenSetStateActive)
}

// Wait for the listen set to become "active".  Returns nil if successfull, or
// errClosed if the dispatcher was closed while waiting.  The caller should be
// holding d.lock.
func (s *listenSet) waitForActive(closed *bool) error {
	for {
		if *closed {
			return errClosed
		}
		if s.state == listenSetStateActive {
			return nil
		}
		s.activeOrClosedCond.Wait()
	}
}

// Try to reap a zombie set after Unlisten().  Returns true if the set should
// be removed, false otherwise.
func (s *listenSet) reap() bool {
	if s.state != listenSetStateZombie {
		panic("unexpected state in reap")
	}

	if len(s.channels) > 0 {
		// we need to be respawned
		s.setState(listenSetStateNewborn)
		return false
	}

	return true
}
<|endoftext|>


<commit_before>package mpb

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"
	"time"
	"unicode/utf8"

	"github.com/vbauerster/mpb/v4/decor"
)

// Filler interface.
// Bar renders by calling Filler's Fill method. You can literally have
// any bar kind, by implementing this interface and passing it to the
// Add method.
type Filler interface {
	Fill(w io.Writer, width int, stat *decor.Statistics)
}

// FillerFunc is function type adapter to convert function into Filler.
type FillerFunc func(w io.Writer, width int, stat *decor.Statistics)

func (f FillerFunc) Fill(w io.Writer, width int, stat *decor.Statistics) {
	f(w, width, stat)
}

// Bar represents a progress Bar.
type Bar struct {
	priority int
	index    int

	runningBar     *Bar
	cacheState     *bState
	operateState   chan func(*bState)
	bFrameCh       chan *bFrame
	syncTableCh    chan [][]chan int
	completed      chan bool
	forceRefreshCh chan time.Time

	// done is closed by Bar's goroutine, after cacheState is written
	done chan struct{}
	// shutdown is closed from master Progress goroutine only
	shutdown chan struct{}

	arbitraryCurrent struct {
		lock    uint32
		current int64
	}
}

type (
	bState struct {
		filler             Filler
		extender           Filler
		id                 int
		width              int
		total              int64
		current            int64
		trimSpace          bool
		toComplete         bool
		removeOnComplete   bool
		barClearOnComplete bool
		completeFlushed    bool
		aDecorators        []decor.Decorator
		pDecorators        []decor.Decorator
		amountReceivers    []decor.AmountReceiver
		shutdownListeners  []decor.ShutdownListener
		bufP, bufB, bufA   *bytes.Buffer
		bufE               *bytes.Buffer
		panicMsg           string

		// following options are assigned to the *Bar
		priority   int
		runningBar *Bar
	}
	bFrame struct {
		rd               io.Reader
		extendedLines    int
		toShutdown       bool
		removeOnComplete bool
	}
)

func newBar(
	ctx context.Context,
	wg *sync.WaitGroup,
	forceRefreshCh chan time.Time,
	filler Filler,
	id, width int,
	total int64,
	options ...BarOption,
) *Bar {
	if total <= 0 {
		total = time.Now().Unix()
	}

	s := &bState{
		filler:   filler,
		id:       id,
		priority: id,
		width:    width,
		total:    total,
	}

	for _, opt := range options {
		if opt != nil {
			opt(s)
		}
	}

	s.bufP = bytes.NewBuffer(make([]byte, 0, width))
	s.bufB = bytes.NewBuffer(make([]byte, 0, width))
	s.bufA = bytes.NewBuffer(make([]byte, 0, width))
	if s.extender != nil {
		s.bufE = bytes.NewBuffer(make([]byte, 0, width))
	}

	b := &Bar{
		priority:       s.priority,
		runningBar:     s.runningBar,
		operateState:   make(chan func(*bState)),
		bFrameCh:       make(chan *bFrame, 1),
		syncTableCh:    make(chan [][]chan int),
		completed:      make(chan bool),
		done:           make(chan struct{}),
		shutdown:       make(chan struct{}),
		forceRefreshCh: forceRefreshCh,
	}

	if b.runningBar != nil {
		b.priority = b.runningBar.priority
	}

	go b.serve(ctx, wg, s)
	return b
}

// RemoveAllPrependers removes all prepend functions.
func (b *Bar) RemoveAllPrependers() {
	select {
	case b.operateState <- func(s *bState) { s.pDecorators = nil }:
	case <-b.done:
	}
}

// RemoveAllAppenders removes all append functions.
func (b *Bar) RemoveAllAppenders() {
	select {
	case b.operateState <- func(s *bState) { s.aDecorators = nil }:
	case <-b.done:
	}
}

// ProxyReader wraps r with metrics required for progress tracking.
func (b *Bar) ProxyReader(r io.Reader) io.ReadCloser {
	if r == nil {
		panic("expect io.Reader, got nil")
	}
	rc, ok := r.(io.ReadCloser)
	if !ok {
		rc = ioutil.NopCloser(r)
	}
	return &proxyReader{rc, b, time.Now()}
}

// ID returs id of the bar.
func (b *Bar) ID() int {
	result := make(chan int)
	select {
	case b.operateState <- func(s *bState) { result <- s.id }:
		return <-result
	case <-b.done:
		return b.cacheState.id
	}
}

// Current returns bar's current number, in other words sum of all increments.
func (b *Bar) Current() int64 {
	result := make(chan int64)
	select {
	case b.operateState <- func(s *bState) { result <- s.current }:
		return <-result
	case <-b.done:
		return b.cacheState.current
	}
}

// SetRefill sets refill, if supported by underlying Filler.
func (b *Bar) SetRefill(upto int) {
	b.operateState <- func(s *bState) {
		if f, ok := s.filler.(interface{ SetRefill(int) }); ok {
			f.SetRefill(upto)
		}
	}
}

// SetTotal sets total dynamically.
// Set final to true, when total is known, it will trigger bar complete event.
func (b *Bar) SetTotal(total int64, final bool) bool {
	select {
	case b.operateState <- func(s *bState) {
		if total > 0 {
			s.total = total
		}
		if final {
			s.current = s.total
			s.toComplete = true
			go b.forceRefresh()
		}
	}:
		return true
	case <-b.done:
		return false
	}
}

// SetCurrent sets progress' current to arbitrary amount.
func (b *Bar) SetCurrent(current int64, wdd ...time.Duration) {
	if current <= 0 {
		return
	}
	for !atomic.CompareAndSwapUint32(&b.arbitraryCurrent.lock, 0, 1) {
		runtime.Gosched()
	}
	last := b.arbitraryCurrent.current
	b.IncrBy(int(current-last), wdd...)
	b.arbitraryCurrent.current = current
	atomic.StoreUint32(&b.arbitraryCurrent.lock, 0)
}

// Increment is a shorthand for b.IncrBy(1).
func (b *Bar) Increment() {
	b.IncrBy(1)
}

// IncrBy increments progress bar by amount of n.
// wdd is optional work duration i.e. time.Since(start), which expected
// to be provided, if any ewma based decorator is used.
func (b *Bar) IncrBy(n int, wdd ...time.Duration) {
	select {
	case b.operateState <- func(s *bState) {
		s.current += int64(n)
		if s.current >= s.total {
			s.current = s.total
			s.toComplete = true
			go b.forceRefresh()
		}
		for _, ar := range s.amountReceivers {
			ar.NextAmount(n, wdd...)
		}
	}:
	case <-b.done:
	}
}

// Completed reports whether the bar is in completed state.
func (b *Bar) Completed() bool {
	// omit select here, because primary usage of the method is for loop
	// condition, like for !bar.Completed() {...} so when toComplete=true
	// it is called once (at which time, the bar is still alive), then
	// quits the loop and never suppose to be called afterwards.
	return <-b.completed
}

func (b *Bar) wSyncTable() [][]chan int {
	select {
	case b.operateState <- func(s *bState) { b.syncTableCh <- s.wSyncTable() }:
		return <-b.syncTableCh
	case <-b.done:
		return b.cacheState.wSyncTable()
	}
}

func (b *Bar) serve(ctx context.Context, wg *sync.WaitGroup, s *bState) {
	defer wg.Done()
	cancel := ctx.Done()
	for {
		select {
		case op := <-b.operateState:
			op(s)
		case b.completed <- s.toComplete:
		case <-cancel:
			s.toComplete = true
			cancel = nil
		case <-b.shutdown:
			b.cacheState = s
			close(b.done)
			// Notifying decorators about shutdown event
			for _, sl := range s.shutdownListeners {
				sl.Shutdown()
			}
			return
		}
	}
}

func (b *Bar) render(debugOut io.Writer, tw int) {
	select {
	case b.operateState <- func(s *bState) {
		defer func() {
			// recovering if user defined decorator panics for example
			if p := recover(); p != nil {
				s.panicMsg = fmt.Sprintf("panic: %v", p)
				fmt.Fprintf(debugOut, "%s %s bar id %02d %v\n", "[mpb]", time.Now(), s.id, s.panicMsg)
				b.bFrameCh <- &bFrame{
					rd:         strings.NewReader(fmt.Sprintf(fmt.Sprintf("%%.%ds\n", tw), s.panicMsg)),
					toShutdown: true,
				}
			}
		}()
		r := s.draw(tw)
		var extendedLines int
		if s.extender != nil {
			s.extender.Fill(s.bufE, tw, newStatistics(s))
			extendedLines = countLines(s.bufE.Bytes())
			r = io.MultiReader(r, s.bufE)
		}
		b.bFrameCh <- &bFrame{
			rd:               r,
			extendedLines:    extendedLines,
			toShutdown:       s.toComplete && !s.completeFlushed,
			removeOnComplete: s.removeOnComplete,
		}
		s.completeFlushed = s.toComplete
	}:
	case <-b.done:
		s := b.cacheState
		r := s.draw(tw)
		var extendedLines int
		if s.extender != nil {
			s.extender.Fill(s.bufE, tw, newStatistics(s))
			extendedLines = countLines(s.bufE.Bytes())
			r = io.MultiReader(r, s.bufE)
		}
		b.bFrameCh <- &bFrame{
			rd:            r,
			extendedLines: extendedLines,
		}
	}
}

func (s *bState) draw(termWidth int) io.Reader {
	if s.panicMsg != "" {
		return strings.NewReader(fmt.Sprintf(fmt.Sprintf("%%.%ds\n", termWidth), s.panicMsg))
	}

	stat := newStatistics(s)

	for _, d := range s.pDecorators {
		s.bufP.WriteString(d.Decor(stat))
	}

	for _, d := range s.aDecorators {
		s.bufA.WriteString(d.Decor(stat))
	}

	if s.barClearOnComplete && s.completeFlushed {
		s.bufA.WriteByte('\n')
		return io.MultiReader(s.bufP, s.bufA)
	}

	prependCount := utf8.RuneCount(s.bufP.Bytes())
	appendCount := utf8.RuneCount(s.bufA.Bytes())

	if !s.trimSpace {
		// reserve space for edge spaces
		termWidth -= 2
		s.bufB.WriteByte(' ')
	}

	calcWidth := s.width
	if prependCount+s.width+appendCount > termWidth {
		calcWidth = termWidth - prependCount - appendCount
	}
	s.filler.Fill(s.bufB, calcWidth, stat)

	if !s.trimSpace {
		s.bufB.WriteByte(' ')
	}

	s.bufA.WriteByte('\n')
	return io.MultiReader(s.bufP, s.bufB, s.bufA)
}

func (s *bState) wSyncTable() [][]chan int {
	columns := make([]chan int, 0, len(s.pDecorators)+len(s.aDecorators))
	var pCount int
	for _, d := range s.pDecorators {
		if ch, ok := d.Sync(); ok {
			columns = append(columns, ch)
			pCount++
		}
	}
	var aCount int
	for _, d := range s.aDecorators {
		if ch, ok := d.Sync(); ok {
			columns = append(columns, ch)
			aCount++
		}
	}
	table := make([][]chan int, 2)
	table[0] = columns[0:pCount]
	table[1] = columns[pCount : pCount+aCount : pCount+aCount]
	return table
}

func (b *Bar) forceRefresh() {
	select {
	case b.forceRefreshCh <- time.Now():
	case <-b.shutdown:
	}
}

func newStatistics(s *bState) *decor.Statistics {
	return &decor.Statistics{
		ID:        s.id,
		Completed: s.completeFlushed,
		Total:     s.total,
		Current:   s.current,
	}
}

func countLines(b []byte) int {
	return bytes.Count(b, []byte("\n"))
}
<commit_msg>SetTotal sets completed once<commit_after>package mpb

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"
	"time"
	"unicode/utf8"

	"github.com/vbauerster/mpb/v4/decor"
)

// Filler interface.
// Bar renders by calling Filler's Fill method. You can literally have
// any bar kind, by implementing this interface and passing it to the
// Add method.
type Filler interface {
	Fill(w io.Writer, width int, stat *decor.Statistics)
}

// FillerFunc is function type adapter to convert function into Filler.
type FillerFunc func(w io.Writer, width int, stat *decor.Statistics)

func (f FillerFunc) Fill(w io.Writer, width int, stat *decor.Statistics) {
	f(w, width, stat)
}

// Bar represents a progress Bar.
type Bar struct {
	priority int
	index    int

	runningBar     *Bar
	cacheState     *bState
	operateState   chan func(*bState)
	bFrameCh       chan *bFrame
	syncTableCh    chan [][]chan int
	completed      chan bool
	forceRefreshCh chan time.Time

	// done is closed by Bar's goroutine, after cacheState is written
	done chan struct{}
	// shutdown is closed from master Progress goroutine only
	shutdown chan struct{}

	arbitraryCurrent struct {
		lock    uint32
		current int64
	}
}

type (
	bState struct {
		filler             Filler
		extender           Filler
		id                 int
		width              int
		total              int64
		current            int64
		trimSpace          bool
		toComplete         bool
		removeOnComplete   bool
		barClearOnComplete bool
		completeFlushed    bool
		aDecorators        []decor.Decorator
		pDecorators        []decor.Decorator
		amountReceivers    []decor.AmountReceiver
		shutdownListeners  []decor.ShutdownListener
		bufP, bufB, bufA   *bytes.Buffer
		bufE               *bytes.Buffer
		panicMsg           string

		// following options are assigned to the *Bar
		priority   int
		runningBar *Bar
	}
	bFrame struct {
		rd               io.Reader
		extendedLines    int
		toShutdown       bool
		removeOnComplete bool
	}
)

func newBar(
	ctx context.Context,
	wg *sync.WaitGroup,
	forceRefreshCh chan time.Time,
	filler Filler,
	id, width int,
	total int64,
	options ...BarOption,
) *Bar {
	if total <= 0 {
		total = time.Now().Unix()
	}

	s := &bState{
		filler:   filler,
		id:       id,
		priority: id,
		width:    width,
		total:    total,
	}

	for _, opt := range options {
		if opt != nil {
			opt(s)
		}
	}

	s.bufP = bytes.NewBuffer(make([]byte, 0, width))
	s.bufB = bytes.NewBuffer(make([]byte, 0, width))
	s.bufA = bytes.NewBuffer(make([]byte, 0, width))
	if s.extender != nil {
		s.bufE = bytes.NewBuffer(make([]byte, 0, width))
	}

	b := &Bar{
		priority:       s.priority,
		runningBar:     s.runningBar,
		operateState:   make(chan func(*bState)),
		bFrameCh:       make(chan *bFrame, 1),
		syncTableCh:    make(chan [][]chan int),
		completed:      make(chan bool),
		done:           make(chan struct{}),
		shutdown:       make(chan struct{}),
		forceRefreshCh: forceRefreshCh,
	}

	if b.runningBar != nil {
		b.priority = b.runningBar.priority
	}

	go b.serve(ctx, wg, s)
	return b
}

// RemoveAllPrependers removes all prepend functions.
func (b *Bar) RemoveAllPrependers() {
	select {
	case b.operateState <- func(s *bState) { s.pDecorators = nil }:
	case <-b.done:
	}
}

// RemoveAllAppenders removes all append functions.
func (b *Bar) RemoveAllAppenders() {
	select {
	case b.operateState <- func(s *bState) { s.aDecorators = nil }:
	case <-b.done:
	}
}

// ProxyReader wraps r with metrics required for progress tracking.
func (b *Bar) ProxyReader(r io.Reader) io.ReadCloser {
	if r == nil {
		panic("expect io.Reader, got nil")
	}
	rc, ok := r.(io.ReadCloser)
	if !ok {
		rc = ioutil.NopCloser(r)
	}
	return &proxyReader{rc, b, time.Now()}
}

// ID returs id of the bar.
func (b *Bar) ID() int {
	result := make(chan int)
	select {
	case b.operateState <- func(s *bState) { result <- s.id }:
		return <-result
	case <-b.done:
		return b.cacheState.id
	}
}

// Current returns bar's current number, in other words sum of all increments.
func (b *Bar) Current() int64 {
	result := make(chan int64)
	select {
	case b.operateState <- func(s *bState) { result <- s.current }:
		return <-result
	case <-b.done:
		return b.cacheState.current
	}
}

// SetRefill sets refill, if supported by underlying Filler.
func (b *Bar) SetRefill(upto int) {
	b.operateState <- func(s *bState) {
		if f, ok := s.filler.(interface{ SetRefill(int) }); ok {
			f.SetRefill(upto)
		}
	}
}

// SetTotal sets total dynamically.
// Set final to true, when total is known, it will trigger bar complete event.
func (b *Bar) SetTotal(total int64, final bool) bool {
	select {
	case b.operateState <- func(s *bState) {
		if total > 0 {
			s.total = total
		}
		if final && !s.toComplete {
			s.current = s.total
			s.toComplete = true
			go b.forceRefresh()
		}
	}:
		return true
	case <-b.done:
		return false
	}
}

// SetCurrent sets progress' current to arbitrary amount.
func (b *Bar) SetCurrent(current int64, wdd ...time.Duration) {
	if current <= 0 {
		return
	}
	for !atomic.CompareAndSwapUint32(&b.arbitraryCurrent.lock, 0, 1) {
		runtime.Gosched()
	}
	last := b.arbitraryCurrent.current
	b.IncrBy(int(current-last), wdd...)
	b.arbitraryCurrent.current = current
	atomic.StoreUint32(&b.arbitraryCurrent.lock, 0)
}

// Increment is a shorthand for b.IncrBy(1).
func (b *Bar) Increment() {
	b.IncrBy(1)
}

// IncrBy increments progress bar by amount of n.
// wdd is optional work duration i.e. time.Since(start), which expected
// to be provided, if any ewma based decorator is used.
func (b *Bar) IncrBy(n int, wdd ...time.Duration) {
	select {
	case b.operateState <- func(s *bState) {
		s.current += int64(n)
		if s.current >= s.total {
			s.current = s.total
			s.toComplete = true
			go b.forceRefresh()
		}
		for _, ar := range s.amountReceivers {
			ar.NextAmount(n, wdd...)
		}
	}:
	case <-b.done:
	}
}

// Completed reports whether the bar is in completed state.
func (b *Bar) Completed() bool {
	// omit select here, because primary usage of the method is for loop
	// condition, like for !bar.Completed() {...} so when toComplete=true
	// it is called once (at which time, the bar is still alive), then
	// quits the loop and never suppose to be called afterwards.
	return <-b.completed
}

func (b *Bar) wSyncTable() [][]chan int {
	select {
	case b.operateState <- func(s *bState) { b.syncTableCh <- s.wSyncTable() }:
		return <-b.syncTableCh
	case <-b.done:
		return b.cacheState.wSyncTable()
	}
}

func (b *Bar) serve(ctx context.Context, wg *sync.WaitGroup, s *bState) {
	defer wg.Done()
	cancel := ctx.Done()
	for {
		select {
		case op := <-b.operateState:
			op(s)
		case b.completed <- s.toComplete:
		case <-cancel:
			s.toComplete = true
			cancel = nil
		case <-b.shutdown:
			b.cacheState = s
			close(b.done)
			// Notifying decorators about shutdown event
			for _, sl := range s.shutdownListeners {
				sl.Shutdown()
			}
			return
		}
	}
}

func (b *Bar) render(debugOut io.Writer, tw int) {
	select {
	case b.operateState <- func(s *bState) {
		defer func() {
			// recovering if user defined decorator panics for example
			if p := recover(); p != nil {
				s.panicMsg = fmt.Sprintf("panic: %v", p)
				fmt.Fprintf(debugOut, "%s %s bar id %02d %v\n", "[mpb]", time.Now(), s.id, s.panicMsg)
				b.bFrameCh <- &bFrame{
					rd:         strings.NewReader(fmt.Sprintf(fmt.Sprintf("%%.%ds\n", tw), s.panicMsg)),
					toShutdown: true,
				}
			}
		}()
		r := s.draw(tw)
		var extendedLines int
		if s.extender != nil {
			s.extender.Fill(s.bufE, tw, newStatistics(s))
			extendedLines = countLines(s.bufE.Bytes())
			r = io.MultiReader(r, s.bufE)
		}
		b.bFrameCh <- &bFrame{
			rd:               r,
			extendedLines:    extendedLines,
			toShutdown:       s.toComplete && !s.completeFlushed,
			removeOnComplete: s.removeOnComplete,
		}
		s.completeFlushed = s.toComplete
	}:
	case <-b.done:
		s := b.cacheState
		r := s.draw(tw)
		var extendedLines int
		if s.extender != nil {
			s.extender.Fill(s.bufE, tw, newStatistics(s))
			extendedLines = countLines(s.bufE.Bytes())
			r = io.MultiReader(r, s.bufE)
		}
		b.bFrameCh <- &bFrame{
			rd:            r,
			extendedLines: extendedLines,
		}
	}
}

func (s *bState) draw(termWidth int) io.Reader {
	if s.panicMsg != "" {
		return strings.NewReader(fmt.Sprintf(fmt.Sprintf("%%.%ds\n", termWidth), s.panicMsg))
	}

	stat := newStatistics(s)

	for _, d := range s.pDecorators {
		s.bufP.WriteString(d.Decor(stat))
	}

	for _, d := range s.aDecorators {
		s.bufA.WriteString(d.Decor(stat))
	}

	if s.barClearOnComplete && s.completeFlushed {
		s.bufA.WriteByte('\n')
		return io.MultiReader(s.bufP, s.bufA)
	}

	prependCount := utf8.RuneCount(s.bufP.Bytes())
	appendCount := utf8.RuneCount(s.bufA.Bytes())

	if !s.trimSpace {
		// reserve space for edge spaces
		termWidth -= 2
		s.bufB.WriteByte(' ')
	}

	calcWidth := s.width
	if prependCount+s.width+appendCount > termWidth {
		calcWidth = termWidth - prependCount - appendCount
	}
	s.filler.Fill(s.bufB, calcWidth, stat)

	if !s.trimSpace {
		s.bufB.WriteByte(' ')
	}

	s.bufA.WriteByte('\n')
	return io.MultiReader(s.bufP, s.bufB, s.bufA)
}

func (s *bState) wSyncTable() [][]chan int {
	columns := make([]chan int, 0, len(s.pDecorators)+len(s.aDecorators))
	var pCount int
	for _, d := range s.pDecorators {
		if ch, ok := d.Sync(); ok {
			columns = append(columns, ch)
			pCount++
		}
	}
	var aCount int
	for _, d := range s.aDecorators {
		if ch, ok := d.Sync(); ok {
			columns = append(columns, ch)
			aCount++
		}
	}
	table := make([][]chan int, 2)
	table[0] = columns[0:pCount]
	table[1] = columns[pCount : pCount+aCount : pCount+aCount]
	return table
}

func (b *Bar) forceRefresh() {
	select {
	case b.forceRefreshCh <- time.Now():
	case <-b.shutdown:
	}
}

func newStatistics(s *bState) *decor.Statistics {
	return &decor.Statistics{
		ID:        s.id,
		Completed: s.completeFlushed,
		Total:     s.total,
		Current:   s.current,
	}
}

func countLines(b []byte) int {
	return bytes.Count(b, []byte("\n"))
}
<|endoftext|>


<commit_before>// Copyright 2013 bee 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.

// Bee is a tool for developling applications based on beego framework.
package main

import (
	"flag"
	"fmt"
	"html/template"
	"io"
	"log"
	"os"
	"strings"
)

const version = "1.1.2"

type Command struct {
	// Run runs the command.
	// The args are the arguments after the command name.
	Run func(cmd *Command, args []string)

	// UsageLine is the one-line usage message.
	// The first word in the line is taken to be the command name.
	UsageLine string

	// Short is the short description shown in the 'go help' output.
	Short template.HTML

	// Long is the long message shown in the 'go help <this-command>' output.
	Long template.HTML

	// Flag is a set of flags specific to this command.
	Flag flag.FlagSet

	// CustomFlags indicates that the command will do its own
	// flag parsing.
	CustomFlags bool
}

// Name returns the command's name: the first word in the usage line.
func (c *Command) Name() string {
	name := c.UsageLine
	i := strings.Index(name, " ")
	if i >= 0 {
		name = name[:i]
	}
	return name
}

func (c *Command) Usage() {
	fmt.Fprintf(os.Stderr, "usage: %s\n\n", c.UsageLine)
	fmt.Fprintf(os.Stderr, "%s\n", strings.TrimSpace(string(c.Long)))
	os.Exit(2)
}

// Runnable reports whether the command can be run; otherwise
// it is a documentation pseudo-command such as importpath.
func (c *Command) Runnable() bool {
	return c.Run != nil
}

var commands = []*Command{
	cmdNew,
	cmdRun,
	cmdPack,
	cmdApiapp,
	cmdRouter,
	cmdTest,
	cmdBale,
	cmdVersion,
	cmdGenerate,
	cmdRundocs,
}

func main() {
	flag.Usage = usage
	flag.Parse()
	log.SetFlags(0)

	args := flag.Args()
	if len(args) < 1 {
		usage()
	}

	if args[0] == "help" {
		help(args[1:])
		return
	}

	for _, cmd := range commands {
		if cmd.Name() == args[0] && cmd.Run != nil {
			cmd.Flag.Usage = func() { cmd.Usage() }
			if cmd.CustomFlags {
				args = args[1:]
			} else {
				cmd.Flag.Parse(args[1:])
				args = cmd.Flag.Args()
			}
			cmd.Run(cmd, args)
			os.Exit(2)
			return
		}
	}

	fmt.Fprintf(os.Stderr, "bee: unknown subcommand %q\nRun 'bee help' for usage.\n", args[0])
	os.Exit(2)
}

var usageTemplate = `Bee is a tool for managing beego framework.

Usage:

	bee command [arguments]

The commands are:
{{range .}}{{if .Runnable}}
    {{.Name | printf "%-11s"}} {{.Short}}{{end}}{{end}}

Use "bee help [command]" for more information about a command.

Additional help topics:
{{range .}}{{if not .Runnable}}
    {{.Name | printf "%-11s"}} {{.Short}}{{end}}{{end}}

Use "bee help [topic]" for more information about that topic.

`

var helpTemplate = `{{if .Runnable}}usage: bee {{.UsageLine}}

{{end}}{{.Long | trim}}
`

func usage() {
	tmpl(os.Stdout, usageTemplate, commands)
	os.Exit(2)
}

func tmpl(w io.Writer, text string, data interface{}) {
	t := template.New("top")
	t.Funcs(template.FuncMap{"trim": func(s template.HTML) template.HTML {
		return template.HTML(strings.TrimSpace(string(s)))
	}})
	template.Must(t.Parse(text))
	if err := t.Execute(w, data); err != nil {
		panic(err)
	}
}

func help(args []string) {
	if len(args) == 0 {
		usage()
		// not exit 2: succeeded at 'go help'.
		return
	}
	if len(args) != 1 {
		fmt.Fprintf(os.Stdout, "usage: bee help command\n\nToo many arguments given.\n")
		os.Exit(2) // failed at 'bee help'
	}

	arg := args[0]

	for _, cmd := range commands {
		if cmd.Name() == arg {
			tmpl(os.Stdout, helpTemplate, cmd)
			// not exit 2: succeeded at 'go help cmd'.
			return
		}
	}

	fmt.Fprintf(os.Stdout, "Unknown help topic %#q.  Run 'bee help'.\n", arg)
	os.Exit(2) // failed at 'bee help cmd'
}
<commit_msg>update version number to 1.3<commit_after>// Copyright 2013 bee 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.

// Bee is a tool for developling applications based on beego framework.
package main

import (
	"flag"
	"fmt"
	"html/template"
	"io"
	"log"
	"os"
	"strings"
)

const version = "1.1.3"

type Command struct {
	// Run runs the command.
	// The args are the arguments after the command name.
	Run func(cmd *Command, args []string)

	// UsageLine is the one-line usage message.
	// The first word in the line is taken to be the command name.
	UsageLine string

	// Short is the short description shown in the 'go help' output.
	Short template.HTML

	// Long is the long message shown in the 'go help <this-command>' output.
	Long template.HTML

	// Flag is a set of flags specific to this command.
	Flag flag.FlagSet

	// CustomFlags indicates that the command will do its own
	// flag parsing.
	CustomFlags bool
}

// Name returns the command's name: the first word in the usage line.
func (c *Command) Name() string {
	name := c.UsageLine
	i := strings.Index(name, " ")
	if i >= 0 {
		name = name[:i]
	}
	return name
}

func (c *Command) Usage() {
	fmt.Fprintf(os.Stderr, "usage: %s\n\n", c.UsageLine)
	fmt.Fprintf(os.Stderr, "%s\n", strings.TrimSpace(string(c.Long)))
	os.Exit(2)
}

// Runnable reports whether the command can be run; otherwise
// it is a documentation pseudo-command such as importpath.
func (c *Command) Runnable() bool {
	return c.Run != nil
}

var commands = []*Command{
	cmdNew,
	cmdRun,
	cmdPack,
	cmdApiapp,
	cmdRouter,
	cmdTest,
	cmdBale,
	cmdVersion,
	cmdGenerate,
	cmdRundocs,
}

func main() {
	flag.Usage = usage
	flag.Parse()
	log.SetFlags(0)

	args := flag.Args()
	if len(args) < 1 {
		usage()
	}

	if args[0] == "help" {
		help(args[1:])
		return
	}

	for _, cmd := range commands {
		if cmd.Name() == args[0] && cmd.Run != nil {
			cmd.Flag.Usage = func() { cmd.Usage() }
			if cmd.CustomFlags {
				args = args[1:]
			} else {
				cmd.Flag.Parse(args[1:])
				args = cmd.Flag.Args()
			}
			cmd.Run(cmd, args)
			os.Exit(2)
			return
		}
	}

	fmt.Fprintf(os.Stderr, "bee: unknown subcommand %q\nRun 'bee help' for usage.\n", args[0])
	os.Exit(2)
}

var usageTemplate = `Bee is a tool for managing beego framework.

Usage:

	bee command [arguments]

The commands are:
{{range .}}{{if .Runnable}}
    {{.Name | printf "%-11s"}} {{.Short}}{{end}}{{end}}

Use "bee help [command]" for more information about a command.

Additional help topics:
{{range .}}{{if not .Runnable}}
    {{.Name | printf "%-11s"}} {{.Short}}{{end}}{{end}}

Use "bee help [topic]" for more information about that topic.

`

var helpTemplate = `{{if .Runnable}}usage: bee {{.UsageLine}}

{{end}}{{.Long | trim}}
`

func usage() {
	tmpl(os.Stdout, usageTemplate, commands)
	os.Exit(2)
}

func tmpl(w io.Writer, text string, data interface{}) {
	t := template.New("top")
	t.Funcs(template.FuncMap{"trim": func(s template.HTML) template.HTML {
		return template.HTML(strings.TrimSpace(string(s)))
	}})
	template.Must(t.Parse(text))
	if err := t.Execute(w, data); err != nil {
		panic(err)
	}
}

func help(args []string) {
	if len(args) == 0 {
		usage()
		// not exit 2: succeeded at 'go help'.
		return
	}
	if len(args) != 1 {
		fmt.Fprintf(os.Stdout, "usage: bee help command\n\nToo many arguments given.\n")
		os.Exit(2) // failed at 'bee help'
	}

	arg := args[0]

	for _, cmd := range commands {
		if cmd.Name() == arg {
			tmpl(os.Stdout, helpTemplate, cmd)
			// not exit 2: succeeded at 'go help cmd'.
			return
		}
	}

	fmt.Fprintf(os.Stdout, "Unknown help topic %#q.  Run 'bee help'.\n", arg)
	os.Exit(2) // failed at 'bee help cmd'
}
<|endoftext|>