<commit_before>// Licensed under the Apache License, Version 2.0
// Details: https://raw.githubusercontent.com/square/quotaservice/master/LICENSE

package quotaservice

import (
	"time"

	"github.com/square/quotaservice/config"
	"github.com/square/quotaservice/logging"
	pbconfig "github.com/square/quotaservice/protos/config"
)

// watcher watches reapableBuckets for activity.
type watcher struct {
	ns           string
	bucketName   string
	identifier   string
	maxIdle      time.Duration
	lastActivity time.Time
	activities   <-chan struct{}
}

// activityDetected tells you if activity has been detected since the last time this method was
// called.
func (w *watcher) activityDetected() bool {
	select {
	case <-w.activities:
		return true
	default:
		return false
	}
}

// tooIdle returns true if a bucket has been idle for longer than its maxIdle
func (w *watcher) tooIdle(now time.Time) bool {
	if w.activityDetected() {
		w.lastActivity = now
		return false
	}
	return now.Sub(w.lastActivity) > w.maxIdle
}

// reapableBucket is a wrapper around a bucket that overrides ReportActivity(), and reports any
// activity on a channel that is monitored by the reaper.
type reapableBucket struct {
	Bucket
	activities chan<- struct{}
}

// ReportActivity is overridden to report activity to the reapableBucket's activity channel.
func (r *reapableBucket) ReportActivity() {
	select {
	case r.activities <- struct{}{}:
	// reported activity
	default:
		// Already reported
	}

	// Now call ReportActivity() on the delegate bucket.
	r.Bucket.ReportActivity()
}

// Destroy closes the reapableBucket's activity channel.
func (r *reapableBucket) Destroy() {
	close(r.activities)

	// Now call Destroy() on the delegate bucket.
	r.Bucket.Destroy()
}

func createWatcher(ns, bucketName string, maxIdle time.Duration, activityChannel <-chan struct{}) *watcher {
	return &watcher{
		ns:         ns,
		bucketName: bucketName,
		identifier: config.FullyQualifiedName(ns, bucketName),
		maxIdle:    maxIdle,
		activities: activityChannel}
}

type reaper struct {
	cfg         config.ReaperConfig
	newWatchers chan<- *watcher
	watchers    map[string]*watcher
}

func newReaper(bc *bucketContainer, r config.ReaperConfig) *reaper {
	watcherChannel := make(chan *watcher, r.BucketWatcherBuffer)
	reaper := &reaper{
		cfg:         r,
		watchers:    make(map[string]*watcher),
		newWatchers: watcherChannel}

	go reaper.reapIdleBuckets(bc, watcherChannel)

	return reaper
}

func (r *reaper) stop() {
	// This will trigger the goroutine waiting on watchers to exit.
	close(r.newWatchers)
}

func (r *reaper) addNewWatcher(w *watcher) {
	r.watchers[w.identifier] = w
	w.lastActivity = time.Now()
}

// checkExpirations checks all watches registered with the reaper, and destroys idle buckets, updating the reaper
// accordingly. Returns the duration after which it should run again.
func (r *reaper) checkExpirations(bc *bucketContainer) time.Duration {
	now := time.Now()
	newSleep := r.cfg.MinFrequency
	var reaped uint64
	for id, w := range r.watchers {
		if w.tooIdle(now) {
			// Reap bucket
			reaped++
			if bc.removeBucket(w.ns, w.bucketName) {
				delete(r.watchers, id)
			}
		} else if w.maxIdle < newSleep {
			// Check if we're sleeping the right amount.
			newSleep = w.maxIdle
		}
	}
	logging.Printf("Reaped %d buckets due to inactivity", reaped)
	return newSleep
}

// watch watches all buckets for activity, deleting the bucket if no activity has been detected
// after a given duration.
func (r *reaper) reapIdleBuckets(bc *bucketContainer, newWatchers <-chan *watcher) {
	sleep := r.cfg.InitSleep
	logging.Printf("reapIdleBuckets started. Initial sleep %v", sleep)
	ticker := time.NewTicker(sleep)

	// Watch on a ticker, or a new watch being created.
	for {
		select {
		case w, ok := <-newWatchers:
			if ok {
				r.addNewWatcher(w)
			} else {
				// newWatchers closed; stop the reaper.
				r.newWatchers = nil
				r.watchers = nil
				return
			}

		case <-ticker.C:
			newSleep := r.checkExpirations(bc)

			if newSleep != sleep {
				logging.Printf("Adjusting ticker to run with duration %v", newSleep)
				// We need a new ticker.
				ticker.Stop()
				ticker = time.NewTicker(newSleep)
				sleep = newSleep
			}
		}
	}
}

// applyWatch decorates a bucket to make it "watchable", if it has a maxIdle and requires garbage
// collection. Callers should ensure they point to the return value of this method when
// referencing their bucket.
func (r *reaper) applyWatch(delegate Bucket, namespace, bucketName string, cfg *pbconfig.BucketConfig) (Bucket, *watcher) {
	if cfg.MaxIdleMillis > 0 {
		activityChannel := make(chan struct{}, 1)
		rb := &reapableBucket{Bucket: delegate, activities: activityChannel}
		w := createWatcher(namespace, bucketName, time.Duration(cfg.MaxIdleMillis)*time.Millisecond, activityChannel)
		r.newWatchers <- w
		return rb, w
	}

	return delegate, nil
}
<commit_msg>Make a minor comment fix to reapIdleBuckets in reaper.go<commit_after>// Licensed under the Apache License, Version 2.0
// Details: https://raw.githubusercontent.com/square/quotaservice/master/LICENSE

package quotaservice

import (
	"time"

	"github.com/square/quotaservice/config"
	"github.com/square/quotaservice/logging"
	pbconfig "github.com/square/quotaservice/protos/config"
)

// watcher watches reapableBuckets for activity.
type watcher struct {
	ns           string
	bucketName   string
	identifier   string
	maxIdle      time.Duration
	lastActivity time.Time
	activities   <-chan struct{}
}

// activityDetected tells you if activity has been detected since the last time this method was
// called.
func (w *watcher) activityDetected() bool {
	select {
	case <-w.activities:
		return true
	default:
		return false
	}
}

// tooIdle returns true if a bucket has been idle for longer than its maxIdle
func (w *watcher) tooIdle(now time.Time) bool {
	if w.activityDetected() {
		w.lastActivity = now
		return false
	}
	return now.Sub(w.lastActivity) > w.maxIdle
}

// reapableBucket is a wrapper around a bucket that overrides ReportActivity(), and reports any
// activity on a channel that is monitored by the reaper.
type reapableBucket struct {
	Bucket
	activities chan<- struct{}
}

// ReportActivity is overridden to report activity to the reapableBucket's activity channel.
func (r *reapableBucket) ReportActivity() {
	select {
	case r.activities <- struct{}{}:
	// reported activity
	default:
		// Already reported
	}

	// Now call ReportActivity() on the delegate bucket.
	r.Bucket.ReportActivity()
}

// Destroy closes the reapableBucket's activity channel.
func (r *reapableBucket) Destroy() {
	close(r.activities)

	// Now call Destroy() on the delegate bucket.
	r.Bucket.Destroy()
}

func createWatcher(ns, bucketName string, maxIdle time.Duration, activityChannel <-chan struct{}) *watcher {
	return &watcher{
		ns:         ns,
		bucketName: bucketName,
		identifier: config.FullyQualifiedName(ns, bucketName),
		maxIdle:    maxIdle,
		activities: activityChannel}
}

type reaper struct {
	cfg         config.ReaperConfig
	newWatchers chan<- *watcher
	watchers    map[string]*watcher
}

func newReaper(bc *bucketContainer, r config.ReaperConfig) *reaper {
	watcherChannel := make(chan *watcher, r.BucketWatcherBuffer)
	reaper := &reaper{
		cfg:         r,
		watchers:    make(map[string]*watcher),
		newWatchers: watcherChannel}

	go reaper.reapIdleBuckets(bc, watcherChannel)

	return reaper
}

func (r *reaper) stop() {
	// This will trigger the goroutine waiting on watchers to exit.
	close(r.newWatchers)
}

func (r *reaper) addNewWatcher(w *watcher) {
	r.watchers[w.identifier] = w
	w.lastActivity = time.Now()
}

// checkExpirations checks all watches registered with the reaper, and destroys idle buckets, updating the reaper
// accordingly. Returns the duration after which it should run again.
func (r *reaper) checkExpirations(bc *bucketContainer) time.Duration {
	now := time.Now()
	newSleep := r.cfg.MinFrequency
	var reaped uint64
	for id, w := range r.watchers {
		if w.tooIdle(now) {
			// Reap bucket
			reaped++
			if bc.removeBucket(w.ns, w.bucketName) {
				delete(r.watchers, id)
			}
		} else if w.maxIdle < newSleep {
			// Check if we're sleeping the right amount.
			newSleep = w.maxIdle
		}
	}
	logging.Printf("Reaped %d buckets due to inactivity", reaped)
	return newSleep
}

// reapIdleBuckets watches all buckets for activity, deleting the bucket if no activity has been detected
// after a given duration.
func (r *reaper) reapIdleBuckets(bc *bucketContainer, newWatchers <-chan *watcher) {
	sleep := r.cfg.InitSleep
	logging.Printf("reapIdleBuckets started. Initial sleep %v", sleep)
	ticker := time.NewTicker(sleep)

	// Watch on a ticker, or a new watch being created.
	for {
		select {
		case w, ok := <-newWatchers:
			if ok {
				r.addNewWatcher(w)
			} else {
				// newWatchers closed; stop the reaper.
				r.newWatchers = nil
				r.watchers = nil
				return
			}

		case <-ticker.C:
			newSleep := r.checkExpirations(bc)

			if newSleep != sleep {
				logging.Printf("Adjusting ticker to run with duration %v", newSleep)
				// We need a new ticker.
				ticker.Stop()
				ticker = time.NewTicker(newSleep)
				sleep = newSleep
			}
		}
	}
}

// applyWatch decorates a bucket to make it "watchable", if it has a maxIdle and requires garbage
// collection. Callers should ensure they point to the return value of this method when
// referencing their bucket.
func (r *reaper) applyWatch(delegate Bucket, namespace, bucketName string, cfg *pbconfig.BucketConfig) (Bucket, *watcher) {
	if cfg.MaxIdleMillis > 0 {
		activityChannel := make(chan struct{}, 1)
		rb := &reapableBucket{Bucket: delegate, activities: activityChannel}
		w := createWatcher(namespace, bucketName, time.Duration(cfg.MaxIdleMillis)*time.Millisecond, activityChannel)
		r.newWatchers <- w
		return rb, w
	}

	return delegate, nil
}
<|endoftext|>


<commit_before>package resolv

const (
	PortDefault = "53"
)

var RootServers = []string{
	"a.root-servers.net.",
	"b.root-servers.net.",
	"c.root-servers.net.",
	"d.root-servers.net.",
	"e.root-servers.net.",
	"f.root-servers.net.",
	"g.root-servers.net.",
	"h.root-servers.net.",
	"i.root-servers.net.",
}
<commit_msg>Add i-m root servers<commit_after>package resolv

const (
	PortDefault = "53"
)

var RootServers = []string{
	"a.root-servers.net.",
	"b.root-servers.net.",
	"c.root-servers.net.",
	"d.root-servers.net.",
	"e.root-servers.net.",
	"f.root-servers.net.",
	"g.root-servers.net.",
	"h.root-servers.net.",
	"i.root-servers.net.",
	"j.root-servers.net.",
	"k.root-servers.net.",
	"l.root-servers.net.",
	"m.root-servers.net.",
}
<|endoftext|>


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

// TODO(adg): accept -a flag on 'create' and 'commit' (like git commit -a)
// TODO(adg): accept -r flag on 'upload' to nominate reviewer
// TODO(adg): support 'create' from non-master branches
// TODO(adg): check style of commit message
// TODO(adg): write doc comment
// TOOD(adg): print gerrit votes on 'pending'

package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"sync"
)

var (
	hookFile = filepath.FromSlash(".git/hooks/commit-msg")
	verbose  = flag.Bool("v", false, "verbose output")
)

const usage = `Usage: %s [-v] <command>
Type "%s help" for more information.
`

const help = `Usage: %s [-v] <command>

The review command is a wrapper for the git command that provides a simple
interface to the "single-commit feature branch" development model.

Available comands:

	create <name>
		Create a local branch with the provided name
		and commit the staged changes to it.

	commit
		Amend local branch HEAD commit with the staged changes.

	diff
		View differences between remote branch HEAD and
		the local branch HEAD.
		(The differences introduced by this change.)

	upload
		Upload HEAD commit to the code review server.

	sync
		Fetch changes from the remote repository and merge them to the
		current branch, rebasing the HEAD commit (if any) on top of
		them. If the HEAD commit has been submitted, switch back to the
		master branch and delete the feature branch.

	pending
		Show local branches and their head commits.

`

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, usage, os.Args[0], os.Args[0])
		os.Exit(2)
	}
	flag.Parse()

	goToRepoRoot()
	installHook()

	switch flag.Arg(0) {
	case "help":
		fmt.Fprintf(os.Stdout, help, os.Args[0])
	case "create", "cr":
		name := flag.Arg(1)
		if name == "" {
			flag.Usage()
		}
		create(name)
	case "commit", "co":
		commit()
	case "diff", "d":
		diff()
	case "upload", "u":
		upload()
	case "sync", "s":
		doSync()
	case "pending", "p":
		pending()
	default:
		flag.Usage()
	}
}

func create(name string) {
	if !hasStagedChanges() {
		dief("no staged changes.\nDid you forget to 'git add'?")
	}
	if b := currentBranch(); b != "master" {
		dief("must run 'create' from the master branch (now on %q).\n"+
			"(Try 'review sync' or 'git checkout master' first.)",
			b)
	}
	run("git", "checkout", "-q", "-b", name)
	if err := runErr("git", "commit", "-q"); err != nil {
		verbosef("Commit failed: %v\nSwitching back to master.\n", err)
		run("git", "checkout", "-q", "master")
		run("git", "branch", "-q", "-d", name)
	}
}

func commit() {
	if !hasStagedChanges() {
		dief("no staged changes. Did you forget to 'git add'?")
	}
	if currentBranch() == "master" {
		dief("can't commit to master branch.")
	}
	run("git", "commit", "-q", "--amend", "-C", "HEAD")
}

func diff() {
	run("git", "diff", "HEAD^", "HEAD")
}

func upload() {
	if currentBranch() == "master" {
		dief("can't upload from master branch.")
	}
	run("git", "push", "-q", "origin", "HEAD:refs/for/master")
}

func doSync() {
	run("git", "fetch", "-q")

	// If we're on master, just fast-forward.
	branch := currentBranch()
	if branch == "master" {
		run("git", "merge", "-q", "--ff-only", "origin/master")
		return
	}

	// Check that exactly this commit was submitted to master. If so,
	// switch back to master, fast-forward, delete the feature branch.
	if branchContains("origin/master", "HEAD") {
		run("git", "checkout", "-q", "master")
		run("git", "merge", "-q", "--ff-only", "origin/master")
		run("git", "branch", "-q", "-d", branch)
		fmt.Printf("Change on %q submitted; branch deleted.\n"+
			"On master branch.\n", branch)
		return
	}

	// Check whether a rebased version of this commit was submitted to
	// master. If so, switch back to master, fast-forward, and
	// provide instructions for deleting the feature branch.
	// (We're not 100% sure that the feature branch HEAD was submitted,
	// so be cautious.)
	if headSubmitted(branch) {
		run("git", "checkout", "-q", "master")
		run("git", "merge", "-q", "--ff-only", "origin/master")
		fmt.Fprintf(os.Stderr,
			"I think the change on %q has been submitted.\n"+
				"If you agree, and no longer need branch %q, "+
				"run:\n\tgit branch -D %v\n"+
				"On master branch.\n",
			branch, branch, branch)
		return
	}

	// Bump master HEAD to that of origin/master, just in case the user
	// switches back to master with "git checkout master" later.
	// TODO(adg): maybe we shouldn't do this at all?
	if !branchContains("origin/master", "master") {
		run("git", "branch", "-f", "master", "origin/master")
	}

	// We have un-submitted changes on this feature branch; rebase.
	run("git", "rebase", "-q", "origin/master")
}

func pending() {
	var (
		wg      sync.WaitGroup
		origin  = originURL()
		current = currentBranch()
	)
	if current == "master" {
		fmt.Println("On master branch.")
	}
	for _, branch := range localBranches() {
		if branch == "master" {
			continue
		}
		wg.Add(1)
		go func(branch string) {
			defer wg.Done()
			p := ""
			if branch == current {
				p = "* "
			}
			id := headChangeId(branch)
			c, err := getChange(origin, id)
			switch err {
			case notFound:
				// TODO(adg): read local commit msg
				var msg string
				fmt.Printf("%v%v:\n\t%v\n\t(not uploaded)\n",
					p, branch, msg)
			case nil:
				status := ""
				switch c.Status {
				case "MERGED":
					status += " [submitted]"
				case "ABANDONED":
					status += " [abandoned]"
				}
				fmt.Printf("%v%v%v:\n\t%v\n\t%v\n",
					p, branch, status, c.Subject, c.URL)
			default:
				fmt.Fprintf(os.Stderr, "fetching change for %q: %v\n", branch, err)
			}
		}(branch)
	}
	wg.Wait()
}

func originURL() string {
	out, err := exec.Command("git", "config", "remote.origin.url").CombinedOutput()
	if err != nil {
		dief("could not find URL for 'origin' remote.\n"+
			"Did you check out from the right place?\n"+
			"git config remote.origin.url: %v\n"+
			"%s", err, out)
	}
	return string(out)
}

func localBranches() (branches []string) {
	for _, s := range runLines("git", "branch", "-l", "-q") {
		branches = append(branches, strings.TrimPrefix(s, "* "))
	}
	return branches
}

func branchContains(branch, rev string) bool {
	for _, s := range runLines("git", "branch", "-r", "--contains", rev) {
		if s == branch {
			return true
		}
	}
	return false
}

var stagedRe = regexp.MustCompile(`^[ACDMR]  `)

func hasStagedChanges() bool {
	for _, s := range gitStatus() {
		if stagedRe.MatchString(s) {
			return true
		}
	}
	return false
}

func currentBranch() string {
	return strings.TrimSpace(runOutput("git", "rev-parse", "--abbrev-ref", "HEAD"))
}

func gitStatus() []string {
	return runLines("git", "status", "-b", "--porcelain")
}

func headSubmitted(branch string) bool {
	s := "Change-Id: " + headChangeId(branch)
	return len(runOutput("git", "log", "--grep", s, "origin/master")) > 0
}

func headChangeId(branch string) string {
	const (
		p = "Change-Id: "
		f = "--format=format:%b"
	)
	for _, s := range runLines("git", "log", "-n", "1", f, branch, "--") {
		if strings.HasPrefix(s, p) {
			return strings.TrimSpace(strings.TrimPrefix(s, p))
		}
	}
	dief("no Change-Id line found in HEAD commit on branch %s.", branch)
	panic("unreachable")
}

func goToRepoRoot() {
	prevDir, err := os.Getwd()
	if err != nil {
		dief("could not get current directory: %v", err)
	}
	for {
		if _, err := os.Stat(".git"); err == nil {
			return
		}
		if err := os.Chdir(".."); err != nil {
			dief("could not chdir: %v", err)
		}
		currentDir, err := os.Getwd()
		if err != nil {
			dief("could not get current directory: %v", err)
		}
		if currentDir == prevDir {
			dief("git root not found.\n" +
				"Run from within the Git tree please.")
		}
		prevDir = currentDir
	}
}

func installHook() {
	_, err := os.Stat(hookFile)
	if err == nil {
		return
	}
	if !os.IsNotExist(err) {
		dief("error checking for hook file: %v", err)
	}
	verbosef("Presubmit hook to add Change-Id to commit messages is missing.\n"+
		"Automatically creating it at %v.\n", hookFile)
	hookContent := []byte(commitMsgHook)
	if err := ioutil.WriteFile(hookFile, hookContent, 0700); err != nil {
		dief("error writing hook file: %v", err)
	}
}

func dief(format string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, "review: "+format+"\n", args...)
	os.Exit(1)
}

func run(command string, args ...string) {
	if err := runErr(command, args...); err != nil {
		if !*verbose {
			// If we're not in verbose mode, print the command
			// before dying to give context to the failure.
			fmt.Fprintln(os.Stderr, commandString(command, args))
		}
		dief("%v", err)
	}
}

func runErr(command string, args ...string) error {
	if *verbose {
		fmt.Fprintln(os.Stderr, commandString(command, args))
	}
	cmd := exec.Command(command, args...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

func verbosef(format string, args ...interface{}) {
	if *verbose {
		fmt.Fprintf(os.Stderr, format, args...)
	}
}

func runOutput(command string, args ...string) string {
	b, err := exec.Command(command, args...).CombinedOutput()
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n%s\n", commandString(command, args), b)
		dief("%v", err)
	}
	return string(b)
}

func runLines(command string, args ...string) []string {
	var s []string
	for _, l := range strings.Split(runOutput(command, args...), "\n") {
		l = strings.TrimSpace(l)
		if l != "" {
			s = append(s, l)
		}
	}
	return s
}

func commandString(command string, args []string) string {
	return strings.Join(append([]string{command}, args...), " ")
}
<commit_msg>Add a couple more todos; very minor refactoring.<commit_after>// Copyright 2014 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// TODO(adg): accept -a flag on 'create' and 'commit' (like git commit -a)
// TODO(adg): accept -r flag on 'upload' to nominate reviewer
// TODO(adg): support 'create' from non-master branches
// TODO(adg): check style of commit message
// TODO(adg): write doc comment
// TOOD(adg): print gerrit votes on 'pending'
// TODO(adg): add gofmt commit hook
// TODO(adg): 'upload' warn about uncommitted changes (maybe commit/create too?)

package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"sync"
)

var (
	hookFile = filepath.FromSlash(".git/hooks/commit-msg")
	verbose  = flag.Bool("v", false, "verbose output")
)

const usage = `Usage: %s [-v] <command>
Type "%s help" for more information.
`

const help = `Usage: %s [-v] <command>

The review command is a wrapper for the git command that provides a simple
interface to the "single-commit feature branch" development model.

Available comands:

	create <name>
		Create a local branch with the provided name
		and commit the staged changes to it.

	commit
		Amend local branch HEAD commit with the staged changes.

	diff
		View differences between remote branch HEAD and
		the local branch HEAD.
		(The differences introduced by this change.)

	upload
		Upload HEAD commit to the code review server.

	sync
		Fetch changes from the remote repository and merge them to the
		current branch, rebasing the HEAD commit (if any) on top of
		them. If the HEAD commit has been submitted, switch back to the
		master branch and delete the feature branch.

	pending
		Show local branches and their head commits.

`

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, usage, os.Args[0], os.Args[0])
		os.Exit(2)
	}
	flag.Parse()

	goToRepoRoot()
	installHook()

	switch flag.Arg(0) {
	case "help":
		fmt.Fprintf(os.Stdout, help, os.Args[0])
	case "create", "cr":
		name := flag.Arg(1)
		if name == "" {
			flag.Usage()
		}
		create(name)
	case "commit", "co":
		commit()
	case "diff", "d":
		diff()
	case "upload", "u":
		upload()
	case "sync", "s":
		doSync()
	case "pending", "p":
		pending()
	default:
		flag.Usage()
	}
}

func create(name string) {
	if !hasStagedChanges() {
		dief("no staged changes.\nDid you forget to 'git add'?")
	}
	if b := currentBranch(); b != "master" {
		dief("must run 'create' from the master branch (now on %q).\n"+
			"(Try 'review sync' or 'git checkout master' first.)",
			b)
	}
	run("git", "checkout", "-q", "-b", name)
	if err := runErr("git", "commit", "-q"); err != nil {
		verbosef("Commit failed: %v\nSwitching back to master.\n", err)
		run("git", "checkout", "-q", "master")
		run("git", "branch", "-q", "-d", name)
	}
}

func commit() {
	if !hasStagedChanges() {
		dief("no staged changes. Did you forget to 'git add'?")
	}
	if currentBranch() == "master" {
		dief("can't commit to master branch.")
	}
	run("git", "commit", "-q", "--amend", "-C", "HEAD")
}

func diff() {
	run("git", "diff", "HEAD^", "HEAD")
}

func upload() {
	if currentBranch() == "master" {
		dief("can't upload from master branch.")
	}
	run("git", "push", "-q", "origin", "HEAD:refs/for/master")
}

func doSync() {
	run("git", "fetch", "-q")

	// If we're on master, just fast-forward.
	branch := currentBranch()
	if branch == "master" {
		run("git", "merge", "-q", "--ff-only", "origin/master")
		return
	}

	// Check that exactly this commit was submitted to master. If so,
	// switch back to master, fast-forward, delete the feature branch.
	if branchContains("origin/master", "HEAD") {
		run("git", "checkout", "-q", "master")
		run("git", "merge", "-q", "--ff-only", "origin/master")
		run("git", "branch", "-q", "-d", branch)
		fmt.Printf("Change on %q submitted; branch deleted.\n"+
			"On master branch.\n", branch)
		return
	}

	// Check whether a rebased version of this commit was submitted to
	// master. If so, switch back to master, fast-forward, and
	// provide instructions for deleting the feature branch.
	// (We're not 100% sure that the feature branch HEAD was submitted,
	// so be cautious.)
	if headSubmitted(branch) {
		run("git", "checkout", "-q", "master")
		run("git", "merge", "-q", "--ff-only", "origin/master")
		fmt.Fprintf(os.Stderr,
			"I think the change on %q has been submitted.\n"+
				"If you agree, and no longer need branch %q, "+
				"run:\n\tgit branch -D %v\n"+
				"On master branch.\n",
			branch, branch, branch)
		return
	}

	// Bump master HEAD to that of origin/master, just in case the user
	// switches back to master with "git checkout master" later.
	// TODO(adg): maybe we shouldn't do this at all?
	if !branchContains("origin/master", "master") {
		run("git", "branch", "-f", "master", "origin/master")
	}

	// We have un-submitted changes on this feature branch; rebase.
	run("git", "rebase", "-q", "origin/master")
}

func pending() {
	var (
		wg      sync.WaitGroup
		origin  = originURL()
		current = currentBranch()
	)
	if current == "master" {
		fmt.Println("On master branch.")
	}
	for _, branch := range localBranches() {
		if branch == "master" {
			continue
		}
		wg.Add(1)
		go func(branch string) {
			defer wg.Done()
			p := ""
			if branch == current {
				p = "* "
			}
			id := headChangeId(branch)
			c, err := getChange(origin, id)
			switch err {
			case notFound:
				// TODO(adg): read local commit msg
				var msg string
				fmt.Printf("%v%v:\n\t%v\n\t(not uploaded)\n",
					p, branch, msg)
			case nil:
				status := ""
				switch c.Status {
				case "MERGED":
					status += " [submitted]"
				case "ABANDONED":
					status += " [abandoned]"
				}
				fmt.Printf("%v%v%v:\n\t%v\n\t%v\n",
					p, branch, status, c.Subject, c.URL)
			default:
				fmt.Fprintf(os.Stderr, "fetching change for %q: %v\n", branch, err)
			}
		}(branch)
	}
	wg.Wait()
}

func originURL() string {
	out, err := exec.Command("git", "config", "remote.origin.url").CombinedOutput()
	if err != nil {
		dief("could not find URL for 'origin' remote.\n"+
			"Did you check out from the right place?\n"+
			"git config remote.origin.url: %v\n"+
			"%s", err, out)
	}
	return string(out)
}

func localBranches() (branches []string) {
	for _, s := range runLines("git", "branch", "-l", "-q") {
		branches = append(branches, strings.TrimPrefix(s, "* "))
	}
	return branches
}

func branchContains(branch, rev string) bool {
	for _, s := range runLines("git", "branch", "-r", "--contains", rev) {
		if s == branch {
			return true
		}
	}
	return false
}

var stagedRe = regexp.MustCompile(`^[ACDMR]  `)

func hasStagedChanges() bool {
	for _, s := range runLines("git", "status", "-b", "--porcelain") {
		if stagedRe.MatchString(s) {
			return true
		}
	}
	return false
}

func currentBranch() string {
	return strings.TrimSpace(runOutput("git", "rev-parse", "--abbrev-ref", "HEAD"))
}

func headSubmitted(branch string) bool {
	s := "Change-Id: " + headChangeId(branch)
	return len(runOutput("git", "log", "--grep", s, "origin/master")) > 0
}

func headChangeId(branch string) string {
	const (
		p = "Change-Id: "
		f = "--format=format:%b"
	)
	for _, s := range runLines("git", "log", "-n", "1", f, branch, "--") {
		if strings.HasPrefix(s, p) {
			return strings.TrimSpace(strings.TrimPrefix(s, p))
		}
	}
	dief("no Change-Id line found in HEAD commit on branch %s.", branch)
	panic("unreachable")
}

func goToRepoRoot() {
	prevDir, err := os.Getwd()
	if err != nil {
		dief("could not get current directory: %v", err)
	}
	for {
		if _, err := os.Stat(".git"); err == nil {
			return
		}
		if err := os.Chdir(".."); err != nil {
			dief("could not chdir: %v", err)
		}
		currentDir, err := os.Getwd()
		if err != nil {
			dief("could not get current directory: %v", err)
		}
		if currentDir == prevDir {
			dief("git root not found.\n" +
				"Run from within the Git tree please.")
		}
		prevDir = currentDir
	}
}

func installHook() {
	_, err := os.Stat(hookFile)
	if err == nil {
		return
	}
	if !os.IsNotExist(err) {
		dief("error checking for hook file: %v", err)
	}
	verbosef("Presubmit hook to add Change-Id to commit messages is missing.\n"+
		"Automatically creating it at %v.\n", hookFile)
	hookContent := []byte(commitMsgHook)
	if err := ioutil.WriteFile(hookFile, hookContent, 0700); err != nil {
		dief("error writing hook file: %v", err)
	}
}

func run(command string, args ...string) {
	if err := runErr(command, args...); err != nil {
		if !*verbose {
			// If we're not in verbose mode, print the command
			// before dying to give context to the failure.
			fmt.Fprintln(os.Stderr, commandString(command, args))
		}
		dief("%v", err)
	}
}

func runErr(command string, args ...string) error {
	if *verbose {
		fmt.Fprintln(os.Stderr, commandString(command, args))
	}
	cmd := exec.Command(command, args...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

func runOutput(command string, args ...string) string {
	b, err := exec.Command(command, args...).CombinedOutput()
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n%s\n", commandString(command, args), b)
		dief("%v", err)
	}
	return string(b)
}

func runLines(command string, args ...string) []string {
	var s []string
	for _, l := range strings.Split(runOutput(command, args...), "\n") {
		l = strings.TrimSpace(l)
		if l != "" {
			s = append(s, l)
		}
	}
	return s
}

func commandString(command string, args []string) string {
	return strings.Join(append([]string{command}, args...), " ")
}

func dief(format string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, "review: "+format+"\n", args...)
	os.Exit(1)
}

func verbosef(format string, args ...interface{}) {
	if *verbose {
		fmt.Fprintf(os.Stderr, format, args...)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"strings"
	"net/http"

  "github.com/asvins/router"
	"github.com/unrolled/render"
	"github.com/asvins/common_interceptors/logger"
)

func DiscoveryHandler(w http.ResponseWriter, req *http.Request) {
  prefix := strings.Join([]string{ServerConfig.Server.Addr, ServerConfig.Server.Port}, ":")
  r := render.New()

	//add discovery links here
  discoveryMap := map[string]string {"discovery": prefix+"/api/discovery"}

	r.JSON(w, http.StatusOK, discoveryMap)
}

func DefRoutes() *router.Router {
  r := router.NewRouter()

	r.AddRoute("/api/discovery", router.GET, DiscoveryHandler)

	// interceptors
	r.AddBaseInterceptor("/", logger.NewLogger())
	return r
}
<commit_msg>Trailing whitespaces<commit_after>package main

import (
	"net/http"
	"strings"

	"github.com/asvins/common_interceptors/logger"
	"github.com/asvins/router"
	"github.com/unrolled/render"
)

func DiscoveryHandler(w http.ResponseWriter, req *http.Request) {
	prefix := strings.Join([]string{ServerConfig.Server.Addr, ServerConfig.Server.Port}, ":")
	r := render.New()

	//add discovery links here
	discoveryMap := map[string]string{"discovery": prefix + "/api/discovery"}

	r.JSON(w, http.StatusOK, discoveryMap)
}

func DefRoutes() *router.Router {
	r := router.NewRouter()

	r.AddRoute("/api/discovery", router.GET, DiscoveryHandler)

	// interceptors
	r.AddBaseInterceptor("/", logger.NewLogger())
	return r
}
<|endoftext|>


<commit_before>package goblin

import (
    "net/http"
    "log"
    "fmt"
    "reflect"
    "regexp"
)

// 路由器, 负责存储路由规则
type Router struct {
    routes []*Route
}

// 路由协议, 负责将 URL 映射到正确的 Handle
type Route struct {
    method          string
    pattern         string
    controllerType  reflect.Type
    actionName      string
    regex           *regexp.Regexp
}

// 路由分发方法
func (m *Router) ServeHTTP(w http.ResponseWriter, r *http.Request) {
    path   := r.URL.Path
    method := r.Method

    // 逐个匹配
    for _, route := range m.routes {
        match, vals := route.route(method, path)
        // 如果匹配成功
        if match {
            log.Println("=> mached route is " + string(route.pattern))

            rc := reflect.New(route.controllerType)
            controller, _ := rc.Interface().(ControllerInterface)   // TODO 没看懂

            // 把请求上下文放入 Context
            ctx := &Context{ Params: vals, Request: r, ResponseWriter: w }

            // 初始化 controller
            controller.Init(ctx, route.controllerType.Name(), route.actionName)

            // Before()
            controller.Before()

            // invoke handler
            method := rc.MethodByName(route.actionName)
            method.Call([]reflect.Value{})

            // After()
            controller.After()

            return
        }
    }
    // TODO 404
    return
}

// 借鉴martini，优化route匹配方法
func (r *Route) route(method string, path string) (bool, map[string]string) {
    if method != r.method {
        return false, nil
    }
    matches := r.regex.FindStringSubmatch(path)

    if len(matches) > 0 && matches[0] == path {
        params := make(map[string]string)
        for i, name := range r.regex.SubexpNames() {
            if len(name) > 0 {
                params[name] = matches[i]    // SubexpNames()返回结果第0个元素永远是空字符串
            }
        }
        return true, params
    }
    return false, nil
}

// 生成regexp对象
// Eg:   pattern = /books/:id/users/:user_id
// Then: route.regex = /books/(?P<id>[^/#?]+)/users/(?P<user_id>[^/#?]+)\/?
var routeReg1 = regexp.MustCompile(`:[^/#?()\.\\]+`)

func newRoute(method string, pattern string, t reflect.Type, action string) *Route {
    route  := Route{method, pattern, t, action, nil}
    pattern = routeReg1.ReplaceAllStringFunc(pattern, func(m string) string {
        return fmt.Sprintf(`(?P<%s>[^/#?]+)`, m[1:])
    })
    pattern += `\/?`
    route.regex = regexp.MustCompile(pattern)
    return &route
}

func (p *Router) AddRoute(method string, pattern string, t reflect.Type, action string) {
    p.routes = append(p.routes, newRoute(method, pattern, t, action))
}

func (p *Router) Get(pattern string, c ControllerInterface, actionName string) {
    t := reflect.Indirect(reflect.ValueOf(c)).Type()
    p.AddRoute("GET", pattern, t, actionName)
    p.AddRoute("HEAD", pattern, t, actionName)
}

func (p *Router) Post(pattern string, c ControllerInterface, actionName string) {
    t := reflect.Indirect(reflect.ValueOf(c)).Type()
    p.AddRoute("POST", pattern, t, actionName)
}

func (p *Router) Put(pattern string, c ControllerInterface, actionName string) {
    t := reflect.Indirect(reflect.ValueOf(c)).Type()
    p.AddRoute("PUT", pattern, t, actionName)
}

func (p *Router) Delete(pattern string, c ControllerInterface, actionName string) {
    t := reflect.Indirect(reflect.ValueOf(c)).Type()
    p.AddRoute("DELETE", pattern, t, actionName)
}

func NewRouter() Router {
    router := Router{routes: make([]*Route, 0, 0)}
    return router
}
<commit_msg>获取请求参数<commit_after>package goblin

import (
    "net/http"
    "log"
    "fmt"
    "reflect"
    "regexp"
    // "io/ioutil"
)

// 路由器, 负责存储路由规则
type Router struct {
    routes []*Route
}

// 路由协议, 负责将 URL 映射到正确的 Handle
type Route struct {
    method          string
    pattern         string
    controllerType  reflect.Type
    actionName      string
    regex           *regexp.Regexp
}

// 路由分发方法
func (p *Router) ServeHTTP(w http.ResponseWriter, r *http.Request) {
    path   := r.URL.Path
    method := r.Method

    // 逐个匹配
    for _, route := range p.routes {
        match, vals := route.route(method, path)
        // 如果匹配成功
        if match {
            log.Println("=> mached route is " + string(route.pattern))

            rc := reflect.New(route.controllerType)
            controller, _ := rc.Interface().(ControllerInterface)   // TODO 没看懂

            // 把请求上下文放入 Context
            ctx := &Context{ Params: vals, Request: r, ResponseWriter: w }

            // 获取请求参数
            r.ParseForm()
            r.ParseMultipartForm(2 * 1024 * 1024)

            for k, v := range r.Form {
                ctx.Params[k] = v[0]
            }

            // 初始化 controller
            controller.Init(ctx, route.controllerType.Name(), route.actionName)

            // Before()
            controller.Before()

            // invoke handler
            method := rc.MethodByName(route.actionName)
            method.Call([]reflect.Value{})

            // After()
            controller.After()

            return
        }
    }
    log.Println("=> 404 not found ")
    p.NotFound(w,r)
}

// 借鉴martini，优化route匹配方法
func (r *Route) route(method string, path string) (bool, map[string]string) {
    if method != r.method {
        return false, nil
    }
    matches := r.regex.FindStringSubmatch(path)

    if len(matches) > 0 && matches[0] == path {
        params := make(map[string]string)
        for i, name := range r.regex.SubexpNames() {
            if len(name) > 0 {
                params[name] = matches[i]    // SubexpNames()返回结果第0个元素永远是空字符串
            }
        }
        return true, params
    }
    return false, nil
}

// 生成regexp对象
// Eg:   pattern = /books/:id/users/:user_id
// Then: route.regex = /books/(?P<id>[^/#?]+)/users/(?P<user_id>[^/#?]+)\/?
var routeReg1 = regexp.MustCompile(`:[^/#?()\.\\]+`)

func newRoute(method string, pattern string, t reflect.Type, action string) *Route {
    route  := Route{method, pattern, t, action, nil}
    pattern = routeReg1.ReplaceAllStringFunc(pattern, func(m string) string {
        return fmt.Sprintf(`(?P<%s>[^/#?]+)`, m[1:])
    })
    pattern += `\/?`
    route.regex = regexp.MustCompile(pattern)
    return &route
}

func (p *Router) AddRoute(method string, pattern string, t reflect.Type, action string) {
    p.routes = append(p.routes, newRoute(method, pattern, t, action))
}

func (p *Router) Get(pattern string, c ControllerInterface, actionName string) {
    t := reflect.Indirect(reflect.ValueOf(c)).Type()
    p.AddRoute("GET", pattern, t, actionName)
    p.AddRoute("HEAD", pattern, t, actionName)
}

func (p *Router) Post(pattern string, c ControllerInterface, actionName string) {
    t := reflect.Indirect(reflect.ValueOf(c)).Type()
    p.AddRoute("POST", pattern, t, actionName)
}

func (p *Router) Put(pattern string, c ControllerInterface, actionName string) {
    t := reflect.Indirect(reflect.ValueOf(c)).Type()
    p.AddRoute("PUT", pattern, t, actionName)
}

func (p *Router) Delete(pattern string, c ControllerInterface, actionName string) {
    t := reflect.Indirect(reflect.ValueOf(c)).Type()
    p.AddRoute("DELETE", pattern, t, actionName)
}

func NewRouter() Router {
    router := Router{routes: make([]*Route, 0, 0)}
    return router
}

func (p *Router) NotFound(w http.ResponseWriter, r *http.Request) {
    w.WriteHeader(404)
    fmt.Fprint(w, "404 Not Found.")
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/censhin/pokedex-api/moves"
	"github.com/censhin/pokedex-api/pokemon"
	"github.com/gorilla/mux"
)

func InitRoutes() *mux.Router {
	router := mux.NewRouter()

	router.HandleFunc("/pokemon", pokemon.CollectionResource).Methods("GET", "POST", "OPTIONS")
	router.HandleFunc("/pokemon/{id}", pokemon.MemberResource).Methods("GET", "PUT", "DELETE", "OPTIONS")
	router.HandleFunc("/moves", moves.CollectionResource).Methods("GET", "OPTIONS")
	router.HandleFunc("/moves/{id}", moves.MemberResource).Methods("GET", "OPTIONS")

	return router
}
<commit_msg>fix(routes): removing OPTIONS, because it literally did nothing here<commit_after>package main

import (
	"github.com/censhin/pokedex-api/moves"
	"github.com/censhin/pokedex-api/pokemon"
	"github.com/gorilla/mux"
)

func InitRoutes() *mux.Router {
	router := mux.NewRouter()

	router.HandleFunc("/pokemon", pokemon.CollectionResource).Methods("GET", "POST")
	router.HandleFunc("/pokemon/{id}", pokemon.MemberResource).Methods("GET", "PUT", "DELETE")
	router.HandleFunc("/moves", moves.CollectionResource).Methods("GET")
	router.HandleFunc("/moves/{id}", moves.MemberResource).Methods("GET")

	return router
}
<|endoftext|>


<commit_before>package boomer

import (
	"fmt"
	"log"
	"math/rand"
	"os"
	"runtime/debug"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

const (
	stateInit     = "ready"
	stateSpawning = "spawning"
	stateRunning  = "running"
	stateStopped  = "stopped"
	stateQuitting = "quitting"
)

const (
	slaveReportInterval = 3 * time.Second
	heartbeatInterval   = 1 * time.Second
)

type runner struct {
	state string

	tasks           []*Task
	totalTaskWeight int

	rateLimiter      RateLimiter
	rateLimitEnabled bool
	stats            *requestStats

	numClients int32
	spawnRate  float64

	// all running workers(goroutines) will select on this channel.
	// close this channel will stop all running workers.
	stopChan chan bool

	// close this channel will stop all goroutines used in runner.
	closeChan chan bool

	outputs []Output
}

// safeRun runs fn and recovers from unexpected panics.
// it prevents panics from Task.Fn crashing boomer.
func (r *runner) safeRun(fn func()) {
	defer func() {
		// don't panic
		err := recover()
		if err != nil {
			stackTrace := debug.Stack()
			errMsg := fmt.Sprintf("%v", err)
			os.Stderr.Write([]byte(errMsg))
			os.Stderr.Write([]byte("\n"))
			os.Stderr.Write(stackTrace)
		}
	}()
	fn()
}

func (r *runner) addOutput(o Output) {
	r.outputs = append(r.outputs, o)
}

func (r *runner) outputOnStart() {
	size := len(r.outputs)
	if size == 0 {
		return
	}
	wg := sync.WaitGroup{}
	wg.Add(size)
	for _, output := range r.outputs {
		go func(o Output) {
			o.OnStart()
			wg.Done()
		}(output)
	}
	wg.Wait()
}

func (r *runner) outputOnEevent(data map[string]interface{}) {
	size := len(r.outputs)
	if size == 0 {
		return
	}
	wg := sync.WaitGroup{}
	wg.Add(size)
	for _, output := range r.outputs {
		go func(o Output) {
			o.OnEvent(data)
			wg.Done()
		}(output)
	}
	wg.Wait()
}

func (r *runner) outputOnStop() {
	size := len(r.outputs)
	if size == 0 {
		return
	}
	wg := sync.WaitGroup{}
	wg.Add(size)
	for _, output := range r.outputs {
		go func(o Output) {
			o.OnStop()
			wg.Done()
		}(output)
	}
	wg.Wait()
}

func (r *runner) spawnWorkers(spawnCount int, quit chan bool, spawnCompleteFunc func()) {
	log.Println("Spawning", spawnCount, "clients at the rate", r.spawnRate, "clients/s...")

	defer func() {
		if spawnCompleteFunc != nil {
			spawnCompleteFunc()
		}
	}()

	for i := 1; i <= spawnCount; i++ {
		sleepTime := time.Duration(1000000/r.spawnRate) * time.Microsecond
		time.Sleep(sleepTime)

		select {
		case <-quit:
			// quit spawning goroutine
			return
		default:
			atomic.AddInt32(&r.numClients, 1)
			go func() {
				for {
					select {
					case <-quit:
						return
					default:
						if r.rateLimitEnabled {
							blocked := r.rateLimiter.Acquire()
							if !blocked {
								task := r.getTask()
								r.safeRun(task.Fn)
							}
						} else {
							task := r.getTask()
							r.safeRun(task.Fn)
						}
					}
				}
			}()
		}
	}
}

// setTasks will set the runner's task list AND the total task weight
// which is used to get a random task later
func (r *runner) setTasks(t []*Task) {
	r.tasks = t

	weightSum := 0
	for _, task := range r.tasks {
		weightSum += task.Weight
	}
	r.totalTaskWeight = weightSum
}

func (r *runner) getTask() *Task {
	tasksCount := len(r.tasks)
	if tasksCount == 1 {
		// Fast path
		return r.tasks[0]
	}

	rs := rand.New(rand.NewSource(time.Now().UnixNano()))

	totalWeight := r.totalTaskWeight
	if totalWeight <= 0 {
		// If all the tasks have not weights defined, they have the same chance to run
		randNum := rs.Intn(tasksCount)
		return r.tasks[randNum]
	}

	randNum := rs.Intn(totalWeight)
	runningSum := 0
	for _, task := range r.tasks {
		runningSum += task.Weight
		if runningSum > randNum {
			return task
		}
	}

	return nil
}

func (r *runner) startSpawning(spawnCount int, spawnRate float64, spawnCompleteFunc func()) {
	Events.Publish("boomer:hatch", spawnCount, spawnRate)
	Events.Publish("boomer:spawn", spawnCount, spawnRate)

	r.stats.clearStatsChan <- true
	r.stopChan = make(chan bool)

	r.spawnRate = spawnRate
	r.numClients = 0

	go r.spawnWorkers(spawnCount, r.stopChan, spawnCompleteFunc)
}

func (r *runner) stop() {
	// publish the boomer stop event
	// user's code can subscribe to this event and do thins like cleaning up
	Events.Publish("boomer:stop")

	// stop previous goroutines without blocking
	// those goroutines will exit when r.safeRun returns
	close(r.stopChan)
	if r.rateLimitEnabled {
		r.rateLimiter.Stop()
	}
}

type localRunner struct {
	runner

	spawnCount int
}

func newLocalRunner(tasks []*Task, rateLimiter RateLimiter, spawnCount int, spawnRate float64) (r *localRunner) {
	r = &localRunner{}
	r.setTasks(tasks)
	r.spawnRate = spawnRate
	r.spawnCount = spawnCount
	r.closeChan = make(chan bool)
	r.addOutput(NewConsoleOutput())

	if rateLimiter != nil {
		r.rateLimitEnabled = true
		r.rateLimiter = rateLimiter
	}

	r.stats = newRequestStats()
	return r
}

func (r *localRunner) run() {
	r.state = stateInit
	r.stats.start()

	wg := sync.WaitGroup{}
	wg.Add(1)
	go func() {
		for {
			select {
			case data := <-r.stats.messageToRunnerChan:
				data["user_count"] = r.numClients
				r.outputOnEevent(data)
			case <-r.closeChan:
				Events.Publish("boomer:quit")
				r.stop()
				wg.Done()
				return
			}
		}
	}()

	if r.rateLimitEnabled {
		r.rateLimiter.Start()
	}
	r.startSpawning(r.spawnCount, r.spawnRate, nil)

	wg.Wait()
}

func (r *localRunner) close() {
	if r.stats != nil {
		r.stats.close()
	}
	close(r.closeChan)
}

// SlaveRunner connects to the master, spawns goroutines and collects stats.
type slaveRunner struct {
	runner

	nodeID     string
	masterHost string
	masterPort int
	client     client
}

func newSlaveRunner(masterHost string, masterPort int, tasks []*Task, rateLimiter RateLimiter) (r *slaveRunner) {
	r = &slaveRunner{}
	r.masterHost = masterHost
	r.masterPort = masterPort
	r.setTasks(tasks)
	r.nodeID = getNodeID()
	r.closeChan = make(chan bool)

	if rateLimiter != nil {
		r.rateLimitEnabled = true
		r.rateLimiter = rateLimiter
	}

	r.stats = newRequestStats()
	return r
}

func (r *slaveRunner) spawnComplete() {
	data := make(map[string]interface{})
	data["count"] = r.numClients
	r.client.sendChannel() <- newMessage("spawning_complete", data, r.nodeID)
	r.state = stateRunning
}

func (r *slaveRunner) onQuiting() {
	if r.state != stateQuitting {
		r.client.sendChannel() <- newMessage("quit", nil, r.nodeID)
	}
}

func (r *slaveRunner) close() {
	if r.stats != nil {
		r.stats.close()
	}
	if r.client != nil {
		r.client.close()
	}
	close(r.closeChan)
}

func (r *slaveRunner) onSpawnMessage(msg *message) {
	r.client.sendChannel() <- newMessage("spawning", nil, r.nodeID)
	rate := msg.Data["spawn_rate"]
	users := msg.Data["num_users"]
	spawnRate := rate.(float64)
	workers := 0
	if _, ok := users.(uint64); ok {
		workers = int(users.(uint64))
	} else {
		workers = int(users.(int64))
	}

	if r.rateLimitEnabled {
		r.rateLimiter.Start()
	}
	r.startSpawning(workers, spawnRate, r.spawnComplete)
}

// Runner acts as a state machine.
func (r *slaveRunner) onMessage(msg *message) {
	if msg.Type == "hatch" {
		log.Println("The master sent a 'hatch' message, you are using an unsupported locust version, please update locust to 1.2.")
		return
	}

	switch r.state {
	case stateInit:
		switch msg.Type {
		case "spawn":
			r.state = stateSpawning
			r.onSpawnMessage(msg)
		case "quit":
			Events.Publish("boomer:quit")
		}
	case stateSpawning:
		fallthrough
	case stateRunning:
		switch msg.Type {
		case "spawn":
			r.state = stateSpawning
			r.stop()
			r.onSpawnMessage(msg)
		case "stop":
			r.stop()
			r.state = stateStopped
			log.Println("Recv stop message from master, all the goroutines are stopped")
			r.client.sendChannel() <- newMessage("client_stopped", nil, r.nodeID)
			r.client.sendChannel() <- newMessage("client_ready", nil, r.nodeID)
			r.state = stateInit
		case "quit":
			r.stop()
			log.Println("Recv quit message from master, all the goroutines are stopped")
			Events.Publish("boomer:quit")
			r.state = stateInit
		}
	case stateStopped:
		switch msg.Type {
		case "spawn":
			r.state = stateSpawning
			r.onSpawnMessage(msg)
		case "quit":
			Events.Publish("boomer:quit")
			r.state = stateInit
		}
	}
}

func (r *slaveRunner) startListener() {
	go func() {
		for {
			select {
			case msg := <-r.client.recvChannel():
				r.onMessage(msg)
			case <-r.closeChan:
				return
			}
		}
	}()
}

func (r *slaveRunner) run() {
	r.state = stateInit
	r.client = newClient(r.masterHost, r.masterPort, r.nodeID)

	err := r.client.connect()
	if err != nil {
		if strings.Contains(err.Error(), "Socket type DEALER is not compatible with PULL") {
			log.Println("Newer version of locust changes ZMQ socket to DEALER and ROUTER, you should update your locust version.")
		} else {
			log.Printf("Failed to connect to master(%s:%d) with error %v\n", r.masterHost, r.masterPort, err)
		}
		return
	}

	// listen to master
	r.startListener()

	r.stats.start()

	// tell master, I'm ready
	r.client.sendChannel() <- newMessage("client_ready", nil, r.nodeID)

	// report to master
	go func() {
		for {
			select {
			case data := <-r.stats.messageToRunnerChan:
				if r.state == stateInit || r.state == stateStopped {
					continue
				}
				data["user_count"] = r.numClients
				r.client.sendChannel() <- newMessage("stats", data, r.nodeID)
				r.outputOnEevent(data)
			case <-r.closeChan:
				return
			}
		}
	}()

	// heartbeat
	// See: https://github.com/locustio/locust/commit/a8c0d7d8c588f3980303358298870f2ea394ab93
	go func() {
		var ticker = time.NewTicker(heartbeatInterval)
		for {
			select {
			case <-ticker.C:
				CPUUsage := GetCurrentCPUUsage()
				data := map[string]interface{}{
					"state":             r.state,
					"current_cpu_usage": CPUUsage,
				}
				r.client.sendChannel() <- newMessage("heartbeat", data, r.nodeID)
			case <-r.closeChan:
				return
			}
		}
	}()

	Events.Subscribe("boomer:quit", r.onQuiting)
}
<commit_msg>FIX: panic on early stop<commit_after>package boomer

import (
	"fmt"
	"log"
	"math/rand"
	"os"
	"runtime/debug"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

const (
	stateInit     = "ready"
	stateSpawning = "spawning"
	stateRunning  = "running"
	stateStopped  = "stopped"
	stateQuitting = "quitting"
)

const (
	slaveReportInterval = 3 * time.Second
	heartbeatInterval   = 1 * time.Second
)

type runner struct {
	state string

	tasks           []*Task
	totalTaskWeight int

	rateLimiter      RateLimiter
	rateLimitEnabled bool
	stats            *requestStats

	numClients int32
	spawnRate  float64

	// all running workers(goroutines) will select on this channel.
	// close this channel will stop all running workers.
	stopChan chan bool

	// close this channel will stop all goroutines used in runner.
	closeChan chan bool

	outputs []Output
}

// safeRun runs fn and recovers from unexpected panics.
// it prevents panics from Task.Fn crashing boomer.
func (r *runner) safeRun(fn func()) {
	defer func() {
		// don't panic
		err := recover()
		if err != nil {
			stackTrace := debug.Stack()
			errMsg := fmt.Sprintf("%v", err)
			os.Stderr.Write([]byte(errMsg))
			os.Stderr.Write([]byte("\n"))
			os.Stderr.Write(stackTrace)
		}
	}()
	fn()
}

func (r *runner) addOutput(o Output) {
	r.outputs = append(r.outputs, o)
}

func (r *runner) outputOnStart() {
	size := len(r.outputs)
	if size == 0 {
		return
	}
	wg := sync.WaitGroup{}
	wg.Add(size)
	for _, output := range r.outputs {
		go func(o Output) {
			o.OnStart()
			wg.Done()
		}(output)
	}
	wg.Wait()
}

func (r *runner) outputOnEevent(data map[string]interface{}) {
	size := len(r.outputs)
	if size == 0 {
		return
	}
	wg := sync.WaitGroup{}
	wg.Add(size)
	for _, output := range r.outputs {
		go func(o Output) {
			o.OnEvent(data)
			wg.Done()
		}(output)
	}
	wg.Wait()
}

func (r *runner) outputOnStop() {
	size := len(r.outputs)
	if size == 0 {
		return
	}
	wg := sync.WaitGroup{}
	wg.Add(size)
	for _, output := range r.outputs {
		go func(o Output) {
			o.OnStop()
			wg.Done()
		}(output)
	}
	wg.Wait()
}

func (r *runner) spawnWorkers(spawnCount int, quit chan bool, spawnCompleteFunc func()) {
	log.Println("Spawning", spawnCount, "clients at the rate", r.spawnRate, "clients/s...")

	for i := 1; i <= spawnCount; i++ {
		sleepTime := time.Duration(1000000/r.spawnRate) * time.Microsecond
		time.Sleep(sleepTime)

		select {
		case <-quit:
			// quit spawning goroutine
			return
		default:
			atomic.AddInt32(&r.numClients, 1)
			go func() {
				for {
					select {
					case <-quit:
						return
					default:
						if r.rateLimitEnabled {
							blocked := r.rateLimiter.Acquire()
							if !blocked {
								task := r.getTask()
								r.safeRun(task.Fn)
							}
						} else {
							task := r.getTask()
							r.safeRun(task.Fn)
						}
					}
				}
			}()
		}
	}

	if spawnCompleteFunc != nil {
		spawnCompleteFunc()
	}
}

// setTasks will set the runner's task list AND the total task weight
// which is used to get a random task later
func (r *runner) setTasks(t []*Task) {
	r.tasks = t

	weightSum := 0
	for _, task := range r.tasks {
		weightSum += task.Weight
	}
	r.totalTaskWeight = weightSum
}

func (r *runner) getTask() *Task {
	tasksCount := len(r.tasks)
	if tasksCount == 1 {
		// Fast path
		return r.tasks[0]
	}

	rs := rand.New(rand.NewSource(time.Now().UnixNano()))

	totalWeight := r.totalTaskWeight
	if totalWeight <= 0 {
		// If all the tasks have not weights defined, they have the same chance to run
		randNum := rs.Intn(tasksCount)
		return r.tasks[randNum]
	}

	randNum := rs.Intn(totalWeight)
	runningSum := 0
	for _, task := range r.tasks {
		runningSum += task.Weight
		if runningSum > randNum {
			return task
		}
	}

	return nil
}

func (r *runner) startSpawning(spawnCount int, spawnRate float64, spawnCompleteFunc func()) {
	Events.Publish("boomer:hatch", spawnCount, spawnRate)
	Events.Publish("boomer:spawn", spawnCount, spawnRate)

	r.stats.clearStatsChan <- true
	r.stopChan = make(chan bool)

	r.spawnRate = spawnRate
	r.numClients = 0

	go r.spawnWorkers(spawnCount, r.stopChan, spawnCompleteFunc)
}

func (r *runner) stop() {
	// publish the boomer stop event
	// user's code can subscribe to this event and do thins like cleaning up
	Events.Publish("boomer:stop")

	// stop previous goroutines without blocking
	// those goroutines will exit when r.safeRun returns
	close(r.stopChan)
	if r.rateLimitEnabled {
		r.rateLimiter.Stop()
	}
}

type localRunner struct {
	runner

	spawnCount int
}

func newLocalRunner(tasks []*Task, rateLimiter RateLimiter, spawnCount int, spawnRate float64) (r *localRunner) {
	r = &localRunner{}
	r.setTasks(tasks)
	r.spawnRate = spawnRate
	r.spawnCount = spawnCount
	r.closeChan = make(chan bool)
	r.addOutput(NewConsoleOutput())

	if rateLimiter != nil {
		r.rateLimitEnabled = true
		r.rateLimiter = rateLimiter
	}

	r.stats = newRequestStats()
	return r
}

func (r *localRunner) run() {
	r.state = stateInit
	r.stats.start()

	wg := sync.WaitGroup{}
	wg.Add(1)
	go func() {
		for {
			select {
			case data := <-r.stats.messageToRunnerChan:
				data["user_count"] = r.numClients
				r.outputOnEevent(data)
			case <-r.closeChan:
				Events.Publish("boomer:quit")
				r.stop()
				wg.Done()
				return
			}
		}
	}()

	if r.rateLimitEnabled {
		r.rateLimiter.Start()
	}
	r.startSpawning(r.spawnCount, r.spawnRate, nil)

	wg.Wait()
}

func (r *localRunner) close() {
	if r.stats != nil {
		r.stats.close()
	}
	close(r.closeChan)
}

// SlaveRunner connects to the master, spawns goroutines and collects stats.
type slaveRunner struct {
	runner

	nodeID     string
	masterHost string
	masterPort int
	client     client
}

func newSlaveRunner(masterHost string, masterPort int, tasks []*Task, rateLimiter RateLimiter) (r *slaveRunner) {
	r = &slaveRunner{}
	r.masterHost = masterHost
	r.masterPort = masterPort
	r.setTasks(tasks)
	r.nodeID = getNodeID()
	r.closeChan = make(chan bool)

	if rateLimiter != nil {
		r.rateLimitEnabled = true
		r.rateLimiter = rateLimiter
	}

	r.stats = newRequestStats()
	return r
}

func (r *slaveRunner) spawnComplete() {
	data := make(map[string]interface{})
	data["count"] = r.numClients
	r.client.sendChannel() <- newMessage("spawning_complete", data, r.nodeID)
	r.state = stateRunning
}

func (r *slaveRunner) onQuiting() {
	if r.state != stateQuitting {
		r.client.sendChannel() <- newMessage("quit", nil, r.nodeID)
	}
}

func (r *slaveRunner) close() {
	if r.stats != nil {
		r.stats.close()
	}
	if r.client != nil {
		r.client.close()
	}
	close(r.closeChan)
}

func (r *slaveRunner) onSpawnMessage(msg *message) {
	r.client.sendChannel() <- newMessage("spawning", nil, r.nodeID)
	rate := msg.Data["spawn_rate"]
	users := msg.Data["num_users"]
	spawnRate := rate.(float64)
	workers := 0
	if _, ok := users.(uint64); ok {
		workers = int(users.(uint64))
	} else {
		workers = int(users.(int64))
	}

	if r.rateLimitEnabled {
		r.rateLimiter.Start()
	}
	r.startSpawning(workers, spawnRate, r.spawnComplete)
}

// Runner acts as a state machine.
func (r *slaveRunner) onMessage(msg *message) {
	if msg.Type == "hatch" {
		log.Println("The master sent a 'hatch' message, you are using an unsupported locust version, please update locust to 1.2.")
		return
	}

	switch r.state {
	case stateInit:
		switch msg.Type {
		case "spawn":
			r.state = stateSpawning
			r.onSpawnMessage(msg)
		case "quit":
			Events.Publish("boomer:quit")
		}
	case stateSpawning:
		fallthrough
	case stateRunning:
		switch msg.Type {
		case "spawn":
			r.state = stateSpawning
			r.stop()
			r.onSpawnMessage(msg)
		case "stop":
			r.stop()
			r.state = stateStopped
			log.Println("Recv stop message from master, all the goroutines are stopped")
			r.client.sendChannel() <- newMessage("client_stopped", nil, r.nodeID)
			r.client.sendChannel() <- newMessage("client_ready", nil, r.nodeID)
			r.state = stateInit
		case "quit":
			r.stop()
			log.Println("Recv quit message from master, all the goroutines are stopped")
			Events.Publish("boomer:quit")
			r.state = stateInit
		}
	case stateStopped:
		switch msg.Type {
		case "spawn":
			r.state = stateSpawning
			r.onSpawnMessage(msg)
		case "quit":
			Events.Publish("boomer:quit")
			r.state = stateInit
		}
	}
}

func (r *slaveRunner) startListener() {
	go func() {
		for {
			select {
			case msg := <-r.client.recvChannel():
				r.onMessage(msg)
			case <-r.closeChan:
				return
			}
		}
	}()
}

func (r *slaveRunner) run() {
	r.state = stateInit
	r.client = newClient(r.masterHost, r.masterPort, r.nodeID)

	err := r.client.connect()
	if err != nil {
		if strings.Contains(err.Error(), "Socket type DEALER is not compatible with PULL") {
			log.Println("Newer version of locust changes ZMQ socket to DEALER and ROUTER, you should update your locust version.")
		} else {
			log.Printf("Failed to connect to master(%s:%d) with error %v\n", r.masterHost, r.masterPort, err)
		}
		return
	}

	// listen to master
	r.startListener()

	r.stats.start()

	// tell master, I'm ready
	r.client.sendChannel() <- newMessage("client_ready", nil, r.nodeID)

	// report to master
	go func() {
		for {
			select {
			case data := <-r.stats.messageToRunnerChan:
				if r.state == stateInit || r.state == stateStopped {
					continue
				}
				data["user_count"] = r.numClients
				r.client.sendChannel() <- newMessage("stats", data, r.nodeID)
				r.outputOnEevent(data)
			case <-r.closeChan:
				return
			}
		}
	}()

	// heartbeat
	// See: https://github.com/locustio/locust/commit/a8c0d7d8c588f3980303358298870f2ea394ab93
	go func() {
		var ticker = time.NewTicker(heartbeatInterval)
		for {
			select {
			case <-ticker.C:
				CPUUsage := GetCurrentCPUUsage()
				data := map[string]interface{}{
					"state":             r.state,
					"current_cpu_usage": CPUUsage,
				}
				r.client.sendChannel() <- newMessage("heartbeat", data, r.nodeID)
			case <-r.closeChan:
				return
			}
		}
	}()

	Events.Subscribe("boomer:quit", r.onQuiting)
}
<|endoftext|>


<commit_before>package main

import (
    "fmt"
    "sync"

    "github.com/aws/aws-sdk-go/aws"
    "github.com/aws/aws-sdk-go/aws/session"
    "github.com/aws/aws-sdk-go/service/s3"
)

const (
  // How many Keys AWS-S3-Api will return in a single request. Note that S3-API
  // has a hard limit of 1000
  // http://docs.aws.amazon.com/AmazonS3/latest/API/RESTBucketGET.html
  S3_MAX_KEYS = 1000
  // How many S3-Result-Sets are buffered until the producer blocks
  CHAN_BUFFER_SIZE = 1000
)

var wg sync.WaitGroup

// Finds every element in a S3-Bucket
// Is meant to run in parallel. Puts resultsets into channel *s3.ListObjectsOutput
// Closes this channel when done
// First parameter is an *s3.S3 Object (http://docs.aws.amazon.com/sdk-for-go/api/service/s3/S3.html)
// Second is the S3-Bucket-Name (Note not the S3-Path like s3://bucket_name)
// Third is the result channel of type *s3.ListObjectsOutput
func list_elements(svc *s3.S3, bucket_name string, result_chan chan *s3.ListObjectsOutput) {
  defer wg.Done()
  defer close(result_chan)

  params := &s3.ListObjectsInput{
    Bucket:       aws.String(bucket_name), // Required
    MaxKeys:      aws.Int64(S3_MAX_KEYS),
  }

  err := svc.ListObjectsPages(params, func(page *s3.ListObjectsOutput, lastPage bool) bool {
    result_chan <- page
    return !lastPage
  })

  if err != nil {
    // Print the error, cast err to awserr.Error to get the Code and
    // Message from an error.
    fmt.Println(err.Error())
    return
  }

  fmt.Printf("%s done\n", bucket_name)
}

func main() {
    // Create an S3 service object in the "eu-west-1" region
    // Note that you can also configure your region globally by
    // exporting the AWS_REGION environment variable
    svc := s3.New(session.New(), &aws.Config{Region: aws.String("eu-west-1")})

    wg.Add(2)

    result_chan_src := make(chan *s3.ListObjectsOutput, CHAN_BUFFER_SIZE)
    result_chan_dest := make(chan *s3.ListObjectsOutput, CHAN_BUFFER_SIZE)

    go list_elements(svc, "knorrtestx1", result_chan_src)
    go list_elements(svc, "knorrtestx2", result_chan_dest)

    wg.Wait()

}
<commit_msg>Basic algorithm done of how to determine difference between Source- and Destination-Bucket<commit_after>package main

import (
    "fmt"
    "sync"

    "github.com/aws/aws-sdk-go/aws"
    "github.com/aws/aws-sdk-go/aws/session"
    "github.com/aws/aws-sdk-go/service/s3"
)

const (
  // How many Keys AWS-S3-Api will return in a single request. Note that S3-API
  // has a hard limit of 1000
  // http://docs.aws.amazon.com/AmazonS3/latest/API/RESTBucketGET.html
  S3_MAX_KEYS = 1000
  // How many S3-Result-Sets are buffered until the producer blocks
  CHAN_BUFFER_SIZE = 1000
)

var wg sync.WaitGroup

// Finds every element in a S3-Bucket
// Is meant to run in parallel. Puts resultsets into channel *s3.ListObjectsOutput
// Closes this channel when done
// First parameter is an *s3.S3 Object (http://docs.aws.amazon.com/sdk-for-go/api/service/s3/S3.html)
// Second is the S3-Bucket-Name (Note not the S3-Path like s3://bucket_name)
// Third is the result channel of type *s3.ListObjectsOutput
func list_elements(svc *s3.S3, bucket_name string, result_chan chan *s3.ListObjectsOutput) {
  defer wg.Done()
  defer close(result_chan)

  params := &s3.ListObjectsInput{
    Bucket:       aws.String(bucket_name), // Required
    MaxKeys:      aws.Int64(S3_MAX_KEYS),
  }

  err := svc.ListObjectsPages(params, func(page *s3.ListObjectsOutput, lastPage bool) bool {
    result_chan <- page
    return !lastPage
  })

  if err != nil {
    // Print the error, cast err to awserr.Error to get the Code and
    // Message from an error.
    fmt.Println(err.Error())
    return
  }

  //fmt.Printf("%s done\n", bucket_name)
}

// Need to convert Chungs of Keys, returned from S3-API into lists
func extract_contents(result_chan chan *s3.ListObjectsOutput, output_chan chan *s3.Object) {
  defer wg.Done()
  for {
    s3_list_chunk, ok := <- result_chan
    if !ok {
      // Channel is closed bei producer
      break
    }

    for _, s3_key := range s3_list_chunk.Contents {
      //fmt.Printf("Key: %v\n", s3_key)
      output_chan <- s3_key
    }

    close(output_chan)

  }
}

// Detects difference between Source- and Destination-Bucket
func find_missing(s3_contents_src, s3_contents_dest chan *s3.Object) {
  defer wg.Done()
  var src, dest *s3.Object

  src = <- s3_contents_src
  dest = <- s3_contents_dest

  for {

    if src == nil && dest == nil {
      // Case A: no elements left in both Buckets
      break

    } else if src != nil && dest != nil {
      // Case B: got an element in both Buckets

      if *src.Key < *dest.Key {
        // src-element is ahead of destination-element. Src-element is missing in Destination-Bucket
        src = <- s3_contents_src

      } else if *src.Key > *dest.Key {
        // src-element is beyond destination-element. Destination-element exists only on Destination-Bucket
        dest = <- s3_contents_dest

      } else {  // src.Key == dest.Key
        // Src and Dest are equal
        src = <- s3_contents_src
        dest = <- s3_contents_dest

      }

    } else if dest == nil {
      // Case C: no element left in Destination-Bucket but in Source-Bucket
      src = <- s3_contents_src

    } else if src == nil {
      // Case D: no elements left in Source-Bucket but in Destination-Bucket
      dest = <- s3_contents_dest

    } else {
      // This means not Case A or not Case B or not Case C or not Case D
      fmt.Printf("This case shouldn't appear. However it did. Src: %v Dest: %v", src, dest)
      // TODO: raise an error. Stop execution
      break

    }
  }
}

func main() {
    // Create an S3 service object in the "eu-west-1" region
    // Note that you can also configure your region globally by
    // exporting the AWS_REGION environment variable
    svc := s3.New(session.New(), &aws.Config{Region: aws.String("eu-west-1")})

    wg.Add(5)

    result_chan_src := make(chan *s3.ListObjectsOutput, CHAN_BUFFER_SIZE)
    result_chan_dest := make(chan *s3.ListObjectsOutput, CHAN_BUFFER_SIZE)
    s3_contents_src := make(chan *s3.Object, CHAN_BUFFER_SIZE)
    s3_contents_dest := make(chan *s3.Object, CHAN_BUFFER_SIZE)

    go list_elements(svc, "knorrtestx1", result_chan_src)
    go list_elements(svc, "knorrtestx2", result_chan_dest)
    go extract_contents(result_chan_src, s3_contents_src)
    go extract_contents(result_chan_dest, s3_contents_dest)
    go find_missing(s3_contents_src, s3_contents_dest)

    wg.Wait()

}
<|endoftext|>


<commit_before>// Copyright (c) 2013-2016 by Michael Dvorkin. All Rights Reserved.
// Use of this source code is governed by a MIT-style license that can
// be found in the LICENSE file.

package mop

import (
	`github.com/michaeldv/termbox-go`
	`strings`
	`time`
)

// Screen is thin wrapper aroung Termbox library to provide basic display
// capabilities as requied by Mop.
type Screen struct {
	width    int        // Current number of columns.
	height   int        // Current number of rows.
	cleared  bool       // True after the screens gets cleared.
	layout   *Layout    // Pointer to layout (gets created by screen).
	markup   *Markup    // Pointer to markup processor (gets created by screen).
	pausedAt *time.Time // Timestamp of the pause request or nil if none.
}

// Initializes Termbox, creates screen along with layout and markup, and
// calculates current screen dimensions. Once initialized the screen is
// ready for display.
func NewScreen() *Screen {
	if err := termbox.Init(); err != nil {
		panic(err)
	}
	screen := &Screen{}
	screen.layout = NewLayout()
	screen.markup = NewMarkup()

	return screen.Resize()
}

// Close gets called upon program termination to close the Termbox.
func (screen *Screen) Close() *Screen {
	termbox.Close()

	return screen
}

// Resize gets called when the screen is being resized. It recalculates screen
// dimensions and requests to clear the screen on next update.
func (screen *Screen) Resize() *Screen {
	screen.width, screen.height = termbox.Size()
	screen.cleared = false

	return screen
}

// Pause is a toggle function that either creates a timestamp of the pause
// request or resets it to nil.
func (screen *Screen) Pause(pause bool) *Screen {
	if pause {
		screen.pausedAt = new(time.Time)
		*screen.pausedAt = time.Now()
	} else {
		screen.pausedAt = nil
	}

	return screen
}

// Clear makes the entire screen blank using default background color.
func (screen *Screen) Clear() *Screen {
	termbox.Clear(termbox.ColorDefault, termbox.ColorDefault)
	screen.cleared = true

	return screen
}

// ClearLine erases the contents of the line starting from (x,y) coordinate
// till the end of the line.
func (screen *Screen) ClearLine(x int, y int) *Screen {
	for i := x; i < screen.width; i++ {
		termbox.SetCell(i, y, ' ', termbox.ColorDefault, termbox.ColorDefault)
	}
	termbox.Flush()

	return screen
}

// Draw accepts variable number of arguments and knows how to display the
// market data, stock quotes, current time, and an arbitrary string.
func (screen *Screen) Draw(objects ...interface{}) *Screen {
	if screen.pausedAt != nil {
		defer screen.DrawLine(0, 0, `<right><r>`+screen.pausedAt.Format(`3:04:05pm PST`)+`</r></right>`)
	}
	for _, ptr := range objects {
		switch ptr.(type) {
		case *Market:
			object := ptr.(*Market)
			screen.draw(screen.layout.Market(object.Fetch()))
		case *Quotes:
			object := ptr.(*Quotes)
			screen.draw(screen.layout.Quotes(object.Fetch()))
		case time.Time:
			timestamp := ptr.(time.Time).Format(`3:04:05pm PST`)
			screen.DrawLine(0, 0, `<right>`+timestamp+`</right>`)
		default:
			screen.draw(ptr.(string))
		}
	}

	return screen
}

// DrawLine takes the incoming string, tokenizes it to extract markup
// elements, and displays it all starting at (x,y) location.
func (screen *Screen) DrawLine(x int, y int, str string) {
	start, column := 0, 0

	for _, token := range screen.markup.Tokenize(str) {
		// First check if it's a tag. Tags are eaten up and not displayed.
		if screen.markup.IsTag(token) {
			continue
		}

		// Here comes the actual text: display it one character at a time.
		for i, char := range token {
			if !screen.markup.RightAligned {
				start = x + column
				column++
			} else {
				start = screen.width - len(token) + i
			}
			termbox.SetCell(start, y, char, screen.markup.Foreground, screen.markup.Background)
		}
	}
	termbox.Flush()
}

// Underlying workhorse function that takes multiline string, splits it into
// lines, and displays them row by row.
func (screen *Screen) draw(str string) {
	if !screen.cleared {
		screen.Clear()
	}
	for row, line := range strings.Split(str, "\n") {
		screen.DrawLine(0, row, line)
	}
}
<commit_msg>Display correct time zone name<commit_after>// Copyright (c) 2013-2016 by Michael Dvorkin. All Rights Reserved.
// Use of this source code is governed by a MIT-style license that can
// be found in the LICENSE file.

package mop

import (
	`github.com/michaeldv/termbox-go`
	`strings`
	`time`
)

// Screen is thin wrapper aroung Termbox library to provide basic display
// capabilities as requied by Mop.
type Screen struct {
	width    int        // Current number of columns.
	height   int        // Current number of rows.
	cleared  bool       // True after the screens gets cleared.
	layout   *Layout    // Pointer to layout (gets created by screen).
	markup   *Markup    // Pointer to markup processor (gets created by screen).
	pausedAt *time.Time // Timestamp of the pause request or nil if none.
}

// Initializes Termbox, creates screen along with layout and markup, and
// calculates current screen dimensions. Once initialized the screen is
// ready for display.
func NewScreen() *Screen {
	if err := termbox.Init(); err != nil {
		panic(err)
	}
	screen := &Screen{}
	screen.layout = NewLayout()
	screen.markup = NewMarkup()

	return screen.Resize()
}

// Close gets called upon program termination to close the Termbox.
func (screen *Screen) Close() *Screen {
	termbox.Close()

	return screen
}

// Resize gets called when the screen is being resized. It recalculates screen
// dimensions and requests to clear the screen on next update.
func (screen *Screen) Resize() *Screen {
	screen.width, screen.height = termbox.Size()
	screen.cleared = false

	return screen
}

// Pause is a toggle function that either creates a timestamp of the pause
// request or resets it to nil.
func (screen *Screen) Pause(pause bool) *Screen {
	if pause {
		screen.pausedAt = new(time.Time)
		*screen.pausedAt = time.Now()
	} else {
		screen.pausedAt = nil
	}

	return screen
}

// Clear makes the entire screen blank using default background color.
func (screen *Screen) Clear() *Screen {
	termbox.Clear(termbox.ColorDefault, termbox.ColorDefault)
	screen.cleared = true

	return screen
}

// ClearLine erases the contents of the line starting from (x,y) coordinate
// till the end of the line.
func (screen *Screen) ClearLine(x int, y int) *Screen {
	for i := x; i < screen.width; i++ {
		termbox.SetCell(i, y, ' ', termbox.ColorDefault, termbox.ColorDefault)
	}
	termbox.Flush()

	return screen
}

// Draw accepts variable number of arguments and knows how to display the
// market data, stock quotes, current time, and an arbitrary string.
func (screen *Screen) Draw(objects ...interface{}) *Screen {
	zonename, _ := time.Now().In(time.Local).Zone()
	if screen.pausedAt != nil {
		defer screen.DrawLine(0, 0, `<right><r>`+screen.pausedAt.Format(`3:04:05pm ` + zonename)+`</r></right>`)
	}
	for _, ptr := range objects {
		switch ptr.(type) {
		case *Market:
			object := ptr.(*Market)
			screen.draw(screen.layout.Market(object.Fetch()))
		case *Quotes:
			object := ptr.(*Quotes)
			screen.draw(screen.layout.Quotes(object.Fetch()))
		case time.Time:
			timestamp := ptr.(time.Time).Format(`3:04:05pm ` + zonename)
			screen.DrawLine(0, 0, `<right>`+timestamp+`</right>`)
		default:
			screen.draw(ptr.(string))
		}
	}

	return screen
}

// DrawLine takes the incoming string, tokenizes it to extract markup
// elements, and displays it all starting at (x,y) location.
func (screen *Screen) DrawLine(x int, y int, str string) {
	start, column := 0, 0

	for _, token := range screen.markup.Tokenize(str) {
		// First check if it's a tag. Tags are eaten up and not displayed.
		if screen.markup.IsTag(token) {
			continue
		}

		// Here comes the actual text: display it one character at a time.
		for i, char := range token {
			if !screen.markup.RightAligned {
				start = x + column
				column++
			} else {
				start = screen.width - len(token) + i
			}
			termbox.SetCell(start, y, char, screen.markup.Foreground, screen.markup.Background)
		}
	}
	termbox.Flush()
}

// Underlying workhorse function that takes multiline string, splits it into
// lines, and displays them row by row.
func (screen *Screen) draw(str string) {
	if !screen.cleared {
		screen.Clear()
	}
	for row, line := range strings.Split(str, "\n") {
		screen.DrawLine(0, row, line)
	}
}
<|endoftext|>


<commit_before>package gocql

import (
	"fmt"
	"math"
	"strconv"
	"sync"
	"sync/atomic"
)

// scyllaSupported represents Scylla connection options as sent in SUPPORTED
// frame.
type scyllaSupported struct {
	shard     int
	nrShards  int
	msbIgnore uint64
}

func parseSupported(supported map[string][]string) scyllaSupported {
	const (
		scyllaShard             = "SCYLLA_SHARD"
		scyllaNrShards          = "SCYLLA_NR_SHARDS"
		scyllaPartitioner       = "SCYLLA_PARTITIONER"
		scyllaShardingAlgorithm = "SCYLLA_SHARDING_ALGORITHM"
		scyllaShardingIgnoreMSB = "SCYLLA_SHARDING_IGNORE_MSB"
	)

	var (
		si  scyllaSupported
		err error
	)

	if s, ok := supported[scyllaShard]; ok {
		if si.shard, err = strconv.Atoi(s[0]); err != nil {
			if GoCQLDebug {
				Logger.Printf("scylla: failed to parse %s value %v: %s", scyllaShard, s, err)
			}
		}
	}
	if s, ok := supported[scyllaNrShards]; ok {
		if si.nrShards, err = strconv.Atoi(s[0]); err != nil {
			if GoCQLDebug {
				Logger.Printf("scylla: failed to parse %s value %v: %s", scyllaNrShards, s, err)
			}
		}
	}
	if s, ok := supported[scyllaShardingIgnoreMSB]; ok {
		if si.msbIgnore, err = strconv.ParseUint(s[0], 10, 64); err != nil {
			if GoCQLDebug {
				Logger.Printf("scylla: failed to parse %s value %v: %s", scyllaShardingIgnoreMSB, s, err)
			}
		}
	}

	var (
		partitioner string
		algorithm   string
	)
	if s, ok := supported[scyllaPartitioner]; ok {
		partitioner = s[0]
	}
	if s, ok := supported[scyllaShardingAlgorithm]; ok {
		algorithm = s[0]
	}

	if partitioner != "org.apache.cassandra.dht.Murmur3Partitioner" || algorithm != "biased-token-round-robin" || si.nrShards == 0 || si.msbIgnore == 0 {
		if GoCQLDebug {
			Logger.Printf("scylla: unsupported sharding configuration")
		}
		return scyllaSupported{}
	}

	return si
}

// isScyllaConn checks if conn is suitable for scyllaConnPicker.
func isScyllaConn(conn *Conn) bool {
	s := parseSupported(conn.supported)
	return s.nrShards != 0
}

// scyllaConnPicker is a specialised ConnPicker that selects connections based
// on token trying to get connection to a shard containing the given token.
type scyllaConnPicker struct {
	muConnsPerShard sync.RWMutex
	connsPerShard   [][]*Conn

	nrConns   int
	nrShards  int
	msbIgnore uint64
	pos       int32
}

func newScyllaConnPicker(conn *Conn) *scyllaConnPicker {
	s := parseSupported(conn.supported)
	if s.nrShards == 0 {
		panic(fmt.Sprintf("scylla: %s not a sharded connection", conn.Address()))
	}

	if GoCQLDebug {
		Logger.Printf("scylla: %s sharding options %+v", conn.Address(), s)
	}
	v := &scyllaConnPicker{
		nrShards:      s.nrShards,
		msbIgnore:     s.msbIgnore,
		connsPerShard: make([][]*Conn, s.nrShards),
	}
	return v
}

func (p *scyllaConnPicker) Remove(conn *Conn) {
	s := parseSupported(conn.supported)
	if s.nrShards == 0 {
		// It is possible for Remove to be called before the connection is added to the pool.
		// Ignoring these connections here is safe.
		if GoCQLDebug {
			Logger.Printf("scylla: %s has unknown sharding state, ignoring it", conn.Address())
		}
		return
	}
	if GoCQLDebug {
		Logger.Printf("scylla: %s remove shard %d connection", conn.Address(), s.shard)
	}

	p.muConnsPerShard.Lock()
	if len(p.connsPerShard[s.shard]) > 0 {
		p.nrConns--
	}
	p.connsPerShard[s.shard] = p.connsPerShard[s.shard][:]
	p.muConnsPerShard.Unlock()
}

func (p *scyllaConnPicker) Close() {
	connsPerShard := p.connsPerShard
	p.connsPerShard = nil

	p.muConnsPerShard.RLock()
	for _, conns := range connsPerShard {
		for _, c := range conns {
			c.Close()
		}
	}
	p.muConnsPerShard.RUnlock()
}

func (p *scyllaConnPicker) Size() (int, int) {
	return p.nrConns, p.nrShards - p.nrConns
}

func (p *scyllaConnPicker) Pick(t token) *Conn {
	if p.nrConns == 0 {
		return nil
	}

	if t == nil {
		idx := int(atomic.AddInt32(&p.pos, 1))
		conn := (*Conn)(nil)
		p.muConnsPerShard.RLock()
		for i := range p.connsPerShard {
			conns := p.connsPerShard[(idx+i)%len(p.connsPerShard)]
			if len(conns) > 0 {
				conn = conns[0]
				break
			}
		}
		p.muConnsPerShard.RUnlock()
		return conn
	}

	mmt, ok := t.(murmur3Token)
	// double check if that's murmur3 token
	if !ok {
		return nil
	}

	idx := p.shardOf(mmt)

	conn := (*Conn)(nil)
	p.muConnsPerShard.RLock()
	conns := p.connsPerShard[idx]
	if len(conns) > 0 {
		conn = conns[0]
	}
	p.muConnsPerShard.RUnlock()
	return conn
}

func (p *scyllaConnPicker) shardOf(token murmur3Token) int {
	shards := uint64(p.nrShards)
	z := uint64(token+math.MinInt64) << p.msbIgnore
	lo := z & 0xffffffff
	hi := (z >> 32) & 0xffffffff
	mul1 := lo * shards
	mul2 := hi * shards
	sum := (mul1 >> 32) + mul2
	return int(sum >> 32)
}

func (p *scyllaConnPicker) Put(conn *Conn) {
	s := parseSupported(conn.supported)
	if s.nrShards == 0 {
		panic(fmt.Sprintf("scylla: %s not a sharded connection", conn.Address()))
	}
	if s.nrShards != p.nrShards {
		panic(fmt.Sprintf("scylla: %s invalid number of shards", conn.Address()))
	}

	p.muConnsPerShard.Lock()
	conns := p.connsPerShard[s.shard]
	wasEmpty := len(conns) == 0
	conns = append(conns, conn)
	p.connsPerShard[s.shard] = conns
	if wasEmpty {
		p.nrConns++
	}
	if p.nrConns == p.nrShards {
		for i, conns := range p.connsPerShard {
			for _, c := range conns[1:] {
				c.Close()
			}
			p.connsPerShard[i] = conns[:1] // keep the first connection
		}
	}
	p.muConnsPerShard.Unlock()

	if GoCQLDebug {
		Logger.Printf("scylla: %s put shard %d connection total: %d missing: %d", conn.Address(), s.shard, p.nrConns, p.nrShards-p.nrConns)
	}
}
<commit_msg>Revert "scylla: keep at least one connection by shard"<commit_after>package gocql

import (
	"fmt"
	"math"
	"strconv"
	"sync/atomic"
)

// scyllaSupported represents Scylla connection options as sent in SUPPORTED
// frame.
type scyllaSupported struct {
	shard     int
	nrShards  int
	msbIgnore uint64
}

func parseSupported(supported map[string][]string) scyllaSupported {
	const (
		scyllaShard             = "SCYLLA_SHARD"
		scyllaNrShards          = "SCYLLA_NR_SHARDS"
		scyllaPartitioner       = "SCYLLA_PARTITIONER"
		scyllaShardingAlgorithm = "SCYLLA_SHARDING_ALGORITHM"
		scyllaShardingIgnoreMSB = "SCYLLA_SHARDING_IGNORE_MSB"
	)

	var (
		si  scyllaSupported
		err error
	)

	if s, ok := supported[scyllaShard]; ok {
		if si.shard, err = strconv.Atoi(s[0]); err != nil {
			if GoCQLDebug {
				Logger.Printf("scylla: failed to parse %s value %v: %s", scyllaShard, s, err)
			}
		}
	}
	if s, ok := supported[scyllaNrShards]; ok {
		if si.nrShards, err = strconv.Atoi(s[0]); err != nil {
			if GoCQLDebug {
				Logger.Printf("scylla: failed to parse %s value %v: %s", scyllaNrShards, s, err)
			}
		}
	}
	if s, ok := supported[scyllaShardingIgnoreMSB]; ok {
		if si.msbIgnore, err = strconv.ParseUint(s[0], 10, 64); err != nil {
			if GoCQLDebug {
				Logger.Printf("scylla: failed to parse %s value %v: %s", scyllaShardingIgnoreMSB, s, err)
			}
		}
	}

	var (
		partitioner string
		algorithm   string
	)
	if s, ok := supported[scyllaPartitioner]; ok {
		partitioner = s[0]
	}
	if s, ok := supported[scyllaShardingAlgorithm]; ok {
		algorithm = s[0]
	}

	if partitioner != "org.apache.cassandra.dht.Murmur3Partitioner" || algorithm != "biased-token-round-robin" || si.nrShards == 0 || si.msbIgnore == 0 {
		if GoCQLDebug {
			Logger.Printf("scylla: unsupported sharding configuration")
		}
		return scyllaSupported{}
	}

	return si
}

// isScyllaConn checks if conn is suitable for scyllaConnPicker.
func isScyllaConn(conn *Conn) bool {
	s := parseSupported(conn.supported)
	return s.nrShards != 0
}

// scyllaConnPicker is a specialised ConnPicker that selects connections based
// on token trying to get connection to a shard containing the given token.
type scyllaConnPicker struct {
	conns     []*Conn
	nrConns   int
	nrShards  int
	msbIgnore uint64
	pos       int32
}

func newScyllaConnPicker(conn *Conn) *scyllaConnPicker {
	s := parseSupported(conn.supported)
	if s.nrShards == 0 {
		panic(fmt.Sprintf("scylla: %s not a sharded connection", conn.Address()))
	}

	if GoCQLDebug {
		Logger.Printf("scylla: %s sharding options %+v", conn.Address(), s)
	}

	return &scyllaConnPicker{
		nrShards:  s.nrShards,
		msbIgnore: s.msbIgnore,
	}
}

func (p *scyllaConnPicker) Remove(conn *Conn) {
	s := parseSupported(conn.supported)
	if s.nrShards == 0 {
		// It is possible for Remove to be called before the connection is added to the pool.
		// Ignoring these connections here is safe.
		if GoCQLDebug {
			Logger.Printf("scylla: %s has unknown sharding state, ignoring it", conn.Address())
		}
		return
	}
	if GoCQLDebug {
		Logger.Printf("scylla: %s remove shard %d connection", conn.Address(), s.shard)
	}

	if p.conns[s.shard] != nil {
		p.conns[s.shard] = nil
		p.nrConns--
	}
}

func (p *scyllaConnPicker) Close() {
	conns := p.conns
	p.conns = nil
	for _, conn := range conns {
		if conn != nil {
			conn.Close()
		}
	}
}

func (p *scyllaConnPicker) Size() (int, int) {
	return p.nrConns, p.nrShards - p.nrConns
}

func (p *scyllaConnPicker) Pick(t token) *Conn {
	if len(p.conns) == 0 {
		return nil
	}

	if t == nil {
		idx := int(atomic.AddInt32(&p.pos, 1))
		for i := 0; i < len(p.conns); i++ {
			if conn := p.conns[(idx+i)%len(p.conns)]; conn != nil {
				return conn
			}
		}
		return nil
	}

	mmt, ok := t.(murmur3Token)
	// double check if that's murmur3 token
	if !ok {
		return nil
	}

	idx := p.shardOf(mmt)
	return p.conns[idx]
}

func (p *scyllaConnPicker) shardOf(token murmur3Token) int {
	shards := uint64(p.nrShards)
	z := uint64(token+math.MinInt64) << p.msbIgnore
	lo := z & 0xffffffff
	hi := (z >> 32) & 0xffffffff
	mul1 := lo * shards
	mul2 := hi * shards
	sum := (mul1 >> 32) + mul2
	return int(sum >> 32)
}

func (p *scyllaConnPicker) Put(conn *Conn) {
	s := parseSupported(conn.supported)
	if s.nrShards == 0 {
		panic(fmt.Sprintf("scylla: %s not a sharded connection", conn.Address()))
	}

	if s.nrShards != len(p.conns) {
		if s.nrShards != p.nrShards {
			panic(fmt.Sprintf("scylla: %s invalid number of shards", conn.Address()))
		}
		conns := p.conns
		p.conns = make([]*Conn, s.nrShards, s.nrShards)
		copy(p.conns, conns)
	}
	if c := p.conns[s.shard]; c != nil {
		conn.Close()
		return
	}
	p.conns[s.shard] = conn
	p.nrConns++
	if GoCQLDebug {
		Logger.Printf("scylla: %s put shard %d connection total: %d missing: %d", conn.Address(), s.shard, p.nrConns, p.nrShards-p.nrConns)
	}
}
<|endoftext|>


<commit_before>package main

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

	"github.com/intervention-engine/fhir/server"
	"github.com/intervention-engine/ie/controllers"
	"github.com/intervention-engine/ie/huddles"
	"github.com/intervention-engine/ie/utilities"
	"github.com/robfig/cron"
)

// Setup the huddle flag to accumulate huddle config file paths
type huddleSlice []string

func (h *huddleSlice) String() string {
	return fmt.Sprint(*h)
}
func (h *huddleSlice) Set(value string) error {
	*h = append(*h, value)
	return nil
}

var huddleFlag huddleSlice

func init() {
	flag.Var(&huddleFlag, "huddle", "path to a huddle configuration file (repeat flag for multiple paths)")
}

func main() {
	// Check for a linked MongoDB container if we are running in Docker
	mongoHost := os.Getenv("MONGO_PORT_27017_TCP_ADDR")
	if mongoHost == "" {
		mongoHost = "localhost"
	}

	codeLookupFlag := flag.Bool("loadCodes", false, "flag to enable download of icd-9 and icd-10 code lookup")
	icd9URL := flag.String("icd9URL", "https://www.cms.gov/Medicare/Coding/ICD9ProviderDiagnosticCodes/Downloads/ICD-9-CM-v32-master-descriptions.zip", "url for icd-9 code definition zip")
	icd10URL := flag.String("icd10URL", "https://www.cms.gov/Medicare/Coding/ICD10/Downloads/2016-Code-Descriptions-in-Tabular-Order.zip", "url for icd-10 code definition zip")
	flag.Parse()

	if *codeLookupFlag {
		utilities.LoadICDFromCMS(mongoHost, *icd9URL, "ICD-9")
		utilities.LoadICDFromCMS(mongoHost, *icd10URL, "ICD-10")
	}

	riskServiceEndpoint := os.Getenv("RISKSERVICE_PORT_9000_TCP_ADDR")
	if riskServiceEndpoint == "" {
		riskServiceEndpoint = "http://localhost:9000"
	} else {
		riskServiceEndpoint = "http://" + riskServiceEndpoint + ":9000"
	}

	subscriptionFlag := flag.Bool("subscriptions", false, "enables limited support for resource subscriptions (default: false)")

	var ip net.IP
	var selfURL string
	host, err := os.Hostname()
	if err != nil {
		panic(err)
	}
	addrs, err := net.LookupIP(host)
	if err != nil {
		fmt.Println("Unable to lookup IP based on hostname, defaulting to localhost.")
		selfURL = "http://localhost:3001"
	}
	for _, addr := range addrs {
		if ipv4 := addr.To4(); ipv4 != nil {
			ip = ipv4
			selfURL = "http://" + ip.String() + ":3001"
		}
	}

	s := server.NewServer(mongoHost)

	// Since the huddle controller needs info from the command line, set it up here.  When IE is refactored
	// to take out globals (and other stuff), this should be rethought.
	c := cron.New()
	huddleController := new(huddles.HuddleSchedulerController)
	for _, huddlePath := range huddleFlag {
		data, err := ioutil.ReadFile(huddlePath)
		if err != nil {
			panic(err)
		}
		config := new(huddles.HuddleConfig)
		if err := json.Unmarshal(data, config); err != nil {
			panic(err)
		}
		huddleController.AddConfig(config)

		// Wait 1 minute before doing initial runs or setting up cron jobs.  This allows the server to get
		// started (since it needs to initiate the db connection, etc).
		time.AfterFunc(1*time.Minute, func() {
			// Do an initial run
			fmt.Println("Initial scheduling for huddle with name ", config.Name)
			huddles.ScheduleHuddles(config)

			// Set up the cron job for future runs
			if config.SchedulerCronSpec != "" {
				err := c.AddFunc(config.SchedulerCronSpec, func() {
					_, err := huddles.ScheduleHuddles(config)
					if err != nil {
						fmt.Printf("ERROR: Could not schedule huddles for huddle with name %s: %v", config.Name, err)
					}
				})
				if err != nil {
					panic(err)
				}
				fmt.Printf("Huddle with name %s scheduled with cron spec: %s\n", config.Name, config.SchedulerCronSpec)
			} else {
				fmt.Printf("Warning: Huddle with name %s is not configured with a scheduler cron job.\n", config.Name)
			}
		})
	}
	s.Engine.GET("/ScheduleHuddles", huddleController.ScheduleHandler)

	if len(huddleFlag) > 0 {
		c.Start()
		defer c.Stop()
	}

	closer := controllers.RegisterRoutes(s, selfURL, riskServiceEndpoint, *subscriptionFlag)
	defer closer()

	s.Run(server.Config{
		UseSmartAuth:         false,
		UseLoggingMiddleware: true,
	})
}
<commit_msg>Add flag to turn on request logging<commit_after>package main

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

	"github.com/intervention-engine/fhir/server"
	"github.com/intervention-engine/ie/controllers"
	"github.com/intervention-engine/ie/huddles"
	"github.com/intervention-engine/ie/utilities"
	"github.com/robfig/cron"
)

// Setup the huddle flag to accumulate huddle config file paths
type huddleSlice []string

func (h *huddleSlice) String() string {
	return fmt.Sprint(*h)
}
func (h *huddleSlice) Set(value string) error {
	*h = append(*h, value)
	return nil
}

var huddleFlag huddleSlice

func init() {
	flag.Var(&huddleFlag, "huddle", "path to a huddle configuration file (repeat flag for multiple paths)")
}

func main() {
	// Check for a linked MongoDB container if we are running in Docker
	mongoHost := os.Getenv("MONGO_PORT_27017_TCP_ADDR")
	if mongoHost == "" {
		mongoHost = "localhost"
	}

	codeLookupFlag := flag.Bool("loadCodes", false, "flag to enable download of icd-9 and icd-10 code lookup")
	icd9URL := flag.String("icd9URL", "https://www.cms.gov/Medicare/Coding/ICD9ProviderDiagnosticCodes/Downloads/ICD-9-CM-v32-master-descriptions.zip", "url for icd-9 code definition zip")
	icd10URL := flag.String("icd10URL", "https://www.cms.gov/Medicare/Coding/ICD10/Downloads/2016-Code-Descriptions-in-Tabular-Order.zip", "url for icd-10 code definition zip")
	flag.Parse()

	if *codeLookupFlag {
		utilities.LoadICDFromCMS(mongoHost, *icd9URL, "ICD-9")
		utilities.LoadICDFromCMS(mongoHost, *icd10URL, "ICD-10")
	}

	riskServiceEndpoint := os.Getenv("RISKSERVICE_PORT_9000_TCP_ADDR")
	if riskServiceEndpoint == "" {
		riskServiceEndpoint = "http://localhost:9000"
	} else {
		riskServiceEndpoint = "http://" + riskServiceEndpoint + ":9000"
	}

	subscriptionFlag := flag.Bool("subscriptions", false, "enables limited support for resource subscriptions (default: false)")

	var ip net.IP
	var selfURL string
	host, err := os.Hostname()
	if err != nil {
		panic(err)
	}
	addrs, err := net.LookupIP(host)
	if err != nil {
		fmt.Println("Unable to lookup IP based on hostname, defaulting to localhost.")
		selfURL = "http://localhost:3001"
	}
	for _, addr := range addrs {
		if ipv4 := addr.To4(); ipv4 != nil {
			ip = ipv4
			selfURL = "http://" + ip.String() + ":3001"
		}
	}

	s := server.NewServer(mongoHost)

	reqLog := flag.Bool("reqlog", false, "Enables request logging -- do NOT use in production")
	if *reqLog {
		s.Engine.Use(server.RequestLoggerHandler)
	}

	// Since the huddle controller needs info from the command line, set it up here.  When IE is refactored
	// to take out globals (and other stuff), this should be rethought.
	c := cron.New()
	huddleController := new(huddles.HuddleSchedulerController)
	for _, huddlePath := range huddleFlag {
		data, err := ioutil.ReadFile(huddlePath)
		if err != nil {
			panic(err)
		}
		config := new(huddles.HuddleConfig)
		if err := json.Unmarshal(data, config); err != nil {
			panic(err)
		}
		huddleController.AddConfig(config)

		// Wait 1 minute before doing initial runs or setting up cron jobs.  This allows the server to get
		// started (since it needs to initiate the db connection, etc).
		time.AfterFunc(1*time.Minute, func() {
			// Do an initial run
			fmt.Println("Initial scheduling for huddle with name ", config.Name)
			huddles.ScheduleHuddles(config)

			// Set up the cron job for future runs
			if config.SchedulerCronSpec != "" {
				err := c.AddFunc(config.SchedulerCronSpec, func() {
					_, err := huddles.ScheduleHuddles(config)
					if err != nil {
						fmt.Printf("ERROR: Could not schedule huddles for huddle with name %s: %v", config.Name, err)
					}
				})
				if err != nil {
					panic(err)
				}
				fmt.Printf("Huddle with name %s scheduled with cron spec: %s\n", config.Name, config.SchedulerCronSpec)
			} else {
				fmt.Printf("Warning: Huddle with name %s is not configured with a scheduler cron job.\n", config.Name)
			}
		})
	}
	s.Engine.GET("/ScheduleHuddles", huddleController.ScheduleHandler)

	if len(huddleFlag) > 0 {
		c.Start()
		defer c.Stop()
	}

	closer := controllers.RegisterRoutes(s, selfURL, riskServiceEndpoint, *subscriptionFlag)
	defer closer()

	s.Run(server.Config{
		UseSmartAuth:         false,
		UseLoggingMiddleware: true,
	})
}
<|endoftext|>