<commit_before>package looli

import (
	"net/http"
)

type (
	Engine struct {
		// A Server defines parameters for running an HTTP server. The zero value for Server is a valid configuration.
		Server *http.Server

		// router with basePath
		RouterPrefix

		// router used to match url
		router *Router

		// when set true, implements a best effort algorithm to return the real client IP, it parses
		// X-Real-IP and X-Forwarded-For in order to work properly with reverse-proxies such us: nginx or haproxy.
		ForwardedByClientIP bool
	}
	HandlerFunc func(*Context)
)

type Handler interface {
	Handle(*Context)
}

func New() *Engine {
	engine := &Engine{
		Server:       &http.Server{},
		RouterPrefix: RouterPrefix{},
		router:       NewRouter(),
	}

	engine.RouterPrefix.engine = engine
	engine.RouterPrefix.router = engine.router
	engine.router.IgnoreCase = false
	engine.router.TrailingSlashRedirect = true
	engine.router.NoRoute = []HandlerFunc{noRoute}
	engine.router.NoMethod = []HandlerFunc{noMethod}
	return engine
}

// noRoute use as a default handler for router not matched
func noRoute(c *Context) {
	c.Status(http.StatusNotFound)
	c.String(default404Body)
}

// noMethod use as a default handler for Method not allowed
func noMethod(c *Context) {
	c.Status(http.StatusMethodNotAllowed)
	c.String(default405Body)
}

// Default return engine instance, add logger, recover handler to it.
func Default() *Engine {
	engine := New()
	engine.RouterPrefix.Use(Logger(), Recover())
	return engine
}

// NoRoute which is called when no matching route is found. If it is not set, noRoute is used.
func (engine *Engine) NoRoute(handlers ...HandlerFunc) {
	if len(handlers) == 0 {
		panic("there must be at least one handler")
	}

	engine.router.NoRoute = handlers
}

// NoMethod which is called when method is not registered. If it is not set, noMethod is used.
func (engine *Engine) NoMethod(handlers ...HandlerFunc) {
	if len(handlers) == 0 {
		panic("there must be at least one handler")
	}

	engine.router.NoMethod = handlers
}

// set IgnoreCase value
func (engine *Engine) SetIgnoreCase(ignoreCase bool) {
	engine.router.IgnoreCase = ignoreCase
}

// set TrailingSlashRedirect value
func (engine *Engine) SetTrailingSlashRedirect(redirect bool) {
	engine.router.TrailingSlashRedirect = redirect
}

// http.Handler interface
func (engine *Engine) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	// engine.router.ServeHTTP(rw, req)
	engine.RouterPrefix.handleRequest(rw, req)
}
<commit_msg>fix: remove server in Engine struct<commit_after>package looli

import (
	"net/http"
)

type (
	Engine struct {
		// router with basePath
		RouterPrefix

		// router used to match url
		router *Router

		// when set true, implements a best effort algorithm to return the real client IP, it parses
		// X-Real-IP and X-Forwarded-For in order to work properly with reverse-proxies such us: nginx or haproxy.
		ForwardedByClientIP bool
	}
	HandlerFunc func(*Context)
)

type Handler interface {
	Handle(*Context)
}

func New() *Engine {
	engine := &Engine{
		RouterPrefix: RouterPrefix{},
		router:       NewRouter(),
	}

	engine.RouterPrefix.engine = engine
	engine.RouterPrefix.router = engine.router
	engine.router.IgnoreCase = false
	engine.router.TrailingSlashRedirect = true
	engine.router.NoRoute = []HandlerFunc{noRoute}
	engine.router.NoMethod = []HandlerFunc{noMethod}
	return engine
}

// noRoute use as a default handler for router not matched
func noRoute(c *Context) {
	c.Status(http.StatusNotFound)
	c.String(default404Body)
}

// noMethod use as a default handler for Method not allowed
func noMethod(c *Context) {
	c.Status(http.StatusMethodNotAllowed)
	c.String(default405Body)
}

// Default return engine instance, add logger, recover handler to it.
func Default() *Engine {
	engine := New()
	engine.RouterPrefix.Use(Logger(), Recover())
	return engine
}

// NoRoute which is called when no matching route is found. If it is not set, noRoute is used.
func (engine *Engine) NoRoute(handlers ...HandlerFunc) {
	if len(handlers) == 0 {
		panic("there must be at least one handler")
	}

	engine.router.NoRoute = handlers
}

// NoMethod which is called when method is not registered. If it is not set, noMethod is used.
func (engine *Engine) NoMethod(handlers ...HandlerFunc) {
	if len(handlers) == 0 {
		panic("there must be at least one handler")
	}

	engine.router.NoMethod = handlers
}

// set IgnoreCase value
func (engine *Engine) SetIgnoreCase(ignoreCase bool) {
	engine.router.IgnoreCase = ignoreCase
}

// set TrailingSlashRedirect value
func (engine *Engine) SetTrailingSlashRedirect(redirect bool) {
	engine.router.TrailingSlashRedirect = redirect
}

// http.Handler interface
func (engine *Engine) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	// engine.router.ServeHTTP(rw, req)
	engine.RouterPrefix.handleRequest(rw, req)
}
<|endoftext|>


<commit_before>package commands

import (
	"fmt"
	"os"
	"os/exec"
	"strings"

	"github.com/leancloud/lean-cli/version"
	"github.com/urfave/cli"
)

// Run the command line
func Run(args []string) {
	app := cli.NewApp()
	app.Name = version.Distribution
	app.Version = version.Version
	app.Usage = fmt.Sprintf("Command line tool to manage and deploy %s apps", version.EngineBrandName)
	app.EnableBashCompletion = true

	app.CommandNotFound = thirdPartyCommand
	app.Commands = []cli.Command{
		{
			Name:   "login",
			Usage:  fmt.Sprintf("Log in to %s", version.BrandName),
			Action: wrapAction(loginAction),
			ArgsUsage: func() string {
				if version.Distribution == "lean" {
					return "[(-u <username>) (-p <password>) [--use-token] (--token <token>)]"
				} else {
					return "[(--token <token>)]"
				}
			}(),
			Flags: func() []cli.Flag {
				regionsString := []string{}

				for _, r := range version.AvailableRegions {
					regionsString = append(regionsString, r.String())
				}

				regions := strings.Join(regionsString, ", ")

				if version.Distribution == "lean" {
					return []cli.Flag{
						cli.StringFlag{
							Name:  "username,u",
							Usage: "Username",
						},
						cli.StringFlag{
							Name:  "password,p",
							Usage: "Password",
						},
						cli.StringFlag{
							Name:  "region,r",
							Usage: "The LeanCloud region to log in to (" + regions + ")",
						},
						cli.BoolFlag{
							Name:  "use-token",
							Usage: "Use AccessToken to log in (ask AccessToken)",
						},
						cli.StringFlag{
							Name:  "token",
							Usage: "AccessToken from LeanCloud Console => your App => LeanEngine => Deploy of your group => Deploy using CLI",
						},
					}
				} else {
					return []cli.Flag{
						cli.StringFlag{
							Name:  "region,r",
							Usage: "The TapTap Developer Services region to log in to (" + regions + ")",
						},
						cli.StringFlag{
							Name:  "token",
							Usage: "AccessToken from TapTap Developer Center => your Game => Game Services => Cloud Services => Cloud Engine => Deploy of your group => Deploy using CLI",
						},
					}
				}
			}(),
		},
		{
			Name:      "switch",
			Usage:     fmt.Sprintf("Change the associated %s app", version.EngineBrandName),
			Action:    wrapAction(switchAction),
			ArgsUsage: "[<appID> | <appName>]",
			Flags: func() []cli.Flag {
				return []cli.Flag{
					cli.StringFlag{
						Name:  "region",
						Usage: fmt.Sprintf("%s region", version.BrandName),
					},
					cli.StringFlag{
						Name:  "group",
						Usage: fmt.Sprintf("%s group", version.EngineBrandName),
					},
				}
			}(),
		},
		{
			Name:   "info",
			Usage:  "Show information about the associated user and app",
			Action: wrapAction(infoAction),
		},
		{
			Name:      "up",
			Usage:     "Start a development instance locally with debug console",
			ArgsUsage: "[(--port <port>) (--console-port <port>) (--cmd <cmd>)]",
			Action:    wrapAction(upAction),
			Flags: []cli.Flag{
				cli.IntFlag{
					Name:  "port,p",
					Usage: "Local port to listen on [default: 3000]",
					Value: 3000,
				},
				cli.IntFlag{
					Name:  "console-port,c",
					Usage: "Port of the debug console",
				},
				cli.StringFlag{
					Name:  "cmd",
					Usage: "Command to start the project, other arguments except --console-port are ignored",
				},
				cli.BoolFlag{
					Name:  "fetch-env",
					Usage: "Fetch environment variables from LeanEngine (secret variables not included)",
				},
			},
		},
		{
			Name:      "new",
			Usage:     fmt.Sprintf("Create a new %s project from official examples", version.EngineBrandName),
			Action:    wrapAction(newAction),
			ArgsUsage: "<path>",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "region",
					Usage: fmt.Sprintf("%s region", version.BrandName),
				},
				cli.StringFlag{
					Name:  "group",
					Usage: fmt.Sprintf("%s group", version.EngineBrandName),
				},
			},
		},
		{
			Name:      "deploy",
			Usage:     fmt.Sprintf("Deploy the project to %s", version.EngineBrandName),
			Action:    wrapAction(deployAction),
			ArgsUsage: "(--prod | --staging) [--no-cache --build-logs --overwrite-functions]",
			Flags: []cli.Flag{
				cli.BoolFlag{
					Name:  "prod",
					Usage: "Deploy to production environment",
				},
				cli.BoolFlag{
					Name:  "staging",
					Usage: "Deploy to staging environment",
				},
				cli.BoolFlag{
					Name:  "build-logs",
					Usage: "Print build logs",
				},
				cli.BoolFlag{
					Name:  "g",
					Usage: "Deploy from git repo",
				},
				cli.BoolFlag{
					Name:  "war",
					Usage: "Deploy .war file for Java project. The first .war file in target/ is used by default",
				},
				cli.BoolFlag{
					Name:  "no-cache",
					Usage: "Disable buliding cache",
				},
				cli.BoolFlag{
					Name:  "overwrite-functions",
					Usage: "Overwrite cloud functions with the same name in other groups",
				},
				cli.StringFlag{
					Name:  "leanignore",
					Usage: "Rule file for ignored files in deployment",
					Value: ".leanignore",
				},
				cli.StringFlag{
					Name:  "message,m",
					Usage: "Comment for this version, only applicable when deploying from local files",
				},
				cli.BoolFlag{
					Name: "keep-deploy-file",
				},
				cli.StringFlag{
					Name:  "revision,r",
					Usage: "Git revision or branch. Only applicable when deploying from Git",
					Value: "master",
				},
				cli.StringFlag{
					Name:  "options",
					Usage: "Send additional deploy options to server, in urlencode format(like `--options build-root=app`)",
				},
				cli.BoolFlag{
					Name:  "direct",
					Usage: "Upload project's tarball to remote directly",
				},
			},
		},
		{
			Name:      "publish",
			Usage:     "Publish the version of staging to production",
			Action:    wrapAction(publishAction),
			ArgsUsage: "[--overwrite-functions]",
			Flags: []cli.Flag{
				cli.BoolFlag{
					Name:  "overwrite-functions",
					Usage: "Overwrite cloud functions with the same name in other groups",
				},
				cli.StringFlag{
					Name:  "options",
					Usage: "Send additional deploy options to server, in urlencode format(like `--options build-root=app`)",
				},
			},
		},
		{
			Name:      "db",
			Usage:     fmt.Sprintf("Access to to %s instances", version.DBBrandName),
			Action:    wrapAction(dbListAction),
			ArgsUsage: "(list | proxy | shell | exec)",
			Subcommands: []cli.Command{
				{
					Name:   "list",
					Usage:  fmt.Sprintf("List %s instances of current app (include share instances)", version.DBBrandName),
					Action: wrapAction(dbListAction),
				},
				{
					Name:      "proxy",
					Usage:     fmt.Sprintf("Proxy %s instance to local port", version.DBBrandName),
					Action:    wrapAction(dbProxyAction),
					ArgsUsage: "<instance-name>",
					Flags: []cli.Flag{
						cli.StringFlag{
							Name:  "app-id",
							Usage: "Specify appId for same name share instance",
						},
						cli.IntFlag{
							Name:  "port, p",
							Usage: "Specify local proxy port [default: 5678]",
							Value: 5678,
						},
					},
				},
				{
					Name:      "shell",
					Usage:     fmt.Sprintf("Enter interactive shell to %s instance", version.DBBrandName),
					Action:    wrapAction(dbShellAction),
					ArgsUsage: "<instance-name>",
					Flags: []cli.Flag{
						cli.StringFlag{
							Name:  "app-id",
							Usage: "Specify appId for same name share instance",
						},
						cli.IntFlag{
							Name:  "port, p",
							Usage: "Specify local proxy port [default: 5678]",
							Value: 5678,
						},
					},
				},
				{
					Name:      "exec",
					Usage:     fmt.Sprintf("Exce commands on %s instance", version.DBBrandName),
					Action:    wrapAction(dbExecAction),
					ArgsUsage: "<instance-name> <db-commands>...",
					Flags: []cli.Flag{
						cli.StringFlag{
							Name:  "app-id",
							Usage: "Specify appId for same name share instance",
						},
						cli.IntFlag{
							Name:  "port, p",
							Usage: "Specify local proxy port [default: 5678]",
							Value: 5678,
						},
					},
				},
			},
		},
		{
			Name:      "file",
			Usage:     "Manage files ('_File' class in Data Storage)",
			ArgsUsage: "(upload)",
			Subcommands: []cli.Command{
				{
					Name:      "upload",
					Usage:     "Upload files",
					Action:    uploadAction,
					ArgsUsage: "<file-path>...",
				},
			},
		},
		{
			Name:      "logs",
			Usage:     fmt.Sprintf("Show %s logs", version.EngineBrandName),
			Action:    wrapAction(logsAction),
			ArgsUsage: "[-f (--env staging|production)]",
			Flags: []cli.Flag{
				cli.BoolFlag{
					Name:  "f",
					Usage: "Wait for and continuously show most recent logs",
				},
				cli.StringFlag{
					Name:  "env,e",
					Usage: "Environment to view (staging / production)",
					Value: "production",
				},
				cli.IntFlag{
					Name:  "limit,l",
					Usage: "Maximum number of lines to show [default: 30]",
					Value: 30,
				},
				cli.StringFlag{
					Name:  "from",
					Usage: "Start date formatted as YYYY-MM-DD (local time) or RFC3339, e.g., 2006-01-02 or 2006-01-02T15:04:05+08:00",
				},
				cli.StringFlag{
					Name:  "to",
					Usage: "End date formated as YYYY-MM-DD (local time) or RFC3339, e.g., 2006-01-02 or 2006-01-02T15:04:05+08:00",
				},
				cli.StringFlag{
					Name:  "format",
					Usage: "Format to use ('default' or 'json')",
					Value: "default",
				},
			},
		},
		{
			Name:      "debug",
			Usage:     "Start the debug console without running the project",
			Action:    wrapAction(debugAction),
			ArgsUsage: "[(--remote <url>) (--port <port>)]",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "remote,r",
					Usage: "URL of target app running on",
					Value: "http://localhost:3000",
				},
				cli.StringFlag{
					Name:  "app-id",
					Usage: "Target App ID, use the App ID of the current project by default",
				},
				cli.IntFlag{
					Name:  "port,p",
					Usage: "Port to listen on [default: 3001]",
					Value: 3001,
				},
			},
		},
		{
			Name:      "env",
			Usage:     fmt.Sprintf("Print custom environment variables on %s (secret variables not included)", version.EngineBrandName),
			Action:    wrapAction(envAction),
			ArgsUsage: "[(set | unset)]",
			Flags: []cli.Flag{
				cli.IntFlag{
					Name:  "port,p",
					Usage: "Local port for the app (affects value of LEANCLOUD_APP_PORT) [default: 3000]",
					Value: 3000,
				},
				cli.StringFlag{
					Name:  "template",
					Usage: "Template for output [default: export {{{name}}}={{{value}}}]",
				},
			},
			Subcommands: []cli.Command{
				{
					Name:      "set",
					Usage:     "Set the value of an environment variable",
					Action:    wrapAction(envSetAction),
					ArgsUsage: "<name> <value>",
				},
				{
					Name:      "unset",
					Usage:     "Delete an environment variable",
					Action:    wrapAction(envUnsetAction),
					ArgsUsage: "<name>",
				},
			},
		},
		{
			Name:      "cql",
			Usage:     "Enter CQL interactive shell (warn: CQL is deprecated)",
			Action:    wrapAction(cqlAction),
			ArgsUsage: "[(--format json) (--eval <cql>)]",
			Flags: []cli.Flag{
				cli.StringFlag{Name: "format,f",
					Usage: "CQL result format can be 'table' or 'json' [default: table]",
					Value: "table",
				},
				cli.StringFlag{
					Name:  "eval",
					Usage: "CQL command to run",
				},
			},
		},
		{
			Name:  "help",
			Usage: "Show usages of all subcommands",
			Action: func(c *cli.Context) error {
				args := c.Args()
				if args.Present() {
					_, err := fmt.Printf("Please use `%s %s --help` for subcommand usage.\n", os.Args[0], args.First())
					return err
				}
				return cli.ShowAppHelp(c)
			},
		},
	}

	app.Before = func(c *cli.Context) error {
		disableGA, ok := os.LookupEnv("NO_ANALYTICS")
		if !ok || disableGA == "false" {
			args := []string{"--_collect-stats"}
			args = append(args, c.Args()...)
			StartBackgroundCommand(exec.Command(os.Args[0], args...))
		}
		return nil
	}

	app.Run(args)
}
<commit_msg>fix: wrap uploadAction<commit_after>package commands

import (
	"fmt"
	"os"
	"os/exec"
	"strings"

	"github.com/leancloud/lean-cli/version"
	"github.com/urfave/cli"
)

// Run the command line
func Run(args []string) {
	app := cli.NewApp()
	app.Name = version.Distribution
	app.Version = version.Version
	app.Usage = fmt.Sprintf("Command line tool to manage and deploy %s apps", version.EngineBrandName)
	app.EnableBashCompletion = true

	app.CommandNotFound = thirdPartyCommand
	app.Commands = []cli.Command{
		{
			Name:   "login",
			Usage:  fmt.Sprintf("Log in to %s", version.BrandName),
			Action: wrapAction(loginAction),
			ArgsUsage: func() string {
				if version.Distribution == "lean" {
					return "[(-u <username>) (-p <password>) [--use-token] (--token <token>)]"
				} else {
					return "[(--token <token>)]"
				}
			}(),
			Flags: func() []cli.Flag {
				regionsString := []string{}

				for _, r := range version.AvailableRegions {
					regionsString = append(regionsString, r.String())
				}

				regions := strings.Join(regionsString, ", ")

				if version.Distribution == "lean" {
					return []cli.Flag{
						cli.StringFlag{
							Name:  "username,u",
							Usage: "Username",
						},
						cli.StringFlag{
							Name:  "password,p",
							Usage: "Password",
						},
						cli.StringFlag{
							Name:  "region,r",
							Usage: "The LeanCloud region to log in to (" + regions + ")",
						},
						cli.BoolFlag{
							Name:  "use-token",
							Usage: "Use AccessToken to log in (ask AccessToken)",
						},
						cli.StringFlag{
							Name:  "token",
							Usage: "AccessToken from LeanCloud Console => your App => LeanEngine => Deploy of your group => Deploy using CLI",
						},
					}
				} else {
					return []cli.Flag{
						cli.StringFlag{
							Name:  "region,r",
							Usage: "The TapTap Developer Services region to log in to (" + regions + ")",
						},
						cli.StringFlag{
							Name:  "token",
							Usage: "AccessToken from TapTap Developer Center => your Game => Game Services => Cloud Services => Cloud Engine => Deploy of your group => Deploy using CLI",
						},
					}
				}
			}(),
		},
		{
			Name:      "switch",
			Usage:     fmt.Sprintf("Change the associated %s app", version.EngineBrandName),
			Action:    wrapAction(switchAction),
			ArgsUsage: "[<appID> | <appName>]",
			Flags: func() []cli.Flag {
				return []cli.Flag{
					cli.StringFlag{
						Name:  "region",
						Usage: fmt.Sprintf("%s region", version.BrandName),
					},
					cli.StringFlag{
						Name:  "group",
						Usage: fmt.Sprintf("%s group", version.EngineBrandName),
					},
				}
			}(),
		},
		{
			Name:   "info",
			Usage:  "Show information about the associated user and app",
			Action: wrapAction(infoAction),
		},
		{
			Name:      "up",
			Usage:     "Start a development instance locally with debug console",
			ArgsUsage: "[(--port <port>) (--console-port <port>) (--cmd <cmd>)]",
			Action:    wrapAction(upAction),
			Flags: []cli.Flag{
				cli.IntFlag{
					Name:  "port,p",
					Usage: "Local port to listen on [default: 3000]",
					Value: 3000,
				},
				cli.IntFlag{
					Name:  "console-port,c",
					Usage: "Port of the debug console",
				},
				cli.StringFlag{
					Name:  "cmd",
					Usage: "Command to start the project, other arguments except --console-port are ignored",
				},
				cli.BoolFlag{
					Name:  "fetch-env",
					Usage: "Fetch environment variables from LeanEngine (secret variables not included)",
				},
			},
		},
		{
			Name:      "new",
			Usage:     fmt.Sprintf("Create a new %s project from official examples", version.EngineBrandName),
			Action:    wrapAction(newAction),
			ArgsUsage: "<path>",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "region",
					Usage: fmt.Sprintf("%s region", version.BrandName),
				},
				cli.StringFlag{
					Name:  "group",
					Usage: fmt.Sprintf("%s group", version.EngineBrandName),
				},
			},
		},
		{
			Name:      "deploy",
			Usage:     fmt.Sprintf("Deploy the project to %s", version.EngineBrandName),
			Action:    wrapAction(deployAction),
			ArgsUsage: "(--prod | --staging) [--no-cache --build-logs --overwrite-functions]",
			Flags: []cli.Flag{
				cli.BoolFlag{
					Name:  "prod",
					Usage: "Deploy to production environment",
				},
				cli.BoolFlag{
					Name:  "staging",
					Usage: "Deploy to staging environment",
				},
				cli.BoolFlag{
					Name:  "build-logs",
					Usage: "Print build logs",
				},
				cli.BoolFlag{
					Name:  "g",
					Usage: "Deploy from git repo",
				},
				cli.BoolFlag{
					Name:  "war",
					Usage: "Deploy .war file for Java project. The first .war file in target/ is used by default",
				},
				cli.BoolFlag{
					Name:  "no-cache",
					Usage: "Disable buliding cache",
				},
				cli.BoolFlag{
					Name:  "overwrite-functions",
					Usage: "Overwrite cloud functions with the same name in other groups",
				},
				cli.StringFlag{
					Name:  "leanignore",
					Usage: "Rule file for ignored files in deployment",
					Value: ".leanignore",
				},
				cli.StringFlag{
					Name:  "message,m",
					Usage: "Comment for this version, only applicable when deploying from local files",
				},
				cli.BoolFlag{
					Name: "keep-deploy-file",
				},
				cli.StringFlag{
					Name:  "revision,r",
					Usage: "Git revision or branch. Only applicable when deploying from Git",
					Value: "master",
				},
				cli.StringFlag{
					Name:  "options",
					Usage: "Send additional deploy options to server, in urlencode format(like `--options build-root=app`)",
				},
				cli.BoolFlag{
					Name:  "direct",
					Usage: "Upload project's tarball to remote directly",
				},
			},
		},
		{
			Name:      "publish",
			Usage:     "Publish the version of staging to production",
			Action:    wrapAction(publishAction),
			ArgsUsage: "[--overwrite-functions]",
			Flags: []cli.Flag{
				cli.BoolFlag{
					Name:  "overwrite-functions",
					Usage: "Overwrite cloud functions with the same name in other groups",
				},
				cli.StringFlag{
					Name:  "options",
					Usage: "Send additional deploy options to server, in urlencode format(like `--options build-root=app`)",
				},
			},
		},
		{
			Name:      "db",
			Usage:     fmt.Sprintf("Access to to %s instances", version.DBBrandName),
			Action:    wrapAction(dbListAction),
			ArgsUsage: "(list | proxy | shell | exec)",
			Subcommands: []cli.Command{
				{
					Name:   "list",
					Usage:  fmt.Sprintf("List %s instances of current app (include share instances)", version.DBBrandName),
					Action: wrapAction(dbListAction),
				},
				{
					Name:      "proxy",
					Usage:     fmt.Sprintf("Proxy %s instance to local port", version.DBBrandName),
					Action:    wrapAction(dbProxyAction),
					ArgsUsage: "<instance-name>",
					Flags: []cli.Flag{
						cli.StringFlag{
							Name:  "app-id",
							Usage: "Specify appId for same name share instance",
						},
						cli.IntFlag{
							Name:  "port, p",
							Usage: "Specify local proxy port [default: 5678]",
							Value: 5678,
						},
					},
				},
				{
					Name:      "shell",
					Usage:     fmt.Sprintf("Enter interactive shell to %s instance", version.DBBrandName),
					Action:    wrapAction(dbShellAction),
					ArgsUsage: "<instance-name>",
					Flags: []cli.Flag{
						cli.StringFlag{
							Name:  "app-id",
							Usage: "Specify appId for same name share instance",
						},
						cli.IntFlag{
							Name:  "port, p",
							Usage: "Specify local proxy port [default: 5678]",
							Value: 5678,
						},
					},
				},
				{
					Name:      "exec",
					Usage:     fmt.Sprintf("Exce commands on %s instance", version.DBBrandName),
					Action:    wrapAction(dbExecAction),
					ArgsUsage: "<instance-name> <db-commands>...",
					Flags: []cli.Flag{
						cli.StringFlag{
							Name:  "app-id",
							Usage: "Specify appId for same name share instance",
						},
						cli.IntFlag{
							Name:  "port, p",
							Usage: "Specify local proxy port [default: 5678]",
							Value: 5678,
						},
					},
				},
			},
		},
		{
			Name:      "file",
			Usage:     "Manage files ('_File' class in Data Storage)",
			ArgsUsage: "(upload)",
			Subcommands: []cli.Command{
				{
					Name:      "upload",
					Usage:     "Upload files",
					Action:    wrapAction(uploadAction),
					ArgsUsage: "<file-path>...",
				},
			},
		},
		{
			Name:      "logs",
			Usage:     fmt.Sprintf("Show %s logs", version.EngineBrandName),
			Action:    wrapAction(logsAction),
			ArgsUsage: "[-f (--env staging|production)]",
			Flags: []cli.Flag{
				cli.BoolFlag{
					Name:  "f",
					Usage: "Wait for and continuously show most recent logs",
				},
				cli.StringFlag{
					Name:  "env,e",
					Usage: "Environment to view (staging / production)",
					Value: "production",
				},
				cli.IntFlag{
					Name:  "limit,l",
					Usage: "Maximum number of lines to show [default: 30]",
					Value: 30,
				},
				cli.StringFlag{
					Name:  "from",
					Usage: "Start date formatted as YYYY-MM-DD (local time) or RFC3339, e.g., 2006-01-02 or 2006-01-02T15:04:05+08:00",
				},
				cli.StringFlag{
					Name:  "to",
					Usage: "End date formated as YYYY-MM-DD (local time) or RFC3339, e.g., 2006-01-02 or 2006-01-02T15:04:05+08:00",
				},
				cli.StringFlag{
					Name:  "format",
					Usage: "Format to use ('default' or 'json')",
					Value: "default",
				},
			},
		},
		{
			Name:      "debug",
			Usage:     "Start the debug console without running the project",
			Action:    wrapAction(debugAction),
			ArgsUsage: "[(--remote <url>) (--port <port>)]",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "remote,r",
					Usage: "URL of target app running on",
					Value: "http://localhost:3000",
				},
				cli.StringFlag{
					Name:  "app-id",
					Usage: "Target App ID, use the App ID of the current project by default",
				},
				cli.IntFlag{
					Name:  "port,p",
					Usage: "Port to listen on [default: 3001]",
					Value: 3001,
				},
			},
		},
		{
			Name:      "env",
			Usage:     fmt.Sprintf("Print custom environment variables on %s (secret variables not included)", version.EngineBrandName),
			Action:    wrapAction(envAction),
			ArgsUsage: "[(set | unset)]",
			Flags: []cli.Flag{
				cli.IntFlag{
					Name:  "port,p",
					Usage: "Local port for the app (affects value of LEANCLOUD_APP_PORT) [default: 3000]",
					Value: 3000,
				},
				cli.StringFlag{
					Name:  "template",
					Usage: "Template for output [default: export {{{name}}}={{{value}}}]",
				},
			},
			Subcommands: []cli.Command{
				{
					Name:      "set",
					Usage:     "Set the value of an environment variable",
					Action:    wrapAction(envSetAction),
					ArgsUsage: "<name> <value>",
				},
				{
					Name:      "unset",
					Usage:     "Delete an environment variable",
					Action:    wrapAction(envUnsetAction),
					ArgsUsage: "<name>",
				},
			},
		},
		{
			Name:      "cql",
			Usage:     "Enter CQL interactive shell (warn: CQL is deprecated)",
			Action:    wrapAction(cqlAction),
			ArgsUsage: "[(--format json) (--eval <cql>)]",
			Flags: []cli.Flag{
				cli.StringFlag{Name: "format,f",
					Usage: "CQL result format can be 'table' or 'json' [default: table]",
					Value: "table",
				},
				cli.StringFlag{
					Name:  "eval",
					Usage: "CQL command to run",
				},
			},
		},
		{
			Name:  "help",
			Usage: "Show usages of all subcommands",
			Action: func(c *cli.Context) error {
				args := c.Args()
				if args.Present() {
					_, err := fmt.Printf("Please use `%s %s --help` for subcommand usage.\n", os.Args[0], args.First())
					return err
				}
				return cli.ShowAppHelp(c)
			},
		},
	}

	app.Before = func(c *cli.Context) error {
		disableGA, ok := os.LookupEnv("NO_ANALYTICS")
		if !ok || disableGA == "false" {
			args := []string{"--_collect-stats"}
			args = append(args, c.Args()...)
			StartBackgroundCommand(exec.Command(os.Args[0], args...))
		}
		return nil
	}

	app.Run(args)
}
<|endoftext|>


<commit_before>// Container provides a Reader and Writer which serialize and deserialize gogen-avro structs to the Avro Object Container File (OCF) format. These are low-level primitives you probably don't need to use directly.
package container

import (
	"bytes"
	"compress/flate"
	"io"

	"github.com/actgardner/gogen-avro/container/avro"
)

// A Codec specifies how the blocks within a container file should be compressed.
type Codec string

const (
	// No compression
	Null Codec = "null"

	// Deflate compression
	Deflate Codec = "deflate"

	// Snappy compression
	Snappy Codec = "snappy"
)

type CloseableResettableWriter interface {
	Close() error
	Reset(io.Writer)
}

// Writer wraps an io.Writer and writes the file and block-level framing required for an OCF file.
// You can create a Writer for a given struct by calling the generated method `New<RecordType>Writer`.
type Writer struct {
	writer           io.Writer
	syncMarker       [16]byte
	codec            Codec
	recordsPerBlock  int64
	blockBuffer      *bytes.Buffer
	compressedWriter io.Writer
	nextBlockRecords int64
}

//  Create a new Writer wrapping the provided io.Writer with the given Codec and number of records per block.
//  The Writer will lazily write the container file header when WriteRecord is called the first time.
//  You must call Flush on the Writer before closing the underlying io.Writer, to ensure the final block is written.
//  A schema string must be passed to ensure that a correct header is written even if no records are written. This
//  is required to produce valid empty Avro container files.
func NewWriter(writer io.Writer, codec Codec, recordsPerBlock int64, schema string) (*Writer, error) {
	blockBytes := make([]byte, 0)
	blockBuffer := bytes.NewBuffer(blockBytes)

	avroWriter := &Writer{
		writer:          writer,
		syncMarker:      [16]byte{'g', 'o', 'g', 'e', 'n', 'a', 'v', 'r', 'o', 'm', 'a', 'g', 'i', 'c', '1', '0'},
		codec:           codec,
		recordsPerBlock: recordsPerBlock,
		blockBuffer:     blockBuffer,
	}
	var err error
	if codec == Deflate {
		avroWriter.compressedWriter, err = flate.NewWriter(avroWriter.blockBuffer, flate.DefaultCompression)
		if err != nil {
			return nil, err
		}
	} else if codec == Snappy {
		avroWriter.compressedWriter = newSnappyWriter(avroWriter.blockBuffer)
	} else {
		avroWriter.compressedWriter = avroWriter.blockBuffer
	}

	err = avroWriter.writeHeader(schema)
	if err != nil {
		return nil, err
	}

	return avroWriter, nil
}

func (avroWriter *Writer) writeHeader(schema string) error {
	header := &avro.AvroContainerHeader{
		Magic: [4]byte{'O', 'b', 'j', 1},
		Meta: &avro.MapBytes{M: map[string][]byte{
			"avro.schema": []byte(schema),
			"avro.codec":  []byte(avroWriter.codec),
		}},
		Sync: avroWriter.syncMarker,
	}
	return header.Serialize(avroWriter.writer)
}

//  Write an AvroRecord to the container file. All gogen-avro generated structs
//  fulfill the AvroRecord interface. Note that all records in a given container file
//  must be of the same Avro type.
func (avroWriter *Writer) WriteRecord(record AvroRecord) error {
	var err error
	// Serialize the new record into the compressed writer
	err = record.Serialize(avroWriter.compressedWriter)
	if err != nil {
		return err
	}
	avroWriter.nextBlockRecords += 1

	// If the block if full, flush and reset the compressed writer,
	// write the header and the block contents
	if avroWriter.nextBlockRecords >= avroWriter.recordsPerBlock {
		return avroWriter.Flush()
	}

	return nil
}

//  Write the current block to the file if it has been filled.  It is
//  best-practise to always call this before the underlying io.Writer is closed.
func (avroWriter *Writer) Flush() error {
	if avroWriter.nextBlockRecords == 0 {
		return nil
	}

	// Write out all of the buffered records as a new block
	// Must be called before closing to ensure the last block is written
	if fwWriter, ok := avroWriter.compressedWriter.(CloseableResettableWriter); ok {
		fwWriter.Close()
		fwWriter.Reset(avroWriter.blockBuffer)
	}

	if avroWriter.nextBlockRecords > 0 {
		block := &avro.AvroContainerBlock{
			NumRecords:  avroWriter.nextBlockRecords,
			RecordBytes: avroWriter.blockBuffer.Bytes(),
			Sync:        avroWriter.syncMarker,
		}
		err := block.Serialize(avroWriter.writer)
		if err != nil {
			return err
		}
	}

	avroWriter.blockBuffer.Reset()
	avroWriter.nextBlockRecords = 0

	return nil
}

// Get the current block buffer size.
// caller might trigger an early Flush if current block size gets huge.
func (avroWriter *Writer) BlockBufferSize() int {
	return avroWriter.blockBuffer.Len()
}
<commit_msg>OCF writer lazily writes header on first write instead of file open<commit_after>// Container provides a Reader and Writer which serialize and deserialize gogen-avro structs to the Avro Object Container File (OCF) format. These are low-level primitives you probably don't need to use directly.
package container

import (
	"bytes"
	"compress/flate"
	"io"

	"github.com/actgardner/gogen-avro/container/avro"
)

// A Codec specifies how the blocks within a container file should be compressed.
type Codec string

const (
	// No compression
	Null Codec = "null"

	// Deflate compression
	Deflate Codec = "deflate"

	// Snappy compression
	Snappy Codec = "snappy"
)

type CloseableResettableWriter interface {
	Close() error
	Reset(io.Writer)
}

// Writer wraps an io.Writer and writes the file and block-level framing required for an OCF file.
// You can create a Writer for a given struct by calling the generated method `New<RecordType>Writer`.
type Writer struct {
	writer           io.Writer
	syncMarker       [16]byte
	codec            Codec
	recordsPerBlock  int64
	blockBuffer      *bytes.Buffer
	compressedWriter io.Writer
	nextBlockRecords int64
	schema           string
	wroteHeader      bool
}

//  Create a new Writer wrapping the provided io.Writer with the given Codec and number of records per block.
//  The Writer will lazily write the container file header when WriteRecord is called the first time.
//  You must call Flush on the Writer before closing the underlying io.Writer, to ensure the final block is written.
//  A schema string must be passed to ensure that a correct header is written even if no records are written. This
//  is required to produce valid empty Avro container files.
func NewWriter(writer io.Writer, codec Codec, recordsPerBlock int64, schema string) (*Writer, error) {
	blockBytes := make([]byte, 0)
	blockBuffer := bytes.NewBuffer(blockBytes)

	avroWriter := &Writer{
		writer:          writer,
		syncMarker:      [16]byte{'g', 'o', 'g', 'e', 'n', 'a', 'v', 'r', 'o', 'm', 'a', 'g', 'i', 'c', '1', '0'},
		codec:           codec,
		recordsPerBlock: recordsPerBlock,
		blockBuffer:     blockBuffer,
		schema:          schema,
	}
	var err error
	if codec == Deflate {
		avroWriter.compressedWriter, err = flate.NewWriter(avroWriter.blockBuffer, flate.DefaultCompression)
		if err != nil {
			return nil, err
		}
	} else if codec == Snappy {
		avroWriter.compressedWriter = newSnappyWriter(avroWriter.blockBuffer)
	} else {
		avroWriter.compressedWriter = avroWriter.blockBuffer
	}

	return avroWriter, nil
}

func (avroWriter *Writer) writeHeader(schema string) error {
	header := &avro.AvroContainerHeader{
		Magic: [4]byte{'O', 'b', 'j', 1},
		Meta: &avro.MapBytes{M: map[string][]byte{
			"avro.schema": []byte(schema),
			"avro.codec":  []byte(avroWriter.codec),
		}},
		Sync: avroWriter.syncMarker,
	}
	return header.Serialize(avroWriter.writer)
}

//  Write an AvroRecord to the container file. All gogen-avro generated structs
//  fulfill the AvroRecord interface. Note that all records in a given container file
//  must be of the same Avro type.
func (avroWriter *Writer) WriteRecord(record AvroRecord) error {
	var err error

	// Lazily write the header if it hasn't been written yet
	if !avroWriter.wroteHeader {
		err = avroWriter.writeHeader(avroWriter.schema)
		if err != nil {
			return err
		}
		avroWriter.wroteHeader = true
	}

	// Serialize the new record into the compressed writer
	err = record.Serialize(avroWriter.compressedWriter)
	if err != nil {
		return err
	}
	avroWriter.nextBlockRecords += 1

	// If the block if full, flush and reset the compressed writer,
	// write the header and the block contents
	if avroWriter.nextBlockRecords >= avroWriter.recordsPerBlock {
		return avroWriter.Flush()
	}

	return nil
}

//  Write the current block to the file if it has been filled.  It is
//  best-practise to always call this before the underlying io.Writer is closed.
func (avroWriter *Writer) Flush() error {
	if avroWriter.nextBlockRecords == 0 {
		return nil
	}

	// Write out all of the buffered records as a new block
	// Must be called before closing to ensure the last block is written
	if fwWriter, ok := avroWriter.compressedWriter.(CloseableResettableWriter); ok {
		fwWriter.Close()
		fwWriter.Reset(avroWriter.blockBuffer)
	}

	if avroWriter.nextBlockRecords > 0 {
		block := &avro.AvroContainerBlock{
			NumRecords:  avroWriter.nextBlockRecords,
			RecordBytes: avroWriter.blockBuffer.Bytes(),
			Sync:        avroWriter.syncMarker,
		}
		err := block.Serialize(avroWriter.writer)
		if err != nil {
			return err
		}
	}

	avroWriter.blockBuffer.Reset()
	avroWriter.nextBlockRecords = 0

	return nil
}

// Get the current block buffer size.
// caller might trigger an early Flush if current block size gets huge.
func (avroWriter *Writer) BlockBufferSize() int {
	return avroWriter.blockBuffer.Len()
}
<|endoftext|>


<commit_before>package controller

import (
	"fmt"
	"net/http"
)

// Error404 handles 404 - Page Not Found
func Error404(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(http.StatusNotFound)
	fmt.Fprint(w, "Not Found 404")
}

// Error500 handles 500 - Internal Server Error
func Error500(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(http.StatusInternalServerError)
	fmt.Fprint(w, "Internal Server Error 500")
}

// InvalidToken handles CSRF attacks
func InvalidToken(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/html")
	w.WriteHeader(http.StatusForbidden)
	fmt.Fprint(w, `Your token <strong>expired</strong>, click <a href="javascript:void(0)" onclick="window.history.back()">here</a> to try again.`)
}
<commit_msg>Updated token invalid page to use refresh code<commit_after>package controller

import (
	"fmt"
	"net/http"
)

// Error404 handles 404 - Page Not Found
func Error404(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(http.StatusNotFound)
	fmt.Fprint(w, "Not Found 404")
}

// Error500 handles 500 - Internal Server Error
func Error500(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(http.StatusInternalServerError)
	fmt.Fprint(w, "Internal Server Error 500")
}

// InvalidToken handles CSRF attacks
func InvalidToken(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/html")
	w.WriteHeader(http.StatusForbidden)
	fmt.Fprint(w, `Your token <strong>expired</strong>, click <a href="javascript:void(0)" onclick="location.replace(document.referrer)">here</a> to try again.`)
}
<|endoftext|>


<commit_before>package coremain

// Various CoreDNS constants.
const (
	CoreVersion = "1.6.5"
	coreName    = "CoreDNS"
	serverType  = "dns"
)
<commit_msg>Bump version to 1.6.6 (#3520)<commit_after>package coremain

// Various CoreDNS constants.
const (
	CoreVersion = "1.6.6"
	coreName    = "CoreDNS"
	serverType  = "dns"
)
<|endoftext|>


<commit_before>package imexport

import (
	"encoding/base32"
	"fmt"
	"io"
	"net/http"
	"os"

	"github.com/cozy/cozy-stack/pkg/crypto"
	"github.com/cozy/cozy-stack/pkg/imexport"
	"github.com/cozy/cozy-stack/pkg/jobs"
	"github.com/cozy/cozy-stack/pkg/vfs"
	workers "github.com/cozy/cozy-stack/pkg/workers/mails"
	"github.com/cozy/cozy-stack/web/middlewares"
	"github.com/labstack/echo"
)

func export(c echo.Context) error {
	instance := middlewares.GetInstance(c)
	domain := instance.Domain

	tab := crypto.GenerateRandomBytes(20)
	id := base32.StdEncoding.EncodeToString(tab)

	w, err := os.Create(fmt.Sprintf("%s-%s.tar.gz", domain, id))
	if err != nil {
		return err
	}
	defer w.Close()

	err = imexport.Tardir(w, instance)
	if err != nil {
		return err
	}

	var objet string
	lien := fmt.Sprintf("http://%s%s%s-%s", domain, c.Path(), domain, id)

	if instance.Locale == "en" {
		objet = "The archive with all your Cozy data is ready"
	} else if instance.Locale == "fr" {
		objet = "L'archive contenant toutes les données de Cozy est prête"
	}

	msg, err := jobs.NewMessage("json", workers.Options{
		Mode:         workers.ModeNoReply,
		Subject:      objet,
		TemplateName: "archiver_" + instance.Locale,
		TemplateValues: map[string]string{
			"RecipientName": domain,
			"Lien":          lien,
		},
	})
	if err != nil {
		return err
	}

	context := jobs.NewWorkerContext(instance.Domain, "abcd")
	err = workers.SendMail(context, msg)
	if err != nil {
		return err
	}

	return c.JSON(http.StatusOK, echo.Map{
		"message": "bienvenue sur la super page",
	})
}
func exportDir(c echo.Context) error {
	domID := c.Param("domain-id")
	fmt.Println(domID)

	src, err := os.Open(fmt.Sprintf("%s.tar.gz", domID))
	if err != nil {
		return err
	}
	dst, err := os.Create("cozy.tar.gz")
	if err != nil {
		return err
	}

	_, err = io.Copy(dst, src)
	if err != nil {
		return err
	}

	err = os.Remove(fmt.Sprintf("%s.tar.gz", domID))
	if err != nil {
		return err
	}

	return c.JSON(http.StatusOK, echo.Map{
		"message": "bienvenue sur la super page bis",
	})
}

func importer(c echo.Context) error {
	instance := middlewares.GetInstance(c)
	fs := instance.VFS()

	r, err := os.Open("cozy.tar.gz")
	if err != nil {
		return err
	}

	rep := c.Param("destination")
	rep = fmt.Sprintf("/%s", rep)

	exist, err := vfs.DirExists(fs, rep)
	if err != nil {
		return err
	}
	var dst *vfs.DirDoc
	if !exist {
		dst, err = vfs.Mkdir(fs, rep, nil)
		if err != nil {
			return err
		}
	} else {
		dst, err = fs.DirByPath(rep)
		if err != nil {
			return err
		}
	}

	err = imexport.Untardir(r, dst.ID(), instance)
	if err != nil {
		return err
	}

	return c.JSON(http.StatusOK, echo.Map{
		"message": "bienvenue sur la super page cozy",
	})
}

// Routes sets the routing for export
func Routes(router *echo.Group) {
	router.GET("/export/", export)
	router.HEAD("/export/", export)

	router.GET("/import/:destination", importer)
	router.HEAD("/import/", importer)

	router.GET("/export/:domain-id", exportDir)

}
<commit_msg>Remove some dead code<commit_after>package imexport

import (
	"encoding/base32"
	"fmt"
	"net/http"
	"os"

	"github.com/cozy/cozy-stack/pkg/crypto"
	"github.com/cozy/cozy-stack/pkg/imexport"
	"github.com/cozy/cozy-stack/pkg/jobs"
	"github.com/cozy/cozy-stack/pkg/vfs"
	workers "github.com/cozy/cozy-stack/pkg/workers/mails"
	"github.com/cozy/cozy-stack/web/middlewares"
	"github.com/labstack/echo"
)

func export(c echo.Context) error {
	instance := middlewares.GetInstance(c)
	domain := instance.Domain

	tab := crypto.GenerateRandomBytes(20)
	id := base32.StdEncoding.EncodeToString(tab)

	w, err := os.Create(fmt.Sprintf("%s-%s.tar.gz", domain, id))
	if err != nil {
		return err
	}
	defer w.Close()

	err = imexport.Tardir(w, instance)
	if err != nil {
		return err
	}

	var objet string
	lien := fmt.Sprintf("http://%s%s%s-%s", domain, c.Path(), domain, id)

	if instance.Locale == "en" {
		objet = "The archive with all your Cozy data is ready"
	} else if instance.Locale == "fr" {
		objet = "L'archive contenant toutes les données de Cozy est prête"
	}

	msg, err := jobs.NewMessage("json", workers.Options{
		Mode:         workers.ModeNoReply,
		Subject:      objet,
		TemplateName: "archiver_" + instance.Locale,
		TemplateValues: map[string]string{
			"RecipientName": domain,
			"Lien":          lien,
		},
	})
	if err != nil {
		return err
	}

	context := jobs.NewWorkerContext(instance.Domain, "abcd")
	err = workers.SendMail(context, msg)
	if err != nil {
		return err
	}

	return c.JSON(http.StatusOK, echo.Map{
		"message": "bienvenue sur la super page",
	})
}

func import(c echo.Context) error {
	instance := middlewares.GetInstance(c)
	fs := instance.VFS()

	r, err := os.Open("cozy.tar.gz")
	if err != nil {
		return err
	}

	rep := c.Param("destination")
	rep = fmt.Sprintf("/%s", rep)

	exist, err := vfs.DirExists(fs, rep)
	if err != nil {
		return err
	}
	var dst *vfs.DirDoc
	if !exist {
		dst, err = vfs.Mkdir(fs, rep, nil)
		if err != nil {
			return err
		}
	} else {
		dst, err = fs.DirByPath(rep)
		if err != nil {
			return err
		}
	}

	err = imexport.Untardir(r, dst.ID(), instance)
	if err != nil {
		return err
	}

	return c.JSON(http.StatusOK, echo.Map{
		"message": "bienvenue sur la super page cozy",
	})
}

// Routes sets the routing for export
func Routes(router *echo.Group) {
	router.GET("/export/", export)
	router.GET("/import/:destination", import)
}
<|endoftext|>


<commit_before>package counters

import (
	"encoding/json"
	"errors"

	"github.com/seiflotfy/counts/counters/abstract"
	"github.com/seiflotfy/counts/counters/wrappers/cuckoofilter"
	"github.com/seiflotfy/counts/counters/wrappers/hllpp"
	"github.com/seiflotfy/counts/storage"
	"github.com/seiflotfy/counts/utils"

	"github.com/hashicorp/golang-lru"
)

/*
ManagerStruct is responsible for manipulating the counters and syncing to disk
*/
type ManagerStruct struct {
	cache *lru.Cache
	info  map[string]abstract.Info
}

var manager *ManagerStruct
var logger = utils.GetLogger()

/*
CreateDomain ...
*/
func (m *ManagerStruct) CreateDomain(domainID string, domainType string, capacity uint64) error {
	//TODO: spit errir uf domainType is invalid
	//FIXME: no hardcoding of immutable here
	info := abstract.Info{ID: domainID, Type: domainType, Capacity: capacity}
	switch domainType {
	case abstract.Default:
		m.cache.Add(info.ID, hllpp.NewDomain(info))
	case abstract.Purgable:
		m.cache.Add(info.ID, cuckoofilter.NewDomain(info))
	default:
		return errors.New("invalid domain type: " + domainType)
	}
	m.dumpInfo(&info)
	return nil
}

/*
DeleteDomain ...
*/
func (m *ManagerStruct) DeleteDomain(domainID string) error {
	m.cache.Remove(domainID)
	return nil
}

/*
GetDomains ...
*/
func (m *ManagerStruct) GetDomains() ([]string, error) {
	// TODO: Remove dummy data and implement proper result
	values := manager.cache.Keys()
	domains := make([]string, len(values), len(values))
	for i, v := range values {
		domains[i] = v.(string)
	}
	return domains, nil
}

/*
AddToDomain ...
*/
func (m *ManagerStruct) AddToDomain(domainID string, values []string) error {
	var val, ok = m.cache.Get(domainID)
	if ok == false {
		return errors.New("No such domain: " + domainID)
	}
	var counter abstract.Counter
	counter = val.(abstract.Counter)

	bytes := make([][]byte, len(values), len(values))
	for i, value := range values {
		bytes[i] = []byte(value)
	}
	counter.AddMultiple(bytes)
	return nil
}

/*
DeleteFromDomain ...
*/
func (m *ManagerStruct) DeleteFromDomain(domainID string, values []string) error {
	var val, ok = m.cache.Get(domainID)
	if ok == false {
		return errors.New("No such domain: " + domainID)
	}
	var counter abstract.Counter
	counter = val.(abstract.Counter)

	bytes := make([][]byte, len(values), len(values))
	for i, value := range values {
		bytes[i] = []byte(value)
	}
	ok, err := counter.RemoveMultiple(bytes)
	return err
}

/*
GetCountForDomain ...
*/
func (m *ManagerStruct) GetCountForDomain(domainID string) (uint, error) {
	var val, ok = m.cache.Get(domainID)
	if ok == false {
		return 0, errors.New("No such domain: " + domainID)
	}
	var counter abstract.Counter
	counter = val.(abstract.Counter)
	count := counter.GetCount()
	return count, nil
}

/*
GetManager returns a singleton Manager
*/
func GetManager() (*ManagerStruct, error) {
	var err error
	if manager == nil {
		manager, err = newManager()
	}
	if err != nil {
		return nil, err
	}
	return manager, nil
}

func newManager() (*ManagerStruct, error) {
	cache, _ := lru.New(100)
	manager = &ManagerStruct{cache, make(map[string]abstract.Info)}
	err := manager.loadInfo()
	if err != nil {
		return nil, err
	}
	err = manager.loadDomains()
	if err != nil {
		return nil, err
	}
	return manager, nil
}

func (m *ManagerStruct) dumpInfo(i *abstract.Info) {
	manager := storage.GetManager()
	infoData, err := json.Marshal(i)
	utils.PanicOnError(err)
	manager.SaveInfo(i.ID, infoData)
}

func (m *ManagerStruct) loadInfo() error {
	manager := storage.GetManager()
	var infoStruct abstract.Info
	infos, err := manager.LoadAllInfo()
	//FIXME: Should we panic?
	if err != nil {
		return err
	}
	for _, infoData := range infos {
		json.Unmarshal(infoData, &infoStruct)
		m.info[infoStruct.ID] = infoStruct
	}
	return nil
}

func (m *ManagerStruct) loadDomains() error {
	strg := storage.GetManager()
	for key, info := range m.info {
		switch info.Type {
		case abstract.Default:
			//FIXME: Do something with this error
			domain, _ := hllpp.NewDomainFromData(info)
			m.cache.Add(info.ID, domain)
		default:
			logger.Info.Println("Invalid counter type", info.Type)
		}
		strg.LoadData(key, 0, 0)
	}
	return nil
}
<commit_msg>Raise error on failure of loading domains<commit_after>package counters

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

	"github.com/seiflotfy/counts/counters/abstract"
	"github.com/seiflotfy/counts/counters/wrappers/cuckoofilter"
	"github.com/seiflotfy/counts/counters/wrappers/hllpp"
	"github.com/seiflotfy/counts/storage"
	"github.com/seiflotfy/counts/utils"

	"github.com/hashicorp/golang-lru"
)

/*
ManagerStruct is responsible for manipulating the counters and syncing to disk
*/
type ManagerStruct struct {
	cache *lru.Cache
	info  map[string]abstract.Info
}

var manager *ManagerStruct
var logger = utils.GetLogger()

/*
CreateDomain ...
*/
func (m *ManagerStruct) CreateDomain(domainID string, domainType string, capacity uint64) error {
	//TODO: spit errir uf domainType is invalid
	//FIXME: no hardcoding of immutable here
	info := abstract.Info{ID: domainID, Type: domainType, Capacity: capacity}
	switch domainType {
	case abstract.Default:
		m.cache.Add(info.ID, hllpp.NewDomain(info))
	case abstract.Purgable:
		m.cache.Add(info.ID, cuckoofilter.NewDomain(info))
	default:
		return errors.New("invalid domain type: " + domainType)
	}
	m.dumpInfo(&info)
	return nil
}

/*
DeleteDomain ...
*/
func (m *ManagerStruct) DeleteDomain(domainID string) error {
	m.cache.Remove(domainID)
	return nil
}

/*
GetDomains ...
*/
func (m *ManagerStruct) GetDomains() ([]string, error) {
	// TODO: Remove dummy data and implement proper result
	values := manager.cache.Keys()
	domains := make([]string, len(values), len(values))
	for i, v := range values {
		domains[i] = v.(string)
	}
	return domains, nil
}

/*
AddToDomain ...
*/
func (m *ManagerStruct) AddToDomain(domainID string, values []string) error {
	var val, ok = m.cache.Get(domainID)
	if ok == false {
		return errors.New("No such domain: " + domainID)
	}
	var counter abstract.Counter
	counter = val.(abstract.Counter)

	bytes := make([][]byte, len(values), len(values))
	for i, value := range values {
		bytes[i] = []byte(value)
	}
	counter.AddMultiple(bytes)
	return nil
}

/*
DeleteFromDomain ...
*/
func (m *ManagerStruct) DeleteFromDomain(domainID string, values []string) error {
	var val, ok = m.cache.Get(domainID)
	if ok == false {
		return errors.New("No such domain: " + domainID)
	}
	var counter abstract.Counter
	counter = val.(abstract.Counter)

	bytes := make([][]byte, len(values), len(values))
	for i, value := range values {
		bytes[i] = []byte(value)
	}
	ok, err := counter.RemoveMultiple(bytes)
	return err
}

/*
GetCountForDomain ...
*/
func (m *ManagerStruct) GetCountForDomain(domainID string) (uint, error) {
	var val, ok = m.cache.Get(domainID)
	if ok == false {
		return 0, errors.New("No such domain: " + domainID)
	}
	var counter abstract.Counter
	counter = val.(abstract.Counter)
	count := counter.GetCount()
	return count, nil
}

/*
GetManager returns a singleton Manager
*/
func GetManager() (*ManagerStruct, error) {
	var err error
	if manager == nil {
		manager, err = newManager()
	}
	if err != nil {
		return nil, err
	}
	return manager, nil
}

func newManager() (*ManagerStruct, error) {
	cache, _ := lru.New(100)
	manager = &ManagerStruct{cache, make(map[string]abstract.Info)}
	err := manager.loadInfo()
	if err != nil {
		return nil, err
	}
	err = manager.loadDomains()
	if err != nil {
		return nil, err
	}
	return manager, nil
}

func (m *ManagerStruct) dumpInfo(i *abstract.Info) {
	manager := storage.GetManager()
	infoData, err := json.Marshal(i)
	utils.PanicOnError(err)
	manager.SaveInfo(i.ID, infoData)
}

func (m *ManagerStruct) loadInfo() error {
	manager := storage.GetManager()
	var infoStruct abstract.Info
	infos, err := manager.LoadAllInfo()
	//FIXME: Should we panic?
	if err != nil {
		return err
	}
	for _, infoData := range infos {
		json.Unmarshal(infoData, &infoStruct)
		m.info[infoStruct.ID] = infoStruct
	}
	return nil
}

func (m *ManagerStruct) loadDomains() error {
	strg := storage.GetManager()
	for key, info := range m.info {
		switch info.Type {
		case abstract.Default:
			domain, err := hllpp.NewDomainFromData(info)
			if err != nil {
				errTxt := fmt.Sprint("Could not load domain ", info, ". Err:", err)
				return errors.New(errTxt)
			}
			m.cache.Add(info.ID, domain)
		default:
			logger.Info.Println("Invalid counter type", info.Type)
		}
		strg.LoadData(key, 0, 0)
	}
	return nil
}
<|endoftext|>


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

package stat

import (
	"github.com/gonum/matrix/mat64"
)

// CovarianceMatrix calculates a covariance matrix (also known as a
// variance-covariance matrix) from a matrix of data, using a two-pass
// algorithm.  It requires a registered BLAS engine in gonum/matrix/mat64.
//
// The matrix returned will be symmetric, square, and positive-semidefinite.
//
// The weights wts should have the same number of elements as the rows in
// input data matrix x.  cov should be a square matrix with the same number of
// columns as the input data matrix x, or if it is nil then a new Dense
// matrix will be constructed.
func CovarianceMatrix(cov *mat64.Dense, x mat64.Matrix, wts []float64) *mat64.Dense {

	// matrix version of the two pass algorithm.  This doesn't use
	// the correction found in the Covariance and Variance functions.

	r, c := x.Dims()

	// determine the mean of each of the columns
	ones := make([]float64, r)
	for i := range ones {
		ones[i] = 1
	}
	b := mat64.NewDense(1, r, ones)
	b.Mul(b, x)
	b.Scale(1/float64(r), b)
	mu := b.RowView(0)

	// subtract the mean from the data
	xc := mat64.DenseCopyOf(x)

	for i := 0; i < r; i++ {
		rv := xc.RowView(i)
		for j, mean := range mu {
			rv[j] -= mean
		}
	}
	var xt mat64.Dense
	xt.TCopy(xc)

	// normalization factor, typical n-1
	var N float64
	if wts != nil {
		if wr := len(wts); wr != r {
			panic("weight vector length mismatch")
		}

		for i, w := range wts {
			rv := xc.RowView(i)
			N += w
			for j := 0; j < c; j++ {
				rv[j] *= w
			}
		}
		N = 1 / (N - 1)
	} else {
		N = 1 / float64(r-1)
	}

	// TODO: indicate that the resulting matrix is symmetric, which
	// should improve performance.
	if cov == nil {
		cov = mat64.NewDense(c, c, nil)
	} else if covr, covc := cov.Dims(); covr != covc || covc != c {
		panic("matrix size mismatch")
	}

	cov.Mul(&xt, xc)
	cov.Scale(N, cov)
	return cov
}
<commit_msg>Improve comment on what algorithm is used<commit_after>// Copyright ©2014 The gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package stat

import (
	"github.com/gonum/matrix/mat64"
)

// CovarianceMatrix calculates a covariance matrix (also known as a
// variance-covariance matrix) from a matrix of data, using a two-pass
// algorithm.  It requires a registered BLAS engine in gonum/matrix/mat64.
//
// The matrix returned will be symmetric, square, and positive-semidefinite.
//
// The weights wts should have the same number of elements as the rows in
// input data matrix x.  cov should be a square matrix with the same number of
// columns as the input data matrix x, or if it is nil then a new Dense
// matrix will be constructed.
func CovarianceMatrix(cov *mat64.Dense, x mat64.Matrix, wts []float64) *mat64.Dense {
	// This is the matrix version of the two-pass algorithm. It doesn't use
	// the correction found in the Covariance and Variance functions.
	r, c := x.Dims()

	// determine the mean of each of the columns
	ones := make([]float64, r)
	for i := range ones {
		ones[i] = 1
	}
	b := mat64.NewDense(1, r, ones)
	b.Mul(b, x)
	b.Scale(1/float64(r), b)
	mu := b.RowView(0)

	// subtract the mean from the data
	xc := mat64.DenseCopyOf(x)

	for i := 0; i < r; i++ {
		rv := xc.RowView(i)
		for j, mean := range mu {
			rv[j] -= mean
		}
	}
	var xt mat64.Dense
	xt.TCopy(xc)

	// normalization factor, typical n-1
	var N float64
	if wts != nil {
		if wr := len(wts); wr != r {
			panic("weight vector length mismatch")
		}

		for i, w := range wts {
			rv := xc.RowView(i)
			N += w
			for j := 0; j < c; j++ {
				rv[j] *= w
			}
		}
		N = 1 / (N - 1)
	} else {
		N = 1 / float64(r-1)
	}

	// TODO: indicate that the resulting matrix is symmetric, which
	// should improve performance.
	if cov == nil {
		cov = mat64.NewDense(c, c, nil)
	} else if covr, covc := cov.Dims(); covr != covc || covc != c {
		panic("matrix size mismatch")
	}

	cov.Mul(&xt, xc)
	cov.Scale(N, cov)
	return cov
}
<|endoftext|>


<commit_before>// Application that captures SKPs from CT's webpage archives.
package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"strconv"
	"sync"
	"time"

	"github.com/skia-dev/glog"

	"go.skia.org/infra/ct/go/util"
	"go.skia.org/infra/ct/go/worker_scripts/worker_common"
	"go.skia.org/infra/go/common"
	skutil "go.skia.org/infra/go/util"
)

const (
	// The number of goroutines that will run in parallel to capture SKPs.
	WORKER_POOL_SIZE = 10
)

var (
	startRange         = flag.Int("start_range", 1, "The number this worker will capture SKPs from.")
	num                = flag.Int("num", 100, "The total number of SKPs to capture starting from the start_range.")
	pagesetType        = flag.String("pageset_type", util.PAGESET_TYPE_MOBILE_10k, "The type of pagesets to create SKPs from. Eg: 10k, Mobile10k, All.")
	chromiumBuild      = flag.String("chromium_build", "", "The chromium build that will be used to create the SKPs.")
	runID              = flag.String("run_id", "", "The unique run id (typically requester + timestamp).")
	targetPlatform     = flag.String("target_platform", util.PLATFORM_LINUX, "The platform the benchmark will run on (Android / Linux).")
	chromeCleanerTimer = flag.Duration("cleaner_timer", 30*time.Minute, "How often all chrome processes will be killed on this slave.")
)

func main() {
	defer common.LogPanic()
	worker_common.Init()
	defer util.TimeTrack(time.Now(), "Capturing SKPs")
	defer glog.Flush()

	// Validate required arguments.
	if *chromiumBuild == "" {
		glog.Fatal("Must specify --chromium_build")
	}
	if *runID == "" {
		glog.Fatal("Must specify --run_id")
	}
	if *targetPlatform == util.PLATFORM_ANDROID {
		glog.Fatal("Android is not yet supported for capturing SKPs.")
	}

	// Reset the local chromium checkout.
	if err := util.ResetCheckout(util.ChromiumSrcDir); err != nil {
		glog.Fatalf("Could not reset %s: %s", util.ChromiumSrcDir, err)
	}
	// Sync the local chromium checkout.
	if err := util.SyncDir(util.ChromiumSrcDir); err != nil {
		glog.Fatalf("Could not gclient sync %s: %s", util.ChromiumSrcDir, err)
	}

	// Instantiate GsUtil object.
	gs, err := util.NewGsUtil(nil)
	if err != nil {
		glog.Fatal(err)
	}

	// Download the specified chromium build.
	if err := gs.DownloadChromiumBuild(*chromiumBuild); err != nil {
		glog.Fatal(err)
	}
	// Delete the chromium build to save space when we are done.
	defer skutil.RemoveAll(filepath.Join(util.ChromiumBuildsDir, *chromiumBuild))
	chromiumBinary := filepath.Join(util.ChromiumBuildsDir, *chromiumBuild, util.BINARY_CHROME)
	if *targetPlatform == util.PLATFORM_ANDROID {
		// Install the APK on the Android device.
		if err := util.InstallChromeAPK(*chromiumBuild); err != nil {
			glog.Fatalf("Could not install the chromium APK: %s", err)
		}
	}

	// Download pagesets if they do not exist locally.
	pathToPagesets := filepath.Join(util.PagesetsDir, *pagesetType)
	if _, err := gs.DownloadSwarmingArtifacts(pathToPagesets, util.PAGESETS_DIR_NAME, *pagesetType, *startRange, *num); err != nil {
		glog.Fatal(err)
	}
	defer skutil.RemoveAll(pathToPagesets)

	// Download archives if they do not exist locally.
	pathToArchives := filepath.Join(util.WebArchivesDir, *pagesetType)
	archivesToIndex, err := gs.DownloadSwarmingArtifacts(pathToArchives, util.WEB_ARCHIVES_DIR_NAME, *pagesetType, *startRange, *num)
	if err != nil {
		glog.Fatal(err)
	}
	defer skutil.RemoveAll(pathToArchives)

	// Create the dir that SKPs will be stored in.
	pathToSkps := filepath.Join(util.SkpsDir, *pagesetType, *chromiumBuild)
	// Delete and remake the local SKPs directory.
	skutil.RemoveAll(pathToSkps)
	skutil.MkdirAll(pathToSkps, 0700)
	defer skutil.RemoveAll(pathToSkps)

	// Construct path to the ct_run_benchmark python script.
	pathToPyFiles := util.GetPathToPyFiles(!*worker_common.Local)

	timeoutSecs := util.PagesetTypeToInfo[*pagesetType].CaptureSKPsTimeoutSecs
	fileInfos, err := ioutil.ReadDir(pathToPagesets)
	if err != nil {
		glog.Fatalf("Unable to read the pagesets dir %s: %s", pathToPagesets, err)
	}

	// Create channel that contains all pageset file names. This channel will
	// be consumed by the worker pool.
	pagesetRequests := util.GetClosedChannelOfPagesets(fileInfos)

	var wg sync.WaitGroup
	// Use a RWMutex for the chromeProcessesCleaner goroutine to communicate to
	// the workers (acting as "readers") when it wants to be the "writer" and
	// kill all zombie chrome processes.
	var mutex sync.RWMutex

	// Loop through workers in the worker pool.
	for i := 0; i < WORKER_POOL_SIZE; i++ {
		// Increment the WaitGroup counter.
		wg.Add(1)

		// Create and run a goroutine closure that captures SKPs.
		go func() {
			// Decrement the WaitGroup counter when the goroutine completes.
			defer wg.Done()

			for pagesetName := range pagesetRequests {

				mutex.RLock()

				// Read the pageset.
				pagesetPath := filepath.Join(pathToPagesets, pagesetName)
				decodedPageset, err := util.ReadPageset(pagesetPath)
				if err != nil {
					glog.Errorf("Could not read %s: %s", pagesetPath, err)
					continue
				}

				glog.Infof("===== Processing %s =====", pagesetPath)

				skutil.LogErr(os.Chdir(pathToPyFiles))
				index, ok := archivesToIndex[decodedPageset.ArchiveDataFile]
				if !ok {
					glog.Errorf("%s not found in the archivesToIndex map", decodedPageset.ArchiveDataFile)
					continue
				}
				args := []string{
					filepath.Join(util.TelemetryBinariesDir, util.BINARY_RUN_BENCHMARK),
					util.BenchmarksToTelemetryName[util.BENCHMARK_SKPICTURE_PRINTER],
					"--also-run-disabled-tests",
					"--page-repeat=1", // Only need one run for SKPs.
					"--skp-outdir=" + path.Join(pathToSkps, strconv.Itoa(index)),
					"--extra-browser-args=" + util.DEFAULT_BROWSER_ARGS,
					"--user-agent=" + decodedPageset.UserAgent,
					"--urls-list=" + decodedPageset.UrlsList,
					"--archive-data-file=" + decodedPageset.ArchiveDataFile,
				}
				// Figure out which browser and device should be used.
				if *targetPlatform == util.PLATFORM_ANDROID {
					args = append(args, "--browser=android-chromium")
				} else {
					args = append(args, "--browser=exact", "--browser-executable="+chromiumBinary)
					args = append(args, "--device=desktop")
				}
				// Set the PYTHONPATH to the pagesets and the telemetry dirs.
				env := []string{
					fmt.Sprintf("PYTHONPATH=%s:%s:%s:%s:$PYTHONPATH", pathToPagesets, util.TelemetryBinariesDir, util.TelemetrySrcDir, util.CatapultSrcDir),
					"DISPLAY=:0",
				}
				// Retry run_benchmark binary 3 times if there are any errors.
				retryAttempts := 3
				for i := 0; ; i++ {
					err = util.ExecuteCmd("python", args, env, time.Duration(timeoutSecs)*time.Second, nil, nil)
					if err == nil {
						break
					}
					if i >= (retryAttempts - 1) {
						glog.Errorf("%s failed inspite of 3 retries. Last error: %s", pagesetPath, err)
						break
					}
					time.Sleep(time.Second)
					glog.Warningf("Retrying due to error: %s", err)
				}

				mutex.RUnlock()
			}
		}()
	}

	if !*worker_common.Local {
		// Start the cleaner.
		go util.ChromeProcessesCleaner(&mutex, *chromeCleanerTimer)
	}

	// Wait for all spawned goroutines to complete.
	wg.Wait()

	// Move and validate all SKP files.
	if err := util.ValidateSKPs(pathToSkps, pathToPyFiles); err != nil {
		glog.Fatal(err)
	}

	// Upload SKPs dir to Google Storage.
	if err := gs.UploadSwarmingArtifacts(util.SKPS_DIR_NAME, filepath.Join(*pagesetType, *chromiumBuild)); err != nil {
		glog.Fatal(err)
	}
}
<commit_msg>Fail capture_skps task if the SKPs dir is empty<commit_after>// Application that captures SKPs from CT's webpage archives.
package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"strconv"
	"sync"
	"time"

	"github.com/skia-dev/glog"

	"go.skia.org/infra/ct/go/util"
	"go.skia.org/infra/ct/go/worker_scripts/worker_common"
	"go.skia.org/infra/go/common"
	skutil "go.skia.org/infra/go/util"
)

const (
	// The number of goroutines that will run in parallel to capture SKPs.
	WORKER_POOL_SIZE = 10
)

var (
	startRange         = flag.Int("start_range", 1, "The number this worker will capture SKPs from.")
	num                = flag.Int("num", 100, "The total number of SKPs to capture starting from the start_range.")
	pagesetType        = flag.String("pageset_type", util.PAGESET_TYPE_MOBILE_10k, "The type of pagesets to create SKPs from. Eg: 10k, Mobile10k, All.")
	chromiumBuild      = flag.String("chromium_build", "", "The chromium build that will be used to create the SKPs.")
	runID              = flag.String("run_id", "", "The unique run id (typically requester + timestamp).")
	targetPlatform     = flag.String("target_platform", util.PLATFORM_LINUX, "The platform the benchmark will run on (Android / Linux).")
	chromeCleanerTimer = flag.Duration("cleaner_timer", 30*time.Minute, "How often all chrome processes will be killed on this slave.")
)

func main() {
	defer common.LogPanic()
	worker_common.Init()
	defer util.TimeTrack(time.Now(), "Capturing SKPs")
	defer glog.Flush()

	// Validate required arguments.
	if *chromiumBuild == "" {
		glog.Fatal("Must specify --chromium_build")
	}
	if *runID == "" {
		glog.Fatal("Must specify --run_id")
	}
	if *targetPlatform == util.PLATFORM_ANDROID {
		glog.Fatal("Android is not yet supported for capturing SKPs.")
	}

	// Reset the local chromium checkout.
	if err := util.ResetCheckout(util.ChromiumSrcDir); err != nil {
		glog.Fatalf("Could not reset %s: %s", util.ChromiumSrcDir, err)
	}
	// Sync the local chromium checkout.
	if err := util.SyncDir(util.ChromiumSrcDir); err != nil {
		glog.Fatalf("Could not gclient sync %s: %s", util.ChromiumSrcDir, err)
	}

	// Instantiate GsUtil object.
	gs, err := util.NewGsUtil(nil)
	if err != nil {
		glog.Fatal(err)
	}

	// Download the specified chromium build.
	if err := gs.DownloadChromiumBuild(*chromiumBuild); err != nil {
		glog.Fatal(err)
	}
	// Delete the chromium build to save space when we are done.
	defer skutil.RemoveAll(filepath.Join(util.ChromiumBuildsDir, *chromiumBuild))
	chromiumBinary := filepath.Join(util.ChromiumBuildsDir, *chromiumBuild, util.BINARY_CHROME)
	if *targetPlatform == util.PLATFORM_ANDROID {
		// Install the APK on the Android device.
		if err := util.InstallChromeAPK(*chromiumBuild); err != nil {
			glog.Fatalf("Could not install the chromium APK: %s", err)
		}
	}

	// Download pagesets if they do not exist locally.
	pathToPagesets := filepath.Join(util.PagesetsDir, *pagesetType)
	if _, err := gs.DownloadSwarmingArtifacts(pathToPagesets, util.PAGESETS_DIR_NAME, *pagesetType, *startRange, *num); err != nil {
		glog.Fatal(err)
	}
	defer skutil.RemoveAll(pathToPagesets)

	// Download archives if they do not exist locally.
	pathToArchives := filepath.Join(util.WebArchivesDir, *pagesetType)
	archivesToIndex, err := gs.DownloadSwarmingArtifacts(pathToArchives, util.WEB_ARCHIVES_DIR_NAME, *pagesetType, *startRange, *num)
	if err != nil {
		glog.Fatal(err)
	}
	defer skutil.RemoveAll(pathToArchives)

	// Create the dir that SKPs will be stored in.
	pathToSkps := filepath.Join(util.SkpsDir, *pagesetType, *chromiumBuild)
	// Delete and remake the local SKPs directory.
	skutil.RemoveAll(pathToSkps)
	skutil.MkdirAll(pathToSkps, 0700)
	defer skutil.RemoveAll(pathToSkps)

	// Construct path to the ct_run_benchmark python script.
	pathToPyFiles := util.GetPathToPyFiles(!*worker_common.Local)

	timeoutSecs := util.PagesetTypeToInfo[*pagesetType].CaptureSKPsTimeoutSecs
	fileInfos, err := ioutil.ReadDir(pathToPagesets)
	if err != nil {
		glog.Fatalf("Unable to read the pagesets dir %s: %s", pathToPagesets, err)
	}

	// Create channel that contains all pageset file names. This channel will
	// be consumed by the worker pool.
	pagesetRequests := util.GetClosedChannelOfPagesets(fileInfos)

	var wg sync.WaitGroup
	// Use a RWMutex for the chromeProcessesCleaner goroutine to communicate to
	// the workers (acting as "readers") when it wants to be the "writer" and
	// kill all zombie chrome processes.
	var mutex sync.RWMutex

	// Loop through workers in the worker pool.
	for i := 0; i < WORKER_POOL_SIZE; i++ {
		// Increment the WaitGroup counter.
		wg.Add(1)

		// Create and run a goroutine closure that captures SKPs.
		go func() {
			// Decrement the WaitGroup counter when the goroutine completes.
			defer wg.Done()

			for pagesetName := range pagesetRequests {

				mutex.RLock()

				// Read the pageset.
				pagesetPath := filepath.Join(pathToPagesets, pagesetName)
				decodedPageset, err := util.ReadPageset(pagesetPath)
				if err != nil {
					glog.Errorf("Could not read %s: %s", pagesetPath, err)
					continue
				}

				glog.Infof("===== Processing %s =====", pagesetPath)

				skutil.LogErr(os.Chdir(pathToPyFiles))
				index, ok := archivesToIndex[decodedPageset.ArchiveDataFile]
				if !ok {
					glog.Errorf("%s not found in the archivesToIndex map", decodedPageset.ArchiveDataFile)
					continue
				}
				args := []string{
					filepath.Join(util.TelemetryBinariesDir, util.BINARY_RUN_BENCHMARK),
					util.BenchmarksToTelemetryName[util.BENCHMARK_SKPICTURE_PRINTER],
					"--also-run-disabled-tests",
					"--page-repeat=1", // Only need one run for SKPs.
					"--skp-outdir=" + path.Join(pathToSkps, strconv.Itoa(index)),
					"--extra-browser-args=" + util.DEFAULT_BROWSER_ARGS,
					"--user-agent=" + decodedPageset.UserAgent,
					"--urls-list=" + decodedPageset.UrlsList,
					"--archive-data-file=" + decodedPageset.ArchiveDataFile,
				}
				// Figure out which browser and device should be used.
				if *targetPlatform == util.PLATFORM_ANDROID {
					args = append(args, "--browser=android-chromium")
				} else {
					args = append(args, "--browser=exact", "--browser-executable="+chromiumBinary)
					args = append(args, "--device=desktop")
				}
				// Set the PYTHONPATH to the pagesets and the telemetry dirs.
				env := []string{
					fmt.Sprintf("PYTHONPATH=%s:%s:%s:%s:$PYTHONPATH", pathToPagesets, util.TelemetryBinariesDir, util.TelemetrySrcDir, util.CatapultSrcDir),
					"DISPLAY=:0",
				}
				// Retry run_benchmark binary 3 times if there are any errors.
				retryAttempts := 3
				for i := 0; ; i++ {
					err = util.ExecuteCmd("python", args, env, time.Duration(timeoutSecs)*time.Second, nil, nil)
					if err == nil {
						break
					}
					if i >= (retryAttempts - 1) {
						glog.Errorf("%s failed inspite of 3 retries. Last error: %s", pagesetPath, err)
						break
					}
					time.Sleep(time.Second)
					glog.Warningf("Retrying due to error: %s", err)
				}

				mutex.RUnlock()
			}
		}()
	}

	if !*worker_common.Local {
		// Start the cleaner.
		go util.ChromeProcessesCleaner(&mutex, *chromeCleanerTimer)
	}

	// Wait for all spawned goroutines to complete.
	wg.Wait()

	// Move and validate all SKP files.
	if err := util.ValidateSKPs(pathToSkps, pathToPyFiles); err != nil {
		glog.Fatal(err)
	}

	// Check to see if there is anything in the pathToSKPs dir.
	skpsEmpty, err := skutil.IsDirEmpty(pathToSkps)
	if err != nil {
		glog.Fatal(err)
	}
	if skpsEmpty {
		glog.Fatalf("Could not create any SKP in %s", pathToSkps)
	}

	// Upload SKPs dir to Google Storage.
	if err := gs.UploadSwarmingArtifacts(util.SKPS_DIR_NAME, filepath.Join(*pagesetType, *chromiumBuild)); err != nil {
		glog.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package cli

import (
	"testing"

	"github.com/libopenstorage/openstorage/api"
	"github.com/stretchr/testify/require"
)

func TestCmdMarshalProto(t *testing.T) {
	volumeSpec := &api.VolumeSpec{
		Size:   64,
		Format: api.FSType_FS_TYPE_EXT4,
	}
	data := cmdMarshalProto(volumeSpec, false)
	require.Equal(
		t,
		`{
 "ephemeral": false,
 "size": "64",
 "format": "ext4",
 "block_size": "0",
 "ha_level": "0",
 "cos": "none",
 "io_profile": "sequential",
 "dedupe": false,
 "snapshot_interval": 0,
 "shared": false,
 "aggregation_level": 0,
 "encrypted": false,
 "passphrase": "",
 "snapshot_schedule": "",
 "scale": 0,
 "sticky": false,
 "max_backups": 0,
 "backup_schedule": "",
 "group_enforced": false
}`,
		data,
	)
}
<commit_msg>Update UT for removed fields<commit_after>package cli

import (
	"testing"

	"github.com/libopenstorage/openstorage/api"
	"github.com/stretchr/testify/require"
)

func TestCmdMarshalProto(t *testing.T) {
	volumeSpec := &api.VolumeSpec{
		Size:   64,
		Format: api.FSType_FS_TYPE_EXT4,
	}
	data := cmdMarshalProto(volumeSpec, false)
	require.Equal(
		t,
		`{
 "ephemeral": false,
 "size": "64",
 "format": "ext4",
 "block_size": "0",
 "ha_level": "0",
 "cos": "none",
 "io_profile": "sequential",
 "dedupe": false,
 "snapshot_interval": 0,
 "shared": false,
 "aggregation_level": 0,
 "encrypted": false,
 "passphrase": "",
 "snapshot_schedule": "",
 "scale": 0,
 "sticky": false,
 "group_enforced": false
}`,
		data,
	)
}
<|endoftext|>


<commit_before>package cluster

import (
	"crypto/tls"
	"errors"
	"fmt"
	"net"
	"strings"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/samalba/dockerclient"
)

const (
	// Force-refresh the state of the node this often.
	stateRefreshPeriod = 30 * time.Second

	// Timeout for requests sent out to the node.
	requestTimeout = 10 * time.Second
)

func NewNode(addr string, overcommitRatio float64) *Node {
	e := &Node{
		Addr:            addr,
		Labels:          make(map[string]string),
		ch:              make(chan bool),
		containers:      make(map[string]*Container),
		healthy:         true,
		overcommitRatio: int64(overcommitRatio * 100),
	}
	return e
}

type Node struct {
	sync.RWMutex

	ID     string
	IP     string
	Addr   string
	Name   string
	Cpus   int64
	Memory int64
	Labels map[string]string

	ch              chan bool
	containers      map[string]*Container
	client          dockerclient.Client
	eventHandler    EventHandler
	healthy         bool
	overcommitRatio int64
}

// Connect will initialize a connection to the Docker daemon running on the
// host, gather machine specs (memory, cpu, ...) and monitor state changes.
func (n *Node) Connect(config *tls.Config) error {
	parts := strings.Split(n.Addr, ":")
	addr, err := net.ResolveIPAddr("ip4", parts[0])
	if err != nil {
		return err
	}
	n.IP = addr.IP.String()

	c, err := dockerclient.NewDockerClientTimeout(n.IP+":"+parts[1], config, time.Duration(requestTimeout))
	if err != nil {
		return err
	}

	return n.connectClient(c)
}

func (n *Node) connectClient(client dockerclient.Client) error {
	n.client = client

	// Fetch the engine labels.
	if err := n.updateSpecs(); err != nil {
		n.client = nil
		return err
	}

	// Force a state update before returning.
	if err := n.refreshContainers(); err != nil {
		n.client = nil
		return err
	}

	// Start the update loop.
	go n.refreshLoop()

	// Start monitoring events from the Node.
	n.client.StartMonitorEvents(n.handler)

	return nil
}

// IsConnected returns true if the engine is connected to a remote docker API
func (n *Node) IsConnected() bool {
	return n.client != nil
}

func (n *Node) IsHealthy() bool {
	return n.healthy
}

// Gather node specs (CPU, memory, constraints, ...).
func (n *Node) updateSpecs() error {
	info, err := n.client.Info()
	if err != nil {
		return err
	}
	// Older versions of Docker don't expose the ID field and are not supported
	// by Swarm.  Catch the error ASAP and refuse to connect.
	if len(info.ID) == 0 {
		return fmt.Errorf("Node %s is running an unsupported version of Docker Engine. Please upgrade.", n.Addr)
	}
	n.ID = info.ID
	n.Name = info.Name
	n.Cpus = info.NCPU
	n.Memory = info.MemTotal
	n.Labels = map[string]string{
		"storagedriver":   info.Driver,
		"executiondriver": info.ExecutionDriver,
		"kernelversion":   info.KernelVersion,
		"operatingsystem": info.OperatingSystem,
	}
	for _, label := range info.Labels {
		kv := strings.SplitN(label, "=", 2)
		n.Labels[kv[0]] = kv[1]
	}
	return nil
}

// Refresh the list and status of containers running on the node.
func (n *Node) refreshContainers() error {
	containers, err := n.client.ListContainers(true, false, "")
	if err != nil {
		return err
	}

	n.Lock()
	defer n.Unlock()

	merged := make(map[string]*Container)
	for _, c := range containers {
		merged, err = n.updateContainer(c, merged)
		if err != nil {
			log.Errorf("[%s/%s] Unable to update state of %s", n.ID, n.Name, c.Id)
		}
	}

	n.containers = merged

	log.Debugf("[%s/%s] Updated state", n.ID, n.Name)
	return nil
}

// Refresh the status of a container running on the node.
func (n *Node) refreshContainer(ID string) error {
	containers, err := n.client.ListContainers(true, false, fmt.Sprintf("{%q:[%q]}", "id", ID))
	if err != nil {
		return err
	}

	if len(containers) > 1 {
		// We expect one container, if we get more than one, trigger a full refresh.
		return n.refreshContainers()
	}

	n.Lock()
	defer n.Unlock()

	if len(containers) == 0 {
		// The container doesn't exist on the node, remove it.
		delete(n.containers, ID)
		return nil
	}

	_, err = n.updateContainer(containers[0], n.containers)
	return err
}

func (n *Node) ForceRefreshContainer(c dockerclient.Container) error {
	return n.inspectContainer(c, n.containers, true)
}

func (n *Node) inspectContainer(c dockerclient.Container, containers map[string]*Container, lock bool) error {

	container := &Container{
		Container: c,
		Node:      n,
	}

	info, err := n.client.InspectContainer(c.Id)
	if err != nil {
		return err
	}
	container.Info = *info

	// real CpuShares -> nb of CPUs
	container.Info.Config.CpuShares = container.Info.Config.CpuShares / 100.0 * n.Cpus

	if lock {
		n.Lock()
		defer n.Unlock()
	}
	containers[container.Id] = container

	return nil
}

func (n *Node) updateContainer(c dockerclient.Container, containers map[string]*Container) (map[string]*Container, error) {
	if current, exists := n.containers[c.Id]; exists {
		// The container exists. Update its state.
		current.Container = c
		containers[current.Id] = current
	} else {
		// This is a brand new container.
		if err := n.inspectContainer(c, containers, false); err != nil {
			return nil, err
		}
	}
	return containers, nil
}

func (n *Node) refreshContainersAsync() {
	n.ch <- true
}

func (n *Node) refreshLoop() {
	for {
		var err error
		select {
		case <-n.ch:
			err = n.refreshContainers()
		case <-time.After(stateRefreshPeriod):
			err = n.refreshContainers()
		}

		if err != nil {
			n.healthy = false
			log.Errorf("[%s/%s] Flagging node as dead. Updated state failed: %v", n.ID, n.Name, err)
		} else {
			if !n.healthy {
				log.Infof("[%s/%s] Node came back to life. Hooray!", n.ID, n.Name)
				n.client.StopAllMonitorEvents()
				n.client.StartMonitorEvents(n.handler)
			}
			n.healthy = true
		}
	}
}

// Return the sum of memory reserved by containers.
func (n *Node) ReservedMemory() int64 {
	var r int64 = 0
	n.RLock()
	for _, c := range n.containers {
		r += c.Info.Config.Memory
	}
	n.RUnlock()
	return r
}

// Return the sum of CPUs reserved by containers.
func (n *Node) ReservedCpus() int64 {
	var r int64 = 0
	n.RLock()
	for _, c := range n.containers {
		r += c.Info.Config.CpuShares
	}
	n.RUnlock()
	return r
}

func (n *Node) UsableMemory() int64 {
	return n.Memory + (n.Memory * n.overcommitRatio / 100)
}

func (n *Node) UsableCpus() int64 {
	return n.Cpus + (n.Cpus * n.overcommitRatio / 100)
}

func (n *Node) Create(config *dockerclient.ContainerConfig, name string, pullImage bool) (*Container, error) {
	var (
		err    error
		id     string
		client = n.client
	)

	newConfig := *config

	// nb of CPUs -> real CpuShares
	newConfig.CpuShares = config.CpuShares * 100 / n.Cpus

	if id, err = client.CreateContainer(&newConfig, name); err != nil {
		// If the error is other than not found, abort immediately.
		if err != dockerclient.ErrNotFound || !pullImage {
			return nil, err
		}
		// Otherwise, try to pull the image...
		if err = n.Pull(config.Image); err != nil {
			return nil, err
		}
		// ...And try again.
		if id, err = client.CreateContainer(&newConfig, name); err != nil {
			return nil, err
		}
	}

	// Register the container immediately while waiting for a state refresh.
	// Force a state refresh to pick up the newly created container.
	n.refreshContainer(id)

	n.RLock()
	defer n.RUnlock()

	return n.containers[id], nil
}

func (n *Node) ListImages() ([]string, error) {
	images, err := n.client.ListImages()
	if err != nil {
		return nil, err
	}

	out := []string{}

	for _, i := range images {
		for _, t := range i.RepoTags {
			out = append(out, t)
		}
	}

	return out, nil
}

// Destroy and remove a container from the node.
func (n *Node) Destroy(container *Container, force bool) error {
	if err := n.client.RemoveContainer(container.Id, force); err != nil {
		return err
	}

	// Remove the container from the state. Eventually, the state refresh loop
	// will rewrite this.
	n.Lock()
	defer n.Unlock()
	delete(n.containers, container.Id)

	return nil
}

func (n *Node) Pull(image string) error {
	if err := n.client.PullImage(image, nil); err != nil {
		return err
	}
	return nil
}

// Register an event handler.
func (n *Node) Events(h EventHandler) error {
	if n.eventHandler != nil {
		return errors.New("event handler already set")
	}
	n.eventHandler = h
	return nil
}

func (n *Node) Containers() []*Container {
	containers := []*Container{}
	n.RLock()
	for _, container := range n.containers {
		containers = append(containers, container)
	}
	n.RUnlock()
	return containers
}

func (n *Node) String() string {
	return fmt.Sprintf("node %s addr %s", n.ID, n.Addr)
}

func (n *Node) handler(ev *dockerclient.Event, args ...interface{}) {
	// Something changed - refresh our internal state.
	n.refreshContainer(ev.Id)

	// If there is no event handler registered, abort right now.
	if n.eventHandler == nil {
		return
	}

	event := &Event{
		Node:  n,
		Event: *ev,
	}

	n.eventHandler.Handle(event)
}

// Inject a container into the internal state.
func (n *Node) AddContainer(container *Container) error {
	n.Lock()
	defer n.Unlock()

	if _, ok := n.containers[container.Id]; ok {
		return errors.New("container already exists")
	}
	n.containers[container.Id] = container
	return nil
}

// Remove a container from the internal test.
func (n *Node) RemoveContainer(container *Container) error {
	n.Lock()
	defer n.Unlock()

	if _, ok := n.containers[container.Id]; !ok {
		return errors.New("container not found")
	}
	delete(n.containers, container.Id)
	return nil
}

// Wipes the internal container state.
func (n *Node) CleanupContainers() {
	n.Lock()
	n.containers = make(map[string]*Container)
	n.Unlock()
}
<commit_msg>emit a custom event when a node dies or comes back<commit_after>package cluster

import (
	"crypto/tls"
	"errors"
	"fmt"
	"net"
	"strings"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/samalba/dockerclient"
)

const (
	// Force-refresh the state of the node this often.
	stateRefreshPeriod = 30 * time.Second

	// Timeout for requests sent out to the node.
	requestTimeout = 10 * time.Second
)

func NewNode(addr string, overcommitRatio float64) *Node {
	e := &Node{
		Addr:            addr,
		Labels:          make(map[string]string),
		ch:              make(chan bool),
		containers:      make(map[string]*Container),
		healthy:         true,
		overcommitRatio: int64(overcommitRatio * 100),
	}
	return e
}

type Node struct {
	sync.RWMutex

	ID     string
	IP     string
	Addr   string
	Name   string
	Cpus   int64
	Memory int64
	Labels map[string]string

	ch              chan bool
	containers      map[string]*Container
	client          dockerclient.Client
	eventHandler    EventHandler
	healthy         bool
	overcommitRatio int64
}

// Connect will initialize a connection to the Docker daemon running on the
// host, gather machine specs (memory, cpu, ...) and monitor state changes.
func (n *Node) Connect(config *tls.Config) error {
	parts := strings.Split(n.Addr, ":")
	addr, err := net.ResolveIPAddr("ip4", parts[0])
	if err != nil {
		return err
	}
	n.IP = addr.IP.String()

	c, err := dockerclient.NewDockerClientTimeout(n.IP+":"+parts[1], config, time.Duration(requestTimeout))
	if err != nil {
		return err
	}

	return n.connectClient(c)
}

func (n *Node) connectClient(client dockerclient.Client) error {
	n.client = client

	// Fetch the engine labels.
	if err := n.updateSpecs(); err != nil {
		n.client = nil
		return err
	}

	// Force a state update before returning.
	if err := n.refreshContainers(); err != nil {
		n.client = nil
		return err
	}

	// Start the update loop.
	go n.refreshLoop()

	// Start monitoring events from the Node.
	n.client.StartMonitorEvents(n.handler)

	return nil
}

// IsConnected returns true if the engine is connected to a remote docker API
func (n *Node) IsConnected() bool {
	return n.client != nil
}

func (n *Node) IsHealthy() bool {
	return n.healthy
}

// Gather node specs (CPU, memory, constraints, ...).
func (n *Node) updateSpecs() error {
	info, err := n.client.Info()
	if err != nil {
		return err
	}
	// Older versions of Docker don't expose the ID field and are not supported
	// by Swarm.  Catch the error ASAP and refuse to connect.
	if len(info.ID) == 0 {
		return fmt.Errorf("Node %s is running an unsupported version of Docker Engine. Please upgrade.", n.Addr)
	}
	n.ID = info.ID
	n.Name = info.Name
	n.Cpus = info.NCPU
	n.Memory = info.MemTotal
	n.Labels = map[string]string{
		"storagedriver":   info.Driver,
		"executiondriver": info.ExecutionDriver,
		"kernelversion":   info.KernelVersion,
		"operatingsystem": info.OperatingSystem,
	}
	for _, label := range info.Labels {
		kv := strings.SplitN(label, "=", 2)
		n.Labels[kv[0]] = kv[1]
	}
	return nil
}

// Refresh the list and status of containers running on the node.
func (n *Node) refreshContainers() error {
	containers, err := n.client.ListContainers(true, false, "")
	if err != nil {
		return err
	}

	n.Lock()
	defer n.Unlock()

	merged := make(map[string]*Container)
	for _, c := range containers {
		merged, err = n.updateContainer(c, merged)
		if err != nil {
			log.Errorf("[%s/%s] Unable to update state of %s", n.ID, n.Name, c.Id)
		}
	}

	n.containers = merged

	log.Debugf("[%s/%s] Updated state", n.ID, n.Name)
	return nil
}

// Refresh the status of a container running on the node.
func (n *Node) refreshContainer(ID string) error {
	containers, err := n.client.ListContainers(true, false, fmt.Sprintf("{%q:[%q]}", "id", ID))
	if err != nil {
		return err
	}

	if len(containers) > 1 {
		// We expect one container, if we get more than one, trigger a full refresh.
		return n.refreshContainers()
	}

	n.Lock()
	defer n.Unlock()

	if len(containers) == 0 {
		// The container doesn't exist on the node, remove it.
		delete(n.containers, ID)
		return nil
	}

	_, err = n.updateContainer(containers[0], n.containers)
	return err
}

func (n *Node) ForceRefreshContainer(c dockerclient.Container) error {
	return n.inspectContainer(c, n.containers, true)
}

func (n *Node) inspectContainer(c dockerclient.Container, containers map[string]*Container, lock bool) error {

	container := &Container{
		Container: c,
		Node:      n,
	}

	info, err := n.client.InspectContainer(c.Id)
	if err != nil {
		return err
	}
	container.Info = *info

	// real CpuShares -> nb of CPUs
	container.Info.Config.CpuShares = container.Info.Config.CpuShares / 100.0 * n.Cpus

	if lock {
		n.Lock()
		defer n.Unlock()
	}
	containers[container.Id] = container

	return nil
}

func (n *Node) updateContainer(c dockerclient.Container, containers map[string]*Container) (map[string]*Container, error) {
	if current, exists := n.containers[c.Id]; exists {
		// The container exists. Update its state.
		current.Container = c
		containers[current.Id] = current
	} else {
		// This is a brand new container.
		if err := n.inspectContainer(c, containers, false); err != nil {
			return nil, err
		}
	}
	return containers, nil
}

func (n *Node) refreshContainersAsync() {
	n.ch <- true
}

func (n *Node) refreshLoop() {
	for {
		var err error
		select {
		case <-n.ch:
			err = n.refreshContainers()
		case <-time.After(stateRefreshPeriod):
			err = n.refreshContainers()
		}

		if err != nil {
			if n.healthy {
				n.emitCustomEvent("node_die")
			}
			n.healthy = false
			log.Errorf("[%s/%s] Flagging node as dead. Updated state failed: %v", n.ID, n.Name, err)
		} else {
			if !n.healthy {
				log.Infof("[%s/%s] Node came back to life. Hooray!", n.ID, n.Name)
				n.client.StopAllMonitorEvents()
				n.client.StartMonitorEvents(n.handler)
				n.emitCustomEvent("node_comes_back")
			}
			n.healthy = true
		}
	}
}

func (n *Node) emitCustomEvent(event string) {
	// If there is no event handler registered, abort right now.
	if n.eventHandler == nil {
		return
	}
	ev := &Event{
		Event: dockerclient.Event{
			Status: event,
			From:   "swarm",
			Time:   time.Now().Unix(),
		},
		Node: n,
	}
	n.eventHandler.Handle(ev)
}

// Return the sum of memory reserved by containers.
func (n *Node) ReservedMemory() int64 {
	var r int64 = 0
	n.RLock()
	for _, c := range n.containers {
		r += c.Info.Config.Memory
	}
	n.RUnlock()
	return r
}

// Return the sum of CPUs reserved by containers.
func (n *Node) ReservedCpus() int64 {
	var r int64 = 0
	n.RLock()
	for _, c := range n.containers {
		r += c.Info.Config.CpuShares
	}
	n.RUnlock()
	return r
}

func (n *Node) UsableMemory() int64 {
	return n.Memory + (n.Memory * n.overcommitRatio / 100)
}

func (n *Node) UsableCpus() int64 {
	return n.Cpus + (n.Cpus * n.overcommitRatio / 100)
}

func (n *Node) Create(config *dockerclient.ContainerConfig, name string, pullImage bool) (*Container, error) {
	var (
		err    error
		id     string
		client = n.client
	)

	newConfig := *config

	// nb of CPUs -> real CpuShares
	newConfig.CpuShares = config.CpuShares * 100 / n.Cpus

	if id, err = client.CreateContainer(&newConfig, name); err != nil {
		// If the error is other than not found, abort immediately.
		if err != dockerclient.ErrNotFound || !pullImage {
			return nil, err
		}
		// Otherwise, try to pull the image...
		if err = n.Pull(config.Image); err != nil {
			return nil, err
		}
		// ...And try again.
		if id, err = client.CreateContainer(&newConfig, name); err != nil {
			return nil, err
		}
	}

	// Register the container immediately while waiting for a state refresh.
	// Force a state refresh to pick up the newly created container.
	n.refreshContainer(id)

	n.RLock()
	defer n.RUnlock()

	return n.containers[id], nil
}

func (n *Node) ListImages() ([]string, error) {
	images, err := n.client.ListImages()
	if err != nil {
		return nil, err
	}

	out := []string{}

	for _, i := range images {
		for _, t := range i.RepoTags {
			out = append(out, t)
		}
	}

	return out, nil
}

// Destroy and remove a container from the node.
func (n *Node) Destroy(container *Container, force bool) error {
	if err := n.client.RemoveContainer(container.Id, force); err != nil {
		return err
	}

	// Remove the container from the state. Eventually, the state refresh loop
	// will rewrite this.
	n.Lock()
	defer n.Unlock()
	delete(n.containers, container.Id)

	return nil
}

func (n *Node) Pull(image string) error {
	if err := n.client.PullImage(image, nil); err != nil {
		return err
	}
	return nil
}

// Register an event handler.
func (n *Node) Events(h EventHandler) error {
	if n.eventHandler != nil {
		return errors.New("event handler already set")
	}
	n.eventHandler = h
	return nil
}

func (n *Node) Containers() []*Container {
	containers := []*Container{}
	n.RLock()
	for _, container := range n.containers {
		containers = append(containers, container)
	}
	n.RUnlock()
	return containers
}

func (n *Node) String() string {
	return fmt.Sprintf("node %s addr %s", n.ID, n.Addr)
}

func (n *Node) handler(ev *dockerclient.Event, args ...interface{}) {
	// Something changed - refresh our internal state.
	n.refreshContainer(ev.Id)

	// If there is no event handler registered, abort right now.
	if n.eventHandler == nil {
		return
	}

	event := &Event{
		Node:  n,
		Event: *ev,
	}

	n.eventHandler.Handle(event)
}

// Inject a container into the internal state.
func (n *Node) AddContainer(container *Container) error {
	n.Lock()
	defer n.Unlock()

	if _, ok := n.containers[container.Id]; ok {
		return errors.New("container already exists")
	}
	n.containers[container.Id] = container
	return nil
}

// Remove a container from the internal test.
func (n *Node) RemoveContainer(container *Container) error {
	n.Lock()
	defer n.Unlock()

	if _, ok := n.containers[container.Id]; !ok {
		return errors.New("container not found")
	}
	delete(n.containers, container.Id)
	return nil
}

// Wipes the internal container state.
func (n *Node) CleanupContainers() {
	n.Lock()
	n.containers = make(map[string]*Container)
	n.Unlock()
}
<|endoftext|>