<commit_before>// Package red implements server speaking redis serialization protocol.
package red

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"runtime/debug"
	"strings"
	"sync"
	"time"

	"github.com/artyom/resp"
)

// HandlerFunc is a type of command processing function. It should return value
// that's passed to resp.Encode function, so it's expected such function would
// return values of resp package types.
type HandlerFunc func(req Request) (interface{}, error)

// Request holds information about single redis command.
type Request struct {
	Name string   // lowercase command itself (first word)
	Args []string // command arguments
}

// NewServer returns initialized server.
func NewServer() *Server {
	return &Server{log: noopLogger{}}
}

// WithLogger configures server to use provided Logger.
func (s *Server) WithLogger(l Logger) {
	if l != nil {
		s.log = l
	}
}

// Handle registers handler for command with given name (case-insensitive)
func (s *Server) Handle(name string, h HandlerFunc) {
	if name == "" {
		panic("Handle called with empty name")
	}
	if h == nil {
		panic("Handle called with nil HandlerFunc")
	}
	if s.handlers == nil {
		s.handlers = make(map[string]HandlerFunc)
	}
	s.handlers[strings.ToLower(name)] = h
}

// Server implements server speaking RESP (REdis Serialization Protocol). Server
// automatically handles MULTI, EXEC, DISCARD commands for transactions, QUIT for
// client-initiated disconnect, other commands are expected to be implemented
// separately and registered with Handle method.
type Server struct {
	log      Logger
	handlers map[string]HandlerFunc
	mu       sync.Mutex // used to serialize transactions
}

// HandleConn processes single client connection, automatically handling
// following commands MULTI/EXEC/DISCARD (transactions), QUIT (client disconnect). It
// calls user-provided handlers for registered commands.
func (s *Server) HandleConn(conn io.ReadWriteCloser) error {
	defer conn.Close()
	rd := bufio.NewReader(conn)
	var tx []Request
	var inTx bool  // if we're inside transaction
	var errTx bool // true if transaction seen error and should be discarded
	var err error
	for {
		if err != nil {
			return err
		}
		req, err := resp.DecodeRequest(rd)
		switch err {
		case nil:
		case resp.ErrInvalidRequest:
			err = resp.Encode(conn, resp.Error("ERR unknown command"))
			continue
		default:
			return err
		}
		cmd := strings.ToLower(req[0])
		switch cmd {
		case "quit":
			return nil
		case "discard":
			if len(req) != 1 {
				if inTx {
					errTx = true
				}
				err = resp.Encode(conn, errWrongArgs(cmd))
				continue
			}
			if !inTx {
				err = resp.Encode(conn, resp.Error("ERR DISCARD without MULTI"))
				continue
			}
			inTx, errTx, tx = false, false, tx[:0]
			err = resp.Encode(conn, resp.OK)
			continue
		case "multi":
			if len(req) != 1 {
				if inTx {
					errTx = true
				}
				err = resp.Encode(conn, errWrongArgs(cmd))
				continue
			}
			if inTx {
				errTx = true
				err = resp.Encode(conn, resp.Error("ERR MULTI calls can not be nested"))
				continue
			}
			inTx, errTx = true, false
			err = resp.Encode(conn, resp.OK)
			continue
		case "exec":
			if len(req) != 1 {
				if inTx {
					errTx = true
				}
				err = resp.Encode(conn, errWrongArgs(cmd))
				continue
			}
			if !inTx {
				err = resp.Encode(conn, resp.Error("ERR EXEC without MULTI"))
				continue
			}
			if errTx {
				inTx, errTx = false, false
				tx = tx[:0]
				err = resp.Encode(conn, resp.Error("EXECABORT Transaction discarded because of previous errors."))
				continue
			}
		default:
			h, ok := s.handlers[cmd]
			if !ok {
				if inTx {
					errTx = true
				}
				err = resp.Encode(conn, errNoCmd(cmd))
				continue
			}
			if inTx {
				if !errTx {
					tx = append(tx, Request{Name: cmd, Args: req[1:]})
				}
				err = resp.Encode(conn, resp.SimpleString("QUEUED"))
				continue
			}
			s.mu.Lock()
			err = resp.Encode(conn, runHandler(h, Request{Name: cmd, Args: req[1:]}))
			s.mu.Unlock()
			continue
		}

		txReplies := make(resp.Array, 0, len(tx))
		s.mu.Lock()
		for _, r := range tx {
			h, ok := s.handlers[r.Name]
			if !ok {
				txReplies = append(txReplies, errNoCmd(r.Name))
				continue
			}
			txReplies = append(txReplies, runHandler(h, r))
		}
		s.mu.Unlock()
		inTx, errTx = false, false
		tx = tx[:0]
		err = resp.Encode(conn, txReplies)
	}
}

// ListenAndServe listens on TCP network address addr and then calls Serve to
// handle requests on incoming connections.
func (s *Server) ListenAndServe(addr string) error {
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	return s.Serve(tcpKeepAliveListener{ln.(*net.TCPListener)})
}

// Serve accepts incoming connections on the Listener l, creating a new service
// goroutine for each.
func (s *Server) Serve(l net.Listener) error {
	defer l.Close()
	for {
		conn, err := l.Accept()
		if err != nil {
			return err
		}
		go func(c net.Conn) {
			switch err := s.HandleConn(c); err {
			case nil, io.EOF:
			default:
				if s.log != nil {
					s.log.Println(err)
				}
			}
		}(conn)
	}
}

// Logger is a set of methods used to log information. *log.Logger implements
// this interface.
type Logger interface {
	Print(v ...interface{})
	Printf(format string, v ...interface{})
	Println(v ...interface{})
}

type noopLogger struct{}

func (noopLogger) Print(v ...interface{})                 {}
func (noopLogger) Printf(format string, v ...interface{}) {}
func (noopLogger) Println(v ...interface{})               {}

func errNoCmd(name string) resp.Error { return resp.Error("ERR unknown command '" + name + "'") }

// ErrWrongArgs are expected to be returned by HandlerFunc implementations when
// number of arguments are wrong. This error is automatically passed to client
// with command name annotated.
var ErrWrongArgs = errors.New("wrong number of arguments")

func errWrongArgs(name string) resp.Error {
	return resp.Error("ERR wrong number of arguments for '" + name + "' command")
}

// runHandler calls HandlerFunc h with given Request r and returns single value
// that is expected to be passed to resp.Encode.
//
// If h returns a non-nil error or panics, then runHandler returns resp.Error
// type, otherwise it's an interface{} returned by h.
func runHandler(h HandlerFunc, r Request) interface{} {
	var panicked bool
	v, err := func() (interface{}, error) {
		defer func() {
			if x := recover(); x != nil {
				panicked = true
				fmt.Fprintln(os.Stderr, "panic recovered:", x)
				debug.PrintStack()
			}
		}()
		return h(r)
	}()
	if panicked {
		return resp.Error("ERR internal server error (panic)")
	}
	if err != nil {
		if err == ErrWrongArgs {
			return errWrongArgs(r.Name)
		}
		text := err.Error()
		if strings.ContainsAny(text, "\r\n") {
			text = strings.Replace(text, "\r", " ", -1)
			text = strings.Replace(text, "\n", " ", -1)
		}
		return resp.Error("ERR " + text)
	}
	return v
}

// tcpKeepAliveListener sets TCP keep-alive timeouts on accepted
// connections. It's used by ListenAndServe and ListenAndServeTLS so
// dead TCP connections (e.g. closing laptop mid-download) eventually
// go away.
type tcpKeepAliveListener struct {
	*net.TCPListener
}

func (ln tcpKeepAliveListener) Accept() (c net.Conn, err error) {
	tc, err := ln.AcceptTCP()
	if err != nil {
		return
	}
	tc.SetKeepAlive(true)
	tc.SetKeepAlivePeriod(3 * time.Minute)
	return tc, nil
}
<commit_msg>Allow disabling transaction safety with WithUnsafeTx()<commit_after>// Package red implements server speaking redis serialization protocol.
package red

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"runtime/debug"
	"strings"
	"sync"
	"time"

	"github.com/artyom/resp"
)

// HandlerFunc is a type of command processing function. It should return value
// that's passed to resp.Encode function, so it's expected such function would
// return values of resp package types.
type HandlerFunc func(req Request) (interface{}, error)

// Request holds information about single redis command.
type Request struct {
	Name string   // lowercase command itself (first word)
	Args []string // command arguments
}

// NewServer returns initialized server.
func NewServer() *Server {
	return &Server{log: noopLogger{}}
}

// WithLogger configures server to use provided Logger.
func (s *Server) WithLogger(l Logger) {
	if l != nil {
		s.log = l
	}
}

// WithUnsafeTx disables command processing serialization that guarantees
// redis-like transaction safety.
func (s *Server) WithUnsafeTx() { s.unsafe = true }

// Handle registers handler for command with given name (case-insensitive)
func (s *Server) Handle(name string, h HandlerFunc) {
	if name == "" {
		panic("Handle called with empty name")
	}
	if h == nil {
		panic("Handle called with nil HandlerFunc")
	}
	if s.handlers == nil {
		s.handlers = make(map[string]HandlerFunc)
	}
	s.handlers[strings.ToLower(name)] = h
}

// Server implements server speaking RESP (REdis Serialization Protocol). Server
// automatically handles MULTI, EXEC, DISCARD commands for transactions, QUIT for
// client-initiated disconnect, other commands are expected to be implemented
// separately and registered with Handle method.
type Server struct {
	log      Logger
	handlers map[string]HandlerFunc
	mu       sync.Mutex // used to serialize transactions
	unsafe   bool       // whether to skip transaction serialization
}

func (s *Server) txLock() {
	if s.unsafe {
		return
	}
	s.mu.Lock()
}
func (s *Server) txUnlock() {
	if s.unsafe {
		return
	}
	s.mu.Unlock()
}

// HandleConn processes single client connection, automatically handling
// following commands MULTI/EXEC/DISCARD (transactions), QUIT (client disconnect). It
// calls user-provided handlers for registered commands.
func (s *Server) HandleConn(conn io.ReadWriteCloser) error {
	defer conn.Close()
	rd := bufio.NewReader(conn)
	var tx []Request
	var inTx bool  // if we're inside transaction
	var errTx bool // true if transaction seen error and should be discarded
	var err error
	for {
		if err != nil {
			return err
		}
		req, err := resp.DecodeRequest(rd)
		switch err {
		case nil:
		case resp.ErrInvalidRequest:
			err = resp.Encode(conn, resp.Error("ERR unknown command"))
			continue
		default:
			return err
		}
		cmd := strings.ToLower(req[0])
		switch cmd {
		case "quit":
			return nil
		case "discard":
			if len(req) != 1 {
				if inTx {
					errTx = true
				}
				err = resp.Encode(conn, errWrongArgs(cmd))
				continue
			}
			if !inTx {
				err = resp.Encode(conn, resp.Error("ERR DISCARD without MULTI"))
				continue
			}
			inTx, errTx, tx = false, false, tx[:0]
			err = resp.Encode(conn, resp.OK)
			continue
		case "multi":
			if len(req) != 1 {
				if inTx {
					errTx = true
				}
				err = resp.Encode(conn, errWrongArgs(cmd))
				continue
			}
			if inTx {
				errTx = true
				err = resp.Encode(conn, resp.Error("ERR MULTI calls can not be nested"))
				continue
			}
			inTx, errTx = true, false
			err = resp.Encode(conn, resp.OK)
			continue
		case "exec":
			if len(req) != 1 {
				if inTx {
					errTx = true
				}
				err = resp.Encode(conn, errWrongArgs(cmd))
				continue
			}
			if !inTx {
				err = resp.Encode(conn, resp.Error("ERR EXEC without MULTI"))
				continue
			}
			if errTx {
				inTx, errTx = false, false
				tx = tx[:0]
				err = resp.Encode(conn, resp.Error("EXECABORT Transaction discarded because of previous errors."))
				continue
			}
		default:
			h, ok := s.handlers[cmd]
			if !ok {
				if inTx {
					errTx = true
				}
				err = resp.Encode(conn, errNoCmd(cmd))
				continue
			}
			if inTx {
				if !errTx {
					tx = append(tx, Request{Name: cmd, Args: req[1:]})
				}
				err = resp.Encode(conn, resp.SimpleString("QUEUED"))
				continue
			}
			s.txLock()
			err = resp.Encode(conn, runHandler(h, Request{Name: cmd, Args: req[1:]}))
			s.txUnlock()
			continue
		}

		txReplies := make(resp.Array, 0, len(tx))
		s.txLock()
		for _, r := range tx {
			h, ok := s.handlers[r.Name]
			if !ok {
				txReplies = append(txReplies, errNoCmd(r.Name))
				continue
			}
			txReplies = append(txReplies, runHandler(h, r))
		}
		s.txUnlock()
		inTx, errTx = false, false
		tx = tx[:0]
		err = resp.Encode(conn, txReplies)
	}
}

// ListenAndServe listens on TCP network address addr and then calls Serve to
// handle requests on incoming connections.
func (s *Server) ListenAndServe(addr string) error {
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	return s.Serve(tcpKeepAliveListener{ln.(*net.TCPListener)})
}

// Serve accepts incoming connections on the Listener l, creating a new service
// goroutine for each.
func (s *Server) Serve(l net.Listener) error {
	defer l.Close()
	for {
		conn, err := l.Accept()
		if err != nil {
			return err
		}
		go func(c net.Conn) {
			switch err := s.HandleConn(c); err {
			case nil, io.EOF:
			default:
				if s.log != nil {
					s.log.Println(err)
				}
			}
		}(conn)
	}
}

// Logger is a set of methods used to log information. *log.Logger implements
// this interface.
type Logger interface {
	Print(v ...interface{})
	Printf(format string, v ...interface{})
	Println(v ...interface{})
}

type noopLogger struct{}

func (noopLogger) Print(v ...interface{})                 {}
func (noopLogger) Printf(format string, v ...interface{}) {}
func (noopLogger) Println(v ...interface{})               {}

func errNoCmd(name string) resp.Error { return resp.Error("ERR unknown command '" + name + "'") }

// ErrWrongArgs are expected to be returned by HandlerFunc implementations when
// number of arguments are wrong. This error is automatically passed to client
// with command name annotated.
var ErrWrongArgs = errors.New("wrong number of arguments")

func errWrongArgs(name string) resp.Error {
	return resp.Error("ERR wrong number of arguments for '" + name + "' command")
}

// runHandler calls HandlerFunc h with given Request r and returns single value
// that is expected to be passed to resp.Encode.
//
// If h returns a non-nil error or panics, then runHandler returns resp.Error
// type, otherwise it's an interface{} returned by h.
func runHandler(h HandlerFunc, r Request) interface{} {
	var panicked bool
	v, err := func() (interface{}, error) {
		defer func() {
			if x := recover(); x != nil {
				panicked = true
				fmt.Fprintln(os.Stderr, "panic recovered:", x)
				debug.PrintStack()
			}
		}()
		return h(r)
	}()
	if panicked {
		return resp.Error("ERR internal server error (panic)")
	}
	if err != nil {
		if err == ErrWrongArgs {
			return errWrongArgs(r.Name)
		}
		text := err.Error()
		if strings.ContainsAny(text, "\r\n") {
			text = strings.Replace(text, "\r", " ", -1)
			text = strings.Replace(text, "\n", " ", -1)
		}
		return resp.Error("ERR " + text)
	}
	return v
}

// tcpKeepAliveListener sets TCP keep-alive timeouts on accepted
// connections. It's used by ListenAndServe and ListenAndServeTLS so
// dead TCP connections (e.g. closing laptop mid-download) eventually
// go away.
type tcpKeepAliveListener struct {
	*net.TCPListener
}

func (ln tcpKeepAliveListener) Accept() (c net.Conn, err error) {
	tc, err := ln.AcceptTCP()
	if err != nil {
		return
	}
	tc.SetKeepAlive(true)
	tc.SetKeepAlivePeriod(3 * time.Minute)
	return tc, nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/hex"
	"fmt"
	"log"

	"github.com/dt/thile/gen"
	"github.com/foursquare/gohfile"
)

type ThriftRpcImpl struct {
	*hfile.CollectionSet
}

func (cs *ThriftRpcImpl) GetValuesSingle(req *gen.SingleHFileKeyRequest) (r *gen.SingleHFileKeyResponse, err error) {
	if Settings.debug {
		log.Printf("[GetValuesSingle] %s (%d keys)\n", *req.HfileName, len(req.SortedKeys))
	}
	reader, err := cs.ScannerFor(*req.HfileName)
	if err != nil {
		return nil, err
	}
	defer reader.Release()

	if req.PerKeyValueLimit != nil {
		// TODO(davidt) impl
		log.Println("[GetValuesSingle] PerKeyValueLimit. oh well.")
	}

	if req.CountOnly != nil {
		// TODO(davidt) impl
		log.Println("[GetValuesSingle] CountOnly. oh well.")
	}

	res := new(gen.SingleHFileKeyResponse)
	res.Values = make(map[int32][]byte)
	found := int32(0)

	for idx, key := range req.SortedKeys {
		if Settings.debug {
			log.Printf("[GetValuesSingle] key: %s\n", hex.EncodeToString(key))
		}
		value, err, ok := reader.GetFirst(key)
		if err != nil {
			return nil, err
		}
		if ok {
			found++
			res.Values[int32(idx)] = value
		}
	}

	if Settings.debug {
		log.Printf("[GetValuesSingle] %s found %d of %d.\n", *req.HfileName, found, len(req.SortedKeys))
	}
	res.KeyCount = &found
	return res, nil
}

func (cs *ThriftRpcImpl) GetValuesMulti(req *gen.SingleHFileKeyRequest) (r *gen.MultiHFileKeyResponse, err error) {
	if Settings.debug {
		log.Println("[GetValuesMulti]", len(req.SortedKeys))
	}

	reader, err := cs.ScannerFor(*req.HfileName)
	if err != nil {
		return nil, err
	}
	defer reader.Release()

	res := new(gen.MultiHFileKeyResponse)
	found := int32(0)

	for idx, key := range req.SortedKeys {
		values, err := reader.GetAll(key)
		if err != nil {
			return nil, err
		}
		if len(values) > 0 {
			found += int32(len(values))
			res.Values[int32(idx)] = values
		}
	}

	res.KeyCount = &found
	return res, nil

}

func (cs *ThriftRpcImpl) GetValuesForPrefixes(req *gen.PrefixRequest) (r *gen.PrefixResponse, err error) {
	res := new(gen.PrefixResponse)
	if i, err := cs.IteratorFor(*req.HfileName); err != nil {
		return nil, err
	} else {
		defer i.Release()
		if res.Values, err = i.AllForPrfixes(req.SortedKeys); err != nil {
			return nil, err
		} else {
			return res, nil
		}
	}
}

func (cs *ThriftRpcImpl) GetValuesMultiSplitKeys(req *gen.MultiHFileSplitKeyRequest) (r *gen.KeyToValuesResponse, err error) {
	res := make(map[string][][]byte)
	scanner, err := cs.ScannerFor(*req.HfileName)
	if err != nil {
		return nil, err
	}

	for _, parts := range RevProduct(req.SplitKey) {
		// TODO(davidt): avoid allocing concated key by adding split-key search lower down.
		key := bytes.Join(parts, nil)

		if values, err := scanner.GetAll(key); err != nil {
			return nil, err
		} else if len(values) > 0 {
			res[string(key)] = values
		}
	}
	return &gen.KeyToValuesResponse{res}, nil
}

func (cs *ThriftRpcImpl) GetIterator(req *gen.IteratorRequest) (*gen.IteratorResponse, error) {
	// 	HfileName     *string `thrift:"hfileName,1" json:"hfileName"`
	// 	IncludeValues *bool   `thrift:"includeValues,2" json:"includeValues"`
	// 	LastKey       []byte  `thrift:"lastKey,3" json:"lastKey"`
	// 	SkipKeys      *int32  `thrift:"skipKeys,4" json:"skipKeys"`
	// 	ResponseLimit *int32  `thrift:"responseLimit,5" json:"responseLimit"`
	// 	EndKey        []byte  `thrift:"endKey,6" json:"endKey"`
	var err error

	if req.ResponseLimit == nil {
		return nil, fmt.Errorf("Missing limit.")
	}
	limit := int(*req.ResponseLimit)

	it, err := cs.IteratorFor(*req.HfileName)
	if err != nil {
		return nil, err
	}
	defer it.Release()

	remaining := false

	if req.LastKey != nil {
		remaining, err = it.Seek(req.LastKey)
	} else {
		remaining, err = it.Next()
	}

	if err != nil {
		return nil, err
	}

	res := new(gen.IteratorResponse)

	if !remaining {
		return res, nil
	}

	skipKeys := int32(0)
	lastKey := it.Key()

	if toSkip := req.GetSkipKeys(); toSkip > 0 {
		for i := int32(0); i < toSkip && remaining; i++ {
			if bytes.Equal(lastKey, it.Key()) {
				skipKeys = skipKeys + 1
			} else {
				skipKeys = 0
			}

			lastKey = it.Key()

			remaining, err = it.Next()
			if err != nil {
				return nil, err
			}
		}
		if !remaining {
			return res, nil
		}
	}

	if req.EndKey != nil {
		remaining = remaining && !hfile.After(it.Key(), req.EndKey)
	}

	r := make([]*gen.KeyValueItem, 0)
	for i := 0; i < limit && remaining; i++ {
		v := []byte{}
		if req.IncludeValues == nil || *req.IncludeValues {
			v = it.Value()
		}
		r = append(r, &gen.KeyValueItem{it.Key(), v})

		if bytes.Equal(lastKey, it.Key()) {
			skipKeys = skipKeys + 1
		} else {
			skipKeys = 0
		}
		lastKey = it.Key()

		remaining, err = it.Next()
		if err != nil {
			return nil, err
		}
		if req.EndKey != nil {
			remaining = remaining && !hfile.After(it.Key(), req.EndKey)
		}
	}
	return &gen.IteratorResponse{r, lastKey, &skipKeys}, nil
}

func (cs *ThriftRpcImpl) GetInfo(req *gen.InfoRequest) (r []*gen.HFileInfo, err error) {
	return nil, fmt.Errorf("Not implemented")
}

func (cs *ThriftRpcImpl) TestTimeout(waitInMillis int32) (r int32, err error) {
	return 0, fmt.Errorf("Not implemented")
}
<commit_msg>switch to upstreams GetScanner/GetIterator on reader<commit_after>package main

import (
	"bytes"
	"encoding/hex"
	"fmt"
	"log"

	"github.com/dt/thile/gen"
	"github.com/foursquare/gohfile"
)

type ThriftRpcImpl struct {
	*hfile.CollectionSet
}

func (cs *ThriftRpcImpl) GetValuesSingle(req *gen.SingleHFileKeyRequest) (r *gen.SingleHFileKeyResponse, err error) {
	if Settings.debug {
		log.Printf("[GetValuesSingle] %s (%d keys)\n", *req.HfileName, len(req.SortedKeys))
	}
	hfile, err := cs.ReaderFor(*req.HfileName)
	if err != nil {
		return nil, err
	}
	reader := hfile.GetScanner()
	defer reader.Release()

	if req.PerKeyValueLimit != nil {
		// TODO(davidt) impl
		log.Println("[GetValuesSingle] PerKeyValueLimit. oh well.")
	}

	if req.CountOnly != nil {
		// TODO(davidt) impl
		log.Println("[GetValuesSingle] CountOnly. oh well.")
	}

	res := new(gen.SingleHFileKeyResponse)
	res.Values = make(map[int32][]byte)
	found := int32(0)

	for idx, key := range req.SortedKeys {
		if Settings.debug {
			log.Printf("[GetValuesSingle] key: %s\n", hex.EncodeToString(key))
		}
		value, err, ok := reader.GetFirst(key)
		if err != nil {
			return nil, err
		}
		if ok {
			found++
			res.Values[int32(idx)] = value
		}
	}

	if Settings.debug {
		log.Printf("[GetValuesSingle] %s found %d of %d.\n", *req.HfileName, found, len(req.SortedKeys))
	}
	res.KeyCount = &found
	return res, nil
}

func (cs *ThriftRpcImpl) GetValuesMulti(req *gen.SingleHFileKeyRequest) (r *gen.MultiHFileKeyResponse, err error) {
	if Settings.debug {
		log.Println("[GetValuesMulti]", len(req.SortedKeys))
	}

	hfile, err := cs.ReaderFor(*req.HfileName)
	if err != nil {
		return nil, err
	}
	reader := hfile.GetScanner()
	defer reader.Release()

	res := new(gen.MultiHFileKeyResponse)
	found := int32(0)

	for idx, key := range req.SortedKeys {
		values, err := reader.GetAll(key)
		if err != nil {
			return nil, err
		}
		if len(values) > 0 {
			found += int32(len(values))
			res.Values[int32(idx)] = values
		}
	}

	res.KeyCount = &found
	return res, nil

}

func (cs *ThriftRpcImpl) GetValuesForPrefixes(req *gen.PrefixRequest) (r *gen.PrefixResponse, err error) {
	res := new(gen.PrefixResponse)
	if reader, err := cs.ReaderFor(*req.HfileName); err != nil {
		return nil, err
	} else {
		i := reader.GetIterator()
		defer i.Release()
		if res.Values, err = i.AllForPrfixes(req.SortedKeys); err != nil {
			return nil, err
		} else {
			return res, nil
		}
	}
}

func (cs *ThriftRpcImpl) GetValuesMultiSplitKeys(req *gen.MultiHFileSplitKeyRequest) (r *gen.KeyToValuesResponse, err error) {
	res := make(map[string][][]byte)
	reader, err := cs.ReaderFor(*req.HfileName)
	if err != nil {
		return nil, err
	}
	scanner := reader.GetScanner()
	defer scanner.Release()

	for _, parts := range RevProduct(req.SplitKey) {
		// TODO(davidt): avoid allocing concated key by adding split-key search lower down.
		key := bytes.Join(parts, nil)

		if values, err := scanner.GetAll(key); err != nil {
			return nil, err
		} else if len(values) > 0 {
			res[string(key)] = values
		}
	}
	return &gen.KeyToValuesResponse{res}, nil
}

func (cs *ThriftRpcImpl) GetIterator(req *gen.IteratorRequest) (*gen.IteratorResponse, error) {
	// 	HfileName     *string `thrift:"hfileName,1" json:"hfileName"`
	// 	IncludeValues *bool   `thrift:"includeValues,2" json:"includeValues"`
	// 	LastKey       []byte  `thrift:"lastKey,3" json:"lastKey"`
	// 	SkipKeys      *int32  `thrift:"skipKeys,4" json:"skipKeys"`
	// 	ResponseLimit *int32  `thrift:"responseLimit,5" json:"responseLimit"`
	// 	EndKey        []byte  `thrift:"endKey,6" json:"endKey"`
	var err error

	if req.ResponseLimit == nil {
		return nil, fmt.Errorf("Missing limit.")
	}
	limit := int(*req.ResponseLimit)

	reader, err := cs.ReaderFor(*req.HfileName)
	if err != nil {
		return nil, err
	}
	it := reader.GetIterator()
	defer it.Release()

	remaining := false

	if req.LastKey != nil {
		remaining, err = it.Seek(req.LastKey)
	} else {
		remaining, err = it.Next()
	}

	if err != nil {
		return nil, err
	}

	res := new(gen.IteratorResponse)

	if !remaining {
		return res, nil
	}

	skipKeys := int32(0)
	lastKey := it.Key()

	if toSkip := req.GetSkipKeys(); toSkip > 0 {
		for i := int32(0); i < toSkip && remaining; i++ {
			if bytes.Equal(lastKey, it.Key()) {
				skipKeys = skipKeys + 1
			} else {
				skipKeys = 0
			}

			lastKey = it.Key()

			remaining, err = it.Next()
			if err != nil {
				return nil, err
			}
		}
		if !remaining {
			return res, nil
		}
	}

	if req.EndKey != nil {
		remaining = remaining && !hfile.After(it.Key(), req.EndKey)
	}

	r := make([]*gen.KeyValueItem, 0)
	for i := 0; i < limit && remaining; i++ {
		v := []byte{}
		if req.IncludeValues == nil || *req.IncludeValues {
			v = it.Value()
		}
		r = append(r, &gen.KeyValueItem{it.Key(), v})

		if bytes.Equal(lastKey, it.Key()) {
			skipKeys = skipKeys + 1
		} else {
			skipKeys = 0
		}
		lastKey = it.Key()

		remaining, err = it.Next()
		if err != nil {
			return nil, err
		}
		if req.EndKey != nil {
			remaining = remaining && !hfile.After(it.Key(), req.EndKey)
		}
	}
	return &gen.IteratorResponse{r, lastKey, &skipKeys}, nil
}

func (cs *ThriftRpcImpl) GetInfo(req *gen.InfoRequest) (r []*gen.HFileInfo, err error) {
	return nil, fmt.Errorf("Not implemented")
}

func (cs *ThriftRpcImpl) TestTimeout(waitInMillis int32) (r int32, err error) {
	return 0, fmt.Errorf("Not implemented")
}
<|endoftext|>


<commit_before>package rye

import (
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"reflect"
	"runtime"
	"strconv"
	"strings"
	"time"

	//log "github.com/sirupsen/logrus"

	"github.com/cactus/go-statsd-client/statsd"
)

//go:generate counterfeiter -o fakes/statsdfakes/fake_statter.go $GOPATH/src/github.com/cactus/go-statsd-client/statsd/client.go Statter
//go:generate perl -pi -e 's/$GOPATH\/src\///g' fakes/statsdfakes/fake_statter.go

// MWHandler struct is used to configure and access rye's basic functionality.
type MWHandler struct {
	Config         Config
	beforeHandlers []Handler
}

// CustomStatter allows the client to log any additional statsD metrics Rye
// computes around the request handler.
type CustomStatter interface {
	ReportStats(string, time.Duration, *http.Request, *Response) error
}

// Config struct allows you to set a reference to a statsd.Statter and include it's stats rate.
type Config struct {
	Statter  statsd.Statter
	StatRate float32

	// toggle types of stats sent
	NoErrStats        bool
	NoDurationStats   bool
	NoStatusCodeStats bool

	// Customer Statter for the client
	CustomStatter CustomStatter
}

// JSONStatus is a simple container used for conveying status messages.
type JSONStatus struct {
	Message string `json:"message"`
	Status  string `json:"status"`
}

// Response struct is utilized by middlewares as a way to share state;
// ie. a middleware can return a *Response as a way to indicate
// that further middleware execution should stop (without an error) or return a
// a hard error by setting `Err` + `StatusCode`.
type Response struct {
	Err           error
	StatusCode    int
	StopExecution bool
	Context       context.Context
}

// Error bubbles a response error providing an implementation of the Error interface.
// It returns the error as a string.
func (r *Response) Error() string {
	return r.Err.Error()
}

// Handler is the primary type that any rye middleware must implement to be called in the Handle() function.
// In order to use this you must return a *rye.Response.
type Handler func(w http.ResponseWriter, r *http.Request) *Response

// Constructor for new instantiating new rye instances
// It returns a constructed *MWHandler instance.
func NewMWHandler(config Config) *MWHandler {
	return &MWHandler{
		Config: config,
	}
}

// Use adds a handler to every request. All handlers set up with use
// are fired first and then any route specific handlers are called
func (m *MWHandler) Use(handler Handler) {
	m.beforeHandlers = append(m.beforeHandlers, handler)
}

// The Handle function is the primary way to set up your chain of middlewares to be called by rye.
// It returns a http.HandlerFunc from net/http that can be set as a route in your http server.
func (m *MWHandler) Handle(customHandlers []Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		exit := false
		for _, handler := range m.beforeHandlers {
			exit, r = m.do(w, r, handler)
			if exit {
				return
			}
		}

		for _, handler := range customHandlers {
			exit, r = m.do(w, r, handler)
			if exit {
				return
			}
		}
	})
}

func (m *MWHandler) do(w http.ResponseWriter, r *http.Request, handler Handler) (bool, *http.Request) {
	var resp *Response

	// Record handler runtime
	func() {
		statusCode := "2xx"
		startTime := time.Now()

		if resp = handler(w, r); resp != nil {
			func() {
				// Stop execution if it's passed
				if resp.StopExecution {
					return
				}

				// If a context is returned, we will
				// replace the current request with a new request
				if resp.Context != nil {
					r = r.WithContext(resp.Context)
					return
				}

				// If there's no error but we have a response
				if resp.Err == nil {
					resp.Err = errors.New("Problem with middleware; neither Err or StopExecution is set")
					resp.StatusCode = http.StatusInternalServerError
				}

				// Now assume we have an error.
				if m.Config.Statter != nil && resp.StatusCode >= 500 {
					go m.reportError()
				}

				// Write the error out
				WriteJSONStatus(w, "error", resp.Error(), resp.StatusCode)
			}()

			if resp.StatusCode > 0 {
				statusCode = strconv.Itoa(resp.StatusCode)
			}
		}

		handlerName := getFuncName(handler)

		if m.Config.Statter != nil {
			// Record runtime metric
			go m.reportDuration(handlerName, startTime)

			// Record status code metric (default 2xx)
			go m.reportStatusCode(handlerName, statusCode)
		}

		// If a CustomStatter is set, send the handler metrics to it.
		// This allows the client to handle these metrics however it wants.
		if m.Config.CustomStatter != nil && resp != nil {
			go m.Config.CustomStatter.ReportStats(handlerName, time.Since(startTime), r, resp)
		}
	}()

	// stop executing rest of the
	// handlers if we encounter an error
	if resp != nil && (resp.StopExecution || resp.Err != nil) {
		return true, r
	}

	return false, r
}

func (m *MWHandler) reportError() {
	if m.Config.NoErrStats {
		return
	}

	m.Config.Statter.Inc("errors", 1, m.Config.StatRate)
}

func (m *MWHandler) reportDuration(handlerName string, startTime time.Time) {
	if m.Config.NoDurationStats {
		return
	}

	m.Config.Statter.TimingDuration(
		"handlers."+handlerName+".runtime",
		time.Since(startTime), // delta
		m.Config.StatRate,
	)
}

func (m *MWHandler) reportStatusCode(handlerName string, statusCode string) {
	if m.Config.NoStatusCodeStats {
		return
	}

	m.Config.Statter.Inc(
		"handlers."+handlerName+"."+statusCode,
		1,
		m.Config.StatRate,
	)
}

// WriteJSONStatus is a wrapper for WriteJSONResponse that returns a marshalled JSONStatus blob
func WriteJSONStatus(rw http.ResponseWriter, status, message string, statusCode int) {
	jsonData, _ := json.Marshal(&JSONStatus{
		Message: message,
		Status:  status,
	})

	WriteJSONResponse(rw, statusCode, jsonData)
}

// WriteJSONResponse writes data and status code to the ResponseWriter
func WriteJSONResponse(rw http.ResponseWriter, statusCode int, content []byte) {
	rw.Header().Set("Content-Type", "application/json")
	rw.WriteHeader(statusCode)
	rw.Write(content)
}

// getFuncName uses reflection to determine a given function name
// It returns a string version of the function name (and performs string cleanup)
func getFuncName(i interface{}) string {
	fullName := runtime.FuncForPC(reflect.ValueOf(i).Pointer()).Name()
	ns := strings.Split(fullName, ".")

	// when we get a method (not a raw function) it comes attached to whatever struct is in its
	// method receiver via a function closure, this is not precisely the same as that method itself
	// so the compiler appends "-fm" so the name of the closure does not conflict with the actual function
	// http://grokbase.com/t/gg/golang-nuts/153jyb5b7p/go-nuts-fm-suffix-in-function-name-what-does-it-mean#20150318ssinqqzrmhx2ep45wjkxsa4rua
	return strings.TrimSuffix(ns[len(ns)-1], ")-fm")
}
<commit_msg>Added named arguments to ReportStats interface<commit_after>package rye

import (
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"reflect"
	"runtime"
	"strconv"
	"strings"
	"time"

	//log "github.com/sirupsen/logrus"

	"github.com/cactus/go-statsd-client/statsd"
)

//go:generate counterfeiter -o fakes/statsdfakes/fake_statter.go $GOPATH/src/github.com/cactus/go-statsd-client/statsd/client.go Statter
//go:generate perl -pi -e 's/$GOPATH\/src\///g' fakes/statsdfakes/fake_statter.go

// MWHandler struct is used to configure and access rye's basic functionality.
type MWHandler struct {
	Config         Config
	beforeHandlers []Handler
}

// CustomStatter allows the client to log any additional statsD metrics Rye
// computes around the request handler.
type CustomStatter interface {
	ReportStats(handlerName string, elapsedTime time.Duration, req *http.Request, resp *Response) error
}

// Config struct allows you to set a reference to a statsd.Statter and include it's stats rate.
type Config struct {
	Statter  statsd.Statter
	StatRate float32

	// toggle types of stats sent
	NoErrStats        bool
	NoDurationStats   bool
	NoStatusCodeStats bool

	// Customer Statter for the client
	CustomStatter CustomStatter
}

// JSONStatus is a simple container used for conveying status messages.
type JSONStatus struct {
	Message string `json:"message"`
	Status  string `json:"status"`
}

// Response struct is utilized by middlewares as a way to share state;
// ie. a middleware can return a *Response as a way to indicate
// that further middleware execution should stop (without an error) or return a
// a hard error by setting `Err` + `StatusCode`.
type Response struct {
	Err           error
	StatusCode    int
	StopExecution bool
	Context       context.Context
}

// Error bubbles a response error providing an implementation of the Error interface.
// It returns the error as a string.
func (r *Response) Error() string {
	return r.Err.Error()
}

// Handler is the primary type that any rye middleware must implement to be called in the Handle() function.
// In order to use this you must return a *rye.Response.
type Handler func(w http.ResponseWriter, r *http.Request) *Response

// Constructor for new instantiating new rye instances
// It returns a constructed *MWHandler instance.
func NewMWHandler(config Config) *MWHandler {
	return &MWHandler{
		Config: config,
	}
}

// Use adds a handler to every request. All handlers set up with use
// are fired first and then any route specific handlers are called
func (m *MWHandler) Use(handler Handler) {
	m.beforeHandlers = append(m.beforeHandlers, handler)
}

// The Handle function is the primary way to set up your chain of middlewares to be called by rye.
// It returns a http.HandlerFunc from net/http that can be set as a route in your http server.
func (m *MWHandler) Handle(customHandlers []Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		exit := false
		for _, handler := range m.beforeHandlers {
			exit, r = m.do(w, r, handler)
			if exit {
				return
			}
		}

		for _, handler := range customHandlers {
			exit, r = m.do(w, r, handler)
			if exit {
				return
			}
		}
	})
}

func (m *MWHandler) do(w http.ResponseWriter, r *http.Request, handler Handler) (bool, *http.Request) {
	var resp *Response

	// Record handler runtime
	func() {
		statusCode := "2xx"
		startTime := time.Now()

		if resp = handler(w, r); resp != nil {
			func() {
				// Stop execution if it's passed
				if resp.StopExecution {
					return
				}

				// If a context is returned, we will
				// replace the current request with a new request
				if resp.Context != nil {
					r = r.WithContext(resp.Context)
					return
				}

				// If there's no error but we have a response
				if resp.Err == nil {
					resp.Err = errors.New("Problem with middleware; neither Err or StopExecution is set")
					resp.StatusCode = http.StatusInternalServerError
				}

				// Now assume we have an error.
				if m.Config.Statter != nil && resp.StatusCode >= 500 {
					go m.reportError()
				}

				// Write the error out
				WriteJSONStatus(w, "error", resp.Error(), resp.StatusCode)
			}()

			if resp.StatusCode > 0 {
				statusCode = strconv.Itoa(resp.StatusCode)
			}
		}

		handlerName := getFuncName(handler)

		if m.Config.Statter != nil {
			// Record runtime metric
			go m.reportDuration(handlerName, startTime)

			// Record status code metric (default 2xx)
			go m.reportStatusCode(handlerName, statusCode)
		}

		// If a CustomStatter is set, send the handler metrics to it.
		// This allows the client to handle these metrics however it wants.
		if m.Config.CustomStatter != nil && resp != nil {
			go m.Config.CustomStatter.ReportStats(handlerName, time.Since(startTime), r, resp)
		}
	}()

	// stop executing rest of the
	// handlers if we encounter an error
	if resp != nil && (resp.StopExecution || resp.Err != nil) {
		return true, r
	}

	return false, r
}

func (m *MWHandler) reportError() {
	if m.Config.NoErrStats {
		return
	}

	m.Config.Statter.Inc("errors", 1, m.Config.StatRate)
}

func (m *MWHandler) reportDuration(handlerName string, startTime time.Time) {
	if m.Config.NoDurationStats {
		return
	}

	m.Config.Statter.TimingDuration(
		"handlers."+handlerName+".runtime",
		time.Since(startTime), // delta
		m.Config.StatRate,
	)
}

func (m *MWHandler) reportStatusCode(handlerName string, statusCode string) {
	if m.Config.NoStatusCodeStats {
		return
	}

	m.Config.Statter.Inc(
		"handlers."+handlerName+"."+statusCode,
		1,
		m.Config.StatRate,
	)
}

// WriteJSONStatus is a wrapper for WriteJSONResponse that returns a marshalled JSONStatus blob
func WriteJSONStatus(rw http.ResponseWriter, status, message string, statusCode int) {
	jsonData, _ := json.Marshal(&JSONStatus{
		Message: message,
		Status:  status,
	})

	WriteJSONResponse(rw, statusCode, jsonData)
}

// WriteJSONResponse writes data and status code to the ResponseWriter
func WriteJSONResponse(rw http.ResponseWriter, statusCode int, content []byte) {
	rw.Header().Set("Content-Type", "application/json")
	rw.WriteHeader(statusCode)
	rw.Write(content)
}

// getFuncName uses reflection to determine a given function name
// It returns a string version of the function name (and performs string cleanup)
func getFuncName(i interface{}) string {
	fullName := runtime.FuncForPC(reflect.ValueOf(i).Pointer()).Name()
	ns := strings.Split(fullName, ".")

	// when we get a method (not a raw function) it comes attached to whatever struct is in its
	// method receiver via a function closure, this is not precisely the same as that method itself
	// so the compiler appends "-fm" so the name of the closure does not conflict with the actual function
	// http://grokbase.com/t/gg/golang-nuts/153jyb5b7p/go-nuts-fm-suffix-in-function-name-what-does-it-mean#20150318ssinqqzrmhx2ep45wjkxsa4rua
	return strings.TrimSuffix(ns[len(ns)-1], ")-fm")
}
<|endoftext|>


<commit_before>package main

// ser.go : serialization

import (
	"bytes"
	"fmt"
	"io"
	"time"

	schema "github.com/glycerine/goq/schema"
	"github.com/glycerine/go-capnproto"
)

func (js *JobServ) ServerToCapnp() (bytes.Buffer, *capn.Segment) {
	seg := capn.NewBuffer(nil)
	z := schema.NewRootZ(seg)

	zjs := schema.NewZgoqserver(seg)
	z.SetGoqserver(zjs)

	// nextjobid
	zjs.SetNextjobid(js.NextJobId)

	var i int
	var j *Job

	// runq
	if len(js.RunQ) > 0 {
		//fmt.Printf("len of Runq: %d, %#v\n", len(js.RunQ), js.RunQ)
		runq := schema.NewZjobList(seg, len(js.RunQ))
		plistRunq := capn.PointerList(runq)
		i = 0
		for _, j = range js.RunQ {
			zjob := JobToCapnpSegment(j, seg, false)
			plistRunq.Set(i, capn.Object(zjob))
			i++
		}
		zjs.SetRunq(runq)
	}

	// waitingjobs
	if len(js.WaitingJobs) > 0 {
		//fmt.Printf("len of WaitingJobs: %d, %#v\n", len(js.WaitingJobs), js.WaitingJobs)
		waitingjobs := schema.NewZjobList(seg, len(js.WaitingJobs))
		plistWaitingjobs := capn.PointerList(waitingjobs)
		i = 0
		for _, j = range js.WaitingJobs {
			zjob := JobToCapnpSegment(j, seg, false)
			plistWaitingjobs.Set(i, capn.Object(zjob))
			i++
		}
		zjs.SetWaitingjobs(waitingjobs)
	}
	// counters
	zjs.SetFinishedjobscount(js.FinishedJobsCount)
	zjs.SetBadsgtcount(js.BadSgtCount)
	zjs.SetCancelledjobcount(js.CancelledJobCount)
	zjs.SetBadnoncecount(js.BadNonceCount)

	buf := bytes.Buffer{}
	seg.WriteTo(&buf)

	return buf, seg
}

func (js *JobServ) SetStateFromCapnp(r io.Reader, fn string) {
	capMsg, err := capn.ReadFromStream(r, nil)
	if err != nil {
		panic(fmt.Errorf("capnp problem reading from file '%s': %s", fn, err))
	}

	z := schema.ReadRootZ(capMsg)
	d3 := z.Which()
	if d3 != schema.Z_GOQSERVER {
		panic(fmt.Sprintf("expected schema.Z_GOQSERVER, got %d", d3))
	}

	zjs := z.Goqserver()
	js.NextJobId = zjs.Nextjobid()

	runqlist := zjs.Runq().ToArray()
	for _, zjob := range runqlist {
		j := CapnpZjobToJob(zjob)

		// don't kill jobs prematurely just because we the server
		// just started back up!
		j.Unansweredping = 0
		j.Lastpingtm = time.Now().UnixNano()

		js.RunQ[j.Id] = j
		js.KnownJobHash[j.Id] = j
		js.RegisterWho(j)
	}

	waitlist := zjs.Waitingjobs().ToArray()
	for _, zjob := range waitlist {
		j := CapnpZjobToJob(zjob)
		js.WaitingJobs = append(js.WaitingJobs, j)
		js.KnownJobHash[j.Id] = j
		js.RegisterWho(j)
	}

	js.FinishedJobsCount = zjs.Finishedjobscount()
	js.BadSgtCount = zjs.Badsgtcount()
	js.CancelledJobCount = zjs.Cancelledjobcount()
	js.BadNonceCount = zjs.Badnoncecount()

}

func CapnpZjobToJob(zj schema.Zjob) *Job {

	return &Job{
		Id:       zj.Id(),
		Msg:      zj.Msg(),
		Aboutjid: zj.Aboutjid(),

		Cmd:  zj.Cmd(),
		Args: zj.Args().ToArray(),
		Out:  zj.Out().ToArray(),
		Env:  zj.Env().ToArray(),

		Host: zj.Host(),
		Stm:  zj.Stm(),
		Etm:  zj.Etm(),

		Elapsec: zj.Elapsec(),
		Status:  zj.Status(),
		Subtime: zj.Subtime(),

		Pid: zj.Pid(),
		Dir: zj.Dir(),

		Submitaddr: zj.Submitaddr(),
		Serveraddr: zj.Serveraddr(),
		Workeraddr: zj.Workeraddr(),
		Finishaddr: zj.Finishaddr().ToArray(),

		Signature: zj.Signature(),
		IsLocal:   zj.Islocal(),
		Cancelled: zj.Cancelled(),

		ArrayId:        zj.Arrayid(),
		GroupId:        zj.Groupid(),
		Delegatetm:     zj.Delegatetm(),
		Lastpingtm:     zj.Lastpingtm(),
		Unansweredping: zj.Unansweredping(),
		Sendernonce:    zj.Sendernonce(),
		Sendtime:       zj.Sendtime(),
	}
}

func CapnpToJob(buf *bytes.Buffer) *Job {
	capMsg, err := capn.ReadFromStream(buf, nil)
	if err != nil {
		panic(err)
	}

	z := schema.ReadRootZ(capMsg)
	d3 := z.Which()
	if d3 != schema.Z_JOB {
		panic(fmt.Sprintf("expected schema.Z_JOB, got %d", d3))
	}

	zj := z.Job()
	job := CapnpZjobToJob(zj)

	return job
}

// once you've already allocated a segment
func JobToCapnpSegment(j *Job, seg *capn.Segment, shorten bool) schema.Zjob {

	zjob := schema.NewZjob(seg)

	zjob.SetId(j.Id)

	if shorten {
		return zjob
	}

	zjob.SetAboutjid(j.Aboutjid)
	zjob.SetMsg(j.Msg)

	zjob.SetCmd(j.Cmd)

	if tl := StringSliceToCapnp(j.Out, seg); tl != nil {
		zjob.SetOut(*tl)
	}

	if tl := StringSliceToCapnp(j.Args, seg); tl != nil {
		zjob.SetArgs(*tl)
	}

	if tl := StringSliceToCapnp(j.Env, seg); tl != nil {
		zjob.SetEnv(*tl)
	}

	zjob.SetHost(j.Host)

	zjob.SetStm(int64(j.Stm))
	zjob.SetEtm(int64(j.Etm))
	zjob.SetElapsec(j.Elapsec)

	zjob.SetStatus(j.Status)
	zjob.SetSubtime(int64(j.Subtime))
	zjob.SetPid(j.Pid)

	zjob.SetDir(j.Dir)

	zjob.SetSubmitaddr(j.Submitaddr)
	zjob.SetServeraddr(j.Serveraddr)
	zjob.SetWorkeraddr(j.Workeraddr)

	if tl := StringSliceToCapnp(j.Finishaddr, seg); tl != nil {
		zjob.SetFinishaddr(*tl)
	}

	zjob.SetSignature(j.Signature)
	zjob.SetIslocal(j.IsLocal)
	zjob.SetCancelled(j.Cancelled)

	zjob.SetArrayid(j.ArrayId)
	zjob.SetGroupid(j.GroupId)
	zjob.SetDelegatetm(j.Delegatetm)
	zjob.SetLastpingtm(j.Lastpingtm)
	zjob.SetUnansweredping(j.Unansweredping)
	zjob.SetSendernonce(j.Sendernonce)
	zjob.SetSendtime(j.Sendtime)

	return zjob
}

func JobToCapnp(j *Job) (bytes.Buffer, *capn.Segment) {
	seg := capn.NewBuffer(nil)
	z := schema.NewRootZ(seg)

	zjob := JobToCapnpSegment(j, seg, false)
	z.SetJob(zjob)

	buf := bytes.Buffer{}
	seg.WriteTo(&buf)

	return buf, seg
}

func StringSliceToCapnp(a []string, seg *capn.Segment) *capn.TextList {
	if a == nil {
		panic("string slice can't be nil")
	}
	// looks like we need a size zero textlist for complete initialization,
	// so do this even if len(a) == 0
	tl := seg.NewTextList(len(a))
	for i := range a {
		tl.Set(i, a[i])
	}
	return &tl
}
<commit_msg>remove debug machinery<commit_after>package main

// ser.go : serialization

import (
	"bytes"
	"fmt"
	"io"
	"time"

	schema "github.com/glycerine/goq/schema"
	"github.com/glycerine/go-capnproto"
)

func (js *JobServ) ServerToCapnp() (bytes.Buffer, *capn.Segment) {
	seg := capn.NewBuffer(nil)
	z := schema.NewRootZ(seg)

	zjs := schema.NewZgoqserver(seg)
	z.SetGoqserver(zjs)

	// nextjobid
	zjs.SetNextjobid(js.NextJobId)

	var i int
	var j *Job

	// runq
	if len(js.RunQ) > 0 {
		//fmt.Printf("len of Runq: %d, %#v\n", len(js.RunQ), js.RunQ)
		runq := schema.NewZjobList(seg, len(js.RunQ))
		plistRunq := capn.PointerList(runq)
		i = 0
		for _, j = range js.RunQ {
			zjob := JobToCapnpSegment(j, seg)
			plistRunq.Set(i, capn.Object(zjob))
			i++
		}
		zjs.SetRunq(runq)
	}

	// waitingjobs
	if len(js.WaitingJobs) > 0 {
		//fmt.Printf("len of WaitingJobs: %d, %#v\n", len(js.WaitingJobs), js.WaitingJobs)
		waitingjobs := schema.NewZjobList(seg, len(js.WaitingJobs))
		plistWaitingjobs := capn.PointerList(waitingjobs)
		i = 0
		for _, j = range js.WaitingJobs {
			zjob := JobToCapnpSegment(j, seg)
			plistWaitingjobs.Set(i, capn.Object(zjob))
			i++
		}
		zjs.SetWaitingjobs(waitingjobs)
	}
	// counters
	zjs.SetFinishedjobscount(js.FinishedJobsCount)
	zjs.SetBadsgtcount(js.BadSgtCount)
	zjs.SetCancelledjobcount(js.CancelledJobCount)
	zjs.SetBadnoncecount(js.BadNonceCount)

	buf := bytes.Buffer{}
	seg.WriteTo(&buf)

	return buf, seg
}

func (js *JobServ) SetStateFromCapnp(r io.Reader, fn string) {
	capMsg, err := capn.ReadFromStream(r, nil)
	if err != nil {
		panic(fmt.Errorf("capnp problem reading from file '%s': %s", fn, err))
	}

	z := schema.ReadRootZ(capMsg)
	d3 := z.Which()
	if d3 != schema.Z_GOQSERVER {
		panic(fmt.Sprintf("expected schema.Z_GOQSERVER, got %d", d3))
	}

	zjs := z.Goqserver()
	js.NextJobId = zjs.Nextjobid()

	runqlist := zjs.Runq().ToArray()
	for _, zjob := range runqlist {
		j := CapnpZjobToJob(zjob)

		// don't kill jobs prematurely just because we the server
		// just started back up!
		j.Unansweredping = 0
		j.Lastpingtm = time.Now().UnixNano()

		js.RunQ[j.Id] = j
		js.KnownJobHash[j.Id] = j
		js.RegisterWho(j)
	}

	waitlist := zjs.Waitingjobs().ToArray()
	for _, zjob := range waitlist {
		j := CapnpZjobToJob(zjob)
		js.WaitingJobs = append(js.WaitingJobs, j)
		js.KnownJobHash[j.Id] = j
		js.RegisterWho(j)
	}

	js.FinishedJobsCount = zjs.Finishedjobscount()
	js.BadSgtCount = zjs.Badsgtcount()
	js.CancelledJobCount = zjs.Cancelledjobcount()
	js.BadNonceCount = zjs.Badnoncecount()

}

func CapnpZjobToJob(zj schema.Zjob) *Job {

	return &Job{
		Id:       zj.Id(),
		Msg:      zj.Msg(),
		Aboutjid: zj.Aboutjid(),

		Cmd:  zj.Cmd(),
		Args: zj.Args().ToArray(),
		Out:  zj.Out().ToArray(),
		Env:  zj.Env().ToArray(),

		Host: zj.Host(),
		Stm:  zj.Stm(),
		Etm:  zj.Etm(),

		Elapsec: zj.Elapsec(),
		Status:  zj.Status(),
		Subtime: zj.Subtime(),

		Pid: zj.Pid(),
		Dir: zj.Dir(),

		Submitaddr: zj.Submitaddr(),
		Serveraddr: zj.Serveraddr(),
		Workeraddr: zj.Workeraddr(),
		Finishaddr: zj.Finishaddr().ToArray(),

		Signature: zj.Signature(),
		IsLocal:   zj.Islocal(),
		Cancelled: zj.Cancelled(),

		ArrayId:        zj.Arrayid(),
		GroupId:        zj.Groupid(),
		Delegatetm:     zj.Delegatetm(),
		Lastpingtm:     zj.Lastpingtm(),
		Unansweredping: zj.Unansweredping(),
		Sendernonce:    zj.Sendernonce(),
		Sendtime:       zj.Sendtime(),
	}
}

func CapnpToJob(buf *bytes.Buffer) *Job {
	capMsg, err := capn.ReadFromStream(buf, nil)
	if err != nil {
		panic(err)
	}

	z := schema.ReadRootZ(capMsg)
	d3 := z.Which()
	if d3 != schema.Z_JOB {
		panic(fmt.Sprintf("expected schema.Z_JOB, got %d", d3))
	}

	zj := z.Job()
	job := CapnpZjobToJob(zj)

	return job
}

// once you've already allocated a segment
func JobToCapnpSegment(j *Job, seg *capn.Segment) schema.Zjob {

	zjob := schema.NewZjob(seg)

	zjob.SetId(j.Id)

	zjob.SetAboutjid(j.Aboutjid)
	zjob.SetMsg(j.Msg)

	zjob.SetCmd(j.Cmd)

	if tl := StringSliceToCapnp(j.Out, seg); tl != nil {
		zjob.SetOut(*tl)
	}

	if tl := StringSliceToCapnp(j.Args, seg); tl != nil {
		zjob.SetArgs(*tl)
	}

	if tl := StringSliceToCapnp(j.Env, seg); tl != nil {
		zjob.SetEnv(*tl)
	}

	zjob.SetHost(j.Host)

	zjob.SetStm(int64(j.Stm))
	zjob.SetEtm(int64(j.Etm))
	zjob.SetElapsec(j.Elapsec)

	zjob.SetStatus(j.Status)
	zjob.SetSubtime(int64(j.Subtime))
	zjob.SetPid(j.Pid)

	zjob.SetDir(j.Dir)

	zjob.SetSubmitaddr(j.Submitaddr)
	zjob.SetServeraddr(j.Serveraddr)
	zjob.SetWorkeraddr(j.Workeraddr)

	if tl := StringSliceToCapnp(j.Finishaddr, seg); tl != nil {
		zjob.SetFinishaddr(*tl)
	}

	zjob.SetSignature(j.Signature)
	zjob.SetIslocal(j.IsLocal)
	zjob.SetCancelled(j.Cancelled)

	zjob.SetArrayid(j.ArrayId)
	zjob.SetGroupid(j.GroupId)
	zjob.SetDelegatetm(j.Delegatetm)
	zjob.SetLastpingtm(j.Lastpingtm)
	zjob.SetUnansweredping(j.Unansweredping)
	zjob.SetSendernonce(j.Sendernonce)
	zjob.SetSendtime(j.Sendtime)

	return zjob
}

func JobToCapnp(j *Job) (bytes.Buffer, *capn.Segment) {
	seg := capn.NewBuffer(nil)
	z := schema.NewRootZ(seg)

	zjob := JobToCapnpSegment(j, seg)
	z.SetJob(zjob)

	buf := bytes.Buffer{}
	seg.WriteTo(&buf)

	return buf, seg
}

func StringSliceToCapnp(a []string, seg *capn.Segment) *capn.TextList {
	if a == nil {
		panic("string slice can't be nil")
	}
	// looks like we need a size zero textlist for complete initialization,
	// so do this even if len(a) == 0
	tl := seg.NewTextList(len(a))
	for i := range a {
		tl.Set(i, a[i])
	}
	return &tl
}
<|endoftext|>


<commit_before>package torpedo_currency_plugin

import (
	"fmt"

	"strings"

	common "github.com/tb0hdan/torpedo_common"
	"github.com/tb0hdan/torpedo_registry"
)

const (
	PBAPIBase  = "https://api.privatbank.ua/p24api"
	PBExchRate = PBAPIBase + "/pubinfo?exchange&json&coursid=11"
)

type PBResponse struct {
	Currency     string `json:"ccy"`
	BaseCurrency string `json:"base_ccy"`
	PurchaseRate string `json:"buy"`
	SaleRate     string `json:"sale"`
}

func PBProcessMessage(api *torpedo_registry.BotAPI, channel interface{}, incoming_message string) {
	var prefix string
	var reply string
	cu := &common.Utils{}
	result := make([]PBResponse, 0)
	err := cu.GetURLUnmarshal(PBExchRate, &result)
	if err != nil {
		api.Bot.PostMessage(channel, fmt.Sprint("An error occured: %+v\n", err), api)
		return
	}
	_, command, _ := common.GetRequestedFeature(incoming_message)

	if command == "" {
		reply = fmt.Sprintf("Usage: %spb [", api.CommandPrefix)
		for idx, item := range result {
			if idx == 0 {
				prefix = ""
			} else {
				prefix = "|"
			}
			reply += fmt.Sprintf("%s%s", prefix, item.Currency)
		}
		reply += "]"
	} else {
		for _, item := range result {
			if item.Currency == strings.ToUpper(command) {
				reply = fmt.Sprintf("Buy/Sale: %s/%s UAH", item.PurchaseRate, item.SaleRate)
			}
		}
	}

	api.Bot.PostMessage(channel, reply, api)
	return
}
<commit_msg>fix PB base currency<commit_after>package torpedo_currency_plugin

import (
	"fmt"

	"strings"

	common "github.com/tb0hdan/torpedo_common"
	"github.com/tb0hdan/torpedo_registry"
)

const (
	PBAPIBase  = "https://api.privatbank.ua/p24api"
	PBExchRate = PBAPIBase + "/pubinfo?exchange&json&coursid=11"
)

type PBResponse struct {
	Currency     string `json:"ccy"`
	BaseCurrency string `json:"base_ccy"`
	PurchaseRate string `json:"buy"`
	SaleRate     string `json:"sale"`
}

func PBProcessMessage(api *torpedo_registry.BotAPI, channel interface{}, incoming_message string) {
	var prefix string
	var reply string
	cu := &common.Utils{}
	result := make([]PBResponse, 0)
	err := cu.GetURLUnmarshal(PBExchRate, &result)
	if err != nil {
		api.Bot.PostMessage(channel, fmt.Sprint("An error occured: %+v\n", err), api)
		return
	}
	_, command, _ := common.GetRequestedFeature(incoming_message)

	if command == "" {
		reply = fmt.Sprintf("Usage: %spb [", api.CommandPrefix)
		for idx, item := range result {
			if idx == 0 {
				prefix = ""
			} else {
				prefix = "|"
			}
			reply += fmt.Sprintf("%s%s", prefix, item.Currency)
		}
		reply += "]"
	} else {
		for _, item := range result {
			if item.Currency == strings.ToUpper(command) {
				reply = fmt.Sprintf("Buy/Sale: %s/%s %s\n", item.PurchaseRate, item.SaleRate, item.BaseCurrency)
			}
		}
	}

	api.Bot.PostMessage(channel, reply, api)
	return
}
<|endoftext|>


<commit_before>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package engine

import (
	"fmt"
	"net/url"

	"github.com/keybase/client/go/libkb"
)

type Prereqs struct {
	Session bool
	Device  bool
}

type Engine interface {
	Run(ctx *Context) error
	Prereqs() Prereqs
	libkb.UIConsumer
	G() *libkb.GlobalContext
}

func runPrereqs(e Engine, ctx *Context) (err error) {
	prq := e.Prereqs()

	if prq.Session {
		var ok bool
		ok, err = IsLoggedIn(e, ctx)
		if !ok {
			urlError, isURLError := err.(*url.Error)
			context := ""
			if isURLError {
				context = fmt.Sprintf("Encountered a network error: %s", urlError.Err)
			}
			err = libkb.LoginRequiredError{Context: context}
		}
		if err != nil {
			return err
		}
	}

	if prq.Device {
		var ok bool
		ok, err = IsProvisioned(e, ctx)
		if err != nil {
			return err
		}
		if !ok {
			err = libkb.DeviceRequiredError{}
			return err
		}
	}

	return

}

func RunEngine(e Engine, ctx *Context) (err error) {
	e.G().Log.Debug("+ RunEngine(%s)", e.Name())
	defer func() { e.G().Log.Debug("- RunEngine(%s) -> %s", e.Name(), libkb.ErrToOk(err)) }()

	// TODO: check secret ui

	if err = check(e, ctx); err != nil {
		return err
	}
	if err = runPrereqs(e, ctx); err != nil {
		return err
	}
	err = e.Run(ctx)
	return err
}

func check(c libkb.UIConsumer, ctx *Context) error {
	if err := checkUI(c, ctx); err != nil {
		return CheckError{fmt.Sprintf("%s: %s", c.Name(), err)}
	}

	for _, sub := range c.SubConsumers() {
		if err := check(sub, ctx); err != nil {
			return CheckError{fmt.Sprintf("%s: %s", sub.Name(), err)}
		}
	}

	return nil
}

func checkUI(c libkb.UIConsumer, ctx *Context) error {
	for _, ui := range c.RequiredUIs() {
		if !ctx.HasUI(ui) {
			return CheckError{fmt.Sprintf("requires ui %q", ui)}
		}
	}
	return nil
}
<commit_msg>Add check in RunEngine for delegate ui<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package engine

import (
	"fmt"
	"net/url"

	"github.com/keybase/client/go/libkb"
)

type Prereqs struct {
	Session bool
	Device  bool
}

type Engine interface {
	Run(ctx *Context) error
	Prereqs() Prereqs
	libkb.UIConsumer
	G() *libkb.GlobalContext
}

func runPrereqs(e Engine, ctx *Context) (err error) {
	prq := e.Prereqs()

	if prq.Session {
		var ok bool
		ok, err = IsLoggedIn(e, ctx)
		if !ok {
			urlError, isURLError := err.(*url.Error)
			context := ""
			if isURLError {
				context = fmt.Sprintf("Encountered a network error: %s", urlError.Err)
			}
			err = libkb.LoginRequiredError{Context: context}
		}
		if err != nil {
			return err
		}
	}

	if prq.Device {
		var ok bool
		ok, err = IsProvisioned(e, ctx)
		if err != nil {
			return err
		}
		if !ok {
			err = libkb.DeviceRequiredError{}
			return err
		}
	}

	return

}

func RunEngine(e Engine, ctx *Context) (err error) {
	e.G().Log.Debug("+ RunEngine(%s)", e.Name())
	defer func() { e.G().Log.Debug("- RunEngine(%s) -> %s", e.Name(), libkb.ErrToOk(err)) }()

	if err = delegateUIs(e, ctx); err != nil {
		return err
	}
	if err = check(e, ctx); err != nil {
		return err
	}
	if err = runPrereqs(e, ctx); err != nil {
		return err
	}

	err = e.Run(ctx)
	return err
}

func delegateUIs(e Engine, ctx *Context) error {
	if e.G().UIRouter == nil {
		return nil
	}

	// currently, only doing this for SecretUI, but in future,
	// perhaps should iterate over all registered UIs in UIRouter.

	if !requiresUI(e, ctx, libkb.SecretUIKind) {
		return nil
	}

	ui, err := e.G().UIRouter.GetSecretUI()
	if err != nil {
		return err
	}

	if ui != nil {
		e.G().Log.Debug("using delegated secret UI for engine %q", e.Name())
		ctx.SecretUI = ui
	}

	return nil
}

func check(c libkb.UIConsumer, ctx *Context) error {
	if err := checkUI(c, ctx); err != nil {
		return CheckError{fmt.Sprintf("%s: %s", c.Name(), err)}
	}

	for _, sub := range c.SubConsumers() {
		if err := check(sub, ctx); err != nil {
			return CheckError{fmt.Sprintf("%s: %s", sub.Name(), err)}
		}
	}

	return nil
}

func checkUI(c libkb.UIConsumer, ctx *Context) error {
	for _, ui := range c.RequiredUIs() {
		if !ctx.HasUI(ui) {
			return CheckError{fmt.Sprintf("requires ui %q", ui)}
		}
	}
	return nil
}

func requiresUI(c libkb.UIConsumer, ctx *Context, kind libkb.UIKind) bool {
	for _, ui := range c.RequiredUIs() {
		if ui == kind {
			return true
		}
	}

	for _, sub := range c.SubConsumers() {
		if requiresUI(sub, ctx, kind) {
			return true
		}
	}

	return false
}
<|endoftext|>


<commit_before>package ledger

import (
	"errors"
	"strconv"
	"strings"
)

type Entry struct {
	Date        string // "Y-m-d"
	Description string
	Change      int // in cents
}

func isValidCurrency(currency string) bool {
	return currency == "USD" || currency == "EUR"
}

func FormatLedger(currency string, locale string, entries []Entry) (string, error) {
	if !isValidCurrency(currency) {
		return "", errors.New("invalid currency")
	}
	var entriesCopy []Entry
	entriesCopy = append(entriesCopy, entries...)

	m1 := map[bool]int{true: 0, false: 1}
	m2 := map[bool]int{true: -1, false: 1}
	es := entriesCopy
	for len(es) > 1 {
		first, rest := es[0], es[1:]
		success := false
		for !success {
			success = true
			for i, e := range rest {
				if (m1[e.Date == first.Date]*m2[e.Date < first.Date]*4 +
					m1[e.Description == first.Description]*m2[e.Description < first.Description]*2 +
					m1[e.Change == first.Change]*m2[e.Change < first.Change]*1) < 0 {
					es[0], es[i+1] = es[i+1], es[0]
					success = false
				}
			}
		}
		es = es[1:]
	}

	var s string
	if locale == "nl-NL" {
		s = "Datum" +
			strings.Repeat(" ", 10-len("Datum")) +
			" | " +
			"Omschrijving" +
			strings.Repeat(" ", 25-len("Omschrijving")) +
			" | " + "Verandering" + "\n"
	} else if locale == "en-US" {
		s = "Date" +
			strings.Repeat(" ", 10-len("Date")) +
			" | " +
			"Description" +
			strings.Repeat(" ", 25-len("Description")) +
			" | " + "Change" + "\n"
	} else {
		return "", errors.New("")
	}
	// Parallelism, always a great idea
	co := make(chan struct {
		i int
		s string
		e error
	})
	for i, et := range entriesCopy {
		go func(i int, entry Entry) {
			if len(entry.Date) != 10 {
				co <- struct {
					i int
					s string
					e error
				}{e: errors.New("")}
			}
			d1, d2, d3, d4, d5 := entry.Date[0:4], entry.Date[4], entry.Date[5:7], entry.Date[7], entry.Date[8:10]
			if d2 != '-' {
				co <- struct {
					i int
					s string
					e error
				}{e: errors.New("")}
			}
			if d4 != '-' {
				co <- struct {
					i int
					s string
					e error
				}{e: errors.New("")}
			}
			de := entry.Description
			if len(de) > 25 {
				de = de[:22] + "..."
			} else {
				de = de + strings.Repeat(" ", 25-len(de))
			}
			var d string
			if locale == "nl-NL" {
				d = d5 + "-" + d3 + "-" + d1
			} else if locale == "en-US" {
				d = d3 + "/" + d5 + "/" + d1
			}
			negative := false
			cents := entry.Change
			if cents < 0 {
				cents = cents * -1
				negative = true
			}
			var a string
			if locale == "nl-NL" {
				if currency == "EUR" {
					a += "€"
				} else if currency == "USD" {
					a += "$"
				} else {
					co <- struct {
						i int
						s string
						e error
					}{e: errors.New("")}
				}
				a += " "
				centsStr := strconv.Itoa(cents)
				switch len(centsStr) {
				case 1:
					centsStr = "00" + centsStr
				case 2:
					centsStr = "0" + centsStr
				}
				rest := centsStr[:len(centsStr)-2]
				var parts []string
				for len(rest) > 3 {
					parts = append(parts, rest[len(rest)-3:])
					rest = rest[:len(rest)-3]
				}
				if len(rest) > 0 {
					parts = append(parts, rest)
				}
				for i := len(parts) - 1; i >= 0; i-- {
					a += parts[i] + "."
				}
				a = a[:len(a)-1]
				a += ","
				a += centsStr[len(centsStr)-2:]
				if negative {
					a += "-"
				} else {
					a += " "
				}
			} else if locale == "en-US" {
				if negative {
					a += "("
				}
				if currency == "EUR" {
					a += "€"
				} else if currency == "USD" {
					a += "$"
				} else {
					co <- struct {
						i int
						s string
						e error
					}{e: errors.New("")}
				}
				centsStr := strconv.Itoa(cents)
				switch len(centsStr) {
				case 1:
					centsStr = "00" + centsStr
				case 2:
					centsStr = "0" + centsStr
				}
				rest := centsStr[:len(centsStr)-2]
				var parts []string
				for len(rest) > 3 {
					parts = append(parts, rest[len(rest)-3:])
					rest = rest[:len(rest)-3]
				}
				if len(rest) > 0 {
					parts = append(parts, rest)
				}
				for i := len(parts) - 1; i >= 0; i-- {
					a += parts[i] + ","
				}
				a = a[:len(a)-1]
				a += "."
				a += centsStr[len(centsStr)-2:]
				if negative {
					a += ")"
				} else {
					a += " "
				}
			} else {
				co <- struct {
					i int
					s string
					e error
				}{e: errors.New("")}
			}
			var al int
			for range a {
				al++
			}
			co <- struct {
				i int
				s string
				e error
			}{i: i, s: d + strings.Repeat(" ", 10-len(d)) + " | " + de + " | " +
				strings.Repeat(" ", 13-al) + a + "\n"}
		}(i, et)
	}
	ss := make([]string, len(entriesCopy))
	for range entriesCopy {
		v := <-co
		if v.e != nil {
			return "", v.e
		}
		ss[v.i] = v.s
	}
	for i := 0; i < len(entriesCopy); i++ {
		s += ss[i]
	}
	return s, nil
}
<commit_msg>Extract isValidLocale<commit_after>package ledger

import (
	"errors"
	"strconv"
	"strings"
)

type Entry struct {
	Date        string // "Y-m-d"
	Description string
	Change      int // in cents
}

func isValidCurrency(currency string) bool {
	return currency == "USD" || currency == "EUR"
}

func isValidLocale(locale string) bool {
	return locale == "en-US" || locale == "nl-NL"
}

func FormatLedger(currency string, locale string, entries []Entry) (string, error) {
	if !isValidCurrency(currency) {
		return "", errors.New("invalid currency")
	}

	if !isValidLocale(locale) {
		return "", errors.New("invalid locale")
	}

	entriesCopy := append([]Entry{}, entries...)

	m1 := map[bool]int{true: 0, false: 1}
	m2 := map[bool]int{true: -1, false: 1}
	es := entriesCopy
	for len(es) > 1 {
		first, rest := es[0], es[1:]
		success := false
		for !success {
			success = true
			for i, e := range rest {
				if (m1[e.Date == first.Date]*m2[e.Date < first.Date]*4 +
					m1[e.Description == first.Description]*m2[e.Description < first.Description]*2 +
					m1[e.Change == first.Change]*m2[e.Change < first.Change]*1) < 0 {
					es[0], es[i+1] = es[i+1], es[0]
					success = false
				}
			}
		}
		es = es[1:]
	}

	var s string
	if locale == "nl-NL" {
		s = "Datum" +
			strings.Repeat(" ", 10-len("Datum")) +
			" | " +
			"Omschrijving" +
			strings.Repeat(" ", 25-len("Omschrijving")) +
			" | " + "Verandering" + "\n"
	} else if locale == "en-US" {
		s = "Date" +
			strings.Repeat(" ", 10-len("Date")) +
			" | " +
			"Description" +
			strings.Repeat(" ", 25-len("Description")) +
			" | " + "Change" + "\n"
	}
	// Parallelism, always a great idea
	co := make(chan struct {
		i int
		s string
		e error
	})
	for i, et := range entriesCopy {
		go func(i int, entry Entry) {
			if len(entry.Date) != 10 {
				co <- struct {
					i int
					s string
					e error
				}{e: errors.New("")}
			}
			d1, d2, d3, d4, d5 := entry.Date[0:4], entry.Date[4], entry.Date[5:7], entry.Date[7], entry.Date[8:10]
			if d2 != '-' {
				co <- struct {
					i int
					s string
					e error
				}{e: errors.New("")}
			}
			if d4 != '-' {
				co <- struct {
					i int
					s string
					e error
				}{e: errors.New("")}
			}
			de := entry.Description
			if len(de) > 25 {
				de = de[:22] + "..."
			} else {
				de = de + strings.Repeat(" ", 25-len(de))
			}
			var d string
			if locale == "nl-NL" {
				d = d5 + "-" + d3 + "-" + d1
			} else if locale == "en-US" {
				d = d3 + "/" + d5 + "/" + d1
			}
			negative := false
			cents := entry.Change
			if cents < 0 {
				cents = cents * -1
				negative = true
			}
			var a string
			if locale == "nl-NL" {
				if currency == "EUR" {
					a += "€"
				} else if currency == "USD" {
					a += "$"
				} else {
					co <- struct {
						i int
						s string
						e error
					}{e: errors.New("")}
				}
				a += " "
				centsStr := strconv.Itoa(cents)
				switch len(centsStr) {
				case 1:
					centsStr = "00" + centsStr
				case 2:
					centsStr = "0" + centsStr
				}
				rest := centsStr[:len(centsStr)-2]
				var parts []string
				for len(rest) > 3 {
					parts = append(parts, rest[len(rest)-3:])
					rest = rest[:len(rest)-3]
				}
				if len(rest) > 0 {
					parts = append(parts, rest)
				}
				for i := len(parts) - 1; i >= 0; i-- {
					a += parts[i] + "."
				}
				a = a[:len(a)-1]
				a += ","
				a += centsStr[len(centsStr)-2:]
				if negative {
					a += "-"
				} else {
					a += " "
				}
			} else if locale == "en-US" {
				if negative {
					a += "("
				}
				if currency == "EUR" {
					a += "€"
				} else if currency == "USD" {
					a += "$"
				} else {
					co <- struct {
						i int
						s string
						e error
					}{e: errors.New("")}
				}
				centsStr := strconv.Itoa(cents)
				switch len(centsStr) {
				case 1:
					centsStr = "00" + centsStr
				case 2:
					centsStr = "0" + centsStr
				}
				rest := centsStr[:len(centsStr)-2]
				var parts []string
				for len(rest) > 3 {
					parts = append(parts, rest[len(rest)-3:])
					rest = rest[:len(rest)-3]
				}
				if len(rest) > 0 {
					parts = append(parts, rest)
				}
				for i := len(parts) - 1; i >= 0; i-- {
					a += parts[i] + ","
				}
				a = a[:len(a)-1]
				a += "."
				a += centsStr[len(centsStr)-2:]
				if negative {
					a += ")"
				} else {
					a += " "
				}
			} else {
				co <- struct {
					i int
					s string
					e error
				}{e: errors.New("")}
			}
			var al int
			for range a {
				al++
			}
			co <- struct {
				i int
				s string
				e error
			}{i: i, s: d + strings.Repeat(" ", 10-len(d)) + " | " + de + " | " +
				strings.Repeat(" ", 13-al) + a + "\n"}
		}(i, et)
	}
	ss := make([]string, len(entriesCopy))
	for range entriesCopy {
		v := <-co
		if v.e != nil {
			return "", v.e
		}
		ss[v.i] = v.s
	}
	for i := 0; i < len(entriesCopy); i++ {
		s += ss[i]
	}
	return s, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Vitess Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package pools

import (
	"fmt"
	"sync"
)

// IDPool is used to ensure that the set of IDs in use concurrently never
// contains any duplicates. The IDs start at 1 and increase without bound, but
// will never be larger than the peak number of concurrent uses.
//
// IDPool's Get() and Set() methods can be used concurrently.
type IDPool struct {
	sync.Mutex

	// used holds the set of values that have been returned to us with Put().
	used map[uint32]bool
	// maxUsed remembers the largest value we've given out.
	maxUsed uint32
}

// NewIDPool creates and initializes an IDPool.
func NewIDPool(initValue uint32) *IDPool {
	return &IDPool{
		used:    make(map[uint32]bool),
		maxUsed: initValue,
	}
}

// Get returns an ID that is unique among currently active users of this pool.
func (pool *IDPool) Get() (id uint32) {
	pool.Lock()
	defer pool.Unlock()

	// Pick a value that's been returned, if any.
	for key := range pool.used {
		delete(pool.used, key)
		return key
	}

	// No recycled IDs are available, so increase the pool size.
	pool.maxUsed++
	return pool.maxUsed
}

// Put recycles an ID back into the pool for others to use. Putting back a value
// or 0, or a value that is not currently "checked out", will result in a panic
// because that should never happen except in the case of a programming error.
func (pool *IDPool) Put(id uint32) {
	pool.Lock()
	defer pool.Unlock()

	if id < 1 || id > pool.maxUsed {
		panic(fmt.Errorf("IDPool.Put(%v): invalid value, must be in the range [1,%v]", id, pool.maxUsed))
	}

	if pool.used[id] {
		panic(fmt.Errorf("IDPool.Put(%v): can't put value that was already recycled", id))
	}

	// If we're recycling maxUsed, just shrink the pool.
	if id == pool.maxUsed {
		pool.maxUsed = id - 1
		return
	}

	// Add it to the set of recycled IDs.
	pool.used[id] = true
}
<commit_msg>doc: fix comment for IDPool<commit_after>/*
Copyright 2019 The Vitess Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package pools

import (
	"fmt"
	"sync"
)

// IDPool is used to ensure that the set of IDs in use concurrently never
// contains any duplicates. The IDs start at 1 and increase without bound, but
// will never be larger than the peak number of concurrent uses.
//
// IDPool's Get() and Put() methods can be used concurrently.
type IDPool struct {
	sync.Mutex

	// used holds the set of values that have been returned to us with Put().
	used map[uint32]bool
	// maxUsed remembers the largest value we've given out.
	maxUsed uint32
}

// NewIDPool creates and initializes an IDPool.
func NewIDPool(initValue uint32) *IDPool {
	return &IDPool{
		used:    make(map[uint32]bool),
		maxUsed: initValue,
	}
}

// Get returns an ID that is unique among currently active users of this pool.
func (pool *IDPool) Get() (id uint32) {
	pool.Lock()
	defer pool.Unlock()

	// Pick a value that's been returned, if any.
	for key := range pool.used {
		delete(pool.used, key)
		return key
	}

	// No recycled IDs are available, so increase the pool size.
	pool.maxUsed++
	return pool.maxUsed
}

// Put recycles an ID back into the pool for others to use. Putting back a value
// or 0, or a value that is not currently "checked out", will result in a panic
// because that should never happen except in the case of a programming error.
func (pool *IDPool) Put(id uint32) {
	pool.Lock()
	defer pool.Unlock()

	if id < 1 || id > pool.maxUsed {
		panic(fmt.Errorf("IDPool.Put(%v): invalid value, must be in the range [1,%v]", id, pool.maxUsed))
	}

	if pool.used[id] {
		panic(fmt.Errorf("IDPool.Put(%v): can't put value that was already recycled", id))
	}

	// If we're recycling maxUsed, just shrink the pool.
	if id == pool.maxUsed {
		pool.maxUsed = id - 1
		return
	}

	// Add it to the set of recycled IDs.
	pool.used[id] = true
}
<|endoftext|>


<commit_before>package goparsetime

import (
	"fmt"
	"testing"
	"time"
)

func TestParsetime(t *testing.T) {

	f := "15:04:05 January 02 2006"
	testDate := time.Now().Truncate(time.Second)
	parseDate, err := Parsetime(testDate.Format(f))
	if err != nil {
		t.Fatal(err)
	}

	if testDate != parseDate {
		t.Errorf("Actual: %#v, expected: %#v\n", parseDate, testDate)
	}

	parseDate, err = Parsetime("18:00 UTC May 5 2003")
	if err != nil {
		t.Fatal(err)
	}
	if parseDate.Unix() != 1052157600 {
		t.Errorf("Actual %d, expected %d", parseDate.Unix(), 1052157600)
	}

	local, _ := time.LoadLocation("Local")
	testDate = time.Date(2015, 12, 25, 10, 0, 0, 0, local)
	parseDate, err = Parsetime(testDate.Format(f))
	if err != nil {
		t.Fatal(err)
	}
	if testDate != parseDate {
		t.Errorf("Actual: %#v, expected: %#v\n", parseDate, testDate)
	}

	// Date addition
	dateAdd := fmt.Sprintf("%s + 5 days", testDate.Format(f))
	parseDate, err = Parsetime(dateAdd)
	if err != nil {
		t.Fatal(err)
	}
	expectedDate := testDate.Add(time.Duration(120 * time.Hour))
	if expectedDate != parseDate {
		t.Errorf("Actual: %#v, expected: %#v\n", parseDate, expectedDate)
	}
}
<commit_msg>Check for error for basic case<commit_after>package goparsetime

import (
	"fmt"
	"testing"
	"time"
)

func TestParsetime(t *testing.T) {

	f := "15:04:05 January 02 2006"
	testDate := time.Now().Truncate(time.Second)
	parseDate, err := Parsetime(testDate.Format(f))
	if err != nil {
		t.Fatal(err)
	}

	if testDate != parseDate {
		t.Errorf("Actual: %#v, expected: %#v\n", parseDate, testDate)
	}

	parseDate, err = Parsetime("18:00 UTC May 5 2003")
	if err != nil {
		t.Fatal(err)
	}
	if parseDate.Unix() != 1052157600 {
		t.Errorf("Actual %d, expected %d", parseDate.Unix(), 1052157600)
	}

	local, _ := time.LoadLocation("Local")
	testDate = time.Date(2015, 12, 25, 10, 0, 0, 0, local)
	parseDate, err = Parsetime(testDate.Format(f))
	if err != nil {
		t.Fatal(err)
	}
	if testDate != parseDate {
		t.Errorf("Actual: %#v, expected: %#v\n", parseDate, testDate)
	}

	// Date addition
	dateAdd := fmt.Sprintf("%s + 5 days", testDate.Format(f))
	parseDate, err = Parsetime(dateAdd)
	if err != nil {
		t.Fatal(err)
	}
	expectedDate := testDate.Add(time.Duration(120 * time.Hour))
	if expectedDate != parseDate {
		t.Errorf("Actual: %#v, expected: %#v\n", parseDate, expectedDate)
	}

	// Invalid date should return an error
	parseDate, err = Parsetime("invalid string")
	if err == nil {
		t.Fatal("Expected error")
	}
}
<|endoftext|>


<commit_before>package libfuse

import (
	"os"
	"syscall"

	"bazil.org/fuse"
	"bazil.org/fuse/fs"
	"github.com/keybase/kbfs/libkbfs"
	"golang.org/x/net/context"
)

// Symlink represents KBFS symlinks.
type Symlink struct {
	// The directory this symlink is in. This should be safe to store
	// here, without fear of Renames etc making it stale, because we
	// never persist a Symlink into Folder.nodes; it has no
	// libkbfs.Node, so that's impossible. This should make FUSE
	// Lookup etc always get new nodes, limiting the lifetime of a
	// single Symlink value.
	parent *Dir
	name   string
}

var _ fs.Node = (*Symlink)(nil)

// Attr implements the fs.Node interface for Symlink
func (s *Symlink) Attr(ctx context.Context, a *fuse.Attr) (err error) {
	ctx = NewContextWithOpID(ctx, s.parent.folder.fs.log)
	s.parent.folder.fs.log.CDebugf(ctx, "Symlink Attr")
	defer func() { s.parent.folder.fs.reportErr(ctx, err) }()

	_, de, err := s.parent.folder.fs.config.KBFSOps().Lookup(ctx, s.parent.node, s.name)
	if err != nil {
		return err
	}

	fillAttr(&de, a)
	a.Mode = os.ModeSymlink | 0777
	return nil
}

var _ fs.NodeReadlinker = (*Symlink)(nil)

// Readlink implements the fs.NodeReadlinker interface for Symlink
func (s *Symlink) Readlink(ctx context.Context, req *fuse.ReadlinkRequest) (link string, err error) {
	ctx = NewContextWithOpID(ctx, s.parent.folder.fs.log)
	s.parent.folder.fs.log.CDebugf(ctx, "Symlink Readlink")
	defer func() { s.parent.folder.fs.reportErr(ctx, err) }()

	_, de, err := s.parent.folder.fs.config.KBFSOps().Lookup(ctx, s.parent.node, s.name)
	if err != nil {
		return "", err
	}

	if de.Type != libkbfs.Sym {
		return "", fuse.Errno(syscall.EINVAL)
	}
	return de.SymPath, nil
}
<commit_msg>FUSE: Symlink.Attr returns ESTALE when appropriate<commit_after>package libfuse

import (
	"os"
	"syscall"

	"bazil.org/fuse"
	"bazil.org/fuse/fs"
	"github.com/keybase/kbfs/libkbfs"
	"golang.org/x/net/context"
)

// Symlink represents KBFS symlinks.
type Symlink struct {
	// The directory this symlink is in. This should be safe to store
	// here, without fear of Renames etc making it stale, because we
	// never persist a Symlink into Folder.nodes; it has no
	// libkbfs.Node, so that's impossible. This should make FUSE
	// Lookup etc always get new nodes, limiting the lifetime of a
	// single Symlink value.
	parent *Dir
	name   string
}

var _ fs.Node = (*Symlink)(nil)

// Attr implements the fs.Node interface for Symlink
func (s *Symlink) Attr(ctx context.Context, a *fuse.Attr) (err error) {
	ctx = NewContextWithOpID(ctx, s.parent.folder.fs.log)
	s.parent.folder.fs.log.CDebugf(ctx, "Symlink Attr")
	defer func() { s.parent.folder.fs.reportErr(ctx, err) }()

	_, de, err := s.parent.folder.fs.config.KBFSOps().Lookup(ctx, s.parent.node, s.name)
	if err != nil {
		if _, ok := err.(libkbfs.NoSuchNameError); ok {
			return fuse.ESTALE
		}
		return err
	}

	fillAttr(&de, a)
	a.Mode = os.ModeSymlink | 0777
	return nil
}

var _ fs.NodeReadlinker = (*Symlink)(nil)

// Readlink implements the fs.NodeReadlinker interface for Symlink
func (s *Symlink) Readlink(ctx context.Context, req *fuse.ReadlinkRequest) (link string, err error) {
	ctx = NewContextWithOpID(ctx, s.parent.folder.fs.log)
	s.parent.folder.fs.log.CDebugf(ctx, "Symlink Readlink")
	defer func() { s.parent.folder.fs.reportErr(ctx, err) }()

	_, de, err := s.parent.folder.fs.config.KBFSOps().Lookup(ctx, s.parent.node, s.name)
	if err != nil {
		return "", err
	}

	if de.Type != libkbfs.Sym {
		return "", fuse.Errno(syscall.EINVAL)
	}
	return de.SymPath, nil
}
<|endoftext|>


<commit_before>// Code generated by pluginator on HelmChartInflationGenerator; DO NOT EDIT.
// pluginator {unknown  1970-01-01T00:00:00Z  }

package builtins

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/imdario/mergo"
	"github.com/pkg/errors"
	"sigs.k8s.io/kustomize/api/resmap"
	"sigs.k8s.io/kustomize/api/types"
	"sigs.k8s.io/yaml"
)

// HelmChartInflationGeneratorPlugin is a plugin to generate resources
// from a remote or local helm chart.
type HelmChartInflationGeneratorPlugin struct {
	h *resmap.PluginHelpers
	types.HelmGlobals
	types.HelmChart
	tmpDir string
}

var KustomizePlugin HelmChartInflationGeneratorPlugin

const (
	valuesMergeOptionMerge    = "merge"
	valuesMergeOptionOverride = "override"
	valuesMergeOptionReplace  = "replace"
)

var legalMergeOptions = []string{
	valuesMergeOptionMerge,
	valuesMergeOptionOverride,
	valuesMergeOptionReplace,
}

// Config uses the input plugin configurations `config` to setup the generator
// options
func (p *HelmChartInflationGeneratorPlugin) Config(
	h *resmap.PluginHelpers, config []byte) (err error) {
	if h.GeneralConfig() == nil {
		return fmt.Errorf("unable to access general config")
	}
	if !h.GeneralConfig().HelmConfig.Enabled {
		return fmt.Errorf("must specify --enable-helm")
	}
	if h.GeneralConfig().HelmConfig.Command == "" {
		return fmt.Errorf("must specify --helm-command")
	}
	p.h = h
	if err = yaml.Unmarshal(config, p); err != nil {
		return
	}
	return p.validateArgs()
}

// This uses the real file system since tmpDir may be used
// by the helm subprocess.  Cannot use a chroot jail or fake
// filesystem since we allow the user to use previously
// downloaded charts.  This is safe since this plugin is
// owned by kustomize.
func (p *HelmChartInflationGeneratorPlugin) establishTmpDir() (err error) {
	if p.tmpDir != "" {
		// already done.
		return nil
	}
	p.tmpDir, err = ioutil.TempDir("", "kustomize-helm-")
	return err
}

func (p *HelmChartInflationGeneratorPlugin) validateArgs() (err error) {
	if p.Name == "" {
		return fmt.Errorf("chart name cannot be empty")
	}

	// ChartHome might be consulted by the plugin (to read
	// values files below it), so it must be located under
	// the loader root (unless root restrictions are
	// disabled, in which case this can be an absolute path).
	if p.ChartHome == "" {
		p.ChartHome = "charts"
	}

	// The ValuesFile may be consulted by the plugin, so it must
	// be under the loader root (unless root restrictions are
	// disabled).
	if p.ValuesFile == "" {
		p.ValuesFile = filepath.Join(p.ChartHome, p.Name, "values.yaml")
	}

	if err = p.errIfIllegalValuesMerge(); err != nil {
		return err
	}

	// ConfigHome is not loaded by the plugin, and can be located anywhere.
	if p.ConfigHome == "" {
		if err = p.establishTmpDir(); err != nil {
			return errors.Wrap(
				err, "unable to create tmp dir for HELM_CONFIG_HOME")
		}
		p.ConfigHome = filepath.Join(p.tmpDir, "helm")
	}
	return nil
}

func (p *HelmChartInflationGeneratorPlugin) errIfIllegalValuesMerge() error {
	if p.ValuesMerge == "" {
		// Use the default.
		p.ValuesMerge = valuesMergeOptionOverride
		return nil
	}
	for _, opt := range legalMergeOptions {
		if p.ValuesMerge == opt {
			return nil
		}
	}
	return fmt.Errorf("valuesMerge must be one of %v", legalMergeOptions)
}

func (p *HelmChartInflationGeneratorPlugin) absChartHome() string {
	if filepath.IsAbs(p.ChartHome) {
		return p.ChartHome
	}
	return filepath.Join(p.h.Loader().Root(), p.ChartHome)
}

func (p *HelmChartInflationGeneratorPlugin) runHelmCommand(
	args []string) ([]byte, error) {
	stdout := new(bytes.Buffer)
	stderr := new(bytes.Buffer)
	cmd := exec.Command(p.h.GeneralConfig().HelmConfig.Command, args...)
	cmd.Stdout = stdout
	cmd.Stderr = stderr
	env := []string{
		fmt.Sprintf("HELM_CONFIG_HOME=%s", p.ConfigHome),
		fmt.Sprintf("HELM_CACHE_HOME=%s/.cache", p.ConfigHome),
		fmt.Sprintf("HELM_DATA_HOME=%s/.data", p.ConfigHome)}
	cmd.Env = append(os.Environ(), env...)
	err := cmd.Run()
	if err != nil {
		helm := p.h.GeneralConfig().HelmConfig.Command
		err = errors.Wrap(
			fmt.Errorf(
				"unable to run: '%s %s' with env=%s (is '%s' installed?)",
				helm, strings.Join(args, " "), env, helm),
			stderr.String(),
		)
	}
	return stdout.Bytes(), err
}

// createNewMergedValuesFile replaces/merges original values file with ValuesInline.
func (p *HelmChartInflationGeneratorPlugin) createNewMergedValuesFile() (
	path string, err error) {
	if p.ValuesMerge == valuesMergeOptionMerge ||
		p.ValuesMerge == valuesMergeOptionOverride {
		if err = p.replaceValuesInline(); err != nil {
			return "", err
		}
	}
	var b []byte
	b, err = yaml.Marshal(p.ValuesInline)
	if err != nil {
		return "", err
	}
	return p.writeValuesBytes(b)
}

func (p *HelmChartInflationGeneratorPlugin) replaceValuesInline() error {
	pValues, err := p.h.Loader().Load(p.ValuesFile)
	if err != nil {
		return err
	}
	chValues := make(map[string]interface{})
	if err = yaml.Unmarshal(pValues, &chValues); err != nil {
		return err
	}
	switch p.ValuesMerge {
	case valuesMergeOptionOverride:
		err = mergo.Merge(
			&chValues, p.ValuesInline, mergo.WithOverride)
	case valuesMergeOptionMerge:
		err = mergo.Merge(&chValues, p.ValuesInline)
	}
	p.ValuesInline = chValues
	return err
}

// copyValuesFile to avoid branching.  TODO: get rid of this.
func (p *HelmChartInflationGeneratorPlugin) copyValuesFile() (string, error) {
	b, err := p.h.Loader().Load(p.ValuesFile)
	if err != nil {
		return "", err
	}
	return p.writeValuesBytes(b)
}

// Write a absolute path file in the tmp file system.
func (p *HelmChartInflationGeneratorPlugin) writeValuesBytes(
	b []byte) (string, error) {
	if err := p.establishTmpDir(); err != nil {
		return "", fmt.Errorf("cannot create tmp dir to write helm values")
	}
	path := filepath.Join(p.tmpDir, p.Name+"-kustomize-values.yaml")
	return path, ioutil.WriteFile(path, b, 0644)
}

func (p *HelmChartInflationGeneratorPlugin) cleanup() {
	if p.tmpDir != "" {
		os.RemoveAll(p.tmpDir)
	}
}

// Generate implements generator
func (p *HelmChartInflationGeneratorPlugin) Generate() (rm resmap.ResMap, err error) {
	defer p.cleanup()
	if err = p.checkHelmVersion(); err != nil {
		return nil, err
	}
	if path, exists := p.chartExistsLocally(); !exists {
		if p.Repo == "" {
			return nil, fmt.Errorf(
				"no repo specified for pull, no chart found at '%s'", path)
		}
		if _, err := p.runHelmCommand(p.pullCommand()); err != nil {
			return nil, err
		}
	}
	if len(p.ValuesInline) > 0 {
		p.ValuesFile, err = p.createNewMergedValuesFile()
	} else {
		p.ValuesFile, err = p.copyValuesFile()
	}
	if err != nil {
		return nil, err
	}
	var stdout []byte
	stdout, err = p.runHelmCommand(p.templateCommand())
	if err != nil {
		return nil, err
	}

	rm, err = p.h.ResmapFactory().NewResMapFromBytes(stdout)
	if err == nil {
		return rm, nil
	}
	// try to remove the contents before first "---" because
	// helm may produce messages to stdout before it
	stdoutStr := string(stdout)
	if idx := strings.Index(stdoutStr, "---"); idx != -1 {
		return p.h.ResmapFactory().NewResMapFromBytes([]byte(stdoutStr[idx:]))
	}
	return nil, err
}

func (p *HelmChartInflationGeneratorPlugin) templateCommand() []string {
	args := []string{"template"}
	if p.ReleaseName != "" {
		args = append(args, p.ReleaseName)
	}
	args = append(args, filepath.Join(p.absChartHome(), p.Name))
	if p.ValuesFile != "" {
		args = append(args, "--values", p.ValuesFile)
	}
	if p.ReleaseName == "" {
		// AFAICT, this doesn't work as intended due to a bug in helm.
		// See https://github.com/helm/helm/issues/6019
		// I've tried placing the flag before and after the name argument.
		args = append(args, "--generate-name")
	}
	return args
}

func (p *HelmChartInflationGeneratorPlugin) pullCommand() []string {
	args := []string{
		"pull",
		"--untar",
		"--untardir", p.absChartHome(),
		"--repo", p.Repo,
		p.Name}
	if p.Version != "" {
		args = append(args, "--version", p.Version)
	}
	return args
}

// chartExistsLocally will return true if the chart does exist in
// local chart home.
func (p *HelmChartInflationGeneratorPlugin) chartExistsLocally() (string, bool) {
	path := filepath.Join(p.absChartHome(), p.Name)
	s, err := os.Stat(path)
	if err != nil {
		return "", false
	}
	return path, s.IsDir()
}

// checkHelmVersion will return an error if the helm version is not V3
func (p *HelmChartInflationGeneratorPlugin) checkHelmVersion() error {
	stdout, err := p.runHelmCommand([]string{"version", "-c", "--short"})
	if err != nil {
		return err
	}
	r, err := regexp.Compile(`v?\d+(\.\d+)+`)
	if err != nil {
		return err
	}
	v := r.FindString(string(stdout))
	if v == "" {
		return fmt.Errorf("cannot find version string in %s", string(stdout))
	}
	if v[0] == 'v' {
		v = v[1:]
	}
	majorVersion := strings.Split(v, ".")[0]
	if majorVersion != "3" {
		return fmt.Errorf("this plugin requires helm V3 but got v%s", v)
	}
	return nil
}

func NewHelmChartInflationGeneratorPlugin() resmap.GeneratorPlugin {
	return &HelmChartInflationGeneratorPlugin{}
}
<commit_msg>Update HelmChartInflationGenerator.go<commit_after>// Code generated by pluginator on HelmChartInflationGenerator; DO NOT EDIT.
// pluginator {unknown  1970-01-01T00:00:00Z  }

package builtins

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/imdario/mergo"
	"github.com/pkg/errors"
	"sigs.k8s.io/kustomize/api/resmap"
	"sigs.k8s.io/kustomize/api/types"
	"sigs.k8s.io/yaml"
)

// HelmChartInflationGeneratorPlugin is a plugin to generate resources
// from a remote or local helm chart.
type HelmChartInflationGeneratorPlugin struct {
	h *resmap.PluginHelpers
	types.HelmGlobals
	types.HelmChart
	tmpDir string
}

var KustomizePlugin HelmChartInflationGeneratorPlugin

const (
	valuesMergeOptionMerge    = "merge"
	valuesMergeOptionOverride = "override"
	valuesMergeOptionReplace  = "replace"
)

var legalMergeOptions = []string{
	valuesMergeOptionMerge,
	valuesMergeOptionOverride,
	valuesMergeOptionReplace,
}

// Config uses the input plugin configurations `config` to setup the generator
// options
func (p *HelmChartInflationGeneratorPlugin) Config(
	h *resmap.PluginHelpers, config []byte) (err error) {
	if h.GeneralConfig() == nil {
		return fmt.Errorf("unable to access general config")
	}
	if !h.GeneralConfig().HelmConfig.Enabled {
		return fmt.Errorf("must specify --enable-helm")
	}
	if h.GeneralConfig().HelmConfig.Command == "" {
		return fmt.Errorf("must specify --helm-command")
	}
	p.h = h
	if err = yaml.Unmarshal(config, p); err != nil {
		return
	}
	return p.validateArgs()
}

// This uses the real file system since tmpDir may be used
// by the helm subprocess.  Cannot use a chroot jail or fake
// filesystem since we allow the user to use previously
// downloaded charts.  This is safe since this plugin is
// owned by kustomize.
func (p *HelmChartInflationGeneratorPlugin) establishTmpDir() (err error) {
	if p.tmpDir != "" {
		// already done.
		return nil
	}
	p.tmpDir, err = ioutil.TempDir("", "kustomize-helm-")
	return err
}

func (p *HelmChartInflationGeneratorPlugin) validateArgs() (err error) {
	if p.Name == "" {
		return fmt.Errorf("chart name cannot be empty")
	}

	// ChartHome might be consulted by the plugin (to read
	// values files below it), so it must be located under
	// the loader root (unless root restrictions are
	// disabled, in which case this can be an absolute path).
	if p.ChartHome == "" {
		p.ChartHome = "charts"
	}

	// The ValuesFile may be consulted by the plugin, so it must
	// be under the loader root (unless root restrictions are
	// disabled).
	if p.ValuesFile == "" {
		p.ValuesFile = filepath.Join(p.ChartHome, p.Name, "values.yaml")
	}

	if err = p.errIfIllegalValuesMerge(); err != nil {
		return err
	}

	// ConfigHome is not loaded by the plugin, and can be located anywhere.
	if p.ConfigHome == "" {
		if err = p.establishTmpDir(); err != nil {
			return errors.Wrap(
				err, "unable to create tmp dir for HELM_CONFIG_HOME")
		}
		p.ConfigHome = filepath.Join(p.tmpDir, "helm")
	}
	return nil
}

func (p *HelmChartInflationGeneratorPlugin) errIfIllegalValuesMerge() error {
	if p.ValuesMerge == "" {
		// Use the default.
		p.ValuesMerge = valuesMergeOptionOverride
		return nil
	}
	for _, opt := range legalMergeOptions {
		if p.ValuesMerge == opt {
			return nil
		}
	}
	return fmt.Errorf("valuesMerge must be one of %v", legalMergeOptions)
}

func (p *HelmChartInflationGeneratorPlugin) absChartHome() string {
	if filepath.IsAbs(p.ChartHome) {
		return p.ChartHome
	}
	return filepath.Join(p.h.Loader().Root(), p.ChartHome)
}

func (p *HelmChartInflationGeneratorPlugin) runHelmCommand(
	args []string) ([]byte, error) {
	stdout := new(bytes.Buffer)
	stderr := new(bytes.Buffer)
	cmd := exec.Command(p.h.GeneralConfig().HelmConfig.Command, args...)
	cmd.Stdout = stdout
	cmd.Stderr = stderr
	env := []string{
		fmt.Sprintf("HELM_CONFIG_HOME=%s", p.ConfigHome),
		fmt.Sprintf("HELM_CACHE_HOME=%s/.cache", p.ConfigHome),
		fmt.Sprintf("HELM_DATA_HOME=%s/.data", p.ConfigHome)}
	cmd.Env = append(os.Environ(), env...)
	err := cmd.Run()
	if err != nil {
		helm := p.h.GeneralConfig().HelmConfig.Command
		err = errors.Wrap(
			fmt.Errorf(
				"unable to run: '%s %s' with env=%s (is '%s' installed?)",
				helm, strings.Join(args, " "), env, helm),
			stderr.String(),
		)
	}
	return stdout.Bytes(), err
}

// createNewMergedValuesFile replaces/merges original values file with ValuesInline.
func (p *HelmChartInflationGeneratorPlugin) createNewMergedValuesFile() (
	path string, err error) {
	if p.ValuesMerge == valuesMergeOptionMerge ||
		p.ValuesMerge == valuesMergeOptionOverride {
		if err = p.replaceValuesInline(); err != nil {
			return "", err
		}
	}
	var b []byte
	b, err = yaml.Marshal(p.ValuesInline)
	if err != nil {
		return "", err
	}
	return p.writeValuesBytes(b)
}

func (p *HelmChartInflationGeneratorPlugin) replaceValuesInline() error {
	pValues, err := p.h.Loader().Load(p.ValuesFile)
	if err != nil {
		return err
	}
	chValues := make(map[string]interface{})
	if err = yaml.Unmarshal(pValues, &chValues); err != nil {
		return err
	}
	switch p.ValuesMerge {
	case valuesMergeOptionOverride:
		err = mergo.Merge(
			&chValues, p.ValuesInline, mergo.WithOverride)
	case valuesMergeOptionMerge:
		err = mergo.Merge(&chValues, p.ValuesInline)
	}
	p.ValuesInline = chValues
	return err
}

// copyValuesFile to avoid branching.  TODO: get rid of this.
func (p *HelmChartInflationGeneratorPlugin) copyValuesFile() (string, error) {
	b, err := p.h.Loader().Load(p.ValuesFile)
	if err != nil {
		return "", err
	}
	return p.writeValuesBytes(b)
}

// Write a absolute path file in the tmp file system.
func (p *HelmChartInflationGeneratorPlugin) writeValuesBytes(
	b []byte) (string, error) {
	if err := p.establishTmpDir(); err != nil {
		return "", fmt.Errorf("cannot create tmp dir to write helm values")
	}
	path := filepath.Join(p.tmpDir, p.Name+"-kustomize-values.yaml")
	return path, ioutil.WriteFile(path, b, 0644)
}

func (p *HelmChartInflationGeneratorPlugin) cleanup() {
	if p.tmpDir != "" {
		os.RemoveAll(p.tmpDir)
	}
}

// Generate implements generator
func (p *HelmChartInflationGeneratorPlugin) Generate() (rm resmap.ResMap, err error) {
	defer p.cleanup()
	if err = p.checkHelmVersion(); err != nil {
		return nil, err
	}
	if path, exists := p.chartExistsLocally(); !exists {
		if p.Repo == "" {
			return nil, fmt.Errorf(
				"no repo specified for pull, no chart found at '%s'", path)
		}
		if _, err := p.runHelmCommand(p.pullCommand()); err != nil {
			return nil, err
		}
	}
	if len(p.ValuesInline) > 0 {
		p.ValuesFile, err = p.createNewMergedValuesFile()
	} else {
		p.ValuesFile, err = p.copyValuesFile()
	}
	if err != nil {
		return nil, err
	}
	var stdout []byte
	stdout, err = p.runHelmCommand(p.templateCommand())
	if err != nil {
		return nil, err
	}

	rm, err = p.h.ResmapFactory().NewResMapFromBytes(stdout)
	if err == nil {
		return rm, nil
	}
	// try to remove the contents before first "---" because
	// helm may produce messages to stdout before it
	stdoutStr := string(stdout)
	if idx := strings.Index(stdoutStr, "---"); idx != -1 {
		return p.h.ResmapFactory().NewResMapFromBytes([]byte(stdoutStr[idx:]))
	}
	return nil, err
}

func (p *HelmChartInflationGeneratorPlugin) templateCommand() []string {
	args := []string{"template"}
	if p.ReleaseName != "" {
		args = append(args, p.ReleaseName)
	}
	if p.Namespace != "" {
		args = append(args, "--namespace", p.Namespace)
	}
	args = append(args, filepath.Join(p.absChartHome(), p.Name))
	if p.ValuesFile != "" {
		args = append(args, "--values", p.ValuesFile)
	}
	if p.ReleaseName == "" {
		// AFAICT, this doesn't work as intended due to a bug in helm.
		// See https://github.com/helm/helm/issues/6019
		// I've tried placing the flag before and after the name argument.
		args = append(args, "--generate-name")
	}
	return args
}

func (p *HelmChartInflationGeneratorPlugin) pullCommand() []string {
	args := []string{
		"pull",
		"--untar",
		"--untardir", p.absChartHome(),
		"--repo", p.Repo,
		p.Name}
	if p.Version != "" {
		args = append(args, "--version", p.Version)
	}
	return args
}

// chartExistsLocally will return true if the chart does exist in
// local chart home.
func (p *HelmChartInflationGeneratorPlugin) chartExistsLocally() (string, bool) {
	path := filepath.Join(p.absChartHome(), p.Name)
	s, err := os.Stat(path)
	if err != nil {
		return "", false
	}
	return path, s.IsDir()
}

// checkHelmVersion will return an error if the helm version is not V3
func (p *HelmChartInflationGeneratorPlugin) checkHelmVersion() error {
	stdout, err := p.runHelmCommand([]string{"version", "-c", "--short"})
	if err != nil {
		return err
	}
	r, err := regexp.Compile(`v?\d+(\.\d+)+`)
	if err != nil {
		return err
	}
	v := r.FindString(string(stdout))
	if v == "" {
		return fmt.Errorf("cannot find version string in %s", string(stdout))
	}
	if v[0] == 'v' {
		v = v[1:]
	}
	majorVersion := strings.Split(v, ".")[0]
	if majorVersion != "3" {
		return fmt.Errorf("this plugin requires helm V3 but got v%s", v)
	}
	return nil
}

func NewHelmChartInflationGeneratorPlugin() resmap.GeneratorPlugin {
	return &HelmChartInflationGeneratorPlugin{}
}
<|endoftext|>