<commit_before>package dns

// A concurrent client implementation.
// Client sends query to a channel which
// will then handle the query. Returned replys
// are return on another channel. Ready for handling --- same
// setup for server - a HANDLER function that gets run
// when the query returns.

// This completely mirrors server.go impl.
import (
	"os"
	//	"net"
)

type QueryHandler interface {
	QueryDNS(w RequestWriter, q *Msg)
}

// A RequestWriter interface is used by an DNS query handler to
// construct an DNS request.
type RequestWriter interface {
	WriteMessages([]*Msg)
        WriteMessage(*Msg)
}

// hijacked connections...?
type reply struct {
	Client *Client
	req    *Msg
}

// QueryMux is an DNS request multiplexer. It matches the
// zone name of each incoming request against a list of
// registered patterns add calls the handler for the pattern
// that most closely matches the zone name.
type QueryMux struct {
	m map[string]QueryHandler
}

// NewQueryMux allocates and returns a new QueryMux.
func NewQueryMux() *QueryMux { return &QueryMux{make(map[string]QueryHandler)} }

// DefaultQueryMux is the default QueryMux used by Query.
var DefaultQueryMux = NewQueryMux()

func newQueryChanSlice() chan []*Msg { return make(chan []*Msg) }
func newQueryChan() chan *Msg { return make(chan *Msg) }

// Default channel to use for the resolver
var DefaultReplyChan = newQueryChanSlice()
var DefaultQueryChan = newQueryChan()

// The HandlerQueryFunc type is an adapter to allow the use of
// ordinary functions as DNS query handlers.  If f is a function
// with the appropriate signature, HandlerQueryFunc(f) is a
// QueryHandler object that calls f.
type HandlerQueryFunc func(RequestWriter, *Msg)

// QueryDNS calls f(w, reg)
func (f HandlerQueryFunc) QueryDNS(w RequestWriter, r *Msg) {
	go f(w, r)
}

func HandleQueryFunc(pattern string, handler func(RequestWriter, *Msg)) {
	DefaultQueryMux.HandleQueryFunc(pattern, handler)
}

// Helper handlers
// Todo

// reusing zoneMatch from server.go
func (mux *QueryMux) match(zone string) QueryHandler {
	var h QueryHandler
	var n = 0
	for k, v := range mux.m {
		if !zoneMatch(k, zone) {
			continue
		}
		if h == nil || len(k) > n {
			n = len(k)
			h = v
		}
	}
	return h
}

func (mux *QueryMux) Handle(pattern string, handler QueryHandler) {
	if pattern == "" {
		panic("dns: invalid pattern " + pattern)
	}
	if pattern[len(pattern)-1] != '.' { // no ending .
		mux.m[pattern+"."] = handler
	} else {
		mux.m[pattern] = handler
	}
}

func (mux *QueryMux) HandleQueryFunc(pattern string, handler func(RequestWriter, *Msg)) {
	mux.Handle(pattern, HandlerQueryFunc(handler))
}

func (mux *QueryMux) QueryDNS(w RequestWriter, request *Msg) {
	h := mux.match(request.Question[0].Name)
	if h == nil {
		//                h = RefusedHandler()
		// something else
	}
	h.QueryDNS(w, request)
}

type Client struct {
	Network      string       // if "tcp" a TCP query will be initiated, otherwise an UDP one
	Attempts     int          // number of attempts
	Retry        bool         // retry with TCP
	ChannelQuery chan *Msg    // read DNS request from this channel
	ChannelReply chan []*Msg    // read DNS request from this channel
	Handler      QueryHandler // handler to invoke, dns.DefaultQueryMux if nil
	ReadTimeout  int64        // the net.Conn.SetReadTimeout value for new connections
	WriteTimeout int64        // the net.Conn.SetWriteTimeout value for new connections
}

// Query accepts incoming DNS request,
// Write to in
// creating a new service thread for each.  The service threads
// read requests and then call handler to reply to them.
// Handler is typically nil, in which case the DefaultServeMux is used.
func Query(c chan *Msg, handler QueryHandler) os.Error {
	client := &Client{ChannelQuery: c, Handler: handler}
	return client.Query()
}

func (c *Client) Query() os.Error {
	handler := c.Handler
	if handler == nil {
		handler = DefaultQueryMux
	}
forever:
	for {
		select {
		case in := <-c.ChannelQuery:
			w := new(reply)
			w.Client = c
			w.req = in
			handler.QueryDNS(w, w.req)
		}
	}
	return nil
}

func (c *Client) ListenAndQuery() os.Error {
	if c.ChannelQuery == nil {
		c.ChannelQuery = DefaultQueryChan
	}
	if c.ChannelReply == nil {
		c.ChannelReply = DefaultReplyChan
	}
	return c.Query()
}

func (c *Client) Do(m *Msg, addr string) {
	// addr !!!
	if c.ChannelQuery == nil {
		DefaultQueryChan <- m
	}
}

func ListenAndQuery(c chan *Msg, handler QueryHandler) {
	client := &Client{ChannelQuery: c, Handler: handler}
	go client.ListenAndQuery()
}

func (w *reply) WriteMessage(m *Msg) {
	// Write to the channel
	w.Client.ChannelReply <- []*Msg{m}
}

func (w *reply) WriteMessages(m []*Msg) {
	// Write to the channel
	w.Client.ChannelReply <- m
}
<commit_msg>Luke, use the slice<commit_after>package dns

// A concurrent client implementation.
// Client sends query to a channel which
// will then handle the query. Returned replys
// are return on another channel. Ready for handling --- same
// setup for server - a HANDLER function that gets run
// when the query returns.

// This completely mirrors server.go impl.
import (
	"os"
	//	"net"
)

type QueryHandler interface {
	QueryDNS(w RequestWriter, q *Msg)
}

// A RequestWriter interface is used by an DNS query handler to
// construct an DNS request.
type RequestWriter interface {
	WriteMessages([]*Msg)
        Write(*Msg)
}

// hijacked connections...?
type reply struct {
	Client *Client
	req    *Msg
}

// QueryMux is an DNS request multiplexer. It matches the
// zone name of each incoming request against a list of
// registered patterns add calls the handler for the pattern
// that most closely matches the zone name.
type QueryMux struct {
	m map[string]QueryHandler
}

// NewQueryMux allocates and returns a new QueryMux.
func NewQueryMux() *QueryMux { return &QueryMux{make(map[string]QueryHandler)} }

// DefaultQueryMux is the default QueryMux used by Query.
var DefaultQueryMux = NewQueryMux()

func newQueryChanSlice() chan []*Msg { return make(chan []*Msg) }
func newQueryChan() chan *Msg { return make(chan *Msg) }

// Default channel to use for the resolver
var DefaultReplyChan = newQueryChanSlice()
var DefaultQueryChan = newQueryChan()

// The HandlerQueryFunc type is an adapter to allow the use of
// ordinary functions as DNS query handlers.  If f is a function
// with the appropriate signature, HandlerQueryFunc(f) is a
// QueryHandler object that calls f.
type HandlerQueryFunc func(RequestWriter, *Msg)

// QueryDNS calls f(w, reg)
func (f HandlerQueryFunc) QueryDNS(w RequestWriter, r *Msg) {
	go f(w, r)
}

func HandleQueryFunc(pattern string, handler func(RequestWriter, *Msg)) {
	DefaultQueryMux.HandleQueryFunc(pattern, handler)
}

// Helper handlers
// Todo

// reusing zoneMatch from server.go
func (mux *QueryMux) match(zone string) QueryHandler {
	var h QueryHandler
	var n = 0
	for k, v := range mux.m {
		if !zoneMatch(k, zone) {
			continue
		}
		if h == nil || len(k) > n {
			n = len(k)
			h = v
		}
	}
	return h
}

func (mux *QueryMux) Handle(pattern string, handler QueryHandler) {
	if pattern == "" {
		panic("dns: invalid pattern " + pattern)
	}
	if pattern[len(pattern)-1] != '.' { // no ending .
		mux.m[pattern+"."] = handler
	} else {
		mux.m[pattern] = handler
	}
}

func (mux *QueryMux) HandleQueryFunc(pattern string, handler func(RequestWriter, *Msg)) {
	mux.Handle(pattern, HandlerQueryFunc(handler))
}

func (mux *QueryMux) QueryDNS(w RequestWriter, request *Msg) {
	h := mux.match(request.Question[0].Name)
	if h == nil {
		//                h = RefusedHandler()
		// something else
	}
	h.QueryDNS(w, request)
}

type Client struct {
	Network      string       // if "tcp" a TCP query will be initiated, otherwise an UDP one
	Attempts     int          // number of attempts
	Retry        bool         // retry with TCP
	ChannelQuery chan *Msg    // read DNS request from this channel
	ChannelReply chan []*Msg    // read DNS request from this channel
	Handler      QueryHandler // handler to invoke, dns.DefaultQueryMux if nil
	ReadTimeout  int64        // the net.Conn.SetReadTimeout value for new connections
	WriteTimeout int64        // the net.Conn.SetWriteTimeout value for new connections
}

// Query accepts incoming DNS request,
// Write to in
// creating a new service thread for each.  The service threads
// read requests and then call handler to reply to them.
// Handler is typically nil, in which case the DefaultServeMux is used.
func Query(c chan *Msg, handler QueryHandler) os.Error {
	client := &Client{ChannelQuery: c, Handler: handler}
	return client.Query()
}

func (c *Client) Query() os.Error {
	handler := c.Handler
	if handler == nil {
		handler = DefaultQueryMux
	}
forever:
	for {
		select {
		case in := <-c.ChannelQuery:
			w := new(reply)
			w.Client = c
			w.req = in
			handler.QueryDNS(w, w.req)
		}
	}
	return nil
}

func (c *Client) ListenAndQuery() os.Error {
	if c.ChannelQuery == nil {
		c.ChannelQuery = DefaultQueryChan
	}
	if c.ChannelReply == nil {
		c.ChannelReply = DefaultReplyChan
	}
	return c.Query()
}

func (c *Client) Do(m *Msg, addr string) {
	// addr !!!
	if c.ChannelQuery == nil {
		DefaultQueryChan <- m
	}
}

func ListenAndQuery(c chan *Msg, handler QueryHandler) {
	client := &Client{ChannelQuery: c, Handler: handler}
	go client.ListenAndQuery()
}

func (w *reply) Write(m *Msg) {
	// Write to the channel
	w.Client.ChannelReply <- []*Msg{w.req, m}
}

func (w *reply) WriteMessages(m []*Msg) {
	// Write to the channel
        m1 := append([]*Msg{w.req}, m...)       // Really the way?
	w.Client.ChannelReply <- m1
}
<|endoftext|>


<commit_before>package wsevent

import (
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"log"
	"net/http"
	"sync"

	ws "github.com/gorilla/websocket"
)

//Client
type Client struct {
	//Session ID
	id string

	conn     *ws.Conn
	connLock *sync.RWMutex
	request  *http.Request
}

type request struct {
	Id   string
	Data json.RawMessage
}

type reply struct {
	Id   string `json:"id"`
	Data string `json:"data,string"`
}

var (
	reqPool   = &sync.Pool{New: func() interface{} { return request{} }}
	replyPool = &sync.Pool{New: func() interface{} { return reply{} }}
)

func genID() string {
	bytes := make([]byte, 32)
	rand.Read(bytes)

	return base64.URLEncoding.EncodeToString(bytes)
}

//Returns the client's unique session ID
func (c *Client) Id() string {
	return c.id
}

// Returns the first http request when established connection.
func (c *Client) Request() *http.Request {
	return c.request
}

func (s *Server) NewClientWithID(upgrader ws.Upgrader, w http.ResponseWriter, r *http.Request, id string) (*Client, error) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		return nil, err
	}

	client := &Client{
		id:       id,
		conn:     conn,
		connLock: new(sync.RWMutex),
		request:  r,
	}
	s.newClient <- client

	return client, nil
}

func (s *Server) NewClient(upgrader ws.Upgrader, w http.ResponseWriter, r *http.Request) (*Client, error) {
	return s.NewClientWithID(upgrader, w, r, genID())
}

func (c *Client) Close() error {
	c.connLock.Lock()
	defer c.connLock.Unlock()
	return c.conn.Close()
}

//A thread-safe variant of WriteMessage
func (c *Client) Emit(data string) error {
	c.connLock.Lock()
	defer c.connLock.Unlock()
	return c.conn.WriteMessage(ws.TextMessage, []byte(data))
}

func (c *Client) cleanup(s *Server) {
	c.conn.Close()

	s.joinedRoomsLock.RLock()
	for _, room := range s.joinedRooms[c.id] {
		//log.Println(room)
		index := -1

		s.roomsLock.Lock()
		for i, client := range s.rooms[room] {
			if client.id == c.id {
				index = i
			}
		}

		s.rooms[room] = append(s.rooms[room][:index], s.rooms[room][index+1:]...)
		s.roomsLock.Unlock()
	}
	s.joinedRoomsLock.RUnlock()

	s.joinedRoomsLock.Lock()
	delete(s.joinedRooms, c.id)
	s.joinedRoomsLock.Unlock()

	if s.OnDisconnect != nil {
		s.OnDisconnect(c.id)
	}
}

func (c *Client) listener(s *Server) {
	for {
		mtype, data, err := c.conn.ReadMessage()
		if err != nil {
			c.cleanup(s)
			return
		}

		js := reqPool.Get().(request)
		err = json.Unmarshal(data, &js)

		if err != nil || mtype != ws.TextMessage {
			log.Println(err)
			continue
		}

		callName := s.Extractor(js.Data)

		s.handlersLock.RLock()
		f, ok := s.handlers[callName]
		s.handlersLock.RUnlock()

		if !ok {
			if s.DefaultHandler != nil {
				f = s.DefaultHandler
				goto call
			}
			continue
		}
	call:
		go func() {
			rtrn := f(s, c, js.Data)

			replyJs := replyPool.Get().(reply)
			replyJs.Id = js.Id
			replyJs.Data = string(rtrn)

			bytes, _ := json.Marshal(replyJs)
			c.Emit(string(bytes))

			reqPool.Put(js)
			replyPool.Put(replyJs)
		}()
	}
}
<commit_msg>Add basic rate limiting measures<commit_after>package wsevent

import (
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"log"
	"net/http"
	"sync"
	"time"

	ws "github.com/gorilla/websocket"
)

//Client
type Client struct {
	//Session ID
	id string

	conn     *ws.Conn
	connLock *sync.RWMutex
	request  *http.Request
}

type request struct {
	Id   string
	Data json.RawMessage
}

type reply struct {
	Id   string `json:"id"`
	Data string `json:"data,string"`
}

var (
	reqPool   = &sync.Pool{New: func() interface{} { return request{} }}
	replyPool = &sync.Pool{New: func() interface{} { return reply{} }}
)

func genID() string {
	bytes := make([]byte, 32)
	rand.Read(bytes)

	return base64.URLEncoding.EncodeToString(bytes)
}

//Returns the client's unique session ID
func (c *Client) Id() string {
	return c.id
}

// Returns the first http request when established connection.
func (c *Client) Request() *http.Request {
	return c.request
}

func (s *Server) NewClientWithID(upgrader ws.Upgrader, w http.ResponseWriter, r *http.Request, id string) (*Client, error) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		return nil, err
	}

	client := &Client{
		id:       id,
		conn:     conn,
		connLock: new(sync.RWMutex),
		request:  r,
	}
	s.newClient <- client

	return client, nil
}

func (s *Server) NewClient(upgrader ws.Upgrader, w http.ResponseWriter, r *http.Request) (*Client, error) {
	return s.NewClientWithID(upgrader, w, r, genID())
}

func (c *Client) Close() error {
	c.connLock.Lock()
	defer c.connLock.Unlock()
	return c.conn.Close()
}

//A thread-safe variant of WriteMessage
func (c *Client) Emit(data string) error {
	c.connLock.Lock()
	defer c.connLock.Unlock()
	return c.conn.WriteMessage(ws.TextMessage, []byte(data))
}

type emitJS struct {
	Id   int         `json:"id"`
	Data interface{} `json:"data"`
}

var emitPool = &sync.Pool{New: func() interface{} { return emitJS{} }}

//A thread-safe variant of EmitJSON
func (c *Client) EmitJSON(v interface{}) error {
	c.connLock.Lock()
	defer c.connLock.Unlock()

	js := emitPool.Get().(emitJS)
	defer emitPool.Put(js)

	js.Id = -1
	js.Data = v

	return c.conn.WriteJSON(js)
}

func (c *Client) cleanup(s *Server) {
	c.conn.Close()

	s.joinedRoomsLock.RLock()
	for _, room := range s.joinedRooms[c.id] {
		//log.Println(room)
		index := -1

		s.roomsLock.Lock()
		for i, client := range s.rooms[room] {
			if client.id == c.id {
				index = i
			}
		}

		s.rooms[room] = append(s.rooms[room][:index], s.rooms[room][index+1:]...)
		s.roomsLock.Unlock()
	}
	s.joinedRoomsLock.RUnlock()

	s.joinedRoomsLock.Lock()
	delete(s.joinedRooms, c.id)
	s.joinedRoomsLock.Unlock()

	if s.OnDisconnect != nil {
		s.OnDisconnect(c.id)
	}
}

func (c *Client) listener(s *Server) {
	throttle := time.NewTicker(time.Millisecond * 10)
	defer throttle.Stop()
	for {
		<-throttle.C
		mtype, data, err := c.conn.ReadMessage()
		if err != nil {
			c.cleanup(s)
			return
		}
		if mtype != ws.TextMessage {
			c.conn.Close()
			return
		}

		js := reqPool.Get().(request)

		if err := json.Unmarshal(data, &js); err != nil {
			continue
		}

		callName := s.Extractor(js.Data)

		s.handlersLock.RLock()
		f, ok := s.handlers[callName]
		s.handlersLock.RUnlock()

		if !ok {
			if s.DefaultHandler != nil {
				f = s.DefaultHandler
				goto call
			}
			continue
		}
	call:
		go func() {
			rtrn := f(s, c, js.Data)
			replyJs := replyPool.Get().(reply)
			replyJs.Id = js.Id
			replyJs.Data = string(rtrn)

			bytes, _ := json.Marshal(replyJs)
			c.Emit(string(bytes))

			reqPool.Put(js)
			replyPool.Put(replyJs)
		}()
	}
}
<|endoftext|>


<commit_before>package groveclient

import (
	"strings"

	pc "github.com/t11e/go-pebbleclient"
)

type Client struct {
	c pc.Client
}

type GetOptions struct {
	Raw *bool
}

type GetManyOptions struct {
	Limit *int
	Raw   *bool
}

type GetManyOutput struct {
	Posts []PostItem `json:"posts"`
}

func New(client pc.Client) (*Client, error) {
	return &Client{client.Options(pc.Options{
		ServiceName: "grove",
		ApiVersion:  1,
	})}, nil
}

func (client *Client) Get(uid string, options GetOptions) (*PostItem, error) {
	params := pc.Params{
		"raw": options.Raw != nil && *options.Raw,
		"uid": uid,
	}

	var out PostItem
	err := client.c.Get("/posts/:uid", &pc.RequestOptions{
		Params: params,
	}, &out)
	if err != nil {
		return nil, err
	}
	return &out, err
}

func (client *Client) GetMany(uids []string, options GetManyOptions) (*GetManyOutput, error) {
	uidList := strings.Join(uids, ",")
	if len(uids) > 1 {
		uidList = uidList + ","
	}

	params := pc.Params{
		"raw":  options.Raw != nil && *options.Raw,
		"uids": uidList,
	}
	if options.Limit != nil {
		params["limit"] = *options.Limit
	}

	var out GetManyOutput
	err := client.c.Get("/posts/:uids", &pc.RequestOptions{
		Params: params,
	}, &out)
	if err != nil {
		return nil, err
	}
	return &out, err
}
<commit_msg>Provide interface so we can mock.<commit_after>package groveclient

import (
	"strings"

	pc "github.com/t11e/go-pebbleclient"
)

type Client interface {
}

type client struct {
	c pc.Client
}

type GetOptions struct {
	Raw *bool
}

type GetManyOptions struct {
	Limit *int
	Raw   *bool
}

type GetManyOutput struct {
	Posts []PostItem `json:"posts"`
}

func New(pebbleClient pc.Client) (Client, error) {
	return &client{pebbleClient.Options(pc.Options{
		ServiceName: "grove",
		ApiVersion:  1,
	})}, nil
}

func (c *client) Get(uid string, options GetOptions) (*PostItem, error) {
	params := pc.Params{
		"raw": options.Raw != nil && *options.Raw,
		"uid": uid,
	}

	var out PostItem
	err := c.c.Get("/posts/:uid", &pc.RequestOptions{
		Params: params,
	}, &out)
	if err != nil {
		return nil, err
	}
	return &out, err
}

func (c *client) GetMany(uids []string, options GetManyOptions) (*GetManyOutput, error) {
	uidList := strings.Join(uids, ",")
	if len(uids) > 1 {
		uidList = uidList + ","
	}

	params := pc.Params{
		"raw":  options.Raw != nil && *options.Raw,
		"uids": uidList,
	}
	if options.Limit != nil {
		params["limit"] = *options.Limit
	}

	var out GetManyOutput
	err := c.c.Get("/posts/:uids", &pc.RequestOptions{
		Params: params,
	}, &out)
	if err != nil {
		return nil, err
	}
	return &out, err
}
<|endoftext|>


<commit_before>package beanspike

import (
	"context"
	"encoding/binary"
	"encoding/json"
	"errors"
	"fmt"
	"math/rand"
	"sync"
	"sync/atomic"
	"time"

	"github.com/redsift/clockwork/workerpool"
	"github.com/redsift/clockwork/workerpool/bdeq"

	srand "crypto/rand"

	"github.com/redsift/clockwork/cron"

	"github.com/bluele/gcache"
)

const (
	IdleInTubeTime = 2 * time.Second
	OutTubeTTR     = 30 * time.Second
)

var (
	ErrAlreadyConnected = errors.New("already connected")
	ErrNotConnected     = errors.New("not connected")
)

type TubeID string

// JobDecoder defines interface app should implement for decoding incoming jobs into the app specific objects.
type JobDecoder interface {
	// Decode unmarshalls job's payload.
	Decode([]byte) (interface{}, error)
}

type JobDecoderFunc func([]byte) (interface{}, error)

func (f JobDecoderFunc) Decode(b []byte) (interface{}, error) { return f(b) }

type JobEventListener interface {
	// OnRelease will be called when handler context canceled
	OnRelease(*Job, interface{})
}

type JobHandler interface {
	// Handle processes an incoming job.
	// Application should watch given context; it is being canceled on job release.
	Handle(context.Context, *ManagedJob, interface{})
}

type JobHandlerFunc func(context.Context, *ManagedJob, interface{})

func (f JobHandlerFunc) Handle(ctx context.Context, job *ManagedJob, v interface{}) { f(ctx, job, v) }

type ManagedJob struct {
	Job
	cancel func()
}

// StopAndRelease provides application way to stop job keeping activity and release the job
func (job *ManagedJob) StopAndRelease() { job.cancel() }

func NewManagedJob(job *Job, cancel func()) *ManagedJob {
	return &ManagedJob{
		Job:    *job,
		cancel: cancel,
	}
}

var retryDuration = func() func(uint) time.Duration {
	rng := &sync.Pool{New: func() interface{} {
		var seed [8]byte
		if _, err := srand.Read(seed[:]); err != nil {
			panic("cannot seed math/rand package with cryptographically secure random number generator")
		}

		return rand.New(rand.NewSource(time.Now().UnixNano() * int64(binary.LittleEndian.Uint64(seed[:]))))
	}}

	return func(c uint) time.Duration {
		if c > 10 {
			c = 10
		}
		rnd := rng.Get().(*rand.Rand)
		d := time.Duration(rnd.Intn(1<<c)) * time.Second
		rng.Put(rnd)
		return d
	}
}()

type Client struct {
	ctx          context.Context
	cancel       context.CancelFunc
	conn         *Conn
	outTubes     gcache.Cache
	cron         *cron.Cron
	reservedJobs map[int64]*ManagedJob
	lock         *sync.RWMutex
	shutdown     int32
}

func NewClient(ctx context.Context) *Client {
	ctx, cancel := context.WithCancel(ctx)

	client := &Client{
		ctx:          ctx,
		cancel:       cancel,
		reservedJobs: make(map[int64]*ManagedJob),
		lock:         new(sync.RWMutex),
		shutdown:     0,
	}

	client.cron = cron.New(cron.HandlerFunc(func(_ time.Time, keys []interface{}) {
		for _, k := range keys {
			client.lock.RLock()
			job, found := client.reservedJobs[k.(int64)]
			client.lock.RUnlock()
			if !found {
				continue
			}
			if err := job.Touch(); err != nil {
				job.StopAndRelease()
			}
		}
	}), cron.WithPrecision(time.Second), cron.WithCalendarMinCap(1<<13))

	// Hope it was a correct answer
	client.outTubes = gcache.New(42).LRU().LoaderFunc(func(k interface{}) (interface{}, error) {
		if client.conn == nil {
			return nil, ErrNotConnected
		}
		return client.conn.Use(k.(string))
	}).Build()

	return client
}

func (c *Client) Connect(host string, port int, statsHandler func(string, string, float64)) error {
	if c.conn != nil {
		return ErrAlreadyConnected
	}
	conn, err := Dial("", host, port, statsHandler)
	if err != nil {
		return err
	}
	c.conn = conn
	return nil
}

func (c *Client) Close() {
	atomic.StoreInt32(&c.shutdown, 1) // disable normal job release
	c.cancel()                        // notify application (all jobs)
	c.cron.Stop()                     // stop keeping

	wp := workerpool.New(128, bdeq.New(1<<12))
	wp.Start()

	var wg sync.WaitGroup

	c.lock.Lock()
	for _, v := range c.reservedJobs {
		job := v
		wg.Add(1)
		wp.Submit(func() {
			// release immediately
			if err := job.Release(0); err != nil {
				DefaultErrorLogger.LogError(fmt.Errorf("stopping: %d release error: %w", job.ID, err))
			}
			wg.Done()
		}, 0)
	}
	c.lock.Unlock()

	wg.Wait()

	// TODO close connection
	// TODO stop workerpool when queue is empty
}

func (c *Client) StartReserveAndKeepLoop(id TubeID, dec JobDecoder, handler JobHandler, batchSize int) error {
	tube, err := c.conn.Use(string(id))
	if err != nil {
		return err
	}

	managingHandler := JobHandlerFunc(func(ctx context.Context, job *ManagedJob, v interface{}) {
		ctx, cancel := context.WithCancel(ctx)

		// wrap job into ManagedJob with a new cancel func
		managedJob := NewManagedJob(&job.Job, func() {
			cancel() // notify application

			// remove from reserved jobs
			c.lock.Lock()
			delete(c.reservedJobs, job.ID)
			c.lock.Unlock()

			// stop cron task
			c.cron.Remove(job.ID)

			// release us usual if not shutdown; otherwise job will be released immediately
			if atomic.LoadInt32(&c.shutdown) == 0 {
				_ = job.ReleaseWithRetry(retryDuration(uint(job.Retries)), true, true)
			}

			// call the original cancel func
			job.StopAndRelease()
		})

		// add to reserved jobs
		c.lock.Lock()
		c.reservedJobs[job.ID] = managedJob
		c.lock.Unlock()

		// schedule regular touch; skip if job has been scheduled already
		_ = c.cron.Add(cron.ScheduleFunc(func(now time.Time) time.Time {
			return now.Add(job.TTR / 2)
		}), job.ID)

		// call the original handler
		handler.Handle(ctx, managedJob, v)
	})

	go func() {
		ticker := time.Tick(IdleInTubeTime)
	LOOP:
		for {
			select {
			case <-c.ctx.Done():
				return
			case <-ticker:
				for c.ctx.Err() == nil {
					n := tube.ReserveBatch(c.ctx, dec, managingHandler, batchSize)
					if n == 0 {
						continue LOOP
					}
				}
			}
		}
	}()

	return nil
}

func (c *Client) Put(tubeID string, v interface{}, metadata string, tob int64) (int64, error) {
	var (
		err  error
		tube interface{}
	)
	tube, err = c.outTubes.Get(tubeID)
	if err != nil {
		return 0, err
	}

	var b []byte
	b, err = json.Marshal(v)
	if err != nil {
		return 0, err
	}

	var id int64
	id, err = tube.(*Tube).Put(b, 0, OutTubeTTR, true, metadata, tob)
	return id, err
}
<commit_msg>[minor] typo and delete unused JobEventListener<commit_after>package beanspike

import (
	"context"
	"encoding/binary"
	"encoding/json"
	"errors"
	"fmt"
	"math/rand"
	"sync"
	"sync/atomic"
	"time"

	"github.com/redsift/clockwork/workerpool"
	"github.com/redsift/clockwork/workerpool/bdeq"

	srand "crypto/rand"

	"github.com/redsift/clockwork/cron"

	"github.com/bluele/gcache"
)

const (
	IdleInTubeTime = 2 * time.Second
	OutTubeTTR     = 30 * time.Second
)

var (
	ErrAlreadyConnected = errors.New("already connected")
	ErrNotConnected     = errors.New("not connected")
)

type TubeID string

// JobDecoder defines interface app should implement for decoding incoming jobs into the app specific objects.
type JobDecoder interface {
	// Decode unmarshalls job's payload.
	Decode([]byte) (interface{}, error)
}

type JobDecoderFunc func([]byte) (interface{}, error)

func (f JobDecoderFunc) Decode(b []byte) (interface{}, error) { return f(b) }

type JobHandler interface {
	// Handle processes an incoming job.
	// Application should watch given context; it is being canceled on job release.
	Handle(context.Context, *ManagedJob, interface{})
}

type JobHandlerFunc func(context.Context, *ManagedJob, interface{})

func (f JobHandlerFunc) Handle(ctx context.Context, job *ManagedJob, v interface{}) { f(ctx, job, v) }

type ManagedJob struct {
	*Job
	cancel func()
}

// StopAndRelease provides application way to stop job keeping activity and release the job
func (job *ManagedJob) StopAndRelease() { job.cancel() }

func NewManagedJob(job *Job, cancel func()) *ManagedJob {
	return &ManagedJob{
		Job:    job,
		cancel: cancel,
	}
}

var retryDuration = func() func(uint) time.Duration {
	rng := &sync.Pool{New: func() interface{} {
		var seed [8]byte
		if _, err := srand.Read(seed[:]); err != nil {
			panic("cannot seed math/rand package with cryptographically secure random number generator")
		}

		return rand.New(rand.NewSource(time.Now().UnixNano() * int64(binary.LittleEndian.Uint64(seed[:]))))
	}}

	return func(c uint) time.Duration {
		if c > 10 {
			c = 10
		}
		rnd := rng.Get().(*rand.Rand)
		d := time.Duration(rnd.Intn(1<<c)) * time.Second
		rng.Put(rnd)
		return d
	}
}()

type Client struct {
	ctx          context.Context
	cancel       context.CancelFunc
	conn         *Conn
	outTubes     gcache.Cache
	cron         *cron.Cron
	reservedJobs map[int64]*ManagedJob
	lock         *sync.RWMutex
	shutdown     int32
}

func NewClient(ctx context.Context) *Client {
	ctx, cancel := context.WithCancel(ctx)

	client := &Client{
		ctx:          ctx,
		cancel:       cancel,
		reservedJobs: make(map[int64]*ManagedJob),
		lock:         new(sync.RWMutex),
		shutdown:     0,
	}

	client.cron = cron.New(cron.HandlerFunc(func(_ time.Time, keys []interface{}) {
		for _, k := range keys {
			client.lock.RLock()
			job, found := client.reservedJobs[k.(int64)]
			client.lock.RUnlock()
			if !found {
				continue
			}
			if err := job.Touch(); err != nil {
				job.StopAndRelease()
			}
		}
	}), cron.WithPrecision(time.Second), cron.WithCalendarMinCap(1<<13))

	// Hope it was a correct answer
	client.outTubes = gcache.New(42).LRU().LoaderFunc(func(k interface{}) (interface{}, error) {
		if client.conn == nil {
			return nil, ErrNotConnected
		}
		return client.conn.Use(k.(string))
	}).Build()

	return client
}

func (c *Client) Connect(host string, port int, statsHandler func(string, string, float64)) error {
	if c.conn != nil {
		return ErrAlreadyConnected
	}
	conn, err := Dial("", host, port, statsHandler)
	if err != nil {
		return err
	}
	c.conn = conn
	return nil
}

func (c *Client) Close() {
	atomic.StoreInt32(&c.shutdown, 1) // disable normal job release
	c.cancel()                        // notify application (all jobs)
	c.cron.Stop()                     // stop keeping

	wp := workerpool.New(128, bdeq.New(1<<12))
	wp.Start()

	var wg sync.WaitGroup

	c.lock.Lock()
	for _, v := range c.reservedJobs {
		job := v
		wg.Add(1)
		wp.Submit(func() {
			// release immediately
			if err := job.Release(0); err != nil {
				DefaultErrorLogger.LogError(fmt.Errorf("stopping: %d release error: %w", job.ID, err))
			}
			wg.Done()
		}, 0)
	}
	c.lock.Unlock()

	wg.Wait()

	// TODO close connection
	// TODO stop worker pool when queue is empty
}

func (c *Client) StartReserveAndKeepLoop(id TubeID, dec JobDecoder, handler JobHandler, batchSize int) error {
	tube, err := c.conn.Use(string(id))
	if err != nil {
		return err
	}

	managingHandler := JobHandlerFunc(func(ctx context.Context, job *ManagedJob, v interface{}) {
		ctx, cancel := context.WithCancel(ctx)

		// wrap job into ManagedJob with a new cancel func
		managedJob := NewManagedJob(job.Job, func() {
			cancel() // notify application

			// remove from reserved jobs
			c.lock.Lock()
			delete(c.reservedJobs, job.ID)
			c.lock.Unlock()

			// stop cron task
			c.cron.Remove(job.ID)

			// release us usual if not shutdown; otherwise job will be released immediately
			if atomic.LoadInt32(&c.shutdown) == 0 {
				_ = job.ReleaseWithRetry(retryDuration(uint(job.Retries)), true, true)
			}

			// call the original cancel func
			job.StopAndRelease()
		})

		// add to reserved jobs
		c.lock.Lock()
		c.reservedJobs[job.ID] = managedJob
		c.lock.Unlock()

		// schedule regular touch; skip if job has been scheduled already
		_ = c.cron.Add(cron.ScheduleFunc(func(now time.Time) time.Time {
			return now.Add(job.TTR / 2)
		}), job.ID)

		// call the original handler
		handler.Handle(ctx, managedJob, v)
	})

	go func() {
		ticker := time.Tick(IdleInTubeTime)
	LOOP:
		for {
			select {
			case <-c.ctx.Done():
				return
			case <-ticker:
				for c.ctx.Err() == nil {
					n := tube.ReserveBatch(c.ctx, dec, managingHandler, batchSize)
					if n == 0 {
						continue LOOP
					}
				}
			}
		}
	}()

	return nil
}

func (c *Client) Put(tubeID string, v interface{}, metadata string, tob int64) (int64, error) {
	var (
		err  error
		tube interface{}
	)
	tube, err = c.outTubes.Get(tubeID)
	if err != nil {
		return 0, err
	}

	var b []byte
	b, err = json.Marshal(v)
	if err != nil {
		return 0, err
	}

	var id int64
	id, err = tube.(*Tube).Put(b, 0, OutTubeTTR, true, metadata, tob)
	return id, err
}
<|endoftext|>


<commit_before>// Copyright 2014, Truveris Inc. All Rights Reserved.
// Use of this source code is governed by the ISC license in the LICENSE file.

package sqs

import (
	"encoding/xml"
	"errors"
	"io/ioutil"
	"net"
	"net/http"
	"strings"
	"time"

	"github.com/mikedewar/aws4"
)

var (
	// Ref:
	// http://docs.aws.amazon.com/general/latest/gr/rande.html#sqs_region
	EndPoints = map[string]string{
		"ap-northeast-1": "https://sqs.ap-northeast-1.amazonaws.com",
		"ap-southeast-1": "https://sqs.ap-southeast-1.amazonaws.com",
		"ap-southeast-2": "https://sqs.ap-southeast-2.amazonaws.com",
		"eu-west-1":      "https://sqs.eu-west-1.amazonaws.com",
		"sa-east-1":      "https://sqs.sa-east-1.amazonaws.com",
		"us-east-1":      "https://sqs.us-east-1.amazonaws.com",
		"us-west-1":      "https://sqs.us-west-1.amazonaws.com",
		"us-west-2":      "https://sqs.us-west-2.amazonaws.com",
	}
)

type Client struct {
	Aws4Client  *aws4.Client
	EndPointURL string
}

func dialTimeout(network, addr string) (net.Conn, error) {
	return net.DialTimeout(network, addr, 25 * time.Second)
}

func NewClient(AccessKey, SecretKey, RegionCode string) (*Client, error) {
	keys := &aws4.Keys{
		AccessKey: AccessKey,
		SecretKey: SecretKey,
	}

	EndPointURL := EndPoints[RegionCode]
	if EndPointURL == "" {
		return nil, errors.New("Unknown region: " + RegionCode)
	}

	transport := &http.Transport{Dial: dialTimeout}
	client := &http.Client{Transport: transport}

	return &Client{
		Aws4Client:  &aws4.Client{Keys: keys, Client: client},
		EndPointURL: EndPointURL,
	}, nil
}

// Simple wrapper around the aws4 client Post() but less verbose.
func (client *Client) Post(queueURL, data string) (*http.Response, error) {
	return client.Aws4Client.Post(queueURL, SQSContentType,
		strings.NewReader(data))
}

// Simple wrapper around the aws4 Get() to keep it consistent.
func (client *Client) Get(url string) (*http.Response, error) {
	return client.Aws4Client.Get(url)
}

// Return a single message body, with its ReceiptHandle. A lack of message is
// not considered an error but *Message will be nil.
func (client *Client) GetSingleMessageFromRequest(request *Request) (*Message, error) {
	var m ReceiveMessageResult

	// These two settings are required for this function to function.
	request.Set("MaxNumberOfMessages", "1")
	request.Set("AttributeName", "SenderId")

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

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode != 200 {
		return nil, errors.New(string(body))
	}

	err = xml.Unmarshal(body, &m)
	if err != nil {
		return nil, err
	}

	// The API call is build to only return one or zero messages.
	if len(m.Bodies) != 1 || len(m.Values) != 1 {
		return nil, nil
	}

	msg := &Message{
		QueueURL:      request.QueueURL,
		Body:          m.Bodies[0],
		ReceiptHandle: m.ReceiptHandles[0],
		UserID:        m.Values[0],
	}

	return msg, nil
}

// Return a single message body, with its ReceiptHandle. A lack of message is
// not considered an error but the return message will be nil.
func (client *Client) GetSingleMessage(url string) (*Message, error) {
	request := NewReceiveMessageRequest(url)
	return client.GetSingleMessageFromRequest(request)
}

// Conduct a DeleteMessage API call on the given queue, using the receipt
// handle from a previously fetched message.
func (client *Client) DeleteMessageFromReceipt(queueURL, receipt string) error {
	url := NewDeleteMessageRequest(queueURL, receipt).URL()

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

	return nil
}

// Conduct a DeleteMessage API call on the given queue, using the receipt
// handle from a previously fetched message.
func (client *Client) DeleteMessage(msg *Message) error {
	url := NewDeleteMessageRequest(msg.QueueURL, msg.ReceiptHandle).URL()

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

	return nil
}

// Conduct a SendMessage API call (POST) on the given queue.
func (client *Client) SendMessage(queueURL, message string) error {
	data := NewSendMessageRequest(queueURL, message).Query()

	resp, err := client.Post(queueURL, data)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}

	if resp.StatusCode != 200 {
		return errors.New(string(body))
	}

	return nil
}

// Get the queue URL from its name.
func (client *Client) GetQueueURL(name string) (string, error) {
	var parsedResponse GetQueueURLResult
	url := NewGetQueueURLRequest(client.EndPointURL, name).URL()

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

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}

	if resp.StatusCode != 200 {
		return "", errors.New(string(body))
	}

	err = xml.Unmarshal(body, &parsedResponse)
	if err != nil {
		return "", err
	}

	return parsedResponse.QueueURL, nil
}

// Create a queue using the provided attributes and return its URL. This
// function can be used to obtain the QueueURL for a queue even if it already
// exists.
func (client *Client) CreateQueueWithAttributes(name string, attributes CreateQueueAttributes) (string, error) {
	var parsedResponse CreateQueueResult
	url := buildCreateQueueURL(client.EndPointURL, name, attributes)

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

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}

	if resp.StatusCode != 200 {
		return "", errors.New(string(body))
	}

	err = xml.Unmarshal(body, &parsedResponse)
	if err != nil {
		return "", err
	}

	return parsedResponse.QueueURL, nil
}

// Create a queue with default parameters and return its URL. This function can
// be used to obtain the QueueURL for a queue even if it already exists.
func (client *Client) CreateQueue(name string) (string, error) {
	url, err := client.CreateQueueWithAttributes(name, CreateQueueAttributes{})
	return url, err
}
<commit_msg>Add ResponseHeaderTimeout to the http client.<commit_after>// Copyright 2014, Truveris Inc. All Rights Reserved.
// Use of this source code is governed by the ISC license in the LICENSE file.

package sqs

import (
	"encoding/xml"
	"errors"
	"io/ioutil"
	"net"
	"net/http"
	"strings"
	"time"

	"github.com/mikedewar/aws4"
)

var (
	// Ref:
	// http://docs.aws.amazon.com/general/latest/gr/rande.html#sqs_region
	EndPoints = map[string]string{
		"ap-northeast-1": "https://sqs.ap-northeast-1.amazonaws.com",
		"ap-southeast-1": "https://sqs.ap-southeast-1.amazonaws.com",
		"ap-southeast-2": "https://sqs.ap-southeast-2.amazonaws.com",
		"eu-west-1":      "https://sqs.eu-west-1.amazonaws.com",
		"sa-east-1":      "https://sqs.sa-east-1.amazonaws.com",
		"us-east-1":      "https://sqs.us-east-1.amazonaws.com",
		"us-west-1":      "https://sqs.us-west-1.amazonaws.com",
		"us-west-2":      "https://sqs.us-west-2.amazonaws.com",
	}

	HTTPTimeout = 25 * time.Second
)

type Client struct {
	Aws4Client  *aws4.Client
	EndPointURL string
}

func dialTimeout(network, addr string) (net.Conn, error) {
	return net.DialTimeout(network, addr, HTTPTimeout)
}

func NewClient(AccessKey, SecretKey, RegionCode string) (*Client, error) {
	keys := &aws4.Keys{
		AccessKey: AccessKey,
		SecretKey: SecretKey,
	}

	EndPointURL := EndPoints[RegionCode]
	if EndPointURL == "" {
		return nil, errors.New("Unknown region: " + RegionCode)
	}

	transport := &http.Transport{Dial: dialTimeout, ResponseHeaderTimeout: HTTPTimeout}
	client := &http.Client{Transport: transport}

	return &Client{
		Aws4Client:  &aws4.Client{Keys: keys, Client: client},
		EndPointURL: EndPointURL,
	}, nil
}

// Simple wrapper around the aws4 client Post() but less verbose.
func (client *Client) Post(queueURL, data string) (*http.Response, error) {
	return client.Aws4Client.Post(queueURL, SQSContentType,
		strings.NewReader(data))
}

// Simple wrapper around the aws4 Get() to keep it consistent.
func (client *Client) Get(url string) (*http.Response, error) {
	return client.Aws4Client.Get(url)
}

// Return a single message body, with its ReceiptHandle. A lack of message is
// not considered an error but *Message will be nil.
func (client *Client) GetSingleMessageFromRequest(request *Request) (*Message, error) {
	var m ReceiveMessageResult

	// These two settings are required for this function to function.
	request.Set("MaxNumberOfMessages", "1")
	request.Set("AttributeName", "SenderId")

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

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode != 200 {
		return nil, errors.New(string(body))
	}

	err = xml.Unmarshal(body, &m)
	if err != nil {
		return nil, err
	}

	// The API call is build to only return one or zero messages.
	if len(m.Bodies) != 1 || len(m.Values) != 1 {
		return nil, nil
	}

	msg := &Message{
		QueueURL:      request.QueueURL,
		Body:          m.Bodies[0],
		ReceiptHandle: m.ReceiptHandles[0],
		UserID:        m.Values[0],
	}

	return msg, nil
}

// Return a single message body, with its ReceiptHandle. A lack of message is
// not considered an error but the return message will be nil.
func (client *Client) GetSingleMessage(url string) (*Message, error) {
	request := NewReceiveMessageRequest(url)
	return client.GetSingleMessageFromRequest(request)
}

// Conduct a DeleteMessage API call on the given queue, using the receipt
// handle from a previously fetched message.
func (client *Client) DeleteMessageFromReceipt(queueURL, receipt string) error {
	url := NewDeleteMessageRequest(queueURL, receipt).URL()

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

	return nil
}

// Conduct a DeleteMessage API call on the given queue, using the receipt
// handle from a previously fetched message.
func (client *Client) DeleteMessage(msg *Message) error {
	url := NewDeleteMessageRequest(msg.QueueURL, msg.ReceiptHandle).URL()

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

	return nil
}

// Conduct a SendMessage API call (POST) on the given queue.
func (client *Client) SendMessage(queueURL, message string) error {
	data := NewSendMessageRequest(queueURL, message).Query()

	resp, err := client.Post(queueURL, data)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}

	if resp.StatusCode != 200 {
		return errors.New(string(body))
	}

	return nil
}

// Get the queue URL from its name.
func (client *Client) GetQueueURL(name string) (string, error) {
	var parsedResponse GetQueueURLResult
	url := NewGetQueueURLRequest(client.EndPointURL, name).URL()

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

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}

	if resp.StatusCode != 200 {
		return "", errors.New(string(body))
	}

	err = xml.Unmarshal(body, &parsedResponse)
	if err != nil {
		return "", err
	}

	return parsedResponse.QueueURL, nil
}

// Create a queue using the provided attributes and return its URL. This
// function can be used to obtain the QueueURL for a queue even if it already
// exists.
func (client *Client) CreateQueueWithAttributes(name string, attributes CreateQueueAttributes) (string, error) {
	var parsedResponse CreateQueueResult
	url := buildCreateQueueURL(client.EndPointURL, name, attributes)

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

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}

	if resp.StatusCode != 200 {
		return "", errors.New(string(body))
	}

	err = xml.Unmarshal(body, &parsedResponse)
	if err != nil {
		return "", err
	}

	return parsedResponse.QueueURL, nil
}

// Create a queue with default parameters and return its URL. This function can
// be used to obtain the QueueURL for a queue even if it already exists.
func (client *Client) CreateQueue(name string) (string, error) {
	url, err := client.CreateQueueWithAttributes(name, CreateQueueAttributes{})
	return url, err
}
<|endoftext|>


<commit_before>package req

import (
	"encoding/json"
	"golang.org/x/net/publicsuffix"
	"io"
	"net/http"
	"net/http/cookiejar"
	"os"
	"time"
)

// DefaultClient returns the global default Client.
func DefaultClient() *Client {
	return defaultClient
}

// SetDefaultClient override the global default Client.
func SetDefaultClient(c *Client) {
	if c != nil {
		defaultClient = c
	}
}

var defaultClient *Client = C()

// Client is the req's http client.
type Client struct {
	log          Logger
	t            *Transport
	t2           *http2Transport
	dumpOptions  *DumpOptions
	httpClient   *http.Client
	jsonDecoder  *json.Decoder
	commonHeader map[string]string
}

func copyCommonHeader(h map[string]string) map[string]string {
	if h == nil {
		return nil
	}
	m := make(map[string]string)
	for k, v := range h {
		m[k] = v
	}
	return m
}

// R create a new request.
func (c *Client) R() *Request {
	req := &http.Request{
		Header:     make(http.Header),
		Proto:      "HTTP/1.1",
		ProtoMajor: 1,
		ProtoMinor: 1,
	}
	return &Request{
		client:      c,
		httpRequest: req,
	}
}

func (c *Client) AutoDiscardResponseBody() *Client {
	c.GetResponseOptions().AutoDiscard = true
	return c
}

// TestMode is like DebugMode, but discard response body, so you can
// dump responses without read response body
func (c *Client) TestMode() *Client {
	return c.DebugMode().AutoDiscardResponseBody()
}

const (
	userAgentFirefox = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10.15; rv:95.0) Gecko/20100101 Firefox/95.0"
	userAgentChrome  = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/97.0.4692.71 Safari/537.36"
)

// DebugMode enables dump for requests and responses, and set user
// agent to pretend to be a web browser, Avoid returning abnormal
// data from some sites.
func (c *Client) DebugMode() *Client {
	return c.AutoDecodeTextType().
		Dump(true).
		SetLogger(NewLogger(os.Stdout)).
		UserAgent(userAgentChrome)
}

// SetLogger set the logger for req.
func (c *Client) SetLogger(log Logger) *Client {
	if log == nil {
		return c
	}
	c.log = log
	return c
}

func (c *Client) GetResponseOptions() *ResponseOptions {
	if c.t.ResponseOptions == nil {
		c.t.ResponseOptions = &ResponseOptions{}
	}
	return c.t.ResponseOptions
}

// ResponseOptions set the ResponseOptions for the underlying Transport.
func (c *Client) SetResponseOptions(opt *ResponseOptions) *Client {
	if opt == nil {
		return c
	}
	c.t.ResponseOptions = opt
	return c
}

// Timeout set the timeout for all requests.
func (c *Client) Timeout(d time.Duration) *Client {
	c.httpClient.Timeout = d
	return c
}

func (c *Client) GetDumpOptions() *DumpOptions {
	if c.dumpOptions == nil {
		c.dumpOptions = newDefaultDumpOptions()
	}
	return c.dumpOptions
}

func (c *Client) enableDump() {
	if c.t.dump != nil { // dump already started
		return
	}
	c.t.EnableDump(c.GetDumpOptions())
}

// DumpToFile indicates that the content should dump to the specified filename.
func (c *Client) DumpToFile(filename string) *Client {
	file, err := os.Create(filename)
	if err != nil {
		logf(c.log, "create dump file error: %v", err)
		return c
	}
	c.GetDumpOptions().Output = file
	return c
}

// DumpTo indicates that the content should dump to the specified destination.
func (c *Client) DumpTo(output io.Writer) *Client {
	c.GetDumpOptions().Output = output
	c.enableDump()
	return c
}

// DumpAsync indicates that the dump should be done asynchronously,
// can be used for debugging in production environment without
// affecting performance.
func (c *Client) DumpAsync() *Client {
	o := c.GetDumpOptions()
	o.Async = true
	c.enableDump()
	return c
}

// DumpOnlyResponse indicates that should dump the responses' head and response.
func (c *Client) DumpOnlyResponse() *Client {
	o := c.GetDumpOptions()
	o.ResponseHead = true
	o.ResponseBody = true
	o.RequestBody = false
	o.RequestHead = false
	c.enableDump()
	return c
}

// DumpOnlyRequest indicates that should dump the requests' head and response.
func (c *Client) DumpOnlyRequest() *Client {
	o := c.GetDumpOptions()
	o.RequestHead = true
	o.RequestBody = true
	o.ResponseBody = false
	o.ResponseHead = false
	c.enableDump()
	return c
}

// DumpOnlyBody indicates that should dump the body of requests and responses.
func (c *Client) DumpOnlyBody() *Client {
	o := c.GetDumpOptions()
	o.RequestBody = true
	o.ResponseBody = true
	o.RequestHead = false
	o.ResponseHead = false
	c.enableDump()
	return c
}

// DumpOnlyHead indicates that should dump the head of requests and responses.
func (c *Client) DumpOnlyHead() *Client {
	o := c.GetDumpOptions()
	o.RequestHead = true
	o.ResponseHead = true
	o.RequestBody = false
	o.ResponseBody = false
	c.enableDump()
	return c
}

// DumpAll indicates that should dump both requests and responses' head and body.
func (c *Client) DumpAll() *Client {
	o := c.GetDumpOptions()
	o.RequestHead = true
	o.RequestBody = true
	o.ResponseHead = true
	o.ResponseBody = true
	c.enableDump()
	return c
}

// NewRequest is the alias of R()
func (c *Client) NewRequest() *Request {
	return c.R()
}

// AutoDecodeAllType indicates that try autodetect and decode all content type.
func (c *Client) AutoDecodeAllType() *Client {
	c.GetResponseOptions().AutoDecodeContentType = func(contentType string) bool {
		return true
	}
	return c
}

// AutoDecodeTextType indicates that only try autodetect and decode the text content type.
func (c *Client) AutoDecodeTextType() *Client {
	c.GetResponseOptions().AutoDecodeContentType = autoDecodeText
	return c
}

// UserAgent set the "User-Agent" header for all requests.
func (c *Client) UserAgent(userAgent string) *Client {
	return c.CommonHeader("User-Agent", userAgent)
}

// CommonHeader set the common header for all requests.
func (c *Client) CommonHeader(key, value string) *Client {
	if c.commonHeader == nil {
		c.commonHeader = make(map[string]string)
	}
	c.commonHeader[key] = value
	return c
}

// Dump if true, enables dump requests and responses,  allowing you
// to clearly see the content of all requests and responses，which
// is very convenient for debugging APIs.
// Dump if false, disable the dump behaviour.
func (c *Client) Dump(enable bool) *Client {
	if !enable {
		c.t.DisableDump()
		return c
	}
	c.enableDump()
	return c
}

// DumpOptions configures the underlying Transport's DumpOptions
func (c *Client) SetDumpOptions(opt *DumpOptions) *Client {
	if opt == nil {
		return c
	}
	c.dumpOptions = opt
	if c.t.dump != nil {
		c.t.dump.DumpOptions = opt
	}
	return c
}

// NewClient is the alias of C
func NewClient() *Client {
	return C()
}

// Clone copy and returns the Client
func (c *Client) Clone() *Client {
	t := c.t.Clone()
	t2, _ := http2ConfigureTransports(t)
	cc := *c.httpClient
	cc.Transport = t
	return &Client{
		httpClient:   &cc,
		t:            t,
		t2:           t2,
		dumpOptions:  c.dumpOptions.Clone(),
		jsonDecoder:  c.jsonDecoder,
		commonHeader: copyCommonHeader(c.commonHeader),
	}
}

// C create a new client.
func C() *Client {
	t := &Transport{
		ForceAttemptHTTP2:     true,
		Proxy:                 http.ProxyFromEnvironment,
		MaxIdleConns:          100,
		IdleConnTimeout:       90 * time.Second,
		TLSHandshakeTimeout:   10 * time.Second,
		ExpectContinueTimeout: 1 * time.Second,
	}
	t2, _ := http2ConfigureTransports(t)
	jar, _ := cookiejar.New(&cookiejar.Options{PublicSuffixList: publicsuffix.List})
	httpClient := &http.Client{
		Transport: t,
		Jar:       jar,
		Timeout:   2 * time.Minute,
	}
	c := &Client{
		log:        &emptyLogger{},
		httpClient: httpClient,
		t:          t,
		t2:         t2,
	}
	return c
}
<commit_msg>add some proxy function<commit_after>package req

import (
	"encoding/json"
	"golang.org/x/net/publicsuffix"
	"io"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"os"
	"time"
)

// DefaultClient returns the global default Client.
func DefaultClient() *Client {
	return defaultClient
}

// SetDefaultClient override the global default Client.
func SetDefaultClient(c *Client) {
	if c != nil {
		defaultClient = c
	}
}

var defaultClient *Client = C()

// Client is the req's http client.
type Client struct {
	log          Logger
	t            *Transport
	t2           *http2Transport
	dumpOptions  *DumpOptions
	httpClient   *http.Client
	jsonDecoder  *json.Decoder
	commonHeader map[string]string
}

func copyCommonHeader(h map[string]string) map[string]string {
	if h == nil {
		return nil
	}
	m := make(map[string]string)
	for k, v := range h {
		m[k] = v
	}
	return m
}

// R create a new request.
func (c *Client) R() *Request {
	req := &http.Request{
		Header:     make(http.Header),
		Proto:      "HTTP/1.1",
		ProtoMajor: 1,
		ProtoMinor: 1,
	}
	return &Request{
		client:      c,
		httpRequest: req,
	}
}

func (c *Client) AutoDiscardResponseBody() *Client {
	c.GetResponseOptions().AutoDiscard = true
	return c
}

// TestMode is like DebugMode, but discard response body, so you can
// dump responses without read response body
func (c *Client) TestMode() *Client {
	return c.DebugMode().AutoDiscardResponseBody()
}

const (
	userAgentFirefox = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10.15; rv:95.0) Gecko/20100101 Firefox/95.0"
	userAgentChrome  = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/97.0.4692.71 Safari/537.36"
)

// DebugMode enables dump for requests and responses, and set user
// agent to pretend to be a web browser, Avoid returning abnormal
// data from some sites.
func (c *Client) DebugMode() *Client {
	return c.AutoDecodeTextType().
		Dump(true).
		SetLogger(NewLogger(os.Stdout)).
		UserAgent(userAgentChrome)
}

// SetLogger set the logger for req.
func (c *Client) SetLogger(log Logger) *Client {
	if log == nil {
		return c
	}
	c.log = log
	return c
}

func (c *Client) GetResponseOptions() *ResponseOptions {
	if c.t.ResponseOptions == nil {
		c.t.ResponseOptions = &ResponseOptions{}
	}
	return c.t.ResponseOptions
}

// ResponseOptions set the ResponseOptions for the underlying Transport.
func (c *Client) SetResponseOptions(opt *ResponseOptions) *Client {
	if opt == nil {
		return c
	}
	c.t.ResponseOptions = opt
	return c
}

// Timeout set the timeout for all requests.
func (c *Client) Timeout(d time.Duration) *Client {
	c.httpClient.Timeout = d
	return c
}

func (c *Client) GetDumpOptions() *DumpOptions {
	if c.dumpOptions == nil {
		c.dumpOptions = newDefaultDumpOptions()
	}
	return c.dumpOptions
}

func (c *Client) enableDump() {
	if c.t.dump != nil { // dump already started
		return
	}
	c.t.EnableDump(c.GetDumpOptions())
}

// DumpToFile indicates that the content should dump to the specified filename.
func (c *Client) DumpToFile(filename string) *Client {
	file, err := os.Create(filename)
	if err != nil {
		logf(c.log, "create dump file error: %v", err)
		return c
	}
	c.GetDumpOptions().Output = file
	return c
}

// DumpTo indicates that the content should dump to the specified destination.
func (c *Client) DumpTo(output io.Writer) *Client {
	c.GetDumpOptions().Output = output
	c.enableDump()
	return c
}

// DumpAsync indicates that the dump should be done asynchronously,
// can be used for debugging in production environment without
// affecting performance.
func (c *Client) DumpAsync() *Client {
	o := c.GetDumpOptions()
	o.Async = true
	c.enableDump()
	return c
}

// DumpOnlyResponse indicates that should dump the responses' head and response.
func (c *Client) DumpOnlyResponse() *Client {
	o := c.GetDumpOptions()
	o.ResponseHead = true
	o.ResponseBody = true
	o.RequestBody = false
	o.RequestHead = false
	c.enableDump()
	return c
}

// DumpOnlyRequest indicates that should dump the requests' head and response.
func (c *Client) DumpOnlyRequest() *Client {
	o := c.GetDumpOptions()
	o.RequestHead = true
	o.RequestBody = true
	o.ResponseBody = false
	o.ResponseHead = false
	c.enableDump()
	return c
}

// DumpOnlyBody indicates that should dump the body of requests and responses.
func (c *Client) DumpOnlyBody() *Client {
	o := c.GetDumpOptions()
	o.RequestBody = true
	o.ResponseBody = true
	o.RequestHead = false
	o.ResponseHead = false
	c.enableDump()
	return c
}

// DumpOnlyHead indicates that should dump the head of requests and responses.
func (c *Client) DumpOnlyHead() *Client {
	o := c.GetDumpOptions()
	o.RequestHead = true
	o.ResponseHead = true
	o.RequestBody = false
	o.ResponseBody = false
	c.enableDump()
	return c
}

// DumpAll indicates that should dump both requests and responses' head and body.
func (c *Client) DumpAll() *Client {
	o := c.GetDumpOptions()
	o.RequestHead = true
	o.RequestBody = true
	o.ResponseHead = true
	o.ResponseBody = true
	c.enableDump()
	return c
}

// NewRequest is the alias of R()
func (c *Client) NewRequest() *Request {
	return c.R()
}

// AutoDecodeAllType indicates that try autodetect and decode all content type.
func (c *Client) AutoDecodeAllType() *Client {
	c.GetResponseOptions().AutoDecodeContentType = func(contentType string) bool {
		return true
	}
	return c
}

// AutoDecodeTextType indicates that only try autodetect and decode the text content type.
func (c *Client) AutoDecodeTextType() *Client {
	c.GetResponseOptions().AutoDecodeContentType = autoDecodeText
	return c
}

// UserAgent set the "User-Agent" header for all requests.
func (c *Client) UserAgent(userAgent string) *Client {
	return c.CommonHeader("User-Agent", userAgent)
}

// CommonHeader set the common header for all requests.
func (c *Client) CommonHeader(key, value string) *Client {
	if c.commonHeader == nil {
		c.commonHeader = make(map[string]string)
	}
	c.commonHeader[key] = value
	return c
}

// Dump if true, enables dump requests and responses,  allowing you
// to clearly see the content of all requests and responses，which
// is very convenient for debugging APIs.
// Dump if false, disable the dump behaviour.
func (c *Client) Dump(enable bool) *Client {
	if !enable {
		c.t.DisableDump()
		return c
	}
	c.enableDump()
	return c
}

// DumpOptions configures the underlying Transport's DumpOptions
func (c *Client) SetDumpOptions(opt *DumpOptions) *Client {
	if opt == nil {
		return c
	}
	c.dumpOptions = opt
	if c.t.dump != nil {
		c.t.dump.DumpOptions = opt
	}
	return c
}

// Proxy set the proxy function.
func (c *Client) Proxy(proxy func(*http.Request) (*url.URL, error)) *Client {
	c.t.Proxy = proxy
	return c
}

func (c *Client) ProxyFromEnv() *Client {
	c.t.Proxy = http.ProxyFromEnvironment
	return c
}

func (c *Client) ProxyURL(proxyUrl string) *Client {
	u, err := url.Parse(proxyUrl)
	if err != nil {
		logf(c.log, "failed to parse proxy url %s: %v", proxyUrl, err)
		return c
	}
	c.t.Proxy = http.ProxyURL(u)
	return c
}

// NewClient is the alias of C
func NewClient() *Client {
	return C()
}

// Clone copy and returns the Client
func (c *Client) Clone() *Client {
	t := c.t.Clone()
	t2, _ := http2ConfigureTransports(t)
	cc := *c.httpClient
	cc.Transport = t
	return &Client{
		httpClient:   &cc,
		t:            t,
		t2:           t2,
		dumpOptions:  c.dumpOptions.Clone(),
		jsonDecoder:  c.jsonDecoder,
		commonHeader: copyCommonHeader(c.commonHeader),
	}
}

// C create a new client.
func C() *Client {
	t := &Transport{
		ForceAttemptHTTP2:     true,
		Proxy:                 http.ProxyFromEnvironment,
		MaxIdleConns:          100,
		IdleConnTimeout:       90 * time.Second,
		TLSHandshakeTimeout:   10 * time.Second,
		ExpectContinueTimeout: 1 * time.Second,
	}
	t2, _ := http2ConfigureTransports(t)
	jar, _ := cookiejar.New(&cookiejar.Options{PublicSuffixList: publicsuffix.List})
	httpClient := &http.Client{
		Transport: t,
		Jar:       jar,
		Timeout:   2 * time.Minute,
	}
	c := &Client{
		log:        &emptyLogger{},
		httpClient: httpClient,
		t:          t,
		t2:         t2,
	}
	return c
}
<|endoftext|>


<commit_before>package gokiq

import (
	"crypto/rand"
	"encoding/json"
	"fmt"
	"io"
	"reflect"

	"github.com/garyburd/redigo/redis"
)

var Client = NewClientConfig()

type jobMap map[reflect.Type]JobConfig

type ClientConfig struct {
	RedisServer    string
	RedisNamespace string
	RedisMaxIdle   int
	Fake           bool

	redisPool   *redis.Pool
	jobMapping  jobMap
	knownQueues map[string]struct{}
}

func NewClientConfig() *ClientConfig {
	return &ClientConfig{
		RedisServer:  defaultRedisServer,
		RedisMaxIdle: 1,
		jobMapping:   make(jobMap),
		knownQueues:  make(map[string]struct{}),
	}
}

func (c *ClientConfig) Register(worker Worker, queue string, retries int) {
	t := workerType(worker)
	c.jobMapping[t] = JobConfig{queue, retries, t.Name()}
	c.trackQueue(queue)
}

func (c *ClientConfig) RegisterName(name string, worker Worker, queue string, retries int) {
	c.jobMapping[workerType(worker)] = JobConfig{queue, retries, name}
	c.trackQueue(queue)
}

func (c *ClientConfig) Connect() {
	// TODO: add a mutex for the redis pool
	if c.redisPool != nil {
		c.redisPool.Close()
	}
	c.redisPool = redis.NewPool(func() (redis.Conn, error) {
		return redis.Dial("tcp", c.RedisServer)
	}, c.RedisMaxIdle)

	queues := make([]interface{}, 1, len(c.knownQueues)+1)
	queues[0] = c.nsKey("queues")
	for queue := range c.knownQueues {
		queues = append(queues, queue)
	}
	c.redisQuery("SADD", queues...)
}

func (c *ClientConfig) QueueJob(worker Worker) error {
	config, ok := c.jobMapping[workerType(worker)]
	if !ok {
		panic(fmt.Errorf("gokiq: Unregistered worker type %T", worker))
	}
	return c.queueJob(worker, config)
}

func (c *ClientConfig) QueueJobWithConfig(worker Worker, config JobConfig) error {
	c.trackQueue(config.Queue)
	return c.queueJob(worker, config)
}

func (c *ClientConfig) queueJob(worker Worker, config JobConfig) error {
	data, err := json.Marshal(worker)
	if err != nil {
		return err
	}
	args := json.RawMessage(data)
	job := &Job{
		Type:  config.name,
		Args:  &args,
		Retry: config.MaxRetries,
		ID:    generateJobID(),
	}
	if c.Fake {
		return worker.Perform()
	}

	_, err = c.redisQuery("RPUSH", c.nsKey("queue:"+config.Queue), job.JSON())
	return err
}

func (c *ClientConfig) trackQueue(queue string) {
	_, known := c.knownQueues[queue]
	if !known {
		c.knownQueues[queue] = struct{}{}
		if c.redisPool != nil {
			c.redisQuery("SADD", c.nsKey("queues"), queue)
		}
	}
}

func (c *ClientConfig) redisQuery(command string, args ...interface{}) (interface{}, error) {
	conn := c.redisPool.Get()
	defer conn.Close()
	return conn.Do(command, args...)
}

func (c *ClientConfig) nsKey(key string) string {
	if c.RedisNamespace != "" {
		return c.RedisNamespace + ":" + key
	}
	return key
}

func generateJobID() string {
	b := make([]byte, 8)
	io.ReadFull(rand.Reader, b)
	return fmt.Sprintf("%x", b)
}

type JobConfig struct {
	Queue      string
	MaxRetries int

	name string
}
<commit_msg>Rename QueueJobWithConfig -> QueueJobConfig<commit_after>package gokiq

import (
	"crypto/rand"
	"encoding/json"
	"fmt"
	"io"
	"reflect"

	"github.com/garyburd/redigo/redis"
)

var Client = NewClientConfig()

type jobMap map[reflect.Type]JobConfig

type ClientConfig struct {
	RedisServer    string
	RedisNamespace string
	RedisMaxIdle   int
	Fake           bool

	redisPool   *redis.Pool
	jobMapping  jobMap
	knownQueues map[string]struct{}
}

func NewClientConfig() *ClientConfig {
	return &ClientConfig{
		RedisServer:  defaultRedisServer,
		RedisMaxIdle: 1,
		jobMapping:   make(jobMap),
		knownQueues:  make(map[string]struct{}),
	}
}

func (c *ClientConfig) Register(worker Worker, queue string, retries int) {
	t := workerType(worker)
	c.jobMapping[t] = JobConfig{queue, retries, t.Name()}
	c.trackQueue(queue)
}

func (c *ClientConfig) RegisterName(name string, worker Worker, queue string, retries int) {
	c.jobMapping[workerType(worker)] = JobConfig{queue, retries, name}
	c.trackQueue(queue)
}

func (c *ClientConfig) Connect() {
	// TODO: add a mutex for the redis pool
	if c.redisPool != nil {
		c.redisPool.Close()
	}
	c.redisPool = redis.NewPool(func() (redis.Conn, error) {
		return redis.Dial("tcp", c.RedisServer)
	}, c.RedisMaxIdle)

	queues := make([]interface{}, 1, len(c.knownQueues)+1)
	queues[0] = c.nsKey("queues")
	for queue := range c.knownQueues {
		queues = append(queues, queue)
	}
	c.redisQuery("SADD", queues...)
}

func (c *ClientConfig) QueueJob(worker Worker) error {
	config, ok := c.jobMapping[workerType(worker)]
	if !ok {
		panic(fmt.Errorf("gokiq: Unregistered worker type %T", worker))
	}
	return c.queueJob(worker, config)
}

func (c *ClientConfig) QueueJobConfig(worker Worker, config JobConfig) error {
	c.trackQueue(config.Queue)
	return c.queueJob(worker, config)
}

func (c *ClientConfig) queueJob(worker Worker, config JobConfig) error {
	data, err := json.Marshal(worker)
	if err != nil {
		return err
	}
	args := json.RawMessage(data)
	job := &Job{
		Type:  config.name,
		Args:  &args,
		Retry: config.MaxRetries,
		ID:    generateJobID(),
	}
	if c.Fake {
		return worker.Perform()
	}

	_, err = c.redisQuery("RPUSH", c.nsKey("queue:"+config.Queue), job.JSON())
	return err
}

func (c *ClientConfig) trackQueue(queue string) {
	_, known := c.knownQueues[queue]
	if !known {
		c.knownQueues[queue] = struct{}{}
		if c.redisPool != nil {
			c.redisQuery("SADD", c.nsKey("queues"), queue)
		}
	}
}

func (c *ClientConfig) redisQuery(command string, args ...interface{}) (interface{}, error) {
	conn := c.redisPool.Get()
	defer conn.Close()
	return conn.Do(command, args...)
}

func (c *ClientConfig) nsKey(key string) string {
	if c.RedisNamespace != "" {
		return c.RedisNamespace + ":" + key
	}
	return key
}

func generateJobID() string {
	b := make([]byte, 8)
	io.ReadFull(rand.Reader, b)
	return fmt.Sprintf("%x", b)
}

type JobConfig struct {
	Queue      string
	MaxRetries int

	name string
}
<|endoftext|>


<commit_before>package i18n

import (
	"errors"
	"net/http"

	"golang.org/x/text/language"

	"github.com/volatile/core"
)

var (
	errUnknownLocale = errors.New("i18n: unknown locale")
)

// ClientLocale returns the current locale used by the client.
// If the locale has not been matched already, it will be done before returning.
func ClientLocale(c *core.Context) string {
	// Use context data to match locale a single time per request.
	if v, ok := c.Data[contextDataKey]; ok {
		return v.(string)
	}

	// Use cookie if exists and valid.
	if useCookie {
		if cookie, err := c.Request.Cookie(cookieName); err == nil && localeExists(cookie.Value) {
			return cookie.Value
		}
	}

	// Match, save and return locale key.
	l := matchLocale(c.Request)
	SetClientLocale(c, l)
	return l
}

// SetClientLocale changes the locale for the actual client, but only if the locale exists.
func SetClientLocale(c *core.Context, l string) error {
	if !localeExists(l) {
		return errUnknownLocale
	}

	if useCookie {
		http.SetCookie(c.ResponseWriter, &http.Cookie{
			Name:   cookieName,
			Value:  l,
			Path:   "/",
			MaxAge: 315569260, // 10 years cookie
		})
	}

	c.Data[contextDataKey] = l
	return nil
}

// matchLocale returns the most appropriate and available locale key for the client.
// Content Language Headers: https://tools.ietf.org/html/rfc3282
func matchLocale(r *http.Request) string {
	tag, _, _ := language.ParseAcceptLanguage(r.Header.Get("Accept-Language"))

	for _, t := range tag {
		b, _ := t.Base()
		if _, ok := (*locales)[b.String()]; ok {
			return b.String()
		}
	}

	return defaultLocale
}
<commit_msg>Export ErrUnknownLocale<commit_after>package i18n

import (
	"errors"
	"net/http"

	"golang.org/x/text/language"

	"github.com/volatile/core"
)

// ErrUnknownLocale is returned when the wanted locale doesn't exists.
var ErrUnknownLocale = errors.New("i18n: unknown locale")

// ClientLocale returns the current locale used by the client.
// If the locale has not been matched already, it will be done before returning.
func ClientLocale(c *core.Context) string {
	// Use context data to match locale a single time per request.
	if v, ok := c.Data[contextDataKey]; ok {
		return v.(string)
	}

	// Use cookie if exists and valid.
	if useCookie {
		if cookie, err := c.Request.Cookie(cookieName); err == nil && localeExists(cookie.Value) {
			return cookie.Value
		}
	}

	// Match, save and return locale key.
	l := matchLocale(c.Request)
	SetClientLocale(c, l)
	return l
}

// SetClientLocale changes the locale for the actual client.
// If the locale l doesn't exists, error ErrUnknownLocale is returned.
func SetClientLocale(c *core.Context, l string) error {
	if !localeExists(l) {
		return ErrUnknownLocale
	}

	if useCookie {
		http.SetCookie(c.ResponseWriter, &http.Cookie{
			Name:   cookieName,
			Value:  l,
			Path:   "/",
			MaxAge: 315569260, // 10 years cookie
		})
	}

	c.Data[contextDataKey] = l
	return nil
}

// matchLocale returns the most appropriate and available locale key for the client.
// Content Language Headers: https://tools.ietf.org/html/rfc3282
func matchLocale(r *http.Request) string {
	tag, _, _ := language.ParseAcceptLanguage(r.Header.Get("Accept-Language"))

	for _, t := range tag {
		b, _ := t.Base()
		if _, ok := (*locales)[b.String()]; ok {
			return b.String()
		}
	}

	return defaultLocale
}
<|endoftext|>


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

// This file implements the web server side for FastCGI
// as specified in http://www.mit.edu/~yandros/doc/specs/fcgi-spec.html

// A part of this file is from golang package net/http/cgi,
// in particular https://golang.org/src/net/http/cgi/host.go

package gofast

import (
	"bufio"
	"fmt"
	"io"
	"net"
	"net/http"
	"strconv"
	"strings"
	"sync"
)

// Request hold information of a standard
// FastCGI request
type Request struct {
	Raw      *http.Request
	ID       uint16
	Params   map[string]string
	Stdin    io.ReadCloser
	KeepConn bool
}

// client is the default implementation of Client
type client struct {
	conn   *conn
	chanID chan uint16
}

// AllocID implements Client.AllocID
func (c *client) AllocID() (reqID uint16) {
	reqID = <-c.chanID
	return
}

// ReleaseID implements Client.ReleaseID
func (c *client) ReleaseID(reqID uint16) {
	go func() {
		// release the ID back to channel for reuse
		// use goroutine to prevent blocking ReleaseID
		c.chanID <- reqID
	}()
}

// writeRequest writes params and stdin to the FastCGI application
func (c *client) writeRequest(resp *ResponsePipe, req *Request) (err error) {

	// FIXME: add other role implementation, add role field to Request
	err = c.conn.writeBeginRequest(req.ID, uint16(roleResponder), 0)
	if err != nil {
		resp.Close()
		return
	}
	err = c.conn.writePairs(typeParams, req.ID, req.Params)
	if err != nil {
		resp.Close()
		return
	}
	if req.Stdin == nil {
		err = c.conn.writeRecord(typeStdin, req.ID, []byte{})
	} else {
		defer req.Stdin.Close()
		p := make([]byte, 1024)
		var count int
		for {
			count, err = req.Stdin.Read(p)
			if err == io.EOF {
				err = nil
			} else if err != nil {
				break
			}
			if count == 0 {
				break
			}

			err = c.conn.writeRecord(typeStdin, req.ID, p[:count])
			if err != nil {
				break
			}
		}
	}

	if err != nil {
		resp.Close()
	}
	return
}

// readResponse read the FastCGI stdout and stderr, then write
// to the response pipe
func (c *client) readResponse(resp *ResponsePipe, req *Request) {
	var rec record

	defer c.ReleaseID(req.ID)
	defer resp.Close()
readLoop:
	for {
		if err := rec.read(c.conn.rwc); err != nil {
			break
		}

		// different output type for different stream
		switch rec.h.Type {
		case typeStdout:
			resp.stdOutWriter.Write(rec.content())
		case typeStderr:
			resp.stdErrWriter.Write(rec.content())
		case typeEndRequest:
			break readLoop
		default:
			panic(fmt.Sprintf("unexpected type %#v in readLoop", rec.h.Type))
		}
	}
}

// Do implements Client.Do
func (c *client) Do(req *Request) (resp *ResponsePipe, err error) {

	resp = NewResponsePipe()

	// FIXME: Should run read and write in parallel.
	//        Specification never said "write before read".
	//        Current workflow may block.

	if err = c.writeRequest(resp, req); err != nil {
		return
	}

	// NOTE: all errors return before readResponse
	go c.readResponse(resp, req)
	return
}

// NewRequest implements Client.NewRequest
func (c *client) NewRequest(r *http.Request) (req *Request) {
	req = &Request{
		Raw:    r,
		ID:     c.AllocID(),
		Params: make(map[string]string),
	}

	// if no http request, return here
	if r == nil {
		return
	}

	// pass body (io.ReadCloser) to stdio
	req.Stdin = r.Body

	return
}

// Client is a client interface of FastCGI
// application process through given
// connection (net.Conn)
type Client interface {

	// Do takes care of a proper FastCGI request
	Do(req *Request) (resp *ResponsePipe, err error)

	// NewRequest returns a standard FastCGI request
	// with a unique request ID allocted by the client
	NewRequest(*http.Request) *Request

	// AllocID allocates a new reqID.
	// It blocks if all possible uint16 IDs are allocated.
	AllocID() uint16

	// ReleaseID releases a reqID.
	// It never blocks.
	ReleaseID(uint16)
}

// NewClient returns a Client of the given
// connection (net.Conn).
//
// limit is the maximum number of request that the
// applcation support. 0 means the maximum number
// available for 16bit request id (65536).
// Default 0.
//
func NewClient(conn net.Conn, limit uint32) Client {
	cid := make(chan uint16)

	if limit == 0 || limit > 65536 {
		limit = 65536
	}
	go func(maxID uint16) {
		for i := uint16(0); i < maxID; i++ {
			cid <- i
		}
		cid <- uint16(maxID)
	}(uint16(limit - 1))

	return &client{
		conn:   newConn(conn),
		chanID: cid,
	}
}

// NewResponsePipe returns an initialized new ResponsePipe struct
func NewResponsePipe() (p *ResponsePipe) {
	p = new(ResponsePipe)
	p.stdOutReader, p.stdOutWriter = io.Pipe()
	p.stdErrReader, p.stdErrWriter = io.Pipe()
	return
}

// ResponsePipe contains readers and writers that handles
// all FastCGI output streams
type ResponsePipe struct {
	stdOutReader io.Reader
	stdOutWriter io.WriteCloser
	stdErrReader io.Reader
	stdErrWriter io.WriteCloser
}

// Close close all writers
func (pipes *ResponsePipe) Close() {
	pipes.stdOutWriter.Close()
	pipes.stdErrWriter.Close()
}

// WriteTo writes the given output into http.ResponseWriter
func (pipes *ResponsePipe) WriteTo(rw http.ResponseWriter, ew io.Writer) (err error) {
	wg := new(sync.WaitGroup)
	wg.Add(2)

	go func() {
		defer wg.Done()
		err = pipes.writeResponse(rw)
	}()

	go func() {
		defer wg.Done()
		err = pipes.writeError(ew)
	}()

	// blocks until all reads and writes are done
	wg.Wait()
	return
}

func (pipes *ResponsePipe) writeError(w io.Writer) (err error) {
	_, err = io.Copy(w, pipes.stdErrReader)
	if err != nil {
		err = fmt.Errorf("gofast: copy error: %v", err.Error())
	}
	return
}

// writeTo writes the given output into http.ResponseWriter
func (pipes *ResponsePipe) writeResponse(w http.ResponseWriter) (err error) {
	linebody := bufio.NewReaderSize(pipes.stdOutReader, 1024)
	headers := make(http.Header)
	statusCode := 0
	headerLines := 0
	sawBlankLine := false

	for {
		var line []byte
		var isPrefix bool
		line, isPrefix, err = linebody.ReadLine()
		if isPrefix {
			w.WriteHeader(http.StatusInternalServerError)
			err = fmt.Errorf("gofast: long header line from subprocess")
			return
		}
		if err == io.EOF {
			break
		}
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			err = fmt.Errorf("gofast: error reading headers: %v", err)
			return
		}
		if len(line) == 0 {
			sawBlankLine = true
			break
		}
		headerLines++
		parts := strings.SplitN(string(line), ":", 2)
		if len(parts) < 2 {
			err = fmt.Errorf("gofast: bogus header line: %s", string(line))
			return
		}
		header, val := parts[0], parts[1]
		header = strings.TrimSpace(header)
		val = strings.TrimSpace(val)
		switch {
		case header == "Status":
			if len(val) < 3 {
				err = fmt.Errorf("gofast: bogus status (short): %q", val)
				return
			}
			var code int
			code, err = strconv.Atoi(val[0:3])
			if err != nil {
				err = fmt.Errorf("gofast: bogus status: %q\nline was %q",
					val, line)
				return
			}
			statusCode = code
		default:
			headers.Add(header, val)
		}
	}
	if headerLines == 0 || !sawBlankLine {
		w.WriteHeader(http.StatusInternalServerError)
		err = fmt.Errorf("gofast: no headers")
		return
	}

	if loc := headers.Get("Location"); loc != "" {
		/*
			if strings.HasPrefix(loc, "/") && h.PathLocationHandler != nil {
				h.handleInternalRedirect(rw, req, loc)
				return
			}
		*/
		if statusCode == 0 {
			statusCode = http.StatusFound
		}
	}

	if statusCode == 0 && headers.Get("Content-Type") == "" {
		w.WriteHeader(http.StatusInternalServerError)
		err = fmt.Errorf("gofast: missing required Content-Type in headers")
		return
	}

	if statusCode == 0 {
		statusCode = http.StatusOK
	}

	// Copy headers to rw's headers, after we've decided not to
	// go into handleInternalRedirect, which won't want its rw
	// headers to have been touched.
	for k, vv := range headers {
		for _, v := range vv {
			w.Header().Add(k, v)
		}
	}

	w.WriteHeader(statusCode)

	_, err = io.Copy(w, linebody)
	if err != nil {
		err = fmt.Errorf("gofast: copy error: %v", err)
	}
	return
}
<commit_msg>fix racing issue in client<commit_after>// Copyright 2016 Yeung Shu Hung and The Go Authors.
// All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This file implements the web server side for FastCGI
// as specified in http://www.mit.edu/~yandros/doc/specs/fcgi-spec.html

// A part of this file is from golang package net/http/cgi,
// in particular https://golang.org/src/net/http/cgi/host.go

package gofast

import (
	"bufio"
	"fmt"
	"io"
	"net"
	"net/http"
	"strconv"
	"strings"
)

// Request hold information of a standard
// FastCGI request
type Request struct {
	Raw      *http.Request
	ID       uint16
	Params   map[string]string
	Stdin    io.ReadCloser
	KeepConn bool
}

// client is the default implementation of Client
type client struct {
	conn   *conn
	chanID chan uint16
}

// AllocID implements Client.AllocID
func (c *client) AllocID() (reqID uint16) {
	reqID = <-c.chanID
	return
}

// ReleaseID implements Client.ReleaseID
func (c *client) ReleaseID(reqID uint16) {
	go func() {
		// release the ID back to channel for reuse
		// use goroutine to prevent blocking ReleaseID
		c.chanID <- reqID
	}()
}

// writeRequest writes params and stdin to the FastCGI application
func (c *client) writeRequest(resp *ResponsePipe, req *Request) (err error) {

	// FIXME: add other role implementation, add role field to Request
	err = c.conn.writeBeginRequest(req.ID, uint16(roleResponder), 0)
	if err != nil {
		resp.Close()
		return
	}
	err = c.conn.writePairs(typeParams, req.ID, req.Params)
	if err != nil {
		resp.Close()
		return
	}
	if req.Stdin == nil {
		err = c.conn.writeRecord(typeStdin, req.ID, []byte{})
	} else {
		defer req.Stdin.Close()
		p := make([]byte, 1024)
		var count int
		for {
			count, err = req.Stdin.Read(p)
			if err == io.EOF {
				err = nil
			} else if err != nil {
				break
			}
			if count == 0 {
				break
			}

			err = c.conn.writeRecord(typeStdin, req.ID, p[:count])
			if err != nil {
				break
			}
		}
	}

	if err != nil {
		resp.Close()
	}
	return
}

// readResponse read the FastCGI stdout and stderr, then write
// to the response pipe
func (c *client) readResponse(resp *ResponsePipe, req *Request) {
	var rec record

	defer c.ReleaseID(req.ID)
	defer resp.Close()
readLoop:
	for {
		if err := rec.read(c.conn.rwc); err != nil {
			break
		}

		// different output type for different stream
		switch rec.h.Type {
		case typeStdout:
			resp.stdOutWriter.Write(rec.content())
		case typeStderr:
			resp.stdErrWriter.Write(rec.content())
		case typeEndRequest:
			break readLoop
		default:
			panic(fmt.Sprintf("unexpected type %#v in readLoop", rec.h.Type))
		}
	}
}

// Do implements Client.Do
func (c *client) Do(req *Request) (resp *ResponsePipe, err error) {

	resp = NewResponsePipe()

	// FIXME: Should run read and write in parallel.
	//        Specification never said "write before read".
	//        Current workflow may block.

	if err = c.writeRequest(resp, req); err != nil {
		return
	}

	// NOTE: all errors return before readResponse
	go c.readResponse(resp, req)
	return
}

// NewRequest implements Client.NewRequest
func (c *client) NewRequest(r *http.Request) (req *Request) {
	req = &Request{
		Raw:    r,
		ID:     c.AllocID(),
		Params: make(map[string]string),
	}

	// if no http request, return here
	if r == nil {
		return
	}

	// pass body (io.ReadCloser) to stdio
	req.Stdin = r.Body

	return
}

// Client is a client interface of FastCGI
// application process through given
// connection (net.Conn)
type Client interface {

	// Do takes care of a proper FastCGI request
	Do(req *Request) (resp *ResponsePipe, err error)

	// NewRequest returns a standard FastCGI request
	// with a unique request ID allocted by the client
	NewRequest(*http.Request) *Request

	// AllocID allocates a new reqID.
	// It blocks if all possible uint16 IDs are allocated.
	AllocID() uint16

	// ReleaseID releases a reqID.
	// It never blocks.
	ReleaseID(uint16)
}

// NewClient returns a Client of the given
// connection (net.Conn).
//
// limit is the maximum number of request that the
// applcation support. 0 means the maximum number
// available for 16bit request id (65536).
// Default 0.
//
func NewClient(conn net.Conn, limit uint32) Client {
	cid := make(chan uint16)

	if limit == 0 || limit > 65536 {
		limit = 65536
	}
	go func(maxID uint16) {
		for i := uint16(0); i < maxID; i++ {
			cid <- i
		}
		cid <- uint16(maxID)
	}(uint16(limit - 1))

	return &client{
		conn:   newConn(conn),
		chanID: cid,
	}
}

// NewResponsePipe returns an initialized new ResponsePipe struct
func NewResponsePipe() (p *ResponsePipe) {
	p = new(ResponsePipe)
	p.stdOutReader, p.stdOutWriter = io.Pipe()
	p.stdErrReader, p.stdErrWriter = io.Pipe()
	return
}

// ResponsePipe contains readers and writers that handles
// all FastCGI output streams
type ResponsePipe struct {
	stdOutReader io.Reader
	stdOutWriter io.WriteCloser
	stdErrReader io.Reader
	stdErrWriter io.WriteCloser
}

// Close close all writers
func (pipes *ResponsePipe) Close() {
	pipes.stdOutWriter.Close()
	pipes.stdErrWriter.Close()
}

// WriteTo writes the given output into http.ResponseWriter
func (pipes *ResponsePipe) WriteTo(rw http.ResponseWriter, ew io.Writer) (err error) {
	chErr := make(chan error, 2)

	go func() {
		chErr <- pipes.writeResponse(rw)
	}()
	go func() {
		chErr <- pipes.writeError(ew)
	}()

	for i := 0; i < 2; i++ {
		if err = <-chErr; err != nil {
			close(chErr)
			return
		}
	}
	return
}

func (pipes *ResponsePipe) writeError(w io.Writer) (err error) {
	_, err = io.Copy(w, pipes.stdErrReader)
	if err != nil {
		err = fmt.Errorf("gofast: copy error: %v", err.Error())
	}
	return
}

// writeTo writes the given output into http.ResponseWriter
func (pipes *ResponsePipe) writeResponse(w http.ResponseWriter) (err error) {
	linebody := bufio.NewReaderSize(pipes.stdOutReader, 1024)
	headers := make(http.Header)
	statusCode := 0
	headerLines := 0
	sawBlankLine := false

	for {
		var line []byte
		var isPrefix bool
		line, isPrefix, err = linebody.ReadLine()
		if isPrefix {
			w.WriteHeader(http.StatusInternalServerError)
			err = fmt.Errorf("gofast: long header line from subprocess")
			return
		}
		if err == io.EOF {
			break
		}
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			err = fmt.Errorf("gofast: error reading headers: %v", err)
			return
		}
		if len(line) == 0 {
			sawBlankLine = true
			break
		}
		headerLines++
		parts := strings.SplitN(string(line), ":", 2)
		if len(parts) < 2 {
			err = fmt.Errorf("gofast: bogus header line: %s", string(line))
			return
		}
		header, val := parts[0], parts[1]
		header = strings.TrimSpace(header)
		val = strings.TrimSpace(val)
		switch {
		case header == "Status":
			if len(val) < 3 {
				err = fmt.Errorf("gofast: bogus status (short): %q", val)
				return
			}
			var code int
			code, err = strconv.Atoi(val[0:3])
			if err != nil {
				err = fmt.Errorf("gofast: bogus status: %q\nline was %q",
					val, line)
				return
			}
			statusCode = code
		default:
			headers.Add(header, val)
		}
	}
	if headerLines == 0 || !sawBlankLine {
		w.WriteHeader(http.StatusInternalServerError)
		err = fmt.Errorf("gofast: no headers")
		return
	}

	if loc := headers.Get("Location"); loc != "" {
		/*
			if strings.HasPrefix(loc, "/") && h.PathLocationHandler != nil {
				h.handleInternalRedirect(rw, req, loc)
				return
			}
		*/
		if statusCode == 0 {
			statusCode = http.StatusFound
		}
	}

	if statusCode == 0 && headers.Get("Content-Type") == "" {
		w.WriteHeader(http.StatusInternalServerError)
		err = fmt.Errorf("gofast: missing required Content-Type in headers")
		return
	}

	if statusCode == 0 {
		statusCode = http.StatusOK
	}

	// Copy headers to rw's headers, after we've decided not to
	// go into handleInternalRedirect, which won't want its rw
	// headers to have been touched.
	for k, vv := range headers {
		for _, v := range vv {
			w.Header().Add(k, v)
		}
	}

	w.WriteHeader(statusCode)

	_, err = io.Copy(w, linebody)
	if err != nil {
		err = fmt.Errorf("gofast: copy error: %v", err)
	}
	return
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"log"
	"os"
	"path"

	"github.com/bitly/go-nsq"
)

const (
	STATUS_OK  string = "OK"
	STATUS_ERR string = "ERR"
)

type Worker struct {
	Consumer       *nsq.Consumer
	ReloadConsumer *nsq.Consumer
	Throughput     int
	QueueAddr      string
	ScriptDir      string
	StoreDir       string
	WorkingDir     string
	WhiteList      map[string]bool
}

func NewWorker(c Config) (Worker, error) {
	fmt.Printf("Creating consumer with topic: %s and channel: %s.\n", c.Topic, c.Worker.Channel)

	var err error
	var worker Worker
	worker.Throughput = c.Worker.Throughput
	worker.QueueAddr = c.QueueAddr
	worker.ScriptDir = c.Worker.ScriptDir
	worker.StoreDir = c.Worker.StoreDir
	worker.WorkingDir = c.Worker.WorkingDir

	conf := nsq.NewConfig()
	conf.Set("max_in_flight", worker.Throughput)
	consumer, err := nsq.NewConsumer(c.Topic, c.Worker.Channel, conf)
	if err != nil {
		log.Println(err)
		return worker, err
	}

	rConsumer, err := nsq.NewConsumer("reload", c.Worker.Channel, nsq.NewConfig())
	if err != nil {
		log.Println(err)
		return worker, err
	}

	worker.Consumer = consumer
	worker.ReloadConsumer = rConsumer

	// Generate whitelist of allowed scripts.
	path := path.Join(config.Worker.ScriptDir, config.Worker.WhiteList)
	err = worker.LoadWhiteList(path)
	if err != nil {
		log.Println(err)
		return worker, err
	}

	fmt.Printf("Worker connecting to %s and running scripts in %s.\n", c.QueueAddr, c.Worker.WorkingDir)
	return worker, nil
}

func (w *Worker) LoadWhiteList(path string) error {
	fmt.Printf("Loading whitelist from %s\n", path)

	file, err := os.Open(path)
	if err != nil {
		return err
	}
	defer file.Close()

	whiteList := make(map[string]bool)

	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		whiteList[scanner.Text()] = true
	}
	err = scanner.Err()
	if err != nil {
		log.Println(err)
		return err
	}

	fmt.Println("Whitelist:")
	for script := range whiteList {
		fmt.Println("	", script)
	}

	w.WhiteList = whiteList
	return nil
}

func (w *Worker) JobRequestHandler(m *nsq.Message) error {
	// Initialize Job from request
	var job Job
	job.Status = STATUS_ERR
	job.ScriptDir = w.ScriptDir
	job.WorkingDir = w.WorkingDir
	job.StoreDir = w.StoreDir

	err := json.Unmarshal(m.Body, &job)
	if err != nil {
		log.Println("Invalid JSON request", err)
		return nil
	}

	// Try and run script
	if w.WhiteList[job.Script] {
		log.Println("Dequeued request as Job", job.ID)

		resultChan := make(chan error, 1)
		go job.Execute(resultChan)
		err := <-resultChan
		if err != nil {
			job.ExecLog = err.Error()
		} else {
			job.Status = STATUS_OK
		}
	} else {
		msg := fmt.Sprintf("%s is not on script whitelist", job.Script)
		job.ExecLog = msg
	}

	log.Println(job.ExecLog)
	job.Log()
	job.Callback()
	return nil
}

func (w *Worker) ReloadRequestHandler(m *nsq.Message) error {
	whitelist := path.Clean(string(m.Body))
	err := w.LoadWhiteList(whitelist)
	if err != nil {
		log.Println("Failed to reload whitelist from", whitelist)
		return err
	}
	log.Println("Reloaded whitelist from", whitelist)
	return nil
}

func (w *Worker) Run() {
	// Set the message handler.
	w.Consumer.SetConcurrentHandlers(nsq.HandlerFunc(w.JobRequestHandler), w.Throughput)
	w.ReloadConsumer.SetHandler(nsq.HandlerFunc(w.ReloadRequestHandler))

	var err error

	// Connect the queue.
	fmt.Println("Connecting to", w.QueueAddr)
	err = w.Consumer.ConnectToNSQD(w.QueueAddr)
	if err != nil {
		log.Println(err)
	}

	err = w.ReloadConsumer.ConnectToNSQD(w.QueueAddr)
	if err != nil {
		log.Println(err)
	}
}

func (w *Worker) Stop() {
	w.Consumer.Stop()
	w.ReloadConsumer.Stop()
	return
}
<commit_msg>Missing WhiteList in config now becomes *.<commit_after>package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"log"
	"os"
	"path"

	"github.com/bitly/go-nsq"
)

const (
	STATUS_OK  string = "OK"
	STATUS_ERR string = "ERR"
)

type Worker struct {
	Consumer       *nsq.Consumer
	ReloadConsumer *nsq.Consumer
	Throughput     int
	QueueAddr      string
	ScriptDir      string
	StoreDir       string
	WorkingDir     string
	WhiteList      map[string]bool
}

func NewWorker(c Config) (Worker, error) {
	fmt.Printf("Creating consumer with topic: %s and channel: %s.\n", c.Topic, c.Worker.Channel)

	var err error
	var worker Worker
	worker.Throughput = c.Worker.Throughput
	worker.QueueAddr = c.QueueAddr
	worker.ScriptDir = c.Worker.ScriptDir
	worker.StoreDir = c.Worker.StoreDir
	worker.WorkingDir = c.Worker.WorkingDir

	conf := nsq.NewConfig()
	conf.Set("max_in_flight", worker.Throughput)
	consumer, err := nsq.NewConsumer(c.Topic, c.Worker.Channel, conf)
	if err != nil {
		log.Println(err)
		return worker, err
	}

	rConsumer, err := nsq.NewConsumer("reload", c.Worker.Channel, nsq.NewConfig())
	if err != nil {
		log.Println(err)
		return worker, err
	}

	worker.Consumer = consumer
	worker.ReloadConsumer = rConsumer

	// Generate whitelist of allowed scripts.
	path := path.Join(config.Worker.ScriptDir, config.Worker.WhiteList)
	err = worker.LoadWhiteList(path)
	if err != nil {
		log.Println(err)
		return worker, err
	}

	fmt.Printf("Worker connecting to %s and running scripts in %s.\n", c.QueueAddr, c.Worker.WorkingDir)
	return worker, nil
}

func (w *Worker) LoadWhiteList(path string) error {

	whiteList := make(map[string]bool)
	var buf bytes.Buffer
	buf.WriteString(fmt.Sprintln("Whitelist:"))

	log.Println("Loading scripts in", path)

	var err error

	fileInfo, err := os.Stat(path)
	if err != nil {
		return err
	}
	if fileInfo.IsDir() {
		buf.WriteString(fmt.Sprintln("All"))
	} else {
		file, err := os.Open(path)
		if err != nil {
			return err
		}
		defer file.Close()

		scanner := bufio.NewScanner(file)

		for scanner.Scan() {
			whiteList[scanner.Text()] = true
		}
		if err != nil {
			log.Println(err)
			return err
		}

		err = scanner.Err()

		for script := range whiteList {
			buf.WriteString(fmt.Sprintln("	", script))
		}
	}

	w.WhiteList = whiteList
	log.Println(buf.String())
	return nil
}

func (w *Worker) JobRequestHandler(m *nsq.Message) error {
	// Initialize Job from request
	var job Job
	job.Status = STATUS_ERR
	job.ScriptDir = w.ScriptDir
	job.WorkingDir = w.WorkingDir
	job.StoreDir = w.StoreDir

	err := json.Unmarshal(m.Body, &job)
	if err != nil {
		log.Println("Invalid JSON request", err)
		return nil
	}

	// Try and run script
	if len(w.WhiteList) == 0 || w.WhiteList[job.Script] {
		log.Println("Dequeued request as Job", job.ID)

		resultChan := make(chan error, 1)
		go job.Execute(resultChan)
		err := <-resultChan
		if err != nil {
			job.ExecLog = err.Error()
		} else {
			job.Status = STATUS_OK
		}
	} else {
		msg := fmt.Sprintf("%s is not on script whitelist", job.Script)
		job.ExecLog = msg
	}

	log.Println(job.ExecLog)
	job.Log()
	job.Callback()
	return nil
}

func (w *Worker) ReloadRequestHandler(m *nsq.Message) error {
	whitelist := path.Clean(string(m.Body))
	err := w.LoadWhiteList(whitelist)
	if err != nil {
		log.Println("Failed to reload whitelist from", whitelist)
		return err
	}
	log.Println("Reloaded whitelist from", whitelist)
	return nil
}

func (w *Worker) Run() {
	// Set the message handler.
	w.Consumer.SetConcurrentHandlers(nsq.HandlerFunc(w.JobRequestHandler), w.Throughput)
	w.ReloadConsumer.SetHandler(nsq.HandlerFunc(w.ReloadRequestHandler))

	var err error

	// Connect the queue.
	fmt.Println("Connecting to", w.QueueAddr)
	err = w.Consumer.ConnectToNSQD(w.QueueAddr)
	if err != nil {
		log.Println(err)
	}

	err = w.ReloadConsumer.ConnectToNSQD(w.QueueAddr)
	if err != nil {
		log.Println(err)
	}
}

func (w *Worker) Stop() {
	w.Consumer.Stop()
	w.ReloadConsumer.Stop()
	return
}
<|endoftext|>


<commit_before>package fwk

import (
	"fmt"
)

type workercontrol struct {
	evts chan int64
	quit chan struct{}
	done chan struct{}
	errc chan error
}

type worker struct {
	slot  int
	keys  []string
	store datastore
	ctxs  []context
	msg   msgstream

	evts <-chan int64
	quit <-chan struct{}
	done chan<- struct{}
	errc chan<- error
}

func newWorker(i int, app *appmgr, ctrl *workercontrol) *worker {
	wrk := &worker{
		slot:  i,
		keys:  app.dflow.keys(),
		store: *app.store,
		ctxs:  make([]context, len(app.tsks)),
		msg:   NewMsgStream(fmt.Sprintf("%s-worker-%03d", app.name, i), app.msg.lvl, nil),
		evts:  ctrl.evts,
		quit:  ctrl.quit,
		done:  ctrl.done,
		errc:  ctrl.errc,
	}
	wrk.store.store = make(map[string]achan, len(wrk.keys))
	for j, tsk := range app.tsks {
		wrk.ctxs[j] = context{
			id:    -1,
			slot:  i,
			store: &wrk.store,
			msg:   NewMsgStream(tsk.Name(), app.msg.lvl, nil),
		}
	}

	go wrk.run(app.tsks)

	return wrk
}

func (wrk *worker) run(tsks []Task) {
	defer func() {
		wrk.done <- struct{}{}
	}()

	for {
		select {
		case ievt, ok := <-wrk.evts:
			if !ok {
				return
			}
			wrk.msg.Infof(">>> running evt=%d...\n", ievt)
			err := wrk.store.reset(wrk.keys)
			if err != nil {
				wrk.errc <- err
				return
			}

			evt := taskrunner{
				ievt: ievt,
				errc: make(chan error, len(tsks)),
				quit: make(chan struct{}),
			}
			for i, tsk := range tsks {
				go evt.run(i, wrk.ctxs[i], tsk)
			}
			ndone := 0
		errloop:
			for {
				select {
				case err, ok := <-evt.errc:
					if !ok {
						return
					}
					ndone++
					if err != nil {
						close(evt.quit)
						wrk.store.close()
						wrk.msg.flush()
						wrk.errc <- err
						return
					}
					if ndone == len(tsks) {
						break errloop
					}
				case <-wrk.quit:
					wrk.store.close()
					close(evt.quit)
					wrk.msg.flush()
					return
				}
			}
			wrk.store.close()
			close(evt.quit)
			wrk.msg.flush()

		case <-wrk.quit:
			wrk.store.close()
			return
		}
	}
}

type taskrunner struct {
	errc chan error
	quit chan struct{}

	ievt int64
}

func (run taskrunner) run(i int, ctx context, tsk Task) {
	ctx.id = run.ievt
	select {
	case run.errc <- tsk.Process(ctx):
		// FIXME(sbinet) dont be so eager to flush...
		ctx.msg.flush()
	case <-run.quit:
		ctx.msg.flush()
	}
}
<commit_msg>worker: less verbose<commit_after>package fwk

import (
	"fmt"
)

type workercontrol struct {
	evts chan int64
	quit chan struct{}
	done chan struct{}
	errc chan error
}

type worker struct {
	slot  int
	keys  []string
	store datastore
	ctxs  []context
	msg   msgstream

	evts <-chan int64
	quit <-chan struct{}
	done chan<- struct{}
	errc chan<- error
}

func newWorker(i int, app *appmgr, ctrl *workercontrol) *worker {
	wrk := &worker{
		slot:  i,
		keys:  app.dflow.keys(),
		store: *app.store,
		ctxs:  make([]context, len(app.tsks)),
		msg:   NewMsgStream(fmt.Sprintf("%s-worker-%03d", app.name, i), app.msg.lvl, nil),
		evts:  ctrl.evts,
		quit:  ctrl.quit,
		done:  ctrl.done,
		errc:  ctrl.errc,
	}
	wrk.store.store = make(map[string]achan, len(wrk.keys))
	for j, tsk := range app.tsks {
		wrk.ctxs[j] = context{
			id:    -1,
			slot:  i,
			store: &wrk.store,
			msg:   NewMsgStream(tsk.Name(), app.msg.lvl, nil),
		}
	}

	go wrk.run(app.tsks)

	return wrk
}

func (wrk *worker) run(tsks []Task) {
	defer func() {
		wrk.done <- struct{}{}
	}()

	for {
		select {
		case ievt, ok := <-wrk.evts:
			if !ok {
				return
			}
			wrk.msg.Debugf(">>> running evt=%d...\n", ievt)
			err := wrk.store.reset(wrk.keys)
			if err != nil {
				wrk.errc <- err
				return
			}

			evt := taskrunner{
				ievt: ievt,
				errc: make(chan error, len(tsks)),
				quit: make(chan struct{}),
			}
			for i, tsk := range tsks {
				go evt.run(i, wrk.ctxs[i], tsk)
			}
			ndone := 0
		errloop:
			for {
				select {
				case err, ok := <-evt.errc:
					if !ok {
						return
					}
					ndone++
					if err != nil {
						close(evt.quit)
						wrk.store.close()
						wrk.msg.flush()
						wrk.errc <- err
						return
					}
					if ndone == len(tsks) {
						break errloop
					}
				case <-wrk.quit:
					wrk.store.close()
					close(evt.quit)
					wrk.msg.flush()
					return
				}
			}
			wrk.store.close()
			close(evt.quit)
			wrk.msg.flush()

		case <-wrk.quit:
			wrk.store.close()
			return
		}
	}
}

type taskrunner struct {
	errc chan error
	quit chan struct{}

	ievt int64
}

func (run taskrunner) run(i int, ctx context, tsk Task) {
	ctx.id = run.ievt
	select {
	case run.errc <- tsk.Process(ctx):
		// FIXME(sbinet) dont be so eager to flush...
		ctx.msg.flush()
	case <-run.quit:
		ctx.msg.flush()
	}
}
<|endoftext|>


<commit_before><commit_msg>No need to repeat iota these days.<commit_after><|endoftext|>