<commit_before>package httptests

import (
	"bytes"
	"encoding/json"
	"io"
	"mime/multipart"
	"net/url"
	"strings"

	"github.com/Nivl/go-params"
	"github.com/Nivl/go-params/formfile"
	"github.com/Nivl/go-rest-tools/router"
	"github.com/gorilla/mux"
)

// RequestInfo represents the params accepted by NewRequest
type RequestInfo struct {
	Endpoint *router.Endpoint

	Params interface{}  // Optional
	Auth   *RequestAuth // Optional
	// Router is used to parse Mux Variables. Default on the api router
	Router *mux.Router

	urlParams   map[string]string
	bodyParams  map[string]interface{} // can be string or []string
	queryParams map[string]interface{} // can be string or []string
	fileParams  map[string]*formfile.FormFile
}

// ParseParams parses the params and copy them in the right list:
// urlParams, bodyParams, and queryParams
func (ri *RequestInfo) ParseParams() {
	ri.urlParams = map[string]string{}
	ri.bodyParams = map[string]interface{}{}
	ri.queryParams = map[string]interface{}{}

	if ri.Params == nil {
		return
	}

	var sources map[string]url.Values
	p := params.New(ri.Params)
	sources, ri.fileParams = p.Extract()

	for k, v := range sources["url"] {
		ri.urlParams[k] = v[0]
	}
	for k, v := range sources["form"] {
		ri.bodyParams[k] = v[0]
		if len(v) > 1 {
			ri.bodyParams[k] = v
		}
	}
	for k, v := range sources["query"] {
		ri.queryParams[k] = v[0]
		if len(v) > 1 {
			ri.queryParams[k] = v
		}
	}
}

// URL returns the full URL
func (ri *RequestInfo) URL() string {
	url := ri.Endpoint.Path
	for param, value := range ri.urlParams {
		url = strings.Replace(url, "{"+param+"}", value, -1)
	}
	return url
}

// PopulateQuery populate the query string of a request
func (ri *RequestInfo) PopulateQuery(qs url.Values) {
	for key, value := range ri.queryParams {
		if slice, isSlice := value.([]string); isSlice {
			qs[key] = slice
		} else {
			qs.Add(key, value.(string))
		}
	}
}

// Body returns the full Body of the request
func (ri *RequestInfo) Body() (mime string, body io.Reader, err error) {
	if ri.fileParams != nil && len(ri.fileParams) > 0 {
		return ri.BodyMultipart()
	}
	return ri.BodyJSON()
}

// BodyJSON returns the body of the request encoded in JSON
func (ri *RequestInfo) BodyJSON() (mime string, body io.Reader, err error) {
	mime = "application/json; charset=utf-8"
	body = bytes.NewBufferString("")

	// Parse the body as a JSON object
	if len(ri.bodyParams) > 0 {
		var jsonDump []byte
		jsonDump, err = json.Marshal(ri.bodyParams)
		if err != nil {
			return
		}
		body = bytes.NewBuffer(jsonDump)
	}
	return
}

// BodyMultipart returns the body of the request as multipart data
func (ri *RequestInfo) BodyMultipart() (mime string, body io.Reader, err error) {
	output := &bytes.Buffer{}
	writer := multipart.NewWriter(output)

	// We attach the files
	for name, f := range ri.fileParams {
		part, err := writer.CreateFormFile(name, f.Header.Filename)
		if err != nil {
			return "", nil, err
		}

		_, err = io.Copy(part, f.File)
		if err != nil {
			return "", nil, err
		}
	}

	// We attach any other form params
	for name, value := range ri.bodyParams {
		if slice, isSlice := value.([]string); isSlice {
			for _, elem := range slice {
				if err = writer.WriteField(name, elem); err != nil {
					return "", nil, err
				}
			}
		} else {
			if err = writer.WriteField(name, value.(string)); err != nil {
				return "", nil, err
			}
		}
	}

	if err := writer.Close(); err != nil {
		return "", nil, err
	}

	return writer.FormDataContentType(), output, nil
}
<commit_msg>style(httptests): add a FIXME note<commit_after>package httptests

import (
	"bytes"
	"encoding/json"
	"io"
	"mime/multipart"
	"net/url"
	"strings"

	"github.com/Nivl/go-params"
	"github.com/Nivl/go-params/formfile"
	"github.com/Nivl/go-rest-tools/router"
	"github.com/gorilla/mux"
)

// RequestInfo represents the params accepted by NewRequest
type RequestInfo struct {
	Endpoint *router.Endpoint

	Params interface{}  // Optional
	Auth   *RequestAuth // Optional
	// Router is used to parse Mux Variables. Default on the api router
	Router *mux.Router

	urlParams   map[string]string
	bodyParams  map[string]interface{} // can be string or []string
	queryParams map[string]interface{} // can be string or []string
	fileParams  map[string]*formfile.FormFile
}

// ParseParams parses the params and copy them in the right list:
// urlParams, bodyParams, and queryParams
func (ri *RequestInfo) ParseParams() {
	ri.urlParams = map[string]string{}
	ri.bodyParams = map[string]interface{}{}
	ri.queryParams = map[string]interface{}{}

	if ri.Params == nil {
		return
	}

	var sources map[string]url.Values
	p := params.New(ri.Params)
	sources, ri.fileParams = p.Extract()

	for k, v := range sources["url"] {
		ri.urlParams[k] = v[0]
	}
	for k, v := range sources["form"] {
		ri.bodyParams[k] = v[0]
		if len(v) > 1 {
			ri.bodyParams[k] = v
		}
	}
	for k, v := range sources["query"] {
		ri.queryParams[k] = v[0]
		if len(v) > 1 {
			ri.queryParams[k] = v
		}
	}
}

// URL returns the full URL
func (ri *RequestInfo) URL() string {
	url := ri.Endpoint.Path
	for param, value := range ri.urlParams {
		url = strings.Replace(url, "{"+param+"}", value, -1)
	}
	return url
}

// PopulateQuery populate the query string of a request
func (ri *RequestInfo) PopulateQuery(qs url.Values) {
	for key, value := range ri.queryParams {
		if slice, isSlice := value.([]string); isSlice {
			qs[key] = slice
		} else {
			qs.Add(key, value.(string))
		}
	}
}

// Body returns the full Body of the request
func (ri *RequestInfo) Body() (mime string, body io.Reader, err error) {
	if ri.fileParams != nil && len(ri.fileParams) > 0 {
		return ri.BodyMultipart()
	}
	return ri.BodyJSON()
}

// BodyJSON returns the body of the request encoded in JSON
// FIXME(melvin): because the data are from a map of string, all the
// 		JSON data will also be string. There's no way to use the output
// 		to recreate a new JSON object containing non string value
func (ri *RequestInfo) BodyJSON() (mime string, body io.Reader, err error) {
	mime = "application/json; charset=utf-8"
	body = bytes.NewBufferString("")

	// Parse the body as a JSON object
	if len(ri.bodyParams) > 0 {
		var jsonDump []byte
		jsonDump, err = json.Marshal(ri.bodyParams)
		if err != nil {
			return
		}
		body = bytes.NewBuffer(jsonDump)
	}
	return
}

// BodyMultipart returns the body of the request as multipart data
func (ri *RequestInfo) BodyMultipart() (mime string, body io.Reader, err error) {
	output := &bytes.Buffer{}
	writer := multipart.NewWriter(output)

	// We attach the files
	for name, f := range ri.fileParams {
		part, err := writer.CreateFormFile(name, f.Header.Filename)
		if err != nil {
			return "", nil, err
		}

		_, err = io.Copy(part, f.File)
		if err != nil {
			return "", nil, err
		}
	}

	// We attach any other form params
	for name, value := range ri.bodyParams {
		if slice, isSlice := value.([]string); isSlice {
			for _, elem := range slice {
				if err = writer.WriteField(name, elem); err != nil {
					return "", nil, err
				}
			}
		} else {
			if err = writer.WriteField(name, value.(string)); err != nil {
				return "", nil, err
			}
		}
	}

	if err := writer.Close(); err != nil {
		return "", nil, err
	}

	return writer.FormDataContentType(), output, nil
}
<|endoftext|>


<commit_before>// Package run is the core logic of get-headers
package run

import (
	"fmt"
	"io"
	"net/http"
	"os"
	"text/tabwriter"
	"time"

	"golang.org/x/sync/errgroup"

	"github.com/carlmjohnson/get-headers/prettyprint"
)

// Don't follow redirects
func checkRedirect(req *http.Request, via []*http.Request) error {
	return http.ErrUseLastResponse
}

// client for all http requests
var client = http.Client{
	CheckRedirect: checkRedirect,
	Transport: &http.Transport{
		DisableCompression: true,
	},
}

// Main takes a list of urls and request parameters, then fetches the URLs and
// outputs the headers to stdout
func Main(cookie, etag string, gzip, ignoreBody bool, urls ...string) error {
	for nurl, url := range urls {
		// Separate subsequent lookups with newline
		if nurl > 0 {
			fmt.Println()
		}

		req, err := http.NewRequest("GET", url, nil)
		if err != nil {
			return err
		}

		if gzip {
			req.Header.Add("Accept-Encoding", "gzip, deflate")
		}

		if etag != "" {
			req.Header.Add("If-None-Match", etag)
		}

		if cookie != "" {
			req.Header.Add("Cookie", cookie)
		}

		start := time.Now()
		resp, err := client.Do(req)
		duration := time.Since(start)

		if err != nil {
			return err
		}

		var (
			n  int64
			eg errgroup.Group
		)

		if !ignoreBody {
			eg.Go(func() error {
				// Copying to /dev/null just to make sure this is real
				n, err = io.Copy(io.Discard, resp.Body)
				duration = time.Since(start)
				if err != nil {
					return err
				}
				return nil
			})
		}

		fmt.Println("GET", url)
		fmt.Println(resp.Proto, resp.Status)
		fmt.Println()
		fmt.Println(prettyprint.ResponseHeader(resp.Header))

		if err := eg.Wait(); err != nil {
			return err
		}

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

		tw := tabwriter.NewWriter(os.Stdout, 0, 0, 2, ' ', 0)
		fmt.Fprintf(tw, "Time\t%s\n", prettyprint.Duration(duration))
		if n != 0 {
			fmt.Fprintf(tw, "Content length\t%s\n", prettyprint.Size(n))
			bps := prettyprint.Size(float64(n) / duration.Seconds())
			fmt.Fprintf(tw, "Speed\t%s/s\n", bps)
		}
		if err := tw.Flush(); err != nil {
			return err
		}
	}

	return nil
}
<commit_msg>Use context for cancellation<commit_after>// Package run is the core logic of get-headers
package run

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/signal"
	"text/tabwriter"
	"time"

	"golang.org/x/sync/errgroup"

	"github.com/carlmjohnson/get-headers/prettyprint"
)

// Don't follow redirects
func checkRedirect(req *http.Request, via []*http.Request) error {
	return http.ErrUseLastResponse
}

// client for all http requests
var client = http.Client{
	CheckRedirect: checkRedirect,
	Transport: &http.Transport{
		DisableCompression: true,
	},
}

// Main takes a list of urls and request parameters, then fetches the URLs and
// outputs the headers to stdout
func Main(cookie, etag string, gzip, ignoreBody bool, urls ...string) error {
	for i, url := range urls {
		// Separate subsequent lookups with newline
		if i > 0 {
			fmt.Println()
		}
		if err := getHeaders(cookie, etag, gzip, ignoreBody, url); err != nil {
			return err
		}
	}
	return nil
}

func getHeaders(cookie, etag string, gzip, ignoreBody bool, url string) error {
	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, os.Kill)
	defer cancel()
	eg, ctx := errgroup.WithContext(ctx)

	req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
	if err != nil {
		return err
	}

	if gzip {
		req.Header.Add("Accept-Encoding", "gzip, deflate")
	}

	if etag != "" {
		req.Header.Add("If-None-Match", etag)
	}

	if cookie != "" {
		req.Header.Add("Cookie", cookie)
	}

	start := time.Now()
	resp, err := client.Do(req)
	duration := time.Since(start)

	if err != nil {
		return err
	}

	var n int64

	if !ignoreBody {
		eg.Go(func() error {
			// Copying to /dev/null just to make sure this is real
			n, err = io.Copy(io.Discard, resp.Body)
			duration = time.Since(start)
			if err != nil {
				return err
			}
			return nil
		})
	}

	fmt.Println("GET", url)
	fmt.Println(resp.Proto, resp.Status)
	fmt.Println()
	fmt.Println(prettyprint.ResponseHeader(resp.Header))

	if err := eg.Wait(); err != nil {
		return err
	}

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

	tw := tabwriter.NewWriter(os.Stdout, 0, 0, 2, ' ', 0)
	fmt.Fprintf(tw, "Time\t%s\n", prettyprint.Duration(duration))
	if n != 0 {
		fmt.Fprintf(tw, "Content length\t%s\n", prettyprint.Size(n))
		bps := prettyprint.Size(float64(n) / duration.Seconds())
		fmt.Fprintf(tw, "Speed\t%s/s\n", bps)
	}
	if err := tw.Flush(); err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package gxutil

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"time"

	stump "github.com/whyrusleeping/stump"
)

type ErrAlreadyInstalled struct {
	pkg string
}

func IsErrAlreadyInstalled(err error) bool {
	_, ok := err.(ErrAlreadyInstalled)
	return ok
}

func (eai ErrAlreadyInstalled) Error() string {
	return fmt.Sprintf("package %s already installed", eai.pkg)
}

func (pm *PM) GetPackageTo(hash, out string) (*Package, error) {
	var pkg Package
	_, err := os.Stat(out)
	if err == nil {
		err := FindPackageInDir(&pkg, out)
		if err == nil {
			return &pkg, nil
		} else if !os.IsNotExist(err) {
			return nil, err
		}

		stump.VLog("Target directory already exists but isn't a valid package, cleaning up...")
		os.RemoveAll(out)
	}

	if err != nil && !os.IsNotExist(err) {
		return nil, err
	}

	outtemp := out + ".part"

	// check if a fetch was previously started and failed, cleanup if found
	_, err = os.Stat(outtemp)
	if err == nil {
		stump.VLog("Found previously failed fetch, cleaning up...")
		if err := os.RemoveAll(outtemp); err != nil {
			stump.Error("cleaning up previous aborted transfer: %s", err)
		}
	}

	begin := time.Now()
	stump.VLog("  - fetching %s via ipfs api", hash)
	tries := 3
	for i := 0; i < tries; i++ {
		if err := pm.Shell().Get(hash, outtemp); err != nil {
			stump.Error("from shell.Get(): %v", err)

			rmerr := os.RemoveAll(outtemp)
			if rmerr != nil {
				stump.Error("cleaning up temp download directory: %s", rmerr)
			}

			if i == tries-1 {
				return nil, err
			}
			stump.Log("retrying fetch %s after a second...", hash)
			time.Sleep(time.Second)
		} else {
			if err := os.Rename(outtemp, out); err != nil {
				return nil, err
			}
			break
		}
	}
	stump.VLog("  - fetch finished in %s", time.Since(begin))

	err = FindPackageInDir(&pkg, out)
	if err != nil {
		return nil, err
	}

	return &pkg, nil
}

func FindPackageInDir(pkg interface{}, dir string) error {
	if err := LoadPackageFile(pkg, filepath.Join(dir, PkgFileName)); err == nil {
		return nil
	}

	name, err := PackageNameInDir(dir)
	if err != nil {
		return err
	}
	return LoadPackageFile(pkg, filepath.Join(dir, name, PkgFileName))
}

func PackageNameInDir(dir string) (string, error) {
	fs, err := ioutil.ReadDir(dir)
	if err != nil {
		return "", err
	}

	if len(fs) == 0 {
		return "", fmt.Errorf("no package found in hashdir: %s", dir)
	}

	if len(fs) > 1 {
		return "", fmt.Errorf("found multiple packages in hashdir: %s", dir)
	}

	return fs[0].Name(), nil
}
<commit_msg>Remove invalid package dir if it exists<commit_after>package gxutil

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"time"

	stump "github.com/whyrusleeping/stump"
)

type ErrAlreadyInstalled struct {
	pkg string
}

func IsErrAlreadyInstalled(err error) bool {
	_, ok := err.(ErrAlreadyInstalled)
	return ok
}

func (eai ErrAlreadyInstalled) Error() string {
	return fmt.Sprintf("package %s already installed", eai.pkg)
}

func (pm *PM) GetPackageTo(hash, out string) (*Package, error) {
	var pkg Package
	_, err := os.Stat(out)
	if err == nil {
		err := FindPackageInDir(&pkg, out)
		if err == nil {
			return &pkg, nil
		}

		stump.VLog("Target directory already exists but isn't a valid package, cleaning up...")
		if oErr := os.RemoveAll(out); oErr != nil {
			stump.Error("cannot purge existing target directory:", oErr)
			return nil, oErr
		}
	}

	if err != nil && !os.IsNotExist(err) {
		return nil, err
	}

	outtemp := out + ".part"

	// check if a fetch was previously started and failed, cleanup if found
	_, err = os.Stat(outtemp)
	if err == nil {
		stump.VLog("Found previously failed fetch, cleaning up...")
		if err := os.RemoveAll(outtemp); err != nil {
			stump.Error("cleaning up previous aborted transfer: %s", err)
		}
	}

	begin := time.Now()
	stump.VLog("  - fetching %s via ipfs api", hash)
	tries := 3
	for i := 0; i < tries; i++ {
		if err := pm.Shell().Get(hash, outtemp); err != nil {
			stump.Error("from shell.Get(): %v", err)

			rmerr := os.RemoveAll(outtemp)
			if rmerr != nil {
				stump.Error("cleaning up temp download directory: %s", rmerr)
			}

			if i == tries-1 {
				return nil, err
			}
			stump.Log("retrying fetch %s after a second...", hash)
			time.Sleep(time.Second)
		} else {
			if err := os.Rename(outtemp, out); err != nil {
				return nil, err
			}
			break
		}
	}
	stump.VLog("  - fetch finished in %s", time.Since(begin))

	err = FindPackageInDir(&pkg, out)
	if err != nil {
		return nil, err
	}

	return &pkg, nil
}

func FindPackageInDir(pkg interface{}, dir string) error {
	if err := LoadPackageFile(pkg, filepath.Join(dir, PkgFileName)); err == nil {
		return nil
	}

	name, err := PackageNameInDir(dir)
	if err != nil {
		return err
	}
	return LoadPackageFile(pkg, filepath.Join(dir, name, PkgFileName))
}

func PackageNameInDir(dir string) (string, error) {
	fs, err := ioutil.ReadDir(dir)
	if err != nil {
		return "", err
	}

	if len(fs) == 0 {
		return "", fmt.Errorf("no package found in hashdir: %s", dir)
	}

	if len(fs) > 1 {
		return "", fmt.Errorf("found multiple packages in hashdir: %s", dir)
	}

	return fs[0].Name(), nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package zanzibar

import (
	"io"
	"net/http"
	"net/http/pprof"
	"os"
	"path/filepath"
	"strconv"
	"sync"
	"time"

	"github.com/pkg/errors"
	"github.com/uber-go/tally"
	"github.com/uber-go/zap"
)

// Clients interface is a placeholder for the generated clients
type Clients interface {
}

// LoggerOptions configures logging in the gateway
type LoggerOptions struct {
	FileName string
}

// MetricsOptions ...
type MetricsOptions struct {
	FlushInterval time.Duration
	Service       string
}

// TChannelOptions ...
type TChannelOptions struct {
	ServiceName string
	ProcessName string
}

// Options configures the gateway
type Options struct {
	IP       string
	Port     int32
	Logger   LoggerOptions
	Metrics  MetricsOptions
	TChannel TChannelOptions

	Clients        Clients
	MetricsBackend tally.CachedStatsReporter
}

// Gateway type
type Gateway struct {
	IP          string
	Port        int32
	RealPort    int32
	RealAddr    string
	WaitGroup   *sync.WaitGroup
	Clients     Clients
	Logger      zap.Logger
	MetricScope tally.Scope

	router            *Router
	loggerFile        *os.File
	metricScopeCloser io.Closer
	metricsBackend    tally.CachedStatsReporter
	server            *HTTPServer
	tchannelServer    *TChannelServer
	// clients?
	//	- logger
	//	- statsd
	//	- panic ???
	//	- http server ?
	//	- process reporter ?
	// HTTP server?
}

// CreateGateway func
func CreateGateway(opts *Options) (*Gateway, error) {
	if opts.MetricsBackend == nil {
		panic("opts.MetricsBackend required")
	}

	if opts.Clients == nil {
		panic("opts.Clients required")
	}

	gateway := &Gateway{
		IP:        opts.IP,
		Port:      opts.Port,
		WaitGroup: &sync.WaitGroup{},
		Clients:   opts.Clients,

		metricsBackend: opts.MetricsBackend,
	}

	gateway.router = NewRouter(gateway)

	if err := gateway.setupLogger(&opts.Logger); err != nil {
		return nil, err
	}

	if err := gateway.setupMetrics(&opts.Metrics); err != nil {
		return nil, err
	}

	if err := gateway.setupHTTPServer(); err != nil {
		return nil, err
	}

	if err := gateway.setupTChannel(&opts.TChannel); err != nil {
		return nil, err
	}

	return gateway, nil
}

// RegisterFn type used to avoid cyclic dependencies
type RegisterFn func(gateway *Gateway, router *Router)

// Bootstrap func
func (gateway *Gateway) Bootstrap(register RegisterFn) error {
	gateway.register(register)

	_, err := gateway.server.JustListen()
	if err != nil {
		gateway.Logger.Error("Error listening on port",
			zap.String("error", err.Error()),
		)
		return errors.Wrap(err, "error listening on port")
	}

	gateway.RealPort = gateway.server.RealPort
	gateway.RealAddr = gateway.server.RealAddr

	gateway.WaitGroup.Add(1)
	go gateway.server.JustServe(gateway.WaitGroup)

	return nil
}

func (gateway *Gateway) register(register RegisterFn) {
	gateway.router.RegisterRaw("GET", "/debug/pprof", pprof.Index)
	gateway.router.RegisterRaw("GET", "/debug/pprof/cmdline", pprof.Cmdline)
	gateway.router.RegisterRaw("GET", "/debug/pprof/profile", pprof.Profile)
	gateway.router.RegisterRaw("GET", "/debug/pprof/symbol", pprof.Symbol)
	gateway.router.RegisterRaw("POST", "/debug/pprof/symbol", pprof.Symbol)
	gateway.router.RegisterRaw(
		"GET", "/debug/pprof/goroutine", pprof.Handler("goroutine").ServeHTTP,
	)
	gateway.router.RegisterRaw(
		"GET", "/debug/pprof/heap", pprof.Handler("heap").ServeHTTP,
	)
	gateway.router.RegisterRaw(
		"GET", "/debug/pprof/threadcreate",
		pprof.Handler("threadcreate").ServeHTTP,
	)
	gateway.router.RegisterRaw(
		"GET", "/debug/pprof/block", pprof.Handler("block").ServeHTTP,
	)

	register(gateway, gateway.router)
}

// Close the http server
func (gateway *Gateway) Close() {
	if gateway.loggerFile != nil {
		_ = gateway.loggerFile.Sync()
		_ = gateway.loggerFile.Close()
	}

	gateway.metricsBackend.Flush()
	_ = gateway.metricScopeCloser.Close()
	gateway.server.Close()
}

// Wait for gateway to close the server
func (gateway *Gateway) Wait() {
	gateway.WaitGroup.Wait()
}

func (gateway *Gateway) setupMetrics(opts *MetricsOptions) error {
	// TODO: decide what default tags we want...
	defaultTags := &map[string]string{}

	prefix := opts.Service + ".production.all-workers"

	flushIntervalConfig := opts.FlushInterval
	var flushInterval time.Duration
	if flushIntervalConfig == 0 {
		flushInterval = 1 * time.Second
	} else {
		flushInterval = flushIntervalConfig
	}

	scope, closer := tally.NewCachedRootScope(
		prefix,
		*defaultTags,
		gateway.metricsBackend,
		flushInterval,
		tally.DefaultSeparator,
	)
	gateway.MetricScope = scope
	gateway.metricScopeCloser = closer

	return nil
}

func (gateway *Gateway) setupLogger(opts *LoggerOptions) error {
	var output zap.Option
	tempLogger := zap.New(
		zap.NewJSONEncoder(),
		zap.Output(os.Stderr),
	)

	loggerFileName := opts.FileName

	if loggerFileName == "" {
		output = zap.Output(os.Stdout)
	} else {
		err := os.MkdirAll(filepath.Dir(loggerFileName), 0777)
		if err != nil {
			tempLogger.Error("Error creating log directory",
				zap.String("error", err.Error()),
			)
			return errors.Wrap(err, "Error creating log directory")
		}

		loggerFile, err := os.OpenFile(
			loggerFileName,
			os.O_APPEND|os.O_WRONLY|os.O_CREATE,
			0666,
		)
		if err != nil {
			tempLogger.Error("Error opening log file",
				zap.String("error", err.Error()),
			)
			return errors.Wrap(err, "Error opening log file")
		}
		gateway.loggerFile = loggerFile
		output = zap.Output(loggerFile)
	}

	// Default to a STDOUT logger
	gateway.Logger = zap.New(
		zap.NewJSONEncoder(
			zap.RFC3339Formatter("ts"),
		),
		output,
	)
	return nil
}

func (gateway *Gateway) setupHTTPServer() error {
	gateway.server = &HTTPServer{
		Server: &http.Server{
			Addr:    gateway.IP + ":" + strconv.FormatInt(int64(gateway.Port), 10),
			Handler: gateway.router,
		},
		Logger: gateway.Logger,
	}

	return nil
}

func (gateway *Gateway) setupTChannel(opts *TChannelOptions) error {
	if opts == nil {
		return errors.New("TChannelOptions were nil")
	}

	tchannelServer, err := NewTChannelServer(
		&TChannelServerOptions{
			ServiceName: opts.ServiceName,
			ProcessName: opts.ProcessName,
		})

	if err != nil {
		return err
	}

	gateway.tchannelServer = tchannelServer
	return nil
}
<commit_msg>runtime: update to use StaticConfig()<commit_after>// Copyright (c) 2017 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package zanzibar

import (
	"io"
	"net/http"
	"net/http/pprof"
	"os"
	"path/filepath"
	"strconv"
	"sync"
	"time"

	"github.com/pkg/errors"
	"github.com/uber-go/tally"
	"github.com/uber-go/zap"
)

// Clients interface is a placeholder for the generated clients
type Clients interface {
}

// Options configures the gateway
type Options struct {
	Clients        Clients
	MetricsBackend tally.CachedStatsReporter
}

// Gateway type
type Gateway struct {
	IP          string
	Port        int32
	RealPort    int32
	RealAddr    string
	WaitGroup   *sync.WaitGroup
	Clients     Clients
	Logger      zap.Logger
	MetricScope tally.Scope

	router            *Router
	loggerFile        *os.File
	metricScopeCloser io.Closer
	metricsBackend    tally.CachedStatsReporter
	server            *HTTPServer
	tchannelServer    *TChannelServer
	// clients?
	//	- panic ???
	//	- process reporter ?
}

// CreateGateway func
func CreateGateway(
	config *StaticConfig, opts *Options,
) (*Gateway, error) {
	if opts.MetricsBackend == nil {
		panic("opts.MetricsBackend required")
	}

	if opts.Clients == nil {
		panic("opts.Clients required")
	}

	gateway := &Gateway{
		IP:        config.GetString("ip"),
		Port:      int32(config.GetInt("port")),
		WaitGroup: &sync.WaitGroup{},
		Clients:   opts.Clients,

		metricsBackend: opts.MetricsBackend,
	}

	gateway.router = NewRouter(gateway)

	if err := gateway.setupLogger(config); err != nil {
		return nil, err
	}

	if err := gateway.setupMetrics(config); err != nil {
		return nil, err
	}

	if err := gateway.setupHTTPServer(); err != nil {
		return nil, err
	}

	if err := gateway.setupTChannel(config); err != nil {
		return nil, err
	}

	return gateway, nil
}

// RegisterFn type used to avoid cyclic dependencies
type RegisterFn func(gateway *Gateway, router *Router)

// Bootstrap func
func (gateway *Gateway) Bootstrap(register RegisterFn) error {
	gateway.register(register)

	_, err := gateway.server.JustListen()
	if err != nil {
		gateway.Logger.Error("Error listening on port",
			zap.String("error", err.Error()),
		)
		return errors.Wrap(err, "error listening on port")
	}

	gateway.RealPort = gateway.server.RealPort
	gateway.RealAddr = gateway.server.RealAddr

	gateway.WaitGroup.Add(1)
	go gateway.server.JustServe(gateway.WaitGroup)

	return nil
}

func (gateway *Gateway) register(register RegisterFn) {
	gateway.router.RegisterRaw("GET", "/debug/pprof", pprof.Index)
	gateway.router.RegisterRaw("GET", "/debug/pprof/cmdline", pprof.Cmdline)
	gateway.router.RegisterRaw("GET", "/debug/pprof/profile", pprof.Profile)
	gateway.router.RegisterRaw("GET", "/debug/pprof/symbol", pprof.Symbol)
	gateway.router.RegisterRaw("POST", "/debug/pprof/symbol", pprof.Symbol)
	gateway.router.RegisterRaw(
		"GET", "/debug/pprof/goroutine", pprof.Handler("goroutine").ServeHTTP,
	)
	gateway.router.RegisterRaw(
		"GET", "/debug/pprof/heap", pprof.Handler("heap").ServeHTTP,
	)
	gateway.router.RegisterRaw(
		"GET", "/debug/pprof/threadcreate",
		pprof.Handler("threadcreate").ServeHTTP,
	)
	gateway.router.RegisterRaw(
		"GET", "/debug/pprof/block", pprof.Handler("block").ServeHTTP,
	)

	register(gateway, gateway.router)
}

// Close the http server
func (gateway *Gateway) Close() {
	if gateway.loggerFile != nil {
		_ = gateway.loggerFile.Sync()
		_ = gateway.loggerFile.Close()
	}

	gateway.metricsBackend.Flush()
	_ = gateway.metricScopeCloser.Close()
	gateway.server.Close()
}

// Wait for gateway to close the server
func (gateway *Gateway) Wait() {
	gateway.WaitGroup.Wait()
}

func (gateway *Gateway) setupMetrics(config *StaticConfig) error {
	// TODO: decide what default tags we want...
	defaultTags := &map[string]string{}

	prefix := config.GetString("metrics.tally.service") +
		".production.all-workers"

	flushIntervalConfig := config.GetInt("metrics.tally.flushInterval")

	scope, closer := tally.NewCachedRootScope(
		prefix,
		*defaultTags,
		gateway.metricsBackend,
		time.Duration(flushIntervalConfig)*time.Millisecond,
		tally.DefaultSeparator,
	)
	gateway.MetricScope = scope
	gateway.metricScopeCloser = closer

	return nil
}

func (gateway *Gateway) setupLogger(config *StaticConfig) error {
	var output zap.Option
	tempLogger := zap.New(
		zap.NewJSONEncoder(),
		zap.Output(os.Stderr),
	)

	loggerFileName := config.GetString("logger.fileName")
	loggerOutput := config.GetString("logger.output")

	if loggerFileName == "" || loggerOutput == "stdout" {
		output = zap.Output(os.Stdout)
	} else {
		err := os.MkdirAll(filepath.Dir(loggerFileName), 0777)
		if err != nil {
			tempLogger.Error("Error creating log directory",
				zap.String("error", err.Error()),
			)
			return errors.Wrap(err, "Error creating log directory")
		}

		loggerFile, err := os.OpenFile(
			loggerFileName,
			os.O_APPEND|os.O_WRONLY|os.O_CREATE,
			0666,
		)
		if err != nil {
			tempLogger.Error("Error opening log file",
				zap.String("error", err.Error()),
			)
			return errors.Wrap(err, "Error opening log file")
		}
		gateway.loggerFile = loggerFile
		output = zap.Output(loggerFile)
	}

	// Default to a STDOUT logger
	gateway.Logger = zap.New(
		zap.NewJSONEncoder(
			zap.RFC3339Formatter("ts"),
		),
		output,
	)
	return nil
}

func (gateway *Gateway) setupHTTPServer() error {
	gateway.server = &HTTPServer{
		Server: &http.Server{
			Addr:    gateway.IP + ":" + strconv.FormatInt(int64(gateway.Port), 10),
			Handler: gateway.router,
		},
		Logger: gateway.Logger,
	}

	return nil
}

func (gateway *Gateway) setupTChannel(config *StaticConfig) error {
	tchannelServer, err := NewTChannelServer(
		&TChannelServerOptions{
			ServiceName: config.GetString("tchannel.serviceName"),
			ProcessName: config.GetString("tchannel.processName"),
		})

	if err != nil {
		return err
	}

	gateway.tchannelServer = tchannelServer
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"net"
	"net/http"
	"strings"
)

// headers is the http.Header used in each target request
// it is defined here to implement the flag.Value interface
// in order to support multiple identical flags for request header
// specification
type headers struct{ http.Header }

func (h headers) String() string {
	buf := &bytes.Buffer{}
	if err := h.Write(buf); err != nil {
		return ""
	}
	return buf.String()
}

// Set implements the flag.Value interface for a map of HTTP Headers.
func (h headers) Set(value string) error {
	parts := strings.SplitN(value, ":", 2)
	if len(parts) != 2 {
		return fmt.Errorf("header '%s' has a wrong format", value)
	}
	key, val := strings.TrimSpace(parts[0]), strings.TrimSpace(parts[1])
	if key == "" || val == "" {
		return fmt.Errorf("header '%s' has a wrong format", value)
	}
	// Add key/value directly to the http.Header (map[string][]string).
	// http.Header.Add() cannonicalizes keys but vegeta is used
	// to test systems that require case-sensitive headers.
	h.Header[key] = append(h.Header[key], val)
	return nil
}

// localAddr implements the Flag interface for parsing net.IPAddr
type localAddr struct{ *net.IPAddr }

func (ip *localAddr) Set(value string) (err error) {
	ip.IPAddr, err = net.ResolveIPAddr("ip", value)
	return
}

// csl implements the flag.Value interface for comma separated lists
type csl []string

func (l *csl) Set(v string) error {
	*l = strings.Split(v, ",")
	return nil
}

func (l csl) String() string { return strings.Join(l, ",") }
<commit_msg>Fix spelling mistake in flags.go<commit_after>package main

import (
	"bytes"
	"fmt"
	"net"
	"net/http"
	"strings"
)

// headers is the http.Header used in each target request
// it is defined here to implement the flag.Value interface
// in order to support multiple identical flags for request header
// specification
type headers struct{ http.Header }

func (h headers) String() string {
	buf := &bytes.Buffer{}
	if err := h.Write(buf); err != nil {
		return ""
	}
	return buf.String()
}

// Set implements the flag.Value interface for a map of HTTP Headers.
func (h headers) Set(value string) error {
	parts := strings.SplitN(value, ":", 2)
	if len(parts) != 2 {
		return fmt.Errorf("header '%s' has a wrong format", value)
	}
	key, val := strings.TrimSpace(parts[0]), strings.TrimSpace(parts[1])
	if key == "" || val == "" {
		return fmt.Errorf("header '%s' has a wrong format", value)
	}
	// Add key/value directly to the http.Header (map[string][]string).
	// http.Header.Add() canonicalizes keys but vegeta is used
	// to test systems that require case-sensitive headers.
	h.Header[key] = append(h.Header[key], val)
	return nil
}

// localAddr implements the Flag interface for parsing net.IPAddr
type localAddr struct{ *net.IPAddr }

func (ip *localAddr) Set(value string) (err error) {
	ip.IPAddr, err = net.ResolveIPAddr("ip", value)
	return
}

// csl implements the flag.Value interface for comma separated lists
type csl []string

func (l *csl) Set(v string) error {
	*l = strings.Split(v, ",")
	return nil
}

func (l csl) String() string { return strings.Join(l, ",") }
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"
	"os/signal"
	"strings"
	"syscall"

	"github.com/coreos/fleet/third_party/github.com/golang/glog"
	"github.com/coreos/fleet/third_party/github.com/rakyll/globalconf"

	"github.com/coreos/fleet/agent"
	"github.com/coreos/fleet/config"
	"github.com/coreos/fleet/server"
	"github.com/coreos/fleet/sign"
	"github.com/coreos/fleet/version"
)

const (
	DefaultConfigFile = "/etc/fleet/fleet.conf"
)

func main() {
	// We use a FlagSets since glog adds a bunch of flags we do not want to publish
	userset := flag.NewFlagSet("fleet", flag.ExitOnError)
	printVersion := userset.Bool("version", false, "Print the version and exit")
	cfgPath := userset.String("config", "", fmt.Sprintf("Path to config file. Fleet will look for a config at %s by default.", DefaultConfigFile))

	// Initialize logging so we have it set up while parsing config information
	config.UpdateLoggingFlagsFromConfig(flag.CommandLine, &config.Config{})

	err := userset.Parse(os.Args[1:])
	if err == flag.ErrHelp {
		userset.Usage()
		syscall.Exit(1)
	}

	if *printVersion {
		fmt.Println("fleet version", version.Version)
		os.Exit(0)
	}

	cfgset := flag.NewFlagSet("fleet", flag.ExitOnError)
	cfgset.Int("verbosity", 0, "Logging level")
	cfgset.Var(&stringSlice{}, "etcd_servers", "List of etcd endpoints")
	cfgset.String("boot_id", "", "Override default BootID of fleet machine")
	cfgset.String("public_ip", "", "IP address that fleet machine should publish")
	cfgset.String("metadata", "", "List of key-value metadata to assign to the fleet machine")
	cfgset.String("unit_prefix", "", "Prefix that should be used for all systemd units")
	cfgset.String("agent_ttl", agent.DefaultTTL, "TTL in seconds of fleet machine state in etcd")
	cfgset.Bool("verify_units", true, "Verify unit file signatures using local SSH identities")
	cfgset.String("authorized_key_file", sign.DefaultAuthorizedKeyFile, "File that contains authorized keys to be used for signature verification")

	globalconf.Register("", cfgset)
	cfg, err := getConfig(cfgset, *cfgPath)
	if err != nil {
		glog.Error(err.Error())
		syscall.Exit(1)
	}

	srv := server.New(*cfg)
	srv.Run()

	reconfigure := func() {
		glog.Infof("Reloading configuration from %s", *cfgPath)

		cfg, err := getConfig(cfgset, *cfgPath)
		if err != nil {
			glog.Errorf(err.Error())
			syscall.Exit(1)
		}

		srv.Stop()

		srv = server.New(*cfg)
		srv.Run()
	}

	shutdown := func() {
		glog.Infof("Gracefully shutting down")
		srv.Stop()
		srv.Purge()
		syscall.Exit(0)
	}

	signals := map[os.Signal]func(){
		syscall.SIGHUP:  reconfigure,
		syscall.SIGTERM: shutdown,
		syscall.SIGINT:  shutdown,
	}

	listenForSignals(signals)
}

func getConfig(flagset *flag.FlagSet, userCfgFile string) (*config.Config, error) {
	opts := globalconf.Options{EnvPrefix: "FLEET_"}

	if userCfgFile != "" {
		// Fail hard if a user-provided config is not usable
		if _, err := os.Stat(userCfgFile); err != nil {
			glog.Errorf("Unable to use config file %s: %v", userCfgFile, err)
			os.Exit(1)
		}

		glog.Infof("Using provided config file %s", userCfgFile)
		opts.Filename = userCfgFile

	} else if _, err := os.Stat(DefaultConfigFile); err == nil {
		glog.Infof("Using default config file %s", DefaultConfigFile)
		opts.Filename = DefaultConfigFile
	} else {
		glog.Infof("Continuing without config file")
	}

	gconf, err := globalconf.NewWithOptions(&opts)
	if err != nil {
		return nil, err
	}

	gconf.ParseSet("", flagset)

	cfg := config.Config{
		Verbosity: (*flagset.Lookup("verbosity")).Value.(flag.Getter).Get().(int),
		EtcdServers: (*flagset.Lookup("etcd_servers")).Value.(flag.Getter).Get().(stringSlice),
		BootId: (*flagset.Lookup("boot_id")).Value.(flag.Getter).Get().(string),
		PublicIP: (*flagset.Lookup("public_ip")).Value.(flag.Getter).Get().(string),
		RawMetadata: (*flagset.Lookup("metadata")).Value.(flag.Getter).Get().(string),
		UnitPrefix: (*flagset.Lookup("unit_prefix")).Value.(flag.Getter).Get().(string),
		AgentTTL: (*flagset.Lookup("agent_ttl")).Value.(flag.Getter).Get().(string),
		VerifyUnits: (*flagset.Lookup("verify_units")).Value.(flag.Getter).Get().(bool),
		AuthorizedKeyFile: (*flagset.Lookup("authorized_key_file")).Value.(flag.Getter).Get().(string),
	}

	config.UpdateLoggingFlagsFromConfig(flag.CommandLine, &cfg)

	return &cfg, nil
}

func listenForSignals(sigmap map[os.Signal]func()) {
	sigchan := make(chan os.Signal, 1)

	for k, _ := range sigmap {
		signal.Notify(sigchan, k)
	}

	for true {
		sig := <-sigchan
		handler, ok := sigmap[sig]
		if ok {
			handler()
		}
	}
}

type stringSlice []string

func (f *stringSlice) Set(value string) error {
	for _, item := range strings.Split(value, ",") {
		item = strings.TrimLeft(item, " [\"")
		item = strings.TrimRight(item, " \"]")
		*f = append(*f, item)
	}

	return nil
}

func (f *stringSlice) String() string {
	return fmt.Sprintf("%v", *f)
}

func (f *stringSlice) Value() []string {
	return *f
}

func (f *stringSlice) Get() interface{} {
	return *f
}
<commit_msg>chore(fleet): set verify_units to be false as default<commit_after>package main

import (
	"flag"
	"fmt"
	"os"
	"os/signal"
	"strings"
	"syscall"

	"github.com/coreos/fleet/third_party/github.com/golang/glog"
	"github.com/coreos/fleet/third_party/github.com/rakyll/globalconf"

	"github.com/coreos/fleet/agent"
	"github.com/coreos/fleet/config"
	"github.com/coreos/fleet/server"
	"github.com/coreos/fleet/sign"
	"github.com/coreos/fleet/version"
)

const (
	DefaultConfigFile = "/etc/fleet/fleet.conf"
)

func main() {
	// We use a FlagSets since glog adds a bunch of flags we do not want to publish
	userset := flag.NewFlagSet("fleet", flag.ExitOnError)
	printVersion := userset.Bool("version", false, "Print the version and exit")
	cfgPath := userset.String("config", "", fmt.Sprintf("Path to config file. Fleet will look for a config at %s by default.", DefaultConfigFile))

	// Initialize logging so we have it set up while parsing config information
	config.UpdateLoggingFlagsFromConfig(flag.CommandLine, &config.Config{})

	err := userset.Parse(os.Args[1:])
	if err == flag.ErrHelp {
		userset.Usage()
		syscall.Exit(1)
	}

	if *printVersion {
		fmt.Println("fleet version", version.Version)
		os.Exit(0)
	}

	cfgset := flag.NewFlagSet("fleet", flag.ExitOnError)
	cfgset.Int("verbosity", 0, "Logging level")
	cfgset.Var(&stringSlice{}, "etcd_servers", "List of etcd endpoints")
	cfgset.String("boot_id", "", "Override default BootID of fleet machine")
	cfgset.String("public_ip", "", "IP address that fleet machine should publish")
	cfgset.String("metadata", "", "List of key-value metadata to assign to the fleet machine")
	cfgset.String("unit_prefix", "", "Prefix that should be used for all systemd units")
	cfgset.String("agent_ttl", agent.DefaultTTL, "TTL in seconds of fleet machine state in etcd")
	cfgset.Bool("verify_units", false, "Verify unit file signatures using local SSH identities")
	cfgset.String("authorized_key_file", sign.DefaultAuthorizedKeyFile, "File that contains authorized keys to be used for signature verification")

	globalconf.Register("", cfgset)
	cfg, err := getConfig(cfgset, *cfgPath)
	if err != nil {
		glog.Error(err.Error())
		syscall.Exit(1)
	}

	srv := server.New(*cfg)
	srv.Run()

	reconfigure := func() {
		glog.Infof("Reloading configuration from %s", *cfgPath)

		cfg, err := getConfig(cfgset, *cfgPath)
		if err != nil {
			glog.Errorf(err.Error())
			syscall.Exit(1)
		}

		srv.Stop()

		srv = server.New(*cfg)
		srv.Run()
	}

	shutdown := func() {
		glog.Infof("Gracefully shutting down")
		srv.Stop()
		srv.Purge()
		syscall.Exit(0)
	}

	signals := map[os.Signal]func(){
		syscall.SIGHUP:  reconfigure,
		syscall.SIGTERM: shutdown,
		syscall.SIGINT:  shutdown,
	}

	listenForSignals(signals)
}

func getConfig(flagset *flag.FlagSet, userCfgFile string) (*config.Config, error) {
	opts := globalconf.Options{EnvPrefix: "FLEET_"}

	if userCfgFile != "" {
		// Fail hard if a user-provided config is not usable
		if _, err := os.Stat(userCfgFile); err != nil {
			glog.Errorf("Unable to use config file %s: %v", userCfgFile, err)
			os.Exit(1)
		}

		glog.Infof("Using provided config file %s", userCfgFile)
		opts.Filename = userCfgFile

	} else if _, err := os.Stat(DefaultConfigFile); err == nil {
		glog.Infof("Using default config file %s", DefaultConfigFile)
		opts.Filename = DefaultConfigFile
	} else {
		glog.Infof("Continuing without config file")
	}

	gconf, err := globalconf.NewWithOptions(&opts)
	if err != nil {
		return nil, err
	}

	gconf.ParseSet("", flagset)

	cfg := config.Config{
		Verbosity: (*flagset.Lookup("verbosity")).Value.(flag.Getter).Get().(int),
		EtcdServers: (*flagset.Lookup("etcd_servers")).Value.(flag.Getter).Get().(stringSlice),
		BootId: (*flagset.Lookup("boot_id")).Value.(flag.Getter).Get().(string),
		PublicIP: (*flagset.Lookup("public_ip")).Value.(flag.Getter).Get().(string),
		RawMetadata: (*flagset.Lookup("metadata")).Value.(flag.Getter).Get().(string),
		UnitPrefix: (*flagset.Lookup("unit_prefix")).Value.(flag.Getter).Get().(string),
		AgentTTL: (*flagset.Lookup("agent_ttl")).Value.(flag.Getter).Get().(string),
		VerifyUnits: (*flagset.Lookup("verify_units")).Value.(flag.Getter).Get().(bool),
		AuthorizedKeyFile: (*flagset.Lookup("authorized_key_file")).Value.(flag.Getter).Get().(string),
	}

	config.UpdateLoggingFlagsFromConfig(flag.CommandLine, &cfg)

	return &cfg, nil
}

func listenForSignals(sigmap map[os.Signal]func()) {
	sigchan := make(chan os.Signal, 1)

	for k, _ := range sigmap {
		signal.Notify(sigchan, k)
	}

	for true {
		sig := <-sigchan
		handler, ok := sigmap[sig]
		if ok {
			handler()
		}
	}
}

type stringSlice []string

func (f *stringSlice) Set(value string) error {
	for _, item := range strings.Split(value, ",") {
		item = strings.TrimLeft(item, " [\"")
		item = strings.TrimRight(item, " \"]")
		*f = append(*f, item)
	}

	return nil
}

func (f *stringSlice) String() string {
	return fmt.Sprintf("%v", *f)
}

func (f *stringSlice) Value() []string {
	return *f
}

func (f *stringSlice) Get() interface{} {
	return *f
}
<|endoftext|>


<commit_before>//
// An interactive shell for DynamoDB
//
package main

import (
	"../../dynago"

	"code.google.com/p/gcfg"
	"github.com/gobs/args"
	"github.com/gobs/cmd"
	"github.com/gobs/httpclient"
	"github.com/gobs/pretty"

	"fmt"
	"log"
	"os"
	"path"
	"strconv"
	"strings"
)

const (
	CONFIG_FILE  = ".dynagorc"
	HISTORY_FILE = ".dynago_history"
)

// the configuration should look like the following
// (with multiple profiles and a selected one)
//
// [dynago]
// profile=xxx
//
// [profile "xxx"]
// region=us-west-1
// accessKey=XXXXXXXX
// secretKey=YYYYYYYY

type Config struct {
	Dynago struct {
		// define default profile
		Profile string
		// enable request debugging
		Debug bool
	}

	// list of named profiles
	Profile map[string]*struct {
		Region    string
		AccessKey string
		SecretKey string
	}
}

// Look for configFile in current directory or home directory.
// No configuration file is NOT an error.
// A malformed configuration file is a FATAL error.

func ReadConfig(configFile string, config *Config) *Config {
	if config == nil {
		config = &Config{}
	}

	// configFile in current directory or full path
	if _, err := os.Stat(configFile); err != nil {
		if strings.Contains(configFile, "/") {
			return config
		}

		// configFile in home directory
		configFile = path.Join(os.Getenv("HOME"), configFile)
		if _, err := os.Stat(configFile); err != nil {
			return config
		}
	}

	err := gcfg.ReadFileInto(config, configFile)
	if err != nil {
		log.Fatal(err)
	}

	return config
}

var (
	// this hold the current list of table names, to be used by the CompletionFunction
	table_list []string
)

func add_to_list(table string) {
	table_list = append(table_list, table)
}

func remove_from_list(table string) {
	for i, t := range table_list {
		if t == table {
			table_list = append(table_list[:i], table_list[i+1:]...)
			return
		}
	}
}

// return list of table names that match the input pattern (table name starts with "text")
func CompletionFunction(text string, line string, start, stop int) []string {
	if len(table_list) > 0 {
		matches := make([]string, 0, len(table_list))

		for _, w := range table_list {
			if strings.HasPrefix(w, text) {
				matches = append(matches, w)
			}
		}

		return matches
	}

	return nil
}

func main() {
	var nextKey dynago.AttributeNameValue
	var selectedTable *dynago.TableInstance

	config := ReadConfig(CONFIG_FILE, nil)
	selected := config.Dynago.Profile

	if len(os.Args) > 1 {
		// there is at least one parameter:
		// override the selected profile
		selected = os.Args[1]
	}

	profile := config.Profile[selected]
	if profile == nil {
		log.Fatal("no profile for ", selected)
	}

	if config.Dynago.Debug {
		httpclient.StartLogging(true, true)
	}

	db := dynago.NewDBClient()

	if len(profile.Region) > 0 {
		db.WithRegion(profile.Region)
	}

	if len(profile.AccessKey) > 0 {
		db.WithCredentials(profile.AccessKey, profile.SecretKey)
	}

	commander := &cmd.Cmd{HistoryFile: HISTORY_FILE, Complete: CompletionFunction, EnableShell: true}
	commander.Prompt = "dynagosh> "
	commander.Init()

	commander.Add(cmd.Command{"config",
		`
		config : display current configuration
		`,
		func(string) (stop bool) {
			pretty.PrettyPrint(config)
			return
		}})

	commander.Add(cmd.Command{"list",
		`
                list : display list of available tables
                `,
		func(string) (stop bool) {
			tables, err := db.ListTables()

			if err != nil {
				fmt.Println(err)
				return
			}

			fmt.Println("Available tables")

			for _, tableName := range tables {
				fmt.Println("  ", tableName)
			}

			if len(tables) > 0 {
				table_list = tables
			}

			return
		}})

	commander.Add(cmd.Command{"describe",
		`
                describe {table} : display table configuration
                `,
		func(line string) (stop bool) {
			tableName := line
			table, err := db.DescribeTable(tableName)
			if err != nil {
				fmt.Println(err)
			} else {
				pretty.PrettyPrint(table)
			}

			return
		}})

	commander.Add(cmd.Command{"use",
		`
                use {table} : select table for queries
                `,
		func(line string) (stop bool) {
			tableName := line
			table, err := db.GetTable(tableName)
			if err != nil {
				fmt.Println(err)
			} else {
				selectedTable = table
				commander.Prompt = "dynagosh: " + tableName + "> "
			}

			return
		}})

	commander.Add(cmd.Command{"create",
		`
		create {tablename} hashKey:hashType [rangeKey:rangeType] [readCapacity] [writeCapacity]
		`,
		func(line string) (stop bool) {
			args := args.GetArgs(line)

			if len(args) < 2 {
				fmt.Println("not enough arguments")
				return
			}

			tableName := args[0]

			hashKey := &dynago.AttributeDefinition{AttributeType: dynago.STRING_ATTRIBUTE}
			var rangeKey *dynago.AttributeDefinition
			rc := 5
			wc := 5

			if strings.Contains(args[1], ":") {
				parts := strings.Split(args[1], ":")
				hashKey.AttributeName = parts[0]
				hashKey.AttributeType = parts[1]
			} else {
				hashKey.AttributeName = args[1]
			}

			if len(args) > 2 {
				rangeKey := &dynago.AttributeDefinition{AttributeType: dynago.STRING_ATTRIBUTE}

				if strings.Contains(args[2], ":") {
					parts := strings.Split(args[2], ":")
					rangeKey.AttributeName = parts[0]
					rangeKey.AttributeType = parts[1]
				} else {
					rangeKey.AttributeName = args[2]
				}
			}

			if table, err := db.CreateTable(tableName, hashKey, rangeKey, rc, wc); err != nil {
				fmt.Println(err)
			} else {
				pretty.PrettyPrint(table)
				add_to_list(tableName)
			}

			return
		}})

	commander.Add(cmd.Command{"delete",
		`
                delete {table} : delete table
                `,
		func(line string) (stop bool) {
			tableName := line
			table, err := db.DeleteTable(tableName)
			if err != nil {
				fmt.Println(err)
			} else {
				pretty.PrettyPrint(table)
				remove_from_list(tableName)
			}

			return
		}})

	commander.Add(cmd.Command{"update",
		`
		update {tablename} readCapacity writeCapacity
		`,
		func(line string) (stop bool) {
			args := args.GetArgs(line)

			if len(args) < 2 {
				fmt.Println("not enough arguments")
				return
			}

			tableName := args[0]
			table, err := db.DescribeTable(tableName)

			if err != nil {
				fmt.Println(err)
				return
			}

			rc := -1 // table.ProvisionedThroughput.ReadCapacityUnits
			wc := -1 // table.ProvisionedThroughput.WriteCapacityUnits

			if v, err := strconv.Atoi(args[1]); err == nil {
				rc = v
			}

			if len(args) > 2 {
				if v, err := strconv.Atoi(args[2]); err == nil {
					wc = v
				}
			}

			if rc <= 0 && wc <= 0 {
				fmt.Println("no valid value for rc or wc")
				return
			}

			if rc <= 0 {
				rc = table.ProvisionedThroughput.ReadCapacityUnits
			}

			if wc <= 0 {
				wc = table.ProvisionedThroughput.WriteCapacityUnits
			}

			if table, err := db.UpdateTable(tableName, rc, wc); err != nil {
				fmt.Println(err)
			} else {
				pretty.PrettyPrint(table)
			}

			return
		}})

	commander.Add(cmd.Command{"get",
		`
		get {tablename} {hashKey} [rangeKey] [attributes]
		`,
		func(line string) (stop bool) {
			args := args.GetArgs(line)

			if len(args) < 2 {
				fmt.Println("not enough arguments")
				return
			}

			tableName := args[0]
			table, err := db.GetTable(tableName)

			if err != nil {
				fmt.Println(err)
				return
			}

			hashKey := args[1]
			var rangeKey string

			if len(args) > 2 {
				rangeKey = args[2]
			}

			var attributes []string

			if len(args) > 3 {
				attributes = args[3:]
			}

			if item, consumed, err := table.GetItem(hashKey, rangeKey, attributes, false, true); err != nil {
				fmt.Println(err)
			} else {
				pretty.PrettyPrint(item)
				fmt.Println("consumed:", consumed)
			}

			return
		}})

	commander.Add(cmd.Command{"query",
		`
		query [--table=tablename] [--limit=pagesize] [--next] [--count] [--consumed] hash-key-value range-key-value
		`,
		func(line string) (stop bool) {
			flags := args.NewFlags("query")

			tableName := flags.String("table", "", "table name")
			limit := flags.Int("limit", 0, "maximum number of items per page")
			count := flags.Bool("count", false, "only return item count")
			next := flags.Bool("next", false, "get next page")
			consumed := flags.Bool("consumed", false, "return consumed capacity")

			args.ParseFlags(flags, line)
			args := flags.Args()

			table := selectedTable

			if len(*tableName) > 1 {
				if t, err := db.GetTable(*tableName); err != nil {
					fmt.Println(err)
					return
				} else {
					table = t
				}
			} else if table == nil {
				fmt.Println("no table selected")
				return
			}

			if len(args) < 1 {
				fmt.Println("not enough arguments")
				return
			}

			query := table.Query(args[0])

			if len(args) > 1 {
				query = query.WithAttrCondition(table.RangeKey().EQ(args[1]))
			}

			if *limit > 0 {
				query = query.WithLimit(*limit)
			}

			if *count {
				query = query.WithSelect(dynago.SELECT_COUNT)
			}

			if *next {
				query = query.WithStartKey(nextKey)
			}

			if *consumed {
				query = query.WithConsumed(true)
			}

			if items, lastKey, consumed, err := query.Exec(nil); err != nil {
				fmt.Println(err)
			} else {
				pretty.PrettyPrint(items)
				fmt.Println("consumed:", consumed)

				nextKey = lastKey
			}

			return
		}})

	commander.Add(cmd.Command{"scan",
		`
		scan [--table=tablename] [--limit=pagesize] [--next] [--count] [--consumed] [--segment=n --total=m]
		`,
		func(line string) (stop bool) {
			flags := args.NewFlags("scan")

			tableName := flags.String("table", "", "table name")
			limit := flags.Int("limit", 0, "maximum number of items per page")
			count := flags.Bool("count", false, "only return item count")
			next := flags.Bool("next", false, "get next page")
			consumed := flags.Bool("consumed", false, "return consumed capacity")
			segment := flags.Int("segment", 0, "segment number")
			total := flags.Int("total", 0, "total segment")

			args.ParseFlags(flags, line)
			//args := flags.Args()

			table := selectedTable

			if len(*tableName) > 1 {
				if t, err := db.GetTable(*tableName); err != nil {
					fmt.Println(err)
					return
				} else {
					table = t
				}
			} else if table == nil {
				fmt.Println("no table selected")
				return
			}

			scan := dynago.ScanTable(table).
				WithSegment(*segment, *total)

			if *limit > 0 {
				scan = scan.WithLimit(*limit)
			}

			if *count {
				scan = scan.WithSelect(dynago.SELECT_COUNT)
			}

			if *next {
				scan = scan.WithStartKey(nextKey)
			}

			if *consumed {
				scan = scan.WithConsumed(true)
			}

			if items, lastKey, consumed, err := scan.Exec(db); err != nil {
				fmt.Println(err)
			} else {
				pretty.PrettyPrint(items)
				fmt.Println("consumed:", consumed)

				nextKey = lastKey
			}

			return
		}})

	commander.Commands["ls"] = commander.Commands["list"]
	commander.Commands["drop"] = commander.Commands["delete"]

	commander.CmdLoop()
}
<commit_msg>- added --hash hash-value parameter - added some --range-{operator} parameters (eq, ne, le, lt, ge, gt)<commit_after>//
// An interactive shell for DynamoDB
//
package main

import (
	"../../dynago"

	"code.google.com/p/gcfg"
	"github.com/gobs/args"
	"github.com/gobs/cmd"
	"github.com/gobs/httpclient"
	"github.com/gobs/pretty"

	"errors"
	"flag"
	"fmt"
	"log"
	"os"
	"path"
	"strconv"
	"strings"
)

const (
	CONFIG_FILE  = ".dynagorc"
	HISTORY_FILE = ".dynago_history"
)

// the configuration should look like the following
// (with multiple profiles and a selected one)
//
// [dynago]
// profile=xxx
//
// [profile "xxx"]
// region=us-west-1
// accessKey=XXXXXXXX
// secretKey=YYYYYYYY

type Config struct {
	Dynago struct {
		// define default profile
		Profile string
		// enable request debugging
		Debug bool
	}

	// list of named profiles
	Profile map[string]*struct {
		Region    string
		AccessKey string
		SecretKey string
	}
}

// Look for configFile in current directory or home directory.
// No configuration file is NOT an error.
// A malformed configuration file is a FATAL error.

func ReadConfig(configFile string, config *Config) *Config {
	if config == nil {
		config = &Config{}
	}

	// configFile in current directory or full path
	if _, err := os.Stat(configFile); err != nil {
		if strings.Contains(configFile, "/") {
			return config
		}

		// configFile in home directory
		configFile = path.Join(os.Getenv("HOME"), configFile)
		if _, err := os.Stat(configFile); err != nil {
			return config
		}
	}

	err := gcfg.ReadFileInto(config, configFile)
	if err != nil {
		log.Fatal(err)
	}

	return config
}

var (
	// this hold the current list of table names, to be used by the CompletionFunction
	table_list []string
)

func add_to_list(table string) {
	table_list = append(table_list, table)
}

func remove_from_list(table string) {
	for i, t := range table_list {
		if t == table {
			table_list = append(table_list[:i], table_list[i+1:]...)
			return
		}
	}
}

// return list of table names that match the input pattern (table name starts with "text")
func CompletionFunction(text string, line string, start, stop int) []string {
	if len(table_list) > 0 {
		matches := make([]string, 0, len(table_list))

		for _, w := range table_list {
			if strings.HasPrefix(w, text) {
				matches = append(matches, w)
			}
		}

		return matches
	}

	return nil
}

type RangeCondition struct {
	Operator string
	Value    string
}

func (cond *RangeCondition) Set(value string) error {
	if len(cond.Value) > 0 {
		return errors.New("range-condition value already set")
	}

	cond.Value = value
	return nil
}

func (cond *RangeCondition) String() string {
	return fmt.Sprintf("%s: %v", cond.Operator, cond.Value)
}

func (cond *RangeCondition) SetOperator(f *flag.Flag) {
	if f.Name == "range" {
		cond.Operator = "EQ"
	} else if strings.HasPrefix(f.Name, "range-") {
		cond.Operator = strings.ToUpper(f.Name[6:])
	}
}

func main() {
	var nextKey dynago.AttributeNameValue
	var selectedTable *dynago.TableInstance

	config := ReadConfig(CONFIG_FILE, nil)
	selected := config.Dynago.Profile

	if len(os.Args) > 1 {
		// there is at least one parameter:
		// override the selected profile
		selected = os.Args[1]
	}

	profile := config.Profile[selected]
	if profile == nil {
		log.Fatal("no profile for ", selected)
	}

	if config.Dynago.Debug {
		httpclient.StartLogging(true, true)
	}

	db := dynago.NewDBClient()

	if len(profile.Region) > 0 {
		db.WithRegion(profile.Region)
	}

	if len(profile.AccessKey) > 0 {
		db.WithCredentials(profile.AccessKey, profile.SecretKey)
	}

	commander := &cmd.Cmd{HistoryFile: HISTORY_FILE, Complete: CompletionFunction, EnableShell: true}
	commander.Prompt = "dynagosh> "
	commander.Init()

	commander.Add(cmd.Command{"config",
		`
		config : display current configuration
		`,
		func(string) (stop bool) {
			pretty.PrettyPrint(config)
			return
		}})

	commander.Add(cmd.Command{"list",
		`
                list : display list of available tables
                `,
		func(string) (stop bool) {
			tables, err := db.ListTables()

			if err != nil {
				fmt.Println(err)
				return
			}

			fmt.Println("Available tables")

			for _, tableName := range tables {
				fmt.Println("  ", tableName)
			}

			if len(tables) > 0 {
				table_list = tables
			}

			return
		}})

	commander.Add(cmd.Command{"describe",
		`
                describe {table} : display table configuration
                `,
		func(line string) (stop bool) {
			tableName := line
			table, err := db.DescribeTable(tableName)
			if err != nil {
				fmt.Println(err)
			} else {
				pretty.PrettyPrint(table)
			}

			return
		}})

	commander.Add(cmd.Command{"use",
		`
                use {table} : select table for queries
                `,
		func(line string) (stop bool) {
			tableName := line
			table, err := db.GetTable(tableName)
			if err != nil {
				fmt.Println(err)
			} else {
				selectedTable = table
				commander.Prompt = "dynagosh: " + tableName + "> "
			}

			return
		}})

	commander.Add(cmd.Command{"create",
		`
		create {tablename} hashKey:hashType [rangeKey:rangeType] [readCapacity] [writeCapacity]
		`,
		func(line string) (stop bool) {
			args := args.GetArgs(line)

			if len(args) < 2 {
				fmt.Println("not enough arguments")
				return
			}

			tableName := args[0]

			hashKey := &dynago.AttributeDefinition{AttributeType: dynago.STRING_ATTRIBUTE}
			var rangeKey *dynago.AttributeDefinition
			rc := 5
			wc := 5

			if strings.Contains(args[1], ":") {
				parts := strings.Split(args[1], ":")
				hashKey.AttributeName = parts[0]
				hashKey.AttributeType = parts[1]
			} else {
				hashKey.AttributeName = args[1]
			}

			if len(args) > 2 {
				rangeKey := &dynago.AttributeDefinition{AttributeType: dynago.STRING_ATTRIBUTE}

				if strings.Contains(args[2], ":") {
					parts := strings.Split(args[2], ":")
					rangeKey.AttributeName = parts[0]
					rangeKey.AttributeType = parts[1]
				} else {
					rangeKey.AttributeName = args[2]
				}
			}

			if table, err := db.CreateTable(tableName, hashKey, rangeKey, rc, wc); err != nil {
				fmt.Println(err)
			} else {
				pretty.PrettyPrint(table)
				add_to_list(tableName)
			}

			return
		}})

	commander.Add(cmd.Command{"delete",
		`
                delete {table} : delete table
                `,
		func(line string) (stop bool) {
			tableName := line
			table, err := db.DeleteTable(tableName)
			if err != nil {
				fmt.Println(err)
			} else {
				pretty.PrettyPrint(table)
				remove_from_list(tableName)
			}

			return
		}})

	commander.Add(cmd.Command{"update",
		`
		update {tablename} readCapacity writeCapacity
		`,
		func(line string) (stop bool) {
			args := args.GetArgs(line)

			if len(args) < 2 {
				fmt.Println("not enough arguments")
				return
			}

			tableName := args[0]
			table, err := db.DescribeTable(tableName)

			if err != nil {
				fmt.Println(err)
				return
			}

			rc := -1 // table.ProvisionedThroughput.ReadCapacityUnits
			wc := -1 // table.ProvisionedThroughput.WriteCapacityUnits

			if v, err := strconv.Atoi(args[1]); err == nil {
				rc = v
			}

			if len(args) > 2 {
				if v, err := strconv.Atoi(args[2]); err == nil {
					wc = v
				}
			}

			if rc <= 0 && wc <= 0 {
				fmt.Println("no valid value for rc or wc")
				return
			}

			if rc <= 0 {
				rc = table.ProvisionedThroughput.ReadCapacityUnits
			}

			if wc <= 0 {
				wc = table.ProvisionedThroughput.WriteCapacityUnits
			}

			if table, err := db.UpdateTable(tableName, rc, wc); err != nil {
				fmt.Println(err)
			} else {
				pretty.PrettyPrint(table)
			}

			return
		}})

	commander.Add(cmd.Command{"get",
		`
		get {tablename} {hashKey} [rangeKey] [attributes]
		`,
		func(line string) (stop bool) {
			args := args.GetArgs(line)

			if len(args) < 2 {
				fmt.Println("not enough arguments")
				return
			}

			tableName := args[0]
			table, err := db.GetTable(tableName)

			if err != nil {
				fmt.Println(err)
				return
			}

			hashKey := args[1]
			var rangeKey string

			if len(args) > 2 {
				rangeKey = args[2]
			}

			var attributes []string

			if len(args) > 3 {
				attributes = args[3:]
			}

			if item, consumed, err := table.GetItem(hashKey, rangeKey, attributes, false, true); err != nil {
				fmt.Println(err)
			} else {
				pretty.PrettyPrint(item)
				fmt.Println("consumed:", consumed)
			}

			return
		}})

	commander.Add(cmd.Command{"query",
		`
		query [--table=tablename] [--limit=pagesize] [--next] [--count] [--consumed] --hash hash-key-value [--range range-key-value]
		`,
		func(line string) (stop bool) {
			flags := args.NewFlags("query")

			tableName := flags.String("table", "", "table name")
			limit := flags.Int("limit", 0, "maximum number of items per page")
			count := flags.Bool("count", false, "only return item count")
			next := flags.Bool("next", false, "get next page")
			consumed := flags.Bool("consumed", false, "return consumed capacity")

			hashKey := flags.String("hash", "", "hash-key value")

			var rangeCond RangeCondition

			flags.Var(&rangeCond, "range", "range-key value")
			flags.Var(&rangeCond, "range-eq", "range-key equal value")
			flags.Var(&rangeCond, "range-ne", "range-key not-equal value")
			flags.Var(&rangeCond, "range-le", "range-key less-or-equal value")
			flags.Var(&rangeCond, "range-lt", "range-key less-than value")
			flags.Var(&rangeCond, "range-ge", "range-key less-or-equal value")
			flags.Var(&rangeCond, "range-gt", "range-key less-than value")

			args.ParseFlags(flags, line)
			args := flags.Args()

			flags.Visit(rangeCond.SetOperator)

			table := selectedTable

			if len(*tableName) > 1 {
				if t, err := db.GetTable(*tableName); err != nil {
					fmt.Println(err)
					return
				} else {
					table = t
				}
			} else if table == nil {
				fmt.Println("no table selected")
				return
			}

			if len(*hashKey) < 1 {
				if len(args) < 1 {
					fmt.Println("not enough arguments")
					return
				}

				*hashKey = args[0]

				if len(rangeCond.Operator) < 1 && len(args) > 1 {
					rangeCond.Operator = "EQ"
					rangeCond.Value = args[1]
				}
			}

			query := table.Query(*hashKey)

			if len(rangeCond.Operator) > 0 {
				query = query.WithAttrCondition(table.RangeKey().Condition(rangeCond.Operator, rangeCond.Value))
			}

			if *limit > 0 {
				query = query.WithLimit(*limit)
			}

			if *count {
				query = query.WithSelect(dynago.SELECT_COUNT)
			}

			if *next {
				query = query.WithStartKey(nextKey)
			}

			if *consumed {
				query = query.WithConsumed(true)
			}

			if items, lastKey, consumed, err := query.Exec(nil); err != nil {
				fmt.Println(err)
			} else {
				pretty.PrettyPrint(items)
				fmt.Println("consumed:", consumed)

				nextKey = lastKey
			}

			return
		}})

	commander.Add(cmd.Command{"scan",
		`
		scan [--table=tablename] [--limit=pagesize] [--next] [--count] [--consumed] [--segment=n --total=m]
		`,
		func(line string) (stop bool) {
			flags := args.NewFlags("scan")

			tableName := flags.String("table", "", "table name")
			limit := flags.Int("limit", 0, "maximum number of items per page")
			count := flags.Bool("count", false, "only return item count")
			next := flags.Bool("next", false, "get next page")
			consumed := flags.Bool("consumed", false, "return consumed capacity")
			segment := flags.Int("segment", 0, "segment number")
			total := flags.Int("total", 0, "total segment")

			args.ParseFlags(flags, line)
			//args := flags.Args()

			table := selectedTable

			if len(*tableName) > 1 {
				if t, err := db.GetTable(*tableName); err != nil {
					fmt.Println(err)
					return
				} else {
					table = t
				}
			} else if table == nil {
				fmt.Println("no table selected")
				return
			}

			scan := dynago.ScanTable(table).
				WithSegment(*segment, *total)

			if *limit > 0 {
				scan = scan.WithLimit(*limit)
			}

			if *count {
				scan = scan.WithSelect(dynago.SELECT_COUNT)
			}

			if *next {
				scan = scan.WithStartKey(nextKey)
			}

			if *consumed {
				scan = scan.WithConsumed(true)
			}

			if items, lastKey, consumed, err := scan.Exec(db); err != nil {
				fmt.Println(err)
			} else {
				pretty.PrettyPrint(items)
				fmt.Println("consumed:", consumed)

				nextKey = lastKey
			}

			return
		}})

	commander.Commands["ls"] = commander.Commands["list"]
	commander.Commands["drop"] = commander.Commands["delete"]

	commander.CmdLoop()
}
<|endoftext|>


<commit_before>package hal

import (
	"fmt"
	"log"
	"sync"
)

// RouterCTX holds the router's context, including input/output chans.
type RouterCTX struct {
	brokers map[string]Broker
	in      chan *Evt     // messages from brokers --> plugins
	out     chan *Evt     // messages from plugins --> brokers
	update  chan struct{} // to notify the router that the instance list changed
	mut     sync.Mutex
	init    sync.Once
}

var routerSingleton RouterCTX

// Router returns the singleton router context. The router is initialized
// on the first call to this function.
func Router() *RouterCTX {
	routerSingleton.init.Do(func() {
		routerSingleton.in = make(chan *Evt, 1000)
		routerSingleton.out = make(chan *Evt, 1000)
		routerSingleton.update = make(chan struct{}, 1)
		routerSingleton.brokers = make(map[string]Broker)
	})

	return &routerSingleton
}

// forward from one (go) channel to another
// TODO: figure out if this needs to check for closed channels, etc.
func forward(from, to chan *Evt) {
	for {
		select {
		case evt := <-from:
			to <- evt
		}
	}
}

// AddBroker adds a broker to the router and starts forwarding
// events between it and the router.
func (r *RouterCTX) AddBroker(b Broker) {
	r.mut.Lock()
	defer r.mut.Unlock()

	if _, exists := r.brokers[b.Name()]; exists {
		panic(fmt.Sprintf("BUG: broker '%s' added > 1 times.", b.Name()))
	}

	b2r := make(chan *Evt, 1000) // messages from the broker to the router

	// start the broker's event stream
	go b.Stream(b2r)

	// forward events from the broker to the router's input channel
	go forward(b2r, r.in)

	r.brokers[b.Name()] = b
}

func (r *RouterCTX) GetBroker(name string) Broker {
	if broker, exists := r.brokers[name]; exists {
		return broker
	}

	return nil
}

// Route is the main method for the router. It blocks and should be run in a goroutine
// exactly once.
func (r *RouterCTX) Route() {
	for {
		select {
		case evt := <-r.in:
			if evt.Broker == nil {
				panic("BUG: received event with nil Broker. This breaks all the things!")
			}

			go r.processEvent(evt)
		}
	}
}

// processEvent processes one event and is intended to run in a goroutine.
func (r *RouterCTX) processEvent(evt *Evt) {
	var pname string // must be in the recovery handler's scope

	// get a snapshot of the instance list
	// TODO: keep an eye on the cost of copying this list for every message
	pr := PluginRegistry()
	instances := pr.InstanceList()

	// if a plugin panics, catch it & log it
	// TODO: report errors to a channel?
	defer func() {
		if r := recover(); r != nil {
			log.Printf("recovered panic in plugin %q\n", pname)
			log.Printf("panic: %q", r)
		}
	}()

	for _, inst := range instances {
		ibname := inst.Broker.Name()
		pname = inst.Plugin.Name // recovery handler ^ will pick this up in a panic

		// a plugin instance matches on broker, channel, and regex
		// first, check if the instance is attached to a specific broker or generic
		if ibname != evt.Broker.Name() && ibname != gBroker.Name() {
			continue
		}

		// if this is a generic broker instance and the event is marked as not
		// generic, skip it
		if ibname == gBroker.Name() && !evt.IsGeneric {
			continue
		}

		// check if it's the correct channel
		if evt.Channel != inst.Channel {
			continue
		}

		// finally, check message text against the regex
		if inst.Regex == "" || inst.regex.MatchString(evt.Body) {
			// this will copy the struct twice. It's intentional to avoid
			// mutating the evt between calls. The plugin func signature
			// forces the second copy.
			evtcpy := *evt

			// pass the plugin instance pointer to the plugin function so
			// it can access its fields for settings, etc.
			evtcpy.instance = inst

			// call the plugin function
			// this may block other plugins from processing the same event but
			// since it's already in a goroutine, other events won't be blocked
			inst.Func(evtcpy)
		}
	}
}
<commit_msg>explicitly use pointers to broker<commit_after>package hal

import (
	"fmt"
	"log"
	"sync"
)

// RouterCTX holds the router's context, including input/output chans.
type RouterCTX struct {
	brokers map[string]*Broker
	in      chan *Evt     // messages from brokers --> plugins
	out     chan *Evt     // messages from plugins --> brokers
	update  chan struct{} // to notify the router that the instance list changed
	mut     sync.Mutex
	init    sync.Once
}

var routerSingleton RouterCTX

// Router returns the singleton router context. The router is initialized
// on the first call to this function.
func Router() *RouterCTX {
	routerSingleton.init.Do(func() {
		routerSingleton.in = make(chan *Evt, 1000)
		routerSingleton.out = make(chan *Evt, 1000)
		routerSingleton.update = make(chan struct{}, 1)
		routerSingleton.brokers = make(map[string]*Broker)
	})

	return &routerSingleton
}

// forward from one (go) channel to another
// TODO: figure out if this needs to check for closed channels, etc.
func forward(from, to chan *Evt) {
	for {
		select {
		case evt := <-from:
			to <- evt
		}
	}
}

// AddBroker adds a broker to the router and starts forwarding
// events between it and the router.
func (r *RouterCTX) AddBroker(b Broker) {
	r.mut.Lock()
	defer r.mut.Unlock()

	if _, exists := r.brokers[b.Name()]; exists {
		panic(fmt.Sprintf("BUG: broker '%s' added > 1 times.", b.Name()))
	}

	b2r := make(chan *Evt, 1000) // messages from the broker to the router

	// start the broker's event stream
	go b.Stream(b2r)

	// forward events from the broker to the router's input channel
	go forward(b2r, r.in)

	r.brokers[b.Name()] = &b
}

func (r *RouterCTX) GetBroker(name string) *Broker {
	if broker, exists := r.brokers[name]; exists {
		return broker
	}

	return nil
}

// Route is the main method for the router. It blocks and should be run in a goroutine
// exactly once.
func (r *RouterCTX) Route() {
	for {
		select {
		case evt := <-r.in:
			if evt.Broker == nil {
				panic("BUG: received event with nil Broker. This breaks all the things!")
			}

			go r.processEvent(evt)
		}
	}
}

// processEvent processes one event and is intended to run in a goroutine.
func (r *RouterCTX) processEvent(evt *Evt) {
	var pname string // must be in the recovery handler's scope

	// get a snapshot of the instance list
	// TODO: keep an eye on the cost of copying this list for every message
	pr := PluginRegistry()
	instances := pr.InstanceList()

	// if a plugin panics, catch it & log it
	// TODO: report errors to a channel?
	defer func() {
		if r := recover(); r != nil {
			log.Printf("recovered panic in plugin %q\n", pname)
			log.Printf("panic: %q", r)
		}
	}()

	for _, inst := range instances {
		ibname := inst.Broker.Name()
		pname = inst.Plugin.Name // recovery handler ^ will pick this up in a panic

		// a plugin instance matches on broker, channel, and regex
		// first, check if the instance is attached to a specific broker or generic
		if ibname != evt.Broker.Name() && ibname != gBroker.Name() {
			continue
		}

		// if this is a generic broker instance and the event is marked as not
		// generic, skip it
		if ibname == gBroker.Name() && !evt.IsGeneric {
			continue
		}

		// check if it's the correct channel
		if evt.Channel != inst.Channel {
			continue
		}

		// finally, check message text against the regex
		if inst.Regex == "" || inst.regex.MatchString(evt.Body) {
			// this will copy the struct twice. It's intentional to avoid
			// mutating the evt between calls. The plugin func signature
			// forces the second copy.
			evtcpy := *evt

			// pass the plugin instance pointer to the plugin function so
			// it can access its fields for settings, etc.
			evtcpy.instance = inst

			// call the plugin function
			// this may block other plugins from processing the same event but
			// since it's already in a goroutine, other events won't be blocked
			inst.Func(evtcpy)
		}
	}
}
<|endoftext|>


<commit_before>//
//   Copyright 2017 Deadlock X42 <deadlock.x42@gmail.com>
//
//   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 generate

import (
	"encoding/json"
	"io"
)

type initializer func(Generator) error

var (
	initializers []initializer
)

// New creates a generate.
func New(r io.Reader) (Generator, error) {
	g := &generator{}
	for {
		err := json.NewDecoder(r).Decode(g)
		if err == io.EOF {
			break
		}
		if err != nil {
			return nil, err
		}
	}
	for _, i := range initializers {
		if err := i(g); err != nil {
			return nil, err
		}
	}
	return g, nil
}
<commit_msg>fix initialzers<commit_after>//
//   Copyright 2017 Deadlock X42 <deadlock.x42@gmail.com>
//
//   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 generate

import (
	"encoding/json"
	"io"
)

type initializer func(Generator) error

var initializers []initializer

// New creates a generate.
func New(r io.Reader) (Generator, error) {
	g := &generator{}
	for {
		err := json.NewDecoder(r).Decode(g)
		if err == io.EOF {
			break
		}
		if err != nil {
			return nil, err
		}
	}
	for _, i := range initializers {
		if err := i(g); err != nil {
			return nil, err
		}
	}
	return g, nil
}
<|endoftext|>


<commit_before>/***** BEGIN LICENSE BLOCK *****

# Author: David Birdsong (david@imgix.com)
# Copyright (c) 2014, Zebrafish Labs Inc.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 	Redistributions of source code must retain the above copyright notice,
# 	this list of conditions and the following disclaimer.
#
# 	Redistributions in binary form must reproduce the above copyright notice,
# 	this list of conditions and the following disclaimer in the documentation
# 	and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
# ***** END LICENSE BLOCK *****/

/*
Package hekametrics adds an output for https://github.com/rcrowley/go-metrics/

hekametrics encodes all metrics from a registry into a Heka protobuf message
and sends to a Heka server on it's native listener port.
*/
package hekametrics

import (
	"code.google.com/p/go-uuid/uuid"
	"fmt"
	"github.com/mozilla-services/heka/client"
	"github.com/mozilla-services/heka/message"
	"github.com/rcrowley/go-metrics"
	"log"
	"net/url"
	"os"
	"time"
)

var logger = log.New(os.Stderr, "[hekametrics]", log.LstdFlags)

type HekaClient struct {
	pid               int32
	hostname, msgtype string

	client    client.Client
	encoder   client.StreamEncoder
	sender    client.Sender
	connect_s *url.URL
}

//NewHekaClient creates and returns a HekaClient
//
//connect string like 'tcp://127.0.0.1:5564' and 'udp://127.0.0.1:5564'
//
//msgtype sets the 'Type' field on a Heka message
func NewHekaClient(connect, msgtype string) (hc *HekaClient, err error) {
	hc = &HekaClient{}
	hc.connect_s, err = url.ParseRequestURI(connect)
	if err != nil {
		return nil, err
	}
	switch hc.connect_s.Scheme {
	case "tcp", "udp":
		logger.Println("good boy")
	default:
		return nil, fmt.Errorf("scheme: '%s' not supported, try 'tcp://<host>:<port>' or 'udp://<host>:<port>'", hc.connect_s.Scheme)
	}
	hc.msgtype = msgtype
	hc.encoder = client.NewProtobufEncoder(nil)
	hc.pid = int32(os.Getpid())
	hc.hostname, err = os.Hostname()
	if err != nil {
		hc.hostname = "<no hostname>"
	}
	return
}

func (hc *HekaClient) write(b []byte) error {
	var err error
	reconnect := func() (e error) {
		if hc.sender != nil {
			hc.sender.Close()
			hc.sender = nil
		}

		logger.Printf("Connecting: %s\n", hc.connect_s)
		hc.sender, e = client.NewNetworkSender(hc.connect_s.Scheme, hc.connect_s.Host)
		if e != nil {
			hc.sender = nil
			logger.Printf("Err Connecting: %s %v\n", hc.connect_s, e)
		}
		return e
	}

	if hc.sender == nil {
		err = reconnect()
		if err != nil {
			return err
		}

	}

	err = hc.sender.SendMessage(b)
	if err != nil {
		logger.Printf("Inject: [error] send message: %s\n", err)
		err = reconnect()
		if err != nil {
			return err
		}
		err = hc.sender.SendMessage(b)
		if err != nil {
			return err
		}

	}

	return err

}

// LogHeka is a blocking exporter function which encodes and sends metrics to a Heka server
//
// all metrics in metrics.Registry r are stored on message.Message.Fields
//
// flushing them every Duration d
func (hc *HekaClient) LogHeka(r metrics.Registry, d time.Duration) {

	var (
		stream []byte
		err    error
	)

	for {
		msg := make_message(r)
		msg.SetTimestamp(time.Now().UnixNano())
		msg.SetUuid(uuid.NewRandom())
		msg.Setlogger("go-metrics")
		msg.SetType(hc.msgtype)
		msg.SetPid(hc.pid)
		msg.SetSeverity(100)
		msg.SetHostname(hc.hostname)
		msg.SetPayload("")

		err = hc.encoder.EncodeMessageStream(msg, &stream)
		if err != nil {
			logger.Printf("Inject: [error] encode message: %s\n", err)
		}
		err = hc.write(stream)
		if err != nil {
			logger.Printf("Inject: [error] send message: %s\n", err)
		}
		time.Sleep(d)

	}

}

func make_message(r metrics.Registry) *message.Message {

	msg := &message.Message{}
	add_float_mapping := func(pref string, names []string, vals []float64) {
		for i, n := range names {

			n = fmt.Sprintf("%s.%s", pref, n)

			if i+1 > len(vals) {
				logger.Println("skipping: %s no value\n", n)
				continue
			}
			f, e := message.NewField(n, vals[i], "")
			if e == nil {
				msg.AddField(f)
			} else {
				logger.Println("skipping: %s %v: %v\n", n, vals[i], e)
			}

		}

	}

	r.Each(func(name string, i interface{}) {

		switch metric := i.(type) {
		case metrics.Counter:
			message.NewInt64Field(msg, name, metric.Count(), "")
		case metrics.Gauge:
			message.NewInt64Field(msg, name, metric.Value(), "")

		case metrics.GaugeFloat64:
			f, e := message.NewField(name, metric.Value(), "")
			if e == nil {
				msg.AddField(f)
			} else {
				logger.Println("skipping: %s %v: %v\n", name, metric.Value(), e)
			}

		case metrics.Histogram:
			h := metric.Snapshot()
			vals_fl := h.Percentiles([]float64{0.5, 0.75, 0.95, 0.99, 0.999})
			vals_fl = append(vals_fl, h.Mean(), h.StdDev())
			names := []string{"50-percentile", "75-percentile", "95-percentile",
				"99-percentile", "999-percentile", "mean", "std-dev"}
			add_float_mapping(fmt.Sprintf("%s.histogram", name), names, vals_fl)

			names = []string{"count", "min", "max"}
			vals_i := []int64{h.Count(), h.Min(), h.Max()}

			for i, n := range names {
				n = fmt.Sprintf("%s.histogram.%s", name, n)
				message.NewInt64Field(msg, n, vals_i[i], n)
			}

		case metrics.Meter:
			m := metric.Snapshot()
			message.NewInt64Field(msg, fmt.Sprintf("%s.count", name),
				m.Count(), "")
			names := []string{"one-minute", "five-minute", "fifteen-minute", "mean"}
			vals_fl := []float64{m.Rate1(), m.Rate5(), m.Rate15(), m.RateMean()}

			add_float_mapping(name, names, vals_fl)
		case metrics.Timer:
			h := metric.Snapshot()
			vals_fl := h.Percentiles([]float64{0.5, 0.75, 0.95, 0.99, 0.999})
			vals_fl = append(vals_fl, h.Mean(), h.StdDev(), h.Rate1(),
				h.Rate5(), h.Rate15(), h.RateMean())
			names := []string{"50-percentile", "75-percentile", "95-percentile",
				"99-percentile", "999-percentile", "mean", "std-dev", "one-minute",
				"five-minute", "fifteen-minute", "mean-rate"}

			add_float_mapping(fmt.Sprintf("%s.timer", name), names, vals_fl)
			names = []string{"count", "min", "max"}
			vals_i := []int64{h.Count(), h.Min(), h.Max()}
			for i, n := range names {
				n = fmt.Sprintf("%s.timer.%s", name, n)
				message.NewInt64Field(msg, n, vals_i[i], "")
			}

		}
	})
	return msg

}
<commit_msg>sp fix<commit_after>/***** BEGIN LICENSE BLOCK *****

# Author: David Birdsong (david@imgix.com)
# Copyright (c) 2014, Zebrafish Labs Inc.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 	Redistributions of source code must retain the above copyright notice,
# 	this list of conditions and the following disclaimer.
#
# 	Redistributions in binary form must reproduce the above copyright notice,
# 	this list of conditions and the following disclaimer in the documentation
# 	and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
# ***** END LICENSE BLOCK *****/

/*
Package hekametrics adds an output for https://github.com/rcrowley/go-metrics/

hekametrics encodes all metrics from a registry into a Heka protobuf message
and sends to a Heka server on it's native listener port.
*/
package hekametrics

import (
	"code.google.com/p/go-uuid/uuid"
	"fmt"
	"github.com/mozilla-services/heka/client"
	"github.com/mozilla-services/heka/message"
	"github.com/rcrowley/go-metrics"
	"log"
	"net/url"
	"os"
	"time"
)

var logger = log.New(os.Stderr, "[hekametrics]", log.LstdFlags)

type HekaClient struct {
	pid               int32
	hostname, msgtype string

	client    client.Client
	encoder   client.StreamEncoder
	sender    client.Sender
	connect_s *url.URL
}

//NewHekaClient creates and returns a HekaClient
//
//connect string like 'tcp://127.0.0.1:5564' and 'udp://127.0.0.1:5564'
//
//msgtype sets the 'Type' field on a Heka message
func NewHekaClient(connect, msgtype string) (hc *HekaClient, err error) {
	hc = &HekaClient{}
	hc.connect_s, err = url.ParseRequestURI(connect)
	if err != nil {
		return nil, err
	}
	switch hc.connect_s.Scheme {
	case "tcp", "udp":
	default:
		return nil, fmt.Errorf("scheme: '%s' not supported, try 'tcp://<host>:<port>' or 'udp://<host>:<port>'", hc.connect_s.Scheme)
	}
	hc.msgtype = msgtype
	hc.encoder = client.NewProtobufEncoder(nil)
	hc.pid = int32(os.Getpid())
	hc.hostname, err = os.Hostname()
	if err != nil {
		hc.hostname = "<no hostname>"
	}
	return
}

func (hc *HekaClient) write(b []byte) error {
	var err error
	reconnect := func() (e error) {
		if hc.sender != nil {
			hc.sender.Close()
			hc.sender = nil
		}

		logger.Printf("Connecting: %s\n", hc.connect_s)
		hc.sender, e = client.NewNetworkSender(hc.connect_s.Scheme, hc.connect_s.Host)
		if e != nil {
			hc.sender = nil
			logger.Printf("Err Connecting: %s %v\n", hc.connect_s, e)
		}
		return e
	}

	if hc.sender == nil {
		err = reconnect()
		if err != nil {
			return err
		}

	}

	err = hc.sender.SendMessage(b)
	if err != nil {
		logger.Printf("Inject: [error] send message: %s\n", err)
		err = reconnect()
		if err != nil {
			return err
		}
		err = hc.sender.SendMessage(b)
		if err != nil {
			return err
		}

	}

	return err

}

// LogHeka is a blocking exporter function which encodes and sends metrics to a Heka server
//
// all metrics in metrics.Registry r are stored on message.Message.Fields
//
// flushing them every Duration d
func (hc *HekaClient) LogHeka(r metrics.Registry, d time.Duration) {

	var (
		stream []byte
		err    error
	)

	for {
		msg := make_message(r)
		msg.SetTimestamp(time.Now().UnixNano())
		msg.SetUuid(uuid.NewRandom())
		msg.SetLogger("go-metrics")
		msg.SetType(hc.msgtype)
		msg.SetPid(hc.pid)
		msg.SetSeverity(100)
		msg.SetHostname(hc.hostname)
		msg.SetPayload("")

		err = hc.encoder.EncodeMessageStream(msg, &stream)
		if err != nil {
			logger.Printf("Inject: [error] encode message: %s\n", err)
		}
		err = hc.write(stream)
		if err != nil {
			logger.Printf("Inject: [error] send message: %s\n", err)
		}
		time.Sleep(d)

	}

}

func make_message(r metrics.Registry) *message.Message {

	msg := &message.Message{}
	add_float_mapping := func(pref string, names []string, vals []float64) {
		for i, n := range names {

			n = fmt.Sprintf("%s.%s", pref, n)

			if i+1 > len(vals) {
				logger.Println("skipping: %s no value\n", n)
				continue
			}
			f, e := message.NewField(n, vals[i], "")
			if e == nil {
				msg.AddField(f)
			} else {
				logger.Println("skipping: %s %v: %v\n", n, vals[i], e)
			}

		}

	}

	r.Each(func(name string, i interface{}) {

		switch metric := i.(type) {
		case metrics.Counter:
			message.NewInt64Field(msg, name, metric.Count(), "")
		case metrics.Gauge:
			message.NewInt64Field(msg, name, metric.Value(), "")

		case metrics.GaugeFloat64:
			f, e := message.NewField(name, metric.Value(), "")
			if e == nil {
				msg.AddField(f)
			} else {
				logger.Println("skipping: %s %v: %v\n", name, metric.Value(), e)
			}

		case metrics.Histogram:
			h := metric.Snapshot()
			vals_fl := h.Percentiles([]float64{0.5, 0.75, 0.95, 0.99, 0.999})
			vals_fl = append(vals_fl, h.Mean(), h.StdDev())
			names := []string{"50-percentile", "75-percentile", "95-percentile",
				"99-percentile", "999-percentile", "mean", "std-dev"}
			add_float_mapping(fmt.Sprintf("%s.histogram", name), names, vals_fl)

			names = []string{"count", "min", "max"}
			vals_i := []int64{h.Count(), h.Min(), h.Max()}

			for i, n := range names {
				n = fmt.Sprintf("%s.histogram.%s", name, n)
				message.NewInt64Field(msg, n, vals_i[i], n)
			}

		case metrics.Meter:
			m := metric.Snapshot()
			message.NewInt64Field(msg, fmt.Sprintf("%s.count", name),
				m.Count(), "")
			names := []string{"one-minute", "five-minute", "fifteen-minute", "mean"}
			vals_fl := []float64{m.Rate1(), m.Rate5(), m.Rate15(), m.RateMean()}

			add_float_mapping(name, names, vals_fl)
		case metrics.Timer:
			h := metric.Snapshot()
			vals_fl := h.Percentiles([]float64{0.5, 0.75, 0.95, 0.99, 0.999})
			vals_fl = append(vals_fl, h.Mean(), h.StdDev(), h.Rate1(),
				h.Rate5(), h.Rate15(), h.RateMean())
			names := []string{"50-percentile", "75-percentile", "95-percentile",
				"99-percentile", "999-percentile", "mean", "std-dev", "one-minute",
				"five-minute", "fifteen-minute", "mean-rate"}

			add_float_mapping(fmt.Sprintf("%s.timer", name), names, vals_fl)
			names = []string{"count", "min", "max"}
			vals_i := []int64{h.Count(), h.Min(), h.Max()}
			for i, n := range names {
				n = fmt.Sprintf("%s.timer.%s", name, n)
				message.NewInt64Field(msg, n, vals_i[i], "")
			}

		}
	})
	return msg

}
<|endoftext|>


<commit_before><commit_msg>Added benchmarks for logging (#54)<commit_after><|endoftext|>


<commit_before>/*
鹅语
===========

  《赤装宁》
   by 弄姥尸
装宁徴宝岚繁勧，
岷崛畑辅逊安刔。
送佩暖锻弌赔仟，
畅酎抹亟赤装宁。

*/

package elang

import (
	"bytes"
	"io"

	"github.com/axgle/mahonia"
)

var (
	// 正转 u | u2j | g2u = u | u2e
	u2g = mahonia.NewEncoder("gbk")
	j2u = mahonia.NewDecoder("euc-jp")

	// 逆转 u | u2g | j2u = u | e2u
	u2j = mahonia.NewEncoder("euc-jp")
	g2u = mahonia.NewDecoder("gbk")
)

type Decoder struct {
	r    io.Reader
	u2jw io.Writer
	g2ur io.Reader
}

func (this *Decoder) Read(p []byte) (n int, err error) {
	io.Copy(this.u2jw, this.r)
	n, err = this.g2ur.Read(p)
	return
}
func NewDecoder(r io.Reader) io.Reader {
	buf := &bytes.Buffer{}
	return &Decoder{r, u2j.NewWriter(buf), g2u.NewReader(buf)}
}

type Encoder struct {
	w    io.Writer
	u2gw io.Writer
	j2ur io.Reader
}

func NewEncoder(w io.Writer) io.WriteCloser {
	buf := &bytes.Buffer{}
	return &Encoder{w, u2g.NewWriter(buf), j2u.NewReader(buf)}
}
func (this *Encoder) Write(p []byte) (n int, err error) {
	n, err = this.u2gw.Write(p)
	io.Copy(this.w, this.j2ur)
	return
}
func (this *Encoder) Close() error {
	return nil
}

// 转换中文到鹅语
func EncodeString(str string) string {
	buf := &bytes.Buffer{}
	e := NewEncoder(buf)
	io.Copy(e, bytes.NewBufferString(str))
	return buf.String()
}

// 转换鹅语到中文
func DecodeString(str string) string {
	buf := &bytes.Buffer{}
	e := NewDecoder(bytes.NewBufferString(str))
	io.Copy(buf, e)
	return buf.String()
}
<commit_msg>reimport "code.google.com/p/mahonia"<commit_after>/*
鹅语
===========

  《赤装宁》
   by 弄姥尸
装宁徴宝岚繁勧，
岷崛畑辅逊安刔。
送佩暖锻弌赔仟，
畅酎抹亟赤装宁。

*/

package elang

import (
	"bytes"
	"io"

	"code.google.com/p/mahonia"
)

var (
	// 正转 u | u2j | g2u = u | u2e
	u2g = mahonia.NewEncoder("gbk")
	j2u = mahonia.NewDecoder("euc-jp")

	// 逆转 u | u2g | j2u = u | e2u
	u2j = mahonia.NewEncoder("euc-jp")
	g2u = mahonia.NewDecoder("gbk")
)

type Decoder struct {
	r    io.Reader
	u2jw io.Writer
	g2ur io.Reader
}

func (this *Decoder) Read(p []byte) (n int, err error) {
	io.Copy(this.u2jw, this.r)
	n, err = this.g2ur.Read(p)
	return
}
func NewDecoder(r io.Reader) io.Reader {
	buf := &bytes.Buffer{}
	return &Decoder{r, u2j.NewWriter(buf), g2u.NewReader(buf)}
}

type Encoder struct {
	w    io.Writer
	u2gw io.Writer
	j2ur io.Reader
}

func NewEncoder(w io.Writer) io.WriteCloser {
	buf := &bytes.Buffer{}
	return &Encoder{w, u2g.NewWriter(buf), j2u.NewReader(buf)}
}
func (this *Encoder) Write(p []byte) (n int, err error) {
	n, err = this.u2gw.Write(p)
	io.Copy(this.w, this.j2ur)
	return
}
func (this *Encoder) Close() error {
	return nil
}

// 转换中文到鹅语
func EncodeString(str string) string {
	buf := &bytes.Buffer{}
	e := NewEncoder(buf)
	io.Copy(e, bytes.NewBufferString(str))
	return buf.String()
}

// 转换鹅语到中文
func DecodeString(str string) string {
	buf := &bytes.Buffer{}
	e := NewDecoder(bytes.NewBufferString(str))
	io.Copy(buf, e)
	return buf.String()
}
<|endoftext|>