<commit_before>package backends

import (
	"encoding/json"
	"fmt"
	"github.com/docker/libswarm/beam"
	"github.com/dotcloud/docker/api"
	"github.com/dotcloud/docker/pkg/version"
	"github.com/dotcloud/docker/utils"
	"github.com/gorilla/mux"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"
	"strconv"
)

func DockerServer() beam.Sender {
	backend := beam.NewServer()
	backend.OnSpawn(beam.Handler(func(ctx *beam.Message) error {
		instance := beam.Task(func(in beam.Receiver, out beam.Sender) {
			url := "tcp://localhost:4243"
			if len(ctx.Args) > 0 {
				url = ctx.Args[0]
			}
			err := listenAndServe(url, out)
			if err != nil {
				fmt.Printf("listenAndServe: %v", err)
			}
		})
		_, err := ctx.Ret.Send(&beam.Message{Verb: beam.Ack, Ret: instance})
		return err
	}))
	return backend
}

type HttpApiFunc func(out beam.Sender, version version.Version, w http.ResponseWriter, r *http.Request, vars map[string]string) error

func listenAndServe(urlStr string, out beam.Sender) error {
	fmt.Println("Starting Docker server...")
	r, err := createRouter(out)
	if err != nil {
		return err
	}

	parsedUrl, err := url.Parse(urlStr)
	if err != nil {
		return err
	}

	var hostAndPath string
	// For Unix sockets we need to capture the path as well as the host
	if parsedUrl.Scheme == "unix" {
		hostAndPath = "/" + parsedUrl.Host + parsedUrl.Path
	} else {
		hostAndPath = parsedUrl.Host
	}

	l, err := net.Listen(parsedUrl.Scheme, hostAndPath)
	if err != nil {
		return err
	}
	httpSrv := http.Server{Addr: hostAndPath, Handler: r}
	return httpSrv.Serve(l)
}

func ping(out beam.Sender, version version.Version, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
	_, err := w.Write([]byte{'O', 'K'})
	return err
}

func getContainersJSON(out beam.Sender, version version.Version, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
	if err := r.ParseForm(); err != nil {
		return err
	}

	o := beam.Obj(out)
	names, err := o.Ls()
	if err != nil {
		return err
	}

	var responses []interface{}

	for _, name := range names {
		_, containerOut, err := o.Attach(name)
		if err != nil {
			return err
		}
		container := beam.Obj(containerOut)
		responseJson, err := container.Get()
		if err != nil {
			return err
		}
		var response struct {
			ID      string
			Created string
			Name    string
			Config  struct {
				Cmd   []string
				Image string
			}
			State struct {
				Running    bool
				StartedAt  string
				FinishedAt string
				ExitCode   int
			}
			NetworkSettings struct {
				Ports map[string][]map[string]string
			}
		}
		if err = json.Unmarshal([]byte(responseJson), &response); err != nil {
			return err
		}
		created, err := time.Parse(time.RFC3339, response.Created)
		if err != nil {
			return err
		}
		var state string
		if response.State.Running {
			state = "Up"
		} else {
			state = fmt.Sprintf("Exited (%d)", response.State.ExitCode)
		}
		type port struct {
			IP string
			PrivatePort int
			PublicPort int
			Type string
		}
		var ports []port
		for p, mappings := range response.NetworkSettings.Ports {
			var portnum int
			var proto string
			_, err := fmt.Sscanf(p, "%d/%s", &portnum, &proto)
			if err != nil {
				return err
			}
			if len(mappings) > 0 {
				for _, mapping := range mappings {
					hostPort, err := strconv.Atoi(mapping["HostPort"])
					if err != nil {
						return err
					}
					newport := port{
						IP: mapping["HostIp"],
						PrivatePort: portnum,
						PublicPort: hostPort,
						Type: proto,
					}
					ports = append(ports, newport)
				}
			} else {
				newport := port{
					PublicPort: portnum,
					Type: proto,
				}
				ports = append(ports, newport)
			}
		}
		responses = append(responses, map[string]interface{}{
			"Id":      response.ID,
			"Command": strings.Join(response.Config.Cmd, " "),
			"Created": created.Unix(),
			"Image":   response.Config.Image,
			"Names":   []string{response.Name},
			"Ports":   ports,
			"Status":  state,
		})
	}

	return writeJSON(w, http.StatusOK, responses)
}

func postContainersCreate(out beam.Sender, version version.Version, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
	if err := r.ParseForm(); err != nil {
		return nil
	}

	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		return err
	}

	container, err := beam.Obj(out).Spawn(string(body))
	if err != nil {
		return err
	}

	responseJson, err := container.Get()
	if err != nil {
		return err
	}

	var response struct{ Id string }
	if err = json.Unmarshal([]byte(responseJson), &response); err != nil {
		return err
	}
	return writeJSON(w, http.StatusCreated, response)
}

func postContainersStart(out beam.Sender, version version.Version, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
	if vars == nil {
		return fmt.Errorf("Missing parameter")
	}

	// TODO: r.Body

	name := vars["name"]
	_, containerOut, err := beam.Obj(out).Attach(name)
	container := beam.Obj(containerOut)
	if err != nil {
		return err
	}
	if err := container.Start(); err != nil {
		return err
	}

	w.WriteHeader(http.StatusNoContent)
	return nil
}

func postContainersStop(out beam.Sender, version version.Version, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
	if vars == nil {
		return fmt.Errorf("Missing parameter")
	}

	name := vars["name"]
	_, containerOut, err := beam.Obj(out).Attach(name)
	container := beam.Obj(containerOut)
	if err != nil {
		return err
	}
	if err := container.Stop(); err != nil {
		return err
	}

	w.WriteHeader(http.StatusNoContent)
	return nil
}

func hijackServer(w http.ResponseWriter) (io.ReadCloser, io.Writer, error) {
	conn, _, err := w.(http.Hijacker).Hijack()
	if err != nil {
		return nil, nil, err
	}
	// Flush the options to make sure the client sets the raw mode
	conn.Write([]byte{})
	return conn, conn, nil
}

func postContainersAttach(out beam.Sender, version version.Version, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
	if err := r.ParseForm(); err != nil {
		return err
	}
	if vars == nil {
		return fmt.Errorf("Missing parameter")
	}

	inStream, outStream, err := hijackServer(w)
	if err != nil {
		return err
	}
	defer func() {
		if tcpc, ok := inStream.(*net.TCPConn); ok {
			tcpc.CloseWrite()
		} else {
			inStream.Close()
		}
	}()
	defer func() {
		if tcpc, ok := outStream.(*net.TCPConn); ok {
			tcpc.CloseWrite()
		} else if closer, ok := outStream.(io.Closer); ok {
			closer.Close()
		}
	}()

	fmt.Fprintf(outStream, "HTTP/1.1 200 OK\r\nContent-Type: application/vnd.docker.raw-stream\r\n\r\n")

	// TODO: if a TTY, then no multiplexing is done
	errStream := utils.NewStdWriter(outStream, utils.Stderr)
	outStream = utils.NewStdWriter(outStream, utils.Stdout)

	_, containerOut, err := beam.Obj(out).Attach(vars["name"])
	if err != nil {
		return err
	}
	container := beam.Obj(containerOut)

	containerR, _, err := container.Attach("")
	var tasks sync.WaitGroup
	go func() {
		defer tasks.Done()
		err := beam.DecodeStream(outStream, containerR, "stdout")
		if err != nil {
			fmt.Printf("decodestream: %v\n", err)
		}
	}()
	tasks.Add(1)
	go func() {
		defer tasks.Done()
		err := beam.DecodeStream(errStream, containerR, "stderr")
		if err != nil {
			fmt.Printf("decodestream: %v\n", err)
		}
	}()
	tasks.Add(1)
	tasks.Wait()

	return nil
}

func postContainersWait(out beam.Sender, version version.Version, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
	if vars == nil {
		return fmt.Errorf("Missing parameter")
	}

	// TODO: this should wait for container to end out output correct
	// exit status
	return writeJSON(w, http.StatusOK, map[string]interface{}{
		"StatusCode": "0",
	})
}

func createRouter(out beam.Sender) (*mux.Router, error) {
	r := mux.NewRouter()
	m := map[string]map[string]HttpApiFunc{
		"GET": {
			"/_ping":           ping,
			"/containers/json": getContainersJSON,
		},
		"POST": {
			"/containers/create":           postContainersCreate,
			"/containers/{name:.*}/attach": postContainersAttach,
			"/containers/{name:.*}/start":  postContainersStart,
			"/containers/{name:.*}/stop":   postContainersStop,
			"/containers/{name:.*}/wait":   postContainersWait,
		},
		"DELETE":  {},
		"OPTIONS": {},
	}

	for method, routes := range m {
		for route, fct := range routes {
			localRoute := route
			localFct := fct
			localMethod := method

			f := makeHttpHandler(out, localMethod, localRoute, localFct, version.Version("0.11.0"))

			// add the new route
			if localRoute == "" {
				r.Methods(localMethod).HandlerFunc(f)
			} else {
				r.Path("/v{version:[0-9.]+}" + localRoute).Methods(localMethod).HandlerFunc(f)
				r.Path(localRoute).Methods(localMethod).HandlerFunc(f)
			}
		}
	}

	return r, nil
}

func makeHttpHandler(out beam.Sender, localMethod string, localRoute string, handlerFunc HttpApiFunc, dockerVersion version.Version) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		// log the request
		fmt.Printf("Calling %s %s\n", localMethod, localRoute)

		version := version.Version(mux.Vars(r)["version"])
		if version == "" {
			version = api.APIVERSION
		}

		if err := handlerFunc(out, version, w, r, mux.Vars(r)); err != nil {
			fmt.Printf("Error: %s", err)
		}
	}
}

func writeJSON(w http.ResponseWriter, code int, v interface{}) error {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(code)
	enc := json.NewEncoder(w)
	return enc.Encode(v)
}
<commit_msg>Adds support for `docker inspect` to dockerserver<commit_after>package backends

import (
	"encoding/json"
	"fmt"
	"github.com/docker/libswarm/beam"
	"github.com/dotcloud/docker/api"
	"github.com/dotcloud/docker/pkg/version"
	dockerContainerConfig "github.com/dotcloud/docker/runconfig"
	"github.com/dotcloud/docker/utils"
	"github.com/gorilla/mux"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"time"
)

func DockerServer() beam.Sender {
	backend := beam.NewServer()
	backend.OnSpawn(beam.Handler(func(ctx *beam.Message) error {
		instance := beam.Task(func(in beam.Receiver, out beam.Sender) {
			url := "tcp://localhost:4243"
			if len(ctx.Args) > 0 {
				url = ctx.Args[0]
			}
			err := listenAndServe(url, out)
			if err != nil {
				fmt.Printf("listenAndServe: %v", err)
			}
		})
		_, err := ctx.Ret.Send(&beam.Message{Verb: beam.Ack, Ret: instance})
		return err
	}))
	return backend
}

type HttpApiFunc func(out beam.Sender, version version.Version, w http.ResponseWriter, r *http.Request, vars map[string]string) error

func listenAndServe(urlStr string, out beam.Sender) error {
	fmt.Println("Starting Docker server...")
	r, err := createRouter(out)
	if err != nil {
		return err
	}

	parsedUrl, err := url.Parse(urlStr)
	if err != nil {
		return err
	}

	var hostAndPath string
	// For Unix sockets we need to capture the path as well as the host
	if parsedUrl.Scheme == "unix" {
		hostAndPath = "/" + parsedUrl.Host + parsedUrl.Path
	} else {
		hostAndPath = parsedUrl.Host
	}

	l, err := net.Listen(parsedUrl.Scheme, hostAndPath)
	if err != nil {
		return err
	}
	httpSrv := http.Server{Addr: hostAndPath, Handler: r}
	return httpSrv.Serve(l)
}

func ping(out beam.Sender, version version.Version, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
	_, err := w.Write([]byte{'O', 'K'})
	return err
}

type containerJson struct {
	Config          *dockerContainerConfig.Config
	HostConfig      *dockerContainerConfig.HostConfig
	ID              string
	Created         string
	Driver          string
	ExecDriver      string
	NetworkSettings struct {
		Bridge      string
		Gateway     string
		IPAddress   string
		IPPrefixLen int
		Ports       map[string][]map[string]string
	}
	State struct {
		Running    bool
		StartedAt  string
		FinishedAt string
		ExitCode   int
		Paused     bool
		Pid        int
	}
	Name           string
	Args           []string
	ProcessLbael   string
	ResolvConfPath string
	HostnamePath   string
	HostsPath      string
	MountLabel     string
	Volumes        map[string]string
	VolumesRW      map[string]bool
}

func getContainerJson(out beam.Sender, containerID string) (containerJson, error) {
	o := beam.Obj(out)

	_, containerOut, err := o.Attach(containerID)
	if err != nil {
		return containerJson{}, err
	}
	container := beam.Obj(containerOut)
	responseJson, err := container.Get()
	if err != nil {
		return containerJson{}, err
	}
	var response containerJson

	if err = json.Unmarshal([]byte(responseJson), &response); err != nil {
		return containerJson{}, err
	}

	return response, nil
}

func getContainerInfo(out beam.Sender, version version.Version, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
	container, err := getContainerJson(out, vars["name"])
	if err != nil {
		return err
	}
	return writeJSON(w, http.StatusOK, container)
}

func getContainersJSON(out beam.Sender, version version.Version, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
	if err := r.ParseForm(); err != nil {
		return err
	}

	o := beam.Obj(out)
	names, err := o.Ls()
	if err != nil {
		return err
	}

	var responses []interface{}

	for _, name := range names {
		response, err := getContainerJson(out, name)
		if err != nil {
			return err
		}
		created, err := time.Parse(time.RFC3339, response.Created)
		if err != nil {
			return err
		}
		var state string
		if response.State.Running {
			state = "Up"
		} else {
			state = fmt.Sprintf("Exited (%d)", response.State.ExitCode)
		}
		type port struct {
			IP          string
			PrivatePort int
			PublicPort  int
			Type        string
		}
		var ports []port
		for p, mappings := range response.NetworkSettings.Ports {
			var portnum int
			var proto string
			_, err := fmt.Sscanf(p, "%d/%s", &portnum, &proto)
			if err != nil {
				return err
			}
			if len(mappings) > 0 {
				for _, mapping := range mappings {
					hostPort, err := strconv.Atoi(mapping["HostPort"])
					if err != nil {
						return err
					}
					newport := port{
						IP:          mapping["HostIp"],
						PrivatePort: portnum,
						PublicPort:  hostPort,
						Type:        proto,
					}
					ports = append(ports, newport)
				}
			} else {
				newport := port{
					PublicPort: portnum,
					Type:       proto,
				}
				ports = append(ports, newport)
			}
		}
		responses = append(responses, map[string]interface{}{
			"Id":      response.ID,
			"Command": strings.Join(response.Config.Cmd, " "),
			"Created": created.Unix(),
			"Image":   response.Config.Image,
			"Names":   []string{response.Name},
			"Ports":   ports,
			"Status":  state,
		})
	}

	return writeJSON(w, http.StatusOK, responses)
}

func postContainersCreate(out beam.Sender, version version.Version, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
	if err := r.ParseForm(); err != nil {
		return nil
	}

	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		return err
	}

	container, err := beam.Obj(out).Spawn(string(body))
	if err != nil {
		return err
	}

	responseJson, err := container.Get()
	if err != nil {
		return err
	}

	var response struct{ Id string }
	if err = json.Unmarshal([]byte(responseJson), &response); err != nil {
		return err
	}
	return writeJSON(w, http.StatusCreated, response)
}

func postContainersStart(out beam.Sender, version version.Version, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
	if vars == nil {
		return fmt.Errorf("Missing parameter")
	}

	// TODO: r.Body

	name := vars["name"]
	_, containerOut, err := beam.Obj(out).Attach(name)
	container := beam.Obj(containerOut)
	if err != nil {
		return err
	}
	if err := container.Start(); err != nil {
		return err
	}

	w.WriteHeader(http.StatusNoContent)
	return nil
}

func postContainersStop(out beam.Sender, version version.Version, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
	if vars == nil {
		return fmt.Errorf("Missing parameter")
	}

	name := vars["name"]
	_, containerOut, err := beam.Obj(out).Attach(name)
	container := beam.Obj(containerOut)
	if err != nil {
		return err
	}
	if err := container.Stop(); err != nil {
		return err
	}

	w.WriteHeader(http.StatusNoContent)
	return nil
}

func hijackServer(w http.ResponseWriter) (io.ReadCloser, io.Writer, error) {
	conn, _, err := w.(http.Hijacker).Hijack()
	if err != nil {
		return nil, nil, err
	}
	// Flush the options to make sure the client sets the raw mode
	conn.Write([]byte{})
	return conn, conn, nil
}

func postContainersAttach(out beam.Sender, version version.Version, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
	if err := r.ParseForm(); err != nil {
		return err
	}
	if vars == nil {
		return fmt.Errorf("Missing parameter")
	}

	inStream, outStream, err := hijackServer(w)
	if err != nil {
		return err
	}
	defer func() {
		if tcpc, ok := inStream.(*net.TCPConn); ok {
			tcpc.CloseWrite()
		} else {
			inStream.Close()
		}
	}()
	defer func() {
		if tcpc, ok := outStream.(*net.TCPConn); ok {
			tcpc.CloseWrite()
		} else if closer, ok := outStream.(io.Closer); ok {
			closer.Close()
		}
	}()

	fmt.Fprintf(outStream, "HTTP/1.1 200 OK\r\nContent-Type: application/vnd.docker.raw-stream\r\n\r\n")

	// TODO: if a TTY, then no multiplexing is done
	errStream := utils.NewStdWriter(outStream, utils.Stderr)
	outStream = utils.NewStdWriter(outStream, utils.Stdout)

	_, containerOut, err := beam.Obj(out).Attach(vars["name"])
	if err != nil {
		return err
	}
	container := beam.Obj(containerOut)

	containerR, _, err := container.Attach("")
	var tasks sync.WaitGroup
	go func() {
		defer tasks.Done()
		err := beam.DecodeStream(outStream, containerR, "stdout")
		if err != nil {
			fmt.Printf("decodestream: %v\n", err)
		}
	}()
	tasks.Add(1)
	go func() {
		defer tasks.Done()
		err := beam.DecodeStream(errStream, containerR, "stderr")
		if err != nil {
			fmt.Printf("decodestream: %v\n", err)
		}
	}()
	tasks.Add(1)
	tasks.Wait()

	return nil
}

func postContainersWait(out beam.Sender, version version.Version, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
	if vars == nil {
		return fmt.Errorf("Missing parameter")
	}

	// TODO: this should wait for container to end out output correct
	// exit status
	return writeJSON(w, http.StatusOK, map[string]interface{}{
		"StatusCode": "0",
	})
}

func createRouter(out beam.Sender) (*mux.Router, error) {
	r := mux.NewRouter()
	m := map[string]map[string]HttpApiFunc{
		"GET": {
			"/_ping":                  ping,
			"/containers/json":        getContainersJSON,
			"/containers/{name}/json": getContainerInfo,
		},
		"POST": {
			"/containers/create":           postContainersCreate,
			"/containers/{name:.*}/attach": postContainersAttach,
			"/containers/{name:.*}/start":  postContainersStart,
			"/containers/{name:.*}/stop":   postContainersStop,
			"/containers/{name:.*}/wait":   postContainersWait,
		},
		"DELETE":  {},
		"OPTIONS": {},
	}

	for method, routes := range m {
		for route, fct := range routes {
			localRoute := route
			localFct := fct
			localMethod := method

			f := makeHttpHandler(out, localMethod, localRoute, localFct, version.Version("0.11.0"))

			// add the new route
			if localRoute == "" {
				r.Methods(localMethod).HandlerFunc(f)
			} else {
				r.Path("/v{version:[0-9.]+}" + localRoute).Methods(localMethod).HandlerFunc(f)
				r.Path(localRoute).Methods(localMethod).HandlerFunc(f)
			}
		}
	}

	return r, nil
}

func makeHttpHandler(out beam.Sender, localMethod string, localRoute string, handlerFunc HttpApiFunc, dockerVersion version.Version) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		// log the request
		fmt.Printf("Calling %s %s\n", localMethod, localRoute)

		version := version.Version(mux.Vars(r)["version"])
		if version == "" {
			version = api.APIVERSION
		}

		if err := handlerFunc(out, version, w, r, mux.Vars(r)); err != nil {
			fmt.Printf("Error: %s", err)
		}
	}
}

func writeJSON(w http.ResponseWriter, code int, v interface{}) error {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(code)
	enc := json.NewEncoder(w)
	return enc.Encode(v)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"math/rand"

	"github.com/golang/protobuf/proto"

	pbd "github.com/brotherlogic/discovery/proto"
	pb "github.com/brotherlogic/gobuildmaster/proto"
	pbs "github.com/brotherlogic/gobuildslave/proto"
)

type getter interface {
	getSlaves() (*pbd.ServiceList, error)
	getJobs(*pbd.RegistryEntry) ([]*pbs.JobAssignment, error)
	getConfig(*pbd.RegistryEntry) ([]*pbs.Requirement, error)
}

type checker interface {
	assess(server string) (*pbs.JobList, *pbs.Config)
	discover() *pbd.ServiceList
	master(entry *pbd.RegistryEntry, master bool) (bool, error)
	setprev([]string)
	getprev() []string
}

func getFleetStatus(c checker) (map[string]*pbs.JobList, map[string]*pbs.Config) {
	resJ := make(map[string]*pbs.JobList)
	resC := make(map[string]*pbs.Config)

	curr := make([]string, 0)
	for _, service := range c.discover().Services {
		if service.Name == "gobuildslave" {
			curr = append(curr, service.Identifier)
			joblist, config := c.assess(service.Identifier)
			resJ[service.Identifier] = joblist
			resC[service.Identifier] = config
		}
	}

	c.setprev(curr)

	return resJ, resC
}

// Find the first available server
func chooseServer(job *pbs.JobSpec, c checker) string {
	services := c.discover().Services
	for _, i := range rand.Perm(len(services)) {
		service := services[i]
		if service.Name == "gobuildslave" && (job.GetServer() == "" || job.GetServer() == service.GetIdentifier()) {
			jobs, sc := c.assess(service.Identifier)

			//Don't accept a server which is already running this job
			jobfine := true
			for _, j := range jobs.Details {
				if j.Spec.Name == job.Name {
					jobfine = false
				}

			}
			if jobfine {
				if sc.Disk > job.Disk && (!job.External || sc.External) && (!job.GetCds() || sc.GetSupportsCds()) {
					return service.Identifier
				}
			}
		}
	}
	return ""
}

// Find the first available server
func (s *Server) selectServer(job *pbs.Job, g getter) string {
	services, _ := g.getSlaves()
	for _, i := range rand.Perm(len(services.Services)) {
		jobs, _ := g.getJobs(services.Services[i])
		//Don't accept a server which is already running this job
		jobfine := true
		for _, j := range jobs {
			if j.Job.Name == job.Name {
				jobfine = false
			}

		}
		if jobfine {
			requirements, err := g.getConfig(services.Services[i])
			if err == nil {
				allmatch := true
				for _, req := range job.Requirements {
					localmatch := false
					for _, r := range requirements {
						if r.Category == req.Category && r.Properties == req.Properties {
							localmatch = true
							s.Log(fmt.Sprintf("MATCH %v and %v for %v on %v", r, req, job.Name, services.Services[i].Identifier))
						}
					}

					if !localmatch {
						allmatch = false
					}
				}
				if allmatch {
					return services.Services[i].Identifier
				}
			}
		}
	}
	return ""
}

func configDiff(cm, cs *pb.Config) *pb.Config {
	retConfig := &pb.Config{}

	for _, entry := range cm.Intents {
		nIntent := &pb.Intent{}
		nIntent.Spec = entry.Spec
		nIntent.Count = entry.Count
		retConfig.Intents = append(retConfig.Intents, nIntent)
	}

	for _, entry := range cs.Intents {
		for _, pair := range retConfig.Intents {
			if entry.Spec.Name == pair.Spec.Name {
				pair.Count -= entry.Count
			}
		}
	}

	return retConfig
}

func loadConfig(f string) (*pb.Config, error) {
	toload := &pb.Config{}
	bytes, _ := ioutil.ReadFile(f)
	proto.UnmarshalText(string(bytes), toload)
	return toload, nil
}

func runJobs(c *pb.Config) []*pbs.JobSpec {
	var jobs []*pbs.JobSpec
	for _, j := range c.Intents {
		for i := 0; i < int(j.Count); i++ {
			jobs = append(jobs, j.Spec)
		}
	}

	return jobs
}
<commit_msg>More details<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"math/rand"

	"github.com/golang/protobuf/proto"

	pbd "github.com/brotherlogic/discovery/proto"
	pb "github.com/brotherlogic/gobuildmaster/proto"
	pbs "github.com/brotherlogic/gobuildslave/proto"
)

type getter interface {
	getSlaves() (*pbd.ServiceList, error)
	getJobs(*pbd.RegistryEntry) ([]*pbs.JobAssignment, error)
	getConfig(*pbd.RegistryEntry) ([]*pbs.Requirement, error)
}

type checker interface {
	assess(server string) (*pbs.JobList, *pbs.Config)
	discover() *pbd.ServiceList
	master(entry *pbd.RegistryEntry, master bool) (bool, error)
	setprev([]string)
	getprev() []string
}

func getFleetStatus(c checker) (map[string]*pbs.JobList, map[string]*pbs.Config) {
	resJ := make(map[string]*pbs.JobList)
	resC := make(map[string]*pbs.Config)

	curr := make([]string, 0)
	for _, service := range c.discover().Services {
		if service.Name == "gobuildslave" {
			curr = append(curr, service.Identifier)
			joblist, config := c.assess(service.Identifier)
			resJ[service.Identifier] = joblist
			resC[service.Identifier] = config
		}
	}

	c.setprev(curr)

	return resJ, resC
}

// Find the first available server
func chooseServer(job *pbs.JobSpec, c checker) string {
	services := c.discover().Services
	for _, i := range rand.Perm(len(services)) {
		service := services[i]
		if service.Name == "gobuildslave" && (job.GetServer() == "" || job.GetServer() == service.GetIdentifier()) {
			jobs, sc := c.assess(service.Identifier)

			//Don't accept a server which is already running this job
			jobfine := true
			for _, j := range jobs.Details {
				if j.Spec.Name == job.Name {
					jobfine = false
				}

			}
			if jobfine {
				if sc.Disk > job.Disk && (!job.External || sc.External) && (!job.GetCds() || sc.GetSupportsCds()) {
					return service.Identifier
				}
			}
		}
	}
	return ""
}

// Find the first available server
func (s *Server) selectServer(job *pbs.Job, g getter) string {
	services, _ := g.getSlaves()
	for _, i := range rand.Perm(len(services.Services)) {
		jobs, _ := g.getJobs(services.Services[i])
		//Don't accept a server which is already running this job
		jobfine := true
		for _, j := range jobs {
			if j.Job.Name == job.Name {
				jobfine = false
			}

		}
		if jobfine {
			requirements, err := g.getConfig(services.Services[i])
			if err == nil {
				allmatch := true
				for _, req := range job.Requirements {
					localmatch := false
					for _, r := range requirements {
						if r.Category == req.Category && r.Properties == req.Properties {
							localmatch = true
							s.Log(fmt.Sprintf("MATCH %v and %v for %v on %v", r, req, job.Name, services.Services[i].Identifier))
						}
					}

					if !localmatch {
						allmatch = false
					}
				}
				if allmatch {
					s.Log(fmt.Sprintf("MATCH %v and %v for %v on %v", job.Requirements, requirements, job.Name, services.Services[i].Identifier))
					return services.Services[i].Identifier
				}
			}
		}
	}
	return ""
}

func configDiff(cm, cs *pb.Config) *pb.Config {
	retConfig := &pb.Config{}

	for _, entry := range cm.Intents {
		nIntent := &pb.Intent{}
		nIntent.Spec = entry.Spec
		nIntent.Count = entry.Count
		retConfig.Intents = append(retConfig.Intents, nIntent)
	}

	for _, entry := range cs.Intents {
		for _, pair := range retConfig.Intents {
			if entry.Spec.Name == pair.Spec.Name {
				pair.Count -= entry.Count
			}
		}
	}

	return retConfig
}

func loadConfig(f string) (*pb.Config, error) {
	toload := &pb.Config{}
	bytes, _ := ioutil.ReadFile(f)
	proto.UnmarshalText(string(bytes), toload)
	return toload, nil
}

func runJobs(c *pb.Config) []*pbs.JobSpec {
	var jobs []*pbs.JobSpec
	for _, j := range c.Intents {
		for i := 0; i < int(j.Count); i++ {
			jobs = append(jobs, j.Spec)
		}
	}

	return jobs
}
<|endoftext|>


<commit_before>package cli

import (
	"errors"
	"io"
	"strings"
	"syscall"
	"testing"
	"time"

	"github.com/elves/elvish/cli/el/codearea"
	"github.com/elves/elvish/cli/el/layout"
	"github.com/elves/elvish/cli/term"
	"github.com/elves/elvish/edit/ui"
	"github.com/elves/elvish/styled"
	"github.com/elves/elvish/sys"
)

const (
	testTTYHeight = 24
	testTTYWidth  = 60
)

func TestReadCode_AbortsOnSetupError(t *testing.T) {
	a, tty := setup()

	setupErr := errors.New("a fake error")
	tty.SetSetup(func() {}, setupErr)

	_, err := a.ReadCode()

	if err != setupErr {
		t.Errorf("ReadCode returns error %v, want %v", err, setupErr)
	}
}

func TestReadCode_CallsRestore(t *testing.T) {
	a, tty := setup()

	restoreCalled := 0
	tty.SetSetup(func() { restoreCalled++ }, nil)
	tty.Inject(term.K('\n'))

	a.ReadCode()

	if restoreCalled != 1 {
		t.Errorf("Restore callback called %d times, want once", restoreCalled)
	}
}

func TestReadCode_ResetsStateBeforeReturn(t *testing.T) {
	a, tty := setupWithSpec(AppSpec{
		CodeAreaState: codearea.State{
			Buffer: codearea.Buffer{Content: "some code"}}})

	tty.Inject(term.K('\n'))

	a.ReadCode()

	if code := a.CodeArea().CopyState().Buffer.Content; code != "" {
		t.Errorf("Editor state has code %q, want empty", code)
	}
}

func TestReadCode_FinalRedrawWithRPrompt(t *testing.T) {
	a, tty := setupWithSpec(AppSpec{
		CodeAreaState: codearea.State{
			Buffer: codearea.Buffer{Content: "code"}},
		RPrompt:           constPrompt{styled.Plain("R")},
		RPromptPersistent: func() bool { return true },
	})

	tty.Inject(term.K('\n'))
	a.ReadCode()

	wantBuf := bb().
		WritePlain("code" + strings.Repeat(" ", testTTYWidth-5) + "R").
		Newline().SetDotToCursor(). // cursor on newline in final redraw
		Buffer()
	tty.TestBuffer(t, wantBuf)
}

func TestReadCode_FinalRedrawWithoutRPrompt(t *testing.T) {
	a, tty := setupWithSpec(AppSpec{
		CodeAreaState: codearea.State{
			Buffer: codearea.Buffer{Content: "code"}},
		RPrompt:           constPrompt{styled.Plain("R")},
		RPromptPersistent: func() bool { return false },
	})

	tty.Inject(term.K('\n'))
	a.ReadCode()

	wantBuf := bb().
		WritePlain("code").         // no rprompt
		Newline().SetDotToCursor(). // cursor on newline in final redraw
		Buffer()
	tty.TestBuffer(t, wantBuf)
}

func TestReadCode_PassesInputEventsToPopup(t *testing.T) {
	/*
		a, tty := setup()

		m := &fakeMode{maxKeys: 3}
		a.state.Raw.Mode = m
		tty.Inject(term.K('a'))
		tty.Inject(term.K('b'))
		tty.Inject(term.K('c'))

		a.ReadCode()

		wantKeysHandled := []ui.Key{
			{Rune: 'a'}, {Rune: 'b'}, {Rune: 'c'},
		}
		if !reflect.DeepEqual(m.keysHandled, wantKeysHandled) {
			t.Errorf("Mode gets keys %v, want %v", m.keysHandled, wantKeysHandled)
		}
	*/
}

func TestReadCode_PassesInputEventsToCodeArea(t *testing.T) {
}

func TestReadCode_CallsBeforeReadlineOnce(t *testing.T) {
	called := 0
	a, tty := setupWithSpec(AppSpec{
		BeforeReadline: []func(){func() { called++ }},
	})

	tty.Inject(term.K('\n'))
	a.ReadCode()

	if called != 1 {
		t.Errorf("BeforeReadline hook called %d times, want 1", called)
	}
}

func TestReadCode_CallsAfterReadlineOnceWithCode(t *testing.T) {
	called := 0
	code := ""
	a, tty := setupWithSpec(AppSpec{
		AfterReadline: []func(string){func(s string) {
			called++
			code = s
		}},
	})

	feedInput(tty, "abc\n")

	a.ReadCode()

	if called != 1 {
		t.Errorf("AfterReadline hook called %d times, want 1", called)
	}
	if code != "abc" {
		t.Errorf("AfterReadline hook called with %q, want %q", code, "abc")
	}
}

func TestReadCode_RespectsMaxHeight(t *testing.T) {
	a, tty := setupWithSpec(AppSpec{
		MaxHeight: func() int { return 2 },
		CodeAreaState: codearea.State{
			// The code needs 3 lines to completely show.
			Buffer: codearea.Buffer{Content: strings.Repeat("a", 15)}}})

	tty.SetSize(10, 5) // Width = 5 to make it easy to test

	codeCh, _ := ReadCodeAsync(a)

	wantBuf := ui.NewBufferBuilder(5).
		WritePlain(strings.Repeat("a", 10)).Buffer()
	tty.TestBuffer(t, wantBuf)

	cleanup(a, codeCh)
}

func TestReadCode_RendersHighlightedCode(t *testing.T) {
	a, tty := setupWithSpec(AppSpec{
		Highlighter: testHighlighter{
			get: func(code string) (styled.Text, []error) {
				return styled.MakeText(code, "red"), nil
			},
		}})

	feedInput(tty, "abc")

	codeCh, _ := ReadCodeAsync(a)

	wantBuf := bb().
		WriteStringSGR("abc", "31" /* SGR for red foreground */).
		SetDotToCursor().Buffer()
	tty.TestBuffer(t, wantBuf)

	cleanup(a, codeCh)
}

func TestReadCode_RendersErrorFromHighlighter(t *testing.T) {
	// TODO
}

func TestReadCode_RedrawsOnHighlighterLateUpdate(t *testing.T) {
	// TODO
}

func TestReadCode_RendersPrompt(t *testing.T) {
	a, tty := setupWithSpec(AppSpec{
		Prompt: constPrompt{styled.Plain("> ")}})

	codeCh, _ := ReadCodeAsync(a)
	defer cleanup(a, codeCh)

	tty.Inject(term.K('a'))

	wantBuf := bb().
		WritePlain("> a").
		SetDotToCursor().Buffer()
	tty.TestBuffer(t, wantBuf)
}

func TestReadCode_RendersRPrompt(t *testing.T) {
	a, tty := setupWithSpec(AppSpec{
		RPrompt: constPrompt{styled.Plain("R")}})

	codeCh, _ := ReadCodeAsync(a)
	defer cleanup(a, codeCh)

	tty.Inject(term.K('a'))

	wantBuf := bb().
		WritePlain("a").SetDotToCursor().
		WritePlain(strings.Repeat(" ", testTTYWidth-2)).
		WritePlain("R").Buffer()
	tty.TestBuffer(t, wantBuf)
}

func TestReadCode_TriggersPrompt(t *testing.T) {
	called := 0
	a, _ := setupWithSpec(AppSpec{
		Prompt: testPrompt{trigger: func(bool) { called++ }}})

	codeCh, _ := ReadCodeAsync(a)
	cleanup(a, codeCh)

	if called != 1 {
		t.Errorf("Prompt.Trigger called %d times, want once", called)
	}
}

func TestReadCode_RedrawsOnPromptLateUpdate(t *testing.T) {
	promptContent := "old"
	prompt := testPrompt{
		get:         func() styled.Text { return styled.Plain(promptContent) },
		lateUpdates: make(chan styled.Text),
	}

	a, tty := setupWithSpec(AppSpec{Prompt: prompt})

	codeCh, _ := ReadCodeAsync(a)
	defer cleanup(a, codeCh)

	bufOldPrompt := bb().
		WritePlain("old").SetDotToCursor().Buffer()
	// Wait until old prompt is rendered
	tty.TestBuffer(t, bufOldPrompt)

	promptContent = "new"
	prompt.lateUpdates <- nil
	bufNewPrompt := bb().
		WritePlain("new").SetDotToCursor().Buffer()
	tty.TestBuffer(t, bufNewPrompt)
}

func TestReadCode_DrawsAndFlushesNotes(t *testing.T) {
	a, tty := setup()

	codeCh, _ := ReadCodeAsync(a)

	// Sanity-check initial state.
	initBuf := bb().Buffer()
	tty.TestBuffer(t, initBuf)

	a.Notify("note")

	wantNotesBuf := bb().WritePlain("note").Buffer()
	tty.TestNotesBuffer(t, wantNotesBuf)

	if n := len(a.CopyState().Notes); n > 0 {
		t.Errorf("State.Notes has %d elements after redrawing, want 0", n)
	}

	cleanup(a, codeCh)
}

func TestReadCode_FinalRedraw(t *testing.T) {
	a, tty := setupWithSpec(AppSpec{
		CodeAreaState: codearea.State{
			Buffer: codearea.Buffer{Content: "some code"}},
		State: State{
			Addon: layout.Label{Content: styled.Plain("addon")}}})

	codeCh, _ := ReadCodeAsync(a)

	// Wait until the initial draw, to ensure that we are indeed observing a
	// different state later.
	wantBuf := bb().
		WritePlain("some code").
		Newline().SetDotToCursor().WritePlain("addon").Buffer()
	tty.TestBuffer(t, wantBuf)

	cleanup(a, codeCh)

	// Final redraw hides the addon, and puts the cursor on a new line.
	wantFinalBuf := bb().
		WritePlain("some code").Newline().SetDotToCursor().Buffer()
	tty.TestBuffer(t, wantFinalBuf)
}

func TestReadCode_QuitsOnSIGHUP(t *testing.T) {
	a, tty := setup()

	tty.Inject(term.K('a'))

	codeCh, errCh := ReadCodeAsync(a)

	wantBuf := bb().WritePlain("a").
		SetDotToCursor().Buffer()
	tty.TestBuffer(t, wantBuf)

	tty.InjectSignal(syscall.SIGHUP)

	select {
	case <-codeCh:
		err := <-errCh
		if err != io.EOF {
			t.Errorf("want ReadCode to return io.EOF on SIGHUP, got %v", err)
		}
	case <-time.After(time.Second):
		t.Errorf("SIGHUP did not cause ReadCode to return")
	}
}

func TestReadCode_ResetsOnSIGINT(t *testing.T) {
	a, tty := setup()

	tty.Inject(term.K('a'))

	codeCh, _ := ReadCodeAsync(a)
	wantBuf := bb().WritePlain("a").
		SetDotToCursor().Buffer()
	tty.TestBuffer(t, wantBuf)

	tty.InjectSignal(syscall.SIGINT)

	wantBuf = bb().Buffer()
	tty.TestBuffer(t, wantBuf)

	cleanup(a, codeCh)
}

func TestReadCode_RedrawsOnSIGWINCH(t *testing.T) {
	content := "1234567890"
	a, tty := setupWithSpec(AppSpec{
		CodeAreaState: codearea.State{
			Buffer: codearea.Buffer{Content: content, Dot: len(content)}}})

	codeCh, _ := ReadCodeAsync(a)
	defer cleanup(a, codeCh)

	wantBuf := bb().WritePlain("1234567890").
		SetDotToCursor().Buffer()
	tty.TestBuffer(t, wantBuf)

	tty.SetSize(24, 4)
	tty.InjectSignal(sys.SIGWINCH)

	wantBuf = ui.NewBufferBuilder(4).WritePlain("1234567890").
		SetDotToCursor().Buffer()
	tty.TestBuffer(t, wantBuf)
}

func setup() (App, TTYCtrl) {
	return setupWithSpec(AppSpec{})
}

func setupWithSpec(spec AppSpec) (App, TTYCtrl) {
	tty, ttyControl := NewFakeTTY()
	ttyControl.SetSize(testTTYHeight, testTTYWidth)
	spec.TTY = tty
	a := NewApp(spec)
	return a, ttyControl
}

func bb() *ui.BufferBuilder {
	return ui.NewBufferBuilder(testTTYWidth)
}

func cleanup(a App, codeCh <-chan string) {
	a.CommitEOF()
	// Make sure that ReadCode has exited
	<-codeCh
}

func feedInput(ttyCtrl TTYCtrl, input string) {
	for _, r := range input {
		ttyCtrl.Inject(term.K(r))
	}
}
<commit_msg>cli/app_test.go: Organize tests.<commit_after>package cli

import (
	"errors"
	"io"
	"reflect"
	"strings"
	"syscall"
	"testing"
	"time"

	"github.com/elves/elvish/cli/el/codearea"
	"github.com/elves/elvish/cli/el/layout"
	"github.com/elves/elvish/cli/term"
	"github.com/elves/elvish/edit/ui"
	"github.com/elves/elvish/styled"
	"github.com/elves/elvish/sys"
)

const (
	testTTYHeight = 24
	testTTYWidth  = 60
)

// Lifecycle aspects.

func TestSetup_ErrorAbortsReadCode(t *testing.T) {
	a, tty := setup()
	setupErr := errors.New("a fake error")
	tty.SetSetup(func() {}, setupErr)

	_, err := a.ReadCode()

	if err != setupErr {
		t.Errorf("ReadCode returns error %v, want %v", err, setupErr)
	}
}

func TestSetup_RestoreIsCalled(t *testing.T) {
	a, tty := setup()
	restoreCalled := 0
	tty.SetSetup(func() { restoreCalled++ }, nil)

	tty.Inject(term.K('\n'))
	a.ReadCode()

	if restoreCalled != 1 {
		t.Errorf("Restore callback called %d times, want once", restoreCalled)
	}
}

func TestState_IsResetBeforeReadCodeReturns(t *testing.T) {
	a, tty := setupWithSpec(AppSpec{
		CodeAreaState: codearea.State{
			Buffer: codearea.Buffer{Content: "some code"}}})

	tty.Inject(term.K('\n'))
	a.ReadCode()

	if code := a.CodeArea().CopyState().Buffer.Content; code != "" {
		t.Errorf("Editor state has code %q, want empty", code)
	}
}

func TestBeforeReadline(t *testing.T) {
	called := 0
	a, tty := setupWithSpec(AppSpec{
		BeforeReadline: []func(){func() { called++ }},
	})

	tty.Inject(term.K('\n'))
	a.ReadCode()

	if called != 1 {
		t.Errorf("BeforeReadline hook called %d times, want 1", called)
	}
}

func TestAfterReadline(t *testing.T) {
	calledWith := []string{}
	a, tty := setupWithSpec(AppSpec{
		AfterReadline: []func(string){func(s string) {
			calledWith = append(calledWith, s)
		}},
	})

	feedInput(tty, "abc\n")

	a.ReadCode()

	wantCalledWith := []string{"abc"}
	if !reflect.DeepEqual(calledWith, wantCalledWith) {
		t.Errorf("AfterReadline hook called with %v, want %v", calledWith, wantCalledWith)
	}
}

func TestFinalRedraw(t *testing.T) {
	a, tty := setupWithSpec(AppSpec{
		CodeAreaState: codearea.State{
			Buffer: codearea.Buffer{Content: "code"}},
		State: State{
			Addon: layout.Label{Content: styled.Plain("addon")}}})
	codeCh, _ := ReadCodeAsync(a)

	// Wait until the stable state.
	wantBuf := bb().
		WritePlain("code").
		Newline().SetDotToCursor().WritePlain("addon").Buffer()
	tty.TestBuffer(t, wantBuf)

	cleanup(a, codeCh)

	// Final redraw hides the addon, and puts the cursor on a new line.
	wantFinalBuf := bb().
		WritePlain("code").Newline().SetDotToCursor().Buffer()
	tty.TestBuffer(t, wantFinalBuf)
}

// Signals.

func TestSIGHUP_ReturnsEOF(t *testing.T) {
	a, tty := setup()

	tty.Inject(term.K('a'))

	_, errCh := ReadCodeAsync(a)
	// Wait until the initial redraw.
	tty.TestBuffer(t, bb().WritePlain("a").SetDotToCursor().Buffer())

	tty.InjectSignal(syscall.SIGHUP)

	select {
	case err := <-errCh:
		if err != io.EOF {
			t.Errorf("want ReadCode to return io.EOF on SIGHUP, got %v", err)
		}
	case <-time.After(time.Second):
		t.Errorf("SIGHUP did not cause ReadCode to return")
	}
}

func TestSIGINT_ResetsState(t *testing.T) {
	a, tty := setup()

	codeCh, _ := ReadCodeAsync(a)
	defer cleanup(a, codeCh)
	// Ensure that the terminal shows an non-empty state.
	feedInput(tty, "code")
	tty.TestBuffer(t, bb().WritePlain("code").SetDotToCursor().Buffer())

	tty.InjectSignal(syscall.SIGINT)

	// Verify that the state has now reset.
	tty.TestBuffer(t, bb().Buffer())
}

func TestSIGWINCH_TriggersRedraw(t *testing.T) {
	a, tty := setup()
	codeCh, _ := ReadCodeAsync(a)
	defer cleanup(a, codeCh)

	// Ensure that the terminal shows the input with the intial width.
	feedInput(tty, "1234567890")
	tty.TestBuffer(t, bb().WritePlain("1234567890").SetDotToCursor().Buffer())

	// Emulate a window size change.
	tty.SetSize(24, 4)
	tty.InjectSignal(sys.SIGWINCH)

	// Test that the editor has redrawn using the new width.
	tty.TestBuffer(t, ui.NewBufferBuilder(4).
		WritePlain("1234567890").SetDotToCursor().Buffer())
}

// Code area.

func TestEvent_HandledByCodeArea(t *testing.T) {
	a, tty := setup()
	codeCh, _ := ReadCodeAsync(a)
	defer cleanup(a, codeCh)

	feedInput(tty, "code")
	tty.TestBuffer(t, bb().WritePlain("code").SetDotToCursor().Buffer())
}

func TestHighlighter(t *testing.T) {
	a, tty := setupWithSpec(AppSpec{
		Highlighter: testHighlighter{
			get: func(code string) (styled.Text, []error) {
				return styled.MakeText(code, "red"), nil
			},
		}})

	codeCh, _ := ReadCodeAsync(a)
	defer cleanup(a, codeCh)
	feedInput(tty, "abc")

	wantBuf := bb().
		WriteStyled(styled.MakeText("abc", "red")).
		SetDotToCursor().Buffer()
	tty.TestBuffer(t, wantBuf)
}

func TestHighlighter_Errors(t *testing.T) {
	// TODO
}

func TestHighlighter_LateUpdate(t *testing.T) {
	// TODO
}

func TestPrompt(t *testing.T) {
	a, tty := setupWithSpec(AppSpec{
		Prompt: constPrompt{styled.Plain("> ")}})

	codeCh, _ := ReadCodeAsync(a)
	defer cleanup(a, codeCh)

	tty.Inject(term.K('a'))

	tty.TestBuffer(t, bb().WritePlain("> a").SetDotToCursor().Buffer())
}

func TestPrompt_Trigger(t *testing.T) {
	called := 0
	a, _ := setupWithSpec(AppSpec{
		Prompt: testPrompt{trigger: func(bool) { called++ }}})

	codeCh, _ := ReadCodeAsync(a)
	cleanup(a, codeCh)

	if called != 1 {
		t.Errorf("Prompt.Trigger called %d times, want once", called)
	}
}

func TestPrompt_LateUpdate(t *testing.T) {
	promptContent := "old"
	prompt := testPrompt{
		get:         func() styled.Text { return styled.Plain(promptContent) },
		lateUpdates: make(chan styled.Text),
	}
	a, tty := setupWithSpec(AppSpec{Prompt: prompt})

	codeCh, _ := ReadCodeAsync(a)
	defer cleanup(a, codeCh)

	// Wait until old prompt is rendered
	tty.TestBuffer(t, bb().WritePlain("old").SetDotToCursor().Buffer())

	promptContent = "new"
	prompt.lateUpdates <- nil
	tty.TestBuffer(t, bb().WritePlain("new").SetDotToCursor().Buffer())
}

func TestRPrompt(t *testing.T) {
	a, tty := setupWithSpec(AppSpec{
		RPrompt: constPrompt{styled.Plain("R")}})

	codeCh, _ := ReadCodeAsync(a)
	defer cleanup(a, codeCh)

	tty.Inject(term.K('a'))

	wantBuf := bb().
		WritePlain("a").SetDotToCursor().
		WritePlain(strings.Repeat(" ", testTTYWidth-2)).
		WritePlain("R").Buffer()
	tty.TestBuffer(t, wantBuf)
}

func TestRPrompt_Persistent(t *testing.T) {
	a, tty := setupWithSpec(AppSpec{
		CodeAreaState: codearea.State{
			Buffer: codearea.Buffer{Content: "code"}},
		RPrompt:           constPrompt{styled.Plain("R")},
		RPromptPersistent: func() bool { return true },
	})

	tty.Inject(term.K('\n'))
	a.ReadCode()

	wantBuf := bb().
		WritePlain("code" + strings.Repeat(" ", testTTYWidth-5) + "R").
		Newline().SetDotToCursor(). // cursor on newline in final redraw
		Buffer()
	tty.TestBuffer(t, wantBuf)
}

func TestRPrompt_NotPersistent(t *testing.T) {
	a, tty := setupWithSpec(AppSpec{
		CodeAreaState: codearea.State{
			Buffer: codearea.Buffer{Content: "code"}},
		RPrompt:           constPrompt{styled.Plain("R")},
		RPromptPersistent: func() bool { return false },
	})

	tty.Inject(term.K('\n'))
	a.ReadCode()

	wantBuf := bb().
		WritePlain("code").         // no rprompt
		Newline().SetDotToCursor(). // cursor on newline in final redraw
		Buffer()
	tty.TestBuffer(t, wantBuf)
}

// Addon.

func TestReadCode_PassesInputEventsToAddon(t *testing.T) {
	/*
		a, tty := setup()

		m := &fakeMode{maxKeys: 3}
		a.state.Raw.Mode = m
		tty.Inject(term.K('a'))
		tty.Inject(term.K('b'))
		tty.Inject(term.K('c'))

		a.ReadCode()

		wantKeysHandled := []ui.Key{
			{Rune: 'a'}, {Rune: 'b'}, {Rune: 'c'},
		}
		if !reflect.DeepEqual(m.keysHandled, wantKeysHandled) {
			t.Errorf("Mode gets keys %v, want %v", m.keysHandled, wantKeysHandled)
		}
	*/
}

// Misc features.

func TestMaxHeight(t *testing.T) {
	a, tty := setupWithSpec(AppSpec{
		MaxHeight: func() int { return 2 },
		CodeAreaState: codearea.State{
			// The code needs 3 lines to completely show.
			Buffer: codearea.Buffer{Content: strings.Repeat("a", 15)}}})
	tty.SetSize(10, 5) // Width = 5 to make it easy to test
	codeCh, _ := ReadCodeAsync(a)
	defer cleanup(a, codeCh)

	wantBuf := ui.NewBufferBuilder(5).
		WritePlain(strings.Repeat("a", 10)). // Only show 2 lines due to MaxHeight.
		Buffer()
	tty.TestBuffer(t, wantBuf)
}

func TestNotes(t *testing.T) {
	a, tty := setup()
	codeCh, _ := ReadCodeAsync(a)
	defer cleanup(a, codeCh)

	// Wait until initial draw.
	tty.TestBuffer(t, bb().Buffer())

	a.Notify("note")

	// Test that the note is rendered onto the notes buffer.
	wantNotesBuf := bb().
		WritePlain("note").Buffer()
	tty.TestNotesBuffer(t, wantNotesBuf)

	// Test that notes are flushed after being rendered.
	if n := len(a.CopyState().Notes); n > 0 {
		t.Errorf("State.Notes has %d elements after redrawing, want 0", n)
	}
}

// Test utilities.

func setup() (App, TTYCtrl) {
	return setupWithSpec(AppSpec{})
}

func setupWithSpec(spec AppSpec) (App, TTYCtrl) {
	tty, ttyControl := NewFakeTTY()
	ttyControl.SetSize(testTTYHeight, testTTYWidth)
	spec.TTY = tty
	a := NewApp(spec)
	return a, ttyControl
}

func bb() *ui.BufferBuilder {
	return ui.NewBufferBuilder(testTTYWidth)
}

func cleanup(a App, codeCh <-chan string) {
	a.CommitEOF()
	// Make sure that ReadCode has exited
	<-codeCh
}

func feedInput(ttyCtrl TTYCtrl, input string) {
	for _, r := range input {
		ttyCtrl.Inject(term.K(r))
	}
}
<|endoftext|>


<commit_before>/*
   Copyright 2020 Docker, Inc.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package cmd

import (
	"context"

	"github.com/containerd/console"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"

	"github.com/docker/api/client"
	"github.com/docker/api/containers"
)

type logsOpts struct {
	Follow bool
	Tail   string
}

// LogsCommand fetches and shows logs of a container
func LogsCommand() *cobra.Command {
	var opts logsOpts
	cmd := &cobra.Command{
		Use:   "logs",
		Short: "Fetch the logs of a container",
		Args:  cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			return runLogs(cmd.Context(), args[0], opts)
		},
	}

	cmd.Flags().BoolVarP(&opts.Follow, "follow", "f", false, "Follow log outut")
	cmd.Flags().StringVar(&opts.Tail, "tail", "all", "Number of lines to show from the end of the logs")

	return cmd
}

func runLogs(ctx context.Context, containerName string, opts logsOpts) error {
	c, err := client.New(ctx)
	if err != nil {
		return errors.Wrap(err, "cannot connect to backend")
	}

	req := containers.LogsRequest{
		Follow: opts.Follow,
		Tail:   opts.Tail,
		Writer: console.Current(),
	}

	return c.ContainerService().Logs(ctx, containerName, req)
}
<commit_msg>Make sure we can have a console<commit_after>/*
   Copyright 2020 Docker, Inc.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package cmd

import (
	"context"
	"io"
	"os"

	"github.com/containerd/console"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"

	"github.com/docker/api/client"
	"github.com/docker/api/containers"
)

type logsOpts struct {
	Follow bool
	Tail   string
}

// LogsCommand fetches and shows logs of a container
func LogsCommand() *cobra.Command {
	var opts logsOpts
	cmd := &cobra.Command{
		Use:   "logs",
		Short: "Fetch the logs of a container",
		Args:  cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			return runLogs(cmd.Context(), args[0], opts)
		},
	}

	cmd.Flags().BoolVarP(&opts.Follow, "follow", "f", false, "Follow log outut")
	cmd.Flags().StringVar(&opts.Tail, "tail", "all", "Number of lines to show from the end of the logs")

	return cmd
}

func runLogs(ctx context.Context, containerName string, opts logsOpts) error {
	c, err := client.New(ctx)
	if err != nil {
		return errors.Wrap(err, "cannot connect to backend")
	}
	var con io.Writer

	con = os.Stdout
	if c, err := console.ConsoleFromFile(os.Stdout); err == nil {
		con = c
	}

	req := containers.LogsRequest{
		Follow: opts.Follow,
		Tail:   opts.Tail,
		Writer: con,
	}

	return c.ContainerService().Logs(ctx, containerName, req)
}
<|endoftext|>


<commit_before>package melody

import (
	"errors"
	"net/http"
	"sync"

	"github.com/gorilla/websocket"
)

// Close codes defined in RFC 6455, section 11.7.
// Duplicate of codes from gorilla/websocket for convenience.
const (
	CloseNormalClosure           = 1000
	CloseGoingAway               = 1001
	CloseProtocolError           = 1002
	CloseUnsupportedData         = 1003
	CloseNoStatusReceived        = 1005
	CloseAbnormalClosure         = 1006
	CloseInvalidFramePayloadData = 1007
	ClosePolicyViolation         = 1008
	CloseMessageTooBig           = 1009
	CloseMandatoryExtension      = 1010
	CloseInternalServerErr       = 1011
	CloseServiceRestart          = 1012
	CloseTryAgainLater           = 1013
	CloseTLSHandshake            = 1015
)

// Duplicate of codes from gorilla/websocket for convenience.
var validReceivedCloseCodes = map[int]bool{
	// see http://www.iana.org/assignments/websocket/websocket.xhtml#close-code-number

	CloseNormalClosure:           true,
	CloseGoingAway:               true,
	CloseProtocolError:           true,
	CloseUnsupportedData:         true,
	CloseNoStatusReceived:        false,
	CloseAbnormalClosure:         false,
	CloseInvalidFramePayloadData: true,
	ClosePolicyViolation:         true,
	CloseMessageTooBig:           true,
	CloseMandatoryExtension:      true,
	CloseInternalServerErr:       true,
	CloseServiceRestart:          true,
	CloseTryAgainLater:           true,
	CloseTLSHandshake:            false,
}

type handleMessageFunc func(*Session, []byte)
type handleErrorFunc func(*Session, error)
type handleCloseFunc func(*Session, int, string) error
type handleSessionFunc func(*Session)
type filterFunc func(*Session) bool

// Melody implements a websocket manager.
type Melody struct {
	Config                   *Config
	Upgrader                 *websocket.Upgrader
	messageHandler           handleMessageFunc
	messageHandlerBinary     handleMessageFunc
	messageSentHandler       handleMessageFunc
	messageSentHandlerBinary handleMessageFunc
	errorHandler             handleErrorFunc
	closeHandler             handleCloseFunc
	connectHandler           handleSessionFunc
	disconnectHandler        handleSessionFunc
	pongHandler              handleSessionFunc
	hub                      *hub
}

// New creates a new melody instance with default Upgrader and Config.
func New() *Melody {
	upgrader := &websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
		CheckOrigin:     func(r *http.Request) bool { return true },
	}

	hub := newHub()

	go hub.run()

	return &Melody{
		Config:                   newConfig(),
		Upgrader:                 upgrader,
		messageHandler:           func(*Session, []byte) {},
		messageHandlerBinary:     func(*Session, []byte) {},
		messageSentHandler:       func(*Session, []byte) {},
		messageSentHandlerBinary: func(*Session, []byte) {},
		errorHandler:             func(*Session, error) {},
		closeHandler:             nil,
		connectHandler:           func(*Session) {},
		disconnectHandler:        func(*Session) {},
		pongHandler:              func(*Session) {},
		hub:                      hub,
	}
}

// HandleConnect fires fn when a session connects.
func (m *Melody) HandleConnect(fn func(*Session)) {
	m.connectHandler = fn
}

// HandleDisconnect fires fn when a session disconnects.
func (m *Melody) HandleDisconnect(fn func(*Session)) {
	m.disconnectHandler = fn
}

// HandlePong fires fn when a pong is received from a session.
func (m *Melody) HandlePong(fn func(*Session)) {
	m.pongHandler = fn
}

// HandleMessage fires fn when a text message comes in.
func (m *Melody) HandleMessage(fn func(*Session, []byte)) {
	m.messageHandler = fn
}

// HandleMessageBinary fires fn when a binary message comes in.
func (m *Melody) HandleMessageBinary(fn func(*Session, []byte)) {
	m.messageHandlerBinary = fn
}

// HandleSentMessage fires fn when a text message is successfully sent.
func (m *Melody) HandleSentMessage(fn func(*Session, []byte)) {
	m.messageSentHandler = fn
}

// HandleSentMessageBinary fires fn when a binary message is successfully sent.
func (m *Melody) HandleSentMessageBinary(fn func(*Session, []byte)) {
	m.messageSentHandlerBinary = fn
}

// HandleError fires fn when a session has an error.
func (m *Melody) HandleError(fn func(*Session, error)) {
	m.errorHandler = fn
}

// HandleClose sets the handler for close messages received from the session.
// The code argument to h is the received close code or CloseNoStatusReceived
// if the close message is empty. The default close handler sends a close frame
// back to the session.
//
// The application must read the connection to process close messages as
// described in the section on Control Frames above.
//
// The connection read methods return a CloseError when a close frame is
// received. Most applications should handle close messages as part of their
// normal error handling. Applications should only set a close handler when the
// application must perform some action before sending a close frame back to
// the session.
func (m *Melody) HandleClose(fn func(*Session, int, string) error) {
	if fn != nil {
		m.closeHandler = fn
	}
}

// HandleRequest upgrades http requests to websocket connections and dispatches them to be handled by the melody instance.
func (m *Melody) HandleRequest(w http.ResponseWriter, r *http.Request) error {
	return m.HandleRequestWithKeys(w, r, nil)
}

// HandleRequestWithKeys does the same as HandleRequest but populates session.Keys with keys.
func (m *Melody) HandleRequestWithKeys(w http.ResponseWriter, r *http.Request, keys map[string]interface{}) error {
	if m.hub.closed() {
		return errors.New("melody instance is closed")
	}

	conn, err := m.Upgrader.Upgrade(w, r, nil)

	if err != nil {
		return err
	}

	session := &Session{
		Request: r,
		Keys:    keys,
		conn:    conn,
		output:  make(chan *envelope, m.Config.MessageBufferSize),
		melody:  m,
		open:    true,
		rwmutex: &sync.RWMutex{},
	}

	m.hub.register <- session

	m.connectHandler(session)

	go session.writePump()

	session.readPump()

	if !m.hub.closed() {
		m.hub.unregister <- session
	}

	session.close()

	m.disconnectHandler(session)

	return nil
}

// Broadcast broadcasts a text message to all sessions.
func (m *Melody) Broadcast(msg []byte) error {
	if m.hub.closed() {
		return errors.New("melody instance is closed")
	}

	message := &envelope{t: websocket.TextMessage, msg: msg}
	m.hub.broadcast <- message

	return nil
}

// BroadcastFilter broadcasts a text message to all sessions that fn returns true for.
func (m *Melody) BroadcastFilter(msg []byte, fn func(*Session) bool) error {
	if m.hub.closed() {
		return errors.New("melody instance is closed")
	}

	message := &envelope{t: websocket.TextMessage, msg: msg, filter: fn}
	m.hub.broadcast <- message

	return nil
}

// BroadcastOthers broadcasts a text message to all sessions except session s.
func (m *Melody) BroadcastOthers(msg []byte, s *Session) error {
	return m.BroadcastFilter(msg, func(q *Session) bool {
		return s != q
	})
}

// BroadcastMultiple broadcasts a text message to multiple sessions given in the sessions slice.
func (m *Melody) BroadcastMultiple(msg []byte, sessions []*Session) error {
	for _, sess := range sessions {
		if writeErr := sess.Write(msg); writeErr != nil {
			return writeErr
		}
	}
	return nil
}

// BroadcastBinary broadcasts a binary message to all sessions.
func (m *Melody) BroadcastBinary(msg []byte) error {
	if m.hub.closed() {
		return errors.New("melody instance is closed")
	}

	message := &envelope{t: websocket.BinaryMessage, msg: msg}
	m.hub.broadcast <- message

	return nil
}

// BroadcastBinaryFilter broadcasts a binary message to all sessions that fn returns true for.
func (m *Melody) BroadcastBinaryFilter(msg []byte, fn func(*Session) bool) error {
	if m.hub.closed() {
		return errors.New("melody instance is closed")
	}

	message := &envelope{t: websocket.BinaryMessage, msg: msg, filter: fn}
	m.hub.broadcast <- message

	return nil
}

// BroadcastBinaryOthers broadcasts a binary message to all sessions except session s.
func (m *Melody) BroadcastBinaryOthers(msg []byte, s *Session) error {
	return m.BroadcastBinaryFilter(msg, func(q *Session) bool {
		return s != q
	})
}

// Close closes the melody instance and all connected sessions.
func (m *Melody) Close() error {
	if m.hub.closed() {
		return errors.New("melody instance is already closed")
	}

	m.hub.exit <- &envelope{t: websocket.CloseMessage, msg: []byte{}}

	return nil
}

// CloseWithMsg closes the melody instance with the given close payload and all connected sessions.
// Use the FormatCloseMessage function to format a proper close message payload.
func (m *Melody) CloseWithMsg(msg []byte) error {
	if m.hub.closed() {
		return errors.New("melody instance is already closed")
	}

	m.hub.exit <- &envelope{t: websocket.CloseMessage, msg: msg}

	return nil
}

// Len return the number of connected sessions.
func (m *Melody) Len() int {
	return m.hub.len()
}

// IsClosed returns the status of the melody instance.
func (m *Melody) IsClosed() bool {
	return m.hub.closed()
}

// FormatCloseMessage formats closeCode and text as a WebSocket close message.
func FormatCloseMessage(closeCode int, text string) []byte {
	return websocket.FormatCloseMessage(closeCode, text)
}
<commit_msg>Support for the subprotocol<commit_after>package melody

import (
	"errors"
	"net/http"
	"sync"

	"github.com/gorilla/websocket"
)

// Close codes defined in RFC 6455, section 11.7.
// Duplicate of codes from gorilla/websocket for convenience.
const (
	CloseNormalClosure           = 1000
	CloseGoingAway               = 1001
	CloseProtocolError           = 1002
	CloseUnsupportedData         = 1003
	CloseNoStatusReceived        = 1005
	CloseAbnormalClosure         = 1006
	CloseInvalidFramePayloadData = 1007
	ClosePolicyViolation         = 1008
	CloseMessageTooBig           = 1009
	CloseMandatoryExtension      = 1010
	CloseInternalServerErr       = 1011
	CloseServiceRestart          = 1012
	CloseTryAgainLater           = 1013
	CloseTLSHandshake            = 1015
)

// Duplicate of codes from gorilla/websocket for convenience.
var validReceivedCloseCodes = map[int]bool{
	// see http://www.iana.org/assignments/websocket/websocket.xhtml#close-code-number

	CloseNormalClosure:           true,
	CloseGoingAway:               true,
	CloseProtocolError:           true,
	CloseUnsupportedData:         true,
	CloseNoStatusReceived:        false,
	CloseAbnormalClosure:         false,
	CloseInvalidFramePayloadData: true,
	ClosePolicyViolation:         true,
	CloseMessageTooBig:           true,
	CloseMandatoryExtension:      true,
	CloseInternalServerErr:       true,
	CloseServiceRestart:          true,
	CloseTryAgainLater:           true,
	CloseTLSHandshake:            false,
}

type handleMessageFunc func(*Session, []byte)
type handleErrorFunc func(*Session, error)
type handleCloseFunc func(*Session, int, string) error
type handleSessionFunc func(*Session)
type filterFunc func(*Session) bool

// Melody implements a websocket manager.
type Melody struct {
	Config                   *Config
	Upgrader                 *websocket.Upgrader
	messageHandler           handleMessageFunc
	messageHandlerBinary     handleMessageFunc
	messageSentHandler       handleMessageFunc
	messageSentHandlerBinary handleMessageFunc
	errorHandler             handleErrorFunc
	closeHandler             handleCloseFunc
	connectHandler           handleSessionFunc
	disconnectHandler        handleSessionFunc
	pongHandler              handleSessionFunc
	hub                      *hub
}

// New creates a new melody instance with default Upgrader and Config.
func New() *Melody {
	upgrader := &websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
		CheckOrigin:     func(r *http.Request) bool { return true },
	}

	hub := newHub()

	go hub.run()

	return &Melody{
		Config:                   newConfig(),
		Upgrader:                 upgrader,
		messageHandler:           func(*Session, []byte) {},
		messageHandlerBinary:     func(*Session, []byte) {},
		messageSentHandler:       func(*Session, []byte) {},
		messageSentHandlerBinary: func(*Session, []byte) {},
		errorHandler:             func(*Session, error) {},
		closeHandler:             nil,
		connectHandler:           func(*Session) {},
		disconnectHandler:        func(*Session) {},
		pongHandler:              func(*Session) {},
		hub:                      hub,
	}
}

// HandleConnect fires fn when a session connects.
func (m *Melody) HandleConnect(fn func(*Session)) {
	m.connectHandler = fn
}

// HandleDisconnect fires fn when a session disconnects.
func (m *Melody) HandleDisconnect(fn func(*Session)) {
	m.disconnectHandler = fn
}

// HandlePong fires fn when a pong is received from a session.
func (m *Melody) HandlePong(fn func(*Session)) {
	m.pongHandler = fn
}

// HandleMessage fires fn when a text message comes in.
func (m *Melody) HandleMessage(fn func(*Session, []byte)) {
	m.messageHandler = fn
}

// HandleMessageBinary fires fn when a binary message comes in.
func (m *Melody) HandleMessageBinary(fn func(*Session, []byte)) {
	m.messageHandlerBinary = fn
}

// HandleSentMessage fires fn when a text message is successfully sent.
func (m *Melody) HandleSentMessage(fn func(*Session, []byte)) {
	m.messageSentHandler = fn
}

// HandleSentMessageBinary fires fn when a binary message is successfully sent.
func (m *Melody) HandleSentMessageBinary(fn func(*Session, []byte)) {
	m.messageSentHandlerBinary = fn
}

// HandleError fires fn when a session has an error.
func (m *Melody) HandleError(fn func(*Session, error)) {
	m.errorHandler = fn
}

// HandleClose sets the handler for close messages received from the session.
// The code argument to h is the received close code or CloseNoStatusReceived
// if the close message is empty. The default close handler sends a close frame
// back to the session.
//
// The application must read the connection to process close messages as
// described in the section on Control Frames above.
//
// The connection read methods return a CloseError when a close frame is
// received. Most applications should handle close messages as part of their
// normal error handling. Applications should only set a close handler when the
// application must perform some action before sending a close frame back to
// the session.
func (m *Melody) HandleClose(fn func(*Session, int, string) error) {
	if fn != nil {
		m.closeHandler = fn
	}
}

// HandleRequest upgrades http requests to websocket connections and dispatches them to be handled by the melody instance.
func (m *Melody) HandleRequest(w http.ResponseWriter, r *http.Request) error {
	return m.HandleRequestWithKeys(w, r, nil)
}

// HandleRequestWithKeys does the same as HandleRequest but populates session.Keys with keys.
func (m *Melody) HandleRequestWithKeys(w http.ResponseWriter, r *http.Request, keys map[string]interface{}) error {
	if m.hub.closed() {
		return errors.New("melody instance is closed")
	}

	conn, err := m.Upgrader.Upgrade(w, r, w.Header())

	if err != nil {
		return err
	}

	session := &Session{
		Request: r,
		Keys:    keys,
		conn:    conn,
		output:  make(chan *envelope, m.Config.MessageBufferSize),
		melody:  m,
		open:    true,
		rwmutex: &sync.RWMutex{},
	}

	m.hub.register <- session

	m.connectHandler(session)

	go session.writePump()

	session.readPump()

	if !m.hub.closed() {
		m.hub.unregister <- session
	}

	session.close()

	m.disconnectHandler(session)

	return nil
}

// Broadcast broadcasts a text message to all sessions.
func (m *Melody) Broadcast(msg []byte) error {
	if m.hub.closed() {
		return errors.New("melody instance is closed")
	}

	message := &envelope{t: websocket.TextMessage, msg: msg}
	m.hub.broadcast <- message

	return nil
}

// BroadcastFilter broadcasts a text message to all sessions that fn returns true for.
func (m *Melody) BroadcastFilter(msg []byte, fn func(*Session) bool) error {
	if m.hub.closed() {
		return errors.New("melody instance is closed")
	}

	message := &envelope{t: websocket.TextMessage, msg: msg, filter: fn}
	m.hub.broadcast <- message

	return nil
}

// BroadcastOthers broadcasts a text message to all sessions except session s.
func (m *Melody) BroadcastOthers(msg []byte, s *Session) error {
	return m.BroadcastFilter(msg, func(q *Session) bool {
		return s != q
	})
}

// BroadcastMultiple broadcasts a text message to multiple sessions given in the sessions slice.
func (m *Melody) BroadcastMultiple(msg []byte, sessions []*Session) error {
	for _, sess := range sessions {
		if writeErr := sess.Write(msg); writeErr != nil {
			return writeErr
		}
	}
	return nil
}

// BroadcastBinary broadcasts a binary message to all sessions.
func (m *Melody) BroadcastBinary(msg []byte) error {
	if m.hub.closed() {
		return errors.New("melody instance is closed")
	}

	message := &envelope{t: websocket.BinaryMessage, msg: msg}
	m.hub.broadcast <- message

	return nil
}

// BroadcastBinaryFilter broadcasts a binary message to all sessions that fn returns true for.
func (m *Melody) BroadcastBinaryFilter(msg []byte, fn func(*Session) bool) error {
	if m.hub.closed() {
		return errors.New("melody instance is closed")
	}

	message := &envelope{t: websocket.BinaryMessage, msg: msg, filter: fn}
	m.hub.broadcast <- message

	return nil
}

// BroadcastBinaryOthers broadcasts a binary message to all sessions except session s.
func (m *Melody) BroadcastBinaryOthers(msg []byte, s *Session) error {
	return m.BroadcastBinaryFilter(msg, func(q *Session) bool {
		return s != q
	})
}

// Close closes the melody instance and all connected sessions.
func (m *Melody) Close() error {
	if m.hub.closed() {
		return errors.New("melody instance is already closed")
	}

	m.hub.exit <- &envelope{t: websocket.CloseMessage, msg: []byte{}}

	return nil
}

// CloseWithMsg closes the melody instance with the given close payload and all connected sessions.
// Use the FormatCloseMessage function to format a proper close message payload.
func (m *Melody) CloseWithMsg(msg []byte) error {
	if m.hub.closed() {
		return errors.New("melody instance is already closed")
	}

	m.hub.exit <- &envelope{t: websocket.CloseMessage, msg: msg}

	return nil
}

// Len return the number of connected sessions.
func (m *Melody) Len() int {
	return m.hub.len()
}

// IsClosed returns the status of the melody instance.
func (m *Melody) IsClosed() bool {
	return m.hub.closed()
}

// FormatCloseMessage formats closeCode and text as a WebSocket close message.
func FormatCloseMessage(closeCode int, text string) []byte {
	return websocket.FormatCloseMessage(closeCode, text)
}
<|endoftext|>


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

package logging

import (
	"sync"
	"sync/atomic"
	"unsafe"
)

// TODO pick one of the memory backends and stick with it or share interface.

// Node is a record node pointing to an optional next node.
type node struct {
	next   *node
	Record *Record
}

// Next returns the next record node. If there's no node available, it will
// return nil.
func (n *node) Next() *node {
	return n.next
}

// MemoryBackend is a simple memory based logging backend that will not produce
// any output but merly keep records, up to the given size, in memory.
type MemoryBackend struct {
	size       int32
	maxSize    int32
	head, tail unsafe.Pointer
}

// NewMemoryBackend creates a simple in-memory logging backend.
func NewMemoryBackend(size int) *MemoryBackend {
	return &MemoryBackend{maxSize: int32(size)}
}

// Log implements the Log method required by Backend.
func (b *MemoryBackend) Log(level Level, calldepth int, rec *Record) error {
	var size int32

	n := &node{Record: rec}
	np := unsafe.Pointer(n)

	// Add the record to the tail. If there's no records available, tail and
	// head will both be nil. When we successfully set the tail and the previous
	// value was nil, it's safe to set the head to the current value too.
	for {
		tailp := b.tail
		swapped := atomic.CompareAndSwapPointer(
			&b.tail,
			tailp,
			np,
		)
		if swapped == true {
			if tailp == nil {
				b.head = np
			} else {
				(*node)(tailp).next = n
			}
			size = atomic.AddInt32(&b.size, 1)
			break
		}
	}

	// Since one record was added, we might have overflowed the list. Remove
	// a record if that is the case. The size will fluctate a bit, but
	// eventual consistent.
	if b.maxSize > 0 && size > b.maxSize {
		for {
			headp := b.head
			head := (*node)(b.head)
			if head.next == nil {
				break
			}
			swapped := atomic.CompareAndSwapPointer(
				&b.head,
				headp,
				unsafe.Pointer(head.next),
			)
			if swapped == true {
				atomic.AddInt32(&b.size, -1)
				break
			}
		}
	}
	return nil
}

// Head returns the oldest record node kept in memory. It can be used to
// iterate over records, one by one, up to the last record.
//
// Note: new records can get added while iterating. Hence the number of records
// iterated over might be larger than the maximum size.
func (b *MemoryBackend) Head() *node {
	return (*node)(b.head)
}

type event int

const (
	eventFlush event = iota
	eventStop
)

// ChannelMemoryBackend is very similar to the MemoryBackend, except that it
// internally utilizes a channel.
type ChannelMemoryBackend struct {
	maxSize    int
	size       int
	incoming   chan *Record
	events     chan event
	mu         sync.Mutex
	running    bool
	flushWg    sync.WaitGroup
	stopWg     sync.WaitGroup
	head, tail *node
}

// NewChannelMemoryBackend creates a simple in-memory logging backend which
// utilizes a go channel for communication.
//
// Start will automatically be called by this function.
func NewChannelMemoryBackend(size int) *ChannelMemoryBackend {
	backend := &ChannelMemoryBackend{
		maxSize:  size,
		incoming: make(chan *Record, 1024),
		events:   make(chan event),
	}
	backend.Start()
	return backend
}

// Start launches the internal goroutine which starts processing data from the
// input channel.
func (b *ChannelMemoryBackend) Start() {
	b.mu.Lock()
	defer b.mu.Unlock()

	// Launch the goroutine unless it's already running.
	if b.running != true {
		b.running = true
		b.stopWg.Add(1)
		go b.process()
	}
}

func (b *ChannelMemoryBackend) process() {
	defer b.stopWg.Done()
	for {
		select {
		case rec := <-b.incoming:
			b.insertRecord(rec)
		case e := <-b.events:
			switch e {
			case eventStop:
				return
			case eventFlush:
				for len(b.incoming) > 0 {
					b.insertRecord(<-b.incoming)
				}
				b.flushWg.Done()
			}
		}
	}
}

func (b *ChannelMemoryBackend) insertRecord(rec *Record) {
	prev := b.tail
	b.tail = &node{Record: rec}
	if prev == nil {
		b.head = b.tail
	} else {
		prev.next = b.tail
	}

	if b.maxSize > 0 && b.size >= b.maxSize {
		b.head = b.head.next
	} else {
		b.size += 1
	}
}

// Flush waits until all records in the buffered channel have been processed.
func (b *ChannelMemoryBackend) Flush() {
	b.flushWg.Add(1)
	b.events <- eventFlush
	b.flushWg.Wait()
}

// Stop signals the internal goroutine to exit and waits until it have.
func (b *ChannelMemoryBackend) Stop() {
	b.mu.Lock()
	if b.running == true {
		b.running = false
		b.events <- eventStop
	}
	b.mu.Unlock()
	b.stopWg.Wait()
}

// Log implements the Log method required by Backend.
func (b *ChannelMemoryBackend) Log(level Level, calldepth int, rec *Record) error {
	b.incoming <- rec
	return nil
}

// Head returns the oldest record node kept in memory. It can be used to
// iterate over records, one by one, up to the last record.
//
// Note: new records can get added while iterating. Hence the number of records
// iterated over might be larger than the maximum size.
func (b *ChannelMemoryBackend) Head() *node {
	return b.head
}
<commit_msg>don't include memory backend on appengine<commit_after>// Copyright 2013, Örjan Persson. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build !appengine

package logging

import (
	"sync"
	"sync/atomic"
	"unsafe"
)

// TODO pick one of the memory backends and stick with it or share interface.

// Node is a record node pointing to an optional next node.
type node struct {
	next   *node
	Record *Record
}

// Next returns the next record node. If there's no node available, it will
// return nil.
func (n *node) Next() *node {
	return n.next
}

// MemoryBackend is a simple memory based logging backend that will not produce
// any output but merly keep records, up to the given size, in memory.
type MemoryBackend struct {
	size       int32
	maxSize    int32
	head, tail unsafe.Pointer
}

// NewMemoryBackend creates a simple in-memory logging backend.
func NewMemoryBackend(size int) *MemoryBackend {
	return &MemoryBackend{maxSize: int32(size)}
}

// Log implements the Log method required by Backend.
func (b *MemoryBackend) Log(level Level, calldepth int, rec *Record) error {
	var size int32

	n := &node{Record: rec}
	np := unsafe.Pointer(n)

	// Add the record to the tail. If there's no records available, tail and
	// head will both be nil. When we successfully set the tail and the previous
	// value was nil, it's safe to set the head to the current value too.
	for {
		tailp := b.tail
		swapped := atomic.CompareAndSwapPointer(
			&b.tail,
			tailp,
			np,
		)
		if swapped == true {
			if tailp == nil {
				b.head = np
			} else {
				(*node)(tailp).next = n
			}
			size = atomic.AddInt32(&b.size, 1)
			break
		}
	}

	// Since one record was added, we might have overflowed the list. Remove
	// a record if that is the case. The size will fluctate a bit, but
	// eventual consistent.
	if b.maxSize > 0 && size > b.maxSize {
		for {
			headp := b.head
			head := (*node)(b.head)
			if head.next == nil {
				break
			}
			swapped := atomic.CompareAndSwapPointer(
				&b.head,
				headp,
				unsafe.Pointer(head.next),
			)
			if swapped == true {
				atomic.AddInt32(&b.size, -1)
				break
			}
		}
	}
	return nil
}

// Head returns the oldest record node kept in memory. It can be used to
// iterate over records, one by one, up to the last record.
//
// Note: new records can get added while iterating. Hence the number of records
// iterated over might be larger than the maximum size.
func (b *MemoryBackend) Head() *node {
	return (*node)(b.head)
}

type event int

const (
	eventFlush event = iota
	eventStop
)

// ChannelMemoryBackend is very similar to the MemoryBackend, except that it
// internally utilizes a channel.
type ChannelMemoryBackend struct {
	maxSize    int
	size       int
	incoming   chan *Record
	events     chan event
	mu         sync.Mutex
	running    bool
	flushWg    sync.WaitGroup
	stopWg     sync.WaitGroup
	head, tail *node
}

// NewChannelMemoryBackend creates a simple in-memory logging backend which
// utilizes a go channel for communication.
//
// Start will automatically be called by this function.
func NewChannelMemoryBackend(size int) *ChannelMemoryBackend {
	backend := &ChannelMemoryBackend{
		maxSize:  size,
		incoming: make(chan *Record, 1024),
		events:   make(chan event),
	}
	backend.Start()
	return backend
}

// Start launches the internal goroutine which starts processing data from the
// input channel.
func (b *ChannelMemoryBackend) Start() {
	b.mu.Lock()
	defer b.mu.Unlock()

	// Launch the goroutine unless it's already running.
	if b.running != true {
		b.running = true
		b.stopWg.Add(1)
		go b.process()
	}
}

func (b *ChannelMemoryBackend) process() {
	defer b.stopWg.Done()
	for {
		select {
		case rec := <-b.incoming:
			b.insertRecord(rec)
		case e := <-b.events:
			switch e {
			case eventStop:
				return
			case eventFlush:
				for len(b.incoming) > 0 {
					b.insertRecord(<-b.incoming)
				}
				b.flushWg.Done()
			}
		}
	}
}

func (b *ChannelMemoryBackend) insertRecord(rec *Record) {
	prev := b.tail
	b.tail = &node{Record: rec}
	if prev == nil {
		b.head = b.tail
	} else {
		prev.next = b.tail
	}

	if b.maxSize > 0 && b.size >= b.maxSize {
		b.head = b.head.next
	} else {
		b.size += 1
	}
}

// Flush waits until all records in the buffered channel have been processed.
func (b *ChannelMemoryBackend) Flush() {
	b.flushWg.Add(1)
	b.events <- eventFlush
	b.flushWg.Wait()
}

// Stop signals the internal goroutine to exit and waits until it have.
func (b *ChannelMemoryBackend) Stop() {
	b.mu.Lock()
	if b.running == true {
		b.running = false
		b.events <- eventStop
	}
	b.mu.Unlock()
	b.stopWg.Wait()
}

// Log implements the Log method required by Backend.
func (b *ChannelMemoryBackend) Log(level Level, calldepth int, rec *Record) error {
	b.incoming <- rec
	return nil
}

// Head returns the oldest record node kept in memory. It can be used to
// iterate over records, one by one, up to the last record.
//
// Note: new records can get added while iterating. Hence the number of records
// iterated over might be larger than the maximum size.
func (b *ChannelMemoryBackend) Head() *node {
	return b.head
}
<|endoftext|>


<commit_before>package middleware

import (
	"net/http"
	"strconv"
	"strings"

	"github.com/aerogo/aero"
	"github.com/aerogo/http/client"
	"github.com/animenotifier/arn"
	"github.com/animenotifier/notify.moe/utils"
	"github.com/fatih/color"
	"github.com/mssola/user_agent"
)

// UserInfo updates user related information after each request.
func UserInfo() aero.Middleware {
	return func(ctx *aero.Context, next func()) {
		next()

		// Ignore non-HTML requests
		if aero.IsMediaType(ctx.Response().Header().Get("Content-Type")) {
			return
		}

		// Ignore API requests
		if strings.HasPrefix(ctx.URI(), "/api/") {
			return
		}

		user := utils.GetUser(ctx)

		// When there's no user logged in, nothing to update
		if user == nil {
			return
		}

		// This works asynchronously so it doesn't block the response
		go updateUserInfo(ctx, user)
	}
}

// Update browser and OS data
func updateUserInfo(ctx *aero.Context, user *arn.User) {
	newIP := ctx.RealIP()
	newUserAgent := ctx.UserAgent()

	if user.UserAgent != newUserAgent {
		user.UserAgent = newUserAgent

		// Parse user agent
		parsed := user_agent.New(user.UserAgent)

		// Browser
		user.Browser.Name, user.Browser.Version = parsed.Browser()

		// OS
		os := parsed.OSInfo()
		user.OS.Name = os.Name
		user.OS.Version = os.Version
	}

	if user.IP != newIP {
		updateUserLocation(user, newIP)
	}

	user.LastSeen = arn.DateTimeUTC()
	user.Save()
}

// Updates the location of the user.
func updateUserLocation(user *arn.User, newIP string) {
	user.IP = newIP
	locationAPI := "https://api.ipinfodb.com/v3/ip-city/?key=" + arn.APIKeys.IPInfoDB.ID + "&ip=" + user.IP + "&format=json"
	response, err := client.Get(locationAPI).End()

	if err != nil {
		color.Red("Couldn't fetch location data | Error: %s | IP: %s", err.Error(), user.IP)
		return
	}

	if response.StatusCode() != http.StatusOK {
		color.Red("Couldn't fetch location data | Status: %d | IP: %s", response.StatusCode, user.IP)
		return
	}

	newLocation := arn.IPInfoDBLocation{}
	response.Unmarshal(&newLocation)

	if newLocation.CountryName != "-" {
		user.Location.CountryName = newLocation.CountryName
		user.Location.CountryCode = newLocation.CountryCode
		user.Location.Latitude, _ = strconv.ParseFloat(newLocation.Latitude, 64)
		user.Location.Longitude, _ = strconv.ParseFloat(newLocation.Longitude, 64)
		user.Location.CityName = newLocation.CityName
		user.Location.RegionName = newLocation.RegionName
		user.Location.TimeZone = newLocation.TimeZone
		user.Location.ZipCode = newLocation.ZipCode
	}
}
<commit_msg>Updated userinfo middleware<commit_after>package middleware

import (
	"net/http"
	"strconv"
	"strings"

	"github.com/aerogo/aero"
	"github.com/aerogo/http/client"
	"github.com/animenotifier/arn"
	"github.com/animenotifier/notify.moe/utils"
	"github.com/fatih/color"
	"github.com/mssola/user_agent"
)

// UserInfo updates user related information after each request.
func UserInfo() aero.Middleware {
	return func(ctx *aero.Context, next func()) {
		next()

		// Ignore non-HTML requests
		contentType := ctx.Response().Header().Get("Content-Type")

		if !strings.HasPrefix(contentType, "text/html") {
			return
		}

		user := utils.GetUser(ctx)

		// When there's no user logged in, nothing to update
		if user == nil {
			return
		}

		// This works asynchronously so it doesn't block the response
		go updateUserInfo(ctx, user)
	}
}

// Update browser and OS data
func updateUserInfo(ctx *aero.Context, user *arn.User) {
	newIP := ctx.RealIP()
	newUserAgent := ctx.UserAgent()

	if user.UserAgent != newUserAgent {
		user.UserAgent = newUserAgent

		// Parse user agent
		parsed := user_agent.New(user.UserAgent)

		// Browser
		user.Browser.Name, user.Browser.Version = parsed.Browser()

		// OS
		os := parsed.OSInfo()
		user.OS.Name = os.Name
		user.OS.Version = os.Version
	}

	if user.IP != newIP {
		updateUserLocation(user, newIP)
	}

	user.LastSeen = arn.DateTimeUTC()
	user.Save()
}

// Updates the location of the user.
func updateUserLocation(user *arn.User, newIP string) {
	user.IP = newIP
	locationAPI := "https://api.ipinfodb.com/v3/ip-city/?key=" + arn.APIKeys.IPInfoDB.ID + "&ip=" + user.IP + "&format=json"
	response, err := client.Get(locationAPI).End()

	if err != nil {
		color.Red("Couldn't fetch location data | Error: %s | IP: %s", err.Error(), user.IP)
		return
	}

	if response.StatusCode() != http.StatusOK {
		color.Red("Couldn't fetch location data | Status: %d | IP: %s", response.StatusCode, user.IP)
		return
	}

	newLocation := arn.IPInfoDBLocation{}
	response.Unmarshal(&newLocation)

	if newLocation.CountryName != "-" {
		user.Location.CountryName = newLocation.CountryName
		user.Location.CountryCode = newLocation.CountryCode
		user.Location.Latitude, _ = strconv.ParseFloat(newLocation.Latitude, 64)
		user.Location.Longitude, _ = strconv.ParseFloat(newLocation.Longitude, 64)
		user.Location.CityName = newLocation.CityName
		user.Location.RegionName = newLocation.RegionName
		user.Location.TimeZone = newLocation.TimeZone
		user.Location.ZipCode = newLocation.ZipCode
	}
}
<|endoftext|>


<commit_before>package workpool

import (
	"context"
	"errors"
	"runtime"
	"sync"
	"time"
)

type WorkFunc func(context.Context)

type entry struct {
	ctx    context.Context
	work   WorkFunc
	time   time.Time
	period time.Duration
	seq    int64
	idx    int
}

type queue []*entry

type worker chan *entry

type WorkPool struct {
	ch     chan worker
	wg     sync.WaitGroup
	once   sync.Once
	ctx    context.Context
	cancel context.CancelFunc
}

func NewPool(size int) *WorkPool {
	if size <= 0 {
		size = runtime.NumCPU()
	}
	ctx, cancel := context.WithCancel(context.TODO())
	return &WorkPool{
		ch:     make(chan worker, size),
		ctx:    ctx,
		cancel: cancel,
	}
}

func (w worker) start(pool *WorkPool) {
	defer func() {
		pool.wg.Done()
		close(w)
	}()

	for {
		select {
		case <-pool.ctx.Done():
			return
		default:
			pool.ch <- w
		}

		select {
		case e := <-w:
			e.work(e.ctx)
		case <-pool.ctx.Done():
			return
		}
	}
}

var ErrClosed = errors.New("workpool has been closed")

func (p *WorkPool) Submit(ctx context.Context, work WorkFunc) (err error) {
	if ctx == nil {
		panic("nil context")
	}

	p.once.Do(func() {
		p.wg.Add(cap(p.ch))
		for i := 0; i < cap(p.ch); i++ {
			w := make(worker)
			go w.start(p)
		}
	})

	select {
	case w := <-p.ch:
		select {
		case <-p.ctx.Done():
			err = ErrClosed
		default:
			w <- &entry{ctx: ctx, work: work}
		}
	case <-ctx.Done():
		err = ctx.Err()
	case <-p.ctx.Done():
		err = ErrClosed
	}
	return
}

func (p *WorkPool) Context() context.Context {
	return p.ctx
}

func (p *WorkPool) Close() (err error) {
	select {
	case <-p.ctx.Done():
		err = ErrClosed
	default:
		p.cancel()
		p.wg.Wait()
		close(p.ch)
	}
	return
}

func (q queue) Len() int {
	return len(q)
}

func (q queue) Less(i, j int) bool {
	switch {
	case q[i].time.Before(q[j].time):
		return true
	case q[i].time.After(q[j].time):
		return false
	case q[i].seq < q[j].seq:
		return true
	default:
		return false
	}
}

func (q queue) Swap(i, j int) {
	q[i], q[j] = q[j], q[i]
	q[i].idx = i
	q[j].idx = j
}

func (q *queue) Push(x interface{}) {
	n := len(*q)
	e := x.(*entry)
	e.idx = n
	*q = append(*q, e)
}

func (q *queue) Pop() interface{} {
	old := *q
	n := len(old)
	e := old[n-1]
	e.idx = -1
	*q = old[0 : n-1]
	return e
}

var defaultPool = NewPool(0)

func Submit(ctx context.Context, work WorkFunc) error {
	return defaultPool.Submit(ctx, work)
}
<commit_msg>implement the Schedule function<commit_after>package workpool

import (
	"container/heap"
	"context"
	"errors"
	"runtime"
	"sync"
	"sync/atomic"
	"time"
)

type WorkFunc func(context.Context)

type task struct {
	ctx    context.Context
	f      WorkFunc
	when   time.Time
	period time.Duration
	seq    uint64
	idx    int
}

type queue []*task

type WorkPool struct {
	l      sync.Mutex
	n      int32
	r      int64
	t      *time.Timer
	q      queue
	qwait  sync.Cond
	ctx    context.Context
	cancel context.CancelFunc
}

func NewPool(size int) *WorkPool {
	if size <= 0 {
		size = runtime.NumCPU()
	}
	p := new(WorkPool)
	p.n = int32(size)
	p.t = time.NewTimer(0)
	p.qwait.L = &p.l
	p.ctx, p.cancel = context.WithCancel(context.Background())

	<-p.t.C

	go func() {
		for {
			select {
			case <-p.ctx.Done():
				return
			case <-p.t.C:
				p.l.Lock()
				p.qwait.Signal()
				p.l.Unlock()
			}
		}
	}()

	return p
}

var seq uint64

func (p *WorkPool) Submit(ctx context.Context, work WorkFunc) error {
	return p.do(&task{ctx: ctx, f: work, when: time.Now(), seq: atomic.AddUint64(&seq, 1)})
}

func (p *WorkPool) Schedule(ctx context.Context, work WorkFunc, delay, period time.Duration) error {
	return p.do(&task{ctx: ctx, f: work, when: time.Now().Add(delay), period: period, seq: atomic.AddUint64(&seq, 1)})
}

var ErrClosed = errors.New("workpool has been closed")

func (p *WorkPool) do(t *task) error {
	select {
	case <-t.ctx.Done():
		return t.ctx.Err()
	case <-p.ctx.Done():
		return ErrClosed
	default:
	}

	p.l.Lock()
	heap.Push(&p.q, t)
	p.qwait.Signal()
	p.l.Unlock()

	for {
		n := atomic.LoadInt32(&p.n)
		if n <= 0 {
			return nil
		}
		if atomic.CompareAndSwapInt32(&p.n, n, n-1) {
			break
		}
	}

	go func() {
		var t *task
		for {
			p.l.Lock()
		wait:
			for {
				select {
				case <-p.ctx.Done():
					p.l.Unlock()
					return
				default:
				}
				if p.q.Len() > 0 {
					t = p.q[0]
					d := t.when.Sub(time.Now())
					if d <= 0 {
						heap.Remove(&p.q, 0)
						p.r = 0
						break wait
					}
					r := t.when.UnixNano()
					if p.r == 0 || p.r > r {
						p.t.Reset(d)
						p.r = r
					}
				}
				p.qwait.Wait()
			}
			if p.q.Len() > 0 {
				p.qwait.Signal()
			}
			p.l.Unlock()
			if t != nil {
				t.f(t.ctx)
				if t.period > 0 {
					t.when = t.when.Add(t.period)
					p.do(t)
				}
			}
		}
	}()

	return nil
}

func (p *WorkPool) Context() context.Context {
	return p.ctx
}

func (p *WorkPool) Close() (err error) {
	select {
	case <-p.ctx.Done():
		err = ErrClosed
	default:
		p.cancel()
		p.l.Lock()
		p.t.Stop()
		p.qwait.Broadcast()
		p.l.Unlock()
	}
	return
}

func (q queue) Len() int {
	return len(q)
}

func (q queue) Less(i, j int) bool {
	switch {
	case q[i].when.Before(q[j].when):
		return true
	case q[i].when.After(q[j].when):
		return false
	case q[i].seq < q[j].seq:
		return true
	default:
		return false
	}
}

func (q queue) Swap(i, j int) {
	q[i], q[j] = q[j], q[i]
	q[i].idx = i
	q[j].idx = j
}

func (q *queue) Push(x interface{}) {
	n := len(*q)
	t := x.(*task)
	t.idx = n
	*q = append(*q, t)
}

func (q *queue) Pop() interface{} {
	old := *q
	n := len(old)
	t := old[n-1]
	t.idx = -1
	*q = old[0 : n-1]
	return t
}
<|endoftext|>