<commit_before>package chromedp

import (
	"context"
	"errors"
	"fmt"
	"log"
	"os"
	"sync"
	"sync/atomic"
	"time"

	easyjson "github.com/mailru/easyjson"

	"github.com/chromedp/cdproto"
	"github.com/chromedp/cdproto/browser"
	"github.com/chromedp/cdproto/cdp"
	"github.com/chromedp/cdproto/runtime"
	"github.com/chromedp/cdproto/target"
)

// Browser is the high-level Chrome DevTools Protocol browser manager, handling
// the browser process runner, WebSocket clients, associated targets, and
// network, page, and DOM events.
type Browser struct {
	// next is the next message id.
	// NOTE: needs to be 64-bit aligned for 32-bit targets too, so be careful when moving this field.
	// This will be eventually done by the compiler once https://github.com/golang/go/issues/599 is fixed.
	next int64

	// LostConnection is closed when the websocket connection to Chrome is
	// dropped. This can be useful to make sure that Browser's context is
	// cancelled (and the handler stopped) once the connection has failed.
	LostConnection chan struct{}

	// closingGracefully is closed by Close before gracefully shutting down
	// the browser. This way, when the connection to the browser is lost and
	// LostConnection is closed, we will know not to immediately kill the
	// Chrome process. This is important to let the browser shut itself off,
	// saving its state to disk.
	closingGracefully chan struct{}

	dialTimeout time.Duration

	// pages keeps track of the attached targets, indexed by each's session
	// ID. The only reaon this is a field is so that the tests can check the
	// map once a browser is closed.
	pages map[target.SessionID]*Target

	listenersMu sync.Mutex
	listeners   []cancelableListener

	conn Transport

	// newTabQueue is the queue used to create new target handlers, once a new
	// tab is created and attached to. The newly created Target is sent back
	// via newTabResult.
	newTabQueue chan *Target

	// cmdQueue is the outgoing command queue.
	cmdQueue chan *cdproto.Message

	// logging funcs
	logf func(string, ...interface{})
	errf func(string, ...interface{})
	dbgf func(string, ...interface{})

	// The optional fields below are helpful for some tests.

	// process can be initialized by the allocators which start a process
	// when allocating a browser.
	process *os.Process

	// userDataDir can be initialized by the allocators which set up user
	// data dirs directly.
	userDataDir string
}

// NewBrowser creates a new browser. Typically, this function wouldn't be called
// directly, as the Allocator interface takes care of it.
func NewBrowser(ctx context.Context, urlstr string, opts ...BrowserOption) (*Browser, error) {
	b := &Browser{
		LostConnection:    make(chan struct{}),
		closingGracefully: make(chan struct{}),

		dialTimeout: 10 * time.Second,

		newTabQueue: make(chan *Target),

		// Fit some jobs without blocking, to reduce blocking in Execute.
		cmdQueue: make(chan *cdproto.Message, 32),

		logf: log.Printf,
	}
	// apply options
	for _, o := range opts {
		o(b)
	}
	// ensure errf is set
	if b.errf == nil {
		b.errf = func(s string, v ...interface{}) { b.logf("ERROR: "+s, v...) }
	}

	dialCtx := ctx
	if b.dialTimeout > 0 {
		var cancel context.CancelFunc
		dialCtx, cancel = context.WithTimeout(ctx, b.dialTimeout)
		defer cancel()
	}

	var err error
	urlstr = forceIP(urlstr)
	b.conn, err = DialContext(dialCtx, urlstr, WithConnDebugf(b.dbgf))
	if err != nil {
		return nil, fmt.Errorf("could not dial %q: %w", urlstr, err)
	}

	go b.run(ctx)
	return b, nil
}

func (b *Browser) newExecutorForTarget(ctx context.Context, targetID target.ID, sessionID target.SessionID) (*Target, error) {
	if targetID == "" {
		return nil, errors.New("empty target ID")
	}
	if sessionID == "" {
		return nil, errors.New("empty session ID")
	}
	t := &Target{
		browser:   b,
		TargetID:  targetID,
		SessionID: sessionID,

		messageQueue: make(chan *cdproto.Message, 1024),
		frames:       make(map[cdp.FrameID]*cdp.Frame),
		execContexts: make(map[cdp.FrameID]runtime.ExecutionContextID),
		cur:          cdp.FrameID(targetID),

		logf: b.logf,
		errf: b.errf,
	}

	// This send should be blocking, to ensure the tab is inserted into the
	// map before any more target events are routed.
	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case b.newTabQueue <- t:
	}
	return t, nil
}

func (b *Browser) Execute(ctx context.Context, method string, params easyjson.Marshaler, res easyjson.Unmarshaler) error {
	// Certain methods aren't available to the user directly.
	if method == browser.CommandClose {
		return fmt.Errorf("to close the browser gracefully, use chromedp.Cancel")
	}
	return b.execute(ctx, method, params, res)
}

func (b *Browser) execute(ctx context.Context, method string, params easyjson.Marshaler, res easyjson.Unmarshaler) error {
	id := atomic.AddInt64(&b.next, 1)
	lctx, cancel := context.WithCancel(ctx)
	ch := make(chan *cdproto.Message, 1)
	fn := func(ev interface{}) {
		if msg, ok := ev.(*cdproto.Message); ok && msg.ID == id {
			select {
			case <-ctx.Done():
			case ch <- msg:
			}
			cancel()
		}
	}
	b.listenersMu.Lock()
	b.listeners = append(b.listeners, cancelableListener{lctx, fn})
	b.listenersMu.Unlock()

	// send command
	var buf []byte
	if params != nil {
		var err error
		buf, err = easyjson.Marshal(params)
		if err != nil {
			return err
		}
	}
	cmd := &cdproto.Message{
		ID:     id,
		Method: cdproto.MethodType(method),
		Params: buf,
	}
	select {
	case <-ctx.Done():
		return ctx.Err()
	case b.cmdQueue <- cmd:
	}

	// wait for result
	select {
	case <-ctx.Done():
		return ctx.Err()
	case msg := <-ch:
		switch {
		case msg == nil:
			return ErrChannelClosed
		case msg.Error != nil:
			return msg.Error
		case res != nil:
			return easyjson.Unmarshal(msg.Result, res)
		}
	}
	return nil
}

func (b *Browser) run(ctx context.Context) {
	defer b.conn.Close()

	// incomingQueue is the queue of incoming target events, to be routed by
	// their session ID.
	incomingQueue := make(chan *cdproto.Message, 1)

	delTabQueue := make(chan target.SessionID, 1)

	// This goroutine continuously reads events from the websocket
	// connection. The separate goroutine is needed since a websocket read
	// is blocking, so it cannot be used in a select statement.
	go func() {
		// Signal to run and exit the browser cleanup goroutine.
		defer close(b.LostConnection)

		for {
			msg := new(cdproto.Message)
			if err := b.conn.Read(ctx, msg); err != nil {
				return
			}

			switch {
			case msg.SessionID != "" && (msg.Method != "" || msg.ID != 0):
				select {
				case <-ctx.Done():
					return
				case incomingQueue <- msg:
				}

			case msg.Method != "":
				ev, err := cdproto.UnmarshalMessage(msg)
				if err != nil {
					b.errf("%s", err)
					continue
				}
				b.listenersMu.Lock()
				b.listeners = runListeners(b.listeners, ev)
				b.listenersMu.Unlock()

				if ev, ok := ev.(*target.EventDetachedFromTarget); ok {
					delTabQueue <- ev.SessionID
				}

			case msg.ID != 0:
				b.listenersMu.Lock()
				b.listeners = runListeners(b.listeners, msg)
				b.listenersMu.Unlock()

			default:
				b.errf("ignoring malformed incoming message (missing id or method): %#v", msg)
			}
		}
	}()

	b.pages = make(map[target.SessionID]*Target, 32)
	for {
		select {
		case <-ctx.Done():
			return

		case msg := <-b.cmdQueue:
			if err := b.conn.Write(ctx, msg); err != nil {
				b.errf("%s", err)
				continue
			}

		case t := <-b.newTabQueue:
			if _, ok := b.pages[t.SessionID]; ok {
				b.errf("executor for %q already exists", t.SessionID)
			}
			b.pages[t.SessionID] = t

		case sessionID := <-delTabQueue:
			if _, ok := b.pages[sessionID]; !ok {
				b.errf("executor for %q doesn't exist", sessionID)
			}
			delete(b.pages, sessionID)

		case m := <-incomingQueue:
			page, ok := b.pages[m.SessionID]
			if !ok {
				// A page we recently closed still sending events.
				continue
			}

			select {
			case <-ctx.Done():
				return
			case page.messageQueue <- m:
			}

		case <-b.LostConnection:
			return // to avoid "write: broken pipe" errors
		}
	}
}

// BrowserOption is a browser option.
type BrowserOption = func(*Browser)

// WithBrowserLogf is a browser option to specify a func to receive general logging.
func WithBrowserLogf(f func(string, ...interface{})) BrowserOption {
	return func(b *Browser) { b.logf = f }
}

// WithBrowserErrorf is a browser option to specify a func to receive error logging.
func WithBrowserErrorf(f func(string, ...interface{})) BrowserOption {
	return func(b *Browser) { b.errf = f }
}

// WithBrowserDebugf is a browser option to specify a func to log actual
// websocket messages.
func WithBrowserDebugf(f func(string, ...interface{})) BrowserOption {
	return func(b *Browser) { b.dbgf = f }
}

// WithConsolef is a browser option to specify a func to receive chrome log events.
//
// Note: NOT YET IMPLEMENTED.
func WithConsolef(f func(string, ...interface{})) BrowserOption {
	return func(b *Browser) {}
}

// WithDialTimeout is a browser option to specify the timeout when dialing a
// browser's websocket address. The default is ten seconds; use a zero duration
// to not use a timeout.
func WithDialTimeout(d time.Duration) BrowserOption {
	return func(b *Browser) { b.dialTimeout = d }
}
<commit_msg>Expose process of the Browser<commit_after>package chromedp

import (
	"context"
	"errors"
	"fmt"
	"log"
	"os"
	"sync"
	"sync/atomic"
	"time"

	easyjson "github.com/mailru/easyjson"

	"github.com/chromedp/cdproto"
	"github.com/chromedp/cdproto/browser"
	"github.com/chromedp/cdproto/cdp"
	"github.com/chromedp/cdproto/runtime"
	"github.com/chromedp/cdproto/target"
)

// Browser is the high-level Chrome DevTools Protocol browser manager, handling
// the browser process runner, WebSocket clients, associated targets, and
// network, page, and DOM events.
type Browser struct {
	// next is the next message id.
	// NOTE: needs to be 64-bit aligned for 32-bit targets too, so be careful when moving this field.
	// This will be eventually done by the compiler once https://github.com/golang/go/issues/599 is fixed.
	next int64

	// LostConnection is closed when the websocket connection to Chrome is
	// dropped. This can be useful to make sure that Browser's context is
	// cancelled (and the handler stopped) once the connection has failed.
	LostConnection chan struct{}

	// closingGracefully is closed by Close before gracefully shutting down
	// the browser. This way, when the connection to the browser is lost and
	// LostConnection is closed, we will know not to immediately kill the
	// Chrome process. This is important to let the browser shut itself off,
	// saving its state to disk.
	closingGracefully chan struct{}

	dialTimeout time.Duration

	// pages keeps track of the attached targets, indexed by each's session
	// ID. The only reaon this is a field is so that the tests can check the
	// map once a browser is closed.
	pages map[target.SessionID]*Target

	listenersMu sync.Mutex
	listeners   []cancelableListener

	conn Transport

	// newTabQueue is the queue used to create new target handlers, once a new
	// tab is created and attached to. The newly created Target is sent back
	// via newTabResult.
	newTabQueue chan *Target

	// cmdQueue is the outgoing command queue.
	cmdQueue chan *cdproto.Message

	// logging funcs
	logf func(string, ...interface{})
	errf func(string, ...interface{})
	dbgf func(string, ...interface{})

	// The optional fields below are helpful for some tests.

	// process can be initialized by the allocators which start a process
	// when allocating a browser.
	process *os.Process

	// userDataDir can be initialized by the allocators which set up user
	// data dirs directly.
	userDataDir string
}

// NewBrowser creates a new browser. Typically, this function wouldn't be called
// directly, as the Allocator interface takes care of it.
func NewBrowser(ctx context.Context, urlstr string, opts ...BrowserOption) (*Browser, error) {
	b := &Browser{
		LostConnection:    make(chan struct{}),
		closingGracefully: make(chan struct{}),

		dialTimeout: 10 * time.Second,

		newTabQueue: make(chan *Target),

		// Fit some jobs without blocking, to reduce blocking in Execute.
		cmdQueue: make(chan *cdproto.Message, 32),

		logf: log.Printf,
	}
	// apply options
	for _, o := range opts {
		o(b)
	}
	// ensure errf is set
	if b.errf == nil {
		b.errf = func(s string, v ...interface{}) { b.logf("ERROR: "+s, v...) }
	}

	dialCtx := ctx
	if b.dialTimeout > 0 {
		var cancel context.CancelFunc
		dialCtx, cancel = context.WithTimeout(ctx, b.dialTimeout)
		defer cancel()
	}

	var err error
	urlstr = forceIP(urlstr)
	b.conn, err = DialContext(dialCtx, urlstr, WithConnDebugf(b.dbgf))
	if err != nil {
		return nil, fmt.Errorf("could not dial %q: %w", urlstr, err)
	}

	go b.run(ctx)
	return b, nil
}

// Process returns the process object of the browser.
//
// It could be nil when the browser is allocated with RemoteAllocator.
// It could be useful for a monitoring system to collect process metrics of the browser process.
// (see https://pkg.go.dev/github.com/prometheus/client_golang/prometheus#NewProcessCollector for an example)
//
// Example:
//     if process := chromedp.FromContext(ctx).Browser.Process(); process != nil {
//         fmt.Printf("Browser PID: %v", process.Pid)
//     }
func (b *Browser) Process() *os.Process {
	return b.process
}

func (b *Browser) newExecutorForTarget(ctx context.Context, targetID target.ID, sessionID target.SessionID) (*Target, error) {
	if targetID == "" {
		return nil, errors.New("empty target ID")
	}
	if sessionID == "" {
		return nil, errors.New("empty session ID")
	}
	t := &Target{
		browser:   b,
		TargetID:  targetID,
		SessionID: sessionID,

		messageQueue: make(chan *cdproto.Message, 1024),
		frames:       make(map[cdp.FrameID]*cdp.Frame),
		execContexts: make(map[cdp.FrameID]runtime.ExecutionContextID),
		cur:          cdp.FrameID(targetID),

		logf: b.logf,
		errf: b.errf,
	}

	// This send should be blocking, to ensure the tab is inserted into the
	// map before any more target events are routed.
	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case b.newTabQueue <- t:
	}
	return t, nil
}

func (b *Browser) Execute(ctx context.Context, method string, params easyjson.Marshaler, res easyjson.Unmarshaler) error {
	// Certain methods aren't available to the user directly.
	if method == browser.CommandClose {
		return fmt.Errorf("to close the browser gracefully, use chromedp.Cancel")
	}
	return b.execute(ctx, method, params, res)
}

func (b *Browser) execute(ctx context.Context, method string, params easyjson.Marshaler, res easyjson.Unmarshaler) error {
	id := atomic.AddInt64(&b.next, 1)
	lctx, cancel := context.WithCancel(ctx)
	ch := make(chan *cdproto.Message, 1)
	fn := func(ev interface{}) {
		if msg, ok := ev.(*cdproto.Message); ok && msg.ID == id {
			select {
			case <-ctx.Done():
			case ch <- msg:
			}
			cancel()
		}
	}
	b.listenersMu.Lock()
	b.listeners = append(b.listeners, cancelableListener{lctx, fn})
	b.listenersMu.Unlock()

	// send command
	var buf []byte
	if params != nil {
		var err error
		buf, err = easyjson.Marshal(params)
		if err != nil {
			return err
		}
	}
	cmd := &cdproto.Message{
		ID:     id,
		Method: cdproto.MethodType(method),
		Params: buf,
	}
	select {
	case <-ctx.Done():
		return ctx.Err()
	case b.cmdQueue <- cmd:
	}

	// wait for result
	select {
	case <-ctx.Done():
		return ctx.Err()
	case msg := <-ch:
		switch {
		case msg == nil:
			return ErrChannelClosed
		case msg.Error != nil:
			return msg.Error
		case res != nil:
			return easyjson.Unmarshal(msg.Result, res)
		}
	}
	return nil
}

func (b *Browser) run(ctx context.Context) {
	defer b.conn.Close()

	// incomingQueue is the queue of incoming target events, to be routed by
	// their session ID.
	incomingQueue := make(chan *cdproto.Message, 1)

	delTabQueue := make(chan target.SessionID, 1)

	// This goroutine continuously reads events from the websocket
	// connection. The separate goroutine is needed since a websocket read
	// is blocking, so it cannot be used in a select statement.
	go func() {
		// Signal to run and exit the browser cleanup goroutine.
		defer close(b.LostConnection)

		for {
			msg := new(cdproto.Message)
			if err := b.conn.Read(ctx, msg); err != nil {
				return
			}

			switch {
			case msg.SessionID != "" && (msg.Method != "" || msg.ID != 0):
				select {
				case <-ctx.Done():
					return
				case incomingQueue <- msg:
				}

			case msg.Method != "":
				ev, err := cdproto.UnmarshalMessage(msg)
				if err != nil {
					b.errf("%s", err)
					continue
				}
				b.listenersMu.Lock()
				b.listeners = runListeners(b.listeners, ev)
				b.listenersMu.Unlock()

				if ev, ok := ev.(*target.EventDetachedFromTarget); ok {
					delTabQueue <- ev.SessionID
				}

			case msg.ID != 0:
				b.listenersMu.Lock()
				b.listeners = runListeners(b.listeners, msg)
				b.listenersMu.Unlock()

			default:
				b.errf("ignoring malformed incoming message (missing id or method): %#v", msg)
			}
		}
	}()

	b.pages = make(map[target.SessionID]*Target, 32)
	for {
		select {
		case <-ctx.Done():
			return

		case msg := <-b.cmdQueue:
			if err := b.conn.Write(ctx, msg); err != nil {
				b.errf("%s", err)
				continue
			}

		case t := <-b.newTabQueue:
			if _, ok := b.pages[t.SessionID]; ok {
				b.errf("executor for %q already exists", t.SessionID)
			}
			b.pages[t.SessionID] = t

		case sessionID := <-delTabQueue:
			if _, ok := b.pages[sessionID]; !ok {
				b.errf("executor for %q doesn't exist", sessionID)
			}
			delete(b.pages, sessionID)

		case m := <-incomingQueue:
			page, ok := b.pages[m.SessionID]
			if !ok {
				// A page we recently closed still sending events.
				continue
			}

			select {
			case <-ctx.Done():
				return
			case page.messageQueue <- m:
			}

		case <-b.LostConnection:
			return // to avoid "write: broken pipe" errors
		}
	}
}

// BrowserOption is a browser option.
type BrowserOption = func(*Browser)

// WithBrowserLogf is a browser option to specify a func to receive general logging.
func WithBrowserLogf(f func(string, ...interface{})) BrowserOption {
	return func(b *Browser) { b.logf = f }
}

// WithBrowserErrorf is a browser option to specify a func to receive error logging.
func WithBrowserErrorf(f func(string, ...interface{})) BrowserOption {
	return func(b *Browser) { b.errf = f }
}

// WithBrowserDebugf is a browser option to specify a func to log actual
// websocket messages.
func WithBrowserDebugf(f func(string, ...interface{})) BrowserOption {
	return func(b *Browser) { b.dbgf = f }
}

// WithConsolef is a browser option to specify a func to receive chrome log events.
//
// Note: NOT YET IMPLEMENTED.
func WithConsolef(f func(string, ...interface{})) BrowserOption {
	return func(b *Browser) {}
}

// WithDialTimeout is a browser option to specify the timeout when dialing a
// browser's websocket address. The default is ten seconds; use a zero duration
// to not use a timeout.
func WithDialTimeout(d time.Duration) BrowserOption {
	return func(b *Browser) { b.dialTimeout = d }
}
<|endoftext|>


<commit_before>package config

import (
	"reflect"
	"testing"
)

func TestVariableDetectWalker(t *testing.T) {
	w := new(variableDetectWalker)

	str := `foo ${var.bar}`
	if err := w.Primitive(reflect.ValueOf(str)); err != nil {
		t.Fatalf("err: %s", err)
	}

	if len(w.Variables) != 1 {
		t.Fatalf("bad: %#v", w.Variables)
	}
	if w.Variables["var.bar"].(*UserVariable).FullKey() != "var.bar" {
		t.Fatalf("bad: %#v", w.Variables)
	}
}

func TestVariableDetectWalker_bad(t *testing.T) {
	w := new(variableDetectWalker)

	str := `foo ${bar}`
	if err := w.Primitive(reflect.ValueOf(str)); err == nil {
		t.Fatal("should error")
	}
}

func TestVariableDetectWalker_escaped(t *testing.T) {
	w := new(variableDetectWalker)

	str := `foo $${var.bar}`
	if err := w.Primitive(reflect.ValueOf(str)); err != nil {
		t.Fatalf("err: %s", err)
	}

	if len(w.Variables) > 0 {
		t.Fatalf("bad: %#v", w.Variables)
	}
}

func TestVariableDetectWalker_empty(t *testing.T) {
	w := new(variableDetectWalker)

	str := `foo`
	if err := w.Primitive(reflect.ValueOf(str)); err != nil {
		t.Fatalf("err: %s", err)
	}

	if len(w.Variables) > 0 {
		t.Fatalf("bad: %#v", w.Variables)
	}
}

<commit_msg>config: add benchmark test<commit_after>package config

import (
	"reflect"
	"testing"
)

func BenchmarkVariableDetectWalker(b *testing.B) {
	w := new(variableDetectWalker)

	str := reflect.ValueOf(`foo ${var.bar} bar ${bar.baz.bing} $${escaped}`)
	for i := 0; i < b.N; i++ {
		w.Variables = nil
		w.Primitive(str)
	}
}

func TestVariableDetectWalker(t *testing.T) {
	w := new(variableDetectWalker)

	str := `foo ${var.bar}`
	if err := w.Primitive(reflect.ValueOf(str)); err != nil {
		t.Fatalf("err: %s", err)
	}

	if len(w.Variables) != 1 {
		t.Fatalf("bad: %#v", w.Variables)
	}
	if w.Variables["var.bar"].(*UserVariable).FullKey() != "var.bar" {
		t.Fatalf("bad: %#v", w.Variables)
	}
}

func TestVariableDetectWalker_bad(t *testing.T) {
	w := new(variableDetectWalker)

	str := `foo ${bar}`
	if err := w.Primitive(reflect.ValueOf(str)); err == nil {
		t.Fatal("should error")
	}
}

func TestVariableDetectWalker_escaped(t *testing.T) {
	w := new(variableDetectWalker)

	str := `foo $${var.bar}`
	if err := w.Primitive(reflect.ValueOf(str)); err != nil {
		t.Fatalf("err: %s", err)
	}

	if len(w.Variables) > 0 {
		t.Fatalf("bad: %#v", w.Variables)
	}
}

func TestVariableDetectWalker_empty(t *testing.T) {
	w := new(variableDetectWalker)

	str := `foo`
	if err := w.Primitive(reflect.ValueOf(str)); err != nil {
		t.Fatalf("err: %s", err)
	}

	if len(w.Variables) > 0 {
		t.Fatalf("bad: %#v", w.Variables)
	}
}

<|endoftext|>


<commit_before>package main

import (
	boshapp "bosh/app"
	boshlog "bosh/logger"
	"os"
)

func main() {
	logger := boshlog.NewLogger(boshlog.LevelDebug)
	defer logger.HandlePanic("Main")
	logger.Debug("main", "Starting agent")

	app := boshapp.New(logger)
	app.Setup(os.Args)
	err := app.Run()

	if err != nil {
		logger.Error("Main", err.Error())
		os.Exit(1)
	}
}
<commit_msg>catch app setup error<commit_after>package main

import (
	"os"

	boshapp "bosh/app"
	boshlog "bosh/logger"
)

const mainLogTag = "main"

func main() {
	logger := boshlog.NewLogger(boshlog.LevelDebug)
	defer logger.HandlePanic("Main")

	logger.Debug(mainLogTag, "Starting agent")

	app := boshapp.New(logger)

	err := app.Setup(os.Args)
	if err != nil {
		logger.Error(mainLogTag, "App setup", err.Error())
		os.Exit(1)
	}

	err = app.Run()
	if err != nil {
		logger.Error(mainLogTag, "App run", err.Error())
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package libkbfs

import (
	"time"

	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/logger"
	keybase1 "github.com/keybase/client/go/protocol"
	rpc "github.com/keybase/go-framed-msgpack-rpc"
	"golang.org/x/net/context"
)

// CryptoClient implements the Crypto interface by sending RPCs to the
// keybase daemon to perform signatures using the device's current
// signing key.
type CryptoClient struct {
	CryptoCommon
	client     keybase1.CryptoClient
	shutdownFn func()
	config     Config
}

var _ Crypto = (*CryptoClient)(nil)

var _ rpc.ConnectionHandler = (*CryptoClient)(nil)

// NewCryptoClient constructs a new CryptoClient.
func NewCryptoClient(config Config, kbCtx *libkb.GlobalContext, log logger.Logger) *CryptoClient {
	c := &CryptoClient{
		CryptoCommon: CryptoCommon{
			codec: config.Codec(),
			log:   log,
		},
		config: config,
	}
	conn := NewSharedKeybaseConnection(kbCtx, config, c)
	c.client = keybase1.CryptoClient{Cli: conn.GetClient()}
	c.shutdownFn = conn.Shutdown
	return c
}

// newCryptoClientWithClient should only be used for testing.
func newCryptoClientWithClient(codec Codec, client rpc.GenericClient,
	log logger.Logger) *CryptoClient {
	return &CryptoClient{
		CryptoCommon: CryptoCommon{
			codec: codec,
			log:   log,
		},
		client: keybase1.CryptoClient{Cli: client},
	}
}

// HandlerName implements the ConnectionHandler interface.
func (CryptoClient) HandlerName() string {
	return "CryptoClient"
}

// OnConnect implements the ConnectionHandler interface.
func (c *CryptoClient) OnConnect(ctx context.Context, conn *rpc.Connection,
	_ rpc.GenericClient, server *rpc.Server) error {
	c.config.KBFSOps().PushConnectionStatusChange(KeybaseServiceName, nil)
	return nil
}

// OnConnectError implements the ConnectionHandler interface.
func (c *CryptoClient) OnConnectError(err error, wait time.Duration) {
	c.log.Warning("CryptoClient: connection error: %q; retrying in %s",
		err, wait)
	c.config.KBFSOps().PushConnectionStatusChange(KeybaseServiceName, err)
}

// OnDoCommandError implements the ConnectionHandler interface.
func (c *CryptoClient) OnDoCommandError(err error, wait time.Duration) {
	c.log.Warning("CryptoClient: docommand error: %q; retrying in %s",
		err, wait)
	c.config.KBFSOps().PushConnectionStatusChange(KeybaseServiceName, err)
}

// OnDisconnected implements the ConnectionHandler interface.
func (c *CryptoClient) OnDisconnected(_ context.Context,
	status rpc.DisconnectStatus) {
	if status == rpc.StartingNonFirstConnection {
		c.log.Warning("CryptoClient is disconnected")
		c.config.KBFSOps().PushConnectionStatusChange(KeybaseServiceName, errDisconnected{})
	}
}

// ShouldRetry implements the ConnectionHandler interface.
func (c *CryptoClient) ShouldRetry(rpcName string, err error) bool {
	return false
}

// ShouldRetryOnConnect implements the ConnectionHandler interface.
func (c *CryptoClient) ShouldRetryOnConnect(err error) bool {
	_, inputCanceled := err.(libkb.InputCanceledError)
	return !inputCanceled
}

// Sign implements the Crypto interface for CryptoClient.
func (c *CryptoClient) Sign(ctx context.Context, msg []byte) (
	sigInfo SignatureInfo, err error) {
	defer func() {
		c.log.CDebugf(ctx, "Signed %d-byte message with %s: err=%v", len(msg),
			sigInfo, err)
	}()

	ed25519SigInfo, err := c.client.SignED25519(ctx, keybase1.SignED25519Arg{
		Msg:    msg,
		Reason: "to use kbfs",
	})
	if err != nil {
		return
	}

	sigInfo = SignatureInfo{
		Version:      SigED25519,
		Signature:    ed25519SigInfo.Sig[:],
		VerifyingKey: MakeVerifyingKey(libkb.NaclSigningKeyPublic(ed25519SigInfo.PublicKey).GetKID()),
	}
	return
}

// SignToString implements the Crypto interface for CryptoClient.
func (c *CryptoClient) SignToString(ctx context.Context, msg []byte) (
	signature string, err error) {
	defer func() {
		c.log.CDebugf(ctx, "Signed %d-byte message: err=%v", len(msg), err)
	}()
	signature, err = c.client.SignToString(ctx, keybase1.SignToStringArg{
		Msg:    msg,
		Reason: "KBFS Authentication",
	})
	return
}

func (c *CryptoClient) prepareTLFCryptKeyClientHalf(encryptedClientHalf EncryptedTLFCryptKeyClientHalf) (
	encryptedData keybase1.EncryptedBytes32, nonce keybase1.BoxNonce, err error) {
	if encryptedClientHalf.Version != EncryptionSecretbox {
		err = UnknownEncryptionVer{encryptedClientHalf.Version}
		return
	}

	if len(encryptedClientHalf.EncryptedData) != len(encryptedData) {
		err = libkb.DecryptionError{}
		return
	}
	copy(encryptedData[:], encryptedClientHalf.EncryptedData)

	if len(encryptedClientHalf.Nonce) != len(nonce) {
		err = InvalidNonceError{encryptedClientHalf.Nonce}
		return
	}
	copy(nonce[:], encryptedClientHalf.Nonce)
	return encryptedData, nonce, err
}

// DecryptTLFCryptKeyClientHalf implements the Crypto interface for
// CryptoClient.
func (c *CryptoClient) DecryptTLFCryptKeyClientHalf(ctx context.Context,
	publicKey TLFEphemeralPublicKey,
	encryptedClientHalf EncryptedTLFCryptKeyClientHalf) (
	clientHalf TLFCryptKeyClientHalf, err error) {
	encryptedData, nonce, err := c.prepareTLFCryptKeyClientHalf(encryptedClientHalf)
	if err != nil {
		return
	}

	decryptedClientHalf, err := c.client.UnboxBytes32(ctx, keybase1.UnboxBytes32Arg{
		EncryptedBytes32: encryptedData,
		Nonce:            nonce,
		PeersPublicKey:   keybase1.BoxPublicKey(publicKey.data),
		Reason:           "to use kbfs",
	})
	if err != nil {
		return
	}

	clientHalf = MakeTLFCryptKeyClientHalf(decryptedClientHalf)
	return
}

// DecryptTLFCryptKeyClientHalfAny implements the Crypto interface for
// CryptoClient.
func (c *CryptoClient) DecryptTLFCryptKeyClientHalfAny(ctx context.Context,
	keys []EncryptedTLFCryptKeyClientAndEphemeral, promptPaper bool) (
	clientHalf TLFCryptKeyClientHalf, index int, err error) {
	if len(keys) == 0 {
		return clientHalf, index, NoKeysError{}
	}
	bundles := make([]keybase1.CiphertextBundle, 0, len(keys))
	errors := make([]error, 0, len(keys))
	indexLookup := make([]int, 0, len(keys))
	for i, k := range keys {
		encryptedData, nonce, err := c.prepareTLFCryptKeyClientHalf(k.ClientHalf)
		if err != nil {
			errors = append(errors, err)
		} else {
			bundles = append(bundles, keybase1.CiphertextBundle{
				Kid:        k.PubKey.kidContainer.kid,
				Ciphertext: encryptedData,
				Nonce:      nonce,
				PublicKey:  keybase1.BoxPublicKey(k.EPubKey.data),
			})
			indexLookup = append(indexLookup, i)
		}
	}
	if len(bundles) == 0 {
		err = errors[0]
		return
	}
	res, err := c.client.UnboxBytes32Any(ctx, keybase1.UnboxBytes32AnyArg{
		Bundles:     bundles,
		Reason:      "to rekey for kbfs",
		PromptPaper: promptPaper,
	})
	if err != nil {
		return
	}
	return MakeTLFCryptKeyClientHalf(res.Plaintext), indexLookup[res.Index], nil
}

// Shutdown implements the Crypto interface for CryptoClient.
func (c *CryptoClient) Shutdown() {
	if c.shutdownFn != nil {
		c.shutdownFn()
	}
}
<commit_msg>crypto_client: log message about long RPCs after 2 minutes (#1038)<commit_after>package libkbfs

import (
	"time"

	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/logger"
	keybase1 "github.com/keybase/client/go/protocol"
	rpc "github.com/keybase/go-framed-msgpack-rpc"
	"golang.org/x/net/context"
)

// CryptoClient implements the Crypto interface by sending RPCs to the
// keybase daemon to perform signatures using the device's current
// signing key.
type CryptoClient struct {
	CryptoCommon
	client     keybase1.CryptoClient
	shutdownFn func()
	config     Config
}

// cryptoRPCWarningTime says how long we should wait before logging a
// message about an RPC taking too long.
const cryptoRPCWarningTime = 2 * time.Minute

var _ Crypto = (*CryptoClient)(nil)

var _ rpc.ConnectionHandler = (*CryptoClient)(nil)

// NewCryptoClient constructs a new CryptoClient.
func NewCryptoClient(config Config, kbCtx *libkb.GlobalContext, log logger.Logger) *CryptoClient {
	c := &CryptoClient{
		CryptoCommon: CryptoCommon{
			codec: config.Codec(),
			log:   log,
		},
		config: config,
	}
	conn := NewSharedKeybaseConnection(kbCtx, config, c)
	c.client = keybase1.CryptoClient{Cli: conn.GetClient()}
	c.shutdownFn = conn.Shutdown
	return c
}

// newCryptoClientWithClient should only be used for testing.
func newCryptoClientWithClient(codec Codec, client rpc.GenericClient,
	log logger.Logger) *CryptoClient {
	return &CryptoClient{
		CryptoCommon: CryptoCommon{
			codec: codec,
			log:   log,
		},
		client: keybase1.CryptoClient{Cli: client},
	}
}

// HandlerName implements the ConnectionHandler interface.
func (CryptoClient) HandlerName() string {
	return "CryptoClient"
}

// OnConnect implements the ConnectionHandler interface.
func (c *CryptoClient) OnConnect(ctx context.Context, conn *rpc.Connection,
	_ rpc.GenericClient, server *rpc.Server) error {
	c.config.KBFSOps().PushConnectionStatusChange(KeybaseServiceName, nil)
	return nil
}

// OnConnectError implements the ConnectionHandler interface.
func (c *CryptoClient) OnConnectError(err error, wait time.Duration) {
	c.log.Warning("CryptoClient: connection error: %q; retrying in %s",
		err, wait)
	c.config.KBFSOps().PushConnectionStatusChange(KeybaseServiceName, err)
}

// OnDoCommandError implements the ConnectionHandler interface.
func (c *CryptoClient) OnDoCommandError(err error, wait time.Duration) {
	c.log.Warning("CryptoClient: docommand error: %q; retrying in %s",
		err, wait)
	c.config.KBFSOps().PushConnectionStatusChange(KeybaseServiceName, err)
}

// OnDisconnected implements the ConnectionHandler interface.
func (c *CryptoClient) OnDisconnected(_ context.Context,
	status rpc.DisconnectStatus) {
	if status == rpc.StartingNonFirstConnection {
		c.log.Warning("CryptoClient is disconnected")
		c.config.KBFSOps().PushConnectionStatusChange(KeybaseServiceName, errDisconnected{})
	}
}

// ShouldRetry implements the ConnectionHandler interface.
func (c *CryptoClient) ShouldRetry(rpcName string, err error) bool {
	return false
}

// ShouldRetryOnConnect implements the ConnectionHandler interface.
func (c *CryptoClient) ShouldRetryOnConnect(err error) bool {
	_, inputCanceled := err.(libkb.InputCanceledError)
	return !inputCanceled
}

func (c *CryptoClient) logAboutLongRPCUnlessCancelled(ctx context.Context,
	method string) *time.Timer {
	return time.AfterFunc(cryptoRPCWarningTime, func() {
		c.log.CInfof(ctx, "%s RPC call took more than %s", method,
			cryptoRPCWarningTime)
	})
}

// Sign implements the Crypto interface for CryptoClient.
func (c *CryptoClient) Sign(ctx context.Context, msg []byte) (
	sigInfo SignatureInfo, err error) {
	c.log.CDebugf(ctx, "Signing %d-byte message", len(msg))
	defer func() {
		c.log.CDebugf(ctx, "Signed %d-byte message with %s: err=%v", len(msg),
			sigInfo, err)
	}()

	timer := c.logAboutLongRPCUnlessCancelled(ctx, "SignED25519")
	defer timer.Stop()
	ed25519SigInfo, err := c.client.SignED25519(ctx, keybase1.SignED25519Arg{
		Msg:    msg,
		Reason: "to use kbfs",
	})
	if err != nil {
		return
	}

	sigInfo = SignatureInfo{
		Version:      SigED25519,
		Signature:    ed25519SigInfo.Sig[:],
		VerifyingKey: MakeVerifyingKey(libkb.NaclSigningKeyPublic(ed25519SigInfo.PublicKey).GetKID()),
	}
	return
}

// SignToString implements the Crypto interface for CryptoClient.
func (c *CryptoClient) SignToString(ctx context.Context, msg []byte) (
	signature string, err error) {
	c.log.CDebugf(ctx, "Signing %d-byte message to string", len(msg))
	defer func() {
		c.log.CDebugf(ctx, "Signed %d-byte message: err=%v", len(msg), err)
	}()

	timer := c.logAboutLongRPCUnlessCancelled(ctx, "SignToString")
	defer timer.Stop()
	signature, err = c.client.SignToString(ctx, keybase1.SignToStringArg{
		Msg:    msg,
		Reason: "KBFS Authentication",
	})
	return
}

func (c *CryptoClient) prepareTLFCryptKeyClientHalf(encryptedClientHalf EncryptedTLFCryptKeyClientHalf) (
	encryptedData keybase1.EncryptedBytes32, nonce keybase1.BoxNonce, err error) {
	if encryptedClientHalf.Version != EncryptionSecretbox {
		err = UnknownEncryptionVer{encryptedClientHalf.Version}
		return
	}

	if len(encryptedClientHalf.EncryptedData) != len(encryptedData) {
		err = libkb.DecryptionError{}
		return
	}
	copy(encryptedData[:], encryptedClientHalf.EncryptedData)

	if len(encryptedClientHalf.Nonce) != len(nonce) {
		err = InvalidNonceError{encryptedClientHalf.Nonce}
		return
	}
	copy(nonce[:], encryptedClientHalf.Nonce)
	return encryptedData, nonce, err
}

// DecryptTLFCryptKeyClientHalf implements the Crypto interface for
// CryptoClient.
func (c *CryptoClient) DecryptTLFCryptKeyClientHalf(ctx context.Context,
	publicKey TLFEphemeralPublicKey,
	encryptedClientHalf EncryptedTLFCryptKeyClientHalf) (
	clientHalf TLFCryptKeyClientHalf, err error) {
	c.log.CDebugf(ctx, "Decrypting TLF client key half")
	defer func() {
		c.log.CDebugf(ctx, "Decrypted TLF client key half: %v", err)
	}()
	encryptedData, nonce, err := c.prepareTLFCryptKeyClientHalf(encryptedClientHalf)
	if err != nil {
		return
	}

	timer := c.logAboutLongRPCUnlessCancelled(ctx, "UnboxBytes32")
	defer timer.Stop()
	decryptedClientHalf, err := c.client.UnboxBytes32(ctx, keybase1.UnboxBytes32Arg{
		EncryptedBytes32: encryptedData,
		Nonce:            nonce,
		PeersPublicKey:   keybase1.BoxPublicKey(publicKey.data),
		Reason:           "to use kbfs",
	})
	if err != nil {
		return
	}

	clientHalf = MakeTLFCryptKeyClientHalf(decryptedClientHalf)
	return
}

// DecryptTLFCryptKeyClientHalfAny implements the Crypto interface for
// CryptoClient.
func (c *CryptoClient) DecryptTLFCryptKeyClientHalfAny(ctx context.Context,
	keys []EncryptedTLFCryptKeyClientAndEphemeral, promptPaper bool) (
	clientHalf TLFCryptKeyClientHalf, index int, err error) {
	c.log.CDebugf(ctx, "Decrypting TLF client key half with any key")
	defer func() {
		c.log.CDebugf(ctx, "Decrypted TLF client key half with any key: %v",
			err)
	}()
	if len(keys) == 0 {
		return clientHalf, index, NoKeysError{}
	}
	bundles := make([]keybase1.CiphertextBundle, 0, len(keys))
	errors := make([]error, 0, len(keys))
	indexLookup := make([]int, 0, len(keys))
	for i, k := range keys {
		encryptedData, nonce, err := c.prepareTLFCryptKeyClientHalf(k.ClientHalf)
		if err != nil {
			errors = append(errors, err)
		} else {
			bundles = append(bundles, keybase1.CiphertextBundle{
				Kid:        k.PubKey.kidContainer.kid,
				Ciphertext: encryptedData,
				Nonce:      nonce,
				PublicKey:  keybase1.BoxPublicKey(k.EPubKey.data),
			})
			indexLookup = append(indexLookup, i)
		}
	}
	if len(bundles) == 0 {
		err = errors[0]
		return
	}
	timer := c.logAboutLongRPCUnlessCancelled(ctx, "UnboxBytes32Any")
	defer timer.Stop()
	res, err := c.client.UnboxBytes32Any(ctx, keybase1.UnboxBytes32AnyArg{
		Bundles:     bundles,
		Reason:      "to rekey for kbfs",
		PromptPaper: promptPaper,
	})
	if err != nil {
		return
	}
	return MakeTLFCryptKeyClientHalf(res.Plaintext), indexLookup[res.Index], nil
}

// Shutdown implements the Crypto interface for CryptoClient.
func (c *CryptoClient) Shutdown() {
	if c.shutdownFn != nil {
		c.shutdownFn()
	}
}
<|endoftext|>


<commit_before>package todotxt

import (
        "time"
        "os"
        "bufio"
        "strings"
        "regexp"
        "sort"
        "unicode"
        "fmt"
)

type Task struct {
        id int
        todo string
        priority byte
        create_date time.Time
        contexts []string
        projects []string
        raw_todo string
        finished bool
        finish_date time.Time
        id_padding int
}

type TaskList []Task

func ParseTask(text string, id int) (Task) {
        var task = Task{}
        task.id = id
        task.raw_todo = text

        splits := strings.Split(text, " ")

        // checking if the task is already finished
        if text[0] == 'x' &&
           text[1] == ' ' &&
           !unicode.IsSpace(rune(text[2])) {
                task.finished = true
                splits = splits[1:]
        }

        date_regexp := "([\\d]{4})-([\\d]{2})-([\\d]{2})"

        // checking for finish date
        if match, _ := regexp.MatchString(date_regexp, splits[0]); match {
                if date, e := time.Parse("2006-01-02", splits[0]); e != nil {
                        panic(e)
                } else {
                        task.finish_date = date
                }

                splits = splits[1:]
        }

        head := splits[0]

        // checking for priority
        if (len(head) == 3) &&
           (head[0] == '(') &&
           (head[2] == ')') &&
           (head[1] >= 65 && head[1] <= 90) { // checking if it's in range [A-Z]
                task.priority = head[1]
                splits = splits[1:]
        }

        // checking for creation date and building the actual todo item
        if match, _ := regexp.MatchString(date_regexp, splits[0]); match {
                if date, e := time.Parse("2006-01-02", splits[0]); e != nil {
                        panic(e)
                } else {
                        task.create_date = date
                }

                task.todo = strings.Join(splits[1:], " ")
        } else {
                task.todo = strings.Join(splits[0:], " ")
        }

        context_regexp, _ := regexp.Compile("@[[:word:]]+")
        contexts := context_regexp.FindAllStringSubmatch(text, -1)
        if len(contexts) != 0 {
                task.contexts = contexts[0]
        }

        project_regexp, _ := regexp.Compile("\\+[[:word:]]+")
        projects := project_regexp.FindAllStringSubmatch(text, -1)
        if len(projects) != 0 {
                task.projects = projects[0]
        }

        return task
}

func LoadTaskList (filename string) (TaskList) {

        var f, err = os.Open(filename)

        if err != nil {
                panic(err)
        }

        defer f.Close()

        var tasklist = TaskList{}

        scanner := bufio.NewScanner(f)

        for scanner.Scan() {
                text := scanner.Text()
                tasklist.Add(text)
        }

        if err := scanner.Err(); err != nil {
                panic(scanner.Err())
        }

        return tasklist
}

type By func(t1, t2 Task) bool

func (by By) Sort(tasks TaskList) {
        ts := &taskSorter{
                tasks: tasks,
                by:    by,
        }
        sort.Sort(ts)
}

type taskSorter struct {
        tasks TaskList
        by func(t1, t2 Task) bool
}

func (s *taskSorter) Len() int {
        return len(s.tasks)
}

func (s *taskSorter) Swap(i, j int) {
        s.tasks[i], s.tasks[j] = s.tasks[j], s.tasks[i]
}

func (s *taskSorter) Less(i, j int) bool {
        return s.by(s.tasks[i], s.tasks[j])
}

func (tasks TaskList) Len() int {
        return len(tasks)
}

func prioCmp(t1, t2 Task) bool {
        return t1.Priority() < t2.Priority()
}

func prioRevCmp(t1, t2 Task) bool {
        return t1.Priority() > t2.Priority()
}

func dateCmp(t1, t2 Task) bool {
        tm1 := t1.CreateDate().Unix()
        tm2 := t2.CreateDate().Unix()

        // if the dates equal, let's use priority
        if tm1 == tm2 {
                return prioCmp(t1, t2)
        } else {
                return tm1 > tm2
        }
}

func dateRevCmp(t1, t2 Task) bool {
        tm1 := t1.CreateDate().Unix()
        tm2 := t2.CreateDate().Unix()

        // if the dates equal, let's use priority
        if tm1 == tm2 {
                return prioCmp(t1, t2)
        } else {
                return tm1 < tm2
        }
}

func lenCmp(t1, t2 Task) bool {
        tl1 := len(t1.raw_todo)
        tl2 := len(t2.raw_todo)
        if tl1 == tl2 {
                return prioCmp(t1, t2)
        } else {
                return tl1 < tl2
        }
}

func lenRevCmp(t1, t2 Task) bool {
        tl1 := len(t1.raw_todo)
        tl2 := len(t2.raw_todo)
        if tl1 == tl2 {
                return prioCmp(t1, t2)
        } else {
                return tl1 > tl2
        }
}

func idCmp(t1, t2 Task) bool {
        return t1.Id() < t2.Id()
}

func (tasks TaskList) Sort(by string) {
        switch by {
        default:
        case "prio":
                By(prioCmp).Sort(tasks)
        case "prio-rev":
                By(prioRevCmp).Sort(tasks)
        case "date":
                By(dateCmp).Sort(tasks)
        case "date-rev":
                By(dateRevCmp).Sort(tasks)
        case "len":
                By(lenCmp).Sort(tasks)
        case "len-rev":
                By(lenRevCmp).Sort(tasks)
        case "id":
                By(idCmp).Sort(tasks)
        }
}

func (tasks TaskList) Save(filename string) {
        tasks.Sort("id")

        f, err := os.Create(filename)
        if err != nil {
                panic(err)
        }

        defer f.Close()

        for _, task := range tasks {
                f.WriteString(task.RawText() + "\n")
        }
        f.Sync()
}

func (tasks *TaskList) Add(todo string) {
        task := ParseTask(todo, tasks.Len())
        *tasks = append(*tasks, task)
}

func (tasks TaskList) Done(id int, finish_date bool) error {
        if id > tasks.Len() || id < 0 {
                return fmt.Errorf("Error: id is %v", id)
        }

        tasks[id].finished = true
        if finish_date {
                t := time.Now()
                tasks[id].raw_todo = "x " + t.Format("2006-01-02") + " " +
                                        tasks[id].raw_todo
        } else {
                tasks[id].raw_todo = "x " + tasks[id].raw_todo
        }

        return nil
}

func (task Task) Id() int {
        return task.id
}

func (task Task) Text() string {
        return task.todo
}

func (task Task) RawText() string {
        return task.raw_todo
}

func (task Task) Priority() byte {
        // if priority is not from [A-Z], let it be 94 (^)
        if task.priority < 65 || task.priority > 90 {
                return 94 // you know, ^
        } else {
                return task.priority
        }
}

func (task Task) Contexts() []string {
        return task.contexts
}

func (task Task) Projects() []string {
        return task.projects
}

func (task Task) CreateDate() time.Time {
        return task.create_date
}

func (task Task) Finished() bool {
        return task.finished
}

func (task Task) FinishDate() time.Time {
        return task.finish_date
}

func (task *Task) SetIdPaddingBy(tasklist TaskList) {
        l := tasklist.Len()

        if l >= 10000 {
                task.id_padding = 5
        } else if l >= 1000 {
                task.id_padding = 4
        } else if l >= 100 {
                task.id_padding = 3
        } else if l >= 10 {
                task.id_padding = 2
        } else {
                task.id_padding = 1
        }
}

func (task *Task) RebuildRawTodo() {
        if task.finished {
                task.raw_todo = task.PrettyPrint("x %P%t")
        } else {
                task.raw_todo = task.PrettyPrint("%P%t")
        }
}

func (task *Task) SetPriority(prio byte) {
        if task.priority < 65 || task.priority > 90 {
                task.priority = '^'
        } else {
                task.priority = prio
        }
}

func (task *Task) SetTodo(todo string) {
        task.todo = todo
}

func (task Task) IdPadding() int {
        return task.id_padding
}

func (task Task) PrettyPrint(pretty string) string {
        rp := regexp.MustCompile("(%[a-zA-Z])")
        out := rp.ReplaceAllStringFunc(pretty, func(s string) string {

                switch s{
                case "%i":
                        str := fmt.Sprintf("%%0%dd", task.IdPadding())
                        return fmt.Sprintf(str, task.Id())
                case "%t":
                        return task.Text()
                case "%T":
                        return task.RawText()
                case "%p":
                        return string(task.Priority())
                case "%P":
                        if task.Priority() != '^' {
                                return "(" + string(task.Priority()) + ") "
                        } else {
                                return ""
                        }
                default:
                        return s
                }
        })
        return out
}
<commit_msg>rand sort<commit_after>package todotxt

import (
        "time"
        "os"
        "bufio"
        "strings"
        "regexp"
        "sort"
        "unicode"
        "fmt"
)

type Task struct {
        id int
        todo string
        priority byte
        create_date time.Time
        contexts []string
        projects []string
        raw_todo string
        finished bool
        finish_date time.Time
        id_padding int
}

type TaskList []Task

func ParseTask(text string, id int) (Task) {
        var task = Task{}
        task.id = id
        task.raw_todo = text

        splits := strings.Split(text, " ")

        // checking if the task is already finished
        if text[0] == 'x' &&
           text[1] == ' ' &&
           !unicode.IsSpace(rune(text[2])) {
                task.finished = true
                splits = splits[1:]
        }

        date_regexp := "([\\d]{4})-([\\d]{2})-([\\d]{2})"

        // checking for finish date
        if match, _ := regexp.MatchString(date_regexp, splits[0]); match {
                if date, e := time.Parse("2006-01-02", splits[0]); e != nil {
                        panic(e)
                } else {
                        task.finish_date = date
                }

                splits = splits[1:]
        }

        head := splits[0]

        // checking for priority
        if (len(head) == 3) &&
           (head[0] == '(') &&
           (head[2] == ')') &&
           (head[1] >= 65 && head[1] <= 90) { // checking if it's in range [A-Z]
                task.priority = head[1]
                splits = splits[1:]
        }

        // checking for creation date and building the actual todo item
        if match, _ := regexp.MatchString(date_regexp, splits[0]); match {
                if date, e := time.Parse("2006-01-02", splits[0]); e != nil {
                        panic(e)
                } else {
                        task.create_date = date
                }

                task.todo = strings.Join(splits[1:], " ")
        } else {
                task.todo = strings.Join(splits[0:], " ")
        }

        context_regexp, _ := regexp.Compile("@[[:word:]]+")
        contexts := context_regexp.FindAllStringSubmatch(text, -1)
        if len(contexts) != 0 {
                task.contexts = contexts[0]
        }

        project_regexp, _ := regexp.Compile("\\+[[:word:]]+")
        projects := project_regexp.FindAllStringSubmatch(text, -1)
        if len(projects) != 0 {
                task.projects = projects[0]
        }

        return task
}

func LoadTaskList (filename string) (TaskList) {

        var f, err = os.Open(filename)

        if err != nil {
                panic(err)
        }

        defer f.Close()

        var tasklist = TaskList{}

        scanner := bufio.NewScanner(f)

        for scanner.Scan() {
                text := scanner.Text()
                tasklist.Add(text)
        }

        if err := scanner.Err(); err != nil {
                panic(scanner.Err())
        }

        return tasklist
}

type By func(t1, t2 Task) bool

func (by By) Sort(tasks TaskList) {
        ts := &taskSorter{
                tasks: tasks,
                by:    by,
        }
        sort.Sort(ts)
}

type taskSorter struct {
        tasks TaskList
        by func(t1, t2 Task) bool
}

func (s *taskSorter) Len() int {
        return len(s.tasks)
}

func (s *taskSorter) Swap(i, j int) {
        s.tasks[i], s.tasks[j] = s.tasks[j], s.tasks[i]
}

func (s *taskSorter) Less(i, j int) bool {
        return s.by(s.tasks[i], s.tasks[j])
}

func (tasks TaskList) Len() int {
        return len(tasks)
}

func prioCmp(t1, t2 Task) bool {
        return t1.Priority() < t2.Priority()
}

func prioRevCmp(t1, t2 Task) bool {
        return t1.Priority() > t2.Priority()
}

func dateCmp(t1, t2 Task) bool {
        tm1 := t1.CreateDate().Unix()
        tm2 := t2.CreateDate().Unix()

        // if the dates equal, let's use priority
        if tm1 == tm2 {
                return prioCmp(t1, t2)
        } else {
                return tm1 > tm2
        }
}

func dateRevCmp(t1, t2 Task) bool {
        tm1 := t1.CreateDate().Unix()
        tm2 := t2.CreateDate().Unix()

        // if the dates equal, let's use priority
        if tm1 == tm2 {
                return prioCmp(t1, t2)
        } else {
                return tm1 < tm2
        }
}

func lenCmp(t1, t2 Task) bool {
        tl1 := len(t1.raw_todo)
        tl2 := len(t2.raw_todo)
        if tl1 == tl2 {
                return prioCmp(t1, t2)
        } else {
                return tl1 < tl2
        }
}

func lenRevCmp(t1, t2 Task) bool {
        tl1 := len(t1.raw_todo)
        tl2 := len(t2.raw_todo)
        if tl1 == tl2 {
                return prioCmp(t1, t2)
        } else {
                return tl1 > tl2
        }
}

func idCmp(t1, t2 Task) bool {
        return t1.Id() < t2.Id()
}

func randCmp(t1, t2 Taks) bool {
        return rand.Intn(t1.Id()) > rand.Intn(t1.Id())
}

func (tasks TaskList) Sort(by string) {
        switch by {
        default:
        case "prio":
                By(prioCmp).Sort(tasks)
        case "prio-rev":
                By(prioRevCmp).Sort(tasks)
        case "date":
                By(dateCmp).Sort(tasks)
        case "date-rev":
                By(dateRevCmp).Sort(tasks)
        case "len":
                By(lenCmp).Sort(tasks)
        case "len-rev":
                By(lenRevCmp).Sort(tasks)
        case "id":
                By(idCmp).Sort(tasks)
        case "rand":
                By(randCmp).Sort(tasks)
        }
}

func (tasks TaskList) Save(filename string) {
        tasks.Sort("id")

        f, err := os.Create(filename)
        if err != nil {
                panic(err)
        }

        defer f.Close()

        for _, task := range tasks {
                f.WriteString(task.RawText() + "\n")
        }
        f.Sync()
}

func (tasks *TaskList) Add(todo string) {
        task := ParseTask(todo, tasks.Len())
        *tasks = append(*tasks, task)
}

func (tasks TaskList) Done(id int, finish_date bool) error {
        if id > tasks.Len() || id < 0 {
                return fmt.Errorf("Error: id is %v", id)
        }

        tasks[id].finished = true
        if finish_date {
                t := time.Now()
                tasks[id].raw_todo = "x " + t.Format("2006-01-02") + " " +
                                        tasks[id].raw_todo
        } else {
                tasks[id].raw_todo = "x " + tasks[id].raw_todo
        }

        return nil
}

func (task Task) Id() int {
        return task.id
}

func (task Task) Text() string {
        return task.todo
}

func (task Task) RawText() string {
        return task.raw_todo
}

func (task Task) Priority() byte {
        // if priority is not from [A-Z], let it be 94 (^)
        if task.priority < 65 || task.priority > 90 {
                return 94 // you know, ^
        } else {
                return task.priority
        }
}

func (task Task) Contexts() []string {
        return task.contexts
}

func (task Task) Projects() []string {
        return task.projects
}

func (task Task) CreateDate() time.Time {
        return task.create_date
}

func (task Task) Finished() bool {
        return task.finished
}

func (task Task) FinishDate() time.Time {
        return task.finish_date
}

func (task *Task) SetIdPaddingBy(tasklist TaskList) {
        l := tasklist.Len()

        if l >= 10000 {
                task.id_padding = 5
        } else if l >= 1000 {
                task.id_padding = 4
        } else if l >= 100 {
                task.id_padding = 3
        } else if l >= 10 {
                task.id_padding = 2
        } else {
                task.id_padding = 1
        }
}

func (task *Task) RebuildRawTodo() {
        if task.finished {
                task.raw_todo = task.PrettyPrint("x %P%t")
        } else {
                task.raw_todo = task.PrettyPrint("%P%t")
        }
}

func (task *Task) SetPriority(prio byte) {
        if task.priority < 65 || task.priority > 90 {
                task.priority = '^'
        } else {
                task.priority = prio
        }
}

func (task *Task) SetTodo(todo string) {
        task.todo = todo
}

func (task Task) IdPadding() int {
        return task.id_padding
}

func (task Task) PrettyPrint(pretty string) string {
        rp := regexp.MustCompile("(%[a-zA-Z])")
        out := rp.ReplaceAllStringFunc(pretty, func(s string) string {

                switch s{
                case "%i":
                        str := fmt.Sprintf("%%0%dd", task.IdPadding())
                        return fmt.Sprintf(str, task.Id())
                case "%t":
                        return task.Text()
                case "%T":
                        return task.RawText()
                case "%p":
                        return string(task.Priority())
                case "%P":
                        if task.Priority() != '^' {
                                return "(" + string(task.Priority()) + ") "
                        } else {
                                return ""
                        }
                default:
                        return s
                }
        })
        return out
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"runtime"
	"runtime/pprof"
	"strings"
	"time"

	"github.com/Dieterbe/profiletrigger/cpu"
	"github.com/Dieterbe/profiletrigger/heap"
	"github.com/raintank/dur"
	"github.com/raintank/statsdaemon"
	"github.com/raintank/statsdaemon/out"

	"net/http"
	_ "net/http/pprof"

	"github.com/rakyll/globalconf"
)

const (
	VERSION = "0.6"
	// number of packets we can read out of udp buffer without processing them
	// statsdaemon doesn't really interrupt the udp reader like some other statsd's do (like on flush)
	// but this can still be useful to deal with traffic bursts.
	// keep in mind that one metric is about 30 to 100 bytes of memory.
	MAX_UNPROCESSED_PACKETS = 1000
)

var (
	listen_addr     = flag.String("listen_addr", ":8125", "listener address for statsd, listens on UDP only.default: :8125")
	admin_addr      = flag.String("admin_addr", ":8126", "listener address for admin port, default: :8126")
	profile_addr    = flag.String("profile_addr", "", "listener address for profiler, must be in format :9000")
	graphite_addr   = flag.String("graphite_addr", "127.0.0.1:2003", "graphite carbon-in url, default: 127.0.0.1:2003")
	flushInterval   = flag.Int("flush_interval", 10, "flush interval in seconds, default: 10")
	processes       = flag.Int("processes", 1, "number of processes to use, default: 1")
	instance        = flag.String("instance", "null", "instance name, defaults to short hostname if not set")
	prefix_counters = flag.String("prefix_counters", "stats_counts.", "counters prefix, default: stats.counters")
	prefix_gauges   = flag.String("prefix_gauges", "stats.gauges.", "gauges prefix, default: stats.gauges")
	prefix_rates    = flag.String("prefix_rates", "stats.", "rates prefix, default: stats. it is recommended that you use stats.rates if possible")
	prefix_timers   = flag.String("prefix_timers", "stats.timers.", "timers prefix, default: stats.timers")

	prefix_m20_counters = flag.String("prefix_m20_counters", "", "counters 2.0 prefix, default: nil")
	prefix_m20_gauges   = flag.String("prefix_m20_gauges", "", "gauges 2.0 prefix, default: nil")
	prefix_m20_rates    = flag.String("prefix_m20_rates", "", "rates 2.0 prefix, default: nil")
	prefix_m20_timers   = flag.String("prefix_m20_timers", "", "timers 2.0 prefix, default: nil")

	legacy_namespace = flag.Bool("legacy_namespace", true, "legacy namespacing (not recommended), default: true")
	flush_rates      = flag.Bool("flush_rates", true, "send count for counters (using prefix_counters), default: true")
	flush_counts     = flag.Bool("flush_counts", false, "send count for counters (using prefix_counters), default: false")

	percentile_thresholds = flag.String("percentile_thresholds", "", "percential thresholds (used by timers), default: nil")
	max_timers_per_s      = flag.Uint64("max_timers_per_s", 1000, "max timers per second, default: 1000")

	proftrigPath = flag.String("proftrigger_path", "/tmp", "profiler file path, default: /tmp") // "path to store triggered profiles"

	proftrigHeapFreqStr    = flag.String("proftrigger_heap_freq", "60s", "profiler heap frequency, default: 60s")   // "inspect status frequency. set to 0 to disable"
	proftrigHeapMinDiffStr = flag.String("proftrigger_heap_min_diff", "1h", "profiler heap min difference, default: 1h") // "minimum time between triggered profiles"
	proftrigHeapThresh     = flag.Int("proftrigger_heap_thresh", 10000000, "profiler heap threshold, default: 10000000")  // "if this many bytes allocated, trigger a profile"

	proftrigCpuFreqStr    = flag.String("proftrigger_cpu_freq", "60s", "profiler cpu frequency, default: 60s")    // "inspect status frequency. set to 0 to disable"
	proftrigCpuMinDiffStr = flag.String("proftrigger_cpu_min_diff", "1h", "profiler cpu min difference, default: 1h") // "minimum time between triggered profiles"
	proftrigCpuDurStr     = flag.String("proftrigger_cpu_dur", "5s", "profiler cpu duration, default 5s")      // "duration of cpu profile"
	proftrigCpuThresh     = flag.Int("proftrigger_cpu_thresh", 80, "profiler cpu threshold, 80")        // "if this much percent cpu used, trigger a profile"

	debug       = flag.Bool("debug", false, "log outgoing metrics, bad lines, and received admin commands")
	showVersion = flag.Bool("version", false, "print version string")
	config_file = flag.String("config_file", "/etc/statsdaemon.ini", "config file location")
	cpuprofile  = flag.String("cpuprofile", "", "write cpu profile to file")
	memprofile  = flag.String("memprofile", "", "write memory profile to this file")
	GitHash     = "(none)"
)

func expand_cfg_vars(in string) (out string) {
	switch in {
	case "HOST":
		hostname, _ := os.Hostname()
		// in case hostname is an fqdn or has dots, only take first part
		parts := strings.SplitN(hostname, ".", 2)
		return parts[0]
	default:
		return ""
	}
}
func main() {
	flag.Parse()

	if *showVersion {
		fmt.Printf("statsdaemon v%s (built w/%s, git hash %s)\n", VERSION, runtime.Version(), GitHash)
		return
	}
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}
	if *memprofile != "" {
		f, err := os.Create(*memprofile)
		if err != nil {
			log.Fatal(err)
		}
		defer f.Close()
		defer pprof.WriteHeapProfile(f)
	}

	path := ""
	if _, err := os.Stat(*config_file); err == nil {
		path = *config_file
        }
	conf, err := globalconf.NewWithOptions(&globalconf.Options{
		Filename:  path,
		EnvPrefix: "SD_",
        })

	conf.ParseAll()

	proftrigHeapFreq := dur.MustParseUsec("proftrigger_heap_freq", *proftrigHeapFreqStr)
	proftrigHeapMinDiff := int(dur.MustParseUNsec("proftrigger_heap_min_diff", *proftrigHeapMinDiffStr))

	proftrigCpuFreq := dur.MustParseUsec("proftrigger_cpu_freq", *proftrigCpuFreqStr)
	proftrigCpuMinDiff := int(dur.MustParseUNsec("proftrigger_cpu_min_diff", *proftrigCpuMinDiffStr))
	proftrigCpuDur := int(dur.MustParseUNsec("proftrigger_cpu_dur", *proftrigCpuDurStr))

	if proftrigHeapFreq > 0 {
		errors := make(chan error)
		trigger, _ := heap.New(*proftrigPath, *proftrigHeapThresh, proftrigHeapMinDiff, time.Duration(proftrigHeapFreq)*time.Second, errors)
		go func() {
			for e := range errors {
				log.Printf("profiletrigger heap: %s", e)
			}
		}()
		go trigger.Run()
	}

	if proftrigCpuFreq > 0 {
		errors := make(chan error)
		freq := time.Duration(proftrigCpuFreq) * time.Second
		duration := time.Duration(proftrigCpuDur) * time.Second
		trigger, _ := cpu.New(*proftrigPath, *proftrigCpuThresh, proftrigCpuMinDiff, freq, duration, errors)
		go func() {
			for e := range errors {
				log.Printf("profiletrigger cpu: %s", e)
			}
		}()
		go trigger.Run()
	}

	runtime.GOMAXPROCS(*processes)
	pct, err := out.NewPercentiles(*percentile_thresholds)
	if err != nil {
		log.Fatal(err)
	}
	inst := os.Expand(*instance, expand_cfg_vars)
	if inst == "" {
		inst = "null"
	}

	signalchan := make(chan os.Signal, 1)
	signal.Notify(signalchan)
	if *profile_addr != "" {
		go func() {
			fmt.Println("Profiling endpoint listening on " + *profile_addr)
			log.Println(http.ListenAndServe(*profile_addr, nil))
		}()
	}

	formatter := out.Formatter{
		PrefixInternal: "service_is_statsdaemon.instance_is_" + inst + ".",

		Legacy_namespace: *legacy_namespace,
		Prefix_counters:  *prefix_counters,
		Prefix_gauges:    *prefix_gauges,
		Prefix_rates:     *prefix_rates,
		Prefix_timers:    *prefix_timers,

		Prefix_m20_counters: *prefix_m20_counters,
		Prefix_m20_gauges:   *prefix_m20_gauges,
		Prefix_m20_rates:    *prefix_m20_rates,
		Prefix_m20_timers:   *prefix_m20_timers,

		Prefix_m20ne_counters: strings.Replace(*prefix_m20_counters, "=", "_is_", -1),
		Prefix_m20ne_gauges:   strings.Replace(*prefix_m20_gauges, "=", "_is_", -1),
		Prefix_m20ne_rates:    strings.Replace(*prefix_m20_rates, "=", "_is_", -1),
		Prefix_m20ne_timers:   strings.Replace(*prefix_m20_timers, "=", "_is_", -1),
	}

	daemon := statsdaemon.New(inst, formatter, *flush_rates, *flush_counts, *pct, *flushInterval, MAX_UNPROCESSED_PACKETS, *max_timers_per_s, *debug, signalchan)
	if *debug {
		consumer := make(chan interface{}, 100)
		daemon.Invalid_lines.Register(consumer)
		go func() {
			for line := range consumer {
				log.Printf("invalid line '%s'\n", line)
			}
		}()
	}
	daemon.Run(*listen_addr, *admin_addr, *graphite_addr)
}
<commit_msg>fix the flag format, these already include default<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"runtime"
	"runtime/pprof"
	"strings"
	"time"

	"github.com/Dieterbe/profiletrigger/cpu"
	"github.com/Dieterbe/profiletrigger/heap"
	"github.com/raintank/dur"
	"github.com/raintank/statsdaemon"
	"github.com/raintank/statsdaemon/out"

	"net/http"
	_ "net/http/pprof"

	"github.com/rakyll/globalconf"
)

const (
	VERSION = "0.6"
	// number of packets we can read out of udp buffer without processing them
	// statsdaemon doesn't really interrupt the udp reader like some other statsd's do (like on flush)
	// but this can still be useful to deal with traffic bursts.
	// keep in mind that one metric is about 30 to 100 bytes of memory.
	MAX_UNPROCESSED_PACKETS = 1000
)

var (
	listen_addr     = flag.String("listen_addr", ":8125", "listener address for statsd, listens on UDP only")
	admin_addr      = flag.String("admin_addr", ":8126", "listener address for admin port")
	profile_addr    = flag.String("profile_addr", "", "listener address for profiler")
	graphite_addr   = flag.String("graphite_addr", "127.0.0.1:2003", "graphite carbon-in url")
	flushInterval   = flag.Int("flush_interval", 10, "flush interval in seconds")
	processes       = flag.Int("processes", 1, "number of processes to use")
	instance        = flag.String("instance", "$HOST", "instance name, defaults to short hostname if not set")
	prefix_counters = flag.String("prefix_counters", "stats_counts.", "counters prefix")
	prefix_gauges   = flag.String("prefix_gauges", "stats.gauges.", "gauges prefix")
	prefix_rates    = flag.String("prefix_rates", "stats.", "rates prefix, it is recommended that you use stats.rates if possible")
	prefix_timers   = flag.String("prefix_timers", "stats.timers.", "timers prefix")

	prefix_m20_counters = flag.String("prefix_m20_counters", "", "counters 2.0 prefix")
	prefix_m20_gauges   = flag.String("prefix_m20_gauges", "", "gauges 2.0 prefix")
	prefix_m20_rates    = flag.String("prefix_m20_rates", "", "rates 2.0 prefix")
	prefix_m20_timers   = flag.String("prefix_m20_timers", "", "timers 2.0 prefix")

	legacy_namespace = flag.Bool("legacy_namespace", true, "legacy namespacing (not recommended)")
	flush_rates      = flag.Bool("flush_rates", true, "send count for counters (using prefix_counters)")
	flush_counts     = flag.Bool("flush_counts", false, "send count for counters (using prefix_counters)")

	percentile_thresholds = flag.String("percentile_thresholds", "", "percential thresholds (used by timers)")
	max_timers_per_s      = flag.Uint64("max_timers_per_s", 1000, "max timers per second")

	proftrigPath = flag.String("proftrigger_path", "/tmp", "profiler file path") // "path to store triggered profiles"

	proftrigHeapFreqStr    = flag.String("proftrigger_heap_freq", "60s", "profiler heap frequency")   // "inspect status frequency. set to 0 to disable"
	proftrigHeapMinDiffStr = flag.String("proftrigger_heap_min_diff", "1h", "profiler heap min difference") // "minimum time between triggered profiles"
	proftrigHeapThresh     = flag.Int("proftrigger_heap_thresh", 10000000, "profiler heap threshold")  // "if this many bytes allocated, trigger a profile"

	proftrigCpuFreqStr    = flag.String("proftrigger_cpu_freq", "60s", "profiler cpu frequency")    // "inspect status frequency. set to 0 to disable"
	proftrigCpuMinDiffStr = flag.String("proftrigger_cpu_min_diff", "1h", "profiler cpu min difference") // "minimum time between triggered profiles"
	proftrigCpuDurStr     = flag.String("proftrigger_cpu_dur", "5s", "profiler cpu duration")      // "duration of cpu profile"
	proftrigCpuThresh     = flag.Int("proftrigger_cpu_thresh", 80, "profiler cpu threshold")        // "if this much percent cpu used, trigger a profile"

	debug       = flag.Bool("debug", false, "log outgoing metrics, bad lines, and received admin commands")
	showVersion = flag.Bool("version", false, "print version string")
	config_file = flag.String("config_file", "/etc/statsdaemon.ini", "config file location")
	cpuprofile  = flag.String("cpuprofile", "", "write cpu profile to file")
	memprofile  = flag.String("memprofile", "", "write memory profile to this file")
	GitHash     = "(none)"
)

func expand_cfg_vars(in string) (out string) {
	switch in {
	case "HOST":
		hostname, _ := os.Hostname()
		// in case hostname is an fqdn or has dots, only take first part
		parts := strings.SplitN(hostname, ".", 2)
		return parts[0]
	default:
		return ""
	}
}
func main() {
	flag.Parse()

	if *showVersion {
		fmt.Printf("statsdaemon v%s (built w/%s, git hash %s)\n", VERSION, runtime.Version(), GitHash)
		return
	}
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}
	if *memprofile != "" {
		f, err := os.Create(*memprofile)
		if err != nil {
			log.Fatal(err)
		}
		defer f.Close()
		defer pprof.WriteHeapProfile(f)
	}

	path := ""
	if _, err := os.Stat(*config_file); err == nil {
		path = *config_file
        }
	conf, err := globalconf.NewWithOptions(&globalconf.Options{
		Filename:  path,
		EnvPrefix: "SD_",
        })

	conf.ParseAll()

	proftrigHeapFreq := dur.MustParseUsec("proftrigger_heap_freq", *proftrigHeapFreqStr)
	proftrigHeapMinDiff := int(dur.MustParseUNsec("proftrigger_heap_min_diff", *proftrigHeapMinDiffStr))

	proftrigCpuFreq := dur.MustParseUsec("proftrigger_cpu_freq", *proftrigCpuFreqStr)
	proftrigCpuMinDiff := int(dur.MustParseUNsec("proftrigger_cpu_min_diff", *proftrigCpuMinDiffStr))
	proftrigCpuDur := int(dur.MustParseUNsec("proftrigger_cpu_dur", *proftrigCpuDurStr))

	if proftrigHeapFreq > 0 {
		errors := make(chan error)
		trigger, _ := heap.New(*proftrigPath, *proftrigHeapThresh, proftrigHeapMinDiff, time.Duration(proftrigHeapFreq)*time.Second, errors)
		go func() {
			for e := range errors {
				log.Printf("profiletrigger heap: %s", e)
			}
		}()
		go trigger.Run()
	}

	if proftrigCpuFreq > 0 {
		errors := make(chan error)
		freq := time.Duration(proftrigCpuFreq) * time.Second
		duration := time.Duration(proftrigCpuDur) * time.Second
		trigger, _ := cpu.New(*proftrigPath, *proftrigCpuThresh, proftrigCpuMinDiff, freq, duration, errors)
		go func() {
			for e := range errors {
				log.Printf("profiletrigger cpu: %s", e)
			}
		}()
		go trigger.Run()
	}

	runtime.GOMAXPROCS(*processes)
	pct, err := out.NewPercentiles(*percentile_thresholds)
	if err != nil {
		log.Fatal(err)
	}
	inst := os.Expand(*instance, expand_cfg_vars)
	if inst == "" {
		inst = "null"
	}

	signalchan := make(chan os.Signal, 1)
	signal.Notify(signalchan)
	if *profile_addr != "" {
		go func() {
			fmt.Println("Profiling endpoint listening on " + *profile_addr)
			log.Println(http.ListenAndServe(*profile_addr, nil))
		}()
	}

	formatter := out.Formatter{
		PrefixInternal: "service_is_statsdaemon.instance_is_" + inst + ".",

		Legacy_namespace: *legacy_namespace,
		Prefix_counters:  *prefix_counters,
		Prefix_gauges:    *prefix_gauges,
		Prefix_rates:     *prefix_rates,
		Prefix_timers:    *prefix_timers,

		Prefix_m20_counters: *prefix_m20_counters,
		Prefix_m20_gauges:   *prefix_m20_gauges,
		Prefix_m20_rates:    *prefix_m20_rates,
		Prefix_m20_timers:   *prefix_m20_timers,

		Prefix_m20ne_counters: strings.Replace(*prefix_m20_counters, "=", "_is_", -1),
		Prefix_m20ne_gauges:   strings.Replace(*prefix_m20_gauges, "=", "_is_", -1),
		Prefix_m20ne_rates:    strings.Replace(*prefix_m20_rates, "=", "_is_", -1),
		Prefix_m20ne_timers:   strings.Replace(*prefix_m20_timers, "=", "_is_", -1),
	}

	daemon := statsdaemon.New(inst, formatter, *flush_rates, *flush_counts, *pct, *flushInterval, MAX_UNPROCESSED_PACKETS, *max_timers_per_s, *debug, signalchan)
	if *debug {
		consumer := make(chan interface{}, 100)
		daemon.Invalid_lines.Register(consumer)
		go func() {
			for line := range consumer {
				log.Printf("invalid line '%s'\n", line)
			}
		}()
	}
	daemon.Run(*listen_addr, *admin_addr, *graphite_addr)
}
<|endoftext|>


<commit_before>package status

import (
	"crypto/md5"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
)

// Checksum validades if a task is up to date by calculating its source
// files checksum
type Checksum struct {
	Dir     string
	Task    string
	Sources []string
}

// IsUpToDate implements the Checker interface
func (c *Checksum) IsUpToDate() (bool, error) {
	checksumFile := filepath.Join(c.Dir, ".task", c.Task)

	data, _ := ioutil.ReadFile(checksumFile)
	oldMd5 := strings.TrimSpace(string(data))

	sources, err := glob(c.Dir, c.Sources)
	if err != nil {
		return false, err
	}

	newMd5, err := c.checksum(sources...)
	if err != nil {
		return false, nil
	}

	_ = os.MkdirAll(filepath.Join(c.Dir, ".task"), 0755)
	if err = ioutil.WriteFile(checksumFile, []byte(newMd5), 0644); err != nil {
		return false, err
	}
	return oldMd5 == newMd5, nil
}

func (c *Checksum) checksum(files ...string) (string, error) {
	h := md5.New()

	for _, f := range files {
		f, err := os.Open(f)
		if err != nil {
			return "", err
		}
		if _, err := io.Copy(h, f); err != nil {
			return "", err
		}
	}

	return fmt.Sprintf("%x", h.Sum(nil)), nil
}

// OnError implements the Checker interface
func (c *Checksum) OnError() error {
	return os.Remove(filepath.Join(c.Dir, ".task", c.Task))
}
<commit_msg>checksum: skip directories<commit_after>package status

import (
	"crypto/md5"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
)

// Checksum validades if a task is up to date by calculating its source
// files checksum
type Checksum struct {
	Dir     string
	Task    string
	Sources []string
}

// IsUpToDate implements the Checker interface
func (c *Checksum) IsUpToDate() (bool, error) {
	checksumFile := filepath.Join(c.Dir, ".task", c.Task)

	data, _ := ioutil.ReadFile(checksumFile)
	oldMd5 := strings.TrimSpace(string(data))

	sources, err := glob(c.Dir, c.Sources)
	if err != nil {
		return false, err
	}

	newMd5, err := c.checksum(sources...)
	if err != nil {
		return false, nil
	}

	_ = os.MkdirAll(filepath.Join(c.Dir, ".task"), 0755)
	if err = ioutil.WriteFile(checksumFile, []byte(newMd5), 0644); err != nil {
		return false, err
	}
	return oldMd5 == newMd5, nil
}

func (c *Checksum) checksum(files ...string) (string, error) {
	h := md5.New()

	for _, f := range files {
		f, err := os.Open(f)
		if err != nil {
			return "", err
		}
		info, err := f.Stat()
		if err != nil {
			return "", err
		}
		if info.IsDir() {
			continue
		}
		if _, err := io.Copy(h, f); err != nil {
			return "", err
		}
	}

	return fmt.Sprintf("%x", h.Sum(nil)), nil
}

// OnError implements the Checker interface
func (c *Checksum) OnError() error {
	return os.Remove(filepath.Join(c.Dir, ".task", c.Task))
}
<|endoftext|>


<commit_before>// Copyright 2020 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package applog

import (
	"sort"
	"sync"
	"sync/atomic"

	"github.com/tsuru/tsuru/provision"
	"github.com/tsuru/tsuru/provision/pool"
	"github.com/tsuru/tsuru/servicemanager"
	appTypes "github.com/tsuru/tsuru/types/app"
)

var _ appTypes.AppLogService = &provisionerWrapper{}

// provisionerWrapper is a layer designed to use provision native logging when is possible,
// otherwise will use backwards compatibility with own tsuru log api.
type provisionerWrapper struct {
	logService        appTypes.AppLogService
	provisionerGetter logsProvisionerGetter
}

func newProvisionerWrapper(logService appTypes.AppLogService) appTypes.AppLogService {
	return &provisionerWrapper{
		logService:        logService,
		provisionerGetter: defaultLogsProvisionerGetter,
	}
}

// Add is uncalled when the target pool uses own provisioner log stack
func (k *provisionerWrapper) Add(appName, message, source, unit string) error {
	return k.logService.Add(appName, message, source, unit)
}

// Enqueue is uncalled when the target pool uses own provisioner log stack
func (k *provisionerWrapper) Enqueue(entry *appTypes.Applog) error {
	return k.logService.Enqueue(entry)
}

func (k *provisionerWrapper) List(args appTypes.ListLogArgs) ([]appTypes.Applog, error) {
	a, err := servicemanager.App.GetByName(args.AppName)
	if err != nil {
		return nil, err
	}
	tsuruLogs, err := k.logService.List(args)
	if err != nil {
		return nil, err
	}
	logsProvisioner, err := k.provisionerGetter(a)
	if err == provision.ErrLogsUnavailable {
		return tsuruLogs, nil
	}
	if err != nil {
		return nil, err
	}

	logs, err := logsProvisioner.ListLogs(a, args)
	if err == provision.ErrLogsUnavailable {
		return tsuruLogs, nil
	}
	if err != nil {
		return nil, err
	}
	logs = append(logs, tsuruLogs...)
	sort.SliceStable(logs, func(i, j int) bool {
		return logs[i].Date.Before(logs[j].Date)
	})
	return logs, err
}

func (k *provisionerWrapper) Watch(args appTypes.ListLogArgs) (appTypes.LogWatcher, error) {
	a, err := servicemanager.App.GetByName(args.AppName)
	if err != nil {
		return nil, err
	}
	tsuruWatcher, err := k.logService.Watch(args)
	if err != nil {
		return nil, err
	}

	logsProvisioner, err := k.provisionerGetter(a)
	if err == provision.ErrLogsUnavailable {
		return tsuruWatcher, nil
	}
	if err != nil {
		return nil, err
	}

	provisionerWatcher, err := logsProvisioner.WatchLogs(a, args)
	if err == provision.ErrLogsUnavailable {
		return tsuruWatcher, nil
	}
	if err != nil {
		return nil, err
	}

	return newMultiWatcher(provisionerWatcher, tsuruWatcher), nil
}

type logsProvisionerGetter func(a appTypes.App) (provision.LogsProvisioner, error)

var defaultLogsProvisionerGetter = func(a appTypes.App) (provision.LogsProvisioner, error) {
	provisioner, err := pool.GetProvisionerForPool(a.GetPool())
	if err != nil {
		return nil, err
	}

	if logsProvisioner, ok := provisioner.(provision.LogsProvisioner); ok {
		return logsProvisioner, nil
	}

	return nil, provision.ErrLogsUnavailable
}

var _ appTypes.LogWatcher = &multiWatcher{}

type multiWatcher struct {
	subWatchers []appTypes.LogWatcher
	ch          chan appTypes.Applog
	close       chan struct{}
	closeCalled int32
	wg          sync.WaitGroup
}

func newMultiWatcher(subWatchers ...appTypes.LogWatcher) *multiWatcher {
	watcher := &multiWatcher{
		subWatchers: subWatchers,
		ch:          make(chan appTypes.Applog, 1000),
		close:       make(chan struct{}),
	}

	watcher.wg.Add(len(subWatchers))
	for _, subWatcher := range subWatchers {
		go watcher.startConsume(subWatcher)
	}

	return watcher
}

func (m *multiWatcher) startConsume(subWatcher appTypes.LogWatcher) {
	defer m.wg.Done()
	c := subWatcher.Chan()
	for {
		select {
		case log, open := <-c:

			if !open {
				return
			}

			m.ch <- log
		case <-m.close:
			return
		}
	}
}
func (m *multiWatcher) Chan() <-chan appTypes.Applog {
	return m.ch
}
func (m *multiWatcher) Close() {
	if atomic.AddInt32(&m.closeCalled, 1) != 1 {
		return
	}

	close(m.close)
	for _, subWatcher := range m.subWatchers {
		subWatcher.Close()
	}
	m.wg.Wait()
	close(m.ch)
}
<commit_msg>applog: use select to prevent channel blocking forever, thanks @cezarsa<commit_after>// Copyright 2020 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package applog

import (
	"sort"
	"sync"
	"sync/atomic"

	"github.com/tsuru/tsuru/provision"
	"github.com/tsuru/tsuru/provision/pool"
	"github.com/tsuru/tsuru/servicemanager"
	appTypes "github.com/tsuru/tsuru/types/app"
)

var _ appTypes.AppLogService = &provisionerWrapper{}

// provisionerWrapper is a layer designed to use provision native logging when is possible,
// otherwise will use backwards compatibility with own tsuru log api.
type provisionerWrapper struct {
	logService        appTypes.AppLogService
	provisionerGetter logsProvisionerGetter
}

func newProvisionerWrapper(logService appTypes.AppLogService) appTypes.AppLogService {
	return &provisionerWrapper{
		logService:        logService,
		provisionerGetter: defaultLogsProvisionerGetter,
	}
}

// Add is uncalled when the target pool uses own provisioner log stack
func (k *provisionerWrapper) Add(appName, message, source, unit string) error {
	return k.logService.Add(appName, message, source, unit)
}

// Enqueue is uncalled when the target pool uses own provisioner log stack
func (k *provisionerWrapper) Enqueue(entry *appTypes.Applog) error {
	return k.logService.Enqueue(entry)
}

func (k *provisionerWrapper) List(args appTypes.ListLogArgs) ([]appTypes.Applog, error) {
	a, err := servicemanager.App.GetByName(args.AppName)
	if err != nil {
		return nil, err
	}
	tsuruLogs, err := k.logService.List(args)
	if err != nil {
		return nil, err
	}
	logsProvisioner, err := k.provisionerGetter(a)
	if err == provision.ErrLogsUnavailable {
		return tsuruLogs, nil
	}
	if err != nil {
		return nil, err
	}

	logs, err := logsProvisioner.ListLogs(a, args)
	if err == provision.ErrLogsUnavailable {
		return tsuruLogs, nil
	}
	if err != nil {
		return nil, err
	}
	logs = append(logs, tsuruLogs...)
	sort.SliceStable(logs, func(i, j int) bool {
		return logs[i].Date.Before(logs[j].Date)
	})
	return logs, err
}

func (k *provisionerWrapper) Watch(args appTypes.ListLogArgs) (appTypes.LogWatcher, error) {
	a, err := servicemanager.App.GetByName(args.AppName)
	if err != nil {
		return nil, err
	}
	tsuruWatcher, err := k.logService.Watch(args)
	if err != nil {
		return nil, err
	}

	logsProvisioner, err := k.provisionerGetter(a)
	if err == provision.ErrLogsUnavailable {
		return tsuruWatcher, nil
	}
	if err != nil {
		return nil, err
	}

	provisionerWatcher, err := logsProvisioner.WatchLogs(a, args)
	if err == provision.ErrLogsUnavailable {
		return tsuruWatcher, nil
	}
	if err != nil {
		return nil, err
	}

	return newMultiWatcher(provisionerWatcher, tsuruWatcher), nil
}

type logsProvisionerGetter func(a appTypes.App) (provision.LogsProvisioner, error)

var defaultLogsProvisionerGetter = func(a appTypes.App) (provision.LogsProvisioner, error) {
	provisioner, err := pool.GetProvisionerForPool(a.GetPool())
	if err != nil {
		return nil, err
	}

	if logsProvisioner, ok := provisioner.(provision.LogsProvisioner); ok {
		return logsProvisioner, nil
	}

	return nil, provision.ErrLogsUnavailable
}

var _ appTypes.LogWatcher = &multiWatcher{}

type multiWatcher struct {
	subWatchers []appTypes.LogWatcher
	ch          chan appTypes.Applog
	close       chan struct{}
	closeCalled int32
	wg          sync.WaitGroup
}

func newMultiWatcher(subWatchers ...appTypes.LogWatcher) *multiWatcher {
	watcher := &multiWatcher{
		subWatchers: subWatchers,
		ch:          make(chan appTypes.Applog, 1000),
		close:       make(chan struct{}),
	}

	watcher.wg.Add(len(subWatchers))
	for _, subWatcher := range subWatchers {
		go watcher.startConsume(subWatcher)
	}

	return watcher
}

func (m *multiWatcher) startConsume(subWatcher appTypes.LogWatcher) {
	defer m.wg.Done()
	c := subWatcher.Chan()
	for {
		select {
		case log, open := <-c:

			if !open {
				return
			}

			select {
			case m.ch <- log:
			case <-m.close:
				return

			}
		case <-m.close:
			return
		}
	}
}
func (m *multiWatcher) Chan() <-chan appTypes.Applog {
	return m.ch
}
func (m *multiWatcher) Close() {
	if atomic.AddInt32(&m.closeCalled, 1) != 1 {
		return
	}

	close(m.close)
	for _, subWatcher := range m.subWatchers {
		subWatcher.Close()
	}
	m.wg.Wait()
	close(m.ch)
}
<|endoftext|>


<commit_before>package admin

import (
	"fmt"
	"path"
	"time"

	"github.com/qor/exchange"
	"github.com/qor/exchange/backends/csv"
	"github.com/qor/media_library"
	"github.com/qor/qor"
	"github.com/qor/qor-example/db"
	"github.com/qor/worker"
)

func getWorker() *worker.Worker {
	Worker := worker.New()

	type sendNewsletterArgument struct {
		Subject      string
		Content      string `sql:"size:65532"`
		SendPassword string
	}

	Worker.RegisterJob(worker.Job{
		Name: "send_newsletter",
		Handler: func(argument interface{}, qorJob worker.QorJobInterface) error {
			qorJob.AddLog("Started sending newsletters...")
			qorJob.AddLog(fmt.Sprintf("Argument: %+v", argument.(*sendNewsletterArgument)))
			for i := 1; i <= 100; i++ {
				time.Sleep(100 * time.Millisecond)
				qorJob.AddLog(fmt.Sprintf("Sending newsletter %v...", i))
				qorJob.SetProgress(uint(i))
			}
			qorJob.AddLog("Finished send newsletters")
			return nil
		},
		Resource: Admin.NewResource(&sendNewsletterArgument{}),
	})

	type importProductArgument struct {
		File media_library.FileSystem
	}

	Worker.RegisterJob(worker.Job{
		Name: "import_products",
		Handler: func(arg interface{}, qorJob worker.QorJobInterface) error {
			argument := arg.(*importProductArgument)

			context := &qor.Context{DB: db.DB}

			ProductExchange.Import(
				csv.New(path.Join("public", argument.File.URL())),
				context,
				func(progress exchange.Progress) error {
					qorJob.SetProgress(uint(float32(progress.Current) / float32(progress.Total) * 100))
					qorJob.AddLog(fmt.Sprintf("Importing product %d", progress.Current))
					return nil
				},
			)

			return nil
		},
		Resource: Admin.NewResource(&importProductArgument{}),
	})

	Worker.RegisterJob(worker.Job{
		Name: "export_products",
		Handler: func(arg interface{}, qorJob worker.QorJobInterface) error {
			qorJob.AddLog("Exporting products...")

			context := &qor.Context{DB: db.DB}
			fileName := fmt.Sprintf("/downloads/products.%v.csv", time.Now().UnixNano())
			ProductExchange.Export(csv.New(path.Join("public", fileName)), context)

			qorJob.SetProgressText(fmt.Sprintf("Download exported products", fileName))
			return nil
		},
	})
	return Worker
}
<commit_msg>Show error table from exchange<commit_after>package admin

import (
	"fmt"
	"path"
	"time"

	"github.com/qor/exchange"
	"github.com/qor/exchange/backends/csv"
	"github.com/qor/media_library"
	"github.com/qor/qor"
	"github.com/qor/qor-example/db"
	"github.com/qor/worker"
)

func getWorker() *worker.Worker {
	Worker := worker.New()

	type sendNewsletterArgument struct {
		Subject      string
		Content      string `sql:"size:65532"`
		SendPassword string
	}

	Worker.RegisterJob(worker.Job{
		Name: "send_newsletter",
		Handler: func(argument interface{}, qorJob worker.QorJobInterface) error {
			qorJob.AddLog("Started sending newsletters...")
			qorJob.AddLog(fmt.Sprintf("Argument: %+v", argument.(*sendNewsletterArgument)))
			for i := 1; i <= 100; i++ {
				time.Sleep(100 * time.Millisecond)
				qorJob.AddLog(fmt.Sprintf("Sending newsletter %v...", i))
				qorJob.SetProgress(uint(i))
			}
			qorJob.AddLog("Finished send newsletters")
			return nil
		},
		Resource: Admin.NewResource(&sendNewsletterArgument{}),
	})

	type importProductArgument struct {
		File media_library.FileSystem
	}

	Worker.RegisterJob(worker.Job{
		Name: "import_products",
		Handler: func(arg interface{}, qorJob worker.QorJobInterface) error {
			argument := arg.(*importProductArgument)

			context := &qor.Context{DB: db.DB}

			ProductExchange.Import(
				csv.New(path.Join("public", argument.File.URL())),
				context,
				func(progress exchange.Progress) error {
					var cells []worker.TableCell
					for _, cell := range progress.Cells {
						var tableCell = worker.TableCell{
							Value: fmt.Sprint(cell.Value),
						}
						if cell.Error != nil {
							tableCell.Error = cell.Error.Error()
							cells = append(cells, tableCell)
						}
					}
					qorJob.AddTableRow(cells...)

					qorJob.SetProgress(uint(float32(progress.Current) / float32(progress.Total) * 100))
					qorJob.AddLog(fmt.Sprintf("Importing product %d", progress.Current))
					return nil
				},
			)

			return nil
		},
		Resource: Admin.NewResource(&importProductArgument{}),
	})

	Worker.RegisterJob(worker.Job{
		Name: "export_products",
		Handler: func(arg interface{}, qorJob worker.QorJobInterface) error {
			qorJob.AddLog("Exporting products...")

			context := &qor.Context{DB: db.DB}
			fileName := fmt.Sprintf("/downloads/products.%v.csv", time.Now().UnixNano())
			ProductExchange.Export(csv.New(path.Join("public", fileName)), context)

			qorJob.SetProgressText(fmt.Sprintf("Download exported products", fileName))
			return nil
		},
	})
	return Worker
}
<|endoftext|>