<commit_before>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"text/tabwriter"
	"time"

	"github.com/influxdb/influxdb/cmd/influx_tsm/b1"
	"github.com/influxdb/influxdb/cmd/influx_tsm/bz1"
	"github.com/influxdb/influxdb/cmd/influx_tsm/tsdb"
)

type ShardReader interface {
	KeyIterator
	Open() error
	Close() error
}

const (
	backupExt = "bak"
	tsmExt    = "tsm"
)

var description = fmt.Sprintf(`
Convert a database from b1 or bz1 format to tsm1 format.

This tool will make backup any directory before conversion. It
is up to the end-user to delete the backup on the disk. Backups are
named by suffixing the database name with '.%s'. The backups will
be ignored by the system since they are not registered with the cluster.

To restore a backup, delete the tsm version, rename the backup and
restart the node.`, backupExt)

var dataPath string
var ds string
var tsmSz uint64
var parallel bool
var disBack bool

const maxTSMSz = 1 * 1024 * 1024 * 1024

func init() {
	flag.StringVar(&ds, "dbs", "", "Comma-delimited list of databases to convert. Default is to convert all")
	flag.Uint64Var(&tsmSz, "sz", maxTSMSz, "Maximum size of individual TSM files.")
	flag.BoolVar(&parallel, "parallel", false, "Perform parallel conversion.")
	flag.BoolVar(&disBack, "nobackup", false, "Disable database backups. Not recommended.")
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [options] <data-path> \n", os.Args[0])
		fmt.Fprintf(os.Stderr, "%s\n\n", description)
		flag.PrintDefaults()
		fmt.Fprintf(os.Stderr, "\n")
	}
}

func main() {
	pg := NewParallelGroup(1)

	flag.Parse()
	if len(flag.Args()) < 1 {
		fmt.Fprintf(os.Stderr, "No data directory specified\n")
		os.Exit(1)
	}
	dataPath = flag.Args()[0]

	if tsmSz > maxTSMSz {
		fmt.Fprintf(os.Stderr, "Maximum TSM file size is %d\n", maxTSMSz)
		os.Exit(1)
	}

	// Check if specific directories were requested.
	reqDs := strings.Split(ds, ",")
	if len(reqDs) == 1 && reqDs[0] == "" {
		reqDs = nil
	}

	// Determine the list of databases
	dbs, err := ioutil.ReadDir(dataPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "failed to access data directory at %s: %s\n", dataPath, err.Error())
		os.Exit(1)
	}
	fmt.Println() // Cleanly separate output from start of program.

	// Dump summary of what is about to happen.
	fmt.Println("b1 and bz1 shard conversion.")
	fmt.Println("-----------------------------------")
	fmt.Println("Data directory is:       ", dataPath)
	fmt.Println("Databases specified:     ", allDBs(reqDs))
	fmt.Println("Database backups enabled:", yesno(!disBack))
	fmt.Println("Parallel mode enabled:   ", yesno(parallel))
	fmt.Println()

	// Get the list of shards for conversion.
	var shards []*tsdb.ShardInfo
	for _, db := range dbs {
		if strings.HasSuffix(db.Name(), backupExt) {
			fmt.Printf("Skipping %s as it looks like a backup.\n", db.Name())
			continue
		}

		d := tsdb.NewDatabase(filepath.Join(dataPath, db.Name()))
		shs, err := d.Shards()
		if err != nil {
			fmt.Fprintf(os.Stderr, "failed to access shards for database %s: %s\n", d.Name(), err.Error())
			os.Exit(1)
		}
		shards = append(shards, shs...)
	}
	sort.Sort(tsdb.ShardInfos(shards))
	usl := len(shards)
	shards = tsdb.ShardInfos(shards).FilterFormat(tsdb.TSM1).ExclusiveDatabases(reqDs)
	sl := len(shards)

	// Anything to convert?
	fmt.Printf("\n%d shard(s) detected, %d non-TSM shards detected.\n", usl, sl)
	if len(shards) == 0 {
		fmt.Printf("Nothing to do.\n")
		os.Exit(0)
	}

	// Display list of convertible shards.
	fmt.Println()
	w := new(tabwriter.Writer)
	w.Init(os.Stdout, 0, 8, 1, '\t', 0)
	fmt.Fprintln(w, "Database\tRetention\tPath\tEngine\tSize")
	for _, si := range shards {
		fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%d\n", si.Database, si.RetentionPolicy, si.FullPath(dataPath), si.FormatAsString(), si.Size)
	}
	w.Flush()

	// Get confirmation from user.
	fmt.Printf("\nThese shards will be converted. Proceed? y/N: ")
	liner := bufio.NewReader(os.Stdin)
	yn, err := liner.ReadString('\n')
	if err != nil {
		fmt.Fprintf(os.Stderr, "failed to read response: %s", err.Error())
		os.Exit(1)
	}
	yn = strings.TrimRight(strings.ToLower(yn), "\n")
	if yn != "y" {
		fmt.Println("Conversion aborted.")
		os.Exit(1)
	}
	fmt.Println("Conversion starting....")

	// Backup each directory.
	if !disBack {
		databases := tsdb.ShardInfos(shards).Databases()
		if parallel {
			pg = NewParallelGroup(len(databases))
		}
		for _, db := range databases {
			pg.Request()
			go func(db string) {
				defer pg.Release()

				start := time.Now()
				err := backupDatabase(filepath.Join(dataPath, db))
				if err != nil {
					fmt.Fprintf(os.Stderr, "Backup of database %s failed: %s\n", db, err.Error())
					os.Exit(1)
				}
				fmt.Printf("Database %s backed up (%v)\n", db, time.Now().Sub(start))
			}(db)
		}
		pg.Wait()
	} else {
		fmt.Println("Database backup disabled.")
	}

	// Convert each shard.
	if parallel {
		pg = NewParallelGroup(len(shards))
	}
	for _, si := range shards {
		pg.Request()
		go func(si *tsdb.ShardInfo) {
			defer pg.Release()

			start := time.Now()
			if err := convertShard(si); err != nil {
				fmt.Fprintf(os.Stderr, "Failed to convert %s: %s\n", si.FullPath(dataPath), err.Error())
				os.Exit(1)
			}
			fmt.Printf("Conversion of %s successful (%s)\n", si.FullPath(dataPath), time.Now().Sub(start))
		}(si)
	}
	pg.Wait()
}

// backupDatabase backs up the database at src.
func backupDatabase(src string) error {
	dest := filepath.Join(src + "." + backupExt)
	if _, err := os.Stat(dest); !os.IsNotExist(err) {
		return fmt.Errorf("backup of %s already exists", src)
	}
	return copyDir(dest, src)
}

// copyDir copies the directory at src to dest. If dest does not exist it
// will be created. It is up to the caller to ensure the paths don't overlap.
func copyDir(dest, src string) error {
	copyFile := func(path string, info os.FileInfo, err error) error {
		// Strip the src from the path and replace with dest.
		toPath := strings.Replace(path, src, dest, 1)

		// Copy it.
		if info.IsDir() {
			if err := os.MkdirAll(toPath, info.Mode()); err != nil {
				return err
			}
		} else {
			err := func() error {
				in, err := os.Open(path)
				if err != nil {
					return err
				}
				defer in.Close()

				out, err := os.OpenFile(toPath, os.O_CREATE|os.O_WRONLY, info.Mode())
				if err != nil {
					return err
				}
				defer out.Close()

				_, err = io.Copy(out, in)
				return err
			}()
			if err != nil {
				return err
			}
		}
		return nil
	}

	return filepath.Walk(src, copyFile)
}

// convertShard converts the shard in-place.
func convertShard(si *tsdb.ShardInfo) error {
	src := si.FullPath(dataPath)
	dst := fmt.Sprintf("%s.%s", src, tsmExt)

	var reader ShardReader
	switch si.Format {
	case tsdb.BZ1:
		reader = bz1.NewReader(src)
	case tsdb.B1:
		reader = b1.NewReader(src)
	default:
		return fmt.Errorf("Unsupported shard format: %s", si.FormatAsString())
	}
	defer reader.Close()

	// Open the shard, and create a converter.
	if err := reader.Open(); err != nil {
		return fmt.Errorf("Failed to open %s for conversion: %s", src, err.Error())
	}
	converter := NewConverter(dst, uint32(tsmSz))

	// Perform the conversion.
	if err := converter.Process(reader); err != nil {
		return fmt.Errorf("Conversion of %s failed: %s", src, err.Error())
	}

	// Delete source shard, and rename new tsm1 shard.
	if err := reader.Close(); err != nil {
		return fmt.Errorf("Conversion of %s failed due to close: %s", src, err.Error())
	}

	if err := os.RemoveAll(si.FullPath(dataPath)); err != nil {
		return fmt.Errorf("Deletion of %s failed: %s", src, err.Error())
	}
	if err := os.Rename(dst, src); err != nil {
		return fmt.Errorf("Rename of %s to %s failed: %s", dst, src, err.Error())
	}

	return nil
}

// ParallelGroup allows the maximum parrallelism of a set of operations to be controlled.
type ParallelGroup struct {
	c  chan struct{}
	wg sync.WaitGroup
}

// NewParallelGroup returns a group which allows n operations to run in parallel. A value of 0
// means no operations will ever run.
func NewParallelGroup(n int) *ParallelGroup {
	return &ParallelGroup{
		c: make(chan struct{}, n),
	}
}

// Request requests permission to start an operation. It will block unless and until
// the parallel requirements would not be violated.
func (p *ParallelGroup) Request() {
	p.wg.Add(1)
	p.c <- struct{}{}
}

// Release informs the group that a previoulsy requested operation has completed.
func (p *ParallelGroup) Release() {
	<-p.c
	p.wg.Done()
}

// Wait blocks until the ParallelGroup has no unreleased operations.
func (p *ParallelGroup) Wait() {
	p.wg.Wait()
}

// yesno returns "yes" for true, "no" for false.
func yesno(b bool) string {
	if b {
		return "yes"
	}
	return "no"
}

// allDBs returns "all" if all databases are requested for conversion.
func allDBs(dbs []string) string {
	if dbs == nil {
		return "all"
	}
	return fmt.Sprintf("%v", dbs)
}
<commit_msg>Correct typo in influx_tsm help<commit_after>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"text/tabwriter"
	"time"

	"github.com/influxdb/influxdb/cmd/influx_tsm/b1"
	"github.com/influxdb/influxdb/cmd/influx_tsm/bz1"
	"github.com/influxdb/influxdb/cmd/influx_tsm/tsdb"
)

type ShardReader interface {
	KeyIterator
	Open() error
	Close() error
}

const (
	backupExt = "bak"
	tsmExt    = "tsm"
)

var description = fmt.Sprintf(`
Convert a database from b1 or bz1 format to tsm1 format.

This tool will backup any directory before conversion. It is up to the
end-user to delete the backup on the disk. Backups are named by suffixing
the database name with '.%s'. The backups will be ignored by the system
since they are not registered with the cluster.

To restore a backup, delete the tsm version, rename the backup and
restart the node.`, backupExt)

var dataPath string
var ds string
var tsmSz uint64
var parallel bool
var disBack bool

const maxTSMSz = 1 * 1024 * 1024 * 1024

func init() {
	flag.StringVar(&ds, "dbs", "", "Comma-delimited list of databases to convert. Default is to convert all")
	flag.Uint64Var(&tsmSz, "sz", maxTSMSz, "Maximum size of individual TSM files.")
	flag.BoolVar(&parallel, "parallel", false, "Perform parallel conversion.")
	flag.BoolVar(&disBack, "nobackup", false, "Disable database backups. Not recommended.")
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [options] <data-path> \n", os.Args[0])
		fmt.Fprintf(os.Stderr, "%s\n\n", description)
		flag.PrintDefaults()
		fmt.Fprintf(os.Stderr, "\n")
	}
}

func main() {
	pg := NewParallelGroup(1)

	flag.Parse()
	if len(flag.Args()) < 1 {
		fmt.Fprintf(os.Stderr, "No data directory specified\n")
		os.Exit(1)
	}
	dataPath = flag.Args()[0]

	if tsmSz > maxTSMSz {
		fmt.Fprintf(os.Stderr, "Maximum TSM file size is %d\n", maxTSMSz)
		os.Exit(1)
	}

	// Check if specific directories were requested.
	reqDs := strings.Split(ds, ",")
	if len(reqDs) == 1 && reqDs[0] == "" {
		reqDs = nil
	}

	// Determine the list of databases
	dbs, err := ioutil.ReadDir(dataPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "failed to access data directory at %s: %s\n", dataPath, err.Error())
		os.Exit(1)
	}
	fmt.Println() // Cleanly separate output from start of program.

	// Dump summary of what is about to happen.
	fmt.Println("b1 and bz1 shard conversion.")
	fmt.Println("-----------------------------------")
	fmt.Println("Data directory is:       ", dataPath)
	fmt.Println("Databases specified:     ", allDBs(reqDs))
	fmt.Println("Database backups enabled:", yesno(!disBack))
	fmt.Println("Parallel mode enabled:   ", yesno(parallel))
	fmt.Println()

	// Get the list of shards for conversion.
	var shards []*tsdb.ShardInfo
	for _, db := range dbs {
		if strings.HasSuffix(db.Name(), backupExt) {
			fmt.Printf("Skipping %s as it looks like a backup.\n", db.Name())
			continue
		}

		d := tsdb.NewDatabase(filepath.Join(dataPath, db.Name()))
		shs, err := d.Shards()
		if err != nil {
			fmt.Fprintf(os.Stderr, "failed to access shards for database %s: %s\n", d.Name(), err.Error())
			os.Exit(1)
		}
		shards = append(shards, shs...)
	}
	sort.Sort(tsdb.ShardInfos(shards))
	usl := len(shards)
	shards = tsdb.ShardInfos(shards).FilterFormat(tsdb.TSM1).ExclusiveDatabases(reqDs)
	sl := len(shards)

	// Anything to convert?
	fmt.Printf("\n%d shard(s) detected, %d non-TSM shards detected.\n", usl, sl)
	if len(shards) == 0 {
		fmt.Printf("Nothing to do.\n")
		os.Exit(0)
	}

	// Display list of convertible shards.
	fmt.Println()
	w := new(tabwriter.Writer)
	w.Init(os.Stdout, 0, 8, 1, '\t', 0)
	fmt.Fprintln(w, "Database\tRetention\tPath\tEngine\tSize")
	for _, si := range shards {
		fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%d\n", si.Database, si.RetentionPolicy, si.FullPath(dataPath), si.FormatAsString(), si.Size)
	}
	w.Flush()

	// Get confirmation from user.
	fmt.Printf("\nThese shards will be converted. Proceed? y/N: ")
	liner := bufio.NewReader(os.Stdin)
	yn, err := liner.ReadString('\n')
	if err != nil {
		fmt.Fprintf(os.Stderr, "failed to read response: %s", err.Error())
		os.Exit(1)
	}
	yn = strings.TrimRight(strings.ToLower(yn), "\n")
	if yn != "y" {
		fmt.Println("Conversion aborted.")
		os.Exit(1)
	}
	fmt.Println("Conversion starting....")

	// Backup each directory.
	if !disBack {
		databases := tsdb.ShardInfos(shards).Databases()
		if parallel {
			pg = NewParallelGroup(len(databases))
		}
		for _, db := range databases {
			pg.Request()
			go func(db string) {
				defer pg.Release()

				start := time.Now()
				err := backupDatabase(filepath.Join(dataPath, db))
				if err != nil {
					fmt.Fprintf(os.Stderr, "Backup of database %s failed: %s\n", db, err.Error())
					os.Exit(1)
				}
				fmt.Printf("Database %s backed up (%v)\n", db, time.Now().Sub(start))
			}(db)
		}
		pg.Wait()
	} else {
		fmt.Println("Database backup disabled.")
	}

	// Convert each shard.
	if parallel {
		pg = NewParallelGroup(len(shards))
	}
	for _, si := range shards {
		pg.Request()
		go func(si *tsdb.ShardInfo) {
			defer pg.Release()

			start := time.Now()
			if err := convertShard(si); err != nil {
				fmt.Fprintf(os.Stderr, "Failed to convert %s: %s\n", si.FullPath(dataPath), err.Error())
				os.Exit(1)
			}
			fmt.Printf("Conversion of %s successful (%s)\n", si.FullPath(dataPath), time.Now().Sub(start))
		}(si)
	}
	pg.Wait()
}

// backupDatabase backs up the database at src.
func backupDatabase(src string) error {
	dest := filepath.Join(src + "." + backupExt)
	if _, err := os.Stat(dest); !os.IsNotExist(err) {
		return fmt.Errorf("backup of %s already exists", src)
	}
	return copyDir(dest, src)
}

// copyDir copies the directory at src to dest. If dest does not exist it
// will be created. It is up to the caller to ensure the paths don't overlap.
func copyDir(dest, src string) error {
	copyFile := func(path string, info os.FileInfo, err error) error {
		// Strip the src from the path and replace with dest.
		toPath := strings.Replace(path, src, dest, 1)

		// Copy it.
		if info.IsDir() {
			if err := os.MkdirAll(toPath, info.Mode()); err != nil {
				return err
			}
		} else {
			err := func() error {
				in, err := os.Open(path)
				if err != nil {
					return err
				}
				defer in.Close()

				out, err := os.OpenFile(toPath, os.O_CREATE|os.O_WRONLY, info.Mode())
				if err != nil {
					return err
				}
				defer out.Close()

				_, err = io.Copy(out, in)
				return err
			}()
			if err != nil {
				return err
			}
		}
		return nil
	}

	return filepath.Walk(src, copyFile)
}

// convertShard converts the shard in-place.
func convertShard(si *tsdb.ShardInfo) error {
	src := si.FullPath(dataPath)
	dst := fmt.Sprintf("%s.%s", src, tsmExt)

	var reader ShardReader
	switch si.Format {
	case tsdb.BZ1:
		reader = bz1.NewReader(src)
	case tsdb.B1:
		reader = b1.NewReader(src)
	default:
		return fmt.Errorf("Unsupported shard format: %s", si.FormatAsString())
	}
	defer reader.Close()

	// Open the shard, and create a converter.
	if err := reader.Open(); err != nil {
		return fmt.Errorf("Failed to open %s for conversion: %s", src, err.Error())
	}
	converter := NewConverter(dst, uint32(tsmSz))

	// Perform the conversion.
	if err := converter.Process(reader); err != nil {
		return fmt.Errorf("Conversion of %s failed: %s", src, err.Error())
	}

	// Delete source shard, and rename new tsm1 shard.
	if err := reader.Close(); err != nil {
		return fmt.Errorf("Conversion of %s failed due to close: %s", src, err.Error())
	}

	if err := os.RemoveAll(si.FullPath(dataPath)); err != nil {
		return fmt.Errorf("Deletion of %s failed: %s", src, err.Error())
	}
	if err := os.Rename(dst, src); err != nil {
		return fmt.Errorf("Rename of %s to %s failed: %s", dst, src, err.Error())
	}

	return nil
}

// ParallelGroup allows the maximum parrallelism of a set of operations to be controlled.
type ParallelGroup struct {
	c  chan struct{}
	wg sync.WaitGroup
}

// NewParallelGroup returns a group which allows n operations to run in parallel. A value of 0
// means no operations will ever run.
func NewParallelGroup(n int) *ParallelGroup {
	return &ParallelGroup{
		c: make(chan struct{}, n),
	}
}

// Request requests permission to start an operation. It will block unless and until
// the parallel requirements would not be violated.
func (p *ParallelGroup) Request() {
	p.wg.Add(1)
	p.c <- struct{}{}
}

// Release informs the group that a previoulsy requested operation has completed.
func (p *ParallelGroup) Release() {
	<-p.c
	p.wg.Done()
}

// Wait blocks until the ParallelGroup has no unreleased operations.
func (p *ParallelGroup) Wait() {
	p.wg.Wait()
}

// yesno returns "yes" for true, "no" for false.
func yesno(b bool) string {
	if b {
		return "yes"
	}
	return "no"
}

// allDBs returns "all" if all databases are requested for conversion.
func allDBs(dbs []string) string {
	if dbs == nil {
		return "all"
	}
	return fmt.Sprintf("%v", dbs)
}
<|endoftext|>


<commit_before>// +build linux darwin freebsd

package main

import (
	"github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/gonuts/flag"
	"github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/commander"

	core "github.com/jbenet/go-ipfs/core"
	ipns "github.com/jbenet/go-ipfs/fuse/ipns"
	mount "github.com/jbenet/go-ipfs/fuse/mount"
	rofs "github.com/jbenet/go-ipfs/fuse/readonly"
)

var cmdIpfsMount = &commander.Command{
	UsageLine: "mount",
	Short:     "Mount an ipfs read-only mountpoint.",
	Long: `ipfs mount <os-path> - Mount an ipfs read-only mountpoint.

    Mount ipfs at a read-only mountpoint on the OS. All ipfs objects
    will be accessible under that directory. Note that the root will
    not be listable, as it is virtual. Accessing known paths directly.

`,
	Run:  mountCmd,
	Flag: *flag.NewFlagSet("ipfs-mount", flag.ExitOnError),
}

func init() {
	cmdIpfsMount.Flag.String("f", "", "specify a mountpoint for ipfs")
	cmdIpfsMount.Flag.String("n", "", "specify a mountpoint for ipns")
}

func mountCmd(c *commander.Command, inp []string) error {
	if err := platformFuseChecks(); err != nil {
		return err
	}

	cc, err := setupCmdContext(c, true)
	if err != nil {
		return err
	}
	defer cc.daemon.Close()

	// update fsdir with flag.
	fsdir := cc.node.Config.Mounts.IPFS
	if val, ok := c.Flag.Lookup("f").Value.Get().(string); ok && val != "" {
		fsdir = val
	}

	// get default mount points
	nsdir := cc.node.Config.Mounts.IPNS
	if val, ok := c.Flag.Lookup("n").Value.Get().(string); ok && val != "" {
		nsdir = val
	}

	return doMount(cc.node, fsdir, nsdir)
}

func doMount(node *core.IpfsNode, fsdir, nsdir string) error {

	// this sync stuff is so that both can be mounted simultaneously.
	var fsmount mount.Mount
	var nsmount mount.Mount
	var err1 error
	var err2 error

	done := make(chan struct{})

	go func() {
		fsmount, err1 = rofs.Mount(node, fsdir)
		done <- struct{}{}
	}()

	go func() {
		nsmount, err2 = ipns.Mount(node, nsdir, fsdir)
		done <- struct{}{}
	}()

	<-done
	<-done

	if err1 != nil || err2 != nil {
		fsmount.Close()
		nsmount.Close()
		if err1 != nil {
			return err1
		} else {
			return err2
		}
	}

	// setup node state, so that it can be cancelled
	node.Mounts.Ipfs = fsmount
	node.Mounts.Ipns = nsmount
	return nil
}

var platformFuseChecks = func() error {
	return nil
}
<commit_msg>ipfs1/mount: text + unmounting<commit_after>// +build linux darwin freebsd

package main

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

	"github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/gonuts/flag"
	"github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/commander"

	core "github.com/jbenet/go-ipfs/core"
	ipns "github.com/jbenet/go-ipfs/fuse/ipns"
	mount "github.com/jbenet/go-ipfs/fuse/mount"
	rofs "github.com/jbenet/go-ipfs/fuse/readonly"
)

var cmdIpfsMount = &commander.Command{
	UsageLine: "mount",
	Short:     "Mount an ipfs read-only mountpoint.",
	Long: `ipfs mount <os-path> - Mount an ipfs read-only mountpoint.

    Mount ipfs at a read-only mountpoint on the OS. All ipfs objects
    will be accessible under that directory. Note that the root will
    not be listable, as it is virtual. Accessing known paths directly.

`,
	Run:  mountCmd,
	Flag: *flag.NewFlagSet("ipfs-mount", flag.ExitOnError),
}

func init() {
	cmdIpfsMount.Flag.String("f", "", "specify a mountpoint for ipfs")
	cmdIpfsMount.Flag.String("n", "", "specify a mountpoint for ipns")
}

func mountCmd(c *commander.Command, inp []string) error {
	if err := platformFuseChecks(); err != nil {
		return err
	}

	cc, err := setupCmdContext(c, true)
	if err != nil {
		return err
	}
	defer cc.daemon.Close()

	// update fsdir with flag.
	fsdir := cc.node.Config.Mounts.IPFS
	if val, ok := c.Flag.Lookup("f").Value.Get().(string); ok && val != "" {
		fsdir = val
	}

	// get default mount points
	nsdir := cc.node.Config.Mounts.IPNS
	if val, ok := c.Flag.Lookup("n").Value.Get().(string); ok && val != "" {
		nsdir = val
	}

	return doMount(cc.node, fsdir, nsdir)
}

func doMount(node *core.IpfsNode, fsdir, nsdir string) error {

	// this sync stuff is so that both can be mounted simultaneously.
	var fsmount mount.Mount
	var nsmount mount.Mount
	var err1 error
	var err2 error

	done := make(chan struct{})

	go func() {
		fsmount, err1 = rofs.Mount(node, fsdir)
		done <- struct{}{}
	}()

	go func() {
		nsmount, err2 = ipns.Mount(node, nsdir, fsdir)
		done <- struct{}{}
	}()

	<-done
	<-done

	if err1 != nil || err2 != nil {
		fsmount.Close()
		nsmount.Close()
		if err1 != nil {
			return err1
		} else {
			return err2
		}
	}

	fmt.Printf("mounted ipfs at %s\n", fsdir)
	fmt.Printf("mounted ipns at %s\n", nsdir)

	// setup node state, so that it can be cancelled
	node.Mounts.Ipfs = fsmount
	node.Mounts.Ipns = nsmount

	// wait until we kill
	sigc := make(chan os.Signal, 1)
	signal.Notify(sigc, syscall.SIGHUP, syscall.SIGINT,
		syscall.SIGTERM, syscall.SIGQUIT)
	<-sigc

	fmt.Println("unmounting...")
	node.Close()
	return nil
}

var platformFuseChecks = func() error {
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package main

import (
	"encoding/base64"
	"fmt"
	"net"

	"launchpad.net/gnuflag"
	"launchpad.net/goyaml"

	"launchpad.net/juju-core/agent"
	"launchpad.net/juju-core/agent/mongo"
	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/constraints"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/environs/config"
	"launchpad.net/juju-core/instance"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/state/api/params"
	"launchpad.net/juju-core/worker/peergrouper"
)

type BootstrapCommand struct {
	cmd.CommandBase
	AgentConf
	EnvConfig   map[string]interface{}
	Constraints constraints.Value
	Hardware    instance.HardwareCharacteristics
	InstanceId  string
}

// Info returns a decription of the command.
func (c *BootstrapCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "bootstrap-state",
		Purpose: "initialize juju state",
	}
}

func (c *BootstrapCommand) SetFlags(f *gnuflag.FlagSet) {
	c.AgentConf.AddFlags(f)
	yamlBase64Var(f, &c.EnvConfig, "env-config", "", "initial environment configuration (yaml, base64 encoded)")
	f.Var(constraints.ConstraintsValue{Target: &c.Constraints}, "constraints", "initial environment constraints (space-separated strings)")
	f.Var(&c.Hardware, "hardware", "hardware characteristics (space-separated strings)")
	f.StringVar(&c.InstanceId, "instance-id", "", "unique instance-id for bootstrap machine")
}

// Init initializes the command for running.
func (c *BootstrapCommand) Init(args []string) error {
	if len(c.EnvConfig) == 0 {
		return requiredError("env-config")
	}
	if c.InstanceId == "" {
		return requiredError("instance-id")
	}
	return c.AgentConf.CheckArgs(args)
}

// Run initializes state for an environment.
func (c *BootstrapCommand) Run(_ *cmd.Context) error {
	envCfg, err := config.New(config.NoDefaults, c.EnvConfig)
	if err != nil {
		return err
	}
	err = c.ReadConfig("machine-0")
	if err != nil {
		return err
	}
	agentConfig := c.CurrentConfig()

	// agent.Jobs is an optional field in the agent config, and was
	// introduced after 1.17.2. We default to allowing units on
	// machine-0 if missing.
	jobs := agentConfig.Jobs()
	if len(jobs) == 0 {
		jobs = []params.MachineJob{
			params.JobManageEnviron,
			params.JobHostUnits,
		}
	}

	// Get the bootstrap machine's addresses from the provider.
	env, err := environs.New(envCfg)
	if err != nil {
		return err
	}
	instanceId := instance.Id(c.InstanceId)
	instances, err := env.Instances([]instance.Id{instanceId})
	if err != nil {
		return err
	}
	addrs, err := instances[0].Addresses()
	if err != nil {
		return err
	}

	// Generate a shared secret for the Mongo replica set, and write it out.
	sharedSecret, err := mongo.GenerateSharedSecret()
	if err != nil {
		return err
	}
	info, ok := agentConfig.StateServingInfo()
	if !ok {
		return fmt.Errorf("bootstrap machine config has no state serving info")
	}
	info.SharedSecret = sharedSecret
	if err := c.ChangeConfig(func(agentConfig agent.ConfigSetter) {
		agentConfig.SetStateServingInfo(info)
	}); err != nil {
		return fmt.Errorf("cannot write agent config: %v", err)
	}
	agentConfig = c.CurrentConfig()

	if err := c.startMongo(addrs, agentConfig); err != nil {
		return err
	}

	logger.Debugf("started mongo")
	// Initialise state, and store any agent config (e.g. password) changes.
	var st *state.State
	err = nil
	writeErr := c.ChangeConfig(func(agentConfig agent.ConfigSetter) {
		st, _, err = agent.InitializeState(
			agentConfig,
			envCfg,
			agent.BootstrapMachineConfig{
				Addresses:       addrs,
				Constraints:     c.Constraints,
				Jobs:            jobs,
				InstanceId:      instanceId,
				Characteristics: c.Hardware,
				SharedSecret:    sharedSecret,
			},
			state.DefaultDialOpts(),
			environs.NewStatePolicy(),
		)
	})
	if writeErr != nil {
		return fmt.Errorf("cannot write initial configuration: %v", err)
	}
	if err != nil {
		return err
	}
	st.Close()
	return nil
}

func (c *BootstrapCommand) startMongo(addrs []instance.Address, agentConfig agent.Config) error {
	logger.Debugf("starting mongo")

	agentConfig := c.CurrentConfig()
	info, ok := agentConfig.StateInfo()
	if !ok {
		return fmt.Errorf("no stateinfo available")
	}
	dialInfo, err := state.DialInfo(info, state.DefaultDialOpts())
	if err != nil {
		return err
	}
	servingInfo, ok := agentConfig.StateServingInfo()
	if !ok {
		return fmt.Errorf("agent config has no state serving info")
	}
	// Use localhost to dial the mongo server, because it's running in
	// auth mode and will refuse to perform any operations unless
	// we dial that address.
	dialInfo.Addrs = []string{
		net.JoinHostPort("127.0.0.1", fmt.Sprint(servingInfo.StatePort)),
	}
	logger.Infof("calling ensureMongoServer")
	err = ensureMongoServer(
		agentConfig.DataDir(),
		agentConfig.Value(agent.Namespace),
		servingInfo)
	if err != nil {
		return err
	}

	peerAddr := mongo.SelectPeerAddress(addrs)
	if peerAddr == "" {
		return fmt.Errorf("no appropriate peer address found in %q", addrs)
	}
	peerHostPort := net.JoinHostPort(peerAddr, fmt.Sprint(servingInfo.StatePort))

	return maybeInitiateMongoServer(peergrouper.InitiateMongoParams{
		DialInfo:       dialInfo,
		MemberHostPort: peerHostPort,
	})
}

// yamlBase64Value implements gnuflag.Value on a map[string]interface{}.
type yamlBase64Value map[string]interface{}

// Set decodes the base64 value into yaml then expands that into a map.
func (v *yamlBase64Value) Set(value string) error {
	decoded, err := base64.StdEncoding.DecodeString(value)
	if err != nil {
		return err
	}
	return goyaml.Unmarshal(decoded, v)
}

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

// yamlBase64Var sets up a gnuflag flag analogous to the FlagSet.*Var methods.
func yamlBase64Var(fs *gnuflag.FlagSet, target *map[string]interface{}, name string, value string, usage string) {
	fs.Var((*yamlBase64Value)(target), name, usage)
}
<commit_msg>cmd/jujud: lose spurious CurrentConfig call<commit_after>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package main

import (
	"encoding/base64"
	"fmt"
	"net"

	"launchpad.net/gnuflag"
	"launchpad.net/goyaml"

	"launchpad.net/juju-core/agent"
	"launchpad.net/juju-core/agent/mongo"
	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/constraints"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/environs/config"
	"launchpad.net/juju-core/instance"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/state/api/params"
	"launchpad.net/juju-core/worker/peergrouper"
)

type BootstrapCommand struct {
	cmd.CommandBase
	AgentConf
	EnvConfig   map[string]interface{}
	Constraints constraints.Value
	Hardware    instance.HardwareCharacteristics
	InstanceId  string
}

// Info returns a decription of the command.
func (c *BootstrapCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "bootstrap-state",
		Purpose: "initialize juju state",
	}
}

func (c *BootstrapCommand) SetFlags(f *gnuflag.FlagSet) {
	c.AgentConf.AddFlags(f)
	yamlBase64Var(f, &c.EnvConfig, "env-config", "", "initial environment configuration (yaml, base64 encoded)")
	f.Var(constraints.ConstraintsValue{Target: &c.Constraints}, "constraints", "initial environment constraints (space-separated strings)")
	f.Var(&c.Hardware, "hardware", "hardware characteristics (space-separated strings)")
	f.StringVar(&c.InstanceId, "instance-id", "", "unique instance-id for bootstrap machine")
}

// Init initializes the command for running.
func (c *BootstrapCommand) Init(args []string) error {
	if len(c.EnvConfig) == 0 {
		return requiredError("env-config")
	}
	if c.InstanceId == "" {
		return requiredError("instance-id")
	}
	return c.AgentConf.CheckArgs(args)
}

// Run initializes state for an environment.
func (c *BootstrapCommand) Run(_ *cmd.Context) error {
	envCfg, err := config.New(config.NoDefaults, c.EnvConfig)
	if err != nil {
		return err
	}
	err = c.ReadConfig("machine-0")
	if err != nil {
		return err
	}
	agentConfig := c.CurrentConfig()

	// agent.Jobs is an optional field in the agent config, and was
	// introduced after 1.17.2. We default to allowing units on
	// machine-0 if missing.
	jobs := agentConfig.Jobs()
	if len(jobs) == 0 {
		jobs = []params.MachineJob{
			params.JobManageEnviron,
			params.JobHostUnits,
		}
	}

	// Get the bootstrap machine's addresses from the provider.
	env, err := environs.New(envCfg)
	if err != nil {
		return err
	}
	instanceId := instance.Id(c.InstanceId)
	instances, err := env.Instances([]instance.Id{instanceId})
	if err != nil {
		return err
	}
	addrs, err := instances[0].Addresses()
	if err != nil {
		return err
	}

	// Generate a shared secret for the Mongo replica set, and write it out.
	sharedSecret, err := mongo.GenerateSharedSecret()
	if err != nil {
		return err
	}
	info, ok := agentConfig.StateServingInfo()
	if !ok {
		return fmt.Errorf("bootstrap machine config has no state serving info")
	}
	info.SharedSecret = sharedSecret
	if err := c.ChangeConfig(func(agentConfig agent.ConfigSetter) {
		agentConfig.SetStateServingInfo(info)
	}); err != nil {
		return fmt.Errorf("cannot write agent config: %v", err)
	}
	agentConfig = c.CurrentConfig()

	if err := c.startMongo(addrs, agentConfig); err != nil {
		return err
	}

	logger.Debugf("started mongo")
	// Initialise state, and store any agent config (e.g. password) changes.
	var st *state.State
	err = nil
	writeErr := c.ChangeConfig(func(agentConfig agent.ConfigSetter) {
		st, _, err = agent.InitializeState(
			agentConfig,
			envCfg,
			agent.BootstrapMachineConfig{
				Addresses:       addrs,
				Constraints:     c.Constraints,
				Jobs:            jobs,
				InstanceId:      instanceId,
				Characteristics: c.Hardware,
				SharedSecret:    sharedSecret,
			},
			state.DefaultDialOpts(),
			environs.NewStatePolicy(),
		)
	})
	if writeErr != nil {
		return fmt.Errorf("cannot write initial configuration: %v", err)
	}
	if err != nil {
		return err
	}
	st.Close()
	return nil
}

func (c *BootstrapCommand) startMongo(addrs []instance.Address, agentConfig agent.Config) error {
	logger.Debugf("starting mongo")
	info, ok := agentConfig.StateInfo()
	if !ok {
		return fmt.Errorf("no stateinfo available")
	}
	dialInfo, err := state.DialInfo(info, state.DefaultDialOpts())
	if err != nil {
		return err
	}
	servingInfo, ok := agentConfig.StateServingInfo()
	if !ok {
		return fmt.Errorf("agent config has no state serving info")
	}
	// Use localhost to dial the mongo server, because it's running in
	// auth mode and will refuse to perform any operations unless
	// we dial that address.
	dialInfo.Addrs = []string{
		net.JoinHostPort("127.0.0.1", fmt.Sprint(servingInfo.StatePort)),
	}
	logger.Infof("calling ensureMongoServer")
	err = ensureMongoServer(
		agentConfig.DataDir(),
		agentConfig.Value(agent.Namespace),
		servingInfo)
	if err != nil {
		return err
	}

	peerAddr := mongo.SelectPeerAddress(addrs)
	if peerAddr == "" {
		return fmt.Errorf("no appropriate peer address found in %q", addrs)
	}
	peerHostPort := net.JoinHostPort(peerAddr, fmt.Sprint(servingInfo.StatePort))

	return maybeInitiateMongoServer(peergrouper.InitiateMongoParams{
		DialInfo:       dialInfo,
		MemberHostPort: peerHostPort,
	})
}

// yamlBase64Value implements gnuflag.Value on a map[string]interface{}.
type yamlBase64Value map[string]interface{}

// Set decodes the base64 value into yaml then expands that into a map.
func (v *yamlBase64Value) Set(value string) error {
	decoded, err := base64.StdEncoding.DecodeString(value)
	if err != nil {
		return err
	}
	return goyaml.Unmarshal(decoded, v)
}

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

// yamlBase64Var sets up a gnuflag flag analogous to the FlagSet.*Var methods.
func yamlBase64Var(fs *gnuflag.FlagSet, target *map[string]interface{}, name string, value string, usage string) {
	fs.Var((*yamlBase64Value)(target), name, usage)
}
<|endoftext|>


<commit_before>package gbdt

import (
	"fmt"
	"log"
	"math"
)

var _ = fmt.Println

const (
	LOG_LIKEHOOD = 1
	LEAST_SQUARE = 2
)

func FxGradient(y int, predict float32) float32 {
	switch Conf.Losstype {
	case LEAST_SQUARE:
		return float32(y) - predict
	case LOG_LIKEHOOD:
		return 2 * float32(y) / float32((1 + math.Exp(float64(2*float32(y)*predict))))
	}
	return 2 * float32(y) / float32((1 + math.Exp(float64(2*float32(y)*predict))))

}

func NodePredictValue(d *DataSet, sample_sequence []int) float32 {
	switch Conf.Losstype {
	case LEAST_SQUARE:
		return LsOptimalValue(d, sample_sequence)
	case LOG_LIKEHOOD:
		return LogitOptimalValue(d, sample_sequence)
	}

	return LogitOptimalValue(d, sample_sequence)
}

func LogitOptimalValue(d *DataSet, sample_sequence []int) float32 {
	var val1 float32 = 0
	var val2 float32 = 0
	if len(sample_sequence) > len(d.Samples) || len(sample_sequence) == 0 {
		return 0.0
	}
	for _, index := range sample_sequence {
		val1 += d.Samples[index].target
		absy := float32(math.Abs(float64(d.Samples[index].target)))
		val2 += absy / (2 - absy)
	}
	return val1 / val2
}

func LsOptimalValue(d *DataSet, sample_sequence []int) float32 {
	var val1 float32 = 0
	var val2 float32 = 0
	if len(sample_sequence) > len(d.Samples) || len(sample_sequence) == 0 {
		return 0.0
	}
	for _, index := range sample_sequence {
		val1 += d.Samples[index].target * d.Samples[index].weight
		val2 += d.Samples[index].weight
	}
	return val1 / val2
}

func LogitCtr(f float32) float32 {
	return 1.0 / (1.0 + float32(math.Exp(-2.0*float64(f))))
}

func SameTarget(d *DataSet, sample_sequence []int) bool {
	if len(sample_sequence) == 0 || len(sample_sequence) > len(d.Samples) {
		fmt.Println(len(sample_sequence))
		log.Fatal("out of index sample_sequence")
	}
	var val float32 = d.Samples[sample_sequence[0]].target
	for i := 1; i < len(sample_sequence); i++ {
		if !Float32Equal(val, d.Samples[sample_sequence[i]].target) {
			return false
		}
	}
	return true
}

func CalculateVariance(d *DataSet, sample_sequence []int) (variance float32) {
	var s, ss, total_weight float64 = 0.0, 0.0, 0.0
	for _, k := range sample_sequence {
		s += float64(d.Samples[k].target * d.Samples[k].weight)
		ss += float64(d.Samples[k].target * d.Samples[k].target * d.Samples[k].weight)
		total_weight += float64(d.Samples[k].weight)
	}
	if total_weight > 1 {
		//variance = float32(ss/total_weight - s*s/total_weight/total_weight)
		variance = float32(ss - s*s/total_weight)
	} else {
		variance = 0
	}
	return
}
<commit_msg>fix bugs for node pred<commit_after>package gbdt

import (
	"fmt"
	"log"
	"math"
)

var _ = fmt.Println

const (
	LOG_LIKEHOOD = 1
	LEAST_SQUARE = 2
)

func FxGradient(y int, predict float32) float32 {
	switch Conf.Losstype {
	case LEAST_SQUARE:
		return float32(y) - predict
	case LOG_LIKEHOOD:
		return 2 * float32(y) / float32((1 + math.Exp(float64(2*float32(y)*predict))))
	}
	return 2 * float32(y) / float32((1 + math.Exp(float64(2*float32(y)*predict))))

}

func NodePredictValue(d *DataSet, sample_sequence []int) float32 {
	switch Conf.Losstype {
	case LEAST_SQUARE:
		return LsOptimalValue(d, sample_sequence)
	case LOG_LIKEHOOD:
		return LogitOptimalValue(d, sample_sequence)
	}

	return LogitOptimalValue(d, sample_sequence)
}

func LogitOptimalValue(d *DataSet, sample_sequence []int) float32 {
	var val1 float32 = 0
	var val2 float32 = 0
	if len(sample_sequence) > len(d.Samples) || len(sample_sequence) == 0 {
		return 0.0
	}
	for _, index := range sample_sequence {
		val1 += d.Samples[index].target
		absy := float32(math.Abs(float64(d.Samples[index].target)))
		val2 += absy * (2 - absy)
	}
	return val1 / val2
}

func LsOptimalValue(d *DataSet, sample_sequence []int) float32 {
	var val1 float32 = 0
	var val2 float32 = 0
	if len(sample_sequence) > len(d.Samples) || len(sample_sequence) == 0 {
		return 0.0
	}
	for _, index := range sample_sequence {
		val1 += d.Samples[index].target * d.Samples[index].weight
		val2 += d.Samples[index].weight
	}
	return val1 / val2
}

func LogitCtr(f float32) float32 {
	return 1.0 / (1.0 + float32(math.Exp(-2.0*float64(f))))
}

func SameTarget(d *DataSet, sample_sequence []int) bool {
	if len(sample_sequence) == 0 || len(sample_sequence) > len(d.Samples) {
		fmt.Println(len(sample_sequence))
		log.Fatal("out of index sample_sequence")
	}
	var val float32 = d.Samples[sample_sequence[0]].target
	for i := 1; i < len(sample_sequence); i++ {
		if !Float32Equal(val, d.Samples[sample_sequence[i]].target) {
			return false
		}
	}
	return true
}

func CalculateVariance(d *DataSet, sample_sequence []int) (variance float32) {
	var s, ss, total_weight float64 = 0.0, 0.0, 0.0
	for _, k := range sample_sequence {
		s += float64(d.Samples[k].target * d.Samples[k].weight)
		ss += float64(d.Samples[k].target * d.Samples[k].target * d.Samples[k].weight)
		total_weight += float64(d.Samples[k].weight)
	}
	if total_weight > 1 {
		//variance = float32(ss/total_weight - s*s/total_weight/total_weight)
		variance = float32(ss - s*s/total_weight)
	} else {
		variance = 0
	}
	return
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"os"
	"os/signal"
	"sync"
	"time"

	"github.com/segmentio/conf"
	nsq "github.com/segmentio/nsq-go"
)

func main() {
	config := struct {
		LookupdHttpAddr        []string `conf:"lookupd-http-addr"         help:"List of nsqlookupd servers"`
		DestinationNsqdTcpAddr []string `conf:"desintation-nsqd-tcp-addr" help:"List of nsqd nodes to publish to"`
		NsqdTcpAddr            string   `conf:"nsqd-tcp-addr"             help:"Addres of the nsqd nodes to consume from"`
		Topic                  string   `conf:"topic"                     help:"Topic to consume messages from"`
		Channel                string   `conf:"channel"                   help:"Channel to consume messages from"`
		DestinationTopic       string   `conf:"destination-topic"         help:"Topic to publish to"`
		RateLimit              int      `conf:"rate-limit"                help:"Maximum number of message per second processed"`
		MaxInFlight            int      `conf:"max-in-flight" help:"Maximum number of in-flight messages"`
	}{
		LookupdHttpAddr:        []string{"localhost:4161"},
		DestinationNsqdTcpAddr: []string{"localhost:4150"},
		NsqdTcpAddr:            "localhost:4150",
		MaxInFlight:            10,
	}

	conf.Load(&config)

	if len(config.Topic) == 0 {
		log.Fatal("error: missing topic")
	}

	if len(config.Channel) == 0 {
		log.Fatal("error: missing channel")
	}

	if len(config.DestinationTopic) == 0 {
		log.Fatal("error: missing destination channel")
	}

	consumer, err := nsq.StartConsumer(nsq.ConsumerConfig{
		Topic:       config.Topic,
		Channel:     config.Channel,
		Lookup:      config.LookupdHttpAddr,
		Address:     config.NsqdTcpAddr,
		MaxInFlight: config.MaxInFlight,
	})
	if err != nil {
		log.Fatal(err)
	}

	join := &sync.WaitGroup{}

	for _, addr := range config.DestinationNsqdTcpAddr {
		producer, err := nsq.StartProducer(nsq.ProducerConfig{
			Topic:   config.DestinationTopic,
			Address: addr,
		})
		if err != nil {
			log.Fatal(err)
		}
		join.Add(1)
		go forward(producer, rateLimit(config.RateLimit, consumer.Messages()), join)
	}

	sigchan := make(chan os.Signal)
	signal.Notify(sigchan, os.Interrupt)
	<-sigchan

	consumer.Stop()
	join.Wait()
}

func forward(producer *nsq.Producer, messages <-chan nsq.Message, join *sync.WaitGroup) {
	defer join.Done()

	for msg := range messages {
		if err := producer.Publish(msg.Body); err != nil {
			msg.Requeue(15 * time.Second)
			log.Print(err)
		} else {
			msg.Finish()
		}
	}
}

func rateLimit(limit int, messages <-chan nsq.Message) <-chan nsq.Message {
	if limit <= 0 {
		return messages
	}

	output := make(chan nsq.Message, cap(messages))

	go func() {
		ticker := time.NewTicker(1 * time.Second)
		defer ticker.Stop()

		input := messages
		count := 0

		for {
			select {
			case <-ticker.C:
				count = 0
				input = messages

			case msg := <-input:
				output <- msg

				if count++; count >= limit {
					input = nil
				}
			}
		}
	}()

	return output
}
<commit_msg>add more logging<commit_after>package main

import (
	"log"
	"os"
	"os/signal"
	"sync"
	"time"

	"github.com/segmentio/conf"
	nsq "github.com/segmentio/nsq-go"
)

func main() {
	config := struct {
		LookupdHttpAddr        []string `conf:"lookupd-http-addr"         help:"List of nsqlookupd servers"`
		DestinationNsqdTcpAddr []string `conf:"desintation-nsqd-tcp-addr" help:"List of nsqd nodes to publish to"`
		NsqdTcpAddr            string   `conf:"nsqd-tcp-addr"             help:"Addres of the nsqd nodes to consume from"`
		Topic                  string   `conf:"topic"                     help:"Topic to consume messages from"`
		Channel                string   `conf:"channel"                   help:"Channel to consume messages from"`
		DestinationTopic       string   `conf:"destination-topic"         help:"Topic to publish to"`
		RateLimit              int      `conf:"rate-limit"                help:"Maximum number of message per second processed"`
		MaxInFlight            int      `conf:"max-in-flight" help:"Maximum number of in-flight messages"`
	}{
		LookupdHttpAddr:        []string{"localhost:4161"},
		DestinationNsqdTcpAddr: []string{"localhost:4150"},
		NsqdTcpAddr:            "localhost:4150",
		MaxInFlight:            10,
	}

	conf.Load(&config)

	if len(config.Topic) == 0 {
		log.Fatal("error: missing topic")
	}

	if len(config.Channel) == 0 {
		log.Fatal("error: missing channel")
	}

	if len(config.DestinationTopic) == 0 {
		log.Fatal("error: missing destination channel")
	}

	consumer, err := nsq.StartConsumer(nsq.ConsumerConfig{
		Topic:       config.Topic,
		Channel:     config.Channel,
		Lookup:      config.LookupdHttpAddr,
		Address:     config.NsqdTcpAddr,
		MaxInFlight: config.MaxInFlight,
	})
	if err != nil {
		log.Fatal(err)
	}

	join := &sync.WaitGroup{}

	for _, addr := range config.DestinationNsqdTcpAddr {
		producer, err := nsq.StartProducer(nsq.ProducerConfig{
			Topic:   config.DestinationTopic,
			Address: addr,
		})
		if err != nil {
			log.Fatal(err)
		}
		join.Add(1)
		go forward(producer, rateLimit(config.RateLimit, consumer.Messages()), join)
	}

	sigchan := make(chan os.Signal)
	signal.Notify(sigchan, os.Interrupt)
	<-sigchan

	consumer.Stop()
	join.Wait()
}

func forward(producer *nsq.Producer, messages <-chan nsq.Message, join *sync.WaitGroup) {
	defer join.Done()

	for msg := range messages {
		if err := producer.Publish(msg.Body); err != nil {
			msg.Requeue(15 * time.Second)
			log.Print(err)
		} else {
			msg.Finish()
		}
	}
}

func rateLimit(limit int, messages <-chan nsq.Message) <-chan nsq.Message {
	if limit <= 0 {
		return messages
	}

	output := make(chan nsq.Message, cap(messages))

	go func() {
		ticker := time.NewTicker(1 * time.Second)
		defer ticker.Stop()

		input := messages
		count := 0

		for {
			select {
			case <-ticker.C:
				log.Printf("publish rate of %d msg/sec", count)
				count = 0
				input = messages

			case msg := <-input:
				output <- msg

				if count++; count >= limit {
					input = nil
					log.Printf("rate limiting to %d msg/sec", limit)
				}
			}
		}
	}()

	return output
}
<|endoftext|>


<commit_before>package operate

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"

	"github.com/dchest/safefile"
	"github.com/go-errors/errors"
	"github.com/itchio/butler/buse"
	"github.com/itchio/butler/cmd/wipe"
	"github.com/itchio/butler/comm"
	"github.com/itchio/butler/mansion"
	"github.com/itchio/wharf/state"
	"github.com/mitchellh/mapstructure"
	"github.com/sourcegraph/jsonrpc2"
)

type OperationContext struct {
	conn        *jsonrpc2.Conn
	ctx         context.Context
	consumer    *state.Consumer
	stageFolder string
	logFile     *os.File

	mansionContext *mansion.Context

	root map[string]interface{}

	// keep track of what we've loaded so far
	// loading more than once is not ok
	loaded map[string]struct{}
}

func LoadContext(conn *jsonrpc2.Conn, mansionContext *mansion.Context, consumer *state.Consumer, stageFolder string) *OperationContext {
	err := os.MkdirAll(stageFolder, 0755)
	if err != nil {
		consumer.Warnf("Could not create operate directory: %s", err.Error())
	}

	logFilePath := filepath.Join(stageFolder, "operate-log.txt")
	logFile, err := os.OpenFile(logFilePath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0644)
	if err != nil {
		consumer.Warnf("Could not open operate log: %s", err.Error())
	}

	ctx := mansionContext.Context()

	subconsumer := &state.Consumer{
		OnMessage: func(level, msg string) {
			if logFile != nil {
				fmt.Fprintf(logFile, "[%s] %s\n", level, msg)
			}
			conn.Notify(ctx, "Log", &buse.LogNotification{
				Level:   level,
				Message: msg,
			})
		},
		OnProgress: func(percent float64) {
			conn.Notify(ctx, "Operation.Progress", &buse.OperationProgressNotification{
				Progress: percent,
			})
		},
		OnProgressLabel:  consumer.OnProgressLabel,
		OnPauseProgress:  consumer.OnPauseProgress,
		OnResumeProgress: consumer.OnResumeProgress,
	}

	oc := &OperationContext{
		consumer:       subconsumer,
		logFile:        logFile,
		stageFolder:    stageFolder,
		ctx:            ctx,
		mansionContext: mansionContext,
		root:           make(map[string]interface{}),
		loaded:         make(map[string]struct{}),
	}

	path := contextPath(stageFolder)

	f, err := os.Open(path)
	if err != nil {
		if os.IsNotExist(err) {
			// empty context, that's fine
		} else {
			oc.consumer.Warnf("While loading context from %s: %s", path, err.Error())
		}
		return oc
	}
	defer f.Close()

	err = json.NewDecoder(f).Decode(&oc.root)
	if err != nil {
		oc.consumer.Warnf("While decoding context from %s: %s", path, err.Error())
	}

	return oc
}

func (oc *OperationContext) Load(s Subcontext) {
	if _, ok := oc.loaded[s.Key()]; ok {
		oc.consumer.Warnf("Refusing to load subcontext %s a second time", s.Key())
		return
	}

	// only load if there's actually something there
	if val, ok := oc.root[s.Key()]; ok {
		dec, err := mapstructure.NewDecoder(&mapstructure.DecoderConfig{
			TagName: "json",
			Result:  s.Data(),
		})
		if err != nil {
			oc.consumer.Warnf("could not load subcontext %s: while configuring decoder, %s", s.Key(), err.Error())
			return
		}

		err = dec.Decode(val)
		if err != nil {
			oc.consumer.Warnf("could not load subcontext %s: while decoding, %s", s.Key(), err.Error())
			return
		}
	}

	oc.loaded[s.Key()] = struct{}{}
}

func (oc *OperationContext) Save(s Subcontext) error {
	oc.root[s.Key()] = s.Data()

	path := contextPath(oc.stageFolder)

	f, err := safefile.Create(path, 0644)
	if err != nil {
		return errors.Wrap(err, 0)
	}

	err = json.NewEncoder(f).Encode(&oc.root)
	if err != nil {
		return errors.Wrap(err, 0)
	}

	err = f.Commit()
	if err != nil {
		return errors.Wrap(err, 0)
	}

	defer f.Close()

	return nil
}

func (oc *OperationContext) Retire() error {
	consumer := oc.Consumer()

	consumer.Infof("Retiring stage folder...")
	err := oc.logFile.Close()
	if err != nil {
		return errors.Wrap(err, 0)
	}

	err = wipe.Do(oc.MansionContext(), comm.NewStateConsumer(), oc.StageFolder())
	if err != nil {
		return errors.Wrap(err, 0)
	}

	return nil
}

func (oc *OperationContext) MansionContext() *mansion.Context {
	return oc.mansionContext
}

func (oc *OperationContext) StageFolder() string {
	return oc.stageFolder
}

func (oc *OperationContext) Consumer() *state.Consumer {
	return oc.consumer
}

func contextPath(stageFolder string) string {
	return filepath.Join(stageFolder, "operate-context.json")
}

type Subcontext interface {
	// Key returns a unique string key used for storing
	// something under the context object
	Key() string

	// Data should return a pointer to the underlying struct
	// of the subcontext
	Data() interface{}
}
<commit_msg>operate context: Don't forget to set conn<commit_after>package operate

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"

	"github.com/dchest/safefile"
	"github.com/go-errors/errors"
	"github.com/itchio/butler/buse"
	"github.com/itchio/butler/cmd/wipe"
	"github.com/itchio/butler/comm"
	"github.com/itchio/butler/mansion"
	"github.com/itchio/wharf/state"
	"github.com/mitchellh/mapstructure"
	"github.com/sourcegraph/jsonrpc2"
)

type OperationContext struct {
	conn        *jsonrpc2.Conn
	ctx         context.Context
	consumer    *state.Consumer
	stageFolder string
	logFile     *os.File

	mansionContext *mansion.Context

	root map[string]interface{}

	// keep track of what we've loaded so far
	// loading more than once is not ok
	loaded map[string]struct{}
}

func LoadContext(conn *jsonrpc2.Conn, mansionContext *mansion.Context, consumer *state.Consumer, stageFolder string) *OperationContext {
	err := os.MkdirAll(stageFolder, 0755)
	if err != nil {
		consumer.Warnf("Could not create operate directory: %s", err.Error())
	}

	logFilePath := filepath.Join(stageFolder, "operate-log.txt")
	logFile, err := os.OpenFile(logFilePath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0644)
	if err != nil {
		consumer.Warnf("Could not open operate log: %s", err.Error())
	}

	ctx := mansionContext.Context()

	subconsumer := &state.Consumer{
		OnMessage: func(level, msg string) {
			if logFile != nil {
				fmt.Fprintf(logFile, "[%s] %s\n", level, msg)
			}
			conn.Notify(ctx, "Log", &buse.LogNotification{
				Level:   level,
				Message: msg,
			})
		},
		OnProgress: func(percent float64) {
			conn.Notify(ctx, "Operation.Progress", &buse.OperationProgressNotification{
				Progress: percent,
			})
		},
		OnProgressLabel:  consumer.OnProgressLabel,
		OnPauseProgress:  consumer.OnPauseProgress,
		OnResumeProgress: consumer.OnResumeProgress,
	}

	oc := &OperationContext{
		consumer:       subconsumer,
		logFile:        logFile,
		stageFolder:    stageFolder,
		ctx:            ctx,
		conn:           conn,
		mansionContext: mansionContext,
		root:           make(map[string]interface{}),
		loaded:         make(map[string]struct{}),
	}

	path := contextPath(stageFolder)

	f, err := os.Open(path)
	if err != nil {
		if os.IsNotExist(err) {
			// empty context, that's fine
		} else {
			oc.consumer.Warnf("While loading context from %s: %s", path, err.Error())
		}
		return oc
	}
	defer f.Close()

	err = json.NewDecoder(f).Decode(&oc.root)
	if err != nil {
		oc.consumer.Warnf("While decoding context from %s: %s", path, err.Error())
	}

	return oc
}

func (oc *OperationContext) Load(s Subcontext) {
	if _, ok := oc.loaded[s.Key()]; ok {
		oc.consumer.Warnf("Refusing to load subcontext %s a second time", s.Key())
		return
	}

	// only load if there's actually something there
	if val, ok := oc.root[s.Key()]; ok {
		dec, err := mapstructure.NewDecoder(&mapstructure.DecoderConfig{
			TagName: "json",
			Result:  s.Data(),
		})
		if err != nil {
			oc.consumer.Warnf("could not load subcontext %s: while configuring decoder, %s", s.Key(), err.Error())
			return
		}

		err = dec.Decode(val)
		if err != nil {
			oc.consumer.Warnf("could not load subcontext %s: while decoding, %s", s.Key(), err.Error())
			return
		}
	}

	oc.loaded[s.Key()] = struct{}{}
}

func (oc *OperationContext) Save(s Subcontext) error {
	oc.root[s.Key()] = s.Data()

	path := contextPath(oc.stageFolder)

	f, err := safefile.Create(path, 0644)
	if err != nil {
		return errors.Wrap(err, 0)
	}

	err = json.NewEncoder(f).Encode(&oc.root)
	if err != nil {
		return errors.Wrap(err, 0)
	}

	err = f.Commit()
	if err != nil {
		return errors.Wrap(err, 0)
	}

	defer f.Close()

	return nil
}

func (oc *OperationContext) Retire() error {
	consumer := oc.Consumer()

	consumer.Infof("Retiring stage folder...")
	err := oc.logFile.Close()
	if err != nil {
		return errors.Wrap(err, 0)
	}

	err = wipe.Do(oc.MansionContext(), comm.NewStateConsumer(), oc.StageFolder())
	if err != nil {
		return errors.Wrap(err, 0)
	}

	return nil
}

func (oc *OperationContext) MansionContext() *mansion.Context {
	return oc.mansionContext
}

func (oc *OperationContext) StageFolder() string {
	return oc.stageFolder
}

func (oc *OperationContext) Consumer() *state.Consumer {
	return oc.consumer
}

func contextPath(stageFolder string) string {
	return filepath.Join(stageFolder, "operate-context.json")
}

type Subcontext interface {
	// Key returns a unique string key used for storing
	// something under the context object
	Key() string

	// Data should return a pointer to the underlying struct
	// of the subcontext
	Data() interface{}
}
<|endoftext|>


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

// The main package for the Prometheus server executable.
package main

import (
	"flag"
	"fmt"
	_ "net/http/pprof" // Comment this line to disable pprof endpoint.
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/log"
	"github.com/prometheus/common/version"
	"github.com/prometheus/prometheus/config"
	"github.com/prometheus/prometheus/notifier"
	"github.com/prometheus/prometheus/promql"
	"github.com/prometheus/prometheus/retrieval"
	"github.com/prometheus/prometheus/rules"
	"github.com/prometheus/prometheus/storage"
	"github.com/prometheus/prometheus/storage/local"
	"github.com/prometheus/prometheus/storage/remote"
	"github.com/prometheus/prometheus/web"
)

func main() {
	os.Exit(Main())
}

var (
	configSuccess = prometheus.NewGauge(prometheus.GaugeOpts{
		Namespace: "prometheus",
		Name:      "config_last_reload_successful",
		Help:      "Whether the last configuration reload attempt was successful.",
	})
	configSuccessTime = prometheus.NewGauge(prometheus.GaugeOpts{
		Namespace: "prometheus",
		Name:      "config_last_reload_success_timestamp_seconds",
		Help:      "Timestamp of the last successful configuration reload.",
	})
)

func init() {
	prometheus.MustRegister(version.NewCollector("prometheus"))
}

// Main manages the startup and shutdown lifecycle of the entire Prometheus server.
func Main() int {
	if err := parse(os.Args[1:]); err != nil {
		log.Error(err)
		return 2
	}

	if cfg.printVersion {
		fmt.Fprintln(os.Stdout, version.Print("prometheus"))
		return 0
	}

	log.Infoln("Starting prometheus", version.Info())
	log.Infoln("Build context", version.BuildContext())

	var reloadables []Reloadable

	var (
		memStorage     = local.NewMemorySeriesStorage(&cfg.storage)
		remoteStorage  = remote.New(&cfg.remote)
		sampleAppender = storage.Fanout{memStorage}
	)
	if remoteStorage != nil {
		sampleAppender = append(sampleAppender, remoteStorage)
		reloadables = append(reloadables, remoteStorage)
	}

	var (
		notifier      = notifier.New(&cfg.notifier)
		targetManager = retrieval.NewTargetManager(sampleAppender)
		queryEngine   = promql.NewEngine(memStorage, &cfg.queryEngine)
	)

	ruleManager := rules.NewManager(&rules.ManagerOptions{
		SampleAppender: sampleAppender,
		Notifier:       notifier,
		QueryEngine:    queryEngine,
		ExternalURL:    cfg.web.ExternalURL,
	})

	flags := map[string]string{}
	cfg.fs.VisitAll(func(f *flag.Flag) {
		flags[f.Name] = f.Value.String()
	})

	version := &web.PrometheusVersion{
		Version:   version.Version,
		Revision:  version.Revision,
		Branch:    version.Branch,
		BuildUser: version.BuildUser,
		BuildDate: version.BuildDate,
		GoVersion: version.GoVersion,
	}

	webHandler := web.New(memStorage, queryEngine, targetManager, ruleManager, version, flags, &cfg.web)

	reloadables = append(reloadables, targetManager, ruleManager, webHandler, notifier)

	if err := reloadConfig(cfg.configFile, reloadables...); err != nil {
		log.Errorf("Error loading config: %s", err)
		return 1
	}

	// Wait for reload or termination signals. Start the handler for SIGHUP as
	// early as possible, but ignore it until we are ready to handle reloading
	// our config.
	hup := make(chan os.Signal)
	hupReady := make(chan bool)
	signal.Notify(hup, syscall.SIGHUP)
	go func() {
		<-hupReady
		for {
			select {
			case <-hup:
				if err := reloadConfig(cfg.configFile, reloadables...); err != nil {
					log.Errorf("Error reloading config: %s", err)
				}
			case rc := <-webHandler.Reload():
				if err := reloadConfig(cfg.configFile, reloadables...); err != nil {
					log.Errorf("Error reloading config: %s", err)
					rc <- err
				} else {
					rc <- nil
				}
			}
		}
	}()

	// Start all components. The order is NOT arbitrary.

	if err := memStorage.Start(); err != nil {
		log.Errorln("Error opening memory series storage:", err)
		return 1
	}
	defer func() {
		if err := memStorage.Stop(); err != nil {
			log.Errorln("Error stopping storage:", err)
		}
	}()

	if remoteStorage != nil {
		prometheus.MustRegister(remoteStorage)

		go remoteStorage.Run()
		defer remoteStorage.Stop()
	}
	// The storage has to be fully initialized before registering.
	prometheus.MustRegister(memStorage)
	prometheus.MustRegister(notifier)
	prometheus.MustRegister(configSuccess)
	prometheus.MustRegister(configSuccessTime)

	// The notifieris a dependency of the rule manager. It has to be
	// started before and torn down afterwards.
	go notifier.Run()
	defer notifier.Stop()

	go ruleManager.Run()
	defer ruleManager.Stop()

	go targetManager.Run()
	defer targetManager.Stop()

	// Shutting down the query engine before the rule manager will cause pending queries
	// to be canceled and ensures a quick shutdown of the rule manager.
	defer queryEngine.Stop()

	go webHandler.Run()

	// Wait for reload or termination signals.
	close(hupReady) // Unblock SIGHUP handler.

	term := make(chan os.Signal)
	signal.Notify(term, os.Interrupt, syscall.SIGTERM)
	select {
	case <-term:
		log.Warn("Received SIGTERM, exiting gracefully...")
	case <-webHandler.Quit():
		log.Warn("Received termination request via web service, exiting gracefully...")
	case err := <-webHandler.ListenError():
		log.Errorln("Error starting web server, exiting gracefully:", err)
	}

	log.Info("See you next time!")
	return 0
}

// Reloadable things can change their internal state to match a new config
// and handle failure gracefully.
type Reloadable interface {
	ApplyConfig(*config.Config) error
}

func reloadConfig(filename string, rls ...Reloadable) (err error) {
	log.Infof("Loading configuration file %s", filename)
	defer func() {
		if err == nil {
			configSuccess.Set(1)
			configSuccessTime.Set(float64(time.Now().Unix()))
		} else {
			configSuccess.Set(0)
		}
	}()

	conf, err := config.LoadFile(filename)
	if err != nil {
		return fmt.Errorf("couldn't load configuration (-config.file=%s): %v", filename, err)
	}

	// Apply all configs and return the first error if there were any.
	for _, rl := range rls {
		if err != nil {
			err = rl.ApplyConfig(conf)
		} else {
			rl.ApplyConfig(conf)
		}
	}
	return err
}
<commit_msg>Fix ApplyConfig() error handling<commit_after>// Copyright 2015 The Prometheus Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// The main package for the Prometheus server executable.
package main

import (
	"flag"
	"fmt"
	_ "net/http/pprof" // Comment this line to disable pprof endpoint.
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/log"
	"github.com/prometheus/common/version"
	"github.com/prometheus/prometheus/config"
	"github.com/prometheus/prometheus/notifier"
	"github.com/prometheus/prometheus/promql"
	"github.com/prometheus/prometheus/retrieval"
	"github.com/prometheus/prometheus/rules"
	"github.com/prometheus/prometheus/storage"
	"github.com/prometheus/prometheus/storage/local"
	"github.com/prometheus/prometheus/storage/remote"
	"github.com/prometheus/prometheus/web"
)

func main() {
	os.Exit(Main())
}

var (
	configSuccess = prometheus.NewGauge(prometheus.GaugeOpts{
		Namespace: "prometheus",
		Name:      "config_last_reload_successful",
		Help:      "Whether the last configuration reload attempt was successful.",
	})
	configSuccessTime = prometheus.NewGauge(prometheus.GaugeOpts{
		Namespace: "prometheus",
		Name:      "config_last_reload_success_timestamp_seconds",
		Help:      "Timestamp of the last successful configuration reload.",
	})
)

func init() {
	prometheus.MustRegister(version.NewCollector("prometheus"))
}

// Main manages the startup and shutdown lifecycle of the entire Prometheus server.
func Main() int {
	if err := parse(os.Args[1:]); err != nil {
		log.Error(err)
		return 2
	}

	if cfg.printVersion {
		fmt.Fprintln(os.Stdout, version.Print("prometheus"))
		return 0
	}

	log.Infoln("Starting prometheus", version.Info())
	log.Infoln("Build context", version.BuildContext())

	var reloadables []Reloadable

	var (
		memStorage     = local.NewMemorySeriesStorage(&cfg.storage)
		remoteStorage  = remote.New(&cfg.remote)
		sampleAppender = storage.Fanout{memStorage}
	)
	if remoteStorage != nil {
		sampleAppender = append(sampleAppender, remoteStorage)
		reloadables = append(reloadables, remoteStorage)
	}

	var (
		notifier      = notifier.New(&cfg.notifier)
		targetManager = retrieval.NewTargetManager(sampleAppender)
		queryEngine   = promql.NewEngine(memStorage, &cfg.queryEngine)
	)

	ruleManager := rules.NewManager(&rules.ManagerOptions{
		SampleAppender: sampleAppender,
		Notifier:       notifier,
		QueryEngine:    queryEngine,
		ExternalURL:    cfg.web.ExternalURL,
	})

	flags := map[string]string{}
	cfg.fs.VisitAll(func(f *flag.Flag) {
		flags[f.Name] = f.Value.String()
	})

	version := &web.PrometheusVersion{
		Version:   version.Version,
		Revision:  version.Revision,
		Branch:    version.Branch,
		BuildUser: version.BuildUser,
		BuildDate: version.BuildDate,
		GoVersion: version.GoVersion,
	}

	webHandler := web.New(memStorage, queryEngine, targetManager, ruleManager, version, flags, &cfg.web)

	reloadables = append(reloadables, targetManager, ruleManager, webHandler, notifier)

	if err := reloadConfig(cfg.configFile, reloadables...); err != nil {
		log.Errorf("Error loading config: %s", err)
		return 1
	}

	// Wait for reload or termination signals. Start the handler for SIGHUP as
	// early as possible, but ignore it until we are ready to handle reloading
	// our config.
	hup := make(chan os.Signal)
	hupReady := make(chan bool)
	signal.Notify(hup, syscall.SIGHUP)
	go func() {
		<-hupReady
		for {
			select {
			case <-hup:
				if err := reloadConfig(cfg.configFile, reloadables...); err != nil {
					log.Errorf("Error reloading config: %s", err)
				}
			case rc := <-webHandler.Reload():
				if err := reloadConfig(cfg.configFile, reloadables...); err != nil {
					log.Errorf("Error reloading config: %s", err)
					rc <- err
				} else {
					rc <- nil
				}
			}
		}
	}()

	// Start all components. The order is NOT arbitrary.

	if err := memStorage.Start(); err != nil {
		log.Errorln("Error opening memory series storage:", err)
		return 1
	}
	defer func() {
		if err := memStorage.Stop(); err != nil {
			log.Errorln("Error stopping storage:", err)
		}
	}()

	if remoteStorage != nil {
		prometheus.MustRegister(remoteStorage)

		go remoteStorage.Run()
		defer remoteStorage.Stop()
	}
	// The storage has to be fully initialized before registering.
	prometheus.MustRegister(memStorage)
	prometheus.MustRegister(notifier)
	prometheus.MustRegister(configSuccess)
	prometheus.MustRegister(configSuccessTime)

	// The notifieris a dependency of the rule manager. It has to be
	// started before and torn down afterwards.
	go notifier.Run()
	defer notifier.Stop()

	go ruleManager.Run()
	defer ruleManager.Stop()

	go targetManager.Run()
	defer targetManager.Stop()

	// Shutting down the query engine before the rule manager will cause pending queries
	// to be canceled and ensures a quick shutdown of the rule manager.
	defer queryEngine.Stop()

	go webHandler.Run()

	// Wait for reload or termination signals.
	close(hupReady) // Unblock SIGHUP handler.

	term := make(chan os.Signal)
	signal.Notify(term, os.Interrupt, syscall.SIGTERM)
	select {
	case <-term:
		log.Warn("Received SIGTERM, exiting gracefully...")
	case <-webHandler.Quit():
		log.Warn("Received termination request via web service, exiting gracefully...")
	case err := <-webHandler.ListenError():
		log.Errorln("Error starting web server, exiting gracefully:", err)
	}

	log.Info("See you next time!")
	return 0
}

// Reloadable things can change their internal state to match a new config
// and handle failure gracefully.
type Reloadable interface {
	ApplyConfig(*config.Config) error
}

func reloadConfig(filename string, rls ...Reloadable) (err error) {
	log.Infof("Loading configuration file %s", filename)
	defer func() {
		if err == nil {
			configSuccess.Set(1)
			configSuccessTime.Set(float64(time.Now().Unix()))
		} else {
			configSuccess.Set(0)
		}
	}()

	conf, err := config.LoadFile(filename)
	if err != nil {
		return fmt.Errorf("couldn't load configuration (-config.file=%s): %v", filename, err)
	}

	failed := false
	for _, rl := range rls {
		if err := rl.ApplyConfig(conf); err != nil {
			log.Error("Failed to apply configuration: ", err)
			failed = true
		}
	}
	if failed {
		return fmt.Errorf("one or more errors occured while applying the new configuration (-config.file=%s)", filename)
	}
	return nil
}
<|endoftext|>


<commit_before>// A simple tool for sending raw messages to an adb server.
package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"strings"

	"github.com/zach-klippenstein/goadb/wire"
)

var port = flag.Int("p", wire.AdbPort, "port the adb server is listening on")

func main() {
	flag.Parse()

	fmt.Println("using port", *port)

	printServerVersion()

	for {
		line := readLine()
		err := doCommand(line)
		if err != nil {
			fmt.Println("error:", err)
		}
	}
}

func printServerVersion() {
	err := doCommand("host:version")
	if err != nil {
		log.Fatal(err)
	}
}

func readLine() string {
	fmt.Print("> ")
	line, err := bufio.NewReader(os.Stdin).ReadString('\n')
	if err != nil && err != io.EOF {
		log.Fatal(err)
	}
	return strings.TrimSpace(line)
}

func doCommand(cmd string) error {
	conn, err := wire.NewDialer("", *port).Dial()
	if err != nil {
		log.Fatal(err)
	}
	defer conn.Close()

	if err := wire.SendMessageString(conn, cmd); err != nil {
		return err
	}

	status, err := conn.ReadStatus()
	if err != nil {
		return err
	}

	var msg string
	for err == nil {
		msg, err = wire.ReadMessageString(conn)
		if err == nil {
			fmt.Printf("%s> %s\n", status, msg)
		}
	}

	if err != io.EOF {
		return err
	}

	return nil
}
<commit_msg>Fixed raw-adb.<commit_after>// A simple tool for sending raw messages to an adb server.
package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"strings"

	"github.com/zach-klippenstein/goadb"
	"github.com/zach-klippenstein/goadb/wire"
)

var port = flag.Int("p", goadb.AdbPort, "port the adb server is listening on")

func main() {
	flag.Parse()

	fmt.Println("using port", *port)

	printServerVersion()

	for {
		line := readLine()
		err := doCommand(line)
		if err != nil {
			fmt.Println("error:", err)
		}
	}
}

func printServerVersion() {
	err := doCommand("host:version")
	if err != nil {
		log.Fatal(err)
	}
}

func readLine() string {
	fmt.Print("> ")
	line, err := bufio.NewReader(os.Stdin).ReadString('\n')
	if err != nil && err != io.EOF {
		log.Fatal(err)
	}
	return strings.TrimSpace(line)
}

func doCommand(cmd string) error {
	conn, err := wire.NewDialer("", *port).Dial()
	if err != nil {
		log.Fatal(err)
	}
	defer conn.Close()

	if err := wire.SendMessageString(conn, cmd); err != nil {
		return err
	}

	status, err := conn.ReadStatus()
	if err != nil {
		return err
	}

	var msg string
	for err == nil {
		msg, err = wire.ReadMessageString(conn)
		if err == nil {
			fmt.Printf("%s> %s\n", status, msg)
		}
	}

	if err != io.EOF {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/url"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"github.com/StackExchange/scollector/collect"
	"github.com/StackExchange/scollector/collectors"
	"github.com/StackExchange/scollector/metadata"
	"github.com/StackExchange/scollector/opentsdb"
	"github.com/StackExchange/scollector/util"
	"github.com/StackExchange/slog"
)

// These constants should remain in source control as their zero values.
const (
	// VersionDate should be set at build time as a date: 20140721184001.
	VersionDate uint64 = 0
	// VersionID should be set at build time as the most recent commit hash.
	VersionID string = ""
)

var (
	flagFilter          = flag.String("f", "", "Filters collectors matching this term. Works with all other arguments.")
	flagList            = flag.Bool("l", false, "List")
	flagPrint           = flag.Bool("p", false, "Print to screen instead of sending to a host")
	flagHost            = flag.String("h", "bosun", `bosun or OpenTSDB host. Ex: "http://tsdb.example.com:4242".`)
	flagColDir          = flag.String("c", "", `External collectors directory.`)
	flagBatchSize       = flag.Int("b", 0, "OpenTSDB batch size. Used for debugging bad data.")
	flagSNMP            = flag.String("s", "", "SNMP host to poll of the format: \"community@host[,community@host...]\".")
	flagICMP            = flag.String("i", "", "ICMP host to ping of the format: \"host[,host...]\".")
	flagVsphere         = flag.String("v", "", `vSphere host to poll of the format: "user:password@host[,user:password@host...]".`)
	flagFake            = flag.Int("fake", 0, "Generates X fake data points on the test.fake metric per second.")
	flagDebug           = flag.Bool("d", false, "Enables debug output.")
	flagJSON            = flag.Bool("j", false, "With -p enabled, prints JSON.")
	flagFullHost        = flag.Bool("u", false, `Enables full hostnames: doesn't truncate to first ".".`)
	flagDisableMetadata = flag.Bool("m", false, "Disable sending of metadata.")
	flagVersion         = flag.Bool("version", false, `Prints the version and exits.`)

	mains []func()
)

func readConf() {
	p, err := exePath()
	if err != nil {
		slog.Error(err)
		return
	}
	dir := filepath.Dir(p)
	p = filepath.Join(dir, "scollector.conf")
	b, err := ioutil.ReadFile(p)
	if err != nil {
		if *flagDebug {
			slog.Error(err)
		}
		return
	}
	for i, line := range strings.Split(string(b), "\n") {
		sp := strings.SplitN(line, "=", 2)
		if len(sp) != 2 {
			if *flagDebug {
				slog.Errorf("expected = in %v:%v", p, i+1)
			}
			continue
		}
		k := strings.TrimSpace(sp[0])
		v := strings.TrimSpace(sp[1])
		f := func(s *string) {
			*s = v
		}
		switch k {
		case "host":
			f(flagHost)
		case "filter":
			f(flagFilter)
		case "coldir":
			f(flagColDir)
		case "snmp":
			f(flagSNMP)
		case "icmp":
			f(flagICMP)
		case "vsphere":
			f(flagVsphere)
		default:
			if *flagDebug {
				slog.Errorf("unknown key in %v:%v", p, i+1)
			}
		}
	}
}

func main() {
	flag.Parse()
	if *flagPrint {
		slog.Set(&slog.StdLog{Log: log.New(os.Stdout, "", log.LstdFlags)})
	}
	if *flagVersion {
		fmt.Printf("scollector version %v (%v)\n", VersionDate, VersionID)
		os.Exit(0)
	}
	for _, m := range mains {
		m()
	}
	readConf()

	util.FullHostname = *flagFullHost
	util.Set()
	if *flagColDir != "" {
		collectors.InitPrograms(*flagColDir)
	}
	if *flagSNMP != "" {
		for _, s := range strings.Split(*flagSNMP, ",") {
			sp := strings.Split(s, "@")
			if len(sp) != 2 {
				slog.Fatal("invalid snmp string:", *flagSNMP)
			}
			collectors.SNMPIfaces(sp[0], sp[1])
			collectors.SNMPCisco(sp[0], sp[1])
		}
	}
	if *flagICMP != "" {
		for _, s := range strings.Split(*flagICMP, ",") {
			collectors.ICMP(s)
		}
	}
	if *flagVsphere != "" {
		for _, s := range strings.Split(*flagVsphere, ",") {
			sp := strings.SplitN(s, ":", 2)
			if len(sp) != 2 {
				slog.Fatal("invalid vsphere string:", *flagVsphere)
			}
			user := sp[0]
			idx := strings.LastIndex(sp[1], "@")
			if idx == -1 {
				slog.Fatal("invalid vsphere string:", *flagVsphere)
			}
			pwd := sp[1][:idx]
			host := sp[1][idx+1:]
			if len(user) == 0 || len(pwd) == 0 || len(host) == 0 {
				slog.Fatal("invalid vsphere string:", *flagVsphere)
			}
			collectors.Vsphere(user, pwd, host)
		}
	}
	if *flagFake > 0 {
		collectors.InitFake(*flagFake)
	}
	collect.Debug = *flagDebug
	c := collectors.Search(*flagFilter)
	for _, col := range c {
		col.Init()
	}
	u, err := parseHost()
	if *flagList {
		list(c)
		return
	} else if err != nil {
		slog.Fatal("invalid host:", *flagHost)
	}
	if *flagPrint {
		collectors.DefaultFreq = time.Second * 3
		slog.Infoln("Set default frequency to", collectors.DefaultFreq)
		*flagDisableMetadata = true
	}
	if !*flagDisableMetadata {
		if err := metadata.Init(u, *flagDebug); err != nil {
			slog.Fatal(err)
		}
	}
	cdp := collectors.Run(c)
	if u != nil && !*flagPrint {
		slog.Infoln("OpenTSDB host:", u)
		if err := collect.InitChan(u, "scollector", cdp); err != nil {
			slog.Fatal(err)
		}
		if VersionDate > 0 {
			if err := collect.Put("version", nil, VersionDate); err != nil {
				slog.Error(err)
			}
		}
		if *flagBatchSize > 0 {
			collect.BatchSize = *flagBatchSize
		}
		go func() {
			const maxMem = 500 * 1024 * 1024 // 500MB
			var m runtime.MemStats
			for _ = range time.Tick(time.Minute) {
				runtime.ReadMemStats(&m)
				if m.Alloc > maxMem {
					panic("memory max reached")
				}
			}
		}()
	} else {
		slog.Infoln("Outputting to screen")
		printPut(cdp)
	}
	select {}
}

func exePath() (string, error) {
	prog := os.Args[0]
	p, err := filepath.Abs(prog)
	if err != nil {
		return "", err
	}
	fi, err := os.Stat(p)
	if err == nil {
		if !fi.Mode().IsDir() {
			return p, nil
		}
		err = fmt.Errorf("%s is directory", p)
	}
	if filepath.Ext(p) == "" {
		p += ".exe"
		fi, err := os.Stat(p)
		if err == nil {
			if !fi.Mode().IsDir() {
				return p, nil
			}
			err = fmt.Errorf("%s is directory", p)
		}
	}
	return "", err
}

func list(cs []collectors.Collector) {
	for _, c := range cs {
		fmt.Println(c.Name())
	}
}

func parseHost() (*url.URL, error) {
	if *flagHost == "" {
		return nil, fmt.Errorf("empty host")
	}
	if !strings.Contains(*flagHost, "//") {
		*flagHost = "http://" + *flagHost
	}
	return url.Parse(*flagHost)
}

func printPut(c chan *opentsdb.DataPoint) {
	for dp := range c {
		if *flagJSON {
			b, _ := json.Marshal(dp)
			slog.Info(string(b))
		} else {
			slog.Info(dp.Telnet())
		}
	}
}
<commit_msg>cmd/scollector: better description for the l flag<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/url"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"github.com/StackExchange/scollector/collect"
	"github.com/StackExchange/scollector/collectors"
	"github.com/StackExchange/scollector/metadata"
	"github.com/StackExchange/scollector/opentsdb"
	"github.com/StackExchange/scollector/util"
	"github.com/StackExchange/slog"
)

// These constants should remain in source control as their zero values.
const (
	// VersionDate should be set at build time as a date: 20140721184001.
	VersionDate uint64 = 0
	// VersionID should be set at build time as the most recent commit hash.
	VersionID string = ""
)

var (
	flagFilter          = flag.String("f", "", "Filters collectors matching this term. Works with all other arguments.")
	flagList            = flag.Bool("l", false, "List available collectors")
	flagPrint           = flag.Bool("p", false, "Print to screen instead of sending to a host")
	flagHost            = flag.String("h", "bosun", `bosun or OpenTSDB host. Ex: "http://tsdb.example.com:4242".`)
	flagColDir          = flag.String("c", "", `External collectors directory.`)
	flagBatchSize       = flag.Int("b", 0, "OpenTSDB batch size. Used for debugging bad data.")
	flagSNMP            = flag.String("s", "", "SNMP host to poll of the format: \"community@host[,community@host...]\".")
	flagICMP            = flag.String("i", "", "ICMP host to ping of the format: \"host[,host...]\".")
	flagVsphere         = flag.String("v", "", `vSphere host to poll of the format: "user:password@host[,user:password@host...]".`)
	flagFake            = flag.Int("fake", 0, "Generates X fake data points on the test.fake metric per second.")
	flagDebug           = flag.Bool("d", false, "Enables debug output.")
	flagJSON            = flag.Bool("j", false, "With -p enabled, prints JSON.")
	flagFullHost        = flag.Bool("u", false, `Enables full hostnames: doesn't truncate to first ".".`)
	flagDisableMetadata = flag.Bool("m", false, "Disable sending of metadata.")
	flagVersion         = flag.Bool("version", false, `Prints the version and exits.`)

	mains []func()
)

func readConf() {
	p, err := exePath()
	if err != nil {
		slog.Error(err)
		return
	}
	dir := filepath.Dir(p)
	p = filepath.Join(dir, "scollector.conf")
	b, err := ioutil.ReadFile(p)
	if err != nil {
		if *flagDebug {
			slog.Error(err)
		}
		return
	}
	for i, line := range strings.Split(string(b), "\n") {
		sp := strings.SplitN(line, "=", 2)
		if len(sp) != 2 {
			if *flagDebug {
				slog.Errorf("expected = in %v:%v", p, i+1)
			}
			continue
		}
		k := strings.TrimSpace(sp[0])
		v := strings.TrimSpace(sp[1])
		f := func(s *string) {
			*s = v
		}
		switch k {
		case "host":
			f(flagHost)
		case "filter":
			f(flagFilter)
		case "coldir":
			f(flagColDir)
		case "snmp":
			f(flagSNMP)
		case "icmp":
			f(flagICMP)
		case "vsphere":
			f(flagVsphere)
		default:
			if *flagDebug {
				slog.Errorf("unknown key in %v:%v", p, i+1)
			}
		}
	}
}

func main() {
	flag.Parse()
	if *flagPrint {
		slog.Set(&slog.StdLog{Log: log.New(os.Stdout, "", log.LstdFlags)})
	}
	if *flagVersion {
		fmt.Printf("scollector version %v (%v)\n", VersionDate, VersionID)
		os.Exit(0)
	}
	for _, m := range mains {
		m()
	}
	readConf()

	util.FullHostname = *flagFullHost
	util.Set()
	if *flagColDir != "" {
		collectors.InitPrograms(*flagColDir)
	}
	if *flagSNMP != "" {
		for _, s := range strings.Split(*flagSNMP, ",") {
			sp := strings.Split(s, "@")
			if len(sp) != 2 {
				slog.Fatal("invalid snmp string:", *flagSNMP)
			}
			collectors.SNMPIfaces(sp[0], sp[1])
			collectors.SNMPCisco(sp[0], sp[1])
		}
	}
	if *flagICMP != "" {
		for _, s := range strings.Split(*flagICMP, ",") {
			collectors.ICMP(s)
		}
	}
	if *flagVsphere != "" {
		for _, s := range strings.Split(*flagVsphere, ",") {
			sp := strings.SplitN(s, ":", 2)
			if len(sp) != 2 {
				slog.Fatal("invalid vsphere string:", *flagVsphere)
			}
			user := sp[0]
			idx := strings.LastIndex(sp[1], "@")
			if idx == -1 {
				slog.Fatal("invalid vsphere string:", *flagVsphere)
			}
			pwd := sp[1][:idx]
			host := sp[1][idx+1:]
			if len(user) == 0 || len(pwd) == 0 || len(host) == 0 {
				slog.Fatal("invalid vsphere string:", *flagVsphere)
			}
			collectors.Vsphere(user, pwd, host)
		}
	}
	if *flagFake > 0 {
		collectors.InitFake(*flagFake)
	}
	collect.Debug = *flagDebug
	c := collectors.Search(*flagFilter)
	for _, col := range c {
		col.Init()
	}
	u, err := parseHost()
	if *flagList {
		list(c)
		return
	} else if err != nil {
		slog.Fatal("invalid host:", *flagHost)
	}
	if *flagPrint {
		collectors.DefaultFreq = time.Second * 3
		slog.Infoln("Set default frequency to", collectors.DefaultFreq)
		*flagDisableMetadata = true
	}
	if !*flagDisableMetadata {
		if err := metadata.Init(u, *flagDebug); err != nil {
			slog.Fatal(err)
		}
	}
	cdp := collectors.Run(c)
	if u != nil && !*flagPrint {
		slog.Infoln("OpenTSDB host:", u)
		if err := collect.InitChan(u, "scollector", cdp); err != nil {
			slog.Fatal(err)
		}
		if VersionDate > 0 {
			if err := collect.Put("version", nil, VersionDate); err != nil {
				slog.Error(err)
			}
		}
		if *flagBatchSize > 0 {
			collect.BatchSize = *flagBatchSize
		}
		go func() {
			const maxMem = 500 * 1024 * 1024 // 500MB
			var m runtime.MemStats
			for _ = range time.Tick(time.Minute) {
				runtime.ReadMemStats(&m)
				if m.Alloc > maxMem {
					panic("memory max reached")
				}
			}
		}()
	} else {
		slog.Infoln("Outputting to screen")
		printPut(cdp)
	}
	select {}
}

func exePath() (string, error) {
	prog := os.Args[0]
	p, err := filepath.Abs(prog)
	if err != nil {
		return "", err
	}
	fi, err := os.Stat(p)
	if err == nil {
		if !fi.Mode().IsDir() {
			return p, nil
		}
		err = fmt.Errorf("%s is directory", p)
	}
	if filepath.Ext(p) == "" {
		p += ".exe"
		fi, err := os.Stat(p)
		if err == nil {
			if !fi.Mode().IsDir() {
				return p, nil
			}
			err = fmt.Errorf("%s is directory", p)
		}
	}
	return "", err
}

func list(cs []collectors.Collector) {
	for _, c := range cs {
		fmt.Println(c.Name())
	}
}

func parseHost() (*url.URL, error) {
	if *flagHost == "" {
		return nil, fmt.Errorf("empty host")
	}
	if !strings.Contains(*flagHost, "//") {
		*flagHost = "http://" + *flagHost
	}
	return url.Parse(*flagHost)
}

func printPut(c chan *opentsdb.DataPoint) {
	for dp := range c {
		if *flagJSON {
			b, _ := json.Marshal(dp)
			slog.Info(string(b))
		} else {
			slog.Info(dp.Telnet())
		}
	}
}
<|endoftext|>