<commit_before>package axe

import (
	"reflect"
	"time"

	"github.com/globalsign/mgo/bson"
)

// Model can be any BSON serializable type.
type Model interface{}

// Task is task that is executed asynchronously.
type Task struct {
	// Name is the unique name of the task.
	Name string

	// Model is the model that holds task related data.
	Model Model

	// Queue is the queue that is used to managed the jobs.
	Queue *Queue

	// Handler is the callback called with tasks.
	Handler func(Model) (bson.M, error)

	// Workers defines the number for spawned workers.
	//
	// Default: 1.
	Workers int

	// MaxAttempts defines the maximum attempts to complete a task.
	//
	// Default: 1
	MaxAttempts int

	// Interval is interval at which the worker will request a job from the queue.
	//
	// Default: 100ms.
	Interval time.Duration
}

func (t *Task) run(p *Pool) {
	// set default workers
	if t.Workers == 0 {
		t.Workers = 1
	}

	// set default max attempts
	if t.MaxAttempts == 0 {
		t.MaxAttempts = 1
	}

	// set default interval
	if t.Interval == 0 {
		t.Interval = 100 * time.Millisecond
	}

	// run workers
	for i := 0; i < t.Workers; i++ {
		go t.worker(p)
	}
}

func (t *Task) worker(p *Pool) {
	// run forever
	for {
		// return if closed
		select {
		case <-p.closed:
			return
		default:
		}

		// attempt to get job from queue
		job := t.Queue.get(t.Name)
		if job == nil {
			// wait some time and try again
			time.Sleep(t.Interval)

			continue
		}

		// execute worker and report errors
		err := t.execute(job)
		if err != nil {
			if p.Reporter != nil {
				p.Reporter(err)
			}
		}
	}
}

func (t *Task) execute(job *Job) error {
	// TODO: Dequeue specified job.

	// get store
	store := t.Queue.Store.Copy()
	defer store.Close()

	// dequeue task
	job, err := dequeue(store, job.ID(), time.Hour)
	if err != nil {
		return err
	}

	// return if missing
	if job == nil {
		return nil
	}

	// instantiate model
	data := reflect.New(reflect.TypeOf(t.Model).Elem()).Interface()

	// unmarshal data
	err = job.Data.Unmarshal(data)
	if err != nil {
		return err
	}

	// run handler
	result, err := t.Handler(data)
	if err != nil {
		return err
	}

	// TODO: Fail task on error.
	// TODO: Cancel task if max attempts has been reached.

	// complete task
	err = complete(store, job.ID(), result)
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>added error handling<commit_after>package axe

import (
	"reflect"
	"time"

	"github.com/globalsign/mgo/bson"
)

// Error is used to signal failed job executions.
type Error struct {
	Reason string
	Retry  bool
}

// Error implements the error interface.
func (c *Error) Error() string {
	return c.Reason
}

// E is a short-hand to construct an error.
func E(reason string, retry bool) *Error {
	return &Error{
		Reason: reason,
		Retry:  retry,
	}
}

// Model can be any BSON serializable type.
type Model interface{}

// Task is task that is executed asynchronously.
type Task struct {
	// Name is the unique name of the task.
	Name string

	// Model is the model that holds task related data.
	Model Model

	// Queue is the queue that is used to managed the jobs.
	Queue *Queue

	// Handler is the callback called with tasks.
	Handler func(Model) (bson.M, error)

	// Workers defines the number for spawned workers.
	//
	// Default: 1.
	Workers int

	// MaxAttempts defines the maximum attempts to complete a task.
	//
	// Default: 1
	MaxAttempts int

	// Interval is interval at which the worker will request a job from the queue.
	//
	// Default: 100ms.
	Interval time.Duration
}

func (t *Task) run(p *Pool) {
	// set default workers
	if t.Workers == 0 {
		t.Workers = 1
	}

	// set default max attempts
	if t.MaxAttempts == 0 {
		t.MaxAttempts = 1
	}

	// set default interval
	if t.Interval == 0 {
		t.Interval = 100 * time.Millisecond
	}

	// run workers
	for i := 0; i < t.Workers; i++ {
		go t.worker(p)
	}
}

func (t *Task) worker(p *Pool) {
	// run forever
	for {
		// return if closed
		select {
		case <-p.closed:
			return
		default:
		}

		// attempt to get job from queue
		job := t.Queue.get(t.Name)
		if job == nil {
			// wait some time and try again
			time.Sleep(t.Interval)

			continue
		}

		// execute worker and report errors
		err := t.execute(job)
		if err != nil {
			if p.Reporter != nil {
				p.Reporter(err)
			}
		}
	}
}

func (t *Task) execute(job *Job) error {
	// get store
	store := t.Queue.Store.Copy()
	defer store.Close()

	// dequeue job
	job, err := dequeue(store, job.ID(), time.Hour)
	if err != nil {
		return err
	}

	// return if missing
	if job == nil {
		return nil
	}

	// instantiate model
	data := reflect.New(reflect.TypeOf(t.Model).Elem()).Interface()

	// unmarshal data
	err = job.Data.Unmarshal(data)
	if err != nil {
		return err
	}

	// run handler
	result, err := t.Handler(data)
	if _, ok := err.(*Error); !ok && err != nil {
		return err
	}

	// check error
	if e, ok := err.(*Error); ok {
		// check retry and attempts
		if !e.Retry || job.Attempts >= t.MaxAttempts {
			// cancel job
			err = cancel(store, job.ID(), e.Reason)
			if err != nil {
				return err
			}

			return nil
		}

		// fail job
		err = fail(store, job.ID(), e.Reason, 0)
		if err != nil {
			return err
		}
	}

	// complete job
	err = complete(store, job.ID(), result)
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2020 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package zanzibar

import (
	"context"
	"fmt"
	"net/http"
	"strings"
	"time"

	"github.com/buger/jsonparser"
	"github.com/pkg/errors"
	"github.com/uber-go/tally"
	"github.com/uber/zanzibar/runtime/jsonwrapper"
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
)

// ServerHTTPResponse struct manages server http response
type ServerHTTPResponse struct {
	Request    *ServerHTTPRequest
	StatusCode int

	responseWriter    http.ResponseWriter
	flushed           bool
	finished          bool
	finishTime        time.Time
	pendingBodyBytes  []byte
	pendingBodyObj    interface{}
	pendingStatusCode int
	logger            Logger
	scope             tally.Scope
	jsonWrapper       jsonwrapper.JSONWrapper
	err               error
}

// NewServerHTTPResponse is helper function to alloc ServerHTTPResponse
func NewServerHTTPResponse(
	w http.ResponseWriter,
	req *ServerHTTPRequest,
) *ServerHTTPResponse {
	return &ServerHTTPResponse{
		Request:        req,
		StatusCode:     200,
		responseWriter: w,
		logger:         req.logger,
		scope:          req.scope,
		jsonWrapper:    req.jsonWrapper,
	}
}

// finish will handle final logic, like metrics
func (res *ServerHTTPResponse) finish(ctx context.Context) {
	logFields := GetLogFieldsFromCtx(ctx)
	if !res.Request.started {
		/* coverage ignore next line */
		res.logger.Error(
			"Forgot to start server response",
			append(logFields, zap.String("path", res.Request.URL.Path))...,
		)
		/* coverage ignore next line */
		return
	}
	if res.finished {
		/* coverage ignore next line */
		res.logger.Error(
			"Finished a server response multiple times",
			append(logFields, zap.String("path", res.Request.URL.Path))...,
		)
		/* coverage ignore next line */
		return
	}
	res.finished = true
	res.finishTime = time.Now()

	_, known := knownStatusCodes[res.StatusCode]
	// no need to put this tag on the context because this is the end of response life cycle
	statusTag := map[string]string{scopeTagStatus: fmt.Sprintf("%d", res.StatusCode)}
	tagged := res.scope.Tagged(statusTag)
	delta := res.finishTime.Sub(res.Request.startTime)
	tagged.Timer(endpointLatency).Record(delta)
	tagged.Histogram(endpointLatencyHist, tally.DefaultBuckets).RecordDuration(delta)
	if !known {
		res.logger.Error(
			"Unknown status code",
			append(logFields, zap.Int("UnknownStatusCode", res.StatusCode))...,
		)
	} else {
		tagged.Counter(endpointStatus).Inc(1)
	}

	logFn := res.logger.Debug
	if !known || res.StatusCode >= 400 && res.StatusCode < 600 {
		tagged.Counter(endpointAppErrors).Inc(1)
		logFn = res.logger.Warn
	}

	span := res.Request.GetSpan()
	if span != nil {
		span.Finish()
	}

	logFn(
		"Finished an incoming server HTTP request",
		append(logFields, serverHTTPLogFields(res.Request, res)...)...,
	)
}

func serverHTTPLogFields(req *ServerHTTPRequest, res *ServerHTTPResponse) []zapcore.Field {
	fields := []zapcore.Field{
		zap.String("method", req.httpRequest.Method),
		zap.String("remoteAddr", req.httpRequest.RemoteAddr),
		zap.String("pathname", req.httpRequest.URL.RequestURI()),
		zap.String("host", req.httpRequest.Host),
		zap.Time("timestamp-started", req.startTime),
		zap.Time("timestamp-finished", res.finishTime),
		zap.Int("statusCode", res.StatusCode),
	}

	for k, v := range res.Headers() {
		if len(v) > 0 {
			fields = append(fields, zap.String(
				fmt.Sprintf("%s-%s", logFieldEndpointResponseHeaderPrefix, k),
				strings.Join(v, ", "),
			))
		}
	}

	if res.err != nil {
		fields = append(fields, zap.Error(res.err))

		cause := errors.Cause(res.err)
		if cause != nil && cause != res.err {
			fields = append(fields, zap.NamedError("errorCause", cause))
		}
	}

	return fields
}

// SendErrorString helper to send an error string
func (res *ServerHTTPResponse) SendErrorString(
	statusCode int, errMsg string,
) {
	res.WriteJSONBytes(statusCode, nil,
		[]byte(`{"error":"`+errMsg+`"}`),
	)
}

// SendError helper to send an server error message, propagates underlying cause to logs etc.
func (res *ServerHTTPResponse) SendError(
	statusCode int, errMsg string, errCause error,
) {
	res.err = errCause
	res.WriteJSONBytes(statusCode, nil,
		[]byte(`{"error":"`+errMsg+`"}`),
	)
}

// WriteBytes writes a byte[] slice that is valid Response
func (res *ServerHTTPResponse) WriteBytes(
	statusCode int, headers Header, bytes []byte,
) {
	if headers != nil {
		for _, k := range headers.Keys() {
			v, ok := headers.Get(k)
			if ok {
				res.responseWriter.Header().Set(k, v)
			}
		}
	}

	res.pendingStatusCode = statusCode
	res.pendingBodyBytes = bytes
}

// WriteJSONBytes writes a byte[] slice that is valid json to Response
func (res *ServerHTTPResponse) WriteJSONBytes(
	statusCode int, headers Header, bytes []byte,
) {
	if headers == nil {
		headers = ServerHTTPHeader{}
	}

	headers.Add("content-type", "application/json")
	res.WriteBytes(statusCode, headers, bytes)
}

// WriteJSON writes a json serializable struct to Response
func (res *ServerHTTPResponse) WriteJSON(
	statusCode int, headers Header, body interface{},
) {
	if body == nil {
		res.SendError(500, "Could not serialize json response", errors.New("No Body JSON"))
		res.logger.Error("Could not serialize nil pointer body")
		return
	}
	bytes, err := res.jsonWrapper.Marshal(body)
	if err != nil {
		res.SendError(500, "Could not serialize json response", err)
		res.logger.Error("Could not serialize json response", zap.Error(err))
		return
	}

	contentTypePresent := false
	if headers != nil {
		for _, k := range headers.Keys() {
			v, ok := headers.Get(k)
			if ok {
				if k == "Content-Type" {
					contentTypePresent = true
				}
				res.responseWriter.Header().Set(k, v)
			}
		}
	}

	// Set the content-type to application/json if not already available
	if !contentTypePresent {
		res.responseWriter.Header().
			Set("content-type", "application/json")
	}
	res.pendingStatusCode = statusCode
	res.pendingBodyBytes = bytes
	res.pendingBodyObj = body
}

// PeekBody allows for inspecting a key path inside the body
// that is not flushed yet. This is useful for response middlewares
// that want to inspect the response body.
func (res *ServerHTTPResponse) PeekBody(
	keys ...string,
) ([]byte, jsonparser.ValueType, error) {
	value, valueType, _, err := jsonparser.Get(
		res.pendingBodyBytes, keys...,
	)

	if err != nil {
		return nil, -1, err
	}

	return value, valueType, nil
}

// Flush will write the body to the response. Before flush is called
// the body is pending. A pending body allows a response middleware to
// write a different body.
func (res *ServerHTTPResponse) flush(ctx context.Context) {
	if res.flushed {
		/* coverage ignore next line */
		res.logger.Error(
			"Flushed a server response multiple times",
			zap.String("path", res.Request.URL.Path),
		)
		/* coverage ignore next line */
		return
	}

	res.flushed = true
	res.writeHeader(res.pendingStatusCode)
	res.writeBytes(res.pendingBodyBytes)
	res.finish(ctx)
}

func (res *ServerHTTPResponse) writeHeader(statusCode int) {
	res.StatusCode = statusCode
	res.responseWriter.WriteHeader(statusCode)
}

// WriteBytes writes raw bytes to output
func (res *ServerHTTPResponse) writeBytes(bytes []byte) {
	_, err := res.responseWriter.Write(bytes)
	if err != nil {
		/* coverage ignore next line */
		res.logger.Error(
			"Could not write string to resp body",
			zap.Error(err),
		)
	}
}

// GetPendingResponse lets you read the pending body bytes, obj and status code
// which isn't sent back yet.
func (res *ServerHTTPResponse) GetPendingResponse() ([]byte, int) {
	return res.pendingBodyBytes, res.pendingStatusCode
}

// Headers returns the underlying http response's headers
func (res *ServerHTTPResponse) Headers() http.Header {
	return res.responseWriter.Header()
}
<commit_msg>Refactor write json (#711)<commit_after>// Copyright (c) 2020 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package zanzibar

import (
	"context"
	"fmt"
	"net/http"
	"strings"
	"time"

	"github.com/buger/jsonparser"
	"github.com/pkg/errors"
	"github.com/uber-go/tally"
	"github.com/uber/zanzibar/runtime/jsonwrapper"
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
)

// ServerHTTPResponse struct manages server http response
type ServerHTTPResponse struct {
	Request    *ServerHTTPRequest
	StatusCode int

	responseWriter    http.ResponseWriter
	flushed           bool
	finished          bool
	finishTime        time.Time
	pendingBodyBytes  []byte
	pendingBodyObj    interface{}
	pendingStatusCode int
	logger            Logger
	scope             tally.Scope
	jsonWrapper       jsonwrapper.JSONWrapper
	err               error
}

// NewServerHTTPResponse is helper function to alloc ServerHTTPResponse
func NewServerHTTPResponse(
	w http.ResponseWriter,
	req *ServerHTTPRequest,
) *ServerHTTPResponse {
	return &ServerHTTPResponse{
		Request:        req,
		StatusCode:     200,
		responseWriter: w,
		logger:         req.logger,
		scope:          req.scope,
		jsonWrapper:    req.jsonWrapper,
	}
}

// finish will handle final logic, like metrics
func (res *ServerHTTPResponse) finish(ctx context.Context) {
	logFields := GetLogFieldsFromCtx(ctx)
	if !res.Request.started {
		/* coverage ignore next line */
		res.logger.Error(
			"Forgot to start server response",
			append(logFields, zap.String("path", res.Request.URL.Path))...,
		)
		/* coverage ignore next line */
		return
	}
	if res.finished {
		/* coverage ignore next line */
		res.logger.Error(
			"Finished a server response multiple times",
			append(logFields, zap.String("path", res.Request.URL.Path))...,
		)
		/* coverage ignore next line */
		return
	}
	res.finished = true
	res.finishTime = time.Now()

	_, known := knownStatusCodes[res.StatusCode]
	// no need to put this tag on the context because this is the end of response life cycle
	statusTag := map[string]string{scopeTagStatus: fmt.Sprintf("%d", res.StatusCode)}
	tagged := res.scope.Tagged(statusTag)
	delta := res.finishTime.Sub(res.Request.startTime)
	tagged.Timer(endpointLatency).Record(delta)
	tagged.Histogram(endpointLatencyHist, tally.DefaultBuckets).RecordDuration(delta)
	if !known {
		res.logger.Error(
			"Unknown status code",
			append(logFields, zap.Int("UnknownStatusCode", res.StatusCode))...,
		)
	} else {
		tagged.Counter(endpointStatus).Inc(1)
	}

	logFn := res.logger.Debug
	if !known || res.StatusCode >= 400 && res.StatusCode < 600 {
		tagged.Counter(endpointAppErrors).Inc(1)
		logFn = res.logger.Warn
	}

	span := res.Request.GetSpan()
	if span != nil {
		span.Finish()
	}

	logFn(
		"Finished an incoming server HTTP request",
		append(logFields, serverHTTPLogFields(res.Request, res)...)...,
	)
}

func serverHTTPLogFields(req *ServerHTTPRequest, res *ServerHTTPResponse) []zapcore.Field {
	fields := []zapcore.Field{
		zap.String("method", req.httpRequest.Method),
		zap.String("remoteAddr", req.httpRequest.RemoteAddr),
		zap.String("pathname", req.httpRequest.URL.RequestURI()),
		zap.String("host", req.httpRequest.Host),
		zap.Time("timestamp-started", req.startTime),
		zap.Time("timestamp-finished", res.finishTime),
		zap.Int("statusCode", res.StatusCode),
	}

	for k, v := range res.Headers() {
		if len(v) > 0 {
			fields = append(fields, zap.String(
				fmt.Sprintf("%s-%s", logFieldEndpointResponseHeaderPrefix, k),
				strings.Join(v, ", "),
			))
		}
	}

	if res.err != nil {
		fields = append(fields, zap.Error(res.err))

		cause := errors.Cause(res.err)
		if cause != nil && cause != res.err {
			fields = append(fields, zap.NamedError("errorCause", cause))
		}
	}

	return fields
}

// SendErrorString helper to send an error string
func (res *ServerHTTPResponse) SendErrorString(
	statusCode int, errMsg string,
) {
	res.WriteJSONBytes(statusCode, nil,
		[]byte(`{"error":"`+errMsg+`"}`),
	)
}

// SendError helper to send an server error message, propagates underlying cause to logs etc.
func (res *ServerHTTPResponse) SendError(
	statusCode int, errMsg string, errCause error,
) {
	res.err = errCause
	res.WriteJSONBytes(statusCode, nil,
		[]byte(`{"error":"`+errMsg+`"}`),
	)
}

// WriteBytes writes a byte[] slice that is valid Response
func (res *ServerHTTPResponse) WriteBytes(
	statusCode int, headers Header, bytes []byte,
) {
	if headers != nil {
		for _, k := range headers.Keys() {
			v, ok := headers.Get(k)
			if ok {
				res.responseWriter.Header().Set(k, v)
			}
		}
	}

	res.pendingStatusCode = statusCode
	res.pendingBodyBytes = bytes
}

// WriteJSONBytes writes a byte[] slice that is valid json to Response
func (res *ServerHTTPResponse) WriteJSONBytes(
	statusCode int, headers Header, bytes []byte,
) {
	if headers == nil {
		headers = ServerHTTPHeader{}
	}

	headers.Add("content-type", "application/json")
	res.WriteBytes(statusCode, headers, bytes)
}

// MarshalResponseJSON serializes a json serializable into bytes
func (res *ServerHTTPResponse) MarshalResponseJSON(body interface{}) []byte {
	if body == nil {
		res.SendError(500, "Could not serialize json response", errors.New("No Body JSON"))
		res.logger.Error("Could not serialize nil pointer body")
		return nil
	}
	bytes, err := res.jsonWrapper.Marshal(body)
	if err != nil {
		res.SendError(500, "Could not serialize json response", err)
		res.logger.Error("Could not serialize json response", zap.Error(err))
		return nil
	}
	return bytes
}

// SendResponse sets content-type if not present and fills Response
func (res *ServerHTTPResponse) SendResponse(statusCode int, headers Header, body interface{}, bytes []byte) {
	contentTypePresent := false
	if headers != nil {
		for _, k := range headers.Keys() {
			v, ok := headers.Get(k)
			if ok {
				if k == "Content-Type" {
					contentTypePresent = true
				}
				res.responseWriter.Header().Set(k, v)
			}
		}
	}

	// Set the content-type to application/json if not already available
	if !contentTypePresent {
		res.responseWriter.Header().
			Set("content-type", "application/json")
	}
	res.pendingStatusCode = statusCode
	res.pendingBodyBytes = bytes
	res.pendingBodyObj = body
}

// WriteJSON writes a json serializable struct to Response
func (res *ServerHTTPResponse) WriteJSON(
	statusCode int, headers Header, body interface{},
) {
	bytes := res.MarshalResponseJSON(body)
	if bytes == nil {
		return
	}
	res.SendResponse(statusCode, headers, body, bytes)
}

// PeekBody allows for inspecting a key path inside the body
// that is not flushed yet. This is useful for response middlewares
// that want to inspect the response body.
func (res *ServerHTTPResponse) PeekBody(
	keys ...string,
) ([]byte, jsonparser.ValueType, error) {
	value, valueType, _, err := jsonparser.Get(
		res.pendingBodyBytes, keys...,
	)

	if err != nil {
		return nil, -1, err
	}

	return value, valueType, nil
}

// Flush will write the body to the response. Before flush is called
// the body is pending. A pending body allows a response middleware to
// write a different body.
func (res *ServerHTTPResponse) flush(ctx context.Context) {
	if res.flushed {
		/* coverage ignore next line */
		res.logger.Error(
			"Flushed a server response multiple times",
			zap.String("path", res.Request.URL.Path),
		)
		/* coverage ignore next line */
		return
	}

	res.flushed = true
	res.writeHeader(res.pendingStatusCode)
	res.writeBytes(res.pendingBodyBytes)
	res.finish(ctx)
}

func (res *ServerHTTPResponse) writeHeader(statusCode int) {
	res.StatusCode = statusCode
	res.responseWriter.WriteHeader(statusCode)
}

// WriteBytes writes raw bytes to output
func (res *ServerHTTPResponse) writeBytes(bytes []byte) {
	_, err := res.responseWriter.Write(bytes)
	if err != nil {
		/* coverage ignore next line */
		res.logger.Error(
			"Could not write string to resp body",
			zap.Error(err),
		)
	}
}

// GetPendingResponse lets you read the pending body bytes, obj and status code
// which isn't sent back yet.
func (res *ServerHTTPResponse) GetPendingResponse() ([]byte, int) {
	return res.pendingBodyBytes, res.pendingStatusCode
}

// Headers returns the underlying http response's headers
func (res *ServerHTTPResponse) Headers() http.Header {
	return res.responseWriter.Header()
}
<|endoftext|>


<commit_before>package deploy

import (
	"fmt"
	"log"
	"strconv"
	"strings"

	"github.com/cenk/backoff"
	"github.com/fatih/color"
	marathon "github.com/gambol99/go-marathon"
	"github.com/hashicorp/consul/api"

	"github.com/InnovaCo/serve/app/build"
	"github.com/InnovaCo/serve/manifest"
	"github.com/InnovaCo/serve/utils"
)

type SiteDeploy struct{}
type SiteRelease struct{}

func (_ SiteDeploy) Run(m *manifest.Manifest, sub *manifest.Manifest) error {
	conf := marathon.NewDefaultConfig()
	conf.URL = fmt.Sprintf("http://%s:8080", m.GetString("marathon.marathon-host"))
	marathonApi, _ := marathon.NewClient(conf)

	name := m.ServiceName() + "-v" + m.BuildVersion()

	app := &marathon.Application{}
	app.Name(m.GetStringOr("info.category", "") + "/" + name)
	app.Command(fmt.Sprintf("serve consul supervisor --service '%s' --port $PORT0 start %s", name, sub.GetStringOr("marathon.cmd", "bin/start")))
	app.Count(sub.GetIntOr("marathon.instances", 1))
	app.Memory(float64(sub.GetIntOr("marathon.mem", 256)))

	app.BackoffSeconds(3)
	app.BackoffFactor(2)
	app.MaxLaunchDelaySeconds(30)

	if cpu, err := strconv.ParseFloat(sub.GetStringOr("marathon.cpu", "0.1"), 64); err == nil {
		app.CPU(cpu)
	}

	if constrs := sub.GetStringOr("marathon.constraints", ""); constrs != "" {
		cs := strings.SplitN(constrs, ":", 2)
		app.AddConstraint(cs[0], "CLUSTER", cs[1])
		app.AddLabel(cs[0], cs[1])
	}

	app.AddEnv("ENV", m.Args("env"))
	app.AddEnv("SERVICE_NAME", m.ServiceName())
	app.AddEnv("MEMORY", sub.GetStringOr("marathon.mem", ""))

	app.AddUris(build.TaskRegistryUrl(m))

	if _, err := marathonApi.UpdateApplication(app, false); err != nil {
		color.Yellow("marathon <- %s", app)
		return err
	}

	color.Green("marathon <- %s", app)

	consul := ConsulClient(m)

	return backoff.Retry(func() error {
		services, _, err := consul.Health().Service(name, "", true, nil)

		if err != nil {
			return err
		}

		if len(services) == 0 {
			log.Printf("Service `%s` not started yet! Retry...", name)
			return fmt.Errorf("Service `%s` not started!", name)
		}

		log.Println(color.GreenString("Service `%s` successfily deloyed!", name))
		return nil
	}, backoff.NewExponentialBackOff())
}

func (_ SiteRelease) Run(m *manifest.Manifest, sub *manifest.Manifest) error {
	log.Println("Release done!", sub)

	conf := api.DefaultConfig()
	conf.Address = m.GetString("consul.consul-host") + ":8500"

	consul, _ := api.NewClient(conf)

	routes := make([]map[string]string, 0)
	for _, route := range sub.Array("routes") {

		// todo: merge with --route flag
		// filter featured: true route
		routes = append(routes, map[string]string{
			"host":     route.GetString("host"),
			"location": route.GetString("location"),
		})
	}

	consul.KV().Put(&api.KVPair{
		Key:   fmt.Sprintf("services/%s/%s/routes", m.ServiceName(), m.BuildVersion()),
		Value: []byte("test"),
	}, nil)

	// находим текущий в консуле и убеждаемся что с ним все ок
	// добавляем ему роуты

	// ищем есть ли старый с такими же роутами:
	//    формируем массив роутов
	//    ищем сервис с таким-же именем но другой версии, и содержащий один из указанных роутов
	//    например в kv можно хранить /kv/services/{name-?branch}/v{version}} и там матчить через compareMaps
	//    если хотябы один роут полностью совпал — это наш кандидат на убивание
	// если есть — убиваем в консуле сразу и через 5 минут в марафоне

	println(utils.MapsEqual(map[string]string{"name": "dima", "version": "1.0"}, map[string]string{"version": "1.0", "name": "dima"}))

	log.Println("route")

	return nil
}

func ConsulClient(m *manifest.Manifest) *api.Client {
	conf := api.DefaultConfig()
	conf.Address = m.GetString("consul.consul-host") + ":8500"

	consul, _ := api.NewClient(conf)
	return consul
}
<commit_msg> = fix marathon backof settings<commit_after>package deploy

import (
	"fmt"
	"log"
	"strconv"
	"strings"

	"github.com/cenk/backoff"
	"github.com/fatih/color"
	marathon "github.com/gambol99/go-marathon"
	"github.com/hashicorp/consul/api"

	"github.com/InnovaCo/serve/app/build"
	"github.com/InnovaCo/serve/manifest"
	"github.com/InnovaCo/serve/utils"
)

type SiteDeploy struct{}
type SiteRelease struct{}

func (_ SiteDeploy) Run(m *manifest.Manifest, sub *manifest.Manifest) error {
	conf := marathon.NewDefaultConfig()
	conf.URL = fmt.Sprintf("http://%s:8080", m.GetString("marathon.marathon-host"))
	marathonApi, _ := marathon.NewClient(conf)

	name := m.ServiceName() + "-v" + m.BuildVersion()

	app := &marathon.Application{
		BackoffSeconds: 3,
		BackoffFactor: 2,
		MaxLaunchDelaySeconds: 30,
	}

	app.Name(m.GetStringOr("info.category", "") + "/" + name)
	app.Command(fmt.Sprintf("serve consul supervisor --service '%s' --port $PORT0 start %s", name, sub.GetStringOr("marathon.cmd", "bin/start")))
	app.Count(sub.GetIntOr("marathon.instances", 1))
	app.Memory(float64(sub.GetIntOr("marathon.mem", 256)))

	if cpu, err := strconv.ParseFloat(sub.GetStringOr("marathon.cpu", "0.1"), 64); err == nil {
		app.CPU(cpu)
	}

	if constrs := sub.GetStringOr("marathon.constraints", ""); constrs != "" {
		cs := strings.SplitN(constrs, ":", 2)
		app.AddConstraint(cs[0], "CLUSTER", cs[1])
		app.AddLabel(cs[0], cs[1])
	}

	app.AddEnv("ENV", m.Args("env"))
	app.AddEnv("SERVICE_NAME", m.ServiceName())
	app.AddEnv("MEMORY", sub.GetStringOr("marathon.mem", ""))

	app.AddUris(build.TaskRegistryUrl(m))

	if _, err := marathonApi.UpdateApplication(app, false); err != nil {
		color.Yellow("marathon <- %s", app)
		return err
	}

	color.Green("marathon <- %s", app)

	consul := ConsulClient(m)

	return backoff.Retry(func() error {
		services, _, err := consul.Health().Service(name, "", true, nil)

		if err != nil {
			return err
		}

		if len(services) == 0 {
			log.Printf("Service `%s` not started yet! Retry...", name)
			return fmt.Errorf("Service `%s` not started!", name)
		}

		log.Println(color.GreenString("Service `%s` successfily deloyed!", name))
		return nil
	}, backoff.NewExponentialBackOff())
}

func (_ SiteRelease) Run(m *manifest.Manifest, sub *manifest.Manifest) error {
	log.Println("Release done!", sub)

	conf := api.DefaultConfig()
	conf.Address = m.GetString("consul.consul-host") + ":8500"

	consul, _ := api.NewClient(conf)

	routes := make([]map[string]string, 0)
	for _, route := range sub.Array("routes") {

		// todo: merge with --route flag
		// filter featured: true route
		routes = append(routes, map[string]string{
			"host":     route.GetString("host"),
			"location": route.GetString("location"),
		})
	}

	consul.KV().Put(&api.KVPair{
		Key:   fmt.Sprintf("services/%s/%s/routes", m.ServiceName(), m.BuildVersion()),
		Value: []byte("test"),
	}, nil)

	// находим текущий в консуле и убеждаемся что с ним все ок
	// добавляем ему роуты

	// ищем есть ли старый с такими же роутами:
	//    формируем массив роутов
	//    ищем сервис с таким-же именем но другой версии, и содержащий один из указанных роутов
	//    например в kv можно хранить /kv/services/{name-?branch}/v{version}} и там матчить через compareMaps
	//    если хотябы один роут полностью совпал — это наш кандидат на убивание
	// если есть — убиваем в консуле сразу и через 5 минут в марафоне

	println(utils.MapsEqual(map[string]string{"name": "dima", "version": "1.0"}, map[string]string{"version": "1.0", "name": "dima"}))

	log.Println("route")

	return nil
}

func ConsulClient(m *manifest.Manifest) *api.Client {
	conf := api.DefaultConfig()
	conf.Address = m.GetString("consul.consul-host") + ":8500"

	consul, _ := api.NewClient(conf)
	return consul
}
<|endoftext|>


<commit_before>// +build windows

package mem

import (
	"context"
	"unsafe"

	"github.com/shirou/gopsutil/internal/common"
	"golang.org/x/sys/windows"
)

var (
	procGlobalMemoryStatusEx = common.Modkernel32.NewProc("GlobalMemoryStatusEx")
	procGetPerformanceInfo   = common.ModPsapi.NewProc("GetPerformanceInfo")
)

type memoryStatusEx struct {
	cbSize                  uint32
	dwMemoryLoad            uint32
	ullTotalPhys            uint64 // in bytes
	ullAvailPhys            uint64
	ullTotalPageFile        uint64
	ullAvailPageFile        uint64
	ullTotalVirtual         uint64
	ullAvailVirtual         uint64
	ullAvailExtendedVirtual uint64
}

func VirtualMemory() (*VirtualMemoryStat, error) {
	return VirtualMemoryWithContext(context.Background())
}

func VirtualMemoryWithContext(ctx context.Context) (*VirtualMemoryStat, error) {
	var memInfo memoryStatusEx
	memInfo.cbSize = uint32(unsafe.Sizeof(memInfo))
	mem, _, _ := procGlobalMemoryStatusEx.Call(uintptr(unsafe.Pointer(&memInfo)))
	if mem == 0 {
		return nil, windows.GetLastError()
	}

	ret := &VirtualMemoryStat{
		Total:       memInfo.ullTotalPhys,
		Available:   memInfo.ullAvailPhys,
		UsedPercent: float64(memInfo.dwMemoryLoad),
	}

	ret.Used = ret.Total - ret.Available
	return ret, nil
}

type performanceInformation struct {
	cb                uint32
	commitTotal       uint64
	commitLimit       uint64
	commitPeak        uint64
	physicalTotal     uint64
	physicalAvailable uint64
	systemCache       uint64
	kernelTotal       uint64
	kernelPaged       uint64
	kernelNonpaged    uint64
	pageSize          uint64
	handleCount       uint32
	processCount      uint32
	threadCount       uint32
}

func SwapMemory() (*SwapMemoryStat, error) {
	return SwapMemoryWithContext(context.Background())
}

func SwapMemoryWithContext(ctx context.Context) (*SwapMemoryStat, error) {
	var perfInfo performanceInformation
	perfInfo.cb = uint32(unsafe.Sizeof(perfInfo))
	mem, _, _ := procGetPerformanceInfo.Call(uintptr(unsafe.Pointer(&perfInfo)), uintptr(perfInfo.cb))
	if mem == 0 {
		return nil, windows.GetLastError()
	}
	tot := perfInfo.commitLimit * perfInfo.pageSize
	used := perfInfo.commitTotal * perfInfo.pageSize
	free := tot - used
	ret := &SwapMemoryStat{
		Total:       tot,
		Used:        used,
		Free:        free,
		UsedPercent: float64(used / tot),
	}

	return ret, nil
}
<commit_msg>Fix float64 casting<commit_after>// +build windows

package mem

import (
	"context"
	"unsafe"

	"github.com/shirou/gopsutil/internal/common"
	"golang.org/x/sys/windows"
)

var (
	procGlobalMemoryStatusEx = common.Modkernel32.NewProc("GlobalMemoryStatusEx")
	procGetPerformanceInfo   = common.ModPsapi.NewProc("GetPerformanceInfo")
)

type memoryStatusEx struct {
	cbSize                  uint32
	dwMemoryLoad            uint32
	ullTotalPhys            uint64 // in bytes
	ullAvailPhys            uint64
	ullTotalPageFile        uint64
	ullAvailPageFile        uint64
	ullTotalVirtual         uint64
	ullAvailVirtual         uint64
	ullAvailExtendedVirtual uint64
}

func VirtualMemory() (*VirtualMemoryStat, error) {
	return VirtualMemoryWithContext(context.Background())
}

func VirtualMemoryWithContext(ctx context.Context) (*VirtualMemoryStat, error) {
	var memInfo memoryStatusEx
	memInfo.cbSize = uint32(unsafe.Sizeof(memInfo))
	mem, _, _ := procGlobalMemoryStatusEx.Call(uintptr(unsafe.Pointer(&memInfo)))
	if mem == 0 {
		return nil, windows.GetLastError()
	}

	ret := &VirtualMemoryStat{
		Total:       memInfo.ullTotalPhys,
		Available:   memInfo.ullAvailPhys,
		UsedPercent: float64(memInfo.dwMemoryLoad),
	}

	ret.Used = ret.Total - ret.Available
	return ret, nil
}

type performanceInformation struct {
	cb                uint32
	commitTotal       uint64
	commitLimit       uint64
	commitPeak        uint64
	physicalTotal     uint64
	physicalAvailable uint64
	systemCache       uint64
	kernelTotal       uint64
	kernelPaged       uint64
	kernelNonpaged    uint64
	pageSize          uint64
	handleCount       uint32
	processCount      uint32
	threadCount       uint32
}

func SwapMemory() (*SwapMemoryStat, error) {
	return SwapMemoryWithContext(context.Background())
}

func SwapMemoryWithContext(ctx context.Context) (*SwapMemoryStat, error) {
	var perfInfo performanceInformation
	perfInfo.cb = uint32(unsafe.Sizeof(perfInfo))
	mem, _, _ := procGetPerformanceInfo.Call(uintptr(unsafe.Pointer(&perfInfo)), uintptr(perfInfo.cb))
	if mem == 0 {
		return nil, windows.GetLastError()
	}
	tot := perfInfo.commitLimit * perfInfo.pageSize
	used := perfInfo.commitTotal * perfInfo.pageSize
	free := tot - used
	ret := &SwapMemoryStat{
		Total:       tot,
		Used:        used,
		Free:        free,
		UsedPercent: float64(used) / float64(tot),
	}

	return ret, nil
}
<|endoftext|>


<commit_before>package saml2aws

import (
	"fmt"
	"strings"
)

// AWSRole aws role attributes
type AWSRole struct {
	RoleARN      string
	PrincipalARN string
	Name         string
}

// ParseAWSRoles parses and splits the roles while also validating the contents
func ParseAWSRoles(roles []string) ([]*AWSRole, error) {
	awsRoles := make([]*AWSRole, len(roles))

	for i, role := range roles {
		awsRole, err := parseRole(role)
		if err != nil {
			return nil, err
		}

		awsRoles[i] = awsRole
	}

	return awsRoles, nil
}

func parseRole(role string) (*AWSRole, error) {
	tokens := strings.Split(role, ",")

	if len(tokens) != 2 {
		return nil, fmt.Errorf("Invalid role string only %d tokens", len(tokens))
	}

	awsRole := &AWSRole{}

	for _, token := range tokens {
		if strings.Contains(token, ":saml-provider") {
			awsRole.PrincipalARN = token
		}
		if strings.Contains(token, ":role") {
			awsRole.RoleARN = token
		}
	}

	if awsRole.PrincipalARN == "" {
		return nil, fmt.Errorf("Unable to locate PrincipalARN in: %s", role)
	}

	if awsRole.RoleARN == "" {
		return nil, fmt.Errorf("Unable to locate RoleARN in: %s", role)
	}

	return awsRole, nil
}
<commit_msg>Trim whitespace on extracted role ARNs<commit_after>package saml2aws

import (
	"fmt"
	"strings"
)

// AWSRole aws role attributes
type AWSRole struct {
	RoleARN      string
	PrincipalARN string
	Name         string
}

// ParseAWSRoles parses and splits the roles while also validating the contents
func ParseAWSRoles(roles []string) ([]*AWSRole, error) {
	awsRoles := make([]*AWSRole, len(roles))

	for i, role := range roles {
		awsRole, err := parseRole(role)
		if err != nil {
			return nil, err
		}

		awsRoles[i] = awsRole
	}

	return awsRoles, nil
}

func parseRole(role string) (*AWSRole, error) {
	tokens := strings.Split(role, ",")

	if len(tokens) != 2 {
		return nil, fmt.Errorf("Invalid role string only %d tokens", len(tokens))
	}

	awsRole := &AWSRole{}

	for _, token := range tokens {
		if strings.Contains(token, ":saml-provider") {
			awsRole.PrincipalARN = strings.TrimSpace(token)
		}
		if strings.Contains(token, ":role") {
			awsRole.RoleARN = strings.TrimSpace(token)
		}
	}

	if awsRole.PrincipalARN == "" {
		return nil, fmt.Errorf("Unable to locate PrincipalARN in: %s", role)
	}

	if awsRole.RoleARN == "" {
		return nil, fmt.Errorf("Unable to locate RoleARN in: %s", role)
	}

	return awsRole, nil
}
<|endoftext|>


<commit_before>package axe

import (
	"context"
	"fmt"
	"reflect"
	"time"

	"gopkg.in/tomb.v2"

	"github.com/256dpi/fire/cinder"
	"github.com/256dpi/fire/coal"
)

// Task describes work that is managed using a job queue.
type Task struct {
	// Name is the unique name of the task.
	Name string

	// Model is the model that holds task related data.
	Model Model

	// Handler is the callback called with jobs for processing. The handler
	// should return errors formatted with E to properly indicate the status of
	// the job. If a task execution is successful the handler may return some
	// data that is attached to the job.
	Handler func(ctx *Context) error

	// Notifier is a callback that is called once after a job has been completed
	// or cancelled.
	Notifier func(ctx *Context, cancelled bool, reason string) error

	// Workers defines the number for spawned workers that dequeue and execute
	// jobs in parallel.
	//
	// Default: 2.
	Workers int

	// MaxAttempts defines the maximum attempts to complete a task. Zero means
	// that the jobs is retried forever. The error retry field will take
	// precedence to this setting and allow retry beyond the configured maximum.
	//
	// Default: 0
	MaxAttempts int

	// Interval defines the rate at which the worker will request a job from the
	// queue.
	//
	// Default: 100ms.
	Interval time.Duration

	// MinDelay is the minimal time after a failed task is retried.
	//
	// Default: 1s.
	MinDelay time.Duration

	// MaxDelay is the maximal time after a failed task is retried.
	//
	// Default: 10m.
	MaxDelay time.Duration

	// DelayFactor defines the exponential increase of the delay after individual
	// attempts.
	//
	// Default: 2.
	DelayFactor float64

	// Lifetime is the time after which the context of a job is cancelled and
	// the execution should be stopped. Should be several minutes less than
	// timeout to prevent race conditions.
	//
	// Default: 5m.
	Lifetime time.Duration

	// Timeout is the time after which a task can be dequeued again in case the
	// worker was not able to set its status.
	//
	// Default: 10m.
	Timeout time.Duration

	// Periodicity may be set to let the system enqueue a job automatically
	// every given interval.
	//
	// Default: 0.
	Periodicity time.Duration

	// PeriodicJob is the blueprint of the job that is periodically enqueued.
	//
	// Default: Blueprint{Name: Task.Name}.
	PeriodicJob Blueprint
}

func (t *Task) start(q *Queue) {
	// set default workers
	if t.Workers == 0 {
		t.Workers = 2
	}

	// set default interval
	if t.Interval == 0 {
		t.Interval = 100 * time.Millisecond
	}

	// set default minimal delay
	if t.MinDelay == 0 {
		t.MinDelay = time.Second
	}

	// set default maximal delay
	if t.MaxDelay == 0 {
		t.MaxDelay = 10 * time.Minute
	}

	// set default delay factor
	if t.DelayFactor < 1 {
		t.DelayFactor = 2
	}

	// set default lifetime
	if t.Lifetime == 0 {
		t.Lifetime = 5 * time.Minute
	}

	// set default timeout
	if t.Timeout == 0 {
		t.Timeout = 10 * time.Minute
	}

	// check timeout
	if t.Lifetime > t.Timeout {
		panic("axe: lifetime must be less than timeout")
	}

	// start workers for queue
	for i := 0; i < t.Workers; i++ {
		q.tomb.Go(func() error {
			return t.worker(q)
		})
	}

	// run periodic enqueuer if interval is given
	if t.Periodicity > 0 {
		q.tomb.Go(func() error {
			return t.enqueuer(q)
		})
	}
}

func (t *Task) worker(q *Queue) error {
	// run forever
	for {
		// return if queue is closed
		if !q.tomb.Alive() {
			return tomb.ErrDying
		}

		// attempt to get job from queue
		job := q.get(t.Name)
		if job == nil {
			// wait some time before trying again
			select {
			case <-time.After(t.Interval):
			case <-q.tomb.Dying():
				return tomb.ErrDying
			}

			continue
		}

		// execute job and report errors
		err := t.execute(q, job)
		if err != nil {
			if q.opts.Reporter != nil {
				q.opts.Reporter(err)
			}
		}
	}
}

func (t *Task) enqueuer(q *Queue) error {
	// prepare blueprint
	blueprint := t.PeriodicJob

	// override task name
	blueprint.Name = t.Name

	for {
		// enqueue task
		_, err := q.Enqueue(blueprint)
		if err != nil && q.opts.Reporter != nil {
			// report error
			q.opts.Reporter(err)

			// wait some time
			select {
			case <-time.After(time.Second):
			case <-q.tomb.Dying():
				return tomb.ErrDying
			}

			continue
		}

		// wait for next interval
		select {
		case <-time.After(t.Periodicity):
		case <-q.tomb.Dying():
			return tomb.ErrDying
		}
	}
}

func (t *Task) execute(q *Queue, job *Job) error {
	// create trace
	trace, outerContext := cinder.CreateTrace(context.Background(), t.Name)
	defer trace.Finish()

	// dequeue job
	job, err := Dequeue(outerContext, q.opts.Store, job.ID(), t.Timeout)
	if err != nil {
		return err
	}

	// return if missing (might be dequeued already by another process)
	if job == nil {
		return nil
	}

	// get time
	start := time.Now()

	// prepare model
	var model Model

	// check model
	if t.Model != nil {
		// instantiate model
		model = reflect.New(reflect.TypeOf(t.Model).Elem()).Interface()

		// unmarshal model
		err = job.Data.Unmarshal(model, coal.TransferBSON)
		if err != nil {
			return err
		}
	}

	// add timeout
	innerContext, cancel := context.WithTimeout(outerContext, t.Lifetime)
	defer cancel()

	// prepare context
	ctx := &Context{
		Context: innerContext,
		Model:   model,
		Attempt: job.Attempts, // incremented when dequeued
		Task:    t,
		Queue:   q,
		Trace:   trace,
	}

	// run handler
	err = t.Handler(ctx)

	// return immediately if lifetime has been reached
	if time.Since(start) > t.Lifetime {
		return fmt.Errorf(`task "%s" ran longer than the specified lifetime`, t.Name)
	}

	// check error
	if e, ok := err.(*Error); ok {
		// check retry
		if e.Retry {
			// fail job
			err = Fail(outerContext, q.opts.Store, job.ID(), e.Reason, Backoff(t.MinDelay, t.MaxDelay, t.DelayFactor, job.Attempts))
			if err != nil {
				return err
			}

			return nil
		}

		// cancel job
		err = Cancel(outerContext, q.opts.Store, job.ID(), e.Reason)
		if err != nil {
			return err
		}

		// call notifier if available
		if t.Notifier != nil {
			err = t.Notifier(ctx, true, e.Reason)
			if err != nil {
				return err
			}
		}

		return nil
	}

	// handle other errors
	if err != nil {
		// check attempts
		if t.MaxAttempts == 0 || job.Attempts < t.MaxAttempts {
			// fail job
			_ = Fail(outerContext, q.opts.Store, job.ID(), err.Error(), Backoff(t.MinDelay, t.MaxDelay, t.DelayFactor, job.Attempts))

			return err
		}

		// cancel job
		_ = Cancel(outerContext, q.opts.Store, job.ID(), err.Error())

		// call notifier if available
		if t.Notifier != nil {
			_ = t.Notifier(ctx, true, err.Error())
		}

		return err
	}

	// complete job
	err = Complete(outerContext, q.opts.Store, job.ID(), ctx.Result)
	if err != nil {
		return err
	}

	// call notifier if available
	if t.Notifier != nil {
		err = t.Notifier(ctx, false, "")
		if err != nil {
			return err
		}
	}

	return nil
}
<commit_msg>renamed parameters<commit_after>package axe

import (
	"context"
	"fmt"
	"reflect"
	"time"

	"gopkg.in/tomb.v2"

	"github.com/256dpi/fire/cinder"
	"github.com/256dpi/fire/coal"
)

// Task describes work that is managed using a job queue.
type Task struct {
	// Name is the unique name of the task.
	Name string

	// Model is the model that holds task related data.
	Model Model

	// Handler is the callback called with jobs for processing. The handler
	// should return errors formatted with E to properly indicate the status of
	// the job. If a task execution is successful the handler may return some
	// data that is attached to the job.
	Handler func(ctx *Context) error

	// Notifier is a callback that is called once after a job has been completed
	// or cancelled.
	Notifier func(ctx *Context, cancelled bool, reason string) error

	// Workers defines the number for spawned workers that dequeue and execute
	// jobs in parallel.
	//
	// Default: 2.
	Workers int

	// MaxAttempts defines the maximum attempts to complete a task. Zero means
	// that the jobs is retried forever. The error retry field will take
	// precedence to this setting and allow retry beyond the configured maximum.
	//
	// Default: 0
	MaxAttempts int

	// Interval defines the rate at which the worker will request a job from the
	// queue.
	//
	// Default: 100ms.
	Interval time.Duration

	// MinDelay is the minimal time after a failed task is retried.
	//
	// Default: 1s.
	MinDelay time.Duration

	// MaxDelay is the maximal time after a failed task is retried.
	//
	// Default: 10m.
	MaxDelay time.Duration

	// DelayFactor defines the exponential increase of the delay after individual
	// attempts.
	//
	// Default: 2.
	DelayFactor float64

	// Lifetime is the time after which the context of a job is cancelled and
	// the execution should be stopped. Should be several minutes less than
	// timeout to prevent race conditions.
	//
	// Default: 5m.
	Lifetime time.Duration

	// Timeout is the time after which a task can be dequeued again in case the
	// worker was not able to set its status.
	//
	// Default: 10m.
	Timeout time.Duration

	// Periodicity may be set to let the system enqueue a job automatically
	// every given interval.
	//
	// Default: 0.
	Periodicity time.Duration

	// PeriodicJob is the blueprint of the job that is periodically enqueued.
	//
	// Default: Blueprint{Name: Task.Name}.
	PeriodicJob Blueprint
}

func (t *Task) start(queue *Queue) {
	// set default workers
	if t.Workers == 0 {
		t.Workers = 2
	}

	// set default interval
	if t.Interval == 0 {
		t.Interval = 100 * time.Millisecond
	}

	// set default minimal delay
	if t.MinDelay == 0 {
		t.MinDelay = time.Second
	}

	// set default maximal delay
	if t.MaxDelay == 0 {
		t.MaxDelay = 10 * time.Minute
	}

	// set default delay factor
	if t.DelayFactor < 1 {
		t.DelayFactor = 2
	}

	// set default lifetime
	if t.Lifetime == 0 {
		t.Lifetime = 5 * time.Minute
	}

	// set default timeout
	if t.Timeout == 0 {
		t.Timeout = 10 * time.Minute
	}

	// check timeout
	if t.Lifetime > t.Timeout {
		panic("axe: lifetime must be less than timeout")
	}

	// start workers for queue
	for i := 0; i < t.Workers; i++ {
		queue.tomb.Go(func() error {
			return t.worker(queue)
		})
	}

	// run periodic enqueuer if interval is given
	if t.Periodicity > 0 {
		queue.tomb.Go(func() error {
			return t.enqueuer(queue)
		})
	}
}

func (t *Task) worker(queue *Queue) error {
	// run forever
	for {
		// return if queue is closed
		if !queue.tomb.Alive() {
			return tomb.ErrDying
		}

		// attempt to get job from queue
		job := queue.get(t.Name)
		if job == nil {
			// wait some time before trying again
			select {
			case <-time.After(t.Interval):
			case <-queue.tomb.Dying():
				return tomb.ErrDying
			}

			continue
		}

		// execute job and report errors
		err := t.execute(queue, job)
		if err != nil {
			if queue.opts.Reporter != nil {
				queue.opts.Reporter(err)
			}
		}
	}
}

func (t *Task) enqueuer(queue *Queue) error {
	// prepare blueprint
	blueprint := t.PeriodicJob

	// override task name
	blueprint.Name = t.Name

	for {
		// enqueue task
		_, err := queue.Enqueue(blueprint)
		if err != nil && queue.opts.Reporter != nil {
			// report error
			queue.opts.Reporter(err)

			// wait some time
			select {
			case <-time.After(time.Second):
			case <-queue.tomb.Dying():
				return tomb.ErrDying
			}

			continue
		}

		// wait for next interval
		select {
		case <-time.After(t.Periodicity):
		case <-queue.tomb.Dying():
			return tomb.ErrDying
		}
	}
}

func (t *Task) execute(queue *Queue, job *Job) error {
	// create trace
	trace, outerContext := cinder.CreateTrace(context.Background(), t.Name)
	defer trace.Finish()

	// dequeue job
	job, err := Dequeue(outerContext, queue.opts.Store, job.ID(), t.Timeout)
	if err != nil {
		return err
	}

	// return if missing (might be dequeued already by another process)
	if job == nil {
		return nil
	}

	// get time
	start := time.Now()

	// prepare model
	var model Model

	// check model
	if t.Model != nil {
		// instantiate model
		model = reflect.New(reflect.TypeOf(t.Model).Elem()).Interface()

		// unmarshal model
		err = job.Data.Unmarshal(model, coal.TransferBSON)
		if err != nil {
			return err
		}
	}

	// add timeout
	innerContext, cancel := context.WithTimeout(outerContext, t.Lifetime)
	defer cancel()

	// prepare context
	ctx := &Context{
		Context: innerContext,
		Model:   model,
		Attempt: job.Attempts, // incremented when dequeued
		Task:    t,
		Queue:   queue,
		Trace:   trace,
	}

	// run handler
	err = t.Handler(ctx)

	// return immediately if lifetime has been reached
	if time.Since(start) > t.Lifetime {
		return fmt.Errorf(`task "%s" ran longer than the specified lifetime`, t.Name)
	}

	// check error
	if e, ok := err.(*Error); ok {
		// check retry
		if e.Retry {
			// fail job
			err = Fail(outerContext, queue.opts.Store, job.ID(), e.Reason, Backoff(t.MinDelay, t.MaxDelay, t.DelayFactor, job.Attempts))
			if err != nil {
				return err
			}

			return nil
		}

		// cancel job
		err = Cancel(outerContext, queue.opts.Store, job.ID(), e.Reason)
		if err != nil {
			return err
		}

		// call notifier if available
		if t.Notifier != nil {
			err = t.Notifier(ctx, true, e.Reason)
			if err != nil {
				return err
			}
		}

		return nil
	}

	// handle other errors
	if err != nil {
		// check attempts
		if t.MaxAttempts == 0 || job.Attempts < t.MaxAttempts {
			// fail job
			_ = Fail(outerContext, queue.opts.Store, job.ID(), err.Error(), Backoff(t.MinDelay, t.MaxDelay, t.DelayFactor, job.Attempts))

			return err
		}

		// cancel job
		_ = Cancel(outerContext, queue.opts.Store, job.ID(), err.Error())

		// call notifier if available
		if t.Notifier != nil {
			_ = t.Notifier(ctx, true, err.Error())
		}

		return err
	}

	// complete job
	err = Complete(outerContext, queue.opts.Store, job.ID(), ctx.Result)
	if err != nil {
		return err
	}

	// call notifier if available
	if t.Notifier != nil {
		err = t.Notifier(ctx, false, "")
		if err != nil {
			return err
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package axe

import (
	"reflect"
	"time"

	"gopkg.in/tomb.v2"

	"github.com/256dpi/fire"
	"github.com/256dpi/fire/coal"
)

// Error is used to signal failed job executions.
type Error struct {
	Reason string
	Retry  bool
}

// Error implements the error interface.
func (c *Error) Error() string {
	return c.Reason
}

// E is a short-hand to construct an error.
func E(reason string, retry bool) *Error {
	return &Error{
		Reason: reason,
		Retry:  retry,
	}
}

// Model can be any BSON serializable type.
type Model interface{}

// Context holds and stores contextual data.
type Context struct {
	// Model is the model carried by the job.
	Model Model

	// Result can be set with a custom result.
	Result coal.Map

	// Task is the task that processes this job.
	//
	// Usage: Read Only
	Task *Task

	// Queue is the queue this job was dequeued from.
	//
	// Usage: Read Only
	Queue *Queue

	// Store is the store used by the queue.
	//
	// Usage: Read Only
	Store *coal.Store

	// The tracer used to trace code execution.
	//
	// Usage: Read Only
	Tracer *fire.Tracer
}

// TC is a shorthand to get a traced collection for the specified model.
func (c *Context) TC(model coal.Model) *coal.TracedCollection {
	return c.Store.TC(c.Tracer, model)
}

// Task describes work that is managed using a job queue.
type Task struct {
	// Name is the unique name of the task.
	Name string

	// Model is the model that holds task related data.
	Model Model

	// Handler is the callback called with jobs for processing. The handler
	// should return errors formatted with E to properly indicate the status of
	// the job. If a task execution is successful the handler may return some
	// data that is attached to the job.
	Handler func(*Context) error

	// Workers defines the number for spawned workers that dequeue and execute
	// jobs in parallel.
	//
	// Default: 2.
	Workers int

	// MaxAttempts defines the maximum attempts to complete a task. Zero means
	// that the jobs is retried forever. The error retry field will take
	// precedence to this setting and allow retry beyond the configured maximum.
	//
	// Default: 0
	MaxAttempts int

	// Interval defines the rate at which the worker will request a job from the
	// queue.
	//
	// Default: 100ms.
	Interval time.Duration

	// MinDelay is the minimal time after a failed task is retried.
	//
	// Default: 1s.
	MinDelay time.Duration

	// MaxDelay is the maximal time after a failed task is retried.
	//
	// Default: 10m.
	MaxDelay time.Duration

	// DelayFactor defines the exponential increase of the delay after individual
	// attempts.
	//
	// Default: 2.
	DelayFactor float64

	// Timeout is the time after which a task can be dequeue again in case the
	// worker was not able to set its status.
	//
	// Default: 10m.
	Timeout time.Duration

	// Periodically may be set to let the system enqueue a job automatically
	// every given interval.
	//
	// Default: 0.
	Periodically time.Duration

	// PeriodicJob is the blueprint of the job that is periodically enqueued.
	//
	// Default: Blueprint{Name: Task.Name}.
	PeriodicJob Blueprint
}

func (t *Task) start(q *Queue) {
	// set default workers
	if t.Workers == 0 {
		t.Workers = 2
	}

	// set default interval
	if t.Interval == 0 {
		t.Interval = 100 * time.Millisecond
	}

	// set default minimal delay
	if t.MinDelay == 0 {
		t.MinDelay = time.Second
	}

	// set default maximal delay
	if t.MaxDelay == 0 {
		t.MaxDelay = 10 * time.Minute
	}

	// set default delay factor
	if t.DelayFactor <= 1 {
		t.DelayFactor = 2
	}

	// set default timeout
	if t.Timeout == 0 {
		t.Timeout = 10 * time.Minute
	}

	// start workers for queue
	for i := 0; i < t.Workers; i++ {
		q.tomb.Go(func() error {
			return t.worker(q)
		})
	}

	// run periodic enqueuer if interval is given
	if t.Periodically > 0 {
		q.tomb.Go(func() error {
			return t.enqueuer(q)
		})
	}
}

func (t *Task) worker(q *Queue) error {
	// run forever
	for {
		// return if queue is closed
		if !q.tomb.Alive() {
			return tomb.ErrDying
		}

		// attempt to get job from queue
		job := q.get(t.Name)
		if job == nil {
			// wait some time before trying again
			select {
			case <-time.After(t.Interval):
			case <-q.tomb.Dying():
				return tomb.ErrDying
			}

			continue
		}

		// execute job and report errors
		err := t.execute(q, job)
		if err != nil {
			if q.reporter != nil {
				q.reporter(err)
			}
		}
	}
}

func (t *Task) enqueuer(q *Queue) error {
	// prepare blueprint
	blueprint := t.PeriodicJob

	// override task name
	blueprint.Name = t.Name

	for {
		// enqueue task
		_, err := q.Enqueue(blueprint)
		if err != nil && q.reporter != nil {
			// report error
			q.reporter(err)

			// wait some time
			select {
			case <-time.After(time.Second):
			case <-q.tomb.Dying():
				return tomb.ErrDying
			}

			continue
		}

		// wait for next interval
		select {
		case <-time.After(t.Periodically):
		case <-q.tomb.Dying():
			return tomb.ErrDying
		}
	}
}

func (t *Task) execute(q *Queue, job *Job) error {
	// dequeue job
	job, err := Dequeue(q.store, job.ID(), t.Timeout)
	if err != nil {
		return err
	}

	// return if missing (might be dequeued already by another process)
	if job == nil {
		return nil
	}

	// prepare model
	var model Model

	// check model
	if t.Model != nil {
		// instantiate model
		model = reflect.New(reflect.TypeOf(t.Model).Elem()).Interface()

		// unmarshal model
		err = job.Data.Unmarshal(model)
		if err != nil {
			return err
		}
	}

	// create tracer
	tracer := fire.NewTracerWithRoot(t.Name)
	defer tracer.Finish(true)

	// prepare context
	ctx := &Context{
		Model:  model,
		Task:   t,
		Queue:  q,
		Store:  q.store,
		Tracer: tracer,
	}

	// run handler
	err = t.Handler(ctx)

	// check error
	if e, ok := err.(*Error); ok {
		// check retry
		if e.Retry {
			// fail job
			err = Fail(q.store, job.ID(), e.Reason, Backoff(t.MinDelay, t.MaxDelay, t.DelayFactor, job.Attempts))
			if err != nil {
				return err
			}

			return nil
		}

		// cancel job
		err = Cancel(q.store, job.ID(), e.Reason)
		if err != nil {
			return err
		}

		return nil
	}

	// handle other errors
	if err != nil {
		// check attempts
		if t.MaxAttempts == 0 || job.Attempts < t.MaxAttempts {
			// fail job
			_ = Fail(q.store, job.ID(), err.Error(), Backoff(t.MinDelay, t.MaxDelay, t.DelayFactor, job.Attempts))

			return err
		}

		// cancel job
		_ = Cancel(q.store, job.ID(), err.Error())

		return err
	}

	// complete job
	err = Complete(q.store, job.ID(), ctx.Result)
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>one is a valid delay factor<commit_after>package axe

import (
	"reflect"
	"time"

	"gopkg.in/tomb.v2"

	"github.com/256dpi/fire"
	"github.com/256dpi/fire/coal"
)

// Error is used to signal failed job executions.
type Error struct {
	Reason string
	Retry  bool
}

// Error implements the error interface.
func (c *Error) Error() string {
	return c.Reason
}

// E is a short-hand to construct an error.
func E(reason string, retry bool) *Error {
	return &Error{
		Reason: reason,
		Retry:  retry,
	}
}

// Model can be any BSON serializable type.
type Model interface{}

// Context holds and stores contextual data.
type Context struct {
	// Model is the model carried by the job.
	Model Model

	// Result can be set with a custom result.
	Result coal.Map

	// Task is the task that processes this job.
	//
	// Usage: Read Only
	Task *Task

	// Queue is the queue this job was dequeued from.
	//
	// Usage: Read Only
	Queue *Queue

	// Store is the store used by the queue.
	//
	// Usage: Read Only
	Store *coal.Store

	// The tracer used to trace code execution.
	//
	// Usage: Read Only
	Tracer *fire.Tracer
}

// TC is a shorthand to get a traced collection for the specified model.
func (c *Context) TC(model coal.Model) *coal.TracedCollection {
	return c.Store.TC(c.Tracer, model)
}

// Task describes work that is managed using a job queue.
type Task struct {
	// Name is the unique name of the task.
	Name string

	// Model is the model that holds task related data.
	Model Model

	// Handler is the callback called with jobs for processing. The handler
	// should return errors formatted with E to properly indicate the status of
	// the job. If a task execution is successful the handler may return some
	// data that is attached to the job.
	Handler func(*Context) error

	// Workers defines the number for spawned workers that dequeue and execute
	// jobs in parallel.
	//
	// Default: 2.
	Workers int

	// MaxAttempts defines the maximum attempts to complete a task. Zero means
	// that the jobs is retried forever. The error retry field will take
	// precedence to this setting and allow retry beyond the configured maximum.
	//
	// Default: 0
	MaxAttempts int

	// Interval defines the rate at which the worker will request a job from the
	// queue.
	//
	// Default: 100ms.
	Interval time.Duration

	// MinDelay is the minimal time after a failed task is retried.
	//
	// Default: 1s.
	MinDelay time.Duration

	// MaxDelay is the maximal time after a failed task is retried.
	//
	// Default: 10m.
	MaxDelay time.Duration

	// DelayFactor defines the exponential increase of the delay after individual
	// attempts.
	//
	// Default: 2.
	DelayFactor float64

	// Timeout is the time after which a task can be dequeue again in case the
	// worker was not able to set its status.
	//
	// Default: 10m.
	Timeout time.Duration

	// Periodically may be set to let the system enqueue a job automatically
	// every given interval.
	//
	// Default: 0.
	Periodically time.Duration

	// PeriodicJob is the blueprint of the job that is periodically enqueued.
	//
	// Default: Blueprint{Name: Task.Name}.
	PeriodicJob Blueprint
}

func (t *Task) start(q *Queue) {
	// set default workers
	if t.Workers == 0 {
		t.Workers = 2
	}

	// set default interval
	if t.Interval == 0 {
		t.Interval = 100 * time.Millisecond
	}

	// set default minimal delay
	if t.MinDelay == 0 {
		t.MinDelay = time.Second
	}

	// set default maximal delay
	if t.MaxDelay == 0 {
		t.MaxDelay = 10 * time.Minute
	}

	// set default delay factor
	if t.DelayFactor < 1 {
		t.DelayFactor = 2
	}

	// set default timeout
	if t.Timeout == 0 {
		t.Timeout = 10 * time.Minute
	}

	// start workers for queue
	for i := 0; i < t.Workers; i++ {
		q.tomb.Go(func() error {
			return t.worker(q)
		})
	}

	// run periodic enqueuer if interval is given
	if t.Periodically > 0 {
		q.tomb.Go(func() error {
			return t.enqueuer(q)
		})
	}
}

func (t *Task) worker(q *Queue) error {
	// run forever
	for {
		// return if queue is closed
		if !q.tomb.Alive() {
			return tomb.ErrDying
		}

		// attempt to get job from queue
		job := q.get(t.Name)
		if job == nil {
			// wait some time before trying again
			select {
			case <-time.After(t.Interval):
			case <-q.tomb.Dying():
				return tomb.ErrDying
			}

			continue
		}

		// execute job and report errors
		err := t.execute(q, job)
		if err != nil {
			if q.reporter != nil {
				q.reporter(err)
			}
		}
	}
}

func (t *Task) enqueuer(q *Queue) error {
	// prepare blueprint
	blueprint := t.PeriodicJob

	// override task name
	blueprint.Name = t.Name

	for {
		// enqueue task
		_, err := q.Enqueue(blueprint)
		if err != nil && q.reporter != nil {
			// report error
			q.reporter(err)

			// wait some time
			select {
			case <-time.After(time.Second):
			case <-q.tomb.Dying():
				return tomb.ErrDying
			}

			continue
		}

		// wait for next interval
		select {
		case <-time.After(t.Periodically):
		case <-q.tomb.Dying():
			return tomb.ErrDying
		}
	}
}

func (t *Task) execute(q *Queue, job *Job) error {
	// dequeue job
	job, err := Dequeue(q.store, job.ID(), t.Timeout)
	if err != nil {
		return err
	}

	// return if missing (might be dequeued already by another process)
	if job == nil {
		return nil
	}

	// prepare model
	var model Model

	// check model
	if t.Model != nil {
		// instantiate model
		model = reflect.New(reflect.TypeOf(t.Model).Elem()).Interface()

		// unmarshal model
		err = job.Data.Unmarshal(model)
		if err != nil {
			return err
		}
	}

	// create tracer
	tracer := fire.NewTracerWithRoot(t.Name)
	defer tracer.Finish(true)

	// prepare context
	ctx := &Context{
		Model:  model,
		Task:   t,
		Queue:  q,
		Store:  q.store,
		Tracer: tracer,
	}

	// run handler
	err = t.Handler(ctx)

	// check error
	if e, ok := err.(*Error); ok {
		// check retry
		if e.Retry {
			// fail job
			err = Fail(q.store, job.ID(), e.Reason, Backoff(t.MinDelay, t.MaxDelay, t.DelayFactor, job.Attempts))
			if err != nil {
				return err
			}

			return nil
		}

		// cancel job
		err = Cancel(q.store, job.ID(), e.Reason)
		if err != nil {
			return err
		}

		return nil
	}

	// handle other errors
	if err != nil {
		// check attempts
		if t.MaxAttempts == 0 || job.Attempts < t.MaxAttempts {
			// fail job
			_ = Fail(q.store, job.ID(), err.Error(), Backoff(t.MinDelay, t.MaxDelay, t.DelayFactor, job.Attempts))

			return err
		}

		// cancel job
		_ = Cancel(q.store, job.ID(), err.Error())

		return err
	}

	// complete job
	err = Complete(q.store, job.ID(), ctx.Result)
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"path/filepath"

	"github.com/xeipuuv/gojsonschema"
)

func main() {
	if len(os.Args[1:]) != 2 {
		fmt.Printf("ERROR: usage is: %s <schema.json> <config.json>\n", os.Args[0])
		os.Exit(1)
	}

	schemaPath, err := filepath.Abs(os.Args[1])
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	documentPath, err := filepath.Abs(os.Args[2])
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	schemaLoader := gojsonschema.NewReferenceLoader("file://" + schemaPath)
	documentLoader := gojsonschema.NewReferenceLoader("file://" + documentPath)

	result, err := gojsonschema.Validate(schemaLoader, documentLoader)
	if err != nil {
		panic(err.Error())
	}

	if result.Valid() {
		fmt.Printf("The document is valid\n")
	} else {
		fmt.Printf("The document is not valid. see errors :\n")
		for _, desc := range result.Errors() {
			fmt.Printf("- %s\n", desc)
		}
		os.Exit(1)
	}
}
<commit_msg>schema/validate: Support reading documents via stdin<commit_after>package main

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

	"github.com/xeipuuv/gojsonschema"
)

func main() {
	nargs := len(os.Args[1:])
	if nargs == 0 || nargs > 2 {
		fmt.Printf("ERROR: usage is: %s <schema.json> [<document.json>]\n", os.Args[0])
		os.Exit(1)
	}

	schemaPath, err := filepath.Abs(os.Args[1])
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	schemaLoader := gojsonschema.NewReferenceLoader("file://" + schemaPath)
	var documentLoader gojsonschema.JSONLoader

	if nargs > 1 {
		documentPath, err := filepath.Abs(os.Args[2])
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		documentLoader = gojsonschema.NewReferenceLoader("file://" + documentPath)
	} else {
		documentBytes, err := ioutil.ReadAll(os.Stdin)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		documentString := string(documentBytes)
		documentLoader = gojsonschema.NewStringLoader(documentString)
	}

	result, err := gojsonschema.Validate(schemaLoader, documentLoader)
	if err != nil {
		panic(err.Error())
	}

	if result.Valid() {
		fmt.Printf("The document is valid\n")
	} else {
		fmt.Printf("The document is not valid. see errors :\n")
		for _, desc := range result.Errors() {
			fmt.Printf("- %s\n", desc)
		}
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package scraper

import (
	"bytes"
	"net/url"
	"os"
	"regexp"

	"github.com/cornelk/goscrape/appcontext"

	"github.com/headzoo/surf"
	"github.com/headzoo/surf/agent"
	"github.com/headzoo/surf/browser"
	"github.com/uber-go/zap"
)

type (
	// Scraper contains all scraping data
	Scraper struct {
		ImageQuality uint
		MaxDepth     uint
		URL          *url.URL

		browser  *browser.Browser
		excludes []*regexp.Regexp
		log      zap.Logger

		assets         map[string]bool
		assetsExternal map[string]bool
		pages          map[string]bool
	}
)

// New creates a new Scraper instance
func New(URL string) (*Scraper, error) {
	u, err := url.Parse(URL)
	if err != nil {
		return nil, err
	}

	b := surf.NewBrowser()
	b.SetUserAgent(agent.GoogleBot())

	s := &Scraper{
		browser:        b,
		log:            appcontext.Logger,
		assets:         make(map[string]bool),
		assetsExternal: make(map[string]bool),
		pages:          make(map[string]bool),
		URL:            u,
	}
	return s, nil
}

// SetExcludes sets and checks the exclusions regular expressions
func (s *Scraper) SetExcludes(excludes []string) error {
	for _, e := range excludes {
		re, err := regexp.Compile(e)
		if err != nil {
			return err
		}

		s.excludes = append(s.excludes, re)
		s.log.Debug("Excluding", zap.Stringer("RE", re))
	}

	return nil
}

// Start starts the scraping
func (s *Scraper) Start() error {
	p := s.URL.Path
	if p == "" {
		p = "/"
	}
	s.pages[p] = false
	return s.scrapeURL(s.URL, 0)
}

func (s *Scraper) scrapeURL(URL *url.URL, currentDepth uint) error {
	s.log.Info("Downloading", zap.Stringer("URL", URL))
	err := s.browser.Open(URL.String())
	if err != nil {
		return err
	}

	buf := &bytes.Buffer{}
	_, err = s.browser.Download(buf)
	if err != nil {
		return err
	}

	html, err := s.fixFileReferences(URL, buf)
	if err != nil {
		return err
	}

	buf = bytes.NewBufferString(html)
	filePath := s.GetFilePath(URL, true)
	err = s.writeFile(filePath, buf) // always update html files, content might have changed
	if err != nil {
		return err
	}

	for _, stylesheet := range s.browser.Stylesheets() {
		err = s.downloadAssetURL(&stylesheet.DownloadableAsset)
		if err != nil {
			return nil
		}
	}
	for _, script := range s.browser.Scripts() {
		err = s.downloadAssetURL(&script.DownloadableAsset)
		if err != nil {
			return nil
		}
	}
	for _, image := range s.browser.Images() {
		err = s.downloadAssetURL(&image.DownloadableAsset)
		if err != nil {
			return nil
		}
	}

	var toScrape []*url.URL
	// check first and download afterwards to not hit max depth limit for start page links because of recursive linking
	for _, link := range s.browser.Links() {
		if s.checkPageURL(link.URL, currentDepth) {
			toScrape = append(toScrape, link.URL)
		}
	}

	for _, URL := range toScrape {
		err = s.scrapeURL(URL, currentDepth+1)
		if err != nil {
			return err
		}
	}
	return nil
}

// checkPageURL checks if a page should be downloaded
func (s *Scraper) checkPageURL(URL *url.URL, currentDepth uint) bool {
	if URL.Host != s.URL.Host {
		s.log.Debug("Skipping external host page", zap.Stringer("URL", URL))
		return false
	}

	p := URL.Path
	if p == "" {
		p = "/"
	}

	_, ok := s.pages[p]
	if ok { // was already downloaded or checked
		s.log.Debug("Skipping already checked page", zap.Stringer("URL", URL))
		return false
	}

	s.pages[p] = false
	if s.MaxDepth != 0 && currentDepth == s.MaxDepth {
		s.log.Debug("Skipping too deep level page", zap.Stringer("URL", URL))
		return false
	}

	if s.isURLExcluded(URL) {
		return false
	}

	s.log.Debug("New page to queue", zap.Stringer("URL", URL))
	return true
}

// downloadAssetURL downloads an asset if it does not exist on disk yet.
func (s *Scraper) downloadAssetURL(asset *browser.DownloadableAsset) error {
	URL := asset.URL

	if URL.Host == s.URL.Host {
		_, ok := s.assets[URL.Path]
		if ok { // was already downloaded or checked
			return nil
		}

		s.assets[URL.Path] = false
	} else {
		if s.isExternalFileChecked(URL) {
			return nil
		}
	}

	if s.isURLExcluded(URL) {
		return nil
	}

	filePath := s.GetFilePath(URL, false)
	if _, err := os.Stat(filePath); !os.IsNotExist(err) {
		return nil
	}

	s.log.Info("Downloading", zap.Stringer("URL", URL))

	buf := &bytes.Buffer{}
	_, err := asset.Download(buf)
	if err != nil {
		return err
	}

	buf = s.checkFileTypeForRecode(filePath, buf)

	return s.writeFile(filePath, buf)
}

func (s *Scraper) isURLExcluded(URL *url.URL) bool {
	for _, re := range s.excludes {
		if re.MatchString(URL.Path) {
			s.log.Info("Skipping URL", zap.Stringer("URL", URL), zap.Stringer("Excluder", re))
			return true
		}
	}
	return false
}

func (s *Scraper) isExternalFileChecked(URL *url.URL) bool {
	if URL.Host == s.URL.Host {
		return false
	}

	fullURL := URL.String()
	_, ok := s.assetsExternal[fullURL]
	if ok { // was already downloaded or checked
		return true
	}

	s.assetsExternal[fullURL] = true
	s.log.Info("External URL", zap.Stringer("URL", URL))

	return false
}
<commit_msg>Handle missing URL schemes and site redirections<commit_after>package scraper

import (
	"bytes"
	"net/url"
	"os"
	"regexp"

	"github.com/cornelk/goscrape/appcontext"

	"github.com/headzoo/surf"
	"github.com/headzoo/surf/agent"
	"github.com/headzoo/surf/browser"
	"github.com/uber-go/zap"
)

type (
	// Scraper contains all scraping data
	Scraper struct {
		ImageQuality uint
		MaxDepth     uint
		URL          *url.URL

		browser  *browser.Browser
		excludes []*regexp.Regexp
		log      zap.Logger

		assets         map[string]bool
		assetsExternal map[string]bool
		pages          map[string]bool
	}
)

// New creates a new Scraper instance
func New(URL string) (*Scraper, error) {
	u, err := url.Parse(URL)
	if err != nil {
		return nil, err
	}

	if u.Scheme == "" {
		u.Scheme = "http" // if no URL scheme was given default to http
	}

	b := surf.NewBrowser()
	b.SetUserAgent(agent.GoogleBot())

	s := &Scraper{
		browser:        b,
		log:            appcontext.Logger,
		assets:         make(map[string]bool),
		assetsExternal: make(map[string]bool),
		pages:          make(map[string]bool),
		URL:            u,
	}
	return s, nil
}

// SetExcludes sets and checks the exclusions regular expressions
func (s *Scraper) SetExcludes(excludes []string) error {
	for _, e := range excludes {
		re, err := regexp.Compile(e)
		if err != nil {
			return err
		}

		s.excludes = append(s.excludes, re)
		s.log.Debug("Excluding", zap.Stringer("RE", re))
	}

	return nil
}

// Start starts the scraping
func (s *Scraper) Start() error {
	p := s.URL.Path
	if p == "" {
		p = "/"
	}
	s.pages[p] = false
	return s.scrapeURL(s.URL, 0)
}

func (s *Scraper) scrapeURL(URL *url.URL, currentDepth uint) error {
	s.log.Info("Downloading", zap.Stringer("URL", URL))
	err := s.browser.Open(URL.String())
	if err != nil {
		return err
	}

	buf := &bytes.Buffer{}
	_, err = s.browser.Download(buf)
	if err != nil {
		return err
	}

	if currentDepth == 0 {
		URL = s.browser.Url() // use the URL that the website returned as new base url for the scrape, in case of a redirect
		s.URL = URL
	}

	html, err := s.fixFileReferences(URL, buf)
	if err != nil {
		return err
	}

	buf = bytes.NewBufferString(html)
	filePath := s.GetFilePath(URL, true)
	err = s.writeFile(filePath, buf) // always update html files, content might have changed
	if err != nil {
		return err
	}

	for _, stylesheet := range s.browser.Stylesheets() {
		err = s.downloadAssetURL(&stylesheet.DownloadableAsset)
		if err != nil {
			return nil
		}
	}
	for _, script := range s.browser.Scripts() {
		err = s.downloadAssetURL(&script.DownloadableAsset)
		if err != nil {
			return nil
		}
	}
	for _, image := range s.browser.Images() {
		err = s.downloadAssetURL(&image.DownloadableAsset)
		if err != nil {
			return nil
		}
	}

	var toScrape []*url.URL
	// check first and download afterwards to not hit max depth limit for start page links because of recursive linking
	for _, link := range s.browser.Links() {
		if s.checkPageURL(link.URL, currentDepth) {
			toScrape = append(toScrape, link.URL)
		}
	}

	for _, URL := range toScrape {
		err = s.scrapeURL(URL, currentDepth+1)
		if err != nil {
			return err
		}
	}
	return nil
}

// checkPageURL checks if a page should be downloaded
func (s *Scraper) checkPageURL(URL *url.URL, currentDepth uint) bool {
	if URL.Host != s.URL.Host {
		s.log.Debug("Skipping external host page", zap.Stringer("URL", URL))
		return false
	}

	p := URL.Path
	if p == "" {
		p = "/"
	}

	_, ok := s.pages[p]
	if ok { // was already downloaded or checked
		s.log.Debug("Skipping already checked page", zap.Stringer("URL", URL))
		return false
	}

	s.pages[p] = false
	if s.MaxDepth != 0 && currentDepth == s.MaxDepth {
		s.log.Debug("Skipping too deep level page", zap.Stringer("URL", URL))
		return false
	}

	if s.isURLExcluded(URL) {
		return false
	}

	s.log.Debug("New page to queue", zap.Stringer("URL", URL))
	return true
}

// downloadAssetURL downloads an asset if it does not exist on disk yet.
func (s *Scraper) downloadAssetURL(asset *browser.DownloadableAsset) error {
	URL := asset.URL

	if URL.Host == s.URL.Host {
		_, ok := s.assets[URL.Path]
		if ok { // was already downloaded or checked
			return nil
		}

		s.assets[URL.Path] = false
	} else {
		if s.isExternalFileChecked(URL) {
			return nil
		}
	}

	if s.isURLExcluded(URL) {
		return nil
	}

	filePath := s.GetFilePath(URL, false)
	if _, err := os.Stat(filePath); !os.IsNotExist(err) {
		return nil
	}

	s.log.Info("Downloading", zap.Stringer("URL", URL))

	buf := &bytes.Buffer{}
	_, err := asset.Download(buf)
	if err != nil {
		return err
	}

	buf = s.checkFileTypeForRecode(filePath, buf)

	return s.writeFile(filePath, buf)
}

func (s *Scraper) isURLExcluded(URL *url.URL) bool {
	for _, re := range s.excludes {
		if re.MatchString(URL.Path) {
			s.log.Info("Skipping URL", zap.Stringer("URL", URL), zap.Stringer("Excluder", re))
			return true
		}
	}
	return false
}

func (s *Scraper) isExternalFileChecked(URL *url.URL) bool {
	if URL.Host == s.URL.Host {
		return false
	}

	fullURL := URL.String()
	_, ok := s.assetsExternal[fullURL]
	if ok { // was already downloaded or checked
		return true
	}

	s.assetsExternal[fullURL] = true
	s.log.Info("External URL", zap.Stringer("URL", URL))

	return false
}
<|endoftext|>


<commit_before>package scraper

import (
	"io/ioutil"
	"log"

	"github.com/PuerkitoBio/goquery"
)

// Scraper struct used for scraping activity.
type Scraper struct {
	Config  *Config
	Results results
}

// results key is tag, value is list of scraped data
type results map[string][]string

type resultsURL struct {
	url     string
	results results
}

// New creates new Scraper and returns pointer to it, with error (if occurred).
func New(configPath string) (*Scraper, error) {
	c, err := newConfig(configPath)
	if err != nil {
		return nil, err
	}

	s := &Scraper{
		Config:  c,
		Results: make(results),
	}

	return s, nil
}

// Start will start scraping in separate goroutine and save results
// when it is done in file that is defined in Config.
func (s *Scraper) Start() {
	resultsURLCh := s.scrapeURLs(s.Config.URLs)

	for range s.Config.URLs {
		resultsURL := <-resultsURLCh
		mergeResults(s.Results, resultsURL.results)
		log.Println("Received results from", resultsURL.url)
	}

	log.Println("Done scraping.")

	err := s.save()
	if err != nil {
		log.Println(err)
		return
	}

	log.Println("Saved to", s.Config.outputPath())
}

func (s *Scraper) scrapeURLs(urls []string) <-chan *resultsURL {
	resultsURLCh := make(chan *resultsURL)

	for _, url := range urls {
		go func(url string) {
			results := make(results)

			// Once individual URL scraping is done, send results back through channel
			defer func() {
				resultsURLCh <- &resultsURL{url: url, results: results}
			}()

			// Construct document for manipulation
			doc, err := goquery.NewDocument(url)
			if err != nil {
				log.Println(err)
				return
			}

			// Process targets
			for _, target := range s.Config.Targets {
				err := s.processTarget(doc, target, results)
				if err != nil {
					log.Println("Error processing target:", err)
				}
			}
		}(url)
	}

	return resultsURLCh
}

func (s *Scraper) processTarget(doc *goquery.Document, target *target, results results) error {
	var selector string
	// If there was no selector given, whole document will be used
	if target.Selector == "" {
		selector = "html"
	} else {
		selector = target.Selector
	}

	var retErr error
	doc.Find(selector).Each(func(i int, sel *goquery.Selection) {
		var value string

		// Handling different types
		switch {
		// Sets value to inner HTML of the node
		case target.Type == "html":
			html, err := sel.Html()
			if err != nil {
				retErr = err
				return
			}
			value = html
		// Sets value to text of the node
		case target.Type == "text":
			value = sel.Text()
		// Sets value to attribute of the node, for example attr:href for href value
		// If attribute value doesn't exist, target is skipped
		case target.attrv != "":
			if attrv, exists := sel.Attr(target.attrv); exists {
				value = attrv
			} else {
				return
			}
		}

		// Submatch regex
		if target.submatchRe != nil {
			matches := target.submatchRe.FindAllStringSubmatch(value, -1)
			for _, match := range matches {
				results[target.Tag] = append(results[target.Tag], match[0])
			}
			return
		}

		results[target.Tag] = append(results[target.Tag], value)
	})

	return retErr
}

func (s *Scraper) save() error {
	data, err := JSONMarshalUnescaped(s.Results)
	if err != nil {
		return err
	}

	err = ioutil.WriteFile(s.Config.outputPath(), data, 0644)
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>Added extra error information<commit_after>package scraper

import (
	"io/ioutil"
	"log"

	"fmt"
	"github.com/PuerkitoBio/goquery"
)

// Scraper struct used for scraping activity.
type Scraper struct {
	Config  *Config
	Results results
}

// results key is tag, value is list of scraped data
type results map[string][]string

type resultsURL struct {
	url     string
	results results
}

// New creates new Scraper and returns pointer to it, with error (if occurred).
func New(configPath string) (*Scraper, error) {
	c, err := newConfig(configPath)
	if err != nil {
		return nil, err
	}

	s := &Scraper{
		Config:  c,
		Results: make(results),
	}

	return s, nil
}

// Start will start scraping in separate goroutine and save results
// when it is done in file that is defined in Config.
func (s *Scraper) Start() {
	resultsURLCh := s.scrapeURLs(s.Config.URLs)

	for range s.Config.URLs {
		resultsURL := <-resultsURLCh
		mergeResults(s.Results, resultsURL.results)
		log.Println("Received results from", resultsURL.url)
	}

	log.Println("Done scraping.")

	err := s.save()
	if err != nil {
		log.Println(err)
		return
	}

	log.Println("Saved to", s.Config.outputPath())
}

func (s *Scraper) scrapeURLs(urls []string) <-chan *resultsURL {
	resultsURLCh := make(chan *resultsURL)

	for _, url := range urls {
		go func(url string) {
			results := make(results)

			// Once individual URL scraping is done, send results back through channel
			defer func() {
				resultsURLCh <- &resultsURL{url: url, results: results}
			}()

			// Construct document for manipulation
			doc, err := goquery.NewDocument(url)
			if err != nil {
				log.Println(err)
				return
			}

			// Process targets
			for _, target := range s.Config.Targets {
				err := s.processTarget(doc, target, results)
				if err != nil {
					log.Println("Error processing target:", err)
				}
			}
		}(url)
	}

	return resultsURLCh
}

func (s *Scraper) processTarget(doc *goquery.Document, target *target, results results) error {
	var selector string
	// If there was no selector given, whole document will be used
	if target.Selector == "" {
		selector = "html"
	} else {
		selector = target.Selector
	}

	var retErr error
	doc.Find(selector).Each(func(i int, sel *goquery.Selection) {
		var value string

		// Handling different types
		switch {
		// Sets value to inner HTML of the node
		case target.Type == "html":
			html, err := sel.Html()
			if err != nil {
				retErr = err
				return
			}
			value = html
		// Sets value to text of the node
		case target.Type == "text":
			value = sel.Text()
		// Sets value to attribute of the node, for example attr:href for href value
		// If attribute value doesn't exist, target is skipped
		case target.attrv != "":
			if attrv, exists := sel.Attr(target.attrv); exists {
				value = attrv
			} else {
				return
			}
		}

		// Submatch regex
		if target.submatchRe != nil {
			matches := target.submatchRe.FindAllStringSubmatch(value, -1)
			for _, match := range matches {
				results[target.Tag] = append(results[target.Tag], match[0])
			}
			return
		}

		results[target.Tag] = append(results[target.Tag], value)
	})

	return retErr
}

func (s *Scraper) save() error {
	data, err := JSONMarshalUnescaped(s.Results)
	if err != nil {
		return fmt.Errorf("error marshaling to json: %v", err)
	}

	err = ioutil.WriteFile(s.Config.outputPath(), data, 0644)
	if err != nil {
		return fmt.Errorf("error writing file: %v", err)
	}

	return nil
}
<|endoftext|>


<commit_before>package mpb_test

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

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

func TestBarCompleted(t *testing.T) {
	p := New(WithOutput(ioutil.Discard))
	total := 80
	bar := p.AddBar(int64(total))

	var count int
	for !bar.Completed() {
		time.Sleep(10 * time.Millisecond)
		bar.Increment()
		count++
	}

	p.Wait()
	if count != total {
		t.Errorf("got count: %d, expected %d\n", count, total)
	}
}

func TestBarID(t *testing.T) {
	p := New(WithOutput(ioutil.Discard))
	total := 100
	wantID := 11
	bar := p.AddBar(int64(total), BarID(wantID))

	go func(total int) {
		for i := 0; i < total; i++ {
			time.Sleep(50 * time.Millisecond)
			bar.Increment()
		}
	}(total)

	gotID := bar.ID()
	if gotID != wantID {
		t.Errorf("Expected bar id: %d, got %d\n", wantID, gotID)
	}

	bar.Abort(true)
	p.Wait()
}

func TestBarSetRefill(t *testing.T) {
	var buf bytes.Buffer

	width := 100
	p := New(WithOutput(&buf), WithWidth(width))

	total := 100
	till := 30
	refillRune, _ := utf8.DecodeLastRuneInString(DefaultBarStyle)

	bar := p.AddBar(int64(total), TrimSpace())

	bar.SetRefill(int64(till))
	bar.IncrBy(till)

	for i := 0; i < total-till; i++ {
		bar.Increment()
		time.Sleep(10 * time.Millisecond)
	}

	p.Wait()

	wantBar := fmt.Sprintf("[%s%s]",
		strings.Repeat(string(refillRune), till-1),
		strings.Repeat("=", total-till-1),
	)

	got := string(getLastLine(buf.Bytes()))

	if !strings.Contains(got, wantBar) {
		t.Errorf("Want bar: %q, got bar: %q\n", wantBar, got)
	}
}

func TestBarHas100PercentWithOnCompleteDecorator(t *testing.T) {
	var buf bytes.Buffer

	p := New(WithOutput(&buf), WithWidth(80))

	total := 50

	bar := p.AddBar(int64(total),
		AppendDecorators(
			decor.OnComplete(
				decor.Percentage(), "done",
			),
		),
	)

	for i := 0; i < total; i++ {
		bar.Increment()
		time.Sleep(10 * time.Millisecond)
	}

	p.Wait()

	hundred := "100 %"
	if !bytes.Contains(buf.Bytes(), []byte(hundred)) {
		t.Errorf("Bar's buffer does not contain: %q\n", hundred)
	}
}

func TestBarHas100PercentWithBarRemoveOnComplete(t *testing.T) {
	var buf bytes.Buffer

	p := New(WithOutput(&buf), WithWidth(80))

	total := 50

	bar := p.AddBar(int64(total),
		BarRemoveOnComplete(),
		AppendDecorators(decor.Percentage()),
	)

	for i := 0; i < total; i++ {
		bar.Increment()
		time.Sleep(10 * time.Millisecond)
	}

	p.Wait()

	hundred := "100 %"
	if !bytes.Contains(buf.Bytes(), []byte(hundred)) {
		t.Errorf("Bar's buffer does not contain: %q\n", hundred)
	}
}

func TestBarStyle(t *testing.T) {
	var buf bytes.Buffer
	customFormat := "╢▌▌░╟"
	total := 80
	p := New(WithOutput(&buf), WithWidth(total))
	bar := p.AddBar(int64(total), BarStyle(customFormat), TrimSpace())

	for i := 0; i < total; i++ {
		bar.Increment()
		time.Sleep(10 * time.Millisecond)
	}

	p.Wait()

	runes := []rune(customFormat)
	wantBar := fmt.Sprintf("%s%s%s",
		string(runes[0]),
		strings.Repeat(string(runes[1]), total-2),
		string(runes[len(runes)-1]),
	)
	got := string(getLastLine(buf.Bytes()))

	if !strings.Contains(got, wantBar) {
		t.Errorf("Want bar: %q:%d, got bar: %q:%d\n", wantBar, utf8.RuneCountInString(wantBar), got, utf8.RuneCountInString(got))
	}
}

func TestBarPanicBeforeComplete(t *testing.T) {
	var buf bytes.Buffer
	p := New(
		WithDebugOutput(&buf),
		WithOutput(ioutil.Discard),
		WithWidth(80),
	)

	total := 100
	panicMsg := "Upps!!!"
	var pCount uint32
	bar := p.AddBar(int64(total),
		PrependDecorators(panicDecorator(panicMsg,
			func(st decor.Statistics) bool {
				if st.Current >= 42 {
					atomic.AddUint32(&pCount, 1)
					return true
				}
				return false
			},
		)),
	)

	for i := 0; i < total; i++ {
		time.Sleep(10 * time.Millisecond)
		bar.Increment()
	}

	p.Wait()

	if pCount != 1 {
		t.Errorf("Decor called after panic %d times\n", pCount-1)
	}

	barStr := buf.String()
	if !strings.Contains(barStr, panicMsg) {
		t.Errorf("%q doesn't contain %q\n", barStr, panicMsg)
	}
}

func TestBarPanicAfterComplete(t *testing.T) {
	var buf bytes.Buffer
	p := New(
		WithDebugOutput(&buf),
		WithOutput(ioutil.Discard),
		WithWidth(80),
	)

	total := 100
	panicMsg := "Upps!!!"
	var pCount uint32
	bar := p.AddBar(int64(total),
		PrependDecorators(panicDecorator(panicMsg,
			func(st decor.Statistics) bool {
				if st.Completed {
					atomic.AddUint32(&pCount, 1)
					return true
				}
				return false
			},
		)),
	)

	for i := 0; i < total; i++ {
		time.Sleep(10 * time.Millisecond)
		bar.Increment()
	}

	p.Wait()

	if pCount != 1 {
		t.Errorf("Decor called after panic %d times\n", pCount-1)
	}

	barStr := buf.String()
	if !strings.Contains(barStr, panicMsg) {
		t.Errorf("%q doesn't contain %q\n", barStr, panicMsg)
	}
}

func panicDecorator(panicMsg string, cond func(decor.Statistics) bool) decor.Decorator {
	return decor.Any(func(st decor.Statistics) string {
		if cond(st) {
			panic(panicMsg)
		}
		return ""
	})
}
<commit_msg>make sure every test has WithWidth<commit_after>package mpb_test

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

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

func TestBarCompleted(t *testing.T) {
	p := New(WithWidth(80), WithOutput(ioutil.Discard))
	total := 80
	bar := p.AddBar(int64(total))

	var count int
	for !bar.Completed() {
		time.Sleep(10 * time.Millisecond)
		bar.Increment()
		count++
	}

	p.Wait()
	if count != total {
		t.Errorf("got count: %d, expected %d\n", count, total)
	}
}

func TestBarID(t *testing.T) {
	p := New(WithWidth(80), WithOutput(ioutil.Discard))
	total := 100
	wantID := 11
	bar := p.AddBar(int64(total), BarID(wantID))

	go func(total int) {
		for i := 0; i < total; i++ {
			time.Sleep(50 * time.Millisecond)
			bar.Increment()
		}
	}(total)

	gotID := bar.ID()
	if gotID != wantID {
		t.Errorf("Expected bar id: %d, got %d\n", wantID, gotID)
	}

	bar.Abort(true)
	p.Wait()
}

func TestBarSetRefill(t *testing.T) {
	var buf bytes.Buffer

	p := New(WithOutput(&buf), WithWidth(100))

	total := 100
	till := 30
	refillRune, _ := utf8.DecodeLastRuneInString(DefaultBarStyle)

	bar := p.AddBar(int64(total), TrimSpace())

	bar.SetRefill(int64(till))
	bar.IncrBy(till)

	for i := 0; i < total-till; i++ {
		bar.Increment()
		time.Sleep(10 * time.Millisecond)
	}

	p.Wait()

	wantBar := fmt.Sprintf("[%s%s]",
		strings.Repeat(string(refillRune), till-1),
		strings.Repeat("=", total-till-1),
	)

	got := string(getLastLine(buf.Bytes()))

	if !strings.Contains(got, wantBar) {
		t.Errorf("Want bar: %q, got bar: %q\n", wantBar, got)
	}
}

func TestBarHas100PercentWithOnCompleteDecorator(t *testing.T) {
	var buf bytes.Buffer

	p := New(WithWidth(80), WithOutput(&buf))

	total := 50

	bar := p.AddBar(int64(total),
		AppendDecorators(
			decor.OnComplete(
				decor.Percentage(), "done",
			),
		),
	)

	for i := 0; i < total; i++ {
		bar.Increment()
		time.Sleep(10 * time.Millisecond)
	}

	p.Wait()

	hundred := "100 %"
	if !bytes.Contains(buf.Bytes(), []byte(hundred)) {
		t.Errorf("Bar's buffer does not contain: %q\n", hundred)
	}
}

func TestBarHas100PercentWithBarRemoveOnComplete(t *testing.T) {
	var buf bytes.Buffer

	p := New(WithWidth(80), WithOutput(&buf))

	total := 50

	bar := p.AddBar(int64(total),
		BarRemoveOnComplete(),
		AppendDecorators(decor.Percentage()),
	)

	for i := 0; i < total; i++ {
		bar.Increment()
		time.Sleep(10 * time.Millisecond)
	}

	p.Wait()

	hundred := "100 %"
	if !bytes.Contains(buf.Bytes(), []byte(hundred)) {
		t.Errorf("Bar's buffer does not contain: %q\n", hundred)
	}
}

func TestBarStyle(t *testing.T) {
	var buf bytes.Buffer
	customFormat := "╢▌▌░╟"
	total := 80
	p := New(WithWidth(total), WithOutput(&buf))
	bar := p.AddBar(int64(total), BarStyle(customFormat), TrimSpace())

	for i := 0; i < total; i++ {
		bar.Increment()
		time.Sleep(10 * time.Millisecond)
	}

	p.Wait()

	runes := []rune(customFormat)
	wantBar := fmt.Sprintf("%s%s%s",
		string(runes[0]),
		strings.Repeat(string(runes[1]), total-2),
		string(runes[len(runes)-1]),
	)
	got := string(getLastLine(buf.Bytes()))

	if !strings.Contains(got, wantBar) {
		t.Errorf("Want bar: %q:%d, got bar: %q:%d\n", wantBar, utf8.RuneCountInString(wantBar), got, utf8.RuneCountInString(got))
	}
}

func TestBarPanicBeforeComplete(t *testing.T) {
	var buf bytes.Buffer
	p := New(
		WithWidth(80),
		WithDebugOutput(&buf),
		WithOutput(ioutil.Discard),
	)

	total := 100
	panicMsg := "Upps!!!"
	var pCount uint32
	bar := p.AddBar(int64(total),
		PrependDecorators(panicDecorator(panicMsg,
			func(st decor.Statistics) bool {
				if st.Current >= 42 {
					atomic.AddUint32(&pCount, 1)
					return true
				}
				return false
			},
		)),
	)

	for i := 0; i < total; i++ {
		time.Sleep(10 * time.Millisecond)
		bar.Increment()
	}

	p.Wait()

	if pCount != 1 {
		t.Errorf("Decor called after panic %d times\n", pCount-1)
	}

	barStr := buf.String()
	if !strings.Contains(barStr, panicMsg) {
		t.Errorf("%q doesn't contain %q\n", barStr, panicMsg)
	}
}

func TestBarPanicAfterComplete(t *testing.T) {
	var buf bytes.Buffer
	p := New(
		WithWidth(80),
		WithDebugOutput(&buf),
		WithOutput(ioutil.Discard),
	)

	total := 100
	panicMsg := "Upps!!!"
	var pCount uint32
	bar := p.AddBar(int64(total),
		PrependDecorators(panicDecorator(panicMsg,
			func(st decor.Statistics) bool {
				if st.Completed {
					atomic.AddUint32(&pCount, 1)
					return true
				}
				return false
			},
		)),
	)

	for i := 0; i < total; i++ {
		time.Sleep(10 * time.Millisecond)
		bar.Increment()
	}

	p.Wait()

	if pCount != 1 {
		t.Errorf("Decor called after panic %d times\n", pCount-1)
	}

	barStr := buf.String()
	if !strings.Contains(barStr, panicMsg) {
		t.Errorf("%q doesn't contain %q\n", barStr, panicMsg)
	}
}

func panicDecorator(panicMsg string, cond func(decor.Statistics) bool) decor.Decorator {
	return decor.Any(func(st decor.Statistics) string {
		if cond(st) {
			panic(panicMsg)
		}
		return ""
	})
}
<|endoftext|>


<commit_before>package main

import (
	"net/http"
	"io/ioutil"
	"strings"
	"path/filepath"
	"log"
	"html/template"
)

func serveBaseHTML(w http.ResponseWriter, r *http.Request) {
	baseHTML := `
			<!DOCTYPE html>

			    <title>Flowbro</title>

			    Flowbro

				{{range $i, $e := .}}

				    {{.Title}}

{{end}}

footer (fixed height)

	`

	files, err := ioutil.ReadDir("webroot/configs")
	if err != nil {
		log.Fatal(err)
	}

	links := []Link{}
	for _, file := range files {
		config := strings.TrimSuffix(file.Name(), filepath.Ext(file.Name()))
		title := strings.Title(strings.Replace(strings.Replace(config, "-", " ", -1), "_", " ", -1))
		links = append(links, Link{
			Url:   "?config=" + config,
			Title: title,
		})
	}

	templ, err := template.New("base").Parse(baseHTML)
	if err != nil {
		log.Fatal(err)
	}

	err = templ.Execute(w, links)
	if err != nil {
		log.Fatal(err)
	}
}

<commit_msg>Lists only js files within the config directory.<commit_after>package main

import (
	"net/http"
	"io/ioutil"
	"strings"
	"path/filepath"
	"log"
	"html/template"
)

func serveBaseHTML(w http.ResponseWriter, r *http.Request) {
	baseHTML := `
			<!DOCTYPE html>

			    <title>Flowbro</title>

			    Flowbro

				{{range $i, $e := .}}

				    {{.Title}}

{{end}}

footer (fixed height)

	`

	files, err := ioutil.ReadDir("webroot/configs")
	if err != nil {
		log.Fatal(err)
	}

	links := []Link{}
	for _, file := range files {
		if filepath.Ext(file.Name()) == ".js" {
			config := strings.TrimSuffix(file.Name(), filepath.Ext(file.Name()))
			title := strings.Title(strings.Replace(strings.Replace(config, "-", " ", -1), "_", " ", -1))
			links = append(links, Link{
				Url:   "?config=" + config,
				Title: title,
			})
		}
	}

	templ, err := template.New("base").Parse(baseHTML)
	if err != nil {
		log.Fatal(err)
	}

	err = templ.Execute(w, links)
	if err != nil {
		log.Fatal(err)
	}
}

<|endoftext|>