<commit_before>// Package soda provides HTTP GET tools for SODA (Socrata Open Data API) webservices, see http://dev.socrata.com/
package soda

import (
	"context"
	"encoding/csv"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"time"
)

type HTTPRequester interface {
	Do(req *http.Request) (*http.Response, error)
}

// GetRequest is a wrapper/container for SODA requests.
// This is NOT safe for use by multiple goroutines as Format, Filters and Query will be overwritten.
// Create a new GetRequest in each goroutine you use or use an OffsetGetRequest
type GetRequest struct {
	apptoken   string
	endpoint   string //endpoint without format (not .json etc at the end)
	Format     string //json, csv etc
	Filters    SimpleFilters
	Query      SoSQL
	Metadata   metadata
	HTTPClient HTTPRequester //For clients who need a custom HTTP client
	context    context.Context
}

// NewGetRequest creates a new GET request, the endpoint must be specified without the format.
// For example https://data.ct.gov/resource/hma6-9xbg
func NewGetRequest(endpoint, apptoken string) *GetRequest {
	return &GetRequest{
		apptoken: apptoken,
		endpoint: endpoint,
		Filters:  make(SimpleFilters),
		Metadata: newMetadata(endpoint),
	}
}

// Get executes the HTTP GET request
func (r *GetRequest) Get() (*http.Response, error) {
	//If offset is used we must specify an order
	if r.Query.Offset > 0 && len(r.Query.Order) == 0 {
		return nil, errors.New("cannot use an offset without setting the order")
	}
	return get(r, r.URLValues().Encode())
}

// GetEndpoint returns the complete SODA URL with format
func (r *GetRequest) GetEndpoint() string {
	if r.Format == "" {
		r.Format = "json"
	}
	return fmt.Sprintf("%s.%s", r.endpoint, r.Format)
}

// URLValues returns the url.Values for the GetRequest
func (r *GetRequest) URLValues() url.Values {
	uv := make(url.Values)
	for key, val := range r.Filters.URLValues() {
		uv[key] = val
	}
	for key, val := range r.Query.URLValues() {
		uv[key] = val
	}
	return uv
}

// Count gets the total number of records in the dataset
// by executing a SODA request
func (r *GetRequest) Count() (uint, error) {

	oldformat := r.Format
	oldorder := r.Query.Order
	oldselect := r.Query.Select
	defer func() {
		r.Format = oldformat
		r.Query.Order = oldorder
		r.Query.Select = oldselect
	}()

	r.Format = "json"
	r.Query.Select = []string{"count(*)"}
	r.Query.ClearOrder()

	resp, err := r.Get()
	if err != nil {
		return 0, err
	}
	defer resp.Body.Close()

	count := make([]struct {
		Count string
	}, 0)
	err = json.NewDecoder(resp.Body).Decode(&count)
	if err != nil {
		return 0, err
	}
	if len(count) == 0 {
		return 0, errors.New("empty count response")
	}
	icount, err := strconv.Atoi(count[0].Count)
	if err != nil {
		return 0, err
	}
	return uint(icount), nil
}

// Fields returns all the fields present in the dataset (ignores select fields).
// Spaces in fieldnames are replaced by underscores.
func (r *GetRequest) Fields() ([]string, error) {

	oldformat := r.Format
	oldorder := r.Query.Order
	oldlimit := r.Query.Limit
	oldselect := r.Query.Select
	defer func() {
		r.Format = oldformat
		r.Query.Select = oldselect
		r.Query.Order = oldorder
		r.Query.Limit = oldlimit
	}()

	r.Format = "csv"
	r.Query.Select = []string{}
	r.Query.Limit = 0
	r.Query.ClearOrder()

	resp, err := r.Get()
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	csvreader := csv.NewReader(resp.Body)
	record, err := csvreader.Read()
	if err != nil && err != io.EOF {
		return nil, err
	}
	fields := make([]string, len(record))
	for i := range record {
		fields[i] = strings.Replace(record[i], " ", "_", -1)
	}
	return fields, nil
}

// Modified returns when the dataset was last updated
func (r *GetRequest) Modified() (time.Time, error) {

	oldformat := r.Format
	oldorder := r.Query.Order
	oldlimit := r.Query.Limit
	oldselect := r.Query.Select
	defer func() {
		r.Format = oldformat
		r.Query.Select = oldselect
		r.Query.Order = oldorder
		r.Query.Limit = oldlimit
	}()

	r.Format = "json"
	r.Query.Select = []string{}
	r.Query.Limit = 0
	r.Query.ClearOrder()

	resp, err := r.Get()
	if err != nil {
		return time.Time{}, err
	}
	defer resp.Body.Close()

	lms := resp.Header.Get("X-Soda2-Truth-Last-Modified")
	if lms == "" {
		lms = resp.Header.Get("Last-Modified")
	}
	if lms == "" {
		return time.Time{}, errors.New("cannot get last modified date, field not present in HTTP header")
	}

	return time.Parse(time.RFC1123, lms)
}

// WithContext sets the context on the http.Request
func (r *GetRequest) WithContext(ctx context.Context) *GetRequest {
	r.context = ctx
	return r
}

// SimpleFilters is the easiest way to filter columns for equality.
// Add the column to filter on a map key and the filter value as map value.
// If you include multiple filters, the filters will be combined using a boolean AND.
// See http://dev.socrata.com/docs/filtering.html
type SimpleFilters map[string]string

//URLValues returns the url.Values for the SimpleFilters
func (sf SimpleFilters) URLValues() url.Values {
	uv := make(url.Values)
	for key, val := range sf {
		uv.Add(key, val)
	}
	return uv
}

// SoSQL implements the Socrata Query Language and is used to build more complex queries.
// See http://dev.socrata.com/docs/queries.html
type SoSQL struct {
	Select []string //The set of columns to be returned. Default: All columns, equivalent to $select=*
	Where  string   //Filters the rows to be returned. Default: No filter, and returning a max of $limit values
	Order  []struct {
		Column string //Column name
		Desc   bool   //Descending. Default: false = Ascending
	} //Specifies the order of results. Default: Unspecified order, but it will be consistent across paging
	Group  string //Column to group results on, similar to SQL Grouping. Default: No grouping
	Limit  uint   //Maximum number of results to return. Default: 1000 (with a maximum of 50,000)
	Offset uint   //Offset count into the results to start at, used for paging. Default: 0
	Q      string //Performs a full text search for a value. Default: No search

}

// Direction is used to set the sort direction to ascending or descending
type Direction bool

const (
	// DirAsc is used to set ascending sort order
	DirAsc Direction = false

	// DirDesc is used to set descending sort order
	DirDesc Direction = true
)

// AddOrder can be called for each field you want to sort the result on.
// If parameter descending is true, the column will be sorted descending, or ascending if false.
func (sq *SoSQL) AddOrder(column string, dir Direction) {
	sq.Order = append(sq.Order, struct {
		Column string
		Desc   bool
	}{column, bool(dir)})
}

// ClearOrder removes all order fields
func (sq *SoSQL) ClearOrder() {
	sq.Order = []struct {
		Column string
		Desc   bool
	}{}
}

// URLValues returns the url.Values for the SoSQL query
func (sq *SoSQL) URLValues() url.Values {
	uv := make(url.Values)
	if len(sq.Select) > 0 {
		uv.Add("$select", strings.Join(sq.Select, ","))
	}
	if len(sq.Where) > 0 {
		uv.Add("$where", sq.Where)
	}
	if len(sq.Order) > 0 {
		order := make([]string, 0)
		for _, o := range sq.Order {
			if o.Desc {
				order = append(order, o.Column+" DESC")
			} else {
				order = append(order, o.Column+" ASC")
			}
		}
		uv.Add("$order", strings.Join(order, ","))
	}
	if len(sq.Q) > 0 {
		uv.Add("$q", sq.Q)
	}
	if len(sq.Group) > 0 {
		uv.Add("$group", sq.Group)
	}
	if sq.Limit > 0 {
		uv.Add("$limit", fmt.Sprintf("%d", sq.Limit))
	}
	if sq.Offset > 0 {
		uv.Add("$offset", fmt.Sprintf("%d", sq.Offset))
	}
	return uv
}

// OffsetGetRequest is a request getter that gets all the records using the filters and limits from gr and
// is safe to use by multiple goroutines, use Next(number) to get the next number of records.
// A sync.WaitGroup is embedded for easy concurrency.
type OffsetGetRequest struct {
	gr     *GetRequest
	m      sync.Mutex
	offset uint
	count  uint
	sync.WaitGroup
}

// ErrDone is returned by OffsetGetRequest.Next when done
var ErrDone = errors.New("Done")

// Next gets the next number of records
func (o *OffsetGetRequest) Next(number uint) (*http.Response, error) {
	o.m.Lock() //lock to protect offset
	if o.IsDone() {
		o.m.Unlock()
		return nil, ErrDone
	}
	if len(o.gr.Query.Order) == 0 { //If offset is used we must specify an order
		return nil, errors.New("cannot use an offset without setting the order")
	}
	if o.offset+number > o.count {
		number = o.count - o.offset
	}
	o.gr.Query.Offset = o.offset
	o.gr.Query.Limit = number
	rawquery := o.gr.URLValues().Encode()
	o.offset += number
	o.m.Unlock() //unlock before the request is done
	return get(o.gr, rawquery)
}

// Count returns the number of records from memory
func (o *OffsetGetRequest) Count() uint {
	return o.count
}

// IsDone returns if we have gotten all records
func (o *OffsetGetRequest) IsDone() bool {
	return o.offset >= o.count
}

// NewOffsetGetRequest creates a new OffsetGetRequest from gr
// and does a count request to determine the number of records to get
func NewOffsetGetRequest(gr *GetRequest) (*OffsetGetRequest, error) {
	count, err := gr.Count()
	if err != nil {
		return nil, err
	}
	return &OffsetGetRequest{gr: gr, offset: 0, count: count}, nil
}

// get is the function that executes the HTTP request
func get(r *GetRequest, rawquery string) (*http.Response, error) {

	var client HTTPRequester = http.DefaultClient
	if r.HTTPClient != nil {
		client = r.HTTPClient
	}

	req, err := http.NewRequest("GET", r.GetEndpoint(), nil)
	if err != nil {
		return nil, err
	}
	if r.context != nil {
		req = req.WithContext(r.context)
	}
	req.URL.RawQuery = rawquery
	req.Header.Set("X-App-Token", r.apptoken)

	// Execute
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode >= 400 {
		errMsg, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return nil, err
		}
		return nil, fmt.Errorf("SODA error %d:\nURL: GET %s\nResponse: %s", resp.StatusCode, req.URL.String(), errMsg)
	}

	return resp, nil
}
<commit_msg>Revert "Add interface wrapper for http.Client (#3)" (#4)<commit_after>// Package soda provides HTTP GET tools for SODA (Socrata Open Data API) webservices, see http://dev.socrata.com/
package soda

import (
	"encoding/csv"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"time"
)

// GetRequest is a wrapper/container for SODA requests.
// This is NOT safe for use by multiple goroutines as Format, Filters and Query will be overwritten.
// Create a new GetRequest in each goroutine you use or use an OffsetGetRequest
type GetRequest struct {
	apptoken   string
	endpoint   string //endpoint without format (not .json etc at the end)
	Format     string //json, csv etc
	Filters    SimpleFilters
	Query      SoSQL
	Metadata   metadata
	HTTPClient *http.Client //For clients who need a custom HTTP client
}

// NewGetRequest creates a new GET request, the endpoint must be specified without the format.
// For example https://data.ct.gov/resource/hma6-9xbg
func NewGetRequest(endpoint, apptoken string) *GetRequest {
	return &GetRequest{
		apptoken: apptoken,
		endpoint: endpoint,
		Filters:  make(SimpleFilters),
		Metadata: newMetadata(endpoint),
	}
}

// Get executes the HTTP GET request
func (r *GetRequest) Get() (*http.Response, error) {
	//If offset is used we must specify an order
	if r.Query.Offset > 0 && len(r.Query.Order) == 0 {
		return nil, errors.New("cannot use an offset without setting the order")
	}
	return get(r, r.URLValues().Encode())
}

// GetEndpoint returns the complete SODA URL with format
func (r *GetRequest) GetEndpoint() string {
	if r.Format == "" {
		r.Format = "json"
	}
	return fmt.Sprintf("%s.%s", r.endpoint, r.Format)
}

// URLValues returns the url.Values for the GetRequest
func (r *GetRequest) URLValues() url.Values {
	uv := make(url.Values)
	for key, val := range r.Filters.URLValues() {
		uv[key] = val
	}
	for key, val := range r.Query.URLValues() {
		uv[key] = val
	}
	return uv
}

// Count gets the total number of records in the dataset
// by executing a SODA request
func (r *GetRequest) Count() (uint, error) {

	oldformat := r.Format
	oldorder := r.Query.Order
	oldselect := r.Query.Select
	defer func() {
		r.Format = oldformat
		r.Query.Order = oldorder
		r.Query.Select = oldselect
	}()

	r.Format = "json"
	r.Query.Select = []string{"count(*)"}
	r.Query.ClearOrder()

	resp, err := r.Get()
	if err != nil {
		return 0, err
	}
	defer resp.Body.Close()

	count := make([]struct {
		Count string
	}, 0)
	err = json.NewDecoder(resp.Body).Decode(&count)
	if err != nil {
		return 0, err
	}
	if len(count) == 0 {
		return 0, errors.New("empty count response")
	}
	icount, err := strconv.Atoi(count[0].Count)
	if err != nil {
		return 0, err
	}
	return uint(icount), nil
}

// Fields returns all the fields present in the dataset (ignores select fields).
// Spaces in fieldnames are replaced by underscores.
func (r *GetRequest) Fields() ([]string, error) {

	oldformat := r.Format
	oldorder := r.Query.Order
	oldlimit := r.Query.Limit
	oldselect := r.Query.Select
	defer func() {
		r.Format = oldformat
		r.Query.Select = oldselect
		r.Query.Order = oldorder
		r.Query.Limit = oldlimit
	}()

	r.Format = "csv"
	r.Query.Select = []string{}
	r.Query.Limit = 0
	r.Query.ClearOrder()

	resp, err := r.Get()
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	csvreader := csv.NewReader(resp.Body)
	record, err := csvreader.Read()
	if err != nil && err != io.EOF {
		return nil, err
	}
	fields := make([]string, len(record))
	for i := range record {
		fields[i] = strings.Replace(record[i], " ", "_", -1)
	}
	return fields, nil
}

// Modified returns when the dataset was last updated
func (r *GetRequest) Modified() (time.Time, error) {

	oldformat := r.Format
	oldorder := r.Query.Order
	oldlimit := r.Query.Limit
	oldselect := r.Query.Select
	defer func() {
		r.Format = oldformat
		r.Query.Select = oldselect
		r.Query.Order = oldorder
		r.Query.Limit = oldlimit
	}()

	r.Format = "json"
	r.Query.Select = []string{}
	r.Query.Limit = 0
	r.Query.ClearOrder()

	resp, err := r.Get()
	if err != nil {
		return time.Time{}, err
	}
	defer resp.Body.Close()

	lms := resp.Header.Get("X-Soda2-Truth-Last-Modified")
	if lms == "" {
		lms = resp.Header.Get("Last-Modified")
	}
	if lms == "" {
		return time.Time{}, errors.New("cannot get last modified date, field not present in HTTP header")
	}

	return time.Parse(time.RFC1123, lms)
}

// SimpleFilters is the easiest way to filter columns for equality.
// Add the column to filter on a map key and the filter value as map value.
// If you include multiple filters, the filters will be combined using a boolean AND.
// See http://dev.socrata.com/docs/filtering.html
type SimpleFilters map[string]string

//URLValues returns the url.Values for the SimpleFilters
func (sf SimpleFilters) URLValues() url.Values {
	uv := make(url.Values)
	for key, val := range sf {
		uv.Add(key, val)
	}
	return uv
}

// SoSQL implements the Socrata Query Language and is used to build more complex queries.
// See http://dev.socrata.com/docs/queries.html
type SoSQL struct {
	Select []string //The set of columns to be returned. Default: All columns, equivalent to $select=*
	Where  string   //Filters the rows to be returned. Default: No filter, and returning a max of $limit values
	Order  []struct {
		Column string //Column name
		Desc   bool   //Descending. Default: false = Ascending
	} //Specifies the order of results. Default: Unspecified order, but it will be consistent across paging
	Group  string //Column to group results on, similar to SQL Grouping. Default: No grouping
	Limit  uint   //Maximum number of results to return. Default: 1000 (with a maximum of 50,000)
	Offset uint   //Offset count into the results to start at, used for paging. Default: 0
	Q      string //Performs a full text search for a value. Default: No search

}

// Direction is used to set the sort direction to ascending or descending
type Direction bool

const (
	// DirAsc is used to set ascending sort order
	DirAsc Direction = false

	// DirDesc is used to set descending sort order
	DirDesc Direction = true
)

// AddOrder can be called for each field you want to sort the result on.
// If parameter descending is true, the column will be sorted descending, or ascending if false.
func (sq *SoSQL) AddOrder(column string, dir Direction) {
	sq.Order = append(sq.Order, struct {
		Column string
		Desc   bool
	}{column, bool(dir)})
}

// ClearOrder removes all order fields
func (sq *SoSQL) ClearOrder() {
	sq.Order = []struct {
		Column string
		Desc   bool
	}{}
}

// URLValues returns the url.Values for the SoSQL query
func (sq *SoSQL) URLValues() url.Values {
	uv := make(url.Values)
	if len(sq.Select) > 0 {
		uv.Add("$select", strings.Join(sq.Select, ","))
	}
	if len(sq.Where) > 0 {
		uv.Add("$where", sq.Where)
	}
	if len(sq.Order) > 0 {
		order := make([]string, 0)
		for _, o := range sq.Order {
			if o.Desc {
				order = append(order, o.Column+" DESC")
			} else {
				order = append(order, o.Column+" ASC")
			}
		}
		uv.Add("$order", strings.Join(order, ","))
	}
	if len(sq.Q) > 0 {
		uv.Add("$q", sq.Q)
	}
	if len(sq.Group) > 0 {
		uv.Add("$group", sq.Group)
	}
	if sq.Limit > 0 {
		uv.Add("$limit", fmt.Sprintf("%d", sq.Limit))
	}
	if sq.Offset > 0 {
		uv.Add("$offset", fmt.Sprintf("%d", sq.Offset))
	}
	return uv
}

// OffsetGetRequest is a request getter that gets all the records using the filters and limits from gr and
// is safe to use by multiple goroutines, use Next(number) to get the next number of records.
// A sync.WaitGroup is embedded for easy concurrency.
type OffsetGetRequest struct {
	gr     *GetRequest
	m      sync.Mutex
	offset uint
	count  uint
	sync.WaitGroup
}

// ErrDone is returned by OffsetGetRequest.Next when done
var ErrDone = errors.New("Done")

// Next gets the next number of records
func (o *OffsetGetRequest) Next(number uint) (*http.Response, error) {
	o.m.Lock() //lock to protect offset
	if o.IsDone() {
		o.m.Unlock()
		return nil, ErrDone
	}
	if len(o.gr.Query.Order) == 0 { //If offset is used we must specify an order
		return nil, errors.New("cannot use an offset without setting the order")
	}
	if o.offset+number > o.count {
		number = o.count - o.offset
	}
	o.gr.Query.Offset = o.offset
	o.gr.Query.Limit = number
	rawquery := o.gr.URLValues().Encode()
	o.offset += number
	o.m.Unlock() //unlock before the request is done
	return get(o.gr, rawquery)
}

// Count returns the number of records from memory
func (o *OffsetGetRequest) Count() uint {
	return o.count
}

// IsDone returns if we have gotten all records
func (o *OffsetGetRequest) IsDone() bool {
	return o.offset >= o.count
}

// NewOffsetGetRequest creates a new OffsetGetRequest from gr
// and does a count request to determine the number of records to get
func NewOffsetGetRequest(gr *GetRequest) (*OffsetGetRequest, error) {
	count, err := gr.Count()
	if err != nil {
		return nil, err
	}
	return &OffsetGetRequest{gr: gr, offset: 0, count: count}, nil
}

// get is the function that executes the HTTP request
func get(r *GetRequest, rawquery string) (*http.Response, error) {

	client := http.DefaultClient
	if r.HTTPClient != nil {
		client = r.HTTPClient
	}

	req, err := http.NewRequest("GET", r.GetEndpoint(), nil)
	if err != nil {
		return nil, err
	}
	req.URL.RawQuery = rawquery
	req.Header.Set("X-App-Token", r.apptoken)

	// Execute
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode >= 400 {
		errMsg, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return nil, err
		}
		return nil, fmt.Errorf("SODA error %d:\nURL: GET %s\nResponse: %s", resp.StatusCode, req.URL.String(), errMsg)
	}

	return resp, nil
}
<|endoftext|>


<commit_before>package opentracing

import (
	"time"

	"github.com/opentracing/opentracing-go/log"
)

// SpanContext represents Span state that must propagate to descendant Spans and across process
// boundaries (e.g., a <trace_id, span_id, sampled> tuple).
type SpanContext interface {
	// ForeachBaggageItem grants access to all baggage items stored in the
	// SpanContext.
	// The handler function will be called for each baggage key/value pair.
	// The ordering of items is not guaranteed.
	//
	// The bool return value indicates if the handler wants to continue iterating
	// through the rest of the baggage items; for example if the handler is trying to
	// find some baggage item by pattern matching the name, it can return false
	// as soon as the item is found to stop further iterations.
	ForeachBaggageItem(handler func(k, v string) bool)
}

// Span represents an active, un-finished span in the OpenTracing system.
//
// Spans are created by the Tracer interface.
type Span interface {
	// Sets the end timestamp and finalizes Span state.
	//
	// With the exception of calls to Context() (which are always allowed),
	// Finish() must be the last call made to any span instance, and to do
	// otherwise leads to undefined behavior.
	Finish()
	// FinishWithOptions is like Finish() but with explicit control over
	// timestamps and log data.
	FinishWithOptions(opts FinishOptions)

	// Context() yields the SpanContext for this Span. Note that the return
	// value of Context() is still valid after a call to Span.Finish(), as is
	// a call to Span.Context() after a call to Span.Finish().
	Context() SpanContext

	// Sets or changes the operation name.
	SetOperationName(operationName string) Span

	// Adds a tag to the span.
	//
	// If there is a pre-existing tag set for `key`, it is overwritten.
	//
	// Tag values can be numeric types, strings, or bools. The behavior of
	// other tag value types is undefined at the OpenTracing level. If a
	// tracing system does not know how to handle a particular value type, it
	// may ignore the tag, but shall not panic.
	SetTag(key string, value interface{}) Span

	// LogFields is an efficient and type-checked way to record key:value
	// logging data about a Span, though the programming interface is a little
	// more verbose than LogKV(). Here's an example:
	//
	//    span.LogFields(
	//        log.String("request_path", request.Path()),
	//        log.Uint32("request_size", request.Size()))
	//
	// Also see Span.FinishWithOptions() and FinishOptions.BulkLogData.
	LogFields(fields ...log.Field)

	// LogKV is a concise, readable way to record key:value logging data about
	// a Span, though unfortunately this also makes it less efficient and less
	// type-safe than LogFields(). Here's an example:
	//
	//    span.LogKV(
	//        "request_path", request.Path(),
	//        "request_size", request.Size())
	//
	// For LogKV (as opposed to LogFields()), the parameters must appear as
	// key-value pairs, like
	//
	//    span.LogKV(key1, val1, key2, val2, key3, val3, ...)
	//
	// The keys must all be strings. The values may be strings, numeric types,
	// bools, Go error instances, or arbitrary structs.
	LogKV(alternatingKeyValues ...interface{})

	// SetBaggageItem sets a key:value pair on this Span and its SpanContext
	// that also propagates to descendants of this Span.
	//
	// SetBaggageItem() enables powerful functionality given a full-stack
	// opentracing integration (e.g., arbitrary application data from a mobile
	// app can make it, transparently, all the way into the depths of a storage
	// system), and with it some powerful costs: use this feature with care.
	//
	// IMPORTANT NOTE #1: SetBaggageItem() will only propagate baggage items to
	// *future* causal descendants of the associated Span.
	//
	// IMPORTANT NOTE #2: Use this thoughtfully and with care. Every key and
	// value is copied into every local *and remote* child of the associated
	// Span, and that can add up to a lot of network and cpu overhead.
	//
	// Returns a reference to this Span for chaining.
	SetBaggageItem(restrictedKey, value string) Span

	// Gets the value for a baggage item given its key. Returns the empty string
	// if the value isn't found in this Span.
	BaggageItem(restrictedKey string) string

	// Provides access to the Tracer that created this Span.
	Tracer() Tracer

	// DEPRECATED
	LogEvent(event string)
	// DEPRECATED
	LogEventWithPayload(event string, payload interface{})
	// DEPRECATED
	Log(data LogData)
}

// LogRecord is data associated with a single Span log. Every LogRecord
// instance must specify at least one Field.
type LogRecord struct {
	Timestamp time.Time
	Fields    []log.Field
}

// FinishOptions allows Span.FinishWithOptions callers to override the finish
// timestamp and provide log data via a bulk interface.
type FinishOptions struct {
	// FinishTime overrides the Span's finish time, or implicitly becomes
	// time.Now() if FinishTime.IsZero().
	//
	// FinishTime must resolve to a timestamp that's >= the Span's StartTime
	// (per StartSpanOptions).
	FinishTime time.Time

	// LogRecords allows the caller to specify the contents of many LogFields()
	// calls with a single slice. May be nil.
	//
	// None of the LogRecord.Timestamp values may be .IsZero() (i.e., they must
	// be set explicitly). Also, they must be >= the Span's start timestamp and
	// <= the FinishTime (or time.Now() if FinishTime.IsZero()). Otherwise the
	// behavior of FinishWithOptions() is undefined.
	//
	// If specified, the caller hands off ownership of LogRecords at
	// FinishWithOptions() invocation time.
	//
	// If specified, the (deprecated) BulkLogData must be nil or empty.
	LogRecords []LogRecord

	// BulkLogData is DEPRECATED.
	BulkLogData []LogData
}

// LogData is DEPRECATED
type LogData struct {
	Timestamp time.Time
	Event     string
	Payload   interface{}
}
<commit_msg>Add a LogData->LogRecord converter<commit_after>package opentracing

import (
	"time"

	"github.com/opentracing/opentracing-go/log"
)

// SpanContext represents Span state that must propagate to descendant Spans and across process
// boundaries (e.g., a <trace_id, span_id, sampled> tuple).
type SpanContext interface {
	// ForeachBaggageItem grants access to all baggage items stored in the
	// SpanContext.
	// The handler function will be called for each baggage key/value pair.
	// The ordering of items is not guaranteed.
	//
	// The bool return value indicates if the handler wants to continue iterating
	// through the rest of the baggage items; for example if the handler is trying to
	// find some baggage item by pattern matching the name, it can return false
	// as soon as the item is found to stop further iterations.
	ForeachBaggageItem(handler func(k, v string) bool)
}

// Span represents an active, un-finished span in the OpenTracing system.
//
// Spans are created by the Tracer interface.
type Span interface {
	// Sets the end timestamp and finalizes Span state.
	//
	// With the exception of calls to Context() (which are always allowed),
	// Finish() must be the last call made to any span instance, and to do
	// otherwise leads to undefined behavior.
	Finish()
	// FinishWithOptions is like Finish() but with explicit control over
	// timestamps and log data.
	FinishWithOptions(opts FinishOptions)

	// Context() yields the SpanContext for this Span. Note that the return
	// value of Context() is still valid after a call to Span.Finish(), as is
	// a call to Span.Context() after a call to Span.Finish().
	Context() SpanContext

	// Sets or changes the operation name.
	SetOperationName(operationName string) Span

	// Adds a tag to the span.
	//
	// If there is a pre-existing tag set for `key`, it is overwritten.
	//
	// Tag values can be numeric types, strings, or bools. The behavior of
	// other tag value types is undefined at the OpenTracing level. If a
	// tracing system does not know how to handle a particular value type, it
	// may ignore the tag, but shall not panic.
	SetTag(key string, value interface{}) Span

	// LogFields is an efficient and type-checked way to record key:value
	// logging data about a Span, though the programming interface is a little
	// more verbose than LogKV(). Here's an example:
	//
	//    span.LogFields(
	//        log.String("request_path", request.Path()),
	//        log.Uint32("request_size", request.Size()))
	//
	// Also see Span.FinishWithOptions() and FinishOptions.BulkLogData.
	LogFields(fields ...log.Field)

	// LogKV is a concise, readable way to record key:value logging data about
	// a Span, though unfortunately this also makes it less efficient and less
	// type-safe than LogFields(). Here's an example:
	//
	//    span.LogKV(
	//        "request_path", request.Path(),
	//        "request_size", request.Size())
	//
	// For LogKV (as opposed to LogFields()), the parameters must appear as
	// key-value pairs, like
	//
	//    span.LogKV(key1, val1, key2, val2, key3, val3, ...)
	//
	// The keys must all be strings. The values may be strings, numeric types,
	// bools, Go error instances, or arbitrary structs.
	//
	// (Note to implementors: consider the log.InterleavedKVToFields() helper)
	LogKV(alternatingKeyValues ...interface{})

	// SetBaggageItem sets a key:value pair on this Span and its SpanContext
	// that also propagates to descendants of this Span.
	//
	// SetBaggageItem() enables powerful functionality given a full-stack
	// opentracing integration (e.g., arbitrary application data from a mobile
	// app can make it, transparently, all the way into the depths of a storage
	// system), and with it some powerful costs: use this feature with care.
	//
	// IMPORTANT NOTE #1: SetBaggageItem() will only propagate baggage items to
	// *future* causal descendants of the associated Span.
	//
	// IMPORTANT NOTE #2: Use this thoughtfully and with care. Every key and
	// value is copied into every local *and remote* child of the associated
	// Span, and that can add up to a lot of network and cpu overhead.
	//
	// Returns a reference to this Span for chaining.
	SetBaggageItem(restrictedKey, value string) Span

	// Gets the value for a baggage item given its key. Returns the empty string
	// if the value isn't found in this Span.
	BaggageItem(restrictedKey string) string

	// Provides access to the Tracer that created this Span.
	Tracer() Tracer

	// DEPRECATED
	LogEvent(event string)
	// DEPRECATED
	LogEventWithPayload(event string, payload interface{})
	// DEPRECATED
	Log(data LogData)
}

// LogRecord is data associated with a single Span log. Every LogRecord
// instance must specify at least one Field.
type LogRecord struct {
	Timestamp time.Time
	Fields    []log.Field
}

// FinishOptions allows Span.FinishWithOptions callers to override the finish
// timestamp and provide log data via a bulk interface.
type FinishOptions struct {
	// FinishTime overrides the Span's finish time, or implicitly becomes
	// time.Now() if FinishTime.IsZero().
	//
	// FinishTime must resolve to a timestamp that's >= the Span's StartTime
	// (per StartSpanOptions).
	FinishTime time.Time

	// LogRecords allows the caller to specify the contents of many LogFields()
	// calls with a single slice. May be nil.
	//
	// None of the LogRecord.Timestamp values may be .IsZero() (i.e., they must
	// be set explicitly). Also, they must be >= the Span's start timestamp and
	// <= the FinishTime (or time.Now() if FinishTime.IsZero()). Otherwise the
	// behavior of FinishWithOptions() is undefined.
	//
	// If specified, the caller hands off ownership of LogRecords at
	// FinishWithOptions() invocation time.
	//
	// If specified, the (deprecated) BulkLogData must be nil or empty.
	LogRecords []LogRecord

	// BulkLogData is DEPRECATED.
	BulkLogData []LogData
}

// LogData is DEPRECATED
type LogData struct {
	Timestamp time.Time
	Event     string
	Payload   interface{}
}

// ToLogRecord converts a deprecated LogData to a non-deprecated LogRecord
func (ld *LogData) ToLogRecord() LogRecord {
	return LogRecord{
		Timestamp: ld.Timestamp,
		Fields: []log.Field{
			log.String("event", ld.Event),
			log.Object("payload", ld.Payload),
		},
	}
}
<|endoftext|>


<commit_before>package ssgo

import (
	"bytes"
	"fmt"
	"net"
	"strconv"
)

type Client struct {
	sock     *net.TCPConn
	recv_buf bytes.Buffer
	pool     *ConPool
	netErr   error
}

func Connect(ip string, port int) (*Client, error) {
	addr, err := net.ResolveTCPAddr("tcp", fmt.Sprintf("%s:%d", ip, port))
	if err != nil {
		return nil, err
	}
	sock, err := net.DialTCP("tcp", nil, addr)
	if err != nil {
		return nil, err
	}
	var c Client
	c.sock = sock
	return &c, nil
}

func (c *Client) Do(args ...interface{}) ([]string, error) {
	err := c.send(args)
	if err != nil {
		return nil, err
	}
	resp, err := c.recv()
	return resp, err
}

func (c *Client) Cmd(args ...interface{}) *Reply {

	r := &Reply{
		State: ReplyError,
		Data:  []string{},
	}

	if err := c.send(args); err != nil {
		return r
	}

	resp, err := c.recv()
	if err != nil || len(resp) < 1 {
		return r
	}

	switch resp[0] {
	case ReplyOK, ReplyNotFound, ReplyError, ReplyFail, ReplyClientError:
		r.State = resp[0]
	}

	if r.State == ReplyOK {
		for k, v := range resp {

			if k == 0 {
				continue
			}

			r.Data = append(r.Data, v)
		}
	}

	return r
}

func (c *Client) Set(key string, val string) (interface{}, error) {
	resp, err := c.Do("set", key, val)
	if err != nil {
		return nil, err
	}
	if len(resp) == 2 && resp[0] == "ok" {
		return true, nil
	}
	return nil, fmt.Errorf("bad response")
}

// TODO: Will somebody write addition semantic methods?
func (c *Client) Get(key string) (interface{}, error) {
	resp, err := c.Do("get", key)
	if err != nil {
		return nil, err
	}
	if len(resp) == 2 && resp[0] == "ok" {
		return resp[1], nil
	}
	if resp[0] == "not_found" {
		return nil, nil
	}
	return nil, fmt.Errorf("bad response")
}

func (c *Client) Del(key string) (interface{}, error) {
	resp, err := c.Do("del", key)
	if err != nil {
		return nil, err
	}

	//response looks like this: [ok 1]
	if len(resp) > 0 && resp[0] == "ok" {
		return true, nil
	}
	return nil, fmt.Errorf("bad response:resp:%v:", resp)
}

func (c *Client) Send(args ...interface{}) error {
	return c.send(args)
}

func (c *Client) send(args []interface{}) error {
	var buf bytes.Buffer
	for _, arg := range args {
		var s string
		switch arg := arg.(type) {
		case string:
			s = arg
		case []byte:
			s = string(arg)
		case []string:
			for _, s := range arg {
				buf.WriteString(fmt.Sprintf("%d", len(s)))
				buf.WriteByte('\n')
				buf.WriteString(s)
				buf.WriteByte('\n')
			}
			continue
		case int:
			s = fmt.Sprintf("%d", arg)
		case int64:
			s = fmt.Sprintf("%d", arg)
		case float64:
			s = fmt.Sprintf("%f", arg)
		case bool:
			if arg {
				s = "1"
			} else {
				s = "0"
			}
		case nil:
			s = ""
		default:
			return fmt.Errorf("bad arguments")
		}
		buf.WriteString(fmt.Sprintf("%d", len(s)))
		buf.WriteByte('\n')
		buf.WriteString(s)
		buf.WriteByte('\n')
	}
	buf.WriteByte('\n')
	_, err := c.sock.Write(buf.Bytes())
	if err != nil {
		c.netErr = err
	}
	return err
}

func (c *Client) Recv() ([]string, error) {
	return c.recv()
}

func (c *Client) recv() ([]string, error) {
	var tmp [1]byte
	for {
		resp := c.parse()
		if resp == nil || len(resp) > 0 {
			return resp, nil
		}
		n, err := c.sock.Read(tmp[0:])
		if err != nil {
			c.netErr = err
			return nil, err
		}
		c.recv_buf.Write(tmp[0:n])
	}
}

func (c *Client) parse() []string {
	resp := []string{}
	buf := c.recv_buf.Bytes()
	var idx, offset int
	idx = 0
	offset = 0

	for {
		idx = bytes.IndexByte(buf[offset:], '\n')
		if idx == -1 {
			break
		}
		p := buf[offset : offset+idx]
		offset += idx + 1
		//fmt.Printf("> [%s]\n", p);
		if len(p) == 0 || (len(p) == 1 && p[0] == '\r') {
			if len(resp) == 0 {
				continue
			} else {
				c.recv_buf.Next(offset)
				return resp
			}
		}

		size, err := strconv.Atoi(string(p))
		if err != nil || size < 0 {
			return nil
		}
		if offset+size >= c.recv_buf.Len() {
			break
		}

		v := buf[offset : offset+size]
		resp = append(resp, string(v))
		offset += size + 1
	}

	//fmt.Printf("buf.size: %d packet not ready...\n", len(buf))
	return []string{}
}

// Close The Client Connection
func (c *Client) Close() error {
	return c.sock.Close()
}

func (c *Client) Release() error {
	if c.netErr != nil {
		// if client have net error, try to close it
		return c.Close()
	} else if c.pool != nil {
		c.pool.push(c)
		return nil
	}
	return c.Close()
}
<commit_msg> Client sent func marshal unkonw arg type by default<commit_after>package ssgo

import (
	"bytes"
	"fmt"
	"net"
	"strconv"
	"encoding/json"
)

type Client struct {
	sock     *net.TCPConn
	recv_buf bytes.Buffer
	pool     *ConPool
	netErr   error
}

func Connect(ip string, port int) (*Client, error) {
	addr, err := net.ResolveTCPAddr("tcp", fmt.Sprintf("%s:%d", ip, port))
	if err != nil {
		return nil, err
	}
	sock, err := net.DialTCP("tcp", nil, addr)
	if err != nil {
		return nil, err
	}
	var c Client
	c.sock = sock
	return &c, nil
}

func (c *Client) Do(args ...interface{}) ([]string, error) {
	err := c.send(args)
	if err != nil {
		return nil, err
	}
	resp, err := c.recv()
	return resp, err
}

func (c *Client) Cmd(args ...interface{}) *Reply {

	r := &Reply{
		State: ReplyError,
		Data:  []string{},
	}

	if err := c.send(args); err != nil {
		return r
	}

	resp, err := c.recv()
	if err != nil || len(resp) < 1 {
		return r
	}

	switch resp[0] {
	case ReplyOK, ReplyNotFound, ReplyError, ReplyFail, ReplyClientError:
		r.State = resp[0]
	}

	if r.State == ReplyOK {
		for k, v := range resp {

			if k == 0 {
				continue
			}

			r.Data = append(r.Data, v)
		}
	}

	return r
}

func (c *Client) Set(key string, val string) (interface{}, error) {
	resp, err := c.Do("set", key, val)
	if err != nil {
		return nil, err
	}
	if len(resp) == 2 && resp[0] == "ok" {
		return true, nil
	}
	return nil, fmt.Errorf("bad response")
}

// TODO: Will somebody write addition semantic methods?
func (c *Client) Get(key string) (interface{}, error) {
	resp, err := c.Do("get", key)
	if err != nil {
		return nil, err
	}
	if len(resp) == 2 && resp[0] == "ok" {
		return resp[1], nil
	}
	if resp[0] == "not_found" {
		return nil, nil
	}
	return nil, fmt.Errorf("bad response")
}

func (c *Client) Del(key string) (interface{}, error) {
	resp, err := c.Do("del", key)
	if err != nil {
		return nil, err
	}

	//response looks like this: [ok 1]
	if len(resp) > 0 && resp[0] == "ok" {
		return true, nil
	}
	return nil, fmt.Errorf("bad response:resp:%v:", resp)
}

func (c *Client) Send(args ...interface{}) error {
	return c.send(args)
}

func (c *Client) send(args []interface{}) error {
	var buf bytes.Buffer
	for _, arg := range args {
		var s string
		switch arg := arg.(type) {
		case string:
			s = arg
		case []byte:
			s = string(arg)
		case []string:
			for _, s := range arg {
				buf.WriteString(fmt.Sprintf("%d", len(s)))
				buf.WriteByte('\n')
				buf.WriteString(s)
				buf.WriteByte('\n')
			}
			continue
		case int:
			s = fmt.Sprintf("%d", arg)
		case int64:
			s = fmt.Sprintf("%d", arg)
		case float64:
			s = fmt.Sprintf("%f", arg)
		case bool:
			if arg {
				s = "1"
			} else {
				s = "0"
			}
		case nil:
			s = ""
		default:
			buf, e := json.Marshal(arg)
			if e !=nil {
				return fmt.Errorf("bad arguments")
			}
			s = string(buf)
		}
		buf.WriteString(fmt.Sprintf("%d", len(s)))
		buf.WriteByte('\n')
		buf.WriteString(s)
		buf.WriteByte('\n')
	}
	buf.WriteByte('\n')
	_, err := c.sock.Write(buf.Bytes())
	if err != nil {
		c.netErr = err
	}
	return err
}

func (c *Client) Recv() ([]string, error) {
	return c.recv()
}

func (c *Client) recv() ([]string, error) {
	var tmp [1]byte
	for {
		resp := c.parse()
		if resp == nil || len(resp) > 0 {
			return resp, nil
		}
		n, err := c.sock.Read(tmp[0:])
		if err != nil {
			c.netErr = err
			return nil, err
		}
		c.recv_buf.Write(tmp[0:n])
	}
}

func (c *Client) parse() []string {
	resp := []string{}
	buf := c.recv_buf.Bytes()
	var idx, offset int
	idx = 0
	offset = 0

	for {
		idx = bytes.IndexByte(buf[offset:], '\n')
		if idx == -1 {
			break
		}
		p := buf[offset : offset+idx]
		offset += idx + 1
		//fmt.Printf("> [%s]\n", p);
		if len(p) == 0 || (len(p) == 1 && p[0] == '\r') {
			if len(resp) == 0 {
				continue
			} else {
				c.recv_buf.Next(offset)
				return resp
			}
		}

		size, err := strconv.Atoi(string(p))
		if err != nil || size < 0 {
			return nil
		}
		if offset+size >= c.recv_buf.Len() {
			break
		}

		v := buf[offset : offset+size]
		resp = append(resp, string(v))
		offset += size + 1
	}

	//fmt.Printf("buf.size: %d packet not ready...\n", len(buf))
	return []string{}
}

// Close The Client Connection
func (c *Client) Close() error {
	return c.sock.Close()
}

func (c *Client) Release() error {
	if c.netErr != nil {
		// if client have net error, try to close it
		return c.Close()
	} else if c.pool != nil {
		c.pool.push(c)
		return nil
	}
	return c.Close()
}
<|endoftext|>


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

package summstat

import (
	"math"
	"sort"
)

// The type of samples we track statistics for
type Sample float64

// Implements sort.Interface so that we can sort samples for percentiles
type sampleSlice []Sample

func (s sampleSlice) Len() int {
	return len(s)
}

func (s sampleSlice) Less(i, j int) bool {
	return s[i] < s[j]
}

func (s sampleSlice) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}

// A Stats represents descriptive statistics about Samples which are being
// added incrementally.
type Stats struct {
	count     int
	sum       Sample
	sum2      Sample
	max       Sample
	min       Sample
	samples   []Sample
	sorted    bool
	bins      []Sample
	binCounts []int
}

// NewStats returns a new Stats
func NewStats() *Stats {
	return &Stats{
		max: -math.MaxFloat64,
		min: math.MaxFloat64,
	}
}

// AddSample adds a sample value and updates the statistics.
func (s *Stats) AddSample(val Sample) {
	s.count++
	s.sum += val
	s.sum2 += val * val
	if val > s.max {
		s.max = val
	}
	if val < s.min {
		s.min = val
	}
	if len(s.bins) > 0 {
		// TODO: use faster lookup method for large bin counts
		for bin, binVal := range s.bins {
			if val <= binVal {
				s.binCounts[bin]++
				break
			}
		}
	} else {
		s.samples = append(s.samples, val)
		s.sorted = false
	}
}

// Count returns the number of samples added.
func (s Stats) Count() int {
	return s.count
}

// Min returns minimal sample value added.
func (s Stats) Min() Sample {
	if s.min > s.max {
		return 0
	}
	return s.min
}

// Max returns the maximal sample value added.
func (s Stats) Max() Sample {
	if s.min > s.max {
		return 0
	}
	return s.max
}

func (s *Stats) sortSamples() {
	if !s.sorted {
		sort.Sort(sampleSlice(s.samples))
		s.sorted = true
	}
}

// Percentile returns the sample value at the given percentile.
//
// It may not be called after CreateBins, which discards the samples from
// which the percentile is calculated.
func (s Stats) Percentile(pct float64) Sample {
	if len(s.bins) > 0 {
		panic("cannot call Percentile() after CreateBins()")
	}
	if len(s.samples) == 0 {
		return 0
	}
	if pct < 0 {
		panic("pct too small")
	}
	if pct > 1 {
		panic("pct too large")
	}
	s.sortSamples()
	// scale pct into int in [0, len-1]
	// Adding 0.5 turns the implicit floor operation of int() into a rounding operation
	i := int(float64(len(s.samples)-1)*pct + 0.5)
	return s.samples[i]
}

// Median returns the median of the samples.
//
// It may not be called after CreateBins, which discards the samples from
// which the percentile is calculated.
func (s Stats) Median() float64 {
	if len(s.bins) > 0 {
		panic("cannot call Percentile() after CreateBins()")
	}
	l := len(s.samples)
	if l == 0 {
		return 0
	}
	s.sortSamples()
	half, rem := l/2, l%2
	if rem == 0 {
		return (float64(s.samples[half]) + float64(s.samples[half-1])) / 2
	}
	return float64(s.samples[half])
}

// Mean returns the mean of the samples.
func (s Stats) Mean() float64 {
	return float64(s.sum) / float64(s.count)
}

// Stddev returns the standard deviation of the samples.
func (s Stats) Stddev() float64 {
	m := s.Mean()
	return math.Sqrt(float64(s.sum2)/float64(s.count) - m*m)
}

// Spread returns the difference of the maximal and minimal sample values.
func (s Stats) Spread() Sample {
	if s.min > s.max {
		return 0
	}
	return s.max - s.min
}

// CreateBins divides the sample space into nbins bins for tracking counts.
//
// As samples are added, the count for the corresponding bin will be
// incremented and the sample value will not be stored.
//
// This saves memory at the expense of granularity. Percentile() and Median()
// cannot be called after CreateBins() because they are no longer meaningful.
// Instead, use Bin(i) to inspect the distribution of data by bin. Any existing
// stored samples are discarded.
//
// The bins created will be:
//   (-Inf,low], (low, s/nmid+low], (s/nmid+low, 2s/nmid], ..., (high,+Inf)
//   where:
//     s = high - low
//     nmid = nbins-2
//
// Thus, the space (high-low) is divided into nbins-2 equally sized pieces
// and the remaining two bins extend from -math.MaxFloat64 to low and high to
// math.MaxFloat64.
//
// Low must be strictly less than high, so nbins must be at least 3.
func (s *Stats) CreateBins(nbins int, low, high Sample) {
	if high <= low {
		panic("high must be greater than low")
	}
	if nbins < 3 {
		panic("Not enough bins")
	}
	spread := high - low
	s.bins = make([]Sample, nbins)
	s.binCounts = make([]int, nbins)
	for i := 0; i < nbins-1; i++ {
		s.bins[i] = Sample(i)*spread/Sample(nbins-2) + low
	}
	s.bins[nbins-1] = math.MaxFloat64
	// save memory: stop storing samples now that we track by bins
	s.samples = []Sample{}
}

// CreateBinsDiscard is shorthand for calling CreateBins(nbins, ...) with low
// value s.Percentile(discardPct) and high value s.Percentile(1-discardPct)
// with a check to make sure enough samples have been collected to make
// discardPct meaningful (1/discardPct samples are required).
func (s *Stats) CreateBinsDiscard(nbins int, discardPct float64) {
	if len(s.samples) < int(1.0/discardPct) {
		panic("Not enough samples")
	}
	s.CreateBins(nbins, s.Percentile(discardPct), s.Percentile(1.0-discardPct))
}

// Returns the count and low and high ends of the i'th bin.
//
// The bin interval is (low,high]
func (s Stats) Bin(i int) (count int, low, high Sample) {
	count = s.binCounts[i]
	high = s.bins[i]
	if i == 0 {
		low = -math.MaxFloat64
	} else {
		low = s.bins[i-1]
	}
	return
}
<commit_msg>Add NBins method.<commit_after>// Copyright 2012 The Summstat Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package summstat

import (
	"math"
	"sort"
)

// The type of samples we track statistics for
type Sample float64

// Implements sort.Interface so that we can sort samples for percentiles
type sampleSlice []Sample

func (s sampleSlice) Len() int {
	return len(s)
}

func (s sampleSlice) Less(i, j int) bool {
	return s[i] < s[j]
}

func (s sampleSlice) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}

// A Stats represents descriptive statistics about Samples which are being
// added incrementally.
type Stats struct {
	count     int
	sum       Sample
	sum2      Sample
	max       Sample
	min       Sample
	samples   []Sample
	sorted    bool
	bins      []Sample
	binCounts []int
}

// NewStats returns a new Stats
func NewStats() *Stats {
	return &Stats{
		max: -math.MaxFloat64,
		min: math.MaxFloat64,
	}
}

// AddSample adds a sample value and updates the statistics.
func (s *Stats) AddSample(val Sample) {
	s.count++
	s.sum += val
	s.sum2 += val * val
	if val > s.max {
		s.max = val
	}
	if val < s.min {
		s.min = val
	}
	if len(s.bins) > 0 {
		// TODO: use faster lookup method for large bin counts
		for bin, binVal := range s.bins {
			if val <= binVal {
				s.binCounts[bin]++
				break
			}
		}
	} else {
		s.samples = append(s.samples, val)
		s.sorted = false
	}
}

// Count returns the number of samples added.
func (s Stats) Count() int {
	return s.count
}

// Min returns minimal sample value added.
func (s Stats) Min() Sample {
	if s.min > s.max {
		return 0
	}
	return s.min
}

// Max returns the maximal sample value added.
func (s Stats) Max() Sample {
	if s.min > s.max {
		return 0
	}
	return s.max
}

func (s *Stats) sortSamples() {
	if !s.sorted {
		sort.Sort(sampleSlice(s.samples))
		s.sorted = true
	}
}

// Percentile returns the sample value at the given percentile.
//
// It may not be called after CreateBins, which discards the samples from
// which the percentile is calculated.
func (s Stats) Percentile(pct float64) Sample {
	if len(s.bins) > 0 {
		panic("cannot call Percentile() after CreateBins()")
	}
	if len(s.samples) == 0 {
		return 0
	}
	if pct < 0 {
		panic("pct too small")
	}
	if pct > 1 {
		panic("pct too large")
	}
	s.sortSamples()
	// scale pct into int in [0, len-1]
	// Adding 0.5 turns the implicit floor operation of int() into a rounding operation
	i := int(float64(len(s.samples)-1)*pct + 0.5)
	return s.samples[i]
}

// Median returns the median of the samples.
//
// It may not be called after CreateBins, which discards the samples from
// which the percentile is calculated.
func (s Stats) Median() float64 {
	if len(s.bins) > 0 {
		panic("cannot call Percentile() after CreateBins()")
	}
	l := len(s.samples)
	if l == 0 {
		return 0
	}
	s.sortSamples()
	half, rem := l/2, l%2
	if rem == 0 {
		return (float64(s.samples[half]) + float64(s.samples[half-1])) / 2
	}
	return float64(s.samples[half])
}

// Mean returns the mean of the samples.
func (s Stats) Mean() float64 {
	return float64(s.sum) / float64(s.count)
}

// Stddev returns the standard deviation of the samples.
func (s Stats) Stddev() float64 {
	m := s.Mean()
	return math.Sqrt(float64(s.sum2)/float64(s.count) - m*m)
}

// Spread returns the difference of the maximal and minimal sample values.
func (s Stats) Spread() Sample {
	if s.min > s.max {
		return 0
	}
	return s.max - s.min
}

// CreateBins divides the sample space into nbins bins for tracking counts.
//
// As samples are added, the count for the corresponding bin will be
// incremented and the sample value will not be stored.
//
// This saves memory at the expense of granularity. Percentile() and Median()
// cannot be called after CreateBins() because they are no longer meaningful.
// Instead, use Bin(i) to inspect the distribution of data by bin. Any existing
// stored samples are discarded.
//
// The bins created will be:
//   (-Inf,low], (low, s/nmid+low], (s/nmid+low, 2s/nmid], ..., (high,+Inf)
//   where:
//     s = high - low
//     nmid = nbins-2
//
// Thus, the space (high-low) is divided into nbins-2 equally sized pieces
// and the remaining two bins extend from -math.MaxFloat64 to low and high to
// math.MaxFloat64.
//
// Low must be strictly less than high, so nbins must be at least 3.
func (s *Stats) CreateBins(nbins int, low, high Sample) {
	if high <= low {
		panic("high must be greater than low")
	}
	if nbins < 3 {
		panic("Not enough bins")
	}
	spread := high - low
	s.bins = make([]Sample, nbins)
	s.binCounts = make([]int, nbins)
	for i := 0; i < nbins-1; i++ {
		s.bins[i] = Sample(i)*spread/Sample(nbins-2) + low
	}
	s.bins[nbins-1] = math.MaxFloat64
	// save memory: stop storing samples now that we track by bins
	s.samples = []Sample{}
}

// CreateBinsDiscard is shorthand for calling CreateBins(nbins, ...) with low
// value s.Percentile(discardPct) and high value s.Percentile(1-discardPct)
// with a check to make sure enough samples have been collected to make
// discardPct meaningful (1/discardPct samples are required).
func (s *Stats) CreateBinsDiscard(nbins int, discardPct float64) {
	if len(s.samples) < int(1.0/discardPct) {
		panic("Not enough samples")
	}
	s.CreateBins(nbins, s.Percentile(discardPct), s.Percentile(1.0-discardPct))
}

// Returns the count and low and high ends of the i'th bin.
//
// The bin interval is (low,high]
func (s Stats) Bin(i int) (count int, low, high Sample) {
	count = s.binCounts[i]
	high = s.bins[i]
	if i == 0 {
		low = -math.MaxFloat64
	} else {
		low = s.bins[i-1]
	}
	return
}

// Returns the number of bins
func (s Stats) NBins() int {
	return len(s.bins)
}
<|endoftext|>


<commit_before>package tail

import (
	"bufio"
	"bytes"
	"context"
	"errors"
	"io"
	"os"
	"path/filepath"
	"sync"
	"time"

	fsnotify "github.com/fsnotify/fsnotify"
)

const (
	openRetryInterval = time.Second
	tailOldFileDelay  = 15 * time.Second
	linesCapacity     = 1024
	errorsCapacity    = 16
)

// Line is a line of the target file.
type Line struct {
	Text string
	Time time.Time
}

// Tail tails a file.
type Tail struct {
	Lines  <-chan *Line
	Errors <-chan error

	lines    chan<- *Line
	errors   chan<- error
	filename string
	wg       sync.WaitGroup
	ctx      context.Context
	cancel   context.CancelFunc
}

type tail struct {
	parent *Tail

	file    *os.File
	reader  *bufio.Reader
	watcher *fsnotify.Watcher
	buf     bytes.Buffer
	ctx     context.Context
	cancel  context.CancelFunc
}

// NewTailFile starts tailing a file
func NewTailFile(filename string) (*Tail, error) {
	filename, err := filepath.Abs(filename)
	if err != nil {
		return nil, err
	}

	ctx, cancel := context.WithCancel(context.Background())
	lines := make(chan *Line, linesCapacity)
	errs := make(chan error, errorsCapacity)
	parent := &Tail{
		Lines:    lines,
		Errors:   errs,
		lines:    lines,
		errors:   errs,
		filename: filename,
		ctx:      ctx,
		cancel:   cancel,
	}

	parent.wg.Add(1)
	go func() {
		defer parent.wg.Done()
		parent.runFile(os.SEEK_END)
	}()
	go parent.wait()

	return parent, nil
}

// NewTailReader starts tailing io.Reader
func NewTailReader(reader io.Reader) (*Tail, error) {
	ctx, cancel := context.WithCancel(context.Background())
	lines := make(chan *Line, linesCapacity)
	errs := make(chan error, errorsCapacity)
	parent := &Tail{
		Lines:  lines,
		Errors: errs,
		lines:  lines,
		errors: errs,
		ctx:    ctx,
		cancel: cancel,
	}
	r := ctxReader{
		ctx: ctx,
		r:   reader,
	}
	t := &tail{
		parent: parent,
		reader: bufio.NewReader(r),
		ctx:    ctx,
		cancel: cancel,
	}

	parent.wg.Add(1)
	go func() {
		defer parent.wg.Done()
		t.runReader()
	}()
	go parent.wait()

	return parent, nil
}

// Close stops tailing the file.
func (t *Tail) Close() error {
	t.cancel()
	t.wg.Wait()
	return nil
}

func (t *Tail) wait() {
	t.wg.Wait()
	close(t.errors)
	close(t.lines)
}

// open opens the target file.
// If it does not exist, wait for creating new file.
func (t *Tail) open(seek int) (*tail, error) {
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		return nil, err
	}
	for {
		file, err := os.Open(t.filename)
		if err == nil {
			// success, seek and watch the file.
			if _, err := file.Seek(0, seek); err != nil {
				file.Close()
				watcher.Close()
				return nil, err
			}
			if err := watcher.Add(t.filename); err != nil {
				file.Close()
				watcher.Close()
				return nil, err
			}
			ctx, cancel := context.WithCancel(t.ctx)
			r := ctxReader{
				ctx: ctx,
				r:   file,
			}
			return &tail{
				parent:  t,
				file:    file,
				reader:  bufio.NewReader(r),
				watcher: watcher,
				ctx:     ctx,
				cancel:  cancel,
			}, nil
		}

		// fail. retry...
		seek = io.SeekStart
		timer := time.NewTimer(openRetryInterval)
		select {
		case <-t.ctx.Done():
			timer.Stop()
			return nil, t.ctx.Err()
		case <-timer.C:
		}
	}
}

// runFile tails target files
func (t *Tail) runFile(seek int) {
	child, err := t.open(seek)
	if err != nil {
		if !errors.Is(err, context.Canceled) {
			t.errors <- err
		}
		return
	}

	t.wg.Add(1)
	go func() {
		defer t.wg.Done()
		child.runFile()
	}()
}

// runFile tails a file
func (t *tail) runFile() {
	defer t.watcher.Close()
	defer t.cancel()

	cherr := make(chan error, 1)
	ch := make(chan struct{}, 1)
	defer close(ch)

	t.parent.wg.Add(1)
	go func() {
		defer t.parent.wg.Done()
		for {
			if err := t.restrict(); err != nil {
				select {
				case cherr <- err:
				case <-t.ctx.Done():
				}
				return
			}
			err := t.tail()
			if err == nil {
				continue
			}

			// wait for writing new lines.
			select {
			case cherr <- err:
			case <-t.ctx.Done():
				return
			}
			select {
			case _, ok := <-ch:
				if !ok {
					return
				}
			case <-t.ctx.Done():
				return
			}
		}
	}()

	var renamed bool
	var waiting bool // waiting for writing new lines?
	for {
		select {
		case event := <-t.watcher.Events:
			if (event.Op & fsnotify.Remove) != 0 {
				// the target file is removed, stop tailing.
				return
			}
			if (event.Op & fsnotify.Rename) != 0 {
				// log rotation is detected.
				if !renamed {
					// start to watch creating new file.
					t.parent.wg.Add(1)
					go func() {
						defer t.parent.wg.Done()
						t.parent.runFile(io.SeekStart)
					}()

					// wait a little, and stop tailing old file.
					go func() {
						timer := time.NewTimer(tailOldFileDelay)
						defer timer.Stop()
						select {
						case <-timer.C:
							t.cancel()
						case <-t.ctx.Done():
						}
					}()
				}
				name, err := getFileName(t.file)
				if err != nil {
					t.parent.errors <- err
					return
				}
				if err := t.watcher.Add(name); err != nil {
					t.parent.errors <- err
					return
				}
				renamed = true
			}

			// notify new lines are wrote.
			if waiting {
				ch <- struct{}{}
				waiting = false
			}
		case err := <-cherr:
			if errors.Is(err, io.EOF) {
				waiting = true
				continue
			}
			t.parent.errors <- err
			return
		case err := <-t.watcher.Errors:
			t.parent.errors <- err
			return
		case <-t.ctx.Done():
			return
		}
	}
}

// runReader tails io.Reader
func (t *tail) runReader() {
	defer t.cancel()
	err := t.tail()
	if errors.Is(err, io.EOF) || errors.Is(err, io.ErrClosedPipe) {
		return
	}
	if err != nil {
		if !errors.Is(err, context.Canceled) {
			t.parent.errors <- err
		}
		return
	}
}

// restrict detects a file that is truncated
func (t *tail) restrict() error {
	stat, err := t.file.Stat()
	if err != nil {
		return err
	}
	pos, err := t.file.Seek(0, io.SeekCurrent)
	if err != nil {
		return err
	}
	if stat.Size() < pos {
		// file is truncated. seek to head of file.
		_, err := t.file.Seek(0, io.SeekStart)
		if err != nil {
			return err
		}
	}
	return nil
}

// tail reads lines until EOF
func (t *tail) tail() error {
	for {
		line, err := t.reader.ReadSlice('\n')
		t.buf.Write(line)
		if err == bufio.ErrBufferFull {
			// the reader cannot find EOL in its buffer.
			// continue to read a line.
			continue
		}
		if err != nil {
			return err
		}
		t.parent.lines <- &Line{t.buf.String(), time.Now()}
		t.buf.Reset()
	}
}
<commit_msg>introduce Op.Has<commit_after>package tail

import (
	"bufio"
	"bytes"
	"context"
	"errors"
	"io"
	"os"
	"path/filepath"
	"sync"
	"time"

	fsnotify "github.com/fsnotify/fsnotify"
)

const (
	openRetryInterval = time.Second
	tailOldFileDelay  = 15 * time.Second
	linesCapacity     = 1024
	errorsCapacity    = 16
)

// Line is a line of the target file.
type Line struct {
	Text string
	Time time.Time
}

// Tail tails a file.
type Tail struct {
	Lines  <-chan *Line
	Errors <-chan error

	lines    chan<- *Line
	errors   chan<- error
	filename string
	wg       sync.WaitGroup
	ctx      context.Context
	cancel   context.CancelFunc
}

type tail struct {
	parent *Tail

	file    *os.File
	reader  *bufio.Reader
	watcher *fsnotify.Watcher
	buf     bytes.Buffer
	ctx     context.Context
	cancel  context.CancelFunc
}

// NewTailFile starts tailing a file
func NewTailFile(filename string) (*Tail, error) {
	filename, err := filepath.Abs(filename)
	if err != nil {
		return nil, err
	}

	ctx, cancel := context.WithCancel(context.Background())
	lines := make(chan *Line, linesCapacity)
	errs := make(chan error, errorsCapacity)
	parent := &Tail{
		Lines:    lines,
		Errors:   errs,
		lines:    lines,
		errors:   errs,
		filename: filename,
		ctx:      ctx,
		cancel:   cancel,
	}

	parent.wg.Add(1)
	go func() {
		defer parent.wg.Done()
		parent.runFile(os.SEEK_END)
	}()
	go parent.wait()

	return parent, nil
}

// NewTailReader starts tailing io.Reader
func NewTailReader(reader io.Reader) (*Tail, error) {
	ctx, cancel := context.WithCancel(context.Background())
	lines := make(chan *Line, linesCapacity)
	errs := make(chan error, errorsCapacity)
	parent := &Tail{
		Lines:  lines,
		Errors: errs,
		lines:  lines,
		errors: errs,
		ctx:    ctx,
		cancel: cancel,
	}
	r := ctxReader{
		ctx: ctx,
		r:   reader,
	}
	t := &tail{
		parent: parent,
		reader: bufio.NewReader(r),
		ctx:    ctx,
		cancel: cancel,
	}

	parent.wg.Add(1)
	go func() {
		defer parent.wg.Done()
		t.runReader()
	}()
	go parent.wait()

	return parent, nil
}

// Close stops tailing the file.
func (t *Tail) Close() error {
	t.cancel()
	t.wg.Wait()
	return nil
}

func (t *Tail) wait() {
	t.wg.Wait()
	close(t.errors)
	close(t.lines)
}

// open opens the target file.
// If it does not exist, wait for creating new file.
func (t *Tail) open(seek int) (*tail, error) {
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		return nil, err
	}
	for {
		file, err := os.Open(t.filename)
		if err == nil {
			// success, seek and watch the file.
			if _, err := file.Seek(0, seek); err != nil {
				file.Close()
				watcher.Close()
				return nil, err
			}
			if err := watcher.Add(t.filename); err != nil {
				file.Close()
				watcher.Close()
				return nil, err
			}
			ctx, cancel := context.WithCancel(t.ctx)
			r := ctxReader{
				ctx: ctx,
				r:   file,
			}
			return &tail{
				parent:  t,
				file:    file,
				reader:  bufio.NewReader(r),
				watcher: watcher,
				ctx:     ctx,
				cancel:  cancel,
			}, nil
		}

		// fail. retry...
		seek = io.SeekStart
		timer := time.NewTimer(openRetryInterval)
		select {
		case <-t.ctx.Done():
			timer.Stop()
			return nil, t.ctx.Err()
		case <-timer.C:
		}
	}
}

// runFile tails target files
func (t *Tail) runFile(seek int) {
	child, err := t.open(seek)
	if err != nil {
		if !errors.Is(err, context.Canceled) {
			t.errors <- err
		}
		return
	}

	t.wg.Add(1)
	go func() {
		defer t.wg.Done()
		child.runFile()
	}()
}

// runFile tails a file
func (t *tail) runFile() {
	defer t.watcher.Close()
	defer t.cancel()

	cherr := make(chan error, 1)
	ch := make(chan struct{}, 1)
	defer close(ch)

	t.parent.wg.Add(1)
	go func() {
		defer t.parent.wg.Done()
		for {
			if err := t.restrict(); err != nil {
				select {
				case cherr <- err:
				case <-t.ctx.Done():
				}
				return
			}
			err := t.tail()
			if err == nil {
				continue
			}

			// wait for writing new lines.
			select {
			case cherr <- err:
			case <-t.ctx.Done():
				return
			}
			select {
			case _, ok := <-ch:
				if !ok {
					return
				}
			case <-t.ctx.Done():
				return
			}
		}
	}()

	var renamed bool
	var waiting bool // waiting for writing new lines?
	for {
		select {
		case event := <-t.watcher.Events:
			if event.Op.Has(fsnotify.Remove) {
				// the target file is removed, stop tailing.
				return
			}
			if event.Op.Has(fsnotify.Rename) {
				// log rotation is detected.
				if !renamed {
					// start to watch creating new file.
					t.parent.wg.Add(1)
					go func() {
						defer t.parent.wg.Done()
						t.parent.runFile(io.SeekStart)
					}()

					// wait a little, and stop tailing old file.
					go func() {
						timer := time.NewTimer(tailOldFileDelay)
						defer timer.Stop()
						select {
						case <-timer.C:
							t.cancel()
						case <-t.ctx.Done():
						}
					}()
				}
				name, err := getFileName(t.file)
				if err != nil {
					t.parent.errors <- err
					return
				}
				if err := t.watcher.Add(name); err != nil {
					t.parent.errors <- err
					return
				}
				renamed = true
			}

			// notify new lines are wrote.
			if waiting {
				ch <- struct{}{}
				waiting = false
			}
		case err := <-cherr:
			if errors.Is(err, io.EOF) {
				waiting = true
				continue
			}
			t.parent.errors <- err
			return
		case err := <-t.watcher.Errors:
			t.parent.errors <- err
			return
		case <-t.ctx.Done():
			return
		}
	}
}

// runReader tails io.Reader
func (t *tail) runReader() {
	defer t.cancel()
	err := t.tail()
	if errors.Is(err, io.EOF) || errors.Is(err, io.ErrClosedPipe) {
		return
	}
	if err != nil {
		if !errors.Is(err, context.Canceled) {
			t.parent.errors <- err
		}
		return
	}
}

// restrict detects a file that is truncated
func (t *tail) restrict() error {
	stat, err := t.file.Stat()
	if err != nil {
		return err
	}
	pos, err := t.file.Seek(0, io.SeekCurrent)
	if err != nil {
		return err
	}
	if stat.Size() < pos {
		// file is truncated. seek to head of file.
		_, err := t.file.Seek(0, io.SeekStart)
		if err != nil {
			return err
		}
	}
	return nil
}

// tail reads lines until EOF
func (t *tail) tail() error {
	for {
		line, err := t.reader.ReadSlice('\n')
		t.buf.Write(line)
		if err == bufio.ErrBufferFull {
			// the reader cannot find EOL in its buffer.
			// continue to read a line.
			continue
		}
		if err != nil {
			return err
		}
		t.parent.lines <- &Line{t.buf.String(), time.Now()}
		t.buf.Reset()
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2013 ActiveState Software Inc. All rights reserved.

package tail

import (
	"bufio"
	"fmt"
	"github.com/ActiveState/tail/watch"
	"io"
	"launchpad.net/tomb"
	"log"
	"os"
	"time"
)

var (
	ErrStop = fmt.Errorf("tail should now stop")
)

type Line struct {
	Text string
	Time time.Time
	Err  error // Error from tail
}

// SeekInfo represents arguments to `os.Seek`
type SeekInfo struct {
	Offset int64
	Whence int // os.SEEK_*
}

// Config is used to specify how a file must be tailed.
type Config struct {
	// File-specifc
	Location  *SeekInfo // Seek to this location before tailing
	ReOpen    bool      // Reopen recreated files (tail -F)
	MustExist bool      // Fail early if the file does not exist
	Poll      bool      // Poll for file changes instead of using inotify
	LimitRate int64     // Maximum read rate (lines per second)

	// Generic IO
	Follow      bool // Continue looking for new lines (tail -f)
	MaxLineSize int  // If non-zero, split longer lines into multiple lines
}

type Tail struct {
	Filename string
	Lines    chan *Line
	Config

	file    *os.File
	reader  *bufio.Reader
	watcher watch.FileWatcher
	changes *watch.FileChanges
	rateMon *RateMonitor

	tomb.Tomb // provides: Done, Kill, Dying
}

// TailFile begins tailing the file. Output stream is made available
// via the `Tail.Lines` channel. To handle errors during tailing,
// invoke the `Wait` or `Err` method after finishing reading from the
// `Lines` channel.
func TailFile(filename string, config Config) (*Tail, error) {
	if config.ReOpen && !config.Follow {
		panic("cannot set ReOpen without Follow.")
	}

	t := &Tail{
		Filename: filename,
		Lines:    make(chan *Line),
		Config:   config}

	t.rateMon = new(RateMonitor)

	if t.Poll {
		t.watcher = watch.NewPollingFileWatcher(filename)
	} else {
		t.watcher = watch.NewInotifyFileWatcher(filename)
	}

	if t.MustExist {
		var err error
		t.file, err = os.Open(t.Filename)
		if err != nil {
			return nil, err
		}
	}

	go t.tailFileSync()

	return t, nil
}

// Return the file's current position, like stdio's ftell().
// But this value is not very accurate.
// it may readed one line in the chan(tail.Lines),
// so it may lost one line.
func (tail *Tail) Tell() (offset int64, err error) {
	if tail.file == nil {
		return
	}
	offset, err = tail.file.Seek(0, os.SEEK_CUR)
	if err == nil {
		offset -= int64(tail.reader.Buffered())
	}
	return
}

// Stop stops the tailing activity.
func (tail *Tail) Stop() error {
	tail.Kill(nil)
	return tail.Wait()
}

func (tail *Tail) close() {
	close(tail.Lines)
	if tail.file != nil {
		tail.file.Close()
	}
}

func (tail *Tail) reopen() error {
	if tail.file != nil {
		tail.file.Close()
	}
	for {
		var err error
		tail.file, err = os.Open(tail.Filename)
		if err != nil {
			if os.IsNotExist(err) {
				log.Printf("Waiting for %s to appear...", tail.Filename)
				if err := tail.watcher.BlockUntilExists(&tail.Tomb); err != nil {
					if err == tomb.ErrDying {
						return err
					}
					return fmt.Errorf("Failed to detect creation of %s: %s", tail.Filename, err)
				}
				continue
			}
			return fmt.Errorf("Unable to open file %s: %s", tail.Filename, err)
		}
		break
	}
	return nil
}

func (tail *Tail) readLine() ([]byte, error) {
	line, isPrefix, err := tail.reader.ReadLine()
	if !isPrefix {
		return line, err
	}

	buf := append([]byte(nil), line...)
	for isPrefix && err == nil {
		line, isPrefix, err = tail.reader.ReadLine()
		buf = append(buf, line...)
	}
	return buf, err
}

func (tail *Tail) tailFileSync() {
	defer tail.Done()
	defer tail.close()

	if !tail.MustExist {
		// deferred first open.
		err := tail.reopen()
		if err != nil {
			if err != tomb.ErrDying {
				tail.Kill(err)
			}
			return
		}
	}

	// Seek to requested location on first open of the file.
	if tail.Location != nil {
		_, err := tail.file.Seek(tail.Location.Offset, tail.Location.Whence)
		// log.Printf("Seeked %s - %+v\n", tail.Filename, tail.Location)
		if err != nil {
			tail.Killf("Seek error on %s: %s", tail.Filename, err)
			return
		}
	}

	tail.reader = bufio.NewReader(tail.file)

	// Read line by line.
	for {
		line, err := tail.readLine()

		switch err {
		case nil:
			if line != nil {
				cooloff := !tail.sendLine(line)
				if cooloff {
					msg := "Too much activity; entering a cool-off period"
					tail.Lines <- &Line{
						msg,
						time.Now(),
						fmt.Errorf(msg)}
					// Wait a second before seeking till the end of
					// file when rate limit is reached.
					select {
					case <-time.After(time.Second):
					case <-tail.Dying():
						return
					}
					_, err := tail.file.Seek(0, 2) // Seek to fine end
					if err != nil {
						tail.Killf("Seek error on %s: %s", tail.Filename, err)
						return
					}
				}
			}
		case io.EOF:
			if !tail.Follow {
				return
			}
			// When EOF is reached, wait for more data to become
			// available. Wait strategy is based on the `tail.watcher`
			// implementation (inotify or polling).
			err := tail.waitForChanges()
			if err != nil {
				if err != ErrStop {
					tail.Kill(err)
				}
				return
			}
		default: // non-EOF error
			tail.Killf("Error reading %s: %s", tail.Filename, err)
			return
		}

		select {
		case <-tail.Dying():
			return
		default:
		}
	}
}

// waitForChanges waits until the file has been appended, deleted,
// moved or truncated. When moved or deleted - the file will be
// reopened if ReOpen is true. Truncated files are always reopened.
func (tail *Tail) waitForChanges() error {
	if tail.changes == nil {
		st, err := tail.file.Stat()
		if err != nil {
			return err
		}
		tail.changes = tail.watcher.ChangeEvents(&tail.Tomb, st)
	}

	select {
	case <-tail.changes.Modified:
		return nil
	case <-tail.changes.Deleted:
		tail.changes = nil
		if tail.ReOpen {
			// XXX: we must not log from a library.
			log.Printf("Re-opening moved/deleted file %s ...", tail.Filename)
			if err := tail.reopen(); err != nil {
				return err
			}
			log.Printf("Successfully reopened %s", tail.Filename)
			tail.reader = bufio.NewReader(tail.file)
			return nil
		} else {
			log.Printf("Stopping tail as file no longer exists: %s", tail.Filename)
			return ErrStop
		}
	case <-tail.changes.Truncated:
		// Always reopen truncated files (Follow is true)
		log.Printf("Re-opening truncated file %s ...", tail.Filename)
		if err := tail.reopen(); err != nil {
			return err
		}
		log.Printf("Successfully reopened truncated %s", tail.Filename)
		tail.reader = bufio.NewReader(tail.file)
		return nil
	case <-tail.Dying():
		return ErrStop
	}
	panic("unreachable")
}

// sendLine sends the line(s) to Lines channel, splitting longer lines
// if necessary. Return false if rate limit is reached.
func (tail *Tail) sendLine(line []byte) bool {
	now := time.Now()
	nowUnix := now.Unix()
	lines := []string{string(line)}

	// Split longer lines
	if tail.MaxLineSize > 0 && len(line) > tail.MaxLineSize {
		lines = partitionString(
			string(line), tail.MaxLineSize)
	}

	for _, line := range lines {
		tail.Lines <- &Line{line, now, nil}
		rate := tail.rateMon.Tick(nowUnix)
		if tail.LimitRate > 0 && rate > tail.LimitRate {
			log.Printf("Rate limit (%v < %v) reached on file (%v); entering 1s cooloff period.\n",
				tail.LimitRate,
				rate,
				tail.Filename)
			return false
		}
	}

	return true
}
<commit_msg>my pkg<commit_after>// Copyright (c) 2013 ActiveState Software Inc. All rights reserved.

package tail

import (
	"bufio"
	"fmt"
	"github.com/funkygao/tail/watch"
	"io"
	"launchpad.net/tomb"
	"log"
	"os"
	"time"
)

var (
	ErrStop = fmt.Errorf("tail should now stop")
)

type Line struct {
	Text string
	Time time.Time
	Err  error // Error from tail
}

// SeekInfo represents arguments to `os.Seek`
type SeekInfo struct {
	Offset int64
	Whence int // os.SEEK_*
}

// Config is used to specify how a file must be tailed.
type Config struct {
	// File-specifc
	Location  *SeekInfo // Seek to this location before tailing
	ReOpen    bool      // Reopen recreated files (tail -F)
	MustExist bool      // Fail early if the file does not exist
	Poll      bool      // Poll for file changes instead of using inotify
	LimitRate int64     // Maximum read rate (lines per second)

	// Generic IO
	Follow      bool // Continue looking for new lines (tail -f)
	MaxLineSize int  // If non-zero, split longer lines into multiple lines
}

type Tail struct {
	Filename string
	Lines    chan *Line
	Config

	file    *os.File
	reader  *bufio.Reader
	watcher watch.FileWatcher
	changes *watch.FileChanges
	rateMon *RateMonitor

	tomb.Tomb // provides: Done, Kill, Dying
}

// TailFile begins tailing the file. Output stream is made available
// via the `Tail.Lines` channel. To handle errors during tailing,
// invoke the `Wait` or `Err` method after finishing reading from the
// `Lines` channel.
func TailFile(filename string, config Config) (*Tail, error) {
	if config.ReOpen && !config.Follow {
		panic("cannot set ReOpen without Follow.")
	}

	t := &Tail{
		Filename: filename,
		Lines:    make(chan *Line),
		Config:   config}

	t.rateMon = new(RateMonitor)

	if t.Poll {
		t.watcher = watch.NewPollingFileWatcher(filename)
	} else {
		t.watcher = watch.NewInotifyFileWatcher(filename)
	}

	if t.MustExist {
		var err error
		t.file, err = os.Open(t.Filename)
		if err != nil {
			return nil, err
		}
	}

	go t.tailFileSync()

	return t, nil
}

// Return the file's current position, like stdio's ftell().
// But this value is not very accurate.
// it may readed one line in the chan(tail.Lines),
// so it may lost one line.
func (tail *Tail) Tell() (offset int64, err error) {
	if tail.file == nil {
		return
	}
	offset, err = tail.file.Seek(0, os.SEEK_CUR)
	if err == nil {
		offset -= int64(tail.reader.Buffered())
	}
	return
}

// Stop stops the tailing activity.
func (tail *Tail) Stop() error {
	tail.Kill(nil)
	return tail.Wait()
}

func (tail *Tail) close() {
	close(tail.Lines)
	if tail.file != nil {
		tail.file.Close()
	}
}

func (tail *Tail) reopen() error {
	if tail.file != nil {
		tail.file.Close()
	}
	for {
		var err error
		tail.file, err = os.Open(tail.Filename)
		if err != nil {
			if os.IsNotExist(err) {
				log.Printf("Waiting for %s to appear...", tail.Filename)
				if err := tail.watcher.BlockUntilExists(&tail.Tomb); err != nil {
					if err == tomb.ErrDying {
						return err
					}
					return fmt.Errorf("Failed to detect creation of %s: %s", tail.Filename, err)
				}
				continue
			}
			return fmt.Errorf("Unable to open file %s: %s", tail.Filename, err)
		}
		break
	}
	return nil
}

func (tail *Tail) readLine() ([]byte, error) {
	line, isPrefix, err := tail.reader.ReadLine()
	if !isPrefix {
		return line, err
	}

	buf := append([]byte(nil), line...)
	for isPrefix && err == nil {
		line, isPrefix, err = tail.reader.ReadLine()
		buf = append(buf, line...)
	}
	return buf, err
}

func (tail *Tail) tailFileSync() {
	defer tail.Done()
	defer tail.close()

	if !tail.MustExist {
		// deferred first open.
		err := tail.reopen()
		if err != nil {
			if err != tomb.ErrDying {
				tail.Kill(err)
			}
			return
		}
	}

	// Seek to requested location on first open of the file.
	if tail.Location != nil {
		_, err := tail.file.Seek(tail.Location.Offset, tail.Location.Whence)
		// log.Printf("Seeked %s - %+v\n", tail.Filename, tail.Location)
		if err != nil {
			tail.Killf("Seek error on %s: %s", tail.Filename, err)
			return
		}
	}

	tail.reader = bufio.NewReader(tail.file)

	// Read line by line.
	for {
		line, err := tail.readLine()

		switch err {
		case nil:
			if line != nil {
				cooloff := !tail.sendLine(line)
				if cooloff {
					msg := "Too much activity; entering a cool-off period"
					tail.Lines <- &Line{
						msg,
						time.Now(),
						fmt.Errorf(msg)}
					// Wait a second before seeking till the end of
					// file when rate limit is reached.
					select {
					case <-time.After(time.Second):
					case <-tail.Dying():
						return
					}
					_, err := tail.file.Seek(0, 2) // Seek to fine end
					if err != nil {
						tail.Killf("Seek error on %s: %s", tail.Filename, err)
						return
					}
				}
			}
		case io.EOF:
			if !tail.Follow {
				return
			}
			// When EOF is reached, wait for more data to become
			// available. Wait strategy is based on the `tail.watcher`
			// implementation (inotify or polling).
			err := tail.waitForChanges()
			if err != nil {
				if err != ErrStop {
					tail.Kill(err)
				}
				return
			}
		default: // non-EOF error
			tail.Killf("Error reading %s: %s", tail.Filename, err)
			return
		}

		select {
		case <-tail.Dying():
			return
		default:
		}
	}
}

// waitForChanges waits until the file has been appended, deleted,
// moved or truncated. When moved or deleted - the file will be
// reopened if ReOpen is true. Truncated files are always reopened.
func (tail *Tail) waitForChanges() error {
	if tail.changes == nil {
		st, err := tail.file.Stat()
		if err != nil {
			return err
		}
		tail.changes = tail.watcher.ChangeEvents(&tail.Tomb, st)
	}

	select {
	case <-tail.changes.Modified:
		return nil
	case <-tail.changes.Deleted:
		tail.changes = nil
		if tail.ReOpen {
			// XXX: we must not log from a library.
			log.Printf("Re-opening moved/deleted file %s ...", tail.Filename)
			if err := tail.reopen(); err != nil {
				return err
			}
			log.Printf("Successfully reopened %s", tail.Filename)
			tail.reader = bufio.NewReader(tail.file)
			return nil
		} else {
			log.Printf("Stopping tail as file no longer exists: %s", tail.Filename)
			return ErrStop
		}
	case <-tail.changes.Truncated:
		// Always reopen truncated files (Follow is true)
		log.Printf("Re-opening truncated file %s ...", tail.Filename)
		if err := tail.reopen(); err != nil {
			return err
		}
		log.Printf("Successfully reopened truncated %s", tail.Filename)
		tail.reader = bufio.NewReader(tail.file)
		return nil
	case <-tail.Dying():
		return ErrStop
	}
	panic("unreachable")
}

// sendLine sends the line(s) to Lines channel, splitting longer lines
// if necessary. Return false if rate limit is reached.
func (tail *Tail) sendLine(line []byte) bool {
	now := time.Now()
	nowUnix := now.Unix()
	lines := []string{string(line)}

	// Split longer lines
	if tail.MaxLineSize > 0 && len(line) > tail.MaxLineSize {
		lines = partitionString(
			string(line), tail.MaxLineSize)
	}

	for _, line := range lines {
		tail.Lines <- &Line{line, now, nil}
		rate := tail.rateMon.Tick(nowUnix)
		if tail.LimitRate > 0 && rate > tail.LimitRate {
			log.Printf("Rate limit (%v < %v) reached on file (%v); entering 1s cooloff period.\n",
				tail.LimitRate,
				rate,
				tail.Filename)
			return false
		}
	}

	return true
}
<|endoftext|>


<commit_before>// Go f*ck your self -- a BrainF*ck interpreter
package main

import (
	"fmt"
	"unicode"
)

const size = 512

type Tape struct {
	cell []int
	pos  int
}

func NewTape() *Tape {
	return &Tape{
		cell: make([]int, size),
		pos:  0,
	}
}

func (t *Tape) Move(n int) {
	t.pos += n
	if t.pos >= len(t.cell) {
		t.cell = append(t.cell, make([]int, size)...)
	} else if t.pos < 0 {
		t.cell = append(make([]int, size), t.cell...)
		t.pos += size
	}
}

func (t *Tape) Add(n int) {
	t.cell[t.pos] += n
}

func (t *Tape) Print() {
	if c := t.cell[t.pos]; c > unicode.MaxASCII {
		fmt.Printf("%d", c)
	} else {
		fmt.Printf("%c", c)
	}
}

func (t *Tape) Scan() {
	fmt.Scanf("%c", &t.cell[t.pos])
}

func (t *Tape) IsZero() bool {
	return t.cell[t.pos] == 0
}

func (t *Tape) Clear() {
	t.cell[t.pos] = 0
}

func (t *Tape) Mult(off, arg int) {
	v := t.cell[t.pos]
	t.cell[t.pos] = 0
	t.Move(off)
	t.cell[t.pos] += v * arg
	t.pos -= off
}
<commit_msg>Reuse code<commit_after>// Go f*ck your self -- a BrainF*ck interpreter
package main

import (
	"fmt"
	"unicode"
)

const size = 512

type Tape struct {
	cell []int
	pos  int
}

func NewTape() *Tape {
	return &Tape{
		cell: make([]int, size),
		pos:  0,
	}
}

func (t *Tape) Move(n int) {
	t.pos += n
	if t.pos >= len(t.cell) {
		t.cell = append(t.cell, make([]int, size)...)
	} else if t.pos < 0 {
		t.cell = append(make([]int, size), t.cell...)
		t.pos += size
	}
}

func (t *Tape) Add(n int) {
	t.cell[t.pos] += n
}

func (t *Tape) Print() {
	if c := t.cell[t.pos]; c > unicode.MaxASCII {
		fmt.Printf("%d", c)
	} else {
		fmt.Printf("%c", c)
	}
}

func (t *Tape) Scan() {
	fmt.Scanf("%c", &t.cell[t.pos])
}

func (t *Tape) IsZero() bool {
	return t.cell[t.pos] == 0
}

func (t *Tape) Clear() {
	t.cell[t.pos] = 0
}

func (t *Tape) Mult(off, arg int) {
	v := t.cell[t.pos]
	t.Move(off)
	t.Add(v * arg)
	t.Move(-off)
	t.Clear()
}
<|endoftext|>


<commit_before>package tasker

import (
	"context"
	"errors"
	"time"

	"github.com/NoneBorder/dora"
	"github.com/astaxie/beego/orm"
)

const (
	TaskStatPending = "pending"
	TaskStatRunning = "running"
	TaskStatRetry   = "retry"
	TaskStatFailed  = "failed"
	TaskStatSuccess = "success"
)

const (
	TaskCtxKeyLog       = "stdout"          // normal log record
	TaskCtxKeyExecAfter = "task-exec-after" // plan task exec time
)

// Task is the core object for tasker package.
type Task struct {
	Id         uint64
	Topic      string
	Status     string
	Timeout    int // 超时ms
	ExecAfter  time.Time `orm:"type(datetime);default(1971-01-01 00:00:00)"`
	Retry      int
	Input      string    `orm:"type(text)"`
	WorkerId   uint64
	WorkerFQDN string    `orm:"column(worker_fqdn);type(text)"`
	Created    time.Time `orm:"auto_now_add"`
	Updated    time.Time `orm:"auto_now"`
	Log        string    `orm:"type(text)"`
	Errlog     string    `orm:"type(text)"`
}

// ConsumeFn represents consume func definition.
type ConsumeFn func(ctx *context.Context, Input string, WorkerID uint64) (err error)

func (self *Task) TableIndex() [][]string {
	return [][]string{
		[]string{"Topic", "Status"},
	}
}

func (self *Task) TableEngine() string {
	return "INNODB"
}

// NewSimpleTask is new task with default settings.
func NewSimpleTask(topic, input string) *Task {
	return &Task{
		Topic:     topic,
		Status:    TaskStatPending,
		Timeout:   5000,
		ExecAfter: zeroTime,
		Retry:     3,
		Input:     input,
		WorkerId:  0,
	}
}

// Publish Task to Msg Queue.
func (self *Task) Publish() (err error) {
	o := orm.NewOrm()
	if _, err = o.Insert(self); err == nil {
		Stat(self.Topic, TaskStatPending, time.Duration(0))
	}
	return
}

func (self *Task) exec(ctx *context.Context, fn ConsumeFn) (err error) {
	errC := make(chan error)
	select {
	case errC <- fn(ctx, self.Input, self.WorkerId):
		err = <-errC

	case <-(*ctx).Done():
		err = (*ctx).Err()
	}

	return
}

func (self *Task) consume(fn ConsumeFn) bool {
	startT := time.Now().Local()
	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(self.Timeout+1)*time.Millisecond)
	defer cancel()

	err := self.exec(&ctx, fn)
	success := err == nil

	if !success {
		// failed
		self.Errlog += err.Error() + "\n"
		self.Status = TaskStatFailed
		self.Retry -= 1
		if self.Retry > 0 {
			self.Status = TaskStatRetry

			// when the job will retry, you can update the next exec time
			if execAfter, ok := ctx.Value(TaskCtxKeyExecAfter).(time.Time); ok {
				self.ExecAfter = execAfter
			}
		}
	} else {
		self.Status = TaskStatSuccess
	}

	// stat task exec
	Stat(self.Topic, self.Status, time.Since(startT))

	// workerid set to 0
	self.WorkerId = 0

	// set stdout log from ctx
	if log, ok := ctx.Value(TaskCtxKeyLog).(string); ok && log != "" {
		self.Log = log
	}

	// update db record
	o := orm.NewOrm()
	for i := 0; i < 3; i++ {
		if _, err := o.Update(self); err != nil {
			dora.Error().Msgf("set task state failed: err=%s, task=%s", err.Error(), self)
		} else {
			break
		}
	}

	return success
}

// Consume a topic task.
func Consume(topic string, fn ConsumeFn) (int, error) {
	o := orm.NewOrm()
	if IsMaster {
		// 仅在 master 上对已经僵死的 task 进行重置，减轻 DB 压力
		if _, err := o.Raw(`UPDATE task SET status=?, retry=retry-1, worker_id=0, updated=NOW() WHERE topic=? AND status=?
		AND TIMESTAMPDIFF(SECOND, updated, now())*1000-5000>timeout`,
			TaskStatRetry, topic, TaskStatRunning,
		).Exec(); err != nil {
			dora.Error().Msgf("update dead running task to waiting failed: %s", err.Error())
		}
	}

	maxConsumeTask := cap(RunningTaskChannel[topic]) - len(RunningTaskChannel[topic])
	if maxConsumeTask == 0 {
		return 0, errors.New("reached max running task limit")
	}

	workerID, err := UniqID.NextID()
	if err != nil {
		return 0, err
	}

	res, err := o.Raw(`UPDATE task SET status=?, worker_id=?, worker_fqdn=concat(worker_fqdn, ?), updated=NOW() WHERE
		topic=? AND worker_id=0 AND status IN (?,?) AND exec_after <= NOW() LIMIT ?`,
		TaskStatRunning, workerID, FQDN()+"\n", topic, TaskStatPending, TaskStatRetry, maxConsumeTask,
	).Exec()
	if err != nil {
		// not update db
		return 0, err
	}

	num, _ := res.RowsAffected()
	if num == 0 {
		// no task for consume
		return 0, nil
	}

	tasks := []*Task{}
	_, err = o.Raw("SELECT * FROM task WHERE `topic`=? AND `status`=? AND `worker_id`=?",
		topic, TaskStatRunning, workerID,
	).QueryRows(&tasks)
	if len(tasks) == 0 || err != nil {
		return 0, err
	}

	// TODO: 如何保证程序退出时正在运行的 task 可用？当前超时后 5s 僵尸 task 会被重新激活，但不能保证原子性
	for _, t := range tasks {
		RunningTaskChannel[topic] <- true
		go func(t *Task) {
			t.consume(fn)
			<-RunningTaskChannel[topic]
		}(t)
	}

	return len(tasks), nil
}

// ConsumeCtxStdout write stdout log to context for record
func ConsumeCtxStdout(ctx *context.Context, log string) {
	*ctx = context.WithValue(*ctx, TaskCtxKeyLog, log)
}

// ConsumeCtxExecAfter update run time for next retry
func ConsumeCtxExecAfter(ctx *context.Context, t time.Time) {
	*ctx = context.WithValue(*ctx, TaskCtxKeyExecAfter, t)
}
<commit_msg>fix task can not finished normal<commit_after>package tasker

import (
	"context"
	"errors"
	"time"

	"github.com/NoneBorder/dora"
	"github.com/astaxie/beego/orm"
)

const (
	TaskStatPending = "pending"
	TaskStatRunning = "running"
	TaskStatRetry   = "retry"
	TaskStatFailed  = "failed"
	TaskStatSuccess = "success"
)

const (
	TaskCtxKeyLog       = "stdout"          // normal log record
	TaskCtxKeyExecAfter = "task-exec-after" // plan task exec time
)

// Task is the core object for tasker package.
type Task struct {
	Id         uint64
	Topic      string
	Status     string
	Timeout    int // 超时ms
	ExecAfter  time.Time `orm:"type(datetime);default(1971-01-01 00:00:00)"`
	Retry      int
	Input      string    `orm:"type(text)"`
	WorkerId   uint64
	WorkerFQDN string    `orm:"column(worker_fqdn);type(text)"`
	Created    time.Time `orm:"auto_now_add"`
	Updated    time.Time `orm:"auto_now"`
	Log        string    `orm:"type(text)"`
	Errlog     string    `orm:"type(text)"`
}

// ConsumeFn represents consume func definition.
type ConsumeFn func(ctx *context.Context, Input string, WorkerID uint64) (err error)

func (self *Task) TableIndex() [][]string {
	return [][]string{
		[]string{"Topic", "Status"},
	}
}

func (self *Task) TableEngine() string {
	return "INNODB"
}

// NewSimpleTask is new task with default settings.
func NewSimpleTask(topic, input string) *Task {
	return &Task{
		Topic:     topic,
		Status:    TaskStatPending,
		Timeout:   5000,
		ExecAfter: zeroTime,
		Retry:     3,
		Input:     input,
		WorkerId:  0,
	}
}

// Publish Task to Msg Queue.
func (self *Task) Publish() (err error) {
	o := orm.NewOrm()
	if _, err = o.Insert(self); err == nil {
		Stat(self.Topic, TaskStatPending, time.Duration(0))
	}
	return
}

func (self *Task) exec(ctx *context.Context, fn ConsumeFn) (err error) {
	// errC in select no other goroutine, must use buffered channel
	errC := make(chan error, 1)
	select {
	case errC <- fn(ctx, self.Input, self.WorkerId):
		err = <-errC

	case <-(*ctx).Done():
		err = (*ctx).Err()
	}

	return
}

func (self *Task) consume(fn ConsumeFn) bool {
	startT := time.Now().Local()
	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(self.Timeout+1)*time.Millisecond)
	defer cancel()

	err := self.exec(&ctx, fn)
	success := err == nil

	if !success {
		// failed
		self.Errlog += err.Error() + "\n"
		self.Status = TaskStatFailed
		self.Retry -= 1
		if self.Retry > 0 {
			self.Status = TaskStatRetry

			// when the job will retry, you can update the next exec time
			if execAfter, ok := ctx.Value(TaskCtxKeyExecAfter).(time.Time); ok {
				self.ExecAfter = execAfter
			}
		}
	} else {
		self.Status = TaskStatSuccess
	}

	// stat task exec
	Stat(self.Topic, self.Status, time.Since(startT))

	// workerid set to 0
	self.WorkerId = 0

	// set stdout log from ctx
	if log, ok := ctx.Value(TaskCtxKeyLog).(string); ok && log != "" {
		self.Log = log
	}

	// update db record
	o := orm.NewOrm()
	for i := 0; i < 3; i++ {
		if _, err := o.Update(self); err != nil {
			dora.Error().Msgf("set task state failed: err=%s, task=%s", err.Error(), self)
		} else {
			break
		}
	}

	return success
}

// Consume a topic task.
func Consume(topic string, fn ConsumeFn) (int, error) {
	o := orm.NewOrm()
	if IsMaster {
		// 仅在 master 上对已经僵死的 task 进行重置，减轻 DB 压力
		if _, err := o.Raw(`UPDATE task SET status=?, retry=retry-1, worker_id=0, updated=NOW() WHERE topic=? AND status=?
		AND TIMESTAMPDIFF(SECOND, updated, now())*1000-5000>timeout`,
			TaskStatRetry, topic, TaskStatRunning,
		).Exec(); err != nil {
			dora.Error().Msgf("update dead running task to waiting failed: %s", err.Error())
		}
	}

	maxConsumeTask := cap(RunningTaskChannel[topic]) - len(RunningTaskChannel[topic])
	if maxConsumeTask == 0 {
		return 0, errors.New("reached max running task limit")
	}

	workerID, err := UniqID.NextID()
	if err != nil {
		return 0, err
	}

	res, err := o.Raw(`UPDATE task SET status=?, worker_id=?, worker_fqdn=concat(worker_fqdn, ?), updated=NOW() WHERE
		topic=? AND worker_id=0 AND status IN (?,?) AND exec_after <= NOW() LIMIT ?`,
		TaskStatRunning, workerID, FQDN()+"\n", topic, TaskStatPending, TaskStatRetry, maxConsumeTask,
	).Exec()
	if err != nil {
		// not update db
		return 0, err
	}

	num, _ := res.RowsAffected()
	if num == 0 {
		// no task for consume
		return 0, nil
	}

	tasks := []*Task{}
	_, err = o.Raw("SELECT * FROM task WHERE `topic`=? AND `status`=? AND `worker_id`=?",
		topic, TaskStatRunning, workerID,
	).QueryRows(&tasks)
	if len(tasks) == 0 || err != nil {
		return 0, err
	}

	// TODO: 如何保证程序退出时正在运行的 task 可用？当前超时后 5s 僵尸 task 会被重新激活，但不能保证原子性
	for _, t := range tasks {
		RunningTaskChannel[topic] <- true
		go func(t *Task) {
			t.consume(fn)
			<-RunningTaskChannel[topic]
		}(t)
	}

	return len(tasks), nil
}

// ConsumeCtxStdout write stdout log to context for record
func ConsumeCtxStdout(ctx *context.Context, log string) {
	*ctx = context.WithValue(*ctx, TaskCtxKeyLog, log)
}

// ConsumeCtxExecAfter update run time for next retry
func ConsumeCtxExecAfter(ctx *context.Context, t time.Time) {
	*ctx = context.WithValue(*ctx, TaskCtxKeyExecAfter, t)
}
<|endoftext|>


<commit_before>//go:generate go run gen.go

package version

// Version tells us the app version string
const Version = "3.3.0"
<commit_msg>Bump to v4.0.0-alpha1<commit_after>//go:generate go run gen.go

package version

// Version tells us the app version string
const Version = "4.0.0-alpha1"
<|endoftext|>