<commit_before>// Package jq provides go bindings for libjq providing a streaming filter of
// JSON documents.
//
// This package provides a thin layer on top of stedolan's libjq -- it would
// likely be helpful to read through the wiki pages about it:
//
// jv: the JSON value type https://github.com/stedolan/jq/wiki/C-API:-jv
//
// libjq: https://github.com/stedolan/jq/wiki/C-API:-libjq
package jq

/*
To install
$ ./configure --disable-maintainer-mode --prefix=$PWD/BUILD
$ make install-libLTLIBRARIES install-includeHEADERS
*/

/*
#cgo LDFLAGS: -ljq
#cgo linux LDFLAGS: -lm

#include <jq.h>
#include <jv.h>

#include <stdlib.h>

void install_jq_error_cb(jq_state *jq, unsigned long long id);
*/
import "C"
import (
	"errors"
	"fmt"
	"sync"
	"sync/atomic"
	"unsafe"
)

// Jq encapsulates the state needed to interface with the libjq C library
type Jq struct {
	_state       *C.struct_jq_state
	errorStoreId uint64
	running      sync.WaitGroup
}

// New initializes a new JQ object and the underlying C library.
func New() (*Jq, error) {
	jq := new(Jq)

	var err error
	jq._state, err = C.jq_init()

	if err != nil {
		return nil, err
	} else if jq == nil {
		return nil, errors.New("jq_init returned nil -- out of memory?")
	}

	return jq, nil
}

// Close the handle to libjq and free C resources.
//
// If Start() has been called this will block until the input Channel it
// returns has been closed.
func (jq *Jq) Close() {
	// If the goroutine from Start() is running we need to make sure it finished cleanly
	// Wait until we aren't running before freeing C things.
	//
	jq.running.Wait()
	if jq._state != nil {
		C.jq_teardown(&jq._state)
		jq._state = nil
	}
	if jq.errorStoreId != 0 {
		globalErrorChannels.Delete(jq.errorStoreId)
		jq.errorStoreId = 0
	}
}

// We cant pass many things over the Go/C boundary, so instead of passing the error channel we pass an opaque indentifier (a 64bit int as it turns out) and use that to look up in a global variable
type errorLookupState struct {
	sync.RWMutex
	idCounter uint64
	channels  map[uint64]chan<- error
}

func (e *errorLookupState) Add(c chan<- error) uint64 {
	newID := atomic.AddUint64(&e.idCounter, 1)
	e.RWMutex.Lock()
	defer e.RWMutex.Unlock()
	e.channels[newID] = c
	return newID
}

func (e *errorLookupState) Get(id uint64) chan<- error {
	e.RWMutex.RLock()
	defer e.RWMutex.RUnlock()
	c, ok := e.channels[id]
	if !ok {
		panic(fmt.Sprintf("Tried to get error channel #%d out of store but it wasn't there!", id))
	}
	return c
}

func (e *errorLookupState) Delete(id uint64) {
	e.RWMutex.Lock()
	defer e.RWMutex.Unlock()
	delete(e.channels, id)
}

// The global state - this also serves to keep the channel in scope by keeping
// a reference to it that the GC can see
var globalErrorChannels = errorLookupState{
	channels: make(map[uint64]chan<- error),
}

//export goLibjqErrorHandler
func goLibjqErrorHandler(id uint64, jv C.jv) {
	ch := globalErrorChannels.Get(id)

	err := _ConvertError(jv)
	ch <- err
}

// Start will compile `program` and return a three channels: input, output and
// error. Sending a jq.Jv* to input cause the program to be run to it and
// one-or-more results returned as jq.Jv* on the output channel, or one or more
// error values sent to the error channel. When you are done sending values
// close the input channel.
//
// args is a list of key/value pairs to bind as variables into the program, and
// must be an array type even if empty. Each element of the array should be an
// object with a "name" and "value" properties. Name should exclude the "$"
// sign. For example this is `[ {"name": "n", "value": 1 } ]` would then be
// `$n` in the programm.
//
// This function is not reentereant -- in that you cannot and should not call
// Start again until you have closed the previous input channel.
//
// If there is a problem compiling the JQ program then the errors will be
// reported on error channel before any input is read so makle sure you account
// for this case.
//
// Any jq.Jv* values passed to the input channel will be owned by the channel.
// If you want to keep them afterwards ensure you Copy() them before passing to
// the channel
func (jq *Jq) Start(program string, args *Jv) (in chan<- *Jv, out <-chan *Jv, errs <-chan error) {
	// Create out two way copy of the channels. We need to be able to recv from
	// input, so need to store the original channel
	cIn := make(chan *Jv)
	cOut := make(chan *Jv)
	cErr := make(chan error)

	// And assign the read/write only versions to the output fars
	in = cIn
	out = cOut
	errs = cErr

	// Before setting up any of the global error handling state, lets check that
	// args is of the right type!
	if args.Kind() != JV_KIND_ARRAY {
		go func() {
			// Take ownership of the inputs
			for jv := range cIn {
				jv.Free()
			}
			cErr <- fmt.Errorf("`args` parameter is of type %s not array!", args.Kind().String())
			args.Free()
			close(cOut)
			close(cErr)
		}()
		return
	}

	if jq.errorStoreId != 0 {
		// We might have called Compile
		globalErrorChannels.Delete(jq.errorStoreId)
	}
	jq.errorStoreId = globalErrorChannels.Add(cErr)

	// Because we can't pass a function pointer to an exported Go func we have to
	// call a C function which uses the exported fund for us.
	// https://github.com/golang/go/wiki/cgo#function-variables
	C.install_jq_error_cb(jq._state, C.ulonglong(jq.errorStoreId))

	jq.running.Add(1)
	go func() {

		if jq._Compile(program, args) == false {
			// Even if compile failed follow the contract. Read any inputs and take
			// ownership of them (aka free them)
			//
			// Errors from compile will be sent to the error channel
			for jv := range cIn {
				jv.Free()
			}
		} else {
			for jv := range cIn {
				results, err := jq.Execute(jv)
				for _, result := range results {
					cOut <- result
				}
				if err != nil {
					cErr <- err
				}
			}
		}
		// Once we've read all the inputs close the output to signal to caller that
		// we are done.
		close(cOut)
		close(cErr)
		C.install_jq_error_cb(jq._state, 0)
		jq.running.Done()
	}()

	return
}

// Execute will run the Compiled() program against a single input and return
// the results.
//
// Using this interface directly is not thread-safe -- it is up to the caller to
// ensure that this is not called from two goroutines concurrently.
func (jq *Jq) Execute(input *Jv) (results []*Jv, err error) {
	flags := C.int(0)
	results = make([]*Jv, 0)

	C.jq_start(jq._state, input.jv, flags)
	result := &Jv{C.jq_next(jq._state)}
	for result.IsValid() {
		results = append(results, result)
		result = &Jv{C.jq_next(jq._state)}
	}
	msg, ok := result.GetInvalidMessageAsString()
	if ok {
		// Uncaught jq exception
		// TODO: get file:line position in input somehow.
		err = errors.New(msg)
	}

	return
}

// Compile the program and make it ready to Execute()
//
// Only a single program can be compiled on a Jq object at once. Calling this
// again a second time will replace the current program.
//
// args is a list of key/value pairs to bind as variables into the program, and
// must be an array type even if empty. Each element of the array should be an
// object with a "name" and "value" properties. Name should exclude the "$"
// sign. For example this is `[ {"name": "n", "value": 1 } ]` would then be
// `$n` in the program.
func (jq *Jq) Compile(prog string, args *Jv) (errs []error) {

	// Before setting up any of the global error handling state, lets check that
	// args is of the right type!
	if args.Kind() != JV_KIND_ARRAY {
		args.Free()
		return []error{fmt.Errorf("`args` parameter is of type %s not array", args.Kind().String())}
	}

	cErr := make(chan error)

	if jq.errorStoreId != 0 {
		// We might have called Compile
		globalErrorChannels.Delete(jq.errorStoreId)
	}
	jq.errorStoreId = globalErrorChannels.Add(cErr)

	C.install_jq_error_cb(jq._state, C.ulonglong(jq.errorStoreId))
	defer C.install_jq_error_cb(jq._state, 0)
	var wg sync.WaitGroup

	wg.Add(1)
	go func() {
		for err := range cErr {
			if err == nil {
				break
			}
			errs = append(errs, err)
		}
		wg.Done()
	}()

	compiled := jq._Compile(prog, args)
	cErr <- nil // Sentinel to break the loop above

	wg.Wait()
	globalErrorChannels.Delete(jq.errorStoreId)
	jq.errorStoreId = 0

	if !compiled && len(errs) == 0 {
		return []error{fmt.Errorf("jq_compile returned error, but no errors were reported. Oops")}
	}
	return errs
}

func (jq *Jq) _Compile(prog string, args *Jv) bool {
	cs := C.CString(prog)
	defer C.free(unsafe.Pointer(cs))

	// If there was an error it will have been sent to errorChannel via the
	// installed error handler
	return C.jq_compile_args(jq._state, cs, args.jv) != 0
}
<commit_msg>Fix shadowing of jv in generated C for go 1.15 (#15)<commit_after>// Package jq provides go bindings for libjq providing a streaming filter of
// JSON documents.
//
// This package provides a thin layer on top of stedolan's libjq -- it would
// likely be helpful to read through the wiki pages about it:
//
// jv: the JSON value type https://github.com/stedolan/jq/wiki/C-API:-jv
//
// libjq: https://github.com/stedolan/jq/wiki/C-API:-libjq
package jq

/*
To install
$ ./configure --disable-maintainer-mode --prefix=$PWD/BUILD
$ make install-libLTLIBRARIES install-includeHEADERS
*/

/*
#cgo LDFLAGS: -ljq
#cgo linux LDFLAGS: -lm

#include <jq.h>
#include <jv.h>

#include <stdlib.h>

void install_jq_error_cb(jq_state *jq, unsigned long long id);
*/
import "C"
import (
	"errors"
	"fmt"
	"sync"
	"sync/atomic"
	"unsafe"
)

// Jq encapsulates the state needed to interface with the libjq C library
type Jq struct {
	_state       *C.struct_jq_state
	errorStoreId uint64
	running      sync.WaitGroup
}

// New initializes a new JQ object and the underlying C library.
func New() (*Jq, error) {
	jq := new(Jq)

	var err error
	jq._state, err = C.jq_init()

	if err != nil {
		return nil, err
	} else if jq == nil {
		return nil, errors.New("jq_init returned nil -- out of memory?")
	}

	return jq, nil
}

// Close the handle to libjq and free C resources.
//
// If Start() has been called this will block until the input Channel it
// returns has been closed.
func (jq *Jq) Close() {
	// If the goroutine from Start() is running we need to make sure it finished cleanly
	// Wait until we aren't running before freeing C things.
	//
	jq.running.Wait()
	if jq._state != nil {
		C.jq_teardown(&jq._state)
		jq._state = nil
	}
	if jq.errorStoreId != 0 {
		globalErrorChannels.Delete(jq.errorStoreId)
		jq.errorStoreId = 0
	}
}

// We cant pass many things over the Go/C boundary, so instead of passing the error channel we pass an opaque indentifier (a 64bit int as it turns out) and use that to look up in a global variable
type errorLookupState struct {
	sync.RWMutex
	idCounter uint64
	channels  map[uint64]chan<- error
}

func (e *errorLookupState) Add(c chan<- error) uint64 {
	newID := atomic.AddUint64(&e.idCounter, 1)
	e.RWMutex.Lock()
	defer e.RWMutex.Unlock()
	e.channels[newID] = c
	return newID
}

func (e *errorLookupState) Get(id uint64) chan<- error {
	e.RWMutex.RLock()
	defer e.RWMutex.RUnlock()
	c, ok := e.channels[id]
	if !ok {
		panic(fmt.Sprintf("Tried to get error channel #%d out of store but it wasn't there!", id))
	}
	return c
}

func (e *errorLookupState) Delete(id uint64) {
	e.RWMutex.Lock()
	defer e.RWMutex.Unlock()
	delete(e.channels, id)
}

// The global state - this also serves to keep the channel in scope by keeping
// a reference to it that the GC can see
var globalErrorChannels = errorLookupState{
	channels: make(map[uint64]chan<- error),
}

//export goLibjqErrorHandler
func goLibjqErrorHandler(id uint64, value C.jv) {
	ch := globalErrorChannels.Get(id)

	err := _ConvertError(value)
	ch <- err
}

// Start will compile `program` and return a three channels: input, output and
// error. Sending a jq.Jv* to input cause the program to be run to it and
// one-or-more results returned as jq.Jv* on the output channel, or one or more
// error values sent to the error channel. When you are done sending values
// close the input channel.
//
// args is a list of key/value pairs to bind as variables into the program, and
// must be an array type even if empty. Each element of the array should be an
// object with a "name" and "value" properties. Name should exclude the "$"
// sign. For example this is `[ {"name": "n", "value": 1 } ]` would then be
// `$n` in the programm.
//
// This function is not reentereant -- in that you cannot and should not call
// Start again until you have closed the previous input channel.
//
// If there is a problem compiling the JQ program then the errors will be
// reported on error channel before any input is read so makle sure you account
// for this case.
//
// Any jq.Jv* values passed to the input channel will be owned by the channel.
// If you want to keep them afterwards ensure you Copy() them before passing to
// the channel
func (jq *Jq) Start(program string, args *Jv) (in chan<- *Jv, out <-chan *Jv, errs <-chan error) {
	// Create out two way copy of the channels. We need to be able to recv from
	// input, so need to store the original channel
	cIn := make(chan *Jv)
	cOut := make(chan *Jv)
	cErr := make(chan error)

	// And assign the read/write only versions to the output fars
	in = cIn
	out = cOut
	errs = cErr

	// Before setting up any of the global error handling state, lets check that
	// args is of the right type!
	if args.Kind() != JV_KIND_ARRAY {
		go func() {
			// Take ownership of the inputs
			for jv := range cIn {
				jv.Free()
			}
			cErr <- fmt.Errorf("`args` parameter is of type %s not array!", args.Kind().String())
			args.Free()
			close(cOut)
			close(cErr)
		}()
		return
	}

	if jq.errorStoreId != 0 {
		// We might have called Compile
		globalErrorChannels.Delete(jq.errorStoreId)
	}
	jq.errorStoreId = globalErrorChannels.Add(cErr)

	// Because we can't pass a function pointer to an exported Go func we have to
	// call a C function which uses the exported fund for us.
	// https://github.com/golang/go/wiki/cgo#function-variables
	C.install_jq_error_cb(jq._state, C.ulonglong(jq.errorStoreId))

	jq.running.Add(1)
	go func() {

		if jq._Compile(program, args) == false {
			// Even if compile failed follow the contract. Read any inputs and take
			// ownership of them (aka free them)
			//
			// Errors from compile will be sent to the error channel
			for jv := range cIn {
				jv.Free()
			}
		} else {
			for jv := range cIn {
				results, err := jq.Execute(jv)
				for _, result := range results {
					cOut <- result
				}
				if err != nil {
					cErr <- err
				}
			}
		}
		// Once we've read all the inputs close the output to signal to caller that
		// we are done.
		close(cOut)
		close(cErr)
		C.install_jq_error_cb(jq._state, 0)
		jq.running.Done()
	}()

	return
}

// Execute will run the Compiled() program against a single input and return
// the results.
//
// Using this interface directly is not thread-safe -- it is up to the caller to
// ensure that this is not called from two goroutines concurrently.
func (jq *Jq) Execute(input *Jv) (results []*Jv, err error) {
	flags := C.int(0)
	results = make([]*Jv, 0)

	C.jq_start(jq._state, input.jv, flags)
	result := &Jv{C.jq_next(jq._state)}
	for result.IsValid() {
		results = append(results, result)
		result = &Jv{C.jq_next(jq._state)}
	}
	msg, ok := result.GetInvalidMessageAsString()
	if ok {
		// Uncaught jq exception
		// TODO: get file:line position in input somehow.
		err = errors.New(msg)
	}

	return
}

// Compile the program and make it ready to Execute()
//
// Only a single program can be compiled on a Jq object at once. Calling this
// again a second time will replace the current program.
//
// args is a list of key/value pairs to bind as variables into the program, and
// must be an array type even if empty. Each element of the array should be an
// object with a "name" and "value" properties. Name should exclude the "$"
// sign. For example this is `[ {"name": "n", "value": 1 } ]` would then be
// `$n` in the program.
func (jq *Jq) Compile(prog string, args *Jv) (errs []error) {

	// Before setting up any of the global error handling state, lets check that
	// args is of the right type!
	if args.Kind() != JV_KIND_ARRAY {
		args.Free()
		return []error{fmt.Errorf("`args` parameter is of type %s not array", args.Kind().String())}
	}

	cErr := make(chan error)

	if jq.errorStoreId != 0 {
		// We might have called Compile
		globalErrorChannels.Delete(jq.errorStoreId)
	}
	jq.errorStoreId = globalErrorChannels.Add(cErr)

	C.install_jq_error_cb(jq._state, C.ulonglong(jq.errorStoreId))
	defer C.install_jq_error_cb(jq._state, 0)
	var wg sync.WaitGroup

	wg.Add(1)
	go func() {
		for err := range cErr {
			if err == nil {
				break
			}
			errs = append(errs, err)
		}
		wg.Done()
	}()

	compiled := jq._Compile(prog, args)
	cErr <- nil // Sentinel to break the loop above

	wg.Wait()
	globalErrorChannels.Delete(jq.errorStoreId)
	jq.errorStoreId = 0

	if !compiled && len(errs) == 0 {
		return []error{fmt.Errorf("jq_compile returned error, but no errors were reported. Oops")}
	}
	return errs
}

func (jq *Jq) _Compile(prog string, args *Jv) bool {
	cs := C.CString(prog)
	defer C.free(unsafe.Pointer(cs))

	// If there was an error it will have been sent to errorChannel via the
	// installed error handler
	return C.jq_compile_args(jq._state, cs, args.jv) != 0
}
<|endoftext|>


<commit_before>package jq

import (
	"bytes"
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"runtime"
	"strings"

	"github.com/jingweno/jqpipe-go"
)

var Path, Version string

func init() {
	var err error

	Path, err = setupJQPath()
	if err != nil {
		log.Fatal(err)
	}

	Version, err = jqVersion()
	if err != nil {
		log.Fatal(err)
	}
}

func setupJQPath() (string, error) {
	pwd, err := os.Getwd()
	if err != nil {
		return "", err
	}

	jqPath := filepath.Join(pwd, "bin", fmt.Sprintf("%s_%s", runtime.GOOS, runtime.GOARCH))
	os.Setenv("PATH", fmt.Sprintf("%s%c%s", jqPath, os.PathListSeparator, os.Getenv("PATH")))

	return filepath.Join(jqPath, "jq"), nil
}

func jqVersion() (string, error) {
	// get version from `jq --help`
	// since `jq --version` diffs between versions
	// e.g., 1.3 & 1.4
	var b bytes.Buffer
	cmd := exec.Command(Path, "--help")
	cmd.Stdout = &b
	cmd.Stderr = &b
	cmd.Run()

	out := bytes.TrimSpace(b.Bytes())
	r := regexp.MustCompile(`\[version (.+)\]`)
	if r.Match(out) {
		m := r.FindSubmatch(out)[1]
		return string(m), nil
	}

	return "", fmt.Errorf("can't find jq version: %s", out)
}

type JQ struct {
	J string          `json:"j"`
	Q string          `json:"q"`
	O map[string]bool `json:"o"`
}

func (j *JQ) Opts() []string {
	opts := []string{}
	for opt, enabled := range j.O {
		if enabled {
			opts = append(opts, fmt.Sprintf("--%s", opt))
		}
	}

	return opts
}

func (j *JQ) Eval() (string, error) {
	seq, err := jq.Eval(j.J, j.Q, j.Opts()...)
	if err != nil {
		return "", err
	}

	result := []string{}
	for _, s := range seq {
		ss := string(s)
		if ss != "" && ss != "null" {
			result = append(result, ss)
		}
	}

	return strings.Join(result, "\n"), nil
}

func (j *JQ) Validate() error {
	errMsgs := []string{}

	if j.Q == "" {
		errMsgs = append(errMsgs, "missing filter")
	}

	if j.J == "" {
		errMsgs = append(errMsgs, "missing JSON")
	}

	if len(errMsgs) > 0 {
		return fmt.Errorf("invalid input: %s", strings.Join(errMsgs, " and "))
	}

	return nil
}

func (j JQ) String() string {
	return fmt.Sprintf("j=%s, q=%s, o=%v", j.J, j.Q, j.Opts())
}
<commit_msg>Timing out `jq` execution.<commit_after>package jq

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"runtime"
	"strings"
	"time"

	"github.com/jingweno/jqpipe-go"
)

const jqExecTimeout = 10

var Path, Version string

func init() {
	var err error

	Path, err = setupJQPath()
	if err != nil {
		log.Fatal(err)
	}

	Version, err = jqVersion()
	if err != nil {
		log.Fatal(err)
	}
}

func setupJQPath() (string, error) {
	pwd, err := os.Getwd()
	if err != nil {
		return "", err
	}

	jqPath := filepath.Join(pwd, "bin", fmt.Sprintf("%s_%s", runtime.GOOS, runtime.GOARCH))
	os.Setenv("PATH", fmt.Sprintf("%s%c%s", jqPath, os.PathListSeparator, os.Getenv("PATH")))

	return filepath.Join(jqPath, "jq"), nil
}

func jqVersion() (string, error) {
	// get version from `jq --help`
	// since `jq --version` diffs between versions
	// e.g., 1.3 & 1.4
	var b bytes.Buffer
	cmd := exec.Command(Path, "--help")
	cmd.Stdout = &b
	cmd.Stderr = &b
	cmd.Run()

	out := bytes.TrimSpace(b.Bytes())
	r := regexp.MustCompile(`\[version (.+)\]`)
	if r.Match(out) {
		m := r.FindSubmatch(out)[1]
		return string(m), nil
	}

	return "", fmt.Errorf("can't find jq version: %s", out)
}

type jqResult struct {
	Seq []json.RawMessage
	Err error
}

func (r *jqResult) Result() (string, error) {
	if r.Err != nil {
		return "", r.Err
	}

	result := []string{}
	for _, s := range r.Seq {
		ss := string(s)
		if ss != "" && ss != "null" {
			result = append(result, ss)
		}
	}

	return strings.Join(result, "\n"), nil
}

type JQ struct {
	J string          `json:"j"`
	Q string          `json:"q"`
	O map[string]bool `json:"o"`
}

func (j *JQ) Opts() []string {
	opts := []string{}
	for opt, enabled := range j.O {
		if enabled {
			opts = append(opts, fmt.Sprintf("--%s", opt))
		}
	}

	return opts
}

// eval `jq` expression with timeout support
func (j *JQ) Eval() (string, error) {
	resultCh := make(chan jqResult, 1)
	go func(js, expr string, opts ...string) {
		seq, err := eval(js, expr, opts...)
		resultCh <- jqResult{seq, err}
	}(j.J, j.Q, j.Opts()...)

	select {
	case r := <-resultCh:
		return r.Result()
	case <-time.After(time.Second * jqExecTimeout):
		return "", fmt.Errorf("jq execution timeout")
	}
}

func (j *JQ) Validate() error {
	errMsgs := []string{}

	if j.Q == "" {
		errMsgs = append(errMsgs, "missing filter")
	}

	if j.J == "" {
		errMsgs = append(errMsgs, "missing JSON")
	}

	if len(errMsgs) > 0 {
		return fmt.Errorf("invalid input: %s", strings.Join(errMsgs, " and "))
	}

	return nil
}

func (j JQ) String() string {
	return fmt.Sprintf("j=%s, q=%s, o=%v", j.J, j.Q, j.Opts())
}

// eval `jq` expression
func eval(js string, expr string, opts ...string) ([]json.RawMessage, error) {
	jq, err := jq.New(bytes.NewReader([]byte(js)), expr, opts...)
	if err != nil {
		return nil, err
	}

	ret := make([]json.RawMessage, 0, 16)
	for {
		next, err := jq.Next()
		switch err {
		case nil:
			ret = append(ret, next)
		case io.EOF:
			return ret, nil
		default:
			return ret, err
		}
	}
}
<|endoftext|>


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

package stats

import (
	"errors"
	"fmt"
	"time"

	"github.com/census-instrumentation/opencensus-go/tags"
	"golang.org/x/net/context"
)

type worker struct {
	measuresByName map[string]Measure
	measures       map[Measure]bool
	viewsByName    map[string]View
	views          map[View]bool

	timer      *time.Ticker
	c          chan command
	quit, done chan bool
}

var defaultWorker *worker

var defaultReportingDuration = 10 * time.Second

// NewMeasureFloat64 creates a new measure of type MeasureFloat64. It returns
// an error if a measure with the same name already exists.
func NewMeasureFloat64(name, description, unit string) (*MeasureFloat64, error) {
	m := &MeasureFloat64{
		name:        name,
		description: description,
		unit:        unit,
		views:       make(map[View]bool),
	}

	req := &registerMeasureReq{
		m:   m,
		err: make(chan error),
	}
	defaultWorker.c <- req
	if err := <-req.err; err != nil {
		return nil, err
	}

	return m, nil
}

// NewMeasureInt64 creates a new measure of type MeasureInt64. It returns an
// error if a measure with the same name already exists.
func NewMeasureInt64(name, description, unit string) (*MeasureInt64, error) {
	m := &MeasureInt64{
		name:        name,
		description: description,
		unit:        unit,
		views:       make(map[View]bool),
	}

	req := &registerMeasureReq{
		m:   m,
		err: make(chan error),
	}
	defaultWorker.c <- req
	if err := <-req.err; err != nil {
		return nil, err
	}

	return m, nil
}

// GetMeasureByName returns the registered measure associated with name.
func GetMeasureByName(name string) (Measure, error) {
	req := &getMeasureByNameReq{
		name: name,
		c:    make(chan *getMeasureByNameResp),
	}
	defaultWorker.c <- req
	resp := <-req.c
	return resp.m, resp.err
}

// DeleteMeasure deletes an existing measure to allow for creation of a new
// measure with the same name. It returns an error if the measure cannot be
// deleted (if one or multiple registered views refer to it).
func DeleteMeasure(m Measure) error {
	req := &deleteMeasureReq{
		m:   m,
		err: make(chan error),
	}
	defaultWorker.c <- req
	return <-req.err
}

// GetViewByName returns the registered view associated with this name.
func GetViewByName(name string) (View, error) {
	req := &getViewByNameReq{
		name: name,
		c:    make(chan *getViewByNameResp),
	}
	defaultWorker.c <- req
	resp := <-req.c
	return resp.v, resp.err
}

// RegisterView registers view. It returns an error if the view cannot be
// registered. Subsequent calls to Record with the same measure as the one in
// the view will NOT cause the usage to be recorded unless a consumer is
// subscribed to the view or ForceCollection for this view is called.
func RegisterView(v View) error {
	if v == nil {
		return errors.New("cannot RegisterView for nil view")
	}

	req := &registerViewReq{
		v:   v,
		err: make(chan error),
	}
	defaultWorker.c <- req
	return <-req.err
}

// UnregisterView deletes the previously registered view. It returns an error
// if the view wasn't registered. All data collected and not reported for the
// corresponding view will be lost. All clients subscribed to this view are
// unsubscribed automatically and their subscriptions channels closed.
func UnregisterView(v View) error {
	if v == nil {
		return errors.New("cannot UnregisterView for nil view")
	}

	req := &unregisterViewReq{
		v:   v,
		err: make(chan error),
	}
	defaultWorker.c <- req
	return <-req.err
}

// SubscribeToView subscribes a client to a View. If the view wasn't already
// registered, it will be automatically registered. It allows for many clients
// to consume the same ViewData with a single registration. -i.e. the aggregate
// of the collected measurements will be reported to the calling code through
// channel c. To avoid data loss, clients must ensure that channel sends
// proceed in a timely manner. The calling code is responsible for using a
// buffered channel or blocking on the channel waiting for the collected data.
func SubscribeToView(v View, c chan *ViewData) error {
	if v == nil {
		return errors.New("cannot SubscribeToView for nil view")
	}

	req := &subscribeToViewReq{
		v:   v,
		c:   c,
		err: make(chan error),
	}
	defaultWorker.c <- req
	return <-req.err
}

// UnsubscribeFromView unsubscribes a previously subscribed channel from the
// View subscriptions. If no more subscriber for v exists and the the ad hoc
// collection for this view isn't active, data stops being collected for this
// view.
func UnsubscribeFromView(v View, c chan *ViewData) error {
	if v == nil {
		return errors.New("cannot UnsubscribeFromView for nil view")
	}

	req := &unsubscribeFromViewReq{
		v:   v,
		c:   c,
		err: make(chan error),
	}
	defaultWorker.c <- req
	return <-req.err
}

// ForceCollection starts data collection for this view even if no
// listeners are subscribed to it.
func ForceCollection(v View) error {
	if v == nil {
		return errors.New("cannot ForceCollection for nil view")
	}

	req := &startForcedCollectionReq{
		v:   v,
		err: make(chan error),
	}
	defaultWorker.c <- req
	return <-req.err
}

// StopForcedCollection stops data collection for this view unless at least
// 1 listener is subscribed to it.
func StopForcedCollection(v View) error {
	if v == nil {
		return errors.New("cannot StopForcedCollection for nil view")
	}

	req := &stopForcedCollectionReq{
		v:   v,
		err: make(chan error),
	}
	defaultWorker.c <- req
	return <-req.err
}

// RetrieveData returns the current collected data for the view.
func RetrieveData(v View) ([]*Row, error) {
	if v == nil {
		return nil, errors.New("cannot retrieve data for nil view")
	}
	req := &retrieveDataReq{
		now: time.Now(),
		v:   v,
		c:   make(chan *retrieveDataResp),
	}
	defaultWorker.c <- req
	resp := <-req.c
	return resp.rows, resp.err
}

// RecordFloat64 records a float64 value against a measure and the tags passed
// as part of the context.
func RecordFloat64(ctx context.Context, mf *MeasureFloat64, v float64) {
	req := &recordFloat64Req{
		now: time.Now(),
		ts:  tags.FromContext(ctx),
		mf:  mf,
		v:   v,
	}
	defaultWorker.c <- req
}

// RecordInt64 records an int64 value against a measure and the tags passed as
// part of the context.
func RecordInt64(ctx context.Context, mi *MeasureInt64, v int64) {
	req := &recordInt64Req{
		now: time.Now(),
		ts:  tags.FromContext(ctx),
		mi:  mi,
		v:   v,
	}
	defaultWorker.c <- req
}

// Record records one or multiple measurements with the same tags at once.
func Record(ctx context.Context, ms ...Measurement) {
	req := &recordReq{
		now: time.Now(),
		ts:  tags.FromContext(ctx),
		ms:  ms,
	}
	defaultWorker.c <- req
}

// SetReportingPeriod sets the interval between reporting aggregated views in
// the program. Calling SetReportingPeriod with duration argument less than or
// equal to zero enables the default behavior.
func SetReportingPeriod(d time.Duration) {
	// TODO(acetechnologist): ensure that the duration d is more than a certain
	// value. e.g. 1s
	req := &setReportingPeriodReq{
		d: d,
		c: make(chan bool),
	}
	defaultWorker.c <- req
	<-req.c // don't return until the timer is set to the new duration.
}

func init() {
	defaultWorker = newWorker()
	go defaultWorker.start()
}

func newWorker() *worker {
	return &worker{
		measuresByName: make(map[string]Measure),
		measures:       make(map[Measure]bool),
		viewsByName:    make(map[string]View),
		views:          make(map[View]bool),
		timer:          time.NewTicker(defaultReportingDuration),
		c:              make(chan command),
		quit:           make(chan bool),
		done:           make(chan bool),
	}
}

func (w *worker) start() {
	for {
		select {
		case cmd := <-w.c:
			if cmd != nil {
				cmd.handleCommand(w)
			}
		case <-w.timer.C:
			w.reportUsage(time.Now())
		case <-w.quit:
			w.timer.Stop()
			close(w.c)
			w.done <- true
			return
		}
	}
}

func (w *worker) stop() {
	w.quit <- true
	_ = <-w.done
}

func (w *worker) tryRegisterMeasure(m Measure) error {
	if x, ok := w.measuresByName[m.Name()]; ok {
		if x != m {
			return fmt.Errorf("cannot register the measure with name '%v' because a different measure with the same name is already registered", m.Name())
		}

		// the measure is already registered so there is nothing to do and the
		// command is considered successful.
		return nil
	}

	w.measuresByName[m.Name()] = m
	w.measures[m] = true
	return nil
}

func (w *worker) tryRegisterView(v View) error {
	if x, ok := w.viewsByName[v.Name()]; ok {
		if x != v {
			return fmt.Errorf("cannot register the view with name '%v' because a different view with the same name is already registered", v.Name())
		}

		// the view is already registered so there is nothing to do and the
		// command is considered successful.
		return nil
	}

	// view is not registered and needs to be registered, but first its measure
	// needs to be registered.
	if err := w.tryRegisterMeasure(v.Measure()); err != nil {
		return fmt.Errorf("%v. Hence cannot register view '%v,", err, v.Name())
	}

	w.viewsByName[v.Name()] = v
	w.views[v] = true
	v.Measure().addView(v)
	return nil
}

func (w *worker) reportUsage(now time.Time) {
	for v := range w.views {
		if v.subscriptionsCount() == 0 {
			continue
		}

		viewData := &ViewData{
			V:    v,
			Rows: v.collectedRows(now),
		}

		for c, s := range v.subscriptions() {
			select {
			case c <- viewData:
				return
			default:
				s.droppedViewData++
			}
		}

		if _, ok := v.Window().(*WindowCumulative); !ok {
			v.clearRows()
		}
	}
}

// RestartWorker is used for testing only. It stops the old worker and creates
// a new worker. It should never be called by production code.
func RestartWorker() {
	defaultWorker.stop()
	defaultWorker = newWorker()
	go defaultWorker.start()
}
<commit_msg>Move init function to the top (#37)<commit_after>// Copyright 2017, OpenCensus Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

package stats

import (
	"errors"
	"fmt"
	"time"

	"github.com/census-instrumentation/opencensus-go/tags"
	"golang.org/x/net/context"
)

func init() {
	defaultWorker = newWorker()
	go defaultWorker.start()
}

type worker struct {
	measuresByName map[string]Measure
	measures       map[Measure]bool
	viewsByName    map[string]View
	views          map[View]bool

	timer      *time.Ticker
	c          chan command
	quit, done chan bool
}

var defaultWorker *worker

var defaultReportingDuration = 10 * time.Second

// NewMeasureFloat64 creates a new measure of type MeasureFloat64. It returns
// an error if a measure with the same name already exists.
func NewMeasureFloat64(name, description, unit string) (*MeasureFloat64, error) {
	m := &MeasureFloat64{
		name:        name,
		description: description,
		unit:        unit,
		views:       make(map[View]bool),
	}

	req := &registerMeasureReq{
		m:   m,
		err: make(chan error),
	}
	defaultWorker.c <- req
	if err := <-req.err; err != nil {
		return nil, err
	}

	return m, nil
}

// NewMeasureInt64 creates a new measure of type MeasureInt64. It returns an
// error if a measure with the same name already exists.
func NewMeasureInt64(name, description, unit string) (*MeasureInt64, error) {
	m := &MeasureInt64{
		name:        name,
		description: description,
		unit:        unit,
		views:       make(map[View]bool),
	}

	req := &registerMeasureReq{
		m:   m,
		err: make(chan error),
	}
	defaultWorker.c <- req
	if err := <-req.err; err != nil {
		return nil, err
	}

	return m, nil
}

// GetMeasureByName returns the registered measure associated with name.
func GetMeasureByName(name string) (Measure, error) {
	req := &getMeasureByNameReq{
		name: name,
		c:    make(chan *getMeasureByNameResp),
	}
	defaultWorker.c <- req
	resp := <-req.c
	return resp.m, resp.err
}

// DeleteMeasure deletes an existing measure to allow for creation of a new
// measure with the same name. It returns an error if the measure cannot be
// deleted (if one or multiple registered views refer to it).
func DeleteMeasure(m Measure) error {
	req := &deleteMeasureReq{
		m:   m,
		err: make(chan error),
	}
	defaultWorker.c <- req
	return <-req.err
}

// GetViewByName returns the registered view associated with this name.
func GetViewByName(name string) (View, error) {
	req := &getViewByNameReq{
		name: name,
		c:    make(chan *getViewByNameResp),
	}
	defaultWorker.c <- req
	resp := <-req.c
	return resp.v, resp.err
}

// RegisterView registers view. It returns an error if the view cannot be
// registered. Subsequent calls to Record with the same measure as the one in
// the view will NOT cause the usage to be recorded unless a consumer is
// subscribed to the view or ForceCollection for this view is called.
func RegisterView(v View) error {
	if v == nil {
		return errors.New("cannot RegisterView for nil view")
	}

	req := &registerViewReq{
		v:   v,
		err: make(chan error),
	}
	defaultWorker.c <- req
	return <-req.err
}

// UnregisterView deletes the previously registered view. It returns an error
// if the view wasn't registered. All data collected and not reported for the
// corresponding view will be lost. All clients subscribed to this view are
// unsubscribed automatically and their subscriptions channels closed.
func UnregisterView(v View) error {
	if v == nil {
		return errors.New("cannot UnregisterView for nil view")
	}

	req := &unregisterViewReq{
		v:   v,
		err: make(chan error),
	}
	defaultWorker.c <- req
	return <-req.err
}

// SubscribeToView subscribes a client to a View. If the view wasn't already
// registered, it will be automatically registered. It allows for many clients
// to consume the same ViewData with a single registration. -i.e. the aggregate
// of the collected measurements will be reported to the calling code through
// channel c. To avoid data loss, clients must ensure that channel sends
// proceed in a timely manner. The calling code is responsible for using a
// buffered channel or blocking on the channel waiting for the collected data.
func SubscribeToView(v View, c chan *ViewData) error {
	if v == nil {
		return errors.New("cannot SubscribeToView for nil view")
	}

	req := &subscribeToViewReq{
		v:   v,
		c:   c,
		err: make(chan error),
	}
	defaultWorker.c <- req
	return <-req.err
}

// UnsubscribeFromView unsubscribes a previously subscribed channel from the
// View subscriptions. If no more subscriber for v exists and the the ad hoc
// collection for this view isn't active, data stops being collected for this
// view.
func UnsubscribeFromView(v View, c chan *ViewData) error {
	if v == nil {
		return errors.New("cannot UnsubscribeFromView for nil view")
	}

	req := &unsubscribeFromViewReq{
		v:   v,
		c:   c,
		err: make(chan error),
	}
	defaultWorker.c <- req
	return <-req.err
}

// ForceCollection starts data collection for this view even if no
// listeners are subscribed to it.
func ForceCollection(v View) error {
	if v == nil {
		return errors.New("cannot ForceCollection for nil view")
	}

	req := &startForcedCollectionReq{
		v:   v,
		err: make(chan error),
	}
	defaultWorker.c <- req
	return <-req.err
}

// StopForcedCollection stops data collection for this view unless at least
// 1 listener is subscribed to it.
func StopForcedCollection(v View) error {
	if v == nil {
		return errors.New("cannot StopForcedCollection for nil view")
	}

	req := &stopForcedCollectionReq{
		v:   v,
		err: make(chan error),
	}
	defaultWorker.c <- req
	return <-req.err
}

// RetrieveData returns the current collected data for the view.
func RetrieveData(v View) ([]*Row, error) {
	if v == nil {
		return nil, errors.New("cannot retrieve data for nil view")
	}
	req := &retrieveDataReq{
		now: time.Now(),
		v:   v,
		c:   make(chan *retrieveDataResp),
	}
	defaultWorker.c <- req
	resp := <-req.c
	return resp.rows, resp.err
}

// RecordFloat64 records a float64 value against a measure and the tags passed
// as part of the context.
func RecordFloat64(ctx context.Context, mf *MeasureFloat64, v float64) {
	req := &recordFloat64Req{
		now: time.Now(),
		ts:  tags.FromContext(ctx),
		mf:  mf,
		v:   v,
	}
	defaultWorker.c <- req
}

// RecordInt64 records an int64 value against a measure and the tags passed as
// part of the context.
func RecordInt64(ctx context.Context, mi *MeasureInt64, v int64) {
	req := &recordInt64Req{
		now: time.Now(),
		ts:  tags.FromContext(ctx),
		mi:  mi,
		v:   v,
	}
	defaultWorker.c <- req
}

// Record records one or multiple measurements with the same tags at once.
func Record(ctx context.Context, ms ...Measurement) {
	req := &recordReq{
		now: time.Now(),
		ts:  tags.FromContext(ctx),
		ms:  ms,
	}
	defaultWorker.c <- req
}

// SetReportingPeriod sets the interval between reporting aggregated views in
// the program. Calling SetReportingPeriod with duration argument less than or
// equal to zero enables the default behavior.
func SetReportingPeriod(d time.Duration) {
	// TODO(acetechnologist): ensure that the duration d is more than a certain
	// value. e.g. 1s
	req := &setReportingPeriodReq{
		d: d,
		c: make(chan bool),
	}
	defaultWorker.c <- req
	<-req.c // don't return until the timer is set to the new duration.
}

func newWorker() *worker {
	return &worker{
		measuresByName: make(map[string]Measure),
		measures:       make(map[Measure]bool),
		viewsByName:    make(map[string]View),
		views:          make(map[View]bool),
		timer:          time.NewTicker(defaultReportingDuration),
		c:              make(chan command),
		quit:           make(chan bool),
		done:           make(chan bool),
	}
}

func (w *worker) start() {
	for {
		select {
		case cmd := <-w.c:
			if cmd != nil {
				cmd.handleCommand(w)
			}
		case <-w.timer.C:
			w.reportUsage(time.Now())
		case <-w.quit:
			w.timer.Stop()
			close(w.c)
			w.done <- true
			return
		}
	}
}

func (w *worker) stop() {
	w.quit <- true
	_ = <-w.done
}

func (w *worker) tryRegisterMeasure(m Measure) error {
	if x, ok := w.measuresByName[m.Name()]; ok {
		if x != m {
			return fmt.Errorf("cannot register the measure with name '%v' because a different measure with the same name is already registered", m.Name())
		}

		// the measure is already registered so there is nothing to do and the
		// command is considered successful.
		return nil
	}

	w.measuresByName[m.Name()] = m
	w.measures[m] = true
	return nil
}

func (w *worker) tryRegisterView(v View) error {
	if x, ok := w.viewsByName[v.Name()]; ok {
		if x != v {
			return fmt.Errorf("cannot register the view with name '%v' because a different view with the same name is already registered", v.Name())
		}

		// the view is already registered so there is nothing to do and the
		// command is considered successful.
		return nil
	}

	// view is not registered and needs to be registered, but first its measure
	// needs to be registered.
	if err := w.tryRegisterMeasure(v.Measure()); err != nil {
		return fmt.Errorf("%v. Hence cannot register view '%v,", err, v.Name())
	}

	w.viewsByName[v.Name()] = v
	w.views[v] = true
	v.Measure().addView(v)
	return nil
}

func (w *worker) reportUsage(now time.Time) {
	for v := range w.views {
		if v.subscriptionsCount() == 0 {
			continue
		}

		viewData := &ViewData{
			V:    v,
			Rows: v.collectedRows(now),
		}

		for c, s := range v.subscriptions() {
			select {
			case c <- viewData:
				return
			default:
				s.droppedViewData++
			}
		}

		if _, ok := v.Window().(*WindowCumulative); !ok {
			v.clearRows()
		}
	}
}

// RestartWorker is used for testing only. It stops the old worker and creates
// a new worker. It should never be called by production code.
func RestartWorker() {
	defaultWorker.stop()
	defaultWorker = newWorker()
	go defaultWorker.start()
}
<|endoftext|>


<commit_before>package statshub

import (
	"appengine"
	"appengine/user"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"github.com/garyburd/redigo/redis"
	"log"
	"net/http"
	"time"
)

var (
	redisConnectTimeout = 10 * time.Second
	redisReadTimeout    = 10 * time.Second
	redisWriteTimeout   = 10 * time.Second
)

type Stats struct {
	UserId      uint64
	Hash        string // sha256(real userid + userid)
	CountryCode string
	Counters    map[string]uint64
	Gauges      map[string]uint64
}

type Response struct {
	Succeeded bool
	Error     string
}

func init() {
	http.HandleFunc("/", handler)
}

func handler(w http.ResponseWriter, r *http.Request) {
	statusCode, err := doHandle(r)
	w.Header().Set("Content-Type", "application/json")
	response := Response{Succeeded: true}
	if err != nil {
		response.Succeeded = false
		response.Error = fmt.Sprintf("%s", err)
	}
	w.WriteHeader(statusCode)
	bytes, err := json.Marshal(&response)
	if err == nil {
		w.Write(bytes)
	}
	if err != nil {
		log.Printf("Unable to respond to client: %s", err)
	}
}

func doHandle(r *http.Request) (statusCode int, err error) {
	context := appengine.NewContext(r)
	user, err := user.CurrentOAuth(context, "")
	if err != nil {
		return 401, fmt.Errorf("Unable to authenticate: %s", err)
	}

	decoder := json.NewDecoder(r.Body)
	stats := &Stats{}
	err = decoder.Decode(stats)
	if err != nil {
		return 400, fmt.Errorf("Unable to decode stats: %s", err)
	}

	hasher := sha256.New()
	hasher.Reset()
	hashInput := fmt.Sprintf("%s%d", user.Email, stats.UserId)
	hasher.Write([]byte(hashInput))
	expectedHash := hex.EncodeToString(hasher.Sum(nil))

	if expectedHash != stats.Hash {
		return 403, fmt.Errorf("Hash mismatch, authentication failure")
	}

	conn, err := connectToRedis()
	if err != nil {
		return 500, fmt.Errorf("Unable to connect to redis: %s", err)
	}

	if err = postStats(conn, stats); err != nil {
		return 500, fmt.Errorf("Unable to post stats: %s", err)
	}

	return 200, nil
}

func connectToRedis() (conn redis.Conn, err error) {
	conn, err = redis.DialTimeout("tcp",
		redisAddr,
		redisConnectTimeout,
		redisReadTimeout,
		redisWriteTimeout,
	)
	if err != nil {
		return
	}
	_, err = conn.Do("AUTH", redisPassword)
	return
}

func postStats(conn redis.Conn, stats *Stats) (err error) {
	redisKeys := func(key string) []string {
		return []string{
			fmt.Sprintf("%d:%s", stats.UserId, key),
			fmt.Sprintf("%s:%s", stats.CountryCode, key),
			fmt.Sprintf("global:%s", key),
		}
	}

	for key, value := range stats.Counters {
		for _, redisKey := range redisKeys(key) {
			if err = conn.Send("INCRBY", redisKey, value); err != nil {
				return
			}
		}
	}

	for key, value := range stats.Gauges {
		for _, redisKey := range redisKeys(key) {
			if err = conn.Send("SET", redisKey, value); err != nil {
				return
			}
		}
	}

	conn.Flush()
	return
}
<commit_msg>Made more restful<commit_after>package statshub

import (
	"appengine"
	"appengine/user"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"github.com/garyburd/redigo/redis"
	"log"
	"net/http"
	"time"
)

var (
	redisConnectTimeout = 10 * time.Second
	redisReadTimeout    = 10 * time.Second
	redisWriteTimeout   = 10 * time.Second
)

type Stats struct {
	UserId      uint64
	Hash        string // sha256(real userid + userid)
	CountryCode string
	Counters    map[string]uint64
	Gauges      map[string]uint64
}

type Response struct {
	Succeeded bool
	Error     string
}

func init() {
	http.HandleFunc("/stats", statsHandler)
}

// statsHandler handles requests to do stuff with stats
func statsHandler(w http.ResponseWriter, r *http.Request) {
	if "POST" == r.Method {
		statusCode, err := postStats(r)
		w.Header().Set("Content-Type", "application/json")
		response := Response{Succeeded: true}
		if err != nil {
			response.Succeeded = false
			response.Error = fmt.Sprintf("%s", err)
		}
		w.WriteHeader(statusCode)
		bytes, err := json.Marshal(&response)
		if err == nil {
			w.Write(bytes)
		}
		if err != nil {
			log.Printf("Unable to respond to client: %s", err)
		}
	} else {
		w.WriteHeader(405)
	}
}

func postStats(r *http.Request) (statusCode int, err error) {
	context := appengine.NewContext(r)
	user, err := user.CurrentOAuth(context, "")
	if err != nil {
		return 401, fmt.Errorf("Unable to authenticate: %s", err)
	}

	decoder := json.NewDecoder(r.Body)
	stats := &Stats{}
	err = decoder.Decode(stats)
	if err != nil {
		return 400, fmt.Errorf("Unable to decode stats: %s", err)
	}

	if statusCode, err = checkHash(user, stats); err != nil {
		return
	}

	conn, err := connectToRedis()
	if err != nil {
		return 500, fmt.Errorf("Unable to connect to redis: %s", err)
	}

	if err = postStatsToRedis(conn, stats); err != nil {
		return 500, fmt.Errorf("Unable to post stats: %s", err)
	}

	return 200, nil
}

func checkHash(user *user.User, stats *Stats) (statusCode int, err error) {
	hasher := sha256.New()
	hasher.Reset()
	hashInput := fmt.Sprintf("%s%d", user.Email, stats.UserId)
	hasher.Write([]byte(hashInput))
	expectedHash := hex.EncodeToString(hasher.Sum(nil))

	if expectedHash != stats.Hash {
		return 403, fmt.Errorf("Hash mismatch, authentication failure")
	} else {
		return
	}
}

func connectToRedis() (conn redis.Conn, err error) {
	conn, err = redis.DialTimeout("tcp",
		redisAddr,
		redisConnectTimeout,
		redisReadTimeout,
		redisWriteTimeout,
	)
	if err != nil {
		return
	}
	_, err = conn.Do("AUTH", redisPassword)
	return
}

func postStatsToRedis(conn redis.Conn, stats *Stats) (err error) {
	redisKeys := func(key string) []string {
		return []string{
			fmt.Sprintf("%d:%s", stats.UserId, key),
			fmt.Sprintf("%s:%s", stats.CountryCode, key),
			fmt.Sprintf("global:%s", key),
		}
	}

	for key, value := range stats.Counters {
		for _, redisKey := range redisKeys(key) {
			if err = conn.Send("INCRBY", redisKey, value); err != nil {
				return
			}
		}
	}

	for key, value := range stats.Gauges {
		for _, redisKey := range redisKeys(key) {
			if err = conn.Send("SET", redisKey, value); err != nil {
				return
			}
		}
	}

	conn.Flush()
	return
}
<|endoftext|>


<commit_before>package storage

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"log"
	"net/http"
	"strings"

	"github.com/googleapis/google-cloud-go-testing/storage/stiface"
	"google.golang.org/api/googleapi"

	"github.com/m-lab/etl/etl"
	"github.com/m-lab/etl/factory"
	"github.com/m-lab/etl/metrics"
	"github.com/m-lab/etl/row"
)

// ObjectWriter creates a writer to a named object.
// It may overwrite an existing object.
// Caller must Close() the writer, or cancel the context.
func ObjectWriter(ctx context.Context, client stiface.Client, bucket string, path string) stiface.Writer {
	b := client.Bucket(bucket)
	o := b.Object(path)
	w := o.NewWriter(ctx)
	// Set smaller chunk size to conserve memory.
	w.SetChunkSize(4 * 1024 * 1024)
	return w
}

// RowWriter implements row.Sink to a GCS file backend.
type RowWriter struct {
	w stiface.Writer

	bucket string
	path   string

	// These act as tokens to serialize access to the writer.
	// This allows concurrent encoding and writing, while ensuring
	// that single client access is correctly ordered.
	encoding chan struct{} // Token required for metric updates.
	writing  chan struct{} // Token required for metric updates.
}

// NewRowWriter creates a RowWriter.
func NewRowWriter(ctx context.Context, client stiface.Client, bucket string, path string) (row.Sink, error) {
	w := ObjectWriter(ctx, client, bucket, path)
	encoding := make(chan struct{}, 1)
	encoding <- struct{}{}
	writing := make(chan struct{}, 1)
	writing <- struct{}{}

	return &RowWriter{bucket: bucket, path: path, w: w, encoding: encoding, writing: writing}, nil
}

// Acquire the encoding token.
// TODO can we allow two encoders, and still sequence the writing?
func (rw *RowWriter) acquireEncodingToken() {
	<-rw.encoding
}

func (rw *RowWriter) releaseEncodingToken() {
	if len(rw.encoding) > 0 {
		log.Println("token error")
		return
	}
	rw.encoding <- struct{}{}
}

// Swap the encoding token for the write token.
// MUST already hold the write token.
func (rw *RowWriter) swapForWritingToken() {
	<-rw.writing
	rw.releaseEncodingToken()
}

func (rw *RowWriter) releaseWritingToken() {
	rw.writing <- struct{}{} // return the token.
}

// Commit commits rows, in order, to the GCS object.
// The GCS object is not available until Close is called, at which
// point the entire object becomes available atomically.
// The returned int is the number of rows written (and pending), or,
// if error is not nil, an estimate of the number of rows written.
func (rw *RowWriter) Commit(rows []interface{}, label string) (int, error) {
	rw.acquireEncodingToken()
	// First, do the encoding.  Other calls to Commit will block here
	// until encoding is done.
	// NOTE: This can cause a fairly hefty memory footprint for
	// large numbers of large rows.
	buf := bytes.NewBuffer(nil)

	for i := range rows {
		j, err := json.Marshal(rows[i])
		if err != nil {
			rw.releaseEncodingToken()
			metrics.BackendFailureCount.WithLabelValues(
				label, "encoding error").Inc()
			return 0, err
		}
		metrics.RowSizeHistogram.WithLabelValues(label).Observe(float64(len(j)))
		buf.Write(j)
		buf.WriteByte('\n')
	}
	numBytes := buf.Len()
	rw.swapForWritingToken()
	defer rw.releaseWritingToken()
	n, err := buf.WriteTo(rw.w) // This is buffered (by 4MB chunks).  Are the writes to GCS synchronous?
	if err != nil {
		switch typedErr := err.(type) {
		case *googleapi.Error:
			metrics.BackendFailureCount.WithLabelValues(
				label, "googleapi.Error").Inc()
			log.Println(typedErr, rw.bucket, rw.path)
			for _, e := range typedErr.Errors {
				log.Println(e)
			}
		default:
			metrics.BackendFailureCount.WithLabelValues(
				label, "other error").Inc()
			log.Println(typedErr, rw.bucket, rw.path)
		}
		// This approximates the number of rows written prior to error.
		// It is unclear whether these rows will actually show up.
		// The caller should likely abandon the archive at this point,
		// as further writing will likely result in a corrupted file.
		// See https://github.com/m-lab/etl/issues/899
		return int(n) * len(rows) / numBytes, err
	}

	// TODO - these may not be committed, so the returned value may be wrong.
	return len(rows), nil
}

// Close synchronizes on the tokens, and closes the backing file.
func (rw *RowWriter) Close() error {
	// Take BOTH tokens, to ensure no other goroutines are still running.
	<-rw.encoding
	<-rw.writing

	close(rw.encoding)
	close(rw.writing)

	log.Println("Closing", rw.bucket, rw.path)
	err := rw.w.Close()
	if err != nil {
		log.Println(err)
	} else {
		log.Println(rw.w.Attrs())
	}
	return err
}

// SinkFactory implements factory.SinkFactory.
type SinkFactory struct {
	client       stiface.Client
	outputBucket string
}

// returns the full path/filename from a gs://bucket/path/filename string.
func pathAndFilename(uri string) (string, error) {
	parts := strings.SplitN(uri, "/", 4)
	if len(parts) != 4 || parts[0] != "gs:" || len(parts[3]) == 0 {
		return "", errors.New("Bad GCS path")
	}
	return parts[3], nil
}

// Get implements factory.SinkFactory
func (sf *SinkFactory) Get(ctx context.Context, path etl.DataPath) (row.Sink, etl.ProcessingError) {
	fn, err := pathAndFilename(path.URI)
	if err != nil {
		return nil, factory.NewError(path.DataType, "InvalidPath",
			http.StatusInternalServerError, err)
	}
	s, err := NewRowWriter(ctx, sf.client, sf.outputBucket, fn+".json")
	if err != nil {
		return nil, factory.NewError(path.DataType, "SinkFactory",
			http.StatusInternalServerError, err)
	}
	return s, nil
}

// NewSinkFactory returns the default SinkFactory
func NewSinkFactory(client stiface.Client, outputBucket string) factory.SinkFactory {
	return &SinkFactory{client: client, outputBucket: outputBucket}
}
<commit_msg>Add row count and error meta-data to json output files (#937)<commit_after>package storage

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net/http"
	"strings"
	"time"

	gcs "cloud.google.com/go/storage"
	"github.com/googleapis/google-cloud-go-testing/storage/stiface"
	"google.golang.org/api/googleapi"

	"github.com/m-lab/etl/etl"
	"github.com/m-lab/etl/factory"
	"github.com/m-lab/etl/metrics"
	"github.com/m-lab/etl/row"
)

// RowWriter implements row.Sink to a GCS file backend.
type RowWriter struct {
	w stiface.Writer
	o stiface.ObjectHandle
	a gcs.ObjectAttrsToUpdate

	rows     int
	writeErr error

	bucket string
	path   string

	// These act as tokens to serialize access to the writer.
	// This allows concurrent encoding and writing, while ensuring
	// that single client access is correctly ordered.
	encoding chan struct{} // Token required for encoding.
	writing  chan struct{} // Token required for writing.
}

// NewRowWriter creates a RowWriter.
func NewRowWriter(ctx context.Context, client stiface.Client, bucket string, path string) (row.Sink, error) {
	b := client.Bucket(bucket)
	o := b.Object(path)
	w := o.NewWriter(ctx)
	// Set smaller chunk size to conserve memory.
	w.SetChunkSize(4 * 1024 * 1024)

	encoding := make(chan struct{}, 1)
	encoding <- struct{}{}
	writing := make(chan struct{}, 1)
	writing <- struct{}{}

	return &RowWriter{bucket: bucket, path: path, o: o, w: w, encoding: encoding, writing: writing}, nil
}

// Acquire the encoding token.
// TODO can we allow two encoders, and still sequence the writing?
func (rw *RowWriter) acquireEncodingToken() {
	<-rw.encoding
}

func (rw *RowWriter) releaseEncodingToken() {
	if len(rw.encoding) > 0 {
		log.Println("token error")
		return
	}
	rw.encoding <- struct{}{}
}

// Swap the encoding token for the write token.
// MUST already hold the write token.
func (rw *RowWriter) swapForWritingToken() {
	<-rw.writing
	rw.releaseEncodingToken()
}

func (rw *RowWriter) releaseWritingToken() {
	rw.writing <- struct{}{} // return the token.
}

// Commit commits rows, in order, to the GCS object.
// The GCS object is not available until Close is called, at which
// point the entire object becomes available atomically.
// The returned int is the number of rows written (and pending), or,
// if error is not nil, an estimate of the number of rows written.
func (rw *RowWriter) Commit(rows []interface{}, label string) (int, error) {
	rw.acquireEncodingToken()
	// First, do the encoding.  Other calls to Commit will block here
	// until encoding is done.
	// NOTE: This can cause a fairly hefty memory footprint for
	// large numbers of large rows.
	buf := bytes.NewBuffer(nil)

	for i := range rows {
		j, err := json.Marshal(rows[i])
		if err != nil {
			rw.releaseEncodingToken()
			metrics.BackendFailureCount.WithLabelValues(
				label, "encoding error").Inc()
			return 0, err
		}
		metrics.RowSizeHistogram.WithLabelValues(label).Observe(float64(len(j)))
		buf.Write(j)
		buf.WriteByte('\n')
	}
	numBytes := buf.Len()
	rw.swapForWritingToken()
	defer rw.releaseWritingToken()
	n, err := buf.WriteTo(rw.w) // This is buffered (by 4MB chunks).  Are the writes to GCS synchronous?
	if err != nil {
		rw.writeErr = err
		switch typedErr := err.(type) {
		case *googleapi.Error:
			metrics.BackendFailureCount.WithLabelValues(
				label, "googleapi.Error").Inc()
			log.Println(typedErr, rw.bucket, rw.path)
			for _, e := range typedErr.Errors {
				log.Println(e)
			}
		default:
			metrics.BackendFailureCount.WithLabelValues(
				label, "other error").Inc()
			log.Println(typedErr, rw.bucket, rw.path)
		}
		// This approximates the number of rows written prior to error.
		// It is unclear whether these rows will actually show up.
		// The caller should likely abandon the archive at this point,
		// as further writing will likely result in a corrupted file.
		// See https://github.com/m-lab/etl/issues/899
		rowEstimate := int(n) * len(rows) / numBytes
		rw.rows += rowEstimate
		return rowEstimate, err
	}

	// TODO - these may not be committed, so the returned value may be wrong.
	rw.rows += len(rows)
	return len(rows), nil
}

// Close synchronizes on the tokens, and closes the backing file.
func (rw *RowWriter) Close() error {
	// Take BOTH tokens, to ensure no other goroutines are still running.
	<-rw.encoding
	<-rw.writing

	close(rw.encoding)
	close(rw.writing)

	log.Println("Closing", rw.bucket, rw.path)
	err := rw.w.Close()
	if err != nil {
		log.Println(err)
		return err
	}

	oa := gcs.ObjectAttrsToUpdate{}
	oa.Metadata = make(map[string]string, 1)
	oa.Metadata["rows"] = fmt.Sprint(rw.rows)
	if rw.writeErr != nil {
		oa.Metadata["writeError"] = rw.writeErr.Error()
	}

	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
	defer cancel()
	attr, err := rw.o.Update(ctx, oa)
	log.Println(attr, err)
	return err

}

// SinkFactory implements factory.SinkFactory.
type SinkFactory struct {
	client       stiface.Client
	outputBucket string
}

// returns the full path/filename from a gs://bucket/path/filename string.
func pathAndFilename(uri string) (string, error) {
	parts := strings.SplitN(uri, "/", 4)
	if len(parts) != 4 || parts[0] != "gs:" || len(parts[3]) == 0 {
		return "", errors.New("Bad GCS path")
	}
	return parts[3], nil
}

// Get implements factory.SinkFactory
func (sf *SinkFactory) Get(ctx context.Context, path etl.DataPath) (row.Sink, etl.ProcessingError) {
	fn, err := pathAndFilename(path.URI)
	if err != nil {
		return nil, factory.NewError(path.DataType, "InvalidPath",
			http.StatusInternalServerError, err)
	}
	s, err := NewRowWriter(ctx, sf.client, sf.outputBucket, fn+".json")
	if err != nil {
		return nil, factory.NewError(path.DataType, "SinkFactory",
			http.StatusInternalServerError, err)
	}
	return s, nil
}

// NewSinkFactory returns the default SinkFactory
func NewSinkFactory(client stiface.Client, outputBucket string) factory.SinkFactory {
	return &SinkFactory{client: client, outputBucket: outputBucket}
}
<|endoftext|>


<commit_before>package store

import (
	"sync"
	"time"

	log "github.com/cihub/seelog"

	"github.com/mattheath/phosphor/domain"
)

type MemoryStore struct {
	sync.RWMutex
	store map[string]*domain.Trace
}

// NewMemoryStore initialises and returns a new MemoryStore
func NewMemoryStore() *MemoryStore {
	s := &MemoryStore{
		store: make(map[string]*domain.Trace),
	}

	// run stats worker
	go s.statsLoop()

	return s
}

// ReadTrace retrieves a full Trace, composed of Frames from the store by ID
func (s *MemoryStore) ReadTrace(id string) (*domain.Trace, error) {
	if s == nil {
		return nil, ErrStoreNotInitialised
	}

	s.RLock()
	defer s.RUnlock()

	return s.store[id], nil
}

// StoreTraceFrame into the store, if the trace doesn't not already exist
// this will be created for the global trace ID
func (s *MemoryStore) StoreFrame(f *domain.Frame) error {
	s.Lock()
	defer s.Unlock()

	if s == nil {
		return ErrStoreNotInitialised
	}
	if f == nil {
		return ErrInvalidFrame
	}
	if f.TraceId == "" {
		return ErrInvalidTraceId
	}

	// Load our current trace
	t := s.store[f.TraceId]

	// Initialise a new trace if we don't have it already
	if t == nil {
		t = domain.NewTrace()
	}

	// Add the new frame to this
	t.AppendFrame(f)

	// Store it back
	s.store[f.TraceId] = t

	return nil
}

// statsLoop loops and outputs stats every 5 seconds
func (s *MemoryStore) statsLoop() {

	tick := time.NewTicker(5 * time.Second)

	// @todo listen for shutdown, stop ticker and exit cleanly
	for {
		<-tick.C // block until tick

		s.printStats()
	}
}

// printStats about the status of the memorystore to stdout
func (s *MemoryStore) printStats() {

	// Get some data while under the mutex
	s.RLock()
	count := len(s.store)
	s.RUnlock()

	// Separate processing and logging outside of mutex
	log.Infof("[MemoryStore] Traces stored: %v", count)
}
<commit_msg>Slight name tweak in memorystore<commit_after>package store

import (
	"sync"
	"time"

	log "github.com/cihub/seelog"

	"github.com/mattheath/phosphor/domain"
)

type MemoryStore struct {
	sync.RWMutex
	traces map[string]*domain.Trace
}

// NewMemoryStore initialises and returns a new MemoryStore
func NewMemoryStore() *MemoryStore {
	s := &MemoryStore{
		traces: make(map[string]*domain.Trace),
	}

	// run stats worker
	go s.statsLoop()

	return s
}

// ReadTrace retrieves a full Trace, composed of Frames from the store by ID
func (s *MemoryStore) ReadTrace(id string) (*domain.Trace, error) {
	if s == nil {
		return nil, ErrStoreNotInitialised
	}

	s.RLock()
	defer s.RUnlock()

	return s.traces[id], nil
}

// StoreTraceFrame into the store, if the trace doesn't not already exist
// this will be created for the global trace ID
func (s *MemoryStore) StoreFrame(f *domain.Frame) error {
	s.Lock()
	defer s.Unlock()

	if s == nil {
		return ErrStoreNotInitialised
	}
	if f == nil {
		return ErrInvalidFrame
	}
	if f.TraceId == "" {
		return ErrInvalidTraceId
	}

	// Load our current trace
	t := s.traces[f.TraceId]

	// Initialise a new trace if we don't have it already
	if t == nil {
		t = domain.NewTrace()
	}

	// Add the new frame to this
	t.AppendFrame(f)

	// Store it back
	s.traces[f.TraceId] = t

	return nil
}

// statsLoop loops and outputs stats every 5 seconds
func (s *MemoryStore) statsLoop() {

	tick := time.NewTicker(5 * time.Second)

	// @todo listen for shutdown, stop ticker and exit cleanly
	for {
		<-tick.C // block until tick

		s.printStats()
	}
}

// printStats about the status of the memorystore to stdout
func (s *MemoryStore) printStats() {

	// Get some data while under the mutex
	s.RLock()
	count := len(s.traces)
	s.RUnlock()

	// Separate processing and logging outside of mutex
	log.Infof("[MemoryStore] Traces stored: %v", count)
}
<|endoftext|>


<commit_before>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

// Package tikv provides tcp connection to kvserver.
package tikv

import (
	"io"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"github.com/grpc-ecosystem/go-grpc-middleware"
	"github.com/grpc-ecosystem/go-grpc-middleware/tracing/opentracing"
	"github.com/grpc-ecosystem/go-grpc-prometheus"
	"github.com/juju/errors"
	"github.com/pingcap/kvproto/pkg/coprocessor"
	"github.com/pingcap/kvproto/pkg/tikvpb"
	"github.com/pingcap/tidb/config"
	"github.com/pingcap/tidb/metrics"
	"github.com/pingcap/tidb/store/tikv/tikvrpc"
	"github.com/pingcap/tidb/terror"
	log "github.com/sirupsen/logrus"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
)

// MaxConnectionCount is the max gRPC connections that will be established with
// each tikv-server.
var MaxConnectionCount uint = 16

// MaxSendMsgSize set max gRPC request message size sent to server. If any request message size is larger than
// current value, an error will be reported from gRPC.
var MaxSendMsgSize = 1<<31 - 1

// MaxCallMsgSize set max gRPC receive message size received from server. If any message size is larger than
// current value, an error will be reported from gRPC.
var MaxCallMsgSize = 1<<31 - 1

// Timeout durations.
const (
	dialTimeout       = 5 * time.Second
	readTimeoutShort  = 20 * time.Second  // For requests that read/write several key-values.
	ReadTimeoutMedium = 60 * time.Second  // For requests that may need scan region.
	ReadTimeoutLong   = 150 * time.Second // For requests that may need scan region multiple times.
	GCTimeout         = 5 * time.Minute

	grpcInitialWindowSize     = 1 << 30
	grpcInitialConnWindowSize = 1 << 30
)

// Client is a client that sends RPC.
// It should not be used after calling Close().
type Client interface {
	// Close should release all data.
	Close() error
	// SendRequest sends Request.
	SendRequest(ctx context.Context, addr string, req *tikvrpc.Request, timeout time.Duration) (*tikvrpc.Response, error)
}

type connArray struct {
	index uint32
	v     []*grpc.ClientConn
	// Bind with a background goroutine to process coprocessor streaming timeout.
	streamTimeout chan *tikvrpc.Lease
}

func newConnArray(maxSize uint, addr string, security config.Security) (*connArray, error) {
	a := &connArray{
		index:         0,
		v:             make([]*grpc.ClientConn, maxSize),
		streamTimeout: make(chan *tikvrpc.Lease, 1024),
	}
	if err := a.Init(addr, security); err != nil {
		return nil, err
	}
	return a, nil
}

func (a *connArray) Init(addr string, security config.Security) error {
	opt := grpc.WithInsecure()
	if len(security.ClusterSSLCA) != 0 {
		tlsConfig, err := security.ToTLSConfig()
		if err != nil {
			return errors.Trace(err)
		}
		opt = grpc.WithTransportCredentials(credentials.NewTLS(tlsConfig))
	}

	unaryInterceptor := grpc_prometheus.UnaryClientInterceptor
	streamInterceptor := grpc_prometheus.StreamClientInterceptor
	cfg := config.GetGlobalConfig()
	if cfg.OpenTracing.Enable {
		unaryInterceptor = grpc_middleware.ChainUnaryClient(
			unaryInterceptor,
			grpc_opentracing.UnaryClientInterceptor(),
		)
		streamInterceptor = grpc_middleware.ChainStreamClient(
			streamInterceptor,
			grpc_opentracing.StreamClientInterceptor(),
		)
	}

	for i := range a.v {

		ctx, cancel := context.WithTimeout(context.Background(), dialTimeout)
		conn, err := grpc.DialContext(
			ctx,
			addr,
			opt,
			grpc.WithInitialWindowSize(grpcInitialWindowSize),
			grpc.WithInitialConnWindowSize(grpcInitialConnWindowSize),
			grpc.WithUnaryInterceptor(unaryInterceptor),
			grpc.WithStreamInterceptor(streamInterceptor),
			grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(MaxCallMsgSize)),
			grpc.WithDefaultCallOptions(grpc.MaxCallSendMsgSize(MaxSendMsgSize)),
		)
		cancel()
		if err != nil {
			// Cleanup if the initialization fails.
			a.Close()
			return errors.Trace(err)
		}
		a.v[i] = conn
	}
	go tikvrpc.CheckStreamTimeoutLoop(a.streamTimeout)

	return nil
}

func (a *connArray) Get() *grpc.ClientConn {
	next := atomic.AddUint32(&a.index, 1) % uint32(len(a.v))
	return a.v[next]
}

func (a *connArray) Close() {
	for i, c := range a.v {
		if c != nil {
			err := c.Close()
			terror.Log(errors.Trace(err))
			a.v[i] = nil
		}
	}
	close(a.streamTimeout)
}

// rpcClient is RPC client struct.
// TODO: Add flow control between RPC clients in TiDB ond RPC servers in TiKV.
// Since we use shared client connection to communicate to the same TiKV, it's possible
// that there are too many concurrent requests which overload the service of TiKV.
// TODO: Implement background cleanup. It adds a background goroutine to periodically check
// whether there is any connection is idle and then close and remove these idle connections.
type rpcClient struct {
	sync.RWMutex
	isClosed bool
	conns    map[string]*connArray
	security config.Security
}

func newRPCClient(security config.Security) *rpcClient {
	return &rpcClient{
		conns:    make(map[string]*connArray),
		security: security,
	}
}

func (c *rpcClient) getConnArray(addr string) (*connArray, error) {
	c.RLock()
	if c.isClosed {
		c.RUnlock()
		return nil, errors.Errorf("rpcClient is closed")
	}
	array, ok := c.conns[addr]
	c.RUnlock()
	if !ok {
		var err error
		array, err = c.createConnArray(addr)
		if err != nil {
			return nil, err
		}
	}
	return array, nil
}

func (c *rpcClient) createConnArray(addr string) (*connArray, error) {
	c.Lock()
	defer c.Unlock()
	array, ok := c.conns[addr]
	if !ok {
		var err error
		array, err = newConnArray(MaxConnectionCount, addr, c.security)
		if err != nil {
			return nil, err
		}
		c.conns[addr] = array
	}
	return array, nil
}

func (c *rpcClient) closeConns() {
	c.Lock()
	if !c.isClosed {
		c.isClosed = true
		// close all connections
		for _, array := range c.conns {
			array.Close()
		}
	}
	c.Unlock()
}

// SendRequest sends a Request to server and receives Response.
func (c *rpcClient) SendRequest(ctx context.Context, addr string, req *tikvrpc.Request, timeout time.Duration) (*tikvrpc.Response, error) {
	start := time.Now()
	reqType := req.Type.String()
	storeID := strconv.FormatUint(req.Context.GetPeer().GetStoreId(), 10)
	defer func() {
		metrics.TiKVSendReqHistogram.WithLabelValues(reqType, storeID).Observe(time.Since(start).Seconds())
	}()

	connArray, err := c.getConnArray(addr)
	if err != nil {
		return nil, errors.Trace(err)
	}
	client := tikvpb.NewTikvClient(connArray.Get())

	if req.Type != tikvrpc.CmdCopStream {
		ctx1, cancel := context.WithTimeout(ctx, timeout)
		defer cancel()
		return tikvrpc.CallRPC(ctx1, client, req)
	}

	// Coprocessor streaming request.
	// Use context to support timeout for grpc streaming client.
	ctx1, cancel := context.WithCancel(ctx)
	resp, err := tikvrpc.CallRPC(ctx1, client, req)
	if err != nil {
		return nil, errors.Trace(err)
	}

	// Put the lease object to the timeout channel, so it would be checked periodically.
	copStream := resp.CopStream
	copStream.Timeout = timeout
	copStream.Lease.Cancel = cancel
	connArray.streamTimeout <- &copStream.Lease

	// Read the first streaming response to get CopStreamResponse.
	// This can make error handling much easier, because SendReq() retry on
	// region error automatically.
	var first *coprocessor.Response
	first, err = copStream.Recv()
	if err != nil {
		if errors.Cause(err) != io.EOF {
			return nil, errors.Trace(err)
		}
		log.Debug("copstream returns nothing for the request.")
	}
	copStream.Response = first
	return resp, nil
}

func (c *rpcClient) Close() error {
	c.closeConns()
	return nil
}
<commit_msg>Add max backoff option (#6770)<commit_after>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

// Package tikv provides tcp connection to kvserver.
package tikv

import (
	"io"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"github.com/grpc-ecosystem/go-grpc-middleware"
	"github.com/grpc-ecosystem/go-grpc-middleware/tracing/opentracing"
	"github.com/grpc-ecosystem/go-grpc-prometheus"
	"github.com/juju/errors"
	"github.com/pingcap/kvproto/pkg/coprocessor"
	"github.com/pingcap/kvproto/pkg/tikvpb"
	"github.com/pingcap/tidb/config"
	"github.com/pingcap/tidb/metrics"
	"github.com/pingcap/tidb/store/tikv/tikvrpc"
	"github.com/pingcap/tidb/terror"
	log "github.com/sirupsen/logrus"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
)

// MaxConnectionCount is the max gRPC connections that will be established with
// each tikv-server.
var MaxConnectionCount uint = 16

// MaxSendMsgSize set max gRPC request message size sent to server. If any request message size is larger than
// current value, an error will be reported from gRPC.
var MaxSendMsgSize = 1<<31 - 1

// MaxCallMsgSize set max gRPC receive message size received from server. If any message size is larger than
// current value, an error will be reported from gRPC.
var MaxCallMsgSize = 1<<31 - 1

// Timeout durations.
const (
	dialTimeout       = 5 * time.Second
	readTimeoutShort  = 20 * time.Second  // For requests that read/write several key-values.
	ReadTimeoutMedium = 60 * time.Second  // For requests that may need scan region.
	ReadTimeoutLong   = 150 * time.Second // For requests that may need scan region multiple times.
	GCTimeout         = 5 * time.Minute

	grpcInitialWindowSize     = 1 << 30
	grpcInitialConnWindowSize = 1 << 30
)

// Client is a client that sends RPC.
// It should not be used after calling Close().
type Client interface {
	// Close should release all data.
	Close() error
	// SendRequest sends Request.
	SendRequest(ctx context.Context, addr string, req *tikvrpc.Request, timeout time.Duration) (*tikvrpc.Response, error)
}

type connArray struct {
	index uint32
	v     []*grpc.ClientConn
	// Bind with a background goroutine to process coprocessor streaming timeout.
	streamTimeout chan *tikvrpc.Lease
}

func newConnArray(maxSize uint, addr string, security config.Security) (*connArray, error) {
	a := &connArray{
		index:         0,
		v:             make([]*grpc.ClientConn, maxSize),
		streamTimeout: make(chan *tikvrpc.Lease, 1024),
	}
	if err := a.Init(addr, security); err != nil {
		return nil, err
	}
	return a, nil
}

func (a *connArray) Init(addr string, security config.Security) error {
	opt := grpc.WithInsecure()
	if len(security.ClusterSSLCA) != 0 {
		tlsConfig, err := security.ToTLSConfig()
		if err != nil {
			return errors.Trace(err)
		}
		opt = grpc.WithTransportCredentials(credentials.NewTLS(tlsConfig))
	}

	unaryInterceptor := grpc_prometheus.UnaryClientInterceptor
	streamInterceptor := grpc_prometheus.StreamClientInterceptor
	cfg := config.GetGlobalConfig()
	if cfg.OpenTracing.Enable {
		unaryInterceptor = grpc_middleware.ChainUnaryClient(
			unaryInterceptor,
			grpc_opentracing.UnaryClientInterceptor(),
		)
		streamInterceptor = grpc_middleware.ChainStreamClient(
			streamInterceptor,
			grpc_opentracing.StreamClientInterceptor(),
		)
	}

	for i := range a.v {
		ctx, cancel := context.WithTimeout(context.Background(), dialTimeout)
		conn, err := grpc.DialContext(
			ctx,
			addr,
			opt,
			grpc.WithInitialWindowSize(grpcInitialWindowSize),
			grpc.WithInitialConnWindowSize(grpcInitialConnWindowSize),
			grpc.WithUnaryInterceptor(unaryInterceptor),
			grpc.WithStreamInterceptor(streamInterceptor),
			grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(MaxCallMsgSize)),
			grpc.WithDefaultCallOptions(grpc.MaxCallSendMsgSize(MaxSendMsgSize)),
			grpc.WithBackoffMaxDelay(time.Second*3),
		)
		cancel()
		if err != nil {
			// Cleanup if the initialization fails.
			a.Close()
			return errors.Trace(err)
		}
		a.v[i] = conn
	}
	go tikvrpc.CheckStreamTimeoutLoop(a.streamTimeout)

	return nil
}

func (a *connArray) Get() *grpc.ClientConn {
	next := atomic.AddUint32(&a.index, 1) % uint32(len(a.v))
	return a.v[next]
}

func (a *connArray) Close() {
	for i, c := range a.v {
		if c != nil {
			err := c.Close()
			terror.Log(errors.Trace(err))
			a.v[i] = nil
		}
	}
	close(a.streamTimeout)
}

// rpcClient is RPC client struct.
// TODO: Add flow control between RPC clients in TiDB ond RPC servers in TiKV.
// Since we use shared client connection to communicate to the same TiKV, it's possible
// that there are too many concurrent requests which overload the service of TiKV.
// TODO: Implement background cleanup. It adds a background goroutine to periodically check
// whether there is any connection is idle and then close and remove these idle connections.
type rpcClient struct {
	sync.RWMutex
	isClosed bool
	conns    map[string]*connArray
	security config.Security
}

func newRPCClient(security config.Security) *rpcClient {
	return &rpcClient{
		conns:    make(map[string]*connArray),
		security: security,
	}
}

func (c *rpcClient) getConnArray(addr string) (*connArray, error) {
	c.RLock()
	if c.isClosed {
		c.RUnlock()
		return nil, errors.Errorf("rpcClient is closed")
	}
	array, ok := c.conns[addr]
	c.RUnlock()
	if !ok {
		var err error
		array, err = c.createConnArray(addr)
		if err != nil {
			return nil, err
		}
	}
	return array, nil
}

func (c *rpcClient) createConnArray(addr string) (*connArray, error) {
	c.Lock()
	defer c.Unlock()
	array, ok := c.conns[addr]
	if !ok {
		var err error
		array, err = newConnArray(MaxConnectionCount, addr, c.security)
		if err != nil {
			return nil, err
		}
		c.conns[addr] = array
	}
	return array, nil
}

func (c *rpcClient) closeConns() {
	c.Lock()
	if !c.isClosed {
		c.isClosed = true
		// close all connections
		for _, array := range c.conns {
			array.Close()
		}
	}
	c.Unlock()
}

// SendRequest sends a Request to server and receives Response.
func (c *rpcClient) SendRequest(ctx context.Context, addr string, req *tikvrpc.Request, timeout time.Duration) (*tikvrpc.Response, error) {
	start := time.Now()
	reqType := req.Type.String()
	storeID := strconv.FormatUint(req.Context.GetPeer().GetStoreId(), 10)
	defer func() {
		metrics.TiKVSendReqHistogram.WithLabelValues(reqType, storeID).Observe(time.Since(start).Seconds())
	}()

	connArray, err := c.getConnArray(addr)
	if err != nil {
		return nil, errors.Trace(err)
	}
	client := tikvpb.NewTikvClient(connArray.Get())

	if req.Type != tikvrpc.CmdCopStream {
		ctx1, cancel := context.WithTimeout(ctx, timeout)
		defer cancel()
		return tikvrpc.CallRPC(ctx1, client, req)
	}

	// Coprocessor streaming request.
	// Use context to support timeout for grpc streaming client.
	ctx1, cancel := context.WithCancel(ctx)
	resp, err := tikvrpc.CallRPC(ctx1, client, req)
	if err != nil {
		return nil, errors.Trace(err)
	}

	// Put the lease object to the timeout channel, so it would be checked periodically.
	copStream := resp.CopStream
	copStream.Timeout = timeout
	copStream.Lease.Cancel = cancel
	connArray.streamTimeout <- &copStream.Lease

	// Read the first streaming response to get CopStreamResponse.
	// This can make error handling much easier, because SendReq() retry on
	// region error automatically.
	var first *coprocessor.Response
	first, err = copStream.Recv()
	if err != nil {
		if errors.Cause(err) != io.EOF {
			return nil, errors.Trace(err)
		}
		log.Debug("copstream returns nothing for the request.")
	}
	copStream.Response = first
	return resp, nil
}

func (c *rpcClient) Close() error {
	c.closeConns()
	return nil
}
<|endoftext|>


<commit_before>package upstream

import (
	"net"
	"net/http"
	"fmt"
	"errors"
	"time"

	"github.com/ngmoco/falcore"
)

type UpstreamTransport struct {
	DNSCacheDuration time.Duration

	host string
	port int

	tcpaddr   *net.TCPAddr
	tcpaddrCacheTime time.Time

	transport *http.Transport
	timeout time.Duration
}

// transport is optional.  We will override Dial
func NewUpstreamTransport(host string, port int, timeout time.Duration, transport *http.Transport)*UpstreamTransport {
	ut := &UpstreamTransport {
		host: host,
		port: port,
		timeout: timeout,
		transport: transport,
	}
	ut.DNSCacheDuration = 15 * time.Minute

	if(ut.transport == nil){
		ut.transport = &http.Transport{}
		ut.transport.MaxIdleConnsPerHost = 15
	}

	ut.transport.Dial = func(n, addr string) (c net.Conn, err error) {
		return ut.dial(n, addr)
	}

	return ut
}

func (t *UpstreamTransport) dial(n, a string)(c net.Conn, err error) {
	var addr *net.TCPAddr
	addr, err = t.lookupIp()

	falcore.Fine("Dialing connection to %v", addr)
	var ctcp *net.TCPConn
	ctcp, err = net.DialTCP("tcp4", nil, addr)
	if err != nil {
		falcore.Error("Dial Failed: %v", err)
		return
	}
	c = &timeoutConnWrapper{conn: ctcp, timeout: t.timeout}
	return
}

func (t *UpstreamTransport) lookupIp()(addr *net.TCPAddr, err error) {
	// Cached tcpaddr
	if t.tcpaddr != nil && t.tcpaddrCacheTime.Add(t.DNSCacheDuration).After(time.Now()) {
		return t.tcpaddr, nil
	}

	ips, err := net.LookupIP(t.host)
	var ip net.IP = nil

	if err != nil {
		return nil, err
	}

	// Find first IPv4 IP
	for i := range ips {
		ip = ips[i].To4()
		if ip != nil {
			break
		}
	}

	if ip != nil {
		t.tcpaddr = &net.TCPAddr{}
		t.tcpaddr.Port = t.port
		t.tcpaddr.IP = ip
		t.tcpaddrCacheTime = time.Now()
		addr = t.tcpaddr
	} else {
		errstr := fmt.Sprintf("Can't get IP addr for %v: %v", t.host, err)
		err = errors.New(errstr)
	}

	return
}

type timeoutConnWrapper struct {
	conn    net.Conn
	timeout time.Duration
}

func (cw *timeoutConnWrapper) Write(b []byte) (int, error) {
	if err := cw.conn.SetDeadline(time.Now().Add(cw.timeout)); err != nil {
		return 0, err
	}
	return cw.conn.Write(b)
}
func (cw *timeoutConnWrapper) Read(b []byte) (n int, err error)   { return cw.conn.Read(b) }
func (cw *timeoutConnWrapper) Close() error                       { return cw.conn.Close() }
func (cw *timeoutConnWrapper) LocalAddr() net.Addr                { return cw.conn.LocalAddr() }
func (cw *timeoutConnWrapper) RemoteAddr() net.Addr               { return cw.conn.RemoteAddr() }
func (cw *timeoutConnWrapper) SetDeadline(t time.Time) error      { return cw.conn.SetDeadline(t) }
func (cw *timeoutConnWrapper) SetReadDeadline(t time.Time) error  { return cw.conn.SetReadDeadline(t) }
func (cw *timeoutConnWrapper) SetWriteDeadline(t time.Time) error { return cw.conn.SetWriteDeadline(t) }
<commit_msg>allow bypass of timeout behavior<commit_after>package upstream

import (
	"net"
	"net/http"
	"fmt"
	"errors"
	"time"

	"github.com/ngmoco/falcore"
)

type UpstreamTransport struct {
	DNSCacheDuration time.Duration

	host string
	port int

	tcpaddr   *net.TCPAddr
	tcpaddrCacheTime time.Time

	transport *http.Transport
	timeout time.Duration
}

// transport is optional.  We will override Dial
func NewUpstreamTransport(host string, port int, timeout time.Duration, transport *http.Transport)*UpstreamTransport {
	ut := &UpstreamTransport {
		host: host,
		port: port,
		timeout: timeout,
		transport: transport,
	}
	ut.DNSCacheDuration = 15 * time.Minute

	if(ut.transport == nil){
		ut.transport = &http.Transport{}
		ut.transport.MaxIdleConnsPerHost = 15
	}

	ut.transport.Dial = func(n, addr string) (c net.Conn, err error) {
		return ut.dial(n, addr)
	}

	return ut
}

func (t *UpstreamTransport) dial(n, a string)(c net.Conn, err error) {
	var addr *net.TCPAddr
	addr, err = t.lookupIp()

	falcore.Fine("Dialing connection to %v", addr)
	var ctcp *net.TCPConn
	ctcp, err = net.DialTCP("tcp4", nil, addr)
	if err != nil {
		falcore.Error("Dial Failed: %v", err)
		return
	}

	// FIXME: Go1 has a race that causes problems with timeouts
	// Recommend disabling until Go1.1
	if t.timeout > 0 {
		c = &timeoutConnWrapper{conn: ctcp, timeout: t.timeout}
	} else {
		c = ctcp
	}

	return
}

func (t *UpstreamTransport) lookupIp()(addr *net.TCPAddr, err error) {
	// Cached tcpaddr
	if t.tcpaddr != nil && t.tcpaddrCacheTime.Add(t.DNSCacheDuration).After(time.Now()) {
		return t.tcpaddr, nil
	}

	ips, err := net.LookupIP(t.host)
	var ip net.IP = nil

	if err != nil {
		return nil, err
	}

	// Find first IPv4 IP
	for i := range ips {
		ip = ips[i].To4()
		if ip != nil {
			break
		}
	}

	if ip != nil {
		t.tcpaddr = &net.TCPAddr{}
		t.tcpaddr.Port = t.port
		t.tcpaddr.IP = ip
		t.tcpaddrCacheTime = time.Now()
		addr = t.tcpaddr
	} else {
		errstr := fmt.Sprintf("Can't get IP addr for %v: %v", t.host, err)
		err = errors.New(errstr)
	}

	return
}

type timeoutConnWrapper struct {
	conn    net.Conn
	timeout time.Duration
}

func (cw *timeoutConnWrapper) Write(b []byte) (int, error) {
	if err := cw.conn.SetDeadline(time.Now().Add(cw.timeout)); err != nil {
		return 0, err
	}
	return cw.conn.Write(b)
}
func (cw *timeoutConnWrapper) Read(b []byte) (n int, err error)   { return cw.conn.Read(b) }
func (cw *timeoutConnWrapper) Close() error                       { return cw.conn.Close() }
func (cw *timeoutConnWrapper) LocalAddr() net.Addr                { return cw.conn.LocalAddr() }
func (cw *timeoutConnWrapper) RemoteAddr() net.Addr               { return cw.conn.RemoteAddr() }
func (cw *timeoutConnWrapper) SetDeadline(t time.Time) error      { return cw.conn.SetDeadline(t) }
func (cw *timeoutConnWrapper) SetReadDeadline(t time.Time) error  { return cw.conn.SetReadDeadline(t) }
func (cw *timeoutConnWrapper) SetWriteDeadline(t time.Time) error { return cw.conn.SetWriteDeadline(t) }
<|endoftext|>


<commit_before>// Copyright 2016 Keybase Inc. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.

// These tests exercise quota reclamation.

package test

import (
	"testing"
	"time"
)

// Check that simple quota reclamation works
func TestQRSimple(t *testing.T) {
	test(t,
		writers("alice"),
		as(alice,
			addTime(1*time.Minute),
			mkfile("a", "hello"),
			rm("a"),
			addTime(2*time.Minute),
			forceQuotaReclamation(),
		),
	)
}

// Check that quota reclamation works, eventually, after enough iterations.
func TestQRLargePointerSet(t *testing.T) {
	var busyWork []fileOp
	iters := 100
	for i := 0; i < iters; i++ {
		busyWork = append(busyWork, mkfile("a", "hello"), rm("a"))
	}
	// 5 unreferenced pointers per iteration -- 3 updates to the root
	// block, one empty file written to, and one non-empty file
	// deleted.
	ptrsPerIter := 5
	var qrOps []optionOp
	// Each reclamation needs a sync after it (e.g., a new "as"
	// clause) to ensure it completes before the next force
	// reclamation.
	for i := 0; i < ptrsPerIter*iters/100; i++ {
		qrOps = append(qrOps, as(alice,
			addTime(2*time.Minute),
			forceQuotaReclamation(),
		))
	}
	totalOps := []optionOp{writers("alice"), as(alice, busyWork...)}
	totalOps = append(totalOps, qrOps...)
	test(t, totalOps...)
}

// Test that quota reclamation handles conflict resolution correctly.
func TestQRAfterCR(t *testing.T) {
	test(t,
		writers("alice", "bob"),
		as(alice,
			mkfile("a/b", "hello"),
		),
		as(bob,
			disableUpdates(),
		),
		as(alice,
			write("a/c", "world"),
		),
		as(bob, noSync(),
			rm("a/b"),
			reenableUpdates(),
		),
		as(alice,
			addTime(2*time.Minute),
			forceQuotaReclamation(),
		),
	)
}

// Check that quota reclamation works on multi-block files
func TestQRWithMultiBlockFiles(t *testing.T) {
	test(t,
		blockSize(20), writers("alice"),
		as(alice,
			addTime(1*time.Minute),
			mkfile("a", ntimesString(15, "0123456789")),
			rm("a"),
			addTime(2*time.Minute),
			forceQuotaReclamation(),
		),
	)
}
<commit_msg>qr_test: bob does the quota reclamation<commit_after>// Copyright 2016 Keybase Inc. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.

// These tests exercise quota reclamation.

package test

import (
	"testing"
	"time"
)

// Check that simple quota reclamation works
func TestQRSimple(t *testing.T) {
	test(t,
		writers("alice"),
		as(alice,
			addTime(1*time.Minute),
			mkfile("a", "hello"),
			rm("a"),
			addTime(2*time.Minute),
			forceQuotaReclamation(),
		),
	)
}

// Check that quota reclamation works, eventually, after enough iterations.
func TestQRLargePointerSet(t *testing.T) {
	var busyWork []fileOp
	iters := 100
	for i := 0; i < iters; i++ {
		busyWork = append(busyWork, mkfile("a", "hello"), rm("a"))
	}
	// 5 unreferenced pointers per iteration -- 3 updates to the root
	// block, one empty file written to, and one non-empty file
	// deleted.
	ptrsPerIter := 5
	var qrOps []optionOp
	// Each reclamation needs a sync after it (e.g., a new "as"
	// clause) to ensure it completes before the next force
	// reclamation.
	for i := 0; i < ptrsPerIter*iters/100; i++ {
		qrOps = append(qrOps, as(alice,
			addTime(2*time.Minute),
			forceQuotaReclamation(),
		))
	}
	totalOps := []optionOp{writers("alice"), as(alice, busyWork...)}
	totalOps = append(totalOps, qrOps...)
	test(t, totalOps...)
}

// Test that quota reclamation handles conflict resolution correctly.
func TestQRAfterCR(t *testing.T) {
	test(t,
		writers("alice", "bob"),
		as(alice,
			mkfile("a/b", "hello"),
		),
		as(bob,
			disableUpdates(),
		),
		as(alice,
			write("a/c", "world"),
		),
		as(bob, noSync(),
			rm("a/b"),
			reenableUpdates(),
		),
		as(bob,
			addTime(2*time.Minute),
			forceQuotaReclamation(),
		),
	)
}

// Check that quota reclamation works on multi-block files
func TestQRWithMultiBlockFiles(t *testing.T) {
	test(t,
		blockSize(20), writers("alice"),
		as(alice,
			addTime(1*time.Minute),
			mkfile("a", ntimesString(15, "0123456789")),
			rm("a"),
			addTime(2*time.Minute),
			forceQuotaReclamation(),
		),
	)
}
<|endoftext|>


<commit_before>package php

import (
	"fmt"
	"strings"
	"unicode/utf8"
)

type lexer struct {
	lastPos int
	pos     int
	start   int
	width   int
	input   string
	items   chan Item // channel of scanned items.
}

func newLexer(input string) *lexer {
	l := &lexer{
		input: input,
		items: make(chan Item),
	}
	go l.run()
	return l
}

// stateFn represents the state of the scanner
// as a function that returns the next state.
type stateFn func(*lexer) stateFn

// run lexes the input by executing state functions until
// the state is nil.
func (l *lexer) run() {
	for state := lexHTML; state != nil; {
		state = state(l)
	}
	close(l.items) // No more tokens will be delivered.
}

func (l *lexer) emit(t ItemType) {
	i := Item{t, Location{Pos: l.start}, l.input[l.start:l.pos]}
	l.items <- i
	l.start = l.pos
}

// nextItem returns the next item from the input.
func (l *lexer) nextItem() Item {
	Item := <-l.items
	l.lastPos = Item.pos.Pos
	return Item
}

// peek returns but does not consume the next rune in the input.
func (l *lexer) peek() rune {
	r := l.next()
	l.backup()
	return r
}

// backup steps back one rune. Can only be called once per call of next.
func (l *lexer) backup() {
	l.pos -= l.width
}

func (l *lexer) accept(valid string) bool {
	if strings.IndexRune(valid, l.next()) >= 0 {
		return true
	}
	l.backup()
	return false
}

// acceptRun consumes a run of runes from the valid set.
func (l *lexer) acceptRun(valid string) {
	for strings.IndexRune(valid, l.next()) >= 0 {
	}
	l.backup()
}

func (l *lexer) next() rune {
	if int(l.pos) >= len(l.input) {
		l.width = 0
		return eof
	}
	r, w := utf8.DecodeRuneInString(l.input[l.pos:])
	l.width = w
	l.pos += l.width
	return r
}

// ignore skips over the pending input before this point.
func (l *lexer) ignore() {
	l.start = l.pos
}

func (l *lexer) skipSpace() {
	r := l.next()
	for isSpace(r) {
		r = l.next()
	}
	l.backup()
	l.ignore()
}

func (l *lexer) errorf(format string, args ...interface{}) stateFn {
	l.items <- Item{itemError, Location{Pos: l.start}, fmt.Sprintf(format, args...)}
	return nil
}

// isSpace reports whether r is a space character.
func isSpace(r rune) bool {
	return r == ' ' || r == '\t' || r == '\n'
}

func isKeyword(i ItemType) bool {
	is, ok := keywordMap[i]
	return is && ok
}

var keywordMap = map[ItemType]bool{
	itemFunction: true,

	itemReturn: true,
	itemEcho:   true,

	itemIf:      true,
	itemElse:    true,
	itemElseIf:  true,
	itemFor:     true,
	itemForeach: true,
	itemWhile:   true,
	itemDo:      true,

	itemTry:     true,
	itemCatch:   true,
	itemFinally: true,

	itemClass:       true,
	itemPrivate:     true,
	itemProtected:   true,
	itemPublic:      true,
	itemInterface:   true,
	itemImplements:  true,
	itemExtends:     true,
	itemNewOperator: true,

	itemInstanceofOperator: true,
}
<commit_msg>Improved location tracking in lexer.<commit_after>package php

import (
	"fmt"
	"strings"
	"unicode/utf8"
)

type lexer struct {
	lastPos int
	pos     int
	line    int
	start   int
	width   int
	input   string
	file    string
	items   chan Item // channel of scanned items.
}

func newLexer(input string) *lexer {
	l := &lexer{
		input: input,
		items: make(chan Item),
	}
	go l.run()
	return l
}

// stateFn represents the state of the scanner
// as a function that returns the next state.
type stateFn func(*lexer) stateFn

// run lexes the input by executing state functions until
// the state is nil.
func (l *lexer) run() {
	for state := lexHTML; state != nil; {
		state = state(l)
	}
	close(l.items) // No more tokens will be delivered.
}

func (l *lexer) emit(t ItemType) {
	i := Item{t, l.currentLocation(), l.input[l.start:l.pos]}
	l.line += strings.Count(i.val, "\n")
	l.items <- i
	l.start = l.pos
}

func (l *lexer) currentLocation() Location {
	return Location{Pos: l.start, Line: l.line, File: l.file}
}

// nextItem returns the next item from the input.
func (l *lexer) nextItem() Item {
	Item := <-l.items
	l.lastPos = Item.pos.Pos
	return Item
}

// peek returns but does not consume the next rune in the input.
func (l *lexer) peek() rune {
	r := l.next()
	l.backup()
	return r
}

// backup steps back one rune. Can only be called once per call of next.
func (l *lexer) backup() {
	l.pos -= l.width
}

func (l *lexer) accept(valid string) bool {
	if strings.IndexRune(valid, l.next()) >= 0 {
		return true
	}
	l.backup()
	return false
}

// acceptRun consumes a run of runes from the valid set.
func (l *lexer) acceptRun(valid string) {
	for strings.IndexRune(valid, l.next()) >= 0 {
	}
	l.backup()
}

func (l *lexer) next() rune {
	if int(l.pos) >= len(l.input) {
		l.width = 0
		return eof
	}
	r, w := utf8.DecodeRuneInString(l.input[l.pos:])
	l.width = w
	l.pos += l.width
	return r
}

// ignore skips over the pending input before this point.
func (l *lexer) ignore() {
	l.start = l.pos
}

func (l *lexer) skipSpace() {
	r := l.next()
	for isSpace(r) {
		r = l.next()
	}
	l.backup()
	l.ignore()
}

func (l *lexer) errorf(format string, args ...interface{}) stateFn {
	i := Item{itemError, l.currentLocation(), fmt.Sprintf(format, args...)}
	l.line += strings.Count(i.val, "\n")
	l.items <- i
	return nil
}

// isSpace reports whether r is a space character.
func isSpace(r rune) bool {
	return r == ' ' || r == '\t' || r == '\n'
}

func isKeyword(i ItemType) bool {
	is, ok := keywordMap[i]
	return is && ok
}

var keywordMap = map[ItemType]bool{
	itemFunction: true,

	itemReturn: true,
	itemEcho:   true,

	itemIf:      true,
	itemElse:    true,
	itemElseIf:  true,
	itemFor:     true,
	itemForeach: true,
	itemWhile:   true,
	itemDo:      true,

	itemTry:     true,
	itemCatch:   true,
	itemFinally: true,

	itemClass:       true,
	itemPrivate:     true,
	itemProtected:   true,
	itemPublic:      true,
	itemInterface:   true,
	itemImplements:  true,
	itemExtends:     true,
	itemNewOperator: true,

	itemInstanceofOperator: true,
}
<|endoftext|>