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

package v3rpc

import (
	"crypto/sha256"
	"io"

	"github.com/coreos/etcd/auth"
	"github.com/coreos/etcd/etcdserver"
	pb "github.com/coreos/etcd/etcdserver/etcdserverpb"
	"github.com/coreos/etcd/mvcc"
	"github.com/coreos/etcd/mvcc/backend"
	"github.com/coreos/etcd/pkg/types"
	"github.com/coreos/etcd/version"
	"golang.org/x/net/context"
)

type KVGetter interface {
	KV() mvcc.ConsistentWatchableKV
}

type BackendGetter interface {
	Backend() backend.Backend
}

type Alarmer interface {
	Alarm(ctx context.Context, ar *pb.AlarmRequest) (*pb.AlarmResponse, error)
}

type RaftStatusGetter interface {
	Index() uint64
	Term() uint64
	Leader() types.ID
}

type AuthGetter interface {
	AuthInfoFromCtx(ctx context.Context) (*auth.AuthInfo, error)
	AuthStore() auth.AuthStore
}

type maintenanceServer struct {
	rg  RaftStatusGetter
	kg  KVGetter
	bg  BackendGetter
	a   Alarmer
	hdr header
}

func NewMaintenanceServer(s *etcdserver.EtcdServer) pb.MaintenanceServer {
	srv := &maintenanceServer{rg: s, kg: s, bg: s, a: s, hdr: newHeader(s)}
	return &authMaintenanceServer{srv, s}
}

func (ms *maintenanceServer) Defragment(ctx context.Context, sr *pb.DefragmentRequest) (*pb.DefragmentResponse, error) {
	plog.Noticef("starting to defragment the storage backend...")
	err := ms.bg.Backend().Defrag()
	if err != nil {
		plog.Errorf("failed to defragment the storage backend (%v)", err)
		return nil, err
	}
	plog.Noticef("finished defragmenting the storage backend")
	return &pb.DefragmentResponse{}, nil
}

func (ms *maintenanceServer) Snapshot(sr *pb.SnapshotRequest, srv pb.Maintenance_SnapshotServer) error {
	snap := ms.bg.Backend().Snapshot()
	pr, pw := io.Pipe()

	defer pr.Close()

	go func() {
		snap.WriteTo(pw)
		if err := snap.Close(); err != nil {
			plog.Errorf("error closing snapshot (%v)", err)
		}
		pw.Close()
	}()

	// send file data
	h := sha256.New()
	br := int64(0)
	buf := make([]byte, 32*1024)
	sz := snap.Size()
	for br < sz {
		n, err := io.ReadFull(pr, buf)
		if err != nil && err != io.EOF && err != io.ErrUnexpectedEOF {
			return togRPCError(err)
		}
		br += int64(n)
		resp := &pb.SnapshotResponse{
			RemainingBytes: uint64(sz - br),
			Blob:           buf[:n],
		}
		if err = srv.Send(resp); err != nil {
			return togRPCError(err)
		}
		h.Write(buf[:n])
	}

	// send sha
	sha := h.Sum(nil)
	hresp := &pb.SnapshotResponse{RemainingBytes: 0, Blob: sha}
	if err := srv.Send(hresp); err != nil {
		return togRPCError(err)
	}

	return nil
}

func (ms *maintenanceServer) Hash(ctx context.Context, r *pb.HashRequest) (*pb.HashResponse, error) {
	h, rev, err := ms.kg.KV().Hash()
	if err != nil {
		return nil, togRPCError(err)
	}
	resp := &pb.HashResponse{Header: &pb.ResponseHeader{Revision: rev}, Hash: h}
	ms.hdr.fill(resp.Header)
	return resp, nil
}

func (ms *maintenanceServer) Alarm(ctx context.Context, ar *pb.AlarmRequest) (*pb.AlarmResponse, error) {
	return ms.a.Alarm(ctx, ar)
}

func (ms *maintenanceServer) Status(ctx context.Context, ar *pb.StatusRequest) (*pb.StatusResponse, error) {
	resp := &pb.StatusResponse{
		Header:    &pb.ResponseHeader{Revision: ms.hdr.rev()},
		Version:   version.Version,
		DbSize:    ms.bg.Backend().Size(),
		Leader:    uint64(ms.rg.Leader()),
		RaftIndex: ms.rg.Index(),
		RaftTerm:  ms.rg.Term(),
	}
	ms.hdr.fill(resp.Header)
	return resp, nil
}

type authMaintenanceServer struct {
	*maintenanceServer
	ag AuthGetter
}

func (ams *authMaintenanceServer) isAuthenticated(ctx context.Context) error {
	authInfo, err := ams.ag.AuthInfoFromCtx(ctx)
	if err != nil {
		return err
	}

	return ams.ag.AuthStore().IsAdminPermitted(authInfo)
}

func (ams *authMaintenanceServer) Defragment(ctx context.Context, sr *pb.DefragmentRequest) (*pb.DefragmentResponse, error) {
	if err := ams.isAuthenticated(ctx); err != nil {
		return nil, err
	}

	return ams.maintenanceServer.Defragment(ctx, sr)
}

func (ams *authMaintenanceServer) Snapshot(sr *pb.SnapshotRequest, srv pb.Maintenance_SnapshotServer) error {
	if err := ams.isAuthenticated(srv.Context()); err != nil {
		return err
	}

	return ams.maintenanceServer.Snapshot(sr, srv)
}

func (ams *authMaintenanceServer) Hash(ctx context.Context, r *pb.HashRequest) (*pb.HashResponse, error) {
	if err := ams.isAuthenticated(ctx); err != nil {
		return nil, err
	}

	return ams.maintenanceServer.Hash(ctx, r)
}

func (ams *authMaintenanceServer) Status(ctx context.Context, ar *pb.StatusRequest) (*pb.StatusResponse, error) {
	return ams.maintenanceServer.Status(ctx, ar)
}
<commit_msg>etcdserver/api/v3rpc: add 'MoveLeader' to 'maintenanceServer'<commit_after>// Copyright 2016 The etcd Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package v3rpc

import (
	"crypto/sha256"
	"io"

	"github.com/coreos/etcd/auth"
	"github.com/coreos/etcd/etcdserver"
	"github.com/coreos/etcd/etcdserver/api/v3rpc/rpctypes"
	pb "github.com/coreos/etcd/etcdserver/etcdserverpb"
	"github.com/coreos/etcd/mvcc"
	"github.com/coreos/etcd/mvcc/backend"
	"github.com/coreos/etcd/pkg/types"
	"github.com/coreos/etcd/version"
	"golang.org/x/net/context"
)

type KVGetter interface {
	KV() mvcc.ConsistentWatchableKV
}

type BackendGetter interface {
	Backend() backend.Backend
}

type Alarmer interface {
	Alarm(ctx context.Context, ar *pb.AlarmRequest) (*pb.AlarmResponse, error)
}

type LeaderTransferrer interface {
	MoveLeader(ctx context.Context, lead, target uint64) error
}

type RaftStatusGetter interface {
	Index() uint64
	Term() uint64
	ID() types.ID
	Leader() types.ID
}

type AuthGetter interface {
	AuthInfoFromCtx(ctx context.Context) (*auth.AuthInfo, error)
	AuthStore() auth.AuthStore
}

type maintenanceServer struct {
	rg  RaftStatusGetter
	kg  KVGetter
	bg  BackendGetter
	a   Alarmer
	lt  LeaderTransferrer
	hdr header
}

func NewMaintenanceServer(s *etcdserver.EtcdServer) pb.MaintenanceServer {
	srv := &maintenanceServer{rg: s, kg: s, bg: s, a: s, lt: s, hdr: newHeader(s)}
	return &authMaintenanceServer{srv, s}
}

func (ms *maintenanceServer) Defragment(ctx context.Context, sr *pb.DefragmentRequest) (*pb.DefragmentResponse, error) {
	plog.Noticef("starting to defragment the storage backend...")
	err := ms.bg.Backend().Defrag()
	if err != nil {
		plog.Errorf("failed to defragment the storage backend (%v)", err)
		return nil, err
	}
	plog.Noticef("finished defragmenting the storage backend")
	return &pb.DefragmentResponse{}, nil
}

func (ms *maintenanceServer) Snapshot(sr *pb.SnapshotRequest, srv pb.Maintenance_SnapshotServer) error {
	snap := ms.bg.Backend().Snapshot()
	pr, pw := io.Pipe()

	defer pr.Close()

	go func() {
		snap.WriteTo(pw)
		if err := snap.Close(); err != nil {
			plog.Errorf("error closing snapshot (%v)", err)
		}
		pw.Close()
	}()

	// send file data
	h := sha256.New()
	br := int64(0)
	buf := make([]byte, 32*1024)
	sz := snap.Size()
	for br < sz {
		n, err := io.ReadFull(pr, buf)
		if err != nil && err != io.EOF && err != io.ErrUnexpectedEOF {
			return togRPCError(err)
		}
		br += int64(n)
		resp := &pb.SnapshotResponse{
			RemainingBytes: uint64(sz - br),
			Blob:           buf[:n],
		}
		if err = srv.Send(resp); err != nil {
			return togRPCError(err)
		}
		h.Write(buf[:n])
	}

	// send sha
	sha := h.Sum(nil)
	hresp := &pb.SnapshotResponse{RemainingBytes: 0, Blob: sha}
	if err := srv.Send(hresp); err != nil {
		return togRPCError(err)
	}

	return nil
}

func (ms *maintenanceServer) Hash(ctx context.Context, r *pb.HashRequest) (*pb.HashResponse, error) {
	h, rev, err := ms.kg.KV().Hash()
	if err != nil {
		return nil, togRPCError(err)
	}
	resp := &pb.HashResponse{Header: &pb.ResponseHeader{Revision: rev}, Hash: h}
	ms.hdr.fill(resp.Header)
	return resp, nil
}

func (ms *maintenanceServer) Alarm(ctx context.Context, ar *pb.AlarmRequest) (*pb.AlarmResponse, error) {
	return ms.a.Alarm(ctx, ar)
}

func (ms *maintenanceServer) Status(ctx context.Context, ar *pb.StatusRequest) (*pb.StatusResponse, error) {
	resp := &pb.StatusResponse{
		Header:    &pb.ResponseHeader{Revision: ms.hdr.rev()},
		Version:   version.Version,
		DbSize:    ms.bg.Backend().Size(),
		Leader:    uint64(ms.rg.Leader()),
		RaftIndex: ms.rg.Index(),
		RaftTerm:  ms.rg.Term(),
	}
	ms.hdr.fill(resp.Header)
	return resp, nil
}

func (ms *maintenanceServer) MoveLeader(ctx context.Context, tr *pb.MoveLeaderRequest) (*pb.MoveLeaderResponse, error) {
	if ms.rg.ID() != ms.rg.Leader() {
		return nil, rpctypes.ErrGRPCNotLeader
	}

	if err := ms.lt.MoveLeader(ctx, uint64(ms.rg.Leader()), tr.TargetID); err != nil {
		return nil, togRPCError(err)
	}
	return &pb.MoveLeaderResponse{}, nil
}

type authMaintenanceServer struct {
	*maintenanceServer
	ag AuthGetter
}

func (ams *authMaintenanceServer) isAuthenticated(ctx context.Context) error {
	authInfo, err := ams.ag.AuthInfoFromCtx(ctx)
	if err != nil {
		return err
	}

	return ams.ag.AuthStore().IsAdminPermitted(authInfo)
}

func (ams *authMaintenanceServer) Defragment(ctx context.Context, sr *pb.DefragmentRequest) (*pb.DefragmentResponse, error) {
	if err := ams.isAuthenticated(ctx); err != nil {
		return nil, err
	}

	return ams.maintenanceServer.Defragment(ctx, sr)
}

func (ams *authMaintenanceServer) Snapshot(sr *pb.SnapshotRequest, srv pb.Maintenance_SnapshotServer) error {
	if err := ams.isAuthenticated(srv.Context()); err != nil {
		return err
	}

	return ams.maintenanceServer.Snapshot(sr, srv)
}

func (ams *authMaintenanceServer) Hash(ctx context.Context, r *pb.HashRequest) (*pb.HashResponse, error) {
	if err := ams.isAuthenticated(ctx); err != nil {
		return nil, err
	}

	return ams.maintenanceServer.Hash(ctx, r)
}

func (ams *authMaintenanceServer) Status(ctx context.Context, ar *pb.StatusRequest) (*pb.StatusResponse, error) {
	return ams.maintenanceServer.Status(ctx, ar)
}

func (ams *authMaintenanceServer) MoveLeader(ctx context.Context, tr *pb.MoveLeaderRequest) (*pb.MoveLeaderResponse, error) {
	return ams.maintenanceServer.MoveLeader(ctx, tr)
}
<|endoftext|>


<commit_before>package bugsnag

import (
	"context"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
	"sync"
	"time"

	"github.com/bugsnag/bugsnag-go/errors"
	"github.com/bugsnag/bugsnag-go/sessions"

	// Fixes a bug with SHA-384 intermediate certs on some platforms.
	// - https://github.com/bugsnag/bugsnag-go/issues/9
	_ "crypto/sha512"
)

// The current version of bugsnag-go.
const VERSION = "1.3.1"
const configuredMultipleTimes = "WARNING: Bugsnag was configured twice. It is recommended to only call bugsnag.Configure once to ensure consistent session tracking behavior"

var once sync.Once
var middleware middlewareStack

// The configuration for the default bugsnag notifier.
var Config Configuration
var sessionTrackingConfig sessions.SessionTrackingConfiguration

// DefaultSessionPublishInterval defines how often sessions should be sent to
// Bugsnag.
// Deprecated: Exposed for developer sanity in testing. Modify at own risk.
var DefaultSessionPublishInterval = 60 * time.Second
var defaultNotifier = Notifier{&Config, nil}
var sessionTracker sessions.SessionTracker

// Configure Bugsnag. The only required setting is the APIKey, which can be
// obtained by clicking on "Settings" in your Bugsnag dashboard. This function
// is also responsible for installing the global panic handler, so it should be
// called as early as possible in your initialization process.
func Configure(config Configuration) {
	Config.update(&config)
	once.Do(Config.PanicHandler)
	startSessionTracking()
}

// StartSession creates a clone of the context.Context instance with Bugsnag
// session data attached.
func StartSession(ctx context.Context) context.Context {
	return sessionTracker.StartSession(ctx)
}

// Notify sends an error to Bugsnag along with the current stack trace. The
// rawData is used to send extra information along with the error. For example
// you can pass the current http.Request to Bugsnag to see information about it
// in the dashboard, or set the severity of the notification.
func Notify(rawData ...interface{}) error {
	return defaultNotifier.Notify(rawData...)
}

// AutoNotify logs a panic on a goroutine and then repanics.
// It should only be used in places that have existing panic handlers further
// up the stack. The rawData is used to send extra information along with any
// panics that are handled this way.
// Usage:
//  go func() {
//		defer bugsnag.AutoNotify()
//      // (possibly crashy code)
//  }()
// See also: bugsnag.Recover()
func AutoNotify(rawData ...interface{}) {
	if err := recover(); err != nil {
		severity := defaultNotifier.getDefaultSeverity(rawData, SeverityError)
		state := HandledState{SeverityReasonHandledPanic, severity, true, ""}
		rawData = append([]interface{}{state}, rawData...)
		defaultNotifier.NotifySync(append(rawData, errors.New(err, 2), true)...)
		panic(err)
	}
}

// Recover logs a panic on a goroutine and then recovers.
// The rawData is used to send extra information along with
// any panics that are handled this way
// Usage: defer bugsnag.Recover()
func Recover(rawData ...interface{}) {
	if err := recover(); err != nil {
		severity := defaultNotifier.getDefaultSeverity(rawData, SeverityWarning)
		state := HandledState{SeverityReasonHandledPanic, severity, false, ""}
		rawData = append([]interface{}{state}, rawData...)
		defaultNotifier.Notify(append(rawData, errors.New(err, 2))...)
	}
}

// OnBeforeNotify adds a callback to be run before a notification is sent to
// Bugsnag.  It can be used to modify the event or its MetaData. Changes made
// to the configuration are local to notifying about this event. To prevent the
// event from being sent to Bugsnag return an error, this error will be
// returned from bugsnag.Notify() and the event will not be sent.
func OnBeforeNotify(callback func(event *Event, config *Configuration) error) {
	middleware.OnBeforeNotify(callback)
}

// Handler creates an http Handler that notifies Bugsnag any panics that
// happen. It then repanics so that the default http Server panic handler can
// handle the panic too. The rawData is used to send extra information along
// with any panics that are handled this way.
func Handler(h http.Handler, rawData ...interface{}) http.Handler {
	notifier := New(rawData...)
	if h == nil {
		h = http.DefaultServeMux
	}

	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		defer notifier.AutoNotify(StartSession(r.Context()), r)
		h.ServeHTTP(w, r)
	})
}

// HandlerFunc creates an http HandlerFunc that notifies Bugsnag about any
// panics that happen. It then repanics so that the default http Server panic
// handler can handle the panic too. The rawData is used to send extra
// information along with any panics that are handled this way. If you have
// already wrapped your http server using bugsnag.Handler() you don't also need
// to wrap each HandlerFunc.
func HandlerFunc(h http.HandlerFunc, rawData ...interface{}) http.HandlerFunc {
	notifier := New(rawData...)

	return func(w http.ResponseWriter, r *http.Request) {
		defer notifier.AutoNotify(r)
		h(w, r)
	}
}

func init() {
	// Set up builtin middlewarez
	OnBeforeNotify(httpRequestMiddleware)

	// Default configuration
	sourceRoot := ""
	if gopath := os.Getenv("GOPATH"); len(gopath) > 0 {
		sourceRoot = filepath.Join(gopath, "src") + "/"
	} else {
		sourceRoot = filepath.Join(runtime.GOROOT(), "src") + "/"
	}
	Config.update(&Configuration{
		APIKey: "",
		Endpoints: Endpoints{
			Notify:   "https://notify.bugsnag.com",
			Sessions: "https://sessions.bugsnag.com",
		},
		Hostname:            "",
		AppType:             "",
		AppVersion:          "",
		AutoCaptureSessions: true,
		ReleaseStage:        "",
		ParamsFilters:       []string{"password", "secret"},
		SourceRoot:          sourceRoot,
		// * for app-engine
		ProjectPackages:     []string{"main*"},
		NotifyReleaseStages: nil,
		Logger:              log.New(os.Stdout, log.Prefix(), log.Flags()),
		PanicHandler:        defaultPanicHandler,
		Transport:           http.DefaultTransport,
	})

	hostname, err := os.Hostname()
	if err == nil {
		Config.Hostname = hostname
	}
}

func startSessionTracking() {
	sessionTrackingConfig.Update(&sessions.SessionTrackingConfiguration{
		APIKey:          Config.APIKey,
		Endpoint:        Config.Endpoints.Sessions,
		Version:         VERSION,
		PublishInterval: DefaultSessionPublishInterval,
		Transport:       Config.Transport,
		ReleaseStage:    Config.ReleaseStage,
		Hostname:        Config.Hostname,
		AppType:         Config.AppType,
		AppVersion:      Config.AppVersion,
		Logger:          Config.Logger,
	})
	if sessionTracker != nil {
		Config.logf(configuredMultipleTimes)
	} else {
		sessionTracker = sessions.NewSessionTracker(&sessionTrackingConfig)
	}
}
<commit_msg>[chore] Improve GoDocs and mention use of context.Context<commit_after>package bugsnag

import (
	"context"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
	"sync"
	"time"

	"github.com/bugsnag/bugsnag-go/errors"
	"github.com/bugsnag/bugsnag-go/sessions"

	// Fixes a bug with SHA-384 intermediate certs on some platforms.
	// - https://github.com/bugsnag/bugsnag-go/issues/9
	_ "crypto/sha512"
)

// VERSION defines the version of this Bugsnag notifier
const VERSION = "1.3.1"
const configuredMultipleTimes = "WARNING: Bugsnag was configured twice. It is recommended to only call bugsnag.Configure once to ensure consistent session tracking behavior"

var once sync.Once
var middleware middlewareStack

// Config is the configuration for the default bugsnag notifier.
var Config Configuration
var sessionTrackingConfig sessions.SessionTrackingConfiguration

// DefaultSessionPublishInterval defines how often sessions should be sent to
// Bugsnag.
// Deprecated: Exposed for developer sanity in testing. Modify at own risk.
var DefaultSessionPublishInterval = 60 * time.Second
var defaultNotifier = Notifier{&Config, nil}
var sessionTracker sessions.SessionTracker

// Configure Bugsnag. The only required setting is the APIKey, which can be
// obtained by clicking on "Settings" in your Bugsnag dashboard. This function
// is also responsible for installing the global panic handler, so it should be
// called as early as possible in your initialization process.
func Configure(config Configuration) {
	Config.update(&config)
	once.Do(Config.PanicHandler)
	startSessionTracking()
}

// StartSession creates new context from the context.Context instance with
// Bugsnag session data attached.
func StartSession(ctx context.Context) context.Context {
	return sessionTracker.StartSession(ctx)
}

// Notify sends an error.Error to Bugsnag along with the current stack trace.
// Although it's not strictly enforced, it's highly recommended to pass a
// context.Context object that has at one-point been returned from
// bugsnag.StartSession. Doing so ensures your stability score remains accurate,
// and future versions of Bugsnag may extract more useful information from this
// context.
// The remaining rawData is used to send extra information along with the
// error. For example you can pass the current http.Request to Bugsnag to see
// information about it in the dashboard, or set the severity of the
// notification.
func Notify(rawData ...interface{}) error {
	return defaultNotifier.Notify(rawData...)
}

// AutoNotify logs a panic on a goroutine and then repanics.
// It should only be used in places that have existing panic handlers further
// up the stack.
// Although it's not strictly enforced, it's highly recommended to pass a
// context.Context object that has at one-point been returned from
// bugsnag.StartSession. Doing so ensures your stability score remains accurate,
// and future versions of Bugsnag may extract more useful information from this
// context.
// The rawData is used to send extra information along with any
// panics that are handled this way.
// Usage:
//  go func() {
//      ctx := bugsnag.StartSession(context.Background())
//		defer bugsnag.AutoNotify(ctx)
//      // (possibly crashy code)
//  }()
// See also: bugsnag.Recover()
func AutoNotify(rawData ...interface{}) {
	if err := recover(); err != nil {
		severity := defaultNotifier.getDefaultSeverity(rawData, SeverityError)
		state := HandledState{SeverityReasonHandledPanic, severity, true, ""}
		rawData = append([]interface{}{state}, rawData...)
		defaultNotifier.NotifySync(append(rawData, errors.New(err, 2), true)...)
		panic(err)
	}
}

// Recover logs a panic on a goroutine and then recovers.
// Although it's not strictly enforced, it's highly recommended to pass a
// context.Context object that has at one-point been returned from
// bugsnag.StartSession. Doing so ensures your stability score remains accurate,
// and future versions of Bugsnag may extract more useful information from this
// context.
// The rawData is used to send extra information along with
// any panics that are handled this way
// Usage:
//  go func() {
//      ctx := bugsnag.StartSession(context.Background())
//		defer bugsnag.AutoNotify(ctx)
//      // (possibly crashy code)
//  }()
// See also: bugsnag.AutoNotify()
func Recover(rawData ...interface{}) {
	if err := recover(); err != nil {
		severity := defaultNotifier.getDefaultSeverity(rawData, SeverityWarning)
		state := HandledState{SeverityReasonHandledPanic, severity, false, ""}
		rawData = append([]interface{}{state}, rawData...)
		defaultNotifier.Notify(append(rawData, errors.New(err, 2))...)
	}
}

// OnBeforeNotify adds a callback to be run before a notification is sent to
// Bugsnag.  It can be used to modify the event or its MetaData. Changes made
// to the configuration are local to notifying about this event. To prevent the
// event from being sent to Bugsnag return an error, this error will be
// returned from bugsnag.Notify() and the event will not be sent.
func OnBeforeNotify(callback func(event *Event, config *Configuration) error) {
	middleware.OnBeforeNotify(callback)
}

// Handler creates an http Handler that notifies Bugsnag any panics that
// happen. It then repanics so that the default http Server panic handler can
// handle the panic too. The rawData is used to send extra information along
// with any panics that are handled this way.
func Handler(h http.Handler, rawData ...interface{}) http.Handler {
	notifier := New(rawData...)
	if h == nil {
		h = http.DefaultServeMux
	}

	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		defer notifier.AutoNotify(StartSession(r.Context()), r)
		h.ServeHTTP(w, r)
	})
}

// HandlerFunc creates an http HandlerFunc that notifies Bugsnag about any
// panics that happen. It then repanics so that the default http Server panic
// handler can handle the panic too. The rawData is used to send extra
// information along with any panics that are handled this way. If you have
// already wrapped your http server using bugsnag.Handler() you don't also need
// to wrap each HandlerFunc.
func HandlerFunc(h http.HandlerFunc, rawData ...interface{}) http.HandlerFunc {
	notifier := New(rawData...)

	return func(w http.ResponseWriter, r *http.Request) {
		defer notifier.AutoNotify(r)
		h(w, r)
	}
}

func init() {
	// Set up builtin middlewarez
	OnBeforeNotify(httpRequestMiddleware)

	// Default configuration
	sourceRoot := ""
	if gopath := os.Getenv("GOPATH"); len(gopath) > 0 {
		sourceRoot = filepath.Join(gopath, "src") + "/"
	} else {
		sourceRoot = filepath.Join(runtime.GOROOT(), "src") + "/"
	}
	Config.update(&Configuration{
		APIKey: "",
		Endpoints: Endpoints{
			Notify:   "https://notify.bugsnag.com",
			Sessions: "https://sessions.bugsnag.com",
		},
		Hostname:            "",
		AppType:             "",
		AppVersion:          "",
		AutoCaptureSessions: true,
		ReleaseStage:        "",
		ParamsFilters:       []string{"password", "secret"},
		SourceRoot:          sourceRoot,
		// * for app-engine
		ProjectPackages:     []string{"main*"},
		NotifyReleaseStages: nil,
		Logger:              log.New(os.Stdout, log.Prefix(), log.Flags()),
		PanicHandler:        defaultPanicHandler,
		Transport:           http.DefaultTransport,
	})

	hostname, err := os.Hostname()
	if err == nil {
		Config.Hostname = hostname
	}
}

func startSessionTracking() {
	sessionTrackingConfig.Update(&sessions.SessionTrackingConfiguration{
		APIKey:          Config.APIKey,
		Endpoint:        Config.Endpoints.Sessions,
		Version:         VERSION,
		PublishInterval: DefaultSessionPublishInterval,
		Transport:       Config.Transport,
		ReleaseStage:    Config.ReleaseStage,
		Hostname:        Config.Hostname,
		AppType:         Config.AppType,
		AppVersion:      Config.AppVersion,
		Logger:          Config.Logger,
	})
	if sessionTracker != nil {
		Config.logf(configuredMultipleTimes)
	} else {
		sessionTracker = sessions.NewSessionTracker(&sessionTrackingConfig)
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2014 The SurgeMQ Authors. All rights reserved.
//
// 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 service

import (
	"bufio"
	"fmt"
	"io"
	"sync/atomic"
	"runtime"
//"github.com/surgemq/message"
	"encoding/binary"
	"strconv"
)

var (
	bufcnt int64
	DefaultBufferSize int64

	DeviceInBufferSize int64
	DeviceOutBufferSize int64

	MasterInBufferSize int64
	MasterOutBufferSize int64
)

const (
	smallReadBlockSize = 512
	defaultReadBlockSize = 8192
	defaultWriteBlockSize = 8192
)

/**
2016.03.03 修改
bingbuffer结构体
 */
type buffer struct {
	readIndex  int64     //读序号
	writeIndex int64     //写序号
	ringBuffer []*[]byte //环形buffer指针数组
	bufferSize int64     //初始化环形buffer指针数组大小
	mask       int64     //掩码：bufferSize-1
	done       int64     //是否完成
}

/**
2016.03.03 添加
初始化ringbuffer
参数bufferSize：初始化环形buffer指针数组大小
 */
func newBuffer(size int64) (*buffer, error) {
	if size < 0 {
		return nil, bufio.ErrNegativeCount
	}
	if size == 0 {
		size = DefaultBufferSize
	}
	if !powerOfTwo64(size) {
		fmt.Printf("Size must be power of two. Try %d.", roundUpPowerOfTwo64(size))
		return nil, fmt.Errorf("Size must be power of two. Try %d.", roundUpPowerOfTwo64(size))
	}

	return &buffer{
		readIndex: int64(0), //读序号
		writeIndex: int64(0), //写序号
		ringBuffer: make([]*[]byte, size), //环形buffer指针数组
		bufferSize: size, //初始化环形buffer指针数组大小
		mask:size - 1,
	}, nil
}

/**
2016.03.03 添加
获取当前读序号
 */
func (this *buffer)GetCurrentReadIndex() (int64) {
	return atomic.LoadInt64(&this.readIndex)
}
/**
2016.03.03 添加
获取当前写序号
 */
func (this *buffer)GetCurrentWriteIndex() (int64) {
	return atomic.LoadInt64(&this.writeIndex)
}

/**
2016.03.03 添加
读取ringbuffer指定的buffer指针，返回该指针并清空ringbuffer该位置存在的指针内容，以及将读序号加1
 */
func (this *buffer)ReadBuffer() (p *[]byte, ok bool) {
	ok = true
	p = nil

	readIndex := atomic.LoadInt64(&this.readIndex)
	writeIndex := atomic.LoadInt64(&this.writeIndex)
	switch  {
	case readIndex >= writeIndex:
		ok = false
	case writeIndex - readIndex > this.bufferSize:
		ok = false
	default:
		//index := buffer.readIndex % buffer.bufferSize
		index := readIndex & this.mask
		p = this.ringBuffer[index][0:]
		this.ringBuffer[index] = nil
		atomic.AddInt64(&this.readIndex, 1)

		if p == nil {
			ok = false
		}
	}
	return p, ok
}

/**
2016.03.03 添加
写入ringbuffer指针，以及将写序号加1
 */
func (this *buffer)WriteBuffer(in *[]byte) (ok bool) {
	ok = true

	readIndex := atomic.LoadInt64(&this.readIndex)
	writeIndex := atomic.LoadInt64(&this.writeIndex)
	switch  {
	case writeIndex - readIndex < 0:
		ok = false
	default:
		//index := buffer.writeIndex % buffer.bufferSize
		index := writeIndex & this.mask
		if this.ringBuffer[index] == nil {
			this.ringBuffer[index] = in
			atomic.AddInt64(&this.writeIndex, 1)
		}else {
			ok = false
		}
	}
	return ok
}

/**
2016.03.03 修改
关闭缓存
 */
func (this *buffer) Close() error {
	atomic.StoreInt64(&this.done, 1)
	return nil
}
/*

/**
2016.03.03 修改
向ringbuffer中写数据（从connection的中向ringbuffer中写）--生产者
*/
func (this *buffer) ReadFrom(r io.Reader) (int64, error) {
	defer this.Close()
	for {
		total := int64(0)
		if this.isDone() {
			return total, io.EOF
		}
		b := make([]byte, 5)
		n, err := r.Read(b[0:1])

		if n > 0 {
			total += int64(n)
			if err != nil {
				return total, err
			}
		}

		/**************************/
		cnt := 1

		// Let's read enough bytes to get the message header (msg type, remaining length)
		for {
			// If we have read 5 bytes and still not done, then there's a problem.
			if cnt > 4 {
				return 0, fmt.Errorf("sendrecv/peekMessageSize: 4th byte of remaining length has continuation bit set")
			}

			// Peek cnt bytes from the input buffer.
			_, err := r.Read(b[cnt:cnt + 1])
			if err != nil {
				return 0, err
			}

			// If we got enough bytes, then check the last byte to see if the continuation
			// bit is set. If so, increment cnt and continue peeking
			if b[cnt] >= 0x80 {
				cnt++
			} else {
				break
			}
		}

		// Get the remaining length of the message
		remlen, _ := binary.Uvarint(b[1:])

		// Total message length is remlen + 1 (msg type) + m (remlen bytes)
		len := int64(len(b))
		remlen_ := int64(remlen)
		total = remlen_ + int64(len)

		//mtype := message.MessageType(b[0] >> 4)
		/****************/
		//var msg message.Message
		//
		//msg, err = mtype.New()
		//if err != nil {
		//	return 0, err
		//}
		b_ := make([]byte, remlen_)
		_, err = r.Read(b_[0:])
		if err != nil {
			return 0, err
		}

		b = append(b, b_...)
		//n, err = msg.Decode(b)
		//if err != nil {
		//	return 0, err
		//}

		/*************************/

		for !this.WriteBuffer(&b) {
			runtime.Gosched()
		}

		return total, nil
	}
}

/**
2016.03.03 修改
 */
func (this *buffer) WriteTo(w io.Writer) (int64, error) {
	defer this.Close()
	total := int64(0)
	for {
		if this.isDone() {
			return total, io.EOF
		}
		p, ok := this.ReadBuffer()
		if !ok {
			runtime.Gosched()
		}
		Log.Debugc(func() string {
			return fmt.Sprintf("p=" + strconv.FormatBool(p == nil))
		})
		Log.Debugc(func() string {
			return fmt.Sprintf("WriteTo函数》》读取*p：" + string(*p))
		})

		Log.Debugc(func() string {
			return fmt.Sprintf(" WriteTo(w io.Writer)(7)")
		})
		//
		//Log.Errorc(func() string {
		//	return fmt.Sprintf("msg::" + msg.Name())
		//})
		//
		//p := make([]byte, msg.Len())
		//_, err := msg.Encode(p)
		//if err != nil {
		//	Log.Errorc(func() string {
		//		return fmt.Sprintf("msg.Encode(p)")
		//	})
		//	return total, io.EOF
		//}
		// There's some data, let's process it first
		if len(*p) > 0 {
			n, err := w.Write(*p)
			total += int64(n)
			Log.Debugc(func() string {
				return fmt.Sprintf("Wrote %d bytes, totaling %d bytes", n, total)
			})

			if err != nil {
				Log.Errorc(func() string {
					return fmt.Sprintf("w.Write(p) error")
				})
				return total, err
			}
		}

		return total, nil
	}
}

/**
2016.03.03 修改
*/
func (this *buffer) isDone() bool {
	if atomic.LoadInt64(&this.done) == 1 {
		return true
	}

	return false
}

func powerOfTwo64(n int64) bool {
	return n != 0 && (n & (n - 1)) == 0
}

func roundUpPowerOfTwo64(n int64) int64 {
	n--
	n |= n >> 1
	n |= n >> 2
	n |= n >> 4
	n |= n >> 8
	n |= n >> 16
	n |= n >> 32
	n++

	return n
}
<commit_msg>添加测试日志<commit_after>// Copyright (c) 2014 The SurgeMQ Authors. All rights reserved.
//
// 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 service

import (
	"bufio"
	"fmt"
	"io"
	"sync/atomic"
	"runtime"
//"github.com/surgemq/message"
	"encoding/binary"
	"strconv"
)

var (
	bufcnt int64
	DefaultBufferSize int64

	DeviceInBufferSize int64
	DeviceOutBufferSize int64

	MasterInBufferSize int64
	MasterOutBufferSize int64
)

const (
	smallReadBlockSize = 512
	defaultReadBlockSize = 8192
	defaultWriteBlockSize = 8192
)

/**
2016.03.03 修改
bingbuffer结构体
 */
type buffer struct {
	readIndex  int64        //读序号
	writeIndex int64        //写序号
	ringBuffer *[]ByteArray //环形buffer指针数组
	bufferSize int64        //初始化环形buffer指针数组大小
	mask       int64        //掩码：bufferSize-1
	done       int64        //是否完成
}

type ByteArray struct {
	bArray []byte
}

func (this *ByteArray)GetArray() ([]byte) {
	return this.bArray
}

/**
2016.03.03 添加
初始化ringbuffer
参数bufferSize：初始化环形buffer指针数组大小
 */
func newBuffer(size int64) (*buffer, error) {
	if size < 0 {
		return nil, bufio.ErrNegativeCount
	}
	if size == 0 {
		size = DefaultBufferSize
	}
	if !powerOfTwo64(size) {
		fmt.Printf("Size must be power of two. Try %d.", roundUpPowerOfTwo64(size))
		return nil, fmt.Errorf("Size must be power of two. Try %d.", roundUpPowerOfTwo64(size))
	}

	return &buffer{
		readIndex: int64(0), //读序号
		writeIndex: int64(0), //写序号
		ringBuffer: make([]ByteArray, size), //环形buffer指针数组
		bufferSize: size, //初始化环形buffer指针数组大小
		mask:size - 1,
	}, nil
}

/**
2016.03.03 添加
获取当前读序号
 */
func (this *buffer)GetCurrentReadIndex() (int64) {
	return atomic.LoadInt64(&this.readIndex)
}
/**
2016.03.03 添加
获取当前写序号
 */
func (this *buffer)GetCurrentWriteIndex() (int64) {
	return atomic.LoadInt64(&this.writeIndex)
}

/**
2016.03.03 添加
读取ringbuffer指定的buffer指针，返回该指针并清空ringbuffer该位置存在的指针内容，以及将读序号加1
 */
func (this *buffer)ReadBuffer() (p *[]byte, ok bool) {
	ok = true
	p = nil

	readIndex := atomic.LoadInt64(&this.readIndex)
	writeIndex := atomic.LoadInt64(&this.writeIndex)
	switch  {
	case readIndex >= writeIndex:
		ok = false
	case writeIndex - readIndex > this.bufferSize:
		ok = false
	default:
		//index := buffer.readIndex % buffer.bufferSize
		index := readIndex & this.mask

		p_ := ByteArray{}(this.ringBuffer[index])
		this.ringBuffer[index] = nil
		atomic.AddInt64(&this.readIndex, 1)
		p = p_.GetArray()[0:]
		if p == nil {
			ok = false
		}
	}
	return p, ok
}

/**
2016.03.03 添加
写入ringbuffer指针，以及将写序号加1
 */
func (this *buffer)WriteBuffer(in []byte) (ok bool) {
	ok = true

	readIndex := atomic.LoadInt64(&this.readIndex)
	writeIndex := atomic.LoadInt64(&this.writeIndex)
	switch  {
	case writeIndex - readIndex < 0:
		ok = false
	default:
		//index := buffer.writeIndex % buffer.bufferSize
		index := writeIndex & this.mask
		if this.ringBuffer[index] == nil {
			this.ringBuffer[index] = &ByteArray{bArray:in}
			atomic.AddInt64(&this.writeIndex, 1)
		}else {
			ok = false
		}
	}
	return ok
}

/**
2016.03.03 修改
关闭缓存
 */
func (this *buffer) Close() error {
	atomic.StoreInt64(&this.done, 1)
	return nil
}
/*

/**
2016.03.03 修改
向ringbuffer中写数据（从connection的中向ringbuffer中写）--生产者
*/
func (this *buffer) ReadFrom(r io.Reader) (int64, error) {
	defer this.Close()
	for {
		total := int64(0)
		if this.isDone() {
			return total, io.EOF
		}
		b := make([]byte, 5)
		n, err := r.Read(b[0:1])

		if n > 0 {
			total += int64(n)
			if err != nil {
				return total, err
			}
		}

		/**************************/
		cnt := 1

		// Let's read enough bytes to get the message header (msg type, remaining length)
		for {
			// If we have read 5 bytes and still not done, then there's a problem.
			if cnt > 4 {
				return 0, fmt.Errorf("sendrecv/peekMessageSize: 4th byte of remaining length has continuation bit set")
			}

			// Peek cnt bytes from the input buffer.
			_, err := r.Read(b[cnt:cnt + 1])
			if err != nil {
				return 0, err
			}

			// If we got enough bytes, then check the last byte to see if the continuation
			// bit is set. If so, increment cnt and continue peeking
			if b[cnt] >= 0x80 {
				cnt++
			} else {
				break
			}
		}

		// Get the remaining length of the message
		remlen, _ := binary.Uvarint(b[1:])

		// Total message length is remlen + 1 (msg type) + m (remlen bytes)
		len := int64(len(b))
		remlen_ := int64(remlen)
		total = remlen_ + int64(len)

		//mtype := message.MessageType(b[0] >> 4)
		/****************/
		//var msg message.Message
		//
		//msg, err = mtype.New()
		//if err != nil {
		//	return 0, err
		//}
		b_ := make([]byte, remlen_)
		_, err = r.Read(b_[0:])
		if err != nil {
			return 0, err
		}

		b = append(b, b_...)
		//n, err = msg.Decode(b)
		//if err != nil {
		//	return 0, err
		//}

		/*************************/

		for !this.WriteBuffer(&b) {
			runtime.Gosched()
		}

		return total, nil
	}
}

/**
2016.03.03 修改
 */
func (this *buffer) WriteTo(w io.Writer) (int64, error) {
	defer this.Close()
	total := int64(0)
	for {
		if this.isDone() {
			return total, io.EOF
		}
		p, ok := this.ReadBuffer()
		if !ok {
			runtime.Gosched()
		}
		Log.Debugc(func() string {
			return fmt.Sprintf("p=" + strconv.FormatBool(p == nil))
		})
		Log.Debugc(func() string {
			return fmt.Sprintf("WriteTo函数》》读取*p：" + string(*p))
		})

		Log.Debugc(func() string {
			return fmt.Sprintf(" WriteTo(w io.Writer)(7)")
		})
		//
		//Log.Errorc(func() string {
		//	return fmt.Sprintf("msg::" + msg.Name())
		//})
		//
		//p := make([]byte, msg.Len())
		//_, err := msg.Encode(p)
		//if err != nil {
		//	Log.Errorc(func() string {
		//		return fmt.Sprintf("msg.Encode(p)")
		//	})
		//	return total, io.EOF
		//}
		// There's some data, let's process it first
		if len(*p) > 0 {
			n, err := w.Write(*p)
			total += int64(n)
			Log.Debugc(func() string {
				return fmt.Sprintf("Wrote %d bytes, totaling %d bytes", n, total)
			})

			if err != nil {
				Log.Errorc(func() string {
					return fmt.Sprintf("w.Write(p) error")
				})
				return total, err
			}
		}

		return total, nil
	}
}

/**
2016.03.03 修改
*/
func (this *buffer) isDone() bool {
	if atomic.LoadInt64(&this.done) == 1 {
		return true
	}

	return false
}

func powerOfTwo64(n int64) bool {
	return n != 0 && (n & (n - 1)) == 0
}

func roundUpPowerOfTwo64(n int64) int64 {
	n--
	n |= n >> 1
	n |= n >> 2
	n |= n >> 4
	n |= n >> 8
	n |= n >> 16
	n |= n >> 32
	n++

	return n
}
<|endoftext|>


<commit_before>package cli

/*
	For debugging exec pipes, try this:
	go run chirp.go -recover=0 -c='puts [exec ls -l | sed {s/[0-9]/#/g} | tr {a-z} {A-Z} ]' 2>/dev/null | od -c
*/

import (
	"bufio"
	"flag"
	. "fmt"
	"github.com/yak-labs/chirp-lang"
	"io/ioutil"
	"os"
	"runtime/pprof"
	"strings"
)

var dFlag = flag.String("d", "", "Debugging flags, each a single letter.")
var cFlag = flag.String("c", "", "Immediate command to execute.")
var recoverFlag = flag.Bool("recover", true, "Set to false to disable recover in the REPL.")
var testFlag = flag.Bool("test", false, "Print test summary at end.")

var scriptName string

func saveArgvStarting(fr *chirp.Frame, i int) {
	argv := []chirp.T{}
	for _, a := range flag.Args() {
		argv = append(argv, chirp.MkString(a))
	}
	fr.SetVar("argv", chirp.MkList(argv)) // Deprecated: argv
	fr.SetVar("Argv", chirp.MkList(argv)) // New: Argv
}

func setEnvironInChirp(fr *chirp.Frame, varName string) {
	h := make(chirp.Hash)
	for _, s := range os.Environ() {
		kv := strings.SplitN(s, "=", 2)
		if len(kv) == 2 {
			h[kv[0]] = chirp.MkString(kv[1])
		}
	}
	fr.SetVar(varName, chirp.MkHash(h))
}

func Main() {
	flag.Parse()
	fr := chirp.NewInterpreter()
	setEnvironInChirp(fr, "Env")

	for _, ch := range *dFlag {
		chirp.Debug[ch] = true
	}

	if cFlag != nil && *cFlag != "" {
		saveArgvStarting(fr, 1)
		fr.Eval(chirp.MkString(*cFlag))
		goto End
	}

	if len(flag.Args()) > 0 {
		// Script mode.
		scriptName = flag.Arg(0)
		contents, err := ioutil.ReadFile(scriptName)
		if err != nil {
			Fprintf(os.Stderr, "Cannot read file %s: %v", scriptName, err)
			os.Exit(2)
			return
		}
		saveArgvStarting(fr, 1)

		fr.Eval(chirp.MkString(string(contents)))
		goto End
	}

	{
		// Interactive mode.
		bio := bufio.NewReader(os.Stdin)
		for {
			Fprint(os.Stderr, "chirp% ") // Prompt to stderr.
			line, isPrefix, err := bio.ReadLine()
			if err != nil {
				if err.Error() == "EOF" { // TODO: better way?
					goto End
				}
				Fprintf(os.Stderr, "ERROR in ReadLine: %s\n", err.Error())
				goto End
			}
			fullLine := line
			for isPrefix {
				line, isPrefix, err = bio.ReadLine()
				if err != nil {
					Fprintf(os.Stderr, "ERROR in ReadLine: %s\n", err.Error())
					goto End
				}
				fullLine = append(fullLine, line...)
			}
			result := EvalStringOrPrintError(fr, string(fullLine))
			if result != "" { // Traditionally, if result is empty, tclsh doesn't print.
				Println(result)
			}
		}
	}

End:
	logAllCounters()
	if chirp.Debug['h'] {
		pprof.Lookup("heap").WriteTo(os.Stderr, 0)
	}
}

func logAllCounters() {
	if chirp.Debug['c'] {
		chirp.LogAllCounters()
	}

	// Print summary for tests.
	if *testFlag {
		chirp.MustMutex.Lock()
		if chirp.MustFails > 0 {
			Fprintf(os.Stderr, "TEST FAILS: %q succeeds=%d fails=%d\n", scriptName, chirp.MustSucceeds, chirp.MustFails)
			os.Exit(1)
		}
		Fprintf(os.Stderr, "Test Done: %q succeeds=%d\n", scriptName, chirp.MustSucceeds)
		chirp.MustMutex.Unlock()
	}
}

func EvalStringOrPrintError(fr *chirp.Frame, cmd string) (out string) {
	if *recoverFlag {
		defer func() {
			if r := recover(); r != nil {
				Fprintln(os.Stderr, "ERROR: ", r) // Error to stderr.
				out = ""
				return
			}
		}()
	}

	return fr.Eval(chirp.MkString(cmd)).String()
}
<commit_msg>Readline for chirp CLI.<commit_after>package cli

/*
	For debugging exec pipes, try this:
	go run chirp.go -recover=0 -c='puts [exec ls -l | sed {s/[0-9]/#/g} | tr {a-z} {A-Z} ]' 2>/dev/null | od -c
*/

import (
	chirp "github.com/yak-labs/chirp-lang"

	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime/pprof"
	"strings"

	"github.com/chzyer/readline"
)

var dFlag = flag.String("d", "", "Debugging flags, each a single letter.")
var cFlag = flag.String("c", "", "Immediate command to execute.")
var recoverFlag = flag.Bool("recover", true, "Set to false to disable recover in the REPL.")
var testFlag = flag.Bool("test", false, "Print test summary at end.")

var scriptName string

func saveArgvStarting(fr *chirp.Frame, i int) {
	argv := []chirp.T{}
	for _, a := range flag.Args() {
		argv = append(argv, chirp.MkString(a))
	}
	fr.SetVar("argv", chirp.MkList(argv)) // Deprecated: argv
	fr.SetVar("Argv", chirp.MkList(argv)) // New: Argv
}

func setEnvironInChirp(fr *chirp.Frame, varName string) {
	h := make(chirp.Hash)
	for _, s := range os.Environ() {
		kv := strings.SplitN(s, "=", 2)
		if len(kv) == 2 {
			h[kv[0]] = chirp.MkString(kv[1])
		}
	}
	fr.SetVar(varName, chirp.MkHash(h))
}

func Main() {
	flag.Parse()
	fr := chirp.NewInterpreter()
	setEnvironInChirp(fr, "Env")

	for _, ch := range *dFlag {
		if ch < 256 {
			chirp.Debug[ch] = true
		}
	}

	if cFlag != nil && *cFlag != "" {
		saveArgvStarting(fr, 1)
		fr.Eval(chirp.MkString(*cFlag))
		goto End
	}

	if len(flag.Args()) > 0 {
		// Script mode.
		scriptName = flag.Arg(0)
		contents, err := ioutil.ReadFile(scriptName)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Cannot read file %s: %v", scriptName, err)
			os.Exit(2)
			return
		}
		saveArgvStarting(fr, 1)

		fr.Eval(chirp.MkString(string(contents)))
		goto End
	}

	{
		// Interactive mode.
		home := os.Getenv("HOME")
		if home == "" {
			home = "."
		}

		rl, err := readline.NewEx(&readline.Config{
			Prompt:          "% ",
			HistoryFile:     filepath.Join(home, ".chirp.history"),
			InterruptPrompt: "*SIGINT*",
			EOFPrompt:       "*EOF*",
			// AutoComplete:    completer,
			// HistorySearchFold:   true,
			// FuncFilterInputRune: filterInput,
		})
		if err != nil {
			panic(err)
		}
		defer rl.Close()

		for {
			fmt.Fprint(os.Stderr, "chirp% ") // Prompt to stderr.
			line, err := rl.Readline()
			if err != nil {
				if err.Error() == "EOF" { // TODO: better way?
					goto End
				}
				fmt.Fprintf(os.Stderr, "ERROR in Readline: %s\n", err.Error())
				goto End
			}
			result := EvalStringOrPrintError(fr, string(line))
			if result != "" { // Traditionally, if result is empty, tclsh doesn't print.
				fmt.Println(result)
			}
		}
	}

End:
	logAllCounters()
	if chirp.Debug['h'] {
		pprof.Lookup("heap").WriteTo(os.Stderr, 0)
	}
}

func logAllCounters() {
	if chirp.Debug['c'] {
		chirp.LogAllCounters()
	}

	// Print summary for tests.
	if *testFlag {
		chirp.MustMutex.Lock()
		if chirp.MustFails > 0 {
			fmt.Fprintf(os.Stderr, "TEST FAILS: %q succeeds=%d fails=%d\n", scriptName, chirp.MustSucceeds, chirp.MustFails)
			os.Exit(1)
		}
		fmt.Fprintf(os.Stderr, "Test Done: %q succeeds=%d\n", scriptName, chirp.MustSucceeds)
		chirp.MustMutex.Unlock()
	}
}

func EvalStringOrPrintError(fr *chirp.Frame, cmd string) (out string) {
	if *recoverFlag {
		defer func() {
			if r := recover(); r != nil {
				fmt.Fprintln(os.Stderr, "ERROR: ", r) // Error to stderr.
				out = ""
				return
			}
		}()
	}

	return fr.Eval(chirp.MkString(cmd)).String()
}
<|endoftext|>


<commit_before>package cli

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"os"
	"regexp"
	"sync"

	"github.com/wanelo/image-server/core"
	"github.com/wanelo/image-server/fetcher"
	"github.com/wanelo/image-server/info"
	"github.com/wanelo/image-server/parser"
	"github.com/wanelo/image-server/processor"
	"github.com/wanelo/image-server/uploader"
)

// Item represents all image properties needed for the result of the processing
type Item struct {
	Hash   string
	URL    string
	Width  int
	Height int
}

// ToTabDelimited creates a tab delimited text representation of an Item
func (i Item) ToTabDelimited() string {
	return fmt.Sprintf("%s\t%s\t%d\t%d\n", i.Hash, i.URL, i.Width, i.Height)
}

// Process instanciates image processing based on the tab delimited input that
// contains source image urls and hashes. Each image is processed by a pool of
// of digesters
func Process(sc *core.ServerConfiguration, namespace string, outputs []string, input io.Reader) error {
	done := make(chan struct{})
	defer close(done)

	idsc := enqueueAll(done, input)

	// Start a fixed number of goroutines to read and digest images.
	c := make(chan result) // HLc
	var wg sync.WaitGroup

	numDigesters := int(sc.ProcessorConcurrency)
	wg.Add(numDigesters)

	for i := 0; i < numDigesters; i++ {
		go func() {
			digester(sc, namespace, outputs, done, idsc, c) // HLc
			wg.Done()
		}()
	}
	go func() {
		wg.Wait()
		close(c) // HLc
	}()
	// End of pipeline. OMIT

	for r := range c {
		log.Printf("Completed processing image %v", r.ID)
	}

	return nil
}

func enqueueAll(done <-chan struct{}, input io.Reader) <-chan *Item {
	idsc := make(chan *Item)
	go func() { // HL
		// Close the ids channel after Walk returns.
		defer close(idsc) // HL

		reader := bufio.NewReader(input)
		for {
			line, err := reader.ReadString('\n')
			if err != nil {
				break
			}

			item, err := lineToItem(line)
			if err == nil {
				idsc <- item
			}

		}
	}()
	return idsc
}

// A result is the product of reading and summing a file using MD5.
type result struct {
	ID  string
	Err error
}

// digester processes image items till done is received.
func digester(sc *core.ServerConfiguration, namespace string, outputs []string, done <-chan struct{}, items <-chan *Item, c chan<- result) {
	for item := range items { // HLpaths
		err := downloadOriginal(sc, namespace, item)
		if err != nil {
			continue
		}

		log.Printf("About to process image: %s", item.Hash)
		localOriginalPath := sc.Adapters.Paths.LocalOriginalPath(namespace, item.Hash)

		for _, filename := range outputs {
			err := processImage(sc, namespace, item.Hash, localOriginalPath, filename)
			if err != nil {
				log.Println(err)
				continue
			}
		}

		select {
		case c <- result{item.Hash, err}:
		case <-done:
			return
		}

		fmt.Fprint(os.Stdout, item.ToTabDelimited())
	}
}

func downloadOriginal(sc *core.ServerConfiguration, namespace string, item *Item) error {
	// Image does not have a hash, need to upload source and get image hash
	f := fetcher.OriginalFetcher{Paths: sc.Adapters.Paths}
	imageDetails, downloaded, err := f.Fetch(namespace, item.URL, item.Hash)

	if err != nil {
		return err
	}

	hash := imageDetails.Hash
	if downloaded {
		localOriginalPath := sc.Adapters.Paths.LocalOriginalPath(namespace, hash)
		uploader := uploader.DefaultUploader(sc.RemoteBasePath)

		err := uploader.CreateDirectory(sc.Adapters.Paths.RemoteImageDirectory(namespace, hash))
		if err != nil {
			return err
		}

		destination := sc.Adapters.Paths.RemoteOriginalPath(namespace, hash)
		localInfoPath := sc.Adapters.Paths.LocalInfoPath(namespace, hash)
		remoteInfoPath := sc.Adapters.Paths.RemoteInfoPath(namespace, hash)

		err = info.SaveImageDetail(imageDetails, localInfoPath)
		if err != nil {
			return err
		}

		// upload info
		err = uploader.Upload(localInfoPath, remoteInfoPath)
		if err != nil {
			return err
		}

		// upload original image
		err = uploader.Upload(localOriginalPath, destination)
		if err != nil {
			return err
		}
	}

	item.Width = imageDetails.Width
	item.Height = imageDetails.Height
	item.Hash = hash
	return nil
}

func processImage(sc *core.ServerConfiguration, namespace string, hash string, localOriginalPath string, filename string) error {
	ic, err := parser.NameToConfiguration(sc, filename)
	if err != nil {
		return fmt.Errorf("Error parsing name: %v\n", err)
	}

	ic.Namespace = namespace
	ic.ID = hash

	// process image
	pchan := &processor.ProcessorChannels{
		ImageProcessed: make(chan *core.ImageConfiguration),
		Skipped:        make(chan string),
	}

	localPath := sc.Adapters.Paths.LocalImagePath(namespace, hash, filename)

	p := processor.Processor{
		Source:             localOriginalPath,
		Destination:        localPath,
		ImageConfiguration: ic,
		Channels:           pchan,
	}

	_, err = p.CreateImage()

	if err != nil {
		return err
	}

	select {
	case <-pchan.ImageProcessed:
		log.Println("about to upload to manta")
		uploader := uploader.DefaultUploader(sc.RemoteBasePath)
		remoteResizedPath := sc.Adapters.Paths.RemoteImagePath(namespace, hash, filename)
		err = uploader.Upload(localPath, remoteResizedPath)
		if err != nil {
			log.Println(err)
		}
	case path := <-pchan.Skipped:
		log.Printf("Skipped processing %s", path)
	}

	return nil
}

func lineToItem(line string) (*Item, error) {
	hr, _ := regexp.Compile("([a-z0-9]{32})")
	ur, _ := regexp.Compile("(htt[^\t\n\f\r ]+)")

	hash := hr.FindString(line)
	url := ur.FindString(line)
	return &Item{hash, url, 0, 0}, nil
}
<commit_msg>CLI process: Skip processing images that are already on manta <commit_after>package cli

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"os"
	"regexp"
	"sync"

	"github.com/wanelo/image-server/core"
	"github.com/wanelo/image-server/fetcher"
	"github.com/wanelo/image-server/info"
	"github.com/wanelo/image-server/parser"
	"github.com/wanelo/image-server/processor"
	"github.com/wanelo/image-server/uploader"
	mantaclient "github.com/wanelo/image-server/uploader/manta/client"
)

// Item represents all image properties needed for the result of the processing
type Item struct {
	Hash   string
	URL    string
	Width  int
	Height int
}

// ToTabDelimited creates a tab delimited text representation of an Item
func (i Item) ToTabDelimited() string {
	return fmt.Sprintf("%s\t%s\t%d\t%d\n", i.Hash, i.URL, i.Width, i.Height)
}

// Process instanciates image processing based on the tab delimited input that
// contains source image urls and hashes. Each image is processed by a pool of
// of digesters
func Process(sc *core.ServerConfiguration, namespace string, outputs []string, input io.Reader) error {
	done := make(chan struct{})
	defer close(done)

	idsc := enqueueAll(done, input)

	// Start a fixed number of goroutines to read and digest images.
	c := make(chan result) // HLc
	var wg sync.WaitGroup

	numDigesters := int(sc.ProcessorConcurrency)
	wg.Add(numDigesters)

	for i := 0; i < numDigesters; i++ {
		go func() {
			digester(sc, namespace, outputs, done, idsc, c) // HLc
			wg.Done()
		}()
	}
	go func() {
		wg.Wait()
		close(c) // HLc
	}()
	// End of pipeline. OMIT

	for r := range c {
		log.Printf("Completed processing image %v", r.ID)
	}

	return nil
}

func enqueueAll(done <-chan struct{}, input io.Reader) <-chan *Item {
	idsc := make(chan *Item)
	go func() { // HL
		// Close the ids channel after Walk returns.
		defer close(idsc) // HL

		reader := bufio.NewReader(input)
		for {
			line, err := reader.ReadString('\n')
			if err != nil {
				break
			}

			item, err := lineToItem(line)
			if err == nil {
				idsc <- item
			}

		}
	}()
	return idsc
}

// A result is the product of reading and summing a file using MD5.
type result struct {
	ID  string
	Err error
}

// digester processes image items till done is received.
func digester(sc *core.ServerConfiguration, namespace string, outputs []string, done <-chan struct{}, items <-chan *Item, c chan<- result) {
	for item := range items {
		var itemOutputs []string
		fetchedExistingOutputs := false

		if item.Hash != "" {
			itemOutputs, _ = calculateMissingOutputs(sc, namespace, item.Hash, outputs)
			fetchedExistingOutputs = true
		}

		err := downloadOriginal(sc, namespace, item)
		if err != nil {
			continue
		}

		// have not tried to retrieve existing outputs
		if !fetchedExistingOutputs {
			itemOutputs, err = calculateMissingOutputs(sc, namespace, item.Hash, outputs)
			log.Println(itemOutputs, err)
			if err != nil {
				// process all outputs
				copy(itemOutputs, outputs)
			}
		}

		log.Printf("About to process image: %s", item.Hash)
		localOriginalPath := sc.Adapters.Paths.LocalOriginalPath(namespace, item.Hash)

		for _, filename := range itemOutputs {
			err := processImage(sc, namespace, item.Hash, localOriginalPath, filename)
			if err != nil {
				log.Println(err)
				continue
			}
		}

		select {
		case c <- result{item.Hash, err}:
		case <-done:
			return
		}

		fmt.Fprint(os.Stdout, item.ToTabDelimited())
	}
}

func calculateMissingOutputs(sc *core.ServerConfiguration, namespace string, imageHash string, outputs []string) ([]string, error) {
	// Determine what versions need to be generated
	var itemOutputs []string
	c := mantaclient.DefaultClient()
	remoteDirectory := sc.Adapters.Paths.RemoteImageDirectory(namespace, imageHash)
	entries, err := c.ListDirectory(remoteDirectory)
	if err == nil {
		m := make(map[string]mantaclient.Entry)
		for _, entry := range entries {
			if entry.Type == "object" {
				m[entry.Name] = entry
			} else {
				// got a directory
			}
		}

		for _, output := range outputs {
			if _, ok := m[output]; ok {
				log.Printf("Skipping %s/%s", remoteDirectory, output)
			} else {
				itemOutputs = append(itemOutputs, output)
			}
		}

	} else {
		return nil, err
	}

	return itemOutputs, nil
}

func downloadOriginal(sc *core.ServerConfiguration, namespace string, item *Item) error {
	// Image does not have a hash, need to upload source and get image hash
	f := fetcher.OriginalFetcher{Paths: sc.Adapters.Paths}
	imageDetails, downloaded, err := f.Fetch(namespace, item.URL, item.Hash)

	if err != nil {
		return err
	}

	hash := imageDetails.Hash
	if downloaded {
		localOriginalPath := sc.Adapters.Paths.LocalOriginalPath(namespace, hash)
		uploader := uploader.DefaultUploader(sc.RemoteBasePath)

		err := uploader.CreateDirectory(sc.Adapters.Paths.RemoteImageDirectory(namespace, hash))
		if err != nil {
			return err
		}

		destination := sc.Adapters.Paths.RemoteOriginalPath(namespace, hash)
		localInfoPath := sc.Adapters.Paths.LocalInfoPath(namespace, hash)
		remoteInfoPath := sc.Adapters.Paths.RemoteInfoPath(namespace, hash)

		err = info.SaveImageDetail(imageDetails, localInfoPath)
		if err != nil {
			return err
		}

		// upload info
		err = uploader.Upload(localInfoPath, remoteInfoPath)
		if err != nil {
			return err
		}

		// upload original image
		err = uploader.Upload(localOriginalPath, destination)
		if err != nil {
			return err
		}
	}

	item.Width = imageDetails.Width
	item.Height = imageDetails.Height
	item.Hash = hash
	return nil
}

func processImage(sc *core.ServerConfiguration, namespace string, hash string, localOriginalPath string, filename string) error {
	ic, err := parser.NameToConfiguration(sc, filename)
	if err != nil {
		return fmt.Errorf("Error parsing name: %v\n", err)
	}

	ic.Namespace = namespace
	ic.ID = hash

	// process image
	pchan := &processor.ProcessorChannels{
		ImageProcessed: make(chan *core.ImageConfiguration),
		Skipped:        make(chan string),
	}

	localPath := sc.Adapters.Paths.LocalImagePath(namespace, hash, filename)

	p := processor.Processor{
		Source:             localOriginalPath,
		Destination:        localPath,
		ImageConfiguration: ic,
		Channels:           pchan,
	}

	_, err = p.CreateImage()

	if err != nil {
		return err
	}

	select {
	case <-pchan.ImageProcessed:
		log.Println("about to upload to manta")
		uploader := uploader.DefaultUploader(sc.RemoteBasePath)
		remoteResizedPath := sc.Adapters.Paths.RemoteImagePath(namespace, hash, filename)
		err = uploader.Upload(localPath, remoteResizedPath)
		if err != nil {
			log.Println(err)
		}
	case path := <-pchan.Skipped:
		log.Printf("Skipped processing %s", path)
	}

	return nil
}

func lineToItem(line string) (*Item, error) {
	hr, _ := regexp.Compile("([a-z0-9]{32})")
	ur, _ := regexp.Compile("(htt[^\t\n\f\r ]+)")

	hash := hr.FindString(line)
	url := ur.FindString(line)
	return &Item{hash, url, 0, 0}, nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"
	"os/signal"
	"sort"
	"strconv"
	"syscall"
	"time"

	"github.com/gophergala2016/goad"
	"github.com/gophergala2016/goad/Godeps/_workspace/src/github.com/nsf/termbox-go"
	"github.com/gophergala2016/goad/queue"
)

var (
	url         string
	concurrency uint
	requests    uint
	timeout     uint
	region      string
)

const coldef = termbox.ColorDefault
const nano = 1000000000

func main() {
	flag.UintVar(&concurrency, "c", 10, "number of concurrent requests")
	flag.UintVar(&requests, "n", 1000, "number of total requests to make")
	flag.UintVar(&timeout, "t", 15, "request timeout in seconds")
	flag.StringVar(&region, "r", "us-east-1", "AWS region")
	flag.Parse()

	if len(flag.Args()) < 1 {
		fmt.Println("You must specify a URL")
		os.Exit(1)
	}

	url = flag.Args()[0]

	test, testerr := goad.NewTest(&goad.TestConfig{
		URL:            url,
		Concurrency:    concurrency,
		TotalRequests:  requests,
		RequestTimeout: time.Duration(timeout) * time.Second,
		Region:         region,
	})
	if testerr != nil {
		fmt.Println(testerr)
		os.Exit(1)
	}

	var finalResult queue.RegionsAggData
	defer printSummary(&finalResult)

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM) // but interrupts from kbd are blocked by termbox

	start(test, &finalResult, sigChan)
}

func start(test *goad.Test, finalResult *queue.RegionsAggData, sigChan chan os.Signal) {
	err := termbox.Init()
	if err != nil {
		panic(err)
	}

	defer termbox.Close()

	resultChan := test.Start()
	termbox.Sync()
	renderString(0, 0, "Launching on AWS...", coldef, coldef)
	_, h := termbox.Size()
	renderString(0, h-1, "Press ctrl-c to interrupt", coldef, coldef)
	termbox.Flush()

	go func() {
		for {
			event := termbox.PollEvent()
			if event.Key == 3 {
				sigChan <- syscall.SIGINT
			}
		}
	}()

outer:
	for {
		select {
		case result, ok := <-resultChan:
			if !ok {
				break outer
			}
			//		result.Regions["eu-west-1"] = result.Regions["us-east-1"]
			// sort so that regions always appear in the same order
			var regions []string
			for key := range result.Regions {
				regions = append(regions, key)
			}
			sort.Strings(regions)
			y := 3
			totalReqs := 0
			for _, region := range regions {
				data := result.Regions[region]
				totalReqs += data.TotalReqs
				y = renderRegion(data, y)
				y++
			}

			y = 0
			percentDone := float64(totalReqs) / float64(result.TotalExpectedRequests)
			drawProgressBar(percentDone, y)

			termbox.Flush()
			finalResult.Regions = result.Regions

		case <-sigChan:
			break outer
		}
	}
}

// renderRegion returns the y for the next empty line
func renderRegion(data queue.AggData, y int) int {
	x := 0
	renderString(x, y, "Region: ", termbox.ColorWhite, termbox.ColorBlue)
	x += 8
	regionStr := fmt.Sprintf("%-10s", data.Region)
	renderString(x, y, regionStr, termbox.ColorWhite|termbox.AttrBold, termbox.ColorBlue)
	x = 0
	y++
	headingStr := "   TotReqs   TotBytes    AveTime   AveReq/s Ave1stByte"
	renderString(x, y, headingStr, coldef|termbox.AttrBold, coldef)
	y++
	resultStr := fmt.Sprintf("%10d %10d   %7.2fs %10.2f   %7.2fs", data.TotalReqs, data.TotBytesRead, float64(data.AveTimeForReq)/nano, data.AveReqPerSec, float64(data.AveTimeToFirst)/nano)
	renderString(x, y, resultStr, coldef, coldef)
	y++
	headingStr = "   Slowest    Fastest     Errors"
	renderString(x, y, headingStr, coldef|termbox.AttrBold, coldef)
	y++
	resultStr = fmt.Sprintf("   %7.2f    %7.2f %10d", float64(data.Slowest)/nano, float64(data.Fastest)/nano, totErrors(&data))
	renderString(x, y, resultStr, coldef, coldef)
	y++

	return y
}

func totErrors(data *queue.AggData) int {
	var okReqs int
	for statusStr, value := range data.Statuses {
		status, _ := strconv.Atoi(statusStr)
		if status >= 200 && status <= 299 {
			okReqs += value
		}
	}
	return data.TotalReqs - okReqs
}

func drawProgressBar(percent float64, y int) {
	x := 0
	percentStr := fmt.Sprintf("%5.1f%%            ", percent*100)
	renderString(x, y, percentStr, coldef, coldef)
	y++

	hashes := int(percent * 50)
	if percent > 0.98 {
		hashes = 50
	}
	renderString(x, y, "[", coldef, coldef)

	for x++; x <= hashes; x++ {
		renderString(x, y, "#", coldef, coldef)
	}
	renderString(51, y, "]", coldef, coldef)
}

func renderString(x int, y int, str string, f termbox.Attribute, b termbox.Attribute) {
	for i, c := range str {
		termbox.SetCell(x+i, y, c, f, b)
	}
}

func boldPrintln(msg string) {
	fmt.Printf("\033[1m%s\033[0m\n", msg)
}

func printSummary(result *queue.RegionsAggData) {
	if len(result.Regions) == 0 {
		boldPrintln("No results received")
		return
	}
	boldPrintln("Regional results")
	fmt.Println("")

	for region, data := range result.Regions {
		fmt.Println("Region: " + region)
		boldPrintln("   TotReqs   TotBytes    AveTime   AveReq/s Ave1stByte")
		fmt.Printf("%10d %10d   %7.2fs %10.2f   %7.2fs\n", data.TotalReqs, data.TotBytesRead, float64(data.AveTimeForReq)/nano, data.AveReqPerSec, float64(data.AveTimeToFirst)/nano)
		boldPrintln("   Slowest    Fastest     Errors")
		fmt.Printf("   %7.2f    %7.2f %10d", float64(data.Slowest)/nano, float64(data.Fastest)/nano, totErrors(&data))
		fmt.Println("")
	}

	overall := queue.SumRegionResults(result)

	fmt.Println("")
	boldPrintln("Overall")
	fmt.Println("")
	boldPrintln("   TotReqs   TotBytes    AveTime   AveReq/s Ave1stByte")
	fmt.Printf("%10d %10d   %7.2fs %10.2f   %7.2fs\n", overall.TotalReqs, overall.TotBytesRead, float64(overall.AveTimeForReq)/nano, overall.AveReqPerSec, float64(overall.AveTimeToFirst)/nano)
	boldPrintln("   Slowest    Fastest     Errors")
	fmt.Printf("   %7.2f    %7.2f %10d", float64(overall.Slowest)/nano, float64(overall.Fastest)/nano, totErrors(overall))
	fmt.Println("")

	boldPrintln("HTTPStatus   Requests")
	for statusStr, value := range overall.Statuses {
		fmt.Printf("%10s %10d\n", statusStr, value)
	}
	fmt.Println("")
}
<commit_msg>Improve CLI output<commit_after>package main

import (
	"flag"
	"fmt"
	"os"
	"os/signal"
	"sort"
	"strconv"
	"syscall"
	"time"

	"github.com/gophergala2016/goad"
	"github.com/gophergala2016/goad/Godeps/_workspace/src/github.com/nsf/termbox-go"
	"github.com/gophergala2016/goad/queue"
)

var (
	url         string
	concurrency uint
	requests    uint
	timeout     uint
	region      string
)

const coldef = termbox.ColorDefault
const nano = 1000000000

func main() {
	flag.UintVar(&concurrency, "c", 10, "number of concurrent requests")
	flag.UintVar(&requests, "n", 1000, "number of total requests to make")
	flag.UintVar(&timeout, "t", 15, "request timeout in seconds")
	flag.StringVar(&region, "r", "us-east-1", "AWS region")
	flag.Parse()

	if len(flag.Args()) < 1 {
		fmt.Println("You must specify a URL")
		os.Exit(1)
	}

	url = flag.Args()[0]

	test, testerr := goad.NewTest(&goad.TestConfig{
		URL:            url,
		Concurrency:    concurrency,
		TotalRequests:  requests,
		RequestTimeout: time.Duration(timeout) * time.Second,
		Region:         region,
	})
	if testerr != nil {
		fmt.Println(testerr)
		os.Exit(1)
	}

	var finalResult queue.RegionsAggData
	defer printSummary(&finalResult)

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM) // but interrupts from kbd are blocked by termbox

	start(test, &finalResult, sigChan)
}

func start(test *goad.Test, finalResult *queue.RegionsAggData, sigChan chan os.Signal) {
	err := termbox.Init()
	if err != nil {
		panic(err)
	}

	defer termbox.Close()

	resultChan := test.Start()
	termbox.Sync()
	renderString(0, 0, "Launching on AWS...", coldef, coldef)
	_, h := termbox.Size()
	renderString(0, h-1, "Press ctrl-c to interrupt", coldef, coldef)
	termbox.Flush()

	go func() {
		for {
			event := termbox.PollEvent()
			if event.Key == 3 {
				sigChan <- syscall.SIGINT
			}
		}
	}()

outer:
	for {
		select {
		case result, ok := <-resultChan:
			if !ok {
				break outer
			}
			//		result.Regions["eu-west-1"] = result.Regions["us-east-1"]
			// sort so that regions always appear in the same order
			var regions []string
			for key := range result.Regions {
				regions = append(regions, key)
			}
			sort.Strings(regions)
			y := 3
			totalReqs := 0
			for _, region := range regions {
				data := result.Regions[region]
				totalReqs += data.TotalReqs
				y = renderRegion(data, y)
				y++
			}

			y = 0
			percentDone := float64(totalReqs) / float64(result.TotalExpectedRequests)
			drawProgressBar(percentDone, y)

			termbox.Flush()
			finalResult.Regions = result.Regions

		case <-sigChan:
			break outer
		}
	}
}

// renderRegion returns the y for the next empty line
func renderRegion(data queue.AggData, y int) int {
	x := 0
	renderString(x, y, "Region: ", termbox.ColorWhite, termbox.ColorBlue)
	x += 8
	regionStr := fmt.Sprintf("%-10s", data.Region)
	renderString(x, y, regionStr, termbox.ColorWhite|termbox.AttrBold, termbox.ColorBlue)
	x = 0
	y++
	headingStr := "   TotReqs   TotBytes    AvgTime   AvgReq/s"
	renderString(x, y, headingStr, coldef|termbox.AttrBold, coldef)
	y++
	resultStr := fmt.Sprintf("%10d %10d   %7.3fs %10.2f", data.TotalReqs, data.TotBytesRead, float64(data.AveTimeForReq)/nano, data.AveReqPerSec)
	renderString(x, y, resultStr, coldef, coldef)
	y++
	headingStr = "   Slowest    Fastest   Timeouts  TotErrors"
	renderString(x, y, headingStr, coldef|termbox.AttrBold, coldef)
	y++
	resultStr = fmt.Sprintf("  %7.3fs   %7.3fs %10d %10d", float64(data.Slowest)/nano, float64(data.Fastest)/nano, data.TotalTimedOut, totErrors(&data))
	renderString(x, y, resultStr, coldef, coldef)
	y++

	return y
}

func totErrors(data *queue.AggData) int {
	var okReqs int
	for statusStr, value := range data.Statuses {
		status, _ := strconv.Atoi(statusStr)
		if status >= 200 && status <= 299 {
			okReqs += value
		}
	}
	return data.TotalReqs - okReqs
}

func drawProgressBar(percent float64, y int) {
	x := 0
	percentStr := fmt.Sprintf("%5.1f%%            ", percent*100)
	renderString(x, y, percentStr, coldef, coldef)
	y++

	hashes := int(percent * 50)
	if percent > 0.98 {
		hashes = 50
	}
	renderString(x, y, "[", coldef, coldef)

	for x++; x <= hashes; x++ {
		renderString(x, y, "#", coldef, coldef)
	}
	renderString(51, y, "]", coldef, coldef)
}

func renderString(x int, y int, str string, f termbox.Attribute, b termbox.Attribute) {
	for i, c := range str {
		termbox.SetCell(x+i, y, c, f, b)
	}
}

func boldPrintln(msg string) {
	fmt.Printf("\033[1m%s\033[0m\n", msg)
}

func printData(data *queue.AggData) {
	boldPrintln("   TotReqs   TotBytes    AvgTime   AvgReq/s")
	fmt.Printf("%10d %10d   %7.3fs %10.2f\n", data.TotalReqs, data.TotBytesRead, float64(data.AveTimeForReq)/nano, data.AveReqPerSec)
	boldPrintln("   Slowest    Fastest   Timeouts  TotErrors")
	fmt.Printf("  %7.3fs   %7.3fs %10d %10d", float64(data.Slowest)/nano, float64(data.Fastest)/nano, data.TotalTimedOut, totErrors(data))
	fmt.Println("")
}

func printSummary(result *queue.RegionsAggData) {
	if len(result.Regions) == 0 {
		boldPrintln("No results received")
		return
	}
	boldPrintln("Regional results")
	fmt.Println("")

	for region, data := range result.Regions {
		fmt.Println("Region: " + region)
		printData(&data)
	}

	overall := queue.SumRegionResults(result)

	fmt.Println("")
	boldPrintln("Overall")
	fmt.Println("")
	printData(overall)

	boldPrintln("HTTPStatus   Requests")
	for statusStr, value := range overall.Statuses {
		fmt.Printf("%10s %10d\n", statusStr, value)
	}
	fmt.Println("")
}
<|endoftext|>


<commit_before>package cli

import (
	"context"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/pkg/errors"

	"golang.org/x/crypto/ssh/terminal"

	"github.com/elpinal/coco3/config"
	"github.com/elpinal/coco3/eval"
	"github.com/elpinal/coco3/gate"
	"github.com/elpinal/coco3/parser"

	"github.com/elpinal/coco3/extra"
	eparser "github.com/elpinal/coco3/extra/parser"

	"github.com/jmoiron/sqlx"
	_ "github.com/mattn/go-sqlite3"
)

type CLI struct {
	In  io.Reader
	Out io.Writer
	Err io.Writer

	*config.Config

	db *sqlx.DB

	exitCh chan int
	doneCh chan struct{} // to ensure exiting just after exitCh received

	execute1 func([]byte) error
}

func (c *CLI) Run(args []string) int {
	f := flag.NewFlagSet("coco3", flag.ContinueOnError)
	f.SetOutput(c.Err)
	f.Usage = func() {
		c.Err.Write([]byte("coco3 is a shell.\n"))
		c.Err.Write([]byte("Usage:\n"))
		f.PrintDefaults()
	}

	if c.Config == nil {
		c.Config = &config.Config{}
	}
	c.Config.Init()
	flagC := f.String("c", "", "take first argument as a command to execute")
	flagE := f.Bool("extra", c.Config.Extra, "switch to extra mode")
	if err := f.Parse(args); err != nil {
		return 2
	}
	return c.run(f.Args(), flagC, flagE)
}

func (c *CLI) run(args []string, flagC *string, flagE *bool) int {
	c.exitCh = make(chan int)
	c.doneCh = make(chan struct{})
	defer close(c.doneCh)

	for _, alias := range c.Config.Alias {
		eval.DefAlias(alias[0], alias[1])
	}

	for k, v := range c.Config.Env {
		err := os.Setenv(k, v)
		if err != nil {
			fmt.Fprintln(c.Err, err)
			return 1
		}
	}

	setpath(c.Config.Paths)

	if *flagE {
		// If -extra flag is on, enable extra mode on any command executions.
		c.execute1 = c.executeExtra
	} else {
		c.execute1 = c.execute
	}

	if len(c.Config.StartUpCommand) > 0 {
		done := make(chan struct{})
		go func() {
			if err := c.execute1(c.Config.StartUpCommand); err != nil {
				c.printExecError(err)
				c.exitCh <- 1
			}
			close(done)
		}()
		select {
		case code := <-c.exitCh:
			return code
		case <-done:
		}
	}

	if *flagC != "" {
		go func() {
			if err := c.execute1([]byte(*flagC)); err != nil {
				c.printExecError(err)
				c.exitCh <- 1
				return
			}
			c.exitCh <- 0
		}()
		return <-c.exitCh
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	if len(args) > 0 {
		go c.runFiles(ctx, args)
		return <-c.exitCh
	}

	histRunes, err := c.getHistory(c.Config.HistFile)
	if err != nil {
		fmt.Fprintln(c.Err, err)
		return 1
	}
	g := gate.NewContext(ctx, c.Config, c.In, c.Out, c.Err, histRunes)
	go func(ctx context.Context) {
		for {
			if err := c.interact(g); err != nil {
				c.printExecError(err)
				g.Clear()
			}
			select {
			case <-ctx.Done():
				return
			default:
			}
		}
	}(ctx)
	return <-c.exitCh
}

func (c *CLI) errorf(s string) {
	fmt.Fprintf(c.Err, s)
}

func (c *CLI) errorln(s string) {
	fmt.Fprintln(c.Err, s)
}

func (c *CLI) getHistory(filename string) ([][]rune, error) {
	db, err := sqlx.Connect("sqlite3", filename)
	if err != nil {
		return nil, errors.Wrap(err, "connecting history file")
	}
	_, err = db.Exec(schema)
	if err != nil {
		return nil, errors.Wrap(err, "initializing history file")
	}
	var history []string
	err = db.Select(&history, "select line from command_info")
	if err != nil {
		return nil, errors.Wrap(err, "restoring history")
	}
	// TODO: Is this way proper?
	c.db = db
	return sanitizeHistory(history), nil
}

func (c *CLI) printExecError(err error) {
	if pe, ok := err.(*eparser.ParseError); ok {
		fmt.Fprintln(c.Err, pe.Verbose())
	} else {
		fmt.Fprintln(c.Err, err)
	}
}

// setpath sets the PATH environment variable.
func setpath(args []string) {
	if len(args) == 0 {
		return
	}
	paths := filepath.SplitList(os.Getenv("PATH"))
	var newPaths []string
	for _, path := range paths {
		if contains(args, path) {
			continue
		}
		newPaths = append(newPaths, path)
	}
	newPaths = append(args, newPaths...)
	os.Setenv("PATH", strings.Join(newPaths, string(filepath.ListSeparator)))
}

func contains(xs []string, s string) bool {
	for _, x := range xs {
		if x == s {
			return true
		}
	}
	return false
}

func sanitizeHistory(history []string) [][]rune {
	histRunes := make([][]rune, 0, len(history))
	for _, line := range history {
		if line == "" {
			continue
		}
		l := len(histRunes)
		s := []rune(line)
		if l > 0 && compareRunes(histRunes[l-1], s) {
			continue
		}
		histRunes = append(histRunes, s)
	}
	return histRunes
}

func compareRunes(r1, r2 []rune) bool {
	if len(r1) != len(r2) {
		return false
	}
	for i, r := range r1 {
		if r2[i] != r {
			return false
		}
	}
	return true
}

func (c *CLI) interact(g gate.Gate) error {
	r, end, err := c.read(g)
	if err != nil {
		return err
	}
	if end {
		c.exitCh <- 0
		<-c.doneCh
		return nil
	}
	go c.writeHistory(r)
	if err := c.execute1([]byte(string(r))); err != nil {
		return err
	}
	g.Clear()
	return nil
}

func (c *CLI) read(g gate.Gate) ([]rune, bool, error) {
	defer c.Out.Write([]byte{'\n'})
	oldState, err := terminal.MakeRaw(0)
	if err != nil {
		panic(err)
	}
	defer func() {
		if err := terminal.Restore(0, oldState); err != nil {
			fmt.Fprintln(c.Err, err)
		}
	}()
	r, end, err := g.Read()
	if err != nil {
		return nil, false, err
	}
	return r, end, nil
}

func (c *CLI) writeHistory(r []rune) {
	startTime := time.Now()
	_, err := c.db.Exec("insert into command_info (time, line) values ($1, $2)", startTime, string(r))
	if err != nil {
		fmt.Fprintf(c.Err, "saving history: %v\n", err)
		c.exitCh <- 1
	}
}

const schema = `
create table if not exists command_info (
    time datetime,
    line text
)`

func (c *CLI) execute(b []byte) error {
	f, err := parser.ParseSrc(b)
	if err != nil {
		return err
	}
	e := eval.New(c.In, c.Out, c.Err, c.db)
	err = e.Eval(f.Lines)
	select {
	case code := <-e.ExitCh:
		c.exitCh <- code
		<-c.doneCh
	default:
	}
	return err
}

func (c *CLI) executeExtra(b []byte) error {
	cmd, err := eparser.Parse(b)
	if err != nil {
		return err
	}
	e := extra.New(extra.Option{DB: c.db})
	err = e.Eval(cmd)
	if err == nil {
		return nil
	}
	if pe, ok := err.(*eparser.ParseError); ok {
		pe.Src = string(b)
	}
	return err
}

func (c *CLI) runFiles(ctx context.Context, files []string) {
	for _, file := range files {
		b, err := ioutil.ReadFile(file)
		if err != nil {
			fmt.Fprintln(c.Err, err)
			c.exitCh <- 1
			return
		}
		if err := c.execute1(b); err != nil {
			fmt.Fprintln(c.Err, err)
			c.exitCh <- 1
			return
		}
		select {
		case <-ctx.Done():
			return
		default:
		}
	}
	c.exitCh <- 0
}
<commit_msg>Add CLI.errorp<commit_after>package cli

import (
	"context"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/pkg/errors"

	"golang.org/x/crypto/ssh/terminal"

	"github.com/elpinal/coco3/config"
	"github.com/elpinal/coco3/eval"
	"github.com/elpinal/coco3/gate"
	"github.com/elpinal/coco3/parser"

	"github.com/elpinal/coco3/extra"
	eparser "github.com/elpinal/coco3/extra/parser"

	"github.com/jmoiron/sqlx"
	_ "github.com/mattn/go-sqlite3"
)

type CLI struct {
	In  io.Reader
	Out io.Writer
	Err io.Writer

	*config.Config

	db *sqlx.DB

	exitCh chan int
	doneCh chan struct{} // to ensure exiting just after exitCh received

	execute1 func([]byte) error
}

func (c *CLI) Run(args []string) int {
	f := flag.NewFlagSet("coco3", flag.ContinueOnError)
	f.SetOutput(c.Err)
	f.Usage = func() {
		c.Err.Write([]byte("coco3 is a shell.\n"))
		c.Err.Write([]byte("Usage:\n"))
		f.PrintDefaults()
	}

	if c.Config == nil {
		c.Config = &config.Config{}
	}
	c.Config.Init()
	flagC := f.String("c", "", "take first argument as a command to execute")
	flagE := f.Bool("extra", c.Config.Extra, "switch to extra mode")
	if err := f.Parse(args); err != nil {
		return 2
	}
	return c.run(f.Args(), flagC, flagE)
}

func (c *CLI) run(args []string, flagC *string, flagE *bool) int {
	c.exitCh = make(chan int)
	c.doneCh = make(chan struct{})
	defer close(c.doneCh)

	for _, alias := range c.Config.Alias {
		eval.DefAlias(alias[0], alias[1])
	}

	for k, v := range c.Config.Env {
		err := os.Setenv(k, v)
		if err != nil {
			fmt.Fprintln(c.Err, err)
			return 1
		}
	}

	setpath(c.Config.Paths)

	if *flagE {
		// If -extra flag is on, enable extra mode on any command executions.
		c.execute1 = c.executeExtra
	} else {
		c.execute1 = c.execute
	}

	if len(c.Config.StartUpCommand) > 0 {
		done := make(chan struct{})
		go func() {
			if err := c.execute1(c.Config.StartUpCommand); err != nil {
				c.printExecError(err)
				c.exitCh <- 1
			}
			close(done)
		}()
		select {
		case code := <-c.exitCh:
			return code
		case <-done:
		}
	}

	if *flagC != "" {
		go func() {
			if err := c.execute1([]byte(*flagC)); err != nil {
				c.printExecError(err)
				c.exitCh <- 1
				return
			}
			c.exitCh <- 0
		}()
		return <-c.exitCh
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	if len(args) > 0 {
		go c.runFiles(ctx, args)
		return <-c.exitCh
	}

	histRunes, err := c.getHistory(c.Config.HistFile)
	if err != nil {
		fmt.Fprintln(c.Err, err)
		return 1
	}
	g := gate.NewContext(ctx, c.Config, c.In, c.Out, c.Err, histRunes)
	go func(ctx context.Context) {
		for {
			if err := c.interact(g); err != nil {
				c.printExecError(err)
				g.Clear()
			}
			select {
			case <-ctx.Done():
				return
			default:
			}
		}
	}(ctx)
	return <-c.exitCh
}

func (c *CLI) errorf(s string) {
	fmt.Fprintf(c.Err, s)
}

func (c *CLI) errorln(s string) {
	fmt.Fprintln(c.Err, s)
}

func (c *CLI) errorp(s string) {
	fmt.Fprint(c.Err, s)
}

func (c *CLI) getHistory(filename string) ([][]rune, error) {
	db, err := sqlx.Connect("sqlite3", filename)
	if err != nil {
		return nil, errors.Wrap(err, "connecting history file")
	}
	_, err = db.Exec(schema)
	if err != nil {
		return nil, errors.Wrap(err, "initializing history file")
	}
	var history []string
	err = db.Select(&history, "select line from command_info")
	if err != nil {
		return nil, errors.Wrap(err, "restoring history")
	}
	// TODO: Is this way proper?
	c.db = db
	return sanitizeHistory(history), nil
}

func (c *CLI) printExecError(err error) {
	if pe, ok := err.(*eparser.ParseError); ok {
		fmt.Fprintln(c.Err, pe.Verbose())
	} else {
		fmt.Fprintln(c.Err, err)
	}
}

// setpath sets the PATH environment variable.
func setpath(args []string) {
	if len(args) == 0 {
		return
	}
	paths := filepath.SplitList(os.Getenv("PATH"))
	var newPaths []string
	for _, path := range paths {
		if contains(args, path) {
			continue
		}
		newPaths = append(newPaths, path)
	}
	newPaths = append(args, newPaths...)
	os.Setenv("PATH", strings.Join(newPaths, string(filepath.ListSeparator)))
}

func contains(xs []string, s string) bool {
	for _, x := range xs {
		if x == s {
			return true
		}
	}
	return false
}

func sanitizeHistory(history []string) [][]rune {
	histRunes := make([][]rune, 0, len(history))
	for _, line := range history {
		if line == "" {
			continue
		}
		l := len(histRunes)
		s := []rune(line)
		if l > 0 && compareRunes(histRunes[l-1], s) {
			continue
		}
		histRunes = append(histRunes, s)
	}
	return histRunes
}

func compareRunes(r1, r2 []rune) bool {
	if len(r1) != len(r2) {
		return false
	}
	for i, r := range r1 {
		if r2[i] != r {
			return false
		}
	}
	return true
}

func (c *CLI) interact(g gate.Gate) error {
	r, end, err := c.read(g)
	if err != nil {
		return err
	}
	if end {
		c.exitCh <- 0
		<-c.doneCh
		return nil
	}
	go c.writeHistory(r)
	if err := c.execute1([]byte(string(r))); err != nil {
		return err
	}
	g.Clear()
	return nil
}

func (c *CLI) read(g gate.Gate) ([]rune, bool, error) {
	defer c.Out.Write([]byte{'\n'})
	oldState, err := terminal.MakeRaw(0)
	if err != nil {
		panic(err)
	}
	defer func() {
		if err := terminal.Restore(0, oldState); err != nil {
			fmt.Fprintln(c.Err, err)
		}
	}()
	r, end, err := g.Read()
	if err != nil {
		return nil, false, err
	}
	return r, end, nil
}

func (c *CLI) writeHistory(r []rune) {
	startTime := time.Now()
	_, err := c.db.Exec("insert into command_info (time, line) values ($1, $2)", startTime, string(r))
	if err != nil {
		fmt.Fprintf(c.Err, "saving history: %v\n", err)
		c.exitCh <- 1
	}
}

const schema = `
create table if not exists command_info (
    time datetime,
    line text
)`

func (c *CLI) execute(b []byte) error {
	f, err := parser.ParseSrc(b)
	if err != nil {
		return err
	}
	e := eval.New(c.In, c.Out, c.Err, c.db)
	err = e.Eval(f.Lines)
	select {
	case code := <-e.ExitCh:
		c.exitCh <- code
		<-c.doneCh
	default:
	}
	return err
}

func (c *CLI) executeExtra(b []byte) error {
	cmd, err := eparser.Parse(b)
	if err != nil {
		return err
	}
	e := extra.New(extra.Option{DB: c.db})
	err = e.Eval(cmd)
	if err == nil {
		return nil
	}
	if pe, ok := err.(*eparser.ParseError); ok {
		pe.Src = string(b)
	}
	return err
}

func (c *CLI) runFiles(ctx context.Context, files []string) {
	for _, file := range files {
		b, err := ioutil.ReadFile(file)
		if err != nil {
			fmt.Fprintln(c.Err, err)
			c.exitCh <- 1
			return
		}
		if err := c.execute1(b); err != nil {
			fmt.Fprintln(c.Err, err)
			c.exitCh <- 1
			return
		}
		select {
		case <-ctx.Done():
			return
		default:
		}
	}
	c.exitCh <- 0
}
<|endoftext|>


<commit_before>package cli

import (
	"context"
	"flag"
	"io"
	"os"

	"github.com/nyaxt/otaru/logger"
)

var Log = logger.Registry().Category("cli")

const (
	BufLen = 32 * 1024
)

func Get(ctx context.Context, cfg *CliConfig, args []string) {
	fset := flag.NewFlagSet("get", flag.ExitOnError)
	fset.Parse(args[1:])

	pathstr := fset.Arg(0)
	w := os.Stdout // FIXME

	r, err := NewReader(pathstr, WithCliConfig(cfg), WithContext(ctx))
	if err != nil {
		logger.Criticalf(Log, "%v", err)
		return
	}
	defer r.Close()

	if _, err := io.Copy(w, r); err != nil {
		logger.Criticalf(Log, "%v", err)
		return
	}
}

func Put(ctx context.Context, cfg *CliConfig, args []string) {
	fset := flag.NewFlagSet("put", flag.ExitOnError)
	fset.Parse(args[1:])

	pathstr, localpathstr := fset.Arg(0), fset.Arg(1)
	// FIXME: pathstr may end in /, in which case should join(pathstr, base(localpathstr))

	f, err := os.Open(localpathstr)
	if err != nil {
		logger.Criticalf(Log, "Failed to open source file: \"%s\". err: %v", localpathstr, err)
		return
	}
	defer f.Close()

	w, err := NewWriter(pathstr, WithCliConfig(cfg), WithContext(ctx))
	if err != nil {
		logger.Criticalf(Log, "%v", err)
		return
	}

	if _, err := io.Copy(w, f); err != nil {
		logger.Criticalf(Log, "%v", err)
		return
	}
}
<commit_msg>otaru-cli get: multi paths<commit_after>package cli

import (
	"context"
	"flag"
	"fmt"
	"io"
	"os"
	"path"

	"github.com/nyaxt/otaru/logger"
)

var Log = logger.Registry().Category("cli")

const (
	BufLen = 32 * 1024
)

func Get(ctx context.Context, cfg *CliConfig, args []string) {
	fset := flag.NewFlagSet("get", flag.ExitOnError)
	flagO := fset.String("o", "", "destination file path")
	flagC := fset.String("C", "", "destination dir path")
	fset.Usage = func() {
		fmt.Printf("Usage of %s get:\n", os.Args[0])
		fmt.Printf(" %s get [PATH]...\n", os.Args[0])
		fset.PrintDefaults()
	}
	fset.Parse(args[1:])

	if fset.NArg() == 0 {
		fset.Usage()
		return
	}
	if *flagO != "" && fset.NArg() != 1 {
		logger.Criticalf(Log, "Only one path is allowed when specified -o option.")
		fset.Usage()
		return
	}

	var destdir string
	if *flagC != "" {
		destdir = *flagC
	} else {
		var err error
		destdir, err = os.Getwd()
		if err != nil {
			logger.Criticalf(Log, "Failed to query current dir: %v", err)
			return
		}
	}
	fi, err := os.Stat(destdir)
	if err != nil {
		logger.Criticalf(Log, "Failed to stat target dir: %v", err)
		return
	}
	if !fi.IsDir() {
		logger.Criticalf(Log, "Specified destination is not a dir")
		return
	}

	for _, srcstr := range fset.Args() {
		r, err := NewReader(srcstr, WithCliConfig(cfg), WithContext(ctx))
		if err != nil {
			logger.Criticalf(Log, "%v", err)
			return
		}

		var dest string
		if *flagO != "" {
			dest = *flagO
		} else {
			dest = path.Join(destdir, path.Base(srcstr))
		}
		w, err := os.OpenFile(dest, os.O_RDWR|os.O_CREATE, 0644) // FIXME: r.Stat().FileMode
		if err != nil {
			logger.Criticalf(Log, "%v", err)
			r.Close()
			return
		}
		logger.Infof(Log, "Remote %s -> Local %s", srcstr, dest)

		if _, err := io.Copy(w, r); err != nil {
			logger.Criticalf(Log, "%v", err)
			r.Close()
			w.Close()
			return
		}
		r.Close()
		w.Close()
	}
}

func Put(ctx context.Context, cfg *CliConfig, args []string) {
	fset := flag.NewFlagSet("put", flag.ExitOnError)
	fset.Parse(args[1:])

	pathstr, localpathstr := fset.Arg(0), fset.Arg(1)
	// FIXME: pathstr may end in /, in which case should join(pathstr, base(localpathstr))

	f, err := os.Open(localpathstr)
	if err != nil {
		logger.Criticalf(Log, "Failed to open source file: \"%s\". err: %v", localpathstr, err)
		return
	}
	defer f.Close()

	w, err := NewWriter(pathstr, WithCliConfig(cfg), WithContext(ctx))
	if err != nil {
		logger.Criticalf(Log, "%v", err)
		return
	}

	if _, err := io.Copy(w, f); err != nil {
		logger.Criticalf(Log, "%v", err)
		return
	}
}
<|endoftext|>