<commit_before>/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"strconv"
	"time"

	"github.com/gogo/protobuf/proto"
	log "github.com/golang/glog"
	"github.com/mesos/mesos-go/auth"
	"github.com/mesos/mesos-go/auth/sasl"
	"github.com/mesos/mesos-go/auth/sasl/mech"
	mesos "github.com/mesos/mesos-go/mesosproto"
	util "github.com/mesos/mesos-go/mesosutil"
	sched "github.com/mesos/mesos-go/scheduler"
	mlog "github.com/yp-engineering/mesos-runonce/mesosutil"
	"golang.org/x/net/context"
)

var eventCh chan *mesos.TaskStatus

var (
	address      = flag.String("address", "127.0.0.1", "Address for mesos to callback on.")
	bindingPort  = flag.Uint("port", 0, "Port for address to use for mesos to callback.")
	authProvider = flag.String("authentication-provider", sasl.ProviderName,
		fmt.Sprintf("Authentication provider to use, default is SASL that supports mechanisms: %+v", mech.ListSupported()))
	master              = flag.String("master", "127.0.0.1:5050", "Master address <ip:port>")
	taskCount           = flag.Int("task-count", 1, "Total task count to run.")
	mesosAuthPrincipal  = flag.String("principal", "", "Mesos authentication principal.")
	mesosAuthSecretFile = flag.String("secret-file", "", "Mesos authentication secret file.")
	mesosRunasUser      = flag.String("user", "root", "Mesos user to run tasks as.")
	dockerImage         = flag.String("docker-image", "", "Docker image to run.")
	dockerCmd           = flag.String("docker-cmd", "", "Docker command to run.")
	dockerEnvVars       = flag.String("env-vars", "", "Docker env vars for the container. E.g. -env-vars='{\"env\":{\"FOO\":\"bar\"}}'")
	dCpus               = flag.Float64("cpus", 1.0, "How many CPUs to use.")
	dMem                = flag.Float64("mem", 10, "How much memory to use.")
)

type ExampleScheduler struct {
	executor      *mesos.ExecutorInfo
	tasksLaunched int
	tasksFinished int
	totalTasks    int
}

type DockerEnvVars struct {
	Env map[string]string `json:"env"`
}

func newExampleScheduler(exec *mesos.ExecutorInfo) *ExampleScheduler {
	return &ExampleScheduler{
		executor:      exec,
		tasksLaunched: 0,
		tasksFinished: 0,
		totalTasks:    *taskCount,
	}
}

func (sched *ExampleScheduler) Registered(driver sched.SchedulerDriver, frameworkId *mesos.FrameworkID, masterInfo *mesos.MasterInfo) {
	log.Infoln("Framework Registered with Master ", masterInfo)
}
func (sched *ExampleScheduler) Reregistered(driver sched.SchedulerDriver, masterInfo *mesos.MasterInfo) {
	log.Infoln("Framework Re-Registered with Master ", masterInfo)
}
func (sched *ExampleScheduler) Disconnected(sched.SchedulerDriver) {
	log.Fatalf("disconnected from master, aborting")
}
func (sched *ExampleScheduler) OfferRescinded(_ sched.SchedulerDriver, oid *mesos.OfferID) {
	log.Errorf("offer rescinded: %v", oid)
}
func (sched *ExampleScheduler) FrameworkMessage(_ sched.SchedulerDriver, eid *mesos.ExecutorID, sid *mesos.SlaveID, msg string) {
	log.Errorf("framework message from executor %q slave %q: %q", eid, sid, msg)
}
func (sched *ExampleScheduler) SlaveLost(_ sched.SchedulerDriver, sid *mesos.SlaveID) {
	log.Errorf("slave lost: %v", sid)
}
func (sched *ExampleScheduler) ExecutorLost(_ sched.SchedulerDriver, eid *mesos.ExecutorID, sid *mesos.SlaveID, code int) {
	log.Errorf("executor %q lost on slave %q code %d", eid, sid, code)
}
func (sched *ExampleScheduler) Error(_ sched.SchedulerDriver, err string) {
	log.Errorf("Scheduler received error: %v", err)
}

func (sched *ExampleScheduler) ResourceOffers(driver sched.SchedulerDriver, offers []*mesos.Offer) {

	if sched.tasksLaunched >= sched.totalTasks {
		log.Info("decline all of the offers since all of our tasks are already launched")
		ids := make([]*mesos.OfferID, len(offers))
		for i, offer := range offers {
			ids[i] = offer.Id
		}
		driver.LaunchTasks(ids, []*mesos.TaskInfo{}, &mesos.Filters{RefuseSeconds: proto.Float64(120)})
		return
	}
	for _, offer := range offers {
		cpuResources := util.FilterResources(offer.Resources, func(res *mesos.Resource) bool {
			return res.GetName() == "cpus"
		})
		cpus := 0.0
		for _, res := range cpuResources {
			cpus += res.GetScalar().GetValue()
		}

		memResources := util.FilterResources(offer.Resources, func(res *mesos.Resource) bool {
			return res.GetName() == "mem"
		})
		mems := 0.0
		for _, res := range memResources {
			mems += res.GetScalar().GetValue()
		}

		log.Infoln("Received Offer <", offer.Id.GetValue(), "> on host ", *offer.Hostname, "with cpus=", cpus, " mem=", mems)

		remainingCpus := cpus
		remainingMems := mems

		var tasks []*mesos.TaskInfo
		for sched.tasksLaunched < sched.totalTasks &&
			*dCpus <= remainingCpus &&
			*dMem <= remainingMems {

			sched.tasksLaunched++

			taskId := &mesos.TaskID{
				Value: proto.String(strconv.Itoa(sched.tasksLaunched)),
			}

			containerType := mesos.ContainerInfo_DOCKER
			task := &mesos.TaskInfo{
				Name:    proto.String("mesos-runonce-" + taskId.GetValue()),
				TaskId:  taskId,
				SlaveId: offer.SlaveId,
				Resources: []*mesos.Resource{
					util.NewScalarResource("cpus", *dCpus),
					util.NewScalarResource("mem", *dMem),
				},
				Command: &mesos.CommandInfo{},
				Container: &mesos.ContainerInfo{
					Type: &containerType,
					Docker: &mesos.ContainerInfo_DockerInfo{
						Image: proto.String(*dockerImage),
					},
				},
			}
			log.Infof("Prepared task: %s with offer %s for launch\n", task.GetName(), offer.Id.GetValue())

			if *dockerEnvVars != "" {
				task.Command.Environment = envVars()
			}

			// Allow arbitrary commands, else just use whatever the image defines in CMD
			if *dockerCmd != "" {
				task.Command.Value = proto.String(*dockerCmd)
			}

			tasks = append(tasks, task)
			remainingCpus -= *dCpus
			remainingMems -= *dMem
		}
		log.Infoln("Launching ", len(tasks), "tasks for offer", offer.Id.GetValue())
		driver.LaunchTasks([]*mesos.OfferID{offer.Id}, tasks, &mesos.Filters{RefuseSeconds: proto.Float64(5)})
	}
}

func (sched *ExampleScheduler) StatusUpdate(driver sched.SchedulerDriver, status *mesos.TaskStatus) {
	log.Infoln("Status update: task", status.TaskId.GetValue(), " is in state ", status.State.Enum().String())
	eventCh <- status

	if status.GetState() == mesos.TaskState_TASK_FINISHED {
		sched.tasksFinished++
	}

	if sched.tasksFinished >= sched.totalTasks {
		log.Infoln("Total tasks completed, stopping framework.")
		driver.Stop(false)
	}

	if status.GetState() == mesos.TaskState_TASK_LOST ||
		status.GetState() == mesos.TaskState_TASK_KILLED ||
		status.GetState() == mesos.TaskState_TASK_FAILED ||
		status.GetState() == mesos.TaskState_TASK_ERROR {
		log.Fatalln(
			"Aborting because task", status.TaskId.GetValue(),
			"is in unexpected state", status.State.String(),
			"with message", status.GetMessage(),
		)
		driver.Abort()
	}
}

// ----------------------- func init() ------------------------- //

func init() {
	flag.Parse()
	log.Infoln("Initializing the Example Scheduler...")
}

func fetchLogs() {
	timer := time.Tick(500 * time.Millisecond)
	finished := false
	offset := 0
	var startedStatus *mesos.TaskStatus
	for {
		select {
		case status := <-eventCh:
			if status.GetState() == mesos.TaskState_TASK_RUNNING {
				startedStatus = status
			}
			if status.GetState() == mesos.TaskState_TASK_FINISHED {
				finished = true
			}
		case <-timer:
			if startedStatus != nil {
				data, err := mlog.FetchLogs(startedStatus, offset)
				if err != nil {
					log.Infof("fetch logs err: %s\n", err)
				} else {
					fmt.Print(string(data))
					if len(data) == 0 && finished {
						return
					}
					offset += len(data)
				}
			}
		}
	}
}
func prepareExecutorInfo() *mesos.ExecutorInfo {
	// Create mesos scheduler driver.
	containerType := mesos.ContainerInfo_DOCKER
	return &mesos.ExecutorInfo{
		ExecutorId: util.NewExecutorID("default"),
		Name:       proto.String("Test Executor (Go)"),
		Source:     proto.String("go_test"),
		Container: &mesos.ContainerInfo{
			Type: &containerType,
		},
	}
}

func parseIP(address string) net.IP {
	addr, err := net.LookupIP(address)
	if err != nil {
		log.Fatal(err)
	}
	if len(addr) < 1 {
		log.Fatalf("failed to parse IP from address '%v'", address)
	}
	return addr[0]
}

func envVars() *mesos.Environment {
	var dEnvVars DockerEnvVars
	err := json.Unmarshal([]byte(*dockerEnvVars), &dEnvVars)
	if err != nil {
		log.Fatalf("JSON error: %#v with unparsable env vars: %+v", err, *dockerEnvVars)
	}

	var variables []*mesos.Environment_Variable
	for key, value := range dEnvVars.Env {
		variables = append(variables, &mesos.Environment_Variable{Name: proto.String(key), Value: proto.String(value)})
	}

	return &mesos.Environment{Variables: variables}
}

// ----------------------- func main() ------------------------- //

func main() {

	// build command executor
	exec := prepareExecutorInfo()

	// the framework
	fwinfo := &mesos.FrameworkInfo{
		User: proto.String(*mesosRunasUser),
		Name: proto.String("mesos-runonce"),
	}

	cred := (*mesos.Credential)(nil)
	if *mesosAuthPrincipal != "" {
		fwinfo.Principal = proto.String(*mesosAuthPrincipal)
		cred = &mesos.Credential{
			Principal: proto.String(*mesosAuthPrincipal),
		}
		if *mesosAuthSecretFile != "" {
			_, err := os.Stat(*mesosAuthSecretFile)
			if err != nil {
				log.Fatal("missing secret file: ", err.Error())
			}
			secret, err := ioutil.ReadFile(*mesosAuthSecretFile)
			if err != nil {
				log.Fatal("failed to read secret file: ", err.Error())
			}
			cred.Secret = proto.String(string(secret))
		}
	}

	publishedAddress := parseIP(*address)
	config := sched.DriverConfig{
		Scheduler:        newExampleScheduler(exec),
		Framework:        fwinfo,
		Master:           *master,
		Credential:       cred,
		PublishedAddress: publishedAddress,
		WithAuthContext: func(ctx context.Context) context.Context {
			ctx = auth.WithLoginProvider(ctx, *authProvider)
			ctx = sasl.WithBindingAddress(ctx, publishedAddress)
			return ctx
		},
	}

	// Allow listening port to be configurable so we can run this inside of
	// mesos if desired.
	if *bindingPort != 0 {
		config.BindingPort = uint16(*bindingPort)
	}

	driver, err := sched.NewMesosSchedulerDriver(config)

	if err != nil {
		log.Errorln("Unable to create a SchedulerDriver ", err.Error())
	}

	eventCh = make(chan *mesos.TaskStatus)
	go fetchLogs()
	if stat, err := driver.Run(); err != nil {
		log.Infof("Framework stopped with status %s and error: %s\n", stat.String(), err.Error())
	}
	log.Infof("framework terminating")
}
<commit_msg>Setting shell to false by default unless we pass in a value to run. Also changing the name of the executor.<commit_after>/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"strconv"
	"time"

	"github.com/gogo/protobuf/proto"
	log "github.com/golang/glog"
	"github.com/mesos/mesos-go/auth"
	"github.com/mesos/mesos-go/auth/sasl"
	"github.com/mesos/mesos-go/auth/sasl/mech"
	mesos "github.com/mesos/mesos-go/mesosproto"
	util "github.com/mesos/mesos-go/mesosutil"
	sched "github.com/mesos/mesos-go/scheduler"
	mlog "github.com/yp-engineering/mesos-runonce/mesosutil"
	"golang.org/x/net/context"
)

var eventCh chan *mesos.TaskStatus

var (
	address      = flag.String("address", "127.0.0.1", "Address for mesos to callback on.")
	bindingPort  = flag.Uint("port", 0, "Port for address to use for mesos to callback.")
	authProvider = flag.String("authentication-provider", sasl.ProviderName,
		fmt.Sprintf("Authentication provider to use, default is SASL that supports mechanisms: %+v", mech.ListSupported()))
	master              = flag.String("master", "127.0.0.1:5050", "Master address <ip:port>")
	taskCount           = flag.Int("task-count", 1, "Total task count to run.")
	mesosAuthPrincipal  = flag.String("principal", "", "Mesos authentication principal.")
	mesosAuthSecretFile = flag.String("secret-file", "", "Mesos authentication secret file.")
	mesosRunasUser      = flag.String("user", "root", "Mesos user to run tasks as.")
	dockerImage         = flag.String("docker-image", "", "Docker image to run.")
	dockerCmd           = flag.String("docker-cmd", "", "Docker command to run.")
	dockerEnvVars       = flag.String("env-vars", "", "Docker env vars for the container. E.g. -env-vars='{\"env\":{\"FOO\":\"bar\"}}'")
	dCpus               = flag.Float64("cpus", 1.0, "How many CPUs to use.")
	dMem                = flag.Float64("mem", 10, "How much memory to use.")
)

type ExampleScheduler struct {
	executor      *mesos.ExecutorInfo
	tasksLaunched int
	tasksFinished int
	totalTasks    int
}

type DockerEnvVars struct {
	Env map[string]string `json:"env"`
}

func newExampleScheduler(exec *mesos.ExecutorInfo) *ExampleScheduler {
	return &ExampleScheduler{
		executor:      exec,
		tasksLaunched: 0,
		tasksFinished: 0,
		totalTasks:    *taskCount,
	}
}

func (sched *ExampleScheduler) Registered(driver sched.SchedulerDriver, frameworkId *mesos.FrameworkID, masterInfo *mesos.MasterInfo) {
	log.Infoln("Framework Registered with Master ", masterInfo)
}
func (sched *ExampleScheduler) Reregistered(driver sched.SchedulerDriver, masterInfo *mesos.MasterInfo) {
	log.Infoln("Framework Re-Registered with Master ", masterInfo)
}
func (sched *ExampleScheduler) Disconnected(sched.SchedulerDriver) {
	log.Fatalf("disconnected from master, aborting")
}
func (sched *ExampleScheduler) OfferRescinded(_ sched.SchedulerDriver, oid *mesos.OfferID) {
	log.Errorf("offer rescinded: %v", oid)
}
func (sched *ExampleScheduler) FrameworkMessage(_ sched.SchedulerDriver, eid *mesos.ExecutorID, sid *mesos.SlaveID, msg string) {
	log.Errorf("framework message from executor %q slave %q: %q", eid, sid, msg)
}
func (sched *ExampleScheduler) SlaveLost(_ sched.SchedulerDriver, sid *mesos.SlaveID) {
	log.Errorf("slave lost: %v", sid)
}
func (sched *ExampleScheduler) ExecutorLost(_ sched.SchedulerDriver, eid *mesos.ExecutorID, sid *mesos.SlaveID, code int) {
	log.Errorf("executor %q lost on slave %q code %d", eid, sid, code)
}
func (sched *ExampleScheduler) Error(_ sched.SchedulerDriver, err string) {
	log.Errorf("Scheduler received error: %v", err)
}

func (sched *ExampleScheduler) ResourceOffers(driver sched.SchedulerDriver, offers []*mesos.Offer) {

	if sched.tasksLaunched >= sched.totalTasks {
		log.Info("decline all of the offers since all of our tasks are already launched")
		ids := make([]*mesos.OfferID, len(offers))
		for i, offer := range offers {
			ids[i] = offer.Id
		}
		driver.LaunchTasks(ids, []*mesos.TaskInfo{}, &mesos.Filters{RefuseSeconds: proto.Float64(120)})
		return
	}
	for _, offer := range offers {
		cpuResources := util.FilterResources(offer.Resources, func(res *mesos.Resource) bool {
			return res.GetName() == "cpus"
		})
		cpus := 0.0
		for _, res := range cpuResources {
			cpus += res.GetScalar().GetValue()
		}

		memResources := util.FilterResources(offer.Resources, func(res *mesos.Resource) bool {
			return res.GetName() == "mem"
		})
		mems := 0.0
		for _, res := range memResources {
			mems += res.GetScalar().GetValue()
		}

		log.Infoln("Received Offer <", offer.Id.GetValue(), "> on host ", *offer.Hostname, "with cpus=", cpus, " mem=", mems)

		remainingCpus := cpus
		remainingMems := mems

		var tasks []*mesos.TaskInfo
		for sched.tasksLaunched < sched.totalTasks &&
			*dCpus <= remainingCpus &&
			*dMem <= remainingMems {

			sched.tasksLaunched++

			taskId := &mesos.TaskID{
				Value: proto.String(strconv.Itoa(sched.tasksLaunched)),
			}

			containerType := mesos.ContainerInfo_DOCKER
			task := &mesos.TaskInfo{
				Name:    proto.String("mesos-runonce-" + taskId.GetValue()),
				TaskId:  taskId,
				SlaveId: offer.SlaveId,
				Resources: []*mesos.Resource{
					util.NewScalarResource("cpus", *dCpus),
					util.NewScalarResource("mem", *dMem),
				},
				Command: &mesos.CommandInfo{
					Shell: proto.Bool(false),
				},
				Container: &mesos.ContainerInfo{
					Type: &containerType,
					Docker: &mesos.ContainerInfo_DockerInfo{
						Image: proto.String(*dockerImage),
					},
				},
			}
			log.Infof("Prepared task: %s with offer %s for launch\n", task.GetName(), offer.Id.GetValue())

			if *dockerEnvVars != "" {
				task.Command.Environment = envVars()
			}

			// Allow arbitrary commands, else just use whatever the image defines in CMD
			if *dockerCmd != "" {
				task.Command.Value = proto.String(*dockerCmd)
				task.Command.Shell = proto.Bool(true)
			}

			tasks = append(tasks, task)
			remainingCpus -= *dCpus
			remainingMems -= *dMem
		}
		log.Infoln("Launching ", len(tasks), "tasks for offer", offer.Id.GetValue())
		driver.LaunchTasks([]*mesos.OfferID{offer.Id}, tasks, &mesos.Filters{RefuseSeconds: proto.Float64(5)})
	}
}

func (sched *ExampleScheduler) StatusUpdate(driver sched.SchedulerDriver, status *mesos.TaskStatus) {
	log.Infoln("Status update: task", status.TaskId.GetValue(), " is in state ", status.State.Enum().String())
	eventCh <- status

	if status.GetState() == mesos.TaskState_TASK_FINISHED {
		sched.tasksFinished++
	}

	if sched.tasksFinished >= sched.totalTasks {
		log.Infoln("Total tasks completed, stopping framework.")
		driver.Stop(false)
	}

	if status.GetState() == mesos.TaskState_TASK_LOST ||
		status.GetState() == mesos.TaskState_TASK_KILLED ||
		status.GetState() == mesos.TaskState_TASK_FAILED ||
		status.GetState() == mesos.TaskState_TASK_ERROR {
		log.Fatalln(
			"Aborting because task", status.TaskId.GetValue(),
			"is in unexpected state", status.State.String(),
			"with message", status.GetMessage(),
		)
		driver.Abort()
	}
}

// ----------------------- func init() ------------------------- //

func init() {
	flag.Parse()
	log.Infoln("Initializing the Example Scheduler...")
}

func fetchLogs() {
	timer := time.Tick(500 * time.Millisecond)
	finished := false
	offset := 0
	var startedStatus *mesos.TaskStatus
	for {
		select {
		case status := <-eventCh:
			if status.GetState() == mesos.TaskState_TASK_RUNNING {
				startedStatus = status
			}
			if status.GetState() == mesos.TaskState_TASK_FINISHED {
				finished = true
			}
		case <-timer:
			if startedStatus != nil {
				data, err := mlog.FetchLogs(startedStatus, offset)
				if err != nil {
					log.Infof("fetch logs err: %s\n", err)
				} else {
					fmt.Print(string(data))
					if len(data) == 0 && finished {
						return
					}
					offset += len(data)
				}
			}
		}
	}
}
func prepareExecutorInfo() *mesos.ExecutorInfo {
	// Create mesos scheduler driver.
	containerType := mesos.ContainerInfo_DOCKER
	return &mesos.ExecutorInfo{
		ExecutorId: util.NewExecutorID("default"),
		Name:       proto.String("mesos-runonce-executor"),
		Source:     proto.String("mesos-runonce-executor"),
		Container: &mesos.ContainerInfo{
			Type: &containerType,
		},
	}
}

func parseIP(address string) net.IP {
	addr, err := net.LookupIP(address)
	if err != nil {
		log.Fatal(err)
	}
	if len(addr) < 1 {
		log.Fatalf("failed to parse IP from address '%v'", address)
	}
	return addr[0]
}

func envVars() *mesos.Environment {
	var dEnvVars DockerEnvVars
	err := json.Unmarshal([]byte(*dockerEnvVars), &dEnvVars)
	if err != nil {
		log.Fatalf("JSON error: %#v with unparsable env vars: %+v", err, *dockerEnvVars)
	}

	var variables []*mesos.Environment_Variable
	for key, value := range dEnvVars.Env {
		variables = append(variables, &mesos.Environment_Variable{Name: proto.String(key), Value: proto.String(value)})
	}

	return &mesos.Environment{Variables: variables}
}

// ----------------------- func main() ------------------------- //

func main() {

	// build command executor
	exec := prepareExecutorInfo()

	// the framework
	fwinfo := &mesos.FrameworkInfo{
		User: proto.String(*mesosRunasUser),
		Name: proto.String("mesos-runonce"),
	}

	cred := (*mesos.Credential)(nil)
	if *mesosAuthPrincipal != "" {
		fwinfo.Principal = proto.String(*mesosAuthPrincipal)
		cred = &mesos.Credential{
			Principal: proto.String(*mesosAuthPrincipal),
		}
		if *mesosAuthSecretFile != "" {
			_, err := os.Stat(*mesosAuthSecretFile)
			if err != nil {
				log.Fatal("missing secret file: ", err.Error())
			}
			secret, err := ioutil.ReadFile(*mesosAuthSecretFile)
			if err != nil {
				log.Fatal("failed to read secret file: ", err.Error())
			}
			cred.Secret = proto.String(string(secret))
		}
	}

	publishedAddress := parseIP(*address)
	config := sched.DriverConfig{
		Scheduler:        newExampleScheduler(exec),
		Framework:        fwinfo,
		Master:           *master,
		Credential:       cred,
		PublishedAddress: publishedAddress,
		WithAuthContext: func(ctx context.Context) context.Context {
			ctx = auth.WithLoginProvider(ctx, *authProvider)
			ctx = sasl.WithBindingAddress(ctx, publishedAddress)
			return ctx
		},
	}

	// Allow listening port to be configurable so we can run this inside of
	// mesos if desired.
	if *bindingPort != 0 {
		config.BindingPort = uint16(*bindingPort)
	}

	driver, err := sched.NewMesosSchedulerDriver(config)

	if err != nil {
		log.Errorln("Unable to create a SchedulerDriver ", err.Error())
	}

	eventCh = make(chan *mesos.TaskStatus)
	go fetchLogs()
	if stat, err := driver.Run(); err != nil {
		log.Infof("Framework stopped with status %s and error: %s\n", stat.String(), err.Error())
	}
	log.Infof("framework terminating")
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/gosuri/uiprogress"
)

const (
	pomodoroTime = 25 * time.Second
	breakTime    = 5 * time.Second
)

type Pomodoro struct {
	bar     *uiprogress.Bar
	running bool
	done    chan struct{}
}

func NewPomodoro() *Pomodoro {
	pomodoro := Pomodoro{
		done: make(chan struct{}, 0),
	}

	sig := make(chan os.Signal, 1)
	signal.Notify(sig, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)

	go func() {
		for {
			select {
			case <-pomodoro.done:
			case <-sig:
				if pomodoro.running {
					pomodoro.Stop()
				} else {
					os.Exit(0)
				}
			}
		}
	}()

	uiprogress.Start()

	return &pomodoro
}

func (p *Pomodoro) Run(duration int) bool {
	tick := time.NewTicker(time.Second / 10)
	p.running = true
	p.bar = uiprogress.AddBar(10 * duration).AppendCompleted().PrependElapsed()

	defer func() {
		tick.Stop()
		p.running = false
		p.done = make(chan struct{}, 0)
	}()

	for {
		select {
		case <-tick.C:
			if p.bar.Current() == duration {
				p.running = false
				close(p.done)
				return true
			}
			p.bar.Set(p.bar.Current() + 1)
		case <-p.done:
			p.running = false
			return false
		}
	}
}

func (p *Pomodoro) Stop() bool {
	if !p.running {
		return false
	}
	close(p.done)
	return true
}

func main() {
	print("\033[H\033[2J")
	stdin := bufio.NewReader(os.Stdin)

	pomodoro := NewPomodoro()
	for {
		// Pomodoro.
		if !pomodoro.Run(25 * 60) {
			print("    Hit enter to start pomodoro..")
			stdin.ReadString('\n')
			print("\033[H\033[2J")
			continue
		}
		print("\a    Hit enter to start break..")
		stdin.ReadString('\n')
		print("\033[H\033[2J")

		// Break.
		pomodoro.Run(5 * 60)
		print("\a    Hit enter to start pomodoro..")
		stdin.ReadString('\n')
		print("\033[H\033[2J")
	}
}
<commit_msg>Fix progress bar<commit_after>package main

import (
	"bufio"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/gosuri/uiprogress"
)

const (
	pomodoroTime = 25 * time.Second
	breakTime    = 5 * time.Second
)

// TODO: progress bar per pomodoro
// TODO: msg per pomodoro
// TODO: remove break pomodoros
type Pomodoro struct {
	bar     *uiprogress.Bar
	running bool
	done    chan struct{}
}

func NewPomodoro() *Pomodoro {
	pomodoro := Pomodoro{
		done: make(chan struct{}, 0),
	}

	sig := make(chan os.Signal, 1)
	signal.Notify(sig, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)

	go func() {
		for {
			select {
			case <-pomodoro.done:
			case <-sig:
				if pomodoro.running {
					pomodoro.Stop()
				} else {
					os.Exit(0)
				}
			}
		}
	}()

	uiprogress.Start()

	return &pomodoro
}

func (p *Pomodoro) Run(duration int) bool {
	tick := time.NewTicker(time.Second / 10)
	p.running = true
	p.bar = uiprogress.AddBar(10 * duration).AppendCompleted().PrependElapsed()
	/*.PrependFunc(func(b *uiprogress.Bar) string {
		return p.msg + ": "
	})*/

	defer func() {
		tick.Stop()
		p.running = false
		p.done = make(chan struct{}, 0)
	}()

	for {
		select {
		case <-tick.C:
			p.bar.Set(p.bar.Current() + 1)
			if p.bar.Current() >= 10*duration {
				p.running = false
				close(p.done)
				return true
			}
		case <-p.done:
			return false
		}
	}
}

func (p *Pomodoro) Stop() bool {
	if !p.running {
		return false
	}
	close(p.done)
	return true
}

func main() {
	print("\033[H\033[2J")
	stdin := bufio.NewReader(os.Stdin)

	pomodoro := NewPomodoro()
	for {
		// Pomodoro.
		if !pomodoro.Run(25 * 60) {
			print("    Hit enter to start pomodoro..")
			stdin.ReadString('\n')
			print("\033[H\033[2J")
			continue
		}
		print("\a    Hit enter to start break..")
		stdin.ReadString('\n')
		print("\033[H\033[2J")

		// Break.
		pomodoro.Run(5 * 60)
		print("\a    Hit enter to start pomodoro..")
		stdin.ReadString('\n')
		print("\033[H\033[2J")
	}
}
<|endoftext|>


<commit_before>package dmuc

import (
	"bytes"
	"flag"
	"os"
)

var l = flag.Bool("l", false, "List files from /usr/local/bin directory")
var a = flag.Bool("a", false, "List files from both /usr/bin and /usr/local/bin directory")
var s = flag.String("s", "", "Applies a grep 'char' filter to the output")
var i = flag.String("i", "", "Apllies a grep 'string' filter to the output")

// Main This function starts the program
func Main() {
	flag.Parse()

	var output bytes.Buffer
	lsArgs, grepArgs := CreateBashCommand(l, a, s, i)
	result := RunCommand(lsArgs, grepArgs, output)
	PrintToScreen(os.Stdout, result)
}
<commit_msg>So packages are not executable. This main file should not have been in the repo. Instead, this file should be a stand alone application that calls the needed functions in the package. The application seems to be working as expected. I need to take some time to add a readme and steps on how to use this application<commit_after><|endoftext|>


<commit_before>// Download Montreal traffic camera images
// Uses open data from the city of Montreal at:
//   - http://donnees.ville.montreal.qc.ca/dataset/cameras-observation-routiere
// See in action via city map application at @ http://ville.montreal.qc.ca/circulation/
package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"os"
	"path"
	"regexp"
	"strconv"
	"sync"
	"time"

	"github.com/sethgrid/multibar"
)

const (
	// See http://donnees.ville.montreal.qc.ca/dataset/cameras-observation-routiere
	URLGeoJSON string = "http://ville.montreal.qc.ca/circulation/sites/ville.montreal.qc.ca.circulation/files/cameras-de-circulation.json"
	URLAbout   string = "http://donnees.ville.montreal.qc.ca/dataset/cameras-observation-routiere"
)

var (
	reDigits = regexp.MustCompile(`\d+`)
)

// Create date/timestampped subdirectories for saving images
func MakeTimeStampDir(parentDir string) string {
	timeStampDir := time.Now().Format(path.Join(parentDir, "060102/150405"))
	err := os.MkdirAll(timeStampDir, 0755)
	if err != nil {
		log.Fatalf("Couldn't create image directory: %s\n", err)
	}

	return timeStampDir
}

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "\n%s: Download Montreal traffic camera images\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "Data source: %s\n\n", URLAbout)
		flag.PrintDefaults()
		fmt.Fprintf(os.Stderr, "\n")
	}

	concurrency := flag.Int("c", 90, "max concurrent downloads")
	parentDir := flag.String("d", "images", "parent directory for downloaded files")
	flag.Parse()

	// timestamped directory to store images in
	tsDir := MakeTimeStampDir(*parentDir)

	// max number of workers (counting semaphore)
	workers := make(chan int, *concurrency)

	var fc FeatureCollection
	geoJSON := download(URLGeoJSON)
	if err := json.Unmarshal(geoJSON, &fc); err != nil {
		log.Printf("Error unmarshalling GeoJSON data: %v\n", err)
	}

	// progress bar
	bars, _ := multibar.New()
	nImgs := len(fc.Features)
	bar1 := bars.MakeBar(nImgs, "Completed Downloads")
	go bars.Listen()

	// FIXME: we are updating shared variable concurrently, use a channel
	// ex: see http://play.golang.org/p/Uc9vlblxMA
	nCompleted := 0

	// wg waits for all downloads to complete
	var wg sync.WaitGroup
	for _, f := range fc.Features {
		imgURL := f.Properties.URLImageEnDirect
		if imgURL == "" {
			continue
		}

		// use one counting semaphore slot (when full, will block until slot free)
		workers <- 1
		wg.Add(1)

		imgFileServer := path.Base(imgURL)
		if imgFileServer == "" {
			log.Fatalf("Couldn't derive filename: %s\n", imgURL)
		}
		imgNum, err := strconv.Atoi(reDigits.FindString(imgFileServer))
		if err != nil {
			log.Printf("Error deriving image number: %s\n", err)
		}
		imgFile := fmt.Sprintf("%03d.jpg", imgNum)

		go func(URL, file, dir string) {
			defer wg.Done()
			img := download(URL)
			saveFile(img, file, dir)

			// update progressbar
			// FIXME (see nCompleted declaration above)
			nCompleted++
			bar1(nCompleted)

			// when download finished, free one slot in workers
			<-workers
		}(imgURL, imgFile, tsDir)
	}

	wg.Wait()
}
<commit_msg>Clean up comments<commit_after>// Download Montreal traffic camera images
package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"os"
	"path"
	"regexp"
	"strconv"
	"sync"
	"time"

	"github.com/sethgrid/multibar"
)

const (
	// See http://donnees.ville.montreal.qc.ca/dataset/cameras-observation-routiere
	URLGeoJSON string = "http://ville.montreal.qc.ca/circulation/sites/ville.montreal.qc.ca.circulation/files/cameras-de-circulation.json"
	URLAbout   string = "http://donnees.ville.montreal.qc.ca/dataset/cameras-observation-routiere"
)

var (
	reDigits = regexp.MustCompile(`\d+`)
)

// Create date/timestampped subdirectories for saving images
func MakeTimeStampDir(parentDir string) string {
	timeStampDir := time.Now().Format(path.Join(parentDir, "060102/150405"))
	err := os.MkdirAll(timeStampDir, 0755)
	if err != nil {
		log.Fatalf("Couldn't create image directory: %s\n", err)
	}

	return timeStampDir
}

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "\n%s: Download Montreal traffic camera images\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "Data source: %s\n\n", URLAbout)
		flag.PrintDefaults()
		fmt.Fprintf(os.Stderr, "\n")
	}

	concurrency := flag.Int("c", 90, "max concurrent downloads")
	parentDir := flag.String("d", "images", "parent directory for downloaded files")
	flag.Parse()

	// timestamped directory to store images in
	tsDir := MakeTimeStampDir(*parentDir)

	// max number of workers (counting semaphore)
	workers := make(chan int, *concurrency)

	var fc FeatureCollection
	geoJSON := download(URLGeoJSON)
	if err := json.Unmarshal(geoJSON, &fc); err != nil {
		log.Printf("Error unmarshalling GeoJSON data: %v\n", err)
	}

	// progress bar
	bars, _ := multibar.New()
	nImgs := len(fc.Features)
	bar1 := bars.MakeBar(nImgs, "Completed Downloads")
	go bars.Listen()

	// FIXME: we are updating shared variable concurrently, use a channel
	// ex: see http://play.golang.org/p/Uc9vlblxMA
	nCompleted := 0

	// wg waits for all downloads to complete
	var wg sync.WaitGroup
	for _, f := range fc.Features {
		imgURL := f.Properties.URLImageEnDirect
		if imgURL == "" {
			continue
		}

		// use one counting semaphore slot (when full, will block until slot free)
		workers <- 1
		wg.Add(1)

		imgFileServer := path.Base(imgURL)
		if imgFileServer == "" {
			log.Fatalf("Couldn't derive filename: %s\n", imgURL)
		}
		imgNum, err := strconv.Atoi(reDigits.FindString(imgFileServer))
		if err != nil {
			log.Printf("Error deriving image number: %s\n", err)
		}
		imgFile := fmt.Sprintf("%03d.jpg", imgNum)

		go func(URL, file, dir string) {
			defer wg.Done()
			img := download(URL)
			saveFile(img, file, dir)

			// update progressbar
			// FIXME (see nCompleted declaration above)
			nCompleted++
			bar1(nCompleted)

			// when download finished, free one slot in workers
			<-workers
		}(imgURL, imgFile, tsDir)
	}

	wg.Wait()
}
<|endoftext|>


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

// Pythia is a web application front-end for the Go source code oracle.
package main

import (
	"flag"
	"fmt"
	"go/build"
	"go/token"
	"net"
	"net/http"
	"os"
	"os/exec"
	"runtime"
	"sort"
	"strings"
	"sync"

	"code.google.com/p/go.tools/importer"
	"code.google.com/p/go.tools/oracle"
)

var (
	httpAddr = flag.String("http", ":8080", "HTTP listen address")
	verbose  = flag.Bool("v", false, "Verbose mode: print incoming queries")
	open     = flag.Bool("open", true, "Try to open browser")
	args     []string
	files    []string
	packages []*importer.PackageInfo
	imp      *importer.Importer
	ora      *oracle.Oracle
	mutex    sync.Mutex
)

func init() {
	if os.Getenv("GOMAXPROCS") == "" {
		n := runtime.NumCPU()
		if n < 4 {
			n = 4
		}
		runtime.GOMAXPROCS(n)
	}
}

const useHelp = "Run 'pythia -help' for more information.\n"

const helpMessage = `Web frontend for the Go source code oracle.
Usage: pythia [<flag> ...] <args> ...

The -http flag specifies the HTTP service address (e.g., ':6060').

The -open flag determines, whether the application should try to
open the browser. It is set to 'true' by default. If set to 'false'
the browser will not be launched.

The -v flag enables verbose mode, in which every incoming query
to the oracle is logged to the standard output.

Examples:

Start pythia with the scope of package oracle:
% pythia code.google.com/p/go.tools/cmd/oracle

Start pythia with the scope of package image/png on port 8081,
but don't open the browser:
% pythia -http=:8081 -open=false image/png
` + importer.InitialPackagesUsage

func main() {
	flag.Usage = func() { fmt.Fprint(os.Stderr, useHelp) }
	flag.CommandLine.Init(os.Args[0], flag.ContinueOnError)
	if err := flag.CommandLine.Parse(os.Args[1:]); err != nil {
		if err == flag.ErrHelp {
			fmt.Println(helpMessage)
			fmt.Println("Flags:")
			flag.PrintDefaults()
		}
		os.Exit(2)
	}
	args = flag.Args()
	if len(args) == 0 {
		fmt.Fprint(os.Stderr, "Error: no package arguments.\n"+useHelp)
		os.Exit(2)
	}

	var err error
	imp = importer.New(&importer.Config{Build: &build.Default})
	ora, err = oracle.New(imp, args, nil, false)
	if err != nil {
		exitError(err)
	}
	files = scopeFiles(imp)
	packages = imp.AllPackages()
	sort.Sort(byPath(packages))

	registerHandlers()

	srv := &http.Server{Addr: *httpAddr}
	l, err := net.Listen("tcp", srv.Addr)
	if err != nil {
		exitError(err)
	}
	if *open {
		url := fmt.Sprintf("http://localhost%s/", *httpAddr)
		if !startBrowser(url) {
			fmt.Println(url)
		}
	}
	exitError(srv.Serve(l))
}

func registerHandlers() {
	http.HandleFunc("/", serveIndex)
	http.HandleFunc("/source", serveSource)
	http.HandleFunc("/file", serveFile)
	http.HandleFunc("/query", serveQuery)
	staticPrefix := "/static/"
	http.Handle(staticPrefix, http.StripPrefix(staticPrefix, http.HandlerFunc(serveStatic)))
}

// byPath makes a slice of package infos sortable by package path.
type byPath []*importer.PackageInfo

func (p byPath) Len() int           { return len(p) }
func (p byPath) Less(i, j int) bool { return p[i].Pkg.Path() < p[j].Pkg.Path() }
func (p byPath) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }

// scopeFiles returns a new slice containing the full paths of all the files
// imported by an importer, sorted in increasing order.
func scopeFiles(imp *importer.Importer) []string {
	files := make([]string, 0)
	imp.Fset.Iterate(func(f *token.File) bool {
		files = append(files, f.Name())
		return true
	})
	sort.Strings(files)
	return files
}

// startBrowser tries to open the URL in a browser
// and reports whether it succeeds.
func startBrowser(url string) bool {
	// try to start the browser
	var args []string
	switch runtime.GOOS {
	case "darwin":
		args = []string{"open"}
	case "windows":
		args = []string{"cmd", "/c", "start"}
	default:
		args = []string{"xdg-open"}
	}
	cmd := exec.Command(args[0], append(args[1:], url)...)
	return cmd.Start() == nil
}

func exitError(err error) {
	fmt.Fprintln(os.Stderr, err)
	os.Exit(1)
}

// cmdLine returns what the command line would look like if the oracle was
// invoked via command line with the given arguments.
func cmdLine(mode, pos, format string, scope []string) string {
	return fmt.Sprintf("oracle -pos=%s -format=%s %s %s",
		pos, format, mode, strings.Join(scope, " "))
}
<commit_msg>Add support for build tags<commit_after>// Copyright 2013 Frederik Zipp.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Pythia is a web application front-end for the Go source code oracle.
package main

import (
	"flag"
	"fmt"
	"go/build"
	"go/token"
	"net"
	"net/http"
	"os"
	"os/exec"
	"runtime"
	"sort"
	"strings"
	"sync"

	"code.google.com/p/go.tools/importer"
	"code.google.com/p/go.tools/oracle"
)

var (
	httpAddr = flag.String("http", ":8080", "HTTP listen address")
	verbose  = flag.Bool("v", false, "Verbose mode: print incoming queries")
	open     = flag.Bool("open", true, "Try to open browser")
	tags     = flag.String("tags", "", "Tags to use when importing packages")
	args     []string
	files    []string
	packages []*importer.PackageInfo
	imp      *importer.Importer
	ora      *oracle.Oracle
	mutex    sync.Mutex
)

func init() {
	if os.Getenv("GOMAXPROCS") == "" {
		n := runtime.NumCPU()
		if n < 4 {
			n = 4
		}
		runtime.GOMAXPROCS(n)
	}
}

const useHelp = "Run 'pythia -help' for more information.\n"

const helpMessage = `Web frontend for the Go source code oracle.
Usage: pythia [<flag> ...] <args> ...

The -http flag specifies the HTTP service address (e.g., ':6060').

The -tags flag specifies comma separated tags to use when importing
code (e.g., 'foo,!darwin').

The -open flag determines, whether the application should try to
open the browser. It is set to 'true' by default. If set to 'false'
the browser will not be launched.

The -v flag enables verbose mode, in which every incoming query
to the oracle is logged to the standard output.

Examples:

Start pythia with the scope of package oracle:
% pythia code.google.com/p/go.tools/cmd/oracle

Start pythia with the scope of package image/png on port 8081,
but don't open the browser:
% pythia -http=:8081 -open=false image/png
` + importer.InitialPackagesUsage

func main() {
	flag.Usage = func() { fmt.Fprint(os.Stderr, useHelp) }
	flag.CommandLine.Init(os.Args[0], flag.ContinueOnError)
	if err := flag.CommandLine.Parse(os.Args[1:]); err != nil {
		if err == flag.ErrHelp {
			fmt.Println(helpMessage)
			fmt.Println("Flags:")
			flag.PrintDefaults()
		}
		os.Exit(2)
	}
	args = flag.Args()
	if len(args) == 0 {
		fmt.Fprint(os.Stderr, "Error: no package arguments.\n"+useHelp)
		os.Exit(2)
	}

	var err error
	settings := build.Default
	settings.BuildTags = strings.Split(*tags, ",")
	imp = importer.New(&importer.Config{Build: &settings})
	ora, err = oracle.New(imp, args, nil, false)
	if err != nil {
		exitError(err)
	}
	files = scopeFiles(imp)
	packages = imp.AllPackages()
	sort.Sort(byPath(packages))

	registerHandlers()

	srv := &http.Server{Addr: *httpAddr}
	l, err := net.Listen("tcp", srv.Addr)
	if err != nil {
		exitError(err)
	}
	if *open {
		url := fmt.Sprintf("http://localhost%s/", *httpAddr)
		if !startBrowser(url) {
			fmt.Println(url)
		}
	}
	exitError(srv.Serve(l))
}

func registerHandlers() {
	http.HandleFunc("/", serveIndex)
	http.HandleFunc("/source", serveSource)
	http.HandleFunc("/file", serveFile)
	http.HandleFunc("/query", serveQuery)
	staticPrefix := "/static/"
	http.Handle(staticPrefix, http.StripPrefix(staticPrefix, http.HandlerFunc(serveStatic)))
}

// byPath makes a slice of package infos sortable by package path.
type byPath []*importer.PackageInfo

func (p byPath) Len() int           { return len(p) }
func (p byPath) Less(i, j int) bool { return p[i].Pkg.Path() < p[j].Pkg.Path() }
func (p byPath) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }

// scopeFiles returns a new slice containing the full paths of all the files
// imported by an importer, sorted in increasing order.
func scopeFiles(imp *importer.Importer) []string {
	files := make([]string, 0)
	imp.Fset.Iterate(func(f *token.File) bool {
		files = append(files, f.Name())
		return true
	})
	sort.Strings(files)
	return files
}

// startBrowser tries to open the URL in a browser
// and reports whether it succeeds.
func startBrowser(url string) bool {
	// try to start the browser
	var args []string
	switch runtime.GOOS {
	case "darwin":
		args = []string{"open"}
	case "windows":
		args = []string{"cmd", "/c", "start"}
	default:
		args = []string{"xdg-open"}
	}
	cmd := exec.Command(args[0], append(args[1:], url)...)
	return cmd.Start() == nil
}

func exitError(err error) {
	fmt.Fprintln(os.Stderr, err)
	os.Exit(1)
}

// cmdLine returns what the command line would look like if the oracle was
// invoked via command line with the given arguments.
func cmdLine(mode, pos, format string, scope []string) string {
	return fmt.Sprintf("oracle -pos=%s -format=%s %s %s",
		pos, format, mode, strings.Join(scope, " "))
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"
	"os/user"
	"path/filepath"

	"time"

	"strings"

	"github.com/claudetech/loggo"
	. "github.com/claudetech/loggo/default"
)

func main() {
	// get the users home dir
	user, err := user.Current()
	if nil != err {
		panic(fmt.Sprintf("Could not read users homedir %v\n", err))
	}

	// parse the command line arguments
	argLogLevel := flag.Int("log-level", 0, "Set the log level (0 = error, 1 = warn, 2 = info, 3 = debug, 4 = trace)")
	argConfigPath := flag.String("config", filepath.Join(user.HomeDir, ".plexdrive"), "The path to the configuration directory")
	argTempPath := flag.String("temp", os.TempDir(), "Path to a temporary directory to store temporary data")
	argChunkSize := flag.Int64("chunk-size", 5*1024*1024, "The size of each chunk that is downloaded (in kb)")
	argRefreshInterval := flag.Duration("refresh-interval", 5*time.Minute, "The number of minutes to wait till checking for changes")
	argClearInterval := flag.Duration("clear-chunk-interval", 1*time.Minute, "The number of minutes to wait till clearing the chunk directory")
	argMountOptions := flag.String("fuse-options", "", "Fuse mount options (e.g. -fuse-options allow_other,...)")
	flag.Parse()

	// check if mountpoint is specified
	argMountPoint := flag.Arg(0)
	if "" == argMountPoint {
		flag.Usage()
		panic(fmt.Errorf("Mountpoint not specified"))
	}

	// parse the mount options
	var mountOptions []string
	if "" != *argMountOptions {
		mountOptions = strings.Split(*argMountOptions, ",")
	}

	// initialize the logger with the specific log level
	var logLevel loggo.Level
	switch *argLogLevel {
	case 0:
		logLevel = loggo.Error
	case 1:
		logLevel = loggo.Warning
	case 2:
		logLevel = loggo.Info
	case 3:
		logLevel = loggo.Debug
	case 4:
		logLevel = loggo.Trace
	default:
		logLevel = loggo.Warning
	}
	Log.SetLevel(logLevel)

	// debug all given parameters
	Log.Debugf("log-level            : %v", logLevel)
	Log.Debugf("config               : %v", *argConfigPath)
	Log.Debugf("temp                 : %v", *argTempPath)
	Log.Debugf("chunk-size           : %v", *argChunkSize)
	Log.Debugf("refresh-interval     : %v", *argRefreshInterval)
	Log.Debugf("clear-chunk-interval : %v", *argClearInterval)
	Log.Debugf("fuse-options         : %v", *argMountOptions)

	// create all directories
	if err := os.MkdirAll(*argConfigPath, 0766); nil != err {
		Log.Errorf("Could not create configuration directory")
		Log.Debugf("%v", err)
		os.Exit(1)
	}
	chunkPath := filepath.Join(*argTempPath, "chunks")
	if err := os.MkdirAll(chunkPath, 0777); nil != err {
		Log.Errorf("Could not create temp chunk directory")
		Log.Debugf("%v", err)
		os.Exit(2)
	}

	// set the global buffer configuration
	SetChunkPath(chunkPath)
	SetChunkSize(*argChunkSize)

	// read the configuration
	config, err := ReadConfig(filepath.Join(*argConfigPath, "config.json"))
	if nil != err {
		Log.Errorf("Could not read configuration")
		Log.Debugf("%v", err)
		os.Exit(3)
	}

	cache, err := NewCache(filepath.Join(*argConfigPath, "cache"))
	if nil != err {
		Log.Errorf("Could not initialize cache")
		Log.Debugf("%v", err)
		os.Exit(4)
	}
	defer cache.Close()

	drive, err := NewDriveClient(config, cache, *argRefreshInterval)
	if nil != err {
		Log.Errorf("Could not initialize Google Drive Client")
		Log.Debugf("%v", err)
		os.Exit(5)
	}

	go CleanChunkDir(chunkPath, *argClearInterval)
	if err := Mount(drive, argMountPoint, mountOptions); nil != err {
		Log.Debugf("%v", err)
		os.Exit(6)
	}
}
<commit_msg>fixed wrong calculation in documentation<commit_after>package main

import (
	"flag"
	"fmt"
	"os"
	"os/user"
	"path/filepath"

	"time"

	"strings"

	"github.com/claudetech/loggo"
	. "github.com/claudetech/loggo/default"
)

func main() {
	// get the users home dir
	user, err := user.Current()
	if nil != err {
		panic(fmt.Sprintf("Could not read users homedir %v\n", err))
	}

	// parse the command line arguments
	argLogLevel := flag.Int("log-level", 0, "Set the log level (0 = error, 1 = warn, 2 = info, 3 = debug, 4 = trace)")
	argConfigPath := flag.String("config", filepath.Join(user.HomeDir, ".plexdrive"), "The path to the configuration directory")
	argTempPath := flag.String("temp", os.TempDir(), "Path to a temporary directory to store temporary data")
	argChunkSize := flag.Int64("chunk-size", 5*1024*1024, "The size of each chunk that is downloaded (in byte)")
	argRefreshInterval := flag.Duration("refresh-interval", 5*time.Minute, "The number of minutes to wait till checking for changes")
	argClearInterval := flag.Duration("clear-chunk-interval", 1*time.Minute, "The number of minutes to wait till clearing the chunk directory")
	argMountOptions := flag.String("fuse-options", "", "Fuse mount options (e.g. -fuse-options allow_other,...)")
	flag.Parse()

	// check if mountpoint is specified
	argMountPoint := flag.Arg(0)
	if "" == argMountPoint {
		flag.Usage()
		panic(fmt.Errorf("Mountpoint not specified"))
	}

	// parse the mount options
	var mountOptions []string
	if "" != *argMountOptions {
		mountOptions = strings.Split(*argMountOptions, ",")
	}

	// initialize the logger with the specific log level
	var logLevel loggo.Level
	switch *argLogLevel {
	case 0:
		logLevel = loggo.Error
	case 1:
		logLevel = loggo.Warning
	case 2:
		logLevel = loggo.Info
	case 3:
		logLevel = loggo.Debug
	case 4:
		logLevel = loggo.Trace
	default:
		logLevel = loggo.Warning
	}
	Log.SetLevel(logLevel)

	// debug all given parameters
	Log.Debugf("log-level            : %v", logLevel)
	Log.Debugf("config               : %v", *argConfigPath)
	Log.Debugf("temp                 : %v", *argTempPath)
	Log.Debugf("chunk-size           : %v", *argChunkSize)
	Log.Debugf("refresh-interval     : %v", *argRefreshInterval)
	Log.Debugf("clear-chunk-interval : %v", *argClearInterval)
	Log.Debugf("fuse-options         : %v", *argMountOptions)

	// create all directories
	if err := os.MkdirAll(*argConfigPath, 0766); nil != err {
		Log.Errorf("Could not create configuration directory")
		Log.Debugf("%v", err)
		os.Exit(1)
	}
	chunkPath := filepath.Join(*argTempPath, "chunks")
	if err := os.MkdirAll(chunkPath, 0777); nil != err {
		Log.Errorf("Could not create temp chunk directory")
		Log.Debugf("%v", err)
		os.Exit(2)
	}

	// set the global buffer configuration
	SetChunkPath(chunkPath)
	SetChunkSize(*argChunkSize)

	// read the configuration
	config, err := ReadConfig(filepath.Join(*argConfigPath, "config.json"))
	if nil != err {
		Log.Errorf("Could not read configuration")
		Log.Debugf("%v", err)
		os.Exit(3)
	}

	cache, err := NewCache(filepath.Join(*argConfigPath, "cache"))
	if nil != err {
		Log.Errorf("Could not initialize cache")
		Log.Debugf("%v", err)
		os.Exit(4)
	}
	defer cache.Close()

	drive, err := NewDriveClient(config, cache, *argRefreshInterval)
	if nil != err {
		Log.Errorf("Could not initialize Google Drive Client")
		Log.Debugf("%v", err)
		os.Exit(5)
	}

	go CleanChunkDir(chunkPath, *argClearInterval)
	if err := Mount(drive, argMountPoint, mountOptions); nil != err {
		Log.Debugf("%v", err)
		os.Exit(6)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"flag"
	"log"
	"net/http"
	"os"
	"os/exec"

	"github.com/stuartnelson3/guac"
)

func main() {
	var (
		port = flag.String("port", "8080", "port to listen on")
	)
	flag.Parse()

	http.HandleFunc("/script.js", func(w http.ResponseWriter, r *http.Request) {
		f, err := os.Open("dist/script.js")
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		defer f.Close()

		fs, err := f.Stat()
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		http.ServeContent(w, r, "script.js", fs.ModTime(), f)
	})

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		f, err := os.Open("index.html")
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		defer f.Close()

		fs, err := f.Stat()
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		http.ServeContent(w, r, "index", fs.ModTime(), f)
	})

	// Recompile the elm code whenever a change is detected.
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	recompileFn := func() error {
		cmd := exec.Command("elm-make", "src/Main.elm", "--output", "dist/script.js")
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		return cmd.Run()
	}

	watcher, err := guac.NewWatcher(ctx, "./src", recompileFn)
	if err != nil {
		log.Fatalf("error watching: %v", err)
	}
	go watcher.Run()

	log.Printf("starting listener on port %s", *port)
	if err := http.ListenAndServe(":"+*port, nil); err != nil {
		log.Fatal(err)
	}
}
<commit_msg>Update API<commit_after>package main

import (
	"context"
	"flag"
	"log"
	"net/http"
	"os"
	"os/exec"

	"github.com/stuartnelson3/guac"
)

func main() {
	var (
		port = flag.String("port", "8080", "port to listen on")
	)
	flag.Parse()

	http.HandleFunc("/script.js", func(w http.ResponseWriter, r *http.Request) {
		http.ServeFile(w, r, "dist/script.js")
	})

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		http.ServeFile(w, r, "index.html")
	})

	http.HandleFunc("/api/v1/silences", func(w http.ResponseWriter, r *http.Request) {
		http.ServeFile(w, r, "silences.json")
	})

	// Recompile the elm code whenever a change is detected.
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	recompileFn := func() error {
		cmd := exec.Command("elm-make", "src/Main.elm", "--output", "dist/script.js")
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		return cmd.Run()
	}

	watcher, err := guac.NewWatcher(ctx, "./src", recompileFn)
	if err != nil {
		log.Fatalf("error watching: %v", err)
	}
	go watcher.Run()

	log.Printf("starting listener on port %s", *port)
	if err := http.ListenAndServe(":"+*port, nil); err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"runtime/pprof"

	log "github.com/Sirupsen/logrus"
)

//GitHash set by ld flags at build time
var GitHash = ""

//Version set by git tag
var Version = ""

//BuildDate set by ld flags at build time
var BuildDate = ""

//AppName Application name
var AppName = "Nitro TFTP"

var debugFlag bool
var configFile string
var configString string
var versionFlag bool
var genconfigFlag bool
var cpuprofile string
var memprofile string

func init() {
	cfgUsage := "Configuration file"
	flag.StringVar(&configFile, "config", "", cfgUsage)
	flag.StringVar(&configFile, "c", "", cfgUsage+" (shorthand)")
	cfgstrUsage := "Configuration string in JSON format"
	flag.StringVar(&configString, "config-string", "", cfgstrUsage)
	flag.StringVar(&configString, "cs", "", cfgstrUsage+" (shorthand)")
	debugUsage := "Enable debug logging"
	flag.BoolVar(&debugFlag, "debug", false, debugUsage)
	flag.BoolVar(&debugFlag, "d", false, debugUsage+" (shorthand)")
	verUsage := "Display Version"
	flag.BoolVar(&versionFlag, "version", false, verUsage)
	flag.BoolVar(&versionFlag, "v", false, verUsage+" (shorthand)")
	gencfgUsage := "Generate example config"
	flag.BoolVar(&genconfigFlag, "gencfg", false, gencfgUsage)
	flag.BoolVar(&genconfigFlag, "g", false, gencfgUsage+" (shorthand)")
	//flag.StringVar(&cpuprofile, "cpuprofile", "", "write cpu profile to file")
	//flag.StringVar(&memprofile, "memprofile", "", "write memory profile to file")
}

func main() {
	flag.Parse()

	// add handlers to help us track memory usage - they don't track memory until they're told to
	//profiler.AddMemoryProfilingHandlers()

	// Uncomment if you want to start profiling automatically
	// profiler.StartProfiling()

	// Enable pprof
	if cpuprofile != "" {
		f, err := os.Create(cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	if memprofile != "" {
		f, err := os.Create(memprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.WriteHeapProfile(f)
		defer f.Close()
	}

	if debugFlag {
		log.SetLevel(log.DebugLevel)
	}

	// listen on port 6060 (pick a port)
	//go http.ListenAndServe(":6060", nil)

	if versionFlag {
		fmt.Printf("Built: %s\nVersion: %s\nGit Commit: %s\n", BuildDate, Version, GitHash)
		os.Exit(0)
	}

	if genconfigFlag {
		cfg := NewConfig()
		c, err := json.MarshalIndent(cfg, "", "\t")
		if err != nil {
			log.Error(err)
			os.Exit(1)
		}
		fmt.Println(string(c))
		os.Exit(0)
	}

	s := &TFTPServer{Debug: debugFlag}

	cfg := NewConfig()

	if configFile != "" && configString == "" {
		err := cfg.Open(configFile)
		if err != nil {
			log.Error(err)
			os.Exit(1)
		}
	} else if configFile == "" && configString != "" {
		err := cfg.StringParse(configString)
		if err != nil {
			log.Error(err)
			os.Exit(1)
		}
	}

	s.LoadConfig(cfg)
	ctrlChan := s.Listen()

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt, os.Kill)
	sig := <-sigChan
	close(sigChan)
	close(ctrlChan)
	log.Println("Caught Signal", sig)
	os.Exit(0)
}
<commit_msg>added basic support for consul registration<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"math/rand"
	"net"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"runtime/pprof"
	"strconv"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/hashicorp/consul/api"
)

//GitHash set by ld flags at build time
var GitHash = ""

//Version set by git tag
var Version = ""

//BuildDate set by ld flags at build time
var BuildDate = ""

//AppName Application name
var AppName = "Nitro TFTP"

var debugFlag bool
var configFile string
var configString string
var consulConfig string
var versionFlag bool
var genconfigFlag bool
var cpuprofile string
var memprofile string

func init() {
	cfgUsage := "Configuration file"
	flag.StringVar(&configFile, "config", "", cfgUsage)
	flag.StringVar(&configFile, "c", "", cfgUsage+" (shorthand)")
	cfgstrUsage := "Configuration string in JSON format"
	flag.StringVar(&configString, "config-string", "", cfgstrUsage)
	flag.StringVar(&configString, "cs", "", cfgstrUsage+" (shorthand)")
	debugUsage := "Enable debug logging"
	flag.BoolVar(&debugFlag, "debug", false, debugUsage)
	flag.BoolVar(&debugFlag, "d", false, debugUsage+" (shorthand)")
	verUsage := "Display Version"
	flag.BoolVar(&versionFlag, "version", false, verUsage)
	flag.BoolVar(&versionFlag, "v", false, verUsage+" (shorthand)")
	gencfgUsage := "Generate example config"
	flag.BoolVar(&genconfigFlag, "gencfg", false, gencfgUsage)
	flag.BoolVar(&genconfigFlag, "g", false, gencfgUsage+" (shorthand)")
	//flag.StringVar(&cpuprofile, "cpuprofile", "", "write cpu profile to file")
	//flag.StringVar(&memprofile, "memprofile", "", "write memory profile to file")
}

func main() {
	flag.Parse()

	// add handlers to help us track memory usage - they don't track memory until they're told to
	//profiler.AddMemoryProfilingHandlers()

	// Uncomment if you want to start profiling automatically
	// profiler.StartProfiling()

	// Enable pprof
	if cpuprofile != "" {
		f, err := os.Create(cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	if memprofile != "" {
		f, err := os.Create(memprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.WriteHeapProfile(f)
		defer f.Close()
	}

	if debugFlag {
		log.SetLevel(log.DebugLevel)
	}

	// listen on port 6060 (pick a port)
	//go http.ListenAndServe(":6060", nil)

	if versionFlag {
		fmt.Printf("Built: %s\nVersion: %s\nGit Commit: %s\n", BuildDate, Version, GitHash)
		os.Exit(0)
	}

	if genconfigFlag {
		cfg := NewConfig()
		c, err := json.MarshalIndent(cfg, "", "\t")
		if err != nil {
			log.Error(err)
			os.Exit(1)
		}
		fmt.Println(string(c))
		os.Exit(0)
	}

	s := &TFTPServer{Debug: debugFlag}

	cfg := NewConfig()

	if configFile != "" && configString == "" {
		err := cfg.Open(configFile)
		if err != nil {
			log.Error(err)
			os.Exit(1)
		}
	} else if configFile == "" && configString != "" {
		err := cfg.StringParse(configString)
		if err != nil {
			log.Error(err)
			os.Exit(1)
		}
	}

	s.LoadConfig(cfg)
	ctrlChan := s.Listen()

	//register as a consul service
	serviceid := strings.Join([]string{AppName, strconv.Itoa(int(time.Now().Unix() * rand.Int63()))}, "_")
	if consulConfig != "" {
		client, err := api.NewClient(api.DefaultConfig())
		if err != nil {
			log.Error(err)
			os.Exit(1)
		}
		agent := client.Agent()
		host, port, err := net.SplitHostPort(s.sock.LocalAddr().String())
		if err != nil {
			log.Error(err)
			os.Exit(1)
		}
		iport, err := strconv.Atoi(port)
		if err != nil {
			log.Error(err)
			os.Exit(1)
		}
		service := &api.AgentServiceRegistration{ID: serviceid, Name: "tftp", Address: host, Port: iport}
		agent.ServiceRegister(service)
	}

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt, os.Kill)
	<-sigChan
	close(sigChan)
	close(ctrlChan)

	if consulConfig != "" {
		client, err := api.NewClient(api.DefaultConfig())
		if err != nil {
			log.Error(err)
			os.Exit(1)
		}
		agent := client.Agent()
		agent.ServiceDeregister(serviceid)
	}
	log.Println("Shutting down $s", AppName)
	os.Exit(0)
}
<|endoftext|>


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

package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"regexp"
	"strings"
	"time"

	"github.com/yarpc/yab/encoding"
	"github.com/yarpc/yab/transport"

	"github.com/casimir/xdg-go"
	"github.com/jessevdk/go-flags"
	"github.com/opentracing/opentracing-go"
	opentracing_ext "github.com/opentracing/opentracing-go/ext"
	"github.com/uber/jaeger-client-go"
	jaeger_config "github.com/uber/jaeger-client-go/config"
	"github.com/uber/tchannel-go"
)

var errHealthAndMethod = errors.New("cannot specify method name and use --health")

func findGroup(parser *flags.Parser, group string) *flags.Group {
	if g := parser.Group.Find(group); g != nil {
		return g
	}

	panic("no group called " + group + " found.")
}

func setGroupDescs(parser *flags.Parser, groupName, shortDesc, longDesc string) {
	g := findGroup(parser, groupName)
	g.ShortDescription = shortDesc
	g.LongDescription = longDesc
}

func fromPositional(args []string, index int, s *string) bool {
	if len(args) <= index {
		return false
	}

	if args[index] != "" {
		*s = args[index]
	}

	return true
}

func main() {
	log.SetFlags(0)
	parseAndRun(consoleOutput{os.Stdout})
}

var errExit = errors.New("sentinel error used to exit cleanly")

func toGroff(s string) string {
	// Expand tabbed lines beginning with "-" as items in a bullet list.
	s = strings.Replace(s, "\n\t* ", "\n.IP \\[bu]\n", -1 /* all occurences */)

	// Two newlines start a new paragraph.
	s = strings.Replace(s, "\n\n", "\n.PP\n", -1)

	// Lines beginning with a tab are interpreted as example code.
	//
	// See http://liw.fi/manpages/ for an explanation of these
	// commands -- tl;dr: turn of paragraph filling and indent the
	// block one level.
	indentRegexp := regexp.MustCompile(`\t(.*)\n`)
	s = indentRegexp.ReplaceAllString(s, ".nf\n.RS\n$1\n.RE\n.fi\n")

	return s
}

func newParser() (*flags.Parser, *Options) {
	opts := newOptions()
	return flags.NewParser(opts, flags.HelpFlag|flags.PassDoubleDash), opts
}

func getOptions(args []string, out output) (*Options, error) {
	parser, opts := newParser()
	parser.Usage = "[<service> <method> <body>] [OPTIONS]"
	parser.ShortDescription = "yet another benchmarker"
	parser.LongDescription = `
yab is a benchmarking tool for TChannel and HTTP applications. It's primarily intended for Thrift applications but supports other encodings like JSON and binary (raw). It can be used in a curl-like fashion when benchmarking features are disabled.
`

	// Read defaults if they're available, before we change the group names.
	if err := parseDefaultConfigs(parser); err != nil {
		return nil, fmt.Errorf("error reading defaults: %v", err)
	}

	overrideDefaults(opts, args)

	setGroupDescs(parser, "request", "Request Options", toGroff(_reqOptsDesc))
	setGroupDescs(parser, "transport", "Transport Options", toGroff(_transportOptsDesc))
	setGroupDescs(parser, "benchmark", "Benchmark Options", toGroff(_benchmarkOptsDesc))

	// If there are no arguments specified, write the help.
	if len(args) == 0 {
		parser.WriteHelp(out)
		return opts, errExit
	}

	remaining, err := parser.ParseArgs(args)
	if err != nil {
		if ferr, ok := err.(*flags.Error); ok {
			if ferr.Type == flags.ErrHelp {
				parser.WriteHelp(out)
				return opts, errExit
			}
		}
		return opts, err
	}
	setEncodingOptions(opts)

	if opts.DisplayVersion {
		out.Printf("yab version %v\n", versionString)
		return opts, errExit
	}

	if opts.ManPage {
		parser.LongDescription += `
Default options can be specified in a ~/.config/yab/defaults.ini file (or ~/Library/Preferences/yab/defaults.ini on Mac) with contents similar to this:

	[request]
	timeout = 2s

	[transport]
	peer-list = "/path/to/peer/list.json"

	[benchmark]
	warmup = 10
`
		parser.LongDescription = toGroff(parser.LongDescription)
		parser.WriteManPage(out)
		return opts, errExit
	}

	fromPositional(remaining, 0, &opts.TOpts.ServiceName)
	fromPositional(remaining, 1, &opts.ROpts.MethodName)

	// We support both:
	// [service] [method] [request]
	// [service] [method] [headers] [request]
	if fromPositional(remaining, 3, &opts.ROpts.RequestJSON) {
		fromPositional(remaining, 2, &opts.ROpts.HeadersJSON)
	} else {
		fromPositional(remaining, 2, &opts.ROpts.RequestJSON)
	}

	return opts, nil
}

// parseAndRun is like main, but uses the given output.
func parseAndRun(out output) {
	opts, err := getOptions(os.Args[1:], out)
	if err != nil {
		if err == errExit {
			return
		}
		out.Fatalf("Failed to parse options: %v", err)
	}
	runWithOptions(*opts, out)
}

// overrideDefaults clears fields in the default options that may
// clash with user-specified options.
// E.g., if the defaults has a peer list file, and the user has specifed
// a peer through the command line, then the final options should only
// contain the peer specified in the args.
func overrideDefaults(defaults *Options, args []string) {
	argsParser, argsOnly := newParser()
	argsParser.ParseArgs(args)

	// Clear default peers if the user has specified peer options in args.
	if len(argsOnly.TOpts.HostPorts) > 0 {
		defaults.TOpts.HostPortFile = ""
	}
	if len(argsOnly.TOpts.HostPortFile) > 0 {
		defaults.TOpts.HostPorts = nil
	}
}

// findBestConfigFile finds the best config file to use. An empty string will be
// returned if no config file should be used.
func findBestConfigFile() string {
	app := xdg.App{Name: "yab"}

	// Find the best config path to use, preferring the user's config path and
	// falling back to the system config path.
	configPaths := []string{app.ConfigPath("defaults.ini")}
	configPaths = append(configPaths, app.SystemConfigPaths("defaults.ini")...)
	var configFile string
	for _, path := range configPaths {
		if _, err := os.Stat(path); err == nil {
			configFile = path
			break
		}
	}
	return configFile
}

// parseDefaultConfigs reads defaults from ~/.config/yab/defaults.ini if they're
// available.
func parseDefaultConfigs(parser *flags.Parser) error {
	configFile := findBestConfigFile()
	if configFile == "" {
		return nil // no defaults file was found
	}

	iniParser := flags.NewIniParser(parser)
	if err := iniParser.ParseFile(configFile); err != nil {
		return fmt.Errorf("couldn't read %v: %v", configFile, err)
	}

	return nil
}

func runWithOptions(opts Options, out output) {
	if opts.ROpts.YamlTemplate != "" {
		if err := readYamlRequest(&opts); err != nil {
			out.Fatalf("Failed while reading yaml template: %v\n", err)
		}
	}

	reqInput, err := getRequestInput(opts.ROpts.RequestJSON, opts.ROpts.RequestFile)
	if err != nil {
		out.Fatalf("Failed while loading body input: %v\n", err)
	}

	headers, err := getHeaders(opts.ROpts.HeadersJSON, opts.ROpts.HeadersFile, opts.ROpts.Headers)
	if err != nil {
		out.Fatalf("Failed while loading headers input: %v\n", err)
	}

	serializer, err := NewSerializer(opts.ROpts)
	if err != nil {
		out.Fatalf("Failed while parsing input: %v\n", err)
	}

	if opts.TOpts.CallerName != "" {
		if opts.BOpts.enabled() {
			out.Fatalf("Cannot override caller name when running benchmarks\n")
		}
	} else {
		opts.TOpts.CallerName = "yab-" + os.Getenv("USER")
	}

	tracer, closer := getTracer(opts, out)
	if closer != nil {
		defer closer.Close()
	}

	// transport abstracts the underlying wire protocol used to make the call.
	transport, err := getTransport(opts.TOpts, serializer.Encoding(), tracer)
	if err != nil {
		out.Fatalf("Failed while parsing options: %v\n", err)
	}

	serializer = withTransportSerializer(transport.Protocol(), serializer, opts.ROpts)

	// req is the transport.Request that will be used to make a call.
	req, err := serializer.Request(reqInput)
	if err != nil {
		out.Fatalf("Failed while parsing request input: %v\n", err)
	}

	req.Headers = headers
	req.TransportHeaders = opts.TOpts.TransportHeaders
	req.Timeout = opts.ROpts.Timeout.Duration()
	if req.Timeout == 0 {
		req.Timeout = time.Second
	}
	req.Baggage = opts.ROpts.Baggage

	// Only make the request if the user hasn't specified 0 warmup.
	if !(opts.BOpts.enabled() && opts.BOpts.WarmupRequests == 0) {
		makeInitialRequest(out, transport, serializer, req)
	}

	runBenchmark(out, opts, benchmarkMethod{
		serializer: serializer,
		req:        req,
	})
}

type noEnveloper interface {
	WithoutEnvelopes() encoding.Serializer
}

func getTracer(opts Options, out output) (opentracing.Tracer, io.Closer) {
	var (
		tracer opentracing.Tracer = opentracing.NoopTracer{}
		closer io.Closer
	)
	if opts.TOpts.Jaeger {
		var jaegerConfig jaeger_config.Configuration
		var err error
		tracer, closer, err = jaegerConfig.New(opts.TOpts.CallerName, jaeger.NullStatsReporter)
		if err != nil {
			out.Fatalf("Failed to create Jaeger tracer: %v\n", err)
		}
	} else if len(opts.ROpts.Baggage) > 0 {
		out.Fatalf("To propagate baggage, you must opt-into a tracing client, i.e., --jaeger")
	}
	return tracer, closer
}

// withTransportSerializer may modify the serializer for the transport used.
// E.g. Thrift payloads are not enveloped when used with TChannel.
func withTransportSerializer(p transport.Protocol, s encoding.Serializer, rOpts RequestOptions) encoding.Serializer {
	switch {
	case p == transport.TChannel && s.Encoding() == encoding.Thrift,
		rOpts.ThriftDisableEnvelopes:
		s = s.(noEnveloper).WithoutEnvelopes()
	}
	return s
}

// makeRequest makes a request using the given transport.
func makeRequest(t transport.Transport, request *transport.Request) (*transport.Response, error) {
	return makeRequestWithTracePriority(t, request, 0)
}

func makeRequestWithTracePriority(t transport.Transport, request *transport.Request, trace uint16) (*transport.Response, error) {
	ctx, cancel := tchannel.NewContext(request.Timeout)
	defer cancel()

	if tracer := t.Tracer(); tracer != nil {
		span := tracer.StartSpan(request.Method)
		opentracing_ext.SamplingPriority.Set(span, trace)
		for k, v := range request.Baggage {
			span = span.SetBaggageItem(k, v)
		}
		ctx = opentracing.ContextWithSpan(ctx, span)
	}

	return t.Call(ctx, request)
}

func makeInitialRequest(out output, transport transport.Transport, serializer encoding.Serializer, req *transport.Request) {
	response, err := makeRequestWithTracePriority(transport, req, 1)
	if err != nil {
		out.Fatalf("Failed while making call: %v\n", err)
	}

	// responseMap converts the Thrift bytes response to a map.
	responseMap, err := serializer.Response(response)
	if err != nil {
		out.Fatalf("Failed while parsing response: %v\n", err)
	}

	// Print the initial output body.
	outSerialized := map[string]interface{}{
		"body": responseMap,
	}
	if len(response.Headers) > 0 {
		outSerialized["headers"] = response.Headers
	}
	for k, v := range response.TransportFields {
		outSerialized[k] = v
	}
	bs, err := json.MarshalIndent(outSerialized, "", "  ")
	if err != nil {
		out.Fatalf("Failed to convert map to JSON: %v\nMap: %+v\n", err, responseMap)
	}
	out.Printf("%s\n\n", bs)
}
<commit_msg>Show defaults in help outputted when no args are specified<commit_after>// Copyright (c) 2016 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"regexp"
	"strings"
	"time"

	"github.com/yarpc/yab/encoding"
	"github.com/yarpc/yab/transport"

	"github.com/casimir/xdg-go"
	"github.com/jessevdk/go-flags"
	"github.com/opentracing/opentracing-go"
	opentracing_ext "github.com/opentracing/opentracing-go/ext"
	"github.com/uber/jaeger-client-go"
	jaeger_config "github.com/uber/jaeger-client-go/config"
	"github.com/uber/tchannel-go"
)

var errHealthAndMethod = errors.New("cannot specify method name and use --health")

func findGroup(parser *flags.Parser, group string) *flags.Group {
	if g := parser.Group.Find(group); g != nil {
		return g
	}

	panic("no group called " + group + " found.")
}

func setGroupDescs(parser *flags.Parser, groupName, shortDesc, longDesc string) {
	g := findGroup(parser, groupName)
	g.ShortDescription = shortDesc
	g.LongDescription = longDesc
}

func fromPositional(args []string, index int, s *string) bool {
	if len(args) <= index {
		return false
	}

	if args[index] != "" {
		*s = args[index]
	}

	return true
}

func main() {
	log.SetFlags(0)
	parseAndRun(consoleOutput{os.Stdout})
}

var errExit = errors.New("sentinel error used to exit cleanly")

func toGroff(s string) string {
	// Expand tabbed lines beginning with "-" as items in a bullet list.
	s = strings.Replace(s, "\n\t* ", "\n.IP \\[bu]\n", -1 /* all occurences */)

	// Two newlines start a new paragraph.
	s = strings.Replace(s, "\n\n", "\n.PP\n", -1)

	// Lines beginning with a tab are interpreted as example code.
	//
	// See http://liw.fi/manpages/ for an explanation of these
	// commands -- tl;dr: turn of paragraph filling and indent the
	// block one level.
	indentRegexp := regexp.MustCompile(`\t(.*)\n`)
	s = indentRegexp.ReplaceAllString(s, ".nf\n.RS\n$1\n.RE\n.fi\n")

	return s
}

func newParser() (*flags.Parser, *Options) {
	opts := newOptions()
	return flags.NewParser(opts, flags.HelpFlag|flags.PassDoubleDash), opts
}

func getOptions(args []string, out output) (*Options, error) {
	parser, opts := newParser()
	parser.Usage = "[<service> <method> <body>] [OPTIONS]"
	parser.ShortDescription = "yet another benchmarker"
	parser.LongDescription = `
yab is a benchmarking tool for TChannel and HTTP applications. It's primarily intended for Thrift applications but supports other encodings like JSON and binary (raw). It can be used in a curl-like fashion when benchmarking features are disabled.
`

	// Read defaults if they're available, before we change the group names.
	if err := parseDefaultConfigs(parser); err != nil {
		return nil, fmt.Errorf("error reading defaults: %v", err)
	}

	overrideDefaults(opts, args)

	setGroupDescs(parser, "request", "Request Options", toGroff(_reqOptsDesc))
	setGroupDescs(parser, "transport", "Transport Options", toGroff(_transportOptsDesc))
	setGroupDescs(parser, "benchmark", "Benchmark Options", toGroff(_benchmarkOptsDesc))

	remaining, err := parser.ParseArgs(args)
	// If there are no arguments specified, write the help.
	// We do this after Parse, otherwise the output doesn't show defaults.
	if len(args) == 0 {
		parser.WriteHelp(out)
		return opts, errExit
	}
	if err != nil {
		if ferr, ok := err.(*flags.Error); ok {
			if ferr.Type == flags.ErrHelp {
				parser.WriteHelp(out)
				return opts, errExit
			}
		}
		return opts, err
	}
	setEncodingOptions(opts)

	if opts.DisplayVersion {
		out.Printf("yab version %v\n", versionString)
		return opts, errExit
	}

	if opts.ManPage {
		parser.LongDescription += `
Default options can be specified in a ~/.config/yab/defaults.ini file (or ~/Library/Preferences/yab/defaults.ini on Mac) with contents similar to this:

	[request]
	timeout = 2s

	[transport]
	peer-list = "/path/to/peer/list.json"

	[benchmark]
	warmup = 10
`
		parser.LongDescription = toGroff(parser.LongDescription)
		parser.WriteManPage(out)
		return opts, errExit
	}

	fromPositional(remaining, 0, &opts.TOpts.ServiceName)
	fromPositional(remaining, 1, &opts.ROpts.MethodName)

	// We support both:
	// [service] [method] [request]
	// [service] [method] [headers] [request]
	if fromPositional(remaining, 3, &opts.ROpts.RequestJSON) {
		fromPositional(remaining, 2, &opts.ROpts.HeadersJSON)
	} else {
		fromPositional(remaining, 2, &opts.ROpts.RequestJSON)
	}

	return opts, nil
}

// parseAndRun is like main, but uses the given output.
func parseAndRun(out output) {
	opts, err := getOptions(os.Args[1:], out)
	if err != nil {
		if err == errExit {
			return
		}
		out.Fatalf("Failed to parse options: %v", err)
	}
	runWithOptions(*opts, out)
}

// overrideDefaults clears fields in the default options that may
// clash with user-specified options.
// E.g., if the defaults has a peer list file, and the user has specifed
// a peer through the command line, then the final options should only
// contain the peer specified in the args.
func overrideDefaults(defaults *Options, args []string) {
	argsParser, argsOnly := newParser()
	argsParser.ParseArgs(args)

	// Clear default peers if the user has specified peer options in args.
	if len(argsOnly.TOpts.HostPorts) > 0 {
		defaults.TOpts.HostPortFile = ""
	}
	if len(argsOnly.TOpts.HostPortFile) > 0 {
		defaults.TOpts.HostPorts = nil
	}
}

// findBestConfigFile finds the best config file to use. An empty string will be
// returned if no config file should be used.
func findBestConfigFile() string {
	app := xdg.App{Name: "yab"}

	// Find the best config path to use, preferring the user's config path and
	// falling back to the system config path.
	configPaths := []string{app.ConfigPath("defaults.ini")}
	configPaths = append(configPaths, app.SystemConfigPaths("defaults.ini")...)
	var configFile string
	for _, path := range configPaths {
		if _, err := os.Stat(path); err == nil {
			configFile = path
			break
		}
	}
	return configFile
}

// parseDefaultConfigs reads defaults from ~/.config/yab/defaults.ini if they're
// available.
func parseDefaultConfigs(parser *flags.Parser) error {
	configFile := findBestConfigFile()
	if configFile == "" {
		return nil // no defaults file was found
	}

	iniParser := flags.NewIniParser(parser)
	if err := iniParser.ParseFile(configFile); err != nil {
		return fmt.Errorf("couldn't read %v: %v", configFile, err)
	}

	return nil
}

func runWithOptions(opts Options, out output) {
	if opts.ROpts.YamlTemplate != "" {
		if err := readYamlRequest(&opts); err != nil {
			out.Fatalf("Failed while reading yaml template: %v\n", err)
		}
	}

	reqInput, err := getRequestInput(opts.ROpts.RequestJSON, opts.ROpts.RequestFile)
	if err != nil {
		out.Fatalf("Failed while loading body input: %v\n", err)
	}

	headers, err := getHeaders(opts.ROpts.HeadersJSON, opts.ROpts.HeadersFile, opts.ROpts.Headers)
	if err != nil {
		out.Fatalf("Failed while loading headers input: %v\n", err)
	}

	serializer, err := NewSerializer(opts.ROpts)
	if err != nil {
		out.Fatalf("Failed while parsing input: %v\n", err)
	}

	if opts.TOpts.CallerName != "" {
		if opts.BOpts.enabled() {
			out.Fatalf("Cannot override caller name when running benchmarks\n")
		}
	} else {
		opts.TOpts.CallerName = "yab-" + os.Getenv("USER")
	}

	tracer, closer := getTracer(opts, out)
	if closer != nil {
		defer closer.Close()
	}

	// transport abstracts the underlying wire protocol used to make the call.
	transport, err := getTransport(opts.TOpts, serializer.Encoding(), tracer)
	if err != nil {
		out.Fatalf("Failed while parsing options: %v\n", err)
	}

	serializer = withTransportSerializer(transport.Protocol(), serializer, opts.ROpts)

	// req is the transport.Request that will be used to make a call.
	req, err := serializer.Request(reqInput)
	if err != nil {
		out.Fatalf("Failed while parsing request input: %v\n", err)
	}

	req.Headers = headers
	req.TransportHeaders = opts.TOpts.TransportHeaders
	req.Timeout = opts.ROpts.Timeout.Duration()
	if req.Timeout == 0 {
		req.Timeout = time.Second
	}
	req.Baggage = opts.ROpts.Baggage

	// Only make the request if the user hasn't specified 0 warmup.
	if !(opts.BOpts.enabled() && opts.BOpts.WarmupRequests == 0) {
		makeInitialRequest(out, transport, serializer, req)
	}

	runBenchmark(out, opts, benchmarkMethod{
		serializer: serializer,
		req:        req,
	})
}

type noEnveloper interface {
	WithoutEnvelopes() encoding.Serializer
}

func getTracer(opts Options, out output) (opentracing.Tracer, io.Closer) {
	var (
		tracer opentracing.Tracer = opentracing.NoopTracer{}
		closer io.Closer
	)
	if opts.TOpts.Jaeger {
		var jaegerConfig jaeger_config.Configuration
		var err error
		tracer, closer, err = jaegerConfig.New(opts.TOpts.CallerName, jaeger.NullStatsReporter)
		if err != nil {
			out.Fatalf("Failed to create Jaeger tracer: %v\n", err)
		}
	} else if len(opts.ROpts.Baggage) > 0 {
		out.Fatalf("To propagate baggage, you must opt-into a tracing client, i.e., --jaeger")
	}
	return tracer, closer
}

// withTransportSerializer may modify the serializer for the transport used.
// E.g. Thrift payloads are not enveloped when used with TChannel.
func withTransportSerializer(p transport.Protocol, s encoding.Serializer, rOpts RequestOptions) encoding.Serializer {
	switch {
	case p == transport.TChannel && s.Encoding() == encoding.Thrift,
		rOpts.ThriftDisableEnvelopes:
		s = s.(noEnveloper).WithoutEnvelopes()
	}
	return s
}

// makeRequest makes a request using the given transport.
func makeRequest(t transport.Transport, request *transport.Request) (*transport.Response, error) {
	return makeRequestWithTracePriority(t, request, 0)
}

func makeRequestWithTracePriority(t transport.Transport, request *transport.Request, trace uint16) (*transport.Response, error) {
	ctx, cancel := tchannel.NewContext(request.Timeout)
	defer cancel()

	if tracer := t.Tracer(); tracer != nil {
		span := tracer.StartSpan(request.Method)
		opentracing_ext.SamplingPriority.Set(span, trace)
		for k, v := range request.Baggage {
			span = span.SetBaggageItem(k, v)
		}
		ctx = opentracing.ContextWithSpan(ctx, span)
	}

	return t.Call(ctx, request)
}

func makeInitialRequest(out output, transport transport.Transport, serializer encoding.Serializer, req *transport.Request) {
	response, err := makeRequestWithTracePriority(transport, req, 1)
	if err != nil {
		out.Fatalf("Failed while making call: %v\n", err)
	}

	// responseMap converts the Thrift bytes response to a map.
	responseMap, err := serializer.Response(response)
	if err != nil {
		out.Fatalf("Failed while parsing response: %v\n", err)
	}

	// Print the initial output body.
	outSerialized := map[string]interface{}{
		"body": responseMap,
	}
	if len(response.Headers) > 0 {
		outSerialized["headers"] = response.Headers
	}
	for k, v := range response.TransportFields {
		outSerialized[k] = v
	}
	bs, err := json.MarshalIndent(outSerialized, "", "  ")
	if err != nil {
		out.Fatalf("Failed to convert map to JSON: %v\nMap: %+v\n", err, responseMap)
	}
	out.Printf("%s\n\n", bs)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/gob"
	"net"
	"os"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/alecthomas/kingpin"
	metrics_graphite "github.com/cyberdelia/go-metrics-graphite"
	gocache "github.com/patrickmn/go-cache"

	"github.com/resourced/resourced-master/application"
	"github.com/resourced/resourced-master/dal"
)

var (
	appConfDirFromEnv  = os.Getenv("RESOURCED_MASTER_CONFIG_DIR")
	appConfDirFromFlag = kingpin.Flag("conf", "Path to config directory").Short('c').String()

	appServerArg  = kingpin.Command("server", "Run resourced-master server.").Default()
	appMigrateArg = kingpin.Command("migrate", "CLI interface for resourced-master database migration.")

	appMigrateUpArg = appMigrateArg.Command("up", "Run all migrations to the most current.").Default()

	appVersion = "4.2.4"
)

func init() {
	gob.Register(&dal.UserRow{})
	gob.Register(&dal.ClusterRow{})
}

func main() {
	kingpin.UsageTemplate(kingpin.CompactUsageTemplate).Version(appVersion).Author("Didip Kerabat")
	parsedCLIArgs := kingpin.Parse()

	if appConfDirFromEnv == "" && *appConfDirFromFlag == "" {
		logrus.Fatal("Path to config directory is required. You must set RESOURCED_MASTER_CONFIG_DIR environment variable or -c flag.")
	}

	configDir := appConfDirFromEnv
	if configDir == "" {
		configDir = *appConfDirFromFlag
	}

	app, err := application.New(configDir)
	if err != nil {
		logrus.Fatal(err)
	}

	switch parsedCLIArgs {
	case "server":
		// Create MessageBus
		bus, err := app.NewMessageBus(app.GeneralConfig)
		if err != nil {
			logrus.Fatal(err)
		}
		app.MessageBus = bus

		go app.MessageBus.ManageClients()

		go app.MessageBus.OnReceive(app.MessageBusHandlers())

		// Broadcast heartbeat
		go app.SendHeartbeat()

		go func() {
			// On boot, assign self to peers map and send a message to RefetchChecksChan
			app.Peers.Set(app.FullAddr(), true, gocache.DefaultExpiration)
			app.RefetchChecksChan <- true
		}()

		// Run all checks
		app.CheckAndRunTriggers()

		// Prune old timeseries data
		go app.PruneAll()

		// Publish metrics to local agent, which is a graphite endpoint.
		addr, err := net.ResolveTCPAddr("tcp", "localhost:"+app.GeneralConfig.LocalAgent.GraphiteTCPPort)
		if err != nil {
			logrus.Fatal(err)
		}
		statsInterval, err := time.ParseDuration(app.GeneralConfig.LocalAgent.ReportMetricsInterval)
		if err != nil {
			logrus.Fatal(err)
		}
		go metrics_graphite.Graphite(app.MetricsRegistry, statsInterval, "ResourcedMaster", addr)

		// Create HTTP server
		srv, err := app.NewHTTPServer()
		if err != nil {
			logrus.Fatal(err)
		}

		if app.GeneralConfig.HTTPS.CertFile != "" && app.GeneralConfig.HTTPS.KeyFile != "" {
			logrus.WithFields(logrus.Fields{"Addr": app.GeneralConfig.Addr}).Info("Running HTTPS server")
			srv.ListenAndServeTLS(app.GeneralConfig.HTTPS.CertFile, app.GeneralConfig.HTTPS.KeyFile)
		} else {
			logrus.WithFields(logrus.Fields{"Addr": app.GeneralConfig.Addr}).Info("Running HTTP server")
			srv.ListenAndServe()
		}

	case "migrate up":
		err := app.MigrateUpAll()
		if err != nil {
			logrus.Fatal(err)
		}
	}
}
<commit_msg>Remove app version, it’s just too much work.<commit_after>package main

import (
	"encoding/gob"
	"net"
	"os"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/alecthomas/kingpin"
	metrics_graphite "github.com/cyberdelia/go-metrics-graphite"
	gocache "github.com/patrickmn/go-cache"

	"github.com/resourced/resourced-master/application"
	"github.com/resourced/resourced-master/dal"
)

var (
	appConfDirFromEnv  = os.Getenv("RESOURCED_MASTER_CONFIG_DIR")
	appConfDirFromFlag = kingpin.Flag("conf", "Path to config directory").Short('c').String()

	appServerArg  = kingpin.Command("server", "Run resourced-master server.").Default()
	appMigrateArg = kingpin.Command("migrate", "CLI interface for resourced-master database migration.")

	appMigrateUpArg = appMigrateArg.Command("up", "Run all migrations to the most current.").Default()
)

func init() {
	gob.Register(&dal.UserRow{})
	gob.Register(&dal.ClusterRow{})
}

func main() {
	kingpin.UsageTemplate(kingpin.CompactUsageTemplate).Author("Didip Kerabat")
	parsedCLIArgs := kingpin.Parse()

	if appConfDirFromEnv == "" && *appConfDirFromFlag == "" {
		logrus.Fatal("Path to config directory is required. You must set RESOURCED_MASTER_CONFIG_DIR environment variable or -c flag.")
	}

	configDir := appConfDirFromEnv
	if configDir == "" {
		configDir = *appConfDirFromFlag
	}

	app, err := application.New(configDir)
	if err != nil {
		logrus.Fatal(err)
	}

	switch parsedCLIArgs {
	case "server":
		// Create MessageBus
		bus, err := app.NewMessageBus(app.GeneralConfig)
		if err != nil {
			logrus.Fatal(err)
		}
		app.MessageBus = bus

		go app.MessageBus.ManageClients()

		go app.MessageBus.OnReceive(app.MessageBusHandlers())

		// Broadcast heartbeat
		go app.SendHeartbeat()

		go func() {
			// On boot, assign self to peers map and send a message to RefetchChecksChan
			app.Peers.Set(app.FullAddr(), true, gocache.DefaultExpiration)
			app.RefetchChecksChan <- true
		}()

		// Run all checks
		app.CheckAndRunTriggers()

		// Prune old timeseries data
		go app.PruneAll()

		// Publish metrics to local agent, which is a graphite endpoint.
		addr, err := net.ResolveTCPAddr("tcp", "localhost:"+app.GeneralConfig.LocalAgent.GraphiteTCPPort)
		if err != nil {
			logrus.Fatal(err)
		}
		statsInterval, err := time.ParseDuration(app.GeneralConfig.LocalAgent.ReportMetricsInterval)
		if err != nil {
			logrus.Fatal(err)
		}
		go metrics_graphite.Graphite(app.MetricsRegistry, statsInterval, "ResourcedMaster", addr)

		// Create HTTP server
		srv, err := app.NewHTTPServer()
		if err != nil {
			logrus.Fatal(err)
		}

		if app.GeneralConfig.HTTPS.CertFile != "" && app.GeneralConfig.HTTPS.KeyFile != "" {
			logrus.WithFields(logrus.Fields{"Addr": app.GeneralConfig.Addr}).Info("Running HTTPS server")
			srv.ListenAndServeTLS(app.GeneralConfig.HTTPS.CertFile, app.GeneralConfig.HTTPS.KeyFile)
		} else {
			logrus.WithFields(logrus.Fields{"Addr": app.GeneralConfig.Addr}).Info("Running HTTP server")
			srv.ListenAndServe()
		}

	case "migrate up":
		err := app.MigrateUpAll()
		if err != nil {
			logrus.Fatal(err)
		}
	}
}
<|endoftext|>