<commit_before>// Copyright 2011 Google Inc. All Rights Reserved.
// This file is available under the Apache license.

package main

import (
	"context"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"runtime"
	"strings"
	"syscall"
	"time"

	"github.com/golang/glog"
	"github.com/google/mtail/internal/exporter"
	"github.com/google/mtail/internal/metrics"
	"github.com/google/mtail/internal/mtail"
	"github.com/google/mtail/internal/waker"
	"go.opencensus.io/trace"
)

type seqStringFlag []string

func (f *seqStringFlag) String() string {
	return fmt.Sprint(*f)
}

func (f *seqStringFlag) Set(value string) error {
	for _, v := range strings.Split(value, ",") {
		*f = append(*f, v)
	}
	return nil
}

var logs seqStringFlag

var (
	port               = flag.String("port", "3903", "HTTP port to listen on.")
	address            = flag.String("address", "", "Host or IP address on which to bind HTTP listener")
	unixSocket         = flag.String("unix_socket", "", "UNIX Socket to listen on")
	progs              = flag.String("progs", "", "Name of the directory containing mtail programs")
	ignoreRegexPattern = flag.String("ignore_filename_regex_pattern", "", "")

	version = flag.Bool("version", false, "Print mtail version information.")

	// Compiler behaviour flags.
	oneShot       = flag.Bool("one_shot", false, "Compile the programs, then read the contents of the provided logs from start until EOF, print the values of the metrics store in the given format and exit. This is a debugging flag only, not for production use.")
	oneShotFormat = flag.String("format", "json", "Format to use with -one_shot. This is a debugging flag only, not for production use. Supported formats: json, prometheus.")
	compileOnly   = flag.Bool("compile_only", false, "Compile programs only, do not load the virtual machine.")
	dumpAst       = flag.Bool("dump_ast", false, "Dump AST of programs after parse (to INFO log).")
	dumpAstTypes  = flag.Bool("dump_ast_types", false, "Dump AST of programs with type annotation after typecheck (to INFO log).")
	dumpBytecode  = flag.Bool("dump_bytecode", false, "Dump bytecode of programs (to INFO log).")

	// VM Runtime behaviour flags.
	syslogUseCurrentYear = flag.Bool("syslog_use_current_year", true, "Patch yearless timestamps with the present year.")
	overrideTimezone     = flag.String("override_timezone", "", "If set, use the provided timezone in timestamp conversion, instead of UTC.")
	emitProgLabel        = flag.Bool("emit_prog_label", true, "Emit the 'prog' label in variable exports.")
	emitMetricTimestamp  = flag.Bool("emit_metric_timestamp", false, "Emit the recorded timestamp of a metric.  If disabled (the default) no explicit timestamp is sent to a collector.")
	logRuntimeErrors     = flag.Bool("vm_logs_runtime_errors", true, "Enables logging of runtime errors to the standard log.  Set to false to only have the errors printed to the HTTP console.")

	// Ops flags.
	pollInterval                = flag.Duration("poll_interval", 250*time.Millisecond, "Set the interval to poll all log files for data; must be positive, or zero to disable polling.  With polling mode, only the files found at mtail startup will be polled.")
	expiredMetricGcTickInterval = flag.Duration("expired_metrics_gc_interval", time.Hour, "interval between expired metric garbage collection runs")
	staleLogGcTickInterval      = flag.Duration("stale_log_gc_interval", time.Hour, "interval between stale log garbage collection runs")
	metricPushInterval          = flag.Duration("metric_push_interval", time.Minute, "interval between metric pushes to passive collectors")
	maxRegexpLength             = flag.Int("max_regexp_length", 1024, "The maximum length a mtail regexp expression can have. Excessively long patterns are likely to cause compilation and runtime performance problems.")
	maxRecursionDepth           = flag.Int("max_recursion_depth", 100, "The maximum length a mtail statement can be, as measured by parsed tokens. Excessively long mtail expressions are likely to cause compilation and runtime performance problems.")

	// Debugging flags.
	blockProfileRate     = flag.Int("block_profile_rate", 0, "Nanoseconds of block time before goroutine blocking events reported. 0 turns off.  See https://golang.org/pkg/runtime/#SetBlockProfileRate")
	mutexProfileFraction = flag.Int("mutex_profile_fraction", 0, "Fraction of mutex contention events reported.  0 turns off.  See http://golang.org/pkg/runtime/#SetMutexProfileFraction")

	// Tracing.
	jaegerEndpoint    = flag.String("jaeger_endpoint", "", "If set, collector endpoint URL of jaeger thrift service")
	traceSamplePeriod = flag.Int("trace_sample_period", 0, "Sample period for traces.  If non-zero, every nth trace will be sampled.")

	// Deprecated.
	_ = flag.Bool("disable_fsnotify", true, "DEPRECATED: this flag is no longer in use.")
	_ = flag.Int("metric_push_interval_seconds", 0, "DEPRECATED: use --metric_push_interval instead")
)

func init() {
	flag.Var(&logs, "logs", "List of log files to monitor, separated by commas.  This flag may be specified multiple times.")
}

var (
	// Branch as well as Version and Revision identifies where in the git
	// history the build came from, as supplied by the linker when copmiled
	// with `make'.  The defaults here indicate that the user did not use
	// `make' as instructed.
	Branch   = "invalid:-use-make-to-build"
	Version  = "invalid:-use-make-to-build"
	Revision = "invalid:-use-make-to-build"
)

func main() {
	buildInfo := mtail.BuildInfo{
		Branch:   Branch,
		Version:  Version,
		Revision: Revision,
	}

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "%s\n", buildInfo.String())
		fmt.Fprintf(os.Stderr, "\nUsage:\n")
		flag.PrintDefaults()
	}
	flag.Parse()
	if *version {
		fmt.Println(buildInfo.String())
		os.Exit(0)
	}
	glog.Info(buildInfo.String())
	glog.Infof("Commandline: %q", os.Args)
	if len(flag.Args()) > 0 {
		glog.Exitf("Too many extra arguments specified: %q\n(the logs flag can be repeated, or the filenames separated by commas.)", flag.Args())
	}
	loc, err := time.LoadLocation(*overrideTimezone)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Couldn't parse timezone %q: %s", *overrideTimezone, err)
		os.Exit(1)
	}
	if *blockProfileRate > 0 {
		glog.Infof("Setting block profile rate to %d", *blockProfileRate)
		runtime.SetBlockProfileRate(*blockProfileRate)
	}
	if *mutexProfileFraction > 0 {
		glog.Infof("Setting mutex profile fraction to %d", *mutexProfileFraction)
		runtime.SetMutexProfileFraction(*mutexProfileFraction)
	}
	if *progs == "" {
		glog.Exitf("mtail requires programs that instruct it how to extract metrics from logs; please use the flag -progs to specify the directory containing the programs.")
	}
	if !(*dumpBytecode || *dumpAst || *dumpAstTypes || *compileOnly) {
		if len(logs) == 0 {
			glog.Exitf("mtail requires the names of logs to follow in order to extract logs from them; please use the flag -logs one or more times to specify glob patterns describing these logs.")
		}
	}

	if *traceSamplePeriod > 0 {
		trace.ApplyConfig(trace.Config{DefaultSampler: trace.ProbabilitySampler(1 / float64(*traceSamplePeriod))})
	}
	if *pollInterval == 0 {
		glog.Infof("no poll interval specified; defaulting to 250ms poll")
		*pollInterval = time.Millisecond * 250
	}

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

	sigint := make(chan os.Signal, 1)
	signal.Notify(sigint, os.Interrupt, syscall.SIGTERM)
	go func() {
		sig := <-sigint
		glog.Infof("Received %+v, exiting...", sig)
		cancel()
	}()

	opts := []mtail.Option{
		mtail.ProgramPath(*progs),
		mtail.LogPathPatterns(logs...),
		mtail.IgnoreRegexPattern(*ignoreRegexPattern),
		mtail.SetBuildInfo(buildInfo),
		mtail.OverrideLocation(loc),
		mtail.MetricPushInterval(*metricPushInterval),
		mtail.MaxRegexpLength(*maxRegexpLength),
		mtail.MaxRecursionDepth(*maxRecursionDepth),
	}
	if *logRuntimeErrors {
		opts = append(opts, mtail.LogRuntimeErrors)
	}
	if *staleLogGcTickInterval > 0 {
		staleLogGcWaker := waker.NewTimed(ctx, *staleLogGcTickInterval)
		opts = append(opts, mtail.StaleLogGcWaker(staleLogGcWaker))
	}
	if *pollInterval > 0 {
		logPatternPollWaker := waker.NewTimed(ctx, *pollInterval)
		opts = append(opts, mtail.LogPatternPollWaker(logPatternPollWaker), mtail.LogstreamPollWaker(logPatternPollWaker))
	}
	if *unixSocket == "" {
		opts = append(opts, mtail.BindAddress(*address, *port))
	} else {
		opts = append(opts, mtail.BindUnixSocket(*unixSocket))
	}
	if *oneShot {
		opts = append(opts, mtail.OneShot)
	}
	if *compileOnly {
		opts = append(opts, mtail.CompileOnly)
	}
	if *dumpAst {
		opts = append(opts, mtail.DumpAst)
	}
	if *dumpAstTypes {
		opts = append(opts, mtail.DumpAstTypes)
	}
	if *dumpBytecode {
		opts = append(opts, mtail.DumpBytecode)
	}
	if *syslogUseCurrentYear {
		opts = append(opts, mtail.SyslogUseCurrentYear)
	}
	if !*emitProgLabel {
		opts = append(opts, mtail.OmitProgLabel)
	}
	if *emitMetricTimestamp {
		opts = append(opts, mtail.EmitMetricTimestamp)
	}
	if *jaegerEndpoint != "" {
		opts = append(opts, mtail.JaegerReporter(*jaegerEndpoint))
	}
	store := metrics.NewStore()
	if *expiredMetricGcTickInterval > 0 {
		store.StartGcLoop(ctx, *expiredMetricGcTickInterval)
	}
	m, err := mtail.New(ctx, store, opts...)
	if err != nil {
		glog.Error(err)
		cancel()
		os.Exit(1) //nolint:gocritic // false positive
	}
	err = m.Run()
	if err != nil {
		glog.Error(err)
		cancel()
		os.Exit(1) //nolint:gocritic // false positive
	}
	if *oneShot {
		switch *oneShotFormat {
		case "prometheus":
			e, err := exporter.New(ctx, nil, store)
			if err != nil {
				glog.Error(err)
				cancel()
				os.Exit(1) //nolint:gocritic // false positive
			}
			err = e.Write(os.Stdout)
			if err != nil {
				glog.Error(err)
				cancel()
				os.Exit(1) //nolint:gocritic // false positive
			}
			cancel()
			os.Exit(0) //nolint:gocritic // false positive
		case "json":
			err = store.WriteMetrics(os.Stdout)
			if err != nil {
				glog.Error(err)
				os.Exit(1) //nolint:gocritic // false positive
			}
			cancel()
			os.Exit(0) //nolint:gocritic // false positive
		default:
			glog.Errorf("unsupported format: %q", *oneShotFormat)
			cancel()
			os.Exit(1) //nolint:gocritic // false positive
		}
	}
}
<commit_msg>Change the flag to -one_shot_format<commit_after>// Copyright 2011 Google Inc. All Rights Reserved.
// This file is available under the Apache license.

package main

import (
	"context"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"runtime"
	"strings"
	"syscall"
	"time"

	"github.com/golang/glog"
	"github.com/google/mtail/internal/exporter"
	"github.com/google/mtail/internal/metrics"
	"github.com/google/mtail/internal/mtail"
	"github.com/google/mtail/internal/waker"
	"go.opencensus.io/trace"
)

type seqStringFlag []string

func (f *seqStringFlag) String() string {
	return fmt.Sprint(*f)
}

func (f *seqStringFlag) Set(value string) error {
	for _, v := range strings.Split(value, ",") {
		*f = append(*f, v)
	}
	return nil
}

var logs seqStringFlag

var (
	port               = flag.String("port", "3903", "HTTP port to listen on.")
	address            = flag.String("address", "", "Host or IP address on which to bind HTTP listener")
	unixSocket         = flag.String("unix_socket", "", "UNIX Socket to listen on")
	progs              = flag.String("progs", "", "Name of the directory containing mtail programs")
	ignoreRegexPattern = flag.String("ignore_filename_regex_pattern", "", "")

	version = flag.Bool("version", false, "Print mtail version information.")

	// Compiler behaviour flags.
	oneShot       = flag.Bool("one_shot", false, "Compile the programs, then read the contents of the provided logs from start until EOF, print the values of the metrics store in the given format and exit. This is a debugging flag only, not for production use.")
	oneShotFormat = flag.String("one_shot_format", "json", "Format to use with -one_shot. This is a debugging flag only, not for production use. Supported formats: json, prometheus.")
	compileOnly   = flag.Bool("compile_only", false, "Compile programs only, do not load the virtual machine.")
	dumpAst       = flag.Bool("dump_ast", false, "Dump AST of programs after parse (to INFO log).")
	dumpAstTypes  = flag.Bool("dump_ast_types", false, "Dump AST of programs with type annotation after typecheck (to INFO log).")
	dumpBytecode  = flag.Bool("dump_bytecode", false, "Dump bytecode of programs (to INFO log).")

	// VM Runtime behaviour flags.
	syslogUseCurrentYear = flag.Bool("syslog_use_current_year", true, "Patch yearless timestamps with the present year.")
	overrideTimezone     = flag.String("override_timezone", "", "If set, use the provided timezone in timestamp conversion, instead of UTC.")
	emitProgLabel        = flag.Bool("emit_prog_label", true, "Emit the 'prog' label in variable exports.")
	emitMetricTimestamp  = flag.Bool("emit_metric_timestamp", false, "Emit the recorded timestamp of a metric.  If disabled (the default) no explicit timestamp is sent to a collector.")
	logRuntimeErrors     = flag.Bool("vm_logs_runtime_errors", true, "Enables logging of runtime errors to the standard log.  Set to false to only have the errors printed to the HTTP console.")

	// Ops flags.
	pollInterval                = flag.Duration("poll_interval", 250*time.Millisecond, "Set the interval to poll all log files for data; must be positive, or zero to disable polling.  With polling mode, only the files found at mtail startup will be polled.")
	expiredMetricGcTickInterval = flag.Duration("expired_metrics_gc_interval", time.Hour, "interval between expired metric garbage collection runs")
	staleLogGcTickInterval      = flag.Duration("stale_log_gc_interval", time.Hour, "interval between stale log garbage collection runs")
	metricPushInterval          = flag.Duration("metric_push_interval", time.Minute, "interval between metric pushes to passive collectors")
	maxRegexpLength             = flag.Int("max_regexp_length", 1024, "The maximum length a mtail regexp expression can have. Excessively long patterns are likely to cause compilation and runtime performance problems.")
	maxRecursionDepth           = flag.Int("max_recursion_depth", 100, "The maximum length a mtail statement can be, as measured by parsed tokens. Excessively long mtail expressions are likely to cause compilation and runtime performance problems.")

	// Debugging flags.
	blockProfileRate     = flag.Int("block_profile_rate", 0, "Nanoseconds of block time before goroutine blocking events reported. 0 turns off.  See https://golang.org/pkg/runtime/#SetBlockProfileRate")
	mutexProfileFraction = flag.Int("mutex_profile_fraction", 0, "Fraction of mutex contention events reported.  0 turns off.  See http://golang.org/pkg/runtime/#SetMutexProfileFraction")

	// Tracing.
	jaegerEndpoint    = flag.String("jaeger_endpoint", "", "If set, collector endpoint URL of jaeger thrift service")
	traceSamplePeriod = flag.Int("trace_sample_period", 0, "Sample period for traces.  If non-zero, every nth trace will be sampled.")

	// Deprecated.
	_ = flag.Bool("disable_fsnotify", true, "DEPRECATED: this flag is no longer in use.")
	_ = flag.Int("metric_push_interval_seconds", 0, "DEPRECATED: use --metric_push_interval instead")
)

func init() {
	flag.Var(&logs, "logs", "List of log files to monitor, separated by commas.  This flag may be specified multiple times.")
}

var (
	// Branch as well as Version and Revision identifies where in the git
	// history the build came from, as supplied by the linker when copmiled
	// with `make'.  The defaults here indicate that the user did not use
	// `make' as instructed.
	Branch   = "invalid:-use-make-to-build"
	Version  = "invalid:-use-make-to-build"
	Revision = "invalid:-use-make-to-build"
)

func main() {
	buildInfo := mtail.BuildInfo{
		Branch:   Branch,
		Version:  Version,
		Revision: Revision,
	}

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "%s\n", buildInfo.String())
		fmt.Fprintf(os.Stderr, "\nUsage:\n")
		flag.PrintDefaults()
	}
	flag.Parse()
	if *version {
		fmt.Println(buildInfo.String())
		os.Exit(0)
	}
	glog.Info(buildInfo.String())
	glog.Infof("Commandline: %q", os.Args)
	if len(flag.Args()) > 0 {
		glog.Exitf("Too many extra arguments specified: %q\n(the logs flag can be repeated, or the filenames separated by commas.)", flag.Args())
	}
	loc, err := time.LoadLocation(*overrideTimezone)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Couldn't parse timezone %q: %s", *overrideTimezone, err)
		os.Exit(1)
	}
	if *blockProfileRate > 0 {
		glog.Infof("Setting block profile rate to %d", *blockProfileRate)
		runtime.SetBlockProfileRate(*blockProfileRate)
	}
	if *mutexProfileFraction > 0 {
		glog.Infof("Setting mutex profile fraction to %d", *mutexProfileFraction)
		runtime.SetMutexProfileFraction(*mutexProfileFraction)
	}
	if *progs == "" {
		glog.Exitf("mtail requires programs that instruct it how to extract metrics from logs; please use the flag -progs to specify the directory containing the programs.")
	}
	if !(*dumpBytecode || *dumpAst || *dumpAstTypes || *compileOnly) {
		if len(logs) == 0 {
			glog.Exitf("mtail requires the names of logs to follow in order to extract logs from them; please use the flag -logs one or more times to specify glob patterns describing these logs.")
		}
	}

	if *traceSamplePeriod > 0 {
		trace.ApplyConfig(trace.Config{DefaultSampler: trace.ProbabilitySampler(1 / float64(*traceSamplePeriod))})
	}
	if *pollInterval == 0 {
		glog.Infof("no poll interval specified; defaulting to 250ms poll")
		*pollInterval = time.Millisecond * 250
	}

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

	sigint := make(chan os.Signal, 1)
	signal.Notify(sigint, os.Interrupt, syscall.SIGTERM)
	go func() {
		sig := <-sigint
		glog.Infof("Received %+v, exiting...", sig)
		cancel()
	}()

	opts := []mtail.Option{
		mtail.ProgramPath(*progs),
		mtail.LogPathPatterns(logs...),
		mtail.IgnoreRegexPattern(*ignoreRegexPattern),
		mtail.SetBuildInfo(buildInfo),
		mtail.OverrideLocation(loc),
		mtail.MetricPushInterval(*metricPushInterval),
		mtail.MaxRegexpLength(*maxRegexpLength),
		mtail.MaxRecursionDepth(*maxRecursionDepth),
	}
	if *logRuntimeErrors {
		opts = append(opts, mtail.LogRuntimeErrors)
	}
	if *staleLogGcTickInterval > 0 {
		staleLogGcWaker := waker.NewTimed(ctx, *staleLogGcTickInterval)
		opts = append(opts, mtail.StaleLogGcWaker(staleLogGcWaker))
	}
	if *pollInterval > 0 {
		logPatternPollWaker := waker.NewTimed(ctx, *pollInterval)
		opts = append(opts, mtail.LogPatternPollWaker(logPatternPollWaker), mtail.LogstreamPollWaker(logPatternPollWaker))
	}
	if *unixSocket == "" {
		opts = append(opts, mtail.BindAddress(*address, *port))
	} else {
		opts = append(opts, mtail.BindUnixSocket(*unixSocket))
	}
	if *oneShot {
		opts = append(opts, mtail.OneShot)
	}
	if *compileOnly {
		opts = append(opts, mtail.CompileOnly)
	}
	if *dumpAst {
		opts = append(opts, mtail.DumpAst)
	}
	if *dumpAstTypes {
		opts = append(opts, mtail.DumpAstTypes)
	}
	if *dumpBytecode {
		opts = append(opts, mtail.DumpBytecode)
	}
	if *syslogUseCurrentYear {
		opts = append(opts, mtail.SyslogUseCurrentYear)
	}
	if !*emitProgLabel {
		opts = append(opts, mtail.OmitProgLabel)
	}
	if *emitMetricTimestamp {
		opts = append(opts, mtail.EmitMetricTimestamp)
	}
	if *jaegerEndpoint != "" {
		opts = append(opts, mtail.JaegerReporter(*jaegerEndpoint))
	}
	store := metrics.NewStore()
	if *expiredMetricGcTickInterval > 0 {
		store.StartGcLoop(ctx, *expiredMetricGcTickInterval)
	}
	m, err := mtail.New(ctx, store, opts...)
	if err != nil {
		glog.Error(err)
		cancel()
		os.Exit(1) //nolint:gocritic // false positive
	}
	err = m.Run()
	if err != nil {
		glog.Error(err)
		cancel()
		os.Exit(1) //nolint:gocritic // false positive
	}
	if *oneShot {
		switch *oneShotFormat {
		case "prometheus":
			e, err := exporter.New(ctx, nil, store)
			if err != nil {
				glog.Error(err)
				cancel()
				os.Exit(1) //nolint:gocritic // false positive
			}
			err = e.Write(os.Stdout)
			if err != nil {
				glog.Error(err)
				cancel()
				os.Exit(1) //nolint:gocritic // false positive
			}
			cancel()
			os.Exit(0) //nolint:gocritic // false positive
		case "json":
			err = store.WriteMetrics(os.Stdout)
			if err != nil {
				glog.Error(err)
				os.Exit(1) //nolint:gocritic // false positive
			}
			cancel()
			os.Exit(0) //nolint:gocritic // false positive
		default:
			glog.Errorf("unsupported format: %q", *oneShotFormat)
			cancel()
			os.Exit(1) //nolint:gocritic // false positive
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"container/heap"
	"fmt"
	"reflect"
	"strconv"

	"github.com/hashicorp/go-immutable-radix"
	commonpb "github.com/relab/raft/raftpb"
	"github.com/relab/rkv/rkvpb"
)

// Store is an implementation of raft.StateMachine. It holds client session
// information, and a key-value store.
type Store struct {
	kvs         *iradix.Tree
	sessions    *iradix.Tree
	pendingCmds map[uint64]*Cmds
}

// NewStore initializes and returns a *Store.
func NewStore() *Store {
	return &Store{
		kvs:         iradix.New(),
		sessions:    iradix.New(),
		pendingCmds: make(map[uint64]*Cmds),
	}
}

// Apply implements raft.StateMachine.
func (s *Store) Apply(entry *commonpb.Entry) interface{} {
	switch entry.EntryType {
	case commonpb.EntryNormal:
		var cmd rkvpb.Cmd
		err := cmd.Unmarshal(entry.Data)

		if err != nil {
			panic(fmt.Sprintf("could not unmarshal %v: %v", entry.Data, err))
		}

		return s.applyStore(entry.Index, &cmd)
	case commonpb.EntryConfChange:
		panic("not implemented yet")
	default:
		panic(fmt.Sprintf("got unknown entry type: %v", entry.EntryType))
	}
}

func (s *Store) applyStore(i uint64, cmd *rkvpb.Cmd) interface{} {
	switch cmd.CmdType {
	case rkvpb.Register:
		id := []byte(strconv.FormatUint(i, 10))
		s.sessions, _, _ = s.sessions.Insert(id, uint64(0))

		var pending Cmds
		heap.Init(&pending)
		s.pendingCmds[i] = &pending

		return i
	case rkvpb.Insert:
		var req rkvpb.InsertRequest
		err := req.Unmarshal(cmd.Data)
		if err != nil {
			panic(fmt.Sprintf("could not unmarshal %v: %v", cmd.Data, err))
		}

		id := []byte(strconv.FormatUint(req.ClientID, 10))
		raw, found := s.sessions.Get(id)

		if !found {
			panic(fmt.Sprintf("clientID %v not found", req.ClientID))
		}

		oldSeq, ok := raw.(uint64)

		if !ok {
			panic(fmt.Sprintf("expected uint64 got %s", reflect.TypeOf(raw)))
		}

		nextSeq := oldSeq + 1

		txn := s.kvs.Txn()

		if req.ClientSeq != nextSeq {
			s.PushRequest(req.ClientID, &req)
		} else {
			txn.Insert([]byte(req.Key), req.Value)
			nextSeq++
		}

		for s.HasRequest(req.ClientID, nextSeq) {
			nextReq := s.PopRequest(req.ClientID)
			txn.Insert([]byte(nextReq.Key), nextReq.Value)
			nextSeq++
		}

		s.kvs = txn.CommitOnly()

		if oldSeq != nextSeq {
			s.sessions, _, _ = s.sessions.Insert(id, nextSeq)
		}

		return true
	case rkvpb.Lookup:
		var req rkvpb.LookupRequest
		err := req.Unmarshal(cmd.Data)
		if err != nil {
			panic(err)
		}

		raw, found := s.kvs.Get([]byte(req.Key))

		if !found {
			return "(none)"
		}

		val, ok := raw.(string)

		if !ok {
			panic(fmt.Sprintf("expected string got %s", reflect.TypeOf(raw)))
		}

		return val
	default:
		panic(fmt.Sprintf("got unknown cmd type: %v", cmd.CmdType))
	}
}

// Cmds implements container/heap.
type Cmds []*rkvpb.InsertRequest

func (c Cmds) Len() int {
	return len(c)
}

func (c Cmds) Less(i, j int) bool {
	return c[i].ClientSeq < c[j].ClientSeq
}

func (c Cmds) Swap(i, j int) {
	c[i].ClientSeq, c[j].ClientSeq = c[j].ClientSeq, c[i].ClientSeq
}

// Pop implements container/heap.
func (c *Cmds) Pop() interface{} {
	n := len(*c)
	x := (*c)[n-1]
	*c = (*c)[:n-1]
	return x
}

// Push implements container/heap.
func (c *Cmds) Push(x interface{}) {
	*c = append(*c, x.(*rkvpb.InsertRequest))
}

// PushRequest command onto queue.
func (s *Store) PushRequest(clientID uint64, req *rkvpb.InsertRequest) {
	heap.Push(s.pendingCmds[clientID], req)
}

// HasRequest returns true if the next command in the queue is the next entry in
// the sequence.
func (s *Store) HasRequest(clientID, clientSeq uint64) bool {
	pending := s.pendingCmds[clientID]

	if len(*pending) == 0 {
		return false
	}

	if (*pending)[0].ClientSeq != clientSeq {
		return false
	}

	return true
}

// PopRequest removes the next command from the queue and returns it. Only call
// PopRequest after a successful call to HasRequest.
func (s *Store) PopRequest(clientID uint64) *rkvpb.InsertRequest {
	return heap.Pop(s.pendingCmds[clientID]).(*rkvpb.InsertRequest)
}
<commit_msg>rkvd/store.go: Rename Request to Cmd<commit_after>package main

import (
	"container/heap"
	"fmt"
	"reflect"
	"strconv"

	"github.com/hashicorp/go-immutable-radix"
	commonpb "github.com/relab/raft/raftpb"
	"github.com/relab/rkv/rkvpb"
)

// Store is an implementation of raft.StateMachine. It holds client session
// information, and a key-value store.
type Store struct {
	kvs         *iradix.Tree
	sessions    *iradix.Tree
	pendingCmds map[uint64]*Cmds
}

// NewStore initializes and returns a *Store.
func NewStore() *Store {
	return &Store{
		kvs:         iradix.New(),
		sessions:    iradix.New(),
		pendingCmds: make(map[uint64]*Cmds),
	}
}

// Apply implements raft.StateMachine.
func (s *Store) Apply(entry *commonpb.Entry) interface{} {
	switch entry.EntryType {
	case commonpb.EntryNormal:
		var cmd rkvpb.Cmd
		err := cmd.Unmarshal(entry.Data)

		if err != nil {
			panic(fmt.Sprintf("could not unmarshal %v: %v", entry.Data, err))
		}

		return s.applyStore(entry.Index, &cmd)
	case commonpb.EntryConfChange:
		panic("not implemented yet")
	default:
		panic(fmt.Sprintf("got unknown entry type: %v", entry.EntryType))
	}
}

func (s *Store) applyStore(i uint64, cmd *rkvpb.Cmd) interface{} {
	switch cmd.CmdType {
	case rkvpb.Register:
		id := []byte(strconv.FormatUint(i, 10))
		s.sessions, _, _ = s.sessions.Insert(id, uint64(0))

		var pending Cmds
		heap.Init(&pending)
		s.pendingCmds[i] = &pending

		return i
	case rkvpb.Insert:
		var req rkvpb.InsertRequest
		err := req.Unmarshal(cmd.Data)
		if err != nil {
			panic(fmt.Sprintf("could not unmarshal %v: %v", cmd.Data, err))
		}

		id := []byte(strconv.FormatUint(req.ClientID, 10))
		raw, found := s.sessions.Get(id)

		if !found {
			panic(fmt.Sprintf("clientID %v not found", req.ClientID))
		}

		oldSeq, ok := raw.(uint64)

		if !ok {
			panic(fmt.Sprintf("expected uint64 got %s", reflect.TypeOf(raw)))
		}

		nextSeq := oldSeq + 1

		txn := s.kvs.Txn()

		if req.ClientSeq != nextSeq {
			s.PushCmd(req.ClientID, &req)
		} else {
			txn.Insert([]byte(req.Key), req.Value)
			nextSeq++
		}

		for s.HasCmd(req.ClientID, nextSeq) {
			nextReq := s.PopCmd(req.ClientID)
			txn.Insert([]byte(nextReq.Key), nextReq.Value)
			nextSeq++
		}

		s.kvs = txn.CommitOnly()

		if oldSeq != nextSeq {
			s.sessions, _, _ = s.sessions.Insert(id, nextSeq)
		}

		return true
	case rkvpb.Lookup:
		var req rkvpb.LookupRequest
		err := req.Unmarshal(cmd.Data)
		if err != nil {
			panic(err)
		}

		raw, found := s.kvs.Get([]byte(req.Key))

		if !found {
			return "(none)"
		}

		val, ok := raw.(string)

		if !ok {
			panic(fmt.Sprintf("expected string got %s", reflect.TypeOf(raw)))
		}

		return val
	default:
		panic(fmt.Sprintf("got unknown cmd type: %v", cmd.CmdType))
	}
}

// Cmds implements container/heap.
type Cmds []*rkvpb.InsertRequest

func (c Cmds) Len() int {
	return len(c)
}

func (c Cmds) Less(i, j int) bool {
	return c[i].ClientSeq < c[j].ClientSeq
}

func (c Cmds) Swap(i, j int) {
	c[i].ClientSeq, c[j].ClientSeq = c[j].ClientSeq, c[i].ClientSeq
}

// Pop implements container/heap.
func (c *Cmds) Pop() interface{} {
	n := len(*c)
	x := (*c)[n-1]
	*c = (*c)[:n-1]
	return x
}

// Push implements container/heap.
func (c *Cmds) Push(x interface{}) {
	*c = append(*c, x.(*rkvpb.InsertRequest))
}

// PushCmd command onto queue.
func (s *Store) PushCmd(clientID uint64, req *rkvpb.InsertRequest) {
	heap.Push(s.pendingCmds[clientID], req)
}

// HasCmd returns true if the next command in the queue is the next entry in
// the sequence.
func (s *Store) HasCmd(clientID, clientSeq uint64) bool {
	pending := s.pendingCmds[clientID]

	if len(*pending) == 0 {
		return false
	}

	if (*pending)[0].ClientSeq != clientSeq {
		return false
	}

	return true
}

// PopCmd removes the next command from the queue and returns it. Only call
// PopCmd after a successful call to HasCmd.
func (s *Store) PopCmd(clientID uint64) *rkvpb.InsertRequest {
	return heap.Pop(s.pendingCmds[clientID]).(*rkvpb.InsertRequest)
}
<|endoftext|>


<commit_before>// Copyright (c) 2016, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

package main

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

	"github.com/mvdan/sh/syntax"
)

var (
	write  = flag.Bool("w", false, "write result to file instead of stdout")
	list   = flag.Bool("l", false, "list files whose formatting differs from shfmt's")
	indent = flag.Int("i", 0, "indent: 0 for tabs (default), >0 for number of spaces")
	posix  = flag.Bool("p", false, "parse POSIX shell code instead of bash")

	parseMode         syntax.ParseMode
	printConfig       syntax.PrintConfig
	readBuf, writeBuf bytes.Buffer

	copyBuf = make([]byte, 32*1024)

	out io.Writer
)

func main() {
	flag.Parse()

	out = os.Stdout
	printConfig.Spaces = *indent
	parseMode |= syntax.ParseComments
	if *posix {
		parseMode |= syntax.PosixConformant
	}
	if flag.NArg() == 0 {
		if err := formatStdin(); err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		return
	}
	anyErr := false
	onError := func(err error) {
		anyErr = true
		fmt.Fprintln(os.Stderr, err)
	}
	for _, path := range flag.Args() {
		walk(path, onError)
	}
	if anyErr {
		os.Exit(1)
	}
}

func formatStdin() error {
	if *write || *list {
		return fmt.Errorf("-w and -l can only be used on files")
	}
	prog, err := syntax.Parse(os.Stdin, "", parseMode)
	if err != nil {
		return err
	}
	return printConfig.Fprint(out, prog)
}

var (
	shellFile    = regexp.MustCompile(`\.(sh|bash)$`)
	validShebang = regexp.MustCompile(`^#!\s?/(usr/)?bin/(env\s+)?(sh|bash)`)
	vcsDir       = regexp.MustCompile(`^\.(git|svn|hg)$`)
)

type shellConfidence int

const (
	notShellFile shellConfidence = iota
	ifValidShebang
	isShellFile
)

func getConfidence(info os.FileInfo) shellConfidence {
	name := info.Name()
	switch {
	case info.IsDir(), name[0] == '.', !info.Mode().IsRegular():
		return notShellFile
	case shellFile.MatchString(name):
		return isShellFile
	case strings.Contains(name, "."):
		return notShellFile // different extension
	case info.Size() < 8:
		return notShellFile // cannot possibly hold valid shebang
	default:
		return ifValidShebang
	}
}

func walk(path string, onError func(error)) {
	info, err := os.Stat(path)
	if err != nil {
		onError(err)
		return
	}
	if !info.IsDir() {
		if err := formatPath(path, false); err != nil {
			onError(err)
		}
		return
	}
	filepath.Walk(path, func(path string, info os.FileInfo, err error) error {
		if info.IsDir() && vcsDir.MatchString(info.Name()) {
			return filepath.SkipDir
		}
		if err != nil {
			onError(err)
			return nil
		}
		conf := getConfidence(info)
		if conf == notShellFile {
			return nil
		}
		err = formatPath(path, conf == ifValidShebang)
		if err != nil && !os.IsNotExist(err) {
			onError(err)
		}
		return nil
	})
}

func empty(f *os.File) error {
	if err := f.Truncate(0); err != nil {
		return err
	}
	_, err := f.Seek(0, 0)
	return err
}

func formatPath(path string, checkShebang bool) error {
	openMode := os.O_RDONLY
	if *write {
		openMode = os.O_RDWR
	}
	f, err := os.OpenFile(path, openMode, 0)
	if err != nil {
		return err
	}
	defer f.Close()
	readBuf.Reset()
	if checkShebang {
		n, err := f.Read(copyBuf[:32])
		if err != nil {
			return err
		}
		if !validShebang.Match(copyBuf[:n]) {
			return nil
		}
		readBuf.Write(copyBuf[:n])
	}
	if _, err := io.CopyBuffer(&readBuf, f, copyBuf); err != nil {
		return err
	}
	src := readBuf.Bytes()
	prog, err := syntax.Parse(&readBuf, path, parseMode)
	if err != nil {
		return err
	}
	writeBuf.Reset()
	printConfig.Fprint(&writeBuf, prog)
	res := writeBuf.Bytes()
	if !bytes.Equal(src, res) {
		if *list {
			fmt.Fprintln(out, path)
		}
		if *write {
			if err := empty(f); err != nil {
				return err
			}
			if _, err := f.Write(res); err != nil {
				return err
			}
		}
	}
	if !*list && !*write {
		if _, err := out.Write(res); err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>cmd/shfmt: don't ignore println and f.Close errors<commit_after>// Copyright (c) 2016, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

package main

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

	"github.com/mvdan/sh/syntax"
)

var (
	write  = flag.Bool("w", false, "write result to file instead of stdout")
	list   = flag.Bool("l", false, "list files whose formatting differs from shfmt's")
	indent = flag.Int("i", 0, "indent: 0 for tabs (default), >0 for number of spaces")
	posix  = flag.Bool("p", false, "parse POSIX shell code instead of bash")

	parseMode         syntax.ParseMode
	printConfig       syntax.PrintConfig
	readBuf, writeBuf bytes.Buffer

	copyBuf = make([]byte, 32*1024)

	out io.Writer
)

func main() {
	flag.Parse()

	out = os.Stdout
	printConfig.Spaces = *indent
	parseMode |= syntax.ParseComments
	if *posix {
		parseMode |= syntax.PosixConformant
	}
	if flag.NArg() == 0 {
		if err := formatStdin(); err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		return
	}
	anyErr := false
	onError := func(err error) {
		anyErr = true
		fmt.Fprintln(os.Stderr, err)
	}
	for _, path := range flag.Args() {
		walk(path, onError)
	}
	if anyErr {
		os.Exit(1)
	}
}

func formatStdin() error {
	if *write || *list {
		return fmt.Errorf("-w and -l can only be used on files")
	}
	prog, err := syntax.Parse(os.Stdin, "", parseMode)
	if err != nil {
		return err
	}
	return printConfig.Fprint(out, prog)
}

var (
	shellFile    = regexp.MustCompile(`\.(sh|bash)$`)
	validShebang = regexp.MustCompile(`^#!\s?/(usr/)?bin/(env\s+)?(sh|bash)`)
	vcsDir       = regexp.MustCompile(`^\.(git|svn|hg)$`)
)

type shellConfidence int

const (
	notShellFile shellConfidence = iota
	ifValidShebang
	isShellFile
)

func getConfidence(info os.FileInfo) shellConfidence {
	name := info.Name()
	switch {
	case info.IsDir(), name[0] == '.', !info.Mode().IsRegular():
		return notShellFile
	case shellFile.MatchString(name):
		return isShellFile
	case strings.Contains(name, "."):
		return notShellFile // different extension
	case info.Size() < 8:
		return notShellFile // cannot possibly hold valid shebang
	default:
		return ifValidShebang
	}
}

func walk(path string, onError func(error)) {
	info, err := os.Stat(path)
	if err != nil {
		onError(err)
		return
	}
	if !info.IsDir() {
		if err := formatPath(path, false); err != nil {
			onError(err)
		}
		return
	}
	filepath.Walk(path, func(path string, info os.FileInfo, err error) error {
		if info.IsDir() && vcsDir.MatchString(info.Name()) {
			return filepath.SkipDir
		}
		if err != nil {
			onError(err)
			return nil
		}
		conf := getConfidence(info)
		if conf == notShellFile {
			return nil
		}
		err = formatPath(path, conf == ifValidShebang)
		if err != nil && !os.IsNotExist(err) {
			onError(err)
		}
		return nil
	})
}

func empty(f *os.File) error {
	if err := f.Truncate(0); err != nil {
		return err
	}
	_, err := f.Seek(0, 0)
	return err
}

func formatPath(path string, checkShebang bool) error {
	openMode := os.O_RDONLY
	if *write {
		openMode = os.O_RDWR
	}
	f, err := os.OpenFile(path, openMode, 0)
	if err != nil {
		return err
	}
	defer f.Close()
	readBuf.Reset()
	if checkShebang {
		n, err := f.Read(copyBuf[:32])
		if err != nil {
			return err
		}
		if !validShebang.Match(copyBuf[:n]) {
			return nil
		}
		readBuf.Write(copyBuf[:n])
	}
	if _, err := io.CopyBuffer(&readBuf, f, copyBuf); err != nil {
		return err
	}
	src := readBuf.Bytes()
	prog, err := syntax.Parse(&readBuf, path, parseMode)
	if err != nil {
		return err
	}
	writeBuf.Reset()
	printConfig.Fprint(&writeBuf, prog)
	res := writeBuf.Bytes()
	if !bytes.Equal(src, res) {
		if *list {
			if _, err := fmt.Fprintln(out, path); err != nil {
				return err
			}
		}
		if *write {
			if err := empty(f); err != nil {
				return err
			}
			if _, err := f.Write(res); err != nil {
				return err
			}
			if err := f.Close(); err != nil {
				return err
			}
		}
	}
	if !*list && !*write {
		if _, err := out.Write(res); err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>// Binary p9ufs provides a local 9P2000.L server for the p9 package.
//
// To use, first start the server:
//     htmpfs xxx.tgz
//
// Then, connect using the Linux 9P filesystem:
//     mount -t 9p -o trans=tcp,port=5641 127.0.0.1 /mnt

package main

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"strings"

	"github.com/hugelgupf/p9/fsimpl/templatefs"
	"github.com/hugelgupf/p9/p9"
	"github.com/u-root/u-root/pkg/ulog"
)

var (
	networktype = flag.String("ntype", "tcp4", "Default network type")
	netaddr     = flag.String("addr", ":5641", "Network address")
	verbose     = flag.Bool("verbose", false, "Verbose")

	fs *fileSystem
)

type fileSystem struct {
	root  *directory
	dirs  []*directory
	files []*file
}

func newFileSystem() *fileSystem {
	return &fileSystem{newDirectory(), []*directory{}, []*file{}}
}

func (fs *fileSystem) addFile(filepath string, file *file) error {
	filecmps := strings.Split(filepath, "/")
	if dir, err := fs.getOrCreateDir(fs.root, filecmps[:len(filecmps)-1]); err != nil {
		return err
	} else {
		return fs.createFile(dir, filecmps[len(filecmps)-1], file)
	}
}

func (fs *fileSystem) getOrCreateDir(d *directory, cmps []string) (*directory, error) {
	if len(cmps) == 0 {
		return d, nil
	}

	cmpname := cmps[0]
	if entry, exists := d.entries[cmpname]; exists {
		if dir, ok := entry.(*directory); ok {
			return fs.getOrCreateDir(dir, cmps[1:])
		} else {
			return nil, fmt.Errorf("File already exists with name %s", cmpname)
		}
	} else {
		newDir := newDirectory()
		d.entries[cmpname] = newDir

		newDir.qid.Type = p9.TypeDir
		newDir.qid.Path = uint64(len(fs.dirs))

		fs.dirs = append(fs.dirs, newDir)

		return fs.getOrCreateDir(newDir, cmps[1:])
	}
}

func (fs *fileSystem) createFile(d *directory, filename string, file *file) error {
	if _, exists := d.entries[filename]; exists {
		return fmt.Errorf("File or directory already exists with name %s", filename)
	}
	d.entries[filename] = file

	file.qid.Type = p9.TypeRegular
	file.qid.Path = uint64(len(fs.files))

	fs.files = append(fs.files, file)

	return nil
}

type entry interface {
}

type file struct {
	templatefs.NotDirectoryFile
	templatefs.ReadOnlyFile

	hdr  *tar.Header
	data *bytes.Buffer
	qid  p9.QID
}

func newFile(hdr *tar.Header) *file {
	return &file{
		hdr:  hdr,
		data: &bytes.Buffer{},
		qid:  p9.QID{},
	}
}

func (f *file) Open(mode p9.OpenFlags) (p9.QID, uint32, error) {
	return f.qid, 4096, nil
}

func (f *file) Close() error {
	return nil
}

func (f *file) GetAttr(req p9.AttrMask) (p9.QID, p9.AttrMask, p9.Attr, error) {
	const blockSize = 4096
	attr := &p9.Attr{
		Mode:             p9.FileMode(f.hdr.Mode),
		UID:              0,
		GID:              0,
		NLink:            0,
		RDev:             0,
		Size:             uint64(f.hdr.Size),
		BlockSize:        blockSize,
		Blocks:           uint64(f.hdr.Size / blockSize),
		ATimeSeconds:     uint64(f.hdr.AccessTime.Unix()),
		ATimeNanoSeconds: uint64(f.hdr.AccessTime.UnixNano()),
		MTimeSeconds:     uint64(f.hdr.ModTime.Unix()),
		MTimeNanoSeconds: uint64(f.hdr.ModTime.UnixNano()),
		CTimeSeconds:     uint64(f.hdr.ChangeTime.Unix()),
		CTimeNanoSeconds: uint64(f.hdr.ChangeTime.UnixNano()),
	}
	return f.qid, req, *attr, nil
}

type directory struct {
	templatefs.IsDir
	templatefs.ReadOnlyDir

	entries map[string]entry
	qid     p9.QID
}

func newDirectory() *directory {
	return &directory{entries: map[string]entry{}}
}

func (d *directory) Open(mode p9.OpenFlags) (p9.QID, uint32, error) {
	return d.qid, 4096, nil
}

func (d *directory) Close() error {
	return nil
}

func (d *directory) GetAttr(req p9.AttrMask) (p9.QID, p9.AttrMask, p9.Attr, error) {
	const blockSize = 4096
	attr := &p9.Attr{
		Mode:             p9.FileMode(f.hdr.Mode),
		UID:              0,
		GID:              0,
		NLink:            0,
		RDev:             0,
		Size:             uint64(f.hdr.Size),
		BlockSize:        blockSize,
		Blocks:           uint64(f.hdr.Size / blockSize),
		ATimeSeconds:     uint64(f.hdr.AccessTime.Unix()),
		ATimeNanoSeconds: uint64(f.hdr.AccessTime.UnixNano()),
		MTimeSeconds:     uint64(f.hdr.ModTime.Unix()),
		MTimeNanoSeconds: uint64(f.hdr.ModTime.UnixNano()),
		CTimeSeconds:     uint64(f.hdr.ChangeTime.Unix()),
		CTimeNanoSeconds: uint64(f.hdr.ChangeTime.UnixNano()),
	}
	return f.qid, req, *attr, nil
}

type attacher struct {
	fs *fileSystem
}

func newAttacher(fs *fileSystem) *attacher {
	return &attacher{fs}
}

func (a *attacher) Attach() (p9.File, error) {
	return a.fs.root, nil
}

func main() {
	flag.Parse()

	// Create and add some files to the archive.
	buf := createTestImage()
	fs = readImage(buf)

	// Bind and listen on the socket.
	listener, err := net.Listen(*networktype, *netaddr)
	if err != nil {
		log.Fatalf("err binding: %v", err)
	}

	var opts []p9.ServerOpt
	if *verbose {
		opts = append(opts, p9.WithServerLogger(ulog.Log))
	}

	// Run the server.
	server := p9.NewServer(newAttacher(fs), opts...)
	server.Serve(listener)
}

// Create and add some files to the archive.
func createTestImage() *bytes.Buffer {
	var buf bytes.Buffer

	gztw := gzip.NewWriter(&buf)
	defer func() {
		if err := gztw.Close(); err != nil {
			log.Fatal(err)
		}
	}()

	tw := tar.NewWriter(gztw)
	defer func() {
		if err := tw.Close(); err != nil {
			log.Fatal(err)
		}
	}()

	var files = []struct {
		Name, Body string
	}{
		{"readme.txt", "This archive contains some text files."},
		{"foo/gopher.txt", "Gopher names:\nGeorge\nGeoffrey\nGonzo"},
		{"bar/todo.txt", "Get animal handling license."},
		{"foo/todo2.txt", "harvey lalal"},
		{"abc/123/sean.txt", "lorem ipshum."},
	}
	for _, file := range files {
		hdr := &tar.Header{
			Name: file.Name,
			Mode: 0600,
			Size: int64(len(file.Body)),
		}
		if err := tw.WriteHeader(hdr); err != nil {
			log.Fatal(err)
		}
		if _, err := tw.Write([]byte(file.Body)); err != nil {
			log.Fatal(err)
		}
	}

	return &buf
}

// Read a compressed tar and produce a file hierarchy
func readImage(buf *bytes.Buffer) *fileSystem {
	gzr, err := gzip.NewReader(buf)
	if err != nil {
		log.Fatal(err)
	}
	defer func() {
		if err := gzr.Close(); err != nil {
			log.Fatal(err)
		}
	}()

	fs := newFileSystem()
	tr := tar.NewReader(gzr)
	for {
		hdr, err := tr.Next()
		if err == io.EOF {
			break // End of archive
		}
		if err != nil {
			log.Fatal(err)
		}

		filename := hdr.Name
		file := newFile(hdr)
		if _, err := io.Copy(file.data, tr); err != nil {
			log.Fatal(err)
		}

		fs.addFile(filename, file)
	}

	return fs
}
<commit_msg>p9 wip<commit_after>// Binary p9ufs provides a local 9P2000.L server for the p9 package.
//
// To use, first start the server:
//     htmpfs xxx.tgz
//
// Then, connect using the Linux 9P filesystem:
//     mount -t 9p -o trans=tcp,port=5641 127.0.0.1 /mnt

package main

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"strings"
	"time"

	"github.com/hugelgupf/p9/fsimpl/templatefs"
	"github.com/hugelgupf/p9/p9"
	"github.com/hugelgupf/p9/sys/linux"
	"github.com/u-root/u-root/pkg/ulog"
)

var (
	networktype = flag.String("ntype", "tcp4", "Default network type")
	netaddr     = flag.String("addr", ":5641", "Network address")
	verbose     = flag.Bool("verbose", false, "Verbose")

	fs *fileSystem
)

type fileSystem struct {
	root     *directory
	dirs     []*directory
	files    []*file
	openTime time.Time
}

func newFileSystem() *fileSystem {
	return &fileSystem{newDirectory(), []*directory{}, []*file{}, time.Now()}
}

func (fs *fileSystem) addFile(filepath string, file *file) error {
	filecmps := strings.Split(filepath, "/")
	if dir, err := fs.getOrCreateDir(fs.root, filecmps[:len(filecmps)-1]); err != nil {
		return err
	} else {
		return fs.createFile(dir, filecmps[len(filecmps)-1], file)
	}
}

func (fs *fileSystem) getOrCreateDir(d *directory, cmps []string) (*directory, error) {
	if len(cmps) == 0 {
		return d, nil
	}

	cmpname := cmps[0]
	if entry, exists := d.entries[cmpname]; exists {
		if dir, ok := entry.(*directory); ok {
			return fs.getOrCreateDir(dir, cmps[1:])
		} else {
			return nil, fmt.Errorf("File already exists with name %s", cmpname)
		}
	} else {
		newDir := newDirectory()
		d.entries[cmpname] = newDir

		newDir.qid.Type = p9.TypeDir
		newDir.qid.Path = uint64(len(fs.dirs))

		fs.dirs = append(fs.dirs, newDir)

		return fs.getOrCreateDir(newDir, cmps[1:])
	}
}

func (fs *fileSystem) createFile(d *directory, filename string, file *file) error {
	if _, exists := d.entries[filename]; exists {
		return fmt.Errorf("File or directory already exists with name %s", filename)
	}
	d.entries[filename] = file

	file.qid.Type = p9.TypeRegular
	file.qid.Path = uint64(len(fs.files))

	fs.files = append(fs.files, file)

	return nil
}

type entry interface {
}

type file struct {
	templatefs.NotDirectoryFile
	templatefs.ReadOnlyFile
	templatefs.NotSymlinkFile
	templatefs.NoopRenamed
	p9.DefaultWalkGetAttr

	hdr  *tar.Header
	data *bytes.Buffer
	qid  p9.QID
}

func newFile(hdr *tar.Header) *file {
	return &file{
		hdr:  hdr,
		data: &bytes.Buffer{},
		qid:  p9.QID{},
	}
}

func (f *file) Open(mode p9.OpenFlags) (p9.QID, uint32, error) {
	return f.qid, 4096, nil
}

func (f *file) Close() error {
	return nil
}

func (f *file) GetAttr(req p9.AttrMask) (p9.QID, p9.AttrMask, p9.Attr, error) {
	const blockSize = 4096
	attr := &p9.Attr{
		Mode:             p9.FileMode(f.hdr.Mode),
		UID:              0,
		GID:              0,
		NLink:            0,
		RDev:             0,
		Size:             uint64(f.hdr.Size),
		BlockSize:        blockSize,
		Blocks:           uint64(f.hdr.Size / blockSize),
		ATimeSeconds:     uint64(f.hdr.AccessTime.Unix()),
		ATimeNanoSeconds: uint64(f.hdr.AccessTime.UnixNano()),
		MTimeSeconds:     uint64(f.hdr.ModTime.Unix()),
		MTimeNanoSeconds: uint64(f.hdr.ModTime.UnixNano()),
		CTimeSeconds:     uint64(f.hdr.ChangeTime.Unix()),
		CTimeNanoSeconds: uint64(f.hdr.ChangeTime.UnixNano()),
	}
	return f.qid, req, *attr, nil
}

func (f *file) StatFS() (p9.FSStat, error) {
	return p9.FSStat{}, linux.ENOSYS
}

func (f *file) Walk(names []string) ([]p9.QID, p9.File, error) {
	/*var qids []p9.QID
	last := &Local{path: l.path}

	// A walk with no names is a copy of self.
	if len(names) == 0 {
		qid, _, err := l.info()
		if err != nil {
			return nil, nil, err
		}
		return []p9.QID{qid}, last, nil
	}

	for _, name := range names {
		c := &Local{path: path.Join(last.path, name)}
		qid, _, err := c.info()
		if err != nil {
			return nil, nil, err
		}
		qids = append(qids, qid)
		last = c
	}
	return qids, last, nil*/
	return nil, nil, nil
}

type directory struct {
	templatefs.IsDir
	templatefs.ReadOnlyDir
	templatefs.NotSymlinkFile
	templatefs.NoopRenamed
	p9.DefaultWalkGetAttr

	entries map[string]entry
	qid     p9.QID
}

func newDirectory() *directory {
	return &directory{entries: map[string]entry{}}
}

func (d *directory) Open(mode p9.OpenFlags) (p9.QID, uint32, error) {
	return d.qid, 4096, nil
}

func (d *directory) Close() error {
	return nil
}

func (d *directory) GetAttr(req p9.AttrMask) (p9.QID, p9.AttrMask, p9.Attr, error) {
	const dirSize = 4096
	const blockSize = 4096
	attr := &p9.Attr{
		Mode:             0445,
		UID:              0,
		GID:              0,
		NLink:            0,
		RDev:             0,
		Size:             4096,
		BlockSize:        blockSize,
		Blocks:           uint64(dirSize / blockSize),
		ATimeSeconds:     uint64(fs.openTime.Unix()),
		ATimeNanoSeconds: uint64(fs.openTime.UnixNano()),
		MTimeSeconds:     uint64(fs.openTime.Unix()),
		MTimeNanoSeconds: uint64(fs.openTime.UnixNano()),
		CTimeSeconds:     uint64(fs.openTime.Unix()),
		CTimeNanoSeconds: uint64(fs.openTime.UnixNano()),
	}
	return d.qid, req, *attr, nil
}

func (d *directory) StatFS() (p9.FSStat, error) {
	return p9.FSStat{}, linux.ENOSYS
}

func (d *directory) Walk(names []string) ([]p9.QID, p9.File, error) {
	/*var qids []p9.QID
	last := &Local{path: l.path}

	// A walk with no names is a copy of self.
	if len(names) == 0 {
		qid, _, err := l.info()
		if err != nil {
			return nil, nil, err
		}
		return []p9.QID{qid}, last, nil
	}

	for _, name := range names {
		c := &Local{path: path.Join(last.path, name)}
		qid, _, err := c.info()
		if err != nil {
			return nil, nil, err
		}
		qids = append(qids, qid)
		last = c
	}
	return qids, last, nil*/
	return nil, nil, nil
}

type attacher struct {
	fs *fileSystem
}

func newAttacher(fs *fileSystem) *attacher {
	return &attacher{fs}
}

func (a *attacher) Attach() (p9.File, error) {
	return a.fs.root, nil
}

func main() {
	flag.Parse()

	// Create and add some files to the archive.
	buf := createTestImage()
	fs = readImage(buf)

	// Bind and listen on the socket.
	listener, err := net.Listen(*networktype, *netaddr)
	if err != nil {
		log.Fatalf("err binding: %v", err)
	}

	var opts []p9.ServerOpt
	if *verbose {
		opts = append(opts, p9.WithServerLogger(ulog.Log))
	}

	// Run the server.
	server := p9.NewServer(newAttacher(fs), opts...)
	server.Serve(listener)
}

// Create and add some files to the archive.
func createTestImage() *bytes.Buffer {
	var buf bytes.Buffer

	gztw := gzip.NewWriter(&buf)
	defer func() {
		if err := gztw.Close(); err != nil {
			log.Fatal(err)
		}
	}()

	tw := tar.NewWriter(gztw)
	defer func() {
		if err := tw.Close(); err != nil {
			log.Fatal(err)
		}
	}()

	var files = []struct {
		Name, Body string
	}{
		{"readme.txt", "This archive contains some text files."},
		{"foo/gopher.txt", "Gopher names:\nGeorge\nGeoffrey\nGonzo"},
		{"bar/todo.txt", "Get animal handling license."},
		{"foo/todo2.txt", "harvey lalal"},
		{"abc/123/sean.txt", "lorem ipshum."},
	}
	for _, file := range files {
		hdr := &tar.Header{
			Name: file.Name,
			Mode: 0600,
			Size: int64(len(file.Body)),
		}
		if err := tw.WriteHeader(hdr); err != nil {
			log.Fatal(err)
		}
		if _, err := tw.Write([]byte(file.Body)); err != nil {
			log.Fatal(err)
		}
	}

	return &buf
}

// Read a compressed tar and produce a file hierarchy
func readImage(buf *bytes.Buffer) *fileSystem {
	gzr, err := gzip.NewReader(buf)
	if err != nil {
		log.Fatal(err)
	}
	defer func() {
		if err := gzr.Close(); err != nil {
			log.Fatal(err)
		}
	}()

	fs := newFileSystem()
	tr := tar.NewReader(gzr)
	for {
		hdr, err := tr.Next()
		if err == io.EOF {
			break // End of archive
		}
		if err != nil {
			log.Fatal(err)
		}

		filename := hdr.Name
		file := newFile(hdr)
		if _, err := io.Copy(file.data, tr); err != nil {
			log.Fatal(err)
		}

		fs.addFile(filename, file)
	}

	return fs
}
<|endoftext|>


<commit_before>package service

import (
	"bytes"
	"io"
	"mime/multipart"
	"net/http"
	"os"
	"time"

	"github.com/builderscon/octav/octav/db"
	"github.com/builderscon/octav/octav/model"
	"github.com/builderscon/octav/octav/tools"
	"github.com/lestrrat/go-pdebug"
	"github.com/pkg/errors"
	"golang.org/x/net/context"
	"golang.org/x/sync/errgroup"
	"google.golang.org/cloud/storage"
)

func (v *Sponsor) getMediaBucketName() string {
	v.bucketOnce.Do(func() {
		if v.MediaBucketName == "" {
			v.MediaBucketName = os.Getenv("GOOGLE_STORAGE_MEDIA_BUCKET")
		}
	})
	return v.MediaBucketName
}

func (v *Sponsor) getStorageClient(ctx context.Context) *storage.Client {
	v.storageOnce.Do(func() {
		if v.Storage == nil {
			client, err := defaultStorageClient(ctx)
			if err != nil {
				panic(err.Error())
			}
			v.Storage = client
		}
	})
	return v.Storage
}

func (v *Sponsor) populateRowForCreate(vdb *db.Sponsor, payload model.CreateSponsorRequest) error {
	vdb.EID = tools.UUID()

	vdb.ConferenceID = payload.ConferenceID
	vdb.Name = payload.Name
	vdb.LogoURL1 = payload.LogoURL1
	vdb.URL = payload.URL
	vdb.GroupName = payload.GroupName
	vdb.SortOrder = payload.SortOrder

	if payload.LogoURL2.Valid() {
		vdb.LogoURL2.Valid = true
		vdb.LogoURL2.String = payload.LogoURL2.String
	}

	if payload.LogoURL3.Valid() {
		vdb.LogoURL3.Valid = true
		vdb.LogoURL3.String = payload.LogoURL3.String
	}

	return nil
}

func (v *Sponsor) populateRowForUpdate(vdb *db.Sponsor, payload model.UpdateSponsorRequest) error {
	if payload.Name.Valid() {
		vdb.Name = payload.Name.String
	}

	if payload.LogoURL1.Valid() {
		vdb.LogoURL1 = payload.LogoURL1.String
	}

	if payload.URL.Valid() {
		vdb.URL = payload.URL.String
	}

	if payload.GroupName.Valid() {
		vdb.GroupName = payload.GroupName.String
	}

	if payload.SortOrder.Valid() {
		vdb.SortOrder = int(payload.SortOrder.Int)
	}

	if payload.LogoURL2.Valid() {
		vdb.LogoURL2.Valid = true
		vdb.LogoURL2.String = payload.LogoURL2.String
	}

	if payload.LogoURL3.Valid() {
		vdb.LogoURL3.Valid = true
		vdb.LogoURL3.String = payload.LogoURL3.String
	}

	return nil
}

type finalizeFunc func() error

func (ff finalizeFunc) FinalizeFunc() func() error {
	return ff
}

// Ignorable always returns true, otherwise the caller will have to
// bail out immediately
func (ff finalizeFunc) Ignorable() bool {
	return true
}

func (ff finalizeFunc) Error() string {
	return "operation needs finalization"
}

func (v *Sponsor) CreateFromPayload(ctx context.Context, tx *db.Tx, payload model.AddSponsorRequest, result *model.Sponsor) (err error) {
	su := User{}
	if err := su.IsConferenceAdministrator(tx, payload.ConferenceID, payload.UserID); err != nil {
		return errors.Wrap(err, "creating a featured speaker requires conference administrator privilege")
	}

	if payload.MultipartForm != nil && payload.MultipartForm.File != nil {
		bucketName := v.getMediaBucketName()
		finalizers := make([]func() error, 0, 3)
		for _, field := range []string{"logo1", "logo2", "logo3"} {
			storagecl := v.getStorageClient(ctx)
			fhs := payload.MultipartForm.File[field]
			if len(fhs) == 0 {
				continue
			}

			var imgf multipart.File
			imgf, err = fhs[0].Open()
			if err != nil {
				return errors.Wrap(err, "failed to open logo file from multipart form")
			}

			var imgbuf bytes.Buffer
			if _, err := io.Copy(&imgbuf, imgf); err != nil {
				return errors.Wrap(err, "failed to copy logo image data to memory")
			}
			imgtyp := http.DetectContentType(imgbuf.Bytes())

			// Only work with image/png or image/jpeg
			var suffix string
			switch imgtyp {
			case "image/png":
				suffix = "png"
			case "image/jpeg":
				suffix = "jpeg"
			default:
				return errors.Errorf("Unsupported image type %s", imgtyp)
			}

			// TODO: Validate the image
			// TODO: Avoid Google Storage hardcoding?
			// Upload this to a temporary location, then upon successful write to DB
			// rename it to $conference_id/$sponsor_id
			tmpname := time.Now().UTC().Format("2006-01-02") + "/" + tools.RandomString(64) + "." + suffix
			wc := storagecl.Bucket(bucketName).Object(tmpname).NewWriter(ctx)
			wc.ContentType = imgtyp
			wc.ACL = []storage.ACLRule{{storage.AllUsers, storage.RoleReader}}
			if _, err := io.Copy(wc, &imgbuf); err != nil {
				return errors.Wrap(err, "failed to write image to temporary location")
			}
			// Note: DO NOT defer wc.Close(), as it's part of the write operation.
			// If wc.Close() does not complete w/o errors. the write failed
			if err := wc.Close(); err != nil {
				return errors.Wrap(err, "failed to write image to temporary location")
			}
			thisfield := field
			finalizers = append(finalizers, func() (err error) {
				if pdebug.Enabled {
					g := pdebug.Marker("finalizeFunc for service.Sponsor.CreateFromPayload").BindError(&err)
					defer g.End()
				}
				dstname := result.ConferenceID + "-" + result.ID + "-" + thisfield + "." + suffix
				src := storagecl.Bucket(bucketName).Object(tmpname)
				dst := storagecl.Bucket(bucketName).Object(dstname)

				if pdebug.Enabled {
					pdebug.Printf("Copying %s to %s", tmpname, dstname)
				}
				if _, err = src.CopyTo(ctx, dst, nil); err != nil {
					return errors.Wrapf(err, "failed to copy from '%s' to '%s'", tmpname, dstname)
				}
				if pdebug.Enabled {
					pdebug.Printf("Deleting %s", tmpname)
				}
				if err := src.Delete(ctx); err != nil {
					return errors.Wrapf(err, "failed to delete '%s'", tmpname)
				}
				return nil
			})
		}

		if len(finalizers) > 0 {
			defer func() {
				if pdebug.Enabled {
					g := pdebug.Marker("deferred function from service.Sponsor.CreateFromPayload")
					defer g.End()
				}

				if err != nil || result == nil {
					// no op
					return
				}

				if pdebug.Enabled {
					pdebug.Printf("Creating finalizeFunc for this logo upload")
				}
				// Even though there was no error, create an error value that has a
				// FinalizeFunc() method, so the callee will recognize it
				err = finalizeFunc(func() error {
					var g errgroup.Group
					for _, f := range finalizers {
						g.Go(f)
					}
					return g.Wait()
				})
			}()
		}
	}

	vdb := db.Sponsor{}
	if err := v.Create(tx, &vdb, model.CreateSponsorRequest{payload}); err != nil {
		return errors.Wrap(err, "failed to store in database")
	}

	c := model.Sponsor{}
	if err := c.FromRow(vdb); err != nil {
		return errors.Wrap(err, "failed to populate model from database")
	}

	*result = c
	return nil
}

func (v *Sponsor) UpdateFromPayload(tx *db.Tx, payload model.UpdateSponsorRequest) (err error) {
	if pdebug.Enabled {
		g := pdebug.Marker("service.Sponsor.UpdateFromPayload").BindError(&err)
		defer g.End()
	}

	vdb := db.Sponsor{}
	if err := vdb.LoadByEID(tx, payload.ID); err != nil {
		return errors.Wrap(err, "failed to load featured speaker from database")
	}

	su := User{}
	if err := su.IsConferenceAdministrator(tx, vdb.ConferenceID, payload.UserID); err != nil {
		return errors.Wrap(err, "updating a featured speaker requires conference administrator privilege")
	}

	return errors.Wrap(v.Update(tx, &vdb, payload), "failed to load featured speaker from database")
}

func (v *Sponsor) DeleteFromPayload(tx *db.Tx, payload model.DeleteSponsorRequest) error {
	var m db.Sponsor
	if err := m.LoadByEID(tx, payload.ID); err != nil {
		return errors.Wrap(err, "failed to load featured speaker from database")
	}

	su := User{}
	if err := su.IsConferenceAdministrator(tx, m.ConferenceID, payload.UserID); err != nil {
		return errors.Wrap(err, "deleting venues require administrator privileges")
	}

	return errors.Wrap(v.Delete(tx, m.EID), "failed to delete from database")
}

func (v *Sponsor) ListFromPayload(tx *db.Tx, result *model.SponsorList, payload model.ListSponsorsRequest) error {
	var vdbl db.SponsorList
	if err := vdbl.LoadByConferenceSinceEID(tx, payload.ConferenceID, payload.Since.String, int(payload.Limit.Int)); err != nil {
		return errors.Wrap(err, "failed to load featured speakers from database")
	}

	l := make(model.SponsorList, len(vdbl))
	for i, vdb := range vdbl {
		if err := (l[i]).FromRow(vdb); err != nil {
			return errors.Wrap(err, "failed to populate model from database")
		}

		if err := v.Decorate(tx, &l[i], payload.Lang.String); err != nil {
			return errors.Wrap(err, "failed to decorate venue with associated data")
		}
	}

	*result = l
	return nil
}

func (v *Sponsor) Decorate(tx *db.Tx, speaker *model.Sponsor, lang string) error {
	if lang == "" {
		return nil
	}

	if err := v.ReplaceL10NStrings(tx, speaker, lang); err != nil {
		return errors.Wrap(err, "failed to replace L10N strings")
	}

	return nil
}
<commit_msg>delay loading the client until the very last minute<commit_after>package service

import (
	"bytes"
	"io"
	"mime/multipart"
	"net/http"
	"os"
	"time"

	"github.com/builderscon/octav/octav/db"
	"github.com/builderscon/octav/octav/model"
	"github.com/builderscon/octav/octav/tools"
	"github.com/lestrrat/go-pdebug"
	"github.com/pkg/errors"
	"golang.org/x/net/context"
	"golang.org/x/sync/errgroup"
	"google.golang.org/cloud/storage"
)

func (v *Sponsor) getMediaBucketName() string {
	v.bucketOnce.Do(func() {
		if v.MediaBucketName == "" {
			v.MediaBucketName = os.Getenv("GOOGLE_STORAGE_MEDIA_BUCKET")
		}
	})
	return v.MediaBucketName
}

func (v *Sponsor) getStorageClient(ctx context.Context) *storage.Client {
	v.storageOnce.Do(func() {
		if v.Storage == nil {
			client, err := defaultStorageClient(ctx)
			if err != nil {
				panic(err.Error())
			}
			v.Storage = client
		}
	})
	return v.Storage
}

func (v *Sponsor) populateRowForCreate(vdb *db.Sponsor, payload model.CreateSponsorRequest) error {
	vdb.EID = tools.UUID()

	vdb.ConferenceID = payload.ConferenceID
	vdb.Name = payload.Name
	vdb.LogoURL1 = payload.LogoURL1
	vdb.URL = payload.URL
	vdb.GroupName = payload.GroupName
	vdb.SortOrder = payload.SortOrder

	if payload.LogoURL2.Valid() {
		vdb.LogoURL2.Valid = true
		vdb.LogoURL2.String = payload.LogoURL2.String
	}

	if payload.LogoURL3.Valid() {
		vdb.LogoURL3.Valid = true
		vdb.LogoURL3.String = payload.LogoURL3.String
	}

	return nil
}

func (v *Sponsor) populateRowForUpdate(vdb *db.Sponsor, payload model.UpdateSponsorRequest) error {
	if payload.Name.Valid() {
		vdb.Name = payload.Name.String
	}

	if payload.LogoURL1.Valid() {
		vdb.LogoURL1 = payload.LogoURL1.String
	}

	if payload.URL.Valid() {
		vdb.URL = payload.URL.String
	}

	if payload.GroupName.Valid() {
		vdb.GroupName = payload.GroupName.String
	}

	if payload.SortOrder.Valid() {
		vdb.SortOrder = int(payload.SortOrder.Int)
	}

	if payload.LogoURL2.Valid() {
		vdb.LogoURL2.Valid = true
		vdb.LogoURL2.String = payload.LogoURL2.String
	}

	if payload.LogoURL3.Valid() {
		vdb.LogoURL3.Valid = true
		vdb.LogoURL3.String = payload.LogoURL3.String
	}

	return nil
}

type finalizeFunc func() error

func (ff finalizeFunc) FinalizeFunc() func() error {
	return ff
}

// Ignorable always returns true, otherwise the caller will have to
// bail out immediately
func (ff finalizeFunc) Ignorable() bool {
	return true
}

func (ff finalizeFunc) Error() string {
	return "operation needs finalization"
}

func (v *Sponsor) CreateFromPayload(ctx context.Context, tx *db.Tx, payload model.AddSponsorRequest, result *model.Sponsor) (err error) {
	su := User{}
	if err := su.IsConferenceAdministrator(tx, payload.ConferenceID, payload.UserID); err != nil {
		return errors.Wrap(err, "creating a featured speaker requires conference administrator privilege")
	}

	if payload.MultipartForm != nil && payload.MultipartForm.File != nil {
		bucketName := v.getMediaBucketName()
		finalizers := make([]func() error, 0, 3)
		for _, field := range []string{"logo1", "logo2", "logo3"} {
			fhs := payload.MultipartForm.File[field]
			if len(fhs) == 0 {
				continue
			}

			var imgf multipart.File
			imgf, err = fhs[0].Open()
			if err != nil {
				return errors.Wrap(err, "failed to open logo file from multipart form")
			}

			var imgbuf bytes.Buffer
			if _, err := io.Copy(&imgbuf, imgf); err != nil {
				return errors.Wrap(err, "failed to copy logo image data to memory")
			}
			imgtyp := http.DetectContentType(imgbuf.Bytes())

			// Only work with image/png or image/jpeg
			var suffix string
			switch imgtyp {
			case "image/png":
				suffix = "png"
			case "image/jpeg":
				suffix = "jpeg"
			default:
				return errors.Errorf("Unsupported image type %s", imgtyp)
			}

			// TODO: Validate the image
			// TODO: Avoid Google Storage hardcoding?
			// Upload this to a temporary location, then upon successful write to DB
			// rename it to $conference_id/$sponsor_id
			storagecl := v.getStorageClient(ctx)
			tmpname := time.Now().UTC().Format("2006-01-02") + "/" + tools.RandomString(64) + "." + suffix
			wc := storagecl.Bucket(bucketName).Object(tmpname).NewWriter(ctx)
			wc.ContentType = imgtyp
			wc.ACL = []storage.ACLRule{{storage.AllUsers, storage.RoleReader}}
			if _, err := io.Copy(wc, &imgbuf); err != nil {
				return errors.Wrap(err, "failed to write image to temporary location")
			}
			// Note: DO NOT defer wc.Close(), as it's part of the write operation.
			// If wc.Close() does not complete w/o errors. the write failed
			if err := wc.Close(); err != nil {
				return errors.Wrap(err, "failed to write image to temporary location")
			}
			thisfield := field
			finalizers = append(finalizers, func() (err error) {
				if pdebug.Enabled {
					g := pdebug.Marker("finalizeFunc for service.Sponsor.CreateFromPayload").BindError(&err)
					defer g.End()
				}
				dstname := result.ConferenceID + "-" + result.ID + "-" + thisfield + "." + suffix
				src := storagecl.Bucket(bucketName).Object(tmpname)
				dst := storagecl.Bucket(bucketName).Object(dstname)

				if pdebug.Enabled {
					pdebug.Printf("Copying %s to %s", tmpname, dstname)
				}
				if _, err = src.CopyTo(ctx, dst, nil); err != nil {
					return errors.Wrapf(err, "failed to copy from '%s' to '%s'", tmpname, dstname)
				}
				if pdebug.Enabled {
					pdebug.Printf("Deleting %s", tmpname)
				}
				if err := src.Delete(ctx); err != nil {
					return errors.Wrapf(err, "failed to delete '%s'", tmpname)
				}
				return nil
			})
		}

		if len(finalizers) > 0 {
			defer func() {
				if pdebug.Enabled {
					g := pdebug.Marker("deferred function from service.Sponsor.CreateFromPayload")
					defer g.End()
				}

				if err != nil || result == nil {
					// no op
					return
				}

				if pdebug.Enabled {
					pdebug.Printf("Creating finalizeFunc for this logo upload")
				}
				// Even though there was no error, create an error value that has a
				// FinalizeFunc() method, so the callee will recognize it
				err = finalizeFunc(func() error {
					var g errgroup.Group
					for _, f := range finalizers {
						g.Go(f)
					}
					return g.Wait()
				})
			}()
		}
	}

	vdb := db.Sponsor{}
	if err := v.Create(tx, &vdb, model.CreateSponsorRequest{payload}); err != nil {
		return errors.Wrap(err, "failed to store in database")
	}

	c := model.Sponsor{}
	if err := c.FromRow(vdb); err != nil {
		return errors.Wrap(err, "failed to populate model from database")
	}

	*result = c
	return nil
}

func (v *Sponsor) UpdateFromPayload(tx *db.Tx, payload model.UpdateSponsorRequest) (err error) {
	if pdebug.Enabled {
		g := pdebug.Marker("service.Sponsor.UpdateFromPayload").BindError(&err)
		defer g.End()
	}

	vdb := db.Sponsor{}
	if err := vdb.LoadByEID(tx, payload.ID); err != nil {
		return errors.Wrap(err, "failed to load featured speaker from database")
	}

	su := User{}
	if err := su.IsConferenceAdministrator(tx, vdb.ConferenceID, payload.UserID); err != nil {
		return errors.Wrap(err, "updating a featured speaker requires conference administrator privilege")
	}

	return errors.Wrap(v.Update(tx, &vdb, payload), "failed to load featured speaker from database")
}

func (v *Sponsor) DeleteFromPayload(tx *db.Tx, payload model.DeleteSponsorRequest) error {
	var m db.Sponsor
	if err := m.LoadByEID(tx, payload.ID); err != nil {
		return errors.Wrap(err, "failed to load featured speaker from database")
	}

	su := User{}
	if err := su.IsConferenceAdministrator(tx, m.ConferenceID, payload.UserID); err != nil {
		return errors.Wrap(err, "deleting venues require administrator privileges")
	}

	return errors.Wrap(v.Delete(tx, m.EID), "failed to delete from database")
}

func (v *Sponsor) ListFromPayload(tx *db.Tx, result *model.SponsorList, payload model.ListSponsorsRequest) error {
	var vdbl db.SponsorList
	if err := vdbl.LoadByConferenceSinceEID(tx, payload.ConferenceID, payload.Since.String, int(payload.Limit.Int)); err != nil {
		return errors.Wrap(err, "failed to load featured speakers from database")
	}

	l := make(model.SponsorList, len(vdbl))
	for i, vdb := range vdbl {
		if err := (l[i]).FromRow(vdb); err != nil {
			return errors.Wrap(err, "failed to populate model from database")
		}

		if err := v.Decorate(tx, &l[i], payload.Lang.String); err != nil {
			return errors.Wrap(err, "failed to decorate venue with associated data")
		}
	}

	*result = l
	return nil
}

func (v *Sponsor) Decorate(tx *db.Tx, speaker *model.Sponsor, lang string) error {
	if lang == "" {
		return nil
	}

	if err := v.ReplaceL10NStrings(tx, speaker, lang); err != nil {
		return errors.Wrap(err, "failed to replace L10N strings")
	}

	return nil
}
<|endoftext|>


<commit_before>package rabbitsmpp

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"time"

	"github.com/streadway/amqp"
	"golang.org/x/net/context"
)

const (
	prefetchCount = 100
	prefetchSize  = 0
	global        = false
)

type Consumer interface {
	Consume() (<-chan Job, <-chan error, error)
	Stop() error
	ID() string
}

type consumer struct {
	*client
	channel Channel
	ctx     context.Context
	cancel  context.CancelFunc
}

func NewConsumer(conf Config) (Consumer, error) {
	c := NewClient(conf).(*client)
	ctx, cancel := context.WithCancel(context.Background())

	return &consumer{
		client:  c,
		channel: nil,
		ctx:     ctx,
		cancel:  cancel,
	}, nil
}

func newConsumerWithContext(ctx context.Context, conf Config) (Consumer, error) {
	c := NewClient(conf).(*client)
	ctx, cancel := context.WithCancel(ctx)

	return &consumer{
		client:  c,
		channel: nil,
		ctx:     ctx,
		cancel:  cancel,
	}, nil
}

func (c *consumer) ID() string {
	return c.client.QueueName()
}

func (c *consumer) bindWithRetry() chan *amqp.Error {
	closeChan, err := c.Bind()
	for err != nil {
		log.Println("Failed to bind consumer:", err)
		time.Sleep(5 * time.Second)
		closeChan, err = c.Bind()
	}
	return closeChan
}

func (c *consumer) getConsumeChannel() (<-chan amqp.Delivery, error) {
	ch, err := c.Channel()
	if err != nil {
		return nil, err
	}
	err = ch.Qos(prefetchCount, prefetchSize, global)
	if err != nil {
		return nil, err
	}

	c.channel = ch

	q, err := c.channel.QueueDeclare(
		c.QueueName(), // name
		true,          // durable
		false,         // delete when unused
		false,         // exclusive
		false,         // no-wait
		nil,           // arguments
	)
	if err != nil {
		return nil, err
	}

	return c.channel.Consume(
		q.Name, // queue
		"",     // consumer
		false,  // auto-ackey
		false,  // exclusive
		false,  // no-local
		false,  // no-wait
		nil,    // args
	)
}

func (c *consumer) Consume() (<-chan Job, <-chan error, error) {
	if c.channel != nil {
		return nil, nil, errors.New("consumer already active")
	}
	closeChan, err := c.Bind()
	if err != nil {
		return nil, nil, err
	}
	dlvChan, err := c.getConsumeChannel()
	if err != nil {
		return nil, nil, err
	}
	jobChan := make(chan Job)
	errChan := make(chan error)

	go func() {
		defer func() {
			_ = c.channel.Close()
			c.channel = nil
		}()
		defer close(jobChan)
		defer close(errChan)
		for {
			err = c.consume(dlvChan, closeChan, jobChan)
			// if consume returns without an error, means that it was terminated
			// properly, otherwise something went wrong and it needs to restart
			if err == nil {
				log.Printf("EOF consuming for: %s", c.ID())
				return
			}
			errChan <- err
			log.Println("stopped consuming jobs:", err)
			closeChan = c.bindWithRetry()
			dlvChan, err = c.getConsumeChannel()
			for err != nil {
				time.Sleep(5 * time.Second)
				dlvChan, err = c.getConsumeChannel()
			}
		}
	}()

	return jobChan, errChan, nil
}

func (c *consumer) consume(dlvChan <-chan amqp.Delivery, closeChan <-chan *amqp.Error, jobChan chan<- Job) error {
	for {
		select {
		case d := <-dlvChan:
			j := Job{}
			err := json.Unmarshal(d.Body, &j)
			if err != nil {
				return fmt.Errorf("failed to unmarshal PDU: %v", err)
			}
			j.delivery = &d
			jobChan <- j
		case err := <-closeChan:
			return err
		case <-c.ctx.Done():
			return nil
		}
	}
}

func (c *consumer) Stop() error {
	if c.channel == nil {
		return nil
	}
	// Sends the stop signal
	c.cancel()
	return nil
}
<commit_msg>Change to prefetch to 20<commit_after>package rabbitsmpp

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"time"

	"github.com/streadway/amqp"
	"golang.org/x/net/context"
)

const (
	prefetchCount = 20
	prefetchSize  = 0
	global        = false
)

type Consumer interface {
	Consume() (<-chan Job, <-chan error, error)
	Stop() error
	ID() string
}

type consumer struct {
	*client
	channel Channel
	ctx     context.Context
	cancel  context.CancelFunc
}

func NewConsumer(conf Config) (Consumer, error) {
	c := NewClient(conf).(*client)
	ctx, cancel := context.WithCancel(context.Background())

	return &consumer{
		client:  c,
		channel: nil,
		ctx:     ctx,
		cancel:  cancel,
	}, nil
}

func newConsumerWithContext(ctx context.Context, conf Config) (Consumer, error) {
	c := NewClient(conf).(*client)
	ctx, cancel := context.WithCancel(ctx)

	return &consumer{
		client:  c,
		channel: nil,
		ctx:     ctx,
		cancel:  cancel,
	}, nil
}

func (c *consumer) ID() string {
	return c.client.QueueName()
}

func (c *consumer) bindWithRetry() chan *amqp.Error {
	closeChan, err := c.Bind()
	for err != nil {
		log.Println("Failed to bind consumer:", err)
		time.Sleep(5 * time.Second)
		closeChan, err = c.Bind()
	}
	return closeChan
}

func (c *consumer) getConsumeChannel() (<-chan amqp.Delivery, error) {
	ch, err := c.Channel()
	if err != nil {
		return nil, err
	}
	err = ch.Qos(prefetchCount, prefetchSize, global)
	if err != nil {
		return nil, err
	}

	c.channel = ch

	q, err := c.channel.QueueDeclare(
		c.QueueName(), // name
		true,          // durable
		false,         // delete when unused
		false,         // exclusive
		false,         // no-wait
		nil,           // arguments
	)
	if err != nil {
		return nil, err
	}

	return c.channel.Consume(
		q.Name, // queue
		"",     // consumer
		false,  // auto-ackey
		false,  // exclusive
		false,  // no-local
		false,  // no-wait
		nil,    // args
	)
}

func (c *consumer) Consume() (<-chan Job, <-chan error, error) {
	if c.channel != nil {
		return nil, nil, errors.New("consumer already active")
	}
	closeChan, err := c.Bind()
	if err != nil {
		return nil, nil, err
	}
	dlvChan, err := c.getConsumeChannel()
	if err != nil {
		return nil, nil, err
	}
	jobChan := make(chan Job)
	errChan := make(chan error)

	go func() {
		defer func() {
			_ = c.channel.Close()
			c.channel = nil
		}()
		defer close(jobChan)
		defer close(errChan)
		for {
			err = c.consume(dlvChan, closeChan, jobChan)
			// if consume returns without an error, means that it was terminated
			// properly, otherwise something went wrong and it needs to restart
			if err == nil {
				log.Printf("EOF consuming for: %s", c.ID())
				return
			}
			errChan <- err
			log.Println("stopped consuming jobs:", err)
			closeChan = c.bindWithRetry()
			dlvChan, err = c.getConsumeChannel()
			for err != nil {
				time.Sleep(5 * time.Second)
				dlvChan, err = c.getConsumeChannel()
			}
		}
	}()

	return jobChan, errChan, nil
}

func (c *consumer) consume(dlvChan <-chan amqp.Delivery, closeChan <-chan *amqp.Error, jobChan chan<- Job) error {
	for {
		select {
		case d := <-dlvChan:
			j := Job{}
			err := json.Unmarshal(d.Body, &j)
			if err != nil {
				return fmt.Errorf("failed to unmarshal PDU: %v", err)
			}
			j.delivery = &d
			jobChan <- j
		case err := <-closeChan:
			return err
		case <-c.ctx.Done():
			return nil
		}
	}
}

func (c *consumer) Stop() error {
	if c.channel == nil {
		return nil
	}
	// Sends the stop signal
	c.cancel()
	return nil
}
<|endoftext|>


<commit_before>package cron

import (
	"log"
	"time"

	"github.com/open-falcon/falcon-plus/modules/aggregator/g"
)

type Worker struct {
	Ticker      *time.Ticker
	ClusterItem *g.Cluster
	Quit        chan struct{}
}

func NewWorker(ci *g.Cluster) Worker {
	w := Worker{}
	w.Ticker = time.NewTicker(time.Duration(ci.Step) * time.Second)
	w.Quit = make(chan struct{})
	w.ClusterItem = ci
	return w
}

func (this Worker) Start() {
	go func() {
		for {
			select {
			case <-this.Ticker.C:
				WorkerRun(this.ClusterItem)
			case <-this.Quit:
				if g.Config().Debug {
					log.Println("[I] drop worker", this.ClusterItem)
				}
				this.Ticker.Stop()
				return
			}
		}
	}()
}

func (this Worker) Drop() {
	close(this.Quit)
}

var Workers = make(map[string]Worker)

func deleteNoUseWorker(m map[string]*g.Cluster) {
	del := []string{}
	for key, worker := range Workers {
		if _, ok := m[key]; !ok {
			worker.Drop()
			del = append(del, key)
		}
	}

	for _, key := range del {
		delete(Workers, key)
	}
}

func createWorkerIfNeed(m map[string]*g.Cluster) {
	for key, item := range m {
		if _, ok := Workers[key]; !ok {
			worker := NewWorker(item)
			Workers[key] = worker
			worker.Start()
		}
	}
}
<commit_msg>fix(aggreator): ticker panic when step <= 0<commit_after>package cron

import (
	"log"
	"time"

	"github.com/open-falcon/falcon-plus/modules/aggregator/g"
)

type Worker struct {
	Ticker      *time.Ticker
	ClusterItem *g.Cluster
	Quit        chan struct{}
}

func NewWorker(ci *g.Cluster) Worker {
	w := Worker{}
	w.Ticker = time.NewTicker(time.Duration(ci.Step) * time.Second)
	w.Quit = make(chan struct{})
	w.ClusterItem = ci
	return w
}

func (this Worker) Start() {
	go func() {
		for {
			select {
			case <-this.Ticker.C:
				WorkerRun(this.ClusterItem)
			case <-this.Quit:
				if g.Config().Debug {
					log.Println("[I] drop worker", this.ClusterItem)
				}
				this.Ticker.Stop()
				return
			}
		}
	}()
}

func (this Worker) Drop() {
	close(this.Quit)
}

var Workers = make(map[string]Worker)

func deleteNoUseWorker(m map[string]*g.Cluster) {
	del := []string{}
	for key, worker := range Workers {
		if _, ok := m[key]; !ok {
			worker.Drop()
			del = append(del, key)
		}
	}

	for _, key := range del {
		delete(Workers, key)
	}
}

func createWorkerIfNeed(m map[string]*g.Cluster) {
	for key, item := range m {
		if _, ok := Workers[key]; !ok {
			if item.Step <= 0 {
				log.Println("[W] invalid cluster(step <= 0):", item)
				continue
			}
			worker := NewWorker(item)
			Workers[key] = worker
			worker.Start()
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"github.com/disorganizer/brig/util"
	"github.com/golang/snappy"
	"io"
	"os"
)

var (
	ErrBadBlockIndex = errors.New("Invalid byte index while reading index.")
)

const (
	MaxBlockSize   = 64 * 1024
	HeaderBufSize  = 32
	IndexBlockSize = 30
)

func openFiles(from, to string) (*os.File, *os.File, error) {
	fdFrom, err := os.Open(from)
	if err != nil {
		return nil, nil, err
	}

	fdTo, err := os.OpenFile(to, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0755)
	if err != nil {
		fdFrom.Close()
		return nil, nil, err
	}

	return fdFrom, fdTo, nil
}

type snappyWriter struct {
	rawW          io.Writer
	zipW          io.Writer
	buf           *bytes.Buffer
	index         []BlockIndex
	headerWritten bool
	compression   uint64
}

type snappyReader struct {
	rawR      io.ReadSeeker
	zipR      io.Reader
	index     []BlockIndex
	headerBuf []byte
	tailBuf   []byte
}

type BlockIndex struct {
	fileOffset uint64
	zipOffset  uint64
	zipSize    uint64
}

func (bl *BlockIndex) marshal(buf []byte) {
	fmt.Println(bl.fileOffset, bl.zipOffset, bl.zipSize)
	binary.PutUvarint(buf[00:10], bl.fileOffset)
	binary.PutUvarint(buf[10:20], bl.zipOffset)
	binary.PutUvarint(buf[20:30], bl.zipSize)
}

func (bl *BlockIndex) unmarshal(buf []byte) error {
	var n int
	if bl.fileOffset, n = binary.Uvarint(buf[00:10]); n <= 0 {
		return ErrBadBlockIndex
	}
	if bl.zipOffset, n = binary.Uvarint(buf[10:20]); n <= 0 {
		return ErrBadBlockIndex
	}
	if bl.zipSize, n = binary.Uvarint(buf[20:30]); n <= 0 {
		return ErrBadBlockIndex
	}
	return nil
}

func (sr *snappyReader) Seek(offset int64, whence int) (int64, error) {
	return offset, nil
}

func (sr *snappyReader) getCurrentIndexBlock(curIndex uint64) uint64 {
	// offset | zipOffset | zipSize
	if len(sr.index) == 0 {
		panic("Index len is zero.")
	}

	fileIdx, zipIdx := uint64(0), uint64(0)
	for _, BlockIndex := range sr.index {
		zipIdx += BlockIndex.zipSize
		if zipIdx >= curIndex {
			return BlockIndex.zipOffset
		}
		fileIdx += MaxBlockSize
	}
	return 0
}

// Read a snappy compressed stream, with random access.
func (sr *snappyReader) Read(p []byte) (int, error) {

	// Do on first read when header buffer is empty.
	fmt.Println(len(sr.headerBuf))
	if len(sr.headerBuf) == 0 {

		if _, err := sr.rawR.Read(sr.headerBuf[:cap(sr.headerBuf)]); err != nil {
			fmt.Println(err)
		}
		if _, err := sr.rawR.Seek(-8, os.SEEK_END); err != nil {
			fmt.Println(err)
			return 0, err
		}
		// Read size of tail.
		var buf = make([]byte, 8)
		if n, err := sr.rawR.Read(buf); err != nil || n != 8 {
			fmt.Println(err)
			return n, err
		}

		tailSize, n := binary.Uvarint(buf)
		if n <= 0 {
			return 0, ErrBadBlockIndex
		}

		sr.tailBuf = make([]byte, tailSize)
		if _, err := sr.rawR.Seek(-(int64(tailSize) + 8), os.SEEK_END); err != nil {
			fmt.Println(err)
			return 0, err
		}
		if _, err := sr.rawR.Read(sr.tailBuf); err != nil {
			fmt.Println(err)
		}

		//Build Index
		for i := uint64(0); i < tailSize; i++ {
			b := BlockIndex{}
			b.unmarshal(sr.tailBuf)
			sr.index = append(sr.index, b)
			sr.tailBuf = sr.tailBuf[30:]
			sr.tailBuf = sr.tailBuf[IndexBlockSize:]
		}
		sr.rawR.Seek(HeaderBufSize, os.SEEK_SET)
	}

	// curOff, err := sr.rawR.Seek(0, os.SEEK_CUR)
	// if err != nil {
	// 	fmt.Println(err)
	// }

	curZipIdx := sr.getCurrentIndexBlock(uint64(curOff))
	if err != nil {
		fmt.Println(err)
	}
	newOffset, err := sr.rawR.Seek(int64(curZipIdx), os.SEEK_SET)
	if err != nil {
		fmt.Println(err)
	}
	fmt.Println(newOffset)
	// Identify current block with getCurrentIndexBlock.
	// Go to the begining of the current block.
	// Read block in MaxBlockSize bytes buffer.
	// Read from buffer into p until full or eof.
	// TODO: Don't let snappy read tail.
	return sr.zipR.Read(p)
}

func (sw *snappyWriter) writeHeaderIfNeeded() error {
	if !sw.headerWritten {
		fmt.Println("writing header")
		buf := [32]byte{}
		binary.PutUvarint(buf[00:16], sw.compression)
		binary.PutUvarint(buf[16:32], MaxBlockSize)
		if _, err := sw.rawW.Write(buf[:]); err != nil {
			return err
		}
	}

	sw.headerWritten = true
	return nil
}

func (sw *snappyWriter) appendToBlockIndex(sizeCompressed int) {
	var fOff, zOff, zBlockSize = uint64(0), uint64(0), uint64(sizeCompressed)
	if len(sw.index) > 0 {
		var prevIdx = sw.index[len(sw.index)-1]
		fOff = prevIdx.fileOffset + MaxBlockSize
		zOff = prevIdx.zipOffset + uint64(sizeCompressed)
	}
	sw.index = append(sw.index, BlockIndex{fOff, zOff, zBlockSize})

}

func (sw *snappyWriter) flushBlock(flushSize int) (int, error) {

	// Compress and flush the current block.
	nc, err := sw.zipW.Write(sw.buf.Next(flushSize))
	if err != nil {
		fmt.Println("flushBlock:", err)
		return nc, err
	}

	// Build and update index for the current block.
	sw.appendToBlockIndex(nc)

	return nc, nil
}

// Write a snappy compressed stream with index.
func (sw *snappyWriter) Write(p []byte) (n int, err error) {

	if err := sw.writeHeaderIfNeeded(); err != nil {
		return 0, err
	}

	// Compress only MaxBlockSize equal chunks.
	for {
		n, _ := sw.buf.Write(p[:util.Min(len(p), MaxBlockSize)])
		fmt.Println(n)

		// Flush the current block.
		if sw.buf.Len() >= MaxBlockSize {
			if n, err := sw.flushBlock(MaxBlockSize); err != nil {
				return n, err
			}
			// Forget flushed input.
			p = p[n:]
			continue
		}
		break
	}

	// Fake bytes written, as expeted by some functions.
	return len(p), nil
}

func NewReader(r io.ReadSeeker) io.ReadSeeker {
	return &snappyReader{
		rawR:      r,
		zipR:      snappy.NewReader(io.Reader(r)),
		headerBuf: make([]byte, 0, HeaderBufSize),
		readBuf:   &bytes.Buffer{},
	}
}

func NewWriter(w io.Writer) io.WriteCloser {
	return &snappyWriter{
		zipW:        snappy.NewWriter(w),
		rawW:        w,
		buf:         &bytes.Buffer{},
		compression: 1,
	}
}

func (sw *snappyWriter) Close() error {

	// Write header on empty files.
	sw.writeHeaderIfNeeded()

	// Write remaining bytes left in buffer.
	nc, err := sw.zipW.Write(sw.buf.Bytes())
	if err != nil {
		fmt.Println("Close():", err)
		return err
	}
	sw.appendToBlockIndex(nc)

	// Write compression index tail and close stream.
	indexSize := uint64(IndexBlockSize * len(sw.index))
	tailBuf := make([]byte, indexSize)
	tailBufStart := tailBuf
	if len(sw.index) > 0 {
		for _, blkidx := range sw.index {
			blkidx.marshal(tailBuf)
			tailBuf = tailBuf[IndexBlockSize:]
		}
	}

	n, err := sw.rawW.Write(tailBufStart)
	if err != nil || uint64(n) != indexSize {
		fmt.Println("Close():", err, "n:", n, "idxSize:", indexSize, "TailBufLen:", len(tailBuf))
		return err
	}

	// Write index tail size at the end of stream.
	var tailSizeBuf = make([]byte, 8)
	binary.PutUvarint(tailSizeBuf, indexSize)
	sw.rawW.Write(tailBuf)

	cl, ok := sw.rawW.(io.Closer)
	if ok {
		return cl.Close()
	}

	return nil
}
<commit_msg>compress.go: SizeAccumulator added to writer.<commit_after>package main

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"github.com/disorganizer/brig/util"
	"github.com/golang/snappy"
	"io"
	"os"
)

var (
	ErrBadBlockIndex = errors.New("Invalid byte index while reading index.")
)

const (
	MaxBlockSize   = 64 * 1024
	HeaderBufSize  = 32
	IndexBlockSize = 30
)

func openFiles(from, to string) (*os.File, *os.File, error) {
	fdFrom, err := os.Open(from)
	if err != nil {
		return nil, nil, err
	}

	fdTo, err := os.OpenFile(to, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0755)
	if err != nil {
		fdFrom.Close()
		return nil, nil, err
	}

	return fdFrom, fdTo, nil
}

type snappyWriter struct {
	sizeAcc       *util.SizeAccumulator
	rawW          io.Writer
	zipW          io.Writer
	buf           *bytes.Buffer
	index         []BlockIndex
	headerWritten bool
	compression   uint64
}

type snappyReader struct {
	rawR      io.ReadSeeker
	zipR      io.Reader
	index     []BlockIndex
	headerBuf []byte
	tailBuf   []byte
}

type BlockIndex struct {
	fileOffset uint64
	zipOffset  uint64
	zipSize    uint64
}

func (bl *BlockIndex) marshal(buf []byte) {
	fmt.Println(bl.fileOffset, bl.zipOffset, bl.zipSize)
	binary.PutUvarint(buf[00:10], bl.fileOffset)
	binary.PutUvarint(buf[10:20], bl.zipOffset)
	binary.PutUvarint(buf[20:30], bl.zipSize)
}

func (bl *BlockIndex) unmarshal(buf []byte) error {
	var n int
	if bl.fileOffset, n = binary.Uvarint(buf[00:10]); n <= 0 {
		return ErrBadBlockIndex
	}
	if bl.zipOffset, n = binary.Uvarint(buf[10:20]); n <= 0 {
		return ErrBadBlockIndex
	}
	if bl.zipSize, n = binary.Uvarint(buf[20:30]); n <= 0 {
		return ErrBadBlockIndex
	}
	return nil
}

func (sr *snappyReader) Seek(offset int64, whence int) (int64, error) {
	return offset, nil
}

func (sr *snappyReader) getCurrentIndexBlock(curIndex uint64) uint64 {
	// offset | zipOffset | zipSize
	if len(sr.index) == 0 {
		panic("Index len is zero.")
	}

	fileIdx, zipIdx := uint64(0), uint64(0)
	for _, BlockIndex := range sr.index {
		zipIdx += BlockIndex.zipSize
		if zipIdx >= curIndex {
			return BlockIndex.zipOffset
		}
		fileIdx += MaxBlockSize
	}
	return 0
}

// Read a snappy compressed stream, with random access.
func (sr *snappyReader) Read(p []byte) (int, error) {

	// Do on first read when header buffer is empty.
	fmt.Println(len(sr.headerBuf))
	if len(sr.headerBuf) == 0 {

		if _, err := sr.rawR.Read(sr.headerBuf[:cap(sr.headerBuf)]); err != nil {
			fmt.Println(err)
		}
		if _, err := sr.rawR.Seek(-8, os.SEEK_END); err != nil {
			fmt.Println(err)
			return 0, err
		}
		// Read size of tail.
		var buf = make([]byte, 8)
		if n, err := sr.rawR.Read(buf); err != nil || n != 8 {
			fmt.Println(err)
			return n, err
		}

		tailSize, n := binary.Uvarint(buf)
		if n <= 0 {
			return 0, ErrBadBlockIndex
		}

		sr.tailBuf = make([]byte, tailSize)
		if _, err := sr.rawR.Seek(-(int64(tailSize) + 8), os.SEEK_END); err != nil {
			fmt.Println(err)
			return 0, err
		}
		if _, err := sr.rawR.Read(sr.tailBuf); err != nil {
			fmt.Println(err)
		}

		//Build Index
		for i := uint64(0); i < tailSize; i++ {
			b := BlockIndex{}
			b.unmarshal(sr.tailBuf)
			sr.index = append(sr.index, b)
			sr.tailBuf = sr.tailBuf[30:]
			sr.tailBuf = sr.tailBuf[IndexBlockSize:]
		}
		sr.rawR.Seek(HeaderBufSize, os.SEEK_SET)
	}

	// curOff, err := sr.rawR.Seek(0, os.SEEK_CUR)
	// if err != nil {
	// 	fmt.Println(err)
	// }

	curZipIdx := sr.getCurrentIndexBlock(uint64(curOff))
	if err != nil {
		fmt.Println(err)
	}
	newOffset, err := sr.rawR.Seek(int64(curZipIdx), os.SEEK_SET)
	if err != nil {
		fmt.Println(err)
	}
	fmt.Println(newOffset)
	// Identify current block with getCurrentIndexBlock.
	// Go to the begining of the current block.
	// Read block in MaxBlockSize bytes buffer.
	// Read from buffer into p until full or eof.
	// TODO: Don't let snappy read tail.
	return sr.zipR.Read(p)
}

func (sw *snappyWriter) writeHeaderIfNeeded() error {
	if !sw.headerWritten {
		fmt.Println("writing header")
		buf := [32]byte{}
		binary.PutUvarint(buf[00:16], sw.compression)
		binary.PutUvarint(buf[16:32], MaxBlockSize)
		if _, err := sw.rawW.Write(buf[:]); err != nil {
			return err
		}
	}

	sw.headerWritten = true
	return nil
}

func (sw *snappyWriter) appendToBlockIndex() {
	var fOff, zOff, zBlockSize = uint64(0), uint64(0), uint64(sw.sizeAcc.Size())
	if len(sw.index) > 0 {
		var prevIdx = sw.index[len(sw.index)-1]
		fOff = prevIdx.fileOffset + MaxBlockSize
		zOff = prevIdx.zipOffset + zBlockSize
	}
	sw.index = append(sw.index, BlockIndex{fOff, zOff, zBlockSize})
	sw.sizeAcc.Reset()

}

func (sw *snappyWriter) flushBlock(flushSize int) (int, error) {

	// Compress and flush the current block.
	nc, err := sw.zipW.Write(sw.buf.Next(flushSize))
	if err != nil {
		fmt.Println("flushBlock:", err)
		return nc, err
	}

	// Build and update index for the current block.
	sw.appendToBlockIndex()

	return nc, nil
}

// Write a snappy compressed stream with index.
func (sw *snappyWriter) Write(p []byte) (n int, err error) {

	if err := sw.writeHeaderIfNeeded(); err != nil {
		return 0, err
	}

	// Compress only MaxBlockSize equal chunks.
	for {
		n, _ := sw.buf.Write(p[:util.Min(len(p), MaxBlockSize)])
		fmt.Println(n)

		// Flush the current block.
		if sw.buf.Len() >= MaxBlockSize {
			if n, err := sw.flushBlock(MaxBlockSize); err != nil {
				return n, err
			}
			// Forget flushed input.
			p = p[n:]
			continue
		}
		break
	}

	// Fake bytes written, as expeted by some functions.
	return len(p), nil
}

func NewReader(r io.ReadSeeker) io.ReadSeeker {
	return &snappyReader{
		rawR:      r,
		zipR:      snappy.NewReader(io.Reader(r)),
		headerBuf: make([]byte, 0, HeaderBufSize),
		readBuf:   &bytes.Buffer{},
	}
}

func NewWriter(w io.Writer) io.WriteCloser {
	s := &util.SizeAccumulator{}
	return &snappyWriter{

		sizeAcc:     s,
		zipW:        snappy.NewWriter(io.MultiWriter(w, s)),
		rawW:        w,
		buf:         &bytes.Buffer{},
		compression: 1,
	}
}

func (sw *snappyWriter) Close() error {

	// Write header on empty files.
	sw.writeHeaderIfNeeded()

	// Write remaining bytes left in buffer.
	nc, err := sw.zipW.Write(sw.buf.Bytes())
	if err != nil {
		fmt.Println("Close():", err)
		return err
	}
	sw.appendToBlockIndex()

	// Write compression index tail and close stream.
	indexSize := uint64(IndexBlockSize * len(sw.index))
	tailBuf := make([]byte, indexSize)
	tailBufStart := tailBuf
	if len(sw.index) > 0 {
		for _, blkidx := range sw.index {
			blkidx.marshal(tailBuf)
			tailBuf = tailBuf[IndexBlockSize:]
		}
	}

	n, err := sw.rawW.Write(tailBufStart)
	if err != nil || uint64(n) != indexSize {
		fmt.Println("Close():", err, "n:", n, "idxSize:", indexSize, "TailBufLen:", len(tailBuf))
		return err
	}

	// Write index tail size at the end of stream.
	var tailSizeBuf = make([]byte, 8)
	binary.PutUvarint(tailSizeBuf, indexSize)
	sw.rawW.Write(tailBuf)

	cl, ok := sw.rawW.(io.Closer)
	if ok {
		return cl.Close()
	}

	return nil
}
<|endoftext|>


<commit_before>package mo

import (
	"runtime"
	"time"

	"github.com/shirou/gopsutil/mem"

	"runtime/pprof"

	n "github.com/shirou/gopsutil/net"
)

//Snapshot Snapshot
type Snapshot struct {
	CPUPercent        float64
	MemPercent        float64
	AverageNetIn      uint64
	AverageNetOut     uint64
	Thread            uint64
	Goroutine         uint64
	TotalConn         uint64
	TotalTopic        uint64
	TotalConsumer     uint64
	MaxConsumer       uint64
	TotalMessageIn    uint64
	AverageMessageIn  uint64
	TotalMessageOut   uint64
	AverageMessageOut uint64
}

//DataSource DataSource
type DataSource interface {
	GetTotalConsumer() map[string]uint64
	GetTotalConn() uint64
	GetTotalTopic() uint64
	GetMessageIn() uint64
	GetMessageOut() uint64
}

//Monitor Monitor
type Monitor struct {
	dataSource DataSource
}

//NewMonitor NewMonitor
func NewMonitor(dataSource DataSource) *Monitor {
	return &Monitor{dataSource: dataSource}
}

//Monitor Monitor
func (mo *Monitor) Monitor(duration time.Duration) *Snapshot {
	snap := Snapshot{}
	netIn1, netOut1 := mo.netMo()
	messageIn := mo.dataSource.GetMessageIn()
	messageOut := mo.dataSource.GetMessageOut()
	snap.CPUPercent = mo.cpuMo(duration)
	netIn2, netOut2 := mo.netMo()
	snap.MemPercent = mo.memMo()
	snap.AverageNetIn = uint64(float64(netIn2-netIn1) / duration.Seconds())
	snap.AverageNetOut = uint64(float64(netOut2-netOut1) / duration.Seconds())
	snap.TotalConn = mo.dataSource.GetTotalConn()
	snap.MaxConsumer, snap.TotalConsumer = mo.consumerMo()
	snap.TotalTopic = mo.dataSource.GetTotalTopic()
	snap.TotalMessageIn = mo.dataSource.GetMessageIn()
	snap.AverageMessageIn = uint64(float64(snap.TotalMessageIn-messageIn) / duration.Seconds())
	snap.TotalMessageOut = mo.dataSource.GetMessageOut()
	snap.AverageMessageOut = uint64(float64(snap.TotalMessageOut-messageOut) / duration.Seconds())
	snap.Goroutine = uint64(runtime.NumGoroutine())
	snap.Thread = uint64(pprof.Lookup("threadcreate").Count())
	return &snap
}

func (mo *Monitor) cpuMo(duration time.Duration) (percent float64) {
	// v, _ := cpu.Percent(duration, false)
	// return v[0]
	return 0
}

func (mo *Monitor) consumerMo() (maximum uint64, total uint64) {
	c := mo.dataSource.GetTotalConsumer()
	for _, consumer := range c {
		total += consumer
		if consumer > maximum {
			maximum = consumer
		}
	}
	return maximum, total
}

func (mo *Monitor) netMo() (in uint64, out uint64) {
	c, _ := n.IOCounters(false)
	if len(c) > 0 {
		return c[0].BytesRecv, c[0].BytesSent
	}
	return 0, 0
}

func (mo *Monitor) memMo() (usedPercent float64) {
	v, _ := mem.VirtualMemory()
	return v.UsedPercent
}
<commit_msg>no message<commit_after>package mo

import (
	"runtime"
	"time"

	"github.com/shirou/gopsutil/mem"

	"runtime/pprof"

	n "github.com/shirou/gopsutil/net"
)

//Snapshot Snapshot
type Snapshot struct {
	CPUPercent        float64
	MemPercent        float64
	AverageNetIn      uint64
	AverageNetOut     uint64
	Thread            uint64
	Goroutine         uint64
	TotalConn         uint64
	TotalTopic        uint64
	TotalConsumer     uint64
	MaxConsumer       uint64
	TotalMessageIn    uint64
	AverageMessageIn  uint64
	TotalMessageOut   uint64
	AverageMessageOut uint64
}

//DataSource DataSource
type DataSource interface {
	GetTotalConsumer() map[string]uint64
	GetTotalConn() uint64
	GetTotalTopic() uint64
	GetMessageIn() uint64
	GetMessageOut() uint64
}

//Monitor Monitor
type Monitor struct {
	dataSource DataSource
}

//NewMonitor NewMonitor
func NewMonitor(dataSource DataSource) *Monitor {
	return &Monitor{dataSource: dataSource}
}

//Monitor Monitor
func (mo *Monitor) Monitor(duration time.Duration) *Snapshot {
	snap := Snapshot{}
	netIn1, netOut1 := mo.netMo()
	messageIn := mo.dataSource.GetMessageIn()
	messageOut := mo.dataSource.GetMessageOut()
	snap.CPUPercent = mo.cpuMo(duration)
	netIn2, netOut2 := mo.netMo()
	snap.MemPercent = mo.memMo()
	snap.AverageNetIn = uint64(float64(netIn2-netIn1) / duration.Seconds())
	snap.AverageNetOut = uint64(float64(netOut2-netOut1) / duration.Seconds())
	snap.TotalConn = mo.dataSource.GetTotalConn()
	snap.MaxConsumer, snap.TotalConsumer = mo.consumerMo()
	snap.TotalTopic = mo.dataSource.GetTotalTopic()
	snap.TotalMessageIn = mo.dataSource.GetMessageIn()
	snap.AverageMessageIn = uint64(float64(snap.TotalMessageIn-messageIn) / duration.Seconds())
	snap.TotalMessageOut = mo.dataSource.GetMessageOut()
	snap.AverageMessageOut = uint64(float64(snap.TotalMessageOut-messageOut) / duration.Seconds())
	snap.Goroutine = uint64(runtime.NumGoroutine())
	snap.Thread = uint64(pprof.Lookup("threadcreate").Count())
	return &snap
}

func (mo *Monitor) cpuMo(duration time.Duration) (percent float64) {
	// v, _ := cpu.Percent(duration, false)
	// return v[0]
	<-time.After(duration)
	return 0
}

func (mo *Monitor) consumerMo() (maximum uint64, total uint64) {
	c := mo.dataSource.GetTotalConsumer()
	for _, consumer := range c {
		total += consumer
		if consumer > maximum {
			maximum = consumer
		}
	}
	return maximum, total
}

func (mo *Monitor) netMo() (in uint64, out uint64) {
	c, _ := n.IOCounters(false)
	if len(c) > 0 {
		return c[0].BytesRecv, c[0].BytesSent
	}
	return 0, 0
}

func (mo *Monitor) memMo() (usedPercent float64) {
	v, _ := mem.VirtualMemory()
	return v.UsedPercent
}
<|endoftext|>