<commit_before>package main

import (
	"encoding/gob"
	"errors"
	"flag"
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"runtime"
	"syscall"
	"unsafe"
)

var (
	modshell32         = syscall.NewLazyDLL("shell32.dll")
	procShellExecuteEx = modshell32.NewProc("ShellExecuteExW")
)

const (
	SEE_MASK_NOCLOSEPROCESS = 0x00000040
)

const (
	ERROR_BAD_FORMAT = 11
)

const (
	SE_ERR_FNF             = 2
	SE_ERR_PNF             = 3
	SE_ERR_ACCESSDENIED    = 5
	SE_ERR_OOM             = 8
	SE_ERR_DLLNOTFOUND     = 32
	SE_ERR_SHARE           = 26
	SE_ERR_ASSOCINCOMPLETE = 27
	SE_ERR_DDETIMEOUT      = 28
	SE_ERR_DDEFAIL         = 29
	SE_ERR_DDEBUSY         = 30
	SE_ERR_NOASSOC         = 31
)

type (
	DWORD     uint32
	HANDLE    uintptr
	HINSTANCE HANDLE
	HKEY      HANDLE
	HWND      HANDLE
	ULONG     uint32
	LPCTSTR   uintptr
	LPVOID    uintptr
)

type SHELLEXECUTEINFO struct {
	cbSize         DWORD
	fMask          ULONG
	hwnd           HWND
	lpVerb         LPCTSTR
	lpFile         LPCTSTR
	lpParameters   LPCTSTR
	lpDirectory    LPCTSTR
	nShow          int
	hInstApp       HINSTANCE
	lpIDList       LPVOID
	lpClass        LPCTSTR
	hkeyClass      HKEY
	dwHotKey       DWORD
	hIconOrMonitor HANDLE
	hProcess       HANDLE
}

func ShellExecuteAndWait(hwnd HWND, lpOperation, lpFile, lpParameters, lpDirectory string, nShowCmd int) error {
	var lpctstrVerb, lpctstrParameters, lpctstrDirectory LPCTSTR
	if len(lpOperation) != 0 {
		lpctstrVerb = LPCTSTR(unsafe.Pointer(syscall.StringToUTF16Ptr(lpOperation)))
	}
	if len(lpParameters) != 0 {
		lpctstrParameters = LPCTSTR(unsafe.Pointer(syscall.StringToUTF16Ptr(lpParameters)))
	}
	if len(lpDirectory) != 0 {
		lpctstrDirectory = LPCTSTR(unsafe.Pointer(syscall.StringToUTF16Ptr(lpDirectory)))
	}
	i := &SHELLEXECUTEINFO{
		fMask:        SEE_MASK_NOCLOSEPROCESS,
		hwnd:         hwnd,
		lpVerb:       lpctstrVerb,
		lpFile:       LPCTSTR(unsafe.Pointer(syscall.StringToUTF16Ptr(lpFile))),
		lpParameters: lpctstrParameters,
		lpDirectory:  lpctstrDirectory,
		nShow:        nShowCmd,
	}
	i.cbSize = DWORD(unsafe.Sizeof(*i))
	return ShellExecuteEx(i)
}

func ShellExecuteEx(pExecInfo *SHELLEXECUTEINFO) error {
	ret, _, _ := procShellExecuteEx.Call(uintptr(unsafe.Pointer(pExecInfo)))
	if ret == 1 && pExecInfo.fMask&SEE_MASK_NOCLOSEPROCESS != 0 {
		s, e := syscall.WaitForSingleObject(syscall.Handle(pExecInfo.hProcess), syscall.INFINITE)
		switch s {
		case syscall.WAIT_OBJECT_0:
			break
		case syscall.WAIT_FAILED:
			return os.NewSyscallError("WaitForSingleObject", e)
		default:
			return errors.New("Unexpected result from WaitForSingleObject")
		}
	}
	errorMsg := ""
	if pExecInfo.hInstApp != 0 && pExecInfo.hInstApp <= 32 {
		switch int(pExecInfo.hInstApp) {
		case SE_ERR_FNF:
			errorMsg = "The specified file was not found"
		case SE_ERR_PNF:
			errorMsg = "The specified path was not found"
		case ERROR_BAD_FORMAT:
			errorMsg = "The .exe file is invalid (non-Win32 .exe or error in .exe image)"
		case SE_ERR_ACCESSDENIED:
			errorMsg = "The operating system denied access to the specified file"
		case SE_ERR_ASSOCINCOMPLETE:
			errorMsg = "The file name association is incomplete or invalid"
		case SE_ERR_DDEBUSY:
			errorMsg = "The DDE transaction could not be completed because other DDE transactions were being processed"
		case SE_ERR_DDEFAIL:
			errorMsg = "The DDE transaction failed"
		case SE_ERR_DDETIMEOUT:
			errorMsg = "The DDE transaction could not be completed because the request timed out"
		case SE_ERR_DLLNOTFOUND:
			errorMsg = "The specified DLL was not found"
		case SE_ERR_NOASSOC:
			errorMsg = "There is no application associated with the given file name extension"
		case SE_ERR_OOM:
			errorMsg = "There was not enough memory to complete the operation"
		case SE_ERR_SHARE:
			errorMsg = "A sharing violation occurred"
		default:
			errorMsg = fmt.Sprintf("Unknown error occurred with error code %v", pExecInfo.hInstApp)
		}
	} else {
		return nil
	}
	return errors.New(errorMsg)
}

type msg struct {
	name string
	exit int
	data []byte
}

func msgWrite(enc *gob.Encoder, typ string) io.WriteCloser {
	r, w := io.Pipe()
	go func() {
		defer r.Close()
		var b [4096]byte
		for {
			n, err := r.Read(b[:])
			if err != nil {
				break
			}
			enc.Encode(&msg{name: typ, data: b[:n]})
		}
	}()
	return w
}

func client(addr string) int {
	// connect to server
	conn, err := net.Dial("tcp", addr)
	if err != nil {
		fmt.Fprintf(os.Stderr, "cannot find executable: %v\n", os.Args[0])
		return 1
	}
	defer conn.Close()
	enc, dec := gob.NewEncoder(conn), gob.NewDecoder(conn)

	cmd := exec.Command(flag.Arg(0), flag.Args()[1:]...)

	// stdin
	inw, err := cmd.StdinPipe()
	if err != nil {
		fmt.Fprintf(os.Stderr, "cannot find executable: %v\n", os.Args[0])
		return 1
	}
	defer inw.Close()

	// stdout
	outw := msgWrite(enc, "stdout")
	defer outw.Close()
	cmd.Stdout = outw

	// stderr
	errw := msgWrite(enc, "stderr")
	defer errw.Close()
	cmd.Stderr = errw

	go func() {
		defer inw.Close()
	in_loop:
		for {
			var m msg
			err = dec.Decode(&m)
			if err != nil {
				return
			}
			switch m.name {
			case "close":
				break in_loop
			case "ctrlc":
				if runtime.GOOS == "windows" {
					// windows doesn't support os.Interrupt
					cmd.Process.Kill()
				} else {
					cmd.Process.Signal(os.Interrupt)
				}
			case "stdin":
				inw.Write(m.data)
			}
		}
	}()

	err = cmd.Run()

	code := 1
	if err != nil {
		if status, ok := cmd.ProcessState.Sys().(syscall.WaitStatus); ok {
			code = status.ExitStatus()
		}
	} else {
		code = 0
	}

	enc.Encode(&msg{name: "exit", exit: code})
	return 0
}

func makeCmdLine(args []string) string {
	var s string
	for _, v := range args {
		if s != "" {
			s += " "
		}
		s += syscall.EscapeArg(v)
	}
	return s
}

func server() int {
	// make listner to communicate child process
	lis, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		fmt.Fprintf(os.Stderr, "cannot find executable: %v\n", os.Args[0])
		return 1
	}
	defer lis.Close()

	// make sure executable name to avoid detecting same executable name
	exe, err := os.Executable()
	if err != nil {
		fmt.Fprintf(os.Stderr, "cannot find executable: %v\n", os.Args[0])
		return 1
	}
	args := []string{"-mode", lis.Addr().String()}
	args = append(args, flag.Args()...)

	go func() {
		err = ShellExecuteAndWait(0, "runas", exe, makeCmdLine(args), "", syscall.SW_HIDE)
		if err != nil {
			lis.Close()
		}
	}()

	conn, err := lis.Accept()
	if err != nil {
		fmt.Fprintf(os.Stderr, "cannot execute command: %v\n", makeCmdLine(flag.Args()))
		return 1
	}
	defer conn.Close()

	enc, dec := gob.NewEncoder(conn), gob.NewDecoder(conn)

	sc := make(chan os.Signal, 1)
	signal.Notify(sc, os.Interrupt)
	go func() {
		for range sc {
			enc.Encode(&msg{name: "ctrlc"})
		}
	}()
	defer close(sc)

	go func() {
		var b [256]byte
		for {
			n, err := os.Stdin.Read(b[:])
			if err != nil {
				// stdin was closed
				enc.Encode(&msg{name: "close"})
				break
			}
			enc.Encode(&msg{name: "stdin", data: b[:n]})
		}
	}()

	for {
		var m msg
		err = dec.Decode(&m)
		if err != nil {
			fmt.Fprintf(os.Stderr, "cannot execute command: %v\n", makeCmdLine(flag.Args()))
			return 1
		}
		switch m.name {
		case "stdout":
			os.Stdout.Write(m.data)
		case "stderr":
			os.Stderr.Write(m.data)
		case "exit":
			return m.exit
		}
	}
	return 0
}

func main() {
	var mode string
	flag.StringVar(&mode, "mode", "", "mode")
	flag.Parse()
	if mode != "" {
		os.Exit(client(mode))
	}
	os.Exit(server())
}
<commit_msg>should be public field<commit_after>package main

import (
	"encoding/gob"
	"errors"
	"flag"
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"runtime"
	"syscall"
	"unsafe"
)

var (
	modshell32         = syscall.NewLazyDLL("shell32.dll")
	procShellExecuteEx = modshell32.NewProc("ShellExecuteExW")
)

const (
	SEE_MASK_NOCLOSEPROCESS = 0x00000040
)

const (
	ERROR_BAD_FORMAT = 11
)

const (
	SE_ERR_FNF             = 2
	SE_ERR_PNF             = 3
	SE_ERR_ACCESSDENIED    = 5
	SE_ERR_OOM             = 8
	SE_ERR_DLLNOTFOUND     = 32
	SE_ERR_SHARE           = 26
	SE_ERR_ASSOCINCOMPLETE = 27
	SE_ERR_DDETIMEOUT      = 28
	SE_ERR_DDEFAIL         = 29
	SE_ERR_DDEBUSY         = 30
	SE_ERR_NOASSOC         = 31
)

type (
	DWORD     uint32
	HANDLE    uintptr
	HINSTANCE HANDLE
	HKEY      HANDLE
	HWND      HANDLE
	ULONG     uint32
	LPCTSTR   uintptr
	LPVOID    uintptr
)

type SHELLEXECUTEINFO struct {
	cbSize         DWORD
	fMask          ULONG
	hwnd           HWND
	lpVerb         LPCTSTR
	lpFile         LPCTSTR
	lpParameters   LPCTSTR
	lpDirectory    LPCTSTR
	nShow          int
	hInstApp       HINSTANCE
	lpIDList       LPVOID
	lpClass        LPCTSTR
	hkeyClass      HKEY
	dwHotKey       DWORD
	hIconOrMonitor HANDLE
	hProcess       HANDLE
}

func ShellExecuteAndWait(hwnd HWND, lpOperation, lpFile, lpParameters, lpDirectory string, nShowCmd int) error {
	var lpctstrVerb, lpctstrParameters, lpctstrDirectory LPCTSTR
	if len(lpOperation) != 0 {
		lpctstrVerb = LPCTSTR(unsafe.Pointer(syscall.StringToUTF16Ptr(lpOperation)))
	}
	if len(lpParameters) != 0 {
		lpctstrParameters = LPCTSTR(unsafe.Pointer(syscall.StringToUTF16Ptr(lpParameters)))
	}
	if len(lpDirectory) != 0 {
		lpctstrDirectory = LPCTSTR(unsafe.Pointer(syscall.StringToUTF16Ptr(lpDirectory)))
	}
	i := &SHELLEXECUTEINFO{
		fMask:        SEE_MASK_NOCLOSEPROCESS,
		hwnd:         hwnd,
		lpVerb:       lpctstrVerb,
		lpFile:       LPCTSTR(unsafe.Pointer(syscall.StringToUTF16Ptr(lpFile))),
		lpParameters: lpctstrParameters,
		lpDirectory:  lpctstrDirectory,
		nShow:        nShowCmd,
	}
	i.cbSize = DWORD(unsafe.Sizeof(*i))
	return ShellExecuteEx(i)
}

func ShellExecuteEx(pExecInfo *SHELLEXECUTEINFO) error {
	ret, _, _ := procShellExecuteEx.Call(uintptr(unsafe.Pointer(pExecInfo)))
	if ret == 1 && pExecInfo.fMask&SEE_MASK_NOCLOSEPROCESS != 0 {
		s, e := syscall.WaitForSingleObject(syscall.Handle(pExecInfo.hProcess), syscall.INFINITE)
		switch s {
		case syscall.WAIT_OBJECT_0:
			break
		case syscall.WAIT_FAILED:
			return os.NewSyscallError("WaitForSingleObject", e)
		default:
			return errors.New("Unexpected result from WaitForSingleObject")
		}
	}
	errorMsg := ""
	if pExecInfo.hInstApp != 0 && pExecInfo.hInstApp <= 32 {
		switch int(pExecInfo.hInstApp) {
		case SE_ERR_FNF:
			errorMsg = "The specified file was not found"
		case SE_ERR_PNF:
			errorMsg = "The specified path was not found"
		case ERROR_BAD_FORMAT:
			errorMsg = "The .exe file is invalid (non-Win32 .exe or error in .exe image)"
		case SE_ERR_ACCESSDENIED:
			errorMsg = "The operating system denied access to the specified file"
		case SE_ERR_ASSOCINCOMPLETE:
			errorMsg = "The file name association is incomplete or invalid"
		case SE_ERR_DDEBUSY:
			errorMsg = "The DDE transaction could not be completed because other DDE transactions were being processed"
		case SE_ERR_DDEFAIL:
			errorMsg = "The DDE transaction failed"
		case SE_ERR_DDETIMEOUT:
			errorMsg = "The DDE transaction could not be completed because the request timed out"
		case SE_ERR_DLLNOTFOUND:
			errorMsg = "The specified DLL was not found"
		case SE_ERR_NOASSOC:
			errorMsg = "There is no application associated with the given file name extension"
		case SE_ERR_OOM:
			errorMsg = "There was not enough memory to complete the operation"
		case SE_ERR_SHARE:
			errorMsg = "A sharing violation occurred"
		default:
			errorMsg = fmt.Sprintf("Unknown error occurred with error code %v", pExecInfo.hInstApp)
		}
	} else {
		return nil
	}
	return errors.New(errorMsg)
}

type msg struct {
	Name string
	Exit int
	Data []byte
}

func msgWrite(enc *gob.Encoder, typ string) io.WriteCloser {
	r, w := io.Pipe()
	go func() {
		defer r.Close()
		var b [4096]byte
		for {
			n, err := r.Read(b[:])
			if err != nil {
				break
			}
			enc.Encode(&msg{Name: typ, Data: b[:n]})
		}
	}()
	return w
}

func client(addr string) int {
	// connect to server
	conn, err := net.Dial("tcp", addr)
	if err != nil {
		fmt.Fprintf(os.Stderr, "cannot find executable: %v\n", os.Args[0])
		return 1
	}
	defer conn.Close()
	enc, dec := gob.NewEncoder(conn), gob.NewDecoder(conn)

	cmd := exec.Command(flag.Arg(0), flag.Args()[1:]...)

	// stdin
	inw, err := cmd.StdinPipe()
	if err != nil {
		fmt.Fprintf(os.Stderr, "cannot find executable: %v\n", os.Args[0])
		return 1
	}
	defer inw.Close()

	// stdout
	outw := msgWrite(enc, "stdout")
	defer outw.Close()
	cmd.Stdout = outw

	// stderr
	errw := msgWrite(enc, "stderr")
	defer errw.Close()
	cmd.Stderr = errw

	go func() {
		defer inw.Close()
	in_loop:
		for {
			var m msg
			err = dec.Decode(&m)
			if err != nil {
				return
			}
			switch m.Name {
			case "close":
				break in_loop
			case "ctrlc":
				if runtime.GOOS == "windows" {
					// windows doesn't support os.Interrupt
					cmd.Process.Kill()
				} else {
					cmd.Process.Signal(os.Interrupt)
				}
			case "stdin":
				inw.Write(m.Data)
			}
		}
	}()

	err = cmd.Run()

	code := 1
	if err != nil {
		if status, ok := cmd.ProcessState.Sys().(syscall.WaitStatus); ok {
			code = status.ExitStatus()
		}
	} else {
		code = 0
	}

	enc.Encode(&msg{Name: "exit", Exit: code})
	return 0
}

func makeCmdLine(args []string) string {
	var s string
	for _, v := range args {
		if s != "" {
			s += " "
		}
		s += syscall.EscapeArg(v)
	}
	return s
}

func server() int {
	// make listner to communicate child process
	lis, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		fmt.Fprintf(os.Stderr, "cannot find executable: %v\n", os.Args[0])
		return 1
	}
	defer lis.Close()

	// make sure executable name to avoid detecting same executable name
	exe, err := os.Executable()
	if err != nil {
		fmt.Fprintf(os.Stderr, "cannot find executable: %v\n", os.Args[0])
		return 1
	}
	args := []string{"-mode", lis.Addr().String()}
	args = append(args, flag.Args()...)

	go func() {
		err = ShellExecuteAndWait(0, "runas", exe, makeCmdLine(args), "", syscall.SW_HIDE)
		if err != nil {
			lis.Close()
		}
	}()

	conn, err := lis.Accept()
	if err != nil {
		fmt.Fprintf(os.Stderr, "cannot execute command: %v\n", makeCmdLine(flag.Args()))
		return 1
	}
	defer conn.Close()

	enc, dec := gob.NewEncoder(conn), gob.NewDecoder(conn)

	sc := make(chan os.Signal, 1)
	signal.Notify(sc, os.Interrupt)
	go func() {
		for range sc {
			enc.Encode(&msg{Name: "ctrlc"})
		}
	}()
	defer close(sc)

	go func() {
		var b [256]byte
		for {
			n, err := os.Stdin.Read(b[:])
			if err != nil {
				// stdin was closed
				enc.Encode(&msg{Name: "close"})
				break
			}
			enc.Encode(&msg{Name: "stdin", Data: b[:n]})
		}
	}()

	for {
		var m msg
		err = dec.Decode(&m)
		if err != nil {
			fmt.Fprintf(os.Stderr, "cannot execute command: %v\n", makeCmdLine(flag.Args()))
			return 1
		}
		switch m.Name {
		case "stdout":
			os.Stdout.Write(m.Data)
		case "stderr":
			os.Stderr.Write(m.Data)
		case "exit":
			return m.Exit
		}
	}
	return 0
}

func main() {
	var mode string
	flag.StringVar(&mode, "mode", "", "mode")
	flag.Parse()
	if mode != "" {
		os.Exit(client(mode))
	}
	os.Exit(server())
}
<|endoftext|>


<commit_before>package main

import (
	"os"
)

var commands = []*Command{
	cmdLogin,
	cmdLogout,
	cmdLogins,
	cmdActive,
	cmdWhoami,
	cmdDescribe,
	cmdSobject,
	cmdField,
	cmdRecord,
	cmdBulk,
	cmdFetch,
	cmdImport,
	cmdExport,
	cmdQuery,
	cmdApex,
	cmdOauth,
	cmdTest,
	cmdSecurity,
	cmdVersion,
	cmdUpdate,
	cmdHelp,
	cmdPush,
	cmdAura,
	cmdPassword,
	cmdNotifySet,
	cmdLimits,
}

func main() {
	args := os.Args[1:]
	if len(args) < 1 {
		usage()
	}

	for _, cmd := range commands {
		if cmd.Name() == args[0] && cmd.Run != nil {
			cmd.Flag.Usage = func() {
				cmd.printUsage()
			}
			if err := cmd.Flag.Parse(args[1:]); err != nil {
				os.Exit(2)
			}
			cmd.Run(cmd, cmd.Flag.Args())
			return
		}
	}
	usage()
}
<commit_msg>Help command moved<commit_after>package main

import (
	"os"
)

var commands = []*Command{
	cmdLogin,
	cmdLogout,
	cmdLogins,
	cmdActive,
	cmdWhoami,
	cmdDescribe,
	cmdSobject,
	cmdField,
	cmdRecord,
	cmdBulk,
	cmdFetch,
	cmdImport,
	cmdExport,
	cmdQuery,
	cmdApex,
	cmdOauth,
	cmdTest,
	cmdSecurity,
	cmdVersion,
	cmdUpdate,
	cmdPush,
	cmdAura,
	cmdPassword,
	cmdNotifySet,
	cmdLimits,
	cmdHelp,
}

func main() {
	args := os.Args[1:]
	if len(args) < 1 {
		usage()
	}

	for _, cmd := range commands {
		if cmd.Name() == args[0] && cmd.Run != nil {
			cmd.Flag.Usage = func() {
				cmd.printUsage()
			}
			if err := cmd.Flag.Parse(args[1:]); err != nil {
				os.Exit(2)
			}
			cmd.Run(cmd, cmd.Flag.Args())
			return
		}
	}
	usage()
}
<|endoftext|>


<commit_before>package main

import (
	"io"
	"log"
	"net/http"
	"os"
	"runtime"
	"sync"
	"time"

	"github.com/couchbaselabs/gateload/api"
	"github.com/couchbaselabs/gateload/workload"
)

const (
	AUTH_TYPE_SESSION = "session"
	AUTH_TYPE_BASIC   = "basic"
)

func main() {

	if os.Getenv("GOMAXPROCS") == "" {
		log.Printf("Setting GOMAXPROCS to %v", runtime.NumCPU())
		runtime.GOMAXPROCS(runtime.NumCPU())
	} else {
		log.Printf("GOMAXPROCS is set at %v", os.Getenv("GOMAXPROCS"))
	}

	// start up an http server, just to serve up expvars
	go http.ListenAndServe(":9876", nil)

	var config workload.Config
	workload.ReadConfig(&config)

	admin := api.SyncGatewayClient{}
	admin.Init(
		config.Hostname,
		config.Database,
		config.Port,
		config.AdminPort,
		config.LogRequests,
	)
	if !admin.Valid() {
		log.Fatalf("unable to connect to sync_gateway, check the hostname and database")
	}

	pendingUsers := make(chan *workload.User)
	users := make([]*workload.User, config.NumPullers+config.NumPushers)

	// start a routine to place pending users into array
	go func() {
		for pendingUser := range pendingUsers {

			// users = append(users, pendingUser)
			users[pendingUser.SeqId-config.UserOffset] = pendingUser
		}
	}()

	rampUpDelay := config.RampUpIntervalMs / (config.NumPullers + config.NumPushers)

	// use a fixed number of workers to create the users/sessions
	userIterator := workload.UserIterator(
		config.NumPullers,
		config.NumPushers,
		config.UserOffset,
		config.ChannelActiveUsers,
		config.ChannelConcurrentUsers,
		config.MinUserOffTimeMs,
		config.MaxUserOffTimeMs,
		rampUpDelay,
		config.RunTimeMs,
	)
	adminWg := sync.WaitGroup{}
	worker := func() {
		defer adminWg.Done()
		for user := range userIterator {
			createSession(&admin, user, config)
			pendingUsers <- user
		}
	}

	for i := 0; i < 200; i++ {
		adminWg.Add(1)
		go worker()
	}

	// wait for all the workers to finish
	adminWg.Wait()
	// close the pending users channel to free that routine
	close(pendingUsers)

	numChannels := (config.NumPullers + config.NumPushers) / config.ChannelActiveUsers
	channelRampUpDelayMs := time.Duration(config.RampUpIntervalMs/numChannels) * time.Millisecond

	wg := sync.WaitGroup{}
	channel := ""
	for _, user := range users {
		nextChannel := user.Channel
		if channel != nextChannel {
			if channel != "" {
				time.Sleep(channelRampUpDelayMs)
			}
			channel = nextChannel
		}
		wg := sync.WaitGroup{}
		go runUser(user, config, &wg)
		wg.Add(1)
	}

	if config.RunTimeMs > 0 {
		time.Sleep(time.Duration(config.RunTimeMs-config.RampUpIntervalMs) * time.Millisecond)
		log.Println("Shutting down clients")
	} else {
		wg.Wait()
	}

	writeExpvarsToFile()

}

func createSession(admin *api.SyncGatewayClient, user *workload.User, config workload.Config) {

	userMeta := api.UserAuth{
		Name:          user.Name,
		Password:      config.Password,
		AdminChannels: []string{user.Channel},
	}
	admin.AddUser(user.Name, userMeta, user.Type)

	if config.AuthType == AUTH_TYPE_SESSION {

		session := api.Session{Name: user.Name, TTL: 2592000} // 1 month
		log.Printf("====== Creating new session for %s (%s)", user.Type, user.Name)
		user.Cookie = admin.CreateSession(user.Name, session)
		log.Printf("====== Done Creating new session for %s (%s)", user.Type, user.Name)

	}

}

func runUser(user *workload.User, config workload.Config, wg *sync.WaitGroup) {
	c := api.SyncGatewayClient{}
	c.Init(
		config.Hostname,
		config.Database,
		config.Port,
		config.AdminPort,
		config.LogRequests,
	)
	if config.AuthType == AUTH_TYPE_SESSION {
		c.AddCookie(&user.Cookie)
	} else {
		c.AddUsername(user.Name)
		c.AddPassword(config.Password)
	}

	log.Printf("Starting new %s (%s)", user.Type, user.Name)
	if user.Type == "pusher" {
		go workload.RunNewPusher(
			user.Schedule,
			user.Name,
			&c,
			user.Channel,
			config.DocSize,
			config.SendAttachment,
			config.DocSizeDistribution,
			user.SeqId,
			config.SleepTimeMs,
			wg,
			config.AddDocToTargetUser,
		)
	} else {
		go workload.RunNewPuller(
			user.Schedule,
			&c,
			user.Channel,
			user.Name,
			config.FeedType,
			wg,
		)
	}
	log.Printf("------ Done Starting new %s (%s)", user.Type, user.Name)

}

// At the end of the run, write the full list of expvars to a file
func writeExpvarsToFile() {

	// read http
	url := "http://localhost:9876/debug/vars"
	resp, err := http.Get(url)
	if err != nil {
		log.Fatalf("Error writing expvars, failed connection to: %v", url)
	}
	defer resp.Body.Close()

	// write to file
	destFileName := "gateload_expvars.json"
	destFile, err := os.Create(destFileName)
	if err != nil {
		log.Fatalf("Error opening file for writing to :%v", destFileName)
	}
	_, err = io.Copy(destFile, resp.Body)
	if err != nil {
		log.Fatalf("Error copying from %v -> %v", url, destFile)
	}

	log.Printf("Wrote results to %v", destFileName)

}
<commit_msg>write expvars to a file more often<commit_after>package main

import (
	"io"
	"log"
	"net/http"
	"os"
	"runtime"
	"sync"
	"time"

	"github.com/couchbaselabs/gateload/api"
	"github.com/couchbaselabs/gateload/workload"
)

const (
	AUTH_TYPE_SESSION = "session"
	AUTH_TYPE_BASIC   = "basic"
)

func main() {

	if os.Getenv("GOMAXPROCS") == "" {
		log.Printf("Setting GOMAXPROCS to %v", runtime.NumCPU())
		runtime.GOMAXPROCS(runtime.NumCPU())
	} else {
		log.Printf("GOMAXPROCS is set at %v", os.Getenv("GOMAXPROCS"))
	}

	// start up a goroutine that will write the expvars to a file
	go func() {
		<-time.After(time.Second * 60)
		writeExpvarsToFile()
	}()

	// start up an http server, just to serve up expvars
	go http.ListenAndServe(":9876", nil)

	var config workload.Config
	workload.ReadConfig(&config)

	admin := api.SyncGatewayClient{}
	admin.Init(
		config.Hostname,
		config.Database,
		config.Port,
		config.AdminPort,
		config.LogRequests,
	)
	if !admin.Valid() {
		log.Fatalf("unable to connect to sync_gateway, check the hostname and database")
	}

	pendingUsers := make(chan *workload.User)
	users := make([]*workload.User, config.NumPullers+config.NumPushers)

	// start a routine to place pending users into array
	go func() {
		for pendingUser := range pendingUsers {

			// users = append(users, pendingUser)
			users[pendingUser.SeqId-config.UserOffset] = pendingUser
		}
	}()

	rampUpDelay := config.RampUpIntervalMs / (config.NumPullers + config.NumPushers)

	// use a fixed number of workers to create the users/sessions
	userIterator := workload.UserIterator(
		config.NumPullers,
		config.NumPushers,
		config.UserOffset,
		config.ChannelActiveUsers,
		config.ChannelConcurrentUsers,
		config.MinUserOffTimeMs,
		config.MaxUserOffTimeMs,
		rampUpDelay,
		config.RunTimeMs,
	)
	adminWg := sync.WaitGroup{}
	worker := func() {
		defer adminWg.Done()
		for user := range userIterator {
			createSession(&admin, user, config)
			pendingUsers <- user
		}
	}

	for i := 0; i < 200; i++ {
		adminWg.Add(1)
		go worker()
	}

	// wait for all the workers to finish
	adminWg.Wait()
	// close the pending users channel to free that routine
	close(pendingUsers)

	numChannels := (config.NumPullers + config.NumPushers) / config.ChannelActiveUsers
	channelRampUpDelayMs := time.Duration(config.RampUpIntervalMs/numChannels) * time.Millisecond

	wg := sync.WaitGroup{}
	channel := ""
	for _, user := range users {
		nextChannel := user.Channel
		if channel != nextChannel {
			if channel != "" {
				time.Sleep(channelRampUpDelayMs)
			}
			channel = nextChannel
		}
		wg := sync.WaitGroup{}
		go runUser(user, config, &wg)
		wg.Add(1)
	}

	if config.RunTimeMs > 0 {
		time.Sleep(time.Duration(config.RunTimeMs-config.RampUpIntervalMs) * time.Millisecond)
		log.Println("Shutting down clients")
	} else {
		wg.Wait()
	}

	// capture final snapshot of expvars
	writeExpvarsToFile()

}

func createSession(admin *api.SyncGatewayClient, user *workload.User, config workload.Config) {

	userMeta := api.UserAuth{
		Name:          user.Name,
		Password:      config.Password,
		AdminChannels: []string{user.Channel},
	}
	admin.AddUser(user.Name, userMeta, user.Type)

	if config.AuthType == AUTH_TYPE_SESSION {

		session := api.Session{Name: user.Name, TTL: 2592000} // 1 month
		log.Printf("====== Creating new session for %s (%s)", user.Type, user.Name)
		user.Cookie = admin.CreateSession(user.Name, session)
		log.Printf("====== Done Creating new session for %s (%s)", user.Type, user.Name)

	}

}

func runUser(user *workload.User, config workload.Config, wg *sync.WaitGroup) {
	c := api.SyncGatewayClient{}
	c.Init(
		config.Hostname,
		config.Database,
		config.Port,
		config.AdminPort,
		config.LogRequests,
	)
	if config.AuthType == AUTH_TYPE_SESSION {
		c.AddCookie(&user.Cookie)
	} else {
		c.AddUsername(user.Name)
		c.AddPassword(config.Password)
	}

	log.Printf("Starting new %s (%s)", user.Type, user.Name)
	if user.Type == "pusher" {
		go workload.RunNewPusher(
			user.Schedule,
			user.Name,
			&c,
			user.Channel,
			config.DocSize,
			config.SendAttachment,
			config.DocSizeDistribution,
			user.SeqId,
			config.SleepTimeMs,
			wg,
			config.AddDocToTargetUser,
		)
	} else {
		go workload.RunNewPuller(
			user.Schedule,
			&c,
			user.Channel,
			user.Name,
			config.FeedType,
			wg,
		)
	}
	log.Printf("------ Done Starting new %s (%s)", user.Type, user.Name)

}

// At the end of the run, write the full list of expvars to a file
func writeExpvarsToFile() {

	// read http
	url := "http://localhost:9876/debug/vars"
	resp, err := http.Get(url)
	if err != nil {
		log.Fatalf("Error writing expvars, failed connection to: %v", url)
	}
	defer resp.Body.Close()

	// write to file
	destFileName := "gateload_expvars.json"
	destFile, err := os.Create(destFileName)
	if err != nil {
		log.Fatalf("Error opening file for writing to :%v", destFileName)
	}
	_, err = io.Copy(destFile, resp.Body)
	if err != nil {
		log.Fatalf("Error copying from %v -> %v", url, destFile)
	}

	log.Printf("Wrote results to %v", destFileName)

}
<|endoftext|>


<commit_before>package main

import (
	"os"
	"runtime"

	"github.com/martin61/i2p-tools/cmd"
	"github.com/codegangsta/cli"
)

func main() {
	// use at most half the cpu cores
	runtime.GOMAXPROCS(runtime.NumCPU() / 2)

	app := cli.NewApp()
	app.Name = "i2p-tools"
	app.Version = "0.1.0"
	app.Usage = "I2P tools and reseed server"
	app.Author = "Matt Drollette"
	app.Email = "matt@drollette.com"
	app.Flags = []cli.Flag{}
	app.Commands = []cli.Command{
		cmd.NewReseedCommand(),
		cmd.NewSu3VerifyCommand(),
		cmd.NewKeygenCommand(),
		// cmd.NewSu3VerifyPublicCommand(),
	}

	if err := app.Run(os.Args); err != nil {
		os.Exit(1)
	}
}
<commit_msg>Update main.go<commit_after>package main

import (
	"os"
	"runtime"

	"github.com/martin61/i2p-tools/cmd"
	"github.com/codegangsta/cli"
)

func main() {
	// use at most half the cpu cores
	runtime.GOMAXPROCS(runtime.NumCPU() / 2)

	app := cli.NewApp()
	app.Name = "i2p-tools"
	app.Version = "0.1.1"
	app.Usage = "I2P tools and reseed server"
	app.Author = "martin61"
	app.Email = "-"
	app.Flags = []cli.Flag{}
	app.Commands = []cli.Command{
		cmd.NewReseedCommand(),
		cmd.NewSu3VerifyCommand(),
		cmd.NewKeygenCommand(),
		// cmd.NewSu3VerifyPublicCommand(),
	}

	if err := app.Run(os.Args); err != nil {
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package couchdb

import (
	"bytes"
	"fmt"
	"net/http"
)

// The database itself.
type Database struct {
	Url string
}

// Constructor for a new Database.
// Accepts optional failure and success messages (in that order).
// Will panic if the database is not running or not reachable.
func Open(url string, msgs ...string) Database {
	db := Database{url}
	// Panic if the database is not running.
	if !db.exists("") {
		panic(fmt.Sprintf("DB is not running at %q", db.Url))
	}
	// Otherwise create the table does not exist.
	if !db.exists(url, msgs...) {
		db.Put("", nil)
	}
	return db
}

// Perform a GET request to the database.
func (d *Database) Get(path string, data []byte) (*http.Response, error) {
	return d.query("GET", path, data)
}

// Perform a PUT request to the database.
func (d *Database) Put(path string, data []byte) (*http.Response, error) {
	return d.query("PUT", path, data)
}

// Perform a DELETE request to the database.
func (d *Database) Delete(path string, data []byte) (*http.Response, error) {
	return d.query("DELETE", path, data)
}

// Simplifies making a request to the database into a single function call.
func (d *Database) query(requestType string, path string, data []byte) (*http.Response, error) {
	url := fmt.Sprintf("%s/%s", d.Url, path)
	req, err := http.NewRequest(requestType, url, bytes.NewBuffer(data))
	if err != nil {
		return nil, err
	}
	client := &http.Client{}
	return client.Do(req)
}

// Checks a given URL to see if the response returns without error.
// Accepts optional failure and success messages (in that order).
func (d *Database) exists(url string, msgs ...string) bool {
	if resp, err := d.Get(url, nil); err != nil || resp.StatusCode != 200 {
		if len(msgs) > 0 {
			fmt.Println(msgs[0])
		}
		if err != nil {
			fmt.Println(err)
		}
		return false
	}

	if len(msgs) > 1 {
		fmt.Println(msgs[1])
	}

	return true
}
<commit_msg>Proper detection and reporting when database is not running.<commit_after>package couchdb

import (
	"bytes"
	"fmt"
	"net/http"
	"os"
)

// The database itself.
type Database struct {
	Url string
}

// Constructor for a new Database.
// Accepts optional failure and success messages (in that order).
func Open(url string, msgs ...string) Database {
	db := Database{url}
	// Otherwise create the table does not exist.
	if !db.exists("", msgs...) {
		db.Put("", nil)
	}
	return db
}

// Perform a GET request to the database.
func (d *Database) Get(path string, data []byte) (*http.Response, error) {
	return d.query("GET", path, data)
}

// Perform a PUT request to the database.
func (d *Database) Put(path string, data []byte) (*http.Response, error) {
	return d.query("PUT", path, data)
}

// Perform a DELETE request to the database.
func (d *Database) Delete(path string, data []byte) (*http.Response, error) {
	return d.query("DELETE", path, data)
}

// Simplifies making a request to the database into a single function call.
func (d *Database) query(requestType string, path string, data []byte) (*http.Response, error) {
	url := fmt.Sprintf("%s/%s", d.Url, path)
	req, err := http.NewRequest(requestType, url, bytes.NewBuffer(data))
	if err != nil {
		return nil, err
	}
	client := &http.Client{}
	return client.Do(req)
}

// Checks a given URL to see if the response returns without error.
// Accepts optional failure and success messages (in that order).
func (d *Database) exists(url string, msgs ...string) bool {
	if resp, err := d.Get(url, nil); err != nil || resp.StatusCode != 200 {
		if len(msgs) > 0 {
			fmt.Println(msgs[0])
		}
		if err != nil {
			fmt.Println(fmt.Sprintf(
				"Database is not running/unreachable at %q", d.Url))
			os.Exit(0)
		}
		return false
	}

	if len(msgs) > 1 {
		fmt.Println(msgs[1])
	}

	return true
}
<|endoftext|>


<commit_before>package main

import (
	"os"

	log "github.com/Sirupsen/logrus"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/defaults"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/s3/s3manager"
	"github.com/codegangsta/cli"
	"github.com/mitchellh/go-homedir"

	"bufio"
	"errors"
	"fmt"
	"github.com/asimihsan/arqinator/arq"
	"github.com/asimihsan/arqinator/connector"
	"runtime"
	"io"
)

func cliSetup(c *cli.Context) error {
	switch c.GlobalString("backup-type") {
	case "s3":
	case "googlecloudstorage":
	default:
		return errors.New("Currently only support backup-type of: ['s3', 'googlecloudstorage']")
	}
	if c.GlobalBool("verbose") {
		log.SetLevel(log.DebugLevel)
	}
	return nil
}

func awsSetup(c *cli.Context) (connector.Connection, error) {
	region := c.GlobalString("s3-region")
	s3BucketName := c.GlobalString("s3-bucket-name")
	cacheDirectory := c.GlobalString("cache-directory")

	defaults.DefaultConfig.Region = aws.String(region)
	svc := s3.New(nil)
	opts := &s3manager.DownloadOptions{
		S3:          svc,
		Concurrency: runtime.GOMAXPROCS(0),
	}
	s3Connection := connector.NewS3Connection(svc, cacheDirectory, s3BucketName, opts)
	return s3Connection, nil
}

func googleCloudStorageSetup(c *cli.Context) (connector.Connection, error) {
	jsonPrivateKeyFilepath := c.GlobalString("gcs-json-private-key-filepath")
	projectID := c.GlobalString("gcs-project-id")
	bucketName := c.GlobalString("gcs-bucket-name")
	cacheDirectory := c.GlobalString("cache-directory")

	connection, err := connector.NewGoogleCloudStorageConnection(jsonPrivateKeyFilepath, projectID, bucketName, cacheDirectory)
	if err != nil {
		return connection, err
	}
	return connection, nil
}

func getConnection(c *cli.Context) (connector.Connection, error) {
	var (
		connection connector.Connection
		err        error
	)
	switch c.GlobalString("backup-type") {
	case "s3":
		connection, err = awsSetup(c)
	case "googlecloudstorage":
		connection, err = googleCloudStorageSetup(c)
	}
	if err != nil {
		log.Debugf("%s", err)
		return nil, err
	}
	return connection, nil
}

func getArqBackupSets(c *cli.Context, connection connector.Connection) ([]*arq.ArqBackupSet, error) {
	password := []byte(os.Getenv("ARQ_ENCRYPTION_PASSWORD"))

	arqBackupSets, err := arq.GetArqBackupSets(connection, password)
	if err != nil {
		log.Debugf("Error during getArqBackupSets: %s", err)
		return nil, err
	}
	return arqBackupSets, nil
}

func listBackupSets(c *cli.Context, connection connector.Connection) error {
	arqBackupSets, err := getArqBackupSets(c, connection)
	if err != nil {
		log.Debugf("Error during listBackupSets: %s", err)
		return nil
	}
	for _, arqBackupSet := range arqBackupSets {
		fmt.Printf("ArqBackupSet\n")
		fmt.Printf("    UUID %s\n", arqBackupSet.UUID)
		fmt.Printf("    ComputerName %s\n", arqBackupSet.ComputerInfo.ComputerName)
		fmt.Printf("    UserName %s\n", arqBackupSet.ComputerInfo.UserName)
		fmt.Printf("    Folders\n")
		for _, bucket := range arqBackupSet.Buckets {
			fmt.Printf("        LocalPath %s\n", bucket.LocalPath)
			fmt.Printf("        UUID %s\n", bucket.UUID)
		}
	}
	return nil
}

func findBucket(c *cli.Context, connection connector.Connection, backupSetUUID string, folderUUID string) (*arq.ArqBucket, error) {
	arqBackupSets, err := getArqBackupSets(c, connection)
	if err != nil {
		log.Debugf("Error during findBucket: %s", err)
		return nil, err
	}
	var bucket *arq.ArqBucket
	for _, arqBackupSet := range arqBackupSets {
		if arqBackupSet.UUID == backupSetUUID {
			for _, folder := range arqBackupSet.Buckets {
				if folder.UUID == folderUUID {
					bucket = folder
				}
			}
		}
	}
	if bucket == nil {
		err := errors.New(fmt.Sprintf("Couldn't find backup set UUID %s, folder UUID %s.", backupSetUUID, folderUUID))
		log.Errorf("%s", err)
		return nil, err
	}
	return bucket, nil
}

func listDirectoryContents(c *cli.Context, connection connector.Connection) error {
	cacheDirectory := c.GlobalString("cache-directory")
	backupSetUUID := c.String("backup-set-uuid")
	folderUUID := c.String("folder-uuid")
	targetPath := c.String("path")

	bucket, err := findBucket(c, connection, backupSetUUID, folderUUID)
	if err != nil {
		err := errors.New(fmt.Sprintf("Couldn't find backup set UUID %s, folder UUID %s.", backupSetUUID, folderUUID))
		log.Errorf("%s", err)
		return err
	}
	backupSet := bucket.ArqBackupSet
	backupSet.CacheTreePackSets()

	tree, node, err := arq.FindNode(cacheDirectory, backupSet, bucket, targetPath)
	if err != nil {
		log.Errorf("Failed to find target path %s: %s", targetPath, err)
		return err
	}
	if node == nil || node.IsTree.IsTrue() {
		if tree == nil {
			err2 := errors.New(fmt.Sprintf("node is tree but no tree found: %s", node))
			log.Errorf("%s", err2)
			return err2
		}
		apsi, _ := arq.NewPackSetIndex(cacheDirectory, backupSet, bucket)
		for _, node := range tree.Nodes {
			if node.IsTree.IsTrue() {
				tree, err := apsi.GetPackFileAsTree(backupSet, bucket, *node.DataBlobKeys[0].SHA1)
				if err != nil {
					log.Debugf("Failed to find tree for node %s: %s", node, err)
					node.PrintOutput()
				} else if tree == nil {
					log.Debugf("directory node %s has no tree", node)
					node.PrintOutput()
				} else {
					tree.PrintOutput(node)
				}
			} else {
				node.PrintOutput()
			}
		}
	} else {
		node.PrintOutput()
	}
	return nil
}

func recover(c *cli.Context, connection connector.Connection) error {
	cacheDirectory := c.GlobalString("cache-directory")
	backupSetUUID := c.String("backup-set-uuid")
	folderUUID := c.String("folder-uuid")
	sourcePath := c.String("source-path")
	destinationPath := c.String("destination-path")

	if _, err := os.Stat(destinationPath); err == nil {
		err := errors.New(fmt.Sprintf("Destination path %s already exists, won't overwrite.", destinationPath))
		log.Errorf("%s", err)
		return err
	}
	bucket, err := findBucket(c, connection, backupSetUUID, folderUUID)
	if err != nil {
		err := errors.New(fmt.Sprintf("Couldn't find backup set UUID %s, folder UUID %s.", backupSetUUID, folderUUID))
		log.Errorf("%s", err)
		return err
	}
	backupSet := bucket.ArqBackupSet
	backupSet.CacheTreePackSets()
	backupSet.CacheBlobPackSets()

	tree, node, err := arq.FindNode(cacheDirectory, backupSet, bucket, sourcePath)
	if err != nil {
		log.Errorf("Failed to find source path %s: %s", sourcePath, err)
		return err
	}
	if node == nil || node.IsTree.IsTrue() {
		log.Errorf("unsupported right now. tree: %s", tree)
		return nil
	}
	apsi, _ := arq.NewPackSetIndex(cacheDirectory, backupSet, bucket)
	f, err := os.Create(destinationPath)
	if err != nil {
		log.Errorf("Failed to open destinationPath %s: %s", destinationPath, err)
		return err
	}
	defer f.Close()
	err = f.Truncate(node.UncompressedDataSize)
	if err != nil {
		log.Errorf("Failed to pre-allocate size of file %s: %s", destinationPath, err)
	}
	w := bufio.NewWriter(f)
	defer w.Flush()
	r, err := arq.GetReaderForBlobKeys(node.DataBlobKeys, apsi, backupSet, bucket)
	if err != nil {
		log.Errorf("Failed during GetReaderForBlobKeys for node %s: %s", node, err)
		return err
	}
	io.Copy(w, r)
	return nil
}

func main() {
	defaultCacheDirectory, err := homedir.Expand("~/.arqinator_cache")
	if err != nil {
		log.Fatal("Failed to get user's home dir: ", err)
	}

	app := cli.NewApp()
	app.Name = "arqinator"
	app.Usage = "restore folders and files from Arq backups"
	app.Version = "0.1"
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "backup-type",
			Usage: "Method used for backup, one of: ['s3', 'googlecloudstorage']",
		},
		cli.StringFlag{
			Name:  "s3-region",
			Usage: "AWS S3 region, e.g. 'us-west-2'.",
		},
		cli.StringFlag{
			Name:  "s3-bucket-name",
			Usage: "AWS S3 bucket name, e.g. 'arq-akiaabdefg-us-west-2'.",
		},
		cli.StringFlag{
			Name:  "gcs-json-private-key-filepath",
			Usage: "Google Cloud Storage JSON private key filepath. See: https://goo.gl/SK5Rb7",
		},
		cli.StringFlag{
			Name:  "gcs-project-id",
			Usage: "Google Cloud Storage project ID.",
		},
		cli.StringFlag{
			Name:  "gcs-bucket-name",
			Usage: "Google Cloud Storage bucket name.",
		},
		cli.StringFlag{
			Name:  "cache-directory",
			Value: defaultCacheDirectory,
			Usage: fmt.Sprintf("Where to cache Arq files for browsing. Default: %s", defaultCacheDirectory),
		},
		cli.BoolFlag{
			Name:  "delete-cache-directory",
			Usage: "Delete cache directory before starting. Useful if seeing errors that could be due to truncated downloads.",
		},
		cli.BoolFlag{
			Name:  "verbose",
			Usage: "Enable verbose logging",
		},
	}
	app.Commands = []cli.Command{
		{
			Name:  "list-backup-sets",
			Usage: "List backup sets in this account.",
			Action: func(c *cli.Context) {
				if err := cliSetup(c); err != nil {
					log.Errorf("%s", err)
					return
				}
				connection, err := getConnection(c)
				if err != nil {
					log.Errorf("%s", err)
					return
				}
				if err := listBackupSets(c, connection); err != nil {
					log.Errorf("%s", err)
					return
				}
			},
		},
		{
			Name:  "list-directory-contents",
			Usage: "List contents of directory in backup.",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "backup-set-uuid",
					Usage: "UUID of backup set. Use 'list-backup-sets' to determine this.",
				},
				cli.StringFlag{
					Name:  "folder-uuid",
					Usage: "UUID of folder. Use 'list-backup-sets' to determine this.",
				},
				cli.StringFlag{
					Name:  "path",
					Usage: "Path of directory or file in backup",
				},
			},
			Action: func(c *cli.Context) {
				if err := cliSetup(c); err != nil {
					log.Errorf("%s", err)
					return
				}
				connection, err := getConnection(c)
				if err != nil {
					log.Errorf("%s", err)
					return
				}
				if err := listDirectoryContents(c, connection); err != nil {
					log.Errorf("%s", err)
					return
				}
			},
		},
		{
			Name:  "recover",
			Usage: "Recover a file or directory from a backup",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "backup-set-uuid",
					Usage: "UUID of backup set. Use 'list-backup-sets' to determine this.",
				},
				cli.StringFlag{
					Name:  "folder-uuid",
					Usage: "UUID of folder. Use 'list-backup-sets' to determine this.",
				},
				cli.StringFlag{
					Name:  "source-path",
					Usage: "Path of directory or file in backup",
				},
				cli.StringFlag{
					Name:  "destination-path",
					Usage: "Path to recover directory or file into. Must not already exist.",
				},
			},
			Action: func(c *cli.Context) {
				if err := cliSetup(c); err != nil {
					log.Errorf("%s", err)
					return
				}
				connection, err := getConnection(c)
				if err != nil {
					log.Errorf("%s", err)
					return
				}
				if err := recover(c, connection); err != nil {
					log.Errorf("%s", err)
					return
				}
			},
		},
	}
	app.Run(os.Args)
}
<commit_msg>repopulate file permissions on recovery<commit_after>package main

import (
	"os"

	log "github.com/Sirupsen/logrus"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/defaults"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/s3/s3manager"
	"github.com/codegangsta/cli"
	"github.com/mitchellh/go-homedir"

	"bufio"
	"errors"
	"fmt"
	"github.com/asimihsan/arqinator/arq"
	"github.com/asimihsan/arqinator/connector"
	"runtime"
	"io"
)

func cliSetup(c *cli.Context) error {
	switch c.GlobalString("backup-type") {
	case "s3":
	case "googlecloudstorage":
	default:
		return errors.New("Currently only support backup-type of: ['s3', 'googlecloudstorage']")
	}
	if c.GlobalBool("verbose") {
		log.SetLevel(log.DebugLevel)
	}
	return nil
}

func awsSetup(c *cli.Context) (connector.Connection, error) {
	region := c.GlobalString("s3-region")
	s3BucketName := c.GlobalString("s3-bucket-name")
	cacheDirectory := c.GlobalString("cache-directory")

	defaults.DefaultConfig.Region = aws.String(region)
	svc := s3.New(nil)
	opts := &s3manager.DownloadOptions{
		S3:          svc,
		Concurrency: runtime.GOMAXPROCS(0),
	}
	s3Connection := connector.NewS3Connection(svc, cacheDirectory, s3BucketName, opts)
	return s3Connection, nil
}

func googleCloudStorageSetup(c *cli.Context) (connector.Connection, error) {
	jsonPrivateKeyFilepath := c.GlobalString("gcs-json-private-key-filepath")
	projectID := c.GlobalString("gcs-project-id")
	bucketName := c.GlobalString("gcs-bucket-name")
	cacheDirectory := c.GlobalString("cache-directory")

	connection, err := connector.NewGoogleCloudStorageConnection(jsonPrivateKeyFilepath, projectID, bucketName, cacheDirectory)
	if err != nil {
		return connection, err
	}
	return connection, nil
}

func getConnection(c *cli.Context) (connector.Connection, error) {
	var (
		connection connector.Connection
		err        error
	)
	switch c.GlobalString("backup-type") {
	case "s3":
		connection, err = awsSetup(c)
	case "googlecloudstorage":
		connection, err = googleCloudStorageSetup(c)
	}
	if err != nil {
		log.Debugf("%s", err)
		return nil, err
	}
	return connection, nil
}

func getArqBackupSets(c *cli.Context, connection connector.Connection) ([]*arq.ArqBackupSet, error) {
	password := []byte(os.Getenv("ARQ_ENCRYPTION_PASSWORD"))

	arqBackupSets, err := arq.GetArqBackupSets(connection, password)
	if err != nil {
		log.Debugf("Error during getArqBackupSets: %s", err)
		return nil, err
	}
	return arqBackupSets, nil
}

func listBackupSets(c *cli.Context, connection connector.Connection) error {
	arqBackupSets, err := getArqBackupSets(c, connection)
	if err != nil {
		log.Debugf("Error during listBackupSets: %s", err)
		return nil
	}
	for _, arqBackupSet := range arqBackupSets {
		fmt.Printf("ArqBackupSet\n")
		fmt.Printf("    UUID %s\n", arqBackupSet.UUID)
		fmt.Printf("    ComputerName %s\n", arqBackupSet.ComputerInfo.ComputerName)
		fmt.Printf("    UserName %s\n", arqBackupSet.ComputerInfo.UserName)
		fmt.Printf("    Folders\n")
		for _, bucket := range arqBackupSet.Buckets {
			fmt.Printf("        LocalPath %s\n", bucket.LocalPath)
			fmt.Printf("        UUID %s\n", bucket.UUID)
		}
	}
	return nil
}

func findBucket(c *cli.Context, connection connector.Connection, backupSetUUID string, folderUUID string) (*arq.ArqBucket, error) {
	arqBackupSets, err := getArqBackupSets(c, connection)
	if err != nil {
		log.Debugf("Error during findBucket: %s", err)
		return nil, err
	}
	var bucket *arq.ArqBucket
	for _, arqBackupSet := range arqBackupSets {
		if arqBackupSet.UUID == backupSetUUID {
			for _, folder := range arqBackupSet.Buckets {
				if folder.UUID == folderUUID {
					bucket = folder
				}
			}
		}
	}
	if bucket == nil {
		err := errors.New(fmt.Sprintf("Couldn't find backup set UUID %s, folder UUID %s.", backupSetUUID, folderUUID))
		log.Errorf("%s", err)
		return nil, err
	}
	return bucket, nil
}

func listDirectoryContents(c *cli.Context, connection connector.Connection) error {
	cacheDirectory := c.GlobalString("cache-directory")
	backupSetUUID := c.String("backup-set-uuid")
	folderUUID := c.String("folder-uuid")
	targetPath := c.String("path")

	bucket, err := findBucket(c, connection, backupSetUUID, folderUUID)
	if err != nil {
		err := errors.New(fmt.Sprintf("Couldn't find backup set UUID %s, folder UUID %s.", backupSetUUID, folderUUID))
		log.Errorf("%s", err)
		return err
	}
	backupSet := bucket.ArqBackupSet
	backupSet.CacheTreePackSets()

	tree, node, err := arq.FindNode(cacheDirectory, backupSet, bucket, targetPath)
	if err != nil {
		log.Errorf("Failed to find target path %s: %s", targetPath, err)
		return err
	}
	if node == nil || node.IsTree.IsTrue() {
		if tree == nil {
			err2 := errors.New(fmt.Sprintf("node is tree but no tree found: %s", node))
			log.Errorf("%s", err2)
			return err2
		}
		apsi, _ := arq.NewPackSetIndex(cacheDirectory, backupSet, bucket)
		for _, node := range tree.Nodes {
			if node.IsTree.IsTrue() {
				tree, err := apsi.GetPackFileAsTree(backupSet, bucket, *node.DataBlobKeys[0].SHA1)
				if err != nil {
					log.Debugf("Failed to find tree for node %s: %s", node, err)
					node.PrintOutput()
				} else if tree == nil {
					log.Debugf("directory node %s has no tree", node)
					node.PrintOutput()
				} else {
					tree.PrintOutput(node)
				}
			} else {
				node.PrintOutput()
			}
		}
	} else {
		node.PrintOutput()
	}
	return nil
}

func recover(c *cli.Context, connection connector.Connection) error {
	cacheDirectory := c.GlobalString("cache-directory")
	backupSetUUID := c.String("backup-set-uuid")
	folderUUID := c.String("folder-uuid")
	sourcePath := c.String("source-path")
	destinationPath := c.String("destination-path")

	if _, err := os.Stat(destinationPath); err == nil {
		err := errors.New(fmt.Sprintf("Destination path %s already exists, won't overwrite.", destinationPath))
		log.Errorf("%s", err)
		return err
	}
	bucket, err := findBucket(c, connection, backupSetUUID, folderUUID)
	if err != nil {
		err := errors.New(fmt.Sprintf("Couldn't find backup set UUID %s, folder UUID %s.", backupSetUUID, folderUUID))
		log.Errorf("%s", err)
		return err
	}
	backupSet := bucket.ArqBackupSet
	backupSet.CacheTreePackSets()
	backupSet.CacheBlobPackSets()

	tree, node, err := arq.FindNode(cacheDirectory, backupSet, bucket, sourcePath)
	if err != nil {
		log.Errorf("Failed to find source path %s: %s", sourcePath, err)
		return err
	}
	if node == nil || node.IsTree.IsTrue() {
		log.Errorf("unsupported right now. tree: %s", tree)
		return nil
	}
	apsi, _ := arq.NewPackSetIndex(cacheDirectory, backupSet, bucket)
	f, err := os.OpenFile(destinationPath, os.O_RDWR|os.O_CREATE|os.O_TRUNC, node.Mode)
	if err != nil {
		log.Errorf("Failed to open destinationPath %s: %s", destinationPath, err)
		return err
	}
	defer f.Close()
	err = f.Truncate(int64(node.UncompressedDataSize))
	if err != nil {
		log.Errorf("Failed to pre-allocate size of file %s: %s", destinationPath, err)
	}
	w := bufio.NewWriter(f)
	defer w.Flush()
	r, err := arq.GetReaderForBlobKeys(node.DataBlobKeys, apsi, backupSet, bucket)
	if err != nil {
		log.Errorf("Failed during GetReaderForBlobKeys for node %s: %s", node, err)
		return err
	}
	io.Copy(w, r)
	return nil
}

func main() {
	defaultCacheDirectory, err := homedir.Expand("~/.arqinator_cache")
	if err != nil {
		log.Fatal("Failed to get user's home dir: ", err)
	}

	app := cli.NewApp()
	app.Name = "arqinator"
	app.Usage = "restore folders and files from Arq backups"
	app.Version = "0.1"
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "backup-type",
			Usage: "Method used for backup, one of: ['s3', 'googlecloudstorage']",
		},
		cli.StringFlag{
			Name:  "s3-region",
			Usage: "AWS S3 region, e.g. 'us-west-2'.",
		},
		cli.StringFlag{
			Name:  "s3-bucket-name",
			Usage: "AWS S3 bucket name, e.g. 'arq-akiaabdefg-us-west-2'.",
		},
		cli.StringFlag{
			Name:  "gcs-json-private-key-filepath",
			Usage: "Google Cloud Storage JSON private key filepath. See: https://goo.gl/SK5Rb7",
		},
		cli.StringFlag{
			Name:  "gcs-project-id",
			Usage: "Google Cloud Storage project ID.",
		},
		cli.StringFlag{
			Name:  "gcs-bucket-name",
			Usage: "Google Cloud Storage bucket name.",
		},
		cli.StringFlag{
			Name:  "cache-directory",
			Value: defaultCacheDirectory,
			Usage: fmt.Sprintf("Where to cache Arq files for browsing. Default: %s", defaultCacheDirectory),
		},
		cli.BoolFlag{
			Name:  "delete-cache-directory",
			Usage: "Delete cache directory before starting. Useful if seeing errors that could be due to truncated downloads.",
		},
		cli.BoolFlag{
			Name:  "verbose",
			Usage: "Enable verbose logging",
		},
	}
	app.Commands = []cli.Command{
		{
			Name:  "list-backup-sets",
			Usage: "List backup sets in this account.",
			Action: func(c *cli.Context) {
				if err := cliSetup(c); err != nil {
					log.Errorf("%s", err)
					return
				}
				connection, err := getConnection(c)
				if err != nil {
					log.Errorf("%s", err)
					return
				}
				if err := listBackupSets(c, connection); err != nil {
					log.Errorf("%s", err)
					return
				}
			},
		},
		{
			Name:  "list-directory-contents",
			Usage: "List contents of directory in backup.",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "backup-set-uuid",
					Usage: "UUID of backup set. Use 'list-backup-sets' to determine this.",
				},
				cli.StringFlag{
					Name:  "folder-uuid",
					Usage: "UUID of folder. Use 'list-backup-sets' to determine this.",
				},
				cli.StringFlag{
					Name:  "path",
					Usage: "Path of directory or file in backup",
				},
			},
			Action: func(c *cli.Context) {
				if err := cliSetup(c); err != nil {
					log.Errorf("%s", err)
					return
				}
				connection, err := getConnection(c)
				if err != nil {
					log.Errorf("%s", err)
					return
				}
				if err := listDirectoryContents(c, connection); err != nil {
					log.Errorf("%s", err)
					return
				}
			},
		},
		{
			Name:  "recover",
			Usage: "Recover a file or directory from a backup",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "backup-set-uuid",
					Usage: "UUID of backup set. Use 'list-backup-sets' to determine this.",
				},
				cli.StringFlag{
					Name:  "folder-uuid",
					Usage: "UUID of folder. Use 'list-backup-sets' to determine this.",
				},
				cli.StringFlag{
					Name:  "source-path",
					Usage: "Path of directory or file in backup",
				},
				cli.StringFlag{
					Name:  "destination-path",
					Usage: "Path to recover directory or file into. Must not already exist.",
				},
			},
			Action: func(c *cli.Context) {
				if err := cliSetup(c); err != nil {
					log.Errorf("%s", err)
					return
				}
				connection, err := getConnection(c)
				if err != nil {
					log.Errorf("%s", err)
					return
				}
				if err := recover(c, connection); err != nil {
					log.Errorf("%s", err)
					return
				}
			},
		},
	}
	app.Run(os.Args)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"net/http"
	"time"

	"github.com/iron-io/iron_go/config"
	"github.com/iron-io/iron_go/mq"
	"github.com/iron-io/iron_go/worker"
)

const (
	interval   = 15 * time.Second
	maxRunTime = 30 * time.Minute
	swapi      = "worker-aws-us-east-1.iron.io"
)

const (
	TriggerFixed = iota
	TriggerProgressive
	TriggerRatio
)

var (
	prev    map[string]int
	codeIds map[string]string
)

type Config struct {
	Environments map[string]config.Settings `json:"envs"`
	Alerts       []QueueWorkerAlert
}

type QueueWorkerAlert struct {
	QueueName  string
	QueueEnv   string
	WorkerName string
	WorkerEnv  string
	Cluster    string
	Triggers   []Trigger
}

type Trigger struct {
	Typ   int
	Value int
}

func queueKey(qw QueueWorkerAlert) string {
	return qw.QueueEnv + "|" + qw.QueueName
}

func main() {
	start := time.Now()
	prev = make(map[string]int)
	codeIds = make(map[string]string)

	// Retrieve configuration
	c := &Config{}
	worker.ParseFlags()
	worker.ConfigFromJSON(c)

	for {
		if time.Since(start) > maxRunTime {
			break
		}

		for _, alert := range c.Alerts {
			queueSize, prevQueueSize := 0, 0
			key := queueKey(alert)

			// Get previous size
			if v, e := prev[key]; e {
				prevQueueSize = v
			}

			queueEnv, exists := c.Environments[alert.QueueEnv]
			if !exists {
				fmt.Printf("Environment %s is not defined for queue %s\n", alert.QueueEnv, alert.QueueName)
				continue
			}
			queueConfig := config.ManualConfig("iron_mq", &queueEnv)
			q := mq.ConfigNew(alert.QueueName, &queueConfig)
			info, err := q.Info()
			if err != nil {
				fmt.Println("Could not get information about", alert.QueueName, err)
				continue
			}

			queueSize = info.Size
			// Update previous size
			prev[key] = info.Size

			workerEnv, exists := c.Environments[alert.WorkerEnv]
			if !exists {
				fmt.Printf("Environment %s is not defined for worker %s\n", alert.WorkerEnv, alert.WorkerName)
				continue
			}
			queued, running, err := workerStats(&workerEnv, alert.WorkerName)
			if err != nil {
				fmt.Printf("Could not get code stats for %s, %v", alert.WorkerName, err)
				continue
			}

			launch := evalTriggers(queued, running, queueSize, prevQueueSize, alert.Triggers)
			_ = launch

			workerConfig := config.ManualConfig("iron_worker", &workerEnv)
			w := &worker.Worker{Settings: workerConfig}

			tasks := make([]worker.Task, launch)
			for x := 0; x < len(tasks); x++ {
				tasks[x].CodeName = alert.WorkerName
				tasks[x].Cluster = alert.Cluster
			}
			_, err = w.TaskQueue(tasks...)
			if err != nil {
				fmt.Println("Could not create tasks for", alert.WorkerName)
				continue
			}
		}

	}
}

func workerKey(projectID, codeName string) string {
	return projectID + "|" + codeName
}

type CodeStats struct {
	Running int `json:"running"`
	Queued  int `json:"queued"`
	// ignore other states
}

func workerStats(env *config.Settings, codeName string) (queued, running int, err error) {
	codeID, exists := codeIds[workerKey(env.ProjectId, codeName)]
	if !exists {
		workerConfig := config.ManualConfig("iron_worker", env)
		w := &worker.Worker{Settings: workerConfig}
		codes, err := w.CodePackageList(0, 100)
		if err != nil {
			return 0, 0, err
		}

		for _, c := range codes {
			codeIds[workerKey(c.ProjectId, c.Name)] = c.Id
			if c.Name == codeName {
				codeID = c.Id
			}
		}
	}

	if len(codeID) == 0 {
		return 0, 0, fmt.Errorf("Could not get id for %s", codeName)
	}
	if len(env.ProjectId) == 0 || len(env.Token) == 0 {
		return 0, 0, fmt.Errorf("Could not get env for %s", codeName)
	}

	resp, err := http.Get(fmt.Sprintf("https://%s/2/projects/%s/codes/%s/stats?oauth=%s", swapi, env.ProjectId, codeID, env.Token))
	if err != nil {
		return 0, 0, err
	}
	defer resp.Body.Close()
	decoder := json.NewDecoder(resp.Body)
	var s CodeStats
	err = decoder.Decode(&s)
	if err != nil {
		return 0, 0, err
	}

	return s.Queued, s.Running, nil
}

func evalTriggers(queued, running, queueSize, prevQueueSize int, triggers []Trigger) (launch int) {
	for _, t := range triggers {
		switch t.Typ {
		case TriggerFixed:
			if queueSize >= t.Value {
				if t.Value <= prevQueueSize {
					continue
				}
				launch++
			}
		case TriggerProgressive:
			if queueSize < t.Value {
				continue
			}

			previous_level := prevQueueSize / t.Value
			current_level := queueSize / t.Value
			if current_level > previous_level {
				launch += current_level - previous_level
			}
		case TriggerRatio:
			expected_runners := (queueSize + t.Value - 1) / t.Value // Only have 0 runners if qsize=0

			diff := expected_runners - (queued + running)
			if diff > 0 {
				launch += diff
			}
		}
	}
	return
}
<commit_msg>Working<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"time"

	"github.com/iron-io/iron_go3/config"
	"github.com/iron-io/iron_go3/mq"
	"github.com/iron-io/iron_go3/worker"
)

const (
	interval   = 5 * time.Second
	maxRunTime = 30 * time.Minute
	swapi      = "worker-aws-us-east-1.iron.io"
)

const (
	TriggerFixed       = "fixed"
	TriggerProgressive = "progressive"
	TriggerRatio       = "ratio"
)

var (
	prev    map[string]int
	codeIds map[string]string
)

type Config struct {
	Environments map[string]config.Settings `json:"envs"`
	Alerts       []QueueWorkerAlert         `json:"alerts"`
}

type QueueWorkerAlert struct {
	QueueName  string    `json:"queueName"`
	QueueEnv   string    `json:"queueEnv"`
	WorkerName string    `json:"workerName"`
	WorkerEnv  string    `json:"workerEnv"`
	Cluster    string    `json:"cluster"`
	Triggers   []Trigger `json:"triggers"`
}

type Trigger struct {
	Typ   string `json:"type"`
	Value int    `json:"value"`
}

func queueKey(qw QueueWorkerAlert) string {
	return qw.QueueEnv + "|" + qw.QueueName
}

func main() {
	start := time.Now()
	prev = make(map[string]int)
	codeIds = make(map[string]string)

	// Retrieve configuration
	c := &Config{}
	worker.ParseFlags()
	err := worker.ConfigFromJSON(c)
	if err != nil {
		log.Fatalln("Could not unparse config", err)
	}

	if len(c.Alerts) == 0 || len(c.Environments) == 0 {
		fmt.Println("No config set")
		return
	}

	for {
		if time.Since(start) > maxRunTime {
			fmt.Println("No triggers specified for an alert")
			break
		}

		for _, alert := range c.Alerts {
			if len(alert.Triggers) == 0 {
				continue
			}

			queueSize, prevQueueSize := 0, 0
			key := queueKey(alert)

			// Get previous size
			if v, e := prev[key]; e {
				prevQueueSize = v
			}

			queueEnv, exists := c.Environments[alert.QueueEnv]
			if !exists {
				fmt.Printf("Environment %q is not defined for queue %q\n", alert.QueueEnv, alert.QueueName)
				continue
			}

			queueConfig := config.ManualConfig("iron_mq", &queueEnv)
			q := mq.ConfigNew(alert.QueueName, &queueConfig)
			info, err := q.Info()
			if err != nil {
				fmt.Println("Could not get information about", alert.QueueName, err)
				continue
			}

			queueSize = info.Size
			// Update previous size
			prev[key] = info.Size

			workerEnv, exists := c.Environments[alert.WorkerEnv]
			if !exists {
				fmt.Printf("Environment %q is not defined for worker %q\n", alert.WorkerEnv, alert.WorkerName)
				continue
			}
			queued, running, err := workerStats(&workerEnv, alert.WorkerName)
			if err != nil {
				fmt.Printf("Could not get code stats for %s, %v", alert.WorkerName, err)
				continue
			}

			launch := evalTriggers(queued, running, queueSize, prevQueueSize, alert.Triggers)
			fmt.Printf("Queue: %s (size=%d, prev=%d), CodeName=%s (queued=%d, running=%d), Launching %d\n", alert.QueueName, queueSize, prevQueueSize, alert.WorkerName, queued, running, launch)

			if launch > 0 {
				workerConfig := config.ManualConfig("iron_worker", &workerEnv)
				w := &worker.Worker{Settings: workerConfig}

				tasks := make([]worker.Task, launch)
				for x := 0; x < len(tasks); x++ {
					tasks[x].CodeName = alert.WorkerName
					tasks[x].Cluster = alert.Cluster
				}

				_, err = w.TaskQueue(tasks...)
				if err != nil {
					fmt.Println("Could not create tasks for", alert.WorkerName)
					continue
				}
			}
		}

		time.Sleep(interval)
	}
}

func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}

func workerKey(projectID, codeName string) string {
	return projectID + "|" + codeName
}

type CodeStats struct {
	Running int `json:"running"`
	Queued  int `json:"queued"`
	// ignore other states
}

func workerStats(env *config.Settings, codeName string) (queued, running int, err error) {
	codeID, exists := codeIds[workerKey(env.ProjectId, codeName)]
	if !exists {
		workerConfig := config.ManualConfig("iron_worker", env)
		w := &worker.Worker{Settings: workerConfig}
		codes, err := w.CodePackageList(0, 100)
		if err != nil {
			return 0, 0, err
		}

		for _, c := range codes {
			codeIds[workerKey(c.ProjectId, c.Name)] = c.Id
			if c.Name == codeName {
				codeID = c.Id
			}
		}
	}

	if len(codeID) == 0 {
		return 0, 0, fmt.Errorf("Could not get id for %s", codeName)
	}
	if len(env.ProjectId) == 0 || len(env.Token) == 0 {
		return 0, 0, fmt.Errorf("Could not get env for %s", codeName)
	}

	url := fmt.Sprintf("https://%s/2/projects/%s/codes/%s/stats?oauth=%s", swapi, env.ProjectId, codeID, env.Token)
	resp, err := http.Get(url)
	if err != nil {
		return 0, 0, err
	}
	defer resp.Body.Close()
	decoder := json.NewDecoder(resp.Body)
	var s CodeStats
	err = decoder.Decode(&s)
	if err != nil {
		return 0, 0, err
	}

	return s.Queued, s.Running, nil
}

func evalTriggers(queued, running, queueSize, prevQueueSize int, triggers []Trigger) (launch int) {
	for _, t := range triggers {
		switch t.Typ {
		case TriggerFixed:
			if queueSize >= t.Value {
				if t.Value <= prevQueueSize {
					continue
				}
				launch = max(launch, 1)
			}
		case TriggerProgressive:
			if queueSize < t.Value {
				continue
			}

			previous_level := prevQueueSize / t.Value
			current_level := queueSize / t.Value
			if current_level > previous_level {
				launch = max(launch, current_level-previous_level)
			}
		case TriggerRatio:
			expected_runners := (queueSize + t.Value - 1) / t.Value // Only have 0 runners if qsize=0

			diff := expected_runners - (queued + running)
			if diff > 0 {
				launch = max(launch, diff)
			}
		}
	}
	return
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"
	"sync"
	"time"

	"github.com/mesosphere/mesos-dns/logging"
	"github.com/mesosphere/mesos-dns/records"
	"github.com/mesosphere/mesos-dns/resolver"

	"github.com/miekg/dns"
)

func main() {
	var wg sync.WaitGroup
	var resolver resolver.Resolver

	versionFlag := false

	cjson := flag.String("config", "config.json", "location of configuration file (json)")
	flag.BoolVar(&logging.VerboseFlag, "e", false, "verbose logging")
	flag.BoolVar(&logging.VeryVerboseFlag, "ee", false, "very verbose logging")
	flag.BoolVar(&versionFlag, "version", false, "output the version")
	flag.Parse()

	if versionFlag {
		fmt.Println(version)
		os.Exit(0)
	}

	logging.SetupLogs()

	resolver.Config = records.SetConfig(*cjson)

	// if ZK is identified, start detector
	if len(resolver.Config.Zk) != 0 {
		dr := make(chan bool)
		go records.ZKdetect(&resolver.Config, dr)
		<-dr

		// wait for the first read from ZK
		//for {
		//	if resolver.Config.StartZk == true {
		//		break
		//	}
		//}
		// horrible hack
		//time.Sleep(100 * time.Millisecond)
	}

	// reload the first time
	resolver.Reload()
	ticker := time.NewTicker(time.Second * time.Duration(resolver.Config.RefreshSeconds))
	go func() {
		for _ = range ticker.C {
			resolver.Reload()
			logging.PrintCurLog()
		}
	}()

	// handle for everything in this domain...
	dns.HandleFunc(resolver.Config.Domain+".", panicRecover(resolver.HandleMesos))
	dns.HandleFunc(".", panicRecover(resolver.HandleNonMesos))

	go resolver.Serve("tcp")
	go resolver.Serve("udp")

	wg.Add(1)
	wg.Wait()
}

// panicRecover catches any panics from the resolvers and sets an error
// code of server failure
func panicRecover(f func(w dns.ResponseWriter, r *dns.Msg)) func(w dns.ResponseWriter, r *dns.Msg) {
	return func(w dns.ResponseWriter, r *dns.Msg) {
		defer func() {
			if rec := recover(); rec != nil {
				m := new(dns.Msg)
				m.SetReply(r)
				m.SetRcode(r, 2)
				_ = w.WriteMsg(m)
				logging.Error.Println(rec)
			}
		}()
		f(w, r)
	}
}
<commit_msg>Synchronization fix.<commit_after>package main

import (
	"flag"
	"fmt"
	"os"
	"sync"
	"time"

	"github.com/mesosphere/mesos-dns/logging"
	"github.com/mesosphere/mesos-dns/records"
	"github.com/mesosphere/mesos-dns/resolver"

	"github.com/miekg/dns"
)

func main() {
	var wg sync.WaitGroup
	var resolver resolver.Resolver

	versionFlag := false

	cjson := flag.String("config", "config.json", "location of configuration file (json)")
	flag.BoolVar(&logging.VerboseFlag, "e", false, "verbose logging")
	flag.BoolVar(&logging.VeryVerboseFlag, "ee", false, "very verbose logging")
	flag.BoolVar(&versionFlag, "version", false, "output the version")
	flag.Parse()

	if versionFlag {
		fmt.Println(version)
		os.Exit(0)
	}

	logging.SetupLogs()

	resolver.Config = records.SetConfig(*cjson)

	// if ZK is identified, start detector and wait for first master
	if len(resolver.Config.Zk) != 0 {
		dr := make(chan bool)
		go records.ZKdetect(&resolver.Config, dr)
		<-dr
		close(dr)
	}

	// reload the first time
	resolver.Reload()
	ticker := time.NewTicker(time.Second * time.Duration(resolver.Config.RefreshSeconds))
	go func() {
		for _ = range ticker.C {
			resolver.Reload()
			logging.PrintCurLog()
		}
	}()

	// handle for everything in this domain...
	dns.HandleFunc(resolver.Config.Domain+".", panicRecover(resolver.HandleMesos))
	dns.HandleFunc(".", panicRecover(resolver.HandleNonMesos))

	go resolver.Serve("tcp")
	go resolver.Serve("udp")

	wg.Add(1)
	wg.Wait()
}

// panicRecover catches any panics from the resolvers and sets an error
// code of server failure
func panicRecover(f func(w dns.ResponseWriter, r *dns.Msg)) func(w dns.ResponseWriter, r *dns.Msg) {
	return func(w dns.ResponseWriter, r *dns.Msg) {
		defer func() {
			if rec := recover(); rec != nil {
				m := new(dns.Msg)
				m.SetReply(r)
				m.SetRcode(r, 2)
				_ = w.WriteMsg(m)
				logging.Error.Println(rec)
			}
		}()
		f(w, r)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"log"
	"os"
	"strconv"
	"strings"
)

var (
	pretty = flag.Bool("p", false, "indent (pretty print) output")
)

var allowedImports = map[string]struct{}{
	"fmt": struct{}{},
	"log": struct{}{},
}

type block struct {
	Statements []statement `json:"statements"`
}

type dataType struct {
	Name string `json:"name"`
}

type parameter struct {
	Name     string   `json:"name"`
	DataType dataType `json:"data-type"`
}

type statement struct {
	Line       int         `json:"line"`
	Type       string      `json:"type"`
	Name       string      `json:"name"`
	ReturnType *dataType   `json:"return-type,omitempty"`
	Parameters []parameter `json:"parameters,omitempty"`
	Arguments  []statement `json:"arguments,omitempty"`
	Left       *statement  `json:"left,omitempty"`
	Right      *statement  `json:"right,omitempty"`
	Block      *block      `json:"block,omitempty"`
}

type superAST struct {
	nodeStack  []ast.Node
	blockStack []*block
	fset       *token.FileSet
}

func newSuperAST(fset *token.FileSet) *superAST {
	a := &superAST{
		fset: fset,
	}
	a.blockStack = append(a.blockStack, new(block))
	return a
}

func (a *superAST) Visit(node ast.Node) ast.Visitor {
	if node == nil {
		popNode := a.nodeStack[len(a.nodeStack)-1]
		switch popNode.(type) {
		case *ast.FuncDecl:
			a.blockStack = a.blockStack[:len(a.blockStack)-1]
		}
		a.nodeStack = a.nodeStack[:len(a.nodeStack)-1]
		log.Printf("%s}", strings.Repeat("  ", len(a.nodeStack)))
		return nil
	}
	curBlock := a.blockStack[len(a.blockStack)-1]
	pos := a.fset.Position(node.Pos())
	log.Printf("%s%T - %#v", strings.Repeat("  ", len(a.nodeStack)), node, pos)
	switch x := node.(type) {
	case *ast.BasicLit:
	case *ast.BlockStmt:
	case *ast.CallExpr:
		call := statement{
			Line: pos.Line,
			Type: "function-call",
			Name: "print",
		}
		curBlock.Statements = append(curBlock.Statements, call)
	case *ast.ExprStmt:
	case *ast.FieldList:
	case *ast.File:
		pname := x.Name.Name
		if pname != "main" {
			log.Fatalf(`Package name is not "main": "%s"`, pname)
		}
		imports := x.Imports
		for _, imp := range imports {
			path, err := strconv.Unquote(imp.Path.Value)
			if err != nil {
				log.Fatalf("Error when unquoting import: %s", err)
			}
			if _, e := allowedImports[path]; !e {
				log.Fatalf(`Import path not allowed: "%s"`, path)
			}
		}
	case *ast.FuncDecl:
		name := x.Name.Name
		/*var params, results []*ast.Field
		if x.Type.Params != nil {
			params = x.Type.Params.List
		}
		if x.Type.Results != nil {
			results = x.Type.Results.List
		}*/
		function := statement{
			Line: pos.Line,
			Type: "function-declaration",
			Name: name,
			ReturnType: &dataType{
				Name: "int",
			},
			Block: new(block),
		}
		curBlock.Statements = append(curBlock.Statements, function)
		a.blockStack = append(a.blockStack, function.Block)
	case *ast.FuncType:
	case *ast.GenDecl:
	case *ast.Ident:
	case *ast.ImportSpec:
	case *ast.SelectorExpr:
	default:
		log.Printf("Uncatched ast.Node type: %T\n", node)
	}
	a.nodeStack = append(a.nodeStack, node)
	return a
}

func main() {
	flag.Parse()
	fset := token.NewFileSet()
	src := `
package main

import "fmt"

func main() {
	fmt.Println("Hello, World!")
}
`
	f, err := parser.ParseFile(fset, "hello_world.go", src, 0)
	if err != nil {
		log.Fatal(err)
	}
	a := newSuperAST(fset)
	ast.Walk(a, f)

	rootBlock := a.blockStack[0]
	if *pretty {
		b, err := json.Marshal(rootBlock)
		if err != nil {
			log.Fatal(err)
		}
		var out bytes.Buffer
		if err := json.Indent(&out, b, "", "  "); err != nil {
			log.Fatal(err)
		}
		if _, err := out.WriteTo(os.Stdout); err != nil {
			log.Fatal(err)
		}
		fmt.Printf("\n")
	} else {
		enc := json.NewEncoder(os.Stdout)
		if err := enc.Encode(rootBlock); err != nil {
			log.Println(err)
		}
	}

}
<commit_msg>Keep a stack of statement slices instead of blocks<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"log"
	"os"
	"strconv"
	"strings"
)

var (
	pretty = flag.Bool("p", false, "indent (pretty print) output")
)

var allowedImports = map[string]struct{}{
	"fmt": struct{}{},
	"log": struct{}{},
}

type block struct {
	Statements []statement `json:"statements"`
}

type dataType struct {
	Name string `json:"name"`
}

type parameter struct {
	Name     string   `json:"name"`
	DataType dataType `json:"data-type"`
}

type statement struct {
	Line       int         `json:"line"`
	Type       string      `json:"type"`
	Name       string      `json:"name"`
	ReturnType *dataType   `json:"return-type,omitempty"`
	Parameters []parameter `json:"parameters,omitempty"`
	Arguments  []statement `json:"arguments,omitempty"`
	Left       *statement  `json:"left,omitempty"`
	Right      *statement  `json:"right,omitempty"`
	Block      *block      `json:"block,omitempty"`
}

type superAST struct {
	RootBlock  *block
	nodeStack  []ast.Node
	stmtsStack []*[]statement
	fset       *token.FileSet
}

func newSuperAST(fset *token.FileSet) *superAST {
	a := &superAST{
		fset:      fset,
		RootBlock: new(block),
	}
	a.stmtsStack = append(a.stmtsStack, &a.RootBlock.Statements)
	return a
}

func (a *superAST) Visit(node ast.Node) ast.Visitor {
	if node == nil {
		popNode := a.nodeStack[len(a.nodeStack)-1]
		switch popNode.(type) {
		case *ast.CallExpr:
			a.stmtsStack = a.stmtsStack[:len(a.stmtsStack)-1]
		case *ast.FuncDecl:
			a.stmtsStack = a.stmtsStack[:len(a.stmtsStack)-1]
		}
		a.nodeStack = a.nodeStack[:len(a.nodeStack)-1]
		log.Printf("%s}", strings.Repeat("  ", len(a.nodeStack)))
		return nil
	}
	curStatements := a.stmtsStack[len(a.stmtsStack)-1]
	pos := a.fset.Position(node.Pos())
	log.Printf("%s%T - %#v", strings.Repeat("  ", len(a.nodeStack)), node, pos)
	switch x := node.(type) {
	case *ast.BasicLit:
	case *ast.BlockStmt:
	case *ast.CallExpr:
		call := statement{
			Line: pos.Line,
			Type: "function-call",
			Name: "print",
		}
		*curStatements = append(*curStatements, call)
	case *ast.ExprStmt:
	case *ast.FieldList:
	case *ast.File:
		pname := x.Name.Name
		if pname != "main" {
			log.Fatalf(`Package name is not "main": "%s"`, pname)
		}
		imports := x.Imports
		for _, imp := range imports {
			path, err := strconv.Unquote(imp.Path.Value)
			if err != nil {
				log.Fatalf("Error when unquoting import: %s", err)
			}
			if _, e := allowedImports[path]; !e {
				log.Fatalf(`Import path not allowed: "%s"`, path)
			}
		}
	case *ast.FuncDecl:
		name := x.Name.Name
		/*var params, results []*ast.Field
		if x.Type.Params != nil {
			params = x.Type.Params.List
		}
		if x.Type.Results != nil {
			results = x.Type.Results.List
		}*/
		fn := statement{
			Line: pos.Line,
			Type: "function-declaration",
			Name: name,
			ReturnType: &dataType{
				Name: "int",
			},
			Block: new(block),
		}
		*curStatements = append(*curStatements, fn)
		a.stmtsStack = append(a.stmtsStack, &fn.Block.Statements)
	case *ast.FuncType:
	case *ast.GenDecl:
	case *ast.Ident:
	case *ast.ImportSpec:
	case *ast.SelectorExpr:
	default:
		log.Printf("Uncatched ast.Node type: %T\n", node)
	}
	a.nodeStack = append(a.nodeStack, node)
	return a
}

func main() {
	flag.Parse()
	fset := token.NewFileSet()
	src := `
package main

import "fmt"

func main() {
	fmt.Println("Hello, World!")
}
`
	f, err := parser.ParseFile(fset, "hello_world.go", src, 0)
	if err != nil {
		log.Fatal(err)
	}
	a := newSuperAST(fset)
	ast.Walk(a, f)

	if *pretty {
		b, err := json.Marshal(a.RootBlock)
		if err != nil {
			log.Fatal(err)
		}
		var out bytes.Buffer
		if err := json.Indent(&out, b, "", "  "); err != nil {
			log.Fatal(err)
		}
		if _, err := out.WriteTo(os.Stdout); err != nil {
			log.Fatal(err)
		}
		fmt.Printf("\n")
	} else {
		enc := json.NewEncoder(os.Stdout)
		if err := enc.Encode(a.RootBlock); err != nil {
			log.Println(err)
		}
	}

}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strings"
	"sync"

	"github.com/asim/mq/go/client"
	"github.com/gorilla/handlers"
	"github.com/gorilla/websocket"
)

type mq struct {
	client client.Client

	sync.RWMutex
	topics map[string][]chan []byte
}

var (
	address = flag.String("address", ":8081", "MQ server address")
	cert    = flag.String("cert_file", "", "TLS certificate file")
	key     = flag.String("key_file", "", "TLS key file")
	proxy   = flag.Bool("proxy", false, "Proxy for an MQ cluster")
	servers = flag.String("servers", "", "Comma separated MQ cluster list used by Proxy")

	defaultMQ *mq
)

func init() {
	flag.Parse()

	if *proxy && len(*servers) == 0 {
		log.Fatal("Proxy enabled without MQ server list")
	}

	defaultMQ = &mq{
		client: client.New(client.WithServers(strings.Split(*servers, ",")...)),
		topics: make(map[string][]chan []byte),
	}
}

func Log(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		log.Printf("%s %s %s", r.RemoteAddr, r.Method, r.URL)
		handler.ServeHTTP(w, r)
	})
}

func (m *mq) pub(topic string, payload []byte) error {
	if *proxy {
		return m.client.Publish(topic, payload)
	}

	m.RLock()
	subscribers, ok := m.topics[topic]
	m.RUnlock()
	if !ok {
		return nil
	}

	go func() {
		for _, subscriber := range subscribers {
			select {
			case subscriber <- payload:
			default:
			}
		}
	}()

	return nil
}

func (m *mq) sub(topic string) (<-chan []byte, error) {
	if *proxy {
		return m.client.Subscribe(topic)
	}

	ch := make(chan []byte, 100)
	m.Lock()
	m.topics[topic] = append(m.topics[topic], ch)
	m.Unlock()
	return ch, nil
}

func (m *mq) unsub(topic string, sub <-chan []byte) error {
	if *proxy {
		return m.client.Unsubscribe(sub)
	}

	m.RLock()
	subscribers, ok := m.topics[topic]
	m.RUnlock()

	if !ok {
		return nil
	}

	var subs []chan []byte
	for _, subscriber := range subscribers {
		if subscriber == sub {
			continue
		}
		subs = append(subs, subscriber)
	}

	m.Lock()
	m.topics[topic] = subs
	m.Unlock()

	return nil
}

func pub(w http.ResponseWriter, r *http.Request) {
	topic := r.URL.Query().Get("topic")
	b, err := ioutil.ReadAll(r.Body)
	if err != nil {
		http.Error(w, "Pub error", http.StatusInternalServerError)
		return
	}
	r.Body.Close()

	err = defaultMQ.pub(topic, b)
	if err != nil {
		http.Error(w, "Pub error", http.StatusInternalServerError)
		return
	}
}

func sub(w http.ResponseWriter, r *http.Request) {
	conn, err := websocket.Upgrade(w, r, w.Header(), 1024, 1024)
	if err != nil {
		log.Println("Failed to open websocket connection")
		http.Error(w, "Could not open websocket connection", http.StatusBadRequest)
		return
	}

	topic := r.URL.Query().Get("topic")
	ch, err := defaultMQ.sub(topic)
	if err != nil {
		log.Printf("Failed to retrieve event for %s topic", topic)
		http.Error(w, "Could not retrieve events", http.StatusInternalServerError)
		return
	}
	defer defaultMQ.unsub(topic, ch)

	for {
		select {
		case e := <-ch:
			if err = conn.WriteMessage(websocket.BinaryMessage, e); err != nil {
				log.Printf("error sending event: %v", err.Error())
				return
			}
		}
	}
}

func main() {
	// MQ Handlers
	http.HandleFunc("/pub", pub)
	http.HandleFunc("/sub", sub)

	// logging handler
	handler := handlers.LoggingHandler(os.Stdout, http.DefaultServeMux)

	if len(*cert) > 0 && len(*key) > 0 {
		log.Println("TLS Enabled")
		log.Println("MQ listening on", *address)
		http.ListenAndServeTLS(*address, *cert, *key, handler)
		return
	}

	if *proxy {
		log.Println("Proxy enabled")
	}

	log.Println("MQ listening on", *address)
	http.ListenAndServe(*address, handler)
}
<commit_msg>remove Log func<commit_after>package main

import (
	"flag"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strings"
	"sync"

	"github.com/asim/mq/go/client"
	"github.com/gorilla/handlers"
	"github.com/gorilla/websocket"
)

type mq struct {
	client client.Client

	sync.RWMutex
	topics map[string][]chan []byte
}

var (
	address = flag.String("address", ":8081", "MQ server address")
	cert    = flag.String("cert_file", "", "TLS certificate file")
	key     = flag.String("key_file", "", "TLS key file")
	proxy   = flag.Bool("proxy", false, "Proxy for an MQ cluster")
	servers = flag.String("servers", "", "Comma separated MQ cluster list used by Proxy")

	defaultMQ *mq
)

func init() {
	flag.Parse()

	if *proxy && len(*servers) == 0 {
		log.Fatal("Proxy enabled without MQ server list")
	}

	defaultMQ = &mq{
		client: client.New(client.WithServers(strings.Split(*servers, ",")...)),
		topics: make(map[string][]chan []byte),
	}
}

func (m *mq) pub(topic string, payload []byte) error {
	if *proxy {
		return m.client.Publish(topic, payload)
	}

	m.RLock()
	subscribers, ok := m.topics[topic]
	m.RUnlock()
	if !ok {
		return nil
	}

	go func() {
		for _, subscriber := range subscribers {
			select {
			case subscriber <- payload:
			default:
			}
		}
	}()

	return nil
}

func (m *mq) sub(topic string) (<-chan []byte, error) {
	if *proxy {
		return m.client.Subscribe(topic)
	}

	ch := make(chan []byte, 100)
	m.Lock()
	m.topics[topic] = append(m.topics[topic], ch)
	m.Unlock()
	return ch, nil
}

func (m *mq) unsub(topic string, sub <-chan []byte) error {
	if *proxy {
		return m.client.Unsubscribe(sub)
	}

	m.RLock()
	subscribers, ok := m.topics[topic]
	m.RUnlock()

	if !ok {
		return nil
	}

	var subs []chan []byte
	for _, subscriber := range subscribers {
		if subscriber == sub {
			continue
		}
		subs = append(subs, subscriber)
	}

	m.Lock()
	m.topics[topic] = subs
	m.Unlock()

	return nil
}

func pub(w http.ResponseWriter, r *http.Request) {
	topic := r.URL.Query().Get("topic")
	b, err := ioutil.ReadAll(r.Body)
	if err != nil {
		http.Error(w, "Pub error", http.StatusInternalServerError)
		return
	}
	r.Body.Close()

	err = defaultMQ.pub(topic, b)
	if err != nil {
		http.Error(w, "Pub error", http.StatusInternalServerError)
		return
	}
}

func sub(w http.ResponseWriter, r *http.Request) {
	conn, err := websocket.Upgrade(w, r, w.Header(), 1024, 1024)
	if err != nil {
		log.Println("Failed to open websocket connection")
		http.Error(w, "Could not open websocket connection", http.StatusBadRequest)
		return
	}

	topic := r.URL.Query().Get("topic")
	ch, err := defaultMQ.sub(topic)
	if err != nil {
		log.Printf("Failed to retrieve event for %s topic", topic)
		http.Error(w, "Could not retrieve events", http.StatusInternalServerError)
		return
	}
	defer defaultMQ.unsub(topic, ch)

	for {
		select {
		case e := <-ch:
			if err = conn.WriteMessage(websocket.BinaryMessage, e); err != nil {
				log.Printf("error sending event: %v", err.Error())
				return
			}
		}
	}
}

func main() {
	// MQ Handlers
	http.HandleFunc("/pub", pub)
	http.HandleFunc("/sub", sub)

	// logging handler
	handler := handlers.LoggingHandler(os.Stdout, http.DefaultServeMux)

	if len(*cert) > 0 && len(*key) > 0 {
		log.Println("TLS Enabled")
		log.Println("MQ listening on", *address)
		http.ListenAndServeTLS(*address, *cert, *key, handler)
		return
	}

	if *proxy {
		log.Println("Proxy enabled")
	}

	log.Println("MQ listening on", *address)
	http.ListenAndServe(*address, handler)
}
<|endoftext|>