<commit_before>// Copyright (c) 2016 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 http

import (
	"context"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"strings"
	"time"

	"go.uber.org/yarpc/internal/errors"
	"go.uber.org/yarpc/peer"
	"go.uber.org/yarpc/peer/hostport"
	"go.uber.org/yarpc/peer/single"
	"go.uber.org/yarpc/transport"

	"github.com/opentracing/opentracing-go"
	"github.com/opentracing/opentracing-go/ext"
	"go.uber.org/atomic"
)

var errOutboundNotStarted = errors.ErrOutboundNotStarted("http.Outbound")

// this ensures the HTTP outbound implements both transport.Outbound interfaces
var (
	_ transport.UnaryOutbound  = (*Outbound)(nil)
	_ transport.OnewayOutbound = (*Outbound)(nil)
)

var defaultURLTemplate, _ = url.Parse("http://localhost")

// NewOutbound builds a new HTTP outbound that sends requests to the given
// URL.
//
// Deprecated: create outbounds through NewPeerListOutbound instead
func NewOutbound(urlStr string, opts ...AgentOption) *Outbound {
	agent := NewAgent(opts...)

	_, hp := parseURL(urlStr)
	peerID := hostport.PeerIdentifier(hp)
	c := single.New(peerID, agent)

	err := c.Start()
	if err != nil {
		// This should never happen, single shouldn't return an error here
		panic(fmt.Sprintf("could not start single peerChooser, err: %s", err))
	}

	return NewChooserOutbound(c).WithURLTemplate(urlStr)
}

func parseURL(urlStr string) (*url.URL, string) {
	parsedURL, err := url.Parse(urlStr)
	if err != nil {
		panic(fmt.Sprintf("invalid url: %s, err: %s", urlStr, err))
	}

	return parsedURL, parsedURL.Host
}

// NewChooserOutbound builds a new HTTP outbound built around a PeerList
// for getting potential downstream hosts.
// Chooser.Choose MUST return *hostport.Peer objects.
// Chooser.Start MUST be called before Outbound.Start
func NewChooserOutbound(chooser peer.Chooser) *Outbound {
	return &Outbound{
		started:     atomic.NewBool(false),
		chooser:     chooser,
		urlTemplate: defaultURLTemplate,
		tracer:      opentracing.GlobalTracer(),
	}
}

// Outbound is an HTTP UnaryOutbound and OnewayOutbound
type Outbound struct {
	started     *atomic.Bool
	chooser     peer.Chooser
	urlTemplate *url.URL
	tracer      opentracing.Tracer
}

// WithURLTemplate configures an alternate URL template.
// The host:port portion of the URL template gets replaced by the chosen peer's
// identifier for each outbound request.
func (o *Outbound) WithURLTemplate(URL string) *Outbound {
	parsedURL, err := url.Parse(URL)
	if err != nil {
		log.Fatalf("failed to configure HTTP outbound: invalid URL template %q: %s", URL, err)
	}
	o.urlTemplate = parsedURL
	return o
}

// WithTracer configures a tracer for the outbound
func (o *Outbound) WithTracer(tracer opentracing.Tracer) *Outbound {
	o.tracer = tracer
	return o
}

// Start the HTTP outbound
func (o *Outbound) Start() error {
	o.started.Swap(true)
	return nil
}

// Stop the HTTP outbound
func (o *Outbound) Stop() error {
	o.started.Swap(false)
	return nil
}

// Call makes a HTTP request
func (o *Outbound) Call(ctx context.Context, treq *transport.Request) (*transport.Response, error) {
	if !o.started.Load() {
		// panic because there's no recovery from this
		panic(errOutboundNotStarted)
	}
	start := time.Now()
	deadline, _ := ctx.Deadline()
	ttl := deadline.Sub(start)

	p, err := o.getPeerForRequest(ctx, treq)
	if err != nil {
		return nil, err
	}
	p.StartRequest(nil)
	defer p.EndRequest(nil)

	req, err := o.createRequest(p, treq)
	if err != nil {
		return nil, err
	}

	req.Header = applicationHeaders.ToHTTPHeaders(treq.Headers, nil)
	ctx, req, span := o.withOpentracingSpan(ctx, req, treq, start)
	defer span.Finish()
	req = o.withCoreHeaders(req, treq, ttl)

	client, err := o.getHTTPClient(p)
	if err != nil {
		return nil, err
	}

	response, err := client.Do(req.WithContext(ctx))

	if err != nil {
		// Workaround borrowed from ctxhttp until
		// https://github.com/golang/go/issues/17711 is resolved.
		select {
		case <-ctx.Done():
			err = ctx.Err()
		default:
		}

		span.SetTag("error", true)
		span.LogEvent(err.Error())
		if err == context.DeadlineExceeded {
			end := time.Now()
			return nil, errors.ClientTimeoutError(treq.Service, treq.Procedure, end.Sub(start))
		}

		return nil, err
	}

	span.SetTag("http.status_code", response.StatusCode)

	if response.StatusCode >= 200 && response.StatusCode < 300 {
		appHeaders := applicationHeaders.FromHTTPHeaders(
			response.Header, transport.NewHeaders())
		return &transport.Response{
			Headers: appHeaders,
			Body:    response.Body,
		}, nil
	}

	return nil, getErrFromResponse(response)
}

type ack struct {
	time time.Time
}

func (a ack) String() string {
	return a.time.String()
}

// CallOneway makes a oneway request
func (o *Outbound) CallOneway(ctx context.Context, treq *transport.Request) (transport.Ack, error) {
	if !o.started.Load() {
		// panic because there's no recovery from this
		panic(errOutboundNotStarted)
	}

	p, err := o.getPeerForRequest(ctx, treq)
	if err != nil {
		return nil, err
	}
	p.StartRequest(nil)
	defer p.EndRequest(nil)

	req, err := o.createRequest(p, treq)
	if err != nil {
		return nil, err
	}

	start := time.Now()
	var ttl time.Duration

	req.Header = applicationHeaders.ToHTTPHeaders(treq.Headers, nil)
	ctx, req, span := o.withOpentracingSpan(ctx, req, treq, start)
	defer span.Finish()
	req = o.withCoreHeaders(req, treq, ttl)

	client, err := o.getHTTPClient(p)
	if err != nil {
		return nil, err
	}

	_, err = client.Do(req.WithContext(ctx))
	if err != nil {
		// Workaround borrowed from ctxhttp until
		// https://github.com/golang/go/issues/17711 is resolved.
		select {
		case <-ctx.Done():
			err = ctx.Err()
		default:
		}
	}

	sent := time.Now()

	if err != nil {
		span.SetTag("error", true)
		span.LogEvent(err.Error())
		return nil, err
	}

	return ack{time: sent}, nil
}

func (o *Outbound) getPeerForRequest(ctx context.Context, treq *transport.Request) (*hostport.Peer, error) {
	p, err := o.chooser.Choose(ctx, treq)
	if err != nil {
		return nil, err
	}

	hpPeer, ok := p.(*hostport.Peer)
	if !ok {
		return nil, peer.ErrInvalidPeerConversion{
			Peer:         p,
			ExpectedType: "*hostport.Peer",
		}
	}

	return hpPeer, nil
}

func (o *Outbound) createRequest(p *hostport.Peer, treq *transport.Request) (*http.Request, error) {
	newURL := *o.urlTemplate
	newURL.Host = p.HostPort()
	return http.NewRequest("POST", newURL.String(), treq.Body)
}

func (o *Outbound) withOpentracingSpan(ctx context.Context, req *http.Request, treq *transport.Request, start time.Time) (context.Context, *http.Request, opentracing.Span) {
	// Apply HTTP Context headers for tracing and baggage carried by tracing.
	tracer := o.tracer
	var parent opentracing.SpanContext // ok to be nil
	if parentSpan := opentracing.SpanFromContext(ctx); parentSpan != nil {
		parent = parentSpan.Context()
	}
	span := tracer.StartSpan(
		treq.Procedure,
		opentracing.StartTime(start),
		opentracing.ChildOf(parent),
		opentracing.Tags{
			"rpc.caller":    treq.Caller,
			"rpc.service":   treq.Service,
			"rpc.encoding":  treq.Encoding,
			"rpc.transport": "http",
		},
	)
	ext.PeerService.Set(span, treq.Service)
	ext.SpanKindRPCClient.Set(span)
	ext.HTTPUrl.Set(span, req.URL.String())
	ctx = opentracing.ContextWithSpan(ctx, span)

	tracer.Inject(
		span.Context(),
		opentracing.HTTPHeaders,
		opentracing.HTTPHeadersCarrier(req.Header),
	)

	return ctx, req, span
}

func (o *Outbound) withCoreHeaders(req *http.Request, treq *transport.Request, ttl time.Duration) *http.Request {
	req.Header.Set(CallerHeader, treq.Caller)
	req.Header.Set(ServiceHeader, treq.Service)
	req.Header.Set(ProcedureHeader, treq.Procedure)
	if ttl != 0 {
		req.Header.Set(TTLMSHeader, fmt.Sprintf("%d", ttl/time.Millisecond))
	}
	if treq.ShardKey != "" {
		req.Header.Set(ShardKeyHeader, treq.ShardKey)
	}
	if treq.RoutingKey != "" {
		req.Header.Set(RoutingKeyHeader, treq.RoutingKey)
	}
	if treq.RoutingDelegate != "" {
		req.Header.Set(RoutingDelegateHeader, treq.RoutingDelegate)
	}

	encoding := string(treq.Encoding)
	if encoding != "" {
		req.Header.Set(EncodingHeader, encoding)
	}

	return req
}

func (o *Outbound) getHTTPClient(p *hostport.Peer) (*http.Client, error) {
	agent, ok := p.Agent().(*Agent)
	if !ok {
		return nil, peer.ErrInvalidAgentConversion{
			Agent:        p.Agent(),
			ExpectedType: "*http.Agent",
		}
	}
	return agent.client, nil
}

func getErrFromResponse(response *http.Response) error {
	// TODO Behavior for 300-range status codes is undefined
	contents, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return err
	}

	if err := response.Body.Close(); err != nil {
		return err
	}

	// Trim the trailing newline from HTTP error messages
	message := strings.TrimSuffix(string(contents), "\n")

	if response.StatusCode >= 400 && response.StatusCode < 500 {
		return errors.RemoteBadRequestError(message)
	}

	if response.StatusCode == http.StatusGatewayTimeout {
		return errors.RemoteTimeoutError(message)
	}

	return errors.RemoteUnexpectedError(message)
}
<commit_msg>[Refactor] Consolidate duplicate code between http.CallOneway and http.Call<commit_after>// Copyright (c) 2016 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 http

import (
	"context"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"strings"
	"time"

	"go.uber.org/yarpc/internal/errors"
	"go.uber.org/yarpc/peer"
	"go.uber.org/yarpc/peer/hostport"
	"go.uber.org/yarpc/peer/single"
	"go.uber.org/yarpc/transport"

	"github.com/opentracing/opentracing-go"
	"github.com/opentracing/opentracing-go/ext"
	"go.uber.org/atomic"
)

var errOutboundNotStarted = errors.ErrOutboundNotStarted("http.Outbound")

// this ensures the HTTP outbound implements both transport.Outbound interfaces
var (
	_ transport.UnaryOutbound  = (*Outbound)(nil)
	_ transport.OnewayOutbound = (*Outbound)(nil)
)

var defaultURLTemplate, _ = url.Parse("http://localhost")

// NewOutbound builds a new HTTP outbound that sends requests to the given
// URL.
//
// Deprecated: create outbounds through NewPeerListOutbound instead
func NewOutbound(urlStr string, opts ...AgentOption) *Outbound {
	agent := NewAgent(opts...)

	_, hp := parseURL(urlStr)
	peerID := hostport.PeerIdentifier(hp)
	c := single.New(peerID, agent)

	err := c.Start()
	if err != nil {
		// This should never happen, single shouldn't return an error here
		panic(fmt.Sprintf("could not start single peerChooser, err: %s", err))
	}

	return NewChooserOutbound(c).WithURLTemplate(urlStr)
}

func parseURL(urlStr string) (*url.URL, string) {
	parsedURL, err := url.Parse(urlStr)
	if err != nil {
		panic(fmt.Sprintf("invalid url: %s, err: %s", urlStr, err))
	}

	return parsedURL, parsedURL.Host
}

// NewChooserOutbound builds a new HTTP outbound built around a PeerList
// for getting potential downstream hosts.
// Chooser.Choose MUST return *hostport.Peer objects.
// Chooser.Start MUST be called before Outbound.Start
func NewChooserOutbound(chooser peer.Chooser) *Outbound {
	return &Outbound{
		started:     atomic.NewBool(false),
		chooser:     chooser,
		urlTemplate: defaultURLTemplate,
		tracer:      opentracing.GlobalTracer(),
	}
}

// Outbound is an HTTP UnaryOutbound and OnewayOutbound
type Outbound struct {
	started     *atomic.Bool
	chooser     peer.Chooser
	urlTemplate *url.URL
	tracer      opentracing.Tracer
}

// WithURLTemplate configures an alternate URL template.
// The host:port portion of the URL template gets replaced by the chosen peer's
// identifier for each outbound request.
func (o *Outbound) WithURLTemplate(URL string) *Outbound {
	parsedURL, err := url.Parse(URL)
	if err != nil {
		log.Fatalf("failed to configure HTTP outbound: invalid URL template %q: %s", URL, err)
	}
	o.urlTemplate = parsedURL
	return o
}

// WithTracer configures a tracer for the outbound
func (o *Outbound) WithTracer(tracer opentracing.Tracer) *Outbound {
	o.tracer = tracer
	return o
}

// Start the HTTP outbound
func (o *Outbound) Start() error {
	o.started.Swap(true)
	return nil
}

// Stop the HTTP outbound
func (o *Outbound) Stop() error {
	o.started.Swap(false)
	return nil
}

// Call makes a HTTP request
func (o *Outbound) Call(ctx context.Context, treq *transport.Request) (*transport.Response, error) {
	if !o.started.Load() {
		// panic because there's no recovery from this
		panic(errOutboundNotStarted)
	}
	start := time.Now()
	deadline, _ := ctx.Deadline()
	ttl := deadline.Sub(start)

	return o.call(ctx, treq, start, ttl)
}

type ack struct {
	time time.Time
}

func (a ack) String() string {
	return a.time.String()
}

// CallOneway makes a oneway request
func (o *Outbound) CallOneway(ctx context.Context, treq *transport.Request) (transport.Ack, error) {
	if !o.started.Load() {
		// panic because there's no recovery from this
		panic(errOutboundNotStarted)
	}
	start := time.Now()
	var ttl time.Duration

	_, err := o.call(ctx, treq, start, ttl)
	if err != nil {
		return nil, err
	}

	return ack{time: time.Now()}, nil
}

func (o *Outbound) call(ctx context.Context, treq *transport.Request, start time.Time, ttl time.Duration) (*transport.Response, error) {
	p, err := o.getPeerForRequest(ctx, treq)
	if err != nil {
		return nil, err
	}
	p.StartRequest(nil)
	defer p.EndRequest(nil)

	req, err := o.createRequest(p, treq)
	if err != nil {
		return nil, err
	}

	req.Header = applicationHeaders.ToHTTPHeaders(treq.Headers, nil)
	ctx, req, span := o.withOpentracingSpan(ctx, req, treq, start)
	defer span.Finish()
	req = o.withCoreHeaders(req, treq, ttl)

	client, err := o.getHTTPClient(p)
	if err != nil {
		return nil, err
	}

	response, err := client.Do(req.WithContext(ctx))

	if err != nil {
		// Workaround borrowed from ctxhttp until
		// https://github.com/golang/go/issues/17711 is resolved.
		select {
		case <-ctx.Done():
			err = ctx.Err()
		default:
		}

		span.SetTag("error", true)
		span.LogEvent(err.Error())
		if err == context.DeadlineExceeded {
			end := time.Now()
			return nil, errors.ClientTimeoutError(treq.Service, treq.Procedure, end.Sub(start))
		}

		return nil, err
	}

	span.SetTag("http.status_code", response.StatusCode)

	if response.StatusCode >= 200 && response.StatusCode < 300 {
		appHeaders := applicationHeaders.FromHTTPHeaders(
			response.Header, transport.NewHeaders())
		return &transport.Response{
			Headers: appHeaders,
			Body:    response.Body,
		}, nil
	}

	return nil, getErrFromResponse(response)
}

func (o *Outbound) getPeerForRequest(ctx context.Context, treq *transport.Request) (*hostport.Peer, error) {
	p, err := o.chooser.Choose(ctx, treq)
	if err != nil {
		return nil, err
	}

	hpPeer, ok := p.(*hostport.Peer)
	if !ok {
		return nil, peer.ErrInvalidPeerConversion{
			Peer:         p,
			ExpectedType: "*hostport.Peer",
		}
	}

	return hpPeer, nil
}

func (o *Outbound) createRequest(p *hostport.Peer, treq *transport.Request) (*http.Request, error) {
	newURL := *o.urlTemplate
	newURL.Host = p.HostPort()
	return http.NewRequest("POST", newURL.String(), treq.Body)
}

func (o *Outbound) withOpentracingSpan(ctx context.Context, req *http.Request, treq *transport.Request, start time.Time) (context.Context, *http.Request, opentracing.Span) {
	// Apply HTTP Context headers for tracing and baggage carried by tracing.
	tracer := o.tracer
	var parent opentracing.SpanContext // ok to be nil
	if parentSpan := opentracing.SpanFromContext(ctx); parentSpan != nil {
		parent = parentSpan.Context()
	}
	span := tracer.StartSpan(
		treq.Procedure,
		opentracing.StartTime(start),
		opentracing.ChildOf(parent),
		opentracing.Tags{
			"rpc.caller":    treq.Caller,
			"rpc.service":   treq.Service,
			"rpc.encoding":  treq.Encoding,
			"rpc.transport": "http",
		},
	)
	ext.PeerService.Set(span, treq.Service)
	ext.SpanKindRPCClient.Set(span)
	ext.HTTPUrl.Set(span, req.URL.String())
	ctx = opentracing.ContextWithSpan(ctx, span)

	tracer.Inject(
		span.Context(),
		opentracing.HTTPHeaders,
		opentracing.HTTPHeadersCarrier(req.Header),
	)

	return ctx, req, span
}

func (o *Outbound) withCoreHeaders(req *http.Request, treq *transport.Request, ttl time.Duration) *http.Request {
	req.Header.Set(CallerHeader, treq.Caller)
	req.Header.Set(ServiceHeader, treq.Service)
	req.Header.Set(ProcedureHeader, treq.Procedure)
	if ttl != 0 {
		req.Header.Set(TTLMSHeader, fmt.Sprintf("%d", ttl/time.Millisecond))
	}
	if treq.ShardKey != "" {
		req.Header.Set(ShardKeyHeader, treq.ShardKey)
	}
	if treq.RoutingKey != "" {
		req.Header.Set(RoutingKeyHeader, treq.RoutingKey)
	}
	if treq.RoutingDelegate != "" {
		req.Header.Set(RoutingDelegateHeader, treq.RoutingDelegate)
	}

	encoding := string(treq.Encoding)
	if encoding != "" {
		req.Header.Set(EncodingHeader, encoding)
	}

	return req
}

func (o *Outbound) getHTTPClient(p *hostport.Peer) (*http.Client, error) {
	agent, ok := p.Agent().(*Agent)
	if !ok {
		return nil, peer.ErrInvalidAgentConversion{
			Agent:        p.Agent(),
			ExpectedType: "*http.Agent",
		}
	}
	return agent.client, nil
}

func getErrFromResponse(response *http.Response) error {
	// TODO Behavior for 300-range status codes is undefined
	contents, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return err
	}

	if err := response.Body.Close(); err != nil {
		return err
	}

	// Trim the trailing newline from HTTP error messages
	message := strings.TrimSuffix(string(contents), "\n")

	if response.StatusCode >= 400 && response.StatusCode < 500 {
		return errors.RemoteBadRequestError(message)
	}

	if response.StatusCode == http.StatusGatewayTimeout {
		return errors.RemoteTimeoutError(message)
	}

	return errors.RemoteUnexpectedError(message)
}
<|endoftext|>


<commit_before>package tsi1

import (
	"container/list"
	"sync"

	"github.com/influxdata/influxdb/tsdb"
	"github.com/prometheus/client_golang/prometheus"
)

// TagValueSeriesIDCache is an LRU cache for series id sets associated with
// name -> key -> value mappings. The purpose of the cache is to provide
// efficient means to get sets of series ids that would otherwise involve merging
// many individual bitmaps at query time.
//
// When initialising a TagValueSeriesIDCache a capacity must be provided. When
// more than c items are added to the cache, the least recently used item is
// evicted from the cache.
//
// A TagValueSeriesIDCache comprises a linked list implementation to track the
// order by which items should be evicted from the cache, and a hashmap implementation
// to provide constant time retrievals of items from the cache.
type TagValueSeriesIDCache struct {
	sync.RWMutex
	cache   map[string]map[string]map[string]*list.Element
	evictor *list.List

	tracker  *cacheTracker
	capacity uint64
}

// NewTagValueSeriesIDCache returns a TagValueSeriesIDCache with capacity c.
func NewTagValueSeriesIDCache(c uint64) *TagValueSeriesIDCache {
	return &TagValueSeriesIDCache{
		cache:    map[string]map[string]map[string]*list.Element{},
		evictor:  list.New(),
		tracker:  newCacheTracker(newCacheMetrics(nil), nil),
		capacity: c,
	}
}

// Get returns the SeriesIDSet associated with the {name, key, value} tuple if it
// exists.
func (c *TagValueSeriesIDCache) Get(name, key, value []byte) *tsdb.SeriesIDSet {
	c.Lock()
	defer c.Unlock()
	return c.get(name, key, value)
}

func (c *TagValueSeriesIDCache) get(name, key, value []byte) *tsdb.SeriesIDSet {
	if mmap, ok := c.cache[string(name)]; ok {
		if tkmap, ok := mmap[string(key)]; ok {
			if ele, ok := tkmap[string(value)]; ok {
				c.tracker.IncGetHit()
				c.evictor.MoveToFront(ele) // This now becomes most recently used.
				return ele.Value.(*seriesIDCacheElement).SeriesIDSet
			}
		}
	}
	c.tracker.IncGetMiss()
	return nil
}

// exists returns true if the an item exists for the tuple {name, key, value}.
func (c *TagValueSeriesIDCache) exists(name, key, value []byte) bool {
	if mmap, ok := c.cache[string(name)]; ok {
		if tkmap, ok := mmap[string(key)]; ok {
			_, ok := tkmap[string(value)]
			return ok
		}
	}
	return false
}

// addToSet adds x to the SeriesIDSet associated with the tuple {name, key, value}
// if it exists. This method takes a lock on the underlying SeriesIDSet.
//
// NB this does not count as an access on the set—therefore the set is not promoted
// within the LRU cache.
func (c *TagValueSeriesIDCache) addToSet(name, key, value []byte, x tsdb.SeriesID) {
	if mmap, ok := c.cache[string(name)]; ok {
		if tkmap, ok := mmap[string(key)]; ok {
			if ele, ok := tkmap[string(value)]; ok {
				ss := ele.Value.(*seriesIDCacheElement).SeriesIDSet
				if ss == nil {
					ele.Value.(*seriesIDCacheElement).SeriesIDSet = tsdb.NewSeriesIDSet(x)
					return
				}
				ele.Value.(*seriesIDCacheElement).SeriesIDSet.Add(x)
			}
		}
	}
}

// measurementContainsSets returns true if there are sets cached for the provided measurement.
func (c *TagValueSeriesIDCache) measurementContainsSets(name []byte) bool {
	_, ok := c.cache[string(name)]
	return ok
}

// Put adds the SeriesIDSet to the cache under the tuple {name, key, value}. If
// the cache is at its limit, then the least recently used item is evicted.
func (c *TagValueSeriesIDCache) Put(name, key, value []byte, ss *tsdb.SeriesIDSet) {
	c.Lock()
	// Check under the write lock if the relevant item is now in the cache.
	if c.exists(name, key, value) {
		c.Unlock()
		c.tracker.IncPutHit()
		return
	}
	defer c.Unlock()

	// Create list item, and add to the front of the eviction list.
	listElement := c.evictor.PushFront(&seriesIDCacheElement{
		name:        name,
		key:         key,
		value:       value,
		SeriesIDSet: ss,
	})

	// Add the listElement to the set of items.
	if mmap, ok := c.cache[string(name)]; ok {
		if tkmap, ok := mmap[string(key)]; ok {
			if _, ok := tkmap[string(value)]; ok {
				goto EVICT
			}

			// Add the set to the map
			tkmap[string(value)] = listElement
			goto EVICT
		}

		// No series set map for the tag key - first tag value for the tag key.
		mmap[string(key)] = map[string]*list.Element{string(value): listElement}
		goto EVICT
	}

	// No map for the measurement - first tag key for the measurment.
	c.cache[string(name)] = map[string]map[string]*list.Element{
		string(key): {string(value): listElement},
	}

EVICT:
	c.checkEviction()
	c.tracker.IncPutMiss()
}

// Delete removes x from the tuple {name, key, value} if it exists.
// This method takes a lock on the underlying SeriesIDSet.
func (c *TagValueSeriesIDCache) Delete(name, key, value []byte, x tsdb.SeriesID) {
	c.Lock()
	c.delete(name, key, value, x)
	c.Unlock()
}

// DeleteMeasurement removes all cached entries for the provided measurement name.
func (c *TagValueSeriesIDCache) DeleteMeasurement(name []byte) {
	c.Lock()
	delete(c.cache, string(name))
	c.Unlock()
}

// delete removes x from the tuple {name, key, value} if it exists.
func (c *TagValueSeriesIDCache) delete(name, key, value []byte, x tsdb.SeriesID) {
	if mmap, ok := c.cache[string(name)]; ok {
		if tkmap, ok := mmap[string(key)]; ok {
			if ele, ok := tkmap[string(value)]; ok {
				if ss := ele.Value.(*seriesIDCacheElement).SeriesIDSet; ss != nil {
					ele.Value.(*seriesIDCacheElement).SeriesIDSet.Remove(x)
					c.tracker.IncDeletesHit()
					return
				}
			}
		}
	}
	c.tracker.IncDeletesMiss()
}

// checkEviction checks if the cache is too big, and evicts the least recently used
// item if it is.
func (c *TagValueSeriesIDCache) checkEviction() {
	l := uint64(c.evictor.Len())
	c.tracker.SetSize(l)
	if l <= c.capacity {
		return
	}

	e := c.evictor.Back() // Least recently used item.
	listElement := e.Value.(*seriesIDCacheElement)
	name := listElement.name
	key := listElement.key
	value := listElement.value

	c.evictor.Remove(e)                                       // Remove from evictor
	delete(c.cache[string(name)][string(key)], string(value)) // Remove from hashmap of items.

	// Check if there are no more tag values for the tag key.
	if len(c.cache[string(name)][string(key)]) == 0 {
		delete(c.cache[string(name)], string(key))
	}

	// Check there are no more tag keys for the measurement.
	if len(c.cache[string(name)]) == 0 {
		delete(c.cache, string(name))
	}
	c.tracker.IncEvictions()
}

func (c *TagValueSeriesIDCache) PrometheusCollectors() []prometheus.Collector {
	var collectors []prometheus.Collector
	collectors = append(collectors, c.tracker.metrics.PrometheusCollectors()...)
	return collectors
}

// seriesIDCacheElement is an item stored within a cache.
type seriesIDCacheElement struct {
	name        []byte
	key         []byte
	value       []byte
	SeriesIDSet *tsdb.SeriesIDSet
}

type cacheTracker struct {
	metrics *cacheMetrics
	labels  prometheus.Labels
	enabled bool
}

func newCacheTracker(metrics *cacheMetrics, defaultLabels prometheus.Labels) *cacheTracker {
	return &cacheTracker{metrics: metrics, labels: defaultLabels, enabled: true}
}

// Labels returns a copy of labels for use with index cache metrics.
func (t *cacheTracker) Labels() prometheus.Labels {
	l := make(map[string]string, len(t.labels))
	for k, v := range t.labels {
		l[k] = v
	}
	return l
}

func (t *cacheTracker) SetSize(sz uint64) {
	if !t.enabled {
		return
	}

	labels := t.Labels()
	t.metrics.Size.With(labels).Set(float64(sz))
}

func (t *cacheTracker) incGet(status string) {
	if !t.enabled {
		return
	}

	labels := t.Labels()
	labels["status"] = status
	t.metrics.Gets.With(labels).Inc()
}

func (t *cacheTracker) IncGetHit()  { t.incGet("hit") }
func (t *cacheTracker) IncGetMiss() { t.incGet("miss") }

func (t *cacheTracker) incPut(status string) {
	if !t.enabled {
		return
	}

	labels := t.Labels()
	labels["status"] = status
	t.metrics.Puts.With(labels).Inc()
}

func (t *cacheTracker) IncPutHit()  { t.incPut("hit") }
func (t *cacheTracker) IncPutMiss() { t.incPut("miss") }

func (t *cacheTracker) incDeletes(status string) {
	if !t.enabled {
		return
	}

	labels := t.Labels()
	labels["status"] = status
	t.metrics.Deletes.With(labels).Inc()
}

func (t *cacheTracker) IncDeletesHit()  { t.incDeletes("hit") }
func (t *cacheTracker) IncDeletesMiss() { t.incDeletes("miss") }

func (t *cacheTracker) IncEvictions() {
	if !t.enabled {
		return
	}

	labels := t.Labels()
	t.metrics.Evictions.With(labels).Inc()
}
<commit_msg>Convert TagValueSeriesIDCache to use string fields<commit_after>package tsi1

import (
	"container/list"
	"sync"

	"github.com/influxdata/influxdb/tsdb"
	"github.com/prometheus/client_golang/prometheus"
)

// TagValueSeriesIDCache is an LRU cache for series id sets associated with
// name -> key -> value mappings. The purpose of the cache is to provide
// efficient means to get sets of series ids that would otherwise involve merging
// many individual bitmaps at query time.
//
// When initialising a TagValueSeriesIDCache a capacity must be provided. When
// more than c items are added to the cache, the least recently used item is
// evicted from the cache.
//
// A TagValueSeriesIDCache comprises a linked list implementation to track the
// order by which items should be evicted from the cache, and a hashmap implementation
// to provide constant time retrievals of items from the cache.
type TagValueSeriesIDCache struct {
	sync.RWMutex
	cache   map[string]map[string]map[string]*list.Element
	evictor *list.List

	tracker  *cacheTracker
	capacity uint64
}

// NewTagValueSeriesIDCache returns a TagValueSeriesIDCache with capacity c.
func NewTagValueSeriesIDCache(c uint64) *TagValueSeriesIDCache {
	return &TagValueSeriesIDCache{
		cache:    map[string]map[string]map[string]*list.Element{},
		evictor:  list.New(),
		tracker:  newCacheTracker(newCacheMetrics(nil), nil),
		capacity: c,
	}
}

// Get returns the SeriesIDSet associated with the {name, key, value} tuple if it
// exists.
func (c *TagValueSeriesIDCache) Get(name, key, value []byte) *tsdb.SeriesIDSet {
	c.Lock()
	defer c.Unlock()
	return c.get(name, key, value)
}

func (c *TagValueSeriesIDCache) get(name, key, value []byte) *tsdb.SeriesIDSet {
	if mmap, ok := c.cache[string(name)]; ok {
		if tkmap, ok := mmap[string(key)]; ok {
			if ele, ok := tkmap[string(value)]; ok {
				c.tracker.IncGetHit()
				c.evictor.MoveToFront(ele) // This now becomes most recently used.
				return ele.Value.(*seriesIDCacheElement).SeriesIDSet
			}
		}
	}
	c.tracker.IncGetMiss()
	return nil
}

// exists returns true if the an item exists for the tuple {name, key, value}.
func (c *TagValueSeriesIDCache) exists(name, key, value []byte) bool {
	if mmap, ok := c.cache[string(name)]; ok {
		if tkmap, ok := mmap[string(key)]; ok {
			_, ok := tkmap[string(value)]
			return ok
		}
	}
	return false
}

// addToSet adds x to the SeriesIDSet associated with the tuple {name, key, value}
// if it exists. This method takes a lock on the underlying SeriesIDSet.
//
// NB this does not count as an access on the set—therefore the set is not promoted
// within the LRU cache.
func (c *TagValueSeriesIDCache) addToSet(name, key, value []byte, x tsdb.SeriesID) {
	if mmap, ok := c.cache[string(name)]; ok {
		if tkmap, ok := mmap[string(key)]; ok {
			if ele, ok := tkmap[string(value)]; ok {
				ss := ele.Value.(*seriesIDCacheElement).SeriesIDSet
				if ss == nil {
					ele.Value.(*seriesIDCacheElement).SeriesIDSet = tsdb.NewSeriesIDSet(x)
					return
				}
				ele.Value.(*seriesIDCacheElement).SeriesIDSet.Add(x)
			}
		}
	}
}

// measurementContainsSets returns true if there are sets cached for the provided measurement.
func (c *TagValueSeriesIDCache) measurementContainsSets(name []byte) bool {
	_, ok := c.cache[string(name)]
	return ok
}

// Put adds the SeriesIDSet to the cache under the tuple {name, key, value}. If
// the cache is at its limit, then the least recently used item is evicted.
func (c *TagValueSeriesIDCache) Put(name, key, value []byte, ss *tsdb.SeriesIDSet) {
	c.Lock()
	// Check under the write lock if the relevant item is now in the cache.
	if c.exists(name, key, value) {
		c.Unlock()
		c.tracker.IncPutHit()
		return
	}
	defer c.Unlock()

	// Create list item, and add to the front of the eviction list.
	listElement := c.evictor.PushFront(&seriesIDCacheElement{
		name:        string(name),
		key:         string(key),
		value:       string(value),
		SeriesIDSet: ss,
	})

	// Add the listElement to the set of items.
	if mmap, ok := c.cache[string(name)]; ok {
		if tkmap, ok := mmap[string(key)]; ok {
			if _, ok := tkmap[string(value)]; ok {
				goto EVICT
			}

			// Add the set to the map
			tkmap[string(value)] = listElement
			goto EVICT
		}

		// No series set map for the tag key - first tag value for the tag key.
		mmap[string(key)] = map[string]*list.Element{string(value): listElement}
		goto EVICT
	}

	// No map for the measurement - first tag key for the measurment.
	c.cache[string(name)] = map[string]map[string]*list.Element{
		string(key): {string(value): listElement},
	}

EVICT:
	c.checkEviction()
	c.tracker.IncPutMiss()
}

// Delete removes x from the tuple {name, key, value} if it exists.
// This method takes a lock on the underlying SeriesIDSet.
func (c *TagValueSeriesIDCache) Delete(name, key, value []byte, x tsdb.SeriesID) {
	c.Lock()
	c.delete(name, key, value, x)
	c.Unlock()
}

// DeleteMeasurement removes all cached entries for the provided measurement name.
func (c *TagValueSeriesIDCache) DeleteMeasurement(name []byte) {
	c.Lock()
	delete(c.cache, string(name))
	c.Unlock()
}

// delete removes x from the tuple {name, key, value} if it exists.
func (c *TagValueSeriesIDCache) delete(name, key, value []byte, x tsdb.SeriesID) {
	if mmap, ok := c.cache[string(name)]; ok {
		if tkmap, ok := mmap[string(key)]; ok {
			if ele, ok := tkmap[string(value)]; ok {
				if ss := ele.Value.(*seriesIDCacheElement).SeriesIDSet; ss != nil {
					ele.Value.(*seriesIDCacheElement).SeriesIDSet.Remove(x)
					c.tracker.IncDeletesHit()
					return
				}
			}
		}
	}
	c.tracker.IncDeletesMiss()
}

// checkEviction checks if the cache is too big, and evicts the least recently used
// item if it is.
func (c *TagValueSeriesIDCache) checkEviction() {
	l := uint64(c.evictor.Len())
	c.tracker.SetSize(l)
	if l <= c.capacity {
		return
	}

	e := c.evictor.Back() // Least recently used item.
	listElement := e.Value.(*seriesIDCacheElement)
	name := listElement.name
	key := listElement.key
	value := listElement.value

	c.evictor.Remove(e)                                       // Remove from evictor
	delete(c.cache[string(name)][string(key)], string(value)) // Remove from hashmap of items.

	// Check if there are no more tag values for the tag key.
	if len(c.cache[string(name)][string(key)]) == 0 {
		delete(c.cache[string(name)], string(key))
	}

	// Check there are no more tag keys for the measurement.
	if len(c.cache[string(name)]) == 0 {
		delete(c.cache, string(name))
	}
	c.tracker.IncEvictions()
}

func (c *TagValueSeriesIDCache) PrometheusCollectors() []prometheus.Collector {
	var collectors []prometheus.Collector
	collectors = append(collectors, c.tracker.metrics.PrometheusCollectors()...)
	return collectors
}

// seriesIDCacheElement is an item stored within a cache.
type seriesIDCacheElement struct {
	name        string
	key         string
	value       string
	SeriesIDSet *tsdb.SeriesIDSet
}

type cacheTracker struct {
	metrics *cacheMetrics
	labels  prometheus.Labels
	enabled bool
}

func newCacheTracker(metrics *cacheMetrics, defaultLabels prometheus.Labels) *cacheTracker {
	return &cacheTracker{metrics: metrics, labels: defaultLabels, enabled: true}
}

// Labels returns a copy of labels for use with index cache metrics.
func (t *cacheTracker) Labels() prometheus.Labels {
	l := make(map[string]string, len(t.labels))
	for k, v := range t.labels {
		l[k] = v
	}
	return l
}

func (t *cacheTracker) SetSize(sz uint64) {
	if !t.enabled {
		return
	}

	labels := t.Labels()
	t.metrics.Size.With(labels).Set(float64(sz))
}

func (t *cacheTracker) incGet(status string) {
	if !t.enabled {
		return
	}

	labels := t.Labels()
	labels["status"] = status
	t.metrics.Gets.With(labels).Inc()
}

func (t *cacheTracker) IncGetHit()  { t.incGet("hit") }
func (t *cacheTracker) IncGetMiss() { t.incGet("miss") }

func (t *cacheTracker) incPut(status string) {
	if !t.enabled {
		return
	}

	labels := t.Labels()
	labels["status"] = status
	t.metrics.Puts.With(labels).Inc()
}

func (t *cacheTracker) IncPutHit()  { t.incPut("hit") }
func (t *cacheTracker) IncPutMiss() { t.incPut("miss") }

func (t *cacheTracker) incDeletes(status string) {
	if !t.enabled {
		return
	}

	labels := t.Labels()
	labels["status"] = status
	t.metrics.Deletes.With(labels).Inc()
}

func (t *cacheTracker) IncDeletesHit()  { t.incDeletes("hit") }
func (t *cacheTracker) IncDeletesMiss() { t.incDeletes("miss") }

func (t *cacheTracker) IncEvictions() {
	if !t.enabled {
		return
	}

	labels := t.Labels()
	t.metrics.Evictions.With(labels).Inc()
}
<|endoftext|>


<commit_before>package log

import (
	// TODO: add a file logger
	"log"
	"sync"
	"testing"
)

func init() {
	level = WarningLevel
	logger = BaseLogger{}
}

// Level of logging trigger
type Level int

// Available logging levels
const (
	DebugLevel Level = iota
	InfoLevel
	NoticeLevel
	WarningLevel
	CriticalLevel
)

var (
	logger Logger
	level  Level
	mu     sync.Mutex
)

// Logger defines the logs levels used by RamSQL engine
type Logger interface {
	Logf(fmt string, values ...interface{})
}

// SetLevel controls the categories of logs written
func SetLevel(lvl Level) {
	mu.Lock()
	level = lvl
	mu.Unlock()
}

func lvl() Level {
	mu.Lock()
	defer mu.Unlock()
	return level
}

// Debug prints debug log
func Debug(format string, values ...interface{}) {
	if lvl() <= DebugLevel {
		logger.Logf("[DEBUG]    "+format, values...)
	}
}

// Info prints information log
func Info(format string, values ...interface{}) {
	if lvl() <= InfoLevel {
		logger.Logf("[INFO]     "+format, values...)
	}
}

// Notice prints information that should be seen
func Notice(format string, values ...interface{}) {
	if lvl() <= NoticeLevel {
		logger.Logf("[NOTICE]   "+format, values...)
	}
}

// Warning prints warnings for user
func Warning(format string, values ...interface{}) {
	if lvl() <= WarningLevel {
		logger.Logf("[WARNING]  "+format, values...)
	}
}

// Critical prints error informations
func Critical(format string, values ...interface{}) {
	mu.Lock()
	logger.Logf("[CRITICAL] "+format, values...)
	mu.Unlock()
}

// BaseLogger logs on stdout
type BaseLogger struct {
}

// Logf logs on stdout
func (l BaseLogger) Logf(fmt string, values ...interface{}) {
	log.Printf(fmt, values...)
}

// TestLogger uses *testing.T as a backend for RamSQL logs
type TestLogger struct {
	t *testing.T
}

// Logf logs in testing log buffer
func (l TestLogger) Logf(fmt string, values ...interface{}) {
	l.t.Logf(fmt, values...)
}

// UseTestLogger should be used only by unit tests
func UseTestLogger(t *testing.T) {
	mu.Lock()
	logger = t
	mu.Unlock()
	SetLevel(WarningLevel)
}
<commit_msg>feat (log): allow use of testing.B<commit_after>package log

import (
	// TODO: add a file logger
	"log"
	"sync"
	"testing"
)

func init() {
	level = WarningLevel
	logger = BaseLogger{}
}

// Level of logging trigger
type Level int

// Available logging levels
const (
	DebugLevel Level = iota
	InfoLevel
	NoticeLevel
	WarningLevel
	CriticalLevel
)

var (
	logger Logger
	level  Level
	mu     sync.Mutex
)

// Logger defines the logs levels used by RamSQL engine
type Logger interface {
	Logf(fmt string, values ...interface{})
}

// SetLevel controls the categories of logs written
func SetLevel(lvl Level) {
	mu.Lock()
	level = lvl
	mu.Unlock()
}

func lvl() Level {
	mu.Lock()
	defer mu.Unlock()
	return level
}

// Debug prints debug log
func Debug(format string, values ...interface{}) {
	if lvl() <= DebugLevel {
		logger.Logf("[DEBUG]    "+format, values...)
	}
}

// Info prints information log
func Info(format string, values ...interface{}) {
	if lvl() <= InfoLevel {
		logger.Logf("[INFO]     "+format, values...)
	}
}

// Notice prints information that should be seen
func Notice(format string, values ...interface{}) {
	if lvl() <= NoticeLevel {
		logger.Logf("[NOTICE]   "+format, values...)
	}
}

// Warning prints warnings for user
func Warning(format string, values ...interface{}) {
	if lvl() <= WarningLevel {
		logger.Logf("[WARNING]  "+format, values...)
	}
}

// Critical prints error informations
func Critical(format string, values ...interface{}) {
	mu.Lock()
	logger.Logf("[CRITICAL] "+format, values...)
	mu.Unlock()
}

// BaseLogger logs on stdout
type BaseLogger struct {
}

// Logf logs on stdout
func (l BaseLogger) Logf(fmt string, values ...interface{}) {
	log.Printf(fmt, values...)
}

// TestLogger uses *testing.T as a backend for RamSQL logs
type TestLogger struct {
	t *testing.T
}

// Logf logs in testing log buffer
func (l TestLogger) Logf(fmt string, values ...interface{}) {
	l.t.Logf(fmt, values...)
}

// UseTestLogger should be used only by unit tests
func UseTestLogger(t testing.TB) {
	mu.Lock()
	logger = t
	mu.Unlock()
	SetLevel(WarningLevel)
}
<|endoftext|>


<commit_before>package cloudpelican

// @author Robin Verlangen
// Tool for logging data to CloudPelican directly from Go

// Imports
import (
    "net"
    "net/http"
    "net/url"
    "log"
    "sync"
    "time"
)

// Settings
var ENDPOINT string = "https://app.cloudpelican.com/api/push/pixel"
var TOKEN string = ""
var backendTimeout = time.Duration(5 * time.Second)
var debugMode = false

// Monitor drain status
var routineQuit chan int = make(chan int)
var startCounter uint64 = uint64(0)
var startCounterMux sync.Mutex
var doneCounter uint64 = uint64(0)
var doneCounterMux sync.Mutex

// Log queue
var writeAheadBufferSize int = 1000
var writeAhead chan string = make(chan string, writeAheadBufferSize)
var writeAheadInit bool
var dropOnFullWriteAheadBuffer bool = true

// Set token
func SetToken(t string) {
    // Validate before setting
    validateToken(t)

    // Store
    TOKEN = t
}

// Set endpoint
func SetEndpoint(e string) {
    // Store
    ENDPOINT = e
}

// Set timeout
func SetBackendTimeout(to time.Duration) {
    backendTimeout = to
}

// Debug
func SetDebugMode(b bool) {
    debugMode = b
}

// Write a message
func LogMessageWithToken(t string, msg string) bool {
    // Create fields map
    var fields map[string]string = make(map[string]string)
    fields["msg"] = msg

    // Push to channel
    return requestAsync(assembleUrl(t, fields))
}

// Write a message
func LogMessage(msg string) bool {
    // Create fields map
    var fields map[string]string = make(map[string]string)
    fields["msg"] = msg

    // Push to channel
    return requestAsync(assembleUrl(TOKEN, fields))
}

// Assemble url
// @return string Url based on the input fields
func assembleUrl(t string, fields map[string]string) string {
    // Token check
    validateToken(t)

    // Baisc query params
    params := url.Values{}
    params.Add("t", t)

    // Fields
    for k, _ := range fields {
        if len(k) == 0 || len(fields[k]) == 0 {
            log.Printf("Skipping invalid field %s with value %s", k, fields[k])
            continue
        }
        params.Add("f[" + k + "]", fields[k])
    }

    // Final url
    return ENDPOINT + "?" + params.Encode()
}

// Request async
func requestAsync(url string) bool {
    // Check amount of open items in the channel, if the channel is full, return false and drop this message
    if dropOnFullWriteAheadBuffer {
        var lwa int = len(writeAhead)
        if lwa == writeAheadBufferSize {
            log.Printf("Write ahead buffer is full and contains %d items. Dropping current log message", lwa)
        }
    }

    // Add counter
    startCounterMux.Lock()
    startCounter++
    startCounterMux.Unlock()

    // Do we have to start a writer?
    if writeAheadInit == false {
        writeAheadInit = true
        backendWriter()
    }

    // Insert into channel
    writeAhead <- url

    // OK
    return true
}

// Backend writer
func backendWriter() {
    go func() {
        // Wait for messages
        for {
            // Read from channel
            var url string
            url = <- writeAhead

            // Client
            transport := &http.Transport{
                Dial: func(netw, addr string) (net.Conn, error) {
                        // we want to wait a maximum of 1.75 seconds...
                        // since we're specifying a 1 second connect timeout and deadline
                        // (read/write timeout) is specified in absolute time we want to
                        // calculate that time first (before connecting)
                        deadline := time.Now().Add(backendTimeout)
                        c, err := net.DialTimeout(netw, addr, time.Second)
                        if err != nil {
                                return nil, err
                        }
                        c.SetDeadline(deadline)
                        return c, nil
                }}
            httpclient := &http.Client{Transport: transport}

            // Make request
            if debugMode {
                log.Printf("Write ahead queue %d\n", len(writeAhead))
                log.Println(url)
            }
            resp, err := httpclient.Get(url)
            if resp != nil && resp.Body != nil {
                defer resp.Body.Close()
            }
            if err != nil {
                log.Printf("Error while forwarding data: %s\n", err)
            }

            // Done counter
            doneCounterMux.Lock()
            doneCounter++
            doneCounterMux.Unlock()
        }
        log.Printf("Stopping backend writer")
    }()
}

// Timeout helper
func dialTimeout(network, addr string) (net.Conn, error) {
    return net.DialTimeout(network, addr, backendTimeout)
}

// Validate the token
func validateToken(t string) {
    if len(t) == 0 {
        log.Println("Please set a valid token with cloudpelican.SetToken(token string)")
    }
}<commit_msg>Bugfix<commit_after>package cloudpelican

// @author Robin Verlangen
// Tool for logging data to CloudPelican directly from Go

// Imports
import (
    "net"
    "net/http"
    "net/url"
    "log"
    "sync"
    "time"
)

// Settings
var ENDPOINT string = "https://app.cloudpelican.com/api/push/pixel"
var TOKEN string = ""
var backendTimeout = time.Duration(5 * time.Second)
var debugMode = false

// Monitor drain status
var routineQuit chan int = make(chan int)
var startCounter uint64 = uint64(0)
var startCounterMux sync.Mutex
var doneCounter uint64 = uint64(0)
var doneCounterMux sync.Mutex

// Log queue
var writeAheadBufferSize int = 1000
var writeAhead chan string = make(chan string, writeAheadBufferSize)
var writeAheadInit bool
var dropOnFullWriteAheadBuffer bool = true

// Set token
func SetToken(t string) {
    // Validate before setting
    validateToken(t)

    // Store
    TOKEN = t
}

// Set endpoint
func SetEndpoint(e string) {
    // Store
    ENDPOINT = e
}

// Set timeout
func SetBackendTimeout(to time.Duration) {
    backendTimeout = to
}

// Debug
func SetDebugMode(b bool) {
    debugMode = b
}

// Write a message
func LogMessageWithToken(t string, msg string) bool {
    // Create fields map
    var fields map[string]string = make(map[string]string)
    fields["msg"] = msg

    // Push to channel
    return requestAsync(assembleUrl(t, fields))
}

// Write a message
func LogMessage(msg string) bool {
    // Create fields map
    var fields map[string]string = make(map[string]string)
    fields["msg"] = msg

    // Push to channel
    return requestAsync(assembleUrl(TOKEN, fields))
}

// Assemble url
// @return string Url based on the input fields
func assembleUrl(t string, fields map[string]string) string {
    // Token check
    validateToken(t)

    // Baisc query params
    params := url.Values{}
    params.Add("t", t)

    // Fields
    for k, _ := range fields {
        if len(k) == 0 || len(fields[k]) == 0 {
            log.Printf("Skipping invalid field %s with value %s", k, fields[k])
            continue
        }
        params.Add("f[" + k + "]", fields[k])
    }

    // Final url
    return ENDPOINT + "?" + params.Encode()
}

// Request async
func requestAsync(url string) bool {
    // Check amount of open items in the channel, if the channel is full, return false and drop this message
    if dropOnFullWriteAheadBuffer {
        var lwa int = len(writeAhead)
        if lwa == writeAheadBufferSize {
            log.Printf("Write ahead buffer is full and contains %d items. Dropping current log message", lwa)
        }
    }

    // Add counter
    startCounterMux.Lock()
    startCounter++
    startCounterMux.Unlock()

    // Do we have to start a writer?
    if writeAheadInit == false {
        writeAheadInit = true
        backendWriter()
    }

    // Insert into channel
    writeAhead <- url

    // OK
    return true
}

// Backend writer
func backendWriter() {
    go func() {
        // Client
        transport := &http.Transport{
            Dial: func(netw, addr string) (net.Conn, error) {
                    // we want to wait a maximum of 1.75 seconds...
                    // since we're specifying a 1 second connect timeout and deadline
                    // (read/write timeout) is specified in absolute time we want to
                    // calculate that time first (before connecting)
                    deadline := time.Now().Add(backendTimeout)
                    c, err := net.DialTimeout(netw, addr, time.Second)
                    if err != nil {
                            return nil, err
                    }
                    c.SetDeadline(deadline)
                    return c, nil
            }}
        httpclient := &http.Client{Transport: transport}

        // Wait for messages
        for {
            // Read from channel
            var url string
            url = <- writeAhead

            // Make request
            if debugMode {
                log.Printf("Write ahead queue %d\n", len(writeAhead))
                log.Println(url)
            }
            resp, err := httpclient.Get(url)
            if err != nil {
                log.Printf("Error while forwarding data: %s\n", err)
            } else {
                defer resp.Body.Close()
            }

            // Done counter
            doneCounterMux.Lock()
            doneCounter++
            doneCounterMux.Unlock()
        }
        log.Printf("Stopping backend writer")
    }()
}

// Timeout helper
func dialTimeout(network, addr string) (net.Conn, error) {
    return net.DialTimeout(network, addr, backendTimeout)
}

// Validate the token
func validateToken(t string) {
    if len(t) == 0 {
        log.Println("Please set a valid token with cloudpelican.SetToken(token string)")
    }
}<|endoftext|>


<commit_before>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package cmd

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"strings"

	"github.com/TheThingsNetwork/ttn/ttnctl/util"
	"github.com/apex/log"
	"github.com/gosuri/uitable"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

type app struct {
	EUI        string   `json:"eui"`
	Name       string   `json:"name"`
	Owner      string   `json:"owner"`
	AccessKeys []string `json:"accessKeys"`
	Valid      bool     `json:"valid"`
}

// applicationsCmd represents the applications command
var applicationsCmd = &cobra.Command{
	Use:   "applications",
	Short: "Show applications",
	Long:  `ttnctl applications retrieves the applications of the logged on user.`,
	Run: func(cmd *cobra.Command, args []string) {
		server := viper.GetString("ttn-account-server")
		req, err := util.NewRequestWithAuth(server, "GET", fmt.Sprintf("%s/applications", server), nil)
		if err != nil {
			ctx.WithError(err).Fatal("Failed to create authenticated request")
		}

		client := &http.Client{}
		resp, err := client.Do(req)
		if err != nil {
			ctx.WithError(err).Fatal("Failed to get applications")
		}
		if resp.StatusCode != http.StatusOK {
			ctx.Fatalf("Failed to get applications: %s", resp.Status)
		}

		defer resp.Body.Close()
		decoder := json.NewDecoder(resp.Body)
		var apps []*app
		err = decoder.Decode(&apps)
		if err != nil {
			ctx.WithError(err).Fatal("Failed to read applications")
		}

		ctx.Infof("Found %d application(s)", len(apps))
		table := uitable.New()
		table.MaxColWidth = 70
		table.AddRow("EUI", "Name", "Owner", "Access Keys", "Valid")
		for _, app := range apps {
			table.AddRow(app.EUI, app.Name, app.Owner, strings.Join(app.AccessKeys, ", "), app.Valid)
		}

		fmt.Println(table)
	},
}

var applicationsCreateCmd = &cobra.Command{
	Use:   "create [eui] [name]",
	Short: "Create a new application",
	Long:  `ttnctl applications create creates a new application.`,
	Run: func(cmd *cobra.Command, args []string) {
		if len(args) < 2 {
			cmd.Help()
			return
		}

		appEUI, err := util.Parse64(args[0])
		if err != nil {
			ctx.Fatalf("Invalid AppEUI: %s", err)
		}

		server := viper.GetString("ttn-account-server")
		values := url.Values{
			"eui":  {fmt.Sprintf("%X", appEUI)},
			"name": {args[1]},
		}
		req, err := util.NewRequestWithAuth(server, "POST", fmt.Sprintf("%s/applications", server), strings.NewReader(values.Encode()))
		if err != nil {
			ctx.WithError(err).Fatal("Failed to create authenticated request")
		}

		req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

		client := &http.Client{}
		resp, err := client.Do(req)
		if err != nil {
			ctx.WithError(err).Fatal("Failed to create application")
		}
		if resp.StatusCode != http.StatusCreated {
			ctx.Fatalf("Failed to create application: %s", resp.Status)
		}

		ctx.Info("Application created successfully")

		// We need to refresh the token to add the new application to the set of
		// claims
		_, err = util.RefreshToken(server)
		if err != nil {
			log.WithError(err).Warn("Failed to refresh token. Please login")
		}
	},
}

var applicationsDeleteCmd = &cobra.Command{
	Use:   "delete [eui]",
	Short: "Delete an application",
	Long:  `ttnctl application delete deletes an existing application.`,
	Run: func(cmd *cobra.Command, args []string) {
		if len(args) < 1 {
			cmd.Help()
			return
		}

		appEUI, err := util.Parse64(args[0])
		if err != nil {
			ctx.Fatalf("Invalid AppEUI: %s", err)
		}

		server := viper.GetString("ttn-account-server")
		req, err := util.NewRequestWithAuth(server, "DELETE", fmt.Sprintf("%s/applications/%s", server, fmt.Sprintf("%X", appEUI)), nil)
		if err != nil {
			ctx.WithError(err).Fatal("Failed to create authenticated request")
		}

		client := &http.Client{}
		resp, err := client.Do(req)
		if err != nil {
			ctx.WithError(err).Fatal("Failed to delete application")
		}
		if resp.StatusCode != http.StatusOK {
			ctx.Fatalf("Failed to delete application: %s", resp.Status)
		}

		ctx.Info("Application deleted successfully")

		// We need to refresh the token to remove the application from the set of
		// claims
		_, err = util.RefreshToken(server)
		if err != nil {
			log.WithError(err).Warn("Failed to refresh token. Please login")
		}
	},
}

func init() {
	RootCmd.AddCommand(applicationsCmd)
	applicationsCmd.AddCommand(applicationsCreateCmd)
	applicationsCmd.AddCommand(applicationsDeleteCmd)
}
<commit_msg>Authorize other user for application<commit_after>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package cmd

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"strings"

	"github.com/TheThingsNetwork/ttn/ttnctl/util"
	"github.com/apex/log"
	"github.com/gosuri/uitable"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

type app struct {
	EUI        string   `json:"eui"`
	Name       string   `json:"name"`
	Owner      string   `json:"owner"`
	AccessKeys []string `json:"accessKeys"`
	Valid      bool     `json:"valid"`
}

// applicationsCmd represents the applications command
var applicationsCmd = &cobra.Command{
	Use:   "applications",
	Short: "Show applications",
	Long:  `ttnctl applications retrieves the applications of the logged on user.`,
	Run: func(cmd *cobra.Command, args []string) {
		server := viper.GetString("ttn-account-server")
		uri := fmt.Sprintf("%s/applications", server)
		req, err := util.NewRequestWithAuth(server, "GET", uri, nil)
		if err != nil {
			ctx.WithError(err).Fatal("Failed to create authenticated request")
		}

		client := &http.Client{}
		resp, err := client.Do(req)
		if err != nil {
			ctx.WithError(err).Fatal("Failed to get applications")
		}
		if resp.StatusCode != http.StatusOK {
			ctx.Fatalf("Failed to get applications: %s", resp.Status)
		}

		defer resp.Body.Close()
		decoder := json.NewDecoder(resp.Body)
		var apps []*app
		err = decoder.Decode(&apps)
		if err != nil {
			ctx.WithError(err).Fatal("Failed to read applications")
		}

		ctx.Infof("Found %d application(s)", len(apps))
		table := uitable.New()
		table.MaxColWidth = 70
		table.AddRow("EUI", "Name", "Owner", "Access Keys", "Valid")
		for _, app := range apps {
			table.AddRow(app.EUI, app.Name, app.Owner, strings.Join(app.AccessKeys, ", "), app.Valid)
		}

		fmt.Println(table)
	},
}

var applicationsCreateCmd = &cobra.Command{
	Use:   "create [eui] [name]",
	Short: "Create a new application",
	Long:  `ttnctl applications create creates a new application.`,
	Run: func(cmd *cobra.Command, args []string) {
		if len(args) < 2 {
			cmd.Help()
			return
		}

		appEUI, err := util.Parse64(args[0])
		if err != nil {
			ctx.Fatalf("Invalid AppEUI: %s", err)
		}

		server := viper.GetString("ttn-account-server")
		uri := fmt.Sprintf("%s/applications", server)
		values := url.Values{
			"eui":  {fmt.Sprintf("%X", appEUI)},
			"name": {args[1]},
		}
		req, err := util.NewRequestWithAuth(server, "POST", uri, strings.NewReader(values.Encode()))
		if err != nil {
			ctx.WithError(err).Fatal("Failed to create authenticated request")
		}

		req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

		client := &http.Client{}
		resp, err := client.Do(req)
		if err != nil {
			ctx.WithError(err).Fatal("Failed to create application")
		}
		if resp.StatusCode != http.StatusCreated {
			ctx.Fatalf("Failed to create application: %s", resp.Status)
		}

		ctx.Info("Application created successfully")

		// We need to refresh the token to add the new application to the set of
		// claims
		_, err = util.RefreshToken(server)
		if err != nil {
			log.WithError(err).Warn("Failed to refresh token. Please login")
		}
	},
}

var applicationsDeleteCmd = &cobra.Command{
	Use:   "delete [eui]",
	Short: "Delete an application",
	Long:  `ttnctl application delete deletes an existing application.`,
	Run: func(cmd *cobra.Command, args []string) {
		if len(args) < 1 {
			cmd.Help()
			return
		}

		appEUI, err := util.Parse64(args[0])
		if err != nil {
			ctx.Fatalf("Invalid AppEUI: %s", err)
		}

		server := viper.GetString("ttn-account-server")
		uri := fmt.Sprintf("%s/applications/%s", server, fmt.Sprintf("%X", appEUI))
		req, err := util.NewRequestWithAuth(server, "DELETE", uri, nil)
		if err != nil {
			ctx.WithError(err).Fatal("Failed to create authenticated request")
		}

		client := &http.Client{}
		resp, err := client.Do(req)
		if err != nil {
			ctx.WithError(err).Fatal("Failed to delete application")
		}
		if resp.StatusCode != http.StatusOK {
			ctx.Fatalf("Failed to delete application: %s", resp.Status)
		}

		ctx.Info("Application deleted successfully")

		// We need to refresh the token to remove the application from the set of
		// claims
		_, err = util.RefreshToken(server)
		if err != nil {
			log.WithError(err).Warn("Failed to refresh token. Please login")
		}
	},
}

var applicationsAuthorizeCmd = &cobra.Command{
	Use:   "authorize [eui] [e-mail]",
	Short: "Authorize a user for the application",
	Long:  `ttnctl applications authorize lets you authorize a user for an application.`,
	Run: func(cmd *cobra.Command, args []string) {
		if len(args) < 2 {
			cmd.Help()
			return
		}

		appEUI, err := util.Parse64(args[0])
		if err != nil {
			ctx.Fatalf("Invalid AppEUI: %s", err)
		}

		server := viper.GetString("ttn-account-server")
		uri := fmt.Sprintf("%s/applications/%s/authorize", server, fmt.Sprintf("%X", appEUI))
		values := url.Values{
			"email": {args[1]},
		}
		req, err := util.NewRequestWithAuth(server, "PUT", uri, strings.NewReader(values.Encode()))
		if err != nil {
			ctx.WithError(err).Fatal("Failed to create authenticated request")
		}

		req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

		client := &http.Client{}
		resp, err := client.Do(req)
		if err != nil {
			ctx.WithError(err).Fatal("Failed to authorize user")
		}
		if resp.StatusCode != http.StatusOK {
			ctx.Fatalf("Failed to authorize user: %s", resp.Status)
		}

		ctx.Info("User authorized successfully")
	},
}

func init() {
	RootCmd.AddCommand(applicationsCmd)
	applicationsCmd.AddCommand(applicationsCreateCmd)
	applicationsCmd.AddCommand(applicationsDeleteCmd)
	applicationsCmd.AddCommand(applicationsAuthorizeCmd)
}
<|endoftext|>


<commit_before>package node

import (
	"fmt"

	"github.com/docker/docker/api/client"
	"github.com/docker/docker/cli"
	"github.com/docker/engine-api/types/swarm"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
)

func newAcceptCommand(dockerCli *client.DockerCli) *cobra.Command {
	var flags *pflag.FlagSet

	cmd := &cobra.Command{
		Use:   "accept NODE [NODE...]",
		Short: "Accept a node in the swarm",
		Args:  cli.RequiresMinArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			return runAccept(dockerCli, flags, args)
		},
	}

	flags = cmd.Flags()
	return cmd
}

func runAccept(dockerCli *client.DockerCli, flags *pflag.FlagSet, args []string) error {
	for _, id := range args {
		if err := runUpdate(dockerCli, id, func(node *swarm.Node) {
			node.Spec.Membership = swarm.NodeMembershipAccepted
		}); err != nil {
			return err
		}
		fmt.Fprintf(dockerCli.Out(), "Node %s accepted in the swarm.", id)
	}

	return nil
}
<commit_msg>Missing line feed after accepting node<commit_after>package node

import (
	"fmt"

	"github.com/docker/docker/api/client"
	"github.com/docker/docker/cli"
	"github.com/docker/engine-api/types/swarm"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
)

func newAcceptCommand(dockerCli *client.DockerCli) *cobra.Command {
	var flags *pflag.FlagSet

	cmd := &cobra.Command{
		Use:   "accept NODE [NODE...]",
		Short: "Accept a node in the swarm",
		Args:  cli.RequiresMinArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			return runAccept(dockerCli, flags, args)
		},
	}

	flags = cmd.Flags()
	return cmd
}

func runAccept(dockerCli *client.DockerCli, flags *pflag.FlagSet, args []string) error {
	for _, id := range args {
		if err := runUpdate(dockerCli, id, func(node *swarm.Node) {
			node.Spec.Membership = swarm.NodeMembershipAccepted
		}); err != nil {
			return err
		}
		fmt.Fprintf(dockerCli.Out(), "Node %s accepted in the swarm.\n", id)
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"sync"

	log "github.com/Sirupsen/logrus"
	"gopkg.in/vmihailenco/msgpack.v2"
)

import (
	"rank/dos"
	"rank/ss"
)

const (
	UPPER_THRESHOLD = 1024 // storage changed to tree when elements exceeds this
	LOWER_THRESHOLD = 512  // storage changed to sortedset when elements below this
)

const (
	SORTEDSET = iota
	RBTREE
)

// a ranking set
type RankSet struct {
	R    dos.Tree        // rbtree
	S    ss.SortedSet    // sorted-set
	M    map[int32]int32 // ID  => SCORE
	Type int
	sync.RWMutex
}

func NewRankSet() *RankSet {
	r := new(RankSet)
	r.M = make(map[int32]int32)
	r.Type = SORTEDSET // default in sortedset
	return r
}

// toggle storage base on Type
func (r *RankSet) toggle() {
	switch r.Type {
	case SORTEDSET:
		for k, v := range r.M {
			r.R.Insert(v, k)
		}
		r.S.Clear()
		r.Type = RBTREE
		log.Debugf("convert sortedset to rbtree %v", len(r.M))
	case RBTREE:
		for k, v := range r.M {
			r.S.Insert(v, k)
		}
		r.R.Clear()
		r.Type = SORTEDSET
		log.Debugf("convert rbtree to sortedset %v", len(r.M))
	}
}

func (r *RankSet) Update(id, newscore int32) {
	r.Lock()
	defer r.Unlock()

	oldscore, ok := r.M[id]
	if !ok { // new element
		if r.Type == SORTEDSET && len(r.M) > UPPER_THRESHOLD { // do convert
			r.toggle()
		}

		switch r.Type {
		case SORTEDSET:
			r.S.Insert(id, newscore)
		case RBTREE:
			r.R.Insert(id, newscore)
		}
	} else {
		switch r.Type {
		case SORTEDSET:
			r.S.Update(id, newscore)
		case RBTREE:
			_, n := r.R.Locate(oldscore, id)
			r.R.Delete(id, n)
			r.R.Insert(newscore, id)
		}
	}
	r.M[id] = newscore
	return
}

func (r *RankSet) Delete(userid int32) {
	r.Lock()
	defer r.Unlock()
	if r.Type == RBTREE && len(r.M) < LOWER_THRESHOLD { // do convert
		r.toggle()
	}

	switch r.Type {
	case SORTEDSET:
		r.S.Delete(userid)
	case RBTREE:
		score := r.M[userid]
		_, n := r.R.Locate(score, userid)
		r.R.Delete(userid, n)
	}
	delete(r.M, userid)
}

func (r *RankSet) Count() int32 {
	r.RLock()
	defer r.RUnlock()
	return int32(len(r.M))
}

// range [A,B]
func (r *RankSet) GetList(A, B int) (ids []int32, scores []int32) {
	if A < 1 || A > B {
		return
	}
	r.RLock()
	defer r.RUnlock()

	if A > len(r.M) {
		return
	}

	if B > len(r.M) {
		B = len(r.M)
	}

	switch r.Type {
	case SORTEDSET:
		ids, scores = r.S.GetList(A, B)
	case RBTREE:
		ids, scores = r.R.GetList(A, B)
	}
	return
}

// rank of a user
func (r *RankSet) Rank(userid int32) (rank int32, score int32) {
	r.RLock()
	defer r.RUnlock()

	switch r.Type {
	case SORTEDSET:
		rankno := r.S.Locate(userid)
		return int32(rankno), r.M[userid]
	case RBTREE:
		rankno, _ := r.R.Locate(r.M[userid], userid)
		return int32(rankno), r.M[userid]
	}
	return
}

// serialization
func (r *RankSet) Marshal() ([]byte, error) {
	r.RLock()
	defer r.RUnlock()
	return msgpack.Marshal(r.M)
}

func (r *RankSet) Unmarshal(bin []byte) error {
	m := make(map[int32]int32)
	r.Lock()
	defer r.Unlock()
	err := msgpack.Unmarshal(bin, &m)
	if err != nil {
		return err
	}

	r.M = m
	if len(r.M) > UPPER_THRESHOLD {
		for id, score := range m {
			r.R.Insert(score, id)
		}
		r.Type = RBTREE
		log.Debugf("rank restored into rbtree %v", len(r.M))
	} else {
		for id, score := range m {
			r.S.Insert(id, score)
		}
		r.Type = SORTEDSET
		log.Debugf("rank restored into sortedset %v", len(r.M))
	}

	return nil
}
<commit_msg>fix rankset.update<commit_after>package main

import (
	"sync"

	log "github.com/Sirupsen/logrus"
	"gopkg.in/vmihailenco/msgpack.v2"
)

import (
	"rank/dos"
	"rank/ss"
)

const (
	UPPER_THRESHOLD = 1024 // storage changed to tree when elements exceeds this
	LOWER_THRESHOLD = 512  // storage changed to sortedset when elements below this
)

const (
	SORTEDSET = iota
	RBTREE
)

// a ranking set
type RankSet struct {
	R    dos.Tree        // rbtree
	S    ss.SortedSet    // sorted-set
	M    map[int32]int32 // ID  => SCORE
	Type int
	sync.RWMutex
}

func NewRankSet() *RankSet {
	r := new(RankSet)
	r.M = make(map[int32]int32)
	r.Type = SORTEDSET // default in sortedset
	return r
}

// toggle storage base on Type
func (r *RankSet) toggle() {
	switch r.Type {
	case SORTEDSET:
		for k, v := range r.M {
			r.R.Insert(v, k)
		}
		r.S.Clear()
		r.Type = RBTREE
		log.Debugf("convert sortedset to rbtree %v", len(r.M))
	case RBTREE:
		for k, v := range r.M {
			r.S.Insert(v, k)
		}
		r.R.Clear()
		r.Type = SORTEDSET
		log.Debugf("convert rbtree to sortedset %v", len(r.M))
	}
}

func (r *RankSet) Update(id, newscore int32) {
	r.Lock()
	defer r.Unlock()

	oldscore, ok := r.M[id]
	if !ok { // new element
		if r.Type == SORTEDSET && len(r.M) > UPPER_THRESHOLD { // do convert
			r.toggle()
		}

		switch r.Type {
		case SORTEDSET:
			r.S.Insert(id, newscore)
		case RBTREE:
			r.R.Insert(newscore, id)
		}
	} else {
		switch r.Type {
		case SORTEDSET:
			r.S.Update(id, newscore)
		case RBTREE:
			_, n := r.R.Locate(oldscore, id)
			r.R.Delete(id, n)
			r.R.Insert(newscore, id)
		}
	}
	r.M[id] = newscore
	return
}

func (r *RankSet) Delete(userid int32) {
	r.Lock()
	defer r.Unlock()
	if r.Type == RBTREE && len(r.M) < LOWER_THRESHOLD { // do convert
		r.toggle()
	}

	switch r.Type {
	case SORTEDSET:
		r.S.Delete(userid)
	case RBTREE:
		score := r.M[userid]
		_, n := r.R.Locate(score, userid)
		r.R.Delete(userid, n)
	}
	delete(r.M, userid)
}

func (r *RankSet) Count() int32 {
	r.RLock()
	defer r.RUnlock()
	return int32(len(r.M))
}

// range [A,B]
func (r *RankSet) GetList(A, B int) (ids []int32, scores []int32) {
	if A < 1 || A > B {
		return
	}
	r.RLock()
	defer r.RUnlock()

	if A > len(r.M) {
		return
	}

	if B > len(r.M) {
		B = len(r.M)
	}

	switch r.Type {
	case SORTEDSET:
		ids, scores = r.S.GetList(A, B)
	case RBTREE:
		ids, scores = r.R.GetList(A, B)
	}
	return
}

// rank of a user
func (r *RankSet) Rank(userid int32) (rank int32, score int32) {
	r.RLock()
	defer r.RUnlock()

	switch r.Type {
	case SORTEDSET:
		rankno := r.S.Locate(userid)
		return int32(rankno), r.M[userid]
	case RBTREE:
		rankno, _ := r.R.Locate(r.M[userid], userid)
		return int32(rankno), r.M[userid]
	}
	return
}

// serialization
func (r *RankSet) Marshal() ([]byte, error) {
	r.RLock()
	defer r.RUnlock()
	return msgpack.Marshal(r.M)
}

func (r *RankSet) Unmarshal(bin []byte) error {
	m := make(map[int32]int32)
	r.Lock()
	defer r.Unlock()
	err := msgpack.Unmarshal(bin, &m)
	if err != nil {
		return err
	}

	r.M = m
	if len(r.M) > UPPER_THRESHOLD {
		for id, score := range m {
			r.R.Insert(score, id)
		}
		r.Type = RBTREE
		log.Debugf("rank restored into rbtree %v", len(r.M))
	} else {
		for id, score := range m {
			r.S.Insert(id, score)
		}
		r.Type = SORTEDSET
		log.Debugf("rank restored into sortedset %v", len(r.M))
	}

	return nil
}
<|endoftext|>


<commit_before>package sse

import (
	"bytes"
	"fmt"
	"unicode/utf8"

	"github.com/pkg/errors"
)

var (
	// standard forms of tokens used in this implementation
	endOfLine  = []byte{'\n'}
	fieldDelim = []byte{':', ' '}

	// Comments are used for keep-alives
	commentMarker = []byte{':'}

	// the spec specifies this as a replacement character for errors
	replacementChar = []byte("\uFFFD")
)

// parser state.  Use a pointer to an integer as a very simple state
// machine.
type parseState int

const (
	startOfStream parseState = iota
	startOfLine
	comment
	readField
	readDelim
	midEOL
	atEndOfLine
	endOfStream
)

func newSplitter() *parseState {
	state := startOfStream
	return &state
}

// SplitFunc returns a bufio.SplitFunc for the text/event-stream MIME
// type.  This is a stateful scanner which tokenizes the event stream.
func SplitFunc() func([]byte, bool) (int, []byte, error) {
	dfa := newSplitter()
	return func(data []byte, atEOF bool) (int, []byte, error) {
		return dfa.scan(data, atEOF)
	}
}

// scan is the function called by Scanner.Split() to find the next token.
func (state *parseState) scan(data []byte, atEOF bool) (advance int, token []byte, err error) {
	if len(data) == 0 {
		if *state == atEndOfLine || *state == midEOL {
			*state = endOfStream
			return 0, endOfLine, nil
		} else {
			return 0, nil, nil
		}
	}
	switch *state {
	case startOfStream:
		advance, token, err = state.scanBOM(data, atEOF)
		if advance != -1 {
			return advance, token, err
		}
		fallthrough
	case startOfLine:
		return state.scanStartOfLine(data, atEOF)
	case midEOL:
		if data[0] == '\n' {
			return 1, nil, nil
		} else {
			*state = startOfLine
			return state.scanStartOfLine(data, atEOF)
		}
	case comment, readDelim:
		return state.scanAnyChar(data, atEOF)
	case readField:
		return state.scanDelim(data, atEOF)
	case atEndOfLine:
		return state.scanEndOfLine(data, atEOF)
	default:
		panic(fmt.Sprintf("bad state %v", state))
	}
}

// scanBOM strips off a useless byte-order mark
// > eventsource ABNF rule: stream        = [ bom ] *event
func (state *parseState) scanBOM(data []byte, atEOF bool) (advance int, token []byte, err error) {
	var maybe, no bool
	if data[0] != '\xEF' {
		no = true
	} else if len(data) == 1 {
		maybe = true
	} else if data[1] != '\xBB' {
		no = true
	} else if len(data) == 2 {
		maybe = true
	} else if data[2] != '\xBF' {
		no = true
	}
	if maybe {
		return 0, nil, nil
	} else {
		*state = startOfLine
		if no {
			return -1, nil, nil
		} else {
			return 3, nil, nil
		}
	}
}

// scanStartOfLine scans most of this rule:
// >   event         = *( comment / field ) end-of-line
func (state *parseState) scanStartOfLine(data []byte, atEOF bool) (advance int, token []byte, err error) {
	switch data[0] {
	case '\r', '\n':
		// an event is complete
		*state = atEndOfLine
		return state.scanEndOfLine(data, atEOF)
	case ':':
		*state = comment
		return 1, commentMarker, nil
	default:
		// field name character
		idx := bytes.IndexAny(data, "\n\r:")
		if idx == -1 {
			if !atEOF {
				return 0, nil, nil
			} else {
				idx = len(data)
			}
		}
		*state = readField
		return idx, validUTF8(data[:idx]), nil
	}
}

// scanAnyChar scans:
// > any-char      = %x0000-0009 / %x000B-000C / %x000E-10FFFF
// >                ; a Unicode character other than U+000A LINE FEED (LF) or U+000D CARRIAGE RETURN (CR)
func (state *parseState) scanAnyChar(data []byte, atEOF bool) (advance int, token []byte, err error) {
	switch data[0] {
	case '\r', '\n':
		// no data.  can't emit an empty token, so...
		*state = atEndOfLine
		return state.scanEndOfLine(data, atEOF)
	default:
		idx := bytes.IndexAny(data, "\n\r")
		if idx == -1 && !atEOF {
			return 0, nil, nil
		} else {
			if idx == -1 {
				idx = len(data)
			}
			*state = atEndOfLine
			return idx, validUTF8(data[:idx]), nil
		}
	}
}

// scanDelim scans the colon and space from:
// > field         = 1*name-char [ colon [ space ] *any-char ] end-of-line
func (state *parseState) scanDelim(data []byte, atEOF bool) (advance int, token []byte, err error) {
	switch data[0] {
	case '\r', '\n':
		// no data.  can't emit an empty token, so...
		*state = atEndOfLine
		return state.scanEndOfLine(data, atEOF)
	case ':':
		if len(data) == 1 {
			if atEOF {
				*state = readDelim
				return 1, fieldDelim, nil
			} else {
				return 0, nil, nil
			}
		} else {
			*state = readDelim
			advance = 1
			if data[1] == ' ' {
				advance = 2
			}
			return advance, fieldDelim, nil
		}
	default:
		// should never happen; can't get into readField state unless
		// next char is '\r', '\n' or ':'
		return 0, nil, errors.Errorf("expected: field delimiter or end of line, saw: %v", data[0])
	}
}

// scanEndOfLine scans a line ending.  Lines may be terminated by \r,
// \n *or* \r\n (or the end of the stream), complicating this somewhat.
// rule:
// > end-of-line   = ( cr lf / cr / lf / eof )
func (state *parseState) scanEndOfLine(data []byte, atEOF bool) (advance int, token []byte, err error) {
	switch data[0] {
	case '\r':
		if len(data) == 1 {
			if !atEOF {
				*state = midEOL
			}
			return 1, endOfLine, nil
		} else if data[1] == '\n' {
			*state = startOfLine
			return 2, endOfLine, nil
		} else {
			*state = startOfLine
			return 1, endOfLine, nil
		}
	case '\n':
		*state = startOfLine
		return 1, endOfLine, nil
	default:
		// should never happen; can't get into readField state unless
		// next char is '\r' or '\n'
		return 0, nil, errors.Errorf("Expected end-of-line, found: %v", data[0])
	}
}

// validUTF8 is a function which ensures this part of the spec:
// > Bytes or sequences of bytes that are not valid UTF-8 sequences must
// > be interpreted as the U+FFFD REPLACEMENT CHARACTER.
func validUTF8(network []byte) []byte {
	if utf8.Valid(network) {
		return network
	}
	// FIXME - can I just return []byte(string([]rune(network)))?
	valid := make([]byte, 0, len(network)+32)
	for len(network) > 0 {
		rune, size := utf8.DecodeRune(network)
		if rune == utf8.RuneError {
			// "or sequences of bytes" - skip to the next rune start
			network = network[1:]
			for len(network) > 0 && !utf8.RuneStart(network[0]) {
				network = network[1:]
			}
			valid = append(valid, []byte(string(rune))...)
		} else {
			valid = append(valid, network[:size]...)
			network = network[size:]
		}
	}
	return valid
}
<commit_msg>Remove unnecessary else blocks #golint<commit_after>package sse

import (
	"bytes"
	"fmt"
	"unicode/utf8"

	"github.com/pkg/errors"
)

var (
	// standard forms of tokens used in this implementation
	endOfLine  = []byte{'\n'}
	fieldDelim = []byte{':', ' '}

	// Comments are used for keep-alives
	commentMarker = []byte{':'}

	// the spec specifies this as a replacement character for errors
	replacementChar = []byte("\uFFFD")
)

// parser state.  Use a pointer to an integer as a very simple state
// machine.
type parseState int

const (
	startOfStream parseState = iota
	startOfLine
	comment
	readField
	readDelim
	midEOL
	atEndOfLine
	endOfStream
)

func newSplitter() *parseState {
	state := startOfStream
	return &state
}

// SplitFunc returns a bufio.SplitFunc for the text/event-stream MIME
// type.  This is a stateful scanner which tokenizes the event stream.
func SplitFunc() func([]byte, bool) (int, []byte, error) {
	dfa := newSplitter()
	return func(data []byte, atEOF bool) (int, []byte, error) {
		return dfa.scan(data, atEOF)
	}
}

// scan is the function called by Scanner.Split() to find the next token.
func (state *parseState) scan(data []byte, atEOF bool) (advance int, token []byte, err error) {
	if len(data) == 0 {
		if *state == atEndOfLine || *state == midEOL {
			*state = endOfStream
			return 0, endOfLine, nil
		}
		return 0, nil, nil
	}
	switch *state {
	case startOfStream:
		advance, token, err = state.scanBOM(data, atEOF)
		if advance != -1 {
			return advance, token, err
		}
		fallthrough
	case startOfLine:
		return state.scanStartOfLine(data, atEOF)
	case midEOL:
		if data[0] == '\n' {
			return 1, nil, nil
		}
		*state = startOfLine
		return state.scanStartOfLine(data, atEOF)
	case comment, readDelim:
		return state.scanAnyChar(data, atEOF)
	case readField:
		return state.scanDelim(data, atEOF)
	case atEndOfLine:
		return state.scanEndOfLine(data, atEOF)
	default:
		panic(fmt.Sprintf("bad state %v", state))
	}
}

// scanBOM strips off a useless byte-order mark
// > eventsource ABNF rule: stream        = [ bom ] *event
func (state *parseState) scanBOM(data []byte, atEOF bool) (advance int, token []byte, err error) {
	var maybe, no bool
	if data[0] != '\xEF' {
		no = true
	} else if len(data) == 1 {
		maybe = true
	} else if data[1] != '\xBB' {
		no = true
	} else if len(data) == 2 {
		maybe = true
	} else if data[2] != '\xBF' {
		no = true
	}
	if maybe {
		return 0, nil, nil
	}
	*state = startOfLine
	if no {
		return -1, nil, nil
	}
	return 3, nil, nil
}

// scanStartOfLine scans most of this rule:
// >   event         = *( comment / field ) end-of-line
func (state *parseState) scanStartOfLine(data []byte, atEOF bool) (advance int, token []byte, err error) {
	switch data[0] {
	case '\r', '\n':
		// an event is complete
		*state = atEndOfLine
		return state.scanEndOfLine(data, atEOF)
	case ':':
		*state = comment
		return 1, commentMarker, nil
	default:
		// field name character
		idx := bytes.IndexAny(data, "\n\r:")
		if idx == -1 {
			if !atEOF {
				return 0, nil, nil
			}
			idx = len(data)
		}
		*state = readField
		return idx, validUTF8(data[:idx]), nil
	}
}

// scanAnyChar scans:
// > any-char      = %x0000-0009 / %x000B-000C / %x000E-10FFFF
// >                ; a Unicode character other than U+000A LINE FEED (LF) or U+000D CARRIAGE RETURN (CR)
func (state *parseState) scanAnyChar(data []byte, atEOF bool) (advance int, token []byte, err error) {
	switch data[0] {
	case '\r', '\n':
		// no data.  can't emit an empty token, so...
		*state = atEndOfLine
		return state.scanEndOfLine(data, atEOF)
	default:
		idx := bytes.IndexAny(data, "\n\r")
		if idx == -1 && !atEOF {
			return 0, nil, nil
		}
		if idx == -1 {
			idx = len(data)
		}
		*state = atEndOfLine
		return idx, validUTF8(data[:idx]), nil
	}
}

// scanDelim scans the colon and space from:
// > field         = 1*name-char [ colon [ space ] *any-char ] end-of-line
func (state *parseState) scanDelim(data []byte, atEOF bool) (advance int, token []byte, err error) {
	switch data[0] {
	case '\r', '\n':
		// no data.  can't emit an empty token, so...
		*state = atEndOfLine
		return state.scanEndOfLine(data, atEOF)
	case ':':
		if len(data) == 1 {
			if atEOF {
				*state = readDelim
				return 1, fieldDelim, nil
			}
			return 0, nil, nil
		}
		*state = readDelim
		advance = 1
		if data[1] == ' ' {
			advance = 2
		}
		return advance, fieldDelim, nil
	default:
		// should never happen; can't get into readField state unless
		// next char is '\r', '\n' or ':'
		return 0, nil, errors.Errorf("expected: field delimiter or end of line, saw: %v", data[0])
	}
}

// scanEndOfLine scans a line ending.  Lines may be terminated by \r,
// \n *or* \r\n (or the end of the stream), complicating this somewhat.
// rule:
// > end-of-line   = ( cr lf / cr / lf / eof )
func (state *parseState) scanEndOfLine(data []byte, atEOF bool) (advance int, token []byte, err error) {
	switch data[0] {
	case '\r':
		if len(data) == 1 {
			if !atEOF {
				*state = midEOL
			}
			return 1, endOfLine, nil
		} else if data[1] == '\n' {
			*state = startOfLine
			return 2, endOfLine, nil
		} else {
			*state = startOfLine
			return 1, endOfLine, nil
		}
	case '\n':
		*state = startOfLine
		return 1, endOfLine, nil
	default:
		// should never happen; can't get into readField state unless
		// next char is '\r' or '\n'
		return 0, nil, errors.Errorf("Expected end-of-line, found: %v", data[0])
	}
}

// validUTF8 is a function which ensures this part of the spec:
// > Bytes or sequences of bytes that are not valid UTF-8 sequences must
// > be interpreted as the U+FFFD REPLACEMENT CHARACTER.
func validUTF8(network []byte) []byte {
	if utf8.Valid(network) {
		return network
	}
	// FIXME - can I just return []byte(string([]rune(network)))?
	valid := make([]byte, 0, len(network)+32)
	for len(network) > 0 {
		rune, size := utf8.DecodeRune(network)
		if rune == utf8.RuneError {
			// "or sequences of bytes" - skip to the next rune start
			network = network[1:]
			for len(network) > 0 && !utf8.RuneStart(network[0]) {
				network = network[1:]
			}
			valid = append(valid, []byte(string(rune))...)
		} else {
			valid = append(valid, network[:size]...)
			network = network[size:]
		}
	}
	return valid
}
<|endoftext|>


<commit_before>package robotstxt

import (
    "container/vector"
    "fmt"
    "go/token"
    "os"
    "strings"
    "utf8"
)

type ByteScanner struct {
    ErrorCount int
    Quiet      bool

    buf       []byte
    pos       token.Position
    lastChunk bool
    ch        int
    //
    //state string
}

const WhitespaceChars = " \t\v"

func NewByteScanner(srcname string, quiet bool) *ByteScanner {
    return &ByteScanner{
        Quiet: quiet,
        ch:    -1,
        pos:   token.Position{Filename: srcname},
        //state: "start",
    }
}

func (s *ByteScanner) Feed(input []byte, end bool) (bool, os.Error) {
    s.buf = input
    s.pos.Offset = 0
    s.pos.Line = 1
    s.pos.Column = 1
    s.lastChunk = end
    // Read first char into look-ahead buffer `s.ch`.
    s.nextChar()
    return false, nil
}

func (s *ByteScanner) GetPosition() token.Position {
    return s.pos
}

func (s *ByteScanner) Scan() (string, os.Error) {
    //println("--- Scan(). Offset / len(s.buf): ", s.pos.Offset, len(s.buf))

    for {
        // Note Offset > len, not >=, so we can Scan last character.
        if s.lastChunk && s.pos.Offset > len(s.buf) {
            return "", os.EOF
        }

        s.skipSpace()

        if s.ch == -1 {
            return "", os.EOF
        }

        // EOL
        if s.isEol() {
            // skip subsequent newline chars
            for s.ch != -1 && s.isEol() {
                s.nextChar()
            }
            // emit newline as separate token
            return "\n", nil
        }

        // skip comments
        if s.ch == '#' {
            s.skipUntilEol()
            //            s.state = "start"
            if s.ch == -1 {
                return "", os.EOF
            }
            // emit newline as separate token
            return "\n", nil
        }

        // else we found something
        break
    }

    /*
       if s.state == "start" {
           s.state = "key"
       }
    */

    tok := string(s.ch)
    s.nextChar()
    for s.ch != -1 && !s.isSpace() && !s.isEol() {
        if s.ch == ':' {
            //            s.state = "pre-value"
            s.nextChar()
            break
        }

        tok += string(s.ch)
        s.nextChar()
    }
    return tok, nil
}

func (s *ByteScanner) ScanAll() ([]string, os.Error) {
    var results vector.StringVector
    for {
        t, err := s.Scan()
        if t != "" {
            results.Push(t)
        }
        if err == os.EOF {
            break
        }
        if err != nil {
            return results, err
        }
    }
    return results, nil
}

func (s *ByteScanner) error(pos token.Position, msg string) {
    s.ErrorCount++
    if !s.Quiet {
        fmt.Fprintf(os.Stderr, "%s: %s\n", pos.String(), msg)
    }
}

func (s *ByteScanner) isEol() bool {
    return s.ch == '\n' || s.ch == '\r'
}

func (s *ByteScanner) isSpace() bool {
    return strings.Index(WhitespaceChars, string(s.ch)) >= 0
}

func (s *ByteScanner) skipSpace() {
    //println("--- string(ch): ", s.ch, ".")
    for s.ch != -1 && s.isSpace() {
        s.nextChar()
    }
}

func (s *ByteScanner) skipUntilEol() {
    //println("--- string(ch): ", s.ch, ".")
    for s.ch != -1 && !s.isEol() {
        s.nextChar()
    }
    // skip subsequent newline chars
    for s.ch != -1 && s.isEol() {
        s.nextChar()
    }
}

// Reads next Unicode char.
func (s *ByteScanner) nextChar() (int, os.Error) {
    //println("--- nextChar(). Offset / len(s.buf): ", s.pos.Offset, len(s.buf))

    if s.pos.Offset >= len(s.buf) {
        s.ch = -1
        return s.ch, os.EOF
    }
    s.pos.Column++
    if s.ch == '\n' {
        s.pos.Line++
        s.pos.Column = 1
    }
    r, w := int(s.buf[s.pos.Offset]), 1
    if r >= 0x80 {
        r, w = utf8.DecodeRune(s.buf[s.pos.Offset:])
        if r == utf8.RuneError && w == 1 {
            s.error(s.pos, "illegal UTF-8 encoding")
        }
    }
    s.pos.Offset += w
    s.ch = r
    return s.ch, nil
}
<commit_msg>scanner: Change error string, added library name.<commit_after>package robotstxt

import (
    "container/vector"
    "fmt"
    "go/token"
    "os"
    "strings"
    "utf8"
)

type ByteScanner struct {
    ErrorCount int
    Quiet      bool

    buf       []byte
    pos       token.Position
    lastChunk bool
    ch        int
    //
    //state string
}

const WhitespaceChars = " \t\v"

func NewByteScanner(srcname string, quiet bool) *ByteScanner {
    return &ByteScanner{
        Quiet: quiet,
        ch:    -1,
        pos:   token.Position{Filename: srcname},
        //state: "start",
    }
}

func (s *ByteScanner) Feed(input []byte, end bool) (bool, os.Error) {
    s.buf = input
    s.pos.Offset = 0
    s.pos.Line = 1
    s.pos.Column = 1
    s.lastChunk = end
    // Read first char into look-ahead buffer `s.ch`.
    s.nextChar()
    return false, nil
}

func (s *ByteScanner) GetPosition() token.Position {
    return s.pos
}

func (s *ByteScanner) Scan() (string, os.Error) {
    //println("--- Scan(). Offset / len(s.buf): ", s.pos.Offset, len(s.buf))

    for {
        // Note Offset > len, not >=, so we can Scan last character.
        if s.lastChunk && s.pos.Offset > len(s.buf) {
            return "", os.EOF
        }

        s.skipSpace()

        if s.ch == -1 {
            return "", os.EOF
        }

        // EOL
        if s.isEol() {
            // skip subsequent newline chars
            for s.ch != -1 && s.isEol() {
                s.nextChar()
            }
            // emit newline as separate token
            return "\n", nil
        }

        // skip comments
        if s.ch == '#' {
            s.skipUntilEol()
            //            s.state = "start"
            if s.ch == -1 {
                return "", os.EOF
            }
            // emit newline as separate token
            return "\n", nil
        }

        // else we found something
        break
    }

    /*
       if s.state == "start" {
           s.state = "key"
       }
    */

    tok := string(s.ch)
    s.nextChar()
    for s.ch != -1 && !s.isSpace() && !s.isEol() {
        if s.ch == ':' {
            //            s.state = "pre-value"
            s.nextChar()
            break
        }

        tok += string(s.ch)
        s.nextChar()
    }
    return tok, nil
}

func (s *ByteScanner) ScanAll() ([]string, os.Error) {
    var results vector.StringVector
    for {
        t, err := s.Scan()
        if t != "" {
            results.Push(t)
        }
        if err == os.EOF {
            break
        }
        if err != nil {
            return results, err
        }
    }
    return results, nil
}

func (s *ByteScanner) error(pos token.Position, msg string) {
    s.ErrorCount++
    if !s.Quiet {
        fmt.Fprintf(os.Stderr, "robotstxt from %s: %s\n", pos.String(), msg)
    }
}

func (s *ByteScanner) isEol() bool {
    return s.ch == '\n' || s.ch == '\r'
}

func (s *ByteScanner) isSpace() bool {
    return strings.Index(WhitespaceChars, string(s.ch)) >= 0
}

func (s *ByteScanner) skipSpace() {
    //println("--- string(ch): ", s.ch, ".")
    for s.ch != -1 && s.isSpace() {
        s.nextChar()
    }
}

func (s *ByteScanner) skipUntilEol() {
    //println("--- string(ch): ", s.ch, ".")
    for s.ch != -1 && !s.isEol() {
        s.nextChar()
    }
    // skip subsequent newline chars
    for s.ch != -1 && s.isEol() {
        s.nextChar()
    }
}

// Reads next Unicode char.
func (s *ByteScanner) nextChar() (int, os.Error) {
    //println("--- nextChar(). Offset / len(s.buf): ", s.pos.Offset, len(s.buf))

    if s.pos.Offset >= len(s.buf) {
        s.ch = -1
        return s.ch, os.EOF
    }
    s.pos.Column++
    if s.ch == '\n' {
        s.pos.Line++
        s.pos.Column = 1
    }
    r, w := int(s.buf[s.pos.Offset]), 1
    if r >= 0x80 {
        r, w = utf8.DecodeRune(s.buf[s.pos.Offset:])
        if r == utf8.RuneError && w == 1 {
            s.error(s.pos, "illegal UTF-8 encoding")
        }
    }
    s.pos.Offset += w
    s.ch = r
    return s.ch, nil
}
<|endoftext|>


<commit_before>// Package scanner provides a recursive file scanner that is useful for
// efficiently processing relatively static datasets.
package scanner

import (
	"io/ioutil"
	"os"
	"strings"
	"sync"
)

// Scanner recursively scans a directory for files.
type Scanner struct {

	// Buffer is length of the files channel buffer.
	Buffer int

	// errs is the channel that errors are sent to.
	errs chan error

	// files is the channel that discovered files are sent to.
	files chan string

	// path is the directory being scanned.
	path string

	// separator is the OS's path separator represented as a string. We
	// store this value in a struct so we don't have to repeatedly convert
	// the rune to a string. This is a micro-optimization, but an
	// optimization none the less.
	separator string
}

// New returns a new Scanner instance.
func New(path string) *Scanner {
	return &Scanner{
		Buffer:    1,
		errs:      make(chan error),
		path:      path,
		separator: string(os.PathSeparator),
	}
}

// Scan recursively scans the directory and sends the files and errors to
// the passed Handler in goroutines.
func (s *Scanner) Scan(h Handler) {
	var wg sync.WaitGroup
	s.files = make(chan string, s.Buffer)

	go func() {
		s.scan(s.path)
		close(s.files)
		close(s.errs)
	}()

	if h == nil {
		return
	}

	wg.Add(2)

	go func() {
		defer wg.Done()
		for err := range s.errs {
			h.HandleError(err)
		}
	}()

	go func() {
		defer wg.Done()
		for f := range s.files {
			h.Handle(f)
		}
	}()

	wg.Wait()
}

// scan recursively scans path and sends the discovered files and errors to
// the built-in channels.
func (s *Scanner) scan(path string) {
	files, err := ioutil.ReadDir(path)
	if err != nil {
		s.errs <- err
		return
	}

	basedir := strings.TrimRight(path, s.separator)
	for _, f := range files {
		file := basedir + s.separator + f.Name()
		if f.IsDir() {
			s.scan(file)
		} else {
			s.files <- file
		}
	}
}

// Handler is an interface for handlers that process scanned files.
type Handler interface {

	// Handle is passed the path to the file that was discovered during the
	// scan operation. Files are passed in a single-threaded manner, so
	// concurrency is the responsibility of the Handler implementation.
	Handle(string)

	// HandleError is passed the errors that occur during scan operations.
	// Like the Handle method, errors are passed in a single-threaded
	// manner.
	HandleError(error)
}

// NewMemoryHandler returns a new MemoryHandler instance, which stores the
// scanned files and errors in the exported Files and Errors fields
// respectively.
func NewMemoryHandler() *MemoryHandler {
	return &MemoryHandler{}
}

// MemoryHandler implements Handler and stores the scanned files and errors
// in slices.
type MemoryHandler struct {
	Files  []string
	Errors []error
}

// Handle implements Handler.Handle and stores file in a struct.
func (h *MemoryHandler) Handle(file string) {
	h.Files = append(h.Files, file)
}

// HandleError implements Handler.HandleError stores err in a struct.
func (h *MemoryHandler) HandleError(err error) {
	h.Errors = append(h.Errors, err)
}
<commit_msg>Ignore symlinks for now<commit_after>// Package scanner provides a recursive file scanner that is useful for
// efficiently processing relatively static datasets.
package scanner

import (
	"io/ioutil"
	"os"
	"strings"
	"sync"
)

// Scanner recursively scans a directory for files.
type Scanner struct {

	// Buffer is length of the files channel buffer.
	Buffer int

	// errs is the channel that errors are sent to.
	errs chan error

	// files is the channel that discovered files are sent to.
	files chan string

	// path is the directory being scanned.
	path string

	// separator is the OS's path separator represented as a string. We
	// store this value in a struct so we don't have to repeatedly convert
	// the rune to a string. This is a micro-optimization, but an
	// optimization none the less.
	separator string
}

// New returns a new Scanner instance.
func New(path string) *Scanner {
	return &Scanner{
		Buffer:    1,
		errs:      make(chan error),
		path:      path,
		separator: string(os.PathSeparator),
	}
}

// Scan recursively scans the directory and sends the files and errors to
// the passed Handler in goroutines.
func (s *Scanner) Scan(h Handler) {
	var wg sync.WaitGroup
	s.files = make(chan string, s.Buffer)

	go func() {
		s.scan(s.path)
		close(s.files)
		close(s.errs)
	}()

	if h == nil {
		return
	}

	wg.Add(2)

	go func() {
		defer wg.Done()
		for err := range s.errs {
			h.HandleError(err)
		}
	}()

	go func() {
		defer wg.Done()
		for f := range s.files {
			h.Handle(f)
		}
	}()

	wg.Wait()
}

// scan recursively scans path and sends the discovered files and errors to
// the built-in channels.
func (s *Scanner) scan(path string) {
	files, err := ioutil.ReadDir(path)
	if err != nil {
		s.errs <- err
		return
	}

	basedir := strings.TrimRight(path, s.separator)
	for _, f := range files {
		file := basedir + s.separator + f.Name()
		if f.IsDir() {
			s.scan(file)
		} else if f.Mode()&os.ModeSymlink == os.ModeSymlink {
			// TODO Figure out how to handle symlinks.
		} else {
			s.files <- file
		}
	}
}

// Handler is an interface for handlers that process scanned files.
type Handler interface {

	// Handle is passed the path to the file that was discovered during the
	// scan operation. Files are passed in a single-threaded manner, so
	// concurrency is the responsibility of the Handler implementation.
	Handle(string)

	// HandleError is passed the errors that occur during scan operations.
	// Like the Handle method, errors are passed in a single-threaded
	// manner.
	HandleError(error)
}

// NewMemoryHandler returns a new MemoryHandler instance, which stores the
// scanned files and errors in the exported Files and Errors fields
// respectively.
func NewMemoryHandler() *MemoryHandler {
	return &MemoryHandler{}
}

// MemoryHandler implements Handler and stores the scanned files and errors
// in slices.
type MemoryHandler struct {
	Files  []string
	Errors []error
}

// Handle implements Handler.Handle and stores file in a struct.
func (h *MemoryHandler) Handle(file string) {
	h.Files = append(h.Files, file)
}

// HandleError implements Handler.HandleError stores err in a struct.
func (h *MemoryHandler) HandleError(err error) {
	h.Errors = append(h.Errors, err)
}
<|endoftext|>